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..7f10d7f 100644 --- a/blueprints/conversation/context.py +++ b/blueprints/conversation/context.py @@ -29,6 +29,10 @@ 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 +# Upper bound on the summary ledger message (summary.STATE_TOKEN_LIMIT + wrapper). +SUMMARY_RESERVE_TOKENS = 2000 def estimate_tokens(value) -> int: @@ -340,46 +344,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 +427,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 +459,68 @@ 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 + # 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 + # 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 + # 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/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..bcc763e 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] @@ -162,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 cbbf825..762c2ae 100644 --- a/tests/unit/test_context_trim.py +++ b/tests/unit/test_context_trim.py @@ -462,3 +462,83 @@ 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 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 cost(out) <= 6000 + assert out[-4:] == [summary, question, messages[-2], messages[-1]] + assert out.count(summary) == 1 + # 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) + + +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(): + 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()