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

pycodex/tools/apply_patch_tool.py

@@ -11,6 +11,7 @@ Expected behavior:
   same success/error text shape Codex expects.
 """
 
+import typing
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -18,7 +19,6 @@ from loguru import logger
 
 from ..protocol import JSONValue
 from .base_tool import BaseTool, ToolContext
-import typing
 
 APPLY_PATCH_LARK_GRAMMAR = """start: begin_patch hunk+ end_patch
 begin_patch: \"*** Begin Patch\" LF
@@ -46,28 +46,36 @@ class ApplyPatchError(RuntimeError):
     pass
 
 
-@dataclass(frozen=True, )
+@dataclass(
+    frozen=True,
+)
 class _AddFileOp:
-    path: 'str'
-    content: 'str'
+    path: "str"
+    content: "str"
 
 
-@dataclass(frozen=True, )
+@dataclass(
+    frozen=True,
+)
 class _DeleteFileOp:
-    path: 'str'
+    path: "str"
 
 
-@dataclass(frozen=True, )
+@dataclass(
+    frozen=True,
+)
 class _UpdateSection:
-    lines: 'typing.Tuple[str, ...]'
-    anchor_end_of_file: 'bool' = False
+    lines: "typing.Tuple[str, ...]"
+    anchor_end_of_file: "bool" = False
 
 
-@dataclass(frozen=True, )
+@dataclass(
+    frozen=True,
+)
 class _UpdateFileOp:
-    path: 'str'
-    move_to: 'typing.Union[str, None]'
-    sections: 'typing.Tuple[_UpdateSection, ...]'
+    path: "str"
+    move_to: "typing.Union[str, None]"
+    sections: "typing.Tuple[_UpdateSection, ...]"
 
 
 class ApplyPatchTool(BaseTool):
@@ -84,20 +92,26 @@ class ApplyPatchTool(BaseTool):
     }
     supports_parallel = False
 
-    def __init__(self, cwd: 'typing.Union[typing.Union[str, Path], None]' = None) -> 'None':
+    def __init__(
+        self, cwd: "typing.Union[typing.Union[str, Path], None]" = None
+    ) -> "None":
         self._workspace_root = Path(cwd or Path.cwd()).resolve()
 
-    async def run(self, context: 'ToolContext', args: 'JSONValue') -> 'JSONValue':
+    async def run(self, context: "ToolContext", args: "JSONValue") -> "JSONValue":
         del context
         patch_text = str(args)
-        logger.debug("apply_patch workspace={} bytes={}", self._workspace_root, len(patch_text))
+        logger.debug(
+            "apply_patch workspace={} bytes={}", self._workspace_root, len(patch_text)
+        )
         try:
             operations = self._parse_patch(patch_text)
             return self._format_result(self._apply_operations(operations), exit_code=0)
         except ApplyPatchError as exc:
             return self._format_result(str(exc), exit_code=1)
 
-    def _parse_patch(self, patch_text: 'str') -> 'typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]':
+    def _parse_patch(
+        self, patch_text: "str"
+    ) -> "typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]":
         lines = patch_text.splitlines()
         if not lines:
             raise ApplyPatchError("patch rejected: empty patch")
@@ -106,7 +120,7 @@ class ApplyPatchTool(BaseTool):
                 "apply_patch verification failed: missing '*** Begin Patch' header"
             )
 
-        operations: 'typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]' = []
+        operations: "typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]" = ([])
         index = 1
         while index < len(lines):
             line = lines[index]
@@ -123,7 +137,7 @@ class ApplyPatchTool(BaseTool):
             if line.startswith("*** Add File: "):
                 path = line[len("*** Add File: ") :]
                 index += 1
-                content_lines: 'typing.List[str]' = []
+                content_lines: "typing.List[str]" = []
                 while index < len(lines) and not lines[index].startswith("*** "):
                     entry = lines[index]
                     if not entry.startswith("+"):
@@ -136,7 +150,9 @@ class ApplyPatchTool(BaseTool):
                     raise ApplyPatchError(
                         f"apply_patch verification failed: add for {path} is missing file content"
                     )
