Skip to main content

ag_ui/server/emit/
reasoning.rs

1//! Streaming one reasoning block.
2
3use crate::{Event, MessageId, ReasoningEncryptedValueSubtype};
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 reasoning block.
11///
12/// Created by [`RunContext::reasoning`](crate::server::RunContext::reasoning). The
13/// `REASONING_*` family nests a message inside a block, so the handle brackets
14/// four events rather than two:
15///
16/// ```text
17/// REASONING_START          ← on creation
18/// REASONING_MESSAGE_START  ← on creation
19/// REASONING_MESSAGE_CONTENT × n
20/// REASONING_MESSAGE_END    ← on end() or Drop
21/// REASONING_END            ← on end() or Drop
22/// ```
23///
24/// Use it for model reasoning you want the client to render. Reasoning a
25/// provider returns only as an opaque blob goes through
26/// [`encrypted_value`](Self::encrypted_value) instead.
27#[derive(Debug)]
28pub struct ReasoningHandle<'a, S> {
29    sink: &'a mut EventSink,
30    state: &'a mut RunState<S>,
31    id: MessageId,
32    ended: bool,
33}
34
35impl<'a, S> ReasoningHandle<'a, S> {
36    /// Emits `REASONING_START` and `REASONING_MESSAGE_START`.
37    pub(crate) fn start(
38        sink: &'a mut EventSink,
39        state: &'a mut RunState<S>,
40        id: MessageId,
41    ) -> Result<Self> {
42        sink.emit(Event::reasoning_start(id.clone()))?;
43        // The block is open now, so from here on a failure must still leave a
44        // handle behind to close it.
45        let handle = Self {
46            sink,
47            state,
48            id,
49            ended: false,
50        };
51        handle
52            .sink
53            .emit(Event::reasoning_message_start(handle.id.clone()))?;
54        Ok(handle)
55    }
56
57    /// The id every event of this block carries.
58    pub fn id(&self) -> &MessageId {
59        &self.id
60    }
61
62    /// Appends reasoning text — `REASONING_MESSAGE_CONTENT`.
63    pub fn delta(&mut self, text: impl Into<String>) -> Result<()> {
64        self.sink
65            .emit(Event::reasoning_message_content(self.id.clone(), text))
66    }
67
68    /// Attaches the provider's opaque reasoning signature —
69    /// `REASONING_ENCRYPTED_VALUE`.
70    ///
71    /// Under zero-data-retention the provider returns no readable reasoning,
72    /// only a blob that must be replayed on the next request for the model to
73    /// stay coherent.
74    pub fn encrypted_value(&mut self, value: impl Into<String>) -> Result<()> {
75        self.sink.emit(Event::reasoning_encrypted_value(
76            ReasoningEncryptedValueSubtype::Message,
77            self.id.clone(),
78            value,
79        ))
80    }
81
82    /// Emits an unrelated event without closing the block. See
83    /// [`MessageHandle::emit`](crate::server::MessageHandle::emit).
84    pub fn emit(&mut self, event: Event) -> Result<()> {
85        self.sink.emit(event)
86    }
87
88    /// Emits `REASONING_MESSAGE_END` then `REASONING_END`, and consumes the
89    /// handle.
90    pub fn end(mut self) -> Result<()> {
91        self.ended = true;
92        self.close()
93    }
94
95    /// Closes both halves. The block terminator is attempted even when the
96    /// message terminator failed, so a half-open block cannot outlive the
97    /// handle.
98    fn close(&mut self) -> Result<()> {
99        let message = self
100            .sink
101            .emit(Event::reasoning_message_end(self.id.clone()));
102        let block = self.sink.emit(Event::reasoning_end(self.id.clone()));
103        message.and(block)
104    }
105}
106
107/// The run's state, reachable while the block is open. See
108/// [`ToolCallHandle`](crate::server::ToolCallHandle), where this matters most.
109impl<S: AgentState> ReasoningHandle<'_, S> {
110    /// The typed state, as of the last publish.
111    pub fn state(&self) -> &S {
112        self.state.get()
113    }
114
115    /// The typed state, mutably. Nothing is emitted until you call
116    /// [`publish_state`](Self::publish_state).
117    pub fn state_mut(&mut self) -> &mut S {
118        self.state.get_mut()
119    }
120
121    /// Publishes whatever [`state_mut`](Self::state_mut) left behind, as a
122    /// `STATE_SNAPSHOT` or a `STATE_DELTA` inside this block's brackets.
123    ///
124    /// A no-op when nothing changed since the last publish.
125    pub fn publish_state(&mut self) -> Result<()> {
126        self.state.publish(self.sink)
127    }
128}
129
130impl<S> Drop for ReasoningHandle<'_, S> {
131    fn drop(&mut self) {
132        if !self.ended {
133            let _ = self.close();
134        }
135    }
136}