diff --git a/python/omp-rpc/src/omp_rpc/client.py b/python/omp-rpc/src/omp_rpc/client.py index 30ac305c7..f87e05f99 100644 --- a/python/omp-rpc/src/omp_rpc/client.py +++ b/python/omp-rpc/src/omp_rpc/client.py @@ -1909,6 +1909,14 @@ class RpcClient: notification = UnknownNotification( _clone_json_object(payload), parse_error=str(exc) ) + if ( + payload_type == "agent_end" + and payload.get("isTerminal") is not False + ): + self._append_async_error( + RpcError(f"Failed to parse terminal agent_end: {exc}") + ) + self._mark_agent_run_completed() self._dispatch_listeners( "notification", notification.type, diff --git a/python/omp-rpc/tests/test_client.py b/python/omp-rpc/tests/test_client.py index cb9eb7356..d6591c47e 100644 --- a/python/omp-rpc/tests/test_client.py +++ b/python/omp-rpc/tests/test_client.py @@ -835,6 +835,19 @@ FORWARD_COMPAT_SERVER = textwrap.dedent( ) if command["type"] != "prompt": continue + if command.get("message") == "malformed terminal": + print( + json.dumps( + { + "type": "agent_end", + "messages": [{"role": "future_role"}], + "isTerminal": True, + } + ), + flush=True, + ) + time.sleep(2) + continue print( json.dumps( { @@ -1333,6 +1346,21 @@ class RpcClientTests(unittest.TestCase): self.assertEqual(len(unknown_errors), 1) self.assertIn("auto_compaction_start.reason", unknown_errors[0] or "") + def test_malformed_terminal_agent_end_wakes_waiter(self) -> None: + unknown_errors: list[str | None] = [] + + with self.make_client(server=FORWARD_COMPAT_SERVER) as client: + client.on_unknown_notification( + lambda event: unknown_errors.append(event.parse_error) + ) + with self.assertRaisesRegex( + RpcError, "Failed to parse terminal agent_end" + ): + client.prompt_and_wait("malformed terminal", timeout=1.0) + + self.assertEqual(len(unknown_errors), 1) + self.assertIn("messages[0].role", unknown_errors[0] or "") + def test_ui_confirmation_and_cancel_round_trip(self) -> None: with self.make_client() as client: client.prompt("needs confirm")