python-codex 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

pycodex/tools/unified_exec_manager.py

@@ -11,8 +11,6 @@ Expected behavior:
 - Return summaries in the same textual shape Codex tools expect.
 """
 
-from __future__ import annotations
-
 import asyncio
 import os
 import shlex
@@ -22,6 +20,9 @@ from pathlib import Path
 
 from loguru import logger
 
+from ..compat import shlex_join, stream_writer_is_closing
+import typing
+
 DEFAULT_EXEC_YIELD_TIME_MS = 10_000
 DEFAULT_WRITE_STDIN_YIELD_TIME_MS = 250
 DEFAULT_MAX_OUTPUT_TOKENS = 10_000
@@ -63,33 +64,33 @@ UNIFIED_EXEC_OUTPUT_SCHEMA = {
 }
 
 
-def _approx_token_count(text: str) -> int:
+def _approx_token_count(text: 'str') -> 'int':
     if not text:
         return 0
     byte_length = len(text.encode("utf-8"))
     return max(1, (byte_length + APPROX_BYTES_PER_TOKEN - 1) // APPROX_BYTES_PER_TOKEN)
 
 
-def _approx_bytes_for_tokens(token_count: int) -> int:
+def _approx_bytes_for_tokens(token_count: 'int') -> 'int':
     return max(token_count, 0) * APPROX_BYTES_PER_TOKEN
 
 
-def _approx_tokens_from_byte_count(byte_count: int) -> int:
+def _approx_tokens_from_byte_count(byte_count: 'int') -> 'int':
     if byte_count <= 0:
         return 0
     return (byte_count + APPROX_BYTES_PER_TOKEN - 1) // APPROX_BYTES_PER_TOKEN
 
 
-def _split_budget(byte_budget: int) -> tuple[int, int]:
+def _split_budget(byte_budget: 'int') -> 'typing.Tuple[int, int]':
     left_budget = byte_budget // 2
     return left_budget, byte_budget - left_budget
 
 
 def _split_string(
-    text: str,
-    beginning_bytes: int,
-    end_bytes: int,
-) -> tuple[str, str]:
+    text: 'str',
+    beginning_bytes: 'int',
+    end_bytes: 'int',
+) -> 'typing.Tuple[str, str]':
     if not text:
         return "", ""
 
@@ -122,7 +123,7 @@ def _split_string(
     return text[:prefix_end], text[suffix_start:]
 
 
-def _truncate_text(text: str, max_tokens: int) -> str:
+def _truncate_text(text: 'str', max_tokens: 'int') -> 'str':
     if not text:
         return ""
 
@@ -141,7 +142,7 @@ def _truncate_text(text: str, max_tokens: int) -> str:
     return f"{prefix}{marker}{suffix}"
 
 
-def _formatted_truncate_text(text: str, max_tokens: int) -> str:
+def _formatted_truncate_text(text: 'str', max_tokens: 'int') -> 'str':
     byte_budget = _approx_bytes_for_tokens(max_tokens)
     if len(text.encode("utf-8")) <= byte_budget:
         return text
@@ -150,13 +151,13 @@ def _formatted_truncate_text(text: str, max_tokens: int) -> str:
     return f"Total output lines: {total_lines}\n\n{_truncate_text(text, max_tokens)}"
 
 
-@dataclass(slots=True)
+@dataclass
 class _HeadTailBuffer:
-    max_bytes: int = UNIFIED_EXEC_OUTPUT_MAX_BYTES
-    head: bytearray = field(default_factory=bytearray)
-    tail: bytearray = field(default_factory=bytearray)
+    max_bytes: 'int' = UNIFIED_EXEC_OUTPUT_MAX_BYTES
+    head: 'bytearray' = field(default_factory=bytearray)
+    tail: 'bytearray' = field(default_factory=bytearray)
 
-    def push_chunk(self, chunk: bytes) -> None:
+    def push_chunk(self, chunk: 'bytes') -> 'None':
         if not chunk or self.max_bytes <= 0:
             return
 
@@ -178,41 +179,45 @@ class _HeadTailBuffer:
             excess = len(self.tail) - tail_budget
             del self.tail[:excess]
 
-    def drain_bytes(self) -> bytes:
+    def drain_bytes(self) -> 'bytes':
         combined = bytes(self.head) + bytes(self.tail)
         self.head.clear()
         self.tail.clear()
         return combined
 
+    def has_data(self) -> 'bool':
+        return bool(self.head or self.tail)
+
 
-@dataclass(slots=True)
+@dataclass
 class UnifiedExecSession:
-    session_id: int
-    process: asyncio.subprocess.Process
-    start_time: float
-    command_display: str
-    tty: bool
-    unread_output: _HeadTailBuffer = field(default_factory=_HeadTailBuffer)
-    reader_task: asyncio.Task | None = None
+    session_id: 'int'
+    process: 'asyncio.subprocess.Process'
+    start_time: 'float'
+    command_display: 'str'
+    tty: 'bool'
+    unread_output: '_HeadTailBuffer' = field(default_factory=_HeadTailBuffer)
+    reader_task: 'typing.Union[asyncio.Task, None]' = None
+    output_event: 'asyncio.Event' = field(default_factory=asyncio.Event)
 
 
 class UnifiedExecManager:
-    def __init__(self, cwd: str | Path | None = None) -> None:
+    def __init__(self, cwd: 'typing.Union[typing.Union[str, Path], None]' = None) -> 'None':
         self._default_cwd = Path(cwd or Path.cwd()).resolve()
         self._next_session_id = DEFAULT_SESSION_ID_START
-        self._sessions: dict[int, UnifiedExecSession] = {}
+        self._sessions: 'typing.Dict[int, UnifiedExecSession]' = {}
         self._lock = asyncio.Lock()
 
     async def exec_command(
         self,
-        cmd: str,
-        workdir: str | None = None,
-        shell: str | None = None,
-        login: bool = DEFAULT_LOGIN,
-        tty: bool = DEFAULT_TTY,
-        yield_time_ms: int = DEFAULT_EXEC_YIELD_TIME_MS,
-        max_output_tokens: int | None = None,
-    ) -> str:
+        cmd: 'str',
+        workdir: 'typing.Union[str, None]' = None,
+        shell: 'typing.Union[str, None]' = None,
+        login: 'bool' = DEFAULT_LOGIN,
+        tty: 'bool' = DEFAULT_TTY,
+        yield_time_ms: 'int' = DEFAULT_EXEC_YIELD_TIME_MS,
+        max_output_tokens: 'typing.Union[int, None]' = None,
+    ) -> 'str':
         session_id = await self._allocate_session_id()
         command = self._build_shell_command(cmd, shell, login)
         cwd = self._resolve_workdir(workdir)
@@ -234,7 +239,7 @@ class UnifiedExecManager:
             session_id=session_id,
             process=process,
             start_time=asyncio.get_running_loop().time(),
-            command_display=shlex.join(command),
+            command_display=shlex_join(command),
             tty=tty,
         )
         session.reader_task = asyncio.create_task(self._pump_output(session))
@@ -250,11 +255,11 @@ class UnifiedExecManager:
 
     async def write_stdin(
         self,
-        session_id: int,
-        chars: str = "",
-        yield_time_ms: int = DEFAULT_WRITE_STDIN_YIELD_TIME_MS,
-        max_output_tokens: int | None = None,
-    ) -> str:
+        session_id: 'int',
+        chars: 'str' = "",
+        yield_time_ms: 'int' = DEFAULT_WRITE_STDIN_YIELD_TIME_MS,
+        max_output_tokens: 'typing.Union[int, None]' = None,
+    ) -> 'str':
         session = await self._get_session(session_id)
         if session is None:
             return f"Error: session_id {session_id} is not running."
@@ -274,31 +279,42 @@ class UnifiedExecManager:
             max_output_tokens,
         )
 
-    async def _allocate_session_id(self) -> int:
+    async def _allocate_session_id(self) -> 'int':
         async with self._lock:
             session_id = self._next_session_id
             self._next_session_id += 1
             return session_id
 
-    async def _get_session(self, session_id: int) -> UnifiedExecSession | None:
+    async def _get_session(self, session_id: 'int') -> 'typing.Union[UnifiedExecSession, None]':
         async with self._lock:
             return self._sessions.get(session_id)
 
     async def _wait_and_snapshot(
         self,
-        session_id: int,
-        yield_time_ms: int,
-        max_output_tokens: int | None,
-    ) -> str:
+        session_id: 'int',
+        yield_time_ms: 'int',
+        max_output_tokens: 'typing.Union[int, None]',
+    ) -> 'str':
         session = await self._get_session(session_id)
         if session is None:
             return f"Error: session_id {session_id} is not running."
 
-        start_wait = asyncio.get_running_loop().time()
+        loop = asyncio.get_running_loop()
+        start_wait = loop.time()
         try:
             await asyncio.wait_for(session.process.wait(), timeout=yield_time_ms / 1000.0)
         except asyncio.TimeoutError:
-            pass
+            remaining_seconds = (yield_time_ms / 1000.0) - (loop.time() - start_wait)
+            if (
+                session.process.returncode is None
+                and not session.unread_output.has_data()
+                and remaining_seconds > 0
+            ):
+                session.output_event.clear()
+                try:
+                    await asyncio.wait_for(session.output_event.wait(), timeout=remaining_seconds)
+                except asyncio.TimeoutError:
+                    pass
 
         if session.reader_task is not None and session.process.returncode is not None:
             await session.reader_task
@@ -328,15 +344,15 @@ class UnifiedExecManager:
 
         return "\n".join(lines)
 
-    async def _close_session(self, session_id: int) -> None:
+    async def _close_session(self, session_id: 'int') -> 'None':
         async with self._lock:
             session = self._sessions.pop(session_id, None)
         if session is None:
             return
-        if session.process.stdin is not None and not session.process.stdin.is_closing():
+        if session.process.stdin is not None and not stream_writer_is_closing(session.process.stdin):
             session.process.stdin.close()
 
-    async def _pump_output(self, session: UnifiedExecSession) -> None:
+    async def _pump_output(self, session: 'UnifiedExecSession') -> 'None':
         stream = session.process.stdout
         if stream is None:
             return
@@ -345,8 +361,10 @@ class UnifiedExecManager:
             if not chunk:
                 break
             session.unread_output.push_chunk(chunk)
+            session.output_event.set()
+        session.output_event.set()
 
-    def _resolve_workdir(self, workdir: str | None) -> Path:
+    def _resolve_workdir(self, workdir: 'typing.Union[str, None]') -> 'Path':
         if not workdir:
             return self._default_cwd
         path = Path(workdir)
@@ -356,10 +374,10 @@ class UnifiedExecManager:
 
     def _build_shell_command(
         self,
-        cmd: str,
-        shell: str | None,
-        login: bool,
-    ) -> list[str]:
+        cmd: 'str',
+        shell: 'typing.Union[str, None]',
+        login: 'bool',
+    ) -> 'typing.List[str]':
         shell_path = shell or os.environ.get("SHELL") or "/bin/bash"
         shell_name = Path(shell_path).name.lower()
         if shell_name in {"cmd", "cmd.exe"}:
@@ -368,13 +386,13 @@ class UnifiedExecManager:
             return [shell_path, "-NoProfile", "-Command", cmd]
         return [shell_path, "-lc" if login else "-c", cmd]
 
-    def _estimate_token_count(self, output: str) -> int | None:
+    def _estimate_token_count(self, output: 'str') -> 'typing.Union[int, None]':
         return _approx_token_count(output)
 
-    def _truncate_output(self, output: str, max_output_tokens: int | None) -> str:
+    def _truncate_output(self, output: 'str', max_output_tokens: 'typing.Union[int, None]') -> 'str':
         token_budget = DEFAULT_MAX_OUTPUT_TOKENS if max_output_tokens is None else max_output_tokens
         return _formatted_truncate_text(output, max(token_budget, 0))
 
-    def _tty_echo(self, chars: str) -> bytes:
+    def _tty_echo(self, chars: 'str') -> 'bytes':
         normalized = chars.replace("\n", "\r\n")
         return normalized.encode("utf-8")

pycodex/tools/update_plan_tool.py

@@ -10,11 +10,10 @@ Expected behavior:
   confirmation text Codex uses.
 """
 
