diff mrjunejune/inference/copilot_sidecar.py @ 260:1f9877b637e9

Add Copilot-powered cyberpunk JRPG chat Integrate the production JRPG chat with Seobeo streaming, Deita persistence, and a Bazel-managed Copilot SDK and LiteLLM inference stack. Co-authored-by: Copilot <[email protected]>
author MrJuneJune <mrjunejune@users.noreply.github.com>
date Wed, 05 Aug 2026 09:19:41 -0700
parents
children 056790c4fb0d
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/mrjunejune/inference/copilot_sidecar.py	Wed Aug 05 09:19:41 2026 -0700
@@ -0,0 +1,744 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import os
+import sys
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, Awaitable, Callable, Protocol
+
+from copilot import CopilotClient, ProviderConfig, RuntimeConnection
+from copilot.rpc import PermissionDecisionReject
+
+
+JsonObject = dict[str, Any]
+Emit = Callable[[JsonObject], Awaitable[None]]
+
+
+class Session(Protocol):
+    session_id: str
+
+    def on(self, handler: Callable[[Any], None]) -> Callable[[], None]: ...
+
+    async def send(self, prompt: str) -> str: ...
+
+    async def abort(self) -> None: ...
+
+    async def disconnect(self) -> None: ...
+
+
+class Client(Protocol):
+    async def start(self) -> None: ...
+
+    async def stop(self) -> None: ...
+
+    async def create_session(self, **kwargs: Any) -> Session: ...
+
+    async def resume_session(self, session_id: str, **kwargs: Any) -> Session: ...
+
+    async def delete_session(self, session_id: str) -> None: ...
+
+
+@dataclass(frozen=True)
+class SidecarConfig:
+    base_url: str
+    model: str
+    wire_api: str
+    base_directory: str
+    api_key: str | None = None
+    idle_timeout_seconds: int = 900
+    max_sessions: int = 64
+
+    @classmethod
+    def from_environment(cls) -> SidecarConfig:
+        base_url = os.environ.get("LITELLM_BASE_URL", "").strip()
+        model = os.environ.get("LITELLM_MODEL", "").strip()
+        wire_api = os.environ.get("LITELLM_WIRE_API", "").strip()
+        base_directory = os.environ.get("COPILOT_SIDECAR_HOME", "").strip()
+        missing = [
+            name
+            for name, value in (
+                ("LITELLM_BASE_URL", base_url),
+                ("LITELLM_MODEL", model),
+                ("LITELLM_WIRE_API", wire_api),
+                ("COPILOT_SIDECAR_HOME", base_directory),
+            )
+            if not value
+        ]
+        if missing:
+            raise ValueError(f"missing required environment: {', '.join(missing)}")
+        if wire_api not in ("completions", "responses"):
+            raise ValueError("LITELLM_WIRE_API must be 'completions' or 'responses'")
+        return cls(
+            base_url=base_url,
+            model=model,
+            wire_api=wire_api,
+            base_directory=os.path.abspath(base_directory),
+            api_key=os.environ.get("LITELLM_API_KEY") or None,
+            idle_timeout_seconds=max(
+                1,
+                int(os.environ.get("COPILOT_SESSION_IDLE_SECONDS", "900")),
+            ),
+            max_sessions=max(
+                1,
+                int(os.environ.get("COPILOT_MAX_SESSIONS", "64")),
+            ),
+        )
+
+
+@dataclass
+class Turn:
+    request_id: str
+    done: bool = False
+
+
+@dataclass
+class Conversation:
+    session: Session
+    unsubscribe: Callable[[], None]
+    active: Turn | None = None
+    last_used: float = field(default_factory=time.monotonic)
+
+
+@dataclass
+class ConversationGate:
+    lock: asyncio.Lock
+    users: int = 0
+
+
+def deny_permission(*_args: Any, **_kwargs: Any) -> PermissionDecisionReject:
+    return PermissionDecisionReject(feedback="The inference sidecar denies all permissions.")
+
+
+class Sidecar:
+    def __init__(self, client: Client, config: SidecarConfig, emit: Emit):
+        self._client = client
+        self._config = config
+        self._emit = emit
+        self._conversations: dict[str, Conversation] = {}
+        self._conversations_lock = asyncio.Lock()
+        self._conversation_gates: dict[str, ConversationGate] = {}
+        self._conversation_gates_lock = asyncio.Lock()
+        self._dispatch_condition = asyncio.Condition()
+        self._active_dispatches = 0
+        self._event_tasks: set[asyncio.Task[None]] = set()
+        self._cleanup_task: asyncio.Task[None] | None = None
+        self._started = False
+        self.shutting_down = False
+
+    async def start(self) -> None:
+        if not self._started:
+            await self._client.start()
+            self._started = True
+            self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+
+    async def announce_ready(self) -> None:
+        await self._send("ready", None, None, status="ok")
+
+    async def dispatch(self, command: JsonObject) -> None:
+        command_name = command.get("command")
+        request_id = command.get("request_id")
+        conversation_id = command.get("conversation_id")
+        if not isinstance(request_id, str) or not request_id:
+            normalized_request_id = request_id if isinstance(request_id, str) else None
+            normalized_conversation_id = (
+                conversation_id if isinstance(conversation_id, str) else None
+            )
+            await self._send(
+                "turn.error",
+                normalized_request_id,
+                normalized_conversation_id,
+                error={"code": "invalid_request", "message": "request_id is required"},
+            )
+            await self._send(
+                "turn.done",
+                normalized_request_id,
+                normalized_conversation_id,
+                failed=True,
+            )
+            return
+
+        if command_name == "shutdown":
+            await self._shutdown(request_id, conversation_id)
+            return
+
+        if not await self._begin_dispatch():
+            await self._fail(
+                request_id,
+                conversation_id if isinstance(conversation_id, str) else None,
+                "shutting_down",
+                "the sidecar is shutting down",
+            )
+            return
+        try:
+            if command_name == "health":
+                await self._send(
+                    "ready",
+                    request_id,
+                    conversation_id if isinstance(conversation_id, str) else None,
+                    status="ok",
+                )
+                return
+            if not isinstance(conversation_id, str) or not conversation_id:
+                await self._fail(
+                    request_id, None, "invalid_request", "conversation_id is required"
+                )
+                return
+
+            await self.evict_idle_sessions(
+                exclude=conversation_id,
+                reserve=1 if command_name == "turn.start" else 0,
+            )
+            async with self._conversation_command_lock(conversation_id):
+                if command_name == "turn.start":
+                    prompt = command.get("prompt")
+                    if not isinstance(prompt, str) or not prompt:
+                        await self._fail(
+                            request_id,
+                            conversation_id,
+                            "invalid_request",
+                            "prompt is required",
+                        )
+                        return
+                    await self._start_turn(request_id, conversation_id, prompt)
+                elif command_name == "turn.abort":
+                    await self._abort_turn(request_id, conversation_id)
+                elif command_name == "conversation.delete":
+                    await self._delete_conversation(request_id, conversation_id)
+                else:
+                    await self._fail(
+                        request_id,
+                        conversation_id,
+                        "unknown_command",
+                        f"unsupported command: {command_name!r}",
+                    )
+        finally:
+            await self._end_dispatch()
+
+    @asynccontextmanager
+    async def _conversation_command_lock(
+        self, conversation_id: str
+    ) -> AsyncIterator[None]:
+        async with self._conversation_gates_lock:
+            gate = self._conversation_gates.get(conversation_id)
+            if gate is None:
+                gate = ConversationGate(lock=asyncio.Lock())
+                self._conversation_gates[conversation_id] = gate
+            gate.users += 1
+        try:
+            async with gate.lock:
+                yield
+        finally:
+            async with self._conversation_gates_lock:
+                gate.users -= 1
+                if (
+                    gate.users == 0
+                    and self._conversation_gates.get(conversation_id) is gate
+                ):
+                    del self._conversation_gates[conversation_id]
+
+    async def _begin_dispatch(self) -> bool:
+        async with self._dispatch_condition:
+            if self.shutting_down:
+                return False
+            self._active_dispatches += 1
+            return True
+
+    async def _end_dispatch(self) -> None:
+        async with self._dispatch_condition:
+            self._active_dispatches -= 1
+            if self._active_dispatches == 0:
+                self._dispatch_condition.notify_all()
+
+    async def _session_options(self) -> JsonObject:
+        provider: ProviderConfig = {
+            "type": "openai",
+            "base_url": self._config.base_url,
+            "wire_api": self._config.wire_api,
+        }
+        if self._config.api_key is not None:
+            provider["api_key"] = self._config.api_key
+        return {
+            "on_permission_request": deny_permission,
+            "model": self._config.model,
+            "provider": provider,
+            "streaming": True,
+            "tools": [],
+            "available_tools": [],
+            "mcp_servers": {},
+            "enable_config_discovery": False,
+            "skip_custom_instructions": True,
+            "enable_skills": False,
+            "enable_session_store": True,
+        }
+
+    async def _get_conversation(self, conversation_id: str) -> Conversation:
+        async with self._conversations_lock:
+            existing = self._conversations.get(conversation_id)
+            if existing is not None:
+                existing.last_used = time.monotonic()
+                return existing
+            if len(self._conversations) >= self._config.max_sessions:
+                raise RuntimeError("Copilot session capacity exhausted")
+
+            options = await self._session_options()
+            try:
+                session = await self._client.resume_session(conversation_id, **options)
+            except Exception:
+                session = await self._client.create_session(
+                    session_id=conversation_id, **options
+                )
+            if session is None:
+                session = await self._client.create_session(
+                    session_id=conversation_id, **options
+                )
+
+            def handle_event(event: Any) -> None:
+                task = asyncio.create_task(self._handle_event(conversation_id, event))
+                self._event_tasks.add(task)
+                task.add_done_callback(self._event_tasks.discard)
+
+            unsubscribe = session.on(handle_event)
+            conversation = Conversation(session=session, unsubscribe=unsubscribe)
+            self._conversations[conversation_id] = conversation
+            return conversation
+
+    async def _cleanup_loop(self) -> None:
+        interval = min(60, self._config.idle_timeout_seconds)
+        try:
+            while not self.shutting_down:
+                await asyncio.sleep(interval)
+                try:
+                    await self.evict_idle_sessions()
+                except Exception as error:
+                    print(
+                        f"copilot-sidecar: session cleanup failed: {error}",
+                        file=sys.stderr,
+                    )
+        except asyncio.CancelledError:
+            return
+
+    async def evict_idle_sessions(
+        self,
+        exclude: str | None = None,
+        reserve: int = 0,
+    ) -> None:
+        now = time.monotonic()
+        async with self._conversations_lock:
+            inactive = sorted(
+                (
+                    (conversation.last_used, conversation_id)
+                    for conversation_id, conversation
+                    in self._conversations.items()
+                    if conversation.active is None
+                    and conversation_id != exclude
+                )
+            )
+            overflow = max(
+                0,
+                len(self._conversations) + reserve -
+                self._config.max_sessions,
+            )
+            candidates = {
+                conversation_id
+                for index, (last_used, conversation_id) in enumerate(inactive)
+                if index < overflow
+                or now - last_used >= self._config.idle_timeout_seconds
+            }
+
+        for conversation_id in candidates:
+            async with self._conversation_command_lock(conversation_id):
+                async with self._conversations_lock:
+                    conversation = self._conversations.get(conversation_id)
+                    if conversation is None or conversation.active is not None:
+                        continue
+                    over_limit = (
+                        len(self._conversations) + reserve >
+                        self._config.max_sessions
+                    )
+                    expired = (
+                        time.monotonic() - conversation.last_used
+                        >= self._config.idle_timeout_seconds
+                    )
+                    if not over_limit and not expired:
+                        continue
+                    self._conversations.pop(conversation_id, None)
+                try:
+                    conversation.unsubscribe()
+                except Exception as error:
+                    print(
+                        f"copilot-sidecar: unsubscribe failed: {error}",
+                        file=sys.stderr,
+                    )
+                try:
+                    await conversation.session.disconnect()
+                except Exception as error:
+                    print(
+                        f"copilot-sidecar: disconnect failed: {error}",
+                        file=sys.stderr,
+                    )
+
+    async def _start_turn(
+        self, request_id: str, conversation_id: str, prompt: str
+    ) -> None:
+        try:
+            conversation = await self._get_conversation(conversation_id)
+            if conversation.active is not None and not conversation.active.done:
+                await self._fail(
+                    request_id,
+                    conversation_id,
+                    "turn_in_progress",
+                    "the conversation already has an active turn",
+                )
+                return
+            conversation.active = Turn(request_id=request_id)
+            await self._send("turn.accepted", request_id, conversation_id)
+            await conversation.session.send(prompt)
+        except Exception as error:
+            await self._finish_with_error(request_id, conversation_id, error)
+
+    async def _abort_turn(self, request_id: str, conversation_id: str) -> None:
+        conversation = self._conversations.get(conversation_id)
+        if conversation is None or conversation.active is None:
+            await self._fail(
+                request_id, conversation_id, "no_active_turn", "no active turn to abort"
+            )
+            return
+        try:
+            target_request_id = conversation.active.request_id
+            await self._send(
+                "turn.accepted",
+                request_id,
+                conversation_id,
+                action="abort",
+                target_request_id=target_request_id,
+            )
+            await conversation.session.abort()
+            await self._send(
+                "turn.done",
+                request_id,
+                conversation_id,
+                action="abort",
+                target_request_id=target_request_id,
+            )
+        except Exception as error:
+            await self._fail(request_id, conversation_id, "sdk_error", str(error))
+
+    async def _delete_conversation(
+        self, request_id: str, conversation_id: str
+    ) -> None:
+        try:
+            async with self._conversations_lock:
+                conversation = self._conversations.pop(conversation_id, None)
+            if conversation is not None:
+                if conversation.active is not None and not conversation.active.done:
+                    await self._finish_turn(
+                        conversation_id, conversation.active, aborted=True, deleted=True
+                    )
+                conversation.unsubscribe()
+                await conversation.session.disconnect()
+                session_id = conversation.session.session_id
+            else:
+                session_id = conversation_id
+            await self._client.delete_session(session_id)
+            await self._send(
+                "turn.done", request_id, conversation_id, action="conversation.delete"
+            )
+        except Exception as error:
+            await self._finish_with_error(request_id, conversation_id, error)
+
+    async def _handle_event(self, conversation_id: str, event: Any) -> None:
+        conversation = self._conversations.get(conversation_id)
+        if conversation is None or conversation.active is None:
+            return
+        turn = conversation.active
+        event_type = getattr(getattr(event, "type", None), "value", None)
+        data = getattr(event, "data", None)
+
+        if event_type == "assistant.message_delta":
+            await self._send(
+                "assistant.delta",
+                turn.request_id,
+                conversation_id,
+                delta=getattr(data, "delta_content", ""),
+                message_id=getattr(data, "message_id", None),
+            )
+        elif event_type == "assistant.message":
+            await self._send(
+                "assistant.completed",
+                turn.request_id,
+                conversation_id,
+                content=getattr(data, "content", ""),
+                message_id=getattr(data, "message_id", None),
+                model=getattr(data, "model", None),
+            )
+        elif event_type == "assistant.usage":
+            usage = {
+                name: getattr(data, name, None)
+                for name in (
+                    "model",
+                    "input_tokens",
+                    "output_tokens",
+                    "reasoning_tokens",
+                    "cache_read_tokens",
+                    "cache_write_tokens",
+                    "finish_reason",
+                    "cost",
+                )
+                if getattr(data, name, None) is not None
+            }
+            await self._send(
+                "assistant.usage",
+                turn.request_id,
+                conversation_id,
+                usage=usage,
+            )
+        elif event_type == "session.error":
+            error = {
+                "code": getattr(data, "error_code", None)
+                or getattr(data, "error_type", "sdk_error"),
+                "message": getattr(data, "message", "Copilot session error"),
+            }
+            if getattr(data, "status_code", None) is not None:
+                error["status_code"] = data.status_code
+            await self._send(
+                "turn.error",
+                turn.request_id,
+                conversation_id,
+                error=error,
+            )
+            await self._finish_turn(conversation_id, turn, failed=True)
+        elif event_type == "session.idle":
+            await self._finish_turn(
+                conversation_id, turn, aborted=bool(getattr(data, "aborted", False))
+            )
+
+    async def _finish_with_error(
+        self, request_id: str, conversation_id: str | None, error: Exception
+    ) -> None:
+        await self._send(
+            "turn.error",
+            request_id,
+            conversation_id,
+            error={"code": "sdk_error", "message": str(error)},
+        )
+        conversation = (
+            self._conversations.get(conversation_id)
+            if conversation_id is not None
+            else None
+        )
+        if conversation is not None and conversation.active is not None:
+            await self._finish_turn(conversation_id, conversation.active, failed=True)
+        else:
+            await self._send("turn.done", request_id, conversation_id, failed=True)
+
+    async def _fail(
+        self,
+        request_id: str,
+        conversation_id: str | None,
+        code: str,
+        message: str,
+    ) -> None:
+        await self._send(
+            "turn.error",
+            request_id,
+            conversation_id,
+            error={"code": code, "message": message},
+        )
+        await self._send("turn.done", request_id, conversation_id, failed=True)
+
+    async def _finish_turn(
+        self, conversation_id: str, turn: Turn, **fields: Any
+    ) -> None:
+        if turn.done:
+            return
+        turn.done = True
+        await self._send("turn.done", turn.request_id, conversation_id, **fields)
+        conversation = self._conversations.get(conversation_id)
+        if conversation is not None and conversation.active is turn:
+            conversation.active = None
+            conversation.last_used = time.monotonic()
+
+    async def _shutdown(
+        self, request_id: str, conversation_id: Any
+    ) -> None:
+        async with self._dispatch_condition:
+            if self.shutting_down:
+                await self._send(
+                    "turn.done",
+                    request_id,
+                    conversation_id if isinstance(conversation_id, str) else None,
+                    action="shutdown",
+                    already_in_progress=True,
+                )
+                return
+            self.shutting_down = True
+            if self._cleanup_task is not None:
+                self._cleanup_task.cancel()
+                try:
+                    await self._cleanup_task
+                except asyncio.CancelledError:
+                    pass
+                self._cleanup_task = None
+            while self._active_dispatches:
+                await self._dispatch_condition.wait()
+
+        await self.drain_events()
+        conversations = list(self._conversations.items())
+        self._conversations.clear()
+        shutdown_errors: list[str] = []
+        for item_conversation_id, conversation in conversations:
+            conversation.unsubscribe()
+            try:
+                if conversation.active is not None:
+                    await self._finish_turn(
+                        item_conversation_id,
+                        conversation.active,
+                        aborted=True,
+                        shutdown=True,
+                    )
+                await conversation.session.disconnect()
+            except Exception as error:
+                shutdown_errors.append(str(error))
+        if self._started:
+            try:
+                await self._client.stop()
+            except Exception as error:
+                shutdown_errors.append(str(error))
+            self._started = False
+        if shutdown_errors:
+            await self._send(
+                "turn.error",
+                request_id,
+                conversation_id if isinstance(conversation_id, str) else None,
+                error={
+                    "code": "shutdown_error",
+                    "message": "; ".join(shutdown_errors),
+                },
+            )
+        await self._send(
+            "turn.done",
+            request_id,
+            conversation_id if isinstance(conversation_id, str) else None,
+            action="shutdown",
+            failed=bool(shutdown_errors),
+        )
+
+    async def drain_events(self) -> None:
+        while self._event_tasks:
+            await asyncio.gather(*tuple(self._event_tasks))
+
+    async def _send(
+        self,
+        event_type: str,
+        request_id: str | None,
+        conversation_id: str | None,
+        **fields: Any,
+    ) -> None:
+        await self._emit(
+            {
+                "type": event_type,
+                "request_id": request_id,
+                "conversation_id": conversation_id,
+                **fields,
+            }
+        )
+
+
+class NdjsonWriter:
+    def __init__(self) -> None:
+        self._lock = asyncio.Lock()
+
+    async def __call__(self, payload: JsonObject) -> None:
+        encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
+        async with self._lock:
+            sys.stdout.write(encoded + "\n")
+            sys.stdout.flush()
+
+
+def build_client(cli_path: str, base_directory: str) -> CopilotClient:
+    resolved_path = os.path.abspath(cli_path)
+    if not os.path.isfile(resolved_path):
+        raise ValueError(f"Copilot CLI does not exist: {resolved_path}")
+    if not os.access(resolved_path, os.X_OK):
+        raise ValueError(f"Copilot CLI is not executable: {resolved_path}")
+    os.makedirs(base_directory, mode=0o700, exist_ok=True)
+    return CopilotClient(
+        connection=RuntimeConnection.for_stdio(path=resolved_path),
+        base_directory=base_directory,
+        use_logged_in_user=False,
+        log_level="error",
+        mode="empty",
+    )
+
+
+async def run(cli_path: str) -> int:
+    writer = NdjsonWriter()
+    config = SidecarConfig.from_environment()
+    sidecar = Sidecar(build_client(cli_path, config.base_directory), config, writer)
+    await sidecar.start()
+    await sidecar.announce_ready()
+
+    tasks: set[asyncio.Task[None]] = set()
+    try:
+        while not sidecar.shutting_down:
+            line = await asyncio.to_thread(sys.stdin.readline)
+            if not line:
+                break
+            try:
+                command = json.loads(line)
+                if not isinstance(command, dict):
+                    raise ValueError("command must be a JSON object")
+            except (json.JSONDecodeError, ValueError) as error:
+                for event_type in ("turn.error", "turn.done"):
+                    payload: JsonObject = {
+                        "type": event_type,
+                        "request_id": None,
+                        "conversation_id": None,
+                    }
+                    if event_type == "turn.error":
+                        payload["error"] = {
+                            "code": "invalid_json",
+                            "message": str(error),
+                        }
+                    else:
+                        payload["failed"] = True
+                    await writer(payload)
+                continue
+            task = asyncio.create_task(sidecar.dispatch(command))
+            tasks.add(task)
+            task.add_done_callback(tasks.discard)
+            if command.get("command") == "shutdown":
+                await task
+                break
+        if tasks:
+            await asyncio.gather(*tasks)
+        await sidecar.drain_events()
+    finally:
+        if not sidecar.shutting_down:
+            await sidecar.dispatch(
+                {
+                    "command": "shutdown",
+                    "request_id": "stdin-eof",
+                    "conversation_id": None,
+                }
+            )
+    return 0
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Copilot SDK NDJSON sidecar")
+    parser.add_argument("copilot_cli", help="path to the Bazel-pinned Copilot CLI")
+    args = parser.parse_args()
+    try:
+        raise SystemExit(asyncio.run(run(args.copilot_cli)))
+    except (OSError, ValueError, RuntimeError) as error:
+        print(f"copilot-sidecar: {error}", file=sys.stderr)
+        raise SystemExit(2) from error
+
+
+if __name__ == "__main__":
+    main()