refactor(prompt-caching): collapse triplicated destination-plan and label parsing
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).
This commit is contained in:
kshitij
parent
7ae4a5efba
commit
af06308425
|
|
@ -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,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -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[<i>](<provider>)``).
|
||||
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.<task>.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[<i>](<provider>)`` — 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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue