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}