diff --git a/crates/switchyard-server/src/sse.rs b/crates/switchyard-server/src/sse.rs index dd2cb1595f..b7794309e1 100644 --- a/crates/switchyard-server/src/sse.rs +++ b/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"); diff --git a/crates/switchyard-server/tests/litellm_gateway.rs b/crates/switchyard-server/tests/litellm_gateway.rs new file mode 100644 index 0000000000..a83d00998f --- /dev/null +++ b/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 = Result>; + +/// 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>, +} + +impl Gateway { + async fn start() -> TestResult { + 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 { + 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 { + 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) -> 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::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) { + 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) { + 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(()) +} diff --git a/crates/switchyard-translation/src/codecs/responses/stream.rs b/crates/switchyard-translation/src/codecs/responses/stream.rs index 6d30a60e2d..35f3a1411a 100644 --- a/crates/switchyard-translation/src/codecs/responses/stream.rs +++ b/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, 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(), }); }; diff --git a/crates/switchyard-translation/src/codecs/stream.rs b/crates/switchyard-translation/src/codecs/stream.rs index db83015a7c..04f3b0f299 100644 --- a/crates/switchyard-translation/src/codecs/stream.rs +++ b/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, diff --git a/crates/switchyard-translation/tests/stream_translation.rs b/crates/switchyard-translation/tests/stream_translation.rs index 26d65adae9..e4dd05f87b 100644 --- a/crates/switchyard-translation/tests/stream_translation.rs +++ b/crates/switchyard-translation/tests/stream_translation.rs @@ -203,6 +203,39 @@ fn preserved_same_format_replay_stops_after_an_error() -> TestResult { Ok(()) } +// A finished item that contradicts text that was already streamed reaches a same-format client +// as an error event in its place. Before, the contradicting event was replayed and the stream +// ended with no error event. +#[test] +fn preserved_same_format_snapshot_conflict_replays_as_an_error() -> TestResult { + let format = WireFormat::OpenAiResponses; + let events = [ + json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "Hello"}), + json!({"type": "response.output_item.done", "output_index": 0, "item": {"type": "message", + "role": "assistant", "content": [{"type": "output_text", "text": "Goodbye"}]}}), + ]; + let engine = TranslationEngine::default(); + let mut decode_state = StreamTranslationState::new(format, format); + let mut encode_state = StreamTranslationState::new(format, format); + let mut replayed = Vec::new(); + for event in events { + let preserved = engine.decode_stream_event(&mut decode_state, format, event)?; + replayed.extend(engine.encode_stream_event(&mut encode_state, format, preserved)?); + } + + let types: Vec<_> = replayed.iter().map(|event| &event["type"]).collect(); + assert_eq!( + types, + [&json!("response.output_text.delta"), &json!("error")], + "{replayed:?}" + ); + assert_eq!( + replayed[1]["message"], + "Responses snapshot conflicts with streamed content" + ); + Ok(()) +} + #[test] fn anthropic_errors_include_type_and_terminate_stream() -> TestResult { let engine = TranslationEngine::default(); @@ -3419,7 +3452,7 @@ fn responses_terminal_snapshots_recover_missing_output_once() -> TestResult { )?; assert!(matches!( decoded.normalized().last(), - Some(LlmResponseChunk::StreamError { .. }) + Some(LlmResponseChunk::DecodeError { .. }) )); assert!( !decoded @@ -3430,3 +3463,65 @@ fn responses_terminal_snapshots_recover_missing_output_once() -> TestResult { } Ok(()) } + +// LiteLLM in front of Bedrock streams a freeform call as function-style argument deltas that +// spell `{"content": ...}`. OpenAI streams it with custom input deltas. Both must decode, +// without a conflict, to one call whose only arguments are the finished input. +#[test] +fn responses_custom_tool_calls_decode_from_the_finished_item() -> TestResult { + let engine = TranslationEngine::default(); + let source = WireFormat::OpenAiResponses; + let done = json!({"type": "custom_tool_call", "id": "ctc_1", "call_id": "call_1", "name": "exec", "input": "ls"}); + let added = json!({"type": "response.output_item.added", "output_index": 1, + "item": {"type": "custom_tool_call", "id": "ctc_1", "call_id": "call_1", "name": "exec", "input": ""}}); + let item_done = + json!({"type": "response.output_item.done", "output_index": 1, "item": done.clone()}); + let completed = json!({"type": "response.completed", "response": {"id": "resp_1", "status": "completed", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}, done]}}); + let litellm = vec![ + added.clone(), + json!({"type": "response.function_call_arguments.delta", "output_index": 1, "delta": "{\"content\": "}), + json!({"type": "response.function_call_arguments.delta", "output_index": 1, "delta": "\"ls\"}"}), + json!({"type": "response.function_call_arguments.done", "output_index": 1, "arguments": "{\"content\": \"ls\"}"}), + item_done.clone(), + completed.clone(), + ]; + let openai = vec![ + added, + json!({"type": "response.custom_tool_call_input.delta", "output_index": 1, "delta": "ls"}), + json!({"type": "response.custom_tool_call_input.done", "output_index": 1, "input": "ls"}), + item_done, + completed, + ]; + for (shape, events) in [("litellm", litellm), ("openai", openai)] { + let mut state = StreamTranslationState::new(source, WireFormat::OpenAiChat); + let mut chunks = Vec::new(); + for event in events { + chunks.extend( + engine + .decode_stream_event(&mut state, source, event)? + .normalized() + .to_vec(), + ); + } + assert!( + !chunks.iter().any(|chunk| matches!( + chunk, + LlmResponseChunk::DecodeError { .. } | LlmResponseChunk::StreamError { .. } + )), + "{shape}: {chunks:?}" + ); + let arguments: String = chunks + .iter() + .filter_map(|chunk| match chunk { + LlmResponseChunk::ToolCallDelta { + arguments_delta: Some(arguments), + .. + } => Some(arguments.as_str()), + _ => None, + }) + .collect(); + assert_eq!(arguments, r#"{"input":"ls"}"#, "{shape}"); + } + Ok(()) +}