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

pycodex/runtime.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import asyncio
+from collections import deque
+from dataclasses import dataclass
+from typing import Literal
+
+from .agent import AgentLoop, EventHandler, NOOP_EVENT_HANDLER, TurnInterrupted
+from .protocol import AgentEvent, Operation, ShutdownOp, Submission, TurnResult, UserTurnOp
+from .utils import uuid7_string
+
+
+@dataclass(slots=True)
+class _QueuedSubmission:
+    submission: Submission
+    turn_id: str
+    futures: list[asyncio.Future[TurnResult | None]]
+
+
+class AgentRuntime:
+    """Thin outer queue that mirrors the Rust `submission_loop` shape."""
+
+    def __init__(self, agent_loop: AgentLoop) -> None:
+        self._agent_loop = agent_loop
+        self._enqueue_queue: deque[_QueuedSubmission] = deque()
+        self._steer_queue: deque[_QueuedSubmission] = deque()
+        self._queue_lock = asyncio.Lock()
+        self._queue_event = asyncio.Event()
+        self._current_submission: _QueuedSubmission | None = None
+        self._current_task: asyncio.Task[TurnResult] | None = None
+        self._event_handler = NOOP_EVENT_HANDLER
+        self._agent_loop.set_event_handler(self._handle_agent_event)
+
+    def set_event_handler(self, event_handler: EventHandler = NOOP_EVENT_HANDLER) -> None:
+        self._event_handler = event_handler
+
+    async def submit_user_turn(self, text: str) -> TurnResult:
+        _submission_id, future = await self.enqueue_user_turn(text, queue="enqueue")
+        result = await future
+        assert result is not None
+        return result
+
+    async def enqueue_user_turn(
+        self,
+        text: str,
+        queue: Literal["enqueue", "steer"] = "enqueue",
+    ) -> tuple[str, asyncio.Future[TurnResult | None]]:
+        future: asyncio.Future[TurnResult | None] = asyncio.get_running_loop().create_future()
+        return await self._enqueue_user_turn_to_queue(
+            text,
+            future,
+            queue=queue,
+        )
+
+    async def shutdown(self) -> None:
+        submission = Submission(id=uuid7_string(), op=ShutdownOp())
+        future: asyncio.Future[TurnResult | None] = asyncio.get_running_loop().create_future()
+        self._enqueue_queue.append(
+            _QueuedSubmission(
+                submission=submission,
+                turn_id=submission.id,
+                futures=[future],
+            )
+        )
+        self._queue_event.set()
+        await future
+
+    async def run_forever(self) -> None:
+        while True:
+            queued = await self._next_submission()
+            submission = queued.submission
+            self._current_submission = queued
+            try:
+                if isinstance(submission.op, UserTurnOp):
+                    self._current_task = asyncio.create_task(
+                        self._agent_loop.run_turn(
+                            list(submission.op.texts),
+                            turn_id=queued.turn_id,
+                        )
+                    )
+                    try:
+                        result = await self._current_task
+                    except TurnInterrupted:
+                        self._finish_submission_exception(
+                            queued,
+                            RuntimeError("submission interrupted"),
+                        )
+                        continue
+                    except asyncio.CancelledError:
+                        self._finish_submission_exception(
+                            queued,
+                            RuntimeError("submission interrupted"),
+                        )
+                        continue
+                    self._finish_submission_result(queued, result)
+                    continue
+
+                if isinstance(submission.op, ShutdownOp):
+                    self._finish_submission_result(queued, None)
+                    break
+
+                self._finish_submission_exception(
+                    queued,
+                    RuntimeError(f"unsupported operation: {type(submission.op).__name__}"),
+                )
+            except Exception as exc:  # pragma: no cover - defensive wrapper
+                self._finish_submission_exception(queued, exc)
+            finally:
+                self._current_task = None
+                self._current_submission = None
+
+    @staticmethod
+    def operation_name(op: Operation) -> str:
+        if isinstance(op, UserTurnOp):
+            return "user_turn"
+        if isinstance(op, ShutdownOp):
+            return "shutdown"
+        return type(op).__name__
+
+    async def _enqueue_user_turn_to_queue(
+        self,
+        text: str,
+        future: asyncio.Future[TurnResult | None],
+        queue: Literal["enqueue", "steer"],
+    ) -> tuple[str, asyncio.Future[TurnResult | None]]:
+        if queue == "steer" and self._has_active_turn():
+            self._agent_loop.interrupt_asap = True
+
+        async with self._queue_lock:
+            if queue == "steer" and self._steer_queue:
+                queued = self._steer_queue[-1]
+                queued.submission.op.texts.append(text)
+                queued.futures.append(future)
+                return queued.submission.id, future
+
+            submission = Submission(id=uuid7_string(), op=UserTurnOp(texts=[text]))
+            current = self._current_submission if self._has_active_turn() else None
+            turn_id = (
+                current.turn_id if queue == "steer" and current is not None else submission.id
+            )
+            queued = _QueuedSubmission(
+                submission=submission,
+                turn_id=turn_id,
+                futures=[future],
+            )
+            if queue == "steer":
+                self._steer_queue.append(queued)
+            else:
+                self._enqueue_queue.append(queued)
+            self._queue_event.set()
+            return submission.id, future
+
+    async def _next_submission(self) -> _QueuedSubmission:
+        while True:
+            async with self._queue_lock:
+                queued: _QueuedSubmission | None = None
+                if self._steer_queue:
+                    queued = self._steer_queue.popleft()
+                elif self._enqueue_queue:
+                    queued = self._enqueue_queue.popleft()
+                if queued is not None:
+                    if not self._steer_queue and not self._enqueue_queue:
+                        self._queue_event.clear()
+                    return queued
+                self._queue_event.clear()
+            await self._queue_event.wait()
+
+    @staticmethod
+    def _finish_submission_result(
+        queued: _QueuedSubmission,
+        result: TurnResult | None,
+    ) -> None:
+        for future in queued.futures:
+            if not future.done():
+                future.set_result(result)
+
+    @staticmethod
+    def _finish_submission_exception(
+        queued: _QueuedSubmission,
+        exc: Exception,
+    ) -> None:
+        for future in queued.futures:
+            if not future.done():
+                future.set_exception(exc)
+
+    def _has_active_turn(self) -> bool:
+        current_task = self._current_task
+        return current_task is not None and not current_task.done()
+
+    def _handle_agent_event(self, event: AgentEvent) -> None:
+        queued = self._current_submission
+        if queued is None:
+            self._event_handler(event)
+            return
+        payload = dict(event.payload)
+        payload.setdefault("submission_id", queued.submission.id)
+        payload.setdefault("turn_id", queued.turn_id)
+        self._event_handler(
+            AgentEvent(kind=event.kind, turn_id=event.turn_id, payload=payload)
+        )

