Add per-alias model concurrency admission (#990)

* feat(models): add per-alias concurrency admission

Add registry-backed FIFO admission limits with queue-aware deadlines and full-stream leases. Expose max_concurrency through storage, admin configuration, OpenAPI, documentation, and diagrams, with role and live backend count coverage.

* fix(api): omit null concurrency schema default

Keep max_concurrency optional for presence-keyed updates without advertising a null default for its non-null integer OpenAPI shape.
This commit is contained in:
Patrick Buckley
2026-08-08 22:01:04 -07:00
committed by GitHub
parent 7a06f5e8bc
commit 98e96ab5f3
43 changed files with 2298 additions and 201 deletions
+18 -2
View File
@@ -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),
+4 -1
View File
@@ -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
+22 -1
View File
@@ -134,6 +134,7 @@ class "ModelLane" as ModelLane <<frozen>> {
+ 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 <<frozen>> {
+ alias: str
+ provider: str
@@ -385,6 +403,7 @@ class "ModelConfig" as ModelCfg <<frozen>> {
+ 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
+6
View File
@@ -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
+2 -2
View File
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:e1431edf3891785b922c52b7897e3af5d39ba9973a815f892d6afdb762c5297b
size 612662
oid sha256:c74e99c530c3a8af9ab35b1e4d8c4fef0ea35c0c04cc35da7cf3588e71382057
size 661175
+2 -2
View File
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:79299b25ccc10484af13684a89ed9457abb9bc1781604bf6a9000ccb84362e55
size 311107
oid sha256:b237e00c28225c847b3e1a083c5c35535d1bf1549f6e6c7e2b6c4fd0e037a980
size 331906
+20
View File
@@ -54,6 +54,23 @@ When a per-model override is `NULL` (empty in the UI), the global default is
used. Switching models via `/model <alias>` 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
+31
View File
@@ -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,
@@ -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.
+23
View File
@@ -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
+66
View File
@@ -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
+129 -16
View File
@@ -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."""
+68
View File
@@ -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()
+199
View File
@@ -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
+462
View File
@@ -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
+19
View File
@@ -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(
+41 -5
View File
@@ -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"
+227 -1
View File
@@ -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
+20
View File
@@ -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.
+6 -7
View File
@@ -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:
+98
View File
@@ -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."""
+69
View File
@@ -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
+5 -7
View File
@@ -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)
+2
View File
@@ -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
+22 -1
View File
@@ -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
+34 -1
View File
@@ -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):
+25 -3
View File
@@ -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 !== "") {
+17
View File
@@ -1695,6 +1695,23 @@
/>
</div>
</div>
<div class="field-pair">
<div>
<label for="model-max-concurrency"
>Max concurrent generations
<span class="label-hint">0 = unlimited</span></label
>
<input
type="number"
id="model-max-concurrency"
value="0"
min="0"
max="2147483647"
step="1"
title="Per-process limit for model generations using this alias"
/>
</div>
</div>
<!-- Hidden outright on a deployment that cannot mint any dynamic
credential and whose row is not already using one; same
treatment the Roles sub-tab gets when its scope is missing. -->
+210
View File
@@ -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()
+69 -2
View File
@@ -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:
+27 -27
View File
@@ -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
+68 -4
View File
@@ -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
+119 -73
View File
@@ -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=(
+6 -5
View File
@@ -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
+17 -26
View File
@@ -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,
+2
View File
@@ -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,
+1
View File
@@ -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."""
...
+1
View File
@@ -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),
+2
View File
@@ -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,
+1
View File
@@ -682,6 +682,7 @@ MODEL_DEFINITION_MUTABLE = frozenset(
"base_url",
"api_key",
"context_window",
"max_concurrency",
"capabilities",
"enabled",
"temperature",
@@ -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")
+11 -11
View File
@@ -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.
+10 -4
View File
@@ -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,