jupyterlab-codex-sidebar 0.1.3 → 0.1.5

package/jupyterlab_codex/protocol.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Literal, Tuple
+
+ProtocolVersion = Literal["1.0.0"]
+
+PROTOCOL_VERSION: ProtocolVersion = "1.0.0"
+
+
+class ProtocolParseError(ValueError):
+    """Raised when an incoming client payload cannot be interpreted."""
+
+
+def _coerce_string(value: Any) -> str:
+    if not isinstance(value, str):
+        return ""
+    return value.strip()
+
+
+def _coerce_string_untrimmed(value: Any) -> str:
+    if not isinstance(value, str):
+        return ""
+    return value
+
+
+def _coerce_bool(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "y", "yes", "on"}
+    return False
+
+
+def parse_client_message(raw: Any) -> Tuple[str, Dict[str, Any]]:
+    """
+    Parse and normalize incoming websocket payloads.
+
+    Returns (msg_type, payload) with normalized string/bool fields.
+    """
+    if not isinstance(raw, dict):
+        raise ProtocolParseError("Invalid message payload")
+
+    msg_type = raw.get("type")
+    if msg_type == "start_session":
+        return (
+            "start_session",
+            {
+                "sessionId": _coerce_string(raw.get("sessionId")),
+                "notebookPath": _coerce_string(raw.get("notebookPath")),
+                "sessionContextKey": _coerce_string(raw.get("sessionContextKey")),
+                "forceNewThread": _coerce_bool(raw.get("forceNewThread")),
+                "commandPath": _coerce_string(raw.get("commandPath")),
+            },
+        )
+
+    if msg_type == "send":
+        return (
+            "send",
+            {
+                "sessionId": _coerce_string(raw.get("sessionId")),
+                "sessionContextKey": _coerce_string(raw.get("sessionContextKey")),
+                "content": _coerce_string(raw.get("content")),
+                "notebookPath": _coerce_string(raw.get("notebookPath")),
+                "commandPath": _coerce_string(raw.get("commandPath")),
+                "model": _coerce_string(raw.get("model")),
+                "reasoningEffort": _coerce_string(raw.get("reasoningEffort")),
+                "sandbox": _coerce_string(raw.get("sandbox")),
+                "selection": _coerce_string_untrimmed(raw.get("selection")),
+                "cellOutput": _coerce_string_untrimmed(raw.get("cellOutput")),
+                "images": raw.get("images") if isinstance(raw.get("images"), list) else [],
+                "uiSelectionPreview": raw.get("uiSelectionPreview"),
+                "uiCellOutputPreview": raw.get("uiCellOutputPreview"),
+                "selectionTruncated": _coerce_bool(raw.get("selectionTruncated")),
+                "cellOutputTruncated": _coerce_bool(raw.get("cellOutputTruncated")),
+            },
+        )
+
+    if msg_type == "delete_session":
+        return ("delete_session", {"sessionId": _coerce_string(raw.get("sessionId"))})
+
+    if msg_type == "delete_all_sessions":
+        return ("delete_all_sessions", {})
+
+    if msg_type == "cancel":
+        return ("cancel", {"runId": _coerce_string(raw.get("runId"))})
+
+    if msg_type == "end_session":
+        return ("end_session", {"sessionId": _coerce_string(raw.get("sessionId"))})
+
+    if msg_type == "refresh_rate_limits":
+        return ("refresh_rate_limits", {})
+
+    raise ProtocolParseError("Unknown message type")
+
+
+def _build_base_message(msg_type: str) -> Dict[str, Any]:
+    return {"type": msg_type, "protocolVersion": PROTOCOL_VERSION}
+
+
+def build_cli_defaults_payload(
+    *,
+    model: str | None = None,
+    reasoning_effort: str | None = None,
+    available_models: list[dict[str, Any]] | None = None,
+) -> Dict[str, Any]:
+    payload = _build_base_message("cli_defaults")
+    if model is not None:
+        payload["model"] = model
+    if reasoning_effort is not None:
+        payload["reasoningEffort"] = reasoning_effort
+    if available_models is not None:
+        payload["availableModels"] = available_models
+    return payload
+
+
+def build_status_payload(
+    *,
+    state: str,
+    run_id: str | None = None,
+    session_id: str | None = None,
+    session_context_key: str | None = None,
+    notebook_path: str | None = None,
+    run_mode: str | None = None,
+    paired_ok: bool | None = None,
+    paired_path: str = "",
+    paired_os_path: str = "",
+    paired_message: str = "",
+    notebook_mode: str = "",
+    effective_sandbox: str | None = None,
+    history: list[dict[str, Any]] | None = None,
+    session_resolution: str | None = None,
+    session_resolution_notice: str | None = None,
+) -> Dict[str, Any]:
+    payload = _build_base_message("status")
+    payload["state"] = state
+
+    if run_id:
+        payload["runId"] = run_id
+    if session_id:
+        payload["sessionId"] = session_id
+    if session_context_key:
+        payload["sessionContextKey"] = session_context_key
+    if notebook_path:
+        payload["notebookPath"] = notebook_path
+    if run_mode:
+        payload["runMode"] = run_mode
+    if paired_ok is not None:
+        payload["pairedOk"] = paired_ok
+    if paired_path:
+        payload["pairedPath"] = paired_path
+    if paired_os_path:
+        payload["pairedOsPath"] = paired_os_path
+    if paired_message:
+        payload["pairedMessage"] = paired_message
+    if notebook_mode:
+        payload["notebookMode"] = notebook_mode
+    if history is not None:
+        payload["history"] = history
+    if session_resolution is not None:
+        payload["sessionResolution"] = session_resolution
+    if session_resolution_notice:
+        payload["sessionResolutionNotice"] = session_resolution_notice
+    if effective_sandbox:
+        payload["effectiveSandbox"] = effective_sandbox
+    return payload
+
+
+def build_output_payload(
+    *,
+    run_id: str,
+    session_id: str,
+    session_context_key: str,
+    notebook_path: str,
+    text: str,
+    role: str = "assistant",
+) -> Dict[str, Any]:
+    return _build_base_message("output") | {
+        "runId": run_id,
+        "sessionId": session_id,
+        "sessionContextKey": session_context_key,
+        "notebookPath": notebook_path,
+        "text": text,
+        "role": role,
+    }
+
+
+def build_event_payload(
+    *,
+    run_id: str,
+    session_id: str,
+    session_context_key: str,
+    notebook_path: str,
+    payload: Dict[str, Any],
+) -> Dict[str, Any]:
+    return _build_base_message("event") | {
+        "runId": run_id,
+        "sessionId": session_id,
+        "sessionContextKey": session_context_key,
+        "notebookPath": notebook_path,
+        "payload": payload,
+    }
+
+
+def build_done_payload(
+    *,
+    run_id: str,
+    session_id: str,
+    session_context_key: str,
+    notebook_path: str,
+    exit_code: int | None,
+    file_changed: bool,
+    run_mode: str,
+    paired_ok: bool,
+    paired_path: str,
+    paired_os_path: str,
+    paired_message: str,
+    notebook_mode: str,
+    cancelled: bool = False,
+) -> Dict[str, Any]:
+    payload = _build_base_message("done") | {
+        "runId": run_id,
+        "sessionId": session_id,
+        "sessionContextKey": session_context_key,
+        "notebookPath": notebook_path,
+        "exitCode": exit_code,
+        "fileChanged": file_changed,
+        "runMode": run_mode,
+        "pairedOk": paired_ok,
+    }
+    payload["pairedPath"] = paired_path
+    payload["pairedOsPath"] = paired_os_path
+    payload["pairedMessage"] = paired_message
+    payload["notebookMode"] = notebook_mode
+    if cancelled:
+        payload["cancelled"] = True
+    return payload
+
+
+def build_error_payload(
+    *,
+    run_id: str | None = None,
+    session_id: str | None = None,
+    session_context_key: str = "",
+    notebook_path: str = "",
+    message: str,
+    run_mode: str | None = None,
+    suggested_command_path: str | None = None,
+    paired_ok: bool | None = None,
+    paired_path: str = "",
+    paired_os_path: str = "",
+    paired_message: str = "",
+    notebook_mode: str = "",
+) -> Dict[str, Any]:
+    payload = _build_base_message("error")
+    if run_id:
+        payload["runId"] = run_id
+    if session_id:
+        payload["sessionId"] = session_id
+    if session_context_key:
+        payload["sessionContextKey"] = session_context_key
+    if notebook_path:
+        payload["notebookPath"] = notebook_path
+    payload["message"] = message
+    if run_mode is not None:
+        payload["runMode"] = run_mode
+    if suggested_command_path:
+        payload["suggestedCommandPath"] = suggested_command_path
+    if paired_ok is not None:
+        payload["pairedOk"] = paired_ok
+    if paired_path:
+        payload["pairedPath"] = paired_path
+    if paired_os_path:
+        payload["pairedOsPath"] = paired_os_path
+    if paired_message:
+        payload["pairedMessage"] = paired_message
+    if notebook_mode:
+        payload["notebookMode"] = notebook_mode
+    return payload
+
+
+def build_rate_limits_payload(snapshot: Any) -> Dict[str, Any]:
+    return _build_base_message("rate_limits") | {"snapshot": snapshot}
+
+
+def build_delete_all_payload(
+    *,
+    ok: bool,
+    deleted_count: int,
+    failed_count: int,
+    message: str,
+) -> Dict[str, Any]:
+    payload = _build_base_message("delete_all_sessions")
+    payload["ok"] = ok
+    payload["deletedCount"] = deleted_count
+    payload["failedCount"] = failed_count
+    payload["message"] = message
+    return payload

