ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: persistent stdio MCP sessions so server state survives across tool calls (#7080)

Browse source 
* Studio: persistent stdio MCP sessions so server state survives across tool calls

call_tool_sync spawned a fresh stdio subprocess per tool call
(keep_alive=False) and tore it down when the call returned, so any stateful
MCP server lost its state between calls: with @playwright/mcp,
browser_navigate opened the page in one subprocess and
browser_take_screenshot ran in a brand-new one, screenshotting about:blank.

Keep one connected client per (command, env) on a dedicated event-loop
thread and reuse it across calls:

- idle sessions are reaped after 5 minutes (in-flight calls excluded) and
  everything closes at exit, preserving the old design's no-orphans property
- a dead subprocess is detected via is_connected() and retried once on a
  fresh session; tool-level errors leave the session alone
- cancel and timeout semantics are unchanged, and a timed-out call does not
  tear the session down
- updating a server's endpoint/env/enabled state or deleting it closes its
  live session
- HTTP/SSE servers stay one-shot per call

* address review feedback

* fix stdio session cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review: per-thread MCP scope, close-during-connect and abort races

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review: unblock no-limit calls on close, drain borrowers before close, scope closes to url+env

* don't retry sessions closed by config changes, re-verify server row before caching, keep env secrets out of generation keys

* fail fast on connect errors and make the stdio key-lock wait cancellable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* quote MCP scope parts so IDs with colons can't collide

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* serialize per-session stdio calls, span one timeout budget across connect and call, hash urls in generation keys

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden persistent stdio MCP sessions: crash recovery, concurrency, scoping

- Evict a stdio session on any transport-level (non-ToolError) call failure and
  do not replay it, so a mid-call subprocess crash can no longer poison the scope.
  Never gate liveness on Client.is_connected() (it only reports that a session
  object exists, not that the subprocess is alive); add a version-adaptive
  dead-transport probe that works on fastmcp 3.0.2 and newer.
- Re-check closed/defunct/config and transport liveness after acquiring the call
  lock, and retire a session before releasing the lock, so a queued same-scope
  caller never reuses a session that another caller's timeout already retired.
- Force a ProactorEventLoop on Windows so the stdio transport can always spawn
  subprocesses regardless of the active event-loop policy.
- Scope stdio sessions per conversation: require thread_id to persist, and tag
  the fields so a session_id and a thread_id with the same value cannot collide.
  A session_id alone is project-wide, so it now falls back to a safe one-shot
  session instead of sharing browser/DB/REPL state across conversations.
- Forward thread_id on the Anthropic Messages path.
- Treat timeout=None as unlimited on connect and the key lock (was capped at 60s).
- Bound the session cache (default 32, override via
  UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS) with LRU eviction of idle sessions.
- Run config_check on cache hits, and log a redacted exe#digest label instead of
  the raw command so credentials in argv never reach the logs.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim the stdio MCP session cache on release and skip close-generation for HTTP servers

Two fixes from review of the persistent stdio session lifecycle:

- Re-enforce the session cap when a session goes idle. A concurrent burst of
  distinct-scope calls can overshoot the cap while every cached session is busy
  (insert-time eviction only reclaims idle sessions), and the overshoot used to
  persist until the 5-minute idle reaper. _release_stdio_session now trims the
  idle overshoot back within the cap, without ever evicting an in-flight call.
- close_stdio_sessions() now no-ops for a specific non-stdio (HTTP/SSE) url.
  Those transports are never cached as stdio sessions, so calling it on every
  HTTP server update or delete used to accrue an unbounded close-generation entry.

Both are covered by regression tests that fail before the change and pass after.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep the live stdio MCP session across a display-name rename

The edit dialog resends url, headers, and use_oauth unchanged whenever a
server is saved, so gating the tool-cache invalidation and stdio session
close on field presence dropped the persistent process on a plain rename
or any no-op edit. Gate on a real value change against the stored row so
only a genuine endpoint, auth, or enable change closes the session.

Regression tests: a rename that resends unchanged url/headers/oauth keeps
the session; a real command change still closes it.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten comments in the stdio MCP session lifecycle

Collapse a few verbose comments to fewer lines with the wording preserved,
and drop one that restated the clear_oauth_tokens_async docstring. Comments
only; no code change.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
Nilay 2026-07-14 14:58:43 +05:30 • committed by GitHub
commit 601155114d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1454 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/core/inference/inference.py
View file

@ -833,6 +833,7 @@ class InferenceBackend:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
):
@ -886,6 +887,7 @@ class InferenceBackend:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
)

2
studio/backend/core/inference/llama_cpp.py
View file

@ -8961,6 +8961,7 @@ class LlamaCppBackend:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
seed: Optional[int] = None,
disable_parallel_tool_use: bool = False,
@ -9983,6 +9984,7 @@ class LlamaCppBackend:
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)

739
studio/backend/core/inference/mcp_client.py
View file

