Mercurial
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()