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 ------------------------------------------------------------