diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 98775e3d3a..944c86a31e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -263,7 +263,11 @@ impl ModelClient { let mut payload_json = serde_json::to_value(&payload)?; if azure_workaround { - attach_item_ids(&mut payload_json, &input_with_instructions); + attach_item_ids( + &mut payload_json, + &input_with_instructions, + &self.conversation_id, + ); } let max_attempts = self.provider.request_max_retries(); @@ -607,7 +611,11 @@ struct ResponseCompletedOutputTokensDetails { reasoning_tokens: i64, } -fn attach_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) { +fn attach_item_ids( + payload_json: &mut Value, + original_items: &[ResponseItem], + conversation_id: &ConversationId, +) { let Some(input_value) = payload_json.get_mut("input") else { return; }; @@ -615,25 +623,47 @@ fn attach_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) { return; }; - for (value, item) in items.iter_mut().zip(original_items.iter()) { - if let ResponseItem::Reasoning { id, .. } - | ResponseItem::Message { id: Some(id), .. } - | ResponseItem::WebSearchCall { id: Some(id), .. } - | ResponseItem::FunctionCall { id: Some(id), .. } - | ResponseItem::LocalShellCall { id: Some(id), .. } - | ResponseItem::CustomToolCall { id: Some(id), .. } = item - { - if id.is_empty() { - continue; - } - - if let Some(obj) = value.as_object_mut() { - obj.insert("id".to_string(), Value::String(id.clone())); - } + for (index, (value, item)) in items.iter_mut().zip(original_items.iter()).enumerate() { + if let Some(obj) = value.as_object_mut() { + let item_id = response_item_id(item) + .filter(|id| !id.is_empty()) + .unwrap_or_else(|| format!("{conversation_id}-input-{index}")); + obj.insert("id".to_string(), Value::String(item_id)); } } } +fn response_item_id(item: &ResponseItem) -> Option { + match item { + ResponseItem::Reasoning { id, .. } => (!id.is_empty()).then_some(id.clone()), + ResponseItem::Message { id, .. } => id.as_ref().filter(|v| !v.is_empty()).cloned(), + ResponseItem::WebSearchCall { id, .. } => id.as_ref().filter(|v| !v.is_empty()).cloned(), + ResponseItem::FunctionCall { id, call_id, .. } => id + .as_ref() + .filter(|v| !v.is_empty()) + .cloned() + .or_else(|| (!call_id.is_empty()).then_some(call_id.clone())), + ResponseItem::LocalShellCall { id, call_id, .. } => id + .as_ref() + .filter(|v| !v.is_empty()) + .cloned() + .or_else(|| call_id.as_ref().filter(|v| !v.is_empty()).cloned()), + ResponseItem::CustomToolCall { id, call_id, .. } => id + .as_ref() + .filter(|v| !v.is_empty()) + .cloned() + .or_else(|| (!call_id.is_empty()).then_some(call_id.clone())), + ResponseItem::FunctionCallOutput { call_id, .. } => { + (!call_id.is_empty()).then_some(call_id.clone()) + } + ResponseItem::CustomToolCallOutput { call_id, .. } => { + (!call_id.is_empty()).then_some(call_id.clone()) + } + ResponseItem::GhostSnapshot { .. } => None, + ResponseItem::Other => None, + } +} + fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { let primary = parse_rate_limit_window( headers, diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index d0aa1d8188..855adce507 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -8,7 +8,6 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; -use codex_protocol::models::WebSearchAction; use codex_protocol::user_input::UserInput; use tracing::warn; use uuid::Uuid; @@ -111,14 +110,17 @@ pub fn parse_turn_item(item: &ResponseItem) -> Option { raw_content, })) } - ResponseItem::WebSearchCall { - id, - action: WebSearchAction::Search { query }, - .. - } => Some(TurnItem::WebSearch(WebSearchItem { - id: id.clone().unwrap_or_default(), - query: query.clone(), - })), + ResponseItem::WebSearchCall { id, action, .. } => { + if action.is_search() + && let Some(query) = action.query() + { + return Some(TurnItem::WebSearch(WebSearchItem { + id: id.clone().unwrap_or_default(), + query: query.to_string(), + })); + } + None + } _ => None, } } @@ -305,9 +307,7 @@ mod tests { let item = ResponseItem::WebSearchCall { id: Some("ws_1".to_string()), status: Some("completed".to_string()), - action: WebSearchAction::Search { - query: "weather".to_string(), - }, + action: WebSearchAction::search("weather"), }; let turn_item = parse_turn_item(&item).expect("expected web search turn item"); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 79c158e084..cb707f19cf 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -991,9 +991,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { prompt.input.push(ResponseItem::WebSearchCall { id: Some("web-search-id".into()), status: Some("completed".into()), - action: WebSearchAction::Search { - query: "weather".into(), - }, + action: WebSearchAction::search("weather"), }); prompt.input.push(ResponseItem::FunctionCall { id: Some("function-id".into()), @@ -1020,6 +1018,13 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { name: "custom_tool".into(), input: "{}".into(), }); + prompt.input.push(ResponseItem::Message { + id: None, + role: "user".into(), + content: vec![ContentItem::InputText { + text: "follow up question".into(), + }], + }); let mut stream = client .stream(&prompt) @@ -1043,13 +1048,18 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { assert_eq!(body["store"], serde_json::Value::Bool(true)); assert_eq!(body["stream"], serde_json::Value::Bool(true)); - assert_eq!(body["input"].as_array().map(Vec::len), Some(6)); + assert_eq!(body["input"].as_array().map(Vec::len), Some(7)); assert_eq!(body["input"][0]["id"].as_str(), Some("reasoning-id")); assert_eq!(body["input"][1]["id"].as_str(), Some("message-id")); assert_eq!(body["input"][2]["id"].as_str(), Some("web-search-id")); assert_eq!(body["input"][3]["id"].as_str(), Some("function-id")); assert_eq!(body["input"][4]["id"].as_str(), Some("local-shell-id")); assert_eq!(body["input"][5]["id"].as_str(), Some("custom-tool-id")); + let expected_user_id = format!("{conversation_id}-input-6"); + assert_eq!( + body["input"][6]["id"].as_str(), + Some(expected_user_id.as_str()) + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index f44d847099..d810f69e4a 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -8,6 +8,8 @@ use serde::Deserialize; use serde::Deserializer; use serde::Serialize; use serde::ser::Serializer; +use serde_json::Value; +use serde_json::json; use ts_rs::TS; use crate::user_input::UserInput; @@ -120,7 +122,7 @@ pub enum ResponseItem { // "action": {"type":"search","query":"weather: San Francisco, CA"} // } WebSearchCall { - #[serde(default, skip_serializing)] + #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(skip)] id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -211,13 +213,37 @@ pub struct LocalShellExecAction { } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, TS)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum WebSearchAction { - Search { - query: String, - }, - #[serde(other)] - Other, +#[serde(transparent)] +#[ts(type = "Record")] +pub struct WebSearchAction(Value); + +impl WebSearchAction { + pub fn search(query: impl Into) -> Self { + Self(json!({ + "type": "search", + "query": query.into(), + })) + } + + fn as_object(&self) -> Option<&serde_json::Map> { + self.0.as_object() + } + + pub fn action_type(&self) -> Option<&str> { + self.as_object() + .and_then(|map| map.get("type")) + .and_then(Value::as_str) + } + + pub fn is_search(&self) -> bool { + matches!(self.action_type(), Some("search")) + } + + pub fn query(&self) -> Option<&str> { + self.as_object() + .and_then(|map| map.get("query")) + .and_then(Value::as_str) + } } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema, TS)] @@ -703,4 +729,17 @@ mod tests { Ok(()) } + + #[test] + fn web_search_call_serializes_id() -> Result<()> { + let item = ResponseItem::WebSearchCall { + id: Some("search-1".into()), + status: Some("completed".into()), + action: WebSearchAction::search("weather"), + }; + + let json = serde_json::to_value(&item)?; + assert_eq!(json.get("id").and_then(|v| v.as_str()), Some("search-1")); + Ok(()) + } }