python-codex 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

pycodex/agent.py

@@ -1,64 +1,78 @@
-
 import asyncio
 import json
-import re
+import typing
+from dataclasses import dataclass, replace
+from pathlib import Path
 from typing import Callable
 
-from .context import ContextManager
-from .model import ModelClient
+from .context import ContextConfig, ContextManager
+from .events import (
+    AutoCompactCompletedEvent,
+    AutoCompactFailedEvent,
+    AutoCompactStartedEvent,
+    CompactCompletedEvent,
+    CompactFailedEvent,
+    CompactStartedEvent,
+    Event,
+    ModelCalledEvent,
+    ModelCompletedEvent,
+    ModelEvent,
+    StreamErrorEvent,
+    TerminalEvent,
+    TokenCountEvent,
+    ToolCompletedEvent,
+    ToolStartedEvent,
+    TurnCompletedEvent,
+    TurnEvent,
+    TurnFailedEvent,
+    TurnInterruptedEvent,
+    TurnStartedEvent,
+)
+from .model import (
+    DEFAULT_ORIGINATOR,
+    ContextLengthExceeded,
+    ModelClient,
+    ModelControl,
+    ResponsesIncompleteError,
+)
 from .protocol import (
-    AgentEvent,
     AssistantMessage,
     ConversationItem,
-    ModelStreamEvent,
+    ModelResponse,
     ReasoningItem,
     ToolCall,
     ToolResult,
     TurnResult,
     UserMessage,
 )
-from .tools import (
-    ClockManager,
-    ClockTool,
-    ExecCommandTool,
-    ToolContext,
-    ToolRegistry,
-    UnifiedExecManager,
-)
-from .utils.truncation import truncate_tool_results_for_history
+from .tools import ToolContext, ToolRegistry
 from .utils import uuid7_string
-import typing
+from .utils.session_persist import (
+    SessionRolloutRecorder,
+    load_resumed_session_path,
+    resolve_codex_home,
+    rollout_path_for_session,
+)
+from .utils.truncation import truncate_tool_result_for_history
 
 if typing.TYPE_CHECKING:
-    from .utils.session_persist import SessionRolloutRecorder
-    from .runtime_services import AgentRuntimeEnvironment
+    from .utils.compactor import CompactResult
 
 
-EventHandler = Callable[[AgentEvent], None]
-BASE_EVENT_HANDLER: 'EventHandler' = lambda _event: None
-_REQUESTED_TOKENS_RE = re.compile(
-    r"requested\s+([0-9,]+)\s+tokens",
-    re.IGNORECASE,
-)
-_REQUESTED_TOKEN_SPLIT_RE = re.compile(
-    r"\(([0-9,]+)\s+in\s+the\s+messages,\s+([0-9,]+)\s+in\s+the\s+completion\)",
-    re.IGNORECASE,
-)
-_MAX_CONTEXT_TOKENS_RE = re.compile(
-    r"maximum\s+context\s+length\s+is\s+([0-9,]+)\s+tokens",
-    re.IGNORECASE,
-)
-_CONTEXT_LENGTH_ERROR_MARKERS = (
-    "context_length_exceeded",
-    "maximum context length",
-    "exceeds the context window",
-    "exceeded the context window",
-)
-TERMINAL_TURN_EVENTS = {"turn_completed", "turn_failed", "turn_interrupted"}
+EventHandler = Callable[[Event], None]
+BASE_EVENT_HANDLER: "EventHandler" = lambda _event: None
 
 
 class TurnInterrupted(RuntimeError):
-    pass
+    def __init__(self):
+        super().__init__("turn interrupted")
+
+
+@dataclass
+class _TurnState:
+    turn_id: "str"
+    iteration: "int" = 0
+    output_text: "typing.Union[str, None]" = None
 
 
 class Agent:
@@ -72,357 +86,392 @@ class Agent:
 
     def __init__(
         self,
-        model_client: 'ModelClient',
-        tool_registry: 'ToolRegistry',
-        context_manager: 'typing.Union[ContextManager, None]' = None,
-        parallel_tool_calls: 'bool' = True,
-        event_handler: 'EventHandler' = BASE_EVENT_HANDLER,
-        initial_history: 'typing.Tuple[ConversationItem, ...]' = (),
-        rollout_recorder: 'typing.Union[SessionRolloutRecorder, None]' = None,
-        runtime_environment: 'AgentRuntimeEnvironment' = None,
-    ) -> 'None':
-        self._model_client = model_client
-        self._tool_registry = tool_registry
-        self._context_manager = context_manager or ContextManager()
-        self._parallel_tool_calls = parallel_tool_calls
-        self._event_handler = event_handler
-        self._history: 'typing.List[ConversationItem]' = list(initial_history)
-        self._rollout_recorder = rollout_recorder
-        self._auto_compact_token_limit = (
-            self._context_manager.resolve_auto_compact_token_limit()
+        model_client: "ModelClient",
+        tool_registry: "ToolRegistry",
+        context_config: "ContextConfig",
+        parallel_tool_calls: "bool" = True,
+        event_handler: "EventHandler" = BASE_EVENT_HANDLER,
+        initial_history: "typing.Tuple[ConversationItem, ...]" = (),
+        session_file_path: "typing.Union[str, Path, None]" = None,
+        session_id: "typing.Union[str, None]" = None,
+    ) -> "None":
+        self.model_client = model_client
+        self.tool_registry = tool_registry
+        self.context_manager = ContextManager(
+            replace(context_config, model=model_client.model)
         )
-        self._last_total_usage_tokens: 'typing.Union[int, None]' = None
-        self.runtime_environment = runtime_environment
-        self.interrupt_asap = False
-        self._turn_running = False
-        exec_command_tool = self._tool_registry.get_tool("exec_command")
-        self._exec_manager = (
-            exec_command_tool._manager
-            if isinstance(exec_command_tool, ExecCommandTool)
-            else None
-        )
-        if self._exec_manager is not None:
-            self._exec_manager.set_notify_hook(self.maybe_invoke)
-        clock_tool = self._tool_registry.get_tool("clock")
-        self._clock_manager: 'typing.Union[ClockManager, None]' = (
-            clock_tool._manager
-            if isinstance(clock_tool, ClockTool)
-            else None
-        )
-        if self._clock_manager is not None:
-            self._clock_manager.set_notify_hook(self.maybe_invoke)
+        self._parallel_tool_calls = parallel_tool_calls
+        self.event_handler = event_handler
+        self._history: "typing.List[ConversationItem]" = list(initial_history)
+        self._configure_recording(session_id or uuid7_string(), session_file_path)
+        self._last_total_usage_tokens: "typing.Union[int, None]" = None
+        self._idle: "typing.Union[asyncio.Event, None]" = None
+        self.is_shutdown = False
+        self.accepts_input = True
+        self._stop_requested = False
+        for tool in self.tool_registry.tools():
+            tool.bind_agent(self)
 
     @property
-    def history(self) -> 'typing.Tuple[ConversationItem, ...]':
+    def history(self) -> "typing.Tuple[ConversationItem, ...]":
         return tuple(self._history)
 
-    def set_event_handler(
-        self, event_handler: 'EventHandler' = BASE_EVENT_HANDLER
-    ) -> 'None':
-        self._event_handler = event_handler
+    @property
+    def session_file_path(self) -> "typing.Union[Path, None]":
+        recorder = self._rollout_recorder
+        return recorder.rollout_path if recorder is not None else None
+
+    @property
+    def recorded_session_file_path(self) -> "typing.Union[Path, None]":
+        recorder = self._rollout_recorder
+        if recorder is None or recorder._session_meta is not None:
+            return None
+        return recorder.rollout_path
 
-    def replace_history(
-        self,
-        history: 'typing.Iterable[ConversationItem]',
-    ) -> 'None':
-        self._history = list(history)
+    @property
+    def is_running(self) -> "bool":
+        return self._idle is not None
 
-    def set_rollout_recorder(
-        self,
-        rollout_recorder: 'typing.Union[SessionRolloutRecorder, None]',
-    ) -> 'None':
-        self._rollout_recorder = rollout_recorder
+    async def wait_until_idle(self) -> "None":
+        while self.is_running:
+            await self._idle.wait()
+
+    def stop_asap(self) -> "None":
+        if self.is_running:
+            self._stop_requested = True
+
+    @property
+    def model_name(self) -> "str":
+        return self.model_client.model
+
+    def set_model(self, model: "str") -> "None":
+        if self.is_running:
+            raise RuntimeError("cannot change model while agent is running")
+        client = typing.cast(ModelControl, self.model_client)
+        client.model = model
+        self.context_manager.set_model(model)
+        self._last_total_usage_tokens = None
 
-    def ask(self, text: 'str') -> 'TurnResult':
+    def resume(
+        self, session_file_path: "typing.Union[str, Path, None]" = None
+    ) -> "typing.Union[typing.Dict[str, object], None]":
+        if self.is_running:
+            raise RuntimeError("cannot restore session while agent is running")
+        resumed = None
+        if session_file_path is not None:
+            resumed = load_resumed_session_path(session_file_path)
+            session_id = str(resumed["session_id"])
+            if not session_id:
+                raise ValueError("session file has no session ID")
+            restored_history = list(resumed["history"])
+            self._configure_recording(session_id, resumed["rollout_path"], resume=True)
+            self._history = restored_history
+            self._last_total_usage_tokens = None
+        self.is_shutdown = False
+        self.accepts_input = True
+        for tool in self.tool_registry.tools():
+            tool.bind_agent(self)
+        return resumed
+
+    def ask(self, text: "str") -> "TurnResult":
         from .utils.async_bridge import run_async
 
         return run_async(self.run_turn([text]))
 
-    def _raise_if_interrupt_requested(
-        self,
-        turn_id: 'str',
-        iteration: 'int',
-        output_text: 'typing.Union[str, None]' = None,
-    ) -> 'None':
-        if self.interrupt_asap:
-            self.interrupt_asap = False
-            payload: 'typing.Dict[str, object]' = {"iteration": iteration}
-            if output_text is not None:
-                payload["output_text"] = output_text
-            self._emit("turn_interrupted", turn_id, **payload)
-            raise TurnInterrupted("turn interrupted")
-
-    async def run_turn(
-        self, texts: 'typing.List[str]', turn_id: 'typing.Union[str, None]' = None
-    ) -> 'TurnResult':
-        self._turn_running = True
-        if self._clock_manager is not None:
-            self._clock_manager.turn_started()
-        turn_id = turn_id or uuid7_string()
-        self.interrupt_asap = False
-        new_user_messages = [UserMessage(text=text) for text in texts]
+    def fork(self) -> "None":
+        if self.is_running:
+            raise RuntimeError("cannot fork session while agent is running")
+        session_id = uuid7_string()
+        path = None
+        if self._rollout_recorder is not None:
+            path = rollout_path_for_session(
+                self.context_manager._config.codex_home or resolve_codex_home(),
+                session_id,
+            )
+        self._configure_recording(session_id, path)
 
-        self._emit(
-            "turn_started",
-            turn_id,
-            user_text="\n".join(texts),
-            user_texts=list(texts),
-        )
-        await self._maybe_auto_compact(turn_id, phase="pre_turn")
-        self._history.extend(new_user_messages)
-        self._persist_history_items(new_user_messages)
+    def _configure_recording(
+        self,
+        session_id: "str",
+        session_file_path: "typing.Union[str, Path, None]" = None,
+        resume: "bool" = False,
+    ) -> "None":
+        if session_file_path is None:
+            recorder = None
+        elif resume:
+            recorder = SessionRolloutRecorder.resume(session_file_path)
+        else:
+            recorder = SessionRolloutRecorder.create(
+                self.context_manager._config.codex_home or resolve_codex_home(),
+                session_id,
+                self.context_manager.cwd,
+                getattr(self.model_client, "_originator", DEFAULT_ORIGINATOR),
+                getattr(
+                    getattr(self.model_client, "_config", None), "provider_name", None
+                ),
+                self.context_manager.resolve_base_instructions(),
+                session_file_path,
+            )
+        if hasattr(self.model_client, "_session_id"):
+            self.model_client._session_id = session_id
+        self.session_id = session_id
+        self._rollout_recorder = recorder
 
-        last_assistant_message: 'typing.Union[str, None]' = None
-        final_response_items: 'typing.Tuple[\n    typing.Union[typing.Union[AssistantMessage, ToolCall], ReasoningItem], ...\n]' = ()
+    async def compact(
+        self,
+        prune_tool_results_on_context_error: "bool" = True,
+    ) -> "typing.Union[CompactResult, None]":
+        if self.is_shutdown:
+            raise RuntimeError("agent is shutdown")
+        if self.is_running:
+            raise RuntimeError("cannot compact while agent is running")
+        self._idle = asyncio.Event()
+        try:
+            return await self._compact_history(
+                uuid7_string(),
+                "manual",
+                None,
+                None,
+                prune_tool_results_on_context_error,
+            )
+        finally:
+            self._idle.set()
+            self._idle = None
 
-        iteration = 0
+    async def run_turn(
+        self, texts: "typing.List[str]", turn_id: "typing.Union[str, None]" = None
+    ) -> "TurnResult":
+        if self.is_shutdown:
+            raise RuntimeError("agent is shutdown")
+        if self.is_running:
+            raise RuntimeError("agent already has an active turn")
+        self._stop_requested = False
+        self._idle = asyncio.Event()
+        turn = _TurnState(turn_id or uuid7_string())
         try:
+            self._emit(TurnStartedEvent(turn.turn_id, tuple(texts)))
             while True:
-                self._raise_if_interrupt_requested(
-                    turn_id,
-                    iteration,
-                    output_text=last_assistant_message,
-                )
-                await self._maybe_auto_compact(turn_id, phase="mid_turn")
-                iteration += 1
-                response = await self._complete_model_request(
-                    turn_id,
-                    iteration,
-                )
-                final_response_items = tuple(response.items)
+                phase = "pre_turn" if turn.iteration == 0 else "mid_turn"
+                await self._maybe_auto_compact(turn.turn_id, phase)
+                if turn.iteration == 0:
+                    self._append_history(UserMessage(text=text) for text in texts)
+                response = await self._sample(turn)
                 self._emit(
-                    "model_completed",
-                    turn_id,
-                    iteration=iteration,
-                    item_count=len(response.items),
+                    ModelCompletedEvent(
+                        turn.turn_id, turn.iteration, len(response.items)
+                    )
                 )
-
-                recorded_items = self._record_model_response_items(response.items)
-                tool_calls = recorded_items[1]
-                if recorded_items[2] is not None:
-                    last_assistant_message = recorded_items[2]
-
+                self._append_history(response.items)
+                tool_calls = []
+                for item in response.items:
+                    if isinstance(item, AssistantMessage):
+                        turn.output_text = item.text
+                    elif isinstance(item, ToolCall):
+                        tool_calls.append(item)
+
+                if tool_calls:
+                    await self._execute_tool_batch(turn.turn_id, tool_calls)
+                if self._stop_requested:
+                    raise TurnInterrupted()
                 if not tool_calls:
-                    self._raise_if_interrupt_requested(
-                        turn_id,
-                        iteration,
-                        output_text=last_assistant_message,
-                    )
-                    self._emit(
-                        "turn_completed",
-                        turn_id,
-                        iteration=iteration,
-                        output_text=last_assistant_message,
-                    )
-                    self._turn_running = False
-                    if self._clock_manager is not None:
-                        self._clock_manager.arm_after_reply()
-                    return TurnResult(
-                        turn_id=turn_id,
-                        output_text=last_assistant_message,
-                        iterations=iteration,
-                        response_items=final_response_items,
-                        history=tuple(self._history),
-                    )
+                    break
 
-                tool_results = await self._execute_tool_batch(turn_id, tool_calls)
-                tool_results = truncate_tool_results_for_history(tool_results)
-                self._history.extend(tool_results)
-                self._persist_history_items(tool_results)
-                follow_up_messages = self._build_follow_up_messages(tool_results)
-                self._history.extend(follow_up_messages)
-                self._persist_history_items(follow_up_messages)
-                self._raise_if_interrupt_requested(
-                    turn_id,
-                    iteration,
-                    output_text=last_assistant_message,
+            self._emit(
+                TurnCompletedEvent(
+                    turn.turn_id,
+                    turn.iteration,
+                    turn.output_text,
+                    self._background_work_count(TurnCompletedEvent),
                 )
+            )
+            result = TurnResult(
+                turn_id=turn.turn_id,
+                output_text=turn.output_text,
+                iterations=turn.iteration,
+                response_items=tuple(response.items),
+                history=self.history,
+            )
+            return result
         except TurnInterrupted:
-            self._turn_running = False
-            raise
-        except asyncio.CancelledError:
-            self._turn_running = False
+            self._emit(
+                TurnInterruptedEvent(
+                    turn.turn_id,
+                    turn.iteration,
+                    turn.output_text,
+                    self._background_work_count(TurnInterruptedEvent),
+                )
+            )
             raise
         except Exception as exc:
-            context_usage = _usage_from_context_length_error(str(exc))
-            if context_usage is not None:
-                self._remember_token_usage(context_usage)
-                self._emit("token_count", turn_id, usage=context_usage)
+            if isinstance(exc, ContextLengthExceeded) and exc.usage is not None:
+                self._remember_token_usage(exc.usage)
+                self._emit(TokenCountEvent(exc.usage, turn.turn_id))
             self._emit(
-                "turn_failed",
-                turn_id,
-                iteration=iteration,
-                error=str(exc),
-                error_type=type(exc).__name__,
+                TurnFailedEvent(
+                    turn.turn_id,
+                    turn.iteration,
+                    str(exc),
+                    type(exc).__name__,
+                    self._background_work_count(TurnFailedEvent),
+                )
             )
-            self._turn_running = False
             raise
+        finally:
+            self._stop_requested = False
+            self._idle.set()
+            self._idle = None
+
+    async def _sample(self, turn: "_TurnState") -> "ModelResponse":
+        for attempt in range(2):
+            if self._stop_requested:
+                raise TurnInterrupted()
+            if attempt == 0:
+                turn.iteration += 1
+            try:
+                return await self._complete_model_request(turn.turn_id, turn.iteration)
+            except ContextLengthExceeded as exc:
+                if attempt == 1:
+                    raise
+                if exc.usage is not None:
+                    self._remember_token_usage(exc.usage)
+                    self._emit(TokenCountEvent(exc.usage, turn.turn_id))
+                await self._compact_history(
+                    turn.turn_id,
+                    phase="context_length_exceeded",
+                    total_tokens=(
+                        exc.usage.get("total_tokens") if exc.usage is not None else None
+                    ),
+                    token_limit=exc.token_limit,
+                    prune_tool_results_on_context_error=True,
+                )
 
-    async def maybe_invoke(self, event: 'typing.Dict[str, object]') -> 'bool':
-        if self._turn_running:
-            return False
-        event_type = event.get("type")
-        if event_type == "exec_command_completed":
-            payload = {
-                "session_id": event.get("session_id"),
-                "exit_code": event.get("exit_code"),
-                "command": event.get("command"),
-            }
-            tag = "exec_command_completed"
-        elif event_type == "clock_tick":
-            payload = {
-                "period_m": event.get("period_m"),
-                "current_time": event.get("current_time"),
-            }
-            tag = "clock_tick"
-        else:
+    async def maybe_invoke(self, event: "typing.Dict[str, object]") -> "bool":
+        if self.is_running or not self.accepts_input:
             return False
+        tag = event["type"]
+        if not isinstance(tag, str) or not tag.isidentifier():
+            raise ValueError("invoke event type must be an identifier")
+        payload = {key: value for key, value in event.items() if key != "type"}
         text = (
             f"<{tag}>\n"
             f"{json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}\n"
             f"</{tag}>"
         )
-        self._turn_running = True
-        task = asyncio.create_task(self.run_turn([text]))
-        task.add_done_callback(
-            lambda task: None if task.cancelled() else task.exception()
-        )
+        await self.run_turn([text])
         return True
 
-    def shutdown(self) -> 'None':
-        if self._clock_manager is not None:
-            self._clock_manager.cancel()
+    def shutdown(self) -> "None":
+        if self.is_shutdown:
+            return
+        self.accepts_input = False
+        self.is_shutdown = True
+        for tool in self.tool_registry.tools():
+            tool.shutdown()
 
     async def _execute_tool_batch(
         self,
-        turn_id: 'str',
-        tool_calls: 'typing.List[ToolCall]',
-    ) -> 'typing.List[ToolResult]':
-        results: 'typing.List[ToolResult]' = []
-        parallel_batch: 'typing.List[ToolCall]' = []
-
+        turn_id: "str",
+        tool_calls: "typing.List[ToolCall]",
+    ) -> "None":
+        batches: "typing.List[typing.List[ToolCall]]" = []
+        parallel_batch: "typing.List[ToolCall]" = []
         for call in tool_calls:
             can_run_parallel = (
                 self._parallel_tool_calls
-                and self._tool_registry.supports_parallel(call.name)
+                and self.tool_registry.supports_parallel(call.name)
             )
             if can_run_parallel:
                 parallel_batch.append(call)
                 continue
 
             if parallel_batch:
-                prior_results = tuple(results)
-                results.extend(
-                    await asyncio.gather(
-                        *(
-                            self._run_single_tool(turn_id, batched_call, prior_results)
-                            for batched_call in parallel_batch
-                        )
-                    )
-                )
+                batches.append(parallel_batch)
             parallel_batch = []
-            results.append(await self._run_single_tool(turn_id, call, tuple(results)))
-
+            batches.append([call])
         if parallel_batch:
-            prior_results = tuple(results)
-            results.extend(
-                await asyncio.gather(
-                    *(
-                        self._run_single_tool(turn_id, batched_call, prior_results)
-                        for batched_call in parallel_batch
-                    )
-                )
+            batches.append(parallel_batch)
+
+        results: "typing.List[ToolResult]" = []
+        for batch in batches:
+            context = ToolContext(
+                turn_id=turn_id,
+                history=self.history,
+            )
+            if len(batch) == 1:
+                results.append(await self._run_single_tool(turn_id, batch[0], context))
+                continue
+            outcomes = await asyncio.gather(
+                *(self._run_single_tool(turn_id, call, context) for call in batch),
+                return_exceptions=True,
             )
-        return results
+            for outcome in outcomes:
+                if isinstance(outcome, BaseException):
+                    raise outcome
+                results.append(outcome)
+        self._append_history(self.tool_registry.follow_up_messages(results))
 
     async def _run_single_tool(
         self,
-        turn_id: 'str',
-        call: 'ToolCall',
-        prior_results: 'typing.Tuple[ToolResult, ...]' = (),
-    ) -> 'ToolResult':
-        payload: 'typing.Dict[str, object]' = {
-            "tool_name": call.name,
-            "call_id": call.call_id,
-            "call": call,
-        }
-        self._emit("tool_started", turn_id, **payload)
-        result = await self._tool_registry.execute(
-            call,
-            ToolContext(
-                turn_id=turn_id,
-                history=tuple(self._history) + prior_results,
-                collaboration_mode=self._context_manager.collaboration_mode,
-            ),
-        )
-        payload["result"] = result
-        payload["is_error"] = result.is_error
-        self._emit("tool_completed", turn_id, **payload)
+        turn_id: "str",
+        call: "ToolCall",
+        context: "ToolContext",
+    ) -> "ToolResult":
+        self._emit(ToolStartedEvent(turn_id, call))
+        result = await self.tool_registry.execute(call, context)
+        self._append_history([truncate_tool_result_for_history(result)])
+        self._emit(ToolCompletedEvent(turn_id, call, result))
         return result
 
-    def _emit(self, kind: 'str', turn_id: 'str', **payload: 'object') -> 'None':
-        if kind in TERMINAL_TURN_EVENTS:
-            payload["background_work_count"] = self._background_work_count(kind)
-        self._event_handler(
-            AgentEvent(kind=kind, turn_id=turn_id, payload=dict(payload))
-        )
+    def _emit(self, event: "TurnEvent") -> "None":
+        handlers = [self.event_handler]
+        if isinstance(event, (TerminalEvent, TurnStartedEvent, CompactStartedEvent)):
+            handlers = [
+                tool.handle_agent_event for tool in self.tool_registry.tools()
+            ] + handlers
+        for handler in handlers:
+            try:
+                handler(event)
+            except Exception as exc:
+                asyncio.get_running_loop().call_exception_handler(
+                    {
+                        "message": "Agent event observer failed: " + event.kind,
+                        "exception": exc,
+                    }
+                )
 
-    def _background_work_count(self, terminal_event: 'str') -> 'int':
-        manager: 'typing.Union[UnifiedExecManager, None]' = self._exec_manager
-        count = 0 if manager is None else manager.running_session_count()
-        clock_manager = self._clock_manager
-        if (
-            terminal_event == "turn_completed"
-            and clock_manager is not None
-            and clock_manager.enabled
-        ):
-            count += 1
-        return count
-
-    def _persist_history_items(
-        self,
-        items: 'typing.Iterable[ConversationItem]',
-    ) -> 'None':
-        recorder = self._rollout_recorder
-        if recorder is None:
-            return
+    def _background_work_count(self, event_type) -> "typing.Union[int, None]":
         try:
-            recorder.append_history_items(items)
-        except Exception:  # pragma: no cover - persistence should not break turns
-            return
+            return sum(
+                tool.background_work_count(
+                    event_type in (TurnCompletedEvent, CompactCompletedEvent)
+                )
+                for tool in self.tool_registry.tools()
+            )
+        except Exception as exc:
+            asyncio.get_running_loop().call_exception_handler(
+                {
+                    "message": "Agent background-work observer failed: "
+                    + event_type.kind,
+                    "exception": exc,
+                }
+            )
+            return None
 
-    def _record_model_response_items(
+    def _append_history(
         self,
-        items: 'typing.Iterable[object]',
-    ) -> 'typing.Tuple[typing.Tuple[ConversationItem, ...], typing.List[ToolCall], typing.Union[str, None]]':
-        persisted_response_items: 'typing.List[ConversationItem]' = []
-        tool_calls: 'typing.List[ToolCall]' = []
-        last_assistant_message = None
-        for item in items:
-            if not isinstance(item, (AssistantMessage, ToolCall, ReasoningItem)):
-                continue
-            self._history.append(item)
-            persisted_response_items.append(item)
-            if isinstance(item, AssistantMessage):
-                last_assistant_message = item.text
-            elif isinstance(item, ToolCall):
-                tool_calls.append(item)
-        self._persist_history_items(persisted_response_items)
-        return tuple(persisted_response_items), tool_calls, last_assistant_message
-
-    def _handle_model_stream_event(self, turn_id: 'str', event: 'ModelStreamEvent') -> 'None':
-        if event.kind == "token_count":
-            self._remember_token_usage(event.payload.get("usage"))
-        if event.kind == "assistant_delta":
-            self._emit("assistant_delta", turn_id, **event.payload)
-        elif event.kind == "tool_call":
-            self._emit("tool_called", turn_id, **event.payload)
-        elif event.kind == "token_count":
-            self._emit("token_count", turn_id, **event.payload)
-        elif event.kind == "stream_error":
-            self._emit("stream_error", turn_id, **event.payload)
-
-    def _remember_token_usage(self, usage: 'object') -> 'None':
+        items: "typing.Iterable[ConversationItem]",
+    ) -> "None":
+        items = tuple(items)
+        if self._rollout_recorder is not None:
+            self._rollout_recorder.append_history_items(items, self._history)
+        self._history.extend(items)
+
+    def _handle_model_stream_event(self, turn_id: "str", event: "ModelEvent") -> "None":
+        if isinstance(event, TokenCountEvent):
+            self._remember_token_usage(event.usage)
+        self._emit(replace(event, turn_id=turn_id))
+
+    def _remember_token_usage(self, usage: "object") -> "None":
         if not isinstance(usage, dict):
             return
         try:
