From 891f44c1a9e52ea3f37b4c958129ac1298889da2 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Mon, 12 Oct 2026 23:01:46 -0400 Subject: [PATCH 1/3] Keep the chat prompt prefix stable for llama-server prompt caching llama-server only reuses a byte-identical prompt prefix. Several things changed the prefix every turn, forcing most of the prompt to be re-prefilled: - Memories: rank still chooses which memories fit the budget, but the chosen ones are emitted in stored order, so the block is identical across queries. - History trimming: the conversation-evidence summary is only included when turns were actually dropped (it duplicated retained turns otherwise). On overflow, old turns are dropped at stable cut points down to ~65% of the budget, so the retained prefix holds for several turns instead of shifting by one turn every request. - Tools: TOOL_ROUTING_MODE=static (default) binds the full role-appropriate tool set in registration order with no discovery tool; admin gating is unchanged. TOOL_ROUTING_MODE=dynamic keeps the keyword router. - Observability: per-call model_calls in the stream_request log now include cached_prompt_tokens when the provider reports cache reads. Co-Authored-By: Claude Opus 5.5 --- .env.example | 3 + blueprints/conversation/__init__.py | 5 + blueprints/conversation/agents.py | 24 ++-- blueprints/conversation/context.py | 150 ++++++++++++++++-------- blueprints/conversation/memory.py | 13 +- blueprints/conversation/tool_routing.py | 17 ++- docker-compose.yml | 1 + tests/unit/test_context_summary.py | 3 +- tests/unit/test_context_trim.py | 49 ++++++++ tests/unit/test_memory_retrieval.py | 31 +++++ tests/unit/test_stream_stall.py | 7 +- tests/unit/test_tool_routing.py | 104 +++++++++++++++- 12 files changed, 338 insertions(+), 69 deletions(-) diff --git a/.env.example b/.env.example index 373eba9..a6f27c3 100644 --- a/.env.example +++ b/.env.example @@ -32,6 +32,9 @@ TTFT_TIMEOUT_SECONDS=30 CONTEXT_WINDOW_TOKENS=32768 PROMPT_BUDGET_TOKENS=16000 RESPONSE_RESERVE_TOKENS=4096 +# static: bind every role-appropriate tool on each call (stable, cacheable prompt +# prefix). dynamic: keyword-route tool groups per call with a discovery tool. +TOOL_ROUTING_MODE=static # Optional durable summary sidecar; defaults to CHECKPOINT_DB_PATH + .summaries.db # Store on the same persistent shared volume as the checkpoints. # CONVERSATION_SUMMARY_DB_PATH=/app/checkpoints/conversation-summaries.db diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index 323711a..e83cd0a 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -75,6 +75,7 @@ def _model_call_summary(call: dict) -> dict: "prefill_ms": round((first_chunk - start) * 1000) if first_chunk else None, "ms": round((end - start) * 1000) if end else None, "prompt_tokens": call.get("prompt_tokens"), + "cached_prompt_tokens": call.get("cached_prompt_tokens"), "output_tokens": call.get("output_tokens"), "model": call.get("model"), "provider": call.get("provider"), @@ -600,6 +601,10 @@ async def stream_query(): if usage: call["prompt_tokens"] = usage.get("input_tokens") call["output_tokens"] = usage.get("output_tokens") + # Prompt-cache hits, when the provider reports them. + call["cached_prompt_tokens"] = ( + usage.get("input_token_details") or {} + ).get("cache_read") call["reasoning_tokens"] = ( usage.get("output_token_details") or {} ).get("reasoning") diff --git a/blueprints/conversation/agents.py b/blueprints/conversation/agents.py index f52ec8e..0581346 100644 --- a/blueprints/conversation/agents.py +++ b/blueprints/conversation/agents.py @@ -29,7 +29,7 @@ from blueprints.conversation.artifacts_util import ( from blueprints.conversation.context import RESPONSE_RESERVE_TOKENS, build_context from blueprints.conversation.memory import save_memory from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT -from blueprints.conversation.tool_routing import ToolRouter +from blueprints.conversation.tool_routing import TOOL_ROUTING_MODES, ToolRouter from blueprints.rag.logic import query_vector_store from blueprints.ynab.sync import ynab_configured from blueprints.ynab.tools import YNAB_TOOLS @@ -1160,16 +1160,24 @@ def build_agent(checkpointer=None, extra_tools=()): """ extra_tools = list(extra_tools) # Keep the guarded Google tools executable for admins even when no MCP tools - # are configured and app.admin_agent aliases app.main_agent. The router never - # binds their schemas for a regular user. Extra tools exist only on the admin - # graph and are independently gated by the same invocation config. + # are configured and app.admin_agent aliases app.main_agent. Neither routing + # mode binds their schemas for a regular user. Extra tools exist only on the + # admin graph and are independently gated by the same invocation config. agent_tools = tools + extra_tools router = ToolRouter(agent_tools, admin_tools=extra_tools) + mode = os.getenv("TOOL_ROUTING_MODE", "static").strip().lower() + if mode not in TOOL_ROUTING_MODES: + raise ValueError( + f"TOOL_ROUTING_MODE must be one of {TOOL_ROUTING_MODES}, got {mode!r}" + ) + dynamic = mode == "dynamic" def routed_model(messages, config: RunnableConfig): - selected = router.select( - messages, - is_admin=(config or {}).get("configurable", {}).get("is_admin", False), + is_admin = (config or {}).get("configurable", {}).get("is_admin", False) + selected = ( + router.select(messages, is_admin=is_admin) + if dynamic + else router.select_static(is_admin=is_admin) ) schemas = [convert_to_openai_tool(t) for t in selected] fallback = openai_fallback.bind_tools(selected) @@ -1188,6 +1196,6 @@ def build_agent(checkpointer=None, extra_tools=()): return create_react_agent( model=select_model, - tools=agent_tools + [router.discovery], + tools=agent_tools + ([router.discovery] if dynamic else []), checkpointer=checkpointer, ) diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py index f90593b..7a5821d 100644 --- a/blueprints/conversation/context.py +++ b/blueprints/conversation/context.py @@ -29,6 +29,8 @@ EARLIER_MESSAGE_MAX_CHARS = 4000 TOOL_RESULT_MAX_CHARS = 12000 TRUNCATION_NOTE = "\n[truncated: earlier context]" IMAGE_PLACEHOLDER = "[image from an earlier message]" +# On overflow, drop old turns down to about this share of the history budget. +TRIM_LOW_WATERMARK = 0.65 def estimate_tokens(value) -> int: @@ -340,46 +342,48 @@ def _summary(turns: list[list[BaseMessage]], limit: int = 2400) -> AIMessage | N ) -def trim_history( - messages: list[BaseMessage], - *, - token_budget: int = 12000, - summary_message: AIMessage | None = None, -) -> list[BaseMessage]: - """Keep the current turn, then recent whole turns, within an estimated budget. +def _stable_cut(costs: list[int], room: int, step: int) -> int: + """Return how many of the oldest prior turns to drop so the rest fit ``room``. - The current user message and tool call arguments are never truncated. Fail - explicitly if these mandatory inputs cannot fit, rather than lose the task. + Cut points sit where the cumulative cost from the oldest turn crosses a + multiple of ``step``. They depend only on turns already in history, so the + cut stays put as new turns arrive and only moves, by about ``step`` tokens, + once the retained turns overflow again. That keeps the prompt prefix + byte-identical across several turns for llama-server's prompt cache. """ - if summary_message is not None: - if not isinstance(summary_message, AIMessage) or summary_message.tool_calls: - raise ValueError( - "Conversation summary must be assistant-role evidence without tool calls" - ) - summary_cost = message_tokens(summary_message) - if summary_cost > token_budget: - logging.info("Omitting optional conversation summary: insufficient budget") - return trim_history(messages, token_budget=token_budget) - try: - trimmed = trim_history(messages, token_budget=token_budget - summary_cost) - except ValueError: - # Historical evidence must never evict a current request that fits - # without it. Required current-turn metadata retains priority. - logging.info( - "Omitting optional conversation summary to preserve current request" - ) - return trim_history(messages, token_budget=token_budget) - if ( - trimmed - and isinstance(trimmed[0], AIMessage) - and isinstance(trimmed[0].content, str) - and trimmed[0].content.startswith("Earlier conversation excerpts") - ): - trimmed = trimmed[1:] - return [summary_message, *trimmed] + remaining = sum(costs) + if remaining <= room: + return 0 + total = 0 + for i, turn_cost in enumerate(costs): + before = total + total += turn_cost + remaining -= turn_cost + if total // step > before // step and remaining <= room: + return i + 1 + # No cut point leaves room for any prior turn (e.g. a heavy current turn). + # Fall back to keeping as many recent whole turns as fit. + remaining = sum(costs) + for i, turn_cost in enumerate(costs): + if remaining <= room: + return i + remaining -= turn_cost + return len(costs) + + +def _trim( + messages: list[BaseMessage], token_budget: int, reserve: int = 0 +) -> tuple[list[BaseMessage], int]: + """Trim to ``token_budget - reserve``; return the messages and dropped turns. + + ``reserve`` shrinks the space without moving the hysteresis cut points, so a + caller-supplied summary of varying size doesn't shift which turns are kept. + """ + step = max(1, int(token_budget * (1 - TRIM_LOW_WATERMARK))) + token_budget -= reserve turns = [_complete_exchanges(turn) for turn in _turns(messages)] if not turns: - return [] + return [], 0 original_current = turns[-1] current = original_current current = [ @@ -421,19 +425,21 @@ def trim_history( else: high = mid - 1 current = compact(low) - selected = [] used = cost(current) - prior = turns[:-1] - # Leave a small share for excerpts of older turns only when necessary. - summary_reserve = min(1000, max(0, token_budget - used)) - for turn in reversed(prior): - shrunk = [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn] - turn_cost = cost(shrunk) - if used + turn_cost > token_budget - summary_reserve: - break - selected.insert(0, shrunk) - used += turn_cost - dropped = prior[: len(prior) - len(selected)] + prior = [ + [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn] + for turn in turns[:-1] + ] + costs = [cost(turn) for turn in prior] + if used + sum(costs) <= token_budget: + cut = 0 + else: + # Leave a small share for excerpts of older turns only when dropping. + summary_reserve = min(1000, max(0, token_budget - used)) + cut = _stable_cut(costs, token_budget - summary_reserve - used, step) + selected = prior[cut:] + used += sum(costs[cut:]) + dropped = turns[:cut] summary = _summary(dropped) result = [m for turn in selected for m in turn] + current if summary and used + message_tokens(summary) <= token_budget: @@ -451,7 +457,53 @@ def trim_history( } ) ) - return result + return result, len(dropped) + + +def trim_history( + messages: list[BaseMessage], + *, + token_budget: int = 12000, + summary_message: AIMessage | None = None, +) -> list[BaseMessage]: + """Keep the current turn, then recent whole turns, within an estimated budget. + + The current user message and tool call arguments are never truncated. Fail + explicitly if these mandatory inputs cannot fit, rather than lose the task. + When the budget overflows, old turns are dropped down to a lower watermark + so the retained prefix stays stable (and cacheable) for several turns. + """ + if summary_message is not None: + if not isinstance(summary_message, AIMessage) or summary_message.tool_calls: + raise ValueError( + "Conversation summary must be assistant-role evidence without tool calls" + ) + trimmed, dropped = _trim(messages, token_budget) + if summary_message is None or not dropped: + # The summary only duplicates retained turns when nothing was dropped, + # and it changes every turn, which would defeat prompt caching. + return trimmed + summary_cost = message_tokens(summary_message) + if summary_cost > token_budget: + logging.info("Omitting optional conversation summary: insufficient budget") + return trimmed + try: + trimmed, _ = _trim(messages, token_budget, reserve=summary_cost) + except ValueError: + # Historical evidence must never evict a current request that fits + # without it. Required current-turn metadata retains priority. + logging.info( + "Omitting optional conversation summary to preserve current request" + ) + return trimmed + if ( + trimmed + and isinstance(trimmed[0], AIMessage) + and isinstance(trimmed[0].content, str) + and trimmed[0].content.startswith("Earlier conversation excerpts") + ): + trimmed = trimmed[1:] + return [summary_message, *trimmed] def build_context( diff --git a/blueprints/conversation/memory.py b/blueprints/conversation/memory.py index 82fa376..baf9fb1 100644 --- a/blueprints/conversation/memory.py +++ b/blueprints/conversation/memory.py @@ -1,3 +1,5 @@ +from collections import Counter + from .models import UserMemory from .prompts import SIMBA_SYSTEM_PROMPT from .context import estimate_tokens @@ -41,16 +43,23 @@ async def build_system_prompt_with_memories(user_id: str, query: str = "") -> st memories = await get_memories_for_user(user_id) if not memories: return SIMBA_SYSTEM_PROMPT - selected = [] header = "\n\nUSER MEMORIES (facts the user has asked you to remember):\n" omission = "\n[Additional saved memories omitted from this prompt.]" # Reserve wrapper and omission costs so the entire injected memory block fits. used = estimate_tokens(header) + estimate_tokens(omission) + chosen = Counter() for memory in rank_memories(memories, query): cost = estimate_tokens(f"- {memory}\n") if used + cost <= MEMORY_BUDGET_TOKENS: - selected.append(memory) + chosen[memory] += 1 used += cost + # Rank only chooses what fits; emit in stored order so the system prompt is + # byte-identical across queries and llama-server can reuse its KV cache. + selected = [] + for memory in memories: + if chosen[memory]: + chosen[memory] -= 1 + selected.append(memory) block = "\n".join(f"- {m}" for m in selected) if len(selected) < len(memories): block += omission diff --git a/blueprints/conversation/tool_routing.py b/blueprints/conversation/tool_routing.py index 2388278..57026d5 100644 --- a/blueprints/conversation/tool_routing.py +++ b/blueprints/conversation/tool_routing.py @@ -1,4 +1,8 @@ -"""Deterministic per-call routing; discovery grants live in tool messages, not globals.""" +"""Tool selection per model call. + +Static mode (the default) binds the full role-appropriate set. Dynamic mode +routes per call by keyword; discovery grants live in tool messages, not globals. +""" import json import re @@ -29,6 +33,7 @@ ADMIN_TOOL_NAMES = { } CORE_TOOL_NAMES = {"get_current_date", "save_user_memory"} DISCOVERY_NAME = "discover_tool_groups" +TOOL_ROUTING_MODES = ("static", "dynamic") def words(text): @@ -138,6 +143,16 @@ class ToolRouter: ) return discover_tool_groups + def select_static(self, *, is_admin=False): + """Every role-appropriate tool in registration order, without discovery. + + The bound set never varies by query or turn, so the rendered tool block + stays byte-identical and llama-server can reuse its cached prefix. + """ + return [ + t for t in self.tools if is_admin or tool_name(t) not in self.admin_names + ] + def select(self, messages, *, is_admin=False): available = self.available_groups(is_admin) human_indices = [ diff --git a/docker-compose.yml b/docker-compose.yml index 59bde2f..c33fe75 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,6 +15,7 @@ x-raggr-common: &raggr-common - JWT_SECRET_KEY=${JWT_SECRET_KEY} - LLAMA_SERVER_URL=${LLAMA_SERVER_URL} - LLAMA_MODEL_NAME=${LLAMA_MODEL_NAME} + - TOOL_ROUTING_MODE=${TOOL_ROUTING_MODE:-static} - OIDC_ISSUER=${OIDC_ISSUER} - OIDC_CLIENT_ID=${OIDC_CLIENT_ID} - OIDC_CLIENT_SECRET=${OIDC_CLIENT_SECRET} diff --git a/tests/unit/test_context_summary.py b/tests/unit/test_context_summary.py index 634ff78..f83a34b 100644 --- a/tests/unit/test_context_summary.py +++ b/tests/unit/test_context_summary.py @@ -102,7 +102,8 @@ async def test_tools_and_incomplete_exchanges_do_not_become_summary_instructions assert "Now continue" not in message.content prompt = build_context(messages, "trusted system", [], summary_message=message) assert prompt[0].content == "trusted system" - assert prompt[1] == message + # Nothing was dropped, so the ledger would only duplicate retained turns. + assert message not in prompt assert prompt[-1] is messages[-1] diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py index cbbf825..4dbeda2 100644 --- a/tests/unit/test_context_trim.py +++ b/tests/unit/test_context_trim.py @@ -462,3 +462,52 @@ def test_focused_note_heading_index_uses_absolute_content_lines(): metadata = json.loads(projected.content)["data"]["markdown_compaction"] assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}] assert msg.artifact["data"]["markdown"] == "full original" + + +def _conversation(turns, size=1000): + messages = [] + for i in range(turns): + messages += [ + HumanMessage(content=f"Question {i}"), + call(str(i)), + result(str(i), "data " * size), + AIMessage(content=f"Decision {i}"), + ] + return messages + + +def test_summary_omitted_when_no_turns_dropped(): + messages = _conversation(2, size=10) + [HumanMessage(content="next")] + summary = AIMessage(content="Earlier: keep dining low") + assert trim_history(messages, summary_message=summary) == trim_history(messages) + assert summary not in trim_history(messages, summary_message=summary) + + +def test_summary_prepended_when_turns_dropped(): + messages = _conversation(12) + [HumanMessage(content="next")] + summary = AIMessage(content="Earlier: keep dining low") + out = trim_history(messages, token_budget=6000, summary_message=summary) + assert out[0] is summary + assert out[-1] is messages[-1] + assert cost(out) <= 6000 + + +def test_trim_drops_to_low_watermark_and_prefix_stays_stable(): + budget = 8000 + messages, oldest, costs = [], [], [] + for i in range(30): + messages.append(HumanMessage(content=f"Question {i}")) + out = trim_history(messages, token_budget=budget) + assert out[-1] is messages[-1] + assert cost(out) <= budget + oldest.append(next(m for m in out if isinstance(m, HumanMessage)).content) + costs.append(cost(out)) + messages += [call(str(i)), result(str(i), "data " * 300), AIMessage("ok")] + cuts = [i for i in range(1, 30) if oldest[i] != oldest[i - 1]] + assert cuts, "history should overflow the budget" + # Each drop goes well below the budget instead of one turn under it... + for i in cuts: + assert costs[i] <= budget * (context.TRIM_LOW_WATERMARK + 0.1) + # ...so the retained prefix is unchanged for several turns in between. + assert all(b - a >= 3 for a, b in zip(cuts, cuts[1:])) + assert len(cuts) <= 5 diff --git a/tests/unit/test_memory_retrieval.py b/tests/unit/test_memory_retrieval.py index 8f54ed3..dcc1249 100644 --- a/tests/unit/test_memory_retrieval.py +++ b/tests/unit/test_memory_retrieval.py @@ -46,3 +46,34 @@ async def test_backwards_compatible_no_query_and_empty_memories(monkeypatch): ) fetch.return_value = ["Likes tea"] assert "Likes tea" in await memory.build_system_prompt_with_memories("user") + + +@pytest.mark.asyncio +async def test_memory_block_is_identical_across_queries_when_all_fit(monkeypatch): + memories = ["Favorite restaurant is Nori", "Cat Simba takes insulin", "Gym is Y"] + monkeypatch.setattr( + memory, "get_memories_for_user", AsyncMock(return_value=memories) + ) + prompts = { + await memory.build_system_prompt_with_memories("user", query=q) + for q in ["", "Simba insulin", "gym", "Nori restaurant"] + } + assert len(prompts) == 1 + prompt = prompts.pop() + assert prompt.index("Nori") < prompt.index("insulin") < prompt.index("Gym") + + +@pytest.mark.asyncio +async def test_rank_chooses_memories_but_stored_order_is_kept(monkeypatch): + filler = "Favorite restaurant is Nori. " * 20 + memories = ["Likes tea", filler, "Simba insulin dose is two units."] + monkeypatch.setattr( + memory, "get_memories_for_user", AsyncMock(return_value=memories) + ) + monkeypatch.setattr(memory, "MEMORY_BUDGET_TOKENS", 110) + prompt = await memory.build_system_prompt_with_memories( + "user", query="Simba insulin" + ) + assert filler not in prompt + # Insulin ranks first but is emitted after the newer memory. + assert prompt.index("Likes tea") < prompt.index("insulin") diff --git a/tests/unit/test_stream_stall.py b/tests/unit/test_stream_stall.py index 1381d56..ee9599e 100644 --- a/tests/unit/test_stream_stall.py +++ b/tests/unit/test_stream_stall.py @@ -61,7 +61,7 @@ def _chunk(run_id, text): } -def _end(run_id, prompt_tokens=None): +def _end(run_id, prompt_tokens=None, cached_tokens=None): usage = ( { "input_tokens": prompt_tokens, @@ -71,6 +71,8 @@ def _end(run_id, prompt_tokens=None): if prompt_tokens else None ) + if usage and cached_tokens is not None: + usage["input_token_details"] = {"cache_read": cached_tokens} return { "event": "on_chat_model_end", "run_id": run_id, @@ -207,7 +209,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog): script = [ (0, _start("m1")), (0, _chunk("m1", "Hi")), - (0, _end("m1", prompt_tokens=27262)), + (0, _end("m1", prompt_tokens=27262, cached_tokens=25000)), (0, _final("Hi")), ] @@ -224,6 +226,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog): assert summary["max_prompt_tokens"] == 27262 (call,) = summary["model_calls"] assert call["prompt_tokens"] == 27262 + assert call["cached_prompt_tokens"] == 25000 assert call["prefill_ms"] is not None and call["ms"] is not None diff --git a/tests/unit/test_tool_routing.py b/tests/unit/test_tool_routing.py index dfb3197..d28cbd9 100644 --- a/tests/unit/test_tool_routing.py +++ b/tests/unit/test_tool_routing.py @@ -136,11 +136,7 @@ def test_discovery_recovery_admin_gating_and_isolation(router): ) -@pytest.mark.asyncio -@pytest.mark.parametrize("primary_fails", [False, True]) -async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset( - monkeypatch, primary_fails -): +def import_agents(): # Constructors are stubbed before agents import: no tokens, services or network. with ( patch("langchain_openai.ChatOpenAI"), @@ -154,7 +150,16 @@ async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset( "utils.obsidian_service.ObsidianService", side_effect=ValueError("offline") ), ): - agents = importlib.import_module("blueprints.conversation.agents") + return importlib.import_module("blueprints.conversation.agents") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("primary_fails", [False, True]) +async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset( + monkeypatch, primary_fails +): + monkeypatch.setenv("TOOL_ROUTING_MODE", "dynamic") + agents = import_agents() bound_names, primary_names, accounted_names = [], [], [] class FakeModel: @@ -238,3 +243,90 @@ def test_schema_dict_tools_remain_routable_and_admin_gated(): router = ToolRouter([extra], admin_tools=[extra]) assert extra not in router.select([HumanMessage(content="extra")]) assert extra in router.select([HumanMessage(content="extra")], is_admin=True) + + +def test_static_selection_is_stable_role_gated_and_has_no_discovery(router): + queries = ["Review spending", "Find Simba's vet records", "Help", "calendar"] + user = [router.select_static() for _ in queries] + assert all(selected == user[0] for selected in user) + assert [t.name for t in user[0]] == [ + "get_current_date", + "save_user_memory", + "ynab_budget_summary", + "mealie_get_recipe", + "obsidian_read_note", + "simba_search", + "web_search", + ] + admin = [t.name for t in router.select_static(is_admin=True)] + assert admin == [t.name for t in router.tools] + assert "discover_tool_groups" not in admin + + +def test_static_selection_gates_schema_dict_admin_tools(): + extra = { + "type": "function", + "function": { + "name": "extra", + "description": "docs", + "parameters": {"type": "object", "properties": {}}, + }, + } + router = ToolRouter([make_tool("web_search"), extra], admin_tools=[extra]) + assert extra not in router.select_static() + assert extra in router.select_static(is_admin=True) + + +@pytest.mark.asyncio +async def test_static_graph_binds_full_role_set_every_call(monkeypatch): + monkeypatch.delenv("TOOL_ROUTING_MODE", raising=False) + agents = import_agents() + bound_names = [] + + class FakeModel: + def bind_tools(self, tools): + bound_names.append([t.name for t in tools]) + + def respond(messages): + if len(bound_names) == 1: + return call("ynab_budget_summary") + return AIMessage(content="done") + + return RunnableLambda(respond) + + monkeypatch.setattr( + agents, + "tools", + [ + make_tool("get_current_date"), + make_tool("ynab_budget_summary"), + make_tool("get_calendar_events"), + ], + ) + monkeypatch.setattr(agents, "llama_chat", None) + monkeypatch.setattr(agents, "openai_fallback", FakeModel()) + monkeypatch.setattr(agents, "build_context", lambda messages, *a, **k: messages) + graph = agents.build_agent() + # A query matching no keyword group still sees every tool, with no + # discovery round trip, and the bound set is identical on every call. + result = await graph.ainvoke({"messages": [HumanMessage(content="Help")]}) + assert result["messages"][-1].content == "done" + assert bound_names == [["get_current_date", "ynab_budget_summary"]] * 2 + + bound_names.clear() + await graph.ainvoke( + {"messages": [HumanMessage(content="Help")]}, + config={"configurable": {"is_admin": True}}, + ) + assert bound_names[0] == [ + "get_current_date", + "ynab_budget_summary", + "get_calendar_events", + ] + + +def test_unknown_routing_mode_fails_explicitly(monkeypatch): + agents = import_agents() + monkeypatch.setenv("TOOL_ROUTING_MODE", "sometimes") + with pytest.raises(ValueError, match="TOOL_ROUTING_MODE"): + agents.build_agent() From 1e638e1e3b6b2dfa6e372dfeb5ace6abcd37a263 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Mon, 12 Oct 2026 23:04:24 -0400 Subject: [PATCH 2/3] Place conversation evidence before the current turn for prompt caching The summary ledger is regenerated every turn. At the head of history it invalidated the cached retained turns on every request once trimming started. Insert it immediately before the current turn's HumanMessage so system, tools and retained history stay byte-identical; excerpts of dropped turns stay at the head since they only change when the cut moves. Reserve a fixed ledger ceiling (SUMMARY_RESERVE_TOKENS) rather than its current size, so a shrinking ledger cannot pull the cut back and change the retained prefix; fall back to the actual size for tight budgets. Co-Authored-By: Claude Opus 5.5 --- blueprints/conversation/context.py | 30 ++++++++++++++++++----------- tests/unit/test_context_summary.py | 14 ++++++++++++++ tests/unit/test_context_trim.py | 31 ++++++++++++++++++++++++++---- 3 files changed, 60 insertions(+), 15 deletions(-) diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py index 7a5821d..a3ad73b 100644 --- a/blueprints/conversation/context.py +++ b/blueprints/conversation/context.py @@ -31,6 +31,8 @@ TRUNCATION_NOTE = "\n[truncated: earlier context]" IMAGE_PLACEHOLDER = "[image from an earlier message]" # On overflow, drop old turns down to about this share of the history budget. TRIM_LOW_WATERMARK = 0.65 +# Upper bound on the summary ledger message (summary.STATE_TOKEN_LIMIT + wrapper). +SUMMARY_RESERVE_TOKENS = 2000 def estimate_tokens(value) -> int: @@ -487,23 +489,29 @@ def trim_history( if summary_cost > token_budget: logging.info("Omitting optional conversation summary: insufficient budget") return trimmed - try: - trimmed, _ = _trim(messages, token_budget, reserve=summary_cost) - except ValueError: + # Reserve a fixed ceiling, not the ledger's current size, so a shrinking + # ledger can't pull the cut back and change the retained prefix. + for reserve in dict.fromkeys( + (max(summary_cost, SUMMARY_RESERVE_TOKENS), summary_cost) + ): + try: + with_room, _ = _trim(messages, token_budget, reserve=reserve) + break + except ValueError: + continue + else: # Historical evidence must never evict a current request that fits # without it. Required current-turn metadata retains priority. logging.info( "Omitting optional conversation summary to preserve current request" ) return trimmed - if ( - trimmed - and isinstance(trimmed[0], AIMessage) - and isinstance(trimmed[0].content, str) - and trimmed[0].content.startswith("Earlier conversation excerpts") - ): - trimmed = trimmed[1:] - return [summary_message, *trimmed] + # The ledger changes every turn, so place it just before the current turn: + # system, tools and retained history then stay byte-identical for the + # prompt cache. Turns start at a HumanMessage, so the last one opens the + # current turn and no tool exchange is split. + start = max(i for i, m in enumerate(with_room) if isinstance(m, HumanMessage)) + return [*with_room[:start], summary_message, *with_room[start:]] def build_context( diff --git a/tests/unit/test_context_summary.py b/tests/unit/test_context_summary.py index f83a34b..bcc763e 100644 --- a/tests/unit/test_context_summary.py +++ b/tests/unit/test_context_summary.py @@ -163,3 +163,17 @@ async def test_deletion_is_scoped_and_survives_reopen(tmp_path): await store.delete("owner", "c") assert (await SummaryStore(path).load("owner", "c")).cursor == 0 assert (await SummaryStore(path).load("other", "c")).cursor == 4 + + +async def test_largest_ledger_fits_the_fixed_summary_reserve(tmp_path): + from blueprints.conversation.context import SUMMARY_RESERVE_TOKENS, message_tokens + + messages = [] + for i in range(100): + messages += [ + HumanMessage(content=f"Must {i}: " + "猫😺" * 1000), + AIMessage(content=f"Decided {i}: " + "budget " * 1000), + ] + messages.append(HumanMessage(content="current")) + state = await SummaryStore(tmp_path / "summary.db").update("u", "c", messages) + assert message_tokens(state.as_message()) <= SUMMARY_RESERVE_TOKENS diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py index 4dbeda2..e30e2b8 100644 --- a/tests/unit/test_context_trim.py +++ b/tests/unit/test_context_trim.py @@ -483,13 +483,36 @@ def test_summary_omitted_when_no_turns_dropped(): assert summary not in trim_history(messages, summary_message=summary) -def test_summary_prepended_when_turns_dropped(): - messages = _conversation(12) + [HumanMessage(content="next")] +def assert_valid_sequence(out): + seen = set() + for m in out: + if isinstance(m, AIMessage): + seen.update(c["id"] for c in m.tool_calls) + if isinstance(m, ToolMessage): + assert m.tool_call_id in seen + + +def test_summary_placed_before_current_turn_when_turns_dropped(): + question = HumanMessage(content="next") + # The current turn has its own tool exchange the summary must not split. + messages = _conversation(12) + [question, call("now"), result("now")] summary = AIMessage(content="Earlier: keep dining low") out = trim_history(messages, token_budget=6000, summary_message=summary) - assert out[0] is summary - assert out[-1] is messages[-1] assert cost(out) <= 6000 + assert out[-4:] == [summary, question, messages[-2], messages[-1]] + assert out.count(summary) == 1 + # Excerpts of dropped turns stay at the head; they only change with the cut. + assert out[0].content.startswith("Earlier conversation excerpts") + assert_valid_sequence(out) + + +def test_changing_summary_leaves_retained_prefix_identical(): + messages = _conversation(12) + [HumanMessage(content="next")] + short = AIMessage(content="Earlier: keep dining low") + long = AIMessage(content="Earlier: " + "constraint " * 400) + a = trim_history(messages, token_budget=6000, summary_message=short) + b = trim_history(messages, token_budget=6000, summary_message=long) + assert a[: a.index(short)] == b[: b.index(long)] def test_trim_drops_to_low_watermark_and_prefix_stays_stable(): From 266c6a8d7919675199f47ac747684e8b551355d9 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Mon, 12 Oct 2026 23:05:21 -0400 Subject: [PATCH 3/3] Drop dropped-turn excerpts when the summary ledger is included The ledger supersedes the "Earlier conversation excerpts" message, as on main; sending both duplicated dropped-turn content. Removal is consistent whenever the ledger is included, so the cached prefix stays stable. Co-Authored-By: Claude Opus 5.5 --- blueprints/conversation/context.py | 9 +++++++++ tests/unit/test_context_trim.py | 12 ++++++++++-- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py index a3ad73b..7f10d7f 100644 --- a/blueprints/conversation/context.py +++ b/blueprints/conversation/context.py @@ -506,6 +506,15 @@ def trim_history( "Omitting optional conversation summary to preserve current request" ) return trimmed + # The ledger supersedes excerpts of dropped turns; whenever it is included + # they are consistently absent, so this doesn't disturb the cached prefix. + if ( + with_room + and isinstance(with_room[0], AIMessage) + and isinstance(with_room[0].content, str) + and with_room[0].content.startswith("Earlier conversation excerpts") + ): + with_room = with_room[1:] # The ledger changes every turn, so place it just before the current turn: # system, tools and retained history then stay byte-identical for the # prompt cache. Turns start at a HumanMessage, so the last one opens the diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py index e30e2b8..762c2ae 100644 --- a/tests/unit/test_context_trim.py +++ b/tests/unit/test_context_trim.py @@ -501,8 +501,16 @@ def test_summary_placed_before_current_turn_when_turns_dropped(): assert cost(out) <= 6000 assert out[-4:] == [summary, question, messages[-2], messages[-1]] assert out.count(summary) == 1 - # Excerpts of dropped turns stay at the head; they only change with the cut. - assert out[0].content.startswith("Earlier conversation excerpts") + # The ledger supersedes excerpts of dropped turns; don't send both. + assert not any( + isinstance(m.content, str) + and m.content.startswith("Earlier conversation excerpts") + for m in out + ) + assert isinstance(out[0], HumanMessage) + # Without a ledger, the excerpts still cover the dropped turns. + plain = trim_history(messages, token_budget=6000) + assert plain[0].content.startswith("Earlier conversation excerpts") assert_valid_sequence(out)