goose_agent/
subagent.rs

1use std::collections::HashSet;
2use std::num::NonZeroUsize;
3use std::sync::Arc;
4
5use anyhow::Result;
6use async_trait::async_trait;
7use futures::stream::{self, SelectAll, StreamExt};
8use futures::FutureExt;
9use goose_provider_types::conversation::message::Message;
10use goose_provider_types::conversation::Conversation;
11use goose_provider_types::maybe_send::{MaybeSend, MaybeSync};
12use tokio::sync::Mutex;
13use tokio_util::sync::CancellationToken;
14
15use crate::operation::{
16    applied, not_applicable, Emitter, Operation, OperationResult, OperationStream,
17};
18
19const OPERATION_NAME: &str = "foreground_subagent";
20const DELIVERED: &str = "delivered";
21const CANCELLED: &str = "cancelled";
22const ENDED_WITHOUT_OUTCOME: &str = "it ended without an outcome";
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum SubagentOutcome {
26    Completed(String),
27    Failed(String),
28    Cancelled,
29}
30
31#[non_exhaustive]
32#[derive(Debug, Clone)]
33pub enum SubagentUpdate {
34    ToParent(Message),
35    Finished(SubagentOutcome),
36}
37
38#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
39#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
40pub trait SubagentRunner<S>: MaybeSend + MaybeSync {
41    /// Returns every subagent session ID that `message` asks to run, even if that subagent
42    /// has already run. `SubagentOperation` calls this for every message in the parent
43    /// conversation and decides which still need to run.
44    fn requested_subagent_session_ids(&self, parent_session: &S, message: &Message) -> Vec<String>;
45
46    async fn start(
47        &self,
48        parent_session: &S,
49        subagent_session_id: &str,
50        cancel: CancellationToken,
51    ) -> OperationStream<'static, SubagentUpdate>;
52
53    fn on_subagent_finished(
54        &self,
55        _subagent_session_id: &str,
56        _outcome: &SubagentOutcome,
57        _remaining: usize,
58        _emit: &Emitter,
59    ) {
60    }
61}
62
63fn finished_subagent_session_ids(message: &Message) -> impl Iterator<Item = &str> {
64    let delivered = message
65        .metadata
66        .operation_note(OPERATION_NAME, DELIVERED)
67        .and_then(serde_json::Value::as_str);
68    let cancelled = message
69        .metadata
70        .operation_note(OPERATION_NAME, CANCELLED)
71        .and_then(serde_json::Value::as_array)
72        .into_iter()
73        .flatten()
74        .filter_map(serde_json::Value::as_str);
75    delivered.into_iter().chain(cancelled)
76}
77
78fn subagent_label(subagent_session_id: &str) -> String {
79    format!("Subagent {subagent_session_id}")
80}
81
82fn subagent_result_message(
83    subagent_session_id: &str,
84    outcome: &SubagentOutcome,
85) -> Option<Message> {
86    let label = subagent_label(subagent_session_id);
87    let text = match outcome {
88        SubagentOutcome::Completed(output) => format!("{label} completed: {output}"),
89        SubagentOutcome::Failed(reason) => format!("{label} failed: {reason}"),
90        SubagentOutcome::Cancelled => {
91            return subagent_cancelled_message(&[subagent_session_id.to_string()])
92        }
93    };
94    let mut message = Message::user().with_text(text).with_visibility(false, true);
95    message.metadata.set_operation_note(
96        OPERATION_NAME,
97        DELIVERED,
98        serde_json::json!(subagent_session_id),
99    );
100    Some(message)
101}
102
103fn subagent_cancelled_message(subagent_session_ids: &[String]) -> Option<Message> {
104    if subagent_session_ids.is_empty() {
105        return None;
106    }
107    let text = subagent_session_ids
108        .iter()
109        .map(|subagent_session_id| {
110            format!(
111                "{} was cancelled before it finished.",
112                subagent_label(subagent_session_id)
113            )
114        })
115        .collect::<Vec<_>>()
116        .join("\n");
117    let mut message = Message::user().with_text(text).with_visibility(false, true);
118    message.metadata.set_operation_note(
119        OPERATION_NAME,
120        CANCELLED,
121        serde_json::json!(subagent_session_ids),
122    );
123    Some(message)
124}
125
126fn tracked_updates(
127    subagent_session_id: String,
128    updates: OperationStream<'static, SubagentUpdate>,
129) -> OperationStream<'static, (String, SubagentUpdate)> {
130    Box::pin(stream::unfold(Some(updates), move |updates| {
131        let subagent_session_id = subagent_session_id.clone();
132        async move {
133            let mut updates = updates?;
134            let update = updates.next().await.unwrap_or_else(|| {
135                SubagentUpdate::Finished(SubagentOutcome::Failed(ENDED_WITHOUT_OUTCOME.to_string()))
136            });
137            let rest = match update {
138                SubagentUpdate::Finished(_) => None,
139                SubagentUpdate::ToParent(_) => Some(updates),
140            };
141            Some(((subagent_session_id, update), rest))
142        }
143    }))
144}
145
146#[derive(Default)]
147struct Running {
148    subagent_session_ids: HashSet<String>,
149    subagent_runs: SelectAll<OperationStream<'static, (String, SubagentUpdate)>>,
150}
151
152pub struct SubagentOperation<S> {
153    runner: Arc<dyn SubagentRunner<S>>,
154    running: Mutex<Running>,
155    max_concurrency: NonZeroUsize,
156}
157
158impl<S> SubagentOperation<S> {
159    pub fn new(runner: Arc<dyn SubagentRunner<S>>) -> Self {
160        Self {
161            runner,
162            running: Mutex::default(),
163            max_concurrency: NonZeroUsize::MAX,
164        }
165    }
166
167    /// Limits how many subagents run at once; additional subagents remain pending.
168    pub fn with_max_concurrency(mut self, max_concurrency: NonZeroUsize) -> Self {
169        self.max_concurrency = max_concurrency;
170        self
171    }
172
173    fn pending_subagents(&self, session: &S, conversation: &Conversation) -> Vec<String> {
174        let mut pending: Vec<String> = Vec::new();
175        for message in conversation.messages() {
176            for subagent_session_id in finished_subagent_session_ids(message) {
177                pending.retain(|pending_id| pending_id != subagent_session_id);
178            }
179            for subagent_session_id in self.runner.requested_subagent_session_ids(session, message)
180            {
181                if !pending.contains(&subagent_session_id) {
182                    pending.push(subagent_session_id);
183                }
184            }
185        }
186        pending
187    }
188
189    fn deliver<E: From<Message>>(
190        &self,
191        subagent_session_id: &str,
192        outcome: &SubagentOutcome,
193        remaining: usize,
194        emit: &Emitter,
195    ) -> Option<E> {
196        let message = subagent_result_message(subagent_session_id, outcome)?;
197        self.runner
198            .on_subagent_finished(subagent_session_id, outcome, remaining, emit);
199        Some(E::from(message))
200    }
201}
202
203#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
204#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
205impl<S, E> Operation<S, E> for SubagentOperation<S>
206where
207    S: MaybeSend + MaybeSync + 'static,
208    E: From<Message> + MaybeSend + 'static,
209{
210    fn name(&self) -> &'static str {
211        OPERATION_NAME
212    }
213
214    async fn finalize_cancellation(
215        &self,
216        session: &S,
217        conversation: &Conversation,
218        emit: &Emitter,
219    ) -> Vec<E> {
220        let mut stopped = std::mem::take(&mut *self.running.lock().await);
221        let pending = self.pending_subagents(session, conversation);
222        let mut effects = Vec::new();
223        let mut delivered = HashSet::new();
224        let finished = std::iter::from_fn(|| loop {
225            match stopped.subagent_runs.next().now_or_never().flatten()? {
226                (subagent_session_id, SubagentUpdate::Finished(outcome)) => {
227                    return Some((subagent_session_id, outcome))
228                }
229                (_, SubagentUpdate::ToParent(_)) => {}
230            }
231        });
232        for (subagent_session_id, outcome) in finished {
233            if outcome == SubagentOutcome::Cancelled {
234                continue;
235            }
236            let remaining = pending.len().saturating_sub(delivered.len() + 1);
237            if let Some(effect) = self.deliver(&subagent_session_id, &outcome, remaining, emit) {
238                delivered.insert(subagent_session_id);
239                effects.push(effect);
240            }
241        }
242        drop(stopped);
243
244        let cancelled: Vec<String> = pending
245            .into_iter()
246            .filter(|subagent_session_id| !delivered.contains(subagent_session_id))
247            .collect();
248        effects.extend(subagent_cancelled_message(&cancelled).map(E::from));
249        effects
250    }
251
252    async fn run(
253        &self,
254        session: &S,
255        conversation: &Conversation,
256        emit: &Emitter,
257    ) -> Result<OperationResult<E>> {
258        let pending = self.pending_subagents(session, conversation);
259        let remaining = pending.len().saturating_sub(1);
260
261        let mut running = self.running.lock().await;
262        for subagent_session_id in &pending {
263            if running.subagent_session_ids.len() >= self.max_concurrency.get() {
264                break;
265            }
266            if running.subagent_session_ids.contains(subagent_session_id) {
267                continue;
268            }
269            let updates = self
270                .runner
271                .start(session, subagent_session_id, emit.cancel_token().clone())
272                .await;
273            running
274                .subagent_session_ids
275                .insert(subagent_session_id.clone());
276            running
277                .subagent_runs
278                .push(tracked_updates(subagent_session_id.clone(), updates));
279        }
280
281        while let Some((subagent_session_id, update)) = running.subagent_runs.next().await {
282            let outcome = match update {
283                SubagentUpdate::ToParent(message) => {
284                    emit.message(message);
285                    continue;
286                }
287                SubagentUpdate::Finished(outcome) => outcome,
288            };
289            running.subagent_session_ids.remove(&subagent_session_id);
290            if let Some(effect) = self.deliver(&subagent_session_id, &outcome, remaining, emit) {
291                return applied([effect]);
292            }
293        }
294        not_applicable()
295    }
296}
goose_agent/subagent.rs

goose_agent/
subagent.rs
