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

pycodex/tools/code_mode_manager.py

@@ -13,6 +13,8 @@ Expected behavior:
 import asyncio
 import json
 import math
+import threading
+import typing
 import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -23,7 +25,6 @@ from ..compat import is_ascii, stream_writer_is_closing
 from ..protocol import JSONDict, JSONValue, ToolCall
 from ..utils.truncation import DEFAULT_MAX_OUTPUT_TOKENS
 from .base_tool import StructuredToolOutput, ToolContext, ToolRegistry
-import typing
 
 DEFAULT_WAIT_YIELD_TIME_MS = 10_000
 CHARS_PER_TOKEN = 4
@@ -33,39 +34,47 @@ WAIT_COMPLETION_GRACE_SECONDS = 0.02
 
 @dataclass
 class ExecCell:
-    cell_id: 'str'
-    process: 'asyncio.subprocess.Process'
-    started_at: 'float'
-    output_items: 'typing.List[JSONDict]' = field(default_factory=list)
-    delivered_count: 'int' = 0
-    reader_task: 'typing.Union[asyncio.Task, None]' = None
-    stderr_task: 'typing.Union[asyncio.Task, None]' = None
-    yield_event: 'asyncio.Event' = field(default_factory=asyncio.Event)
-    output_event: 'asyncio.Event' = field(default_factory=asyncio.Event)
-    done_event: 'asyncio.Event' = field(default_factory=asyncio.Event)
-    completed: 'bool' = False
-    terminated: 'bool' = False
-    error_text: 'typing.Union[str, None]' = None
-    stderr_chunks: 'typing.List[str]' = field(default_factory=list)
-
-
-@dataclass(frozen=True, )
+    cell_id: "str"
+    process: "asyncio.subprocess.Process"
+    started_at: "float"
+    output_items: "typing.List[JSONDict]" = field(default_factory=list)
+    delivered_count: "int" = 0
+    reader_task: "typing.Union[asyncio.Task, None]" = None
+    stderr_task: "typing.Union[asyncio.Task, None]" = None
+    yield_event: "asyncio.Event" = field(default_factory=asyncio.Event)
+    output_event: "asyncio.Event" = field(default_factory=asyncio.Event)
+    done_event: "asyncio.Event" = field(default_factory=asyncio.Event)
+    completed: "bool" = False
+    terminated: "bool" = False
+    error_text: "typing.Union[str, None]" = None
+    stderr_chunks: "typing.List[str]" = field(default_factory=list)
+
+
+@dataclass(
+    frozen=True,
+)
 class ParsedExecSource:
-    code: 'str'
-    yield_time_ms: 'typing.Union[int, None]'
-    max_output_tokens: 'typing.Union[int, None]'
+    code: "str"
+    yield_time_ms: "typing.Union[int, None]"
+    max_output_tokens: "typing.Union[int, None]"
 
 
 class CodeModeManager:
-    def __init__(self, registry: 'ToolRegistry', cwd: 'typing.Union[typing.Union[str, Path], None]' = None) -> 'None':
+    def __init__(
+        self,
+        registry: "ToolRegistry",
+        cwd: "typing.Union[typing.Union[str, Path], None]" = None,
+    ) -> "None":
         self._registry = registry
         self._default_cwd = Path(cwd or Path.cwd()).resolve()
         self._runtime_script = Path(__file__).with_name("exec_runtime.js")
-        self._stored_values: 'typing.Dict[str, JSONValue]' = {}
-        self._cells: 'typing.Dict[str, ExecCell]' = {}
-        self._lock = asyncio.Lock()
+        self._stored_values: "typing.Dict[str, JSONValue]" = {}
+        self._cells: "typing.Dict[str, ExecCell]" = {}
+        self._lock = threading.Lock()
 
-    async def exec(self, source: 'str', context: 'ToolContext') -> 'typing.Union[StructuredToolOutput, str]':
+    async def exec(
+        self, source: "str", context: "ToolContext"
+    ) -> "typing.Union[StructuredToolOutput, str]":
         try:
             parsed = self._parse_exec_source(source)
         except ValueError as exc:
