diff --git a/pycodeloop/cli/serve.py b/pycodeloop/cli/serve.py index 23bc3c9..82efbd4 100644 --- a/pycodeloop/cli/serve.py +++ b/pycodeloop/cli/serve.py @@ -72,6 +72,12 @@ def _respond_error(self, request_id, code: int, message: str) -> None: def _wire_callbacks(self) -> None: agent = self.flow.agent + def on_request(message_count: int, tool_count: int) -> None: + self._notify( + "chat/request", + {"messageCount": message_count, "toolCount": tool_count}, + ) + def on_text_delta(delta: str) -> None: self._notify("chat/textDelta", {"delta": delta}) @@ -135,6 +141,7 @@ def confirm(name: str, preview: str) -> bool | str: finally: self._confirm_waiters.pop(request_id, None) + agent.on_request = on_request agent.on_text_delta = on_text_delta agent.on_turn_end = on_turn_end agent.on_tool_call = on_tool_call diff --git a/tests/functional/test_serve_end_to_end.py b/tests/functional/test_serve_end_to_end.py index 2cd83c5..988dfa7 100644 --- a/tests/functional/test_serve_end_to_end.py +++ b/tests/functional/test_serve_end_to_end.py @@ -143,6 +143,37 @@ def test_chat_send_streams_tool_call_and_returns_final_text(self): response["result"]["text"], "Here's the directory listing." ) + def test_chat_send_emits_on_request_notification(self): + self.server = FakeLLMServer([chat_completion(text="hi")]) + self.client = ServeSubprocess(self.server.url, self.tmp_path) + self.client.next_message() # ready + + self.client.send( + { + "jsonrpc": "2.0", + "id": "1", + "method": "chat/send", + "params": {"prompt": "hello"}, + } + ) + + messages = [] + response = None + while response is None: + message = self.client.next_message() + if message.get("id") == "1" and "result" in message: + response = message + else: + messages.append(message) + + request_notification = next( + m for m in messages if m.get("method") == "chat/request" + ) + self.assertIsInstance( + request_notification["params"]["messageCount"], int + ) + self.assertIsInstance(request_notification["params"]["toolCount"], int) + def test_dangerous_tool_blocks_on_confirm_round_trip(self): self.server = FakeLLMServer( [