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

feat: surface output-token-limit info (#10831)

Browse Source
This commit is contained in:
Lifei Zhou
2026-08-04 22:44:08 +10:00
committed by GitHub
parent 6a1344ba49
commit bb539f7d6f
17 changed files with 1471 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/output.rs
+18 -1
View File
@@ -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<T: Display>(value: T) -> StyledObject<T> {
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(), &notification.msg);
crates/goose-provider-types/src/conversation.rs
+60 -9
View File
@@ -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;
crates/goose-provider-types/src/conversation/message.rs
+8 -1
View File
@@ -665,6 +665,9 @@ pub struct MessageMetadata {
pub agent_visible: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub inference: Option<InferenceMetadata>,
/// 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]
crates/goose-provider-types/src/formats/anthropic.rs
+39
View File
@@ -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<String>,
tool_calls: Vec<String>,
tool_errors: Vec<String>,
output_token_limit_message_ids: Vec<Option<String>>,
}
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,
crates/goose-provider-types/src/formats/openai.rs
+323 -19
View File
@@ -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<Vec<Value>> {
/// Convert OpenAI's API response to internal Message format
pub fn response_to_message(response: &Value) -> anyhow::Result<Message> {
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<Message> {
})
.filter(|m: &serde_json::Map<String, Value>| !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<Message> {
}
}
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<StreamingChunk, ProviderError> {
})
}
fn output_token_limit_marker(id: Option<String>) -> 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<S>(
mut stream: S,
) -> impl Stream<Item = anyhow::Result<(Option<Message>, Option<ProviderUsage>)>> + '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<String> = None;
let mut last_response_id: Option<String> = 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<ProviderUsage>,
tool_calls: Vec<String>,
has_text_content: bool,
text: String,
output_token_limit_message_ids: Vec<Option<String>>,
}
async fn run_streaming_test(response_lines: &str) -> anyhow::Result<StreamingUsageTestResult> {
@@ -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\":\"<think>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\":\"<thi\"},\"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?;
usage_count += usize::from(usage.is_some());
if let Some(message) = message {
streamed_messages.push(message);
}
}
assert_eq!(usage_count, 1);
let marked_messages: Vec<_> = 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(), "<thi");
Ok(())
}
#[tokio::test]
async fn test_streaming_suppresses_inline_think_when_structured_reasoning_follows(
) -> anyhow::Result<()> {
crates/goose-provider-types/src/formats/openai_responses.rs
+160
View File
@@ -107,6 +107,26 @@ pub struct ResponseReasoningInfo {
pub summary: Option<String>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ResponseIncompleteDetails {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reason: Option<String>,
}
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<ResponseUsage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning: Option<ResponseReasoningInfo>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub incomplete_details: Option<ResponseIncompleteDetails>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
@@ -866,6 +894,15 @@ fn process_streaming_output_items(
Ok(content)
}
fn output_token_limit_marker(id: Option<String>) -> 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<S>(
mut stream: S,
) -> impl Stream<Item = anyhow::Result<(Option<Message>, Option<ProviderUsage>)>> + 'static
@@ -881,6 +918,7 @@ where
let mut final_usage: Option<ProviderUsage> = None;
let mut output_items: Vec<ResponseOutputItemInfo> = 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<ProviderUsage> = 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<ProviderUsage> = 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![
crates/goose/src/acp/server.rs
+92 -102
View File
@@ -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<Usage> {
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<Agent>,
tool_requests: &HashMap<String, ToolRequest>,
cx: &ConnectionTo<Client>,
) -> 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<Agent>,
cx: &ConnectionTo<Client>,
) -> 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(
crates/goose/src/acp/server/load_session.rs
+107 -169
View File
@@ -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<String, serde_json::Value> {
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<Meta>, 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<Message> {
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<Client>,
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<_>>(),
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,
})),
}),
);
}
}
crates/goose/src/acp/server/message_meta.rs
+313
View File
@@ -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<String, serde_json::Value> {
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,
})),
);
}
}
crates/goose/src/acp/server/tool_calls/conversion.rs
+63 -5
View File
@@ -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<Meta> {
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::*;
crates/goose/src/agents/agent.rs
+132 -5
View File
@@ -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<String> = None;
let mut pending_turn_usage: Option<ProviderUsage> = None;
let mut preferred_turn_usage_message_id: Option<String> = 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<MessageStream, ProviderError> {
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()?;
crates/goose/src/agents/reply_parts.rs
+5 -1
View File
@@ -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],
ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts
+94
View File
@@ -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', () => {
ui/desktop/src/acp/adapter/messages.ts
+31 -12
View File
@@ -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);
ui/desktop/src/acp/adapter/shared.ts
+6
View File
@@ -26,6 +26,8 @@ export type ToolCallState = Omit<ToolCallUpdate, '_meta'>;
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,
};
}
ui/desktop/src/components/GooseMessage.tsx
+18 -2
View File
@@ -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<HTMLDivElement | null>(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({
</div>
)}
{outputTokenLimitReached && (
<div className="mt-2 flex items-start gap-1.5 text-xs text-text-secondary">
<AlertTriangle className="mt-0.5 h-3.5 w-3.5 shrink-0 text-yellow-600 dark:text-yellow-400" />
<span>{outputTokenLimitNotice}</span>
</div>
)}
{hasToolConfirmation && !toolConfirmationShownInline && (
<ToolCallConfirmation
sessionId={sessionId}
ui/desktop/src/types/message.ts
+2
View File
@@ -186,7 +186,9 @@ export type MessageUsage = {
export type MessageMetadata = {
agentVisible: boolean;
fallbackContent?: boolean;
inference?: InferenceMetadata | null;
outputTokenLimitReached?: boolean;
steer?: boolean;
usage?: MessageUsage | null;
userVisible: boolean;