diff --git a/agent/context_engine.py b/agent/context_engine.py index 2225f25473d7..b772125c0cbe 100644 --- a/agent/context_engine.py +++ b/agent/context_engine.py @@ -210,6 +210,123 @@ def prune_tool_results_only( """ return messages, 0 + # -- Optional: per-turn context selection (distinct from compression) -- + + def select_context( + self, + request_messages: List[Dict[str, Any]], + *, + conversation_messages: List[Dict[str, Any]] = None, + incoming_message: Dict[str, Any] = None, + budget_tokens: int = 0, + ) -> List[Dict[str, Any]]: + """Optionally choose/replace the context for THIS request, pre-generation. + + Called every turn after the request message list is assembled and + before it is dispatched to the provider — independent of + ``should_compress()``. This lets an engine *select* which context + enters the prompt (retrieval, topic routing, role/branch switching) + rather than *shrink* context that is already there. The two verbs are + orthogonal: + + - ``compress()`` : context is too long -> make it shorter. + - ``select_context()``: this turn belongs to a different context + -> use that one instead. + + Without this hook, engines that need per-turn access to the message + list have to force ``should_compress()`` to return ``True`` so that + ``compress()`` is invoked every turn purely as a callback — which + conflates selection with compression and degrades behaviour when the + engine's backend is unavailable. ``select_context()`` removes the need + for that workaround. + + The returned list is request-only: it replaces the messages sent to + the provider for this single call and MUST NOT be treated as persisted + transcript state. The conversation history in the session DB is left + untouched, so nothing leaks across turns. Return ``None`` to leave the + request unchanged. + + Unlike the ``pre_llm_call`` plugin hook (which appends to the user + message and intentionally never rewrites the list, to preserve the + cache prefix), ``select_context()`` may *replace* the message list. + + Ordering / cache contract: the host runs this hook **before** prompt + cache-control and **before** every request sanitizer (orphaned-tool + cleanup, thinking-only/role normalization, whitespace/JSON + normalization). So (a) whatever the hook returns still passes through + the same validation as any request — a malformed replacement cannot + reach the provider — and (b) prompt-cache stability (an AGENTS.md + invariant) is preserved: the default no-op leaves the request + byte-identical, so cache behaviour is unchanged for the built-in + compressor and any non-implementing engine. An engine that *does* + replace the list changes its own cache prefix by definition; that is + the engine's concern, and cache-control breakpoints are re-derived on + the selected list. The hook is evaluated per provider request (so it + re-runs on retries within a turn), consistent with "select the context + for THIS request". + + Args: + request_messages: The assembled request message list (system + prompt + history + any ephemeral prefill), in OpenAI format. + conversation_messages: The unmodified persisted conversation + history, for reference only (do not mutate). + incoming_message: The current turn's user message, if available. + budget_tokens: The active model's context length, or 0 if unknown. + + Default returns ``None`` (no-op) — zero impact on the built-in + compressor or any existing engine. + """ + return None + + def on_turn_complete( + self, + messages: List[Dict[str, Any]], + usage: Dict[str, Any] = None, + **kwargs: Any, + ) -> None: + """Observe a finished user turn (post-turn ingestion / observation). + + Called from the standard turn-finalization path once the assistant/tool + loop completes, with the finalized in-memory transcript snapshot. This + is the complement to ``select_context()``: selection happens *before* + the request, while observation happens *after* the turn. It lets an + engine ingest, index, summarize, or update routing / topic / session + state from what actually happened — so the next ``select_context()`` + can act on it. + + Coverage: this fires from the normal finalization seam. Some abnormal + early-return paths in the loop (e.g. a content-policy block or a + provider terminal failure) persist and return without routing through + finalization, and therefore do not currently emit this hook. Treat it + as a best-effort post-turn observation for completed turns, not a + guaranteed callback for every possible early exit; unifying all + terminal paths behind one finalization seam is a separate follow-up. + + Together the two hooks remove the need to abuse ``should_compress()`` / + ``compress()`` as a generic per-turn callback just to observe history, + and they cover the case where a turn finishes and there may be no next + request from which to infer the previous turn. + + ``messages`` is a shallow copy and should be treated as read-only: + return values are ignored and this hook must not rely on transcript + mutation for persistence. ``kwargs`` may include ``turn_id``, + ``task_id``, ``api_call_count``, ``interrupted``, ``failed``, and + ``turn_exit_reason``. + + ``usage`` carries the completed turn's canonical token usage (the same + dict shape passed to ``update_from_response`` — ``prompt_tokens`` / + ``completion_tokens`` / ``total_tokens`` plus the canonical + ``input_tokens`` / ``output_tokens`` / ``cache_read_tokens`` / + ``cache_write_tokens`` / ``reasoning_tokens`` buckets) so an engine can + weigh how large/expensive the selected context actually was when + deciding the next ``select_context()``. It is ``None`` on finalized + turns that never reached a provider response (e.g. interrupt); engines + must treat it as optional. + + Default is a no-op. + """ + return None + # -- Optional: pre-flight check ---------------------------------------- def should_compress_preflight(self, messages: List[Dict[str, Any]]) -> bool: diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index a764501febab..a234213e32ca 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -716,6 +716,136 @@ def _sync_failover_system_message(agent, api_messages, active_system_prompt): return sp +def _apply_context_engine_selection( + agent: Any, + api_messages: List[Dict[str, Any]], + conversation_messages: List[Dict[str, Any]], + incoming_message: Optional[Dict[str, Any]], + *, + logger: Any, +) -> List[Dict[str, Any]]: + """Run the optional per-turn ``ContextEngine.select_context()`` hook. + + Returns the (possibly replaced) request message list. The hook is for + context *selection / routing* (retrieval, topic routing, role switching), + which is distinct from compression and fires every turn independent of + ``should_compress()``. + + Fail-open by design: a missing hook, any exception, or an invalid return + value yields the unmodified ``api_messages``. The result is request-only — + persisted conversation history is never mutated here. + """ + engine = getattr(agent, "context_compressor", None) + if engine is None or not hasattr(engine, "select_context"): + return api_messages + + # Skip the no-op base implementation so non-implementing engines — + # including the built-in ContextCompressor — pay nothing per request: + # no history copies below, no call. ``hasattr`` alone is not enough, + # because the ABC defines a default ``select_context`` that every engine + # inherits. Mirrors the base-method short-circuit in + # ``_notify_context_engine_turn_complete``. Lazy import avoids any import + # cycle with agent.context_engine. + try: + from agent.context_engine import ContextEngine as _CE + if getattr(engine.select_context, "__func__", None) is _CE.select_context: + return api_messages + except Exception: + pass + + session_label = getattr(agent, "session_id", None) or "-" + # Pass shallow copies of the reference-only inputs so an engine that + # mutates them in place cannot alter persisted transcript state. Only + # ``request_messages`` (the per-call request list) is meant to be acted on, + # and it may be replaced wholesale via the return value — never mutated in + # place either. ``conversation_messages`` / ``incoming_message`` are + # read-only context; copying enforces the request-only contract rather than + # merely documenting it. + _conv_copy = [dict(m) if isinstance(m, dict) else m for m in conversation_messages] \ + if conversation_messages is not None else None + _incoming_copy = dict(incoming_message) if isinstance(incoming_message, dict) else incoming_message + try: + selected = engine.select_context( + api_messages, + conversation_messages=_conv_copy, + incoming_message=_incoming_copy, + budget_tokens=getattr(engine, "context_length", 0) or 0, + ) + except Exception: + logger.warning( + "Context engine select_context hook failed; using unmodified " + "request messages (session=%s)", + session_label, + exc_info=True, + ) + return api_messages + + if selected is None: + return api_messages + # Require a NON-EMPTY list of dicts. An empty list must fall open to the + # original request: ``all([])`` is ``True``, so without the emptiness check + # a ``[]`` returned by a buggy/failing engine would replace a valid request + # with an empty message list that the downstream sanitizers cannot restore, + # reaching the provider as an invalid request instead of failing open. + if isinstance(selected, list) and selected and all(isinstance(m, dict) for m in selected): + return selected + + logger.warning( + "Context engine select_context returned an invalid value " + "(not a non-empty list of dicts); ignoring (session=%s)", + session_label, + ) + return api_messages + + +def _notify_context_engine_turn_complete( + agent: Any, + messages: List[Dict[str, Any]], + *, + usage: Optional[Dict[str, Any]] = None, + logger: Any, + **meta: Any, +) -> None: + """Notify the active context engine that a user turn has finished. + + Calls the optional ``ContextEngine.on_turn_complete()`` observation hook + once per turn, after the assistant/tool loop has produced the finalized + transcript. The complement to ``select_context()`` (pre-request selection): + this lets an engine ingest / index / summarize the completed turn. + + Fail-open: a missing or no-op hook, or any exception, is swallowed. + ``messages`` is passed as a shallow copy so the engine cannot mutate the + persisted transcript. + """ + engine = getattr(agent, "context_compressor", None) + hook = getattr(engine, "on_turn_complete", None) + if engine is None or not callable(hook): + return + + # Skip the no-op base implementation so non-implementing engines (incl. + # the built-in compressor) pay nothing per turn. Lazy import avoids any + # import cycle with agent.context_engine. + try: + from agent.context_engine import ContextEngine as _CE + if getattr(hook, "__func__", None) is _CE.on_turn_complete: + return + except Exception: + pass + + try: + hook( + [dict(m) if isinstance(m, dict) else m for m in messages], + usage=usage, + **meta, + ) + except Exception: + logger.warning( + "Context engine on_turn_complete hook failed (session=%s)", + getattr(agent, "session_id", None) or "-", + exc_info=True, + ) + + def run_conversation( agent, user_message: Any, @@ -855,6 +985,13 @@ def run_conversation( # over instead of spinning. Reset here so each turn starts fresh. See #26080. agent._auth_pool_refresh_counts = {} + # Reset the per-turn usage holder forwarded to the context engine's + # on_turn_complete() observation hook. Set after each successful provider + # response (see below); left as None on turns that never reach a response + # (early failure / interrupt) so the hook receives None rather than a + # stale prior turn's usage. + agent._last_turn_usage = None + # Optional opt-in runtime: if api_mode == codex_app_server, hand the # turn to the codex app-server subprocess (terminal/file ops/patching # all run inside Codex). Default Hermes path is bypassed entirely. @@ -1208,6 +1345,26 @@ def run_conversation( for idx, pfm in enumerate(agent.prefill_messages): api_messages.insert(sys_offset + idx, pfm.copy()) + # Per-turn context selection hook (additive, no-op by default). + # Lets a context engine select/replace which context enters the + # prompt for THIS call only — retrieval, topic routing, role/branch + # switching — distinct from compression and independent of + # should_compress(). Request-only: persisted history is untouched, so + # caching/sanitization below operate on whatever the engine selected. + # Fail-open (see _apply_context_engine_selection). + _sel_incoming = ( + messages[current_turn_user_idx] + if 0 <= current_turn_user_idx < len(messages) + else None + ) + api_messages = _apply_context_engine_selection( + agent, + api_messages, + messages, + _sel_incoming, + logger=request_logger, + ) + # Apply Anthropic prompt caching for Claude models on native # Anthropic, OpenRouter, and third-party Anthropic-compatible # gateways. Auto-detected: if ``_use_prompt_caching`` is set, @@ -2620,6 +2777,14 @@ def _perform_api_call(next_api_kwargs): "reasoning_tokens": canonical_usage.reasoning_tokens, } agent.context_compressor.update_from_response(usage_dict) + + # Stash this response's canonical usage so the post-turn + # on_turn_complete() observation hook can forward it (the + # same dict shape passed to update_from_response). A turn + # may make several API calls; the engine's per-turn signal + # of interest is the cost/size of the latest assembled + # request, so we keep the most recent call's usage. + agent._last_turn_usage = dict(usage_dict) elif getattr( agent.context_compressor, "awaiting_real_usage_after_compression", diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index 2126a9afdc26..4e2d318b2e2c 100644 --- a/agent/turn_finalizer.py +++ b/agent/turn_finalizer.py @@ -525,6 +525,34 @@ def finalize_turn( except Exception as exc: logger.warning("post_llm_call hook failed: %s", exc) + # Context engine observation hook: notify the active engine that this + # turn has finished, with the finalized transcript. Complements the + # per-request select_context() hook (selection before the request; + # observation after the turn). No-op default, fail-open. + try: + from agent.conversation_loop import _notify_context_engine_turn_complete + # Forward the turn's canonical usage when the host has it. The loop + # stashes the most recent API response's usage dict (the same + # canonical buckets fed to ``update_from_response``) on the agent as + # ``_last_turn_usage``. It is ``None`` on turns that never reached a + # provider response (early failure / interrupt), which is exactly the + # contract: real usage when available, ``None`` otherwise. + _turn_usage = getattr(agent, "_last_turn_usage", None) + _notify_context_engine_turn_complete( + agent, + messages, + usage=_turn_usage, + logger=logger, + turn_id=turn_id, + task_id=effective_task_id, + api_call_count=api_call_count, + interrupted=interrupted, + failed=failed, + turn_exit_reason=_turn_exit_reason, + ) + except Exception as exc: + logger.warning("on_turn_complete notification failed: %s", exc) + # Extract reasoning from the CURRENT turn only. Walk backwards # but stop at the user message that started this turn — anything # earlier is from a prior turn and must not leak into the reasoning diff --git a/contributors/emails/chaosxinglong@gmail.com b/contributors/emails/chaosxinglong@gmail.com new file mode 100644 index 000000000000..dac570b17399 --- /dev/null +++ b/contributors/emails/chaosxinglong@gmail.com @@ -0,0 +1 @@ +chaos-xxl diff --git a/tests/agent/test_context_engine_on_turn_complete_usage.py b/tests/agent/test_context_engine_on_turn_complete_usage.py new file mode 100644 index 000000000000..b962c983c450 --- /dev/null +++ b/tests/agent/test_context_engine_on_turn_complete_usage.py @@ -0,0 +1,155 @@ +"""Integration test: ``finalize_turn`` forwards the turn's real usage to the +``ContextEngine.on_turn_complete()`` observation hook. + +The hook is the engine's post-turn observation point, so it must receive the +completed turn's canonical token usage (prompt/completion/total + the canonical +``input_tokens`` / ``output_tokens`` / ``cache_read_tokens`` / +``cache_write_tokens`` / ``reasoning_tokens`` buckets) when the host has it — +not a hardcoded ``None`` — so the engine can weigh how large/expensive the +selected context was before the next ``select_context()``. + +The conversation loop stashes the most recent provider response's usage on the +agent as ``_last_turn_usage`` (the same dict shape fed to +``update_from_response``); ``finalize_turn`` forwards it. On turns that never +reach a provider response (early failure / interrupt) it stays ``None`` and the +hook receives ``None``. These tests pin both ends of that contract through the +real ``finalize_turn`` call site (the path that previously passed +``usage=None``). +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from agent.context_engine import ContextEngine + +# Reuse the minimal agent harness that exercises the real finalize_turn path. +from tests.agent.test_turn_finalizer_cleanup_guard import _StubAgent, _run + + +class _CapturingEngine(ContextEngine): + """Engine that records what on_turn_complete() receives.""" + + last_prompt_tokens = 0 + + def __init__(self) -> None: + self.captured: Dict[str, Any] = {} + + @property + def name(self) -> str: + return "capturing" + + def update_from_response(self, usage: Dict[str, Any]) -> None: + pass + + def should_compress(self, prompt_tokens: int = None) -> bool: + return False + + def compress( + self, + messages: List[Dict[str, Any]], + current_tokens: int = None, + focus_topic: str = None, + ) -> List[Dict[str, Any]]: + return messages + + def on_turn_complete(self, messages, usage=None, **kwargs): + self.captured["seen"] = True + self.captured["usage"] = usage + self.captured["kwargs"] = kwargs + + +CANONICAL_USAGE = { + "prompt_tokens": 1200, + "completion_tokens": 80, + "total_tokens": 1280, + "input_tokens": 1200, + "output_tokens": 80, + "cache_read_tokens": 1024, + "cache_write_tokens": 0, + "reasoning_tokens": 16, +} + + +def _agent_with_engine() -> _StubAgent: + agent = _StubAgent(raise_in=()) + agent.context_compressor = _CapturingEngine() + return agent + + +def test_finalize_turn_forwards_canonical_usage_when_available(): + """A completed turn forwards the stashed canonical usage dict intact.""" + agent = _agent_with_engine() + agent._last_turn_usage = dict(CANONICAL_USAGE) + + _run(agent, final_response="done") + + captured = agent.context_compressor.captured + assert captured.get("seen") is True + # The full canonical bucket set is forwarded unchanged — the engine relies + # on cache_read/write + reasoning, not just the legacy aggregate keys. + assert captured["usage"] == CANONICAL_USAGE + # Turn metadata still rides alongside usage. + assert captured["kwargs"]["turn_id"] == "turn-1" + + +def test_finalize_turn_forwards_none_when_no_response_usage(): + """An early-failure/interrupt turn (no stashed usage) forwards None.""" + agent = _agent_with_engine() + # _last_turn_usage left unset, mirroring a turn that never reached a + # provider response. + if hasattr(agent, "_last_turn_usage"): + delattr(agent, "_last_turn_usage") + + _run(agent, final_response="done") + + captured = agent.context_compressor.captured + assert captured.get("seen") is True + assert captured["usage"] is None + + +def test_finalization_seam_observes_interrupted_turn_with_none_usage(): + """Pins the documented coverage contract for on_turn_complete(). + + on_turn_complete() fires from the turn-finalization seam and reports + ``usage=None`` on a finalized turn that never reached a provider response + (e.g. interrupt), forwarding the ``interrupted`` flag. This is the testable + (positive) half of the contract. + + The negative half — abnormal early-return paths in ``run_conversation`` + (content-policy block, provider terminal failure, etc.) bypass finalization + and therefore do NOT emit the hook — is documented as best-effort coverage. + It is intentionally not pinned here: exercising those inline early returns + requires a full ``run_conversation`` harness, and unifying all terminal + paths behind one seam is a separate follow-up. + """ + from agent.turn_finalizer import finalize_turn + + agent = _agent_with_engine() + if hasattr(agent, "_last_turn_usage"): + delattr(agent, "_last_turn_usage") # never reached a provider response + + finalize_turn( + agent, + final_response="interrupted mid-turn", + api_call_count=1, + interrupted=True, + failed=False, + messages=[ + {"role": "user", "content": "do a thing"}, + {"role": "assistant", "content": "partial"}, + ], + conversation_history=None, + effective_task_id="task-1", + turn_id="turn-int", + user_message="do a thing", + original_user_message="do a thing", + _should_review_memory=False, + _turn_exit_reason="interrupt", + ) + + captured = agent.context_compressor.captured + assert captured.get("seen") is True + assert captured["usage"] is None + assert captured["kwargs"]["interrupted"] is True + assert captured["kwargs"]["turn_id"] == "turn-int" diff --git a/tests/agent/test_context_engine_select_context.py b/tests/agent/test_context_engine_select_context.py new file mode 100644 index 000000000000..c8523943708e --- /dev/null +++ b/tests/agent/test_context_engine_select_context.py @@ -0,0 +1,400 @@ +"""Tests for the per-turn ``ContextEngine.select_context()`` hook. + +``select_context()`` is the *selection / routing* verb — distinct from +compression — that lets an external context engine replace which context +enters the prompt for a single request, every turn, independent of +``should_compress()``. It is additive and no-op by default, and the host +call site (``_apply_context_engine_selection``) is fail-open: a missing hook, +an exception, or an invalid return value must leave the assembled request +untouched and must never mutate persisted history. + +This pins the contract that engines such as retrieval-augmented, topic-routed, +and role-switching engines rely on (RFC #36765), consolidating the per-turn +request-assembly surface proposed across #41918, #24949, #47109, and #50053. +""" + +from __future__ import annotations + +from typing import Any, Dict, List +from unittest.mock import MagicMock + +from agent.context_engine import ContextEngine +from agent.conversation_loop import ( + _apply_context_engine_selection, + _notify_context_engine_turn_complete, +) + + +class _MinimalEngine(ContextEngine): + """Concrete engine implementing only the abstract methods.""" + + @property + def name(self) -> str: + return "minimal" + + def update_from_response(self, usage: Dict[str, Any]) -> None: + pass + + def should_compress(self, prompt_tokens: int = None) -> bool: + return False + + def compress( + self, + messages: List[Dict[str, Any]], + current_tokens: int = None, + focus_topic: str = None, + ) -> List[Dict[str, Any]]: + return messages + + +def _agent_with(engine) -> Any: + agent = MagicMock() + agent.session_id = "test-session" + agent.context_compressor = engine + return agent + + +REQUEST = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hello"}, +] +HISTORY = [{"role": "user", "content": "hello"}] + + +# -- ABC default ----------------------------------------------------------- + +def test_default_select_context_is_noop(): + """The base implementation returns None (no replacement).""" + engine = _MinimalEngine() + assert ( + engine.select_context( + REQUEST, + conversation_messages=HISTORY, + incoming_message=HISTORY[-1], + budget_tokens=0, + ) + is None + ) + + +# -- Host call site: _apply_context_engine_selection ----------------------- + +def test_none_return_leaves_request_unchanged(): + """An engine returning None falls through to the assembled request.""" + engine = _MinimalEngine() # default select_context -> None + agent = _agent_with(engine) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is REQUEST + + +def test_base_noop_select_context_is_short_circuited_not_called(): + """Non-implementing engines skip the hook entirely (no call, no copies). + + The built-in ContextCompressor — and any engine that merely inherits the + ABC default — must keep the default request path byte-identical AND pay + nothing per request. ``hasattr`` alone cannot distinguish "inherits the + no-op default" from "implements the hook" because the ABC defines + ``select_context`` on every engine; the host therefore identity-checks the + bound method against ``ContextEngine.select_context`` and short-circuits + WITHOUT calling it or building the shallow reference copies. This pins + that: even a base implementation patched to raise is never invoked. + """ + from unittest.mock import patch as _patch + + def _explode(self, request_messages, **kwargs): + raise AssertionError("base select_context must not be invoked") + + engine = _MinimalEngine() # inherits the ABC default + agent = _agent_with(engine) + logger = MagicMock() + with _patch.object(ContextEngine, "select_context", _explode): + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=logger + ) + assert out is REQUEST + assert not logger.warning.called + + +def test_builtin_compressor_inherits_base_select_context(): + """The built-in ContextCompressor must NOT implement the new verbs. + + Guards the default-path byte-identity contract: if someone overrides + ``select_context`` / ``on_turn_complete`` on ContextCompressor, the host + short-circuits no longer skip it and the default request pipeline gains a + per-request call — update this pin only together with that decision. + """ + from agent.context_compressor import ContextCompressor + + assert "select_context" not in ContextCompressor.__dict__ + assert "on_turn_complete" not in ContextCompressor.__dict__ + + +def test_missing_hook_leaves_request_unchanged(): + """An engine without select_context (older/stub base) is a no-op.""" + engine = object() # no select_context attribute + agent = _agent_with(engine) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is REQUEST + + +def test_no_engine_leaves_request_unchanged(): + agent = MagicMock() + agent.session_id = "test-session" + agent.context_compressor = None + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is REQUEST + + +def test_valid_list_replaces_request(): + """A valid list of dicts replaces the request messages for this call.""" + replacement = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "routed-context"}, + ] + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + return replacement + + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is replacement + + +def test_exception_fails_open(): + """A raising hook is swallowed; the unmodified request is used.""" + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + raise RuntimeError("backend offline") + + logger = MagicMock() + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=logger + ) + assert out is REQUEST + assert logger.warning.called + + +def test_non_list_return_is_ignored(): + """A non-list return value is rejected and logged, request unchanged.""" + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + return {"role": "user", "content": "oops not a list"} + + logger = MagicMock() + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=logger + ) + assert out is REQUEST + assert logger.warning.called + + +def test_list_of_non_dicts_is_ignored(): + """A list that isn't all dicts is rejected, request unchanged.""" + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + return ["not", "dicts"] + + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is REQUEST + + +def test_empty_list_keeps_original_request(): + """An empty list must fall open to the original request. + + ``all([])`` is ``True``, so without an emptiness check a ``[]`` returned by + a failing/buggy engine would replace a valid assembled request with an + empty message list the downstream sanitizers cannot restore — reaching the + provider as an invalid request instead of failing open. Guards the fail-open + contract. + """ + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + return [] + + logger = MagicMock() + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=logger + ) + assert out is REQUEST + assert logger.warning.called + + +def test_engine_mutating_inputs_cannot_corrupt_persisted_state(): + """An engine that mutates its read-only inputs in place must not affect the + persisted conversation history / incoming message. + + ``conversation_messages`` and ``incoming_message`` are reference-only + context. The host passes shallow copies, so even a misbehaving engine that + appends to / edits them in ``select_context()`` cannot alter the live + persisted objects. Enforces the request-only contract (not just documents). + """ + history = [{"role": "user", "content": "hello"}] + incoming = history[-1] + history_snapshot = [dict(m) for m in history] + incoming_snapshot = dict(incoming) + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, *, conversation_messages=None, + incoming_message=None, **kwargs): + # Misbehaving engine: mutate the read-only inputs in place. + if conversation_messages is not None: + conversation_messages.append({"role": "user", "content": "INJECTED"}) + if conversation_messages and isinstance(conversation_messages[0], dict): + conversation_messages[0]["content"] = "TAMPERED" + if isinstance(incoming_message, dict): + incoming_message["content"] = "TAMPERED" + return None + + agent = _agent_with(_Engine()) + _apply_context_engine_selection( + agent, REQUEST, history, incoming, logger=MagicMock() + ) + # Persisted history + incoming message are untouched despite the engine's + # in-place mutation of the copies it received. + assert history == history_snapshot + assert incoming == incoming_snapshot + + +def test_persisted_history_not_mutated(): + """The hook must not mutate the persisted conversation history.""" + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, *, conversation_messages=None, **kwargs): + # Even a misbehaving engine touching its inputs must not affect + # what the host persists — the host passes the live list, so we + # assert the host contract by checking the engine received it and + # the canonical copy is unchanged after the call. + return list(request_messages) + + history_snapshot = [dict(m) for m in HISTORY] + agent = _agent_with(_Engine()) + _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert HISTORY == history_snapshot + + +# -- cache-stability + downstream-sanitizer contract ----------------------- + +def test_noop_preserves_request_byte_stable_for_cache(): + """No-op default must leave the request byte-identical. + + Prompt-cache stability is a host invariant (AGENTS.md): the hook runs + before cache-control, so a no-op engine must not perturb the list — + otherwise cache breakpoints would shift for every existing engine. The + host returns the *same object*, so cache-control sees identical input. + """ + snapshot = [dict(m) for m in REQUEST] + agent = _agent_with(_MinimalEngine()) # default select_context -> None + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is REQUEST # same object -> byte-stable for cache-control + assert REQUEST == snapshot # unperturbed + + +def test_role_unusual_replacement_passed_through_for_downstream_sanitizers(): + """The hook does structural validation only; role/tool normalization is + deferred to the existing downstream sanitizers. + + A `system -> user -> user` replacement (the exact shape flagged as a + role-alternation risk on sibling PRs) is well-formed structurally, so the + host returns it verbatim. Role-pairing/orphaned-tool cleanup runs *after* + this hook in the request pipeline (`_sanitize_api_messages`, + `_drop_thinking_only_and_merge_users`), so select_context cannot emit a + malformed request that bypasses validation. + """ + role_unusual = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "a"}, + {"role": "user", "content": "b"}, + ] + + class _Engine(_MinimalEngine): + def select_context(self, request_messages, **kwargs): + return role_unusual + + agent = _agent_with(_Engine()) + out = _apply_context_engine_selection( + agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() + ) + assert out is role_unusual # accepted structurally; downstream sanitizers normalize + + +# -- on_turn_complete (post-turn observation) ------------------------------ + +def test_default_on_turn_complete_is_noop(): + """The base on_turn_complete returns None and does nothing.""" + assert _MinimalEngine().on_turn_complete(HISTORY, usage=None) is None + + +def test_on_turn_complete_called_with_snapshot_and_meta(): + """Host forwards a transcript copy + metadata; base no-op is skipped.""" + captured = {} + + class _Engine(_MinimalEngine): + def on_turn_complete(self, messages, usage=None, **kwargs): + captured["messages"] = messages + captured["usage"] = usage + captured["kwargs"] = kwargs + + agent = _agent_with(_Engine()) + _notify_context_engine_turn_complete( + agent, HISTORY, usage={"total_tokens": 12}, logger=MagicMock(), + turn_id="t1", api_call_count=1, + ) + assert captured["messages"] == HISTORY + assert captured["messages"] is not HISTORY # shallow copy + assert captured["usage"] == {"total_tokens": 12} + assert captured["kwargs"]["turn_id"] == "t1" + assert captured["kwargs"]["api_call_count"] == 1 + + +def test_on_turn_complete_base_noop_is_skipped(): + """An engine that only inherits the base no-op is handled safely. + + The helper short-circuits the base implementation (so non-implementing + engines pay nothing), and in any case must not raise. + """ + agent = _agent_with(_MinimalEngine()) # inherits base on_turn_complete + _notify_context_engine_turn_complete(agent, HISTORY, logger=MagicMock()) + + +def test_on_turn_complete_fails_open(): + """A raising observation hook is swallowed and logged.""" + + class _Engine(_MinimalEngine): + def on_turn_complete(self, messages, usage=None, **kwargs): + raise RuntimeError("indexing backend down") + + logger = MagicMock() + agent = _agent_with(_Engine()) + _notify_context_engine_turn_complete(agent, HISTORY, logger=logger) + assert logger.warning.called + + +def test_on_turn_complete_missing_engine_is_safe(): + agent = MagicMock() + agent.session_id = "s" + agent.context_compressor = None + # No engine -> silent return, no raise. + _notify_context_engine_turn_complete(agent, HISTORY, logger=MagicMock()) diff --git a/website/docs/developer-guide/context-engine-plugin.md b/website/docs/developer-guide/context-engine-plugin.md index a6e53de9dbf4..169b2bf7cc68 100644 --- a/website/docs/developer-guide/context-engine-plugin.md +++ b/website/docs/developer-guide/context-engine-plugin.md @@ -97,6 +97,67 @@ These have sensible defaults in the ABC. Override as needed: | `handle_tool_call(name, args, **kwargs)` | Returns error JSON | You implement tool handlers | | `should_compress_preflight(messages)` | Returns `False` | You can do a cheap pre-API-call estimate | | `get_status()` | Standard token/threshold dict | You have custom metrics to expose | +| `select_context(request_messages, *, conversation_messages, incoming_message, budget_tokens)` | Returns `None` (no-op) | You select/route which context enters **this** request (retrieval, topic routing) — see below | +| `on_turn_complete(messages, usage=None, **kwargs)` | No-op | You ingest/index/observe the finished turn — see below | + +## Per-turn context selection and observation + +`compress()` answers "context is too long → make it shorter". Two optional, +no-op-default hooks cover the orthogonal *selection / observation* axis, so an +engine no longer has to force `should_compress()` to `True` and abuse +`compress()` as a per-turn callback: + +```python +def select_context(self, request_messages, *, conversation_messages=None, + incoming_message=None, budget_tokens=0): + """Choose/replace the context for THIS request, before dispatch. + + Return a new message list to use for this one provider call (retrieval, + topic routing, role/branch switching), or None to leave it unchanged. + Request-only: the persisted conversation history is never mutated. + """ + +def on_turn_complete(self, messages, usage=None, **kwargs): + """Observe a finished turn after the assistant/tool loop completes. + + Receives a shallow copy of the finalized transcript plus the turn's + canonical usage dict (or None if no provider response was reached), so the + engine can ingest/index/summarize for the next select_context(). The return + value is ignored. + """ +``` + +Contract: + +- **No-op by default, fail-open.** Both default to `return None`. A missing hook, an exception, or an invalid return value leaves the request untouched — so a failing engine is never worse than not installing one. The host also identity-checks for the inherited ABC default and skips it entirely, so non-implementing engines (including the built-in compressor) pay no per-request work at all. +- **`select_context()` is request-only.** The returned list replaces the messages for a single provider call; persisted history is never written. Returning `None`, `[]`, a non-list, or a list containing non-dicts all fall open to the unmodified request. +- **Ordering / cache stability.** The hook runs **before** prompt cache-control and every request sanitizer, so (a) a replacement still passes the same validation as any request, and (b) the no-op default leaves the request byte-identical — prompt-cache behaviour is unchanged for non-implementing engines. An engine that replaces the list changes only its own cache prefix. Evaluated per provider request (re-runs on retries). +- **`on_turn_complete()`** is post-turn observation only; treat `messages` as read-only. **Coverage is best-effort:** it fires from the standard turn-finalization seam. Some abnormal early-return paths in the loop (e.g. a content-policy block or a provider terminal failure) persist and return without routing through finalization, so they do not currently emit this hook — treat it as a best-effort observation for completed turns, not a guaranteed callback for every early exit. Unifying all terminal paths behind one finalization seam is a separate follow-up. + +### When to use these hooks — and when NOT to + +- **Implement `select_context()` only when your engine must *replace* the + per-request context** — retrieval-augmented selection, topic/branch routing, + role switching. It is the only verb that can swap which messages enter a + request: the `pre_llm_call` plugin hook is inject-only by documented design + (it appends to the user message and never rewrites the list, to preserve the + prompt-cache prefix). If you don't need replacement, don't implement it. +- **If your plugin only needs post-turn observation / ingestion** (indexing, + memory sync, analytics), implement a **memory provider** (`sync_turn()` — + see [Memory Provider Plugins](./memory-provider-plugin.md)) instead of a + context engine. A context engine takes ownership of the session's compaction + policy; a memory provider observes turns without owning anything. + `on_turn_complete()` exists as the observation mirror for engines that + *already* need `select_context()` — so the same component can learn from the + turn it just routed — not as a general-purpose turn callback. +- **Prompt-cache impact of a real `select_context()`.** A non-no-op selection + naturally changes the prompt-cache prefix for the turns where it changes the + selection — that request's prefix no longer matches the provider's cached + prefix, so those turns re-write cache instead of reading it. Engines should + return **stable selections when nothing has changed** (same object or an + equal list), and only reshape the context when the routing decision actually + differs; a selection that shuffles per turn silently forfeits cache reuse + every turn. ## Engine tools