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

pycodex/tools/agent_tool_schemas.py

@@ -0,0 +1,70 @@
+"""Shared schemas for Codex-aligned collaboration tools."""
+
+from __future__ import annotations
+
+COLLAB_INPUT_ITEMS_SCHEMA = {
+    "type": "array",
+    "description": (
+        "Structured input items. Use this to pass explicit mentions (for "
+        "example app:// connector paths)."
+    ),
+    "items": {
+        "type": "object",
+        "properties": {
+            "type": {
+                "type": "string",
+                "description": "Input item type: text, image, local_image, skill, or mention.",
+            },
+            "text": {
+                "type": "string",
+                "description": "Text content when type is text.",
+            },
+            "image_url": {
+                "type": "string",
+                "description": "Image URL when type is image.",
+            },
+            "path": {
+                "type": "string",
+                "description": (
+                    "Path when type is local_image/skill, or structured mention target "
+                    "such as app://<connector-id> or plugin://<plugin-name>@<marketplace-name> "
+                    "when type is mention."
+                ),
+            },
+            "name": {
+                "type": "string",
+                "description": "Display name when type is skill or mention.",
+            },
+        },
+        "additionalProperties": False,
+    },
+}
+
+AGENT_STATUS_SCHEMA = {
+    "oneOf": [
+        {
+            "type": "string",
+            "enum": ["pending_init", "running", "shutdown", "not_found"],
+        },
+        {
+            "type": "object",
+            "properties": {
+                "completed": {
+                    "type": ["string", "null"],
+                }
+            },
+            "required": ["completed"],
+            "additionalProperties": False,
+        },
+        {
+            "type": "object",
+            "properties": {
+                "errored": {
+                    "type": "string",
+                }
+            },
+            "required": ["errored"],
+            "additionalProperties": False,
+        },
+    ]
+}

pycodex/tools/apply_patch_tool.py

