mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-13 07:22:24 -06:00
Compare commits
132 Commits
stable/1.7
...
v1.5.8
| Author | SHA1 | Date | |
|---|---|---|---|
| 8ce4c8e737 | |||
| 95e67dc768 | |||
| dc35cbc7bf | |||
| 79f4d0030d | |||
| 14e504db1f | |||
| 0abe0cb77d | |||
| 610513398b | |||
| 2d6519f9a8 | |||
| a99ce49311 | |||
| 46f3571c93 | |||
| 53cabe7e20 | |||
| 1d9fd94e23 | |||
| eb89ddab1e | |||
| eb92e61755 | |||
| 0e2ea122eb | |||
| eb9dd2402a | |||
| 0c58910c4b | |||
| 2e393d76b4 | |||
| dec175f176 | |||
| 592433b46d | |||
| da5321eb88 | |||
| baa2214f96 | |||
| 3b60a69e4f | |||
| 5d1213d3dc | |||
| 9ae2b376c7 | |||
| b368bdeecc | |||
| d767aca784 | |||
| 12bc580dee | |||
| aa1446364b | |||
| 21507dc02e | |||
| a0ed4b9897 | |||
| ab8ee0d759 | |||
| d34f6cd0b1 | |||
| b219c47ba8 | |||
| d770a811a8 | |||
| 912e9c57b0 | |||
| d7c6053441 | |||
| e7a17a20b0 | |||
| 931a1eca9d | |||
| 048285a423 | |||
| 481347eb17 | |||
| 31a554a4bd | |||
| af2c0ae13a | |||
| 0808dc0af0 | |||
| cfc8a6c8c0 | |||
| 266e3536aa | |||
| 42bf9aecaf | |||
| 191775dd7e | |||
| c41fd2be2e | |||
| 1787fb5c11 | |||
| 814c42763d | |||
| 242596ced3 | |||
| bde0913442 | |||
| 570b198f1b | |||
| 96d935f1f7 | |||
| 1a1043c4df | |||
| 55aab54774 | |||
| 0f8c8b38a3 | |||
| b0f7029ff1 | |||
| a4c335d7bf | |||
| 21663d1567 | |||
| c823156af5 | |||
| bace928477 | |||
| d16c911750 | |||
| b8fadad94f | |||
| 63aecdf2fa | |||
| cbe8940b30 | |||
| 1dcd1e2ec4 | |||
| 32e29ff255 | |||
| 3e2fe0bc9d | |||
| 5f5eee4aab | |||
| 6d532ed776 | |||
| c3d9cdae82 | |||
| 366d316941 | |||
| 3e87f4262e | |||
| f50b559792 | |||
| c6b3c0bc5f | |||
| cefb74a226 | |||
| 273d547f4e | |||
| bbb404c363 | |||
| 072113f7ca | |||
| 7f1b0acf7a | |||
| 6904bd8f39 | |||
| 4d677d1ebf | |||
| 35f462a46d | |||
| ec74334e74 | |||
| 2fd0c29a92 | |||
| 1207d27363 | |||
| 733c9818d4 | |||
| 28a2779c10 | |||
| 56364b0b5b | |||
| afb5804a7c | |||
| 4b508a1319 | |||
| 9c2cb185e1 | |||
| 0519b847bd | |||
| bbc8b99a9f | |||
| bd9f780b21 | |||
| 5d14b5f675 | |||
| 4693fa95f1 | |||
| c3423d6606 | |||
| 53f1222c22 | |||
| 802d87a57f | |||
| 8349d9994d | |||
| ac1fd67137 | |||
| 1b40ae79f9 | |||
| b078ddccf0 | |||
| 4b6c93a0e9 | |||
| 9d283e951f | |||
| 4e407e7d4f | |||
| 7ab24e500b | |||
| 5bcbcb73b9 | |||
| af6749421a | |||
| 4d6cb77075 | |||
| 5c225ef39b | |||
| f5a843f44a | |||
| cf44841624 | |||
| 7ffab6a272 | |||
| ba3bc9d989 | |||
| dbe023b4dd | |||
| 3d3a8b7367 | |||
| 99eff73a97 | |||
| 99fcd30299 | |||
| 423c2e80b7 | |||
| a0eb77360d | |||
| 3dd0e196fe | |||
| 19c3db5329 | |||
| 0bea72019e | |||
| b9ff52d582 | |||
| 961f999c93 | |||
| 8bdb916064 | |||
| 25fe4e728a | |||
| 0d1a32ff65 |
@@ -47,9 +47,9 @@ jobs:
|
||||
# explicit setup, that suite silently skips if the runner
|
||||
# image happens not to ship Node, masking regressions in
|
||||
# the browser-side renderer.
|
||||
- uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5
|
||||
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6
|
||||
with:
|
||||
node-version: "20"
|
||||
node-version: "24"
|
||||
- run: pip install -e ".[test]"
|
||||
- run: pytest tests/ -m "not live" --cov=turnstone --cov-report=term-missing --cov-report=xml -q
|
||||
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
|
||||
@@ -79,9 +79,9 @@ jobs:
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
with:
|
||||
python-version: "3.14"
|
||||
- uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5
|
||||
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6
|
||||
with:
|
||||
node-version: "20"
|
||||
node-version: "24"
|
||||
- run: pip install -e ".[test,postgres]"
|
||||
- run: pytest tests/ -m "not live" --storage-backend=postgresql -q
|
||||
env:
|
||||
|
||||
+6
-3
@@ -8,14 +8,17 @@ FROM python:3.14-slim
|
||||
LABEL org.opencontainers.image.title="turnstone" \
|
||||
org.opencontainers.image.description="Multi-node AI orchestration platform"
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.11.7 /uv /usr/local/bin/uv
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.11.8 /uv /usr/local/bin/uv
|
||||
|
||||
# Remove the slim image's man page exclusion so man-db has actual content
|
||||
RUN rm -f /etc/dpkg/dpkg.cfg.d/docker
|
||||
|
||||
# System dependencies: psycopg (libpq5), developer tooling for agent workflows
|
||||
# System dependencies: psycopg (libpq5), developer tooling for agent workflows.
|
||||
# ripgrep is the preferred backend for the search tool — natively bounds
|
||||
# per-line, per-file, and per-filesize so pathological inputs (minified
|
||||
# bundles, training-data JSONL with multi-MB single records) can't OOM us.
|
||||
RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-recommends \
|
||||
libpq5 git curl jq man-db manpages procps file \
|
||||
libpq5 git curl jq man-db manpages procps file ripgrep \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Node.js LTS (for npx-based MCP servers like @modelcontextprotocol/server-github)
|
||||
|
||||
@@ -546,11 +546,9 @@ adds, removes, or reconnects servers as needed.
|
||||
6. `_exec_mcp_tool()` calls `call_tool_sync()` which dispatches to the async loop
|
||||
via `asyncio.run_coroutine_threadsafe()`
|
||||
|
||||
**Tool refresh:** Three mechanisms keep tools up-to-date without restart:
|
||||
**Tool refresh:** Two mechanisms keep tools up-to-date without restart:
|
||||
- **Push:** Servers declaring `tools.listChanged` send `ToolListChangedNotification`;
|
||||
the registered `message_handler` triggers immediate single-server refresh.
|
||||
- **Periodic:** Servers without push support are polled on a staggered interval
|
||||
(default 4 h, configurable via `[mcp] refresh_interval` or `--mcp-refresh-interval`).
|
||||
- **Manual:** `/mcp refresh [server]` calls `refresh_sync()` for on-demand refresh
|
||||
(also attempts reconnection for disconnected servers).
|
||||
|
||||
@@ -572,10 +570,11 @@ from a healthy connection do not trip the breaker. When the cooldown expires
|
||||
(`call_tool_sync`, `read_resource_sync`, `get_prompt_sync`, `refresh_sync`)
|
||||
cancel orphaned futures on timeout to prevent coroutine accumulation on the
|
||||
background event loop. Push notification refreshes are debounced (5 s per
|
||||
server) to protect against notification storms. The periodic refresh loop
|
||||
attempts reconnection for disconnected servers with exponential backoff
|
||||
(60 s–1 h). Transport stream references are pre-closed before stack teardown to
|
||||
work around the MCP SDK's anyio cancel-scope CPU busy-loop (SDK #2147).
|
||||
server) to protect against notification storms. Operators can force a
|
||||
catalog refresh or full reconnect from the admin panel; reconnects clear
|
||||
the circuit breaker and run a fresh handshake. Transport stream references
|
||||
are pre-closed before stack teardown to work around the MCP SDK's anyio
|
||||
cancel-scope CPU busy-loop (SDK #2147).
|
||||
|
||||
**Error isolation:** Per-server connection/refresh failures are caught and logged; other
|
||||
servers are unaffected. Tool execution errors return error strings to the LLM
|
||||
|
||||
@@ -40,7 +40,7 @@ package "turnstone/core/" <<Rectangle>> {
|
||||
component [auth.py\nAuthentication] as auth <<core>>
|
||||
component [healthcheck.py\nBackendHealthMonitor] as healthcheck <<core>>
|
||||
component [ratelimit.py\nRateLimiter] as ratelimit <<core>>
|
||||
component [mcp_client.py\nMCPClientManager\n(push + periodic refresh)] as mcp <<core>>
|
||||
component [mcp_client.py\nMCPClientManager\n(push + manual refresh)] as mcp <<core>>
|
||||
component [tool_search.py\nToolSearchManager, BM25] as toolsearch <<core>>
|
||||
component [model_registry.py\nModelRegistry] as registry <<core>>
|
||||
}
|
||||
|
||||
@@ -253,7 +253,7 @@ class "MCPClientManager" as MCPMgr {
|
||||
Background asyncio event loop
|
||||
bridges async MCP SDK to
|
||||
sync ChatSession dispatch.
|
||||
Push + periodic + manual refresh.
|
||||
Push + manual refresh.
|
||||
Resources + prompts discovered
|
||||
alongside tools at startup.
|
||||
--
|
||||
|
||||
@@ -190,21 +190,25 @@ group Push Notifications (debounced 5s per server)
|
||||
MCPMgr -> Storage : sync_prompts_to_storage()
|
||||
end
|
||||
|
||||
group Periodic Polling (default 4h)
|
||||
MCPMgr -> MCPMgr : _periodic_refresh()
|
||||
group Manual Refresh
|
||||
Session -> MCPMgr : refresh_sync()
|
||||
note right
|
||||
Only polls capabilities
|
||||
without push support.
|
||||
Staggered per-server.
|
||||
Disconnected servers get
|
||||
reconnect attempts with
|
||||
exponential backoff (60s-1h).
|
||||
/mcp refresh [server] —
|
||||
re-fetches catalog and
|
||||
attempts reconnect for
|
||||
disconnected servers.
|
||||
end note
|
||||
end
|
||||
|
||||
group Manual Refresh
|
||||
Session -> MCPMgr : refresh_sync()
|
||||
note right: /mcp refresh [server]
|
||||
group Manual Reconnect
|
||||
Session -> MCPMgr : reconnect_sync(name)
|
||||
note right
|
||||
Operator-driven via the
|
||||
console admin panel —
|
||||
tears down session, clears
|
||||
circuit breaker, runs a
|
||||
fresh handshake.
|
||||
end note
|
||||
end
|
||||
|
||||
== Policy Evaluation ==
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a3b5c59403a6febd81667fc8fd2a7d22bc59da6130eba0dea5449c42668d0ede
|
||||
size 387044
|
||||
oid sha256:95dd5ebc899a1261d516686a5aa3319a7f45015d411302825fa28afbfc82e1ce
|
||||
size 326766
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:474b900448ec04d1117b48a2b55614524721b2f04ac4bda66170bd0a06aae0f2
|
||||
size 624573
|
||||
oid sha256:25b5448bbb7da8ddafe4f65c6c5e6cbcaa9cb9f31746ca46d3a2241bc47b1956
|
||||
size 259687
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7623df33be9baf7647ca1c2450640df57e1cd73e8be1f8168aae16e546ad683c
|
||||
size 459941
|
||||
oid sha256:d6aff446a062aa08f316985d00c2183148694f786d7f22172bc50b30046c728b
|
||||
size 379259
|
||||
|
||||
+82
-13
@@ -39,18 +39,19 @@ are set.
|
||||
| `TURNSTONE_OIDC_ROLE_CLAIM` | No | — | ID token claim containing role/group values (see [Role Mapping](#role-mapping)) |
|
||||
| `TURNSTONE_OIDC_ROLE_MAP` | No | — | Mapping from claim values to Turnstone role IDs (see [Role Mapping](#role-mapping)) |
|
||||
| `TURNSTONE_OIDC_PASSWORD_ENABLED` | No | `true` | Set to `false` to hide the password form and block all username/password logins (including admin). API tokens continue to work. |
|
||||
| `TURNSTONE_OIDC_REDIRECT_BASE` | No | — | Externally-reachable origin for the OIDC redirect URI (e.g. `https://app.example.com`). Recommended when running behind a reverse proxy. When unset, derived from the request Host header. |
|
||||
| `TURNSTONE_OIDC_REDIRECT_BASE` | Yes | — | Externally-reachable origin for the OIDC redirect URI (e.g. `https://app.example.com`). Without this, OIDC will refuse to start. The previous Host-header fallback was unsafe under permissive reverse proxies. |
|
||||
| `TURNSTONE_OIDC_TRUSTED_ENDPOINT_HOSTS` | No | — | Comma-separated list of additional hostnames whose endpoints the IdP discovery document is allowed to reference. See [Cross-host endpoints](#cross-host-endpoints). |
|
||||
|
||||
OIDC is enabled when all three required fields (issuer, client ID, client
|
||||
secret) are non-empty. If any is missing, OIDC is silently disabled and
|
||||
the login screen shows only the password form.
|
||||
All four required fields — issuer, client ID, client secret, and
|
||||
`TURNSTONE_OIDC_REDIRECT_BASE` — must be set. If any are missing OIDC
|
||||
is disabled at startup (an error is logged when only `redirect_base`
|
||||
is missing) and the login screen shows only the password form.
|
||||
|
||||
### Reverse Proxy / Load Balancer
|
||||
### Redirect base (required)
|
||||
|
||||
When Turnstone runs behind a reverse proxy, the internal `Host` header may
|
||||
not match the externally-reachable URL. Set `TURNSTONE_OIDC_REDIRECT_BASE`
|
||||
to the public origin so the redirect URI sent to the identity provider is
|
||||
correct:
|
||||
`TURNSTONE_OIDC_REDIRECT_BASE` pins the redirect URI sent to the identity
|
||||
provider to a known externally-visible origin. Set it to the public origin
|
||||
of your Turnstone deployment:
|
||||
|
||||
```bash
|
||||
TURNSTONE_OIDC_REDIRECT_BASE=https://app.example.com
|
||||
@@ -60,6 +61,44 @@ The resulting callback URL will be
|
||||
`https://app.example.com/v1/api/auth/oidc/callback` — register this as the
|
||||
authorized redirect URI in your identity provider.
|
||||
|
||||
OIDC will refuse to start when this variable is unset. There is no
|
||||
Host-header fallback: a permissive reverse proxy or direct backend access
|
||||
would otherwise let an attacker spoof `Host` and steer the IdP redirect
|
||||
to a callback origin they control.
|
||||
|
||||
### Cross-host endpoints
|
||||
|
||||
By default, every endpoint in the IdP discovery document
|
||||
(`token_endpoint`, `jwks_uri`, `userinfo_endpoint`) must share the
|
||||
issuer's `(scheme, host, port)`. This prevents a hostile or compromised
|
||||
IdP from redirecting the token-exchange POST (which carries
|
||||
`client_secret`) to an arbitrary host, and prevents JWKS fetches from
|
||||
being aimed at internal services.
|
||||
|
||||
A few public IdPs legitimately split endpoints across hostnames. Google
|
||||
is the canonical example:
|
||||
|
||||
| Field | Hostname |
|
||||
|-------|----------|
|
||||
| issuer | `accounts.google.com` |
|
||||
| token_endpoint | `oauth2.googleapis.com` |
|
||||
| jwks_uri | `www.googleapis.com` |
|
||||
| userinfo_endpoint | `openidconnect.googleapis.com` |
|
||||
|
||||
Google's set is built in — operators using `https://accounts.google.com`
|
||||
need no extra configuration.
|
||||
|
||||
For other IdPs whose discovery document references a non-issuer host,
|
||||
extend the allow-list explicitly:
|
||||
|
||||
```bash
|
||||
TURNSTONE_OIDC_TRUSTED_ENDPOINT_HOSTS=token.example.com,keys.example.com
|
||||
```
|
||||
|
||||
The same scheme / no-userinfo / SSRF rules apply to allow-listed hosts —
|
||||
this knob only relaxes the same-origin check, not the security gates.
|
||||
Each entry is a hostname (no scheme, no path).
|
||||
|
||||
### config.toml alternative
|
||||
|
||||
```toml
|
||||
@@ -198,6 +237,19 @@ TURNSTONE_OIDC_ROLE_MAP="admin:builtin-admin,engineering:builtin-operator,viewer
|
||||
the user authenticates via OIDC, so new group memberships are picked
|
||||
up on the next login.
|
||||
|
||||
### `assigned_by` markers
|
||||
|
||||
Role assignments record an `assigned_by` value that controls how the
|
||||
sync logic treats them. OIDC-driven flows use two distinct markers:
|
||||
|
||||
- `oidc` — set by claim-driven role mapping; revoked automatically on
|
||||
the next login when the corresponding claim value is no longer
|
||||
present.
|
||||
- `oidc-default` — applied to brand-new OIDC users who have no
|
||||
claim-mapped roles, as a safety net so they still get
|
||||
`builtin-viewer` access on first login. Survives subsequent logins
|
||||
regardless of claim contents and is never revoked by `apply_role_mapping`.
|
||||
|
||||
### Built-in Roles
|
||||
|
||||
| Role ID | Permissions |
|
||||
@@ -375,10 +427,27 @@ callback validation. Entries are automatically cleaned up after 5 minutes.
|
||||
|
||||
### "OIDC not configured"
|
||||
|
||||
All three required environment variables must be set:
|
||||
`TURNSTONE_OIDC_ISSUER`, `TURNSTONE_OIDC_CLIENT_ID`, and
|
||||
`TURNSTONE_OIDC_CLIENT_SECRET`. Check that none are empty or
|
||||
whitespace-only.
|
||||
All four required environment variables must be set:
|
||||
`TURNSTONE_OIDC_ISSUER`, `TURNSTONE_OIDC_CLIENT_ID`,
|
||||
`TURNSTONE_OIDC_CLIENT_SECRET`, and `TURNSTONE_OIDC_REDIRECT_BASE`.
|
||||
Check that none are empty or whitespace-only.
|
||||
|
||||
### "OIDC enabled but TURNSTONE_OIDC_REDIRECT_BASE is unset"
|
||||
|
||||
This error is logged when the three credential variables are set but
|
||||
`TURNSTONE_OIDC_REDIRECT_BASE` is missing. OIDC is disabled at startup
|
||||
to prevent Host-header-derived redirect URI spoofing. Set the variable
|
||||
to your service's externally-visible origin (e.g.
|
||||
`https://app.example.com`) and restart the server. See
|
||||
[Redirect base](#redirect-base-required) for the rationale.
|
||||
|
||||
### Discovery silently disables OIDC with "host does not match issuer"
|
||||
|
||||
The IdP discovery document points `token_endpoint`, `jwks_uri`, or
|
||||
`userinfo_endpoint` at a hostname that doesn't share the issuer's
|
||||
origin. If the IdP is legitimate, add the additional hostname(s) to
|
||||
`TURNSTONE_OIDC_TRUSTED_ENDPOINT_HOSTS`. Google is allow-listed
|
||||
automatically; see [Cross-host endpoints](#cross-host-endpoints).
|
||||
|
||||
### "Login session expired"
|
||||
|
||||
|
||||
+1
-1
@@ -100,7 +100,7 @@ initialization:
|
||||
| `tools` | timeout, truncation, agent_max_turns, skip_permissions, search, search_threshold, search_max_results |
|
||||
| `server` | workstream_idle_timeout, max_workstreams |
|
||||
| `cluster` | node_fan_out_limit, mcp_max_servers |
|
||||
| `mcp` | config_path, refresh_interval, registry_url |
|
||||
| `mcp` | config_path, registry_url |
|
||||
| `ratelimit` | enabled, requests_per_second, burst, trusted_proxies |
|
||||
| `health` | backend_probe_interval, backend_probe_timeout, circuit_breaker_threshold, circuit_breaker_cooldown |
|
||||
| `judge` | enabled, model, provider, base_url, api_key, confidence_threshold, max_context_ratio, timeout, read_only_tools, output_guard, redact_secrets, cancel_on_approval |
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
---
|
||||
name: import-conversation-history
|
||||
description: Use this skill when the user wants to import or migrate conversation history from another LLM chat or coding tool (e.g. ChatGPT, Claude.ai, Cursor, Copilot Chat, Aider, Gemini, a custom JSON export) into Turnstone. The skill teaches Turnstone's destination contracts — workstream identity, the OpenAI-shaped message rows, tool-call/result pairing, provider-fidelity blobs, attachments, and archive-vs-resumable choice — so the agent can map any source format onto them. Trigger phrases: "import my chats", "migrate this transcript into Turnstone", "bring my Claude.ai history over", "load this export as a workstream".
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# Importing Conversation History into Turnstone
|
||||
|
||||
## Overview
|
||||
|
||||
Source formats vary; the destination does not. Your job is to translate whatever the user hands you (JSON dump, ZIP export, scraped HTML, screenshot OCR, raw transcript) into Turnstone's internal shape: **one workstream row** plus an ordered sequence of **conversation rows** in OpenAI message format. This skill documents the destination so you can write a correct mapper for any source.
|
||||
|
||||
Two questions to settle with the user before writing anything:
|
||||
|
||||
1. **Archive or resumable?** An archive ("saved" workstream — `state="closed"`) is read-only history. A resumable workstream (`state="idle"`) lets the user continue the conversation; this only works cleanly when the source LLM matches a Turnstone-supported provider/model and tool definitions still resolve.
|
||||
2. **One workstream per source thread, or merge?** Default to one-to-one unless the user explicitly asks to merge.
|
||||
|
||||
Default to **archive** when in doubt — resuming a foreign transcript with mismatched tool schemas or stale provider signatures will fail at the next turn.
|
||||
|
||||
## Turnstone Data Model (the destination)
|
||||
|
||||
Two tables carry the conversation:
|
||||
|
||||
### `workstreams` (one row per imported thread)
|
||||
|
||||
| Column | Required | Notes |
|
||||
|---|---|---|
|
||||
| `ws_id` | yes | 32-char lowercase hex. Auto-generate with `secrets.token_hex(16)` if you don't already have one. **First 4 hex chars are the routing bucket** — see "Identity & Routing" below. |
|
||||
| `name` | yes | Short title. Pull from source thread title; fall back to first ~60 chars of first user message. |
|
||||
| `state` | yes | `"closed"` for archive, `"idle"` for resumable. Never set `"running"` on import. |
|
||||
| `kind` | yes | `"interactive"` for normal threads. Do NOT use `"coordinator"` for imports — that's reserved for cluster-spawned coordinator workstreams. |
|
||||
| `parent_ws_id` | no | Leave NULL. Only set if you're importing a coordinator-spawned subtree and re-parenting it; rare. |
|
||||
| `user_id` | yes | Owner. Must exist in `users`; importer must know which Turnstone user owns the imported history. |
|
||||
| `node_id` | yes (multi-node) | Denormalized cache of the node that owns this `ws_id`'s bucket. Single-node deployments can leave it NULL or set it to the only node. |
|
||||
| `alias` | no | Human-typeable short name. Optional; must be unique cluster-wide if set. |
|
||||
| `title` | no | Auto-titled later by the LLM; safe to leave NULL on import. |
|
||||
| `skill_id`, `skill_version` | yes | Default `""` and `0` unless the source thread was scoped to a Turnstone skill. |
|
||||
| `created`, `updated` | yes | ISO8601 strings. Use the source's first/last message timestamps when available. |
|
||||
|
||||
### `conversations` (many rows per thread, ordered by `id`/`timestamp`)
|
||||
|
||||
| Column | Notes |
|
||||
|---|---|
|
||||
| `ws_id` | The workstream this row belongs to. |
|
||||
| `timestamp` | ISO8601 string. Preserve source timestamps; fall back to monotonically increasing values if unknown. **Order is canonical via `id` (autoincrement), not `timestamp`** — but always insert in conversational order so both agree. |
|
||||
| `role` | One of `system`, `user`, `assistant`, `tool`, `developer`. See role mapping below. |
|
||||
| `content` | Text. May be NULL for assistant rows that are *only* tool calls. |
|
||||
| `tool_name` | Set on `role="tool"` rows (the tool whose result this is). NULL otherwise. |
|
||||
| `tool_call_id` | Set on `role="tool"` rows (matches the assistant row's `tool_calls[].id`). NULL otherwise. |
|
||||
| `tool_calls` | JSON-encoded list, on `role="assistant"` rows that issued tool calls. OpenAI shape — see "Tool Calls" below. |
|
||||
| `provider_data` | JSON blob preserving provider-native content blocks (Anthropic `signature`, Gemini `thought_signature`, etc.). Optional; only matters for **resumable** imports against the same provider. Skip for archives. |
|
||||
|
||||
The internal format is **OpenAI-shaped**, even when the source was Anthropic or Gemini. Providers translate at their own API boundary; storage stays uniform.
|
||||
|
||||
## Identity & Routing (`ws_id`)
|
||||
|
||||
- `ws_id` is **32-char lowercase hex** (i.e. `secrets.token_hex(16)`).
|
||||
- The **routing bucket** is `int(ws_id[:4], 16)` — the first 4 hex chars place this workstream on a specific node via the consistent hash ring.
|
||||
- For multi-node imports: either insert through the console's routing proxy (which forwards to the owning node), or generate `ws_id`s and write directly to each node's database in batches grouped by bucket.
|
||||
- For single-node imports: bucket math is irrelevant; any `ws_id` works.
|
||||
- **Do not reuse the source platform's IDs as `ws_id`** unless they happen to be 32-char hex. Generate fresh; if you need the old ID for traceability, store it in `workstream_config` under a key like `import.source_id`.
|
||||
|
||||
## Recommended Import Path
|
||||
|
||||
Three options, in order of preference:
|
||||
|
||||
### 1. Storage protocol (recommended for full history)
|
||||
|
||||
Use `turnstone.core.storage.Storage.save_messages_bulk(rows)`. This is the canonical bulk-insert primitive and bypasses the LLM round-trip entirely.
|
||||
|
||||
```python
|
||||
from turnstone.core.storage import get_storage # construct via the same path the server uses
|
||||
|
||||
storage = get_storage(...) # see turnstone.core.storage.__init__ for the project's wiring
|
||||
|
||||
storage.create_workstream( # or whatever the project's exposed creator is — check turnstone/core/storage/_protocol.py
|
||||
ws_id=ws_id,
|
||||
user_id=user_id,
|
||||
name=name,
|
||||
state="closed",
|
||||
kind="interactive",
|
||||
...
|
||||
)
|
||||
|
||||
storage.save_messages_bulk([
|
||||
{"ws_id": ws_id, "role": "user", "content": "Hello"},
|
||||
{"ws_id": ws_id, "role": "assistant", "content": "Hi! What can I help with?"},
|
||||
{"ws_id": ws_id, "role": "assistant", "content": None,
|
||||
"tool_calls": json.dumps([{"id": "call_1", "type": "function",
|
||||
"function": {"name": "search", "arguments": "{\"q\":\"x\"}"}}])},
|
||||
{"ws_id": ws_id, "role": "tool", "tool_name": "search", "tool_call_id": "call_1",
|
||||
"content": "result text"},
|
||||
# ...
|
||||
])
|
||||
```
|
||||
|
||||
`save_messages_bulk` handles `timestamp` and the workstream's `updated` column internally, so you don't need to compute them per row. **Verify the exact creator signature** by reading `turnstone/core/storage/_protocol.py` — table layout has shifted across migrations and the Storage protocol is the source of truth.
|
||||
|
||||
### 2. SDK `create_workstream(resume_ws=...)` (when the source is already a Turnstone workstream)
|
||||
|
||||
Only useful for *Turnstone → Turnstone* re-parenting. Not relevant for foreign sources.
|
||||
|
||||
### 3. SDK `create_workstream(initial_message=...)` + `send()` per turn (last resort)
|
||||
|
||||
Only fits archives where the source had **no tool calls** and you don't care about preserving assistant turns verbatim. Each `send()` triggers a real LLM round-trip, which is expensive and rewrites assistant content. Don't use this for full history.
|
||||
|
||||
## Role Mapping
|
||||
|
||||
Common source-role conventions and how they map to Turnstone:
|
||||
|
||||
| Source role | Turnstone `role` | Notes |
|
||||
|---|---|---|
|
||||
| `user`, `human` | `user` | Direct map. |
|
||||
| `assistant`, `ai`, `model`, `bot` | `assistant` | Direct map. |
|
||||
| `system` | `system` | Preserve only if it's content the user wrote (custom instructions). Drop boilerplate provider preambles — Turnstone composes its own system message. |
|
||||
| `developer` (OpenAI o-series) | `developer` | Preserve. |
|
||||
| `tool`, `function`, `tool_result` | `tool` | Must carry `tool_name` and `tool_call_id` matching the prior assistant row's `tool_calls[].id`. |
|
||||
| `tool_use` (Anthropic) | `assistant` with `tool_calls` | Anthropic emits tool calls *inside* an assistant message; flatten to OpenAI shape. |
|
||||
| `human_feedback`, `revision` | `user` | Treat as a follow-up user turn. |
|
||||
|
||||
## Tool Calls (the most error-prone part)
|
||||
|
||||
Turnstone stores tool calls in OpenAI's nested-function shape on the assistant row, and matches them with `role="tool"` result rows by `tool_call_id`.
|
||||
|
||||
### Assistant row with tool calls
|
||||
|
||||
```json
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": null,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_abc123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search_web",
|
||||
"arguments": "{\"query\":\"turnstone import\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`tool_calls[].function.arguments` is **a JSON-encoded string**, not an object. Source formats commonly get this wrong — Anthropic stores arguments as a parsed object, Gemini as a struct. Always re-serialize to a string.
|
||||
|
||||
### Tool result row
|
||||
|
||||
```json
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_name": "search_web",
|
||||
"tool_call_id": "call_abc123",
|
||||
"content": "..."
|
||||
}
|
||||
```
|
||||
|
||||
Pairing rules:
|
||||
- Every assistant `tool_calls[].id` MUST be followed by exactly one `role="tool"` row with the matching `tool_call_id`, before the next user/assistant turn.
|
||||
- If the source dropped the tool result (cut-off transcript), insert a synthetic `role="tool"` row with `content="[tool result missing in source]"` to keep the chain valid. An assistant row with an unanswered `tool_calls[].id` will break replay and any LLM round-trip.
|
||||
- Multi-tool assistant turns: one `role="tool"` row per call, in any order, all before the next non-tool row.
|
||||
|
||||
### Tool ID generation
|
||||
|
||||
If the source used opaque tool IDs that aren't unique within a thread (some platforms reuse them), regenerate with a stable scheme like `f"call_{i}"` where `i` is a per-thread counter. Update both the assistant and tool rows together.
|
||||
|
||||
## Provider Fidelity (`provider_data`)
|
||||
|
||||
Skip this entirely for **archive** imports.
|
||||
|
||||
For **resumable** imports against the same provider, populate `provider_data` to preserve provider-specific tool-call metadata that the next API round-trip will require:
|
||||
|
||||
- **Anthropic**: `signature` field on thinking blocks; required for round-tripping extended-thinking responses.
|
||||
- **Gemini**: `thought_signature` on tool calls; required for fidelity.
|
||||
- **OpenAI**: typically nothing to preserve.
|
||||
|
||||
The runtime-side dict key is `_provider_content` (a list of provider-native blocks); the persisted column is `provider_data` (the same list, JSON-encoded). If you don't have provider-native blocks from the source — and you usually won't, because a foreign export won't include them — leave `provider_data` NULL. The first new turn will succeed without it, but the previous assistant turn's reasoning won't replay back to the model.
|
||||
|
||||
## Attachments
|
||||
|
||||
If the source thread had image or file attachments:
|
||||
|
||||
- **Size limits**: images ≤ 4 MiB, text documents ≤ 512 KiB. Reject or downsample anything bigger.
|
||||
- **Allowed types**: server validates magic bytes for images and UTF-8-decodes for text. Binary blobs that aren't images won't pass.
|
||||
- **Lifecycle**: pending → reserved → consumed. For imports, the cleanest path is to upload as pending and immediately consume by attaching to the relevant `conversations.id`.
|
||||
|
||||
Two import paths:
|
||||
|
||||
1. **Bulk-insert + post-attach**: insert messages first, get back the assistant/user `conversations.id`, then write `workstream_attachments` rows linking the file to `message_id`.
|
||||
2. **SDK multipart create**: `create_workstream(attachments=[...], initial_message=...)` for the *first* turn only — the server reserves and consumes them onto that turn. Doesn't help for mid-thread attachments.
|
||||
|
||||
For full-history imports with multiple attachments at different turns, path (1) is the only option.
|
||||
|
||||
## Validation Checklist
|
||||
|
||||
Before declaring success, verify:
|
||||
|
||||
- [ ] `ws_id` is 32-char lowercase hex.
|
||||
- [ ] `workstreams` row exists with the right `user_id`, `state`, `kind`.
|
||||
- [ ] Conversation rows are inserted **in order** (autoincrement `id` will reflect insert order).
|
||||
- [ ] Every assistant `tool_calls[].id` has a matching `role="tool"` row with the same `tool_call_id`.
|
||||
- [ ] `tool_calls[].function.arguments` is a JSON-encoded **string**, not a parsed object.
|
||||
- [ ] First message is typically `role="user"` (not `system`) — Turnstone composes its own system prompt at runtime.
|
||||
- [ ] No empty assistant rows (`content=NULL` AND `tool_calls=NULL` is invalid).
|
||||
- [ ] If multi-node: the `ws_id`'s bucket maps to a node that exists; `workstreams.node_id` matches.
|
||||
- [ ] Round-trip test: run `Storage.load_messages(ws_id)` and confirm the reconstructed list matches what you inserted (modulo timestamps).
|
||||
|
||||
## Anti-patterns
|
||||
|
||||
- **Don't import the source provider's system prompt verbatim.** Provider boilerplate ("You are Claude...", "You are ChatGPT...") will conflict with Turnstone's composed system message and confuse the model on resume. Drop it; preserve only user-authored custom instructions.
|
||||
- **Don't preserve foreign tool definitions as Turnstone tools.** If the source had custom tools that don't exist in Turnstone, the assistant rows that called them are still valid history (archive), but the workstream is **not resumable** — mark `state="closed"`.
|
||||
- **Don't fabricate `tool_call_id`s without re-pairing.** Mismatched ids silently break the replay chain on the next turn.
|
||||
- **Don't skip the `tool_name` field on `role="tool"` rows.** Some load paths use it for display and audit; NULL there will render as "unknown tool".
|
||||
- **Don't write through the LLM (`send()` per turn) for full history.** It's expensive, rewrites assistant turns, and rate-limits will bite long imports.
|
||||
|
||||
## Quick Reference
|
||||
|
||||
| Task | Path |
|
||||
|---|---|
|
||||
| Generate ws_id | `secrets.token_hex(16)` |
|
||||
| Bulk insert messages | `Storage.save_messages_bulk(rows)` |
|
||||
| Archive (read-only) | `state="closed"`, skip `provider_data` |
|
||||
| Resumable | `state="idle"`, populate `provider_data` if same provider |
|
||||
| Tool call id | OpenAI shape: `{"id": ..., "type": "function", "function": {"name": ..., "arguments": "<json string>"}}` |
|
||||
| Tool result row | `role="tool"`, `tool_name`, `tool_call_id`, `content` |
|
||||
| Source role → Turnstone role | See "Role Mapping" table |
|
||||
| Per-thread metadata | Store source IDs in `workstream_config` under `import.*` keys |
|
||||
|
||||
## Files to read before writing the importer
|
||||
|
||||
- `turnstone/core/storage/_schema.py` — authoritative table definitions.
|
||||
- `turnstone/core/storage/_protocol.py` — `save_message`, `save_messages_bulk`, `load_messages` signatures.
|
||||
- `turnstone/core/session.py` (around the message-save section) — how the runtime constructs in-memory message dicts; mirror this shape on import to round-trip cleanly.
|
||||
- `turnstone/api/server_schemas.py` — Pydantic shapes for the SDK paths if you go through HTTP.
|
||||
+5
-14
@@ -758,22 +758,18 @@ MCP tools (3):
|
||||
|
||||
### Dynamic tool refresh
|
||||
|
||||
MCP tool lists stay up-to-date without restart through three mechanisms:
|
||||
MCP tool lists stay up-to-date without restart through two mechanisms:
|
||||
|
||||
1. **Push notifications** -- MCP servers that declare `tools.listChanged: true` in
|
||||
their capabilities send `notifications/tools/list_changed` when their tool list
|
||||
changes. `MCPClientManager` registers a `message_handler` on each `ClientSession`
|
||||
that triggers an immediate refresh for that server.
|
||||
|
||||
2. **Periodic timer** -- Servers that do *not* support push notifications are polled
|
||||
on a configurable interval (default 4 hours). The timer is staggered using a
|
||||
launch-time seed (`monotonic_ns ^ pid`) so cluster nodes don't all hit MCP
|
||||
servers simultaneously. Configure via `[mcp] refresh_interval` in `config.toml`
|
||||
or `--mcp-refresh-interval SECONDS` on the CLI. Set to `0` to disable.
|
||||
|
||||
3. **Manual** -- `/mcp refresh` re-fetches tools from all servers immediately.
|
||||
2. **Manual** -- `/mcp refresh` re-fetches tools from all servers immediately.
|
||||
`/mcp refresh <server>` targets a single server. If a server has disconnected,
|
||||
manual refresh attempts reconnection.
|
||||
manual refresh attempts reconnection. The console admin panel exposes the
|
||||
same controls (refresh / reconnect buttons per server) for cluster-wide
|
||||
fan-out.
|
||||
|
||||
When tools change, `MCPClientManager` rebuilds its merged tool list using copy-on-write
|
||||
(new list/dict objects assigned atomically) and notifies all active `ChatSession`
|
||||
@@ -781,11 +777,6 @@ instances via registered listener callbacks. Each session rebuilds its `_tools`,
|
||||
`_task_tools`, `_agent_tools`, and reconstructs its `ToolSearchManager` (if active),
|
||||
preserving the set of previously expanded (discovered) tools.
|
||||
|
||||
```toml
|
||||
[mcp]
|
||||
refresh_interval = 14400 # seconds (default 4h), 0 to disable
|
||||
```
|
||||
|
||||
```
|
||||
/mcp refresh
|
||||
MCP refresh complete:
|
||||
|
||||
+3
-2
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "turnstone"
|
||||
version = "1.5.0"
|
||||
version = "1.5.8"
|
||||
description = "Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."
|
||||
readme = "README.md"
|
||||
license = "BUSL-1.1"
|
||||
@@ -24,7 +24,7 @@ classifiers = [
|
||||
dependencies = [
|
||||
"openai>=2.24",
|
||||
"httpx>=0.28",
|
||||
"mcp>=1.6",
|
||||
"mcp>=1.27",
|
||||
"starlette>=0.45",
|
||||
"uvicorn>=0.34",
|
||||
"sse-starlette>=2.0",
|
||||
@@ -35,6 +35,7 @@ dependencies = [
|
||||
"structlog>=24.1",
|
||||
"PyJWT>=2.8",
|
||||
"bcrypt>=4.0",
|
||||
"cryptography>=42",
|
||||
"python-frontmatter>=1.0",
|
||||
]
|
||||
|
||||
|
||||
Generated
+12
-12
@@ -373,9 +373,9 @@
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/@tybys/wasm-util": {
|
||||
"version": "0.10.1",
|
||||
"resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.1.tgz",
|
||||
"integrity": "sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg==",
|
||||
"version": "0.10.2",
|
||||
"resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.2.tgz",
|
||||
"integrity": "sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
@@ -902,9 +902,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/nanoid": {
|
||||
"version": "3.3.11",
|
||||
"resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.11.tgz",
|
||||
"integrity": "sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w==",
|
||||
"version": "3.3.12",
|
||||
"resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz",
|
||||
"integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
@@ -959,9 +959,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/postcss": {
|
||||
"version": "8.5.12",
|
||||
"resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.12.tgz",
|
||||
"integrity": "sha512-W62t/Se6rA0Az3DfCL0AqJwXuKwBeYg6nOaIgzP+xZ7N5BFCI7DYi1qs6ygUYT6rvfi6t9k65UMLJC+PHZpDAA==",
|
||||
"version": "8.5.13",
|
||||
"resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.13.tgz",
|
||||
"integrity": "sha512-qif0+jGGZoLWdHey3UFHHWP0H7Gbmsk8T5VEqyYFbWqPr1XqvLGBbk/sl8V5exGmcYJklJOhOQq1pV9IcsiFag==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
@@ -1060,9 +1060,9 @@
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/tinyexec": {
|
||||
"version": "1.1.1",
|
||||
"resolved": "https://registry.npmjs.org/tinyexec/-/tinyexec-1.1.1.tgz",
|
||||
"integrity": "sha512-VKS/ZaQhhkKFMANmAOhhXVoIfBXblQxGX1myCQ2faQrfmobMftXeJPcZGp0gS07ocvGJWDLZGyOZDadDBqYIJg==",
|
||||
"version": "1.1.2",
|
||||
"resolved": "https://registry.npmjs.org/tinyexec/-/tinyexec-1.1.2.tgz",
|
||||
"integrity": "sha512-dAqSqE/RabpBKI8+h26GfLq6Vb3JVXs30XYQjdMjaj/c2tS8IYYMbIzP599KtRj7c57/wYApb3QjgRgXmrCukA==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
|
||||
@@ -35,11 +35,10 @@ def _seed_children(
|
||||
The production path populates the registry via the cluster-event
|
||||
fan-out thread observing ``ws_created`` events. These tests just
|
||||
need a known-children set for the endpoint handlers to iterate —
|
||||
inject directly under ``_children_lock`` rather than spinning up
|
||||
the collector + fan-out plumbing.
|
||||
inject directly via the registry's bulk-merge surface rather than
|
||||
spinning up the collector + fan-out plumbing.
|
||||
"""
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked(coord_ws_id, child_ws_ids)
|
||||
adapter._registry.merge_children(coord_ws_id, child_ws_ids)
|
||||
|
||||
|
||||
class _AuthMiddleware(BaseHTTPMiddleware):
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Shared test helpers — kept out of conftest.py since these are factories,
|
||||
not fixtures, and several test files want to import them directly."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
def make_chat_session(**overrides: Any) -> Any:
|
||||
"""Build a minimal ``ChatSession`` with sane test defaults.
|
||||
|
||||
Caller passes any constructor arg as a kwarg to override the default —
|
||||
e.g. ``make_chat_session(memory_config=MemoryConfig(fetch_limit=5))``.
|
||||
"""
|
||||
from turnstone.core.session import ChatSession
|
||||
|
||||
defaults: dict[str, Any] = {
|
||||
"client": MagicMock(),
|
||||
"model": "test-model",
|
||||
"ui": MagicMock(),
|
||||
"instructions": None,
|
||||
"temperature": 0.5,
|
||||
"max_tokens": 4096,
|
||||
"tool_timeout": 30,
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return ChatSession(**defaults)
|
||||
|
||||
|
||||
def patch_session_storage(
|
||||
monkeypatch: Any,
|
||||
*,
|
||||
active: bool = True,
|
||||
raise_on_is_active: bool = False,
|
||||
) -> list[str]:
|
||||
"""Patch ``session.get_storage`` to a stub whose ``is_watch_active``
|
||||
returns *active* (or raises if *raise_on_is_active*). Returns the
|
||||
list of ``watch_id``s the predicate was called with.
|
||||
"""
|
||||
from turnstone.core import session as session_mod
|
||||
|
||||
calls: list[str] = []
|
||||
|
||||
class _Stub:
|
||||
def is_watch_active(self, watch_id: str) -> bool:
|
||||
calls.append(watch_id)
|
||||
if raise_on_is_active:
|
||||
raise RuntimeError("storage down")
|
||||
return active
|
||||
|
||||
monkeypatch.setattr(session_mod, "get_storage", lambda: _Stub())
|
||||
return calls
|
||||
@@ -1,10 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.core.mcp_client import MCPClientManager, StaticServerState
|
||||
from turnstone.core.mcp_crypto import MCPTokenCipher
|
||||
from turnstone.core.oidc import OIDCConfig
|
||||
|
||||
|
||||
def make_mcp_token_cipher() -> MCPTokenCipher:
|
||||
"""Build a single-key MCP token cipher for tests.
|
||||
|
||||
Used by test files that need to exercise ``MCPTokenStore`` round-
|
||||
trips without the lifespan-side configuration loader; centralised
|
||||
here so the key/material defaults stay aligned across files.
|
||||
"""
|
||||
import base64
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from turnstone.core.mcp_crypto import MCPTokenCipher, MCPTokenCipherConfig
|
||||
|
||||
raw = base64.urlsafe_b64decode(Fernet.generate_key())
|
||||
return MCPTokenCipher(MCPTokenCipherConfig(keys=(raw,)))
|
||||
|
||||
|
||||
def _seed_static_state(mgr: MCPClientManager, name: str, **overrides: Any) -> StaticServerState:
|
||||
"""Get-or-create a ``StaticServerState`` on ``mgr`` and apply ``overrides``.
|
||||
|
||||
Shared across MCP test files so the helper stays in one place. Imported
|
||||
where needed; ``StaticServerState`` is constructed lazily so non-MCP
|
||||
tests don't pay the import cost.
|
||||
"""
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
state = mgr._static_servers.get(name)
|
||||
if state is None:
|
||||
state = StaticServerState(name=name)
|
||||
mgr._static_servers[name] = state
|
||||
for k, v in overrides.items():
|
||||
setattr(state, k, v)
|
||||
return state
|
||||
|
||||
|
||||
def make_oidc_test_config(**overrides: Any) -> OIDCConfig:
|
||||
"""Build a test ``OIDCConfig`` with sensible defaults.
|
||||
|
||||
Shared between ``test_oidc.py`` and ``test_oidc_handlers.py`` so the
|
||||
defaults (including the now-required ``redirect_base``) stay aligned.
|
||||
"""
|
||||
from turnstone.core.oidc import OIDCConfig
|
||||
|
||||
defaults: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"issuer": "https://idp.example.com",
|
||||
"client_id": "my-client",
|
||||
"client_secret": "my-secret",
|
||||
"scopes": "openid email profile",
|
||||
"provider_name": "TestIDP",
|
||||
"role_claim": "",
|
||||
"role_map": {},
|
||||
"password_enabled": True,
|
||||
"redirect_base": "https://app.example.com",
|
||||
"authorization_endpoint": "https://idp.example.com/authorize",
|
||||
"token_endpoint": "https://idp.example.com/token",
|
||||
"userinfo_endpoint": "https://idp.example.com/userinfo",
|
||||
"jwks_uri": "https://idp.example.com/.well-known/jwks.json",
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return OIDCConfig(**defaults)
|
||||
|
||||
|
||||
def pytest_addoption(parser: pytest.Parser) -> None:
|
||||
parser.addoption(
|
||||
|
||||
@@ -0,0 +1,391 @@
|
||||
"""Spike 1 — validate MCP SDK behavior for the per-(user, server) session pool.
|
||||
|
||||
Three scenarios:
|
||||
|
||||
1. N=20 concurrent ClientSession instances to the same URL.
|
||||
Verifies: no FD blow-up, no shared transport state, each session's
|
||||
tools/list returns independently.
|
||||
|
||||
2. Two concurrent tools/call on a shared ClientSession with interleaving
|
||||
payloads. Verifies: request_id demux works under contention.
|
||||
|
||||
3. Per-session Authorization header isolation. Verifies: different Bearer
|
||||
tokens per ClientSession reach the server with the expected
|
||||
Authorization header — i.e. httpx connection pooling does not cross
|
||||
headers between sessions.
|
||||
|
||||
Run: uv run python tests/spike_sdk_concurrency.py
|
||||
|
||||
Outcome gates Phase 5's pool architecture; if any scenario fails, fall
|
||||
back to per-call header injection (Alternative F in the OAuth-MCP RFC).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import uvicorn
|
||||
from mcp import ClientSession
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
# Reduce uvicorn / mcp log noise so spike output is readable.
|
||||
logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
logging.getLogger("mcp").setLevel(logging.WARNING)
|
||||
|
||||
# Records (auth_header, tool_name) per request — populated by the
|
||||
# AuthHeaderRecorder middleware below. Indexed by call sequence.
|
||||
SERVER_OBSERVATIONS: list[tuple[str | None, str | None]] = []
|
||||
# Tool-call payloads observed (for request_id demux verification).
|
||||
TOOL_CALL_PAYLOADS: list[str] = []
|
||||
|
||||
|
||||
class AuthHeaderRecorder(BaseHTTPMiddleware):
|
||||
"""Records the Authorization header on every request the server sees."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
||||
auth = request.headers.get("authorization")
|
||||
# We only record the auth header here; tool name comes from the
|
||||
# body payload which we can't read non-destructively. The tool
|
||||
# handler logs the payload it received.
|
||||
SERVER_OBSERVATIONS.append((auth, None))
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
"""Bind to port 0, return the assigned port."""
|
||||
s = socket.socket()
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
s.close()
|
||||
return port
|
||||
|
||||
|
||||
def build_server(port: int) -> uvicorn.Server:
|
||||
"""Create a minimal FastMCP server with one echo tool."""
|
||||
mcp = FastMCP(name="spike-target", streamable_http_path="/mcp")
|
||||
|
||||
@mcp.tool()
|
||||
async def echo(payload: str) -> str:
|
||||
"""Echo the payload back. Records the payload server-side."""
|
||||
TOOL_CALL_PAYLOADS.append(payload)
|
||||
# Add a small await so two concurrent calls can interleave
|
||||
# on the wire if the SDK pools the requests.
|
||||
await asyncio.sleep(0.05)
|
||||
return f"echoed:{payload}"
|
||||
|
||||
app = mcp.streamable_http_app()
|
||||
app.add_middleware(AuthHeaderRecorder)
|
||||
|
||||
config = uvicorn.Config(
|
||||
app,
|
||||
host="127.0.0.1",
|
||||
port=port,
|
||||
log_level="warning",
|
||||
access_log=False,
|
||||
)
|
||||
return uvicorn.Server(config)
|
||||
|
||||
|
||||
def run_server_in_thread(server: uvicorn.Server) -> threading.Thread:
|
||||
"""Boot the server in a background thread on its own asyncio loop."""
|
||||
|
||||
def _run() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(server.serve())
|
||||
|
||||
t = threading.Thread(target=_run, daemon=True, name="spike-server")
|
||||
t.start()
|
||||
return t
|
||||
|
||||
|
||||
async def wait_for_server_ready(url: str, timeout: float = 5.0) -> None:
|
||||
"""Poll the server until it accepts connections."""
|
||||
import urllib.parse
|
||||
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
reader, writer = await asyncio.open_connection(parsed.hostname, parsed.port)
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
return
|
||||
except OSError:
|
||||
await asyncio.sleep(0.05)
|
||||
raise TimeoutError(f"server at {url} not ready within {timeout}s")
|
||||
|
||||
|
||||
def fd_count() -> int:
|
||||
"""Count open file descriptors for the current process."""
|
||||
try:
|
||||
return len(os.listdir(f"/proc/{os.getpid()}/fd"))
|
||||
except OSError:
|
||||
return -1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scenario 1: N=20 concurrent ClientSession instances
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def scenario_1_concurrent_sessions(url: str, n: int = 20) -> dict:
|
||||
"""Open N concurrent ClientSession instances and call tools/list on each."""
|
||||
print(f"\n=== Scenario 1: {n} concurrent ClientSession instances ===")
|
||||
fd_before = fd_count()
|
||||
|
||||
async def one_session(idx: int) -> dict:
|
||||
headers = {"Authorization": f"Bearer test-token-{idx}"}
|
||||
async with (
|
||||
streamablehttp_client(url=url, headers=headers) as (read, write, _),
|
||||
ClientSession(read, write) as session,
|
||||
):
|
||||
await session.initialize()
|
||||
tools = await session.list_tools()
|
||||
return {
|
||||
"idx": idx,
|
||||
"tool_count": len(tools.tools),
|
||||
"tool_names": [t.name for t in tools.tools],
|
||||
}
|
||||
|
||||
start = time.monotonic()
|
||||
results = await asyncio.gather(*[one_session(i) for i in range(n)], return_exceptions=True)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
fd_after = fd_count()
|
||||
# Allow some settling time for FDs to release.
|
||||
await asyncio.sleep(0.5)
|
||||
fd_settled = fd_count()
|
||||
|
||||
successes = [r for r in results if isinstance(r, dict)]
|
||||
failures = [r for r in results if isinstance(r, Exception)]
|
||||
|
||||
# Verify every session got the same tool catalog.
|
||||
catalog_consistent = (
|
||||
len(successes) == n and len({tuple(r["tool_names"]) for r in successes}) == 1
|
||||
)
|
||||
|
||||
return {
|
||||
"scenario": "concurrent_sessions",
|
||||
"n": n,
|
||||
"successes": len(successes),
|
||||
"failures": len(failures),
|
||||
"elapsed_seconds": round(elapsed, 3),
|
||||
"fd_before": fd_before,
|
||||
"fd_during_peak": fd_after,
|
||||
"fd_settled": fd_settled,
|
||||
"fd_growth_during": fd_after - fd_before,
|
||||
"fd_growth_settled": fd_settled - fd_before,
|
||||
"catalog_consistent": catalog_consistent,
|
||||
"first_failure": str(failures[0]) if failures else None,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scenario 2: 2 concurrent tools/call on a shared session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def scenario_2_concurrent_calls_shared_session(url: str) -> dict:
|
||||
"""Two concurrent tools/call on one ClientSession with interleaving payloads.
|
||||
|
||||
The echo tool sleeps 50ms, so concurrent calls overlap on the wire.
|
||||
Each call passes a distinct payload (~10KB) to make request bodies
|
||||
spannable across multiple stream frames.
|
||||
"""
|
||||
print("\n=== Scenario 2: 2 concurrent tools/call on shared session ===")
|
||||
|
||||
# Generous-size payloads so both bodies live during the await.
|
||||
payload_a = "A" * 10000
|
||||
payload_b = "B" * 10000
|
||||
|
||||
headers = {"Authorization": "Bearer shared-session-token"}
|
||||
async with (
|
||||
streamablehttp_client(url=url, headers=headers) as (read, write, _),
|
||||
ClientSession(read, write) as session,
|
||||
):
|
||||
await session.initialize()
|
||||
|
||||
TOOL_CALL_PAYLOADS.clear()
|
||||
|
||||
start = time.monotonic()
|
||||
results = await asyncio.gather(
|
||||
session.call_tool("echo", {"payload": payload_a}),
|
||||
session.call_tool("echo", {"payload": payload_b}),
|
||||
return_exceptions=True,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
successes = [r for r in results if not isinstance(r, Exception)]
|
||||
failures = [r for r in results if isinstance(r, Exception)]
|
||||
|
||||
# Each result.content[0].text should be "echoed:{payload}".
|
||||
response_payloads: list[str] = []
|
||||
if len(successes) == 2:
|
||||
for r in successes:
|
||||
text = r.content[0].text if r.content else ""
|
||||
response_payloads.append(text)
|
||||
|
||||
# Order may not match call order — what matters is both payloads echo.
|
||||
expected = {f"echoed:{payload_a}", f"echoed:{payload_b}"}
|
||||
received = set(response_payloads)
|
||||
demux_ok = received == expected
|
||||
|
||||
# Did both calls actually overlap? If sequential, elapsed ~= 0.1+s;
|
||||
# if concurrent, ~0.05s.
|
||||
concurrent_observed = elapsed < 0.09
|
||||
|
||||
return {
|
||||
"scenario": "concurrent_calls_shared_session",
|
||||
"successes": len(successes),
|
||||
"failures": len(failures),
|
||||
"elapsed_seconds": round(elapsed, 3),
|
||||
"demux_ok": demux_ok,
|
||||
"expected_payloads_received": list(received) if demux_ok else None,
|
||||
"actual_payloads_received_count": len(received),
|
||||
"appears_concurrent_on_wire": concurrent_observed,
|
||||
"first_failure": str(failures[0]) if failures else None,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scenario 3: per-session header isolation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def scenario_3_header_isolation(url: str, n: int = 5) -> dict:
|
||||
"""Open N sessions with distinct Authorization headers, call echo on each.
|
||||
|
||||
Verifies the server sees each session's own header — i.e. httpx
|
||||
connection pooling does not cross headers between concurrent
|
||||
ClientSession instances against the same URL.
|
||||
"""
|
||||
print(f"\n=== Scenario 3: {n}-session Authorization-header isolation ===")
|
||||
SERVER_OBSERVATIONS.clear()
|
||||
|
||||
async def one_session(idx: int) -> str | None:
|
||||
headers = {"Authorization": f"Bearer iso-token-{idx}"}
|
||||
async with (
|
||||
streamablehttp_client(url=url, headers=headers) as (read, write, _),
|
||||
ClientSession(read, write) as session,
|
||||
):
|
||||
await session.initialize()
|
||||
# One call per session.
|
||||
await session.call_tool("echo", {"payload": f"session-{idx}"})
|
||||
return f"Bearer iso-token-{idx}"
|
||||
|
||||
start = time.monotonic()
|
||||
expected_tokens = await asyncio.gather(*[one_session(i) for i in range(n)])
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Tally observed Authorization headers, ignoring None entries (initial
|
||||
# handshake sometimes lacks auth).
|
||||
observed_auth = [auth for auth, _ in SERVER_OBSERVATIONS if auth]
|
||||
expected_set = set(expected_tokens)
|
||||
observed_set = set(observed_auth)
|
||||
|
||||
# Every expected token must show up at least once on the server.
|
||||
all_present = expected_set.issubset(observed_set)
|
||||
# No spurious tokens.
|
||||
no_extras = observed_set.issubset(expected_set)
|
||||
# Frequency: at least one observation per token.
|
||||
counts = defaultdict(int)
|
||||
for a in observed_auth:
|
||||
counts[a] += 1
|
||||
each_seen = all(counts[t] >= 1 for t in expected_tokens)
|
||||
|
||||
return {
|
||||
"scenario": "header_isolation",
|
||||
"n": n,
|
||||
"elapsed_seconds": round(elapsed, 3),
|
||||
"expected_tokens": sorted(expected_set),
|
||||
"observed_tokens": sorted(observed_set),
|
||||
"all_expected_present": all_present,
|
||||
"no_extra_tokens_observed": no_extras,
|
||||
"each_token_seen_at_least_once": each_seen,
|
||||
"header_counts_per_token": dict(counts),
|
||||
"total_requests_observed": len(observed_auth),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Driver
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
port = find_free_port()
|
||||
url = f"http://127.0.0.1:{port}/mcp"
|
||||
|
||||
server = build_server(port)
|
||||
server_thread = run_server_in_thread(server)
|
||||
try:
|
||||
await wait_for_server_ready(url)
|
||||
print(f"server up at {url}\n")
|
||||
|
||||
result_1 = await scenario_1_concurrent_sessions(url, n=20)
|
||||
print_scenario_result(result_1)
|
||||
|
||||
result_2 = await scenario_2_concurrent_calls_shared_session(url)
|
||||
print_scenario_result(result_2)
|
||||
|
||||
result_3 = await scenario_3_header_isolation(url, n=5)
|
||||
print_scenario_result(result_3)
|
||||
|
||||
# Final verdict
|
||||
verdict_1 = (
|
||||
result_1["successes"] == result_1["n"]
|
||||
and result_1["catalog_consistent"]
|
||||
and result_1["fd_growth_settled"] < 30 # 20 sessions, generous bound
|
||||
)
|
||||
verdict_2 = result_2["demux_ok"] and result_2["successes"] == 2
|
||||
verdict_3 = (
|
||||
result_3["all_expected_present"]
|
||||
and result_3["no_extra_tokens_observed"]
|
||||
and result_3["each_token_seen_at_least_once"]
|
||||
)
|
||||
|
||||
print("\n=== VERDICT ===")
|
||||
print(f" Scenario 1 (concurrent sessions): {'PASS' if verdict_1 else 'FAIL'}")
|
||||
print(f" Scenario 2 (concurrent calls shared): {'PASS' if verdict_2 else 'FAIL'}")
|
||||
print(f" Scenario 3 (header isolation): {'PASS' if verdict_3 else 'FAIL'}")
|
||||
all_pass = verdict_1 and verdict_2 and verdict_3
|
||||
print(
|
||||
f"\n Phase 5 per-(user, server) pool architecture: "
|
||||
f"{'VIABLE' if all_pass else 'NEEDS REWORK (Alternative F fallback)'}"
|
||||
)
|
||||
sys.exit(0 if all_pass else 1)
|
||||
finally:
|
||||
server.should_exit = True
|
||||
server_thread.join(timeout=5)
|
||||
|
||||
|
||||
def print_scenario_result(result: dict) -> None:
|
||||
print(f"\nresult[{result['scenario']}]:")
|
||||
for k, v in result.items():
|
||||
if k == "scenario":
|
||||
continue
|
||||
print(f" {k}: {v}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with contextlib.suppress(KeyboardInterrupt):
|
||||
asyncio.run(main())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -92,3 +92,366 @@ def test_tool_error_does_not_overwrite_approval_badge() -> None:
|
||||
"badge instead so the approval verdict stays visible alongside "
|
||||
"the error."
|
||||
)
|
||||
|
||||
|
||||
def test_replay_history_renders_content_before_tool_block() -> None:
|
||||
"""In ``replayHistory``'s ``role === "assistant"`` branch, the
|
||||
``msg.content`` render must precede the ``msg.tool_calls`` render.
|
||||
|
||||
Two reasons, both load-bearing:
|
||||
|
||||
1. **Structural** — the next loop iteration's ``role === "tool"``
|
||||
message anchors to ``lastToolBlock``. The tool-block branch sets
|
||||
that anchor; the content branch clears it. If content runs after
|
||||
the tool block, the clear silently drops the upcoming tool
|
||||
result. Pre-fix, every interactive tool result was missing from
|
||||
saved-workstream replays whenever the assistant turn carried
|
||||
both narration and tool calls (very common output shape).
|
||||
|
||||
2. **Visual** — the live SSE path renders content first
|
||||
(``stream_text`` streams before ``tool_info`` /
|
||||
``approve_request``), so replay should match.
|
||||
|
||||
The test pins the order via the offsets of the ``msg.content`` and
|
||||
``msg.tool_calls`` branch headers inside the function body."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("Pane.prototype.replayHistory = function")
|
||||
end = body.index("Pane.prototype._attachRetryToLastAssistant", start)
|
||||
fn = body[start:end]
|
||||
# Locate the assistant branch and bound the search to its body —
|
||||
# the function also handles user / tool roles which would otherwise
|
||||
# confuse the offset comparison.
|
||||
asst_start = fn.index('msg.role === "assistant"')
|
||||
asst_end = fn.index('msg.role === "tool"', asst_start)
|
||||
asst = fn[asst_start:asst_end]
|
||||
# ``if (msg.content && msg.content.trim())`` guards against a
|
||||
# whitespace-only content row (Qwen-style "\n\n" left over after a
|
||||
# reasoning-parser model strips ``<think>…</think>`` and emits
|
||||
# nothing else before the tool call). Pre-trim guard, those rows
|
||||
# rendered as a visible-but-empty ``.msg.assistant`` card on
|
||||
# replay. Match the substring up to ``msg.content`` so the test
|
||||
# tolerates either guard shape without locking the trim() in.
|
||||
content_idx = asst.index("if (msg.content")
|
||||
tool_calls_idx = asst.index("if (msg.tool_calls && msg.tool_calls.length)")
|
||||
assert content_idx < tool_calls_idx, (
|
||||
"replayHistory must render msg.content BEFORE msg.tool_calls "
|
||||
"inside the assistant branch — otherwise the lastToolBlock "
|
||||
"anchor is clobbered before the next iteration's tool result "
|
||||
"can attach to it (and the visual order also drifts from the "
|
||||
"live SSE flow)."
|
||||
)
|
||||
|
||||
|
||||
def test_replay_history_renders_persisted_verdict_badge() -> None:
|
||||
"""Saved-workstream replays must paint the persisted intent verdict
|
||||
next to each tool div, using the same ``renderVerdictBadge`` helper
|
||||
the live ``showInlineToolBlock`` path uses. Pre-fix the audit trail
|
||||
was complete in storage (``intent_verdicts`` table) but never
|
||||
surfaced on replay — operators reviewing a saved workstream
|
||||
couldn't see what the heuristic / LLM judge thought of any tool
|
||||
call. This test pins the call site so a refactor that drops the
|
||||
decoration regresses the audit surface."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("Pane.prototype.replayHistory = function")
|
||||
end = body.index("Pane.prototype._attachRetryToLastAssistant", start)
|
||||
fn = body[start:end]
|
||||
# Match a `renderVerdictBadge(<something>.verdict, ...)` call inside
|
||||
# the replay loop. Loose on whitespace + identifier so a future
|
||||
# rename of the iteration variable doesn't trip CI.
|
||||
badge_call_re = re.compile(
|
||||
r"renderVerdictBadge\(\s*\w+\.verdict\b",
|
||||
)
|
||||
assert badge_call_re.search(fn), (
|
||||
"replayHistory must call renderVerdictBadge(tc.verdict, ...) "
|
||||
"when a persisted verdict is attached to a tool_call entry — "
|
||||
"otherwise the audit-trail data persisted to intent_verdicts "
|
||||
"doesn't surface on saved-workstream replays."
|
||||
)
|
||||
|
||||
|
||||
def test_shared_utils_defines_replay_advisories_after_tool() -> None:
|
||||
"""The shared ``replayAdvisoriesAfterTool`` helper in
|
||||
``shared_static/utils.js`` is the single source of advisory-walk +
|
||||
type-filter logic for both ``app.js`` (interactive) and
|
||||
``coordinator.js`` (coord). A refactor that drops the helper
|
||||
breaks both surfaces, so guard its definition + filter shape here.
|
||||
"""
|
||||
utils_js = Path(__file__).resolve().parent.parent / "turnstone/shared_static/utils.js"
|
||||
body = utils_js.read_text(encoding="utf-8")
|
||||
assert "function replayAdvisoriesAfterTool" in body, (
|
||||
"shared/utils.js must define replayAdvisoriesAfterTool — "
|
||||
"interactive and coord both invoke it."
|
||||
)
|
||||
# The type filter — ``adv.type !== 'user_interjection'`` — must
|
||||
# remain in the helper so a future advisory shape (output_guard,
|
||||
# metacognitive nudge, etc.) doesn't silently render as a user
|
||||
# bubble.
|
||||
assert 'adv.type !== "user_interjection"' in body, (
|
||||
"replayAdvisoriesAfterTool must filter by advisory type so a "
|
||||
"future non-user_interjection advisory shape doesn't silently "
|
||||
"render as a user bubble."
|
||||
)
|
||||
|
||||
|
||||
def test_replay_renders_user_interjection_advisory_after_tool_block() -> None:
|
||||
"""Queued user messages spliced into the last tool-result envelope
|
||||
of a batch (Seam 1) persist on the tool DB row as a wrapped
|
||||
``<tool_output>`` envelope. ``decorate_history_messages`` extracts
|
||||
the advisory back out and the wire layer projects it onto
|
||||
``msg.advisories``; ``replayHistory`` must invoke the shared
|
||||
``replayAdvisoriesAfterTool`` helper (defined in
|
||||
``shared/utils.js``) so each ``user_interjection`` renders through
|
||||
``addUserMessage`` and the bubble looks identical to a Seam 2/3
|
||||
user row.
|
||||
|
||||
This test pins the call site so a refactor that drops the helper
|
||||
invocation regresses the queued-during-batch replay shape
|
||||
silently."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("Pane.prototype.replayHistory = function")
|
||||
end = body.index("Pane.prototype._attachRetryToLastAssistant", start)
|
||||
fn = body[start:end]
|
||||
# The replay loop must invoke the shared helper, passing
|
||||
# ``msg.advisories`` and a renderer that routes through
|
||||
# ``addUserMessage``. The helper itself filters on
|
||||
# ``adv.type !== "user_interjection"``; that branch lives in
|
||||
# ``shared/utils.js`` (test_shared_utils_js or runtime smoke covers
|
||||
# the helper's body).
|
||||
assert "replayAdvisoriesAfterTool(msg.advisories" in fn, (
|
||||
"replayHistory must invoke replayAdvisoriesAfterTool with "
|
||||
"msg.advisories so queued messages spliced into the tool "
|
||||
"envelope render as user bubbles after the tool block."
|
||||
)
|
||||
assert "addUserMessage(text" in fn, (
|
||||
"replayHistory's renderer callback must route the extracted "
|
||||
"advisory text through addUserMessage so the rendered bubble "
|
||||
"matches a normal user-row replay."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 8 — Chunk D: MCP error embed + settings panel UX
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_INDEX_HTML = Path(__file__).resolve().parent.parent / "turnstone/ui/static/index.html"
|
||||
_STYLE_CSS = Path(__file__).resolve().parent.parent / "turnstone/ui/static/style.css"
|
||||
|
||||
# The Phase-8 D-chunk pins the absence of an unsafe DOM-write API
|
||||
# in two regions of app.js. Spell the property name out of literal
|
||||
# concatenation so the tooling that flags occurrences in code
|
||||
# strings doesn't false-positive on the test source.
|
||||
_UNSAFE_DOM_WRITE_RE = re.compile(r"\.inner" + r"HTML\s*=")
|
||||
|
||||
|
||||
def test_phase8_mcp_error_helpers_defined_in_app_js() -> None:
|
||||
"""The Phase 8 dashboard renderer adds three load-bearing helpers
|
||||
next to the existing media-embed pattern: ``tryParseMcpError``
|
||||
(envelope detector), ``buildMcpErrorEmbed`` (interactive card),
|
||||
and the ``_pendingConsentServers`` set that drives the gear-icon
|
||||
badge. A regression that drops any of them silently degrades the
|
||||
OAuth consent UX to a plain JSON dump, so guard their existence
|
||||
here."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
assert "function tryParseMcpError" in body, (
|
||||
"tryParseMcpError must remain defined — appendToolOutput's "
|
||||
"error branch depends on it to detect the MCP error envelope."
|
||||
)
|
||||
assert "function buildMcpErrorEmbed" in body, (
|
||||
"buildMcpErrorEmbed must remain defined — it renders the "
|
||||
"interactive consent / forbidden / operator card."
|
||||
)
|
||||
assert "_pendingConsentServers" in body, (
|
||||
"_pendingConsentServers state must remain — it backs the "
|
||||
"gear-icon badge so a user who scrolls past a consent prompt "
|
||||
"still has a stable signal that consent is pending."
|
||||
)
|
||||
# The buildMcpErrorEmbed pattern must also wire the "actionable"
|
||||
# branch (consent_required / insufficient_scope) into the badge
|
||||
# via _onConsentDetected; pin the helper name.
|
||||
assert "_onConsentDetected" in body, (
|
||||
"_onConsentDetected must remain — buildMcpErrorEmbed calls it "
|
||||
"for the actionable category to surface the gear-icon badge."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_settings_panel_handlers_defined() -> None:
|
||||
"""The settings modal exposes four entry points that the inline
|
||||
``onclick`` attributes in index.html depend on. Renaming or
|
||||
deleting any of them breaks the modal silently (the buttons are
|
||||
still rendered but click-to-action is dead). Catch that here."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
for name in [
|
||||
"function openSettingsPanel",
|
||||
"function closeSettingsPanel",
|
||||
"function confirmRevokeMcp",
|
||||
"function cancelRevokeMcp",
|
||||
]:
|
||||
assert name in body, f"Missing required handler: {name}"
|
||||
# The connections list is fetched against the Phase-7 endpoint —
|
||||
# pin the URL so a server-side rename forces an explicit UI bump.
|
||||
assert "/v1/api/mcp/oauth/connections" in body, (
|
||||
"Settings panel must fetch /v1/api/mcp/oauth/connections — "
|
||||
"a server-side rename needs an explicit UI update."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_appendtooloutput_dispatches_mcp_error_before_renderer() -> None:
|
||||
"""``appendToolOutput`` must call ``tryParseMcpError`` inside its
|
||||
``isError`` branch BEFORE falling through to the plain
|
||||
``renderToolOutput`` path. The ordering is what makes the
|
||||
interactive consent card replace the JSON dump; reverse the calls
|
||||
and the user sees the raw error envelope as text again."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("Pane.prototype.appendToolOutput = function")
|
||||
end = body.index("Pane.prototype.", start + 10)
|
||||
fn = body[start:end]
|
||||
parse_idx = fn.find("tryParseMcpError(")
|
||||
render_idx = fn.find("renderToolOutput(")
|
||||
assert parse_idx >= 0, (
|
||||
"appendToolOutput must call tryParseMcpError on the error path "
|
||||
"before renderToolOutput, otherwise the consent card never "
|
||||
"replaces the plain JSON output."
|
||||
)
|
||||
assert render_idx >= 0, "renderToolOutput call must remain present"
|
||||
assert parse_idx < render_idx, (
|
||||
"tryParseMcpError must run BEFORE renderToolOutput so the "
|
||||
"interactive card path takes precedence over plain rendering."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_no_unsafe_dom_write_in_settings_panel() -> None:
|
||||
"""Defensive XSS guard: the settings panel renders user-controlled
|
||||
server names, scope strings, and timestamp values into the DOM.
|
||||
The whole section MUST go through ``textContent``-style APIs; an
|
||||
unsafe-DOM-write assignment would be a regression vector. Bound
|
||||
the check to the section 15 body to avoid false positives
|
||||
elsewhere."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("// 15. MCP server connections settings panel")
|
||||
# Bound to the full settings section (terminates at the next
|
||||
# top-level keydown handler block).
|
||||
end = body.index('document.addEventListener("keydown"', start)
|
||||
section = body[start:end]
|
||||
assert not _UNSAFE_DOM_WRITE_RE.search(section), (
|
||||
"Section 15 must not assign to the unsafe DOM-write property — "
|
||||
"server names and scope values flow through here and would be "
|
||||
"XSS-injectable. Use textContent / DOM APIs instead."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_settings_button_in_index_html() -> None:
|
||||
"""The gear-icon entry-point for the settings panel must remain
|
||||
in the appbar's actions span. The console proxy IIFE prepends a
|
||||
node pill to ``header.firstChild`` (turnstone/console/server.py:
|
||||
202); our button is appended inside ``<span class='appbar-actions'>``
|
||||
on the right, so they don't collide. Pin both shape constraints
|
||||
here so a future appbar refactor keeps them disjoint."""
|
||||
body = _INDEX_HTML.read_text(encoding="utf-8")
|
||||
assert 'id="settings-btn"' in body, (
|
||||
"index.html must keep the #settings-btn — onclick handlers "
|
||||
"and the consent badge target it by id."
|
||||
)
|
||||
assert 'onclick="openSettingsPanel()"' in body, (
|
||||
"settings-btn must wire onclick=openSettingsPanel() — losing "
|
||||
"the binding leaves the panel unreachable."
|
||||
)
|
||||
# The button must live inside <span class="appbar-actions"> so the
|
||||
# console proxy's header.insertBefore(pill, header.firstChild)
|
||||
# leaves it untouched.
|
||||
actions_open = body.index('class="appbar-actions"')
|
||||
actions_close = body.index("</span>", actions_open)
|
||||
assert 'id="settings-btn"' in body[actions_open:actions_close], (
|
||||
"settings-btn must be inside <span class='appbar-actions'> "
|
||||
"so the console proxy's firstChild prepend doesn't shift it."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_settings_modal_in_index_html() -> None:
|
||||
"""Both the settings overlay and the revoke-confirmation overlay
|
||||
must remain in the modal area. The Escape-key deferral list in
|
||||
app.js targets these ids, so removing them silently breaks the
|
||||
handler chain."""
|
||||
body = _INDEX_HTML.read_text(encoding="utf-8")
|
||||
assert 'id="settings-overlay"' in body
|
||||
assert 'id="revoke-mcp-overlay"' in body
|
||||
# Each overlay must have role="dialog" + aria-modal="true" so
|
||||
# screen readers and the existing modal-deferral handlers can
|
||||
# treat them like the rest of the modal stack.
|
||||
for overlay_id in ("settings-overlay", "revoke-mcp-overlay"):
|
||||
idx = body.index(f'id="{overlay_id}"')
|
||||
# Bound to ~600 chars after the open tag so we only check this
|
||||
# overlay's attributes.
|
||||
chunk = body[idx : idx + 600]
|
||||
assert 'role="dialog"' in chunk, f"{overlay_id} missing role=dialog"
|
||||
assert 'aria-modal="true"' in chunk, f"{overlay_id} missing aria-modal=true"
|
||||
|
||||
|
||||
def test_phase8_xss_safe_render_in_build_mcp_error_embed() -> None:
|
||||
"""Adversarial input — the renderer for an MCP error envelope
|
||||
must use ``textContent`` (not the unsafe DOM-write API) for every
|
||||
field that flows from the server: ``err.detail``, ``err.server``,
|
||||
scopes list. The card builder uses createElement + textContent
|
||||
throughout so a script-tag server name renders harmlessly. Pin
|
||||
the absence of the unsafe-write inside ``buildMcpErrorEmbed``."""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
start = body.index("function buildMcpErrorEmbed(")
|
||||
# Bound to the function body — find its closing brace at column 0.
|
||||
rest = body[start:]
|
||||
# Closing function brace at line start (matches existing functions)
|
||||
end_match = re.search(r"\n}\n", rest)
|
||||
assert end_match is not None
|
||||
fn = rest[: end_match.end()]
|
||||
assert not _UNSAFE_DOM_WRITE_RE.search(fn), (
|
||||
"buildMcpErrorEmbed must not use the unsafe-DOM-write API — "
|
||||
"server names and detail strings flow through here. An "
|
||||
"adversarial server name must render harmlessly via "
|
||||
"textContent."
|
||||
)
|
||||
|
||||
|
||||
def test_phase8_css_classes_present_in_stylesheet() -> None:
|
||||
"""The card / badge / modal classes referenced from app.js must
|
||||
have CSS rules. Without them the DOM still works but the visual
|
||||
treatment is gone, which would silently degrade the consent UX."""
|
||||
css = _STYLE_CSS.read_text(encoding="utf-8")
|
||||
for selector in [
|
||||
".mcp-error-card",
|
||||
".mcp-error-icon",
|
||||
".mcp-error-action-btn",
|
||||
".mcp-scope-pill",
|
||||
"#settings-overlay",
|
||||
"#settings-box",
|
||||
".settings-revoke-btn",
|
||||
".settings-consent-badge",
|
||||
"#revoke-mcp-overlay",
|
||||
]:
|
||||
assert selector in css, f"Missing CSS rule for {selector}"
|
||||
|
||||
|
||||
def test_phase8_consent_url_prefix_check_in_click_handler() -> None:
|
||||
"""Defence-in-depth: the consent button's click handler must reject
|
||||
any ``consent_url`` that doesn't start with the dispatcher's known
|
||||
prefix (``/v1/api/mcp/oauth/start``). ``_build_consent_url`` always
|
||||
emits a path-relative URL with that exact prefix; a non-prefix
|
||||
value implies the producer drifted (or was compromised) and a
|
||||
``window.open("javascript:...")`` would be catastrophic.
|
||||
|
||||
The renderer is the last line of defence before ``window.open`` and
|
||||
must not rely on the producer-side guarantee alone. Pin the prefix
|
||||
string and the ``startsWith`` form so a future refactor can't
|
||||
silently weaken the guard.
|
||||
"""
|
||||
body = _APP_JS.read_text(encoding="utf-8")
|
||||
# Bound the search to the click handler region (between the
|
||||
# ``buildMcpErrorEmbed`` function and the next top-level helper) to
|
||||
# avoid false positives from unrelated string occurrences.
|
||||
start = body.index("function buildMcpErrorEmbed(")
|
||||
end = body.index("\n}\n", start) + 1
|
||||
fn = body[start:end]
|
||||
assert 'consentUrl.startsWith("/v1/api/mcp/oauth/start")' in fn, (
|
||||
"Click handler must guard window.open with "
|
||||
'consentUrl.startsWith("/v1/api/mcp/oauth/start"). Without it '
|
||||
"a future producer drift to a non-path-relative URL (or a "
|
||||
'"javascript:" injection) would be passed straight to '
|
||||
"window.open."
|
||||
)
|
||||
|
||||
@@ -207,6 +207,30 @@ class TestRequiredScope:
|
||||
"""Only POST is elevated — GET falls through to read."""
|
||||
assert required_scope("GET", "/api/_internal/mcp-reload") == "read"
|
||||
|
||||
def test_internal_mcp_refresh_one_needs_approve(self):
|
||||
assert required_scope("POST", "/api/_internal/mcp-refresh/srv") == "approve"
|
||||
|
||||
def test_v1_internal_mcp_refresh_one_needs_approve(self):
|
||||
assert required_scope("POST", "/v1/api/_internal/mcp-refresh/srv") == "approve"
|
||||
|
||||
def test_proxy_internal_mcp_refresh_one_needs_approve(self):
|
||||
assert required_scope("POST", "/node/n1/v1/api/_internal/mcp-refresh/srv") == "approve"
|
||||
|
||||
def test_proxy_no_v1_internal_mcp_refresh_one_needs_approve(self):
|
||||
assert required_scope("POST", "/node/n1/api/_internal/mcp-refresh/srv") == "approve"
|
||||
|
||||
def test_internal_mcp_reconnect_one_needs_approve(self):
|
||||
assert required_scope("POST", "/api/_internal/mcp-reconnect/srv") == "approve"
|
||||
|
||||
def test_v1_internal_mcp_reconnect_one_needs_approve(self):
|
||||
assert required_scope("POST", "/v1/api/_internal/mcp-reconnect/srv") == "approve"
|
||||
|
||||
def test_proxy_internal_mcp_reconnect_one_needs_approve(self):
|
||||
assert required_scope("POST", "/node/n1/v1/api/_internal/mcp-reconnect/srv") == "approve"
|
||||
|
||||
def test_proxy_no_v1_internal_mcp_reconnect_one_needs_approve(self):
|
||||
assert required_scope("POST", "/node/n1/api/_internal/mcp-reconnect/srv") == "approve"
|
||||
|
||||
# Workstream sub-resource mutations (parametric paths)
|
||||
def test_ws_delete_needs_write(self):
|
||||
assert required_scope("POST", "/api/workstreams/abc123/delete") == "write"
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Tests for ``turnstone.server._build_history`` reminder + source surfacing.
|
||||
|
||||
The replay path (``_build_history``) projects the ``_source`` and
|
||||
``_reminders`` side-channels onto the wire entry the frontend
|
||||
consumes. Persisted via migration 050 (Commit 1) so multi-tab /
|
||||
multi-device replay sees the same metacognitive bubble shape the
|
||||
originating tab saw live.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
from turnstone.server import _build_history
|
||||
|
||||
|
||||
def _make_stub_session(messages: list[dict[str, Any]]) -> Any:
|
||||
"""Minimal ChatSession-shaped stub. ``_build_history`` only reads
|
||||
``session.messages`` plus calls ``_load_verdict_indexes(ws_id)`` —
|
||||
the latter we patch out below.
|
||||
"""
|
||||
return SimpleNamespace(messages=messages, _ws_id="ws-test")
|
||||
|
||||
|
||||
def _build(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Run ``_build_history`` against a stub session, bypassing the
|
||||
verdicts / output-assessment storage round-trip (no tool_calls in
|
||||
these tests, so the indexes are unused anyway).
|
||||
"""
|
||||
session = _make_stub_session(messages)
|
||||
with patch(
|
||||
"turnstone.server._load_verdict_indexes",
|
||||
return_value=({}, {}),
|
||||
):
|
||||
return _build_history(session)
|
||||
|
||||
|
||||
class TestSourceSurfacing:
|
||||
def test_source_surfaces_when_set(self) -> None:
|
||||
msg = {
|
||||
"role": "user",
|
||||
"content": "",
|
||||
"_source": "system_nudge",
|
||||
}
|
||||
history = _build([msg])
|
||||
assert len(history) == 1
|
||||
assert history[0]["source"] == "system_nudge"
|
||||
|
||||
def test_source_absent_when_unset(self) -> None:
|
||||
msg = {"role": "user", "content": "hello"}
|
||||
history = _build([msg])
|
||||
assert "source" not in history[0]
|
||||
|
||||
|
||||
class TestRemindersWidening:
|
||||
def test_watch_triggered_optional_fields_propagate(self) -> None:
|
||||
"""The widened payload (Commit 2) carries watch_name / command /
|
||||
poll_count / max_polls / is_final on each ``watch_triggered``
|
||||
reminder so the frontend renders ``.msg.watch-result``.
|
||||
"""
|
||||
msg = {
|
||||
"role": "user",
|
||||
"content": "",
|
||||
"_source": "system_nudge",
|
||||
"_reminders": [
|
||||
{
|
||||
"type": "watch_triggered",
|
||||
"text": "$ ls\nfile.txt",
|
||||
"watch_name": "w1",
|
||||
"command": "ls",
|
||||
"poll_count": 2,
|
||||
"max_polls": 100,
|
||||
"is_final": False,
|
||||
}
|
||||
],
|
||||
}
|
||||
history = _build([msg])
|
||||
assert history[0]["source"] == "system_nudge"
|
||||
assert history[0]["reminders"] == [
|
||||
{
|
||||
"type": "watch_triggered",
|
||||
"text": "$ ls\nfile.txt",
|
||||
"watch_name": "w1",
|
||||
"command": "ls",
|
||||
"poll_count": 2,
|
||||
"max_polls": 100,
|
||||
"is_final": False,
|
||||
}
|
||||
]
|
||||
|
||||
def test_legacy_two_field_reminders_still_work(self) -> None:
|
||||
"""Producers without optional fields (correction / denial /
|
||||
idle_children) keep the legacy ``{type, text}`` shape — the
|
||||
widened filter just doesn't add anything beyond that."""
|
||||
msg = {
|
||||
"role": "user",
|
||||
"content": "noted",
|
||||
"_reminders": [{"type": "correction", "text": "watch out"}],
|
||||
}
|
||||
history = _build([msg])
|
||||
assert history[0]["reminders"] == [{"type": "correction", "text": "watch out"}]
|
||||
|
||||
def test_unknown_keys_are_dropped(self) -> None:
|
||||
"""The wire-layer filter projects on a known set of keys so a
|
||||
future producer accidentally stuffing arbitrary fields can't
|
||||
leak them through replay.
|
||||
"""
|
||||
msg = {
|
||||
"role": "user",
|
||||
"content": "x",
|
||||
"_reminders": [
|
||||
{
|
||||
"type": "correction",
|
||||
"text": "hi",
|
||||
"secret": "leak-me",
|
||||
"internal_id": 42,
|
||||
}
|
||||
],
|
||||
}
|
||||
history = _build([msg])
|
||||
clean = history[0]["reminders"][0]
|
||||
assert "secret" not in clean
|
||||
assert "internal_id" not in clean
|
||||
assert clean == {"type": "correction", "text": "hi"}
|
||||
|
||||
def test_malformed_reminder_skipped(self) -> None:
|
||||
"""A non-dict / empty entry is filtered out instead of breaking
|
||||
the rest of the list (mirrors the defensive filter in
|
||||
``_apply_reminders_for_provider``).
|
||||
"""
|
||||
msg = {
|
||||
"role": "user",
|
||||
"content": "x",
|
||||
"_reminders": [
|
||||
"garbage string",
|
||||
{"type": "", "text": ""}, # empty type + text → drop
|
||||
{"type": "denial", "text": "ok"},
|
||||
],
|
||||
}
|
||||
history = _build([msg])
|
||||
assert history[0]["reminders"] == [{"type": "denial", "text": "ok"}]
|
||||
@@ -0,0 +1,270 @@
|
||||
"""Unit tests for :mod:`turnstone.core.child_source`.
|
||||
|
||||
Covers both strategies in isolation against fakes — no live collector,
|
||||
no live SessionManager. Adapter-level integration coverage continues to
|
||||
live in ``test_coordinator_adapter.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from turnstone.core.child_source import ClusterChildSource, SameNodeChildSource
|
||||
from turnstone.core.children_registry import ChildrenRegistry
|
||||
from turnstone.core.workstream import WorkstreamState
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import queue
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SameNodeChildSource
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeManager:
|
||||
"""Minimal SessionManager stand-in implementing the subscribe API."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.subscribers: list[Any] = []
|
||||
|
||||
def subscribe_to_state(self, callback: Any) -> None:
|
||||
self.subscribers.append(callback)
|
||||
|
||||
def unsubscribe_from_state(self, callback: Any) -> None:
|
||||
with contextlib.suppress(ValueError):
|
||||
self.subscribers.remove(callback)
|
||||
|
||||
def fire(self, ws_id: str, state: WorkstreamState) -> None:
|
||||
for cb in self.subscribers:
|
||||
cb(ws_id, state)
|
||||
|
||||
|
||||
class TestSameNodeChildSource:
|
||||
def test_start_subscribes_to_manager(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
sink_calls: list[dict[str, Any]] = []
|
||||
src.start(sink=sink_calls.append)
|
||||
assert len(mgr.subscribers) == 1
|
||||
|
||||
def test_state_change_for_known_child_pushes_to_sink(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
registry.install("p1", object())
|
||||
registry.add_child("p1", "c1")
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
sink_calls: list[dict[str, Any]] = []
|
||||
src.start(sink=sink_calls.append)
|
||||
|
||||
mgr.fire("c1", WorkstreamState.RUNNING)
|
||||
|
||||
assert len(sink_calls) == 1
|
||||
ev = sink_calls[0]
|
||||
assert ev["type"] == "cluster_state"
|
||||
assert ev["ws_id"] == "c1"
|
||||
assert ev["state"] == "running"
|
||||
|
||||
def test_state_change_for_unknown_workstream_is_dropped(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
sink_calls: list[dict[str, Any]] = []
|
||||
src.start(sink=sink_calls.append)
|
||||
|
||||
# No registry entry — pre-filter drops the event without
|
||||
# invoking the sink.
|
||||
mgr.fire("ws-unknown", WorkstreamState.IDLE)
|
||||
assert sink_calls == []
|
||||
|
||||
def test_shutdown_unsubscribes(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
src.start(sink=lambda ev: None)
|
||||
assert len(mgr.subscribers) == 1
|
||||
src.shutdown()
|
||||
assert mgr.subscribers == []
|
||||
|
||||
def test_start_is_idempotent(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
src.start(sink=lambda ev: None)
|
||||
src.start(sink=lambda ev: None)
|
||||
# Second start is a no-op; only one subscription.
|
||||
assert len(mgr.subscribers) == 1
|
||||
|
||||
def test_sink_exception_does_not_propagate(self) -> None:
|
||||
mgr = _FakeManager()
|
||||
registry = ChildrenRegistry()
|
||||
registry.install("p1", object())
|
||||
registry.add_child("p1", "c1")
|
||||
src = SameNodeChildSource(mgr, registry)
|
||||
|
||||
def bad_sink(ev: dict[str, Any]) -> None:
|
||||
raise RuntimeError("sink boom")
|
||||
|
||||
src.start(sink=bad_sink)
|
||||
# Should not raise — the strategy catches sink failures and logs.
|
||||
mgr.fire("c1", WorkstreamState.RUNNING)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ClusterChildSource
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeCollector:
|
||||
"""Minimal ClusterCollector stand-in providing the listener API."""
|
||||
|
||||
def __init__(self, snapshot: dict[str, Any] | None = None) -> None:
|
||||
self._snapshot = snapshot or {"nodes": []}
|
||||
self.queues: list[queue.Queue[dict[str, Any]]] = []
|
||||
self.unregistered: list[queue.Queue[dict[str, Any]]] = []
|
||||
|
||||
def get_snapshot_and_register(self, q: queue.Queue[dict[str, Any]]) -> dict[str, Any]:
|
||||
self.queues.append(q)
|
||||
return self._snapshot
|
||||
|
||||
def unregister_listener(self, q: queue.Queue[dict[str, Any]]) -> None:
|
||||
self.unregistered.append(q)
|
||||
|
||||
def emit(self, event: dict[str, Any]) -> None:
|
||||
"""Push an event to all registered listener queues."""
|
||||
for q in self.queues:
|
||||
q.put(event)
|
||||
|
||||
|
||||
class TestClusterChildSource:
|
||||
def test_start_subscribes_to_collector(self) -> None:
|
||||
coll = _FakeCollector()
|
||||
registry = ChildrenRegistry()
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=list,
|
||||
)
|
||||
try:
|
||||
src.start(sink=lambda ev: None)
|
||||
assert len(coll.queues) == 1
|
||||
finally:
|
||||
src.shutdown()
|
||||
|
||||
def test_start_primes_registry_from_snapshot(self) -> None:
|
||||
snapshot = {
|
||||
"nodes": [
|
||||
{
|
||||
"workstreams": [
|
||||
{"id": "c1", "parent_ws_id": "p1"},
|
||||
{"id": "c2", "parent_ws_id": "p1"},
|
||||
# Unknown parent — dropped
|
||||
{"id": "x", "parent_ws_id": "p-unknown"},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
coll = _FakeCollector(snapshot)
|
||||
registry = ChildrenRegistry()
|
||||
registry.install("p1", object())
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=lambda: ["p1"],
|
||||
)
|
||||
try:
|
||||
src.start(sink=lambda ev: None)
|
||||
assert set(registry.children_of("p1")) == {"c1", "c2"}
|
||||
assert registry.parent_for("x") is None
|
||||
finally:
|
||||
src.shutdown()
|
||||
|
||||
def test_event_dispatched_to_sink(self) -> None:
|
||||
coll = _FakeCollector()
|
||||
registry = ChildrenRegistry()
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=list,
|
||||
)
|
||||
sink_calls: list[dict[str, Any]] = []
|
||||
|
||||
try:
|
||||
src.start(sink=sink_calls.append)
|
||||
coll.emit({"type": "cluster_state", "ws_id": "c1", "state": "running"})
|
||||
# Daemon thread loop has 1.0s queue timeout; poll briefly.
|
||||
for _ in range(20):
|
||||
if sink_calls:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert len(sink_calls) == 1
|
||||
assert sink_calls[0]["ws_id"] == "c1"
|
||||
finally:
|
||||
src.shutdown()
|
||||
|
||||
def test_shutdown_unregisters_and_joins_thread(self) -> None:
|
||||
coll = _FakeCollector()
|
||||
registry = ChildrenRegistry()
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=list,
|
||||
)
|
||||
src.start(sink=lambda ev: None)
|
||||
src.shutdown()
|
||||
assert coll.unregistered == coll.queues
|
||||
# Second shutdown is a no-op (idempotent).
|
||||
src.shutdown()
|
||||
|
||||
def test_start_is_idempotent(self) -> None:
|
||||
coll = _FakeCollector()
|
||||
registry = ChildrenRegistry()
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=list,
|
||||
)
|
||||
try:
|
||||
src.start(sink=lambda ev: None)
|
||||
src.start(sink=lambda ev: None)
|
||||
assert len(coll.queues) == 1
|
||||
finally:
|
||||
src.shutdown()
|
||||
|
||||
def test_sink_exception_does_not_kill_thread(self) -> None:
|
||||
coll = _FakeCollector()
|
||||
registry = ChildrenRegistry()
|
||||
src = ClusterChildSource(
|
||||
collector=coll,
|
||||
registry=registry,
|
||||
parents_provider=list,
|
||||
)
|
||||
survived_calls: list[dict[str, Any]] = []
|
||||
call_count = [0]
|
||||
|
||||
def flaky_sink(ev: dict[str, Any]) -> None:
|
||||
call_count[0] += 1
|
||||
if call_count[0] == 1:
|
||||
raise RuntimeError("first one boom")
|
||||
survived_calls.append(ev)
|
||||
|
||||
try:
|
||||
src.start(sink=flaky_sink)
|
||||
coll.emit({"type": "cluster_state", "ws_id": "c1", "state": "x"})
|
||||
coll.emit({"type": "cluster_state", "ws_id": "c2", "state": "y"})
|
||||
for _ in range(40):
|
||||
if survived_calls:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert len(survived_calls) == 1
|
||||
assert survived_calls[0]["ws_id"] == "c2"
|
||||
finally:
|
||||
src.shutdown()
|
||||
|
||||
|
||||
# Multi-subscriber observer tests for ``SessionManager.subscribe_to_state``
|
||||
# / ``unsubscribe_from_state`` live in ``test_session_manager.py`` where
|
||||
# the proper FakeAdapter / FakeStorage construction helpers already exist.
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Unit tests for :class:`turnstone.core.children_registry.ChildrenRegistry`.
|
||||
|
||||
The registry was lifted from ``CoordinatorAdapter`` in Stage 3 Step 1.
|
||||
Adapter-level coverage for the integrated behavior already lives in
|
||||
``test_coordinator_adapter.py``; this file pins the data structure
|
||||
invariants in isolation so the registry can be reused by future
|
||||
``ChildSource`` strategies (Step 2) without re-deriving the behavior
|
||||
from the adapter test surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.children_registry import ChildrenRegistry
|
||||
|
||||
|
||||
class _Sentinel:
|
||||
"""Lightweight UI stand-in; identity-comparable, no behavior."""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry() -> ChildrenRegistry:
|
||||
return ChildrenRegistry()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# install / uninstall
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInstallUninstall:
|
||||
def test_install_seeds_empty_child_set_and_presence(self, registry: ChildrenRegistry) -> None:
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
assert registry.children_of("p1") == []
|
||||
assert registry.ui_for("p1") is ui
|
||||
assert registry.parents() == ["p1"]
|
||||
|
||||
def test_install_is_idempotent_repoints_ui_keeps_children(
|
||||
self, registry: ChildrenRegistry
|
||||
) -> None:
|
||||
ui_a = _Sentinel()
|
||||
ui_b = _Sentinel()
|
||||
registry.install("p1", ui_a)
|
||||
registry.merge_children("p1", ["c1", "c2"])
|
||||
registry.install("p1", ui_b)
|
||||
assert registry.ui_for("p1") is ui_b
|
||||
assert set(registry.children_of("p1")) == {"c1", "c2"}
|
||||
|
||||
def test_uninstall_clears_forward_reverse_and_presence(
|
||||
self, registry: ChildrenRegistry
|
||||
) -> None:
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
registry.merge_children("p1", ["c1", "c2"])
|
||||
registry.uninstall("p1")
|
||||
assert registry.children_of("p1") == []
|
||||
assert registry.ui_for("p1") is None
|
||||
assert registry.parents() == []
|
||||
assert registry.parent_for("c1") is None
|
||||
assert registry.parent_for("c2") is None
|
||||
|
||||
def test_uninstall_unknown_parent_is_noop(self, registry: ChildrenRegistry) -> None:
|
||||
registry.uninstall("never-installed") # must not raise
|
||||
|
||||
def test_uninstall_does_not_clobber_other_parents(self, registry: ChildrenRegistry) -> None:
|
||||
registry.install("p1", _Sentinel())
|
||||
registry.install("p2", _Sentinel())
|
||||
registry.merge_children("p1", ["c1"])
|
||||
registry.merge_children("p2", ["c2"])
|
||||
registry.uninstall("p1")
|
||||
assert registry.parent_for("c1") is None
|
||||
assert registry.parent_for("c2") == "p2"
|
||||
assert registry.parents() == ["p2"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# add_child — atomic check-and-route
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAddChild:
|
||||
def test_add_child_returns_ui_on_success(self, registry: ChildrenRegistry) -> None:
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
assert registry.add_child("p1", "c1") is ui
|
||||
assert registry.parent_for("c1") == "p1"
|
||||
assert registry.children_of("p1") == ["c1"]
|
||||
|
||||
def test_add_child_returns_none_when_parent_not_installed(
|
||||
self, registry: ChildrenRegistry
|
||||
) -> None:
|
||||
assert registry.add_child("absent", "c1") is None
|
||||
assert registry.parent_for("c1") is None
|
||||
|
||||
def test_add_child_returns_none_on_duplicate(self, registry: ChildrenRegistry) -> None:
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
assert registry.add_child("p1", "c1") is ui
|
||||
# second add for same child returns None — caller must not
|
||||
# double-dispatch.
|
||||
assert registry.add_child("p1", "c1") is None
|
||||
assert registry.children_of("p1") == ["c1"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# merge_children — bulk seeding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMergeChildren:
|
||||
def test_merge_seeds_forward_and_reverse(self, registry: ChildrenRegistry) -> None:
|
||||
registry.merge_children("p1", ["c1", "c2", "c3"])
|
||||
assert set(registry.children_of("p1")) == {"c1", "c2", "c3"}
|
||||
for cid in ("c1", "c2", "c3"):
|
||||
assert registry.parent_for(cid) == "p1"
|
||||
|
||||
def test_merge_is_idempotent(self, registry: ChildrenRegistry) -> None:
|
||||
registry.merge_children("p1", ["c1"])
|
||||
registry.merge_children("p1", ["c1"])
|
||||
assert registry.children_of("p1") == ["c1"]
|
||||
|
||||
def test_merge_skips_empty_or_falsy_ids(self, registry: ChildrenRegistry) -> None:
|
||||
registry.merge_children("p1", ["", "c1", "", "c2"])
|
||||
assert set(registry.children_of("p1")) == {"c1", "c2"}
|
||||
|
||||
def test_merge_does_not_require_install(self, registry: ChildrenRegistry) -> None:
|
||||
# Snapshot-priming may run before the parent's install fires —
|
||||
# the merge still seeds the forward set so the install picks
|
||||
# the children up. (Storage-seeded rebuild relies on this.)
|
||||
registry.merge_children("p1", ["c1"])
|
||||
assert registry.children_of("p1") == ["c1"]
|
||||
# ui_for is still None because install hasn't run
|
||||
assert registry.ui_for("p1") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lookups — return copies, not live refs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLookups:
|
||||
def test_children_of_returns_copy(self, registry: ChildrenRegistry) -> None:
|
||||
registry.install("p1", _Sentinel())
|
||||
registry.merge_children("p1", ["c1", "c2"])
|
||||
snap = registry.children_of("p1")
|
||||
snap.append("c3-injected")
|
||||
assert "c3-injected" not in registry.children_of("p1")
|
||||
|
||||
def test_children_of_unknown_parent_returns_empty(self, registry: ChildrenRegistry) -> None:
|
||||
assert registry.children_of("absent") == []
|
||||
|
||||
def test_parent_for_unknown_child_returns_none(self, registry: ChildrenRegistry) -> None:
|
||||
assert registry.parent_for("absent") is None
|
||||
|
||||
def test_parents_returns_copy(self, registry: ChildrenRegistry) -> None:
|
||||
registry.install("p1", _Sentinel())
|
||||
snap = registry.parents()
|
||||
snap.append("p2-injected")
|
||||
assert "p2-injected" not in registry.parents()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Concurrency — concurrent add_child must not exceed the unique-set
|
||||
# invariant or leave a half-installed reverse-index entry.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConcurrency:
|
||||
def test_concurrent_add_child_returns_ui_exactly_once_per_unique(
|
||||
self, registry: ChildrenRegistry
|
||||
) -> None:
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
results: list[object] = []
|
||||
results_lock = threading.Lock()
|
||||
|
||||
def attempt_add(child_id: str) -> None:
|
||||
r = registry.add_child("p1", child_id)
|
||||
with results_lock:
|
||||
results.append(r)
|
||||
|
||||
threads = [threading.Thread(target=attempt_add, args=("c1",)) for _ in range(20)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Exactly one thread sees the UI; the remaining 19 see None
|
||||
# (duplicate). The forward + reverse indexes carry exactly one
|
||||
# entry for c1.
|
||||
successes = [r for r in results if r is ui]
|
||||
nones = [r for r in results if r is None]
|
||||
assert len(successes) == 1
|
||||
assert len(nones) == 19
|
||||
assert registry.children_of("p1") == ["c1"]
|
||||
assert registry.parent_for("c1") == "p1"
|
||||
|
||||
def test_concurrent_install_and_add_child_no_resurrect(
|
||||
self, registry: ChildrenRegistry
|
||||
) -> None:
|
||||
# add_child racing with uninstall: either lands first (registry
|
||||
# populated) or the parent is gone (returns None). Must NOT
|
||||
# leave a forward-set entry without presence — that would be
|
||||
# the "resurrected after close" leak the locked dispatch path
|
||||
# was guarding against.
|
||||
ui = _Sentinel()
|
||||
registry.install("p1", ui)
|
||||
|
||||
outcomes: list[object] = []
|
||||
|
||||
def adder() -> None:
|
||||
outcomes.append(registry.add_child("p1", "c1"))
|
||||
|
||||
def uninstaller() -> None:
|
||||
registry.uninstall("p1")
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=adder),
|
||||
threading.Thread(target=uninstaller),
|
||||
]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# If add_child landed first: c1 is in the forward set, then
|
||||
# uninstall clears everything. End state: nothing.
|
||||
# If uninstall landed first: add_child sees no presence,
|
||||
# returns None, no entry added. End state: nothing.
|
||||
# Either way, the leak invariant holds: child set is empty or
|
||||
# parent is gone, never "child set populated but no presence".
|
||||
children = registry.children_of("p1")
|
||||
ui_present = registry.ui_for("p1") is not None
|
||||
if children:
|
||||
assert ui_present, "registry leaked: children set without presence"
|
||||
@@ -30,9 +30,25 @@ def _full_hdr() -> dict[str, str]:
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_metrics(monkeypatch):
|
||||
"""Swap ``turnstone.server._metrics`` for a fresh collector
|
||||
per-test, with auto-restore.
|
||||
|
||||
Bare ``srv_mod._metrics = MetricsCollector()`` (the prior
|
||||
pattern) leaks into any test file that already bound the name
|
||||
via ``from turnstone.server import _metrics`` at import time —
|
||||
those tests' patches then operate on a different instance from
|
||||
the one the live ``_publish_models_metadata`` reads, and the
|
||||
monkeypatch silently no-ops. ``monkeypatch.setattr`` restores
|
||||
after the test, so the leak is contained.
|
||||
"""
|
||||
fresh = MetricsCollector()
|
||||
fresh.model = "test-model"
|
||||
monkeypatch.setattr(srv_mod, "_metrics", fresh)
|
||||
|
||||
|
||||
def _make_app(storage: Any) -> TestClient:
|
||||
srv_mod._metrics = MetricsCollector()
|
||||
srv_mod._metrics.model = "test-model"
|
||||
mock_session = MagicMock()
|
||||
mock_ws = MagicMock()
|
||||
mock_ws.id = "ws-target"
|
||||
|
||||
+279
-19
@@ -314,6 +314,46 @@ class TestCollectorSnapshot:
|
||||
assert event["ws_id"] == "ws1"
|
||||
assert event["state"] == "running"
|
||||
|
||||
def test_apply_snapshot_state_change_does_not_carry_pending_approval_detail(self):
|
||||
"""Stage 3 cleanup — the snapshot-resync cluster_state event no
|
||||
longer piggybacks ``pending_approval_detail`` (the field is
|
||||
gone from cluster_state entirely). On reconnect the browser's
|
||||
bulk fetch — triggered by the ``activity_state="approval"``
|
||||
transition in the reducer — pulls the items directly from
|
||||
``ui.serialize_pending_approval_detail()`` via the dashboard
|
||||
endpoint."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "same", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_snapshot(
|
||||
"node-a",
|
||||
{
|
||||
"type": "node_snapshot",
|
||||
"node_id": "node-a",
|
||||
"workstreams": [
|
||||
{
|
||||
"id": "ws1",
|
||||
"name": "same",
|
||||
"state": "running",
|
||||
"activity_state": "approval",
|
||||
}
|
||||
],
|
||||
"health": {},
|
||||
"aggregate": {},
|
||||
},
|
||||
)
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["type"] == "cluster_state"
|
||||
assert event["activity_state"] == "approval"
|
||||
assert "pending_approval_detail" not in event
|
||||
|
||||
def test_apply_snapshot_skips_empty_id_workstream(self):
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080")
|
||||
@@ -359,6 +399,36 @@ class TestCollectorDelta:
|
||||
# Verify in-memory state was updated
|
||||
assert c._nodes["node-a"].workstreams["ws1"]["state"] == "running"
|
||||
|
||||
def test_apply_delta_ws_state_does_not_carry_pending_approval_detail(self):
|
||||
"""Stage 3 cleanup — ``cluster_state`` no longer carries the
|
||||
``pending_approval_detail`` piggyback. Approval items now arrive
|
||||
via bulk fetch on activity_state transition; verdicts via the
|
||||
explicit ``intent_verdict`` event class. Symmetric event flow,
|
||||
no piggyback to dedupe against."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_delta(
|
||||
"node-a",
|
||||
{
|
||||
"type": "ws_state",
|
||||
"ws_id": "ws1",
|
||||
"state": "running",
|
||||
"activity_state": "approval",
|
||||
},
|
||||
)
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["type"] == "cluster_state"
|
||||
assert event["activity_state"] == "approval"
|
||||
assert "pending_approval_detail" not in event
|
||||
|
||||
def test_apply_delta_ws_created(self):
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080")
|
||||
@@ -404,6 +474,139 @@ class TestCollectorDelta:
|
||||
assert event["name"] == "new-name"
|
||||
assert c._nodes["node-a"].workstreams["ws1"]["name"] == "new-name"
|
||||
|
||||
def test_apply_delta_intent_verdict_forwards_verbatim(self):
|
||||
"""Stage 3 Step 5 — node-emitted intent_verdict events flow
|
||||
through _apply_delta to cluster fan-out so coord adapters can
|
||||
re-emit as child_ws_intent_verdict on the parent's SSE."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
verdict = {
|
||||
"call_id": "c1",
|
||||
"risk_level": "low",
|
||||
"confidence": 0.9,
|
||||
"recommendation": "approve",
|
||||
}
|
||||
c._apply_delta(
|
||||
"node-a",
|
||||
{"type": "intent_verdict", "ws_id": "ws1", "verdict": verdict},
|
||||
)
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["type"] == "intent_verdict"
|
||||
assert event["ws_id"] == "ws1"
|
||||
assert event["node_id"] == "node-a"
|
||||
assert event["verdict"] == verdict
|
||||
|
||||
def test_apply_delta_intent_verdict_drops_when_ws_id_missing(self):
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080")
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_delta("node-a", {"type": "intent_verdict", "verdict": {}})
|
||||
|
||||
assert q.empty()
|
||||
|
||||
def test_apply_delta_approval_resolved_forwards_verbatim(self):
|
||||
"""Stage 3 Step 5 — paired with intent_verdict; clears the
|
||||
coord tree's pending-approval pill in lockstep with the
|
||||
actual decision rather than waiting for the state-change
|
||||
piggyback."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_delta(
|
||||
"node-a",
|
||||
{
|
||||
"type": "approval_resolved",
|
||||
"ws_id": "ws1",
|
||||
"approved": True,
|
||||
"feedback": "lgtm",
|
||||
"always": False,
|
||||
},
|
||||
)
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["type"] == "approval_resolved"
|
||||
assert event["ws_id"] == "ws1"
|
||||
assert event["node_id"] == "node-a"
|
||||
assert event["approved"] is True
|
||||
assert event["feedback"] == "lgtm"
|
||||
assert event["always"] is False
|
||||
|
||||
def test_apply_delta_approve_request_forwards_detail(self):
|
||||
"""Push path for the initial approval items — eliminates the
|
||||
bulk-fetch race that left the coord row stuck on a loading
|
||||
placeholder when the bulk fetch landed in the gap between
|
||||
_emit_state(ATTENTION) and approve_tools setting _pending_approval."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
detail = {
|
||||
"type": "approve_request",
|
||||
"items": [{"call_id": "c1", "header": "tool x"}],
|
||||
"judge_pending": True,
|
||||
}
|
||||
c._apply_delta(
|
||||
"node-a",
|
||||
{"type": "approve_request", "ws_id": "ws1", "detail": detail},
|
||||
)
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["type"] == "approve_request"
|
||||
assert event["ws_id"] == "ws1"
|
||||
assert event["node_id"] == "node-a"
|
||||
assert event["detail"] == detail
|
||||
|
||||
def test_apply_delta_approve_request_drops_when_ws_id_missing(self):
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080")
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_delta("node-a", {"type": "approve_request", "detail": {}})
|
||||
|
||||
assert q.empty()
|
||||
|
||||
def test_apply_delta_approval_resolved_coerces_missing_fields(self):
|
||||
"""Defensive: ``approved`` / ``always`` / ``feedback`` may be
|
||||
omitted by older nodes mid-rolling-upgrade; collector coerces
|
||||
to safe defaults."""
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
node_id="node-a",
|
||||
server_url="http://a:8080",
|
||||
workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}},
|
||||
)
|
||||
q: queue.Queue[dict] = queue.Queue()
|
||||
c.register_listener(q)
|
||||
|
||||
c._apply_delta("node-a", {"type": "approval_resolved", "ws_id": "ws1"})
|
||||
|
||||
event = q.get_nowait()
|
||||
assert event["approved"] is False
|
||||
assert event["feedback"] == ""
|
||||
assert event["always"] is False
|
||||
|
||||
def test_apply_delta_health_changed(self):
|
||||
c = _make_collector()
|
||||
c._nodes["node-a"] = NodeSnapshot(
|
||||
@@ -1309,11 +1512,62 @@ class TestProxyRewriting:
|
||||
assert "window.fetch" in _JS_PROXY_SHIM
|
||||
assert "window.EventSource" in _JS_PROXY_SHIM
|
||||
|
||||
def test_console_banner_contains_placeholder(self):
|
||||
from turnstone.console.server import _CONSOLE_BANNER_TEMPLATE
|
||||
def test_js_shim_carries_node_id_placeholder(self):
|
||||
"""The picker reads the current node_id from the shim's _nodeId
|
||||
closure variable; the placeholder must be present and substitutable."""
|
||||
from turnstone.console.server import _JS_PROXY_SHIM
|
||||
|
||||
assert "NODE_ID_PLACEHOLDER" in _CONSOLE_BANNER_TEMPLATE
|
||||
assert "Console" in _CONSOLE_BANNER_TEMPLATE
|
||||
assert "NODE_ID_PLACEHOLDER" in _JS_PROXY_SHIM
|
||||
replaced = _JS_PROXY_SHIM.replace("NODE_ID_PLACEHOLDER", "node-a")
|
||||
assert "node-a" in replaced
|
||||
assert "NODE_ID_PLACEHOLDER" not in replaced
|
||||
|
||||
def test_js_shim_includes_picker_pieces(self):
|
||||
"""Picker logic ships in the same IIFE as the prefix shim — verify
|
||||
the moving parts are present so a future refactor doesn't silently
|
||||
drop them. /v1/api/cluster/nodes is the lazy-fetch target;
|
||||
#ui-header is the DOM anchor; console-node-pill is the trigger
|
||||
class; ws-tab-dropdown is the menu shell we share with the
|
||||
workstream chevron menu (style + behaviour parity); ArrowDown is
|
||||
the keyboard-nav primitive that disambiguates this from a plain
|
||||
click-only menu."""
|
||||
from turnstone.console.server import _JS_PROXY_SHIM
|
||||
|
||||
# limit=1000 matches the collector's hard cap; without it the
|
||||
# picker would silently drop nodes past the 100-default in
|
||||
# clusters with >100 nodes.
|
||||
assert "/v1/api/cluster/nodes?limit=1000" in _JS_PROXY_SHIM
|
||||
assert "ui-header" in _JS_PROXY_SHIM
|
||||
assert "console-node-pill" in _JS_PROXY_SHIM
|
||||
assert "ws-tab-dropdown" in _JS_PROXY_SHIM
|
||||
assert "ArrowDown" in _JS_PROXY_SHIM
|
||||
assert "DOMContentLoaded" in _JS_PROXY_SHIM
|
||||
|
||||
def test_proxy_style_drops_banner_styles(self):
|
||||
"""The legacy banner CSS classes (.console-banner, .ts-header-back-link
|
||||
offsets, .dashboard-overlay top:32px hack) should be gone — the new
|
||||
picker lives inside #ui-header and doesn't need overlay offsets."""
|
||||
from turnstone.console.server import _CONSOLE_PROXY_STYLE
|
||||
|
||||
assert ".console-banner" not in _CONSOLE_PROXY_STYLE
|
||||
assert "dashboard-overlay" not in _CONSOLE_PROXY_STYLE
|
||||
assert ".console-node-pill" in _CONSOLE_PROXY_STYLE
|
||||
assert ".console-node-menu" in _CONSOLE_PROXY_STYLE
|
||||
|
||||
def test_proxy_style_uses_canonical_degraded_color(self):
|
||||
"""Degraded health dot must use --accent (the canonical "needs
|
||||
attention" token used by the cluster-overview node table at
|
||||
console/static/style.css:548) and not --yellow. Yellow is reserved
|
||||
for the dash-state attention dot, a stronger signal."""
|
||||
from turnstone.console.server import _CONSOLE_PROXY_STYLE
|
||||
|
||||
assert "console-node-menu-item-dot--degraded" in _CONSOLE_PROXY_STYLE
|
||||
# The degraded rule sits on its own line; assert it uses --accent
|
||||
# by checking the CSS substring has --accent and not --yellow.
|
||||
idx = _CONSOLE_PROXY_STYLE.find("console-node-menu-item-dot--degraded")
|
||||
rule = _CONSOLE_PROXY_STYLE[idx : idx + 200]
|
||||
assert "var(--accent)" in rule
|
||||
assert "var(--yellow)" not in rule
|
||||
|
||||
def test_html_rewriting_changes_static_paths(self):
|
||||
"""Simulate the proxy_index rewriting logic."""
|
||||
@@ -1330,16 +1584,24 @@ class TestProxyRewriting:
|
||||
assert 'href="/static/' not in rewritten
|
||||
assert 'src="/static/' not in rewritten
|
||||
|
||||
def test_banner_injection_after_body(self):
|
||||
"""Simulate the banner injection logic."""
|
||||
from turnstone.console.server import _CONSOLE_BANNER_TEMPLATE
|
||||
def test_shim_injection_after_body(self):
|
||||
"""Simulate the proxy shim injection — the shim ships the node-id
|
||||
and prefix as JS literals and renders the picker at runtime, so
|
||||
we assert the substituted JS literals land in the page."""
|
||||
from turnstone.console.server import _CONSOLE_PROXY_STYLE, _JS_PROXY_SHIM
|
||||
|
||||
sample_html = "<html><body><div>content</div></body></html>"
|
||||
banner = _CONSOLE_BANNER_TEMPLATE.replace("NODE_ID_PLACEHOLDER", "node-a")
|
||||
result = sample_html.replace("<body>", "<body>" + banner, 1)
|
||||
assert "node-a" in result
|
||||
assert "Console" in result
|
||||
assert result.startswith("<html><body><div")
|
||||
prefix = "/node/node-a"
|
||||
shim_js = _JS_PROXY_SHIM.replace('"PREFIX_PLACEHOLDER"', json.dumps(prefix)).replace(
|
||||
'"NODE_ID_PLACEHOLDER"', json.dumps("node-a")
|
||||
)
|
||||
injection = _CONSOLE_PROXY_STYLE + "<script>" + shim_js + "</script>"
|
||||
result = sample_html.replace("<body>", "<body>" + injection, 1)
|
||||
assert '"node-a"' in result
|
||||
assert '"/node/node-a"' in result
|
||||
assert "PREFIX_PLACEHOLDER" not in result
|
||||
assert "NODE_ID_PLACEHOLDER" not in result
|
||||
assert result.startswith("<html><body><style>")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1574,17 +1836,15 @@ class TestProxySharedStatic:
|
||||
def test_proxy_shim_injected_in_html(self):
|
||||
"""Verify shim is injected as inline script in proxied HTML."""
|
||||
|
||||
from turnstone.console.server import _CONSOLE_BANNER_TEMPLATE, _JS_PROXY_SHIM
|
||||
from turnstone.console.server import _JS_PROXY_SHIM
|
||||
|
||||
sample_html = "<html><body><div>content</div></body></html>"
|
||||
prefix = "/node/test-node"
|
||||
banner = _CONSOLE_BANNER_TEMPLATE.replace("NODE_ID_PLACEHOLDER", "test-node")
|
||||
shim = (
|
||||
"<script>"
|
||||
+ _JS_PROXY_SHIM.replace('"PREFIX_PLACEHOLDER"', json.dumps(prefix))
|
||||
+ "</script>"
|
||||
shim_js = _JS_PROXY_SHIM.replace('"PREFIX_PLACEHOLDER"', json.dumps(prefix)).replace(
|
||||
'"NODE_ID_PLACEHOLDER"', json.dumps("test-node")
|
||||
)
|
||||
result = sample_html.replace("<body>", "<body>" + banner + shim, 1)
|
||||
shim = "<script>" + shim_js + "</script>"
|
||||
result = sample_html.replace("<body>", "<body>" + shim, 1)
|
||||
assert "<script>" in result
|
||||
assert "/node/test-node" in result
|
||||
assert "window.fetch" in result
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Tests for the console's coordinator idle-cleanup thread helper.
|
||||
|
||||
The helper itself is a tiny loop wrapping ``mgr.close_idle``; the heavy
|
||||
lifting is in ``SessionManager.close_idle`` (covered in
|
||||
``test_session_manager.py``) and ``bulk_close_stale_orphans`` (covered
|
||||
in ``test_storage_sqlite.py``). These tests verify the glue:
|
||||
|
||||
- the helper runs an initial sweep BEFORE its first sleep (cold-start
|
||||
cleanup without blocking the lifespan),
|
||||
- the helper swallows exceptions so a transient DB blip can't kill the
|
||||
daemon thread,
|
||||
- the helper exits cleanly when ``stop_event`` is set.
|
||||
|
||||
The ``stop_event`` parameter is exclusively for tests — production
|
||||
callers pass ``None`` and the daemon runs for process lifetime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from turnstone.console.server import _coord_idle_cleanup_thread
|
||||
|
||||
|
||||
class _StubMgr:
|
||||
def __init__(
|
||||
self, *, stop_event: threading.Event, expected_calls: int, raise_after: int = -1
|
||||
) -> None:
|
||||
self.calls: list[float] = []
|
||||
self.sleep_calls_at_each_close: list[int] = []
|
||||
self._stop_event = stop_event
|
||||
self._expected = expected_calls
|
||||
self._raise_after = raise_after
|
||||
self._sleep_count = 0
|
||||
|
||||
def close_idle(self, timeout_sec: float) -> list[str]:
|
||||
# Snapshot how many sleeps preceded this close — lets the
|
||||
# "initial sweep" test verify the first close_idle ran with
|
||||
# zero preceding sleeps.
|
||||
self.sleep_calls_at_each_close.append(self._sleep_count)
|
||||
self.calls.append(timeout_sec)
|
||||
try:
|
||||
if 0 <= self._raise_after < len(self.calls):
|
||||
raise RuntimeError("simulated DB blip")
|
||||
finally:
|
||||
# Set stop after the helper has been exercised enough,
|
||||
# regardless of whether this call raised.
|
||||
if len(self.calls) >= self._expected:
|
||||
self._stop_event.set()
|
||||
return []
|
||||
|
||||
def record_sleep(self, _seconds: float) -> None:
|
||||
self._sleep_count += 1
|
||||
|
||||
|
||||
def _run_until_done(mgr: _StubMgr, stop_event: threading.Event, timeout_sec: float) -> None:
|
||||
with patch("turnstone.console.server.time.sleep", mgr.record_sleep):
|
||||
thread = threading.Thread(
|
||||
target=_coord_idle_cleanup_thread,
|
||||
args=(mgr, timeout_sec, stop_event),
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
thread.join(timeout=2.0)
|
||||
assert not thread.is_alive(), "helper failed to exit on stop_event"
|
||||
|
||||
|
||||
def test_coord_idle_cleanup_runs_initial_sweep_before_sleep() -> None:
|
||||
"""The first close_idle call must happen BEFORE the first time.sleep —
|
||||
otherwise cold-start orphans wait one ``check_every`` interval (~30 min
|
||||
on default 2h timeout) for the first reap. Crucial because the
|
||||
lifespan no longer does a synchronous initial sweep."""
|
||||
stop_event = threading.Event()
|
||||
mgr = _StubMgr(stop_event=stop_event, expected_calls=1)
|
||||
_run_until_done(mgr, stop_event, timeout_sec=120.0)
|
||||
assert mgr.sleep_calls_at_each_close == [0], "first close_idle should run before any sleep"
|
||||
|
||||
|
||||
def test_coord_idle_cleanup_calls_close_idle_each_tick() -> None:
|
||||
stop_event = threading.Event()
|
||||
mgr = _StubMgr(stop_event=stop_event, expected_calls=3)
|
||||
_run_until_done(mgr, stop_event, timeout_sec=120.0)
|
||||
assert len(mgr.calls) == 3
|
||||
assert all(t == 120.0 for t in mgr.calls)
|
||||
|
||||
|
||||
def test_coord_idle_cleanup_survives_close_idle_exceptions() -> None:
|
||||
"""A transient DB error must not kill the daemon thread — the next
|
||||
tick should still fire close_idle. Without the try/except, a single
|
||||
blip would silently leak orphans forever."""
|
||||
stop_event = threading.Event()
|
||||
mgr = _StubMgr(stop_event=stop_event, expected_calls=4, raise_after=1)
|
||||
_run_until_done(mgr, stop_event, timeout_sec=120.0)
|
||||
# All four calls must have fired despite calls 2-4 raising.
|
||||
assert len(mgr.calls) == 4
|
||||
|
||||
|
||||
def test_coord_idle_cleanup_exits_cleanly_on_stop_event() -> None:
|
||||
"""The stop_event mechanism is the test contract; verify the thread
|
||||
actually exits when the event is set, without needing exceptions or
|
||||
daemon-process termination."""
|
||||
stop_event = threading.Event()
|
||||
mgr = _StubMgr(stop_event=stop_event, expected_calls=2)
|
||||
_run_until_done(mgr, stop_event, timeout_sec=120.0)
|
||||
assert stop_event.is_set()
|
||||
@@ -464,3 +464,148 @@ def test_coord_budget_override_survives_wildcard_allow_policy() -> None:
|
||||
assert approved is True
|
||||
types = [e.get("type") for e in captured_events]
|
||||
assert "approve_request" in types, "Wildcard allow must not strip the budget-override prompt"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cluster-bus broadcast hooks — _broadcast_intent_verdict / _approval_resolved
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBroadcastIntentVerdict:
|
||||
"""``ConsoleCoordinatorUI._broadcast_intent_verdict`` overrides the
|
||||
no-op base hook to push the verdict onto the cluster bus via
|
||||
``ClusterCollector.emit_console_ws_intent_verdict``. The far more
|
||||
common path is the per-node ``WebUI`` override (covered in
|
||||
test_webui_content.py); this lights up the rare coord-self path
|
||||
(a coord that runs its own LLM judge).
|
||||
"""
|
||||
|
||||
def test_calls_collector_emit_with_ws_id_and_verdict(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
verdict = {
|
||||
"call_id": "c1",
|
||||
"risk_level": "high",
|
||||
"confidence": 0.91,
|
||||
}
|
||||
ui._broadcast_intent_verdict(verdict)
|
||||
collector.emit_console_ws_intent_verdict.assert_called_once_with(
|
||||
"coord-a",
|
||||
verdict,
|
||||
)
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
def test_no_op_when_collector_unset(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
# Doesn't raise.
|
||||
ui._broadcast_intent_verdict({"call_id": "c1"})
|
||||
|
||||
def test_collector_exception_swallowed(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
collector.emit_console_ws_intent_verdict.side_effect = RuntimeError("boom")
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
# Doesn't raise — collector failures are observational only.
|
||||
ui._broadcast_intent_verdict({"call_id": "c1"})
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
|
||||
class TestBroadcastApprovalResolved:
|
||||
"""``ConsoleCoordinatorUI._broadcast_approval_resolved`` overrides
|
||||
the base hook to push the resolution onto the cluster bus via
|
||||
``ClusterCollector.emit_console_ws_approval_resolved``."""
|
||||
|
||||
def test_calls_collector_with_decision_fields(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
ui._broadcast_approval_resolved(True, "lgtm", always=True)
|
||||
collector.emit_console_ws_approval_resolved.assert_called_once_with(
|
||||
"coord-a",
|
||||
approved=True,
|
||||
feedback="lgtm",
|
||||
always=True,
|
||||
)
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
def test_normalises_none_feedback_to_empty_string(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
ui._broadcast_approval_resolved(False, None)
|
||||
collector.emit_console_ws_approval_resolved.assert_called_once_with(
|
||||
"coord-a",
|
||||
approved=False,
|
||||
feedback="",
|
||||
always=False,
|
||||
)
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
def test_no_op_when_collector_unset(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approval_resolved(True, None)
|
||||
|
||||
def test_collector_exception_swallowed(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
collector.emit_console_ws_approval_resolved.side_effect = RuntimeError("boom")
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approval_resolved(True, "ok")
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
|
||||
class TestBroadcastApproveRequest:
|
||||
"""Coord-side override for the approve_request push. Same rationale
|
||||
as the WebUI override — the coord-self path is rare today, but
|
||||
parity keeps the override symmetric with the rest of the broadcast
|
||||
family."""
|
||||
|
||||
def test_calls_collector_emit_with_ws_id_and_detail(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
detail = {
|
||||
"type": "approve_request",
|
||||
"items": [{"call_id": "c1", "header": "tool x"}],
|
||||
"judge_pending": True,
|
||||
}
|
||||
ui._broadcast_approve_request(detail)
|
||||
collector.emit_console_ws_approve_request.assert_called_once_with(
|
||||
"coord-a",
|
||||
detail,
|
||||
)
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
def test_no_op_when_collector_unset(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approve_request({"items": []})
|
||||
|
||||
def test_collector_exception_swallowed(self) -> None:
|
||||
ui = ConsoleCoordinatorUI(ws_id="coord-a", user_id="u1")
|
||||
collector = MagicMock()
|
||||
collector.emit_console_ws_approve_request.side_effect = RuntimeError("boom")
|
||||
ConsoleCoordinatorUI._collector = collector
|
||||
try:
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approve_request({"items": []})
|
||||
finally:
|
||||
ConsoleCoordinatorUI._collector = None
|
||||
|
||||
@@ -378,13 +378,21 @@ class TestCoordinatorAdapterWorkerDispatch:
|
||||
|
||||
|
||||
class TestCoordinatorAdapterChildrenRegistry:
|
||||
def test_emit_created_seeds_empty_children_set(self) -> None:
|
||||
"""Adapter-level integration with :class:`ChildrenRegistry`.
|
||||
|
||||
Pure-registry invariants (forward/reverse consistency, idempotent
|
||||
merge, locking) live in ``test_children_registry.py``. These
|
||||
tests cover the adapter's wiring: that ``emit_*`` paths drive the
|
||||
registry correctly and that the snapshot-priming bridge between
|
||||
a collector snapshot and the registry preserves merge semantics.
|
||||
"""
|
||||
|
||||
def test_emit_created_installs_parent(self) -> None:
|
||||
adapter, _ = _make_adapter()
|
||||
ws = _make_ws()
|
||||
adapter.emit_created(ws)
|
||||
assert ws.id in adapter._children
|
||||
assert adapter._children[ws.id] == set()
|
||||
assert adapter._active_coords[ws.id] is ws.ui
|
||||
assert adapter._registry.children_of(ws.id) == []
|
||||
assert adapter._registry.ui_for(ws.id) is ws.ui
|
||||
|
||||
def test_emit_rehydrated_calls_rebuild(self) -> None:
|
||||
adapter, _ = _make_adapter()
|
||||
@@ -398,42 +406,36 @@ class TestCoordinatorAdapterChildrenRegistry:
|
||||
adapter.emit_rehydrated(ws)
|
||||
assert calls == [ws.id]
|
||||
|
||||
def test_emit_closed_clears_forward_and_reverse_indexes(self) -> None:
|
||||
def test_emit_closed_uninstalls_parent_and_clears_children(self) -> None:
|
||||
adapter, _ = _make_adapter()
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-a1", "child-a2"])
|
||||
adapter._merge_child_ids_locked("coord-b", ["child-b1"])
|
||||
adapter._active_coords["coord-a"] = object()
|
||||
adapter._active_coords["coord-b"] = object()
|
||||
adapter._registry.install("coord-a", object())
|
||||
adapter._registry.install("coord-b", object())
|
||||
adapter._registry.merge_children("coord-a", ["child-a1", "child-a2"])
|
||||
adapter._registry.merge_children("coord-b", ["child-b1"])
|
||||
|
||||
adapter.emit_closed("coord-a")
|
||||
|
||||
assert "coord-a" not in adapter._children
|
||||
assert "coord-a" not in adapter._active_coords
|
||||
assert "child-a1" not in adapter._child_to_coord
|
||||
assert "child-a2" not in adapter._child_to_coord
|
||||
assert adapter._registry.ui_for("coord-a") is None
|
||||
assert adapter._registry.children_of("coord-a") == []
|
||||
assert adapter._registry.parent_for("child-a1") is None
|
||||
assert adapter._registry.parent_for("child-a2") is None
|
||||
# coord-b untouched
|
||||
assert adapter._child_to_coord["child-b1"] == "coord-b"
|
||||
assert "coord-b" in adapter._children
|
||||
|
||||
def test_merge_child_ids_locked_is_idempotent(self) -> None:
|
||||
adapter, _ = _make_adapter()
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-1"])
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-1"])
|
||||
assert adapter._children["coord-a"] == {"child-1"}
|
||||
assert adapter._child_to_coord == {"child-1": "coord-a"}
|
||||
assert adapter._registry.parent_for("child-b1") == "coord-b"
|
||||
assert adapter._registry.ui_for("coord-b") is not None
|
||||
|
||||
def test_prime_children_from_snapshot_merges_without_overwriting(self) -> None:
|
||||
# Snapshot priming now lives on ClusterChildSource (production
|
||||
# path). The adapter no longer carries its own duplicate copy.
|
||||
from turnstone.core.child_source import ClusterChildSource
|
||||
|
||||
adapter, _ = _make_adapter()
|
||||
# Seed one in-memory coord + one existing child
|
||||
coord_ws = _make_ws()
|
||||
coord_ws.id = "coord-a"
|
||||
mgr = MagicMock()
|
||||
mgr.list_all.return_value = [coord_ws]
|
||||
adapter.attach(mgr)
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-a1"])
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
|
||||
source = ClusterChildSource(
|
||||
collector=MagicMock(),
|
||||
registry=adapter._registry,
|
||||
parents_provider=lambda: ["coord-a"],
|
||||
)
|
||||
|
||||
snapshot = {
|
||||
"nodes": [
|
||||
@@ -448,10 +450,13 @@ class TestCoordinatorAdapterChildrenRegistry:
|
||||
},
|
||||
],
|
||||
}
|
||||
adapter._prime_children_from_snapshot(snapshot)
|
||||
assert adapter._children["coord-a"] == {"child-a1", "child-a2"}
|
||||
assert adapter._child_to_coord["child-a2"] == "coord-a"
|
||||
assert "child-x" not in adapter._child_to_coord
|
||||
source._prime_from_snapshot(snapshot)
|
||||
assert set(adapter._registry.children_of("coord-a")) == {
|
||||
"child-a1",
|
||||
"child-a2",
|
||||
}
|
||||
assert adapter._registry.parent_for("child-a2") == "coord-a"
|
||||
assert adapter._registry.parent_for("child-x") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -478,9 +483,7 @@ class TestCoordinatorAdapterDispatchChildEvent:
|
||||
coord_ws.id = coord_id
|
||||
recorder = _UIRecorder()
|
||||
coord_ws.ui = recorder # type: ignore[assignment]
|
||||
with adapter._children_lock:
|
||||
adapter._children.setdefault(coord_id, set())
|
||||
adapter._active_coords[coord_id] = recorder
|
||||
adapter._registry.install(coord_id, recorder)
|
||||
adapter.attach(_StubManager(coord_ws)) # type: ignore[arg-type]
|
||||
return adapter, recorder, coord_ws
|
||||
|
||||
@@ -510,12 +513,11 @@ class TestCoordinatorAdapterDispatchChildEvent:
|
||||
assert payload["child_ws_id"] == "child-a1"
|
||||
assert payload["parent_ws_id"] == "coord-a"
|
||||
# Reverse index updated for subsequent cluster_state events.
|
||||
assert adapter._child_to_coord["child-a1"] == "coord-a"
|
||||
assert adapter._registry.parent_for("child-a1") == "coord-a"
|
||||
|
||||
def test_dispatch_cluster_state_routes_via_reverse_index(self) -> None:
|
||||
adapter, recorder, _ = self._setup()
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-a1"])
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "cluster_state",
|
||||
@@ -533,8 +535,7 @@ class TestCoordinatorAdapterDispatchChildEvent:
|
||||
|
||||
def test_dispatch_ws_closed_routes_to_parent_coord(self) -> None:
|
||||
adapter, recorder, _ = self._setup()
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-a1"])
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event(
|
||||
{"type": "ws_closed", "ws_id": "child-a1", "reason": "evicted"}
|
||||
)
|
||||
@@ -548,8 +549,7 @@ class TestCoordinatorAdapterDispatchChildEvent:
|
||||
"""perf-6: _enqueue_on_ui mutates the payload dict in place with
|
||||
the coord's ws_id so the browser can discriminate child events."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
with adapter._children_lock:
|
||||
adapter._merge_child_ids_locked("coord-a", ["child-a1"])
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "cluster_state",
|
||||
@@ -558,3 +558,161 @@ class TestCoordinatorAdapterDispatchChildEvent:
|
||||
}
|
||||
)
|
||||
assert recorder.enqueued[0]["ws_id"] == "coord-a"
|
||||
|
||||
def test_dispatch_cluster_state_does_not_carry_pending_approval_detail(
|
||||
self,
|
||||
) -> None:
|
||||
"""Stage 3 cleanup — the ``pending_approval_detail`` piggyback
|
||||
on ``cluster_state`` is gone. Approval items now arrive via
|
||||
bulk fetch (triggered by ``activity_state="approval"`` in the
|
||||
browser); verdicts via ``child_ws_intent_verdict``; resolution
|
||||
via ``child_ws_approval_resolved``. The state event carries
|
||||
only state + activity_state — no detail field."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "cluster_state",
|
||||
"ws_id": "child-a1",
|
||||
"state": "running",
|
||||
"activity_state": "approval",
|
||||
}
|
||||
)
|
||||
assert len(recorder.enqueued) == 1
|
||||
payload = recorder.enqueued[0]
|
||||
assert payload["type"] == "child_ws_state"
|
||||
assert payload["activity_state"] == "approval"
|
||||
assert "pending_approval_detail" not in payload
|
||||
|
||||
def test_dispatch_intent_verdict_emits_child_ws_intent_verdict(self) -> None:
|
||||
"""Stage 3 Step 6 — explicit verdict events are re-emitted as
|
||||
child_ws_intent_verdict on the parent's SSE so the tree UI
|
||||
renders the risk pill without polling."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
verdict = {
|
||||
"call_id": "c1",
|
||||
"risk_level": "low",
|
||||
"confidence": 0.92,
|
||||
"recommendation": "approve",
|
||||
}
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "intent_verdict",
|
||||
"ws_id": "child-a1",
|
||||
"node_id": "node-1",
|
||||
"verdict": verdict,
|
||||
}
|
||||
)
|
||||
assert len(recorder.enqueued) == 1
|
||||
payload = recorder.enqueued[0]
|
||||
assert payload["type"] == "child_ws_intent_verdict"
|
||||
assert payload["child_ws_id"] == "child-a1"
|
||||
assert payload["parent_ws_id"] == "coord-a"
|
||||
assert payload["node_id"] == "node-1"
|
||||
assert payload["verdict"] == verdict
|
||||
|
||||
def test_dispatch_intent_verdict_unknown_child_drops(self) -> None:
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "intent_verdict",
|
||||
"ws_id": "ws-orphan",
|
||||
"verdict": {"call_id": "c1"},
|
||||
}
|
||||
)
|
||||
assert recorder.enqueued == []
|
||||
|
||||
def test_dispatch_approval_resolved_emits_child_ws_approval_resolved(
|
||||
self,
|
||||
) -> None:
|
||||
"""Stage 3 Step 6 — paired with intent_verdict; clears the
|
||||
pending-approval pill on the parent's tree UI in lockstep
|
||||
with the actual decision."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "approval_resolved",
|
||||
"ws_id": "child-a1",
|
||||
"node_id": "node-1",
|
||||
"approved": True,
|
||||
"feedback": "lgtm",
|
||||
"always": False,
|
||||
}
|
||||
)
|
||||
assert len(recorder.enqueued) == 1
|
||||
payload = recorder.enqueued[0]
|
||||
assert payload["type"] == "child_ws_approval_resolved"
|
||||
assert payload["child_ws_id"] == "child-a1"
|
||||
assert payload["parent_ws_id"] == "coord-a"
|
||||
assert payload["approved"] is True
|
||||
assert payload["feedback"] == "lgtm"
|
||||
assert payload["always"] is False
|
||||
|
||||
def test_dispatch_approval_resolved_coerces_missing_fields(self) -> None:
|
||||
"""Older nodes mid-rolling-upgrade may omit approved / always /
|
||||
feedback; dispatch coerces to safe defaults."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
adapter._dispatch_child_event({"type": "approval_resolved", "ws_id": "child-a1"})
|
||||
assert len(recorder.enqueued) == 1
|
||||
payload = recorder.enqueued[0]
|
||||
assert payload["approved"] is False
|
||||
assert payload["feedback"] == ""
|
||||
assert payload["always"] is False
|
||||
|
||||
def test_dispatch_approval_resolved_unknown_child_drops(self) -> None:
|
||||
"""Symmetric to the intent_verdict drop test — events for
|
||||
ws_ids the registry doesn't know about silently drop instead
|
||||
of fanning out to a parent that has no business seeing them."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "approval_resolved",
|
||||
"ws_id": "ws-orphan",
|
||||
"approved": True,
|
||||
},
|
||||
)
|
||||
assert recorder.enqueued == []
|
||||
|
||||
def test_dispatch_approve_request_emits_child_ws_approve_request(
|
||||
self,
|
||||
) -> None:
|
||||
"""Push path for the initial approval items — eliminates the
|
||||
bulk-fetch race that left the coord row stuck on a loading
|
||||
placeholder when the bulk fetch landed in the gap between
|
||||
_emit_state(ATTENTION) and approve_tools setting _pending_approval."""
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._registry.merge_children("coord-a", ["child-a1"])
|
||||
detail = {
|
||||
"type": "approve_request",
|
||||
"items": [{"call_id": "c1", "header": "tool x"}],
|
||||
"judge_pending": True,
|
||||
}
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "approve_request",
|
||||
"ws_id": "child-a1",
|
||||
"node_id": "node-1",
|
||||
"detail": detail,
|
||||
},
|
||||
)
|
||||
assert len(recorder.enqueued) == 1
|
||||
payload = recorder.enqueued[0]
|
||||
assert payload["type"] == "child_ws_approve_request"
|
||||
assert payload["child_ws_id"] == "child-a1"
|
||||
assert payload["parent_ws_id"] == "coord-a"
|
||||
assert payload["node_id"] == "node-1"
|
||||
assert payload["detail"] == detail
|
||||
|
||||
def test_dispatch_approve_request_unknown_child_drops(self) -> None:
|
||||
adapter, recorder, _ = self._setup()
|
||||
adapter._dispatch_child_event(
|
||||
{
|
||||
"type": "approve_request",
|
||||
"ws_id": "ws-orphan",
|
||||
"detail": {"items": []},
|
||||
},
|
||||
)
|
||||
assert recorder.enqueued == []
|
||||
|
||||
@@ -949,6 +949,137 @@ def test_list_nodes_empty_on_no_matching_filters(storage_with_nodes):
|
||||
assert result["truncated"] is False
|
||||
|
||||
|
||||
def test_list_nodes_surfaces_healthy_model_aliases(tmp_path):
|
||||
"""The node's heartbeat loop projects its registry into a ``models``
|
||||
metadata entry shaped like ``[{alias, provider, healthy}, ...]``.
|
||||
``list_nodes`` flattens that to the healthy-alias list at the top
|
||||
level (under ``model_aliases``) so a coordinator can pass aliases
|
||||
straight to ``spawn_workstream(model=)`` without having to
|
||||
introspect the metadata blob. The provider-side model identifier
|
||||
(``cfg.model``) is intentionally NOT in the payload — coords kept
|
||||
reaching for it when they should pass the local alias."""
|
||||
st = SQLiteBackend(str(tmp_path / "nodes.db"))
|
||||
_set_meta(
|
||||
st,
|
||||
"node-x",
|
||||
[
|
||||
("arch", "x86_64", "auto"),
|
||||
(
|
||||
"models",
|
||||
[
|
||||
{"alias": "gpt5", "provider": "openai", "healthy": True},
|
||||
{"alias": "claude-opus-47", "provider": "anthropic", "healthy": True},
|
||||
{"alias": "broken", "provider": "openai", "healthy": False},
|
||||
],
|
||||
"auto",
|
||||
),
|
||||
],
|
||||
)
|
||||
_register_service(st, "node-x")
|
||||
client = _make_read_client(st)
|
||||
result = client.list_nodes()
|
||||
node = result["nodes"][0]
|
||||
assert node["model_aliases"] == ["gpt5", "claude-opus-47"]
|
||||
# Full per-alias info still available under metadata for callers
|
||||
# that want provider / healthy detail (e.g. surfacing degraded
|
||||
# aliases in a UI).
|
||||
full = node["metadata"]["models"]["value"]
|
||||
assert {row["alias"] for row in full} == {"gpt5", "claude-opus-47", "broken"}
|
||||
# ``model`` (the provider-side identifier) is intentionally absent
|
||||
# — keep the payload to the three values a coord actually uses.
|
||||
for row in full:
|
||||
assert "model" not in row
|
||||
|
||||
|
||||
def test_list_nodes_model_aliases_distinct_from_metadata_models(tmp_path):
|
||||
"""Pin the naming distinction explicitly: the top-level shortlist
|
||||
(``model_aliases``, list of strings) and the rich metadata blob
|
||||
(``metadata.models.value``, list of dicts) live under different
|
||||
keys so a caller that confuses them gets a clear KeyError rather
|
||||
than a silent shape mismatch."""
|
||||
st = SQLiteBackend(str(tmp_path / "nodes.db"))
|
||||
_set_meta(
|
||||
st,
|
||||
"node-x",
|
||||
[
|
||||
(
|
||||
"models",
|
||||
[{"alias": "a", "provider": "openai", "healthy": True}],
|
||||
"auto",
|
||||
),
|
||||
],
|
||||
)
|
||||
_register_service(st, "node-x")
|
||||
client = _make_read_client(st)
|
||||
node = client.list_nodes()["nodes"][0]
|
||||
# No top-level ``models`` field — only ``model_aliases``.
|
||||
assert "models" not in node
|
||||
assert node["model_aliases"] == ["a"]
|
||||
# Rich shape stays under metadata.
|
||||
assert isinstance(node["metadata"]["models"]["value"], list)
|
||||
assert isinstance(node["metadata"]["models"]["value"][0], dict)
|
||||
|
||||
|
||||
def test_list_nodes_model_aliases_empty_when_node_has_not_published(tmp_path):
|
||||
"""Nodes from older builds — or a node mid-startup before its first
|
||||
metadata write — won't have a ``models`` entry. The top-level
|
||||
``model_aliases`` field defaults to ``[]`` rather than being
|
||||
omitted so coordinators can rely on the key being present."""
|
||||
st = SQLiteBackend(str(tmp_path / "nodes.db"))
|
||||
_set_meta(st, "node-y", [("arch", "x86_64", "auto")])
|
||||
_register_service(st, "node-y")
|
||||
client = _make_read_client(st)
|
||||
result = client.list_nodes()
|
||||
assert result["nodes"][0]["model_aliases"] == []
|
||||
|
||||
|
||||
def test_list_nodes_models_tolerates_malformed_entries(tmp_path):
|
||||
"""If a node ever stores a malformed ``models`` entry (wrong outer
|
||||
type, missing alias, non-bool healthy), the projection drops the
|
||||
bad rows rather than raising — the rest of the response should
|
||||
still be useful."""
|
||||
st = SQLiteBackend(str(tmp_path / "nodes.db"))
|
||||
_set_meta(
|
||||
st,
|
||||
"node-z",
|
||||
[
|
||||
(
|
||||
"models",
|
||||
[
|
||||
{"alias": "ok", "provider": "p", "healthy": True},
|
||||
"not-a-dict",
|
||||
{"provider": "p", "healthy": True}, # missing alias
|
||||
{"alias": "", "healthy": True}, # empty alias
|
||||
{"alias": "degraded", "healthy": False},
|
||||
{"alias": 42, "healthy": True}, # non-string alias
|
||||
],
|
||||
"auto",
|
||||
),
|
||||
],
|
||||
)
|
||||
_register_service(st, "node-z")
|
||||
client = _make_read_client(st)
|
||||
result = client.list_nodes()
|
||||
assert result["nodes"][0]["model_aliases"] == ["ok"]
|
||||
|
||||
|
||||
def test_list_nodes_models_handles_non_list_payload(tmp_path):
|
||||
"""A node with a corrupted models entry (dict, scalar, null) shouldn't
|
||||
blow up the whole list_nodes call. ``model_aliases`` falls back to ``[]``."""
|
||||
st = SQLiteBackend(str(tmp_path / "nodes.db"))
|
||||
_set_meta(
|
||||
st,
|
||||
"node-w",
|
||||
[
|
||||
("models", {"oops": "not a list"}, "auto"),
|
||||
],
|
||||
)
|
||||
_register_service(st, "node-w")
|
||||
client = _make_read_client(st)
|
||||
result = client.list_nodes()
|
||||
assert result["nodes"][0]["model_aliases"] == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# list_skills
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,489 @@
|
||||
"""Unit tests for :class:`CoordinatorIdleObserver`.
|
||||
|
||||
Drives a fake :class:`SessionManager` that mirrors the real one's
|
||||
``subscribe_to_state`` / ``get`` contract, plus a fake storage with the
|
||||
``list_workstreams`` slice the observer queries.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import threading
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.console.coordinator_idle_observer import CoordinatorIdleObserver
|
||||
from turnstone.core.nudge_queue import NudgeQueue
|
||||
from turnstone.core.workstream import WorkstreamKind, WorkstreamState
|
||||
|
||||
|
||||
class _FakeRow:
|
||||
"""SQLAlchemy-Row-like wrapper exposing ``_mapping``."""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self._mapping = kwargs
|
||||
|
||||
|
||||
class _FakeStorage:
|
||||
def __init__(self) -> None:
|
||||
self.children: list[dict[str, Any]] = []
|
||||
self.list_calls: list[dict[str, Any]] = []
|
||||
self.count_calls: list[dict[str, Any]] = []
|
||||
self.list_raises: bool = False
|
||||
self.count_raises: bool = False
|
||||
|
||||
def list_workstreams(
|
||||
self,
|
||||
node_id: str | None = None,
|
||||
limit: int = 100,
|
||||
*,
|
||||
parent_ws_id: str | None = None,
|
||||
kind: WorkstreamKind | str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> list[Any]:
|
||||
self.list_calls.append(
|
||||
{
|
||||
"limit": limit,
|
||||
"parent_ws_id": parent_ws_id,
|
||||
"kind": kind,
|
||||
"user_id": user_id,
|
||||
}
|
||||
)
|
||||
if self.list_raises:
|
||||
raise RuntimeError("storage forced failure")
|
||||
return [_FakeRow(**c) for c in self.children]
|
||||
|
||||
def count_workstreams_by_state(
|
||||
self,
|
||||
*,
|
||||
parent_ws_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
self.count_calls.append({"parent_ws_id": parent_ws_id, "user_id": user_id})
|
||||
if self.count_raises:
|
||||
raise RuntimeError("count forced failure")
|
||||
counts: dict[str, int] = {}
|
||||
for c in self.children:
|
||||
counts[c["state"]] = counts.get(c["state"], 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self._nudge_queue = NudgeQueue()
|
||||
self.messages: list[dict[str, Any]] = []
|
||||
self._wake_source_tag: str = ""
|
||||
self._metacog_state: dict[str, float] = {}
|
||||
self._mem_cfg = MagicMock(nudge_cooldown=300)
|
||||
|
||||
def _visible_memory_count(self) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
class _FakeWorkstream:
|
||||
def __init__(
|
||||
self,
|
||||
ws_id: str = "ws-coord",
|
||||
kind: WorkstreamKind = WorkstreamKind.COORDINATOR,
|
||||
user_id: str = "u1",
|
||||
) -> None:
|
||||
self.id = ws_id
|
||||
self.kind = kind
|
||||
self.user_id = user_id
|
||||
self.session: _FakeSession | None = _FakeSession()
|
||||
|
||||
|
||||
class _FakeManager:
|
||||
def __init__(self) -> None:
|
||||
self._workstreams: dict[str, _FakeWorkstream] = {}
|
||||
self._subscribers: list[Any] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def add_ws(self, ws: _FakeWorkstream) -> None:
|
||||
self._workstreams[ws.id] = ws
|
||||
|
||||
def remove_ws(self, ws_id: str) -> None:
|
||||
self._workstreams.pop(ws_id, None)
|
||||
|
||||
def get(self, ws_id: str) -> _FakeWorkstream | None:
|
||||
return self._workstreams.get(ws_id)
|
||||
|
||||
def subscribe_to_state(self, callback: Any) -> None:
|
||||
with self._lock:
|
||||
self._subscribers.append(callback)
|
||||
|
||||
def unsubscribe_from_state(self, callback: Any) -> None:
|
||||
with self._lock, contextlib.suppress(ValueError):
|
||||
self._subscribers.remove(callback)
|
||||
|
||||
def fire_state(self, ws_id: str, state: WorkstreamState) -> None:
|
||||
with self._lock:
|
||||
subs = list(self._subscribers)
|
||||
for cb in subs:
|
||||
with contextlib.suppress(Exception):
|
||||
cb(ws_id, state)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def coord_setup() -> tuple[_FakeManager, _FakeStorage, _FakeWorkstream]:
|
||||
mgr = _FakeManager()
|
||||
storage = _FakeStorage()
|
||||
ws = _FakeWorkstream()
|
||||
mgr.add_ws(ws)
|
||||
return mgr, storage, ws
|
||||
|
||||
|
||||
def _add_active_child(storage: _FakeStorage, **overrides: Any) -> None:
|
||||
storage.children.append(
|
||||
{
|
||||
"ws_id": overrides.get("ws_id", "child-1"),
|
||||
"name": overrides.get("name", "research"),
|
||||
"state": overrides.get("state", "running"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TestEnqueueOnIdle:
|
||||
def test_idle_with_active_children_enqueues(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage, ws_id="child-a", state="running")
|
||||
_add_active_child(storage, ws_id="child-b", state="thinking")
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
|
||||
snap = ws.session._nudge_queue.pending("any")
|
||||
assert len(snap) == 1
|
||||
nudge_type, text = snap[0]
|
||||
assert nudge_type == "idle_children"
|
||||
assert "child-a" in text
|
||||
assert "child-b" in text
|
||||
|
||||
def test_idle_with_no_active_children_no_enqueue(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
# storage.children is empty
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
def test_idle_only_idle_state_children_no_enqueue(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
# All children "idle" — terminal-from-coord-perspective; not active.
|
||||
_add_active_child(storage, state="idle")
|
||||
_add_active_child(storage, state="closed")
|
||||
_add_active_child(storage, state="error")
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
def test_non_idle_state_no_enqueue(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
for state in (
|
||||
WorkstreamState.RUNNING,
|
||||
WorkstreamState.THINKING,
|
||||
WorkstreamState.ATTENTION,
|
||||
WorkstreamState.ERROR,
|
||||
):
|
||||
mgr.fire_state(ws.id, state)
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
|
||||
class TestKindFilter:
|
||||
def test_interactive_workstream_skipped(self):
|
||||
mgr = _FakeManager()
|
||||
storage = _FakeStorage()
|
||||
_add_active_child(storage)
|
||||
ws = _FakeWorkstream(kind=WorkstreamKind.INTERACTIVE)
|
||||
mgr.add_ws(ws)
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
# Observer ignored the non-coord workstream entirely.
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
# Storage was NOT queried — kind check happens before list_workstreams.
|
||||
assert storage.list_calls == []
|
||||
|
||||
|
||||
class TestWaitForWorkstreamSkip:
|
||||
def test_skips_when_last_assistant_used_wait(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "kick off"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-1",
|
||||
"function": {"name": "wait_for_workstream", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
# Don't pile on — model is already using the right tool.
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
def test_fires_when_last_assistant_used_different_tool(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call-1", "function": {"name": "spawn_workstream", "arguments": "{}"}}
|
||||
],
|
||||
},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 1
|
||||
|
||||
|
||||
class TestHardCap:
|
||||
def test_hard_cap_blocks_after_n_fires(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
|
||||
# Bypass cooldown for this test: each call burns a per-type slot
|
||||
# in ``_metacog_state`` so we need to clear it between fires.
|
||||
for _ in range(3):
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
|
||||
# Cap = 3 fires. Even with cooldown bypassed, the 4th doesn't fire.
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
# We enqueued 3 entries total; cap blocked the 4th.
|
||||
snap = ws.session._nudge_queue.pending("any")
|
||||
assert len(snap) == 3
|
||||
|
||||
def test_cap_resets_when_state_leaves_idle_without_wake(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
|
||||
# Burn the cap.
|
||||
for _ in range(3):
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 3
|
||||
|
||||
# Drain the queue (simulate the watcher delivering them).
|
||||
ws.session._nudge_queue.drain({"any"})
|
||||
|
||||
# Real (non-wake) leave-IDLE: tag is empty. Cap resets.
|
||||
ws.session._wake_source_tag = ""
|
||||
mgr.fire_state(ws.id, WorkstreamState.RUNNING)
|
||||
|
||||
# New IDLE — cap is fresh, fires again.
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 1
|
||||
|
||||
def test_cap_does_not_reset_during_wake_driven_exit(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
|
||||
# Burn the cap.
|
||||
for _ in range(3):
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
ws.session._nudge_queue.drain({"any"})
|
||||
|
||||
# Wake-driven leave-IDLE: tag is set during the wake send.
|
||||
ws.session._wake_source_tag = "system_nudge"
|
||||
mgr.fire_state(ws.id, WorkstreamState.RUNNING)
|
||||
ws.session._wake_source_tag = "" # tag cleared at end of wake send
|
||||
|
||||
# Cap should NOT have reset — re-IDLE shouldn't fire.
|
||||
ws.session._metacog_state.clear()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 0
|
||||
|
||||
|
||||
class TestCooldown:
|
||||
def test_cooldown_blocks_within_window(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 1
|
||||
|
||||
# Drain so the queue isn't the gate.
|
||||
ws.session._nudge_queue.drain({"any"})
|
||||
|
||||
# Second fire within the cooldown window → should_nudge returns False.
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 0
|
||||
|
||||
|
||||
class TestStorageFailure:
|
||||
def test_storage_exception_is_swallowed(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
storage.list_raises = True
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
# Must not raise / propagate.
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
|
||||
class TestValidUntilPredicate:
|
||||
def test_predicate_drops_when_children_finish_before_drain(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage, ws_id="child-a", state="running")
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 1
|
||||
|
||||
# Children now complete (storage shows none active).
|
||||
storage.children.clear()
|
||||
|
||||
# Drain at the user seam — predicate re-queries, finds 0 active,
|
||||
# drops the entry without delivering.
|
||||
from turnstone.core.nudge_queue import USER_DRAIN
|
||||
|
||||
delivered = ws.session._nudge_queue.drain(USER_DRAIN)
|
||||
assert delivered == []
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
def test_predicate_delivers_when_children_still_active(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage, ws_id="child-a", state="running")
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
|
||||
# Children still active → predicate returns True → entry delivers.
|
||||
from turnstone.core.nudge_queue import USER_DRAIN
|
||||
|
||||
delivered = ws.session._nudge_queue.drain(USER_DRAIN)
|
||||
assert len(delivered) == 1
|
||||
assert delivered[0][0] == "idle_children"
|
||||
|
||||
def test_predicate_drops_on_storage_failure(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
|
||||
# Storage failure at drain time. Predicate treats raises as
|
||||
# "no longer valid" (drop) — see NudgeQueue.drain's predicate
|
||||
# exception handling.
|
||||
storage.count_raises = True
|
||||
|
||||
from turnstone.core.nudge_queue import USER_DRAIN
|
||||
|
||||
delivered = ws.session._nudge_queue.drain(USER_DRAIN)
|
||||
assert delivered == []
|
||||
|
||||
|
||||
class TestLifecycle:
|
||||
def test_start_idempotent(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
observer.start() # no-op
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
# Double-subscribe would have produced 2 entries.
|
||||
assert len(ws.session._nudge_queue.pending("any")) == 1
|
||||
|
||||
def test_shutdown_unsubscribes(self, coord_setup):
|
||||
mgr, storage, ws = coord_setup
|
||||
_add_active_child(storage)
|
||||
# ≥2 messages so should_nudge's message_count > 1 gate clears.
|
||||
ws.session.messages = [
|
||||
{"role": "user", "content": "go"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
]
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
observer.shutdown()
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
def test_shutdown_idempotent(self, coord_setup):
|
||||
mgr, _storage, _ws = coord_setup
|
||||
observer = CoordinatorIdleObserver(mgr, _storage)
|
||||
observer.start()
|
||||
observer.shutdown()
|
||||
observer.shutdown() # no error
|
||||
+206
-16
@@ -79,9 +79,10 @@ def test_coordinator_js_exposes_inline_approval_helpers():
|
||||
assert "function submitChildApproval" in body or "submitChildApproval(" in body
|
||||
# The shared approve POST helper (parameterized for child ws_ids)
|
||||
assert "function approveWorkstream" in body or "approveWorkstream(" in body
|
||||
# The urgent live-bulk fetch option that fires on activity_state
|
||||
# transitions in/out of "approval"
|
||||
assert "{ urgent: true }" in body or "urgent: true" in body
|
||||
# The 409 stale-call_id retry path uses invalidateLiveBadge +
|
||||
# scheduleLiveFetch (Stage 3 cleanup removed the urgent flag —
|
||||
# cache invalidation makes the TTL gate fall through naturally).
|
||||
assert "invalidateLiveBadge(targetWsId)" in body
|
||||
# Server-side payload field — drift here means the JS reads stale keys
|
||||
assert "pending_approval_detail" in body
|
||||
# Reconnect parity (chunk 4): the SSE re-open handler must drop
|
||||
@@ -90,10 +91,10 @@ def test_coordinator_js_exposes_inline_approval_helpers():
|
||||
# can't render zombie approve/deny buttons on a row whose
|
||||
# approval was resolved during the gap. The implementation
|
||||
# iterates the cache and deletes only !permanent entries —
|
||||
# asserting the literal Map iteration form keeps a refactor
|
||||
# back to liveBadgeCache.clear() (which would re-pay 403s on
|
||||
# every reconnect for denied ids) from sneaking in.
|
||||
assert "liveBadgeCache.delete" in body
|
||||
# asserting the literal helper call keeps a refactor back to
|
||||
# _liveBadgeCacheClear() (which would re-pay 403s on every
|
||||
# reconnect for denied ids) from sneaking in.
|
||||
assert "_liveBadgeCacheDelete" in body
|
||||
# Edge-case matrix sentinel labels — POLICY-BLOCKED renders when
|
||||
# an item has error set + needs_approval=False (server-side
|
||||
# tool policy already blocked the call); "(judge unavailable)"
|
||||
@@ -113,15 +114,19 @@ def test_coordinator_js_exposes_inline_approval_helpers():
|
||||
# coord-self ws_id (the coord lives on the console process).
|
||||
# Children live on cluster nodes and 404 without the prefix.
|
||||
assert "/v1/api/route/workstreams/" in body
|
||||
# Late-judge polling — the LLM judge runs async on the child
|
||||
# node and never pushes a signal that reaches the coord, so
|
||||
# the row's pending_approval_detail with judge_pending=true
|
||||
# would freeze on heuristic verdicts forever without this
|
||||
# poll loop. The poller is GLOBAL (not per-row) so off-screen
|
||||
# rows still refresh — a per-row poller's scheduleLiveFetch
|
||||
# call short-circuits on non-visible rows, leaving them stuck.
|
||||
assert "_maybeStartJudgePoll" in body
|
||||
assert "_judgePollTick" in body
|
||||
# Late-arriving LLM judge verdicts — Stage 3 Step 5 promoted
|
||||
# ``intent_verdict`` and ``approval_resolved`` to first-class
|
||||
# cluster-bus event types, so the coord adapter dispatches them
|
||||
# as ``child_ws_intent_verdict`` / ``child_ws_approval_resolved``
|
||||
# on the parent's SSE stream. The browser handlers write
|
||||
# directly to liveBadgeCache (bypassing scheduleLiveFetch's
|
||||
# visibility gate cleanly) so off-screen rows pick up verdicts
|
||||
# without polling. Replaced the old ``_judgePollTick`` 90-second
|
||||
# global poll loop and its visibility-gate-bypass workaround.
|
||||
assert "handleChildIntentVerdict" in body
|
||||
assert "handleChildApprovalResolved" in body
|
||||
assert "child_ws_intent_verdict" in body
|
||||
assert "child_ws_approval_resolved" in body
|
||||
# Reload parity for the coord-self approval gate: init() must
|
||||
# consume the authoritative GET /workstreams snapshot's
|
||||
# pending_approval_detail so a freshly opened tab can render
|
||||
@@ -152,3 +157,188 @@ def test_coordinator_js_exposes_inline_approval_helpers():
|
||||
# any prior denial. bug-1 / bug-3 from the second /review pass.
|
||||
assert "Denied by user" in body
|
||||
assert "callOutcomes" in body
|
||||
# User-message attachment pills — both live send (coordSend) and
|
||||
# history replay route through appendUserMessageWithAttachments.
|
||||
# Renaming or dropping the helper would silently regress the
|
||||
# attachment affordance to the pre-fix plain-text bubble, which
|
||||
# would only surface in manual testing of an attached-file flow.
|
||||
# The CSS class is the visual anchor (coordinator.css) — keeping
|
||||
# both literals in the smoke layer covers JS↔CSS drift in either
|
||||
# direction.
|
||||
assert "function appendUserMessageWithAttachments" in body
|
||||
assert "msg-user-attach" in body
|
||||
# PR #487 — whitespace-only assistant content (Qwen3 with vLLM
|
||||
# ``--reasoning-parser`` strips ``<think>…</think>`` and emits only
|
||||
# ``"\n\n"`` as content before a tool call) must be skipped on
|
||||
# history replay or the empty ``.msg.assistant`` card surfaces as
|
||||
# a phantom row. The literal substring ``content.trim()`` is the
|
||||
# single-line guard the rendering branch uses; a refactor that
|
||||
# drops the trim() (e.g. simplifies to ``if (!content)``) silently
|
||||
# regresses the phantom-card fix on the multi-node coord path.
|
||||
# Mirrors ``test_app_js.py``'s same-shape pin on ``app.js``.
|
||||
assert "content.trim()" in body
|
||||
# PR #487 — coord history replay must render the assistant content
|
||||
# card BEFORE the tool batch, not after, so DOM order matches the
|
||||
# chronological order the model emitted (text → dispatch → results).
|
||||
# Pre-fix the tool_calls branch sat at the role-agnostic top of the
|
||||
# loop and rendered ahead of the assistant text that announced the
|
||||
# batch, putting parallel fan-outs visually above their narrating
|
||||
# message. The fix hoisted the synthesis into ``renderAssistantToolBatch``
|
||||
# called from inside the assistant branch AFTER the content card —
|
||||
# asserting the helper name lets a refactor that re-inlines or
|
||||
# renames it surface here instead of via manual reload testing.
|
||||
assert "function renderAssistantToolBatch" in body
|
||||
assert "renderAssistantToolBatch(m)" in body
|
||||
|
||||
|
||||
def test_coordinator_js_handle_child_state_no_longer_reads_sse_pending_approval_detail():
|
||||
"""Stage 3 cleanup — ``pending_approval_detail`` is no longer
|
||||
piggybacked on child_ws_state events. Approval items now arrive
|
||||
via bulk fetch on the activity_state="approval" transition;
|
||||
verdicts via the explicit ``child_ws_intent_verdict`` event class;
|
||||
resolution via ``child_ws_approval_resolved``. A refactor that
|
||||
re-introduces the piggyback would silently re-open the
|
||||
duplicate-path race the dedicated event classes were added to
|
||||
eliminate.
|
||||
|
||||
Structural assertions (regex against multi-line source) — symbol-
|
||||
presence alone wouldn't catch a guard that keeps the names but
|
||||
inverts the comparison or drops the ``prev.live`` check. This
|
||||
codebase has no JS test framework, so locking the guard's shape
|
||||
here is the next-best thing to a behavioral test."""
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
coord_js = Path(__file__).resolve().parent.parent / (
|
||||
"turnstone/console/static/coordinator/coordinator.js"
|
||||
)
|
||||
body = coord_js.read_text(encoding="utf-8")
|
||||
|
||||
# The piggyback read is gone from handleChildState. (The string
|
||||
# may still appear elsewhere — e.g. handleChildIntentVerdict
|
||||
# reading from cache, or comments — but never as ``ev.pending_approval_detail``.)
|
||||
assert "ev.pending_approval_detail" not in body
|
||||
# The pre-fix urgent-fetch on activity_state transitions is gone.
|
||||
assert "enteredApproval" not in body
|
||||
assert "leftApproval" not in body
|
||||
# ``pendingApproval`` flag derivation must check BOTH state and
|
||||
# activity_state. The worker thread can fire the state transition
|
||||
# to "attention" before approve_tools updates activity_state, so
|
||||
# checking only activity_state misses children that legitimately
|
||||
# need approval. Pin the disjunction so the regression doesn't
|
||||
# silently re-introduce.
|
||||
assert re.search(
|
||||
r'existing\.state\s*===\s*"attention"\s*\|\|\s*'
|
||||
r'existing\.activity_state\s*===\s*"approval"',
|
||||
body,
|
||||
), (
|
||||
"handleChildState must derive pendingApproval from "
|
||||
"(state==='attention' || activity_state==='approval')"
|
||||
)
|
||||
|
||||
# SSE-authoritative window constant is defined and used.
|
||||
assert re.search(r"\bconst\s+SSE_AUTHORITATIVE_MS\s*=\s*\d+", body), (
|
||||
"SSE_AUTHORITATIVE_MS constant must be defined as a numeric literal"
|
||||
)
|
||||
|
||||
# SSE writers tag entries with sseUpdatedAt: Date.now() so the
|
||||
# merge guard in flushLiveFetches preserves them against stale
|
||||
# bulk-fetch responses. handleChildState only stamps when it
|
||||
# AUTHORITATIVELY clears the detail (off-approval transition);
|
||||
# writers that stamp unconditionally are intent_verdict (verdict
|
||||
# stamp), approval_resolved (clear), and the optimistic-clear
|
||||
# path in submitChildApproval. Pinning the literal Date.now()
|
||||
# call keeps a refactor that drops the SSE-source tag entirely
|
||||
# from sneaking in.
|
||||
assert re.search(
|
||||
r"sseUpdatedAt:\s*Date\.now\(\)",
|
||||
body,
|
||||
), "Critical SSE writers must stamp sseUpdatedAt: Date.now()"
|
||||
|
||||
# flushLiveFetches' merge guard structure: SSE-set pending_approval
|
||||
# / _detail wins over a stale bulk-poll snapshot when (live) AND
|
||||
# (prev exists) AND (prev.sseUpdatedAt set) AND (within window)
|
||||
# AND (prev.live exists). Inverting the comparison or dropping
|
||||
# any of these guards reopens the clobber bug.
|
||||
merge_guard = re.search(
|
||||
r"if\s*\(\s*live\s*&&\s*prev\s*&&\s*prev\.sseUpdatedAt\s*&&\s*"
|
||||
r"now\s*-\s*prev\.sseUpdatedAt\s*<\s*SSE_AUTHORITATIVE_MS\s*&&\s*"
|
||||
r"prev\.live\s*\)",
|
||||
body,
|
||||
)
|
||||
assert merge_guard is not None, (
|
||||
"flushLiveFetches merge guard must be the conjunction "
|
||||
"(live && prev && prev.sseUpdatedAt && now - prev.sseUpdatedAt < "
|
||||
"SSE_AUTHORITATIVE_MS && prev.live). An inverted comparison or "
|
||||
"missing prev.live check would let a stale bulk-poll clobber a "
|
||||
"fresh SSE-set approval."
|
||||
)
|
||||
|
||||
# The merge body must preserve BOTH pending_approval and
|
||||
# pending_approval_detail from prev — preserving only one would
|
||||
# render a row with a phantom badge but no buttons (or vice versa).
|
||||
merge_body = re.search(
|
||||
r"mergedLive\s*=\s*Object\.assign\(\s*\{\}\s*,\s*live\s*,\s*\{"
|
||||
r"[^}]*pending_approval:\s*prev\.live\.pending_approval[^}]*"
|
||||
r"pending_approval_detail:\s*prev\.live\.pending_approval_detail",
|
||||
body,
|
||||
)
|
||||
assert merge_body is not None, (
|
||||
"Merge body must preserve both pending_approval AND "
|
||||
"pending_approval_detail from prev.live — preserving only one "
|
||||
"creates a half-rendered approval row."
|
||||
)
|
||||
|
||||
# flushLiveFetches must forward sseUpdatedAt onto the new cache
|
||||
# entry so the SSE-source tag survives the bulk-poll write back —
|
||||
# without this, every bulk-poll resets the window and the next
|
||||
# late-arriving poll silently clobbers.
|
||||
assert re.search(
|
||||
r"sseUpdatedAt:\s*prev\s*\?\s*prev\.sseUpdatedAt",
|
||||
body,
|
||||
), (
|
||||
"flushLiveFetches must forward prev.sseUpdatedAt onto the new "
|
||||
"cache entry (preserving the SSE-source window across bulk-poll "
|
||||
"cycles) — without this, the second bulk-poll after an SSE "
|
||||
"transition silently clobbers."
|
||||
)
|
||||
|
||||
|
||||
def test_coord_history_renders_user_interjection_advisory_after_tool_block():
|
||||
"""Queued user messages spliced into the last tool-result envelope
|
||||
of a batch (Seam 1) persist on the tool DB row as a wrapped
|
||||
``<tool_output>`` envelope. ``decorate_history_messages`` extracts
|
||||
the advisory back out and the wire layer projects it onto
|
||||
``m.advisories``; the coord history loop must invoke the shared
|
||||
``replayAdvisoriesAfterTool`` helper (defined in
|
||||
``shared/utils.js``) so each ``user_interjection`` renders through
|
||||
``appendUserMessageWithAttachments`` and the bubble looks identical
|
||||
to a Seam 2/3 user row.
|
||||
|
||||
This test pins the call site so a refactor that drops the helper
|
||||
invocation regresses the queued-during-batch replay shape silently.
|
||||
Mirrors ``test_app_js.py``'s same-shape pin on interactive's
|
||||
``replayHistory``."""
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
coord_js = Path(__file__).resolve().parent.parent / (
|
||||
"turnstone/console/static/coordinator/coordinator.js"
|
||||
)
|
||||
body = coord_js.read_text(encoding="utf-8")
|
||||
|
||||
assert "replayAdvisoriesAfterTool(m.advisories" in body, (
|
||||
"Coord history loop must invoke replayAdvisoriesAfterTool with "
|
||||
"m.advisories so queued messages spliced into the tool envelope "
|
||||
"render as user bubbles after the tool block."
|
||||
)
|
||||
# The renderer callback routes through appendUserMessageWithAttachments
|
||||
# so the bubble matches a normal user-row replay.
|
||||
assert re.search(
|
||||
r"appendUserMessageWithAttachments\(\s*text",
|
||||
body,
|
||||
), (
|
||||
"Coord history loop's renderer callback must route the extracted "
|
||||
"advisory text through appendUserMessageWithAttachments so the "
|
||||
"rendered bubble matches a normal user-row replay."
|
||||
)
|
||||
|
||||
@@ -15,6 +15,16 @@ class TestIsSecret:
|
||||
assert _is_secret("TURNSTONE_JWT_SECRET") is True
|
||||
assert _is_secret("AWS_SECRET_ACCESS_KEY") is True
|
||||
|
||||
def test_tool_config_paths_scrubbed(self):
|
||||
"""Tool-config env vars whose target files load executable
|
||||
directives must be scrubbed even though they don't match a
|
||||
secret-suffix pattern. Defence-in-depth alongside on-CLI
|
||||
``--no-config`` for ripgrep and friends."""
|
||||
assert _is_secret("RIPGREP_CONFIG_PATH") is True
|
||||
assert _is_secret("GIT_CONFIG") is True
|
||||
assert _is_secret("GIT_CONFIG_GLOBAL") is True
|
||||
assert _is_secret("GIT_CONFIG_SYSTEM") is True
|
||||
|
||||
def test_suffix_matching(self):
|
||||
assert _is_secret("MY_CUSTOM_API_KEY") is True
|
||||
assert _is_secret("DB_PASSWORD") is True
|
||||
|
||||
@@ -0,0 +1,455 @@
|
||||
"""Unit tests for ``turnstone.core.history_decoration``.
|
||||
|
||||
The decoration helpers are shared between two surfaces — interactive's
|
||||
SSE replay (``_build_history``) and the lifted ``/history`` REST
|
||||
endpoint (``make_history_handler``, used by both interactive and
|
||||
coord). Pinning the wire shape here lets a future schema/projection
|
||||
change land in one file rather than spread across the two surfaces.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from turnstone.core.history_decoration import (
|
||||
build_output_assessment_payload,
|
||||
build_verdict_payload,
|
||||
decorate_history_messages,
|
||||
decorate_tool_call,
|
||||
)
|
||||
|
||||
|
||||
class TestBuildVerdictPayload:
|
||||
"""The wire-shape projection that's the single source of truth for
|
||||
what intent_verdict fields ship to the client."""
|
||||
|
||||
def test_skips_unflagged_baseline(self) -> None:
|
||||
"""``risk_level`` "none" is the unflagged-tool baseline; the
|
||||
client filters those anyway, so projecting None at the wire
|
||||
layer keeps the payload tight on long workstreams."""
|
||||
row = {"risk_level": "none", "recommendation": "approve", "tier": "heuristic"}
|
||||
assert build_verdict_payload(row) is None
|
||||
|
||||
def test_drops_call_id_and_func_name(self) -> None:
|
||||
"""The client already has these on ``tc.id`` / ``tc.name``;
|
||||
re-shipping them per-tool_call would balloon long replays."""
|
||||
row = {
|
||||
"call_id": "call_abc",
|
||||
"func_name": "bash",
|
||||
"risk_level": "medium",
|
||||
"recommendation": "review",
|
||||
"confidence": 0.8,
|
||||
"intent_summary": "summary",
|
||||
"tier": "heuristic",
|
||||
}
|
||||
out = build_verdict_payload(row)
|
||||
assert out is not None
|
||||
assert "call_id" not in out
|
||||
assert "func_name" not in out
|
||||
# Sanity — the kept fields are the ones renderVerdictBadge reads.
|
||||
assert out["risk_level"] == "medium"
|
||||
assert out["recommendation"] == "review"
|
||||
assert out["confidence"] == 0.8
|
||||
assert out["intent_summary"] == "summary"
|
||||
assert out["tier"] == "heuristic"
|
||||
|
||||
def test_includes_reasoning_for_either_tier_when_present(self) -> None:
|
||||
"""Heuristic verdicts in this project emit structured
|
||||
rationales (one per matched pattern) — e.g.
|
||||
``policy.py`` writes a reasoning string per heuristic hit.
|
||||
Ship the field for either tier when it has content; only
|
||||
omit when the row didn't write one."""
|
||||
for tier in ("heuristic", "llm"):
|
||||
row = {
|
||||
"risk_level": "high",
|
||||
"tier": tier,
|
||||
"reasoning": "The command exfiltrates ~/.ssh/id_rsa over an external connection.",
|
||||
}
|
||||
out = build_verdict_payload(row)
|
||||
assert out is not None
|
||||
assert "id_rsa" in out["reasoning"]
|
||||
|
||||
def test_omits_reasoning_when_empty(self) -> None:
|
||||
"""An absent / empty reasoning string shouldn't ship as
|
||||
``reasoning: ""`` — the rationale ``<details>`` block on the
|
||||
client renders an empty disclosure when the field is present
|
||||
but empty."""
|
||||
row = {"risk_level": "high", "tier": "heuristic", "reasoning": ""}
|
||||
out = build_verdict_payload(row)
|
||||
assert out is not None
|
||||
assert "reasoning" not in out
|
||||
|
||||
def test_includes_judge_model_when_present(self) -> None:
|
||||
"""``judge_model`` rides through so the batch tier badge can
|
||||
render ``⚖ llm:claude-haiku-4`` on history-only replays
|
||||
rather than the bare ``⚖ llm`` label."""
|
||||
row = {"risk_level": "high", "tier": "llm", "judge_model": "claude-haiku-4"}
|
||||
out = build_verdict_payload(row)
|
||||
assert out is not None
|
||||
assert out["judge_model"] == "claude-haiku-4"
|
||||
|
||||
def test_omits_judge_model_when_empty(self) -> None:
|
||||
row = {"risk_level": "medium", "tier": "heuristic", "judge_model": ""}
|
||||
out = build_verdict_payload(row)
|
||||
assert out is not None
|
||||
assert "judge_model" not in out
|
||||
|
||||
|
||||
class TestBuildOutputAssessmentPayload:
|
||||
"""Output-guard wire shape — flags decoded from JSON string at
|
||||
this layer so the client never has to parse twice."""
|
||||
|
||||
def test_skips_unflagged_baseline(self) -> None:
|
||||
row = {"risk_level": "none", "flags": "[]"}
|
||||
assert build_output_assessment_payload(row) is None
|
||||
|
||||
def test_decodes_flags_from_json(self) -> None:
|
||||
row = {"risk_level": "high", "flags": '["api_key","email"]', "redacted": 1}
|
||||
out = build_output_assessment_payload(row)
|
||||
assert out is not None
|
||||
assert out["flags"] == ["api_key", "email"]
|
||||
assert out["redacted"] is True
|
||||
assert out["risk_level"] == "high"
|
||||
|
||||
def test_handles_malformed_flags_json(self) -> None:
|
||||
"""Bad JSON in ``flags`` must not block the rest of the
|
||||
assessment from rendering — degrade to empty list."""
|
||||
row = {"risk_level": "medium", "flags": "not-json", "redacted": 0}
|
||||
out = build_output_assessment_payload(row)
|
||||
assert out is not None
|
||||
assert out["flags"] == []
|
||||
assert out["redacted"] is False
|
||||
|
||||
|
||||
class TestDecorateToolCall:
|
||||
"""In-place mutation of either OpenAI-format or flattened tool_call
|
||||
entries — both shapes carry ``id`` at the top level."""
|
||||
|
||||
def test_attaches_verdict_when_present(self) -> None:
|
||||
tc: dict[str, object] = {"id": "call_1", "function": {"name": "bash", "arguments": "{}"}}
|
||||
verdicts = {
|
||||
"call_1": {
|
||||
"risk_level": "medium",
|
||||
"recommendation": "review",
|
||||
"confidence": 0.7,
|
||||
"intent_summary": "summary",
|
||||
"tier": "heuristic",
|
||||
}
|
||||
}
|
||||
decorate_tool_call(tc, verdicts, {})
|
||||
assert "verdict" in tc
|
||||
assert tc["verdict"]["risk_level"] == "medium" # type: ignore[index]
|
||||
|
||||
def test_skips_when_no_call_id_match(self) -> None:
|
||||
tc: dict[str, object] = {"id": "call_other", "name": "bash"}
|
||||
verdicts = {
|
||||
"call_1": {"risk_level": "medium", "tier": "heuristic"},
|
||||
}
|
||||
decorate_tool_call(tc, verdicts, {})
|
||||
assert "verdict" not in tc
|
||||
|
||||
def test_skips_unflagged_verdict(self) -> None:
|
||||
"""``build_verdict_payload`` returns None for unflagged rows;
|
||||
decorate_tool_call must not stamp ``verdict`` in that case."""
|
||||
tc: dict[str, object] = {"id": "call_1", "name": "bash"}
|
||||
verdicts = {"call_1": {"risk_level": "none", "tier": "heuristic"}}
|
||||
decorate_tool_call(tc, verdicts, {})
|
||||
assert "verdict" not in tc
|
||||
|
||||
def test_handles_empty_id(self) -> None:
|
||||
"""A tool_call with no id can't be paired against the lookup
|
||||
table — must not raise (or stamp the wrong row's verdict)."""
|
||||
tc: dict[str, object] = {"id": "", "name": "bash"}
|
||||
verdicts = {"call_1": {"risk_level": "high", "tier": "heuristic"}}
|
||||
decorate_tool_call(tc, verdicts, {})
|
||||
assert "verdict" not in tc
|
||||
|
||||
|
||||
class TestDecorateHistoryMessages:
|
||||
"""End-to-end mutation of a /history-shaped message list — covers
|
||||
the full transform applied by ``make_history_handler``."""
|
||||
|
||||
def test_decorates_tool_calls_with_verdict_and_assessment(self) -> None:
|
||||
verdicts = {
|
||||
"call_a": {
|
||||
"risk_level": "high",
|
||||
"recommendation": "deny",
|
||||
"confidence": 0.95,
|
||||
"intent_summary": "exfil",
|
||||
"tier": "llm",
|
||||
"reasoning": "ssh key access",
|
||||
}
|
||||
}
|
||||
assessments = {
|
||||
"call_a": {"risk_level": "high", "flags": '["secret"]', "redacted": 1},
|
||||
}
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "running",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_a",
|
||||
"function": {"name": "bash", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": "long output"},
|
||||
{"role": "tool", "tool_call_id": "call_b", "content": "short"},
|
||||
]
|
||||
decorate_history_messages(messages, verdicts, assessments)
|
||||
# Assistant tool_calls got both decorations.
|
||||
tc = messages[1]["tool_calls"][0] # type: ignore[index]
|
||||
assert tc["verdict"]["risk_level"] == "high"
|
||||
assert tc["verdict"]["tier"] == "llm"
|
||||
assert "reasoning" in tc["verdict"]
|
||||
assert tc["output_assessment"]["flags"] == ["secret"]
|
||||
assert tc["output_assessment"]["redacted"] is True
|
||||
# Plain tool content (no envelope) is left intact and no
|
||||
# advisories key is set.
|
||||
assert messages[2]["content"] == "long output"
|
||||
assert "advisories" not in messages[2]
|
||||
assert messages[3]["content"] == "short"
|
||||
assert "advisories" not in messages[3]
|
||||
|
||||
def test_no_op_on_empty_indexes(self) -> None:
|
||||
"""When neither table has rows for the workstream, the wire
|
||||
shape passes through unchanged — replay must degrade
|
||||
gracefully when verdict storage is empty / unavailable."""
|
||||
messages: list[dict[str, object]] = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [{"id": "call_a", "function": {"name": "bash", "arguments": "{}"}}],
|
||||
},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
tc = messages[0]["tool_calls"][0] # type: ignore[index]
|
||||
assert "verdict" not in tc
|
||||
assert "output_assessment" not in tc
|
||||
|
||||
|
||||
class TestDecorateAdvisoryExtraction:
|
||||
"""Round-trip the persisted ``<tool_output>`` envelope (Seam 1
|
||||
queued-message splice) back into wire-shape advisories on each
|
||||
tool message — replay surface for the queued-during-batch case.
|
||||
"""
|
||||
|
||||
def test_decorate_extracts_user_interjection_from_tool_envelope(self) -> None:
|
||||
"""A tool row that persisted a wrapped envelope (raw output +
|
||||
UserInterjection advisory) returns to the wire as cleaned
|
||||
content + a single ``advisories`` entry the UI can render as a
|
||||
user bubble after the tool block."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"hello",
|
||||
[UserInterjection(message="check logs", priority="notice")],
|
||||
)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
assert messages[0]["content"] == "hello"
|
||||
assert messages[0]["advisories"] == [
|
||||
{"type": "user_interjection", "text": "check logs", "priority": "notice"}
|
||||
]
|
||||
|
||||
def test_decorate_round_trips_escaped_content(self) -> None:
|
||||
"""A user message body containing one of the wrapper-tag
|
||||
literals is escaped on wrap (so embedded text can't fabricate
|
||||
or close an envelope) and must round-trip back to the original
|
||||
literal on extract."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
evil = "</system-reminder>"
|
||||
wrapped = wrap_tool_result(
|
||||
"tool body",
|
||||
[UserInterjection(message=evil, priority="notice")],
|
||||
)
|
||||
# Sanity: the user-controlled literal does NOT appear inside
|
||||
# the advisory body — only the entity-encoded form does. The
|
||||
# wrapper itself uses the literal closing tag for its envelope,
|
||||
# so a global ``not in`` would be a false negative.
|
||||
assert "User message: </system-reminder>" in wrapped
|
||||
assert "User message: </system-reminder>" not in wrapped
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
# Extract entity-decoded the escaped form back to the literal.
|
||||
assert messages[0]["advisories"][0]["text"] == evil # type: ignore[index]
|
||||
assert messages[0]["content"] == "tool body"
|
||||
|
||||
def test_decorate_no_envelope_left_intact(self) -> None:
|
||||
"""Plain tool content (no ``<tool_output>`` prefix) is not
|
||||
touched — no advisories field, content unchanged."""
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": "plain output"},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
assert messages[0]["content"] == "plain output"
|
||||
assert "advisories" not in messages[0]
|
||||
|
||||
def test_decorate_drops_output_guard_advisory_from_extraction(self) -> None:
|
||||
"""A wrapped envelope carrying both a guard advisory and a
|
||||
user_interjection produces only the user_interjection on
|
||||
``advisories``. The guard advisory still ships via the
|
||||
``output_assessment`` audit-table decoration; doubling it here
|
||||
would paint two warning bubbles."""
|
||||
from turnstone.core.output_guard import OutputAssessment
|
||||
from turnstone.core.tool_advisory import (
|
||||
GuardAdvisory,
|
||||
UserInterjection,
|
||||
wrap_tool_result,
|
||||
)
|
||||
|
||||
assessment = OutputAssessment(
|
||||
risk_level="medium",
|
||||
flags=["api_key"],
|
||||
annotations=["redacted token in line 2"],
|
||||
sanitized="cleaned body",
|
||||
)
|
||||
wrapped = wrap_tool_result(
|
||||
"raw body",
|
||||
[
|
||||
GuardAdvisory(assessment=assessment, func_name="bash"),
|
||||
UserInterjection(message="and here", priority="notice"),
|
||||
],
|
||||
)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
adv = messages[0]["advisories"]
|
||||
assert len(adv) == 1 # type: ignore[arg-type]
|
||||
assert adv[0]["type"] == "user_interjection" # type: ignore[index]
|
||||
|
||||
def test_decorate_handles_important_priority(self) -> None:
|
||||
"""The MUST-address preamble round-trips to ``priority=important``."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"out",
|
||||
[UserInterjection(message="urgent", priority="important")],
|
||||
)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
adv = messages[0]["advisories"][0] # type: ignore[index]
|
||||
assert adv["priority"] == "important"
|
||||
assert adv["text"] == "urgent"
|
||||
|
||||
def test_decorate_suppresses_empty_advisory_body(self) -> None:
|
||||
"""``queue_message`` doesn't reject empty / whitespace-only
|
||||
text, so an advisory with an empty body can round-trip through
|
||||
``wrap_tool_result``. ``_classify_advisory`` must filter those
|
||||
out so replay doesn't paint a featureless empty user bubble.
|
||||
|
||||
Removing the ``if not body.strip(): return None`` guard in
|
||||
``_classify_advisory`` breaks this test."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"tool body",
|
||||
[UserInterjection(message="", priority="notice")],
|
||||
)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
# Envelope is still stripped from content (the cleaning side
|
||||
# of decoration runs unconditionally), but no advisories
|
||||
# surface — the empty body is filtered.
|
||||
assert messages[0]["content"] == "tool body"
|
||||
assert "advisories" not in messages[0]
|
||||
|
||||
def test_decorate_suppresses_whitespace_only_advisory_body(self) -> None:
|
||||
"""Whitespace-only bodies are similarly suppressed — same
|
||||
reasoning as the empty-body case."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"tool body",
|
||||
[UserInterjection(message=" \n\t ", priority="notice")],
|
||||
)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": wrapped},
|
||||
]
|
||||
decorate_history_messages(messages, {}, {})
|
||||
assert messages[0]["content"] == "tool body"
|
||||
assert "advisories" not in messages[0]
|
||||
|
||||
def test_wrap_extract_round_trips_preexisting_entities(self) -> None:
|
||||
"""A user message body containing literal HTML-entity references
|
||||
matching the wrapper-escape forms must round-trip identically
|
||||
through ``wrap_tool_result + extract_advisories_from_tool_envelope``.
|
||||
Without escaping ``&`` first in the encode step, encode→decode
|
||||
would produce the bare wrapper tag, fabricating an envelope the
|
||||
wrapper layer never produced.
|
||||
"""
|
||||
from turnstone.core.history_decoration import (
|
||||
extract_advisories_from_tool_envelope,
|
||||
)
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
tricky = "I describe XML tags like <tool_output> in my docs."
|
||||
wrapped = wrap_tool_result(
|
||||
"tool body",
|
||||
[UserInterjection(message=tricky, priority="notice")],
|
||||
)
|
||||
result = extract_advisories_from_tool_envelope(wrapped)
|
||||
assert result is not None
|
||||
cleaned, advisories = result
|
||||
assert cleaned == "tool body"
|
||||
assert len(advisories) == 1
|
||||
# The original literal entity-reference text round-trips
|
||||
# identically — the parser does not silently turn it into a
|
||||
# bare wrapper tag.
|
||||
assert advisories[0]["text"] == tricky
|
||||
|
||||
def test_save_load_decorate_round_trips_envelope(self, backend) -> None:
|
||||
"""End-to-end round-trip pinning the persisted-envelope
|
||||
contract. Persists a wrapped tool-output envelope via
|
||||
``save_message``, loads via ``load_messages``, runs
|
||||
``decorate_history_messages``, asserts the wire shape carries
|
||||
the extracted advisory + cleaned content. Pins the contract
|
||||
every component in the chain participates in (persistence
|
||||
layer ↔ in-memory replay ↔ wire projection) so a schema drift,
|
||||
an envelope-format change, or a parser regression surfaces
|
||||
here rather than only in production.
|
||||
"""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"command output",
|
||||
[UserInterjection(message="check the logs", priority="notice")],
|
||||
)
|
||||
backend.register_workstream("ws_rt_1")
|
||||
backend.save_message("ws_rt_1", "user", "go")
|
||||
backend.save_message(
|
||||
"ws_rt_1",
|
||||
"assistant",
|
||||
None,
|
||||
tool_calls='[{"id":"call_a","type":"function","function":{"name":"bash","arguments":"{}"}}]',
|
||||
)
|
||||
backend.save_message(
|
||||
"ws_rt_1",
|
||||
"tool",
|
||||
wrapped,
|
||||
tool_call_id="call_a",
|
||||
)
|
||||
msgs = backend.load_messages("ws_rt_1")
|
||||
# Persisted shape — content survives the storage layer
|
||||
# untouched. Symmetry with in-memory ``self.messages[i]['content']``
|
||||
# is what makes envelope extraction lossless on replay.
|
||||
tool_msg = next(m for m in msgs if m["role"] == "tool")
|
||||
assert tool_msg["content"] == wrapped
|
||||
# Decorate (the /history shared transform) — extracts the
|
||||
# advisory and strips the envelope.
|
||||
decorate_history_messages(msgs, {}, {})
|
||||
tool_msg = next(m for m in msgs if m["role"] == "tool")
|
||||
assert tool_msg["content"] == "command output"
|
||||
assert tool_msg["advisories"] == [
|
||||
{"type": "user_interjection", "text": "check the logs", "priority": "notice"}
|
||||
]
|
||||
@@ -0,0 +1,385 @@
|
||||
"""Boundary-crossing integration test for the wake trigger pipeline.
|
||||
|
||||
Drives a *real* :class:`SessionManager` + a *real* :class:`ChatSession`
|
||||
+ a *real* :class:`IdleNudgeWatcher` end-to-end. The only stub is the
|
||||
LLM provider (patched ``_create_stream_with_retry``); every other layer
|
||||
is production code:
|
||||
|
||||
* ``SessionManager.set_state`` snapshotting + iterating subscribers
|
||||
* ``IdleNudgeWatcher._on_state`` peeking the queue
|
||||
* ``session_worker.send`` atomic-spawn + daemon thread
|
||||
* ``ChatSession.deliver_wake_nudge_from_queue`` opening / closing
|
||||
``_wake_source_tag``
|
||||
* ``ChatSession.send`` chat loop short-circuiting metacog detection
|
||||
* ``_append_user_turn`` stamping ``_source = "system_nudge"``
|
||||
* ``_attach_pending_user_reminders`` draining ``USER_DRAIN``
|
||||
* ``_apply_reminders_for_provider`` splicing the rendered envelope
|
||||
onto empty content
|
||||
|
||||
Per ``feedback_tests_through_boundaries.md``: direct injection tests
|
||||
that bypass these boundaries silently mask wiring bugs. This test is
|
||||
the structural integration gate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.test_session_manager import FakeStorage
|
||||
from turnstone.core.idle_nudge_watcher import IdleNudgeWatcher
|
||||
from turnstone.core.session import ChatSession
|
||||
from turnstone.core.session_manager import SessionManager
|
||||
from turnstone.core.workstream import Workstream, WorkstreamKind, WorkstreamState
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Minimal fake adapter / UI for this integration test. Storage reuses
|
||||
# the canonical FakeStorage from test_session_manager.py to avoid the
|
||||
# drift risk of a parallel fake.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeUI:
|
||||
"""Minimal UI surface for ChatSession + SessionManager.cleanup_ui."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.events: list[tuple[str, Any]] = []
|
||||
|
||||
def _unblock(self) -> None: # SessionManager.close calls this
|
||||
pass
|
||||
|
||||
def broadcast_ws_closed(self) -> None:
|
||||
pass
|
||||
|
||||
# ChatSession callbacks (no-op for this test)
|
||||
def on_thinking_start(self) -> None:
|
||||
pass
|
||||
|
||||
def on_thinking_end(self) -> None:
|
||||
pass
|
||||
|
||||
def on_state_change(self, state: str) -> None:
|
||||
self.events.append(("state", state))
|
||||
|
||||
def on_user_reminder(self, reminders: Any, source: str | None = None) -> None:
|
||||
self.events.append(("user_reminder", reminders, source))
|
||||
|
||||
def on_error(self, message: str) -> None:
|
||||
pass
|
||||
|
||||
def on_rename(self, name: str) -> None:
|
||||
pass
|
||||
|
||||
def on_output_warning(self, call_id: Any, assessment: Any) -> None:
|
||||
pass
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Catch-all for any UI hook not enumerated above so the chat
|
||||
# loop's ``self.ui.<something>()`` call doesn't blow up.
|
||||
return MagicMock()
|
||||
|
||||
|
||||
class _BuildRealSessionAdapter:
|
||||
"""Adapter that returns a real :class:`ChatSession` instead of a stub.
|
||||
|
||||
Tracks emit_* events the integration test asserts on. Mirrors the
|
||||
``SessionKindAdapter`` + ``SessionEventEmitter`` Protocol surface
|
||||
that production ``WebUI`` / coord adapters expose.
|
||||
"""
|
||||
|
||||
def __init__(self, kind: WorkstreamKind = WorkstreamKind.INTERACTIVE) -> None:
|
||||
self.kind = kind
|
||||
self.events: list[str] = []
|
||||
self.cleaned_up: list[str] = []
|
||||
|
||||
def emit_created(self, ws: Workstream) -> None:
|
||||
self.events.append(f"created:{ws.id}")
|
||||
|
||||
def emit_rehydrated(self, ws: Workstream) -> None:
|
||||
self.events.append(f"rehydrated:{ws.id}")
|
||||
|
||||
def emit_state(self, ws: Workstream, state: WorkstreamState) -> None:
|
||||
self.events.append(f"state:{ws.id}:{state.value}")
|
||||
|
||||
def emit_closed(self, ws_id: str, *, reason: str = "closed", name: str = "") -> None:
|
||||
self.events.append(f"closed:{ws_id}")
|
||||
|
||||
def cleanup_ui(self, ws: Workstream) -> None:
|
||||
# Real production cleanup_ui calls ws.session.cancel() + close().
|
||||
# We don't need that here — the test exits cleanly via pytest
|
||||
# teardown without exercising the cleanup path. Just record
|
||||
# the call for any test that wants to assert on it.
|
||||
self.cleaned_up.append(ws.id)
|
||||
|
||||
def build_ui(self, ws: Workstream) -> Any:
|
||||
return _FakeUI()
|
||||
|
||||
def build_session(
|
||||
self,
|
||||
ws: Workstream,
|
||||
*,
|
||||
skill: Any = None,
|
||||
model: Any = None,
|
||||
client_type: Any = None,
|
||||
**extra: Any,
|
||||
) -> Any:
|
||||
# Mirror SessionManager.create's keyword set so config-threading
|
||||
# bugs surface here rather than being silently swallowed by
|
||||
# **kwargs. ``model`` flows to the real ChatSession; the rest
|
||||
# are accepted but not used by this test.
|
||||
client = MagicMock()
|
||||
return ChatSession(
|
||||
client=client,
|
||||
model=str(model) if model else "test-model",
|
||||
ui=ws.ui,
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def real_mgr() -> tuple[SessionManager, _BuildRealSessionAdapter]:
|
||||
"""Real SessionManager wired to an adapter that builds real ChatSessions.
|
||||
|
||||
No StateWriter is wired so ``set_state`` writes directly to storage
|
||||
on the calling thread (we want subscriber dispatch to fire in the
|
||||
same thread the test invokes ``set_state`` on).
|
||||
"""
|
||||
adapter = _BuildRealSessionAdapter()
|
||||
storage = FakeStorage()
|
||||
mgr = SessionManager(
|
||||
adapter,
|
||||
storage=storage,
|
||||
max_active=5,
|
||||
event_emitter=adapter,
|
||||
)
|
||||
return mgr, adapter
|
||||
|
||||
|
||||
def _wait_for_worker_done(ws: Workstream, timeout: float = 5.0) -> None:
|
||||
"""Poll ``ws._worker_running`` until it clears or timeout elapses."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
with ws._lock:
|
||||
if not ws._worker_running:
|
||||
return
|
||||
time.sleep(0.01)
|
||||
raise AssertionError(f"worker thread for ws={ws.id[:8]} didn't exit within {timeout}s")
|
||||
|
||||
|
||||
def test_idle_event_through_real_session_manager_drives_wake_send(real_mgr, tmp_db):
|
||||
"""The full wake pipeline, no direct-injection shortcuts.
|
||||
|
||||
Boundary path under test:
|
||||
enqueue → mgr.set_state(IDLE)
|
||||
→ SessionManager._state_subscribers iteration (real)
|
||||
→ IdleNudgeWatcher._on_state (real)
|
||||
→ session_worker.send (real)
|
||||
→ real daemon thread
|
||||
→ ChatSession.deliver_wake_nudge_from_queue (real)
|
||||
→ ChatSession.send("") (real, with patched LLM stream)
|
||||
→ _append_user_turn stamps ``_source``
|
||||
→ _attach_pending_user_reminders drains ``{"user","any"}``
|
||||
→ _apply_reminders_for_provider splices envelope onto empty content
|
||||
"""
|
||||
mgr, _adapter = real_mgr
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
|
||||
try:
|
||||
ws = mgr.create(user_id="u1", name="wake-int", skill=None)
|
||||
assert ws.session is not None
|
||||
# Patch the LLM-facing surface so send() runs the chat loop end-to-end
|
||||
# without any real provider. We patch on the just-built ChatSession;
|
||||
# the patches are reverted by the `with` block.
|
||||
with (
|
||||
patch.object(ws.session, "_create_stream_with_retry", return_value=iter([])),
|
||||
patch.object(
|
||||
ws.session,
|
||||
"_stream_response",
|
||||
return_value={"role": "assistant", "content": "ok"},
|
||||
),
|
||||
patch.object(ws.session, "_update_token_table"),
|
||||
patch.object(ws.session, "_print_status_line"),
|
||||
patch.object(ws.session, "_visible_memory_count", return_value=0),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
# Suppress the auto-title side-thread; orthogonal to wake.
|
||||
ws.session._title_generated = True
|
||||
|
||||
# Enqueue an any-channel nudge — the future ``idle_children`` shape.
|
||||
ws.session._nudge_queue.enqueue("idle_children", "your kids", "any")
|
||||
assert len(ws.session._nudge_queue) == 1
|
||||
|
||||
# Trigger IDLE. This runs subscriber dispatch synchronously on
|
||||
# the calling thread → IdleNudgeWatcher._on_state → session_worker.send
|
||||
# → spawn daemon thread → deliver_wake_nudge_from_queue.
|
||||
mgr.set_state(ws.id, WorkstreamState.IDLE)
|
||||
|
||||
# Wait for the daemon thread to clear ``_worker_running`` so the
|
||||
# post-conditions are stable.
|
||||
_wait_for_worker_done(ws)
|
||||
|
||||
# Queue fully drained by the wake.
|
||||
assert len(ws.session._nudge_queue) == 0
|
||||
|
||||
# The synthesized empty user message landed in history with the
|
||||
# ``_source`` audit tag and the reminder side-channel populated.
|
||||
user_msgs = [m for m in ws.session.messages if m.get("role") == "user"]
|
||||
assert user_msgs, "expected a synthesized user message from the wake"
|
||||
wake_msg = user_msgs[-1]
|
||||
assert wake_msg["content"] == ""
|
||||
assert wake_msg.get("_source") == "system_nudge"
|
||||
assert wake_msg.get("_reminders") == [{"type": "idle_children", "text": "your kids"}]
|
||||
|
||||
# The wake-source tag is reset post-send so subsequent activity
|
||||
# behaves normally.
|
||||
assert ws.session._wake_source_tag == ""
|
||||
finally:
|
||||
watcher.shutdown()
|
||||
|
||||
|
||||
def test_idle_event_with_empty_queue_does_not_dispatch_wake(real_mgr, tmp_db):
|
||||
"""Non-empty queue is the gate. An IDLE event on a workstream with
|
||||
nothing queued must NOT call ``session_worker.send``.
|
||||
|
||||
Patches the dispatch primitive directly rather than racing a
|
||||
``time.sleep`` against an erroneous spawn — the question is
|
||||
whether the watcher's gate fired, which is a deterministic
|
||||
decision the patch captures.
|
||||
"""
|
||||
mgr, _adapter = real_mgr
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
|
||||
try:
|
||||
ws = mgr.create(user_id="u1", name="empty-int", skill=None)
|
||||
# No enqueue.
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.set_state(ws.id, WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 0, "wake must not dispatch for an empty queue"
|
||||
finally:
|
||||
watcher.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def coord_mgr() -> tuple[SessionManager, _BuildRealSessionAdapter, FakeStorage]:
|
||||
"""Real coord-side SessionManager with the adapter's kind set to
|
||||
COORDINATOR. Same shape as ``real_mgr`` but for the coord half of
|
||||
the lifespan. No StateWriter wired so subscriber dispatch fires
|
||||
synchronously on the test thread.
|
||||
"""
|
||||
adapter = _BuildRealSessionAdapter(kind=WorkstreamKind.COORDINATOR)
|
||||
storage = FakeStorage()
|
||||
mgr = SessionManager(
|
||||
adapter,
|
||||
storage=storage,
|
||||
max_active=5,
|
||||
event_emitter=adapter,
|
||||
)
|
||||
return mgr, adapter, storage
|
||||
|
||||
|
||||
def test_coord_idle_with_active_children_emits_envelope_via_real_managers(coord_mgr, tmp_db):
|
||||
"""Full coord-path integration test (matches design doc §7.4).
|
||||
|
||||
Drives the production install order — ``CoordinatorIdleObserver``
|
||||
registered FIRST, then ``IdleNudgeWatcher`` — and asserts the
|
||||
full chain: observer enqueues on IDLE → watcher peeks → wake
|
||||
spawns a worker → ``deliver_wake_nudge_from_queue`` drains and
|
||||
runs the synthetic empty-user turn → reminder envelope reaches
|
||||
the synthesized user message via the side-channel.
|
||||
|
||||
The boundary-crossing path tested here mirrors what
|
||||
``console/server.py``'s lifespan does at production startup; if
|
||||
the install order is ever reversed, this test fails.
|
||||
"""
|
||||
from turnstone.console.coordinator_idle_observer import CoordinatorIdleObserver
|
||||
from turnstone.core.workstream import WorkstreamKind as _Kind
|
||||
|
||||
mgr, adapter, storage = coord_mgr
|
||||
# Observer FIRST, then watcher. Same order as
|
||||
# ``console/server.py:4435-4443`` — production correctness depends
|
||||
# on subscribers firing in registration order on the same IDLE.
|
||||
observer = CoordinatorIdleObserver(mgr, storage)
|
||||
observer.start()
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
|
||||
try:
|
||||
coord = mgr.create(user_id="u1", name="parent-coord", skill=None)
|
||||
assert coord.session is not None
|
||||
|
||||
# Two interactive children of the coord, both running. Use
|
||||
# the storage's register_workstream API so the rows match
|
||||
# production shape (the observer queries via list_workstreams).
|
||||
storage.register_workstream(
|
||||
"child-a",
|
||||
user_id="u1",
|
||||
name="research-pricing",
|
||||
kind=_Kind.INTERACTIVE,
|
||||
parent_ws_id=coord.id,
|
||||
state="running",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"child-b",
|
||||
user_id="u1",
|
||||
name="draft-rfc",
|
||||
kind=_Kind.INTERACTIVE,
|
||||
parent_ws_id=coord.id,
|
||||
state="thinking",
|
||||
)
|
||||
|
||||
# Pretend the coord has already had a real conversation so
|
||||
# ``should_nudge``'s message_count > 1 gate passes.
|
||||
coord.session.messages.append({"role": "user", "content": "spawn 2"})
|
||||
coord.session.messages.append({"role": "assistant", "content": "ok"})
|
||||
|
||||
with (
|
||||
patch.object(coord.session, "_create_stream_with_retry", return_value=iter([])),
|
||||
patch.object(
|
||||
coord.session,
|
||||
"_stream_response",
|
||||
return_value={"role": "assistant", "content": "ack"},
|
||||
),
|
||||
patch.object(coord.session, "_full_messages", return_value=[]),
|
||||
patch.object(coord.session, "_update_token_table"),
|
||||
patch.object(coord.session, "_print_status_line"),
|
||||
patch.object(coord.session, "_visible_memory_count", return_value=0),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
coord.session._title_generated = True
|
||||
mgr.set_state(coord.id, WorkstreamState.IDLE)
|
||||
_wait_for_worker_done(coord)
|
||||
|
||||
# Queue drained — the wake delivered the observer's enqueue.
|
||||
assert len(coord.session._nudge_queue) == 0
|
||||
# The synthetic empty-user turn landed with a reminder containing
|
||||
# both children.
|
||||
user_msgs = [m for m in coord.session.messages if m.get("role") == "user"]
|
||||
# Two real msgs (user + assistant context above) plus the wake.
|
||||
wake_msg = user_msgs[-1]
|
||||
assert wake_msg["content"] == ""
|
||||
assert wake_msg.get("_source") == "system_nudge"
|
||||
reminders = wake_msg.get("_reminders") or []
|
||||
assert len(reminders) == 1
|
||||
assert reminders[0]["type"] == "idle_children"
|
||||
text = reminders[0]["text"]
|
||||
assert "research-pricing" in text
|
||||
assert "draft-rfc" in text
|
||||
assert "child-a" in text
|
||||
assert "child-b" in text
|
||||
assert "wait_for_workstream" in text
|
||||
finally:
|
||||
watcher.shutdown()
|
||||
observer.shutdown()
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Unit tests for :class:`IdleNudgeWatcher`.
|
||||
|
||||
Drives a fake :class:`SessionManager` that mimics the real one's
|
||||
``subscribe_to_state`` / ``get`` contract. The watcher itself
|
||||
dispatches via ``turnstone.core.session_worker.send``; we patch that
|
||||
module-level function to capture calls without spawning real threads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import threading
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.idle_nudge_watcher import IdleNudgeWatcher
|
||||
from turnstone.core.nudge_queue import NudgeQueue
|
||||
from turnstone.core.workstream import WorkstreamState
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self._nudge_queue = NudgeQueue()
|
||||
self.deliver_wake_nudge_from_queue_called = 0
|
||||
|
||||
def deliver_wake_nudge_from_queue(self) -> None:
|
||||
self.deliver_wake_nudge_from_queue_called += 1
|
||||
|
||||
|
||||
class _FakeWorkstream:
|
||||
def __init__(self, ws_id: str = "ws-test") -> None:
|
||||
self.id = ws_id
|
||||
self.session: _FakeSession | None = _FakeSession()
|
||||
self._lock = threading.Lock()
|
||||
self._worker_running = False
|
||||
self._closed = False
|
||||
self.worker_thread: Any = None
|
||||
|
||||
|
||||
class _FakeManager:
|
||||
"""Mimics SessionManager's subscribe-to-state surface without a DB."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._workstreams: dict[str, _FakeWorkstream] = {}
|
||||
self._subscribers: list[Any] = []
|
||||
self._subscribers_lock = threading.Lock()
|
||||
|
||||
def add_ws(self, ws: _FakeWorkstream) -> None:
|
||||
self._workstreams[ws.id] = ws
|
||||
|
||||
def get(self, ws_id: str) -> _FakeWorkstream | None:
|
||||
return self._workstreams.get(ws_id)
|
||||
|
||||
def subscribe_to_state(self, callback: Any) -> None:
|
||||
with self._subscribers_lock:
|
||||
self._subscribers.append(callback)
|
||||
|
||||
def unsubscribe_from_state(self, callback: Any) -> None:
|
||||
with self._subscribers_lock, contextlib.suppress(ValueError):
|
||||
self._subscribers.remove(callback)
|
||||
|
||||
def fire_state(self, ws_id: str, state: WorkstreamState) -> None:
|
||||
"""Mirror SessionManager.set_state's subscriber-fan-out behaviour."""
|
||||
with self._subscribers_lock:
|
||||
subs = list(self._subscribers)
|
||||
for cb in subs:
|
||||
# Match contextlib.suppress(Exception) in real SessionManager.
|
||||
with contextlib.suppress(Exception):
|
||||
cb(ws_id, state)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_mgr_and_ws() -> tuple[_FakeManager, _FakeWorkstream]:
|
||||
mgr = _FakeManager()
|
||||
ws = _FakeWorkstream()
|
||||
mgr.add_ws(ws)
|
||||
return mgr, ws
|
||||
|
||||
|
||||
class TestIdleNudgeWatcher:
|
||||
def test_idle_event_with_empty_queue_no_op(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 0
|
||||
|
||||
def test_idle_event_with_pending_nudge_dispatches(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
ws.session._nudge_queue.enqueue("idle_children", "your kids", "any")
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 1
|
||||
kwargs = mock_send.call_args.kwargs
|
||||
# `enqueue=lambda: None` — verify by calling and checking no-op.
|
||||
assert kwargs["enqueue"]() is None
|
||||
# `run` should call deliver_wake_nudge_from_queue when invoked.
|
||||
kwargs["run"]()
|
||||
assert ws.session.deliver_wake_nudge_from_queue_called == 1
|
||||
assert kwargs["thread_name"].startswith("wake-nudge-")
|
||||
|
||||
def test_non_idle_state_ignored(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
ws.session._nudge_queue.enqueue("foo", "bar", "any")
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
for state in (
|
||||
WorkstreamState.RUNNING,
|
||||
WorkstreamState.THINKING,
|
||||
WorkstreamState.ATTENTION,
|
||||
WorkstreamState.ERROR,
|
||||
):
|
||||
mgr.fire_state(ws.id, state)
|
||||
assert mock_send.call_count == 0
|
||||
|
||||
def test_unknown_ws_ignored(self, fake_mgr_and_ws):
|
||||
mgr, _ws = fake_mgr_and_ws
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state("ghost", WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 0
|
||||
|
||||
def test_session_none_ignored(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
ws.session = None # workstream loaded but session not yet built
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 0
|
||||
|
||||
def test_start_is_idempotent(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
watcher.start() # no-op
|
||||
ws.session._nudge_queue.enqueue("foo", "bar", "any")
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
# Only one subscriber was registered despite the double-start.
|
||||
assert mock_send.call_count == 1
|
||||
|
||||
def test_shutdown_unsubscribes(self, fake_mgr_and_ws):
|
||||
mgr, ws = fake_mgr_and_ws
|
||||
ws.session._nudge_queue.enqueue("foo", "bar", "any")
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
watcher.shutdown()
|
||||
with patch("turnstone.core.session_worker.send") as mock_send:
|
||||
mgr.fire_state(ws.id, WorkstreamState.IDLE)
|
||||
assert mock_send.call_count == 0
|
||||
|
||||
def test_shutdown_is_idempotent(self, fake_mgr_and_ws):
|
||||
mgr, _ws = fake_mgr_and_ws
|
||||
watcher = IdleNudgeWatcher(mgr)
|
||||
watcher.start()
|
||||
watcher.shutdown()
|
||||
watcher.shutdown() # no error
|
||||
@@ -420,8 +420,9 @@ class TestSkillCatalogDisclosure:
|
||||
session.system_messages = []
|
||||
session._agent_system_messages = []
|
||||
session.reasoning_effort = "medium"
|
||||
session._pending_tool_advisories = []
|
||||
session._pending_user_advisories = []
|
||||
from turnstone.core.nudge_queue import NudgeQueue
|
||||
|
||||
session._nudge_queue = NudgeQueue()
|
||||
session._tool_search = None
|
||||
session._mcp_client = None
|
||||
session._notify_on_complete = "{}"
|
||||
|
||||
+1049
-2
File diff suppressed because it is too large
Load Diff
+786
-148
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,117 @@
|
||||
"""Structural gate against the Phase 7b sibling-bug pattern.
|
||||
|
||||
Phase 7b's bug-1 was a single ``f"MCP X error: {e}"`` site dropping a
|
||||
structured-error JSON. Phase 8 introduces the ``consent_url`` field on
|
||||
the same JSON envelope: every ``_structured_error(...)`` invocation
|
||||
that emits ``mcp_consent_required`` or ``mcp_insufficient_scope`` MUST
|
||||
also pass a ``consent_url=`` kwarg, otherwise the dashboard renderer
|
||||
can't surface a re-consent button.
|
||||
|
||||
This test is purely structural — it scans the source of
|
||||
:mod:`turnstone.core.mcp_client` and asserts every consent-required /
|
||||
insufficient-scope ``_structured_error`` call carries
|
||||
``consent_url=``. It catches future regressions where a new exec path
|
||||
adds a fourth call site and forgets the kwarg.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import turnstone.core.mcp_client as _mcp_client_module
|
||||
|
||||
_USER_ACTIONABLE_CODES = ("mcp_consent_required", "mcp_insufficient_scope")
|
||||
|
||||
|
||||
def _read_source() -> str:
|
||||
path = Path(_mcp_client_module.__file__)
|
||||
return path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def _find_structured_error_blocks(source: str) -> list[tuple[int, str]]:
|
||||
"""Return ``(line_no, block)`` pairs for every ``_structured_error(...)``.
|
||||
|
||||
Each block is the call's argument list expanded across however many
|
||||
lines the formatter chose. Uses a paren-counting walk so multi-line
|
||||
kwargs and nested expressions are captured correctly.
|
||||
"""
|
||||
blocks: list[tuple[int, str]] = []
|
||||
needle = "_structured_error("
|
||||
idx = 0
|
||||
while True:
|
||||
loc = source.find(needle, idx)
|
||||
if loc < 0:
|
||||
break
|
||||
# Skip the function definition itself.
|
||||
if source[loc - 4 : loc] == "def ":
|
||||
idx = loc + len(needle)
|
||||
continue
|
||||
line_no = source.count("\n", 0, loc) + 1
|
||||
depth = 1
|
||||
end = loc + len(needle)
|
||||
while end < len(source) and depth > 0:
|
||||
ch = source[end]
|
||||
if ch == "(":
|
||||
depth += 1
|
||||
elif ch == ")":
|
||||
depth -= 1
|
||||
end += 1
|
||||
blocks.append((line_no, source[loc:end]))
|
||||
idx = end
|
||||
return blocks
|
||||
|
||||
|
||||
def test_every_user_actionable_structured_error_passes_consent_url() -> None:
|
||||
source = _read_source()
|
||||
blocks = _find_structured_error_blocks(source)
|
||||
user_actionable_blocks = [
|
||||
(ln, blk)
|
||||
for ln, blk in blocks
|
||||
if any(f'code="{code}"' in blk for code in _USER_ACTIONABLE_CODES)
|
||||
]
|
||||
|
||||
# Sanity check: ensure we actually scanned the file the audit cares
|
||||
# about (a stale path or import would otherwise silently pass with
|
||||
# zero matches).
|
||||
assert user_actionable_blocks, (
|
||||
"No mcp_consent_required / mcp_insufficient_scope _structured_error "
|
||||
"call sites found — has the audit been pointed at the wrong file?"
|
||||
)
|
||||
|
||||
missing: list[tuple[int, str]] = []
|
||||
for ln, blk in user_actionable_blocks:
|
||||
if "consent_url=" not in blk:
|
||||
# Strip whitespace and truncate so the failure message is
|
||||
# readable in CI.
|
||||
collapsed = re.sub(r"\s+", " ", blk).strip()
|
||||
missing.append((ln, collapsed[:200]))
|
||||
|
||||
assert not missing, (
|
||||
"Sibling-bug regression: the following consent-required / "
|
||||
"insufficient-scope _structured_error sites are missing the "
|
||||
"consent_url= kwarg.\n" + "\n".join(f" line {ln}: {snippet}" for ln, snippet in missing)
|
||||
)
|
||||
|
||||
|
||||
def test_audit_finds_all_known_user_actionable_sites() -> None:
|
||||
"""Lock the count so accidental deletions are caught.
|
||||
|
||||
There are 13 user-actionable ``_structured_error`` call sites today
|
||||
(4 each in the tool / resource / prompt token-classify branches +
|
||||
3 in the post-retry-failed branches + 1 in ``_handle_auth_403``'s
|
||||
insufficient-scope branch). If a new exec path is added the count
|
||||
can rise; if a branch is removed the count can fall — both are
|
||||
fine, but require an intentional bump of this number to confirm
|
||||
the change went through review.
|
||||
"""
|
||||
source = _read_source()
|
||||
blocks = _find_structured_error_blocks(source)
|
||||
user_actionable_count = sum(
|
||||
1 for _, blk in blocks if any(f'code="{code}"' in blk for code in _USER_ACTIONABLE_CODES)
|
||||
)
|
||||
assert user_actionable_count == 13, (
|
||||
f"Expected 13 user-actionable _structured_error sites, got "
|
||||
f"{user_actionable_count}. If this is intentional, bump the "
|
||||
f"expected count and document why in the commit message."
|
||||
)
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Tests for ``turnstone.core.mcp_crypto`` cipher + config loading.
|
||||
|
||||
Covers token-at-rest encryption for OAuth-MCP.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from turnstone.core.mcp_crypto import (
|
||||
MCPTokenCipher,
|
||||
MCPTokenCipherConfig,
|
||||
MCPTokenDecryptError,
|
||||
MCPTokenKeyConfigError,
|
||||
_key_fingerprint,
|
||||
_validate_key,
|
||||
load_mcp_token_cipher_config,
|
||||
)
|
||||
|
||||
|
||||
def _new_raw_key() -> bytes:
|
||||
"""Return a fresh 32-byte Fernet key as raw bytes (post-base64-decode)."""
|
||||
return base64.urlsafe_b64decode(Fernet.generate_key())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cipher round-trip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCipherRoundTrip:
|
||||
def test_round_trip_single_key(self) -> None:
|
||||
cipher = MCPTokenCipher(MCPTokenCipherConfig(keys=(_new_raw_key(),)))
|
||||
plaintext = b"access_token_12345"
|
||||
ct = cipher.encrypt(plaintext)
|
||||
assert ct != plaintext
|
||||
assert cipher.decrypt(ct) == plaintext
|
||||
|
||||
def test_round_trip_unicode_token(self) -> None:
|
||||
cipher = MCPTokenCipher(MCPTokenCipherConfig(keys=(_new_raw_key(),)))
|
||||
# Tokens may legitimately carry UTF-8 bytes (e.g. JWT with
|
||||
# non-ASCII claim values). Round-trip a multi-byte sequence.
|
||||
plaintext = "tok_é中💯".encode()
|
||||
ct = cipher.encrypt(plaintext)
|
||||
assert cipher.decrypt(ct) == plaintext
|
||||
|
||||
def test_wrong_key_raises_decrypt_error(self) -> None:
|
||||
cipher_a = MCPTokenCipher(MCPTokenCipherConfig(keys=(_new_raw_key(),)))
|
||||
cipher_b = MCPTokenCipher(MCPTokenCipherConfig(keys=(_new_raw_key(),)))
|
||||
ct = cipher_a.encrypt(b"secret")
|
||||
with pytest.raises(MCPTokenDecryptError) as exc_info:
|
||||
cipher_b.decrypt(ct)
|
||||
# Audit-trail correlation: error must carry the fingerprints of
|
||||
# the keys actually attempted, not a placeholder.
|
||||
assert exc_info.value.key_fingerprints_attempted
|
||||
assert exc_info.value.key_fingerprints_attempted == cipher_b.key_fingerprints
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Rotation (MultiFernet behavior)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRotation:
|
||||
def test_rotation_forward(self) -> None:
|
||||
"""Encrypt with a new-only cipher, decrypt with a [v2, v1] cluster.
|
||||
|
||||
Mirrors the operational situation where a node already has the
|
||||
rotated key list installed and a peer just wrote a row under v2.
|
||||
"""
|
||||
v1 = _new_raw_key()
|
||||
v2 = _new_raw_key()
|
||||
new_only = MCPTokenCipher(MCPTokenCipherConfig(keys=(v2,)))
|
||||
cluster = MCPTokenCipher(MCPTokenCipherConfig(keys=(v2, v1)))
|
||||
ct = new_only.encrypt(b"hello")
|
||||
assert cluster.decrypt(ct) == b"hello"
|
||||
|
||||
def test_rotation_backward_keeps_old_decryptable(self) -> None:
|
||||
"""A row written under the OLD key (v1) must still decrypt after
|
||||
rotation places v2 first and keeps v1 as fallback."""
|
||||
v1 = _new_raw_key()
|
||||
v2 = _new_raw_key()
|
||||
old_only = MCPTokenCipher(MCPTokenCipherConfig(keys=(v1,)))
|
||||
rotated = MCPTokenCipher(MCPTokenCipherConfig(keys=(v2, v1)))
|
||||
ct = old_only.encrypt(b"legacy")
|
||||
assert rotated.decrypt(ct) == b"legacy"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config loader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _patch_load_config(monkeypatch: pytest.MonkeyPatch, payload: dict) -> None:
|
||||
"""Override ``turnstone.core.config.load_config`` to return ``payload``
|
||||
when the ``"security"`` section is requested."""
|
||||
|
||||
def fake(section: str | None = None) -> dict:
|
||||
if section == "security":
|
||||
return payload
|
||||
return {}
|
||||
|
||||
import turnstone.core.config as cfg_mod
|
||||
|
||||
monkeypatch.setattr(cfg_mod, "load_config", fake)
|
||||
|
||||
|
||||
class TestLoadConfig:
|
||||
def test_load_singular_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
key = Fernet.generate_key().decode()
|
||||
_patch_load_config(monkeypatch, {"mcp_token_encryption_key": key})
|
||||
cfg = load_mcp_token_cipher_config()
|
||||
assert cfg is not None
|
||||
assert len(cfg.keys) == 1
|
||||
|
||||
def test_load_plural_overrides_singular(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
plural = [Fernet.generate_key().decode(), Fernet.generate_key().decode()]
|
||||
_patch_load_config(
|
||||
monkeypatch,
|
||||
{
|
||||
"mcp_token_encryption_keys": plural,
|
||||
"mcp_token_encryption_key": Fernet.generate_key().decode(),
|
||||
},
|
||||
)
|
||||
cfg = load_mcp_token_cipher_config()
|
||||
assert cfg is not None
|
||||
assert len(cfg.keys) == 2 # plural wins, singular ignored
|
||||
|
||||
def test_load_returns_none_when_absent(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_load_config(monkeypatch, {})
|
||||
assert load_mcp_token_cipher_config() is None
|
||||
|
||||
def test_load_empty_plural_falls_through_to_singular(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Operator wrote ``mcp_token_encryption_keys = []`` AND set a
|
||||
singular value: empty plural is treated as absent."""
|
||||
key = Fernet.generate_key().decode()
|
||||
_patch_load_config(
|
||||
monkeypatch,
|
||||
{"mcp_token_encryption_keys": [], "mcp_token_encryption_key": key},
|
||||
)
|
||||
cfg = load_mcp_token_cipher_config()
|
||||
assert cfg is not None
|
||||
assert len(cfg.keys) == 1
|
||||
|
||||
def test_load_invalid_base64_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_load_config(monkeypatch, {"mcp_token_encryption_key": "###not-base64###"})
|
||||
with pytest.raises(MCPTokenKeyConfigError) as exc_info:
|
||||
load_mcp_token_cipher_config()
|
||||
# Operator-facing hint is part of every error message.
|
||||
assert "regenerate with:" in str(exc_info.value)
|
||||
|
||||
def test_load_wrong_length_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# 24 raw bytes → 32 base64 chars; not 32 raw bytes after decode.
|
||||
short_key = base64.urlsafe_b64encode(b"\x00" * 24).decode()
|
||||
_patch_load_config(monkeypatch, {"mcp_token_encryption_key": short_key})
|
||||
with pytest.raises(MCPTokenKeyConfigError) as exc_info:
|
||||
load_mcp_token_cipher_config()
|
||||
assert "32 bytes" in str(exc_info.value)
|
||||
|
||||
def test_load_non_list_plural_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_load_config(monkeypatch, {"mcp_token_encryption_keys": "single-string-not-list"})
|
||||
with pytest.raises(MCPTokenKeyConfigError) as exc_info:
|
||||
load_mcp_token_cipher_config()
|
||||
assert "list" in str(exc_info.value).lower()
|
||||
|
||||
def test_load_non_string_in_plural_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_patch_load_config(monkeypatch, {"mcp_token_encryption_keys": [12345]})
|
||||
with pytest.raises(MCPTokenKeyConfigError):
|
||||
load_mcp_token_cipher_config()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fingerprint stability
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFingerprint:
|
||||
def test_key_fingerprint_stable_and_short(self) -> None:
|
||||
key = _new_raw_key()
|
||||
fp1 = _key_fingerprint(key)
|
||||
fp2 = _key_fingerprint(key)
|
||||
assert fp1 == fp2
|
||||
# 8 bytes -> 16 hex characters.
|
||||
assert len(fp1) == 16
|
||||
assert all(c in "0123456789abcdef" for c in fp1)
|
||||
|
||||
def test_different_keys_have_different_fingerprints(self) -> None:
|
||||
fp1 = _key_fingerprint(_new_raw_key())
|
||||
fp2 = _key_fingerprint(_new_raw_key())
|
||||
assert fp1 != fp2
|
||||
|
||||
def test_cipher_fingerprints_match_keys(self) -> None:
|
||||
v1 = _new_raw_key()
|
||||
v2 = _new_raw_key()
|
||||
cipher = MCPTokenCipher(MCPTokenCipherConfig(keys=(v1, v2)))
|
||||
assert cipher.key_fingerprints == (
|
||||
_key_fingerprint(v1),
|
||||
_key_fingerprint(v2),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Direct ``_validate_key`` — exercises edge cases not reachable via loader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestValidateKey:
|
||||
def test_empty_string_rejected(self) -> None:
|
||||
with pytest.raises(MCPTokenKeyConfigError):
|
||||
_validate_key("", label="x")
|
||||
|
||||
def test_whitespace_only_rejected(self) -> None:
|
||||
with pytest.raises(MCPTokenKeyConfigError):
|
||||
_validate_key(" ", label="x")
|
||||
|
||||
def test_label_propagated_in_error(self) -> None:
|
||||
with pytest.raises(MCPTokenKeyConfigError) as exc_info:
|
||||
_validate_key("###", label="my_label_42")
|
||||
assert "my_label_42" in str(exc_info.value)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCPTokenCipher constructor guard
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCipherConstructorGuard:
|
||||
def test_empty_keys_rejected(self) -> None:
|
||||
with pytest.raises(MCPTokenKeyConfigError):
|
||||
MCPTokenCipher(MCPTokenCipherConfig(keys=()))
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tests.conftest import _seed_static_state
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -105,14 +106,18 @@ class TestRemoveServerSync:
|
||||
"""remove_server_sync cleans up all per-server state dicts."""
|
||||
mgr = MCPClientManager({"test": {"command": "echo"}})
|
||||
# Simulate state as if the server was connected
|
||||
mgr._per_server_tools["test"] = [_fake_openai_tool()]
|
||||
mgr._per_server_resources["test"] = [_fake_resource_dict()]
|
||||
mgr._per_server_prompts["test"] = [_fake_prompt_dict()]
|
||||
mgr._supports_list_changed["test"] = True
|
||||
mgr._supports_resources["test"] = True
|
||||
mgr._supports_resource_list_changed["test"] = True
|
||||
mgr._supports_prompts["test"] = True
|
||||
mgr._supports_prompt_list_changed["test"] = True
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"test",
|
||||
tools=[_fake_openai_tool()],
|
||||
resources=[_fake_resource_dict()],
|
||||
prompts=[_fake_prompt_dict()],
|
||||
supports_list_changed=True,
|
||||
supports_resources=True,
|
||||
supports_resource_list_changed=True,
|
||||
supports_prompts=True,
|
||||
supports_prompt_list_changed=True,
|
||||
)
|
||||
mgr._rebuild_tools()
|
||||
mgr._rebuild_resources()
|
||||
mgr._rebuild_prompts()
|
||||
@@ -127,14 +132,7 @@ class TestRemoveServerSync:
|
||||
assert len(mgr.get_tools()) == 0
|
||||
assert mgr.resource_count == 0
|
||||
assert mgr.prompt_count == 0
|
||||
assert "test" not in mgr._per_server_tools
|
||||
assert "test" not in mgr._per_server_resources
|
||||
assert "test" not in mgr._per_server_prompts
|
||||
assert "test" not in mgr._supports_list_changed
|
||||
assert "test" not in mgr._supports_resources
|
||||
assert "test" not in mgr._supports_resource_list_changed
|
||||
assert "test" not in mgr._supports_prompts
|
||||
assert "test" not in mgr._supports_prompt_list_changed
|
||||
assert "test" not in mgr._static_servers
|
||||
|
||||
def test_removes_config_to_prevent_reconnect(self) -> None:
|
||||
"""remove_server_sync removes from _server_configs to prevent reconnect."""
|
||||
@@ -146,8 +144,8 @@ class TestRemoveServerSync:
|
||||
def test_preserves_other_servers(self) -> None:
|
||||
"""Removing one server does not affect another server's state."""
|
||||
mgr = MCPClientManager({"srv_a": {}, "srv_b": {}})
|
||||
mgr._per_server_tools["srv_a"] = [_fake_openai_tool("mcp__srv_a__foo")]
|
||||
mgr._per_server_tools["srv_b"] = [_fake_openai_tool("mcp__srv_b__bar")]
|
||||
_seed_static_state(mgr, "srv_a", tools=[_fake_openai_tool("mcp__srv_a__foo")])
|
||||
_seed_static_state(mgr, "srv_b", tools=[_fake_openai_tool("mcp__srv_b__bar")])
|
||||
mgr._rebuild_tools()
|
||||
|
||||
assert len(mgr.get_tools()) == 2
|
||||
@@ -179,13 +177,17 @@ class TestGetServerStatus:
|
||||
"""Status of a connected server reports correct tool/resource/prompt counts."""
|
||||
mgr = MCPClientManager({"test": {}})
|
||||
# Simulate connected state
|
||||
mgr._sessions["test"] = object() # any truthy value
|
||||
mgr._per_server_tools["test"] = [
|
||||
_fake_openai_tool("mcp__test__a"),
|
||||
_fake_openai_tool("mcp__test__b"),
|
||||
]
|
||||
mgr._per_server_resources["test"] = [_fake_resource_dict()]
|
||||
mgr._per_server_prompts["test"] = [_fake_prompt_dict()]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"test",
|
||||
session=object(), # any truthy value
|
||||
tools=[
|
||||
_fake_openai_tool("mcp__test__a"),
|
||||
_fake_openai_tool("mcp__test__b"),
|
||||
],
|
||||
resources=[_fake_resource_dict()],
|
||||
prompts=[_fake_prompt_dict()],
|
||||
)
|
||||
|
||||
status = mgr.get_server_status("test")
|
||||
assert status["connected"] is True
|
||||
@@ -225,8 +227,7 @@ class TestGetAllServerStatus:
|
||||
def test_mixed_connected_and_disconnected(self) -> None:
|
||||
"""Status correctly reflects a mix of connected and disconnected servers."""
|
||||
mgr = MCPClientManager({"up": {}, "down": {}})
|
||||
mgr._sessions["up"] = object()
|
||||
mgr._per_server_tools["up"] = [_fake_openai_tool("mcp__up__x")]
|
||||
_seed_static_state(mgr, "up", session=object(), tools=[_fake_openai_tool("mcp__up__x")])
|
||||
|
||||
statuses = mgr.get_all_server_status()
|
||||
assert statuses["up"]["connected"] is True
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Unit tests for ``turnstone.core.mcp_http_parsers``.
|
||||
|
||||
The parser replaces the prior hand-rolled scanners that used
|
||||
``header.lower().find("scope")`` to locate parameter names — that approach
|
||||
misparsed ``scope`` embedded inside other tokens (``xscope``) or inside
|
||||
quoted-string values of preceding params. Each adversarial case below
|
||||
asserts the new tokenizer respects RFC 7235 ``challenge → auth-param``
|
||||
boundaries; the docstrings document the equivalent input that broke the
|
||||
naive parser. Negative-test verification: temporarily reverting
|
||||
``parse_www_authenticate_scope`` to delegate to ``header.lower().find("scope")``
|
||||
makes ``test_scope_inside_realm_value`` and ``test_scope_inside_xscope`` fail.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.mcp_http_parsers import (
|
||||
parse_www_authenticate_bearer,
|
||||
parse_www_authenticate_error,
|
||||
parse_www_authenticate_scope,
|
||||
)
|
||||
|
||||
|
||||
class TestParseScope:
|
||||
def test_basic_scope(self) -> None:
|
||||
header = 'Bearer error="insufficient_scope", scope="files:read mail:send"'
|
||||
assert parse_www_authenticate_scope(header) == ("files:read", "mail:send")
|
||||
|
||||
def test_no_scope_param(self) -> None:
|
||||
assert parse_www_authenticate_scope('Bearer error="invalid_token"') == ()
|
||||
|
||||
def test_unterminated_quoted_string_returns_empty(self) -> None:
|
||||
assert parse_www_authenticate_scope('Bearer scope="files:read') == ()
|
||||
|
||||
def test_escaped_chars_in_value_drops_invalid_scope_token(self) -> None:
|
||||
# RFC 7230 §3.2.6 backslash escapes decode the literal scope to
|
||||
# ``files:read "weird"``. RFC 6749 §3.3 ``scope-token`` forbids
|
||||
# ``"``, so ``"weird"`` is dropped and only ``files:read``
|
||||
# survives the post-split validation.
|
||||
header = r'Bearer scope="files:read \"weird\""'
|
||||
assert parse_www_authenticate_scope(header) == ("files:read",)
|
||||
|
||||
def test_empty_string(self) -> None:
|
||||
assert parse_www_authenticate_scope("") == ()
|
||||
|
||||
def test_unquoted_scope_value(self) -> None:
|
||||
# Unquoted single token.
|
||||
assert parse_www_authenticate_scope("Bearer scope=files:read") == ("files:read",)
|
||||
|
||||
# --- the four headline misparse cases ---
|
||||
|
||||
def test_scope_inside_xscope(self) -> None:
|
||||
"""``Bearer xscope="value"`` must NOT be read as ``scope``.
|
||||
|
||||
The naive ``find("scope")`` matched at position 7 inside
|
||||
``xscope`` and returned ``("value",)``.
|
||||
"""
|
||||
assert parse_www_authenticate_scope('Bearer xscope="value"') == ()
|
||||
|
||||
def test_scope_inside_realm_value(self) -> None:
|
||||
"""``Bearer realm="my scope=fake", scope="real"`` must return ``("real",)``.
|
||||
|
||||
The naive parser found ``scope=`` inside the quoted ``realm``
|
||||
value first and returned ``("fake",)``.
|
||||
"""
|
||||
header = 'Bearer realm="my scope=fake", scope="real"'
|
||||
assert parse_www_authenticate_scope(header) == ("real",)
|
||||
|
||||
def test_scope_inside_quoted_realm_with_escaped_quotes(self) -> None:
|
||||
"""``Bearer realm="foo scope=\\"admin:write\\" bar"`` returns ``()``.
|
||||
|
||||
The inner ``scope=`` is wholly inside the quoted-string value of
|
||||
``realm`` — there is no top-level ``scope`` auth-param, so the
|
||||
result is empty.
|
||||
"""
|
||||
header = r'Bearer realm="foo scope=\"admin:write\" bar"'
|
||||
assert parse_www_authenticate_scope(header) == ()
|
||||
|
||||
def test_scope_token_validation_drops_control_bytes(self) -> None:
|
||||
"""Tokens containing CR / LF / tab / DEL / quote are dropped.
|
||||
|
||||
RFC 6749 §3.3 restricts ``scope-token`` to visible ASCII
|
||||
excluding ``"`` and ``\\``. The splitter applies that
|
||||
validation so a malicious AS cannot smuggle CRLF (or the like)
|
||||
through a future log / notification path that prints the scope
|
||||
list verbatim. ``"a\\rb"`` and ``"\\nc"`` fail validation;
|
||||
``"d"`` survives. The legitimate space separator splits ``d``
|
||||
into its own token.
|
||||
"""
|
||||
# Build via concatenation so the assertion stays intelligible.
|
||||
header = 'Bearer scope="a\rb \nc d"'
|
||||
assert parse_www_authenticate_scope(header) == ("d",)
|
||||
|
||||
|
||||
class TestParseError:
|
||||
def test_basic_quoted_error(self) -> None:
|
||||
assert (
|
||||
parse_www_authenticate_error('Bearer error="insufficient_scope"')
|
||||
== "insufficient_scope"
|
||||
)
|
||||
|
||||
def test_other_quoted_error_tokens(self) -> None:
|
||||
assert parse_www_authenticate_error('Bearer error="invalid_token"') == "invalid_token"
|
||||
assert parse_www_authenticate_error('Bearer error="invalid_request"') == "invalid_request"
|
||||
|
||||
def test_no_error_param(self) -> None:
|
||||
assert parse_www_authenticate_error("Bearer realm=foo") is None
|
||||
|
||||
def test_error_description_does_not_match_error(self) -> None:
|
||||
"""``error_description`` is its own auth-param key, not ``error``.
|
||||
|
||||
The tokenizer reads ``_`` as part of the token (RFC 7230 ``tchar``),
|
||||
so ``error_description`` becomes one key, ``error`` another.
|
||||
"""
|
||||
assert parse_www_authenticate_error('Bearer error_description="bad"') is None
|
||||
|
||||
def test_unquoted_error(self) -> None:
|
||||
# Some ASes don't quote the error token.
|
||||
assert (
|
||||
parse_www_authenticate_error("Bearer error=insufficient_scope") == "insufficient_scope"
|
||||
)
|
||||
|
||||
def test_empty_string(self) -> None:
|
||||
assert parse_www_authenticate_error("") is None
|
||||
|
||||
def test_error_inside_realm_value(self) -> None:
|
||||
"""``Bearer realm="my error=fake", error="real"`` must return ``"real"``.
|
||||
|
||||
Naive parser grabbed ``fake`` from inside the ``realm`` quoted
|
||||
value.
|
||||
"""
|
||||
header = 'Bearer realm="my error=fake", error="real"'
|
||||
assert parse_www_authenticate_error(header) == "real"
|
||||
|
||||
|
||||
class TestBearerDict:
|
||||
def test_returns_lowercased_keys(self) -> None:
|
||||
header = 'Bearer Realm="x", Error="y", Scope="a b"'
|
||||
params = parse_www_authenticate_bearer(header)
|
||||
assert params == {"realm": "x", "error": "y", "scope": "a b"}
|
||||
|
||||
def test_non_bearer_scheme_returns_empty(self) -> None:
|
||||
assert parse_www_authenticate_bearer('Basic realm="x"') == {}
|
||||
|
||||
def test_no_scheme(self) -> None:
|
||||
assert parse_www_authenticate_bearer('realm="x"') == {}
|
||||
|
||||
def test_bearer_only_no_params(self) -> None:
|
||||
assert parse_www_authenticate_bearer("Bearer ") == {}
|
||||
|
||||
def test_bearer_with_no_space_returns_empty(self) -> None:
|
||||
# ``BearerToken`` is not a Bearer challenge (no separator).
|
||||
assert parse_www_authenticate_bearer("BearerToken") == {}
|
||||
|
||||
def test_first_value_wins_on_duplicate(self) -> None:
|
||||
# If a malformed AS sends two ``scope=`` params we keep the first.
|
||||
# The earlier ``find()``-based scanner would have returned the
|
||||
# last; either choice is legal for malformed input but we need
|
||||
# to be consistent.
|
||||
header = 'Bearer scope="first", scope="second"'
|
||||
assert parse_www_authenticate_bearer(header) == {"scope": "first"}
|
||||
|
||||
def test_trailing_comma(self) -> None:
|
||||
header = 'Bearer error="x",'
|
||||
assert parse_www_authenticate_bearer(header) == {"error": "x"}
|
||||
|
||||
def test_multiple_commas(self) -> None:
|
||||
header = 'Bearer ,, error="x",,, scope="y"'
|
||||
assert parse_www_authenticate_bearer(header) == {"error": "x", "scope": "y"}
|
||||
|
||||
def test_embedded_escaped_quote(self) -> None:
|
||||
header = r'Bearer realm="he said \"hi\""'
|
||||
assert parse_www_authenticate_bearer(header) == {"realm": 'he said "hi"'}
|
||||
|
||||
def test_param_without_value_skipped(self) -> None:
|
||||
header = 'Bearer realm, error="x"'
|
||||
# ``realm`` without ``=`` is dropped; ``error`` survives.
|
||||
assert parse_www_authenticate_bearer(header) == {"error": "x"}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"header,expected",
|
||||
[
|
||||
("", {}),
|
||||
("Bearer", {}),
|
||||
('Bearer realm=""', {"realm": ""}),
|
||||
('Bearer realm="", scope=""', {"realm": "", "scope": ""}),
|
||||
],
|
||||
)
|
||||
def test_edge_cases(self, header: str, expected: dict[str, str]) -> None:
|
||||
assert parse_www_authenticate_bearer(header) == expected
|
||||
|
||||
|
||||
class TestPathologicalInput:
|
||||
def test_oversized_pathological_input_rejected_under_50ms(self) -> None:
|
||||
"""Headers longer than the defensive cap return ``{}`` immediately.
|
||||
|
||||
The cap is set to 4096 bytes — real ASes emit a few hundred bytes
|
||||
at most. This guards both ``parse_www_authenticate_bearer``
|
||||
callers against pathological input from a misbehaving server.
|
||||
The previous ``header.lower().find("scope", i)`` loop was
|
||||
O(N**2) — a 100 KB header with no ``=`` took ~330 ms because
|
||||
each ``find`` rescanned the entire suffix. The single-pass
|
||||
tokenizer (capped at 4 KB) reduces this to a one-shot length
|
||||
check that returns ``{}`` in microseconds, so the budget is
|
||||
generous regardless of which side of the cap was hit.
|
||||
"""
|
||||
big = "Bearer scope=" + "a" * 10_000
|
||||
start = time.perf_counter()
|
||||
result = parse_www_authenticate_scope(big)
|
||||
elapsed = time.perf_counter() - start
|
||||
assert result == ()
|
||||
assert elapsed < 0.05, f"oversized-header reject took {elapsed * 1000:.1f}ms"
|
||||
|
||||
def test_within_cap_long_header_under_50ms(self) -> None:
|
||||
"""A 4 KB header with thousands of ``find`` candidates still parses fast.
|
||||
|
||||
Stays under the cap so the tokenizer actually runs end to end —
|
||||
the goal is to prove the inner loop is O(N), not just that the
|
||||
cap rejects oversized input.
|
||||
"""
|
||||
# Pack the header right up to the cap with non-matching
|
||||
# auth-params, then put the real ``scope`` at the end.
|
||||
filler_parts = []
|
||||
size = len("Bearer ")
|
||||
i = 0
|
||||
while size < 3900:
|
||||
part = f'xscope{i}="ignore", '
|
||||
if size + len(part) > 3900:
|
||||
break
|
||||
filler_parts.append(part)
|
||||
size += len(part)
|
||||
i += 1
|
||||
header = "Bearer " + "".join(filler_parts) + 'scope="real"'
|
||||
assert len(header) <= 4096
|
||||
start = time.perf_counter()
|
||||
result = parse_www_authenticate_scope(header)
|
||||
elapsed = time.perf_counter() - start
|
||||
assert result == ("real",)
|
||||
assert elapsed < 0.05, f"4kb tokenize took {elapsed * 1000:.1f}ms"
|
||||
@@ -14,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import _seed_static_state
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
@@ -109,13 +110,15 @@ class TestFullLifecycleResourcesPrompts:
|
||||
|
||||
def test_rebuild_resources_produces_merged_state(self, mgr: MCPClientManager) -> None:
|
||||
"""_rebuild_resources merges per-server resources into a unified list."""
|
||||
mgr._per_server_resources["alpha"] = [
|
||||
_make_resource("file:///a.txt", "a", "alpha"),
|
||||
_make_resource("file:///b.txt", "b", "alpha"),
|
||||
]
|
||||
mgr._per_server_resources["beta"] = [
|
||||
_make_resource("file:///c.txt", "c", "beta"),
|
||||
]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
resources=[
|
||||
_make_resource("file:///a.txt", "a", "alpha"),
|
||||
_make_resource("file:///b.txt", "b", "alpha"),
|
||||
],
|
||||
)
|
||||
_seed_static_state(mgr, "beta", resources=[_make_resource("file:///c.txt", "c", "beta")])
|
||||
|
||||
mgr._rebuild_resources()
|
||||
|
||||
@@ -130,13 +133,19 @@ class TestFullLifecycleResourcesPrompts:
|
||||
|
||||
def test_rebuild_prompts_produces_merged_state(self, mgr: MCPClientManager) -> None:
|
||||
"""_rebuild_prompts merges per-server prompts into a unified list."""
|
||||
mgr._per_server_prompts["alpha"] = [
|
||||
_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello"),
|
||||
]
|
||||
mgr._per_server_prompts["beta"] = [
|
||||
_make_prompt("mcp__beta__summarize", "summarize", "beta", "Summarize text"),
|
||||
_make_prompt("mcp__beta__translate", "translate", "beta", "Translate text"),
|
||||
]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello")],
|
||||
)
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"beta",
|
||||
prompts=[
|
||||
_make_prompt("mcp__beta__summarize", "summarize", "beta", "Summarize text"),
|
||||
_make_prompt("mcp__beta__translate", "translate", "beta", "Translate text"),
|
||||
],
|
||||
)
|
||||
|
||||
mgr._rebuild_prompts()
|
||||
|
||||
@@ -164,10 +173,12 @@ class TestFullLifecycleResourcesPrompts:
|
||||
try:
|
||||
# Populate session and resource map
|
||||
session = _make_mock_session()
|
||||
mgr._sessions["alpha"] = session
|
||||
mgr._per_server_resources["alpha"] = [
|
||||
_make_resource("file:///readme.md", "readme", "alpha"),
|
||||
]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
session=session,
|
||||
resources=[_make_resource("file:///readme.md", "readme", "alpha")],
|
||||
)
|
||||
mgr._rebuild_resources()
|
||||
|
||||
result = mgr.read_resource_sync("file:///readme.md", timeout=5)
|
||||
@@ -194,18 +205,22 @@ class TestFullLifecycleResourcesPrompts:
|
||||
|
||||
try:
|
||||
session = _make_mock_session()
|
||||
mgr._sessions["alpha"] = session
|
||||
# Register a template resource (no concrete resources)
|
||||
mgr._per_server_resources["alpha"] = [
|
||||
{
|
||||
"uri": "db://tables/{table}/rows/{id}",
|
||||
"name": "row",
|
||||
"description": "Fetch a row",
|
||||
"mimeType": "application/json",
|
||||
"server": "alpha",
|
||||
"template": True,
|
||||
},
|
||||
]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
session=session,
|
||||
resources=[
|
||||
{
|
||||
"uri": "db://tables/{table}/rows/{id}",
|
||||
"name": "row",
|
||||
"description": "Fetch a row",
|
||||
"mimeType": "application/json",
|
||||
"server": "alpha",
|
||||
"template": True,
|
||||
},
|
||||
],
|
||||
)
|
||||
mgr._rebuild_resources()
|
||||
|
||||
# Template should not be in _resource_map
|
||||
@@ -230,10 +245,12 @@ class TestFullLifecycleResourcesPrompts:
|
||||
|
||||
try:
|
||||
session = _make_mock_session()
|
||||
mgr._sessions["alpha"] = session
|
||||
mgr._per_server_prompts["alpha"] = [
|
||||
_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello"),
|
||||
]
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
session=session,
|
||||
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello")],
|
||||
)
|
||||
mgr._rebuild_prompts()
|
||||
|
||||
messages = mgr.get_prompt_sync(
|
||||
@@ -314,40 +331,42 @@ class TestFullLifecycleResourcesPrompts:
|
||||
def test_shutdown_clears_all_state(self, mgr: MCPClientManager) -> None:
|
||||
"""shutdown() clears sessions, tools, resources, prompts, and listeners."""
|
||||
# Populate state
|
||||
mgr._sessions["alpha"] = MagicMock()
|
||||
mgr._per_server_tools["alpha"] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "mcp__alpha__search",
|
||||
"description": "Search",
|
||||
"parameters": {},
|
||||
_seed_static_state(
|
||||
mgr,
|
||||
"alpha",
|
||||
session=MagicMock(),
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "mcp__alpha__search",
|
||||
"description": "Search",
|
||||
"parameters": {},
|
||||
},
|
||||
}
|
||||
],
|
||||
resources=[
|
||||
_make_resource("file:///a.txt", "a", "alpha"),
|
||||
{
|
||||
"uri": "db://tables/{table}",
|
||||
"name": "table",
|
||||
"description": "",
|
||||
"mimeType": "",
|
||||
"server": "alpha",
|
||||
"template": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
],
|
||||
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha")],
|
||||
)
|
||||
mgr._rebuild_tools()
|
||||
mgr._per_server_resources["alpha"] = [
|
||||
_make_resource("file:///a.txt", "a", "alpha"),
|
||||
{
|
||||
"uri": "db://tables/{table}",
|
||||
"name": "table",
|
||||
"description": "",
|
||||
"mimeType": "",
|
||||
"server": "alpha",
|
||||
"template": True,
|
||||
},
|
||||
]
|
||||
mgr._rebuild_resources()
|
||||
mgr._per_server_prompts["alpha"] = [
|
||||
_make_prompt("mcp__alpha__greet", "greet", "alpha"),
|
||||
]
|
||||
mgr._rebuild_prompts()
|
||||
mgr._listeners.append(lambda: None)
|
||||
mgr._resource_listeners.append(lambda: None)
|
||||
mgr._prompt_listeners.append(lambda: None)
|
||||
|
||||
# Verify populated
|
||||
assert len(mgr._sessions) == 1
|
||||
assert len(mgr._static_servers) == 1
|
||||
assert len(mgr._tools) == 1
|
||||
assert len(mgr._resources) == 2 # 1 concrete + 1 template
|
||||
assert len(mgr._template_prefixes) == 1
|
||||
@@ -355,7 +374,7 @@ class TestFullLifecycleResourcesPrompts:
|
||||
|
||||
mgr.shutdown()
|
||||
|
||||
assert len(mgr._sessions) == 0
|
||||
assert len(mgr._static_servers) == 0
|
||||
assert len(mgr._tools) == 0
|
||||
assert len(mgr._tool_map) == 0
|
||||
assert len(mgr._resources) == 0
|
||||
@@ -376,19 +395,15 @@ class TestFullLifecycleResourcesPrompts:
|
||||
mgr.add_resource_listener(lambda: resource_fired.append(1))
|
||||
mgr.add_prompt_listener(lambda: prompt_fired.append(1))
|
||||
|
||||
mgr._per_server_tools["alpha"] = []
|
||||
_seed_static_state(mgr, "alpha", tools=[])
|
||||
mgr._rebuild_tools()
|
||||
assert len(tool_fired) == 1
|
||||
|
||||
mgr._per_server_resources["alpha"] = [
|
||||
_make_resource("file:///x.txt", "x", "alpha"),
|
||||
]
|
||||
_seed_static_state(mgr, "alpha", resources=[_make_resource("file:///x.txt", "x", "alpha")])
|
||||
mgr._rebuild_resources()
|
||||
assert len(resource_fired) == 1
|
||||
|
||||
mgr._per_server_prompts["alpha"] = [
|
||||
_make_prompt("mcp__alpha__p1", "p1", "alpha"),
|
||||
]
|
||||
_seed_static_state(mgr, "alpha", prompts=[_make_prompt("mcp__alpha__p1", "p1", "alpha")])
|
||||
mgr._rebuild_prompts()
|
||||
assert len(prompt_fired) == 1
|
||||
|
||||
|
||||
@@ -0,0 +1,780 @@
|
||||
"""Integration tests for the MCP OAuth ``/connections`` endpoints.
|
||||
|
||||
Covers the list and revoke handlers that surface user-owned MCP server
|
||||
consents to the settings UI:
|
||||
|
||||
* ``GET /v1/api/mcp/oauth/connections`` — non-secret projection only.
|
||||
* ``DELETE /v1/api/mcp/oauth/connections/{server_name}`` — best-effort
|
||||
upstream revoke (RFC 7009) followed by the authoritative local
|
||||
delete; cross-user attempts return 404 with the exact same body
|
||||
shape as a never-existed row to avoid leaking tenant existence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.routing import Mount, Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.auth import AuthResult
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.mcp_oauth import (
|
||||
handle_mcp_oauth_list_connections,
|
||||
handle_mcp_oauth_revoke_connection,
|
||||
)
|
||||
from turnstone.core.oidc import OIDCConfig
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures + helpers (mirror tests/test_mcp_oauth_handlers.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _InjectAuthMiddleware(BaseHTTPMiddleware):
|
||||
"""Stamp a fixed authenticated user on every request."""
|
||||
|
||||
def __init__(self, app: Any, user_id: str = "user-1") -> None:
|
||||
super().__init__(app)
|
||||
self._user_id = user_id
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id=self._user_id,
|
||||
scopes=frozenset({"write"}),
|
||||
token_source="config",
|
||||
permissions=frozenset({"read", "write"}),
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
class _NoAuthMiddleware(BaseHTTPMiddleware):
|
||||
"""Leave ``request.state.auth_result`` unset so handlers see anon."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
async def _list_handler(request: Request) -> Response:
|
||||
return await handle_mcp_oauth_list_connections(request)
|
||||
|
||||
|
||||
async def _revoke_handler(request: Request) -> Response:
|
||||
return await handle_mcp_oauth_revoke_connection(request)
|
||||
|
||||
|
||||
def _build_app(
|
||||
*,
|
||||
storage: SQLiteBackend,
|
||||
http_client: httpx.AsyncClient | MagicMock,
|
||||
token_store: MCPTokenStore | None,
|
||||
user_id: str = "user-1",
|
||||
mcp_client: Any = None,
|
||||
authenticated: bool = True,
|
||||
) -> Starlette:
|
||||
middleware: list[Middleware]
|
||||
if authenticated:
|
||||
middleware = [Middleware(_InjectAuthMiddleware, user_id=user_id)]
|
||||
else:
|
||||
middleware = [Middleware(_NoAuthMiddleware)]
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Mount(
|
||||
"/v1",
|
||||
routes=[
|
||||
Route("/api/mcp/oauth/connections", _list_handler),
|
||||
Route(
|
||||
"/api/mcp/oauth/connections/{server_name}",
|
||||
_revoke_handler,
|
||||
methods=["DELETE"],
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
middleware=middleware,
|
||||
)
|
||||
app.state.auth_storage = storage
|
||||
app.state.mcp_token_store = token_store
|
||||
app.state.mcp_oauth_http_client = http_client
|
||||
app.state.mcp_oauth_refresh_locks = {}
|
||||
app.state.mcp_oauth_dcr_locks = {}
|
||||
app.state.mcp_oauth_metadata_cache = {}
|
||||
app.state.mcp_oauth_last_cleanup_monotonic = 0.0
|
||||
app.state.oidc_config = OIDCConfig(enabled=False, redirect_base="https://testserver")
|
||||
if mcp_client is not None:
|
||||
app.state.mcp_client = mcp_client
|
||||
return app
|
||||
|
||||
|
||||
def _make_token_store(backend: SQLiteBackend) -> MCPTokenStore:
|
||||
return MCPTokenStore(backend, make_mcp_token_cipher(), node_id="test")
|
||||
|
||||
|
||||
def _seed_oauth_user_server(
|
||||
backend: SQLiteBackend,
|
||||
*,
|
||||
name: str = "srv-oauth",
|
||||
server_id: str = "srv-id-1",
|
||||
cached_issuer: str | None = "https://as.example.com",
|
||||
) -> str:
|
||||
backend.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="client-abc",
|
||||
oauth_scopes="openid profile",
|
||||
oauth_audience="https://mcp.example.com",
|
||||
oauth_authorization_server_url=None,
|
||||
)
|
||||
if cached_issuer is not None:
|
||||
backend.update_mcp_server(server_id, oauth_as_issuer_cached=cached_issuer)
|
||||
return server_id
|
||||
|
||||
|
||||
def _seed_user_token(
|
||||
token_store: MCPTokenStore,
|
||||
*,
|
||||
user_id: str = "user-1",
|
||||
server_name: str = "srv-oauth",
|
||||
refresh_token: str | None = "refresh-secret",
|
||||
) -> None:
|
||||
token_store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token="access-secret",
|
||||
refresh_token=refresh_token,
|
||||
expires_at="2099-12-31T00:00:00",
|
||||
scopes="openid profile",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
|
||||
def _good_as_metadata_doc(
|
||||
*, revocation_endpoint: str | None = "https://as.example.com/revoke"
|
||||
) -> dict[str, Any]:
|
||||
doc: dict[str, Any] = {
|
||||
"issuer": "https://as.example.com",
|
||||
"authorization_endpoint": "https://as.example.com/authorize",
|
||||
"token_endpoint": "https://as.example.com/token",
|
||||
"registration_endpoint": "https://as.example.com/register",
|
||||
"jwks_uri": "https://as.example.com/jwks",
|
||||
"code_challenge_methods_supported": ["S256"],
|
||||
"token_endpoint_auth_methods_supported": ["none", "client_secret_basic"],
|
||||
}
|
||||
if revocation_endpoint is not None:
|
||||
doc["revocation_endpoint"] = revocation_endpoint
|
||||
return doc
|
||||
|
||||
|
||||
def _mk_response(
|
||||
status_code: int = 200,
|
||||
json_body: Any = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> MagicMock:
|
||||
import json as _json
|
||||
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
resp.headers = headers or {}
|
||||
body_str = _json.dumps(json_body) if json_body is not None else ""
|
||||
resp.content = body_str.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_str
|
||||
return resp
|
||||
|
||||
|
||||
def _public_addr_patch():
|
||||
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
|
||||
|
||||
|
||||
def _drain_revoke_upstream_tasks(client: TestClient, timeout: float = 2.0) -> None:
|
||||
"""Block until all in-flight upstream-revoke tasks complete.
|
||||
|
||||
Phase 8 perf-1 made the RFC 7009 AS round-trip a fire-and-forget
|
||||
task so the user-visible 204 isn't gated on the AS. The tasks were
|
||||
scheduled on the TestClient's portal loop; we re-enter that loop
|
||||
via :attr:`TestClient.portal` to await them. Tests that assert
|
||||
against the upstream POST must call this helper before the
|
||||
assertion.
|
||||
"""
|
||||
from turnstone.core.mcp_oauth import _revoke_upstream_tasks
|
||||
|
||||
portal = getattr(client, "portal", None)
|
||||
if portal is None:
|
||||
return
|
||||
|
||||
async def _drain() -> None:
|
||||
pending = list(_revoke_upstream_tasks)
|
||||
if pending:
|
||||
async with asyncio.timeout(timeout):
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
portal.call(_drain)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
backend = SQLiteBackend(str(tmp_path / "test.db"))
|
||||
backend.create_user("user-1", "user1", "User One", "hash")
|
||||
backend.create_user("user-2", "user2", "User Two", "hash")
|
||||
return backend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def http_client_mock() -> MagicMock:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock()
|
||||
client.post = AsyncMock()
|
||||
return client
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /connections
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestListConnections:
|
||||
def test_list_connections_unauthenticated_401(
|
||||
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,
|
||||
authenticated=False,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
assert resp.status_code == 401
|
||||
assert resp.json() == {"error": "Authentication required"}
|
||||
|
||||
def test_list_connections_no_token_store_503(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=None)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_list_connections_empty_user_returns_empty_list(
|
||||
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)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"connections": []}
|
||||
|
||||
def test_list_connections_returns_users_consents(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage, name="srv-a", server_id="srv-id-a")
|
||||
_seed_oauth_user_server(storage, name="srv-b", server_id="srv-id-b")
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, server_name="srv-a")
|
||||
_seed_user_token(token_store, server_name="srv-b")
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert "connections" in body
|
||||
servers = sorted(row["server_name"] for row in body["connections"])
|
||||
assert servers == ["srv-a", "srv-b"]
|
||||
|
||||
def test_list_connections_isolates_by_user(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, user_id="user-1", server_name="srv-oauth")
|
||||
_seed_user_token(token_store, user_id="user-2", server_name="srv-oauth")
|
||||
|
||||
# User-1 sees only user-1's row.
|
||||
app = _build_app(
|
||||
storage=storage, http_client=http_client_mock, token_store=token_store, user_id="user-1"
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
rows = resp.json()["connections"]
|
||||
assert all(row["user_id"] == "user-1" for row in rows)
|
||||
assert len(rows) == 1
|
||||
|
||||
# User-2 sees only user-2's row.
|
||||
app2 = _build_app(
|
||||
storage=storage, http_client=http_client_mock, token_store=token_store, user_id="user-2"
|
||||
)
|
||||
client2 = TestClient(app2, raise_server_exceptions=False)
|
||||
resp2 = client2.get("/v1/api/mcp/oauth/connections")
|
||||
rows2 = resp2.json()["connections"]
|
||||
assert all(row["user_id"] == "user-2" for row in rows2)
|
||||
assert len(rows2) == 1
|
||||
|
||||
def test_list_connections_does_not_leak_secret_fields(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/v1/api/mcp/oauth/connections")
|
||||
rows = resp.json()["connections"]
|
||||
assert rows
|
||||
for row in rows:
|
||||
for forbidden in (
|
||||
"access_token",
|
||||
"refresh_token",
|
||||
"access_token_ct",
|
||||
"refresh_token_ct",
|
||||
):
|
||||
assert forbidden not in row, f"secret field {forbidden!r} leaked in {row!r}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DELETE /connections/{server_name}
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRevokeConnection:
|
||||
def test_revoke_connection_unauthenticated_401(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store)
|
||||
app = _build_app(
|
||||
storage=storage,
|
||||
http_client=http_client_mock,
|
||||
token_store=token_store,
|
||||
authenticated=False,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_revoke_connection_missing_row_404(
|
||||
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)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-nonexistent")
|
||||
assert resp.status_code == 404
|
||||
assert resp.json() == {"error": "No such connection"}
|
||||
|
||||
def test_revoke_connection_local_delete_succeeds_204(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
# No refresh token → upstream revoke is skipped entirely.
|
||||
_seed_user_token(token_store, refresh_token=None)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 204
|
||||
# Local row is gone.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
# Upstream not contacted.
|
||||
http_client_mock.post.assert_not_called()
|
||||
|
||||
def test_revoke_connection_with_revocation_endpoint_calls_upstream(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token="refresh-secret")
|
||||
|
||||
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
|
||||
http_client_mock.post.return_value = _mk_response(200)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
# ``with TestClient(...)`` keeps a persistent portal so the
|
||||
# fire-and-forget upstream-revoke task isn't cancelled when
|
||||
# the request handler returns. See ``_drain_revoke_upstream_tasks``.
|
||||
# The SSRF-validator's ``socket.getaddrinfo`` patch must wrap
|
||||
# the drain too — the discovery call now runs on the background
|
||||
# task and resolves the AS hostname after the request returns.
|
||||
with (
|
||||
TestClient(app, raise_server_exceptions=False) as client,
|
||||
_public_addr_patch(),
|
||||
):
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
|
||||
assert resp.status_code == 204
|
||||
# Local row is gone.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
# The upstream RFC 7009 POST is fire-and-forget post-Phase-8 perf-1
|
||||
# so the test must drain the in-flight task set before asserting.
|
||||
_drain_revoke_upstream_tasks(client)
|
||||
# Upstream POSTed to revocation_endpoint with refresh-token grant.
|
||||
assert http_client_mock.post.await_count == 1
|
||||
call = http_client_mock.post.await_args
|
||||
assert call.args[0] == "https://as.example.com/revoke"
|
||||
data = call.kwargs.get("data") or {}
|
||||
assert data.get("token") == "refresh-secret"
|
||||
assert data.get("token_type_hint") == "refresh_token"
|
||||
assert data.get("client_id") == "client-abc"
|
||||
|
||||
def test_revoke_connection_without_revocation_endpoint_skips_upstream(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token="refresh-secret")
|
||||
|
||||
http_client_mock.get.return_value = _mk_response(
|
||||
200, _good_as_metadata_doc(revocation_endpoint=None)
|
||||
)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
# ``with TestClient(...)`` keeps the portal alive for the
|
||||
# background task drain.
|
||||
with (
|
||||
TestClient(app, raise_server_exceptions=False) as client,
|
||||
_public_addr_patch(),
|
||||
):
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
|
||||
assert resp.status_code == 204
|
||||
# Local row gone, upstream POST never made.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
# Drain the fire-and-forget discovery task before asserting on
|
||||
# the AS POST — the task runs ``discover_authorization_server``
|
||||
# but does NOT proceed to POST because revocation_endpoint is
|
||||
# absent.
|
||||
_drain_revoke_upstream_tasks(client)
|
||||
http_client_mock.post.assert_not_called()
|
||||
|
||||
def test_revoke_connection_upstream_failure_still_204(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token="refresh-secret")
|
||||
|
||||
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
|
||||
# AS returns 500 — local delete must still succeed.
|
||||
http_client_mock.post.return_value = _mk_response(500)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
with _public_addr_patch():
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
|
||||
assert resp.status_code == 204
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
|
||||
def test_revoke_connection_audit_event_emitted_with_user_revoked_reason(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token=None)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 204
|
||||
|
||||
# Audit row was written via the storage API (tests don't poke at
|
||||
# the SQLite schema directly — the table name is an internal
|
||||
# detail).
|
||||
events = storage.list_audit_events(action="mcp_server.oauth.token_revoked")
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["user_id"] == "user-1"
|
||||
# resource_id is the immutable server_id PK, not the name.
|
||||
assert ev["resource_id"] == "srv-id-1"
|
||||
import json as _json
|
||||
|
||||
detail = _json.loads(ev["detail"]) if isinstance(ev["detail"], str) else ev["detail"]
|
||||
assert detail["reason"] == "user_revoked"
|
||||
assert detail["upstream_revoke_outcome"] == "no_refresh_token"
|
||||
assert detail["server_name"] == "srv-oauth"
|
||||
|
||||
def test_revoke_connection_cross_user_attempt_404(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
# Owned by user-2, not user-1.
|
||||
_seed_user_token(token_store, user_id="user-2", server_name="srv-oauth")
|
||||
|
||||
app = _build_app(
|
||||
storage=storage, http_client=http_client_mock, token_store=token_store, user_id="user-1"
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
# Cross-user attempt MUST surface as a generic 404, byte-identical
|
||||
# body to the never-existed case (no tenant existence leak).
|
||||
assert resp.status_code == 404
|
||||
assert resp.json() == {"error": "No such connection"}
|
||||
# Drain pending tasks defensively, then confirm the upstream
|
||||
# endpoint was NEVER contacted on the 404-cross-user path. A
|
||||
# bug that scheduled the AS round-trip before the cross-user
|
||||
# check would leak existence via the AS-side 200/4xx response.
|
||||
_drain_revoke_upstream_tasks(client)
|
||||
http_client_mock.post.assert_not_called()
|
||||
# User-2's row is untouched.
|
||||
assert token_store.get_user_token("user-2", "srv-oauth") is not None
|
||||
|
||||
def test_revoke_connection_evicts_pool_session(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token=None)
|
||||
|
||||
mcp_client_mock = MagicMock()
|
||||
# ``evict_user_session`` is the public sync surface on
|
||||
# MCPClientManager; mirror its signature here so the handler's
|
||||
# ``hasattr`` gate triggers.
|
||||
mcp_client_mock.evict_user_session = MagicMock(return_value=None)
|
||||
|
||||
app = _build_app(
|
||||
storage=storage,
|
||||
http_client=http_client_mock,
|
||||
token_store=token_store,
|
||||
mcp_client=mcp_client_mock,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 204
|
||||
|
||||
mcp_client_mock.evict_user_session.assert_called_once_with("user-1", "srv-oauth")
|
||||
|
||||
def test_revoke_connection_pool_eviction_failure_does_not_block_204(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token=None)
|
||||
|
||||
mcp_client_mock = MagicMock()
|
||||
mcp_client_mock.evict_user_session = MagicMock(side_effect=RuntimeError("loop closed"))
|
||||
|
||||
app = _build_app(
|
||||
storage=storage,
|
||||
http_client=http_client_mock,
|
||||
token_store=token_store,
|
||||
mcp_client=mcp_client_mock,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 204
|
||||
# Local delete still happened.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
|
||||
def test_revoke_connection_204_not_gated_on_slow_upstream(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
"""The user-visible 204 must return promptly even when the
|
||||
upstream AS round-trip is slow / hanging. Pre-perf-1 the
|
||||
handler awaited ``revoke_token_at_as`` synchronously, so a
|
||||
stuck AS could block the user's revoke confirmation. The
|
||||
fire-and-forget refactor moves the call onto a background task
|
||||
so the 204 returns in well under 1s regardless of AS latency.
|
||||
Bound is conservative for CI runner jitter.
|
||||
"""
|
||||
import time
|
||||
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token="refresh-secret")
|
||||
|
||||
http_client_mock.get.return_value = _mk_response(200, _good_as_metadata_doc())
|
||||
|
||||
async def _slow_post(*_args: Any, **_kwargs: Any) -> Any:
|
||||
# Simulate a slow / unreachable AS — must NOT gate the
|
||||
# user-visible 204 on this round-trip.
|
||||
await asyncio.sleep(5.0)
|
||||
return _mk_response(200)
|
||||
|
||||
http_client_mock.post = AsyncMock(side_effect=_slow_post)
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
with _public_addr_patch():
|
||||
start = time.monotonic()
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert resp.status_code == 204
|
||||
# 1s ceiling — the 204 must return on the local-delete path
|
||||
# without waiting on the AS POST (which sleeps 5s above). Bound
|
||||
# is intentionally generous for CI runner jitter; the actual
|
||||
# path is on the order of milliseconds.
|
||||
assert elapsed < 1.0, (
|
||||
f"204 returned in {elapsed:.3f}s — should be <1s; the "
|
||||
"fire-and-forget upstream revoke isn't decoupled from the "
|
||||
"response."
|
||||
)
|
||||
# The local row IS gone — the authoritative delete ran before
|
||||
# the 204 returned, even though the AS round-trip is still
|
||||
# in flight.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
# Cancel any in-flight tasks so the test client can exit cleanly.
|
||||
from turnstone.core.mcp_oauth import _revoke_upstream_tasks
|
||||
|
||||
portal = getattr(client, "portal", None)
|
||||
if portal is not None:
|
||||
for task in list(_revoke_upstream_tasks):
|
||||
portal.call(task.cancel)
|
||||
|
||||
def test_revoke_connection_sheds_upstream_when_task_set_full(
|
||||
self, storage: SQLiteBackend, http_client_mock: MagicMock
|
||||
) -> None:
|
||||
"""Round-2 q-2 regression: the soft cap on ``_revoke_upstream_tasks``
|
||||
is the only protection against unbounded background-task pile-up
|
||||
under a coordinated mass-revoke. When the set is full, the local
|
||||
delete still runs but no upstream task is scheduled; the audit
|
||||
detail records ``upstream_revoke_outcome="shed_by_cap"`` and
|
||||
the AS endpoint is never contacted.
|
||||
"""
|
||||
from turnstone.core.mcp_oauth import (
|
||||
_REVOKE_UPSTREAM_TASKS_MAX,
|
||||
_revoke_upstream_tasks,
|
||||
)
|
||||
|
||||
_seed_oauth_user_server(storage)
|
||||
token_store = _make_token_store(storage)
|
||||
_seed_user_token(token_store, refresh_token="refresh-secret")
|
||||
|
||||
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
|
||||
sentinel_event_holder: dict[str, asyncio.Event] = {}
|
||||
|
||||
# Use ``with TestClient(...)`` so the portal stays alive — we
|
||||
# need to schedule sentinel tasks on the portal's loop and the
|
||||
# tasks must outlive the request to actually fill the set.
|
||||
with (
|
||||
TestClient(app, raise_server_exceptions=False) as client,
|
||||
_public_addr_patch(),
|
||||
):
|
||||
portal = client.portal
|
||||
assert portal is not None
|
||||
|
||||
async def _create_sentinel_event() -> asyncio.Event:
|
||||
event = asyncio.Event()
|
||||
sentinel_event_holder["event"] = event
|
||||
return event
|
||||
|
||||
sentinel_event = portal.call(_create_sentinel_event)
|
||||
|
||||
async def _wait_on_event() -> None:
|
||||
await sentinel_event.wait()
|
||||
|
||||
async def _fill_task_set() -> list[asyncio.Task[None]]:
|
||||
tasks: list[asyncio.Task[None]] = []
|
||||
for _ in range(_REVOKE_UPSTREAM_TASKS_MAX):
|
||||
t = asyncio.create_task(_wait_on_event())
|
||||
_revoke_upstream_tasks.add(t)
|
||||
tasks.append(t)
|
||||
return tasks
|
||||
|
||||
sentinels = portal.call(_fill_task_set)
|
||||
assert len(_revoke_upstream_tasks) >= _REVOKE_UPSTREAM_TASKS_MAX
|
||||
|
||||
try:
|
||||
resp = client.delete("/v1/api/mcp/oauth/connections/srv-oauth")
|
||||
assert resp.status_code == 204
|
||||
# Local row is still gone — authoritative delete ran.
|
||||
assert token_store.get_user_token("user-1", "srv-oauth") is None
|
||||
# AS endpoint MUST NOT have been contacted.
|
||||
http_client_mock.post.assert_not_called()
|
||||
# Audit detail records the categorical shed outcome.
|
||||
events = storage.list_audit_events(action="mcp_server.oauth.token_revoked")
|
||||
assert len(events) == 1
|
||||
detail = events[0]["detail"]
|
||||
if isinstance(detail, str):
|
||||
import json as _json
|
||||
|
||||
detail = _json.loads(detail)
|
||||
assert detail["upstream_revoke_outcome"] == "shed_by_cap"
|
||||
finally:
|
||||
# Release sentinels so the portal can shut down cleanly.
|
||||
async def _release() -> None:
|
||||
sentinel_event.set()
|
||||
for t in sentinels:
|
||||
t.cancel()
|
||||
await asyncio.gather(*sentinels, return_exceptions=True)
|
||||
|
||||
portal.call(_release)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# evict_user_session helper sanity checks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEvictUserSession:
|
||||
def test_evict_user_session_no_loop_is_silent_noop(self) -> None:
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
|
||||
mgr = MCPClientManager.__new__(MCPClientManager)
|
||||
mgr._loop = None # type: ignore[attr-defined]
|
||||
# Must not raise.
|
||||
mgr.evict_user_session("user-1", "srv-oauth")
|
||||
|
||||
def test_evict_user_session_dispatches_to_loop(self) -> None:
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
|
||||
mgr = MCPClientManager.__new__(MCPClientManager)
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
mgr._loop = loop # type: ignore[attr-defined]
|
||||
mgr._user_pool_entries = {} # type: ignore[attr-defined]
|
||||
mgr._last_pool_notification_refresh = {} # type: ignore[attr-defined]
|
||||
evicted: list[tuple[str, str]] = []
|
||||
|
||||
def _fake_evict(key: tuple[str, str]) -> None:
|
||||
evicted.append(key)
|
||||
|
||||
mgr._evict_session = _fake_evict # type: ignore[method-assign]
|
||||
|
||||
# Run the dispatch on a separate thread so the loop can drain.
|
||||
import threading
|
||||
|
||||
done = threading.Event()
|
||||
|
||||
def _run_loop() -> None:
|
||||
loop.call_later(0.05, loop.stop)
|
||||
loop.run_forever()
|
||||
done.set()
|
||||
|
||||
t = threading.Thread(target=_run_loop, daemon=True)
|
||||
t.start()
|
||||
mgr.evict_user_session("user-1", "srv-oauth")
|
||||
done.wait(timeout=1.0)
|
||||
|
||||
assert evicted == [("user-1", "srv-oauth")]
|
||||
finally:
|
||||
if not loop.is_closed():
|
||||
loop.close()
|
||||
@@ -0,0 +1,626 @@
|
||||
"""Discovery tests for the per-(user, server) MCP OAuth flow.
|
||||
|
||||
Covers PRM (RFC 9728) and AS metadata (RFC 8414) discovery, including:
|
||||
- override URL takes precedence
|
||||
- PRM happy path: server URL -> .well-known/oauth-protected-resource
|
||||
-> ``authorization_servers[0]``
|
||||
- PRM 401 + ``WWW-Authenticate: Bearer resource_metadata="..."`` follows
|
||||
the URL.
|
||||
- AS metadata without S256 -> :class:`MCPOAuthDiscoveryError`.
|
||||
- SSRF rejection on AS issuer URL.
|
||||
- In-memory cache hit/miss + persistent cache write to
|
||||
``mcp_servers.oauth_as_issuer_cached``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from turnstone.core.mcp_oauth import (
|
||||
ASMetadata,
|
||||
MCPOAuthDiscoveryError,
|
||||
_parse_prm_url_from_www_authenticate,
|
||||
discover_authorization_server,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _mk_response(
|
||||
status_code: int = 200,
|
||||
json_body: Any = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> MagicMock:
|
||||
"""Build a MagicMock that quacks like ``httpx.Response``."""
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
resp.headers = headers or {}
|
||||
resp.content = (str(json_body) if json_body is not None else "").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 = str(json_body) if json_body is not None else ""
|
||||
return resp
|
||||
|
||||
|
||||
def _good_as_metadata_doc() -> dict[str, Any]:
|
||||
return {
|
||||
"issuer": "https://as.example.com",
|
||||
"authorization_endpoint": "https://as.example.com/authorize",
|
||||
"token_endpoint": "https://as.example.com/token",
|
||||
"jwks_uri": "https://as.example.com/jwks",
|
||||
"code_challenge_methods_supported": ["S256"],
|
||||
"token_endpoint_auth_methods_supported": ["none", "client_secret_basic"],
|
||||
"registration_endpoint": "https://as.example.com/register",
|
||||
}
|
||||
|
||||
|
||||
def _public_addr_patch():
|
||||
return patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("93.184.216.34", 0))])
|
||||
|
||||
|
||||
def _mk_storage_mock(server_id: str = "srv-id") -> MagicMock:
|
||||
storage = MagicMock()
|
||||
storage.update_mcp_server.return_value = True
|
||||
return storage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PRM parser
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParsePRMUrl:
|
||||
def test_extracts_resource_metadata_url(self) -> None:
|
||||
header = (
|
||||
'Bearer error="invalid_token", '
|
||||
'resource_metadata="https://srv.example.com/.well-known/oauth-protected-resource"'
|
||||
)
|
||||
url = _parse_prm_url_from_www_authenticate(header)
|
||||
assert url == "https://srv.example.com/.well-known/oauth-protected-resource"
|
||||
|
||||
def test_returns_none_when_absent(self) -> None:
|
||||
assert _parse_prm_url_from_www_authenticate('Bearer realm="x"') is None
|
||||
|
||||
def test_handles_empty_header(self) -> None:
|
||||
assert _parse_prm_url_from_www_authenticate("") is None
|
||||
|
||||
def test_handles_escaped_quote_in_value(self) -> None:
|
||||
"""RFC 7230 quoted-string allows ``\\"`` — naive ``[^"]+`` truncates.
|
||||
|
||||
A malicious or buggy resource server could send an embedded
|
||||
escaped quote; the parser must yield the unescaped value, not
|
||||
the prefix up to the escaped quote.
|
||||
"""
|
||||
header = 'Bearer resource_metadata="https://srv.example.com/with\\"quote"'
|
||||
url = _parse_prm_url_from_www_authenticate(header)
|
||||
assert url == 'https://srv.example.com/with"quote'
|
||||
|
||||
def test_handles_escaped_backslash(self) -> None:
|
||||
header = 'Bearer resource_metadata="https://srv.example.com/back\\\\slash"'
|
||||
url = _parse_prm_url_from_www_authenticate(header)
|
||||
assert url == "https://srv.example.com/back\\slash"
|
||||
|
||||
def test_unterminated_quoted_string_returns_none(self) -> None:
|
||||
# Closing quote missing — naive regex would still match, but
|
||||
# the proper parser should reject malformed input.
|
||||
header = 'Bearer resource_metadata="https://srv.example.com/no-close'
|
||||
assert _parse_prm_url_from_www_authenticate(header) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# discover_authorization_server happy paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDiscoveryOverride:
|
||||
def test_override_url_skips_prm(self) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert isinstance(meta, ASMetadata)
|
||||
assert meta.token_endpoint == "https://as.example.com/token"
|
||||
# Only the AS metadata URL was hit, not PRM.
|
||||
called_urls = [c.args[0] for c in client.get.call_args_list]
|
||||
assert all("oauth-authorization-server" in u for u in called_urls)
|
||||
|
||||
|
||||
class TestDiscoveryPRM:
|
||||
def test_prm_happy_path(self) -> None:
|
||||
async def _get(url, *args, **kwargs):
|
||||
if url.endswith("/oauth-protected-resource"):
|
||||
return _mk_response(
|
||||
200,
|
||||
{
|
||||
"resource": "https://mcp.example.com",
|
||||
"authorization_servers": ["https://as.example.com"],
|
||||
},
|
||||
)
|
||||
if url.endswith("/oauth-authorization-server"):
|
||||
return _mk_response(200, _good_as_metadata_doc())
|
||||
raise AssertionError(f"unexpected URL: {url}")
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(side_effect=_get)
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert meta.issuer == "https://as.example.com"
|
||||
|
||||
def test_prm_401_follows_www_authenticate(self) -> None:
|
||||
async def _get(url, *args, **kwargs):
|
||||
if url == "https://mcp.example.com/.well-known/oauth-protected-resource":
|
||||
return _mk_response(
|
||||
401,
|
||||
headers={
|
||||
"www-authenticate": (
|
||||
'Bearer error="invalid_token", '
|
||||
"resource_metadata="
|
||||
'"https://meta.example.com/prm"'
|
||||
)
|
||||
},
|
||||
json_body=None,
|
||||
)
|
||||
if url == "https://meta.example.com/prm":
|
||||
return _mk_response(
|
||||
200,
|
||||
{
|
||||
"authorization_servers": ["https://as.example.com"],
|
||||
},
|
||||
)
|
||||
if url.endswith("/oauth-authorization-server"):
|
||||
return _mk_response(200, _good_as_metadata_doc())
|
||||
raise AssertionError(f"unexpected URL: {url}")
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(side_effect=_get)
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert meta.token_endpoint == "https://as.example.com/token"
|
||||
|
||||
def test_prm_401_without_resource_metadata_raises(self) -> None:
|
||||
async def _get(url, *args, **kwargs):
|
||||
return _mk_response(401, headers={"www-authenticate": "Basic realm=x"})
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(side_effect=_get)
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError, match="resource_metadata"):
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AS metadata validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestASMetadataValidation:
|
||||
def test_no_s256_raises(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
doc["code_challenge_methods_supported"] = ["plain"]
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError, match="S256"):
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_missing_endpoints_raises(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
del doc["token_endpoint"]
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError, match="missing required"):
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_third_party_endpoint_rejected(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
doc["token_endpoint"] = "https://attacker.example.com/token"
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError, match="token_endpoint"):
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_ssrf_on_override_rejected(self) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock()
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
# Resolve to private 10.x — SSRF guard fires before any HTTP call.
|
||||
with patch(
|
||||
"socket.getaddrinfo",
|
||||
return_value=[(2, 1, 6, "", ("10.0.0.1", 0))],
|
||||
):
|
||||
await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://internal.corp.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError):
|
||||
asyncio.run(_run())
|
||||
client.get.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Caching
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMetadataCache:
|
||||
def test_cache_miss_then_hit(self) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
|
||||
storage = _mk_storage_mock()
|
||||
cache: dict[str, tuple[ASMetadata, float]] = {}
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
first = await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
metadata_cache=cache,
|
||||
)
|
||||
second = await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer="https://as.example.com",
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
metadata_cache=cache,
|
||||
)
|
||||
return first, second
|
||||
|
||||
first, second = asyncio.run(_run())
|
||||
assert first.token_endpoint == second.token_endpoint
|
||||
# First call hit AS metadata; second call hit the cache.
|
||||
assert client.get.call_count == 1
|
||||
|
||||
def test_cache_expiry_refetches(self) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
|
||||
storage = _mk_storage_mock()
|
||||
# Pre-populate cache with a very stale entry.
|
||||
stale_meta = ASMetadata(
|
||||
issuer="https://as.example.com",
|
||||
authorization_endpoint="https://as.example.com/authorize",
|
||||
token_endpoint="https://as.example.com/token",
|
||||
registration_endpoint=None,
|
||||
revocation_endpoint=None,
|
||||
jwks_uri=None,
|
||||
code_challenge_methods_supported=("S256",),
|
||||
token_endpoint_auth_methods_supported=(),
|
||||
)
|
||||
cache = {"https://as.example.com": (stale_meta, time.monotonic() - 10**6)}
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
metadata_cache=cache,
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
# Stale entry was bypassed -> we hit the network.
|
||||
assert client.get.call_count == 1
|
||||
assert meta.token_endpoint == "https://as.example.com/token"
|
||||
|
||||
def test_persistent_cache_write_on_first_resolution(self) -> None:
|
||||
async def _get(url, *args, **kwargs):
|
||||
if url.endswith("/oauth-protected-resource"):
|
||||
return _mk_response(200, {"authorization_servers": ["https://as.example.com"]})
|
||||
return _mk_response(200, _good_as_metadata_doc())
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(side_effect=_get)
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
asyncio.run(_run())
|
||||
# update_mcp_server was called once with the cached issuer.
|
||||
storage.update_mcp_server.assert_called_once_with(
|
||||
"srv-id", oauth_as_issuer_cached="https://as.example.com"
|
||||
)
|
||||
|
||||
def test_persistent_cache_skip_when_already_cached(self) -> None:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, _good_as_metadata_doc()))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run():
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer="https://as.example.com",
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
asyncio.run(_run())
|
||||
storage.update_mcp_server.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sec-3 — cached_issuer re-validated on read
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCachedIssuerSSRFRevalidation:
|
||||
"""A cached issuer URL must still pass SSRF validation on every read.
|
||||
|
||||
Defense-in-depth: an admin who points ``oauth_as_issuer_cached`` at a
|
||||
private address (or a hostname that has rebound to one) should not
|
||||
bypass the guard just because the value was already in the row.
|
||||
"""
|
||||
|
||||
def test_cached_issuer_rejected_clears_row_and_falls_through_to_prm(self) -> None:
|
||||
async def _get(url: str, *args: Any, **kwargs: Any) -> MagicMock:
|
||||
if url.endswith("/oauth-protected-resource"):
|
||||
return _mk_response(200, {"authorization_servers": ["https://as.example.com"]})
|
||||
if url.endswith("/oauth-authorization-server"):
|
||||
return _mk_response(200, _good_as_metadata_doc())
|
||||
raise AssertionError(f"unexpected URL: {url}")
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(side_effect=_get)
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
# cached_issuer points at a private host. SSRF guard fires on
|
||||
# the cached value first, the row is cleared, and PRM
|
||||
# discovery runs as a fallback.
|
||||
async def _run() -> Any:
|
||||
with patch(
|
||||
"socket.getaddrinfo",
|
||||
# Private resolution for "internal.corp", public for everything else.
|
||||
side_effect=lambda host, *a, **kw: [
|
||||
(2, 1, 6, "", ("10.0.0.1" if "internal" in host else "93.184.216.34", 0))
|
||||
],
|
||||
):
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url=None,
|
||||
cached_issuer="https://internal.corp.example.com",
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert meta.token_endpoint == "https://as.example.com/token"
|
||||
# The bad cached_issuer was cleared from the row.
|
||||
clear_calls = [
|
||||
c
|
||||
for c in storage.update_mcp_server.call_args_list
|
||||
if c.kwargs.get("oauth_as_issuer_cached") is None
|
||||
]
|
||||
assert clear_calls, "cached_issuer should have been cleared"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# revocation_endpoint parsing (RFC 8414)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestASMetadataRevocationEndpoint:
|
||||
def test_as_metadata_parses_revocation_endpoint(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
doc["revocation_endpoint"] = "https://as.example.com/revoke"
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run() -> ASMetadata:
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert meta.revocation_endpoint == "https://as.example.com/revoke"
|
||||
|
||||
def test_as_metadata_revocation_endpoint_absent(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
doc.pop("revocation_endpoint", None)
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run() -> ASMetadata:
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
meta = asyncio.run(_run())
|
||||
assert meta.revocation_endpoint is None
|
||||
|
||||
def test_as_metadata_revocation_endpoint_rejected_when_cross_origin(self) -> None:
|
||||
doc = _good_as_metadata_doc()
|
||||
doc["revocation_endpoint"] = "https://attacker.example.com/revoke"
|
||||
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = AsyncMock(return_value=_mk_response(200, doc))
|
||||
storage = _mk_storage_mock()
|
||||
|
||||
async def _run() -> ASMetadata:
|
||||
with _public_addr_patch():
|
||||
return await discover_authorization_server(
|
||||
server_name="srv-x",
|
||||
server_url="https://mcp.example.com/sse",
|
||||
override_url="https://as.example.com",
|
||||
cached_issuer=None,
|
||||
http_client=client,
|
||||
storage=storage,
|
||||
server_id="srv-id",
|
||||
trusted_hosts=frozenset(),
|
||||
)
|
||||
|
||||
with pytest.raises(MCPOAuthDiscoveryError, match="revocation_endpoint"):
|
||||
asyncio.run(_run())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,133 @@
|
||||
"""Smoke tests for the new OAuth-MCP storage tables.
|
||||
|
||||
Phase 2 only adds the schema — token CRUD lands in Phase 3 and pending-
|
||||
state CRUD in Phase 4. These tests verify the tables exist after
|
||||
``init_storage`` and accept the documented row shape via raw SQL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from turnstone.core.storage._schema import mcp_oauth_pending, mcp_user_tokens
|
||||
|
||||
|
||||
class TestMcpUserTokensTable:
|
||||
def test_table_exists_and_accepts_row(self, backend) -> None:
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.insert(mcp_user_tokens),
|
||||
{
|
||||
"user_id": "u1",
|
||||
"server_name": "srv-a",
|
||||
"access_token_ct": b"\x00ciphertext-a",
|
||||
"refresh_token_ct": b"\x00ciphertext-r",
|
||||
"expires_at": "2026-05-04T12:00:00",
|
||||
"scopes": "openid profile",
|
||||
"as_issuer": "https://auth.example.com",
|
||||
"audience": "https://mcp.example.com",
|
||||
"created": "2026-05-04T11:00:00",
|
||||
"last_refreshed": None,
|
||||
},
|
||||
)
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
sa.select(mcp_user_tokens).where(
|
||||
(mcp_user_tokens.c.user_id == "u1") & (mcp_user_tokens.c.server_name == "srv-a")
|
||||
)
|
||||
).one()
|
||||
assert row.access_token_ct == b"\x00ciphertext-a"
|
||||
assert row.refresh_token_ct == b"\x00ciphertext-r"
|
||||
assert row.scopes == "openid profile"
|
||||
assert row.audience == "https://mcp.example.com"
|
||||
|
||||
def test_composite_pk_distinguishes_user_server(self, backend) -> None:
|
||||
"""Same user, different server => two rows; same (user, server) => conflict."""
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.insert(mcp_user_tokens),
|
||||
[
|
||||
{
|
||||
"user_id": "u1",
|
||||
"server_name": "srv-a",
|
||||
"access_token_ct": b"a",
|
||||
"refresh_token_ct": None,
|
||||
"expires_at": None,
|
||||
"scopes": None,
|
||||
"as_issuer": "https://auth.example.com",
|
||||
"audience": "https://a.example.com",
|
||||
"created": "2026-05-04T11:00:00",
|
||||
"last_refreshed": None,
|
||||
},
|
||||
{
|
||||
"user_id": "u1",
|
||||
"server_name": "srv-b",
|
||||
"access_token_ct": b"b",
|
||||
"refresh_token_ct": None,
|
||||
"expires_at": None,
|
||||
"scopes": None,
|
||||
"as_issuer": "https://auth.example.com",
|
||||
"audience": "https://b.example.com",
|
||||
"created": "2026-05-04T11:00:00",
|
||||
"last_refreshed": None,
|
||||
},
|
||||
],
|
||||
)
|
||||
conn.commit()
|
||||
count = conn.execute(sa.select(sa.func.count()).select_from(mcp_user_tokens)).scalar()
|
||||
assert count == 2
|
||||
|
||||
|
||||
class TestMcpOauthPendingTable:
|
||||
def test_table_exists_and_accepts_row(self, backend) -> None:
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.insert(mcp_oauth_pending),
|
||||
{
|
||||
"state": "rand-state-xyz",
|
||||
"user_id": "u1",
|
||||
"server_name": "srv-a",
|
||||
"code_verifier": "verifier-blob",
|
||||
"return_url": "/admin/mcp-servers",
|
||||
"created_at": "2026-05-04T11:00:00",
|
||||
},
|
||||
)
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
sa.select(mcp_oauth_pending).where(mcp_oauth_pending.c.state == "rand-state-xyz")
|
||||
).one()
|
||||
assert row.user_id == "u1"
|
||||
assert row.server_name == "srv-a"
|
||||
assert row.return_url == "/admin/mcp-servers"
|
||||
|
||||
def test_state_pk_unique(self, backend) -> None:
|
||||
"""A second insert with the same state value raises IntegrityError."""
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.insert(mcp_oauth_pending),
|
||||
{
|
||||
"state": "dup-state",
|
||||
"user_id": "u1",
|
||||
"server_name": "srv-a",
|
||||
"code_verifier": "v",
|
||||
"return_url": "/x",
|
||||
"created_at": "2026-05-04T11:00:00",
|
||||
},
|
||||
)
|
||||
conn.commit()
|
||||
import pytest
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
with pytest.raises(IntegrityError), backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.insert(mcp_oauth_pending),
|
||||
{
|
||||
"state": "dup-state",
|
||||
"user_id": "u2",
|
||||
"server_name": "srv-b",
|
||||
"code_verifier": "v",
|
||||
"return_url": "/y",
|
||||
"created_at": "2026-05-04T11:01:00",
|
||||
},
|
||||
)
|
||||
conn.commit()
|
||||
@@ -0,0 +1,53 @@
|
||||
"""PKCE pair-generation tests for the MCP OAuth flow.
|
||||
|
||||
Verifies the contract documented in RFC 7636 §4.1 and §4.2:
|
||||
|
||||
- ``code_verifier`` is a high-entropy 43..128 character urlsafe-base64 string.
|
||||
- ``code_challenge`` is the BASE64URL-NO-PADDING encoding of
|
||||
``SHA256(verifier)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import string
|
||||
|
||||
from turnstone.core.mcp_oauth import generate_pkce_pair
|
||||
|
||||
_URLSAFE_CHARS = set(string.ascii_letters + string.digits + "-_")
|
||||
|
||||
|
||||
class TestGeneratePkcePair:
|
||||
def test_returns_tuple_of_strings(self) -> None:
|
||||
verifier, challenge = generate_pkce_pair()
|
||||
assert isinstance(verifier, str)
|
||||
assert isinstance(challenge, str)
|
||||
|
||||
def test_verifier_length_in_rfc_range(self) -> None:
|
||||
for _ in range(20):
|
||||
verifier, _ = generate_pkce_pair()
|
||||
assert 43 <= len(verifier) <= 128
|
||||
|
||||
def test_verifier_is_urlsafe(self) -> None:
|
||||
for _ in range(20):
|
||||
verifier, _ = generate_pkce_pair()
|
||||
assert all(ch in _URLSAFE_CHARS for ch in verifier)
|
||||
|
||||
def test_challenge_matches_sha256_of_verifier(self) -> None:
|
||||
for _ in range(20):
|
||||
verifier, challenge = generate_pkce_pair()
|
||||
digest = hashlib.sha256(verifier.encode("ascii")).digest()
|
||||
expected = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
|
||||
assert challenge == expected
|
||||
|
||||
def test_challenge_has_no_padding(self) -> None:
|
||||
for _ in range(20):
|
||||
_, challenge = generate_pkce_pair()
|
||||
assert "=" not in challenge
|
||||
|
||||
def test_pairs_are_unique(self) -> None:
|
||||
pairs = {generate_pkce_pair() for _ in range(50)}
|
||||
# 50 random draws shouldn't collide; if they do we have a much
|
||||
# bigger problem than this assertion.
|
||||
assert len(pairs) == 50
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,399 @@
|
||||
"""Tests for :func:`turnstone.core.mcp_oauth.revoke_token_at_as`.
|
||||
|
||||
The helper is best-effort RFC 7009 token revocation. It must:
|
||||
- skip cleanly when the AS metadata doesn't advertise a revocation endpoint
|
||||
- POST the form body when one is present (with optional client_secret)
|
||||
- never raise on non-2xx, network errors, or timeouts — caller doesn't
|
||||
want try/except in cleanup paths
|
||||
- never use ``exc_info=True`` — chained ``__context__`` may carry an
|
||||
``httpx.Request`` whose ``Authorization`` header holds a bearer; the
|
||||
bearer-leak invariant requires structured fields with type names only
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
||||
from turnstone.core.mcp_oauth import (
|
||||
ASMetadata,
|
||||
MCPOAuthDiscoveryError,
|
||||
_attempt_upstream_revoke,
|
||||
revoke_token_at_as,
|
||||
)
|
||||
|
||||
|
||||
def _make_as_metadata(
|
||||
*,
|
||||
revocation_endpoint: str | None = "https://as.example.com/revoke",
|
||||
) -> ASMetadata:
|
||||
return ASMetadata(
|
||||
issuer="https://as.example.com",
|
||||
authorization_endpoint="https://as.example.com/authorize",
|
||||
token_endpoint="https://as.example.com/token",
|
||||
registration_endpoint=None,
|
||||
revocation_endpoint=revocation_endpoint,
|
||||
jwks_uri=None,
|
||||
code_challenge_methods_supported=("S256",),
|
||||
token_endpoint_auth_methods_supported=("client_secret_basic",),
|
||||
)
|
||||
|
||||
|
||||
def _mk_response(status_code: int) -> MagicMock:
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = status_code
|
||||
return resp
|
||||
|
||||
|
||||
class TestRevocationUnsupported:
|
||||
def test_revoke_token_skipped_when_revocation_endpoint_none(self) -> None:
|
||||
as_meta = _make_as_metadata(revocation_endpoint=None)
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock()
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
client.post.assert_not_called()
|
||||
info_events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_unsupported" in info_events
|
||||
|
||||
|
||||
class TestRevocationSuccess:
|
||||
def test_revoke_token_succeeds_on_200(self) -> None:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(return_value=_mk_response(200))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret="s-secret",
|
||||
)
|
||||
)
|
||||
|
||||
# POST shape — URL + form body keys.
|
||||
client.post.assert_awaited_once()
|
||||
call_args = client.post.call_args
|
||||
assert call_args.args[0] == "https://as.example.com/revoke"
|
||||
body = call_args.kwargs["data"]
|
||||
assert body == {
|
||||
"token": "r-secret",
|
||||
"token_type_hint": "refresh_token",
|
||||
"client_id": "client-1",
|
||||
"client_secret": "s-secret",
|
||||
}
|
||||
info_events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_succeeded" in info_events
|
||||
|
||||
def test_revoke_token_omits_client_secret_when_none(self) -> None:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(return_value=_mk_response(200))
|
||||
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
body = client.post.call_args.kwargs["data"]
|
||||
assert "client_secret" not in body
|
||||
assert body["token"] == "r-secret"
|
||||
assert body["token_type_hint"] == "refresh_token"
|
||||
assert body["client_id"] == "client-1"
|
||||
|
||||
def test_revoke_token_succeeds_on_204(self) -> None:
|
||||
# RFC 7009 says the AS MAY return any 2xx; treat the whole range
|
||||
# as success.
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(return_value=_mk_response(204))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
info_events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_succeeded" in info_events
|
||||
|
||||
|
||||
class TestRevocationFailureLogged:
|
||||
def _run_and_capture(self, status: int) -> list[Any]:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(return_value=_mk_response(status))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
return mock_log.info.call_args_list
|
||||
|
||||
def test_revoke_token_logs_on_400_does_not_raise(self) -> None:
|
||||
calls = self._run_and_capture(400)
|
||||
events = [c.args[0] for c in calls]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
# Must include status field.
|
||||
failed_call = next(c for c in calls if c.args[0] == "mcp_server.oauth.revocation_failed")
|
||||
assert failed_call.kwargs.get("status") == 400
|
||||
|
||||
def test_revoke_token_logs_on_401_does_not_raise(self) -> None:
|
||||
calls = self._run_and_capture(401)
|
||||
events = [c.args[0] for c in calls]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
failed_call = next(c for c in calls if c.args[0] == "mcp_server.oauth.revocation_failed")
|
||||
assert failed_call.kwargs.get("status") == 401
|
||||
|
||||
def test_revoke_token_logs_on_403_does_not_raise(self) -> None:
|
||||
calls = self._run_and_capture(403)
|
||||
events = [c.args[0] for c in calls]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
|
||||
def test_revoke_token_logs_on_5xx_does_not_raise(self) -> None:
|
||||
calls = self._run_and_capture(500)
|
||||
events = [c.args[0] for c in calls]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
failed_call = next(c for c in calls if c.args[0] == "mcp_server.oauth.revocation_failed")
|
||||
assert failed_call.kwargs.get("status") == 500
|
||||
|
||||
|
||||
class TestRevocationExceptionPaths:
|
||||
def test_revoke_token_handles_network_error(self) -> None:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(side_effect=httpx.ConnectError("boom"))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
failed_call = next(
|
||||
c
|
||||
for c in mock_log.info.call_args_list
|
||||
if c.args[0] == "mcp_server.oauth.revocation_failed"
|
||||
)
|
||||
assert failed_call.kwargs.get("error") == "ConnectError"
|
||||
|
||||
def test_revoke_token_handles_httpx_timeout(self) -> None:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(side_effect=httpx.TimeoutException("slow"))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
failed_call = next(
|
||||
c
|
||||
for c in mock_log.info.call_args_list
|
||||
if c.args[0] == "mcp_server.oauth.revocation_failed"
|
||||
)
|
||||
assert failed_call.kwargs.get("error") == "TimeoutException"
|
||||
|
||||
def test_revoke_token_handles_asyncio_timeout(self) -> None:
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
async def _slow(*_args: Any, **_kwargs: Any) -> Any:
|
||||
await asyncio.sleep(10.0)
|
||||
raise AssertionError("should have timed out")
|
||||
|
||||
client.post = AsyncMock(side_effect=_slow)
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
timeout_seconds=0.05,
|
||||
)
|
||||
)
|
||||
|
||||
events = [c.args[0] for c in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.revocation_failed" in events
|
||||
failed_call = next(
|
||||
c
|
||||
for c in mock_log.info.call_args_list
|
||||
if c.args[0] == "mcp_server.oauth.revocation_failed"
|
||||
)
|
||||
# ``asyncio.timeout`` raises ``TimeoutError`` (Python's builtin)
|
||||
# on cancellation.
|
||||
assert failed_call.kwargs.get("error") == "TimeoutError"
|
||||
|
||||
def test_revoke_token_no_exc_info_in_logs(self) -> None:
|
||||
"""Bearer-leak invariant: the revoke path must NEVER set
|
||||
``exc_info=True``. Chained ``__context__`` may include an
|
||||
``httpx.Request`` whose ``Authorization`` header holds a
|
||||
bearer; the traceback formatter would render it.
|
||||
"""
|
||||
as_meta = _make_as_metadata()
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = AsyncMock(side_effect=httpx.ConnectError("boom"))
|
||||
|
||||
with patch("turnstone.core.mcp_oauth.log") as mock_log:
|
||||
asyncio.run(
|
||||
revoke_token_at_as(
|
||||
as_metadata=as_meta,
|
||||
http_client=client,
|
||||
refresh_token="r-secret",
|
||||
client_id="client-1",
|
||||
client_secret=None,
|
||||
)
|
||||
)
|
||||
|
||||
# No info call may carry exc_info.
|
||||
for call in mock_log.info.call_args_list:
|
||||
assert "exc_info" not in call.kwargs, (
|
||||
f"mcp_server.oauth log info({call.args[0]!r}) used exc_info — "
|
||||
"this violates the bearer-leak invariant"
|
||||
)
|
||||
# Defensively: also check warning + exception levels for the
|
||||
# same call site.
|
||||
for call in mock_log.warning.call_args_list:
|
||||
assert "exc_info" not in call.kwargs
|
||||
mock_log.exception.assert_not_called()
|
||||
|
||||
|
||||
class TestAttemptUpstreamRevokeNeverRaises:
|
||||
"""Round-2 q-3 regression: ``_attempt_upstream_revoke``'s docstring
|
||||
claims ``Never raises``. Background-task semantics make this load-
|
||||
bearing — a propagated exception logs ``Task exception was never
|
||||
retrieved`` because the ``set.discard`` done-callback doesn't read
|
||||
``task.exception()``.
|
||||
|
||||
The wrapper's narrow inner ``except`` clauses (``MCPOAuthDiscoveryError``,
|
||||
``MCPTokenDecryptError``) leave room for any other exception type
|
||||
raised by ``discover_authorization_server`` /
|
||||
``storage.get_mcp_oauth_client_secret_ct`` / ``token_store.cipher.decrypt``
|
||||
to escape. The outer ``try/except Exception`` is what keeps the
|
||||
contract honest. These tests pin that gate.
|
||||
"""
|
||||
|
||||
def _build_args(self) -> dict[str, Any]:
|
||||
token_store = MagicMock()
|
||||
token_store.cipher = MagicMock()
|
||||
token_store.cipher.decrypt.return_value = b"shh"
|
||||
storage = MagicMock()
|
||||
storage.get_mcp_oauth_client_secret_ct.return_value = None
|
||||
return {
|
||||
"http_client": MagicMock(spec=httpx.AsyncClient),
|
||||
"metadata_cache": None,
|
||||
"storage": storage,
|
||||
"token_store": token_store,
|
||||
"server_name": "srv-oauth",
|
||||
"server_row": {
|
||||
"url": "https://mcp.example.com",
|
||||
"oauth_client_id": "client-1",
|
||||
"oauth_authorization_server_url": None,
|
||||
"oauth_as_issuer_cached": None,
|
||||
},
|
||||
"server_id_for_audit": "srv-id-1",
|
||||
"refresh_token": "r-secret",
|
||||
}
|
||||
|
||||
def test_attempt_upstream_revoke_swallows_unexpected_exception(self) -> None:
|
||||
"""A generic exception from a path the inner handlers don't
|
||||
cover MUST be caught at the outer boundary and logged with type
|
||||
name only (no exc_info=True per the bearer-leak invariant).
|
||||
"""
|
||||
args = self._build_args()
|
||||
|
||||
async def _boom(*_a: Any, **_kw: Any) -> Any:
|
||||
raise RuntimeError("network blew up")
|
||||
|
||||
with (
|
||||
patch("turnstone.core.mcp_oauth.discover_authorization_server", side_effect=_boom),
|
||||
patch("turnstone.core.mcp_oauth.log") as mock_log,
|
||||
):
|
||||
# MUST NOT raise.
|
||||
asyncio.run(_attempt_upstream_revoke(**args))
|
||||
|
||||
events = [call.args[0] for call in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.upstream_revoke_failed" in events, (
|
||||
"outer try/except must log mcp_server.oauth.upstream_revoke_failed "
|
||||
"with the exception type name when an unexpected exception escapes "
|
||||
"the narrow inner handlers"
|
||||
)
|
||||
for call in mock_log.info.call_args_list:
|
||||
assert "exc_info" not in call.kwargs, (
|
||||
"outer-block log must not use exc_info=True — chained "
|
||||
"__context__ may carry an httpx.Request bearer"
|
||||
)
|
||||
|
||||
def test_attempt_upstream_revoke_logs_discovery_failure(self) -> None:
|
||||
"""Round-2 bug-1: ``MCPOAuthDiscoveryError`` MUST emit
|
||||
``upstream_revoke_discovery_failed`` so operators have visibility
|
||||
into a silent-discovery-failure path that previously logged
|
||||
nothing while the audit row recorded ``upstream_revoke_outcome=scheduled``.
|
||||
"""
|
||||
args = self._build_args()
|
||||
|
||||
async def _disc_fail(*_a: Any, **_kw: Any) -> Any:
|
||||
raise MCPOAuthDiscoveryError("PRM fetch 503")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_oauth.discover_authorization_server",
|
||||
side_effect=_disc_fail,
|
||||
),
|
||||
patch("turnstone.core.mcp_oauth.log") as mock_log,
|
||||
):
|
||||
asyncio.run(_attempt_upstream_revoke(**args))
|
||||
|
||||
events = [call.args[0] for call in mock_log.info.call_args_list]
|
||||
assert "mcp_server.oauth.upstream_revoke_discovery_failed" in events
|
||||
assert "mcp_server.oauth.upstream_revoke_failed" not in events
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Storage CRUD tests for the per-(user, server) MCP OAuth pending-state table.
|
||||
|
||||
Validates the storage-protocol additions for the per-(user, server)
|
||||
OAuth flow:
|
||||
|
||||
- ``create_mcp_oauth_pending_state``
|
||||
- ``pop_mcp_oauth_pending_state`` (atomic, with TTL)
|
||||
- ``cleanup_expired_mcp_oauth_pending_states``
|
||||
- ``get_mcp_oauth_client_secret_ct``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
class TestCreateAndPop:
|
||||
def test_round_trip(self, backend) -> None:
|
||||
backend.create_mcp_oauth_pending_state(
|
||||
"state-1",
|
||||
"user-a",
|
||||
"srv-x",
|
||||
"verifier-blob",
|
||||
"/admin/mcp-servers",
|
||||
)
|
||||
row = backend.pop_mcp_oauth_pending_state("state-1", max_age_seconds=600)
|
||||
assert row is not None
|
||||
assert row["state"] == "state-1"
|
||||
assert row["user_id"] == "user-a"
|
||||
assert row["server_name"] == "srv-x"
|
||||
assert row["code_verifier"] == "verifier-blob"
|
||||
assert row["return_url"] == "/admin/mcp-servers"
|
||||
|
||||
def test_pop_consumes_row(self, backend) -> None:
|
||||
backend.create_mcp_oauth_pending_state("s2", "u", "s", "v", "/r")
|
||||
first = backend.pop_mcp_oauth_pending_state("s2")
|
||||
assert first is not None
|
||||
# Second pop must miss — row was consumed.
|
||||
second = backend.pop_mcp_oauth_pending_state("s2")
|
||||
assert second is None
|
||||
|
||||
def test_pop_missing_returns_none(self, backend) -> None:
|
||||
assert backend.pop_mcp_oauth_pending_state("never-existed") is None
|
||||
|
||||
|
||||
class TestTTL:
|
||||
def test_pop_rejects_expired_row(self, backend) -> None:
|
||||
backend.create_mcp_oauth_pending_state("old-state", "u", "s", "v", "/r")
|
||||
# Backdate it so it's older than the TTL window.
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"UPDATE mcp_oauth_pending SET created_at = '2020-01-01T00:00:00' "
|
||||
"WHERE state = 'old-state'"
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
# Default TTL is 600s — the row is decades old.
|
||||
row = backend.pop_mcp_oauth_pending_state("old-state")
|
||||
assert row is None
|
||||
|
||||
# Even though pop returned None, the row must have been wiped — a
|
||||
# second pop with a giant TTL must still see nothing.
|
||||
again = backend.pop_mcp_oauth_pending_state("old-state", max_age_seconds=10**9)
|
||||
assert again is None
|
||||
|
||||
def test_pop_accepts_fresh_row(self, backend) -> None:
|
||||
backend.create_mcp_oauth_pending_state("fresh", "u", "s", "v", "/r")
|
||||
row = backend.pop_mcp_oauth_pending_state("fresh", max_age_seconds=600)
|
||||
assert row is not None
|
||||
assert row["state"] == "fresh"
|
||||
|
||||
|
||||
class TestCleanup:
|
||||
def test_cleanup_deletes_only_expired(self, backend) -> None:
|
||||
backend.create_mcp_oauth_pending_state("old", "u", "s", "v", "/r")
|
||||
backend.create_mcp_oauth_pending_state("new", "u", "s", "v", "/r")
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"UPDATE mcp_oauth_pending SET created_at = '2020-01-01T00:00:00' "
|
||||
"WHERE state = 'old'"
|
||||
)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
deleted = backend.cleanup_expired_mcp_oauth_pending_states(max_age_seconds=600)
|
||||
assert deleted == 1
|
||||
# Old gone, new still around.
|
||||
assert backend.pop_mcp_oauth_pending_state("old") is None
|
||||
survivor = backend.pop_mcp_oauth_pending_state("new")
|
||||
assert survivor is not None
|
||||
|
||||
def test_cleanup_no_rows(self, backend) -> None:
|
||||
assert backend.cleanup_expired_mcp_oauth_pending_states() == 0
|
||||
|
||||
|
||||
class TestGetOAuthClientSecretCt:
|
||||
def test_returns_none_when_unset(self, backend) -> None:
|
||||
backend.create_mcp_server(
|
||||
server_id="srv-id",
|
||||
name="srv-x",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
)
|
||||
assert backend.get_mcp_oauth_client_secret_ct("srv-id") is None
|
||||
|
||||
def test_returns_ciphertext_after_set(self, backend) -> None:
|
||||
backend.create_mcp_server(
|
||||
server_id="srv-id",
|
||||
name="srv-x",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
)
|
||||
ct = b"\x00\xff\x42encrypted-blob"
|
||||
ok = backend.set_mcp_oauth_client_secret_ct("srv-id", ct)
|
||||
assert ok is True
|
||||
out = backend.get_mcp_oauth_client_secret_ct("srv-id")
|
||||
assert out == ct
|
||||
|
||||
def test_returns_none_for_missing_server(self, backend) -> None:
|
||||
assert backend.get_mcp_oauth_client_secret_ct("does-not-exist") is None
|
||||
|
||||
|
||||
def _create_user_token_row(
|
||||
backend,
|
||||
*,
|
||||
user_id: str,
|
||||
server_name: str,
|
||||
created: str,
|
||||
) -> None:
|
||||
"""Insert a token row + backdate ``created`` so ordering is deterministic.
|
||||
|
||||
The storage helper stamps ``created`` from ``datetime.now(UTC)``; for
|
||||
multi-row ordering tests we backdate via raw SQL so the inserts stay
|
||||
independent of clock resolution.
|
||||
"""
|
||||
backend.create_mcp_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token_ct=b"ct-access",
|
||||
refresh_token_ct=b"ct-refresh",
|
||||
expires_at="2026-05-04T12:00:00",
|
||||
scopes="openid",
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"UPDATE mcp_user_tokens SET created = :created "
|
||||
"WHERE user_id = :uid AND server_name = :sn"
|
||||
),
|
||||
{"created": created, "uid": user_id, "sn": server_name},
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
class TestListMCPUserTokenMetadataByUser:
|
||||
def test_list_mcp_user_token_metadata_by_user_empty(self, backend) -> None:
|
||||
assert backend.list_mcp_user_token_metadata_by_user("nobody") == []
|
||||
|
||||
def test_list_mcp_user_token_metadata_by_user_single_server(self, backend) -> None:
|
||||
_create_user_token_row(
|
||||
backend, user_id="u1", server_name="srv-a", created="2026-05-01T00:00:00"
|
||||
)
|
||||
rows = backend.list_mcp_user_token_metadata_by_user("u1")
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["user_id"] == "u1"
|
||||
assert rows[0]["server_name"] == "srv-a"
|
||||
assert rows[0]["as_issuer"] == "https://auth.example.com"
|
||||
assert rows[0]["audience"] == "https://mcp.example.com"
|
||||
assert rows[0]["scopes"] == "openid"
|
||||
# Projection MUST omit ciphertext columns — the SQL no longer
|
||||
# selects them, so the TypedDict has no key.
|
||||
assert "access_token_ct" not in rows[0]
|
||||
assert "refresh_token_ct" not in rows[0]
|
||||
|
||||
def test_list_mcp_user_token_metadata_by_user_multiple_servers(self, backend) -> None:
|
||||
_create_user_token_row(
|
||||
backend, user_id="u1", server_name="srv-c", created="2026-05-03T00:00:00"
|
||||
)
|
||||
_create_user_token_row(
|
||||
backend, user_id="u1", server_name="srv-a", created="2026-05-01T00:00:00"
|
||||
)
|
||||
_create_user_token_row(
|
||||
backend, user_id="u1", server_name="srv-b", created="2026-05-02T00:00:00"
|
||||
)
|
||||
rows = backend.list_mcp_user_token_metadata_by_user("u1")
|
||||
assert [r["server_name"] for r in rows] == ["srv-a", "srv-b", "srv-c"]
|
||||
|
||||
def test_list_mcp_user_token_metadata_by_user_isolates_by_user(self, backend) -> None:
|
||||
_create_user_token_row(
|
||||
backend, user_id="user-a", server_name="srv-a", created="2026-05-01T00:00:00"
|
||||
)
|
||||
_create_user_token_row(
|
||||
backend, user_id="user-a", server_name="srv-b", created="2026-05-02T00:00:00"
|
||||
)
|
||||
_create_user_token_row(
|
||||
backend, user_id="user-b", server_name="srv-a", created="2026-05-03T00:00:00"
|
||||
)
|
||||
rows_a = backend.list_mcp_user_token_metadata_by_user("user-a")
|
||||
assert {r["server_name"] for r in rows_a} == {"srv-a", "srv-b"}
|
||||
assert all(r["user_id"] == "user-a" for r in rows_a)
|
||||
|
||||
rows_b = backend.list_mcp_user_token_metadata_by_user("user-b")
|
||||
assert len(rows_b) == 1
|
||||
assert rows_b[0]["user_id"] == "user-b"
|
||||
assert rows_b[0]["server_name"] == "srv-a"
|
||||
@@ -0,0 +1,912 @@
|
||||
"""Phase 6 integration tests — real-transport drives 401/403 through the SDK.
|
||||
|
||||
These are the structural exit criterion for Phase 6. They MUST drive
|
||||
through the real ``streamablehttp_client``, the real httpx response-hook
|
||||
path, and a REAL upstream MCP server (a ``FastMCP`` in-process subprocess
|
||||
with a starlette middleware that programmatically returns 401/403 with
|
||||
crafted ``WWW-Authenticate`` headers).
|
||||
|
||||
Direct ``httpx.HTTPStatusError`` injection is FORBIDDEN here — Phase 5
|
||||
bug-1 was masked precisely by that pattern (the production code path
|
||||
was structurally unreachable, but the unit-test injection bypassed the
|
||||
SDK's swallow). The integration tests gate that the production path
|
||||
actually receives the carrier signal end-to-end.
|
||||
|
||||
The fixture upstream is built in-thread (uvicorn on its own asyncio
|
||||
loop in a background thread) — same pattern as
|
||||
``tests/spike_sdk_concurrency.py``. Per the orchestrator's startup-cost
|
||||
note, measured at ~0.05s per fixture spin-up locally; well under the
|
||||
2s threshold for default-collection inclusion.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import uvicorn
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.mcp_oauth import TokenLookupResult
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
# Quiet noisy logs during tests.
|
||||
logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
logging.getLogger("mcp").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixture upstream — programmable BehaviorMiddleware
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BehaviorMiddleware(BaseHTTPMiddleware):
|
||||
"""Inspects per-request behaviour state and returns 401/403 on demand.
|
||||
|
||||
The behaviour is steered by a mutable ``behaviour`` dict on the
|
||||
middleware instance; tests mutate it via the fixture handle.
|
||||
Records every request's Authorization header for assertion.
|
||||
|
||||
Behaviour semantics:
|
||||
* ``"once_401"``: return 401 once, then 200 thereafter.
|
||||
* ``"always_401"``: always return 401.
|
||||
* ``"once_403_insufficient"``: return 403 with insufficient_scope once.
|
||||
* ``"once_403_generic"``: return 403 without error param once.
|
||||
* ``"once_multi_www_auth_403"``: return 403 with TWO
|
||||
``WWW-Authenticate`` headers — first ``Bearer`` challenge
|
||||
carries the SAFE scopes, second carries INJECTED scopes. The
|
||||
dispatcher must report only the first.
|
||||
* ``"never"`` (default): pass through to the real handler.
|
||||
|
||||
``www_authenticate`` overrides the default header crafted per shape.
|
||||
"""
|
||||
|
||||
def __init__(self, app: Any, behaviour: dict[str, Any]) -> None:
|
||||
super().__init__(app)
|
||||
self._behaviour = behaviour
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response:
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
|
||||
# Record the Authorization header for assertion. POST is the
|
||||
# tools/call request the dispatcher sends.
|
||||
if request.method == "POST" and "/mcp" in str(request.url):
|
||||
self._behaviour.setdefault("post_auth_headers", []).append(
|
||||
request.headers.get("authorization")
|
||||
)
|
||||
|
||||
mode = self._behaviour.get("mode", "never")
|
||||
if mode == "once_401":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "always_401":
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_insufficient":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate",
|
||||
'Bearer error="insufficient_scope", scope="files:write mail:send"',
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_generic" and not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get("www_authenticate", "Bearer realm=mcp")
|
||||
},
|
||||
)
|
||||
elif mode == "once_multi_www_auth_403" and not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
# Two ``WWW-Authenticate: Bearer ...`` challenges. The
|
||||
# first carries ``error=insufficient_scope`` but NO
|
||||
# ``scope=`` parameter; the second carries the INJECTED
|
||||
# scopes the dispatcher must NOT report. The first
|
||||
# challenge intentionally lacks ``scope`` because
|
||||
# ``parse_www_authenticate_bearer`` uses ``setdefault`` —
|
||||
# if the first challenge HAD a scope, ``setdefault`` would
|
||||
# already win on first-occurrence. The vector this test
|
||||
# guards is the case where a defended absence becomes a
|
||||
# silent presence: a hook regression to ``get(...)`` joins
|
||||
# repeated headers with ``, `` and the parser then folds
|
||||
# the second challenge's scope into the first challenge's
|
||||
# params dict because there is no first-occurrence to
|
||||
# protect.
|
||||
response = StarletteResponse("forbidden", status_code=403)
|
||||
response.headers.append(
|
||||
"www-authenticate",
|
||||
'Bearer realm="legit", error="insufficient_scope"',
|
||||
)
|
||||
response.headers.append(
|
||||
"www-authenticate",
|
||||
'Bearer error="insufficient_scope", scope="org:admin db:write"',
|
||||
)
|
||||
return response
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
s = socket.socket()
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
s.close()
|
||||
return port
|
||||
|
||||
|
||||
def _build_server(port: int, behaviour: dict[str, Any]) -> uvicorn.Server:
|
||||
mcp = FastMCP(name="phase6-target", streamable_http_path="/mcp")
|
||||
|
||||
@mcp.tool()
|
||||
async def echo(payload: str = "default") -> str:
|
||||
return f"echoed:{payload}"
|
||||
|
||||
app = mcp.streamable_http_app()
|
||||
app.add_middleware(BehaviorMiddleware, behaviour=behaviour)
|
||||
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning", access_log=False)
|
||||
return uvicorn.Server(config)
|
||||
|
||||
|
||||
def _wait_ready(port: int, timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError(f"upstream at 127.0.0.1:{port} not ready after {timeout}s")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upstream():
|
||||
"""Boot a FastMCP fixture upstream in a background thread.
|
||||
|
||||
Yields ``(url, behaviour)`` where ``behaviour`` is a mutable dict
|
||||
the test mutates to steer the middleware (set ``mode`` to one of
|
||||
the BehaviorMiddleware shapes).
|
||||
"""
|
||||
port = _find_free_port()
|
||||
behaviour: dict[str, Any] = {}
|
||||
server = _build_server(port, behaviour)
|
||||
|
||||
def _run() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(server.serve())
|
||||
|
||||
t = threading.Thread(target=_run, daemon=True, name="phase6-upstream")
|
||||
t.start()
|
||||
try:
|
||||
_wait_ready(port)
|
||||
yield f"http://127.0.0.1:{port}/mcp", behaviour
|
||||
finally:
|
||||
server.should_exit = True
|
||||
t.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def _seed_oauth_server(
|
||||
storage: SQLiteBackend,
|
||||
*,
|
||||
name: str = "pool-srv",
|
||||
server_id: str = "srv-pool",
|
||||
url: str = "https://mcp.example.com/sse",
|
||||
) -> None:
|
||||
storage.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url=url,
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="client-abc",
|
||||
oauth_scopes="openid",
|
||||
oauth_audience=url,
|
||||
)
|
||||
|
||||
|
||||
def _seed_user_token(
|
||||
storage: SQLiteBackend,
|
||||
cipher: Any,
|
||||
*,
|
||||
user_id: str = "user-1",
|
||||
server_name: str = "pool-srv",
|
||||
expires_in_seconds: int = 3600,
|
||||
access_token: str = "access-aaa",
|
||||
refresh_token: str | None = "refresh-rrr",
|
||||
) -> None:
|
||||
expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
)
|
||||
store = MCPTokenStore(storage, cipher, node_id="test")
|
||||
store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_at=expires_at,
|
||||
scopes="openid",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
|
||||
def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"),
|
||||
mcp_oauth_http_client=MagicMock(),
|
||||
mcp_oauth_refresh_locks={},
|
||||
mcp_oauth_metadata_cache={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_loop_mgr():
|
||||
cfg: dict[str, Any] = {}
|
||||
mgr = MCPClientManager(cfg)
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop")
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
yield mgr, loop, thread
|
||||
finally:
|
||||
|
||||
async def _drain(m: MCPClientManager) -> None:
|
||||
task = m._user_pool_eviction_task
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
m._user_pool_eviction_task = None
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 21: 401 → refresh-and-retry → success
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_401_refresh_and_retry_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Real upstream returns 401 once with ``WWW-Authenticate: Bearer
|
||||
error="invalid_token"``, then 200. Dispatcher carrier captures the
|
||||
401, ``force_refresh=True`` mints a new bearer (stubbed), retry
|
||||
succeeds. Hard invariant 3: breaker counter remains 0.
|
||||
|
||||
Drives through the REAL ``streamablehttp_client`` and a REAL
|
||||
upstream subprocess (no ``httpx.HTTPStatusError`` injection). This
|
||||
is the structural exit gate for Phase 6 — the equivalent unit
|
||||
tests CANNOT prove the production wiring works because the SDK
|
||||
swallows the underlying exception.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
# Override URL to point at the local upstream (loopback http:// is
|
||||
# exempt from the URL-validator).
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
result = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__echo", {"payload": "hi"}, user_id="user-1", timeout=15
|
||||
)
|
||||
|
||||
assert "echoed:hi" in result
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
# Server saw at least 2 POSTs to /mcp (initial + retry).
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}"
|
||||
# Retry carries a different bearer than the initial.
|
||||
initial = post_headers[0]
|
||||
retry = post_headers[1]
|
||||
assert initial != retry, (
|
||||
"retry attached the same bearer as the initial; the dispatcher "
|
||||
"did not pick up the refreshed token."
|
||||
)
|
||||
# Pool entry has a session after the successful retry.
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 22: 401 + refresh failure → mcp_consent_required
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_401_with_refresh_failure_emits_consent_required(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="refresh_failed")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
|
||||
|
||||
# Structured-error envelopes flow back via ``RuntimeError(json_str)``
|
||||
# so the session-layer ``except Exception`` handler routes the
|
||||
# consent card uniformly across tool / resource / prompt dispatchers.
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert payload["error"]["server"] == "pool-srv"
|
||||
# Phase 8 — consent_url surfaces a /start URL the dashboard can open
|
||||
# in a popup. URL-encoded server name; no scopes baked in (the AS
|
||||
# picks up the configured scopes server-side at /start).
|
||||
assert payload["error"]["consent_url"] == "/v1/api/mcp/oauth/start?server=pool-srv"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 23: 403 + insufficient_scope → mcp_insufficient_scope with parsed scopes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_403_insufficient_scope_emits_structured_error(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_insufficient"
|
||||
behaviour["www_authenticate"] = (
|
||||
'Bearer error="insufficient_scope", scope="files:write mail:send"'
|
||||
)
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_insufficient_scope"
|
||||
assert payload["error"]["scopes_required"] == ["files:write", "mail:send"]
|
||||
# Phase 8 — consent_url carries the step-up scopes URL-encoded so the
|
||||
# dashboard can union them with the configured set at /start.
|
||||
assert payload["error"]["consent_url"] == (
|
||||
"/v1/api/mcp/oauth/start?server=pool-srv&scopes=files%3Awrite%20mail%3Asend"
|
||||
)
|
||||
# No retry — exactly ONE POST attempted before the structured error.
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) == 1, (
|
||||
f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs"
|
||||
)
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 24: 403 without insufficient_scope → generic forbidden
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_403_no_insufficient_scope_emits_generic_forbidden(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_generic"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_tool_call_forbidden"
|
||||
assert "scopes_required" not in payload["error"]
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) == 1, (
|
||||
f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sec-1: multi-WWW-Authenticate header injection — only the FIRST
|
||||
# Bearer challenge feeds the structured-error / audit emission.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_403_multi_www_authenticate_drops_injected_scopes(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Upstream returns a 403 with TWO ``WWW-Authenticate: Bearer ...``
|
||||
challenges. The first carries ``error=insufficient_scope`` but NO
|
||||
``scope=`` parameter; the second carries INJECTED scopes
|
||||
(``["org:admin", "db:write"]``). The dispatcher must report
|
||||
``scopes_required == []`` — derived from the first challenge alone
|
||||
— never the second challenge's injected scopes.
|
||||
|
||||
Two layers of defence cooperate (either alone neutralises the
|
||||
vector; both run together so a regression in one cannot silently
|
||||
re-open it):
|
||||
|
||||
1. ``_make_capturing_http_factory._hook`` reads
|
||||
``response.headers.get_list("www-authenticate")[0]`` rather than
|
||||
``response.headers.get(...)`` — the latter joins repeated
|
||||
headers with ``", "`` which the RFC 7235 tokenizer would
|
||||
otherwise consume as a continuation of the first challenge.
|
||||
2. ``parse_www_authenticate_bearer`` stops at the first ``Bearer``
|
||||
challenge boundary even if the input was already joined, so a
|
||||
hook regression to ``get(...)`` would NOT re-open the vector.
|
||||
|
||||
The first challenge intentionally lacks ``scope=`` — the parser
|
||||
uses ``setdefault`` so a first-occurrence ``scope`` would already
|
||||
win and mask a single-layer regression. The undefended-absence
|
||||
case is what proves both layers actually do their job.
|
||||
|
||||
Negative-test (CRITICAL — Phase 5 lesson): verified by reverting
|
||||
the hook to ``response.headers.get("www-authenticate")`` AND
|
||||
removing the ``_looks_like_bearer_challenge_start`` guard in
|
||||
``parse_www_authenticate_bearer``. The test then fails because
|
||||
``scopes_required`` becomes ``["org:admin", "db:write"]`` — the
|
||||
injected scopes from the second challenge silently fold into the
|
||||
first challenge's params dict via httpx's comma-joined header
|
||||
value (the absence of a first-occurrence scope means nothing
|
||||
blocks the fold).
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_multi_www_auth_403"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_insufficient_scope", (
|
||||
f"expected mcp_insufficient_scope; got {payload!r}"
|
||||
)
|
||||
# ``scopes_required`` derives from the FIRST challenge alone, which
|
||||
# carries no ``scope=`` parameter. The injected second challenge
|
||||
# MUST NOT appear here.
|
||||
assert payload["error"]["scopes_required"] == [], (
|
||||
"Multi-header injection slipped through: dispatcher reported "
|
||||
"scopes from the SECOND Bearer challenge. Got "
|
||||
f"{payload['error']['scopes_required']!r}; expected []."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 25: 401 retry ceiling — never recurse
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_401_retry_ceiling(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Upstream always returns 401; refresh stub keeps minting tokens.
|
||||
After exactly ONE retry, dispatcher emits ``mcp_consent_required``.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
refresh_count = 0
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
nonlocal refresh_count
|
||||
if kwargs.get("force_refresh"):
|
||||
refresh_count += 1
|
||||
return TokenLookupResult(kind="token", token=f"refreshed-{refresh_count}")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
# Exactly ONE refresh round-trip.
|
||||
assert refresh_count == 1, f"expected exactly 1 refresh round-trip; got {refresh_count}"
|
||||
# Server saw EXACTLY 2 POSTs (initial + 1 retry).
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) == 2, (
|
||||
f"expected exactly 2 POSTs (initial + 1 retry); got {len(post_headers)}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 26: breaker unaffected by repeated auth failures (slow — 50 cycles)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_breaker_unaffected_by_auth_failures(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""50 sequential dispatches all hit 401 with refresh-failed → 50
|
||||
cycles of ``mcp_consent_required``. ``_consecutive_failures`` MUST
|
||||
stay at 0 throughout (hard invariant 3 verified end-to-end).
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="refresh_failed")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
for _ in range(50):
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
|
||||
)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 27: static path unaffected by Phase 6 changes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_static_path_unaffected(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Static-path connect against an unauthed upstream succeeds without
|
||||
going through the capturing factory. This is the integration-level
|
||||
mirror of ``test_reconnect_preserves_static_state_identity``.
|
||||
|
||||
Drives the static path against the same fixture upstream (with
|
||||
``behaviour={}`` so middleware passes through) — confirms the
|
||||
static path's session lifecycle is byte-identical even when the
|
||||
pool path's auth introspection is wired up.
|
||||
"""
|
||||
url, _behaviour = upstream
|
||||
# No mode → middleware passes through to FastMCP.
|
||||
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
|
||||
# Manually configure mgr with a static-path server pointing at the
|
||||
# fixture upstream. Use _connect_one (not the pool path).
|
||||
cfg = {"type": "streamable-http", "url": url}
|
||||
|
||||
async def _connect_static() -> None:
|
||||
await mgr._connect_one("static-srv", cfg)
|
||||
|
||||
fut = asyncio.run_coroutine_threadsafe(_connect_static(), loop)
|
||||
fut.result(timeout=15)
|
||||
|
||||
state_before = mgr._static_servers.get("static-srv")
|
||||
assert state_before is not None
|
||||
assert state_before.session is not None
|
||||
# Snapshot identity.
|
||||
state_id_before = id(state_before)
|
||||
session_before = state_before.session
|
||||
|
||||
# Reconnect — the canonical regression check is that the
|
||||
# StaticServerState object identity is preserved.
|
||||
fut = asyncio.run_coroutine_threadsafe(_connect_static(), loop)
|
||||
fut.result(timeout=15)
|
||||
|
||||
state_after = mgr._static_servers.get("static-srv")
|
||||
assert state_after is not None
|
||||
assert id(state_after) == state_id_before, (
|
||||
"Static path StaticServerState identity changed across reconnect; "
|
||||
"hard invariant 1 violated."
|
||||
)
|
||||
assert state_after.session is not None
|
||||
assert state_after.session is not session_before, (
|
||||
"Reconnect did not actually replace the session"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 27b: static dispatch unaffected by Phase 8 consent_url kwarg
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_static_dispatch_unaffected_by_consent_url_kwarg(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Static-auth tool dispatch must be byte-identical post-Phase 8.
|
||||
|
||||
The Phase 8 changes only ADD a ``consent_url`` kwarg to
|
||||
``_structured_error`` invocations on the pool path. Static dispatch
|
||||
must not pick up the field — there's no consent flow for
|
||||
``auth_type='none'`` / ``'static'`` servers, and exposing one would
|
||||
confuse the dashboard renderer. Asserts a successful tool result is
|
||||
a plain string with no JSON envelope and no ``consent_url`` substring.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "never" # passthrough — succeeds
|
||||
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cfg = {"type": "streamable-http", "url": url}
|
||||
|
||||
async def _connect_static() -> None:
|
||||
await mgr._connect_one("static-srv", cfg)
|
||||
|
||||
fut = asyncio.run_coroutine_threadsafe(_connect_static(), loop)
|
||||
fut.result(timeout=15)
|
||||
|
||||
# Drive call_tool_sync without a user_id — the static path is taken.
|
||||
result = mgr.call_tool_sync("mcp__static-srv__echo", {"payload": "static-x"}, timeout=15)
|
||||
|
||||
# Static path returns the FastMCP fixture's echo string.
|
||||
assert "echoed:static-x" in result
|
||||
# No JSON envelope leaked through; specifically no consent_url field.
|
||||
assert "consent_url" not in result, (
|
||||
f"Static-auth tool dispatch surfaced a consent_url; result: {result!r}"
|
||||
)
|
||||
# Defensive: result is not a JSON-encoded structured error.
|
||||
try:
|
||||
parsed = json.loads(result)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
parsed = None
|
||||
if isinstance(parsed, dict):
|
||||
assert "error" not in parsed, (
|
||||
f"Static-auth dispatch returned a structured-error envelope; got {parsed!r}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 28: pool reuse — 401 on a SECOND dispatch (carrier owned by entry)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_integration_pool_reuse_401_refresh_and_retry_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Reused pool sessions still capture 401 correctly.
|
||||
|
||||
Dispatch 1 hits a passthrough upstream (200) and populates
|
||||
``entry.session``. Dispatch 2 reuses that session — no fresh
|
||||
connect, so a per-dispatch ``_AuthCapture`` would never reach
|
||||
the httpx response hook (the hook closes over the carrier passed
|
||||
at first connect, which lives on the entry). A correctly-wired
|
||||
entry-owned carrier is the only shape that lets dispatch 2's 401
|
||||
surface to the dispatcher.
|
||||
|
||||
Two independent production bugs gate this test passing; both must
|
||||
hold for reused-session 401 recovery to work end-to-end:
|
||||
|
||||
1. The carrier must live on the pool entry (not per-dispatch) so
|
||||
the response hook bound at first connect writes to the same
|
||||
object the dispatcher reads across reuse. Verified by reverting
|
||||
``PoolEntryState.auth_capture`` to a per-dispatch
|
||||
``_AuthCapture()`` allocation: the carrier-fired event never
|
||||
reaches the dispatcher and the test times out.
|
||||
|
||||
2. The dispatcher must race ``call_tool`` against the carrier's
|
||||
fired event. The SDK's ``_receive_loop`` runs in BaseSession's
|
||||
TaskGroup nested inside ``streamablehttp_client``'s TaskGroup;
|
||||
when an upstream 4xx fires, the outer TaskGroup cancels
|
||||
``_receive_loop`` mid-finally before it can deliver
|
||||
``CONNECTION_CLOSED`` to the response stream's waiting
|
||||
receiver. anyio's ``send_nowait`` skips waiters with pending
|
||||
cancellation — but our dispatch task (created via
|
||||
``run_coroutine_threadsafe`` for the reused-session case) has
|
||||
NO pending cancellation, so the send delivers but the receiver
|
||||
never wakes (the waiter's Event is set on stale state). Result:
|
||||
a forever-hung ``response_stream_reader.receive()``. Verified
|
||||
by reverting the ``asyncio.wait({call_task, fired_task})``
|
||||
race in ``_dispatch_pool_with_entry`` to a bare ``await
|
||||
session.call_tool(...)``: the test times out.
|
||||
|
||||
This test is the structural gate against the per-dispatch carrier
|
||||
pattern: it looks right in code review and passes single-dispatch
|
||||
integration tests, but breaks silently on session reuse — and the
|
||||
SDK-level hang the carrier fix exposes silently strands the
|
||||
dispatcher even when the carrier is correct.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "never" # passthrough — dispatch 1 succeeds
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
# Dispatch 1: passthrough success. Establishes the pooled session.
|
||||
result1 = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__echo", {"payload": "first"}, user_id="user-1", timeout=15
|
||||
)
|
||||
assert "echoed:first" in result1
|
||||
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
session_after_first = entry.session
|
||||
assert session_after_first is not None, (
|
||||
"test setup: dispatch 1 did not populate entry.session; "
|
||||
"subsequent dispatch will not exercise the reuse path"
|
||||
)
|
||||
|
||||
# Reconfigure upstream to 401 once on the next call. Reset the
|
||||
# auth-headers log so we can count dispatch-2's POSTs cleanly.
|
||||
behaviour["post_auth_headers"] = []
|
||||
behaviour["mode"] = "once_401"
|
||||
behaviour["_fired"] = False
|
||||
|
||||
# Dispatch 2: same (user, server). The hook from dispatch 1's
|
||||
# connect is still bound to entry.auth_capture. The 401 fires;
|
||||
# the dispatcher's auth_401 path triggers refresh-and-retry.
|
||||
result2 = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__echo", {"payload": "second"}, user_id="user-1", timeout=15
|
||||
)
|
||||
|
||||
assert "echoed:second" in result2, (
|
||||
f"reused-session 401 retry did not succeed. result: {result2!r}. "
|
||||
"If this is JSON with mcp_consent_required, the dispatcher "
|
||||
"fell through to consent_required emission; if a generic "
|
||||
"tool error, the carrier was empty (auth branch unreachable)."
|
||||
)
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0, (
|
||||
"auth failures must not trip the per-server breaker"
|
||||
)
|
||||
|
||||
# Dispatch 2 produces multiple POSTs: the original 401 with the
|
||||
# rejected bearer, then the retry's full connect handshake
|
||||
# (initialize + notifications/initialized + tools/list) followed by
|
||||
# the actual tools/call — all under the refreshed bearer. The retry
|
||||
# reconnects because the auth_401 handler evicted the broken session.
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) >= 2, (
|
||||
f"expected >=2 POSTs after dispatch 2 (401 + retry); "
|
||||
f"got {len(post_headers)}: {post_headers}"
|
||||
)
|
||||
# First POST is the original bearer that got 401'd.
|
||||
assert post_headers[0] == "Bearer access-aaa", (
|
||||
f"first POST was {post_headers[0]!r}; expected the original bearer"
|
||||
)
|
||||
# Every subsequent POST carries the refreshed bearer (the retry
|
||||
# ran with force_refresh=True and reconnected with the new token).
|
||||
refreshed = post_headers[1:]
|
||||
assert all(h == "Bearer refreshed-bearer" for h in refreshed), (
|
||||
f"retry POSTs carried unexpected bearer(s); observed: {post_headers}"
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,720 @@
|
||||
"""Phase 7b integration tests — real-transport prompt get 401/403/etc.
|
||||
|
||||
Mirror of :mod:`tests.test_mcp_pool_auth_resource_integration` for the
|
||||
prompt path (RFC §3.3). Drives through the real ``streamablehttp_client``,
|
||||
real httpx response-hook plumbing, and a real upstream subprocess
|
||||
(``FastMCP`` with a programmable ``BehaviorMiddleware``). Direct
|
||||
``httpx.HTTPStatusError`` injection is forbidden (invariant 14).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import uvicorn
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.mcp_oauth import TokenLookupResult
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
logging.getLogger("mcp").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
class BehaviorMiddleware(BaseHTTPMiddleware):
|
||||
"""Programmable upstream behaviour — see
|
||||
:mod:`tests.test_mcp_pool_auth_integration` for the semantics. This
|
||||
copy serves the prompt integration tests.
|
||||
"""
|
||||
|
||||
def __init__(self, app: Any, behaviour: dict[str, Any]) -> None:
|
||||
super().__init__(app)
|
||||
self._behaviour = behaviour
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response:
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
|
||||
if request.method == "POST" and "/mcp" in str(request.url):
|
||||
self._behaviour.setdefault("post_auth_headers", []).append(
|
||||
request.headers.get("authorization")
|
||||
)
|
||||
|
||||
mode = self._behaviour.get("mode", "never")
|
||||
if mode == "once_401":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "always_401":
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_insufficient":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate",
|
||||
'Bearer error="insufficient_scope", scope="prompts:read"',
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_generic" and not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get("www_authenticate", "Bearer realm=mcp")
|
||||
},
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
s = socket.socket()
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
s.close()
|
||||
return port
|
||||
|
||||
|
||||
def _build_server(port: int, behaviour: dict[str, Any]) -> uvicorn.Server:
|
||||
mcp = FastMCP(name="phase7b-prompt-target", streamable_http_path="/mcp")
|
||||
|
||||
@mcp.prompt()
|
||||
def greet(who: str = "world") -> str:
|
||||
return f"Hello, {who}!"
|
||||
|
||||
@mcp.prompt()
|
||||
def summarize(topic: str = "today") -> str:
|
||||
return f"Please summarize {topic}."
|
||||
|
||||
app = mcp.streamable_http_app()
|
||||
app.add_middleware(BehaviorMiddleware, behaviour=behaviour)
|
||||
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning", access_log=False)
|
||||
return uvicorn.Server(config)
|
||||
|
||||
|
||||
def _wait_ready(port: int, timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError(f"upstream at 127.0.0.1:{port} not ready after {timeout}s")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upstream():
|
||||
port = _find_free_port()
|
||||
behaviour: dict[str, Any] = {}
|
||||
server = _build_server(port, behaviour)
|
||||
|
||||
def _run() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(server.serve())
|
||||
|
||||
t = threading.Thread(target=_run, daemon=True, name="phase7b-prompt-upstream")
|
||||
t.start()
|
||||
try:
|
||||
_wait_ready(port)
|
||||
yield f"http://127.0.0.1:{port}/mcp", behaviour
|
||||
finally:
|
||||
server.should_exit = True
|
||||
t.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def _seed_oauth_server(
|
||||
storage: SQLiteBackend,
|
||||
*,
|
||||
name: str = "pool-srv",
|
||||
server_id: str = "srv-pool",
|
||||
url: str = "https://mcp.example.com/sse",
|
||||
) -> None:
|
||||
storage.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url=url,
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="client-abc",
|
||||
oauth_scopes="openid",
|
||||
oauth_audience=url,
|
||||
)
|
||||
|
||||
|
||||
def _seed_user_token(
|
||||
storage: SQLiteBackend,
|
||||
cipher: Any,
|
||||
*,
|
||||
user_id: str = "user-1",
|
||||
server_name: str = "pool-srv",
|
||||
expires_in_seconds: int = 3600,
|
||||
access_token: str = "access-aaa",
|
||||
refresh_token: str | None = "refresh-rrr",
|
||||
) -> None:
|
||||
expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
)
|
||||
store = MCPTokenStore(storage, cipher, node_id="test")
|
||||
store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_at=expires_at,
|
||||
scopes="openid",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
|
||||
def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"),
|
||||
mcp_oauth_http_client=MagicMock(),
|
||||
mcp_oauth_refresh_locks={},
|
||||
mcp_oauth_metadata_cache={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_loop_mgr():
|
||||
cfg: dict[str, Any] = {}
|
||||
mgr = MCPClientManager(cfg)
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop")
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
yield mgr, loop, thread
|
||||
finally:
|
||||
|
||||
async def _drain(m: MCPClientManager) -> None:
|
||||
task = m._user_pool_eviction_task
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
m._user_pool_eviction_task = None
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
def _seed_pool_prompt_map(
|
||||
mgr: MCPClientManager,
|
||||
user_id: str,
|
||||
server_name: str,
|
||||
prefixed_name: str,
|
||||
original_name: str,
|
||||
) -> None:
|
||||
"""Pre-seed ``_user_prompt_map`` so ``_resolve_pool_target_prompt``
|
||||
finds the prefixed name. Production wires this through
|
||||
``_connect_one_pool``; the integration tests seed it directly so the
|
||||
test focuses on the dispatch behaviour after resolution succeeds.
|
||||
"""
|
||||
|
||||
async def _seed() -> None:
|
||||
entry = await mgr._ensure_pool_entry((user_id, server_name))
|
||||
entry.prompts = [
|
||||
{
|
||||
"name": prefixed_name,
|
||||
"original_name": original_name,
|
||||
"server": server_name,
|
||||
"description": "",
|
||||
"arguments": [],
|
||||
}
|
||||
]
|
||||
mgr._rebuild_user_prompt_map(user_id)
|
||||
|
||||
assert mgr._loop is not None
|
||||
asyncio.run_coroutine_threadsafe(_seed(), mgr._loop).result(timeout=5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-1: 401 → refresh → retry → success (prompt path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_401_refresh_and_retry_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Real upstream returns 401 once, then 200. Carrier captures 401,
|
||||
force_refresh=True mints a new bearer, retry returns the prompt
|
||||
messages. Hard invariant 3: breaker counter remains 0.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
messages = mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "everyone"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "user"
|
||||
assert "everyone" in messages[0]["content"]
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}"
|
||||
assert post_headers[0] != post_headers[1], (
|
||||
"retry attached the same bearer as the initial; the dispatcher "
|
||||
"did not pick up the refreshed token."
|
||||
)
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-2: persistent 401 → mcp_consent_required (prompt path) → RuntimeError
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_persistent_401_emits_consent_required(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert payload["error"]["server"] == "pool-srv"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-3: 403 + insufficient_scope → mcp_insufficient_scope (prompt path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_403_insufficient_scope_emits_structured_error(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_insufficient"
|
||||
behaviour["www_authenticate"] = 'Bearer error="insufficient_scope", scope="prompts:read"'
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_insufficient_scope"
|
||||
assert payload["error"]["scopes_required"] == ["prompts:read"]
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) == 1, (
|
||||
f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs"
|
||||
)
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-3b: 403 generic → mcp_prompt_get_forbidden
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_403_generic_forbidden(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_generic"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
# Per the kind="prompt" wiring of `_handle_auth_403`, the
|
||||
# operation-specific code surfaces here rather than the tool path's
|
||||
# generic mcp_tool_call_forbidden.
|
||||
assert payload["error"]["code"] == "mcp_prompt_get_forbidden"
|
||||
assert "scopes_required" not in payload["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-6: breaker isolation — auth failures NEVER trip the breaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_breaker_unaffected_by_auth_failures(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Repeated 401 + refresh-failed cycles leave breaker at 0
|
||||
(hard invariant 3 verified end-to-end for the prompt path)."""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="refresh_failed")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
# Re-seed each iteration: symmetric eviction (Phase 7b) clears
|
||||
# ``_user_prompt_map`` on auth failure so the next dispatch's
|
||||
# resolver would miss without a fresh seed. Production reconnect
|
||||
# repopulates this; the test simulates that out-of-band.
|
||||
for _ in range(10):
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Negative tests — token lookup edge cases (prompt path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_missing_token_emits_consent_required(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, _behaviour = upstream
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="missing")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
|
||||
|
||||
def test_prompt_get_decrypt_failure_emits_token_undecryptable(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, _behaviour = upstream
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="decrypt_failure")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_token_undecryptable_key_unknown"
|
||||
|
||||
|
||||
def test_prompt_get_http_url_emits_url_insecure(
|
||||
running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""An ``http://`` (non-loopback) oauth_user URL must surface
|
||||
``mcp_oauth_url_insecure`` BEFORE the bearer is attached.
|
||||
"""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url="http://example.com/mcp")
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_prompt_map(mgr, "user-1", "pool-srv", "mcp__pool-srv__greet", "greet")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as excinfo,
|
||||
):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
payload = json.loads(str(excinfo.value))
|
||||
assert payload["error"]["code"] == "mcp_oauth_url_insecure"
|
||||
|
||||
|
||||
def test_prompt_get_unknown_name_raises_value_error(
|
||||
running_loop_mgr: Any,
|
||||
) -> None:
|
||||
"""When the prefixed name doesn't resolve to either pool or static,
|
||||
the static-path code raises ``ValueError``. Per-user-first
|
||||
resolution (scope decision 0.1) means user_id-bearing callers still
|
||||
hit this path when their pool catalog doesn't carry the name."""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
with pytest.raises(ValueError, match="Unknown MCP prompt"):
|
||||
mgr.get_prompt_sync(
|
||||
"mcp__nonexistent__missing",
|
||||
None,
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-PR-E2E: real discovery + dispatch in same connect (no _seed_pool_prompt_map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prompt_get_e2e_discovery_then_dispatch_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Drive REAL discovery + dispatch end-to-end through the pool path.
|
||||
|
||||
Mirror of the tool path's
|
||||
``test_integration_pool_reuse_401_refresh_and_retry_succeeds``: skips
|
||||
the ``_seed_pool_prompt_map`` shortcut and lets ``_connect_one_pool``
|
||||
populate ``_user_prompt_map`` from the real ``prompts/list``
|
||||
upstream response. Verifies that the entry's discovered prompts
|
||||
match what the FastMCP fixture advertises AND that
|
||||
``_user_prompt_map[user_id]`` is populated with the prefixed name
|
||||
after dispatch — proving the discovery path actually fired.
|
||||
|
||||
This is the structural gate against a regression where prompt
|
||||
dispatch silently bypasses discovery (e.g., a mis-wired resolver
|
||||
that finds the (server, original) via prefix-parsing alone never
|
||||
populates the per-user catalog).
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "never" # passthrough — discovery + dispatch both succeed
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
# NB: no `_seed_pool_prompt_map` — the resolver finds (server, original)
|
||||
# via the `mcp__{server}__{prompt}` prefix and hands off to
|
||||
# ``_dispatch_pool_prompt_sync``, which lazy-connects via
|
||||
# ``_connect_one_pool``. The connect runs the real ``prompts/list``
|
||||
# against the FastMCP fixture and populates the per-user catalog.
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
messages = mgr.get_prompt_sync(
|
||||
"mcp__pool-srv__greet",
|
||||
{"who": "world"},
|
||||
user_id="user-1",
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "user"
|
||||
assert "world" in messages[0]["content"]
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
# Discovery populated the entry's prompts with both fixtures
|
||||
# (``greet`` and ``summarize``) — proves real ``prompts/list``
|
||||
# ran during the connect, not just the targeted ``prompts/get``.
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is not None
|
||||
assert entry.prompts is not None
|
||||
discovered_names = {p["name"] for p in entry.prompts}
|
||||
assert "mcp__pool-srv__greet" in discovered_names
|
||||
assert "mcp__pool-srv__summarize" in discovered_names
|
||||
|
||||
# ``_rebuild_user_prompt_map`` ran during the connect, populating the
|
||||
# per-user catalog. This is the signal that discovery wired into the
|
||||
# routing tables — without it, a follow-up ``get_prompt_sync`` would
|
||||
# need to re-resolve via prefix parsing every time.
|
||||
user_prompt_map = mgr._user_prompt_map.get("user-1") or {}
|
||||
assert "mcp__pool-srv__greet" in user_prompt_map
|
||||
assert "mcp__pool-srv__summarize" in user_prompt_map
|
||||
@@ -0,0 +1,690 @@
|
||||
"""Phase 7b integration tests — real-transport resource read 401/403/etc.
|
||||
|
||||
Mirror of :mod:`tests.test_mcp_pool_auth_integration` for the resource
|
||||
path (RFC §3.2). Drives through the real ``streamablehttp_client``,
|
||||
real httpx response-hook plumbing, and a real upstream subprocess
|
||||
(``FastMCP`` with a programmable ``BehaviorMiddleware``). Direct
|
||||
``httpx.HTTPStatusError`` injection is forbidden (invariant 14).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import uvicorn
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.mcp_oauth import TokenLookupResult
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
logging.getLogger("mcp").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
class BehaviorMiddleware(BaseHTTPMiddleware):
|
||||
"""Programmable upstream behaviour — see
|
||||
:mod:`tests.test_mcp_pool_auth_integration` for the semantics. This
|
||||
copy serves the resource integration tests.
|
||||
"""
|
||||
|
||||
def __init__(self, app: Any, behaviour: dict[str, Any]) -> None:
|
||||
super().__init__(app)
|
||||
self._behaviour = behaviour
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable[..., Any]) -> Response:
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
|
||||
if request.method == "POST" and "/mcp" in str(request.url):
|
||||
self._behaviour.setdefault("post_auth_headers", []).append(
|
||||
request.headers.get("authorization")
|
||||
)
|
||||
|
||||
mode = self._behaviour.get("mode", "never")
|
||||
if mode == "once_401":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "always_401":
|
||||
return StarletteResponse(
|
||||
"unauthorized",
|
||||
status_code=401,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate", 'Bearer error="invalid_token"'
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_insufficient":
|
||||
if not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get(
|
||||
"www_authenticate",
|
||||
'Bearer error="insufficient_scope", scope="files:read"',
|
||||
)
|
||||
},
|
||||
)
|
||||
elif mode == "once_403_generic" and not self._behaviour.get("_fired"):
|
||||
self._behaviour["_fired"] = True
|
||||
return StarletteResponse(
|
||||
"forbidden",
|
||||
status_code=403,
|
||||
headers={
|
||||
"www-authenticate": self._behaviour.get("www_authenticate", "Bearer realm=mcp")
|
||||
},
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
s = socket.socket()
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
s.close()
|
||||
return port
|
||||
|
||||
|
||||
def _build_server(port: int, behaviour: dict[str, Any]) -> uvicorn.Server:
|
||||
mcp = FastMCP(name="phase7b-resource-target", streamable_http_path="/mcp")
|
||||
|
||||
@mcp.resource("res://hello")
|
||||
def hello() -> str:
|
||||
return "world"
|
||||
|
||||
@mcp.resource("res://json/data")
|
||||
def jdata() -> str:
|
||||
return '{"k": 1}'
|
||||
|
||||
# Echo tool exists so the e2e test can trigger ``_connect_one_pool``
|
||||
# (and the full tool + resource + prompt discovery) via prefix-parsed
|
||||
# ``call_tool_sync`` BEFORE the resource read. The other tests in this
|
||||
# module use ``_seed_pool_resource_map`` and never invoke tools, so
|
||||
# adding the tool is invisible to them.
|
||||
@mcp.tool()
|
||||
async def echo(payload: str = "default") -> str:
|
||||
return f"echoed:{payload}"
|
||||
|
||||
app = mcp.streamable_http_app()
|
||||
app.add_middleware(BehaviorMiddleware, behaviour=behaviour)
|
||||
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning", access_log=False)
|
||||
return uvicorn.Server(config)
|
||||
|
||||
|
||||
def _wait_ready(port: int, timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError(f"upstream at 127.0.0.1:{port} not ready after {timeout}s")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upstream():
|
||||
port = _find_free_port()
|
||||
behaviour: dict[str, Any] = {}
|
||||
server = _build_server(port, behaviour)
|
||||
|
||||
def _run() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(server.serve())
|
||||
|
||||
t = threading.Thread(target=_run, daemon=True, name="phase7b-resource-upstream")
|
||||
t.start()
|
||||
try:
|
||||
_wait_ready(port)
|
||||
yield f"http://127.0.0.1:{port}/mcp", behaviour
|
||||
finally:
|
||||
server.should_exit = True
|
||||
t.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def _seed_oauth_server(
|
||||
storage: SQLiteBackend,
|
||||
*,
|
||||
name: str = "pool-srv",
|
||||
server_id: str = "srv-pool",
|
||||
url: str = "https://mcp.example.com/sse",
|
||||
) -> None:
|
||||
storage.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url=url,
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="client-abc",
|
||||
oauth_scopes="openid",
|
||||
oauth_audience=url,
|
||||
)
|
||||
|
||||
|
||||
def _seed_user_token(
|
||||
storage: SQLiteBackend,
|
||||
cipher: Any,
|
||||
*,
|
||||
user_id: str = "user-1",
|
||||
server_name: str = "pool-srv",
|
||||
expires_in_seconds: int = 3600,
|
||||
access_token: str = "access-aaa",
|
||||
refresh_token: str | None = "refresh-rrr",
|
||||
) -> None:
|
||||
expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
)
|
||||
store = MCPTokenStore(storage, cipher, node_id="test")
|
||||
store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_at=expires_at,
|
||||
scopes="openid",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
|
||||
def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"),
|
||||
mcp_oauth_http_client=MagicMock(),
|
||||
mcp_oauth_refresh_locks={},
|
||||
mcp_oauth_metadata_cache={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_loop_mgr():
|
||||
cfg: dict[str, Any] = {}
|
||||
mgr = MCPClientManager(cfg)
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop")
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
yield mgr, loop, thread
|
||||
finally:
|
||||
|
||||
async def _drain(m: MCPClientManager) -> None:
|
||||
task = m._user_pool_eviction_task
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
m._user_pool_eviction_task = None
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
def _seed_pool_resource_map(
|
||||
mgr: MCPClientManager, user_id: str, server_name: str, uri: str
|
||||
) -> None:
|
||||
"""Pre-seed ``_user_resource_map`` so ``_resolve_pool_target_resource``
|
||||
finds the URI. Production wires this through ``_connect_one_pool``;
|
||||
the integration tests seed it directly so the test focuses on the
|
||||
dispatch behaviour after resolution succeeds.
|
||||
"""
|
||||
|
||||
async def _seed() -> None:
|
||||
entry = await mgr._ensure_pool_entry((user_id, server_name))
|
||||
entry.resources = [
|
||||
{
|
||||
"uri": uri,
|
||||
"name": "",
|
||||
"description": "",
|
||||
"mimeType": "",
|
||||
"server": server_name,
|
||||
}
|
||||
]
|
||||
mgr._rebuild_user_resource_map(user_id)
|
||||
|
||||
assert mgr._loop is not None
|
||||
asyncio.run_coroutine_threadsafe(_seed(), mgr._loop).result(timeout=5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-1: 401 → refresh → retry → success (resource path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_401_refresh_and_retry_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Real upstream returns 401 once, then 200. Carrier captures 401,
|
||||
force_refresh=True mints a new bearer, retry returns the resource.
|
||||
Hard invariant 3: breaker counter remains 0.
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
|
||||
assert result == "world"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) >= 2, f"expected >=2 POSTs; got {len(post_headers)}"
|
||||
assert post_headers[0] != post_headers[1], (
|
||||
"retry attached the same bearer as the initial; the dispatcher "
|
||||
"did not pick up the refreshed token."
|
||||
)
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-2: persistent 401 → mcp_consent_required (resource path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_persistent_401_emits_consent_required(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="token", token="refreshed-bearer")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert payload["error"]["server"] == "pool-srv"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-3: 403 + insufficient_scope → mcp_insufficient_scope (resource path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_403_insufficient_scope_emits_structured_error(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_insufficient"
|
||||
behaviour["www_authenticate"] = 'Bearer error="insufficient_scope", scope="files:read"'
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_insufficient_scope"
|
||||
assert payload["error"]["scopes_required"] == ["files:read"]
|
||||
post_headers = behaviour.get("post_auth_headers", [])
|
||||
assert len(post_headers) == 1, (
|
||||
f"403 must NOT trigger a retry; observed {len(post_headers)} POSTs"
|
||||
)
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-3b: 403 generic → mcp_resource_read_forbidden
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_403_generic_forbidden(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "once_403_generic"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
# Per the kind="resource" wiring of `_handle_auth_403`, the
|
||||
# operation-specific code surfaces here rather than the tool path's
|
||||
# generic mcp_tool_call_forbidden.
|
||||
assert payload["error"]["code"] == "mcp_resource_read_forbidden"
|
||||
assert "scopes_required" not in payload["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-6: breaker isolation — auth failures NEVER trip the breaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_breaker_unaffected_by_auth_failures(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Repeated 401 + refresh-failed cycles leave breaker at 0
|
||||
(hard invariant 3 verified end-to-end for the resource path)."""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "always_401"
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
|
||||
if kwargs.get("force_refresh"):
|
||||
return TokenLookupResult(kind="refresh_failed")
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
# Re-seed each iteration: symmetric eviction (Phase 7b) clears
|
||||
# ``_user_resource_map`` on auth failure so the next dispatch's
|
||||
# resolver would miss without a fresh seed. Production reconnect
|
||||
# repopulates this; the test simulates that out-of-band.
|
||||
for _ in range(10):
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Negative tests — token lookup edge cases (resource path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_missing_token_emits_consent_required(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, _behaviour = upstream
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="missing")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
|
||||
|
||||
def test_resource_read_decrypt_failure_emits_token_undecryptable(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
url, _behaviour = upstream
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="decrypt_failure")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_token_undecryptable_key_unknown"
|
||||
|
||||
|
||||
def test_resource_read_http_url_emits_url_insecure(
|
||||
running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""An ``http://`` (non-loopback) oauth_user URL must surface
|
||||
``mcp_oauth_url_insecure`` BEFORE the bearer is attached.
|
||||
"""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url="http://example.com/mcp")
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
),
|
||||
pytest.raises(RuntimeError) as exc_info,
|
||||
):
|
||||
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=5)
|
||||
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_oauth_url_insecure"
|
||||
|
||||
|
||||
def test_resource_read_unknown_uri_raises_value_error(
|
||||
running_loop_mgr: Any,
|
||||
) -> None:
|
||||
"""When the URI doesn't resolve to either pool or static, the
|
||||
static-path code raises ``ValueError``. Per-user-first resolution
|
||||
(scope decision 0.1) means user_id-bearing callers still hit this
|
||||
path when their pool catalog doesn't carry the URI."""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
with pytest.raises(ValueError, match="Unknown MCP resource"):
|
||||
mgr.read_resource_sync("res://nonexistent", user_id="user-1", timeout=5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# I-RP-E2E: real discovery + dispatch in same connect (no _seed_pool_resource_map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resource_read_e2e_discovery_then_dispatch_succeeds(
|
||||
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Drive REAL discovery + dispatch end-to-end through the pool path.
|
||||
|
||||
Mirror of the tool path's
|
||||
``test_integration_pool_reuse_401_refresh_and_retry_succeeds``: skips
|
||||
the ``_seed_pool_resource_map`` shortcut and lets ``_connect_one_pool``
|
||||
populate ``_user_resource_map`` from the real ``resources/list``
|
||||
upstream response. Verifies that the entry's discovered resources
|
||||
match what the FastMCP fixture advertises AND that
|
||||
``_user_resource_map[user_id]`` is populated with the URI(s) after
|
||||
discovery — proving the discovery path actually fired.
|
||||
|
||||
Resource URIs do NOT carry a server-name prefix (unlike tools and
|
||||
prompts), so the resource resolver cannot derive (server, uri) by
|
||||
parsing alone. The test triggers the connect via a prefix-parsed
|
||||
``call_tool_sync`` first (which runs the full
|
||||
tools+resources+prompts discovery against the FastMCP fixture),
|
||||
then drives ``read_resource_sync`` against a URI that the
|
||||
upstream advertised — proving that real discovery wired the URI
|
||||
into the per-user catalog.
|
||||
|
||||
Structural gate against a regression where resource discovery is
|
||||
silently skipped (e.g., a capability-gating bug that drops the
|
||||
``resources/list`` call but keeps the connect succeeding).
|
||||
"""
|
||||
url, behaviour = upstream
|
||||
behaviour["mode"] = "never" # passthrough — discovery + dispatch both succeed
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url=url)
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
|
||||
# NB: no `_seed_pool_resource_map` — the connect runs the real
|
||||
# ``resources/list`` against the FastMCP fixture and populates the
|
||||
# per-user catalog. The tool call below triggers that connect because
|
||||
# ``_resolve_pool_target`` derives (server, original) from the
|
||||
# ``mcp__pool-srv__echo`` prefix and lazy-connects via
|
||||
# ``_connect_one_pool``.
|
||||
|
||||
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
|
||||
return TokenLookupResult(kind="token", token="access-aaa")
|
||||
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.get_user_access_token_classified",
|
||||
side_effect=_fake_classified,
|
||||
):
|
||||
# Step 1: trigger the connect via prefix-parsed tool dispatch.
|
||||
# Discovery (tools + resources + prompts) populates the per-user
|
||||
# catalogs.
|
||||
tool_result = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__echo", {"payload": "ignite"}, user_id="user-1", timeout=15
|
||||
)
|
||||
assert "echoed:ignite" in tool_result
|
||||
|
||||
# Step 2: now that discovery has populated ``_user_resource_map``,
|
||||
# the resource resolver finds ``res://hello`` and dispatches the
|
||||
# read on the SAME pool entry / session.
|
||||
result = mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
|
||||
|
||||
assert result == "world"
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 0
|
||||
|
||||
# Discovery populated the entry's resources with both fixtures
|
||||
# (``res://hello`` and ``res://json/data``) — proves real
|
||||
# ``resources/list`` ran during the connect, not just the targeted
|
||||
# ``resources/read``.
|
||||
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
|
||||
assert entry.session is not None
|
||||
assert entry.resources is not None
|
||||
discovered_uris = {r["uri"] for r in entry.resources if not r.get("template")}
|
||||
assert "res://hello" in discovered_uris
|
||||
assert "res://json/data" in discovered_uris
|
||||
|
||||
# ``_rebuild_user_resource_map`` ran during the connect, populating
|
||||
# the per-user catalog. This is the signal that discovery wired into
|
||||
# the routing tables — without it, ``read_resource_sync`` would have
|
||||
# raised ValueError because the resolver had no entry for the URI.
|
||||
user_resource_map = mgr._user_resource_map.get("user-1") or {}
|
||||
assert "res://hello" in user_resource_map
|
||||
assert "res://json/data" in user_resource_map
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Tests for ``MCPTokenStore`` ciphertext-aware CRUD.
|
||||
|
||||
Phase 3 of the OAuth-MCP RFC: validates the encrypt/decrypt boundary
|
||||
between :class:`MCPTokenStore` and the storage protocol's ciphertext-only
|
||||
columns. Exercises the row-not-deleted-on-decrypt-failure invariant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from turnstone.core.mcp_crypto import (
|
||||
MCPTokenCipher,
|
||||
MCPTokenCipherConfig,
|
||||
MCPTokenDecryptError,
|
||||
MCPTokenStore,
|
||||
)
|
||||
|
||||
|
||||
def _make_cipher() -> MCPTokenCipher:
|
||||
raw = base64.urlsafe_b64decode(Fernet.generate_key())
|
||||
return MCPTokenCipher(MCPTokenCipherConfig(keys=(raw,)))
|
||||
|
||||
|
||||
def _make_store(backend, *, audit: bool = False) -> tuple[MCPTokenStore, MCPTokenCipher]:
|
||||
cipher = _make_cipher()
|
||||
store = MCPTokenStore(
|
||||
backend,
|
||||
cipher,
|
||||
node_id="test-node",
|
||||
audit_storage=backend if audit else None,
|
||||
)
|
||||
return store, cipher
|
||||
|
||||
|
||||
def _seed_server(backend, *, server_id: str = "srv-id-1", name: str = "srv-a") -> str:
|
||||
backend.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
)
|
||||
return server_id
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User-token CRUD
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestUserTokenCRUD:
|
||||
def test_create_and_get_round_trip(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
store.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="access-aaa",
|
||||
refresh_token="refresh-bbb",
|
||||
expires_at="2026-05-04T12:00:00",
|
||||
scopes="openid profile",
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
plain = store.get_user_token("u1", "srv-a")
|
||||
assert plain is not None
|
||||
assert plain["user_id"] == "u1"
|
||||
assert plain["server_name"] == "srv-a"
|
||||
assert plain["access_token"] == "access-aaa"
|
||||
assert plain["refresh_token"] == "refresh-bbb"
|
||||
assert plain["scopes"] == "openid profile"
|
||||
assert plain["audience"] == "https://mcp.example.com"
|
||||
|
||||
def test_create_with_no_refresh_token(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
store.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="access-only",
|
||||
refresh_token=None,
|
||||
expires_at=None,
|
||||
scopes=None,
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
plain = store.get_user_token("u1", "srv-a")
|
||||
assert plain is not None
|
||||
assert plain["access_token"] == "access-only"
|
||||
assert plain["refresh_token"] is None
|
||||
|
||||
def test_get_missing_returns_none(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
assert store.get_user_token("nobody", "srv-a") is None
|
||||
|
||||
def test_update_after_refresh(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
store.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="old-access",
|
||||
refresh_token="old-refresh",
|
||||
expires_at="2026-05-04T12:00:00",
|
||||
scopes="openid",
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
ok = store.update_user_token_after_refresh(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="new-access",
|
||||
refresh_token="new-refresh",
|
||||
expires_at="2026-05-04T13:00:00",
|
||||
)
|
||||
assert ok is True
|
||||
plain = store.get_user_token("u1", "srv-a")
|
||||
assert plain is not None
|
||||
assert plain["access_token"] == "new-access"
|
||||
assert plain["refresh_token"] == "new-refresh"
|
||||
assert plain["expires_at"] == "2026-05-04T13:00:00"
|
||||
# Preserved columns:
|
||||
assert plain["scopes"] == "openid"
|
||||
assert plain["as_issuer"] == "https://auth.example.com"
|
||||
# last_refreshed got stamped:
|
||||
assert plain["last_refreshed"] is not None
|
||||
|
||||
def test_update_after_refresh_missing_row_returns_false(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
ok = store.update_user_token_after_refresh(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="x",
|
||||
refresh_token=None,
|
||||
expires_at=None,
|
||||
)
|
||||
assert ok is False
|
||||
|
||||
def test_delete(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
store.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="a",
|
||||
refresh_token=None,
|
||||
expires_at=None,
|
||||
scopes=None,
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
assert store.delete_user_token("u1", "srv-a") is True
|
||||
assert store.get_user_token("u1", "srv-a") is None
|
||||
# Idempotent: deleting again returns False.
|
||||
assert store.delete_user_token("u1", "srv-a") is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Client-secret writer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClientSecretWriter:
|
||||
def test_set_oauth_client_secret_round_trip(self, backend) -> None:
|
||||
store, cipher = _make_store(backend)
|
||||
server_id = _seed_server(backend)
|
||||
ok = store.set_oauth_client_secret(server_id, "plaintext-secret")
|
||||
assert ok is True
|
||||
# Read raw via get_mcp_server: ciphertext != plaintext, decrypts back.
|
||||
raw = backend.get_mcp_server(server_id)
|
||||
assert raw is not None
|
||||
ct = raw["oauth_client_secret_ct"]
|
||||
assert isinstance(ct, (bytes, bytearray, memoryview))
|
||||
ct_bytes = bytes(ct)
|
||||
assert ct_bytes != b"plaintext-secret"
|
||||
assert cipher.decrypt(ct_bytes) == b"plaintext-secret"
|
||||
|
||||
def test_set_oauth_client_secret_clear_with_none(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
server_id = _seed_server(backend)
|
||||
store.set_oauth_client_secret(server_id, "x")
|
||||
assert store.set_oauth_client_secret(server_id, None) is True
|
||||
raw = backend.get_mcp_server(server_id)
|
||||
assert raw is not None
|
||||
assert raw["oauth_client_secret_ct"] is None
|
||||
|
||||
def test_set_oauth_client_secret_missing_server_returns_false(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
ok = store.set_oauth_client_secret("does-not-exist", "x")
|
||||
assert ok is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Decrypt failure: row preservation invariant
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDecryptFailureInvariant:
|
||||
def test_get_user_token_with_wrong_key_raises_decrypt_error(self, backend) -> None:
|
||||
"""CRITICAL: when no installed key can decrypt a stored row,
|
||||
``get_user_token`` MUST NOT auto-delete the row. The row is
|
||||
still valid; this node just doesn't have the right key.
|
||||
"""
|
||||
# Write under cipher A.
|
||||
store_a, _cipher_a = _make_store(backend)
|
||||
store_a.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="secret-access",
|
||||
refresh_token="secret-refresh",
|
||||
expires_at="2026-05-04T12:00:00",
|
||||
scopes="openid",
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
raw_before = backend.get_mcp_user_token("u1", "srv-a")
|
||||
assert raw_before is not None
|
||||
ct_before = bytes(raw_before["access_token_ct"])
|
||||
|
||||
# Read under cipher B (different key).
|
||||
store_b, cipher_b = _make_store(backend)
|
||||
with pytest.raises(MCPTokenDecryptError) as exc_info:
|
||||
store_b.get_user_token("u1", "srv-a")
|
||||
# The exception carries the keys we tried — useful for audit.
|
||||
assert exc_info.value.key_fingerprints_attempted == cipher_b.key_fingerprints
|
||||
|
||||
# Row MUST still exist with ciphertext intact.
|
||||
raw_after = backend.get_mcp_user_token("u1", "srv-a")
|
||||
assert raw_after is not None
|
||||
assert bytes(raw_after["access_token_ct"]) == ct_before
|
||||
|
||||
def test_decrypt_failure_emits_audit_when_configured(self, backend) -> None:
|
||||
"""When ``audit_storage`` is set, decrypt failures emit a
|
||||
``mcp_server.oauth.token_decrypt_failure`` audit event."""
|
||||
store_a, _ = _make_store(backend)
|
||||
store_a.create_user_token(
|
||||
"u1",
|
||||
"srv-a",
|
||||
access_token="x",
|
||||
refresh_token=None,
|
||||
expires_at=None,
|
||||
scopes=None,
|
||||
as_issuer="https://a",
|
||||
audience="https://m",
|
||||
)
|
||||
|
||||
store_b, cipher_b = _make_store(backend, audit=True)
|
||||
with pytest.raises(MCPTokenDecryptError):
|
||||
store_b.get_user_token("u1", "srv-a")
|
||||
|
||||
events = backend.list_audit_events(limit=10)
|
||||
actions = {ev.get("action") for ev in events}
|
||||
assert "mcp_server.oauth.token_decrypt_failure" in actions
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Client-secret reader — q-9
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClientSecretReader:
|
||||
def test_get_oauth_client_secret_returns_none_when_row_absent(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
assert store.get_oauth_client_secret("does-not-exist") is None
|
||||
|
||||
def test_get_oauth_client_secret_returns_none_when_column_null(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
server_id = _seed_server(backend)
|
||||
# No set_oauth_client_secret call — column stays NULL.
|
||||
assert store.get_oauth_client_secret(server_id) is None
|
||||
|
||||
def test_get_oauth_client_secret_round_trip(self, backend) -> None:
|
||||
store, _ = _make_store(backend)
|
||||
server_id = _seed_server(backend)
|
||||
store.set_oauth_client_secret(server_id, "shhh-its-secret")
|
||||
assert store.get_oauth_client_secret(server_id) == "shhh-its-secret"
|
||||
|
||||
def test_get_oauth_client_secret_raises_on_key_mismatch(self, backend) -> None:
|
||||
store_a, _ = _make_store(backend)
|
||||
server_id = _seed_server(backend)
|
||||
store_a.set_oauth_client_secret(server_id, "secret-under-key-a")
|
||||
|
||||
# Cipher B has a different key — decrypt fails loudly.
|
||||
store_b, _ = _make_store(backend)
|
||||
with pytest.raises(MCPTokenDecryptError):
|
||||
store_b.get_oauth_client_secret(server_id)
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Tests for ``MCPTokenStore.list_user_token_metadata``.
|
||||
|
||||
Validates the non-secret projection used by the settings UI: ciphertext
|
||||
columns are stripped, ordering is preserved, and the empty case returns
|
||||
``[]``. Decrypt is intentionally skipped — the list view must never need
|
||||
the access/refresh secrets.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import sqlalchemy as sa
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from turnstone.core.mcp_crypto import (
|
||||
MCPTokenCipher,
|
||||
MCPTokenCipherConfig,
|
||||
MCPTokenStore,
|
||||
)
|
||||
|
||||
|
||||
def _make_cipher() -> MCPTokenCipher:
|
||||
raw = base64.urlsafe_b64decode(Fernet.generate_key())
|
||||
return MCPTokenCipher(MCPTokenCipherConfig(keys=(raw,)))
|
||||
|
||||
|
||||
def _make_store(backend) -> MCPTokenStore:
|
||||
return MCPTokenStore(backend, _make_cipher(), node_id="test-node")
|
||||
|
||||
|
||||
def _seed_token(
|
||||
store: MCPTokenStore,
|
||||
backend,
|
||||
*,
|
||||
user_id: str,
|
||||
server_name: str,
|
||||
created: str,
|
||||
) -> None:
|
||||
"""Create a token via the store and backdate ``created`` for ordering."""
|
||||
store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token="access-secret",
|
||||
refresh_token="refresh-secret",
|
||||
expires_at="2026-05-04T12:00:00",
|
||||
scopes="openid profile",
|
||||
as_issuer="https://auth.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"UPDATE mcp_user_tokens SET created = :created "
|
||||
"WHERE user_id = :uid AND server_name = :sn"
|
||||
),
|
||||
{"created": created, "uid": user_id, "sn": server_name},
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
class TestListUserTokenMetadata:
|
||||
def test_list_user_token_metadata_returns_non_secret_fields_only(self, backend) -> None:
|
||||
store = _make_store(backend)
|
||||
_seed_token(
|
||||
store, backend, user_id="u1", server_name="srv-a", created="2026-05-01T00:00:00"
|
||||
)
|
||||
rows = store.list_user_token_metadata("u1")
|
||||
assert len(rows) == 1
|
||||
meta = rows[0]
|
||||
# Secrets MUST be absent.
|
||||
assert "access_token" not in meta
|
||||
assert "refresh_token" not in meta
|
||||
assert "access_token_ct" not in meta
|
||||
assert "refresh_token_ct" not in meta
|
||||
# Non-secret columns surface verbatim.
|
||||
assert meta["user_id"] == "u1"
|
||||
assert meta["server_name"] == "srv-a"
|
||||
assert meta["scopes"] == "openid profile"
|
||||
assert meta["as_issuer"] == "https://auth.example.com"
|
||||
assert meta["audience"] == "https://mcp.example.com"
|
||||
assert meta["expires_at"] == "2026-05-04T12:00:00"
|
||||
assert meta["created"] == "2026-05-01T00:00:00"
|
||||
assert meta["last_refreshed"] is None
|
||||
|
||||
def test_list_user_token_metadata_empty(self, backend) -> None:
|
||||
store = _make_store(backend)
|
||||
assert store.list_user_token_metadata("nobody") == []
|
||||
|
||||
def test_list_user_token_metadata_preserves_creation_order(self, backend) -> None:
|
||||
store = _make_store(backend)
|
||||
_seed_token(
|
||||
store, backend, user_id="u1", server_name="srv-c", created="2026-05-03T00:00:00"
|
||||
)
|
||||
_seed_token(
|
||||
store, backend, user_id="u1", server_name="srv-a", created="2026-05-01T00:00:00"
|
||||
)
|
||||
_seed_token(
|
||||
store, backend, user_id="u1", server_name="srv-b", created="2026-05-02T00:00:00"
|
||||
)
|
||||
rows = store.list_user_token_metadata("u1")
|
||||
assert [r["server_name"] for r in rows] == ["srv-a", "srv-b", "srv-c"]
|
||||
assert [r["created"] for r in rows] == [
|
||||
"2026-05-01T00:00:00",
|
||||
"2026-05-02T00:00:00",
|
||||
"2026-05-03T00:00:00",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,972 @@
|
||||
"""Tests for the per-(user, server) MCP session pool.
|
||||
|
||||
Covers Phase 5 of the OAuth-MCP rollout: pool data structures,
|
||||
``_ensure_pool_entry`` lazy allocation, ``_connect_one_pool`` plumbing,
|
||||
the dispatch state machine in ``_dispatch_pool``, idle / LRU eviction,
|
||||
failure classification, and ``user_id`` thread-through.
|
||||
|
||||
The static path (``auth_type ∈ {none, static}``) MUST stay
|
||||
byte-identical — see ``test_mcp_client.py``'s
|
||||
``test_reconnect_preserves_static_state_identity``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from contextlib import AsyncExitStack
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import make_mcp_token_cipher
|
||||
from turnstone.core.mcp_client import MCPClientManager, PoolEntryState
|
||||
from turnstone.core.mcp_crypto import MCPTokenStore
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures and helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
"""A fresh SQLite backend per test (not the shared singleton)."""
|
||||
return SQLiteBackend(str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def _seed_oauth_server(
|
||||
storage: SQLiteBackend,
|
||||
*,
|
||||
name: str = "pool-srv",
|
||||
server_id: str = "srv-pool",
|
||||
url: str = "https://mcp.example.com/sse",
|
||||
) -> None:
|
||||
storage.create_mcp_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url=url,
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="client-abc",
|
||||
oauth_scopes="openid",
|
||||
oauth_audience=url,
|
||||
)
|
||||
|
||||
|
||||
def _seed_user_token(
|
||||
storage: SQLiteBackend,
|
||||
cipher: Any,
|
||||
*,
|
||||
user_id: str = "user-1",
|
||||
server_name: str = "pool-srv",
|
||||
expires_in_seconds: int = 3600,
|
||||
access_token: str = "access-aaa",
|
||||
) -> None:
|
||||
expires_at = (datetime.now(UTC) + timedelta(seconds=expires_in_seconds)).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
)
|
||||
store = MCPTokenStore(storage, cipher, node_id="test")
|
||||
store.create_user_token(
|
||||
user_id,
|
||||
server_name,
|
||||
access_token=access_token,
|
||||
refresh_token="refresh-rrr",
|
||||
expires_at=expires_at,
|
||||
scopes="openid",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
|
||||
def _make_app_state(storage: SQLiteBackend, *, cipher: Any) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
auth_storage=storage,
|
||||
mcp_token_store=MCPTokenStore(storage, cipher, node_id="test"),
|
||||
mcp_oauth_http_client=MagicMock(),
|
||||
mcp_oauth_refresh_locks={},
|
||||
mcp_oauth_metadata_cache={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_loop_mgr():
|
||||
"""Background-loop fixture matching the static-path test convention.
|
||||
|
||||
Tests that need a wired-up app_state assign it via ``mgr.set_app_state``.
|
||||
"""
|
||||
cfg: dict[str, Any] = {}
|
||||
mgr = MCPClientManager(cfg)
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True, name="mcp-pool-test-loop")
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
yield mgr, loop, thread
|
||||
finally:
|
||||
# Drain the eviction task before stopping the loop so its log/stream
|
||||
# handlers don't fire after pytest has torn its handlers down. Mirrors
|
||||
# the production ``shutdown()`` shape.
|
||||
async def _drain(m: MCPClientManager) -> None:
|
||||
task = m._user_pool_eviction_task
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
m._user_pool_eviction_task = None
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_drain(mgr), loop).result(timeout=2)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
def _run_on_loop(loop: asyncio.AbstractEventLoop, coro: Any) -> Any:
|
||||
"""Submit *coro* to *loop*, wait for the result with a 5s timeout."""
|
||||
fut = asyncio.run_coroutine_threadsafe(coro, loop)
|
||||
return fut.result(timeout=5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pool data structures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPoolDataStructures:
|
||||
"""``_user_pool_entries``, ``_user_pool_locks``, eviction-task state."""
|
||||
|
||||
def test_pool_state_starts_empty(self) -> None:
|
||||
mgr = MCPClientManager({})
|
||||
assert mgr._user_pool_entries == {}
|
||||
assert mgr._user_pool_last_used == {}
|
||||
assert mgr._user_pool_locks == {}
|
||||
assert mgr._user_pool_eviction_task is None
|
||||
|
||||
def test_set_app_state_persists(self) -> None:
|
||||
mgr = MCPClientManager({})
|
||||
sentinel = SimpleNamespace(token_store=object())
|
||||
mgr.set_app_state(sentinel)
|
||||
assert mgr._app_state is sentinel
|
||||
|
||||
def test_ensure_pool_entry_allocates_lock_on_loop(self, running_loop_mgr) -> None:
|
||||
"""``asyncio.Lock`` MUST be created on the mcp-loop (RFC §2.0 #2)."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
key = ("user-A", "pool-srv")
|
||||
entry = _run_on_loop(loop, mgr._ensure_pool_entry(key))
|
||||
assert isinstance(entry, PoolEntryState)
|
||||
assert entry.key == key
|
||||
assert isinstance(entry.open_lock, asyncio.Lock)
|
||||
# Calling again returns the same entry / lock object.
|
||||
entry2 = _run_on_loop(loop, mgr._ensure_pool_entry(key))
|
||||
assert entry2 is entry
|
||||
assert entry2.open_lock is entry.open_lock
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lazy connect (`_connect_one_pool`)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _AsyncCM:
|
||||
"""Awaitable async context manager that returns ``value`` from __aenter__."""
|
||||
|
||||
def __init__(self, value: Any) -> None:
|
||||
self._value = value
|
||||
|
||||
async def __aenter__(self) -> Any:
|
||||
return self._value
|
||||
|
||||
async def __aexit__(self, *exc: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class TestLazyConnect:
|
||||
def test_connect_pool_injects_authorization_header(self, running_loop_mgr) -> None:
|
||||
from unittest.mock import patch
|
||||
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
|
||||
observed_kwargs: dict[str, Any] = {}
|
||||
|
||||
async def _probe(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
fake_session = MagicMock()
|
||||
fake_session.initialize = AsyncMock(return_value=None)
|
||||
# Phase 7b: ``_connect_one_pool`` discovers tools, resources,
|
||||
# and prompts after ``initialize()`` returns (resources/prompts
|
||||
# capability-gated). The capability stub returns a tools-only
|
||||
# advertisement so the test can keep its narrow focus on the
|
||||
# bearer-injection contract; resources/prompts paths are
|
||||
# exercised by the real-transport tests in
|
||||
# ``tests/test_mcp_user_catalog.py``.
|
||||
fake_caps = MagicMock()
|
||||
fake_caps.resources = None
|
||||
fake_caps.prompts = None
|
||||
fake_session.get_server_capabilities = MagicMock(return_value=fake_caps)
|
||||
fake_session.list_tools = AsyncMock(return_value=MagicMock(tools=[]))
|
||||
|
||||
def _stream_factory(*, url: str, headers: dict[str, str]) -> _AsyncCM:
|
||||
observed_kwargs["url"] = url
|
||||
observed_kwargs["headers"] = dict(headers)
|
||||
return _AsyncCM((AsyncMock(), AsyncMock(), lambda: None))
|
||||
|
||||
with (
|
||||
patch("turnstone.core.mcp_client.streamablehttp_client", side_effect=_stream_factory),
|
||||
patch.object(mgr, "_tcp_probe", side_effect=_probe),
|
||||
patch("turnstone.core.mcp_client.ClientSession", return_value=_AsyncCM(fake_session)),
|
||||
):
|
||||
cfg = {
|
||||
"type": "streamable-http",
|
||||
"url": "https://mcp.example.com/sse",
|
||||
"headers": {},
|
||||
}
|
||||
entry = _run_on_loop(
|
||||
loop,
|
||||
mgr._connect_one_pool(("user-1", "pool-srv"), cfg, "access-aaa"),
|
||||
)
|
||||
|
||||
assert entry.session is fake_session
|
||||
assert observed_kwargs["headers"]["Authorization"] == "Bearer access-aaa"
|
||||
|
||||
def test_connect_pool_rejects_non_http_transport(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cfg = {"type": "stdio", "command": "echo"}
|
||||
with pytest.raises(RuntimeError, match="streamable-http"):
|
||||
_run_on_loop(
|
||||
loop,
|
||||
mgr._connect_one_pool(("user-1", "pool-srv"), cfg, "access-aaa"),
|
||||
)
|
||||
|
||||
def test_pool_path_does_not_touch_static_servers(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
# Pre-seed a static-path entry so accidental writes are observable.
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
sentinel = StaticServerState(name="static-srv", session=MagicMock())
|
||||
mgr._static_servers["static-srv"] = sentinel
|
||||
|
||||
async def _seed_pool() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-1", "pool-srv"))
|
||||
entry.session = MagicMock()
|
||||
entry.last_used = time.monotonic()
|
||||
|
||||
_run_on_loop(loop, _seed_pool())
|
||||
# Pool side has its own state; the static dict is untouched.
|
||||
assert mgr._static_servers["static-srv"] is sentinel
|
||||
assert mgr._user_pool_entries[("user-1", "pool-srv")].session is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Eviction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEviction:
|
||||
def test_idle_eviction_closes_stale_entries(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
mgr._user_pool_idle_ttl_s = 0.0 # everything is stale
|
||||
|
||||
async def _seed() -> list[PoolEntryState]:
|
||||
entries = []
|
||||
for i in range(3):
|
||||
entry = await mgr._ensure_pool_entry((f"u{i}", "pool-srv"))
|
||||
entry.session = MagicMock()
|
||||
entries.append(entry)
|
||||
return entries
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
|
||||
async def _evict() -> None:
|
||||
await mgr._evict_idle_pool_entries()
|
||||
|
||||
_run_on_loop(loop, _evict())
|
||||
assert mgr._user_pool_entries == {}
|
||||
|
||||
def test_eviction_skips_locked_entries(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
mgr._user_pool_idle_ttl_s = 0.0
|
||||
|
||||
async def _seed_and_lock() -> tuple[asyncio.Lock, asyncio.Event]:
|
||||
entry = await mgr._ensure_pool_entry(("u-busy", "pool-srv"))
|
||||
entry.session = MagicMock()
|
||||
held = asyncio.Event()
|
||||
|
||||
async def _hold() -> None:
|
||||
async with entry.open_lock:
|
||||
held.set()
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
asyncio.create_task(_hold())
|
||||
await held.wait()
|
||||
return entry.open_lock, held
|
||||
|
||||
_run_on_loop(loop, _seed_and_lock())
|
||||
|
||||
async def _evict() -> None:
|
||||
await mgr._evict_idle_pool_entries()
|
||||
|
||||
_run_on_loop(loop, _evict())
|
||||
# Entry survives because eviction skipped the locked key.
|
||||
assert ("u-busy", "pool-srv") in mgr._user_pool_entries
|
||||
|
||||
def test_lru_cap_evicts_oldest(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
mgr._user_pool_idle_ttl_s = 999_999.0 # TTL effectively disabled
|
||||
mgr._user_pool_lru_max = 2
|
||||
|
||||
async def _seed() -> None:
|
||||
base = time.monotonic()
|
||||
for i in range(5):
|
||||
key = (f"u{i}", "pool-srv")
|
||||
entry = await mgr._ensure_pool_entry(key)
|
||||
entry.session = MagicMock()
|
||||
# Recent timestamps so TTL doesn't fire — only LRU should.
|
||||
entry.last_used = base + i
|
||||
mgr._user_pool_last_used[key] = base + i
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
|
||||
async def _evict() -> None:
|
||||
await mgr._evict_idle_pool_entries()
|
||||
|
||||
_run_on_loop(loop, _evict())
|
||||
assert len(mgr._user_pool_entries) <= 2
|
||||
# The two newest survive (u3, u4).
|
||||
assert ("u4", "pool-srv") in mgr._user_pool_entries
|
||||
assert ("u3", "pool-srv") in mgr._user_pool_entries
|
||||
|
||||
def test_eviction_resilient_to_close_errors(self, running_loop_mgr) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
mgr._user_pool_idle_ttl_s = 0.0
|
||||
|
||||
broken_stack = MagicMock(spec=AsyncExitStack)
|
||||
broken_stack.aclose = AsyncMock(side_effect=RuntimeError("close failed"))
|
||||
|
||||
async def _seed() -> None:
|
||||
for i in range(2):
|
||||
entry = await mgr._ensure_pool_entry((f"u{i}", "pool-srv"))
|
||||
entry.session = MagicMock()
|
||||
entry.stack = broken_stack
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
|
||||
async def _evict() -> None:
|
||||
await mgr._evict_idle_pool_entries()
|
||||
|
||||
# Eviction must not raise even if close fails.
|
||||
_run_on_loop(loop, _evict())
|
||||
# All entries removed from the dict regardless.
|
||||
assert mgr._user_pool_entries == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dispatch state machine
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDispatchStateMachine:
|
||||
"""One row per state in the §1.5 / RFC §6 state machine."""
|
||||
|
||||
def _wire_pool(
|
||||
self, mgr: MCPClientManager, storage: SQLiteBackend, cipher: Any
|
||||
) -> SimpleNamespace:
|
||||
mgr.set_storage(storage)
|
||||
state = _make_app_state(storage, cipher=cipher)
|
||||
mgr.set_app_state(state)
|
||||
return state
|
||||
|
||||
def test_no_token_emits_consent_required(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
self._wire_pool(mgr, storage, cipher)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
assert payload["error"]["server"] == "pool-srv"
|
||||
|
||||
def test_decrypt_failure_does_not_emit_consent(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
from turnstone.core.mcp_crypto import MCPTokenDecryptError
|
||||
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
_seed_user_token(storage, cipher)
|
||||
state = self._wire_pool(mgr, storage, cipher)
|
||||
|
||||
def _raise(*args, **kwargs):
|
||||
raise MCPTokenDecryptError(
|
||||
"no installed key can decrypt",
|
||||
key_fingerprints_attempted=("aabbccdd",),
|
||||
)
|
||||
|
||||
state.mcp_token_store.get_user_token = _raise
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_token_undecryptable_key_unknown"
|
||||
# Operator fingerprints stay server-side (audit log + structured log);
|
||||
# the agent-facing payload must NOT carry them onward to the LLM
|
||||
# provider.
|
||||
assert "key_fingerprints_attempted" not in payload["error"]
|
||||
|
||||
def test_refresh_failure_emits_consent(self, running_loop_mgr, storage: SQLiteBackend) -> None:
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
# Seed an expired token with no refresh — the classified getter
|
||||
# treats this as "refresh_failed" (deletes the row, returns the
|
||||
# tagged result).
|
||||
_seed_user_token(storage, cipher, expires_in_seconds=-1000)
|
||||
state = self._wire_pool(mgr, storage, cipher)
|
||||
# Drop the refresh token to force the no-refresh-token branch.
|
||||
state.mcp_token_store.delete_user_token("user-1", "pool-srv")
|
||||
state.mcp_token_store.create_user_token(
|
||||
"user-1",
|
||||
"pool-srv",
|
||||
access_token="access-aaa",
|
||||
refresh_token=None,
|
||||
expires_at=(datetime.now(UTC) - timedelta(seconds=1000)).strftime("%Y-%m-%dT%H:%M:%S"),
|
||||
scopes="openid",
|
||||
as_issuer="https://as.example.com",
|
||||
audience="https://mcp.example.com",
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_consent_required"
|
||||
|
||||
def test_token_present_dispatches_to_session(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
_seed_user_token(storage, cipher, expires_in_seconds=3600)
|
||||
self._wire_pool(mgr, storage, cipher)
|
||||
|
||||
# Pre-seed a connected pool entry so dispatch never touches the
|
||||
# SDK or the network.
|
||||
fake_session = MagicMock()
|
||||
|
||||
async def _call_tool(name, args):
|
||||
content = MagicMock()
|
||||
content.text = "tool-result"
|
||||
res = MagicMock()
|
||||
res.content = [content]
|
||||
res.isError = False
|
||||
return res
|
||||
|
||||
fake_session.call_tool = _call_tool
|
||||
|
||||
async def _seed_entry() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-1", "pool-srv"))
|
||||
entry.session = fake_session
|
||||
|
||||
_run_on_loop(loop, _seed_entry())
|
||||
|
||||
result = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{"q": "hi"},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
assert result == "tool-result"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Failure classification
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClassifyFailure:
|
||||
def test_transport_failure_classified_as_transport(self) -> None:
|
||||
mgr = MCPClientManager({})
|
||||
for exc in (
|
||||
BrokenPipeError(),
|
||||
ConnectionResetError(),
|
||||
EOFError(),
|
||||
TimeoutError("net"),
|
||||
):
|
||||
assert mgr._classify_failure(exc) == "transport"
|
||||
|
||||
def test_protocol_error_classified_as_protocol(self) -> None:
|
||||
from mcp import McpError
|
||||
from mcp.types import ErrorData
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
err = McpError(ErrorData(code=-32600, message="bad request"))
|
||||
assert mgr._classify_failure(err) == "protocol"
|
||||
|
||||
def test_other_classified_as_other(self) -> None:
|
||||
mgr = MCPClientManager({})
|
||||
assert mgr._classify_failure(ValueError("nope")) == "other"
|
||||
|
||||
def test_http_401_classified_as_auth_401(self) -> None:
|
||||
"""Defense-in-depth: ``HTTPStatusError`` classification still works
|
||||
even though Phase 6 normally consults the carrier instead.
|
||||
|
||||
Phase 6 split ``"auth"`` into ``"auth_401"`` / ``"auth_403"``
|
||||
so the dispatcher can refresh-and-retry only on 401.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
req = httpx.Request("POST", "https://mcp.example.com/sse")
|
||||
resp = httpx.Response(401, request=req)
|
||||
exc = httpx.HTTPStatusError("unauthorized", request=req, response=resp)
|
||||
assert mgr._classify_failure(exc) == "auth_401"
|
||||
|
||||
def test_http_403_classified_as_auth_403(self) -> None:
|
||||
import httpx
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
req = httpx.Request("POST", "https://mcp.example.com/sse")
|
||||
resp = httpx.Response(403, request=req)
|
||||
exc = httpx.HTTPStatusError("forbidden", request=req, response=resp)
|
||||
assert mgr._classify_failure(exc) == "auth_403"
|
||||
|
||||
def test_http_500_not_classified_as_auth(self) -> None:
|
||||
import httpx
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
req = httpx.Request("POST", "https://mcp.example.com/sse")
|
||||
resp = httpx.Response(500, request=req)
|
||||
exc = httpx.HTTPStatusError("server", request=req, response=resp)
|
||||
# 5xx is not auth — falls through to "other".
|
||||
assert mgr._classify_failure(exc) == "other"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wired-failure paths in _dispatch_pool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDispatchFailureWiring:
|
||||
"""``_classify_failure`` is consulted in production, not just tests."""
|
||||
|
||||
def _wire_pool(
|
||||
self, mgr: MCPClientManager, storage: SQLiteBackend, cipher: Any
|
||||
) -> SimpleNamespace:
|
||||
mgr.set_storage(storage)
|
||||
state = _make_app_state(storage, cipher=cipher)
|
||||
mgr.set_app_state(state)
|
||||
return state
|
||||
|
||||
def _seed_connected_session(
|
||||
self, mgr: MCPClientManager, loop: asyncio.AbstractEventLoop, exc: BaseException
|
||||
) -> None:
|
||||
async def _seed() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-1", "pool-srv"))
|
||||
sess = MagicMock()
|
||||
|
||||
async def _raise(*_args: Any, **_kwargs: Any) -> Any:
|
||||
raise exc
|
||||
|
||||
sess.call_tool = _raise
|
||||
entry.session = sess
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
|
||||
def test_dispatch_pool_transport_failure_trips_breaker(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
_seed_user_token(storage, cipher)
|
||||
self._wire_pool(mgr, storage, cipher)
|
||||
|
||||
self._seed_connected_session(mgr, loop, BrokenPipeError("dead"))
|
||||
|
||||
with pytest.raises(BrokenPipeError):
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
# Transport failure ticks the breaker.
|
||||
assert mgr._consecutive_failures.get("pool-srv", 0) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTPS enforcement (sec-1)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHttpsEnforcement:
|
||||
def test_pool_rejects_http_url_for_oauth_user(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url="http://insecure.example.com/sse")
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
state = _make_app_state(storage, cipher=cipher)
|
||||
mgr.set_app_state(state)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
payload = json.loads(str(exc_info.value))
|
||||
assert payload["error"]["code"] == "mcp_oauth_url_insecure"
|
||||
assert payload["error"]["server"] == "pool-srv"
|
||||
|
||||
def test_pool_accepts_loopback_http(self, running_loop_mgr, storage: SQLiteBackend) -> None:
|
||||
"""``http://127.0.0.1`` and ``http://localhost`` should not be blocked."""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv", url="http://127.0.0.1:8000/sse")
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
state = _make_app_state(storage, cipher=cipher)
|
||||
mgr.set_app_state(state)
|
||||
|
||||
# Pre-seed a connected pool entry so dispatch succeeds without
|
||||
# touching the network.
|
||||
fake_session = MagicMock()
|
||||
|
||||
async def _call_tool(name, args):
|
||||
content = MagicMock()
|
||||
content.text = "ok"
|
||||
res = MagicMock()
|
||||
res.content = [content]
|
||||
res.isError = False
|
||||
return res
|
||||
|
||||
fake_session.call_tool = _call_tool
|
||||
|
||||
async def _seed_entry() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-1", "pool-srv"))
|
||||
entry.session = fake_session
|
||||
|
||||
_run_on_loop(loop, _seed_entry())
|
||||
|
||||
result = mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
# Loopback URL not rejected — dispatch reaches the (fake) session.
|
||||
assert result == "ok"
|
||||
|
||||
def test_validate_oauth_user_url_helper(self) -> None:
|
||||
from turnstone.core.mcp_client import _validate_oauth_user_url
|
||||
|
||||
# Acceptable: https + the exact loopback hostnames.
|
||||
_validate_oauth_user_url("https://mcp.example.com/sse")
|
||||
_validate_oauth_user_url("http://localhost/sse")
|
||||
_validate_oauth_user_url("http://127.0.0.1:9000/sse")
|
||||
_validate_oauth_user_url("http://[::1]/sse")
|
||||
|
||||
# Rejected: non-https + non-loopback. The ``*.localhost`` suffix
|
||||
# bypass is intentionally NOT honored (RFC 6761 localhost-zone
|
||||
# resolution is configuration-dependent — custom resolvers,
|
||||
# /etc/hosts, Docker overlays may map ``foo.localhost`` to
|
||||
# non-loopback IPs).
|
||||
for bad in (
|
||||
"http://mcp.example.com/sse",
|
||||
"http://app.localhost/sse",
|
||||
"ws://mcp.example.com/sse",
|
||||
"ftp://mcp.example.com/sse",
|
||||
"//mcp.example.com/sse",
|
||||
):
|
||||
with pytest.raises(ValueError, match="https://"):
|
||||
_validate_oauth_user_url(bad)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _resolve_pool_target parser (q-9)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolvePoolTarget:
|
||||
def _make_mgr_with_oauth_server(
|
||||
self, storage: SQLiteBackend, *, name: str = "pool-srv"
|
||||
) -> MCPClientManager:
|
||||
_seed_oauth_server(storage, name=name)
|
||||
mgr = MCPClientManager({})
|
||||
mgr.set_storage(storage)
|
||||
return mgr
|
||||
|
||||
def test_malformed_prefix(self, storage: SQLiteBackend) -> None:
|
||||
mgr = self._make_mgr_with_oauth_server(storage)
|
||||
# Wrong prefix.
|
||||
assert mgr._resolve_pool_target("xyz__pool-srv__t", None, None) is None
|
||||
|
||||
def test_too_few_separators(self, storage: SQLiteBackend) -> None:
|
||||
mgr = self._make_mgr_with_oauth_server(storage)
|
||||
# mcp__server with no original_name segment.
|
||||
assert mgr._resolve_pool_target("mcp__pool-srv", None, None) is None
|
||||
|
||||
def test_empty_server_segment(self, storage: SQLiteBackend) -> None:
|
||||
mgr = self._make_mgr_with_oauth_server(storage)
|
||||
# mcp____tool — server segment is empty.
|
||||
assert mgr._resolve_pool_target("mcp____tool", None, None) is None
|
||||
|
||||
def test_original_with_double_underscore_round_trips(self, storage: SQLiteBackend) -> None:
|
||||
mgr = self._make_mgr_with_oauth_server(storage)
|
||||
target = mgr._resolve_pool_target("mcp__pool-srv__do__thing", None, None)
|
||||
assert target is not None
|
||||
assert target[0] == "pool-srv"
|
||||
# Original-name keeps its embedded ``__``.
|
||||
assert target[1] == "do__thing"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRU + lock interlock (q-7)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLruInterlock:
|
||||
def test_lru_cap_skips_locked_oldest(self, running_loop_mgr) -> None:
|
||||
"""LRU eviction must skip a locked entry the same way TTL does."""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
mgr._user_pool_idle_ttl_s = 999_999.0 # disable TTL
|
||||
mgr._user_pool_lru_max = 2
|
||||
|
||||
async def _seed_and_lock_oldest() -> tuple[asyncio.Lock, asyncio.Event]:
|
||||
base = time.monotonic()
|
||||
for i in range(3):
|
||||
key = (f"u{i}", "pool-srv")
|
||||
entry = await mgr._ensure_pool_entry(key)
|
||||
entry.session = MagicMock()
|
||||
# Older index ⇒ older timestamp.
|
||||
entry.last_used = base + i
|
||||
mgr._user_pool_last_used[key] = base + i
|
||||
# Lock the oldest (u0) so eviction must skip it and pick a younger one.
|
||||
oldest = mgr._user_pool_entries[("u0", "pool-srv")]
|
||||
held = asyncio.Event()
|
||||
|
||||
async def _hold() -> None:
|
||||
async with oldest.open_lock:
|
||||
held.set()
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
asyncio.create_task(_hold())
|
||||
await held.wait()
|
||||
return oldest.open_lock, held
|
||||
|
||||
_run_on_loop(loop, _seed_and_lock_oldest())
|
||||
|
||||
async def _evict() -> None:
|
||||
await mgr._evict_idle_pool_entries()
|
||||
|
||||
_run_on_loop(loop, _evict())
|
||||
# Locked u0 must survive.
|
||||
assert ("u0", "pool-srv") in mgr._user_pool_entries
|
||||
# The oldest unlocked entry (u1) was evicted to bring count down to cap.
|
||||
assert ("u1", "pool-srv") not in mgr._user_pool_entries
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Concurrent dispatch on shared session (M4 / perf-1)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConcurrentDispatch:
|
||||
def test_pool_concurrent_dispatch_to_same_user_server_is_serialized(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Phase 6: two tool calls on the SAME (user, server) MUST serialize
|
||||
on ``open_lock`` so the auth-introspection carrier never crosses
|
||||
between concurrent dispatches.
|
||||
|
||||
Phase 5 perf-1 released ``open_lock`` before ``call_tool`` so two
|
||||
concurrent same-key calls multiplexed on a shared
|
||||
``ClientSession``. Phase 6 reverts that for the auth-aware path
|
||||
because the per-dispatch ``_AuthCapture`` is keyed off the
|
||||
``httpx.AsyncClient`` event hook — releasing the lock would let
|
||||
a concurrent dispatch overwrite the carrier mid-flight,
|
||||
attributing one caller's 401 to another (a security bug).
|
||||
|
||||
Verified by reverting ``_dispatch_pool_with_entry`` to the
|
||||
Phase 5 shape (release ``open_lock`` before ``call_tool`` —
|
||||
i.e. move the ``in_flight += 1`` / ``call_tool`` / decrement
|
||||
block out of the ``async with`` body) and confirming this test
|
||||
observes ``max_concurrency == 2``.
|
||||
"""
|
||||
mgr, loop, _ = running_loop_mgr
|
||||
cipher = make_mcp_token_cipher()
|
||||
_seed_oauth_server(storage, name="pool-srv")
|
||||
_seed_user_token(storage, cipher)
|
||||
mgr.set_storage(storage)
|
||||
state = _make_app_state(storage, cipher=cipher)
|
||||
mgr.set_app_state(state)
|
||||
|
||||
observed_max_concurrency = 0
|
||||
in_flight = 0
|
||||
in_flight_lock = threading.Lock()
|
||||
|
||||
async def _call_tool(name, args):
|
||||
nonlocal observed_max_concurrency, in_flight
|
||||
with in_flight_lock:
|
||||
in_flight += 1
|
||||
observed_max_concurrency = max(observed_max_concurrency, in_flight)
|
||||
try:
|
||||
# Hold a moment so concurrent calls would overlap if
|
||||
# they weren't serialized on ``open_lock``.
|
||||
await asyncio.sleep(0.1)
|
||||
content = MagicMock()
|
||||
content.text = "ok"
|
||||
res = MagicMock()
|
||||
res.content = [content]
|
||||
res.isError = False
|
||||
return res
|
||||
finally:
|
||||
with in_flight_lock:
|
||||
in_flight -= 1
|
||||
|
||||
fake_session = MagicMock()
|
||||
fake_session.call_tool = _call_tool
|
||||
|
||||
async def _seed_entry() -> None:
|
||||
entry = await mgr._ensure_pool_entry(("user-1", "pool-srv"))
|
||||
entry.session = fake_session
|
||||
|
||||
_run_on_loop(loop, _seed_entry())
|
||||
|
||||
results: list[str] = []
|
||||
errors: list[Exception] = []
|
||||
|
||||
def _dispatch() -> None:
|
||||
try:
|
||||
results.append(
|
||||
mgr.call_tool_sync(
|
||||
"mcp__pool-srv__do_thing",
|
||||
{},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
)
|
||||
except Exception as exc: # pragma: no cover — diagnostic only
|
||||
errors.append(exc)
|
||||
|
||||
t1 = threading.Thread(target=_dispatch)
|
||||
t2 = threading.Thread(target=_dispatch)
|
||||
t1.start()
|
||||
t2.start()
|
||||
t1.join(timeout=5)
|
||||
t2.join(timeout=5)
|
||||
|
||||
assert errors == []
|
||||
assert results == ["ok", "ok"]
|
||||
# ``open_lock`` held across ``call_tool`` — the second dispatch
|
||||
# waits for the first to release before entering call_tool.
|
||||
assert observed_max_concurrency == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# user_id thread-through (signature)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestUserIdThreadThrough:
|
||||
def test_default_user_id_takes_static_path(self, running_loop_mgr) -> None:
|
||||
"""``user_id=None`` must leave the static-path call byte-identical."""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
# Static-path tool registered the standard way.
|
||||
mgr._tool_map["mcp__static__t"] = ("static-srv", "t")
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
fake_session = MagicMock()
|
||||
|
||||
async def _call_tool(name, args):
|
||||
content = MagicMock()
|
||||
content.text = "static-output"
|
||||
res = MagicMock()
|
||||
res.content = [content]
|
||||
res.isError = False
|
||||
return res
|
||||
|
||||
fake_session.call_tool = _call_tool
|
||||
mgr._static_servers["static-srv"] = StaticServerState(
|
||||
name="static-srv", session=fake_session
|
||||
)
|
||||
|
||||
# No user_id, no app_state — pool branch is skipped entirely.
|
||||
result = mgr.call_tool_sync("mcp__static__t", {"q": "hi"}, user_id=None, timeout=5)
|
||||
assert result == "static-output"
|
||||
|
||||
def test_user_id_with_static_path_does_not_use_pool(
|
||||
self, running_loop_mgr, storage: SQLiteBackend
|
||||
) -> None:
|
||||
"""Caller passes user_id but the resolved server is static — pool
|
||||
branch must not run because ``_lookup_server_row`` reports
|
||||
``auth_type != 'oauth_user'``."""
|
||||
mgr, _loop, _ = running_loop_mgr
|
||||
storage.create_mcp_server(
|
||||
server_id="srv-static",
|
||||
name="static-srv",
|
||||
transport="stdio",
|
||||
url="",
|
||||
command="echo",
|
||||
auth_type="static",
|
||||
)
|
||||
mgr.set_storage(storage)
|
||||
mgr.set_app_state(SimpleNamespace())
|
||||
|
||||
mgr._tool_map["mcp__static-srv__t"] = ("static-srv", "t")
|
||||
from turnstone.core.mcp_client import StaticServerState
|
||||
|
||||
fake_session = MagicMock()
|
||||
|
||||
async def _call_tool(name, args):
|
||||
content = MagicMock()
|
||||
content.text = "static-output"
|
||||
res = MagicMock()
|
||||
res.content = [content]
|
||||
res.isError = False
|
||||
return res
|
||||
|
||||
fake_session.call_tool = _call_tool
|
||||
mgr._static_servers["static-srv"] = StaticServerState(
|
||||
name="static-srv", session=fake_session
|
||||
)
|
||||
|
||||
result = mgr.call_tool_sync(
|
||||
"mcp__static-srv__t",
|
||||
{"q": "hi"},
|
||||
user_id="user-1",
|
||||
timeout=5,
|
||||
)
|
||||
assert result == "static-output"
|
||||
# No pool entries were created.
|
||||
assert mgr._user_pool_entries == {}
|
||||
@@ -1,6 +1,9 @@
|
||||
"""Tests for turnstone.core.memory_relevance — scoring, formatting, context extraction."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from turnstone.core.memory_relevance import (
|
||||
MemoryConfig,
|
||||
build_memory_context,
|
||||
extract_recent_context,
|
||||
score_memories,
|
||||
@@ -192,3 +195,265 @@ class TestExtractRecentContext:
|
||||
|
||||
def test_empty_messages(self):
|
||||
assert extract_recent_context([]) == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Composition candidate-selection (_init_system_messages)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_mem(name: str, content: str = "", memory_id: str | None = None) -> dict[str, str]:
|
||||
return {
|
||||
"name": name,
|
||||
"memory_id": memory_id or f"mid_{name}",
|
||||
"type": "project",
|
||||
"scope": "global",
|
||||
"scope_id": "",
|
||||
"description": "",
|
||||
"content": content or name,
|
||||
"updated": "2024-01-01T00:00:00",
|
||||
}
|
||||
|
||||
|
||||
def _make_session(fetch_limit: int = 5, relevance_k: int = 3, **kwargs: object):
|
||||
"""Composition tests need a real ChatSession (constructor calls
|
||||
``_init_system_messages`` once, unpatched, before the test gets a chance
|
||||
to install patches). ``tmp_db`` initializes the storage singleton that
|
||||
constructor needs; tests then patch the visibility helpers and call
|
||||
``_init_system_messages`` a second time to exercise the new logic.
|
||||
"""
|
||||
from tests._helpers import make_chat_session
|
||||
|
||||
return make_chat_session(
|
||||
memory_config=MemoryConfig(fetch_limit=fetch_limit, relevance_k=relevance_k),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class TestCompositionCandidateSelection:
|
||||
"""Verify the query-aware candidate set in _init_system_messages."""
|
||||
|
||||
def test_recency_ceiling_regression(self, tmp_db):
|
||||
"""Old relevant memory not in recency top-N still injected via search path."""
|
||||
session = _make_session(fetch_limit=5, relevance_k=3)
|
||||
session.messages = [{"role": "user", "content": "postgres database configuration"}]
|
||||
|
||||
old_mem = _make_mem(
|
||||
"ancient_db_config",
|
||||
content="postgres database configuration connection host port",
|
||||
memory_id="m_old",
|
||||
)
|
||||
# Recency top-5 do not include old_mem
|
||||
recent = [_make_mem(f"recent_{i}", memory_id=f"mr{i}") for i in range(5)]
|
||||
|
||||
with (
|
||||
patch.object(session, "_search_visible_memories", return_value=[old_mem]),
|
||||
patch.object(session, "_list_visible_memories", return_value=recent),
|
||||
):
|
||||
session._init_system_messages()
|
||||
|
||||
joined = "\n".join(m["content"] for m in session.system_messages if m["role"] == "system")
|
||||
# With the fix, old_mem enters the candidate pool via search and wins BM25
|
||||
assert "ancient_db_config" in joined
|
||||
|
||||
def test_empty_query_falls_back_to_recency(self, tmp_db):
|
||||
"""No user messages → empty context → recency path, search never called."""
|
||||
session = _make_session()
|
||||
session.messages = [] # extract_recent_context returns ""
|
||||
|
||||
recency = [_make_mem("note_alpha"), _make_mem("note_beta")]
|
||||
|
||||
with (
|
||||
patch.object(session, "_list_visible_memories", return_value=recency),
|
||||
patch.object(session, "_search_visible_memories") as search_mock,
|
||||
):
|
||||
session._init_system_messages()
|
||||
|
||||
search_mock.assert_not_called()
|
||||
joined = "\n".join(m["content"] for m in session.system_messages if m["role"] == "system")
|
||||
assert "note_alpha" in joined
|
||||
|
||||
def test_sparse_match_union_fills_candidate_pool(self, tmp_db):
|
||||
"""Search returning < fetch_limit results unions with recency fillers."""
|
||||
session = _make_session(fetch_limit=5, relevance_k=4)
|
||||
session.messages = [{"role": "user", "content": "unique_term xyzzy"}]
|
||||
|
||||
hit_a = _make_mem("hit_alpha", content="unique_term xyzzy alpha", memory_id="m_ha")
|
||||
hit_b = _make_mem("hit_beta", content="unique_term xyzzy beta", memory_id="m_hb")
|
||||
search_hits = [hit_a, hit_b] # 2 < fetch_limit=5 → triggers union
|
||||
|
||||
# Recency overlaps on hit_a/hit_b and adds 3 fillers
|
||||
filler = [_make_mem(f"filler_{i}", memory_id=f"mf{i}") for i in range(3)]
|
||||
recency = [hit_a, hit_b] + filler
|
||||
|
||||
with (
|
||||
patch.object(session, "_search_visible_memories", return_value=search_hits),
|
||||
patch.object(session, "_list_visible_memories", return_value=recency),
|
||||
):
|
||||
session._init_system_messages()
|
||||
|
||||
joined = "\n".join(m["content"] for m in session.system_messages if m["role"] == "system")
|
||||
# Both hits match "unique_term xyzzy" well → appear after BM25 ranking
|
||||
assert "hit_alpha" in joined
|
||||
assert "hit_beta" in joined
|
||||
|
||||
def test_recency_preserved_when_search_returns_noise_above_relevance_k(self, tmp_db):
|
||||
"""Pool guarantee: recency-50 always reaches BM25, even when search
|
||||
returns enough noise hits to clear ``relevance_k``.
|
||||
|
||||
Closes the narrow regression vs. the original bug — without the
|
||||
``fetch_limit`` threshold, a stopword-dominated cap-search that
|
||||
returned >= relevance_k irrelevant hits would short-circuit and
|
||||
evict the recency-only memory the bug had been surfacing.
|
||||
"""
|
||||
session = _make_session(fetch_limit=10, relevance_k=3)
|
||||
session.messages = [{"role": "user", "content": "configure host"}]
|
||||
|
||||
# Search returns relevance_k=3 noise hits — enough to skip recency
|
||||
# under the OLD threshold, not enough to fill fetch_limit=10.
|
||||
noise = [
|
||||
_make_mem(f"noise_{i}", content="generic content", memory_id=f"mn{i}") for i in range(3)
|
||||
]
|
||||
# The memory the user actually wants — distinctive, in recency,
|
||||
# but its content doesn't share any token with the noise hits.
|
||||
wanted = _make_mem(
|
||||
"host_config_v2",
|
||||
content="host=localhost port=5432 db=production",
|
||||
memory_id="m_wanted",
|
||||
)
|
||||
recency = [wanted] + [_make_mem(f"recent_{i}", memory_id=f"mr{i}") for i in range(5)]
|
||||
|
||||
with (
|
||||
patch.object(session, "_search_visible_memories", return_value=noise),
|
||||
patch.object(session, "_list_visible_memories", return_value=recency),
|
||||
):
|
||||
session._init_system_messages()
|
||||
|
||||
joined = "\n".join(m["content"] for m in session.system_messages if m["role"] == "system")
|
||||
# ``wanted`` reached BM25 via the union and matched "host" → injected.
|
||||
assert "host_config_v2" in joined
|
||||
|
||||
def test_recency_tail_preserved_when_search_adds_distinct_hits(self, tmp_db):
|
||||
"""SUPERSET invariant: every recency item is in the candidate pool
|
||||
when search adds hits, even if the resulting union exceeds
|
||||
fetch_limit. Truncating the union at fetch_limit (the prior
|
||||
behavior) evicted the recency tail — which is exactly where
|
||||
ancient-but-recently-touched memories live, the recall this PR
|
||||
sets out to improve.
|
||||
"""
|
||||
session = _make_session(fetch_limit=10, relevance_k=3)
|
||||
session.messages = [{"role": "user", "content": "alpha"}]
|
||||
|
||||
# 5 search hits, none of which appear in recency.
|
||||
search_hits = [
|
||||
_make_mem(f"search_{i}", content="alpha", memory_id=f"ms{i}") for i in range(5)
|
||||
]
|
||||
# 10 recency items; without the union uncap, the 5 oldest of these
|
||||
# would be displaced by the 5 search hits.
|
||||
recency = [_make_mem(f"recency_{i}", memory_id=f"mr{i}") for i in range(10)]
|
||||
|
||||
with (
|
||||
patch.object(session, "_search_visible_memories", return_value=search_hits),
|
||||
patch.object(session, "_list_visible_memories", return_value=recency),
|
||||
):
|
||||
candidates, source = session._select_memory_candidates("alpha")
|
||||
|
||||
candidate_ids = {c["memory_id"] for c in candidates}
|
||||
# Pool is search_hits ∪ recency — 15 items, no truncation.
|
||||
assert len(candidates) == 15
|
||||
assert source == "union"
|
||||
# Every recency item present (no tail eviction).
|
||||
for i in range(10):
|
||||
assert f"mr{i}" in candidate_ids, f"recency item {i} evicted"
|
||||
# And every search hit is also in the pool.
|
||||
for i in range(5):
|
||||
assert f"ms{i}" in candidate_ids, f"search hit {i} missing"
|
||||
|
||||
def test_coord_scope_isolated_visibility(self, tmp_db):
|
||||
"""Coord composition queries the coord scope alone, never the
|
||||
global/workstream/user union."""
|
||||
from turnstone.core.workstream import WorkstreamKind
|
||||
|
||||
coord = _make_session(
|
||||
fetch_limit=5,
|
||||
relevance_k=3,
|
||||
ws_id="coord-1",
|
||||
user_id="user-1",
|
||||
kind=WorkstreamKind.COORDINATOR,
|
||||
)
|
||||
scopes = coord._visible_scopes()
|
||||
assert scopes == [("coordinator", "coord-1")]
|
||||
# And: search uses those same scopes (no global/user fan-in)
|
||||
coord.messages = [{"role": "user", "content": "anything"}]
|
||||
with patch(
|
||||
"turnstone.core.session.search_visible_structured_memories",
|
||||
return_value=[],
|
||||
) as search_mock:
|
||||
coord._search_visible_memories("anything", limit=5)
|
||||
search_mock.assert_called_once()
|
||||
# Second positional arg is the scopes list
|
||||
assert search_mock.call_args.args[1] == [("coordinator", "coord-1")]
|
||||
|
||||
|
||||
class TestMemorySearchToolExecution:
|
||||
"""End-to-end test of ``memory(action='search')`` through _exec_memory.
|
||||
|
||||
Drives the actual tool dispatch (not just the storage facade) so the
|
||||
OR-of-terms fix and the coalesced ``memory.search`` log get exercised
|
||||
together.
|
||||
"""
|
||||
|
||||
def test_search_action_returns_or_of_terms_results(self, tmp_db):
|
||||
"""Multi-word query returns rows where ANY term matches — not all."""
|
||||
from turnstone.core.memory import save_structured_memory
|
||||
|
||||
save_structured_memory("postgres_notes", "host=localhost port=5432")
|
||||
save_structured_memory("redis_notes", "host=redis port=6379")
|
||||
save_structured_memory("unrelated", "completely different")
|
||||
|
||||
session = _make_session()
|
||||
item = session._prepare_memory(
|
||||
"call-1",
|
||||
{"action": "search", "query": "postgres no_such_word_a no_such_word_b"},
|
||||
)
|
||||
# Sanity: prepare returned a search-ready dispatch (not an error item)
|
||||
assert item.get("action") == "search"
|
||||
|
||||
call_id, msg = session._exec_memory(item)
|
||||
assert call_id == "call-1"
|
||||
assert "postgres_notes" in msg
|
||||
# Other memories don't match any query term
|
||||
assert "unrelated" not in msg
|
||||
|
||||
|
||||
class TestPerTurnSearchCache:
|
||||
"""The per-turn cache spares redundant SQL across mid-turn rebuilds."""
|
||||
|
||||
def test_repeated_search_in_same_turn_hits_cache(self, tmp_db):
|
||||
from turnstone.core.memory import save_structured_memory
|
||||
|
||||
save_structured_memory("hello_mem", "alpha beta gamma")
|
||||
session = _make_session()
|
||||
with patch(
|
||||
"turnstone.core.session.search_visible_structured_memories",
|
||||
return_value=[],
|
||||
) as backend_mock:
|
||||
session._search_visible_memories("alpha beta", limit=5)
|
||||
session._search_visible_memories("alpha beta", limit=5)
|
||||
session._search_visible_memories("alpha beta", limit=5)
|
||||
# 3 calls but only 1 backend hit — cache absorbed the rest
|
||||
assert backend_mock.call_count == 1
|
||||
|
||||
def test_user_turn_invalidates_cache(self, tmp_db):
|
||||
from turnstone.core.memory import save_structured_memory
|
||||
|
||||
save_structured_memory("hello_mem", "alpha")
|
||||
session = _make_session()
|
||||
with patch(
|
||||
"turnstone.core.session.search_visible_structured_memories",
|
||||
return_value=[],
|
||||
) as backend_mock:
|
||||
session._search_visible_memories("alpha", limit=5)
|
||||
session._invalidate_memory_cache() # simulates new user turn
|
||||
session._search_visible_memories("alpha", limit=5)
|
||||
assert backend_mock.call_count == 2
|
||||
|
||||
@@ -4,12 +4,16 @@ from turnstone.core.metacognition import (
|
||||
NUDGE_COMPLETION,
|
||||
NUDGE_CORRECTION,
|
||||
NUDGE_DENIAL,
|
||||
NUDGE_IDLE_CHILDREN_DISPLAY_CAP,
|
||||
NUDGE_IDLE_CHILDREN_WAIT_CAP,
|
||||
NUDGE_REPEAT,
|
||||
NUDGE_RESUME,
|
||||
NUDGE_START,
|
||||
NUDGE_TOOL_ERROR,
|
||||
RepeatDetector,
|
||||
detect_completion,
|
||||
detect_correction,
|
||||
format_idle_children_nudge,
|
||||
format_nudge,
|
||||
should_nudge,
|
||||
)
|
||||
@@ -308,3 +312,276 @@ class TestRepeatNudge:
|
||||
"""Repeat nudge should fire even with zero memories."""
|
||||
state: dict[str, float] = {}
|
||||
assert should_nudge("repeat", state, message_count=5, memory_count=0) is True
|
||||
|
||||
|
||||
class TestRepeatDetector:
|
||||
"""Repeat-detection streak machine — fires only when the same signature
|
||||
is recorded ``threshold`` times *consecutively* (default 3). Recording
|
||||
any different signature resets the streak, so an interrupted repeat
|
||||
isn't flagged as a stuck loop."""
|
||||
|
||||
def test_below_threshold_does_not_fire(self):
|
||||
det = RepeatDetector()
|
||||
assert det.record("a") is False
|
||||
assert det.record("a") is False # second call still under threshold
|
||||
|
||||
def test_at_threshold_fires(self):
|
||||
det = RepeatDetector()
|
||||
det.record("a")
|
||||
det.record("a")
|
||||
assert det.record("a") is True
|
||||
|
||||
def test_continues_to_fire_past_threshold(self):
|
||||
# Caller is responsible for clearing after a fire — until they do,
|
||||
# subsequent identical calls keep returning True.
|
||||
det = RepeatDetector()
|
||||
det.record("a")
|
||||
det.record("a")
|
||||
assert det.record("a") is True
|
||||
assert det.record("a") is True
|
||||
|
||||
def test_clear_resets_count(self):
|
||||
det = RepeatDetector()
|
||||
det.record("a")
|
||||
det.record("a")
|
||||
det.clear()
|
||||
assert det.record("a") is False # back to 1 after clear
|
||||
|
||||
def test_intervening_sig_resets_streak(self):
|
||||
# The streak is consecutive: recording any other sig mid-streak
|
||||
# discards the in-progress count. An alternating pattern like
|
||||
# [A, A, B, A, A] is two short streaks of 2, not a streak of 4.
|
||||
det = RepeatDetector()
|
||||
det.record("a")
|
||||
det.record("a")
|
||||
assert det.record("b") is False # b at count 1; a's streak is gone
|
||||
assert det.record("a") is False # a starts fresh at 1
|
||||
assert det.record("a") is False # a at 2
|
||||
assert det.record("a") is True # a hits 3 — fresh streak completes
|
||||
|
||||
def test_errored_signature_counts_toward_repeat(self):
|
||||
# Regression: when metacog was split out of the system message,
|
||||
# the error-output skip got reintroduced and stuck-loop detection
|
||||
# silently broke for tools that kept failing. Detector itself is
|
||||
# signature-only — error vs. success is the caller's policy.
|
||||
det = RepeatDetector()
|
||||
# Caller records an errored call's sig the same as a successful one;
|
||||
# the streak is what matters.
|
||||
for _ in range(3):
|
||||
last = det.record("bash:ls /nonexistent")
|
||||
assert last is True
|
||||
|
||||
def test_custom_threshold(self):
|
||||
det = RepeatDetector(threshold=2)
|
||||
assert det.record("a") is False
|
||||
assert det.record("a") is True
|
||||
|
||||
def test_threshold_one_fires_immediately(self):
|
||||
det = RepeatDetector(threshold=1)
|
||||
assert det.record("a") is True
|
||||
|
||||
|
||||
class TestFormatIdleChildrenNudge:
|
||||
"""``format_idle_children_nudge`` renders the wake-driven idle_children
|
||||
body — no ``<system-reminder>`` envelope (the side-channel splice
|
||||
wraps it at the wire boundary).
|
||||
"""
|
||||
|
||||
def test_empty_list_returns_empty_string(self):
|
||||
# Caller short-circuits on `if not text: return` — so empty
|
||||
# input MUST produce empty output, not a header-only stub.
|
||||
assert format_idle_children_nudge([]) == ""
|
||||
|
||||
def test_single_child_renders(self):
|
||||
children = [{"ws_id": "ws-abc12345", "name": "research-task", "state": "running"}]
|
||||
text = format_idle_children_nudge(children)
|
||||
assert "ws-abc12" in text # short-id form (8 chars)
|
||||
assert "research-task" in text
|
||||
assert "running" in text
|
||||
assert "wait_for_workstream" in text
|
||||
assert "ws-abc12345" in text # full id appears in the suggestion's ws_ids list
|
||||
|
||||
def test_under_display_cap_no_overflow_line(self):
|
||||
children = [
|
||||
{"ws_id": f"ws-{i:08d}", "name": f"task-{i}", "state": "running"} for i in range(3)
|
||||
]
|
||||
text = format_idle_children_nudge(children)
|
||||
assert "...and" not in text
|
||||
for i in range(3):
|
||||
assert f"task-{i}" in text
|
||||
|
||||
def test_over_display_cap_renders_overflow_line(self):
|
||||
n = NUDGE_IDLE_CHILDREN_DISPLAY_CAP + 4
|
||||
children = [
|
||||
{"ws_id": f"ws-{i:08d}", "name": f"task-{i}", "state": "thinking"} for i in range(n)
|
||||
]
|
||||
text = format_idle_children_nudge(children)
|
||||
assert f"...and {n - NUDGE_IDLE_CHILDREN_DISPLAY_CAP} more" in text
|
||||
# First N children are inline; later ones are folded into "...and N more".
|
||||
for i in range(NUDGE_IDLE_CHILDREN_DISPLAY_CAP):
|
||||
assert f"task-{i}" in text
|
||||
for i in range(NUDGE_IDLE_CHILDREN_DISPLAY_CAP, n):
|
||||
# Names beyond the display cap aren't visible; only counted.
|
||||
assert f"task-{i}" not in text
|
||||
|
||||
def test_over_wait_cap_truncates_suggestion_ws_ids(self):
|
||||
n = NUDGE_IDLE_CHILDREN_WAIT_CAP + 5
|
||||
children = [
|
||||
{"ws_id": f"ws-{i:08d}", "name": f"task-{i}", "state": "running"} for i in range(n)
|
||||
]
|
||||
text = format_idle_children_nudge(children)
|
||||
# The first WAIT_CAP ids appear in the suggestion; later ones don't.
|
||||
first_in_suggestion = f"ws-{NUDGE_IDLE_CHILDREN_WAIT_CAP - 1:08d}"
|
||||
first_excluded = f"ws-{NUDGE_IDLE_CHILDREN_WAIT_CAP:08d}"
|
||||
assert first_in_suggestion in text
|
||||
assert first_excluded not in text
|
||||
|
||||
def test_unnamed_child_falls_back(self):
|
||||
children = [{"ws_id": "ws-deadbeef", "name": "", "state": "attention"}]
|
||||
text = format_idle_children_nudge(children)
|
||||
assert "(unnamed)" in text
|
||||
assert "attention" in text
|
||||
|
||||
def test_newline_in_name_does_not_forge_extra_bullet(self):
|
||||
"""A workstream name with embedded ``\\n`` / ``\\t`` / ``\\r`` MUST
|
||||
NOT break the bullet structure — :func:`sanitize_name`'s strict
|
||||
regex strips control chars (incl. TAB/LF/CR) so the name stays
|
||||
on a single line under its own bullet. Without this, a
|
||||
malicious child name like ``"foo\\n - ws-fake (running): bar"``
|
||||
would forge a fake sibling row in the rendered list.
|
||||
"""
|
||||
children = [
|
||||
{"ws_id": "ws-real0001", "name": "real", "state": "running"},
|
||||
{
|
||||
"ws_id": "ws-evil0002",
|
||||
"name": "evil\n - ws-fake (running): forged",
|
||||
"state": "thinking",
|
||||
},
|
||||
{"ws_id": "ws-real0003", "name": "tail", "state": "running"},
|
||||
]
|
||||
text = format_idle_children_nudge(children)
|
||||
bullet_rows = [ln for ln in text.splitlines() if ln.startswith(" - ")]
|
||||
assert len(bullet_rows) == 3, (
|
||||
f"expected 3 bullet rows; got {len(bullet_rows)}: {bullet_rows!r}"
|
||||
)
|
||||
evil_row = next(row for row in bullet_rows if "ws-evil" in row)
|
||||
assert "\n" not in evil_row
|
||||
assert "\t" not in evil_row
|
||||
assert "\r" not in evil_row
|
||||
assert "evil" in evil_row
|
||||
assert "ws-real" in bullet_rows[2]
|
||||
assert "tail" in bullet_rows[2]
|
||||
|
||||
def test_missing_state_renders_question_mark(self):
|
||||
children = [{"ws_id": "ws-12345678", "name": "x"}]
|
||||
text = format_idle_children_nudge(children)
|
||||
# Defensive default — exotic state keys / partial dicts shouldn't crash.
|
||||
assert "?" in text
|
||||
|
||||
def test_no_system_reminder_envelope(self):
|
||||
# The side-channel ``_apply_reminders_for_provider`` splice
|
||||
# adds ``<system-reminder>`` at the wire boundary; the formatter
|
||||
# MUST NOT wrap, or the model would see a doubled envelope.
|
||||
text = format_idle_children_nudge([{"ws_id": "ws-x", "name": "y", "state": "running"}])
|
||||
assert "<system-reminder>" not in text
|
||||
assert "</system-reminder>" not in text
|
||||
|
||||
def test_format_nudge_returns_empty_for_idle_children(self):
|
||||
# The static map's idle_children entry is the empty string by
|
||||
# design — format_idle_children_nudge produces the real body.
|
||||
assert format_nudge("idle_children") == ""
|
||||
|
||||
def test_should_nudge_recognises_idle_children_type(self, monkeypatch):
|
||||
# Type registration in ``_NUDGE_MAP`` makes ``should_nudge``
|
||||
# recognise it for cooldown gating; without the entry it would
|
||||
# silently return False on every call.
|
||||
state: dict[str, float] = {}
|
||||
# message_count > 1 to clear the first-message gate.
|
||||
assert should_nudge("idle_children", state, message_count=4, memory_count=0) is True
|
||||
# Cooldown set on success → second immediate call returns False.
|
||||
assert should_nudge("idle_children", state, message_count=5, memory_count=0) is False
|
||||
|
||||
|
||||
class TestSanitizeName:
|
||||
"""Strict sanitiser for single-line user-controlled name fields
|
||||
(used by :func:`format_idle_children_nudge` for the workstream
|
||||
``name``). Strips ASCII control chars **including** TAB/LF/CR
|
||||
plus Unicode steering vectors and angle-bracket tag breakers.
|
||||
"""
|
||||
|
||||
def test_empty_input_returns_empty(self):
|
||||
from turnstone.core.metacognition import sanitize_name
|
||||
|
||||
assert sanitize_name("") == ""
|
||||
|
||||
def test_strips_tab_lf_cr(self):
|
||||
"""Strict variant: TAB/LF/CR are stripped so a hostile name with
|
||||
an embedded newline can't break a bullet's one-line structure.
|
||||
"""
|
||||
from turnstone.core.metacognition import sanitize_name
|
||||
|
||||
# All three become spaces (then collapsed to one inline space
|
||||
# by the trailing ``strip()``-on-leading/trailing-only step
|
||||
# — interior runs stay as multiple spaces, that's fine for a
|
||||
# one-line name).
|
||||
assert sanitize_name("a\tb") == "a b"
|
||||
assert sanitize_name("a\nb") == "a b"
|
||||
assert sanitize_name("a\rb") == "a b"
|
||||
|
||||
def test_strips_other_ascii_control_chars(self):
|
||||
from turnstone.core.metacognition import sanitize_name
|
||||
|
||||
assert sanitize_name("a\x07b\x0bc\x0cd") == "a b c d"
|
||||
assert sanitize_name("a\x7fb") == "a b"
|
||||
|
||||
def test_strips_angle_bracket_tag_breakers(self):
|
||||
from turnstone.core.metacognition import sanitize_name
|
||||
|
||||
assert sanitize_name("a</thinking>b") == "a/thinkingb"
|
||||
|
||||
|
||||
class TestSanitizePayload:
|
||||
"""Permissive sanitiser used by the ``watch_triggered`` producer.
|
||||
Strips ASCII control chars (except TAB/LF/CR), Unicode steering
|
||||
vectors (bidi, zero-width, BOM, tag chars), and angle-bracket
|
||||
tag breakers — keeps everything else intact, so multi-line shell
|
||||
output retains its line structure.
|
||||
"""
|
||||
|
||||
def test_empty_input_returns_empty(self):
|
||||
from turnstone.core.metacognition import sanitize_payload
|
||||
|
||||
assert sanitize_payload("") == ""
|
||||
|
||||
def test_strips_ascii_control_chars(self):
|
||||
"""``\\x00``-``\\x1f`` minus TAB/LF/CR plus ``\\x7f`` (DEL) become spaces."""
|
||||
from turnstone.core.metacognition import sanitize_payload
|
||||
|
||||
# BEL (0x07), VT (0x0b), FF (0x0c) — all in strip set.
|
||||
assert sanitize_payload("a\x07b\x0bc\x0cd") == "a b c d"
|
||||
# DEL (0x7f).
|
||||
assert sanitize_payload("a\x7fb") == "a b"
|
||||
|
||||
def test_preserves_tab_lf_cr(self):
|
||||
"""TAB / LF / CR are intentionally preserved so multi-line shell
|
||||
output keeps its line structure when sanitised as a watch payload.
|
||||
"""
|
||||
from turnstone.core.metacognition import sanitize_payload
|
||||
|
||||
# Newlines kept; only the leading + trailing strip happens.
|
||||
out = sanitize_payload("line1\nline2\n\tindented\rline3")
|
||||
assert out == "line1\nline2\n\tindented\rline3"
|
||||
|
||||
def test_strips_bidi_and_zero_width(self):
|
||||
from turnstone.core.metacognition import sanitize_payload
|
||||
|
||||
# U+202E RIGHT-TO-LEFT OVERRIDE; U+200B ZERO WIDTH SPACE.
|
||||
assert sanitize_payload("abc") == "a b c"
|
||||
|
||||
def test_strips_angle_bracket_tag_breakers(self):
|
||||
from turnstone.core.metacognition import sanitize_payload
|
||||
|
||||
# "<" / ">" go away entirely (not replaced with space) so a name
|
||||
# like "</thinking>" doesn't leave a hole the model can read as
|
||||
# a structural marker.
|
||||
assert sanitize_payload("a</thinking>b") == "a/thinkingb"
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Tests for alembic migration 049 (OAuth-MCP schema).
|
||||
|
||||
Drives ``command.upgrade`` from a programmatic Alembic config against
|
||||
an isolated SQLite database per test, then asserts:
|
||||
|
||||
* the two new tables (``mcp_user_tokens``, ``mcp_oauth_pending``) exist,
|
||||
* the eight new ``mcp_servers`` columns exist,
|
||||
* the post-upgrade ``UPDATE mcp_servers`` normalization rewrites rows
|
||||
with empty / missing headers to ``auth_type='none'`` while leaving
|
||||
rows with non-empty headers at ``auth_type='static'``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
|
||||
_MIGRATIONS_DIR = str(
|
||||
Path(__file__).resolve().parent.parent / "turnstone" / "core" / "storage" / "migrations"
|
||||
)
|
||||
|
||||
|
||||
def _alembic_cfg(db_path: Path) -> Config:
|
||||
cfg = Config()
|
||||
cfg.set_main_option("script_location", _MIGRATIONS_DIR)
|
||||
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
|
||||
return cfg
|
||||
|
||||
|
||||
class TestMigration049:
|
||||
def test_creates_new_tables_and_columns(self, tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "049.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
|
||||
# Walk forward through 048 first, then explicitly to 049 so we
|
||||
# exercise the *upgrade* function (not just the schema's `head`).
|
||||
command.upgrade(cfg, "048")
|
||||
command.upgrade(cfg, "049")
|
||||
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
inspector = sa.inspect(engine)
|
||||
tables = set(inspector.get_table_names())
|
||||
assert "mcp_user_tokens" in tables
|
||||
assert "mcp_oauth_pending" in tables
|
||||
|
||||
mcp_cols = {c["name"] for c in inspector.get_columns("mcp_servers")}
|
||||
new_cols = {
|
||||
"auth_type",
|
||||
"oauth_client_id",
|
||||
"oauth_client_secret_ct",
|
||||
"oauth_scopes",
|
||||
"oauth_audience",
|
||||
"oauth_registration_mode",
|
||||
"oauth_authorization_server_url",
|
||||
"oauth_as_issuer_cached",
|
||||
}
|
||||
assert new_cols.issubset(mcp_cols), new_cols - mcp_cols
|
||||
|
||||
# Index check on mcp_oauth_pending.
|
||||
indexes = {ix["name"] for ix in inspector.get_indexes("mcp_oauth_pending")}
|
||||
assert "idx_mcp_pending_created" in indexes
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_normalizes_empty_headers_to_none(self, tmp_path: Path) -> None:
|
||||
"""Streamable-http rows with NULL / '' / '{}' headers become
|
||||
auth_type='none'; rows with non-empty headers stay 'static'.
|
||||
Stdio rows always stay 'static' regardless of headers — the
|
||||
column value is opaque when there is no HTTP transport."""
|
||||
db_path = tmp_path / "049-norm.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
|
||||
# Apply everything up to 048, seed rows, then apply 049.
|
||||
command.upgrade(cfg, "048")
|
||||
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"""
|
||||
INSERT INTO mcp_servers (
|
||||
server_id, name, transport, command, args, url,
|
||||
headers, env, auto_approve, enabled, created_by,
|
||||
registry_name, registry_version, registry_meta,
|
||||
created, updated
|
||||
) VALUES (
|
||||
:sid, :name, :transport, '', '[]',
|
||||
'https://x', :headers, '{}', 0, 1, '', NULL, '',
|
||||
'{}', '2026-05-04T11:00:00', '2026-05-04T11:00:00'
|
||||
)
|
||||
"""
|
||||
),
|
||||
[
|
||||
{
|
||||
"sid": "s-empty-str",
|
||||
"name": "empty-str",
|
||||
"transport": "streamable-http",
|
||||
"headers": "",
|
||||
},
|
||||
{
|
||||
"sid": "s-empty-obj",
|
||||
"name": "empty-obj",
|
||||
"transport": "streamable-http",
|
||||
"headers": "{}",
|
||||
},
|
||||
{
|
||||
"sid": "s-with-headers",
|
||||
"name": "with-headers",
|
||||
"transport": "streamable-http",
|
||||
"headers": '{"Authorization":"Bearer x"}',
|
||||
},
|
||||
# Stdio rows must keep the 'static' default, even
|
||||
# though their headers are empty — auth_type is
|
||||
# opaque for stdio.
|
||||
{
|
||||
"sid": "s-stdio-empty",
|
||||
"name": "stdio-empty",
|
||||
"transport": "stdio",
|
||||
"headers": "{}",
|
||||
},
|
||||
{
|
||||
"sid": "s-stdio-null",
|
||||
"name": "stdio-null",
|
||||
"transport": "stdio",
|
||||
"headers": "",
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
command.upgrade(cfg, "049")
|
||||
|
||||
with engine.connect() as conn:
|
||||
rows = dict(conn.execute(sa.text("SELECT name, auth_type FROM mcp_servers")).all())
|
||||
assert rows["empty-str"] == "none"
|
||||
assert rows["empty-obj"] == "none"
|
||||
assert rows["with-headers"] == "static"
|
||||
# Stdio rows must remain at the 'static' column default even
|
||||
# when headers are empty — the migration only touches HTTP
|
||||
# rows where auth_type is semantically meaningful.
|
||||
assert rows["stdio-empty"] == "static"
|
||||
assert rows["stdio-null"] == "static"
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
def test_full_chain_to_head(self, tmp_path: Path) -> None:
|
||||
"""Sanity: running ``upgrade head`` on a fresh DB yields the
|
||||
same end-state column set as ``_schema.metadata``."""
|
||||
db_path = tmp_path / "049-head.db"
|
||||
cfg = _alembic_cfg(db_path)
|
||||
command.upgrade(cfg, "head")
|
||||
|
||||
engine = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
from turnstone.core.storage._schema import mcp_servers
|
||||
|
||||
inspector = sa.inspect(engine)
|
||||
actual = {c["name"] for c in inspector.get_columns("mcp_servers")}
|
||||
expected = {c.name for c in mcp_servers.columns}
|
||||
assert expected.issubset(actual), expected - actual
|
||||
finally:
|
||||
engine.dispose()
|
||||
+144
-12
@@ -770,16 +770,76 @@ class TestRegistryReload:
|
||||
assert reg.has_alias("b")
|
||||
assert reg.default == "b"
|
||||
|
||||
def test_reload_clears_clients(self) -> None:
|
||||
models = {"a": ModelConfig("a", "http://x/v1", "key", "m")}
|
||||
def test_reload_keeps_clients_when_connection_target_unchanged(self) -> None:
|
||||
"""Selective teardown: a model edit that leaves base_url / api_key /
|
||||
provider intact (e.g. admin tweaks the underlying ``model`` name or
|
||||
``temperature``) keeps the cached HTTP client warm — no need to
|
||||
re-establish TLS+pool when the endpoint is the same."""
|
||||
models = {"a": ModelConfig("a", "http://x/v1", "key", "m1", provider="openai")}
|
||||
reg = ModelRegistry(models=models, default="a")
|
||||
# Force client creation
|
||||
reg.get_client("a")
|
||||
assert "a" in reg._clients
|
||||
client_before = reg._clients["a"]
|
||||
provider_before = reg.get_provider("a")
|
||||
|
||||
# Same endpoint (base_url, api_key, provider), only ``model`` changed.
|
||||
new_models = {"a": ModelConfig("a", "http://x/v1", "key", "m2", provider="openai")}
|
||||
reg.reload(new_models, "a")
|
||||
|
||||
assert "a" in reg._clients
|
||||
assert reg._clients["a"] is client_before
|
||||
assert "a" in reg._providers
|
||||
assert reg._providers["a"] is provider_before
|
||||
|
||||
def test_reload_drops_client_when_base_url_changes(self) -> None:
|
||||
"""A ``base_url`` change drops the cached client (different
|
||||
endpoint = new connection) but keeps the cached provider —
|
||||
``LLMProvider`` is keyed only on the provider string, which
|
||||
didn't change."""
|
||||
models = {"a": ModelConfig("a", "http://x/v1", "key", "m", provider="openai")}
|
||||
reg = ModelRegistry(models=models, default="a")
|
||||
reg.get_client("a")
|
||||
provider_before = reg.get_provider("a")
|
||||
|
||||
new_models = {"a": ModelConfig("a", "http://y/v1", "key", "m", provider="openai")}
|
||||
reg.reload(new_models, "a")
|
||||
|
||||
# Reload with same models — clients should be cleared
|
||||
reg.reload(dict(models), "a")
|
||||
assert "a" not in reg._clients
|
||||
assert "a" in reg._providers
|
||||
assert reg._providers["a"] is provider_before
|
||||
|
||||
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
|
||||
``LLMProvider`` implementation against the new SDK."""
|
||||
models = {"a": ModelConfig("a", "http://x/v1", "key", "m", provider="openai")}
|
||||
reg = ModelRegistry(models=models, default="a")
|
||||
reg.get_client("a")
|
||||
reg.get_provider("a")
|
||||
|
||||
new_models = {"a": ModelConfig("a", "http://x/v1", "key", "m", provider="anthropic")}
|
||||
reg.reload(new_models, "a")
|
||||
|
||||
assert "a" not in reg._clients
|
||||
assert "a" not in reg._providers
|
||||
|
||||
def test_reload_drops_clients_for_removed_aliases(self) -> None:
|
||||
"""Aliases removed from the registry must release their cached
|
||||
clients — otherwise a deleted endpoint's connection pool would
|
||||
outlive the alias indefinitely."""
|
||||
models = {
|
||||
"a": ModelConfig("a", "http://x/v1", "key", "m"),
|
||||
"b": ModelConfig("b", "http://y/v1", "key", "m"),
|
||||
}
|
||||
reg = ModelRegistry(models=models, default="a")
|
||||
reg.get_client("a")
|
||||
reg.get_client("b")
|
||||
|
||||
# Drop "b" entirely.
|
||||
new_models = {"a": ModelConfig("a", "http://x/v1", "key", "m")}
|
||||
reg.reload(new_models, "a")
|
||||
|
||||
assert "a" in reg._clients # unchanged endpoint, kept warm
|
||||
assert "b" not in reg._clients
|
||||
|
||||
def test_reload_validates_default(self) -> None:
|
||||
models_a = {"a": ModelConfig("a", "x", "x", "m")}
|
||||
@@ -1066,24 +1126,57 @@ class TestSessionAgentModel:
|
||||
def _captured_effort(captured: dict[str, Any]) -> str | None:
|
||||
"""Pull reasoning_effort out of provider-specific shapes.
|
||||
|
||||
openai-compatible servers receive it via extra_body.chat_template_kwargs;
|
||||
commercial providers receive it as a top-level kwarg.
|
||||
Chat Completions delivers it as a top-level ``reasoning_effort`` kwarg
|
||||
(when the model's caps permit it). Operators who route reasoning_effort
|
||||
through ``chat_template_kwargs`` (gpt-oss-style local templates) get
|
||||
it inside ``extra_body.chat_template_kwargs``.
|
||||
"""
|
||||
if "reasoning_effort" in captured:
|
||||
return captured["reasoning_effort"]
|
||||
eb = captured.get("extra_body") or {}
|
||||
ctk = eb.get("chat_template_kwargs") or {}
|
||||
return ctk.get("reasoning_effort") or captured.get("reasoning_effort")
|
||||
return ctk.get("reasoning_effort")
|
||||
|
||||
@staticmethod
|
||||
def _effort_caps() -> dict[str, Any]:
|
||||
"""Capabilities that allow Chat-Completions reasoning_effort to flow."""
|
||||
return {
|
||||
"reasoning_effort_values": [
|
||||
"minimal",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"max",
|
||||
],
|
||||
}
|
||||
|
||||
def _three_model_registry(self, **kwargs: Any) -> ModelRegistry:
|
||||
caps = self._effort_caps()
|
||||
return ModelRegistry(
|
||||
models={
|
||||
"main": ModelConfig(
|
||||
"main", "http://m/v1", "k", "main-model", provider="openai-compatible"
|
||||
"main",
|
||||
"http://m/v1",
|
||||
"k",
|
||||
"main-model",
|
||||
provider="openai-compatible",
|
||||
capabilities=dict(caps),
|
||||
),
|
||||
"smart": ModelConfig(
|
||||
"smart", "http://s/v1", "k", "smart-model", provider="openai-compatible"
|
||||
"smart",
|
||||
"http://s/v1",
|
||||
"k",
|
||||
"smart-model",
|
||||
provider="openai-compatible",
|
||||
capabilities=dict(caps),
|
||||
),
|
||||
"fast": ModelConfig(
|
||||
"fast", "http://f/v1", "k", "fast-model", provider="openai-compatible"
|
||||
"fast",
|
||||
"http://f/v1",
|
||||
"k",
|
||||
"fast-model",
|
||||
provider="openai-compatible",
|
||||
capabilities=dict(caps),
|
||||
),
|
||||
},
|
||||
default="main",
|
||||
@@ -1189,6 +1282,45 @@ class TestSessionAgentModel:
|
||||
session._run_agent([{"role": "user", "content": "x"}], label="plan", agent_alias="fast")
|
||||
assert captured["model"] == "fast-model"
|
||||
|
||||
def test_session_fallback_inherits_primary_alias_for_caps(self) -> None:
|
||||
"""When _run_agent has no registry agent route, it must fall back to
|
||||
the session's primary alias for capability and server_compat lookup —
|
||||
otherwise per-model caps (reasoning_effort_values, server_compat) get
|
||||
silently dropped on the agent path."""
|
||||
reg = self._three_model_registry() # no agent_model / plan_model set
|
||||
session = _make_session(registry=reg, model_alias="main")
|
||||
# Probe what _run_agent passes to _provider_extra_params and
|
||||
# _resolve_capabilities by recording the model_alias on each call.
|
||||
captured_extra_alias: list[str | None] = []
|
||||
captured_resolve_alias: list[str | None] = []
|
||||
original_extra = session._provider_extra_params
|
||||
original_resolve = session._resolve_capabilities
|
||||
|
||||
def spy_extra(*args: Any, **kwargs: Any) -> Any:
|
||||
captured_extra_alias.append(kwargs.get("model_alias"))
|
||||
return original_extra(*args, **kwargs)
|
||||
|
||||
def spy_resolve(*args: Any, **kwargs: Any) -> Any:
|
||||
# _resolve_capabilities(provider, model, alias)
|
||||
alias = args[2] if len(args) >= 3 else kwargs.get("alias")
|
||||
captured_resolve_alias.append(alias)
|
||||
return original_resolve(*args, **kwargs)
|
||||
|
||||
session._provider_extra_params = spy_extra # type: ignore[method-assign]
|
||||
session._resolve_capabilities = spy_resolve # type: ignore[method-assign]
|
||||
|
||||
self._capture_on(session.client) # patch client.chat.completions.create
|
||||
session._run_agent([{"role": "user", "content": "x"}], label="plan")
|
||||
|
||||
assert captured_extra_alias and captured_extra_alias[-1] == "main", (
|
||||
f"agent fallback path did not inherit primary alias for extra_params: "
|
||||
f"{captured_extra_alias!r}"
|
||||
)
|
||||
assert captured_resolve_alias and captured_resolve_alias[-1] == "main", (
|
||||
f"agent fallback path did not inherit primary alias for caps: "
|
||||
f"{captured_resolve_alias!r}"
|
||||
)
|
||||
|
||||
def test_invalid_alias_raises_in_run_agent(self) -> None:
|
||||
"""Defence-in-depth: _prepare_* validates first, but _run_agent
|
||||
rejects unknown aliases too rather than silently falling back."""
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
"""``models_changed`` SSE fanout coverage.
|
||||
|
||||
The console pushes a ``models_changed`` cluster event whenever a model
|
||||
definition is created / updated / deleted / reloaded, or whenever a
|
||||
setting in :data:`turnstone.console.server._MODEL_AFFECTING_SETTING_KEYS`
|
||||
is updated or reset. Connected browsers refetch ``/v1/api/models`` on
|
||||
receipt so the home composer dropdown + admin Models → Roles sub-tab
|
||||
reflect alias edits without a manual reload.
|
||||
|
||||
These tests pin two contracts:
|
||||
|
||||
- every model-definition CRUD path emits exactly one ``models_changed``
|
||||
fanout (so the browser stays in sync with the DB);
|
||||
- settings PUT / DELETE only emit the fanout when the key is
|
||||
model-affecting — unrelated keys (e.g. ``session.retention_days``)
|
||||
must not trigger spurious dropdown re-renders across the cluster.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.routing import Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from tests._coord_test_helpers import _AuthMiddleware
|
||||
from turnstone.console.server import (
|
||||
_MODEL_AFFECTING_SETTING_KEYS,
|
||||
admin_create_model_definition,
|
||||
admin_delete_model_definition,
|
||||
admin_delete_setting,
|
||||
admin_model_reload,
|
||||
admin_update_model_definition,
|
||||
admin_update_setting,
|
||||
)
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path: Any) -> SQLiteBackend:
|
||||
return SQLiteBackend(str(tmp_path / "models_changed.db"))
|
||||
|
||||
|
||||
def _seed(storage: SQLiteBackend, *, definition_id: str, alias: str) -> None:
|
||||
storage.create_model_definition(
|
||||
definition_id=definition_id,
|
||||
alias=alias,
|
||||
model="model-x",
|
||||
provider="openai-compatible",
|
||||
base_url="http://localhost:8000/v1",
|
||||
api_key="sk-test",
|
||||
context_window=8192,
|
||||
capabilities="{}",
|
||||
enabled=True,
|
||||
created_by="admin",
|
||||
)
|
||||
|
||||
|
||||
def _make_client(storage: SQLiteBackend) -> tuple[TestClient, MagicMock]:
|
||||
"""Build a TestClient + return the stub collector for assertion.
|
||||
|
||||
Wires the four model-definition CRUD/reload routes plus the two
|
||||
settings mutation routes. Collector is a MagicMock so each
|
||||
``emit_models_changed`` call lands as a recorded call without
|
||||
spinning up the full SSE listener queue.
|
||||
"""
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route(
|
||||
"/v1/api/admin/model-definitions",
|
||||
admin_create_model_definition,
|
||||
methods=["POST"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/admin/model-definitions/reload",
|
||||
admin_model_reload,
|
||||
methods=["POST"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/admin/model-definitions/{definition_id}",
|
||||
admin_update_model_definition,
|
||||
methods=["PUT"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/admin/model-definitions/{definition_id}",
|
||||
admin_delete_model_definition,
|
||||
methods=["DELETE"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/admin/settings/{key:path}",
|
||||
admin_update_setting,
|
||||
methods=["PUT"],
|
||||
),
|
||||
Route(
|
||||
"/v1/api/admin/settings/{key:path}",
|
||||
admin_delete_setting,
|
||||
methods=["DELETE"],
|
||||
),
|
||||
],
|
||||
middleware=[Middleware(_AuthMiddleware)],
|
||||
)
|
||||
app.state.auth_storage = storage
|
||||
app.state.coord_registry = None # CRUD endpoints handle this gracefully
|
||||
collector = MagicMock()
|
||||
collector.get_all_nodes.return_value = []
|
||||
app.state.collector = collector
|
||||
app.state.proxy_client = MagicMock()
|
||||
app.state.config_store = MagicMock()
|
||||
client = TestClient(app)
|
||||
client.headers.update(
|
||||
{
|
||||
"X-Test-User": "admin",
|
||||
"X-Test-Perms": "admin.models,admin.settings",
|
||||
}
|
||||
)
|
||||
return client, collector
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model-definition CRUD endpoints fan out ``models_changed``
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_emits_models_changed(storage: SQLiteBackend) -> None:
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.post(
|
||||
"/v1/api/admin/model-definitions",
|
||||
json={
|
||||
"alias": "fast",
|
||||
"model": "fast-model",
|
||||
"provider": "openai-compatible",
|
||||
"base_url": "http://localhost:9000/v1",
|
||||
"api_key": "sk-x",
|
||||
"context_window": 4096,
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
def test_update_emits_models_changed(storage: SQLiteBackend) -> None:
|
||||
_seed(storage, definition_id="m1", alias="local")
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.put(
|
||||
"/v1/api/admin/model-definitions/m1",
|
||||
json={"model": "swapped-model"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
def test_update_with_empty_body_does_not_emit(storage: SQLiteBackend) -> None:
|
||||
"""Empty-body PUT writes no rows + skips the registry refresh — no
|
||||
SSE fanout either, since nothing actually changed."""
|
||||
_seed(storage, definition_id="m1", alias="local")
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.put("/v1/api/admin/model-definitions/m1", json={})
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 0
|
||||
|
||||
|
||||
def test_delete_emits_models_changed(storage: SQLiteBackend) -> None:
|
||||
_seed(storage, definition_id="m1", alias="local")
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.delete("/v1/api/admin/model-definitions/m1")
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
def test_reload_emits_models_changed(storage: SQLiteBackend) -> None:
|
||||
_seed(storage, definition_id="m1", alias="local")
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.post("/v1/api/admin/model-definitions/reload")
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Settings PUT / DELETE only emit for model-affecting keys
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Pinned snapshot of the role-related keys we expect the allowlist to
|
||||
# cover today. The frozenset itself is asserted further down so a
|
||||
# stray addition doesn't silently bypass coverage.
|
||||
_EXPECTED_AFFECTING_KEYS = frozenset(
|
||||
{
|
||||
"model.default_alias",
|
||||
"model.plan_alias",
|
||||
"model.plan_effort",
|
||||
"model.task_alias",
|
||||
"model.task_effort",
|
||||
"coordinator.model_alias",
|
||||
"coordinator.reasoning_effort",
|
||||
"judge.model",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _value_for_key(key: str) -> str:
|
||||
"""Return a registry-valid value for ``key``.
|
||||
|
||||
``reasoning_effort`` keys have a fixed choice list; alias-shaped
|
||||
keys accept arbitrary strings. Avoids per-key custom payloads.
|
||||
"""
|
||||
if (
|
||||
key.endswith("reasoning_effort")
|
||||
or key.endswith("plan_effort")
|
||||
or key.endswith("task_effort")
|
||||
):
|
||||
return "low"
|
||||
return "anything"
|
||||
|
||||
|
||||
def test_affecting_keys_set_matches_expected() -> None:
|
||||
"""Lock in the allowlist so an unintentional removal is caught."""
|
||||
assert _MODEL_AFFECTING_SETTING_KEYS == _EXPECTED_AFFECTING_KEYS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", sorted(_EXPECTED_AFFECTING_KEYS))
|
||||
def test_settings_put_emits_for_model_affecting_key(storage: SQLiteBackend, key: str) -> None:
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.put(
|
||||
f"/v1/api/admin/settings/{key}",
|
||||
json={"value": _value_for_key(key)},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", sorted(_EXPECTED_AFFECTING_KEYS))
|
||||
def test_settings_delete_emits_for_model_affecting_key(storage: SQLiteBackend, key: str) -> None:
|
||||
client, collector = _make_client(storage)
|
||||
# Seed a row so DELETE has something to remove (otherwise 404).
|
||||
client.put(
|
||||
f"/v1/api/admin/settings/{key}",
|
||||
json={"value": _value_for_key(key)},
|
||||
)
|
||||
collector.emit_models_changed.reset_mock()
|
||||
resp = client.delete(f"/v1/api/admin/settings/{key}")
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 1
|
||||
|
||||
|
||||
def test_settings_put_does_not_emit_for_unrelated_key(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""Updating a non-model setting (here: a session retention knob)
|
||||
must not trigger a cluster-wide dropdown refresh."""
|
||||
client, collector = _make_client(storage)
|
||||
resp = client.put(
|
||||
"/v1/api/admin/settings/session.retention_days",
|
||||
json={"value": 30},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 0
|
||||
|
||||
|
||||
def test_settings_delete_does_not_emit_for_unrelated_key(
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
client, collector = _make_client(storage)
|
||||
client.put(
|
||||
"/v1/api/admin/settings/session.retention_days",
|
||||
json={"value": 30},
|
||||
)
|
||||
collector.emit_models_changed.reset_mock()
|
||||
resp = client.delete("/v1/api/admin/settings/session.retention_days")
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert collector.emit_models_changed.call_count == 0
|
||||
@@ -0,0 +1,533 @@
|
||||
"""Unit tests for :class:`NudgeQueue`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.nudge_queue import TOOL_DRAIN, USER_DRAIN, NudgeQueue
|
||||
|
||||
|
||||
class TestEnqueueDrain:
|
||||
def test_enqueue_drain_fifo_order(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
q.enqueue("c", "3", "any")
|
||||
# Drain everything regardless of channel — preserves insertion order.
|
||||
out = q.drain({"user", "tool", "any"})
|
||||
# Drain returns ``(nudge_type, text, metadata)``; producers
|
||||
# without ``metadata`` see ``None`` in the third slot.
|
||||
assert out == [("a", "1", None), ("b", "2", None), ("c", "3", None)]
|
||||
assert len(q) == 0
|
||||
|
||||
def test_drain_filter_keeps_non_matching(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "x", "user")
|
||||
q.enqueue("b", "y", "tool")
|
||||
# Drain only user → tool entry stays.
|
||||
out = q.drain(USER_DRAIN)
|
||||
assert out == [("a", "x", None)]
|
||||
assert len(q) == 1
|
||||
# Now drain tool — gets the remaining entry.
|
||||
out = q.drain(TOOL_DRAIN)
|
||||
assert out == [("b", "y", None)]
|
||||
assert len(q) == 0
|
||||
|
||||
def test_any_channel_drains_on_either_seam(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("c", "z", "any")
|
||||
# User-seam drain pulls "any".
|
||||
assert q.drain(USER_DRAIN) == [("c", "z", None)]
|
||||
assert len(q) == 0
|
||||
# Re-enqueue and prove tool-seam also drains "any".
|
||||
q.enqueue("d", "w", "any")
|
||||
assert q.drain(TOOL_DRAIN) == [("d", "w", None)]
|
||||
assert len(q) == 0
|
||||
|
||||
def test_drain_empty_filter_no_op(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
# Empty filter drains nothing.
|
||||
assert q.drain(set()) == []
|
||||
assert len(q) == 1
|
||||
|
||||
def test_drain_empty_queue_returns_empty_list(self):
|
||||
q = NudgeQueue()
|
||||
# Fast-path: no items → no kept-deque allocation, just `[]`.
|
||||
assert q.drain(USER_DRAIN) == []
|
||||
assert q.drain({"user", "tool", "any"}) == []
|
||||
assert len(q) == 0
|
||||
|
||||
def test_drain_preserves_order_across_partial_drain(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
q.enqueue("c", "3", "user")
|
||||
q.enqueue("d", "4", "tool")
|
||||
# Drain user — should get "a" then "c" in order; "b","d" stay.
|
||||
assert q.drain({"user"}) == [("a", "1", None), ("c", "3", None)]
|
||||
# Tool drain follows insertion order on remaining.
|
||||
assert q.drain({"tool"}) == [("b", "2", None), ("d", "4", None)]
|
||||
|
||||
|
||||
class TestLenAndClear:
|
||||
def test_len_does_not_mutate(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
assert len(q) == 1
|
||||
assert len(q) == 1 # second call still 1; not consumed
|
||||
assert q.pending() == [("a", "1")]
|
||||
|
||||
def test_len_empty_is_zero(self):
|
||||
q = NudgeQueue()
|
||||
assert len(q) == 0
|
||||
|
||||
def test_clear_returns_count(self):
|
||||
q = NudgeQueue()
|
||||
assert q.clear() == 0
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
q.enqueue("c", "3", "any")
|
||||
assert q.clear() == 3
|
||||
assert len(q) == 0
|
||||
|
||||
def test_clear_empty_returns_zero(self):
|
||||
q = NudgeQueue()
|
||||
assert q.clear() == 0
|
||||
|
||||
|
||||
class TestDropOldestByType:
|
||||
def test_drop_oldest_by_type_removes_earliest_match(self):
|
||||
"""Drop the FIRST entry of the matching type; later matches stay."""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("other", "first", "any")
|
||||
q.enqueue("target", "older", "any")
|
||||
q.enqueue("target", "newer", "any")
|
||||
# "older" is the earliest target — drop it.
|
||||
assert q.drop_oldest_by_type("target") is True
|
||||
assert q.pending() == [("other", "first"), ("target", "newer")]
|
||||
|
||||
def test_drop_oldest_by_type_no_match_returns_false(self):
|
||||
"""Empty queue and unmatched-type cases both return False."""
|
||||
q = NudgeQueue()
|
||||
# Empty.
|
||||
assert q.drop_oldest_by_type("target") is False
|
||||
# Non-matching items only.
|
||||
q.enqueue("other", "1", "any")
|
||||
q.enqueue("other", "2", "tool")
|
||||
assert q.drop_oldest_by_type("target") is False
|
||||
# Queue is unaffected.
|
||||
assert q.pending() == [("other", "1"), ("other", "2")]
|
||||
|
||||
def test_drop_oldest_by_type_only_drops_one(self):
|
||||
"""Multiple matching entries → only the first is removed."""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("target", "1", "any")
|
||||
q.enqueue("target", "2", "any")
|
||||
q.enqueue("target", "3", "any")
|
||||
assert q.drop_oldest_by_type("target") is True
|
||||
assert q.pending() == [("target", "2"), ("target", "3")]
|
||||
|
||||
def test_drop_oldest_by_type_channel_filter(self):
|
||||
"""With ``channel`` set, drop walks only that channel. Pairs with
|
||||
:meth:`count_by_type(..., channel=...)` so producer-side soft caps
|
||||
operate on a consistent entry set.
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("target", "user-1", "user")
|
||||
q.enqueue("target", "any-1", "any")
|
||||
q.enqueue("target", "any-2", "any")
|
||||
# Drop the oldest "any"-channel target — leaves the user one
|
||||
# untouched even though it's earlier in insertion order.
|
||||
assert q.drop_oldest_by_type("target", channel="any") is True
|
||||
assert q.pending() == [
|
||||
("target", "user-1"),
|
||||
("target", "any-2"),
|
||||
]
|
||||
# And a channel with no matches returns False without touching
|
||||
# the queue.
|
||||
assert q.drop_oldest_by_type("target", channel="tool") is False
|
||||
assert q.pending() == [
|
||||
("target", "user-1"),
|
||||
("target", "any-2"),
|
||||
]
|
||||
|
||||
|
||||
class TestCapAtOrDropOldest:
|
||||
def test_below_cap_no_drop(self):
|
||||
q = NudgeQueue()
|
||||
for i in range(3):
|
||||
q.enqueue("target", f"t-{i}", "any")
|
||||
# 3 entries, cap=5 → no drop.
|
||||
assert q.cap_at_or_drop_oldest("target", 5, channel="any") is False
|
||||
assert q.count_by_type("target") == 3
|
||||
|
||||
def test_at_cap_drops_oldest(self):
|
||||
q = NudgeQueue()
|
||||
for i in range(5):
|
||||
q.enqueue("target", f"t-{i}", "any")
|
||||
# 5 entries, cap=5 → drop the oldest ("t-0"), leaving 4.
|
||||
assert q.cap_at_or_drop_oldest("target", 5, channel="any") is True
|
||||
remaining = q.pending()
|
||||
assert ("target", "t-0") not in remaining
|
||||
assert len(remaining) == 4
|
||||
assert remaining[0] == ("target", "t-1") # FIFO drop-oldest preserved
|
||||
|
||||
def test_above_cap_drops_only_one(self):
|
||||
q = NudgeQueue()
|
||||
for i in range(7):
|
||||
q.enqueue("target", f"t-{i}", "any")
|
||||
# 7 entries, cap=5 → drop only ONE per call (soft-cap regulates over time).
|
||||
assert q.cap_at_or_drop_oldest("target", 5, channel="any") is True
|
||||
assert q.count_by_type("target") == 6
|
||||
|
||||
def test_channel_filter_respected(self):
|
||||
q = NudgeQueue()
|
||||
for i in range(3):
|
||||
q.enqueue("target", f"any-{i}", "any")
|
||||
for i in range(3):
|
||||
q.enqueue("target", f"user-{i}", "user")
|
||||
# 3 "any"-channel entries; cap=3 on channel="any" → drop oldest "any" only.
|
||||
assert q.cap_at_or_drop_oldest("target", 3, channel="any") is True
|
||||
# User-channel entries untouched.
|
||||
assert q.count_by_type("target", channel="user") == 3
|
||||
assert q.count_by_type("target", channel="any") == 2
|
||||
|
||||
def test_other_types_ignored(self):
|
||||
q = NudgeQueue()
|
||||
for i in range(5):
|
||||
q.enqueue("other", f"o-{i}", "any")
|
||||
q.enqueue("target", "t-0", "any")
|
||||
# Only one "target" entry; cap=1 on "target" → drop it. "other"
|
||||
# entries are untouched even though queue holds 6 total.
|
||||
assert q.cap_at_or_drop_oldest("target", 1, channel="any") is True
|
||||
assert q.count_by_type("target") == 0
|
||||
assert q.count_by_type("other") == 5
|
||||
|
||||
def test_zero_or_negative_cap_no_op(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("target", "t-0", "any")
|
||||
assert q.cap_at_or_drop_oldest("target", 0, channel="any") is False
|
||||
assert q.cap_at_or_drop_oldest("target", -1, channel="any") is False
|
||||
assert q.count_by_type("target") == 1
|
||||
|
||||
def test_no_match_returns_false(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("other", "o-0", "any")
|
||||
assert q.cap_at_or_drop_oldest("target", 1, channel="any") is False
|
||||
assert q.count_by_type("other") == 1
|
||||
|
||||
|
||||
class TestCountByType:
|
||||
def test_count_by_type_no_channel(self):
|
||||
"""Count across all channels with ``channel=None``."""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("target", "1", "user")
|
||||
q.enqueue("other", "x", "any")
|
||||
q.enqueue("target", "2", "any")
|
||||
q.enqueue("target", "3", "tool")
|
||||
assert q.count_by_type("target") == 3
|
||||
assert q.count_by_type("other") == 1
|
||||
assert q.count_by_type("missing") == 0
|
||||
|
||||
def test_count_by_type_with_channel_filter(self):
|
||||
"""Filter narrows the count to one channel — used by producer-side
|
||||
soft caps that pair with ``drop_oldest_by_type(..., channel=...)``.
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("target", "u-1", "user")
|
||||
q.enqueue("target", "a-1", "any")
|
||||
q.enqueue("target", "a-2", "any")
|
||||
q.enqueue("target", "t-1", "tool")
|
||||
assert q.count_by_type("target", channel="any") == 2
|
||||
assert q.count_by_type("target", channel="user") == 1
|
||||
assert q.count_by_type("target", channel="tool") == 1
|
||||
|
||||
def test_count_by_type_empty_queue(self):
|
||||
q = NudgeQueue()
|
||||
assert q.count_by_type("anything") == 0
|
||||
assert q.count_by_type("anything", channel="any") == 0
|
||||
|
||||
|
||||
class TestPending:
|
||||
def test_pending_no_filter_returns_all_in_order(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
q.enqueue("c", "3", "any")
|
||||
# All three, in insertion order, as (nudge_type, text) tuples.
|
||||
assert q.pending() == [("a", "1"), ("b", "2"), ("c", "3")]
|
||||
|
||||
def test_pending_channel_filter(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
q.enqueue("c", "3", "user")
|
||||
q.enqueue("d", "4", "any")
|
||||
assert q.pending("user") == [("a", "1"), ("c", "3")]
|
||||
assert q.pending("tool") == [("b", "2")]
|
||||
assert q.pending("any") == [("d", "4")]
|
||||
|
||||
def test_pending_does_not_mutate(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
# Two pending calls return same content; nothing consumed.
|
||||
first = q.pending()
|
||||
second = q.pending()
|
||||
assert first == second
|
||||
assert len(q) == 2
|
||||
|
||||
|
||||
class TestMetadata:
|
||||
"""Producer-supplied ``metadata`` rides alongside ``(type, text)`` on
|
||||
drain. Today only ``watch_triggered`` populates it; the wire shape
|
||||
accommodates future producers (e.g. structured tool_error context)
|
||||
without another schema bump.
|
||||
"""
|
||||
|
||||
def test_drain_returns_metadata_when_set(self):
|
||||
q = NudgeQueue()
|
||||
meta = {"watch_name": "w1", "command": "ls", "poll_count": 2}
|
||||
q.enqueue("watch_triggered", "$ ls\nfile.txt", "any", metadata=meta)
|
||||
out = q.drain({"any"})
|
||||
assert out == [("watch_triggered", "$ ls\nfile.txt", meta)]
|
||||
|
||||
def test_drain_returns_none_when_metadata_unset(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("idle_children", "kids", "any") # no metadata kwarg
|
||||
out = q.drain({"any"})
|
||||
assert out == [("idle_children", "kids", None)]
|
||||
|
||||
def test_pending_with_metadata_projects_third_field(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("watch_triggered", "out", "any", metadata={"watch_name": "w"})
|
||||
snapshot = q.pending_with_metadata()
|
||||
assert snapshot == [
|
||||
("a", "1", None),
|
||||
("watch_triggered", "out", {"watch_name": "w"}),
|
||||
]
|
||||
# ``pending`` (without metadata) keeps the legacy 2-tuple shape.
|
||||
assert q.pending() == [("a", "1"), ("watch_triggered", "out")]
|
||||
|
||||
def test_metadata_survives_partial_drain(self):
|
||||
"""A ``user``-channel drain leaves an unaffected ``tool``-channel
|
||||
entry — its metadata must still be present on the next drain."""
|
||||
q = NudgeQueue()
|
||||
q.enqueue("user_thing", "u", "user")
|
||||
q.enqueue("watch_triggered", "w-out", "tool", metadata={"watch_name": "w1"})
|
||||
# User drain doesn't touch the tool entry.
|
||||
assert q.drain({"user"}) == [("user_thing", "u", None)]
|
||||
# Tool drain still has the metadata.
|
||||
assert q.drain({"tool"}) == [
|
||||
("watch_triggered", "w-out", {"watch_name": "w1"}),
|
||||
]
|
||||
|
||||
def test_metadata_with_valid_until_predicate(self):
|
||||
"""Metadata + ``valid_until`` co-exist on the same entry; the
|
||||
predicate gate runs as before, and on a True result the metadata
|
||||
rides the drained tuple.
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
q.enqueue(
|
||||
"watch_triggered",
|
||||
"w-out",
|
||||
"any",
|
||||
valid_until=lambda: True,
|
||||
metadata={"watch_name": "w1", "is_final": True},
|
||||
)
|
||||
out = q.drain({"any"})
|
||||
assert out == [
|
||||
("watch_triggered", "w-out", {"watch_name": "w1", "is_final": True}),
|
||||
]
|
||||
|
||||
|
||||
class TestHasPending:
|
||||
def test_has_pending_returns_false_on_empty_queue(self):
|
||||
q = NudgeQueue()
|
||||
assert q.has_pending({"user", "any"}) is False
|
||||
assert q.has_pending({"tool"}) is False
|
||||
|
||||
def test_has_pending_short_circuits_on_first_match(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "tool")
|
||||
q.enqueue("b", "2", "user")
|
||||
# First entry doesn't match, second does — true after walking 2.
|
||||
assert q.has_pending({"user"}) is True
|
||||
|
||||
def test_has_pending_returns_false_when_no_match(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "tool")
|
||||
q.enqueue("b", "2", "tool")
|
||||
assert q.has_pending({"user", "any"}) is False
|
||||
|
||||
def test_has_pending_matches_any_channel(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "any")
|
||||
# USER_DRAIN-shaped filter pulls "any" entries.
|
||||
assert q.has_pending(USER_DRAIN) is True
|
||||
# TOOL_DRAIN-shaped filter also pulls "any" entries.
|
||||
assert q.has_pending(TOOL_DRAIN) is True
|
||||
|
||||
def test_has_pending_does_not_mutate(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "user")
|
||||
q.enqueue("b", "2", "tool")
|
||||
before = q.pending()
|
||||
q.has_pending({"user"})
|
||||
q.has_pending({"tool"})
|
||||
q.has_pending(set())
|
||||
assert q.pending() == before
|
||||
|
||||
|
||||
class TestValidation:
|
||||
def test_invalid_channel_raises(self):
|
||||
q = NudgeQueue()
|
||||
with pytest.raises(ValueError, match="channel"):
|
||||
q.enqueue("a", "1", "wake") # type: ignore[arg-type]
|
||||
with pytest.raises(ValueError):
|
||||
q.enqueue("b", "2", "") # type: ignore[arg-type]
|
||||
# Queue is unaffected by the failed enqueues.
|
||||
assert len(q) == 0
|
||||
|
||||
def test_channel_is_required(self):
|
||||
q = NudgeQueue()
|
||||
# No default — caller MUST pick a seam consciously.
|
||||
with pytest.raises(TypeError):
|
||||
q.enqueue("a", "1") # type: ignore[call-arg]
|
||||
|
||||
|
||||
class TestValidUntil:
|
||||
"""``valid_until`` predicate: drain re-checks freshness; falsy /
|
||||
raising predicates drop the entry without delivery.
|
||||
"""
|
||||
|
||||
def test_valid_until_true_delivers(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "any", valid_until=lambda: True)
|
||||
out = q.drain({"any"})
|
||||
assert out == [("a", "1", None)]
|
||||
|
||||
def test_valid_until_false_drops_silently(self):
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "any", valid_until=lambda: False)
|
||||
out = q.drain({"any"})
|
||||
assert out == []
|
||||
# Already removed from queue (drain partition removes BEFORE
|
||||
# predicate check — falsy doesn't return to queue).
|
||||
assert len(q) == 0
|
||||
|
||||
def test_valid_until_exception_drops_silently(self):
|
||||
q = NudgeQueue()
|
||||
|
||||
def boom() -> bool:
|
||||
raise RuntimeError("predicate crash")
|
||||
|
||||
q.enqueue("a", "1", "any", valid_until=boom)
|
||||
out = q.drain({"any"})
|
||||
assert out == []
|
||||
# Crash-on-predicate is treated as "no longer valid" — drop, not propagate.
|
||||
assert len(q) == 0
|
||||
|
||||
def test_valid_until_evaluated_outside_lock(self):
|
||||
"""The predicate may do non-trivial work (e.g. storage I/O)
|
||||
without blocking other producers. Verify the predicate runs
|
||||
outside the queue's internal lock by enqueueing from inside
|
||||
the predicate — would deadlock if the lock was still held.
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
|
||||
def reentrant() -> bool:
|
||||
# If the lock is held during predicate eval, this enqueue
|
||||
# would block forever (RLock would let it through, but the
|
||||
# queue uses a plain Lock).
|
||||
q.enqueue("inner", "from-predicate", "any", valid_until=lambda: True)
|
||||
return True
|
||||
|
||||
q.enqueue("outer", "1", "any", valid_until=reentrant)
|
||||
out = q.drain({"any"})
|
||||
# Outer's predicate ran outside the lock, enqueued "inner";
|
||||
# outer's True return delivered "outer". "inner" was enqueued
|
||||
# AFTER the partition snapshot, so it stays in the queue.
|
||||
assert out == [("outer", "1", None)]
|
||||
assert q.pending() == [("inner", "from-predicate")]
|
||||
|
||||
def test_valid_until_only_evaluated_for_matching_channel(self):
|
||||
"""A non-matching entry's predicate must NOT fire — that would
|
||||
be wasted work (or worse, a side-effecting predicate would run
|
||||
when the entry is supposed to stay queued).
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
calls = []
|
||||
|
||||
def track() -> bool:
|
||||
calls.append(1)
|
||||
return True
|
||||
|
||||
# Tool-channel entry; we drain user-channel. Predicate must not run.
|
||||
q.enqueue("a", "1", "tool", valid_until=track)
|
||||
q.drain({"user", "any"})
|
||||
assert calls == []
|
||||
# Entry stays queued.
|
||||
assert q.pending("tool") == [("a", "1")]
|
||||
|
||||
def test_valid_until_default_none_always_delivers(self):
|
||||
# No predicate → entry behaves identically to pre-PR-3 entries.
|
||||
q = NudgeQueue()
|
||||
q.enqueue("a", "1", "any") # no valid_until kwarg
|
||||
assert q.drain({"any"}) == [("a", "1", None)]
|
||||
|
||||
|
||||
class TestConcurrency:
|
||||
def test_concurrent_enqueue_drain_no_loss(self):
|
||||
"""16 producer threads × 64 nudges = 1024 total; one consumer
|
||||
drains in a loop until producers finish + queue empty. Verify
|
||||
every produced item is observed exactly once.
|
||||
"""
|
||||
q = NudgeQueue()
|
||||
producers = 16
|
||||
per_producer = 64
|
||||
total = producers * per_producer
|
||||
|
||||
produced: set[tuple[str, str]] = set()
|
||||
produced_lock = threading.Lock()
|
||||
observed: list[tuple[str, str]] = []
|
||||
observed_lock = threading.Lock()
|
||||
done_event = threading.Event()
|
||||
|
||||
def produce(pid: int) -> None:
|
||||
for i in range(per_producer):
|
||||
key = (f"p{pid}", f"i{i}")
|
||||
with produced_lock:
|
||||
produced.add(key)
|
||||
q.enqueue(key[0], key[1], "user")
|
||||
|
||||
def consume() -> None:
|
||||
while not done_event.is_set() or len(q) > 0:
|
||||
drained = q.drain({"user"})
|
||||
if drained:
|
||||
with observed_lock:
|
||||
# Drop the trailing ``metadata`` slot — every
|
||||
# entry here was enqueued without metadata, so
|
||||
# the comparison set / count matches the produced
|
||||
# ``(type, text)`` shape.
|
||||
observed.extend((nt, txt) for nt, txt, _meta in drained)
|
||||
|
||||
consumer = threading.Thread(target=consume, daemon=True)
|
||||
consumer.start()
|
||||
threads = [threading.Thread(target=produce, args=(i,)) for i in range(producers)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
done_event.set()
|
||||
consumer.join(timeout=5.0)
|
||||
assert not consumer.is_alive(), "consumer didn't finish in time"
|
||||
|
||||
# Every produced key observed; no duplicates.
|
||||
assert set(observed) == produced
|
||||
assert len(observed) == total
|
||||
assert len(q) == 0
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Direct tests for the shared SSRF helpers in :mod:`turnstone.core.oauth_ssrf`.
|
||||
|
||||
The OIDC test suite already exercises these via the OIDC adapter
|
||||
(``OIDCError`` re-raises). This file pins the canonical
|
||||
:class:`OAuthSSRFError` exception so callers that don't go through OIDC
|
||||
(notably ``mcp_oauth``) can rely on a stable contract.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import urllib.parse
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.oauth_ssrf import (
|
||||
OAuthSSRFError,
|
||||
effective_port,
|
||||
is_localhost,
|
||||
validate_discovered_endpoint,
|
||||
validate_url_no_ssrf,
|
||||
)
|
||||
|
||||
|
||||
class TestIsLocalhost:
|
||||
def test_loopback_names(self) -> None:
|
||||
assert is_localhost("localhost")
|
||||
assert is_localhost("127.0.0.1")
|
||||
assert is_localhost("::1")
|
||||
assert is_localhost("foo.localhost")
|
||||
|
||||
def test_non_loopback(self) -> None:
|
||||
assert not is_localhost("example.com")
|
||||
assert not is_localhost("internal.corp")
|
||||
|
||||
|
||||
class TestEffectivePort:
|
||||
def test_explicit_port(self) -> None:
|
||||
p = urllib.parse.urlparse("https://idp.example.com:9443/foo")
|
||||
assert effective_port(p) == 9443
|
||||
|
||||
def test_default_https(self) -> None:
|
||||
p = urllib.parse.urlparse("https://idp.example.com/foo")
|
||||
assert effective_port(p) == 443
|
||||
|
||||
def test_default_http(self) -> None:
|
||||
p = urllib.parse.urlparse("http://idp.example.com/foo")
|
||||
assert effective_port(p) == 80
|
||||
|
||||
def test_unknown_scheme(self) -> None:
|
||||
p = urllib.parse.urlparse("ftp://idp.example.com/foo")
|
||||
assert effective_port(p) is None
|
||||
|
||||
|
||||
class TestValidateUrlNoSSRF:
|
||||
_PUBLIC_ADDR = [(2, 1, 6, "", ("93.184.216.34", 0))]
|
||||
_PRIVATE_ADDR = [(2, 1, 6, "", ("10.0.0.1", 0))]
|
||||
_LOOPBACK_ADDR = [(2, 1, 6, "", ("127.0.0.1", 0))]
|
||||
|
||||
def test_valid_https(self) -> None:
|
||||
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
|
||||
parsed = validate_url_no_ssrf("https://idp.example.com/foo", allow_http=False)
|
||||
assert parsed.scheme == "https"
|
||||
assert parsed.hostname == "idp.example.com"
|
||||
|
||||
def test_rejects_http_when_not_allowed(self) -> None:
|
||||
with pytest.raises(OAuthSSRFError, match="must use HTTPS"):
|
||||
validate_url_no_ssrf("http://idp.example.com", allow_http=False)
|
||||
|
||||
def test_allows_http_localhost_with_flag(self) -> None:
|
||||
with patch("socket.getaddrinfo", return_value=self._LOOPBACK_ADDR):
|
||||
validate_url_no_ssrf("http://localhost:8080", allow_http=True)
|
||||
|
||||
def test_rejects_http_non_localhost_even_with_flag(self) -> None:
|
||||
with pytest.raises(OAuthSSRFError, match="must use HTTPS"):
|
||||
validate_url_no_ssrf("http://idp.example.com", allow_http=True)
|
||||
|
||||
def test_rejects_userinfo(self) -> None:
|
||||
with pytest.raises(OAuthSSRFError, match="embedded credentials"):
|
||||
validate_url_no_ssrf("https://user:pass@idp.example.com", allow_http=False)
|
||||
|
||||
def test_rejects_private_address(self) -> None:
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=self._PRIVATE_ADDR),
|
||||
pytest.raises(OAuthSSRFError, match="non-public address"),
|
||||
):
|
||||
validate_url_no_ssrf("https://corp.example.com", allow_http=False)
|
||||
|
||||
def test_rejects_unresolvable(self) -> None:
|
||||
import socket
|
||||
|
||||
with (
|
||||
patch("socket.getaddrinfo", side_effect=socket.gaierror("fail")),
|
||||
pytest.raises(OAuthSSRFError, match="cannot be resolved"),
|
||||
):
|
||||
validate_url_no_ssrf("https://no.such.host.invalid", allow_http=False)
|
||||
|
||||
|
||||
class TestValidateDiscoveredEndpoint:
|
||||
_PUBLIC_ADDR = [(2, 1, 6, "", ("93.184.216.34", 0))]
|
||||
|
||||
def test_same_origin_passes(self) -> None:
|
||||
issuer = urllib.parse.urlparse("https://idp.example.com")
|
||||
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
|
||||
validate_discovered_endpoint(
|
||||
"https://idp.example.com/token",
|
||||
issuer,
|
||||
allow_http=False,
|
||||
trusted_endpoint_hosts=frozenset(),
|
||||
)
|
||||
|
||||
def test_third_party_host_rejected(self) -> None:
|
||||
issuer = urllib.parse.urlparse("https://idp.example.com")
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR),
|
||||
pytest.raises(OAuthSSRFError, match="not trusted"),
|
||||
):
|
||||
validate_discovered_endpoint(
|
||||
"https://attacker.example.com/token",
|
||||
issuer,
|
||||
allow_http=False,
|
||||
trusted_endpoint_hosts=frozenset(),
|
||||
)
|
||||
|
||||
def test_trusted_endpoint_host_passes(self) -> None:
|
||||
issuer = urllib.parse.urlparse("https://idp.example.com")
|
||||
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
|
||||
validate_discovered_endpoint(
|
||||
"https://shard.example.com/token",
|
||||
issuer,
|
||||
allow_http=False,
|
||||
trusted_endpoint_hosts=frozenset({"shard.example.com"}),
|
||||
)
|
||||
|
||||
def test_known_google_alias_passes(self) -> None:
|
||||
"""The hard-coded Google alias map covers oauth2.googleapis.com."""
|
||||
issuer = urllib.parse.urlparse("https://accounts.google.com")
|
||||
with patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR):
|
||||
validate_discovered_endpoint(
|
||||
"https://oauth2.googleapis.com/token",
|
||||
issuer,
|
||||
allow_http=False,
|
||||
trusted_endpoint_hosts=frozenset(),
|
||||
)
|
||||
|
||||
def test_scheme_mismatch_rejected(self) -> None:
|
||||
# When the issuer is http://localhost (allow_http=True), an
|
||||
# https:// endpoint must still be rejected as a scheme mismatch.
|
||||
issuer = urllib.parse.urlparse("http://localhost:8080")
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=[(2, 1, 6, "", ("127.0.0.1", 0))]),
|
||||
pytest.raises(OAuthSSRFError, match="scheme"),
|
||||
):
|
||||
validate_discovered_endpoint(
|
||||
"https://localhost:8080/token",
|
||||
issuer,
|
||||
allow_http=True,
|
||||
trusted_endpoint_hosts=frozenset(),
|
||||
)
|
||||
|
||||
def test_port_mismatch_rejected(self) -> None:
|
||||
issuer = urllib.parse.urlparse("https://idp.example.com")
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=self._PUBLIC_ADDR),
|
||||
pytest.raises(OAuthSSRFError, match="port"),
|
||||
):
|
||||
validate_discovered_endpoint(
|
||||
"https://idp.example.com:9443/token",
|
||||
issuer,
|
||||
allow_http=False,
|
||||
trusted_endpoint_hosts=frozenset(),
|
||||
)
|
||||
+1432
-149
File diff suppressed because it is too large
Load Diff
+259
-45
@@ -18,10 +18,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.routing import Mount, Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
from tests.conftest import make_oidc_test_config as _make_oidc_config
|
||||
from turnstone.console.server import (
|
||||
admin_delete_oidc_identity,
|
||||
admin_list_oidc_identities,
|
||||
@@ -32,34 +29,12 @@ from turnstone.core.auth import (
|
||||
handle_oidc_authorize,
|
||||
handle_oidc_callback,
|
||||
)
|
||||
from turnstone.core.oidc import OIDCConfig, OIDCError
|
||||
from turnstone.core.oidc import OIDCConfig, OIDCError, OIDCKeyNotFoundError
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_oidc_config(**overrides: Any) -> OIDCConfig:
|
||||
"""Build a test OIDCConfig with sensible defaults."""
|
||||
defaults: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"issuer": "https://idp.example.com",
|
||||
"client_id": "my-client",
|
||||
"client_secret": "my-secret",
|
||||
"scopes": "openid email profile",
|
||||
"provider_name": "TestIDP",
|
||||
"role_claim": "",
|
||||
"role_map": {},
|
||||
"password_enabled": True,
|
||||
"authorization_endpoint": "https://idp.example.com/authorize",
|
||||
"token_endpoint": "https://idp.example.com/token",
|
||||
"userinfo_endpoint": "https://idp.example.com/userinfo",
|
||||
"jwks_uri": "https://idp.example.com/.well-known/jwks.json",
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return OIDCConfig(**defaults)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Thin handler wrappers — match the pattern used in server.py / console
|
||||
@@ -290,9 +265,9 @@ class TestOIDCCallback:
|
||||
) -> None:
|
||||
storage.create_oidc_pending_state(state, nonce, code_verifier, audience)
|
||||
|
||||
@patch("turnstone.core.oidc.provision_oidc_user")
|
||||
@patch("turnstone.core.oidc.validate_id_token")
|
||||
@patch("turnstone.core.oidc.exchange_code", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_happy_path(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
@@ -379,7 +354,7 @@ class TestOIDCCallback:
|
||||
assert resp.status_code == 302
|
||||
assert "Login+session+expired" in resp.headers["location"]
|
||||
|
||||
@patch("turnstone.core.oidc.exchange_code", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_code_exchange_failure(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
@@ -396,8 +371,8 @@ class TestOIDCCallback:
|
||||
assert resp.status_code == 302
|
||||
assert "Authentication+failed" in resp.headers["location"]
|
||||
|
||||
@patch("turnstone.core.oidc.validate_id_token")
|
||||
@patch("turnstone.core.oidc.exchange_code", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_token_validation_failure(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
@@ -416,10 +391,10 @@ class TestOIDCCallback:
|
||||
assert resp.status_code == 302
|
||||
assert "Authentication+failed" in resp.headers["location"]
|
||||
|
||||
@patch("turnstone.core.oidc.provision_oidc_user")
|
||||
@patch("turnstone.core.oidc.validate_id_token")
|
||||
@patch("turnstone.core.oidc.fetch_jwks", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.oidc.exchange_code", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.fetch_jwks", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_jwks_key_rotation_retry(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
@@ -429,13 +404,13 @@ class TestOIDCCallback:
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""First validate raises 'kid not found in JWKS', fetch_jwks retried, second validate succeeds."""
|
||||
"""First validate raises kid-not-found, fetch_jwks retried, second validate succeeds."""
|
||||
self._seed_pending_state(storage)
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
|
||||
# First call raises kid-not-found; second call (after JWKS refresh) succeeds
|
||||
mock_validate.side_effect = [
|
||||
OIDCError("Signing key 'new-kid' not found in JWKS"),
|
||||
OIDCKeyNotFoundError("Signing key 'new-kid' not found in JWKS"),
|
||||
{"sub": "user123", "email": "u@example.com", "nonce": "test-nonce"},
|
||||
]
|
||||
mock_fetch_jwks.return_value = {"keys": [{"kid": "new-kid", "kty": "RSA"}]}
|
||||
@@ -450,9 +425,61 @@ class TestOIDCCallback:
|
||||
mock_fetch_jwks.assert_called_once()
|
||||
assert mock_validate.call_count == 2
|
||||
|
||||
@patch("turnstone.core.oidc.provision_oidc_user")
|
||||
@patch("turnstone.core.oidc.validate_id_token")
|
||||
@patch("turnstone.core.oidc.exchange_code", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.fetch_jwks", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_callback_uses_keynotfound_for_jwks_retry(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
mock_fetch_jwks: AsyncMock,
|
||||
mock_validate: Any,
|
||||
mock_provision: Any,
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""Retry path keys off the OIDCKeyNotFoundError type, not message substring."""
|
||||
self._seed_pending_state(storage)
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
|
||||
# First raises subclass; rephrased message must not affect retry behaviour.
|
||||
mock_validate.side_effect = [
|
||||
OIDCKeyNotFoundError("rotated key absent from cached set"),
|
||||
{"sub": "user123", "email": "u@example.com", "nonce": "test-nonce"},
|
||||
]
|
||||
mock_fetch_jwks.return_value = {"keys": [{"kid": "new-kid", "kty": "RSA"}]}
|
||||
mock_provision.return_value = {"user_id": "test-admin", "username": "testadmin"}
|
||||
|
||||
resp = authorize_client.get(
|
||||
"/v1/api/auth/oidc/callback?code=authcode&state=valid-state",
|
||||
follow_redirects=False,
|
||||
)
|
||||
assert resp.status_code == 302
|
||||
assert "oidc_success=1" in resp.headers["location"]
|
||||
mock_fetch_jwks.assert_called_once()
|
||||
assert mock_validate.call_count == 2
|
||||
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_callback_returns_authentication_failed_on_missing_id_token(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""A token endpoint response without id_token must redirect with auth-failed."""
|
||||
self._seed_pending_state(storage)
|
||||
mock_exchange.return_value = {"access_token": "x"}
|
||||
|
||||
resp = authorize_client.get(
|
||||
"/v1/api/auth/oidc/callback?code=authcode&state=valid-state",
|
||||
follow_redirects=False,
|
||||
)
|
||||
assert resp.status_code == 302
|
||||
assert "oidc_error=Authentication+failed" in resp.headers["location"]
|
||||
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_no_users_after_oidc_success_redirects_setup(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
@@ -507,6 +534,193 @@ class TestOIDCCallback:
|
||||
assert "oidc_error" in resp.headers["location"]
|
||||
assert "Too+many" in resp.headers["location"]
|
||||
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_setup_gate_uses_count_users_not_full_scan(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
mock_validate: Any,
|
||||
mock_provision: Any,
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""Callback's setup-complete gate must call count_users, not list_users."""
|
||||
from unittest.mock import patch as obj_patch
|
||||
|
||||
self._seed_pending_state(storage)
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
mock_validate.return_value = {
|
||||
"sub": "u1",
|
||||
"email": "u@example.com",
|
||||
"nonce": "test-nonce",
|
||||
}
|
||||
mock_provision.return_value = {"user_id": "test-admin", "username": "testadmin"}
|
||||
|
||||
with (
|
||||
obj_patch.object(storage, "count_users", wraps=storage.count_users) as count_spy,
|
||||
obj_patch.object(storage, "list_users", wraps=storage.list_users) as list_spy,
|
||||
):
|
||||
resp = authorize_client.get(
|
||||
"/v1/api/auth/oidc/callback?code=authcode&state=valid-state",
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
assert resp.status_code == 302
|
||||
assert "oidc_success=1" in resp.headers["location"]
|
||||
count_spy.assert_called_once_with()
|
||||
list_spy.assert_not_called()
|
||||
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_state_cleanup_is_gated(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
mock_validate: Any,
|
||||
mock_provision: Any,
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""Cleanup runs once per cleanup-interval window, not every callback."""
|
||||
from unittest.mock import patch as obj_patch
|
||||
|
||||
# First call seeds the cleanup timestamp; subsequent calls within
|
||||
# _OIDC_STATE_CLEANUP_INTERVAL_S must NOT trigger cleanup again.
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
mock_validate.return_value = {
|
||||
"sub": "u1",
|
||||
"email": "u@example.com",
|
||||
"nonce": "test-nonce",
|
||||
}
|
||||
mock_provision.return_value = {"user_id": "test-admin", "username": "testadmin"}
|
||||
|
||||
with obj_patch.object(
|
||||
storage, "cleanup_expired_oidc_states", wraps=storage.cleanup_expired_oidc_states
|
||||
) as cleanup_spy:
|
||||
for state in ("s1", "s2", "s3"):
|
||||
self._seed_pending_state(storage, state=state, nonce="test-nonce")
|
||||
authorize_client.get(
|
||||
f"/v1/api/auth/oidc/callback?code=c&state={state}",
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
assert cleanup_spy.call_count == 1
|
||||
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_callback_uses_pending_audience_not_handler_audience(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
mock_validate: Any,
|
||||
mock_provision: Any,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""JWT ``aud`` claim must come from the audience stored at /authorize,
|
||||
not the audience the callback handler was invoked with.
|
||||
|
||||
Regression for the cross-service audience-confusion concern: a
|
||||
login flow opened against the server (audience ``"turnstone-server"``)
|
||||
must not be silently re-targeted to ``"turnstone-console"`` when
|
||||
the callback runs through the console's handler wrapper.
|
||||
"""
|
||||
import jwt as pyjwt
|
||||
|
||||
# Seed pending state with the SERVER audience.
|
||||
storage.create_oidc_pending_state(
|
||||
"audience-state",
|
||||
"audience-nonce",
|
||||
"audience-verifier",
|
||||
"turnstone-server",
|
||||
)
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
mock_validate.return_value = {
|
||||
"sub": "user-aud",
|
||||
"email": "u@example.com",
|
||||
"nonce": "audience-nonce",
|
||||
}
|
||||
mock_provision.return_value = {"user_id": "test-admin", "username": "testadmin"}
|
||||
|
||||
# Wire a callback bound to the CONSOLE audience. After bug-3 the
|
||||
# stored audience must take precedence.
|
||||
async def _console_callback(request: Request) -> Response:
|
||||
return await handle_oidc_callback(request, "turnstone-console")
|
||||
|
||||
jwt_secret = "test-jwt-secret-key-padded-32b!!"
|
||||
app = Starlette(
|
||||
routes=[Mount("/v1", routes=[Route("/api/auth/oidc/callback", _console_callback)])]
|
||||
)
|
||||
app.state.oidc_config = _make_oidc_config()
|
||||
app.state.auth_storage = storage
|
||||
app.state.jwt_secret = jwt_secret
|
||||
app.state.jwks_data = {"keys": []}
|
||||
app.state.login_limiter = None
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
resp = client.get(
|
||||
"/v1/api/auth/oidc/callback?code=authcode&state=audience-state",
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
assert resp.status_code == 302
|
||||
assert "oidc_success=1" in resp.headers["location"]
|
||||
|
||||
# Extract the JWT from the Set-Cookie header and decode it.
|
||||
set_cookie = resp.headers["set-cookie"]
|
||||
cookie_kv = set_cookie.split(";", 1)[0]
|
||||
name, _, token = cookie_kv.partition("=")
|
||||
assert name == "turnstone_auth"
|
||||
assert token
|
||||
|
||||
# Decoding without audience verification first to inspect the claim.
|
||||
claims = pyjwt.decode(
|
||||
token, jwt_secret, algorithms=["HS256"], options={"verify_aud": False}
|
||||
)
|
||||
assert claims["aud"] == "turnstone-server"
|
||||
assert claims["aud"] != "turnstone-console"
|
||||
|
||||
@patch("turnstone.core.auth.provision_oidc_user")
|
||||
@patch("turnstone.core.auth.validate_id_token")
|
||||
@patch("turnstone.core.auth.fetch_jwks", new_callable=AsyncMock)
|
||||
@patch("turnstone.core.auth.exchange_code", new_callable=AsyncMock)
|
||||
def test_jwks_refetch_dedup_when_kid_appears(
|
||||
self,
|
||||
mock_exchange: AsyncMock,
|
||||
mock_fetch_jwks: AsyncMock,
|
||||
mock_validate: Any,
|
||||
mock_provision: Any,
|
||||
authorize_client: TestClient,
|
||||
storage: SQLiteBackend,
|
||||
) -> None:
|
||||
"""If a concurrent caller already refreshed JWKS, second caller skips fetch."""
|
||||
from unittest.mock import patch as obj_patch
|
||||
|
||||
self._seed_pending_state(storage)
|
||||
mock_exchange.return_value = {"id_token": "fake.jwt.token"}
|
||||
mock_provision.return_value = {"user_id": "test-admin", "username": "testadmin"}
|
||||
|
||||
# First validate raises kid-not-found; second succeeds.
|
||||
mock_validate.side_effect = [
|
||||
OIDCKeyNotFoundError("Signing key 'k-rotated' not found"),
|
||||
{"sub": "u1", "email": "u@example.com", "nonce": "test-nonce"},
|
||||
]
|
||||
|
||||
# Pre-populate the JWKS cache so the rotated kid is already
|
||||
# present — analog of a concurrent caller having won the lock.
|
||||
# The retry path must short-circuit and skip the network fetch.
|
||||
authorize_client.app.state.jwks_data = {"keys": [{"kid": "k-rotated", "kty": "RSA"}]}
|
||||
|
||||
with obj_patch("jwt.get_unverified_header", return_value={"kid": "k-rotated"}):
|
||||
resp = authorize_client.get(
|
||||
"/v1/api/auth/oidc/callback?code=authcode&state=valid-state",
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
assert resp.status_code == 302
|
||||
assert "oidc_success=1" in resp.headers["location"]
|
||||
mock_fetch_jwks.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Admin OIDC identity endpoint tests
|
||||
|
||||
@@ -6,6 +6,85 @@ import time
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.storage import StorageConflictError
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Atomic OIDC user provisioning
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCreateOIDCUser:
|
||||
def test_create_oidc_user_success(self, db):
|
||||
"""Both rows present after one atomic call."""
|
||||
db.create_oidc_user(
|
||||
user_id="u-new",
|
||||
username="alice",
|
||||
display_name="Alice",
|
||||
password_hash="!oidc",
|
||||
issuer="https://idp.example.com",
|
||||
subject="sub-1",
|
||||
email="alice@example.com",
|
||||
)
|
||||
|
||||
user = db.get_user("u-new")
|
||||
assert user is not None
|
||||
assert user["username"] == "alice"
|
||||
assert user["password_hash"] == "!oidc"
|
||||
|
||||
identity = db.get_oidc_identity("https://idp.example.com", "sub-1")
|
||||
assert identity is not None
|
||||
assert identity["user_id"] == "u-new"
|
||||
assert identity["email"] == "alice@example.com"
|
||||
|
||||
def test_create_oidc_user_username_conflict_rolls_back(self, db):
|
||||
"""Pre-existing username -> StorageConflictError; identity NOT inserted."""
|
||||
db.create_user("u-existing", "alice", "Alice", "$2b$12$hash")
|
||||
|
||||
with pytest.raises(StorageConflictError, match="username"):
|
||||
db.create_oidc_user(
|
||||
user_id="u-new",
|
||||
username="alice",
|
||||
display_name="Alice2",
|
||||
password_hash="!oidc",
|
||||
issuer="https://idp.example.com",
|
||||
subject="sub-1",
|
||||
email="alice2@example.com",
|
||||
)
|
||||
|
||||
# The new user_id row must not exist.
|
||||
assert db.get_user("u-new") is None
|
||||
# The identity row must not exist.
|
||||
assert db.get_oidc_identity("https://idp.example.com", "sub-1") is None
|
||||
# The pre-existing user is untouched.
|
||||
existing = db.get_user("u-existing")
|
||||
assert existing is not None
|
||||
assert existing["password_hash"] == "$2b$12$hash"
|
||||
|
||||
def test_create_oidc_user_identity_conflict_rolls_back(self, db):
|
||||
"""Pre-existing (issuer, subject) -> StorageConflictError; user row rolled back."""
|
||||
db.create_user("u-other", "other", "Other", "!oidc")
|
||||
db.create_oidc_identity("https://idp.example.com", "sub-1", "u-other", "other@example.com")
|
||||
|
||||
with pytest.raises(StorageConflictError, match="OIDC identity"):
|
||||
db.create_oidc_user(
|
||||
user_id="u-new",
|
||||
username="bob",
|
||||
display_name="Bob",
|
||||
password_hash="!oidc",
|
||||
issuer="https://idp.example.com",
|
||||
subject="sub-1",
|
||||
email="bob@example.com",
|
||||
)
|
||||
|
||||
# The candidate user row was rolled back.
|
||||
assert db.get_user("u-new") is None
|
||||
assert db.get_user_by_username("bob") is None
|
||||
# The pre-existing identity still points at the original user.
|
||||
identity = db.get_oidc_identity("https://idp.example.com", "sub-1")
|
||||
assert identity is not None
|
||||
assert identity["user_id"] == "u-other"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OIDC Identity CRUD
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -304,3 +383,217 @@ class TestOIDCPendingState:
|
||||
.where(oidc_pending_states.c.state == "state-cleanup")
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# count_users / find_existing_usernames
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCountUsers:
|
||||
def test_count_users_empty(self, db):
|
||||
assert db.count_users() == 0
|
||||
|
||||
def test_count_users_after_inserts(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h1")
|
||||
db.create_user("u2", "bob", "Bob", "h2")
|
||||
db.create_user("u3", "carol", "Carol", "h3")
|
||||
assert db.count_users() == 3
|
||||
|
||||
|
||||
class TestFindExistingUsernames:
|
||||
def test_empty_input_returns_empty_set(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h1")
|
||||
assert db.find_existing_usernames([]) == set()
|
||||
|
||||
def test_returns_subset_present_in_db(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h1")
|
||||
db.create_user("u2", "bob", "Bob", "h2")
|
||||
|
||||
existing = db.find_existing_usernames(["alice", "bob", "carol", "dave"])
|
||||
assert existing == {"alice", "bob"}
|
||||
|
||||
def test_no_matches_returns_empty_set(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h1")
|
||||
assert db.find_existing_usernames(["bob", "carol"]) == set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# replace_oidc_roles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestReplaceOIDCRoles:
|
||||
def _seed_role(self, db, role_id):
|
||||
db.create_role(role_id, role_id, role_id, "perm.read", False, "")
|
||||
|
||||
def test_inserts_added_roles(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
self._seed_role(db, "role-b")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a", "role-b"})
|
||||
|
||||
assert added == {"role-a", "role-b"}
|
||||
assert removed == set()
|
||||
roles = {r["role_id"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a", "role-b"}
|
||||
|
||||
def test_removes_stale_oidc_roles(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
self._seed_role(db, "role-b")
|
||||
db.assign_role("u1", "role-a", "oidc")
|
||||
db.assign_role("u1", "role-b", "oidc")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == {"role-b"}
|
||||
roles = {r["role_id"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a"}
|
||||
|
||||
def test_preserves_non_oidc_roles(self, db):
|
||||
"""Manually-assigned and oidc-default rows are NOT touched."""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-manual")
|
||||
self._seed_role(db, "role-default")
|
||||
self._seed_role(db, "role-oidc-old")
|
||||
db.assign_role("u1", "role-manual", "admin-ui")
|
||||
db.assign_role("u1", "role-default", "oidc-default")
|
||||
db.assign_role("u1", "role-oidc-old", "oidc")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", set())
|
||||
|
||||
# Only the oidc-assigned row was diffed
|
||||
assert added == set()
|
||||
assert removed == {"role-oidc-old"}
|
||||
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {
|
||||
"role-manual": "admin-ui",
|
||||
"role-default": "oidc-default",
|
||||
}
|
||||
|
||||
def test_no_op_when_desired_matches_current(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
db.assign_role("u1", "role-a", "oidc")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
assert {r["role_id"] for r in db.list_user_roles("u1")} == {"role-a"}
|
||||
|
||||
def test_empty_user_no_oidc_history(self, db):
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", set())
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
|
||||
def test_desired_role_blocked_by_admin_ui_assignment(self, db):
|
||||
"""Desired role already held via admin-ui: untouched, no PK conflict."""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
db.assign_role("u1", "role-a", "admin-ui")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a": "admin-ui"}
|
||||
|
||||
def test_desired_role_blocked_by_oidc_default_assignment(self, db):
|
||||
"""Desired role already held via oidc-default fallback: untouched."""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
db.assign_role("u1", "role-a", "oidc-default")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a": "oidc-default"}
|
||||
|
||||
def test_desired_role_added_alongside_blocked_role(self, db):
|
||||
"""Mixed case: one desired role is blocked (admin-ui), the other inserts cleanly."""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
self._seed_role(db, "role-b")
|
||||
db.assign_role("u1", "role-a", "admin-ui")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a", "role-b"})
|
||||
|
||||
assert added == {"role-b"}
|
||||
assert removed == set()
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a": "admin-ui", "role-b": "oidc"}
|
||||
|
||||
def test_revoke_only_oidc_assigned_roles(self, db):
|
||||
"""OIDC-assigned roles get revoked when not in desired; admin-ui rows survive."""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-manual")
|
||||
self._seed_role(db, "role-oidc-old")
|
||||
self._seed_role(db, "role-default")
|
||||
db.assign_role("u1", "role-manual", "admin-ui")
|
||||
db.assign_role("u1", "role-oidc-old", "oidc")
|
||||
db.assign_role("u1", "role-default", "oidc-default")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", set())
|
||||
|
||||
assert added == set()
|
||||
assert removed == {"role-oidc-old"}
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-manual": "admin-ui", "role-default": "oidc-default"}
|
||||
|
||||
def test_replace_oidc_roles_no_op_steady_state(self, db):
|
||||
"""Steady-state re-login: claims unchanged, function must short-circuit.
|
||||
|
||||
This pins the contract that drives the SQLite optimistic-read fast
|
||||
path — the common case (token refresh with identical role claims)
|
||||
must not acquire a write lock.
|
||||
"""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
self._seed_role(db, "role-b")
|
||||
db.assign_role("u1", "role-a", "oidc")
|
||||
db.assign_role("u1", "role-b", "oidc")
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a", "role-b"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
# All rows still oidc-assigned with identical membership.
|
||||
roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")}
|
||||
assert roles == {"role-a": "oidc", "role-b": "oidc"}
|
||||
|
||||
def test_replace_oidc_roles_returns_post_lock_diff(self, db):
|
||||
"""Returned (added, removed) reflects the post-lock state, not the optimistic read.
|
||||
|
||||
The SQLite implementation re-reads under the write lock to defend
|
||||
against races; the values returned must come from that re-read so
|
||||
callers (apply_role_mapping audit logs) see the actual transition
|
||||
that hit the table. Steady-state input must collapse to empty
|
||||
sets and leave row timestamps unchanged.
|
||||
"""
|
||||
db.create_user("u1", "alice", "Alice", "h")
|
||||
self._seed_role(db, "role-a")
|
||||
db.assign_role("u1", "role-a", "oidc")
|
||||
|
||||
before = db.list_user_roles("u1")
|
||||
assert len(before) == 1
|
||||
original_created = before[0]["assignment_created"]
|
||||
|
||||
added, removed = db.replace_oidc_roles("u1", {"role-a"})
|
||||
|
||||
assert added == set()
|
||||
assert removed == set()
|
||||
# No write occurred — the assignment row's timestamp is untouched.
|
||||
after = db.list_user_roles("u1")
|
||||
assert len(after) == 1
|
||||
assert after[0]["assignment_created"] == original_created
|
||||
|
||||
@@ -224,6 +224,25 @@ class TestOpenAIProvider:
|
||||
sanitize_messages([original])
|
||||
assert original["content"] is None
|
||||
|
||||
def test_sanitize_messages_strips_underscore_sibling_keys(self) -> None:
|
||||
"""Internal sibling metadata (``_reminders``, ``_reminders_delivered``,
|
||||
``_attachments_meta``, ``_provider_content``) must be stripped
|
||||
before the wire — the OpenAI-compat APIs reject unknown fields."""
|
||||
msgs = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "hi",
|
||||
"_reminders": [{"type": "correction", "text": "watch"}],
|
||||
"_reminders_delivered": True,
|
||||
"_attachments_meta": [{"kind": "image"}],
|
||||
}
|
||||
]
|
||||
result = sanitize_messages(msgs)
|
||||
assert result == [{"role": "user", "content": "hi"}]
|
||||
assert "_reminders" not in result[0]
|
||||
assert "_reminders_delivered" not in result[0]
|
||||
assert "_attachments_meta" not in result[0]
|
||||
|
||||
# -- sanitize_messages: orphan detection -----------------------------------
|
||||
|
||||
def test_sanitize_orphaned_tool_call_synthesized(self) -> None:
|
||||
@@ -1364,6 +1383,34 @@ class TestProviderFactory:
|
||||
p2 = create_provider("openai")
|
||||
assert p1 is p2
|
||||
|
||||
def test_create_provider_compat_responses_surface(self) -> None:
|
||||
"""openai-compatible + api_surface=responses returns the Responses provider."""
|
||||
from turnstone.core.providers import OpenAIResponsesProvider, create_provider
|
||||
|
||||
provider = create_provider("openai-compatible", api_surface="responses")
|
||||
assert isinstance(provider, OpenAIResponsesProvider)
|
||||
|
||||
def test_create_provider_compat_chat_surface_default(self) -> None:
|
||||
"""openai-compatible defaults to Chat Completions."""
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
for surface in (None, "", "chat"):
|
||||
provider = create_provider("openai-compatible", api_surface=surface)
|
||||
assert isinstance(provider, OpenAIChatCompletionsProvider)
|
||||
|
||||
def test_create_provider_invalid_api_surface(self) -> None:
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown api_surface"):
|
||||
create_provider("openai-compatible", api_surface="bogus")
|
||||
|
||||
def test_create_provider_openai_ignores_api_surface(self) -> None:
|
||||
"""Cloud OpenAI is always Responses regardless of api_surface."""
|
||||
from turnstone.core.providers import OpenAIResponsesProvider, create_provider
|
||||
|
||||
provider = create_provider("openai", api_surface="chat")
|
||||
assert isinstance(provider, OpenAIResponsesProvider)
|
||||
|
||||
# -- Google provider -------------------------------------------------------
|
||||
|
||||
def test_create_provider_google(self) -> None:
|
||||
|
||||
@@ -15,9 +15,26 @@ def _row(
|
||||
tc_id=None,
|
||||
pdata=None,
|
||||
tool_calls=None,
|
||||
source=None,
|
||||
reminders=None,
|
||||
):
|
||||
"""Build a 7-element conversation row tuple (id, role, ...)."""
|
||||
return (next(_row_ids), role, content, tool_name, tc_id, pdata, tool_calls)
|
||||
"""Build a 9-element conversation row tuple (id, role, ...).
|
||||
|
||||
Trailing ``source`` / ``reminders`` mirror the persisted twins of
|
||||
the in-memory ``_source`` / ``_reminders`` side-channels added in
|
||||
migration 050.
|
||||
"""
|
||||
return (
|
||||
next(_row_ids),
|
||||
role,
|
||||
content,
|
||||
tool_name,
|
||||
tc_id,
|
||||
pdata,
|
||||
tool_calls,
|
||||
source,
|
||||
reminders,
|
||||
)
|
||||
|
||||
|
||||
class TestAssistantWithToolCalls:
|
||||
|
||||
@@ -10,6 +10,7 @@ from __future__ import annotations
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import uuid
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -707,22 +708,26 @@ class TestQueuedAttachmentReservation:
|
||||
mgr.get.return_value = ws
|
||||
return ws, session
|
||||
|
||||
def _queue_with_attachment(self, client, mgr, ws_id: str, filename: str = "q.md"):
|
||||
def _reserve_attachment(self, client, mgr, ws_id: str, filename: str = "q.md"):
|
||||
"""Set up a reserved attachment for the busy-worker tests below.
|
||||
|
||||
The queue-with-attachments path was removed (queued user turns
|
||||
can't carry attachments — see ``AttachmentsNotQueueableError``),
|
||||
so the tests reserve directly via ``reserve_attachments`` to
|
||||
produce the same on-disk state without going through the
|
||||
rejected route path.
|
||||
"""
|
||||
from turnstone.core.memory import reserve_attachments
|
||||
|
||||
aid = _upload(client, ws_id, "userA", filename, b"Q", "text/markdown")
|
||||
ws, session = self._wire_busy_ws(mgr, ws_id)
|
||||
resp = client.post(
|
||||
f"/v1/api/workstreams/{ws_id}/send",
|
||||
json={"message": "queued", "attachment_ids": [aid]},
|
||||
headers=_auth("userA"),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["status"] == "queued"
|
||||
return aid, body["msg_id"], session
|
||||
msg_id = uuid.uuid4().hex
|
||||
reserve_attachments([aid], msg_id, ws_id, "userA")
|
||||
return aid, msg_id, session
|
||||
|
||||
def test_reserved_attachment_hidden_from_pending_listing(self, app_client):
|
||||
client, mgr = app_client
|
||||
aid, _mid, _session = self._queue_with_attachment(client, mgr, "ws-A")
|
||||
aid, _mid, _session = self._reserve_attachment(client, mgr, "ws-A")
|
||||
resp = client.get("/v1/api/workstreams/ws-A/attachments", headers=_auth("userA"))
|
||||
# Reserved attachment is not in the pending listing
|
||||
ids = [a["attachment_id"] for a in resp.json()["attachments"]]
|
||||
@@ -730,7 +735,7 @@ class TestQueuedAttachmentReservation:
|
||||
|
||||
def test_reserved_attachment_cannot_be_deleted(self, app_client):
|
||||
client, mgr = app_client
|
||||
aid, _mid, _session = self._queue_with_attachment(client, mgr, "ws-A")
|
||||
aid, _mid, _session = self._reserve_attachment(client, mgr, "ws-A")
|
||||
resp = client.delete(
|
||||
f"/v1/api/workstreams/ws-A/attachments/{aid}",
|
||||
headers=_auth("userA"),
|
||||
@@ -745,7 +750,7 @@ class TestQueuedAttachmentReservation:
|
||||
|
||||
def test_reserved_attachment_not_auto_consumed_by_later_send(self, app_client):
|
||||
client, mgr = app_client
|
||||
aid, _mid, session = self._queue_with_attachment(client, mgr, "ws-A")
|
||||
aid, _mid, session = self._reserve_attachment(client, mgr, "ws-A")
|
||||
|
||||
# Swap the busy worker for an idle one and capture the next
|
||||
# session.send call so we can assert on its attachment list.
|
||||
@@ -781,7 +786,7 @@ class TestQueuedAttachmentReservation:
|
||||
|
||||
def test_reserved_attachment_rejected_in_explicit_ids(self, app_client):
|
||||
client, mgr = app_client
|
||||
aid, _mid, session = self._queue_with_attachment(client, mgr, "ws-A")
|
||||
aid, _mid, session = self._reserve_attachment(client, mgr, "ws-A")
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
@@ -813,29 +818,26 @@ class TestQueuedAttachmentReservation:
|
||||
if atts is not None:
|
||||
assert aid not in [a.attachment_id for a in atts]
|
||||
|
||||
def test_dequeue_releases_reservation(self, app_client):
|
||||
def test_send_with_attachments_to_busy_worker_returns_attachments_busy(self, app_client):
|
||||
"""An attempt to attach mid-tool-call returns ``attachments_busy``;
|
||||
attachments stay pending so the client can retry once idle."""
|
||||
client, mgr = app_client
|
||||
aid, mid, session = self._queue_with_attachment(client, mgr, "ws-A")
|
||||
|
||||
# Cancel the queued message — DELETE /api/send with msg_id
|
||||
resp = client.request(
|
||||
"DELETE",
|
||||
aid = _upload(client, "ws-A", "userA", "x.md", b"X", "text/markdown")
|
||||
self._wire_busy_ws(mgr, "ws-A")
|
||||
resp = client.post(
|
||||
"/v1/api/workstreams/ws-A/send",
|
||||
json={"msg_id": mid},
|
||||
json={"message": "with file", "attachment_ids": [aid]},
|
||||
headers=_auth("userA"),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.json().get("status") == "removed"
|
||||
|
||||
# Attachment is back to pending — visible + deletable
|
||||
body = resp.json()
|
||||
assert body["status"] == "attachments_busy"
|
||||
assert body["attached_ids"] == []
|
||||
assert body["dropped_attachment_ids"] == [aid]
|
||||
# Reservation released — attachment is still pending and visible.
|
||||
resp = client.get("/v1/api/workstreams/ws-A/attachments", headers=_auth("userA"))
|
||||
ids = [a["attachment_id"] for a in resp.json()["attachments"]]
|
||||
assert aid in ids
|
||||
resp = client.delete(
|
||||
f"/v1/api/workstreams/ws-A/attachments/{aid}",
|
||||
headers=_auth("userA"),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
|
||||
class TestReserveThenDispatchRace:
|
||||
|
||||
+80
-35
@@ -89,6 +89,27 @@ class TestSuggestProfile:
|
||||
p = suggest_profile("vllm", "Google/GEMMA-4-31B-IT")
|
||||
assert p["capabilities"]["thinking_mode"] == "manual"
|
||||
|
||||
def test_vllm_mistral_medium_not_auto_suggested(self) -> None:
|
||||
"""Mistral medium falls back to the generic vLLM profile.
|
||||
|
||||
We don't auto-suggest the Responses surface for Mistral medium because
|
||||
vLLM's Responses API tool-call parser isn't wired up for it yet —
|
||||
operators who want per-request reasoning effort must pick "Responses
|
||||
API" manually in the admin UI and accept the tool-calling limitation.
|
||||
"""
|
||||
p = suggest_profile("vllm", "mistralai/Mistral-Medium-3-Instruct")
|
||||
assert p["server_compat"]["server_type"] == "vllm"
|
||||
assert "api_surface" not in p["server_compat"]
|
||||
assert "capabilities" not in p
|
||||
|
||||
def test_vllm_mistral_medium_profile_still_available(self) -> None:
|
||||
"""The vllm-mistral-medium profile remains in _PROFILES so an operator
|
||||
who explicitly opts in via the admin UI gets the Responses surface."""
|
||||
from turnstone.core.server_compat import _PROFILES
|
||||
|
||||
assert "vllm-mistral-medium" in _PROFILES
|
||||
assert _PROFILES["vllm-mistral-medium"]["server_compat"]["api_surface"] == "responses"
|
||||
|
||||
def test_holo_requires_holo2(self) -> None:
|
||||
"""Short 'holo' prefix shouldn't false-match; 'holo2' should match."""
|
||||
p_short = suggest_profile("vllm", "some-org/hologram-7b")
|
||||
@@ -110,32 +131,47 @@ class TestSuggestProfile:
|
||||
|
||||
|
||||
class TestMergeServerCompat:
|
||||
def test_empty_compat_returns_base_only(self) -> None:
|
||||
def test_empty_base_and_compat_is_empty(self) -> None:
|
||||
"""No base, no compat → no extra_body needed."""
|
||||
assert merge_server_compat(None, {}) == {}
|
||||
assert merge_server_compat({}, {}) == {}
|
||||
|
||||
def test_explicit_base_passes_through(self) -> None:
|
||||
"""Explicit chat_template_kwargs base is forwarded as-is."""
|
||||
base = {"reasoning_effort": "medium"}
|
||||
result = merge_server_compat(base, {})
|
||||
assert result == {"chat_template_kwargs": {"reasoning_effort": "medium"}}
|
||||
|
||||
def test_extra_body_merged_top_level(self) -> None:
|
||||
base = {"reasoning_effort": "medium"}
|
||||
compat = {"extra_body": {"skip_special_tokens": False}}
|
||||
result = merge_server_compat(base, compat)
|
||||
assert result["skip_special_tokens"] is False
|
||||
assert "chat_template_kwargs" in result
|
||||
def test_extra_body_merged_top_level_no_base(self) -> None:
|
||||
"""Server-level overrides forward without a chat_template_kwargs wrapper."""
|
||||
result = merge_server_compat(None, {"extra_body": {"skip_special_tokens": False}})
|
||||
assert result == {"skip_special_tokens": False}
|
||||
|
||||
def test_full_vllm_gemma_compat(self) -> None:
|
||||
base = {"reasoning_effort": "medium"}
|
||||
def test_full_vllm_gemma_compat_no_base(self) -> None:
|
||||
"""vLLM workaround forwards on its own."""
|
||||
compat = {
|
||||
"server_type": "vllm",
|
||||
"extra_body": {"skip_special_tokens": False},
|
||||
}
|
||||
result = merge_server_compat(base, compat)
|
||||
result = merge_server_compat(None, compat)
|
||||
assert result == {"skip_special_tokens": False}
|
||||
|
||||
def test_operator_chat_template_kwargs_only(self) -> None:
|
||||
"""Operator can set chat_template_kwargs explicitly without seeding the base."""
|
||||
compat = {
|
||||
"extra_body": {
|
||||
"chat_template_kwargs": {"reasoning_effort": "high"},
|
||||
"skip_special_tokens": False,
|
||||
},
|
||||
}
|
||||
result = merge_server_compat(None, compat)
|
||||
assert result == {
|
||||
"chat_template_kwargs": {"reasoning_effort": "medium"},
|
||||
"chat_template_kwargs": {"reasoning_effort": "high"},
|
||||
"skip_special_tokens": False,
|
||||
}
|
||||
|
||||
def test_extra_body_chat_template_kwargs_deep_merged(self) -> None:
|
||||
"""chat_template_kwargs in extra_body is deep-merged, operator wins."""
|
||||
def test_extra_body_chat_template_kwargs_deep_merged_with_base(self) -> None:
|
||||
"""Operator chat_template_kwargs deep-merges over the seeded base."""
|
||||
base = {"reasoning_effort": "medium"}
|
||||
compat = {
|
||||
"extra_body": {
|
||||
@@ -144,17 +180,15 @@ class TestMergeServerCompat:
|
||||
},
|
||||
}
|
||||
result = merge_server_compat(base, compat)
|
||||
# Operator values win over base
|
||||
assert result["chat_template_kwargs"]["custom_flag"] is True
|
||||
# Operator value wins over seeded base
|
||||
assert result["chat_template_kwargs"]["reasoning_effort"] == "high"
|
||||
assert result["skip_special_tokens"] is False
|
||||
|
||||
def test_extra_body_chat_template_kwargs_non_dict_ignored(self) -> None:
|
||||
"""Non-dict chat_template_kwargs in extra_body is safely ignored."""
|
||||
base = {"reasoning_effort": "medium"}
|
||||
compat = {"extra_body": {"chat_template_kwargs": "bad"}}
|
||||
result = merge_server_compat(base, compat)
|
||||
assert result["chat_template_kwargs"] == {"reasoning_effort": "medium"}
|
||||
assert merge_server_compat(None, compat) == {}
|
||||
|
||||
def test_base_not_mutated(self) -> None:
|
||||
base = {"reasoning_effort": "medium"}
|
||||
@@ -164,9 +198,7 @@ class TestMergeServerCompat:
|
||||
|
||||
def test_non_dict_extra_body_ignored(self) -> None:
|
||||
"""Gracefully handle malformed server_compat."""
|
||||
base = {"reasoning_effort": "medium"}
|
||||
result = merge_server_compat(base, {"extra_body": 42})
|
||||
assert result == {"chat_template_kwargs": {"reasoning_effort": "medium"}}
|
||||
assert merge_server_compat(None, {"extra_body": 42}) == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -178,45 +210,58 @@ class TestEndToEndRequestShaping:
|
||||
"""Compose both layers — session builds extra_params, provider applies thinking."""
|
||||
|
||||
def test_vllm_gemma_full_flow(self) -> None:
|
||||
"""Session merges server workarounds, provider adds thinking param."""
|
||||
"""Session forwards server workarounds, provider adds thinking param."""
|
||||
caps = ModelCapabilities(thinking_mode="manual", thinking_param="enable_thinking")
|
||||
base_ctk = {"reasoning_effort": "medium"}
|
||||
server_compat = {
|
||||
"server_type": "vllm",
|
||||
"extra_body": {"skip_special_tokens": False},
|
||||
}
|
||||
# Step 1: session merges
|
||||
extra_params = merge_server_compat(base_ctk, server_compat)
|
||||
# Step 2: provider finalises
|
||||
# Step 1: session forwards (no auto-injection of reasoning_effort).
|
||||
extra_params = merge_server_compat(None, server_compat)
|
||||
# Step 2: provider injects thinking param into chat_template_kwargs.
|
||||
extra_body = dict(extra_params)
|
||||
OpenAIChatCompletionsProvider._apply_thinking_mode(extra_body, caps)
|
||||
|
||||
assert extra_body == {
|
||||
"chat_template_kwargs": {
|
||||
"reasoning_effort": "medium",
|
||||
"enable_thinking": True,
|
||||
},
|
||||
"chat_template_kwargs": {"enable_thinking": True},
|
||||
"skip_special_tokens": False,
|
||||
}
|
||||
|
||||
def test_granite_thinking_key(self) -> None:
|
||||
"""Granite uses 'thinking' instead of 'enable_thinking'."""
|
||||
caps = ModelCapabilities(thinking_mode="manual", thinking_param="thinking")
|
||||
extra_params = merge_server_compat({"reasoning_effort": "low"}, {})
|
||||
extra_params = merge_server_compat(None, {})
|
||||
extra_body = dict(extra_params)
|
||||
OpenAIChatCompletionsProvider._apply_thinking_mode(extra_body, caps)
|
||||
|
||||
assert extra_body["chat_template_kwargs"]["thinking"] is True
|
||||
assert "enable_thinking" not in extra_body["chat_template_kwargs"]
|
||||
assert extra_body == {"chat_template_kwargs": {"thinking": True}}
|
||||
|
||||
def test_non_thinking_model_no_injection(self) -> None:
|
||||
"""Non-thinking model gets no thinking params."""
|
||||
"""Non-thinking model gets no chat_template_kwargs at all."""
|
||||
caps = ModelCapabilities() # thinking_mode="none"
|
||||
extra_params = merge_server_compat({"reasoning_effort": "medium"}, {})
|
||||
extra_params = merge_server_compat(None, {})
|
||||
extra_body = dict(extra_params)
|
||||
OpenAIChatCompletionsProvider._apply_thinking_mode(extra_body, caps)
|
||||
|
||||
assert extra_body == {"chat_template_kwargs": {"reasoning_effort": "medium"}}
|
||||
assert extra_body == {}
|
||||
|
||||
def test_operator_reasoning_effort_passthrough(self) -> None:
|
||||
"""Operator-supplied reasoning_effort under chat_template_kwargs is preserved."""
|
||||
caps = ModelCapabilities(thinking_mode="manual", thinking_param="enable_thinking")
|
||||
compat = {
|
||||
"server_type": "vllm",
|
||||
"extra_body": {"chat_template_kwargs": {"reasoning_effort": "high"}},
|
||||
}
|
||||
extra_params = merge_server_compat(None, compat)
|
||||
extra_body = dict(extra_params)
|
||||
OpenAIChatCompletionsProvider._apply_thinking_mode(extra_body, caps)
|
||||
|
||||
assert extra_body == {
|
||||
"chat_template_kwargs": {
|
||||
"reasoning_effort": "high",
|
||||
"enable_thinking": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Tests for ``initialize_mcp_crypto_state`` startup gate.
|
||||
|
||||
Phase 3 of the OAuth-MCP RFC: validates fail-loud behavior when an
|
||||
operator forgets the encryption key on a node that hosts OAuth-protected
|
||||
MCP server rows.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import types
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
import turnstone.core.config as cfg_mod
|
||||
from turnstone.core.mcp_crypto import (
|
||||
MCPTokenCipher,
|
||||
MCPTokenStore,
|
||||
initialize_mcp_crypto_state,
|
||||
)
|
||||
|
||||
|
||||
def _patch_security(monkeypatch: pytest.MonkeyPatch, payload: dict) -> None:
|
||||
"""Override ``load_config('security')`` to return ``payload``."""
|
||||
|
||||
def fake(section: str | None = None) -> dict:
|
||||
if section == "security":
|
||||
return payload
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(cfg_mod, "load_config", fake)
|
||||
|
||||
|
||||
class TestInitializeMcpCryptoState:
|
||||
def test_startup_succeeds_with_no_oauth_user_rows_and_no_key(
|
||||
self, backend, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Common case: no key, no oauth_user rows -> sentinels installed."""
|
||||
_patch_security(monkeypatch, {})
|
||||
|
||||
state = types.SimpleNamespace()
|
||||
initialize_mcp_crypto_state(state, node_id="n1")
|
||||
|
||||
assert state.mcp_token_cipher is None
|
||||
assert state.mcp_token_store is None
|
||||
|
||||
def test_startup_succeeds_with_key_and_oauth_user_row(
|
||||
self, backend, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Operator has wired a key and at least one oauth_user row.
|
||||
|
||||
Cipher + store should land on app_state.
|
||||
"""
|
||||
# Plant an oauth_user row.
|
||||
backend.create_mcp_server(
|
||||
server_id="srv-1",
|
||||
name="oauth-srv",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
)
|
||||
_patch_security(
|
||||
monkeypatch,
|
||||
{"mcp_token_encryption_key": Fernet.generate_key().decode()},
|
||||
)
|
||||
|
||||
state = types.SimpleNamespace()
|
||||
initialize_mcp_crypto_state(state, node_id="n1")
|
||||
|
||||
assert isinstance(state.mcp_token_cipher, MCPTokenCipher)
|
||||
assert isinstance(state.mcp_token_store, MCPTokenStore)
|
||||
|
||||
def test_startup_aborts_with_oauth_user_row_and_no_key(
|
||||
self, backend, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Misconfiguration: oauth_user row exists, no key -> SystemExit(1)."""
|
||||
backend.create_mcp_server(
|
||||
server_id="srv-1",
|
||||
name="oauth-srv",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
)
|
||||
_patch_security(monkeypatch, {})
|
||||
|
||||
state = types.SimpleNamespace()
|
||||
with (
|
||||
caplog.at_level("ERROR", logger="turnstone.core.mcp_crypto"),
|
||||
pytest.raises(SystemExit) as exc_info,
|
||||
):
|
||||
initialize_mcp_crypto_state(state, node_id="n1")
|
||||
assert exc_info.value.code == 1
|
||||
# Operator-actionable error message names BOTH supported config-key
|
||||
# forms so an operator using the rotation list (plural) is not
|
||||
# misled into thinking only the singular form is valid.
|
||||
messages = " ".join(record.message for record in caplog.records)
|
||||
assert "mcp_token_encryption_keys" in messages
|
||||
assert re.search(r"mcp_token_encryption_key(?!s)", messages) is not None
|
||||
|
||||
def test_startup_aborts_with_invalid_key(
|
||||
self, backend, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Malformed key material should fail loud at startup, not at first use."""
|
||||
_patch_security(monkeypatch, {"mcp_token_encryption_key": "###not-base64###"})
|
||||
|
||||
state = types.SimpleNamespace()
|
||||
with (
|
||||
caplog.at_level("ERROR", logger="turnstone.core.mcp_crypto"),
|
||||
pytest.raises(SystemExit) as exc_info,
|
||||
):
|
||||
initialize_mcp_crypto_state(state, node_id="n1")
|
||||
assert exc_info.value.code == 1
|
||||
@@ -0,0 +1,346 @@
|
||||
"""Tests for the per-node ``models`` metadata pipeline.
|
||||
|
||||
Two helpers in ``server.py`` carry the load:
|
||||
|
||||
- ``_collect_node_models_metadata`` projects the live ``ModelRegistry``
|
||||
into the node_metadata row shape ``[{alias, provider, healthy}, ...]``.
|
||||
- ``_publish_models_metadata`` short-circuits redundant writes via a
|
||||
payload cache on ``app_state`` and is the helper called from both
|
||||
the heartbeat loop and ``internal_model_reload``.
|
||||
|
||||
These tests pin the projection shape, the health-flag wiring, the
|
||||
cache short-circuit, and the model-reload integration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from turnstone.core.healthcheck import HealthTrackerRegistry
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
from turnstone.server import (
|
||||
_collect_node_models_metadata,
|
||||
_publish_models_metadata,
|
||||
)
|
||||
|
||||
|
||||
def _registry(*aliases_with_url: tuple[str, str]) -> ModelRegistry:
|
||||
"""Build a registry from ``(alias, base_url)`` pairs.
|
||||
|
||||
Two aliases sharing a ``base_url`` deliberately share a tracker —
|
||||
that's the contract the cluster-level health surface needs to
|
||||
preserve, and it's worth pinning in a test.
|
||||
"""
|
||||
models = {
|
||||
alias: ModelConfig(alias=alias, base_url=url, api_key="k", model=alias, provider="openai")
|
||||
for alias, url in aliases_with_url
|
||||
}
|
||||
default = aliases_with_url[0][0]
|
||||
return ModelRegistry(models, default=default)
|
||||
|
||||
|
||||
def test_returns_none_when_registry_missing():
|
||||
state = SimpleNamespace()
|
||||
assert _collect_node_models_metadata(state) is None
|
||||
|
||||
|
||||
def test_projects_all_aliases_with_default_healthy_when_no_tracker():
|
||||
"""Without a ``health_registry`` (or before any request has flowed
|
||||
through a backend), every alias surfaces as ``healthy=True`` —
|
||||
operators shouldn't get an empty ``models`` list on a freshly
|
||||
started node just because the backends haven't been exercised."""
|
||||
reg = _registry(("a", "http://x"), ("b", "http://y"))
|
||||
state = SimpleNamespace(registry=reg)
|
||||
entry = _collect_node_models_metadata(state)
|
||||
assert entry is not None
|
||||
key, value, source = entry
|
||||
assert key == "models"
|
||||
assert source == "auto"
|
||||
rows = json.loads(value)
|
||||
assert len(rows) == 2
|
||||
aliases = {r["alias"] for r in rows}
|
||||
assert aliases == {"a", "b"}
|
||||
assert all(r["healthy"] is True for r in rows)
|
||||
assert all(r["provider"] == "openai" for r in rows)
|
||||
# Provider-side model identifier intentionally omitted — coords
|
||||
# kept passing it as ``spawn_workstream(model=...)`` when they
|
||||
# should have passed the local alias. Lock the projected keys
|
||||
# so a future contributor doesn't reintroduce the footgun.
|
||||
for row in rows:
|
||||
assert set(row.keys()) == {"alias", "provider", "healthy"}
|
||||
|
||||
|
||||
def test_health_flag_reflects_tracker_state():
|
||||
reg = _registry(("a", "http://x"), ("b", "http://y"))
|
||||
health_reg = HealthTrackerRegistry(failure_threshold=2)
|
||||
# Seed the tracker for "a"'s backend and drive it into the degraded
|
||||
# state — two consecutive failures cross the threshold.
|
||||
bad_tracker = health_reg.get_tracker(provider="openai", base_url="http://x")
|
||||
bad_tracker.record_failure()
|
||||
bad_tracker.record_failure()
|
||||
assert bad_tracker.is_degraded
|
||||
# "b" gets a tracker that has only seen successes.
|
||||
good_tracker = health_reg.get_tracker(provider="openai", base_url="http://y")
|
||||
good_tracker.record_success()
|
||||
state = SimpleNamespace(registry=reg, health_registry=health_reg)
|
||||
rows = json.loads(_collect_node_models_metadata(state)[1])
|
||||
by_alias = {r["alias"]: r for r in rows}
|
||||
assert by_alias["a"]["healthy"] is False
|
||||
assert by_alias["b"]["healthy"] is True
|
||||
|
||||
|
||||
def test_two_aliases_sharing_a_backend_share_a_tracker():
|
||||
"""Two aliases that point at the same ``(provider, base_url)``
|
||||
share a single :class:`BackendHealthTracker` — degrading one is
|
||||
expected to surface as degraded on the other. The list_nodes
|
||||
projection should respect that, otherwise a coord could see
|
||||
``alias-a`` healthy and ``alias-b`` degraded for the same
|
||||
backend."""
|
||||
reg = _registry(("alpha", "http://shared"), ("beta", "http://shared"))
|
||||
health_reg = HealthTrackerRegistry(failure_threshold=1)
|
||||
tracker = health_reg.get_tracker(provider="openai", base_url="http://shared")
|
||||
tracker.record_failure() # threshold=1 — degraded immediately
|
||||
state = SimpleNamespace(registry=reg, health_registry=health_reg)
|
||||
rows = json.loads(_collect_node_models_metadata(state)[1])
|
||||
assert {r["alias"]: r["healthy"] for r in rows} == {"alpha": False, "beta": False}
|
||||
|
||||
|
||||
def test_alias_with_no_tracker_yet_defaults_to_healthy():
|
||||
"""An alias the registry knows about but whose backend hasn't been
|
||||
invoked yet has no tracker. Default to healthy so a brand-new
|
||||
alias is immediately visible to coordinators rather than waiting
|
||||
for the first request to seed a tracker.
|
||||
|
||||
The collector calls ``health_reg.get_tracker(...)`` which mints a
|
||||
fresh tracker on first lookup — that's the path under test here.
|
||||
The freshly minted tracker reports ``is_healthy=True`` (default
|
||||
state), so the projection labels the alias healthy.
|
||||
"""
|
||||
reg = _registry(("a", "http://x"))
|
||||
health_reg = HealthTrackerRegistry() # empty — no trackers seeded
|
||||
state = SimpleNamespace(registry=reg, health_registry=health_reg)
|
||||
rows = json.loads(_collect_node_models_metadata(state)[1])
|
||||
assert rows[0]["healthy"] is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _publish_models_metadata — cache short-circuit + projection wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _publish_state() -> SimpleNamespace:
|
||||
"""Build an ``app_state`` with a minimal registry + health surface."""
|
||||
reg = _registry(("a", "http://x"))
|
||||
return SimpleNamespace(registry=reg, health_registry=HealthTrackerRegistry())
|
||||
|
||||
|
||||
def test_publish_writes_when_payload_changes():
|
||||
"""First publish has nothing in the cache — write happens; cache
|
||||
fills. Second publish on the same unchanged registry skips the
|
||||
write entirely."""
|
||||
state = _publish_state()
|
||||
storage = MagicMock()
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 1
|
||||
cached = state._last_models_payload
|
||||
assert isinstance(cached, str) and "alias" in cached
|
||||
# Second call, same registry, same health: cached payload matches
|
||||
# — write must be skipped to avoid the per-30s UPSERT churn.
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 1
|
||||
|
||||
|
||||
def test_publish_records_metric_outcome(monkeypatch):
|
||||
"""The publish helper feeds ``record_node_models_publish`` so
|
||||
Prometheus can expose the hit-rate. Storage failures must NOT
|
||||
record either outcome — counters should reflect actual cache
|
||||
decisions, not transient DB errors that will retry.
|
||||
|
||||
Replaces the module-level ``turnstone.server._metrics`` binding
|
||||
via string-form monkeypatch (with auto-restore) rather than
|
||||
patching an instance attribute on the imported singleton. Other
|
||||
tests in the suite reassign ``srv_mod._metrics`` (some without
|
||||
using monkeypatch), so an instance captured at import time can
|
||||
diverge from the binding the live ``_publish_models_metadata``
|
||||
reads on each call.
|
||||
"""
|
||||
state = _publish_state()
|
||||
storage = MagicMock()
|
||||
calls: list[bool] = []
|
||||
|
||||
class _FakeMetrics:
|
||||
def record_node_models_publish(self, *, written: bool) -> None:
|
||||
calls.append(written)
|
||||
|
||||
monkeypatch.setattr("turnstone.server._metrics", _FakeMetrics())
|
||||
|
||||
_publish_models_metadata(state, storage, "node-a") # first → write
|
||||
_publish_models_metadata(state, storage, "node-a") # second → skip
|
||||
assert calls == [True, False]
|
||||
|
||||
# Storage error: no metric recorded.
|
||||
storage.set_node_metadata_bulk.side_effect = RuntimeError("db down")
|
||||
state._last_models_payload = None # invalidate cache to force a write attempt
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert calls == [True, False] # unchanged
|
||||
|
||||
|
||||
def test_publish_rewrites_when_health_flips():
|
||||
"""A health-tracker state change must invalidate the cache and
|
||||
drive a fresh write — otherwise the discovery surface would lag
|
||||
a flip indefinitely."""
|
||||
state = _publish_state()
|
||||
storage = MagicMock()
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 1
|
||||
# Drive the only tracker to degraded.
|
||||
tracker = state.health_registry.get_tracker(provider="openai", base_url="http://x")
|
||||
for _ in range(10):
|
||||
tracker.record_failure()
|
||||
assert tracker.is_degraded
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 2
|
||||
|
||||
|
||||
def test_publish_swallows_storage_error_without_updating_cache():
|
||||
"""A storage failure must NOT poison the cache — the next call
|
||||
should retry the write rather than think it succeeded."""
|
||||
state = _publish_state()
|
||||
storage = MagicMock()
|
||||
storage.set_node_metadata_bulk.side_effect = RuntimeError("db down")
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 1
|
||||
assert getattr(state, "_last_models_payload", None) is None
|
||||
# Recover: a subsequent successful call writes again.
|
||||
storage.set_node_metadata_bulk.side_effect = None
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 2
|
||||
assert state._last_models_payload is not None
|
||||
|
||||
|
||||
def test_publish_skips_when_registry_missing():
|
||||
"""Without a registry there's nothing to project; nothing should
|
||||
be written and the cache must not be set."""
|
||||
state = SimpleNamespace()
|
||||
storage = MagicMock()
|
||||
_publish_models_metadata(state, storage, "node-a")
|
||||
assert storage.set_node_metadata_bulk.call_count == 0
|
||||
assert getattr(state, "_last_models_payload", None) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# internal_model_reload — integration: registry change must rewrite the row
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_model_reload_endpoint_rewrites_models_metadata(monkeypatch, tmp_path):
|
||||
"""A successful ``internal_model_reload`` must refresh
|
||||
``node_metadata.models`` so a coordinator sees the new alias on
|
||||
its next ``list_nodes`` without waiting up to 30s for the
|
||||
heartbeat tick.
|
||||
|
||||
The endpoint pulls a fresh registry from
|
||||
``load_model_registry(...)`` and reloads in-place — we stub the
|
||||
loader to return a registry with a different alias set so the
|
||||
publish-cache invalidation is exercised end-to-end.
|
||||
"""
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
from turnstone.server import internal_model_reload
|
||||
|
||||
storage = SQLiteBackend(str(tmp_path / "reload.db"))
|
||||
|
||||
# Old registry — single alias "a".
|
||||
old_reg = _registry(("a", "http://x"))
|
||||
# New registry that ``load_model_registry`` will return — adds "b".
|
||||
new_reg = ModelRegistry(
|
||||
{
|
||||
"a": ModelConfig(
|
||||
alias="a", base_url="http://x", api_key="k", model="a", provider="openai"
|
||||
),
|
||||
"b": ModelConfig(
|
||||
alias="b", base_url="http://y", api_key="k", model="b", provider="openai"
|
||||
),
|
||||
},
|
||||
default="a",
|
||||
)
|
||||
health_reg = HealthTrackerRegistry()
|
||||
|
||||
app_state = SimpleNamespace(
|
||||
registry=old_reg,
|
||||
health_registry=health_reg,
|
||||
cli_model_args={
|
||||
"base_url": "",
|
||||
"api_key": "",
|
||||
"model": "",
|
||||
"context_window": 0,
|
||||
"provider": "openai",
|
||||
},
|
||||
config_store=None,
|
||||
node_id="node-a",
|
||||
)
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=app_state))
|
||||
|
||||
# Patch the loader and storage accessors used inside the endpoint.
|
||||
# ``internal_model_reload`` does ``from turnstone.core.storage._registry
|
||||
# import get_storage`` inline, so patching the symbol on that module
|
||||
# is what intercepts the call.
|
||||
monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", lambda **_kw: new_reg)
|
||||
monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage)
|
||||
# The endpoint also broadcasts schema refreshes to active sessions
|
||||
# — stub this out, it's irrelevant to the metadata-write path.
|
||||
monkeypatch.setattr("turnstone.server._broadcast_agent_tool_schema_refresh", lambda _s: None)
|
||||
|
||||
response = internal_model_reload(request) # type: ignore[arg-type]
|
||||
assert response.status_code == 200
|
||||
|
||||
rows = storage.get_node_metadata("node-a")
|
||||
by_key = {r["key"]: r for r in rows}
|
||||
assert "models" in by_key
|
||||
payload = json.loads(by_key["models"]["value"])
|
||||
assert {r["alias"] for r in payload} == {"a", "b"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shutdown race: heartbeat write must NOT resurrect post-shutdown delete
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_heartbeat_write_awaits_before_shutdown_delete():
|
||||
"""Pin the shutdown-race fix.
|
||||
|
||||
Before the fix, the lifespan shutdown sequence was:
|
||||
|
||||
1. ``_heartbeat_task.cancel()`` — fire-and-forget
|
||||
2. ``delete_node_metadata_by_source(node_id, "auto")``
|
||||
|
||||
A heartbeat tick already inside ``asyncio.to_thread(...)`` for
|
||||
the ``set_node_metadata_bulk`` call would complete AFTER step 2,
|
||||
resurrecting the deleted ``models`` row. The fix awaits the
|
||||
cancelled task with ``contextlib.suppress(...)`` between (1) and
|
||||
(2), so the in-flight write lands first.
|
||||
|
||||
We verify the fix by introspecting ``server.py`` source — the
|
||||
real lifespan is hard to test deterministically without a full
|
||||
Starlette app, but the textual ordering between
|
||||
``_heartbeat_task.cancel()`` and the delete is a stable contract
|
||||
that catches the regression cheaply.
|
||||
"""
|
||||
import inspect
|
||||
import sys
|
||||
|
||||
src = inspect.getsource(sys.modules[_collect_node_models_metadata.__module__])
|
||||
cancel_idx = src.find("_heartbeat_task.cancel()")
|
||||
delete_idx = src.find('delete_node_metadata_by_source, _svc_node_id, "auto"')
|
||||
await_idx = src.find("await _heartbeat_task", cancel_idx)
|
||||
assert cancel_idx != -1
|
||||
assert delete_idx != -1
|
||||
assert await_idx != -1
|
||||
# The fix-line must sit BETWEEN the cancel and the delete.
|
||||
assert cancel_idx < await_idx < delete_idx, (
|
||||
"Shutdown race regression: "
|
||||
"_heartbeat_task.cancel() must be followed by `await _heartbeat_task` "
|
||||
"BEFORE delete_node_metadata_by_source(..., 'auto') so an in-flight "
|
||||
"set_node_metadata_bulk lands before the delete."
|
||||
)
|
||||
+2606
-222
File diff suppressed because it is too large
Load Diff
@@ -4,6 +4,8 @@ from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.attachments import Attachment
|
||||
from turnstone.core.memory import (
|
||||
get_attachment,
|
||||
@@ -273,149 +275,29 @@ class TestProviderIntegration:
|
||||
assert "DO THE THING" in parts[1]["text"]
|
||||
|
||||
|
||||
class TestQueuedWithAttachments:
|
||||
"""Queued user turns must carry their attachments through to dequeue."""
|
||||
class TestQueuedAttachmentsRejected:
|
||||
"""Queued user messages can't carry attachments — see
|
||||
:class:`AttachmentsNotQueueableError` for the role-ordering reason
|
||||
(an attachment-bearing queued item would have to be appended as a
|
||||
separate user turn, injecting ``user`` between
|
||||
``assistant(tool_calls)`` and ``tool``)."""
|
||||
|
||||
def test_queue_message_rejects_attachments(self, tmp_db, mock_openai_client):
|
||||
from turnstone.core.session import AttachmentsNotQueueableError
|
||||
|
||||
def test_queue_message_stores_attachment_ids(self, tmp_db, mock_openai_client):
|
||||
s = _make_session(mock_openai_client)
|
||||
# Seed a pending attachment owned by the session user
|
||||
save_attachment("a-q1", s._ws_id, "u1", "q.md", "text/markdown", 1, "text", b"q")
|
||||
cleaned, priority, msg_id = s.queue_message("queued text", attachment_ids=["a-q1"])
|
||||
assert cleaned == "queued text"
|
||||
with pytest.raises(AttachmentsNotQueueableError):
|
||||
s.queue_message("queued text", attachment_ids=["a-q1"])
|
||||
# Queue stayed empty — nothing partially committed.
|
||||
assert s._queued_messages == {}
|
||||
|
||||
def test_queue_message_accepts_text_only(self, tmp_db, mock_openai_client):
|
||||
s = _make_session(mock_openai_client)
|
||||
cleaned, priority, msg_id = s.queue_message("plain text")
|
||||
assert cleaned == "plain text"
|
||||
with s._queued_lock:
|
||||
entry = s._queued_messages[msg_id]
|
||||
# Entry shape is (cleaned, priority, attachment_ids_tuple)
|
||||
assert entry[0] == "queued text"
|
||||
assert entry[2] == ("a-q1",)
|
||||
|
||||
def test_flush_queued_injects_multipart_user_turn(self, tmp_db, mock_openai_client):
|
||||
from turnstone.core.memory import reserve_attachments
|
||||
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-f1", s._ws_id, "u1", "f.md", "text/markdown", 3, "text", b"DAT")
|
||||
_c, _p, msg_id = s.queue_message("please review", attachment_ids=["a-f1"])
|
||||
# Server-side would have reserved before queueing; mirror that
|
||||
# so consume's token match succeeds on flush.
|
||||
reserve_attachments(["a-f1"], msg_id, s._ws_id, "u1")
|
||||
s._flush_queued_messages()
|
||||
|
||||
msgs = s.messages
|
||||
assert len(msgs) == 1
|
||||
msg = msgs[0]
|
||||
assert msg["role"] == "user"
|
||||
# Multipart shape — text + document parts
|
||||
assert isinstance(msg["content"], list)
|
||||
assert msg["content"][0] == {"type": "text", "text": "please review"}
|
||||
doc = msg["content"][1]
|
||||
assert doc["type"] == "document"
|
||||
assert doc["document"]["name"] == "f.md"
|
||||
assert doc["document"]["data"] == "DAT"
|
||||
# And the attachment is now consumed (not pending)
|
||||
assert get_attachment("a-f1")["message_id"] is not None
|
||||
assert list_pending_attachments(s._ws_id, "u1") == []
|
||||
|
||||
def test_flush_mixed_attachment_and_text_items(self, tmp_db, mock_openai_client):
|
||||
# Text-only items should combine into one turn while
|
||||
# attachment-bearing items flush as separate multipart turns.
|
||||
from turnstone.core.memory import reserve_attachments
|
||||
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-mx", s._ws_id, "u1", "x.md", "text/markdown", 1, "text", b"x")
|
||||
s.queue_message("first plain")
|
||||
_c, _p, mid = s.queue_message("with file", attachment_ids=["a-mx"])
|
||||
reserve_attachments(["a-mx"], mid, s._ws_id, "u1")
|
||||
s.queue_message("another plain")
|
||||
s._flush_queued_messages()
|
||||
|
||||
# We expect at least two user messages: one combining the plain
|
||||
# items flanking the multipart turn is allowed, but the
|
||||
# multipart turn must remain its own message.
|
||||
user_msgs = [m for m in s.messages if m.get("role") == "user"]
|
||||
multipart = [m for m in user_msgs if isinstance(m["content"], list)]
|
||||
assert len(multipart) == 1
|
||||
assert "with file" in multipart[0]["content"][0]["text"]
|
||||
|
||||
def test_flush_drops_cross_user_attachment_silently(self, tmp_db, mock_openai_client):
|
||||
# A forged attachment_id belonging to another user must not
|
||||
# produce an attached part — dequeue resolution re-scopes.
|
||||
s = _make_session(mock_openai_client, user_id="u1")
|
||||
save_attachment("a-other", s._ws_id, "u2", "other.md", "text/plain", 1, "text", b"o")
|
||||
s.queue_message("hi", attachment_ids=["a-other"])
|
||||
s._flush_queued_messages()
|
||||
# Flushed as plain text-only turn — the forged id was scope-dropped.
|
||||
msgs = s.messages
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["content"] == "hi"
|
||||
|
||||
|
||||
class TestQueueReservationLifecycle:
|
||||
"""session.queue_message + dequeue_message lifecycle with reservations."""
|
||||
|
||||
def test_dequeue_unreserves_attachments(self, tmp_db, mock_openai_client):
|
||||
from turnstone.core.memory import get_attachment, reserve_attachments
|
||||
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-deq", s._ws_id, "u1", "x.md", "text/plain", 1, "text", b"x")
|
||||
_cleaned, _priority, msg_id = s.queue_message("queued", attachment_ids=["a-deq"])
|
||||
# Simulate the server reserving after queue_message
|
||||
reserve_attachments(["a-deq"], msg_id, s._ws_id, "u1")
|
||||
assert get_attachment("a-deq")["reserved_for_msg_id"] == msg_id
|
||||
|
||||
# Dequeue (user cancelled the queued send)
|
||||
assert s.dequeue_message(msg_id) is True
|
||||
# Reservation is released — back to pending
|
||||
assert get_attachment("a-deq")["reserved_for_msg_id"] is None
|
||||
assert len(list_pending_attachments(s._ws_id, "u1")) == 1
|
||||
|
||||
def test_flush_consumes_reserved_attachment(self, tmp_db, mock_openai_client):
|
||||
from turnstone.core.memory import get_attachment, reserve_attachments
|
||||
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-flush", s._ws_id, "u1", "y.md", "text/plain", 1, "text", b"y")
|
||||
_c, _p, msg_id = s.queue_message("go", attachment_ids=["a-flush"])
|
||||
reserve_attachments(["a-flush"], msg_id, s._ws_id, "u1")
|
||||
|
||||
# Flush — queue drain must accept the reserved-for-this-msg attachment
|
||||
s._flush_queued_messages()
|
||||
row = get_attachment("a-flush")
|
||||
assert row["message_id"] is not None
|
||||
assert row["reserved_for_msg_id"] is None # cleared on consume
|
||||
# And the in-memory message is multipart with the doc attached
|
||||
assert isinstance(s.messages[-1]["content"], list)
|
||||
assert any(p.get("type") == "document" for p in s.messages[-1]["content"])
|
||||
|
||||
def test_resolve_rejects_reservation_for_other_msg(self, tmp_db, mock_openai_client):
|
||||
from turnstone.core.memory import reserve_attachments
|
||||
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-other", s._ws_id, "u1", "z.md", "text/plain", 1, "text", b"z")
|
||||
reserve_attachments(["a-other"], "q-OTHER", s._ws_id, "u1")
|
||||
# allow_reserved_for=None (default) → reserved rows are skipped
|
||||
assert s._resolve_attachment_ids(["a-other"]) == []
|
||||
# allow_reserved_for matches → accepted
|
||||
out = s._resolve_attachment_ids(["a-other"], allow_reserved_for="q-OTHER")
|
||||
assert [a.attachment_id for a in out] == ["a-other"]
|
||||
|
||||
|
||||
class TestExplicitAttachmentIdsOrderPreserved:
|
||||
"""session._resolve_attachment_ids must honour request order."""
|
||||
|
||||
def test_resolve_preserves_request_order(self, tmp_db, mock_openai_client):
|
||||
s = _make_session(mock_openai_client)
|
||||
# Insert in one order, request in the reverse order — resolver
|
||||
# must reflect the request, not the DB's INSERT order.
|
||||
save_attachment("a-1", s._ws_id, "u1", "first.md", "text/plain", 1, "text", b"1")
|
||||
save_attachment("a-2", s._ws_id, "u1", "second.md", "text/plain", 1, "text", b"2")
|
||||
save_attachment("a-3", s._ws_id, "u1", "third.md", "text/plain", 1, "text", b"3")
|
||||
|
||||
out = s._resolve_attachment_ids(["a-3", "a-1", "a-2"])
|
||||
assert [a.attachment_id for a in out] == ["a-3", "a-1", "a-2"]
|
||||
|
||||
def test_resolve_skips_unknown_and_keeps_order(self, tmp_db, mock_openai_client):
|
||||
s = _make_session(mock_openai_client)
|
||||
save_attachment("a-k", s._ws_id, "u1", "k.md", "text/plain", 1, "text", b"k")
|
||||
out = s._resolve_attachment_ids(["unknown", "a-k", ""])
|
||||
assert [a.attachment_id for a in out] == ["a-k"]
|
||||
assert s._queued_messages[msg_id] == ("plain text", priority)
|
||||
|
||||
|
||||
class TestTokenAccounting:
|
||||
|
||||
@@ -18,13 +18,22 @@ from __future__ import annotations
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.session_manager import SessionKindAdapter, SessionManager
|
||||
from turnstone.core.workstream import Workstream, WorkstreamKind, WorkstreamState
|
||||
from turnstone.core.workstream import (
|
||||
BULK_CLOSE_STATE_VALUES,
|
||||
Workstream,
|
||||
WorkstreamKind,
|
||||
WorkstreamState,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test fixtures
|
||||
@@ -92,6 +101,7 @@ class FakeAdapter:
|
||||
self.cleaned_up: list[str] = []
|
||||
self.build_session_calls = 0
|
||||
self.build_session_raises = build_session_raises
|
||||
self.last_build_model: object | None = None
|
||||
# Slow down session build so concurrent tests can race.
|
||||
self.build_session_delay = 0.0
|
||||
|
||||
@@ -135,8 +145,13 @@ class FakeAdapter:
|
||||
def build_ui(self, ws: Workstream) -> Any:
|
||||
return FakeUI()
|
||||
|
||||
def build_session(self, ws: Workstream, **_: object) -> Any:
|
||||
def build_session(self, ws: Workstream, **kwargs: object) -> Any:
|
||||
self.build_session_calls += 1
|
||||
# Record the ``model`` kwarg (None on fresh-create, the saved
|
||||
# alias on rehydrate) so tests can assert SessionManager.open()
|
||||
# threads the persisted alias through to construction instead
|
||||
# of letting the adapter resolve the *current* default alias.
|
||||
self.last_build_model = kwargs.get("model")
|
||||
if self.build_session_delay:
|
||||
time.sleep(self.build_session_delay)
|
||||
if self.build_session_raises:
|
||||
@@ -148,6 +163,28 @@ class FakeAdapter:
|
||||
return [e for e in self.events if e.kind == kind]
|
||||
|
||||
|
||||
class _FakeRowMapping:
|
||||
"""SQLAlchemy-Row-like wrapper exposing ``_mapping`` over a ``_Row``.
|
||||
|
||||
The real backends return ``Row`` objects with a ``_mapping`` attribute;
|
||||
consumers (e.g. ``CoordinatorIdleObserver._active_children``) prefer
|
||||
``row._mapping[<col>]`` access. This shim mirrors that contract so
|
||||
fakes are interchangeable with real Rows in tests.
|
||||
"""
|
||||
|
||||
def __init__(self, row: _Row) -> None:
|
||||
self._mapping = {
|
||||
"ws_id": row.ws_id,
|
||||
"user_id": row.user_id,
|
||||
"name": row.name,
|
||||
"kind": row.kind,
|
||||
"state": row.state,
|
||||
"parent_ws_id": row.parent_ws_id,
|
||||
"updated": row.updated,
|
||||
"node_id": row.node_id,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Row:
|
||||
ws_id: str
|
||||
@@ -156,6 +193,8 @@ class _Row:
|
||||
kind: str
|
||||
state: str = "idle"
|
||||
parent_ws_id: str | None = None
|
||||
updated: str = ""
|
||||
node_id: str | None = None
|
||||
|
||||
|
||||
class FakeStorage:
|
||||
@@ -164,8 +203,26 @@ class FakeStorage:
|
||||
def __init__(self) -> None:
|
||||
self.rows: dict[str, _Row] = {}
|
||||
self.state_updates: list[tuple[str, str]] = []
|
||||
self.touch_calls: list[str] = []
|
||||
self.register_raises = False
|
||||
self.lock = threading.Lock()
|
||||
# Live-services lookup target for close_idle pass 2. Map
|
||||
# service_type → list of live service_ids. Tests that exercise
|
||||
# liveness scoping populate this directly; default empty means
|
||||
# "no peers alive" (every row unprotected by liveness).
|
||||
self.live_services: dict[str, list[str]] = {}
|
||||
self.list_services_raises = False
|
||||
# Per-ws config (model_alias, temperature, …). Populated by
|
||||
# tests that exercise the rehydrate-preserves-config path; the
|
||||
# SessionManager.open() rehydrate path reads this through
|
||||
# ``self._storage.load_workstream_config`` so it can pass the
|
||||
# saved alias into ``build_session`` and avoid clobbering the
|
||||
# original on construction.
|
||||
self.ws_config: dict[str, dict[str, str]] = {}
|
||||
|
||||
@staticmethod
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
|
||||
def register_workstream(
|
||||
self,
|
||||
@@ -178,6 +235,8 @@ class FakeStorage:
|
||||
parent_ws_id: str | None = None,
|
||||
skill_id: str = "",
|
||||
skill_version: int = 0,
|
||||
state: str = "idle",
|
||||
updated: str | None = None,
|
||||
) -> None:
|
||||
if self.register_raises:
|
||||
raise RuntimeError("register forced failure")
|
||||
@@ -188,14 +247,66 @@ class FakeStorage:
|
||||
user_id=user_id or "",
|
||||
name=name,
|
||||
kind=kind_str,
|
||||
state=state,
|
||||
parent_ws_id=parent_ws_id,
|
||||
updated=updated if updated is not None else self._now_iso(),
|
||||
node_id=node_id,
|
||||
)
|
||||
|
||||
def touch_workstream(self, ws_id: str) -> None:
|
||||
with self.lock:
|
||||
self.touch_calls.append(ws_id)
|
||||
if ws_id in self.rows:
|
||||
self.rows[ws_id].updated = self._now_iso()
|
||||
|
||||
def update_workstream_state(self, ws_id: str, state: str) -> None:
|
||||
with self.lock:
|
||||
self.state_updates.append((ws_id, state))
|
||||
if ws_id in self.rows:
|
||||
self.rows[ws_id].state = state
|
||||
self.rows[ws_id].updated = self._now_iso()
|
||||
|
||||
def bulk_close_stale_orphans(
|
||||
self,
|
||||
kind: WorkstreamKind | str,
|
||||
cutoff: str,
|
||||
exclude_ws_ids: list[str],
|
||||
live_node_ids: list[str] | None = None,
|
||||
) -> list[str]:
|
||||
kind_str = kind.value if isinstance(kind, WorkstreamKind) else str(kind)
|
||||
excluded = set(exclude_ws_ids)
|
||||
live_set = set(live_node_ids) if live_node_ids else set()
|
||||
now = self._now_iso()
|
||||
closed: list[str] = []
|
||||
with self.lock:
|
||||
for ws_id, row in self.rows.items():
|
||||
if (
|
||||
row.kind == kind_str
|
||||
and row.state in BULK_CLOSE_STATE_VALUES
|
||||
and row.updated < cutoff
|
||||
and ws_id not in excluded
|
||||
):
|
||||
# Liveness gate: when live_node_ids was provided AND
|
||||
# non-empty, protect rows whose owner is in the live
|
||||
# set. NULL node_id is always eligible. When
|
||||
# live_node_ids is None or empty, no protection
|
||||
# (mirror of the real backends).
|
||||
if live_node_ids and row.node_id is not None and row.node_id in live_set:
|
||||
continue
|
||||
row.state = "closed"
|
||||
row.updated = now
|
||||
self.state_updates.append((ws_id, "closed"))
|
||||
closed.append(ws_id)
|
||||
return closed
|
||||
|
||||
def list_services(self, service_type: str, max_age_seconds: int = 120) -> list[dict[str, str]]:
|
||||
if self.list_services_raises:
|
||||
raise RuntimeError("list_services forced failure")
|
||||
with self.lock:
|
||||
return [
|
||||
{"service_id": sid, "service_type": service_type}
|
||||
for sid in self.live_services.get(service_type, [])
|
||||
]
|
||||
|
||||
def get_workstream(self, ws_id: str) -> dict[str, Any] | None:
|
||||
with self.lock:
|
||||
@@ -211,6 +322,47 @@ class FakeStorage:
|
||||
"parent_ws_id": row.parent_ws_id,
|
||||
}
|
||||
|
||||
def list_workstreams(
|
||||
self,
|
||||
node_id: str | None = None,
|
||||
limit: int = 100,
|
||||
*,
|
||||
parent_ws_id: str | None = None,
|
||||
kind: WorkstreamKind | str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> list[Any]:
|
||||
kind_str = kind.value if isinstance(kind, WorkstreamKind) else kind
|
||||
with self.lock:
|
||||
matched: list[_FakeRowMapping] = []
|
||||
for row in self.rows.values():
|
||||
if parent_ws_id is not None and row.parent_ws_id != parent_ws_id:
|
||||
continue
|
||||
if kind_str is not None and row.kind != kind_str:
|
||||
continue
|
||||
if user_id is not None and row.user_id != user_id:
|
||||
continue
|
||||
matched.append(_FakeRowMapping(row))
|
||||
# Order by updated DESC so the consumer's LIMIT semantics match
|
||||
# production (storage backends order this way).
|
||||
matched.sort(key=lambda r: r._mapping["updated"], reverse=True)
|
||||
return matched[:limit]
|
||||
|
||||
def count_workstreams_by_state(
|
||||
self,
|
||||
*,
|
||||
parent_ws_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
counts: dict[str, int] = {}
|
||||
with self.lock:
|
||||
for row in self.rows.values():
|
||||
if parent_ws_id is not None and row.parent_ws_id != parent_ws_id:
|
||||
continue
|
||||
if user_id is not None and row.user_id != user_id:
|
||||
continue
|
||||
counts[row.state] = counts.get(row.state, 0) + 1
|
||||
return counts
|
||||
|
||||
def delete_workstream(self, ws_id: str) -> None:
|
||||
with self.lock:
|
||||
self.rows.pop(ws_id, None)
|
||||
@@ -218,6 +370,18 @@ class FakeStorage:
|
||||
def count_skill_versions(self, template_id: str) -> int:
|
||||
return 0
|
||||
|
||||
def load_workstream_config(self, ws_id: str) -> dict[str, str]:
|
||||
with self.lock:
|
||||
return dict(self.ws_config.get(ws_id, {}))
|
||||
|
||||
def save_workstream_config(self, ws_id: str, config: dict[str, str]) -> None:
|
||||
# Mirrors the real backend's INSERT OR REPLACE per-key semantics
|
||||
# — callers expect a partial save to overwrite only the keys
|
||||
# they pass, not the whole row.
|
||||
with self.lock:
|
||||
row = self.ws_config.setdefault(ws_id, {})
|
||||
row.update(config)
|
||||
|
||||
|
||||
_EMITTER_DEFAULT = object()
|
||||
|
||||
@@ -228,6 +392,8 @@ def _make_manager(
|
||||
max_active: int = 5,
|
||||
storage: FakeStorage | None = None,
|
||||
event_emitter: Any = _EMITTER_DEFAULT,
|
||||
node_id: str | None = None,
|
||||
model_validator: Callable[[str], bool] | None = None,
|
||||
) -> tuple[SessionManager, FakeAdapter, FakeStorage]:
|
||||
"""Build a SessionManager wired to a FakeAdapter for both Protocols.
|
||||
|
||||
@@ -246,6 +412,8 @@ def _make_manager(
|
||||
storage=storage,
|
||||
max_active=max_active,
|
||||
event_emitter=emitter,
|
||||
node_id=node_id,
|
||||
model_validator=model_validator,
|
||||
)
|
||||
return mgr, adapter, storage
|
||||
|
||||
@@ -579,6 +747,120 @@ def test_open_resurrects_closed_state() -> None:
|
||||
assert ws_id in [e.ws_id for e in adapter.events_of("rehydrated")]
|
||||
|
||||
|
||||
def test_open_threads_saved_model_alias_into_build_session() -> None:
|
||||
"""Reopening a closed ws must build the session with the *original*
|
||||
model alias, not the current registry default.
|
||||
|
||||
Without this, ``build_session(ws)`` is called with ``model=None`` →
|
||||
the production session_factory resolves ``_effective_default_alias()``
|
||||
→ ChatSession's ``__init__`` writes those defaults to
|
||||
``workstream_config`` (INSERT OR REPLACE) → the subsequent
|
||||
``resume()`` restores what is now the default. Net effect: every
|
||||
persisted knob (model, temperature, reasoning_effort, max_tokens,
|
||||
skill, creative_mode, instructions, …) silently resets on every
|
||||
reopen and on every service restart.
|
||||
"""
|
||||
mgr, adapter, storage = _make_manager()
|
||||
ws = mgr.create(user_id="u1")
|
||||
ws_id = ws.id
|
||||
# Pretend the user set a non-default alias when the ws was created;
|
||||
# the real path goes through ChatSession._save_config but the
|
||||
# FakeSession in this suite doesn't model that, so seed directly.
|
||||
storage.ws_config[ws_id] = {"model_alias": "gpt-5-pro"}
|
||||
mgr.close(ws_id)
|
||||
adapter.last_build_model = "<unset>" # sentinel — must be overwritten
|
||||
|
||||
reopened = mgr.open(ws_id)
|
||||
|
||||
assert reopened is not None
|
||||
assert adapter.last_build_model == "gpt-5-pro"
|
||||
|
||||
|
||||
def test_open_drops_saved_alias_when_validator_rejects() -> None:
|
||||
"""When the persisted alias is no longer in the registry, the
|
||||
manager must drop it before reaching ``build_session``. The
|
||||
factory still raises on unknown aliases on the fresh-create path
|
||||
(so a typo in body.model surfaces as 503), so the rehydrate path
|
||||
has to filter the alias here rather than relying on factory-side
|
||||
fallback. Without this filter, every reopen of a workstream pinned
|
||||
to a since-removed alias 500s."""
|
||||
mgr, adapter, storage = _make_manager(
|
||||
# Validator says "alias is no longer in the registry".
|
||||
model_validator=lambda alias: False,
|
||||
)
|
||||
ws = mgr.create(user_id="u1")
|
||||
ws_id = ws.id
|
||||
storage.ws_config[ws_id] = {"model_alias": "since-removed-alias"}
|
||||
mgr.close(ws_id)
|
||||
adapter.last_build_model = "<unset>"
|
||||
|
||||
reopened = mgr.open(ws_id)
|
||||
|
||||
assert reopened is not None
|
||||
assert adapter.last_build_model is None # alias dropped before reaching build_session
|
||||
|
||||
|
||||
def test_open_keeps_saved_alias_when_validator_accepts() -> None:
|
||||
"""Sanity: an alias that still resolves must be passed through
|
||||
unchanged. Filter only fires for stale aliases."""
|
||||
accepted: list[str] = []
|
||||
|
||||
def validator(alias: str) -> bool:
|
||||
accepted.append(alias)
|
||||
return True
|
||||
|
||||
mgr, adapter, storage = _make_manager(model_validator=validator)
|
||||
ws = mgr.create(user_id="u1")
|
||||
ws_id = ws.id
|
||||
storage.ws_config[ws_id] = {"model_alias": "still-live"}
|
||||
mgr.close(ws_id)
|
||||
adapter.last_build_model = "<unset>"
|
||||
|
||||
reopened = mgr.open(ws_id)
|
||||
|
||||
assert reopened is not None
|
||||
assert accepted == ["still-live"]
|
||||
assert adapter.last_build_model == "still-live"
|
||||
|
||||
|
||||
def test_open_falls_back_to_none_when_no_saved_alias() -> None:
|
||||
"""Reopening a ws with no saved alias must pass ``model=None`` to
|
||||
``build_session`` so the adapter's session_factory can fall back to
|
||||
the current default — matching the user's intent: best effort
|
||||
restore, default when the original is gone."""
|
||||
mgr, adapter, storage = _make_manager()
|
||||
ws = mgr.create(user_id="u1")
|
||||
ws_id = ws.id
|
||||
# No ws_config row — simulates "alias was never saved" or "saved
|
||||
# alias was empty string".
|
||||
assert ws_id not in storage.ws_config
|
||||
mgr.close(ws_id)
|
||||
adapter.last_build_model = "<unset>"
|
||||
|
||||
reopened = mgr.open(ws_id)
|
||||
|
||||
assert reopened is not None
|
||||
assert adapter.last_build_model is None
|
||||
|
||||
|
||||
def test_open_touches_workstream_on_rehydrate() -> None:
|
||||
"""Rehydrating a workstream must bump its ``updated`` so a concurrent
|
||||
close_idle pass-2 in this same process can't clobber the freshly-loaded
|
||||
row to ``closed`` because its DB ``updated`` is older than the cutoff.
|
||||
The touch is best-effort (try/except in open()) but must fire on the
|
||||
happy path."""
|
||||
mgr, _, storage = _make_manager()
|
||||
ws = mgr.create(user_id="u1")
|
||||
ws_id = ws.id
|
||||
mgr.close(ws_id)
|
||||
storage.touch_calls.clear() # only care about touches from rehydrate
|
||||
|
||||
reopened = mgr.open(ws_id)
|
||||
|
||||
assert reopened is not None
|
||||
assert ws_id in storage.touch_calls
|
||||
|
||||
|
||||
def test_open_ignores_owner_mismatch() -> None:
|
||||
# Turnstone is a trusted-team tool; row-level ownership is
|
||||
# metadata, not an access boundary. ``open`` no longer cares
|
||||
@@ -827,6 +1109,197 @@ def test_close_idle_on_empty_manager_returns_empty_list() -> None:
|
||||
assert mgr.close_idle(max_age_seconds=1.0) == []
|
||||
|
||||
|
||||
def test_close_idle_runs_db_orphan_pass() -> None:
|
||||
"""DB rows of this kind that aren't loaded into the manager get
|
||||
bulk-closed when their ``updated`` is older than the cutoff. Catches
|
||||
the orphan-after-process-restart case the original close_idle missed."""
|
||||
mgr, _, storage = _make_manager()
|
||||
# Orphan rows live in storage but were never loaded via mgr.create.
|
||||
storage.register_workstream(
|
||||
"orphan-1",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"orphan-2",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
state="thinking",
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert set(closed) == {"orphan-1", "orphan-2"}
|
||||
assert ("orphan-1", "closed") in storage.state_updates
|
||||
assert ("orphan-2", "closed") in storage.state_updates
|
||||
assert storage.rows["orphan-1"].state == "closed"
|
||||
assert storage.rows["orphan-2"].state == "closed"
|
||||
|
||||
|
||||
def test_close_idle_excludes_loaded_workstreams_from_db_pass() -> None:
|
||||
"""A workstream loaded into memory must NOT be reaped by the DB
|
||||
orphan pass even when its storage ``updated`` is stale — the
|
||||
in-memory pass owns those. Verifies the exclude_ws_ids plumbing."""
|
||||
mgr, _, storage = _make_manager()
|
||||
ws = mgr.create(user_id="u1")
|
||||
# Force the storage row's ``updated`` to look stale. In practice
|
||||
# ``set_state`` would bump it, but we're simulating a long-running
|
||||
# active workstream whose updated drifted older than the cutoff.
|
||||
storage.rows[ws.id].updated = "2020-01-01T00:00:00"
|
||||
|
||||
# Huge timeout so the in-memory IDLE pass skips it (stays loaded).
|
||||
closed = mgr.close_idle(max_age_seconds=10_000.0)
|
||||
|
||||
assert ws.id not in closed
|
||||
assert mgr.get(ws.id) is not None
|
||||
assert storage.rows[ws.id].state == "idle"
|
||||
|
||||
|
||||
def test_close_idle_filters_db_orphans_by_kind() -> None:
|
||||
"""An interactive manager's close_idle must not touch coordinator
|
||||
rows in storage and vice versa. Without this filter, both managers
|
||||
would race to close each other's rows."""
|
||||
mgr, _, storage = _make_manager() # interactive by default
|
||||
storage.register_workstream(
|
||||
"coord-orphan",
|
||||
kind=WorkstreamKind.COORDINATOR,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"interactive-orphan",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert "interactive-orphan" in closed
|
||||
assert "coord-orphan" not in closed
|
||||
assert storage.rows["coord-orphan"].state == "idle"
|
||||
assert storage.rows["interactive-orphan"].state == "closed"
|
||||
|
||||
|
||||
def test_close_idle_protects_rows_owned_by_live_services() -> None:
|
||||
"""Multi-node correctness: rows whose ``node_id`` matches a service
|
||||
with a recent heartbeat must NOT be reaped, even when *this* manager
|
||||
is on a different node — the alive peer may legitimately have them
|
||||
loaded. Liveness is the rendezvous router's primitive (post-PR-#384);
|
||||
using it here keeps reap scoping aligned with routing.
|
||||
|
||||
Default ``_make_manager`` uses an INTERACTIVE adapter, which derives
|
||||
``service_type='server'`` — so live_services seeded under "server"
|
||||
are what the manager queries."""
|
||||
mgr, _, storage = _make_manager()
|
||||
storage.live_services["server"] = ["node-b"] # only node-b is alive
|
||||
storage.register_workstream(
|
||||
"ours-from-dead-node",
|
||||
node_id="node-a", # dead pod (not in live_services)
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"theirs-still-alive",
|
||||
node_id="node-b",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert closed == ["ours-from-dead-node"]
|
||||
assert storage.rows["ours-from-dead-node"].state == "closed"
|
||||
assert storage.rows["theirs-still-alive"].state == "idle"
|
||||
|
||||
|
||||
def test_close_idle_protects_live_services_for_coordinator_kind() -> None:
|
||||
"""Coord-side parity: a coordinator manager derives
|
||||
``service_type='console'``, so live_services seeded under "console"
|
||||
are what gets queried. Mirrors the interactive test to ensure both
|
||||
halves of the production wiring are exercised."""
|
||||
coord_adapter = FakeAdapter(kind=WorkstreamKind.COORDINATOR)
|
||||
mgr, _, storage = _make_manager(coord_adapter)
|
||||
storage.live_services["console"] = ["console"] # console is alive
|
||||
storage.register_workstream(
|
||||
"alive-console-coord",
|
||||
node_id="console",
|
||||
kind=WorkstreamKind.COORDINATOR,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"dead-console-coord",
|
||||
node_id="dead-console-instance", # not in live set
|
||||
kind=WorkstreamKind.COORDINATOR,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert closed == ["dead-console-coord"]
|
||||
assert storage.rows["alive-console-coord"].state == "idle"
|
||||
assert storage.rows["dead-console-coord"].state == "closed"
|
||||
|
||||
|
||||
def test_close_idle_reaps_rows_with_null_node_id() -> None:
|
||||
"""A row with no ``node_id`` has no owner identity — age alone gates
|
||||
the reap. Defends against a NULL silently propagating through ``NOT
|
||||
IN (live)`` and protecting orphans forever."""
|
||||
mgr, _, storage = _make_manager()
|
||||
storage.live_services["server"] = ["node-a"]
|
||||
storage.register_workstream(
|
||||
"no-owner",
|
||||
node_id=None,
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert closed == ["no-owner"]
|
||||
|
||||
|
||||
def test_close_idle_reaps_all_orphans_when_no_peers_alive() -> None:
|
||||
"""When ``list_services`` returns an empty list (no heartbeating
|
||||
peers), every stale orphan is unprotected and gets reaped. This is
|
||||
the cold-start / single-process / dead-cluster-recovery case."""
|
||||
mgr, _, storage = _make_manager()
|
||||
# storage.live_services["server"] left empty — no peers heartbeating
|
||||
storage.register_workstream(
|
||||
"any-node-1",
|
||||
node_id="node-a",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
storage.register_workstream(
|
||||
"any-node-2",
|
||||
node_id="node-b",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert set(closed) == {"any-node-1", "any-node-2"}
|
||||
|
||||
|
||||
def test_close_idle_skips_pass_2_when_list_services_fails() -> None:
|
||||
"""Conservative fallback: if list_services fails we can't enumerate
|
||||
live owners safely, so pass 2 must skip rather than reap blind. Pass
|
||||
1 (in-memory IDLE) still runs."""
|
||||
mgr, _, storage = _make_manager()
|
||||
storage.list_services_raises = True
|
||||
storage.register_workstream(
|
||||
"would-be-orphan",
|
||||
node_id="node-a",
|
||||
kind=WorkstreamKind.INTERACTIVE,
|
||||
updated="2020-01-01T00:00:00",
|
||||
)
|
||||
|
||||
closed = mgr.close_idle(max_age_seconds=0.0)
|
||||
|
||||
assert closed == []
|
||||
assert storage.rows["would-be-orphan"].state == "idle"
|
||||
|
||||
|
||||
def test_list_all_returns_creation_order() -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
a = mgr.create(user_id="u1")
|
||||
@@ -1084,3 +1557,127 @@ class TestSessionManagerWithStateWriter:
|
||||
assert "running" not in ws_writes, (
|
||||
f"set_state after close enqueued through buffer: {ws_writes}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-subscriber observer — subscribe_to_state / unsubscribe_from_state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStateSubscribers:
|
||||
"""Multi-subscriber observer for ``set_state``.
|
||||
|
||||
Used by the CLI's background-attention notifier and by
|
||||
``SameNodeChildSource``. Subscribe / unsubscribe must be safe under
|
||||
concurrent dispatch, and dispatch must not skip / repeat callbacks
|
||||
when subscribers register or unregister mid-iteration.
|
||||
"""
|
||||
|
||||
def test_subscribe_fires_on_set_state(self) -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
ws = mgr.create(user_id="u1", name="ws", skill=None)
|
||||
events: list[tuple[str, str]] = []
|
||||
|
||||
def cb(ws_id: str, state: WorkstreamState) -> None:
|
||||
events.append((ws_id, state.value))
|
||||
|
||||
mgr.subscribe_to_state(cb)
|
||||
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
||||
assert events == [(ws.id, "running")]
|
||||
|
||||
def test_unsubscribe_stops_firing(self) -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
ws = mgr.create(user_id="u1", name="ws", skill=None)
|
||||
events: list[str] = []
|
||||
|
||||
def cb(_ws_id: str, state: WorkstreamState) -> None:
|
||||
events.append(state.value)
|
||||
|
||||
mgr.subscribe_to_state(cb)
|
||||
mgr.unsubscribe_from_state(cb)
|
||||
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
||||
assert events == []
|
||||
|
||||
def test_unsubscribe_unknown_is_noop(self) -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
# Doesn't raise.
|
||||
mgr.unsubscribe_from_state(lambda *_: None)
|
||||
|
||||
def test_multiple_subscribers_fire_in_registration_order(self) -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
ws = mgr.create(user_id="u1", name="ws", skill=None)
|
||||
order: list[int] = []
|
||||
|
||||
def make(i: int) -> Callable[[str, WorkstreamState], None]:
|
||||
def cb(_ws_id: str, _state: WorkstreamState) -> None:
|
||||
order.append(i)
|
||||
|
||||
return cb
|
||||
|
||||
mgr.subscribe_to_state(make(1))
|
||||
mgr.subscribe_to_state(make(2))
|
||||
mgr.subscribe_to_state(make(3))
|
||||
mgr.set_state(ws.id, WorkstreamState.IDLE)
|
||||
assert order == [1, 2, 3]
|
||||
|
||||
def test_subscriber_exception_does_not_block_others(self) -> None:
|
||||
mgr, _, _ = _make_manager()
|
||||
ws = mgr.create(user_id="u1", name="ws", skill=None)
|
||||
survived: list[str] = []
|
||||
|
||||
def boom(*_: Any) -> None:
|
||||
raise RuntimeError("subscriber crash")
|
||||
|
||||
def good(_ws_id: str, state: WorkstreamState) -> None:
|
||||
survived.append(state.value)
|
||||
|
||||
mgr.subscribe_to_state(boom)
|
||||
mgr.subscribe_to_state(good)
|
||||
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
||||
assert survived == ["running"]
|
||||
|
||||
@pytest.mark.parametrize("n_threads", [10, 50])
|
||||
def test_concurrent_subscribe(self, n_threads: int) -> None:
|
||||
"""Subscribe from many threads; all callbacks land in the list.
|
||||
|
||||
Validates the lock around mutation — without it the underlying
|
||||
list.append could lose entries under contention.
|
||||
"""
|
||||
mgr, _, _ = _make_manager()
|
||||
callbacks = [lambda *_, i=i: None for i in range(n_threads)]
|
||||
threads = [threading.Thread(target=mgr.subscribe_to_state, args=(cb,)) for cb in callbacks]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
# Snapshot under the lock to read the count safely.
|
||||
with mgr._state_subscribers_lock:
|
||||
assert len(mgr._state_subscribers) == n_threads
|
||||
|
||||
def test_subscribe_during_dispatch_does_not_corrupt_iteration(self) -> None:
|
||||
"""A subscriber that calls subscribe_to_state during its own
|
||||
callback must not affect the in-flight dispatch (snapshot
|
||||
isolation). This is the bug-1 invariant: mutation during
|
||||
iteration can't shift the iterator's index because dispatch
|
||||
iterates a snapshot, not the live list.
|
||||
"""
|
||||
mgr, _, _ = _make_manager()
|
||||
ws = mgr.create(user_id="u1", name="ws", skill=None)
|
||||
fired: list[str] = []
|
||||
|
||||
def late(_ws_id: str, state: WorkstreamState) -> None:
|
||||
fired.append("late:" + state.value)
|
||||
|
||||
def first(_ws_id: str, state: WorkstreamState) -> None:
|
||||
fired.append("first:" + state.value)
|
||||
mgr.subscribe_to_state(late) # mid-dispatch addition
|
||||
|
||||
mgr.subscribe_to_state(first)
|
||||
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
||||
# ``late`` was added during dispatch but the snapshot was
|
||||
# already frozen — so it doesn't fire on this round.
|
||||
assert fired == ["first:running"]
|
||||
# Next round it does fire, in registration order after first.
|
||||
fired.clear()
|
||||
mgr.set_state(ws.id, WorkstreamState.IDLE)
|
||||
assert fired == ["first:idle", "late:idle"]
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Tests for ``_format_mcp_dispatch_error`` and the three MCP exec sites.
|
||||
|
||||
The Phase 7b pool dispatcher signals user-actionable failures (consent
|
||||
required, insufficient scope) via ``RuntimeError(json_str)`` where
|
||||
``json_str`` is the structured-error payload built by
|
||||
:func:`turnstone.core.mcp_client._structured_error`. The exec sites in
|
||||
:mod:`turnstone.core.session` previously wrapped that JSON in
|
||||
``f"MCP X error: {e}"``, destroying the structured shape the dashboard
|
||||
renderer keys on. The helper preserves the JSON when the exception
|
||||
text decodes to a structured-error envelope and prefixes otherwise.
|
||||
|
||||
Sibling-bug coverage: every exec site (tool / read_resource /
|
||||
use_prompt) gets two assertions — JSON preserved on a consent-required
|
||||
exception, JSON-prefixed on a generic transport failure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from tests.test_session import _make_session
|
||||
from turnstone.core.session import _format_mcp_dispatch_error
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for the helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFormatMcpDispatchError:
|
||||
def test_preserves_consent_required_payload(self) -> None:
|
||||
payload = json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"code": "mcp_consent_required",
|
||||
"server": "srv-x",
|
||||
"detail": "No token for user. Consent flow required.",
|
||||
"consent_url": "/v1/api/mcp/oauth/start?server=srv-x",
|
||||
}
|
||||
}
|
||||
)
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError(payload))
|
||||
assert out == payload
|
||||
|
||||
def test_preserves_insufficient_scope_payload(self) -> None:
|
||||
payload = json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"code": "mcp_insufficient_scope",
|
||||
"server": "srv-x",
|
||||
"detail": "Tool requires elevated scopes.",
|
||||
"scopes_required": ["read", "write"],
|
||||
"consent_url": "/v1/api/mcp/oauth/start?server=srv-x&scopes=read+write",
|
||||
}
|
||||
}
|
||||
)
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError(payload))
|
||||
assert out == payload
|
||||
|
||||
def test_prefixes_generic_runtime_error(self) -> None:
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError("connection lost"))
|
||||
assert out == "MCP tool error: connection lost"
|
||||
|
||||
def test_prefixes_value_error(self) -> None:
|
||||
out = _format_mcp_dispatch_error("MCP tool error", ValueError("bad input"))
|
||||
assert out == "MCP tool error: bad input"
|
||||
|
||||
def test_prefixes_random_json_without_mcp_code(self) -> None:
|
||||
# JSON that isn't a structured-error envelope must NOT be passed
|
||||
# through verbatim — the helper only opens the gate for codes
|
||||
# prefixed ``mcp_``.
|
||||
payload = json.dumps({"foo": "bar"})
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError(payload))
|
||||
assert out == f"MCP tool error: {payload}"
|
||||
|
||||
def test_prefixes_envelope_with_non_mcp_code(self) -> None:
|
||||
payload = json.dumps({"error": {"code": "other_error", "server": "x", "detail": "y"}})
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError(payload))
|
||||
assert out == f"MCP tool error: {payload}"
|
||||
|
||||
def test_prefixes_envelope_without_dict_error(self) -> None:
|
||||
payload = json.dumps({"error": "plain string"})
|
||||
out = _format_mcp_dispatch_error("MCP tool error", RuntimeError(payload))
|
||||
assert out == f"MCP tool error: {payload}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests against the three MCP exec sites
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_CONSENT_REQUIRED_JSON = json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"code": "mcp_consent_required",
|
||||
"server": "srv-oauth",
|
||||
"detail": "No token for user. Consent flow required.",
|
||||
"consent_url": "/v1/api/mcp/oauth/start?server=srv-oauth",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _record_outputs(session) -> list[tuple[str, str, str, bool]]:
|
||||
"""Patch ``_report_tool_result`` to capture (call_id, name, output, is_error)."""
|
||||
captures: list[tuple[str, str, str, bool]] = []
|
||||
|
||||
def _capture(call_id: str, name: str, output: str, *, is_error: bool = False) -> None:
|
||||
captures.append((call_id, name, output, is_error))
|
||||
|
||||
session._report_tool_result = _capture # type: ignore[method-assign]
|
||||
return captures
|
||||
|
||||
|
||||
class TestExecMcpToolDispatchError:
|
||||
def test_exec_mcp_tool_preserves_structured_error_json(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.call_tool_sync.side_effect = RuntimeError(_CONSENT_REQUIRED_JSON)
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "tc_1",
|
||||
"mcp_func_name": "mcp__srv-oauth__do",
|
||||
"mcp_args": {},
|
||||
}
|
||||
session._exec_mcp_tool(item)
|
||||
|
||||
assert len(captures) == 1
|
||||
_, _, output, is_error = captures[0]
|
||||
assert output == _CONSENT_REQUIRED_JSON
|
||||
assert is_error is True
|
||||
|
||||
def test_exec_mcp_tool_prefixes_non_structured_error(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.call_tool_sync.side_effect = RuntimeError("connection lost")
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "tc_2",
|
||||
"mcp_func_name": "mcp__srv-oauth__do",
|
||||
"mcp_args": {},
|
||||
}
|
||||
session._exec_mcp_tool(item)
|
||||
|
||||
assert captures[0][2] == "MCP tool error: connection lost"
|
||||
assert captures[0][3] is True
|
||||
|
||||
|
||||
class TestExecReadResourceDispatchError:
|
||||
def test_exec_read_resource_preserves_structured_error_json(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.read_resource_sync.side_effect = RuntimeError(_CONSENT_REQUIRED_JSON)
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "rc_1",
|
||||
"resource_uri": "https://example.com/r",
|
||||
}
|
||||
# The exec site emits a ``log.warning`` (no ``exc_info`` — bearer-leak
|
||||
# invariant) on failure. Patch the logger so the test doesn't emit
|
||||
# noise to the captured stderr — assertions don't depend on log
|
||||
# output.
|
||||
with patch("turnstone.core.session.log"):
|
||||
session._exec_read_resource(item)
|
||||
|
||||
assert len(captures) == 1
|
||||
assert captures[0][2] == _CONSENT_REQUIRED_JSON
|
||||
assert captures[0][3] is True
|
||||
|
||||
def test_exec_read_resource_prefixes_non_structured_error(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.read_resource_sync.side_effect = RuntimeError("connection lost")
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "rc_2",
|
||||
"resource_uri": "https://example.com/r",
|
||||
}
|
||||
with patch("turnstone.core.session.log"):
|
||||
session._exec_read_resource(item)
|
||||
|
||||
assert captures[0][2] == "MCP resource error: connection lost"
|
||||
assert captures[0][3] is True
|
||||
|
||||
|
||||
class TestExecUsePromptDispatchError:
|
||||
def test_exec_use_prompt_preserves_structured_error_json(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_prompt_sync.side_effect = RuntimeError(_CONSENT_REQUIRED_JSON)
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "pc_1",
|
||||
"prompt_name": "mcp__srv-oauth__greet",
|
||||
"prompt_arguments": {},
|
||||
}
|
||||
with patch("turnstone.core.session.log"):
|
||||
session._exec_use_prompt(item)
|
||||
|
||||
assert len(captures) == 1
|
||||
assert captures[0][2] == _CONSENT_REQUIRED_JSON
|
||||
assert captures[0][3] is True
|
||||
|
||||
def test_exec_use_prompt_prefixes_non_structured_error(self, tmp_db) -> None:
|
||||
session = _make_session()
|
||||
captures = _record_outputs(session)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_prompt_sync.side_effect = RuntimeError("connection lost")
|
||||
session._mcp_client = mock_client
|
||||
|
||||
item = {
|
||||
"call_id": "pc_2",
|
||||
"prompt_name": "mcp__srv-oauth__greet",
|
||||
"prompt_arguments": {},
|
||||
}
|
||||
with patch("turnstone.core.session.log"):
|
||||
session._exec_use_prompt(item)
|
||||
|
||||
assert captures[0][2] == "MCP prompt error: connection lost"
|
||||
assert captures[0][3] is True
|
||||
+160
-15
@@ -523,8 +523,15 @@ class TestWorkstreamConfig:
|
||||
assert session.instructions == "be concise"
|
||||
assert session.creative_mode is True
|
||||
|
||||
def test_resume_restores_model(self, tmp_db):
|
||||
"""ChatSession.resume() should restore the model from workstream config."""
|
||||
def test_resume_keeps_defaults_when_alias_unresolvable(self, tmp_db):
|
||||
"""When the saved alias is empty or no longer in the registry,
|
||||
``resume()`` must NOT copy ``saved_model`` onto the constructor's
|
||||
default provider. Pairing a removed model name with a default
|
||||
provider that doesn't know about it produces a broken session
|
||||
whose next API call fails — the exact regression Copilot flagged
|
||||
on PR #465. The constructor already resolved a coherent default
|
||||
(provider + model + capabilities); resume should leave it intact
|
||||
and just log the unreachable saved values."""
|
||||
client = MagicMock()
|
||||
client.models.list.return_value.data = [MagicMock(id="test-model")]
|
||||
ui = MagicMock()
|
||||
@@ -533,13 +540,14 @@ class TestWorkstreamConfig:
|
||||
ui.on_state_change = MagicMock()
|
||||
ui.on_rename = MagicMock()
|
||||
|
||||
# Create a workstream that was using a specific model
|
||||
register_workstream("model_ws")
|
||||
save_message("model_ws", "user", "hello")
|
||||
save_message("model_ws", "assistant", "hi")
|
||||
# Empty alias + an orphan model name — same shape resume sees
|
||||
# when an operator removes an alias from the registry that the
|
||||
# workstream was originally pinned to.
|
||||
save_workstream_config("model_ws", {"model": "gpt-5", "model_alias": ""})
|
||||
|
||||
# Resume into a session that was created with a different model
|
||||
session = ChatSession(
|
||||
client=client,
|
||||
model="gpt-5-nano",
|
||||
@@ -552,7 +560,102 @@ class TestWorkstreamConfig:
|
||||
assert session.model == "gpt-5-nano"
|
||||
result = session.resume("model_ws")
|
||||
assert result is True
|
||||
assert session.model == "gpt-5"
|
||||
# Constructor's coherent default is preserved — saved orphan
|
||||
# model name is NOT copied over.
|
||||
assert session.model == "gpt-5-nano"
|
||||
|
||||
def test_init_does_not_clobber_existing_config(self, tmp_db):
|
||||
"""ChatSession.__init__ must NOT overwrite existing
|
||||
``workstream_config`` keys when constructing for an already-
|
||||
persisted ws_id.
|
||||
|
||||
This is the fix for the rehydrate bug: ``SessionManager.open()``
|
||||
builds a ChatSession with the persisted ws_id; the legacy
|
||||
``__init__`` unconditionally called ``_save_config()`` which is
|
||||
``INSERT OR REPLACE`` per-key — silently resetting model_alias,
|
||||
temperature, reasoning_effort, max_tokens, skill, creative_mode,
|
||||
and instructions to the constructor defaults *before*
|
||||
``resume()`` got a chance to read them back.
|
||||
"""
|
||||
client = MagicMock()
|
||||
client.models.list.return_value.data = [MagicMock(id="test-model")]
|
||||
ui = MagicMock()
|
||||
ui.on_info = MagicMock()
|
||||
ui.on_error = MagicMock()
|
||||
ui.on_state_change = MagicMock()
|
||||
ui.on_rename = MagicMock()
|
||||
|
||||
register_workstream("rehydrate_ws")
|
||||
save_workstream_config(
|
||||
"rehydrate_ws",
|
||||
{
|
||||
"model": "gpt-5-pro",
|
||||
"model_alias": "gpt-5-pro",
|
||||
"temperature": "0.2",
|
||||
"reasoning_effort": "high",
|
||||
"max_tokens": "8192",
|
||||
"creative_mode": "True",
|
||||
"instructions": "preserve me",
|
||||
},
|
||||
)
|
||||
|
||||
ChatSession(
|
||||
client=client,
|
||||
model="some-default-model",
|
||||
ui=ui,
|
||||
instructions=None,
|
||||
temperature=0.7,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
reasoning_effort="medium",
|
||||
ws_id="rehydrate_ws",
|
||||
)
|
||||
|
||||
loaded = load_workstream_config("rehydrate_ws")
|
||||
assert loaded["model"] == "gpt-5-pro"
|
||||
assert loaded["model_alias"] == "gpt-5-pro"
|
||||
assert loaded["temperature"] == "0.2"
|
||||
assert loaded["reasoning_effort"] == "high"
|
||||
assert loaded["max_tokens"] == "8192"
|
||||
assert loaded["creative_mode"] == "True"
|
||||
assert loaded["instructions"] == "preserve me"
|
||||
|
||||
def test_init_writes_config_on_fresh_create(self, tmp_db):
|
||||
"""The opposite half of the contract: when no config row exists
|
||||
yet, ``__init__`` must still persist the constructor's values so
|
||||
a later resume can find them. This is the path that previously
|
||||
worked — the fix must not break it."""
|
||||
client = MagicMock()
|
||||
client.models.list.return_value.data = [MagicMock(id="test-model")]
|
||||
ui = MagicMock()
|
||||
ui.on_info = MagicMock()
|
||||
ui.on_error = MagicMock()
|
||||
ui.on_state_change = MagicMock()
|
||||
ui.on_rename = MagicMock()
|
||||
|
||||
# No save_workstream_config() before ChatSession() — this is
|
||||
# the fresh-create path the SessionManager.create() flow takes.
|
||||
register_workstream("fresh_ws")
|
||||
assert load_workstream_config("fresh_ws") == {}
|
||||
|
||||
ChatSession(
|
||||
client=client,
|
||||
model="gpt-5-mini",
|
||||
ui=ui,
|
||||
instructions="be terse",
|
||||
temperature=0.4,
|
||||
max_tokens=2048,
|
||||
tool_timeout=30,
|
||||
reasoning_effort="low",
|
||||
ws_id="fresh_ws",
|
||||
)
|
||||
|
||||
loaded = load_workstream_config("fresh_ws")
|
||||
assert loaded["model"] == "gpt-5-mini"
|
||||
assert loaded["temperature"] == "0.4"
|
||||
assert loaded["reasoning_effort"] == "low"
|
||||
assert loaded["max_tokens"] == "2048"
|
||||
assert loaded["instructions"] == "be terse"
|
||||
|
||||
|
||||
# ── Prune workstreams ─────────────────────────────────────────────────
|
||||
@@ -811,8 +914,11 @@ class TestMCPToolGating:
|
||||
"""read_resource excluded when MCP client has no resources."""
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
mcp_client.resource_count = 0
|
||||
mcp_client.prompt_count = 2
|
||||
# Phase 7b: gating uses ``*_count_for_user`` so the test mocks
|
||||
# the per-user variant (the property remains for static-only
|
||||
# admin paths). Returning 0 / 2 mirrors the prior contract.
|
||||
mcp_client.resource_count_for_user.return_value = 0
|
||||
mcp_client.prompt_count_for_user.return_value = 2
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
@@ -834,8 +940,8 @@ class TestMCPToolGating:
|
||||
"""use_prompt excluded when MCP client has no prompts."""
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
mcp_client.resource_count = 3
|
||||
mcp_client.prompt_count = 0
|
||||
mcp_client.resource_count_for_user.return_value = 3
|
||||
mcp_client.prompt_count_for_user.return_value = 0
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
@@ -857,8 +963,8 @@ class TestMCPToolGating:
|
||||
"""Both tools present when MCP client has resources and prompts."""
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
mcp_client.resource_count = 1
|
||||
mcp_client.prompt_count = 1
|
||||
mcp_client.resource_count_for_user.return_value = 1
|
||||
mcp_client.prompt_count_for_user.return_value = 1
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
@@ -880,8 +986,8 @@ class TestMCPToolGating:
|
||||
"""Gating applies even when tool_search is active (client-side path)."""
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
mcp_client.resource_count = 0
|
||||
mcp_client.prompt_count = 0
|
||||
mcp_client.resource_count_for_user.return_value = 0
|
||||
mcp_client.prompt_count_for_user.return_value = 0
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
@@ -909,8 +1015,8 @@ class TestMCPToolGating:
|
||||
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
mcp_client.resource_count = 0
|
||||
mcp_client.prompt_count = 0
|
||||
mcp_client.resource_count_for_user.return_value = 0
|
||||
mcp_client.prompt_count_for_user.return_value = 0
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
@@ -931,3 +1037,42 @@ class TestMCPToolGating:
|
||||
names = [t.get("function", {}).get("name") for t in tools]
|
||||
assert "read_resource" not in names
|
||||
assert "use_prompt" not in names
|
||||
|
||||
def test_pool_only_user_keeps_read_resource_and_use_prompt(self, tmp_db, mock_openai_client):
|
||||
"""Phase 7b canary: a pool-only user (static catalog empty) still
|
||||
sees ``read_resource`` and ``use_prompt`` because the gating
|
||||
consults ``*_count_for_user`` (scope decision 0.2).
|
||||
|
||||
Drives ``resource_count = prompt_count = 0`` (the static-only
|
||||
properties are zero) but ``*_count_for_user(uid) > 0`` because
|
||||
the user has pool entries; the tools must remain visible.
|
||||
"""
|
||||
mcp_client = MagicMock()
|
||||
mcp_client.get_tools.return_value = []
|
||||
# Static catalog is empty; admin-style legacy properties say 0.
|
||||
mcp_client.resource_count = 0
|
||||
mcp_client.prompt_count = 0
|
||||
# Per-user variant reports the user's pool entries.
|
||||
mcp_client.resource_count_for_user.return_value = 2
|
||||
mcp_client.prompt_count_for_user.return_value = 1
|
||||
|
||||
session = ChatSession(
|
||||
client=mock_openai_client,
|
||||
model="local-model",
|
||||
ui=MagicMock(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=1000,
|
||||
tool_timeout=10,
|
||||
mcp_client=mcp_client,
|
||||
user_id="pool-only-user",
|
||||
)
|
||||
|
||||
tools = session._get_active_tools()
|
||||
names = [t.get("function", {}).get("name") for t in tools]
|
||||
assert "read_resource" in names
|
||||
assert "use_prompt" in names
|
||||
# Verify the per-user gate was actually consulted with the
|
||||
# session's ``user_id`` (sanity-check on the wiring).
|
||||
mcp_client.resource_count_for_user.assert_any_call("pool-only-user")
|
||||
mcp_client.prompt_count_for_user.assert_any_call("pool-only-user")
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Tests for the SKILL.md parse admin API endpoint.
|
||||
|
||||
The endpoint is a thin permission-checked wrapper around
|
||||
``turnstone.core.skill_parser.parse_skill_md``. These tests cover the
|
||||
routing, auth, and error-handling layers — parser semantics live in
|
||||
``test_skill_parser.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.routing import Mount, Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
from turnstone.console.server import admin_parse_skill
|
||||
from turnstone.core.auth import AuthResult
|
||||
|
||||
|
||||
class _InjectAuthMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="test-user",
|
||||
scopes=frozenset({"approve"}),
|
||||
token_source="config",
|
||||
permissions=frozenset({"read", "write", "approve", "admin.skills"}),
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
class _InjectAuthNoSkillsMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="test-user",
|
||||
scopes=frozenset({"approve"}),
|
||||
token_source="jwt",
|
||||
permissions=frozenset({"read", "write", "approve"}),
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
_ROUTES = [
|
||||
Mount(
|
||||
"/v1",
|
||||
routes=[
|
||||
Route("/api/admin/skills/parse", admin_parse_skill, methods=["POST"]),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> TestClient:
|
||||
app = Starlette(
|
||||
routes=_ROUTES,
|
||||
middleware=[Middleware(_InjectAuthMiddleware)],
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_no_perm() -> TestClient:
|
||||
app = Starlette(
|
||||
routes=_ROUTES,
|
||||
middleware=[Middleware(_InjectAuthNoSkillsMiddleware)],
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
_FULL_SKILL = """\
|
||||
---
|
||||
name: code-review
|
||||
description: Automated code review skill
|
||||
author: Test Author
|
||||
version: 2.0.0
|
||||
tags: [python, review, quality]
|
||||
allowed-tools: [read_file, list_directory]
|
||||
license: MIT
|
||||
compatibility: ">=0.7"
|
||||
---
|
||||
|
||||
# Code Review
|
||||
|
||||
Review code for best practices.
|
||||
"""
|
||||
|
||||
_MINIMAL_SKILL = """\
|
||||
---
|
||||
name: minimal
|
||||
---
|
||||
|
||||
Just some content.
|
||||
"""
|
||||
|
||||
|
||||
class TestParseSkill:
|
||||
def test_parses_full_frontmatter(self, client: TestClient) -> None:
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": _FULL_SKILL})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["name"] == "code-review"
|
||||
assert data["description"] == "Automated code review skill"
|
||||
assert data["author"] == "Test Author"
|
||||
assert data["version"] == "2.0.0"
|
||||
assert data["tags"] == ["python", "review", "quality"]
|
||||
assert data["allowed_tools"] == ["read_file", "list_directory"]
|
||||
assert data["license"] == "MIT"
|
||||
assert data["compatibility"] == ">=0.7"
|
||||
assert "# Code Review" in data["content"]
|
||||
# Frontmatter should not leak into the body.
|
||||
assert "name: code-review" not in data["content"]
|
||||
# ParsedSkill carries raw_frontmatter (the full YAML dict) but the
|
||||
# handler whitelists fields by hand to avoid leaking arbitrary keys.
|
||||
# Pin that contract — a future refactor to dataclasses.asdict would
|
||||
# silently break it without this assertion.
|
||||
assert "raw_frontmatter" not in data
|
||||
|
||||
def test_parses_minimal_frontmatter(self, client: TestClient) -> None:
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": _MINIMAL_SKILL})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["name"] == "minimal"
|
||||
assert data["description"] == "Just some content."
|
||||
assert data["version"] == "1.0.0"
|
||||
assert data["tags"] == []
|
||||
assert data["allowed_tools"] == []
|
||||
assert data["license"] == ""
|
||||
|
||||
def test_anthropic_nested_metadata_tags(self, client: TestClient) -> None:
|
||||
# Anthropic-style skill puts tags under metadata.tags rather than
|
||||
# at the top level — the parser must handle both layouts.
|
||||
raw = """\
|
||||
---
|
||||
name: nested-meta
|
||||
description: A skill using nested metadata
|
||||
metadata:
|
||||
tags: [alpha, beta]
|
||||
author: Anthropic
|
||||
version: 3.1.4
|
||||
---
|
||||
|
||||
Body.
|
||||
"""
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": raw})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["tags"] == ["alpha", "beta"]
|
||||
assert data["author"] == "Anthropic"
|
||||
assert data["version"] == "3.1.4"
|
||||
|
||||
def test_unquoted_colon_in_description(self, client: TestClient) -> None:
|
||||
# Common cross-client mistake: ``description: Use when: the user...``
|
||||
# The parser retries with the description value quoted.
|
||||
raw = """\
|
||||
---
|
||||
name: colon-desc
|
||||
description: Use when: the user asks for a review
|
||||
---
|
||||
|
||||
Body.
|
||||
"""
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": raw})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["name"] == "colon-desc"
|
||||
assert "Use when" in data["description"]
|
||||
|
||||
def test_missing_name_returns_400(self, client: TestClient) -> None:
|
||||
raw = """\
|
||||
---
|
||||
description: No name field
|
||||
---
|
||||
|
||||
Body.
|
||||
"""
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": raw})
|
||||
assert resp.status_code == 400
|
||||
assert "name" in resp.json()["error"].lower()
|
||||
|
||||
def test_missing_raw_returns_400(self, client: TestClient) -> None:
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={})
|
||||
assert resp.status_code == 400
|
||||
assert "raw" in resp.json()["error"].lower()
|
||||
|
||||
def test_blank_raw_returns_400(self, client: TestClient) -> None:
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": " \n"})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_invalid_yaml_returns_400(self, client: TestClient) -> None:
|
||||
# YAML that the malformed-description retry can't fix.
|
||||
raw = "---\nname: [not, valid, here\n---\nBody.\n"
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": raw})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_requires_admin_skills_permission(self, client_no_perm: TestClient) -> None:
|
||||
resp = client_no_perm.post("/v1/api/admin/skills/parse", json={"raw": _MINIMAL_SKILL})
|
||||
assert resp.status_code == 403
|
||||
|
||||
def test_oversized_content_length_returns_413(self, client: TestClient) -> None:
|
||||
# Content-Length pre-check rejects oversized bodies before they're
|
||||
# buffered into memory. Caps worker memory against an admin-token
|
||||
# holder spraying multi-GB JSON. The threshold is generous (~4×
|
||||
# the per-string cap) so payload here must clearly exceed it.
|
||||
oversized = "a" * 200_000
|
||||
resp = client.post("/v1/api/admin/skills/parse", json={"raw": oversized})
|
||||
assert resp.status_code == 413
|
||||
|
||||
def test_oversized_raw_chunked_returns_413(self, client: TestClient) -> None:
|
||||
# When the client sends Transfer-Encoding: chunked there is no
|
||||
# Content-Length header, so the pre-check is skipped and the
|
||||
# application-layer cap is the only line of defence. httpx switches
|
||||
# to chunked when the body is a generator.
|
||||
def _gen() -> Iterator[bytes]:
|
||||
yield b'{"raw":"' + b"a" * 33_000 + b'"}'
|
||||
|
||||
resp = client.post(
|
||||
"/v1/api/admin/skills/parse",
|
||||
content=_gen(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
assert resp.status_code == 413
|
||||
assert "raw" in resp.json()["error"].lower()
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Tests for ``_source`` / ``_reminders`` round-tripping through both
|
||||
storage backends.
|
||||
|
||||
Persisting the in-memory side-channels lets multi-tab / multi-device
|
||||
replay show the same metacognitive bubble shape the originating tab
|
||||
saw live — see ``docs/design/watch-card-ux.md`` §1.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from turnstone.core.storage._schema import conversations
|
||||
|
||||
|
||||
class TestSourceRoundtrip:
|
||||
def test_source_roundtrip(self, backend):
|
||||
backend.register_workstream("s1")
|
||||
backend.save_message("s1", "user", "", source="system_nudge")
|
||||
msgs = backend.load_messages("s1")
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["role"] == "user"
|
||||
assert msgs[0]["content"] == ""
|
||||
assert msgs[0].get("_source") == "system_nudge"
|
||||
|
||||
def test_source_absent_when_not_set(self, backend):
|
||||
backend.register_workstream("s1")
|
||||
backend.save_message("s1", "user", "hello")
|
||||
msgs = backend.load_messages("s1")
|
||||
assert "_source" not in msgs[0]
|
||||
|
||||
|
||||
class TestRemindersRoundtrip:
|
||||
def test_reminders_roundtrip(self, backend):
|
||||
backend.register_workstream("s1")
|
||||
payload = [
|
||||
{
|
||||
"type": "watch_triggered",
|
||||
"text": "$ ls\nfile.txt\n",
|
||||
"watch_name": "w1",
|
||||
"command": "ls",
|
||||
"poll_count": 2,
|
||||
"max_polls": 100,
|
||||
"is_final": False,
|
||||
}
|
||||
]
|
||||
backend.save_message(
|
||||
"s1",
|
||||
"user",
|
||||
"",
|
||||
source="system_nudge",
|
||||
reminders=json.dumps(payload, separators=(",", ":")),
|
||||
)
|
||||
msgs = backend.load_messages("s1")
|
||||
assert msgs[0].get("_reminders") == payload
|
||||
# Optional fields preserved verbatim.
|
||||
rem = msgs[0]["_reminders"][0]
|
||||
assert rem["watch_name"] == "w1"
|
||||
assert rem["command"] == "ls"
|
||||
assert rem["poll_count"] == 2
|
||||
assert rem["max_polls"] == 100
|
||||
assert rem["is_final"] is False
|
||||
|
||||
def test_reminders_null_renders_as_no_key(self, backend):
|
||||
"""Absent vs. empty-list should map to the same shape on the
|
||||
load side: ``_reminders`` simply not present in the dict.
|
||||
Mirrors the ``_attachments_meta`` precedent in
|
||||
``reconstruct_messages``.
|
||||
"""
|
||||
backend.register_workstream("s1")
|
||||
backend.save_message("s1", "user", "hello")
|
||||
msgs = backend.load_messages("s1")
|
||||
assert "_reminders" not in msgs[0]
|
||||
|
||||
def test_tool_reminders_roundtrip(self, backend):
|
||||
backend.register_workstream("s1")
|
||||
# Build a minimal valid history: assistant turn with one
|
||||
# tool_call followed by the tool result that carries the
|
||||
# tool-channel reminder. Without the assistant turn the
|
||||
# tool row would be orphaned and stripped by the repair pass.
|
||||
tc_json = json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "c1",
|
||||
"type": "function",
|
||||
"function": {"name": "bash", "arguments": "{}"},
|
||||
}
|
||||
]
|
||||
)
|
||||
backend.save_message("s1", "user", "go")
|
||||
backend.save_message("s1", "assistant", None, tool_calls=tc_json)
|
||||
payload = [{"type": "tool_error", "text": "command failed"}]
|
||||
backend.save_message(
|
||||
"s1",
|
||||
"tool",
|
||||
"boom",
|
||||
tool_call_id="c1",
|
||||
reminders=json.dumps(payload, separators=(",", ":")),
|
||||
)
|
||||
msgs = backend.load_messages("s1")
|
||||
# Find the tool message and assert reminders survived load.
|
||||
tool_msgs = [m for m in msgs if m.get("role") == "tool"]
|
||||
assert len(tool_msgs) == 1
|
||||
assert tool_msgs[0].get("_reminders") == payload
|
||||
|
||||
def test_nul_bytes_stripped_from_source_and_reminders(self, backend):
|
||||
"""NUL bytes must be stripped at the storage layer.
|
||||
|
||||
Producers (``sanitize_payload`` on the watch dispatch path,
|
||||
constants for non-watch nudges) already strip NUL today so
|
||||
nothing in production reaches this clamp — but the layer is
|
||||
the tripwire if a future producer forgets, mirroring how
|
||||
``content`` and ``provider_data`` are sanitized. PostgreSQL
|
||||
TEXT columns reject NUL outright, so the sanitization is also
|
||||
a hard correctness invariant on that backend.
|
||||
|
||||
``json.dumps`` already escapes NUL inside string values to
|
||||
``\\u0000`` so a real NUL byte can't enter ``_reminders`` via
|
||||
the normal encode path — the test feeds a raw NUL directly to
|
||||
cover the bypass case (a future producer that hand-builds the
|
||||
column string).
|
||||
"""
|
||||
backend.register_workstream("s1")
|
||||
backend.save_message(
|
||||
"s1",
|
||||
"user",
|
||||
"",
|
||||
source="system_nudge\x00",
|
||||
reminders='[{"type":"watch_triggered","text":"ok\x00bad"}]',
|
||||
)
|
||||
msgs = backend.load_messages("s1")
|
||||
assert msgs[0].get("_source") == "system_nudge"
|
||||
assert msgs[0].get("_reminders") == [{"type": "watch_triggered", "text": "okbad"}]
|
||||
|
||||
def test_malformed_reminders_json_does_not_crash_load(self, backend):
|
||||
"""A garbage string in the column must not abort the whole
|
||||
load — mirrors the ``provider_data`` JSON-decode-suppress
|
||||
pattern. Concretely: write a row with valid columns BUT a
|
||||
corrupted ``_reminders`` value via raw SQL, then verify the
|
||||
load returns the message with no ``_reminders`` key (rather
|
||||
than raising or surfacing the garbage).
|
||||
"""
|
||||
backend.register_workstream("s1")
|
||||
msg_id = backend.save_message("s1", "user", "hello")
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.update(conversations)
|
||||
.where(conversations.c.id == msg_id)
|
||||
.values(_reminders="this is not json {{")
|
||||
)
|
||||
conn.commit()
|
||||
msgs = backend.load_messages("s1")
|
||||
assert len(msgs) == 1
|
||||
# Garbage suppressed silently — key absent, content intact.
|
||||
assert "_reminders" not in msgs[0]
|
||||
assert msgs[0]["content"] == "hello"
|
||||
@@ -4,6 +4,10 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from turnstone.core.storage._schema import workstreams
|
||||
|
||||
# -- Workstream registration ---------------------------------------------------
|
||||
|
||||
|
||||
@@ -664,6 +668,335 @@ class TestBatchPrimitives:
|
||||
assert result == {"never-seen": 0}
|
||||
|
||||
|
||||
# -- bulk_close_stale_orphans --------------------------------------------------
|
||||
|
||||
|
||||
def _force_updated(backend: Any, ws_id: str, updated: str) -> None:
|
||||
"""Stamp a workstream row's ``updated`` column directly.
|
||||
|
||||
The public surface only sets ``updated`` to ``now``, which makes it
|
||||
impossible to fabricate a stale row through register/update calls.
|
||||
Reaches into ``backend._engine`` — same access pattern conftest uses
|
||||
for cross-backend cleanup.
|
||||
"""
|
||||
with backend._engine.connect() as conn:
|
||||
conn.execute(
|
||||
sa.update(workstreams).where(workstreams.c.ws_id == ws_id).values(updated=updated)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
class TestBulkCloseStaleOrphans:
|
||||
def test_closes_stale_non_terminal_rows_of_kind(self, backend):
|
||||
backend.register_workstream("stale-idle", kind="interactive")
|
||||
backend.register_workstream("stale-thinking", kind="interactive")
|
||||
backend.update_workstream_state("stale-thinking", "thinking")
|
||||
backend.register_workstream("fresh-idle", kind="interactive")
|
||||
_force_updated(backend, "stale-idle", "2020-01-01T00:00:00")
|
||||
_force_updated(backend, "stale-thinking", "2020-01-01T00:00:00")
|
||||
# fresh-idle stays at registration time (effectively now)
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert set(closed) == {"stale-idle", "stale-thinking"}
|
||||
rows = backend.get_workstreams_batch(["stale-idle", "stale-thinking", "fresh-idle"])
|
||||
assert rows["stale-idle"]["state"] == "closed"
|
||||
assert rows["stale-thinking"]["state"] == "closed"
|
||||
assert rows["fresh-idle"]["state"] == "idle"
|
||||
|
||||
def test_skips_already_closed(self, backend):
|
||||
backend.register_workstream("already-closed", kind="interactive")
|
||||
backend.update_workstream_state("already-closed", "closed")
|
||||
_force_updated(backend, "already-closed", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert closed == []
|
||||
|
||||
def test_filters_by_kind(self, backend):
|
||||
backend.register_workstream("interactive-stale", kind="interactive")
|
||||
backend.register_workstream("coord-stale", kind="coordinator")
|
||||
_force_updated(backend, "interactive-stale", "2020-01-01T00:00:00")
|
||||
_force_updated(backend, "coord-stale", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert closed == ["interactive-stale"]
|
||||
rows = backend.get_workstreams_batch(["interactive-stale", "coord-stale"])
|
||||
assert rows["interactive-stale"]["state"] == "closed"
|
||||
assert rows["coord-stale"]["state"] == "idle"
|
||||
|
||||
def test_excludes_loaded_ws_ids(self, backend):
|
||||
backend.register_workstream("ws-keep", kind="interactive")
|
||||
backend.register_workstream("ws-close", kind="interactive")
|
||||
_force_updated(backend, "ws-keep", "2020-01-01T00:00:00")
|
||||
_force_updated(backend, "ws-close", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=["ws-keep"]
|
||||
)
|
||||
|
||||
assert closed == ["ws-close"]
|
||||
rows = backend.get_workstreams_batch(["ws-keep", "ws-close"])
|
||||
assert rows["ws-keep"]["state"] == "idle"
|
||||
assert rows["ws-close"]["state"] == "closed"
|
||||
|
||||
def test_empty_exclude_list_does_not_break_sql(self, backend):
|
||||
backend.register_workstream("orphan", kind="interactive")
|
||||
_force_updated(backend, "orphan", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert closed == ["orphan"]
|
||||
|
||||
def test_no_orphans_returns_empty(self, backend):
|
||||
backend.register_workstream("fresh", kind="interactive")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert closed == []
|
||||
|
||||
def test_closes_all_non_terminal_states(self, backend):
|
||||
for ws_id, state in [
|
||||
("o-idle", "idle"),
|
||||
("o-thinking", "thinking"),
|
||||
("o-attention", "attention"),
|
||||
("o-running", "running"),
|
||||
]:
|
||||
backend.register_workstream(ws_id, kind="interactive")
|
||||
if state != "idle":
|
||||
backend.update_workstream_state(ws_id, state)
|
||||
_force_updated(backend, ws_id, "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert set(closed) == {"o-idle", "o-thinking", "o-attention", "o-running"}
|
||||
|
||||
def test_bumps_updated_on_close(self, backend):
|
||||
stale_updated = "2020-01-01T00:00:00"
|
||||
backend.register_workstream("orphan", kind="interactive")
|
||||
_force_updated(backend, "orphan", stale_updated)
|
||||
|
||||
backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
# ``updated`` must change away from the forced stale value. Asserting
|
||||
# inequality from the seed (rather than ``> "2024-01-01..."``) keeps
|
||||
# the test independent of wall-clock date.
|
||||
with backend._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
sa.select(workstreams.c.updated).where(workstreams.c.ws_id == "orphan")
|
||||
).one()
|
||||
assert row[0] != stale_updated
|
||||
|
||||
def test_protects_rows_owned_by_live_services(self, backend):
|
||||
"""Liveness scoping (post-#384 rendezvous-routing world): rows
|
||||
whose ``node_id`` matches a heartbeating service must NOT be
|
||||
reaped, because that owner may legitimately have them loaded on
|
||||
another worker. Rows whose ``node_id`` matches a dead service
|
||||
ARE eligible — that's how dead-pod orphans get reclaimed in
|
||||
containerized deployments with dynamic hostnames."""
|
||||
backend.register_workstream("dead-node", node_id="dead-pod-x4k2", kind="interactive")
|
||||
backend.register_workstream("alive-node", node_id="alive-pod-y9p3", kind="interactive")
|
||||
_force_updated(backend, "dead-node", "2020-01-01T00:00:00")
|
||||
_force_updated(backend, "alive-node", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive",
|
||||
cutoff="2024-01-01T00:00:00",
|
||||
exclude_ws_ids=[],
|
||||
live_node_ids=["alive-pod-y9p3"],
|
||||
)
|
||||
|
||||
assert closed == ["dead-node"]
|
||||
rows = backend.get_workstreams_batch(["dead-node", "alive-node"])
|
||||
assert rows["dead-node"]["state"] == "closed"
|
||||
assert rows["alive-node"]["state"] == "idle"
|
||||
|
||||
def test_null_node_id_always_eligible(self, backend):
|
||||
"""A row with NULL ``node_id`` has no owner identity — age alone
|
||||
gates the reap. Belt-and-suspenders against ``NULL NOT IN (...)``
|
||||
evaluating to NULL (not TRUE) and silently protecting orphans
|
||||
forever."""
|
||||
backend.register_workstream("no-owner", node_id=None, kind="interactive")
|
||||
_force_updated(backend, "no-owner", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive",
|
||||
cutoff="2024-01-01T00:00:00",
|
||||
exclude_ws_ids=[],
|
||||
live_node_ids=["some-other-node"],
|
||||
)
|
||||
|
||||
assert closed == ["no-owner"]
|
||||
|
||||
def test_live_node_ids_none_skips_filter(self, backend):
|
||||
"""``live_node_ids=None`` is the single-process / operator-backfill
|
||||
mode — all rows of *kind* are eligible regardless of node_id."""
|
||||
backend.register_workstream("node-a", node_id="node-a", kind="interactive")
|
||||
backend.register_workstream("node-b", node_id="node-b", kind="interactive")
|
||||
_force_updated(backend, "node-a", "2020-01-01T00:00:00")
|
||||
_force_updated(backend, "node-b", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive", cutoff="2024-01-01T00:00:00", exclude_ws_ids=[]
|
||||
)
|
||||
|
||||
assert set(closed) == {"node-a", "node-b"}
|
||||
|
||||
def test_empty_live_node_ids_treats_all_as_dead(self, backend):
|
||||
"""Empty list ``live_node_ids=[]`` means "no nodes alive" — every
|
||||
row's owner is unprotected. Useful for operator scripts that
|
||||
want to reap regardless of liveness."""
|
||||
backend.register_workstream("any", node_id="node-a", kind="interactive")
|
||||
_force_updated(backend, "any", "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive",
|
||||
cutoff="2024-01-01T00:00:00",
|
||||
exclude_ws_ids=[],
|
||||
live_node_ids=[],
|
||||
)
|
||||
|
||||
assert closed == ["any"]
|
||||
|
||||
def test_combines_live_node_ids_and_exclude_ws_ids(self, backend):
|
||||
"""Both filters stack as AND clauses on the UPDATE. Covers the
|
||||
full 2x2 matrix to catch a future edit that replaces an AND with
|
||||
an OR or drops one of the filters: only the (orphan + dead-node)
|
||||
cell should be reaped."""
|
||||
# All four registered with the same stale ``updated``.
|
||||
for ws_id, node in [
|
||||
("loaded-alive", "alive-node"),
|
||||
("loaded-dead", "dead-node"),
|
||||
("orphan-alive", "alive-node"),
|
||||
("orphan-dead", "dead-node"),
|
||||
]:
|
||||
backend.register_workstream(ws_id, node_id=node, kind="interactive")
|
||||
_force_updated(backend, ws_id, "2020-01-01T00:00:00")
|
||||
|
||||
closed = backend.bulk_close_stale_orphans(
|
||||
"interactive",
|
||||
cutoff="2024-01-01T00:00:00",
|
||||
exclude_ws_ids=["loaded-alive", "loaded-dead"],
|
||||
live_node_ids=["alive-node"],
|
||||
)
|
||||
|
||||
# Only orphan-dead is unprotected by both filters.
|
||||
assert closed == ["orphan-dead"]
|
||||
rows = backend.get_workstreams_batch(
|
||||
["loaded-alive", "loaded-dead", "orphan-alive", "orphan-dead"]
|
||||
)
|
||||
assert rows["loaded-alive"]["state"] == "idle"
|
||||
assert rows["loaded-dead"]["state"] == "idle"
|
||||
assert rows["orphan-alive"]["state"] == "idle"
|
||||
assert rows["orphan-dead"]["state"] == "closed"
|
||||
|
||||
|
||||
# -- touch_workstream ----------------------------------------------------------
|
||||
|
||||
|
||||
class TestTouchWorkstream:
|
||||
def test_bumps_updated_only(self, backend):
|
||||
"""Used by ``open()`` on rehydrate to defend against the orphan
|
||||
reaper clobbering a freshly-loaded row. Must not change ``state``
|
||||
(the open() path explicitly avoids state writes to dodge a race
|
||||
with concurrent close())."""
|
||||
stale_updated = "2020-01-01T00:00:00"
|
||||
backend.register_workstream("ws-touch", kind="interactive")
|
||||
backend.update_workstream_state("ws-touch", "closed") # simulate prior close
|
||||
_force_updated(backend, "ws-touch", stale_updated)
|
||||
|
||||
backend.touch_workstream("ws-touch")
|
||||
|
||||
with backend._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
sa.select(workstreams.c.state, workstreams.c.updated).where(
|
||||
workstreams.c.ws_id == "ws-touch"
|
||||
)
|
||||
).one()
|
||||
assert row[0] == "closed", "state must not be modified by touch"
|
||||
# Compare against the forced stale value rather than a fixed calendar
|
||||
# date so the test is independent of wall-clock time.
|
||||
assert row[1] != stale_updated, "updated must be bumped"
|
||||
|
||||
def test_unknown_id_is_noop(self, backend):
|
||||
"""Touch on a missing id must not raise — open()'s exception
|
||||
handler is best-effort."""
|
||||
backend.touch_workstream("nonexistent") # must not raise
|
||||
|
||||
|
||||
# -- MCP OAuth columns ---------------------------------------------------------
|
||||
|
||||
|
||||
class TestMcpServerOauthColumns:
|
||||
def test_mcp_servers_oauth_columns_round_trip(self, backend: Any) -> None:
|
||||
"""An oauth_user row round-trips through create -> get with all
|
||||
seven OAuth text columns intact."""
|
||||
sid = "oauth-srv-1"
|
||||
backend.create_mcp_server(
|
||||
server_id=sid,
|
||||
name="oauth-srv",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="cli_abc123",
|
||||
oauth_scopes="openid profile",
|
||||
oauth_audience="https://mcp.example.com",
|
||||
oauth_registration_mode="preregistered",
|
||||
oauth_authorization_server_url="https://auth.example.com",
|
||||
oauth_as_issuer_cached="https://auth.example.com",
|
||||
)
|
||||
s = backend.get_mcp_server(sid)
|
||||
assert s is not None
|
||||
assert s["auth_type"] == "oauth_user"
|
||||
assert s["oauth_client_id"] == "cli_abc123"
|
||||
assert s["oauth_scopes"] == "openid profile"
|
||||
assert s["oauth_audience"] == "https://mcp.example.com"
|
||||
assert s["oauth_registration_mode"] == "preregistered"
|
||||
assert s["oauth_authorization_server_url"] == "https://auth.example.com"
|
||||
assert s["oauth_as_issuer_cached"] == "https://auth.example.com"
|
||||
# Phase 2 leaves the ciphertext slot NULL even when other oauth
|
||||
# fields are populated; Phase 3 wires the encryption write path.
|
||||
assert s["oauth_client_secret_ct"] is None
|
||||
|
||||
def test_update_auth_type_static_to_oauth(self, backend: Any) -> None:
|
||||
sid = "oauth-srv-2"
|
||||
backend.create_mcp_server(
|
||||
server_id=sid,
|
||||
name="static-then-oauth",
|
||||
transport="streamable-http",
|
||||
url="https://mcp.example.com/sse",
|
||||
)
|
||||
assert backend.get_mcp_server(sid)["auth_type"] == "static"
|
||||
|
||||
ok = backend.update_mcp_server(
|
||||
sid,
|
||||
auth_type="oauth_user",
|
||||
oauth_client_id="cli_after",
|
||||
oauth_audience="https://mcp.example.com",
|
||||
)
|
||||
assert ok is True
|
||||
s = backend.get_mcp_server(sid)
|
||||
assert s is not None
|
||||
assert s["auth_type"] == "oauth_user"
|
||||
assert s["oauth_client_id"] == "cli_after"
|
||||
assert s["oauth_audience"] == "https://mcp.example.com"
|
||||
|
||||
|
||||
# -- Lifecycle -----------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
@@ -67,6 +67,42 @@ class TestSearchStructuredMemories:
|
||||
assert len(results) >= 1
|
||||
assert any(r["name"] == "db_host" for r in results)
|
||||
|
||||
def test_multiword_or_matches_partial(self, tmp_db):
|
||||
"""OR-of-terms: memory matching only 1 of 3 query terms is returned."""
|
||||
save_structured_memory("postgres_config", "host=localhost port=5432")
|
||||
save_structured_memory("redis_config", "host=redis port=6379")
|
||||
save_structured_memory("unrelated", "nothing relevant here")
|
||||
|
||||
# "postgres missing_word_a missing_word_b": only postgres_config matches "postgres"
|
||||
results = search_structured_memories("postgres missing_word_a missing_word_b")
|
||||
names = {r["name"] for r in results}
|
||||
assert "postgres_config" in names
|
||||
assert "unrelated" not in names
|
||||
|
||||
def test_multiword_or_multiple_partial_matches(self, tmp_db):
|
||||
"""Multiple memories each matching different terms are all returned."""
|
||||
save_structured_memory("key_alpha", "alpha content here")
|
||||
save_structured_memory("key_beta", "beta content here")
|
||||
save_structured_memory("key_other", "completely different")
|
||||
|
||||
results = search_structured_memories("alpha beta")
|
||||
names = {r["name"] for r in results}
|
||||
assert "key_alpha" in names
|
||||
assert "key_beta" in names
|
||||
assert "key_other" not in names
|
||||
|
||||
def test_search_scope_filtering_preserved(self, tmp_db):
|
||||
"""Search with scope filter only returns memories in that scope."""
|
||||
save_structured_memory("ws1_fact", "alpha info", scope="workstream", scope_id="ws1")
|
||||
save_structured_memory("ws2_fact", "alpha info", scope="workstream", scope_id="ws2")
|
||||
save_structured_memory("global_fact", "alpha info", scope="global")
|
||||
|
||||
results = search_structured_memories("alpha", scope="workstream", scope_id="ws1")
|
||||
names = {r["name"] for r in results}
|
||||
assert "ws1_fact" in names
|
||||
assert "ws2_fact" not in names
|
||||
assert "global_fact" not in names
|
||||
|
||||
|
||||
class TestGetStructuredMemoryByName:
|
||||
def test_get_existing(self, tmp_db):
|
||||
|
||||
@@ -126,3 +126,148 @@ class TestCount:
|
||||
backend.create_structured_memory("m2", "b", "", "project", "workstream", "ws1", "2")
|
||||
assert backend.count_structured_memories(scope="global") == 1
|
||||
assert backend.count_structured_memories(scope="workstream") == 1
|
||||
|
||||
|
||||
class TestSearchOrOfTerms:
|
||||
"""Verify that multi-word search uses OR-of-terms (any term matches → row included)."""
|
||||
|
||||
def test_single_matching_term_in_multi_word_query(self, backend):
|
||||
"""Memory with content 'apple' found when query is 'apple banana cherry'."""
|
||||
backend.create_structured_memory("m1", "apple_mem", "", "project", "global", "", "apple")
|
||||
backend.create_structured_memory("m2", "other_mem", "", "project", "global", "", "grape")
|
||||
|
||||
results = backend.search_structured_memories("apple banana cherry")
|
||||
names = {r["name"] for r in results}
|
||||
assert "apple_mem" in names # matches "apple" — OR-of-terms keeps it
|
||||
assert "other_mem" not in names # "grape" matches nothing in the query
|
||||
|
||||
def test_partial_overlap_across_memories(self, backend):
|
||||
"""Each memory matches one of three terms; all three are returned."""
|
||||
backend.create_structured_memory("m1", "alpha_doc", "", "project", "global", "", "alpha")
|
||||
backend.create_structured_memory("m2", "beta_doc", "", "project", "global", "", "beta")
|
||||
backend.create_structured_memory("m3", "gamma_doc", "", "project", "global", "", "gamma")
|
||||
backend.create_structured_memory("m4", "unrelated", "", "project", "global", "", "delta")
|
||||
|
||||
results = backend.search_structured_memories("alpha beta gamma")
|
||||
names = {r["name"] for r in results}
|
||||
assert "alpha_doc" in names
|
||||
assert "beta_doc" in names
|
||||
assert "gamma_doc" in names
|
||||
assert "unrelated" not in names # "delta" doesn't appear in the query
|
||||
|
||||
def test_scope_filter_preserved(self, backend):
|
||||
"""OR-of-terms search still respects scope / scope_id filters."""
|
||||
backend.create_structured_memory(
|
||||
"m1", "ws1_note", "", "project", "workstream", "ws1", "info"
|
||||
)
|
||||
backend.create_structured_memory(
|
||||
"m2", "ws2_note", "", "project", "workstream", "ws2", "info"
|
||||
)
|
||||
backend.create_structured_memory("m3", "global_note", "", "project", "global", "", "info")
|
||||
|
||||
results = backend.search_structured_memories("info", scope="workstream", scope_id="ws1")
|
||||
names = {r["name"] for r in results}
|
||||
assert "ws1_note" in names
|
||||
assert "ws2_note" not in names
|
||||
assert "global_note" not in names
|
||||
|
||||
def test_term_cap_normalizes_unbounded_query(self, backend):
|
||||
"""A multi-KB query collapses to <= MAX terms (de-dupe + length filter)."""
|
||||
backend.create_structured_memory("m1", "alpha_doc", "", "project", "global", "", "alpha")
|
||||
backend.create_structured_memory(
|
||||
"m2", "other_doc", "", "project", "global", "", "irrelevant"
|
||||
)
|
||||
|
||||
# Build a noisy query: same word repeated, plus 1-char tokens that
|
||||
# the normalizer drops, plus the actual signal "alpha".
|
||||
noisy = " ".join(["x"] * 100 + ["alpha"] * 50)
|
||||
results = backend.search_structured_memories(noisy)
|
||||
names = {r["name"] for r in results}
|
||||
assert "alpha_doc" in names
|
||||
|
||||
|
||||
class TestVisibleStructuredMemories:
|
||||
"""Single-query union helpers used by the composition path."""
|
||||
|
||||
def test_list_visible_unions_global_workstream_user(self, backend):
|
||||
backend.create_structured_memory("m1", "g_note", "", "project", "global", "", "g")
|
||||
backend.create_structured_memory("m2", "ws_note", "", "project", "workstream", "ws1", "w")
|
||||
backend.create_structured_memory("m3", "u_note", "", "project", "user", "u1", "u")
|
||||
backend.create_structured_memory("m4", "other_ws", "", "project", "workstream", "ws2", "x")
|
||||
|
||||
scopes = [("global", ""), ("workstream", "ws1"), ("user", "u1")]
|
||||
rows = backend.list_visible_structured_memories(scopes)
|
||||
names = {r["name"] for r in rows}
|
||||
assert names == {"g_note", "ws_note", "u_note"} # ws2 excluded
|
||||
|
||||
def test_search_visible_unions_scopes_and_terms(self, backend):
|
||||
backend.create_structured_memory("m1", "g_alpha", "", "project", "global", "", "alpha")
|
||||
backend.create_structured_memory(
|
||||
"m2", "ws_beta", "", "project", "workstream", "ws1", "beta"
|
||||
)
|
||||
backend.create_structured_memory(
|
||||
"m3", "ws_other", "", "project", "workstream", "ws2", "alpha"
|
||||
)
|
||||
|
||||
scopes = [("global", ""), ("workstream", "ws1")]
|
||||
rows = backend.search_visible_structured_memories("alpha beta", scopes)
|
||||
names = {r["name"] for r in rows}
|
||||
assert "g_alpha" in names # global, matches "alpha"
|
||||
assert "ws_beta" in names # ws1, matches "beta"
|
||||
assert "ws_other" not in names # ws2 -> outside visibility
|
||||
|
||||
def test_visible_helpers_handle_empty_scopes(self, backend):
|
||||
backend.create_structured_memory("m1", "anything", "", "project", "global", "", "x")
|
||||
assert backend.list_visible_structured_memories([]) == []
|
||||
assert backend.search_visible_structured_memories("x", []) == []
|
||||
|
||||
|
||||
class TestStableOrderingOnTimestampTies:
|
||||
"""When two memories share an `updated` timestamp, secondary sort on
|
||||
memory_id keeps the order deterministic across calls.
|
||||
|
||||
`updated` is second-precision, and touch_structured_memories() can bump
|
||||
a batch to identical timestamps — without a tie-breaker BM25 input
|
||||
order shuffles run-to-run, busting the LLM-side prompt cache.
|
||||
"""
|
||||
|
||||
def _seed_with_shared_timestamp(self, backend):
|
||||
# Create three memories then force their `updated` columns equal —
|
||||
# mirrors the real-world case where a touch_structured_memories
|
||||
# batch lands them in the same second.
|
||||
for mid in ("zebra_id", "apple_id", "mango_id"):
|
||||
backend.create_structured_memory(
|
||||
mid, f"name_{mid}", "", "project", "global", "", "shared content"
|
||||
)
|
||||
import sqlalchemy as sa
|
||||
|
||||
with backend._conn() as conn:
|
||||
conn.execute(sa.text("UPDATE structured_memories SET updated = '2024-01-01T00:00:00'"))
|
||||
conn.commit()
|
||||
|
||||
def test_list_stable_order_under_tied_updated(self, backend):
|
||||
self._seed_with_shared_timestamp(backend)
|
||||
first = [r["memory_id"] for r in backend.list_structured_memories()]
|
||||
second = [r["memory_id"] for r in backend.list_structured_memories()]
|
||||
# Deterministic across calls AND sorted by memory_id ASC for ties
|
||||
assert first == second
|
||||
assert first == ["apple_id", "mango_id", "zebra_id"]
|
||||
|
||||
def test_search_stable_order_under_tied_updated(self, backend):
|
||||
self._seed_with_shared_timestamp(backend)
|
||||
first = [r["memory_id"] for r in backend.search_structured_memories("shared")]
|
||||
second = [r["memory_id"] for r in backend.search_structured_memories("shared")]
|
||||
assert first == second
|
||||
assert first == ["apple_id", "mango_id", "zebra_id"]
|
||||
|
||||
def test_visible_search_stable_order_under_tied_updated(self, backend):
|
||||
self._seed_with_shared_timestamp(backend)
|
||||
scopes = [("global", "")]
|
||||
first = [
|
||||
r["memory_id"] for r in backend.search_visible_structured_memories("shared", scopes)
|
||||
]
|
||||
second = [
|
||||
r["memory_id"] for r in backend.search_visible_structured_memories("shared", scopes)
|
||||
]
|
||||
assert first == second
|
||||
assert first == ["apple_id", "mango_id", "zebra_id"]
|
||||
|
||||
@@ -7,6 +7,7 @@ from turnstone.core.tool_advisory import (
|
||||
GuardAdvisory,
|
||||
MetacognitiveAdvisory,
|
||||
UserInterjection,
|
||||
escape_wrapper_tags,
|
||||
parse_priority,
|
||||
render_system_reminder,
|
||||
wrap_tool_result,
|
||||
@@ -224,6 +225,63 @@ class TestMetacognitiveAdvisory:
|
||||
assert "don't repeat tool calls" in result
|
||||
|
||||
|
||||
class TestEscapeWrapperTags:
|
||||
"""``escape_wrapper_tags`` must round-trip through
|
||||
``_entity_decode_wrapper_tags`` for any input — not just text that
|
||||
happens to contain only wrapper tags.
|
||||
"""
|
||||
|
||||
def test_short_circuit_passes_through_plain_text(self) -> None:
|
||||
"""No ``<`` and no ``&`` — ``escape_wrapper_tags`` must avoid
|
||||
the four ``replace`` chains. Common case for most tool outputs;
|
||||
the short-circuit keeps wrap_tool_result's overhead near zero."""
|
||||
text = "plain text without any markup"
|
||||
assert escape_wrapper_tags(text) == text
|
||||
|
||||
def test_escape_wrapper_tags_round_trips_preexisting_entities(self) -> None:
|
||||
"""Asymmetry guard — a tool output that happens to contain the
|
||||
literal string ``<tool_output>`` (e.g. documentation
|
||||
describing the wrapper format) must round-trip identically.
|
||||
Without escaping ``&`` first, encode→decode would produce the
|
||||
bare ``<tool_output>`` tag, fabricating an envelope the wrapper
|
||||
layer never produced."""
|
||||
from turnstone.core.history_decoration import _entity_decode_wrapper_tags
|
||||
|
||||
text = "I describe XML tags like <tool_output> in my docs."
|
||||
encoded = escape_wrapper_tags(text)
|
||||
# Sanity: the original literal got escaped to a sentinel form
|
||||
# that can't collide with our wrapper-tag escapes.
|
||||
assert "&lt;tool_output&gt;" in encoded
|
||||
assert "<tool_output>" not in encoded
|
||||
# Round-trip back to the literal source.
|
||||
assert _entity_decode_wrapper_tags(encoded) == text
|
||||
|
||||
def test_escape_wrapper_tags_round_trips_real_wrapper_tag(self) -> None:
|
||||
"""A literal ``<tool_output>`` in source text round-trips back
|
||||
correctly — encoding produces ``<tool_output>`` (no
|
||||
``&`` prefix because there was no pre-existing entity), and
|
||||
decoding restores the literal."""
|
||||
from turnstone.core.history_decoration import _entity_decode_wrapper_tags
|
||||
|
||||
text = "Here is a literal <tool_output> tag in my doc."
|
||||
encoded = escape_wrapper_tags(text)
|
||||
assert "<tool_output>" not in encoded
|
||||
assert "<tool_output>" in encoded
|
||||
assert _entity_decode_wrapper_tags(encoded) == text
|
||||
|
||||
def test_escape_wrapper_tags_round_trips_mixed_content(self) -> None:
|
||||
"""Mixed: literal wrapper tags AND pre-existing entity
|
||||
references — both round-trip."""
|
||||
from turnstone.core.history_decoration import _entity_decode_wrapper_tags
|
||||
|
||||
text = (
|
||||
"Mixed: literal <tool_output> next to escaped <system-reminder> "
|
||||
"and a stray & on its own."
|
||||
)
|
||||
encoded = escape_wrapper_tags(text)
|
||||
assert _entity_decode_wrapper_tags(encoded) == text
|
||||
|
||||
|
||||
class TestRenderSystemReminder:
|
||||
"""render_system_reminder builds a standalone <system-reminder> envelope."""
|
||||
|
||||
|
||||
+82
-5
@@ -9,6 +9,7 @@ import pytest
|
||||
|
||||
from turnstone.core.watch import (
|
||||
WatchRunner,
|
||||
build_watch_reminder,
|
||||
evaluate_condition,
|
||||
format_interval,
|
||||
format_watch_message,
|
||||
@@ -301,6 +302,62 @@ class TestFormatWatchMessage:
|
||||
assert "max polls" in msg.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_watch_reminder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuildWatchReminder:
|
||||
"""The structured-reminder builder lifts ``format_watch_message``'s
|
||||
args into a dict the dispatch closure can pass to
|
||||
``WatchRunner._dispatch_result``. ``text`` matches the formatter's
|
||||
output verbatim (so compaction / channel adapters / wire splice
|
||||
keep their behaviour), and the optional fields ride alongside for
|
||||
the frontend's ``.msg.watch-result`` card.
|
||||
"""
|
||||
|
||||
def test_emits_text_body_and_fields(self):
|
||||
kwargs = dict(
|
||||
name="pr-review",
|
||||
command="gh pr view --json state",
|
||||
output='{"state": "MERGED"}',
|
||||
poll_count=5,
|
||||
max_polls=100,
|
||||
elapsed_secs=1500,
|
||||
stop_on='data["state"] == "MERGED"',
|
||||
is_final=True,
|
||||
reason='condition met: data["state"] == "MERGED"',
|
||||
)
|
||||
reminder = build_watch_reminder(**kwargs)
|
||||
# Round-trip with format_watch_message — text is the same body
|
||||
# the wire splice + channel adapters have always seen.
|
||||
assert reminder["text"] == format_watch_message(**kwargs)
|
||||
# Optional fields ride alongside.
|
||||
assert reminder["type"] == "watch_triggered"
|
||||
assert reminder["watch_name"] == "pr-review"
|
||||
assert reminder["command"] == "gh pr view --json state"
|
||||
assert reminder["poll_count"] == 5
|
||||
assert reminder["max_polls"] == 100
|
||||
assert reminder["is_final"] is True
|
||||
|
||||
def test_non_final_carries_is_final_false(self):
|
||||
reminder = build_watch_reminder(
|
||||
name="deploy",
|
||||
command="curl -s http://localhost/health",
|
||||
output="ok",
|
||||
poll_count=3,
|
||||
max_polls=50,
|
||||
elapsed_secs=90,
|
||||
stop_on=None,
|
||||
is_final=False,
|
||||
reason="",
|
||||
)
|
||||
assert reminder["is_final"] is False
|
||||
assert reminder["poll_count"] == 3
|
||||
# No "auto-cancelled" body for non-final fires.
|
||||
assert "auto-cancelled" not in reminder["text"].lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# WatchRunner
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -450,13 +507,17 @@ class TestWatchRunner:
|
||||
runner.set_dispatch_fn("ws-1", fn1)
|
||||
runner.set_dispatch_fn("ws-2", fn2)
|
||||
|
||||
runner._dispatch_result("ws-1", "msg1")
|
||||
fn1.assert_called_once_with("msg1")
|
||||
# ``_dispatch_result`` takes a structured reminder dict, not a
|
||||
# bare string.
|
||||
reminder1 = {"type": "watch_triggered", "text": "msg1"}
|
||||
runner._dispatch_result("ws-1", reminder1, "watch-a")
|
||||
fn1.assert_called_once_with(reminder1, "watch-a")
|
||||
fn2.assert_not_called()
|
||||
|
||||
runner.remove_dispatch_fn("ws-1")
|
||||
# After removal, dispatch should try restore_fn
|
||||
runner._dispatch_result("ws-1", "msg2")
|
||||
reminder2 = {"type": "watch_triggered", "text": "msg2"}
|
||||
runner._dispatch_result("ws-1", reminder2, "watch-b")
|
||||
fn1.assert_called_once() # still just the one call
|
||||
|
||||
def test_restore_fn_called_for_evicted(self):
|
||||
@@ -464,9 +525,25 @@ class TestWatchRunner:
|
||||
restore_fn = MagicMock(return_value=restored_fn)
|
||||
runner = self._make_runner(restore_fn=restore_fn)
|
||||
|
||||
runner._dispatch_result("ws-evicted", "hello")
|
||||
reminder = {"type": "watch_triggered", "text": "hello"}
|
||||
runner._dispatch_result("ws-evicted", reminder, "watch-x")
|
||||
restore_fn.assert_called_once_with("ws-evicted")
|
||||
restored_fn.assert_called_once_with("hello")
|
||||
restored_fn.assert_called_once_with(reminder, "watch-x")
|
||||
|
||||
def test_get_dispatch_fn_returns_registered_fn(self):
|
||||
"""``get_dispatch_fn`` is the public accessor used by the
|
||||
server-side restore path to retrieve the per-ws closure that
|
||||
``set_watch_runner`` constructed during workstream rehydrate.
|
||||
"""
|
||||
runner = self._make_runner()
|
||||
fn = MagicMock()
|
||||
runner.set_dispatch_fn("ws-1", fn)
|
||||
assert runner.get_dispatch_fn("ws-1") is fn
|
||||
# Unknown ws → None.
|
||||
assert runner.get_dispatch_fn("ws-missing") is None
|
||||
# After removal → None.
|
||||
runner.remove_dispatch_fn("ws-1")
|
||||
assert runner.get_dispatch_fn("ws-1") is None
|
||||
|
||||
def test_run_command_success(self):
|
||||
runner = self._make_runner()
|
||||
|
||||
+402
-199
@@ -1,206 +1,409 @@
|
||||
"""Tests for _make_watch_dispatch error/cancel handling and concurrency guards."""
|
||||
"""Tests for the watch dispatch closure built inside ``set_watch_runner``.
|
||||
|
||||
The closure routes watch results onto the per-session :class:`NudgeQueue`
|
||||
under the unified pull-model surface. Each test focuses on one
|
||||
assertion: enqueue shape, sanitisation, soft-cap drop-oldest,
|
||||
``valid_until`` predicate, and concurrent-enqueue safety.
|
||||
|
||||
Tests in this file replace the pre-switchover suite that pinned the
|
||||
``_make_watch_dispatch`` worker-spawn / ``_watch_pending`` machinery —
|
||||
the contracts those tests pinned no longer exist. See
|
||||
``tests/test_watch.py`` for the still-relevant ``WatchRunner``
|
||||
mechanics tests, and ``tests/test_watch_integration.py`` for the
|
||||
boundary-crossing integration test covering the chat-loop drain.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from turnstone.core.session import GenerationCancelled
|
||||
from turnstone.core.workstream import Workstream
|
||||
from turnstone.server import _make_watch_dispatch
|
||||
import pytest
|
||||
|
||||
from tests._helpers import patch_session_storage
|
||||
from turnstone.core.session import _WATCH_QUEUE_SOFT_CAP, ChatSession
|
||||
|
||||
class _StubSession:
|
||||
"""Minimal ChatSession stand-in with controllable send() behaviour."""
|
||||
|
||||
def __init__(self, *, side_effect=None):
|
||||
self._watch_pending: queue.Queue = queue.Queue(maxsize=20)
|
||||
self._side_effect = side_effect
|
||||
class _NullUI:
|
||||
"""UI adapter that discards all output — local to this test module
|
||||
to avoid a cross-test-file import (mirrors the pattern in
|
||||
test_session.py / test_rewind_retry.py).
|
||||
"""
|
||||
|
||||
def send(self, msg: str) -> None:
|
||||
if self._side_effect is not None:
|
||||
raise self._side_effect
|
||||
|
||||
|
||||
class _RecordingUI:
|
||||
"""Track calls made by the dispatch error handlers."""
|
||||
|
||||
def __init__(self):
|
||||
self.errors: list[str] = []
|
||||
self.state_changes: list[str] = []
|
||||
self.stream_end_calls: int = 0
|
||||
|
||||
# -- SessionUI protocol stubs used by the dispatch code --
|
||||
|
||||
def on_error(self, message: str) -> None:
|
||||
self.errors.append(message)
|
||||
|
||||
def on_state_change(self, state: str) -> None:
|
||||
self.state_changes.append(state)
|
||||
|
||||
def on_stream_end(self) -> None:
|
||||
self.stream_end_calls += 1
|
||||
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _wait_for_worker(ws: Workstream, timeout: float = 2.0) -> None:
|
||||
"""Block until the worker thread started by dispatch() finishes."""
|
||||
t = ws.worker_thread
|
||||
if t is not None:
|
||||
t.join(timeout)
|
||||
assert not t.is_alive(), "worker thread did not finish in time"
|
||||
|
||||
|
||||
# ── GenerationCancelled path ────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_cancelled_emits_stream_end_and_idle():
|
||||
session = _StubSession(side_effect=GenerationCancelled())
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("hello")
|
||||
_wait_for_worker(ws)
|
||||
|
||||
assert ui.stream_end_calls == 1
|
||||
assert ui.state_changes == ["idle"]
|
||||
assert ui.errors == []
|
||||
|
||||
|
||||
# ── Generic exception path ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_exception_emits_stream_end_and_error():
|
||||
session = _StubSession(side_effect=RuntimeError("boom"))
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("hello")
|
||||
_wait_for_worker(ws)
|
||||
|
||||
assert ui.stream_end_calls == 1
|
||||
assert ui.state_changes == ["error"]
|
||||
assert len(ui.errors) == 1
|
||||
assert "boom" in ui.errors[0]
|
||||
|
||||
|
||||
# ── Worker-thread identity guard ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_abandoned_thread_emits_no_events():
|
||||
"""After force-cancel sets worker_thread=None, the old thread must not
|
||||
emit stream_end or state changes."""
|
||||
barrier = threading.Event()
|
||||
|
||||
class _BlockingSession(_StubSession):
|
||||
def send(self, msg: str) -> None:
|
||||
barrier.wait(timeout=5)
|
||||
raise RuntimeError("late error")
|
||||
|
||||
session = _BlockingSession()
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("hello")
|
||||
|
||||
# Simulate force-cancel: clear the worker_thread reference.
|
||||
ws.worker_thread = None
|
||||
barrier.set()
|
||||
|
||||
# Wait for the thread to actually complete (it's still running).
|
||||
time.sleep(0.3)
|
||||
|
||||
assert ui.stream_end_calls == 0
|
||||
assert ui.state_changes == []
|
||||
assert ui.errors == []
|
||||
|
||||
|
||||
# ── Path A: busy workstream enqueue ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_busy_workstream_enqueues_message():
|
||||
"""When the workstream already has a live worker, dispatch enqueues."""
|
||||
session = _StubSession()
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
# Simulate a live worker — session_worker.send gates on
|
||||
# ``_worker_running``, not ``Thread.is_alive``.
|
||||
ws._worker_running = True
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("queued msg")
|
||||
|
||||
item = session._watch_pending.get_nowait()
|
||||
assert item == {"message": "queued msg"}
|
||||
|
||||
|
||||
def test_busy_workstream_drops_on_full_queue():
|
||||
"""When the pending queue is full, dispatch drops the message."""
|
||||
session = _StubSession()
|
||||
# Fill the queue to capacity.
|
||||
for i in range(20):
|
||||
session._watch_pending.put_nowait({"message": f"msg{i}"})
|
||||
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
ws._worker_running = True # simulate a live worker
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
# Should not block or raise — just log a warning and drop.
|
||||
dispatch("overflow msg")
|
||||
|
||||
assert session._watch_pending.full()
|
||||
|
||||
|
||||
# ── Lock guard ───────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_dispatch_holds_lock_during_thread_start():
|
||||
"""Dispatch acquires ws._lock before checking/starting the worker."""
|
||||
session = _StubSession()
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
acquire_count = 0
|
||||
inner = ws._lock
|
||||
|
||||
class _CountingLock:
|
||||
def __enter__(self):
|
||||
nonlocal acquire_count
|
||||
acquire_count += 1
|
||||
return inner.__enter__()
|
||||
|
||||
def __exit__(self, *args):
|
||||
return inner.__exit__(*args)
|
||||
|
||||
ws._lock = _CountingLock() # type: ignore[assignment]
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("hello")
|
||||
_wait_for_worker(ws)
|
||||
|
||||
assert acquire_count >= 1
|
||||
|
||||
|
||||
# ── Happy path ───────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_successful_send_no_error_events():
|
||||
"""Normal send() completion should not trigger error/cancel events."""
|
||||
session = _StubSession() # send() does nothing (success)
|
||||
ws = Workstream()
|
||||
ui = _RecordingUI()
|
||||
|
||||
dispatch = _make_watch_dispatch(ws, session, ui)
|
||||
dispatch("hello")
|
||||
_wait_for_worker(ws)
|
||||
|
||||
assert ui.stream_end_calls == 0
|
||||
assert ui.state_changes == []
|
||||
assert ui.errors == []
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Catch-all: any UI hook the chat loop calls becomes a no-op.
|
||||
return MagicMock()
|
||||
|
||||
|
||||
def _make_session_for_dispatch(**kwargs: Any) -> ChatSession:
|
||||
"""ChatSession built with the same minimal harness used elsewhere
|
||||
in the test suite, scoped down to what the dispatch closure needs.
|
||||
"""
|
||||
client = MagicMock()
|
||||
defaults = dict(
|
||||
client=client,
|
||||
model="test-model",
|
||||
ui=_NullUI(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return ChatSession(**defaults)
|
||||
|
||||
|
||||
def _register_runner(session: ChatSession) -> tuple[Any, Any]:
|
||||
"""Attach a minimal stub ``WatchRunner`` to *session* and return the
|
||||
``(runner, dispatch_fn)`` pair captured by ``set_dispatch_fn``.
|
||||
"""
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
class _StubRunner:
|
||||
def set_dispatch_fn(self, ws_id: str, fn: Any) -> None:
|
||||
captured["fn"] = fn
|
||||
|
||||
runner = _StubRunner()
|
||||
session.set_watch_runner(runner)
|
||||
return runner, captured["fn"]
|
||||
|
||||
|
||||
def _reminder(text: str, **extra: Any) -> dict[str, Any]:
|
||||
"""Build a structured ``watch_triggered`` reminder dict for tests.
|
||||
|
||||
Mirrors the shape produced by :func:`turnstone.core.watch.build_watch_reminder`
|
||||
— ``text`` is the formatted body, optional fields ride alongside.
|
||||
Tests that don't care about the optional fields can call with
|
||||
``text`` only.
|
||||
"""
|
||||
out: dict[str, Any] = {"type": "watch_triggered", "text": text}
|
||||
out.update(extra)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Enqueue shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnqueueShape:
|
||||
"""``set_watch_runner``'s closure produces a single
|
||||
``("watch_triggered", text, "any")`` entry per fire.
|
||||
"""
|
||||
|
||||
def test_dispatch_enqueues_watch_triggered_with_any_channel(self, tmp_db):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
dispatch(_reminder("watch fired body"), "watch-1")
|
||||
|
||||
# One entry, "watch_triggered" type, on "any" channel.
|
||||
assert len(session._nudge_queue) == 1
|
||||
assert session._nudge_queue.pending(channel="any") == [
|
||||
("watch_triggered", "watch fired body")
|
||||
]
|
||||
# NOT on "user" or "tool" channels.
|
||||
assert session._nudge_queue.pending(channel="user") == []
|
||||
assert session._nudge_queue.pending(channel="tool") == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sanitisation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSanitisation:
|
||||
"""``sanitize_payload`` runs producer-side over the formatted message
|
||||
before it ever reaches the queue. The wire-boundary
|
||||
``escape_wrapper_tags`` only protects ``<system-reminder>`` /
|
||||
``<tool_output>`` envelopes; this layer covers everything else.
|
||||
"""
|
||||
|
||||
def test_dispatch_sanitizes_payload_before_enqueue(self, tmp_db):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# Build a payload with: BEL (\x07), zero-width space (U+200B),
|
||||
# bidi RTL override (U+202E), and angle-bracket tag breakers.
|
||||
raw = "before\x07middleaftermore<thinking>tail"
|
||||
dispatch(_reminder(raw), "watch-1")
|
||||
|
||||
pending = session._nudge_queue.pending(channel="any")
|
||||
assert len(pending) == 1
|
||||
sanitized = pending[0][1]
|
||||
# Control / steering chars become spaces; angle brackets vanish.
|
||||
assert "\x07" not in sanitized
|
||||
assert "" not in sanitized
|
||||
assert "" not in sanitized
|
||||
assert "<" not in sanitized
|
||||
assert ">" not in sanitized
|
||||
# Real content survives.
|
||||
assert "before" in sanitized
|
||||
assert "thinking" in sanitized
|
||||
|
||||
def test_dispatch_preserves_newlines_for_multiline_output(self, tmp_db):
|
||||
"""Multi-line shell output must keep its layout — TAB / LF / CR
|
||||
are intentionally preserved by ``sanitize_payload`` (R8).
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
dispatch(_reminder("line1\nline2\n\tindented\n\rline3"), "watch-1")
|
||||
|
||||
pending = session._nudge_queue.pending(channel="any")
|
||||
assert len(pending) == 1
|
||||
text = pending[0][1]
|
||||
# Lines stay separated; tab kept.
|
||||
assert "\n" in text
|
||||
assert "\t" in text
|
||||
|
||||
def test_dispatch_drops_empty_after_sanitization(self, tmp_db):
|
||||
"""A payload that's all control chars sanitises to "" — no enqueue."""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# All-control + DEL + zero-width — strips to empty.
|
||||
dispatch(_reminder("\x07\x0b\x7f"), "watch-1")
|
||||
|
||||
assert len(session._nudge_queue) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Soft cap
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSoftCap:
|
||||
"""When ``"watch_triggered"`` saturates at :data:`_WATCH_QUEUE_SOFT_CAP`,
|
||||
the closure drops the OLDEST entry of that type and enqueues the new
|
||||
one — so the queue stays ≤ cap with the most recent watch outputs.
|
||||
"""
|
||||
|
||||
def test_dispatch_drop_oldest_at_soft_cap(self, tmp_db, caplog):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# Pre-fill at the cap. Each entry has a unique body so we can
|
||||
# tell which one(s) survived a drop.
|
||||
for i in range(_WATCH_QUEUE_SOFT_CAP):
|
||||
dispatch(_reminder(f"body-{i}"), "watch-1")
|
||||
assert len(session._nudge_queue) == _WATCH_QUEUE_SOFT_CAP
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
dispatch(_reminder("overflow"), "watch-1")
|
||||
|
||||
# Total stays at cap (one dropped, one added).
|
||||
assert len(session._nudge_queue) == _WATCH_QUEUE_SOFT_CAP
|
||||
bodies = [text for _t, text in session._nudge_queue.pending(channel="any")]
|
||||
# Oldest ("body-0") gone; newest ("overflow") present.
|
||||
assert "body-0" not in bodies
|
||||
assert "overflow" in bodies
|
||||
# Warning logged.
|
||||
assert any("watch_dispatch.queue_full" in r.message for r in caplog.records), (
|
||||
"expected a watch_dispatch.queue_full warning record"
|
||||
)
|
||||
|
||||
def test_dispatch_soft_cap_does_not_evict_other_types(self, tmp_db):
|
||||
"""A watch saturation drop must only target watch-typed entries.
|
||||
Other producers (idle_children, advisories) have their own
|
||||
rate limiters and must not be collateral damage.
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# Mix in a few non-watch entries on the same queue.
|
||||
session._nudge_queue.enqueue("idle_children", "ic-1", "any")
|
||||
session._nudge_queue.enqueue("idle_children", "ic-2", "any")
|
||||
|
||||
# Saturate watches up to cap (queue holds cap+2 total).
|
||||
for i in range(_WATCH_QUEUE_SOFT_CAP):
|
||||
dispatch(_reminder(f"body-{i}"), "watch-1")
|
||||
# One more triggers drop-oldest of a "watch_triggered" entry.
|
||||
dispatch(_reminder("overflow"), "watch-1")
|
||||
|
||||
# Both idle_children entries survived — no collateral eviction.
|
||||
idle_bodies = [
|
||||
text
|
||||
for nt, text in session._nudge_queue.pending(channel="any")
|
||||
if nt == "idle_children"
|
||||
]
|
||||
assert idle_bodies == ["ic-1", "ic-2"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# valid_until predicate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestValidUntil:
|
||||
"""The ``valid_until`` predicate captured at dispatch time re-checks
|
||||
the watch's ``active`` flag at drain time, so a cancelled watch's
|
||||
last splat doesn't ride out a future wake.
|
||||
"""
|
||||
|
||||
def test_valid_until_drops_when_watch_inactive(self, tmp_db, monkeypatch):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# Storage stub returns False at drain time.
|
||||
is_active_calls = patch_session_storage(monkeypatch, active=False)
|
||||
|
||||
dispatch(_reminder("body"), "watch-1")
|
||||
# Drain fires the predicate; entry should NOT be delivered.
|
||||
out = session._nudge_queue.drain({"any"})
|
||||
assert out == []
|
||||
# Predicate ran once with the dispatched watch_id.
|
||||
assert is_active_calls == ["watch-1"]
|
||||
|
||||
def test_valid_until_drops_when_storage_raises(self, tmp_db, monkeypatch):
|
||||
"""The closure's broad-except in the predicate translates a
|
||||
storage-layer exception to ``False`` so the drain doesn't
|
||||
propagate; the predicate captured ``watch_id`` correctly
|
||||
(otherwise storage wouldn't even be touched).
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
patch_session_storage(monkeypatch, raise_on_is_active=True)
|
||||
|
||||
dispatch(_reminder("body"), "watch-bound-id")
|
||||
out = session._nudge_queue.drain({"any"})
|
||||
assert out == []
|
||||
|
||||
def test_valid_until_delivers_when_watch_active(self, tmp_db, monkeypatch):
|
||||
"""Happy-path counter-test for the predicate above: the entry
|
||||
DOES drain when the watch is still active.
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
patch_session_storage(monkeypatch, active=True)
|
||||
|
||||
dispatch(_reminder("body"), "watch-1")
|
||||
out = session._nudge_queue.drain({"any"})
|
||||
assert len(out) == 1
|
||||
assert out[0][0] == "watch_triggered"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Concurrency
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConcurrency:
|
||||
"""Two threads each fire 100 dispatches against the same session;
|
||||
the soft-cap read-then-mutate window stays bounded and the queue
|
||||
settles in a consistent state.
|
||||
|
||||
Per the plan's risk register R2: in production only one daemon
|
||||
thread (``WatchRunner``'s ``_run``) ever calls a session's dispatch
|
||||
fn, so the 3-acquisition non-atomicity is harmless. This test
|
||||
pins lock-correctness anyway against the broader race window.
|
||||
"""
|
||||
|
||||
def test_dispatch_concurrent_enqueues_thread_safe(self, tmp_db, monkeypatch):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
# Bypass the storage-touching valid_until predicate: count cap
|
||||
# behaviour, not storage round-trips.
|
||||
patch_session_storage(monkeypatch, active=True)
|
||||
|
||||
per_thread = 100
|
||||
labels = ("a", "b")
|
||||
|
||||
def fire(label: str) -> None:
|
||||
for i in range(per_thread):
|
||||
dispatch(_reminder(f"{label}-{i}"), f"watch-{label}")
|
||||
|
||||
threads = [threading.Thread(target=fire, args=(label,), daemon=True) for label in labels]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join(timeout=5.0)
|
||||
for t in threads:
|
||||
assert not t.is_alive(), "dispatch thread did not finish in time"
|
||||
|
||||
# The non-atomic count-then-drop window admits at most one "slip"
|
||||
# per concurrent thread above the cap (each thread can observe a
|
||||
# sub-cap count and append before another thread's drop runs).
|
||||
depth = len(session._nudge_queue)
|
||||
assert depth <= len(threads) * per_thread
|
||||
assert depth <= _WATCH_QUEUE_SOFT_CAP + len(threads)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Empty-input / multi-call invariants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", ["", " ", "\x07\x0b"])
|
||||
def test_dispatch_no_op_for_empty_payloads(tmp_db, payload: str):
|
||||
"""Whitespace-only / pure-control payloads sanitise to empty and
|
||||
do not produce a queue entry — silent drop.
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
dispatch(_reminder(payload), "watch-1")
|
||||
|
||||
assert len(session._nudge_queue) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Metadata propagation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMetadataPropagation:
|
||||
"""The dispatch closure pulls optional fields out of the structured
|
||||
``reminder`` dict and attaches them to the queue entry's
|
||||
``metadata``. Drain seams later merge ``metadata`` into the
|
||||
rendered reminder dict so the frontend can display a structured
|
||||
``.msg.watch-result`` card.
|
||||
"""
|
||||
|
||||
def test_dispatch_attaches_watch_metadata_on_enqueue(self, tmp_db):
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
reminder = _reminder(
|
||||
"$ ls\nfile.txt",
|
||||
watch_name="my-watch",
|
||||
command="ls",
|
||||
poll_count=2,
|
||||
max_polls=100,
|
||||
is_final=False,
|
||||
)
|
||||
dispatch(reminder, "watch-1")
|
||||
|
||||
# Snapshot via ``pending_with_metadata`` to inspect the full
|
||||
# entry shape. Exactly one entry, with the optional fields
|
||||
# carried verbatim onto ``metadata``.
|
||||
snapshot = session._nudge_queue.pending_with_metadata(channel="any")
|
||||
assert len(snapshot) == 1
|
||||
nt, _text, meta = snapshot[0]
|
||||
assert nt == "watch_triggered"
|
||||
assert meta == {
|
||||
"watch_name": "my-watch",
|
||||
"command": "ls",
|
||||
"poll_count": 2,
|
||||
"max_polls": 100,
|
||||
"is_final": False,
|
||||
}
|
||||
|
||||
def test_dispatch_omits_metadata_when_optional_fields_missing(self, tmp_db):
|
||||
"""A bare ``{type, text}`` reminder produces an entry with no
|
||||
metadata — the closure builds an empty dict, sees nothing to
|
||||
carry, and falls through to ``metadata=None``.
|
||||
"""
|
||||
session = _make_session_for_dispatch()
|
||||
_runner, dispatch = _register_runner(session)
|
||||
|
||||
dispatch(_reminder("just a body"), "watch-1")
|
||||
|
||||
snapshot = session._nudge_queue.pending_with_metadata(channel="any")
|
||||
assert len(snapshot) == 1
|
||||
_nt, _text, meta = snapshot[0]
|
||||
assert meta is None
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
"""Boundary-crossing integration test for the watch switchover pipeline.
|
||||
|
||||
Drives a real :class:`ChatSession` + a real :class:`WatchRunner` (with
|
||||
its daemon thread skipped — we call ``_dispatch_result`` directly to
|
||||
avoid the timer dependency) end-to-end through the chat-loop drain
|
||||
seam. The only stub is the LLM provider (patched
|
||||
``_create_stream_with_retry``); every other layer is production code:
|
||||
|
||||
* ``WatchRunner._dispatch_result`` releasing the dispatch lock before
|
||||
fan-out
|
||||
* the closure built inside ``ChatSession.set_watch_runner`` —
|
||||
``sanitize_payload`` + soft-cap check + ``valid_until`` predicate +
|
||||
``NudgeQueue.enqueue("watch_triggered", ..., "any", ...)``
|
||||
* ``ChatSession.send`` chat loop short-circuiting metacog detection
|
||||
* ``_attach_pending_user_reminders`` draining ``USER_DRAIN`` (which
|
||||
matches ``"any"``)
|
||||
* ``_apply_reminders_for_provider`` splicing the rendered envelope onto
|
||||
the user message before the wire boundary
|
||||
|
||||
Per ``feedback_tests_through_boundaries.md``: direct injection tests
|
||||
that bypass these boundaries silently mask wiring bugs. This test is
|
||||
the structural integration gate for the watch switchover.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from tests._helpers import patch_session_storage
|
||||
from turnstone.core.session import ChatSession
|
||||
from turnstone.core.watch import WatchRunner
|
||||
|
||||
|
||||
class _NullUI:
|
||||
"""UI adapter that no-ops every chat-loop hook the test triggers."""
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return MagicMock()
|
||||
|
||||
|
||||
def _make_session() -> 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).
|
||||
"""
|
||||
return ChatSession(
|
||||
client=MagicMock(),
|
||||
model="test-model",
|
||||
ui=_NullUI(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
)
|
||||
|
||||
|
||||
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 ``<system-reminder>`` envelope path:
|
||||
|
||||
1. WatchRunner.dispatch fires watch text against the session's
|
||||
registered closure (synchronously — no daemon thread).
|
||||
2. NudgeQueue holds one ``"watch_triggered"`` entry on ``"any"``.
|
||||
3. session.send("ok") runs the chat loop with a stubbed LLM.
|
||||
4. The drain seam drains the watch entry; the wire payload's user
|
||||
message has the watch text spliced into a ``<system-reminder>``
|
||||
envelope.
|
||||
"""
|
||||
session = _make_session()
|
||||
|
||||
# Bypass the storage-touching predicate — we want to assert the
|
||||
# envelope splice, not exercise a fresh sqlite watch row.
|
||||
patch_session_storage(monkeypatch, active=True)
|
||||
|
||||
# Real WatchRunner; we don't ``start()`` the daemon thread (that
|
||||
# would race with the test's deterministic order). Direct call
|
||||
# to ``_dispatch_result`` exercises the same dispatch path the
|
||||
# daemon would invoke. Runner-side ``storage`` is unused on this
|
||||
# path (only the polling loop touches it); a MagicMock placeholder
|
||||
# keeps the constructor signature happy.
|
||||
runner = WatchRunner(storage=MagicMock(), node_id="test-node")
|
||||
session.set_watch_runner(runner)
|
||||
|
||||
# 1. Fire a watch result synchronously.
|
||||
runner._dispatch_result(
|
||||
session._ws_id,
|
||||
{"type": "watch_triggered", "text": "watch payload body"},
|
||||
"watch-1",
|
||||
)
|
||||
|
||||
# 2. The queue holds one entry on the "any" channel.
|
||||
assert len(session._nudge_queue) == 1
|
||||
pending = session._nudge_queue.pending(channel="any")
|
||||
assert pending == [("watch_triggered", "watch payload body")]
|
||||
|
||||
# 3. Run the chat loop with the LLM patched. We don't care about
|
||||
# the assistant turn's content; only the wire payload sent to the
|
||||
# provider matters.
|
||||
with (
|
||||
patch.object(session, "_create_stream_with_retry", return_value=iter([])),
|
||||
patch.object(
|
||||
session,
|
||||
"_stream_response",
|
||||
return_value={"role": "assistant", "content": "ok"},
|
||||
),
|
||||
patch.object(session, "_update_token_table"),
|
||||
patch.object(session, "_print_status_line"),
|
||||
patch.object(session, "_visible_memory_count", return_value=0),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
session._title_generated = True # suppress orthogonal title side-thread
|
||||
session.send("ok")
|
||||
|
||||
# 4. Queue fully drained by the user-message attach seam.
|
||||
assert len(session._nudge_queue) == 0
|
||||
|
||||
# The user message that drove the assistant turn has the watch
|
||||
# text in its ``_reminders`` side-channel — the production
|
||||
# ``_apply_reminders_for_provider`` splice consumes that to wrap
|
||||
# the content in ``<system-reminder>`` at the wire boundary.
|
||||
user_msgs = [m for m in session.messages if m.get("role") == "user"]
|
||||
assert user_msgs, "expected a user message in history"
|
||||
last_user = user_msgs[-1]
|
||||
reminders = last_user.get("_reminders") or []
|
||||
assert any(
|
||||
r.get("type") == "watch_triggered" and "watch payload body" in r.get("text", "")
|
||||
for r in reminders
|
||||
), f"expected watch_triggered reminder on user message; got {reminders!r}"
|
||||
|
||||
|
||||
def test_three_back_to_back_watch_fires_drain_into_one_turn(tmp_db, monkeypatch):
|
||||
"""Behavioural delta from plan section 3.4 / risk register R3.
|
||||
|
||||
N back-to-back watch fires used to produce N successive
|
||||
``send()`` turns (each a separate model invocation, capped at
|
||||
``_MAX_WATCH_CHAIN = 5``). After the switchover, the N entries
|
||||
drain into ONE envelope splice on the next drain seam — one
|
||||
assistant turn responding to all N watch results. Pinning this
|
||||
behavioural delta protects against accidental regression to
|
||||
the old per-fire-turn shape.
|
||||
"""
|
||||
session = _make_session()
|
||||
|
||||
patch_session_storage(monkeypatch, active=True)
|
||||
|
||||
runner = WatchRunner(storage=MagicMock(), node_id="test-node")
|
||||
session.set_watch_runner(runner)
|
||||
|
||||
runner._dispatch_result(
|
||||
session._ws_id, {"type": "watch_triggered", "text": "fire one"}, "watch-1"
|
||||
)
|
||||
runner._dispatch_result(
|
||||
session._ws_id, {"type": "watch_triggered", "text": "fire two"}, "watch-1"
|
||||
)
|
||||
runner._dispatch_result(
|
||||
session._ws_id, {"type": "watch_triggered", "text": "fire three"}, "watch-1"
|
||||
)
|
||||
assert len(session._nudge_queue) == 3
|
||||
|
||||
with (
|
||||
patch.object(session, "_create_stream_with_retry", return_value=iter([])),
|
||||
patch.object(
|
||||
session,
|
||||
"_stream_response",
|
||||
return_value={"role": "assistant", "content": "got it"},
|
||||
),
|
||||
patch.object(session, "_update_token_table"),
|
||||
patch.object(session, "_print_status_line"),
|
||||
patch.object(session, "_visible_memory_count", return_value=0),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
session._title_generated = True
|
||||
session.send("user")
|
||||
|
||||
# All three drained into the single user-message attach.
|
||||
user_msgs = [m for m in session.messages if m.get("role") == "user"]
|
||||
last_user = user_msgs[-1]
|
||||
reminders = last_user.get("_reminders") or []
|
||||
watch_reminders = [r for r in reminders if r.get("type") == "watch_triggered"]
|
||||
assert len(watch_reminders) == 3
|
||||
bodies = [r.get("text", "") for r in watch_reminders]
|
||||
assert any("fire one" in b for b in bodies)
|
||||
assert any("fire two" in b for b in bodies)
|
||||
assert any("fire three" in b for b in bodies)
|
||||
# And there's exactly ONE assistant turn (not three).
|
||||
assistant_turns = [m for m in session.messages if m.get("role") == "assistant"]
|
||||
assert len(assistant_turns) == 1
|
||||
|
||||
|
||||
def test_watch_dispatch_through_restore_fn_lands_on_rehydrated_session(tmp_db, monkeypatch):
|
||||
"""Cover the production ``_watch_restore_fn`` closure surface.
|
||||
|
||||
Path under test:
|
||||
WatchRunner._dispatch_result(ws_id, msg, watch_id)
|
||||
no dispatch fn registered (original session evicted)
|
||||
restore_fn(ws_id) constructs a fresh ChatSession,
|
||||
calls session.resume(ws_id) to adopt the original ws_id,
|
||||
re-registers the dispatch closure via session.set_watch_runner,
|
||||
returns runner.get_dispatch_fn(session._ws_id)
|
||||
runner invokes the returned fn with (msg, watch_id)
|
||||
watch payload lands on the rehydrated session's NudgeQueue
|
||||
|
||||
Construction inside ``server.py``'s ``_watch_restore_fn`` is the new
|
||||
contract surface introduced by the switchover; this test pins that
|
||||
contract so a future refactor of the closure (e.g. swapping
|
||||
``manager.create + session.resume`` for ``manager.open``) doesn't
|
||||
silently break the watch-restore pipeline.
|
||||
"""
|
||||
from turnstone.core import session as session_mod
|
||||
|
||||
patch_session_storage(monkeypatch, active=True)
|
||||
|
||||
# Stage 1 — build the original session and persist a message so
|
||||
# ``session.resume`` finds the ws_id in storage.
|
||||
original = _make_session()
|
||||
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).
|
||||
session_mod.save_message(original_ws_id, "user", "kickoff message")
|
||||
|
||||
# Stage 2 — runner with NO dispatch fn registered (simulates the
|
||||
# original session being evicted between watch fire and dispatch).
|
||||
# The restore_fn captures *which* fresh ChatSession got built so the
|
||||
# test can assert the queue landed on it (not on the original).
|
||||
rehydrated_holder: dict[str, ChatSession] = {}
|
||||
|
||||
def _restore_fn(ws_id: str) -> Any:
|
||||
"""Mirror the production ``_watch_restore_fn`` closure shape:
|
||||
construct a fresh session, resume the persisted ws_id (so the
|
||||
new session adopts the original ws_id), wire the dispatch
|
||||
closure, return the dispatch fn.
|
||||
"""
|
||||
new_session = _make_session()
|
||||
ok = new_session.resume(ws_id)
|
||||
assert ok, "resume should succeed against a non-empty message log"
|
||||
new_session.set_watch_runner(runner)
|
||||
rehydrated_holder["session"] = new_session
|
||||
return runner.get_dispatch_fn(new_session._ws_id)
|
||||
|
||||
runner = WatchRunner(
|
||||
storage=MagicMock(),
|
||||
node_id="test-node",
|
||||
restore_fn=_restore_fn,
|
||||
)
|
||||
|
||||
# Sanity: no dispatch fn registered yet for the original ws_id.
|
||||
assert runner.get_dispatch_fn(original_ws_id) is None
|
||||
|
||||
# Stage 3 — fire a watch result. ``_dispatch_result`` should fall
|
||||
# through to the restore branch. The dispatch surface takes a
|
||||
# structured reminder dict.
|
||||
runner._dispatch_result(
|
||||
original_ws_id,
|
||||
{"type": "watch_triggered", "text": "post-restore body"},
|
||||
"watch-1",
|
||||
)
|
||||
|
||||
# The restore fn ran exactly once and produced a fresh session that
|
||||
# adopted the original ws_id.
|
||||
assert "session" in rehydrated_holder, "restore_fn was not invoked"
|
||||
rehydrated = rehydrated_holder["session"]
|
||||
assert rehydrated is not original
|
||||
assert rehydrated._ws_id == original_ws_id
|
||||
|
||||
# The watch payload landed on the rehydrated session's queue, not on
|
||||
# the (now-evicted) original session's queue.
|
||||
assert len(rehydrated._nudge_queue) == 1
|
||||
assert rehydrated._nudge_queue.pending(channel="any") == [
|
||||
("watch_triggered", "post-restore body")
|
||||
]
|
||||
# Original session's queue stays empty — the dispatch did NOT
|
||||
# accidentally route back to it.
|
||||
assert len(original._nudge_queue) == 0
|
||||
@@ -72,6 +72,20 @@ class TestWatchCRUD:
|
||||
assert db.delete_watch("nope") is False
|
||||
|
||||
|
||||
class TestIsWatchActive:
|
||||
def test_active_row_returns_true(self, db):
|
||||
db.create_watch(**_make_watch_kwargs())
|
||||
assert db.is_watch_active("watch_001") is True
|
||||
|
||||
def test_inactive_row_returns_false(self, db):
|
||||
db.create_watch(**_make_watch_kwargs())
|
||||
db.update_watch("watch_001", active=False)
|
||||
assert db.is_watch_active("watch_001") is False
|
||||
|
||||
def test_missing_row_returns_false(self, db):
|
||||
assert db.is_watch_active("nope") is False
|
||||
|
||||
|
||||
class TestWatchListQueries:
|
||||
def test_list_for_ws(self, db):
|
||||
db.create_watch(**_make_watch_kwargs(watch_id="w1", ws_id="ws-1", name="a"))
|
||||
|
||||
@@ -173,3 +173,50 @@ class TestResolveClient:
|
||||
def test_unknown_backend_returns_none(self):
|
||||
client = resolve_web_search_client("typo_backend", tavily_key="key")
|
||||
assert client is None
|
||||
|
||||
def test_resolve_web_search_client_rejects_oauth_user_backend(self):
|
||||
"""A web_search backend pointing at an ``auth_type=oauth_user``
|
||||
MCP server MUST be rejected at boot — per-node web_search
|
||||
cannot carry per-user tokens, so resolving the backend would
|
||||
guarantee a 401-on-call instead of a clean disablement.
|
||||
|
||||
Phase 7 invariant 8 corollary: pool tools are user-scoped;
|
||||
every entry point that lacks per-user identity (web_search
|
||||
boot resolver, eval harness, CLI default) MUST refuse them
|
||||
rather than silently produce a broken client.
|
||||
|
||||
Verified by reverting the ``server_auth_type(...) == 'oauth_user'``
|
||||
guard in ``resolve_web_search_client``: the resolver returns
|
||||
an ``MCPSearchClient`` whose ``call_tool_sync`` would surface
|
||||
a 401 / consent_required structured error on every search.
|
||||
"""
|
||||
mcp = MagicMock()
|
||||
mcp.is_mcp_tool.return_value = True # name resolves
|
||||
mcp.server_auth_type.return_value = "oauth_user"
|
||||
client = resolve_web_search_client(
|
||||
"mcp:oauth-search:search", tavily_key=None, mcp_client=mcp
|
||||
)
|
||||
assert client is None, (
|
||||
"oauth_user-backed web_search backend resolved to a non-None client; "
|
||||
"boot-time guard missing or regressed."
|
||||
)
|
||||
# Per-turn callers must read from the in-memory cache, never
|
||||
# the SQL helper — perf regression guard.
|
||||
mcp.server_auth_type.assert_called_with("oauth-search")
|
||||
assert not mcp._lookup_server_row.called, (
|
||||
"resolver issued a SQL roundtrip via _lookup_server_row; "
|
||||
"per-turn web_search backend resolution must use the "
|
||||
"in-memory server_auth_type accessor."
|
||||
)
|
||||
|
||||
def test_resolve_web_search_client_accepts_static_backend(self):
|
||||
"""Static-path (``auth_type=none`` or ``static``) MCP backends
|
||||
still resolve cleanly — the new guard ONLY rejects oauth_user.
|
||||
"""
|
||||
mcp = MagicMock()
|
||||
mcp.is_mcp_tool.return_value = True
|
||||
mcp.server_auth_type.return_value = None
|
||||
client = resolve_web_search_client(
|
||||
"mcp:static-search:search", tavily_key=None, mcp_client=mcp
|
||||
)
|
||||
assert isinstance(client, MCPSearchClient)
|
||||
|
||||
@@ -155,3 +155,179 @@ class TestContentAccumulation:
|
||||
assert len(idle_events) == 1
|
||||
# Content should be capped, not contain everything
|
||||
assert len(idle_events[0]["content"]) <= _MAX_TURN_CONTENT_CHARS + 1024
|
||||
|
||||
|
||||
class TestPendingApprovalDetailNotPiggybacked:
|
||||
"""Stage 3 cleanup — ``pending_approval_detail`` is no longer
|
||||
piggybacked on ``ws_state`` events. Approval items now arrive via
|
||||
bulk fetch when the coord tree's reducer sees the
|
||||
``activity_state="approval"`` transition; verdicts via the explicit
|
||||
``intent_verdict`` event class; resolution via
|
||||
``approval_resolved``. These tests lock the no-piggyback contract
|
||||
down so a future regression doesn't silently re-introduce the
|
||||
duplicated path."""
|
||||
|
||||
def test_state_broadcast_omits_field_when_no_approval_pending(self):
|
||||
ui = _make_ui()
|
||||
assert ui._pending_approval is None
|
||||
ui._broadcast_state("running")
|
||||
|
||||
events = _drain_global()
|
||||
running_events = [e for e in events if e.get("state") == "running"]
|
||||
assert len(running_events) == 1
|
||||
assert "pending_approval_detail" not in running_events[0]
|
||||
|
||||
def test_state_broadcast_omits_field_even_when_approval_pending(self):
|
||||
"""The piggyback is gone: even when ``_pending_approval`` is set,
|
||||
the state broadcast must NOT carry ``pending_approval_detail``.
|
||||
The browser triggers a bulk fetch off the
|
||||
``activity_state="approval"`` transition to get the items."""
|
||||
ui = _make_ui()
|
||||
ui._pending_approval = {
|
||||
"type": "approve_request",
|
||||
"items": [
|
||||
{
|
||||
"call_id": "c1",
|
||||
"header": "tool x",
|
||||
"func_args": "{}",
|
||||
"intent_summary": "do x",
|
||||
"needs_approval": True,
|
||||
}
|
||||
],
|
||||
"judge_pending": False,
|
||||
}
|
||||
ui._broadcast_state("attention")
|
||||
|
||||
events = _drain_global()
|
||||
attn = [e for e in events if e.get("state") == "attention"]
|
||||
assert len(attn) == 1
|
||||
assert "pending_approval_detail" not in attn[0]
|
||||
|
||||
def test_field_stays_absent_after_approval_resolves(self):
|
||||
ui = _make_ui()
|
||||
ui._pending_approval = {
|
||||
"type": "approve_request",
|
||||
"items": [{"call_id": "c1", "header": "x"}],
|
||||
"judge_pending": False,
|
||||
}
|
||||
ui._broadcast_state("attention")
|
||||
_drain_global()
|
||||
|
||||
ui._pending_approval = None
|
||||
ui._broadcast_state("running")
|
||||
events = _drain_global()
|
||||
running = [e for e in events if e.get("state") == "running"]
|
||||
assert len(running) == 1
|
||||
assert "pending_approval_detail" not in running[0]
|
||||
|
||||
|
||||
class TestBroadcastIntentVerdict:
|
||||
"""Producer-side coverage for ``WebUI._broadcast_intent_verdict``.
|
||||
|
||||
The collector-side test (``test_apply_delta_intent_verdict_*`` in
|
||||
test_console.py) covers consumption; this pins the event shape the
|
||||
producer puts on the global queue. A field rename or missed key
|
||||
here would slip past the consumer test because the consumer reads
|
||||
via ``data.get(...)``.
|
||||
"""
|
||||
|
||||
def test_pushes_intent_verdict_event_to_global_queue(self):
|
||||
ui = _make_ui()
|
||||
verdict = {
|
||||
"call_id": "c1",
|
||||
"risk_level": "low",
|
||||
"confidence": 0.92,
|
||||
"recommendation": "approve",
|
||||
"reasoning": "tool reads only",
|
||||
}
|
||||
ui._broadcast_intent_verdict(verdict)
|
||||
|
||||
events = _drain_global()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["type"] == "intent_verdict"
|
||||
assert ev["ws_id"] == "ws-test"
|
||||
assert ev["verdict"] == verdict
|
||||
|
||||
def test_no_op_when_global_queue_unset(self):
|
||||
WebUI._global_queue = None
|
||||
ui = _make_ui()
|
||||
# Doesn't raise.
|
||||
ui._broadcast_intent_verdict({"call_id": "c1"})
|
||||
|
||||
def test_queue_full_swallowed(self):
|
||||
# Force a tiny queue then fill it so the next put_nowait
|
||||
# raises queue.Full — the broadcast must absorb it without
|
||||
# propagating (matches _broadcast_state's queue.Full handling).
|
||||
WebUI._global_queue = queue.Queue(maxsize=1)
|
||||
WebUI._global_queue.put_nowait({"sentinel": True})
|
||||
ui = _make_ui()
|
||||
# Doesn't raise.
|
||||
ui._broadcast_intent_verdict({"call_id": "c1"})
|
||||
|
||||
|
||||
class TestBroadcastApprovalResolved:
|
||||
"""Producer-side coverage for ``WebUI._broadcast_approval_resolved``."""
|
||||
|
||||
def test_pushes_approval_resolved_event_to_global_queue(self):
|
||||
ui = _make_ui()
|
||||
ui._broadcast_approval_resolved(True, "lgtm", always=False)
|
||||
|
||||
events = _drain_global()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["type"] == "approval_resolved"
|
||||
assert ev["ws_id"] == "ws-test"
|
||||
assert ev["approved"] is True
|
||||
assert ev["feedback"] == "lgtm"
|
||||
assert ev["always"] is False
|
||||
|
||||
def test_normalises_none_feedback_to_empty_string(self):
|
||||
ui = _make_ui()
|
||||
ui._broadcast_approval_resolved(False, None)
|
||||
|
||||
events = _drain_global()
|
||||
assert events[0]["feedback"] == ""
|
||||
assert events[0]["approved"] is False
|
||||
assert events[0]["always"] is False
|
||||
|
||||
def test_always_kwarg_propagates(self):
|
||||
ui = _make_ui()
|
||||
ui._broadcast_approval_resolved(True, "ok", always=True)
|
||||
events = _drain_global()
|
||||
assert events[0]["always"] is True
|
||||
|
||||
def test_no_op_when_global_queue_unset(self):
|
||||
WebUI._global_queue = None
|
||||
ui = _make_ui()
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approval_resolved(True, None)
|
||||
|
||||
|
||||
class TestBroadcastApproveRequest:
|
||||
"""Producer-side coverage for ``WebUI._broadcast_approve_request`` —
|
||||
push path for the initial approval items so a coord parent's tree
|
||||
UI can render the inline approve/deny block immediately without
|
||||
waiting for a bulk-fetch round-trip."""
|
||||
|
||||
def test_pushes_approve_request_event_to_global_queue(self):
|
||||
ui = _make_ui()
|
||||
detail = {
|
||||
"type": "approve_request",
|
||||
"items": [{"call_id": "c1", "header": "tool x"}],
|
||||
"judge_pending": True,
|
||||
}
|
||||
ui._broadcast_approve_request(detail)
|
||||
|
||||
events = _drain_global()
|
||||
assert len(events) == 1
|
||||
ev = events[0]
|
||||
assert ev["type"] == "approve_request"
|
||||
assert ev["ws_id"] == "ws-test"
|
||||
assert ev["detail"] == detail
|
||||
|
||||
def test_no_op_when_global_queue_unset(self):
|
||||
WebUI._global_queue = None
|
||||
ui = _make_ui()
|
||||
# Doesn't raise.
|
||||
ui._broadcast_approve_request({"items": []})
|
||||
|
||||
@@ -899,6 +899,356 @@ class TestHistoryInteractive:
|
||||
assert client.get(base, params={"limit": 999}).status_code == 200
|
||||
|
||||
|
||||
class TestBuildHistoryReminderPropagation:
|
||||
"""``_build_history`` must surface the ``_reminders`` side-channel on
|
||||
each entry so a tab reconnecting via ``/history`` renders the same
|
||||
metacognitive nudge bubble the originating tab saw via the live
|
||||
``user_reminder`` SSE event.
|
||||
"""
|
||||
|
||||
def _session_with_messages(self, messages: list[dict]) -> MagicMock:
|
||||
session = MagicMock()
|
||||
session.messages = messages
|
||||
return session
|
||||
|
||||
def test_reminders_sidechannel_surfaces_on_entry(self):
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "ah no",
|
||||
"_reminders": [{"type": "correction", "text": "watch out"}],
|
||||
}
|
||||
]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == "ah no"
|
||||
assert history[0]["reminders"] == [{"type": "correction", "text": "watch out"}]
|
||||
|
||||
def test_no_reminders_key_when_sidechannel_absent(self):
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages([{"role": "user", "content": "just a message"}])
|
||||
history = _build_history(session)
|
||||
assert "reminders" not in history[0]
|
||||
|
||||
def test_no_reminders_key_when_sidechannel_empty(self):
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages([{"role": "user", "content": "hi", "_reminders": []}])
|
||||
history = _build_history(session)
|
||||
assert "reminders" not in history[0]
|
||||
|
||||
def test_multiple_reminders_preserved_in_order(self):
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "x",
|
||||
"_reminders": [
|
||||
{"type": "denial", "text": "FIRST"},
|
||||
{"type": "correction", "text": "SECOND"},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["reminders"] == [
|
||||
{"type": "denial", "text": "FIRST"},
|
||||
{"type": "correction", "text": "SECOND"},
|
||||
]
|
||||
|
||||
def test_reminders_coexist_with_attachments(self):
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "look"},
|
||||
{"type": "image_url", "image_url": {"url": "data:..."}},
|
||||
],
|
||||
"_reminders": [{"type": "correction", "text": "watch"}],
|
||||
}
|
||||
]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == "look"
|
||||
assert history[0]["attachments"] == [{"kind": "image", "filename": "", "mime_type": ""}]
|
||||
assert history[0]["reminders"] == [{"type": "correction", "text": "watch"}]
|
||||
|
||||
def test_malformed_reminders_filtered_out(self):
|
||||
"""Defensive: a non-dict element in the list (corruption / bug)
|
||||
is dropped rather than crashing the history serialisation."""
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "x",
|
||||
"_reminders": [
|
||||
{"type": "correction", "text": "ok"},
|
||||
"not-a-dict",
|
||||
{"type": "denial"}, # missing text
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
history = _build_history(session)
|
||||
# Non-dicts dropped; missing-text fills with empty string.
|
||||
assert history[0]["reminders"] == [
|
||||
{"type": "correction", "text": "ok"},
|
||||
{"type": "denial", "text": ""},
|
||||
]
|
||||
|
||||
def test_clean_message_passes_through_unchanged(self):
|
||||
"""No reminders, plain content — _build_history is a no-op for the
|
||||
reminder field and ``content`` rides through verbatim."""
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[{"role": "user", "content": "just a normal message"}]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == "just a normal message"
|
||||
assert "reminders" not in history[0]
|
||||
|
||||
def test_assistant_content_with_literal_reminder_tag_unchanged(self):
|
||||
"""Assistant output may legitimately reference the tag (e.g. when
|
||||
the model is explaining the reminder system itself). No
|
||||
transformation should ever apply to assistant content."""
|
||||
from turnstone.server import _build_history
|
||||
|
||||
content = "Here is a <system-reminder> tag in assistant output."
|
||||
session = self._session_with_messages([{"role": "assistant", "content": content}])
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == content
|
||||
|
||||
|
||||
class TestBuildHistoryAdvisoryRoundTrip:
|
||||
"""``_build_history`` must round-trip the persisted
|
||||
``<tool_output>`` envelope (Seam 1 queued-message splice) to
|
||||
cleaned content + a wire-shape ``advisories`` array.
|
||||
|
||||
Production realism note: ``session.messages`` never carries an
|
||||
``advisories`` key — only ``decorate_history_messages`` mutates
|
||||
dicts to add it for the REST ``/history`` path, and the SSE replay
|
||||
surface bypasses that decoration entirely. The earlier
|
||||
``TestBuildHistoryAdvisoryPropagation`` class pre-populated
|
||||
``advisories`` directly on the session messages, which tested a
|
||||
passthrough that doesn't exist in production — the SSE replay code
|
||||
path silently dropped queued messages despite the green tests.
|
||||
These round-trip tests exercise the production shape (wrapped
|
||||
envelope on the tool row's ``content``) so a regression in the
|
||||
inline ``extract_advisories_from_tool_envelope`` call inside
|
||||
``_build_history`` surfaces here.
|
||||
"""
|
||||
|
||||
def _session_with_messages(self, messages: list[dict]) -> MagicMock:
|
||||
session = MagicMock()
|
||||
session.messages = messages
|
||||
return session
|
||||
|
||||
def test_build_history_round_trips_envelope_to_advisories(self):
|
||||
"""The production-realistic shape: a tool row whose ``content``
|
||||
is the wrapped ``<tool_output>`` envelope (no ``advisories``
|
||||
key set — that's the bug-1 footprint). ``_build_history``
|
||||
must extract the advisory back out and ship it on the wire as
|
||||
cleaned content + ``advisories``.
|
||||
|
||||
Reverting the inline ``extract_advisories_from_tool_envelope``
|
||||
call in ``server._build_history``'s tool-message branch breaks
|
||||
this test.
|
||||
"""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
from turnstone.server import _build_history
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"tool body",
|
||||
[UserInterjection(message="check logs", priority="notice")],
|
||||
)
|
||||
session = self._session_with_messages(
|
||||
[
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_a",
|
||||
"content": wrapped,
|
||||
}
|
||||
]
|
||||
)
|
||||
history = _build_history(session)
|
||||
# Cleaned content rides on the wire — envelope stripped.
|
||||
assert history[0]["content"] == "tool body"
|
||||
# Advisory survives as a wire-shape entry the JS can render
|
||||
# as a user bubble after the tool block.
|
||||
assert history[0]["advisories"] == [
|
||||
{"type": "user_interjection", "text": "check logs", "priority": "notice"}
|
||||
]
|
||||
|
||||
def test_build_history_round_trips_important_priority(self):
|
||||
"""The ``important`` priority preamble round-trips — pin both
|
||||
the priority detection in the parser and the projection through
|
||||
to the wire shape."""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
from turnstone.server import _build_history
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"out",
|
||||
[UserInterjection(message="urgent", priority="important")],
|
||||
)
|
||||
session = self._session_with_messages(
|
||||
[{"role": "tool", "tool_call_id": "call_a", "content": wrapped}]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == "out"
|
||||
assert history[0]["advisories"] == [
|
||||
{"type": "user_interjection", "text": "urgent", "priority": "important"}
|
||||
]
|
||||
|
||||
def test_build_history_no_envelope_passes_through_unchanged(self):
|
||||
"""Plain tool content (no ``<tool_output>`` prefix) — no
|
||||
advisories field, content unchanged."""
|
||||
from turnstone.server import _build_history
|
||||
|
||||
session = self._session_with_messages(
|
||||
[{"role": "tool", "tool_call_id": "call_a", "content": "plain output"}]
|
||||
)
|
||||
history = _build_history(session)
|
||||
assert history[0]["content"] == "plain output"
|
||||
assert "advisories" not in history[0]
|
||||
|
||||
def test_build_history_round_trip_through_full_decoration_chain(self):
|
||||
"""End-to-end pin: persist a wrapped envelope into ``messages``,
|
||||
run the full decoration chain (``decorate_history_messages``
|
||||
followed by ``_build_history``), assert the wire shape carries
|
||||
the advisory. This pins the contract every component in the
|
||||
chain participates in — REST ``/history`` callers go through
|
||||
``decorate_history_messages``, and SSE replay goes through
|
||||
``_build_history`` — both must produce the same wire shape.
|
||||
"""
|
||||
from turnstone.core.history_decoration import decorate_history_messages
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
from turnstone.server import _build_history
|
||||
|
||||
wrapped = wrap_tool_result(
|
||||
"raw",
|
||||
[UserInterjection(message="hi", priority="notice")],
|
||||
)
|
||||
# Decorate first — REST /history shape.
|
||||
rest_messages: list[dict] = [{"role": "tool", "tool_call_id": "call_a", "content": wrapped}]
|
||||
decorate_history_messages(rest_messages, {}, {})
|
||||
# And separately drive _build_history with a fresh undecorated
|
||||
# message — SSE replay shape.
|
||||
session = self._session_with_messages(
|
||||
[{"role": "tool", "tool_call_id": "call_a", "content": wrapped}]
|
||||
)
|
||||
sse_history = _build_history(session)
|
||||
# Both surfaces produce the same advisory + cleaned content.
|
||||
assert rest_messages[0]["content"] == "raw"
|
||||
assert rest_messages[0]["advisories"] == [
|
||||
{"type": "user_interjection", "text": "hi", "priority": "notice"}
|
||||
]
|
||||
assert sse_history[0]["content"] == "raw"
|
||||
assert sse_history[0]["advisories"] == [
|
||||
{"type": "user_interjection", "text": "hi", "priority": "notice"}
|
||||
]
|
||||
|
||||
def test_build_history_extracts_advisories_from_list_content_text_part(self):
|
||||
"""List-typed tool output (image / structured MCP results)
|
||||
with a Seam 1 splice carries the wrap envelope as a separate
|
||||
text part (``session.py``'s tool-result loop appends
|
||||
``{"type": "text", "text": wrap_tool_result("", advisories)}``
|
||||
when ``output`` is a list). ``_build_history`` must walk the
|
||||
list parts, extract advisories from any wrap-envelope text
|
||||
part, and DROP that text part from the projected list — the
|
||||
cleaned inner content is empty by construction, and leaving
|
||||
the part would cause the JS replay to render the literal
|
||||
envelope text as a chunk inside the tool block AND fail to
|
||||
render the queued message as a user bubble.
|
||||
|
||||
Removing the list-content branch in ``_build_history``'s tool-
|
||||
message advisory extraction breaks this test.
|
||||
"""
|
||||
from turnstone.core.tool_advisory import UserInterjection, wrap_tool_result
|
||||
from turnstone.server import _build_history
|
||||
|
||||
wrap_text = wrap_tool_result(
|
||||
"",
|
||||
[UserInterjection(message="inspect histogram", priority="notice")],
|
||||
)
|
||||
list_content = [
|
||||
{"type": "text", "text": "the chart shows X"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,xxx"}},
|
||||
{"type": "text", "text": wrap_text},
|
||||
]
|
||||
session = self._session_with_messages(
|
||||
[{"role": "tool", "tool_call_id": "call_a", "content": list_content}]
|
||||
)
|
||||
history = _build_history(session)
|
||||
# Wire-shape content keeps the original text + image parts but
|
||||
# has the wrap text-part dropped.
|
||||
wire_content = history[0]["content"]
|
||||
assert isinstance(wire_content, list)
|
||||
assert len(wire_content) == 2
|
||||
assert wire_content[0] == {"type": "text", "text": "the chart shows X"}
|
||||
assert wire_content[1] == {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,xxx"},
|
||||
}
|
||||
# Advisory rides on the wire so JS replay renders the user
|
||||
# bubble after the tool block — same contract as the string-
|
||||
# content path.
|
||||
assert history[0]["advisories"] == [
|
||||
{
|
||||
"type": "user_interjection",
|
||||
"text": "inspect histogram",
|
||||
"priority": "notice",
|
||||
}
|
||||
]
|
||||
|
||||
def test_build_history_keeps_legitimate_envelope_text_part_with_body(self):
|
||||
"""A tool that legitimately produces output containing a
|
||||
well-formed ``<tool_output>`` envelope as a text part (e.g.
|
||||
documentation viewer, code analyzer demoing the wrapper, an
|
||||
echo tool) must NOT have that part dropped on replay. The
|
||||
list-content drop heuristic must require both an empty cleaned
|
||||
inner body AND at least one extracted advisory — the
|
||||
signature of the injected ``wrap_tool_result("", advisories)``
|
||||
carrier. A legitimate tool envelope has non-empty inner body
|
||||
OR no advisories, and stays in the projected list verbatim.
|
||||
|
||||
Removing the ``not cleaned_text and advisories_from_part``
|
||||
guard breaks this test (the legitimate envelope gets dropped
|
||||
from the wire content)."""
|
||||
from turnstone.server import _build_history
|
||||
|
||||
legit_envelope_text = (
|
||||
"<tool_output>\nThis is what a tool_output envelope looks like.\n</tool_output>"
|
||||
)
|
||||
list_content = [
|
||||
{"type": "text", "text": "doc preview:"},
|
||||
{"type": "text", "text": legit_envelope_text},
|
||||
]
|
||||
session = self._session_with_messages(
|
||||
[{"role": "tool", "tool_call_id": "call_a", "content": list_content}]
|
||||
)
|
||||
history = _build_history(session)
|
||||
# All parts survive — none dropped.
|
||||
wire_content = history[0]["content"]
|
||||
assert isinstance(wire_content, list)
|
||||
assert len(wire_content) == 2
|
||||
assert wire_content[1]["text"] == legit_envelope_text
|
||||
# No advisories surfaced (no system-reminder blocks were
|
||||
# extracted from the legitimate envelope).
|
||||
assert "advisories" not in history[0]
|
||||
|
||||
|
||||
class TestDetailInteractive:
|
||||
"""Interactive parity for the lifted ``GET /v1/api/workstreams/{ws_id}``.
|
||||
|
||||
|
||||
@@ -117,7 +117,6 @@
|
||||
|
||||
[mcp]
|
||||
# config_path = "" # Path to MCP servers config file (JSON)
|
||||
# refresh_interval = 14400 # Refresh interval in seconds (default: 4h)
|
||||
|
||||
# --- Server (node, console) ---
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
"""turnstone - Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."""
|
||||
|
||||
__version__ = "1.5.0"
|
||||
__version__ = "1.5.8"
|
||||
|
||||
@@ -674,6 +674,16 @@ class McpServerInfo(BaseModel):
|
||||
registry_name: str | None = None
|
||||
registry_version: str = ""
|
||||
registry_meta: str = "{}"
|
||||
auth_type: str = "static"
|
||||
oauth_client_id: str | None = None
|
||||
oauth_scopes: str | None = None
|
||||
oauth_audience: str | None = None
|
||||
oauth_registration_mode: str | None = None
|
||||
oauth_authorization_server_url: str | None = None
|
||||
oauth_as_issuer_cached: str | None = None
|
||||
# Fernet ciphertext; never decrypted on the read path. Responses
|
||||
# carry the masked ``"***"`` sentinel via ``_mask_mcp_secrets``.
|
||||
oauth_client_secret_ct: str | None = None
|
||||
created: str
|
||||
updated: str
|
||||
|
||||
@@ -704,6 +714,16 @@ class CreateMcpServerRequest(BaseModel):
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
auto_approve: bool = False
|
||||
enabled: bool = True
|
||||
# OAuth-MCP: one of 'none' | 'static' | 'oauth_user'.
|
||||
# ``oauth_client_secret`` is plaintext input; never persisted,
|
||||
# redacted in audit log.
|
||||
auth_type: str = "static"
|
||||
oauth_client_id: str | None = None
|
||||
oauth_client_secret: str | None = None
|
||||
oauth_scopes: str | None = None
|
||||
oauth_audience: str | None = None
|
||||
oauth_registration_mode: str | None = None
|
||||
oauth_authorization_server_url: str | None = None
|
||||
|
||||
|
||||
class UpdateMcpServerRequest(BaseModel):
|
||||
@@ -716,6 +736,13 @@ class UpdateMcpServerRequest(BaseModel):
|
||||
env: dict[str, str] | None = None
|
||||
auto_approve: bool | None = None
|
||||
enabled: bool | None = None
|
||||
auth_type: str | None = None
|
||||
oauth_client_id: str | None = None
|
||||
oauth_client_secret: str | None = None
|
||||
oauth_scopes: str | None = None
|
||||
oauth_audience: str | None = None
|
||||
oauth_registration_mode: str | None = None
|
||||
oauth_authorization_server_url: str | None = None
|
||||
|
||||
|
||||
class ListMcpServersResponse(BaseModel):
|
||||
@@ -765,6 +792,34 @@ class SkillDiscoverResponse(BaseModel):
|
||||
skills: list[SkillDiscoverListing]
|
||||
|
||||
|
||||
class ParseSkillRequest(BaseModel):
|
||||
raw: str = Field(
|
||||
min_length=1,
|
||||
max_length=32_768,
|
||||
description=(
|
||||
"Raw SKILL.md text — YAML frontmatter delimited by ``---`` "
|
||||
"followed by the markdown body. Capped at 32 KiB to match "
|
||||
"``admin_create_skill``'s ``content`` ceiling and to bound "
|
||||
"the synchronous YAML parser's worst-case CPU cost. The "
|
||||
"handler reuses the Python parser at "
|
||||
"``turnstone.core.skill_parser`` so admin UIs and external "
|
||||
"import paths agree on field extraction."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ParseSkillResponse(BaseModel):
|
||||
name: str
|
||||
description: str
|
||||
content: str
|
||||
tags: list[str] = Field(default_factory=list)
|
||||
author: str = ""
|
||||
version: str = "1.0.0"
|
||||
allowed_tools: list[str] = Field(default_factory=list)
|
||||
license: str = ""
|
||||
compatibility: str = ""
|
||||
|
||||
|
||||
class SkillInstallRequest(BaseModel):
|
||||
source: str # "skills.sh" or "github"
|
||||
skill_id: str = "" # for skills.sh
|
||||
|
||||
@@ -77,6 +77,8 @@ from turnstone.api.console_schemas import (
|
||||
NodeMetadataResponse,
|
||||
OrgInfo,
|
||||
OutputAssessmentInfo,
|
||||
ParseSkillRequest,
|
||||
ParseSkillResponse,
|
||||
RegistryInstallRequest,
|
||||
RegistrySearchResponse,
|
||||
RoleInfo,
|
||||
@@ -552,6 +554,15 @@ CONSOLE_ENDPOINTS: list[EndpointSpec] = [
|
||||
error_codes=[400, 404, 409, 502],
|
||||
tags=["Admin"],
|
||||
),
|
||||
EndpointSpec(
|
||||
"/v1/api/admin/skills/parse",
|
||||
"POST",
|
||||
"Parse a SKILL.md document and return its frontmatter fields and body",
|
||||
request_model=ParseSkillRequest,
|
||||
response_model=ParseSkillResponse,
|
||||
error_codes=[400, 413],
|
||||
tags=["Admin"],
|
||||
),
|
||||
# --- Governance: Skills ---
|
||||
EndpointSpec(
|
||||
"/v1/api/admin/skills",
|
||||
@@ -1594,6 +1605,8 @@ _ALL_MODELS: list[type[BaseModel]] = [
|
||||
SkillInstallResponse,
|
||||
SkillInfo,
|
||||
SkillVersionInfo,
|
||||
ParseSkillRequest,
|
||||
ParseSkillResponse,
|
||||
CreateSkillRequest,
|
||||
UpdateSkillRequest,
|
||||
ListSkillsResponse,
|
||||
|
||||
+26
-11
@@ -312,6 +312,31 @@ class TerminalUI(SessionUI):
|
||||
sys.stdout.write(f"{RED}{message}{RESET}\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def _print_reminder(self, reminders: list[dict[str, Any]]) -> None:
|
||||
"""Render a metacognitive reminder list as ``[metacognition · type] text``
|
||||
lines in the terminal — the CLI's equivalent of the web UI's
|
||||
yellow themed bubble. Used by both ``on_user_reminder`` and
|
||||
``on_tool_reminder``; the rendering is identical because
|
||||
terminal output is anchor-by-flow rather than DOM-by-anchor.
|
||||
"""
|
||||
for r in reminders:
|
||||
nt = str(r.get("type", "") or "")
|
||||
text = str(r.get("text", "") or "")
|
||||
label = "metacognition" + (f" · {nt}" if nt else "")
|
||||
sys.stdout.write(f"{YELLOW}[{label}]{RESET} {text}\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def on_user_reminder(self, reminders: list[dict[str, Any]], source: str | None = None) -> None:
|
||||
# ``source`` ignored — the CLI doesn't render a wake marker
|
||||
# (terminal output is anchored by sequence, not anchor element).
|
||||
self._print_reminder(reminders)
|
||||
|
||||
def on_tool_reminder(self, reminders: list[dict[str, Any]], tool_call_id: str) -> None:
|
||||
# tool_call_id ignored — the CLI anchors by output sequence
|
||||
# (the line lands directly after the tool result that
|
||||
# triggered the batch's reminder).
|
||||
self._print_reminder(reminders)
|
||||
|
||||
def on_state_change(self, state: str) -> None:
|
||||
pass # base TerminalUI ignores state changes
|
||||
|
||||
@@ -992,15 +1017,6 @@ def main() -> None:
|
||||
help="Path to MCP server config file (standard mcpServers JSON format)",
|
||||
)
|
||||
|
||||
from turnstone.core.config import nonneg_float
|
||||
|
||||
parser.add_argument(
|
||||
"--mcp-refresh-interval",
|
||||
type=nonneg_float,
|
||||
default=14400,
|
||||
metavar="SECONDS",
|
||||
help="Periodic MCP tool refresh interval for servers without push notifications (default: 14400 = 4h, 0 to disable)",
|
||||
)
|
||||
judge_group = parser.add_argument_group("Judge options")
|
||||
judge_group.add_argument(
|
||||
"--judge",
|
||||
@@ -1113,7 +1129,6 @@ def main() -> None:
|
||||
|
||||
mcp_client = create_mcp_client(
|
||||
getattr(args, "mcp_config", None),
|
||||
refresh_interval=getattr(args, "mcp_refresh_interval", 14400),
|
||||
storage=_get_storage(),
|
||||
)
|
||||
|
||||
@@ -1219,7 +1234,7 @@ def main() -> None:
|
||||
)
|
||||
sys.stderr.flush()
|
||||
|
||||
manager._on_state_change = _bg_attention_notify
|
||||
manager.subscribe_to_state(_bg_attention_notify)
|
||||
|
||||
# Print banner
|
||||
print(f"\n{bold('Chat')} with {cyan(model)}")
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user