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

pycodex/portable_server.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from urllib.parse import unquote, urlparse
+
+from .portable import (
+    DEFAULT_STORAGE_SERVER,
+    HEALTHCHECK_PATH,
+    STORAGE_API_PREFIX,
+    _call_id_from_payload,
+)
+
+
+class CodexStorageServer:
+    def __init__(
+        self,
+        root: str | Path,
+        host: str = "127.0.0.1",
+        port: int = 5577,
+    ) -> None:
+        self._root = Path(root).resolve()
+        self._root.mkdir(parents=True, exist_ok=True)
+        self._objects_dir = self._root / "objects"
+        self._objects_dir.mkdir(parents=True, exist_ok=True)
+        self._server = ThreadingHTTPServer((host, port), self._build_handler())
+        self._thread: threading.Thread | None = None
+
+    @property
+    def host(self) -> str:
+        return str(self._server.server_address[0])
+
+    @property
+    def port(self) -> int:
+        return int(self._server.server_address[1])
+
+    @property
+    def server_address(self) -> str:
+        return f"{self.host}:{self.port}"
+
+    @property
+    def base_url(self) -> str:
+        return f"http://{self.server_address}{STORAGE_API_PREFIX}"
+
+    @property
+    def root(self) -> Path:
+        return self._root
+
+    def start(self) -> None:
+        if self._thread is not None:
+            return
+        self._thread = threading.Thread(
+            target=self._server.serve_forever,
+            name="pycodex-storage-server",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._server.shutdown()
+        self._server.server_close()
+        if self._thread is not None:
+            self._thread.join(timeout=5.0)
+            self._thread = None
+
+    def _build_handler(self):
+        server = self
+
+        class Handler(BaseHTTPRequestHandler):
+            def do_GET(self) -> None:  # noqa: N802
+                path = urlparse(self.path).path
+                if path == HEALTHCHECK_PATH:
+                    self._send_json(200, {"ok": True})
+                    return
+                if not path.startswith(f"{STORAGE_API_PREFIX}/call/"):
+                    self._send_json(404, {"error": "not found"})
+                    return
+                call_id = unquote(path[len(f"{STORAGE_API_PREFIX}/call/") :]).strip()
+                if not call_id:
+                    self._send_json(400, {"error": "missing call_id"})
+                    return
+                object_path = server._object_path(call_id)
+                if not object_path.is_file():
+                    self._send_json(404, {"error": "not found"})
+                    return
+                payload = object_path.read_bytes()
+                print(
+                    "[server] call: "
+                    f"client={self.client_address[0]} call_id={call_id} path={object_path}",
+                    flush=True,
+                )
+                self.send_response(200)
+                self.send_header("Content-Type", "application/octet-stream")
+                self.send_header("Content-Length", str(len(payload)))
+                self.send_header("X-Pycodex-Sha256", hashlib.sha256(payload).hexdigest())
+                self.send_header("X-Pycodex-Call-Id", call_id)
+                self.end_headers()
+                self.wfile.write(payload)
+
+            def do_POST(self) -> None:  # noqa: N802
+                path = urlparse(self.path).path
+                if path != f"{STORAGE_API_PREFIX}/put":
+                    self._send_json(404, {"error": "not found"})
+                    return
+                content_length = int(self.headers.get("Content-Length", "0") or "0")
+                if content_length <= 0:
+                    self._send_json(400, {"error": "empty body"})
+                    return
+                payload = self.rfile.read(content_length)
+                sha256 = hashlib.sha256(payload).hexdigest()
+                expected_sha256 = self.headers.get("X-Pycodex-Sha256", "").strip().lower()
+                if expected_sha256 and expected_sha256 != sha256:
+                    self._send_json(400, {"error": "checksum mismatch"})
+                    return
+                call_id = _call_id_from_payload(payload)
+                object_path = server._object_path(call_id)
+                if not object_path.is_file():
+                    object_path.write_bytes(payload)
+                    status = "stored"
+                else:
+                    status = "reused"
+                print(
+                    "[server] put: "
+                    f"client={self.client_address[0]} "
+                    f"call_id={call_id} status={status} path={object_path}",
+                    flush=True,
+                )
+                host_header = self.headers.get("Host", server.server_address).strip() or server.server_address
+                self._send_json(
+                    200,
+                    {
+                        "call_id": call_id,
+                        "call": f"{call_id}@{host_header}",
+                    },
+                )
+
+            def log_message(self, _format: str, *_args) -> None:
+                return
+
+            def _send_json(self, status: int, payload: dict[str, object]) -> None:
+                body = json.dumps(payload).encode("utf-8")
+                self.send_response(status)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+        return Handler
+
+    def _object_path(self, call_id: str) -> Path:
+        return self._objects_dir / f"{call_id}.bin"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m pycodex.portable_server",
+        description="Run a pycodex remote storage service for --put/--call testing.",
+    )
+    parser.add_argument(
+        "--root",
+        default=str(Path(".tmp") / "pycodex_storage"),
+        help="Directory used to store uploaded encrypted bundles.",
+    )
+    parser.add_argument(
+        "--host",
+        default=DEFAULT_STORAGE_SERVER.split(":", 1)[0],
+        help="Host interface to bind.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(DEFAULT_STORAGE_SERVER.split(":", 1)[1]),
+        help="Port to bind.",
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    server = CodexStorageServer(args.root, host=args.host, port=args.port)
+    server.start()
+    print(f"storage server listening on {server.base_url}", flush=True)
+    print(f"storage root: {server.root}", flush=True)
+    print(f"put current home: pycodex --put @{server.server_address}", flush=True)
+    print(
+        f"put custom home: pycodex --put /data/.codex/@{server.server_address}",
+        flush=True,
+    )
+    try:
+        if server._thread is not None:
+            server._thread.join()
+    except KeyboardInterrupt:
+        return 130
+    finally:
+        server.stop()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

pycodex/runtime.py

@@ -3,12 +3,15 @@ from __future__ import annotations
 import asyncio
 from collections import deque
 from dataclasses import dataclass
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 from .agent import AgentLoop, EventHandler, NOOP_EVENT_HANDLER, TurnInterrupted
 from .protocol import AgentEvent, Operation, ShutdownOp, Submission, TurnResult, UserTurnOp
 from .utils import uuid7_string
 
+if TYPE_CHECKING:
+    from .runtime_services import RuntimeEnvironment
+
 
 @dataclass(slots=True)
 class _QueuedSubmission:
@@ -20,8 +23,9 @@ class _QueuedSubmission:
 class AgentRuntime:
     """Thin outer queue that mirrors the Rust `submission_loop` shape."""
 
-    def __init__(self, agent_loop: AgentLoop) -> None:
+    def __init__(self, agent_loop: AgentLoop, runtime_environment: RuntimeEnvironment | None = None) -> None:
         self._agent_loop = agent_loop
+        self.runtime_environment = runtime_environment
         self._enqueue_queue: deque[_QueuedSubmission] = deque()
         self._steer_queue: deque[_QueuedSubmission] = deque()
         self._queue_lock = asyncio.Lock()

pycodex/runtime_services.py

@@ -346,6 +346,9 @@ class SubAgentManager:
             managed.state = "completed"
         finally:
             managed.pending_submission_ids.discard(submission_id)
+            if managed.pending_submission_ids and managed.error_message is None:
+                managed.completed_message = None
+                managed.state = "running"
             async with self._condition:
                 self._condition.notify_all()
 
@@ -397,11 +400,12 @@ class RuntimeEnvironment:
         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)
 