@@ -0,0 +1,363 @@
+"""`apply_patch` tool for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the original Codex `apply_patch` freeform/custom tool.
+
+Expected behavior:
+- Accept the same freeform patch envelope used by Codex main.
+- Verify the whole patch before mutating the filesystem so failed patches do not
+  leave partial edits behind.
+- Apply add/delete/update/move operations inside the workspace and return the
+  same success/error text shape Codex expects.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from loguru import logger
+
+from ..protocol import JSONValue
+from .base_tool import BaseTool, ToolContext
+
+APPLY_PATCH_LARK_GRAMMAR = """start: begin_patch hunk+ end_patch
+begin_patch: \"*** Begin Patch\" LF
+end_patch: \"*** End Patch\" LF?
+
+hunk: add_hunk | delete_hunk | update_hunk
+add_hunk: \"*** Add File: \" filename LF add_line+
+delete_hunk: \"*** Delete File: \" filename LF
+update_hunk: \"*** Update File: \" filename LF change_move? change?
+
+filename: /(.+)/
+add_line: \"+\" /(.*)/ LF -> line
+
+change_move: \"*** Move to: \" filename LF
+change: (change_context | change_line)+ eof_line?
+change_context: (\"@@\" | \"@@ \" /(.+)/) LF
+change_line: (\"+\" | \"-\" | \" \" ) /(.*)/ LF
+eof_line: \"*** End of File\" LF
+
+%import common.LF
+"""
+
+
+class ApplyPatchError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _AddFileOp:
+    path: str
+    content: str
+
+
+@dataclass(frozen=True, slots=True)
+class _DeleteFileOp:
+    path: str
+
+
+@dataclass(frozen=True, slots=True)
+class _UpdateSection:
+    lines: tuple[str, ...]
+    anchor_end_of_file: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class _UpdateFileOp:
+    path: str
+    move_to: str | None
+    sections: tuple[_UpdateSection, ...]
+
+
+class ApplyPatchTool(BaseTool):
+    name = "apply_patch"
+    description = (
+        "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, "
+        "so do not wrap the patch in JSON."
+    )
+    tool_type = "custom"
+    format = {
+        "type": "grammar",
+        "syntax": "lark",
+        "definition": APPLY_PATCH_LARK_GRAMMAR,
+    }
+    supports_parallel = False
+
+    def __init__(self, cwd: str | Path | None = None) -> None:
+        self._workspace_root = Path(cwd or Path.cwd()).resolve()
+
+    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))
+        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) -> list[_AddFileOp | _DeleteFileOp | _UpdateFileOp]:
+        lines = patch_text.splitlines()
+        if not lines:
+            raise ApplyPatchError("patch rejected: empty patch")
+        if lines[0] != "*** Begin Patch":
+            raise ApplyPatchError(
+                "apply_patch verification failed: missing '*** Begin Patch' header"
+            )
+
+        operations: list[_AddFileOp | _DeleteFileOp | _UpdateFileOp] = []
+        index = 1
+        while index < len(lines):
+            line = lines[index]
+            if line == "*** End Patch":
+                if not operations:
+                    raise ApplyPatchError("patch rejected: empty patch")
+                for trailing in lines[index + 1 :]:
+                    if trailing.strip():
+                        raise ApplyPatchError(
+                            "apply_patch verification failed: unexpected content after '*** End Patch'"
+                        )
+                return operations
+
+            if line.startswith("*** Add File: "):
+                path = line.removeprefix("*** Add File: ")
+                index += 1
+                content_lines: list[str] = []
+                while index < len(lines) and not lines[index].startswith("*** "):
+                    entry = lines[index]
+                    if not entry.startswith("+"):
+                        raise ApplyPatchError(
+                            f"apply_patch verification failed: {entry!r} is not a valid add line"
+                        )
+                    content_lines.append(entry[1:])
+                    index += 1
+                if not content_lines:
+                    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)))
+                continue
+
+            if line.startswith("*** Delete File: "):
+                path = line.removeprefix("*** Delete File: ")
+                operations.append(_DeleteFileOp(path=path))
+                index += 1
+                continue
+
+            if line.startswith("*** Update File: "):
+                path = line.removeprefix("*** Update File: ")
+                index += 1
+                move_to = None
+                if index < len(lines) and lines[index].startswith("*** Move to: "):
+                    move_to = lines[index].removeprefix("*** Move to: ")
+                    index += 1
+
+                sections: list[_UpdateSection] = []
+                current_lines: list[str] = []
+                saw_hunk_header = False
+                anchor_end_of_file = False
+                while index < len(lines):
+                    entry = lines[index]
+                    if entry.startswith("*** ") and entry != "*** End of File":
+                        break
+                    if entry == "@@" or entry.startswith("@@ "):
+                        if saw_hunk_header:
+                            sections.append(
+                                _UpdateSection(
+                                    lines=tuple(current_lines),
+                                    anchor_end_of_file=anchor_end_of_file,
+                                )
+                            )
+                            current_lines = []
+                            anchor_end_of_file = False
+                        saw_hunk_header = True
+                        index += 1
+                        continue
+                    if entry == "*** End of File":
+                        if not saw_hunk_header:
+                            raise ApplyPatchError(
+                                "apply_patch verification failed: '*** End of File' must follow a hunk"
+                            )
+                        anchor_end_of_file = True
+                        index += 1
+                        continue
+                    if not saw_hunk_header:
+                        raise ApplyPatchError(
+                            f"apply_patch verification failed: {entry!r} is not a valid hunk header"
+                        )
+                    if not entry or entry[0] not in {" ", "+", "-"}:
+                        raise ApplyPatchError(
+                            f"apply_patch verification failed: {entry!r} is not a valid change line"
+                        )
+                    current_lines.append(entry)
+                    index += 1
+
+                if not saw_hunk_header:
+                    raise ApplyPatchError(
+                        f"apply_patch verification failed: update for {path} is missing a hunk"
+                    )
+                sections.append(
+                    _UpdateSection(
+                        lines=tuple(current_lines),
+                        anchor_end_of_file=anchor_end_of_file,
+                    )
+                )
+                operations.append(
+                    _UpdateFileOp(
+                        path=path,
+                        move_to=move_to,
+                        sections=tuple(sections),
+                    )
+                )
+                continue
+
+            raise ApplyPatchError(
+                f"apply_patch verification failed: {line!r} is not a valid hunk header"
+            )
+
+        raise ApplyPatchError("apply_patch verification failed: missing '*** End Patch' footer")
+
+    def _apply_operations(
+        self,
+        operations: list[_AddFileOp | _DeleteFileOp | _UpdateFileOp],
+    ) -> str:
+        preview: dict[Path, str | None] = {}
+        summaries: dict[Path, str] = {}
+
+        for operation in operations:
+            if isinstance(operation, _AddFileOp):
+                path = self._resolve_workspace_path(operation.path)
+                preview[path] = operation.content
+                summaries[path] = "A"
+                continue
+
+            if isinstance(operation, _DeleteFileOp):
+                path = self._resolve_workspace_path(operation.path)
+                self._read_preview_file(path, preview)
+                preview[path] = None
+                summaries[path] = "D"
+                continue
+
+            path = self._resolve_workspace_path(operation.path)
+            original = self._read_preview_file(path, preview)
+            updated = self._apply_update(path, original, operation.sections)
+            destination = path
+            if operation.move_to is not None:
+                destination = self._resolve_workspace_path(operation.move_to)
+                preview[path] = None
+                summaries.pop(path, None)
+            preview[destination] = updated
+            summaries[destination] = "M"
+
+        self._write_preview(preview)
+        return self._format_success(summaries)
+
+    def _read_preview_file(self, path: Path, preview: dict[Path, str | None]) -> str:
+        if path in preview:
+            content = preview[path]
+            if content is None:
+                raise ApplyPatchError(
+                    f"apply_patch verification failed: Failed to read {path.relative_to(self._workspace_root)}"
+                )
+            return content
+
+        if not path.exists() or not path.is_file():
+            raise ApplyPatchError(
+                f"apply_patch verification failed: Failed to read {path.relative_to(self._workspace_root)}"
+            )
+        return path.read_text(encoding="utf-8", errors="replace")
+
+    def _apply_update(
+        self,
+        path: Path,
+        original_text: str,
+        sections: tuple[_UpdateSection, ...],
+    ) -> str:
+        lines = original_text.splitlines()
+        cursor = 0
+        for section in sections:
+            old_block = [line[1:] for line in section.lines if line[:1] in {" ", "-"}]
+            new_block = [line[1:] for line in section.lines if line[:1] in {" ", "+"}]
+            if not old_block and not new_block:
+                continue
+
+            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 "
+                    f"{path.relative_to(self._workspace_root)}"
+                )
+            lines[match_index : match_index + len(old_block)] = new_block
+            cursor = match_index + len(new_block)
+        return self._join_lines(lines)
+
+    def _find_match(
+        self,
+        lines: list[str],
+        old_block: list[str],
+        cursor: int,
+        anchor_end_of_file: bool,
+    ) -> 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:
+                return start
+            return None
+
+        if not old_block:
+            return cursor
+
+        limit = len(lines) - len(old_block) + 1
+        for start in range(max(cursor, 0), max(limit, 0)):
+            if lines[start : start + len(old_block)] == old_block:
+                return start
+        for start in range(0, max(limit, 0)):
+            if lines[start : start + len(old_block)] == old_block:
+                return start
+        return None
+
+    def _write_preview(self, preview: dict[Path, str | None]) -> None:
+        for path, content in preview.items():
+            if content is None:
+                if path.exists():
+                    path.unlink()
+                continue
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_text(content, encoding="utf-8")
+
+    def _format_success(self, summaries: dict[Path, str]) -> str:
+        buckets = {"A": [], "M": [], "D": []}
+        for path, status in summaries.items():
+            buckets[status].append(path.relative_to(self._workspace_root).as_posix())
+        lines = ["Success. Updated the following files:"]
+        for status in ("A", "M", "D"):
+            for rel_path in sorted(buckets[status]):
+                lines.append(f"{status} {rel_path}")
+        return "\n".join(lines) + "\n"
+
+    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}"
+        )
+
+    def _resolve_workspace_path(self, path_text: str) -> Path:
+        path = Path(path_text)
+        resolved = path if path.is_absolute() else self._workspace_root / path
+        resolved = resolved.resolve()
+        try:
+            resolved.relative_to(self._workspace_root)
+        except ValueError as exc:
+            raise ApplyPatchError(
+                "patch rejected: writing outside of the project; rejected by user approval settings"
+            ) from exc
+        return resolved
+
+    def _join_lines(self, lines: list[str]) -> str:
+        if not lines:
+            return ""
+        return "\n".join(lines) + "\n"

pycodex/tools/base_tool.py

@@ -0,0 +1,168 @@
+"""Shared tool abstractions for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the registry/handler layer behind Rust Codex tool routing
+  rather than one single end-user tool.
+
+Expected behavior:
+- `BaseTool` defines the contract every local Python tool must implement.
+- `ToolRegistry` stores concrete tool instances, exposes tool specs to the
+  model, and dispatches `ToolCall` executions back into `ToolResult`s.
+"""
+
+from __future__ import annotations
+
+import inspect
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import lru_cache
+import json
+from pathlib import Path
+
+from ..protocol import ConversationItem, JSONDict, JSONValue, ToolCall, ToolResult, ToolSpec
+
+EXEC_TOOLS_SNAPSHOT_PATH = (
+    Path(__file__).resolve().parent.parent / "prompts" / "exec_tools.json"
+)
+
+
+@lru_cache(maxsize=1)
+def _load_exec_tool_payloads() -> dict[str, JSONDict]:
+    payloads: dict[str, JSONDict] = {}
+    for payload in json.loads(EXEC_TOOLS_SNAPSHOT_PATH.read_text()):
+        if not isinstance(payload, dict):
+            continue
+        name = payload.get("name")
+        if isinstance(name, str):
+            payloads[name] = payload
+            continue
+        if payload.get("type") == "web_search":
+            payloads["web_search"] = payload
+    return payloads
+
+
+@dataclass(frozen=True, slots=True)
+class ToolContext:
+    turn_id: str
+    history: tuple[ConversationItem, ...]
+    collaboration_mode: str = "default"
+
+
+class StructuredToolOutput:
+    def __init__(
+        self,
+        output: JSONValue,
+        content_items: tuple[JSONDict, ...] | list[JSONDict] | None = None,
+        success: 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: JSONDict | None = None
+    tool_type: str = "function"
+    format: JSONDict | None = None
+    options: JSONDict | None = None
+    output_schema: JSONDict | None = None
+    supports_parallel: bool = True
+
+    def spec(self) -> ToolSpec:
+        return ToolSpec(
+            name=self.name,
+            description=self.description,
+            input_schema=self.input_schema,
+            tool_type=self.tool_type,
+            format=self.format,
+            options=self.options,
+            output_schema=self.output_schema,
+            supports_parallel=self.supports_parallel,
+            raw_payload=self.raw_payload(),
+        )
+
+    def serialize(self) -> JSONDict:
+        return self.spec().serialize()
+
+    def raw_payload(self) -> JSONDict | None:
+        return _load_exec_tool_payloads().get(self.name)
+
+    @abstractmethod
+    async def run(self, context: ToolContext, args: JSONValue) -> JSONValue:
+        raise NotImplementedError
+
+
+class ToolRegistry:
+    def __init__(self) -> None:
+        self._tools: dict[str, BaseTool] = {}
+
+    def register(self, tool: BaseTool) -> None:
+        self._tools[tool.name] = tool
+
+    def model_visible_specs(self) -> list[ToolSpec]:
+        return [tool.spec() for tool in self._tools.values()]
+
+    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:
+        tool = self._tools.get(call.name)
+        if tool is None:
+            return ToolResult(
+                call_id=call.call_id,
+                name=call.name,
+                output={"error": f"unknown tool: {call.name}"},
+                is_error=True,
+                tool_type=call.tool_type,
+            )
+
+        try:
+            maybe_result = tool.run(context, call.arguments)
+            if inspect.isawaitable(maybe_result):
+                output = await maybe_result
+            else:
+                output = maybe_result
+            if isinstance(output, StructuredToolOutput):
+                return ToolResult(
+                    call_id=call.call_id,
+                    name=call.name,
+                    output=output.output,
+                    content_items=output.content_items,
+                    success=output.success,
+                    tool_type=call.tool_type,
+                )
+            return ToolResult(
+                call_id=call.call_id,
+                name=call.name,
+                output=output,
+                tool_type=call.tool_type,
+            )
+        except Exception as exc:  # pragma: no cover - defensive wrapper
+            return ToolResult(
+                call_id=call.call_id,
+                name=call.name,
+                output={"error": f"{type(exc).__name__}: {exc}"},
+                is_error=True,
+                tool_type=call.tool_type,
+            )
+
+    def __contains__(self, tool_name: str) -> bool:
+        return tool_name in self._tools
+
+    def __len__(self) -> int:
+        return len(self._tools)
+
+    def names(self) -> tuple[str, ...]:
+        return tuple(self._tools)
+
+    def get_tool(self, tool_name: str) -> BaseTool | None:
+        return self._tools.get(tool_name)
+
+    def tools(self) -> tuple[BaseTool, ...]:
+        return tuple(self._tools.values())
+
+
+Registry = ToolRegistry

pycodex/tools/close_agent_tool.py

@@ -0,0 +1,55 @@
+"""`close_agent` tool for the Python Codex prototype.
+
+Original Codex mapping:
+- Corresponds to the original Codex `close_agent` collaboration tool.
+
+Expected behavior:
+- Shut down a spawned agent when it is no longer needed.
+- Return the agent status observed at close time.
+"""
+
+from __future__ import annotations
+
+from ..protocol import JSONDict, JSONValue
+from ..runtime_services import SubAgentManager
+from .agent_tool_schemas import AGENT_STATUS_SCHEMA
+from .base_tool import BaseTool, ToolContext
+
+CLOSE_AGENT_OUTPUT_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "status": AGENT_STATUS_SCHEMA,
+    },
+    "required": ["status"],
+    "additionalProperties": False,
+}
+
+
+class CloseAgentTool(BaseTool):
+    name = "close_agent"
+    description = (
+        "Close an agent when it is no longer needed and return its status."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "id": {
+                "type": "string",
+                "description": "Agent id to close (from spawn_agent).",
+            }
+        },
+        "required": ["id"],
+        "additionalProperties": False,
+    }
+    output_schema = CLOSE_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:
+        del context
+        agent_id = str(args.get("id", "")).strip()
+        if not agent_id:
+            return "Error: `id` is required."
+        return await self._subagent_manager.close_agent(agent_id)