fix(translation): finish Codex tool calls through a LiteLLM gateway

crates/switchyard-server/src/sse.rs

@@ -39,6 +39,9 @@ pub(crate) fn frame_stream(
                 // Preserve the upstream error's fields, apart from credential redaction,
                 // rather than replacing it with a synthesized error.
                 Err(LlmStreamError::Upstream(value)) => {
+                    // Log only the event type; the error message can quote request content.
+                    let event_type = value.get("type").and_then(Value::as_str).unwrap_or("error");
+                    tracing::warn!(event_type, "stream ended with an in-band error event");
                     failed = true;
                     frame_event(target_format, value.clone(), &redactor).unwrap_or_else(|error| {
                         tracing::warn!(error = %error, "in-band error event could not be framed");

crates/switchyard-server/tests/litellm_gateway.rs

@@ -0,0 +1,223 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Regression test for Codex traffic sent through a LiteLLM gateway to Bedrock Claude.
+//!
+//! The stub gateway streams a freeform tool call the way the real gateway does: as
+//! function-style argument deltas whose JSON differs from the finished call. The test runs the
+//! production server, config loader, router, and HTTP client against that stub.
+
+use std::convert::Infallible;
+use std::error::Error;
+use std::io::Write;
+
+use axum::body::Body;
+use axum::http::{Request, StatusCode};
+use axum::response::sse::{Event, Sse};
+use axum::response::{IntoResponse, Response};
+use axum::routing::post;
+use axum::{Json, Router};
+use http_body_util::BodyExt;
+use serde_json::{Value, json};
+use switchyard_server::{build_switchyard_router, config::load_server_state};
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
+use tower::ServiceExt;
+
+type TestResult<T = ()> = Result<T, Box<dyn Error + Send + Sync>>;
+
+/// Both route types send the test's requests to the `claude` target.
+const CLAUDE_ROUTES: [&str; 2] = ["gateway/passthrough", "gateway/stage"];
+
+// `Gateway::drop` stops the server task even when an assertion fails.
+struct Gateway {
+    root: String,
+    task: JoinHandle<std::io::Result<()>>,
+}
+
+impl Gateway {
+    async fn start() -> TestResult<Self> {
+        let app = Router::new().route("/v1/responses", post(gateway));
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let root = format!("http://{}", listener.local_addr()?);
+        let task = tokio::spawn(async move { axum::serve(listener, app).await });
+        Ok(Self { root, task })
+    }
+}
+
+impl Drop for Gateway {
+    fn drop(&mut self) {
+        self.task.abort();
+    }
+}
+
+/// Builds the server from a config file, as `switchyard-server --config` does.
+fn app(gateway: &Gateway) -> TestResult<Router> {
+    let mut config = tempfile::Builder::new().suffix(".toml").tempfile()?;
+    write!(
+        config,
+        r#"
+schema_version = 1
+
+[llm_clients.gateway]
+format = "openai_responses"
+base_url = "{root}/v1"
+max_retries = 0
+
+[targets.claude]
+id = "model/claude"
+llm_client = "gateway"
+
+[targets.gpt]
+id = "model/gpt"
+llm_client = "gateway"
+
+[routes.passthrough]
+id = "gateway/passthrough"
+type = "passthrough"
+target = "claude"
+
+[routes.stage]
+id = "gateway/stage"
+type = "stage_router"
+capable_target = "gpt"
+efficient_target = "claude"
+picker = "efficient_first"
+confidence_threshold = 0.5
+"#,
+        root = gateway.root,
+    )?;
+    Ok(build_switchyard_router(load_server_state(config.path())?))
+}
+
+async fn post_responses(app: &Router, body: &Value) -> TestResult<(StatusCode, String)> {
+    let request = Request::builder()
+        .method("POST")
+        .uri("/v1/responses")
+        .header("content-type", "application/json")
+        .body(Body::from(serde_json::to_vec(body)?))?;
+    let response = app.clone().oneshot(request).await?;
+    let status = response.status();
+    let bytes = response.into_body().collect().await?.to_bytes();
+    Ok((status, String::from_utf8(bytes.to_vec())?))
+}
+
+/// Every `data:` frame of an SSE body, parsed as JSON.
+fn sse_events(body: &str) -> Vec<Value> {
+    body.lines()
+        .filter_map(|line| line.strip_prefix("data: "))
+        .filter_map(|data| serde_json::from_str(data).ok())
+        .collect()
+}
+
+fn event_types(events: &[Value]) -> Vec<&str> {
+    events
+        .iter()
+        .map(|event| event["type"].as_str().unwrap_or_default())
+        .collect()
+}
+
+/// A streamed Codex turn that offers a freeform tool and a function tool.
+fn codex_turn(route: &str) -> Value {
+    json!({
+        "model": route,
+        "stream": true,
+        "tools": [
+            {"type": "custom", "name": "exec", "description": "Run JavaScript", "format": {"type": "text"}},
+            {"type": "function", "name": "shell", "description": "Run a command",
+                "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}}}
+        ],
+        "tool_choice": "auto",
+        "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "List files"}]}]
+    })
+}
+
+/// Answers like LiteLLM in front of Bedrock Claude: a freeform call, then a function call.
+async fn gateway(Json(body): Json<Value>) -> Response {
+    let (custom_events, custom) = litellm_custom_call(0, "ls");
+    let (function_events, function) = function_call(1, r#"{"cmd": "pwd"}"#);
+    let mut events = [custom_events, function_events].concat();
+    let response = json!({
+        "id": "resp_gateway", "object": "response", "model": body["model"], "status": "completed",
+        "output": [custom, function], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
+    });
+    let mut created = response.clone();
+    created["status"] = json!("in_progress");
+    created["output"] = json!([]);
+    events.insert(0, json!({"type": "response.created", "response": created}));
+    events.push(json!({"type": "response.completed", "response": response}));
+    // The real gateway sends data-only frames without `event:` names.
+    let frames = events
+        .into_iter()
+        .map(|event| Ok::<Event, Infallible>(Event::default().data(event.to_string())));
+    Sse::new(futures_util::stream::iter(frames)).into_response()
+}
+
+/// A freeform call in the gateway's shape: function-style argument deltas that spell
+/// `{"content": ...}`, then a finished `custom_tool_call` whose `input` is the bare text.
+fn litellm_custom_call(index: u64, input: &str) -> (Vec<Value>, Value) {
+    let added = json!({"type": "custom_tool_call", "id": "ctc_1", "call_id": "call_1", "name": "exec", "input": ""});
+    let mut done = added.clone();
+    done["input"] = json!(input);
+    let arguments = json!({"content": input}).to_string();
+    let (first, second) = arguments.split_at(arguments.len() / 2);
+    let events = vec![
+        json!({"type": "response.output_item.added", "output_index": index, "item": added}),
+        json!({"type": "response.function_call_arguments.delta", "item_id": "ctc_1", "output_index": index, "delta": first}),
+        json!({"type": "response.function_call_arguments.delta", "item_id": "ctc_1", "output_index": index, "delta": second}),
+        json!({"type": "response.function_call_arguments.done", "item_id": "ctc_1", "output_index": index, "arguments": arguments}),
+        json!({"type": "response.output_item.done", "output_index": index, "item": done}),
+    ];
+    (events, done)
+}
+
+fn function_call(index: u64, arguments: &str) -> (Vec<Value>, Value) {
+    let added = json!({"type": "function_call", "id": "fc_2", "call_id": "call_2", "name": "shell", "arguments": ""});
+    let mut done = added.clone();
+    done["arguments"] = json!(arguments);
+    let events = vec![
+        json!({"type": "response.output_item.added", "output_index": index, "item": added}),
+        json!({"type": "response.function_call_arguments.delta", "item_id": "fc_2", "output_index": index, "delta": arguments}),
+        json!({"type": "response.output_item.done", "output_index": index, "item": done}),
+    ];
+    (events, done)
+}
+
+/// Codex showed "Reconnecting... (stream closed before response.completed)" after every tool
+/// call because the server ended the stream after the gateway's freeform call.
+#[tokio::test]
+async fn codex_tool_call_streams_reach_response_completed() -> TestResult {
+    let gateway = Gateway::start().await?;
+    let app = app(&gateway)?;
+    for route in CLAUDE_ROUTES {
+        let (status, body) = post_responses(&app, &codex_turn(route)).await?;
+        assert_eq!(status, StatusCode::OK, "{route}: {body}");
+        let events = sse_events(&body);
+        let types = event_types(&events);
+        assert_eq!(
+            types.last(),
+            Some(&"response.completed"),
+            "{route}: {types:?}"
+        );
+        let calls: Vec<_> = events
+            .iter()
+            .filter(|event| event["type"] == "response.output_item.done")
+            .map(|event| {
+                let item = &event["item"];
+                (
+                    &item["name"],
+                    item.get("input").unwrap_or(&item["arguments"]),
+                )
+            })
+            .collect();
+        assert_eq!(
+            calls,
+            [
+                (&json!("exec"), &json!("ls")),
+                (&json!("shell"), &json!(r#"{"cmd": "pwd"}"#))
+            ],
+            "{route}"
+        );
+    }
+    Ok(())
+}

crates/switchyard-translation/src/codecs/responses/stream.rs

@@ -188,6 +188,13 @@ fn decode_responses_stream(
             event
                 .get("delta")
                 .and_then(Value::as_str)
+                // The decoder reads a custom call's input only from its completed item.
+                .filter(|_| {
+                    !state
+                        .tool_states
+                        .get(&(output_index as usize))
+                        .is_some_and(|tool| tool.decoded_custom_call)
+                })
                 .map(|delta| {
                     state.decoded_tool_call = true;
                     // Recorded so `response.output_item.done`, which repeats
@@ -269,7 +276,7 @@ fn decode_responses_stream(
                 for (position, item) in items.iter().enumerate() {
                     if let Some(item) = item.as_object() {
                         out.extend(decode_responses_completed_item(item, position, state));
-                        if matches!(out.last(), Some(LlmResponseChunk::StreamError { .. })) {
+                        if matches!(out.last(), Some(LlmResponseChunk::DecodeError { .. })) {
                             return out;
                         }
                     }
@@ -670,15 +677,14 @@ fn decode_responses_output_item_added(
         .or_else(|| item.get("id"))
         .and_then(Value::as_str);
     let name = item.get("name").and_then(Value::as_str);
-    state
-        .tool_states
-        .entry(index)
-        .or_default()
-        .has_decoded_identity |=
+    let tool = state.tool_states.entry(index).or_default();
+    tool.has_decoded_identity |=
         id.is_some_and(|id| !id.is_empty()) && name.is_some_and(|name| !name.is_empty());
+    let custom = item_type == Some("custom_tool_call");
+    tool.decoded_custom_call |= custom;
     // A freeform call's `input` becomes the single `input` argument; it is only complete on
     // the done event, so nothing is emitted for it here beyond id and name.
-    let arguments_delta = if item_type == Some("custom_tool_call") {
+    let arguments_delta = if custom {
         None
     } else {
         item.get("arguments")
@@ -687,12 +693,7 @@ fn decode_responses_output_item_added(
             .map(ToOwned::to_owned)
     };
     if let Some(arguments) = arguments_delta.as_deref() {
-        state
-            .tool_states
-            .entry(index)
-            .or_default()
-            .decoded_arguments
-            .push_str(arguments);
+        tool.decoded_arguments.push_str(arguments);
     }
     vec![LlmResponseChunk::ToolCallDelta {
         index,
@@ -789,13 +790,14 @@ fn decode_responses_completed_item(
     Vec::new()
 }
 
-// A snapshot may extend streamed content, but cannot retract content already sent.
+// A snapshot may extend streamed content, but cannot retract content already sent. A
+// conflict is invalid provider data, so it is a `DecodeError`.
 fn snapshot_suffix(
     decoded: &mut String,
     snapshot: &str,
 ) -> Result<Option<String>, LlmResponseChunk> {
     let Some(suffix) = snapshot.strip_prefix(decoded.as_str()) else {
-        return Err(LlmResponseChunk::StreamError {
+        return Err(LlmResponseChunk::DecodeError {
             message: "Responses snapshot conflicts with streamed content".to_string(),
         });
     };

crates/switchyard-translation/src/codecs/stream.rs

@@ -131,6 +131,11 @@ pub(crate) struct StreamToolState {
     pub(crate) decoded_arguments: String,
     #[serde(default)]
     pub(crate) has_decoded_identity: bool,
+    /// True when the source announced a `custom_tool_call`. The decoder then reads the call's
+    /// input only from the completed item, because some providers also stream argument deltas
+    /// whose JSON does not match that input.
+    #[serde(default)]
+    pub(crate) decoded_custom_call: bool,
     pub(crate) pending_arguments: String,
     pub(crate) started: bool,
     pub(crate) content_index: Option<usize>,