import asyncio import types import unittest from dataclasses import replace from mrjunejune.inference.copilot_sidecar import Sidecar, SidecarConfig def event(event_type, **data): return types.SimpleNamespace( type=types.SimpleNamespace(value=event_type), data=types.SimpleNamespace(**data), ) class FakeSession: def __init__(self, session_id, behavior="complete"): self.session_id = session_id self.behavior = behavior self.handlers = [] self.prompts = [] self.abort_calls = 0 self.disconnect_calls = 0 self.disconnect_error = False def on(self, handler): self.handlers.append(handler) def unsubscribe(): if handler in self.handlers: self.handlers.remove(handler) return unsubscribe def emit(self, value): for handler in list(self.handlers): handler(value) async def send(self, prompt): self.prompts.append(prompt) message_id = f"message-{len(self.prompts)}" if self.behavior == "complete": self.emit( event( "assistant.message_delta", delta_content=f"{prompt}-delta", message_id=message_id, ) ) self.emit( event( "assistant.message", content=f"{prompt}-answer", message_id=message_id, model="test-model", ) ) self.emit( event( "assistant.usage", model="test-model", input_tokens=3, output_tokens=5, ) ) self.emit(event("session.idle", aborted=False)) elif self.behavior == "error": self.emit( event( "session.error", error_type="provider", error_code="upstream_error", message="upstream failed", status_code=502, ) ) return message_id async def abort(self): self.abort_calls += 1 if self.behavior == "abort_error": raise RuntimeError("abort failed") self.emit(event("session.idle", aborted=True)) async def disconnect(self): self.disconnect_calls += 1 if self.disconnect_error: raise RuntimeError("disconnect failed") class FakeClient: def __init__(self, behavior="complete"): self.behavior = behavior self.sessions = {} self.create_calls = [] self.resume_calls = [] self.delete_calls = [] self.started = False self.stopped = False self.resume_started = None self.resume_release = None async def start(self): self.started = True async def stop(self): self.stopped = True async def resume_session(self, session_id, **kwargs): self.resume_calls.append((session_id, kwargs)) if self.resume_started is not None: self.resume_started.set() if self.resume_release is not None: await self.resume_release.wait() if session_id not in self.sessions: raise LookupError(session_id) return self.sessions[session_id] async def create_session(self, session_id, **kwargs): self.create_calls.append((session_id, kwargs)) session = FakeSession(session_id, self.behavior) self.sessions[session_id] = session return session async def delete_session(self, session_id): self.delete_calls.append(session_id) self.sessions.pop(session_id, None) class SidecarTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.output = [] async def capture(payload): self.output.append(payload) self.capture = capture self.config = SidecarConfig( base_url="http://litellm.invalid/v1", model="test-model", wire_api="responses", base_directory="/not-used-by-fake", ) async def make_sidecar(self, behavior="complete"): client = FakeClient(behavior) sidecar = Sidecar(client, self.config, self.capture) await sidecar.start() return sidecar, client async def test_health_and_invalid_protocol(self): sidecar, _ = await self.make_sidecar() await sidecar.announce_ready() await sidecar.dispatch({"command": "health", "request_id": "health-1"}) await sidecar.dispatch({"command": "unknown", "request_id": "bad-1"}) self.assertEqual(self.output[0]["type"], "ready") self.assertIsNone(self.output[0]["request_id"]) self.assertEqual(self.output[1]["request_id"], "health-1") self.assertEqual( [item["type"] for item in self.output[-2:]], ["turn.error", "turn.done"], ) async def test_multi_turn_reuses_one_session_and_configures_sdk(self): sidecar, client = await self.make_sidecar() for request_id, prompt in (("r1", "first"), ("r2", "second")): await sidecar.dispatch( { "command": "turn.start", "request_id": request_id, "conversation_id": "conversation-a", "prompt": prompt, } ) await sidecar.drain_events() self.assertEqual(len(client.create_calls), 1) self.assertEqual(len(client.resume_calls), 1) session = client.sessions["conversation-a"] self.assertEqual(session.prompts, ["first", "second"]) options = client.create_calls[0][1] self.assertEqual(options["provider"]["type"], "openai") self.assertEqual( options["provider"]["base_url"], "http://litellm.invalid/v1" ) self.assertEqual(options["provider"]["wire_api"], "responses") self.assertEqual(options["model"], "test-model") self.assertEqual(options["available_tools"], []) self.assertTrue(options["streaming"]) self.assertEqual( [item["request_id"] for item in self.output if item["type"] == "turn.done"], ["r1", "r2"], ) async def test_concurrent_conversations_keep_correlation(self): sidecar, _ = await self.make_sidecar() await asyncio.gather( sidecar.dispatch( { "command": "turn.start", "request_id": "left-request", "conversation_id": "left", "prompt": "left", } ), sidecar.dispatch( { "command": "turn.start", "request_id": "right-request", "conversation_id": "right", "prompt": "right", } ), ) await sidecar.drain_events() correlated = { (item["request_id"], item["conversation_id"]) for item in self.output if item["type"] in ("assistant.delta", "assistant.completed") } self.assertEqual( correlated, {("left-request", "left"), ("right-request", "right")}, ) async def test_abort_finishes_active_turn(self): sidecar, client = await self.make_sidecar("pending") await sidecar.dispatch( { "command": "turn.start", "request_id": "turn-request", "conversation_id": "abort-me", "prompt": "wait", } ) await sidecar.dispatch( { "command": "turn.abort", "request_id": "abort-request", "conversation_id": "abort-me", } ) await sidecar.drain_events() self.assertEqual(client.sessions["abort-me"].abort_calls, 1) done = [ item for item in self.output if item["type"] == "turn.done" and item["request_id"] == "turn-request" ] self.assertEqual(done[0]["aborted"], True) abort_accept = [ item for item in self.output if item["type"] == "turn.accepted" and item["request_id"] == "abort-request" ] self.assertEqual(abort_accept[0]["target_request_id"], "turn-request") async def test_session_error_is_terminal(self): sidecar, _ = await self.make_sidecar("error") await sidecar.dispatch( { "command": "turn.start", "request_id": "error-request", "conversation_id": "error-conversation", "prompt": "fail", } ) await sidecar.drain_events() errors = [item for item in self.output if item["type"] == "turn.error"] done = [item for item in self.output if item["type"] == "turn.done"] self.assertEqual(errors[0]["error"]["code"], "upstream_error") self.assertEqual(errors[0]["error"]["status_code"], 502) self.assertEqual(len(done), 1) self.assertTrue(done[0]["failed"]) async def test_abort_error_terminates_abort_request(self): sidecar, _ = await self.make_sidecar("abort_error") await sidecar.dispatch( { "command": "turn.start", "request_id": "active-request", "conversation_id": "abort-error", "prompt": "wait", } ) await sidecar.dispatch( { "command": "turn.abort", "request_id": "failed-abort", "conversation_id": "abort-error", } ) abort_events = [ item for item in self.output if item["request_id"] == "failed-abort" ] self.assertEqual( [item["type"] for item in abort_events], ["turn.accepted", "turn.error", "turn.done"], ) self.assertTrue(abort_events[-1]["failed"]) async def test_delete_and_shutdown_release_resources(self): sidecar, client = await self.make_sidecar("pending") await sidecar.dispatch( { "command": "turn.start", "request_id": "active", "conversation_id": "delete-me", "prompt": "wait", } ) session = client.sessions["delete-me"] await sidecar.dispatch( { "command": "conversation.delete", "request_id": "delete-request", "conversation_id": "delete-me", } ) await sidecar.dispatch( { "command": "shutdown", "request_id": "shutdown-request", "conversation_id": None, } ) self.assertEqual(session.disconnect_calls, 1) self.assertEqual(client.delete_calls, ["delete-me"]) self.assertTrue(client.stopped) self.assertTrue(sidecar.shutting_down) shutdown = [ item for item in self.output if item["request_id"] == "shutdown-request" ] self.assertEqual(shutdown[0]["type"], "turn.done") async def test_idle_and_overflow_sessions_are_evicted(self): self.config = replace( self.config, idle_timeout_seconds=3600, max_sessions=1, ) sidecar, client = await self.make_sidecar() for conversation_id in ("old", "new"): await sidecar.dispatch( { "command": "turn.start", "request_id": f"request-{conversation_id}", "conversation_id": conversation_id, "prompt": conversation_id, } ) await sidecar.drain_events() await asyncio.sleep(0.01) old_session = client.sessions["old"] old_session.disconnect_error = True await sidecar.evict_idle_sessions() self.assertEqual(old_session.disconnect_calls, 1) self.assertNotIn("old", sidecar._conversations) self.assertIn("new", sidecar._conversations) sidecar._conversations["new"].last_used -= 4000 await sidecar.evict_idle_sessions() self.assertEqual(client.sessions["new"].disconnect_calls, 1) self.assertEqual(sidecar._conversations, {}) await sidecar.dispatch( { "command": "shutdown", "request_id": "shutdown-eviction", "conversation_id": None, } ) async def test_shutdown_waits_for_inflight_session_creation(self): client = FakeClient("pending") client.resume_started = asyncio.Event() client.resume_release = asyncio.Event() sidecar = Sidecar(client, self.config, self.capture) await sidecar.start() start_task = asyncio.create_task( sidecar.dispatch( { "command": "turn.start", "request_id": "starting", "conversation_id": "race", "prompt": "wait", } ) ) await client.resume_started.wait() shutdown_task = asyncio.create_task( sidecar.dispatch( { "command": "shutdown", "request_id": "shutdown", "conversation_id": None, } ) ) await asyncio.sleep(0) self.assertTrue(sidecar.shutting_down) client.resume_release.set() await asyncio.gather(start_task, shutdown_task) self.assertTrue(client.stopped) self.assertEqual(client.sessions["race"].disconnect_calls, 1) await sidecar.dispatch( { "command": "turn.start", "request_id": "too-late", "conversation_id": "late", "prompt": "no", } ) late_error = [ item for item in self.output if item["request_id"] == "too-late" and item["type"] == "turn.error" ] self.assertEqual(late_error[0]["error"]["code"], "shutting_down") async def test_command_gates_are_released(self): sidecar, _ = await self.make_sidecar() for index in range(20): await sidecar.dispatch( { "command": "unknown", "request_id": f"unknown-{index}", "conversation_id": f"conversation-{index}", } ) self.assertEqual(sidecar._conversation_gates, {}) if __name__ == "__main__": unittest.main()