cortex-loop 0.1.0a1__py3-none-any.whl

cortex_ops_cli/_openai_bridge_protocol.py

@@ -0,0 +1,469 @@
+from __future__ import annotations
+
+import json
+import select
+import subprocess
+import time
+from pathlib import Path
+from typing import Any, Callable
+
+COMMAND_APPROVAL_METHOD = "item/commandExecution/requestApproval"
+FILE_CHANGE_APPROVAL_METHOD = "item/fileChange/requestApproval"
+NOTIFICATION_ITEM_COMPLETED = "item/completed"
+NOTIFICATION_TURN_COMPLETED = "turn/completed"
+NOTIFICATION_TASK_COMPLETE = "codex/event/task_complete"
+NOTIFICATION_THREAD_STATUS_CHANGED = "thread/status/changed"
+DEFAULT_APPROVAL_POLICY_CANDIDATES = ("untrusted", "unlessTrusted", "on-request")
+
+
+class BridgeError(RuntimeError):
+    pass
+
+
+class AppServerClient:
+    def __init__(self, *, codex_bin: str, cwd: Path, timeout_seconds: float) -> None:
+        cmd = [codex_bin, "app-server", "--listen", "stdio://"]
+        self.proc = subprocess.Popen(
+            cmd,
+            cwd=str(cwd),
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+        self.timeout_seconds = timeout_seconds
+        self._next_request_id = 1
+        self._pending_messages: list[dict[str, Any]] = []
+
+    def close(self) -> None:
+        if self.proc.poll() is not None:
+            return
+        try:
+            self.proc.terminate()
+            self.proc.wait(timeout=1.0)
+        except Exception:
+            self.proc.kill()
+            self.proc.wait(timeout=1.0)
+
+    def _write(self, payload: dict[str, Any]) -> None:
+        if self.proc.stdin is None:
+            raise BridgeError("codex app-server stdin unavailable")
+        self.proc.stdin.write(json.dumps(payload) + "\n")
+        self.proc.stdin.flush()
+
+    def _read_line(self, timeout_seconds: float) -> str:
+        if self.proc.stdout is None:
+            raise BridgeError("codex app-server stdout unavailable")
+        fd = self.proc.stdout.fileno()
+        ready, _, _ = select.select([fd], [], [], timeout_seconds)
+        if not ready:
+            raise BridgeError("Timed out waiting for codex app-server response")
+        line = self.proc.stdout.readline()
+        if not line:
+            stderr = ""
+            if self.proc.stderr is not None:
+                try:
+                    stderr = self.proc.stderr.read().strip()
+                except Exception:
+                    stderr = ""
+            raise BridgeError(f"codex app-server closed stdout. stderr={stderr[:4000]}")
+        return line
+
+    def _read_message(self, timeout_seconds: float, *, allow_pending: bool = True) -> dict[str, Any]:
+        if allow_pending and self._pending_messages:
+            return self._pending_messages.pop(0)
+        line = self._read_line(timeout_seconds)
+        try:
+            decoded = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise BridgeError(f"Invalid JSON-RPC payload from app-server: {exc}") from exc
+        if not isinstance(decoded, dict):
+            raise BridgeError("Invalid JSON-RPC payload from app-server: expected object")
+        return decoded
+
+    def _pop_pending_response(self, request_id: int) -> dict[str, Any] | None:
+        for idx, msg in enumerate(self._pending_messages):
+            if msg.get("id") == request_id and ("result" in msg or "error" in msg):
+                return self._pending_messages.pop(idx)
+        return None
+
+    def request(
+        self,
+        method: str,
+        params: dict[str, Any] | None = None,
+        *,
+        timeout_seconds: float | None = None,
+    ) -> Any:
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._write({"id": request_id, "method": method, "params": params or {}})
+        deadline = time.time() + (self.timeout_seconds if timeout_seconds is None else timeout_seconds)
+        while True:
+            pending_response = self._pop_pending_response(request_id)
+            if pending_response is not None:
+                if "error" in pending_response:
+                    raise BridgeError(f"JSON-RPC {method} failed: {pending_response['error']}")
+                return pending_response.get("result")
+
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                raise BridgeError(f"Timed out waiting for JSON-RPC response to {method}")
+
+            msg = self._read_message(max(0.01, remaining), allow_pending=False)
+            if msg.get("id") == request_id and ("result" in msg or "error" in msg):
+                if "error" in msg:
+                    raise BridgeError(f"JSON-RPC {method} failed: {msg['error']}")
+                return msg.get("result")
+            self._pending_messages.append(msg)
+
+    def send_notification(self, method: str, params: dict[str, Any] | None = None) -> None:
+        self._write({"method": method, "params": params or {}})
+
+    def send_server_request_result(self, request_id: Any, result: dict[str, Any]) -> None:
+        self._write({"id": request_id, "result": result})
+
+    def next_message(self, *, timeout_seconds: float | None = None) -> dict[str, Any]:
+        return self._read_message(self.timeout_seconds if timeout_seconds is None else timeout_seconds)
+
+
+class BridgeRunResult(dict):
+    pass
+
+
+def classify_command_surface(
+    *,
+    command_items: list[dict[str, Any]],
+    approval_requests: list[dict[str, Any]],
+) -> dict[str, Any]:
+    approved_item_ids = {
+        str(item.get("item_id") or "").strip()
+        for item in approval_requests
+        if isinstance(item, dict) and str(item.get("method")) == COMMAND_APPROVAL_METHOD
+    }
+    approved_item_ids.discard("")
+    command_item_ids = {
+        str(item.get("item_id") or "").strip()
+        for item in command_items
+        if isinstance(item, dict)
+    }
+    command_item_ids.discard("")
+
+    with_approval = sorted(command_item_ids & approved_item_ids)
+    without_approval = sorted(command_item_ids - approved_item_ids)
+    declined_item_ids = {
+        str(item.get("item_id") or "").strip()
+        for item in approval_requests
+        if isinstance(item, dict)
+        and str(item.get("method")) == COMMAND_APPROVAL_METHOD
+        and str(item.get("decision")) == "decline"
+    }
+    declined_item_ids.discard("")
+
+    nonblocking_declines: list[str] = []
+    for item in command_items:
+        if not isinstance(item, dict):
+            continue
+        item_id = str(item.get("item_id") or "").strip()
+        if item_id not in declined_item_ids:
+            continue
+        status = str(item.get("status") or "").strip().lower()
+        has_exec_payload = item.get("exit_code") is not None or bool(
+            str(item.get("aggregated_output") or "").strip()
+        )
+        if status != "declined" or has_exec_payload:
+            nonblocking_declines.append(item_id or "<missing_item_id>")
+
+    return {
+        "command_item_ids": sorted(command_item_ids),
+        "approved_item_ids": sorted(approved_item_ids),
+        "command_items_with_approval": with_approval,
+        "command_items_without_approval": without_approval,
+        "declined_item_ids": sorted(declined_item_ids),
+        "nonblocking_declines": sorted(set(nonblocking_declines)),
+    }
+
+
+def initialize_app_server(*, client: AppServerClient) -> None:
+    init_result = client.request(
+        "initialize",
+        {
+            "clientInfo": {
+                "name": "cortex-openai-bridge",
+                "title": "Cortex OpenAI Bridge",
+                "version": "1.0.0",
+            },
+            "capabilities": {
+                "experimentalApi": False,
+                "optOutNotificationMethods": [],
+            },
+        },
+    )
+    if not isinstance(init_result, dict):
+        raise BridgeError("initialize did not return a JSON object")
+    client.send_notification("initialized", {})
+
+
+def is_policy_compat_error(exc: BridgeError) -> bool:
+    message = str(exc).lower()
+    return (
+        "approvalpolicy" in message
+        or "askforapproval" in message
+        or "unknown variant" in message
+        or "invalid type" in message
+        or "invalid value" in message
+        or "unlesstrusted" in message
+    )
+
+
+def start_thread_with_policy_fallback(
+    *,
+    client: AppServerClient,
+    cwd: Path,
+    model: str | None,
+    approval_policy_candidates: tuple[str, ...] = DEFAULT_APPROVAL_POLICY_CANDIDATES,
+) -> tuple[dict[str, Any], str]:
+    if not approval_policy_candidates:
+        raise BridgeError("No approval policy candidates configured.")
+
+    last_error: BridgeError | None = None
+    for idx, policy in enumerate(approval_policy_candidates):
+        params: dict[str, Any] = {
+            "cwd": str(cwd),
+            "approvalPolicy": policy,
+            "sandbox": "workspace-write",
+        }
+        if model:
+            params["model"] = model
+        try:
+            result = client.request("thread/start", params)
+            if not isinstance(result, dict):
+                raise BridgeError("thread/start did not return a JSON object")
+            return result, policy
+        except BridgeError as exc:
+            last_error = exc
+            if idx < len(approval_policy_candidates) - 1 and is_policy_compat_error(exc):
+                continue
+            raise
+
+    if last_error is not None:
+        raise last_error
+    raise BridgeError("thread/start failed: no approval policy candidate succeeded")
+
+
+def _bridge_turn_result(
+    *,
+    thread_id: str,
+    turn_id: str,
+    approval_policy_used: str,
+    approval_policy_candidates: tuple[str, ...],
+    final_text: str,
+    approval_requests: list[dict[str, Any]],
+    command_completion_items: list[dict[str, Any]],
+    coverage_gaps: list[str],
+    duplicate_turn_completed_count: int,
+    elapsed_seconds: float,
+) -> BridgeRunResult:
+    return BridgeRunResult(
+        {
+            "ok": True,
+            "thread_id": thread_id,
+            "turn_id": turn_id,
+            "approval_policy_used": approval_policy_used,
+            "approval_policy_candidates": list(approval_policy_candidates),
+            "text": final_text,
+            "response_present": bool(final_text.strip()),
+            "approval_requests": approval_requests,
+            "command_completion_items": command_completion_items,
+            "command_surface": classify_command_surface(
+                command_items=command_completion_items,
+                approval_requests=approval_requests,
+            ),
+            "coverage_gaps": sorted(set(coverage_gaps)),
+            "duplicate_turn_completed_count": duplicate_turn_completed_count,
+            "elapsed_seconds": elapsed_seconds,
+        }
+    )
+
+
+def execute_turn(
+    *,
+    codex_bin: str,
+    cwd: Path,
+    prompt: str,
+    model: str | None,
+    timeout_seconds: float,
+    approval_policy_candidates: tuple[str, ...] = DEFAULT_APPROVAL_POLICY_CANDIDATES,
+    approval_handler: Callable[[str, dict[str, Any]], tuple[str, dict[str, Any]]],
+    client_cls: type[AppServerClient] = AppServerClient,
+) -> BridgeRunResult:
+    client = client_cls(codex_bin=codex_bin, cwd=cwd, timeout_seconds=timeout_seconds)
+    start_time = time.time()
+    try:
+        initialize_app_server(client=client)
+        thread_result, approval_policy_used = start_thread_with_policy_fallback(
+            client=client,
+            cwd=cwd,
+            model=model,
+            approval_policy_candidates=approval_policy_candidates,
+        )
+        thread_id = str(((thread_result or {}).get("thread") or {}).get("id") or "").strip()
+        if not thread_id:
+            raise BridgeError("thread/start did not return thread.id")
+
+        turn_result = client.request(
+            "turn/start",
+            {
+                "threadId": thread_id,
+                "input": [{"type": "text", "text": prompt, "textElements": []}],
+                "cwd": str(cwd),
+            },
+        )
+        turn_id = str(((turn_result or {}).get("turn") or {}).get("id") or "").strip()
+        if not turn_id:
+            raise BridgeError("turn/start did not return turn.id")
+
+        final_agent_messages: dict[str, str] = {}
+        completed_turns: set[str] = set()
+        duplicate_turn_completed_count = 0
+        command_completion_items: list[dict[str, Any]] = []
+        approval_requests: list[dict[str, Any]] = []
+        coverage_gaps: list[str] = []
+        task_complete_turns: set[str] = set()
+        task_complete_last_messages: dict[str, str] = {}
+
+        while True:
+            msg = client.next_message()
+            method = str(msg.get("method") or "")
+
+            if method and "id" in msg and "result" not in msg and "error" not in msg:
+                request_id = msg.get("id")
+                params = msg.get("params") if isinstance(msg.get("params"), dict) else {}
+                if method in {COMMAND_APPROVAL_METHOD, FILE_CHANGE_APPROVAL_METHOD}:
+                    decision, metadata = approval_handler(method, params)
+                    if decision not in {"accept", "decline", "cancel"}:
+                        decision = "decline"
+                    approval_requests.append(
+                        {
+                            "method": method,
+                            "item_id": str(params.get("itemId") or ""),
+                            "command": str(params.get("command") or ""),
+                            "cwd": str(params.get("cwd") or ""),
+                            "decision": decision,
+                            "metadata": metadata,
+                        }
+                    )
+                    client.send_server_request_result(request_id, {"decision": decision})
+                    continue
+                coverage_gaps.append(f"unhandled_server_request:{method}")
+                client.send_server_request_result(request_id, {"decision": "decline"})
+                continue
+
+            if not method:
+                continue
+
+            params = msg.get("params") if isinstance(msg.get("params"), dict) else {}
+            if method == NOTIFICATION_ITEM_COMPLETED:
+                item = params.get("item") if isinstance(params.get("item"), dict) else {}
+                item_type = str(item.get("type") or "")
+                item_turn_id = str(params.get("turnId") or "")
+                if item_type == "agentMessage" and item_turn_id:
+                    final_agent_messages[item_turn_id] = str(item.get("text") or "")
+                if item_type == "commandExecution":
+                    command_completion_items.append(
+                        {
+                            "item_id": str(item.get("id") or ""),
+                            "turn_id": item_turn_id,
+                            "status": str(item.get("status") or ""),
+                            "exit_code": item.get("exitCode"),
+                            "aggregated_output": item.get("aggregatedOutput"),
+                            "command": str(item.get("command") or ""),
+                            "cwd": str(item.get("cwd") or ""),
+                        }
+                    )
+                continue
+
+            if method == NOTIFICATION_TURN_COMPLETED:
+                notification_turn = params.get("turn") if isinstance(params.get("turn"), dict) else {}
+                completed_turn_id = str(notification_turn.get("id") or "")
+                if not completed_turn_id:
+                    coverage_gaps.append("turn_completed_missing_turn_id")
+                    continue
+                if completed_turn_id in completed_turns:
+                    duplicate_turn_completed_count += 1
+                    continue
+                completed_turns.add(completed_turn_id)
+                if completed_turn_id != turn_id:
+                    continue
+
+                final_text = str(final_agent_messages.get(turn_id, "") or "")
+                elapsed_seconds = round(time.time() - start_time, 6)
+                return _bridge_turn_result(
+                    thread_id=thread_id,
+                    turn_id=turn_id,
+                    approval_policy_used=approval_policy_used,
+                    approval_policy_candidates=approval_policy_candidates,
+                    final_text=final_text,
+                    approval_requests=approval_requests,
+                    command_completion_items=command_completion_items,
+                    coverage_gaps=coverage_gaps,
+                    duplicate_turn_completed_count=duplicate_turn_completed_count,
+                    elapsed_seconds=elapsed_seconds,
+                )
+
+            if method == NOTIFICATION_TASK_COMPLETE:
+                task_msg = params.get("msg") if isinstance(params.get("msg"), dict) else {}
+                completed_turn_id = str(task_msg.get("turn_id") or params.get("id") or "").strip()
+                if completed_turn_id:
+                    task_complete_turns.add(completed_turn_id)
+                    last_message = str(task_msg.get("last_agent_message") or "")
+                    if last_message:
+                        task_complete_last_messages[completed_turn_id] = last_message
+                    if completed_turn_id == turn_id and completed_turn_id not in completed_turns:
+                        coverage_gaps.append("turn_completed_missing_used_task_complete_fallback")
+                        final_text = str(
+                            final_agent_messages.get(turn_id, "")
+                            or task_complete_last_messages.get(turn_id, "")
+                            or ""
+                        )
+                        elapsed_seconds = round(time.time() - start_time, 6)
+                        return _bridge_turn_result(
+                            thread_id=thread_id,
+                            turn_id=turn_id,
+                            approval_policy_used=approval_policy_used,
+                            approval_policy_candidates=approval_policy_candidates,
+                            final_text=final_text,
+                            approval_requests=approval_requests,
+                            command_completion_items=command_completion_items,
+                            coverage_gaps=coverage_gaps,
+                            duplicate_turn_completed_count=duplicate_turn_completed_count,
+                            elapsed_seconds=elapsed_seconds,
+                        )
+                continue
+
+            if method == NOTIFICATION_THREAD_STATUS_CHANGED:
+                event_thread_id = str(params.get("threadId") or "")
+                status = params.get("status") if isinstance(params.get("status"), dict) else {}
+                status_type = str(status.get("type") or "")
+                if event_thread_id == thread_id and status_type == "idle" and turn_id in task_complete_turns:
+                    coverage_gaps.append("turn_completed_missing_used_task_complete_fallback")
+                    final_text = str(
+                        final_agent_messages.get(turn_id, "")
+                        or task_complete_last_messages.get(turn_id, "")
+                        or ""
+                    )
+                    elapsed_seconds = round(time.time() - start_time, 6)
+                    return _bridge_turn_result(
+                        thread_id=thread_id,
+                        turn_id=turn_id,
+                        approval_policy_used=approval_policy_used,
+                        approval_policy_candidates=approval_policy_candidates,
+                        final_text=final_text,
+                        approval_requests=approval_requests,
+                        command_completion_items=command_completion_items,
+                        coverage_gaps=coverage_gaps,
+                        duplicate_turn_completed_count=duplicate_turn_completed_count,
+                        elapsed_seconds=elapsed_seconds,
+                    )
+    finally:
+        client.close()