pycodex/runtime_services.py

@@ -0,0 +1,408 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import random
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Literal
+
+from .protocol import ConversationItem, TurnResult
+from .utils import uuid7_string
+
+if TYPE_CHECKING:
+    from .runtime import AgentRuntime
+
+PlanStatus = Literal["pending", "in_progress", "completed"]
+PlanListener = Callable[[dict[str, object]], None]
+RuntimeBuilder = Callable[
+    [str | None, str | None, tuple[ConversationItem, ...], str],
+    "AgentRuntime",
+]
+AsyncJSONHandler = Callable[[dict[str, object]], Awaitable[dict[str, object] | None]]
+
+DEFAULT_AGENT_NICKNAME_CANDIDATES = (
+    "Bacon",
+    "Descartes",
+    "Pascal",
+    "Fermat",
+    "Huygens",
+    "Leibniz",
+    "Newton",
+    "Halley",
+    "Euler",
+    "Lagrange",
+    "Laplace",
+    "Volta",
+    "Gauss",
+    "Ampere",
+    "Faraday",
+    "Darwin",
+    "Lovelace",
+    "Boole",
+    "Pasteur",
+    "Maxwell",
+    "Mendel",
+    "Curie",
+    "Planck",
+    "Tesla",
+    "Poincare",
+    "Noether",
+    "Hilbert",
+    "Einstein",
+    "Raman",
+    "Bohr",
+    "Turing",
+    "Hubble",
+    "Feynman",
+    "Franklin",
+    "McClintock",
+    "Meitner",
+    "Herschel",
+    "Linnaeus",
+    "Wegener",
+    "Chandrasekhar",
+    "Sagan",
+    "Goodall",
+    "Carson",
+    "Carver",
+    "Socrates",
+    "Plato",
+    "Aristotle",
+    "Epicurus",
+    "Cicero",
+    "Confucius",
+    "Mencius",
+    "Zeno",
+    "Locke",
+    "Hume",
+    "Kant",
+    "Hegel",
+    "Kierkegaard",
+    "Mill",
+    "Nietzsche",
+    "Peirce",
+    "James",
+    "Dewey",
+    "Russell",
+    "Popper",
+    "Sartre",
+    "Beauvoir",
+    "Arendt",
+    "Rawls",
+    "Singer",
+    "Anscombe",
+    "Parfit",
+    "Kuhn",
+    "Boyle",
+    "Hooke",
+    "Harvey",
+    "Dalton",
+    "Helmholtz",
+    "Gibbs",
+    "Lorentz",
+    "Schrodinger",
+    "Heisenberg",
+    "Pauli",
+    "Dirac",
+    "Bernoulli",
+    "Godel",
+    "Nash",
+    "Banach",
+    "Ramanujan",
+    "Erdos",
+)
+
+
+@dataclass(frozen=True, slots=True)
+class PlanItem:
+    step: str
+    status: PlanStatus
+
+
+class PlanStore:
+    def __init__(self) -> None:
+        self._explanation: str | None = None
+        self._plan: tuple[PlanItem, ...] = ()
+        self._listener: PlanListener = lambda _payload: None
+
+    def set_listener(self, listener: PlanListener | None) -> None:
+        self._listener = listener or (lambda _payload: None)
+
+    def update(self, explanation: str | None, plan: tuple[PlanItem, ...]) -> None:
+        in_progress = sum(1 for item in plan if item.status == "in_progress")
+        if in_progress > 1:
+            raise ValueError("at most one plan step can be in_progress")
+        self._explanation = explanation
+        self._plan = plan
+        self._listener(
+            {
+                "explanation": explanation,
+                "plan": [
+                    {"step": item.step, "status": item.status}
+                    for item in self._plan
+                ],
+            }
+        )
+
+    def snapshot(self) -> dict[str, object]:
+        return {
+            "explanation": self._explanation,
+            "plan": [
+                {"step": item.step, "status": item.status}
+                for item in self._plan
+            ],
+        }
+
+
+class RequestUserInputManager:
+    def __init__(self) -> None:
+        self._handler: AsyncJSONHandler | None = None
+
+    def set_handler(self, handler: AsyncJSONHandler | None) -> None:
+        self._handler = handler
+
+    async def request(self, payload: dict[str, object]) -> dict[str, object] | None:
+        handler = self._handler
+        if handler is None:
+            return None
+        return await handler(payload)
+
+
+class RequestPermissionsManager:
+    def __init__(self) -> None:
+        self._handler: AsyncJSONHandler | None = None
+
+    def set_handler(self, handler: AsyncJSONHandler | None) -> None:
+        self._handler = handler
+
+    async def request(self, payload: dict[str, object]) -> dict[str, object] | None:
+        handler = self._handler
+        if handler is None:
+            return None
+        return await handler(payload)
+
+
+@dataclass(slots=True)
+class ManagedAgent:
+    agent_id: str
+    runtime: "AgentRuntime"
+    worker_task: asyncio.Task[None]
+    nickname: str | None = None
+    state: str = "pending_init"
+    completed_message: str | None = None
+    error_message: str | None = None
+    pending_submission_ids: set[str] = field(default_factory=set)
+
+
+class SubAgentManager:
+    def __init__(self) -> None:
+        self._runtime_builder: RuntimeBuilder | None = None
+        self._agents: dict[str, ManagedAgent] = {}
+        self._condition = asyncio.Condition()
+        self._available_nicknames: list[str] = []
+        self._nickname_random = random.Random()
+
+    def set_runtime_builder(self, builder: RuntimeBuilder | None) -> None:
+        self._runtime_builder = builder
+
+    async def spawn_agent(
+        self,
+        message: str | None,
+        items: list[dict[str, object]] | None,
+        agent_type: str | None,
+        fork_context: bool,
+        model: str | None,
+        reasoning_effort: str | None,
+        history: tuple[ConversationItem, ...],
+    ) -> dict[str, object]:
+        builder = self._runtime_builder
+        if builder is None:
+            raise RuntimeError("spawn_agent is unavailable before runtime initialization")
+
+        initial_history = history if fork_context else ()
+        agent_id = uuid7_string()
+        runtime = builder(model, reasoning_effort, initial_history, agent_id)
+        worker_task = asyncio.create_task(runtime.run_forever())
+        nickname = self._next_nickname()
+        managed = ManagedAgent(
+            agent_id=agent_id,
+            runtime=runtime,
+            worker_task=worker_task,
+            nickname=nickname,
+        )
+        async with self._condition:
+            self._agents[agent_id] = managed
+            self._condition.notify_all()
+
+        initial_prompt = self._compose_prompt(message, items)
+        if initial_prompt:
+            await self.send_input(agent_id, initial_prompt, interrupt=False)
+
+        return {
+            "agent_id": agent_id,
+            "nickname": nickname,
+        }
+
+    async def send_input(
+        self,
+        agent_id: str,
+        prompt_text: str,
+        interrupt: bool,
+    ) -> dict[str, object]:
+        managed = self._agents.get(agent_id)
+        if managed is None:
+            raise RuntimeError(f"unknown agent: {agent_id}")
+        if managed.state == "shutdown":
+            raise RuntimeError(f"agent is shutdown: {agent_id}")
+
+        submission_id, future = await managed.runtime.enqueue_user_turn(
+            prompt_text,
+            queue="steer" if interrupt else "enqueue",
+        )
+        managed.state = "running"
+        managed.completed_message = None
+        managed.error_message = None
+        managed.pending_submission_ids.add(submission_id)
+        asyncio.create_task(self._track_submission(managed, submission_id, future))
+        async with self._condition:
+            self._condition.notify_all()
+        return {"submission_id": submission_id}
+
+    async def resume_agent(self, agent_id: str) -> dict[str, object]:
+        managed = self._agents.get(agent_id)
+        if managed is None:
+            return {"status": "not_found"}
+        if managed.worker_task.done():
+            managed.worker_task = asyncio.create_task(managed.runtime.run_forever())
+            managed.state = "pending_init"
+            managed.completed_message = None
+            managed.error_message = None
+        async with self._condition:
+            self._condition.notify_all()
+        return {"status": self._status_payload(managed)}
+
+    async def close_agent(self, agent_id: str) -> dict[str, object]:
+        managed = self._agents.get(agent_id)
+        if managed is None:
+            return {"status": "not_found"}
+        previous_status = self._status_payload(managed)
+        if not managed.worker_task.done():
+            managed.runtime._agent_loop.interrupt_asap = True
+            await managed.runtime.shutdown()
+            await managed.worker_task
+        managed.state = "shutdown"
+        managed.pending_submission_ids.clear()
+        async with self._condition:
+            self._condition.notify_all()
+        return {"status": previous_status}
+
+    def _next_nickname(self) -> str:
+        if not self._available_nicknames:
+            self._available_nicknames = list(DEFAULT_AGENT_NICKNAME_CANDIDATES)
+            self._nickname_random.shuffle(self._available_nicknames)
+        return self._available_nicknames.pop()
+
+    async def wait_agents(
+        self,
+        agent_ids: list[str],
+        timeout_ms: int = 30_000,
+    ) -> dict[str, object]:
+        timeout_seconds = max(timeout_ms, 1) / 1000.0
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + timeout_seconds
+
+        while True:
+            snapshot = {
+                agent_id: self._status_payload(self._agents.get(agent_id))
+                for agent_id in agent_ids
+            }
+            if any(self._is_final_status(status) for status in snapshot.values()):
+                return {"status": snapshot, "timed_out": False}
+
+            remaining = deadline - loop.time()
+            if remaining <= 0:
+                return {"status": {}, "timed_out": True}
+
+            async with self._condition:
+                try:
+                    await asyncio.wait_for(self._condition.wait(), timeout=remaining)
+                except asyncio.TimeoutError:
+                    return {"status": {}, "timed_out": True}
+
+    async def _track_submission(
+        self,
+        managed: ManagedAgent,
+        submission_id: str,
+        future: asyncio.Future[TurnResult | None],
+    ) -> None:
+        try:
+            result = await future
+        except Exception as exc:  # pragma: no cover - background safety
+            managed.error_message = f"{type(exc).__name__}: {exc}"
+            managed.state = "errored"
+        else:
+            managed.completed_message = None if result is None else result.output_text
+            managed.state = "completed"
+        finally:
+            managed.pending_submission_ids.discard(submission_id)
+            async with self._condition:
+                self._condition.notify_all()
+
+    def _compose_prompt(
+        self,
+        message: str | None,
+        items: list[dict[str, object]] | None,
+    ) -> str:
+        parts: list[str] = []
+        if message:
+            parts.append(message.strip())
+        for item in items or []:
+            item_type = str(item.get("type", ""))
+            if item_type == "text":
+                text = str(item.get("text", "")).strip()
+                if text:
+                    parts.append(text)
+            elif item_type == "image":
+                image_url = str(item.get("image_url", "")).strip()
+                if image_url:
+                    parts.append(f"[image] {image_url}")
+            else:
+                parts.append(json.dumps(item, ensure_ascii=False))
+        return "\n\n".join(part for part in parts if part)
+
+    def _status_payload(self, managed: ManagedAgent | None) -> object:
+        if managed is None:
+            return "not_found"
+        if managed.error_message is not None:
+            return {"errored": managed.error_message}
+        if managed.state == "completed":
+            return {"completed": managed.completed_message}
+        if managed.state in {"pending_init", "running", "shutdown"}:
+            return managed.state
+        return managed.state
+
+    def _is_final_status(self, status: object) -> bool:
+        if isinstance(status, str):
+            return status in {"shutdown", "not_found"}
+        if isinstance(status, dict):
+            return "completed" in status or "errored" in status
+        return False
+
+
+class RuntimeEnvironment:
+    def __init__(self) -> None:
+        self.plan_store = PlanStore()
+        self.subagent_manager = SubAgentManager()
+        self.request_user_input_manager = RequestUserInputManager()
+        self.request_permissions_manager = RequestPermissionsManager()
+
+    def configure_runtime_builder(self, builder: RuntimeBuilder | None) -> None:
+        self.subagent_manager.set_runtime_builder(builder)
+
+
+_RUNTIME_ENV = RuntimeEnvironment()
+
+
+def get_runtime_environment() -> RuntimeEnvironment:
+    return _RUNTIME_ENV