-from __future__ import annotations
-
 from ..protocol import JSONDict, JSONValue
 from ..runtime_services import PlanItem, PlanStore
 from .base_tool import BaseTool, ToolContext
+import typing
 
 VALID_PLAN_STATUSES = {"pending", "in_progress", "completed"}
 
@@ -52,16 +51,16 @@ class UpdatePlanTool(BaseTool):
     }
     supports_parallel = False
 
-    def __init__(self, plan_store: PlanStore) -> None:
+    def __init__(self, plan_store: 'PlanStore') -> 'None':
         self._plan_store = plan_store
 
-    async def run(self, context: ToolContext, args: JSONDict) -> JSONValue:
+    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
         del context
         raw_plan = args.get("plan")
         if not isinstance(raw_plan, list):
             return "Error: `plan` must be a list."
 
-        plan_items: list[PlanItem] = []
+        plan_items: 'typing.List[PlanItem]' = []
         for item in raw_plan:
             if not isinstance(item, dict):
                 return "Error: each `plan` item must be an object."

pycodex/tools/view_image_tool.py

@@ -12,14 +12,13 @@ Expected behavior:
   item that Codex uses when feeding image tool output back to the model.
 """
 
-from __future__ import annotations
-
 import base64
 import mimetypes
 from pathlib import Path
 
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, StructuredToolOutput, ToolContext
+import typing
 
 VIEW_IMAGE_OUTPUT_SCHEMA = {
     "type": "object",
@@ -62,10 +61,10 @@ class ViewImageTool(BaseTool):
     }
     output_schema = VIEW_IMAGE_OUTPUT_SCHEMA
 
-    def __init__(self, cwd: 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: JSONDict) -> JSONValue:
+    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
         del context
         path_value = str(args.get("path", "")).strip()
         if not path_value:
@@ -102,7 +101,7 @@ class ViewImageTool(BaseTool):
             "image_url": image_url,
             "detail": detail,
         }
-        image_item: JSONDict = {
+        image_item: 'JSONDict' = {
             "type": "input_image",
             "image_url": image_url,
         }

pycodex/tools/wait_agent_tool.py

@@ -9,8 +9,6 @@ Expected behavior:
   wait times out.
 """
 