-                operations.append(_AddFileOp(path=path, content=self._join_lines(content_lines)))
+                operations.append(
+                    _AddFileOp(path=path, content=self._join_lines(content_lines))
+                )
                 continue
 
             if line.startswith("*** Delete File: "):
@@ -153,8 +169,8 @@ class ApplyPatchTool(BaseTool):
                     move_to = lines[index][len("*** Move to: ") :]
                     index += 1
 
-                sections: 'typing.List[_UpdateSection]' = []
-                current_lines: 'typing.List[str]' = []
+                sections: "typing.List[_UpdateSection]" = []
+                current_lines: "typing.List[str]" = []
                 saw_hunk_header = False
                 anchor_end_of_file = False
                 while index < len(lines):
@@ -216,14 +232,16 @@ class ApplyPatchTool(BaseTool):
                 f"apply_patch verification failed: {line!r} is not a valid hunk header"
             )
 
-        raise ApplyPatchError("apply_patch verification failed: missing '*** End Patch' footer")
+        raise ApplyPatchError(
+            "apply_patch verification failed: missing '*** End Patch' footer"
+        )
 
     def _apply_operations(
         self,
-        operations: 'typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]',
-    ) -> 'str':
-        preview: 'typing.Dict[Path, typing.Union[str, None]]' = {}
-        summaries: 'typing.Dict[Path, str]' = {}
+        operations: "typing.List[typing.Union[typing.Union[_AddFileOp, _DeleteFileOp], _UpdateFileOp]]",
+    ) -> "str":
+        preview: "typing.Dict[Path, typing.Union[str, None]]" = {}
+        summaries: "typing.Dict[Path, str]" = {}
 
         for operation in operations:
             if isinstance(operation, _AddFileOp):
@@ -253,7 +271,9 @@ class ApplyPatchTool(BaseTool):
         self._write_preview(preview)
         return self._format_success(summaries)
 
-    def _read_preview_file(self, path: 'Path', preview: 'typing.Dict[Path, typing.Union[str, None]]') -> 'str':
+    def _read_preview_file(
+        self, path: "Path", preview: "typing.Dict[Path, typing.Union[str, None]]"
+    ) -> "str":
         if path in preview:
             content = preview[path]
             if content is None:
@@ -270,10 +290,10 @@ class ApplyPatchTool(BaseTool):
 
     def _apply_update(
         self,
-        path: 'Path',
-        original_text: 'str',
-        sections: 'typing.Tuple[_UpdateSection, ...]',
-    ) -> 'str':
+        path: "Path",
+        original_text: "str",
+        sections: "typing.Tuple[_UpdateSection, ...]",
+    ) -> "str":
         lines = original_text.splitlines()
         cursor = 0
         for section in sections:
@@ -282,7 +302,9 @@ class ApplyPatchTool(BaseTool):
             if not old_block and not new_block:
                 continue
 
-            match_index = self._find_match(lines, old_block, cursor, section.anchor_end_of_file)
+            match_index = self._find_match(
+                lines, old_block, cursor, section.anchor_end_of_file
+            )
             if match_index is None:
                 raise ApplyPatchError(
                     "apply_patch verification failed: Failed to find expected lines in "
@@ -294,11 +316,11 @@ class ApplyPatchTool(BaseTool):
 
     def _find_match(
         self,
-        lines: 'typing.List[str]',
-        old_block: 'typing.List[str]',
-        cursor: 'int',
-        anchor_end_of_file: 'bool',
-    ) -> 'typing.Union[int, None]':
+        lines: "typing.List[str]",
+        old_block: "typing.List[str]",
+        cursor: "int",
+        anchor_end_of_file: "bool",
+    ) -> "typing.Union[int, None]":
         if anchor_end_of_file:
             start = len(lines) - len(old_block)
             if start >= 0 and lines[start : start + len(old_block)] == old_block:
@@ -317,7 +339,9 @@ class ApplyPatchTool(BaseTool):
                 return start
         return None
 
