fix: preserve Codex MCP namespaces through translation

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<RawResponse> {
         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<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": "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]

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());
     }

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<String>,
+    request_extensions: ProviderExtensions,
 ) -> 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 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())
+        }
     }
 }

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, 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(),
@@ -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<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 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::<Vec<_>>()
+        })
+        .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(())
+}