From af06308425a01805584fff9cdf0e909657aa5266 Mon Sep 17 00:00:00 2001 From: kshitij <82637225+kshitijk4poor@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:17:50 +0530 Subject: [PATCH] refactor(prompt-caching): collapse triplicated destination-plan and label parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three copies of the same logic landed with #76032: - MoA's _call_prepared_aggregator and auxiliary_client's _replan_synchronous_cache_sections both implemented stub → policy → strip → plan for a resolved destination. Extract plan_cache_sections_for_destination() into agent_runtime_helpers (which already owns the policy functions) and route both through it. Also removes a redundant full-transcript deepcopy+strip per request (the caller pre-stripped what build_prompt_cache_plan strips again). - The fallback_chain[N] label regex + chain-entry lookup lived in _fallback_entry_timeout AND _fallback_destination. Extract _fallback_chain_entry() and reuse. MoA's cache-plan failure log is promoted debug → warning: the call-block site skips MoA, so this block is the aggregator's only decoration path — a silent failure ships an undecorated request (the 0%-cache MoA bug class). Behavior-preserving; 195 targeted tests green. Follow-up to #76032 (#20880). --- agent/agent_runtime_helpers.py | 56 +++++++++++++++++++++ agent/auxiliary_client.py | 90 +++++++++++++--------------------- agent/moa_loop.py | 55 +++++++-------------- 3 files changed, 108 insertions(+), 93 deletions(-) diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 2aba754b5fd7c..072fe507f3323 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1857,6 +1857,62 @@ def _direct_native_anthropic_tool_cache_capability( ) +def plan_cache_sections_for_destination( + messages: list, + tools: Optional[list], + *, + provider: str, + base_url: str, + api_mode: str, + model: str, +) -> Tuple[list, list]: + """Plan request-local cache sections for one resolved destination. + + Shared core of the synchronous acting-aggregator (MoA) and auxiliary + fallback senders: resolve the cache policy for the destination's real + provider/base_url/api_mode/model, then either return stripped canonical + copies (non-caching route) or a :func:`build_prompt_cache_plan` layout + (caching route, with the direct-native tool marker when the destination + is api.anthropic.com on the Messages wire). + + Never mutates ``messages`` or ``tools`` — both return values are + request-local copies. + """ + from types import SimpleNamespace + + from agent.prompt_caching import ( + build_prompt_cache_plan, + strip_anthropic_cache_control, + strip_anthropic_tool_cache_control, + ) + + stub = SimpleNamespace(provider="", base_url="", api_mode="", model="") + should_cache, native_layout = anthropic_prompt_cache_policy( + stub, + provider=provider, + base_url=base_url, + api_mode=api_mode, + model=model, + ) + if not should_cache: + canonical_messages = copy.deepcopy(messages or []) + strip_anthropic_cache_control(canonical_messages) + return canonical_messages, strip_anthropic_tool_cache_control(tools) + plan = build_prompt_cache_plan( + messages, + tools, + native_anthropic=native_layout, + direct_native_tool_cache=_direct_native_anthropic_tool_cache_capability( + stub, + provider=provider, + base_url=base_url, + api_mode=api_mode, + model=model, + ), + ) + return plan.messages, plan.tools + + def anthropic_prompt_cache_policy( agent, *, diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 9d26d1bf611b4..60ebe1b552e4e 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -4181,6 +4181,27 @@ def _auth_refresh_provider_for_route( return normalized +def _fallback_chain_entry(task: Optional[str], fb_label: str) -> Optional[Dict[str, Any]]: + """Resolve the configured ``fallback_chain`` entry a label points at. + + Labels minted by :func:`_try_configured_fallback_chain` carry the entry + index in our own stable format (``fallback_chain[]()``). + Returns ``None`` when the label is not a configured-chain candidate or + the index no longer resolves to a dict entry. + """ + if not task or not fb_label: + return None + m = re.match(r"fallback_chain\[(\d+)\]", fb_label) + if not m: + return None + try: + chain = _get_auxiliary_task_config(task).get("fallback_chain") + entry = chain[int(m.group(1))] if isinstance(chain, list) else None + except Exception: + return None + return entry if isinstance(entry, dict) else None + + def _fallback_entry_timeout(task: Optional[str], fb_label: str) -> Optional[float]: """Resolve a per-entry ``timeout`` for a configured fallback candidate. @@ -4192,24 +4213,13 @@ def _fallback_entry_timeout(task: Optional[str], fb_label: str) -> Optional[floa primary's 30s deadline every turn (#62452). Entries in ``auxiliary..fallback_chain`` may declare their own - ``timeout`` (seconds). This helper reads it by parsing the entry index - out of the label minted by :func:`_try_configured_fallback_chain` - (``fallback_chain[]()`` — our own stable format). Returns - ``None`` when the label is not a configured-chain candidate, the entry - has no ``timeout``, or the value is invalid — callers then keep the - task-level timeout, preserving existing behavior. + ``timeout`` (seconds). Returns ``None`` when the label is not a + configured-chain candidate, the entry has no ``timeout``, or the value + is invalid — callers then keep the task-level timeout, preserving + existing behavior. """ - if not task or not fb_label: - return None - m = re.match(r"fallback_chain\[(\d+)\]", fb_label) - if not m: - return None - try: - chain = _get_auxiliary_task_config(task).get("fallback_chain") - entry = chain[int(m.group(1))] if isinstance(chain, list) else None - raw = entry.get("timeout") if isinstance(entry, dict) else None - except Exception: - return None + entry = _fallback_chain_entry(task, fb_label) + raw = entry.get("timeout") if entry else None if isinstance(raw, (int, float)) and not isinstance(raw, bool) and raw > 0: return float(raw) return None @@ -4284,15 +4294,9 @@ def _fallback_destination( api_mode = None model = fb_model - match = re.match(r"fallback_chain\[(\d+)\]", fb_label or "") - if match and task: - try: - chain = _get_auxiliary_task_config(task).get("fallback_chain") - entry = chain[int(match.group(1))] if isinstance(chain, list) else None - except Exception: - entry = None - if isinstance(entry, dict): - return _fallback_destination_from_entry(entry, fb_client, fb_model) + entry = _fallback_chain_entry(task, fb_label) + if entry is not None: + return _fallback_destination_from_entry(entry, fb_client, fb_model) return _complete_fallback_destination(provider, base_url, api_mode, model) @@ -4304,42 +4308,16 @@ def _replan_synchronous_cache_sections( destination: _FallbackDestination, ) -> tuple[list, list]: """Strip source decoration and plan one synchronous destination locally.""" - from agent.agent_runtime_helpers import ( - _direct_native_anthropic_tool_cache_capability, - anthropic_prompt_cache_policy, - ) - from agent.prompt_caching import ( - build_prompt_cache_plan, - strip_anthropic_cache_control, - strip_anthropic_tool_cache_control, - ) + from agent.agent_runtime_helpers import plan_cache_sections_for_destination - canonical_messages = copy.deepcopy(messages or []) - strip_anthropic_cache_control(canonical_messages) - canonical_tools = strip_anthropic_tool_cache_control(tools) - stub = SimpleNamespace(provider="", base_url="", api_mode="", model="") - should_cache, native_layout = anthropic_prompt_cache_policy( - stub, + return plan_cache_sections_for_destination( + messages, + tools, provider=destination.provider, base_url=destination.base_url, api_mode=destination.api_mode or "", model=destination.model or "", ) - if not should_cache: - return canonical_messages, canonical_tools - plan = build_prompt_cache_plan( - canonical_messages, - canonical_tools, - native_anthropic=native_layout, - direct_native_tool_cache=_direct_native_anthropic_tool_cache_capability( - stub, - provider=destination.provider, - base_url=destination.base_url, - api_mode=destination.api_mode or "", - model=destination.model or "", - ), - ) - return plan.messages, plan.tools def _call_fallback_candidate_sync( diff --git a/agent/moa_loop.py b/agent/moa_loop.py index 671827e7fd81e..e149ea9f08ca6 100644 --- a/agent/moa_loop.py +++ b/agent/moa_loop.py @@ -8,7 +8,6 @@ iteration. from __future__ import annotations -import copy import hashlib import logging import re @@ -1660,57 +1659,39 @@ class MoAChatCompletions: extra_body: Any = agg_kwargs.get("extra_body") agg_runtime = _slot_runtime(aggregator) try: - from types import SimpleNamespace - from agent.agent_runtime_helpers import ( - _direct_native_anthropic_tool_cache_capability, - anthropic_prompt_cache_policy, - ) - from agent.prompt_caching import ( - build_prompt_cache_plan, - strip_anthropic_cache_control, - strip_anthropic_tool_cache_control, + plan_cache_sections_for_destination, ) guidance = prepared.get("guidance") - canonical_messages = copy.deepcopy(agg_messages) + planning_messages = agg_messages if guidance: - canonical_messages = peel_reference_guidance( - canonical_messages, + planning_messages = peel_reference_guidance( + agg_messages, str(guidance), ) - strip_anthropic_cache_control(canonical_messages) - canonical_tools = strip_anthropic_tool_cache_control(tools) - cache_stub = SimpleNamespace(provider="", base_url="", api_mode="", model="") - should_cache, native_layout = anthropic_prompt_cache_policy( - cache_stub, + # plan_cache_sections_for_destination never mutates its inputs + # and always returns request-local copies, so the prepared + # state stays canonical. + agg_messages, tools = plan_cache_sections_for_destination( + planning_messages, + tools, provider=agg_runtime.get("provider") or "", base_url=agg_runtime.get("base_url") or "", api_mode=agg_runtime.get("api_mode") or "", model=agg_runtime.get("model") or "", ) - if should_cache: - plan = build_prompt_cache_plan( - canonical_messages, - canonical_tools, - native_anthropic=native_layout, - direct_native_tool_cache=_direct_native_anthropic_tool_cache_capability( - cache_stub, - provider=agg_runtime.get("provider") or "", - base_url=agg_runtime.get("base_url") or "", - api_mode=agg_runtime.get("api_mode") or "", - model=agg_runtime.get("model") or "", - ), - ) - agg_messages = plan.messages - tools = plan.tools - else: - agg_messages = canonical_messages - tools = canonical_tools if guidance: _attach_reference_guidance(agg_messages, str(guidance)) except Exception as exc: # pragma: no cover - cache planning must not block MoA - logger.debug("MoA aggregator cache plan skipped: %s", exc) + # Warning, not debug: since the call-block site skips MoA, this + # block is the aggregator's ONLY decoration path — a silent + # failure here ships an undecorated request and regresses the + # exact 0%-cache MoA failure the planning exists to prevent. + logger.warning( + "MoA aggregator cache plan failed — sending undecorated " + "request (cache misses expected): %s", exc, + ) # Record the exact aggregator INPUT (incl. the injected reference # context) into the pending trace so a trace captures what the # aggregator actually saw, not a reconstruction. Traces are a