diff --git a/Cargo.lock b/Cargo.lock index bca11447c..96c855b34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5333,11 +5333,14 @@ version = "0.1.0-alpha.6" dependencies = [ "anyhow", "async-trait", + "futures", "goose-provider-types", "rmcp 3.1.4", "serde_json", "tokio", "tokio-util", + "tracing", + "tracing-futures", "uuid", ] diff --git a/crates/goose-agent/Cargo.toml b/crates/goose-agent/Cargo.toml index ab7e88964..fdc988ad3 100644 --- a/crates/goose-agent/Cargo.toml +++ b/crates/goose-agent/Cargo.toml @@ -20,3 +20,6 @@ goose-provider-types = { version = "0.1.0-alpha.6", path = "../goose-provider-ty rmcp = { workspace = true } uuid = { workspace = true } serde_json = { workspace = true } +futures = { workspace = true } +tracing = { workspace = true } +tracing-futures = { workspace = true } diff --git a/crates/goose-agent/src/inference.rs b/crates/goose-agent/src/inference.rs new file mode 100644 index 000000000..4551c492b --- /dev/null +++ b/crates/goose-agent/src/inference.rs @@ -0,0 +1,600 @@ +//! Provider inference operation for the unrolled agent loop. + +use std::sync::Arc; + +use anyhow::Result; +use async_trait::async_trait; +use futures::StreamExt; +use goose_provider_types::base::Provider; +use goose_provider_types::conversation::message::{InferenceMetadata, Message, MessageContent}; +use goose_provider_types::conversation::token_usage::ProviderUsage; +use goose_provider_types::conversation::{ + effective_role, fix_conversation, merge_consecutive_messages_for_request, Conversation, + EffectiveRole, +}; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::model::ModelConfig; +use tracing_futures::Instrument; + +use crate::operation::{ + applied, messages_since_kickoff, not_applicable, trailing_error, yielded_with, Emitter, + Inference, InferenceInput, Operation, OperationResult, +}; + +pub struct PreparedInferenceRequest { + pub system_prompt: String, + pub tools: Vec, + pub additional_messages: Vec, +} + +#[async_trait] +pub trait InferenceRequestPreparer: Send + Sync { + async fn prepare( + &self, + session: &S, + conversation: &Conversation, + input: InferenceInput, + ) -> Result; +} + +pub struct IdentityInferenceRequestPreparer; + +#[async_trait] +impl InferenceRequestPreparer for IdentityInferenceRequestPreparer { + async fn prepare( + &self, + _session: &S, + _conversation: &Conversation, + input: InferenceInput, + ) -> Result { + Ok(PreparedInferenceRequest { + system_prompt: input + .prompt_parts + .into_iter() + .map(|(_, part)| part) + .collect::>() + .join("\n\n"), + tools: input.tools, + additional_messages: Vec::new(), + }) + } +} + +pub trait InferenceEffect: From + Send + 'static { + fn record_usage(usage: ProviderUsage) -> Self; +} + +const EMPTY_RESPONSE_MESSAGE: &str = + "The model returned an empty response. Please resend your message to continue."; +const CANCELLED_TOOL_RESPONSE: &str = "Tool call was cancelled before execution"; + +fn is_thinking(content: &MessageContent) -> bool { + matches!( + content, + MessageContent::Thinking(_) | MessageContent::RedactedThinking(_) + ) +} + +fn normalize_tool_call_thinking(accumulator: &mut Conversation, chunk: &mut Message) { + if !chunk + .content + .iter() + .any(|content| matches!(content, MessageContent::ToolRequest(_))) + { + return; + } + + let has_direct_thinking = chunk.content.iter().any(is_thinking); + let mut prior_thinking = Vec::new(); + for message in accumulator.messages_mut() { + if message.role != chunk.role + || message + .content + .iter() + .any(|content| matches!(content, MessageContent::ToolRequest(_))) + { + continue; + } + prior_thinking.extend( + message + .content + .iter() + .filter(|content| is_thinking(content)) + .cloned(), + ); + message.content.retain(|content| !is_thinking(content)); + } + accumulator + .messages_mut() + .retain(|message| !message.content.is_empty()); + + if !has_direct_thinking && !prior_thinking.is_empty() { + if let Some(tool_request) = chunk + .content + .iter() + .position(|content| matches!(content, MessageContent::ToolRequest(_))) + { + chunk + .content + .splice(tool_request..tool_request, prior_thinking); + } + } +} + +pub fn chat_span( + provider: &dyn Provider, + model_config: &ModelConfig, + session_id: &str, + purpose: &'static str, +) -> tracing::Span { + let span = tracing::info_span!( + target: "goose::state_machine", + "chat", + "gen_ai.operation.name" = "chat", + "gen_ai.provider.name" = %provider.get_name(), + "gen_ai.request.model" = %model_config.model_name, + "gen_ai.request.temperature" = tracing::field::Empty, + "gen_ai.request.max_tokens" = tracing::field::Empty, + "gen_ai.response.model" = tracing::field::Empty, + "gen_ai.response.finish_reasons" = tracing::field::Empty, + "gen_ai.response.id" = tracing::field::Empty, + "gen_ai.usage.input_tokens" = tracing::field::Empty, + "gen_ai.usage.output_tokens" = tracing::field::Empty, + "goose.chat.purpose" = purpose, + "error.type" = tracing::field::Empty, + session.id = %session_id, + ); + record_request_params(&span, model_config); + span +} + +fn record_request_params(span: &tracing::Span, model_config: &ModelConfig) { + if let Some(temperature) = model_config.temperature { + span.record("gen_ai.request.temperature", temperature as f64); + } + if let Some(max_tokens) = model_config.max_tokens { + span.record("gen_ai.request.max_tokens", max_tokens as i64); + } +} + +pub fn record_chat_usage(span: &tracing::Span, usage: &ProviderUsage) { + span.record("gen_ai.response.model", usage.model.as_str()); + if let Some(tokens) = usage.usage.input_tokens { + span.record("gen_ai.usage.input_tokens", tokens); + } + if let Some(tokens) = usage.usage.output_tokens { + span.record("gen_ai.usage.output_tokens", tokens); + } + if let Some(tokens) = usage.usage.cache_read_input_tokens { + span.record("gen_ai.usage.cache_read.input_tokens", tokens); + } + if let Some(tokens) = usage.usage.cache_write_input_tokens { + span.record("gen_ai.usage.cache_creation.input_tokens", tokens); + } + if let Some(reasons) = &usage.finish_reasons { + let reasons_json = serde_json::to_string(reasons).unwrap_or_default(); + span.record("gen_ai.response.finish_reasons", reasons_json.as_str()); + } + if let Some(id) = &usage.response_id { + span.record("gen_ai.response.id", id.as_str()); + } +} + +pub struct InferenceRunner<'a, S, E> { + provider: Arc, + model_config: ModelConfig, + request_preparer: Arc + 'a>, + effect: std::marker::PhantomData E>, +} + +/// The agent-visible conversation as the provider sees it: tool requests left +/// unanswered by an earlier turn are dropped, since nothing will answer them now. +fn messages_for_provider(conversation: &Conversation, turn: &[Message]) -> Vec { + let answered: std::collections::HashSet<&str> = conversation + .messages() + .iter() + .flat_map(|message| message.get_tool_response_ids()) + .collect(); + let start = conversation.len() - turn.len(); + conversation + .messages() + .iter() + .enumerate() + .filter(|(_, message)| message.is_agent_visible()) + .map(|(index, message)| { + let mut message = message.agent_visible_content(); + if index < start { + message.content.retain(|content| match content { + MessageContent::ToolRequest(request) => answered.contains(request.id.as_str()), + _ => true, + }); + } + message + }) + .filter(|message| !message.content.is_empty()) + .collect() +} + +fn latest_provider_session_id<'a>( + conversation: &'a Conversation, + provider: &str, +) -> Option<&'a str> { + conversation + .messages() + .iter() + .rev() + .find_map(|message| message.metadata.inference.as_ref()) + .filter(|inference| inference.provider == provider) + .and_then(|inference| inference.provider_session_id.as_deref()) +} + +fn ends_with_provider_turn(messages: &[Message]) -> bool { + messages.last().is_some_and(|message| { + matches!( + effective_role(message), + EffectiveRole::User | EffectiveRole::Tool + ) + }) +} + +fn cancellation_response(persisted: &[Message], pending: &[Message]) -> Option { + let mut answered = persisted + .iter() + .chain(pending) + .flat_map(Message::get_tool_response_ids) + .collect::>(); + let mut request_ids = std::collections::HashSet::new(); + let mut response = Message::user(); + for request in persisted + .iter() + .chain(pending) + .flat_map(|message| &message.content) + .filter_map(MessageContent::as_tool_request) + { + if request_ids.insert(request.id.as_str()) && !answered.remove(request.id.as_str()) { + response.add_tool_response_with_metadata( + request.id.clone(), + Ok(rmcp::model::CallToolResult::error(vec![ + rmcp::model::ContentBlock::text(CANCELLED_TOOL_RESPONSE), + ])), + request.metadata.as_ref(), + ); + } + } + (!response.get_tool_response_ids().is_empty()).then_some(response) +} + +fn inference_span(provider: &dyn Provider, model_config: &ModelConfig) -> tracing::Span { + let span = tracing::info_span!( + target: "goose::state_machine", + "chat", + "gen_ai.operation.name" = "chat", + "gen_ai.provider.name" = %provider.get_name(), + "gen_ai.request.model" = %model_config.model_name, + "gen_ai.request.temperature" = tracing::field::Empty, + "gen_ai.request.max_tokens" = tracing::field::Empty, + "gen_ai.response.model" = tracing::field::Empty, + "gen_ai.response.finish_reasons" = tracing::field::Empty, + "gen_ai.response.id" = tracing::field::Empty, + "gen_ai.usage.input_tokens" = tracing::field::Empty, + "gen_ai.usage.output_tokens" = tracing::field::Empty, + "error.type" = tracing::field::Empty, + ); + record_request_params(&span, model_config); + span +} + +impl<'a, S: Sync, E: InferenceEffect> InferenceRunner<'a, S, E> { + pub fn new(provider: Arc, model_config: ModelConfig) -> Self { + Self { + provider, + model_config, + request_preparer: Arc::new(IdentityInferenceRequestPreparer), + effect: std::marker::PhantomData, + } + } + + pub fn with_request_preparer( + mut self, + request_preparer: Arc + 'a>, + ) -> Self { + self.request_preparer = request_preparer; + self + } + + async fn error_outcome(&self, err: &ProviderError, emit: &Emitter) -> Vec { + tracing::Span::current().record("error.type", err.telemetry_type()); + tracing::error!("LLM provider error: {err}"); + let message = Message::from_provider_error(err); + let message = emit.message(message).await; + vec![E::from(message)] + } +} + +#[async_trait] +impl Operation for InferenceRunner<'_, S, E> { + fn name(&self) -> &'static str { + "llm" + } + + async fn cancel( + &self, + _session: &S, + conversation: &Conversation, + result: OperationResult, + emit: &Emitter, + ) -> Result> { + let OperationResult::NotApplicable = result else { + return Ok(result); + }; + let Some(response) = cancellation_response(messages_since_kickoff(conversation)?, &[]) + else { + return Ok(OperationResult::NotApplicable); + }; + let response = emit.message(response).await; + applied([E::from(response)]) + } +} + +#[async_trait] +impl Inference for InferenceRunner<'_, S, E> { + fn applies(&self, conversation: &Conversation) -> bool { + let Ok(turn) = messages_since_kickoff(conversation) else { + return false; + }; + trailing_error(conversation).is_none() + && ends_with_provider_turn(&messages_for_provider(conversation, turn)) + } + + async fn infer( + &self, + session: &S, + conversation: &Conversation, + input: InferenceInput, + emit: &Emitter, + ) -> Result> { + let messages = messages_since_kickoff(conversation)?; + if trailing_error(conversation).is_some() { + return not_applicable(); + } + + let mut messages_for_provider = messages_for_provider(conversation, messages); + if !ends_with_provider_turn(&messages_for_provider) { + return not_applicable(); + } + + let span = inference_span(self.provider.as_ref(), &self.model_config); + + async { + let PreparedInferenceRequest { + system_prompt, + tools, + additional_messages, + } = self + .request_preparer + .prepare(session, conversation, input) + .await?; + + for message in &additional_messages { + messages_for_provider.push(message.clone()); + } + let mut usage_effects: Vec = additional_messages.into_iter().map(E::from).collect(); + + let provider_name = self.provider.get_name(); + if let Some(session_id) = latest_provider_session_id(conversation, provider_name) { + if let Err(error) = self.provider.resume(session_id).await { + tracing::warn!( + provider = provider_name, + %error, + "Could not resume provider session; continuing with a handoff" + ); + } + } + + let projected = + Conversation::new_unvalidated(messages_for_provider).agent_visible_messages(); + let (fixed, _) = fix_conversation(Conversation::new_unvalidated(projected)); + let conversation_for_provider = Conversation::new_unvalidated( + merge_consecutive_messages_for_request(fixed.messages().clone()), + ); + let stream = self + .provider + .stream( + &self.model_config, + &system_prompt, + conversation_for_provider.messages(), + &tools, + ) + .await; + + let mut stream = match stream { + Ok(stream) => stream, + Err(err) => { + usage_effects.extend(self.error_outcome(&err, emit).await); + return applied(usage_effects); + } + }; + + let requested_model = self.model_config.model_name.clone(); + let resolved_model = self + .provider + .fetch_model_info(&requested_model) + .await + .ok() + .and_then(|model_info| model_info.resolved_model); + let provider_session_id = self.provider.provider_session_id(); + let inference = Some(InferenceMetadata { + provider: self.provider.get_name().to_string(), + requested_model, + resolved_model, + provider_session_id, + }); + + let mut accumulator = Conversation::empty(); + let mut tool_request_ids = std::collections::HashSet::new(); + let mut provider_usage = None; + let mut cancelled = false; + loop { + tokio::select! { + biased; + _ = emit.cancelled() => { + cancelled = true; + break; + }, + next = stream.next() => { + let Some(result) = next else { break }; + let (msg_opt, usage_opt) = match result { + Ok(chunk) => chunk, + Err(err) => { + if let Some(usage) = provider_usage { + usage_effects.push(E::record_usage(usage)); + } + usage_effects.extend(accumulator.into_iter().map(E::from)); + usage_effects.extend(self.error_outcome(&err, emit).await); + return applied(usage_effects); + } + }; + if let Some(usage) = usage_opt { + let span = tracing::Span::current(); + record_chat_usage(&span, &usage); + provider_usage = Some(usage); + } + if let Some(mut chunk) = msg_opt { + if let Some(inference) = &inference { + chunk = chunk.with_inference_if_assistant(inference.clone()); + } + chunk.content.retain(|content| match content { + MessageContent::ToolRequest(request) => { + tool_request_ids.insert(request.id.clone()) + } + _ => true, + }); + normalize_tool_call_thinking(&mut accumulator, &mut chunk); + if chunk.content.is_empty() { + if chunk.metadata.output_token_limit_reached { + chunk = emit.message(chunk).await; + } + accumulator.push(chunk); + continue; + } + let chunk = emit.message(chunk).await; + accumulator.push(chunk); + } + } + } + } + + if let Some(usage) = provider_usage { + usage_effects.push(E::record_usage(usage)); + } + + if cancelled || emit.cancel_token().is_cancelled() { + if let Some(response) = cancellation_response(messages, accumulator.messages()) { + let response = emit.message(response).await; + accumulator.push(response); + } + } + + let empty_response = !cancelled + && !accumulator + .iter() + .any(|message| message.metadata.output_token_limit_reached) + && accumulator.iter().all(|message| { + message.content.iter().all(|content| match content { + MessageContent::Text(text) => text.text.trim().is_empty(), + MessageContent::Thinking(thinking) => thinking.thinking.trim().is_empty(), + _ => false, + }) + }); + if empty_response { + let message = Message::assistant().with_text(EMPTY_RESPONSE_MESSAGE); + let message = emit.message(message).await; + usage_effects.push(E::from(message)); + return yielded_with(usage_effects); + } + + usage_effects.extend(accumulator.into_iter().map(|message| E::from(message))); + applied(usage_effects) + } + .instrument(span) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn provider_session_id_comes_only_from_latest_inference() { + let conversation = Conversation::new_unvalidated([ + Message::assistant().with_inference(InferenceMetadata { + provider: "provider-a".to_string(), + requested_model: "model".to_string(), + resolved_model: None, + provider_session_id: Some("session-a".to_string()), + }), + Message::assistant().with_inference(InferenceMetadata { + provider: "provider-b".to_string(), + requested_model: "model".to_string(), + resolved_model: None, + provider_session_id: Some("session-b".to_string()), + }), + ]); + + assert_eq!( + latest_provider_session_id(&conversation, "provider-b"), + Some("session-b") + ); + assert_eq!( + latest_provider_session_id(&conversation, "provider-a"), + None + ); + } + + #[test] + fn cancellation_response_includes_requests_from_unconverted_messages() { + let persisted = [Message::user().with_text("run it")]; + let pending = [Message::assistant().with_tool_request( + "pending-call", + Ok(rmcp::model::CallToolRequestParams::new("tool")), + )]; + + let response = cancellation_response(&persisted, &pending).expect("cancellation response"); + + assert_eq!( + response.get_tool_response_ids(), + std::collections::HashSet::from(["pending-call"]) + ); + let cancellation_text = response + .content + .iter() + .filter_map(MessageContent::as_tool_response) + .flat_map(|response| { + response + .tool_result + .as_ref() + .expect("tool result") + .content + .iter() + }) + .filter_map(|content| content.as_text()) + .map(|text| text.text.as_str()) + .collect::>(); + assert_eq!(cancellation_text, vec![CANCELLED_TOOL_RESPONSE]); + } + + #[test] + fn cancellation_response_skips_answered_requests() { + let request = Message::assistant().with_tool_request( + "answered-call", + Ok(rmcp::model::CallToolRequestParams::new("tool")), + ); + let response = Message::user().with_tool_response( + "answered-call", + Ok(rmcp::model::CallToolResult::success(vec![])), + ); + + assert!(cancellation_response(&[request, response], &[]).is_none()); + } +} diff --git a/crates/goose-agent/src/lib.rs b/crates/goose-agent/src/lib.rs index 1336ff9a9..b1f818333 100644 --- a/crates/goose-agent/src/lib.rs +++ b/crates/goose-agent/src/lib.rs @@ -1,3 +1,4 @@ pub mod events; +pub mod inference; pub mod machine; pub mod operation; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 42794f68b..9406aff02 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -31,11 +31,11 @@ use crate::agents::prompt_manager::PromptManager; use crate::agents::retry::{RetryManager, RetryResult}; use crate::agents::state_machine::{ run_goose, BangShellOperation, CompactionOperation, DoctorOperation, Emitter, - EntryHookOperation, ExitOnErrorOperation, GooseEffect, InferenceRunner, MaxTurnsOperation, - Operation, ProjectOperation, RecipeOperation, RetryOperation, SkillOperation, - SlashCommandOperation, StateMachine, SteerOperation, SteerQueue, Step, StopHookOperation, - ToolApprovalOperation, ToolExecutionOperation, ToolPairCompactionOperation, - UnknownToolOperation, MAX_TURNS_MESSAGE, + EntryHookOperation, ExitOnErrorOperation, GooseEffect, GooseInferenceProvider, + GooseInferenceRequestPreparer, InferenceRunner, MaxTurnsOperation, Operation, ProjectOperation, + RecipeOperation, RetryOperation, SkillOperation, SlashCommandOperation, StateMachine, + StatusOperation, SteerOperation, SteerQueue, Step, StopHookOperation, ToolApprovalOperation, + ToolExecutionOperation, ToolPairCompactionOperation, UnknownToolOperation, MAX_TURNS_MESSAGE, }; use crate::agents::types::{ FrontendTool, SessionConfig, SharedProvider, ToolResultReceiver, @@ -1750,17 +1750,24 @@ impl Agent { Arc::new(ExitOnErrorOperation), ]; operations.extend(remaining_operations); - let inference = Arc::new(InferenceRunner::new( - provider, - model_config, - self.extension_manager.clone(), - &self.current_goose_mode, - &self.prompt_manager, - &self.tool_inspection_manager, - &self.frontend_instructions, - )); + let request_preparer = GooseInferenceRequestPreparer { + #[cfg(feature = "code-mode")] + extension_manager: self.extension_manager.clone(), + goose_mode: &self.current_goose_mode, + prompt_manager: &self.prompt_manager, + tool_inspection_manager: &self.tool_inspection_manager, + frontend_instructions: &self.frontend_instructions, + context_limit, + }; + let status_operation = + Arc::new(StatusOperation::new(provider.clone(), model_config.clone())); + let inference_provider = Arc::new(GooseInferenceProvider::new(provider)); + let inference = Arc::new( + InferenceRunner::new(inference_provider, model_config) + .with_request_preparer(Arc::new(request_preparer)), + ); let mut command_handlers = operations.clone(); - command_handlers.push(inference.clone()); + command_handlers.push(status_operation); let command_operation: Arc + '_> = Arc::new(SlashCommandOperation::new(command_handlers)); let operations: Vec<_> = @@ -1862,7 +1869,10 @@ impl Agent { let (tx, mut rx) = mpsc::channel::(32); let emit = Emitter::new(tx, cancel.clone()); let result = { - let run = run_goose(&machine, session_manager.as_ref(), &session_id, &emit); + let run = crate::session_context::with_session_id( + Some(session_id.clone()), + run_goose(&machine, session_manager.as_ref(), &session_id, &emit), + ); tokio::pin!(run); loop { tokio::select! { diff --git a/crates/goose/src/agents/state_machine/inference_preparation.rs b/crates/goose/src/agents/state_machine/inference_preparation.rs new file mode 100644 index 000000000..c976914d6 --- /dev/null +++ b/crates/goose/src/agents/state_machine/inference_preparation.rs @@ -0,0 +1,91 @@ +//! Goose-specific inference request preparation. + +#[cfg(feature = "code-mode")] +use crate::agents::ExtensionManager; +use crate::agents::PromptManager; +use crate::config::GooseMode; +use crate::session::Session; +use crate::tool_inspection::ToolInspectionManager; +use anyhow::Result; +use async_trait::async_trait; +use goose_agent::inference::{InferenceRequestPreparer, PreparedInferenceRequest}; +use goose_agent::operation::{messages_since_kickoff, InferenceInput}; +use goose_providers::conversation::message::Message; +use goose_providers::conversation::Conversation; +#[cfg(feature = "code-mode")] +use std::sync::Arc; +use tokio::sync::Mutex; + +pub struct GooseInferenceRequestPreparer<'a> { + #[cfg(feature = "code-mode")] + pub(crate) extension_manager: Arc, + pub(crate) goose_mode: &'a Mutex, + pub(crate) prompt_manager: &'a Mutex, + pub(crate) tool_inspection_manager: &'a ToolInspectionManager, + pub(crate) frontend_instructions: &'a Mutex>, + pub(crate) context_limit: usize, +} + +#[async_trait] +impl InferenceRequestPreparer for GooseInferenceRequestPreparer<'_> { + async fn prepare( + &self, + session: &Session, + conversation: &Conversation, + mut input: InferenceInput, + ) -> Result { + #[cfg(feature = "code-mode")] + let code_execution_mode = self + .extension_manager + .is_extension_enabled( + crate::agents::platform_extensions::code_execution::EXTENSION_NAME, + ) + .await; + #[cfg(not(feature = "code-mode"))] + let code_execution_mode = false; + + let goose_mode = *self.goose_mode.lock().await; + if goose_mode == GooseMode::SmartApprove { + self.tool_inspection_manager + .apply_tool_annotations(&input.tools); + } + let tools = + crate::agents::reply_parts::prepare_inference_tools(input.tools, code_execution_mode); + if let Some(frontend_instructions) = self.frontend_instructions.lock().await.clone() { + input + .prompt_parts + .push(("frontend".to_string(), frontend_instructions)); + } + let system_prompt = self.prompt_manager.lock().await.build_system_prompt( + &session.working_dir, + input.prompt_parts, + goose_mode, + ); + let turn = messages_since_kickoff(conversation)?; + let turn_start = turn + .first() + .and_then(|message| chrono::DateTime::from_timestamp(message.created, 0)) + .map(|timestamp| timestamp.with_timezone(&chrono::Local)) + .unwrap_or_else(chrono::Local::now); + let last = turn + .iter() + .rev() + .find(|message| message.is_turn_context()) + .map(Message::as_concat_text); + let context_limit = Some(self.context_limit); + let additional_messages = crate::agents::moim::turn_context_event( + &session.working_dir, + context_limit, + input.moim_parts, + turn_start, + ) + .filter(|event| Some(event.as_concat_text()) != last) + .into_iter() + .collect(); + Ok(PreparedInferenceRequest { + system_prompt, + tools, + additional_messages, + }) + } +} diff --git a/crates/goose/src/agents/state_machine/mod.rs b/crates/goose/src/agents/state_machine/mod.rs index fe6825c5a..802f4cef8 100644 --- a/crates/goose/src/agents/state_machine/mod.rs +++ b/crates/goose/src/agents/state_machine/mod.rs @@ -6,6 +6,7 @@ //! configuration is part of `Agent::reply`, not the state-machine protocol. mod effects; +mod inference_preparation; mod ops_bang_shell; mod ops_compaction; mod ops_doctor; @@ -18,6 +19,7 @@ mod ops_recipe; mod ops_retry; mod ops_skills; mod ops_slash_command; +mod ops_status; mod ops_steer; mod ops_stop_hook; mod ops_tool_approval; @@ -41,18 +43,20 @@ pub use goose_agent::operation::{ InferenceInput, MachineEffect, Operation, OperationResult, SlashCommand, StepResult, }; +pub(super) use inference_preparation::GooseInferenceRequestPreparer; pub(super) use ops_bang_shell::{bang_shell_command, BangShellOperation}; pub(super) use ops_compaction::CompactionOperation; pub(super) use ops_doctor::DoctorOperation; pub(super) use ops_entry_hook::EntryHookOperation; pub(super) use ops_exit_on_error::ExitOnErrorOperation; -pub(super) use ops_llm::InferenceRunner; +pub(super) use ops_llm::{GooseInferenceProvider, InferenceRunner}; pub(super) use ops_maxturns::{MaxTurnsOperation, MAX_TURNS_MESSAGE}; pub(super) use ops_project::ProjectOperation; pub(super) use ops_recipe::RecipeOperation; pub(super) use ops_retry::RetryOperation; pub(super) use ops_skills::SkillOperation; pub(super) use ops_slash_command::SlashCommandOperation; +pub(super) use ops_status::StatusOperation; pub(super) use ops_steer::{SteerOperation, SteerQueue}; pub(super) use ops_stop_hook::StopHookOperation; pub(super) use ops_tool_approval::ToolApprovalOperation; diff --git a/crates/goose/src/agents/state_machine/ops_compaction.rs b/crates/goose/src/agents/state_machine/ops_compaction.rs index 7e4dfed24..08166b2a9 100644 --- a/crates/goose/src/agents/state_machine/ops_compaction.rs +++ b/crates/goose/src/agents/state_machine/ops_compaction.rs @@ -76,6 +76,13 @@ impl CompactionOperation { (tokens as f64 / self.context_limit as f64) > self.threshold } + async fn context_tokens(&self, session: &Session, conversation: &Conversation) -> Result { + match session.usage.total_tokens { + Some(tokens) => Ok(tokens), + None => crate::context_mgmt::count_context_tokens(conversation).await, + } + } + async fn command_error( conversation: &Conversation, message: String, @@ -195,13 +202,13 @@ impl Operation for CompactionOperation { async fn moim_parts( &self, session: &Session, - _conversation: &Conversation, + conversation: &Conversation, ) -> Result> { if self.manages_own_context { return Ok(Vec::new()); } Ok(compaction_part( - session.usage.total_tokens, + Some(self.context_tokens(session, conversation).await?), self.context_limit, self.threshold, ) @@ -240,9 +247,9 @@ impl Operation for CompactionOperation { if last_effective_role(messages)? != EffectiveRole::User { return not_applicable(); } - match session.usage.total_tokens { - Some(tokens) if tokens > 0 && self.over_threshold(tokens as usize) => {} - _ => return not_applicable(), + let tokens = self.context_tokens(session, conversation).await?; + if tokens <= 0 || !self.over_threshold(tokens as usize) { + return not_applicable(); } } diff --git a/crates/goose/src/agents/state_machine/ops_llm.rs b/crates/goose/src/agents/state_machine/ops_llm.rs index dbcf25e77..9347e0bd2 100644 --- a/crates/goose/src/agents/state_machine/ops_llm.rs +++ b/crates/goose/src/agents/state_machine/ops_llm.rs @@ -1,612 +1,118 @@ -//! Builds a provider request and streams the next assistant response. +//! Goose integration for the reusable inference operation. use std::sync::Arc; -use crate::agents::state_machine::ops_unknown_tool::UNCLAIMED_TOOL_ERROR; -use crate::agents::state_machine::{ - applied, messages_since_kickoff, not_applicable, trailing_error, yielded_with, - ConversationEffect, Emitter, GooseEffect, Inference, InferenceInput, Operation, - OperationResult, SlashCommand, -}; -use crate::agents::{ExtensionManager, PromptManager}; -use crate::config::GooseMode; -use crate::conversation::message::{InferenceMetadata, Message, MessageContent}; -use crate::conversation::{effective_role, Conversation, EffectiveRole}; -use crate::providers::base::{Provider, ProviderUsage}; -use crate::session::Session; -use crate::tool_inspection::ToolInspectionManager; -use anyhow::{anyhow, Result}; use async_trait::async_trait; -use futures::StreamExt; +use goose_agent::inference::InferenceEffect; +pub use goose_agent::inference::InferenceRunner; +use goose_providers::base::{MessageStream, ModelInfo, Provider}; +use goose_providers::conversation::message::Message; +use goose_providers::conversation::token_usage::ProviderUsage; use goose_providers::errors::ProviderError; use goose_providers::model::ModelConfig; -use tokio::sync::Mutex; -use tracing_futures::Instrument; -const EMPTY_RESPONSE_MESSAGE: &str = - "The model returned an empty response. Please resend your message to continue."; -const CANCELLED_TOOL_RESPONSE: &str = "Tool call was cancelled before execution"; +use crate::agents::state_machine::GooseEffect; -fn is_thinking(content: &MessageContent) -> bool { - matches!( - content, - MessageContent::Thinking(_) | MessageContent::RedactedThinking(_) - ) +pub(super) use goose_agent::inference::{chat_span, record_chat_usage}; + +pub struct GooseInferenceProvider { + inner: Arc, } -fn normalize_tool_call_thinking(accumulator: &mut Conversation, chunk: &mut Message) { - if !chunk - .content +impl GooseInferenceProvider { + pub fn new(inner: Arc) -> Self { + Self { inner } + } +} + +impl InferenceEffect for GooseEffect { + fn record_usage(usage: ProviderUsage) -> Self { + GooseEffect::RecordUsage(usage) + } +} + +fn enrich_unclaimed_tool_errors(messages: &[Message], tools: &[rmcp::model::Tool]) -> Vec { + let mut available_tools = tools .iter() - .any(|content| matches!(content, MessageContent::ToolRequest(_))) - { - return; - } - - let has_direct_thinking = chunk.content.iter().any(is_thinking); - let mut prior_thinking = Vec::new(); - for message in accumulator.messages_mut() { - if message.role != chunk.role - || message - .content - .iter() - .any(|content| matches!(content, MessageContent::ToolRequest(_))) - { - continue; - } - prior_thinking.extend( - message - .content - .iter() - .filter(|content| is_thinking(content)) - .cloned(), - ); - message.content.retain(|content| !is_thinking(content)); - } - accumulator - .messages_mut() - .retain(|message| !message.content.is_empty()); - - if !has_direct_thinking && !prior_thinking.is_empty() { - if let Some(tool_request) = chunk - .content - .iter() - .position(|content| matches!(content, MessageContent::ToolRequest(_))) - { - chunk - .content - .splice(tool_request..tool_request, prior_thinking); - } - } -} - -pub(super) fn chat_span( - provider: &dyn Provider, - model_config: &ModelConfig, - session_id: &str, - purpose: &'static str, -) -> tracing::Span { - let span = tracing::info_span!( - target: "goose::state_machine", - "chat", - "gen_ai.operation.name" = "chat", - "gen_ai.provider.name" = %provider.get_name(), - "gen_ai.request.model" = %model_config.model_name, - "gen_ai.request.temperature" = tracing::field::Empty, - "gen_ai.request.max_tokens" = tracing::field::Empty, - "gen_ai.response.model" = tracing::field::Empty, - "gen_ai.response.finish_reasons" = tracing::field::Empty, - "gen_ai.response.id" = tracing::field::Empty, - "gen_ai.usage.input_tokens" = tracing::field::Empty, - "gen_ai.usage.output_tokens" = tracing::field::Empty, - "goose.chat.purpose" = purpose, - "error.type" = tracing::field::Empty, - session.id = %session_id, - ); - super::super::gen_ai_telemetry::record_request_params(&span, model_config); - span -} - -pub(super) fn record_chat_usage(span: &tracing::Span, usage: &ProviderUsage) { - super::super::gen_ai_telemetry::record_provider_usage(span, usage); -} - -pub struct InferenceRunner<'a> { - provider: Arc, - model_config: ModelConfig, - #[cfg(feature = "code-mode")] - extension_manager: Arc, - goose_mode: &'a Mutex, - prompt_manager: &'a Mutex, - tool_inspection_manager: &'a ToolInspectionManager, - frontend_instructions: &'a Mutex>, -} - -/// The agent-visible conversation as the provider sees it: tool requests left -/// unanswered by an earlier turn are dropped, since nothing will answer them now. -fn messages_for_provider(conversation: &Conversation, turn: &[Message]) -> Vec { - let answered: std::collections::HashSet<&str> = conversation - .messages() - .iter() - .flat_map(|message| message.get_tool_response_ids()) - .collect(); - let start = conversation.len() - turn.len(); - conversation - .messages() - .iter() - .enumerate() - .filter(|(_, message)| message.is_agent_visible()) - .map(|(index, message)| { - let mut message = message.agent_visible_content(); - if index < start { - message.content.retain(|content| match content { - MessageContent::ToolRequest(request) => answered.contains(request.id.as_str()), - _ => true, - }); - } - message - }) - .filter(|message| !message.content.is_empty()) - .collect() -} - -fn ends_with_provider_turn(messages: &[Message]) -> bool { - messages.last().is_some_and(|message| { - matches!( - effective_role(message), - EffectiveRole::User | EffectiveRole::Tool - ) - }) -} - -impl<'a> InferenceRunner<'a> { - #[allow(clippy::too_many_arguments)] - pub fn new( - provider: Arc, - model_config: ModelConfig, - #[cfg(feature = "code-mode")] extension_manager: Arc, - #[cfg(not(feature = "code-mode"))] _extension_manager: Arc, - goose_mode: &'a Mutex, - prompt_manager: &'a Mutex, - tool_inspection_manager: &'a ToolInspectionManager, - frontend_instructions: &'a Mutex>, - ) -> Self { - Self { - provider, - model_config, - #[cfg(feature = "code-mode")] - extension_manager, - goose_mode, - prompt_manager, - tool_inspection_manager, - frontend_instructions, - } - } - - async fn error_outcome(&self, err: &ProviderError, emit: &Emitter) -> Vec { - #[cfg(feature = "telemetry")] - crate::posthog::emit_error(err.telemetry_type(), &err.to_string()); - tracing::Span::current().record("error.type", err.telemetry_type()); - tracing::error!("LLM provider error: {err}"); - let message = Message::from_provider_error(err); - let message = emit.message(message).await; - vec![message.into()] - } -} - -#[async_trait] -impl Operation for InferenceRunner<'_> { - fn name(&self) -> &'static str { - "llm" - } - - async fn cancel( - &self, - _session: &Session, - conversation: &Conversation, - result: OperationResult, - emit: &Emitter, - ) -> Result> { - let mut answered = conversation - .messages() - .iter() - .flat_map(Message::get_tool_response_ids) - .map(str::to_string) - .collect::>(); - let mut requests = Vec::new(); - let mut request_ids = std::collections::HashSet::new(); - - let mut collect = |message: &Message| { - for content in &message.content { - match content { - MessageContent::ToolRequest(request) => { - if request_ids.insert(request.id.clone()) { - requests.push(request.clone()); - } - } - MessageContent::ToolResponse(response) => { - answered.insert(response.id.clone()); - } - _ => {} - } - } - }; - for message in messages_since_kickoff(conversation)? { - collect(message); - } - if let OperationResult::Applied(step) = &result { - for effect in &step.effects { - if let GooseEffect::Conversation(ConversationEffect::AppendMessage(message)) = - effect - { - collect(message); - } - } - } - - let mut response = Message::user(); - for request in requests { - if !answered.contains(&request.id) { - response.add_tool_response_with_metadata( - request.id, - Ok(rmcp::model::CallToolResult::error(vec![ - rmcp::model::ContentBlock::text(CANCELLED_TOOL_RESPONSE), - ])), - request.metadata.as_ref(), - ); - } - } - if response.get_tool_response_ids().is_empty() { - return Ok(result); - } - - let response = emit.message(response).await; - match result { - OperationResult::NotApplicable => applied([response.into()]), - OperationResult::Applied(mut step) => { - step.effects.push(response.into()); - Ok(OperationResult::Applied(step)) - } - } - } - - async fn run_command( - &self, - command: &SlashCommand<'_>, - session: &Session, - conversation: &Conversation, - emit: &Emitter, - ) -> Result> { - if command.command != "status" { - return not_applicable(); - } - - let context_limit = self - .provider - .get_context_limit(&self.model_config) - .await - .unwrap_or_else(|_| self.model_config.context_limit()); - let context_tokens = session.usage.total_tokens.unwrap_or(0).max(0) as usize; - let lifetime_tokens = session.accumulated_usage.total_tokens.unwrap_or(0).max(0) as usize; - let context_pct = if context_limit > 0 { - let pct = ((context_tokens as f64 / context_limit as f64) * 100.0).round() as usize; - format!("{}%", pct.min(100)) - } else { - "N/A".to_string() - }; - let response = Message::assistant() - .with_text(format!( - "**Session status**\n\n\ - - Model: {}\n\ - - Provider: {}\n\ - - Mode: {}\n\ - - Tokens (lifetime): {}\n\ - - Context: {} / {} tokens ({})", - self.model_config.model_name, - self.provider.get_name(), - session.goose_mode, - lifetime_tokens, - context_tokens, - context_limit, - context_pct, - )) - .with_visibility(true, false); - let command_message = messages_since_kickoff(conversation)? - .first() - .cloned() - .ok_or_else(|| anyhow!("status command conversation has no kickoff message"))?; - let message_id = command_message - .id - .clone() - .ok_or_else(|| anyhow!("Persisted slash command message has no id"))?; - emit.message(command_message.with_visibility(true, false)) - .await; - let response = emit.message(response).await; - yielded_with([ - ConversationEffect::SetMessageVisibility { - message_id, - user_visible: true, - agent_visible: false, - } - .into(), - response.into(), - ]) - } -} - -#[async_trait] -impl Inference for InferenceRunner<'_> { - fn applies(&self, conversation: &Conversation) -> bool { - let Ok(turn) = messages_since_kickoff(conversation) else { - return false; - }; - trailing_error(conversation).is_none() - && ends_with_provider_turn(&messages_for_provider(conversation, turn)) - } - - async fn infer( - &self, - session: &Session, - conversation: &Conversation, - mut input: InferenceInput, - emit: &Emitter, - ) -> Result> { - let messages = messages_since_kickoff(conversation)?; - if trailing_error(conversation).is_some() { - return not_applicable(); - } - - let mut messages_for_provider = messages_for_provider(conversation, messages); - if !ends_with_provider_turn(&messages_for_provider) { - return not_applicable(); - } - - let span = chat_span( - self.provider.as_ref(), - &self.model_config, - &session.id, - "inference", - ); - - async { - #[cfg(feature = "code-mode")] - let code_execution_mode = self - .extension_manager - .is_extension_enabled( - crate::agents::platform_extensions::code_execution::EXTENSION_NAME, - ) - .await; - #[cfg(not(feature = "code-mode"))] - let code_execution_mode = false; - - let goose_mode = *self.goose_mode.lock().await; - if goose_mode == GooseMode::SmartApprove { - self.tool_inspection_manager - .apply_tool_annotations(&input.tools); - } - let tools = crate::agents::reply_parts::prepare_inference_tools( - input.tools, - code_execution_mode, - ); - if let Some(frontend_instructions) = self.frontend_instructions.lock().await.clone() { - input - .prompt_parts - .push(("frontend".to_string(), frontend_instructions)); - } - let system_prompt = self.prompt_manager.lock().await.build_system_prompt( - &session.working_dir, - input.prompt_parts, - goose_mode, - ); - let (tools, toolshim_tools, system_prompt) = - crate::agents::reply_parts::prepare_tools_for_provider( - tools, - system_prompt, - &self.model_config, - ); - let mut available_tools = tools - .iter() - .chain(toolshim_tools.iter()) - .map(|tool| tool.name.as_ref()) - .collect::>(); - available_tools.sort_unstable(); - available_tools.dedup(); - let available_tools = available_tools.join(", "); - for message in &mut messages_for_provider { - for content in &mut message.content { - let MessageContent::ToolResponse(response) = content else { - continue; - }; - let Some(metadata) = &mut response.metadata else { - continue; - }; - if metadata.remove(UNCLAIMED_TOOL_ERROR).is_none() { - continue; - } - let Ok(result) = &mut response.tool_result else { - continue; - }; - result.content.push(rmcp::model::ContentBlock::text(format!( - "Available tools: [{}].", - available_tools - ))); - } - } - - let context_limit = self - .provider - .get_context_limit(&self.model_config) - .await - .unwrap_or_else(|_| self.model_config.context_limit()); - let provider_name = self.provider.get_name(); - if let Some(session_id) = super::super::latest_provider_session_id( - conversation.messages(), - provider_name, - ) { - if let Err(error) = self.provider.resume(session_id).await { - tracing::warn!( - provider = provider_name, - %error, - "Could not resume provider session; continuing with a handoff" - ); - } - } - let turn = messages_since_kickoff(conversation)?; - let turn_start = turn - .first() - .and_then(|message| chrono::DateTime::from_timestamp(message.created, 0)) - .map(|timestamp| timestamp.with_timezone(&chrono::Local)) - .unwrap_or_else(chrono::Local::now); - // Persist a fresh turn-context event only when its bytes changed - // (e.g. the turn budget ticked); earlier events stay in place. - let last_turn_context = turn - .iter() - .rev() - .find(|message| message.is_turn_context()) - .map(Message::as_concat_text); - let turn_context = crate::agents::moim::turn_context_event( - &session.working_dir, - Some(context_limit), - input.moim_parts, - turn_start, - ) - .filter(|event| Some(event.as_concat_text()) != last_turn_context); - if let Some(event) = &turn_context { - messages_for_provider.push(event.clone()); - } - let conversation_for_provider = Conversation::new_unvalidated(messages_for_provider); - let mut usage_effects: Vec = - turn_context.into_iter().map(GooseEffect::from).collect(); - - let stream = crate::agents::reply_parts::stream_response_from_provider( - self.provider.clone(), - self.model_config.clone(), - &session.id, - &system_prompt, - conversation_for_provider.messages(), - &tools, - &toolshim_tools, - ) - .await; - - let mut stream = match stream { - Ok(stream) => stream, - Err(err) => { - usage_effects.extend(self.error_outcome(&err, emit).await); - return applied(usage_effects); - } + .map(|tool| tool.name.as_ref()) + .collect::>(); + available_tools.sort_unstable(); + available_tools.dedup(); + let available_tools = available_tools.join(", "); + let mut messages = messages.to_vec(); + for message in &mut messages { + for content in &mut message.content { + let goose_providers::conversation::message::MessageContent::ToolResponse(response) = + content + else { + continue; }; - - let requested_model = self.model_config.model_name.clone(); - let resolved_model = self - .provider - .fetch_model_info(&requested_model) - .await - .ok() - .and_then(|model_info| model_info.resolved_model); - let provider_session_id = self.provider.provider_session_id(); - let inference = Some(InferenceMetadata { - provider: self.provider.get_name().to_string(), - requested_model, - resolved_model, - provider_session_id, - }); - - let mut accumulator = Conversation::empty(); - let mut tool_request_ids = std::collections::HashSet::new(); - loop { - tokio::select! { - biased; - _ = emit.cancelled() => break, - next = stream.next() => { - let Some(result) = next else { break }; - let (msg_opt, usage_opt) = match result { - Ok(chunk) => chunk, - Err(err) => { - usage_effects.extend(accumulator.into_iter().map(GooseEffect::from)); - usage_effects.extend(self.error_outcome(&err, emit).await); - return applied(usage_effects); - } - }; - if let Some(usage) = usage_opt { - let span = tracing::Span::current(); - record_chat_usage(&span, &usage); - usage_effects.push(GooseEffect::RecordUsage(usage)); - } - if let Some(mut chunk) = msg_opt { - if let Some(inference) = &inference { - chunk = chunk.with_inference_if_assistant(inference.clone()); - } - chunk.content.retain(|content| match content { - MessageContent::ToolRequest(request) => { - tool_request_ids.insert(request.id.clone()) - } - _ => true, - }); - normalize_tool_call_thinking(&mut accumulator, &mut chunk); - if chunk.content.is_empty() { - if chunk.metadata.output_token_limit_reached { - chunk = emit.message(chunk).await; - } - accumulator.push(chunk); - continue; - } - let chunk = emit.message(chunk).await; - accumulator.push(chunk); - } - } - } + let Some(metadata) = &mut response.metadata else { + continue; + }; + if metadata + .remove(super::ops_unknown_tool::UNCLAIMED_TOOL_ERROR) + .is_none() + { + continue; } - - let empty_response = !accumulator - .iter() - .any(|message| message.metadata.output_token_limit_reached) - && accumulator.iter().all(|message| { - message.content.iter().all(|content| match content { - MessageContent::Text(text) => text.text.trim().is_empty(), - MessageContent::Thinking(thinking) => thinking.thinking.trim().is_empty(), - _ => false, - }) - }); - if empty_response { - let message = Message::assistant().with_text(EMPTY_RESPONSE_MESSAGE); - let message = emit.message(message).await; - usage_effects.push(message.into()); - return yielded_with(usage_effects); - } - - let has_recorded_tokens = usage_effects.iter().any(|effect| { - matches!( - effect, - GooseEffect::RecordUsage(usage) - if usage.usage.input_tokens.is_some() - || usage.usage.output_tokens.is_some() - || usage.usage.total_tokens.is_some() - ) - }); - if !has_recorded_tokens { - let mut usage = ProviderUsage::new( - self.model_config.model_name.clone(), - goose_providers::conversation::token_usage::Usage::default(), - ); - if let Some(response) = accumulator.last() { - crate::providers::usage_estimator::ensure_usage_tokens( - &mut usage, - &system_prompt, - conversation_for_provider.messages(), - response, - &tools, - ) - .await?; - record_chat_usage(&tracing::Span::current(), &usage); - usage_effects.push(GooseEffect::RecordUsage(usage)); - } - } - - usage_effects.extend(accumulator.into_iter().map(Into::into)); - applied(usage_effects) + let Ok(result) = &mut response.tool_result else { + continue; + }; + result.content.push(rmcp::model::ContentBlock::text(format!( + "Available tools: [{available_tools}]." + ))); } - .instrument(span) + } + messages +} + +#[async_trait] +impl Provider for GooseInferenceProvider { + fn get_name(&self) -> &str { + self.inner.get_name() + } + + fn provider_session_id(&self) -> Option { + self.inner.provider_session_id() + } + + async fn resume(&self, session_id: &str) -> Result<(), ProviderError> { + self.inner.resume(session_id).await + } + + async fn stream( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[rmcp::model::Tool], + ) -> Result { + let messages = enrich_unclaimed_tool_errors(messages, tools); + let (tools, toolshim_tools, system_prompt) = + crate::agents::reply_parts::prepare_tools_for_provider( + tools.to_vec(), + system.to_string(), + model_config, + ); + let session_id = crate::session_context::current_session_id().unwrap_or_default(); + crate::agents::reply_parts::stream_response_from_provider( + self.inner.clone(), + model_config.clone(), + &session_id, + &system_prompt, + &messages, + &tools, + &toolshim_tools, + ) .await } + + async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { + self.inner.get_context_limit(model_config).await + } + + async fn fetch_model_info(&self, model_name: &str) -> Result { + self.inner.fetch_model_info(model_name).await + } } diff --git a/crates/goose/src/agents/state_machine/ops_status.rs b/crates/goose/src/agents/state_machine/ops_status.rs new file mode 100644 index 000000000..f2e8e5066 --- /dev/null +++ b/crates/goose/src/agents/state_machine/ops_status.rs @@ -0,0 +1,85 @@ +use std::sync::Arc; + +use anyhow::{anyhow, Result}; +use async_trait::async_trait; +use goose_providers::base::Provider; +use goose_providers::model::ModelConfig; + +use crate::agents::state_machine::{ + messages_since_kickoff, not_applicable, yielded_with, ConversationEffect, Emitter, GooseEffect, + Operation, OperationResult, SlashCommand, +}; +use crate::conversation::message::Message; +use crate::conversation::Conversation; +use crate::session::Session; + +pub struct StatusOperation { + provider: Arc, + model_config: ModelConfig, +} + +impl StatusOperation { + pub fn new(provider: Arc, model_config: ModelConfig) -> Self { + Self { + provider, + model_config, + } + } +} + +#[async_trait] +impl Operation for StatusOperation { + fn name(&self) -> &'static str { + "status" + } + + async fn run_command( + &self, + command: &SlashCommand<'_>, + session: &Session, + conversation: &Conversation, + emit: &Emitter, + ) -> Result> { + if command.command != "status" { + return not_applicable(); + } + let context_limit = self + .provider + .get_context_limit(&self.model_config) + .await + .unwrap_or_else(|_| self.model_config.context_limit()); + let context_tokens = session.usage.total_tokens.unwrap_or(0); + let lifetime_tokens = session.accumulated_usage.total_tokens.unwrap_or(0); + let context_pct = if context_limit > 0 { + format!( + "{}%", + ((context_tokens as f64 / context_limit as f64) * 100.0) + .round() + .min(100.0) as usize + ) + } else { + "N/A".to_string() + }; + let response = Message::assistant().with_text(format!("**Session status**\n\n- Model: {}\n- Provider: {}\n- Mode: {}\n- Tokens (lifetime): {}\n- Context: {} / {} tokens ({})", self.model_config.model_name, self.provider.get_name(), session.goose_mode, lifetime_tokens, context_tokens, context_limit, context_pct)).with_visibility(true, false); + let command_message = messages_since_kickoff(conversation)? + .first() + .cloned() + .ok_or_else(|| anyhow!("status command conversation has no kickoff message"))?; + let message_id = command_message + .id + .clone() + .ok_or_else(|| anyhow!("Persisted slash command message has no id"))?; + emit.message(command_message.with_visibility(true, false)) + .await; + let response = emit.message(response).await; + yielded_with([ + ConversationEffect::SetMessageVisibility { + message_id, + user_visible: true, + agent_visible: false, + } + .into(), + response.into(), + ]) + } +} diff --git a/crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs b/crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs index aebde4919..907f60c8f 100644 --- a/crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs +++ b/crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs @@ -135,6 +135,39 @@ async fn proactive_and_manual_compaction_continue_with_replaced_usage() -> Resul Ok(()) } +#[tokio::test] +async fn tokenless_provider_compacts_estimated_context() -> Result<()> { + let (pipeline, api) = pipeline::test_pipeline_with(ProviderFeatures { + reports_usage: false, + ..ProviderFeatures::default() + }) + .await?; + let pipeline = pipeline + .with_model_config( + goose_providers::model::ModelConfig::new("gpt-4.1").with_context_limit(Some(200)), + ) + .await; + let large_context = (0..500) + .map(|index| format!("token-{index}")) + .collect::>() + .join(" "); + pipeline + .seed([Message::user().with_text(large_context)]) + .await?; + assert!(pipeline.session().await?.usage.total_tokens.is_none()); + + api.on(SUMMARIZE_HISTORY).reply("summary"); + api.on("Your context was compacted") + .reply("continued after estimated compaction"); + + let compacted = pipeline.run(["continue"]).await?; + compacted.assert_message(-1, Agent, "continued after estimated compaction"); + compacted.assert_emitted("Performing auto-compaction"); + assert_eq!(compacted.history_replacements(), 1); + + Ok(()) +} + #[tokio::test] async fn a_failed_compact_command_reports_the_error_and_keeps_working() -> Result<()> { let (pipeline, api) = test_pipeline().await?; diff --git a/crates/goose/src/agents/state_machine/tests/pipeline.rs b/crates/goose/src/agents/state_machine/tests/pipeline.rs index 9aaf22bde..0a8237aab 100644 --- a/crates/goose/src/agents/state_machine/tests/pipeline.rs +++ b/crates/goose/src/agents/state_machine/tests/pipeline.rs @@ -16,9 +16,10 @@ use crate::agents::mcp_client::McpClientTrait; use crate::agents::prompt_manager::PromptManager; use crate::agents::state_machine::{ BangShellOperation, CompactionOperation, DoctorOperation, Emitter, EntryHookOperation, - ExitOnErrorOperation, GooseEffect, InferenceRunner, MaxTurnsOperation, Operation, - ProjectOperation, RecipeOperation, RetryOperation, SkillOperation, SlashCommandOperation, - StateMachine, SteerOperation, SteerQueue, Step, StopHookOperation, ToolApprovalOperation, + ExitOnErrorOperation, GooseEffect, GooseInferenceProvider, GooseInferenceRequestPreparer, + InferenceRunner, MaxTurnsOperation, Operation, ProjectOperation, RecipeOperation, + RetryOperation, SkillOperation, SlashCommandOperation, StateMachine, StatusOperation, + SteerOperation, SteerQueue, Step, StopHookOperation, ToolApprovalOperation, ToolExecutionOperation, ToolPairCompactionOperation, UnknownToolOperation, }; use crate::agents::AgentEvent; @@ -173,17 +174,26 @@ impl TestPipeline { Arc::new(ExitOnErrorOperation), ]; operations.extend(remaining_operations); - let inference = Arc::new(InferenceRunner::new( - provider, + let request_preparer = GooseInferenceRequestPreparer { + #[cfg(feature = "code-mode")] + extension_manager: self.extension_manager.clone(), + goose_mode: &self.goose_mode, + prompt_manager: &self.prompt_manager, + tool_inspection_manager: &self.tool_inspection_manager, + frontend_instructions: &self.frontend_instructions, + context_limit: self.model_config.context_limit(), + }; + let status_operation = Arc::new(StatusOperation::new( + provider.clone(), self.model_config.clone(), - self.extension_manager.clone(), - &self.goose_mode, - &self.prompt_manager, - &self.tool_inspection_manager, - &self.frontend_instructions, )); + let inference_provider = Arc::new(GooseInferenceProvider::new(provider)); + let inference = Arc::new( + InferenceRunner::new(inference_provider, self.model_config.clone()) + .with_request_preparer(Arc::new(request_preparer)), + ); let mut command_handlers = operations.clone(); - command_handlers.push(inference.clone()); + command_handlers.push(status_operation); let command_operation: Arc + '_> = Arc::new(SlashCommandOperation::new(command_handlers)); let steps = std::iter::once(Arc::new(EntryHookOperation::new(self.hook_manager.clone())) diff --git a/crates/goose/src/agents/state_machine/tests/provider_lifecycle.rs b/crates/goose/src/agents/state_machine/tests/provider_lifecycle.rs index b9b604bd4..0f1801abb 100644 --- a/crates/goose/src/agents/state_machine/tests/provider_lifecycle.rs +++ b/crates/goose/src/agents/state_machine/tests/provider_lifecycle.rs @@ -174,20 +174,15 @@ async fn provider_lifecycle() -> Result<()> { result.assert_message(-2, ToolResponse, "result: 6"); result.assert_message(-1, Agent, "The total is 6."); - let conversation = result.session.conversation.as_ref().unwrap(); - let first_total = result + assert!(result.session.usage.total_tokens.is_none()); + assert!(result .session - .usage - .total_tokens - .expect("estimated session usage"); - assert!(first_total > 0); - assert_eq!( - conversation - .last() - .and_then(|message| message.metadata.usage.as_ref()) - .and_then(|usage| usage.total_tokens), - Some(first_total) - ); + .conversation + .as_ref() + .and_then(Conversation::last) + .and_then(|message| message.metadata.usage.as_ref()) + .and_then(|usage| usage.total_tokens) + .is_none()); api.on("return no choices").no_choices(); let result = pipeline.run(["return no choices"]).await?; @@ -224,11 +219,7 @@ async fn provider_lifecycle() -> Result<()> { .reply("recovered from server error"); let result = pipeline.run(["after server error"]).await?; result.assert_message(-1, Agent, "recovered from server error"); - assert!(result - .session - .usage - .total_tokens - .is_some_and(|total| total > first_total)); + assert!(result.session.usage.total_tokens.is_none()); assert!(api .calls() .iter()