mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
feat(models): per-alias backend auth via Entra OBO and app identity (#898)
Adds a per-alias `auth_mode` on model definitions so a model backend can authenticate to an Entra-fronted gateway with a per-request minted token instead of one shared static API key, letting the gateway attribute calls to the actual user or to the app as a machine identity. - `static` (default, unchanged) sends the stored `api_key`. - `entra_obo` mints a per-user On-Behalf-Of token for `obo_audience` from the caller's captured refresh credential. - `entra_app` mints an app-identity token via the client-credentials grant, and covers userless turns that OBO cannot. Reuses the existing OBO grant legs, refresh-token rotation CAS, cluster advisory lock and the `mcp_user_tokens` mint-cache, keyed under synthetic `__model_obo__:<audience>` / `__model_app__:<audience>` rows. The token binds at the call site through `client.with_options(api_key=...)` so each SDK emits it on its own auth path rather than through header injection. Migration 068 adds `auth_mode` and `obo_audience`. Both are additive and existing rows default to `static`, so behaviour is unchanged unless an alias opts in. Operator controls: `model.auth_audience_allowlist` is an exact-match allow-list that gates which audiences may be configured and denies all by default, and changing a mode or audience requires `admin.mcp`. `model.auth_fail_closed` decides whether a failed mint may fall back to an explicitly configured static key. A delegated call with no user, or a dynamic alias with no real static key, always refuses. Two changes here apply regardless of whether any alias opts in: - Storage and app state are now wired into the console MCP client manager. This fixes per-user `oauth_user` / `oauth_obo` dispatch for coordinator-hosted sessions, which previously raised `RuntimeError` on first call because `set_app_state` was only ever called on the node. - Unattended watch restores and `--resume` resolve the persisted workstream owner instead of constructing the session under an empty principal. A workstream with no owner is now a permanent refusal rather than an anonymous, auto-approved run.
This commit is contained in:
@@ -18,6 +18,19 @@ Earlier stable lines (`stable/1.6`, `stable/1.5`) are frozen.
|
||||
|
||||
### Added
|
||||
|
||||
- **Per-model Entra gateway authentication.** Model definitions can bind either
|
||||
a caller-delegated OBO token (`entra_obo`) or a shared app-identity token
|
||||
(`entra_app`) through the provider SDK credential surface. Mints reuse the
|
||||
encrypted cluster token cache, refresh-rotation CAS, and advisory locking;
|
||||
add a host-local memo, failure cooldown, long-lived mint HTTP client, audience
|
||||
allow-list/permission boundary, identity-unlink purge, and optional
|
||||
`model.auth_fail_closed` refusal policy. Delegated identity now propagates
|
||||
through judge, output-guard, and principal-scoped perception lanes, and
|
||||
unattended watch restoration reacquires the persisted workstream owner.
|
||||
Ownerless OBO calls and dynamic aliases without a real static fallback always
|
||||
fail closed; grant modes are never silently switched. Static authentication
|
||||
remains the default.
|
||||
|
||||
- **Compaction is visible now: lifecycle events, a progress bar, and a
|
||||
persistent transcript card.** Context compaction (manual `/compact` and
|
||||
auto) emits a first-class `compaction` SSE event
|
||||
|
||||
+2
-1
@@ -111,6 +111,7 @@ In the admin MCP form, choose **Sign-in passthrough** and set **Audience** (requ
|
||||
The captured credential is a single per-user secret that can mint for every `oauth_obo` server, so treat it like any long-lived credential:
|
||||
|
||||
- **Cut off one user:** unlink their OIDC identity in the admin console (**Users → OIDC identities → delete**). This revokes the captured credential **and** purges their minted cache rows, so future mints fail and cached tokens are dropped. (Warmed in-memory sessions on server nodes self-expire at the access-token TTL; there is no cross-node per-user session-kill.) Removing the user's access at the IdP is the authoritative cut-off.
|
||||
- The same unlink also purges that user's synthetic `__model_obo__:` gateway-token rows and requests eviction from every registered host's in-process mint memo. Shared `entra_app` model tokens live under the `__app__` pseudo-user and are intentionally not user-deprovisioned; revoking the app credential prevents new mints, while a cached app bearer lasts until `expires_at`.
|
||||
- **Flush a server's minted tokens** (e.g. after narrowing its audience): the server row's **flush cache** action drops all users' cached tokens for that server. This is **not** a revocation — users re-mint on next use from their still-valid sign-in. It is surfaced honestly (audit `mcp_server.oauth.obo_cache_flushed`, response `effect: cache_flush_remints`) so it is never mistaken for cutting access.
|
||||
- Per-server revocation in the `oauth_user` sense does not exist for `oauth_obo` — the credential is issuer-scoped and IdP-governed. Revoke at the IdP.
|
||||
|
||||
@@ -128,7 +129,7 @@ The captured credential is a single per-user secret that can mint for every `oau
|
||||
|
||||
4. **Step-up scope**: when a tool call hits `403` with `WWW-Authenticate: error="insufficient_scope"`, Turnstone emits `mcp_insufficient_scope` with the parsed scope set; the dashboard offers a "Connect with additional scopes" affordance that opens `/v1/api/mcp/oauth/start?server=<name>&scopes=<extra>` so the union of original + new scopes flows into the AS authorize request.
|
||||
|
||||
5. **User revoke** (settings modal): `DELETE /v1/api/mcp/oauth/connections/{server_name}` runs the authoritative local delete + best-effort RFC 7009 upstream revoke (fire-and-forget, capped at 256 concurrent in-flight tasks). `oauth_obo` servers are excluded: their rows are mint cache, not consents — deleting one only forces a re-mint — so the connections list hides them and the endpoint refuses them with `409` (revocation for sign-in passthrough happens at the identity layer: unlink the identity or revoke at the IdP).
|
||||
5. **User revoke** (settings modal): `DELETE /v1/api/mcp/oauth/connections/{server_name}` runs the authoritative local delete + best-effort RFC 7009 upstream revoke (fire-and-forget, capped at 256 concurrent in-flight tasks). `oauth_obo` servers and synthetic model-auth rows are excluded: their rows are mint caches, not consents — deleting one only forces a re-mint — so the connections list hides them and the endpoint refuses them with `409` (revocation for sign-in passthrough happens at the identity layer: unlink the identity or revoke at the IdP).
|
||||
|
||||
6. **Admin bulk-revoke** (Phase 9): `POST /v1/api/admin/mcp-servers/{name}/bulk-revoke` drops every user's token for the server. Upstream RFC 7009 revoke is intentionally **not** attempted in bulk (avoids N upstream HTTP calls per admin click); tokens at the AS expire naturally. Use the per-user revoke endpoint if you need guaranteed upstream invalidation.
|
||||
|
||||
|
||||
@@ -134,6 +134,24 @@ This knob only affects the login-flow IdP configured here. OAuth
|
||||
endpoints advertised by remote MCP servers are untrusted input and are
|
||||
always held to the strict public-address rule.
|
||||
|
||||
### Model gateway credentials
|
||||
|
||||
The same OIDC registration can authenticate model gateways. A model definition
|
||||
with `auth_mode = "entra_obo"` redeems the driving user's captured credential
|
||||
for its exact `obo_audience`; `auth_mode = "entra_app"` uses the registration's
|
||||
client ID and secret with Entra client credentials. Both bind the result through
|
||||
the provider SDK's native credential option rather than injecting an override
|
||||
header. The grant mode is never inferred: missing user context or a failed OBO
|
||||
mint cannot switch an `entra_obo` definition to client credentials.
|
||||
|
||||
`entra_obo` requires `capture_user_credential = true`, the MCP encryption key,
|
||||
and delegated/admin-consented permission to the audience. `entra_app` requires
|
||||
`obo_grant_profile = "entra"` and a confidential-client secret; RFC 8693
|
||||
client-credentials is not implemented. Configure the permitted resource IDs in
|
||||
the runtime setting `model.auth_audience_allowlist` before saving dynamic model
|
||||
definitions. See [Settings](settings.md#model-backend-authentication) for
|
||||
permissions, failure policy, and lane identity rules.
|
||||
|
||||
### config.toml alternative
|
||||
|
||||
```toml
|
||||
|
||||
+46
-1
@@ -54,6 +54,51 @@ 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.
|
||||
|
||||
### Model backend authentication
|
||||
|
||||
Model definitions support three backend credential modes:
|
||||
|
||||
| `auth_mode` | Identity sent to the model gateway |
|
||||
|-------------|------------------------------------|
|
||||
| `static` | The definition's stored `api_key`. |
|
||||
| `entra_obo` | A caller-delegated Entra access token minted from that user's captured OIDC credential. |
|
||||
| `entra_app` | A shared app-identity token minted with Turnstone's OIDC client credentials. |
|
||||
|
||||
Dynamic modes require an exact `obo_audience` resource App ID URI. Before an
|
||||
admin can save one, an operator must add that literal audience to
|
||||
`model.auth_audience_allowlist` (comma- or newline-separated). Wildcards and
|
||||
base-URL host matching are intentionally unsupported. Changing dynamic auth,
|
||||
its audience, or the gateway `base_url` also requires `admin.mcp`; service
|
||||
tokens do not bypass this capability-escalation gate.
|
||||
|
||||
`entra_app` is supported only with `[oidc] obo_grant_profile = "entra"`.
|
||||
Judge, output-guard, perception, utility, and sub-agent lanes inherit the
|
||||
session's effective user for `entra_obo`. The perception memo is partitioned by
|
||||
that principal as well as alias and content hash, so a result authorized as one
|
||||
user cannot be served to another. Scheduled and wake-driven work retains the
|
||||
workstream owner even when no user is connected. Eval and optimizer lanes are
|
||||
registry-less development tools and therefore do not use dynamic model
|
||||
authentication.
|
||||
|
||||
`entra_app` is an explicit model-definition choice; Turnstone never changes a
|
||||
failed or ownerless `entra_obo` call into a client-credentials grant. An
|
||||
`entra_obo` call with no effective user always refuses. A dynamic alias without
|
||||
a real static key also always refuses instead of issuing its SDK-construction
|
||||
placeholder. When a real static key is explicitly configured, mint failures
|
||||
may use it by default; set `model.auth_fail_closed = true` to prohibit even that
|
||||
fallback. A refusal is not routed through the model fallback chain.
|
||||
|
||||
Dynamic token caches are encrypted in `mcp_user_tokens`, shared across nodes,
|
||||
and memoized on each host. Unlinking a user's OIDC identity purges their
|
||||
`entra_obo` rows and memo entries. `entra_app` rows belong to the shared
|
||||
`__app__` identity and are not user-deprovisioned; after client-credential
|
||||
revocation, an already-minted app bearer remains usable until its recorded
|
||||
expiry.
|
||||
|
||||
`obo_audience` is literal and capped at 2048 characters. Environment-variable
|
||||
expansion is deliberately not applied, so the allow-list decision cannot vary
|
||||
by node or expand beyond the persisted boundary.
|
||||
|
||||
### Responses output controls (per-model)
|
||||
|
||||
Models whose capability table declares Responses output controls expose two
|
||||
@@ -124,7 +169,7 @@ initialization:
|
||||
|
||||
| Section | Settings |
|
||||
|---------|----------|
|
||||
| `model` | default_alias, temperature, max_tokens, reasoning_effort, task_alias, task_effort |
|
||||
| `model` | default_alias, auth_audience_allowlist, auth_fail_closed, temperature, max_tokens, reasoning_effort, task_alias, task_effort |
|
||||
| `session` | instructions, retention_days, compact_max_tokens, auto_compact_pct |
|
||||
| `tools` | timeout, truncation, agent_max_turns, skip_permissions, search, search_threshold, search_max_results |
|
||||
| `server` | workstream_idle_timeout, max_workstreams |
|
||||
|
||||
@@ -10824,6 +10824,16 @@
|
||||
"title": "Replay Reasoning To Model",
|
||||
"type": "boolean"
|
||||
},
|
||||
"auth_mode": {
|
||||
"default": "static",
|
||||
"title": "Auth Mode",
|
||||
"type": "string"
|
||||
},
|
||||
"obo_audience": {
|
||||
"default": "",
|
||||
"title": "Obo Audience",
|
||||
"type": "string"
|
||||
},
|
||||
"source": {
|
||||
"default": "",
|
||||
"title": "Source",
|
||||
@@ -10938,6 +10948,16 @@
|
||||
"default": false,
|
||||
"title": "Replay Reasoning To Model",
|
||||
"type": "boolean"
|
||||
},
|
||||
"auth_mode": {
|
||||
"default": "static",
|
||||
"title": "Auth Mode",
|
||||
"type": "string"
|
||||
},
|
||||
"obo_audience": {
|
||||
"default": "",
|
||||
"title": "Obo Audience",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
@@ -11105,6 +11125,30 @@
|
||||
],
|
||||
"default": null,
|
||||
"title": "Replay Reasoning To Model"
|
||||
},
|
||||
"auth_mode": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"title": "Auth Mode"
|
||||
},
|
||||
"obo_audience": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"default": null,
|
||||
"title": "Obo Audience"
|
||||
}
|
||||
},
|
||||
"title": "UpdateModelDefinitionRequest",
|
||||
|
||||
@@ -86,6 +86,8 @@ def _seed_model_def(
|
||||
model: str,
|
||||
base_url: str = "http://localhost:8000/v1",
|
||||
enabled: bool = True,
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
) -> None:
|
||||
"""Insert a model definition row directly via the storage API."""
|
||||
storage.create_model_definition(
|
||||
@@ -99,6 +101,8 @@ def _seed_model_def(
|
||||
capabilities="{}",
|
||||
enabled=enabled,
|
||||
created_by="admin",
|
||||
auth_mode=auth_mode,
|
||||
obo_audience=obo_audience,
|
||||
)
|
||||
|
||||
|
||||
@@ -838,11 +842,198 @@ def _make_client(storage: SQLiteBackend, registry: ModelRegistry | None) -> Test
|
||||
app.state.collector.get_all_nodes.return_value = []
|
||||
app.state.proxy_client = MagicMock()
|
||||
app.state.config_store = MagicMock()
|
||||
app.state.config_store.get.side_effect = lambda key, default=None: (
|
||||
"api://approved" if key == "model.auth_audience_allowlist" else default
|
||||
)
|
||||
app.state.oidc_config = SimpleNamespace(obo_grant_profile="entra")
|
||||
client = TestClient(app)
|
||||
client.headers.update({"X-Test-User": "admin", "X-Test-Perms": "admin.models"})
|
||||
return client
|
||||
|
||||
|
||||
def test_create_rejects_unknown_auth_mode(storage: SQLiteBackend) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={"alias": "bad-auth", "model": "x", "auth_mode": "bogus"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "auth_mode" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_create_rejects_entra_obo_without_audience(storage: SQLiteBackend) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={"alias": "missing-aud", "model": "x", "auth_mode": "entra_obo"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "obo_audience" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_update_rejects_entra_obo_when_stored_audience_empty(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"auth_mode": "entra_obo"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "obo_audience" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_update_rejects_clearing_audience_on_entra_obo(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
_seed_model_def(
|
||||
storage,
|
||||
definition_id="m1",
|
||||
alias="local",
|
||||
model="m",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience="api://approved",
|
||||
)
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"obo_audience": ""},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "obo_audience" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_dynamic_auth_create_requires_admin_mcp(storage: SQLiteBackend) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={
|
||||
"alias": "gateway",
|
||||
"model": "x",
|
||||
"auth_mode": "entra_obo",
|
||||
"obo_audience": "api://approved",
|
||||
},
|
||||
)
|
||||
|
||||
assert resp.status_code == 403, resp.text
|
||||
assert "admin.mcp" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_dynamic_auth_create_rejects_unapproved_audience(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
client.headers.update({"X-Test-User": "admin", "X-Test-Perms": "admin.models,admin.mcp"})
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={
|
||||
"alias": "gateway",
|
||||
"model": "x",
|
||||
"auth_mode": "entra_obo",
|
||||
"obo_audience": "api://not-approved",
|
||||
},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "allowlist" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_dynamic_alias_base_url_change_requires_admin_mcp(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
_seed_model_def(
|
||||
storage,
|
||||
definition_id="m1",
|
||||
alias="local",
|
||||
model="m",
|
||||
base_url="https://approved.example/v1",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience="api://approved",
|
||||
)
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"base_url": "https://attacker.example/v1"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 403, resp.text
|
||||
assert "admin.mcp" in resp.json()["error"]
|
||||
assert storage.get_model_definition("m1")["base_url"] == "https://approved.example/v1"
|
||||
|
||||
|
||||
def test_entra_app_create_rejects_non_entra_profile(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
_seed_model_def(storage, definition_id="m1", alias="local", model="m")
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
client.app.state.oidc_config = SimpleNamespace(obo_grant_profile="rfc8693")
|
||||
client.headers.update({"X-Test-User": "admin", "X-Test-Perms": "admin.models,admin.mcp"})
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={
|
||||
"alias": "gateway",
|
||||
"model": "x",
|
||||
"auth_mode": "entra_app",
|
||||
"obo_audience": "api://approved",
|
||||
},
|
||||
)
|
||||
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "obo_grant_profile='entra'" in resp.json()["error"]
|
||||
|
||||
|
||||
def test_unchanged_dynamic_auth_fields_do_not_require_admin_mcp(
|
||||
storage: SQLiteBackend,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The admin form always submits both fields; equality, not presence,
|
||||
decides whether the capability-escalation permission is needed."""
|
||||
from turnstone.console import server as server_module
|
||||
|
||||
_seed_model_def(
|
||||
storage,
|
||||
definition_id="m1",
|
||||
alias="local",
|
||||
model="m",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience="api://approved",
|
||||
)
|
||||
client = _make_client(storage, _make_registry(alias="local", model="m"))
|
||||
monkeypatch.setattr(
|
||||
server_module,
|
||||
"_ensure_console_mcp_client",
|
||||
lambda _app: {"skipped": "test"},
|
||||
)
|
||||
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={
|
||||
"auth_mode": "entra_obo",
|
||||
"obo_audience": "api://approved",
|
||||
"temperature": 0.4,
|
||||
},
|
||||
)
|
||||
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
|
||||
def test_create_endpoint_refreshes_registry(storage: SQLiteBackend) -> None:
|
||||
"""POST /api/admin/model-definitions bumps the in-process registry
|
||||
so newly-spawned coord sessions see the new alias immediately."""
|
||||
|
||||
@@ -170,6 +170,16 @@ class TestRequiredScope:
|
||||
def test_post_unknown_path_needs_read(self):
|
||||
assert required_scope("POST", "/api/unknown") == "read"
|
||||
|
||||
def test_model_auth_cache_invalidation_needs_approve(self):
|
||||
assert required_scope("POST", "/api/_internal/model-auth-cache-invalidate") == "approve"
|
||||
assert (
|
||||
required_scope(
|
||||
"POST",
|
||||
"/v1/api/_internal/model-auth-cache-invalidate",
|
||||
)
|
||||
== "approve"
|
||||
)
|
||||
|
||||
def test_v1_post_send_needs_write(self):
|
||||
assert required_scope("POST", "/v1/api/workstreams/abc/send") == "write"
|
||||
|
||||
|
||||
+108
-3
@@ -89,6 +89,19 @@ class _InjectAuthNoMcpMiddleware(BaseHTTPMiddleware):
|
||||
return resp
|
||||
|
||||
|
||||
class _InjectServiceAuthMiddleware(BaseHTTPMiddleware):
|
||||
"""Inject the cluster service identity used for internal cache eviction."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="console-proxy",
|
||||
scopes=frozenset({"read", "approve", "service"}),
|
||||
token_source="console",
|
||||
)
|
||||
resp: Response = await call_next(request)
|
||||
return resp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -149,6 +162,7 @@ def _routes_with_internal() -> list[Mount]:
|
||||
internal_mcp_refresh_one,
|
||||
internal_mcp_reload,
|
||||
internal_mcp_status,
|
||||
internal_model_auth_cache_invalidate,
|
||||
)
|
||||
|
||||
return [
|
||||
@@ -168,6 +182,11 @@ def _routes_with_internal() -> list[Mount]:
|
||||
internal_mcp_reconnect_one,
|
||||
methods=["POST"],
|
||||
),
|
||||
Route(
|
||||
"/api/_internal/model-auth-cache-invalidate",
|
||||
internal_model_auth_cache_invalidate,
|
||||
methods=["POST"],
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
@@ -2466,6 +2485,62 @@ class TestInternalMcpReloadEndpoint:
|
||||
assert data["updated"] == ["c"]
|
||||
|
||||
|
||||
class TestInternalModelAuthCacheInvalidateEndpoint:
|
||||
def test_service_identity_evicts_only_delegated_model_memo(self) -> None:
|
||||
app = Starlette(
|
||||
routes=_routes_with_internal(),
|
||||
middleware=[Middleware(_InjectServiceAuthMiddleware)],
|
||||
)
|
||||
app.state.mcp_client = MagicMock()
|
||||
app.state.mcp_client.invalidate_model_mint_memo_sync.return_value = 2
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
response = client.post(
|
||||
"/v1/api/_internal/model-auth-cache-invalidate",
|
||||
json={"user_id": "user-x"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "ok", "evicted": 2}
|
||||
app.state.mcp_client.invalidate_model_mint_memo_sync.assert_called_once_with(
|
||||
user_id="user-x",
|
||||
server_prefix="__model_obo__:",
|
||||
)
|
||||
|
||||
def test_non_service_identity_is_rejected(self) -> None:
|
||||
app = Starlette(
|
||||
routes=_routes_with_internal(),
|
||||
middleware=[Middleware(_InjectAuthMiddleware)],
|
||||
)
|
||||
app.state.mcp_client = MagicMock()
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
response = client.post(
|
||||
"/v1/api/_internal/model-auth-cache-invalidate",
|
||||
json={"user_id": "user-x"},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
app.state.mcp_client.invalidate_model_mint_memo_sync.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("body", [{}, {"user_id": ""}, {"user_id": "bad\nid"}])
|
||||
def test_invalid_user_id_is_rejected(self, body: dict[str, str]) -> None:
|
||||
app = Starlette(
|
||||
routes=_routes_with_internal(),
|
||||
middleware=[Middleware(_InjectServiceAuthMiddleware)],
|
||||
)
|
||||
app.state.mcp_client = MagicMock()
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
response = client.post(
|
||||
"/v1/api/_internal/model-auth-cache-invalidate",
|
||||
json=body,
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
app.state.mcp_client.invalidate_model_mint_memo_sync.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _notify_nodes_mcp_refresh_one / _notify_nodes_mcp_reconnect_one
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -3229,11 +3304,19 @@ class TestEnsureConsoleMcpClient:
|
||||
lazy-construct/reconcile path, shared by the reload fan-out (all
|
||||
admin-write producers) and the operator POST /reload."""
|
||||
|
||||
def _app(self, *, manager: Any = None, config_path: Any = None) -> Any:
|
||||
def _app(
|
||||
self,
|
||||
*,
|
||||
manager: Any = None,
|
||||
config_path: Any = None,
|
||||
dynamic_model_auth: bool = False,
|
||||
) -> Any:
|
||||
import types
|
||||
|
||||
state = types.SimpleNamespace()
|
||||
state.auth_storage = MagicMock()
|
||||
state.coord_registry = MagicMock()
|
||||
state.coord_registry.has_dynamic_auth.return_value = dynamic_model_auth
|
||||
cs = MagicMock()
|
||||
cs.get.side_effect = lambda k, d=None: config_path if k == "mcp.config_path" else d
|
||||
state.config_store = cs
|
||||
@@ -3263,13 +3346,30 @@ class TestEnsureConsoleMcpClient:
|
||||
app = self._app(config_path="/etc/turnstone/mcp.json")
|
||||
out = _ensure_console_mcp_client(app)
|
||||
create.assert_called_once_with(
|
||||
"/etc/turnstone/mcp.json", storage=app.state.auth_storage
|
||||
"/etc/turnstone/mcp.json",
|
||||
storage=app.state.auth_storage,
|
||||
required=False,
|
||||
)
|
||||
inst = create.return_value
|
||||
assert app.state.mcp_client is inst
|
||||
inst.set_storage.assert_called_once_with(app.state.auth_storage)
|
||||
inst.set_app_state.assert_called_once_with(app.state)
|
||||
inst.reconcile_sync.assert_called_once_with(app.state.auth_storage)
|
||||
assert out is inst.reconcile_sync.return_value
|
||||
|
||||
def test_dynamic_model_auth_requires_manager_without_mcp_servers(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch("turnstone.core.mcp_client.create_mcp_client") as create:
|
||||
app = self._app(dynamic_model_auth=True)
|
||||
_ensure_console_mcp_client(app)
|
||||
|
||||
create.assert_called_once_with(
|
||||
None,
|
||||
storage=app.state.auth_storage,
|
||||
required=True,
|
||||
)
|
||||
|
||||
def test_nothing_configured_skips(self):
|
||||
"""create_mcp_client returning None (no DB rows, no file config)
|
||||
is a skip, not an error — and nothing is stored on app.state."""
|
||||
@@ -3293,7 +3393,12 @@ class TestEnsureConsoleMcpClient:
|
||||
|
||||
constructed: list[Any] = []
|
||||
|
||||
def _slow_create(config_path: Any = None, *, storage: Any = None) -> Any:
|
||||
def _slow_create(
|
||||
config_path: Any = None,
|
||||
*,
|
||||
storage: Any = None,
|
||||
required: bool = False,
|
||||
) -> Any:
|
||||
time.sleep(0.05)
|
||||
mgr = MagicMock()
|
||||
mgr.reconcile_sync.return_value = {"added": [], "removed": [], "updated": []}
|
||||
|
||||
@@ -974,6 +974,22 @@ class TestCreateMcpClient:
|
||||
with patch("turnstone.core.mcp_client.load_mcp_config", return_value={}):
|
||||
assert create_mcp_client(storage=storage) is None
|
||||
|
||||
def test_required_constructs_empty_manager_for_model_auth(self):
|
||||
"""Dynamic model auth needs the mint loop even with zero MCP servers."""
|
||||
from turnstone.core.mcp_client import create_mcp_client
|
||||
|
||||
storage = MagicMock()
|
||||
storage.list_mcp_servers.return_value = []
|
||||
with (
|
||||
patch("turnstone.core.mcp_client.load_mcp_config", return_value={}),
|
||||
patch("turnstone.core.mcp_client.MCPClientManager") as cls,
|
||||
):
|
||||
result = create_mcp_client(storage=storage, required=True)
|
||||
|
||||
assert result is cls.return_value
|
||||
cls.assert_called_once_with({})
|
||||
cls.return_value.start.assert_called_once_with()
|
||||
|
||||
def test_pool_only_rows_construct_manager(self):
|
||||
"""Pool-backed rows alone must construct an empty-config manager.
|
||||
|
||||
|
||||
@@ -380,6 +380,25 @@ class TestListConnections:
|
||||
names = [c["server_name"] for c in resp.json()["connections"]]
|
||||
assert names == ["srv-oauth"]
|
||||
|
||||
def test_list_connections_hides_synthetic_model_cache_rows(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_oauth_user_server(storage)
|
||||
_seed_user_token(token_store)
|
||||
_seed_user_token(
|
||||
token_store,
|
||||
server_name="__model_obo__:api-model",
|
||||
refresh_token=None,
|
||||
)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
with TestClient(app) as client:
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert [row["server_name"] for row in resp.json()["connections"]] == ["srv-oauth"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DELETE /connections/{server_name}
|
||||
@@ -387,6 +406,24 @@ class TestListConnections:
|
||||
|
||||
|
||||
class TestRevokeConnection:
|
||||
def test_revoke_synthetic_model_cache_is_always_409_without_existence_oracle(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
token_store = _make_token_store(storage)
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
with TestClient(app) as client:
|
||||
absent = client.delete("/v1/api/mcp/oauth/connections/__model_obo__:absent")
|
||||
_seed_user_token(
|
||||
token_store,
|
||||
server_name="__model_obo__:present",
|
||||
refresh_token=None,
|
||||
)
|
||||
present = client.delete("/v1/api/mcp/oauth/connections/__model_obo__:present")
|
||||
|
||||
assert absent.status_code == present.status_code == 409
|
||||
assert absent.json() == present.json()
|
||||
assert token_store.get_user_token("user-1", "__model_obo__:present") is not None
|
||||
|
||||
def test_revoke_connection_obo_server_409_and_keeps_row(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
|
||||
@@ -2407,6 +2407,10 @@ class TestUserTokenFreshnessSweep:
|
||||
property — total invisibility to static / no-auth deployments."""
|
||||
|
||||
def _wire(self, mgr: MCPClientManager, storage: SQLiteBackend, cipher: Any) -> None:
|
||||
# The storage drive-set now authoritatively joins mcp_servers and keeps
|
||||
# oauth_user rows; mirror production instead of relying on orphan token
|
||||
# rows that the pre-model-cache query happened to enumerate.
|
||||
_seed_oauth_server(storage)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
mgr._oauth_user_server_names = {"pool-srv"}
|
||||
@@ -2772,6 +2776,8 @@ class TestUserTokenFreshnessSweep:
|
||||
|
||||
def test_reconcile_targets_pairs_expiry_unfiltered_with_last_exercised(self, storage) -> None:
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="srv-a", server_id="srv-a-id")
|
||||
_seed_oauth_server(storage, name="srv-b", server_id="srv-b-id")
|
||||
# alice consents to two servers → two rows.
|
||||
_seed_user_token(storage, cipher, user_id="alice", server_name="srv-a")
|
||||
_seed_user_token(storage, cipher, user_id="alice", server_name="srv-b")
|
||||
@@ -2780,6 +2786,14 @@ class TestUserTokenFreshnessSweep:
|
||||
_seed_user_token(
|
||||
storage, cipher, user_id="bob", server_name="srv-a", expires_in_seconds=-999
|
||||
)
|
||||
# Synthetic model mint-cache rows share the table but never drive the
|
||||
# oauth_user refresh-token keepalive sweep.
|
||||
_seed_user_token(
|
||||
storage,
|
||||
cipher,
|
||||
user_id="alice",
|
||||
server_name="__model_obo__:api-model",
|
||||
)
|
||||
targets = storage.list_mcp_user_token_reconcile_targets()
|
||||
# (user, server) identity, all three grants present regardless of expiry.
|
||||
assert sorted((u, s) for u, s, _ in targets) == [
|
||||
|
||||
@@ -0,0 +1,840 @@
|
||||
"""Tests for per-user OBO auth on model backends (auth_mode='entra_obo').
|
||||
|
||||
Covers the whole thin feature that lets a model backend authenticate to its
|
||||
gateway with a per-user Entra On-Behalf-Of access token instead of one static
|
||||
``api_key``:
|
||||
|
||||
* migration 068 — the two ``model_definitions`` columns, defaulting existing
|
||||
rows to the pre-feature ``static`` behaviour;
|
||||
* storage + admin-load round-trip of ``auth_mode`` / ``obo_audience``;
|
||||
* :func:`mint_obo_access_token` — the model-provider mint (reuses the MCP OBO
|
||||
grant legs + rotation write-back, but with an in-process token cache and no
|
||||
per-server machinery);
|
||||
* ``ModelRegistry.get_client`` constructing an OBO backend that has no static
|
||||
fallback key;
|
||||
* ``ChatSession._model_backend_auth_token`` — the per-call token resolve at the
|
||||
model call site. Ownerless OBO and keyless dynamic-auth failures refuse;
|
||||
explicit static keys remain an operator-controlled mint-failure fallback.
|
||||
The token is bound via ``client.with_options(api_key=...)`` at the call site
|
||||
(not injected as a header — the Anthropic SDK ignores an ``extra_headers``
|
||||
``x-api-key`` override).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.judge import JudgeConfig
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.mcp_oauth import mint_app_access_token, mint_obo_access_token
|
||||
from turnstone.core.model_registry import (
|
||||
ModelAuthConfigError,
|
||||
ModelConfig,
|
||||
ModelRegistry,
|
||||
load_model_registry,
|
||||
)
|
||||
from turnstone.core.oidc import OIDCConfig
|
||||
from turnstone.core.session import BackendAuthUnavailableError, ChatSession
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
USER = "user-1"
|
||||
ISSUER = "https://idp.test"
|
||||
TOKEN_ENDPOINT = "https://idp.test/token"
|
||||
AUDIENCE = "https://models.example.com"
|
||||
|
||||
_MIGRATIONS_DIR = str(
|
||||
Path(__file__).resolve().parent.parent / "turnstone" / "core" / "storage" / "migrations"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Migration 068
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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 TestMigration068:
|
||||
def test_upgrade_adds_auth_columns(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "068-up.db"
|
||||
command.upgrade(_alembic_cfg(db_path), "068")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
cols = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")}
|
||||
assert {"auth_mode", "obo_audience"} <= cols
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_preexisting_row_defaults_to_static(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "068-default.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "067")
|
||||
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', 'gpt', 'gpt-5', "
|
||||
"'2026-01-01T00:00:00', '2026-01-01T00:00:00')"
|
||||
)
|
||||
)
|
||||
command.upgrade(cfg, "068")
|
||||
with engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
sa.text(
|
||||
"SELECT auth_mode, obo_audience FROM model_definitions "
|
||||
"WHERE definition_id = 'd1'"
|
||||
)
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
# A pre-068 row keeps byte-identical behaviour: static, no audience.
|
||||
assert row[0] == "static" and row[1] == ""
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_downgrade_removes_auth_columns(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "068-down.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "068")
|
||||
command.downgrade(cfg, "067")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
cols = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")}
|
||||
assert "auth_mode" not in cols and "obo_audience" not in cols
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_downgrade_then_upgrade_round_trip(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "068-roundtrip.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "068")
|
||||
command.downgrade(cfg, "067")
|
||||
command.upgrade(cfg, "068")
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
cols = {c["name"] for c in sa.inspect(engine).get_columns("model_definitions")}
|
||||
assert {"auth_mode", "obo_audience"} <= cols
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Storage + admin-load round-trip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
class TestModelDefinitionStorage:
|
||||
def test_create_and_read_back_obo_fields(self, storage: SQLiteBackend) -> None:
|
||||
storage.create_model_definition(
|
||||
definition_id="d1",
|
||||
alias="tf-opus",
|
||||
model="vmg/opus",
|
||||
provider="anthropic",
|
||||
base_url="https://gateway.example.com",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
row = storage.get_model_definition_by_alias("tf-opus")
|
||||
assert row is not None
|
||||
assert row["auth_mode"] == "entra_obo"
|
||||
assert row["obo_audience"] == AUDIENCE
|
||||
|
||||
def test_defaults_static_when_unspecified(self, storage: SQLiteBackend) -> None:
|
||||
storage.create_model_definition(definition_id="d2", alias="plain", model="gpt-5")
|
||||
row = storage.get_model_definition_by_alias("plain")
|
||||
assert row is not None
|
||||
assert row["auth_mode"] == "static"
|
||||
assert row["obo_audience"] == ""
|
||||
|
||||
def test_update_toggles_auth_mode(self, storage: SQLiteBackend) -> None:
|
||||
storage.create_model_definition(definition_id="d3", alias="m3", model="gpt-5")
|
||||
assert storage.update_model_definition("d3", auth_mode="entra_obo", obo_audience=AUDIENCE)
|
||||
row = storage.get_model_definition("d3")
|
||||
assert row is not None
|
||||
assert row["auth_mode"] == "entra_obo" and row["obo_audience"] == AUDIENCE
|
||||
|
||||
def test_load_model_registry_carries_obo_fields(self, storage: SQLiteBackend) -> None:
|
||||
storage.create_model_definition(
|
||||
definition_id="d4",
|
||||
alias="tf",
|
||||
model="vmg/opus",
|
||||
provider="anthropic",
|
||||
base_url="https://gateway.example.com",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
registry = load_model_registry(storage=storage, allow_empty=True)
|
||||
cfg = registry.get_config("tf")
|
||||
assert cfg.auth_mode == "entra_obo"
|
||||
assert cfg.obo_audience == AUDIENCE
|
||||
|
||||
def test_invalid_db_auth_mode_is_not_swallowed_as_storage_failure(
|
||||
self,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
storage.create_model_definition(
|
||||
definition_id="bad-mode",
|
||||
alias="bad",
|
||||
model="m",
|
||||
auth_mode="bogus",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
|
||||
with pytest.raises(ModelAuthConfigError, match="invalid auth_mode"):
|
||||
load_model_registry(storage=storage, allow_empty=True)
|
||||
|
||||
def test_runtime_audience_rejects_control_characters(
|
||||
self,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
storage.create_model_definition(
|
||||
definition_id="bad-audience",
|
||||
alias="bad",
|
||||
model="m",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience="api://gateway\ninjected",
|
||||
)
|
||||
|
||||
with pytest.raises(ModelAuthConfigError, match="control characters"):
|
||||
load_model_registry(storage=storage, allow_empty=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ModelRegistry.get_client — OBO backend with no static fallback key
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetClientKeyInjection:
|
||||
"""get_client feeds a placeholder key ONLY for an entra_obo backend with no
|
||||
static fallback — everything else passes ``cfg.api_key`` through unchanged.
|
||||
Spies on ``create_client`` so it's independent of SDK/env behaviour."""
|
||||
|
||||
def _seen_api_key(self, cfg: ModelConfig, monkeypatch: Any) -> str:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def _spy(provider: str, *, base_url: str, api_key: str) -> object:
|
||||
seen["api_key"] = api_key
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr("turnstone.core.model_registry.create_client", _spy)
|
||||
ModelRegistry(models={cfg.alias: cfg}, default=cfg.alias).get_client(cfg.alias)
|
||||
return seen["api_key"]
|
||||
|
||||
def test_obo_blank_key_gets_placeholder(self, monkeypatch: Any) -> None:
|
||||
cfg = ModelConfig(
|
||||
alias="tf",
|
||||
base_url="u",
|
||||
api_key="", # no static fallback — real credential injected per call
|
||||
model="m",
|
||||
provider="anthropic",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
assert self._seen_api_key(cfg, monkeypatch) == "backend-auth-placeholder-unused"
|
||||
|
||||
def test_obo_with_static_key_keeps_it(self, monkeypatch: Any) -> None:
|
||||
cfg = ModelConfig(
|
||||
alias="tf",
|
||||
base_url="u",
|
||||
api_key="real-key",
|
||||
model="m",
|
||||
provider="anthropic",
|
||||
auth_mode="entra_obo",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
assert self._seen_api_key(cfg, monkeypatch) == "real-key"
|
||||
|
||||
def test_static_blank_key_unchanged(self, monkeypatch: Any) -> None:
|
||||
cfg = ModelConfig(alias="a", base_url="u", api_key="", model="m", provider="anthropic")
|
||||
# No placeholder for a static alias — the empty key rides through exactly
|
||||
# as before (create_client then coerces it to an env-var fallback).
|
||||
assert self._seen_api_key(cfg, monkeypatch) == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mint_obo_access_token — the model-provider mint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_oidc_config(**overrides: Any) -> OIDCConfig:
|
||||
defaults: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"issuer": ISSUER,
|
||||
"client_id": "cid",
|
||||
"client_secret": "csecret",
|
||||
"token_endpoint": TOKEN_ENDPOINT,
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return OIDCConfig(**defaults)
|
||||
|
||||
|
||||
def _make_app_state(
|
||||
storage: SQLiteBackend, *, http_client: httpx.AsyncClient, oidc_config: OIDCConfig
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, make_mcp_token_cipher(), node_id="test"),
|
||||
oidc_config=oidc_config,
|
||||
obo_http_client=http_client,
|
||||
mcp_oauth_refresh_locks={},
|
||||
)
|
||||
|
||||
|
||||
def _mk_response(status_code: int = 200, json_body: Any = None) -> MagicMock:
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
resp.headers = {}
|
||||
body = "" if json_body is None else str(json_body)
|
||||
resp.content = body.encode("utf-8")
|
||||
if json_body is not None:
|
||||
resp.json.return_value = json_body
|
||||
else:
|
||||
resp.json.side_effect = ValueError("no body")
|
||||
resp.text = body
|
||||
return resp
|
||||
|
||||
|
||||
def _seed_credential(state: SimpleNamespace, *, refresh_token: str = "rt-1") -> None:
|
||||
state.mcp_token_store.upsert_oidc_credential(USER, ISSUER, refresh_token=refresh_token)
|
||||
|
||||
|
||||
def _mint(state: SimpleNamespace, **kwargs: Any) -> Any:
|
||||
async def _run() -> Any:
|
||||
return await mint_obo_access_token(
|
||||
app_state=state, user_id=USER, audience=AUDIENCE, **kwargs
|
||||
)
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
class TestMintOboAccessToken:
|
||||
def test_happy_path_redeems_default_scope_and_caches(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(200, {"access_token": "at-minted", "expires_in": 3600})
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
_seed_credential(state)
|
||||
state.mcp_token_store.get_user_token = MagicMock( # type: ignore[method-assign]
|
||||
wraps=state.mcp_token_store.get_user_token
|
||||
)
|
||||
|
||||
token = _mint(state)
|
||||
|
||||
assert token == "at-minted"
|
||||
# Exact entra wire shape — scope pins <audience>/.default.
|
||||
assert client.post.call_count == 1
|
||||
call = client.post.call_args
|
||||
assert call.args == (TOKEN_ENDPOINT,)
|
||||
assert call.kwargs["data"] == {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": "rt-1",
|
||||
"client_id": "cid",
|
||||
"client_secret": "csecret",
|
||||
"scope": f"{AUDIENCE}/.default",
|
||||
}
|
||||
reads_after_mint = state.mcp_token_store.get_user_token.call_count
|
||||
# Second call serves the mcp-loop memo — no DB decrypt or IdP trip.
|
||||
token2 = _mint(state)
|
||||
assert token2 == "at-minted"
|
||||
assert client.post.call_count == 1
|
||||
assert state.mcp_token_store.get_user_token.call_count == reads_after_mint
|
||||
|
||||
def test_minted_token_cached_in_db_and_shared_across_nodes(
|
||||
self, storage: SQLiteBackend
|
||||
) -> None:
|
||||
# One shared enc key, as a cluster shares MCP_ENC_KEY across workers.
|
||||
cipher = make_mcp_token_cipher()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(200, {"access_token": "at-minted", "expires_in": 3600})
|
||||
)
|
||||
node_a = SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="A"),
|
||||
oidc_config=_make_oidc_config(),
|
||||
obo_http_client=client,
|
||||
mcp_oauth_refresh_locks={},
|
||||
)
|
||||
node_a.mcp_token_store.upsert_oidc_credential(USER, ISSUER, refresh_token="rt-1")
|
||||
|
||||
assert _mint(node_a) == "at-minted"
|
||||
assert client.post.call_count == 1
|
||||
|
||||
# Persisted as a "cache, not custody" row (refresh_token NULL), decodable.
|
||||
cache_server = f"__model_obo__:{AUDIENCE}"
|
||||
raw = storage.get_mcp_user_token(USER, cache_server)
|
||||
assert raw is not None and raw["refresh_token_ct"] is None
|
||||
plain = node_a.mcp_token_store.get_user_token(USER, cache_server)
|
||||
assert plain is not None
|
||||
assert plain["access_token"] == "at-minted"
|
||||
assert plain["audience"] == AUDIENCE
|
||||
|
||||
# A DIFFERENT worker (same DB + enc key) serves the cached token with NO
|
||||
# new IdP round-trip — no needless per-worker re-mint.
|
||||
node_b = SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="B"),
|
||||
oidc_config=_make_oidc_config(),
|
||||
obo_http_client=client,
|
||||
mcp_oauth_refresh_locks={},
|
||||
)
|
||||
assert _mint(node_b) == "at-minted"
|
||||
assert client.post.call_count == 1
|
||||
|
||||
def test_rotated_refresh_token_persisted_to_credential(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(
|
||||
200, {"access_token": "at", "expires_in": 3600, "refresh_token": "rt-2"}
|
||||
)
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
_seed_credential(state, refresh_token="rt-1")
|
||||
|
||||
assert _mint(state) == "at"
|
||||
cred = state.mcp_token_store.get_oidc_credential(USER, ISSUER)
|
||||
assert cred is not None and cred["refresh_token"] == "rt-2"
|
||||
|
||||
def test_force_refresh_bypasses_cache(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
side_effect=[
|
||||
_mk_response(200, {"access_token": "at-1", "expires_in": 3600}),
|
||||
_mk_response(200, {"access_token": "at-2", "expires_in": 3600}),
|
||||
]
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
_seed_credential(state)
|
||||
|
||||
assert _mint(state) == "at-1"
|
||||
assert _mint(state, force_refresh=True) == "at-2"
|
||||
assert client.post.call_count == 2
|
||||
|
||||
def test_missing_credential_returns_none_no_http(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
# No captured credential seeded.
|
||||
assert _mint(state) is None
|
||||
assert client.post.call_count == 0
|
||||
|
||||
def test_oidc_disabled_returns_none_no_http(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
state = _make_app_state(
|
||||
storage, http_client=client, oidc_config=_make_oidc_config(enabled=False)
|
||||
)
|
||||
_seed_credential(state)
|
||||
assert _mint(state) is None
|
||||
assert client.post.call_count == 0
|
||||
|
||||
def test_unusable_profile_returns_none_no_http(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
state = _make_app_state(
|
||||
storage, http_client=client, oidc_config=_make_oidc_config(obo_grant_profile="")
|
||||
)
|
||||
_seed_credential(state)
|
||||
assert _mint(state) is None
|
||||
assert client.post.call_count == 0
|
||||
|
||||
def test_permanent_rejection_returns_none(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(
|
||||
400, {"error": "invalid_grant", "error_description": "AADSTS65001"}
|
||||
)
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
_seed_credential(state)
|
||||
# A failed mint falls back to the static credential (None), and never
|
||||
# auto-deletes the shared credential.
|
||||
assert _mint(state) is None
|
||||
assert state.mcp_token_store.get_oidc_credential(USER, ISSUER) is not None
|
||||
# The audience-scoped cooldown suppresses a dead-grant retry storm.
|
||||
assert _mint(state) is None
|
||||
assert client.post.call_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mint_app_access_token — app-identity (client-credentials) mint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mint_app(state: SimpleNamespace, **kwargs: Any) -> Any:
|
||||
async def _run() -> Any:
|
||||
return await mint_app_access_token(app_state=state, audience=AUDIENCE, **kwargs)
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
class TestMintAppAccessToken:
|
||||
def test_happy_path_client_credentials_and_caches(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(200, {"access_token": "app-at", "expires_in": 3600})
|
||||
)
|
||||
# NOTE: no captured user credential seeded — app identity needs none.
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
|
||||
assert _mint_app(state) == "app-at"
|
||||
# Exact client-credentials wire shape — scope pins <audience>/.default.
|
||||
assert client.post.call_count == 1
|
||||
call = client.post.call_args
|
||||
assert call.args == (TOKEN_ENDPOINT,)
|
||||
assert call.kwargs["data"] == {
|
||||
"grant_type": "client_credentials",
|
||||
"client_id": "cid",
|
||||
"client_secret": "csecret",
|
||||
"scope": f"{AUDIENCE}/.default",
|
||||
}
|
||||
# Cached in the DB under the synthetic __app__ user — second call, no IdP.
|
||||
assert _mint_app(state) == "app-at"
|
||||
assert client.post.call_count == 1
|
||||
cache_server = f"__model_app__:{AUDIENCE}"
|
||||
raw = storage.get_mcp_user_token("__app__", cache_server)
|
||||
assert raw is not None and raw["refresh_token_ct"] is None
|
||||
|
||||
def test_works_with_zero_user_credentials(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(200, {"access_token": "app-at", "expires_in": 3600})
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
# The credential store is empty; the app grant still succeeds.
|
||||
assert state.mcp_token_store.get_oidc_credential(USER, ISSUER) is None
|
||||
assert _mint_app(state) == "app-at"
|
||||
|
||||
def test_oidc_disabled_returns_none_no_http(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
state = _make_app_state(
|
||||
storage, http_client=client, oidc_config=_make_oidc_config(enabled=False)
|
||||
)
|
||||
assert _mint_app(state) is None
|
||||
assert client.post.call_count == 0
|
||||
|
||||
def test_rejected_grant_returns_none(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
return_value=_mk_response(
|
||||
400, {"error": "invalid_client", "error_description": "AADSTS7000215"}
|
||||
)
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
assert _mint_app(state) is None
|
||||
assert _mint_app(state) is None
|
||||
assert client.post.call_count == 1
|
||||
|
||||
def test_non_entra_profile_is_explicitly_refused(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
state = _make_app_state(
|
||||
storage,
|
||||
http_client=client,
|
||||
oidc_config=_make_oidc_config(obo_grant_profile="rfc8693"),
|
||||
)
|
||||
|
||||
assert _mint_app(state) is None
|
||||
client.post.assert_not_called()
|
||||
|
||||
def test_force_refresh_bypasses_cache(self, storage: SQLiteBackend) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(
|
||||
side_effect=[
|
||||
_mk_response(200, {"access_token": "app-1", "expires_in": 3600}),
|
||||
_mk_response(200, {"access_token": "app-2", "expires_in": 3600}),
|
||||
]
|
||||
)
|
||||
state = _make_app_state(storage, http_client=client, oidc_config=_make_oidc_config())
|
||||
assert _mint_app(state) == "app-1"
|
||||
assert _mint_app(state, force_refresh=True) == "app-2"
|
||||
assert client.post.call_count == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ChatSession._model_backend_auth_token — resolve at the model call site
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _fake_session(
|
||||
*,
|
||||
registry: ModelRegistry | None,
|
||||
user_id: str | None,
|
||||
mint_token: str | None,
|
||||
app_token: str | None = None,
|
||||
) -> SimpleNamespace:
|
||||
"""Minimal stand-in exposing exactly what the backend-auth resolver reads."""
|
||||
mcp = SimpleNamespace(
|
||||
mint_model_obo_token_sync=MagicMock(return_value=mint_token),
|
||||
mint_app_token_sync=MagicMock(return_value=app_token),
|
||||
)
|
||||
return SimpleNamespace(
|
||||
_registry=registry,
|
||||
_mcp_mint_client=mcp,
|
||||
_mcp_effective_user_id=user_id,
|
||||
_config_store=None,
|
||||
)
|
||||
|
||||
|
||||
def _registry_with(cfg: ModelConfig) -> ModelRegistry:
|
||||
return ModelRegistry(models={cfg.alias: cfg}, default=cfg.alias)
|
||||
|
||||
|
||||
class TestModelOboToken:
|
||||
def _obo_cfg(
|
||||
self,
|
||||
provider: str = "anthropic",
|
||||
*,
|
||||
api_key: str = "static-fallback",
|
||||
) -> ModelConfig:
|
||||
return ModelConfig(
|
||||
alias="tf",
|
||||
base_url="https://gateway.example.com",
|
||||
api_key=api_key,
|
||||
model="vmg/opus",
|
||||
provider=provider,
|
||||
auth_mode="entra_obo",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
|
||||
def test_obo_alias_with_user_returns_token(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg())
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token="minted-jwt")
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") == "minted-jwt"
|
||||
sess._mcp_mint_client.mint_model_obo_token_sync.assert_called_once_with(
|
||||
user_id=USER, audience=AUDIENCE
|
||||
)
|
||||
|
||||
def test_token_is_provider_agnostic(self) -> None:
|
||||
# The raw token is returned regardless of provider surface — the caller
|
||||
# binds it via ``with_options(api_key=...)``, so there is no per-provider
|
||||
# header shaping here. (The old header-injection path returned x-api-key
|
||||
# for anthropic, which the SDK silently ignored → the prod 401.)
|
||||
for provider in ("anthropic", "openai-compatible"):
|
||||
reg = _registry_with(self._obo_cfg(provider=provider))
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token="minted-jwt")
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") == "minted-jwt"
|
||||
|
||||
def test_auxiliary_judges_inherit_the_session_obo_resolver(
|
||||
self,
|
||||
mock_openai_client: Any,
|
||||
) -> None:
|
||||
"""Judge lanes must not quietly regress to app-only authentication."""
|
||||
reg = _registry_with(self._obo_cfg(provider="openai"))
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
model="vmg/opus",
|
||||
ui=MagicMock(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=1024,
|
||||
tool_timeout=30,
|
||||
registry=reg,
|
||||
model_alias="tf",
|
||||
judge_config=JudgeConfig(
|
||||
enabled=True,
|
||||
output_guard_llm=True,
|
||||
),
|
||||
user_id=USER,
|
||||
)
|
||||
session._mcp_mint_client = SimpleNamespace(
|
||||
mint_model_obo_token_sync=MagicMock(return_value="minted-jwt"),
|
||||
mint_app_token_sync=MagicMock(return_value="app-jwt"),
|
||||
)
|
||||
|
||||
intent_judge = session._ensure_judge()
|
||||
output_guard = session._ensure_output_guard_judge()
|
||||
|
||||
assert intent_judge is not None
|
||||
assert output_guard is not None
|
||||
assert intent_judge._backend_auth_resolver == session._model_backend_auth_token
|
||||
assert output_guard._backend_auth_resolver == session._model_backend_auth_token
|
||||
assert intent_judge._backend_auth_resolver("tf") == "minted-jwt"
|
||||
session._mcp_mint_client.mint_model_obo_token_sync.assert_called_once_with(
|
||||
user_id=USER,
|
||||
audience=AUDIENCE,
|
||||
)
|
||||
session._mcp_mint_client.mint_app_token_sync.assert_not_called()
|
||||
|
||||
def test_primary_stream_binds_backend_token_once_before_retry_loop(self) -> None:
|
||||
"""The main streaming path mirrors model_turn's SDK binding."""
|
||||
sess = MagicMock()
|
||||
sess._MAX_RETRIES = 2
|
||||
sess._provider = MagicMock()
|
||||
sess._provider.create_streaming.return_value = iter(())
|
||||
sess._get_capabilities.return_value = SimpleNamespace(default_reasoning_effort=None)
|
||||
sess._maybe_attach_vllm_chat_reasoning.side_effect = lambda messages, _provider, _alias: (
|
||||
messages
|
||||
)
|
||||
sess._model_backend_auth_token.return_value = "minted-jwt"
|
||||
sess._cancel_ref = []
|
||||
sess._get_active_tools.return_value = []
|
||||
sess.max_tokens = 1024
|
||||
sess.temperature = None
|
||||
sess.reasoning_effort = None
|
||||
sess._provider_extra_params.return_value = None
|
||||
sess._get_deferred_names.return_value = frozenset()
|
||||
sess._resolve_replay_reasoning_to_model.return_value = False
|
||||
base_client = MagicMock()
|
||||
base_client.base_url = "https://gateway.example.com"
|
||||
bound_client = object()
|
||||
base_client.with_options.return_value = bound_client
|
||||
|
||||
stream = ChatSession._try_stream(
|
||||
sess,
|
||||
base_client,
|
||||
"vmg/opus",
|
||||
[{"role": "user", "content": "hi"}],
|
||||
model_alias="tf",
|
||||
)
|
||||
|
||||
assert list(stream) == []
|
||||
sess._model_backend_auth_token.assert_called_once_with("tf")
|
||||
base_client.with_options.assert_called_once_with(api_key="minted-jwt")
|
||||
assert sess._provider.create_streaming.call_args.kwargs["client"] is bound_client
|
||||
|
||||
def test_primary_stream_forwards_alias_for_obo(self) -> None:
|
||||
# Regression: the primary _create_stream_with_retry call must pass
|
||||
# model_alias, or the backend-auth resolver can't resolve the OBO token and an
|
||||
# entra_obo main turn goes out on the static client key. The fallback
|
||||
# path and utility (title) completions always passed the alias; the
|
||||
# primary path silently didn't.
|
||||
sess = MagicMock()
|
||||
sess._model_alias = "oboagent"
|
||||
ChatSession._create_stream_with_retry(sess, [{"role": "user", "content": "hi"}])
|
||||
sess._try_stream.assert_called_once()
|
||||
assert sess._try_stream.call_args.kwargs.get("model_alias") == "oboagent"
|
||||
|
||||
def test_fail_closed_refusal_never_enters_model_fallback_chain(self) -> None:
|
||||
sess = MagicMock()
|
||||
sess._model_alias = "oboagent"
|
||||
sess._try_stream.side_effect = BackendAuthUnavailableError("mint failed")
|
||||
sess._registry.fallback = ["static-backup"]
|
||||
sess._get_health_tracker.return_value = None
|
||||
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._create_stream_with_retry(sess, [{"role": "user", "content": "hi"}])
|
||||
|
||||
sess._try_fallback.assert_not_called()
|
||||
|
||||
def test_static_alias_returns_none_and_never_mints(self) -> None:
|
||||
static_cfg = ModelConfig(
|
||||
alias="plain",
|
||||
base_url="",
|
||||
api_key="k",
|
||||
model="gpt-5",
|
||||
provider="openai",
|
||||
)
|
||||
reg = _registry_with(static_cfg)
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token="unused")
|
||||
assert ChatSession._model_backend_auth_token(sess, "plain") is None
|
||||
sess._mcp_mint_client.mint_model_obo_token_sync.assert_not_called()
|
||||
|
||||
def test_no_user_context_refuses_and_never_mints(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg())
|
||||
sess = _fake_session(registry=reg, user_id="", mint_token="unused")
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._model_backend_auth_token(sess, "tf")
|
||||
sess._mcp_mint_client.mint_model_obo_token_sync.assert_not_called()
|
||||
|
||||
def test_failed_mint_falls_back_to_static(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg())
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token=None)
|
||||
# Mint returned None (no credential / rejected) → None so the static
|
||||
# client credential stands.
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") is None
|
||||
|
||||
def test_failed_mint_refuses_when_operator_enables_fail_closed(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg())
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token=None)
|
||||
sess._config_store = SimpleNamespace(get=lambda key: key == "model.auth_fail_closed")
|
||||
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._model_backend_auth_token(sess, "tf")
|
||||
|
||||
def test_failed_mint_without_real_static_key_always_refuses(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg(api_key=""))
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token=None)
|
||||
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._model_backend_auth_token(sess, "tf")
|
||||
|
||||
def test_missing_mint_host_without_real_static_key_always_refuses(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg(api_key=""))
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token=None)
|
||||
sess._mcp_mint_client = None
|
||||
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._model_backend_auth_token(sess, "tf")
|
||||
|
||||
def test_unknown_alias_returns_none(self) -> None:
|
||||
reg = _registry_with(self._obo_cfg())
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token="x")
|
||||
assert ChatSession._model_backend_auth_token(sess, "does-not-exist") is None
|
||||
|
||||
# -- entra_app (app-identity / client-credentials) --------------------------
|
||||
|
||||
def _app_cfg(self, *, api_key: str = "static-fallback") -> ModelConfig:
|
||||
return ModelConfig(
|
||||
alias="tf",
|
||||
base_url="https://gateway.example.com",
|
||||
api_key=api_key,
|
||||
model="vmg/opus",
|
||||
provider="anthropic",
|
||||
auth_mode="entra_app",
|
||||
obo_audience=AUDIENCE,
|
||||
)
|
||||
|
||||
def test_app_alias_mints_without_user(self) -> None:
|
||||
# entra_app is an explicit model-definition choice for a service
|
||||
# principal; it is never inferred from a missing OBO user.
|
||||
reg = _registry_with(self._app_cfg())
|
||||
sess = _fake_session(registry=reg, user_id="", mint_token=None, app_token="app-jwt")
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") == "app-jwt"
|
||||
sess._mcp_mint_client.mint_app_token_sync.assert_called_once_with(audience=AUDIENCE)
|
||||
sess._mcp_mint_client.mint_model_obo_token_sync.assert_not_called()
|
||||
|
||||
def test_app_alias_uses_app_identity_even_with_user(self) -> None:
|
||||
# A user is present, but entra_app deliberately uses the app identity,
|
||||
# not per-user OBO.
|
||||
reg = _registry_with(self._app_cfg())
|
||||
sess = _fake_session(registry=reg, user_id=USER, mint_token="obo-jwt", app_token="app-jwt")
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") == "app-jwt"
|
||||
sess._mcp_mint_client.mint_model_obo_token_sync.assert_not_called()
|
||||
|
||||
def test_app_failed_mint_falls_back_to_static(self) -> None:
|
||||
reg = _registry_with(self._app_cfg())
|
||||
sess = _fake_session(registry=reg, user_id="", mint_token=None, app_token=None)
|
||||
assert ChatSession._model_backend_auth_token(sess, "tf") is None
|
||||
|
||||
def test_app_failed_mint_without_real_static_key_always_refuses(self) -> None:
|
||||
reg = _registry_with(self._app_cfg(api_key=""))
|
||||
sess = _fake_session(registry=reg, user_id="", mint_token=None, app_token=None)
|
||||
|
||||
with pytest.raises(BackendAuthUnavailableError):
|
||||
ChatSession._model_backend_auth_token(sess, "tf")
|
||||
@@ -1001,6 +1001,37 @@ class TestRegistryReload:
|
||||
assert "a" in reg._providers
|
||||
assert reg._providers["a"] is provider_before
|
||||
|
||||
def test_reload_drops_client_when_auth_mode_changes(self) -> None:
|
||||
"""Client construction chooses a placeholder from auth_mode, so a mode
|
||||
change must rebuild even when URL and stored api_key are unchanged."""
|
||||
models = {
|
||||
"a": ModelConfig(
|
||||
"a",
|
||||
"http://x/v1",
|
||||
"",
|
||||
"m",
|
||||
provider="openai",
|
||||
auth_mode="static",
|
||||
)
|
||||
}
|
||||
reg = ModelRegistry(models=models, default="a")
|
||||
reg._clients["a"] = MagicMock()
|
||||
|
||||
new_models = {
|
||||
"a": ModelConfig(
|
||||
"a",
|
||||
"http://x/v1",
|
||||
"",
|
||||
"m",
|
||||
provider="openai",
|
||||
auth_mode="entra_app",
|
||||
obo_audience="api://gateway",
|
||||
)
|
||||
}
|
||||
reg.reload(new_models, "a")
|
||||
|
||||
assert "a" not in reg._clients
|
||||
|
||||
def test_reload_drops_provider_when_provider_string_changes(self) -> None:
|
||||
"""A provider-type swap (e.g. openai → anthropic) drops both the
|
||||
client AND the provider so the next resolve picks up the right
|
||||
|
||||
@@ -76,6 +76,48 @@ def _lane(provider: _FakeProvider, **kw: Any) -> ModelLane:
|
||||
return ModelLane(provider=provider, client=object(), model="m", **kw)
|
||||
|
||||
|
||||
def test_backend_auth_token_binds_sdk_credential_once() -> None:
|
||||
"""Dynamic credentials use SDK with_options, not an override header."""
|
||||
provider = _FakeProvider([CompletionResult(content="ok")])
|
||||
base_client = MagicMock()
|
||||
bound_client = object()
|
||||
base_client.with_options.return_value = bound_client
|
||||
lane = ModelLane(provider=provider, client=base_client, model="m", alias="gateway")
|
||||
|
||||
result = model_turn(
|
||||
lane,
|
||||
[Turn.user("hello")],
|
||||
backend_auth_token="minted-token",
|
||||
)
|
||||
|
||||
assert result.content == "ok"
|
||||
base_client.with_options.assert_called_once_with(api_key="minted-token")
|
||||
assert provider.calls[0]["client"] is bound_client
|
||||
assert "extra_headers" not in provider.calls[0]
|
||||
|
||||
|
||||
def test_entra_app_lane_resolver_never_issues_placeholder_client() -> None:
|
||||
"""A resolver-carrying lane binds its app token before the provider call."""
|
||||
provider = _FakeProvider([CompletionResult(content="ok")])
|
||||
placeholder_client = MagicMock(name="backend-auth-placeholder-unused")
|
||||
bound_client = object()
|
||||
placeholder_client.with_options.return_value = bound_client
|
||||
resolver = MagicMock(return_value="app-token")
|
||||
lane = ModelLane(
|
||||
provider=provider,
|
||||
client=placeholder_client,
|
||||
model="m",
|
||||
alias="app-gateway",
|
||||
backend_auth_resolver=resolver,
|
||||
)
|
||||
|
||||
model_turn(lane, [Turn.user("hello")])
|
||||
|
||||
resolver.assert_called_once_with("app-gateway")
|
||||
placeholder_client.with_options.assert_called_once_with(api_key="app-token")
|
||||
assert provider.calls[0]["client"] is bound_client
|
||||
|
||||
|
||||
class _FlakyProvider:
|
||||
"""Scripted drain-time deaths: each script entry is either a
|
||||
``CompletionResult`` (streamed normally) or an exception instance
|
||||
|
||||
@@ -10,7 +10,7 @@ from __future__ import annotations
|
||||
import urllib.parse
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
@@ -1144,6 +1144,16 @@ class TestAdminOIDCIdentities:
|
||||
)
|
||||
app.state.auth_storage = storage
|
||||
app.state.mcp_token_store = store
|
||||
app.state.mcp_client = MagicMock()
|
||||
app.state.collector = MagicMock()
|
||||
app.state.collector.get_all_nodes.return_value = [
|
||||
{"node_id": "node-1", "server_url": "https://node-1.example"}
|
||||
]
|
||||
app.state.proxy_client = MagicMock()
|
||||
app.state.proxy_client.post = AsyncMock(return_value=SimpleNamespace(status_code=200))
|
||||
app.state.proxy_token_mgr = SimpleNamespace(
|
||||
bearer_header={"Authorization": "Bearer service-token"}
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
storage.create_oidc_identity(issuer, "sub-1", "user-x", "x@example.com")
|
||||
@@ -1166,15 +1176,50 @@ class TestAdminOIDCIdentities:
|
||||
as_issuer=issuer,
|
||||
audience="api://mcp-a",
|
||||
)
|
||||
store.create_user_token(
|
||||
"user-x",
|
||||
"__model_obo__:api://model-a",
|
||||
access_token="model-at",
|
||||
refresh_token=None,
|
||||
expires_at="2026-12-31T00:00:00",
|
||||
scopes=None,
|
||||
as_issuer=issuer,
|
||||
audience="api://model-a",
|
||||
)
|
||||
# Shared app-identity rows are deliberately not tied to this user.
|
||||
store.create_user_token(
|
||||
"__app__",
|
||||
"__model_app__:api://model-a",
|
||||
access_token="app-at",
|
||||
refresh_token=None,
|
||||
expires_at="2026-12-31T00:00:00",
|
||||
scopes=None,
|
||||
as_issuer=issuer,
|
||||
audience="api://model-a",
|
||||
)
|
||||
|
||||
resp = client.delete(f"/v1/api/admin/oidc-identities?issuer={issuer}&subject=sub-1")
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
assert body["obo_credential_revoked"] is True
|
||||
assert body["obo_cache_rows_purged"] == 1
|
||||
assert body["obo_cache_rows_purged"] == 2
|
||||
# Credential gone → no future mints; cache row gone → no live cached bearer.
|
||||
assert store.get_oidc_credential("user-x", issuer) is None
|
||||
assert storage.get_mcp_user_token("user-x", "obo-srv") is None
|
||||
assert storage.get_mcp_user_token("user-x", "__model_obo__:api://model-a") is None
|
||||
assert storage.get_mcp_user_token("__app__", "__model_app__:api://model-a") is not None
|
||||
app.state.mcp_client.invalidate_model_mint_memo_sync.assert_called_once_with(
|
||||
user_id="user-x",
|
||||
server_prefix="__model_obo__:",
|
||||
)
|
||||
app.state.proxy_client.post.assert_awaited_once_with(
|
||||
"https://node-1.example/v1/api/_internal/model-auth-cache-invalidate",
|
||||
headers={"Authorization": "Bearer service-token"},
|
||||
json={"user_id": "user-x"},
|
||||
timeout=5,
|
||||
)
|
||||
assert body["model_memo_nodes_invalidated"] == 1
|
||||
assert body["model_memo_nodes_failed"] == 0
|
||||
|
||||
def test_delete_nonexistent_returns_404(self, admin_client: TestClient) -> None:
|
||||
resp = admin_client.delete(
|
||||
|
||||
@@ -87,12 +87,13 @@ def test_describe_empty_parts_skips_backend() -> None:
|
||||
assert prov.calls == 0
|
||||
|
||||
|
||||
def test_describe_cached_memoizes_by_alias_and_hash() -> None:
|
||||
def test_describe_cached_memoizes_by_principal_alias_and_hash() -> None:
|
||||
prov = _StubProvider(content="desc")
|
||||
kw: dict[str, Any] = {
|
||||
"provider": prov,
|
||||
"client": object(),
|
||||
"model": "m",
|
||||
"principal_id": "user-a",
|
||||
"alias": "omni",
|
||||
"content_hash": "h1",
|
||||
"parts": _parts(),
|
||||
@@ -102,6 +103,8 @@ def test_describe_cached_memoizes_by_alias_and_hash() -> None:
|
||||
assert prov.calls == 1 # second served from cache
|
||||
perception.describe_cached(**{**kw, "content_hash": "h2"})
|
||||
assert prov.calls == 2 # distinct hash → fresh perceive
|
||||
perception.describe_cached(**{**kw, "principal_id": "user-b"})
|
||||
assert prov.calls == 3 # same content under another user's grant → fresh perceive
|
||||
|
||||
|
||||
def test_describe_cached_does_not_cache_failures() -> None:
|
||||
@@ -110,6 +113,7 @@ def test_describe_cached_does_not_cache_failures() -> None:
|
||||
"provider": prov,
|
||||
"client": object(),
|
||||
"model": "m",
|
||||
"principal_id": "user-a",
|
||||
"alias": "omni",
|
||||
"content_hash": "h",
|
||||
"parts": _parts(),
|
||||
@@ -120,7 +124,14 @@ def test_describe_cached_does_not_cache_failures() -> None:
|
||||
|
||||
|
||||
def test_describe_peek_returns_none_when_absent() -> None:
|
||||
assert perception.describe_peek(alias="omni", content_hash="missing") is None
|
||||
assert (
|
||||
perception.describe_peek(
|
||||
principal_id="user-a",
|
||||
alias="omni",
|
||||
content_hash="missing",
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_describe_peek_returns_cached_without_recompute() -> None:
|
||||
@@ -129,6 +140,7 @@ def test_describe_peek_returns_cached_without_recompute() -> None:
|
||||
"provider": prov,
|
||||
"client": object(),
|
||||
"model": "m",
|
||||
"principal_id": "user-a",
|
||||
"alias": "omni",
|
||||
"content_hash": "h",
|
||||
"parts": _parts(),
|
||||
@@ -137,5 +149,20 @@ def test_describe_peek_returns_cached_without_recompute() -> None:
|
||||
assert prov.calls == 1
|
||||
# Peek serves the memoized text and never re-invokes the backend — this is
|
||||
# what lets the wire resolver skip the PDF rasterize on a cross-send hit.
|
||||
assert perception.describe_peek(alias="omni", content_hash="h") == "desc"
|
||||
assert (
|
||||
perception.describe_peek(
|
||||
principal_id="user-a",
|
||||
alias="omni",
|
||||
content_hash="h",
|
||||
)
|
||||
== "desc"
|
||||
)
|
||||
assert (
|
||||
perception.describe_peek(
|
||||
principal_id="user-b",
|
||||
alias="omni",
|
||||
content_hash="h",
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert prov.calls == 1
|
||||
|
||||
@@ -50,8 +50,11 @@ def _vllm_registry(*, replay: bool = True, alias: str = "qwen3") -> Any:
|
||||
replay_reasoning_to_model=replay,
|
||||
capabilities={},
|
||||
server_compat={"server_type": "vllm"},
|
||||
auth_mode="static",
|
||||
obo_audience="",
|
||||
)
|
||||
return SimpleNamespace(
|
||||
has_alias=lambda a: a == alias,
|
||||
get_config=lambda a: cfg if a == alias else (_ for _ in ()).throw(KeyError(a)),
|
||||
)
|
||||
|
||||
@@ -61,8 +64,11 @@ def _registry_with_server_type(server_type: str, *, replay: bool = True) -> Any:
|
||||
replay_reasoning_to_model=replay,
|
||||
capabilities={},
|
||||
server_compat={"server_type": server_type},
|
||||
auth_mode="static",
|
||||
obo_audience="",
|
||||
)
|
||||
return SimpleNamespace(
|
||||
has_alias=lambda _alias: True,
|
||||
get_config=lambda _alias: cfg,
|
||||
)
|
||||
|
||||
|
||||
@@ -32,7 +32,8 @@ from tests._helpers import patch_session_storage
|
||||
from turnstone.core.session import ChatSession
|
||||
from turnstone.core.storage import get_storage
|
||||
from turnstone.core.trajectory import dicts_from_turns
|
||||
from turnstone.core.watch import WatchRunner
|
||||
from turnstone.core.watch import WatchRunner, WatchWorkstreamUnrestorable
|
||||
from turnstone.server import _watch_restore_owner
|
||||
|
||||
|
||||
class _NullUI:
|
||||
@@ -42,7 +43,7 @@ class _NullUI:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
def _make_session() -> ChatSession:
|
||||
def _make_session(*, user_id: str = "") -> ChatSession:
|
||||
"""Real ChatSession with the same minimal setup the unit-test suite
|
||||
uses; no LLM calls happen until a chat-loop method is exercised
|
||||
(and even then the LLM provider is patched).
|
||||
@@ -55,9 +56,28 @@ def _make_session() -> ChatSession:
|
||||
temperature=0.5,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
|
||||
def test_watch_restore_owner_requires_persisted_principal() -> None:
|
||||
storage = MagicMock()
|
||||
storage.get_workstream_owner.return_value = "user-123"
|
||||
assert _watch_restore_owner(storage, "ws-1") == "user-123"
|
||||
|
||||
storage.get_workstream_owner.return_value = ""
|
||||
with pytest.raises(WatchWorkstreamUnrestorable):
|
||||
_watch_restore_owner(storage, "ws-unowned")
|
||||
|
||||
storage.get_workstream_owner.return_value = None
|
||||
with pytest.raises(WatchWorkstreamUnrestorable):
|
||||
_watch_restore_owner(storage, "ws-missing")
|
||||
|
||||
storage.get_workstream_owner.side_effect = RuntimeError("storage unavailable")
|
||||
with pytest.raises(RuntimeError, match="storage unavailable"):
|
||||
_watch_restore_owner(storage, "ws-transient")
|
||||
|
||||
|
||||
def test_watch_fires_then_user_send_drains_envelope(tmp_db, monkeypatch):
|
||||
"""Pin the cross-PR concern that watch text reaches the model via
|
||||
the unified operator-context ``system`` turn path:
|
||||
@@ -217,7 +237,8 @@ def test_watch_dispatch_through_restore_fn_lands_on_rehydrated_session(tmp_db, m
|
||||
|
||||
# Stage 1 — build the original session and persist a message so
|
||||
# ``session.resume`` finds the ws_id in storage.
|
||||
original = _make_session()
|
||||
owner_id = "user-123"
|
||||
original = _make_session(user_id=owner_id)
|
||||
original_ws_id = original._ws_id
|
||||
# Persist a stub user message so ``load_messages(original_ws_id)``
|
||||
# returns something non-empty (resume short-circuits on empty).
|
||||
@@ -235,7 +256,11 @@ def test_watch_dispatch_through_restore_fn_lands_on_rehydrated_session(tmp_db, m
|
||||
new session adopts the original ws_id), wire the dispatch
|
||||
closure, return the dispatch fn.
|
||||
"""
|
||||
new_session = _make_session()
|
||||
owner_storage = MagicMock()
|
||||
owner_storage.get_workstream_owner.return_value = owner_id
|
||||
new_session = _make_session(
|
||||
user_id=_watch_restore_owner(owner_storage, ws_id),
|
||||
)
|
||||
ok = new_session.resume(ws_id)
|
||||
assert ok, "resume should succeed against a non-empty message log"
|
||||
new_session.set_watch_runner(runner)
|
||||
@@ -266,6 +291,7 @@ def test_watch_dispatch_through_restore_fn_lands_on_rehydrated_session(tmp_db, m
|
||||
rehydrated = rehydrated_holder["session"]
|
||||
assert rehydrated is not original
|
||||
assert rehydrated._ws_id == original_ws_id
|
||||
assert rehydrated._mcp_effective_user_id == owner_id
|
||||
|
||||
# The watch payload landed on the rehydrated session's queue, not on
|
||||
# the (now-evicted) original session's queue.
|
||||
|
||||
@@ -1011,6 +1011,10 @@ class ModelDefinitionInfo(BaseModel):
|
||||
reasoning_effort: str | None = None
|
||||
surface_persisted_reasoning: bool = True
|
||||
replay_reasoning_to_model: bool = False
|
||||
# "static" (send api_key), "entra_obo" (delegated-user token), or
|
||||
# "entra_app" (shared app token) for obo_audience at call time.
|
||||
auth_mode: str = "static"
|
||||
obo_audience: str = ""
|
||||
source: str = ""
|
||||
created_by: str = ""
|
||||
created: str = ""
|
||||
@@ -1031,6 +1035,8 @@ class CreateModelDefinitionRequest(BaseModel):
|
||||
reasoning_effort: str | None = None
|
||||
surface_persisted_reasoning: bool = True
|
||||
replay_reasoning_to_model: bool = False
|
||||
auth_mode: str = "static"
|
||||
obo_audience: str = ""
|
||||
|
||||
|
||||
class UpdateModelDefinitionRequest(BaseModel):
|
||||
@@ -1047,6 +1053,8 @@ class UpdateModelDefinitionRequest(BaseModel):
|
||||
reasoning_effort: str | None = None
|
||||
surface_persisted_reasoning: bool | None = None
|
||||
replay_reasoning_to_model: bool | None = None
|
||||
auth_mode: str | None = None
|
||||
obo_audience: str | None = None
|
||||
|
||||
|
||||
class ListModelDefinitionsResponse(BaseModel):
|
||||
|
||||
@@ -7,6 +7,7 @@ model auto-detection, workstream management, and the main() REPL entry point.
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
@@ -1242,7 +1243,41 @@ def main() -> None:
|
||||
mcp_client = create_mcp_client(
|
||||
getattr(args, "mcp_config", None),
|
||||
storage=_get_storage(),
|
||||
required=registry.has_dynamic_auth(),
|
||||
)
|
||||
if mcp_client is not None:
|
||||
cli_auth_storage = _get_storage()
|
||||
mcp_client.set_storage(cli_auth_storage)
|
||||
if registry.has_dynamic_auth():
|
||||
# The CLI has no ASGI lifespan, but app-identity model auth needs
|
||||
# the same discovered OIDC config and encrypted token store as the
|
||||
# server hosts. The mint HTTP client itself belongs to mcp-loop.
|
||||
from types import SimpleNamespace
|
||||
|
||||
from turnstone.core.mcp_crypto import initialize_mcp_crypto_state
|
||||
from turnstone.core.oidc import (
|
||||
close_oidc_state,
|
||||
initialize_oidc_state,
|
||||
load_oidc_config,
|
||||
)
|
||||
|
||||
cli_auth_state = SimpleNamespace(
|
||||
auth_storage=cli_auth_storage,
|
||||
oidc_config=load_oidc_config(),
|
||||
registry=registry,
|
||||
)
|
||||
|
||||
async def _initialize_cli_auth_state() -> None:
|
||||
await initialize_oidc_state(cli_auth_state)
|
||||
# Login/JWKS callbacks do not exist in the CLI. Close their
|
||||
# client while retaining the discovered token endpoint.
|
||||
await close_oidc_state(cli_auth_state)
|
||||
|
||||
asyncio.run(_initialize_cli_auth_state())
|
||||
initialize_mcp_crypto_state(cli_auth_state, node_id="cli")
|
||||
cli_auth_state.mcp_oauth_refresh_locks = {}
|
||||
cli_auth_state.mcp_oauth_refresh_backoff = {}
|
||||
mcp_client.set_app_state(cli_auth_state)
|
||||
|
||||
# apply_config() merges [judge] config.toml values into args as
|
||||
# Output_guard and redact_secrets default to True, enabling the heuristic
|
||||
|
||||
+238
-8
@@ -60,6 +60,7 @@ from turnstone.core.deadline import DeadlineExceededError, run_with_deadline
|
||||
from turnstone.core.mcp_crypto import is_user_scoped_auth
|
||||
from turnstone.core.memory import get_workstream_display_names
|
||||
from turnstone.core.metacognition import field_str, sanitize_display
|
||||
from turnstone.core.model_registry import MODEL_AUTH_MODES as _MODEL_AUTH_MODES
|
||||
from turnstone.core.rendezvous import NoAvailableNodeError
|
||||
from turnstone.core.session_replay import session_replay_preamble
|
||||
from turnstone.core.session_routes import (
|
||||
@@ -5297,7 +5298,13 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
create_mcp_client,
|
||||
config_store.get("mcp.config_path") or None,
|
||||
storage=storage,
|
||||
required=bool(
|
||||
app.state.coord_registry and app.state.coord_registry.has_dynamic_auth()
|
||||
),
|
||||
)
|
||||
if app.state.mcp_client is not None:
|
||||
app.state.mcp_client.set_storage(storage)
|
||||
app.state.mcp_client.set_app_state(app.state)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"console MCP manager boot failed — coordinators get no MCP "
|
||||
@@ -5796,10 +5803,57 @@ async def admin_list_oidc_identities(request: Request) -> JSONResponse:
|
||||
return JSONResponse({"oidc_identities": identities})
|
||||
|
||||
|
||||
async def _invalidate_model_auth_memos_cluster(request: Request, user_id: str) -> tuple[int, int]:
|
||||
"""Best-effort eviction of one user's delegated model memo on every node."""
|
||||
collector = getattr(request.app.state, "collector", None)
|
||||
client = getattr(request.app.state, "proxy_client", None)
|
||||
token_mgr = getattr(request.app.state, "proxy_token_mgr", None)
|
||||
if collector is None or client is None or token_mgr is None:
|
||||
return 0, 0
|
||||
|
||||
headers = dict(token_mgr.bearer_header)
|
||||
nodes = collector.get_all_nodes()
|
||||
sem = asyncio.Semaphore(_get_fan_out_limit(request))
|
||||
|
||||
async def _invalidate(node: dict[str, Any]) -> bool | None:
|
||||
node_id = str(node.get("node_id") or "")
|
||||
url = str(node.get("server_url") or "").rstrip("/")
|
||||
if not url:
|
||||
return None
|
||||
try:
|
||||
async with sem:
|
||||
response = await client.post(
|
||||
f"{url}/v1/api/_internal/model-auth-cache-invalidate",
|
||||
headers=headers,
|
||||
json={"user_id": user_id},
|
||||
timeout=5,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
log.warning(
|
||||
"admin.oidc_identity.model_memo_invalidate_rejected user=%s node=%s status=%s",
|
||||
user_id,
|
||||
node_id,
|
||||
response.status_code,
|
||||
)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"admin.oidc_identity.model_memo_invalidate_failed user=%s node=%s",
|
||||
user_id,
|
||||
node_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
outcomes = await asyncio.gather(*(_invalidate(node) for node in nodes))
|
||||
return outcomes.count(True), outcomes.count(False)
|
||||
|
||||
|
||||
async def admin_delete_oidc_identity(request: Request) -> JSONResponse:
|
||||
"""DELETE /v1/api/admin/oidc-identities?issuer=...&subject=... — unlink OIDC identity."""
|
||||
from turnstone.core.audit import record_audit
|
||||
from turnstone.core.auth import require_permission
|
||||
from turnstone.core.mcp_oauth import MODEL_OBO_CACHE_PREFIX
|
||||
from turnstone.core.web_helpers import require_storage_or_503
|
||||
|
||||
storage, err = require_storage_or_503(request)
|
||||
@@ -5861,6 +5915,44 @@ async def admin_delete_oidc_identity(request: Request) -> JSONResponse:
|
||||
server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
# Model-provider OBO rows use synthetic server keys and therefore are
|
||||
# not present in mcp_servers/obo_server_names. Purge only delegated
|
||||
# model rows for this user; __model_app__ rows belong to the shared
|
||||
# __app__ identity and intentionally survive user deprovisioning.
|
||||
try:
|
||||
metadata = token_store.list_user_token_metadata(user_id)
|
||||
except Exception:
|
||||
metadata = []
|
||||
log.warning(
|
||||
"admin.oidc_identity.model_obo_cache_list_failed user=%s",
|
||||
user_id,
|
||||
exc_info=True,
|
||||
)
|
||||
for row in metadata:
|
||||
server_name = str(row.get("server_name") or "")
|
||||
if not server_name.startswith(MODEL_OBO_CACHE_PREFIX):
|
||||
continue
|
||||
try:
|
||||
if token_store.delete_user_token(user_id, server_name):
|
||||
obo_cache_purged += 1
|
||||
except Exception:
|
||||
log.warning(
|
||||
"admin.oidc_identity.model_obo_cache_purge_failed user=%s server=%s",
|
||||
user_id,
|
||||
server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
# The console-hosted coordinator has its own manager/memo and is not in the
|
||||
# node collector, so invalidate it directly as well as fanning out below.
|
||||
mcp_client = getattr(request.app.state, "mcp_client", None)
|
||||
if mcp_client is not None and hasattr(mcp_client, "invalidate_model_mint_memo_sync"):
|
||||
mcp_client.invalidate_model_mint_memo_sync(
|
||||
user_id=user_id,
|
||||
server_prefix=MODEL_OBO_CACHE_PREFIX,
|
||||
)
|
||||
memo_nodes_invalidated, memo_nodes_failed = await _invalidate_model_auth_memos_cluster(
|
||||
request, user_id
|
||||
)
|
||||
|
||||
audit_uid, ip = _audit_context(request)
|
||||
record_audit(
|
||||
@@ -5873,20 +5965,19 @@ async def admin_delete_oidc_identity(request: Request) -> JSONResponse:
|
||||
"user_id": user_id,
|
||||
"obo_credential_revoked": credential_revoked,
|
||||
"obo_cache_rows_purged": obo_cache_purged,
|
||||
"model_memo_nodes_invalidated": memo_nodes_invalidated,
|
||||
"model_memo_nodes_failed": memo_nodes_failed,
|
||||
},
|
||||
ip,
|
||||
)
|
||||
|
||||
# Note: warmed per-user pool sessions on server nodes hold the bearer
|
||||
# in-memory and self-clear at token TTL / idle eviction; the console has no
|
||||
# per-user cross-node eviction primitive. Credential + cache purge stop all
|
||||
# FUTURE dispatch (next call re-reads the now-empty cache → mint → missing
|
||||
# credential → re-login), bounding residual access to the current token TTL.
|
||||
return JSONResponse(
|
||||
{
|
||||
"status": "ok",
|
||||
"obo_credential_revoked": credential_revoked,
|
||||
"obo_cache_rows_purged": obo_cache_purged,
|
||||
"model_memo_nodes_invalidated": memo_nodes_invalidated,
|
||||
"model_memo_nodes_failed": memo_nodes_failed,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -9572,7 +9663,9 @@ def _clean_oauth_text(value: Any, *, max_length: int = 512) -> str | None:
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
# OAuth identifiers/URLs have no valid C0 controls. Removing them here
|
||||
# prevents log/header ambiguity and benefits both MCP and model auth.
|
||||
text = re.sub(r"[\x00-\x1f\x7f]", "", str(value)).strip()
|
||||
if not text:
|
||||
return None
|
||||
return text[:max_length]
|
||||
@@ -10770,11 +10863,28 @@ def _ensure_console_mcp_client(app: Any) -> dict[str, Any]:
|
||||
|
||||
cs = getattr(app.state, "config_store", None)
|
||||
cfg_path = cs.get("mcp.config_path") if cs is not None else None
|
||||
mgr = create_mcp_client(cfg_path or None, storage=storage)
|
||||
registry = getattr(app.state, "coord_registry", None)
|
||||
mgr = create_mcp_client(
|
||||
cfg_path or None,
|
||||
storage=storage,
|
||||
required=bool(registry and registry.has_dynamic_auth()),
|
||||
)
|
||||
if mgr is None:
|
||||
return {"skipped": "no MCP servers configured"}
|
||||
app.state.mcp_client = mgr
|
||||
return mgr.reconcile_sync(storage)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(app.state)
|
||||
result = mgr.reconcile_sync(storage)
|
||||
coord_mgr = getattr(app.state, "coord_mgr", None)
|
||||
if coord_mgr is not None:
|
||||
try:
|
||||
for ws in coord_mgr.list_all():
|
||||
session = getattr(ws, "session", None)
|
||||
if session is not None:
|
||||
session.set_model_mint_client(mgr)
|
||||
except Exception:
|
||||
log.debug("console.model_mint_client_session_refresh_failed", exc_info=True)
|
||||
return result
|
||||
|
||||
|
||||
async def _notify_nodes_mcp_reload(request: Request) -> dict[str, Any]:
|
||||
@@ -11330,6 +11440,47 @@ _REASONING_EFFORT_CHOICES = frozenset(
|
||||
_API_SURFACE_CHOICES = frozenset({"chat", "responses"})
|
||||
|
||||
|
||||
def _model_auth_audience_allowlist(request: Request) -> frozenset[str]:
|
||||
"""Exact operator-approved resource audiences for dynamic model auth."""
|
||||
config_store = getattr(request.app.state, "config_store", None)
|
||||
raw = config_store.get("model.auth_audience_allowlist") if config_store is not None else ""
|
||||
return frozenset(item.strip() for item in re.split(r"[,\n]", str(raw or "")) if item.strip())
|
||||
|
||||
|
||||
def _validate_dynamic_model_auth(
|
||||
request: Request,
|
||||
*,
|
||||
auth_mode: str,
|
||||
audience: str,
|
||||
) -> JSONResponse | None:
|
||||
"""Validate allow-list/profile constraints for a changed dynamic config."""
|
||||
if auth_mode == "static":
|
||||
return None
|
||||
if audience not in _model_auth_audience_allowlist(request):
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
"obo_audience is not in the operator-configured model.auth_audience_allowlist"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
profile = str(
|
||||
getattr(getattr(request.app.state, "oidc_config", None), "obo_grant_profile", "") or ""
|
||||
)
|
||||
if auth_mode == "entra_app" and profile != "entra":
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
"auth_mode 'entra_app' requires [oidc] obo_grant_profile='entra'; "
|
||||
"RFC 8693 client-credentials is not supported"
|
||||
)
|
||||
},
|
||||
status_code=400,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _validate_api_surface(caps: Any) -> str | None:
|
||||
"""Return an error message if ``caps["server_compat"]["api_surface"]`` is invalid.
|
||||
|
||||
@@ -11618,6 +11769,8 @@ async def admin_list_model_definitions(request: Request) -> JSONResponse:
|
||||
cfg_temperature = None
|
||||
cfg_max_tokens = None
|
||||
cfg_reasoning_effort = None
|
||||
cfg_auth_mode = "static"
|
||||
cfg_obo_audience = ""
|
||||
for node_models in node_statuses.values():
|
||||
nm = node_models.get(alias)
|
||||
if nm:
|
||||
@@ -11627,6 +11780,8 @@ async def admin_list_model_definitions(request: Request) -> JSONResponse:
|
||||
cfg_temperature = nm.get("temperature")
|
||||
cfg_max_tokens = nm.get("max_tokens")
|
||||
cfg_reasoning_effort = nm.get("reasoning_effort")
|
||||
cfg_auth_mode = nm.get("auth_mode", "static")
|
||||
cfg_obo_audience = nm.get("obo_audience", "")
|
||||
break
|
||||
result.append(
|
||||
{
|
||||
@@ -11642,6 +11797,8 @@ async def admin_list_model_definitions(request: Request) -> JSONResponse:
|
||||
"temperature": cfg_temperature,
|
||||
"max_tokens": cfg_max_tokens,
|
||||
"reasoning_effort": cfg_reasoning_effort,
|
||||
"auth_mode": cfg_auth_mode,
|
||||
"obo_audience": cfg_obo_audience,
|
||||
"source": "config",
|
||||
"created_by": "",
|
||||
"created": "",
|
||||
@@ -11767,6 +11924,30 @@ async def admin_create_model_definition(request: Request) -> JSONResponse:
|
||||
surface_persisted_reasoning = bool(body.get("surface_persisted_reasoning", True))
|
||||
replay_reasoning_to_model = bool(body.get("replay_reasoning_to_model", False))
|
||||
|
||||
auth_mode = str(body.get("auth_mode", "static")).strip() or "static"
|
||||
if auth_mode not in _MODEL_AUTH_MODES:
|
||||
return JSONResponse({"error": f"Invalid auth_mode: {auth_mode!r}"}, status_code=400)
|
||||
obo_audience = _clean_oauth_text(body.get("obo_audience"), max_length=2048) or ""
|
||||
if auth_mode in ("entra_obo", "entra_app") and not obo_audience:
|
||||
return JSONResponse(
|
||||
{"error": "obo_audience is required when auth_mode is 'entra_obo' or 'entra_app'"},
|
||||
status_code=400,
|
||||
)
|
||||
dynamic_auth_error = _validate_dynamic_model_auth(
|
||||
request,
|
||||
auth_mode=auth_mode,
|
||||
audience=obo_audience,
|
||||
)
|
||||
if dynamic_auth_error is not None:
|
||||
return dynamic_auth_error
|
||||
if auth_mode != "static":
|
||||
# Redeeming an operator-chosen audience is the same capability as
|
||||
# configuring oauth_audience on MCP. Service credentials do not bypass
|
||||
# this capability-escalation gate.
|
||||
err = require_permission(request, "admin.mcp", allow_service_bypass=False)
|
||||
if err:
|
||||
return err
|
||||
|
||||
storage.create_model_definition(
|
||||
definition_id=definition_id,
|
||||
alias=alias,
|
||||
@@ -11783,6 +11964,8 @@ async def admin_create_model_definition(request: Request) -> JSONResponse:
|
||||
reasoning_effort=reasoning_effort,
|
||||
surface_persisted_reasoning=surface_persisted_reasoning,
|
||||
replay_reasoning_to_model=replay_reasoning_to_model,
|
||||
auth_mode=auth_mode,
|
||||
obo_audience=obo_audience,
|
||||
)
|
||||
|
||||
record_audit(
|
||||
@@ -11800,6 +11983,7 @@ async def admin_create_model_definition(request: Request) -> JSONResponse:
|
||||
# model promotes the coord subsystem from "not initialized" to ready
|
||||
# without a console restart. No-op when already built.
|
||||
await asyncio.to_thread(_maybe_bootstrap_coord_subsystem, request.app, storage)
|
||||
await asyncio.to_thread(_ensure_console_mcp_client, request.app)
|
||||
_emit_models_changed(request)
|
||||
|
||||
created = storage.get_model_definition(definition_id)
|
||||
@@ -11949,6 +12133,51 @@ async def admin_update_model_definition(request: Request) -> JSONResponse:
|
||||
if "replay_reasoning_to_model" in body:
|
||||
updates["replay_reasoning_to_model"] = bool(body["replay_reasoning_to_model"])
|
||||
|
||||
if "auth_mode" in body:
|
||||
am = str(body["auth_mode"]).strip() or "static"
|
||||
if am not in _MODEL_AUTH_MODES:
|
||||
return JSONResponse({"error": f"Invalid auth_mode: {am!r}"}, status_code=400)
|
||||
updates["auth_mode"] = am
|
||||
if "obo_audience" in body:
|
||||
updates["obo_audience"] = _clean_oauth_text(body["obo_audience"], max_length=2048) or ""
|
||||
# Cross-field: entra_obo needs an audience. Validate the POST-merge state so
|
||||
# a request that touches only one of the pair still checks against the other.
|
||||
eff_auth_mode = updates.get("auth_mode", existing.get("auth_mode", "static"))
|
||||
eff_audience = updates.get("obo_audience", existing.get("obo_audience", ""))
|
||||
if eff_auth_mode in ("entra_obo", "entra_app") and not eff_audience:
|
||||
return JSONResponse(
|
||||
{"error": "obo_audience is required when auth_mode is 'entra_obo' or 'entra_app'"},
|
||||
status_code=400,
|
||||
)
|
||||
old_auth_mode = str(existing.get("auth_mode") or "static")
|
||||
old_audience = str(existing.get("obo_audience") or "")
|
||||
old_base_url = str(existing.get("base_url") or "")
|
||||
eff_base_url = str(updates.get("base_url", old_base_url))
|
||||
auth_config_changed = (
|
||||
str(eff_auth_mode) != old_auth_mode
|
||||
or str(eff_audience) != old_audience
|
||||
# A base-URL change on either side of a dynamic configuration can
|
||||
# redirect a valid bearer even when the auth fields are unchanged.
|
||||
or (
|
||||
eff_base_url != old_base_url
|
||||
and (
|
||||
old_auth_mode in ("entra_obo", "entra_app")
|
||||
or eff_auth_mode in ("entra_obo", "entra_app")
|
||||
)
|
||||
)
|
||||
)
|
||||
if auth_config_changed:
|
||||
dynamic_auth_error = _validate_dynamic_model_auth(
|
||||
request,
|
||||
auth_mode=str(eff_auth_mode),
|
||||
audience=str(eff_audience),
|
||||
)
|
||||
if dynamic_auth_error is not None:
|
||||
return dynamic_auth_error
|
||||
err = require_permission(request, "admin.mcp", allow_service_bypass=False)
|
||||
if err:
|
||||
return err
|
||||
|
||||
if updates:
|
||||
storage.update_model_definition(definition_id, **updates)
|
||||
|
||||
@@ -11969,6 +12198,7 @@ async def admin_update_model_definition(request: Request) -> JSONResponse:
|
||||
if updates:
|
||||
await asyncio.to_thread(_refresh_coord_registry, request.app.state, storage)
|
||||
await asyncio.to_thread(_maybe_bootstrap_coord_subsystem, request.app, storage)
|
||||
await asyncio.to_thread(_ensure_console_mcp_client, request.app)
|
||||
_emit_models_changed(request)
|
||||
|
||||
model_def = storage.get_model_definition(definition_id)
|
||||
|
||||
@@ -7159,6 +7159,10 @@ function _renderModels(items) {
|
||||
// operator opt-in). Default values are silent.
|
||||
if (m.surface_persisted_reasoning === false) overrides.push("surface=off");
|
||||
if (m.replay_reasoning_to_model === true) overrides.push("replay=on");
|
||||
// Per-user OBO / app-identity backend auth surfaces as a hint (static is
|
||||
// the default).
|
||||
if (m.auth_mode === "entra_obo") overrides.push("obo");
|
||||
else if (m.auth_mode === "entra_app") overrides.push("app");
|
||||
if (overrides.length) {
|
||||
const ovrSpan = document.createElement("span");
|
||||
ovrSpan.className = "model-overrides-hint";
|
||||
@@ -7361,6 +7365,8 @@ function showCreateModelModal() {
|
||||
document.getElementById("model-enabled").checked = true;
|
||||
document.getElementById("model-surface-persisted-reasoning").checked = true;
|
||||
document.getElementById("model-replay-reasoning").checked = false;
|
||||
document.getElementById("model-auth-mode").value = "static";
|
||||
document.getElementById("model-obo-audience").value = "";
|
||||
document.getElementById("model-detect-result").hidden = true;
|
||||
document.getElementById("model-detect-btn").disabled = false;
|
||||
document.getElementById("model-detect-btn").textContent = "Detect";
|
||||
@@ -7414,6 +7420,10 @@ function showEditModelModal(definitionId) {
|
||||
m.max_tokens != null ? m.max_tokens : "";
|
||||
document.getElementById("model-reasoning-effort").value =
|
||||
m.reasoning_effort != null ? m.reasoning_effort : "";
|
||||
document.getElementById("model-auth-mode").value =
|
||||
m.auth_mode || "static";
|
||||
document.getElementById("model-obo-audience").value =
|
||||
m.obo_audience || "";
|
||||
// Parse capabilities JSON and extract server_compat for structured fields
|
||||
let capsObj = {};
|
||||
try {
|
||||
@@ -7710,6 +7720,26 @@ function submitCreateModel() {
|
||||
"model-replay-reasoning",
|
||||
).checked;
|
||||
|
||||
// Backend auth: entra_obo mints a per-user OBO token for obo_audience at
|
||||
// call time; entra_app mints an app-identity (client-credentials) token from
|
||||
// Turnstone's SSO app reg. (The server re-validates the same pairing.)
|
||||
const authMode =
|
||||
document.getElementById("model-auth-mode").value || "static";
|
||||
const oboAudience = document
|
||||
.getElementById("model-obo-audience")
|
||||
.value.trim();
|
||||
if (
|
||||
(authMode === "entra_obo" || authMode === "entra_app") &&
|
||||
oboAudience === ""
|
||||
) {
|
||||
_showModelError(
|
||||
"OBO audience is required when auth mode is 'entra_obo' or 'entra_app'",
|
||||
);
|
||||
return;
|
||||
}
|
||||
form.auth_mode = authMode;
|
||||
form.obo_audience = oboAudience;
|
||||
|
||||
const apiKey = document.getElementById("model-api-key").value;
|
||||
if (apiKey) form.api_key = apiKey;
|
||||
|
||||
|
||||
@@ -1695,6 +1695,35 @@
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div class="sh-section">Backend auth</div>
|
||||
<div class="field-pair">
|
||||
<div>
|
||||
<label for="model-auth-mode">Auth mode</label>
|
||||
<select id="model-auth-mode">
|
||||
<option value="static">static (API key)</option>
|
||||
<option value="entra_obo">
|
||||
entra_obo (per-user OBO)
|
||||
</option>
|
||||
<option value="entra_app">
|
||||
entra_app (app identity)
|
||||
</option>
|
||||
</select>
|
||||
</div>
|
||||
<div>
|
||||
<label for="model-obo-audience"
|
||||
>OBO audience
|
||||
<span class="label-hint"
|
||||
>resource App ID URI; required for entra_obo /
|
||||
entra_app</span
|
||||
></label
|
||||
>
|
||||
<input
|
||||
type="text"
|
||||
id="model-obo-audience"
|
||||
placeholder="https://your-resource.example.com"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<label class="toggle-switch">
|
||||
<input type="checkbox" id="model-enabled" checked />
|
||||
<span class="toggle-track" aria-hidden="true"></span>
|
||||
|
||||
@@ -700,6 +700,7 @@ APPROVE_PATHS: frozenset[str] = frozenset(
|
||||
{
|
||||
"/api/_internal/config-reload",
|
||||
"/api/_internal/mcp-reload",
|
||||
"/api/_internal/model-auth-cache-invalidate",
|
||||
"/api/_internal/model-reload",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -979,6 +979,7 @@ class IntentJudge:
|
||||
model_registry: Any | None = None,
|
||||
session_model_alias: str = "",
|
||||
config_store: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None,
|
||||
) -> None:
|
||||
self._config = config
|
||||
self._rule_registry = rule_registry
|
||||
@@ -987,6 +988,7 @@ class IntentJudge:
|
||||
# global ``model.temperature``) resolve like every other lane.
|
||||
self._model_registry = model_registry
|
||||
self._config_store = config_store
|
||||
self._backend_auth_resolver = backend_auth_resolver
|
||||
# The caller (ChatSession) resolves the session model's real caps from
|
||||
# _get_capabilities (config/registry-aware) and passes them in; they are
|
||||
# this judge's wire capabilities and window when it inherits the session
|
||||
@@ -1353,6 +1355,7 @@ class IntentJudge:
|
||||
registry=self._model_registry,
|
||||
capabilities=self._capabilities,
|
||||
config_store=self._config_store,
|
||||
backend_auth_resolver=self._backend_auth_resolver,
|
||||
)
|
||||
|
||||
# Multi-turn judge loop
|
||||
|
||||
@@ -57,7 +57,10 @@ from turnstone.core.mcp_oauth import (
|
||||
emit_oauth_failure_audit,
|
||||
get_obo_access_token_classified,
|
||||
get_user_access_token_classified,
|
||||
invalidate_model_mint_memo,
|
||||
is_user_scoped_auth,
|
||||
mint_app_access_token,
|
||||
mint_obo_access_token,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -881,6 +884,10 @@ class MCPClientManager:
|
||||
# asserts non-None when it actually runs, so static-only callers
|
||||
# never hit it.
|
||||
self._app_state: Any = None
|
||||
# Long-lived HTTP client created and closed on the mcp-loop. Model-token
|
||||
# mints run on that loop and must not borrow the lifespan-loop OAuth
|
||||
# client or pay a DNS/TLS setup on every turn.
|
||||
self._model_auth_http_client: httpx.AsyncClient | None = None
|
||||
|
||||
# In-memory cache of server names whose ``auth_type='oauth_user'``.
|
||||
# ``_db_servers_to_config`` strips oauth_user rows on the way into
|
||||
@@ -1008,6 +1015,104 @@ class MCPClientManager:
|
||||
deployments may leave it unset.
|
||||
"""
|
||||
self._app_state = app_state
|
||||
if self._model_auth_http_client is not None:
|
||||
app_state.obo_http_client = self._model_auth_http_client
|
||||
|
||||
def mint_model_obo_token_sync(
|
||||
self, *, user_id: str, audience: str, timeout: float = 20.0
|
||||
) -> str | None:
|
||||
"""Resolve a per-user model-provider OBO access token synchronously.
|
||||
|
||||
Bridges the sync agent/model loop to :func:`mint_obo_access_token` on
|
||||
the mcp-loop (same thread that owns the token store's asyncio locks),
|
||||
mirroring ``call_tool_sync``. Returns ``None`` on any failure — no
|
||||
credential, mint rejected, OAuth unwired, timeout — so the model call
|
||||
falls back to the backend's static credential.
|
||||
"""
|
||||
if not user_id or not audience:
|
||||
return None
|
||||
loop = self._loop
|
||||
if loop is None or self._app_state is None:
|
||||
return None
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
mint_obo_access_token(app_state=self._app_state, user_id=user_id, audience=audience),
|
||||
loop,
|
||||
)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
log.warning(
|
||||
"model obo token mint timed out user=%s audience=%s",
|
||||
user_id,
|
||||
audience,
|
||||
)
|
||||
return None
|
||||
except Exception:
|
||||
log.debug(
|
||||
"model obo token mint failed user=%s audience=%s",
|
||||
user_id,
|
||||
audience,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
def mint_app_token_sync(self, *, audience: str, timeout: float = 20.0) -> str | None:
|
||||
"""Resolve an app-identity (client-credentials) model token synchronously.
|
||||
|
||||
The ``auth_mode='entra_app'`` sibling of ``mint_model_obo_token_sync``:
|
||||
bridges to :func:`mint_app_access_token` on the mcp-loop. No user needed
|
||||
— Turnstone's own SSO app registration is the identity. Returns ``None``
|
||||
on any failure so the model call falls back to the static credential.
|
||||
"""
|
||||
if not audience:
|
||||
return None
|
||||
loop = self._loop
|
||||
if loop is None or self._app_state is None:
|
||||
return None
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
mint_app_access_token(app_state=self._app_state, audience=audience),
|
||||
loop,
|
||||
)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
log.warning("model app token mint timed out audience=%s", audience)
|
||||
return None
|
||||
except Exception:
|
||||
log.debug("model app token mint failed audience=%s", audience, exc_info=True)
|
||||
return None
|
||||
|
||||
def invalidate_model_mint_memo_sync(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
server_prefix: str,
|
||||
timeout: float = 5.0,
|
||||
) -> int:
|
||||
"""Invalidate model-token memo entries on their owning MCP loop."""
|
||||
loop = self._loop
|
||||
if loop is None or self._app_state is None:
|
||||
return 0
|
||||
|
||||
async def _invalidate() -> int:
|
||||
return invalidate_model_mint_memo(
|
||||
self._app_state,
|
||||
user_id=user_id,
|
||||
server_prefix=server_prefix,
|
||||
)
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(_invalidate(), loop)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
log.warning("model token memo invalidation timed out user=%s", user_id)
|
||||
return 0
|
||||
except Exception:
|
||||
log.warning("model token memo invalidation failed user=%s", user_id, exc_info=True)
|
||||
return 0
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
@@ -1028,6 +1133,9 @@ class MCPClientManager:
|
||||
|
||||
async def _connect_all(self) -> None:
|
||||
"""Connect to every configured server (runs on the background loop)."""
|
||||
self._model_auth_http_client = httpx.AsyncClient(timeout=10.0)
|
||||
if self._app_state is not None:
|
||||
self._app_state.obo_http_client = self._model_auth_http_client
|
||||
self._exit_stack = AsyncExitStack()
|
||||
await self._exit_stack.__aenter__()
|
||||
|
||||
@@ -5421,6 +5529,20 @@ class MCPClientManager:
|
||||
except Exception:
|
||||
log.debug("Error closing MCP exit stack", exc_info=True)
|
||||
|
||||
if self._loop and self._model_auth_http_client is not None:
|
||||
mint_client = self._model_auth_http_client
|
||||
future = asyncio.run_coroutine_threadsafe(mint_client.aclose(), self._loop)
|
||||
try:
|
||||
future.result(timeout=10)
|
||||
except Exception:
|
||||
log.debug("Error closing model-auth HTTP client", exc_info=True)
|
||||
if (
|
||||
self._app_state is not None
|
||||
and getattr(self._app_state, "obo_http_client", None) is mint_client
|
||||
):
|
||||
self._app_state.obo_http_client = None
|
||||
self._model_auth_http_client = None
|
||||
|
||||
if self._loop:
|
||||
self._loop.call_soon_threadsafe(self._loop.stop)
|
||||
if self._thread:
|
||||
@@ -9213,11 +9335,14 @@ def create_mcp_client(
|
||||
config_path: str | None = None,
|
||||
*,
|
||||
storage: Any = None,
|
||||
required: bool = False,
|
||||
) -> MCPClientManager | None:
|
||||
"""Create and start an MCP client manager.
|
||||
|
||||
Returns *None* if nothing is configured — no static servers from any
|
||||
source AND no pool-backed (``oauth_user``/``oauth_obo``) DB rows.
|
||||
source AND no pool-backed (``oauth_user``/``oauth_obo``) DB rows — unless
|
||||
``required`` is true. Dynamic model authentication uses an empty-config
|
||||
manager solely for its mint loop and therefore sets ``required``.
|
||||
Pool-backed rows alone construct an empty-config manager: their
|
||||
connections form lazily per user, so they contribute nothing to the
|
||||
static *servers* dict, but the host still needs a running manager
|
||||
@@ -9243,7 +9368,7 @@ def create_mcp_client(
|
||||
log.warning("Failed to load DB-managed MCP servers", exc_info=True)
|
||||
|
||||
servers = load_mcp_config(config_path, storage=storage)
|
||||
if not servers and not oauth_user_names and not obo_names:
|
||||
if not required and not servers and not oauth_user_names and not obo_names:
|
||||
return None
|
||||
|
||||
mgr = MCPClientManager(servers)
|
||||
|
||||
@@ -626,11 +626,20 @@ def initialize_mcp_crypto_state(app_state: object, *, node_id: str = "") -> None
|
||||
user_scoped_count = sum(
|
||||
1 for row in storage.list_mcp_servers() if is_user_scoped_auth(row.get("auth_type"))
|
||||
)
|
||||
model_registry = getattr(app_state, "registry", None) or getattr(
|
||||
app_state, "coord_registry", None
|
||||
)
|
||||
dynamic_model_auth = bool(
|
||||
model_registry is not None
|
||||
and hasattr(model_registry, "has_dynamic_auth")
|
||||
and model_registry.has_dynamic_auth()
|
||||
)
|
||||
|
||||
if user_scoped_count > 0 and cipher_cfg is None:
|
||||
if (user_scoped_count > 0 or dynamic_model_auth) and cipher_cfg is None:
|
||||
log.error(
|
||||
"mcp.oauth: %d server(s) configured with auth_type='oauth_user'/'oauth_obo' but %s",
|
||||
"mcp.oauth: %d user-scoped MCP server(s), dynamic_model_auth=%s, but %s",
|
||||
user_scoped_count,
|
||||
dynamic_model_auth,
|
||||
_STARTUP_KEY_REQUIRED_HINT,
|
||||
)
|
||||
raise SystemExit(1)
|
||||
|
||||
+539
-1
@@ -2275,6 +2275,8 @@ async def _read_obo_credential(
|
||||
user_id: str,
|
||||
server_name: str,
|
||||
issuer: str,
|
||||
*,
|
||||
prune_on_missing: bool = True,
|
||||
) -> TokenLookupResult | OIDCCredentialPlain:
|
||||
"""Read the captured IdP credential, classifying absence and undecryptability.
|
||||
|
||||
@@ -2284,10 +2286,24 @@ async def _read_obo_credential(
|
||||
``get_oidc_credential`` raises ``MCPTokenDecryptError``; catching it here
|
||||
keeps the mint path's classified-result contract intact (a raw exception
|
||||
would escape ``_dispatch_pool`` into the session's generic error path).
|
||||
Model-backend mints pass ``prune_on_missing=False`` because they arm a
|
||||
cooldown on this outcome and are still holding their synthetic cache lock;
|
||||
the classified MCP path retains the default cleanup semantics.
|
||||
"""
|
||||
try:
|
||||
credential = await asyncio.to_thread(token_store.get_oidc_credential, user_id, issuer)
|
||||
except MCPTokenDecryptError as exc:
|
||||
if not prune_on_missing:
|
||||
log.warning(
|
||||
"mcp_server.oauth.obo_credential_decrypt_failed",
|
||||
user_id=user_id,
|
||||
server_name=server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
return TokenLookupResult(
|
||||
kind="decrypt_failure",
|
||||
decrypt_fingerprints=tuple(exc.key_fingerprints_attempted),
|
||||
)
|
||||
return _decrypt_failure_result(
|
||||
app_state,
|
||||
user_id,
|
||||
@@ -2297,6 +2313,8 @@ async def _read_obo_credential(
|
||||
)
|
||||
if credential is None:
|
||||
# No captured credential → the consent affordance is a re-login.
|
||||
if not prune_on_missing:
|
||||
return TokenLookupResult(kind="missing")
|
||||
return _no_token_result(app_state, user_id, server_name, TokenLookupResult(kind="missing"))
|
||||
return credential
|
||||
|
||||
@@ -2681,6 +2699,504 @@ async def get_obo_access_token_classified(
|
||||
return _token_result(app_state, user_id, server_name, access_token)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model-provider dynamic authentication
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deliberately lighter than get_obo_access_token_classified: a model backend is
|
||||
# ONE resource addressed by a single audience, redeemed identically for every
|
||||
# alias that points at it — not a per-server grant graph. So this owns no
|
||||
# dead-grant classification or re-consent affordance. It reuses the same mint
|
||||
# legs, credential store, RT-rotation CAS, cluster credential lock, AND the same
|
||||
# ``mcp_user_tokens`` mint-cache row the classified path uses (refresh_token=NULL,
|
||||
# "cache, not custody") — keyed under a synthetic ``__model_obo__:<audience>``
|
||||
# server name. The DB row shares the token across workers; a loop-local memo
|
||||
# avoids a SQL read + decrypt on every warm model turn.
|
||||
|
||||
MODEL_OBO_CACHE_PREFIX = "__model_obo__:"
|
||||
MODEL_APP_CACHE_PREFIX = "__model_app__:"
|
||||
_SYNTHETIC_TOKEN_PREFIXES = (MODEL_OBO_CACHE_PREFIX, MODEL_APP_CACHE_PREFIX)
|
||||
|
||||
|
||||
def _model_mint_memo(app_state: Any) -> dict[tuple[str, str], MCPUserTokenPlain]:
|
||||
"""Return the mcp-loop-owned model-token memo."""
|
||||
memo = getattr(app_state, "model_auth_token_cache", None)
|
||||
if not isinstance(memo, dict):
|
||||
memo = {}
|
||||
app_state.model_auth_token_cache = memo
|
||||
return memo
|
||||
|
||||
|
||||
def invalidate_model_mint_memo(
|
||||
app_state: Any,
|
||||
*,
|
||||
user_id: str,
|
||||
server_prefix: str = MODEL_OBO_CACHE_PREFIX,
|
||||
) -> int:
|
||||
"""Remove memo entries for one principal and synthetic-key prefix.
|
||||
|
||||
This must run on the manager's MCP loop. OIDC unlink schedules it there
|
||||
alongside deleting the corresponding DB rows, so the in-process fast path
|
||||
cannot extend a revoked bearer beyond the purge.
|
||||
"""
|
||||
memo = _model_mint_memo(app_state)
|
||||
keys = [key for key in memo if key[0] == user_id and key[1].startswith(server_prefix)]
|
||||
for key in keys:
|
||||
memo.pop(key, None)
|
||||
return len(keys)
|
||||
|
||||
|
||||
async def _serve_fresh_mint_cache(
|
||||
*,
|
||||
app_state: Any,
|
||||
token_store: MCPTokenStore,
|
||||
user_id: str,
|
||||
cache_server: str,
|
||||
audience: str,
|
||||
scopes: str,
|
||||
) -> str | None:
|
||||
"""Serve a fresh model-mint token from the loop memo or encrypted DB row."""
|
||||
key = (user_id, cache_server)
|
||||
memo = _model_mint_memo(app_state)
|
||||
plain = memo.get(key)
|
||||
if _is_fresh_obo_cache_row(plain, audience, scopes) and plain is not None:
|
||||
return plain["access_token"]
|
||||
memo.pop(key, None)
|
||||
try:
|
||||
plain = await asyncio.to_thread(token_store.get_user_token, user_id, cache_server)
|
||||
except MCPTokenDecryptError:
|
||||
# A mint-cache row encrypted under a retired key is only a cache miss.
|
||||
return None
|
||||
if not _is_fresh_obo_cache_row(plain, audience, scopes) or plain is None:
|
||||
return None
|
||||
memo[key] = plain
|
||||
return plain["access_token"]
|
||||
|
||||
|
||||
def _memoize_minted_token(
|
||||
app_state: Any,
|
||||
*,
|
||||
user_id: str,
|
||||
cache_server: str,
|
||||
access_token: str,
|
||||
expires_at: str | None,
|
||||
scopes: str,
|
||||
issuer: str,
|
||||
audience: str,
|
||||
) -> None:
|
||||
"""Install a freshly minted token in the loop-local memo."""
|
||||
now = datetime.now(UTC).replace(tzinfo=None).isoformat(timespec="seconds")
|
||||
_model_mint_memo(app_state)[(user_id, cache_server)] = {
|
||||
"user_id": user_id,
|
||||
"server_name": cache_server,
|
||||
"access_token": access_token,
|
||||
"refresh_token": None,
|
||||
"expires_at": expires_at,
|
||||
"scopes": scopes or None,
|
||||
"as_issuer": issuer,
|
||||
"audience": audience,
|
||||
"created": now,
|
||||
"last_refreshed": now,
|
||||
}
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _enter_mint_client(app_state: Any) -> Any:
|
||||
"""Yield the MCP-loop-owned mint client, with a temporary test fallback."""
|
||||
injected_client: httpx.AsyncClient | None = getattr(app_state, "obo_http_client", None)
|
||||
if injected_client is not None:
|
||||
yield injected_client
|
||||
return
|
||||
async with httpx.AsyncClient(timeout=_DEFAULT_HTTP_TIMEOUT) as mint_client:
|
||||
yield mint_client
|
||||
|
||||
|
||||
def _prune_model_mint_lock_when_idle(
|
||||
app_state: Any,
|
||||
user_id: str,
|
||||
cache_server: str,
|
||||
lock: asyncio.Lock,
|
||||
) -> None:
|
||||
"""Prune a synthetic lock after queued waiters have had a chance to acquire it."""
|
||||
|
||||
def _drop_if_idle() -> None:
|
||||
locks = getattr(app_state, "mcp_oauth_refresh_locks", None)
|
||||
waiters = getattr(lock, "_waiters", None)
|
||||
if (
|
||||
isinstance(locks, dict)
|
||||
and locks.get((user_id, cache_server)) is lock
|
||||
and not lock.locked()
|
||||
and not waiters
|
||||
):
|
||||
locks.pop((user_id, cache_server), None)
|
||||
|
||||
_drop_if_idle()
|
||||
# A released lock with a queued waiter is intentionally retained. Give the
|
||||
# waiter priority, then let the last participant prune on its own return.
|
||||
if getattr(app_state, "mcp_oauth_refresh_locks", {}).get((user_id, cache_server)) is lock:
|
||||
asyncio.get_running_loop().call_soon(_drop_if_idle)
|
||||
|
||||
|
||||
def _model_obo_cache_server(audience: str) -> str:
|
||||
"""Synthetic ``mcp_user_tokens`` server key for a model audience's mint-cache row.
|
||||
|
||||
The table's PK is ``(user_id, server_name)`` and model audiences vary, so the
|
||||
audience is embedded in the key — mirroring the in-process ``(user, audience)``
|
||||
key and the ``__model_obo__:<audience>`` single-flight lock. The ``__model_obo__:``
|
||||
prefix keeps these cache rows distinguishable from real oauth_user server rows.
|
||||
"""
|
||||
return f"{MODEL_OBO_CACHE_PREFIX}{audience}"
|
||||
|
||||
|
||||
async def mint_obo_access_token(
|
||||
*,
|
||||
app_state: Any,
|
||||
user_id: str,
|
||||
audience: str,
|
||||
force_refresh: bool = False,
|
||||
) -> str | None:
|
||||
"""Per-user Entra OBO access token for an arbitrary resource *audience*.
|
||||
|
||||
Redeems the user's captured refresh credential (``oidc_user_credentials``)
|
||||
for *audience* via the configured OBO grant profile, persists any rotated
|
||||
refresh token (value CAS, cluster-locked exactly like the MCP mint), and
|
||||
caches the minted access token in an ``mcp_user_tokens`` mint-cache row
|
||||
(``refresh_token=NULL``) keyed ``__model_obo__:<audience>`` — the same
|
||||
"cache, not custody" row the classified path uses — so the token is shared
|
||||
across worker nodes and inspectable, until shortly before expiry.
|
||||
|
||||
Returns ``None`` — the signal for callers to fall back to their static
|
||||
credential — when OIDC is disabled/unconfigured, the profile has no mint
|
||||
leg, the user has no captured credential (or it won't decrypt), or the mint
|
||||
is rejected. This is the model-provider entry point; the MCP-server path
|
||||
uses :func:`get_obo_access_token_classified`, which additionally owns
|
||||
per-server cache rows, dead-grant classification, and consent affordances
|
||||
this helper deliberately omits.
|
||||
"""
|
||||
if not user_id or not audience:
|
||||
return None
|
||||
oidc_config = getattr(app_state, "oidc_config", None)
|
||||
if oidc_config is None or not getattr(oidc_config, "enabled", False):
|
||||
return None
|
||||
token_store: MCPTokenStore | None = getattr(app_state, "mcp_token_store", None)
|
||||
storage = _get_storage(app_state)
|
||||
if token_store is None or storage is None:
|
||||
return None
|
||||
profile = str(getattr(oidc_config, "obo_grant_profile", "") or "")
|
||||
mint = _OBO_MINT_LEGS.get(profile)
|
||||
if mint is None:
|
||||
log.warning("model_obo.unsupported_grant_profile", profile=profile)
|
||||
return None
|
||||
issuer = str(getattr(oidc_config, "issuer", "") or "")
|
||||
cache_server = _model_obo_cache_server(audience)
|
||||
cached_token = await _serve_fresh_mint_cache(
|
||||
app_state=app_state,
|
||||
token_store=token_store,
|
||||
user_id=user_id,
|
||||
cache_server=cache_server,
|
||||
audience=audience,
|
||||
scopes="",
|
||||
)
|
||||
if cached_token and not force_refresh:
|
||||
_clear_refresh_backoff(app_state, user_id, cache_server)
|
||||
return cached_token
|
||||
# The model cache key is intentionally the cooldown key. The shared
|
||||
# ``__obo__:<issuer>`` credential key is also used by MCP OBO dispatch;
|
||||
# arming it here would let one broken model audience suppress every OBO
|
||||
# tool for this user.
|
||||
if not cached_token and _refresh_in_cooldown(app_state, user_id, cache_server):
|
||||
return None
|
||||
|
||||
# Single-flight the mint: a per-(user, audience) asyncio lock for local
|
||||
# coalescing, then the SAME per-(user, issuer) credential lock + cluster
|
||||
# advisory lock the MCP mint takes — the refresh credential is the shared
|
||||
# mutable resource (rotation write-back), so a model mint and an MCP mint for
|
||||
# the same user serialise on it cluster-wide. Order is always
|
||||
# audience → credential and nothing takes the reverse, so no deadlock.
|
||||
lock = _refresh_lock_for(app_state, user_id, cache_server)
|
||||
credential_key = f"__obo__:{issuer}"
|
||||
credential_lock = _refresh_lock_for(app_state, user_id, credential_key)
|
||||
pg_lock = await _acquire_pg_refresh_lock(storage, user_id, credential_key)
|
||||
try:
|
||||
async with lock, credential_lock, pg_lock:
|
||||
cached_token = await _serve_fresh_mint_cache(
|
||||
app_state=app_state,
|
||||
token_store=token_store,
|
||||
user_id=user_id,
|
||||
cache_server=cache_server,
|
||||
audience=audience,
|
||||
scopes="",
|
||||
)
|
||||
if cached_token and not force_refresh:
|
||||
_clear_refresh_backoff(app_state, user_id, cache_server)
|
||||
return cached_token
|
||||
if not cached_token and _refresh_in_cooldown(app_state, user_id, cache_server):
|
||||
return None
|
||||
|
||||
credential = await _read_obo_credential(
|
||||
app_state,
|
||||
token_store,
|
||||
user_id,
|
||||
cache_server,
|
||||
issuer,
|
||||
prune_on_missing=False,
|
||||
)
|
||||
if isinstance(credential, TokenLookupResult):
|
||||
_arm_cooldown(app_state, user_id, cache_server)
|
||||
return None
|
||||
|
||||
async def _persist_rotation(new_credential_rt: str) -> None:
|
||||
# Best-effort, same contract as the classified path: the mint
|
||||
# already produced a working token, so rotation write-back
|
||||
# failure must not discard it.
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
token_store.update_oidc_credential_after_redeem,
|
||||
user_id,
|
||||
issuer,
|
||||
refresh_token=new_credential_rt,
|
||||
expected_current=credential["refresh_token"],
|
||||
)
|
||||
except Exception:
|
||||
log.error(
|
||||
"model_obo.rotation_persist_failed",
|
||||
user_id=user_id,
|
||||
audience=audience,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
try:
|
||||
async with _enter_mint_client(app_state) as mint_client:
|
||||
tokens = await mint(
|
||||
oidc_config=oidc_config,
|
||||
credential_refresh_token=credential["refresh_token"],
|
||||
audience=audience,
|
||||
scopes="",
|
||||
http_client=mint_client,
|
||||
persist_rotation=_persist_rotation,
|
||||
)
|
||||
except MCPOAuthRefreshFailed:
|
||||
_arm_cooldown(app_state, user_id, cache_server)
|
||||
log.warning(
|
||||
"model_obo.mint_failed",
|
||||
user_id=user_id,
|
||||
audience=audience,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
access_token = tokens.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
_arm_cooldown(app_state, user_id, cache_server)
|
||||
log.warning(
|
||||
"model_obo.mint_missing_access_token",
|
||||
user_id=user_id,
|
||||
audience=audience,
|
||||
)
|
||||
return None
|
||||
expires_at = _expires_at_from_response(
|
||||
tokens, default_ttl_seconds=_OBO_DEFAULT_TTL_SECONDS
|
||||
)
|
||||
try:
|
||||
await _persist_obo_cache_row(
|
||||
token_store,
|
||||
user_id,
|
||||
cache_server,
|
||||
access_token=access_token,
|
||||
expires_at=expires_at,
|
||||
scopes="",
|
||||
issuer=issuer,
|
||||
audience=audience,
|
||||
)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"model_obo.cache_persist_failed",
|
||||
user_id=user_id,
|
||||
audience=audience,
|
||||
exc_info=True,
|
||||
)
|
||||
_memoize_minted_token(
|
||||
app_state,
|
||||
user_id=user_id,
|
||||
cache_server=cache_server,
|
||||
access_token=access_token,
|
||||
expires_at=expires_at,
|
||||
scopes="",
|
||||
issuer=issuer,
|
||||
audience=audience,
|
||||
)
|
||||
_clear_refresh_backoff(app_state, user_id, cache_server)
|
||||
log.info(
|
||||
"model_obo.minted",
|
||||
user_id=user_id,
|
||||
audience=audience,
|
||||
cache_server=cache_server,
|
||||
)
|
||||
return access_token
|
||||
finally:
|
||||
_prune_model_mint_lock_when_idle(app_state, user_id, cache_server, lock)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# App-identity (client-credentials) model token — Turnstone's own SSO app reg
|
||||
# ---------------------------------------------------------------------------
|
||||
# The ``auth_mode='entra_app'`` sibling of the OBO path: instead of a per-user
|
||||
# On-Behalf-Of token it mints an APP token from the ``[oidc]`` client id + secret
|
||||
# via the client-credentials grant. No user, no captured refresh token, no
|
||||
# rotation — one token per audience, shared by everyone, so a gateway resolves it
|
||||
# to a single machine (virtual-account) identity with no per-user attribution. It
|
||||
# reuses the same DB mint-cache under a synthetic ``__app__`` user.
|
||||
|
||||
_APP_CACHE_USER = "__app__"
|
||||
|
||||
|
||||
def _model_app_cache_server(audience: str) -> str:
|
||||
"""Synthetic ``mcp_user_tokens`` key for an app-credential mint-cache row.
|
||||
|
||||
App tokens carry no user, so they cache once per audience under the shared
|
||||
``__app__`` pseudo-user; the ``__model_app__:`` prefix keeps them distinct
|
||||
from per-user OBO rows (``__model_obo__:``) and real oauth_user server rows.
|
||||
"""
|
||||
return f"{MODEL_APP_CACHE_PREFIX}{audience}"
|
||||
|
||||
|
||||
async def mint_app_access_token(
|
||||
*,
|
||||
app_state: Any,
|
||||
audience: str,
|
||||
force_refresh: bool = False,
|
||||
) -> str | None:
|
||||
"""App-identity Entra access token for *audience* via client-credentials.
|
||||
|
||||
Uses Turnstone's own SSO app registration (``[oidc]`` ``client_id`` +
|
||||
``client_secret``) — no user, no captured refresh token, no rotation. One
|
||||
token per audience, shared by every caller and cached in an
|
||||
``mcp_user_tokens`` row under the synthetic ``__app__`` user until shortly
|
||||
before expiry. This is the ``auth_mode='entra_app'`` backend entry point —
|
||||
the "we already have SSO, let the app call the gateway as its own managed
|
||||
identity" path (a gateway resolves it to one virtual account, no per-user
|
||||
attribution). Because it needs no user context it also serves utility /
|
||||
coordinator / service / CLI turns that OBO cannot. Returns ``None`` — the
|
||||
signal to fall back to the static credential — when OIDC is
|
||||
disabled/unconfigured, the app has no secret, or the grant is rejected.
|
||||
"""
|
||||
if not audience:
|
||||
return None
|
||||
oidc_config = getattr(app_state, "oidc_config", None)
|
||||
if oidc_config is None or not getattr(oidc_config, "enabled", False):
|
||||
return None
|
||||
profile = str(getattr(oidc_config, "obo_grant_profile", "") or "")
|
||||
if profile != "entra":
|
||||
log.warning("model_app.unsupported_grant_profile", profile=profile)
|
||||
return None
|
||||
client_id = str(getattr(oidc_config, "client_id", "") or "")
|
||||
client_secret = str(getattr(oidc_config, "client_secret", "") or "")
|
||||
token_endpoint = str(getattr(oidc_config, "token_endpoint", "") or "")
|
||||
token_store: MCPTokenStore | None = getattr(app_state, "mcp_token_store", None)
|
||||
storage = _get_storage(app_state)
|
||||
if token_store is None or storage is None:
|
||||
return None
|
||||
issuer = str(getattr(oidc_config, "issuer", "") or "")
|
||||
|
||||
cache_server = _model_app_cache_server(audience)
|
||||
cached_token = await _serve_fresh_mint_cache(
|
||||
app_state=app_state,
|
||||
token_store=token_store,
|
||||
user_id=_APP_CACHE_USER,
|
||||
cache_server=cache_server,
|
||||
audience=audience,
|
||||
scopes="",
|
||||
)
|
||||
if cached_token and not force_refresh:
|
||||
_clear_refresh_backoff(app_state, _APP_CACHE_USER, cache_server)
|
||||
return cached_token
|
||||
if not cached_token and _refresh_in_cooldown(app_state, _APP_CACHE_USER, cache_server):
|
||||
return None
|
||||
if not (client_id and client_secret and token_endpoint):
|
||||
_arm_cooldown(app_state, _APP_CACHE_USER, cache_server)
|
||||
log.warning(
|
||||
"model_app.credentials_unavailable",
|
||||
has_client_id=bool(client_id),
|
||||
has_client_secret=bool(client_secret),
|
||||
has_token_endpoint=bool(token_endpoint),
|
||||
)
|
||||
return None
|
||||
|
||||
# Single-flight the mint. No per-user credential to rotate, so only the
|
||||
# per-audience local + cluster lock is taken (no credential lock).
|
||||
lock = _refresh_lock_for(app_state, _APP_CACHE_USER, cache_server)
|
||||
pg_lock = await _acquire_pg_refresh_lock(storage, _APP_CACHE_USER, cache_server)
|
||||
try:
|
||||
async with lock, pg_lock:
|
||||
cached_token = await _serve_fresh_mint_cache(
|
||||
app_state=app_state,
|
||||
token_store=token_store,
|
||||
user_id=_APP_CACHE_USER,
|
||||
cache_server=cache_server,
|
||||
audience=audience,
|
||||
scopes="",
|
||||
)
|
||||
if cached_token and not force_refresh:
|
||||
_clear_refresh_backoff(app_state, _APP_CACHE_USER, cache_server)
|
||||
return cached_token
|
||||
if not cached_token and _refresh_in_cooldown(app_state, _APP_CACHE_USER, cache_server):
|
||||
return None
|
||||
|
||||
try:
|
||||
async with _enter_mint_client(app_state) as mint_client:
|
||||
tokens = await _obo_token_post(
|
||||
token_endpoint=token_endpoint,
|
||||
data={
|
||||
"grant_type": "client_credentials",
|
||||
"client_id": client_id,
|
||||
"client_secret": client_secret,
|
||||
"scope": f"{audience}/.default",
|
||||
},
|
||||
http_client=mint_client,
|
||||
leg="client-credentials",
|
||||
)
|
||||
except MCPOAuthRefreshFailed:
|
||||
_arm_cooldown(app_state, _APP_CACHE_USER, cache_server)
|
||||
log.warning("model_app.mint_failed", audience=audience, exc_info=True)
|
||||
return None
|
||||
|
||||
access_token = tokens.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
_arm_cooldown(app_state, _APP_CACHE_USER, cache_server)
|
||||
log.warning("model_app.mint_missing_access_token", audience=audience)
|
||||
return None
|
||||
expires_at = _expires_at_from_response(
|
||||
tokens, default_ttl_seconds=_OBO_DEFAULT_TTL_SECONDS
|
||||
)
|
||||
try:
|
||||
await _persist_obo_cache_row(
|
||||
token_store,
|
||||
_APP_CACHE_USER,
|
||||
cache_server,
|
||||
access_token=access_token,
|
||||
expires_at=expires_at,
|
||||
scopes="",
|
||||
issuer=issuer,
|
||||
audience=audience,
|
||||
)
|
||||
except Exception:
|
||||
log.warning("model_app.cache_persist_failed", audience=audience, exc_info=True)
|
||||
_memoize_minted_token(
|
||||
app_state,
|
||||
user_id=_APP_CACHE_USER,
|
||||
cache_server=cache_server,
|
||||
access_token=access_token,
|
||||
expires_at=expires_at,
|
||||
scopes="",
|
||||
issuer=issuer,
|
||||
audience=audience,
|
||||
)
|
||||
_clear_refresh_backoff(app_state, _APP_CACHE_USER, cache_server)
|
||||
log.info("model_app.minted", audience=audience, cache_server=cache_server)
|
||||
return access_token
|
||||
finally:
|
||||
_prune_model_mint_lock_when_idle(app_state, _APP_CACHE_USER, cache_server, lock)
|
||||
|
||||
|
||||
def _token_needs_refresh(expires_at: str | None) -> bool:
|
||||
"""Return True when *expires_at* is missing, malformed, or within the skew window."""
|
||||
if not expires_at:
|
||||
@@ -3744,7 +4260,16 @@ async def _handle_mcp_oauth_list_connections_inner(request: Request) -> Response
|
||||
obo_names = await asyncio.to_thread(obo_server_names, storage)
|
||||
except Exception:
|
||||
obo_names = set()
|
||||
return JSONResponse({"connections": [r for r in rows if r["server_name"] not in obo_names]})
|
||||
return JSONResponse(
|
||||
{
|
||||
"connections": [
|
||||
r
|
||||
for r in rows
|
||||
if r["server_name"] not in obo_names
|
||||
and not str(r["server_name"]).startswith(_SYNTHETIC_TOKEN_PREFIXES)
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def handle_mcp_oauth_revoke_connection(request: Request) -> Response:
|
||||
@@ -3877,6 +4402,19 @@ async def _handle_mcp_oauth_revoke_connection_inner(request: Request) -> Respons
|
||||
server_name = request.path_params.get("server_name", "").strip()
|
||||
if not server_name:
|
||||
return JSONResponse({"error": "Missing server_name"}, status_code=400)
|
||||
# Synthetic model rows are mint caches, not user-revocable MCP
|
||||
# connections. Check before row existence to avoid a 404/409 oracle for
|
||||
# whether this user currently has a token for a guessed audience.
|
||||
if server_name.startswith(_SYNTHETIC_TOKEN_PREFIXES):
|
||||
return JSONResponse(
|
||||
{
|
||||
"error": (
|
||||
"This is an internal model-authentication cache, not an MCP "
|
||||
"connection. It cannot be disconnected from this endpoint."
|
||||
)
|
||||
},
|
||||
status_code=409,
|
||||
)
|
||||
|
||||
storage = _get_storage(request.app.state)
|
||||
if storage is None:
|
||||
|
||||
@@ -17,6 +17,12 @@ from turnstone.core.providers import LLMProvider, create_client, create_provider
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
MODEL_AUTH_MODES = frozenset({"static", "entra_obo", "entra_app"})
|
||||
|
||||
|
||||
class ModelAuthConfigError(ValueError):
|
||||
"""A model definition contains unsafe or internally inconsistent auth settings."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model configuration
|
||||
@@ -48,6 +54,39 @@ class ModelConfig:
|
||||
# Server compatibility settings for openai-compatible backends.
|
||||
# Populated from capabilities["server_compat"] during load.
|
||||
server_compat: dict[str, Any] = field(default_factory=dict)
|
||||
# Backend credential mode. ``static`` (default) sends ``api_key`` unchanged;
|
||||
# ``entra_obo`` mints a caller-delegated Entra token; ``entra_app`` mints a
|
||||
# shared app-identity token. Dynamic tokens are bound as the SDK credential
|
||||
# (x-api-key for Anthropic surfaces, Authorization: Bearer for OpenAI-style).
|
||||
# ``obo_audience`` is the exact operator-approved resource App ID URI.
|
||||
auth_mode: str = "static"
|
||||
obo_audience: str = ""
|
||||
|
||||
|
||||
def _normalize_auth_mode(alias: str, mode: Any, audience: Any) -> tuple[str, str]:
|
||||
"""Validate and normalize one model's backend-auth configuration.
|
||||
|
||||
DB and config.toml rows share this path so a typo cannot silently downgrade
|
||||
dynamic authentication to a static key. Audience values remain literal:
|
||||
environment expansion would make authorization node-dependent and could
|
||||
bypass the admin allow-list and length boundary.
|
||||
"""
|
||||
normalized_mode = str(mode or "static").strip() or "static"
|
||||
normalized_audience = str(audience or "").strip()
|
||||
if normalized_mode not in MODEL_AUTH_MODES:
|
||||
raise ModelAuthConfigError(
|
||||
f"Model '{alias}' has invalid auth_mode {normalized_mode!r}; "
|
||||
f"expected one of {sorted(MODEL_AUTH_MODES)}"
|
||||
)
|
||||
if normalized_mode != "static" and not normalized_audience:
|
||||
raise ModelAuthConfigError(
|
||||
f"Model '{alias}' requires obo_audience when auth_mode is {normalized_mode!r}"
|
||||
)
|
||||
if any(ord(ch) < 32 or ord(ch) == 127 for ch in normalized_audience):
|
||||
raise ModelAuthConfigError(f"Model '{alias}' obo_audience contains control characters")
|
||||
if len(normalized_audience) > 2048:
|
||||
raise ModelAuthConfigError(f"Model '{alias}' obo_audience exceeds 2048 characters")
|
||||
return normalized_mode, normalized_audience
|
||||
|
||||
|
||||
def _api_surface_of(cfg: ModelConfig) -> str | None:
|
||||
@@ -145,9 +184,18 @@ class ModelRegistry:
|
||||
raise ValueError(f"Unknown model alias: {alias}")
|
||||
if alias not in self._clients:
|
||||
cfg = self._models[alias]
|
||||
# An entra_obo / entra_app backend authenticates per-call via a
|
||||
# minted token bound with ``client.with_options(api_key=...)``, so
|
||||
# the cached client only needs to CONSTRUCT — feed a placeholder
|
||||
# when no static fallback key is set (the SDKs reject an empty key
|
||||
# that also has no env fallback). The real credential is supplied
|
||||
# per call and never rides on this base client object.
|
||||
client_key = cfg.api_key
|
||||
if not client_key and cfg.auth_mode in ("entra_obo", "entra_app"):
|
||||
client_key = "backend-auth-placeholder-unused"
|
||||
try:
|
||||
self._clients[alias] = create_client(
|
||||
cfg.provider, base_url=cfg.base_url, api_key=cfg.api_key
|
||||
cfg.provider, base_url=cfg.base_url, api_key=client_key
|
||||
)
|
||||
except ValueError:
|
||||
# create_client's own misconfig errors already carry
|
||||
@@ -198,6 +246,10 @@ class ModelRegistry:
|
||||
"""Check if *alias* exists in the registry."""
|
||||
return alias in self._models
|
||||
|
||||
def has_dynamic_auth(self) -> bool:
|
||||
"""Return whether any alias needs a runtime-minted backend credential."""
|
||||
return any(cfg.auth_mode != "static" for cfg in self._models.values())
|
||||
|
||||
def list_aliases(self) -> list[str]:
|
||||
"""Return all registered model aliases."""
|
||||
return list(self._models.keys())
|
||||
@@ -272,8 +324,8 @@ class ModelRegistry:
|
||||
self.task_model = task_model
|
||||
self.task_effort = task_effort
|
||||
# Selective teardown — close + drop only clients whose
|
||||
# connection target actually changed (alias removed, or
|
||||
# base_url / api_key / provider differs). Keeps connection
|
||||
# construction/connection target changed (alias removed, or
|
||||
# base_url / api_key / provider / auth_mode differs). Keeps connection
|
||||
# pools warm for the common admin-edit case where only
|
||||
# ``model`` / ``temperature`` / ``context_window`` changed.
|
||||
for alias, client in list(self._clients.items()):
|
||||
@@ -285,6 +337,7 @@ class ModelRegistry:
|
||||
or old_cfg.base_url != new_cfg.base_url
|
||||
or old_cfg.api_key != new_cfg.api_key
|
||||
or old_cfg.provider != new_cfg.provider
|
||||
or old_cfg.auth_mode != new_cfg.auth_mode
|
||||
):
|
||||
if hasattr(client, "close"):
|
||||
client.close()
|
||||
@@ -432,6 +485,13 @@ def load_model_registry(
|
||||
# pre-052 row missing these columns degrades gracefully.
|
||||
row_surface_persisted_reasoning = bool(row.get("surface_persisted_reasoning", True))
|
||||
row_replay_reasoning = bool(row.get("replay_reasoning_to_model", False))
|
||||
# Per-user OBO auth (defaults match a pre-068 row missing the
|
||||
# columns → "static", no audience → unchanged behaviour).
|
||||
row_auth_mode, row_obo_audience = _normalize_auth_mode(
|
||||
alias,
|
||||
row.get("auth_mode"),
|
||||
row.get("obo_audience"),
|
||||
)
|
||||
configs[alias] = ModelConfig(
|
||||
alias=alias,
|
||||
base_url=row_base_url,
|
||||
@@ -449,7 +509,14 @@ def load_model_registry(
|
||||
surface_persisted_reasoning=row_surface_persisted_reasoning,
|
||||
replay_reasoning_to_model=row_replay_reasoning,
|
||||
server_compat=row_server_compat,
|
||||
auth_mode=row_auth_mode,
|
||||
obo_audience=row_obo_audience,
|
||||
)
|
||||
except ModelAuthConfigError:
|
||||
# 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.
|
||||
raise
|
||||
except Exception:
|
||||
if strict:
|
||||
raise
|
||||
@@ -502,6 +569,11 @@ def load_model_registry(
|
||||
entry_server_compat = entry_caps.pop("server_compat", {})
|
||||
if not isinstance(entry_server_compat, dict):
|
||||
entry_server_compat = {}
|
||||
entry_auth_mode, entry_obo_audience = _normalize_auth_mode(
|
||||
alias,
|
||||
entry.get("auth_mode", "static"),
|
||||
entry.get("obo_audience", ""),
|
||||
)
|
||||
configs[alias] = ModelConfig(
|
||||
alias=alias,
|
||||
base_url=entry_base_url,
|
||||
@@ -520,6 +592,8 @@ def load_model_registry(
|
||||
max_tokens=entry_max_tokens,
|
||||
reasoning_effort=entry_effort,
|
||||
server_compat=entry_server_compat,
|
||||
auth_mode=entry_auth_mode,
|
||||
obo_audience=entry_obo_audience,
|
||||
)
|
||||
|
||||
# 3. Back-compat shim: synthesize a "default" alias from CLI/auto-detected
|
||||
|
||||
@@ -374,6 +374,9 @@ class ModelLane:
|
||||
registry: ModelRegistry | None = None
|
||||
temperature: float | None = None
|
||||
reasoning_effort: str | None = None
|
||||
# Runtime credential resolver supplied by the host that owns OAuth state.
|
||||
# Kept on the lane because this synchronous module has no manager singleton.
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None
|
||||
|
||||
|
||||
def resolve_lane(
|
||||
@@ -386,6 +389,7 @@ def resolve_lane(
|
||||
capabilities: ModelCapabilities | None = None,
|
||||
extra_params: dict[str, Any] | None | EllipsisType = ...,
|
||||
config_store: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None,
|
||||
) -> ModelLane:
|
||||
"""Build a :class:`ModelLane`, resolving what the caller didn't supply.
|
||||
|
||||
@@ -428,6 +432,7 @@ def resolve_lane(
|
||||
registry=registry,
|
||||
temperature=resolve_temperature_setting(cfg, config_store),
|
||||
reasoning_effort=resolve_effort_setting(cfg, config_store),
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
)
|
||||
|
||||
|
||||
@@ -636,6 +641,7 @@ def model_turn(
|
||||
wire_id_map: dict[str, str] | None = None,
|
||||
resolve_attachments: Callable[[list[str]], dict[str, Any]] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
backend_auth_token: str | None = None,
|
||||
) -> ModelTurnResult:
|
||||
"""Advance a trajectory by one model turn: lower, sample, re-ingest.
|
||||
|
||||
@@ -698,6 +704,18 @@ def model_turn(
|
||||
policy inside ``create_streaming`` and propagate unchanged). An
|
||||
aborted *cancel_ref* suppresses retries — a deadline that closed the
|
||||
stream must not have the request resurrected behind its back.
|
||||
|
||||
*backend_auth_token* is a delegated-user or app-identity credential for a
|
||||
dynamically authenticated backend. When set, the call is issued on
|
||||
``lane.client.with_options(api_key=...)`` — a copy that reuses the
|
||||
client's connection pool but swaps the credential, so the SDK emits it as
|
||||
its own auth header (``x-api-key`` for Anthropic, ``Authorization: Bearer``
|
||||
for OpenAI-style). This is deliberately NOT header injection via
|
||||
``extra_headers``: the Anthropic SDK does not let ``extra_headers``
|
||||
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.
|
||||
"""
|
||||
if mint is not None and wire_id_map is None:
|
||||
raise ValueError(
|
||||
@@ -725,6 +743,16 @@ def model_turn(
|
||||
or (lane.capabilities.default_reasoning_effort if lane.capabilities else None)
|
||||
or None
|
||||
)
|
||||
resolved_backend_auth = backend_auth_token
|
||||
if resolved_backend_auth is None and lane.backend_auth_resolver is not None:
|
||||
resolved_backend_auth = lane.backend_auth_resolver(lane.alias)
|
||||
# Bind once: SDK ``with_options`` preserves the base transport/pool while
|
||||
# replacing only the provider credential. A drain retry reuses this client.
|
||||
call_client = (
|
||||
lane.client.with_options(api_key=resolved_backend_auth)
|
||||
if resolved_backend_auth
|
||||
else lane.client
|
||||
)
|
||||
attempt = 0
|
||||
while True:
|
||||
# ``create_streaming`` stays OUTSIDE the try: every adapter issues
|
||||
@@ -732,8 +760,11 @@ def model_turn(
|
||||
# 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=lane.client,
|
||||
client=call_client,
|
||||
model=lane.model,
|
||||
messages=wire,
|
||||
tools=tools,
|
||||
|
||||
@@ -58,6 +58,7 @@ from turnstone.core.trajectory import Turn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
|
||||
from turnstone.core.judge import JudgeConfig
|
||||
from turnstone.core.providers._protocol import LLMProvider, ModelCapabilities
|
||||
@@ -271,6 +272,7 @@ class OutputGuardJudge:
|
||||
session_capabilities: ModelCapabilities | None = None,
|
||||
session_model_alias: str = "",
|
||||
config_store: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None,
|
||||
) -> None:
|
||||
self._config = config
|
||||
# Carried into the per-evaluation ModelLane so extra_params, the
|
||||
@@ -278,6 +280,7 @@ class OutputGuardJudge:
|
||||
# registry like every other lane.
|
||||
self._model_registry = model_registry
|
||||
self._config_store = config_store
|
||||
self._backend_auth_resolver = backend_auth_resolver
|
||||
# Caller's resolved session-model caps (config/registry-aware): the wire
|
||||
# capabilities + window when this judge inherits the session model, and
|
||||
# the alias path's window fallback. The window comes ONLY from these
|
||||
@@ -538,6 +541,7 @@ class OutputGuardJudge:
|
||||
registry=self._model_registry,
|
||||
capabilities=self._capabilities,
|
||||
config_store=self._config_store,
|
||||
backend_auth_resolver=self._backend_auth_resolver,
|
||||
)
|
||||
# Sampling deliberately not pinned (house rule) — the lane inherits
|
||||
# the guard model's full assignment scheme, effort included: a
|
||||
|
||||
@@ -37,6 +37,8 @@ from turnstone.core.model_turn import model_turn, resolve_lane
|
||||
from turnstone.core.trajectory import AttachmentRef, Role, TextBlock, Turn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from turnstone.core.providers._protocol import LLMProvider
|
||||
|
||||
log = get_logger(__name__)
|
||||
@@ -79,6 +81,7 @@ def describe(
|
||||
registry: Any | None = None,
|
||||
config_store: Any | None = None,
|
||||
capabilities: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None,
|
||||
) -> str:
|
||||
"""Perceive ``parts`` via the perception model, returning the text.
|
||||
|
||||
@@ -120,6 +123,7 @@ def describe(
|
||||
registry=registry,
|
||||
config_store=config_store,
|
||||
capabilities=capabilities,
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
)
|
||||
try:
|
||||
result = model_turn(
|
||||
@@ -140,7 +144,12 @@ def describe(
|
||||
# subsequent turn.
|
||||
_CACHE_MAX = 256
|
||||
_cache_lock = threading.Lock()
|
||||
_cache: dict[str, str] = {}
|
||||
_cache: dict[tuple[str, str, str], str] = {}
|
||||
|
||||
|
||||
def _cache_key(*, principal_id: str, alias: str, content_hash: str) -> tuple[str, str, str]:
|
||||
"""Partition perceived content by the identity that authorized the call."""
|
||||
return (principal_id, alias, content_hash)
|
||||
|
||||
|
||||
def _clear_perception_cache_for_test() -> None:
|
||||
@@ -153,6 +162,7 @@ def describe_cached(
|
||||
provider: LLMProvider,
|
||||
client: Any,
|
||||
model: str,
|
||||
principal_id: str,
|
||||
alias: str,
|
||||
content_hash: str,
|
||||
parts: list[dict[str, Any]],
|
||||
@@ -160,14 +170,17 @@ def describe_cached(
|
||||
registry: Any | None = None,
|
||||
config_store: Any | None = None,
|
||||
capabilities: Any | None = None,
|
||||
backend_auth_resolver: Callable[[str], str | None] | None = None,
|
||||
) -> str:
|
||||
"""Memoized, non-raising :func:`describe` for the wire fallback.
|
||||
|
||||
Keyed by ``(alias, content_hash)``. Returns ``""`` on a backend failure (a
|
||||
placeholder is rendered upstream) and does *not* cache failures, so a
|
||||
transient outage doesn't poison the memo.
|
||||
Keyed by ``(principal_id, alias, content_hash)``. The principal partition is
|
||||
load-bearing for delegated backend authentication: a description produced
|
||||
under one user's OBO grant must never be served to another user without a
|
||||
call authorized as that user. Returns ``""`` on a backend failure (a
|
||||
placeholder is rendered upstream) and does *not* cache failures.
|
||||
"""
|
||||
key = f"{alias}:{content_hash}"
|
||||
key = _cache_key(principal_id=principal_id, alias=alias, content_hash=content_hash)
|
||||
with _cache_lock:
|
||||
if key in _cache:
|
||||
return _cache[key]
|
||||
@@ -182,6 +195,7 @@ def describe_cached(
|
||||
registry=registry,
|
||||
config_store=config_store,
|
||||
capabilities=capabilities,
|
||||
backend_auth_resolver=backend_auth_resolver,
|
||||
)
|
||||
except PerceptionBackendError as exc:
|
||||
log.warning("perception fallback failed (alias=%s): %s", alias, exc)
|
||||
@@ -193,13 +207,14 @@ def describe_cached(
|
||||
return text
|
||||
|
||||
|
||||
def describe_peek(*, alias: str, content_hash: str) -> str | None:
|
||||
"""Return the memoized description for ``(alias, content_hash)`` without
|
||||
computing, or ``None`` if absent.
|
||||
def describe_peek(*, principal_id: str, alias: str, content_hash: str) -> str | None:
|
||||
"""Return the principal-scoped memoized description without computing.
|
||||
|
||||
Lets the wire resolver skip the expensive parts build (a PDF rasterize) when
|
||||
the description is already memoized from an earlier send — :func:`describe_cached`
|
||||
ignores ``parts`` on a hit, so building them first would be pure waste.
|
||||
"""
|
||||
with _cache_lock:
|
||||
return _cache.get(f"{alias}:{content_hash}")
|
||||
return _cache.get(
|
||||
_cache_key(principal_id=principal_id, alias=alias, content_hash=content_hash)
|
||||
)
|
||||
|
||||
+149
-3
@@ -1431,6 +1431,10 @@ def _tool_turn_meta(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BackendAuthUnavailableError(RuntimeError):
|
||||
"""A fail-closed dynamic model credential could not be resolved."""
|
||||
|
||||
|
||||
class ChatSession:
|
||||
# The mid-turn interjection queue's cap — an ALIAS of the shared
|
||||
# per-workstream backpressure bound (see workstream.PENDING_SENDS_MAX):
|
||||
@@ -1839,6 +1843,10 @@ class ChatSession:
|
||||
# _task_tools, no listeners, no resource/prompt catalogs, and the
|
||||
# refresh callbacks stay inert (they all guard on _mcp_client).
|
||||
# Task agents keep their native tools; only the MCP surface closes.
|
||||
# Model authentication is host infrastructure, not an MCP tool-surface
|
||||
# capability. Preserve the raw manager even when the persona gate hides
|
||||
# MCP tools, resources, and prompts or a resume drops that surface.
|
||||
self._mcp_mint_client = mcp_client
|
||||
self._mcp_client = mcp_client if self._persona_mcp else None
|
||||
# True when a real client was withheld by the persona gate (as
|
||||
# opposed to no MCP in the deployment at all). Mid-session
|
||||
@@ -3283,6 +3291,10 @@ class ChatSession:
|
||||
self._set_session_tools([])
|
||||
self._render_agent_tool_descriptions()
|
||||
|
||||
def set_model_mint_client(self, client: MCPClientManager | None) -> None:
|
||||
"""Update the ungated model-auth manager after a runtime registry reload."""
|
||||
self._mcp_mint_client = client
|
||||
|
||||
def _handle_mcp_refresh(self, arg: str) -> None:
|
||||
"""Handle ``/mcp refresh [server]``."""
|
||||
assert self._mcp_client is not None
|
||||
@@ -4582,12 +4594,17 @@ class ChatSession:
|
||||
return None
|
||||
from turnstone.core.perception import describe_cached, describe_peek
|
||||
|
||||
# Peek the (alias, content_hash) memo BEFORE building parts: for a PDF,
|
||||
# Peek the (principal, alias, content_hash) memo BEFORE building parts: for a PDF,
|
||||
# _perception_parts rasterizes every page, but describe_cached returns a
|
||||
# memoized description without touching parts on a hit — so on a cross-send
|
||||
# hit the rasterize would be pure waste.
|
||||
content_hash = str(att.get("attachment_id"))
|
||||
text = describe_peek(alias=alias, content_hash=content_hash)
|
||||
principal_id = (self._mcp_effective_user_id or "").strip()
|
||||
text = describe_peek(
|
||||
principal_id=principal_id,
|
||||
alias=alias,
|
||||
content_hash=content_hash,
|
||||
)
|
||||
if text is None:
|
||||
parts = self._perception_parts(att, kind)
|
||||
if not parts:
|
||||
@@ -4596,6 +4613,7 @@ class ChatSession:
|
||||
provider=provider,
|
||||
client=client,
|
||||
model=model,
|
||||
principal_id=principal_id,
|
||||
alias=alias,
|
||||
content_hash=content_hash,
|
||||
parts=parts,
|
||||
@@ -4608,6 +4626,9 @@ class ChatSession:
|
||||
registry=self._registry,
|
||||
config_store=self._config_store,
|
||||
capabilities=caps,
|
||||
# The memo is partitioned by the same effective principal used
|
||||
# by the resolver, so delegated output cannot cross users.
|
||||
backend_auth_resolver=self._model_backend_auth_token,
|
||||
)
|
||||
if not text:
|
||||
return None
|
||||
@@ -5204,6 +5225,7 @@ class ChatSession:
|
||||
registry=self._registry,
|
||||
capabilities=caps,
|
||||
config_store=self._config_store,
|
||||
backend_auth_resolver=self._model_backend_auth_token,
|
||||
)
|
||||
result = model_turn(
|
||||
lane,
|
||||
@@ -5439,10 +5461,14 @@ class ChatSession:
|
||||
tracker = self._get_health_tracker()
|
||||
|
||||
try:
|
||||
result = self._try_stream(self.client, self.model, msgs)
|
||||
result = self._try_stream(self.client, self.model, msgs, model_alias=self._model_alias)
|
||||
if tracker:
|
||||
tracker.record_success()
|
||||
return result
|
||||
except BackendAuthUnavailableError:
|
||||
# Explicit fail-closed policy: never reinterpret an authentication
|
||||
# refusal as backend health and never route it to a static fallback.
|
||||
raise
|
||||
except Exception as primary_err:
|
||||
if tracker:
|
||||
tracker.record_failure()
|
||||
@@ -5502,6 +5528,8 @@ class ChatSession:
|
||||
if fb_tracker:
|
||||
fb_tracker.record_success()
|
||||
return result
|
||||
except BackendAuthUnavailableError:
|
||||
raise
|
||||
except Exception as fb_err:
|
||||
if fb_tracker:
|
||||
fb_tracker.record_failure()
|
||||
@@ -5550,6 +5578,14 @@ class ChatSession:
|
||||
role_counts,
|
||||
)
|
||||
msgs = self._maybe_attach_vllm_chat_reasoning(msgs, prov, model_alias)
|
||||
# Dynamic backend auth: bind the minted token as the SDK client's
|
||||
# api_key (with_options reuses the connection pool) so it becomes the
|
||||
# provider's own credential header — extra_headers can't override the
|
||||
# Anthropic SDK's x-api-key. None → the backend's static key stands.
|
||||
# Resolved once, outside the retry loop.
|
||||
backend_auth_token = self._model_backend_auth_token(model_alias or "")
|
||||
if backend_auth_token:
|
||||
client = client.with_options(api_key=backend_auth_token)
|
||||
last_err: Exception | None = None
|
||||
for attempt in range(self._MAX_RETRIES + 1):
|
||||
self._check_cancelled()
|
||||
@@ -6034,6 +6070,109 @@ class ChatSession:
|
||||
"""
|
||||
return self._acting_user_id or self._mcp_user_id
|
||||
|
||||
def _model_backend_auth_token(self, alias: str) -> str | None:
|
||||
"""Resolve the delegated-user or app-identity credential for *alias*.
|
||||
|
||||
The caller binds the returned token as the SDK client's ``api_key`` via
|
||||
``with_options``, which preserves the connection pool and lets each SDK
|
||||
emit its native credential header. Header injection is deliberately not
|
||||
used: Anthropic does not allow ``extra_headers`` to replace ``x-api-key``.
|
||||
|
||||
Every session-owned lane, including judge, output guard, and perception,
|
||||
resolves through this same effective principal. ``entra_app`` remains an
|
||||
explicit model-definition choice; it is never inferred from a missing
|
||||
user or failed OBO mint.
|
||||
|
||||
A dynamic alias with no static key always fails closed rather than
|
||||
issuing the SDK-construction placeholder. When a real static key exists,
|
||||
mint failures retain that explicit fallback unless the operator enables
|
||||
``model.auth_fail_closed``. An ``entra_obo`` call with no user always
|
||||
fails closed regardless of fallback policy.
|
||||
"""
|
||||
registry = self._registry
|
||||
if registry is None or not alias:
|
||||
return None
|
||||
try:
|
||||
cfg = registry.get_config(alias)
|
||||
except (KeyError, ValueError):
|
||||
return None
|
||||
mode = getattr(cfg, "auth_mode", "static")
|
||||
if mode not in ("entra_obo", "entra_app") or not cfg.obo_audience:
|
||||
return None
|
||||
has_static_key = bool(getattr(cfg, "api_key", ""))
|
||||
configured_fail_closed = bool(
|
||||
self._config_store is not None and self._config_store.get("model.auth_fail_closed")
|
||||
)
|
||||
must_fail_closed = configured_fail_closed or not has_static_key
|
||||
user_id = ""
|
||||
if mode == "entra_obo":
|
||||
user_id = (self._mcp_effective_user_id or "").strip()
|
||||
if not user_id:
|
||||
log.warning(
|
||||
"model_obo.no_user_context",
|
||||
alias=alias,
|
||||
audience=cfg.obo_audience,
|
||||
has_static_key=has_static_key,
|
||||
)
|
||||
raise BackendAuthUnavailableError(
|
||||
f"Delegated backend authentication has no user for model alias {alias!r}"
|
||||
)
|
||||
mcp = self._mcp_mint_client
|
||||
if mcp is None:
|
||||
log.warning(
|
||||
"model_backend_auth.mint_client_unavailable",
|
||||
alias=alias,
|
||||
auth_mode=mode,
|
||||
audience=cfg.obo_audience,
|
||||
has_static_key=has_static_key,
|
||||
)
|
||||
if must_fail_closed:
|
||||
raise BackendAuthUnavailableError(
|
||||
f"Dynamic backend authentication unavailable for model alias {alias!r}"
|
||||
)
|
||||
return None
|
||||
if mode == "entra_app":
|
||||
# App/managed identity via client-credentials — Turnstone's own SSO
|
||||
# app reg. This is used only when the model definition explicitly
|
||||
# selects entra_app; missing OBO context never switches grant modes.
|
||||
# The gateway resolves it to one shared virtual account (no per-user
|
||||
# attribution).
|
||||
token = mcp.mint_app_token_sync(audience=cfg.obo_audience)
|
||||
if not token:
|
||||
log.warning(
|
||||
"model_app.fallback_to_static",
|
||||
alias=alias,
|
||||
audience=cfg.obo_audience,
|
||||
has_static_key=has_static_key,
|
||||
)
|
||||
if must_fail_closed:
|
||||
raise BackendAuthUnavailableError(
|
||||
f"App backend authentication unavailable for model alias {alias!r}"
|
||||
)
|
||||
return None
|
||||
return token
|
||||
# entra_obo — per-user On-Behalf-Of.
|
||||
token = mcp.mint_model_obo_token_sync(user_id=user_id, audience=cfg.obo_audience)
|
||||
if not token:
|
||||
# A user IS driving but the mint yielded nothing (no captured
|
||||
# credential, decrypt failure, or the AS rejected the grant). Never
|
||||
# silent: when the operator explicitly configured a static key and
|
||||
# left fail-closed off, that key stands; a keyless alias raises below
|
||||
# and can never issue its SDK-construction placeholder.
|
||||
log.warning(
|
||||
"model_obo.fallback_to_static",
|
||||
alias=alias,
|
||||
user_id=user_id,
|
||||
audience=cfg.obo_audience,
|
||||
has_static_key=has_static_key,
|
||||
)
|
||||
if must_fail_closed:
|
||||
raise BackendAuthUnavailableError(
|
||||
f"Delegated backend authentication unavailable for model alias {alias!r}"
|
||||
)
|
||||
return None
|
||||
return token
|
||||
|
||||
def _history_scope_user_id(self) -> str | None:
|
||||
"""Identity that scopes conversation-history reads (recall tool,
|
||||
``/history``).
|
||||
@@ -8910,6 +9049,9 @@ class ChatSession:
|
||||
session_model_alias=self._model_alias or "",
|
||||
# For the temperature ladder's global rung (model.temperature).
|
||||
config_store=self._config_store,
|
||||
# The verdict belongs to this turn and carries the same acting
|
||||
# principal as the model/tool activity it evaluates.
|
||||
backend_auth_resolver=self._model_backend_auth_token,
|
||||
)
|
||||
except Exception:
|
||||
log.warning("judge.init_failed", exc_info=True)
|
||||
@@ -8947,6 +9089,7 @@ class ChatSession:
|
||||
session_model_alias=self._model_alias or "",
|
||||
# For the temperature ladder's global rung (model.temperature).
|
||||
config_store=self._config_store,
|
||||
backend_auth_resolver=self._model_backend_auth_token,
|
||||
)
|
||||
except Exception:
|
||||
log.warning("output_guard_judge.init_failed", exc_info=True)
|
||||
@@ -16763,6 +16906,8 @@ class ChatSession:
|
||||
capabilities=agent_caps,
|
||||
config_store=self._config_store,
|
||||
)
|
||||
# Resolve once per sub-agent run, outside its request retry loop.
|
||||
agent_backend_auth_token = self._model_backend_auth_token(lane.alias)
|
||||
|
||||
def _api_call(
|
||||
turns: list[Turn],
|
||||
@@ -16797,6 +16942,7 @@ 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,
|
||||
)
|
||||
# Sub-agent turns bypass on_status — record per-turn so
|
||||
# task-agent spend is visible in the dashboard, attributed
|
||||
|
||||
@@ -52,6 +52,30 @@ def _build_registry() -> dict[str, SettingDef]:
|
||||
"Change this at runtime to switch all new sessions to a different model "
|
||||
"without restarting.",
|
||||
),
|
||||
SettingDef(
|
||||
"model.auth_audience_allowlist",
|
||||
"str",
|
||||
"",
|
||||
"Comma- or newline-separated model gateway audience allow-list",
|
||||
"model",
|
||||
help="Exact Entra resource App ID URIs that model definitions may redeem. "
|
||||
"Required before an administrator can save entra_obo or entra_app auth. "
|
||||
"Use literal values such as api://<application-id>; no wildcard or host "
|
||||
"matching is performed.",
|
||||
),
|
||||
SettingDef(
|
||||
"model.auth_fail_closed",
|
||||
"bool",
|
||||
False,
|
||||
"Refuse dynamic-auth model calls when token minting fails",
|
||||
"model",
|
||||
help="When enabled, an entra_obo or entra_app model call is refused if its "
|
||||
"runtime credential cannot be minted. This also prevents routing that "
|
||||
"failure into the model fallback chain. A delegated call without a user, "
|
||||
"or any dynamic alias without a real static key, always refuses regardless "
|
||||
"of this setting. Leave disabled only to permit fallback to an explicitly "
|
||||
"configured static key after a mint failure.",
|
||||
),
|
||||
SettingDef(
|
||||
"model.temperature",
|
||||
"float",
|
||||
|
||||
@@ -4822,8 +4822,9 @@ class PostgreSQLBackend:
|
||||
token row — the freshness sweep's drive set + keepalive-refresh signal.
|
||||
|
||||
Unfiltered by expiry: an expired access token with a live refresh token
|
||||
is still a consented grant the sweep must keep hot. Only ``oauth_user``
|
||||
servers write these rows; no ciphertext is projected.
|
||||
is still a consented grant the sweep must keep hot. Joined to
|
||||
``mcp_servers`` so only ``oauth_user`` grants drive the sweep; synthetic
|
||||
model mint-cache rows are excluded. No ciphertext is projected.
|
||||
"""
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
@@ -4832,6 +4833,13 @@ class PostgreSQLBackend:
|
||||
mcp_user_tokens.c.server_name,
|
||||
sa.func.coalesce(mcp_user_tokens.c.last_refreshed, mcp_user_tokens.c.created),
|
||||
)
|
||||
.select_from(
|
||||
mcp_user_tokens.join(
|
||||
mcp_servers,
|
||||
mcp_servers.c.name == mcp_user_tokens.c.server_name,
|
||||
)
|
||||
)
|
||||
.where(mcp_servers.c.auth_type == "oauth_user")
|
||||
).fetchall()
|
||||
return [(row[0], row[1], row[2]) for row in rows]
|
||||
|
||||
@@ -5071,6 +5079,8 @@ class PostgreSQLBackend:
|
||||
reasoning_effort: str | None = None,
|
||||
surface_persisted_reasoning: bool = True,
|
||||
replay_reasoning_to_model: bool = False,
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
) -> None:
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
@@ -5093,6 +5103,8 @@ class PostgreSQLBackend:
|
||||
reasoning_effort=reasoning_effort,
|
||||
surface_persisted_reasoning=1 if surface_persisted_reasoning else 0,
|
||||
replay_reasoning_to_model=1 if replay_reasoning_to_model else 0,
|
||||
auth_mode=auth_mode,
|
||||
obo_audience=obo_audience,
|
||||
created_by=created_by,
|
||||
created=now,
|
||||
updated=now,
|
||||
|
||||
@@ -2257,9 +2257,11 @@ class StorageBackend(Protocol):
|
||||
an *idle* refresh token can't expire one between a user's real sessions.
|
||||
Deliberately UNFILTERED by ``expires_at``: an expired access token backed
|
||||
by a live refresh token is still a consented, reconcilable grant. Only
|
||||
``auth_type='oauth_user'`` servers ever write these rows, so a static /
|
||||
no-auth server is structurally absent; ciphertext columns are never
|
||||
touched — only the identity + timestamp cross the wire.
|
||||
The storage query joins ``mcp_servers`` and keeps only
|
||||
``auth_type='oauth_user'`` rows. Other auth paths also use the token
|
||||
table as a mint cache, but they are not refresh-grant sweep targets.
|
||||
Ciphertext columns are never touched — only identity + timestamp cross
|
||||
the wire.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -2417,6 +2419,8 @@ class StorageBackend(Protocol):
|
||||
reasoning_effort: str | None = None,
|
||||
surface_persisted_reasoning: bool = True,
|
||||
replay_reasoning_to_model: bool = False,
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
) -> None:
|
||||
"""Create a model definition. No-op if definition_id already exists."""
|
||||
...
|
||||
|
||||
@@ -840,6 +840,11 @@ model_definitions = sa.Table(
|
||||
sa.Column("reasoning_effort", sa.Text, nullable=True),
|
||||
sa.Column("surface_persisted_reasoning", sa.Integer, nullable=False, server_default="1"),
|
||||
sa.Column("replay_reasoning_to_model", sa.Integer, nullable=False, server_default="0"),
|
||||
# Backend-gateway credential mode. "static" sends api_key, "entra_obo"
|
||||
# mints a delegated-user token, and "entra_app" mints a shared app token
|
||||
# for ``obo_audience`` at call time (migration 068).
|
||||
sa.Column("auth_mode", sa.Text, nullable=False, server_default="static"),
|
||||
sa.Column("obo_audience", sa.Text, nullable=False, server_default=""),
|
||||
sa.Column("created_by", sa.Text, nullable=False, server_default=""),
|
||||
sa.Column("created", sa.Text, nullable=False),
|
||||
sa.Column("updated", sa.Text, nullable=False),
|
||||
|
||||
@@ -4976,8 +4976,9 @@ class SQLiteBackend:
|
||||
token row — the freshness sweep's drive set + keepalive-refresh signal.
|
||||
|
||||
Unfiltered by expiry: an expired access token with a live refresh token
|
||||
is still a consented grant the sweep must keep hot. Only ``oauth_user``
|
||||
servers write these rows; no ciphertext is projected.
|
||||
is still a consented grant the sweep must keep hot. Joined to
|
||||
``mcp_servers`` so only ``oauth_user`` grants drive the sweep; synthetic
|
||||
model mint-cache rows are excluded. No ciphertext is projected.
|
||||
"""
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
@@ -4986,6 +4987,13 @@ class SQLiteBackend:
|
||||
mcp_user_tokens.c.server_name,
|
||||
sa.func.coalesce(mcp_user_tokens.c.last_refreshed, mcp_user_tokens.c.created),
|
||||
)
|
||||
.select_from(
|
||||
mcp_user_tokens.join(
|
||||
mcp_servers,
|
||||
mcp_servers.c.name == mcp_user_tokens.c.server_name,
|
||||
)
|
||||
)
|
||||
.where(mcp_servers.c.auth_type == "oauth_user")
|
||||
).fetchall()
|
||||
return [(row[0], row[1], row[2]) for row in rows]
|
||||
|
||||
@@ -5225,6 +5233,8 @@ class SQLiteBackend:
|
||||
reasoning_effort: str | None = None,
|
||||
surface_persisted_reasoning: bool = True,
|
||||
replay_reasoning_to_model: bool = False,
|
||||
auth_mode: str = "static",
|
||||
obo_audience: str = "",
|
||||
) -> None:
|
||||
|
||||
now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
@@ -5246,6 +5256,8 @@ class SQLiteBackend:
|
||||
"reasoning_effort": reasoning_effort,
|
||||
"surface_persisted_reasoning": 1 if surface_persisted_reasoning else 0,
|
||||
"replay_reasoning_to_model": (1 if replay_reasoning_to_model else 0),
|
||||
"auth_mode": auth_mode,
|
||||
"obo_audience": obo_audience,
|
||||
"created_by": created_by,
|
||||
"created": now,
|
||||
"updated": now,
|
||||
|
||||
@@ -642,6 +642,8 @@ MODEL_DEFINITION_MUTABLE = frozenset(
|
||||
"reasoning_effort",
|
||||
"surface_persisted_reasoning",
|
||||
"replay_reasoning_to_model",
|
||||
"auth_mode",
|
||||
"obo_audience",
|
||||
}
|
||||
)
|
||||
PROJECT_MUTABLE = frozenset({"name", "visibility", "state", "parent_project_id"})
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Add dynamic backend-auth columns to model_definitions.
|
||||
|
||||
Lets a model backend authenticate with a caller-delegated Entra token
|
||||
(``entra_obo``), a shared app-identity token (``entra_app``), or the existing
|
||||
static ``api_key``. Dynamic modes mint for ``obo_audience`` at call time and
|
||||
bind that credential through the provider SDK. ``static`` remains the default,
|
||||
so existing rows are untouched.
|
||||
|
||||
Revision ID: 068
|
||||
Revises: 067
|
||||
Create Date: 2026-07-21
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "068"
|
||||
down_revision = "067"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"model_definitions",
|
||||
sa.Column("auth_mode", sa.Text, nullable=False, server_default="static"),
|
||||
)
|
||||
op.add_column(
|
||||
"model_definitions",
|
||||
sa.Column("obo_audience", sa.Text, nullable=False, server_default=""),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("model_definitions", "obo_audience")
|
||||
op.drop_column("model_definitions", "auth_mode")
|
||||
+120
-2
@@ -816,6 +816,30 @@ def _watch_fire_wake_fn(ws: Workstream) -> Callable[[], object]:
|
||||
return lambda: wake_workstream_if_pending(ws, trigger="watch-fire")
|
||||
|
||||
|
||||
def _watch_restore_owner(storage: Any, ws_id: str) -> str:
|
||||
"""Resolve the principal an unattended watch restore must execute as.
|
||||
|
||||
``None`` means the target workstream no longer exists and ``""`` means a
|
||||
legacy/unowned row. Neither may be turned into an anonymous, auto-approved
|
||||
model call. Storage errors deliberately propagate so the restore caller can
|
||||
retry them as transient failures.
|
||||
"""
|
||||
from turnstone.core.watch import WatchWorkstreamUnrestorable
|
||||
|
||||
owner = storage.get_workstream_owner(ws_id)
|
||||
if owner is None:
|
||||
log.warning("watch_restore: ws %s no longer exists", ws_id)
|
||||
raise WatchWorkstreamUnrestorable(ws_id)
|
||||
resolved = str(owner).strip()
|
||||
if not resolved:
|
||||
log.error(
|
||||
"watch_restore: refusing unattended restore for unowned ws %s",
|
||||
ws_id,
|
||||
)
|
||||
raise WatchWorkstreamUnrestorable(ws_id)
|
||||
return resolved
|
||||
|
||||
|
||||
def _interactive_open_post_load(request: Request, ws: Workstream) -> None:
|
||||
"""Post-load hook for the lifted interactive ``open`` body.
|
||||
|
||||
@@ -4166,6 +4190,28 @@ def internal_model_reload(request: Request) -> JSONResponse:
|
||||
finally:
|
||||
new_registry.shutdown()
|
||||
|
||||
# A model may switch from static to dynamic auth while the node has no MCP
|
||||
# servers. Ensure the dedicated mint loop exists after the registry reload;
|
||||
# model auth must not depend on an unrelated MCP catalog being configured.
|
||||
if registry.has_dynamic_auth() and getattr(request.app.state, "mcp_client", None) is None:
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
|
||||
mcp_mgr = MCPClientManager({})
|
||||
mcp_mgr.start()
|
||||
mcp_mgr.set_storage(storage)
|
||||
mcp_mgr.set_app_state(request.app.state)
|
||||
request.app.state.mcp_client = mcp_mgr
|
||||
mcp_ref = getattr(request.app.state, "mcp_ref", None)
|
||||
if mcp_ref is not None:
|
||||
mcp_ref[0] = mcp_mgr
|
||||
workstreams = getattr(request.app.state, "workstreams", None)
|
||||
if workstreams is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
for ws in workstreams.list_all():
|
||||
session = getattr(ws, "session", None)
|
||||
if session is not None:
|
||||
session.set_model_mint_client(mcp_mgr)
|
||||
|
||||
# Ensure health trackers exist for any newly-added backends
|
||||
health_reg = getattr(request.app.state, "health_registry", None)
|
||||
if health_reg:
|
||||
@@ -4205,10 +4251,50 @@ def internal_model_status(request: Request) -> JSONResponse:
|
||||
"temperature": cfg.temperature,
|
||||
"max_tokens": cfg.max_tokens,
|
||||
"reasoning_effort": cfg.reasoning_effort,
|
||||
"auth_mode": cfg.auth_mode,
|
||||
"obo_audience": cfg.obo_audience,
|
||||
}
|
||||
return JSONResponse({"models": models})
|
||||
|
||||
|
||||
async def internal_model_auth_cache_invalidate(request: Request) -> JSONResponse:
|
||||
"""Evict one user's delegated model-token memo from this node.
|
||||
|
||||
Identity unlink deletes the authoritative cache rows in shared storage, then
|
||||
the console fans this service-only request to every node. Without the
|
||||
in-process eviction, a warm node could keep using the deleted bearer until
|
||||
its access-token expiry.
|
||||
"""
|
||||
auth_result = getattr(getattr(request, "state", None), "auth_result", None)
|
||||
if auth_result is None or not auth_result.has_scope("service"):
|
||||
return JSONResponse({"error": "Service scope required"}, status_code=403)
|
||||
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return JSONResponse({"error": "Invalid JSON body"}, status_code=400)
|
||||
user_id = body.get("user_id") if isinstance(body, dict) else None
|
||||
if (
|
||||
not isinstance(user_id, str)
|
||||
or not user_id
|
||||
or len(user_id) > 512
|
||||
or any(ord(ch) < 32 or ord(ch) == 127 for ch in user_id)
|
||||
):
|
||||
return JSONResponse({"error": "Invalid user_id"}, status_code=400)
|
||||
|
||||
mcp_mgr = getattr(request.app.state, "mcp_client", None)
|
||||
if mcp_mgr is None:
|
||||
return JSONResponse({"status": "noop", "evicted": 0})
|
||||
|
||||
from turnstone.core.mcp_oauth import MODEL_OBO_CACHE_PREFIX
|
||||
|
||||
evicted = mcp_mgr.invalidate_model_mint_memo_sync(
|
||||
user_id=user_id,
|
||||
server_prefix=MODEL_OBO_CACHE_PREFIX,
|
||||
)
|
||||
return JSONResponse({"status": "ok", "evicted": evicted})
|
||||
|
||||
|
||||
def _collect_node_models_metadata(app_state: Any) -> tuple[str, str, str] | None:
|
||||
"""Build the ``("models", json_value, "auto")`` node_metadata entry.
|
||||
|
||||
@@ -5131,6 +5217,11 @@ def create_app(
|
||||
methods=["POST"],
|
||||
),
|
||||
Route("/api/_internal/model-status", internal_model_status),
|
||||
Route(
|
||||
"/api/_internal/model-auth-cache-invalidate",
|
||||
internal_model_auth_cache_invalidate,
|
||||
methods=["POST"],
|
||||
),
|
||||
],
|
||||
),
|
||||
Route("/health", health),
|
||||
@@ -5431,6 +5522,7 @@ def main() -> None:
|
||||
mcp_client = create_mcp_client(
|
||||
mcp_config_cli or config_store.get("mcp.config_path") or None,
|
||||
storage=_get_storage(),
|
||||
required=registry.has_dynamic_auth(),
|
||||
)
|
||||
# Mutable ref so session_factory always sees the latest MCP client,
|
||||
# including ones created by internal_mcp_reload after startup.
|
||||
@@ -5731,7 +5823,21 @@ def main() -> None:
|
||||
raise WatchWorkstreamUnrestorable(ws_id) from exc
|
||||
|
||||
try:
|
||||
ws = manager.create(user_id="", name="watch-restore", **persona_kwargs)
|
||||
owner_id = _watch_restore_owner(_get_storage(), ws_id)
|
||||
except WatchWorkstreamUnrestorable:
|
||||
raise
|
||||
except Exception:
|
||||
# An owner read is authoritative for the execution principal. A
|
||||
# storage blip is retryable; anonymous execution is not.
|
||||
log.warning(
|
||||
"watch_restore: owner lookup failed for ws %s (treating as transient)",
|
||||
ws_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
ws = manager.create(user_id=owner_id, name="watch-restore", **persona_kwargs)
|
||||
except RuntimeError:
|
||||
# TRANSIENT: all restore slots active right now. Return None so
|
||||
# the runner holds the reminder and retries on a later tick.
|
||||
@@ -5828,7 +5934,19 @@ def main() -> None:
|
||||
if not target_id:
|
||||
log.error("Workstream not found: %s", args.resume)
|
||||
sys.exit(1)
|
||||
ws = manager.create(user_id="", name="resumed", **_resume_persona_kwargs(target_id))
|
||||
try:
|
||||
resume_owner = _watch_restore_owner(_get_storage(), target_id)
|
||||
except WatchWorkstreamUnrestorable:
|
||||
log.error("Cannot resume unowned workstream: %s", target_id)
|
||||
sys.exit(1)
|
||||
except Exception:
|
||||
log.exception("Cannot resolve owner for workstream: %s", target_id)
|
||||
sys.exit(1)
|
||||
ws = manager.create(
|
||||
user_id=resume_owner,
|
||||
name="resumed",
|
||||
**_resume_persona_kwargs(target_id),
|
||||
)
|
||||
if not isinstance(ws.ui, WebUI):
|
||||
raise TypeError(f"Expected WebUI, got {type(ws.ui).__name__}")
|
||||
if args.skip_permissions or config_store.get("tools.skip_permissions"):
|
||||
|
||||
Reference in New Issue
Block a user