can1357
528d8225b7
- Added error handling callbacks (on_protocol_error, on_listener_error) and error history tracking (protocol_errors, listener_errors properties) for RPC failures. - Added max_event_history and max_stderr_chunks parameters to bound retained history in long-lived hosts, preventing unbounded memory growth. - Added assistant_text_with_thinking() and message_text_with_thinking() helper functions for backward compatibility with thinking block inclusion. - Improved error correlation for id-less parse/unknown-command failures to waiting requests and changed prompt scheduling failures to raise instead of timing out. - Refactored event dispatch to _dispatch_listeners() with exception handling and error recording for robust listener failure management.
176 lines
6.1 KiB
Python
176 lines
6.1 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
from omp_rpc import (
|
|
AgentEndEvent,
|
|
ExtensionUiRequest,
|
|
SessionState,
|
|
TodoReminderEvent,
|
|
assistant_text,
|
|
assistant_text_with_thinking,
|
|
parse_notification,
|
|
parse_session_state,
|
|
)
|
|
|
|
|
|
class ProtocolParsingTests(unittest.TestCase):
|
|
def test_parse_session_state(self) -> None:
|
|
state = parse_session_state(
|
|
{
|
|
"model": {
|
|
"id": "claude-sonnet-4-5",
|
|
"name": "Claude Sonnet 4.5",
|
|
"api": "anthropic-messages",
|
|
"provider": "anthropic",
|
|
"baseUrl": "https://api.anthropic.com",
|
|
"reasoning": True,
|
|
"input": ["text", "image"],
|
|
"cost": {
|
|
"input": 1.0,
|
|
"output": 2.0,
|
|
"cacheRead": 0.1,
|
|
"cacheWrite": 0.2,
|
|
},
|
|
"contextWindow": 200000,
|
|
"maxTokens": 8192,
|
|
"thinking": {
|
|
"minLevel": "minimal",
|
|
"maxLevel": "high",
|
|
"mode": "effort",
|
|
},
|
|
},
|
|
"thinkingLevel": "medium",
|
|
"isStreaming": False,
|
|
"isCompacting": False,
|
|
"steeringMode": "one-at-a-time",
|
|
"followUpMode": "all",
|
|
"interruptMode": "immediate",
|
|
"sessionFile": "/tmp/test.jsonl",
|
|
"sessionId": "session-123",
|
|
"sessionName": "Scratchpad",
|
|
"autoCompactionEnabled": True,
|
|
"messageCount": 4,
|
|
"queuedMessageCount": 1,
|
|
"todoPhases": [
|
|
{
|
|
"id": "phase-1",
|
|
"name": "Todos",
|
|
"tasks": [
|
|
{
|
|
"id": "task-1",
|
|
"content": "Map tools",
|
|
"status": "in_progress",
|
|
"details": "Inspect read and edit first.",
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"systemPrompt": "You are useful.",
|
|
"dumpTools": [
|
|
{
|
|
"name": "read",
|
|
"description": "Read files",
|
|
"parameters": {"type": "object"},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
self.assertIsInstance(state, SessionState)
|
|
self.assertEqual(state.session_id, "session-123")
|
|
self.assertEqual(state.follow_up_mode, "all")
|
|
self.assertEqual(state.model.id if state.model else None, "claude-sonnet-4-5")
|
|
self.assertEqual(state.todo_phases[0].tasks[0].status, "in_progress")
|
|
self.assertEqual(state.dump_tools[0].name, "read")
|
|
|
|
def test_parse_agent_end_notification(self) -> None:
|
|
notification = parse_notification(
|
|
{
|
|
"type": "agent_end",
|
|
"messages": [
|
|
{
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "hello"}],
|
|
"api": "anthropic-messages",
|
|
"provider": "anthropic",
|
|
"model": "claude-sonnet-4-5",
|
|
"usage": {
|
|
"input": 1,
|
|
"output": 1,
|
|
"cacheRead": 0,
|
|
"cacheWrite": 0,
|
|
"totalTokens": 2,
|
|
"cost": {
|
|
"input": 0.0,
|
|
"output": 0.0,
|
|
"cacheRead": 0.0,
|
|
"cacheWrite": 0.0,
|
|
"total": 0.0,
|
|
},
|
|
},
|
|
"stopReason": "stop",
|
|
"timestamp": 1,
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
self.assertIsInstance(notification, AgentEndEvent)
|
|
self.assertEqual(assistant_text(notification.messages[0]), "hello")
|
|
|
|
def test_parse_extension_ui_request(self) -> None:
|
|
notification = parse_notification(
|
|
{
|
|
"type": "extension_ui_request",
|
|
"id": "ui-1",
|
|
"method": "confirm",
|
|
"title": "Confirm",
|
|
"message": "Continue?",
|
|
"timeout": 1000,
|
|
}
|
|
)
|
|
|
|
self.assertIsInstance(notification, ExtensionUiRequest)
|
|
self.assertEqual(notification.method, "confirm")
|
|
self.assertEqual(notification.message, "Continue?")
|
|
self.assertTrue(notification.is_interactive())
|
|
self.assertTrue(notification.requires_response())
|
|
self.assertFalse(notification.is_passive())
|
|
|
|
def test_parse_todo_reminder_notification(self) -> None:
|
|
notification = parse_notification(
|
|
{
|
|
"type": "todo_reminder",
|
|
"attempt": 1,
|
|
"maxAttempts": 3,
|
|
"todos": [
|
|
{
|
|
"id": "task-1",
|
|
"content": "Map tools",
|
|
"status": "pending",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
self.assertIsInstance(notification, TodoReminderEvent)
|
|
self.assertEqual(notification.todos[0].content, "Map tools")
|
|
self.assertEqual(notification.todos[0].status, "pending")
|
|
|
|
def test_assistant_text_excludes_thinking_by_default(self) -> None:
|
|
message = {
|
|
"role": "assistant",
|
|
"content": [
|
|
{"type": "thinking", "thinking": "internal"},
|
|
{"type": "text", "text": "visible"},
|
|
],
|
|
}
|
|
|
|
self.assertEqual(assistant_text(message), "visible")
|
|
self.assertEqual(assistant_text_with_thinking(message), "internalvisible")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|