@ -4,11 +4,16 @@
from __future__ import annotations
import asyncio
import atexit
import concurrent.futures
import hashlib
import json
import os
import shlex
import sys
import threading
import time
import uuid
from typing import Any, Optional
from loggers import get_logger
@ -115,6 +120,18 @@ def join_stdio_command(parts: list[str]) -> str:
return shlex.join(parts)
def _stdio_log_id(url: str) -> str:
"""A non-secret label for logs. stdio commands can embed credentials in argv
(e.g. ``npx server --token sk-...``), so never log the raw command; use the
executable basename plus a short digest of the full command instead."""
try:
parts = parse_stdio_command(url)
exe = os.path.basename(parts[0]) if parts else "<empty>"
except Exception: # noqa: BLE001
exe = "<invalid>"
return f"{exe}#{hashlib.sha256(url.encode()).hexdigest()[:12]}"
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (bypassing the
sandbox), so allowed only when the host is the user's own machine. On startup
@ -192,7 +209,6 @@ async def clear_oauth_tokens_async(url: str) -> None:
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
# Cleanup is best-effort; the row delete still wins.
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
@ -237,6 +253,502 @@ def _client(
return Client(transport_cls(url = url, headers = headers or None, auth = auth))
# Persistent stdio sessions: a stdio MCP server owns live state (a browser, a
# DB handle), so keep one connected client per (command, env, chat session) on
# a dedicated event-loop thread instead of respawning per call.
_STDIO_SESSION_IDLE_TTL = 300.0
_STDIO_SESSION_REAP_INTERVAL = 30.0
_STDIO_CONNECT_TIMEOUT = 60.0 # allows first-run `npx -y ...` package download
_STDIO_CLOSE_TIMEOUT = 10.0
_STDIO_WEDGE_MARGIN = 15.0
# Cap concurrent persistent sessions: each owns a subprocess + loop thread, and
# the scope includes a caller-supplied thread_id, so an unbounded cache is a
# resource-exhaustion surface. Overridable via env for large deployments.
try:
_STDIO_MAX_SESSIONS = max(1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32")))
except ValueError:
_STDIO_MAX_SESSIONS = 32
def _is_tool_error(exc: BaseException) -> bool:
"""A tool-level failure (the tool ran and errored) leaves the transport alive,
so the session is kept; fastmcp raises ToolError for these. Anything else from
call_tool is transport-level. Version-safe (fastmcp 3.0.2 has no dead probe)."""
try:
from fastmcp.exceptions import ToolError
except Exception: # noqa: BLE001
return False
return isinstance(exc, ToolError)
def _transport_dead(session) -> bool:
"""Best-effort, version-adaptive liveness probe for a cached stdio client.
``Client.is_connected()`` only checks a session object exists, not that the
subprocess is alive, so it is never used here. Returns True only when the
transport is positively gone; unknown returns False (the call surfaces it)."""
client = getattr(session, "client", None)
if client is None:
return True
transport = getattr(client, "transport", None)
probe = getattr(transport, "_is_session_dead", None)
if callable(probe):
try:
if probe():
return True
except Exception: # noqa: BLE001
pass
connect_task = getattr(transport, "_connect_task", None)
if connect_task is not None:
try:
if connect_task.done():
return True
except Exception: # noqa: BLE001
pass
return False
class _SessionWedged(Exception):
pass
class _SessionClosed(Exception):
"""The session was closed (server update/delete/shutdown) mid-call."""
def _abort_future(future) -> None:
# Let the cancelled coroutine unwind before its loop is stopped.
future.cancel()
try:
future.result(1.0)
except BaseException: # noqa: BLE001
pass
class _StdioSession:
def __init__(self, url: str, headers: Optional[dict]):
self.url = url
self.headers = headers
self.client = None
self.closed = threading.Event()
self.defunct = False # discarded; close once in_flight drains (see _retire)
self._close_lock = threading.Lock()
self.call_lock = threading.Lock() # serializes tool calls on this session
self.last_used = time.monotonic()
self.in_flight = 0 # guarded by _stdio_sessions_lock
# On Windows a bare new_event_loop() can be a SelectorEventLoop (if any
# component set that policy), which cannot spawn subprocesses natively;
# force a ProactorEventLoop so the stdio transport always works.
if sys.platform == "win32":
self.loop = asyncio.ProactorEventLoop()
else:
self.loop = asyncio.new_event_loop()
self._thread = threading.Thread(
target = self._run_loop, name = "mcp-stdio-session", daemon = True
)
self._thread.start()
def _run_loop(self) -> None:
asyncio.set_event_loop(self.loop)
try:
self.loop.run_forever()
finally:
self.loop.close()
def connect(self, timeout: Optional[float], cancel_event) -> None:
async def _open():
client = _client(self.url, self.headers)
await client.__aenter__()
# Publish on the loop thread with no await in between: if an abort
# races a just-completed connect, close() still sees the client and
# __aexit__s it instead of orphaning the subprocess.
self.client = client
return client
future = asyncio.run_coroutine_threadsafe(_open(), self.loop)
# timeout=None means unlimited (no connect deadline); a finite caller
# timeout still bounds connect by min(timeout, _STDIO_CONNECT_TIMEOUT).
window = None if timeout is None else min(timeout, _STDIO_CONNECT_TIMEOUT)
deadline = None if window is None else time.monotonic() + window
while True:
if cancel_event is not None and cancel_event.is_set():
_abort_future(future)
raise _MCPCancelled
try:
future.result(0.05)
return
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the connect itself failed fast; don't wait out the window
if deadline is not None and time.monotonic() >= deadline:
_abort_future(future)
raise asyncio.TimeoutError
def is_connected(self) -> bool:
client = self.client
if client is None:
return False
probe = getattr(client, "is_connected", None)
try:
return bool(probe()) if callable(probe) else True
except Exception:
return False
def run(self, coro, timeout: Optional[float]):
self.last_used = time.monotonic()
future = asyncio.run_coroutine_threadsafe(coro, self.loop)
# The coroutine enforces the tool timeout; the margin only catches a
# wedged loop. No deadline at all when the caller set none -- but poll
# so a session closed under us (server update/delete) can't hang the
# request thread forever on a stopped loop.
deadline = None if timeout is None else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN
try:
while True:
try:
return future.result(0.25)
except concurrent.futures.CancelledError:
# Only close() cancels in-flight tasks (in _shutdown).
raise _SessionClosed
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the call's own timeout; the session stays usable
if self.closed.is_set():
future.cancel()
raise _SessionClosed
if deadline is not None and time.monotonic() >= deadline:
future.cancel()
raise _SessionWedged
finally:
self.last_used = time.monotonic()
def close(self) -> None:
# Idempotent: a discard racing close_stdio_sessions() may close twice.
# Setting `closed` first also unblocks run() waiters (they poll it).
with self._close_lock:
if self.closed.is_set():
return
self.closed.set()
loop = getattr(self, "loop", None)
loop_alive = loop is not None and not loop.is_closed()
if loop_alive:
async def _shutdown() -> None:
# Runs on the loop thread, so it serializes with an aborted
# connect() that finished anyway and just published its client.
client, self.client = self.client, None
if client is not None:
await client.__aexit__(None, None, None)
# Cancel in-flight calls so they unwind before loop.stop
# (their run() waiters have already been released via `closed`).
for task in asyncio.all_tasks():
if task is not asyncio.current_task():
task.cancel()
try:
asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(_STDIO_CLOSE_TIMEOUT)
except Exception as exc: # noqa: BLE001
logger.warning(
"MCP stdio session close failed for %s: %s",
_stdio_log_id(getattr(self, "url", "")),
exc,
)
try:
loop.call_soon_threadsafe(loop.stop)
except RuntimeError:
pass
else:
self.client = None
thread = getattr(self, "_thread", None)
if thread is not None:
thread.join(timeout = 5.0)
_stdio_sessions: dict[tuple, _StdioSession] = {}
# Per-key locks so a slow connect/close never blocks unrelated servers; the
# global lock only guards the dicts.
class _StdioKeyLock:
"""A per-key lock that can be removed once nobody references it."""
def __init__(self) -> None:
self.lock = threading.Lock()
self.users = 0 # guarded by _stdio_sessions_lock
_stdio_key_locks: dict[tuple, _StdioKeyLock] = {}
_stdio_sessions_lock = threading.Lock()
_stdio_reaper_started = False
# close_stdio_sessions() can only close sessions already published in
# _stdio_sessions; one still inside connect() would be missed and cached
# stale. Bump a generation on every close so that connect discards its
# session instead of publishing it. Guarded by _stdio_sessions_lock.
_stdio_close_all_gen = 0
_stdio_url_close_gen: dict[str, int] = {}
_stdio_cfg_close_gen: dict[tuple, int] = {}
# close_stdio_sessions(url): match any env for that command.
_ANY_HEADERS = object()
def _headers_key(headers: Optional[dict]) -> tuple:
return tuple(sorted((headers or {}).items()))
def _url_close_key(url: str) -> str:
# Commands/URLs (token args, embedded credentials) and env values can hold
# secrets and these maps are never pruned; key by digest so closed/edited
# configs don't retain them in memory forever.
return hashlib.sha256(url.encode()).hexdigest()
def _cfg_close_key(url: str, headers: Optional[dict]) -> str:
return hashlib.sha256(repr((url, _headers_key(headers))).encode()).hexdigest()
def _stdio_close_generation(url: str, headers: Optional[dict]) -> tuple[int, int, int]:
return (
_stdio_close_all_gen,
_stdio_url_close_gen.get(_url_close_key(url), 0),
_stdio_cfg_close_gen.get(_cfg_close_key(url, headers), 0),
)
def _session_key(url: str, headers: Optional[dict], scope: Optional[str]) -> tuple:
return (url, _headers_key(headers), scope or "")
def _checkout_stdio_session(key: tuple) -> Optional[_StdioSession]:
session = _stdio_sessions.get(key)
if session is not None and session.is_connected():
session.last_used = time.monotonic()
session.in_flight += 1
return session
return None
def _borrow_stdio_key_lock(key: tuple) -> _StdioKeyLock:
"""Return a stable per-key lock while a caller waits for/connects it."""
key_lock = _stdio_key_locks.setdefault(key, _StdioKeyLock())
key_lock.users += 1
return key_lock
def _discard_stdio_key_lock(key: tuple) -> None:
key_lock = _stdio_key_locks.get(key)
if key_lock is not None and key_lock.users == 0 and key not in _stdio_sessions:
_stdio_key_locks.pop(key, None)
def _return_stdio_key_lock(key: tuple, key_lock: _StdioKeyLock) -> None:
with _stdio_sessions_lock:
key_lock.users -= 1
_discard_stdio_key_lock(key)
def _get_stdio_session(
url: str, headers: Optional[dict], scope: Optional[str], deadline, cancel_event, config_check
) -> _StdioSession:
"""``deadline`` is the caller's absolute monotonic budget (None = no limit):
the key-lock wait and the connect share it, so a slow startup can't stack
full timeout windows (see _call_stdio_tool)."""
global _stdio_reaper_started
key = _session_key(url, headers, scope)
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
key_lock = _borrow_stdio_key_lock(key)
try:
# Poll the acquire with connect()'s deadline/cancel semantics: a second
# same-scope call must not block uncancellably behind another caller's
# slow startup (e.g. a first-run npx download).
remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
# timeout=None means no key-lock deadline (only cancel unblocks it).
window = None if remaining is None else min(remaining, _STDIO_CONNECT_TIMEOUT)
lock_deadline = None if window is None else time.monotonic() + window
while not key_lock.lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
if lock_deadline is not None and time.monotonic() >= lock_deadline:
raise asyncio.TimeoutError
try:
stale = None
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
if key in _stdio_sessions:
stale = _stdio_sessions.pop(key)
generation = _stdio_close_generation(url, headers)
if stale is not None:
_retire_stdio_session(stale)
session = _StdioSession(url, headers)
try:
session.connect(
None if deadline is None else max(0.0, deadline - time.monotonic()),
cancel_event,
)
except Exception:
session.close()
raise
# A caller can read the server row, then lose to an update/delete whose close ran
# before our generation snapshot. Re-verify the row after connect; the generation check
# below covers a close landing between this check and publish.
if config_check is not None:
try:
current = bool(config_check())
except Exception: # noqa: BLE001
current = False
if not current:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
evicted: list = []
with _stdio_sessions_lock:
closed_while_connecting = _stdio_close_generation(url, headers) != generation
if not closed_while_connecting:
session.in_flight = 1
evicted = _evict_stdio_lru_locked() # bound the cache (LRU idle)
_stdio_sessions[key] = session
if not _stdio_reaper_started:
_stdio_reaper_started = True
threading.Thread(
target = _stdio_session_reaper, name = "mcp-stdio-reaper", daemon = True
).start()
atexit.register(close_stdio_sessions)
for victim in evicted:
logger.info("Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url))
victim.close()
if closed_while_connecting:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
return session
finally:
key_lock.lock.release()
finally:
_return_stdio_key_lock(key, key_lock)
def _release_stdio_session(session: _StdioSession) -> None:
victims: list = []
with _stdio_sessions_lock:
session.in_flight = max(0, session.in_flight - 1)
session.last_used = time.monotonic()
close_now = session.defunct and session.in_flight == 0
# Re-enforce the cap once a burst's sessions go idle. Insert-time eviction
# only trims idle sessions, so it can overshoot while every cached session
# is busy; reclaim that overshoot here instead of waiting for the idle
# reaper. Never evict the session we just used (its last_used is newest).
while len(_stdio_sessions) > _STDIO_MAX_SESSIONS:
idle = [
(s.last_used, k)
for k, s in _stdio_sessions.items()
if s.in_flight == 0 and s is not session
]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
if close_now:
session.close()
for victim in victims:
victim.close()
def _retire_stdio_session(session: _StdioSession) -> None:
"""Close a discarded session, but only once no other borrower is mid-call
on it -- overlapping same-scope calls share one client, and one call's
timeout must not kill another's in-flight request. The last borrower's
_release_stdio_session() performs the deferred close."""
with _stdio_sessions_lock:
session.defunct = True
busy = session.in_flight > 0
if not busy:
session.close()
def _drop_stdio_session(key: tuple, session: _StdioSession) -> None:
with _stdio_sessions_lock:
if _stdio_sessions.get(key) is session:
_stdio_sessions.pop(key)
_discard_stdio_key_lock(key)
_retire_stdio_session(session)
def _evict_stdio_lru_locked() -> list:
"""Caller holds _stdio_sessions_lock. Evict least-recently-used *idle*
sessions until the cache is under the cap. Returns the evicted sessions so
the caller can close them OUTSIDE the lock. If every session is busy the
cache may transiently overshoot rather than kill an in-flight call."""
victims: list = []
while len(_stdio_sessions) >= _STDIO_MAX_SESSIONS:
idle = [(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
return victims
def close_stdio_sessions(url: Optional[str] = None, headers = _ANY_HEADERS) -> None:
"""Close persistent stdio sessions: all of them (``url`` None), every env
for one command (``headers`` omitted), or one server config (url + headers).
Two server rows can share a command with different envs; editing one must
not kill the other's live state, so the routes pass the edited row's env."""
global _stdio_close_all_gen
# HTTP/SSE servers are never cached as stdio sessions, so a specific non-stdio
# url has nothing to close and must not accrue a close-generation entry.
if url is not None and not is_stdio(url):
return
hk = None if headers is _ANY_HEADERS else _headers_key(headers)
with _stdio_sessions_lock:
if url is None:
_stdio_close_all_gen += 1
elif hk is None:
uk = _url_close_key(url)
_stdio_url_close_gen[uk] = _stdio_url_close_gen.get(uk, 0) + 1
else:
cfg = _cfg_close_key(url, headers)
_stdio_cfg_close_gen[cfg] = _stdio_cfg_close_gen.get(cfg, 0) + 1
keys = [
k
for k in _stdio_sessions
if (url is None or k[0] == url) and (hk is None or k[1] == hk)
]
sessions = [_stdio_sessions.pop(k) for k in keys]
for key in keys:
_discard_stdio_key_lock(key)
for session in sessions:
session.close()
def _reap_idle_stdio_sessions(now: Optional[float] = None) -> None:
now = time.monotonic() if now is None else now
with _stdio_sessions_lock:
expired = [
key
for key, session in _stdio_sessions.items()
if session.in_flight == 0 and now - session.last_used >= _STDIO_SESSION_IDLE_TTL
]
sessions = [_stdio_sessions.pop(key) for key in expired]
for key in expired:
_discard_stdio_key_lock(key)
for session in sessions:
logger.info("Closing idle stdio MCP session: %s", _stdio_log_id(session.url))
session.close()
def _stdio_session_reaper() -> None:
while True:
time.sleep(_STDIO_SESSION_REAP_INTERVAL)
try:
_reap_idle_stdio_sessions()
except Exception as exc: # noqa: BLE001
logger.debug("stdio session reaper iteration failed: %s", exc)
async def list_tools_async(
url: str,
headers: Optional[dict] = None,
@ -251,11 +763,10 @@ async def list_tools_async(
return await asyncio.wait_for(_fetch(), timeout = timeout)
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools()
# probes a server only on a cache miss, keeping MCP discovery off the chat
# send's critical path -- tool schemas are stable within a session. The
# /refresh route warms it; a URL/header/OAuth change or a delete evicts it.
# Successful probes are cached indefinitely.
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() probes a server only
# on a cache miss, keeping MCP discovery off the chat send's critical path -- tool schemas are
# stable within a session. The /refresh route warms it; a URL/header/OAuth change or a delete
# evicts it. Successful probes are cached indefinitely.
_tool_cache: dict[str, list[dict]] = {}
# server_id -> monotonic time before which a failed server must not be
@ -343,6 +854,165 @@ def _flatten_result(result: Any) -> str:
return body
async def _race_tool_call(call_coro, timeout: Optional[float], cancel_event) -> Any:
"""Await ``call_coro`` under ``timeout``, polling ``cancel_event`` so a
/cancel POST interrupts even mid-network-read."""
async def _watch_cancel() -> None:
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
if cancel_event is not None and cancel_event.is_set():
call_coro.close()
raise _MCPCancelled
call_task = asyncio.create_task(call_coro)
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
def _call_stdio_tool(
url: str,
headers: Optional[dict],
name: str,
args: dict,
timeout,
cancel_event,
scope: Optional[str],
config_check,
) -> Any:
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
# One deadline covers the key-lock wait, connect, call-lock wait, and the
# call itself, matching the one-shot/HTTP paths where the timeout wrapped
# connect plus call in a single window.
deadline = None if timeout is None else time.monotonic() + timeout
def _remaining() -> Optional[float]:
return None if deadline is None else max(0.0, deadline - time.monotonic())
# Callers without a Studio session id must retain the former one-shot
# behavior: no browser/cookie/tool state can leak into another request.
# Use an ephemeral key (and close it below) rather than the shared empty
# scope that the persistent-session cache used previously.
def _config_ok() -> bool:
if config_check is None:
return True
try:
return bool(config_check())
except Exception: # noqa: BLE001
return False
ephemeral = not scope
if ephemeral:
scope = f"request-{uuid.uuid4().hex}"
key = _session_key(url, headers, scope)
# attempt 0 may find the cached session stale/dead *before* dispatch and
# reconnect once (safe); attempt 1 is a freshly connected session.
for attempt in (0, 1):
session = _get_stdio_session(url, headers, scope, deadline, cancel_event, config_check)
try:
# Serialize calls per session: overlapping same-scope calls must
# not interleave operations on one stateful server (browser, REPL).
while not session.call_lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
rem = _remaining()
if rem is not None and rem <= 0:
raise asyncio.TimeoutError
except BaseException:
# Never touched the transport: keep the session for its borrower.
_release_stdio_session(session)
if ephemeral:
_drop_stdio_session(key, session)
raise
discard_session = ephemeral
retry = False
try:
# We may have waited on the call lock while another caller's timeout retired this
# session, a server update/delete invalidated it, or a reused subprocess died. Re-check
# all three before dispatch so we never run on a retired/dead client or a stale config.
if session.closed.is_set():
# Intentional close (server update/delete/shutdown): don't retry on stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif session.defunct:
# A concurrent same-scope caller's timeout retired this session;
# move to a fresh one instead of reusing the retired client.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server session was retired during the call")
elif not _config_ok():
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif _transport_dead(session):
# Dead BEFORE dispatch: no request was sent, so reconnect + retry.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server connection is not available")
else:
rem = _remaining()
coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event)
return session.run(coro, rem)
except (_MCPCancelled, asyncio.TimeoutError):
# _race_tool_call cancels the pending call but cancellation is
# cooperative. Never return this client to the cache while the
# timed-out/cancelled operation might still run on its transport.
discard_session = True
raise
except _SessionWedged:
discard_session = True
raise asyncio.TimeoutError
except _SessionClosed:
# close_stdio_sessions() shut this session mid-call (server
# update/delete/shutdown); don't retry on the stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
except Exception as exc:
if session.closed.is_set():
# An intentional close (server update/delete) can surface as a plain transport
# error or AttributeError instead of _SessionClosed; don't mistake it for a crash.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
# ToolError leaves the transport alive -> keep the session so its state
# survives. Any other exception is transport-level (dead subprocess,
# broken pipe): evict so it can't poison the scope, but DO NOT replay
# (the tool may already have run); the next call opens a fresh session.
if not _is_tool_error(exc):
discard_session = True
raise
finally:
# Set defunct + remove from the cache BEFORE releasing the call lock,
# so a queued same-scope borrower observes the retirement and opens a
# fresh session instead of reusing this one.
_release_stdio_session(session)
if discard_session:
_drop_stdio_session(key, session)
session.call_lock.release()
if not retry:
break
raise RuntimeError("unreachable")
def call_tool_sync(
url: str,
headers: Optional[dict],
@ -351,58 +1021,35 @@ def call_tool_sync(
timeout: Optional[float] = 300.0,
use_oauth: bool = False,
cancel_event = None,
scope: Optional[str] = None,
config_check = None,
) -> str:
"""Synchronously call an MCP tool.
"""Synchronously call an MCP tool. stdio servers reuse a persistent session
keyed by (command, env, scope) only when ``scope`` is provided; calls
without one stay one-shot. HTTP servers always stay one-shot.
``cancel_event`` (threading.Event) cancels the in-flight call when set.
``config_check`` (callable -> bool) re-validates the caller's server config
before a fresh stdio session is cached; False fails the call."""
``cancel_event``: optional ``threading.Event``. When set, the in-flight call is
cancelled and a cancellation Error returned. Polled alongside the tool call via
``asyncio.wait`` so a /cancel POST interrupts even mid-network-read.
"""
async def _call() -> Any:
async def _one_shot() -> Any:
async with _client(url, headers, use_oauth) as client:
# raise_on_error=False lets an is_error result (which may still carry
# image content) reach _flatten_result instead of FastMCP raising ToolError
# and dropping the images. Transport failures still raise (handled below).
return await client.call_tool(name, args, raise_on_error = False)
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches routes/inference.py's cancel watcher cadence.
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
async def _race() -> Any:
# Check cancellation before spawning the call task so a pre-set event
# short-circuits before opening the transport / HTTP connection.
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
try:
result = asyncio.run(_race())
if is_stdio(url):
result = _call_stdio_tool(
url, headers, name, args, timeout, cancel_event, scope, config_check
)
else:
result = asyncio.run(_race_tool_call(_one_shot(), timeout, cancel_event))
except _MCPCancelled:
return f"Error: MCP tool '{name}' cancelled"
except asyncio.TimeoutError:
return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
suffix = f" after {timeout:g}s" if timeout is not None else ""
return f"Error: MCP tool '{name}' timed out{suffix}"
except Exception as exc:
logger.exception("MCP call_tool failed for %s: %s", name, exc)
return f"Error: MCP tool '{name}' failed: {exc}"

2
studio/backend/core/inference/orchestrator.py
View file

@ -1293,6 +1293,7 @@ class InferenceOrchestrator:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
@ -1359,6 +1360,7 @@ class InferenceOrchestrator:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,

2
studio/backend/core/inference/safetensors_agentic.py
View file

@ -424,6 +424,7 @@ def run_safetensors_tool_loop(
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
@ -1116,6 +1117,7 @@ def run_safetensors_tool_loop(
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)

35
studio/backend/core/inference/tools.py
View file

@ -20,6 +20,7 @@ import subprocess
import sys
import tempfile
import threading
import urllib.parse
import urllib.request
from core.inference.mcp_client import (
@ -971,6 +972,7 @@ def execute_tool(
cancel_event = None,
timeout: int | None = _TIMEOUT_UNSET,
session_id: str | None = None,
thread_id: str | None = None,
rag_scope: dict | None = None,
disable_sandbox: bool = False,
) -> str:
@ -978,6 +980,8 @@ def execute_tool(
``timeout``: int seconds, ``None`` = no limit, unset = ``_EXEC_TIMEOUT``.
``session_id``: optional ID for per-conversation sandbox isolation.
``thread_id``: optional conversation ID; scopes stateful MCP stdio sessions
per thread (session_id alone can be shared project-wide).
``rag_scope``: hidden per-request RAG context the model never sees; consumed
by ``search_knowledge_base``.
``disable_sandbox``: Bypass Permissions; run python/terminal without the
@ -1002,14 +1006,41 @@ def execute_tool(
return f"Error: MCP server '{server_id}' is disabled"
if is_stdio(server["url"]) and not stdio_mcp_enabled():
return f"Error: stdio MCP server '{server_id}' is disabled on this host"
# Persist a stateful stdio session only per conversation (thread_id).
# session_id is the project-wide sandbox id, so scoping by it alone leaks
# browser/DB/REPL state across conversations; fall back to one-shot. Tag +
# percent-quote the parts so ids can't collide or ":" merge conversations.
if thread_id:
mcp_scope = "s={}:t={}".format(
urllib.parse.quote(session_id or "", safe = ""),
urllib.parse.quote(thread_id, safe = ""),
)
else:
mcp_scope = None
headers = parse_server_headers(server)
url = server["url"]
def _config_current() -> bool:
# Re-read before a stdio session is cached: this call may have read
# the row just before an update/delete closed its sessions.
row = mcp_servers_db.get_server(server_id)
return (
row is not None
and bool(row.get("is_enabled"))
and row.get("url") == url
and parse_server_headers(row) == headers
)
return call_tool_sync(
url = server["url"],
headers = parse_server_headers(server),
url = url,
headers = headers,
name = tool_name,
args = arguments,
timeout = effective_timeout,
use_oauth = bool(server.get("use_oauth")),
cancel_event = cancel_event,
scope = mcp_scope,
config_check = _config_current,
)
if name == "web_search":
return _web_search(

8
studio/backend/models/inference.py
View file

@ -815,6 +815,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
rag_scope: Optional[dict] = Field(
None,
description = (
@ -1682,6 +1686,10 @@ class AnthropicMessagesRequest(BaseModel):
enable_tools: Optional[bool] = None
enabled_tools: Optional[list[str]] = None
session_id: Optional[str] = None
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
cancel_id: Optional[str] = None
bypass_permissions: Optional[bool] = Field(
False,

3
studio/backend/routes/inference.py
View file

@ -6929,6 +6929,7 @@ async def openai_chat_completions(
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
thread_id = payload.thread_id,
rag_scope = payload.rag_scope,
disable_parallel_tool_use = payload.parallel_tool_calls is False,
# Bypass Permissions takes precedence over the confirm gate:
@ -8225,6 +8226,7 @@ async def openai_chat_completions(
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
thread_id = payload.thread_id,
rag_scope = payload.rag_scope,
# Bypass Permissions takes precedence over the confirm gate:
# never prompt while bypassing.
@ -11755,6 +11757,7 @@ async def anthropic_messages(
nudge_tool_calls = payload.nudge_tool_calls,
tool_call_timeout = 300,
session_id = payload.session_id,
thread_id = payload.thread_id,
# Anthropic passthrough has no rag_scope field (RAG is local-only).
rag_scope = getattr(payload, "rag_scope", None),
disable_parallel_tool_use = _disable_parallel,

15
studio/backend/routes/mcp_servers.py
View file

@ -1,6 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import json
import uuid
from urllib.parse import urlparse
@ -13,6 +14,7 @@ from core.inference.mcp_client import (
TOOL_CACHE_INVALIDATING_FIELDS,
cache_tools,
clear_oauth_tokens_async,
close_stdio_sessions,
invalidate_tool_cache,
is_stdio,
list_tools_async,
@ -206,11 +208,15 @@ async def update_mcp_server(
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
# A new endpoint/auth makes cached tools wrong and disabling makes them
# unreachable, so drop them and let the next send re-probe; a rename
# leaves them valid.
if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS:
# A new endpoint/auth makes cached tools wrong and disabling makes them unreachable, so drop
# them and let the next send re-probe; a rename leaves them valid. Live stdio sessions for the
# old endpoint close too. Gate on a real value change, not mere presence: the edit dialog
# resends url/headers/oauth unchanged on a rename, which must not drop the session.
if any(changes[k] != old.get(k) for k in changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS):
invalidate_tool_cache(server_id)
# Narrow to this row's env: another server row sharing the command but
# with a different env keeps its live sessions.
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
return _row_to_response(mcp_servers_db.get_server(server_id))
@ -223,6 +229,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.delete_server(server_id)
invalidate_tool_cache(server_id)
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)

2
studio/backend/tests/test_bypass_permissions.py
View file

@ -257,6 +257,7 @@ def test_loop_forwards_disable_sandbox_and_does_not_gate():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):
@ -290,6 +291,7 @@ def test_loop_bypass_overrides_confirm_for_direct_callers():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

62
studio/backend/tests/test_mcp_servers.py
View file

@ -797,6 +797,68 @@ def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch):
assert mcp_client.get_cached_tools("s1") == cached
def test_update_rename_keeps_stdio_session(tmp_path, monkeypatch):
"""The edit dialog resends url/headers/oauth unchanged on a rename, so gating
the close on field presence would drop the live stdio session. Only a real
endpoint/auth change may close it."""
import asyncio
import json
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
headers_json = json.dumps({"API_KEY": "x"}),
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(
display_name = "B",
url = "npx demo-server",
headers = {"API_KEY": "x"},
use_oauth = False,
),
current_subject = "u",
)
)
assert closed == []
assert mcp_servers_db.get_server("s1")["display_name"] == "B"
def test_update_stdio_command_change_closes_session(tmp_path, monkeypatch):
"""A real command change must still close the old stdio session."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx other-server"), current_subject = "u"
)
)
assert len(closed) == 1
def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch):
"""Disabling a server must drop its cached tools, not leave them unread."""
import asyncio

629
studio/backend/tests/test_mcp_stdio_sessions.py Normal file
View file

@ -0,0 +1,629 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference import mcp_client
from core.inference.mcp_client import call_tool_sync, close_stdio_sessions
STDIO_URL = "npx fake-stateful-server"
HTTP_URL = "https://mcp.example.test/mcp"
def _result(text: str) -> SimpleNamespace:
return SimpleNamespace(
content = [SimpleNamespace(type = "text", text = text)],
is_error = False,
structured_content = None,
)
class FakeClient:
instances: list["FakeClient"] = []
def __init__(self, url: str):
self.url = url
self.entered = 0
self.exited = 0
self.calls: list[tuple[str, dict]] = []
self.connected = False
self.fail_next = False
self.call_delay = 0.0
# Models a dead stdio transport: real Client.is_connected() stays True
# after the subprocess dies, so liveness is probed via the transport.
self.dead = False
self.transport = SimpleNamespace(_is_session_dead = lambda: self.dead)
FakeClient.instances.append(self)
async def __aenter__(self):
self.entered += 1
self.connected = True
return self
async def __aexit__(self, *exc):
self.exited += 1
self.connected = False
def is_connected(self) -> bool:
return self.connected
async def call_tool(self, name: str, args: dict):
if self.call_delay:
await asyncio.sleep(self.call_delay)
if self.fail_next:
self.fail_next = False
self.connected = False
raise RuntimeError("transport closed")
self.calls.append((name, args))
return _result(f"call-{len(self.calls)}")
@pytest.fixture
def fake_clients(monkeypatch):
FakeClient.instances = []
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: FakeClient(url)
)
yield FakeClient.instances
close_stdio_sessions()
def test_stdio_call_without_scope_is_one_shot(fake_clients):
r1 = call_tool_sync(STDIO_URL, None, "browser_navigate", {"url": "https://x.test"})
r2 = call_tool_sync(STDIO_URL, None, "browser_take_screenshot", {})
assert r1 == "call-1"
assert r2 == "call-1"
assert len(fake_clients) == 2
assert all(client.entered == 1 and client.exited == 1 for client in fake_clients)
def test_stdio_sessions_keyed_by_url_and_env(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {})
call_tool_sync("npx other-server", None, "t", {})
call_tool_sync(STDIO_URL, {"ENV_VAR": "1"}, "t", {})
assert len(fake_clients) == 3
def test_stdio_sessions_scoped_per_chat(fake_clients):
# Two conversations must not share one stateful server process.
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-a")
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-b")
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat-a")
assert len(fake_clients) == 2
assert fake_clients[0].calls and len(fake_clients[0].calls) == 2
def test_dead_stdio_session_recovers(fake_clients):
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
# Subprocess dies between calls: the dead transport is detected before the
# next dispatch, so the call reconnects on a fresh session instead of failing.
fake_clients[0].dead = True
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
assert fake_clients[0].exited == 1
def test_tool_error_does_not_recycle_session(fake_clients, monkeypatch):
from fastmcp.exceptions import ToolError
class ToolFailure(FakeClient):
async def call_tool(self, name, args):
if name == "boom":
raise ToolError("tool exploded") # tool-level: session stays connected
return await super().call_tool(name, args)
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: ToolFailure(url)
)
assert call_tool_sync(STDIO_URL, None, "boom", {}, scope = "chat").startswith("Error: MCP tool")
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 1
def test_http_stays_one_shot(fake_clients):
call_tool_sync(HTTP_URL, None, "t", {})
call_tool_sync(HTTP_URL, None, "t", {})
assert len(fake_clients) == 2
assert all(c.entered == 1 and c.exited == 1 for c in fake_clients)
def test_timeout_discards_stdio_session(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
fake_clients[0].call_delay = 0.5
out = call_tool_sync(
STDIO_URL,
None,
"slow",
{},
timeout = 0.05,
cancel_event = threading.Event(),
scope = "chat",
)
assert "timed out" in out
assert fake_clients[0].exited == 1
assert key not in mcp_client._stdio_key_locks
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
def test_no_timeout_allows_long_call(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
fake_clients[0].call_delay = 0.2
# timeout=None means no deadline: the call must not be treated as wedged.
assert call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat") == "call-2"
def test_connect_races_cancel_event(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(5.0)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
ev = threading.Event()
threading.Timer(0.1, ev.set).start()
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev)
assert out == "Error: MCP tool 't' cancelled"
assert time.monotonic() - start < 3.0
assert mcp_client._stdio_sessions == {}
def test_connect_respects_caller_timeout(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(5.0)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.2)
assert "timed out" in out
assert time.monotonic() - start < 3.0
assert mcp_client._stdio_sessions == {}
def test_connect_failure_timeout_surfaces_immediately(fake_clients, monkeypatch):
class InitTimeout(FakeClient):
async def __aenter__(self):
raise asyncio.TimeoutError # e.g. fastmcp's own init timeout
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: InitTimeout(url)
)
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 30.0)
assert "timed out" in out
# Must fail fast, not wait out the 30s/60s connect window.
assert time.monotonic() - start < 5.0
assert mcp_client._stdio_sessions == {}
def test_key_lock_wait_honors_cancel_and_timeout(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(1.5)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
first = threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
first.start()
key = mcp_client._session_key(STDIO_URL, None, "chat")
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
key_lock = mcp_client._stdio_key_locks.get(key)
if key_lock is not None and key_lock.lock.locked():
break
time.sleep(0.01)
# Second same-scope call is stuck behind the first slow connect: Stop must
# interrupt the key-lock wait, and a short tool timeout must bound it.
ev = threading.Event()
threading.Timer(0.2, ev.set).start()
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev, scope = "chat")
assert out == "Error: MCP tool 't' cancelled"
assert time.monotonic() - start < 1.0
start = time.monotonic()
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.2, scope = "chat")
assert "timed out" in out
assert time.monotonic() - start < 1.0
first.join(10.0)
assert not first.is_alive()
def test_cancel_pre_set_spawns_nothing(fake_clients):
ev = threading.Event()
ev.set()
out = call_tool_sync(STDIO_URL, None, "t", {}, cancel_event = ev)
assert out == "Error: MCP tool 't' cancelled"
assert fake_clients == []
def test_idle_reap_closes_session(fake_clients, monkeypatch):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
assert key in mcp_client._stdio_key_locks
monkeypatch.setattr(mcp_client, "_STDIO_SESSION_IDLE_TTL", 0.0)
mcp_client._reap_idle_stdio_sessions()
assert fake_clients[0].exited == 1
assert mcp_client._stdio_sessions == {}
assert key not in mcp_client._stdio_key_locks
# Next call transparently opens a fresh session.
assert call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat") == "call-1"
assert len(fake_clients) == 2
def test_reap_skips_in_flight_session(fake_clients, monkeypatch):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
monkeypatch.setattr(mcp_client, "_STDIO_SESSION_IDLE_TTL", 0.0)
session = next(iter(mcp_client._stdio_sessions.values()))
with mcp_client._stdio_sessions_lock:
session.in_flight = 1
try:
mcp_client._reap_idle_stdio_sessions()
assert fake_clients[0].exited == 0
finally:
with mcp_client._stdio_sessions_lock:
session.in_flight = 0
def test_close_during_connect_is_not_cached(fake_clients, monkeypatch):
class SlowStart(FakeClient):
async def __aenter__(self):
await asyncio.sleep(0.5)
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowStart(url))
results: list[str] = []
worker = threading.Thread(
target = lambda: results.append(call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
)
worker.start()
deadline = time.monotonic() + 5.0
while not fake_clients and time.monotonic() < deadline:
time.sleep(0.01)
assert fake_clients # connect is in progress
# Server deleted/updated mid-connect: the session must not be cached after.
close_stdio_sessions(STDIO_URL)
worker.join(10.0)
assert results and results[0].startswith("Error: MCP tool 't' failed")
assert mcp_client._stdio_sessions == {}
assert fake_clients[0].exited == 1
def test_connect_abort_race_still_closes_client(fake_clients, monkeypatch):
class WinsRace(FakeClient):
async def __aenter__(self):
try:
await asyncio.sleep(5.0)
except asyncio.CancelledError:
pass # connect finishes just as the abort lands
return await super().__aenter__()
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: WinsRace(url))
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.1)
assert "timed out" in out
assert fake_clients[0].entered == 1
assert fake_clients[0].exited == 1 # no orphaned subprocess
assert mcp_client._stdio_sessions == {}
def test_close_unblocks_no_limit_call(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 30.0
results: list[str] = []
worker = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
worker.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# Server deleted while a no-limit call is in flight: the request thread
# must not hang forever on the stopped session loop.
close_stdio_sessions(STDIO_URL)
worker.join(5.0)
assert not worker.is_alive()
assert results and results[0].startswith("Error: MCP tool 'slow' failed")
def test_lock_wait_timeout_spares_the_borrowed_session(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 1.0
results: list[str] = []
slow = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
slow.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# A second same-scope call times out waiting for the call lock; it never
# touched the transport, so the shared session must stay alive and cached.
out = call_tool_sync(STDIO_URL, None, "fast", {}, timeout = 0.05, scope = "chat")
assert "timed out" in out
assert fake_clients[0].exited == 0
slow.join(10.0)
assert results == ["call-2"]
assert fake_clients[0].exited == 0
assert len(mcp_client._stdio_sessions) == 1
def test_stale_session_close_deferred_until_borrower_drains(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
fake_clients[0].call_delay = 0.8
results: list[str] = []
slow = threading.Thread(
target = lambda: results.append(
call_tool_sync(STDIO_URL, None, "slow", {}, timeout = None, scope = "chat")
)
)
slow.start()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
with mcp_client._stdio_sessions_lock:
if session.in_flight >= 1:
break
time.sleep(0.01)
# The subprocess "dies" mid-call: a new caller replaces the stale session,
# but its close must wait for the slow borrower instead of killing its call.
fake_clients[0].connected = False
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
assert out == "call-1"
assert len(fake_clients) == 2
assert fake_clients[0].exited == 0
slow.join(10.0)
assert results == ["call-2"]
assert fake_clients[0].exited == 1 # last borrower performed the deferred close
assert len(mcp_client._stdio_sessions) == 1
def test_error_on_closed_session_does_not_retry(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
session = next(iter(mcp_client._stdio_sessions.values()))
# A close can surface at the borrower as a plain transport error instead
# of _SessionClosed; that must not be treated as a crash and retried.
fake_clients[0].fail_next = True
session.closed.set()
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
assert out == "Error: MCP tool 't' failed: MCP server was updated or removed during the call"
assert len(fake_clients) == 1 # no respawn for the removed config
def test_config_check_blocks_stale_publish(fake_clients):
# Simulates a caller that read the server row before an update/delete:
# the row re-check runs after connect and must block caching.
out = call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat", config_check = lambda: False)
assert out.startswith("Error: MCP tool 't' failed")
assert mcp_client._stdio_sessions == {}
assert fake_clients[0].exited == 1
def test_close_generation_keys_hold_no_secrets(fake_clients):
secret_url = "npx server --token sk-url-secret"
close_stdio_sessions(secret_url, {"API_KEY": "sk-env-secret"})
close_stdio_sessions(secret_url)
gen_keys = list(mcp_client._stdio_cfg_close_gen) + list(mcp_client._stdio_url_close_gen)
assert gen_keys
# These maps are never pruned: neither command/URL nor env may persist.
assert all("sk-url-secret" not in repr(k) and "sk-env-secret" not in repr(k) for k in gen_keys)
def test_overlapping_calls_serialize_on_shared_session(fake_clients, monkeypatch):
class OverlapDetect(FakeClient):
active = 0
max_active = 0
async def call_tool(self, name, args):
OverlapDetect.active += 1
OverlapDetect.max_active = max(OverlapDetect.max_active, OverlapDetect.active)
try:
await asyncio.sleep(0.2)
return await super().call_tool(name, args)
finally:
OverlapDetect.active -= 1
monkeypatch.setattr(
mcp_client, "_client", lambda url, headers, use_oauth = False: OverlapDetect(url)
)
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
workers = [
threading.Thread(target = lambda: call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat"))
for _ in range(2)
]
for worker in workers:
worker.start()
for worker in workers:
worker.join(10.0)
# A stateful server must never see interleaved same-scope operations.
assert OverlapDetect.max_active == 1
assert len(fake_clients) == 1
def test_timeout_budget_spans_connect_and_call(fake_clients, monkeypatch):
class SlowBoth(FakeClient):
async def __aenter__(self):
await asyncio.sleep(0.4)
return await super().__aenter__()
async def call_tool(self, name, args):
await asyncio.sleep(0.5)
return await super().call_tool(name, args)
monkeypatch.setattr(mcp_client, "_client", lambda url, headers, use_oauth = False: SlowBoth(url))
start = time.monotonic()
# 0.4s connect + 0.5s call vs a 0.6s budget: the call must inherit only
# the remaining ~0.2s, not a fresh full window.
out = call_tool_sync(STDIO_URL, None, "t", {}, timeout = 0.6, scope = "chat")
assert "timed out" in out
assert time.monotonic() - start < 2.0
def test_close_narrowed_by_headers_spares_other_env(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
call_tool_sync(STDIO_URL, {"ENV_VAR": "b"}, "t", {}, scope = "chat")
# Two server rows can share a command with different envs; editing one
# must only close its own sessions.
close_stdio_sessions(STDIO_URL, None)
assert fake_clients[0].exited == 1
assert fake_clients[1].exited == 0
assert len(mcp_client._stdio_sessions) == 1
close_stdio_sessions(STDIO_URL) # headers omitted: any env for the command
assert fake_clients[1].exited == 1
assert mcp_client._stdio_sessions == {}
def test_close_stdio_sessions_by_url(fake_clients):
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
call_tool_sync("npx other-server", None, "t", {}, scope = "chat")
key = mcp_client._session_key(STDIO_URL, None, "chat")
close_stdio_sessions(STDIO_URL)
assert fake_clients[0].exited == 1
assert fake_clients[1].exited == 0
assert len(mcp_client._stdio_sessions) == 1
assert key not in mcp_client._stdio_key_locks
def test_execute_tool_mcp_scope_is_per_thread(tmp_path, monkeypatch):
# session_id is the sandbox id and can be shared project-wide; the stdio
# session scope must also carry the per-conversation thread id.
from core.inference import tools as tools_mod
from storage import mcp_servers_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True)
mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True)
scopes: list = []
def fake_call_tool_sync(**kwargs):
scopes.append(kwargs["scope"])
return "ok"
monkeypatch.setattr(tools_mod, "call_tool_sync", fake_call_tool_sync)
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-a")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "project-p1", thread_id = "thread-b")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "sess-only")
tools_mod.execute_tool("mcp__s1__t", {}, thread_id = "thread-a")
# Persist only with a thread_id; session_id alone stays one-shot (None) so a
# project-wide id can't leak state across conversations. Fields are tagged.
assert scopes == ["s=project-p1:t=thread-a", "s=project-p1:t=thread-b", None, "s=:t=thread-a"]
# IDs containing ":" must not collapse distinct conversations into one scope,
# and a session-only id must never collide with a thread-only id.
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "a:b", thread_id = "c")
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "a", thread_id = "b:c")
assert scopes[-2] != scopes[-1]
tools_mod.execute_tool("mcp__s1__t", {}, session_id = "same")
tools_mod.execute_tool("mcp__s1__t", {}, thread_id = "same")
assert scopes[-2] != scopes[-1] # session-only "same" != thread-only "same"
def test_execute_tool_config_check_tracks_row(tmp_path, monkeypatch):
from core.inference import tools as tools_mod
from storage import mcp_servers_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
monkeypatch.setattr(tools_mod, "stdio_mcp_enabled", lambda: True)
mcp_servers_db.create_server(id = "s1", display_name = "S", url = STDIO_URL, is_enabled = True)
captured: dict = {}
monkeypatch.setattr(tools_mod, "call_tool_sync", lambda **kw: captured.update(kw) or "ok")
tools_mod.execute_tool("mcp__s1__t", {})
check = captured["config_check"]
assert check() is True
mcp_servers_db.update_server("s1", {"url": "npx different-server"})
assert check() is False
def test_multi_block_result_flattens_through_session(fake_clients):
async def _rich_call(name, args):
return SimpleNamespace(
content = [
SimpleNamespace(type = "text", text = "### Page"),
SimpleNamespace(type = "text", text = "- Page URL: https://example.com/"),
],
is_error = False,
structured_content = None,
)
call_tool_sync(STDIO_URL, None, "t", {}, scope = "chat")
fake_clients[0].call_tool = _rich_call
out = call_tool_sync(STDIO_URL, None, "browser_snapshot", {}, scope = "chat")
assert out == "### Page\n- Page URL: https://example.com/"
def test_stdio_cache_trims_overshoot_after_burst(fake_clients, monkeypatch):
# A concurrent burst of distinct-scope calls can overshoot the cap while every
# session is busy (insert-time eviction only reclaims idle sessions). Once the
# calls finish, release-time trimming must bring the cache back within cap.
monkeypatch.setattr(mcp_client, "_STDIO_MAX_SESSIONS", 2)
def slow_client(
url,
headers,
use_oauth = False,
):
client = FakeClient(url)
client.call_delay = 0.5 # keep every session in-flight during the burst
return client
monkeypatch.setattr(mcp_client, "_client", slow_client)
errors: list = []
def worker(i: int):
try:
call_tool_sync(STDIO_URL, None, "t", {}, scope = f"chat-{i}")
except Exception as exc: # noqa: BLE001
errors.append(exc)
threads = [threading.Thread(target = worker, args = (i,)) for i in range(5)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(10.0)
assert not errors, errors
assert len(mcp_client._stdio_sessions) <= 2
def test_close_http_server_creates_no_stdio_tombstone(fake_clients):
# HTTP/SSE servers are never cached as stdio sessions, so closing one on
# update/delete must not accrue a close-generation entry (an unbounded leak).
before_cfg = len(mcp_client._stdio_cfg_close_gen)
before_url = len(mcp_client._stdio_url_close_gen)
for i in range(50):
close_stdio_sessions(f"https://mcp-{i}.example/mcp", {"K": str(i)})
close_stdio_sessions(f"https://mcp-{i}.example/mcp")
assert len(mcp_client._stdio_cfg_close_gen) == before_cfg
assert len(mcp_client._stdio_url_close_gen) == before_url
# a real stdio command still registers a generation (the guard is non-stdio only)
close_stdio_sessions(STDIO_URL, {"K": "v"})
assert len(mcp_client._stdio_cfg_close_gen) == before_cfg + 1

1
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -1179,6 +1179,7 @@ class FakeExecuteTool:
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

1
studio/backend/tests/test_safetensors_toolcall_wiring.py
View file

@ -91,6 +91,7 @@ class StubExecutor:
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

1
studio/backend/tests/test_tool_confirm_loop.py
View file

@ -42,6 +42,7 @@ class _FakeExecuteTool:
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

1
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -2922,6 +2922,7 @@ export function createOpenAIStreamAdapter(
audio_base64: audioBase64,
cancel_id: cancelId,
...(sandboxSessionId ? { session_id: sandboxSessionId } : {}),
...(resolvedThreadId ? { thread_id: resolvedThreadId } : {}),
...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
...(supportsReasoning
? reasoningStyle === "enable_thinking_effort"