ag_ui/server/emit/
step.rs

1//! Bracketing a named step.
2
3use std::ops::{Deref, DerefMut};
4
5use crate::{Event, StepName};
6
7use crate::server::context::RunContext;
8use crate::server::error::Result;
9
10/// One open step.
11///
12/// Created by [`RunContext::step`](crate::server::RunContext::step). `STEP_STARTED` has
13/// already gone out; `Drop` emits `STEP_FINISHED`, including on the early
14/// return that a `?` produces.
15///
16/// A step is a *scope*, not a stream, so unlike the message and tool-call
17/// handles this one dereferences to the run context — everything nests inside
18/// it, steps included:
19///
20/// ```
21/// # use ag_ui::RunAgentInput;
22/// # use ag_ui::server::RunContext;
23/// # let (mut ctx, mut events) = RunContext::<()>::new(RunAgentInput::new("t", "r"))?;
24/// {
25///     let mut step = ctx.step("research")?;
26///     step.say("looking it up")?;   // through Deref
27/// }                                 // STEP_FINISHED here
28/// assert_eq!(events.drain().len(), 5);
29/// # Ok::<(), ag_ui::server::Error>(())
30/// ```
31#[derive(Debug)]
32pub struct StepGuard<'a, S> {
33    ctx: &'a mut RunContext<S>,
34    name: StepName,
35    ended: bool,
36}
37
38impl<'a, S> StepGuard<'a, S> {
39    /// Emits `STEP_STARTED` and takes the step.
40    pub(crate) fn start(ctx: &'a mut RunContext<S>, name: StepName) -> Result<Self> {
41        ctx.emit(Event::step_started(name.clone()))?;
42        Ok(Self {
43            ctx,
44            name,
45            ended: false,
46        })
47    }
48
49    /// The step's name.
50    pub fn name(&self) -> &StepName {
51        &self.name
52    }
53
54    /// Emits `STEP_FINISHED` and consumes the guard.
55    ///
56    /// Only worth calling over letting the guard drop when you want to see the
57    /// error.
58    pub fn finish(mut self) -> Result<()> {
59        self.ended = true;
60        self.ctx.emit(Event::step_finished(self.name.clone()))
61    }
62}
63
64impl<S> Deref for StepGuard<'_, S> {
65    type Target = RunContext<S>;
66
67    fn deref(&self) -> &Self::Target {
68        self.ctx
69    }
70}
71
72impl<S> DerefMut for StepGuard<'_, S> {
73    fn deref_mut(&mut self) -> &mut Self::Target {
74        self.ctx
75    }
76}
77
78impl<S> Drop for StepGuard<'_, S> {
79    fn drop(&mut self) {
80        if !self.ended {
81            let _ = self.ctx.emit(Event::step_finished(self.name.clone()));
82        }
83    }
84}
ag_ui/server/emit/step.rs

ag_ui/server/emit/
step.rs
