From 114250f8caf14c812be4c1d14c86abd438438fb6 Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Mon, 31 Aug 2026 16:26:08 -0400 Subject: [PATCH] fix(gdk): preserve tool-call indices in streaming responses (#11637) --- .../src/conversation/message.rs | 22 +++ .../src/conversation/tool_request.rs | 12 +- .../src/formats/anthropic.rs | 169 +++++++++++++++--- .../src/formats/openai.rs | 73 +++++++- crates/goose-sdk/src/bindings.rs | 44 +++-- documentation/docs/gdk/index.md | 2 +- documentation/src/data/gdk-api.json | 6 + 7 files changed, 279 insertions(+), 49 deletions(-) diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index 1d5a3cf75..55bae8044 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -161,6 +161,12 @@ pub const TOOL_META_EXTERNAL_DISPATCH_KEY: &str = "goose.external_dispatch"; /// for this tool call. Used to make the title survive session reload. pub const TOOL_META_TITLE_KEY: &str = "goose.toolSummary.title"; +/// Key under `ToolRequest.tool_meta` storing the provider-reported index of the +/// tool call within the streamed response. Streaming clients need this to +/// correlate incremental argument fragments with the right call when a model +/// emits several tool calls in parallel. +pub const TOOL_META_PROVIDER_INDEX_KEY: &str = "goose.toolCall.providerIndex"; + /// Key under `ToolRequest.tool_meta` storing the LLM-generated chain summary /// for the chain that starts at this tool request. Shape: `{ "summary": String, /// "count": u64 }`. Only attached to the FIRST tool request in a chain. @@ -460,6 +466,22 @@ impl MessageContentBlock { }) } + pub fn tool_request_with_provider_index>( + id: S, + tool_call: ToolResult, + metadata: Option<&ProviderMetadata>, + provider_index: i32, + ) -> Self { + MessageContentBlock::ToolRequest(ToolRequest { + id: id.into(), + tool_call, + metadata: metadata.cloned(), + tool_meta: Some(serde_json::json!({ + TOOL_META_PROVIDER_INDEX_KEY: provider_index, + })), + }) + } + pub fn tool_response>(id: S, tool_result: ToolResult) -> Self { MessageContentBlock::ToolResponse(ToolResponse { id: id.into(), diff --git a/crates/goose-provider-types/src/conversation/tool_request.rs b/crates/goose-provider-types/src/conversation/tool_request.rs index f8064b004..f5dc55b0c 100644 --- a/crates/goose-provider-types/src/conversation/tool_request.rs +++ b/crates/goose-provider-types/src/conversation/tool_request.rs @@ -1,6 +1,6 @@ use super::message::{ ToolChainSummary, ToolNameParts, ToolRequest, TOOL_META_CHAIN_SUMMARY_KEY, - TOOL_META_EXTERNAL_DISPATCH_KEY, TOOL_META_TITLE_KEY, + TOOL_META_EXTERNAL_DISPATCH_KEY, TOOL_META_PROVIDER_INDEX_KEY, TOOL_META_TITLE_KEY, }; impl<'a> From<&'a str> for ToolNameParts<'a> { @@ -70,6 +70,16 @@ impl ToolRequest { .and_then(|v| v.as_str()) } + /// Provider-reported index of this tool call within the streamed response. + /// See [`TOOL_META_PROVIDER_INDEX_KEY`]. + pub fn provider_index(&self) -> Option { + self.tool_meta + .as_ref() + .and_then(|v| v.get(TOOL_META_PROVIDER_INDEX_KEY)) + .and_then(|v| v.as_i64()) + .map(|index| index as i32) + } + pub fn generated_chain_summary(&self) -> Option { let obj = self .tool_meta diff --git a/crates/goose-provider-types/src/formats/anthropic.rs b/crates/goose-provider-types/src/formats/anthropic.rs index de34e39f1..bc714ce12 100644 --- a/crates/goose-provider-types/src/formats/anthropic.rs +++ b/crates/goose-provider-types/src/formats/anthropic.rs @@ -901,9 +901,21 @@ where signature: String, } + fn block_index(event_data: &Value) -> Option { + event_data + .get("index") + .and_then(|v| v.as_i64()) + .map(|index| index as i32) + } + try_stream! { - let mut accumulated_tool_calls: std::collections::HashMap = std::collections::HashMap::new(); - let mut current_tool_id: Option = None; + struct StreamingToolCall { + id: String, + name: String, + arguments: String, + } + + let mut accumulated_tool_calls: std::collections::HashMap = std::collections::HashMap::new(); let mut final_usage: Option = None; let mut message_id: Option = None; let mut thinking: Option = None; @@ -958,11 +970,16 @@ where if let Some(content_block) = event.data.get("content_block") { match content_block.get(TYPE_FIELD).and_then(|v| v.as_str()) { Some(TOOL_USE_TYPE) => { - if let Some(id) = content_block.get("id").and_then(|v| v.as_str()) { - current_tool_id = Some(id.to_string()); - if let Some(name) = content_block.get("name").and_then(|v| v.as_str()) { - accumulated_tool_calls.insert(id.to_string(), (name.to_string(), String::new())); - } + if let (Some(index), Some(id), Some(name)) = ( + block_index(&event.data), + content_block.get("id").and_then(|v| v.as_str()), + content_block.get("name").and_then(|v| v.as_str()), + ) { + accumulated_tool_calls.insert(index, StreamingToolCall { + id: id.to_string(), + name: name.to_string(), + arguments: String::new(), + }); } } Some(THINKING_TYPE) => { @@ -1003,10 +1020,10 @@ where yield (Some(message), None); } Ok(ContentBlockDelta::InputJsonDelta { partial_json }) => { - if let Some(tool_id) = ¤t_tool_id { - if let Some((_name, args)) = accumulated_tool_calls.get_mut(tool_id) { - args.push_str(&partial_json); - } + if let Some(call) = block_index(&event.data) + .and_then(|index| accumulated_tool_calls.get_mut(&index)) + { + call.arguments.push_str(&partial_json); } } Ok(ContentBlockDelta::ThinkingDelta { thinking: t }) => { @@ -1035,17 +1052,18 @@ where yield (Some(message), None); } } - if let Some(tool_id) = current_tool_id.take() { - if let Some((name, args)) = accumulated_tool_calls.remove(&tool_id) { - let parsed_args = if args.is_empty() { + if let Some(index) = block_index(&event.data) { + if let Some(call) = accumulated_tool_calls.remove(&index) { + let StreamingToolCall { id, name, arguments } = call; + let parsed_args = if arguments.is_empty() { json!({}) } else { - match crate::json::parse_tool_arguments(&args) { + match crate::json::parse_tool_arguments(&arguments) { Some(parsed) => parsed, None => { - let message_text = crate::json::truncation_error_message(&args) + let message_text = crate::json::truncation_error_message(&arguments) .unwrap_or_else(|| { - format!("Could not parse tool arguments: {args}") + format!("Could not parse tool arguments: {arguments}") }); let error = ErrorData::new( ErrorCode::INVALID_PARAMS, @@ -1055,7 +1073,7 @@ where let mut message = Message::new( Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContentBlock::tool_request(tool_id, Err(error))], + vec![MessageContentBlock::tool_request_with_provider_index(id, Err(error), None, index)], ); message.id = message_id.clone(); yield (Some(message), None); @@ -1069,7 +1087,7 @@ where let mut message = Message::new( rmcp::model::Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContentBlock::tool_request(tool_id, Ok(tool_call))], + vec![MessageContentBlock::tool_request_with_provider_index(id, Ok(tool_call), None, index)], ); message.id = message_id.clone(); yield (Some(message), None); @@ -1158,10 +1176,10 @@ where // content_block_stop, so its args are truncated rather than complete. if !accumulated_tool_calls.is_empty() { let truncated_by_limit = stop_reason.as_deref() == Some("max_tokens"); - let mut ids: Vec = accumulated_tool_calls.keys().cloned().collect(); - ids.sort(); - for id in ids { - if let Some((_name, args)) = accumulated_tool_calls.remove(&id) { + let mut indices: Vec = accumulated_tool_calls.keys().copied().collect(); + indices.sort(); + for index in indices { + if let Some(StreamingToolCall { id, arguments: args, .. }) = accumulated_tool_calls.remove(&index) { let guidance = if truncated_by_limit { "The model's response was truncated — it hit the output token limit while generating this tool call. \ Try increasing max_tokens for this provider or breaking the task into smaller steps." @@ -1178,7 +1196,7 @@ where let mut message = Message::new( Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContentBlock::tool_request(id, Err(error))], + vec![MessageContentBlock::tool_request_with_provider_index(id, Err(error), None, index)], ); message.id = message_id.clone(); yield (Some(message), None); @@ -2427,6 +2445,74 @@ mod tests { response_to_streaming_message(stream).collect().await } + #[tokio::test] + async fn test_streaming_reassembles_interleaved_parallel_tool_calls() { + let events = concat!( + r#"data: {"type":"message_start","message":{"id":"msg_par","role":"assistant","content":[],"model":"claude-opus-4-6","usage":{"input_tokens":5,"output_tokens":0}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool_a","name":"search","input":{}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_b","name":"write","input":{}}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"path\":"}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"query\":"}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"/tmp/a.md\"}"}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"rust\"}"}}"#, + "\n", + r#"data: {"type":"content_block_stop","index":1}"#, + "\n", + r#"data: {"type":"content_block_stop","index":0}"#, + "\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":20}}"#, + "\n", + r#"data: {"type":"message_stop"}"#, + ); + + let requests = collect_tool_requests(events).await; + + assert_eq!(requests.len(), 2); + let write = requests + .iter() + .find(|r| r.id == "tool_b") + .expect("write tool request"); + assert_eq!(write.provider_index(), Some(1)); + let write_call = write.tool_call.as_ref().expect("write args parsed"); + assert_eq!(write_call.name, "write"); + assert_eq!( + write_call.arguments.as_ref().unwrap()["path"], + json!("/tmp/a.md") + ); + + let search = requests + .iter() + .find(|r| r.id == "tool_a") + .expect("search tool request"); + assert_eq!(search.provider_index(), Some(0)); + let search_call = search.tool_call.as_ref().expect("search args parsed"); + assert_eq!(search_call.name, "search"); + assert_eq!( + search_call.arguments.as_ref().unwrap()["query"], + json!("rust") + ); + } + + async fn collect_tool_requests(events: &str) -> Vec { + let mut requests = Vec::new(); + for result in collect_stream_results(events).await { + if let Ok((Some(msg), _usage)) = result { + for content in &msg.content { + if let MessageContentBlock::ToolRequest(req) = content { + requests.push(req.clone()); + } + } + } + } + requests + } + #[tokio::test] async fn test_streaming_preserves_cache_tokens_through_delta_merge() { let events = concat!( @@ -2739,6 +2825,41 @@ mod tests { ); } + #[tokio::test] + async fn test_streaming_unfinished_tool_calls_keep_provider_indices() { + let events = concat!( + r#"data: {"type":"message_start","message":{"id":"msg_t3","role":"assistant","content":[],"model":"claude-opus-4-6","usage":{"input_tokens":10,"output_tokens":0}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool_open_a","name":"search","input":{}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_open_b","name":"write","input":{}}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"query\":\"ru"}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"path\":\"/re"}}"#, + "\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":8192}}"#, + "\n", + r#"data: {"type":"message_stop"}"#, + ); + + let requests = collect_tool_requests(events).await; + + assert_eq!(requests.len(), 2); + let indexed: Vec<(String, Option)> = requests + .iter() + .map(|r| (r.id.clone(), r.provider_index())) + .collect(); + assert_eq!( + indexed, + vec![ + ("tool_open_a".to_string(), Some(0)), + ("tool_open_b".to_string(), Some(1)), + ] + ); + assert!(requests.iter().all(|r| r.tool_call.is_err())); + } + #[tokio::test] async fn test_streaming_complete_tool_call_unaffected() { // Regression guard: a normal, complete tool call must still parse and diff --git a/crates/goose-provider-types/src/formats/openai.rs b/crates/goose-provider-types/src/formats/openai.rs index 0c657f848..c78676d62 100644 --- a/crates/goose-provider-types/src/formats/openai.rs +++ b/crates/goose-provider-types/src/formats/openai.rs @@ -1457,23 +1457,26 @@ where }; let content = if output_token_limit_reached { - MessageContentBlock::tool_request_with_metadata( + MessageContentBlock::tool_request_with_provider_index( id.clone(), Err(output_token_limit_tool_error(function_name, id)), metadata.as_ref(), + index, ) } else if arguments.is_empty() { - MessageContentBlock::tool_request_with_metadata( + MessageContentBlock::tool_request_with_provider_index( id.clone(), Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(json!({})))), metadata.as_ref(), + index, ) } else { match parse_tool_arguments(arguments) { - Some(params) if params.is_object() => MessageContentBlock::tool_request_with_metadata( + Some(params) if params.is_object() => MessageContentBlock::tool_request_with_provider_index( id.clone(), Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))), metadata.as_ref(), + index, ), // Valid JSON but NOT an object (a bare array/string/number). // Surface a tool error so the model retries instead of @@ -1488,7 +1491,7 @@ where )), data: None, }; - MessageContentBlock::tool_request_with_metadata(id.clone(), Err(error), metadata.as_ref()) + MessageContentBlock::tool_request_with_provider_index(id.clone(), Err(error), metadata.as_ref(), index) } None => { let message_text = truncation_error_message(arguments) @@ -1500,7 +1503,7 @@ where message: Cow::from(message_text), data: None, }; - MessageContentBlock::tool_request_with_metadata(id.clone(), Err(error), metadata.as_ref()) + MessageContentBlock::tool_request_with_provider_index(id.clone(), Err(error), metadata.as_ref(), index) } } }; @@ -3315,6 +3318,66 @@ mod tests { assert_eq!(usage.usage.total_tokens, Some(expected_total)); } + async fn collect_streamed_tool_requests( + response_lines: &str, + ) -> Vec { + let lines: Vec = response_lines.lines().map(|s| s.to_string()).collect(); + let response_stream = tokio_stream::iter(lines.into_iter().map(Ok)); + let messages = response_to_streaming_message(response_stream); + pin!(messages); + + let mut requests = Vec::new(); + while let Some(Ok((message, _usage))) = messages.next().await { + if let Some(msg) = message { + for content in &msg.content { + if let MessageContentBlock::ToolRequest(req) = content { + requests.push(req.clone()); + } + } + } + } + requests + } + + #[tokio::test] + async fn test_streaming_reassembles_interleaved_parallel_tool_calls() { + let response_lines = concat!( + r#"data: {"id":"chatcmpl-par","model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"search","arguments":""}},{"index":1,"id":"call_b","type":"function","function":{"name":"write","arguments":""}}]},"finish_reason":null}]}"#, + "\n", + r#"data: {"id":"chatcmpl-par","model":"test-model","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{\"path\":"}}]},"finish_reason":null}]}"#, + "\n", + r#"data: {"id":"chatcmpl-par","model":"test-model","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"query\":"}}]},"finish_reason":null}]}"#, + "\n", + r#"data: {"id":"chatcmpl-par","model":"test-model","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"/tmp/a.md\"}"}}]},"finish_reason":null}]}"#, + "\n", + r#"data: {"id":"chatcmpl-par","model":"test-model","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"rust\"}"}}]},"finish_reason":"tool_calls"}]}"#, + "\n", + "data: [DONE]", + ); + + let requests = collect_streamed_tool_requests(response_lines).await; + + assert_eq!(requests.len(), 2); + assert_eq!( + requests + .iter() + .map(|r| r.provider_index()) + .collect::>(), + vec![Some(0), Some(1)] + ); + + let search = requests[0].tool_call.as_ref().expect("search args parsed"); + assert_eq!(search.name, "search"); + assert_eq!(search.arguments.as_ref().unwrap()["query"], json!("rust")); + + let write = requests[1].tool_call.as_ref().expect("write args parsed"); + assert_eq!(write.name, "write"); + assert_eq!( + write.arguments.as_ref().unwrap()["path"], + json!("/tmp/a.md") + ); + } + #[tokio::test] async fn test_streaming_marks_length_on_empty_terminal_chunk() -> anyhow::Result<()> { let response_lines = r#" diff --git a/crates/goose-sdk/src/bindings.rs b/crates/goose-sdk/src/bindings.rs index 6ce9ab478..967b7b94c 100644 --- a/crates/goose-sdk/src/bindings.rs +++ b/crates/goose-sdk/src/bindings.rs @@ -535,6 +535,7 @@ pub enum StreamChunk { id: String, name: String, arguments_json: String, + index: Option, #[uniffi(default = None)] provider_metadata_json: Option, }, @@ -1239,25 +1240,32 @@ fn message_to_chunks(message: Message) -> Vec { text: text.text.clone(), }) } - GooseMessageContent::ToolRequest(request) => match request.tool_call { - Ok(tool_call) => Some(StreamChunk::ToolChunk { - id: request.id, - name: tool_call.name.to_string(), - arguments_json: serde_json::to_string(&tool_call.arguments.unwrap_or_default()) + GooseMessageContent::ToolRequest(request) => { + let index = request.provider_index(); + let provider_metadata_json = request + .metadata + .as_ref() + .and_then(|metadata| serde_json::to_string(metadata).ok()); + match request.tool_call { + Ok(tool_call) => Some(StreamChunk::ToolChunk { + index, + id: request.id, + name: tool_call.name.to_string(), + arguments_json: serde_json::to_string( + &tool_call.arguments.unwrap_or_default(), + ) .unwrap_or_else(|_| "{}".to_string()), - provider_metadata_json: request - .metadata - .as_ref() - .and_then(|metadata| serde_json::to_string(metadata).ok()), - }), - Err(error) => Some(StreamChunk::ErrorChunk { - error: GooseStreamError { - kind: GooseStreamErrorKind::Generic, - message: error.to_string(), - retry_after_ms: None, - }, - }), - }, + provider_metadata_json, + }), + Err(error) => Some(StreamChunk::ErrorChunk { + error: GooseStreamError { + kind: GooseStreamErrorKind::Generic, + message: error.to_string(), + retry_after_ms: None, + }, + }), + } + } GooseMessageContent::Thinking(thinking) => Some(StreamChunk::ThinkingChunk { thinking: thinking.thinking, signature: thinking.signature, diff --git a/documentation/docs/gdk/index.md b/documentation/docs/gdk/index.md index decae1647..4dc3583ec 100644 --- a/documentation/docs/gdk/index.md +++ b/documentation/docs/gdk/index.md @@ -223,7 +223,7 @@ resolved when the provider is constructed. | Chunk | Meaning | | --- | --- | | `TextChunk` | Assistant text | -| `ToolChunk` | A tool call request with JSON arguments | +| `ToolChunk` | A tool call request with JSON arguments and the provider's tool-call `index` | | `ThinkingChunk` / `RedactedThinkingChunk` | Reasoning output | | `EndChunk` | Stream finished, carries final token `Usage` | | `ErrorChunk` | Mid-stream failure, carries a `GooseStreamError` | diff --git a/documentation/src/data/gdk-api.json b/documentation/src/data/gdk-api.json index adc65ddeb..f373936e9 100644 --- a/documentation/src/data/gdk-api.json +++ b/documentation/src/data/gdk-api.json @@ -688,6 +688,12 @@ "default": null, "docs": "" }, + { + "name": "index", + "type": "Option", + "default": null, + "docs": "" + }, { "name": "provider_metadata_json", "type": "Option",