From 601155114d0ff29ae4070007f9a334416ebdb031 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 14 Jul 2026 14:58:43 +0530 Subject: [PATCH] Studio: persistent stdio MCP sessions so server state survives across tool calls (#7080) * 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 --- studio/backend/core/inference/inference.py | 2 + studio/backend/core/inference/llama_cpp.py | 2 + studio/backend/core/inference/mcp_client.py | 739 ++++++++++++++++-- studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 2 + studio/backend/core/inference/tools.py | 35 +- studio/backend/models/inference.py | 8 + studio/backend/routes/inference.py | 3 + studio/backend/routes/mcp_servers.py | 15 +- .../backend/tests/test_bypass_permissions.py | 2 + studio/backend/tests/test_mcp_servers.py | 62 ++ .../backend/tests/test_mcp_stdio_sessions.py | 629 +++++++++++++++ .../tests/test_safetensors_tool_loop.py | 1 + .../tests/test_safetensors_toolcall_wiring.py | 1 + .../backend/tests/test_tool_confirm_loop.py | 1 + .../src/features/chat/api/chat-adapter.ts | 1 + 16 files changed, 1453 insertions(+), 52 deletions(-) create mode 100644 studio/backend/tests/test_mcp_stdio_sessions.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index bb4be39ef6..172e9e5546 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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, ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b06c6eb5cb..d7e3173f37 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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, ) diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index c6b8acfdc4..6b5ce02216 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -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 "" + except Exception: # noqa: BLE001 + exe = "" + 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}" diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4fa0d3ed26..2b1ceca75a 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 81c25b777e..a18d2758ba 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -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, ) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 82c50933fc..8d8e7ef3dd 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -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( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 53b0f14b09..53441bdb3a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ec5d309810..8141ff073a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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, diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 80cd8e1437..71f0fd2874 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -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) diff --git a/studio/backend/tests/test_bypass_permissions.py b/studio/backend/tests/test_bypass_permissions.py index d92509a5fe..585aa207a7 100644 --- a/studio/backend/tests/test_bypass_permissions.py +++ b/studio/backend/tests/test_bypass_permissions.py @@ -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, ): diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 24784ca34c..6432ffb8e1 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -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 diff --git a/studio/backend/tests/test_mcp_stdio_sessions.py b/studio/backend/tests/test_mcp_stdio_sessions.py new file mode 100644 index 0000000000..d714d9d640 --- /dev/null +++ b/studio/backend/tests/test_mcp_stdio_sessions.py @@ -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 diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index a8546b82c4..63fdbbd8e9 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -1179,6 +1179,7 @@ class FakeExecuteTool: cancel_event = None, timeout = None, session_id = None, + thread_id = None, rag_scope = None, disable_sandbox = False, ): diff --git a/studio/backend/tests/test_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py index 5c298a7966..8909e0c0c5 100644 --- a/studio/backend/tests/test_safetensors_toolcall_wiring.py +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -91,6 +91,7 @@ class StubExecutor: cancel_event = None, timeout = None, session_id = None, + thread_id = None, rag_scope = None, disable_sandbox = False, ): diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py index 17ef697674..3db591f542 100644 --- a/studio/backend/tests/test_tool_confirm_loop.py +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -42,6 +42,7 @@ class _FakeExecuteTool: cancel_event = None, timeout = None, session_id = None, + thread_id = None, rag_scope = None, disable_sandbox = False, ): diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 4fb7d9bd7e..9ed11deccc 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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"