ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-27 06:14:48 -06:00
Code Issues Packages Projects Releases Wiki Activity

Address remaining blockers to entra id on behalf of flow for user impersonation to mcp servers (#706)

Browse Source 
* This is a collection of little snippits to resolve all the OBO flow problems required to get this talking to entra id for on behalf of user impersonating to protected mcp servers. we make sure turnstone checks these mcp servers on startup, and address some of microsofts opinionated implementations of oauth2/oidc and metadata provided by the identity provider.

* minor token timeout bugfix

---------

Co-authored-by: root <root@pow3rtools>
This commit is contained in:
metaclassing
2026-06-25 21:26:47 -05:00
committed by GitHub
parent 1c746af36f
commit 76e9d0a4f1
3 changed files with 226 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
turnstone/core/mcp_client.py
+162
View File
@@ -406,6 +406,12 @@ class PoolEntryState:
prompts: list[dict[str, Any]] | None = None
last_used: float = 0.0
in_flight: int = 0
# Access token this session's httpx client was connected with. The bearer is
# frozen into the client's STATIC headers at connect (_connect_one_pool), so
# when the stored token later refreshes we compare against this to detect a
# stale session and reconnect (rebind the current token) proactively —
# instead of replaying the stale bearer and eating a guaranteed upstream 401.
bound_token: str | None = None
auth_capture: _AuthCapture = field(default_factory=_AuthCapture)
# Set by the response hook when the carrier captures a 4xx; awaited
# by ``_dispatch_pool_with_entry``'s race against ``call_tool``.
@@ -1568,6 +1574,7 @@ class MCPClientManager:
# recovery path. Same ordering invariant covers resources/prompts
# (R16).
entry.session = session
entry.bound_token = access_token # remember the bearer this session carries
entry.last_used = time.monotonic()
self._user_pool_last_used[key] = entry.last_used
@@ -1585,6 +1592,137 @@ class MCPClientManager:
self._notify_user_prompt_listeners(user_id)
return entry
# -- pool priming ---------------------------------------------------------
async def _prime_user_server(
self, key: tuple[str, str], cfg: dict[str, Any], access_token: str
) -> int:
"""Proactively connect a pool entry so its catalog populates into
``get_tools(user_id)`` WITHOUT waiting for a tool dispatch.
``auth_type='oauth_user'`` tools are per-user and were previously
discovered ONLY lazily, on first dispatch (``_dispatch_pool_with_entry``
at the ``session is None`` branch). That creates a chicken-and-egg:
the model can't emit a call for a tool it can't see, but the tool
only appears after a call connects the pool — so the per-user
catalog stays empty and the server is stuck "connecting". This
connects at a known-good moment (OAuth consent completion), commits
the catalog and fires the per-user tool/resource/prompt listeners so
any live :class:`ChatSession` refreshes its tool list.
Idempotent: a no-op if a dispatch (or an earlier prime) already
established the session. MUST run on the mcp-loop; takes
``entry.open_lock`` exactly like dispatch so it can't race a
concurrent connect/eviction. Returns the discovered tool count.
"""
entry = await self._ensure_pool_entry(key)
async with entry.open_lock:
if entry.session is not None:
return len(entry.tools or [])
fresh = await self._connect_one_pool(
key,
cfg,
access_token,
auth_capture=entry.auth_capture,
auth_fired_event=entry.auth_fired_event,
)
return len(fresh.tools or [])
async def prime_user_server(
self,
*,
user_id: str,
server_name: str,
access_token: str,
server_row: dict[str, Any],
timeout: float = 20.0,
) -> bool:
"""Best-effort warm of a ``(user, server)`` pool so the user's tool
catalog populates immediately (e.g. right after OAuth consent).
Returns ``False`` (no-op) for non-``oauth_user`` servers, before the
mcp-loop is running, or with no token. NEVER raises: a prime failure
must not change the user-observable consent/redirect outcome, and the
lazy dispatch path remains the backstop. Schedules the connect onto
the mcp-loop (this is called from the request loop) and awaits it
without blocking that loop.
"""
if server_name not in self._oauth_user_server_names:
return False
loop = self._loop
if loop is None or not access_token:
return False
cfg = _pool_cfg_from_row(server_row)
key = (user_id, server_name)
try:
fut = asyncio.run_coroutine_threadsafe(
self._prime_user_server(key, cfg, access_token), loop
)
count = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout)
log.info(
"mcp pool primed user=%s server=%s tools=%d", user_id, server_name, count
)
return True
except Exception as exc:
log.warning(
"mcp pool prime failed user=%s server=%s: %s", user_id, server_name, exc
)
return False
def prime_user_pools(self, user_id: str) -> None:
"""Fire-and-forget: warm THIS user's consented ``oauth_user`` pools.
Called at ChatSession start so a per-user OAuth server's tools are
present automatically — no manual reconnect after a reboot/upgrade, and
no chicken-and-egg (the model can't dispatch a tool it can't see). Only
touches servers the user already has a stored token for; skips servers
already connected. Non-blocking: schedules onto the mcp-loop and returns
immediately. The per-user tool listeners (registered by ChatSession)
deliver the catalog to the live session when each prime completes.
"""
if not user_id or self._loop is None or not self._oauth_user_server_names:
return
if self._app_state is None or self._storage is None:
return
# run_coroutine_threadsafe keeps the task referenced by the loop while
# it runs, so no strong-ref bookkeeping is needed here.
asyncio.run_coroutine_threadsafe(self._prime_user_pools(user_id), self._loop)
async def _prime_user_pools(self, user_id: str) -> None:
"""Per-server body of :meth:`prime_user_pools` (runs on the mcp-loop)."""
from turnstone.core.mcp_oauth import get_user_access_token_classified
for server_name in list(self._oauth_user_server_names):
try:
key = (user_id, server_name)
entry = self._user_pool_entries.get(key)
if entry is not None and entry.session is not None:
continue # already connected — nothing to do
lookup = await get_user_access_token_classified(
app_state=self._app_state, user_id=user_id, server_name=server_name
)
if lookup.kind != "token" or not lookup.token:
continue # no usable token (not consented) — lazy paths handle it
server_row = await asyncio.to_thread(
self._storage.get_mcp_server_by_name, server_name
)
if not server_row:
continue
cfg = _pool_cfg_from_row(server_row)
await self._prime_user_server(key, cfg, lookup.token)
log.info(
"mcp pool auto-primed at session start user=%s server=%s",
user_id,
server_name,
)
except Exception:
log.debug(
"mcp pool auto-prime failed user=%s server=%s",
user_id,
server_name,
exc_info=True,
)
# -- pool eviction --------------------------------------------------------
async def _user_pool_eviction_loop(self) -> None:
@@ -4714,6 +4852,30 @@ class MCPClientManager:
entry.auth_capture.www_authenticate = None
entry.auth_fired_event.clear()
session = entry.session
if (
session is not None
and entry.bound_token is not None
and entry.bound_token != access_token
):
# The stored token refreshed since this session connected. The
# pooled httpx client's Authorization header is frozen at connect
# (_connect_one_pool), so a warm session would replay the now-
# stale bearer and eat a guaranteed upstream 401 before the
# auth_401 retry could heal it. Proactively reconnect to rebind
# the current token: _connect_one_pool pre-closes the old
# stack/streams, and entry.tools is retained (catalog intact), so
# this is a transparent in-place token rotation — attempt #0 now
# carries a valid bearer.
#
# ``bound_token is not None`` gate: only rotate when this session
# was established through ``_connect_one_pool`` (which records the
# bearer). A directly-injected session (None bind token) is left
# to the existing auth_401 retry path — and is the shape unit
# tests use, so this avoids spurious reconnects there.
log.debug(
"mcp_pool.token_rotated_reconnect user=%s server=%s", key[0], key[1]
)
session = None
if session is None:
# Lazy connect — also covers post-eviction recovery.
fresh = await self._connect_one_pool(
turnstone/core/mcp_oauth.py
+53 -8
View File
@@ -249,14 +249,29 @@ async def _fetch_as_metadata(
except OAuthSSRFError as exc:
raise MCPOAuthDiscoveryError(f"AS issuer URL rejected: {exc}") from exc
metadata_url = issuer.rstrip("/") + "/.well-known/oauth-authorization-server"
try:
resp = await http_client.get(metadata_url, timeout=_DEFAULT_HTTP_TIMEOUT)
except httpx.HTTPError as exc:
raise MCPOAuthDiscoveryError(f"AS metadata fetch failed: {exc}") from exc
if resp.status_code != 200:
raise MCPOAuthDiscoveryError(f"AS metadata returned HTTP {resp.status_code}")
# MCP auth permits OpenID Connect discovery as a fallback to RFC 8414.
# Major IdPs (notably Microsoft Entra) serve ONLY the OIDC document
# (.well-known/openid-configuration) and 404 the oauth-authorization-server
# path — refusing to support them would lock out the most common
# enterprise AS. Try RFC 8414 first (preferred), then OIDC.
base = issuer.rstrip("/")
metadata_candidates = (
base + "/.well-known/oauth-authorization-server",
base + "/.well-known/openid-configuration",
)
resp = None
last_status: int | None = None
for metadata_url in metadata_candidates:
try:
r = await http_client.get(metadata_url, timeout=_DEFAULT_HTTP_TIMEOUT)
except httpx.HTTPError as exc:
raise MCPOAuthDiscoveryError(f"AS metadata fetch failed: {exc}") from exc
if r.status_code == 200:
resp = r
break
last_status = r.status_code
if resp is None:
raise MCPOAuthDiscoveryError(f"AS metadata returned HTTP {last_status}")
if len(resp.content) > _MAX_DISCOVERY_BODY_BYTES:
raise MCPOAuthDiscoveryError("AS metadata response body exceeds size limit")
@@ -321,6 +336,14 @@ async def _fetch_as_metadata(
if not isinstance(code_methods_raw, list):
code_methods_raw = []
code_methods = tuple(str(m) for m in code_methods_raw)
if not code_methods:
# The OIDC discovery document does not require advertising
# code_challenge_methods_supported, and some IdPs (Entra) omit it
# despite fully supporting S256. Absence is not a denial — assume
# S256 (mandated by OAuth 2.1 / MCP auth) rather than locking the
# AS out. A NON-empty list missing S256 is still a hard refusal.
log.info("mcp_server.oauth.s256_assumed_absent_advertisement")
code_methods = ("S256",)
if "S256" not in code_methods:
raise MCPOAuthDiscoveryError(
"AS metadata does not advertise S256 PKCE — refusing to proceed"
@@ -2350,6 +2373,28 @@ async def _handle_mcp_oauth_callback_inner(request: Request) -> Response:
},
)
# Prime the per-user pool so the just-consented server's tools populate
# into this user's catalog immediately. Without this, oauth_user tools
# are discovered only lazily on first dispatch — but the agent can't
# emit a call for a tool it can't yet see, so the catalog stays empty
# and the server is stuck "connecting". Best-effort: a prime failure
# does not change consent success; lazy dispatch remains the backstop.
mcp_client = getattr(request.app.state, "mcp_client", None)
if mcp_client is not None and hasattr(mcp_client, "prime_user_server"):
try:
await mcp_client.prime_user_server(
user_id=user_id,
server_name=server_name,
access_token=access_token,
server_row=server_row,
)
except Exception:
log.debug(
"mcp_server.oauth.pool_prime_failed",
server_name=server_name,
exc_info=True,
)
# Phase 9 — clear any deferred-consent records for this (user,
# server) now that consent has completed. Best-effort: a storage
# failure here doesn't change the user-observable callback success;
turnstone/core/session.py
+11
View File
@@ -1203,6 +1203,17 @@ class ChatSession:
# for OTHER users do not wake this session.
self._mcp_prompt_cb = self._on_mcp_prompts_changed
mcp_client.add_prompt_listener(self._mcp_prompt_cb, user_id=self._mcp_user_id)
# Proactively warm this user's per-user OAuth (oauth_user) pools so
# their tools are present without a manual reconnect (e.g. after a
# reboot/upgrade, or right after consent). Fire-and-forget — the
# listeners registered just above deliver the catalog to this
# session once each prime completes. No-op for users with no
# consented oauth_user servers.
if self._mcp_user_id and hasattr(mcp_client, "prime_user_pools"):
try:
mcp_client.prime_user_pools(self._mcp_user_id)
except Exception:
pass
else:
self._tools = INTERACTIVE_TOOLS
self._task_tools = TASK_AGENT_TOOLS