diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 2efad4e9580..2ad7ee51199 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -487,8 +487,27 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history) ) # First turn of a new session (or recovering from a broken stored - # prompt) — build from scratch. - agent._cached_system_prompt = agent._build_system_prompt(system_message) + # prompt) — build from scratch. One exception: a prompt-cache prewarm + # (agent.prompt_prewarm) already built and SENT a system prompt for this + # session before the first user message. Reuse those exact bytes when the + # runtime identity still matches — a rebuild here could differ in the + # volatile tail (timestamp) and split the just-warmed provider prefix. + _prewarmed = getattr(agent, "_prewarmed_system_prompt", None) + if ( + not conversation_history + and system_message is None + and isinstance(_prewarmed, str) + and _prewarmed + and _stored_prompt_matches_runtime(agent, _prewarmed) + ): + agent._cached_system_prompt = _prewarmed + from agent.system_prompt import reconstruct_static_prefix + + reconstruct_static_prefix(agent, log_label="prewarm") + else: + agent._cached_system_prompt = agent._build_system_prompt(system_message) + # One-shot: never reuse across /new, compression rebuilds, or model swaps. + agent._prewarmed_system_prompt = None # Plugin hook: on_session_start — fired once when a brand-new # session is created (not on continuation). Plugins can use this diff --git a/agent/prompt_prewarm.py b/agent/prompt_prewarm.py new file mode 100644 index 00000000000..5d9ba510af7 --- /dev/null +++ b/agent/prompt_prewarm.py @@ -0,0 +1,167 @@ +"""Prompt-cache prewarm — pay the provider-side cache write before turn 1. + +The first API call of a fresh session ingests the entire uncached prefix +(system prompt + tool schemas — commonly 50-70k tokens), which shows up as +10-20s of first-message latency on Anthropic-cached routes. Every later call +rides the prompt cache and drops to TTFT + generation. + +``prewarm_prompt_cache`` issues one minimal, non-streaming request (same +tools, same system prompt, ``max_tokens=1``) whose only purpose is the +provider-side cache write at the injected ``cache_control`` breakpoints. +The first real user turn then *reads* the prefix cache instead of writing +it cold. + +Cost note: the cache write (1.25x input for the 5m TTL) is paid by the +first real call today anyway — prewarming only moves it earlier. The extra +spend is one cache *read* (0.1x) on the first real turn, plus the full +write being wasted when a session is opened but never used. That trade-off +is why the feature is config-gated (``agent.prewarm_prompt_cache``, +default off). + +Everything here is fail-open: a prewarm failure must never surface to the +user or block the session — the worst case is the status quo (cold first +turn). + +Pure helper — no threads. Callers (the TUI/desktop gateway) own scheduling. +""" + +import logging +import time + +logger = logging.getLogger(__name__) + +# Transports whose request shape we can safely reproduce outside the main +# conversation loop. Codex/Responses, Bedrock Converse, ACP, and the MoA +# facade have bespoke client/stream ownership — and (except Bedrock) no +# Anthropic cache_control semantics — so they are excluded. +_PREWARM_API_MODES = {"chat_completions", "anthropic_messages"} + + +def prewarm_supported(agent) -> bool: + """True when a prewarm request would actually warm a provider cache.""" + if not getattr(agent, "_use_prompt_caching", False): + return False + if getattr(agent, "api_mode", None) not in _PREWARM_API_MODES: + return False + if getattr(agent, "provider", None) == "moa": + return False + return True + + +def prewarm_prompt_cache(agent) -> bool: + """Issue one minimal request so the provider writes the prompt-prefix cache. + + Builds the exact prefix the first real turn will send — same tool + schemas (``agent.tools`` via ``_build_api_kwargs``), same system prompt + with the same ``[static, volatile]`` cache_control layout — followed by + a throwaway user message. The trailing user message differs from the + real first message, but Anthropic caching is prefix-based: the + breakpoints on the static system prefix and full system prompt land + identically, which is where the tens of thousands of tokens live. + + Returns True when the prewarm request completed, False when skipped or + failed. Never raises. + """ + if not prewarm_supported(agent): + return False + + try: + system_prompt = getattr(agent, "_cached_system_prompt", None) + if not system_prompt: + # Same builder the first turn uses; also populates + # ``_cached_system_prompt_static`` as a side effect. The real + # turn rebuilds through ``_restore_or_build_system_prompt`` — + # only the volatile tail (timestamp) can differ, and that sits + # after the static-prefix breakpoint. + system_prompt = agent._build_system_prompt() + static_prefix = getattr(agent, "_cached_system_prompt_static", None) + + from agent.prompt_caching import apply_anthropic_cache_control + + api_messages = apply_anthropic_cache_control( + [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": "ping"}, + ], + cache_ttl=getattr(agent, "_cache_ttl", "5m") or "5m", + native_anthropic=getattr(agent, "_use_native_cache_layout", False), + static_system_prefix=( + static_prefix if isinstance(static_prefix, str) else None + ), + ) + + api_kwargs = agent._build_api_kwargs(api_messages) + # Cheapest possible generation: the request exists only for its + # prompt-ingestion side effect. Both output-cap spellings are + # forced so no provider default (or profile-injected cap) makes + # the throwaway completion generate real output. + if "max_completion_tokens" in api_kwargs: + api_kwargs["max_completion_tokens"] = 1 + else: + api_kwargs["max_tokens"] = 1 + api_kwargs.pop("stream", None) + # Extended thinking demands max_tokens > budget_tokens, which a + # 1-token request violates (400). Strip thinking/reasoning knobs: + # per Anthropic caching semantics, thinking-parameter changes only + # invalidate message-level cache blocks — the system-prompt and + # tools cache blocks (the entire point of the prewarm) survive. + api_kwargs.pop("thinking", None) + api_kwargs.pop("reasoning_effort", None) + extra_body = api_kwargs.get("extra_body") + if isinstance(extra_body, dict): + extra_body.pop("reasoning", None) + extra_body.pop("thinking", None) + + from agent.chat_completion_helpers import _dispatch_nonstreaming_api_request + + created: list[tuple[object, str]] = [] + + def _make_client(reason: str, kind: str = "openai"): + client = ( + agent._create_request_anthropic_client(reason=reason) + if kind == "anthropic_messages" + else agent._create_request_openai_client( + reason=reason, api_kwargs=api_kwargs + ) + ) + created.append((client, kind)) + return client + + started = time.monotonic() + try: + _dispatch_nonstreaming_api_request( + agent, api_kwargs, make_client=_make_client + ) + finally: + for client, kind in created: + try: + if kind == "anthropic_messages": + close = getattr(client, "close", None) + if callable(close): + close() + else: + agent._close_request_openai_client( + client, reason="prompt_prewarm" + ) + except Exception: + logger.debug( + "prompt prewarm client close failed", exc_info=True + ) + + logger.info( + "Prompt-cache prewarm complete: model=%s provider=%s %.1fs", + getattr(agent, "model", "") or "", + getattr(agent, "provider", "") or "", + time.monotonic() - started, + ) + # Hand the exact sent bytes to the first real turn: + # ``_restore_or_build_system_prompt`` adopts this instead of + # rebuilding, so the volatile tail (timestamp) can't drift between + # the prewarmed prefix and the first real request. + agent._prewarmed_system_prompt = system_prompt + return True + except Exception: + # Fail-open by contract: a failed prewarm just means the first real + # turn pays the cache write itself (the status quo). + logger.info("Prompt-cache prewarm failed (fail-open)", exc_info=True) + return False diff --git a/hermes_cli/config.py b/hermes_cli/config.py index 3e8b6e5f07d..6dbe922f4ba 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -974,6 +974,17 @@ DEFAULT_CONFIG = { # on a genuinely hung build. Raise it for deployments with many slow # or unreachable MCP servers. "build_wait_timeout": 600, + # Prompt-cache prewarm (TUI/desktop): right after a session's agent is + # built, issue one minimal max_tokens=1 request so the provider writes + # the prompt-prefix cache (system prompt + tool schemas) BEFORE the + # first user message. Cuts cold first-message latency from provider + # ingestion of a 50-70k-token uncached prefix (10-20s observed) down + # to a cache read. Off by default: the write (1.25x input for the 5m + # TTL) is wasted whenever a session is opened but never used, and the + # first real turn additionally pays one cache read (0.1x). Only + # applies to prompt-caching routes (Anthropic-compatible); no-op + # elsewhere. + "prewarm_prompt_cache": False, # Max app-level retry attempts for API errors (connection drops, # provider timeouts, 5xx, etc.) before the agent surfaces the # failure. The OpenAI SDK already does its own low-level retries diff --git a/tests/agent/test_prompt_prewarm.py b/tests/agent/test_prompt_prewarm.py new file mode 100644 index 00000000000..f010cf610a2 --- /dev/null +++ b/tests/agent/test_prompt_prewarm.py @@ -0,0 +1,238 @@ +"""Tests for ``agent.prompt_prewarm`` — the prompt-cache prewarm request. + +The prewarm exists to pay the provider-side prompt-cache write (system +prompt + tool schemas) before the first user message, so the first real +turn reads a warm prefix instead of ingesting 50-70k uncached tokens. + +Behavior contracts covered: + + * ``prewarm_supported`` gates on prompt caching being active and on a + reproducible transport (chat_completions / anthropic_messages, not MoA). + * ``prewarm_prompt_cache`` sends the SAME system prompt bytes the first + real turn will send, with cache_control markers, capped at 1 output + token, non-streaming, and with thinking/reasoning knobs stripped. + * The request client it creates is always closed (success and failure). + * Fail-open: any error → returns False, never raises. + * The sent prompt is handed to the first real turn via + ``agent._prewarmed_system_prompt`` and adopted by + ``_restore_or_build_system_prompt`` so the volatile tail (timestamp) + cannot drift between the prewarm and the first real request. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from agent.prompt_prewarm import prewarm_prompt_cache, prewarm_supported + + +SYSTEM_PROMPT = "You are Hermes Agent.\n\nSTATIC PART\n\nVolatile tail" +STATIC_PREFIX = "You are Hermes Agent.\n\nSTATIC PART" + + +def _make_agent(api_mode: str = "chat_completions"): + agent = MagicMock() + agent._use_prompt_caching = True + agent._use_native_cache_layout = False + agent._cache_ttl = "5m" + agent.api_mode = api_mode + agent.provider = "nous" + agent.model = "anthropic/claude-fable-5" + agent._cached_system_prompt = SYSTEM_PROMPT + agent._cached_system_prompt_static = STATIC_PREFIX + agent._prewarmed_system_prompt = None + + # _build_api_kwargs echoes the messages it was given, like the real one. + def _build_api_kwargs(api_messages): + return { + "model": agent.model, + "messages": api_messages, + "max_tokens": 8192, + "stream": False, + } + + agent._build_api_kwargs = MagicMock(side_effect=_build_api_kwargs) + + client = MagicMock() + agent._create_request_openai_client = MagicMock(return_value=client) + agent._create_request_anthropic_client = MagicMock(return_value=client) + agent._request_client = client + return agent + + +class TestPrewarmSupported: + def test_supported_on_chat_completions_with_caching(self): + assert prewarm_supported(_make_agent()) is True + + def test_supported_on_anthropic_messages(self): + assert prewarm_supported(_make_agent(api_mode="anthropic_messages")) is True + + def test_not_supported_without_prompt_caching(self): + agent = _make_agent() + agent._use_prompt_caching = False + assert prewarm_supported(agent) is False + + def test_not_supported_on_bespoke_transports(self): + for mode in ("codex_responses", "bedrock_converse", "acp"): + assert prewarm_supported(_make_agent(api_mode=mode)) is False + + def test_not_supported_for_moa(self): + agent = _make_agent() + agent.provider = "moa" + assert prewarm_supported(agent) is False + + +class TestPrewarmRequest: + def test_sends_cached_system_prompt_with_markers(self): + agent = _make_agent() + assert prewarm_prompt_cache(agent) is True + + create = agent._request_client.chat.completions.create + assert create.call_count == 1 + kwargs = create.call_args.kwargs + messages = kwargs["messages"] + + assert messages[0]["role"] == "system" + # The static prefix split produced the two-part [static, volatile] + # layout, each part carrying a cache_control marker, and the joined + # bytes are exactly the prompt the first real turn will send. + system_content = messages[0]["content"] + assert isinstance(system_content, list) + assert "".join(p["text"] for p in system_content) == SYSTEM_PROMPT + assert all("cache_control" in p for p in system_content) + assert messages[1]["role"] == "user" + + def test_output_capped_to_one_token_and_non_streaming(self): + agent = _make_agent() + prewarm_prompt_cache(agent) + kwargs = agent._request_client.chat.completions.create.call_args.kwargs + assert kwargs["max_tokens"] == 1 + assert "stream" not in kwargs + + def test_thinking_and_reasoning_knobs_stripped(self): + agent = _make_agent() + + def _build_api_kwargs(api_messages): + return { + "model": agent.model, + "messages": api_messages, + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "reasoning_effort": "high", + "extra_body": {"reasoning": {"effort": "high"}, "keep": 1}, + } + + agent._build_api_kwargs = MagicMock(side_effect=_build_api_kwargs) + assert prewarm_prompt_cache(agent) is True + kwargs = agent._request_client.chat.completions.create.call_args.kwargs + assert "thinking" not in kwargs + assert "reasoning_effort" not in kwargs + assert "reasoning" not in kwargs["extra_body"] + assert kwargs["extra_body"]["keep"] == 1 + + def test_builds_prompt_when_not_cached(self): + agent = _make_agent() + agent._cached_system_prompt = None + agent._build_system_prompt = MagicMock(return_value=SYSTEM_PROMPT) + assert prewarm_prompt_cache(agent) is True + agent._build_system_prompt.assert_called_once_with() + + def test_hands_sent_prompt_to_first_real_turn(self): + agent = _make_agent() + prewarm_prompt_cache(agent) + assert agent._prewarmed_system_prompt == SYSTEM_PROMPT + + def test_request_client_closed_on_success(self): + agent = _make_agent() + prewarm_prompt_cache(agent) + agent._close_request_openai_client.assert_called_once() + + def test_request_client_closed_on_request_failure(self): + agent = _make_agent() + agent._request_client.chat.completions.create.side_effect = RuntimeError( + "provider 500" + ) + assert prewarm_prompt_cache(agent) is False + agent._close_request_openai_client.assert_called_once() + + def test_fail_open_never_raises(self): + agent = _make_agent() + agent._build_api_kwargs = MagicMock(side_effect=RuntimeError("boom")) + assert prewarm_prompt_cache(agent) is False + assert agent._prewarmed_system_prompt is None + + def test_skips_unsupported_agent(self): + agent = _make_agent() + agent._use_prompt_caching = False + assert prewarm_prompt_cache(agent) is False + agent._build_api_kwargs.assert_not_called() + + def test_anthropic_messages_uses_anthropic_client(self): + agent = _make_agent(api_mode="anthropic_messages") + agent._anthropic_messages_create = MagicMock(return_value=MagicMock()) + assert prewarm_prompt_cache(agent) is True + agent._create_request_anthropic_client.assert_called_once() + agent._anthropic_messages_create.assert_called_once() + + +class TestFirstTurnAdoption: + """The first real turn must reuse the exact prewarmed bytes.""" + + def _restore_agent(self, prewarmed): + agent = MagicMock() + agent._cached_system_prompt = None + agent.session_id = "sid" + agent.model = "anthropic/claude-fable-5" + agent.provider = "nous" + agent.platform = "desktop" + agent._session_db = None + agent._use_prompt_caching = False + agent._prewarmed_system_prompt = prewarmed + agent._build_system_prompt = MagicMock(return_value="FRESH_BUILD") + return agent + + def test_first_turn_reuses_prewarmed_prompt(self): + from agent.conversation_loop import _restore_or_build_system_prompt + + agent = self._restore_agent(SYSTEM_PROMPT) + _restore_or_build_system_prompt(agent, None, None) + assert agent._cached_system_prompt == SYSTEM_PROMPT + agent._build_system_prompt.assert_not_called() + # One-shot: consumed after adoption. + assert agent._prewarmed_system_prompt is None + + def test_prewarmed_prompt_rejected_on_model_switch(self): + """A /model switch between prewarm and first message must rebuild.""" + from agent.conversation_loop import _restore_or_build_system_prompt + + stale = SYSTEM_PROMPT + "\nModel: anthropic/claude-old\nProvider: nous" + agent = self._restore_agent(stale) + agent.model = "openai/gpt-6" + _restore_or_build_system_prompt(agent, None, None) + assert agent._cached_system_prompt == "FRESH_BUILD" + agent._build_system_prompt.assert_called_once_with(None) + assert agent._prewarmed_system_prompt is None + + def test_prewarmed_prompt_ignored_with_custom_system_message(self): + from agent.conversation_loop import _restore_or_build_system_prompt + + agent = self._restore_agent(SYSTEM_PROMPT) + _restore_or_build_system_prompt(agent, "custom system", None) + assert agent._cached_system_prompt == "FRESH_BUILD" + agent._build_system_prompt.assert_called_once_with("custom system") + + def test_prewarmed_prompt_ignored_on_continuing_session(self): + from agent.conversation_loop import _restore_or_build_system_prompt + + agent = self._restore_agent(SYSTEM_PROMPT) + _restore_or_build_system_prompt( + agent, None, [{"role": "user", "content": "hi"}] + ) + assert agent._cached_system_prompt == "FRESH_BUILD" + + def test_no_prewarm_builds_fresh(self): + from agent.conversation_loop import _restore_or_build_system_prompt + + agent = self._restore_agent(None) + _restore_or_build_system_prompt(agent, None, None) + assert agent._cached_system_prompt == "FRESH_BUILD" diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 5ab0dd44676..9fb589774c2 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1969,6 +1969,9 @@ def _start_agent_build(sid: str, session: dict) -> None: # was built without those tools. Catch up once they land — see # _schedule_mcp_late_refresh. Cache-safe (pre-first-turn only). _schedule_mcp_late_refresh(sid, agent) + # Optionally pay the provider-side prompt-cache write now so the + # first real message reads a warm prefix (agent.prewarm_prompt_cache). + _schedule_prompt_prewarm(sid, agent) except Exception as e: current["agent_error"] = str(e) _emit("error", sid, {"message": f"agent init failed: {e}"}) @@ -5595,6 +5598,74 @@ def _schedule_mcp_late_refresh(sid: str, agent) -> None: ).start() +def _schedule_prompt_prewarm(sid: str, agent) -> None: + """Warm the provider prompt cache for a freshly built session agent. + + The first API call of a session pays provider-side ingestion of the whole + uncached prefix (system prompt + tool schemas, commonly 50-70k tokens) — + observed as 10-20s of first-message latency on Anthropic-cached routes. + When ``agent.prewarm_prompt_cache`` is enabled, issue one minimal request + off the response path right after the agent is built so the first real + turn *reads* the prefix cache instead of writing it cold. + + Ordering: if MCP discovery is still in flight, wait for it (the late + refresh rebuilds ``agent.tools``, and the tools block is part of the + cached prefix — prewarming before it lands would warm a prefix the real + turn no longer sends). Skipped silently the moment the user has already + started the conversation; the in-flight real call warms the cache itself. + """ + try: + agent_cfg = _load_cfg().get("agent") or {} + if not ( + isinstance(agent_cfg, dict) + and is_truthy_value(agent_cfg.get("prewarm_prompt_cache", False)) + ): + return + from agent.prompt_prewarm import prewarm_supported + + if not prewarm_supported(agent): + return + except Exception: + return + + def _wait_then_prewarm() -> None: + try: + from tui_gateway.entry import ( + mcp_discovery_in_flight, + join_mcp_discovery, + ) + + if mcp_discovery_in_flight(): + # Same bound as the late MCP refresh; a server slower than + # this is dead-ish and the tool snapshot is already frozen. + join_mcp_discovery(timeout=30.0) + # Give _schedule_mcp_late_refresh's rebuild a beat to land so + # the prewarmed tools block matches the refreshed snapshot. + time.sleep(0.5) + except Exception: + pass + session = _sessions.get(sid) + if session is None or session.get("agent") is not agent: + return + # The user beat us to it — their real first call is already warming + # the cache; a duplicate prefix ingestion would only add cost. + if ( + session.get("running") + or int(getattr(agent, "_user_turn_count", 0) or 0) > 0 + or int(getattr(agent, "_api_call_count", 0) or 0) > 0 + ): + return + from agent.prompt_prewarm import prewarm_prompt_cache + + prewarm_prompt_cache(agent) + + threading.Thread( + target=_wait_then_prewarm, + name=f"tui-prompt-prewarm-{sid}", + daemon=True, + ).start() + + class _RuntimeFallbackResolution(NamedTuple): runtime: dict selected_model: str | None @@ -5962,6 +6033,9 @@ def _init_session( _notify_session_boundary("on_session_reset", key, _session_source(_sessions.get(sid, {}))) _emit("session.info", sid, _session_info(agent, _sessions.get(sid, {}))) _schedule_mcp_late_refresh(sid, agent) + # Optionally pay the provider-side prompt-cache write now so the first + # real message reads a warm prefix (agent.prewarm_prompt_cache). + _schedule_prompt_prewarm(sid, agent) def _new_session_key() -> str: