python-codex 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

pycodex/__init__.py

@@ -2,12 +2,13 @@ from .compat import patch_asyncio
 
 patch_asyncio()
 
-from .agent import AgentLoop
+from .agent import Agent
 from .context import ContextConfig, ContextManager
 from .model import (
     ModelClient,
     NOOP_MODEL_STREAM_EVENT_HANDLER,
     ResponsesApiError,
+    ResponsesIncompleteError,
     ResponsesModelClient,
     ResponsesProviderConfig,
 )
@@ -26,14 +27,14 @@ from .protocol import (
     TurnResult,
     UserMessage,
 )
-from .runtime import AgentRuntime
+from .runtime import CliSubmissionQueue
 from .runtime_services import (
     PlanStore,
     RequestPermissionsManager,
     RequestUserInputManager,
     SubAgentManager,
-    create_runtime_environment,
-    get_runtime_environment,
+    create_agent_runtime_environment,
+    get_agent_runtime_environment,
 )
 from .tools import (
     ApplyPatchTool,
@@ -90,13 +91,13 @@ def debug(stop: 'bool' = False):
 
 __all__ = [
     "AgentEvent",
-    "AgentLoop",
-    "AgentRuntime",
+    "Agent",
+    "CliSubmissionQueue",
     "ApplyPatchTool",
     "AssistantMessage",
     "BaseTool",
     "CloseAgentTool",
-    "create_runtime_environment",
+    "create_agent_runtime_environment",
     "CodeModeManager",
     "ContextConfig",
     "ContextManager",
@@ -120,6 +121,7 @@ __all__ = [
     "RequestUserInputManager",
     "ResumeAgentTool",
     "ResponsesApiError",
+    "ResponsesIncompleteError",
     "ResponsesModelClient",
     "ResponsesProviderConfig",
     "SendInputTool",
@@ -142,5 +144,5 @@ __all__ = [
     "WaitTool",
     "WebSearchTool",
     "WriteStdinTool",
-    "get_runtime_environment",
+    "get_agent_runtime_environment",
 ]

pycodex/agent.py

@@ -1,6 +1,7 @@
 
 import asyncio
 import json
+import re
 from typing import Callable
 
 from .context import ContextManager
@@ -22,17 +23,36 @@ import typing
 
 if typing.TYPE_CHECKING:
     from .utils.session_persist import SessionRolloutRecorder
+    from .runtime_services import AgentRuntimeEnvironment
 
 
 EventHandler = Callable[[AgentEvent], None]
-NOOP_EVENT_HANDLER: 'EventHandler' = lambda _event: 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",
+)
 
 
 class TurnInterrupted(RuntimeError):
     pass
 
 
-class AgentLoop:
+class Agent:
     """Minimal Python port of Codex's turn loop.
 
     The core idea mirrors the Rust implementation:
@@ -47,9 +67,10 @@ class AgentLoop:
         tool_registry: 'ToolRegistry',
         context_manager: 'typing.Union[ContextManager, None]' = None,
         parallel_tool_calls: 'bool' = True,
-        event_handler: 'EventHandler' = NOOP_EVENT_HANDLER,
+        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
@@ -58,6 +79,11 @@ class AgentLoop:
         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()
+        )
+        self._last_total_usage_tokens: 'typing.Union[int, None]' = None
+        self.runtime_environment = runtime_environment
         self.interrupt_asap = False
 
     @property
@@ -65,7 +91,7 @@ class AgentLoop:
         return tuple(self._history)
 
     def set_event_handler(
-        self, event_handler: 'EventHandler' = NOOP_EVENT_HANDLER
+        self, event_handler: 'EventHandler' = BASE_EVENT_HANDLER
     ) -> 'None':
         self._event_handler = event_handler
 
@@ -81,6 +107,11 @@ class AgentLoop:
     ) -> 'None':
         self._rollout_recorder = rollout_recorder
 
+    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',
@@ -101,8 +132,6 @@ class AgentLoop:
         turn_id = turn_id or uuid7_string()
         self.interrupt_asap = False
         new_user_messages = [UserMessage(text=text) for text in texts]
-        self._history.extend(new_user_messages)
-        self._persist_history_items(new_user_messages)
 
         self._emit(
             "turn_started",
@@ -110,6 +139,9 @@ class AgentLoop:
             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)
 
         last_assistant_message: 'typing.Union[str, None]' = None
         final_response_items: 'typing.Tuple[\n    typing.Union[typing.Union[AssistantMessage, ToolCall], ReasoningItem], ...\n]' = ()
@@ -122,23 +154,11 @@ class AgentLoop:
                     iteration,
                     output_text=last_assistant_message,
                 )
+                await self._maybe_auto_compact(turn_id, phase="mid_turn")
                 iteration += 1
-                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",
+                response = await self._complete_model_request(
                     turn_id,
-                    iteration=iteration,
-                    history_size=len(prompt.input),
-                    tool_count=len(prompt.tools),
-                )
-                response = await self._model_client.complete(
-                    prompt,
-                    lambda event: self._handle_model_stream_event(turn_id, event),
+                    iteration,
                 )
                 final_response_items = tuple(response.items)
                 self._emit(
@@ -193,6 +213,10 @@ class AgentLoop:
         except TurnInterrupted:
             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)
             self._emit(
                 "turn_failed",
                 turn_id,
@@ -287,6 +311,8 @@ class AgentLoop:
             return
 
     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":
@@ -296,6 +322,149 @@ class AgentLoop:
         elif event.kind == "stream_error":
             self._emit("stream_error", turn_id, **event.payload)
 
+    def _remember_token_usage(self, usage: 'object') -> 'None':
+        if not isinstance(usage, dict):
+            return
+        try:
+            self._last_total_usage_tokens = int(usage["total_tokens"])
+        except (KeyError, TypeError, ValueError):
+            return
+
+    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),
+            )
+            try:
+                return await self._model_client.complete(
+                    prompt,
+                    lambda event: self._handle_model_stream_event(turn_id, event),
+                )
+            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)
+
+    async def _maybe_auto_compact(
+        self,
+        turn_id: 'str',
+        phase: 'str',
+    ) -> 'None':
+        limit = self._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(
+            turn_id,
+            phase=phase,
+            total_tokens=total_tokens,
+            token_limit=limit,
+            prune_tool_results_on_context_error=True,
+        )
+
+    async def _run_auto_compact(
+        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,
+        )
+
+        def handle_compact_stream_event(event: 'ModelStreamEvent') -> 'None':
+            if event.kind == "stream_error":
+                self._emit("stream_error", turn_id, **event.payload)
+
+        try:
+            compact_result = await compact_agent(
+                self,
+                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",
+                turn_id,
+                **failed_payload,
+            )
+            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(),
+            }
+        )
+        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]',
@@ -326,3 +495,39 @@ class AgentLoop:
                             )
                         )
         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(",", ""))