diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index bf8ec1e54..d0fec5067 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -18,7 +18,7 @@ use switchyard_protocol::{ }; use switchyard_translation::{ WireFormat, decode_aggregated_response, decode_request, decode_stream, - encode_aggregated_response, encode_request, encode_stream, + encode_aggregated_response_with_extensions, encode_request, encode_stream_with_extensions, }; use tracing::Instrument; @@ -487,6 +487,7 @@ impl TranslatingLlmClient { ) -> Result { let llm_request = decode_request(wire_format, &raw_http_request) .map_err(|error| LlmClientError::RequestTranslation(error.to_string()))?; + let request_extensions = llm_request.extensions.clone(); // The model that serves the call — the rewrite target when the caller pinned // one, else the request's own model. Mirrors `call_rewrite_model`'s own // resolution so the response names whoever answered. @@ -512,13 +513,22 @@ impl TranslatingLlmClient { match response.llm_response { LlmResponse::Agg(agg) => { - let body = - encode_aggregated_response(&agg, wire_format, served_model.as_deref()) - .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; + let body = encode_aggregated_response_with_extensions( + &agg, + wire_format, + served_model.as_deref(), + &request_extensions, + ) + .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; Ok(RawResponse::Buffered(body)) } LlmResponse::Stream(chunks) => { - let events = encode_stream(chunks, wire_format, served_model)?; + let events = encode_stream_with_extensions( + chunks, + wire_format, + served_model, + &request_extensions, + )?; Ok(RawResponse::Stream(events)) } } @@ -1956,6 +1966,80 @@ mod tests { Ok(()) } + // A Codex namespace is folded into the upstream tool name, then split back + // into name and namespace on the Responses call that returns to Codex. + #[tokio::test] + async fn call_rewrite_model_raw_restores_codex_mcp_namespace() + -> std::result::Result<(), Box> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(wiremock::matchers::body_partial_json(json!({ + "tools": [{ + "type": "function", + "function": {"name": "mcp__open_websearch__search"} + }] + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-1", + "model": "gpt", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": { + "name": "mcp__open_websearch__search", + "arguments": "{\"q\":\"rust\"}" + } + }] + }, + "finish_reason": "tool_calls" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; + let raw = json!({ + "model": "client-facing", + "input": "Search for Rust.", + "tools": [{ + "type": "namespace", + "name": "mcp__open_websearch", + "tools": [{ + "type": "function", + "name": "search", + "description": "Search the web", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}} + }] + }] + }); + + let RawResponse::Buffered(body) = client + .call_rewrite_model_raw( + raw, + None, + Some(&ModelId::from("gpt")), + WireFormat::OpenAiResponses, + ) + .await? + else { + panic!("expected a buffered response"); + }; + + assert_eq!(body["output"][0]["type"], "function_call"); + assert_eq!(body["output"][0]["name"], "search"); + assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch"); + // Arguments are parsed and re-serialized, so the spacing is normalized. + assert_eq!(body["output"][0]["arguments"], "{\"q\": \"rust\"}"); + Ok(()) + } + // Raw path, streaming: an inbound `stream: true` request yields an unframed stream // of OpenAI Chat chunk objects whose deltas reassemble the completion. #[tokio::test] diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 483c1f6a6..004aa3a8a 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -692,6 +692,9 @@ async fn decision( Ok(aggregate) => aggregate, Err(error) => return client_error(&error), }; + // The request moved into the decision run, so its namespace mapping + // is gone by here. A Codex tool call in this preview keeps its + // qualified name. match encode_aggregated_response( &aggregate, input_format, @@ -924,6 +927,8 @@ async fn handle_llm_request( Ok(resolved) => resolved, Err(response) => return response, }; + // Only the Codex namespace mapping is needed downstream, not the whole request. + let request_extensions = request.llm_request.extensions.clone(); let algorithm = Arc::clone(&route.algorithm); let client_router = route.target_clients.clone(); let observer = stats_observer( @@ -956,10 +961,11 @@ async fn handle_llm_request( }; let response_model = served_model.as_ref().map(ToString::to_string); - let mut response = match into_http_response(response, wire_format, response_model) { - Ok(response) => response, - Err(error) => return server_error(error.to_string()), - }; + let mut response = + match into_http_response(response, wire_format, response_model, request_extensions) { + Ok(response) => response, + Err(error) => return server_error(error.to_string()), + }; if let Some(served_model) = served_model.as_ref() { attach_routing_headers(&mut response, served_model.as_str()); } diff --git a/crates/switchyard-server/src/response.rs b/crates/switchyard-server/src/response.rs index 2149f7036..950ec6399 100644 --- a/crates/switchyard-server/src/response.rs +++ b/crates/switchyard-server/src/response.rs @@ -7,8 +7,10 @@ use std::error::Error; use axum::Json; use axum::response::{IntoResponse, Response as HttpResponse}; -use switchyard_protocol::{LlmResponse, Response as AlgorithmResponse}; -use switchyard_translation::{WireFormat, encode_aggregated_response, encode_stream}; +use switchyard_protocol::{LlmResponse, ProviderExtensions, Response as AlgorithmResponse}; +use switchyard_translation::{ + WireFormat, encode_aggregated_response_with_extensions, encode_stream_with_extensions, +}; use crate::sse::frame_stream; @@ -21,18 +23,26 @@ pub(crate) fn into_http_response( response: AlgorithmResponse, target_format: WireFormat, served_model: Option, + request_extensions: ProviderExtensions, ) -> Result { match response.llm_response { - LlmResponse::Agg(response) => Ok(Json(encode_aggregated_response( - &response, - target_format, - served_model.as_deref(), - )?) - .into_response()), - LlmResponse::Stream(stream) => Ok(frame_stream( - encode_stream(stream, target_format, served_model)?, - target_format, - ) - .into_response()), + LlmResponse::Agg(response) => { + let body = encode_aggregated_response_with_extensions( + &response, + target_format, + served_model.as_deref(), + &request_extensions, + )?; + Ok(Json(body).into_response()) + } + LlmResponse::Stream(stream) => { + let events = encode_stream_with_extensions( + stream, + target_format, + served_model, + &request_extensions, + )?; + Ok(frame_stream(events, target_format).into_response()) + } } } diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 21b1e46da..e68890411 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -149,6 +149,27 @@ async fn upstream_chat( .into_response(); } if body["stream"].as_bool() == Some(true) { + // Streamed tool call, for the namespace-on-every-event assertions. The + // model calls a tool by the name it was given, so echo that name back. + if body["messages"][0]["content"] == "mcp-tool-call" { + let called = body["tool_choice"]["function"]["name"] + .as_str() + .or_else(|| body["tools"][0]["function"]["name"].as_str()) + .unwrap_or("search") + .to_string(); + let events = [ + json!({"id": "chatcmpl-mcp", "model": model, "choices": [{"index": 0, "delta": {"role": "assistant", "tool_calls": [{"index": 0, "id": "call_1", "type": "function", "function": {"name": called, "arguments": ""}}]}}]}).to_string(), + json!({"id": "chatcmpl-mcp", "model": model, "choices": [{"index": 0, "delta": {"tool_calls": [{"index": 0, "function": {"arguments": "{\"q\":\"rust\"}"}}]}}]}).to_string(), + json!({"id": "chatcmpl-mcp", "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}], "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}}).to_string(), + "[DONE]".to_string(), + ]; + let stream = futures_util::stream::iter( + events + .into_iter() + .map(|data| Ok::(Event::default().data(data))), + ); + return Sse::new(stream).into_response(); + } if body["messages"][0]["content"] == "stream-error" { let events = [ json!({"id": "chatcmpl-stream-error", "model": model, "choices": [{"index": 0, "delta": {"role": "assistant"}}]}).to_string(), @@ -210,6 +231,35 @@ async fn upstream_chat( .into_response(); } + // Buffered tool call, the non-streaming counterpart of the branch above. + if body["messages"][0]["content"] == "mcp-tool-call" { + let called = body["tool_choice"]["function"]["name"] + .as_str() + .or_else(|| body["tools"][0]["function"]["name"].as_str()) + .unwrap_or("search") + .to_string(); + return Json(json!({ + "id": "chatcmpl-mcp", + "object": "chat.completion", + "model": model, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": called, "arguments": "{\"q\":\"rust\"}"} + }] + }, + "finish_reason": "tool_calls" + }], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7} + })) + .into_response(); + } + let custom_target_schema = body .pointer("/response_format/json_schema/schema/properties/decision/properties/target") .is_some(); @@ -3220,3 +3270,101 @@ async fn advisor_route_redo_fail_open_and_stats_projection() -> TestResult { assert_eq!(gate_count(&stats, &["discarded", "turns"]), 0); Ok(()) } + +// Returns every `data:` frame of an SSE body as JSON, skipping `[DONE]`. +fn sse_events(body: &str) -> Vec { + body.lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|data| *data != "[DONE]") + .filter_map(|data| serde_json::from_str(data).ok()) + .collect() +} + +// The end-to-end contract for Codex tool namespaces, in one request. +// +// Two MCP servers expose the same tool name, so the flat upstream can only tell +// them apart by the namespace folded into each name. Everything naming a tool — +// the definitions, the recorded call in history, and the forced tool choice — +// has to use that same spelling, and the response has to split it back into the +// name and namespace Codex dispatches on. +#[tokio::test] +async fn responses_round_trips_codex_tool_namespaces() -> TestResult { + const MODEL: &str = "model/mcp-namespaces"; + let (upstream, app) = test_app(&[(ROUTE_MODEL, &[MODEL])]).await?; + + let response = send( + &app, + "POST", + "/v1/responses", + Some(json!({ + "model": ROUTE_MODEL, + "stream": true, + "input": [ + {"type": "message", "role": "user", + "content": [{"type": "input_text", "text": "mcp-tool-call"}]}, + {"type": "function_call", "call_id": "call_prior", "name": "search", + "namespace": "mcp__b", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_prior", "output": "earlier"} + ], + "tool_choice": {"type": "function", "name": "search", "namespace": "mcp__b"}, + "tools": [ + {"type": "namespace", "name": "mcp__a", "tools": [{ + "type": "function", "name": "search", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}} + }]}, + {"type": "namespace", "name": "mcp__b", "tools": [{ + "type": "function", "name": "search", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}} + }]} + ] + })), + ) + .await?; + assert_eq!(response.status, StatusCode::OK); + + // The upstream sees two distinct tools, and every reference to the forced + // one uses the qualified spelling. + let calls = upstream.calls.lock().await; + let sent = &calls[0]; + let offered = sent["tools"] + .as_array() + .map(|tools| { + tools + .iter() + .filter_map(|tool| tool["function"]["name"].as_str()) + .collect::>() + }) + .unwrap_or_default(); + assert_eq!(offered, vec!["mcp__a__search", "mcp__b__search"]); + assert_eq!(sent["tool_choice"]["function"]["name"], "mcp__b__search"); + let recorded = sent["messages"] + .as_array() + .and_then(|messages| { + messages + .iter() + .find_map(|message| message["tool_calls"][0]["function"]["name"].as_str()) + }) + .ok_or("no recorded tool call reached the upstream")?; + assert_eq!(recorded, "mcp__b__search"); + drop(calls); + + // The upstream answers with the qualified name; Codex must receive the tool + // name and the namespace it dispatches on, on every event that names a call. + let events = sse_events(response.text()?); + for event_type in ["response.output_item.added", "response.output_item.done"] { + let item = events + .iter() + .find(|event| event["type"] == event_type) + .map(|event| event["item"].clone()) + .ok_or(format!("stream produced no {event_type}"))?; + assert_eq!(item["name"], "search", "{event_type}"); + assert_eq!(item["namespace"], "mcp__b", "{event_type}"); + } + let completed = events + .iter() + .find(|event| event["type"] == "response.completed") + .ok_or("stream produced no response.completed event")?; + assert_eq!(completed["response"]["output"][0]["name"], "search"); + assert_eq!(completed["response"]["output"][0]["namespace"], "mcp__b"); + Ok(()) +} diff --git a/crates/switchyard-translation/src/codecs/responses/buffered.rs b/crates/switchyard-translation/src/codecs/responses/buffered.rs index db5ccab95..32d062cb5 100644 --- a/crates/switchyard-translation/src/codecs/responses/buffered.rs +++ b/crates/switchyard-translation/src/codecs/responses/buffered.rs @@ -87,7 +87,8 @@ impl FormatCodec for OpenAiResponsesCodec { &mut diagnostics, policy, )?; - request.tools = decode_responses_tools(body.get("tools")); + let mut tool_namespaces = Map::new(); + request.tools = decode_responses_tools(body.get("tools"), &mut tool_namespaces); request.tool_choice = body .get("tool_choice") .and_then(decode_responses_tool_choice); @@ -107,6 +108,7 @@ impl FormatCodec for OpenAiResponsesCodec { "stream", ], ); + crate::codex_namespaces::attach_tool_namespaces(&mut request.extensions, tool_namespaces); Ok(DecodedRequest { request, diagnostics, @@ -147,13 +149,27 @@ impl FormatCodec for OpenAiResponsesCodec { } body.insert( "input".to_string(), - encode_responses_input(&request.messages, &mut diagnostics, _policy)?, + encode_responses_input( + &request.messages, + &mut diagnostics, + _policy, + crate::codex_namespaces::tool_namespaces(&request.extensions), + )?, ); if !request.tools.is_empty() { - body.insert("tools".to_string(), encode_responses_tools(&request.tools)); + body.insert( + "tools".to_string(), + encode_responses_tools( + &request.tools, + crate::codex_namespaces::tool_namespaces(&request.extensions), + ), + ); } if let Some(choice) = &request.tool_choice - && let Some(choice) = encode_responses_tool_choice(choice) + && let Some(choice) = encode_responses_tool_choice( + choice, + crate::codex_namespaces::tool_namespaces(&request.extensions), + ) { body.insert("tool_choice".to_string(), choice); } @@ -414,13 +430,27 @@ fn decode_responses_input( } DeterministicIdPolicy::Preserve => String::new(), }); + // History has to spell a tool the way its definition + // does, or the transcript teaches the model a name the + // upstream was never offered. + let name = item + .get("name") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + let name = match item + .get("namespace") + .and_then(Value::as_str) + .filter(|namespace| !namespace.is_empty()) + { + Some(namespace) => { + crate::codex_namespaces::qualified_tool_name(namespace, &name) + } + None => name, + }; pending_tool_calls.push(ToolCall { id, - name: item - .get("name") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(), + name, arguments: item.get("arguments").cloned().unwrap_or_else(|| json!({})), }); } @@ -735,7 +765,16 @@ fn request_role_from_responses(role: Option<&str>, path: &str) -> Result { } // Decodes Responses tool shapes, including Codex-style tool entries. -fn decode_responses_tools(value: Option<&Value>) -> Vec { +// Decodes Responses tool shapes, including Codex-style tool entries. +// +// Codex groups tools into non-standard ``namespace`` containers that +// OpenAI-compatible upstreams do not accept. Each child is exposed under a +// `__` name, and `namespaces` records the mapping so the +// response can split it back apart. +fn decode_responses_tools( + value: Option<&Value>, + namespaces: &mut Map, +) -> Vec { let Some(tools) = value.and_then(Value::as_array) else { return Vec::new(); }; @@ -744,7 +783,30 @@ fn decode_responses_tools(value: Option<&Value>) -> Vec { let Some(tool) = tool.as_object() else { continue; }; - if tool.get("type").and_then(Value::as_str) == Some("function") { + if tool.get("type").and_then(Value::as_str) == Some("namespace") { + // An unnamed container carries nothing to dispatch on, so its + // children stay bare rather than gaining an empty prefix. + let container = tool + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()); + for mut child in decode_responses_tools(tool.get("tools"), namespaces) { + // A nested container already qualified its own children, and the + // innermost name is the one that identifies the tool. + let already_qualified = namespaces.contains_key(&child.name); + if let Some(container) = container + && !already_qualified + { + let qualified = + crate::codex_namespaces::qualified_tool_name(container, &child.name); + crate::codex_namespaces::record_tool_namespace( + namespaces, &qualified, container, + ); + child.name = qualified; + } + out.push(child); + } + } else if tool.get("type").and_then(Value::as_str) == Some("function") { if let Some(function) = tool.get("function").and_then(Value::as_object) { if let Some(name) = function.get("name").and_then(Value::as_str) && !name.is_empty() @@ -836,12 +898,22 @@ fn decode_responses_tool_choice(value: &Value) -> Option { Value::String(text) if text == "required" => Some(ToolChoice::Required), Value::String(text) if text == "none" => Some(ToolChoice::None), Value::Object(object) if object.get("type").and_then(Value::as_str) == Some("function") => { - object - .get("name") - .and_then(Value::as_str) - .map(|name| ToolChoice::Tool { - name: name.to_string(), - }) + // A forced tool has to name the same thing its definition does, or + // the upstream rejects a choice naming a tool it was never offered. + object.get("name").and_then(Value::as_str).map(|name| { + match object + .get("namespace") + .and_then(Value::as_str) + .filter(|namespace| !namespace.is_empty()) + { + Some(namespace) => ToolChoice::Tool { + name: crate::codex_namespaces::qualified_tool_name(namespace, name), + }, + None => ToolChoice::Tool { + name: name.to_string(), + }, + } + }) } Value::Object(_) => None, _ => Some(ToolChoice::Raw(value.clone())), @@ -923,6 +995,7 @@ fn encode_responses_input( messages: &[Message], diagnostics: &mut Vec, policy: &TranslationPolicy, + namespaces: Option<&Map>, ) -> Result { if messages.len() == 1 && matches!(messages[0].role, Role::User) @@ -958,13 +1031,17 @@ fn encode_responses_input( ContentBlock::ToolCall(_) | ContentBlock::ToolResult(_) ) }) { - encoded.extend(content.iter().filter_map(encode_responses_special_input)); + encoded.extend( + content + .iter() + .filter_map(|block| encode_responses_special_input(block, namespaces)), + ); continue; } let mut visible_content = Vec::new(); let mut emitted_special = false; for block in &content { - if let Some(item) = encode_responses_special_input(block) { + if let Some(item) = encode_responses_special_input(block, namespaces) { encoded.push(item); emitted_special = true; } else { @@ -984,7 +1061,10 @@ fn encode_responses_input( } // Encodes IR blocks that Responses represents as top-level input items. -fn encode_responses_special_input(block: &ContentBlock) -> Option { +fn encode_responses_special_input( + block: &ContentBlock, + namespaces: Option<&Map>, +) -> Option { match block { ContentBlock::Reasoning { text, @@ -995,12 +1075,27 @@ fn encode_responses_special_input(block: &ContentBlock) -> Option { "content": [{"type": "reasoning_text", "text": text}], "summary": [], })), - ContentBlock::ToolCall(call) => Some(json!({ - "type": "function_call", - "call_id": call.id, - "name": call.name, - "arguments": json_string(&call.arguments), - })), + ContentBlock::ToolCall(call) => { + // A Responses client dispatches on name plus namespace, so undo the + // qualification this request applied for a flat upstream. + let (name, namespace) = namespaces + .and_then(|namespaces| { + crate::codex_namespaces::split_qualified_name(namespaces, &call.name) + }) + .map_or((call.name.clone(), None), |(name, namespace)| { + (name, Some(namespace)) + }); + let mut item = json!({ + "type": "function_call", + "call_id": call.id, + "name": name, + "arguments": json_string(&call.arguments), + }); + if let Some(namespace) = namespace { + item["namespace"] = Value::String(namespace); + } + Some(item) + } ContentBlock::ToolResult(result) => Some(json!({ "type": "function_call_output", "call_id": result.tool_call_id, @@ -1091,33 +1186,73 @@ fn encode_responses_content( } // Encodes normalized tool definitions into Responses tool JSON. -fn encode_responses_tools(tools: &[ToolDefinition]) -> Value { - Value::Array( - tools - .iter() - .map(|tool| { - let mut item = json!({ - "type": "function", - "name": tool.name, - "description": tool.description.clone().unwrap_or_default(), - "parameters": tool.parameters, - }); - if let Some(strict) = tool.strict { - item["strict"] = Value::Bool(strict); +// Encodes normalized tools as Responses entries. +// +// A Responses client understands Codex containers, so a tool whose name this +// request qualified is regrouped under the container it came from. +fn encode_responses_tools( + tools: &[ToolDefinition], + namespaces: Option<&Map>, +) -> Value { + let mut out: Vec = Vec::new(); + let mut containers: Vec<(String, Vec)> = Vec::new(); + for tool in tools { + let mut item = json!({ + "type": "function", + "name": tool.name, + "description": tool.description.clone().unwrap_or_default(), + "parameters": tool.parameters, + }); + if let Some(strict) = tool.strict { + item["strict"] = Value::Bool(strict); + } + let split = namespaces.and_then(|namespaces| { + crate::codex_namespaces::split_qualified_name(namespaces, &tool.name) + }); + match split { + None => out.push(item), + Some((name, namespace)) => { + item["name"] = Value::String(name); + match containers + .iter_mut() + .find(|(existing, _)| *existing == namespace) + { + Some((_, children)) => children.push(item), + None => containers.push((namespace, vec![item])), } - item - }) - .collect(), - ) + } + } + } + for (namespace, children) in containers { + out.push(json!({ + "type": "namespace", + "name": namespace, + "tools": children, + })); + } + Value::Array(out) } // Encodes normalized tool choice into Responses JSON. -fn encode_responses_tool_choice(choice: &ToolChoice) -> Option { +fn encode_responses_tool_choice( + choice: &ToolChoice, + namespaces: Option<&Map>, +) -> Option { match choice { ToolChoice::Auto => Some(Value::String("auto".to_string())), ToolChoice::Required => Some(Value::String("required".to_string())), ToolChoice::None => Some(Value::String("none".to_string())), - ToolChoice::Tool { name } => Some(json!({"type": "function", "name": name})), + ToolChoice::Tool { name } => { + let split = namespaces.and_then(|namespaces| { + crate::codex_namespaces::split_qualified_name(namespaces, name) + }); + Some(match split { + Some((tool, namespace)) => { + json!({"type": "function", "name": tool, "namespace": namespace}) + } + None => json!({"type": "function", "name": name}), + }) + } ToolChoice::Raw(value) => Some(value.clone()), } } diff --git a/crates/switchyard-translation/src/codex_namespaces.rs b/crates/switchyard-translation/src/codex_namespaces.rs new file mode 100644 index 000000000..faaff5f77 --- /dev/null +++ b/crates/switchyard-translation/src/codex_namespaces.rs @@ -0,0 +1,247 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Codex tool namespace preservation across a Responses/Chat translation. +//! +//! Codex groups tools into non-standard ``namespace`` containers — MCP servers, +//! usually behind an `mcp__` prefix, plus builtin groups such as +//! `multi_agent_v1` — and dispatches on the pair of tool name and namespace. It +//! therefore expects the namespace back on the call it receives: +//! `{"type": "function_call", "name": "search", "namespace": "mcp__docs"}`. +//! +//! OpenAI-compatible upstreams accept only flat `function` tools, so the request +//! codec flattens the containers and names each child `__`. +//! Two tools differing only by namespace therefore stay distinct upstream, and +//! the conversation history uses the same spelling so the transcript never +//! teaches the model a name the upstream was not offered. +//! +//! The qualified-name to namespace mapping rides in the request's +//! [`ProviderExtensions`], so no provider-neutral type grows a Codex-specific +//! field. No codec copies unknown extension keys into an outbound body, so the +//! mapping reaches neither the upstream nor the client. +//! +//! No container is filtered by its name, because Codex resolves a call with no +//! namespace against its default group: dropping a builtin group's namespace +//! would look the call up in the wrong place. + +use std::collections::{HashMap, HashSet}; + +use serde_json::{Map, Value}; +use switchyard_protocol::ProviderExtensions; + +/// Separator between a namespace and a tool name in a qualified wire name. +pub const NAMESPACE_SEPARATOR: &str = "__"; + +/// Request extension key holding the qualified-name to namespace mapping. +/// +/// Prefixed so it cannot collide with a real provider field, and so a codec that +/// allowlists provider fields never forwards it. +pub const TOOL_NAMESPACES_KEY: &str = "switchyard_codex_tool_namespaces"; + +/// Joins a namespace and a tool name into the name used on the wire. +pub fn qualified_tool_name(namespace: &str, tool: &str) -> String { + format!("{namespace}{NAMESPACE_SEPARATOR}{tool}") +} + +/// Records that `qualified` was produced by flattening a tool out of `namespace`. +pub fn record_tool_namespace( + namespaces: &mut Map, + qualified: &str, + namespace: &str, +) { + namespaces.insert(qualified.to_string(), Value::String(namespace.to_string())); +} + +/// Stores a collected mapping on a request's extensions, when it has entries. +pub fn attach_tool_namespaces(extensions: &mut ProviderExtensions, namespaces: Map) { + if !namespaces.is_empty() { + extensions + .fields + .insert(TOOL_NAMESPACES_KEY.to_string(), Value::Object(namespaces)); + } +} + +/// Reads the mapping back off a request's extensions. +pub fn tool_namespaces(extensions: &ProviderExtensions) -> Option<&Map> { + extensions + .fields + .get(TOOL_NAMESPACES_KEY) + .and_then(Value::as_object) +} + +/// Splits a qualified wire name back into its tool name and namespace. +/// +/// Returns `None` for a name the request never qualified, so an unrecognized +/// call is left alone rather than attributed to the wrong namespace. The tool +/// name may itself contain the separator, so the namespace is matched as a +/// prefix rather than by splitting on it. +pub fn split_qualified_name( + namespaces: &Map, + qualified: &str, +) -> Option<(String, String)> { + let namespace = namespaces.get(qualified).and_then(Value::as_str)?; + let tool = qualified + .strip_prefix(namespace)? + .strip_prefix(NAMESPACE_SEPARATOR)?; + Some((tool.to_string(), namespace.to_string())) +} + +/// Reverse map from an upstream tool name to its Codex tool name and namespace. +/// +/// The exact qualified name is always registered. A model often returns a near +/// miss instead, so two fallback spellings are registered when neither can be +/// confused with another tool: +/// +/// * the qualified name without the `mcp__` prefix +/// * the bare tool name, only when exactly one namespace claims it +pub fn qualified_tool_origins( + extensions: &ProviderExtensions, +) -> HashMap { + let Some(namespaces) = tool_namespaces(extensions) else { + return HashMap::new(); + }; + let split: Vec<(String, String, String)> = namespaces + .keys() + .filter_map(|qualified| { + let (tool, namespace) = split_qualified_name(namespaces, qualified)?; + Some((qualified.clone(), tool, namespace)) + }) + .collect(); + + // A bare name claimed by more than one namespace cannot be guessed, and a + // fallback must never shadow a name that is itself a real qualified tool. + let mut bare_claims: HashMap<&str, usize> = HashMap::new(); + for (_, tool, _) in &split { + *bare_claims.entry(tool.as_str()).or_default() += 1; + } + let qualified_names: HashSet<&str> = split + .iter() + .map(|(qualified, _, _)| qualified.as_str()) + .collect(); + + let mut origins = HashMap::new(); + for (qualified, tool, namespace) in &split { + let origin = (tool.clone(), namespace.clone()); + if let Some(stripped) = namespace.strip_prefix("mcp__") { + let spelling = qualified_tool_name(stripped, tool); + if !qualified_names.contains(spelling.as_str()) { + origins.entry(spelling).or_insert_with(|| origin.clone()); + } + } + if bare_claims.get(tool.as_str()) == Some(&1) && !qualified_names.contains(tool.as_str()) { + origins + .entry(tool.clone()) + .or_insert_with(|| origin.clone()); + } + // The exact spelling always wins over a fallback. + origins.insert(qualified.clone(), origin); + } + origins +} + +/// Rewrite `function_call` names back to the Codex tool name plus namespace. +/// +/// Walks the whole value, covering a buffered body and each streaming event, +/// where the item is nested under `item` (`response.output_item.added` / +/// `.done`) or `response.output` (`response.completed`). An existing +/// `namespace` is never overwritten. +pub fn restore_qualified_tool_names(body: &mut Value, origins: &HashMap) { + if origins.is_empty() { + return; + } + match body { + Value::Array(values) => { + for value in values { + restore_qualified_tool_names(value, origins); + } + } + Value::Object(object) => { + if object.get("type").and_then(Value::as_str) == Some("function_call") + && let Some(name) = object.get("name").and_then(Value::as_str) + && let Some((tool, namespace)) = origins.get(name) + { + object.insert("name".to_string(), Value::String(tool.clone())); + object + .entry("namespace".to_string()) + .or_insert_with(|| Value::String(namespace.clone())); + } + for value in object.values_mut() { + restore_qualified_tool_names(value, origins); + } + } + _ => {} + } +} + +#[cfg(test)] +mod tests { + use serde_json::{Map, json}; + use switchyard_protocol::ProviderExtensions; + + use super::{ + attach_tool_namespaces, qualified_tool_name, qualified_tool_origins, record_tool_namespace, + restore_qualified_tool_names, split_qualified_name, tool_namespaces, + }; + + fn extensions(pairs: &[(&str, &str)]) -> ProviderExtensions { + let mut namespaces = Map::new(); + for (namespace, tool) in pairs { + record_tool_namespace( + &mut namespaces, + &qualified_tool_name(namespace, tool), + namespace, + ); + } + let mut extensions = ProviderExtensions::default(); + attach_tool_namespaces(&mut extensions, namespaces); + extensions + } + + // A tool name may itself contain the separator, so the namespace is matched + // as a prefix rather than by splitting on the first `__`. + #[test] + fn splits_a_qualified_name_back_apart() { + let simple = extensions(&[("mcp__docs", "search")]); + let simple = tool_namespaces(&simple).expect("mapping present"); + assert_eq!( + split_qualified_name(simple, "mcp__docs__search"), + Some(("search".to_string(), "mcp__docs".to_string())) + ); + assert_eq!(split_qualified_name(simple, "unknown_tool"), None); + + let nested = extensions(&[("mcp__docs", "fetch__raw")]); + let nested = tool_namespaces(&nested).expect("mapping present"); + assert_eq!( + split_qualified_name(nested, "mcp__docs__fetch__raw"), + Some(("fetch__raw".to_string(), "mcp__docs".to_string())) + ); + } + + // A bare name claimed by two namespaces must not be guessed: a wrong guess + // dispatches the call to the wrong server. + #[test] + fn leaves_an_ambiguous_bare_name_alone() { + let origins = + qualified_tool_origins(&extensions(&[("mcp__a", "search"), ("mcp__b", "search")])); + let mut response = json!({"output": [{"type": "function_call", "name": "search"}]}); + let before = response.clone(); + + restore_qualified_tool_names(&mut response, &origins); + + assert_eq!(response, before); + } + + // Codex namespaces builtin groups too, so nothing may key on `mcp__`. + #[test] + fn resolves_namespaces_that_are_not_mcp_servers() { + let origins = qualified_tool_origins(&extensions(&[("multi_agent_v1", "spawn_agent")])); + let mut response = json!({ + "output": [{"type": "function_call", "name": "multi_agent_v1__spawn_agent"}] + }); + + restore_qualified_tool_names(&mut response, &origins); + + assert_eq!(response["output"][0]["name"], "spawn_agent"); + assert_eq!(response["output"][0]["namespace"], "multi_agent_v1"); + } +} diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs index f1fc7c321..106701b22 100644 --- a/crates/switchyard-translation/src/helpers.rs +++ b/crates/switchyard-translation/src/helpers.rs @@ -55,6 +55,24 @@ pub fn encode_aggregated_response( agg: &AggLlmResponse, wire_format: WireFormat, served_model: Option<&str>, +) -> Result { + encode_aggregated_response_with_extensions( + agg, + wire_format, + served_model, + &switchyard_protocol::ProviderExtensions::default(), + ) +} + +/// Encodes an aggregated response, honouring request extensions. +/// +/// Identical to [`encode_aggregated_response`] except that Codex tool +/// namespaces recorded on the request are restored on the encoded body. +pub fn encode_aggregated_response_with_extensions( + agg: &AggLlmResponse, + wire_format: WireFormat, + served_model: Option<&str>, + request_extensions: &switchyard_protocol::ProviderExtensions, ) -> Result { let mut body = DEFAULT_TRANSLATION_ENGINE .encode_response(wire_format, agg, &DEFAULT_TRANSLATION_POLICY)? @@ -62,6 +80,10 @@ pub fn encode_aggregated_response( if let (Some(model), Value::Object(object)) = (served_model, &mut body) { object.insert("model".to_string(), Value::String(model.to_string())); } + crate::codex_namespaces::restore_qualified_tool_names( + &mut body, + &crate::codex_namespaces::qualified_tool_origins(request_extensions), + ); Ok(body) } @@ -86,6 +108,25 @@ pub fn encode_stream( target: WireFormat, served_model: Option, ) -> std::result::Result { + encode_stream_with_extensions( + chunks, + target, + served_model, + &switchyard_protocol::ProviderExtensions::default(), + ) +} + +/// Encodes a response stream, honouring request extensions. +/// +/// Identical to [`encode_stream`] except that Codex tool namespaces recorded on +/// the request are restored on each encoded event. +pub fn encode_stream_with_extensions( + chunks: LlmResponseStream, + target: WireFormat, + served_model: Option, + request_extensions: &switchyard_protocol::ProviderExtensions, +) -> std::result::Result { + let origins = crate::codex_namespaces::qualified_tool_origins(request_extensions); let target_format: FormatId = target.into(); // The target is always a built-in wire format, so this lookup cannot fail; a // failure returns as an `Err` rather than a panic. @@ -115,6 +156,7 @@ pub fn encode_stream( target, served_model_for_events.as_deref(), ); + crate::codex_namespaces::restore_qualified_tool_names(&mut value, &origins); yield value; } if state.errored { @@ -127,6 +169,7 @@ pub fn encode_stream( target, served_model_for_events.as_deref(), ); + crate::codex_namespaces::restore_qualified_tool_names(&mut value, &origins); yield value; } }; diff --git a/crates/switchyard-translation/src/lib.rs b/crates/switchyard-translation/src/lib.rs index d07a3497d..bd5d5e060 100644 --- a/crates/switchyard-translation/src/lib.rs +++ b/crates/switchyard-translation/src/lib.rs @@ -8,6 +8,7 @@ //! servers, Python objects, or FFI bindings. pub mod codecs; +pub(crate) mod codex_namespaces; pub mod diagnostic; pub mod engine; pub mod error;