mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
feat: add streaming support for Claude Code CLI provider (#6833)
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
@@ -1,16 +1,21 @@
|
||||
use anyhow::Result;
|
||||
use async_stream::try_stream;
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use rmcp::model::Role;
|
||||
use rmcp::model::{Role, Tool};
|
||||
use serde_json::{json, Value};
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use tempfile::NamedTempFile;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
|
||||
use super::base::{ConfigKey, Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage};
|
||||
use super::base::{
|
||||
stream_from_single_message, ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata,
|
||||
ProviderUsage, Usage,
|
||||
};
|
||||
use super::errors::ProviderError;
|
||||
use super::utils::{filter_extensions_from_system_prompt, RequestLog};
|
||||
use crate::config::base::ClaudeCodeCommand;
|
||||
@@ -20,7 +25,8 @@ use crate::config::{Config, ExtensionConfig, GooseMode};
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::subprocess::configure_subprocess;
|
||||
use rmcp::model::Tool;
|
||||
|
||||
use super::cli_common::{error_from_event, extract_usage_tokens};
|
||||
|
||||
const CLAUDE_CODE_PROVIDER_NAME: &str = "claude-code";
|
||||
pub const CLAUDE_CODE_DEFAULT_MODEL: &str = "default";
|
||||
@@ -35,6 +41,7 @@ struct CliProcess {
|
||||
current_model: String,
|
||||
log_model_update: bool,
|
||||
next_request_id: u64,
|
||||
needs_drain: bool,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CliProcess {
|
||||
@@ -129,10 +136,58 @@ impl CliProcess {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn drain_pending_response(&mut self) {
|
||||
if !self.needs_drain {
|
||||
return;
|
||||
}
|
||||
tracing::debug!("Draining cancelled response from CLI process");
|
||||
|
||||
let drain = async {
|
||||
let mut line = String::new();
|
||||
loop {
|
||||
line.clear();
|
||||
match self.reader.read_line(&mut line).await {
|
||||
Ok(0) => break,
|
||||
Ok(_) => {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
|
||||
match parsed.get("type").and_then(|t| t.as_str()) {
|
||||
Some("result") | Some("error") => break,
|
||||
_ => continue,
|
||||
}
|
||||
} else {
|
||||
tracing::trace!(line = trimmed, "Non-JSON line during drain");
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
if tokio::time::timeout(DRAIN_TIMEOUT, drain).await.is_err() {
|
||||
// CLI is still producing the old response. Leave needs_drain
|
||||
// true so the next call retries — by then the old response
|
||||
// likely completed and drain will succeed quickly.
|
||||
tracing::warn!(
|
||||
"Drain did not complete in {DRAIN_TIMEOUT:?}; \
|
||||
will retry on next request"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
self.needs_drain = false;
|
||||
tracing::debug!("Drain complete, protocol re-synced");
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CliProcess {
|
||||
fn drop(&mut self) {
|
||||
self.stderr_handle.abort();
|
||||
let _ = self.child.start_kill();
|
||||
}
|
||||
}
|
||||
@@ -152,7 +207,7 @@ pub struct ClaudeCodeProvider {
|
||||
#[serde(skip)]
|
||||
mcp_config_file: Option<NamedTempFile>,
|
||||
#[serde(skip)]
|
||||
cli_process: tokio::sync::OnceCell<tokio::sync::Mutex<CliProcess>>,
|
||||
cli_process: tokio::sync::OnceCell<Arc<tokio::sync::Mutex<CliProcess>>>,
|
||||
}
|
||||
|
||||
impl ClaudeCodeProvider {
|
||||
@@ -245,14 +300,11 @@ impl ClaudeCodeProvider {
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
GooseMode::Chat => {
|
||||
// Chat mode doesn't need permission flags
|
||||
}
|
||||
GooseMode::Chat => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Parse NDJSON stream-json response from Claude CLI
|
||||
fn parse_claude_response(
|
||||
&self,
|
||||
json_lines: &[String],
|
||||
@@ -265,7 +317,6 @@ impl ClaudeCodeProvider {
|
||||
match parsed.get("type").and_then(|t| t.as_str()) {
|
||||
Some("assistant") => {
|
||||
if let Some(message) = parsed.get("message") {
|
||||
// Extract text content from this assistant message
|
||||
if let Some(content) = message.get("content").and_then(|c| c.as_array())
|
||||
{
|
||||
for item in content {
|
||||
@@ -276,65 +327,33 @@ impl ClaudeCodeProvider {
|
||||
all_text_content.push(text.to_string());
|
||||
}
|
||||
}
|
||||
// Skip tool_use - those are claude CLI's internal tools
|
||||
}
|
||||
}
|
||||
|
||||
// Extract usage information
|
||||
if let Some(usage_info) = message.get("usage") {
|
||||
usage.input_tokens = usage_info
|
||||
.get("input_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|v| v as i32);
|
||||
usage.output_tokens = usage_info
|
||||
.get("output_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|v| v as i32);
|
||||
|
||||
if usage.total_tokens.is_none() {
|
||||
if let (Some(input), Some(output)) =
|
||||
(usage.input_tokens, usage.output_tokens)
|
||||
{
|
||||
usage.total_tokens = Some(input + output);
|
||||
}
|
||||
}
|
||||
usage = extract_usage_tokens(usage_info);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("result") => {
|
||||
// Extract additional usage info from result if available
|
||||
if let Some(result_usage) = parsed.get("usage") {
|
||||
if usage.input_tokens.is_none() {
|
||||
usage.input_tokens = result_usage
|
||||
.get("input_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|v| v as i32);
|
||||
}
|
||||
if usage.output_tokens.is_none() {
|
||||
usage.output_tokens = result_usage
|
||||
.get("output_tokens")
|
||||
.and_then(|v| v.as_i64())
|
||||
.map(|v| v as i32);
|
||||
}
|
||||
let new = extract_usage_tokens(result_usage);
|
||||
usage = Usage::new(
|
||||
usage.input_tokens.or(new.input_tokens),
|
||||
usage.output_tokens.or(new.output_tokens),
|
||||
None,
|
||||
);
|
||||
}
|
||||
}
|
||||
Some("error") => {
|
||||
let error_msg = parsed
|
||||
.get("error")
|
||||
.and_then(|e| e.as_str())
|
||||
.unwrap_or("Unknown error");
|
||||
return Err(ProviderError::RequestFailed(format!(
|
||||
"Claude CLI error: {}",
|
||||
error_msg
|
||||
)));
|
||||
return Err(error_from_event("Claude CLI", &parsed));
|
||||
}
|
||||
Some("system") => {} // Ignore system init events
|
||||
_ => {} // Ignore other event types
|
||||
Some("system") => {}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Combine all text content into a single message
|
||||
let combined_text = all_text_content.join("\n\n");
|
||||
if combined_text.is_empty() {
|
||||
return Err(ProviderError::RequestFailed(
|
||||
@@ -353,6 +372,73 @@ impl ClaudeCodeProvider {
|
||||
Ok((response_message, usage))
|
||||
}
|
||||
|
||||
fn spawn_process(&self, filtered_system: &str) -> Result<CliProcess, ProviderError> {
|
||||
let mut cmd = self.build_stream_json_command();
|
||||
|
||||
if let Some(f) = &self.mcp_config_file {
|
||||
cmd.arg("--mcp-config").arg(f.path());
|
||||
cmd.arg("--strict-mcp-config");
|
||||
}
|
||||
|
||||
cmd.arg("--include-partial-messages")
|
||||
.arg("--system-prompt")
|
||||
.arg(filtered_system)
|
||||
.arg("--model")
|
||||
.arg(&self.model.model_name);
|
||||
|
||||
Self::apply_permission_flags(&mut cmd)?;
|
||||
|
||||
let mut child = cmd.spawn().map_err(|e| {
|
||||
ProviderError::RequestFailed(format!(
|
||||
"Failed to spawn Claude CLI command '{:?}': {}.",
|
||||
self.command, e
|
||||
))
|
||||
})?;
|
||||
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdin".to_string()))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdout".to_string()))?;
|
||||
|
||||
let stderr = child.stderr.take();
|
||||
let stderr_handle = tokio::spawn(async move {
|
||||
let mut output = String::new();
|
||||
if let Some(mut stderr) = stderr {
|
||||
use tokio::io::AsyncReadExt;
|
||||
let _ = stderr.read_to_string(&mut output).await;
|
||||
}
|
||||
output
|
||||
});
|
||||
|
||||
Ok(CliProcess {
|
||||
child,
|
||||
stdin: Box::new(stdin),
|
||||
reader: BufReader::new(Box::new(stdout)),
|
||||
stderr_handle,
|
||||
current_model: self.model.model_name.clone(),
|
||||
log_model_update: false,
|
||||
next_request_id: 0,
|
||||
needs_drain: false,
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_or_init_process(
|
||||
&self,
|
||||
filtered_system: &str,
|
||||
) -> Result<&Arc<tokio::sync::Mutex<CliProcess>>, ProviderError> {
|
||||
self.cli_process
|
||||
.get_or_try_init(|| async {
|
||||
Ok(Arc::new(tokio::sync::Mutex::new(
|
||||
self.spawn_process(filtered_system)?,
|
||||
)))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn execute_command(
|
||||
&self,
|
||||
system: &str,
|
||||
@@ -363,74 +449,19 @@ impl ClaudeCodeProvider {
|
||||
) -> Result<Vec<String>, ProviderError> {
|
||||
let filtered_system = filter_extensions_from_system_prompt(system);
|
||||
|
||||
if std::env::var("GOOSE_CLAUDE_CODE_DEBUG").is_ok() {
|
||||
println!("=== CLAUDE CODE PROVIDER DEBUG ===");
|
||||
println!("Command: {:?}", self.command);
|
||||
println!("Original system prompt length: {} chars", system.len());
|
||||
println!(
|
||||
"Filtered system prompt length: {} chars",
|
||||
filtered_system.len()
|
||||
);
|
||||
println!("Filtered system prompt: {}", filtered_system);
|
||||
println!("================================");
|
||||
}
|
||||
|
||||
// Spawn lazily on first call (OnceCell ensures exactly once)
|
||||
let process_mutex = self
|
||||
.cli_process
|
||||
.get_or_try_init(|| async {
|
||||
let mut cmd = self.build_stream_json_command();
|
||||
if let Some(f) = &self.mcp_config_file {
|
||||
cmd.arg("--mcp-config").arg(f.path());
|
||||
cmd.arg("--strict-mcp-config");
|
||||
}
|
||||
// System prompt is set once at process start and cannot be updated at runtime.
|
||||
cmd.arg("--system-prompt").arg(&filtered_system);
|
||||
|
||||
// The initial model can be updated later.
|
||||
cmd.arg("--model").arg(&self.model.model_name);
|
||||
|
||||
Self::apply_permission_flags(&mut cmd)?;
|
||||
|
||||
let mut child = cmd.spawn().map_err(|e| {
|
||||
ProviderError::RequestFailed(format!(
|
||||
"Failed to spawn Claude CLI command '{:?}': {}.",
|
||||
self.command, e
|
||||
))
|
||||
})?;
|
||||
|
||||
let stdin = child.stdin.take().ok_or_else(|| {
|
||||
ProviderError::RequestFailed("Failed to capture stdin".to_string())
|
||||
})?;
|
||||
let stdout = child.stdout.take().ok_or_else(|| {
|
||||
ProviderError::RequestFailed("Failed to capture stdout".to_string())
|
||||
})?;
|
||||
|
||||
// Drain stderr concurrently to prevent pipe buffer deadlock
|
||||
let stderr = child.stderr.take();
|
||||
let stderr_handle = tokio::spawn(async move {
|
||||
let mut output = String::new();
|
||||
if let Some(mut stderr) = stderr {
|
||||
use tokio::io::AsyncReadExt;
|
||||
let _ = stderr.read_to_string(&mut output).await;
|
||||
}
|
||||
output
|
||||
});
|
||||
|
||||
Ok::<_, ProviderError>(tokio::sync::Mutex::new(CliProcess {
|
||||
child,
|
||||
stdin: Box::new(stdin),
|
||||
reader: BufReader::new(Box::new(stdout)),
|
||||
stderr_handle,
|
||||
current_model: String::new(),
|
||||
log_model_update: false,
|
||||
next_request_id: 0,
|
||||
}))
|
||||
})
|
||||
.await?;
|
||||
tracing::debug!(
|
||||
command = ?self.command,
|
||||
system_prompt_len = system.len(),
|
||||
filtered_system_prompt_len = filtered_system.len(),
|
||||
"Executing Claude CLI command"
|
||||
);
|
||||
|
||||
let process_mutex = self.get_or_init_process(&filtered_system).await?;
|
||||
let mut process = process_mutex.lock().await;
|
||||
|
||||
// Drain any pending response from a cancelled stream
|
||||
process.drain_pending_response().await;
|
||||
|
||||
// Switch model if it differs from what the CLI is currently using.
|
||||
process.send_set_model(model).await?;
|
||||
|
||||
@@ -449,7 +480,7 @@ impl ClaudeCodeProvider {
|
||||
ProviderError::RequestFailed(format!("Failed to write newline to stdin: {}", e))
|
||||
})?;
|
||||
|
||||
// Read lines until we see a "result" event
|
||||
// Read lines until we see a "result" or "error" event
|
||||
let mut lines = Vec::new();
|
||||
let mut line = String::new();
|
||||
|
||||
@@ -457,7 +488,6 @@ impl ClaudeCodeProvider {
|
||||
line.clear();
|
||||
match process.reader.read_line(&mut line).await {
|
||||
Ok(0) => {
|
||||
// EOF means the process died
|
||||
return Err(ProviderError::RequestFailed(
|
||||
"Claude CLI process terminated unexpectedly".to_string(),
|
||||
));
|
||||
@@ -467,31 +497,31 @@ impl ClaudeCodeProvider {
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
lines.push(trimmed.to_string());
|
||||
|
||||
if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
|
||||
match parsed.get("type").and_then(|t| t.as_str()) {
|
||||
Some("result") => break,
|
||||
Some("error") => break,
|
||||
Some("stream_event") => continue,
|
||||
Some("result") | Some("error") => {
|
||||
lines.push(trimmed.to_string());
|
||||
break;
|
||||
}
|
||||
// The system init with the resolved model arrives here,
|
||||
// not in send_set_model (which only sees control_response):
|
||||
// send_set_model: {"type":"control_response",...}
|
||||
// execute_command: {"type":"system",...,"model":"claude-sonnet-4-5-20250929",...}
|
||||
// not in send_set_model (which only sees control_response).
|
||||
Some("system") if process.log_model_update => {
|
||||
if let Some(resolved) = parsed.get("model").and_then(|m| m.as_str())
|
||||
{
|
||||
if std::env::var("GOOSE_CLAUDE_CODE_DEBUG").is_ok() {
|
||||
println!(
|
||||
"set_model: {} resolved to {}",
|
||||
process.current_model, resolved
|
||||
);
|
||||
}
|
||||
tracing::debug!(
|
||||
from = %process.current_model,
|
||||
to = %resolved,
|
||||
"set_model resolved"
|
||||
);
|
||||
}
|
||||
process.log_model_update = false;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
lines.push(trimmed.to_string());
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(ProviderError::RequestFailed(format!(
|
||||
@@ -503,57 +533,8 @@ impl ClaudeCodeProvider {
|
||||
}
|
||||
|
||||
tracing::debug!("Command executed successfully, got {} lines", lines.len());
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
tracing::debug!("Line {}: {}", i, line);
|
||||
}
|
||||
|
||||
Ok(lines)
|
||||
}
|
||||
|
||||
/// Generate a simple session description without calling subprocess
|
||||
fn generate_simple_session_description(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Extract the first user message text
|
||||
let description = messages
|
||||
.iter()
|
||||
.find(|m| m.role == Role::User)
|
||||
.and_then(|m| {
|
||||
m.content.iter().find_map(|c| match c {
|
||||
MessageContent::Text(text_content) => Some(&text_content.text),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|text| {
|
||||
// Take first few words, limit to 4 words
|
||||
text.split_whitespace()
|
||||
.take(4)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
})
|
||||
.unwrap_or_else(|| "Simple task".to_string());
|
||||
|
||||
if std::env::var("GOOSE_CLAUDE_CODE_DEBUG").is_ok() {
|
||||
println!("=== CLAUDE CODE PROVIDER DEBUG ===");
|
||||
println!("Generated simple session description: {}", description);
|
||||
println!("Skipped subprocess call for session description");
|
||||
println!("================================");
|
||||
}
|
||||
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
vec![MessageContent::text(description.clone())],
|
||||
);
|
||||
|
||||
let usage = Usage::default();
|
||||
|
||||
Ok((
|
||||
message,
|
||||
ProviderUsage::new(self.model.model_name.clone(), usage),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract model aliases from the CLI's initialize control_response.
|
||||
@@ -708,7 +689,6 @@ impl Provider for ClaudeCodeProvider {
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
// Return the model config with appropriate context limit for Claude models
|
||||
self.model.clone()
|
||||
}
|
||||
|
||||
@@ -785,9 +765,11 @@ impl Provider for ClaudeCodeProvider {
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Check if this is a session description request (short system prompt asking for 4 words or less)
|
||||
if system.contains("four words or less") || system.contains("4 words or less") {
|
||||
return self.generate_simple_session_description(messages);
|
||||
if super::cli_common::is_session_description_request(system) {
|
||||
return super::cli_common::generate_simple_session_description(
|
||||
&model_config.model_name,
|
||||
messages,
|
||||
);
|
||||
}
|
||||
|
||||
// session_id is None before a session is created (e.g. model listing).
|
||||
@@ -798,7 +780,6 @@ impl Provider for ClaudeCodeProvider {
|
||||
|
||||
let (message, usage) = self.parse_claude_response(&json_lines)?;
|
||||
|
||||
// Create a dummy payload for debug tracing
|
||||
let payload = json!({
|
||||
"command": self.command,
|
||||
"model": model_config.model_name,
|
||||
@@ -819,6 +800,172 @@ impl Provider for ClaudeCodeProvider {
|
||||
ProviderUsage::new(model_config.model_name.clone(), usage),
|
||||
))
|
||||
}
|
||||
|
||||
fn supports_streaming(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
session_id: &str,
|
||||
system: &str,
|
||||
messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
if super::cli_common::is_session_description_request(system) {
|
||||
let (message, usage) = super::cli_common::generate_simple_session_description(
|
||||
&self.model.model_name,
|
||||
messages,
|
||||
)?;
|
||||
return Ok(stream_from_single_message(message, usage));
|
||||
}
|
||||
|
||||
let filtered_system = filter_extensions_from_system_prompt(system);
|
||||
let process_arc = Arc::clone(self.get_or_init_process(&filtered_system).await?);
|
||||
|
||||
// Prepare the payload outside the lock — these don't need the process.
|
||||
let blocks = self.last_user_content_blocks(messages);
|
||||
let ndjson_line = build_stream_json_input(&blocks, session_id);
|
||||
let model_name = self.model.model_name.clone();
|
||||
let message_id = uuid::Uuid::new_v4().to_string();
|
||||
|
||||
Ok(Box::pin(try_stream! {
|
||||
// Single lock acquisition covers write-to-stdin and read-from-stdout,
|
||||
// eliminating the race window between the two.
|
||||
let mut process = process_arc.lock_owned().await;
|
||||
process.drain_pending_response().await;
|
||||
process.send_set_model(&model_name).await?;
|
||||
|
||||
process
|
||||
.stdin
|
||||
.write_all(ndjson_line.as_bytes())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ProviderError::RequestFailed(format!("Failed to write to stdin: {}", e))
|
||||
})?;
|
||||
process.stdin.write_all(b"\n").await.map_err(|e| {
|
||||
ProviderError::RequestFailed(format!("Failed to write newline to stdin: {}", e))
|
||||
})?;
|
||||
|
||||
process.needs_drain = true;
|
||||
let mut line = String::new();
|
||||
let mut accumulated_usage = Usage::default();
|
||||
let mut stream_error: Option<ProviderError> = None;
|
||||
let stream_timestamp = chrono::Utc::now().timestamp();
|
||||
|
||||
loop {
|
||||
line.clear();
|
||||
match process.reader.read_line(&mut line).await {
|
||||
Ok(0) => {
|
||||
process.needs_drain = false;
|
||||
stream_error = Some(ProviderError::RequestFailed(
|
||||
"Claude CLI process terminated unexpectedly".to_string(),
|
||||
));
|
||||
break;
|
||||
}
|
||||
Ok(_) => {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
|
||||
match parsed.get("type").and_then(|t| t.as_str()) {
|
||||
Some("stream_event") => {
|
||||
if let Some(event) = parsed.get("event") {
|
||||
match event.get("type").and_then(|t| t.as_str()) {
|
||||
Some("content_block_delta") => {
|
||||
if let Some(text) = event
|
||||
.get("delta")
|
||||
.filter(|d| {
|
||||
d.get("type").and_then(|t| t.as_str())
|
||||
== Some("text_delta")
|
||||
})
|
||||
.and_then(|d| d.get("text"))
|
||||
.and_then(|t| t.as_str())
|
||||
{
|
||||
let mut partial_message = Message::new(
|
||||
Role::Assistant,
|
||||
stream_timestamp,
|
||||
vec![MessageContent::text(text)],
|
||||
);
|
||||
partial_message.id =
|
||||
Some(message_id.clone());
|
||||
yield (Some(partial_message), None);
|
||||
}
|
||||
}
|
||||
Some("message_start") => {
|
||||
if let Some(usage_info) = event
|
||||
.get("message")
|
||||
.and_then(|m| m.get("usage"))
|
||||
{
|
||||
let new = extract_usage_tokens(usage_info);
|
||||
if let Some(i) = new.input_tokens {
|
||||
accumulated_usage.input_tokens = Some(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("message_delta") => {
|
||||
if let Some(usage_info) = event.get("usage") {
|
||||
let new = extract_usage_tokens(usage_info);
|
||||
if let Some(o) = new.output_tokens {
|
||||
accumulated_usage.output_tokens = Some(o);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("result") => {
|
||||
process.needs_drain = false;
|
||||
if let Some(usage_info) = parsed.get("usage") {
|
||||
let new = extract_usage_tokens(usage_info);
|
||||
accumulated_usage = Usage::new(
|
||||
new.input_tokens.or(accumulated_usage.input_tokens),
|
||||
new.output_tokens.or(accumulated_usage.output_tokens),
|
||||
None,
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Some("error") => {
|
||||
process.needs_drain = false;
|
||||
stream_error = Some(error_from_event("Claude CLI", &parsed));
|
||||
break;
|
||||
}
|
||||
Some("system") if process.log_model_update => {
|
||||
if let Some(resolved) = parsed.get("model").and_then(|m| m.as_str()) {
|
||||
tracing::debug!(
|
||||
from = %process.current_model,
|
||||
to = %resolved,
|
||||
"set_model resolved"
|
||||
);
|
||||
}
|
||||
process.log_model_update = false;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
process.needs_drain = false;
|
||||
stream_error = Some(ProviderError::RequestFailed(format!(
|
||||
"Failed to read streaming output: {e}"
|
||||
)));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(err) = stream_error {
|
||||
Err(err)?;
|
||||
}
|
||||
|
||||
let provider_usage = ProviderUsage::new(model_name, accumulated_usage);
|
||||
yield (None, Some(provider_usage));
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -833,6 +980,44 @@ mod tests {
|
||||
use tempfile::tempdir;
|
||||
use test_case::test_case;
|
||||
|
||||
#[test_case(
|
||||
json!({"input_tokens": 100, "output_tokens": 50}),
|
||||
Some(100), Some(50)
|
||||
; "both_tokens"
|
||||
)]
|
||||
#[test_case(json!({"input_tokens": 100}), Some(100), None ; "input_only")]
|
||||
#[test_case(json!({}), None, None ; "empty_usage")]
|
||||
fn test_extract_usage_tokens(
|
||||
usage_json: Value,
|
||||
expected_input: Option<i32>,
|
||||
expected_output: Option<i32>,
|
||||
) {
|
||||
let usage = extract_usage_tokens(&usage_json);
|
||||
assert_eq!(usage.input_tokens, expected_input);
|
||||
assert_eq!(usage.output_tokens, expected_output);
|
||||
}
|
||||
|
||||
#[test_case(
|
||||
r#"{"type":"error","error":"context window exceeded"}"#,
|
||||
true
|
||||
; "context_exceeded"
|
||||
)]
|
||||
#[test_case(
|
||||
r#"{"type":"error","error":"Model not supported"}"#,
|
||||
false
|
||||
; "generic_error_from_event"
|
||||
)]
|
||||
#[test_case(r#"{"type":"error"}"#, false ; "missing_error_field")]
|
||||
fn test_error_from_event(line: &str, is_context_exceeded: bool) {
|
||||
let parsed: Value = serde_json::from_str(line).unwrap();
|
||||
let err = error_from_event("Claude CLI", &parsed);
|
||||
if is_context_exceeded {
|
||||
assert!(matches!(err, ProviderError::ContextLengthExceeded(_)));
|
||||
} else {
|
||||
assert!(matches!(err, ProviderError::RequestFailed(_)));
|
||||
}
|
||||
}
|
||||
|
||||
/// (role, text, optional (image_data, mime_type))
|
||||
type MsgSpec<'a> = (&'a str, &'a str, Option<(&'a str, &'a str)>);
|
||||
|
||||
@@ -958,6 +1143,17 @@ mod tests {
|
||||
Some(3), Some(3)
|
||||
; "system_init_filtered"
|
||||
)]
|
||||
#[test_case(
|
||||
&[
|
||||
r#"{"type":"stream_event","event":{"type":"content_block_delta","delta":{"type":"text_delta","text":"He"}}}"#,
|
||||
r#"{"type":"stream_event","event":{"type":"content_block_delta","delta":{"type":"text_delta","text":"llo"}}}"#,
|
||||
r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}],"usage":{"input_tokens":50,"output_tokens":10}}}"#,
|
||||
r#"{"type":"result","subtype":"success","result":"Hello","session_id":"abc"}"#,
|
||||
],
|
||||
"Hello",
|
||||
Some(50), Some(10)
|
||||
; "streaming_events_ignored_by_parse"
|
||||
)]
|
||||
fn test_parse_claude_response_ok(
|
||||
lines: &[&str],
|
||||
expected_text: &str,
|
||||
@@ -984,7 +1180,7 @@ mod tests {
|
||||
)]
|
||||
#[test_case(
|
||||
&[r#"{"type":"error","error":"context window exceeded"}"#],
|
||||
ProviderError::RequestFailed("Claude CLI error: context window exceeded".into())
|
||||
ProviderError::ContextLengthExceeded("context window exceeded".into())
|
||||
; "context_length"
|
||||
)]
|
||||
#[test_case(
|
||||
@@ -1160,6 +1356,7 @@ mod tests {
|
||||
current_model: String::new(),
|
||||
log_model_update: false,
|
||||
next_request_id: 0,
|
||||
needs_drain: false,
|
||||
};
|
||||
(process, stdin_reader)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use super::base::{ProviderUsage, Usage};
|
||||
use super::errors::ProviderError;
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use rmcp::model::Role;
|
||||
|
||||
pub(crate) fn extract_usage_tokens(usage_info: &Value) -> Usage {
|
||||
let get = |key: &str| {
|
||||
usage_info
|
||||
.get(key)
|
||||
.and_then(|v| v.as_i64())
|
||||
.and_then(|v| i32::try_from(v).ok())
|
||||
};
|
||||
Usage::new(
|
||||
get("input_tokens"),
|
||||
get("output_tokens"),
|
||||
get("total_tokens"),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn error_from_event(provider_name: &str, parsed: &Value) -> ProviderError {
|
||||
let error_msg = parsed
|
||||
.get("error")
|
||||
.and_then(|e| e.as_str())
|
||||
.or_else(|| parsed.get("message").and_then(|m| m.as_str()))
|
||||
.unwrap_or("Unknown error");
|
||||
if error_msg.contains("context window exceeded") {
|
||||
ProviderError::ContextLengthExceeded(error_msg.to_string())
|
||||
} else {
|
||||
ProviderError::RequestFailed(format!("{provider_name} error: {error_msg}"))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_session_description_request(system: &str) -> bool {
|
||||
system.contains("four words or less") || system.contains("4 words or less")
|
||||
}
|
||||
|
||||
pub(crate) fn generate_simple_session_description(
|
||||
model_name: &str,
|
||||
messages: &[Message],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
let description = messages
|
||||
.iter()
|
||||
.find(|m| m.role == Role::User)
|
||||
.and_then(|m| {
|
||||
m.content.iter().find_map(|c| match c {
|
||||
MessageContent::Text(text_content) => Some(&text_content.text),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|text| {
|
||||
text.split_whitespace()
|
||||
.take(4)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
})
|
||||
.unwrap_or_else(|| "Simple task".to_string());
|
||||
|
||||
tracing::debug!(
|
||||
description = %description,
|
||||
"Generated simple session description, skipped subprocess"
|
||||
);
|
||||
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
vec![MessageContent::text(description)],
|
||||
);
|
||||
|
||||
Ok((
|
||||
message,
|
||||
ProviderUsage::new(model_name.to_string(), Usage::default()),
|
||||
))
|
||||
}
|
||||
@@ -409,51 +409,6 @@ impl CodexProvider {
|
||||
|
||||
Ok((message, usage))
|
||||
}
|
||||
|
||||
/// Generate a simple session description without calling subprocess
|
||||
fn generate_simple_session_description(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Extract the first user message text
|
||||
let description = messages
|
||||
.iter()
|
||||
.find(|m| m.role == Role::User)
|
||||
.and_then(|m| {
|
||||
m.content.iter().find_map(|c| match c {
|
||||
MessageContent::Text(text_content) => Some(&text_content.text),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|text| {
|
||||
// Take first few words, limit to 4 words
|
||||
text.split_whitespace()
|
||||
.take(4)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
})
|
||||
.unwrap_or_else(|| "Simple task".to_string());
|
||||
|
||||
if std::env::var("GOOSE_CODEX_DEBUG").is_ok() {
|
||||
println!("=== CODEX PROVIDER DEBUG ===");
|
||||
println!("Generated simple session description: {}", description);
|
||||
println!("Skipped subprocess call for session description");
|
||||
println!("============================");
|
||||
}
|
||||
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
vec![MessageContent::text(description.clone())],
|
||||
);
|
||||
|
||||
let usage = Usage::default();
|
||||
|
||||
Ok((
|
||||
message,
|
||||
ProviderUsage::new(self.model.model_name.clone(), usage),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds the text prompt and extracts images to temp files in a single pass.
|
||||
@@ -724,9 +679,11 @@ impl Provider for CodexProvider {
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Check if this is a session description request
|
||||
if system.contains("four words or less") || system.contains("4 words or less") {
|
||||
return self.generate_simple_session_description(messages);
|
||||
if super::cli_common::is_session_description_request(system) {
|
||||
return super::cli_common::generate_simple_session_description(
|
||||
&model_config.model_name,
|
||||
messages,
|
||||
);
|
||||
}
|
||||
|
||||
let lines = self.execute_command(system, messages, tools).await?;
|
||||
@@ -1192,15 +1149,6 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_session_description_generation() {
|
||||
let provider = CodexProvider {
|
||||
command: PathBuf::from("codex"),
|
||||
model: ModelConfig::new("gpt-5.2-codex").unwrap(),
|
||||
name: "codex".to_string(),
|
||||
reasoning_effort: "high".to_string(),
|
||||
skip_git_check: false,
|
||||
mcp_config_overrides: Vec::new(),
|
||||
};
|
||||
|
||||
let messages = vec![Message::new(
|
||||
Role::User,
|
||||
chrono::Utc::now().timestamp(),
|
||||
@@ -1209,12 +1157,15 @@ mod tests {
|
||||
)],
|
||||
)];
|
||||
|
||||
let result = provider.generate_simple_session_description(&messages);
|
||||
let result = crate::providers::cli_common::generate_simple_session_description(
|
||||
"gpt-5.2-codex",
|
||||
&messages,
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (message, _usage) = result.unwrap();
|
||||
let (message, usage) = result.unwrap();
|
||||
assert_eq!(usage.model, "gpt-5.2-codex");
|
||||
if let MessageContent::Text(text) = &message.content[0] {
|
||||
// Should be truncated to 4 words
|
||||
let word_count = text.text.split_whitespace().count();
|
||||
assert!(word_count <= 4);
|
||||
} else {
|
||||
@@ -1224,18 +1175,12 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_session_description_empty_messages() {
|
||||
let provider = CodexProvider {
|
||||
command: PathBuf::from("codex"),
|
||||
model: ModelConfig::new("gpt-5.2-codex").unwrap(),
|
||||
name: "codex".to_string(),
|
||||
reasoning_effort: "high".to_string(),
|
||||
skip_git_check: false,
|
||||
mcp_config_overrides: Vec::new(),
|
||||
};
|
||||
|
||||
let messages: Vec<Message> = vec![];
|
||||
|
||||
let result = provider.generate_simple_session_description(&messages);
|
||||
let result = crate::providers::cli_common::generate_simple_session_description(
|
||||
"gpt-5.2-codex",
|
||||
&messages,
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let (message, _usage) = result.unwrap();
|
||||
|
||||
@@ -275,51 +275,6 @@ impl CursorAgentProvider {
|
||||
|
||||
Ok(lines)
|
||||
}
|
||||
|
||||
/// Generate a simple session description without calling subprocess
|
||||
fn generate_simple_session_description(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Extract the first user message text
|
||||
let description = messages
|
||||
.iter()
|
||||
.find(|m| m.role == Role::User)
|
||||
.and_then(|m| {
|
||||
m.content.iter().find_map(|c| match c {
|
||||
MessageContent::Text(text_content) => Some(&text_content.text),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|text| {
|
||||
// Take first few words, limit to 4 words
|
||||
text.split_whitespace()
|
||||
.take(4)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
})
|
||||
.unwrap_or_else(|| "Simple task".to_string());
|
||||
|
||||
if std::env::var("GOOSE_CURSOR_AGENT_DEBUG").is_ok() {
|
||||
println!("=== CURSOR AGENT PROVIDER DEBUG ===");
|
||||
println!("Generated simple session description: {}", description);
|
||||
println!("Skipped subprocess call for session description");
|
||||
println!("================================");
|
||||
}
|
||||
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
vec![MessageContent::text(description.clone())],
|
||||
);
|
||||
|
||||
let usage = Usage::default();
|
||||
|
||||
Ok((
|
||||
message,
|
||||
ProviderUsage::new(self.model.model_name.clone(), usage),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for CursorAgentProvider {
|
||||
@@ -378,9 +333,11 @@ impl Provider for CursorAgentProvider {
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
// Check if this is a session description request (short system prompt asking for 4 words or less)
|
||||
if system.contains("four words or less") || system.contains("4 words or less") {
|
||||
return self.generate_simple_session_description(messages);
|
||||
if super::cli_common::is_session_description_request(system) {
|
||||
return super::cli_common::generate_simple_session_description(
|
||||
&model_config.model_name,
|
||||
messages,
|
||||
);
|
||||
}
|
||||
|
||||
let lines = self.execute_command(system, messages, tools).await?;
|
||||
|
||||
@@ -8,6 +8,7 @@ use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
|
||||
use super::base::{Provider, ProviderDef, ProviderMetadata, ProviderUsage, Usage};
|
||||
use super::cli_common::{error_from_event, extract_usage_tokens};
|
||||
use super::errors::ProviderError;
|
||||
use super::utils::{filter_extensions_from_system_prompt, RequestLog};
|
||||
use crate::config::base::GeminiCliCommand;
|
||||
@@ -31,29 +32,6 @@ pub const GEMINI_CLI_KNOWN_MODELS: &[&str] = &[
|
||||
|
||||
pub const GEMINI_CLI_DOC_URL: &str = "https://ai.google.dev/gemini-api/docs";
|
||||
|
||||
fn extract_usage_from_stats(stats: &Value) -> Usage {
|
||||
let get = |key: &str| {
|
||||
stats
|
||||
.get(key)
|
||||
.and_then(|v| v.as_i64())
|
||||
.and_then(|v| i32::try_from(v).ok())
|
||||
};
|
||||
Usage::new(
|
||||
get("input_tokens"),
|
||||
get("output_tokens"),
|
||||
get("total_tokens"),
|
||||
)
|
||||
}
|
||||
|
||||
fn error_from_event(parsed: &Value) -> ProviderError {
|
||||
let error_msg = parsed
|
||||
.get("error")
|
||||
.and_then(|e| e.as_str())
|
||||
.or_else(|| parsed.get("message").and_then(|m| m.as_str()))
|
||||
.unwrap_or("Unknown error");
|
||||
ProviderError::RequestFailed(format!("Gemini CLI error: {error_msg}"))
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
pub struct GeminiCliProvider {
|
||||
command: PathBuf,
|
||||
@@ -95,48 +73,6 @@ impl GeminiCliProvider {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn is_session_description_request(system: &str) -> bool {
|
||||
system.contains("four words or less") || system.contains("4 words or less")
|
||||
}
|
||||
|
||||
fn generate_simple_session_description(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
let description = messages
|
||||
.iter()
|
||||
.find(|m| m.role == Role::User)
|
||||
.and_then(|m| {
|
||||
m.content.iter().find_map(|c| match c {
|
||||
MessageContent::Text(text_content) => Some(&text_content.text),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|text| {
|
||||
text.split_whitespace()
|
||||
.take(4)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
})
|
||||
.unwrap_or_else(|| "Simple task".to_string());
|
||||
|
||||
tracing::debug!(
|
||||
description = %description,
|
||||
"Generated simple session description, skipped subprocess"
|
||||
);
|
||||
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
vec![MessageContent::text(description)],
|
||||
);
|
||||
|
||||
Ok((
|
||||
message,
|
||||
ProviderUsage::new(self.model.model_name.clone(), Usage::default()),
|
||||
))
|
||||
}
|
||||
|
||||
/// Build the prompt for the CLI invocation. When resuming a session the CLI
|
||||
/// maintains conversation context internally, so only the latest user
|
||||
/// message is needed. On the first turn (no session yet) the system prompt
|
||||
@@ -300,11 +236,11 @@ impl GeminiCliProvider {
|
||||
}
|
||||
Some("result") => {
|
||||
if let Some(stats) = parsed.get("stats") {
|
||||
usage = extract_usage_from_stats(stats);
|
||||
usage = extract_usage_tokens(stats);
|
||||
}
|
||||
}
|
||||
Some("error") => {
|
||||
return Err(error_from_event(parsed));
|
||||
return Err(error_from_event("Gemini CLI", parsed));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
@@ -380,8 +316,11 @@ impl Provider for GeminiCliProvider {
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
if Self::is_session_description_request(system) {
|
||||
return self.generate_simple_session_description(messages);
|
||||
if super::cli_common::is_session_description_request(system) {
|
||||
return super::cli_common::generate_simple_session_description(
|
||||
&model_config.model_name,
|
||||
messages,
|
||||
);
|
||||
}
|
||||
|
||||
let payload = json!({
|
||||
|
||||
@@ -8,6 +8,7 @@ pub mod bedrock;
|
||||
pub mod canonical;
|
||||
pub mod chatgpt_codex;
|
||||
pub mod claude_code;
|
||||
pub(crate) mod cli_common;
|
||||
pub mod codex;
|
||||
pub mod cursor_agent;
|
||||
pub mod databricks;
|
||||
|
||||
Reference in New Issue
Block a user