@@ -77,11 +86,11 @@ class CodeModeManager:
 
     async def wait(
         self,
-        cell_id: 'str',
-        yield_time_ms: 'int',
-        max_tokens: 'typing.Union[int, None]',
-        terminate: 'bool',
-    ) -> 'typing.Union[StructuredToolOutput, str]':
+        cell_id: "str",
+        yield_time_ms: "int",
+        max_tokens: "typing.Union[int, None]",
+        terminate: "bool",
+    ) -> "typing.Union[StructuredToolOutput, str]":
         cell = self._cells.get(cell_id)
         if cell is None:
             return f"Error: unknown exec cell `{cell_id}`."
@@ -98,8 +107,8 @@ class CodeModeManager:
         await self._wait_for_wait(cell, yield_time_ms)
         return await self._snapshot_cell(cell, max_tokens)
 
-    def enabled_tools(self) -> 'typing.List[typing.Dict[str, str]]':
-        enabled: 'typing.List[typing.Dict[str, str]]' = []
+    def enabled_tools(self) -> "typing.List[typing.Dict[str, str]]":
+        enabled: "typing.List[typing.Dict[str, str]]" = []
         for tool in self._registry.tools():
             if tool.name in {"exec", "wait"}:
                 continue
@@ -116,7 +125,7 @@ class CodeModeManager:
         enabled.sort(key=lambda item: item["tool_name"])
         return enabled
 
-    async def _start_cell(self, code: 'str', context: 'ToolContext') -> 'ExecCell':
+    async def _start_cell(self, code: "str", context: "ToolContext") -> "ExecCell":
         cell_id = uuid.uuid4().hex[:10]
         process = await asyncio.create_subprocess_exec(
             "node",
@@ -147,7 +156,7 @@ class CodeModeManager:
         logger.debug("exec start cell_id={} cwd={}", cell_id, self._default_cwd)
         return cell
 
-    async def _read_stdout(self, cell: 'ExecCell', context: 'ToolContext') -> 'None':
+    async def _read_stdout(self, cell: "ExecCell", context: "ToolContext") -> "None":
         stream = cell.process.stdout
         if stream is None:
             cell.error_text = "missing stdout pipe"
@@ -181,7 +190,7 @@ class CodeModeManager:
                 cell.output_event.set()
                 continue
             if msg_type == "output_image":
-                image_item: 'JSONDict' = {
+                image_item: "JSONDict" = {
                     "type": "input_image",
                     "image_url": str(message.get("image_url", "")),
                 }
@@ -202,7 +211,7 @@ class CodeModeManager:
                 cell.error_text = self._coerce_optional_text(message.get("error_text"))
                 stored_values = message.get("stored_values")
                 if isinstance(stored_values, dict):
-                    async with self._lock:
+                    with self._lock:
                         self._stored_values = stored_values
                 cell.done_event.set()
                 cell.output_event.set()
@@ -220,7 +229,7 @@ class CodeModeManager:
             cell.done_event.set()
             cell.output_event.set()
 
-    async def _read_stderr(self, cell: 'ExecCell') -> 'None':
+    async def _read_stderr(self, cell: "ExecCell") -> "None":
         stream = cell.process.stderr
         if stream is None:
             return
@@ -232,10 +241,10 @@ class CodeModeManager:
 
     async def _handle_nested_tool_call(
         self,
-        cell: 'ExecCell',
-        context: 'ToolContext',
-        message: 'JSONDict',
-    ) -> 'None':
+        cell: "ExecCell",
+        context: "ToolContext",
+        message: "JSONDict",
+    ) -> "None":
         tool_name = str(message.get("tool_name", ""))
         request_id = str(message.get("id", ""))
         tool = self._registry.get_tool(tool_name)
@@ -287,14 +296,16 @@ class CodeModeManager:
             }
         await self._send_message(cell, payload)
 
