feat(context-engine): select_context() + on_turn_complete() ABC verbs — salvage of #51226 (@chaos-xxl, RFC #36765)

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:

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",

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

contributors/emails/chaosxinglong@gmail.com

@@ -0,0 +1 @@
+chaos-xxl