Mercurial

comparison mrjunejune/inference/copilot_sidecar_test.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 import asyncio
2 import types
3 import unittest
4 from dataclasses import replace
5
6 from mrjunejune.inference.copilot_sidecar import Sidecar, SidecarConfig
7
8
9 def event(event_type, **data):
10 return types.SimpleNamespace(
11 type=types.SimpleNamespace(value=event_type),
12 data=types.SimpleNamespace(**data),
13 )
14
15
16 class FakeSession:
17 def __init__(self, session_id, behavior="complete"):
18 self.session_id = session_id
19 self.behavior = behavior
20 self.handlers = []
21 self.prompts = []
22 self.abort_calls = 0
23 self.disconnect_calls = 0
24 self.disconnect_error = False
25
26 def on(self, handler):
27 self.handlers.append(handler)
28
29 def unsubscribe():
30 if handler in self.handlers:
31 self.handlers.remove(handler)
32
33 return unsubscribe
34
35 def emit(self, value):
36 for handler in list(self.handlers):
37 handler(value)
38
39 async def send(self, prompt):
40 self.prompts.append(prompt)
41 message_id = f"message-{len(self.prompts)}"
42 if self.behavior == "complete":
43 self.emit(
44 event(
45 "assistant.message_delta",
46 delta_content=f"{prompt}-delta",
47 message_id=message_id,
48 )
49 )
50 self.emit(
51 event(
52 "assistant.message",
53 content=f"{prompt}-answer",
54 message_id=message_id,
55 model="test-model",
56 )
57 )
58 self.emit(
59 event(
60 "assistant.usage",
61 model="test-model",
62 input_tokens=3,
63 output_tokens=5,
64 )
65 )
66 self.emit(event("session.idle", aborted=False))
67 elif self.behavior == "error":
68 self.emit(
69 event(
70 "session.error",
71 error_type="provider",
72 error_code="upstream_error",
73 message="upstream failed",
74 status_code=502,
75 )
76 )
77 return message_id
78
79 async def abort(self):
80 self.abort_calls += 1
81 if self.behavior == "abort_error":
82 raise RuntimeError("abort failed")
83 self.emit(event("session.idle", aborted=True))
84
85 async def disconnect(self):
86 self.disconnect_calls += 1
87 if self.disconnect_error:
88 raise RuntimeError("disconnect failed")
89
90
91 class FakeClient:
92 def __init__(self, behavior="complete"):
93 self.behavior = behavior
94 self.sessions = {}
95 self.create_calls = []
96 self.resume_calls = []
97 self.delete_calls = []
98 self.started = False
99 self.stopped = False
100 self.resume_started = None
101 self.resume_release = None
102
103 async def start(self):
104 self.started = True
105
106 async def stop(self):
107 self.stopped = True
108
109 async def resume_session(self, session_id, **kwargs):
110 self.resume_calls.append((session_id, kwargs))
111 if self.resume_started is not None:
112 self.resume_started.set()
113 if self.resume_release is not None:
114 await self.resume_release.wait()
115 if session_id not in self.sessions:
116 raise LookupError(session_id)
117 return self.sessions[session_id]
118
119 async def create_session(self, session_id, **kwargs):
120 self.create_calls.append((session_id, kwargs))
121 session = FakeSession(session_id, self.behavior)
122 self.sessions[session_id] = session
123 return session
124
125 async def delete_session(self, session_id):
126 self.delete_calls.append(session_id)
127 self.sessions.pop(session_id, None)
128
129
130 class SidecarTest(unittest.IsolatedAsyncioTestCase):
131 async def asyncSetUp(self):
132 self.output = []
133
134 async def capture(payload):
135 self.output.append(payload)
136
137 self.capture = capture
138 self.config = SidecarConfig(
139 base_url="http://litellm.invalid/v1",
140 model="test-model",
141 wire_api="responses",
142 base_directory="/not-used-by-fake",
143 )
144
145 async def make_sidecar(self, behavior="complete"):
146 client = FakeClient(behavior)
147 sidecar = Sidecar(client, self.config, self.capture)
148 await sidecar.start()
149 return sidecar, client
150
151 async def test_health_and_invalid_protocol(self):
152 sidecar, _ = await self.make_sidecar()
153 await sidecar.announce_ready()
154 await sidecar.dispatch({"command": "health", "request_id": "health-1"})
155 await sidecar.dispatch({"command": "unknown", "request_id": "bad-1"})
156
157 self.assertEqual(self.output[0]["type"], "ready")
158 self.assertIsNone(self.output[0]["request_id"])
159 self.assertEqual(self.output[1]["request_id"], "health-1")
160 self.assertEqual(
161 [item["type"] for item in self.output[-2:]],
162 ["turn.error", "turn.done"],
163 )
164
165 async def test_multi_turn_reuses_one_session_and_configures_sdk(self):
166 sidecar, client = await self.make_sidecar()
167 for request_id, prompt in (("r1", "first"), ("r2", "second")):
168 await sidecar.dispatch(
169 {
170 "command": "turn.start",
171 "request_id": request_id,
172 "conversation_id": "conversation-a",
173 "prompt": prompt,
174 }
175 )
176 await sidecar.drain_events()
177
178 self.assertEqual(len(client.create_calls), 1)
179 self.assertEqual(len(client.resume_calls), 1)
180 session = client.sessions["conversation-a"]
181 self.assertEqual(session.prompts, ["first", "second"])
182 options = client.create_calls[0][1]
183 self.assertEqual(options["provider"]["type"], "openai")
184 self.assertEqual(
185 options["provider"]["base_url"], "http://litellm.invalid/v1"
186 )
187 self.assertEqual(options["provider"]["wire_api"], "responses")
188 self.assertEqual(options["model"], "test-model")
189 self.assertEqual(options["available_tools"], [])
190 self.assertTrue(options["streaming"])
191 self.assertEqual(
192 [item["request_id"] for item in self.output if item["type"] == "turn.done"],
193 ["r1", "r2"],
194 )
195
196 async def test_concurrent_conversations_keep_correlation(self):
197 sidecar, _ = await self.make_sidecar()
198 await asyncio.gather(
199 sidecar.dispatch(
200 {
201 "command": "turn.start",
202 "request_id": "left-request",
203 "conversation_id": "left",
204 "prompt": "left",
205 }
206 ),
207 sidecar.dispatch(
208 {
209 "command": "turn.start",
210 "request_id": "right-request",
211 "conversation_id": "right",
212 "prompt": "right",
213 }
214 ),
215 )
216 await sidecar.drain_events()
217
218 correlated = {
219 (item["request_id"], item["conversation_id"])
220 for item in self.output
221 if item["type"] in ("assistant.delta", "assistant.completed")
222 }
223 self.assertEqual(
224 correlated,
225 {("left-request", "left"), ("right-request", "right")},
226 )
227
228 async def test_abort_finishes_active_turn(self):
229 sidecar, client = await self.make_sidecar("pending")
230 await sidecar.dispatch(
231 {
232 "command": "turn.start",
233 "request_id": "turn-request",
234 "conversation_id": "abort-me",
235 "prompt": "wait",
236 }
237 )
238 await sidecar.dispatch(
239 {
240 "command": "turn.abort",
241 "request_id": "abort-request",
242 "conversation_id": "abort-me",
243 }
244 )
245 await sidecar.drain_events()
246
247 self.assertEqual(client.sessions["abort-me"].abort_calls, 1)
248 done = [
249 item
250 for item in self.output
251 if item["type"] == "turn.done"
252 and item["request_id"] == "turn-request"
253 ]
254 self.assertEqual(done[0]["aborted"], True)
255 abort_accept = [
256 item
257 for item in self.output
258 if item["type"] == "turn.accepted"
259 and item["request_id"] == "abort-request"
260 ]
261 self.assertEqual(abort_accept[0]["target_request_id"], "turn-request")
262
263 async def test_session_error_is_terminal(self):
264 sidecar, _ = await self.make_sidecar("error")
265 await sidecar.dispatch(
266 {
267 "command": "turn.start",
268 "request_id": "error-request",
269 "conversation_id": "error-conversation",
270 "prompt": "fail",
271 }
272 )
273 await sidecar.drain_events()
274
275 errors = [item for item in self.output if item["type"] == "turn.error"]
276 done = [item for item in self.output if item["type"] == "turn.done"]
277 self.assertEqual(errors[0]["error"]["code"], "upstream_error")
278 self.assertEqual(errors[0]["error"]["status_code"], 502)
279 self.assertEqual(len(done), 1)
280 self.assertTrue(done[0]["failed"])
281
282 async def test_abort_error_terminates_abort_request(self):
283 sidecar, _ = await self.make_sidecar("abort_error")
284 await sidecar.dispatch(
285 {
286 "command": "turn.start",
287 "request_id": "active-request",
288 "conversation_id": "abort-error",
289 "prompt": "wait",
290 }
291 )
292 await sidecar.dispatch(
293 {
294 "command": "turn.abort",
295 "request_id": "failed-abort",
296 "conversation_id": "abort-error",
297 }
298 )
299
300 abort_events = [
301 item for item in self.output if item["request_id"] == "failed-abort"
302 ]
303 self.assertEqual(
304 [item["type"] for item in abort_events],
305 ["turn.accepted", "turn.error", "turn.done"],
306 )
307 self.assertTrue(abort_events[-1]["failed"])
308
309 async def test_delete_and_shutdown_release_resources(self):
310 sidecar, client = await self.make_sidecar("pending")
311 await sidecar.dispatch(
312 {
313 "command": "turn.start",
314 "request_id": "active",
315 "conversation_id": "delete-me",
316 "prompt": "wait",
317 }
318 )
319 session = client.sessions["delete-me"]
320 await sidecar.dispatch(
321 {
322 "command": "conversation.delete",
323 "request_id": "delete-request",
324 "conversation_id": "delete-me",
325 }
326 )
327 await sidecar.dispatch(
328 {
329 "command": "shutdown",
330 "request_id": "shutdown-request",
331 "conversation_id": None,
332 }
333 )
334
335 self.assertEqual(session.disconnect_calls, 1)
336 self.assertEqual(client.delete_calls, ["delete-me"])
337 self.assertTrue(client.stopped)
338 self.assertTrue(sidecar.shutting_down)
339 shutdown = [
340 item
341 for item in self.output
342 if item["request_id"] == "shutdown-request"
343 ]
344 self.assertEqual(shutdown[0]["type"], "turn.done")
345
346 async def test_idle_and_overflow_sessions_are_evicted(self):
347 self.config = replace(
348 self.config,
349 idle_timeout_seconds=3600,
350 max_sessions=1,
351 )
352 sidecar, client = await self.make_sidecar()
353 for conversation_id in ("old", "new"):
354 await sidecar.dispatch(
355 {
356 "command": "turn.start",
357 "request_id": f"request-{conversation_id}",
358 "conversation_id": conversation_id,
359 "prompt": conversation_id,
360 }
361 )
362 await sidecar.drain_events()
363 await asyncio.sleep(0.01)
364
365 old_session = client.sessions["old"]
366 old_session.disconnect_error = True
367 await sidecar.evict_idle_sessions()
368 self.assertEqual(old_session.disconnect_calls, 1)
369 self.assertNotIn("old", sidecar._conversations)
370 self.assertIn("new", sidecar._conversations)
371
372 sidecar._conversations["new"].last_used -= 4000
373 await sidecar.evict_idle_sessions()
374 self.assertEqual(client.sessions["new"].disconnect_calls, 1)
375 self.assertEqual(sidecar._conversations, {})
376 await sidecar.dispatch(
377 {
378 "command": "shutdown",
379 "request_id": "shutdown-eviction",
380 "conversation_id": None,
381 }
382 )
383
384 async def test_shutdown_waits_for_inflight_session_creation(self):
385 client = FakeClient("pending")
386 client.resume_started = asyncio.Event()
387 client.resume_release = asyncio.Event()
388 sidecar = Sidecar(client, self.config, self.capture)
389 await sidecar.start()
390
391 start_task = asyncio.create_task(
392 sidecar.dispatch(
393 {
394 "command": "turn.start",
395 "request_id": "starting",
396 "conversation_id": "race",
397 "prompt": "wait",
398 }
399 )
400 )
401 await client.resume_started.wait()
402 shutdown_task = asyncio.create_task(
403 sidecar.dispatch(
404 {
405 "command": "shutdown",
406 "request_id": "shutdown",
407 "conversation_id": None,
408 }
409 )
410 )
411 await asyncio.sleep(0)
412 self.assertTrue(sidecar.shutting_down)
413 client.resume_release.set()
414 await asyncio.gather(start_task, shutdown_task)
415
416 self.assertTrue(client.stopped)
417 self.assertEqual(client.sessions["race"].disconnect_calls, 1)
418 await sidecar.dispatch(
419 {
420 "command": "turn.start",
421 "request_id": "too-late",
422 "conversation_id": "late",
423 "prompt": "no",
424 }
425 )
426 late_error = [
427 item
428 for item in self.output
429 if item["request_id"] == "too-late" and item["type"] == "turn.error"
430 ]
431 self.assertEqual(late_error[0]["error"]["code"], "shutting_down")
432
433 async def test_command_gates_are_released(self):
434 sidecar, _ = await self.make_sidecar()
435 for index in range(20):
436 await sidecar.dispatch(
437 {
438 "command": "unknown",
439 "request_id": f"unknown-{index}",
440 "conversation_id": f"conversation-{index}",
441 }
442 )
443 self.assertEqual(sidecar._conversation_gates, {})
444
445
446 if __name__ == "__main__":
447 unittest.main()