@@ -432,67 +481,45 @@ class Agent:
 
     async def _complete_model_request(
         self,
-        turn_id: 'str',
-        iteration: 'int',
-    ) -> 'typing.Any':
-        attempted_context_compact = False
-        while True:
-            prompt = self._context_manager.build_prompt(
-                self._history,
-                self._tool_registry.model_visible_specs(),
-                self._parallel_tool_calls,
-                turn_id=turn_id,
-            )
-            self._emit(
-                "model_called",
-                turn_id,
-                iteration=iteration,
-                history_size=len(prompt.input),
-                tool_count=len(prompt.tools),
+        turn_id: "str",
+        iteration: "int",
+    ) -> "ModelResponse":
+        prompt = self.context_manager.build_prompt(
+            self._history,
+            self.tool_registry.model_visible_specs(),
+            self._parallel_tool_calls,
+            turn_id=turn_id,
+        )
+        self._emit(
+            ModelCalledEvent(turn_id, iteration, len(prompt.input), len(prompt.tools))
+        )
+        try:
+            return await self.model_client.complete(
+                prompt,
+                lambda event: self._handle_model_stream_event(turn_id, event),
             )
-            try:
-                return await self._model_client.complete(
-                    prompt,
-                    lambda event: self._handle_model_stream_event(turn_id, event),
+        except ResponsesIncompleteError as exc:
+            if exc.reason == "max_output_tokens":
+                self._append_history(
+                    item
+                    for item in exc.partial_items
+                    if isinstance(item, (AssistantMessage, ReasoningItem))
                 )
-            except Exception as exc:
-                error_message = str(exc)
-                if (
-                    not _is_context_length_error_message(error_message)
-                    or attempted_context_compact
-                ):
-                    raise
-                attempted_context_compact = True
-                context_usage = _usage_from_context_length_error(error_message)
-                if context_usage is not None:
-                    self._remember_token_usage(context_usage)
-                    self._emit("token_count", turn_id, usage=context_usage)
-                await self._run_auto_compact(
-                    turn_id,
-                    phase="context_length_exceeded",
-                    total_tokens=(
-                        context_usage.get("total_tokens")
-                        if context_usage is not None
-                        else None
-                    ),
-                    token_limit=_context_length_error_token_limit(error_message),
-                    prune_tool_results_on_context_error=True,
-                )
-                self._raise_if_interrupt_requested(turn_id, iteration)
+            raise
 
     async def _maybe_auto_compact(
         self,
-        turn_id: 'str',
-        phase: 'str',
-    ) -> 'None':
-        limit = self._auto_compact_token_limit
+        turn_id: "str",
+        phase: "str",
+    ) -> "None":
+        limit = self.context_manager.resolve_auto_compact_token_limit()
         total_tokens = self._last_total_usage_tokens
         if limit is None or total_tokens is None:
             return
         if total_tokens < limit or not self._history:
             return
 
