mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-19 18:41:00 -06:00
Patrick Buckley
dad98c062d
Light up production reachability of pool dispatch (RFC §3, invariant 8)
by widening the public catalog API to optionally take a ``user_id``:
- ``MCPClientManager.get_tools(user_id=None)`` returns the merged
static + per-user pool view when ``user_id`` is supplied; the default
preserves the legacy global-only contract.
- ``is_mcp_tool(name, *, user_id=None)`` extends the lookup to the
per-user ``_user_tool_map``. Pool tools become reachable from
``ChatSession._prepare_tool`` only when the session-bound user_id
flows through — flipping invariant 8 from "must hold" to "satisfied".
- Listener identity becomes ``(user_id, callback)``. Static-path
changes fire ALL listeners (admin + every user); pool-entry
changes fire only matching-user + admin (``None``) listeners.
RFC §3.3.
- Pool sessions discover their tool list on first connect
(``_connect_one_pool`` → ``await session.list_tools()``); the
notification closure binds to ``(user_id, server_name)`` so
push-driven ``list_changed`` updates target the correct user's
catalog. R6 verified empirically: ``list_tools()`` 401 propagates
through anyio TaskGroup unwinding, no hang — plain ``await`` is
fine, no carrier-race shape needed for discovery.
- ``_evict_session`` drops ``entry.tools`` and rebuilds the user's
index so an evicted-then-reconnected session doesn't carry
stale catalog state.
- ``web_search.resolve_web_search_client`` refuses
``auth_type=oauth_user`` backends (per-node web search can't
carry per-user tokens).
Resources / prompts pool dispatch deferred to Phase 7b — invariant 8
is satisfied by the tool path alone, and the resource/prompt path
needs sibling ``_dispatch_pool_resource_sync`` /
``_dispatch_pool_prompt_sync`` helpers each with their own
carrier-race plumbing (~400 LOC). Phase 7b will follow the patterns
established here.
CLI sessions default ``user_id=""`` and so cannot use oauth_user
MCP servers — documented limitation; users must use the web UI.
Round-1 review fixes (4-finder review applied, no push yet):
- bug-1: get_tools(user_id) was iterating _user_pool_entries from sync
threads while the mcp-loop concurrently mutated it (RuntimeError:
dictionary changed size during iteration). Now reads from a sibling
_user_tools dict updated atomically by _rebuild_user_tool_map.
- bug-2: _close_pool_entry_if_idle (LRU/TTL eviction) skipped the
catalog cleanup that _evict_session does — stale tools persisted
in _user_tool_map and ChatSession's tool list never rebuilt. Now
mirrors _evict_session.
- perf-1: _last_pool_notification_refresh debounce dict was never
pruned in either eviction path. Now popped alongside the entry.
- perf-3: web_search resolver was issuing a sync SQL query per LLM
turn to gate oauth_user backends. Now reads from the cached
in-memory config.
- sec-1: bearer token could leak into exc_info-rendered tracebacks
via Sentry/faulthandler. log.debug now uses structured fields,
not exc_info.
- sec-2: tools-per-server response now capped at 1000 (defensive,
mirrors _MAX_ERROR_LEN / _MAX_INSUFFICIENT_SCOPE_REPORTED).
- Test cleanup: dropped two listener fan-out tests duplicating
test_mcp_client.py coverage; renamed test_pool_session_notification_handler
to match its actual scope (_refresh_pool_server_tools); removed
stale comments referencing /tmp/r6-spike*.py scratchpads and a
misleading "copy-on-write" comment.
Round-2 pre-push review fixes (focused single-pass review applied):
- round2-1: bug-2's catalog-cleanup block in _close_pool_entry_if_idle
had no integration test (exactly the failure mode flagged in
feedback_tests_through_boundaries.md). Added
test_close_pool_entry_if_idle_clears_catalog_and_fires_listener
driving the LRU/TTL eviction path through real streamablehttp_client +
MockTransport. Negative-test verified: reverting the
_rebuild_user_tool_map / _notify_user_tool_listeners calls makes
the new test fail.
- round2-3: documented the _oauth_user_server_names cache invariant
in add_server_sync / remove_server_sync docstrings. Cache is
reconcile_sync's sole owner — direct callers leave it stale, but
_db_servers_to_config strips oauth_user rows so production paths
are unaffected. Static→oauth_user transitions correctly leave the
name in the cache because remove_server_sync drops the static
connection, not the cache identity.
- round2-6: strengthened test_rebuild_user_tool_map_populates and
test_rebuild_user_tool_map_drops_empty_user to assert on the
_user_tools sibling cache (bug-1 fix). Without this, a future
revert dropping the sibling write would still pass the unit
tests because get_tools coverage lives in separate tests.
Round-3 full-stack review fixes (multi-stage review on the final
state caught what the layered apply passes missed):
- q-1 REGRESSION: pool tool-discovery used asyncio.wait_for around
session.list_tools(), the exact pattern the f6a3b66 fix (and
feedback_asyncio_timeout_vs_wait_for.md) put in place to avoid.
Python 3.11's asyncio.wait_for wraps the inner coroutine in a
fresh task → cross-task scope-exit when the SDK's anyio TaskGroup
unwinds on a 401. Switched to `async with asyncio.timeout(...):`
pattern used by _safe_close_stack.
- sec-2: TOCTOU in _connect_one_pool — entry.tools was published
(via _rebuild_user_tool_map + listener fan-out) BEFORE entry.session
was assigned. A sync-thread reader could observe a tool whose
backing entry has session=None. Defence-in-depth — dispatch
re-fetches its own token and lazy-reconnects on session=None — but
reordering catches the race at the source. entry.session now
publishes BEFORE catalog visibility.
- bug-1: _close_pool_entry_if_idle's _user_pool_locks.pop ran
unconditionally after the try/finally, but the early-return
branches (entry None on re-check, in_flight > 0 under lock) skip
it via Python's return-through-finally semantics. The lock was
never popped on those paths. Now gated behind an `evicted` flag
set only on the success path; in_flight > 0 leaves the lock for
the active dispatcher to reuse, entry-None races leave the lock
for re-allocation by _ensure_pool_entry. Comment now describes
the actual semantics, not the original promise.
- bug-2: softened the _rebuild_user_tool_map docstring's atomicity
claim. The two-dict write is technically non-atomic across Python
statements; in practice the window is sub-microsecond on the
mcp-loop with no awaits between writes, and the listener fan-out
fires AFTER both writes complete. Docstring now says "back-to-back
on the mcp-loop" instead of "atomically alongside".
- q-3: dropped `hasattr(mcp_client, "server_auth_type")` defensive
check in web_search.py. The method ships in this commit; the
hasattr created a silent fallthrough that would let a future
rename silently re-enable oauth_user backends.
- q-4: surfaced the CLI / empty-user_id limitation in a docstring
comment at ChatSession.__init__'s self._user_id assignment. The
note previously lived only inside is_mcp_tool's docstring — a
future maintainer wiring CLI features against MCP pool servers
wouldn't think to read is_mcp_tool to find the constraint.
- q-2 + q-5: deleted a tautological duplicate test in
test_mcp_user_catalog.py whose docstring claimed to test
ChatSession.close but never instantiated a ChatSession (the
manager-level identity semantics are already covered by
test_listener_identity_includes_user_id in the same file and by
test_session_close_removes_listener_with_same_user_id in
test_mcp_client.py which DOES drive a ChatSession). Reworded a
misleading "fixture provides only 5s" comment to point at the
actual `_run_on_loop(..., timeout=5)` site.
- q-6: the `self._user_id or None` collapse repeated at 8 sites
across session.py. Cached once at __init__ as
``self._mcp_user_id`` (since ``_user_id`` is set once and never
mutated); 8 call sites now read the cached value. The empty-
string-is-CLI-sentinel invariant is documented at the assignment
site, not re-asserted at each consumer.
Deferred to follow-up:
- sec-1: a hostile MCP server bound to user-A could craft a
tool.name containing `__` to synthesize a prefixed-name collision
in user-A's own catalog. Bounded impact: cross-tenant dispatch is
prevented by the per-tenant token gate in _dispatch_pool, and
user-B's get_tools(user_id="B") never includes user-A's pool
entries. The fix needs policy decisions (reject vs. sanitize)
and touches _mcp_to_openai which is shared between static and
pool paths; better discussed in its own follow-up where the
policy applies uniformly to static-path servers too. The threat
model already requires user-A to have consented to a malicious
server, who has many more dangerous vectors than tool-name
shenanigans.
Test count delta: +31 tests (5435 → 5466, ``-m "not live"``; one
test deleted in round-3 apply per q-2):
- ``tests/test_mcp_client.py`` +20 (per-user catalog state, listener
identity, session thread-through)
- ``tests/test_mcp_user_catalog.py`` +9 NEW (integration tests
driving real ``streamablehttp_client`` + ``httpx.MockTransport`` per
invariant 14: discovery on connect, user isolation, eviction +
reconnect, LRU/TTL eviction (round2-1), R6 401-propagation
regression, static byte-identical canonical regression; review
passes dropped duplicate listener fan-out tests from earlier
drafts whose coverage lived in test_mcp_client.py)
- ``tests/test_web_search.py`` +2 (oauth_user backend rejection +
static backend acceptance regression; updated to use the new
``server_auth_type`` in-memory accessor)