diff --git a/docs/architecture.md b/docs/architecture.md index 0fc30f0b..36be0ea4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -920,6 +920,7 @@ configurations so workstreams can use different LLM backends. [models.local] base_url = "http://localhost:8000/v1" model = "qwen3-32b" +max_concurrency = 1 # provider defaults to "openai" [models.claude] @@ -958,6 +959,20 @@ binding on mismatch. Failed reload validation changes neither maps nor generation. If only non-transport fields changed, compatible client pools can remain warm, but the session still receives a new coherent config/lane. +**Per-alias admission:** Each registry alias owns a stable, hot-resizable +`ModelAdmission`. `max_concurrency = 0` is unlimited; a positive value limits +simultaneous generations for that alias in one process. Every registry-backed +role carries the same gate on its `ModelLane`, so main turns, judges, task +agents, perception, compaction, and background generation coordinate through +one FIFO. Two aliases never share a gate implicitly, even when their URLs are +identical. `model_turn()` materializes attachment fallbacks before admission, +then holds one lease across eager stream creation and the complete drain, +releasing before retry backoff. Admission wait is credited out of deadline +accounting, preventing queued judges from spending their request budget before +dispatch. The gate survives cap-only reloads in place; the field is excluded +from semantic `ModelConfig` equality so a capacity edit does not reset judges +or output-guard state. + Primary loops, recursive compaction, judges, title generation, audio, and task agents all consume `ModelLane` rather than inspecting provider/client handles. Fallback is a lane change, so retry classification and result provenance come @@ -1196,8 +1211,9 @@ Verified quirks of vLLM's Anthropic endpoint: **Database model definitions:** On server entry points, models can also be defined in the `model_definitions` table (admin Models tab). DB models support -the same per-model sampling overrides. Config.toml models override DB models -with the same alias in-memory (the DB rows are never modified). +the same per-model sampling overrides and per-alias `max_concurrency`. +Config.toml models override DB models with the same alias in-memory (the DB +rows are never modified). **Lifecycle:** 1. `load_model_registry()` loads DB model definitions (if storage available), diff --git a/docs/console.md b/docs/console.md index 15998e22..6fca14d8 100644 --- a/docs/console.md +++ b/docs/console.md @@ -455,7 +455,10 @@ configuration editor. The **Channels** tab links users to either a Discord or Slack account via a per-row channel-type selector. The **Models** tab is a CRUD editor for `model_definitions`, including static and dynamic backend-auth -modes. Model edits rebind existing workstreams at their next send while +modes and a per-process **Max concurrent generations** limit for each alias +(`0` means unlimited). The limit is shared by every model-backed role using +that alias and a streaming generation holds its slot through the full decode. +Model edits rebind existing workstreams at their next send while in-flight requests keep their original definition snapshot; see [Settings](settings.md#model-definition-reloads) for the full contract. The **Nodes** tab edits per-node metadata, and the **TLS** tab manages CA and leaf certificates for the diff --git a/docs/diagrams/03-core-engine-classes.puml b/docs/diagrams/03-core-engine-classes.puml index 8545b539..d2cbd288 100644 --- a/docs/diagrams/03-core-engine-classes.puml +++ b/docs/diagrams/03-core-engine-classes.puml @@ -134,6 +134,7 @@ class "ModelLane" as ModelLane <> { + capabilities: ModelCapabilities + extra_params: dict | None + registry: ModelRegistry | None + + admission: ModelAdmission | None + backend_auth_config: ModelConfig | None + backend_auth_resolver: Callable | None } @@ -357,12 +358,13 @@ class "ModelRegistry" as ModelReg { - _models: dict[str, ModelConfig] - _clients: dict[str, Any] - _providers: dict[str, LLMProvider] + - _admissions: dict[str, ModelAdmission] - _client_lock: Lock + default: str + fallback: list[str] + agent_model: str | None -- - + resolve_binding(alias) → (client, model, config, provider, generation) + + resolve_binding(alias) → (client, model, config, provider, admission, generation) + get_client(alias) → Any + get_provider(alias) → LLMProvider + has_alias(alias) → bool @@ -376,6 +378,22 @@ class "ModelRegistry" as ModelReg { core/model_registry.py } +class "ModelAdmission" as ModelAdmission { + - alias: str + - _limit: int + - _in_flight: int + - _waiters: deque + + acquire(cancel_ref) → AdmissionLease + + set_limit(limit) + + snapshot() → AdmissionSnapshot + -- + Per-process FIFO generation gate. + Stable across alias hot reloads; + queue time is deadline credit. + -- + core/admission.py +} + class "ModelConfig" as ModelCfg <> { + alias: str + provider: str @@ -385,6 +403,7 @@ class "ModelConfig" as ModelCfg <> { + temperature: float | None + max_tokens: int | None + reasoning_effort: str | None + + max_concurrency: int + auth_mode: str + obo_audience: str + obo_scopes: str @@ -475,12 +494,14 @@ KindAdapter ..> ChatSession : constructs ModelReg --> "*" ModelCfg : holds ModelReg --> "*" LLMProvider : caches +ModelReg --> "*" ModelAdmission : owns per alias LLMProvider --> ModelCaps : returns ModelReg --> ResolvedBinding : resolves atomically ResolvedBinding --> ModelLane ModelLane --> LLMProvider ModelLane --> ModelCaps ModelLane --> ModelCfg : auth/config snapshot +ModelLane --> ModelAdmission : admission lease ModelTurnFn --> ModelLane ModelTurnFn --> ModelTurnResult ModelTurnFn ..> BackendAuth : per-call resolver diff --git a/docs/diagrams/04-conversation-turn.puml b/docs/diagrams/04-conversation-turn.puml index 03e0ec78..8e2fa414 100644 --- a/docs/diagrams/04-conversation-turn.puml +++ b/docs/diagrams/04-conversation-turn.puml @@ -10,6 +10,7 @@ participant "SessionManager" as Manager participant "ChatSession" as Session participant "SessionUIBase" as UI participant "model_turn()\n+ lowering" as Plant +participant "ModelAdmission\n(per alias)" as Admission participant "LLM provider" as Provider participant "Tool workers" as Tools database "StorageBackend\n(SQLite / PostgreSQL)" as Storage @@ -54,6 +55,9 @@ loop until final answer and no queued input Session -> Plant : model_turn(active ModelLane, Turns,\n tools, cancel_ref, on_chunk) activate Plant Plant -> Plant : canonical Turns → provider wire\nrestore ids + repair + lane-specific fold + Plant -> Plant : materialize attachment refs\n(nested perception before outer slot) + Plant -> Admission : acquire(cancel_ref) + activate Admission Plant -> Plant : resolve per-call backend credential\nfrom lane's pinned ModelConfig Plant -> Provider : create_streaming(...) activate Provider @@ -68,6 +72,8 @@ loop until final answer and no queued input Provider --> Plant : finish + usage + native blocks deactivate Provider Plant -> Plant : drain + re-ingest assistant Turn\nwith serving-lane provenance + Plant -> Admission : release before retry backoff + deactivate Admission Plant --> Session : ModelTurnResult deactivate Plant diff --git a/docs/diagrams/png/03-core-engine-classes.png b/docs/diagrams/png/03-core-engine-classes.png index 810c4143..11a2e909 100644 --- a/docs/diagrams/png/03-core-engine-classes.png +++ b/docs/diagrams/png/03-core-engine-classes.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e1431edf3891785b922c52b7897e3af5d39ba9973a815f892d6afdb762c5297b -size 612662 +oid sha256:c74e99c530c3a8af9ab35b1e4d8c4fef0ea35c0c04cc35da7cf3588e71382057 +size 661175 diff --git a/docs/diagrams/png/04-conversation-turn.png b/docs/diagrams/png/04-conversation-turn.png index 1f7a331a..c08b8035 100644 --- a/docs/diagrams/png/04-conversation-turn.png +++ b/docs/diagrams/png/04-conversation-turn.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:79299b25ccc10484af13684a89ed9457abb9bc1781604bf6a9000ccb84362e55 -size 311107 +oid sha256:b237e00c28225c847b3e1a083c5c35535d1bf1549f6e6c7e2b6c4fd0e037a980 +size 331906 diff --git a/docs/settings.md b/docs/settings.md index 51111744..5efd49f8 100644 --- a/docs/settings.md +++ b/docs/settings.md @@ -54,6 +54,23 @@ When a per-model override is `NULL` (empty in the UI), the global default is used. Switching models via `/model ` re-resolves sampling parameters from the new model's overrides or global defaults. +### Per-model concurrency + +Each model definition may set `max_concurrency` to limit simultaneous model +generations for that alias in one Turnstone process. `0` or an omitted value +means unlimited. The gate is shared by every role using the alias—interactive +turns, coordinators, task agents, judges, output guards, perception, compaction, +and title generation—and a streaming generation holds its slot until the +stream is fully drained or closed. + +Admission is strictly per alias. Two aliases remain independent even when they +point to the same URL; Turnstone does not infer shared capacity from endpoint +text. Queue time is excluded from judge/output-guard deadline accounting, and +each retry releases its slot before backoff and reacquires for the next wire +attempt. The cap is local to each process, not cluster-wide; account for the +number of nodes targeting the same inference server. Direct STT/TTS protocol +calls and Cohere/Jina reranking do not currently consume this generation cap. + ### Model backend authentication Model definitions support four backend credential modes: @@ -445,6 +462,9 @@ send. Endpoint, provider, backend model ID, capabilities, extra parameters, and backend-auth configuration are replaced as one immutable binding. In-flight turns, judges, and task agents finish or cancel against the binding they started with; an admin edit never tears one request across two definitions. +The alias's admission gate is retained and resized in place, so a concurrency +edit preserves in-flight accounting and does not reset cached judges or the +output-guard rate limiter. Sampling and other saved workstream configuration remain workstream state. A model-definition edit does not silently rewrite a live workstream's chosen diff --git a/sdk/typescript/openapi-console.json b/sdk/typescript/openapi-console.json index 862aeb11..3e2300d4 100644 --- a/sdk/typescript/openapi-console.json +++ b/sdk/typescript/openapi-console.json @@ -11880,6 +11880,14 @@ "title": "Context Window", "type": "integer" }, + "max_concurrency": { + "default": 0, + "description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.", + "maximum": 2147483647, + "minimum": 0, + "title": "Max Concurrency", + "type": "integer" + }, "capabilities": { "default": "{}", "title": "Capabilities", @@ -12015,6 +12023,14 @@ "title": "Context Window", "type": "integer" }, + "max_concurrency": { + "default": 0, + "description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.", + "maximum": 2147483647, + "minimum": 0, + "title": "Max Concurrency", + "type": "integer" + }, "capabilities": { "default": "{}", "title": "Capabilities", @@ -12151,6 +12167,14 @@ "title": "Context Window", "type": "integer" }, + "max_concurrency": { + "default": 0, + "description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.", + "maximum": 2147483647, + "minimum": 0, + "title": "Max Concurrency", + "type": "integer" + }, "capabilities": { "additionalProperties": true, "title": "Capabilities", @@ -12304,6 +12328,13 @@ "default": null, "title": "Context Window" }, + "max_concurrency": { + "description": "Maximum concurrent model generations for this alias in one process; zero means unlimited.", + "maximum": 2147483647, + "minimum": 0, + "title": "Max Concurrency", + "type": "integer" + }, "capabilities": { "additionalProperties": true, "default": null, diff --git a/tests/test_admin_model_registry_refresh.py b/tests/test_admin_model_registry_refresh.py index 307caf71..86aaffc4 100644 --- a/tests/test_admin_model_registry_refresh.py +++ b/tests/test_admin_model_registry_refresh.py @@ -92,6 +92,7 @@ def _seed_model_def( obo_audience: str = "", obo_scopes: str = "", capabilities: str = "{}", + max_concurrency: int = 0, ) -> None: """Insert a model definition row directly via the storage API.""" storage.create_model_definition( @@ -108,6 +109,7 @@ def _seed_model_def( auth_mode=auth_mode, obo_audience=obo_audience, obo_scopes=obo_scopes, + max_concurrency=max_concurrency, ) @@ -2431,6 +2433,91 @@ def test_update_endpoint_refreshes_registry(storage: SQLiteBackend) -> None: assert registry.get_config("local").model == "new-model" +def test_create_and_update_max_concurrency_refresh_registry(storage: SQLiteBackend) -> None: + _seed_model_def(storage, definition_id="m1", alias="local", model="m") + registry = _make_registry(alias="local", model="m") + client = _make_client(storage, registry) + + created = client.post( + "/v1/api/admin/model-definitions", + json={"alias": "limited", "model": "m2", "max_concurrency": 3}, + ) + assert created.status_code == 200, created.text + assert created.json()["max_concurrency"] == 3 + assert registry.get_config("limited").max_concurrency == 3 + + definition_id = created.json()["definition_id"] + updated = client.put( + f"/v1/api/admin/model-definitions/{definition_id}", + json={"max_concurrency": 0}, + ) + assert updated.status_code == 200, updated.text + assert updated.json()["max_concurrency"] == 0 + assert registry.get_config("limited").max_concurrency == 0 + + +def test_update_omission_preserves_max_concurrency(storage: SQLiteBackend) -> None: + _seed_model_def( + storage, + definition_id="m1", + alias="local", + model="m", + max_concurrency=2, + ) + registry = _make_registry(alias="local", model="m") + client = _make_client(storage, registry) + + resp = client.put( + "/v1/api/admin/model-definitions/m1", + json={"model": "m2"}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["max_concurrency"] == 2 + assert storage.get_model_definition("m1")["max_concurrency"] == 2 + + +@pytest.mark.parametrize("invalid", [None, True, "1", 1.0, -1, 2_147_483_648]) +def test_create_rejects_invalid_max_concurrency( + storage: SQLiteBackend, + invalid: Any, +) -> None: + client = _make_client(storage, _make_registry()) + + resp = client.post( + "/v1/api/admin/model-definitions", + json={"alias": "invalid", "model": "m", "max_concurrency": invalid}, + ) + + assert resp.status_code == 400, resp.text + assert "max_concurrency" in resp.json()["error"] + assert storage.get_model_definition_by_alias("invalid") is None + + +@pytest.mark.parametrize("invalid", [None, True, "1", 1.0, -1, 2_147_483_648]) +def test_update_rejects_invalid_max_concurrency( + storage: SQLiteBackend, + invalid: Any, +) -> None: + _seed_model_def( + storage, + definition_id="m1", + alias="local", + model="m", + max_concurrency=2, + ) + client = _make_client(storage, _make_registry(alias="local", model="m")) + + resp = client.put( + "/v1/api/admin/model-definitions/m1", + json={"max_concurrency": invalid}, + ) + + assert resp.status_code == 400, resp.text + assert "max_concurrency" in resp.json()["error"] + assert storage.get_model_definition("m1")["max_concurrency"] == 2 + + def test_update_endpoint_skips_refresh_on_empty_body( storage: SQLiteBackend, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -3809,6 +3896,7 @@ def test_every_mutable_column_probes_its_classification( "base_url": "https://other.example/v1", "api_key": "sk-new", "context_window": 4096, + "max_concurrency": 2, "capabilities": {"note": "probe"}, # Arm direction: the loop seeds THIS row disabled, so the probe is the # gated false→true flip rather than the carved-out disarm. diff --git a/tests/test_app_js.py b/tests/test_app_js.py index 0bfa9a55..d5212c2d 100644 --- a/tests/test_app_js.py +++ b/tests/test_app_js.py @@ -909,6 +909,29 @@ def test_model_response_controls_are_capability_driven_and_sparse() -> None: assert 'apiSurfEl.addEventListener("change", _onModelFieldChange)' in admin +def test_model_max_concurrency_form_round_trips_strict_integer() -> None: + html = _CONSOLE_INDEX.read_text(encoding="utf-8") + admin = _CONSOLE_ADMIN_JS.read_text(encoding="utf-8") + + assert 'id="model-max-concurrency"' in html + assert 'max="2147483647"' in html + assert "0 = unlimited" in html + + create = _slice_function_body(admin, "showCreateModelModal") + edit = _slice_function_body(admin, "showEditModelModal") + render = _slice_function_body(admin, "_renderModels") + assert create is not None and edit is not None and render is not None + assert 'getElementById("model-max-concurrency").value = "0"' in create + assert "m.max_concurrency != null ? m.max_concurrency : 0" in edit + # submitCreateModel is longer than the balanced-slice helper's bounded + # window; these names are unique to that form path, so whole-file pins are + # both stable and unambiguous. + assert "Number.isInteger(maxConcurrency)" in admin + assert "maxConcurrency > 2147483647" in admin + assert "form.max_concurrency = maxConcurrency" in admin + assert 'overrides.push("limit=" + m.max_concurrency)' in render + + def test_shared_utils_defines_set_safe_html_helper() -> None: """``setSafeHtml`` in ``shared/utils.js`` is the single audited entry point for installing trusted HTML strings into a DOM element outside diff --git a/tests/test_deadline.py b/tests/test_deadline.py index 2860cdaa..b9a0ad22 100644 --- a/tests/test_deadline.py +++ b/tests/test_deadline.py @@ -16,6 +16,8 @@ import pytest from turnstone.core.deadline import ( DeadlineCancelledError, DeadlineExceededError, + StreamAbortRef, + run_abortable_with_deadline, run_with_deadline, ) @@ -111,6 +113,70 @@ def test_on_abandon_errors_do_not_mask_the_deadline_error() -> None: ) +def test_abortable_deadline_credits_only_active_admission_wait() -> None: + def _work(ref: StreamAbortRef) -> str: + ref.begin_admission_wait() + time.sleep(0.12) + ref.end_admission_wait() + time.sleep(0.02) + return "ok" + + start = time.monotonic() + assert ( + run_abortable_with_deadline( + _work, + timeout=0.05, + poll=0.005, + thread_name="dl-admission-credit", + ) + == "ok" + ) + assert time.monotonic() - start >= 0.12 + + +def test_provider_time_still_expires_after_admission_credit() -> None: + provider_started = threading.Event() + + def _work(ref: StreamAbortRef) -> None: + ref.begin_admission_wait() + time.sleep(0.08) + ref.end_admission_wait() + provider_started.set() + time.sleep(1.0) + + with pytest.raises(DeadlineExceededError): + run_abortable_with_deadline( + _work, + timeout=0.05, + poll=0.005, + thread_name="dl-provider-after-admission", + ) + + # The admission interval was credited (the call reached provider work), + # but that work consumed the unchanged logical deadline and timed out. + assert provider_started.is_set() + + +def test_dispatch_marker_is_observability_not_a_clock_reset() -> None: + captured: list[StreamAbortRef] = [] + + def _work(ref: StreamAbortRef) -> None: + captured.append(ref) + time.sleep(0.06) + ref.mark_dispatch() + time.sleep(0.08) + + with pytest.raises(DeadlineExceededError): + run_abortable_with_deadline( + _work, + timeout=0.1, + poll=0.005, + thread_name="dl-dispatch-marker", + ) + assert captured[0].dispatch_count == 1 + assert captured[0].last_dispatch_at is not None + + class TestStreamAbortRef: def test_abort_closes_captured_stream(self) -> None: from unittest.mock import MagicMock diff --git a/tests/test_judge.py b/tests/test_judge.py index 295185c7..bb53c4a5 100644 --- a/tests/test_judge.py +++ b/tests/test_judge.py @@ -5,18 +5,24 @@ from __future__ import annotations import json import threading import time +from dataclasses import replace from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from unittest.mock import MagicMock from tests._session_helpers import as_stream from tests._session_helpers import mock_completion_result as _mock_result +from turnstone.core.admission import ModelAdmission from turnstone.core.judge import IntentJudge, IntentVerdict, JudgeConfig, evaluate_heuristic +from turnstone.core.model_backend_auth import BackendAuthUnavailableError from turnstone.core.model_registry import ModelConfig from turnstone.core.model_turn import ModelLane, ResolvedModelBinding -from turnstone.core.providers._protocol import ModelCapabilities +from turnstone.core.providers._protocol import IncompleteStreamError, ModelCapabilities from turnstone.core.trajectory import Role +if TYPE_CHECKING: + import pytest + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -412,33 +418,140 @@ class TestCancelEventSemantics: assert all(v.tier == "llm" for v in results) assert provider.create_streaming.call_count == 3 - def test_batch_backend_auth_resolves_once_and_reuses_token(self): - """One judge batch owns one delegated credential snapshot.""" + def test_queued_backend_auth_resolves_after_admission_per_attempt( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A queued judge never ages and reuses a pre-admission token.""" + monkeypatch.setattr("turnstone.core.model_turn._DRAIN_RETRY_BASE_DELAY", 0.0) provider = _make_mock_provider(_good_verdict_json()) + provider.retryable_error_names = frozenset({"IncompleteStreamError"}) + + def _incomplete_stream() -> Any: + yield from () + raise IncompleteStreamError("retry after admission") + + provider.create_streaming.side_effect = [ + _incomplete_stream(), + as_stream(_mock_result(_good_verdict_json())), + ] judge = _make_judge(provider) + admission = ModelAdmission("judge", 1) + auth_config = MagicMock(name="pinned-auth-config") + judge._lane = replace( + judge._lane, + alias="judge", + admission=admission, + backend_auth_config=auth_config, + ) + batch_client = MagicMock() - bound_client = object() - batch_client.with_options.return_value = bound_client + batch_client.with_options.side_effect = lambda *, api_key: f"client:{api_key}" judge._create_client = MagicMock(return_value=batch_client) # type: ignore[method-assign] - resolver = MagicMock(return_value="user-a-token") + + token_epoch = "stale" + resolutions: list[tuple[str, Any, str]] = [] + + def _resolve(alias: str, config: ModelConfig | None) -> str: + token = f"{token_epoch}-{len(resolutions) + 1}" + resolutions.append((alias, config, token)) + return token + results: list[IntentVerdict] = [] - items = [_make_item(call_id=f"tc_{i}") for i in range(2)] + done = threading.Event() + holder = admission.acquire() judge.evaluate( - items, + [_make_item()], [{"role": "user", "content": "test"}], results.append, - backend_auth_resolver=resolver, + done_callback=done.set, + backend_auth_resolver=_resolve, ) - _wait_for(results, 2) - resolver.assert_called_once_with("", None) + deadline = time.monotonic() + 2.0 + while admission.snapshot().queued != 1 and time.monotonic() < deadline: + time.sleep(0.005) + queued = admission.snapshot().queued == 1 + resolutions_while_queued = list(resolutions) + token_epoch = "fresh" + holder.release() + + assert done.wait(5.0) + assert queued + assert resolutions_while_queued == [] + assert resolutions == [ + ("judge", auth_config, "fresh-1"), + ("judge", auth_config, "fresh-2"), + ] assert batch_client.with_options.call_count == 2 - assert all( - call.kwargs["client"] is bound_client - for call in provider.create_streaming.call_args_list + assert [call.kwargs["api_key"] for call in batch_client.with_options.call_args_list] == [ + "fresh-1", + "fresh-2", + ] + assert [call.kwargs["client"] for call in provider.create_streaming.call_args_list] == [ + "client:fresh-1", + "client:fresh-2", + ] + assert len(results) == 1 + assert results[0].tier == "llm" + + def test_lazy_auth_failure_falls_back_remaining_batch_without_dispatch(self, caplog) -> None: + """A fail-closed mint after admission aborts the whole judge batch.""" + provider = _make_mock_provider(_good_verdict_json()) + judge = _make_judge(provider) + admission = ModelAdmission("judge", 1) + auth_config = MagicMock(name="pinned-auth-config") + judge._lane = replace( + judge._lane, + alias="judge", + admission=admission, + backend_auth_config=auth_config, ) - assert [verdict.call_id for verdict in results] == ["tc_0", "tc_1"] + + batch_client = MagicMock() + judge._create_client = MagicMock(return_value=batch_client) # type: ignore[method-assign] + resolver = MagicMock(side_effect=BackendAuthUnavailableError("mint unavailable")) + items = [_make_item(call_id=f"tc_{i}") for i in range(3)] + results: list[IntentVerdict] = [] + done = threading.Event() + holder = admission.acquire() + + with caplog.at_level("ERROR", logger="turnstone.core.judge"): + judge.evaluate( + items, + [{"role": "user", "content": "test"}], + results.append, + done_callback=done.set, + backend_auth_resolver=resolver, + ) + + deadline = time.monotonic() + 2.0 + while admission.snapshot().queued != 1 and time.monotonic() < deadline: + time.sleep(0.005) + queued = admission.snapshot().queued == 1 + resolutions_while_queued = resolver.call_count + holder.release() + finished = done.wait(5.0) + + assert queued + assert resolutions_while_queued == 0 + assert finished + resolver.assert_called_once_with("judge", auth_config) + batch_client.with_options.assert_not_called() + provider.create_streaming.assert_not_called() + assert admission.snapshot().in_flight == 0 + assert [verdict.call_id for verdict in results] == ["tc_0", "tc_1", "tc_2"] + assert all(verdict.tier == "llm_fallback" for verdict in results) + assert all( + "judge backend authentication failed" in verdict.reasoning for verdict in results + ) + auth_logs = [ + record + for record in caplog.records + if "Judge backend authentication failed" in record.getMessage() + ] + assert len(auth_logs) == 1 def test_cancelled_batch_skips_backend_auth_resolution(self): """The daemon checks cancellation before doing a credential mint.""" diff --git a/tests/test_migration_070.py b/tests/test_migration_070.py new file mode 100644 index 00000000..b28e8482 --- /dev/null +++ b/tests/test_migration_070.py @@ -0,0 +1,68 @@ +"""Migration coverage for per-alias model concurrency.""" + +from __future__ import annotations + +from pathlib import Path + +import sqlalchemy as sa +from alembic import command +from alembic.config import Config + +_MIGRATIONS_DIR = str( + Path(__file__).resolve().parent.parent / "turnstone" / "core" / "storage" / "migrations" +) + + +def _alembic_cfg(db_path: Path) -> Config: + cfg = Config() + cfg.set_main_option("script_location", _MIGRATIONS_DIR) + cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}") + return cfg + + +class TestMigration070: + def test_upgrade_defaults_preexisting_rows_to_unlimited(self, tmp_path: Path) -> None: + db_path = tmp_path / "070-up.db" + cfg = _alembic_cfg(db_path) + command.upgrade(cfg, "069") + engine = sa.create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as conn: + conn.execute( + sa.text( + "INSERT INTO model_definitions " + "(definition_id, alias, model, created, updated) " + "VALUES ('d1', 'local', 'm', " + "'2026-01-01T00:00:00', '2026-01-01T00:00:00')" + ) + ) + command.upgrade(cfg, "070") + with engine.connect() as conn: + value = conn.execute( + sa.text( + "SELECT max_concurrency FROM model_definitions WHERE definition_id = 'd1'" + ) + ).scalar_one() + assert value == 0 + finally: + engine.dispose() + + def test_downgrade_then_upgrade_round_trip(self, tmp_path: Path) -> None: + db_path = tmp_path / "070-roundtrip.db" + cfg = _alembic_cfg(db_path) + command.upgrade(cfg, "070") + command.downgrade(cfg, "069") + engine = sa.create_engine(f"sqlite:///{db_path}") + try: + columns = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")} + assert "max_concurrency" not in columns + finally: + engine.dispose() + + command.upgrade(cfg, "070") + engine = sa.create_engine(f"sqlite:///{db_path}") + try: + columns = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")} + assert "max_concurrency" in columns + finally: + engine.dispose() diff --git a/tests/test_model_admission.py b/tests/test_model_admission.py new file mode 100644 index 00000000..3fbb8ee7 --- /dev/null +++ b/tests/test_model_admission.py @@ -0,0 +1,199 @@ +"""Focused contract tests for per-alias model admission.""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import pytest + +import turnstone.core.admission as admission_mod +from turnstone.core.admission import ModelAdmission +from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef +from turnstone.core.model_registry import ( + KEY_GUARD_DEFERRED_TO_LIFESPAN, + ModelConfig, + ModelRegistry, +) + + +def _wait_until(predicate: Any, *, timeout: float = 2.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError("condition did not become true") + time.sleep(0.005) + + +def _config(alias: str, limit: int, *, base_url: str | None = None) -> ModelConfig: + return ModelConfig( + alias=alias, + base_url=base_url or f"http://{alias}.example/v1", + api_key="test", + model="model", + max_concurrency=limit, + ) + + +def test_unlimited_holders_are_counted_and_live_narrowing_drains() -> None: + gate = ModelAdmission("primary", 0) + first = gate.acquire() + second = gate.acquire() + assert gate.snapshot().in_flight == 2 + + gate.set_limit(1) + acquired = threading.Event() + release_waiter = threading.Event() + + def _waiter() -> None: + with gate.acquire(): + acquired.set() + release_waiter.wait(2.0) + + thread = threading.Thread(target=_waiter, daemon=True) + thread.start() + _wait_until(lambda: gate.snapshot().queued == 1) + + first.release() + assert not acquired.wait(0.05) + second.release() + assert acquired.wait(1.0) + + release_waiter.set() + thread.join(1.0) + assert not thread.is_alive() + assert gate.snapshot().in_flight == 0 + + +def test_fifo_waiters_and_hot_widening() -> None: + gate = ModelAdmission("primary", 1) + original = gate.acquire() + acquired_order: list[int] = [] + acquired = [threading.Event(), threading.Event()] + releases = [threading.Event(), threading.Event()] + + def _waiter(index: int) -> None: + with gate.acquire(): + acquired_order.append(index) + acquired[index].set() + releases[index].wait(2.0) + + threads: list[threading.Thread] = [] + for index in range(2): + thread = threading.Thread(target=_waiter, args=(index,), daemon=True) + threads.append(thread) + thread.start() + _wait_until(lambda expected=index + 1: gate.snapshot().queued == expected) + + gate.set_limit(2) + assert acquired[0].wait(1.0) + assert not acquired[1].wait(0.05) + + original.release() + assert acquired[1].wait(1.0) + assert acquired_order == [0, 1] + + for event in releases: + event.set() + for thread in threads: + thread.join(1.0) + assert not thread.is_alive() + + +def test_cancelled_waiter_is_removed_and_never_admitted() -> None: + gate = ModelAdmission("primary", 1) + holder = gate.acquire() + cancel_ref = StreamAbortRef() + finished = threading.Event() + errors: list[BaseException] = [] + + def _waiter() -> None: + try: + gate.acquire(cancel_ref=cancel_ref) + except BaseException as exc: # test records the worker's exact exit + errors.append(exc) + finally: + finished.set() + + thread = threading.Thread(target=_waiter, daemon=True) + thread.start() + _wait_until(lambda: gate.snapshot().queued == 1) + cancel_ref.abort() + + assert finished.wait(1.0) + assert len(errors) == 1 + assert isinstance(errors[0], DeadlineCancelledError) + assert gate.snapshot().queued == 0 + holder.release() + assert gate.snapshot().in_flight == 0 + + +def test_registry_keeps_one_gate_per_alias_across_resize_remove_and_readd() -> None: + registry = ModelRegistry( + { + "alpha": _config("alpha", 1, base_url="http://shared.example/v1"), + "beta": _config("beta", 3, base_url="http://shared.example/v1"), + }, + default="alpha", + ) + alpha = registry.get_admission("alpha") + beta = registry.get_admission("beta") + assert alpha is not beta + + registry.reload( + {"alpha": _config("alpha", 2)}, + default="alpha", + app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN, + ) + assert registry.get_admission("alpha") is alpha + assert alpha.limit == 2 + + registry.reload( + {}, + default="", + app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN, + ) + registry.reload( + {"alpha": _config("alpha", 4)}, + default="alpha", + app_state=KEY_GUARD_DEFERRED_TO_LIFESPAN, + ) + assert registry.get_admission("alpha") is alpha + assert alpha.limit == 4 + + +def test_wait_stall_and_resize_logs_expose_queue_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[tuple[str, dict[str, Any]]] = [] + monkeypatch.setattr(admission_mod, "_CANCEL_POLL_SECONDS", 0.005) + monkeypatch.setattr(admission_mod, "_STALL_WARNING_SECONDS", 0.01) + monkeypatch.setattr( + admission_mod, + "log", + SimpleNamespace( + info=lambda event, **fields: events.append((event, fields)), + warning=lambda event, **fields: events.append((event, fields)), + ), + ) + gate = ModelAdmission("alpha", 1) + holder = gate.acquire() + + thread = threading.Thread(target=lambda: gate.acquire().release(), daemon=True) + thread.start() + _wait_until(lambda: gate.snapshot().queued == 1) + _wait_until(lambda: any(event == "model.admission_stalled" for event, _ in events)) + holder.release() + thread.join(1.0) + gate.set_limit(2) + + by_name = {event: fields for event, fields in events} + assert by_name["model.admission_wait"]["alias"] == "alpha" + assert by_name["model.admission_wait"]["queued_ahead"] == 0 + assert by_name["model.admission_stalled"]["in_flight"] == 1 + assert by_name["model.admission_stalled"]["queued"] == 1 + assert by_name["model.admission_resized"]["previous_limit"] == 1 + assert by_name["model.admission_resized"]["limit"] == 2 diff --git a/tests/test_model_admission_live.py b/tests/test_model_admission_live.py new file mode 100644 index 00000000..81ae983e --- /dev/null +++ b/tests/test_model_admission_live.py @@ -0,0 +1,462 @@ +"""Live request-count assertions for per-alias model admission. + +These tests exercise the real OpenAI-compatible backend configured by the +same environment contract as :mod:`tests.test_server_live`. A local threaded +reverse proxy sits between ``model_turn`` and that backend solely to count +HTTP requests. It buffers each complete upstream response, then relays the +unchanged SSE payload to Turnstone's real streaming provider path. + +Run explicitly (a backend must be listening at ``TURNSTONE_TEST_BASE_URL``, +which defaults to ``http://localhost:8000/v1``):: + + pytest tests/test_model_admission_live.py -m live -v +""" + +from __future__ import annotations + +import dataclasses +import json +import os +import threading +import time +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import TYPE_CHECKING, Any + +import httpx +import pytest +from openai import OpenAI + +from turnstone.core.model_registry import ModelConfig, ModelRegistry +from turnstone.core.model_turn import ModelLane, model_turn, resolve_model_binding +from turnstone.core.trajectory import Turn + +if TYPE_CHECKING: + from types import TracebackType + + +_LIVE_BASE_URL = os.environ.get("TURNSTONE_TEST_BASE_URL", "http://localhost:8000/v1") +_LIVE_API_KEY = os.environ.get("TURNSTONE_TEST_API_KEY", "not-needed") or "not-needed" +_HOP_BY_HOP_HEADERS = frozenset( + { + "connection", + "content-encoding", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + } +) + + +@dataclass(frozen=True) +class _ObservedCounts: + requests: int + forwarded: int + completed: int + failed: int + active: int + peak: int + forward_active: int + forward_peak: int + requests_by_alias: dict[str, int] + forwarded_by_alias: dict[str, int] + completed_by_alias: dict[str, int] + peak_by_alias: dict[str, int] + forward_peak_by_alias: dict[str, int] + rendezvous_timed_out: bool + + +class _ConcurrencyCounter: + """Thread-safe request counts with one bounded overlap rendezvous. + + The first ``rendezvous_size`` admitted requests wait for each other before + any is forwarded upstream. A short settling hold after the rendezvous + makes an over-admission regression observable as a peak above the cap, + independent of how quickly the live backend generates the tiny response. + """ + + def __init__( + self, + *, + rendezvous_size: int, + rendezvous_timeout: float = 3.0, + settling_hold: float = 0.5, + ) -> None: + self._lock = threading.Lock() + self._rendezvous = threading.Event() + self._rendezvous_size = rendezvous_size + self._rendezvous_timeout = rendezvous_timeout + self._settling_hold = settling_hold + self._requests = 0 + self._forwarded = 0 + self._completed = 0 + self._failed = 0 + self._active = 0 + self._peak = 0 + self._forward_active = 0 + self._forward_peak = 0 + self._active_by_alias: dict[str, int] = defaultdict(int) + self._forward_active_by_alias: dict[str, int] = defaultdict(int) + self._requests_by_alias: dict[str, int] = defaultdict(int) + self._forwarded_by_alias: dict[str, int] = defaultdict(int) + self._completed_by_alias: dict[str, int] = defaultdict(int) + self._peak_by_alias: dict[str, int] = defaultdict(int) + self._forward_peak_by_alias: dict[str, int] = defaultdict(int) + self._rendezvous_timed_out = False + + def enter(self, alias: str) -> None: + with self._lock: + self._requests += 1 + self._requests_by_alias[alias] += 1 + self._active += 1 + self._active_by_alias[alias] += 1 + self._peak = max(self._peak, self._active) + self._peak_by_alias[alias] = max( + self._peak_by_alias[alias], self._active_by_alias[alias] + ) + if self._active >= self._rendezvous_size: + self._rendezvous.set() + + if not self._rendezvous.wait(self._rendezvous_timeout): + # Release this and every later request so a broken cap fails by + # count instead of leaving live-test worker threads parked. + with self._lock: + self._rendezvous_timed_out = True + self._rendezvous.set() + time.sleep(self._settling_hold) + + def begin_forward(self, alias: str) -> None: + """Record one HTTP request entering the real upstream backend.""" + with self._lock: + self._forwarded += 1 + self._forwarded_by_alias[alias] += 1 + self._forward_active += 1 + self._forward_active_by_alias[alias] += 1 + self._forward_peak = max(self._forward_peak, self._forward_active) + self._forward_peak_by_alias[alias] = max( + self._forward_peak_by_alias[alias], + self._forward_active_by_alias[alias], + ) + + def end_forward(self, alias: str) -> None: + with self._lock: + self._forward_active -= 1 + self._forward_active_by_alias[alias] -= 1 + + def leave(self, alias: str, *, completed: bool) -> None: + with self._lock: + self._active -= 1 + self._active_by_alias[alias] -= 1 + if completed: + self._completed += 1 + self._completed_by_alias[alias] += 1 + else: + self._failed += 1 + + def snapshot(self) -> _ObservedCounts: + with self._lock: + return _ObservedCounts( + requests=self._requests, + forwarded=self._forwarded, + completed=self._completed, + failed=self._failed, + active=self._active, + peak=self._peak, + forward_active=self._forward_active, + forward_peak=self._forward_peak, + requests_by_alias=dict(self._requests_by_alias), + forwarded_by_alias=dict(self._forwarded_by_alias), + completed_by_alias=dict(self._completed_by_alias), + peak_by_alias=dict(self._peak_by_alias), + forward_peak_by_alias=dict(self._forward_peak_by_alias), + rendezvous_timed_out=self._rendezvous_timed_out, + ) + + +class _CountingProxyServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, counter: _ConcurrencyCounter) -> None: + super().__init__(("127.0.0.1", 0), _CountingProxyHandler) + self.counter = counter + self.upstream_base_url = _LIVE_BASE_URL.rstrip("/") + self.upstream_authorization = f"Bearer {_LIVE_API_KEY}" + self.alias_by_authorization: dict[str, str] = {} + + +class _CountingProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + @property + def _proxy(self) -> _CountingProxyServer: + if not isinstance(self.server, _CountingProxyServer): + raise TypeError("counting handler requires _CountingProxyServer") + return self.server + + def do_POST(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API + alias = self._proxy.alias_by_authorization.get(self.headers.get("Authorization", "")) + if alias is None: + self._write_json_error(401, "unknown live-test alias credential") + return + + raw_length = self.headers.get("Content-Length", "0") + try: + content_length = int(raw_length) + except ValueError: + self._write_json_error(400, "invalid content length") + return + body = self.rfile.read(content_length) + counter = self._proxy.counter + counter.enter(alias) + completed = False + try: + headers = { + key: value + for key, value in self.headers.items() + if key.lower() not in _HOP_BY_HOP_HEADERS + and key.lower() not in {"authorization", "host"} + } + # The temporary per-alias credentials are observation tags only; + # the real backend sees exactly the configured live credential. + headers["Authorization"] = self._proxy.upstream_authorization + headers["Accept-Encoding"] = "identity" + counter.begin_forward(alias) + try: + upstream = httpx.post( + f"{self._proxy.upstream_base_url}{self.path}", + content=body, + headers=headers, + timeout=httpx.Timeout(60.0, connect=5.0), + ) + finally: + counter.end_forward(alias) + payload = upstream.content + self.send_response(upstream.status_code) + for key, value in upstream.headers.multi_items(): + if key.lower() not in _HOP_BY_HOP_HEADERS: + self.send_header(key, value) + self.send_header("Content-Length", str(len(payload))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(payload) + self.wfile.flush() + completed = upstream.status_code < 400 + except Exception: + # The worker observes a normal 502; its exception and the failed + # count retain the useful signal without leaking backend details. + self._write_json_error(502, "live backend proxy failure") + finally: + counter.leave(alias, completed=completed) + + def log_message(self, _format: str, *args: Any) -> None: + """Keep expected local proxy traffic out of pytest output.""" + + def _write_json_error(self, status: int, detail: str) -> None: + payload = json.dumps({"error": {"message": detail}}).encode() + try: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(payload) + self.wfile.flush() + except (BrokenPipeError, ConnectionError): + pass + + +class _CountingProxy: + def __init__(self, counter: _ConcurrencyCounter) -> None: + self._server = _CountingProxyServer(counter) + self._thread = threading.Thread( + target=self._server.serve_forever, + name="model-admission-live-proxy", + daemon=True, + ) + + @property + def base_url(self) -> str: + address = self._server.server_address + host = address[0] + port = address[1] + if not isinstance(host, str) or not isinstance(port, int): + raise TypeError("counting proxy did not bind an IPv4 TCP address") + return f"http://{host}:{port}" + + def credential_for(self, alias: str) -> str: + credential = f"turnstone-live-admission-{alias}" + self._server.alias_by_authorization[f"Bearer {credential}"] = alias + return credential + + def __enter__(self) -> _CountingProxy: + self._thread.start() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + del exc_type, exc_value, traceback + self._server.shutdown() + self._server.server_close() + self._thread.join(timeout=5.0) + if self._thread.is_alive(): + raise AssertionError("live counting proxy did not stop") + + +@pytest.fixture(scope="module") +def live_model_id() -> str: + """Auto-detect the model using the established live-backend contract.""" + client = OpenAI(base_url=_LIVE_BASE_URL, api_key=_LIVE_API_KEY) + try: + models = client.models.list() + finally: + client.close() + ids = [model.id for model in models.data] + assert ids, "No models found on the live backend" + return ids[0] + + +def _live_lane( + proxy: _CountingProxy, + model: str, + *, + alias: str, + limit: int, +) -> tuple[ModelRegistry, ModelLane]: + config = ModelConfig( + alias=alias, + base_url=proxy.base_url, + api_key=proxy.credential_for(alias), + model=model, + provider="openai-compatible", + max_concurrency=limit, + ) + registry = ModelRegistry({alias: config}, default=alias) + try: + binding = resolve_model_binding(registry, alias) + # Bound each live call and disable SDK retries: the proxy's exact + # request count should equal the number of model_turn invocations, + # while model_turn still exercises the real provider and admission + # lifecycle. + client = binding.lane.client.with_options(timeout=60.0, max_retries=0) + return registry, dataclasses.replace(binding.lane, client=client) + except BaseException: + registry.shutdown() + raise + + +def _run_parallel(lanes: list[ModelLane], model: str) -> None: + barrier = threading.Barrier(len(lanes)) + + def _one_turn(lane: ModelLane) -> str: + barrier.wait(timeout=10.0) + result = model_turn( + lane, + [Turn.user("Reply with OK.")], + max_tokens=8, + ) + return result.serving_model + + with ThreadPoolExecutor(max_workers=len(lanes)) as pool: + futures = [pool.submit(_one_turn, lane) for lane in lanes] + observed_models = [future.result(timeout=75.0) for future in futures] + assert observed_models == [model] * len(lanes) + + +@pytest.mark.live +class TestLiveModelAdmissionCounts: + """Count genuine backend requests at the per-alias admission boundary.""" + + def test_alias_cap_bounds_exact_peak_and_completes_every_request( + self, live_model_id: str + ) -> None: + counter = _ConcurrencyCounter(rendezvous_size=2) + with _CountingProxy(counter) as proxy: + registry, lane = _live_lane(proxy, live_model_id, alias="limited", limit=2) + try: + _run_parallel([lane, lane, lane, lane], live_model_id) + finally: + try: + lane.client.close() + finally: + registry.shutdown() + + counts = counter.snapshot() + assert counts.requests == 4 + assert counts.forwarded == 4 + assert counts.completed == 4 + assert counts.failed == 0 + assert counts.active == 0 + assert counts.peak == 2 + assert counts.forward_active == 0 + assert counts.forward_peak == 2 + assert counts.requests_by_alias == {"limited": 4} + assert counts.forwarded_by_alias == {"limited": 4} + assert counts.completed_by_alias == {"limited": 4} + assert counts.peak_by_alias == {"limited": 2} + assert counts.forward_peak_by_alias == {"limited": 2} + assert counts.rendezvous_timed_out is False + + def test_aliases_sharing_endpoint_have_independent_caps(self, live_model_id: str) -> None: + counter = _ConcurrencyCounter(rendezvous_size=2) + with _CountingProxy(counter) as proxy: + configs = { + alias: ModelConfig( + alias=alias, + base_url=proxy.base_url, + api_key=proxy.credential_for(alias), + model=live_model_id, + provider="openai-compatible", + max_concurrency=1, + ) + for alias in ("alpha", "beta") + } + # The two registry aliases deliberately name the exact same URL; + # both requests are then forwarded to the same physical backend. + assert configs["alpha"].base_url == configs["beta"].base_url + registry = ModelRegistry(configs, default="alpha") + lanes: list[ModelLane] = [] + try: + alpha = resolve_model_binding(registry, "alpha").lane + beta = resolve_model_binding(registry, "beta").lane + alpha = dataclasses.replace( + alpha, client=alpha.client.with_options(timeout=60.0, max_retries=0) + ) + lanes.append(alpha) + beta = dataclasses.replace( + beta, client=beta.client.with_options(timeout=60.0, max_retries=0) + ) + lanes.append(beta) + _run_parallel(lanes, live_model_id) + finally: + try: + for lane in lanes: + lane.client.close() + finally: + registry.shutdown() + + counts = counter.snapshot() + assert counts.requests == 2 + assert counts.forwarded == 2 + assert counts.completed == 2 + assert counts.failed == 0 + assert counts.active == 0 + assert counts.peak == 2 + assert counts.forward_active == 0 + assert counts.forward_peak == 2 + assert counts.requests_by_alias == {"alpha": 1, "beta": 1} + assert counts.forwarded_by_alias == {"alpha": 1, "beta": 1} + assert counts.completed_by_alias == {"alpha": 1, "beta": 1} + assert counts.peak_by_alias == {"alpha": 1, "beta": 1} + assert counts.forward_peak_by_alias == {"alpha": 1, "beta": 1} + assert counts.rendezvous_timed_out is False diff --git a/tests/test_model_definition_storage.py b/tests/test_model_definition_storage.py index bc44776b..7a85dc4c 100644 --- a/tests/test_model_definition_storage.py +++ b/tests/test_model_definition_storage.py @@ -33,6 +33,7 @@ class TestModelDefinitionStorage: assert m["base_url"] == "https://api.openai.com/v1" assert m["api_key"] == "sk-test" assert m["context_window"] == 128000 + assert m["max_concurrency"] == 0 assert m["capabilities"] == "{}" assert m["enabled"] is True @@ -158,6 +159,7 @@ class TestModelDefinitionStorage: assert m["base_url"] == "" assert m["api_key"] == "" assert m["context_window"] == 32768 + assert m["max_concurrency"] == 0 assert m["capabilities"] == "{}" assert m["enabled"] is True assert m["created_by"] == "" @@ -166,6 +168,23 @@ class TestModelDefinitionStorage: assert m["max_tokens"] is None assert m["reasoning_effort"] is None + def test_create_update_and_list_max_concurrency(self, db: SQLiteBackend) -> None: + did = _make_id() + db.create_model_definition( + definition_id=did, + alias="limited", + model="gpt-5", + max_concurrency=3, + ) + assert db.get_model_definition(did)["max_concurrency"] == 3 + + assert db.update_model_definition(did, max_concurrency=1) + assert db.get_model_definition_by_alias("limited")["max_concurrency"] == 1 + assert db.list_model_definitions()[0]["max_concurrency"] == 1 + + assert db.update_model_definition(did, max_concurrency=0) + assert db.get_model_definition(did)["max_concurrency"] == 0 + def test_create_with_sampling_params(self, db: SQLiteBackend) -> None: did = _make_id() db.create_model_definition( diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 718bdb47..c751e34e 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -17,6 +17,7 @@ from turnstone.core import model_registry as mr_module from turnstone.core.model_registry import ( KEY_GUARD_DEFERRED_TO_LIFESPAN, DynamicAuthKeyError, + ModelConcurrencyConfigError, ModelConfig, ModelRegistry, UnknownModelAliasError, @@ -50,6 +51,7 @@ class TestModelConfig: assert cfg.alias == "local" assert cfg.model == "qwen3-32b" assert cfg.context_window == 32768 # default + assert cfg.max_concurrency == 0 def test_custom_context_window(self) -> None: cfg = ModelConfig( @@ -112,6 +114,16 @@ class TestModelConfig: assert cfg.surface_persisted_reasoning is False assert cfg.replay_reasoning_to_model is True + def test_max_concurrency_is_strict_and_not_binding_identity(self) -> None: + unlimited = ModelConfig(alias="x", base_url="x", api_key="x", model="x") + limited = dataclasses.replace(unlimited, max_concurrency=2) + assert limited.max_concurrency == 2 + assert unlimited == limited + + for invalid in (-1, 2_147_483_648, True, 1.0, "1", None): + with pytest.raises(ModelConcurrencyConfigError, match="max_concurrency"): + dataclasses.replace(unlimited, max_concurrency=invalid) # type: ignore[arg-type] + # --------------------------------------------------------------------------- # ModelRegistry @@ -386,6 +398,7 @@ class TestLoadModelRegistry: "api_key": "sk-test", "model": "gpt-4o", "context_window": 128000, + "max_concurrency": 2, }, }, "model": { @@ -404,8 +417,26 @@ class TestLoadModelRegistry: assert reg.has_alias("openai") assert not reg.has_alias("default") assert reg.default == "openai" - _, model, _, _ = reg.resolve() + _, model, cfg, _ = reg.resolve() assert model == "gpt-4o" + assert cfg.max_concurrency == 2 + + @pytest.mark.parametrize("invalid", [-1, 2_147_483_648, True, 1.0, "1", None]) + def test_config_rejects_invalid_max_concurrency(self, invalid: Any) -> None: + fake_cfg = { + "models": { + "local": { + "base_url": "http://localhost:8000/v1", + "model": "m", + "max_concurrency": invalid, + } + } + } + with ( + patch("turnstone.core.model_registry.load_config", return_value=fake_cfg), + pytest.raises(ModelConcurrencyConfigError, match="max_concurrency"), + ): + load_model_registry() def test_config_context_window_zero_inherits_detected(self) -> None: """``context_window = 0`` in a [models.*] entry is the auto-detect @@ -794,6 +825,7 @@ class TestLoadModelRegistryWithDB: "temperature": 1.5, "max_tokens": 4096, "reasoning_effort": "high", + "max_concurrency": 4, } ] ) @@ -803,6 +835,7 @@ class TestLoadModelRegistryWithDB: assert cfg.temperature == 1.5 assert cfg.max_tokens == 4096 assert cfg.reasoning_effort == "high" + assert cfg.max_concurrency == 4 def test_db_sampling_params_null_means_none(self) -> None: """NULL sampling params in DB map to None (use global default).""" @@ -3638,8 +3671,8 @@ class TestSessionAgentModel: session._run_agent([Turn.user("x")], label="task", agent_alias="fast") assert captured["model"] == "fast-model" - def test_session_fallback_uses_exact_primary_lane_and_caps(self) -> None: - """A sub-agent inherits the primary lane with auth already consumed.""" + def test_session_fallback_uses_exact_primary_lane_and_pinned_auth(self) -> None: + """A sub-agent keeps the primary lane and pins its auth resolver.""" import turnstone.core.session as session_module reg = self._three_model_registry() # no agent_model / plan_model set @@ -3658,8 +3691,11 @@ class TestSessionAgentModel: assert model_turn_spy.call_count == 1 used_lane = model_turn_spy.call_args.args[0] - assert used_lane == dataclasses.replace(primary_lane, backend_auth_resolver=None) - assert used_lane.backend_auth_resolver is None + assert used_lane == dataclasses.replace( + primary_lane, + backend_auth_resolver=used_lane.backend_auth_resolver, + ) + assert used_lane.backend_auth_resolver is not None assert used_lane.capabilities is primary_caps assert used_lane.client is _client(session) assert used_lane.alias == "main" diff --git a/tests/test_model_turn.py b/tests/test_model_turn.py index 7106a0ea..5644f16d 100644 --- a/tests/test_model_turn.py +++ b/tests/test_model_turn.py @@ -12,6 +12,8 @@ import ast import inspect import logging import textwrap +import threading +import time from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock @@ -39,7 +41,7 @@ from turnstone.core.providers._protocol import ( serialized_tool_chars, ) from turnstone.core.session import ChatSession -from turnstone.core.trajectory import Role, ToolCall, Turn +from turnstone.core.trajectory import AttachmentRef, Role, ToolCall, Turn class _FakeProvider: @@ -230,6 +232,230 @@ def test_resolve_model_binding_canonicalizes_empty_alias_to_default() -> None: assert binding.registry_generation == 7 +def test_model_turn_materializes_before_admission_and_mints_inside_hold() -> None: + order: list[str] = [] + + class _Gate: + held = False + + def acquire(self, *, cancel_ref: Any = None) -> Any: + del cancel_ref + order.append("acquire") + gate = self + + class _Lease: + def __enter__(self) -> None: + gate.held = True + order.append("enter") + + def __exit__(self, *_exc: object) -> None: + gate.held = False + order.append("release") + + return _Lease() + + gate = _Gate() + bound_client = object() + base_client = MagicMock() + base_client.with_options.return_value = bound_client + + def _resolve(ids: list[str]) -> dict[str, Any]: + assert ids == ["image-1"] + assert not gate.held + order.append("materialize") + return { + "image-1": { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc"}, + } + } + + def _auth(_alias: str, _cfg: Any) -> str: + assert gate.held + order.append("auth") + return "minted-token" + + class _Provider(_FakeProvider): + def create_streaming(self, **kwargs: Any) -> Any: + assert gate.held + assert kwargs["client"] is bound_client + assert kwargs["resolve_attachments"] is None + order.append("dispatch") + self.calls.append(kwargs) + + def _stream() -> Any: + assert gate.held + order.append("drain") + yield from as_stream(CompletionResult(content="ok")) + + return _stream() + + provider = _Provider([]) + lane = ModelLane( + provider=provider, + client=base_client, + model="m", + alias="primary", + backend_auth_resolver=_auth, + admission=gate, # type: ignore[arg-type] + ) + turn = Turn(Role.USER, (AttachmentRef(attachment_id="image-1", kind="image"),)) + + result = model_turn(lane, [turn], resolve_attachments=_resolve) + + assert result.content == "ok" + assert order == ["materialize", "acquire", "enter", "auth", "dispatch", "drain", "release"] + assert result.wire_msgs == [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc"}, + } + ], + } + ] + + +def test_model_turn_releases_admission_before_retry_backoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from turnstone.core.deadline import StreamAbortRef + + order: list[str] = [] + + class _Gate: + held = False + acquire_calls = 0 + + def acquire(self, *, cancel_ref: Any = None) -> Any: + del cancel_ref + self.acquire_calls += 1 + gate = self + + class _Lease: + def __enter__(self) -> None: + assert not gate.held + gate.held = True + order.append("enter") + + def __exit__(self, *_exc: object) -> None: + gate.held = False + order.append("release") + + return _Lease() + + gate = _Gate() + provider = _FlakyProvider([IncompleteStreamError("retry"), CompletionResult(content="ok")]) + dispatch = provider.create_streaming + + def _dispatch(**kwargs: Any) -> Any: + assert gate.held + order.append("dispatch") + return dispatch(**kwargs) + + provider.create_streaming = _dispatch # type: ignore[method-assign] + + def _sleep(_delay: float) -> None: + assert not gate.held + order.append("backoff") + + monkeypatch.setattr(model_turn_mod, "time", SimpleNamespace(sleep=_sleep)) + lane = ModelLane( + provider=provider, + client=object(), + model="m", + admission=gate, # type: ignore[arg-type] + ) + ref = StreamAbortRef() + + result = model_turn(lane, [Turn.user("x")], cancel_ref=ref) + + assert result.content == "ok" + assert gate.acquire_calls == 2 + assert ref.dispatch_count == 2 + assert order == ["enter", "dispatch", "release", "backoff", "enter", "dispatch", "release"] + + +def test_same_alias_attachment_work_completes_before_outer_admission() -> None: + from turnstone.core.admission import ModelAdmission + + gate = ModelAdmission("primary", 1) + nested_completed = False + + def _resolve(ids: list[str]) -> dict[str, Any]: + nonlocal nested_completed + assert ids == ["image-1"] + # Models a nested perception call using the same alias. This would + # block forever if the outer model_turn had already taken the slot. + with gate.acquire(): + nested_completed = True + return {"image-1": {"type": "image_url", "image_url": {"url": "data:x"}}} + + provider = _FakeProvider([CompletionResult(content="ok")]) + lane = ModelLane(provider=provider, client=object(), model="m", admission=gate) + turn = Turn(Role.USER, (AttachmentRef(attachment_id="image-1", kind="image"),)) + + assert model_turn(lane, [turn], resolve_attachments=_resolve).content == "ok" + assert nested_completed + assert gate.snapshot().in_flight == 0 + + +def test_model_turn_releases_admission_when_eager_create_fails() -> None: + from turnstone.core.admission import ModelAdmission + + class _CreateFailureProvider(_FakeProvider): + def create_streaming(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + raise RuntimeError("connect failed") + + gate = ModelAdmission("primary", 1) + provider = _CreateFailureProvider([]) + lane = ModelLane(provider=provider, client=object(), model="m", admission=gate) + + with pytest.raises(RuntimeError, match="connect failed"): + model_turn(lane, [Turn.user("x")]) + + assert len(provider.calls) == 1 + assert gate.snapshot().in_flight == 0 + + +def test_model_turn_cancelled_while_queued_never_dispatches() -> None: + from turnstone.core.admission import ModelAdmission + from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef + + gate = ModelAdmission("primary", 1) + holder = gate.acquire() + provider = _FakeProvider([CompletionResult(content="never")]) + lane = ModelLane(provider=provider, client=object(), model="m", admission=gate) + cancel_ref = StreamAbortRef() + errors: list[BaseException] = [] + + def _call() -> None: + try: + model_turn(lane, [Turn.user("x")], cancel_ref=cancel_ref) + except BaseException as exc: # test records the worker's exact exit + errors.append(exc) + + thread = threading.Thread(target=_call, daemon=True) + thread.start() + deadline = time.monotonic() + 1.0 + while gate.snapshot().queued != 1: + if time.monotonic() >= deadline: + holder.release() + raise AssertionError("model turn did not queue") + time.sleep(0.005) + cancel_ref.abort() + thread.join(1.0) + holder.release() + + assert not thread.is_alive() + assert len(errors) == 1 + assert isinstance(errors[0], DeadlineCancelledError) + assert provider.calls == [] + + class _FlakyProvider: """Scripted drain-time deaths: each script entry is either a ``CompletionResult`` (streamed normally) or an exception instance diff --git a/tests/test_openapi.py b/tests/test_openapi.py index e79d1f69..c0b8f5c9 100644 --- a/tests/test_openapi.py +++ b/tests/test_openapi.py @@ -297,6 +297,26 @@ class TestConsoleSpec: ) +def test_model_max_concurrency_schema_is_strict_non_nullable_integer() -> None: + from turnstone.api.console_spec import build_console_spec + + schemas = build_console_spec()["components"]["schemas"] + for name in ("ModelDefinitionInfo", "CreateModelDefinitionRequest"): + prop = schemas[name]["properties"]["max_concurrency"] + assert prop["type"] == "integer" + assert prop["default"] == 0 + assert prop["minimum"] == 0 + assert prop["maximum"] == 2_147_483_647 + assert "anyOf" not in prop + + update_prop = schemas["UpdateModelDefinitionRequest"]["properties"]["max_concurrency"] + assert update_prop["type"] == "integer" + assert update_prop["minimum"] == 0 + assert update_prop["maximum"] == 2_147_483_647 + assert "anyOf" not in update_prop + assert "default" not in update_prop + + class TestCheckedInArtifactFreshness: """The checked-in `sdk/typescript/*.json` specs must match their source. diff --git a/tests/test_perception.py b/tests/test_perception.py index b0b685dc..3b4e7e55 100644 --- a/tests/test_perception.py +++ b/tests/test_perception.py @@ -19,8 +19,8 @@ class _StubProvider: ``describe`` routes through ``model_turn``, so the stub carries the lane surface (``provider_name``, ``get_capabilities``) and returns a full - ``CompletionResult`` shape, and it records the ``resolve_attachments`` - callback the translator would use to materialize the by-reference parts. + ``CompletionResult`` shape, and it records that ``model_turn`` already + materialized by-reference parts and cleared the provider callback. """ provider_name = "openai-compatible" @@ -100,11 +100,10 @@ def test_describe_lowers_prompt_then_by_reference_parts() -> None: assert prov.last_messages is not None content = prov.last_messages[0]["content"] assert content[0]["type"] == "text" # prompt leads - # The attachment rides by reference; the translator materializes it via - # the threaded resolver, which must return the prebuilt parts verbatim. - assert content[1]["attachment_id"] == "perception-input" - assert prov.last_resolve is not None - assert prov.last_resolve(["perception-input"]) == {"perception-input": _parts()} + # model_turn materializes before admission, then hands the provider the + # prebuilt inline part with no resolver left to invoke under the gate. + assert content[1] == _parts()[0] + assert prov.last_resolve is None def test_describe_empty_parts_skips_backend() -> None: diff --git a/tests/test_server_node_models_metadata.py b/tests/test_server_node_models_metadata.py index ee53ae99..fbd908af 100644 --- a/tests/test_server_node_models_metadata.py +++ b/tests/test_server_node_models_metadata.py @@ -139,6 +139,7 @@ def test_model_status_route_carries_backend_auth_fields(): auth_mode="rfc8693_obo", obo_audience="api://gw", obo_scopes="aud-gw openid", + max_concurrency=3, ) reg = ModelRegistry(models={"gw": cfg}, default="gw") request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(registry=reg))) @@ -147,6 +148,7 @@ def test_model_status_route_carries_backend_auth_fields(): assert entry["auth_mode"] == "rfc8693_obo" assert entry["obo_audience"] == "api://gw" assert entry["obo_scopes"] == "aud-gw openid" + assert entry["max_concurrency"] == 3 # --------------------------------------------------------------------------- @@ -326,6 +328,64 @@ def test_model_reload_endpoint_rewrites_models_metadata(monkeypatch, tmp_path): assert {r["alias"] for r in payload} == {"a", "b"} +def test_model_reload_cap_only_change_resizes_the_stable_gate(monkeypatch): + """Operational capacity participates in the endpoint's no-op check even + though it deliberately does not participate in ModelConfig identity.""" + from turnstone.server import internal_model_reload + + old_reg = ModelRegistry( + { + "a": ModelConfig( + alias="a", + base_url="http://x", + api_key="k", + model="a", + max_concurrency=1, + ) + }, + default="a", + ) + new_reg = ModelRegistry( + { + "a": ModelConfig( + alias="a", + base_url="http://x", + api_key="k", + model="a", + max_concurrency=3, + ) + }, + default="a", + ) + original_gate = old_reg.get_admission("a") + app_state = SimpleNamespace( + registry=old_reg, + cli_model_args={ + "base_url": "", + "api_key": "", + "model": "", + "context_window": 0, + "provider": "openai", + }, + config_store=None, + node_id="", + ) + request = SimpleNamespace(app=SimpleNamespace(state=app_state)) + + monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", lambda **_kw: new_reg) + monkeypatch.setattr("turnstone.core.storage._registry.get_storage", MagicMock) + monkeypatch.setattr("turnstone.server._broadcast_agent_tool_schema_refresh", lambda _s: None) + + response = internal_model_reload(request) # type: ignore[arg-type] + body = json.loads(response.body) + + assert response.status_code == 200 + assert body.get("noop", False) is False + assert old_reg.get_admission("a") is original_gate + assert original_gate.limit == 3 + assert old_reg.generation == 1 + + def test_model_reload_refuses_dynamic_auth_without_key(monkeypatch, tmp_path, caplog): """A keyless node cannot acquire a dynamic alias via model-reload: 503, a deployment fault, not the 422 bad-arguments exit.""" @@ -420,6 +480,44 @@ def test_model_reload_maps_auth_config_error_to_422(monkeypatch, tmp_path): assert old_reg.has_alias("a") +def test_model_reload_maps_concurrency_config_error_to_422(monkeypatch, tmp_path): + """A corrupt persisted cap uses the same structured config-error exit.""" + from turnstone.core.model_registry import ModelConcurrencyConfigError + from turnstone.core.storage._sqlite import SQLiteBackend + from turnstone.server import internal_model_reload + + storage = SQLiteBackend(str(tmp_path / "reload.db")) + old_reg = _registry(("a", "http://x")) + app_state = SimpleNamespace( + registry=old_reg, + cli_model_args={ + "base_url": "", + "api_key": "", + "model": "", + "context_window": 0, + "provider": "openai", + }, + config_store=None, + node_id="node-a", + ) + request = SimpleNamespace(app=SimpleNamespace(state=app_state)) + + def _raise_concurrency_config(**_kw): + raise ModelConcurrencyConfigError("Model 'gw' max_concurrency must be an integer") + + monkeypatch.setattr( + "turnstone.core.model_registry.load_model_registry", + _raise_concurrency_config, + ) + monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage) + + response = internal_model_reload(request) # type: ignore[arg-type] + + assert response.status_code == 422 + assert "max_concurrency" in json.loads(response.body)["reason"] + assert old_reg.has_alias("a") + + def test_config_reload_maps_dynamic_auth_key_error_to_503(monkeypatch, tmp_path, caplog): """Every caller of the reload chokepoint handles its refusal the same way: the settings fan-out answers 503, not an unhandled 500.""" diff --git a/tests/test_session.py b/tests/test_session.py index 25769537..0950040b 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -10057,6 +10057,75 @@ def test_task_agent_static_auth_fallback_never_reresolves_as_successor(): lane.client.with_options.assert_not_called() +def test_task_agent_defers_pinned_auth_resolution_until_model_admission(): + """A child mints for its initiating user only after admission.""" + from dataclasses import replace + + from turnstone.core.admission import ModelAdmission + from turnstone.core.model_turn import model_turn as real_model_turn + + session = _make_session() + session._acting_user_id = "user-b" + provider = seam_provider("done", provider_name="openai-compatible") + lane = replace_session_lane(session, provider=provider, alias="task-gateway") + auth_config = MagicMock(name="pinned-auth-config") + gate = ModelAdmission("task-gateway", 1) + stale_live_resolver = MagicMock(return_value="token-for-user-b") + lane = replace( + lane, + backend_auth_resolver=stale_live_resolver, + backend_auth_config=auth_config, + admission=gate, + ) + session._model_binding = replace(session._model_binding, lane=lane) + bound_client = object() + lane.client.with_options.return_value = bound_client + mint_in_flight: list[int] = [] + + def _resolve_for_principal(alias, config, *, principal_id): + assert alias == "task-gateway" + assert config is auth_config + assert session._acting_user_id == "user-b" + assert principal_id == "user-a" + mint_in_flight.append(gate.snapshot().in_flight) + return "token-for-user-a" + + pinned_resolver = MagicMock(side_effect=_resolve_for_principal) + session._model_backend_auth_token_for_principal = pinned_resolver + stream = provider.create_streaming.return_value + + def _dispatch(**kwargs): + assert gate.snapshot().in_flight == 1 + assert kwargs["client"] is bound_client + return stream + + provider.create_streaming.side_effect = _dispatch + + with patch("turnstone.core.session.model_turn", wraps=real_model_turn) as plant_call: + result = session._run_agent( + [Turn.user("finish the task")], + tools=[], + auto_tools=set(), + principal_id="user-a", + ) + + assert result == "done" + assert plant_call.call_count == 1 + called_lane = plant_call.call_args.args[0] + assert called_lane.admission is gate + assert called_lane.backend_auth_resolver is not None + assert "backend_auth_token" not in plant_call.call_args.kwargs + pinned_resolver.assert_called_once_with( + lane.alias, + auth_config, + principal_id="user-a", + ) + stale_live_resolver.assert_not_called() + lane.client.with_options.assert_called_once_with(api_key="token-for-user-a") + assert mint_in_flight == [1] + assert gate.snapshot().in_flight == 0 + + def test_already_cancelled_task_agent_does_not_resolve_backend_auth(): from turnstone.core.session import GenerationCancelled diff --git a/tests/test_session_attachments.py b/tests/test_session_attachments.py index 4b5fb13e..10a366c0 100644 --- a/tests/test_session_attachments.py +++ b/tests/test_session_attachments.py @@ -782,15 +782,13 @@ class TestPerceptionFallback: ) assert part["type"] == "text" assert "DESCRIPTION" in part["text"] - # the perception model was handed the rasterized pages, not the raw - # PDF: the wire carries the prompt + a by-reference placeholder, and - # the threaded resolver materializes the page parts at the translator. + # The perception model was handed the rasterized pages, not the raw + # PDF: model_turn expands the by-reference placeholder before taking + # admission, then the provider receives inline pages and no resolver. sent = prov.create_streaming.call_args.kwargs["messages"][0]["content"] assert sent[0]["type"] == "text" - assert sent[1]["attachment_id"] == "perception-input" - resolver = prov.create_streaming.call_args.kwargs["resolve_attachments"] - pages = resolver(["perception-input"])["perception-input"] - assert [p["type"] for p in pages] == ["image_url", "image_url"] + assert [p["type"] for p in sent[1:]] == ["image_url", "image_url"] + assert prov.create_streaming.call_args.kwargs["resolve_attachments"] is None def test_audio_perception_when_omni_and_no_stt(self, tmp_db, mock_openai_client): s = _make_session(mock_openai_client) diff --git a/turnstone.example.toml b/turnstone.example.toml index 7646876f..11cb989c 100644 --- a/turnstone.example.toml +++ b/turnstone.example.toml @@ -44,6 +44,8 @@ # provider = "openai" # base_url = "http://localhost:8000/v1" # context_window = 8192 +# max_concurrency = 1 # Per-process generation cap for this alias; 0 = unlimited. +# # For llama.cpp, start with its usable -np slot count. # # [models.local.capabilities] # supports_vision = false diff --git a/turnstone/api/console_schemas.py b/turnstone/api/console_schemas.py index 3a2d9028..7b4031c0 100644 --- a/turnstone/api/console_schemas.py +++ b/turnstone/api/console_schemas.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, Literal +from typing import Annotated, Any, Literal, TypeAlias from pydantic import BaseModel, Field @@ -12,6 +12,7 @@ from pydantic import BaseModel, Field from pydantic.json_schema import SkipJsonSchema # noqa: TC002 from turnstone.api.server_schemas import CreateWorkstreamRequest, CreateWorkstreamResponse +from turnstone.core.model_registry import MAX_MODEL_CONCURRENCY from turnstone.core.skill_kind import SkillKind from turnstone.core.skill_parser import MAX_SKILL_DESCRIPTION_LEN @@ -1002,6 +1003,18 @@ class RegistryInstallRequest(BaseModel): # Admin: Model Definitions # --------------------------------------------------------------------------- +ModelMaxConcurrency: TypeAlias = Annotated[ + int, + Field( + strict=True, + ge=0, + le=MAX_MODEL_CONCURRENCY, + description=( + "Maximum concurrent model generations for this alias in one process; zero means unlimited." + ), + ), +] + class ModelDefinitionInfo(BaseModel): definition_id: str @@ -1011,6 +1024,7 @@ class ModelDefinitionInfo(BaseModel): base_url: str = "" api_key: str = "" context_window: int = 32768 + max_concurrency: ModelMaxConcurrency = 0 capabilities: str = "{}" enabled: bool = True temperature: float | None = None @@ -1059,6 +1073,7 @@ class CreateModelDefinitionRequest(BaseModel): base_url: str = "" api_key: str = "" context_window: int = 32768 + max_concurrency: ModelMaxConcurrency = 0 capabilities: dict[str, Any] = Field(default_factory=dict) enabled: bool = True temperature: float | None = None @@ -1078,6 +1093,12 @@ class UpdateModelDefinitionRequest(BaseModel): base_url: str | None = None api_key: str | None = None context_window: int | None = None + # The runtime update handler is presence-keyed. Keep null out of the + # advertised union because explicit JSON null is refused; clients clear a + # limit by sending the canonical unlimited value, zero. + max_concurrency: ModelMaxConcurrency | SkipJsonSchema[None] = Field( + default_factory=lambda: None + ) # SkipJsonSchema drops the null member from the ADVERTISED union while the # Python type still tolerates None: the presence-keyed update handler # refuses an explicit JSON null, so advertising null would let generated diff --git a/turnstone/console/server.py b/turnstone/console/server.py index 54fbae2e..6c0f2492 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -64,6 +64,7 @@ from turnstone.core.metacognition import field_str, sanitize_display from turnstone.core.model_registry import ( APP_IDENTITY_MODEL_AUTH_MODES, DYNAMIC_MODEL_AUTH_MODES, + MAX_MODEL_CONCURRENCY, MODEL_AUTH_MODE_PROFILES, MODEL_AUTH_TEXT_MAX_LEN, SCOPES_MODEL_AUTH_MODES, @@ -11964,10 +11965,12 @@ def _oidc_configured_for_model_auth(request: Request) -> bool: # "changing this column can neither redirect where a minted credential is # sent nor re-arm minting that a disable or revocation stopped": the four # sampling/shaping knobs hit the same endpoint with the same credential, -# and the two reasoning toggles only select what history surfaces. +# the admission knob only limits callers of that alias, and the two reasoning +# toggles only select what history surfaces. MODEL_AUTH_NEUTRAL_FIELDS = frozenset( { "context_window", + "max_concurrency", "temperature", "max_tokens", "reasoning_effort", @@ -11976,6 +11979,26 @@ MODEL_AUTH_NEUTRAL_FIELDS = frozenset( } ) + +def _parse_model_max_concurrency(raw: Any) -> tuple[int, JSONResponse | None]: + """Parse the strict model-definition concurrency scalar. + + JSON booleans are integer subclasses in Python, so exact type equality is + load-bearing. Strings and integral floats are also refused rather than + silently normalized into a materially different admission policy. + """ + if type(raw) is not int or raw < 0 or raw > MAX_MODEL_CONCURRENCY: + return 0, JSONResponse( + { + "error": ( + f"max_concurrency must be an integer between 0 and {MAX_MODEL_CONCURRENCY}" + ) + }, + status_code=400, + ) + return raw, None + + # The two cross-field refusal messages, shared verbatim by the create and # update twins (their guard CONDITIONS differ — raw body vs post-merge pair — # but the text must not drift). Mode lists come from the frozenset so a @@ -12728,6 +12751,7 @@ async def admin_list_model_definitions(request: Request) -> JSONResponse: "base_url": "", "api_key": "", "context_window": nm.get("context_window", 0), + "max_concurrency": nm.get("max_concurrency", 0), "capabilities": "{}", "enabled": True, "temperature": nm.get("temperature"), @@ -12875,6 +12899,9 @@ async def admin_create_model_definition(request: Request) -> JSONResponse: if isinstance(ctx_raw, float) and not math.isfinite(ctx_raw): return JSONResponse({"error": "context_window must be a finite number"}, status_code=400) context_window = max(0, int(ctx_raw)) if isinstance(ctx_raw, (int, float)) else 0 + max_concurrency, concurrency_err = _parse_model_max_concurrency(body.get("max_concurrency", 0)) + if concurrency_err is not None: + return concurrency_err caps = body.get("capabilities", {}) if not isinstance(caps, dict): # Mirror of the update twin's refusal: coercing a null or the STRING @@ -12977,6 +13004,7 @@ async def admin_create_model_definition(request: Request) -> JSONResponse: base_url=base_url, api_key=api_key, context_window=context_window, + max_concurrency=max_concurrency, capabilities=capabilities, enabled=enabled, created_by=audit_uid, @@ -13120,6 +13148,11 @@ async def admin_update_model_definition(request: Request) -> JSONResponse: {"error": "context_window must be a finite number"}, status_code=400 ) updates["context_window"] = max(0, int(ctx_raw)) if isinstance(ctx_raw, (int, float)) else 0 + if "max_concurrency" in body: + max_concurrency, concurrency_err = _parse_model_max_concurrency(body["max_concurrency"]) + if concurrency_err is not None: + return concurrency_err + updates["max_concurrency"] = max_concurrency if "capabilities" in body: caps = body["capabilities"] if not isinstance(caps, dict): diff --git a/turnstone/console/static/admin.js b/turnstone/console/static/admin.js index 1cd8afab..80602bb0 100644 --- a/turnstone/console/static/admin.js +++ b/turnstone/console/static/admin.js @@ -7525,8 +7525,9 @@ function _renderModels(items) { colAlias.appendChild(document.createTextNode(" ")); colAlias.appendChild(defBadge); } - // Per-model sampling override indicators + // Non-default per-model setting indicators const overrides = []; + if (m.max_concurrency > 0) overrides.push("limit=" + m.max_concurrency); if (m.temperature != null) overrides.push("temp=" + m.temperature); if (m.max_tokens != null) overrides.push("max_tok=" + m.max_tokens); if (m.reasoning_effort != null) @@ -7569,10 +7570,10 @@ function _renderModels(items) { const ovrSpan = document.createElement("span"); ovrSpan.className = "model-overrides-hint"; ovrSpan.textContent = overrides.join(", "); - ovrSpan.title = "Per-model overrides (override global defaults)"; + ovrSpan.title = "Per-model settings"; ovrSpan.setAttribute( "aria-label", - "Per-model overrides: " + overrides.join(", "), + "Per-model settings: " + overrides.join(", "), ); colAlias.appendChild(document.createElement("br")); colAlias.appendChild(ovrSpan); @@ -7746,6 +7747,7 @@ function showCreateModelModal() { document.getElementById("model-api-key").value = ""; document.getElementById("model-api-key").placeholder = "sk-..."; document.getElementById("model-ctx-window").value = "0"; + document.getElementById("model-max-concurrency").value = "0"; document.getElementById("model-temperature").value = ""; document.getElementById("model-max-tokens").value = ""; document.getElementById("model-reasoning-effort").value = ""; @@ -7855,6 +7857,8 @@ function showEditModelModal(definitionId) { "\u2022\u2022\u2022 (leave blank to keep existing)"; document.getElementById("model-ctx-window").value = m.context_window != null ? m.context_window : 0; + document.getElementById("model-max-concurrency").value = + m.max_concurrency != null ? m.max_concurrency : 0; document.getElementById("model-temperature").value = m.temperature != null ? m.temperature : ""; document.getElementById("model-max-tokens").value = @@ -8142,6 +8146,24 @@ function submitCreateModel() { enabled: document.getElementById("model-enabled").checked, }; + // Per-alias admission. Empty and zero are the canonical unlimited value; + // every other spelling must be an exact non-negative integer. + const concurrencyText = document + .getElementById("model-max-concurrency") + .value.trim(); + const maxConcurrency = concurrencyText === "" ? 0 : Number(concurrencyText); + if ( + !Number.isInteger(maxConcurrency) || + maxConcurrency < 0 || + maxConcurrency > 2147483647 + ) { + _showModelError( + "Max concurrent generations must be a whole number from 0 to 2147483647", + ); + return; + } + form.max_concurrency = maxConcurrency; + // Per-model sampling overrides — null when empty (use global default) const tempVal = document.getElementById("model-temperature").value.trim(); if (tempVal !== "") { diff --git a/turnstone/console/static/index.html b/turnstone/console/static/index.html index 2138be4f..78fb7d7b 100644 --- a/turnstone/console/static/index.html +++ b/turnstone/console/static/index.html @@ -1695,6 +1695,23 @@ /> +
+
+ + +
+
diff --git a/turnstone/core/admission.py b/turnstone/core/admission.py new file mode 100644 index 00000000..05b49e4f --- /dev/null +++ b/turnstone/core/admission.py @@ -0,0 +1,210 @@ +"""Per-alias admission control for model dispatch. + +The registry owns one :class:`ModelAdmission` object per alias and keeps that +object stable across hot reloads. A zero limit is unlimited, but calls are +still counted so a live ``0 -> N`` resize can drain already-running work before +admitting more. +""" + +from __future__ import annotations + +import contextlib +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Any + +from turnstone.core.deadline import DeadlineCancelledError +from turnstone.core.log import get_logger + +log = get_logger(__name__) + +_CANCEL_POLL_SECONDS = 0.25 +_STALL_WARNING_SECONDS = 5.0 + + +def _call_hook(target: Any, name: str) -> None: + hook = getattr(target, name, None) + if callable(hook): + with contextlib.suppress(Exception): + hook() + + +@dataclass(frozen=True, slots=True) +class AdmissionSnapshot: + """Non-sensitive instantaneous state for logs and tests.""" + + alias: str + limit: int + in_flight: int + queued: int + + +class AdmissionLease: + """One idempotently releasable admission hold.""" + + __slots__ = ("_gate", "wait_seconds", "queued_ahead", "_released") + + def __init__( + self, + gate: ModelAdmission, + *, + wait_seconds: float, + queued_ahead: int, + ) -> None: + self._gate = gate + self.wait_seconds = wait_seconds + self.queued_ahead = queued_ahead + self._released = False + + def release(self) -> None: + if self._released: + return + self._released = True + self._gate._release() + + def __enter__(self) -> AdmissionLease: + return self + + def __exit__(self, *_exc: object) -> None: + self.release() + + +class ModelAdmission: + """FIFO, hot-resizable in-flight gate for one model alias. + + ``limit == 0`` preserves unlimited behavior. Unlimited calls still take + leases and increment ``in_flight`` so narrowing the gate during a hot reload + observes and drains them rather than starting from a false zero. + """ + + def __init__(self, alias: str, limit: int = 0) -> None: + if type(limit) is not int or limit < 0: + raise ValueError("model admission limit must be a non-negative integer") + self.alias = alias + self._limit = limit + self._in_flight = 0 + self._waiters: deque[object] = deque() + self._cv = threading.Condition() + + @property + def limit(self) -> int: + with self._cv: + return self._limit + + def snapshot(self) -> AdmissionSnapshot: + with self._cv: + return AdmissionSnapshot( + alias=self.alias, + limit=self._limit, + in_flight=self._in_flight, + queued=len(self._waiters), + ) + + def set_limit(self, limit: int) -> None: + """Resize in place; narrowing lets current holders drain naturally.""" + if type(limit) is not int or limit < 0: + raise ValueError("model admission limit must be a non-negative integer") + with self._cv: + if limit == self._limit: + return + previous = self._limit + self._limit = limit + self._cv.notify_all() + log.info( + "model.admission_resized", + alias=self.alias, + previous_limit=previous, + limit=limit, + in_flight=self._in_flight, + queued=len(self._waiters), + ) + + def _available(self) -> bool: + return self._limit == 0 or self._in_flight < self._limit + + def acquire(self, *, cancel_ref: Any = None) -> AdmissionLease: + """Wait FIFO for a slot, abandoning promptly when *cancel_ref* aborts.""" + if bool(getattr(cancel_ref, "aborted", False)): + raise DeadlineCancelledError("cancel_ref aborted during model admission") + + started = time.monotonic() + ticket = object() + queued_ahead = 0 + waiting = False + did_wait = False + warned = False + try: + with self._cv: + if self._waiters or not self._available(): + queued_ahead = len(self._waiters) + self._waiters.append(ticket) + waiting = True + did_wait = True + _call_hook(cancel_ref, "begin_admission_wait") + + while waiting: + if bool(getattr(cancel_ref, "aborted", False)): + with contextlib.suppress(ValueError): + self._waiters.remove(ticket) + self._cv.notify_all() + raise DeadlineCancelledError("cancel_ref aborted during model admission") + if self._waiters[0] is ticket and self._available(): + self._waiters.popleft() + waiting = False + break + waited = time.monotonic() - started + if not warned and waited >= _STALL_WARNING_SECONDS: + warned = True + log.warning( + "model.admission_stalled", + alias=self.alias, + limit=self._limit, + in_flight=self._in_flight, + queued=len(self._waiters), + wait_seconds=round(waited, 3), + ) + self._cv.wait(_CANCEL_POLL_SECONDS) + + self._in_flight += 1 + # When the limit is wider than one, let the new FIFO head claim + # the next already-free slot without waiting for this holder to + # release first. + self._cv.notify_all() + finally: + if waiting: + # An unexpected hook/condition failure must not strand its FIFO + # ticket. The normal cancellation path already removed it. + with self._cv: + with contextlib.suppress(ValueError): + self._waiters.remove(ticket) + self._cv.notify_all() + if did_wait: + _call_hook(cancel_ref, "end_admission_wait") + + waited = time.monotonic() - started if did_wait else 0.0 + lease = AdmissionLease( + self, + wait_seconds=waited, + queued_ahead=queued_ahead, + ) + if bool(getattr(cancel_ref, "aborted", False)): + lease.release() + raise DeadlineCancelledError("cancel_ref aborted during model admission") + if waited > 0: + log.info( + "model.admission_wait", + alias=self.alias, + limit=self.limit, + queued_ahead=queued_ahead, + wait_seconds=round(waited, 3), + ) + return lease + + def _release(self) -> None: + with self._cv: + if self._in_flight <= 0: + raise RuntimeError("model admission lease released without a holder") + self._in_flight -= 1 + self._cv.notify_all() diff --git a/turnstone/core/deadline.py b/turnstone/core/deadline.py index 1db7b279..34ae7d3e 100644 --- a/turnstone/core/deadline.py +++ b/turnstone/core/deadline.py @@ -57,12 +57,25 @@ class StreamAbortRef(list[Any]): fix here must be mirrored there. """ - __slots__ = ("_aborted", "_cancel_event") + __slots__ = ( + "_aborted", + "_cancel_event", + "_timing_lock", + "_admission_wait_started", + "_admission_wait_credit", + "_dispatch_count", + "_last_dispatch_at", + ) def __init__(self, cancel_event: threading.Event | None = None) -> None: super().__init__() self._aborted = False self._cancel_event = cancel_event + self._timing_lock = threading.Lock() + self._admission_wait_started: float | None = None + self._admission_wait_credit = 0.0 + self._dispatch_count = 0 + self._last_dispatch_at: float | None = None def append(self, stream: Any) -> None: super().append(stream) @@ -77,6 +90,49 @@ class StreamAbortRef(list[Any]): with contextlib.suppress(Exception): stream.close() + def begin_admission_wait(self) -> None: + """Freeze this call's deadline while it waits for model admission.""" + with self._timing_lock: + if self._admission_wait_started is None: + self._admission_wait_started = time.monotonic() + + def end_admission_wait(self) -> None: + """Resume the deadline and retain the elapsed admission credit.""" + now = time.monotonic() + with self._timing_lock: + started = self._admission_wait_started + if started is None: + return + self._admission_wait_credit += max(0.0, now - started) + self._admission_wait_started = None + + def admission_wait_credit(self) -> float: + """Return completed plus currently accruing admission-wait time.""" + now = time.monotonic() + with self._timing_lock: + credit = self._admission_wait_credit + if self._admission_wait_started is not None: + credit += max(0.0, now - self._admission_wait_started) + return credit + + def mark_dispatch(self) -> None: + """Record a provider dispatch without changing deadline semantics.""" + with self._timing_lock: + self._dispatch_count += 1 + self._last_dispatch_at = time.monotonic() + + @property + def dispatch_count(self) -> int: + """Number of provider dispatch attempts observed by this ref.""" + with self._timing_lock: + return self._dispatch_count + + @property + def last_dispatch_at(self) -> float | None: + """Monotonic timestamp of the latest provider dispatch, if any.""" + with self._timing_lock: + return self._last_dispatch_at + @property def aborted(self) -> bool: """Whether :meth:`abort` has fired. @@ -123,6 +179,7 @@ def run_abortable_with_deadline( poll=poll, thread_name=thread_name, on_abandon=abort_ref.abort, + deadline_credit=abort_ref.admission_wait_credit, ) @@ -134,6 +191,7 @@ def run_with_deadline( poll: float = 1.0, thread_name: str = "deadline-worker", on_abandon: Callable[[], None] | None = None, + deadline_credit: Callable[[], float] | None = None, ) -> _T: """Run ``fn()`` on a daemon thread, bounded by ``timeout``/``cancel_event``. @@ -153,6 +211,11 @@ def run_with_deadline( ``poll`` bounds how often ``cancel_event`` is checked (and thus the worst- case latency from a cancel to this function returning). + + ``deadline_credit`` may dynamically extend the original deadline. The + abortable wrapper uses it only for time spent queued at model admission; + lowering, attachment materialization, credential minting, dispatch, + draining, and retry backoff continue to consume the original budget. """ box: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1) @@ -186,7 +249,11 @@ def run_with_deadline( if cancel_event is not None and cancel_event.is_set(): _abandon(DeadlineCancelledError()) - remaining = deadline - time.monotonic() + credit = 0.0 + if deadline_credit is not None: + with contextlib.suppress(Exception): + credit = max(0.0, float(deadline_credit())) + remaining = deadline + credit - time.monotonic() if remaining <= 0: _abandon(DeadlineExceededError()) try: diff --git a/turnstone/core/judge.py b/turnstone/core/judge.py index ee365369..1eb324b8 100644 --- a/turnstone/core/judge.py +++ b/turnstone/core/judge.py @@ -25,6 +25,7 @@ from turnstone.core.deadline import ( run_abortable_with_deadline, ) from turnstone.core.log import get_logger +from turnstone.core.model_backend_auth import BackendAuthUnavailableError from turnstone.core.model_registry import ModelClientConstructionError from turnstone.core.model_turn import ( ModelLane, @@ -1318,9 +1319,10 @@ class IntentJudge: live set (parallel task agents each spawn a generation; ``close()`` aborts whatever is still live). backend_auth_resolver: Batch-scoped resolver whose closure pins - the initiating principal. It is invoked once by the daemon, - after its initial cancellation check, and the resulting token - is reused for every item and evidence turn in this batch. + the initiating principal. The resolver remains on the + batch's lane so each plant-call attempt mints after acquiring + alias admission. The mint cache keeps this inexpensive while + preventing a queued bearer from expiring before dispatch. Returns: List of heuristic verdicts (one per item), available immediately. @@ -1398,28 +1400,17 @@ class IntentJudge: ) return - # Resolve delegated credentials exactly once for the batch. The - # caller-supplied closure has already captured the initiating - # principal, so a later shared-workstream handoff cannot mint a - # successor user's token for this payload. - backend_auth_token: str | None = None + # Preserve the caller-pinned principal on the batch lane, but let + # model_turn resolve credentials only after alias admission. An + # admission backlog can outlive a bearer token; carrying the + # resolver refreshes near-expiry tokens at the dispatch boundary + # without allowing a later shared-workstream actor to take over. batch_lane = self._lane if backend_auth_resolver is not None: - try: - backend_auth_token = backend_auth_resolver( - self._lane.alias, - self._lane.backend_auth_config, - ) - except Exception: - log.exception("Judge backend authentication failed") - self._deliver_fallbacks( - items, - heuristic_verdicts, - callback, - "judge backend authentication failed", - ) - return - batch_lane = replace(batch_lane, backend_auth_resolver=None) + batch_lane = replace( + batch_lane, + backend_auth_resolver=backend_auth_resolver, + ) if cancel_event and cancel_event.is_set(): self._deliver_fallbacks( @@ -1434,7 +1425,8 @@ class IntentJudge: # One lane derivative for the whole batch: only the judge-owned # fresh client differs from the immutable constructor binding. # Every item and evidence turn therefore stays on one provider, - # model, capability, config, and credential snapshot. + # model, capability, auth-config, and initiating-principal + # binding while credentials refresh per admitted attempt. batch_lane = replace(batch_lane, client=client) for idx, (item, h_verdict) in enumerate(zip(items, heuristic_verdicts, strict=True)): if cancel_event and cancel_event.is_set(): @@ -1453,7 +1445,6 @@ class IntentJudge: cancel_event, client, lane=batch_lane, - backend_auth_token=backend_auth_token, ) if llm_verdict: log.info( @@ -1497,6 +1488,15 @@ class IntentJudge: "judge cancelled before evaluating this call", ) return + except BackendAuthUnavailableError: + log.exception("Judge backend authentication failed") + self._deliver_fallbacks( + items[idx:], + heuristic_verdicts[idx:], + callback, + "judge backend authentication failed", + ) + return except Exception: log.exception( "Judge evaluation failed for %s", @@ -1549,7 +1549,6 @@ class IntentJudge: client: Any | None, *, lane: ModelLane | None = None, - backend_auth_token: str | None = None, ) -> IntentVerdict | None: """Run LLM judge for a single tool call. Returns verdict or None.""" if lane is None: @@ -1654,7 +1653,6 @@ class IntentJudge: tools=_tools, max_tokens=2048, cancel_ref=ref, - backend_auth_token=backend_auth_token, ) result = run_abortable_with_deadline( @@ -1681,6 +1679,8 @@ class IntentJudge: log.info("judge.verdict.from_partial", turn=turn + 1) return verdict return None + except BackendAuthUnavailableError: + raise except Exception as e: log.info("judge.turn.failed", turn=turn + 1, error=str(e)) return None diff --git a/turnstone/core/model_registry.py b/turnstone/core/model_registry.py index 7cfec0a8..b428fe2c 100644 --- a/turnstone/core/model_registry.py +++ b/turnstone/core/model_registry.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from collections.abc import Mapping +from turnstone.core.admission import ModelAdmission from turnstone.core.config import load_config from turnstone.core.log import get_logger from turnstone.core.providers import LLMProvider, create_client, create_provider @@ -29,6 +30,12 @@ MODEL_AUTH_MODES = frozenset({"static", "entra_obo", "entra_app", "rfc8693_obo"} # drift into "console stores it, registry refuses to load it". MODEL_AUTH_TEXT_MAX_LEN = 2048 +# ``model_definitions.max_concurrency`` is an ``INTEGER`` on both supported +# databases. Keep the public/config/API bound aligned with PostgreSQL's +# signed 32-bit representation so a value accepted on SQLite cannot fail when +# the same definition is moved to PostgreSQL. +MAX_MODEL_CONCURRENCY = 2_147_483_647 + # Derived, never hand-listed (fail-safe defaults): any mode later added to # MODEL_AUTH_MODES lands in this set BY CONSTRUCTION unless it is literally # "static", so membership tests fail CLOSED for modes nobody classified. @@ -80,6 +87,10 @@ class ModelAuthConfigError(ValueError): """A model definition contains unsafe or internally inconsistent auth settings.""" +class ModelConcurrencyConfigError(ValueError): + """A model definition has an invalid per-alias concurrency limit.""" + + class ModelClientConstructionError(ValueError): """A registry alias exists but its binding could not be constructed. @@ -263,6 +274,24 @@ class ModelConfig: auth_mode: str = "static" obo_audience: str = "" obo_scopes: str = "" + # Per-process admission limit for this alias. Zero preserves the + # historical unlimited behavior. Operational admission changes do not + # change binding identity, hence ``compare=False``. + max_concurrency: int = field(default=0, compare=False) + + def __post_init__(self) -> None: + # ``bool`` is an ``int`` subclass, so use exact type equality. A + # permissive coercion here could turn ``true`` into a one-request cap + # or a typo into unlimited operation. + if ( + type(self.max_concurrency) is not int + or self.max_concurrency < 0 + or self.max_concurrency > MAX_MODEL_CONCURRENCY + ): + raise ModelConcurrencyConfigError( + f"Model {self.alias!r} max_concurrency must be an integer " + f"between 0 and {MAX_MODEL_CONCURRENCY}" + ) def strip_control_characters(value: str) -> str: @@ -453,6 +482,10 @@ class ModelRegistry: self.task_effort = task_effort self._clients: dict[str, Any] = {} self._providers: dict[str, LLMProvider] = {} + self._admissions = { + alias: ModelAdmission(alias, int(getattr(cfg, "max_concurrency", 0))) + for alias, cfg in self._models.items() + } self._client_lock = threading.Lock() # Monotone count of completed reload() swaps. A counter, not a field # diff: sessions re-resolve on ANY difference at the next send, so an @@ -536,6 +569,14 @@ class ModelRegistry: raise UnknownModelAliasError(alias) return self._models[alias] + def get_admission(self, alias: str) -> ModelAdmission: + """Return the stable per-alias admission gate.""" + with self._client_lock: + gate = self._admissions.get(alias) + if gate is None: + raise UnknownModelAliasError(alias) + return gate + def has_alias(self, alias: str) -> bool: """Check if *alias* exists in the registry.""" return alias in self._models @@ -565,8 +606,8 @@ class ModelRegistry: def resolve_binding( self, alias: str | None = None - ) -> tuple[Any, str, ModelConfig, LLMProvider, int]: - """Resolve *alias* to ``(client, model_name, config, provider, generation)``. + ) -> tuple[Any, str, ModelConfig, LLMProvider, ModelAdmission, int]: + """Resolve client, model, config, provider, admission, and generation. The session bind primitive: everything a rebind commits, read under ONE lock acquisition, so a :meth:`reload` landing between separate @@ -592,7 +633,14 @@ class ModelRegistry: raise except ValueError as exc: raise ModelClientConstructionError(str(exc)) from exc - return (client, cfg.model, cfg, provider, self._generation) + return ( + client, + cfg.model, + cfg, + provider, + self._admissions[alias], + self._generation, + ) def resolve_agent_alias(self, kind: str) -> str | None: """Return the configured alias for a sub-agent ``kind``. @@ -703,6 +751,20 @@ class ModelRegistry: self.agent_model = agent_model self.task_model = task_model self.task_effort = task_effort + # Admission is strictly per alias. Resize surviving gates in + # place so live lanes and new resolutions coordinate through the + # same FIFO even when the alias moves to a different endpoint. + for alias, cfg in self._models.items(): + limit = int(getattr(cfg, "max_concurrency", 0)) + gate = self._admissions.get(alias) + if gate is None: + self._admissions[alias] = ModelAdmission(alias, limit) + else: + gate.set_limit(limit) + # Removed aliases remain as tombstones for this registry's + # lifetime. A stale lane may still hold or queue on that object; + # re-adding the alias must reconfigure the same gate rather than + # split old and new work across two independent limits. # Selective teardown — close + drop only clients whose # construction/connection target changed (alias removed, or # base_url / api_key / provider / auth_mode differs). Keeps connection @@ -894,8 +956,9 @@ def load_model_registry( auth_mode=row_auth_mode, obo_audience=row_obo_audience, obo_scopes=row_obo_scopes, + max_concurrency=row.get("max_concurrency", 0), ) - except ModelAuthConfigError: + except (ModelAuthConfigError, ModelConcurrencyConfigError): # Configuration errors are authoritative row content, not a # transient storage-read failure. Never degrade past them into a # config-only registry or provider SDK environment credentials. @@ -979,6 +1042,7 @@ def load_model_registry( auth_mode=entry_auth_mode, obo_audience=entry_obo_audience, obo_scopes=entry_obo_scopes, + max_concurrency=entry.get("max_concurrency", 0), ) # 3. Back-compat shim: synthesize a "default" alias from CLI/auto-detected diff --git a/turnstone/core/model_turn.py b/turnstone/core/model_turn.py index e7738eec..d61f4bdf 100644 --- a/turnstone/core/model_turn.py +++ b/turnstone/core/model_turn.py @@ -38,6 +38,7 @@ Contract, held deliberately narrow: from __future__ import annotations +import contextlib import random import time import uuid @@ -50,6 +51,7 @@ if TYPE_CHECKING: from turnstone.core.model_registry import ModelConfig, ModelRegistry +from turnstone.core.admission import ModelAdmission from turnstone.core.deadline import DeadlineCancelledError from turnstone.core.history_decoration import attach_vllm_chat_reasoning_field from turnstone.core.log import get_logger @@ -103,7 +105,13 @@ from turnstone.core.storage._utils import ( _CLIENT_TOOL_CALL_BLOCK_TYPES, strip_orphan_client_tool_blocks, ) -from turnstone.core.trajectory import ProviderNative, ToolCall, Turn, dicts_from_turns +from turnstone.core.trajectory import ( + ProviderNative, + ToolCall, + Turn, + dicts_from_turns, + materialize_attachments, +) log = get_logger(__name__) @@ -481,6 +489,9 @@ class ModelLane: # deployment-wide model.auth_fail_closed switch remains a live per-mint # policy read. backend_auth_config: ModelConfig | None = None + # Stable per-alias registry gate. The gate object survives hot-resizes, + # so old and newly resolved lanes coordinate through one FIFO. + admission: ModelAdmission | None = None @dataclass(frozen=True) @@ -554,6 +565,7 @@ def same_model_lane_binding(left: ModelLane, right: ModelLane) -> bool: and left.provider is right.provider and left.model == right.model and left.alias == right.alias + and left.admission is right.admission ) @@ -673,6 +685,7 @@ def resolve_lane( cfg: ModelConfig | None | EllipsisType = ..., config_store: Any | None = None, backend_auth_resolver: Callable[[str, ModelConfig | None], str | None] | None = None, + admission: ModelAdmission | None = None, ) -> ModelLane: """Build a :class:`ModelLane`, resolving what the caller didn't supply. @@ -705,6 +718,16 @@ def resolve_lane( if extra_params is ... else extra_params ) + if admission is None and registry is not None and alias: + # Direct registry-backed lane rebuilds (notably judge sampling lanes) + # still join the alias's stable gate. Duck-typed test registries may + # manufacture attributes dynamically, hence the concrete type check. + try: + candidate = registry.get_admission(alias) + except Exception: + candidate = None + if isinstance(candidate, ModelAdmission): + admission = candidate return ModelLane( provider=provider, client=client, @@ -717,6 +740,7 @@ def resolve_lane( reasoning_effort=resolve_effort_setting(resolved_cfg, config_store), backend_auth_resolver=backend_auth_resolver, backend_auth_config=resolved_cfg, + admission=admission, ) @@ -733,7 +757,14 @@ def resolve_model_binding( # empty alias on a default binding disables registry-backed live flags and # delegated backend authentication on every later plant call. effective_alias = alias or registry.default - client, model, cfg, provider, generation = registry.resolve_binding(effective_alias) + resolved: Any = registry.resolve_binding(effective_alias) + if len(resolved) == 6: + client, model, cfg, provider, admission, generation = resolved + else: + # Compatibility for lightweight registry fakes that predate admission; + # real ModelRegistry snapshots always take the six-value branch. + client, model, cfg, provider, generation = resolved + admission = None lane = resolve_lane( provider, client, @@ -743,6 +774,7 @@ def resolve_model_binding( cfg=cfg, config_store=config_store, backend_auth_resolver=backend_auth_resolver, + admission=admission, ) return ResolvedModelBinding(lane=lane, config=cfg, registry_generation=generation) @@ -1087,11 +1119,11 @@ def model_turn( operator- or user-resolved knob (the session's own knobs, a CLI flag). *resolve_attachments* materializes by-reference ``AttachmentRef`` - content at the provider translator (``{type: kind, attachment_id}`` + content immediately before admission (``{type: kind, attachment_id}`` placeholders → inline parts; one id may expand to several parts, e.g. - a rasterized PDF). Turn IR never carries inline media bytes — a lane - with non-text content passes refs plus this resolver, exactly like the - main loop's wire path. + a rasterized PDF). This ordering keeps any nested perception/audio work + outside the outer alias's gate, avoiding self-deadlock at a limit of one. + Turn IR itself never carries inline media bytes. *mint* rewrites each returned tool call's id (provider-original → caller-scoped) before the Turn is built; the native blocks keep the @@ -1190,7 +1222,8 @@ def model_turn( override its ``x-api-key``, so an injected header is silently dropped. ``None`` leaves the lane's static client credential in place. When the explicit argument is absent, ``lane.backend_auth_resolver`` resolves it - once before the drain-retry loop. + after admission for each transport attempt, so a queued call cannot age a + minted credential before it reaches the wire. """ if mint is not None and wire_id_map is None: raise ValueError( @@ -1235,81 +1268,94 @@ def model_turn( or (lane.capabilities.default_reasoning_effort if lane.capabilities else None) or None ) - # A Stop can land while deterministic wire preparation runs. Re-read the - # abort signal before a possibly networked credential mint; the later read - # remains necessary for cancellation that lands while the mint itself is - # blocked. - call_client = lane_call_client( - lane, - backend_auth_token=backend_auth_token, - cancel_ref=cancel_ref, - ) + # Materialization may perform storage reads and nested perception/audio + # sampling. Complete it before taking the outer alias's admission slot so + # a cap of one cannot deadlock on a nested call that needs the same alias. + served_wire = materialize_attachments(wire, resolve_attachments) # A partially-surfaced stream is never silently re-issued — the # streaming caller owns re-issue. drain_retries = 0 if on_chunk is not None else _DRAIN_RETRIES attempt = 0 request_metrics: list[ProviderRequestMetrics] = [] while True: - # Last read before the wire — it covers everything the entry read is - # too early to see: the lowering, and on a delegated-auth alias the - # credential resolve, which can block. _raise_if_aborted(cancel_ref, lane) - # ``create_streaming`` stays OUTSIDE the try: every adapter issues - # the HTTP request eagerly in its body (inside the SDK's own - # request-level retry), so an exception from it is a request-time - # failure that already got its retries; only drain-time failures - # are mid-stream deaths the SDK could never see. - # Dynamic backends bind the token as the client's api_key so the SDK - # emits it as its own auth header; with_options reuses the pool. - # (extra_headers can't override the Anthropic SDK's x-api-key.) - chunks = lane.provider.create_streaming( - client=call_client, - model=lane.model, - messages=wire, - tools=tools, - max_tokens=max_tokens, - temperature=temperature if temperature is not None else lane.temperature, - reasoning_effort=effective_effort, - extra_params=lane.extra_params, - deferred_names=deferred_names, - cancel_ref=cancel_ref, - capabilities=lane.capabilities, - replay_reasoning_to_model=resolve_replay_reasoning_to_model( - lane.registry, lane.alias, caps=lane.capabilities, cfg=cfg - ), - resolve_attachments=resolve_attachments, - request_metrics_ref=request_metrics, - ) - try: - result = drain_stream( - _tee_chunks(chunks, on_chunk) if on_chunk else chunks, - scan_inline_reasoning=lane_scans_inline_reasoning(lane), + lease = lane.admission.acquire(cancel_ref=cancel_ref) if lane.admission else None + drain_error: Exception | None = None + with lease if lease is not None else contextlib.nullcontext(): + # Admission precedes a dynamic credential mint. This work and the + # full create+drain remain inside the hold; the context exits before + # any retry backoff below. + call_client = lane_call_client( + lane, + backend_auth_token=backend_auth_token, + cancel_ref=cancel_ref, ) - break - except Exception as exc: - attempt += 1 - if ( - attempt > drain_retries - or bool(getattr(cancel_ref, "aborted", False)) - or type(exc).__name__ not in lane.provider.retryable_error_names - ): - raise - delay = _DRAIN_RETRY_BASE_DELAY * (2 ** (attempt - 1)) * (0.5 + random.random()) - log.warning( - "model_turn.drain_retry", - error_type=type(exc).__name__, - attempt=attempt, + _raise_if_aborted(cancel_ref, lane) + mark_dispatch = getattr(cancel_ref, "mark_dispatch", None) + if callable(mark_dispatch): + with contextlib.suppress(Exception): + mark_dispatch() + # ``create_streaming`` stays OUTSIDE the drain-error catch: every + # adapter issues eagerly in its body, so a request-time failure has + # already received the SDK's own retries and propagates unchanged. + chunks = lane.provider.create_streaming( + client=call_client, model=lane.model, - retry_in=round(delay, 2), + messages=served_wire, + tools=tools, + max_tokens=max_tokens, + temperature=temperature if temperature is not None else lane.temperature, + reasoning_effort=effective_effort, + extra_params=lane.extra_params, + deferred_names=deferred_names, + cancel_ref=cancel_ref, + capabilities=lane.capabilities, + replay_reasoning_to_model=resolve_replay_reasoning_to_model( + lane.registry, lane.alias, caps=lane.capabilities, cfg=cfg + ), + # Already materialized before admission; provider translators + # retain their no-op fallback for direct callers. + resolve_attachments=None, + request_metrics_ref=request_metrics, ) - if delay > 0: - time.sleep(delay) - if bool(getattr(cancel_ref, "aborted", False)): - # The deadline abandoned this worker while it was backing off. - # The loop-top read would stop the re-issue anyway; this arm - # exists to die with the ORIGINAL transport failure rather than - # the abandonment error, so the cause of the death survives. - raise + try: + result = drain_stream( + _tee_chunks(chunks, on_chunk) if on_chunk else chunks, + scan_inline_reasoning=lane_scans_inline_reasoning(lane), + ) + except Exception as exc: + drain_error = exc + finally: + close = getattr(chunks, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + if drain_error is None: + break + attempt += 1 + if ( + attempt > drain_retries + or bool(getattr(cancel_ref, "aborted", False)) + or type(drain_error).__name__ not in lane.provider.retryable_error_names + ): + raise drain_error + delay = _DRAIN_RETRY_BASE_DELAY * (2 ** (attempt - 1)) * (0.5 + random.random()) + log.warning( + "model_turn.drain_retry", + error_type=type(drain_error).__name__, + attempt=attempt, + model=lane.model, + retry_in=round(delay, 2), + ) + if delay > 0: + time.sleep(delay) + if bool(getattr(cancel_ref, "aborted", False)): + # The deadline abandoned this worker while it was backing off. + # The loop-top read would stop the re-issue anyway; this arm + # exists to die with the ORIGINAL transport failure rather than + # the abandonment error, so the cause of the death survives. + raise drain_error raw_calls: list[dict[str, Any]] = list(result.tool_calls or []) # Record blanks BEFORE the uuid back-fill: a back-filled id exists only @@ -1369,7 +1415,7 @@ def model_turn( finish_reason=result.finish_reason, usage=result.usage, tool_calls=raw_calls, - wire_msgs=wire, + wire_msgs=served_wire, producer=lane.provider.provider_name, serving_model=lane.model, tool_def_chars=( diff --git a/turnstone/core/perception.py b/turnstone/core/perception.py index 397ad796..28c33c39 100644 --- a/turnstone/core/perception.py +++ b/turnstone/core/perception.py @@ -23,8 +23,9 @@ one alias; a vision-only model covers image/PDF and is simply skipped for audio. The call goes through :func:`turnstone.core.model_turn.model_turn` (the shared plant-call seam, #827), so any provider works: the trajectory carries the attachment by reference and the pre-built OpenAI-shaped parts (``image_url`` / -``input_audio``) materialize at the provider translator via the -``resolve_attachments`` callback, exactly like the main loop's wire path. +``input_audio``) materialize in ``model_turn`` via the ``resolve_attachments`` +callback before the alias admission slot is acquired, exactly like the main loop's +wire path. """ from __future__ import annotations @@ -78,9 +79,9 @@ def describe( ``parts`` are pre-built OpenAI-shaped content parts — ``image_url`` for image/PDF-page perception, ``input_audio`` for audio. The trajectory - carries them by reference; ``model_turn`` hands the resolver to the - provider translator, which materializes the placeholder into these exact - parts (one ref may expand to many, e.g. a rasterized PDF). + carries them by reference; ``model_turn`` uses the resolver before model + admission to materialize the placeholder into these exact parts (one ref + may expand to many, e.g. a rasterized PDF). ``lane`` is the caller's already-resolved binding snapshot, so the modality gate and the plant call cannot observe different registry diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 4017a40c..c952df54 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -979,8 +979,9 @@ def _neutralize_untrusted_fences(text: str) -> str: def _neutralize_attachment_part(part: Any) -> Any: """Return an attachment part with textual trust-marker forgeries defanged. - Attachment placeholders are materialized after ordinary message folding, - so their text cannot rely on ``fold_system_turns`` for this boundary. Only + Attachment placeholders are materialized after ordinary message folding and + before model admission, so their text cannot rely on ``fold_system_turns`` + for this boundary. Only model-visible text and document strings are inspected; binary image/audio data and base64 PDF payloads retain their exact bytes. """ @@ -5921,8 +5922,8 @@ class ChatSession: ``self.messages`` is the canonical ``Turn`` trajectory; lowering it here emits by-reference attachments as ``{type: kind, attachment_id}`` - placeholders. The provider translator materializes them to inline bytes - via :meth:`_resolve_attachments` — resolution lives at the C layer.""" + placeholders. ``model_turn`` materializes them to inline wire parts via + :meth:`_resolve_attachments` before acquiring the alias admission slot.""" return self.system_messages + dicts_from_turns(self.messages) def _resolve_attachments( @@ -5936,10 +5937,10 @@ class ChatSession: """Resolve content-addressed attachment ids to inline wire content parts. The send-time materialization of the by-reference content lane: handed to - the provider translator, which calls it with the placeholder ids it finds - and expands each to the inline part the wire needs. Blobs are - batch-fetched from the content-addressed store; a pruned id resolves to - nothing and the translator drops its placeholder. + ``model_turn``, which calls it with the placeholder ids it finds and + expands each to the inline part the wire needs before model admission. + Blobs are batch-fetched from the content-addressed store; a pruned id + resolves to nothing and the shared materializer drops its placeholder. Kinds the active model can't ingest natively (pdf without ``supports_pdf``, audio without ``supports_audio_input``) are converted client-side here — @@ -12672,8 +12673,8 @@ class ChatSession: # A shared workstream may bind a new actor while this daemon is still # evaluating later items. Capture the initiating identity now; the - # daemon resolves one token after its initial cancellation check and - # reuses it for the whole batch. + # daemon carries this resolver into each admitted plant call so token + # refresh cannot switch to a successor principal. judge_principal = ( (self._mcp_effective_user_id or "") if principal_id is None else principal_id ).strip() @@ -20594,22 +20595,13 @@ class ChatSession: # invocation; the native lane carried here serves the WITHIN-RUN # reasoning continuity of the agent's own tool loop. same_lane = (lane.alias or "") == (primary_lane.alias or "") - # Resolve once per sub-agent run, outside its request retry loop. A - # fail-open ``None`` is an intentional static-client result, not an - # invitation for ``model_turn`` to resolve again through the lane's - # mutable session-principal callback after a shared-user handoff. + # Keep the initiating principal pinned for the whole sub-agent run, + # but defer each credential mint until model_turn has acquired this + # alias's admission slot. A long queue can outlive a bearer token; + # the principal-bound resolver refreshes at the dispatch boundary + # without consulting a later shared-workstream actor. cancel_scope.check() - try: - agent_backend_auth_token = self._model_backend_auth_token_for_principal( - lane.alias, - lane.backend_auth_config, - principal_id=agent_principal, - ) - except Exception: - cancel_scope.check() - raise - cancel_scope.check() - lane = dataclasses.replace(lane, backend_auth_resolver=None) + lane = self._lane_for_backend_auth_principal(lane, agent_principal) def _api_call( turns: list[Turn], @@ -20645,7 +20637,6 @@ class ChatSession: or (self.reasoning_effort if same_lane else None), mint=mint, wire_id_map=wire_id_map, - backend_auth_token=agent_backend_auth_token, cancel_ref=cancel_scope.cancel_ref, prepare_wire=lambda wire, serving_lane: self._prepare_lowered_wire_messages( wire, diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 32844757..60f74853 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -5537,6 +5537,7 @@ class PostgreSQLBackend: auth_mode: str = "static", obo_audience: str = "", obo_scopes: str = "", + max_concurrency: int = 0, ) -> None: from sqlalchemy.dialects import postgresql @@ -5562,6 +5563,7 @@ class PostgreSQLBackend: auth_mode=auth_mode, obo_audience=obo_audience, obo_scopes=obo_scopes, + max_concurrency=max_concurrency, created_by=created_by, created=now, updated=now, diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index 284d3d48..2c299ffd 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -2640,6 +2640,7 @@ class StorageBackend(Protocol): auth_mode: str = "static", obo_audience: str = "", obo_scopes: str = "", + max_concurrency: int = 0, ) -> None: """Create a model definition. No-op if definition_id already exists.""" ... diff --git a/turnstone/core/storage/_schema.py b/turnstone/core/storage/_schema.py index 513d0265..32e52dc7 100644 --- a/turnstone/core/storage/_schema.py +++ b/turnstone/core/storage/_schema.py @@ -833,6 +833,7 @@ model_definitions = sa.Table( sa.Column("base_url", sa.Text, nullable=False, server_default=""), sa.Column("api_key", sa.Text, nullable=False, server_default=""), sa.Column("context_window", sa.Integer, nullable=False, server_default="32768"), + sa.Column("max_concurrency", sa.Integer, nullable=False, server_default="0"), sa.Column("capabilities", sa.Text, nullable=False, server_default="{}"), sa.Column("enabled", sa.Integer, nullable=False, server_default="1"), sa.Column("temperature", sa.Float, nullable=True), diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index f5266580..aa8f31ab 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -5677,6 +5677,7 @@ class SQLiteBackend: auth_mode: str = "static", obo_audience: str = "", obo_scopes: str = "", + max_concurrency: int = 0, ) -> None: now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") @@ -5701,6 +5702,7 @@ class SQLiteBackend: "auth_mode": auth_mode, "obo_audience": obo_audience, "obo_scopes": obo_scopes, + "max_concurrency": max_concurrency, "created_by": created_by, "created": now, "updated": now, diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index b9fc199a..bdcbace8 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -682,6 +682,7 @@ MODEL_DEFINITION_MUTABLE = frozenset( "base_url", "api_key", "context_window", + "max_concurrency", "capabilities", "enabled", "temperature", diff --git a/turnstone/core/storage/migrations/versions/070_model_max_concurrency.py b/turnstone/core/storage/migrations/versions/070_model_max_concurrency.py new file mode 100644 index 00000000..67ef8748 --- /dev/null +++ b/turnstone/core/storage/migrations/versions/070_model_max_concurrency.py @@ -0,0 +1,29 @@ +"""Add the per-alias model concurrency limit. + +``max_concurrency`` bounds concurrent model generations for one alias in one +process. Zero preserves the pre-feature unlimited behavior for every +existing definition. + +Revision ID: 070 +Revises: 069 +Create Date: 2026-08-08 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "070" +down_revision = "069" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column( + "model_definitions", + sa.Column("max_concurrency", sa.Integer, nullable=False, server_default="0"), + ) + + +def downgrade() -> None: + op.drop_column("model_definitions", "max_concurrency") diff --git a/turnstone/core/trajectory.py b/turnstone/core/trajectory.py index 87719047..8f01d69c 100644 --- a/turnstone/core/trajectory.py +++ b/turnstone/core/trajectory.py @@ -12,7 +12,7 @@ Field set and rationale: ``docs/design/canonical-trajectory-ideal-target.md`` § NOTE: non-text content rides as ``AttachmentRef`` — a reference to a content-addressed blob in ``workstream_attachments``. ``Turn``s never carry bytes, and the dict bridge carries only the ``{type: kind, attachment_id}`` placeholder. Each output boundary (the -provider translator, the ``/history`` display, export) materializes the placeholder to an +model-dispatch path, the ``/history`` display, export) materializes the placeholder to an inline part by point-lookup against the blob store, via :func:`resolve_attachment_parts`. """ @@ -68,9 +68,9 @@ class TextBlock: class AttachmentRef: """A reference to attachment bytes held in the content-addressed blob store. - Non-text content is carried *by reference* (never inline bytes): the translator - resolves ``attachment_id`` to bytes and expands it to the provider's native - format at wire time. ``kind`` is the by-reference placeholder type — + Non-text content is carried *by reference* (never inline bytes): the model-dispatch + path resolves ``attachment_id`` to bytes and expands it to the provider-neutral + inline wire part before provider admission. ``kind`` is the placeholder type — ``"image"``, ``"document"`` (text docs), ``"pdf"``, or ``"audio"``. The dict-bridge keys off ``attachment_id`` and is kind-agnostic, so new kinds need no change here. @@ -437,10 +437,10 @@ def resolve_attachment_parts( ``{type: kind, attachment_id}`` placeholders in a message's list content; *parts_by_id* maps an id to its inline content part — or a *list* of parts (one placeholder may expand to several, e.g. a PDF rasterized to one image - per page for a vision model) — built from the content-addressed blob. This is the materialization the - translator — and reconstruct, for display — runs at its output boundary: a - placeholder whose blob is missing (pruned) is dropped, so a consumer never - sees an unresolved reference. Identity-preserving when no message carries a + per page for a vision model) — built from the content-addressed blob. This + shared substitution runs at each output boundary (model dispatch or display): + a placeholder whose blob is missing (pruned) is dropped, so a consumer never + sees an unresolved reference. Identity-preserving when no message carries a placeholder; never mutates the input. """ @@ -480,9 +480,9 @@ def materialize_attachments( ) -> list[dict[str, Any]]: """Expand by-reference attachment placeholders to inline parts at the wire. - The translator's entry point for the by-reference content lane: collect the - placeholder ids across *messages*, ask *resolve* (a storage point-lookup the - session hands down) for their inline content parts, and substitute via + The shared wire-boundary entry point for the by-reference content lane: + collect the placeholder ids across *messages*, ask *resolve* (a storage + point-lookup the session hands down) for their inline content parts, and substitute via :func:`resolve_attachment_parts`. A ``None`` resolver (no storage — e.g. a unit test or an in-memory sub-agent whose media is already inline) or a placeholder-free trajectory is a no-op, so the common path is allocation-free. diff --git a/turnstone/server.py b/turnstone/server.py index d9a0e915..b3a60e77 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -4484,6 +4484,7 @@ def internal_model_reload(request: Request) -> JSONResponse: from turnstone.core.model_registry import ( DynamicAuthKeyError, ModelAuthConfigError, + ModelConcurrencyConfigError, load_model_registry, ) from turnstone.core.storage._registry import get_storage @@ -4503,10 +4504,10 @@ def internal_model_reload(request: Request) -> JSONResponse: provider=cli_args["provider"], storage=storage, ) - except ModelAuthConfigError as exc: - # A row whose auth fields violate _normalize_auth_mode, reachable via - # direct SQL, a migration mishap, or console version skew (console - # writes are validated). The loader deliberately propagates it; this + except (ModelAuthConfigError, ModelConcurrencyConfigError) as exc: + # A row whose auth or concurrency fields violate registry validation, + # reachable via direct SQL, a migration mishap, or console version skew + # (console writes are validated). The loader deliberately propagates it; this # arm keeps the node from answering a bare 500 where the console's # refresh path catches the same row. Same structured 422 contract as # the bad-arguments arm below: the reason names the row and field. @@ -4530,8 +4531,12 @@ def internal_model_reload(request: Request) -> JSONResponse: # No-op fast path: skip reload when nothing changed (avoids client churn # on broadcast model-reloads where this node has no pending changes). + admission_unchanged = { + alias: cfg.max_concurrency for alias, cfg in new_registry.models.items() + } == {alias: cfg.max_concurrency for alias, cfg in registry.models.items()} unchanged = ( new_registry.models == registry.models + and admission_unchanged and new_registry.fallback == registry.fallback and new_registry.agent_model == registry.agent_model and eff_default == registry.default @@ -4624,6 +4629,7 @@ def internal_model_status(request: Request) -> JSONResponse: "provider": cfg.provider, "source": cfg.source, "context_window": cfg.context_window, + "max_concurrency": cfg.max_concurrency, "enabled": True, "temperature": cfg.temperature, "max_tokens": cfg.max_tokens,