-        await self._run_auto_compact(
+        await self._compact_history(
             turn_id,
             phase=phase,
             total_tokens=total_tokens,
@@ -500,134 +527,73 @@ class Agent:
             prune_tool_results_on_context_error=True,
         )
 
-    async def _run_auto_compact(
+    async def _compact_history(
         self,
-        turn_id: 'str',
-        phase: 'str',
-        total_tokens: 'typing.Union[int, None]' = None,
-        token_limit: 'typing.Union[int, None]' = None,
-        prune_tool_results_on_context_error: 'bool' = False,
-    ) -> 'None':
-        from .utils.compactor import compact_agent
-
-        payload: 'typing.Dict[str, object]' = {"phase": phase}
-        if total_tokens is not None:
-            payload["total_tokens"] = total_tokens
-        if token_limit is not None:
-            payload["token_limit"] = token_limit
-        self._emit(
-            "auto_compact_started",
-            turn_id,
-            **payload,
+        turn_id: "str",
+        phase: "str",
+        total_tokens: "typing.Union[int, None]" = None,
+        token_limit: "typing.Union[int, None]" = None,
+        prune_tool_results_on_context_error: "bool" = False,
+    ) -> "typing.Union[CompactResult, None]":
+        from .utils.compactor import compact_history
+
+        if not self._history:
+            return None
+        details = (turn_id, phase, total_tokens, token_limit)
+        started_type = (
+            CompactStartedEvent if phase == "manual" else AutoCompactStartedEvent
         )
+        self._emit(started_type(*details))
 
-        def handle_compact_stream_event(event: 'ModelStreamEvent') -> 'None':
-            if event.kind == "stream_error":
-                self._emit("stream_error", turn_id, **event.payload)
+        def handle_compact_stream_event(event: "ModelEvent") -> "None":
+            if isinstance(event, StreamErrorEvent) or (
+                phase == "manual" and isinstance(event, TokenCountEvent)
+            ):
+                self._emit(replace(event, turn_id=turn_id))
 
         try:
-            compact_result = await compact_agent(
-                self,
+            recorder = self._rollout_recorder
+            compact_result = await compact_history(
+                self.history,
+                self.model_client,
+                self.context_manager,
                 handle_compact_stream_event,
                 prune_tool_results_on_context_error,
-            )
-        except Exception as exc:
-            failed_payload = dict(payload)
-            failed_payload.update(
-                {
-                    "error": str(exc),
-                    "error_type": type(exc).__name__,
-                }
-            )
-            self._emit(
-                "auto_compact_failed",
+                str(recorder.rollout_path) if recorder is not None else None,
                 turn_id,
-                **failed_payload,
             )
+            if recorder is not None:
+                recorder.append_compacted_history(compact_result.history, self._history)
+            self._history = list(compact_result.history)
+            self._last_total_usage_tokens = None
+        except Exception as exc:
+            if phase == "manual":
+                self._emit(
+                    CompactFailedEvent(
+                        *details,
+                        str(exc),
+                        type(exc).__name__,
+                        self._background_work_count(CompactFailedEvent),
+                    )
+                )
+            else:
+                self._emit(
+                    AutoCompactFailedEvent(*details, str(exc), type(exc).__name__)
+                )
             raise
 
-        self._last_total_usage_tokens = None
-        if compact_result is None:
-            return
-        completed_payload = dict(payload)
-        completed_payload.update(
-            {
-                "original_item_count": compact_result.original_item_count,
-                "retained_item_count": compact_result.retained_item_count,
-                "summary": compact_result.display_text(),
-            }
+        completed = details + (
+            compact_result.original_item_count,
+            compact_result.retained_item_count,
+            compact_result.pruned_tool_results,
         )
-        if compact_result.pruned_tool_results:
-            completed_payload["pruned_tool_results"] = compact_result.pruned_tool_results
-        self._emit(
-            "auto_compact_completed",
-            turn_id,
-            **completed_payload,
-        )
-
-    def _build_follow_up_messages(
-        self,
-        tool_results: 'typing.List[ToolResult]',
-    ) -> 'typing.List[UserMessage]':
-        follow_ups: 'typing.List[UserMessage]' = []
-        for result in tool_results:
-            statuses = None
-            if (
-                result.name == "wait_agent"
-                and not result.is_error
-                and isinstance(result.output, dict)
-            ):
-                statuses = result.output.get("status")
-            if isinstance(statuses, dict):
-                for agent_id, status in statuses.items():
-                    if isinstance(agent_id, str) and isinstance(status, dict):
-                        payload = {
-                            "agent_id": agent_id,
-                            "status": status,
-                        }
-                        follow_ups.append(
-                            UserMessage(
-                                text=(
-                                    "<subagent_notification>\n"
-                                    f"{json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}\n"
-                                    "</subagent_notification>"
-                                )
-                            )
-                        )
-        return follow_ups
-
-
-def _usage_from_context_length_error(
-    message: 'str',
-) -> 'typing.Union[typing.Dict[str, int], None]':
-    if not _is_context_length_error_message(message):
-        return None
-
-    requested_match = _REQUESTED_TOKENS_RE.search(message)
-    if requested_match is None:
-        return None
-
-    usage = {"total_tokens": _parse_token_count(requested_match.group(1))}
-    split_match = _REQUESTED_TOKEN_SPLIT_RE.search(message)
-    if split_match is not None:
-        usage["input_tokens"] = _parse_token_count(split_match.group(1))
-        usage["output_tokens"] = _parse_token_count(split_match.group(2))
-    else:
-        usage["input_tokens"] = usage["total_tokens"]
-    return usage
-
-
-def _is_context_length_error_message(message: 'str') -> 'bool':
-    lower = message.lower()
-    return any(marker in lower for marker in _CONTEXT_LENGTH_ERROR_MARKERS)
-
-
-def _context_length_error_token_limit(message: 'str') -> 'typing.Union[int, None]':
-    limit_match = _MAX_CONTEXT_TOKENS_RE.search(message)
-    if limit_match is None:
-        return None
-    return _parse_token_count(limit_match.group(1))
-
-
-def _parse_token_count(value: 'str') -> 'int':
-    return int(value.replace(",", ""))
+        if phase == "manual":
+            self._emit(
+                CompactCompletedEvent(
+                    *completed,
+                    self._background_work_count(CompactCompletedEvent),
+                )
+            )
+        else:
+            self._emit(AutoCompactCompletedEvent(*completed))
+        return compact_result