package/jupyterlab_codex/runner.py

@@ -11,6 +11,7 @@ class CodexRunner:
     def __init__(self, command: str = "codex", args: List[str] | None = None):
         configured_command = os.environ.get("JUPYTERLAB_CODEX_COMMAND", "").strip()
         self._command = self._resolve_command(configured_command or command)
+        self._model_catalog_cache: dict[str, tuple[float, list[dict[str, Any]]]] = {}
         if args is not None:
             self._raw_args = list(args)
             self._common_args: List[str] = []
@@ -32,24 +33,24 @@ class CodexRunner:
             "never",
             "--skip-git-repo-check",
         ]
-        self._model_catalog_cache: list[dict[str, Any]] = []
-        self._model_catalog_cache_time = 0.0
 
-    async def list_available_models(self, command: str | None = None) -> list[dict[str, Any]]:
+    async def list_available_models(
+        self, command: str | None = None, force_refresh: bool = False
+    ) -> list[dict[str, Any]]:
+        command_to_run = self._resolve_command((command or "").strip() or self._command)
         now = time.monotonic()
-        if self._model_catalog_cache and now - self._model_catalog_cache_time < 600:
-            return list(self._model_catalog_cache)
+        cached = self._model_catalog_cache.get(command_to_run)
+        if not force_refresh and cached and now - cached[0] < 600:
+            return list(cached[1])
 