-from __future__ import annotations
-
 from ..protocol import JSONDict, JSONValue
 from ..runtime_services import SubAgentManager
 from .agent_tool_schemas import AGENT_STATUS_SCHEMA
@@ -60,10 +58,10 @@ class WaitAgentTool(BaseTool):
     output_schema = WAIT_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
         ids = args.get("ids")
         if not isinstance(ids, list) or not ids:

pycodex/tools/wait_tool.py

@@ -9,11 +9,10 @@ Expected behavior:
 - Return only the new output since the previous `exec` / `wait` snapshot.
 """
 
-from __future__ import annotations
-
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
 from .code_mode_manager import DEFAULT_WAIT_YIELD_TIME_MS, CodeModeManager
+import typing
 
 
 class WaitTool(BaseTool):
@@ -46,10 +45,10 @@ class WaitTool(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: JSONDict) -> JSONValue:
+    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
         del context
         cell_id = str(args.get("cell_id", "")).strip()
         if not cell_id:
@@ -61,7 +60,7 @@ class WaitTool(BaseTool):
             terminate=bool(args.get("terminate", False)),
         )
 
-    def _optional_int(self, args: JSONDict, key: str) -> int | None:
+    def _optional_int(self, args: 'JSONDict', key: 'str') -> 'typing.Union[int, None]':
         value = args.get(key)
         if value in (None, ""):
             return None

pycodex/tools/web_search_tool.py

@@ -10,8 +10,6 @@ Expected behavior:
 - Never expect a local tool-call round-trip result from the model.
 """
 
-from __future__ import annotations
-
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
 
@@ -25,6 +23,6 @@ class WebSearchTool(BaseTool):
     }
     supports_parallel = False
 
-    async def run(self, context: ToolContext, args: JSONValue) -> JSONValue:
+    async def run(self, context: 'ToolContext', args: 'JSONValue') -> 'JSONValue':
         del context, args
         return "Error: web_search is provider-native and should not be executed locally."

pycodex/tools/write_stdin_tool.py

@@ -10,8 +10,6 @@ Expected behavior:
 - Reuse the same `session_id` until the process exits.
 """
 
-from __future__ import annotations
-
 from ..protocol import JSONDict, JSONValue
 from .base_tool import BaseTool, ToolContext
 from .unified_exec_manager import (
@@ -19,6 +17,7 @@ from .unified_exec_manager import (
     UNIFIED_EXEC_OUTPUT_SCHEMA,
     UnifiedExecManager,
 )
+import typing
 
 
 class WriteStdinTool(BaseTool):
@@ -50,10 +49,10 @@ class WriteStdinTool(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:
+    async def run(self, context: 'ToolContext', args: 'JSONDict') -> 'JSONValue':
         del context
         session_id = args.get("session_id")
         if session_id is None:
@@ -68,7 +67,7 @@ class WriteStdinTool(BaseTool):
             max_output_tokens=self._optional_int(args, "max_output_tokens"),
         )
 
-    def _optional_int(self, args: JSONDict, key: str) -> int | None:
+    def _optional_int(self, args: 'JSONDict', key: 'str') -> 'typing.Union[int, None]':
         value = args.get(key)
         if value in (None, ""):
             return None

pycodex/utils/dotenv.py

@@ -1,14 +1,14 @@
-from __future__ import annotations
 
 import os
 from pathlib import Path
+import typing
 
 ILLEGAL_ENV_VAR_PREFIX = "CODEX_"
 DOTENV_FILENAME = ".env"
-_LOADED_CODEX_DOTENV_HOMES: set[str] = set()
+_LOADED_CODEX_DOTENV_HOMES: 'typing.Set[str]' = set()
 
 
-def load_codex_dotenv(config_path: str | Path) -> None:
+def load_codex_dotenv(config_path: 'typing.Union[str, Path]') -> 'None':
     codex_home = str(Path(config_path).resolve().parent)
     if codex_home in _LOADED_CODEX_DOTENV_HOMES:
         return
@@ -26,8 +26,8 @@ def load_codex_dotenv(config_path: str | Path) -> None:
     _LOADED_CODEX_DOTENV_HOMES.add(codex_home)
 
 
-def parse_dotenv(text: str) -> dict[str, str]:
-    values: dict[str, str] = {}
+def parse_dotenv(text: 'str') -> 'typing.Dict[str, str]':
+    values: 'typing.Dict[str, str]' = {}
     for raw_line in text.splitlines():
         line = raw_line.strip()
         if not line or line.startswith("#"):
@@ -45,7 +45,7 @@ def parse_dotenv(text: str) -> dict[str, str]:
     return values
 
 
-def parse_dotenv_value(raw_value: str) -> str:
+def parse_dotenv_value(raw_value: 'str') -> 'str':
     if not raw_value:
         return ""