fix: preserve Codex MCP namespaces through translation

Signed-off-by: Brian Grinstead <briangrinstead@gmail.com>

Author: bgrins

crates/libsy-llm-client/src/client.rs

@@ -17,7 +17,8 @@ use switchyard_protocol::{
 };
 use switchyard_translation::{
     WireFormat, decode_aggregated_response, decode_request, decode_stream,
-    encode_aggregated_response, encode_request, encode_stream,
+    encode_aggregated_response, encode_request, encode_stream, responses_tool_namespaces,
+    restore_responses_tool_namespaces,
 };
 use tracing::Instrument;
 
@@ -463,6 +464,7 @@ impl TranslatingLlmClient {
         model: Option<&str>,
         wire_format: WireFormat,
     ) -> Result<RawResponse> {
+        let response_namespaces = responses_tool_namespaces(&raw_http_request, wire_format);
         let llm_request = decode_request(wire_format, &raw_http_request)
             .map_err(|error| LlmClientError::RequestTranslation(error.to_string()))?;
         // The model that serves the call — the rewrite target when the caller pinned
@@ -490,14 +492,25 @@ impl TranslatingLlmClient {
 
         match response.llm_response {
             LlmResponse::Agg(agg) => {
-                let body =
+                let mut body =
                     encode_aggregated_response(&agg, wire_format, served_model.as_deref())
                         .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
+                restore_responses_tool_namespaces(&mut body, &response_namespaces);
                 Ok(RawResponse::Buffered(body))
             }
             LlmResponse::Stream(chunks) => {
                 let events = encode_stream(chunks, wire_format, served_model)?;
-                Ok(RawResponse::Stream(events))
+                if response_namespaces.is_empty() {
+                    Ok(RawResponse::Stream(events))
+                } else {
+                    let events = events.map(move |event| {
+                        event.map(|mut value| {
+                            restore_responses_tool_namespaces(&mut value, &response_namespaces);
+                            value
+                        })
+                    });
+                    Ok(RawResponse::Stream(Box::pin(events)))
+                }
             }
         }
     }
@@ -1873,6 +1886,72 @@ mod tests {
         Ok(())
     }
 
+    // A Codex MCP namespace is flattened for the Chat upstream, then restored
+    // on the Responses function call that comes back to Codex.
+    #[tokio::test]
+    async fn call_rewrite_model_raw_restores_codex_mcp_namespace()
+    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
+        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": "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": "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("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]

crates/switchyard-server/src/lib.rs

@@ -42,7 +42,7 @@ use tokio::net::{TcpListener, TcpSocket};
 use tokio::task;
 use tracing::{Instrument, Level};
 
-use switchyard_translation::{WireFormat, decode_request};
+use switchyard_translation::{WireFormat, decode_request, responses_tool_namespaces};
 
 use crate::response::into_http_response;
 use crate::stats::{StatsAccumulator, StatsSnapshot, prefix_probe, tracking_enabled_from_env};
@@ -686,6 +686,7 @@ async fn handle_llm_request(
     wire_format: WireFormat,
     routing_log_context: Option<routing_log::RoutingLogContext>,
 ) -> Response {
+    let response_namespaces = responses_tool_namespaces(&body, wire_format);
     let cache_probe = state.track_cache_eligibility.then(|| prefix_probe(&body));
     let (route, request) = match resolve_route(&state, metadata, body, wire_format) {
         Ok(resolved) => resolved,
@@ -728,10 +729,11 @@ async fn handle_llm_request(
     };
 
     let served_model = decision.map(|decision| decision.selected_model_id().to_string());
-    let mut response = match into_http_response(response, wire_format, served_model) {
-        Ok(response) => response,
-        Err(error) => return server_error(error.to_string()),
-    };
+    let mut response =
+        match into_http_response(response, wire_format, served_model, response_namespaces) {
+            Ok(response) => response,
+            Err(error) => return server_error(error.to_string()),
+        };
     if let Some(decision) = decision {
         attach_routing_headers(&mut response, decision);
     }

crates/switchyard-server/src/response.rs

@@ -3,12 +3,16 @@
 
 //! Response encoding glue for libsy server endpoints.
 
+use std::collections::HashMap;
 use std::error::Error;
 
 use axum::Json;
 use axum::response::{IntoResponse, Response as HttpResponse};
+use futures_util::StreamExt;
 use switchyard_protocol::{LlmResponse, Response as AlgorithmResponse};
-use switchyard_translation::{WireFormat, encode_aggregated_response, encode_stream};
+use switchyard_translation::{
+    WireFormat, encode_aggregated_response, encode_stream, restore_responses_tool_namespaces,
+};
 
 use crate::sse::frame_stream;
 
@@ -21,18 +25,24 @@ pub(crate) fn into_http_response(
     response: AlgorithmResponse,
     target_format: WireFormat,
     served_model: Option<String>,
+    response_namespaces: HashMap<String, String>,
 ) -> Result<HttpResponse, BoxError> {
     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 mut body =
+                encode_aggregated_response(&response, target_format, served_model.as_deref())?;
+            restore_responses_tool_namespaces(&mut body, &response_namespaces);
+            Ok(Json(body).into_response())
+        }
+        LlmResponse::Stream(stream) => {
+            let events = encode_stream(stream, target_format, served_model)?;
+            let events = events.map(move |event| {
+                event.map(|mut value| {
+                    restore_responses_tool_namespaces(&mut value, &response_namespaces);
+                    value
+                })
+            });
+            Ok(frame_stream(Box::pin(events), target_format).into_response())
+        }
     }
 }

crates/switchyard-server/tests/server.rs

@@ -123,6 +123,21 @@ async fn upstream_chat(
             .into_response();
     }
     if body["stream"].as_bool() == Some(true) {
+        // Streamed tool call, for the namespace-on-every-event assertions.
+        if body["messages"][0]["content"] == "mcp-tool-call" {
+            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": "search", "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, Infallible>(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(),
@@ -153,6 +168,30 @@ async fn upstream_chat(
         return Sse::new(stream).into_response();
     }
 
+    // Buffered tool call, the non-streaming counterpart of the branch above.
+    if body["messages"][0]["content"] == "mcp-tool-call" {
+        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": "search", "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();
@@ -2138,3 +2177,119 @@ async fn request_and_upstream_errors_use_the_inbound_wire_format() -> TestResult
     );
     Ok(())
 }
+
+// Returns every `data:` frame of an SSE body as JSON, skipping `[DONE]`.
+fn sse_events(body: &str) -> Vec<Value> {
+    body.lines()
+        .filter_map(|line| line.strip_prefix("data: "))
+        .filter(|data| *data != "[DONE]")
+        .filter_map(|data| serde_json::from_str(data).ok())
+        .collect()
+}
+
+// The Codex request shape: MCP tools wrapped in a `namespace` container.
+fn codex_mcp_responses_request(stream: bool) -> Value {
+    json!({
+        "model": ROUTE_MODEL,
+        "input": "mcp-tool-call",
+        "stream": stream,
+        "tools": [{
+            "type": "namespace",
+            "name": "mcp__open_websearch__",
+            "description": "Web search MCP tools",
+            "tools": [{
+                "type": "function",
+                "name": "search",
+                "description": "Search the web",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"q": {"type": "string"}},
+                    "required": ["q"]
+                }
+            }]
+        }]
+    })
+}
+
+// The container is flattened for a Chat-only upstream, and the namespace is
+// restored on the returned function call.
+#[tokio::test]
+async fn responses_buffered_restores_codex_mcp_namespace() -> TestResult {
+    const MODEL: &str = "model/mcp-buffered";
+    let (upstream, app) = test_app(&[(ROUTE_MODEL, &[MODEL])]).await?;
+
+    let response = send(
+        &app,
+        "POST",
+        "/v1/responses",
+        Some(codex_mcp_responses_request(false)),
+    )
+    .await?;
+
+    assert_eq!(response.status, StatusCode::OK);
+    let body = response.json()?;
+    assert_eq!(body["output"][0]["type"], "function_call");
+    assert_eq!(body["output"][0]["name"], "search");
+    assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch__");
+
+    // The upstream must never see the `namespace` container itself.
+    let calls = upstream.calls.lock().await;
+    let tools = calls[0]["tools"]
+        .as_array()
+        .ok_or("upstream received no tools")?;
+    assert_eq!(tools.len(), 1);
+    assert_eq!(tools[0]["type"], "function");
+    assert_eq!(tools[0]["function"]["name"], "search");
+    assert!(
+        calls[0]["tools"][0].get("namespace").is_none(),
+        "namespace container leaked upstream"
+    );
+    Ok(())
+}
+
+// The namespace has to survive on every output-item event, not only on the
+// terminal aggregate.
+#[tokio::test]
+async fn responses_stream_restores_codex_mcp_namespace() -> TestResult {
+    const MODEL: &str = "model/mcp-stream";
+    let (_upstream, app) = test_app(&[(ROUTE_MODEL, &[MODEL])]).await?;
+
+    let response = send(
+        &app,
+        "POST",
+        "/v1/responses",
+        Some(codex_mcp_responses_request(true)),
+    )
+    .await?;
+
+    assert_eq!(response.status, StatusCode::OK);
+    let events = sse_events(response.text()?);
+
+    let namespace_of = |event_type: &str| -> Option<Value> {
+        events
+            .iter()
+            .find(|event| event["type"] == event_type)
+            .map(|event| event["item"]["namespace"].clone())
+    };
+    assert_eq!(
+        namespace_of("response.output_item.added"),
+        Some(json!("mcp__open_websearch__")),
+        "namespace missing from response.output_item.added"
+    );
+    assert_eq!(
+        namespace_of("response.output_item.done"),
+        Some(json!("mcp__open_websearch__")),
+        "namespace missing from response.output_item.done"
+    );
+
+    let completed = events
+        .iter()
+        .find(|event| event["type"] == "response.completed")
+        .ok_or("stream produced no response.completed event")?;
+    assert_eq!(
+        completed["response"]["output"][0]["namespace"], "mcp__open_websearch__",
+        "namespace missing from the response.completed aggregate"
+    );
+    assert_eq!(completed["response"]["output"][0]["name"], "search");
+    Ok(())
+}