diff --git a/agent/context_engine.py b/agent/context_engine.py index 20134b680cb3..17480f578260 100644 --- a/agent/context_engine.py +++ b/agent/context_engine.py @@ -263,6 +263,36 @@ class ContextEngine(ABC): """ 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 once after the assistant/tool loop completes for a turn, 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. + + 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``. + + 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 6f93f38bb69c..e164b0aa5934 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -769,6 +769,54 @@ def _apply_context_engine_selection( 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, diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index 2126a9afdc26..f2ddb947ba48 100644 --- a/agent/turn_finalizer.py +++ b/agent/turn_finalizer.py @@ -525,6 +525,27 @@ 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 + _notify_context_engine_turn_complete( + agent, + messages, + usage=None, + 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/tests/agent/test_context_engine_select_context.py b/tests/agent/test_context_engine_select_context.py index d2b686fe2fdc..55319d2f33d5 100644 --- a/tests/agent/test_context_engine_select_context.py +++ b/tests/agent/test_context_engine_select_context.py @@ -19,7 +19,10 @@ 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 +from agent.conversation_loop import ( + _apply_context_engine_selection, + _notify_context_engine_turn_complete, +) class _MinimalEngine(ContextEngine): @@ -187,3 +190,63 @@ def test_persisted_history_not_mutated(): agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock() ) assert HISTORY == history_snapshot + + +# -- 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())