goose_agent/
operation.rs

1use anyhow::{anyhow, Result};
2use async_trait::async_trait;
3use futures::Stream;
4use std::future::Future;
5use std::pin::Pin;
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use crate::events::AgentEvent;
10use goose_provider_types::conversation::message::{Message, MessageContent, MessageErrorKind};
11use goose_provider_types::conversation::{effective_role, Conversation, EffectiveRole};
12use goose_provider_types::maybe_send::{MaybeSend, MaybeSync};
13use rmcp::model::Tool;
14
15#[cfg(not(target_arch = "wasm32"))]
16pub type OperationFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
17
18#[cfg(target_arch = "wasm32")]
19pub type OperationFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
20
21#[cfg(not(target_arch = "wasm32"))]
22pub type OperationStream<'a, T> = Pin<Box<dyn Stream<Item = T> + Send + 'a>>;
23
24#[cfg(target_arch = "wasm32")]
25pub type OperationStream<'a, T> = Pin<Box<dyn Stream<Item = T> + 'a>>;
26
27pub const CLIENT_LOG: &str = "client_log";
28
29pub struct SlashCommand<'a> {
30    pub command: &'a str,
31    pub params_str: &'a str,
32}
33
34pub fn messages_since_kickoff(conversation: &Conversation) -> Result<&[Message]> {
35    let messages = conversation.messages();
36    let is_kickoff = |msg: &Message| msg.role == rmcp::model::Role::User && !msg.is_tool_response();
37    let start = messages
38        .iter()
39        .rposition(|msg| is_kickoff(msg) && msg.is_user_visible())
40        .or_else(|| {
41            messages
42                .iter()
43                .rposition(|msg| is_kickoff(msg) && msg.is_agent_visible())
44        })
45        .ok_or_else(|| anyhow!("state machine conversation has no kickoff message"))?;
46    Ok(&messages[start..])
47}
48
49pub fn trailing_error(conversation: &Conversation) -> Option<MessageErrorKind> {
50    conversation.last().and_then(Message::error_kind)
51}
52
53pub fn last_effective_role(messages: &[Message]) -> Result<EffectiveRole> {
54    messages
55        .last()
56        .map(effective_role)
57        .ok_or_else(|| anyhow!("cannot determine the role of an empty conversation"))
58}
59
60pub fn assistant_turn_count(messages: &[Message]) -> u32 {
61    let mut turns = 0;
62    let mut in_assistant_block = false;
63    for message in messages.iter().rev() {
64        if message.role == rmcp::model::Role::Assistant {
65            if !in_assistant_block {
66                turns += 1;
67                in_assistant_block = true;
68            }
69        } else {
70            in_assistant_block = false;
71        }
72    }
73    turns
74}
75
76pub fn ends_turn(messages: &[Message]) -> bool {
77    messages.last().is_some_and(|last| {
78        last.role == rmcp::model::Role::Assistant
79            && last.error_kind().is_none()
80            && !last.content.iter().any(|content| {
81                matches!(
82                    content,
83                    MessageContent::ToolRequest(_) | MessageContent::ActionRequired(_)
84                )
85            })
86    })
87}
88
89#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
90#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
91pub trait Operation<S, E: MaybeSend + 'static = ConversationEffect>: MaybeSend + MaybeSync {
92    fn name(&self) -> &'static str;
93
94    /// Note on a message something this operation did, so that a pipeline rebuilt
95    /// from the persisted conversation reaches the same conclusion. Notes record
96    /// past actions only — anything an operation would have to compute again does
97    /// not belong here.
98    fn set_message_meta(&self, message: &mut Message, key: &str, value: serde_json::Value) {
99        message.metadata.set_operation_note(self.name(), key, value);
100    }
101
102    fn message_meta<'a>(&self, message: &'a Message, key: &str) -> Option<&'a serde_json::Value> {
103        message.metadata.operation_note(self.name(), key)
104    }
105
106    async fn finalize_cancellation(
107        &self,
108        _session: &S,
109        _conversation: &Conversation,
110        _emit: &Emitter,
111    ) -> Vec<E> {
112        Vec::new()
113    }
114
115    async fn run_command(
116        &self,
117        _command: &SlashCommand<'_>,
118        _session: &S,
119        _conversation: &Conversation,
120        _emit: &Emitter,
121    ) -> Result<OperationResult<E>> {
122        not_applicable()
123    }
124
125    async fn inference_tools(&self, _session: &S) -> Result<Vec<Tool>> {
126        Ok(Vec::new())
127    }
128
129    async fn prompt_parts(
130        &self,
131        _session: &S,
132        _conversation: &Conversation,
133    ) -> Result<Vec<(String, String)>> {
134        Ok(Vec::new())
135    }
136
137    async fn moim_parts(&self, _session: &S, _conversation: &Conversation) -> Result<Vec<String>> {
138        Ok(Vec::new())
139    }
140
141    async fn run(
142        &self,
143        _session: &S,
144        _conversation: &Conversation,
145        _emit: &Emitter,
146    ) -> Result<OperationResult<E>> {
147        not_applicable()
148    }
149}
150
151#[derive(Default)]
152pub struct InferenceInput {
153    pub tools: Vec<Tool>,
154    pub prompt_parts: Vec<(String, String)>,
155    pub moim_parts: Vec<String>,
156}
157
158#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
159#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
160pub trait Inference<S, E: MaybeSend + 'static = ConversationEffect>: Operation<S, E> {
161    /// Whether the next step would reach the provider. The machine asks before
162    /// firing the hooks that mark the start of a turn.
163    fn applies(&self, conversation: &Conversation) -> bool;
164
165    async fn prepare_session(&self, _session: &S) -> Result<Option<S>> {
166        Ok(None)
167    }
168
169    async fn infer(
170        &self,
171        session: &S,
172        conversation: &Conversation,
173        input: InferenceInput,
174        emit: &Emitter,
175    ) -> Result<OperationResult<E>>;
176}
177
178#[derive(Debug, Clone, Copy, PartialEq, Eq)]
179pub enum RunStatus {
180    Continuing,
181    Yielded,
182    Cancelled,
183}
184
185pub struct StepResult<E = ConversationEffect> {
186    pub effects: Vec<E>,
187    pub applied_step: Option<&'static str>,
188    pub status: RunStatus,
189}
190
191pub enum OperationResult<E = ConversationEffect> {
192    NotApplicable,
193    Applied(StepResult<E>),
194}
195
196pub fn not_applicable<E>() -> Result<OperationResult<E>> {
197    Ok(OperationResult::NotApplicable)
198}
199
200pub fn applied<E>(effects: impl IntoIterator<Item = E>) -> Result<OperationResult<E>> {
201    Ok(OperationResult::Applied(StepResult {
202        effects: effects.into_iter().collect(),
203        applied_step: None,
204        status: RunStatus::Continuing,
205    }))
206}
207
208pub fn yielded<E>() -> Result<OperationResult<E>> {
209    Ok(OperationResult::Applied(StepResult {
210        effects: Vec::new(),
211        applied_step: None,
212        status: RunStatus::Yielded,
213    }))
214}
215
216pub fn yielded_with<E>(effects: impl IntoIterator<Item = E>) -> Result<OperationResult<E>> {
217    Ok(OperationResult::Applied(StepResult {
218        effects: effects.into_iter().collect(),
219        applied_step: None,
220        status: RunStatus::Yielded,
221    }))
222}
223
224pub trait MachineEffect {
225    fn ensure_message_ids(&mut self);
226}
227
228pub enum ConversationEffect {
229    AppendMessage(Message),
230    ReplaceConversation(Conversation),
231    PatchToolRequestMeta {
232        tool_call_id: String,
233        patch: serde_json::Value,
234    },
235    SetMessageOperationNote {
236        message_id: String,
237        operation: String,
238        key: String,
239        value: serde_json::Value,
240    },
241    SetMessageVisibility {
242        message_id: String,
243        user_visible: bool,
244        agent_visible: bool,
245    },
246}
247
248impl MachineEffect for ConversationEffect {
249    fn ensure_message_ids(&mut self) {
250        let messages = match self {
251            ConversationEffect::AppendMessage(message) => std::slice::from_mut(message),
252            ConversationEffect::ReplaceConversation(conversation) => {
253                conversation.messages_mut().as_mut_slice()
254            }
255            _ => return,
256        };
257        for message in messages {
258            if message.id.is_none() {
259                message.id = Some(format!("msg_{}", uuid::Uuid::new_v4()));
260            }
261        }
262    }
263}
264
265impl From<Message> for ConversationEffect {
266    fn from(message: Message) -> Self {
267        ConversationEffect::AppendMessage(message)
268    }
269}
270
271impl From<Conversation> for ConversationEffect {
272    fn from(conversation: Conversation) -> Self {
273        ConversationEffect::ReplaceConversation(conversation)
274    }
275}
276
277pub struct Emitter {
278    tx: mpsc::UnboundedSender<AgentEvent>,
279    cancel: CancellationToken,
280}
281
282impl Emitter {
283    pub fn new(tx: mpsc::UnboundedSender<AgentEvent>, cancel: CancellationToken) -> Self {
284        Self { tx, cancel }
285    }
286
287    pub fn emit(&self, event: AgentEvent) {
288        let _ = self.tx.send(event);
289    }
290
291    pub fn message(&self, message: Message) -> Message {
292        let message = message.with_generated_id_if_missing();
293        self.emit(AgentEvent::Message(message.clone()));
294        message
295    }
296
297    pub fn cancel_token(&self) -> &CancellationToken {
298        &self.cancel
299    }
300
301    pub async fn cancelled(&self) {
302        self.cancel.cancelled().await
303    }
304}
goose_agent/operation.rs

goose_agent/
operation.rs