-        command_to_run = self._resolve_command((command or "").strip() or self._command)
         try:
             models = await self._load_available_models(command_to_run)
         except Exception:
-            return []
+            return list(cached[1]) if cached else []
 
         if not isinstance(models, list) or not models:
-            return []
-        self._model_catalog_cache = models
-        self._model_catalog_cache_time = now
+            return list(cached[1]) if cached else []
+        self._model_catalog_cache[command_to_run] = (now, models)
         return list(models)
 
     async def _load_available_models(self, command_to_run: str) -> list[dict[str, Any]]:
@@ -60,25 +61,61 @@ class CodexRunner:
             "stdio://",
             stdin=asyncio.subprocess.PIPE,
             stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.DEVNULL,
         )
-        if proc.stdin is None or proc.stdout is None or proc.stderr is None:
+        if proc.stdin is None or proc.stdout is None:
             raise RuntimeError("Failed to open app-server subprocess streams")
 
-        async def read_message() -> dict[str, Any] | None:
-            raw = await asyncio.wait_for(proc.stdout.readline(), timeout=3.0)
-            if not raw:
+        buffer = bytearray()
+        max_message_bytes = 1024 * 1024
+
+        def _pop_message_from_buffer() -> dict[str, Any] | None:
+            separator_index = buffer.find(b"\n")
+            if separator_index < 0:
                 return None