-    async def _send_message(self, cell: 'ExecCell', payload: 'JSONDict') -> 'None':
+    async def _send_message(self, cell: "ExecCell", payload: "JSONDict") -> "None":
         stdin = cell.process.stdin
         if stdin is None or stream_writer_is_closing(stdin):
             return
         stdin.write((json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8"))
         await stdin.drain()
 
-    async def _wait_for_exec(self, cell: 'ExecCell', yield_time_ms: 'typing.Union[int, None]') -> 'None':
+    async def _wait_for_exec(
+        self, cell: "ExecCell", yield_time_ms: "typing.Union[int, None]"
+    ) -> "None":
         done_task = asyncio.create_task(cell.done_event.wait())
         yield_task = asyncio.create_task(cell.yield_event.wait())
         tasks = {done_task, yield_task}
@@ -314,7 +325,7 @@ class CodeModeManager:
                     task.cancel()
         cell.yield_event.clear()
 
-    async def _wait_for_wait(self, cell: 'ExecCell', yield_time_ms: 'int') -> 'None':
+    async def _wait_for_wait(self, cell: "ExecCell", yield_time_ms: "int") -> "None":
         loop = asyncio.get_running_loop()
         deadline = loop.time() + max(yield_time_ms, 1) / 1000.0
         initial_count = cell.delivered_count
@@ -357,9 +368,9 @@ class CodeModeManager:
 
     async def _wait_for_completion_grace(
         self,
-        cell: 'ExecCell',
-        timeout_seconds: 'float',
-    ) -> 'None':
+        cell: "ExecCell",
+        timeout_seconds: "float",
+    ) -> "None":
         if timeout_seconds <= 0:
             return
         done_task = asyncio.create_task(cell.done_event.wait())
@@ -378,9 +389,9 @@ class CodeModeManager:
 
     async def _snapshot_cell(
         self,
-        cell: 'ExecCell',
-        max_tokens: 'typing.Union[int, None]',
-    ) -> 'StructuredToolOutput':
+        cell: "ExecCell",
+        max_tokens: "typing.Union[int, None]",
+    ) -> "StructuredToolOutput":
         if cell.process.returncode is not None and cell.reader_task is not None:
             await cell.reader_task
 
@@ -417,17 +428,19 @@ class CodeModeManager:
         if cell.done_event.is_set():
             self._cells.pop(cell.cell_id, None)
 
-        return StructuredToolOutput(output=output_text, content_items=tuple(content_items))
+        return StructuredToolOutput(
+            output=output_text, content_items=tuple(content_items)
+        )
 
     def _truncate_content_items(
         self,
-        items: 'typing.List[JSONDict]',
-        max_tokens: 'typing.Union[int, None]',
-    ) -> 'typing.List[JSONDict]':
+        items: "typing.List[JSONDict]",
+        max_tokens: "typing.Union[int, None]",
+    ) -> "typing.List[JSONDict]":
         token_budget = DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens
         max_chars = max(1, token_budget) * CHARS_PER_TOKEN
         total_chars = 0
-        truncated: 'typing.List[JSONDict]' = []
+        truncated: "typing.List[JSONDict]" = []
         for item in items:
             if item.get("type") != "input_text":
                 truncated.append(item)
@@ -449,7 +462,7 @@ class CodeModeManager:
             total_chars += len(text)
         return truncated
 
-    def _status_text(self, cell: 'ExecCell') -> 'str':
+    def _status_text(self, cell: "ExecCell") -> "str":
         if cell.terminated:
             return "Script terminated"
         if not cell.done_event.is_set():
@@ -458,11 +471,9 @@ class CodeModeManager:
             return "Script failed"
         return "Script completed"
 
-    def _parse_exec_source(self, input_text: 'str') -> 'ParsedExecSource':
+    def _parse_exec_source(self, input_text: "str") -> "ParsedExecSource":
         if not input_text.strip():
-            raise ValueError(
-                "exec expects raw JavaScript source text (non-empty)."
-            )
+            raise ValueError("exec expects raw JavaScript source text (non-empty).")
         code = input_text
         yield_time_ms = None
         max_output_tokens = None
@@ -499,13 +510,15 @@ class CodeModeManager:
             max_output_tokens=max_output_tokens,
         )
 
