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

feat: option to stream json - jsonl really (#6228)

Browse Source
This commit is contained in:
Michael Neale
2025-12-23 10:24:19 +11:00
committed by GitHub
parent 91a6b21f39
commit 1a79f28ad8
2 changed files with 107 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+3 -3
View File
@@ -753,13 +753,13 @@ enum Command {
)] )]
additional_sub_recipes: Vec<String>, additional_sub_recipes: Vec<String>,
/// Output format (text, json) /// Output format (text, json, stream-json)
#[arg( #[arg(
long = "output-format", long = "output-format",
value_name = "FORMAT", value_name = "FORMAT",
help = "Output format (text, json)", help = "Output format (text, json, stream-json)",
default_value = "text", default_value = "text",
value_parser = clap::builder::PossibleValuesParser::new(["text", "json"]) value_parser = clap::builder::PossibleValuesParser::new(["text", "json", "stream-json"])
)] )]
output_format: String, output_format: String,
crates/goose-cli/src/session/mod.rs
+104 -25
View File
@@ -65,6 +65,42 @@ struct JsonMetadata {
status: String, status: String,
} }
#[derive(Serialize, Debug)]
#[serde(tag = "type", rename_all = "snake_case")]
enum StreamEvent {
Message {
message: Message,
},
Notification {
extension_id: String,
#[serde(flatten)]
data: NotificationData,
},
ModelChange {
model: String,
mode: String,
},
Error {
error: String,
},
Complete {
total_tokens: Option<i32>,
},
}
#[derive(Serialize, Debug)]
#[serde(rename_all = "snake_case")]
enum NotificationData {
Log {
message: String,
},
Progress {
progress: f64,
total: Option<f64>,
message: Option<String>,
},
}
pub enum RunMode { pub enum RunMode {
Normal, Normal,
Plan, Plan,
@@ -812,8 +848,15 @@ impl CliSession {
interactive: bool, interactive: bool,
cancel_token: CancellationToken, cancel_token: CancellationToken,
) -> Result<()> { ) -> Result<()> {
// Cache the output format check to avoid repeated string comparisons in the hot loop
let is_json_mode = self.output_format == "json"; let is_json_mode = self.output_format == "json";
let is_stream_json_mode = self.output_format == "stream-json";
// Helper to emit a streaming JSON event
let emit_stream_event = |event: &StreamEvent| {
if let Ok(json) = serde_json::to_string(event) {
println!("{}", json);
}
};
let session_config = SessionConfig { let session_config = SessionConfig {
id: self.session_id.clone(), id: self.session_id.clone(),
@@ -1027,13 +1070,15 @@ impl CliSession {
if interactive {output::hide_thinking()}; if interactive {output::hide_thinking()};
let _ = progress_bars.hide(); let _ = progress_bars.hide();
// Don't render in JSON mode // Handle different output formats
if !is_json_mode { if is_stream_json_mode {
emit_stream_event(&StreamEvent::Message { message: message.clone() });
} else if !is_json_mode {
output::render_message(&message, self.debug); output::render_message(&message, self.debug);
} }
} }
} }
Some(Ok(AgentEvent::McpNotification((_id, message)))) => { Some(Ok(AgentEvent::McpNotification((extension_id, message)))) => {
match &message { match &message {
ServerNotification::LoggingMessageNotification(notification) => { ServerNotification::LoggingMessageNotification(notification) => {
let data = &notification.params.data; let data = &notification.params.data;
@@ -1101,9 +1146,14 @@ impl CliSession {
}, },
}; };
if is_stream_json_mode {
emit_stream_event(&StreamEvent::Notification {
extension_id: extension_id.clone(),
data: NotificationData::Log { message: formatted_message.clone() },
});
}
// Handle subagent notifications - show immediately // Handle subagent notifications - show immediately
if let Some(_id) = subagent_id { else if let Some(_id) = subagent_id {
// TODO: proper display for subagent notifications
if interactive { if interactive {
let _ = progress_bars.hide(); let _ = progress_bars.hide();
if !is_json_mode { if !is_json_mode {
@@ -1125,7 +1175,6 @@ impl CliSession {
std::io::stdout().flush().unwrap(); std::io::stdout().flush().unwrap();
} }
} else if notification_type == "shell_output" { } else if notification_type == "shell_output" {
// Hide spinner, print shell output, spinner will resume
if interactive { if interactive {
let _ = progress_bars.hide(); let _ = progress_bars.hide();
} }
@@ -1145,12 +1194,24 @@ impl CliSession {
let text = notification.params.message.as_deref(); let text = notification.params.message.as_deref();
let total = notification.params.total; let total = notification.params.total;
let token = &notification.params.progress_token; let token = &notification.params.progress_token;
progress_bars.update(
&token.0.to_string(), if is_stream_json_mode {
progress, emit_stream_event(&StreamEvent::Notification {
total, extension_id: extension_id.clone(),
text, data: NotificationData::Progress {
); progress,
total,
message: text.map(String::from),
},
});
} else {
progress_bars.update(
&token.0.to_string(),
progress,
total,
text,
);
}
}, },
_ => (), _ => (),
} }
@@ -1159,32 +1220,44 @@ impl CliSession {
self.messages = updated_conversation; self.messages = updated_conversation;
} }
Some(Ok(AgentEvent::ModelChange { model, mode })) => { Some(Ok(AgentEvent::ModelChange { model, mode })) => {
// Log model change if in debug mode if is_stream_json_mode {
if self.debug { emit_stream_event(&StreamEvent::ModelChange {
model: model.clone(),
mode: mode.clone(),
});
} else if self.debug {
eprintln!("Model changed to {} in {} mode", model, mode); eprintln!("Model changed to {} in {} mode", model, mode);
} }
} }
Some(Err(e)) => { Some(Err(e)) => {
// TODO(Douwe): Delete this let error_msg = e.to_string();
// Check if it's a ProviderError::ContextLengthExceeded
if is_stream_json_mode {
emit_stream_event(&StreamEvent::Error { error: error_msg.clone() });
}
if e.downcast_ref::<goose::providers::errors::ProviderError>() if e.downcast_ref::<goose::providers::errors::ProviderError>()
.map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_))) .map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_)))
.unwrap_or(false) { .unwrap_or(false) {
output::render_text( if !is_stream_json_mode {
"Compaction requested. Should have happened in the agent!", output::render_text(
Some(Color::Yellow), "Compaction requested. Should have happened in the agent!",
true Some(Color::Yellow),
); true
);
}
warn!("Compaction requested. Should have happened in the agent!"); warn!("Compaction requested. Should have happened in the agent!");
} }
eprintln!("Error: {}", e); if !is_stream_json_mode {
eprintln!("Error: {}", error_msg);
}
cancel_token_clone.cancel(); cancel_token_clone.cancel();
drop(stream); drop(stream);
if let Err(e) = self.handle_interrupted_messages(false).await { if let Err(e) = self.handle_interrupted_messages(false).await {
eprintln!("Error handling interruption: {}", e); eprintln!("Error handling interruption: {}", e);
} else { } else if !is_stream_json_mode {
output::render_error( output::render_error(
"The error above was an exception we were not able to handle.\n\ "The error above was an exception we were not able to handle.\n\
These errors are often related to connection or authentication\n\ These errors are often related to connection or authentication\n\
@@ -1207,7 +1280,7 @@ impl CliSession {
} }
} }
// Output JSON if requested // Output based on format
if is_json_mode { if is_json_mode {
let metadata = match SessionManager::get_session(&self.session_id, false).await { let metadata = match SessionManager::get_session(&self.session_id, false).await {
Ok(session) => JsonMetadata { Ok(session) => JsonMetadata {
@@ -1226,6 +1299,12 @@ impl CliSession {
}; };
println!("{}", serde_json::to_string_pretty(&json_output)?); println!("{}", serde_json::to_string_pretty(&json_output)?);
} else if is_stream_json_mode {
let total_tokens = SessionManager::get_session(&self.session_id, false)
.await
.ok()
.and_then(|s| s.total_tokens);
emit_stream_event(&StreamEvent::Complete { total_tokens });
} else { } else {
println!(); println!();
} }