beaver/claude-code-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(stream): warm-up turn, because MCP servers aren't up on turn one

Browse Source 
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.
This commit is contained in:
hh
2026-07-28 03:22:30 +02:00
parent b50911c3ed
commit 621da90623
3 changed files with 184 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/claude_code_api/backend.py
+24 -1
View File
@@ -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:
src/claude_code_api/stream.py
+92
View File
@@ -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.
tests/test_stream.py
+68
View File
@@ -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 ------------------------------------------------------------