diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index bb43231730..a8d6cc2a9e 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -74,15 +74,24 @@ impl ToolOutput { } } - pub fn into_response(self, call_id: &str) -> ResponseInputItem { + pub fn into_response(self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { match self { - ToolOutput::Function { content, success } => ResponseInputItem::FunctionCallOutput { - call_id: call_id.to_string(), - output: FunctionCallOutputPayload { - content, - success: Some(success), - }, - }, + ToolOutput::Function { content, success } => { + if matches!(payload, ToolPayload::Custom { .. }) { + ResponseInputItem::CustomToolCallOutput { + call_id: call_id.to_string(), + output: content, + } + } else { + ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload { + content, + success: Some(success), + }, + } + } + } ToolOutput::Mcp { result } => ResponseInputItem::McpToolCallOutput { call_id: call_id.to_string(), result, @@ -91,6 +100,53 @@ impl ToolOutput { } } +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn custom_tool_calls_should_roundtrip_as_custom_outputs() { + let payload = ToolPayload::Custom { + input: "patch".to_string(), + }; + let response = ToolOutput::Function { + content: "patched".to_string(), + success: true, + } + .into_response("call-42", &payload); + + match response { + ResponseInputItem::CustomToolCallOutput { call_id, output } => { + assert_eq!(call_id, "call-42"); + assert_eq!(output, "patched"); + } + other => panic!("expected CustomToolCallOutput, got {other:?}"), + } + } + + #[test] + fn function_payloads_remain_function_outputs() { + let payload = ToolPayload::Function { + arguments: "{}".to_string(), + }; + let response = ToolOutput::Function { + content: "ok".to_string(), + success: true, + } + .into_response("fn-1", &payload); + + match response { + ResponseInputItem::FunctionCallOutput { call_id, output } => { + assert_eq!(call_id, "fn-1"); + assert_eq!(output.content, "ok"); + assert_eq!(output.success, Some(true)); + } + other => panic!("expected FunctionCallOutput, got {other:?}"), + } + } +} + #[derive(Clone, Debug)] pub(crate) struct ExecCommandContext { pub(crate) sub_id: String, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 24b62445fa..dc8e21d23b 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -64,7 +64,8 @@ impl ToolRegistry { let tool_name = invocation.tool_name.clone(); let call_id_owned = invocation.call_id.clone(); let otel = invocation.turn.client.get_otel_event_manager(); - let log_payload = invocation.payload.log_payload().into_owned(); + let payload_for_response = invocation.payload.clone(); + let log_payload = payload_for_response.log_payload().into_owned(); let handler = match self.handler(tool_name.as_ref()) { Some(handler) => handler, @@ -127,7 +128,7 @@ impl ToolRegistry { let output = guard.take().ok_or_else(|| { FunctionCallError::RespondToModel("tool produced no output".to_string()) })?; - Ok(output.into_response(&call_id_owned)) + Ok(output.into_response(&call_id_owned, &payload_for_response)) } Err(err) => Err(err), }