Mercurial

comparison mrjunejune/inference/mock_sidecar.py @ 261:b401627fc49e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add JRPG mock flows and interactive previews Add scripted mock SSE commands, custom event forwarding, animated chat turns, full-height message navigation, and a cyberpunk resume dossier. Co-authored-by: Copilot <[email protected]>
author MrJuneJune <mrjunejune@users.noreply.github.com>
date Wed, 05 Aug 2026 20:38:32 -0700
parents
children 056790c4fb0d
comparison
equal deleted inserted replaced
260:1f9877b637e9 261:b401627fc49e
1 from __future__ import annotations
2
3 import argparse
4 import asyncio
5 import json
6 import os
7 import pathlib
8 import re
9 import sys
10 from dataclasses import dataclass
11 from typing import Any, Awaitable, Callable
12
13
14 JsonObject = dict[str, Any]
15 Emit = Callable[[JsonObject], Awaitable[None]]
16 COMMAND_PATTERN = re.compile(r"^![a-z][a-z0-9_-]*$")
17 EVENT_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
18 RESERVED_EVENT_TYPES = {"bridge.closed", "ready", "turn.accepted", "turn.done"}
19 RESERVED_EVENT_FIELDS = {
20 "conversation_id",
21 "mock",
22 "mock_command",
23 "request_id",
24 }
25 MAX_DELAY_MS = 60_000
26
27
28 @dataclass(frozen=True)
29 class MockEvent:
30 event_type: str
31 delay_ms: int | None
32 payload: JsonObject
33
34
35 @dataclass(frozen=True)
36 class MockCommand:
37 name: str
38 description: str
39 events: tuple[MockEvent, ...]
40 failed: bool
41
42
43 @dataclass(frozen=True)
44 class MockConfig:
45 delay_ms: int
46 fallback: str
47 commands: dict[str, MockCommand]
48
49 def select(self, prompt: str) -> MockCommand:
50 first_token = prompt.strip().split(maxsplit=1)[0].lower()
51 return self.commands.get(first_token, self.commands[self.fallback])
52
53
54 @dataclass
55 class MockTurn:
56 request_id: str
57 task: asyncio.Task[None]
58
59
60 def _require_non_negative_integer(value: Any, description: str) -> int:
61 if not isinstance(value, int) or isinstance(value, bool) or value < 0:
62 raise ValueError(f"{description} must be a non-negative integer")
63 if value > MAX_DELAY_MS:
64 raise ValueError(f"{description} must not exceed {MAX_DELAY_MS}")
65 return value
66
67
68 def _validate_event(command_name: str, index: int, raw: Any) -> MockEvent:
69 if not isinstance(raw, dict):
70 raise ValueError(f"{command_name} event {index} must be an object")
71 event_type = raw.get("type")
72 if not isinstance(event_type, str) or not EVENT_PATTERN.fullmatch(event_type):
73 raise ValueError(
74 f"{command_name} event {index} requires a safe event type"
75 )
76 if event_type in RESERVED_EVENT_TYPES:
77 raise ValueError(
78 f"{command_name} cannot script lifecycle event {event_type}"
79 )
80 reserved = set(raw) & RESERVED_EVENT_FIELDS
81 if reserved:
82 raise ValueError(
83 f"{command_name} event {index} uses reserved fields: "
84 f"{', '.join(sorted(reserved))}"
85 )
86 delay_ms = raw.get("delay_ms")
87 if delay_ms is not None:
88 delay_ms = _require_non_negative_integer(
89 delay_ms,
90 f"{command_name} event {index} delay_ms",
91 )
92 payload = {
93 key: value
94 for key, value in raw.items()
95 if key not in {"delay_ms", "type"}
96 }
97 if event_type == "assistant.delta":
98 if not isinstance(payload.get("delta"), str) or not payload["delta"]:
99 raise ValueError(
100 f"{command_name} assistant.delta requires non-empty delta"
101 )
102 elif event_type == "assistant.completed":
103 if not isinstance(payload.get("content"), str):
104 raise ValueError(
105 f"{command_name} assistant.completed requires content"
106 )
107 elif event_type == "assistant.usage":
108 usage = payload.get("usage")
109 if not isinstance(usage, dict):
110 raise ValueError(f"{command_name} assistant.usage requires usage")
111 for key in ("input_tokens", "output_tokens"):
112 _require_non_negative_integer(
113 usage.get(key),
114 f"{command_name} assistant.usage {key}",
115 )
116 elif event_type == "turn.error":
117 error = payload.get("error")
118 if not isinstance(error, dict):
119 raise ValueError(f"{command_name} turn.error requires error")
120 if (
121 not isinstance(error.get("code"), str)
122 or not error["code"]
123 or not isinstance(error.get("message"), str)
124 or not error["message"]
125 ):
126 raise ValueError(
127 f"{command_name} turn.error requires error code and message"
128 )
129 return MockEvent(event_type, delay_ms, payload)
130
131
132 def _validate_command(name: str, raw: Any) -> MockCommand:
133 if not COMMAND_PATTERN.fullmatch(name):
134 raise ValueError(f"invalid mock command: {name}")
135 if not isinstance(raw, dict):
136 raise ValueError(f"{name} must contain an object")
137 unknown = set(raw) - {"description", "events"}
138 if unknown:
139 raise ValueError(
140 f"{name} has unknown keys: {', '.join(sorted(unknown))}"
141 )
142 description = raw.get("description", "")
143 if not isinstance(description, str):
144 raise ValueError(f"{name} description must be a string")
145 raw_events = raw.get("events")
146 if not isinstance(raw_events, list) or not raw_events:
147 raise ValueError(f"{name} events must be a non-empty array")
148 if len(raw_events) > 256:
149 raise ValueError(f"{name} has too many events")
150 events = tuple(
151 _validate_event(name, index, event)
152 for index, event in enumerate(raw_events)
153 )
154 errors = [event for event in events if event.event_type == "turn.error"]
155 completed = [
156 event for event in events if event.event_type == "assistant.completed"
157 ]
158 if errors:
159 if len(errors) != 1 or errors[0] is not events[-1]:
160 raise ValueError(f"{name} turn.error must be the final scripted event")
161 if completed:
162 raise ValueError(f"{name} cannot complete and fail the same turn")
163 elif len(completed) != 1:
164 raise ValueError(f"{name} requires one assistant.completed event")
165
166 deltas = "".join(
167 event.payload["delta"]
168 for event in events
169 if event.event_type == "assistant.delta"
170 )
171 if completed:
172 completed_index = events.index(completed[0])
173 if any(
174 event.event_type == "assistant.delta"
175 for event in events[completed_index + 1:]
176 ):
177 raise ValueError(
178 f"{name} assistant.completed must follow all deltas"
179 )
180 if completed and deltas and deltas != completed[0].payload["content"]:
181 raise ValueError(f"{name} deltas must join to completed content")
182 return MockCommand(name, description, events, failed=bool(errors))
183
184
185 def load_mock_config(path: str | os.PathLike[str]) -> MockConfig:
186 with open(path, encoding="utf-8") as config_file:
187 payload = json.load(config_file)
188 if not isinstance(payload, dict):
189 raise ValueError("mock response file must contain a JSON object")
190 unknown = set(payload) - {"commands", "delay_ms", "fallback"}
191 if unknown:
192 raise ValueError(f"unknown mock config keys: {', '.join(sorted(unknown))}")
193 delay_ms = _require_non_negative_integer(
194 payload.get("delay_ms", 55),
195 "delay_ms",
196 )
197 fallback = payload.get("fallback")
198 if not isinstance(fallback, str) or not COMMAND_PATTERN.fullmatch(fallback):
199 raise ValueError("fallback must be a !command")
200 raw_commands = payload.get("commands")
201 if not isinstance(raw_commands, dict) or not raw_commands:
202 raise ValueError("commands must be a non-empty object")
203 commands = {
204 name.lower(): _validate_command(name.lower(), raw)
205 for name, raw in raw_commands.items()
206 }
207 if fallback.lower() not in commands:
208 raise ValueError(f"fallback command is not defined: {fallback}")
209 return MockConfig(delay_ms, fallback.lower(), commands)
210
211
212 class MockSidecar:
213 def __init__(self, emit: Emit, config: MockConfig) -> None:
214 self._emit = emit
215 self._config = config
216 self._active: dict[str, MockTurn] = {}
217 self.shutting_down = False
218
219 async def announce_ready(self) -> None:
220 await self._send("ready", None, None, status="ok", mock=True)
221
222 async def dispatch(self, command: JsonObject) -> None:
223 command_name = command.get("command")
224 request_id = command.get("request_id")
225 conversation_id = command.get("conversation_id")
226 if not isinstance(request_id, str) or not request_id:
227 await self._fail(
228 request_id if isinstance(request_id, str) else None,
229 conversation_id if isinstance(conversation_id, str) else None,
230 "invalid_request",
231 "request_id is required",
232 )
233 return
234 if command_name == "shutdown":
235 await self._shutdown(request_id, conversation_id)
236 return
237 if command_name == "health":
238 await self._send(
239 "ready",
240 request_id,
241 conversation_id if isinstance(conversation_id, str) else None,
242 status="ok",
243 mock=True,
244 )
245 return
246 if not isinstance(conversation_id, str) or not conversation_id:
247 await self._fail(
248 request_id,
249 None,
250 "invalid_request",
251 "conversation_id is required",
252 )
253 return
254 if command_name == "turn.start":
255 prompt = command.get("prompt")
256 if not isinstance(prompt, str) or not prompt:
257 await self._fail(
258 request_id,
259 conversation_id,
260 "invalid_request",
261 "prompt is required",
262 )
263 return
264 await self._start_turn(request_id, conversation_id, prompt)
265 elif command_name == "turn.abort":
266 await self._abort_turn(request_id, conversation_id)
267 elif command_name == "conversation.delete":
268 await self._delete_conversation(request_id, conversation_id)
269 else:
270 await self._fail(
271 request_id,
272 conversation_id,
273 "unknown_command",
274 f"unsupported command: {command_name!r}",
275 )
276
277 async def wait_for_idle(self) -> None:
278 tasks = [turn.task for turn in self._active.values()]
279 if tasks:
280 await asyncio.gather(*tasks, return_exceptions=True)
281
282 async def _start_turn(
283 self,
284 request_id: str,
285 conversation_id: str,
286 prompt: str,
287 ) -> None:
288 if conversation_id in self._active:
289 await self._fail(
290 request_id,
291 conversation_id,
292 "turn_in_progress",
293 "the conversation already has an active turn",
294 )
295 return
296 scripted_command = self._config.select(prompt)
297 await self._send(
298 "turn.accepted",
299 request_id,
300 conversation_id,
301 mock=True,
302 mock_command=scripted_command.name,
303 )
304 task = asyncio.create_task(
305 self._stream_turn(request_id, conversation_id, scripted_command)
306 )
307 self._active[conversation_id] = MockTurn(request_id, task)
308
309 async def _stream_turn(
310 self,
311 request_id: str,
312 conversation_id: str,
313 scripted_command: MockCommand,
314 ) -> None:
315 try:
316 for event in scripted_command.events:
317 delay = (
318 self._config.delay_ms
319 if event.delay_ms is None
320 else event.delay_ms
321 )
322 await asyncio.sleep(delay / 1000)
323 await self._send(
324 event.event_type,
325 request_id,
326 conversation_id,
327 **event.payload,
328 mock=True,
329 mock_command=scripted_command.name,
330 )
331 await self._send(
332 "turn.done",
333 request_id,
334 conversation_id,
335 failed=scripted_command.failed,
336 mock=True,
337 mock_command=scripted_command.name,
338 )
339 except asyncio.CancelledError:
340 await self._send(
341 "turn.done",
342 request_id,
343 conversation_id,
344 aborted=True,
345 mock=True,
346 mock_command=scripted_command.name,
347 )
348 except Exception as error:
349 await self._send(
350 "turn.error",
351 request_id,
352 conversation_id,
353 error={
354 "code": "mock_script_failed",
355 "message": str(error),
356 },
357 mock_command=scripted_command.name,
358 )
359 await self._send(
360 "turn.done",
361 request_id,
362 conversation_id,
363 failed=True,
364 mock=True,
365 mock_command=scripted_command.name,
366 )
367 finally:
368 active = self._active.get(conversation_id)
369 if active is not None and active.request_id == request_id:
370 self._active.pop(conversation_id, None)
371
372 async def _abort_turn(
373 self,
374 request_id: str,
375 conversation_id: str,
376 ) -> None:
377 active = self._active.get(conversation_id)
378 if active is None:
379 await self._fail(
380 request_id,
381 conversation_id,
382 "no_active_turn",
383 "no active turn to abort",
384 )
385 return
386 await self._send(
387 "turn.accepted",
388 request_id,
389 conversation_id,
390 action="abort",
391 target_request_id=active.request_id,
392 mock=True,
393 )
394 active.task.cancel()
395 await asyncio.gather(active.task, return_exceptions=True)
396 await self._send(
397 "turn.done",
398 request_id,
399 conversation_id,
400 action="abort",
401 target_request_id=active.request_id,
402 mock=True,
403 )
404
405 async def _delete_conversation(
406 self,
407 request_id: str,
408 conversation_id: str,
409 ) -> None:
410 active = self._active.get(conversation_id)
411 if active is not None:
412 active.task.cancel()
413 await asyncio.gather(active.task, return_exceptions=True)
414 await self._send(
415 "turn.done",
416 request_id,
417 conversation_id,
418 action="conversation.delete",
419 mock=True,
420 )
421
422 async def _shutdown(
423 self,
424 request_id: str,
425 conversation_id: Any,
426 ) -> None:
427 self.shutting_down = True
428 tasks = [turn.task for turn in self._active.values()]
429 for task in tasks:
430 task.cancel()
431 if tasks:
432 await asyncio.gather(*tasks, return_exceptions=True)
433 await self._send(
434 "turn.done",
435 request_id,
436 conversation_id if isinstance(conversation_id, str) else None,
437 action="shutdown",
438 mock=True,
439 )
440
441 async def _fail(
442 self,
443 request_id: str | None,
444 conversation_id: str | None,
445 code: str,
446 message: str,
447 ) -> None:
448 await self._send(
449 "turn.error",
450 request_id,
451 conversation_id,
452 error={"code": code, "message": message},
453 )
454 await self._send(
455 "turn.done",
456 request_id,
457 conversation_id,
458 failed=True,
459 mock=True,
460 )
461
462 async def _send(
463 self,
464 event_type: str,
465 request_id: str | None,
466 conversation_id: str | None,
467 **fields: Any,
468 ) -> None:
469 await self._emit(
470 {
471 "type": event_type,
472 "request_id": request_id,
473 "conversation_id": conversation_id,
474 **fields,
475 }
476 )
477
478
479 async def run(config: MockConfig) -> None:
480 async def emit(payload: JsonObject) -> None:
481 sys.stdout.write(json.dumps(payload, separators=(",", ":")) + "\n")
482 sys.stdout.flush()
483
484 sidecar = MockSidecar(emit, config)
485 await sidecar.announce_ready()
486 while not sidecar.shutting_down:
487 line = await asyncio.to_thread(sys.stdin.readline)
488 if not line:
489 break
490 try:
491 command = json.loads(line)
492 if not isinstance(command, dict):
493 raise ValueError("command must be a JSON object")
494 except (json.JSONDecodeError, ValueError) as error:
495 await sidecar._fail(None, None, "invalid_json", str(error))
496 continue
497 await sidecar.dispatch(command)
498 await sidecar.wait_for_idle()
499
500
501 def main() -> None:
502 packaged_responses = pathlib.Path(__file__).with_name("mock_responses.json")
503 parser = argparse.ArgumentParser(description="Scripted JRPG mock sidecar")
504 parser.add_argument("copilot_cli", nargs="?")
505 parser.add_argument("--responses", default=str(packaged_responses))
506 args = parser.parse_args()
507 responses_path = os.environ.get("MRJUNEJUNE_MOCK_RESPONSES") or args.responses
508 config = load_mock_config(responses_path)
509 delay_override = os.environ.get("MRJUNEJUNE_MOCK_DELAY_MS")
510 if delay_override is not None:
511 try:
512 delay_ms = int(delay_override)
513 except ValueError as error:
514 raise ValueError(
515 "MRJUNEJUNE_MOCK_DELAY_MS must be a non-negative integer"
516 ) from error
517 delay_ms = _require_non_negative_integer(
518 delay_ms,
519 "MRJUNEJUNE_MOCK_DELAY_MS",
520 )
521 config = MockConfig(delay_ms, config.fallback, config.commands)
522 asyncio.run(run(config))
523
524
525 if __name__ == "__main__":
526 main()