+def create_runtime_environment() -> RuntimeEnvironment:
+    return RuntimeEnvironment()
 
-_RUNTIME_ENV = RuntimeEnvironment()
+
+_RUNTIME_ENV = create_runtime_environment()
 
 
 def get_runtime_environment() -> RuntimeEnvironment:

pycodex/tools/exec_tool.py

@@ -17,7 +17,7 @@ from ..protocol import JSONValue
 from .base_tool import BaseTool, ToolContext
 from .code_mode_manager import CodeModeManager
 
-EXEC_FREEFORM_GRAMMAR = """start: pragma_source | plain_source
+EXEC_FREEFORM_GRAMMAR = r"""start: pragma_source | plain_source
 pragma_source: PRAGMA_LINE NEWLINE SOURCE
 plain_source: SOURCE
 

pycodex/tools/unified_exec_manager.py

@@ -184,6 +184,9 @@ class _HeadTailBuffer:
         self.tail.clear()
         return combined
 
+    def has_data(self) -> bool:
+        return bool(self.head or self.tail)
+
 
 @dataclass(slots=True)
 class UnifiedExecSession:
@@ -194,6 +197,7 @@ class UnifiedExecSession:
     tty: bool
     unread_output: _HeadTailBuffer = field(default_factory=_HeadTailBuffer)
     reader_task: asyncio.Task | None = None