pycodex/tools/__init__.py

@@ -0,0 +1,58 @@
+"""Tool package for the Python Codex prototype.
+
+This package groups the local tool abstractions and concrete tool
+implementations that back `pycodex`.
+"""
+
+from .base_tool import BaseTool, Registry, ToolContext, ToolRegistry
+from .apply_patch_tool import ApplyPatchTool
+from .close_agent_tool import CloseAgentTool
+from .code_mode_manager import CodeModeManager
+from .exec_command_tool import ExecCommandTool
+from .exec_tool import ExecTool
+from .grep_files_tool import GrepFilesTool
+from .list_dir_tool import ListDirTool
+from .read_file_tool import ReadFileTool
+from .request_permissions_tool import RequestPermissionsTool
+from .request_user_input_tool import RequestUserInputTool
+from .resume_agent_tool import ResumeAgentTool
+from .send_input_tool import SendInputTool
+from .shell_command_tool import ShellCommandTool
+from .shell_tool import ShellTool
+from .spawn_agent_tool import SpawnAgentTool
+from .unified_exec_manager import UnifiedExecManager
+from .update_plan_tool import UpdatePlanTool
+from .view_image_tool import ViewImageTool
+from .wait_agent_tool import WaitAgentTool
+from .wait_tool import WaitTool
+from .web_search_tool import WebSearchTool
+from .write_stdin_tool import WriteStdinTool
+
+__all__ = [
+    "ApplyPatchTool",
+    "BaseTool",
+    "CloseAgentTool",
+    "CodeModeManager",
+    "ExecTool",
+    "ExecCommandTool",
+    "GrepFilesTool",
+    "ListDirTool",
+    "ReadFileTool",
+    "Registry",
+    "RequestPermissionsTool",
+    "RequestUserInputTool",
+    "ResumeAgentTool",
+    "SendInputTool",
+    "ShellCommandTool",
+    "ShellTool",
+    "SpawnAgentTool",
+    "ToolContext",
+    "ToolRegistry",
+    "UnifiedExecManager",
+    "UpdatePlanTool",
+    "ViewImageTool",
+    "WaitAgentTool",
+    "WaitTool",
+    "WebSearchTool",
+    "WriteStdinTool",
+]