ag_ui/server/emit/
message.rs

1//! Streaming one text message.
2
3use crate::{Event, MessageId, TextMessageRole};
4
5use crate::server::agent::AgentState;
6use crate::server::emit::EventSink;
7use crate::server::error::Result;
8use crate::server::state::RunState;
9
10/// One open text message.
11///
12/// Created by [`RunContext::assistant_message`](crate::server::RunContext::assistant_message).
13/// `TEXT_MESSAGE_START` has already gone out by the time you hold one; `Drop`
14/// emits `TEXT_MESSAGE_END` if [`end`](Self::end) was not called.
15///
16/// ```
17/// # use ag_ui::RunAgentInput;
18/// # use ag_ui::server::RunContext;
19/// # let (mut ctx, mut events) = RunContext::<()>::new(RunAgentInput::new("t", "r"))?;
20/// let mut message = ctx.assistant_message()?;
21/// for word in ["Hello", ", ", "world"] {
22///     message.delta(word)?;
23/// }
24/// message.end()?;
25/// assert_eq!(events.drain().len(), 5);
26/// # Ok::<(), ag_ui::server::Error>(())
27/// ```
28#[derive(Debug)]
29pub struct MessageHandle<'a, S> {
30    sink: &'a mut EventSink,
31    state: &'a mut RunState<S>,
32    id: MessageId,
33    ended: bool,
34}
35
36impl<'a, S> MessageHandle<'a, S> {
37    /// Emits `TEXT_MESSAGE_START` and takes the message.
38    ///
39    /// Returns `Err` without producing a handle when the start could not be
40    /// emitted, so a failed open never leaves a terminator to be emitted for a
41    /// message that does not exist.
42    pub(crate) fn start(
43        sink: &'a mut EventSink,
44        state: &'a mut RunState<S>,
45        id: MessageId,
46        role: TextMessageRole,
47    ) -> Result<Self> {
48        sink.emit(Event::text_message_start(id.clone(), role))?;
49        Ok(Self {
50            sink,
51            state,
52            id,
53            ended: false,
54        })
55    }
56
57    /// The id every event of this message carries.
58    pub fn id(&self) -> &MessageId {
59        &self.id
60    }
61
62    /// Appends text — `TEXT_MESSAGE_CONTENT`.
63    ///
64    /// No `.await`: see the [module docs](crate::server::emit) for why the emit path is
65    /// synchronous.
66    pub fn delta(&mut self, text: impl Into<String>) -> Result<()> {
67        self.sink
68            .emit(Event::text_message_content(self.id.clone(), text))
69    }
70
71    /// Emits an unrelated event without closing the message.
72    ///
73    /// For the unordered families — `STATE_*`, `ACTIVITY_*`, `CUSTOM`, `RAW` —
74    /// which may legally interleave with a message. Opening a second message
75    /// through here is a protocol violation the verifier will reject.
76    pub fn emit(&mut self, event: Event) -> Result<()> {
77        self.sink.emit(event)
78    }
79
80    /// Emits `TEXT_MESSAGE_END` and consumes the handle.
81    ///
82    /// Only worth calling over letting the handle drop when you want to see the
83    /// error: `Drop` cannot report one.
84    pub fn end(mut self) -> Result<()> {
85        self.ended = true;
86        self.sink.emit(Event::text_message_end(self.id.clone()))
87    }
88}
89
90/// The run's state, reachable while the message is open — an agent that
91/// narrates what it is doing changes both in the same breath. See
92/// [`ToolCallHandle`](crate::server::ToolCallHandle), where this matters most.
93impl<S: AgentState> MessageHandle<'_, S> {
94    /// The typed state, as of the last publish.
95    pub fn state(&self) -> &S {
96        self.state.get()
97    }
98
99    /// The typed state, mutably. Nothing is emitted until you call
100    /// [`publish_state`](Self::publish_state).
101    pub fn state_mut(&mut self) -> &mut S {
102        self.state.get_mut()
103    }
104
105    /// Publishes whatever [`state_mut`](Self::state_mut) left behind, as a
106    /// `STATE_SNAPSHOT` or a `STATE_DELTA` inside this message's brackets.
107    ///
108    /// A no-op when nothing changed since the last publish.
109    pub fn publish_state(&mut self) -> Result<()> {
110        self.state.publish(self.sink)
111    }
112}
113
114impl<S> Drop for MessageHandle<'_, S> {
115    fn drop(&mut self) {
116        if !self.ended {
117            // Nowhere to report a failure to; a dead channel or a cancelled run
118            // makes the terminator moot anyway.
119            let _ = self.sink.emit(Event::text_message_end(self.id.clone()));
120        }
121    }
122}
ag_ui/server/emit/message.rs

ag_ui/server/emit/
message.rs