-    def _write_preview(self, preview: 'typing.Dict[Path, typing.Union[str, None]]') -> 'None':
+    def _write_preview(
+        self, preview: "typing.Dict[Path, typing.Union[str, None]]"
+    ) -> "None":
         for path, content in preview.items():
             if content is None:
                 if path.exists():
@@ -326,7 +350,7 @@ class ApplyPatchTool(BaseTool):
             path.parent.mkdir(parents=True, exist_ok=True)
             path.write_text(content, encoding="utf-8")
 
-    def _format_success(self, summaries: 'typing.Dict[Path, str]') -> 'str':
+    def _format_success(self, summaries: "typing.Dict[Path, str]") -> "str":
         buckets = {"A": [], "M": [], "D": []}
         for path, status in summaries.items():
             buckets[status].append(self._display_path(path))
@@ -336,26 +360,23 @@ class ApplyPatchTool(BaseTool):
                 lines.append(f"{status} {rel_path}")
         return " ".join(lines) + "\n"
 
-    def _format_result(self, output: 'str', exit_code: 'int') -> 'str':
+    def _format_result(self, output: "str", exit_code: "int") -> "str":
         return (
-            f"Exit code: {exit_code}\n"
-            "Wall time: 0 seconds\n"
-            "Output:\n"
-            f"{output}"
+            f"Exit code: {exit_code}\n" "Wall time: 0 seconds\n" "Output:\n" f"{output}"
         )
 
-    def _resolve_workspace_path(self, path_text: 'str') -> 'Path':
+    def _resolve_workspace_path(self, path_text: "str") -> "Path":
         path = Path(path_text).expanduser()
         resolved = path if path.is_absolute() else self._workspace_root / path
         return resolved.resolve()
 
-    def _display_path(self, path: 'Path') -> 'str':
+    def _display_path(self, path: "Path") -> "str":
         try:
             return path.relative_to(self._workspace_root).as_posix()
         except ValueError:
             return str(path)
 
-    def _join_lines(self, lines: 'typing.List[str]') -> 'str':
+    def _join_lines(self, lines: "typing.List[str]") -> "str":
         if not lines:
             return ""
         return "\n".join(lines) + "\n"

pycodex/tools/base_tool.py

