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

pycodex/tools/unified_exec_manager.py

@@ -21,15 +21,18 @@ from pathlib import Path
 from loguru import logger
 
 from ..compat import shlex_join, stream_writer_is_closing
+from ..utils.truncation import (
+    DEFAULT_MAX_OUTPUT_TOKENS,
+    approx_token_count,
+    formatted_truncate_text,
+)
 import typing
 
 DEFAULT_EXEC_YIELD_TIME_MS = 10_000
 DEFAULT_WRITE_STDIN_YIELD_TIME_MS = 250
-DEFAULT_MAX_OUTPUT_TOKENS = 10_000
 DEFAULT_LOGIN = True
 DEFAULT_TTY = False
 DEFAULT_SESSION_ID_START = 1000
-APPROX_BYTES_PER_TOKEN = 4
 UNIFIED_EXEC_OUTPUT_MAX_BYTES = 1024 * 1024
 UNIFIED_EXEC_OUTPUT_SCHEMA = {
     "type": "object",
@@ -63,94 +66,6 @@ UNIFIED_EXEC_OUTPUT_SCHEMA = {
     "additionalProperties": False,
 }
 
-
-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':
-    return max(token_count, 0) * APPROX_BYTES_PER_TOKEN
-
-
-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') -> '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',
-) -> 'typing.Tuple[str, str]':
-    if not text:
-        return "", ""
-
-    total_bytes = len(text.encode("utf-8"))
-    tail_start_target = max(total_bytes - end_bytes, 0)
-    prefix_end = 0
-    suffix_start = len(text)
-    suffix_started = False
-    current_byte = 0
-
-    for index, char in enumerate(text):
-        char_bytes = len(char.encode("utf-8"))
-        char_start = current_byte
-        char_end = current_byte + char_bytes
-        if char_end <= beginning_bytes:
-            prefix_end = index + 1
-            current_byte = char_end
-            continue
-        if char_start >= tail_start_target:
-            if not suffix_started:
-                suffix_start = index
-                suffix_started = True
-            current_byte = char_end
-            continue
-        current_byte = char_end
-
-    if suffix_start < prefix_end:
-        suffix_start = prefix_end
-
-    return text[:prefix_end], text[suffix_start:]
-
-
-def _truncate_text(text: 'str', max_tokens: 'int') -> 'str':
-    if not text:
-        return ""
-
-    max_bytes = _approx_bytes_for_tokens(max_tokens)
-    total_bytes = len(text.encode("utf-8"))
-    if total_bytes <= max_bytes:
-        return text
-
-    removed_tokens = _approx_tokens_from_byte_count(total_bytes - max_bytes)
-    marker = f"\u2026{removed_tokens} tokens truncated\u2026"
-    if max_bytes == 0:
-        return marker
-
-    left_budget, right_budget = _split_budget(max_bytes)
-    prefix, suffix = _split_string(text, left_budget, right_budget)
-    return f"{prefix}{marker}{suffix}"
-
-
-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
-
-    total_lines = len(text.splitlines())
-    return f"Total output lines: {total_lines}\n\n{_truncate_text(text, max_tokens)}"
-
-
 @dataclass
 class _HeadTailBuffer:
     max_bytes: 'int' = UNIFIED_EXEC_OUTPUT_MAX_BYTES
@@ -207,6 +122,20 @@ class UnifiedExecManager:
         self._next_session_id = DEFAULT_SESSION_ID_START
         self._sessions: 'typing.Dict[int, UnifiedExecSession]' = {}
         self._lock = asyncio.Lock()
+        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':
+        self._notify_hook = callback
+
+    def running_session_count(self) -> 'int':
+        return sum(
+            1
+            for session in self._sessions.values()
+            if session.process.returncode is None
+        )
 
     async def exec_command(
         self,
@@ -250,11 +179,15 @@ class UnifiedExecManager:
         async with self._lock:
             self._sessions[session_id] = session
 
-        return await self._wait_and_snapshot(
+        output = await self._wait_and_snapshot(
             session_id,
             max(yield_time_ms, 1),
             max_output_tokens,
         )
+        session = await self._get_session(session_id)
+        if session is not None:
+            asyncio.create_task(self._notify_when_session_completes(session_id))
+        return output
 
     async def write_stdin(
         self,
@@ -367,6 +300,29 @@ class UnifiedExecManager:
             session.output_event.set()
         session.output_event.set()
 
+    async def _notify_when_session_completes(self, session_id: 'int') -> 'None':
+        callback = self._notify_hook
+        if callback is None:
+            return
+        session = await self._get_session(session_id)
+        if session is None:
+            return
+        await session.process.wait()
+        async with self._lock:
+            if self._sessions.get(session_id) is not session:
+                return
+        try:
+            await callback(
+                {
+                    "type": "exec_command_completed",
+                    "session_id": session_id,
+                    "exit_code": session.process.returncode,
+                    "command": session.command_display,
+                }
+            )
+        except Exception:  # pragma: no cover - background notification must not break tools
+            return
+
     def _resolve_workdir(self, workdir: 'typing.Union[str, None]') -> 'Path':
         if not workdir:
             return self._default_cwd
@@ -390,11 +346,11 @@ class UnifiedExecManager:
         return [shell_path, "-lc" if login else "-c", cmd]
 
     def _estimate_token_count(self, output: 'str') -> 'typing.Union[int, None]':
-        return _approx_token_count(output)
+        return approx_token_count(output)
 
     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))
+        return formatted_truncate_text(output, max(token_budget, 0))
 
     def _tty_echo(self, chars: 'str') -> 'bytes':
         normalized = chars.replace("\n", "\r\n")

pycodex/tools/update_plan_tool.py

@@ -21,24 +21,32 @@ VALID_PLAN_STATUSES = {"pending", "in_progress", "completed"}
 class UpdatePlanTool(BaseTool):
     name = "update_plan"
     description = (
-        "Updates the task plan. Provide an optional explanation and a list of "
-        "plan items, each with a step and status. At most one step can be "
-        "in_progress at a time."
+        "Updates the task plan.\n"
+        "Provide an optional explanation and a list of plan items, each with a "
+        "step and status.\n"
+        "At most one step can be in_progress at a time.\n"
     )
     input_schema = {
         "type": "object",
         "properties": {
-            "explanation": {"type": "string"},
+            "explanation": {
+                "type": "string",
+                "description": "Optional explanation for this plan update.",
+            },
             "plan": {
                 "type": "array",
                 "description": "The list of steps",
                 "items": {
                     "type": "object",
                     "properties": {
-                        "step": {"type": "string"},
+                        "step": {
+                            "type": "string",
+                            "description": "Task step text.",
+                        },
                         "status": {
                             "type": "string",
-                            "description": "One of: pending, in_progress, completed",
+                            "enum": ["pending", "in_progress", "completed"],
+                            "description": "Step status.",
                         },
                     },
                     "required": ["step", "status"],

pycodex/tools/view_image_tool.py

@@ -6,8 +6,8 @@ Original Codex mapping:
 Expected behavior:
 - Load a local image file and turn it into a data URL that can be attached back
   into the next model request.
-- Accept the documented `path` argument plus the optional `detail: "original"`
-  hint.
+- Accept the documented `path` argument plus optional `detail: "high" |
+  "original"` hint.
 - Return both the JSON object result and the structured `input_image` content
   item that Codex uses when feeding image tool output back to the model.
 """
@@ -28,8 +28,9 @@ VIEW_IMAGE_OUTPUT_SCHEMA = {
             "description": "Data URL for the loaded image.",
         },
         "detail": {
-            "type": ["string", "null"],
-            "description": "Image detail hint returned by view_image. Returns `original` when original resolution is preserved, otherwise `null`.",
+            "type": "string",
+            "enum": ["high", "original"],
+            "description": "Image detail hint returned by view_image. Returns `high` for default resized behavior or `original` when original resolution is preserved.",
         },
     },
     "required": ["image_url", "detail"],
@@ -40,20 +41,20 @@ VIEW_IMAGE_OUTPUT_SCHEMA = {
 class ViewImageTool(BaseTool):
     name = "view_image"
     description = (
-        "View a local image from the filesystem (only use if given a full "
-        "filepath by the user, and the image isn't already attached to the "
-        "thread context within <image ...> tags)."
+        "View a local image file from the filesystem when visual inspection is "
+        "needed. Use this for images already available on disk."
     )
     input_schema = {
         "type": "object",
         "properties": {
             "path": {
                 "type": "string",
-                "description": "Local filesystem path to an image file",
+                "description": "Local filesystem path to an image file.",
             },
             "detail": {
                 "type": "string",
-                "description": "Optional detail override. The only supported value is `original`; omit this field for default resized behavior. Use `original` to preserve the file's original resolution instead of resizing to fit.",
+                "enum": ["high", "original"],
+                "description": "Image detail level. Defaults to `high`; use `original` to preserve exact resolution.",
             },
         },
         "required": ["path"],
@@ -72,13 +73,14 @@ class ViewImageTool(BaseTool):
 
         detail_value = args.get("detail")
         if detail_value in (None, ""):
-            detail = None
-        elif detail_value == "original":
-            detail = "original"
+            detail = "high"
+        elif detail_value in ("high", "original"):
+            detail = str(detail_value)
         else:
             return (
-                "Error: `detail` only supports `original`; omit `detail` for default "
-                f"behavior, got `{detail_value}`."
+                "Error: `detail` only supports `high` or `original`; omit "
+                "`detail` for default high resized behavior, got "
+                f"`{detail_value}`."
             )
 
         path = Path(path_value)
@@ -104,7 +106,6 @@ class ViewImageTool(BaseTool):
         image_item: 'JSONDict' = {
             "type": "input_image",
             "image_url": image_url,
+            "detail": detail,
         }
-        if detail is not None:
-            image_item["detail"] = detail
         return StructuredToolOutput(output=output, content_items=(image_item,))

pycodex/tools/wait_agent_tool.py

@@ -14,6 +14,10 @@ from ..runtime_services import SubAgentManager
 from .agent_tool_schemas import AGENT_STATUS_SCHEMA
 from .base_tool import BaseTool, ToolContext
 
+DEFAULT_WAIT_AGENT_TIMEOUT_MS = 30_000
+MIN_WAIT_AGENT_TIMEOUT_MS = 10_000
+MAX_WAIT_AGENT_TIMEOUT_MS = 3_600_000
+
 WAIT_AGENT_OUTPUT_SCHEMA = {
     "type": "object",
     "properties": {
@@ -37,7 +41,8 @@ class WaitAgentTool(BaseTool):
     description = (
         "Wait for agents to reach a final status. Completed statuses may "
         "include the agent's final message. Returns empty status when timed "
-        "out."
+        "out. Once the agent reaches a final status, a notification message "
+        "will be received containing the same completed status."
     )
     input_schema = {
         "type": "object",
@@ -49,7 +54,7 @@ class WaitAgentTool(BaseTool):
             },
             "timeout_ms": {
                 "type": "integer",
-                "description": "Optional timeout in milliseconds.",
+                "description": "Timeout in milliseconds. Defaults to 30000, min 10000, max 3600000. Prefer longer waits (minutes) to avoid busy polling.",
             },
         },
         "required": ["ids"],
@@ -69,5 +74,12 @@ class WaitAgentTool(BaseTool):
         agent_ids = [str(item).strip() for item in ids if str(item).strip()]
         if not agent_ids:
             return "Error: `ids` must include at least one non-empty id."
-        timeout_ms = int(args.get("timeout_ms", 30_000))
+        timeout_ms = self._timeout_ms(args)
         return await self._subagent_manager.wait_agents(agent_ids, timeout_ms)
+
+    def _timeout_ms(self, args: 'JSONDict') -> 'int':
+        value = int(args.get("timeout_ms", DEFAULT_WAIT_AGENT_TIMEOUT_MS))
+        return min(
+            max(value, MIN_WAIT_AGENT_TIMEOUT_MS),
+            MAX_WAIT_AGENT_TIMEOUT_MS,
+        )

pycodex/tools/wait_tool.py

@@ -18,7 +18,21 @@ import typing
 class WaitTool(BaseTool):
     name = "wait"
     description = (
-        "Waits on a yielded `exec` cell and returns new output or completion."
+        "Waits on a yielded `exec` cell and returns new output or completion.\n"
+        "- Use `wait` only after `exec` returns `Script running with cell ID ...`.\n"
+        "- `cell_id` identifies the running `exec` cell to resume.\n"
+        "- `yield_time_ms` controls how long to wait for more output before "
+        "yielding again. Defaults to 10000 ms.\n"
+        "- `max_tokens` limits how much new output this wait call returns. "
+        "Defaults to 10000 tokens.\n"
+        "- `terminate: true` stops the running cell; false or omitted waits "
+        "for output.\n"
+        "- `wait` returns only the new output since the last yield, or the "
+        "final completion or termination result for that cell.\n"
+        "- If the cell is still running, `wait` may yield again with the same "
+        "`cell_id`.\n"
+        "- If the cell has already finished, `wait` returns the completed "
+        "result and closes the cell."
     )
     input_schema = {
         "type": "object",
@@ -29,15 +43,15 @@ class WaitTool(BaseTool):
             },
             "yield_time_ms": {
                 "type": "integer",
-                "description": "How long to wait (in milliseconds) for more output before yielding again.",
+                "description": "Wait before yielding more output. Defaults to 10000 ms.",
             },
             "max_tokens": {
                 "type": "integer",
-                "description": "Maximum number of output tokens to return for this wait call.",
+                "description": "Output token budget for this wait call. Defaults to 10000 tokens.",
             },
             "terminate": {
                 "type": "boolean",
-                "description": "Whether to terminate the running exec cell.",
+                "description": "True stops the running exec cell; false or omitted waits for output.",
             },
         },
         "required": ["cell_id"],

pycodex/tools/web_search_tool.py

@@ -16,10 +16,11 @@ from .base_tool import BaseTool, ToolContext
 
 class WebSearchTool(BaseTool):
     name = "web_search"
-    description = "Provider-native web search tool declaration."
+    description = ""
     tool_type = "web_search"
     options = {
         "external_web_access": True,
+        "search_content_types": ["text", "image"],
     }
     supports_parallel = False
 

pycodex/tools/write_stdin_tool.py

@@ -19,6 +19,11 @@ from .unified_exec_manager import (
 )
 import typing
 
+MIN_WRITE_YIELD_TIME_MS = 250
+MAX_WRITE_YIELD_TIME_MS = 30_000
+DEFAULT_WRITE_STDIN_POLL_YIELD_TIME_MS = 5_000
+MAX_WRITE_STDIN_POLL_YIELD_TIME_MS = 300_000
+
 
 class WriteStdinTool(BaseTool):
     name = "write_stdin"
@@ -27,20 +32,20 @@ class WriteStdinTool(BaseTool):
         "type": "object",
         "properties": {
             "session_id": {
-                "type": "integer",
+                "type": "number",
                 "description": "Identifier of the running unified exec session.",
             },
             "chars": {
                 "type": "string",
-                "description": "Bytes to write to stdin (may be empty to poll).",
+                "description": "Bytes to write to stdin. Defaults to empty, which polls without writing.",
             },
             "yield_time_ms": {
-                "type": "integer",
-                "description": "How long to wait (in milliseconds) for output before yielding.",
+                "type": "number",
+                "description": "Wait before yielding output. Non-empty writes default to 250 ms and cap at 30000 ms; empty polls wait 5000-300000 ms by default.",
             },
             "max_output_tokens": {
-                "type": "integer",
-                "description": "Maximum number of tokens to return. Excess output will be truncated.",
+                "type": "number",
+                "description": "Output token budget. Defaults to 10000 tokens; larger requests may be capped by policy.",
             },
         },
         "required": ["session_id"],
@@ -57,18 +62,45 @@ class WriteStdinTool(BaseTool):
         session_id = args.get("session_id")
         if session_id is None:
             return "Error: `session_id` is required."
+        chars = str(args.get("chars", ""))
 
         return await self._manager.write_stdin(
             session_id=int(session_id),
-            chars=str(args.get("chars", "")),
-            yield_time_ms=int(
-                args.get("yield_time_ms", DEFAULT_WRITE_STDIN_YIELD_TIME_MS)
-            ),
+            chars=chars,
+            yield_time_ms=self._yield_time_ms(args, chars),
             max_output_tokens=self._optional_int(args, "max_output_tokens"),
         )
 
+    def _yield_time_ms(self, args: 'JSONDict', chars: 'str') -> 'int':
+        if chars:
+            return self._bounded_int(
+                args,
+                "yield_time_ms",
+                DEFAULT_WRITE_STDIN_YIELD_TIME_MS,
+                MIN_WRITE_YIELD_TIME_MS,
+                MAX_WRITE_YIELD_TIME_MS,
+            )
+        return self._bounded_int(
+            args,
+            "yield_time_ms",
+            DEFAULT_WRITE_STDIN_POLL_YIELD_TIME_MS,
+            DEFAULT_WRITE_STDIN_POLL_YIELD_TIME_MS,
+            MAX_WRITE_STDIN_POLL_YIELD_TIME_MS,
+        )
+
     def _optional_int(self, args: 'JSONDict', key: 'str') -> 'typing.Union[int, None]':
         value = args.get(key)
         if value in (None, ""):
             return None
         return int(value)
+
+    def _bounded_int(
+        self,
+        args: 'JSONDict',
+        key: 'str',
+        default: 'int',
+        minimum: 'int',
+        maximum: 'int',
+    ) -> 'int':
+        value = int(args.get(key, default))
+        return min(max(value, minimum), maximum)

pycodex/utils/compactor.py

@@ -37,6 +37,7 @@ SUMMARY_PREFIX = (
 COMPACT_USER_MESSAGE_MAX_TOKENS = 20_000
 _APPROX_CHARS_PER_TOKEN = 4
 _SUBAGENT_NOTIFICATION_PREFIX = "<subagent_notification>\n"
+_EXEC_COMMAND_COMPLETED_PREFIX = "<exec_command_completed>\n"
 
 
 @dataclass(frozen=True)
@@ -253,7 +254,12 @@ def _pluralize(noun: 'str', count: 'int') -> 'str':
 
 
 def _is_synthetic_user_message(text: 'str') -> 'bool':
-    return text.startswith(_SUBAGENT_NOTIFICATION_PREFIX)
+    return text.startswith(
+        (
+            _SUBAGENT_NOTIFICATION_PREFIX,
+            _EXEC_COMMAND_COMPLETED_PREFIX,
+        )
+    )
 
 
 def _is_context_length_error(message: 'str') -> 'bool':

pycodex/utils/session_persist.py

@@ -198,7 +198,18 @@ def load_resumed_session(
     session = sessions[resume_index - 1]
     thread_id = session["thread_id"]
     rollout_path = Path(session["rollout_path"])
-    thread_name = _latest_thread_names_by_id(codex_home).get(thread_id)
+    return load_resumed_session_path(
+        rollout_path,
+        thread_name=_latest_thread_names_by_id(codex_home).get(thread_id),
+    )
+
+
+def load_resumed_session_path(
+    rollout_path: 'typing.Union[str, Path]',
+    thread_name: 'typing.Union[str, None]' = None,
+) -> 'typing.Dict[str, object]':
+    rollout_path = Path(rollout_path)
+    thread_id = _thread_id_from_rollout_path(rollout_path) or ""
     session_id = thread_id
     history: 'typing.List[ConversationItem]' = []
     saw_user_turn = False
@@ -245,6 +256,10 @@ def load_resumed_session(
     if not history:
         raise ValueError(f"No resumable history found in {rollout_path}")
 
+    history = _trim_incomplete_tool_call_tail(history)
+    if not history:
+        raise ValueError(f"No resumable history found in {rollout_path}")
+
     turns = conversation_history_to_turns(history)
     title = thread_name or (shorten_title(turns[0][0]) if turns else thread_id)
     return {
@@ -277,6 +292,32 @@ def conversation_history_to_turns(
     return tuple(turns)
 
 
+def _trim_incomplete_tool_call_tail(
+    history: 'typing.List[ConversationItem]',
+) -> 'typing.List[ConversationItem]':
+    pending_call_ids: 'typing.Set[str]' = set()
+    call_indexes: 'typing.Dict[str, int]' = {}
+
+    for index, item in enumerate(history):
+        if isinstance(item, ToolCall):
+            pending_call_ids.add(item.call_id)
+            call_indexes[item.call_id] = index
+            continue
+        if isinstance(item, ToolResult):
+            pending_call_ids.discard(item.call_id)
+
+    if not pending_call_ids:
+        return history
+
+    trim_start = min(call_indexes[call_id] for call_id in pending_call_ids)
+    while trim_start > 0 and isinstance(
+        history[trim_start - 1],
+        (AssistantMessage, ReasoningItem, ToolCall),
+    ):
+        trim_start -= 1
+    return history[:trim_start]
+
+
 def _latest_thread_names_by_id(codex_home: 'Path') -> 'typing.Dict[str, str]':
     index_path = codex_home / SESSION_INDEX_FILENAME
     if not index_path.exists():