diff --git a/apis/src/anthropic/stream_events/config.rs b/apis/src/anthropic/stream_events/config.rs index 0e3840ef5b..6c3121b10b 100644 --- a/apis/src/anthropic/stream_events/config.rs +++ b/apis/src/anthropic/stream_events/config.rs @@ -22,6 +22,15 @@ pub(crate) struct AnthropicStreamEventsConfig { /// Maximum incomplete SSE event bytes retained between chunks. #[serde(default = "default_max_partial_event_bytes")] pub max_partial_event_bytes: usize, + + /// Maximum number of distinct streaming tool-call content blocks + /// retained per response. Each tool-call index an upstream streams + /// pins per-block state for the response's lifetime, so an upstream + /// that emits unbounded unique indices would grow memory without + /// limit. Exceeding this cap fails the stream closed. Default: + /// 10,000. + #[serde(default = "default_max_tool_blocks")] + pub max_tool_blocks: usize, } /// Default maximum partial event bytes. @@ -29,6 +38,14 @@ fn default_max_partial_event_bytes() -> usize { DEFAULT_JSON_BODY_MAX_BYTES } +/// Default maximum retained streaming tool-call content blocks. +fn default_max_tool_blocks() -> usize { + DEFAULT_MAX_TOOL_BLOCKS +} + +/// Default cap on distinct streaming tool-call content blocks per response. +const DEFAULT_MAX_TOOL_BLOCKS: usize = 10_000; + // ----------------------------------------------------------------------------- // Config Validation // ----------------------------------------------------------------------------- @@ -36,9 +53,19 @@ fn default_max_partial_event_bytes() -> usize { /// Validate the parsed configuration. pub(crate) fn build_config(cfg: AnthropicStreamEventsConfig) -> Result { validate_max_partial_event_bytes(cfg.max_partial_event_bytes)?; + validate_max_tool_blocks(cfg.max_tool_blocks)?; Ok(cfg) } +/// Validate the maximum retained tool-call content block count. +fn validate_max_tool_blocks(value: usize) -> Result<(), FilterError> { + if value == 0 { + return Err("anthropic_stream_events: 'max_tool_blocks' must be greater than 0".into()); + } + + Ok(()) +} + /// Validate the maximum partial SSE event byte limit. fn validate_max_partial_event_bytes(value: usize) -> Result<(), FilterError> { if value == 0 { diff --git a/apis/src/anthropic/stream_events/mod.rs b/apis/src/anthropic/stream_events/mod.rs index c224fc5985..5e509b9810 100644 --- a/apis/src/anthropic/stream_events/mod.rs +++ b/apis/src/anthropic/stream_events/mod.rs @@ -52,6 +52,14 @@ const TOOL_BLOCK_INDEX_SUFFIX: &str = ".index"; /// Metadata key suffix tracking whether a tool call's content block is open. const TOOL_BLOCK_OPEN_SUFFIX: &str = ".open"; +/// Metadata key counting distinct tool-call content blocks opened so far. +/// +/// Each opened block pins per-block state (index and open/closed flag) +/// for the response's lifetime; this count bounds that growth against +/// `max_tool_blocks`. The trailing token is `tool_block_count`, not +/// `tool_block.`, so it never matches [`TOOL_BLOCK_KEY_PREFIX`]. +const TOOL_BLOCK_COUNT_KEY: &str = "anthropic_stream.tool_block_count"; + /// Metadata key for the finish reason from the upstream provider. const FINISH_REASON_KEY: &str = "anthropic_stream.finish_reason"; @@ -92,6 +100,7 @@ const ARMED_KEY: &str = "anthropic_stream.armed"; /// ```yaml /// filter: anthropic_stream_events /// max_partial_event_bytes: 10485760 +/// max_tool_blocks: 10000 /// ``` pub struct AnthropicStreamEventsFilter { /// Parsed and validated configuration. @@ -109,6 +118,28 @@ impl AnthropicStreamEventsFilter { let validated = build_config(cfg)?; Ok(Box::new(Self { config: validated })) } + + /// Decode and transform one response body chunk under the filter's + /// configured partial-event and tool-block limits. + /// + /// # Errors + /// + /// Returns [`FilterError`] if a partial SSE event or the retained + /// tool-call block count exceeds its configured limit. + fn process_response_chunk( + &self, + ctx: &mut HttpFilterContext<'_>, + bytes: &Bytes, + end_of_stream: bool, + ) -> Result, FilterError> { + decode_and_process_chunk( + ctx, + bytes, + end_of_stream, + self.config.max_partial_event_bytes, + self.config.max_tool_blocks, + ) + } } #[async_trait] @@ -167,8 +198,8 @@ impl HttpFilter for AnthropicStreamEventsFilter { let Some(bytes) = body.as_ref() else { if end_of_stream { - let empty = Bytes::new(); - let output = decode_and_process_chunk(ctx, &empty, true, self.config.max_partial_event_bytes)? + let output = self + .process_response_chunk(ctx, &Bytes::new(), true)? .unwrap_or_default(); if !output.is_empty() { *body = Some(output); @@ -177,8 +208,7 @@ impl HttpFilter for AnthropicStreamEventsFilter { return Ok(FilterAction::Continue); }; - let Some(output) = decode_and_process_chunk(ctx, bytes, end_of_stream, self.config.max_partial_event_bytes)? - else { + let Some(output) = self.process_response_chunk(ctx, bytes, end_of_stream)? else { *body = Some(Bytes::new()); return Ok(FilterAction::Continue); }; @@ -200,6 +230,7 @@ fn decode_and_process_chunk( bytes: &Bytes, end_of_stream: bool, max_partial_event_bytes: usize, + max_tool_blocks: usize, ) -> Result, FilterError> { let combined = combine_pending_utf8(ctx, bytes); let Some(valid_up_to) = valid_utf8_prefix_len(ctx, combined.as_slice(), end_of_stream) else { @@ -216,7 +247,7 @@ fn decode_and_process_chunk( return Ok(None); } - process_sse_chunk(ctx, chunk_str, end_of_stream, max_partial_event_bytes).map(Some) + process_sse_chunk(ctx, chunk_str, end_of_stream, max_partial_event_bytes, max_tool_blocks).map(Some) } /// Prefix any incomplete UTF-8 bytes retained from the previous chunk. @@ -302,6 +333,7 @@ fn process_sse_chunk( chunk_str: &str, end_of_stream: bool, max_partial_event_bytes: usize, + max_tool_blocks: usize, ) -> Result { let leftover = ctx.filter_metadata.get(LINE_BUFFER_KEY).cloned().unwrap_or_default(); let combined = format!("{leftover}{chunk_str}"); @@ -322,7 +354,7 @@ fn process_sse_chunk( while let Some((event_block, rest)) = remaining.split_once("\n\n") { remaining = rest; - process_event_block(ctx, event_block, &mut output); + process_event_block(ctx, event_block, &mut output, max_tool_blocks)?; } let to_buffer = if pending_cr { @@ -410,7 +442,12 @@ fn is_streaming_request(ctx: &HttpFilterContext<'_>) -> bool { /// Collects all `data` fields into one newline-delimited payload before /// processing it. Accepts bare `data`, `data: value`, and `data:value` /// per the SSE specification. -fn process_event_block(ctx: &mut HttpFilterContext<'_>, block: &str, output: &mut Vec) { +fn process_event_block( + ctx: &mut HttpFilterContext<'_>, + block: &str, + output: &mut Vec, + max_tool_blocks: usize, +) -> Result<(), FilterError> { let mut event_data = None::>; for line in block.lines() { @@ -436,9 +473,11 @@ fn process_event_block(ctx: &mut HttpFilterContext<'_>, block: &str, output: &mu if data == OPENAI_DONE_SENTINEL { emit_done(ctx, output); } else if let Ok(chunk) = serde_json::from_str::(&data) { - transform_chunk(ctx, &chunk, output); + transform_chunk(ctx, &chunk, output, max_tool_blocks)?; } } + + Ok(()) } // ----------------------------------------------------------------------------- @@ -446,7 +485,12 @@ fn process_event_block(ctx: &mut HttpFilterContext<'_>, block: &str, output: &mu // ----------------------------------------------------------------------------- /// Transform a single `OpenAI` SSE chunk into Anthropic events. -fn transform_chunk(ctx: &mut HttpFilterContext<'_>, chunk: &Value, output: &mut Vec) { +fn transform_chunk( + ctx: &mut HttpFilterContext<'_>, + chunk: &Value, + output: &mut Vec, + max_tool_blocks: usize, +) -> Result<(), FilterError> { let started = ctx .filter_metadata .get(STREAM_STATE_KEY) @@ -458,7 +502,7 @@ fn transform_chunk(ctx: &mut HttpFilterContext<'_>, chunk: &Value, output: &mut if let Some(choice) = extract_first_choice(chunk) { if let Some(delta) = choice.get("delta") { - transform_delta(ctx, delta, output); + transform_delta(ctx, delta, output, max_tool_blocks)?; } if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) { ctx.set_metadata(FINISH_REASON_KEY, reason.to_owned()); @@ -472,6 +516,8 @@ fn transform_chunk(ctx: &mut HttpFilterContext<'_>, chunk: &Value, output: &mut { ctx.set_metadata(OUTPUT_TOKENS_KEY, ot.to_string()); } + + Ok(()) } /// Emit the initial `message_start` event and mark the stream as started. @@ -522,7 +568,12 @@ fn generate_timestamp_id() -> u128 { // ----------------------------------------------------------------------------- /// Transform a delta object from a streaming chunk. -fn transform_delta(ctx: &mut HttpFilterContext<'_>, delta: &Value, output: &mut Vec) { +fn transform_delta( + ctx: &mut HttpFilterContext<'_>, + delta: &Value, + output: &mut Vec, + max_tool_blocks: usize, +) -> Result<(), FilterError> { if let Some(content) = delta.get("content").and_then(Value::as_str) { emit_text_delta(ctx, content, output); } @@ -530,9 +581,11 @@ fn transform_delta(ctx: &mut HttpFilterContext<'_>, delta: &Value, output: &mut if let Some(Value::Array(tool_calls)) = delta.get("tool_calls") { close_text_block_if_open(ctx, output); for tc in tool_calls { - transform_tool_delta(ctx, tc, output); + transform_tool_delta(ctx, tc, output, max_tool_blocks)?; } } + + Ok(()) } /// Emit a text content delta, opening a new block if needed. @@ -569,16 +622,21 @@ fn emit_text_delta(ctx: &mut HttpFilterContext<'_>, content: &str, output: &mut // ----------------------------------------------------------------------------- /// Transform a tool call delta into Anthropic content block events. -fn transform_tool_delta(ctx: &mut HttpFilterContext<'_>, tc: &Value, output: &mut Vec) { +fn transform_tool_delta( + ctx: &mut HttpFilterContext<'_>, + tc: &Value, + output: &mut Vec, + max_tool_blocks: usize, +) -> Result<(), FilterError> { let tool_call_key = tool_call_key(tc); - if let Some(id) = tc.get("id").and_then(Value::as_str) - && !is_tool_block_open(ctx, &tool_call_key) - { - emit_tool_block_start(ctx, &tool_call_key, tc, id, output); + if tc.get("id").and_then(Value::as_str).is_some() && !is_tool_block_open(ctx, &tool_call_key) { + emit_tool_block_start(ctx, &tool_call_key, tc, output, max_tool_blocks)?; } emit_tool_arguments_delta(ctx, &tool_call_key, tc, output); + + Ok(()) } /// Close any open text content block and advance the block index. @@ -605,14 +663,29 @@ fn tool_call_key(tc: &Value) -> String { } /// Emit a `content_block_start` for a tool-use block. +/// +/// # Errors +/// +/// Fails closed with a [`FilterError`] before opening the +/// `max_tool_blocks + 1`th block, bounding the per-response tool-call +/// state that would otherwise grow with every unique upstream index. fn emit_tool_block_start( ctx: &mut HttpFilterContext<'_>, tool_call_key: &str, tc: &Value, - id: &str, output: &mut Vec, -) { + max_tool_blocks: usize, +) -> Result<(), FilterError> { + let opened = get_tool_block_count(ctx); + if opened >= max_tool_blocks { + return Err(format!( + "anthropic_stream_events: streaming tool-call content blocks exceed max_tool_blocks ({max_tool_blocks})" + ) + .into()); + } + let idx = get_block_index(ctx); + let id = tc.get("id").and_then(Value::as_str).unwrap_or_default(); let name = tc .get("function") .and_then(|f| f.get("name")) @@ -632,6 +705,9 @@ fn emit_tool_block_start( set_tool_block_index(ctx, tool_call_key, idx); set_tool_block_open(ctx, tool_call_key, true); increment_block_index(ctx); + ctx.set_metadata(TOOL_BLOCK_COUNT_KEY, (opened + 1).to_string()); + + Ok(()) } /// Emit an `input_json_delta` if the tool call has non-empty arguments. @@ -773,6 +849,14 @@ fn increment_block_index(ctx: &mut HttpFilterContext<'_>) { ctx.set_metadata(BLOCK_INDEX_KEY, (current + 1).to_string()); } +/// Return how many tool-call content blocks have opened this response. +fn get_tool_block_count(ctx: &HttpFilterContext<'_>) -> usize { + ctx.filter_metadata + .get(TOOL_BLOCK_COUNT_KEY) + .and_then(|v| v.parse().ok()) + .unwrap_or(0) +} + /// Build the metadata key for a tool call's Anthropic block index. fn tool_block_index_key(tool_call_key: &str) -> String { format!("{TOOL_BLOCK_KEY_PREFIX}{tool_call_key}{TOOL_BLOCK_INDEX_SUFFIX}") @@ -1697,6 +1781,51 @@ mod tests { ); } + #[test] + fn custom_max_tool_blocks_parses() { + let yaml: serde_yaml::Value = serde_yaml::from_str("max_tool_blocks: 5").unwrap(); + let result = AnthropicStreamEventsFilter::from_config(&yaml); + + assert!( + result.is_ok(), + "streaming filter should accept a custom max_tool_blocks" + ); + } + + #[test] + fn zero_max_tool_blocks_rejected() { + let yaml: serde_yaml::Value = serde_yaml::from_str("max_tool_blocks: 0").unwrap(); + let result = AnthropicStreamEventsFilter::from_config(&yaml); + + assert!(result.is_err(), "streaming filter should reject a zero max_tool_blocks"); + } + + #[test] + fn exceeding_max_tool_blocks_fails_closed() { + let (filter, mut ctx) = make_filter_and_context_from_yaml("max_tool_blocks: 2"); + + let block = |index: u64| { + format!( + "data: {{\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":{index},\"id\":\"call_{index}\",\"function\":{{\"name\":\"f{index}\",\"arguments\":\"{{}}\"}}}}]}},\"index\":0}}]}}\n\n" + ) + }; + + let mut body0 = Some(Bytes::from(block(0))); + drop(filter.on_response_body(&mut ctx, &mut body0, false).unwrap()); + + let mut body1 = Some(Bytes::from(block(1))); + drop(filter.on_response_body(&mut ctx, &mut body1, false).unwrap()); + + let mut body2 = Some(Bytes::from(block(2))); + let result = filter.on_response_body(&mut ctx, &mut body2, false); + + let err = result.unwrap_err(); + assert!( + err.to_string().contains("max_tool_blocks"), + "exceeding the tool-block cap should fail closed and mention max_tool_blocks, got: {err}" + ); + } + // Test Utilities fn event_data(output: &str, event_type: &str) -> Value { diff --git a/docs/filters/anthropic_stream_events.md b/docs/filters/anthropic_stream_events.md index e102df0fce..1614352259 100644 --- a/docs/filters/anthropic_stream_events.md +++ b/docs/filters/anthropic_stream_events.md @@ -14,6 +14,7 @@ Arms automatically when an upstream classifier or transform filter sets `anthrop | Field | Type | Required | Description | |-------|------|---------|-------------| | `max_partial_event_bytes` | integer | no | Maximum incomplete SSE event bytes retained between chunks. | +| `max_tool_blocks` | integer | no | Maximum number of distinct streaming tool-call content blocks retained per response. Each tool-call index an upstream streams pins per-block state for the response's lifetime, so an upstream that emits unbounded unique indices would grow memory without limit. Exceeding this cap fails the stream closed. Default: 10,000. | ## Examples @@ -28,4 +29,5 @@ filter: anthropic_stream_events ```yaml filter: anthropic_stream_events max_partial_event_bytes: 10485760 +max_tool_blocks: 10000 ``` diff --git a/examples/configs/anthropic/messages-to-openai.yaml b/examples/configs/anthropic/messages-to-openai.yaml index e363cf9a19..7d5b419e56 100644 --- a/examples/configs/anthropic/messages-to-openai.yaml +++ b/examples/configs/anthropic/messages-to-openai.yaml @@ -31,6 +31,11 @@ filter_chains: # skips filter dispatch entirely for non-SSE responses. - filter: anthropic_stream_events max_partial_event_bytes: 10485760 + # Cap on distinct streaming tool-call content blocks retained per + # response. Bounds per-response memory when an upstream streams + # many unique tool-call indices; exceeding it fails the stream + # closed. + max_tool_blocks: 10000 response_conditions: - when: headers: diff --git a/tests/integration/tests/suite/anthropic_messages.rs b/tests/integration/tests/suite/anthropic_messages.rs index 0a4a2afc82..83a5630183 100644 --- a/tests/integration/tests/suite/anthropic_messages.rs +++ b/tests/integration/tests/suite/anthropic_messages.rs @@ -204,6 +204,60 @@ fn streaming_collects_full_text() { assert!(!full_text.is_empty(), "collected text should not be empty"); } +#[test] +fn streaming_tool_calls_within_cap_completes() { + let backend = Backend::fixed(&tool_call_stream_sse()) + .header("content-type", "text/event-stream") + .header("cache-control", "no-cache") + .start_with_shutdown(); + let proxy_port = free_port(); + let config = Config::from_yaml(&transform_yaml(proxy_port, backend.port(), 5)).unwrap(); + let proxy = start_proxy(&config); + + let raw = http_send( + proxy.addr(), + &anthropic_post( + "/v1/messages", + r#"{"model":"mock-model","messages":[{"role":"user","content":"call tools"}],"max_tokens":64,"stream":true}"#, + ), + ); + let body = parse_body(&raw); + + // Three distinct tool-call indices open three blocks, all under the + // cap of 5, so the transform runs to completion. + assert!( + body.contains("event: message_stop"), + "a tool-call stream within max_tool_blocks should complete with message_stop; body: {body}" + ); +} + +#[test] +fn streaming_tool_calls_exceeding_cap_fails_closed() { + let backend = Backend::fixed(&tool_call_stream_sse()) + .header("content-type", "text/event-stream") + .header("cache-control", "no-cache") + .start_with_shutdown(); + let proxy_port = free_port(); + // Same stream as the within-cap test; only the cap changes. The third + // distinct tool-call index exceeds max_tool_blocks and fails closed. + let config = Config::from_yaml(&transform_yaml(proxy_port, backend.port(), 2)).unwrap(); + let proxy = start_proxy(&config); + + let raw = http_send( + proxy.addr(), + &anthropic_post( + "/v1/messages", + r#"{"model":"mock-model","messages":[{"role":"user","content":"call tools"}],"max_tokens":64,"stream":true}"#, + ), + ); + let body = parse_body(&raw); + + assert!( + !body.contains("event: message_stop"), + "exceeding max_tool_blocks should fail the stream closed before message_stop; body: {body}" + ); +} + #[test] fn non_streaming_with_temperature() { let recording = Recording::load("anthropic/messages/temperature.json"); @@ -443,6 +497,58 @@ insecure_options: ) } +fn transform_yaml(proxy_port: u16, backend_port: u16, max_tool_blocks: usize) -> String { + format!( + r#" +listeners: + - name: test + address: "127.0.0.1:{proxy_port}" + filter_chains: [transform] + +filter_chains: + - name: transform + filters: + - filter: anthropic_messages_format + on_invalid: continue + - filter: anthropic_to_openai + max_body_bytes: 1048576 + - filter: anthropic_stream_events + max_tool_blocks: {max_tool_blocks} + - filter: router + routes: + - path_prefix: "/" + cluster: mock + - filter: load_balancer + clusters: + - name: mock + endpoints: + - "127.0.0.1:{backend_port}" + +insecure_options: + allow_private_endpoints: true +"# + ) +} + +/// OpenAI Chat Completions SSE with three tool-call deltas at distinct +/// indices, a finish chunk, and the `[DONE]` sentinel. Each distinct +/// index opens a new Anthropic tool-use content block in the transform, +/// so the stream pins three blocks of per-block metadata. +fn tool_call_stream_sse() -> String { + let block = |index: u64| { + format!( + "data: {{\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":{index},\"id\":\"call_{index}\",\"function\":{{\"name\":\"f{index}\",\"arguments\":\"{{}}\"}}}}]}},\"index\":0}}]}}\n\n" + ) + }; + + format!( + "{}{}{}data: {{\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{{\"delta\":{{}},\"index\":0,\"finish_reason\":\"tool_calls\"}}]}}\n\ndata: [DONE]\n\n", + block(0), + block(1), + block(2), + ) +} + fn anthropic_post(path: &str, body: &str) -> String { format!( "POST {path} HTTP/1.1\r\n\