-    def _normalize_identifier(self, tool_name: 'str') -> 'str':
+    def _normalize_identifier(self, tool_name: "str") -> "str":
         identifier = []
         for index, char in enumerate(tool_name):
             is_valid = (
                 char == "_"
                 or char == "$"
-                or (is_ascii(char) and char.isalnum() and (index != 0 or char.isalpha()))
+                or (
+                    is_ascii(char) and char.isalnum() and (index != 0 or char.isalpha())
+                )
             )
             if is_valid:
                 identifier.append(char)
@@ -513,7 +526,7 @@ class CodeModeManager:
                 identifier.append("_")
         return "".join(identifier) or "_"
 
-    def _coerce_optional_text(self, value: 'JSONValue') -> 'typing.Union[str, None]':
+    def _coerce_optional_text(self, value: "JSONValue") -> "typing.Union[str, None]":
         if value in (None, ""):
             return None
         return str(value)

pycodex/tools/exec_command_tool.py

@@ -10,6 +10,8 @@ Expected behavior:
   `write_stdin`.
 """
 
+import typing
+
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
 from .unified_exec_manager import (
@@ -19,7 +21,9 @@ from .unified_exec_manager import (
     UNIFIED_EXEC_OUTPUT_SCHEMA,
     UnifiedExecManager,
 )
-import typing
+
+if typing.TYPE_CHECKING:
+    from ..agent import Agent
 
 MIN_EXEC_YIELD_TIME_MS = 250
 MAX_EXEC_YIELD_TIME_MS = 30_000
@@ -70,10 +74,19 @@ class ExecCommandTool(BaseTool):
     output_schema = UNIFIED_EXEC_OUTPUT_SCHEMA
     supports_parallel = False
 
-    def __init__(self, manager: 'UnifiedExecManager') -> 'None':
+    def __init__(self, manager: "UnifiedExecManager") -> "None":
         self._manager = manager
 
-    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
+    def bind_agent(self, agent: "Agent") -> "None":
+        self._manager.set_notify_hook(agent.maybe_invoke)
+
+    def shutdown(self) -> "None":
+        self._manager.set_notify_hook(None)
+
+    def background_work_count(self, after_reply: "bool") -> "int":
+        return self._manager.running_session_count()
+
+    async def run(self, context: "ToolContext", args: "JSONDict") -> "JSONValue":
         del context
         cmd = str(args.get("cmd", "")).strip()
         if not cmd:
@@ -95,13 +108,15 @@ class ExecCommandTool(BaseTool):
             max_output_tokens=self._optional_int(args, "max_output_tokens"),
         )
 
-    def _optional_string(self, args: 'JSONDict', key: 'str') -> 'typing.Union[str, None]':
+    def _optional_string(
+        self, args: "JSONDict", key: "str"
+    ) -> "typing.Union[str, None]":
         value = args.get(key)
         if value in (None, ""):
             return None
         return str(value)
 
-    def _optional_int(self, args: 'JSONDict', key: 'str') -> 'typing.Union[int, None]':
+    def _optional_int(self, args: "JSONDict", key: "str") -> "typing.Union[int, None]":
         value = args.get(key)
         if value in (None, ""):
             return None
@@ -109,11 +124,11 @@ class ExecCommandTool(BaseTool):
 
     def _bounded_int(
         self,
-        args: 'JSONDict',
-        key: 'str',
-        default: 'int',
-        minimum: 'int',
-        maximum: 'int',
-    ) -> 'int':
+        args: "JSONDict",
+        key: "str",
+        default: "int",
+        minimum: "int",
+        maximum: "int",
+    ) -> "int":
         value = int(args.get(key, default))
         return min(max(value, minimum), maximum)

pycodex/tools/exec_tool.py

@@ -40,8 +40,8 @@ class ExecTool(BaseTool):
         "- Accepts raw JavaScript source text, not JSON, quoted strings, or "
         "markdown code fences.\n"
         "- You may optionally start the tool input with a first-line pragma "
-        "like `// @exec: {\"yield_time_ms\": 10000, "
-        "\"max_output_tokens\": 1000}`.\n"
+        'like `// @exec: {"yield_time_ms": 10000, '
+        '"max_output_tokens": 1000}`.\n'
         "- `yield_time_ms` asks `exec` to yield early if the script is still "
         "running. Defaults to 10000 ms.\n"
         "- `max_output_tokens` sets the token budget for direct `exec` results. "
@@ -55,8 +55,8 @@ class ExecTool(BaseTool):
     }
     supports_parallel = False
 
-    def __init__(self, manager: 'CodeModeManager') -> 'None':
+    def __init__(self, manager: "CodeModeManager") -> "None":
         self._manager = manager
 
-    async def run(self, context: 'ToolContext', args: 'JSONValue') -> 'JSONValue':
+    async def run(self, context: "ToolContext", args: "JSONValue") -> "JSONValue":
         return await self._manager.exec(str(args), context)

pycodex/tools/grep_files_tool.py

@@ -12,11 +12,11 @@ Expected behavior:
 import asyncio
 import fnmatch
 import re
+import typing
 from pathlib import Path
 
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
-import typing
 
 DEFAULT_LIMIT = 100
 MAX_LIMIT = 2000
@@ -53,10 +53,12 @@ class GrepFilesTool(BaseTool):
         "additionalProperties": 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._working_directory = Path(cwd or Path.cwd()).resolve()
 
-    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
+    async def run(self, context: "ToolContext", args: "JSONDict") -> "JSONValue":
         del context
         pattern = str(args.get("pattern", "")).strip()
         include = str(args.get("include", "")).strip() or None
@@ -123,13 +125,13 @@ class GrepFilesTool(BaseTool):
 
     def _search_with_python(
         self,
-        pattern: 'str',
-        include: 'typing.Union[str, None]',
-        search_path: 'Path',
-        limit: 'int',
-    ) -> 'typing.List[str]':
+        pattern: "str",
+        include: "typing.Union[str, None]",
+        search_path: "Path",
+        limit: "int",
+    ) -> "typing.List[str]":
         regex = re.compile(pattern)
-        candidates: 'typing.List[Path]' = []
+        candidates: "typing.List[Path]" = []
 
         if search_path.is_file():
             candidates = [search_path]
@@ -153,7 +155,7 @@ class GrepFilesTool(BaseTool):
         matches.sort(key=lambda path: path.stat().st_mtime, reverse=True)
         return [str(path) for path in matches[:limit]]
 
-    def _resolve_path(self, path_arg) -> 'Path':
+    def _resolve_path(self, path_arg) -> "Path":
         if path_arg in (None, ""):
             return self._working_directory
         path = Path(str(path_arg))

pycodex/tools/ipython_tool.py

@@ -1,6 +1,7 @@
 """Optional tool that executes code in the current IPython kernel."""
 
-from ..utils.toolcall_visualize import colorize_tool_message, tool_summary
+from ..events import ToolCompletedEvent
+from ..utils.event_helpers import colorize_tool_message
 from .base_tool import BaseTool, ToolContext
 
 
@@ -28,7 +29,7 @@ class IPythonTool(BaseTool):
         self._shell = shell
         self.name = name
 
-    async def run(self, context: 'ToolContext', args):
+    async def run(self, context: "ToolContext", args):
         del context
         if not isinstance(args, dict):
             return {"error": "arguments must be an object"}
@@ -55,9 +56,7 @@ class IPythonTool(BaseTool):
             data = getattr(item, "data", None)
             if isinstance(data, dict):
                 display_outputs.append(
-                    data.get("text/plain")
-                    or data.get("text/html")
-                    or repr(data)
+                    data.get("text/plain") or data.get("text/html") or repr(data)
                 )
             else:
                 display_outputs.append(repr(item))
@@ -106,9 +105,7 @@ def _install_agent_shortcut(shell):
         if not prompt.strip():
             return [f"{indent}print('Usage: @{agent_name} <prompt>')\n"]
 
-        return [
-            f"{indent}print({agent_name}.ask({prompt!r}).output_text)\n"
-        ]
+        return [f"{indent}print({agent_name}.ask({prompt!r}).output_text)\n"]
 
     shell.input_transformers_cleanup.append(transform)
     shell.user_ns["_pycodex_agent_shortcut_transform"] = transform
@@ -117,13 +114,13 @@ def _install_agent_shortcut(shell):
 
 def attach_ipython_event_printer(agent, color=True):
     def handle_event(event):
-        if event.kind == "tool_completed":
-            tool_name = str(event.payload.get("tool_name", "")).strip()
-            message = tool_summary(event.payload)
+        if isinstance(event, ToolCompletedEvent):
+            tool_name = event.call.name
+            message = event.visualize()
             for line in message.splitlines() or [""]:
                 print(colorize_tool_message(line, color, tool_name), flush=True)
 
-    agent.set_event_handler(handle_event)
+    agent.event_handler = handle_event
     return handle_event
 
 
@@ -137,7 +134,7 @@ def attach_ipython_tool(
         raise RuntimeError("not running inside IPython")
 
     tool = IPythonTool(shell, name=name)
-    agent._tool_registry.register(tool)
+    agent.tool_registry.register(tool)
     if shortcut:
         _install_agent_shortcut(shell)
     if print_tool_events:

pycodex/tools/list_dir_tool.py

@@ -9,12 +9,12 @@ Expected behavior:
   commands like `find` or `ls -R`.
 """
 
