mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-31 19:16:29 +00:00
feat: prompt-cache prewarm for TUI/desktop sessions (agent.prewarm_prompt_cache)
The first API call of a fresh desktop/TUI session pays provider-side ingestion of the entire uncached prefix (system prompt + tool schemas, commonly 50-70k tokens) — observed as ~20s first-message latency on Anthropic-cached routes, vs 4-7s on every later 100%-cache-hit turn. When agent.prewarm_prompt_cache is enabled (config.yaml, default off), the gateway issues one minimal non-streaming max_tokens=1 request right after the session agent is built — same tool schemas, same system prompt with the same [static, volatile] cache_control layout — so the provider writes the prompt-prefix cache BEFORE the first user message. The first real turn then reads a warm prefix instead of writing it cold. Measured on a live desktop-shaped session (nous / claude-fable-5, 67k prefix): prewarm 2.7s off the response path, first real turn 3.8s with cache=67045/67131 (100%) — down from 20.4s cold. Details: - agent/prompt_prewarm.py: pure helper; supported only where the request shape is reproducible (chat_completions / anthropic_messages, not MoA/ Codex/Bedrock/ACP) and _use_prompt_caching is on. Thinking/reasoning knobs are stripped (max_tokens=1 violates budget_tokens; thinking changes don't invalidate system/tools cache blocks). Fail-open: any failure returns False and the first real turn pays the write itself. - The exact sent prompt is handed to the first real turn via _prewarmed_system_prompt; _restore_or_build_system_prompt adopts it (gated on runtime-identity match, no history, no custom system message) so the volatile tail can't drift and split the just-warmed prefix. One-shot — cleared after every first-turn resolution. - tui_gateway/server.py: _schedule_prompt_prewarm fires from both agent build sites, waits for late MCP discovery first (tools are part of the cached prefix), and skips if the user already started the conversation. Cost note: the cache write (1.25x input, 5m TTL) is paid by the first real call today anyway; prewarming moves it earlier. Extra spend is one 0.1x cache read per session plus wasted writes for sessions opened but never used — which is why it ships default-off.
This commit is contained in:
parent
d83e858507
commit
e2dfa843ef
5 changed files with 511 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
167
agent/prompt_prewarm.py
Normal file
167
agent/prompt_prewarm.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
238
tests/agent/test_prompt_prewarm.py
Normal file
238
tests/agent/test_prompt_prewarm.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue