diff --git a/src/providers/codex/continuation.rs b/src/providers/codex/continuation.rs index 0ea93181..2cb2d8d7 100644 --- a/src/providers/codex/continuation.rs +++ b/src/providers/codex/continuation.rs @@ -537,15 +537,52 @@ fn input_suffix_after_prefix( return None; } for i in 0..prefix.len() { - let a = serde_json::to_value(&input[i]).unwrap_or_default(); - let b = serde_json::to_value(&prefix[i]).unwrap_or_default(); - if a != b { + if !input_items_equivalent(&input[i], &prefix[i]) { return None; } } Some(input[prefix.len()..].to_vec()) } +/// Compares two transcript items for continuation prefix purposes. +/// +/// `function_call.arguments` is a JSON string that is asymmetric across the two +/// sides: the recorded transcript keeps the raw text the model streamed, while +/// the next request re-serializes the client's parsed tool input (sorted keys, +/// no whitespace). Compare those two as JSON values so a purely textual +/// difference does not break an otherwise append-only prefix. +fn input_items_equivalent(a: &ResponsesInputItem, b: &ResponsesInputItem) -> bool { + match (a, b) { + ( + ResponsesInputItem::FunctionCall { + call_id: a_call_id, + name: a_name, + arguments: a_arguments, + }, + ResponsesInputItem::FunctionCall { + call_id: b_call_id, + name: b_name, + arguments: b_arguments, + }, + ) => { + a_call_id == b_call_id + && a_name == b_name + && (a_arguments == b_arguments + || matches!( + ( + serde_json::from_str::(a_arguments), + serde_json::from_str::(b_arguments), + ), + (Ok(a_value), Ok(b_value)) if a_value == b_value + )) + } + _ => { + serde_json::to_value(a).unwrap_or_default() + == serde_json::to_value(b).unwrap_or_default() + } + } +} + fn prompt_signature(body: &ResponsesRequest) -> String { let value = serde_json::to_value(body).unwrap_or_default(); let obj = match value.as_object() { @@ -618,6 +655,7 @@ fn evict_oldest(registry: &mut ContinuationRegistry) { #[cfg(test)] mod tests { + use super::super::translate::request::ResponsesFunctionCallOutput; use super::*; use serde_json::json; @@ -946,4 +984,182 @@ mod tests { ); assert!(!has_continuation_for_owner_for_tests(&owner)); } + + fn function_call(call_id: &str, name: &str, arguments: &str) -> ResponsesInputItem { + ResponsesInputItem::FunctionCall { + call_id: call_id.to_string(), + name: name.to_string(), + arguments: arguments.to_string(), + } + } + + fn function_call_output(call_id: &str, output: &str) -> ResponsesInputItem { + ResponsesInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: ResponsesFunctionCallOutput::Text(output.to_string()), + } + } + + fn record_with_output_items( + owner: &ConversationIdentity, + request: &ResponsesRequest, + response_id: &str, + output_items: &[ResponsesInputItem], + ) { + let reservation = continuation_candidate_for_owner(Some(owner), request, true); + record_continuation_for_owner( + &reservation, + request, + Some(response_id), + Some(1), + output_items, + ); + } + + #[test] + fn function_call_arguments_with_reordered_keys_keep_continuation() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + let recorded = function_call( + "call_1", + "Edit", + "{\"file_path\":\"/tmp/x\",\"old_string\":\"a\",\"new_string\":\"b\"}", + ); + record_with_output_items(&owner, &request, "resp_1", &[recorded]); + + let next = request_with_input( + vec![ + input("one"), + function_call( + "call_1", + "Edit", + "{\"file_path\":\"/tmp/x\",\"new_string\":\"b\",\"old_string\":\"a\"}", + ), + function_call_output("call_1", "done"), + ], + None, + ); + let reservation = continuation_candidate_for_owner(Some(&owner), &next, true); + assert_eq!(reservation.candidate().disabled_reason, None); + assert_eq!( + reservation.candidate().previous_response_id.as_deref(), + Some("resp_1") + ); + assert_eq!(reservation.candidate().input_delta_count, 1); + let delta = reservation.candidate().input_delta.clone().unwrap(); + assert_eq!(delta.len(), 1); + assert!(matches!( + delta[0], + ResponsesInputItem::FunctionCallOutput { .. } + )); + } + + #[test] + fn function_call_arguments_with_different_values_are_not_append_only() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + let recorded = function_call( + "call_1", + "Edit", + "{\"file_path\":\"/tmp/x\",\"old_string\":\"a\",\"new_string\":\"b\"}", + ); + record_with_output_items(&owner, &request, "resp_1", &[recorded]); + + let next = request_with_input( + vec![ + input("one"), + function_call( + "call_1", + "Edit", + "{\"file_path\":\"/tmp/x\",\"new_string\":\"c\",\"old_string\":\"a\"}", + ), + function_call_output("call_1", "done"), + ], + None, + ); + let reservation = continuation_candidate_for_owner(Some(&owner), &next, true); + assert_eq!( + reservation.candidate().disabled_reason.as_deref(), + Some("not_append_only") + ); + assert_eq!(reservation.candidate().previous_response_id, None); + } + + #[test] + fn unparseable_function_call_arguments_require_exact_text() { + let _registry_guard = lock_registry(); + let owner = main_owner("session-a"); + let request = request_with_input(vec![input("one")], None); + record_with_output_items( + &owner, + &request, + "resp_1", + &[function_call("call_1", "Edit", "{not json")], + ); + + let identical = request_with_input( + vec![ + input("one"), + function_call("call_1", "Edit", "{not json"), + function_call_output("call_1", "done"), + ], + None, + ); + let reservation = continuation_candidate_for_owner(Some(&owner), &identical, true); + assert_eq!(reservation.candidate().disabled_reason, None); + assert_eq!( + reservation.candidate().previous_response_id.as_deref(), + Some("resp_1") + ); + assert_eq!(reservation.candidate().input_delta_count, 1); + + record_with_output_items( + &owner, + &request, + "resp_2", + &[function_call("call_1", "Edit", "{not json")], + ); + let spaced = request_with_input( + vec![ + input("one"), + function_call("call_1", "Edit", "{not json "), + function_call_output("call_1", "done"), + ], + None, + ); + let reservation = continuation_candidate_for_owner(Some(&owner), &spaced, true); + assert_eq!( + reservation.candidate().disabled_reason.as_deref(), + Some("not_append_only") + ); + assert_eq!(reservation.candidate().previous_response_id, None); + } + + #[test] + fn input_items_equivalent_compares_arguments_as_json() { + assert!(input_items_equivalent( + &function_call("call_1", "Edit", "{\"a\": 1}"), + &function_call("call_1", "Edit", "{\"a\":1}"), + )); + assert!(!input_items_equivalent( + &function_call("call_1", "Edit", "{\"a\":1}"), + &function_call("call_1", "Edit", "{not json"), + )); + assert!(!input_items_equivalent( + &function_call("call_1", "Edit", "{\"a\":1}"), + &function_call("call_2", "Edit", "{\"a\":1}"), + )); + assert!(!input_items_equivalent( + &function_call("call_1", "Edit", "{\"a\":1}"), + &function_call("call_1", "Write", "{\"a\":1}"), + )); + assert!(input_items_equivalent(&input("one"), &input("one"))); + assert!(!input_items_equivalent(&input("one"), &input("two"))); + assert!(!input_items_equivalent( + &input("one"), + &function_call("call_1", "Edit", "{\"a\":1}"), + )); + } } diff --git a/tests/codex_agent_continuation.rs b/tests/codex_agent_continuation.rs index 49c301a4..0bf6b2dd 100644 --- a/tests/codex_agent_continuation.rs +++ b/tests/codex_agent_continuation.rs @@ -197,12 +197,24 @@ enum MockOutcome { close_after: bool, acknowledged: oneshot::Sender<()>, }, + FunctionCall { + response_id: String, + call: FunctionCallReply, + acknowledged: oneshot::Sender<()>, + }, RawEvent { event: Value, acknowledged: oneshot::Sender<()>, }, } +#[derive(Clone)] +struct FunctionCallReply { + call_id: String, + name: String, + arguments: String, +} + struct PendingRequest { captured: CapturedRequest, outcome: oneshot::Sender, @@ -231,6 +243,31 @@ impl PendingRequest { self.captured } + async fn respond_with_function_call( + self, + response_id: &str, + call: FunctionCallReply, + ) -> CapturedRequest { + let (acknowledged, acknowledgement) = oneshot::channel(); + self.outcome + .send(MockOutcome::FunctionCall { + response_id: response_id.to_string(), + call, + acknowledged, + }) + .unwrap_or_else(|_| { + panic!( + "upstream socket closed before responding to {}", + self.captured.marker() + ) + }); + tokio::time::timeout(REQUEST_TIMEOUT, acknowledgement) + .await + .expect("mock response acknowledgement timed out") + .expect("mock response acknowledgement sender dropped"); + self.captured + } + async fn respond_rate_limited(self) -> CapturedRequest { let (acknowledged, acknowledgement) = oneshot::channel(); self.outcome @@ -514,6 +551,20 @@ async fn handle_socket( } let _ = acknowledged.send(()); } + MockOutcome::FunctionCall { + response_id, + call, + acknowledged, + } => { + if emit_function_call(&mut websocket, &response_id, &call) + .await + .is_err() + { + let _ = acknowledged.send(()); + return; + } + let _ = acknowledged.send(()); + } MockOutcome::RawEvent { event, acknowledged, @@ -534,6 +585,7 @@ async fn handle_socket( } } +#[allow(clippy::result_large_err)] async fn emit_completion( websocket: &mut WebSocketStream, response_id: &str, @@ -569,6 +621,54 @@ async fn emit_completion( Ok(()) } +#[allow(clippy::result_large_err)] +async fn emit_function_call( + websocket: &mut WebSocketStream, + response_id: &str, + call: &FunctionCallReply, +) -> Result<(), tokio_tungstenite::tungstenite::Error> { + let item_id = format!("fc-{response_id}"); + let events = [ + json!({ + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "function_call", + "id": item_id, + "call_id": call.call_id, + "name": call.name + } + }), + json!({ + "type": "response.function_call_arguments.delta", + "output_index": 0, + "delta": call.arguments + }), + json!({ + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "function_call", + "id": item_id, + "call_id": call.call_id, + "name": call.name, + "arguments": call.arguments + } + }), + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "usage": {"input_tokens": 5, "output_tokens": 2} + } + }), + ]; + for event in events { + websocket.send(Message::Text(event.to_string())).await?; + } + Ok(()) +} + #[derive(Debug, Clone, Default)] struct IdentityHeaders { values: Vec<(&'static str, String)>, @@ -712,6 +812,23 @@ impl TestHarness { resolve_request(pending, request, response_id, reply, false).await } + /// Round trip for turns whose final input item carries no text marker, + /// such as a turn that ends in a tool result. + async fn round_trip_any( + &mut self, + body: Value, + identity: IdentityHeaders, + response_id: &str, + reply: &str, + ) -> CapturedRequest { + let request = self.start_request(body, identity.clone()); + let pending = self + .upstream + .next_any_request(identity.upstream_session.as_deref()) + .await; + resolve_request(pending, request, response_id, reply, false).await + } + async fn shutdown(mut self) { drop(self.client); if let Some(shutdown) = self.server_shutdown.take() { @@ -1899,3 +2016,97 @@ async fn cross_owner_completion_order_is_independent() { assert_eq!(harness.upstream.snapshot().len(), 6); harness.shutdown().await; } + +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn tool_turn_with_noncanonical_arguments_continues_with_delta() { + let _environment_lock = env_lock(); + let mut harness = TestHarness::start().await; + let case = unique("tool-turn"); + let session = tagged(&case, "session"); + let headers = IdentityHeaders::main(&session); + let first_prompt = tagged(&case, "prompt-1"); + let first_response = tagged(&case, "resp-1"); + let call_id = tagged(&case, "call-1"); + // Streamed in the model's own key order, which differs from the canonical + // serialization of the tool input the client replays on the next turn. + let arguments = r#"{"file_path":"/tmp/x","old_string":"a","new_string":"b"}"#; + + let tool_request = harness.start_request( + messages_body(false, vec![message("user", &first_prompt)]), + headers.clone(), + ); + let tool_pending = harness.pending(&first_prompt, &headers).await; + let first_capture = tool_pending + .respond_with_function_call( + &first_response, + FunctionCallReply { + call_id: call_id.clone(), + name: "Edit".to_string(), + arguments: arguments.to_string(), + }, + ) + .await; + let tool_response = tokio::time::timeout(REQUEST_TIMEOUT, tool_request) + .await + .expect("downstream tool request timed out") + .expect("downstream tool request failed"); + assert_eq!(tool_response.status, StatusCode::OK); + assert!( + tool_response.body.contains(&call_id), + "downstream response did not carry the tool call: {}", + tool_response.body + ); + + let second_capture = harness + .round_trip_any( + messages_body( + false, + vec![ + message("user", &first_prompt), + json!({ + "role": "assistant", + "content": [{ + "type": "tool_use", + "id": call_id, + "name": "Edit", + "input": {"file_path": "/tmp/x", "old_string": "a", "new_string": "b"} + }] + }), + json!({ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": call_id, + "content": "done" + }] + }), + ], + ), + headers, + &tagged(&case, "resp-2"), + &tagged(&case, "reply-2"), + ) + .await; + + assert_full_input(&first_capture, &[("user", &first_prompt)]); + assert_eq!( + second_capture.previous_response_id(), + Some(first_response.as_str()), + "tool turn must continue from the recorded response" + ); + assert_eq!(second_capture.socket_ordinal, first_capture.socket_ordinal); + let delta = second_capture.body["input"] + .as_array() + .expect("continuation delta input"); + assert_eq!( + delta.len(), + 1, + "tool turn must append only the tool result: {}", + second_capture.body + ); + assert_eq!(delta[0]["type"], "function_call_output"); + assert_eq!(delta[0]["call_id"], call_id.as_str()); + assert_eq!(harness.upstream.snapshot().len(), 2); + harness.shutdown().await; +}