+    output_event: asyncio.Event = field(default_factory=asyncio.Event)
 
 
 class UnifiedExecManager:
@@ -294,11 +298,22 @@ class UnifiedExecManager:
         if session is None:
             return f"Error: session_id {session_id} is not running."
 
-        start_wait = asyncio.get_running_loop().time()
+        loop = asyncio.get_running_loop()
+        start_wait = loop.time()
         try:
             await asyncio.wait_for(session.process.wait(), timeout=yield_time_ms / 1000.0)
         except asyncio.TimeoutError:
-            pass
+            remaining_seconds = (yield_time_ms / 1000.0) - (loop.time() - start_wait)
+            if (
+                session.process.returncode is None
+                and not session.unread_output.has_data()
+                and remaining_seconds > 0
+            ):
+                session.output_event.clear()
+                try:
+                    await asyncio.wait_for(session.output_event.wait(), timeout=remaining_seconds)
+                except asyncio.TimeoutError:
+                    pass
 
         if session.reader_task is not None and session.process.returncode is not None:
             await session.reader_task
@@ -345,6 +360,8 @@ class UnifiedExecManager:
             if not chunk:
                 break
             session.unread_output.push_chunk(chunk)
+            session.output_event.set()
+        session.output_event.set()
 
     def _resolve_workdir(self, workdir: str | None) -> Path:
         if not workdir:

pycodex/utils/get_env.py

@@ -83,10 +83,15 @@ def get_package_version() -> str:
     detected = _detect_upstream_codex_version()
     if detected is not None:
         return detected
-    try:
-        return importlib.metadata.version("pycodex")
-    except importlib.metadata.PackageNotFoundError:
-        return "0.1.0"
+    for distribution_name in ("python-codex", "pycodex"):
+        try:
+            return importlib.metadata.version(distribution_name)
+        except importlib.metadata.PackageNotFoundError:
+            continue
+    local_version = _read_local_package_version()
+    if local_version is not None:
+        return local_version
+    return "0.1.0"
 
 
 def get_os_info() -> tuple[str, str]:
@@ -178,6 +183,20 @@ def _normalize_os_version(version: str) -> str:
     return version
 
 
+def _read_local_package_version() -> str | None:
+    pyproject_path = Path(__file__).resolve().parents[2] / "pyproject.toml"
+    if not pyproject_path.is_file():
+        return None
+    match = re.search(
+        r'^\s*version\s*=\s*"([^"]+)"\s*$',
+        pyproject_path.read_text(encoding="utf-8"),
+        flags=re.MULTILINE,
+    )
+    if match is None:
+        return None
+    return match.group(1).strip() or None
+
+
 def _tmux_display_message(fmt: str) -> str | None:
     try:
         output = subprocess.run(