From bd030904bc2a5cf01742ebb485a8766c53090951 Mon Sep 17 00:00:00 2001 From: Franklin Date: Fri, 6 Feb 2026 19:16:53 -0500 Subject: [PATCH] pi tests --- server/packages/sandbox-agent/src/router.rs | 54 +++--- .../tests/agent-flows/pi_rpc_integration.rs | 5 +- .../tests/http/agent_endpoints.rs | 5 +- ...nts_snapshots@agent_install_codex.snap.new | 6 + ..._session_snapshot@multi_turn_mock.snap.new | 35 +--- ...n_snapshot@permission_events_mock.snap.new | 168 ------------------ ...apshot@question_reply_events_mock.snap.new | 42 ++++- ..._snapshot@concurrency_events_mock.snap.new | 74 ++++---- ..._events_snapshot@http_events_mock.snap.new | 35 +++- ...e_events_snapshot@sse_events_mock.snap.new | 75 ++++++++ .../universal-agent-schema/src/agents/pi.rs | 27 ++- .../tests/pi_conversion.rs | 107 +++++++++++ 12 files changed, 371 insertions(+), 262 deletions(-) create mode 100644 server/packages/sandbox-agent/tests/http/snapshots/http_endpoints__agent_endpoints__agent_endpoints_snapshots@agent_install_codex.snap.new create mode 100644 server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_sse_events_snapshot@sse_events_mock.snap.new diff --git a/server/packages/sandbox-agent/src/router.rs b/server/packages/sandbox-agent/src/router.rs index d22ce4e..8bf43d7 100644 --- a/server/packages/sandbox-agent/src/router.rs +++ b/server/packages/sandbox-agent/src/router.rs @@ -3863,11 +3863,12 @@ impl SessionManager { "message": prompt }); - let response_rx = runtime - .send_request(id, &request) - .ok_or_else(|| SandboxError::StreamError { - message: "failed to send pi prompt request".to_string(), - })?; + let response_rx = + runtime + .send_request(id, &request) + .ok_or_else(|| SandboxError::StreamError { + message: "failed to send pi prompt request".to_string(), + })?; let response = tokio::time::timeout(Duration::from_secs(30), response_rx) .await .map_err(|_| SandboxError::StreamError { @@ -3884,7 +3885,12 @@ impl SessionManager { let detail = response .get("error") .cloned() - .or_else(|| response.get("data").and_then(|data| data.get("error")).cloned()) + .or_else(|| { + response + .get("data") + .and_then(|data| data.get("error")) + .cloned() + }) .unwrap_or_else(|| response.clone()); return Err(SandboxError::InvalidRequest { message: format!("pi prompt failed: {detail}"), @@ -3927,7 +3933,12 @@ impl SessionManager { let detail = response .get("error") .cloned() - .or_else(|| response.get("data").and_then(|data| data.get("error")).cloned()) + .or_else(|| { + response + .get("data") + .and_then(|data| data.get("error")) + .cloned() + }) .unwrap_or_else(|| response.clone()); return Err(SandboxError::InvalidRequest { message: format!("pi set_thinking_level failed for '{level}': {detail}"), @@ -5529,7 +5540,8 @@ fn parse_pi_models_output(output: &str) -> AgentModelsResponse { value.eq_ignore_ascii_case("yes") || value.eq_ignore_ascii_case("no") }) }); - let supports_thinking = thinking_value.is_some_and(|value| value.eq_ignore_ascii_case("yes")); + let supports_thinking = + thinking_value.is_some_and(|value| value.eq_ignore_ascii_case("yes")); let (variants, default_variant) = if supports_thinking { (Some(pi_variants()), Some("medium".to_string())) } else { @@ -6590,7 +6602,10 @@ anthropic claude-sonnet-4-5-20250929 sonnet no .iter() .find(|model| model.id == "anthropic/claude-sonnet-4-5-20250929") .expect("anthropic model"); - assert_eq!(anthropic.variants.as_deref(), Some(&["off".to_string()][..])); + assert_eq!( + anthropic.variants.as_deref(), + Some(&["off".to_string()][..]) + ); assert_eq!(anthropic.default_variant.as_deref(), Some("off")); let openai = parsed @@ -6621,10 +6636,7 @@ groq llama-3.3-70b-versatile alias no .map(|model| (model.id.as_str(), model.default_variant.as_deref())) .collect::>(); - assert_eq!( - models, - vec![("groq/llama-3.3-70b-versatile", Some("off"))] - ); + assert_eq!(models, vec![("groq/llama-3.3-70b-versatile", Some("off"))]); } #[test] @@ -6646,11 +6658,7 @@ groq llama-3.3-70b-versatile alias no assert_eq!( models, - vec![( - "openrouter/qwen/qwen3-32b", - pi_variants(), - Some("medium") - )] + vec![("openrouter/qwen/qwen3-32b", pi_variants(), Some("medium"))] ); } @@ -6987,7 +6995,10 @@ mod pi_runtime_tests { let prompt_line = stdin_rx.recv().await.expect("prompt request"); let prompt_request: Value = serde_json::from_str(&prompt_line).expect("json request"); - assert_eq!(prompt_request.get("type").and_then(Value::as_str), Some("prompt")); + assert_eq!( + prompt_request.get("type").and_then(Value::as_str), + Some("prompt") + ); assert_eq!( prompt_request.get("message").and_then(Value::as_str), Some("Hello") @@ -7027,7 +7038,10 @@ mod pi_runtime_tests { let prompt_line = stdin_rx.recv().await.expect("prompt request"); let prompt_request: Value = serde_json::from_str(&prompt_line).expect("json request"); - assert_eq!(prompt_request.get("type").and_then(Value::as_str), Some("prompt")); + assert_eq!( + prompt_request.get("type").and_then(Value::as_str), + Some("prompt") + ); let prompt_id = prompt_request .get("id") .and_then(Value::as_i64) diff --git a/server/packages/sandbox-agent/tests/agent-flows/pi_rpc_integration.rs b/server/packages/sandbox-agent/tests/agent-flows/pi_rpc_integration.rs index 993836c..3113ae7 100644 --- a/server/packages/sandbox-agent/tests/agent-flows/pi_rpc_integration.rs +++ b/server/packages/sandbox-agent/tests/agent-flows/pi_rpc_integration.rs @@ -183,7 +183,10 @@ async fn pi_variant_high_applies_for_thinking_model() { create_pi_session(&app.app, session_id, Some(&model_id), Some("high")).await; let events = read_turn_stream_events(&app.app, session_id, Duration::from_secs(120)).await; - assert!(!events.is_empty(), "no events from pi thinking-variant stream"); + assert!( + !events.is_empty(), + "no events from pi thinking-variant stream" + ); assert!( !events.iter().any(is_unparsed_event), "agent.unparsed event encountered for thinking-variant session" diff --git a/server/packages/sandbox-agent/tests/http/agent_endpoints.rs b/server/packages/sandbox-agent/tests/http/agent_endpoints.rs index 7c7623d..de8b536 100644 --- a/server/packages/sandbox-agent/tests/http/agent_endpoints.rs +++ b/server/packages/sandbox-agent/tests/http/agent_endpoints.rs @@ -241,7 +241,10 @@ async fn pi_capabilities_and_models_expose_variants() { .iter() .filter_map(Value::as_str) .collect::>(); - assert!(!variant_ids.is_empty(), "pi model {model_id} has no variants"); + assert!( + !variant_ids.is_empty(), + "pi model {model_id} has no variants" + ); if variant_ids == vec!["off"] { assert_eq!( default_variant, "off", diff --git a/server/packages/sandbox-agent/tests/http/snapshots/http_endpoints__agent_endpoints__agent_endpoints_snapshots@agent_install_codex.snap.new b/server/packages/sandbox-agent/tests/http/snapshots/http_endpoints__agent_endpoints__agent_endpoints_snapshots@agent_install_codex.snap.new new file mode 100644 index 0000000..d01df04 --- /dev/null +++ b/server/packages/sandbox-agent/tests/http/snapshots/http_endpoints__agent_endpoints__agent_endpoints_snapshots@agent_install_codex.snap.new @@ -0,0 +1,6 @@ +--- +source: server/packages/sandbox-agent/tests/http/agent_endpoints.rs +assertion_line: 145 +expression: snapshot_status(status) +--- +status: 204 diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__multi_turn__assert_session_snapshot@multi_turn_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__multi_turn__assert_session_snapshot@multi_turn_mock.snap.new index 2a091af..d6e25cf 100644 --- a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__multi_turn__assert_session_snapshot@multi_turn_mock.snap.new +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__multi_turn__assert_session_snapshot@multi_turn_mock.snap.new @@ -99,33 +99,10 @@ second: status: in_progress seq: 4 type: item.started - - delta: - delta: "" - item_id: "" - native_item_id: "" + - item: + content_types: [] + kind: message + role: assistant + status: completed seq: 5 - type: item.delta - - delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 6 - type: item.delta - - delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 7 - type: item.delta - - delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 8 - type: item.delta - - delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 9 - type: item.delta + type: item.completed diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__permissions__assert_session_snapshot@permission_events_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__permissions__assert_session_snapshot@permission_events_mock.snap.new index d5c1b20..82e9c6c 100644 --- a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__permissions__assert_session_snapshot@permission_events_mock.snap.new +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__permissions__assert_session_snapshot@permission_events_mock.snap.new @@ -67,171 +67,3 @@ expression: value native_item_id: "" seq: 10 type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 11 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 12 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 13 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 14 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 15 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 16 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 17 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 18 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 19 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 20 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 21 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 22 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 23 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 24 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 25 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 26 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 27 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 28 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 29 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 30 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 31 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 32 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 33 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 34 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 35 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 36 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 37 - type: item.delta -- delta: - delta: "" - item_id: "" - native_item_id: "" - seq: 38 - type: item.delta diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__questions__assert_session_snapshot@question_reply_events_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__questions__assert_session_snapshot@question_reply_events_mock.snap.new index f414271..0428c57 100644 --- a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__questions__assert_session_snapshot@question_reply_events_mock.snap.new +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__questions__assert_session_snapshot@question_reply_events_mock.snap.new @@ -97,9 +97,41 @@ expression: value native_item_id: "" seq: 15 type: item.delta -- question: - id: "" - options: 4 - status: requested +- delta: + delta: "" + item_id: "" + native_item_id: "" seq: 16 - type: question.requested + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 17 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 18 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 19 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 20 + type: item.delta +- item: + content_types: + - text + kind: message + role: assistant + status: completed + seq: 21 + type: item.completed diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__assert_session_snapshot@concurrency_events_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__assert_session_snapshot@concurrency_events_mock.snap.new index a6e0065..38d2285 100644 --- a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__assert_session_snapshot@concurrency_events_mock.snap.new +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__assert_session_snapshot@concurrency_events_mock.snap.new @@ -44,6 +44,47 @@ session_a: native_item_id: "" seq: 6 type: item.delta +session_b: + - metadata: true + seq: 1 + session: started + type: session.started + - item: + content_types: + - text + kind: message + role: user + status: in_progress + seq: 2 + type: item.started + - delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 3 + type: item.delta + - item: + content_types: + - text + kind: message + role: user + status: completed + seq: 4 + type: item.completed + - item: + content_types: + - text + kind: message + role: assistant + status: in_progress + seq: 5 + type: item.started + - delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 6 + type: item.delta - delta: delta: "" item_id: "" @@ -62,44 +103,15 @@ session_a: native_item_id: "" seq: 9 type: item.delta -session_b: - - metadata: true - seq: 1 - session: started - type: session.started - - item: - content_types: - - text - kind: message - role: user - status: in_progress - seq: 2 - type: item.started - delta: delta: "" item_id: "" native_item_id: "" - seq: 3 + seq: 10 type: item.delta - - item: - content_types: - - text - kind: message - role: user - status: completed - seq: 4 - type: item.completed - - item: - content_types: - - text - kind: message - role: assistant - status: in_progress - seq: 5 - type: item.started - delta: delta: "" item_id: "" native_item_id: "" - seq: 6 + seq: 11 type: item.delta diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_http_events_snapshot@http_events_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_http_events_snapshot@http_events_mock.snap.new index da365cc..158b730 100644 --- a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_http_events_snapshot@http_events_mock.snap.new +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_http_events_snapshot@http_events_mock.snap.new @@ -37,10 +37,33 @@ expression: normalized status: in_progress seq: 5 type: item.started -- item: - content_types: [] - kind: message - role: assistant - status: completed +- delta: + delta: "" + item_id: "" + native_item_id: "" seq: 6 - type: item.completed + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 7 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 8 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 9 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 10 + type: item.delta diff --git a/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_sse_events_snapshot@sse_events_mock.snap.new b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_sse_events_snapshot@sse_events_mock.snap.new new file mode 100644 index 0000000..8df795c --- /dev/null +++ b/server/packages/sandbox-agent/tests/sessions/snapshots/sessions__sessions__session_lifecycle__run_sse_events_snapshot@sse_events_mock.snap.new @@ -0,0 +1,75 @@ +--- +source: server/packages/sandbox-agent/tests/sessions/../common/http.rs +assertion_line: 1039 +expression: normalized +--- +- metadata: true + seq: 1 + session: started + type: session.started +- item: + content_types: + - text + kind: message + role: user + status: in_progress + seq: 2 + type: item.started +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 3 + type: item.delta +- item: + content_types: + - text + kind: message + role: user + status: completed + seq: 4 + type: item.completed +- item: + content_types: + - text + kind: message + role: assistant + status: in_progress + seq: 5 + type: item.started +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 6 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 7 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 8 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 9 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 10 + type: item.delta +- delta: + delta: "" + item_id: "" + native_item_id: "" + seq: 11 + type: item.delta diff --git a/server/packages/universal-agent-schema/src/agents/pi.rs b/server/packages/universal-agent-schema/src/agents/pi.rs index a863c43..7aa9932 100644 --- a/server/packages/universal-agent-schema/src/agents/pi.rs +++ b/server/packages/universal-agent-schema/src/agents/pi.rs @@ -12,6 +12,7 @@ use crate::{ pub struct PiEventConverter { tool_result_buffers: HashMap, tool_result_started: HashSet, + message_completed: HashSet, message_errors: HashSet, message_reasoning: HashMap, message_text: HashMap, @@ -121,6 +122,7 @@ impl PiEventConverter { return Ok(Vec::new()); } let message_id = self.ensure_message_id(extract_message_id(raw)); + self.message_completed.remove(&message_id); self.message_started.insert(message_id.clone()); let content = message.and_then(parse_message_content).unwrap_or_default(); let entry = self.message_text.entry(message_id.clone()).or_default(); @@ -210,15 +212,24 @@ impl PiEventConverter { self.clear_last_message_id(Some(&message_id)); return Ok(Vec::new()); } + if self.message_completed.contains(&message_id) { + self.clear_last_message_id(Some(&message_id)); + return Ok(Vec::new()); + } let message = raw .get("message") .or_else(|| assistant_event.get("message")); let conversion = self.complete_message(Some(message_id.clone()), message); + self.message_completed.insert(message_id.clone()); self.clear_last_message_id(Some(&message_id)); Ok(vec![conversion]) } "error" => { let message_id = self.ensure_message_id(message_id); + if self.message_completed.contains(&message_id) { + self.clear_last_message_id(Some(&message_id)); + return Ok(Vec::new()); + } let error_text = assistant_event .get("error") .or_else(|| raw.get("error")) @@ -228,6 +239,7 @@ impl PiEventConverter { self.message_text.remove(&message_id); self.message_errors.insert(message_id.clone()); self.message_started.remove(&message_id); + self.message_completed.insert(message_id.clone()); self.clear_last_message_id(Some(&message_id)); let item = UniversalItem { item_id: String::new(), @@ -261,7 +273,12 @@ impl PiEventConverter { self.clear_last_message_id(Some(&message_id)); return Ok(Vec::new()); } + if self.message_completed.contains(&message_id) { + self.clear_last_message_id(Some(&message_id)); + return Ok(Vec::new()); + } let conversion = self.complete_message(Some(message_id.clone()), message); + self.message_completed.insert(message_id.clone()); self.clear_last_message_id(Some(&message_id)); Ok(vec![conversion]) } @@ -479,7 +496,7 @@ fn status_event(label: &str, raw: &Value) -> EventConversion { kind: ItemKind::Status, role: Some(ItemRole::System), content: vec![ContentPart::Status { - label: format!("pi.{label}"), + label: pi_status_label(label), detail, }], status: ItemStatus::Completed, @@ -490,6 +507,14 @@ fn status_event(label: &str, raw: &Value) -> EventConversion { ) } +fn pi_status_label(label: &str) -> String { + match label { + "turn_end" => "turn.completed".to_string(), + "agent_end" => "session.idle".to_string(), + _ => format!("pi.{label}"), + } +} + fn item_delta(message_id: Option, delta: String) -> EventConversion { EventConversion::new( UniversalEventType::ItemDelta, diff --git a/server/packages/universal-agent-schema/tests/pi_conversion.rs b/server/packages/universal-agent-schema/tests/pi_conversion.rs index 971fde8..eb4cd72 100644 --- a/server/packages/universal-agent-schema/tests/pi_conversion.rs +++ b/server/packages/universal-agent-schema/tests/pi_conversion.rs @@ -216,6 +216,56 @@ fn pi_unknown_event_returns_error() { assert!(converter.event_to_universal(&event).is_err()); } +#[test] +fn pi_turn_and_agent_end_emit_terminal_status_labels() { + let mut converter = PiEventConverter::default(); + + let turn_end = parse_event(json!({ + "type": "turn_end", + "sessionId": "session-1" + })); + let turn_events = converter + .event_to_universal(&turn_end) + .expect("turn_end conversions"); + assert_eq!(turn_events[0].event_type, UniversalEventType::ItemCompleted); + if let UniversalEventData::Item(item) = &turn_events[0].data { + assert_eq!(item.item.kind, ItemKind::Status); + assert!( + matches!( + item.item.content.first(), + Some(ContentPart::Status { label, .. }) if label == "turn.completed" + ), + "turn_end should map to turn.completed status" + ); + } else { + panic!("expected item event"); + } + + let agent_end = parse_event(json!({ + "type": "agent_end", + "sessionId": "session-1" + })); + let agent_events = converter + .event_to_universal(&agent_end) + .expect("agent_end conversions"); + assert_eq!( + agent_events[0].event_type, + UniversalEventType::ItemCompleted + ); + if let UniversalEventData::Item(item) = &agent_events[0].data { + assert_eq!(item.item.kind, ItemKind::Status); + assert!( + matches!( + item.item.content.first(), + Some(ContentPart::Status { label, .. }) if label == "session.idle" + ), + "agent_end should map to session.idle status" + ); + } else { + panic!("expected item event"); + } +} + #[test] fn pi_message_done_completes_without_message_end() { let mut converter = PiEventConverter::default(); @@ -263,6 +313,63 @@ fn pi_message_done_completes_without_message_end() { } } +#[test] +fn pi_message_done_then_message_end_does_not_double_complete() { + 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 _ = 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 _ = 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.len(), 1); + assert_eq!(done_events[0].event_type, UniversalEventType::ItemCompleted); + + 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!( + end_events.is_empty(), + "message_end after done should not emit a second completion" + ); +} + #[test] fn pi_message_end_error_surfaces_failed_status_and_error_text() { let mut converter = PiEventConverter::default();