@@ -10,47 +10,61 @@ Expected behavior:
   model, and dispatches `ToolCall` executions back into `ToolResult`s.
 """
 
-import inspect
+import asyncio
 import json
+import traceback
+import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-import traceback
 
-from ..protocol import ConversationItem, JSONDict, JSONValue, ToolCall, ToolResult, ToolSpec
+from ..events import Event
+from ..protocol import (
+    ConversationItem,
+    JSONDict,
+    JSONValue,
+    ToolCall,
+    ToolResult,
+    ToolSpec,
+    UserMessage,
+)
+from ..runtime_services import AgentRuntimeEnvironment
 from ..utils import get_debug_dir
-import typing
 
+if typing.TYPE_CHECKING:
+    from ..agent import Agent
 
-@dataclass(frozen=True, )
+
+@dataclass(
+    frozen=True,
+)
 class ToolContext:
-    turn_id: 'str'
-    history: 'typing.Tuple[ConversationItem, ...]'
-    collaboration_mode: 'str' = "default"
+    turn_id: "str"
+    history: "typing.Tuple[ConversationItem, ...]"
 
 
 class StructuredToolOutput:
     def __init__(
         self,
-        output: 'JSONValue',
-        content_items: 'typing.Union[typing.Union[typing.Tuple[JSONDict, ...], typing.List[JSONDict]], None]' = None,
-        success: 'typing.Union[bool, None]' = None,
-    ) -> 'None':
+        output: "JSONValue",
+        content_items: "typing.Union[typing.Union[typing.Tuple[JSONDict, ...], typing.List[JSONDict]], None]" = None,
+        success: "typing.Union[bool, None]" = None,
+    ) -> "None":
         self.output = output
         self.content_items = None if content_items is None else tuple(content_items)
         self.success = success
 
 
 class BaseTool(ABC):
-    name: 'str'
-    description: 'str'
-    input_schema: 'typing.Union[JSONDict, None]' = None
-    tool_type: 'str' = "function"
-    format: 'typing.Union[JSONDict, None]' = None
-    options: 'typing.Union[JSONDict, None]' = None
-    output_schema: 'typing.Union[JSONDict, None]' = None
-    supports_parallel: 'bool' = True
-
-    def spec(self) -> 'ToolSpec':
+    name: "str"
+    description: "str"
+    input_schema: "typing.Union[JSONDict, None]" = None
+    tool_type: "str" = "function"
+    format: "typing.Union[JSONDict, None]" = None
+    options: "typing.Union[JSONDict, None]" = None
+    output_schema: "typing.Union[JSONDict, None]" = None
+    supports_parallel: "bool" = True
+
+    def spec(self) -> "ToolSpec":
         return ToolSpec(
             name=self.name,
             description=self.description,
@@ -62,29 +76,51 @@ class BaseTool(ABC):
             supports_parallel=self.supports_parallel,
         )
 
-    def serialize(self) -> 'JSONDict':
+    def serialize(self) -> "JSONDict":
         return self.spec().serialize()
 
     @abstractmethod
-    async def run(self, context: 'ToolContext', args: 'JSONValue') -> 'JSONValue':
+    async def run(
+        self, context: "ToolContext", args: "JSONValue"
+    ) -> "typing.Union[JSONValue, StructuredToolOutput]":
         raise NotImplementedError
 
+    def follow_up_messages(
+        self, output: "JSONValue"
+    ) -> "typing.Tuple[UserMessage, ...]":
+        return ()
+
+    def bind_agent(self, agent: "Agent") -> "None":
+        pass
+
+    def handle_agent_event(self, event: "Event") -> "None":
+        pass
+
+    def shutdown(self) -> "None":
+        pass
+
+    def background_work_count(self, after_reply: "bool") -> "int":
+        return 0
+
 
 class ToolRegistry:
-    def __init__(self) -> 'None':
-        self._tools: 'typing.Dict[str, BaseTool]' = {}
+    def __init__(
+        self, runtime_environment: "typing.Union[AgentRuntimeEnvironment, None]" = None
+    ) -> "None":
+        self._tools: "typing.Dict[str, BaseTool]" = {}
+        self.runtime_environment = runtime_environment or AgentRuntimeEnvironment()
 
-    def register(self, tool: 'BaseTool') -> 'None':
+    def register(self, tool: "BaseTool") -> "None":
         self._tools[tool.name] = tool
 
-    def model_visible_specs(self) -> 'typing.List[ToolSpec]':
+    def model_visible_specs(self) -> "typing.List[ToolSpec]":
         return [tool.spec() for tool in self._tools.values()]
 
-    def supports_parallel(self, tool_name: 'str') -> 'bool':
+    def supports_parallel(self, tool_name: "str") -> "bool":
         tool = self._tools.get(tool_name)
         return False if tool is None else tool.supports_parallel
 
-    async def execute(self, call: 'ToolCall', context: 'ToolContext') -> 'ToolResult':
+    async def execute(self, call: "ToolCall", context: "ToolContext") -> "ToolResult":
         tool = self._tools.get(call.name)
         if tool is None:
             return ToolResult(
@@ -96,11 +132,7 @@ class ToolRegistry:
             )
 
         try:
-            maybe_result = tool.run(context, call.arguments)
-            if inspect.isawaitable(maybe_result):
-                output = await maybe_result
-            else:
-                output = maybe_result
+            output = await tool.run(context, call.arguments)
             if isinstance(output, StructuredToolOutput):
                 return ToolResult(
                     call_id=call.call_id,
@@ -116,10 +148,14 @@ class ToolRegistry:
                 output=output,
                 tool_type=call.tool_type,
             )
-        except Exception as exc:  # pragma: no cover - defensive wrapper
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
             debug_dir = get_debug_dir()
             if debug_dir is not None:
-                with (debug_dir / "tool_errors.jsonl").open("a", encoding="utf-8") as handle:
+                with (debug_dir / "tool_errors.jsonl").open(
+                    "a", encoding="utf-8"
+                ) as handle:
                     handle.write(
                         json.dumps(
                             {
@@ -140,19 +176,30 @@ class ToolRegistry:
                 tool_type=call.tool_type,
             )
 
-    def __contains__(self, tool_name: 'str') -> 'bool':
+    def follow_up_messages(
+        self, results: "typing.Iterable[ToolResult]"
+    ) -> "typing.Tuple[UserMessage, ...]":
+        messages = []
+        for result in results:
+            if not result.is_error:
+                messages.extend(
+                    self._tools[result.name].follow_up_messages(result.output)
+                )
+        return tuple(messages)
+
+    def __contains__(self, tool_name: "str") -> "bool":
         return tool_name in self._tools
 
-    def __len__(self) -> 'int':
+    def __len__(self) -> "int":
         return len(self._tools)
 
-    def names(self) -> 'typing.Tuple[str, ...]':
+    def names(self) -> "typing.Tuple[str, ...]":
         return tuple(self._tools)
 
-    def get_tool(self, tool_name: 'str') -> 'typing.Union[BaseTool, None]':
+    def get_tool(self, tool_name: "str") -> "typing.Union[BaseTool, None]":
         return self._tools.get(tool_name)
 
-    def tools(self) -> 'typing.Tuple[BaseTool, ...]':
+    def tools(self) -> "typing.Tuple[BaseTool, ...]":
         return tuple(self._tools.values())
 
 

pycodex/tools/clock_tool.py

@@ -6,13 +6,19 @@ cancels the pending countdown, and every successful reply starts a fresh one.
 """
 
 import asyncio
