feat: add streaming support for Claude Code CLI provider (#6833)

Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-02-13 20:58:32 +05:30
committed by GitHub
parent 6a7adb4e94
commit ed3c08590a
6 changed files with 486 additions and 372 deletions
+382 -185
View File
@@ -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)
}
+75
View File
@@ -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()),
))
}
+15 -70
View File
@@ -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();
+5 -48
View File
@@ -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 -69
View File
@@ -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!({
+1
View File
@@ -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;