From ed674e534b63085c1aa032d3872ea820e995e5f5 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 23 Jul 2026 11:52:27 -0700 Subject: [PATCH 1/3] fix: avoid duplicate ATIF user steps on continuations Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 194 ++++++++++++++++++++++++-- crates/core/tests/unit/atif_tests.rs | 125 +++++++++++++++-- 2 files changed, 299 insertions(+), 20 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 83c28dd24..e9e78cbf2 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -18,7 +18,7 @@ //! //! | NeMo Relay Event | ATIF Step | Notes | //! |-----------------|-------------------------|--------------------------------------| -//! | LLM Start | `user` step | Messages extracted from LlmRequest | +//! | LLM Start | `user` step | Fresh user turns only; same-turn continuations stay on the agent step | //! | LLM End | `agent` step | Response content, tool_calls promoted| //! | Tool Start | *(skipped)* | tool_calls come from LLM End instead | //! | Tool End | agent observation | Correlated by `source_call_id` | @@ -38,7 +38,7 @@ use uuid::Uuid; use crate::api::event::{Event, EventNormalizationExt}; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; -use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use crate::codec::request::{AnnotatedLlmRequest, ContentPart, Message, MessageContent}; use crate::codec::response::AnnotatedLlmResponse; use crate::error::Result; use crate::json::Json; @@ -287,6 +287,12 @@ pub struct AtifStepExtra { pub tool_invocations: Option>, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RequestTurnState { + FreshUser, + Continuation, +} + /// A complete ATIF trajectory. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AtifTrajectory { @@ -930,9 +936,10 @@ fn extract_reasoning_content(output: &Json) -> Option { /// Extract the latest user-facing message from an LLM request payload. /// /// LLM start inputs typically contain `{ "messages": [...], "model": "...", -/// "max_tokens": ..., "tools": [...], "stream": ... }`. For the ATIF user step -/// we emit a schema-compatible message value (string or content-part array) -/// and preserve the full LLM request in `Step.extra.llm_request`. +/// "max_tokens": ..., "tools": [...], "stream": ... }`. For ATIF we emit a +/// schema-compatible message value (string or content-part array) and preserve +/// the full LLM request in `Step.extra.llm_request`, either on the user step or +/// on the matching agent step for same-turn continuations. /// /// Returns the latest user message content if present, a prompt if present, or /// a stringified representation of the input as a last resort. @@ -964,6 +971,94 @@ fn extract_user_messages(input: &Json) -> Json { atif_content_value(input) } +fn llm_start_user_step_message(event: &Event, input: &Json, has_paired_end: bool) -> Option { + let turn_state = event + .annotated_request() + .and_then(|request| request_turn_state_from_annotation(request.as_ref())) + .or_else(|| request_turn_state_from_raw(input)); + + if has_paired_end && matches!(turn_state, Some(RequestTurnState::Continuation)) { + return None; + } + + event + .annotated_request() + .and_then(|request| atif_message_from_annotated_request(request.as_ref())) + .or_else(|| Some(extract_user_messages(input))) +} + +fn request_turn_state_from_raw(input: &Json) -> Option { + let object = input.as_object()?; + if let Some(messages) = object.get("messages").and_then(Json::as_array) + && let Some(state) = chat_messages_turn_state(messages) + { + return Some(state); + } + if let Some(input_items) = object.get("input") + && let Some(state) = openai_responses_turn_state(input_items) + { + return Some(state); + } + object.get("prompt").map(|_| RequestTurnState::FreshUser) +} + +fn request_turn_state_from_annotation(request: &AnnotatedLlmRequest) -> Option { + request + .messages + .iter() + .rev() + .find_map(annotated_message_turn_state) +} + +fn chat_messages_turn_state(messages: &[Json]) -> Option { + messages + .iter() + .rev() + .filter_map(Json::as_object) + .find_map(|message| match message.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") | None => Some(RequestTurnState::FreshUser), + Some(_) => Some(RequestTurnState::Continuation), + }) +} + +fn openai_responses_turn_state(input: &Json) -> Option { + if input.is_string() { + return Some(RequestTurnState::FreshUser); + } + + input + .as_array()? + .iter() + .rev() + .filter_map(Json::as_object) + .find_map(openai_responses_item_turn_state) +} + +fn openai_responses_item_turn_state( + item: &serde_json::Map, +) -> Option { + match item.get("type").and_then(Json::as_str) { + Some("message") => match item.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") | None => Some(RequestTurnState::FreshUser), + Some(_) => Some(RequestTurnState::Continuation), + }, + Some( + "function_call" + | "custom_tool_call" + | "mcp_tool_call" + | "tool_call" + | "function_call_output" + | "custom_tool_call_output" + | "mcp_tool_result" + | "tool_result" + | "reasoning", + ) => Some(RequestTurnState::Continuation), + _ => None, + } +} + fn openai_responses_input_message(input: &Json) -> Option { if input.is_string() { return Some(atif_content_value(input)); @@ -1084,6 +1179,64 @@ fn atif_message_from_annotated_request(request: &AnnotatedLlmRequest) -> Option< } } +fn annotated_message_turn_state(message: &Message) -> Option { + match message { + Message::System { .. } | Message::Developer { .. } => None, + Message::User { content, .. } => Some(if annotated_content_starts_new_turn(content) { + RequestTurnState::FreshUser + } else { + RequestTurnState::Continuation + }), + Message::Assistant { .. } + | Message::Tool { .. } + | Message::Function { .. } + | Message::ToolCallItem { .. } + | Message::ToolResultItem { .. } => Some(RequestTurnState::Continuation), + Message::ProviderNative { value, .. } => provider_native_turn_state(value), + } +} + +fn annotated_content_starts_new_turn(content: &MessageContent) -> bool { + match content { + MessageContent::Text(_) => true, + MessageContent::Parts(parts) => { + let has_user_input = parts.iter().any(|part| { + matches!( + part, + ContentPart::Text { .. } + | ContentPart::ImageUrl { .. } + | ContentPart::Image { .. } + | ContentPart::Audio { .. } + | ContentPart::File { .. } + ) + }); + if has_user_input { + return true; + } + let has_tool_continuation = parts.iter().any(|part| { + matches!( + part, + ContentPart::ToolUse { .. } | ContentPart::ToolResult { .. } + ) + }); + !has_tool_continuation + } + } +} + +fn provider_native_turn_state(value: &Json) -> Option { + let object = value.as_object()?; + if object.get("type").is_some() { + return openai_responses_item_turn_state(object); + } + match object.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") => Some(RequestTurnState::FreshUser), + Some(_) => Some(RequestTurnState::Continuation), + None => None, + } +} + fn atif_message_from_annotated_response(response: &AnnotatedLlmResponse) -> Option { match &response.message { Some(MessageContent::Text(text)) => Some(Json::String(text.clone())), @@ -1380,6 +1533,7 @@ fn json_string_at(value: &Json, path: &[&str]) -> Option { struct EventLookupMaps { name_map: std::collections::HashMap, start_ts_map: std::collections::HashMap>, + llm_end_uuids: HashSet, llm_start_model_names: HashMap, tool_call_ids: std::collections::HashMap, suppressed_llm_events: HashSet, @@ -1407,6 +1561,7 @@ impl EventLookupMaps { ) -> Self { let mut name_map = std::collections::HashMap::new(); let mut start_ts_map = std::collections::HashMap::new(); + let mut llm_end_uuids = HashSet::new(); let mut llm_start_model_names = HashMap::new(); for event in events { if is_start_event(event) { @@ -1418,11 +1573,17 @@ impl EventLookupMaps { llm_start_model_names.insert(event.uuid(), model_name); } } + if event.scope_category() == Some(crate::api::event::ScopeCategory::End) + && event.category().map(|category| category.as_str()) == Some("llm") + { + llm_end_uuids.insert(event.uuid()); + } } let llm_dedupe = build_llm_dedupe(llm_dedupe_events); Self { name_map, start_ts_map, + llm_end_uuids, llm_start_model_names, tool_call_ids: build_tool_call_correlations(tool_correlation_events), suppressed_llm_events: llm_dedupe.suppressed_events, @@ -1928,6 +2089,7 @@ struct PendingAgentStep { step_idx: Option, ancestry: Option, invocation: Option, + llm_request: Option, llm_response: Option, tool_ancestry: Vec, tool_invocations: Vec, @@ -1947,7 +2109,7 @@ impl PendingAgentStep { let extra = AtifStepExtra { ancestry, invocation: self.invocation.take(), - llm_request: None, + llm_request: self.llm_request.take(), llm_response: self.llm_response.take(), event_payload: None, tool_ancestry: std::mem::take(&mut self.tool_ancestry), @@ -1977,6 +2139,10 @@ impl PendingAgentStep { self.tool_call_order = tool_call_order; } + fn stash_llm_request(&mut self, llm_request: Json) { + self.llm_request = Some(llm_request); + } + fn push_tool_metadata(&mut self, ancestry: AtifAncestry, invocation: AtifInvocationInfo) { self.tool_ancestry.push(ancestry); self.tool_invocations.push(invocation); @@ -2299,6 +2465,11 @@ impl StepConversionState { }; let content = unwrap_llm_request(input); self.current_reasoning_effort = extract_reasoning_effort(&content); + let has_paired_end = lookups.llm_end_uuids.contains(&event.uuid()); + let Some(message) = llm_start_user_step_message(event, &content, has_paired_end) else { + self.current_agent.stash_llm_request(content); + return; + }; let extra = AtifStepExtra { ancestry: build_ancestry(event, &lookups.name_map), invocation: None, @@ -2311,10 +2482,7 @@ impl StepConversionState { self.steps.push(AtifStep { step_id: 0, source: "user".to_string(), - message: event - .annotated_request() - .and_then(|request| atif_message_from_annotated_request(request)) - .unwrap_or_else(|| extract_user_messages(&content)), + message, timestamp: Some(event.timestamp().to_rfc3339()), model_name: None, reasoning_effort: None, @@ -3382,8 +3550,10 @@ fn events_to_steps_for_agent( /// Mapping logic: /// 1. Sort events by timestamp. /// 2. For each LLM pair: -/// - Start event → user step (message = extracted `messages` array from -/// unwrapped LlmRequest content, stripping `max_tokens`/`model`/etc.) +/// - Start event → user step when the request begins a fresh user turn +/// - Start events that continue the same turn after tool work stash +/// `llm_request` on the matching agent step instead of repeating the user +/// message /// - End event → agent step (message = extracted content, metrics from /// token_usage, tool_calls promoted to AtifToolCall entries with parsed /// JSON arguments) diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 10f203bc3..b095f65c7 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -4890,7 +4890,8 @@ fn test_exporter_user_message_extraction() { #[test] fn test_exporter_full_agent_loop() { // Simulate a complete agent loop: LLM→tool_calls→observations→LLM→final answer - // This should produce 5 steps: user, agent+tool_calls, merged obs, user, agent + // The second request continues the same user turn after tool work, so it + // should not create a duplicate user step. let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); let llm1_uuid = Uuid::now_v7(); let llm2_uuid = Uuid::now_v7(); @@ -4984,8 +4985,8 @@ fn test_exporter_full_agent_loop() { let trajectory = exporter.export().unwrap(); assert_atif_v17_shape(&trajectory); - // Expected: user, agent+tool_calls+observations, user, agent - assert_eq!(trajectory.steps.len(), 4); + // Expected: user, agent+tool_calls+observations, agent + assert_eq!(trajectory.steps.len(), 3); assert_eq!(trajectory.steps[0].source, "user"); assert_eq!(trajectory.steps[0].step_id, 1); @@ -4997,17 +4998,23 @@ fn test_exporter_full_agent_loop() { assert_eq!(tcs[0].function_name, "get_weather"); assert_eq!(tcs[1].function_name, "get_population"); - assert_eq!(trajectory.steps[2].source, "user"); - assert_eq!(trajectory.steps[2].step_id, 3); let obs = trajectory.steps[1].observation.as_ref().unwrap(); assert_eq!(obs.results.len(), 2); - assert_eq!(trajectory.steps[3].source, "agent"); - assert_eq!(trajectory.steps[3].step_id, 4); + assert_eq!(trajectory.steps[2].source, "agent"); + assert_eq!(trajectory.steps[2].step_id, 3); assert_eq!( - trajectory.steps[3].message, + trajectory.steps[2].message, json!("The weather in SF is 62°F and foggy. Population is 873,965.") ); + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["messages"][0]["content"], + json!("What is the weather and population of SF?") + ); + assert_eq!(llm_request["messages"][3]["tool_call_id"], json!("c2")); // Final metrics should aggregate both LLM calls let fm = trajectory.final_metrics.as_ref().unwrap(); @@ -5015,6 +5022,108 @@ fn test_exporter_full_agent_loop() { assert_eq!(fm.total_completion_tokens, Some(80)); } +#[test] +fn test_exporter_skips_duplicate_user_step_for_openai_responses_tool_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + })) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "shell", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "pip is missing" + } + ] + })) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_2", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I found that pip is missing."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Fix pip")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["input"][2]["type"], + json!("function_call_output") + ); + assert_eq!(llm_request["input"][2]["output"], json!("pip is missing")); +} + #[test] fn test_reasoning_content_extracted() { // When an LLM End event carries output["reasoning"], the agent step From 7d07861085b0ebeca30b050cb895228773f37750 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 23 Jul 2026 12:49:46 -0700 Subject: [PATCH 2/3] fix: classify OpenAI Responses continuations in ATIF Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 12 +-- crates/core/tests/unit/atif_tests.rs | 106 ++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 11 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index e9e78cbf2..c4d1ca2ee 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -1044,17 +1044,7 @@ fn openai_responses_item_turn_state( Some("user") | None => Some(RequestTurnState::FreshUser), Some(_) => Some(RequestTurnState::Continuation), }, - Some( - "function_call" - | "custom_tool_call" - | "mcp_tool_call" - | "tool_call" - | "function_call_output" - | "custom_tool_call_output" - | "mcp_tool_result" - | "tool_result" - | "reasoning", - ) => Some(RequestTurnState::Continuation), + Some(_) => Some(RequestTurnState::Continuation), _ => None, } } diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index b095f65c7..d7d4171d8 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -5124,6 +5124,112 @@ fn test_exporter_skips_duplicate_user_step_for_openai_responses_tool_continuatio assert_eq!(llm_request["input"][2]["output"], json!("pip is missing")); } +#[test] +fn test_exporter_skips_duplicate_user_step_for_openai_responses_shell_call_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + let codec = OpenAIResponsesCodec; + + let llm1_start_input = json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + }); + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(llm1_start_input) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + + let llm2_start_input = json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "shell_call", + "id": "sh_1", + "call_id": "sh_call", + "action": {"commands": ["which pip"]} + } + ] + }); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(llm2_start_input.clone()) + .annotated_request( + codec + .decode(&LlmRequest { + headers: serde_json::Map::new(), + content: llm2_start_input, + }) + .unwrap(), + ) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_2", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I found the missing pip binary."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Fix pip")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!(llm_request["input"][1]["type"], json!("shell_call")); + assert_eq!(llm_request["input"][1]["call_id"], json!("sh_call")); +} + #[test] fn test_reasoning_content_extracted() { // When an LLM End event carries output["reasoning"], the agent step From 31b10795c5ccb882df5d265211f9070cd9776ba8 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 23 Jul 2026 14:30:25 -0700 Subject: [PATCH 3/3] fix: cover Anthropic continuations and unpaired ATIF ends Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 62 ++++++- crates/core/tests/unit/atif_tests.rs | 237 ++++++++++++++++++++++++++ 2 files changed, 294 insertions(+), 5 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index c4d1ca2ee..5571dfd81 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -1015,11 +1015,62 @@ fn chat_messages_turn_state(messages: &[Json]) -> Option { .iter() .rev() .filter_map(Json::as_object) - .find_map(|message| match message.get("role").and_then(Json::as_str) { - Some("system" | "developer") => None, - Some("user") | None => Some(RequestTurnState::FreshUser), - Some(_) => Some(RequestTurnState::Continuation), - }) + .find_map(raw_chat_message_turn_state) +} + +fn raw_chat_message_turn_state( + message: &serde_json::Map, +) -> Option { + match message.get("role").and_then(Json::as_str) { + Some("system" | "developer") => None, + Some("user") | None => Some(match message.get("content") { + Some(content) if raw_content_starts_new_turn(content) => RequestTurnState::FreshUser, + Some(_) => RequestTurnState::Continuation, + None => RequestTurnState::FreshUser, + }), + Some(_) => Some(RequestTurnState::Continuation), + } +} + +fn raw_content_starts_new_turn(content: &Json) -> bool { + match content { + Json::String(_) => true, + Json::Array(parts) => { + if parts.iter().any(raw_content_part_is_user_input) { + return true; + } + !parts.iter().any(raw_content_part_is_tool_continuation) + } + Json::Object(_) => { + raw_content_part_is_user_input(content) + || !raw_content_part_is_tool_continuation(content) + } + _ => true, + } +} + +fn raw_content_part_is_user_input(part: &Json) -> bool { + matches!( + part.get("type").and_then(Json::as_str), + Some( + "text" + | "input_text" + | "image_url" + | "image" + | "input_image" + | "audio" + | "input_audio" + | "file" + | "document" + ) + ) +} + +fn raw_content_part_is_tool_continuation(part: &Json) -> bool { + matches!( + part.get("type").and_then(Json::as_str), + Some("tool_result" | "tool_use") + ) } fn openai_responses_turn_state(input: &Json) -> Option { @@ -1565,6 +1616,7 @@ impl EventLookupMaps { } if event.scope_category() == Some(crate::api::event::ScopeCategory::End) && event.category().map(|category| category.as_str()) == Some("llm") + && event.data().is_some() { llm_end_uuids.insert(event.uuid()); } diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index d7d4171d8..050134f3a 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -5230,6 +5230,243 @@ fn test_exporter_skips_duplicate_user_step_for_openai_responses_shell_call_conti assert_eq!(llm_request["input"][1]["call_id"], json!("sh_call")); } +#[test] +fn test_exporter_skips_duplicate_user_step_for_anthropic_tool_result_continuation() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "claude-sonnet-4", + "system": "Be concise.", + "messages": [{"role": "user", "content": "Find the file."}] + })) + .model_name("claude-sonnet-4") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [ + {"type": "text", "text": "I will search for it."}, + { + "type": "tool_use", + "id": "toolu_01", + "name": "search", + "input": {"query": "README.md"} + } + ], + "stop_reason": "tool_use" + })) + .model_name("claude-sonnet-4") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "claude-sonnet-4", + "system": "Be concise.", + "messages": [ + {"role": "user", "content": "Find the file."}, + {"role": "assistant", "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "search", + "input": {"query": "README.md"} + } + ]}, + {"role": "user", "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": "README.md", + "is_error": false + } + ]} + ] + })) + .model_name("claude-sonnet-4") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("anthropic.messages") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "msg_02", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [{"type": "text", "text": "README.md is the relevant file."}], + "stop_reason": "end_turn" + })) + .model_name("claude-sonnet-4") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state + .events + .extend([llm1_start, llm1_end, llm2_start, llm2_end]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 3); + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 1); + assert_eq!(user_steps[0].message, json!("Find the file.")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[2].extra.clone().unwrap()).unwrap(); + let llm_request = final_extra.llm_request.unwrap(); + assert_eq!( + llm_request["messages"][2]["content"][0]["type"], + json!("tool_result") + ); + assert_eq!( + llm_request["messages"][2]["content"][0]["tool_use_id"], + json!("toolu_01") + ); +} + +#[test] +fn test_exporter_does_not_stash_unpaired_continuation_request_on_next_agent_step() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm1_uuid = Uuid::now_v7(); + let llm2_uuid = Uuid::now_v7(); + let llm3_uuid = Uuid::now_v7(); + + let llm1_start = event_builder(llm1_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }] + })) + .model_name("switchyard") + .build(); + let llm1_end = event_builder(llm1_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_1", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will inspect the environment."}] + }] + })) + .model_name("switchyard") + .build(); + let llm2_start = event_builder(llm2_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Fix pip"}] + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "shell", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "pip is missing" + } + ] + })) + .model_name("switchyard") + .build(); + let llm2_end = event_builder(llm2_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .model_name("switchyard") + .build(); + let llm3_start = event_builder(llm3_uuid, EventType::Start) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "switchyard", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Try ensurepip instead"}] + }] + })) + .model_name("switchyard") + .build(); + let llm3_end = event_builder(llm3_uuid, EventType::End) + .name("openai.responses") + .scope_type(ScopeType::Llm) + .output(json!({ + "id": "resp_3", + "model": "switchyard", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Use ensurepip to restore pip."}] + }] + })) + .model_name("switchyard") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.extend([ + llm1_start, llm1_end, llm2_start, llm2_end, llm3_start, llm3_end, + ]); + } + + let trajectory = exporter.export().unwrap(); + assert_atif_v17_shape(&trajectory); + assert_eq!(trajectory.steps.len(), 5); + + let user_steps = trajectory + .steps + .iter() + .filter(|step| step.source == "user") + .collect::>(); + assert_eq!(user_steps.len(), 3); + assert_eq!(user_steps[0].message, json!("Fix pip")); + assert_eq!(user_steps[1].message, json!("Fix pip")); + assert_eq!(user_steps[2].message, json!("Try ensurepip instead")); + + let final_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[4].extra.clone().unwrap()).unwrap(); + assert!(final_extra.llm_request.is_none()); + + let llm3_user_extra: AtifStepExtra = + serde_json::from_value(trajectory.steps[3].extra.clone().unwrap()).unwrap(); + let llm3_request = llm3_user_extra.llm_request.unwrap(); + assert_eq!( + llm3_request["input"][0]["content"][0]["text"], + json!("Try ensurepip instead") + ); +} + #[test] fn test_reasoning_content_extracted() { // When an LLM End event carries output["reasoning"], the agent step