-from datetime import datetime
 import math
+import typing
+from datetime import datetime
 
+from ..events import (
+    CompactCompletedEvent,
+    CompactStartedEvent,
+    Event,
+    TurnCompletedEvent,
+    TurnStartedEvent,
+)
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
-import typing
-
 
 CLOCK_STATE_OUTPUT_SCHEMA = {
     "type": "object",
@@ -34,25 +40,25 @@ CLOCK_STATE_OUTPUT_SCHEMA = {
 }
 
 
-def _current_time() -> 'str':
+def _current_time() -> "str":
     return datetime.now().astimezone().isoformat(timespec="seconds")
 
 
 class ClockManager:
-    def __init__(self, seconds_per_minute: 'float' = 60.0) -> 'None':
+    def __init__(self, seconds_per_minute: "float" = 60.0) -> "None":
         self._seconds_per_minute = seconds_per_minute
-        self._period_m: 'typing.Union[float, None]' = None
-        self._timer_task: 'typing.Union[asyncio.Task, None]' = None
+        self._period_m: "typing.Union[float, None]" = None
+        self._timer_task: "typing.Union[asyncio.Task, None]" = None
         self._generation = 0
-        self._notify_hook: 'typing.Union[typing.Callable[[typing.Dict[str, object]], typing.Awaitable[typing.Any]], None]' = None
+        self._notify_hook: "typing.Union[typing.Callable[[typing.Dict[str, object]], typing.Awaitable[typing.Any]], None]" = (None)
 
     def set_notify_hook(
         self,
-        callback: 'typing.Union[typing.Callable[[typing.Dict[str, object]], typing.Awaitable[typing.Any]], None]',
-    ) -> 'None':
+        callback: "typing.Union[typing.Callable[[typing.Dict[str, object]], typing.Awaitable[typing.Any]], None]",
+    ) -> "None":
         self._notify_hook = callback
 
-    def set_period(self, value: 'object') -> 'JSONDict':
+    def set_period(self, value: "object") -> "JSONDict":
         self._cancel_pending()
         if value is None:
             self._period_m = None
@@ -65,20 +71,20 @@ class ClockManager:
         self._period_m = period_m
         return self.snapshot()
 