-            line = raw.decode("utf-8", errors="replace").strip()
+            raw_line = bytes(buffer[:separator_index])
+            del buffer[:separator_index + 1]
+            line = raw_line.decode("utf-8", errors="replace").strip()
             if not line:
-                return None
+                return {}
             try:
                 payload = json.loads(line)
             except json.JSONDecodeError:
-                return None
+                return {}
             if isinstance(payload, dict):
                 return payload
-            return None
+            return {}
+
+        async def read_message() -> dict[str, Any] | None:
+            deadline = time.monotonic() + 3.0
+            while True:
+                message = _pop_message_from_buffer()
+                if message is not None:
+                    if message:
+                        return message
+                    continue
+
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise RuntimeError("App-server response timed out")
+
+                chunk = await asyncio.wait_for(proc.stdout.read(4096), timeout=remaining)
+                if not chunk:
+                    if buffer:
+                        line = buffer.decode("utf-8", errors="replace").strip()
+                        buffer.clear()
+                        if not line:
+                            return None
+                        try:
+                            payload = json.loads(line)
+                        except json.JSONDecodeError:
+                            return None
+                        if isinstance(payload, dict):
+                            return payload
+                    return None
+                buffer.extend(chunk)
+                if len(buffer) > max_message_bytes and b"\n" not in buffer:
+                    raise RuntimeError("App-server emitted an oversized response line")
 
         async def read_response(expected_id: int) -> dict[str, Any]:
             while True:
@@ -239,10 +276,15 @@ class CodexRunner:
         sandbox: str | None = None,
         images: List[str] | None = None,
         command: str | None = None,
