From 621da90623279949fbe47b80bfa57a7c8c6bf74b Mon Sep 17 00:00:00 2001 From: h Date: Tue, 28 Jul 2026 03:22:30 +0200 Subject: [PATCH] feat(stream): warm-up turn, because MCP servers aren't up on turn one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CLI starts connecting its MCP servers when a turn starts and does not wait for them, so the first turn of a session sees only built-in tools. Measured on the pi with two HTTP MCP servers: turn 1 reports both `pending` with 29 tools available, turn 2 reports both `connected` with 90. For an agent whose job is those tools, that first reply is silently wrong — the model doesn't see them and answers as best it can. The PTY transport hid this: its multi-second wait for the TUI to settle happened to cover the connect. Nothing about the headless path does, and nothing cheap fixes it — an 8s pause before the first prompt changed nothing, a `/status` slash command returns in 134ms without touching the MCP client, and a `control_request`/`initialize` handshake answers with the command list and leaves the servers pending. Only a real turn does it, so `warmup_turn` spends one deliberately: a few tokens and a couple of seconds, once per session, against sessions that are pooled for the whole conversation. The price is two short messages at the head of the transcript, which is why the default prompt reads as procedural rather than conversational. Also log, at WARNING, any MCP server a turn starts without — this is a failure with no other symptom, and it should be one grep away rather than a mystery about the agent forgetting a tool it has. --- src/claude_code_api/backend.py | 25 ++++++++- src/claude_code_api/stream.py | 92 ++++++++++++++++++++++++++++++++++ tests/test_stream.py | 68 +++++++++++++++++++++++++ 3 files changed, 184 insertions(+), 1 deletion(-) diff --git a/src/claude_code_api/backend.py b/src/claude_code_api/backend.py index 61cfb30..9ab2999 100644 --- a/src/claude_code_api/backend.py +++ b/src/claude_code_api/backend.py @@ -49,7 +49,11 @@ from claude_code_api.injection import ( ) from claude_code_api.paths import resolve_jsonl_path from claude_code_api.pty import PtyClaudeProcess, PtyProcessOptions -from claude_code_api.stream import StreamClaudeProcess, StreamTurnManager +from claude_code_api.stream import ( + DEFAULT_WARMUP_PROMPT, + StreamClaudeProcess, + StreamTurnManager, +) from claude_code_api.turn import TurnManager from claude_code_api.watcher import JsonlWatcher @@ -107,6 +111,23 @@ class BackendOptions: left on the dataclass rather than rejected so a config can flip ``transport`` without being rewritten.""" + warmup_turn: bool = False + """Spend one throwaway turn when a session is created. + + Set this whenever ``mcp_servers`` is configured and the agent is + expected to have those tools on its *first* reply. The CLI starts + connecting MCP servers when a turn starts and does not wait for + them, so turn one otherwise runs with the built-in tools only — + see :meth:`claude_code_api.stream.StreamTurnManager.warmup`, which + documents what was measured and which cheaper alternatives don't + work. ``stream_json`` only; the PTY transport's readiness wait + happens to cover this already.""" + + warmup_prompt: str = DEFAULT_WARMUP_PROMPT + """The throwaway prompt used by ``warmup_turn``. It lands at the head + of the session transcript, so keep it short and obviously + procedural.""" + include_partial_messages: bool = False """Emit :class:`~claude_code_api.events.StreamEvent` for token-level deltas. Requires ``transport="stream_json"``; silently inert on the @@ -563,6 +584,8 @@ class ClaudeCodeBackend: session_id, proc.pid, ) + if self._opts.warmup_turn: + await tm.warmup(self._opts.warmup_prompt) return _LiveSession(pty=proc, watcher=None, tm=tm) def _build_pty_options(self, *, session_id: str, resume: bool) -> PtyProcessOptions: diff --git a/src/claude_code_api/stream.py b/src/claude_code_api/stream.py index 9e2fdc3..f227fd9 100644 --- a/src/claude_code_api/stream.py +++ b/src/claude_code_api/stream.py @@ -122,6 +122,14 @@ _TUI_ONLY_FLAGS: frozenset[str] = frozenset( # Anything else is a failure the caller needs to see as an exception. _OK_RESULT_SUBTYPES: frozenset[str] = frozenset({"success"}) +DEFAULT_WARMUP_PROMPT = "Session warm-up. Reply with exactly: ready" +"""See :meth:`StreamTurnManager.warmup` for why a turn has to be spent. + +Kept short and obviously procedural: it lands at the head of the +session transcript, so the model should read it as handshake noise +rather than as the user opening the conversation. +""" + def strip_tui_only_flags(extra_args: Iterable[str]) -> tuple[str, ...]: """Drop interactive-only flags from `extra_args`, loudly. @@ -575,6 +583,60 @@ class StreamTurnManager: self._started = True _log.info("start: session_id=%s READY", self._proc.session_id) + async def warmup(self, prompt: str = DEFAULT_WARMUP_PROMPT) -> None: + """Spend one throwaway turn so MCP servers are connected for the next. + + The CLI does not connect its MCP servers at spawn. It starts + connecting when a turn starts, and does not wait for the result — + so the *first* turn of a session sees only the built-in tools, + however long you wait before sending it. Measured on the pi with + two HTTP MCP servers: turn 1 reports them ``pending`` with 29 + tools available, turn 2 reports ``connected`` with 90. + + Sleeping does not help (an 8s pause before the first prompt + changed nothing), and neither does anything free: a ``/status`` + slash command returns in 134ms without touching the MCP client, + and a ``control_request``/``initialize`` handshake answers with + the command list and likewise leaves the servers pending. Only a + real turn does it, so this spends one deliberately. + + The cost is a few tokens and a couple of seconds, once per + session — and sessions are pooled across a conversation, not + per message. The visible price is two short messages at the head + of the transcript. + + Failures are logged and swallowed: a warm-up is an optimization, + and whatever broke it will resurface on the real turn with a + message that actually relates to what the caller asked for. + """ + _log.info( + "warmup: session_id=%s spending a turn to connect MCP servers", + self._proc.session_id, + ) + t0 = time.monotonic() + # The warm-up is not part of the conversation the caller thinks + # it is having, so it must not advance the counter that + # `ResultMessage.num_turns` reports. Restore rather than + # decrement: a failure can leave the counter either way. + before = self._turn_count + try: + async for _ in self.send_user_message(prompt): + pass + _log.info( + "warmup: session_id=%s done in %.2fs", + self._proc.session_id, + time.monotonic() - t0, + ) + except Exception as exc: + _log.warning( + "warmup: session_id=%s failed (%s) — continuing; the real " + "turn will report the underlying problem", + self._proc.session_id, + exc, + ) + finally: + self._turn_count = before + async def send_user_message(self, text: str) -> AsyncIterator[Event]: """Send `text` and stream typed events until the turn's ``result``.""" if not self._started: @@ -686,6 +748,8 @@ class StreamTurnManager: continue if rtype == "system" and record.get("subtype") in _BOOKKEEPING_SUBTYPES: + if record.get("subtype") == "init": + _warn_on_unconnected_mcp(record, sid) continue try: @@ -771,6 +835,34 @@ class StreamTurnManager: await self.aclose() +def _warn_on_unconnected_mcp(record: Mapping[str, Any], session_id: str) -> None: + """Log the MCP servers this turn is starting without. + + A turn that begins while a server is still ``pending`` runs without + that server's tools — silently, from the model's point of view: it + simply doesn't see them and answers as best it can. That is exactly + the failure :meth:`StreamTurnManager.warmup` exists to prevent, so + when it happens anyway it should be one grep away rather than a + mystery about why the agent "forgot" it had a tool. + """ + servers = record.get("mcp_servers") + if not isinstance(servers, list): + return + stragglers = [ + f"{s.get('name')}:{s.get('status')}" + for s in servers + if isinstance(s, Mapping) and s.get("status") not in ("connected", None) + ] + if stragglers: + _log.warning( + "session_id=%s turn starting without %d MCP server(s): %s — " + "their tools are unavailable for this turn", + session_id, + len(stragglers), + ", ".join(stragglers), + ) + + def _adapt_envelope(record: Mapping[str, Any]) -> dict[str, Any]: """Bridge stdout's envelope naming to the JSONL naming `normalize` knows. diff --git a/tests/test_stream.py b/tests/test_stream.py index f5a0e62..551803f 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -311,6 +311,74 @@ async def test_partial_message_indices_are_rebased_across_requests() -> None: ) +async def test_warmup_spends_a_turn_without_counting_it() -> None: + """The warm-up must not be visible in the conversation's numbering. + + It is a handshake with the CLI, not something the caller asked for, + so `ResultMessage.num_turns` on the first real reply should say 1. + """ + proc = FakeStreamProcess( + [ + _assistant({"type": "text", "text": "ready"}), + _result(), + _assistant({"type": "text", "text": "real answer"}), + _result(), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + await tm.warmup("warm me") + assert proc.sent == ["warm me"] + assert tm.turn_count == 0 + + await _drain(tm, "the real question") + assert proc.sent == ["warm me", "the real question"] + assert tm.turn_count == 1 + + +async def test_warmup_failure_is_swallowed() -> None: + """A failed warm-up must not take the session down with it.""" + proc = FakeStreamProcess([], stderr="boom") + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + await tm.warmup() # would raise ProcessError if it propagated + assert tm.turn_count == 0 + + +async def test_unconnected_mcp_servers_are_logged( + caplog: Any, +) -> None: + """A turn starting without its MCP tools has to be greppable. + + Nothing else surfaces it: the model simply doesn't see those tools + and answers anyway, which reads as the agent forgetting a capability. + """ + import logging + + proc = FakeStreamProcess( + [ + { + "type": "system", + "subtype": "init", + "session_id": "s", + "tools": ["Bash"], + "mcp_servers": [ + {"name": "telegram", "status": "pending"}, + {"name": "firefly", "status": "connected"}, + ], + }, + _assistant({"type": "text", "text": "x"}), + _result(), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + with caplog.at_level(logging.WARNING, logger="claude_code_api.stream"): + await _drain(tm) + assert "telegram:pending" in caplog.text + assert "firefly" not in caplog.text + + # --- failures ------------------------------------------------------------