tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

fix(gdk): preserve tool-call indices in streaming responses (#11637)

Browse Source
This commit is contained in:
Alex Hancock
2026-08-31 16:26:08 -04:00
committed by GitHub
parent 1d0146631c
commit 114250f8ca
7 changed files with 279 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-provider-types/src/conversation/message.rs
+22
View File
@@ -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<S: Into<String>>(
id: S,
tool_call: ToolResult<CallToolRequestParams>,
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<S: Into<String>>(id: S, tool_result: ToolResult<CallToolResult>) -> Self {
MessageContentBlock::ToolResponse(ToolResponse {
id: id.into(),
crates/goose-provider-types/src/conversation/tool_request.rs
+11 -1
View File
@@ -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<i32> {
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<ToolChainSummary> {
let obj = self
.tool_meta
crates/goose-provider-types/src/formats/anthropic.rs
+145 -24
View File
@@ -901,9 +901,21 @@ where
signature: String,
}
fn block_index(event_data: &Value) -> Option<i32> {
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<String, (String, String)> = std::collections::HashMap::new();
let mut current_tool_id: Option<String> = None;
struct StreamingToolCall {
id: String,
name: String,
arguments: String,
}
let mut accumulated_tool_calls: std::collections::HashMap<i32, StreamingToolCall> = std::collections::HashMap::new();
let mut final_usage: Option<ProviderUsage> = None;
let mut message_id: Option<String> = None;
let mut thinking: Option<ThinkingState> = 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) = &current_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<String> = 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<i32> = 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<crate::conversation::message::ToolRequest> {
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<i32>)> = 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
crates/goose-provider-types/src/formats/openai.rs
+68 -5
View File
@@ -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<crate::conversation::message::ToolRequest> {
let lines: Vec<String> = 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<_>>(),
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#"
crates/goose-sdk/src/bindings.rs
+26 -18
View File
@@ -535,6 +535,7 @@ pub enum StreamChunk {
id: String,
name: String,
arguments_json: String,
index: Option<i32>,
#[uniffi(default = None)]
provider_metadata_json: Option<String>,
},
@@ -1239,25 +1240,32 @@ fn message_to_chunks(message: Message) -> Vec<StreamChunk> {
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,
documentation/docs/gdk/index.md
+1 -1
View File
@@ -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` |
documentation/src/data/gdk-api.json
+6
View File
@@ -688,6 +688,12 @@
"default": null,
"docs": ""
},
{
"name": "index",
"type": "Option<i32>",
"default": null,
"docs": ""
},
{
"name": "provider_metadata_json",
"type": "Option<String>",