agentlink-cli 0.1.0__py3-none-any.whl

connector/codex/rpc.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import shutil
+import sys
+import time
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import Any
+
+from connector.launch import launch_target
+from connector.logging import logger
+from connector.version import connector_version
+
+NotificationHandler = Callable[[dict[str, Any]], Awaitable[None]]
+APP_SERVER_STREAM_LIMIT = 64 * 1024 * 1024
+
+
+class JsonRpcStdioClient:
+    """Line-delimited JSON-RPC client for `codex app-server --listen stdio://`."""
+
+    def __init__(self, command: list[str] | None = None) -> None:
+        self.command = command or _resolve_codex_command()
+        self.process: asyncio.subprocess.Process | None = None
+        self._start_lock = asyncio.Lock()
+        self._next_id = 1
+        self._pending: dict[int | str, asyncio.Future[dict[str, Any]]] = {}
+        self._server_request_ids: set[int | str] = set()
+        self._notification_handler: NotificationHandler | None = None
+        self._initialized = False
+
+    async def start(self, handler: NotificationHandler) -> None:
+        async with self._start_lock:
+            if self.process and self._initialized:
+                self._notification_handler = handler
+                return
+
+            self._notification_handler = handler
+            if self.process is None:
+                logger.info("starting codex app-server command={}", self.command)
+                self.process = await asyncio.create_subprocess_exec(
+                    *self.command,
+                    stdin=asyncio.subprocess.PIPE,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    limit=APP_SERVER_STREAM_LIMIT,
+                )
+                self._track_reader(asyncio.create_task(self._read_stdout(self.process)), "stdout")
+                self._track_reader(asyncio.create_task(self._read_stderr(self.process)), "stderr")
+
+            await self.request(
+                "initialize",
+                {
+                    "clientInfo": {
+                        "name": "agent-server-connector",
+                        "title": "Agent Server Connector",
+                        "version": connector_version(),
+                    },
+                    "capabilities": {
+                        "experimentalApi": True,
+                        "requestAttestation": False,
+                    },
+                },
+            )
+            await self.notify("initialized")
+            self._initialized = True
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Codex app-server is not started")
+
+        request_id = self._next_id
+        self._next_id += 1
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[dict[str, Any]] = loop.create_future()
+        self._pending[request_id] = future
+        payload: dict[str, Any] = {
+            "jsonrpc": "2.0",
+            "id": request_id,
+            "method": method,
+            "params": params or {},
+        }
+
+        started = time.perf_counter()
+        self.process.stdin.write((json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8"))
+        await self.process.stdin.drain()
+        try:
+            return await future
+        finally:
+            elapsed_ms = (time.perf_counter() - started) * 1000
+            logger.trace("codex rpc method={} id={} elapsed_ms={:.1f}", method, request_id, elapsed_ms)
+
+    async def notify(self, method: str, params: dict[str, Any] | None = None) -> None:
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Codex app-server is not started")
+
+        payload: dict[str, Any] = {"jsonrpc": "2.0", "method": method, "params": params or {}}
+        self.process.stdin.write((json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8"))
+        await self.process.stdin.drain()
+
+    async def respond(self, request_id: str | int, result: dict[str, Any] | None = None) -> None:
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Codex app-server is not started")
+
+        response_id = self._response_id_for(request_id)
+        payload: dict[str, Any] = {"jsonrpc": "2.0", "id": response_id, "result": result or {}}
+        self.process.stdin.write((json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8"))
+        await self.process.stdin.drain()
+
+    async def close(self) -> None:
+        if self.process is None:
+            return
+        self.process.terminate()
+        try:
+            await asyncio.wait_for(self.process.wait(), timeout=5)
+        except TimeoutError:
+            self.process.kill()
+            await self.process.wait()
+        finally:
+            self.process = None
+            self._initialized = False
+
+    async def _read_stdout(self, process: asyncio.subprocess.Process) -> None:
+        assert process.stdout
+        while line := await process.stdout.readline():
+            try:
+                payload = json.loads(line)
+            except json.JSONDecodeError:
+                logger.warning("codex app-server emitted non-json stdout: {}", line.decode(errors="replace").strip())
+                continue
+
+            request_id = payload.get("id")
+            if request_id in self._pending and ("result" in payload or "error" in payload):
+                future = self._pending.pop(request_id)
+                self._settle_pending_future(future, payload)
+                continue
+
+            if request_id is not None and isinstance(payload.get("method"), str):
+                self._server_request_ids.add(request_id)
+
+            if self._notification_handler is not None:
+                await self._notification_handler(payload)
+
+    async def _read_stderr(self, process: asyncio.subprocess.Process) -> None:
+        assert process.stderr
+        while line := await process.stderr.readline():
+            logger.trace("codex app-server stderr: {}", line.decode(errors="replace").rstrip())
+
+    def _track_reader(self, task: asyncio.Task[None], name: str) -> None:
+        def done(completed: asyncio.Task[None]) -> None:
+            try:
+                completed.result()
+            except asyncio.CancelledError:
+                return
+            except Exception:
+                logger.exception("codex app-server {} reader stopped unexpectedly", name)
+
+        task.add_done_callback(done)
+
+    def _settle_pending_future(self, future: asyncio.Future[dict[str, Any]], payload: dict[str, Any]) -> None:
+        if future.done():
+            logger.trace("codex rpc received response for completed request id={}", payload.get("id"))
+            return
+        if "error" in payload:
+            future.set_exception(RuntimeError(json.dumps(payload["error"], ensure_ascii=False)))
+        else:
+            future.set_result(payload.get("result") or {})
+
+    def _response_id_for(self, request_id: str | int) -> str | int:
+        if request_id in self._server_request_ids:
+            self._server_request_ids.remove(request_id)
+            return request_id
+        if isinstance(request_id, str):
+            try:
+                numeric_request_id = int(request_id)
+            except ValueError:
+                numeric_request_id = None
+            if numeric_request_id is not None and numeric_request_id in self._server_request_ids:
+                self._server_request_ids.remove(numeric_request_id)
+                logger.trace(
+                    "codex rpc coerced approval response id from string to number request_id={}",
+                    request_id,
+                )
+                return numeric_request_id
+        logger.warning("codex rpc responding to unknown server request id={}", request_id)
+        return request_id
+
+
+def _resolve_codex_bin() -> str:
+    for candidate in codex_candidate_paths():
+        if candidate["source"] == "custom":
+            return candidate["path"]
+        path = Path(candidate["path"])
+        if path.is_file():
+            return str(path)
+    return "codex"
+
+
+def _resolve_codex_command() -> list[str]:
+    path = _resolve_codex_bin()
+    return launch_target("cli", path).command(["app-server", "--listen", "stdio://"])
+
+
+def codex_candidate_paths() -> list[dict[str, str]]:
+    if sys.platform == "win32":
+        home = Path.home()
+        appdata = os.environ.get("APPDATA", str(home / "AppData" / "Roaming"))
+        candidates = [
+            {"source": "custom", "path": os.getenv("CODEX_BIN", "")},
+            *[
+                {"source": "nvm", "path": str(Path("C:/nvm4w/nodejs") / name)}
+                for name in ("codex.cmd", "codex.ps1", "codex.exe")
+            ],
+            {"source": "cli", "path": shutil.which("codex") or ""},
+            *[
+                {"source": "npm", "path": str(Path(appdata) / "npm" / name)}
+                for name in ("codex.cmd", "codex.ps1", "codex.exe")
+            ],
+            *[
+                {"source": "npm", "path": str(home / ".npm-global" / "bin" / name)}
+                for name in ("codex.cmd", "codex.ps1", "codex.exe")
+            ],
+            *[
+                {"source": "cli", "path": str(home / ".local" / "bin" / name)}
+                for name in ("codex.exe", "codex.cmd", "codex.ps1")
+            ],
+            *[
+                {"source": "scoop", "path": str(home / "scoop" / "shims" / name)}
+                for name in ("codex.exe", "codex.cmd", "codex.ps1")
+            ],
+        ]
+        seen: set[str] = set()
+        out: list[dict[str, str]] = []
+        for candidate in candidates:
+            path = candidate.get("path") or ""
+            if not path or path in seen:
+                continue
+            seen.add(path)
+            out.append(candidate)
+        return out
+
+    candidates = [
+        {"source": "custom", "path": os.getenv("CODEX_BIN", "")},
+        {"source": "app", "path": "/Applications/Codex.app/Contents/Resources/codex"},
+        {"source": "app", "path": str(Path.home() / "Applications" / "Codex.app" / "Contents" / "Resources" / "codex")},
+        {"source": "cli", "path": shutil.which("codex") or ""},
+        {"source": "cli", "path": "/opt/homebrew/bin/codex"},
+        {"source": "cli", "path": "/usr/local/bin/codex"},
+    ]
+
+    seen: set[str] = set()
+    out: list[dict[str, str]] = []
+    for candidate in candidates:
+        path = candidate.get("path") or ""
+        if not path or path in seen:
+            continue
+        seen.add(path)
+        out.append(candidate)
+    return out

connector/control.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from connector.local_runtime import (
+    ConnectorAlreadyRunningError,
+    assert_can_start,
+    clear_runtime,
+    runtime_path,
+    write_runtime,
+)
+from connector.logging import logger
+from connector.runtime import BackendRpcClient, ConnectorAuthenticationError, ConnectorConfig
+
+
+ControlNotifier = Callable[[str, Any], Awaitable[None]]
+
+
+class ConnectorController:
+    def __init__(
+        self,
+        *,
+        config_path: str | Path | None = None,
+        notifier: ControlNotifier | None = None,
+        client_factory: Callable[[ConnectorConfig], BackendRpcClient] = BackendRpcClient,
+    ) -> None:
+        self.config_path = Path(config_path) if config_path is not None else ConnectorConfig.default_path()
+        self.notifier = notifier
+        self.client_factory = client_factory
+        self.runtime_path = runtime_path(self.config_path)
+        self._runtime_task: asyncio.Task[None] | None = None
+        self._pairing_task: asyncio.Task[None] | None = None
+        self._last_error: str | None = None
+        self._auth_failed = False
+
+    def get_state(self, _params: Any = None) -> dict[str, Any]:
+        return {
+            "status": self._status(),
+            "running": self._runtime_task is not None and not self._runtime_task.done(),
+            "pairing": self._pairing_task is not None and not self._pairing_task.done(),
+            "authFailed": self._auth_failed,
+            "lastError": self._last_error,
+            "configPath": str(self.config_path),
+            "runtimePath": str(self.runtime_path),
+            "hasConfig": self.config_path.exists(),
+        }
+
+    def get_paths(self, _params: Any = None) -> dict[str, str]:
+        return {
+            "configPath": str(self.config_path),
+            "configDir": str(self.config_path.parent),
+            "runtimePath": str(self.runtime_path),
+        }
+
+    def get_config(self, _params: Any = None) -> dict[str, Any]:
+        if not self.config_path.exists():
+            return default_config_payload()
+        return config_to_payload(ConnectorConfig.load(self.config_path))
+
+    async def save_config(self, params: Any) -> dict[str, Any]:
+        config = config_from_params(params)
+        saved_path = config.save(self.config_path)
+        self._auth_failed = False
+        self._last_error = None
+        logger.info("saved connector config path={}", saved_path)
+        await self._emit_state()
+        return config_to_payload(config)
+
+    async def start(self, params: Any = None) -> dict[str, Any]:
+        if self._runtime_task is not None and not self._runtime_task.done():
+            return self.get_state()
+
+        config = config_from_params(params) if isinstance(params, dict) and params else ConnectorConfig.load(self.config_path)
+        self._last_error = None
+        self._auth_failed = False
+        try:
+            assert_can_start(self.runtime_path, config)
+        except ConnectorAlreadyRunningError as exc:
+            self._last_error = str(exc)
+            await self._emit_state()
+            raise
+        write_runtime(self.runtime_path, config, kind="desktop")
+        self._runtime_task = asyncio.create_task(self._run_runtime(config))
+        logger.info("starting connector runtime")
+        await self._emit_state()
+        return self.get_state()
+
+    async def stop(self, _params: Any = None) -> dict[str, Any]:
+        if self._runtime_task is not None and not self._runtime_task.done():
+            self._runtime_task.cancel()
+            try:
+                await self._runtime_task
+            except asyncio.CancelledError:
+                pass
+        self._runtime_task = None
+        clear_runtime(self.runtime_path)
+        logger.info("stopped connector runtime")
+        await self._emit_state()
+        return self.get_state()
+
+    async def restart(self, params: Any = None) -> dict[str, Any]:
+        await self.stop()
+        return await self.start(params)
+
+    async def start_pairing(self, params: Any) -> dict[str, Any]:
+        if self._pairing_task is not None and not self._pairing_task.done():
+            self._pairing_task.cancel()
+        server = str_param(params, "server") or str_param(params, "serverUrl")
+        server_url = await resolve_pair_server_url(server, timeout=float_param(params, "resolveTimeout", 10))
+        timeout = float_param(params, "timeout", 600)
+        poll_interval = float_param(params, "pollInterval", 2)
+        self._pairing_task = asyncio.create_task(self._run_pairing(server_url, timeout=timeout, poll_interval=poll_interval))
+        payload = {"status": "starting", "serverUrl": server_url}
+        await self._emit_pairing(payload)
+        await self._emit_state()
+        return payload
+
+    async def cancel_pairing(self, _params: Any = None) -> dict[str, Any]:
+        if self._pairing_task is not None and not self._pairing_task.done():
+            self._pairing_task.cancel()
+        self._pairing_task = None
+        payload = {"status": "cancelled"}
+        await self._emit_pairing(payload)
+        await self._emit_state()
+        return payload
+
+    async def shutdown(self) -> None:
+        if self._pairing_task is not None and not self._pairing_task.done():
+            self._pairing_task.cancel()
+        await self.stop()
+
+    async def _run_runtime(self, config: ConnectorConfig) -> None:
+        try:
+            await self.client_factory(config).run_forever()
+        except asyncio.CancelledError:
+            raise
+        except ConnectorAuthenticationError as exc:
+            self._auth_failed = True
+            self._last_error = str(exc)
+            logger.error("connector authentication failed: {}", exc)
+        except Exception as exc:
+            self._last_error = str(exc) or exc.__class__.__name__
+            logger.exception("connector runtime failed")
+        finally:
+            clear_runtime(self.runtime_path)
+            if self._runtime_task is asyncio.current_task():
+                self._runtime_task = None
+            await self._emit_state()
+
+    async def _run_pairing(self, server_url: str, *, timeout: float, poll_interval: float) -> None:
+        try:
+            async with httpx.AsyncClient(timeout=30) as client:
+                start_response = await client.post(
+                    f"{server_url}/pairing/start",
+                    json={"serverUrl": server_url, "ttlSeconds": int(timeout)},
+                )
+                start_response.raise_for_status()
+                pairing = start_response.json()
+                pairing_id = pairing["pairingId"]
+                code = pairing["code"]
+                await self._emit_pairing(
+                    {
+                        "status": "waiting",
+                        "serverUrl": server_url,
+                        "pairingId": pairing_id,
+                        "code": code,
+                    }
+                )
+
+                deadline = asyncio.get_running_loop().time() + timeout
+                while asyncio.get_running_loop().time() < deadline:
+                    poll_response = await client.post(f"{server_url}/pairing/poll", json={"pairingId": pairing_id})
+                    poll_response.raise_for_status()
+                    payload = poll_response.json()
+                    if payload["status"] == "claimed" and payload.get("config"):
+                        config = ConnectorConfig.from_mapping(payload["config"])
+                        config.save(self.config_path)
+                        await self._emit_pairing({"status": "claimed", "config": config_to_payload(config)})
+                        await self.start()
+                        return
+                    if payload["status"] in {"expired", "consumed"}:
+                        await self._emit_pairing({"status": payload["status"]})
+                        return
+                    await asyncio.sleep(poll_interval)
+            await self._emit_pairing({"status": "expired"})
+        except asyncio.CancelledError:
+            await self._emit_pairing({"status": "cancelled"})
+            raise
+        except Exception as exc:
+            self._last_error = str(exc) or exc.__class__.__name__
+            await self._emit_pairing({"status": "error", "error": self._last_error})
+            await self._emit_state()
+
+    def _status(self) -> str:
+        if self._runtime_task is not None and not self._runtime_task.done():
+            return "running"
+        if self._auth_failed:
+            return "expired credential"
+        if self._last_error:
+            return "error"
+        return "stopped"
+
+    async def _emit_state(self) -> None:
+        await self._notify("connector/state", self.get_state())
+
+    async def _emit_pairing(self, payload: dict[str, Any]) -> None:
+        await self._notify("connector/pairing", payload)
+
+    async def _notify(self, method: str, params: Any) -> None:
+        if self.notifier is not None:
+            await self.notifier(method, params)
+
+
+def default_config_payload() -> dict[str, Any]:
+    return {
+        "serverUrl": "",
+        "connectorId": "",
+        "connectorToken": "",
+        "heartbeatSeconds": 20,
+        "reconnectSeconds": 3,
+        "syncExistingOnConnect": True,
+        "syncIntervalSeconds": 30,
+        "stateDbPath": None,
+    }
+
+
+def config_to_payload(config: ConnectorConfig) -> dict[str, Any]:
+    return {
+        "serverUrl": config.server_url,
+        "connectorId": config.connector_id,
+        "connectorToken": config.connector_token,
+        "heartbeatSeconds": config.heartbeat_seconds,
+        "reconnectSeconds": config.reconnect_seconds,
+        "syncExistingOnConnect": config.sync_existing_on_connect,
+        "syncIntervalSeconds": config.sync_interval_seconds,
+        "stateDbPath": config.state_db_path,
+    }
+
+
+def config_from_params(params: Any) -> ConnectorConfig:
+    if not isinstance(params, dict):
+        raise ValueError("config params must be an object")
+    server_url = str(params.get("serverUrl") or "").strip().rstrip("/")
+    connector_id = str(params.get("connectorId") or "").strip()
+    connector_token = str(params.get("connectorToken") or "").strip()
+    if not server_url or not connector_id or not connector_token:
+        raise ValueError("serverUrl, connectorId, and connectorToken are required")
+    return ConnectorConfig(
+        server_url=server_url,
+        connector_id=connector_id,
+        connector_token=connector_token,
+        heartbeat_seconds=float(params.get("heartbeatSeconds", 20)),
+        reconnect_seconds=float(params.get("reconnectSeconds", 3)),
+        sync_existing_on_connect=bool(params.get("syncExistingOnConnect", True)),
+        sync_interval_seconds=float(params.get("syncIntervalSeconds", 30)),
+        state_db_path=params.get("stateDbPath") if isinstance(params.get("stateDbPath"), str) else None,
+    )
+
+
+async def resolve_pair_server_url(value: str | None, *, timeout: float = 10) -> str:
+    normalized = str(value or "").strip().rstrip("/")
+    if not normalized:
+        raise ValueError("server is required")
+    if normalized.startswith(("http://", "https://")):
+        return normalized
+    candidates = [f"https://{normalized}", f"http://{normalized}"]
+    errors: list[str] = []
+    for candidate in candidates:
+        try:
+            async with httpx.AsyncClient(timeout=timeout) as client:
+                response = await client.get(f"{candidate}/health")
+                if response.status_code < 500:
+                    return candidate
+                errors.append(f"{candidate}: HTTP {response.status_code}")
+        except httpx.RequestError as exc:
+            errors.append(f"{candidate}: {exc}")
+    raise ValueError(f"could not reach server over https or http ({'; '.join(errors)})")
+
+
+def str_param(params: Any, key: str) -> str | None:
+    if not isinstance(params, dict):
+        return None
+    value = params.get(key)
+    return value if isinstance(value, str) and value.strip() else None
+
+
+def float_param(params: Any, key: str, default: float) -> float:
+    if not isinstance(params, dict):
+        return default
+    try:
+        return float(params.get(key, default))
+    except (TypeError, ValueError):
+        return default