Mercurial

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

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
comparison
equal deleted inserted replaced
259:667156fcd3e3 260:1f9877b637e9
1 from __future__ import annotations
2
3 import argparse
4 import asyncio
5 import json
6 import os
7 import sys
8 import time
9 from contextlib import asynccontextmanager
10 from dataclasses import dataclass, field
11 from typing import Any, AsyncIterator, Awaitable, Callable, Protocol
12
13 from copilot import CopilotClient, ProviderConfig, RuntimeConnection
14 from copilot.rpc import PermissionDecisionReject
15
16
17 JsonObject = dict[str, Any]
18 Emit = Callable[[JsonObject], Awaitable[None]]
19
20
21 class Session(Protocol):
22 session_id: str
23
24 def on(self, handler: Callable[[Any], None]) -> Callable[[], None]: ...
25
26 async def send(self, prompt: str) -> str: ...
27
28 async def abort(self) -> None: ...
29
30 async def disconnect(self) -> None: ...
31
32
33 class Client(Protocol):
34 async def start(self) -> None: ...
35
36 async def stop(self) -> None: ...
37
38 async def create_session(self, **kwargs: Any) -> Session: ...
39
40 async def resume_session(self, session_id: str, **kwargs: Any) -> Session: ...
41
42 async def delete_session(self, session_id: str) -> None: ...
43
44
45 @dataclass(frozen=True)
46 class SidecarConfig:
47 base_url: str
48 model: str
49 wire_api: str
50 base_directory: str
51 api_key: str | None = None
52 idle_timeout_seconds: int = 900
53 max_sessions: int = 64
54
55 @classmethod
56 def from_environment(cls) -> SidecarConfig:
57 base_url = os.environ.get("LITELLM_BASE_URL", "").strip()
58 model = os.environ.get("LITELLM_MODEL", "").strip()
59 wire_api = os.environ.get("LITELLM_WIRE_API", "").strip()
60 base_directory = os.environ.get("COPILOT_SIDECAR_HOME", "").strip()
61 missing = [
62 name
63 for name, value in (
64 ("LITELLM_BASE_URL", base_url),
65 ("LITELLM_MODEL", model),
66 ("LITELLM_WIRE_API", wire_api),
67 ("COPILOT_SIDECAR_HOME", base_directory),
68 )
69 if not value
70 ]
71 if missing:
72 raise ValueError(f"missing required environment: {', '.join(missing)}")
73 if wire_api not in ("completions", "responses"):
74 raise ValueError("LITELLM_WIRE_API must be 'completions' or 'responses'")
75 return cls(
76 base_url=base_url,
77 model=model,
78 wire_api=wire_api,
79 base_directory=os.path.abspath(base_directory),
80 api_key=os.environ.get("LITELLM_API_KEY") or None,
81 idle_timeout_seconds=max(
82 1,
83 int(os.environ.get("COPILOT_SESSION_IDLE_SECONDS", "900")),
84 ),
85 max_sessions=max(
86 1,
87 int(os.environ.get("COPILOT_MAX_SESSIONS", "64")),
88 ),
89 )
90
91
92 @dataclass
93 class Turn:
94 request_id: str
95 done: bool = False
96
97
98 @dataclass
99 class Conversation:
100 session: Session
101 unsubscribe: Callable[[], None]
102 active: Turn | None = None
103 last_used: float = field(default_factory=time.monotonic)
104
105
106 @dataclass
107 class ConversationGate:
108 lock: asyncio.Lock
109 users: int = 0
110
111
112 def deny_permission(*_args: Any, **_kwargs: Any) -> PermissionDecisionReject:
113 return PermissionDecisionReject(feedback="The inference sidecar denies all permissions.")
114
115
116 class Sidecar:
117 def __init__(self, client: Client, config: SidecarConfig, emit: Emit):
118 self._client = client
119 self._config = config
120 self._emit = emit
121 self._conversations: dict[str, Conversation] = {}
122 self._conversations_lock = asyncio.Lock()
123 self._conversation_gates: dict[str, ConversationGate] = {}
124 self._conversation_gates_lock = asyncio.Lock()
125 self._dispatch_condition = asyncio.Condition()
126 self._active_dispatches = 0
127 self._event_tasks: set[asyncio.Task[None]] = set()
128 self._cleanup_task: asyncio.Task[None] | None = None
129 self._started = False
130 self.shutting_down = False
131
132 async def start(self) -> None:
133 if not self._started:
134 await self._client.start()
135 self._started = True
136 self._cleanup_task = asyncio.create_task(self._cleanup_loop())
137
138 async def announce_ready(self) -> None:
139 await self._send("ready", None, None, status="ok")
140
141 async def dispatch(self, command: JsonObject) -> None:
142 command_name = command.get("command")
143 request_id = command.get("request_id")
144 conversation_id = command.get("conversation_id")
145 if not isinstance(request_id, str) or not request_id:
146 normalized_request_id = request_id if isinstance(request_id, str) else None
147 normalized_conversation_id = (
148 conversation_id if isinstance(conversation_id, str) else None
149 )
150 await self._send(
151 "turn.error",
152 normalized_request_id,
153 normalized_conversation_id,
154 error={"code": "invalid_request", "message": "request_id is required"},
155 )
156 await self._send(
157 "turn.done",
158 normalized_request_id,
159 normalized_conversation_id,
160 failed=True,
161 )
162 return
163
164 if command_name == "shutdown":
165 await self._shutdown(request_id, conversation_id)
166 return
167
168 if not await self._begin_dispatch():
169 await self._fail(
170 request_id,
171 conversation_id if isinstance(conversation_id, str) else None,
172 "shutting_down",
173 "the sidecar is shutting down",
174 )
175 return
176 try:
177 if command_name == "health":
178 await self._send(
179 "ready",
180 request_id,
181 conversation_id if isinstance(conversation_id, str) else None,
182 status="ok",
183 )
184 return
185 if not isinstance(conversation_id, str) or not conversation_id:
186 await self._fail(
187 request_id, None, "invalid_request", "conversation_id is required"
188 )
189 return
190
191 await self.evict_idle_sessions(
192 exclude=conversation_id,
193 reserve=1 if command_name == "turn.start" else 0,
194 )
195 async with self._conversation_command_lock(conversation_id):
196 if command_name == "turn.start":
197 prompt = command.get("prompt")
198 if not isinstance(prompt, str) or not prompt:
199 await self._fail(
200 request_id,
201 conversation_id,
202 "invalid_request",
203 "prompt is required",
204 )
205 return
206 await self._start_turn(request_id, conversation_id, prompt)
207 elif command_name == "turn.abort":
208 await self._abort_turn(request_id, conversation_id)
209 elif command_name == "conversation.delete":
210 await self._delete_conversation(request_id, conversation_id)
211 else:
212 await self._fail(
213 request_id,
214 conversation_id,
215 "unknown_command",
216 f"unsupported command: {command_name!r}",
217 )
218 finally:
219 await self._end_dispatch()
220
221 @asynccontextmanager
222 async def _conversation_command_lock(
223 self, conversation_id: str
224 ) -> AsyncIterator[None]:
225 async with self._conversation_gates_lock:
226 gate = self._conversation_gates.get(conversation_id)
227 if gate is None:
228 gate = ConversationGate(lock=asyncio.Lock())
229 self._conversation_gates[conversation_id] = gate
230 gate.users += 1
231 try:
232 async with gate.lock:
233 yield
234 finally:
235 async with self._conversation_gates_lock:
236 gate.users -= 1
237 if (
238 gate.users == 0
239 and self._conversation_gates.get(conversation_id) is gate
240 ):
241 del self._conversation_gates[conversation_id]
242
243 async def _begin_dispatch(self) -> bool:
244 async with self._dispatch_condition:
245 if self.shutting_down:
246 return False
247 self._active_dispatches += 1
248 return True
249
250 async def _end_dispatch(self) -> None:
251 async with self._dispatch_condition:
252 self._active_dispatches -= 1
253 if self._active_dispatches == 0:
254 self._dispatch_condition.notify_all()
255
256 async def _session_options(self) -> JsonObject:
257 provider: ProviderConfig = {
258 "type": "openai",
259 "base_url": self._config.base_url,
260 "wire_api": self._config.wire_api,
261 }
262 if self._config.api_key is not None:
263 provider["api_key"] = self._config.api_key
264 return {
265 "on_permission_request": deny_permission,
266 "model": self._config.model,
267 "provider": provider,
268 "streaming": True,
269 "tools": [],
270 "available_tools": [],
271 "mcp_servers": {},
272 "enable_config_discovery": False,
273 "skip_custom_instructions": True,
274 "enable_skills": False,
275 "enable_session_store": True,
276 }
277
278 async def _get_conversation(self, conversation_id: str) -> Conversation:
279 async with self._conversations_lock:
280 existing = self._conversations.get(conversation_id)
281 if existing is not None:
282 existing.last_used = time.monotonic()
283 return existing
284 if len(self._conversations) >= self._config.max_sessions:
285 raise RuntimeError("Copilot session capacity exhausted")
286
287 options = await self._session_options()
288 try:
289 session = await self._client.resume_session(conversation_id, **options)
290 except Exception:
291 session = await self._client.create_session(
292 session_id=conversation_id, **options
293 )
294 if session is None:
295 session = await self._client.create_session(
296 session_id=conversation_id, **options
297 )
298
299 def handle_event(event: Any) -> None:
300 task = asyncio.create_task(self._handle_event(conversation_id, event))
301 self._event_tasks.add(task)
302 task.add_done_callback(self._event_tasks.discard)
303
304 unsubscribe = session.on(handle_event)
305 conversation = Conversation(session=session, unsubscribe=unsubscribe)
306 self._conversations[conversation_id] = conversation
307 return conversation
308
309 async def _cleanup_loop(self) -> None:
310 interval = min(60, self._config.idle_timeout_seconds)
311 try:
312 while not self.shutting_down:
313 await asyncio.sleep(interval)
314 try:
315 await self.evict_idle_sessions()
316 except Exception as error:
317 print(
318 f"copilot-sidecar: session cleanup failed: {error}",
319 file=sys.stderr,
320 )
321 except asyncio.CancelledError:
322 return
323
324 async def evict_idle_sessions(
325 self,
326 exclude: str | None = None,
327 reserve: int = 0,
328 ) -> None:
329 now = time.monotonic()
330 async with self._conversations_lock:
331 inactive = sorted(
332 (
333 (conversation.last_used, conversation_id)
334 for conversation_id, conversation
335 in self._conversations.items()
336 if conversation.active is None
337 and conversation_id != exclude
338 )
339 )
340 overflow = max(
341 0,
342 len(self._conversations) + reserve -
343 self._config.max_sessions,
344 )
345 candidates = {
346 conversation_id
347 for index, (last_used, conversation_id) in enumerate(inactive)
348 if index < overflow
349 or now - last_used >= self._config.idle_timeout_seconds
350 }
351
352 for conversation_id in candidates:
353 async with self._conversation_command_lock(conversation_id):
354 async with self._conversations_lock:
355 conversation = self._conversations.get(conversation_id)
356 if conversation is None or conversation.active is not None:
357 continue
358 over_limit = (
359 len(self._conversations) + reserve >
360 self._config.max_sessions
361 )
362 expired = (
363 time.monotonic() - conversation.last_used
364 >= self._config.idle_timeout_seconds
365 )
366 if not over_limit and not expired:
367 continue
368 self._conversations.pop(conversation_id, None)
369 try:
370 conversation.unsubscribe()
371 except Exception as error:
372 print(
373 f"copilot-sidecar: unsubscribe failed: {error}",
374 file=sys.stderr,
375 )
376 try:
377 await conversation.session.disconnect()
378 except Exception as error:
379 print(
380 f"copilot-sidecar: disconnect failed: {error}",
381 file=sys.stderr,
382 )
383
384 async def _start_turn(
385 self, request_id: str, conversation_id: str, prompt: str
386 ) -> None:
387 try:
388 conversation = await self._get_conversation(conversation_id)
389 if conversation.active is not None and not conversation.active.done:
390 await self._fail(
391 request_id,
392 conversation_id,
393 "turn_in_progress",
394 "the conversation already has an active turn",
395 )
396 return
397 conversation.active = Turn(request_id=request_id)
398 await self._send("turn.accepted", request_id, conversation_id)
399 await conversation.session.send(prompt)
400 except Exception as error:
401 await self._finish_with_error(request_id, conversation_id, error)
402
403 async def _abort_turn(self, request_id: str, conversation_id: str) -> None:
404 conversation = self._conversations.get(conversation_id)
405 if conversation is None or conversation.active is None:
406 await self._fail(
407 request_id, conversation_id, "no_active_turn", "no active turn to abort"
408 )
409 return
410 try:
411 target_request_id = conversation.active.request_id
412 await self._send(
413 "turn.accepted",
414 request_id,
415 conversation_id,
416 action="abort",
417 target_request_id=target_request_id,
418 )
419 await conversation.session.abort()
420 await self._send(
421 "turn.done",
422 request_id,
423 conversation_id,
424 action="abort",
425 target_request_id=target_request_id,
426 )
427 except Exception as error:
428 await self._fail(request_id, conversation_id, "sdk_error", str(error))
429
430 async def _delete_conversation(
431 self, request_id: str, conversation_id: str
432 ) -> None:
433 try:
434 async with self._conversations_lock:
435 conversation = self._conversations.pop(conversation_id, None)
436 if conversation is not None:
437 if conversation.active is not None and not conversation.active.done:
438 await self._finish_turn(
439 conversation_id, conversation.active, aborted=True, deleted=True
440 )
441 conversation.unsubscribe()
442 await conversation.session.disconnect()
443 session_id = conversation.session.session_id
444 else:
445 session_id = conversation_id
446 await self._client.delete_session(session_id)
447 await self._send(
448 "turn.done", request_id, conversation_id, action="conversation.delete"
449 )
450 except Exception as error:
451 await self._finish_with_error(request_id, conversation_id, error)
452
453 async def _handle_event(self, conversation_id: str, event: Any) -> None:
454 conversation = self._conversations.get(conversation_id)
455 if conversation is None or conversation.active is None:
456 return
457 turn = conversation.active
458 event_type = getattr(getattr(event, "type", None), "value", None)
459 data = getattr(event, "data", None)
460
461 if event_type == "assistant.message_delta":
462 await self._send(
463 "assistant.delta",
464 turn.request_id,
465 conversation_id,
466 delta=getattr(data, "delta_content", ""),
467 message_id=getattr(data, "message_id", None),
468 )
469 elif event_type == "assistant.message":
470 await self._send(
471 "assistant.completed",
472 turn.request_id,
473 conversation_id,
474 content=getattr(data, "content", ""),
475 message_id=getattr(data, "message_id", None),
476 model=getattr(data, "model", None),
477 )
478 elif event_type == "assistant.usage":
479 usage = {
480 name: getattr(data, name, None)
481 for name in (
482 "model",
483 "input_tokens",
484 "output_tokens",
485 "reasoning_tokens",
486 "cache_read_tokens",
487 "cache_write_tokens",
488 "finish_reason",
489 "cost",
490 )
491 if getattr(data, name, None) is not None
492 }
493 await self._send(
494 "assistant.usage",
495 turn.request_id,
496 conversation_id,
497 usage=usage,
498 )
499 elif event_type == "session.error":
500 error = {
501 "code": getattr(data, "error_code", None)
502 or getattr(data, "error_type", "sdk_error"),
503 "message": getattr(data, "message", "Copilot session error"),
504 }
505 if getattr(data, "status_code", None) is not None:
506 error["status_code"] = data.status_code
507 await self._send(
508 "turn.error",
509 turn.request_id,
510 conversation_id,
511 error=error,
512 )
513 await self._finish_turn(conversation_id, turn, failed=True)
514 elif event_type == "session.idle":
515 await self._finish_turn(
516 conversation_id, turn, aborted=bool(getattr(data, "aborted", False))
517 )
518
519 async def _finish_with_error(
520 self, request_id: str, conversation_id: str | None, error: Exception
521 ) -> None:
522 await self._send(
523 "turn.error",
524 request_id,
525 conversation_id,
526 error={"code": "sdk_error", "message": str(error)},
527 )
528 conversation = (
529 self._conversations.get(conversation_id)
530 if conversation_id is not None
531 else None
532 )
533 if conversation is not None and conversation.active is not None:
534 await self._finish_turn(conversation_id, conversation.active, failed=True)
535 else:
536 await self._send("turn.done", request_id, conversation_id, failed=True)
537
538 async def _fail(
539 self,
540 request_id: str,
541 conversation_id: str | None,
542 code: str,
543 message: str,
544 ) -> None:
545 await self._send(
546 "turn.error",
547 request_id,
548 conversation_id,
549 error={"code": code, "message": message},
550 )
551 await self._send("turn.done", request_id, conversation_id, failed=True)
552
553 async def _finish_turn(
554 self, conversation_id: str, turn: Turn, **fields: Any
555 ) -> None:
556 if turn.done:
557 return
558 turn.done = True
559 await self._send("turn.done", turn.request_id, conversation_id, **fields)
560 conversation = self._conversations.get(conversation_id)
561 if conversation is not None and conversation.active is turn:
562 conversation.active = None
563 conversation.last_used = time.monotonic()
564
565 async def _shutdown(
566 self, request_id: str, conversation_id: Any
567 ) -> None:
568 async with self._dispatch_condition:
569 if self.shutting_down:
570 await self._send(
571 "turn.done",
572 request_id,
573 conversation_id if isinstance(conversation_id, str) else None,
574 action="shutdown",
575 already_in_progress=True,
576 )
577 return
578 self.shutting_down = True
579 if self._cleanup_task is not None:
580 self._cleanup_task.cancel()
581 try:
582 await self._cleanup_task
583 except asyncio.CancelledError:
584 pass
585 self._cleanup_task = None
586 while self._active_dispatches:
587 await self._dispatch_condition.wait()
588
589 await self.drain_events()
590 conversations = list(self._conversations.items())
591 self._conversations.clear()
592 shutdown_errors: list[str] = []
593 for item_conversation_id, conversation in conversations:
594 conversation.unsubscribe()
595 try:
596 if conversation.active is not None:
597 await self._finish_turn(
598 item_conversation_id,
599 conversation.active,
600 aborted=True,
601 shutdown=True,
602 )
603 await conversation.session.disconnect()
604 except Exception as error:
605 shutdown_errors.append(str(error))
606 if self._started:
607 try:
608 await self._client.stop()
609 except Exception as error:
610 shutdown_errors.append(str(error))
611 self._started = False
612 if shutdown_errors:
613 await self._send(
614 "turn.error",
615 request_id,
616 conversation_id if isinstance(conversation_id, str) else None,
617 error={
618 "code": "shutdown_error",
619 "message": "; ".join(shutdown_errors),
620 },
621 )
622 await self._send(
623 "turn.done",
624 request_id,
625 conversation_id if isinstance(conversation_id, str) else None,
626 action="shutdown",
627 failed=bool(shutdown_errors),
628 )
629
630 async def drain_events(self) -> None:
631 while self._event_tasks:
632 await asyncio.gather(*tuple(self._event_tasks))
633
634 async def _send(
635 self,
636 event_type: str,
637 request_id: str | None,
638 conversation_id: str | None,
639 **fields: Any,
640 ) -> None:
641 await self._emit(
642 {
643 "type": event_type,
644 "request_id": request_id,
645 "conversation_id": conversation_id,
646 **fields,
647 }
648 )
649
650
651 class NdjsonWriter:
652 def __init__(self) -> None:
653 self._lock = asyncio.Lock()
654
655 async def __call__(self, payload: JsonObject) -> None:
656 encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
657 async with self._lock:
658 sys.stdout.write(encoded + "\n")
659 sys.stdout.flush()
660
661
662 def build_client(cli_path: str, base_directory: str) -> CopilotClient:
663 resolved_path = os.path.abspath(cli_path)
664 if not os.path.isfile(resolved_path):
665 raise ValueError(f"Copilot CLI does not exist: {resolved_path}")
666 if not os.access(resolved_path, os.X_OK):
667 raise ValueError(f"Copilot CLI is not executable: {resolved_path}")
668 os.makedirs(base_directory, mode=0o700, exist_ok=True)
669 return CopilotClient(
670 connection=RuntimeConnection.for_stdio(path=resolved_path),
671 base_directory=base_directory,
672 use_logged_in_user=False,
673 log_level="error",
674 mode="empty",
675 )
676
677
678 async def run(cli_path: str) -> int:
679 writer = NdjsonWriter()
680 config = SidecarConfig.from_environment()
681 sidecar = Sidecar(build_client(cli_path, config.base_directory), config, writer)
682 await sidecar.start()
683 await sidecar.announce_ready()
684
685 tasks: set[asyncio.Task[None]] = set()
686 try:
687 while not sidecar.shutting_down:
688 line = await asyncio.to_thread(sys.stdin.readline)
689 if not line:
690 break
691 try:
692 command = json.loads(line)
693 if not isinstance(command, dict):
694 raise ValueError("command must be a JSON object")
695 except (json.JSONDecodeError, ValueError) as error:
696 for event_type in ("turn.error", "turn.done"):
697 payload: JsonObject = {
698 "type": event_type,
699 "request_id": None,
700 "conversation_id": None,
701 }
702 if event_type == "turn.error":
703 payload["error"] = {
704 "code": "invalid_json",
705 "message": str(error),
706 }
707 else:
708 payload["failed"] = True
709 await writer(payload)
710 continue
711 task = asyncio.create_task(sidecar.dispatch(command))
712 tasks.add(task)
713 task.add_done_callback(tasks.discard)
714 if command.get("command") == "shutdown":
715 await task
716 break
717 if tasks:
718 await asyncio.gather(*tasks)
719 await sidecar.drain_events()
720 finally:
721 if not sidecar.shutting_down:
722 await sidecar.dispatch(
723 {
724 "command": "shutdown",
725 "request_id": "stdin-eof",
726 "conversation_id": None,
727 }
728 )
729 return 0
730
731
732 def main() -> None:
733 parser = argparse.ArgumentParser(description="Copilot SDK NDJSON sidecar")
734 parser.add_argument("copilot_cli", help="path to the Bazel-pinned Copilot CLI")
735 args = parser.parse_args()
736 try:
737 raise SystemExit(asyncio.run(run(args.copilot_cli)))
738 except (OSError, ValueError, RuntimeError) as error:
739 print(f"copilot-sidecar: {error}", file=sys.stderr)
740 raise SystemExit(2) from error
741
742
743 if __name__ == "__main__":
744 main()