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:
Teknium 2026-07-27 18:04:16 -07:00
parent d83e858507
commit e2dfa843ef
5 changed files with 511 additions and 2 deletions

View file

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