ag_ui/server/emit/
subagent.rs

1//! Scoping a subagent invocation.
2
3use std::ops::{Deref, DerefMut};
4
5use crate::{Event, SubagentFinishedEvent, SubagentOutcome, SubagentRunId, SubagentStartedEvent};
6use serde_json::Value;
7
8use crate::server::context::RunContext;
9use crate::server::error::Result;
10
11/// One open subagent invocation.
12///
13/// Created by [`RunContext::subagent`](crate::server::RunContext::subagent).
14/// `SUBAGENT_STARTED` has already gone out; `Drop` emits `SUBAGENT_FINISHED`
15/// with a success outcome, including on the early return that a `?`
16/// produces.
17///
18/// A subagent is a *scope*, like a step: everything emitted through the
19/// handle comes out attributed to it. The handle therefore dereferences to
20/// the run context — messages, tool calls, reasoning, steps and nested
21/// subagents all open through it, and every event they produce carries this
22/// invocation's `subagentRunId` without the agent saying so:
23///
24/// ```
25/// # use ag_ui::{Event, EventType, RunAgentInput, TextMessageRole};
26/// # use ag_ui::server::RunContext;
27/// # let (mut ctx, mut events) = RunContext::<()>::new(RunAgentInput::new("t", "r"))?;
28/// {
29///     let mut researcher = ctx.subagent("researcher")?;
30///     researcher.say("Three sources found.")?;   // attributed, through Deref
31///     researcher.finish_with(serde_json::json!({ "sources": 3 }))?;
32/// }
33/// ctx.say("Thanks.")?;                            // the parent's own, untagged
34///
35/// let events = events.drain();
36/// assert_eq!(events[0].event_type(), EventType::SubagentStarted);
37/// assert_eq!(events[1].subagent_run_id().map(|id| id.as_str()), Some("r-sub-1"));
38/// assert_eq!(events[4].event_type(), EventType::SubagentFinished);
39/// assert_eq!(events[5].subagent_run_id(), None);
40/// # Ok::<(), ag_ui::server::Error>(())
41/// ```
42///
43/// Nesting is automatic: a subagent opened through a handle gets the handle's
44/// id as its `parentSubagentRunId`. Two subagents cannot be open at once
45/// through handles — the second `subagent()` is a borrow-check error, as
46/// everything overlapping is here. For subagents that genuinely stream
47/// concurrently, tag events yourself and emit them interleaved; see the
48/// [module docs](crate::server::emit#subagents).
49///
50/// # Ending it
51///
52/// The terminator names the subagent it closes and is not itself attributed
53/// to it, so every method here restores the enclosing attribution *before*
54/// emitting. `Drop` cannot tell success from failure: on the error path you
55/// care about, call [`fail`](Self::fail) — or [`suspend`](Self::suspend) when
56/// the run is about to pause on an interrupt the subagent raised.
57#[derive(Debug)]
58pub struct SubagentHandle<'a, S> {
59    ctx: &'a mut RunContext<S>,
60    id: SubagentRunId,
61    name: String,
62    /// The attribution in force before this subagent, restored on close.
63    previous: Option<SubagentRunId>,
64    ended: bool,
65}
66
67impl<'a, S> SubagentHandle<'a, S> {
68    /// Emits `SUBAGENT_STARTED` and scopes the context to the new subagent.
69    ///
70    /// A `parent_subagent_run_id` the caller left absent is filled from the
71    /// enclosing scope, which is what makes nesting through `Deref` correct
72    /// without the agent naming the parent.
73    pub(crate) fn start(
74        ctx: &'a mut RunContext<S>,
75        mut started: SubagentStartedEvent,
76    ) -> Result<Self> {
77        if started.parent_subagent_run_id.is_none() {
78            started.parent_subagent_run_id = ctx.subagent_run_id().cloned();
79        }
80        let id = started.subagent_run_id.clone();
81        let name = started.name.clone();
82        ctx.emit(started.into())?;
83        let previous = ctx.set_attribution(Some(id.clone()));
84        Ok(Self {
85            ctx,
86            id,
87            name,
88            previous,
89            ended: false,
90        })
91    }
92
93    /// The id every event emitted through this handle carries.
94    pub fn id(&self) -> &SubagentRunId {
95        &self.id
96    }
97
98    /// The subagent's declared name, as announced.
99    pub fn name(&self) -> &str {
100        &self.name
101    }
102
103    /// Leaves the scope: the enclosing attribution is back in force, and the
104    /// terminator about to be emitted belongs to it.
105    fn leave(&mut self) {
106        self.ended = true;
107        let previous = self.previous.take();
108        self.ctx.set_attribution(previous);
109    }
110
111    /// Emits `SUBAGENT_FINISHED` with a success outcome and consumes the
112    /// handle.
113    ///
114    /// Only worth calling over letting the handle drop when you want to see
115    /// the error: `Drop` cannot report one.
116    pub fn finish(mut self) -> Result<()> {
117        self.leave();
118        self.ctx
119            .emit(Event::subagent_finished_success(self.id.clone()))
120    }
121
122    /// Emits `SUBAGENT_FINISHED` carrying a completion payload — the
123    /// subagent's counterpart of `RUN_FINISHED.result`.
124    pub fn finish_with(mut self, result: impl Into<Value>) -> Result<()> {
125        self.leave();
126        self.ctx.emit(
127            SubagentFinishedEvent::new(self.id.clone())
128                .with_result(result)
129                .with_outcome(SubagentOutcome::Success)
130                .into(),
131        )
132    }
133
134    /// Emits `SUBAGENT_FINISHED` with a suspended outcome: the subagent is
135    /// waiting on `interrupt_ids`, which the run is about to return in an
136    /// interrupt outcome.
137    ///
138    /// Build each interrupt with
139    /// [`Interrupt::with_subagent_run_id`](crate::Interrupt::with_subagent_run_id)
140    /// so a client can render it inside this subagent's group, and announce
141    /// the same id again on the resuming run to continue the invocation.
142    pub fn suspend(mut self, interrupt_ids: impl Into<Vec<String>>) -> Result<()> {
143        self.leave();
144        self.ctx.emit(Event::subagent_finished_suspended(
145            self.id.clone(),
146            interrupt_ids,
147        ))
148    }
149
150    /// Emits `SUBAGENT_ERROR` and consumes the handle.
151    pub fn fail(mut self, message: impl Into<String>) -> Result<()> {
152        self.leave();
153        self.ctx
154            .emit(Event::subagent_error(self.id.clone(), message))
155    }
156
157    /// Emits `SUBAGENT_ERROR` with a machine-readable code.
158    pub fn fail_with_code(
159        mut self,
160        message: impl Into<String>,
161        code: impl Into<String>,
162    ) -> Result<()> {
163        self.leave();
164        self.ctx.emit(
165            crate::SubagentErrorEvent::new(self.id.clone(), message)
166                .with_code(code)
167                .into(),
168        )
169    }
170}
171
172impl<S> Deref for SubagentHandle<'_, S> {
173    type Target = RunContext<S>;
174
175    fn deref(&self) -> &Self::Target {
176        self.ctx
177    }
178}
179
180impl<S> DerefMut for SubagentHandle<'_, S> {
181    fn deref_mut(&mut self) -> &mut Self::Target {
182        self.ctx
183    }
184}
185
186impl<S> Drop for SubagentHandle<'_, S> {
187    fn drop(&mut self) {
188        if !self.ended {
189            self.leave();
190            // Nowhere to report a failure to; a dead channel or a cancelled run
191            // makes the terminator moot anyway.
192            let _ = self
193                .ctx
194                .emit(Event::subagent_finished_success(self.id.clone()));
195        }
196    }
197}
ag_ui/server/emit/subagent.rs

ag_ui/server/emit/
subagent.rs
