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

make the unrolled loop state machine generic (#11139)

Browse Source 
Signed-off-by: Jack Amadeo <jackamadeo@squareup.com>
This commit is contained in:
Jack Amadeo
2026-08-13 20:56:59 +00:00
committed by GitHub
parent 2508223fc0
commit d2389ef4ad
30 changed files with 710 additions and 587 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+17 -24
View File
@@ -27,16 +27,18 @@ use crate::agents::platform_extensions::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::prompt_manager::PromptManager;
use crate::agents::retry::{RetryManager, RetryResult};
use crate::agents::state_machine::{
BangShellOperation, CompactionOperation, DoctorOperation, Emitter, ExitOnErrorOperation,
InferenceRunner, MaxTurnsOperation, Operation, ProjectOperation, RecipeOperation,
RetryOperation, SkillOperation, SlashCommandOperation, StateMachine, SteerOperation,
SteerQueue, Step, StopHookOperation, ToolApprovalOperation, ToolExecutionOperation,
ToolPairCompactionOperation, UnknownToolOperation, MAX_TURNS_MESSAGE,
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,
};
use crate::agents::types::{
FrontendTool, SessionConfig, SharedProvider, ToolResultReceiver,
DEFAULT_ON_FAILURE_TIMEOUT_SECONDS, DEFAULT_RETRY_TIMEOUT_SECONDS,
};
use crate::agents::AgentEvent;
use crate::config::extensions::name_to_key;
use crate::config::permission::PermissionManager;
use crate::config::{get_enabled_extensions, Config, GooseMode};
@@ -71,7 +73,7 @@ use goose_providers::thinking::ThinkingEffort;
use regex::Regex;
use rmcp::model::{
CallToolRequestParams, CallToolResult, ContentBlock, ElicitationAction, ErrorCode, ErrorData,
GetPromptResult, Prompt, ServerNotification, Tool,
GetPromptResult, Prompt, Tool,
};
use serde_json::Value;
use tokio::sync::{mpsc, Mutex};
@@ -282,18 +284,6 @@ pub struct Agent {
steer_queues: Mutex<HashMap<String, SteerQueue>>,
}
#[derive(Clone, Debug)]
pub enum AgentEvent {
Message(Message),
Usage(crate::providers::base::ProviderUsage),
MessageUsage {
message_id: Option<String>,
usage: MessageUsage,
},
McpNotification((String, ServerNotification)),
HistoryReplaced(Conversation),
}
fn ensure_message_event_id(event: AgentEvent) -> AgentEvent {
match event {
AgentEvent::Message(message) => AgentEvent::Message(message.with_generated_id_if_missing()),
@@ -1600,7 +1590,7 @@ impl Agent {
max_turns: Option<u32>,
cancel: CancellationToken,
steer_queue: SteerQueue,
) -> StateMachine<'_> {
) -> StateMachine<'_, Session, GooseEffect> {
let max_turns = max_turns.unwrap_or_else(|| {
Config::global()
.get_param::<u32>("GOOSE_MAX_TURNS")
@@ -1633,7 +1623,7 @@ impl Agent {
let tool_pair_compaction_enabled = crate::context_mgmt::tool_pair_summarization_enabled()
&& !provider.manages_own_context();
let operations: Vec<Arc<dyn Operation + '_>> = vec![
let operations: Vec<Arc<dyn Operation<Session, GooseEffect> + '_>> = vec![
Arc::new(SteerOperation::new(steer_queue, self.hook_manager.clone())),
Arc::new(MaxTurnsOperation::new(max_turns)),
Arc::new(BangShellOperation::new()),
@@ -1686,9 +1676,12 @@ impl Agent {
));
let mut command_handlers = operations.clone();
command_handlers.push(inference.clone());
let command_operation: Arc<dyn Operation + '_> =
let command_operation: Arc<dyn Operation<Session, GooseEffect> + '_> =
Arc::new(SlashCommandOperation::new(command_handlers));
let operations: Vec<_> = std::iter::once(command_operation)
let operations: Vec<_> =
std::iter::once(Arc::new(EntryHookOperation::new(self.hook_manager.clone()))
as Arc<dyn Operation<Session, GooseEffect> + '_>)
.chain(std::iter::once(command_operation))
.chain(operations)
.collect();
@@ -1698,7 +1691,7 @@ impl Agent {
.chain(std::iter::once(Step::Inference(inference)))
.collect();
StateMachine::new(steps, cancel).with_hook_manager(self.hook_manager.clone())
StateMachine::new(steps, cancel)
}
pub(crate) async fn reply_with_state_machine(
@@ -1784,7 +1777,7 @@ impl Agent {
let (tx, mut rx) = mpsc::channel::<AgentEvent>(32);
let emit = Emitter::new(tx, cancel.clone());
let result = {
let run = machine.run(session_manager.as_ref(), &session_id, &emit);
let run = run_goose(&machine, session_manager.as_ref(), &session_id, &emit);
tokio::pin!(run);
loop {
tokio::select! {
crates/goose/src/agents/events.rs
+16
View File
@@ -0,0 +1,16 @@
use rmcp::model::ServerNotification;
use crate::conversation::message::{Message, MessageUsage};
use crate::conversation::Conversation;
#[derive(Clone, Debug)]
pub enum AgentEvent {
Message(Message),
Usage(crate::providers::base::ProviderUsage),
MessageUsage {
message_id: Option<String>,
usage: MessageUsage,
},
McpNotification((String, ServerNotification)),
HistoryReplaced(Conversation),
}
crates/goose/src/agents/mod.rs
+4 -2
View File
@@ -1,11 +1,12 @@
mod agent;
pub mod container;
mod events;
pub mod execute_commands;
pub mod extension;
pub mod extension_malware_check;
pub mod extension_manager;
pub mod final_output_tool;
mod gen_ai_telemetry;
pub(crate) mod gen_ai_telemetry;
mod large_response_handler;
pub mod mcp_client;
pub mod moim;
@@ -25,8 +26,9 @@ mod tool_schema_normalize;
pub mod types;
pub mod validate_extensions;
pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult, GoosePlatform};
pub use agent::{Agent, AgentConfig, ExtensionLoadResult, GoosePlatform};
pub use container::Container;
pub use events::AgentEvent;
pub use execute_commands::COMPACT_TRIGGERS;
pub use extension::{ExtensionConfig, ExtensionError};
pub use extension_manager::ExtensionManager;
crates/goose/src/agents/state_machine/effects.rs
+51
View File
@@ -0,0 +1,51 @@
use crate::agents::state_machine::operation::{ConversationEffect, MachineEffect};
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::providers::base::ProviderUsage;
use crate::recipe::Recipe;
use crate::session::ExtensionData;
pub enum GooseEffect {
Conversation(ConversationEffect),
ReplaceConversation {
conversation: Conversation,
usage: Option<ProviderUsage>,
},
SetRecipe(Box<Option<Recipe>>),
SetExtensionData(ExtensionData),
RecordUsage(ProviderUsage),
}
impl MachineEffect for GooseEffect {
fn ensure_message_ids(&mut self) {
match self {
GooseEffect::Conversation(effect) => effect.ensure_message_ids(),
GooseEffect::ReplaceConversation { conversation, .. } => {
for message in conversation.messages_mut() {
if message.id.is_none() {
message.id = Some(format!("msg_{}", uuid::Uuid::new_v4()));
}
}
}
_ => {}
}
}
}
impl From<ConversationEffect> for GooseEffect {
fn from(effect: ConversationEffect) -> Self {
GooseEffect::Conversation(effect)
}
}
impl From<Message> for GooseEffect {
fn from(message: Message) -> Self {
ConversationEffect::from(message).into()
}
}
impl From<Conversation> for GooseEffect {
fn from(conversation: Conversation) -> Self {
ConversationEffect::from(conversation).into()
}
}
crates/goose/src/agents/state_machine/machine.rs
+65 -239
View File
@@ -1,26 +1,43 @@
use std::sync::Arc;
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;
use crate::agents::state_machine::operation::{
messages_since_kickoff, Emitter, Inference, InferenceInput, Operation, OperationFuture,
OperationResult, StateEffect, StepResult,
ConversationEffect, Emitter, Inference, InferenceInput, MachineEffect, Operation,
OperationFuture, OperationResult, StepResult,
};
use crate::agents::state_machine::usage;
use crate::agents::AgentEvent;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::hooks::{HookContext, HookEvent, HookManager};
use crate::session::{Session, SessionManager};
pub enum Step<'a> {
Operation(Arc<dyn Operation + 'a>),
Inference(Arc<dyn Inference + 'a>),
pub trait MachineSession: Send + Sync {
fn id(&self) -> &str;
fn conversation(&self) -> Option<&Conversation>;
}
impl Step<'_> {
fn operation(&self) -> &dyn Operation {
#[async_trait]
pub trait SessionLoader<S>: Send + Sync {
async fn load(&self, session_id: &str) -> Result<S>;
}
#[async_trait]
pub trait EffectHandler<S, E>: Send + Sync {
async fn apply_effects(&self, session: &S, effects: &mut [E], emit: &Emitter) -> Result<()>;
}
pub trait EffectUsage<E>: Send + Sync {
fn usage(&self, _effect: &E) -> Option<goose_providers::conversation::token_usage::Usage> {
None
}
}
pub enum Step<'a, S, E = ConversationEffect> {
Operation(Arc<dyn Operation<S, E> + 'a>),
Inference(Arc<dyn Inference<S, E> + 'a>),
}
impl<S, E: Send> Step<'_, S, E> {
fn operation(&self) -> &dyn Operation<S, E> {
match self {
Step::Operation(operation) => operation.as_ref(),
Step::Inference(inference) => inference.as_ref(),
@@ -28,87 +45,31 @@ impl Step<'_> {
}
}
pub struct StateMachine<'a> {
steps: Vec<Step<'a>>,
pub struct StateMachine<'a, S, E = ConversationEffect> {
steps: Vec<Step<'a, S, E>>,
cancel: CancellationToken,
hook_manager: HookManager,
}
impl<'a> StateMachine<'a> {
pub fn new(steps: Vec<Step<'a>>, cancel: CancellationToken) -> Self {
Self {
steps,
cancel,
hook_manager: HookManager::default(),
}
impl<'a, S, E> StateMachine<'a, S, E>
where
S: MachineSession,
E: MachineEffect + Send + 'static,
{
pub fn new(steps: Vec<Step<'a, S, E>>, cancel: CancellationToken) -> Self {
Self { steps, cancel }
}
pub fn with_hook_manager(mut self, hook_manager: HookManager) -> Self {
self.hook_manager = hook_manager;
self
}
async fn emit_entry_hooks(&self, session: &Session, conversation: &Conversation) -> Result<()> {
let messages = messages_since_kickoff(conversation)?;
if Self::has_agent_reply(messages) {
return Ok(());
}
let messages_before_kickoff =
&conversation.messages()[..conversation.len() - messages.len()];
if !messages_before_kickoff.iter().any(|message| {
message.role == rmcp::model::Role::User
&& message.is_user_visible()
&& !message.is_tool_response()
}) {
self.hook_manager
.emit(
HookEvent::SessionStart,
HookContext::new(HookEvent::SessionStart, &session.id)
.with_working_dir(session.working_dir.to_string_lossy().to_string()),
)
.await;
}
let prompt = messages
.first()
.map(Message::as_concat_text)
.unwrap_or_default();
if !prompt.is_empty() {
self.hook_manager
.emit(
HookEvent::UserPromptSubmit,
HookContext::new(HookEvent::UserPromptSubmit, &session.id)
.with_message(prompt)
.with_working_dir(session.working_dir.to_string_lossy().to_string()),
)
.await;
}
Ok(())
}
fn has_agent_reply(messages: &[Message]) -> bool {
messages.iter().any(|message| {
message.role == rmcp::model::Role::Assistant
&& ((message.is_user_visible() && message.is_agent_visible())
|| message.error_kind().is_some())
})
}
pub async fn step(&self, session: &Session, emit: &Emitter) -> Result<Option<StepResult>> {
pub async fn step(&self, session: &S, emit: &Emitter) -> Result<Option<StepResult<E>>> {
let conversation = session
.conversation
.as_ref()
.conversation()
.ok_or_else(|| anyhow!("state-machine session loaded without conversation"))?;
self.emit_entry_hooks(session, conversation).await?;
for step in &self.steps {
let name = step.operation().name();
let result = if self.cancel.is_cancelled() {
OperationResult::NotApplicable
} else {
let step_fut: OperationFuture<'_, Result<OperationResult>> = match step {
let step_fut: OperationFuture<'_, Result<OperationResult<E>>> = match step {
Step::Operation(operation) => operation.run(session, conversation, emit),
Step::Inference(inference) => {
let mut input = InferenceInput::default();
@@ -140,13 +101,13 @@ impl<'a> StateMachine<'a> {
match result {
OperationResult::NotApplicable => {}
OperationResult::Applied(mut result) => {
// `step` and `apply` are separate entry points; a caller that
// drives steps itself still gets effects it can persist.
result.ensure_message_ids();
result.applied_step = Some(name);
for effect in &mut result.effects {
effect.ensure_message_ids();
}
if cancelled {
result.yield_to_client = true;
}
tracing::debug!(target: "goose::state_machine", step = name, "applied step");
return Ok(Some(result));
}
}
@@ -155,173 +116,38 @@ impl<'a> StateMachine<'a> {
Ok(None)
}
pub async fn apply(
pub async fn apply<R>(
&self,
session_manager: &SessionManager,
session: &Session,
result: &mut StepResult,
runtime: &R,
session: &S,
result: &mut StepResult<E>,
emit: &Emitter,
) -> Result<()> {
result.ensure_message_ids();
usage::enrich(session, &mut result.effects);
for effect in &result.effects {
match effect {
StateEffect::AppendMessage(message) => {
session_manager.add_message(&session.id, message).await?;
}
StateEffect::ReplaceConversation {
conversation,
usage: replacement_usage,
} => {
if let Some(usage) = replacement_usage {
usage::record(session_manager, session, usage, true).await?;
}
session_manager
.replace_conversation(&session.id, conversation)
.await?;
session_manager
.update(&session.id)
.usage(usage::estimate_context(conversation).await?)
.apply()
.await?;
}
StateEffect::PatchToolRequestMeta {
tool_call_id,
patch,
} => {
session_manager
.update_tool_request_meta(&session.id, tool_call_id, patch.clone())
.await?;
}
StateEffect::SetMessageVisibility {
message_id,
user_visible,
agent_visible,
} => {
session_manager
.update_message_metadata(&session.id, message_id, |mut metadata| {
metadata.user_visible = *user_visible;
metadata.agent_visible = *agent_visible;
metadata
})
.await?;
}
StateEffect::SetRecipe(recipe) => {
session_manager
.update(&session.id)
.recipe(recipe.as_ref().clone())
.apply()
.await?;
}
StateEffect::SetExtensionData(extension_data) => {
session_manager
.update(&session.id)
.extension_data(extension_data.clone())
.apply()
.await?;
}
StateEffect::RecordUsage(usage) => {
usage::record(session_manager, session, usage, false).await?;
}
}
) -> Result<()>
where
R: EffectHandler<S, E>,
{
for effect in &mut result.effects {
effect.ensure_message_ids();
}
for effect in &result.effects {
match effect {
StateEffect::AppendMessage(message) => {
if let Some(usage) = message
.metadata
.usage
.as_deref()
.filter(|_| !message.user_visible_content().content.is_empty())
.cloned()
{
emit.emit(AgentEvent::MessageUsage {
message_id: message.id.clone(),
usage,
})
.await;
}
}
StateEffect::ReplaceConversation { conversation, .. } => {
emit.emit(AgentEvent::HistoryReplaced(conversation.clone()))
.await;
}
StateEffect::RecordUsage(usage) => {
emit.emit(AgentEvent::Usage(usage.clone())).await
}
_ => {}
}
}
Ok(())
runtime
.apply_effects(session, &mut result.effects, emit)
.await
}
pub async fn run(
&self,
session_manager: &SessionManager,
session_id: &str,
emit: &Emitter,
) -> Result<Session> {
let entry_session = session_manager.get_session(session_id, true).await?;
if let Some(input) = entry_session
.conversation
.as_ref()
.and_then(|conversation| messages_since_kickoff(conversation).ok())
.and_then(|messages| messages.first())
.map(Message::user_visible_content)
.map(|message| message.as_concat_text())
.filter(|text| !text.is_empty())
{
tracing::Span::current().record("trace_input", input.as_str());
}
let mut turn_usage = goose_providers::conversation::token_usage::Usage::default();
pub async fn run<R>(&self, runtime: &R, session_id: &str, emit: &Emitter) -> Result<S>
where
R: SessionLoader<S> + EffectHandler<S, E>,
{
loop {
let session = session_manager.get_session(session_id, true).await?;
let session = runtime.load(session_id).await?;
let Some(mut result) = self.step(&session, emit).await? else {
break;
};
for effect in &result.effects {
match effect {
StateEffect::RecordUsage(usage)
| StateEffect::ReplaceConversation {
usage: Some(usage), ..
} => turn_usage += usage.usage,
_ => {}
}
}
self.apply(session_manager, &session, &mut result, emit)
.await?;
self.apply(runtime, &session, &mut result, emit).await?;
if result.yield_to_client {
break;
}
}
let session = session_manager.get_session(session_id, true).await?;
let last_assistant_text = session
.conversation
.as_ref()
.and_then(|conversation| messages_since_kickoff(conversation).ok())
.into_iter()
.flatten()
.rev()
.filter(|message| message.role == rmcp::model::Role::Assistant)
.map(Message::user_visible_content)
.map(|message| message.as_concat_text())
.find(|text| !text.is_empty())
.unwrap_or_default();
if !last_assistant_text.is_empty() {
let span = tracing::Span::current();
span.record("trace_output", last_assistant_text.as_str());
if crate::agents::gen_ai_telemetry::capture_message_content() {
let output =
crate::agents::gen_ai_telemetry::simple_output_json(&last_assistant_text);
span.record("gen_ai.output.messages", output.as_str());
}
}
crate::agents::gen_ai_telemetry::record_usage(&tracing::Span::current(), &turn_usage);
Ok(session)
runtime.load(session_id).await
}
}
crates/goose/src/agents/state_machine/mod.rs
+10 -3
View File
@@ -5,11 +5,13 @@
//! `StateMachine::run`. Goose's concrete operations remain internal because their
//! configuration is part of `Agent::reply`, not the state-machine protocol.
mod effects;
mod machine;
mod operation;
mod ops_bang_shell;
mod ops_compaction;
mod ops_doctor;
mod ops_entry_hook;
mod ops_exit_on_error;
mod ops_llm;
mod ops_maxturns;
@@ -24,20 +26,25 @@ mod ops_tool_approval;
mod ops_tool_pair_compaction;
mod ops_toolcalling;
mod ops_unknown_tool;
mod session;
pub(crate) use session::run as run_goose;
mod usage;
#[cfg(test)]
mod tests;
pub use machine::{StateMachine, Step};
pub use effects::GooseEffect;
pub use machine::{EffectHandler, EffectUsage, MachineSession, SessionLoader, StateMachine, Step};
pub use operation::{
applied, not_applicable, yielded, yielded_with, Emitter, Inference, InferenceInput, Operation,
OperationResult, SlashCommand, StateEffect, StepResult,
applied, assistant_turn_count, ends_turn, last_effective_role, messages_since_kickoff,
not_applicable, trailing_error, yielded, yielded_with, ConversationEffect, Emitter, Inference,
InferenceInput, MachineEffect, Operation, OperationResult, SlashCommand, StepResult,
};
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_maxturns::{MaxTurnsOperation, MAX_TURNS_MESSAGE};
crates/goose/src/agents/state_machine/operation.rs
+40 -56
View File
@@ -8,9 +8,6 @@ use tokio_util::sync::CancellationToken;
use crate::agents::AgentEvent;
use crate::conversation::message::{Message, MessageContent, MessageErrorKind};
use crate::conversation::{effective_role, Conversation, EffectiveRole};
use crate::providers::base::ProviderUsage;
use crate::recipe::Recipe;
use crate::session::{ExtensionData, Session};
use rmcp::model::Tool;
pub type OperationFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -76,7 +73,7 @@ pub fn ends_turn(messages: &[Message]) -> bool {
}
#[async_trait]
pub trait Operation: Send + Sync {
pub trait Operation<S, E: Send + 'static = ConversationEffect>: Send + Sync {
fn name(&self) -> &'static str;
/// Note on a message something this operation did, so that a pipeline rebuilt
@@ -93,50 +90,46 @@ pub trait Operation: Send + Sync {
async fn cancel(
&self,
_session: &Session,
_session: &S,
_conversation: &Conversation,
result: OperationResult,
result: OperationResult<E>,
_emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<E>> {
Ok(result)
}
async fn run_command(
&self,
_command: &SlashCommand<'_>,
_session: &Session,
_session: &S,
_conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<E>> {
not_applicable()
}
async fn inference_tools(&self, _session: &Session) -> Result<Vec<Tool>> {
async fn inference_tools(&self, _session: &S) -> Result<Vec<Tool>> {
Ok(Vec::new())
}
async fn prompt_parts(
&self,
_session: &Session,
_session: &S,
_conversation: &Conversation,
) -> Result<Vec<(String, String)>> {
Ok(Vec::new())
}
async fn moim_parts(
&self,
_session: &Session,
_conversation: &Conversation,
) -> Result<Vec<String>> {
async fn moim_parts(&self, _session: &S, _conversation: &Conversation) -> Result<Vec<String>> {
Ok(Vec::new())
}
async fn run(
&self,
_session: &Session,
_session: &S,
_conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<E>> {
not_applicable()
}
}
@@ -149,69 +142,66 @@ pub struct InferenceInput {
}
#[async_trait]
pub trait Inference: Operation {
pub trait Inference<S, E: Send + 'static = ConversationEffect>: Operation<S, E> {
/// Whether the next step would reach the provider. The machine asks before
/// firing the hooks that mark the start of a turn.
fn applies(&self, conversation: &Conversation) -> bool;
async fn infer(
&self,
session: &Session,
session: &S,
conversation: &Conversation,
input: InferenceInput,
emit: &Emitter,
) -> Result<OperationResult>;
) -> Result<OperationResult<E>>;
}
pub struct StepResult {
pub effects: Vec<StateEffect>,
pub struct StepResult<E = ConversationEffect> {
pub effects: Vec<E>,
pub applied_step: Option<&'static str>,
pub yield_to_client: bool,
}
impl StepResult {
pub(super) fn ensure_message_ids(&mut self) {
for effect in &mut self.effects {
effect.ensure_message_ids();
}
}
}
pub enum OperationResult {
pub enum OperationResult<E = ConversationEffect> {
NotApplicable,
Applied(StepResult),
Applied(StepResult<E>),
}
pub fn not_applicable() -> Result<OperationResult> {
pub fn not_applicable<E>() -> Result<OperationResult<E>> {
Ok(OperationResult::NotApplicable)
}
pub fn applied(effects: impl IntoIterator<Item = StateEffect>) -> Result<OperationResult> {
pub fn applied<E>(effects: impl IntoIterator<Item = E>) -> Result<OperationResult<E>> {
Ok(OperationResult::Applied(StepResult {
effects: effects.into_iter().collect(),
applied_step: None,
yield_to_client: false,
}))
}
pub fn yielded() -> Result<OperationResult> {
pub fn yielded<E>() -> Result<OperationResult<E>> {
Ok(OperationResult::Applied(StepResult {
effects: Vec::new(),
applied_step: None,
yield_to_client: true,
}))
}
pub fn yielded_with(effects: impl IntoIterator<Item = StateEffect>) -> Result<OperationResult> {
pub fn yielded_with<E>(effects: impl IntoIterator<Item = E>) -> Result<OperationResult<E>> {
Ok(OperationResult::Applied(StepResult {
effects: effects.into_iter().collect(),
applied_step: None,
yield_to_client: true,
}))
}
pub enum StateEffect {
pub trait MachineEffect {
fn ensure_message_ids(&mut self);
}
pub enum ConversationEffect {
AppendMessage(Message),
ReplaceConversation {
conversation: Conversation,
usage: Option<ProviderUsage>,
},
ReplaceConversation(Conversation),
PatchToolRequestMeta {
tool_call_id: String,
patch: serde_json::Value,
@@ -221,16 +211,13 @@ pub enum StateEffect {
user_visible: bool,
agent_visible: bool,
},
SetRecipe(Box<Option<Recipe>>),
SetExtensionData(ExtensionData),
RecordUsage(ProviderUsage),
}
impl StateEffect {
pub(super) fn ensure_message_ids(&mut self) {
impl MachineEffect for ConversationEffect {
fn ensure_message_ids(&mut self) {
let messages = match self {
StateEffect::AppendMessage(message) => std::slice::from_mut(message),
StateEffect::ReplaceConversation { conversation, .. } => {
ConversationEffect::AppendMessage(message) => std::slice::from_mut(message),
ConversationEffect::ReplaceConversation(conversation) => {
conversation.messages_mut().as_mut_slice()
}
_ => return,
@@ -243,18 +230,15 @@ impl StateEffect {
}
}
impl From<Message> for StateEffect {
impl From<Message> for ConversationEffect {
fn from(message: Message) -> Self {
StateEffect::AppendMessage(message)
ConversationEffect::AppendMessage(message)
}
}
impl From<Conversation> for StateEffect {
impl From<Conversation> for ConversationEffect {
fn from(conversation: Conversation) -> Self {
StateEffect::ReplaceConversation {
conversation,
usage: None,
}
ConversationEffect::ReplaceConversation(conversation)
}
}
crates/goose/src/agents/state_machine/ops_bang_shell.rs
+4 -3
View File
@@ -4,7 +4,8 @@ use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::CallToolRequestParams;
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
applied, last_effective_role, messages_since_kickoff, not_applicable, yielded, Emitter,
Operation, OperationResult,
};
@@ -31,7 +32,7 @@ impl BangShellOperation {
}
#[async_trait]
impl Operation for BangShellOperation {
impl Operation<Session, GooseEffect> for BangShellOperation {
fn name(&self) -> &'static str {
"bang_shell"
}
@@ -41,7 +42,7 @@ impl Operation for BangShellOperation {
_session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
let Some(kickoff) = messages.first() else {
return not_applicable();
crates/goose/src/agents/state_machine/ops_compaction.rs
+18 -13
View File
@@ -6,11 +6,12 @@ use anyhow::{anyhow, Result};
use async_trait::async_trait;
use tracing_futures::Instrument;
use crate::agents::state_machine::operation::{
applied, last_effective_role, messages_since_kickoff, not_applicable, trailing_error, yielded,
yielded_with, Emitter, Operation, OperationResult, SlashCommand, StateEffect,
};
use crate::agents::state_machine::ops_llm::{chat_span, record_chat_usage};
use crate::agents::state_machine::{
applied, last_effective_role, messages_since_kickoff, not_applicable, trailing_error, yielded,
yielded_with, ConversationEffect, Emitter, GooseEffect, Operation, OperationResult,
SlashCommand,
};
use crate::context_mgmt::compact_messages;
use crate::conversation::message::{Message, MessageErrorKind, SystemNotificationType};
use crate::conversation::{Conversation, EffectiveRole};
@@ -79,7 +80,7 @@ impl CompactionOperation {
conversation: &Conversation,
message: String,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let command = messages_since_kickoff(conversation)?
.first()
.cloned()
@@ -95,16 +96,20 @@ impl CompactionOperation {
emit.message(command).await;
let response = emit.message(response).await;
yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
])
}
async fn clear(conversation: &Conversation, emit: &Emitter) -> Result<OperationResult> {
async fn clear(
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult<GooseEffect>> {
let command = messages_since_kickoff(conversation)?
.first()
.cloned()
@@ -124,7 +129,7 @@ impl CompactionOperation {
}
#[async_trait]
impl Operation for CompactionOperation {
impl Operation<Session, GooseEffect> for CompactionOperation {
fn name(&self) -> &'static str {
"compaction"
}
@@ -135,7 +140,7 @@ impl Operation for CompactionOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
match command.command {
"clear" => return Self::clear(conversation, emit).await,
"compact" => {}
@@ -179,7 +184,7 @@ impl Operation for CompactionOperation {
emit.message(command).await;
let response = emit.message(response).await;
yielded_with([
StateEffect::ReplaceConversation {
GooseEffect::ReplaceConversation {
conversation: compacted,
usage: Some(usage),
},
@@ -209,7 +214,7 @@ impl Operation for CompactionOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if self.manages_own_context {
return not_applicable();
}
@@ -294,7 +299,7 @@ impl Operation for CompactionOperation {
"Compaction complete",
))
.await;
applied([StateEffect::ReplaceConversation {
applied([GooseEffect::ReplaceConversation {
conversation: compacted,
usage: Some(usage),
}])
crates/goose/src/agents/state_machine/ops_doctor.rs
+11 -9
View File
@@ -4,9 +4,9 @@ use anyhow::{anyhow, Result};
use async_trait::async_trait;
use rmcp::model::Role;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, yielded_with, Emitter, Operation,
OperationResult, SlashCommand, StateEffect,
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, yielded_with, ConversationEffect, Emitter,
GooseEffect, Operation, OperationResult, SlashCommand,
};
use crate::conversation::message::Message;
use crate::conversation::Conversation;
@@ -15,7 +15,7 @@ use crate::session::Session;
pub struct DoctorOperation;
#[async_trait]
impl Operation for DoctorOperation {
impl Operation<Session, GooseEffect> for DoctorOperation {
fn name(&self) -> &'static str {
"doctor"
}
@@ -26,7 +26,7 @@ impl Operation for DoctorOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if command.command != "doctor" {
return not_applicable();
}
@@ -55,21 +55,23 @@ impl Operation for DoctorOperation {
emit.message(command_message).await;
let result = emit.message(result).await;
return yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
result.into(),
]);
}
applied([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
result.with_visibility(false, true).into(),
])
}
crates/goose/src/agents/state_machine/ops_entry_hook.rs
+77
View File
@@ -0,0 +1,77 @@
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult,
};
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::hooks::{HookContext, HookEvent, HookManager};
use crate::session::Session;
pub struct EntryHookOperation {
hook_manager: HookManager,
}
impl EntryHookOperation {
pub fn new(hook_manager: HookManager) -> Self {
Self { hook_manager }
}
}
#[async_trait]
impl Operation<Session, GooseEffect> for EntryHookOperation {
fn name(&self) -> &'static str {
"entry_hook"
}
async fn run(
&self,
session: &Session,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
if messages.iter().any(|message| {
message.role == rmcp::model::Role::Assistant
&& ((message.is_user_visible() && message.is_agent_visible())
|| message.error_kind().is_some())
}) {
return not_applicable();
}
let messages_before_kickoff =
&conversation.messages()[..conversation.len() - messages.len()];
if !messages_before_kickoff.iter().any(|message| {
message.role == rmcp::model::Role::User
&& message.is_user_visible()
&& !message.is_tool_response()
}) {
self.hook_manager
.emit(
HookEvent::SessionStart,
HookContext::new(HookEvent::SessionStart, &session.id)
.with_working_dir(session.working_dir.to_string_lossy().to_string()),
)
.await;
}
let prompt = messages
.first()
.map(Message::as_concat_text)
.unwrap_or_default();
if !prompt.is_empty() {
self.hook_manager
.emit(
HookEvent::UserPromptSubmit,
HookContext::new(HookEvent::UserPromptSubmit, &session.id)
.with_message(prompt)
.with_working_dir(session.working_dir.to_string_lossy().to_string()),
)
.await;
}
not_applicable()
}
}
crates/goose/src/agents/state_machine/ops_exit_on_error.rs
+4 -3
View File
@@ -3,7 +3,8 @@
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
not_applicable, trailing_error, yielded, Emitter, Operation, OperationResult,
};
use crate::conversation::Conversation;
@@ -12,7 +13,7 @@ use crate::session::Session;
pub struct ExitOnErrorOperation;
#[async_trait]
impl Operation for ExitOnErrorOperation {
impl Operation<Session, GooseEffect> for ExitOnErrorOperation {
fn name(&self) -> &'static str {
"exit_on_error"
}
@@ -22,7 +23,7 @@ impl Operation for ExitOnErrorOperation {
_session: &Session,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if trailing_error(conversation).is_none() {
return not_applicable();
}
crates/goose/src/agents/state_machine/ops_llm.rs
+24 -20
View File
@@ -2,11 +2,12 @@
use std::sync::Arc;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, trailing_error, yielded_with, Emitter,
Inference, InferenceInput, Operation, OperationResult, SlashCommand, StateEffect,
};
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};
@@ -182,7 +183,7 @@ impl<'a> InferenceRunner<'a> {
}
}
async fn error_outcome(&self, err: &ProviderError, emit: &Emitter) -> Vec<StateEffect> {
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());
@@ -194,7 +195,7 @@ impl<'a> InferenceRunner<'a> {
}
#[async_trait]
impl Operation for InferenceRunner<'_> {
impl Operation<Session, GooseEffect> for InferenceRunner<'_> {
fn name(&self) -> &'static str {
"llm"
}
@@ -203,9 +204,9 @@ impl Operation for InferenceRunner<'_> {
&self,
_session: &Session,
conversation: &Conversation,
result: OperationResult,
result: OperationResult<GooseEffect>,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let mut answered = conversation
.messages()
.iter()
@@ -235,7 +236,9 @@ impl Operation for InferenceRunner<'_> {
}
if let OperationResult::Applied(step) = &result {
for effect in &step.effects {
if let StateEffect::AppendMessage(message) = effect {
if let GooseEffect::Conversation(ConversationEffect::AppendMessage(message)) =
effect
{
collect(message);
}
}
@@ -273,7 +276,7 @@ impl Operation for InferenceRunner<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if command.command != "status" {
return not_applicable();
}
@@ -320,18 +323,19 @@ impl Operation for InferenceRunner<'_> {
.await;
let response = emit.message(response).await;
yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
])
}
}
#[async_trait]
impl Inference for InferenceRunner<'_> {
impl Inference<Session, GooseEffect> for InferenceRunner<'_> {
fn applies(&self, conversation: &Conversation) -> bool {
let Ok(turn) = messages_since_kickoff(conversation) else {
return false;
@@ -346,7 +350,7 @@ impl Inference for InferenceRunner<'_> {
conversation: &Conversation,
mut input: InferenceInput,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
if trailing_error(conversation).is_some() {
return not_applicable();
@@ -471,8 +475,8 @@ impl Inference for InferenceRunner<'_> {
messages_for_provider.push(event.clone());
}
let conversation_for_provider = Conversation::new_unvalidated(messages_for_provider);
let mut usage_effects: Vec<StateEffect> =
turn_context.into_iter().map(StateEffect::from).collect();
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(),
@@ -519,7 +523,7 @@ impl Inference for InferenceRunner<'_> {
let (msg_opt, usage_opt) = match result {
Ok(chunk) => chunk,
Err(err) => {
usage_effects.extend(accumulator.into_iter().map(StateEffect::from));
usage_effects.extend(accumulator.into_iter().map(GooseEffect::from));
usage_effects.extend(self.error_outcome(&err, emit).await);
return applied(usage_effects);
}
@@ -527,7 +531,7 @@ impl Inference for InferenceRunner<'_> {
if let Some(usage) = usage_opt {
let span = tracing::Span::current();
record_chat_usage(&span, &usage);
usage_effects.push(StateEffect::RecordUsage(usage));
usage_effects.push(GooseEffect::RecordUsage(usage));
}
if let Some(mut chunk) = msg_opt {
if let Some(inference) = &inference {
@@ -573,7 +577,7 @@ impl Inference for InferenceRunner<'_> {
let has_recorded_usage = usage_effects
.iter()
.any(|effect| matches!(effect, StateEffect::RecordUsage(_)));
.any(|effect| matches!(effect, GooseEffect::RecordUsage(_)));
if !has_recorded_usage {
let mut usage = ProviderUsage::new(
self.model_config.model_name.clone(),
@@ -589,7 +593,7 @@ impl Inference for InferenceRunner<'_> {
)
.await?;
record_chat_usage(&tracing::Span::current(), &usage);
usage_effects.push(StateEffect::RecordUsage(usage));
usage_effects.push(GooseEffect::RecordUsage(usage));
}
}
crates/goose/src/agents/state_machine/ops_maxturns.rs
+4 -3
View File
@@ -1,6 +1,7 @@
//! Ends the turn when the agent has used its autonomous turn budget.
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
assistant_turn_count, messages_since_kickoff, not_applicable, yielded_with, Emitter, Operation,
OperationResult,
};
@@ -33,7 +34,7 @@ impl MaxTurnsOperation {
}
#[async_trait]
impl Operation for MaxTurnsOperation {
impl Operation<Session, GooseEffect> for MaxTurnsOperation {
fn name(&self) -> &'static str {
"max_turns"
}
@@ -54,7 +55,7 @@ impl Operation for MaxTurnsOperation {
_session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
if assistant_turn_count(messages) < self.max_turns {
return not_applicable();
crates/goose/src/agents/state_machine/ops_project.rs
+2 -1
View File
@@ -3,6 +3,7 @@
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::Operation;
use crate::conversation::Conversation;
use crate::session::Session;
@@ -10,7 +11,7 @@ use crate::session::Session;
pub struct ProjectOperation;
#[async_trait]
impl Operation for ProjectOperation {
impl Operation<Session, GooseEffect> for ProjectOperation {
fn name(&self) -> &'static str {
"project"
}
crates/goose/src/agents/state_machine/ops_recipe.rs
+15 -13
View File
@@ -10,13 +10,13 @@ use tracing_futures::Instrument;
use crate::agents::final_output_tool::{
FinalOutputTool, FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME,
};
use crate::agents::state_machine::operation::{
applied, ends_turn, last_effective_role, messages_since_kickoff, not_applicable, yielded_with,
Emitter, Operation, OperationResult, SlashCommand, StateEffect,
};
use crate::agents::state_machine::ops_toolcalling::{
pending_tool_requests, tool_span, ToolDisposition,
};
use crate::agents::state_machine::{
applied, ends_turn, last_effective_role, messages_since_kickoff, not_applicable, yielded_with,
ConversationEffect, Emitter, GooseEffect, Operation, OperationResult, SlashCommand,
};
use crate::conversation::message::{Message, MessageContent};
use crate::conversation::{Conversation, EffectiveRole};
use crate::session::Session;
@@ -77,7 +77,7 @@ impl RecipeOperation {
conversation: &Conversation,
message: String,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let command = messages_since_kickoff(conversation)?
.first()
.cloned()
@@ -93,18 +93,19 @@ impl RecipeOperation {
emit.message(command).await;
let response = emit.message(response).await;
yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
])
}
}
#[async_trait]
impl Operation for RecipeOperation {
impl Operation<Session, GooseEffect> for RecipeOperation {
fn name(&self) -> &'static str {
"recipe"
}
@@ -115,7 +116,7 @@ impl Operation for RecipeOperation {
_session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let (recipe, prompt) = match crate::slash_commands::recipe_slash_command::resolve_command(
command.command,
command.params_str,
@@ -144,12 +145,13 @@ impl Operation for RecipeOperation {
.clone()
.ok_or_else(|| anyhow!("Persisted slash command message has no id"))?;
applied([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
StateEffect::SetRecipe(Box::new(Some(recipe))),
}
.into(),
GooseEffect::SetRecipe(Box::new(Some(recipe))),
Message::user()
.with_text(prompt)
.with_visibility(false, true)
@@ -182,7 +184,7 @@ impl Operation for RecipeOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let Some(mut final_output) = Self::final_output(session)? else {
return not_applicable();
};
crates/goose/src/agents/state_machine/ops_retry.rs
+9 -8
View File
@@ -7,9 +7,9 @@ use std::time::Duration;
use crate::agents::retry::{
execute_on_failure_command_with_timeout, execute_success_checks_with_timeout,
};
use crate::agents::state_machine::operation::{
applied, ends_turn, messages_since_kickoff, not_applicable, yielded_with, Emitter, Operation,
OperationResult, SlashCommand, StateEffect,
use crate::agents::state_machine::{
applied, ends_turn, messages_since_kickoff, not_applicable, yielded_with, ConversationEffect,
Emitter, GooseEffect, Operation, OperationResult, SlashCommand,
};
use crate::agents::types::RetryConfig;
use crate::conversation::message::{Message, MessageErrorKind, SystemNotificationType};
@@ -83,7 +83,7 @@ impl<'a> RetryOperation<'a> {
}
#[async_trait]
impl Operation for RetryOperation<'_> {
impl Operation<Session, GooseEffect> for RetryOperation<'_> {
fn name(&self) -> &'static str {
"retry"
}
@@ -94,7 +94,7 @@ impl Operation for RetryOperation<'_> {
_session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let target = match command.command {
"goal" => &self.goal,
"grind" => &self.grind,
@@ -152,11 +152,12 @@ impl Operation for RetryOperation<'_> {
let response = emit.message(response).await;
let mut effects = vec![
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
];
if starts_turn {
@@ -179,7 +180,7 @@ impl Operation for RetryOperation<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
if !ends_turn(messages) {
return not_applicable();
crates/goose/src/agents/state_machine/ops_skills.rs
+14 -12
View File
@@ -10,13 +10,13 @@ use schemars::{schema_for, JsonSchema};
use serde::Deserialize;
use serde_json::Value;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, yielded_with, Emitter, Operation,
OperationResult, SlashCommand, StateEffect,
};
use crate::agents::state_machine::ops_toolcalling::{
pending_tool_requests, tool_span, ToolDisposition,
};
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, yielded_with, ConversationEffect, Emitter,
GooseEffect, Operation, OperationResult, SlashCommand,
};
use crate::agents::tool_execution::{CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::config::GooseMode;
use crate::conversation::message::Message;
@@ -206,7 +206,7 @@ impl SkillOperation {
conversation: &Conversation,
message: String,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let command = messages_since_kickoff(conversation)?
.first()
.cloned()
@@ -222,18 +222,19 @@ impl SkillOperation {
emit.message(command).await;
let response = emit.message(response).await;
yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
])
}
}
#[async_trait]
impl Operation for SkillOperation {
impl Operation<Session, GooseEffect> for SkillOperation {
fn name(&self) -> &'static str {
"skills"
}
@@ -244,7 +245,7 @@ impl Operation for SkillOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if command.command == "skills" {
return Self::command_response(
conversation,
@@ -273,11 +274,12 @@ impl Operation for SkillOperation {
.clone()
.ok_or_else(|| anyhow!("Persisted slash command message has no id"))?;
applied([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
Message::user()
.with_text(prompt)
.with_visibility(false, true)
@@ -305,7 +307,7 @@ impl Operation for SkillOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let pending: Vec<_> = pending_tool_requests(messages_since_kickoff(conversation)?)
.into_iter()
.filter(|(request, _)| {
crates/goose/src/agents/state_machine/ops_slash_command.rs
+6 -5
View File
@@ -5,14 +5,15 @@ use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult, SlashCommand,
};
use crate::conversation::Conversation;
use crate::session::Session;
pub struct SlashCommandOperation<'a> {
operations: Vec<Arc<dyn Operation + 'a>>,
operations: Vec<Arc<dyn Operation<Session, GooseEffect> + 'a>>,
}
fn parse_slash_command(message: &str) -> Option<SlashCommand<'_>> {
@@ -35,13 +36,13 @@ fn parse_slash_command(message: &str) -> Option<SlashCommand<'_>> {
}
impl<'a> SlashCommandOperation<'a> {
pub fn new(operations: Vec<Arc<dyn Operation + 'a>>) -> Self {
pub fn new(operations: Vec<Arc<dyn Operation<Session, GooseEffect> + 'a>>) -> Self {
Self { operations }
}
}
#[async_trait]
impl Operation for SlashCommandOperation<'_> {
impl Operation<Session, GooseEffect> for SlashCommandOperation<'_> {
fn name(&self) -> &'static str {
"slash_command"
}
@@ -51,7 +52,7 @@ impl Operation for SlashCommandOperation<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
let Some(user_message) = messages.first() else {
return not_applicable();
crates/goose/src/agents/state_machine/ops_steer.rs
+4 -3
View File
@@ -7,7 +7,8 @@ use anyhow::Result;
use async_trait::async_trait;
use tokio::sync::Mutex;
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
applied, ends_turn, last_effective_role, messages_since_kickoff, not_applicable, Emitter,
Operation, OperationResult,
};
@@ -33,7 +34,7 @@ impl SteerOperation {
}
#[async_trait]
impl Operation for SteerOperation {
impl Operation<Session, GooseEffect> for SteerOperation {
fn name(&self) -> &'static str {
"steer"
}
@@ -43,7 +44,7 @@ impl Operation for SteerOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
let between_turns =
ends_turn(messages) || last_effective_role(messages)? == EffectiveRole::Tool;
crates/goose/src/agents/state_machine/ops_stop_hook.rs
+4 -3
View File
@@ -3,7 +3,8 @@
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::operation::{
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::{
applied, ends_turn, messages_since_kickoff, not_applicable, yielded, yielded_with, Emitter,
Operation, OperationResult,
};
@@ -57,7 +58,7 @@ impl StopHookOperation {
}
#[async_trait]
impl Operation for StopHookOperation {
impl Operation<Session, GooseEffect> for StopHookOperation {
fn name(&self) -> &'static str {
"stop_hook"
}
@@ -67,7 +68,7 @@ impl Operation for StopHookOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let messages = messages_since_kickoff(conversation)?;
if !ends_turn(messages) {
return not_applicable();
crates/goose/src/agents/state_machine/ops_tool_approval.rs
+8 -7
View File
@@ -5,9 +5,9 @@ use std::collections::{HashMap, HashSet};
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult,
StateEffect,
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, ConversationEffect, Emitter, GooseEffect,
Operation, OperationResult,
};
use crate::config::permission::PermissionLevel;
use crate::config::GooseMode;
@@ -40,7 +40,7 @@ impl<'a> ToolApprovalOperation<'a> {
}
#[async_trait]
impl Operation for ToolApprovalOperation<'_> {
impl Operation<Session, GooseEffect> for ToolApprovalOperation<'_> {
fn name(&self) -> &'static str {
"tool_approval"
}
@@ -50,7 +50,7 @@ impl Operation for ToolApprovalOperation<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let goose_mode = *self.goose_mode.lock().await;
if goose_mode == GooseMode::Chat {
return not_applicable();
@@ -259,9 +259,10 @@ fn permission_allows(permission: &Permission) -> bool {
matches!(permission, Permission::AllowOnce | Permission::AlwaysAllow)
}
fn mark_executable(tool_call_id: &str, executable: bool) -> StateEffect {
StateEffect::PatchToolRequestMeta {
fn mark_executable(tool_call_id: &str, executable: bool) -> GooseEffect {
ConversationEffect::PatchToolRequestMeta {
tool_call_id: tool_call_id.to_string(),
patch: serde_json::json!({ TOOL_EXECUTABLE_KEY: executable }),
}
.into()
}
crates/goose/src/agents/state_machine/ops_tool_pair_compaction.rs
+15 -12
View File
@@ -6,11 +6,11 @@ use anyhow::Result;
use async_trait::async_trait;
use tracing_futures::Instrument;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult,
StateEffect,
};
use crate::agents::state_machine::ops_llm::chat_span;
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, ConversationEffect, Emitter, GooseEffect,
Operation, OperationResult,
};
use crate::context_mgmt::{summarize_tool_call, tool_ids_to_summarize};
use crate::conversation::message::MessageContent;
use crate::conversation::Conversation;
@@ -42,7 +42,7 @@ impl ToolPairCompactionOperation {
}
#[async_trait]
impl Operation for ToolPairCompactionOperation {
impl Operation<Session, GooseEffect> for ToolPairCompactionOperation {
fn name(&self) -> &'static str {
"tool_pair_compaction"
}
@@ -52,7 +52,7 @@ impl Operation for ToolPairCompactionOperation {
session: &Session,
conversation: &Conversation,
_emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
if !self.enabled {
return not_applicable();
}
@@ -66,7 +66,7 @@ impl Operation for ToolPairCompactionOperation {
if tool_ids.is_empty() {
return not_applicable();
}
let mut effects: Vec<StateEffect> = Vec::new();
let mut effects: Vec<GooseEffect> = Vec::new();
let mut hidden_messages: std::collections::HashSet<String> = Default::default();
for tool_id in tool_ids {
let pair: Vec<_> = conversation
@@ -146,11 +146,14 @@ impl Operation for ToolPairCompactionOperation {
continue;
};
hidden_messages.insert(message_id.clone());
effects.push(StateEffect::SetMessageVisibility {
message_id,
user_visible: message.is_user_visible(),
agent_visible: false,
});
effects.push(
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: message.is_user_visible(),
agent_visible: false,
}
.into(),
);
}
effects.push(summary.into());
}
crates/goose/src/agents/state_machine/ops_toolcalling.rs
+18 -16
View File
@@ -9,11 +9,11 @@ use rmcp::model::{CallToolRequestParams, CallToolResult, ContentBlock, ErrorData
use crate::agents::extension_manager::ExtensionManager;
use crate::agents::platform_extensions::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, yielded_with, Emitter, Operation,
OperationResult, SlashCommand, StateEffect,
};
use crate::agents::state_machine::ops_tool_approval::request_executable;
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, yielded_with, ConversationEffect, Emitter,
GooseEffect, Operation, OperationResult, SlashCommand,
};
use crate::agents::tool_execution::{
tool_stream, ToolCallResult, ToolStreamItem, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE,
};
@@ -276,19 +276,19 @@ impl<'a> ToolExecutionOperation<'a> {
.await
}
async fn extension_state_effect(&self, session: &Session) -> Result<StateEffect> {
async fn extension_state_effect(&self, session: &Session) -> Result<GooseEffect> {
let extension_configs = self.extension_manager.get_extension_configs().await;
let extensions_state = EnabledExtensionsState::new(extension_configs);
let mut extension_data = session.extension_data.clone();
extensions_state.to_extension_data(&mut extension_data)?;
Ok(StateEffect::SetExtensionData(extension_data))
Ok(GooseEffect::SetExtensionData(extension_data))
}
async fn command_response(
conversation: &Conversation,
message: String,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let command = messages_since_kickoff(conversation)?
.first()
.cloned()
@@ -304,11 +304,12 @@ impl<'a> ToolExecutionOperation<'a> {
emit.message(command).await;
let response = emit.message(response).await;
yielded_with([
StateEffect::SetMessageVisibility {
ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
},
}
.into(),
response.into(),
])
}
@@ -319,7 +320,7 @@ impl<'a> ToolExecutionOperation<'a> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let prompts = match self
.extension_manager
.list_prompts(&session.id, emit.cancel_token().clone())
@@ -373,7 +374,7 @@ impl<'a> ToolExecutionOperation<'a> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let params: Vec<_> = command.params_str.split_whitespace().collect();
let Some(prompt_name) = params.first() else {
return Self::command_response(
@@ -462,11 +463,12 @@ impl<'a> ToolExecutionOperation<'a> {
.id
.clone()
.ok_or_else(|| anyhow!("Persisted slash command message has no id"))?;
let mut effects = vec![StateEffect::SetMessageVisibility {
let mut effects = vec![ConversationEffect::SetMessageVisibility {
message_id,
user_visible: true,
agent_visible: false,
}];
}
.into()];
for (index, prompt_message) in result.messages.into_iter().enumerate() {
let message = Message::from(prompt_message);
let expected_role = if index % 2 == 0 {
@@ -573,7 +575,7 @@ fn approval_denied(permission: Option<&crate::permission::Permission>) -> bool {
}
#[async_trait]
impl Operation for ToolExecutionOperation<'_> {
impl Operation<Session, GooseEffect> for ToolExecutionOperation<'_> {
fn name(&self) -> &'static str {
"tool_execution"
}
@@ -584,7 +586,7 @@ impl Operation for ToolExecutionOperation<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
match command.command {
"prompts" => {
self.list_prompts(command, session, conversation, emit)
@@ -676,7 +678,7 @@ impl Operation for ToolExecutionOperation<'_> {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let mut pending = pending_tool_requests(messages_since_kickoff(conversation)?);
if pending.is_empty() {
return not_applicable();
crates/goose/src/agents/state_machine/ops_unknown_tool.rs
+6 -5
View File
@@ -4,12 +4,13 @@ use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{CallToolResult, ContentBlock};
use crate::agents::state_machine::operation::{
applied, messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult,
};
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::ops_toolcalling::{
pending_tool_requests, tool_span, ToolDisposition,
};
use crate::agents::state_machine::{
applied, messages_since_kickoff, not_applicable, Emitter, Operation, OperationResult,
};
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::session::Session;
@@ -19,7 +20,7 @@ pub(super) const UNCLAIMED_TOOL_ERROR: &str = "goose.unclaimed_tool";
pub struct UnknownToolOperation;
#[async_trait]
impl Operation for UnknownToolOperation {
impl Operation<Session, GooseEffect> for UnknownToolOperation {
fn name(&self) -> &'static str {
"unknown_tool"
}
@@ -29,7 +30,7 @@ impl Operation for UnknownToolOperation {
session: &Session,
conversation: &Conversation,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
let pending = pending_tool_requests(messages_since_kickoff(conversation)?);
if pending.is_empty() {
return not_applicable();
crates/goose/src/agents/state_machine/session.rs
+218
View File
@@ -0,0 +1,218 @@
use anyhow::Result;
use async_trait::async_trait;
use crate::agents::state_machine::effects::GooseEffect;
use crate::agents::state_machine::machine::{
EffectHandler, EffectUsage, MachineSession, SessionLoader,
};
use crate::agents::state_machine::operation::{ConversationEffect, Emitter, MachineEffect};
use crate::agents::state_machine::usage;
use crate::agents::AgentEvent;
use crate::conversation::Conversation;
use crate::session::{Session, SessionManager};
impl MachineSession for Session {
fn id(&self) -> &str {
&self.id
}
fn conversation(&self) -> Option<&Conversation> {
self.conversation.as_ref()
}
}
#[async_trait]
impl SessionLoader<Session> for SessionManager {
async fn load(&self, session_id: &str) -> Result<Session> {
self.get_session(session_id, true).await
}
}
#[async_trait]
impl EffectHandler<Session, GooseEffect> for SessionManager {
async fn apply_effects(
&self,
session: &Session,
effects: &mut [GooseEffect],
emit: &Emitter,
) -> Result<()> {
for effect in effects.iter_mut() {
effect.ensure_message_ids();
}
usage::enrich(session, effects);
for effect in effects.iter_mut() {
match effect {
GooseEffect::Conversation(ConversationEffect::AppendMessage(message)) => {
self.add_message(&session.id, message).await?;
}
GooseEffect::Conversation(ConversationEffect::ReplaceConversation(
conversation,
)) => {
self.replace_conversation(&session.id, conversation).await?;
self.update(&session.id)
.usage(usage::estimate_context(conversation).await?)
.apply()
.await?;
}
GooseEffect::ReplaceConversation {
conversation,
usage: replacement_usage,
} => {
if let Some(provider_usage) = replacement_usage {
usage::record(self, session, provider_usage, true).await?;
}
self.replace_conversation(&session.id, conversation).await?;
self.update(&session.id)
.usage(usage::estimate_context(conversation).await?)
.apply()
.await?;
}
GooseEffect::Conversation(ConversationEffect::PatchToolRequestMeta {
tool_call_id,
patch,
}) => {
self.update_tool_request_meta(&session.id, tool_call_id, patch.clone())
.await?;
}
GooseEffect::Conversation(ConversationEffect::SetMessageVisibility {
message_id,
user_visible,
agent_visible,
}) => {
self.update_message_metadata(&session.id, message_id, |mut metadata| {
metadata.user_visible = *user_visible;
metadata.agent_visible = *agent_visible;
metadata
})
.await?;
}
GooseEffect::SetRecipe(recipe) => {
self.update(&session.id)
.recipe(recipe.as_ref().clone())
.apply()
.await?;
}
GooseEffect::SetExtensionData(extension_data) => {
self.update(&session.id)
.extension_data(extension_data.clone())
.apply()
.await?;
}
GooseEffect::RecordUsage(provider_usage) => {
usage::record(self, session, provider_usage, false).await?;
}
}
}
for effect in effects {
match effect {
GooseEffect::Conversation(ConversationEffect::AppendMessage(message)) => {
if let Some(usage) = message
.metadata
.usage
.as_deref()
.filter(|_| !message.user_visible_content().content.is_empty())
.cloned()
{
emit.emit(AgentEvent::MessageUsage {
message_id: message.id.clone(),
usage,
})
.await;
}
}
GooseEffect::Conversation(ConversationEffect::ReplaceConversation(
conversation,
))
| GooseEffect::ReplaceConversation { conversation, .. } => {
emit.emit(AgentEvent::HistoryReplaced(conversation.clone()))
.await;
}
GooseEffect::RecordUsage(usage) => {
emit.emit(AgentEvent::Usage(usage.clone())).await
}
_ => {}
}
}
Ok(())
}
}
impl EffectUsage<GooseEffect> for SessionManager {
fn usage(
&self,
effect: &GooseEffect,
) -> Option<goose_providers::conversation::token_usage::Usage> {
match effect {
GooseEffect::RecordUsage(usage)
| GooseEffect::ReplaceConversation {
usage: Some(usage), ..
} => Some(usage.usage),
_ => None,
}
}
}
pub(crate) async fn run(
machine: &crate::agents::state_machine::StateMachine<'_, Session, GooseEffect>,
runtime: &SessionManager,
session_id: &str,
emit: &Emitter,
) -> Result<Session> {
let entry_session = runtime.load(session_id).await?;
if let Some(input) = entry_session
.conversation()
.and_then(|conversation| {
crate::agents::state_machine::messages_since_kickoff(conversation).ok()
})
.and_then(|messages| messages.first())
.map(crate::conversation::message::Message::user_visible_content)
.map(|message| message.as_concat_text())
.filter(|text| !text.is_empty())
{
tracing::Span::current().record("trace_input", input.as_str());
}
let mut turn_usage = goose_providers::conversation::token_usage::Usage::default();
loop {
let session = runtime.load(session_id).await?;
let Some(mut result) = machine.step(&session, emit).await? else {
break;
};
tracing::debug!(target: "goose::state_machine", step = result.applied_step, "applied step");
for effect in &result.effects {
if let Some(usage) = runtime.usage(effect) {
turn_usage += usage;
}
}
machine.apply(runtime, &session, &mut result, emit).await?;
if result.yield_to_client {
break;
}
}
let session = runtime.load(session_id).await?;
let last_assistant_text = session
.conversation()
.and_then(|conversation| {
crate::agents::state_machine::messages_since_kickoff(conversation).ok()
})
.into_iter()
.flatten()
.rev()
.filter(|message| message.role == rmcp::model::Role::Assistant)
.map(crate::conversation::message::Message::user_visible_content)
.map(|message| message.as_concat_text())
.find(|text| !text.is_empty())
.unwrap_or_default();
if !last_assistant_text.is_empty() {
let span = tracing::Span::current();
span.record("trace_output", last_assistant_text.as_str());
if crate::agents::gen_ai_telemetry::capture_message_content() {
let output = crate::agents::gen_ai_telemetry::simple_output_json(&last_assistant_text);
span.record("gen_ai.output.messages", output.as_str());
}
}
crate::agents::gen_ai_telemetry::record_usage(&tracing::Span::current(), &turn_usage);
Ok(session)
}
crates/goose/src/agents/state_machine/tests/compaction_lifecycle.rs
+3 -2
View File
@@ -81,10 +81,11 @@ async fn proactive_and_manual_compaction_continue_with_replaced_usage() -> Resul
state_machine::StateMachine::new(Vec::new(), tokio_util::sync::CancellationToken::new());
let (tx, _rx) = tokio::sync::mpsc::channel(4);
let emit = state_machine::Emitter::new(tx, tokio_util::sync::CancellationToken::new());
let apply = async |effects: Vec<state_machine::StateEffect>| -> Result<()> {
let apply = async |effects: Vec<state_machine::GooseEffect>| -> Result<()> {
let session = pipeline.session().await?;
let mut result = state_machine::StepResult {
effects,
applied_step: None,
yield_to_client: false,
};
machine
@@ -112,7 +113,7 @@ async fn proactive_and_manual_compaction_continue_with_replaced_usage() -> Resul
);
apply(vec![
replacement.into(),
state_machine::StateEffect::RecordUsage(usage),
state_machine::GooseEffect::RecordUsage(usage),
Message::assistant()
.with_text("response after replacement")
.into(),
crates/goose/src/agents/state_machine/tests/pipeline.rs
+19 -14
View File
@@ -15,11 +15,11 @@ use crate::agents::extension_manager::{ExtensionManager, ExtensionManagerCapabil
use crate::agents::mcp_client::McpClientTrait;
use crate::agents::prompt_manager::PromptManager;
use crate::agents::state_machine::{
BangShellOperation, CompactionOperation, DoctorOperation, Emitter, ExitOnErrorOperation,
InferenceRunner, MaxTurnsOperation, Operation, ProjectOperation, RecipeOperation,
RetryOperation, SkillOperation, SlashCommandOperation, StateMachine, SteerOperation,
SteerQueue, Step, StopHookOperation, ToolApprovalOperation, ToolExecutionOperation,
ToolPairCompactionOperation, UnknownToolOperation,
BangShellOperation, CompactionOperation, DoctorOperation, Emitter, EntryHookOperation,
ExitOnErrorOperation, GooseEffect, InferenceRunner, MaxTurnsOperation, Operation,
ProjectOperation, RecipeOperation, RetryOperation, SkillOperation, SlashCommandOperation,
StateMachine, SteerOperation, SteerQueue, Step, StopHookOperation, ToolApprovalOperation,
ToolExecutionOperation, ToolPairCompactionOperation, UnknownToolOperation,
};
use crate::agents::AgentEvent;
use crate::config::permission::{PermissionLevel, PermissionManager};
@@ -105,13 +105,16 @@ pub(super) struct TestPipeline {
}
impl TestPipeline {
pub(super) fn machine(&self, cancel: CancellationToken) -> StateMachine<'_> {
pub(super) fn machine(
&self,
cancel: CancellationToken,
) -> StateMachine<'_, Session, GooseEffect> {
let provider = self.provider.clone();
let tool_call_cutoff = crate::context_mgmt::compute_tool_call_cutoff(
self.model_config.context_limit(),
COMPACTION_THRESHOLD,
);
let operations: Vec<Arc<dyn Operation + '_>> = vec![
let operations: Vec<Arc<dyn Operation<Session, GooseEffect> + '_>> = vec![
Arc::new(SteerOperation::new(
self.steer_queue.clone(),
self.hook_manager.clone(),
@@ -167,15 +170,17 @@ impl TestPipeline {
));
let mut command_handlers = operations.clone();
command_handlers.push(inference.clone());
let command_operation: Arc<dyn Operation + '_> =
let command_operation: Arc<dyn Operation<Session, GooseEffect> + '_> =
Arc::new(SlashCommandOperation::new(command_handlers));
let steps = std::iter::once(command_operation)
.chain(operations)
.map(Step::Operation)
.chain(std::iter::once(Step::Inference(inference)))
.collect();
let steps = std::iter::once(Arc::new(EntryHookOperation::new(self.hook_manager.clone()))
as Arc<dyn Operation<Session, GooseEffect> + '_>)
.chain(std::iter::once(command_operation))
.chain(operations)
.map(Step::Operation)
.chain(std::iter::once(Step::Inference(inference)))
.collect();
StateMachine::new(steps, cancel).with_hook_manager(self.hook_manager.clone())
StateMachine::new(steps, cancel)
}
pub(super) async fn with_goose_mode(self, mode: GooseMode) -> Self {
crates/goose/src/agents/state_machine/usage.rs
+8 -8
View File
@@ -1,14 +1,14 @@
use anyhow::Result;
use goose_providers::conversation::token_usage::{CostSource, ProviderUsage, Usage as TokenUsage};
use crate::agents::state_machine::operation::StateEffect;
use crate::agents::state_machine::{ConversationEffect, GooseEffect};
use crate::conversation::message::MessageUsage;
use crate::conversation::Conversation;
use crate::session::{Session, SessionManager};
fn attach_to_last_assistant(effects: &mut [StateEffect], usage: &ProviderUsage) {
fn attach_to_last_assistant(effects: &mut [GooseEffect], usage: &ProviderUsage) {
let Some(message) = effects.iter_mut().rev().find_map(|effect| match effect {
StateEffect::AppendMessage(message)
GooseEffect::Conversation(ConversationEffect::AppendMessage(message))
if message.role == rmcp::model::Role::Assistant && message.error_kind().is_none() =>
{
Some(message)
@@ -20,11 +20,11 @@ fn attach_to_last_assistant(effects: &mut [StateEffect], usage: &ProviderUsage)
message.metadata.usage = Some(Box::new(MessageUsage::from_provider_usage(usage, false)));
}
pub(super) fn enrich(session: &Session, effects: &mut [StateEffect]) {
pub(super) fn enrich(session: &Session, effects: &mut [GooseEffect]) {
for index in 0..effects.len() {
let (usage, replaces_conversation) = match &effects[index] {
StateEffect::RecordUsage(usage) => (usage.clone(), false),
StateEffect::ReplaceConversation {
GooseEffect::RecordUsage(usage) => (usage.clone(), false),
GooseEffect::ReplaceConversation {
usage: Some(usage), ..
} => (usage.clone(), true),
_ => continue,
@@ -53,8 +53,8 @@ pub(super) fn enrich(session: &Session, effects: &mut [StateEffect]) {
attach_to_last_assistant(effects, &enriched);
}
match &mut effects[index] {
StateEffect::RecordUsage(usage) => *usage = enriched,
StateEffect::ReplaceConversation { usage, .. } => *usage = Some(enriched),
GooseEffect::RecordUsage(usage) => *usage = enriched,
GooseEffect::ReplaceConversation { usage, .. } => *usage = Some(enriched),
_ => {}
}
}
crates/goose/tests/state_machine_api.rs
+16 -103
View File
@@ -1,10 +1,9 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use goose::agents::state_machine::{
yielded_with, Emitter, Inference, InferenceInput, Operation, OperationResult, StateEffect,
yielded_with, Emitter, GooseEffect, Inference, InferenceInput, Operation, OperationResult,
StateMachine, Step,
};
use goose::agents::AgentEvent;
@@ -18,18 +17,11 @@ use goose_providers::conversation::token_usage::Usage;
use goose_providers::model::ModelConfig;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::field::{Field, Visit};
use tracing::span::{Attributes, Id, Record};
use tracing::Subscriber;
use tracing_futures::Instrument;
use tracing_subscriber::layer::{Context, SubscriberExt};
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::Layer;
struct PromptPart;
#[async_trait]
impl Operation for PromptPart {
impl Operation<Session, GooseEffect> for PromptPart {
fn name(&self) -> &'static str {
"prompt_part"
}
@@ -46,14 +38,14 @@ impl Operation for PromptPart {
struct TestInference;
#[async_trait]
impl Operation for TestInference {
impl Operation<Session, GooseEffect> for TestInference {
fn name(&self) -> &'static str {
"test_inference"
}
}
#[async_trait]
impl Inference for TestInference {
impl Inference<Session, GooseEffect> for TestInference {
fn applies(&self, _conversation: &Conversation) -> bool {
true
}
@@ -64,7 +56,7 @@ impl Inference for TestInference {
conversation: &Conversation,
input: InferenceInput,
emit: &Emitter,
) -> Result<OperationResult> {
) -> Result<OperationResult<GooseEffect>> {
assert_eq!(
input.prompt_parts,
[("test".to_string(), "custom context".to_string())]
@@ -80,8 +72,8 @@ impl Inference for TestInference {
.message(Message::assistant().with_text(format!("{prompt} answered")))
.await;
yielded_with([
StateEffect::from(message),
StateEffect::RecordUsage(ProviderUsage::new(
GooseEffect::from(message),
GooseEffect::RecordUsage(ProviderUsage::new(
"test-model".to_string(),
Usage::new(Some(5), Some(7), Some(12)),
)),
@@ -89,52 +81,8 @@ impl Inference for TestInference {
}
}
#[derive(Clone, Default)]
struct TraceFields(Arc<Mutex<HashMap<String, String>>>);
impl<S> Layer<S> for TraceFields
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
if context
.span(id)
.is_some_and(|span| span.metadata().name() == "state_machine_test")
{
attributes.record(&mut FieldVisitor(self.0.clone()));
}
}
fn on_record(&self, id: &Id, values: &Record<'_>, context: Context<'_, S>) {
if context
.span(id)
.is_some_and(|span| span.metadata().name() == "state_machine_test")
{
values.record(&mut FieldVisitor(self.0.clone()));
}
}
}
struct FieldVisitor(Arc<Mutex<HashMap<String, String>>>);
impl Visit for FieldVisitor {
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
self.0
.lock()
.unwrap()
.insert(field.name().to_string(), format!("{value:?}"));
}
fn record_str(&mut self, field: &Field, value: &str) {
self.0
.lock()
.unwrap()
.insert(field.name().to_string(), value.to_string());
}
}
#[tokio::test]
async fn custom_pipeline_supports_step_apply_run_tracing_and_usage() -> Result<()> {
async fn custom_pipeline_supports_step_apply_run_and_usage() -> Result<()> {
let _env = goose_test_support::otel::clear_otel_env(&[(
"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
"true",
@@ -171,6 +119,12 @@ async fn custom_pipeline_supports_step_apply_run_tracing_and_usage() -> Result<(
let session = session_manager.get_session(&session.id, true).await?;
let mut result = machine.step(&session, &emit).await?.unwrap();
assert!(matches!(
result.effects.first(),
Some(GooseEffect::Conversation(
goose::agents::state_machine::ConversationEffect::AppendMessage(message)
)) if message.id.is_some()
));
machine
.apply(&session_manager, &session, &mut result, &emit)
.await?;
@@ -191,21 +145,7 @@ async fn custom_pipeline_supports_step_apply_run_tracing_and_usage() -> Result<(
session_manager
.add_message(&session.id, &Message::user().with_text("second turn"))
.await?;
let fields = TraceFields::default();
let subscriber = tracing_subscriber::registry().with(fields.clone());
let _guard = tracing::subscriber::set_default(subscriber);
let span = tracing::info_span!(
"state_machine_test",
trace_input = tracing::field::Empty,
trace_output = tracing::field::Empty,
gen_ai.output.messages = tracing::field::Empty,
gen_ai.usage.input_tokens = tracing::field::Empty,
gen_ai.usage.output_tokens = tracing::field::Empty,
);
let session = machine
.run(&session_manager, &session.id, &emit)
.instrument(span)
.await?;
let session = machine.run(&session_manager, &session.id, &emit).await?;
assert_eq!(
session
@@ -231,32 +171,5 @@ async fn custom_pipeline_supports_step_apply_run_tracing_and_usage() -> Result<(
assert_eq!(terminal_usage.total_tokens, 12);
assert_eq!(terminal_usage.accumulated_total_tokens, 24);
let fields = fields.0.lock().unwrap();
assert_eq!(
fields.get("trace_input").map(String::as_str),
Some("second turn")
);
assert_eq!(
fields.get("trace_output").map(String::as_str),
Some("second turn answered")
);
assert_eq!(
fields.get("gen_ai.usage.input_tokens").map(String::as_str),
Some("5")
);
assert_eq!(
fields.get("gen_ai.usage.output_tokens").map(String::as_str),
Some("7")
);
let output: serde_json::Value = serde_json::from_str(&fields["gen_ai.output.messages"])?;
assert_eq!(
output,
serde_json::json!([{
"role": "assistant",
"content": "second turn answered",
"finish_reason": "stop",
}])
);
Ok(())
}