+import typing
 from collections import deque
 from pathlib import Path
 
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
-import typing
 
 MAX_ENTRY_LENGTH = 500
 INDENTATION_SPACES = 2
@@ -50,7 +50,7 @@ class ListDirTool(BaseTool):
         "additionalProperties": False,
     }
 
-    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
+    async def run(self, context: "ToolContext", args: "JSONDict") -> "JSONValue":
         del context
         dir_path = Path(str(args.get("dir_path", "")))
         offset = int(args.get("offset", 1))
@@ -86,15 +86,19 @@ class ListDirTool(BaseTool):
             lines.append(f"More than {len(selected)} entries found")
         return "\n".join(lines)
 
-    def _collect_entries(self, root: 'Path', depth: 'int') -> 'typing.List[typing.Dict[str, object]]':
-        entries: 'typing.List[typing.Dict[str, object]]' = []
+    def _collect_entries(
+        self, root: "Path", depth: "int"
+    ) -> "typing.List[typing.Dict[str, object]]":
+        entries: "typing.List[typing.Dict[str, object]]" = []
         queue = deque([(root, Path(), depth)])
 
         while queue:
             current_dir, prefix, remaining_depth = queue.popleft()
             dir_entries = []
             for child in current_dir.iterdir():
-                relative_path = prefix / child.name if prefix.parts else Path(child.name)
+                relative_path = (
+                    prefix / child.name if prefix.parts else Path(child.name)
+                )
                 kind = self._entry_kind(child)
                 dir_entries.append(
                     (
@@ -118,7 +122,7 @@ class ListDirTool(BaseTool):
         entries.sort(key=lambda entry: entry["name"])
         return entries
 
-    def _entry_kind(self, path: 'Path') -> 'str':
+    def _entry_kind(self, path: "Path") -> "str":
         if path.is_symlink():
             return "symlink"
         if path.is_dir():
@@ -127,14 +131,14 @@ class ListDirTool(BaseTool):
             return "file"
         return "other"
 
-    def _format_entry_name(self, path: 'Path') -> 'str':
+    def _format_entry_name(self, path: "Path") -> "str":
         text = path.as_posix()
         return text[:MAX_ENTRY_LENGTH]
 
-    def _format_component(self, name: 'str') -> 'str':
+    def _format_component(self, name: "str") -> "str":
         return name[:MAX_ENTRY_LENGTH]
 
-    def _format_entry_line(self, entry: 'typing.Dict[str, object]') -> 'str':
+    def _format_entry_line(self, entry: "typing.Dict[str, object]") -> "str":
         indent = " " * (int(entry["depth"]) * INDENTATION_SPACES)
         name = str(entry["display_name"])
         kind = str(entry["kind"])