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

move inference operation into goose-agent (#11294)

Browse Source 
Signed-off-by: Jack Amadeo <jackamadeo@squareup.com>
Signed-off-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Jack Amadeo
2026-08-25 17:58:15 +00:00
committed by GitHub
parent 5679655955
commit 724250c424
13 changed files with 990 additions and 646 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+3
View File
@@ -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",
]
crates/goose-agent/Cargo.toml
+3
View File
@@ -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 }
crates/goose-agent/src/inference.rs
+600
View File
@@ -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<rmcp::model::Tool>,
pub additional_messages: Vec<Message>,
}
#[async_trait]
pub trait InferenceRequestPreparer<S>: Send + Sync {
async fn prepare(
&self,
session: &S,
conversation: &Conversation,
input: InferenceInput,
) -> Result<PreparedInferenceRequest>;
}
pub struct IdentityInferenceRequestPreparer;
#[async_trait]
impl<S: Sync> InferenceRequestPreparer<S> for IdentityInferenceRequestPreparer {
async fn prepare(
&self,
_session: &S,
_conversation: &Conversation,
input: InferenceInput,
) -> Result<PreparedInferenceRequest> {
Ok(PreparedInferenceRequest {
system_prompt: input
.prompt_parts
.into_iter()
.map(|(_, part)| part)
.collect::<Vec<_>>()
.join("\n\n"),
tools: input.tools,
additional_messages: Vec::new(),
})
}
}
pub trait InferenceEffect: From<Message> + 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<dyn Provider>,
model_config: ModelConfig,
request_preparer: Arc<dyn InferenceRequestPreparer<S> + 'a>,
effect: std::marker::PhantomData<fn() -> 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<Message> {
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<Message> {
let mut answered = persisted
.iter()
.chain(pending)
.flat_map(Message::get_tool_response_ids)
.collect::<std::collections::HashSet<_>>();
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<dyn Provider>, 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<dyn InferenceRequestPreparer<S> + 'a>,
) -> Self {
self.request_preparer = request_preparer;
self
}
async fn error_outcome(&self, err: &ProviderError, emit: &Emitter) -> Vec<E> {
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<S: Sync, E: InferenceEffect> Operation<S, E> for InferenceRunner<'_, S, E> {
fn name(&self) -> &'static str {
"llm"
}
async fn cancel(
&self,
_session: &S,
conversation: &Conversation,
result: OperationResult<E>,
emit: &Emitter,
) -> Result<OperationResult<E>> {
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<S: Sync, E: InferenceEffect> Inference<S, E> 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<OperationResult<E>> {
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<E> = 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::<Vec<_>>();
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());
}
}
crates/goose-agent/src/lib.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod events;
pub mod inference;
pub mod machine;
pub mod operation;
crates/goose/src/agents/agent.rs
+26 -16
View File
@@ -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<dyn Operation<Session, GooseEffect> + '_> =
Arc::new(SlashCommandOperation::new(command_handlers));
let operations: Vec<_> =
@@ -1862,7 +1869,10 @@ impl Agent {
let (tx, mut rx) = mpsc::channel::<AgentEvent>(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! {
crates/goose/src/agents/state_machine/inference_preparation.rs
+91
View File
@@ -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<ExtensionManager>,
pub(crate) goose_mode: &'a Mutex<GooseMode>,
pub(crate) prompt_manager: &'a Mutex<PromptManager>,
pub(crate) tool_inspection_manager: &'a ToolInspectionManager,
pub(crate) frontend_instructions: &'a Mutex<Option<String>>,
pub(crate) context_limit: usize,
}
#[async_trait]
impl InferenceRequestPreparer<Session> for GooseInferenceRequestPreparer<'_> {
async fn prepare(
&self,
session: &Session,
conversation: &Conversation,
mut input: InferenceInput,
) -> Result<PreparedInferenceRequest> {
#[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,
})
}
}
crates/goose/src/agents/state_machine/mod.rs
+5 -1
View File
@@ -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;
crates/goose/src/agents/state_machine/ops_compaction.rs
+12 -5
View File
@@ -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<i32> {
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<Session, GooseEffect> for CompactionOperation {
async fn moim_parts(
&self,
session: &Session,
_conversation: &Conversation,
conversation: &Conversation,
) -> Result<Vec<String>> {
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<Session, GooseEffect> 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();
}
}
crates/goose/src/agents/state_machine/ops_llm.rs
+101 -595
View File
@@ -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<dyn Provider>,
}
fn normalize_tool_call_thinking(accumulator: &mut Conversation, chunk: &mut Message) {
if !chunk
.content
impl GooseInferenceProvider {
pub fn new(inner: Arc<dyn Provider>) -> 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<Message> {
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<dyn Provider>,
model_config: ModelConfig,
#[cfg(feature = "code-mode")]
extension_manager: Arc<ExtensionManager>,
goose_mode: &'a Mutex<GooseMode>,
prompt_manager: &'a Mutex<PromptManager>,
tool_inspection_manager: &'a ToolInspectionManager,
frontend_instructions: &'a Mutex<Option<String>>,
}
/// 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<Message> {
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<dyn Provider>,
model_config: ModelConfig,
#[cfg(feature = "code-mode")] extension_manager: Arc<ExtensionManager>,
#[cfg(not(feature = "code-mode"))] _extension_manager: Arc<ExtensionManager>,
goose_mode: &'a Mutex<GooseMode>,
prompt_manager: &'a Mutex<PromptManager>,
tool_inspection_manager: &'a ToolInspectionManager,
frontend_instructions: &'a Mutex<Option<String>>,
) -> 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<GooseEffect> {
#[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<Session, GooseEffect> for InferenceRunner<'_> {
fn name(&self) -> &'static str {
"llm"
}
async fn cancel(
&self,
_session: &Session,
conversation: &Conversation,
result: OperationResult<GooseEffect>,
emit: &Emitter,
) -> Result<OperationResult<GooseEffect>> {
let mut answered = conversation
.messages()
.iter()
.flat_map(Message::get_tool_response_ids)
.map(str::to_string)
.collect::<std::collections::HashSet<_>>();
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<OperationResult<GooseEffect>> {
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<Session, GooseEffect> 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<OperationResult<GooseEffect>> {
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::<Vec<_>>();
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<GooseEffect> =
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::<Vec<_>>();
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<String> {
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<MessageStream, ProviderError> {
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<usize, ProviderError> {
self.inner.get_context_limit(model_config).await
}
async fn fetch_model_info(&self, model_name: &str) -> Result<ModelInfo, ProviderError> {
self.inner.fetch_model_info(model_name).await
}
}
crates/goose/src/agents/state_machine/ops_status.rs
+85
View File
@@ -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<dyn Provider>,
model_config: ModelConfig,
}
impl StatusOperation {
pub fn new(provider: Arc<dyn Provider>, model_config: ModelConfig) -> Self {
Self {
provider,
model_config,
}
}
}
#[async_trait]
impl Operation<Session, GooseEffect> for StatusOperation {
fn name(&self) -> &'static str {
"status"
}
async fn run_command(
&self,
command: &SlashCommand<'_>,
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult<GooseEffect>> {
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(),
])
}
}
crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs
+33
View File
@@ -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::<Vec<_>>()
.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?;
crates/goose/src/agents/state_machine/tests/pipeline.rs
+21 -11
View File
@@ -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<dyn Operation<Session, GooseEffect> + '_> =
Arc::new(SlashCommandOperation::new(command_handlers));
let steps = std::iter::once(Arc::new(EntryHookOperation::new(self.hook_manager.clone()))
crates/goose/src/agents/state_machine/tests/provider_lifecycle.rs
+9 -18
View File
@@ -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()