+        resume_session_id: str | None = None,
     ) -> int:
         command_to_run = self._resolve_command((command or "").strip() or self._command)
         args = self._args_for_options(
-            model=model, reasoning_effort=reasoning_effort, sandbox=sandbox, images=images
+            model=model,
+            reasoning_effort=reasoning_effort,
+            sandbox=sandbox,
+            images=images,
+            resume_session_id=resume_session_id,
         )
 
         proc = await asyncio.create_subprocess_exec(
@@ -263,15 +305,39 @@ class CodexRunner:
         proc.stdin.close()
 
         async def _read_stdout() -> None:
-            async for raw in proc.stdout:
-                line = raw.decode("utf-8", errors="replace").strip()
-                if not line:
-                    continue
-                try:
-                    event = json.loads(line)
-                except json.JSONDecodeError:
-                    event = {"type": "raw", "text": line}
-                await on_event(event)
+            buffer = bytearray()
+            max_event_line_bytes = 1024 * 1024
+            while True:
+                chunk = await proc.stdout.read(8192)
+                if not chunk:
+                    if buffer:
+                        line = buffer.decode("utf-8", errors="replace").strip()
+                        if line:
+                            try:
+                                event = json.loads(line)
+                            except json.JSONDecodeError:
+                                event = {"type": "raw", "text": line}
+                            await on_event(event)
+                    break
+
+                buffer.extend(chunk)
+                if len(buffer) > max_event_line_bytes and b"\n" not in buffer:
+                    raise RuntimeError("Codex emitted an oversized unterminated stdout line")
+                while True:
+                    separator_index = buffer.find(b"\n")
+                    if separator_index < 0:
+                        break
+                    raw_line = bytes(buffer[:separator_index])
+                    del buffer[:separator_index + 1]
+
+                    line = raw_line.decode("utf-8", errors="replace").strip()
+                    if not line:
+                        continue
+                    try:
+                        event = json.loads(line)
+                    except json.JSONDecodeError:
+                        event = {"type": "raw", "text": line}
+                    await on_event(event)
 
         async def _read_stderr() -> None:
             # Stream stderr so users can see prompts/errors even if Codex blocks.
@@ -287,6 +353,9 @@ class CodexRunner:
         except asyncio.CancelledError:
             await self._terminate_process(proc)
             raise
+        except Exception:
+            await self._terminate_process(proc)
+            raise
 
     async def _terminate_process(self, proc: asyncio.subprocess.Process) -> None:
         if proc.returncode is not None:
@@ -308,11 +377,13 @@ class CodexRunner:
         reasoning_effort: str | None,
         sandbox: str | None,
         images: List[str] | None,
+        resume_session_id: str | None = None,
     ) -> List[str]:
         requested_model = (model or "").strip()
         requested_reasoning_effort = (reasoning_effort or "").strip()
         requested_sandbox = (sandbox or "").strip()
         requested_images = [p for p in (images or []) if isinstance(p, str) and p.strip()]
+        requested_resume_session_id = (resume_session_id or "").strip()
 
         if self._raw_args is not None:
             args = list(self._raw_args)
@@ -336,7 +407,17 @@ class CodexRunner:
                 cleaned.append(token)
                 idx += 1
 
-            insertion_index = cleaned.index("-") if "-" in cleaned else len(cleaned)
+            stdin_insertion_index = cleaned.index("-") if "-" in cleaned else len(cleaned)
+            if requested_resume_session_id and "resume" not in cleaned:
+                cleaned[stdin_insertion_index:stdin_insertion_index] = [
+                    "resume",
+                    requested_resume_session_id,
+                ]
+            options_insertion_index = (
+                cleaned.index("resume")
+                if "resume" in cleaned
+                else (cleaned.index("-") if "-" in cleaned else len(cleaned))
+            )
             to_insert: List[str] = []
             if requested_images:
                 to_insert.extend(["--image", *requested_images])
@@ -347,11 +428,36 @@ class CodexRunner:
             if requested_reasoning_effort:
                 to_insert.extend(["-c", f'model_reasoning_effort="{requested_reasoning_effort}"'])
 
-            cleaned[insertion_index:insertion_index] = to_insert
+            cleaned[options_insertion_index:options_insertion_index] = to_insert
             return cleaned
 
         effective_model = requested_model or self._default_model
         effective_sandbox = requested_sandbox or self._default_sandbox
+
+        if requested_resume_session_id:
+            # `codex exec resume` does not accept `--sandbox` after `resume`.
+            # Sandbox must be passed as an `exec` option before the `resume` subcommand.
+            args: List[str] = [
+                "--ask-for-approval",
+                "never",
+                "exec",
+                "--json",
+                "--color",
+                "never",
+                "--skip-git-repo-check",
+            ]
+            if effective_sandbox:
+                args.extend(["--sandbox", effective_sandbox])
+            args.append("resume")
+            if effective_model:
+                args.extend(["-m", effective_model])
+            if requested_reasoning_effort:
+                args.extend(["-c", f'model_reasoning_effort="{requested_reasoning_effort}"'])
+            if requested_images:
+                args.extend(["--image", *requested_images])
+            args.extend([requested_resume_session_id, "-"])
+            return args
+
         args = list(self._common_args)
         if effective_sandbox:
             args.extend(["--sandbox", effective_sandbox])