-    def snapshot(self) -> 'JSONDict':
+    def snapshot(self) -> "JSONDict":
         return {
             "enabled": self.enabled,
             "period_m": self._period_m,
         }
 
     @property
-    def enabled(self) -> 'bool':
+    def enabled(self) -> "bool":
         return self._period_m is not None
 
-    def turn_started(self) -> 'None':
+    def turn_started(self) -> "None":
         self._cancel_pending()
 
-    def arm_after_reply(self) -> 'None':
+    def arm_after_reply(self) -> "None":
         self._cancel_pending()
         period_m = self._period_m
         if period_m is None or self._notify_hook is None:
@@ -91,17 +97,19 @@ class ClockManager:
             lambda task: None if task.cancelled() else task.exception()
         )
 
-    def cancel(self) -> 'None':
+    def cancel(self) -> "None":
         self.set_period(None)
 
-    def _cancel_pending(self) -> 'None':
+    def _cancel_pending(self) -> "None":
         self._generation += 1
         task = self._timer_task
         self._timer_task = None
         if task is not None and not task.done():
             task.cancel()
 
-    async def _wait_and_notify(self, generation: 'int', period_m: 'float') -> 'None':
+    async def _wait_and_notify(self, generation: "int", period_m: "float") -> "None":
+        from ..agent import TurnInterrupted
+
         try:
             await asyncio.sleep(period_m * self._seconds_per_minute)
         except asyncio.CancelledError:
@@ -121,8 +129,19 @@ class ClockManager:
                     "current_time": _current_time(),
                 }
             )
-        except Exception:
-            started = False
+        except TurnInterrupted:
+            # The next successful reply will arm the clock after a steer.
+            return
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            asyncio.get_running_loop().call_exception_handler(
+                {
+                    "message": "Clock notification failed",
+                    "exception": exc,
+                }
+            )
+            return
 
         if (
             not started
@@ -147,9 +166,7 @@ class ClockTool(BaseTool):
                     {"type": "number"},
                     {"type": "null"},
                 ],
-                "description": (
-                    "Positive period in minutes, or null to cancel."
-                ),
+                "description": ("Positive period in minutes, or null to cancel."),
             },
         },
         "required": ["period_m"],
@@ -158,10 +175,26 @@ class ClockTool(BaseTool):
     output_schema = CLOCK_STATE_OUTPUT_SCHEMA
     supports_parallel = False
 
-    def __init__(self, manager: 'ClockManager') -> 'None':
+    def __init__(self, manager: "ClockManager") -> "None":
         self._manager = manager
 
-    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
+    def bind_agent(self, agent) -> "None":
+        self._manager.set_notify_hook(agent.maybe_invoke)
+
+    def handle_agent_event(self, event: "Event") -> "None":
+        if isinstance(event, (TurnStartedEvent, CompactStartedEvent)):
+            self._manager.turn_started()
+        elif isinstance(event, (TurnCompletedEvent, CompactCompletedEvent)):
+            self._manager.arm_after_reply()
+
+    def shutdown(self) -> "None":
+        self._manager.cancel()
+        self._manager.set_notify_hook(None)
+
+    def background_work_count(self, after_reply: "bool") -> "int":
+        return int(after_reply and self._manager.enabled)
+
+    async def run(self, context: "ToolContext", args: "JSONDict") -> "JSONValue":
         del context
         if not isinstance(args, dict) or "period_m" not in args:
             raise ValueError("clock requires period_m")

pycodex/tools/close_agent_tool.py

@@ -49,10 +49,10 @@ class CloseAgentTool(BaseTool):
     output_schema = CLOSE_AGENT_OUTPUT_SCHEMA
     supports_parallel = False
 
-    def __init__(self, subagent_manager: 'SubAgentManager') -> 'None':
+    def __init__(self, subagent_manager: "SubAgentManager") -> "None":
         self._subagent_manager = subagent_manager
 
-    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
+    async def run(self, context: "ToolContext", args: "JSONDict") -> "JSONValue":
         del context
         agent_id = str(args.get("id", "")).strip()
         if not agent_id: