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

pycodex/tools/shell_command_tool.py

@@ -0,0 +1,107 @@
+"""`shell_command` tool for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the original Codex `shell_command` tool.
+
+Expected behavior:
+- Execute a shell-script string in the user's shell context.
+- Accept the same core input shape as Codex: `command: string`, plus `workdir`,
+  `timeout_ms`, and `login`.
+- Return a concise text summary including working directory, exit status,
+  stdout, and stderr.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from ..protocol import JSONDict, JSONValue
+from .base_tool import BaseTool, ToolContext
+
+DEFAULT_SHELL_TIMEOUT_MS = 30_000
+MAX_OUTPUT_CHARS = 12_000
+
+
+class ShellCommandTool(BaseTool):
+    name = "shell_command"
+    description = "Runs a shell command string and returns its output."
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "command": {"type": "string"},
+            "workdir": {"type": "string"},
+            "timeout_ms": {"type": "integer"},
+            "login": {"type": "boolean"},
+        },
+        "required": ["command"],
+    }
+    supports_parallel = False
+
+    def __init__(self, cwd: str | Path | None = None) -> None:
+        self._working_directory = Path(cwd or Path.cwd()).resolve()
+
+    async def run(self, context: ToolContext, args: JSONDict) -> JSONValue:
+        del context
+        command = str(args.get("command", "")).strip()
+        timeout_ms = int(args.get("timeout_ms", DEFAULT_SHELL_TIMEOUT_MS))
+        if not command:
+            return "Error: `command` is required."
+
+        login = bool(args.get("login", True))
+        workdir_arg = args.get("workdir")
+        working_directory = self._resolve_workdir(workdir_arg)
+        shell_args = ["bash", "-lc" if login else "-c", command]
+
+        process = await asyncio.create_subprocess_exec(
+            *shell_args,
+            cwd=str(working_directory),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+
+        try:
+            stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                process.communicate(),
+                timeout=max(timeout_ms, 1) / 1000.0,
+            )
+            timed_out = False
+        except asyncio.TimeoutError:
+            process.kill()
+            stdout_bytes, stderr_bytes = await process.communicate()
+            timed_out = True
+
+        stdout = stdout_bytes.decode("utf-8", errors="replace")
+        stderr = stderr_bytes.decode("utf-8", errors="replace")
+        pieces = [f"Working directory: {working_directory}"]
+
+        if timed_out:
+            pieces.append(f"Timeout: exceeded {timeout_ms} ms")
+        else:
+            pieces.append(f"Exit code: {process.returncode}")
+
+        stdout = self._clip_output(stdout)
+        stderr = self._clip_output(stderr)
+
+        if stdout:
+            pieces.append("Stdout:")
+            pieces.append(stdout)
+
+        if stderr:
+            pieces.append("Stderr:")
+            pieces.append(stderr)
+
+        return "\n".join(pieces)
+
+    def _resolve_workdir(self, workdir_arg) -> Path:
+        if workdir_arg in (None, ""):
+            return self._working_directory
+        workdir = Path(str(workdir_arg))
+        if not workdir.is_absolute():
+            workdir = self._working_directory / workdir
+        return workdir.resolve()
+
+    def _clip_output(self, text: str) -> str:
+        if len(text) <= MAX_OUTPUT_CHARS:
+            return text
+        return text[:MAX_OUTPUT_CHARS] + "\n...[truncated]..."

pycodex/tools/shell_tool.py

@@ -0,0 +1,112 @@
+"""`shell` tool for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the original Codex `shell` tool.
+
+Expected behavior:
+- Execute a command as argv, not as a shell-script string.
+- Accept the same core input shape as Codex: `command: string[]`, plus
+  `workdir` and `timeout_ms`.
+- Return a concise text summary including working directory, exit status,
+  stdout, and stderr.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from ..protocol import JSONDict, JSONValue
+from .base_tool import BaseTool, ToolContext
+
+DEFAULT_SHELL_TIMEOUT_MS = 30_000
+MAX_OUTPUT_CHARS = 12_000
+
+
+class ShellTool(BaseTool):
+    name = "shell"
+    description = (
+        "Runs a shell command and returns its output. The command must be passed "
+        "as argv, typically prefixed with ['bash', '-lc'] for shell syntax."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "command": {
+                "type": "array",
+                "items": {"type": "string"},
+            },
+            "workdir": {"type": "string"},
+            "timeout_ms": {"type": "integer"},
+        },
+        "required": ["command"],
+    }
+    supports_parallel = False
+
+    def __init__(self, cwd: str | Path | None = None) -> None:
+        self._working_directory = Path(cwd or Path.cwd()).resolve()
+
+    async def run(self, context: ToolContext, args: JSONDict) -> JSONValue:
+        del context
+        command = args.get("command")
+        timeout_ms = int(args.get("timeout_ms", DEFAULT_SHELL_TIMEOUT_MS))
+        if not isinstance(command, list) or not command:
+            return "Error: `command` must be a non-empty string array."
+        if not all(isinstance(part, str) and part for part in command):
+            return "Error: each `command` entry must be a non-empty string."
+
+        workdir_arg = args.get("workdir")
+        working_directory = self._resolve_workdir(workdir_arg)
+
+        process = await asyncio.create_subprocess_exec(
+            *command,
+            cwd=str(working_directory),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+
+        try:
+            stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                process.communicate(),
+                timeout=max(timeout_ms, 1) / 1000.0,
+            )
+            timed_out = False
+        except asyncio.TimeoutError:
+            process.kill()
+            stdout_bytes, stderr_bytes = await process.communicate()
+            timed_out = True
+
+        stdout = stdout_bytes.decode("utf-8", errors="replace")
+        stderr = stderr_bytes.decode("utf-8", errors="replace")
+        pieces = [f"Working directory: {working_directory}"]
+
+        if timed_out:
+            pieces.append(f"Timeout: exceeded {timeout_ms} ms")
+        else:
+            pieces.append(f"Exit code: {process.returncode}")
+
+        stdout = self._clip_output(stdout)
+        stderr = self._clip_output(stderr)
+
+        if stdout:
+            pieces.append("Stdout:")
+            pieces.append(stdout)
+
+        if stderr:
+            pieces.append("Stderr:")
+            pieces.append(stderr)
+
+        return "\n".join(pieces)
+
+    def _resolve_workdir(self, workdir_arg) -> Path:
+        if workdir_arg in (None, ""):
+            return self._working_directory
+        workdir = Path(str(workdir_arg))
+        if not workdir.is_absolute():
+            workdir = self._working_directory / workdir
+        return workdir.resolve()
+
+    def _clip_output(self, text: str) -> str:
+        if len(text) <= MAX_OUTPUT_CHARS:
+            return text
+        return text[:MAX_OUTPUT_CHARS] + "\n...[truncated]..."

pycodex/tools/spawn_agent_tool.py

@@ -0,0 +1,97 @@
+"""`spawn_agent` tool for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the original Codex `spawn_agent` collaboration tool.
+
+Expected behavior:
+- Spawn a sibling agent runtime that can work in parallel with the caller.
+- Optionally seed the spawned agent with the current thread history.
+- Return the new agent identifier plus any user-facing nickname.
+"""
+
+from __future__ import annotations
+
+from ..protocol import JSONDict, JSONValue
+from ..runtime_services import SubAgentManager
+from .agent_tool_schemas import COLLAB_INPUT_ITEMS_SCHEMA
+from .base_tool import BaseTool, ToolContext
+
+SPAWN_AGENT_OUTPUT_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "agent_id": {
+            "type": "string",
+            "description": "Thread identifier for the spawned agent.",
+        },
+        "nickname": {
+            "type": ["string", "null"],
+            "description": "User-facing nickname for the spawned agent when available.",
+        },
+    },
+    "required": ["agent_id", "nickname"],
+    "additionalProperties": False,
+}
+
+
+class SpawnAgentTool(BaseTool):
+    name = "spawn_agent"
+    description = (
+        "Spawn a sub-agent for a well-scoped task. Returns the agent id (and "
+        "user-facing nickname when available) to use to communicate with this "
+        "agent."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "message": {
+                "type": "string",
+                "description": "Initial plain-text task for the new agent. Use either message or items.",
+            },
+            "items": COLLAB_INPUT_ITEMS_SCHEMA,
+            "agent_type": {
+                "type": "string",
+                "description": "Optional type name for the new agent.",
+            },
+            "fork_context": {
+                "type": "boolean",
+                "description": "When true, fork the current thread history into the new agent before sending the initial prompt.",
+            },
+            "model": {
+                "type": "string",
+                "description": "Optional model override for the new agent.",
+            },
+            "reasoning_effort": {
+                "type": "string",
+                "description": "Optional reasoning effort override for the new agent.",
+            },
+        },
+        "additionalProperties": False,
+    }
+    output_schema = SPAWN_AGENT_OUTPUT_SCHEMA
+    supports_parallel = False
+
+    def __init__(self, subagent_manager: SubAgentManager) -> None:
+        self._subagent_manager = subagent_manager
+
+    async def run(self, context: ToolContext, args: JSONDict) -> JSONValue:
+        message = self._optional_string(args, "message")
+        items = args.get("items")
+        if items is not None and not isinstance(items, list):
+            return "Error: `items` must be a list when provided."
+        if message is None and not items:
+            return "Provide one of: message or items"
+        return await self._subagent_manager.spawn_agent(
+            message=message,
+            items=items,
+            agent_type=self._optional_string(args, "agent_type"),
+            fork_context=bool(args.get("fork_context", False)),
+            model=self._optional_string(args, "model"),
+            reasoning_effort=self._optional_string(args, "reasoning_effort"),
+            history=context.history,
+        )
+
+    def _optional_string(self, args: JSONDict, key: str) -> str | None:
+        value = args.get(key)
+        if value in (None, ""):
+            return None
+        return str(value)

pycodex/tools/unified_exec_manager.py

@@ -0,0 +1,380 @@
+"""Shared runtime for `exec_command` / `write_stdin`.
+
+Original Codex mapping:
+- Corresponds to the shared unified-exec session manager behind the original
+  Codex `exec_command` and `write_stdin` tools.
+
+Expected behavior:
+- Start a long-lived command session for `exec_command`.
+- Preserve unread output between calls.
+- Accept follow-up stdin writes and polling through `write_stdin`.
+- Return summaries in the same textual shape Codex tools expect.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import shlex
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from loguru import logger
+
+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",
+    "properties": {
+        "chunk_id": {
+            "type": "string",
+            "description": "Chunk identifier included when the response reports one.",
+        },
+        "wall_time_seconds": {
+            "type": "number",
+            "description": "Elapsed wall time spent waiting for output in seconds.",
+        },
+        "exit_code": {
+            "type": "number",
+            "description": "Process exit code when the command finished during this call.",
+        },
+        "session_id": {
+            "type": "number",
+            "description": "Session identifier to pass to write_stdin when the process is still running.",
+        },
+        "original_token_count": {
+            "type": "number",
+            "description": "Approximate token count before output truncation.",
+        },
+        "output": {
+            "type": "string",
+            "description": "Command output text, possibly truncated.",
+        },
+    },
+    "required": ["wall_time_seconds", "output"],
+    "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) -> 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]:
+    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(slots=True)
+class _HeadTailBuffer:
+    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:
+        if not chunk or self.max_bytes <= 0:
+            return
+
+        head_budget = self.max_bytes // 2
+        tail_budget = self.max_bytes - head_budget
+        remaining = bytes(chunk)
+
+        if len(self.head) < head_budget:
+            head_room = head_budget - len(self.head)
+            head_part = remaining[:head_room]
+            self.head.extend(head_part)
+            remaining = remaining[len(head_part) :]
+
+        if not remaining or tail_budget <= 0:
+            return
+
+        self.tail.extend(remaining)
+        if len(self.tail) > tail_budget:
+            excess = len(self.tail) - tail_budget
+            del self.tail[:excess]
+
+    def drain_bytes(self) -> bytes:
+        combined = bytes(self.head) + bytes(self.tail)
+        self.head.clear()
+        self.tail.clear()
+        return combined
+
+
+@dataclass(slots=True)
+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
+
+
+class UnifiedExecManager:
+    def __init__(self, cwd: 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._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:
+        session_id = await self._allocate_session_id()
+        command = self._build_shell_command(cmd, shell, login)
+        cwd = self._resolve_workdir(workdir)
+        logger.debug(
+            "exec_command start session_id={} shell_command={} cwd={}",
+            session_id,
+            command,
+            cwd,
+        )
+
+        process = await asyncio.create_subprocess_exec(
+            *command,
+            cwd=str(cwd),
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        session = UnifiedExecSession(
+            session_id=session_id,
+            process=process,
+            start_time=asyncio.get_running_loop().time(),
+            command_display=shlex.join(command),
+            tty=tty,
+        )
+        session.reader_task = asyncio.create_task(self._pump_output(session))
+
+        async with self._lock:
+            self._sessions[session_id] = session
+
+        return await self._wait_and_snapshot(
+            session_id,
+            max(yield_time_ms, 1),
+            max_output_tokens,
+        )
+
+    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 = await self._get_session(session_id)
+        if session is None:
+            return f"Error: session_id {session_id} is not running."
+
+        if chars:
+            if session.process.stdin is None:
+                return f"Error: session_id {session_id} does not accept stdin."
+            logger.debug("write_stdin session_id={} chars_len={}", session_id, len(chars))
+            if session.tty:
+                session.unread_output.push_chunk(self._tty_echo(chars))
+            session.process.stdin.write(chars.encode("utf-8"))
+            await session.process.stdin.drain()
+
+        return await self._wait_and_snapshot(
+            session_id,
+            max(yield_time_ms, 1),
+            max_output_tokens,
+        )
+
+    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 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 = 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()
+        try:
+            await asyncio.wait_for(session.process.wait(), timeout=yield_time_ms / 1000.0)
+        except asyncio.TimeoutError:
+            pass
+
+        if session.reader_task is not None and session.process.returncode is not None:
+            await session.reader_task
+
+        wall_time = asyncio.get_running_loop().time() - start_wait
+        output_bytes = session.unread_output.drain_bytes()
+        output_text = output_bytes.decode("utf-8", errors="replace")
+        original_token_count = self._estimate_token_count(output_text)
+        output_text = self._truncate_output(output_text, max_output_tokens)
+
+        lines = [
+            f"Command: {session.command_display}",
+            f"Chunk ID: {uuid.uuid4().hex[:6]}",
+            f"Wall time: {wall_time:.4f} seconds",
+        ]
+        if session.process.returncode is None:
+            lines.append(f"Process running with session ID {session_id}")
+        else:
+            lines.append(f"Process exited with code {session.process.returncode}")
+        if original_token_count is not None:
+            lines.append(f"Original token count: {original_token_count}")
+        lines.append("Output:")
+        lines.append(output_text)
+
+        if session.process.returncode is not None:
+            await self._close_session(session_id)
+
+        return "\n".join(lines)
+
+    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():
+            session.process.stdin.close()
+
+    async def _pump_output(self, session: UnifiedExecSession) -> None:
+        stream = session.process.stdout
+        if stream is None:
+            return
+        while True:
+            chunk = await stream.read(4096)
+            if not chunk:
+                break
+            session.unread_output.push_chunk(chunk)
+
+    def _resolve_workdir(self, workdir: str | None) -> Path:
+        if not workdir:
+            return self._default_cwd
+        path = Path(workdir)
+        if not path.is_absolute():
+            path = self._default_cwd / path
+        return path.resolve()
+
+    def _build_shell_command(
+        self,
+        cmd: str,
+        shell: str | None,
+        login: bool,
+    ) -> 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"}:
+            return [shell_path, "/C", cmd]
+        if "powershell" in shell_name:
+            return [shell_path, "-NoProfile", "-Command", cmd]
+        return [shell_path, "-lc" if login else "-c", cmd]
+
+    def _estimate_token_count(self, output: str) -> int | None:
+        return _approx_token_count(output)
+
+    def _truncate_output(self, output: str, max_output_tokens: 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:
+        normalized = chars.replace("\n", "\r\n")
+        return normalized.encode("utf-8")