From dec464c35141a59b8742c02b4f35553c3dd3cdaa Mon Sep 17 00:00:00 2001 From: xue xinglong Date: Tue, 23 Jun 2026 14:21:12 +0800 Subject: [PATCH] feat(context-engine): add select_context() per-turn selection hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an optional, no-op-default select_context() hook to the ContextEngine ABC, called every turn after the request messages are assembled and before provider dispatch — independent of should_compress(). Lets an engine select or replace which context enters the prompt for a single request (retrieval, topic routing, role/branch switching) without mutating persisted history, removing the need to abuse should_compress()=True as a per-turn callback. The host call site (_apply_context_engine_selection) is fail-open: a missing hook, an exception, or an invalid return value leaves the assembled request untouched. Additive and non-breaking: the built-in compressor and every existing engine are unaffected. Consolidates the per-turn request-assembly surface proposed across #41918, Related: #36765 #41918 #24949 #47109 #50053 #23837 #25115 #29370 --- agent/context_engine.py | 53 +++++ agent/conversation_loop.py | 73 +++++++ .../test_context_engine_select_context.py | 189 ++++++++++++++++++ 3 files changed, 315 insertions(+) create mode 100644 tests/agent/test_context_engine_select_context.py diff --git a/agent/context_engine.py b/agent/context_engine.py index 2225f25473d..20134b680cb 100644 --- a/agent/context_engine.py +++ b/agent/context_engine.py @@ -210,6 +210,59 @@ class ContextEngine(ABC): """ 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. + + 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 + # -- 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 a764501feba..6f93f38bb69 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -716,6 +716,59 @@ 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 + + session_label = getattr(agent, "session_id", None) or "-" + try: + selected = engine.select_context( + api_messages, + conversation_messages=conversation_messages, + incoming_message=incoming_message, + 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 + if isinstance(selected, list) and all(isinstance(m, dict) for m in selected): + return selected + + logger.warning( + "Context engine select_context returned a non-list of dicts; " + "ignoring (session=%s)", + session_label, + ) + return api_messages + + def run_conversation( agent, user_message: Any, @@ -1208,6 +1261,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, 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 00000000000..d2b686fe2fd --- /dev/null +++ b/tests/agent/test_context_engine_select_context.py @@ -0,0 +1,189 @@ +"""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 + + +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_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_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