From bb539f7d6f950b3d3b17a540c7dc954f3f03c554 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Tue, 4 Aug 2026 22:44:08 +1000 Subject: [PATCH] feat: surface output-token-limit info (#10831) --- crates/goose-cli/src/session/output.rs | 19 +- .../goose-provider-types/src/conversation.rs | 69 +++- .../src/conversation/message.rs | 9 +- .../src/formats/anthropic.rs | 39 ++ .../src/formats/openai.rs | 342 +++++++++++++++++- .../src/formats/openai_responses.rs | 160 ++++++++ crates/goose/src/acp/server.rs | 194 +++++----- crates/goose/src/acp/server/load_session.rs | 276 ++++++-------- crates/goose/src/acp/server/message_meta.rs | 313 ++++++++++++++++ .../src/acp/server/tool_calls/conversion.rs | 68 +++- crates/goose/src/agents/agent.rs | 137 ++++++- crates/goose/src/agents/reply_parts.rs | 6 +- .../sessionNotificationAdapter.test.ts | 94 +++++ ui/desktop/src/acp/adapter/messages.ts | 43 ++- ui/desktop/src/acp/adapter/shared.ts | 6 + ui/desktop/src/components/GooseMessage.tsx | 20 +- ui/desktop/src/types/message.ts | 2 + 17 files changed, 1471 insertions(+), 326 deletions(-) create mode 100644 crates/goose/src/acp/server/message_meta.rs diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index ad1a3bd39..84508eb4e 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -12,7 +12,7 @@ use goose::subprocess::SubprocessExt; use goose::utils::safe_truncate; use goose_providers::conversation::token_usage::Usage; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; -use rmcp::model::{CallToolRequestParams, JsonObject, PromptArgument}; +use rmcp::model::{CallToolRequestParams, JsonObject, PromptArgument, Role}; use serde_json::Value; use std::cell::RefCell; use std::collections::{HashMap, VecDeque}; @@ -26,6 +26,8 @@ use super::streaming_buffer::MarkdownBuffer; pub const DEFAULT_MIN_PRIORITY: f32 = 0.0; pub const DEFAULT_CLI_LIGHT_THEME: &str = "GitHub"; pub const DEFAULT_CLI_DARK_THEME: &str = "zenburn"; +const OUTPUT_TOKEN_LIMIT_WARNING: &str = + "Warning: Response reached the model's output-token limit and may be incomplete."; fn accent(value: T) -> StyledObject { style(value).cyan() @@ -289,6 +291,9 @@ pub fn render_message(message: &Message, debug: bool) { } } + if reached_output_token_limit(&message) { + render_output_token_limit_warning(); + } let _ = std::io::stdout().flush(); } @@ -379,9 +384,21 @@ pub fn render_message_streaming( } } + if reached_output_token_limit(&message) { + flush_markdown_buffer(buffer, theme); + render_output_token_limit_warning(); + } let _ = std::io::stdout().flush(); } +fn reached_output_token_limit(message: &Message) -> bool { + message.role == Role::Assistant && message.metadata.output_token_limit_reached +} + +fn render_output_token_limit_warning() { + println!("\n{}", warning(OUTPUT_TOKEN_LIMIT_WARNING)); +} + fn render_credits_exhausted_notification(notification: &SystemNotificationContent) { hide_thinking(); println!("\n{} {}", warning("warning:").bold(), ¬ification.msg); diff --git a/crates/goose-provider-types/src/conversation.rs b/crates/goose-provider-types/src/conversation.rs index 052080dae..d6be8dff5 100644 --- a/crates/goose-provider-types/src/conversation.rs +++ b/crates/goose-provider-types/src/conversation.rs @@ -48,14 +48,26 @@ impl Conversation { } pub fn push(&mut self, message: Message) { - if message.content.is_empty() && message.metadata.inference.is_some() { - if let Some(existing) = self - .0 - .iter_mut() - .rev() - .find(|m| m.role == message.role && m.is_user_visible()) - { - existing.metadata.inference = message.metadata.inference; + let output_token_limit_reached = message.metadata.output_token_limit_reached; + if message.content.is_empty() + && (message.metadata.inference.is_some() || output_token_limit_reached) + { + if let Some(existing) = self.0.iter_mut().rev().find(|existing| { + existing.role == message.role + && existing.is_user_visible() + && (!output_token_limit_reached + || (message.id.is_some() + && existing.id.as_deref() == message.id.as_deref())) + }) { + if let Some(inference) = message.metadata.inference.clone() { + existing.metadata.inference = Some(inference); + } + existing.metadata.output_token_limit_reached |= output_token_limit_reached; + return; + } + + if output_token_limit_reached { + self.0.push(message.with_visibility(true, false)); } return; } @@ -68,6 +80,7 @@ impl Conversation { if message.metadata.inference.is_some() { last.metadata.inference = message.metadata.inference.clone(); } + last.metadata.output_token_limit_reached |= message.metadata.output_token_limit_reached; match (last.content.last_mut(), message.content.last()) { ( Some(MessageContentBlock::Text(ref mut last)), @@ -683,7 +696,7 @@ pub fn debug_conversation_fix( #[cfg(test)] mod tests { - use crate::conversation::message::{Message, MessageContentBlock}; + use crate::conversation::message::{InferenceMetadata, Message, MessageContentBlock}; use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation}; use rmcp::model::{CallToolRequestParams, Role}; use rmcp::object; @@ -1746,6 +1759,44 @@ mod tests { } } + #[test] + fn test_push_merges_empty_output_token_limit_update_by_id() { + let mut conv = Conversation::empty(); + conv.push(Message::assistant().with_text("first").with_id("turn-1")); + + let inference = InferenceMetadata { + provider: "test-provider".to_string(), + requested_model: "test-model".to_string(), + resolved_model: None, + }; + let mut limited = Message::assistant() + .with_id("turn-1") + .with_inference(inference.clone()); + limited.metadata.output_token_limit_reached = true; + conv.push(limited); + + assert_eq!(conv.messages().len(), 1); + assert!(conv.messages()[0].metadata.output_token_limit_reached); + assert_eq!(conv.messages()[0].metadata.inference, Some(inference)); + } + + #[test] + fn test_push_retains_unmatched_output_token_limit_update_for_user_only() { + let mut conv = Conversation::empty(); + let mut limited = Message::assistant().with_id("turn-1"); + limited.metadata.output_token_limit_reached = true; + + conv.push(limited); + + assert_eq!(conv.messages().len(), 1); + let persisted = &conv.messages()[0]; + assert!(persisted.content.is_empty()); + assert!(persisted.metadata.user_visible); + assert!(!persisted.metadata.agent_visible); + assert!(persisted.metadata.output_token_limit_reached); + assert!(conv.agent_visible_messages().is_empty()); + } + #[test] fn test_push_thinking_adopts_signature_on_closing_delta() { use crate::conversation::message::MessageContentBlock; diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index d6bfda109..4b91f41a5 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -665,6 +665,9 @@ pub struct MessageMetadata { pub agent_visible: bool, #[serde(skip_serializing_if = "Option::is_none")] pub inference: Option, + /// Whether the provider stopped generating because it reached its output token limit + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub output_token_limit_reached: bool, /// Whether this message is a steer injected into an active run. UI-only: /// surfaced as `_meta.goose.steer` so clients can mark the steer boundary /// without matching user-visible text. Never sent to providers. @@ -680,6 +683,7 @@ impl Default for MessageMetadata { user_visible: true, agent_visible: true, inference: None, + output_token_limit_reached: false, steer: false, usage: None, } @@ -1619,15 +1623,17 @@ mod tests { #[test] fn test_message_metadata_serialization() { - let message = Message::user() + let mut message = Message::user() .with_text("Test message") .with_visibility(false, true); + message.metadata.output_token_limit_reached = true; let json_str = serde_json::to_string(&message).unwrap(); let value: Value = serde_json::from_str(&json_str).unwrap(); assert_eq!(value["metadata"]["userVisible"], false); assert_eq!(value["metadata"]["agentVisible"], true); + assert_eq!(value["metadata"]["outputTokenLimitReached"], true); } #[test] @@ -1649,6 +1655,7 @@ mod tests { let message: Message = serde_json::from_str(json_with_metadata).unwrap(); assert!(!message.is_user_visible()); assert!(message.is_agent_visible()); + assert!(!message.metadata.output_token_limit_reached); } #[test] diff --git a/crates/goose-provider-types/src/formats/anthropic.rs b/crates/goose-provider-types/src/formats/anthropic.rs index e114319ed..26303d5eb 100644 --- a/crates/goose-provider-types/src/formats/anthropic.rs +++ b/crates/goose-provider-types/src/formats/anthropic.rs @@ -1161,6 +1161,13 @@ where } } + if stop_reason.as_deref() == Some("max_tokens") { + let mut message = Message::assistant(); + message.id = message_id; + message.metadata.output_token_limit_reached = true; + yield (Some(message), None); + } + if let Some(usage) = final_usage { yield (None, Some(usage)); } @@ -2065,6 +2072,7 @@ mod tests { text: Vec, tool_calls: Vec, tool_errors: Vec, + output_token_limit_message_ids: Vec>, } async fn collect_stream(events: &str) -> StreamedParts { @@ -2072,6 +2080,9 @@ mod tests { for result in collect_stream_results(events).await { if let Ok((Some(msg), _usage)) = result { + if msg.metadata.output_token_limit_reached { + parts.output_token_limit_message_ids.push(msg.id.clone()); + } for c in &msg.content { match c { MessageContentBlock::Thinking(t) => { @@ -2097,6 +2108,30 @@ mod tests { parts } + #[tokio::test] + async fn test_streaming_marks_max_tokens() { + let events = concat!( + r#"data: {"type":"message_start","message":{"id":"msg_limit","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":"text","text":""}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Partial answer"}}"#, + "\n", + r#"data: {"type":"content_block_stop","index":0}"#, + "\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":25}}"#, + "\n", + r#"data: {"type":"message_stop"}"#, + ); + + let parts = collect_stream(events).await; + assert_eq!(parts.text, vec!["Partial answer"]); + assert_eq!( + parts.output_token_limit_message_ids, + vec![Some("msg_limit".to_string())] + ); + } + #[tokio::test] async fn test_streaming_thinking_and_text() { let events = concat!( @@ -2467,6 +2502,10 @@ mod tests { ); let parts = collect_stream(events).await; + assert_eq!( + parts.output_token_limit_message_ids, + vec![Some("msg_t2".to_string())] + ); assert_eq!( parts.tool_errors.len(), 1, diff --git a/crates/goose-provider-types/src/formats/openai.rs b/crates/goose-provider-types/src/formats/openai.rs index ba55b552d..70571b81f 100644 --- a/crates/goose-provider-types/src/formats/openai.rs +++ b/crates/goose-provider-types/src/formats/openai.rs @@ -49,6 +49,16 @@ fn describe_json_value(value: &Value) -> &'static str { } } +fn output_token_limit_tool_error(function_name: &str, id: &str) -> ErrorData { + ErrorData { + code: ErrorCode::INVALID_PARAMS, + message: Cow::from(format!( + "Tool arguments for {function_name} (id {id}) were truncated because the model reached its output token limit" + )), + data: None, + } +} + pub fn is_reserved_request_param_key(key: &str) -> bool { matches!(key, "messages" | "model" | "stream" | "stream_options") } @@ -671,6 +681,11 @@ pub fn format_tools(tools: &[Tool]) -> anyhow::Result> { /// Convert OpenAI's API response to internal Message format pub fn response_to_message(response: &Value) -> anyhow::Result { + let output_token_limit_reached = response + .pointer("/choices/0/finish_reason") + .and_then(Value::as_str) + == Some("length"); + let Some(original) = response .get("choices") .and_then(|c| c.get(0)) @@ -751,6 +766,16 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { }) .filter(|m: &serde_json::Map| !m.is_empty()); + if output_token_limit_reached { + let error = output_token_limit_tool_error(&function_name, &id); + content.push(MessageContentBlock::tool_request_with_metadata( + id, + Err(error), + metadata.as_ref(), + )); + continue; + } + if function_name.is_empty() { let error = ErrorData { code: ErrorCode::INVALID_REQUEST, @@ -815,11 +840,9 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { } } - Ok(Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - content, - )) + let mut message = Message::new(Role::Assistant, chrono::Utc::now().timestamp(), content); + message.metadata.output_token_limit_reached = output_token_limit_reached; + Ok(message) } pub fn get_usage(usage: &Value) -> Usage { @@ -1063,6 +1086,15 @@ fn parse_streaming_chunk(line: &str) -> Result { }) } +fn output_token_limit_marker(id: Option) -> Message { + let mut message = Message::assistant(); + if let Some(id) = id { + message = message.with_id(id); + } + message.metadata.output_token_limit_reached = true; + message +} + pub fn response_to_streaming_message( mut stream: S, ) -> impl Stream, Option)>> + 'static @@ -1083,6 +1115,9 @@ where // reasoning_content in a later chunk would produce duplicated reasoning. let mut pending_inline_thinking = String::new(); let mut last_seen_model: Option = None; + let mut last_response_id: Option = None; + let mut output_token_limit_reached = false; + let mut output_token_limit_metadata_emitted = false; 'outer: while let Some(response) = stream.next().await { let response_str = response?; @@ -1102,6 +1137,9 @@ where if let Some(model) = &chunk.model { last_seen_model = Some(model.clone()); } + if let Some(id) = &chunk.id { + last_response_id = Some(id.clone()); + } if !chunk.choices.is_empty() { if let Some(details) = &chunk.choices[0].delta.reasoning_details { @@ -1117,6 +1155,11 @@ where } let mut usage = extract_usage_with_output_tokens(&chunk, last_seen_model.as_deref()); + output_token_limit_reached |= chunk + .choices + .first() + .and_then(|choice| choice.finish_reason.as_deref()) + == Some("length"); if chunk.choices.is_empty() { yield (None, usage) @@ -1132,7 +1175,10 @@ where } } - let is_complete = chunk.choices[0].finish_reason == Some("tool_calls".to_string()); + let is_complete = matches!( + chunk.choices[0].finish_reason.as_deref(), + Some("tool_calls" | "length") + ); if !is_complete { let mut done = false; @@ -1148,12 +1194,19 @@ where if let Some(model) = &tool_chunk.model { last_seen_model = Some(model.clone()); } + if let Some(id) = &tool_chunk.id { + last_response_id = Some(id.clone()); + } if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk, last_seen_model.as_deref()) { usage = Some(chunk_usage); } if !tool_chunk.choices.is_empty() { + output_token_limit_reached |= + tool_chunk.choices[0].finish_reason.as_deref() + == Some("length"); + if let Some(details) = &tool_chunk.choices[0].delta.reasoning_details { accumulated_reasoning.extend(details.iter().cloned()); } @@ -1262,7 +1315,13 @@ where extra_fields.as_ref().filter(|m| !m.is_empty()).cloned() }; - let content = if arguments.is_empty() { + let content = if output_token_limit_reached { + MessageContentBlock::tool_request_with_metadata( + id.clone(), + Err(output_token_limit_tool_error(function_name, id)), + metadata.as_ref(), + ) + } else if arguments.is_empty() { MessageContentBlock::tool_request_with_metadata( id.clone(), Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(json!({})))), @@ -1319,6 +1378,8 @@ where if let Some(id) = chunk.id { msg = msg.with_id(id); } + msg.metadata.output_token_limit_reached = output_token_limit_reached; + output_token_limit_metadata_emitted |= output_token_limit_reached; yield ( Some(msg), @@ -1397,14 +1458,23 @@ where content.push(MessageContentBlock::thinking(trailing_thinking, "")); } - yield ( - Some(Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - content, - )), - None, - ) + let mut message = Message::new( + Role::Assistant, + chrono::Utc::now().timestamp(), + content, + ); + if let Some(id) = last_response_id.clone() { + message = message.with_id(id); + } + message.metadata.output_token_limit_reached = + output_token_limit_reached && !output_token_limit_metadata_emitted; + output_token_limit_metadata_emitted |= message.metadata.output_token_limit_reached; + + yield (Some(message), None) + } + + if output_token_limit_reached && !output_token_limit_metadata_emitted { + yield (Some(output_token_limit_marker(last_response_id)), None) } } } @@ -2174,6 +2244,76 @@ mod tests { Ok(()) } + #[test] + fn test_response_to_message_marks_length_finish_reason() -> anyhow::Result<()> { + let response = json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": "Partial answer" + }, + "finish_reason": "length" + }] + }); + + let message = response_to_message(&response)?; + + assert_eq!(message.as_concat_text(), "Partial answer"); + assert!(message.metadata.output_token_limit_reached); + + Ok(()) + } + + #[test] + fn test_response_to_message_rejects_length_terminated_tool_calls() -> anyhow::Result<()> { + let response = json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_empty", + "type": "function", + "function": { + "name": "empty_tool", + "arguments": "" + } + }, + { + "id": "call_valid", + "type": "function", + "function": { + "name": "valid_tool", + "arguments": "{\"value\":true}" + } + } + ] + }, + "finish_reason": "length" + }] + }); + + let message = response_to_message(&response)?; + + assert!(message.metadata.output_token_limit_reached); + assert_eq!(message.content.len(), 2); + for (content, expected_id) in message.content.iter().zip(["call_empty", "call_valid"]) { + let MessageContentBlock::ToolRequest(request) = content else { + panic!("expected tool request"); + }; + assert_eq!(request.id, expected_id); + let error = request + .tool_call + .as_ref() + .expect_err("length-terminated tool call must not be executable"); + assert_eq!(error.code, ErrorCode::INVALID_PARAMS); + assert!(error.message.contains("output token limit")); + } + + Ok(()) + } + #[test] fn test_response_to_message_valid_toolrequest() -> anyhow::Result<()> { let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; @@ -2756,6 +2896,8 @@ mod tests { usage: Option, tool_calls: Vec, has_text_content: bool, + text: String, + output_token_limit_message_ids: Vec>, } async fn run_streaming_test(response_lines: &str) -> anyhow::Result { @@ -2769,6 +2911,8 @@ mod tests { usage: None, tool_calls: Vec::new(), has_text_content: false, + text: String::new(), + output_token_limit_message_ids: Vec::new(), }; while let Some(Ok((message, usage))) = messages.next().await { @@ -2777,6 +2921,9 @@ mod tests { result.usage = Some(u); } if let Some(msg) = message { + if msg.metadata.output_token_limit_reached { + result.output_token_limit_message_ids.push(msg.id.clone()); + } for content in &msg.content { match content { MessageContentBlock::ToolRequest(req) => { @@ -2786,6 +2933,7 @@ mod tests { } MessageContentBlock::Text(text) if !text.text.is_empty() => { result.has_text_content = true; + result.text.push_str(&text.text); } _ => {} } @@ -2814,6 +2962,26 @@ mod tests { assert_eq!(usage.usage.total_tokens, Some(expected_total)); } + #[tokio::test] + async fn test_streaming_marks_length_on_empty_terminal_chunk() -> anyhow::Result<()> { + let response_lines = r#" +data: {"id":"chatcmpl-limit","model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"Partial answer"},"finish_reason":null}]} +data: {"id":"chatcmpl-limit","model":"test-model","choices":[{"index":0,"delta":{},"finish_reason":"length"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}} +data: [DONE] +"#; + + let result = run_streaming_test(response_lines).await?; + + assert_eq!(result.text, "Partial answer"); + assert_eq!( + result.output_token_limit_message_ids, + vec![Some("chatcmpl-limit".to_string())] + ); + assert_usage_yielded_once(&result, 10, 5, 15); + + Ok(()) + } + #[test] fn test_get_usage_preserves_provider_totals_with_cache_fields() { let usage = get_usage(&json!({ @@ -3134,7 +3302,8 @@ data: [DONE]"#; async fn test_streaming_non_object_arguments_does_not_panic() -> anyhow::Result<()> { // Streamed tool call whose arguments are valid JSON but NOT an object. // Must yield an INVALID_PARAMS tool error, not panic via rmcp `object()`. - let response_lines = r#"data: {"model":"test-model","choices":[{"delta":{"role":"assistant","tool_calls":[{"id":"call_bad","function":{"name":"test_tool","arguments":"[1, 2, 3]"},"type":"function","index":0}]},"index":0,"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"object":"chat.completion.chunk","id":"test-id","created":1234567890} + let response_lines = r#"data: {"model":"test-model","choices":[{"delta":{"role":"assistant","tool_calls":[{"id":"call_bad","function":{"name":"test_tool","arguments":"[1, "},"type":"function","index":0}]},"index":0,"finish_reason":null}],"object":"chat.completion.chunk","id":"test-id","created":1234567890} +data: {"model":"test-model","choices":[{"delta":{"tool_calls":[{"function":{"arguments":"2, 3]"},"index":0}]},"index":0,"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"object":"chat.completion.chunk","id":"test-id","created":1234567890} data: [DONE]"#; let response_stream = @@ -3142,9 +3311,17 @@ data: [DONE]"#; let messages = response_to_streaming_message(response_stream); pin!(messages); - while let Some(Ok((message, _usage))) = messages.next().await { + let mut found_tool_error = false; + let mut usage_count = 0; + while let Some(result) = messages.next().await { + let (message, usage) = result?; + if usage.is_some() { + usage_count += 1; + } if let Some(msg) = message { if let MessageContentBlock::ToolRequest(request) = &msg.content[0] { + assert!(!msg.metadata.output_token_limit_reached); + assert_eq!(msg.id.as_deref(), Some("test-id")); match &request.tool_call { Err(ErrorData { code: ErrorCode::INVALID_PARAMS, @@ -3156,14 +3333,60 @@ data: [DONE]"#; m.contains("test_tool"), "error must name the original tool so the model can retry it: {m}" ); - return Ok(()); + found_tool_error = true; } _ => panic!("expected INVALID_PARAMS for non-object streamed args"), } } } } - panic!("expected a tool request message"); + + assert!(found_tool_error, "expected a tool request message"); + assert_eq!(usage_count, 1); + Ok(()) + } + + #[tokio::test] + async fn test_streaming_length_rejects_tool_calls_regardless_of_arguments() -> anyhow::Result<()> + { + let response_lines = r#"data: {"model":"test-model","choices":[{"delta":{"role":"assistant","tool_calls":[{"id":"call_empty","function":{"name":"empty_tool","arguments":""},"type":"function","index":0},{"id":"call_valid","function":{"name":"valid_tool","arguments":"{\"value\":"},"type":"function","index":1}]},"index":0,"finish_reason":null}],"object":"chat.completion.chunk","id":"test-id","created":1234567890} +data: {"model":"test-model","choices":[{"delta":{"tool_calls":[{"function":{"arguments":""},"index":0},{"function":{"arguments":"true}"},"index":1}]},"index":0,"finish_reason":"length"}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"object":"chat.completion.chunk","id":"test-id","created":1234567890} +data: [DONE]"#; + + let response_stream = + tokio_stream::iter(response_lines.lines().map(|line| Ok(line.to_string()))); + let messages = response_to_streaming_message(response_stream); + pin!(messages); + + let mut tool_error_ids = Vec::new(); + let mut usage_count = 0; + while let Some(result) = messages.next().await { + let (message, usage) = result?; + if usage.is_some() { + usage_count += 1; + } + if let Some(msg) = message { + assert!(msg.metadata.output_token_limit_reached); + assert_eq!(msg.id.as_deref(), Some("test-id")); + for content in msg.content { + if let MessageContentBlock::ToolRequest(request) = content { + let request_id = request.id.clone(); + match request.tool_call { + Err(error) => { + assert_eq!(error.code, ErrorCode::INVALID_PARAMS); + assert!(error.message.contains("output token limit")); + tool_error_ids.push(request_id); + } + Ok(_) => panic!("length-terminated tool call must not be executable"), + } + } + } + } + } + + assert_eq!(tool_error_ids, vec!["call_empty", "call_valid"]); + assert_eq!(usage_count, 1); + Ok(()) } #[tokio::test] @@ -3205,6 +3428,87 @@ data: [DONE]"#; Ok(()) } + #[tokio::test] + async fn test_streaming_truncated_inline_think_preserves_output_token_limit( + ) -> anyhow::Result<()> { + let response_lines = concat!( + "data: {\"id\":\"chunk-1\",\"model\":\"test-model\",\"choices\":[{\"delta\":{\"content\":\"unfinished reasoning\"},\"index\":0,\"finish_reason\":null}]}\n", + "data: {\"id\":\"chunk-1\",\"choices\":[{\"delta\":{},\"index\":0,\"finish_reason\":\"length\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n", + "data: [DONE]\n" + ); + + let response_stream = + tokio_stream::iter(response_lines.lines().map(|line| Ok(line.to_string()))); + let mut messages = std::pin::pin!(response_to_streaming_message(response_stream)); + let mut streamed_messages = Vec::new(); + let mut usage_count = 0; + + while let Some(result) = messages.next().await { + let (message, usage) = result?; + if usage.is_some() { + usage_count += 1; + } + if let Some(message) = message { + streamed_messages.push(message); + } + } + + assert_eq!(usage_count, 1); + assert_eq!( + streamed_messages + .iter() + .filter(|message| message.metadata.output_token_limit_reached) + .count(), + 1 + ); + let trailing_message = streamed_messages + .last() + .expect("expected trailing thinking"); + assert!(trailing_message.metadata.output_token_limit_reached); + assert!(matches!( + trailing_message.content.as_slice(), + [MessageContentBlock::Thinking(thinking)] + if thinking.thinking == "unfinished reasoning" + )); + + Ok(()) + } + + #[tokio::test] + async fn test_streaming_partial_think_tag_emits_one_output_limit_marker() -> anyhow::Result<()> + { + let response_lines = concat!( + "data: {\"id\":\"chunk-1\",\"model\":\"test-model\",\"choices\":[{\"delta\":{\"content\":\" = streamed_messages + .iter() + .filter(|message| message.metadata.output_token_limit_reached) + .collect(); + assert_eq!(marked_messages.len(), 1); + assert_eq!(marked_messages[0].id.as_deref(), Some("chunk-1")); + assert_eq!(marked_messages[0].as_concat_text(), " anyhow::Result<()> { diff --git a/crates/goose-provider-types/src/formats/openai_responses.rs b/crates/goose-provider-types/src/formats/openai_responses.rs index a4facaf96..685b34097 100644 --- a/crates/goose-provider-types/src/formats/openai_responses.rs +++ b/crates/goose-provider-types/src/formats/openai_responses.rs @@ -107,6 +107,26 @@ pub struct ResponseReasoningInfo { pub summary: Option, } +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ResponseIncompleteDetails { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reason: Option, +} + +fn is_output_token_limit_incomplete_reason(reason: &str) -> bool { + matches!(reason, "max_output_tokens" | "max_tokens") +} + +fn response_reached_output_token_limit( + status: &str, + incomplete_details: Option<&ResponseIncompleteDetails>, +) -> bool { + status == "incomplete" + && incomplete_details + .and_then(|details| details.reason.as_deref()) + .is_some_and(is_output_token_limit_incomplete_reason) +} + #[derive(Debug, Serialize, Deserialize)] pub struct ResponseUsage { pub input_tokens: i32, @@ -207,6 +227,11 @@ pub enum ResponsesStreamEvent { sequence_number: i32, response: ResponseMetadata, }, + #[serde(rename = "response.incomplete")] + ResponseIncomplete { + sequence_number: i32, + response: ResponseMetadata, + }, #[serde(rename = "response.failed")] ResponseFailed { sequence_number: i32, error: Value }, #[serde(rename = "response.function_call_arguments.delta")] @@ -262,6 +287,7 @@ fn is_known_responses_stream_event_type(event_type: &str) -> bool { | "response.content_part.done" | "response.output_text.done" | "response.completed" + | "response.incomplete" | "response.failed" | "response.function_call_arguments.delta" | "response.function_call_arguments.done" @@ -310,6 +336,8 @@ pub struct ResponseMetadata { pub usage: Option, #[serde(skip_serializing_if = "Option::is_none")] pub reasoning: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub incomplete_details: Option, } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -866,6 +894,15 @@ fn process_streaming_output_items( Ok(content) } +fn output_token_limit_marker(id: Option) -> Message { + let mut message = Message::assistant(); + if let Some(id) = id { + message = message.with_id(id); + } + message.metadata.output_token_limit_reached = true; + message +} + pub fn responses_api_to_streaming_message( mut stream: S, ) -> impl Stream, Option)>> + 'static @@ -881,6 +918,7 @@ where let mut final_usage: Option = None; let mut output_items: Vec = Vec::new(); let mut is_text_response = false; + let mut output_token_limit_reached = false; 'outer: while let Some(response) = stream.next().await { let response_str = response?; @@ -968,6 +1006,35 @@ where break 'outer; } + ResponsesStreamEvent::ResponseIncomplete { response, .. } => { + let model = model_name.as_ref().unwrap_or(&response.model); + let usage = response.usage.as_ref().map_or_else( + Usage::default, + ResponseUsage::to_usage, + ); + final_usage = Some(ProviderUsage::new(model.clone(), usage)); + response_id = Some(response.id.clone()); + output_token_limit_reached = response_reached_output_token_limit( + &response.status, + response.incomplete_details.as_ref(), + ); + + if !response.output.is_empty() { + output_items = response + .output + .into_iter() + .filter(|item| match item { + ResponseOutputItemInfo::FunctionCall { status, .. } => { + status.as_deref() == Some("completed") + } + _ => true, + }) + .collect(); + } + + break 'outer; + } + ResponsesStreamEvent::FunctionCallArgumentsDelta { .. } => { // Function call arguments are being streamed, but we'll get the complete // arguments in the OutputItemDone event, so we can ignore deltas for now @@ -1028,7 +1095,10 @@ where if let Some(id) = response_id { message = message.with_id(id); } + message.metadata.output_token_limit_reached = output_token_limit_reached; yield (Some(message), final_usage); + } else if output_token_limit_reached { + yield (Some(output_token_limit_marker(response_id)), final_usage); } else if let Some(usage) = final_usage { yield (None, Some(usage)); } @@ -1129,6 +1199,96 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_responses_stream_marks_output_token_limit_when_incomplete() -> anyhow::Result<()> + { + let lines = vec![ + r#"data: {"type":"response.created","sequence_number":1,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"in_progress","model":"gpt-5.2-pro","output":[]}}"#.to_string(), + r#"data: {"type":"response.output_text.delta","sequence_number":2,"item_id":"msg_1","output_index":0,"content_index":0,"delta":"Partial response"}"#.to_string(), + r#"data: {"type":"response.incomplete","sequence_number":3,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"incomplete","model":"gpt-5.2-pro","output":[],"incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#.to_string(), + "data: [DONE]".to_string(), + ]; + + let response_stream = tokio_stream::iter(lines.into_iter().map(Ok)); + let messages = responses_api_to_streaming_message(response_stream); + futures::pin_mut!(messages); + + let mut text_parts = Vec::new(); + let mut output_token_limit_markers = Vec::new(); + let mut usage: Option = None; + + while let Some(item) = messages.next().await { + let (message, maybe_usage) = item?; + if let Some(msg) = message { + if msg.metadata.output_token_limit_reached { + output_token_limit_markers.push((msg.id.clone(), msg.content.is_empty())); + } + for content in msg.content { + if let MessageContentBlock::Text(text) = content { + text_parts.push(text.text); + } + } + } + if let Some(final_usage) = maybe_usage { + usage = Some(final_usage); + } + } + + assert_eq!(text_parts.concat(), "Partial response"); + assert_eq!( + output_token_limit_markers, + vec![(Some("resp_1".to_string()), true)] + ); + let usage = usage.expect("usage should be present when the response is incomplete"); + assert_eq!(usage.model, "gpt-5.2-pro"); + assert_eq!(usage.usage.input_tokens, Some(10)); + assert_eq!(usage.usage.output_tokens, Some(5)); + assert_eq!(usage.usage.total_tokens, Some(15)); + + Ok(()) + } + + #[tokio::test] + async fn test_responses_stream_drops_incomplete_function_calls() -> anyhow::Result<()> { + let lines = vec![ + r#"data: {"type":"response.created","sequence_number":1,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"in_progress","model":"gpt-5.2-pro","output":[]}}"#.to_string(), + r#"data: {"type":"response.incomplete","sequence_number":2,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"incomplete","model":"gpt-5.2-pro","output":[{"type":"function_call","id":"fc_complete","status":"completed","call_id":"call_complete","name":"read_file","arguments":"{\"path\":\"README.md\"}"},{"type":"function_call","id":"fc_partial","status":"in_progress","call_id":"call_partial","name":"shell","arguments":"{\"command\":\"echo"}],"incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#.to_string(), + "data: [DONE]".to_string(), + ]; + + let response_stream = tokio_stream::iter(lines.into_iter().map(Ok)); + let messages = responses_api_to_streaming_message(response_stream); + futures::pin_mut!(messages); + + let mut tool_request_ids = Vec::new(); + let mut output_token_limit_reached = false; + let mut usage: Option = None; + + while let Some(item) = messages.next().await { + let (message, maybe_usage) = item?; + if let Some(msg) = message { + output_token_limit_reached |= msg.metadata.output_token_limit_reached; + for content in msg.content { + if let MessageContentBlock::ToolRequest(request) = content { + tool_request_ids.push(request.id); + } + } + } + if let Some(final_usage) = maybe_usage { + usage = Some(final_usage); + } + } + + assert_eq!(tool_request_ids, vec!["call_complete"]); + assert!(output_token_limit_reached); + let usage = usage.expect("usage should be present when the response is incomplete"); + assert_eq!(usage.usage.input_tokens, Some(10)); + assert_eq!(usage.usage.output_tokens, Some(5)); + assert_eq!(usage.usage.total_tokens, Some(15)); + + Ok(()) + } + #[tokio::test] async fn test_responses_stream_allows_message_output_without_id_status() -> anyhow::Result<()> { let lines = vec![ diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 9a0f69356..c3554b601 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -41,17 +41,17 @@ use crate::utils::sanitize_unicode_tags; use agent_client_protocol::schema::v1::{ AgentCapabilities, Annotations, AuthMethod, AuthMethodAgent, AuthenticateRequest, AuthenticateResponse, CancelNotification, CloseSessionRequest, CloseSessionResponse, - ConfigOptionUpdate, ContentBlock, ContentChunk, Cost, CurrentModeUpdate, - EmbeddedResourceResource, FileSystemCapabilities, ForkSessionRequest, ForkSessionResponse, - ImageContent, Implementation, InitializeRequest, InitializeResponse, ListSessionsRequest, - ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, - MessageId, Meta, NewSessionRequest, NewSessionResponse, PermissionOption, PermissionOptionKind, - PromptCapabilities, PromptRequest, PromptResponse, RequestPermissionOutcome, - RequestPermissionRequest, ResourceLink, SessionCapabilities, SessionCloseCapabilities, - SessionConfigOption, SessionId, SessionInfoUpdate, SessionListCapabilities, - SessionNotification, SessionUpdate, SetSessionConfigOptionRequest, - SetSessionConfigOptionResponse, SetSessionModeRequest, SetSessionModeResponse, StopReason, - TextContent, ToolCallId, ToolCallUpdate, Usage, UsageUpdate, + ConfigOptionUpdate, ContentBlock, Cost, CurrentModeUpdate, EmbeddedResourceResource, + FileSystemCapabilities, ForkSessionRequest, ForkSessionResponse, ImageContent, Implementation, + InitializeRequest, InitializeResponse, ListSessionsRequest, ListSessionsResponse, + LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, Meta, NewSessionRequest, + NewSessionResponse, PermissionOption, PermissionOptionKind, PromptCapabilities, PromptRequest, + PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, ResourceLink, + SessionCapabilities, SessionCloseCapabilities, SessionConfigOption, SessionId, + SessionInfoUpdate, SessionListCapabilities, SessionNotification, SessionUpdate, + SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, SetSessionModeRequest, + SetSessionModeResponse, StopReason, TextContent, ToolCallId, ToolCallUpdate, Usage, + UsageUpdate, }; use agent_client_protocol::util::MatchDispatchFrom; use agent_client_protocol::{ @@ -75,9 +75,12 @@ use tracing::{debug, error, info, warn}; use url::Url; use uuid::Uuid; +use self::message_meta::{ + content_chunk_for_message, message_meta_without_steer, populate_output_token_limit_content, +}; use self::tool_calls::chain::{breaks_consecutive_tool_calls, ReadyToolChain, ToolChainTracker}; use self::tool_calls::conversion::{ - build_initial_tool_call, build_permission_tool_call_update, + build_initial_tool_call_with_message_meta, build_permission_tool_call_update, tool_call_update_fields_from_response, trusted_update_meta, }; use self::tool_calls::enrichment::{spawn_chain_summary_enrichment, spawn_tool_title_enrichment}; @@ -98,6 +101,7 @@ mod list_sessions; mod load_session; mod local_inference; mod manage_sessions; +mod message_meta; mod new_session; mod onboarding; mod prompts; @@ -513,6 +517,22 @@ fn build_prompt_usage(session: &Session) -> Option { Some(Usage::new(total, input, output)) } +fn prompt_stop_reason(was_cancelled: bool, output_token_limit_reached: bool) -> StopReason { + if was_cancelled { + StopReason::Cancelled + } else if output_token_limit_reached { + StopReason::MaxTokens + } else { + StopReason::EndTurn + } +} + +fn update_output_token_limit_reached(output_token_limit_reached: &mut bool, message: &Message) { + if message.role == Role::Assistant { + *output_token_limit_reached = message.metadata.output_token_limit_reached; + } +} + pub(super) struct UsageUpdates { pub(super) custom: GooseSessionNotification, pub(super) standard: UsageUpdate, @@ -1007,27 +1027,22 @@ impl GooseAcpAgent { message } - #[allow(clippy::too_many_arguments)] async fn handle_message_content( &self, content_item: &MessageContent, + message: &Message, session_id: &SessionId, - session_id_str: &str, - message_id: Option<&str>, - message_created: i64, - role: &Role, - steer: bool, agent: &Arc, tool_requests: &HashMap, cx: &ConnectionTo, ) -> Result<(), agent_client_protocol::Error> { + let role = &message.role; + match content_item { MessageContent::Text(text) => { let chunk = content_chunk_for_message( + message, ContentBlock::Text(TextContent::new(text.text.clone())), - message_id, - message_created, - steer, ); let update = match role { Role::User => SessionUpdate::UserMessageChunk(chunk), @@ -1036,7 +1051,7 @@ impl GooseAcpAgent { cx.send_notification(SessionNotification::new(session_id.clone(), update))?; } MessageContent::ToolRequest(tool_request) => { - self.handle_tool_request(tool_request, session_id, session_id_str, agent, cx) + self.handle_tool_request(tool_request, message, session_id, agent, cx) .await?; } MessageContent::ToolResponse(tool_response) => { @@ -1052,10 +1067,8 @@ impl GooseAcpAgent { cx.send_notification(SessionNotification::new( session_id.clone(), SessionUpdate::AgentThoughtChunk(content_chunk_for_message( + message, ContentBlock::Text(TextContent::new(thinking.thinking.clone())), - message_id, - message_created, - steer, )), ))?; } @@ -1078,16 +1091,16 @@ impl GooseAcpAgent { } ActionRequiredData::Elicitation { id, - message, + message: elicitation_message, requested_schema, } => { self.handle_form_elicitation( cx, session_id, id, - message, + elicitation_message, requested_schema, - message_update_meta(message_id, message_created, false), + message_meta_without_steer(message), ) .await?; } @@ -1112,12 +1125,7 @@ impl GooseAcpAgent { ), ); } - let chunk = content_chunk_for_message( - ContentBlock::Image(image_content), - message_id, - message_created, - steer, - ); + let chunk = content_chunk_for_message(message, ContentBlock::Image(image_content)); let update = match role { Role::User => SessionUpdate::UserMessageChunk(chunk), Role::Assistant => SessionUpdate::AgentMessageChunk(chunk), @@ -1161,14 +1169,17 @@ impl GooseAcpAgent { async fn handle_tool_request( &self, tool_request: &ToolRequest, + message: &Message, session_id: &SessionId, - session_id_for_persist: &str, agent: &Arc, cx: &ConnectionTo, ) -> Result<(), agent_client_protocol::Error> { let client_requests_label_enrichment = self.requests_tool_call_label_enrichment(); - let initial_tool_call = - build_initial_tool_call(tool_request, client_requests_label_enrichment); + let initial_tool_call = build_initial_tool_call_with_message_meta( + tool_request, + message, + client_requests_label_enrichment, + ); let tool_call_notifier = ToolCallNotifier::new(cx, session_id); tool_call_notifier.send_initial(initial_tool_call)?; @@ -1181,7 +1192,7 @@ impl GooseAcpAgent { agent, tool_call_notifier, &self.session_manager, - session_id_for_persist, + session_id.0.as_ref(), tool_request, ); } @@ -1412,35 +1423,6 @@ fn message_usage_update( } } -fn message_update_meta(message_id: Option<&str>, created: i64, steer: bool) -> Meta { - let mut goose = serde_json::Map::new(); - goose.insert("created".to_string(), serde_json::json!(created)); - if let Some(id) = message_id { - goose.insert("messageId".to_string(), serde_json::json!(id)); - } - if steer { - goose.insert("steer".to_string(), serde_json::json!(true)); - } - - let mut meta = serde_json::Map::new(); - meta.insert("goose".to_string(), serde_json::Value::Object(goose)); - meta -} - -fn content_chunk_for_message( - content: ContentBlock, - message_id: Option<&str>, - created: i64, - steer: bool, -) -> ContentChunk { - let mut chunk = - ContentChunk::new(content).meta(message_update_meta(message_id, created, steer)); - if let Some(message_id) = message_id { - chunk = chunk.message_id(MessageId::new(message_id)); - } - chunk -} - impl GooseAcpAgent { async fn on_initialize( &self, @@ -1805,6 +1787,7 @@ impl GooseAcpAgent { }; let mut was_cancelled = false; + let mut output_token_limit_reached = false; let mut tool_requests = HashMap::new(); let mut chain_tracker = ToolChainTracker::default(); let mut stream_error = None; @@ -1816,9 +1799,8 @@ impl GooseAcpAgent { } match event { - Ok(crate::agents::AgentEvent::Message(message)) => { - // Agent persists messages via session_manager.add_message() internally. - let stored_message_id = message.id.clone(); + Ok(crate::agents::AgentEvent::Message(mut message)) => { + update_output_token_limit_reached(&mut output_token_limit_reached, &message); let sessions = self.sessions.lock().await; if !sessions.contains_key(&session_id) { @@ -1829,6 +1811,7 @@ impl GooseAcpAgent { break; } + populate_output_token_limit_content(&mut message); for content_item in &message.content { if let Some(error) = prompt_error_from_message_content(content_item) { stream_error = Some(error); @@ -1842,12 +1825,8 @@ impl GooseAcpAgent { if let Err(error) = self .handle_message_content( content_item, + &message, &args.session_id, - &session_id, - stored_message_id.as_deref(), - message.created, - &message.role, - message.metadata.steer, &agent, &tool_requests, cx, @@ -1876,6 +1855,7 @@ impl GooseAcpAgent { self.spawn_ready_chain_summary(chain, &agent, &args.session_id, cx); } } + if stream_error.is_some() { break; } @@ -1943,11 +1923,7 @@ impl GooseAcpAgent { ))?; } - let stop_reason = if was_cancelled { - StopReason::Cancelled - } else { - StopReason::EndTurn - }; + let stop_reason = prompt_stop_reason(was_cancelled, output_token_limit_reached); let mut response = PromptResponse::new(stop_reason); if let Some(usage) = build_prompt_usage(&session) { @@ -2566,28 +2542,6 @@ print(\"hello, world\") assert_eq!(outcome_to_confirmation(&input), expected); } - #[test] - fn test_message_update_meta_includes_created_and_message_id() { - let meta = message_update_meta(Some("msg_live"), 1_700_000_000, false); - - assert_eq!( - meta.get("goose"), - Some(&serde_json::json!({ - "created": 1_700_000_000, - "messageId": "msg_live", - })), - ); - - let chunk = content_chunk_for_message( - ContentBlock::Text(TextContent::new("hello")), - Some("msg_live"), - 1_700_000_000, - true, - ); - - assert_eq!(chunk.message_id, Some(MessageId::new("msg_live"))); - } - #[test] fn test_credits_exhausted_system_notification_maps_to_prompt_error() { let content = MessageContent::SystemNotification(SystemNotificationContent { @@ -2676,6 +2630,42 @@ print(\"hello, world\") assert!(build_prompt_usage(&session).is_none()); } + #[test_case(false, false, StopReason::EndTurn; "normal completion")] + #[test_case(false, true, StopReason::MaxTokens; "output token limit")] + #[test_case(true, true, StopReason::Cancelled; "cancellation takes precedence")] + fn test_prompt_stop_reason( + was_cancelled: bool, + output_token_limit_reached: bool, + expected: StopReason, + ) { + assert_eq!( + prompt_stop_reason(was_cancelled, output_token_limit_reached), + expected + ); + } + + #[test] + fn test_output_token_limit_state_tracks_latest_assistant_message() { + let mut output_token_limit_reached = false; + let mut marker = Message::assistant(); + marker.metadata.output_token_limit_reached = true; + + update_output_token_limit_reached(&mut output_token_limit_reached, &marker); + assert!(output_token_limit_reached); + + update_output_token_limit_reached( + &mut output_token_limit_reached, + &Message::user().with_text("continue"), + ); + assert!(output_token_limit_reached); + + update_output_token_limit_reached( + &mut output_token_limit_reached, + &Message::assistant().with_text("Complete response"), + ); + assert!(!output_token_limit_reached); + } + #[test] fn test_build_usage_update_clamps_negative_used_to_zero() { let mut session = make_session_with_usage( diff --git a/crates/goose/src/acp/server/load_session.rs b/crates/goose/src/acp/server/load_session.rs index dfd755e06..20403e1be 100644 --- a/crates/goose/src/acp/server/load_session.rs +++ b/crates/goose/src/acp/server/load_session.rs @@ -1,48 +1,14 @@ +use super::message_meta::{ + content_chunk_for_message, merge_message_meta, populate_output_token_limit_content, +}; use super::tool_calls::conversion::{ - build_initial_tool_call, tool_call_update_fields_from_response, trusted_update_meta, + build_initial_tool_call_with_message_meta, tool_call_update_fields_from_response, + trusted_update_meta, }; use super::tool_calls::enrichment::tool_chain_summary; use super::*; -use agent_client_protocol::schema::v1::{MessageId, ToolCall}; - -fn replay_message_meta(message: &Message) -> Meta { - let mut meta = serde_json::Map::new(); - meta.insert( - "goose".to_string(), - serde_json::Value::Object(replay_message_goose_meta(message)), - ); - meta -} - -fn replay_message_goose_meta(message: &Message) -> serde_json::Map { - let mut goose = serde_json::Map::new(); - goose.insert("created".to_string(), serde_json::json!(message.created)); - if let Some(id) = &message.id { - goose.insert("messageId".to_string(), serde_json::json!(id)); - } - if message.metadata.steer { - goose.insert("steer".to_string(), serde_json::json!(true)); - } - goose -} - -fn merge_replay_message_meta(meta: Option, message: &Message) -> Meta { - let replay_goose = replay_message_goose_meta(message); - let mut meta = meta.unwrap_or_default(); - let goose_value = meta - .entry("goose".to_string()) - .or_insert_with(|| serde_json::Value::Object(serde_json::Map::new())); - - if let serde_json::Value::Object(goose) = goose_value { - for (key, value) in replay_goose { - goose.insert(key, value); - } - } else { - *goose_value = serde_json::Value::Object(replay_goose); - } - - meta -} +use crate::conversation::Conversation; +use agent_client_protocol::schema::v1::ToolCall; fn replay_audience_annotations(audience: &[Role]) -> Annotations { Annotations::new().audience( @@ -56,13 +22,27 @@ fn replay_audience_annotations(audience: &[Role]) -> Annotations { ) } +fn messages_for_acp_replay(conversation: &Conversation) -> Vec { + conversation + .messages() + .iter() + .filter(|message| message.is_user_visible()) + .map(Message::user_visible_content) + .map(|mut message| { + populate_output_token_limit_content(&mut message); + message + }) + .filter(|message| !message.content.is_empty()) + .collect() +} + fn send_replay_content_chunk( cx: &ConnectionTo, session_id: &SessionId, message: &Message, content: ContentBlock, ) -> std::result::Result<(), agent_client_protocol::Error> { - let chunk = replay_content_chunk_for_message(message, content); + let chunk = content_chunk_for_message(message, content); let update = match message.role { Role::User => SessionUpdate::UserMessageChunk(chunk), Role::Assistant => SessionUpdate::AgentMessageChunk(chunk), @@ -70,20 +50,16 @@ fn send_replay_content_chunk( cx.send_notification(SessionNotification::new(session_id.clone(), update)) } -fn replay_content_chunk_for_message(message: &Message, content: ContentBlock) -> ContentChunk { - let mut chunk = ContentChunk::new(content).meta(replay_message_meta(message)); - if let Some(message_id) = message.id.as_deref() { - chunk = chunk.message_id(MessageId::new(message_id)); - } - chunk -} - fn build_replayed_tool_call( tool_request: &ToolRequest, + message: &Message, client_requests_tool_call_label_enrichment: bool, ) -> ToolCall { - let mut tool_call = - build_initial_tool_call(tool_request, client_requests_tool_call_label_enrichment); + let mut tool_call = build_initial_tool_call_with_message_meta( + tool_request, + message, + client_requests_tool_call_label_enrichment, + ); if !client_requests_tool_call_label_enrichment { return tool_call; @@ -120,7 +96,7 @@ fn replay_conversation_to_client( let messages = session .conversation .as_ref() - .map(|c| c.user_visible_messages()) + .map(messages_for_acp_replay) .unwrap_or_default(); let mut replay_tool_requests = HashMap::new(); @@ -156,12 +132,11 @@ fn replay_conversation_to_client( MessageContent::ToolRequest(tool_request) => { replay_tool_requests.insert(tool_request.id.clone(), tool_request.clone()); - let mut tool_call = build_replayed_tool_call( + let tool_call = build_replayed_tool_call( tool_request, + message, client_requests_tool_call_label_enrichment, ); - let meta = tool_call.meta.take(); - let tool_call = tool_call.meta(merge_replay_message_meta(meta, message)); tool_call_notifier.send_initial(tool_call)?; } @@ -171,19 +146,17 @@ fn replay_conversation_to_client( replay_tool_requests.get(&tool_response.id), true, ); + let meta = trusted_update_meta(tool_response).unwrap_or_default(); let update = ToolCallUpdate::new(ToolCallId::new(tool_response.id.clone()), fields) - .meta(merge_replay_message_meta( - trusted_update_meta(tool_response), - message, - )); + .meta(merge_message_meta(meta, message)); tool_call_notifier.send_update(update)?; } MessageContent::Thinking(thinking) => { cx.send_notification(SessionNotification::new( session_id.clone(), - SessionUpdate::AgentThoughtChunk(replay_content_chunk_for_message( + SessionUpdate::AgentThoughtChunk(content_chunk_for_message( message, ContentBlock::Text(TextContent::new(thinking.thinking.clone())), )), @@ -278,6 +251,49 @@ mod tests { use super::*; use rmcp::model::CallToolRequestParams; + #[test] + fn acp_replay_populates_only_empty_marked_assistant_messages() { + let visible_message = Message::assistant() + .with_text("visible") + .with_id("msg_visible"); + let empty_message = Message::assistant().with_id("msg_empty"); + + let mut marked_message = Message::assistant().with_id("msg_limited"); + marked_message.metadata.output_token_limit_reached = true; + + let mut marked_user_message = Message::user().with_id("msg_user"); + marked_user_message.metadata.output_token_limit_reached = true; + + let mut hidden_marked_message = Message::assistant() + .with_id("msg_hidden") + .with_visibility(false, false); + hidden_marked_message.metadata.output_token_limit_reached = true; + + let conversation = Conversation::new_unvalidated([ + visible_message, + empty_message, + marked_message, + marked_user_message, + hidden_marked_message, + ]); + + let messages = messages_for_acp_replay(&conversation); + + assert_eq!( + messages + .iter() + .filter_map(|message| message.id.as_deref()) + .collect::>(), + vec!["msg_visible", "msg_limited"] + ); + assert_eq!( + messages[1].as_concat_text(), + "Response stopped because the model reached its output-token limit." + ); + assert!(messages[1].metadata.output_token_limit_reached); + assert!(conversation.messages()[2].content.is_empty()); + } + fn persisted_enriched_tool_request() -> ToolRequest { ToolRequest { id: "req_first".to_string(), @@ -295,7 +311,11 @@ mod tests { #[test] fn replay_includes_persisted_enrichment_when_requested() { - let tool_call = build_replayed_tool_call(&persisted_enriched_tool_request(), true); + let mut message = + Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_replay"); + message.metadata.output_token_limit_reached = true; + let tool_call = + build_replayed_tool_call(&persisted_enriched_tool_request(), &message, true); let goose = tool_call .meta .as_ref() @@ -304,20 +324,28 @@ mod tests { assert_eq!(tool_call.title, "applied dark mode polish"); assert_eq!( - goose.get("toolCall"), - Some( - &serde_json::json!({ "toolName": "developer__shell", "extensionName": "developer" }) - ), - ); - assert_eq!( - goose.get("toolChainSummary"), - Some(&serde_json::json!({ "summary": "applied dark mode polish", "count": 3 })), + goose, + &serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_replay", + "outputTokenLimitReached": true, + "toolCall": { + "toolName": "developer__shell", + "extensionName": "developer", + }, + "toolChainSummary": { + "summary": "applied dark mode polish", + "count": 3, + }, + }), ); } #[test] fn replay_omits_persisted_enrichment_when_not_requested() { - let tool_call = build_replayed_tool_call(&persisted_enriched_tool_request(), false); + let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_replay"); + let tool_call = + build_replayed_tool_call(&persisted_enriched_tool_request(), &message, false); let goose = tool_call .meta .as_ref() @@ -326,105 +354,15 @@ mod tests { assert_eq!(tool_call.title, "developer: shell"); assert_eq!( - goose.get("toolCall"), - Some( - &serde_json::json!({ "toolName": "developer__shell", "extensionName": "developer" }) - ), - ); - assert_eq!(goose.get("toolChainSummary"), None); - } - - #[test] - fn merge_replay_message_meta_preserves_existing_goose_meta() { - let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_1"); - let existing = serde_json::from_value(serde_json::json!({ - "goose": { - "mcpApp": { - "resourceUri": "ui://trusted/app", - "extensionName": "weather", - "toolName": "weather__render", + goose, + &serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_replay", + "toolCall": { + "toolName": "developer__shell", + "extensionName": "developer", }, - } - })) - .unwrap(); - - let merged = merge_replay_message_meta(Some(existing), &message); - - assert_eq!( - merged.get("goose"), - Some(&serde_json::json!({ - "created": 1_700_000_000, - "messageId": "msg_1", - "mcpApp": { - "resourceUri": "ui://trusted/app", - "extensionName": "weather", - "toolName": "weather__render", - }, - })), - ); - } - - #[test] - fn merge_replay_message_meta_creates_fresh_when_none() { - let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_2"); - - let merged = merge_replay_message_meta(None, &message); - - assert_eq!( - merged.get("goose"), - Some(&serde_json::json!({ - "created": 1_700_000_000, - "messageId": "msg_2", - })), - ); - - let chunk = replay_content_chunk_for_message( - &message, - ContentBlock::Text(TextContent::new("replayed text")), - ); - - assert_eq!(chunk.message_id, Some(MessageId::new("msg_2"))); - } - - #[test] - fn merge_replay_message_meta_includes_steer_marker() { - let message = Message::new(Role::User, 1_700_000_000, vec![]) - .with_id("msg_steer") - .with_steer(); - - let merged = merge_replay_message_meta(None, &message); - - assert_eq!( - merged.get("goose"), - Some(&serde_json::json!({ - "created": 1_700_000_000, - "messageId": "msg_steer", - "steer": true, - })), - "replay must carry the steer marker so the boundary survives reload" - ); - } - - #[test] - fn merge_replay_message_meta_omits_steer_when_not_set() { - let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_plain"); - - let merged = merge_replay_message_meta(None, &message); - - assert_eq!(merged.get("goose").and_then(|g| g.get("steer")), None); - } - - #[test] - fn merge_replay_message_meta_omits_message_id_when_none() { - let message = Message::new(Role::Assistant, 1_700_000_000, vec![]); - - let merged = merge_replay_message_meta(None, &message); - - assert_eq!( - merged.get("goose"), - Some(&serde_json::json!({ - "created": 1_700_000_000, - })), + }), ); } } diff --git a/crates/goose/src/acp/server/message_meta.rs b/crates/goose/src/acp/server/message_meta.rs new file mode 100644 index 000000000..85765a8c4 --- /dev/null +++ b/crates/goose/src/acp/server/message_meta.rs @@ -0,0 +1,313 @@ +use crate::conversation::message::{Message, MessageContent}; +use agent_client_protocol::schema::v1::{ContentBlock, ContentChunk, MessageId, Meta}; +use serde::Serialize; + +const OUTPUT_TOKEN_LIMIT_TEXT: &str = + "Response stopped because the model reached its output-token limit."; + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct GooseMessageMeta<'a> { + created: i64, + #[serde(skip_serializing_if = "Option::is_none")] + message_id: Option<&'a str>, + #[serde(skip_serializing_if = "std::ops::Not::not")] + steer: bool, + #[serde(skip_serializing_if = "std::ops::Not::not")] + output_token_limit_reached: bool, + #[serde(skip_serializing_if = "std::ops::Not::not")] + fallback_content: bool, +} + +fn goose_message_meta( + message: &Message, + steer: bool, +) -> serde_json::Map { + let message_meta = GooseMessageMeta { + created: message.created, + message_id: message.id.as_deref(), + steer, + output_token_limit_reached: message.metadata.output_token_limit_reached, + fallback_content: has_output_token_limit_fallback_content(message), + }; + + match serde_json::to_value(message_meta) { + Ok(serde_json::Value::Object(meta)) => meta, + _ => serde_json::Map::new(), + } +} + +fn extend_message_meta(meta: &mut Meta, message: &Message, steer: bool) { + let message_goose = goose_message_meta(message, steer); + let goose_value = meta + .entry("goose".to_string()) + .or_insert_with(|| serde_json::Value::Object(serde_json::Map::new())); + + if let serde_json::Value::Object(goose) = goose_value { + goose.extend(message_goose); + } else { + *goose_value = serde_json::Value::Object(message_goose); + } +} + +fn message_meta_with_steer(message: &Message, steer: bool) -> Meta { + let mut meta = Meta::new(); + extend_message_meta(&mut meta, message, steer); + meta +} + +fn message_meta(message: &Message) -> Meta { + message_meta_with_steer(message, message.metadata.steer) +} + +pub(super) fn message_meta_without_steer(message: &Message) -> Meta { + message_meta_with_steer(message, false) +} + +pub(super) fn merge_message_meta(mut meta: Meta, message: &Message) -> Meta { + extend_message_meta(&mut meta, message, message.metadata.steer); + meta +} + +pub(super) fn content_chunk_for_message(message: &Message, content: ContentBlock) -> ContentChunk { + let mut chunk = ContentChunk::new(content).meta(message_meta(message)); + if let Some(message_id) = message.id.as_deref() { + chunk = chunk.message_id(MessageId::new(message_id)); + } + chunk +} + +pub(super) fn populate_output_token_limit_content(message: &mut Message) { + if message.role != rmcp::model::Role::Assistant + || !message.content.is_empty() + || !message.metadata.output_token_limit_reached + { + return; + } + + message + .content + .push(MessageContent::text(OUTPUT_TOKEN_LIMIT_TEXT)); +} + +fn has_output_token_limit_fallback_content(message: &Message) -> bool { + message.role == rmcp::model::Role::Assistant + && message.metadata.output_token_limit_reached + && matches!( + message.content.as_slice(), + [MessageContent::Text(text)] if text.text == OUTPUT_TOKEN_LIMIT_TEXT + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_client_protocol::schema::v1::TextContent; + use rmcp::model::Role; + + #[test] + fn message_meta_serializes_message_fields() { + let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_live"); + assert_eq!( + message_meta(&message).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + })), + ); + + let steer_message = message.clone().with_steer(); + assert_eq!( + message_meta(&steer_message).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "steer": true, + })), + ); + assert_eq!( + message_meta_without_steer(&steer_message).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + })), + ); + + let message_without_id = Message::new(Role::Assistant, 1_700_000_000, vec![]); + assert_eq!( + message_meta(&message_without_id).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + })), + ); + + let mut limited_message = message.clone(); + limited_message.metadata.output_token_limit_reached = true; + assert_eq!( + message_meta(&limited_message).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "outputTokenLimitReached": true, + })), + ); + + populate_output_token_limit_content(&mut limited_message); + assert_eq!( + message_meta(&limited_message).get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "outputTokenLimitReached": true, + "fallbackContent": true, + })), + ); + } + + #[test] + fn content_chunk_carries_message_id_and_metadata() { + let mut message = Message::new(Role::Assistant, 1_700_000_000, vec![]) + .with_id("msg_live") + .with_steer(); + message.metadata.output_token_limit_reached = true; + populate_output_token_limit_content(&mut message); + let chunk = content_chunk_for_message( + &message, + ContentBlock::Text(TextContent::new(OUTPUT_TOKEN_LIMIT_TEXT)), + ); + + assert_eq!(chunk.message_id, Some(MessageId::new("msg_live"))); + assert_eq!( + chunk.meta.as_ref().and_then(|meta| meta.get("goose")), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "outputTokenLimitReached": true, + "fallbackContent": true, + "steer": true, + })), + ); + } + + #[test] + fn populate_output_token_limit_content_only_changes_empty_marked_assistant_messages() { + let mut message = Message::new(Role::Assistant, 1_700_000_000, vec![]) + .with_id("msg_limited") + .with_steer(); + message.metadata.output_token_limit_reached = true; + + populate_output_token_limit_content(&mut message); + + assert!(matches!( + message.content.as_slice(), + [MessageContent::Text(text)] if text.text == OUTPUT_TOKEN_LIMIT_TEXT + )); + assert!(has_output_token_limit_fallback_content(&message)); + + let mut unmarked_message = Message::new(Role::Assistant, 1_700_000_000, vec![]); + populate_output_token_limit_content(&mut unmarked_message); + assert!(unmarked_message.content.is_empty()); + assert!(!has_output_token_limit_fallback_content(&unmarked_message)); + + let mut marked_message_with_content = + Message::new(Role::Assistant, 1_700_000_000, vec![]).with_text("partial response"); + marked_message_with_content + .metadata + .output_token_limit_reached = true; + populate_output_token_limit_content(&mut marked_message_with_content); + assert!(matches!( + marked_message_with_content.content.as_slice(), + [MessageContent::Text(text)] if text.text == "partial response" + )); + assert!(!has_output_token_limit_fallback_content( + &marked_message_with_content + )); + + let mut marked_user_message = Message::new(Role::User, 1_700_000_000, vec![]); + marked_user_message.metadata.output_token_limit_reached = true; + populate_output_token_limit_content(&mut marked_user_message); + assert!(marked_user_message.content.is_empty()); + assert!(!has_output_token_limit_fallback_content( + &marked_user_message + )); + } + + #[test] + fn merge_message_meta_preserves_existing_metadata() { + let mut message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_1"); + message.metadata.output_token_limit_reached = true; + let existing = serde_json::from_value(serde_json::json!({ + "goose": { + "created": 1, + "messageId": "old", + "toolCall": { + "toolName": "weather__render", + "extensionName": "weather", + }, + "mcpApp": { + "resourceUri": "ui://trusted/app", + "extensionName": "weather", + "toolName": "weather__render", + }, + }, + "otherNamespace": { + "preserve": true, + }, + })) + .unwrap(); + + let merged = merge_message_meta(existing, &message); + + assert_eq!( + merged.get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_1", + "outputTokenLimitReached": true, + "toolCall": { + "toolName": "weather__render", + "extensionName": "weather", + }, + "mcpApp": { + "resourceUri": "ui://trusted/app", + "extensionName": "weather", + "toolName": "weather__render", + }, + })), + ); + assert_eq!( + merged.get("otherNamespace"), + Some(&serde_json::json!({ + "preserve": true, + })), + ); + } + + #[test] + fn merge_message_meta_replaces_non_object_goose_metadata() { + let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_1"); + let existing = serde_json::from_value(serde_json::json!({ + "goose": "invalid", + "otherNamespace": { + "preserve": true, + }, + })) + .unwrap(); + + let merged = merge_message_meta(existing, &message); + + assert_eq!( + merged.get("goose"), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_1", + })), + ); + assert_eq!( + merged.get("otherNamespace"), + Some(&serde_json::json!({ + "preserve": true, + })), + ); + } +} diff --git a/crates/goose/src/acp/server/tool_calls/conversion.rs b/crates/goose/src/acp/server/tool_calls/conversion.rs index 49674f1d8..a41d06e18 100644 --- a/crates/goose/src/acp/server/tool_calls/conversion.rs +++ b/crates/goose/src/acp/server/tool_calls/conversion.rs @@ -1,6 +1,6 @@ use crate::acp::tools::AcpAwareToolMeta; use crate::agents::extension_manager::TRUSTED_TOOL_UPDATE_META_KEY; -use crate::conversation::message::{ToolNameParts, ToolRequest, ToolResponse}; +use crate::conversation::message::{Message, ToolNameParts, ToolRequest, ToolResponse}; use crate::mcp_utils::ToolResult; use agent_client_protocol::schema::v1::{ BlobResourceContents, Content, ContentBlock, EmbeddedResource, EmbeddedResourceResource, @@ -9,6 +9,8 @@ use agent_client_protocol::schema::v1::{ }; use rmcp::model::{CallToolResult, ContentBlock as RmcpContentBlock, ResourceContents}; +use super::super::message_meta::merge_message_meta; + pub(crate) fn format_tool_name(tool_name: &str) -> String { let parts = ToolNameParts::from(tool_name); if let Some(extension_name) = parts.extension_name { @@ -82,10 +84,7 @@ pub(crate) fn goose_tool_call_meta(tool_request: &ToolRequest) -> Option { Some(meta) } -pub(crate) fn build_initial_tool_call( - tool_request: &ToolRequest, - include_generated_title: bool, -) -> ToolCall { +fn build_initial_tool_call(tool_request: &ToolRequest, include_generated_title: bool) -> ToolCall { let tool_name = match &tool_request.tool_call { Ok(tool_call) => tool_call.name.to_string(), Err(_) => "error".to_string(), @@ -117,6 +116,16 @@ pub(crate) fn build_initial_tool_call( tool_call.meta(goose_meta) } +pub(crate) fn build_initial_tool_call_with_message_meta( + tool_request: &ToolRequest, + message: &Message, + include_generated_title: bool, +) -> ToolCall { + let mut tool_call = build_initial_tool_call(tool_request, include_generated_title); + let meta = tool_call.meta.take().unwrap_or_default(); + tool_call.meta(merge_message_meta(meta, message)) +} + pub(crate) fn build_permission_tool_call_update( request_id: &str, tool_name: &str, @@ -453,6 +462,55 @@ mod tests { } } + mod build_initial_tool_call_with_message_meta { + use super::*; + use rmcp::model::Role; + + #[test] + fn merges_message_and_tool_metadata() { + let request = ToolRequest { + id: "req_1".to_string(), + tool_call: Ok(CallToolRequestParams::new("developer__shell")), + metadata: None, + tool_meta: None, + }; + let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_live"); + + let tool_call = build_initial_tool_call_with_message_meta(&request, &message, false); + assert_eq!( + tool_call.meta.as_ref().and_then(|meta| meta.get("goose")), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "toolCall": { + "extensionName": "developer", + "toolName": "developer__shell", + }, + })), + ); + + let mut limited_message = message; + limited_message.metadata.output_token_limit_reached = true; + let limited_tool_call = + build_initial_tool_call_with_message_meta(&request, &limited_message, false); + assert_eq!( + limited_tool_call + .meta + .as_ref() + .and_then(|meta| meta.get("goose")), + Some(&serde_json::json!({ + "created": 1_700_000_000, + "messageId": "msg_live", + "outputTokenLimitReached": true, + "toolCall": { + "extensionName": "developer", + "toolName": "developer__shell", + }, + })), + ); + } + } + mod build_permission_tool_call_update { use super::*; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index a75d2c94d..beac05982 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -2230,6 +2230,7 @@ impl Agent { let mut exit_chat = false; let mut provider_errored = false; let mut provider_produced_content = false; + let mut provider_reached_output_token_limit = false; let mut pending_final_output: Option = None; let mut pending_turn_usage: Option = None; let mut preferred_turn_usage_message_id: Option = None; @@ -2256,11 +2257,13 @@ impl Agent { } if let Some(response) = response { + provider_reached_output_token_limit |= + response.metadata.output_token_limit_reached; + if !response.content.is_empty() - && response - .content - .iter() - .all(|content| matches!(content, MessageContent::SystemNotification(_))) + && response.content.iter().all(|content| { + matches!(content, MessageContent::SystemNotification(_)) + }) { yield AgentEvent::Message(response); tokio::task::yield_now().await; @@ -2318,7 +2321,9 @@ impl Agent { }, ); - if !filtered_response.content.is_empty() { + if !filtered_response.content.is_empty() + || filtered_response.metadata.output_token_limit_reached + { yield AgentEvent::Message(filtered_response.clone()); tokio::task::yield_now().await; } @@ -2871,6 +2876,7 @@ impl Agent { && !exit_chat && !provider_errored && !did_recovery_compact_this_iteration + && !provider_reached_output_token_limit && !provider_produced_content && last_assistant_text.is_empty(); @@ -4224,6 +4230,55 @@ echo start >> "$PLUGIN_ROOT/hook.log" } } + struct OutputLimitMarkerProvider { + include_content: bool, + call_count: AtomicUsize, + } + + impl OutputLimitMarkerProvider { + fn new(include_content: bool) -> Self { + Self { + include_content, + call_count: AtomicUsize::new(0), + } + } + + fn call_count(&self) -> usize { + self.call_count.load(Ordering::SeqCst) + } + } + + #[async_trait::async_trait] + impl crate::providers::base::Provider for OutputLimitMarkerProvider { + async fn stream( + &self, + _model_config: &goose_providers::model::ModelConfig, + _system_prompt: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result { + self.call_count.fetch_add(1, Ordering::SeqCst); + let message_id = "provider-output-limit"; + let content = Message::assistant() + .with_text("Partial answer") + .with_id(message_id); + let mut marker = Message::assistant().with_id(message_id); + marker.metadata.output_token_limit_reached = true; + let usage = ProviderUsage::new("mock-model".to_string(), Usage::default()); + + let mut events = Vec::new(); + if self.include_content { + events.push(Ok((Some(content), None))); + } + events.push(Ok((Some(marker), Some(usage)))); + Ok(Box::pin(futures::stream::iter(events))) + } + + fn get_name(&self) -> &str { + "output-limit-marker" + } + } + struct RefusingProvider { call_count: AtomicUsize, } @@ -4383,6 +4438,78 @@ echo start >> "$PLUGIN_ROOT/hook.log" .collect() } + #[tokio::test] + async fn output_limit_marker_is_emitted_and_persisted() -> Result<()> { + let temp_dir = tempfile::tempdir()?; + let hook_manager = crate::hooks::HookManager::from_plugins_for_test(vec![]); + let provider = Arc::new(OutputLimitMarkerProvider::new(true)); + let (agent, session_id) = + create_test_agent(temp_dir.path().join("data"), hook_manager, provider).await?; + + let messages = run_stop_hook_test_turn(&agent, &session_id, "hello").await?; + let marker = messages + .iter() + .find(|message| message.metadata.output_token_limit_reached) + .expect("output-limit marker should be emitted"); + assert!(marker.content.is_empty()); + assert_eq!(marker.id.as_deref(), Some("provider-output-limit")); + + let session = agent + .config + .session_manager + .get_session(&session_id, true) + .await?; + let conversation = session + .conversation + .expect("session should have a conversation"); + let persisted = conversation + .messages() + .iter() + .find(|message| message.id.as_deref() == Some("provider-output-limit")) + .expect("provider response should be persisted"); + assert_eq!(persisted.as_concat_text(), "Partial answer"); + assert!(persisted.metadata.output_token_limit_reached); + assert!(persisted.metadata.usage.is_some()); + + Ok(()) + } + + #[tokio::test] + async fn zero_content_output_limit_is_persisted_without_empty_response_retry() -> Result<()> { + let temp_dir = tempfile::tempdir()?; + let hook_manager = crate::hooks::HookManager::from_plugins_for_test(vec![]); + let provider = Arc::new(OutputLimitMarkerProvider::new(false)); + let (agent, session_id) = + create_test_agent(temp_dir.path().join("data"), hook_manager, provider.clone()).await?; + + run_stop_hook_test_turn(&agent, &session_id, "hello").await?; + + assert_eq!(provider.call_count(), 1); + let session = agent + .config + .session_manager + .get_session(&session_id, true) + .await?; + let conversation = session + .conversation + .expect("session should have a conversation"); + let persisted = conversation + .messages() + .iter() + .find(|message| message.id.as_deref() == Some("provider-output-limit")) + .expect("zero-content output-limit marker should be persisted"); + assert!(persisted.content.is_empty()); + assert!(persisted.metadata.user_visible); + assert!(!persisted.metadata.agent_visible); + assert!(persisted.metadata.output_token_limit_reached); + assert!(conversation + .agent_visible_messages() + .iter() + .all(|message| message.id.as_deref() != Some("provider-output-limit"))); + + Ok(()) + } + #[tokio::test] async fn session_start_hook_emits_once_for_first_reply_turn() -> Result<()> { let env = SessionStartHookTestEnv::new()?; diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 27ce3f5b0..6a94fbc6d 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -607,6 +607,8 @@ impl Agent { let mut filtered_message = Message::new(response.role.clone(), response.created, filtered_content); + filtered_message.metadata.output_token_limit_reached = + response.metadata.output_token_limit_reached; // Preserve the ID if it exists if let Some(id) = response.id.clone() { @@ -1421,17 +1423,19 @@ mod tests { #[tokio::test] async fn categorize_tool_requests_keeps_thinking_when_not_previously_streamed() { let agent = crate::agents::Agent::new(); - let response = Message::assistant() + let mut response = Message::assistant() .with_thinking("final-only reasoning", "") .with_tool_request( "tool-1", Ok(rmcp::model::CallToolRequestParams::new("test_tool")), ); + response.metadata.output_token_limit_reached = true; let (_frontend_requests, other_requests, filtered_message) = agent.categorize_tool_requests(&response, &[], false).await; assert_eq!(other_requests.len(), 1); + assert!(filtered_message.metadata.output_token_limit_reached); assert_eq!(filtered_message.content.len(), 2); assert!(matches!( filtered_message.content[0], diff --git a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts index 321f7a4ab..55f221ba0 100644 --- a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts +++ b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts @@ -18,6 +18,8 @@ const DEFAULT_TOOL_CALL_UPDATE = { sessionUpdate: 'tool_call_update', toolCallId: 'tool-1', } as const; +const OUTPUT_TOKEN_LIMIT_FALLBACK_TEXT = + 'Response stopped because the model reached its output-token limit.'; function acpUpdate(update: SessionNotification['update']): SessionNotification { return { @@ -145,6 +147,42 @@ describe('createAcpSessionNotificationAdapter', () => { expect(firstContent(messages[0])).toMatchObject({ type: 'text', text: 'Hell' }); }); + it('keeps streamed content visible when an output-limit fallback chunk arrives later', () => { + const adapter = createAcpSessionNotificationAdapter(); + + adapter.apply( + acpUpdate({ + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: 'Partial response' }, + _meta: { goose: { messageId: 'msg-1' } }, + } as SessionNotification['update']) + ); + + const messages = expectOnlyMessagesChange( + adapter.apply( + acpUpdate({ + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: OUTPUT_TOKEN_LIMIT_FALLBACK_TEXT }, + _meta: { + goose: { + messageId: 'msg-1', + outputTokenLimitReached: true, + fallbackContent: true, + }, + }, + } as SessionNotification['update']) + ) + ); + + expect(messages).toHaveLength(1); + expect(firstContent(messages[0])).toMatchObject({ + type: 'text', + text: 'Partial response', + }); + expect(messages[0].metadata.outputTokenLimitReached).toBe(true); + expect(messages[0].metadata.fallbackContent).toBeUndefined(); + }); + it('reconciles locally rendered steer text with server chunks', () => { const adapter = createAcpSessionNotificationAdapter([ { @@ -296,6 +334,62 @@ describe('createAcpSessionNotificationAdapter', () => { signature: '', }); }); + + it('preserves output-limit metadata when creating a thought message', () => { + const adapter = createAcpSessionNotificationAdapter(); + + const messages = expectOnlyMessagesChange( + adapter.apply( + acpUpdate({ + sessionUpdate: 'agent_thought_chunk', + content: { type: 'text', text: 'Truncated thinking' }, + _meta: { + goose: { + messageId: 'thought-1', + outputTokenLimitReached: true, + }, + }, + } as SessionNotification['update']) + ) + ); + + expect(messages).toHaveLength(1); + expect(messages[0].metadata.outputTokenLimitReached).toBe(true); + }); + + it('preserves output-limit metadata when updating a thought message', () => { + const adapter = createAcpSessionNotificationAdapter(); + + adapter.apply( + acpUpdate({ + sessionUpdate: 'agent_thought_chunk', + content: { type: 'text', text: 'Truncated ' }, + _meta: { goose: { messageId: 'thought-1' } }, + } as SessionNotification['update']) + ); + + const messages = expectOnlyMessagesChange( + adapter.apply( + acpUpdate({ + sessionUpdate: 'agent_thought_chunk', + content: { type: 'text', text: 'thinking' }, + _meta: { + goose: { + messageId: 'thought-1', + outputTokenLimitReached: true, + }, + }, + } as SessionNotification['update']) + ) + ); + + expect(messages).toHaveLength(1); + expect(firstContent(messages[0])).toMatchObject({ + type: 'thinking', + thinking: 'Truncated thinking', + }); + expect(messages[0].metadata.outputTokenLimitReached).toBe(true); + }); }); describe('tools', () => { diff --git a/ui/desktop/src/acp/adapter/messages.ts b/ui/desktop/src/acp/adapter/messages.ts index a5634b685..a91ab1cc4 100644 --- a/ui/desktop/src/acp/adapter/messages.ts +++ b/ui/desktop/src/acp/adapter/messages.ts @@ -31,6 +31,20 @@ export function applyContentChunk( const existing = findMessageForChunk(state, role, messageId, gooseMeta.created); if (existing) { + const isOutputLimitFallbackChunk = + gooseMeta.outputTokenLimitReached === true && gooseMeta.fallbackContent === true; + const existingMessageHasContent = existing.content.length > 0; + const shouldSkipFallbackChunk = isOutputLimitFallbackChunk && existingMessageHasContent; + + existing.metadata.outputTokenLimitReached = gooseMeta.outputTokenLimitReached; + existing.metadata.fallbackContent = shouldSkipFallbackChunk + ? undefined + : gooseMeta.fallbackContent; + + if (shouldSkipFallbackChunk) { + return messagesChangeWithLocalSteerConfirmation(state, existing, gooseMeta.steer); + } + const lastContent = existing.content[existing.content.length - 1]; if (reconcileLocalSteerTextChunk(state, existing, content, gooseMeta.steer)) { return messagesChangeWithLocalSteerConfirmation(state, existing, gooseMeta.steer); @@ -54,6 +68,8 @@ export function applyContentChunk( metadata: { ...DEFAULT_VISIBLE_MESSAGE_METADATA, ...(gooseMeta.steer ? { steer: true } : {}), + outputTokenLimitReached: gooseMeta.outputTokenLimitReached, + fallbackContent: gooseMeta.fallbackContent, }, }); } @@ -71,23 +87,26 @@ export function applyThoughtChunk( const gooseMeta = getGooseMessageMeta(update); const messageId = update.messageId ?? gooseMeta.messageId; - const existing = findMessageForChunk(state, 'assistant', messageId, gooseMeta.created); + let message = findMessageForChunk(state, 'assistant', messageId, gooseMeta.created); - if (existing) { - const lastContent = existing.content[existing.content.length - 1]; - if (lastContent?.type === 'thinking') { - lastContent.thinking += update.content.text; - } else { - existing.content.push({ type: 'thinking', thinking: update.content.text, signature: '' }); - } - } else { - state.messages.push({ + if (!message) { + message = { ...(messageId ? { id: messageId } : {}), role: 'assistant', created: gooseMeta.created ?? Math.floor(Date.now() / 1000), - content: [{ type: 'thinking', thinking: update.content.text, signature: '' }], + content: [], metadata: { ...DEFAULT_VISIBLE_MESSAGE_METADATA }, - }); + }; + state.messages.push(message); + } + + message.metadata.outputTokenLimitReached = gooseMeta.outputTokenLimitReached; + + const lastContent = message.content[message.content.length - 1]; + if (lastContent?.type === 'thinking') { + lastContent.thinking += update.content.text; + } else { + message.content.push({ type: 'thinking', thinking: update.content.text, signature: '' }); } return messagesChange(state); diff --git a/ui/desktop/src/acp/adapter/shared.ts b/ui/desktop/src/acp/adapter/shared.ts index aa5f033f8..3a94611fe 100644 --- a/ui/desktop/src/acp/adapter/shared.ts +++ b/ui/desktop/src/acp/adapter/shared.ts @@ -26,6 +26,8 @@ export type ToolCallState = Omit; export interface GooseMessageMeta { messageId?: string; created?: number; + outputTokenLimitReached?: boolean; + fallbackContent?: boolean; steer?: boolean; } @@ -65,9 +67,13 @@ export function getGooseMessageMeta(update: { _meta?: unknown }): GooseMessageMe return {}; } + const outputTokenLimitReached = goose.outputTokenLimitReached === true; + return { created: typeof goose.created === 'number' ? goose.created : undefined, messageId: typeof goose.messageId === 'string' ? goose.messageId : undefined, + outputTokenLimitReached: outputTokenLimitReached ? true : undefined, + fallbackContent: goose.fallbackContent === true ? true : undefined, steer: goose.steer === true ? true : undefined, }; } diff --git a/ui/desktop/src/components/GooseMessage.tsx b/ui/desktop/src/components/GooseMessage.tsx index dbad9f23f..5997b4fa6 100644 --- a/ui/desktop/src/components/GooseMessage.tsx +++ b/ui/desktop/src/components/GooseMessage.tsx @@ -1,4 +1,5 @@ import { useMemo, useRef } from 'react'; +import { AlertTriangle } from 'lucide-react'; import ImagePreview from './ImagePreview'; import { formatMessageTimestamp } from '../utils/timeUtils'; import MarkdownContent from './MarkdownContent'; @@ -49,8 +50,13 @@ export default function GooseMessage({ }: GooseMessageProps) { const contentRef = useRef(null); - const { textContent: displayText, imagePaths } = getTextAndImageContent(message); - const thinkingContent = getThinkingContent(message); + const outputTokenLimitReached = message.metadata.outputTokenLimitReached === true; + const isOutputTokenLimitFallback = + outputTokenLimitReached && message.metadata.fallbackContent === true; + const { textContent, imagePaths: allImagePaths } = getTextAndImageContent(message); + const displayText = isOutputTokenLimitFallback ? '' : textContent; + const imagePaths = isOutputTokenLimitFallback ? [] : allImagePaths; + const thinkingContent = isOutputTokenLimitFallback ? null : getThinkingContent(message); const timestamp = useMemo(() => formatMessageTimestamp(message.created), [message.created]); const toolRequests = getToolRequests(message); @@ -76,6 +82,9 @@ export default function GooseMessage({ ); const hasToolConfirmation = toolConfirmationContent !== undefined; const hasElicitation = elicitationContent !== undefined; + const outputTokenLimitNotice = isOutputTokenLimitFallback + ? "Response reached the model's output-token limit before returning content." + : "Response reached the model's output-token limit and may be incomplete."; const elicitationData = elicitationContent?.data.actionType === 'elicitation' ? (elicitationContent.data as typeof elicitationContent.data & { @@ -219,6 +228,13 @@ export default function GooseMessage({ )} + {outputTokenLimitReached && ( +
+ + {outputTokenLimitNotice} +
+ )} + {hasToolConfirmation && !toolConfirmationShownInline && (