support pi

server/packages/universal-agent-schema/tests/pi_conversion.rs

@@ -0,0 +1,264 @@
+use sandbox_agent_universal_agent_schema::convert_pi::PiEventConverter;
+use sandbox_agent_universal_agent_schema::pi as pi_schema;
+use sandbox_agent_universal_agent_schema::{
+    ContentPart, ItemKind, ItemRole, ItemStatus, UniversalEventData, UniversalEventType,
+};
+use serde_json::json;
+
+fn parse_event(value: serde_json::Value) -> pi_schema::RpcEvent {
+    serde_json::from_value(value).expect("pi event")
+}
+
+#[test]
+fn pi_message_flow_converts() {
+    let mut converter = PiEventConverter::default();
+
+    let start_event = parse_event(json!({
+        "type": "message_start",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "message": {
+            "role": "assistant",
+            "content": [{ "type": "text", "text": "Hello" }]
+        }
+    }));
+    let start_events = converter
+        .event_to_universal(&start_event)
+        .expect("start conversions");
+    assert_eq!(start_events[0].event_type, UniversalEventType::ItemStarted);
+    if let UniversalEventData::Item(item) = &start_events[0].data {
+        assert_eq!(item.item.kind, ItemKind::Message);
+        assert_eq!(item.item.role, Some(ItemRole::Assistant));
+        assert_eq!(item.item.status, ItemStatus::InProgress);
+    } else {
+        panic!("expected item event");
+    }
+
+    let update_event = parse_event(json!({
+        "type": "message_update",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "assistantMessageEvent": { "type": "text_delta", "delta": " world" }
+    }));
+    let update_events = converter
+        .event_to_universal(&update_event)
+        .expect("update conversions");
+    assert_eq!(update_events[0].event_type, UniversalEventType::ItemDelta);
+
+    let end_event = parse_event(json!({
+        "type": "message_end",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "message": {
+            "role": "assistant",
+            "content": [{ "type": "text", "text": "Hello world" }]
+        }
+    }));
+    let end_events = converter
+        .event_to_universal(&end_event)
+        .expect("end conversions");
+    assert_eq!(end_events[0].event_type, UniversalEventType::ItemCompleted);
+    if let UniversalEventData::Item(item) = &end_events[0].data {
+        assert_eq!(item.item.kind, ItemKind::Message);
+        assert_eq!(item.item.role, Some(ItemRole::Assistant));
+        assert_eq!(item.item.status, ItemStatus::Completed);
+    } else {
+        panic!("expected item event");
+    }
+}
+
+#[test]
+fn pi_user_message_echo_is_skipped() {
+    let mut converter = PiEventConverter::default();
+
+    // Pi may echo the user message as a message_start with role "user".
+    // The daemon already records synthetic user events, so the converter
+    // must skip these to avoid a duplicate assistant-looking bubble.
+    let start_event = parse_event(json!({
+        "type": "message_start",
+        "sessionId": "session-1",
+        "messageId": "user-msg-1",
+        "message": {
+            "role": "user",
+            "content": [{ "type": "text", "text": "hello!" }]
+        }
+    }));
+    let events = converter
+        .event_to_universal(&start_event)
+        .expect("user message_start should not error");
+    assert!(
+        events.is_empty(),
+        "user message_start should produce no events, got {}",
+        events.len()
+    );
+
+    let end_event = parse_event(json!({
+        "type": "message_end",
+        "sessionId": "session-1",
+        "messageId": "user-msg-1",
+        "message": {
+            "role": "user",
+            "content": [{ "type": "text", "text": "hello!" }]
+        }
+    }));
+    let events = converter
+        .event_to_universal(&end_event)
+        .expect("user message_end should not error");
+    assert!(
+        events.is_empty(),
+        "user message_end should produce no events, got {}",
+        events.len()
+    );
+
+    // A subsequent assistant message should still work normally.
+    let assistant_start = parse_event(json!({
+        "type": "message_start",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "message": {
+            "role": "assistant",
+            "content": [{ "type": "text", "text": "Hello! How can I help?" }]
+        }
+    }));
+    let events = converter
+        .event_to_universal(&assistant_start)
+        .expect("assistant message_start");
+    assert_eq!(events.len(), 1);
+    assert_eq!(events[0].event_type, UniversalEventType::ItemStarted);
+    if let UniversalEventData::Item(item) = &events[0].data {
+        assert_eq!(item.item.role, Some(ItemRole::Assistant));
+    } else {
+        panic!("expected item event");
+    }
+}
+
+#[test]
+fn pi_tool_execution_converts_with_partial_deltas() {
+    let mut converter = PiEventConverter::default();
+
+    let start_event = parse_event(json!({
+        "type": "tool_execution_start",
+        "sessionId": "session-1",
+        "toolCallId": "call-1",
+        "toolName": "bash",
+        "args": { "command": "ls" }
+    }));
+    let start_events = converter
+        .event_to_universal(&start_event)
+        .expect("tool start");
+    assert_eq!(start_events[0].event_type, UniversalEventType::ItemStarted);
+    if let UniversalEventData::Item(item) = &start_events[0].data {
+        assert_eq!(item.item.kind, ItemKind::ToolCall);
+        assert_eq!(item.item.role, Some(ItemRole::Assistant));
+        match &item.item.content[0] {
+            ContentPart::ToolCall { name, .. } => assert_eq!(name, "bash"),
+            _ => panic!("expected tool call content"),
+        }
+    }
+
+    let update_event = parse_event(json!({
+        "type": "tool_execution_update",
+        "sessionId": "session-1",
+        "toolCallId": "call-1",
+        "partialResult": "foo"
+    }));
+    let update_events = converter
+        .event_to_universal(&update_event)
+        .expect("tool update");
+    assert!(update_events
+        .iter()
+        .any(|event| event.event_type == UniversalEventType::ItemDelta));
+
+    let update_event2 = parse_event(json!({
+        "type": "tool_execution_update",
+        "sessionId": "session-1",
+        "toolCallId": "call-1",
+        "partialResult": "foobar"
+    }));
+    let update_events2 = converter
+        .event_to_universal(&update_event2)
+        .expect("tool update 2");
+    let delta = update_events2
+        .iter()
+        .find_map(|event| match &event.data {
+            UniversalEventData::ItemDelta(data) => Some(data.delta.clone()),
+            _ => None,
+        })
+        .unwrap_or_default();
+    assert_eq!(delta, "bar");
+
+    let end_event = parse_event(json!({
+        "type": "tool_execution_end",
+        "sessionId": "session-1",
+        "toolCallId": "call-1",
+        "result": { "type": "text", "content": "done" },
+        "isError": false
+    }));
+    let end_events = converter.event_to_universal(&end_event).expect("tool end");
+    assert_eq!(end_events[0].event_type, UniversalEventType::ItemCompleted);
+    if let UniversalEventData::Item(item) = &end_events[0].data {
+        assert_eq!(item.item.kind, ItemKind::ToolResult);
+        assert_eq!(item.item.role, Some(ItemRole::Tool));
+        match &item.item.content[0] {
+            ContentPart::ToolResult { output, .. } => assert_eq!(output, "done"),
+            _ => panic!("expected tool result content"),
+        }
+    }
+}
+
+#[test]
+fn pi_unknown_event_returns_error() {
+    let mut converter = PiEventConverter::default();
+    let event = parse_event(json!({
+        "type": "unknown_event",
+        "sessionId": "session-1"
+    }));
+    assert!(converter.event_to_universal(&event).is_err());
+}
+
+#[test]
+fn pi_message_done_completes_without_message_end() {
+    let mut converter = PiEventConverter::default();
+
+    let start_event = parse_event(json!({
+        "type": "message_start",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "message": {
+            "role": "assistant",
+            "content": [{ "type": "text", "text": "Hello" }]
+        }
+    }));
+    let _start_events = converter
+        .event_to_universal(&start_event)
+        .expect("start conversions");
+
+    let update_event = parse_event(json!({
+        "type": "message_update",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "assistantMessageEvent": { "type": "text_delta", "delta": " world" }
+    }));
+    let _update_events = converter
+        .event_to_universal(&update_event)
+        .expect("update conversions");
+
+    let done_event = parse_event(json!({
+        "type": "message_update",
+        "sessionId": "session-1",
+        "messageId": "msg-1",
+        "assistantMessageEvent": { "type": "done" }
+    }));
+    let done_events = converter
+        .event_to_universal(&done_event)
+        .expect("done conversions");
+    assert_eq!(done_events[0].event_type, UniversalEventType::ItemCompleted);
+    if let UniversalEventData::Item(item) = &done_events[0].data {
+        assert_eq!(item.item.status, ItemStatus::Completed);
+        assert!(
+            matches!(item.item.content.get(0), Some(ContentPart::Text { text }) if text == "Hello world")
+        );
+    } else {
+        panic!("expected item event");
+    }
+}