ag_ui/server/
context.rs

1//! What an agent is handed for one run.
2
3use std::future::Future;
4
5use crate::{
6    Context, Event, Message, MessageId, ResumeEntry, RunAgentInput, RunId, StepName, SubagentRunId,
7    SubagentStartedEvent, TextMessageRole, ThreadId, Tool, ToolCallId,
8};
9use futures_channel::mpsc;
10use futures_util::future::{Either, select};
11use serde_json::Value;
12
13use crate::server::agent::AgentState;
14use crate::server::cancel::{CancellationToken, Cancelled};
15use crate::server::emit::{
16    EventReceiver, EventSink, MessageHandle, ReasoningHandle, StepGuard, SubagentHandle,
17    ToolCallHandle,
18};
19use crate::server::error::{Error, Result};
20use crate::server::state::RunState;
21use crate::server::transform::TransformerChain;
22
23/// The request, the state, the event sink and the cancellation flag — one
24/// run's whole world.
25///
26/// An agent gets `&mut RunContext<S>` and emits through it. Every emitter takes
27/// `&mut self`, which is what makes two overlapping messages a borrow-check
28/// error rather than a protocol violation discovered by a confused frontend.
29///
30/// ```
31/// # use ag_ui::RunAgentInput;
32/// # use ag_ui::server::RunContext;
33/// # let (mut ctx, _events) = RunContext::<()>::new(RunAgentInput::new("thread-1", "run-1"))?;
34/// assert_eq!(ctx.thread_id().as_str(), "thread-1");
35/// assert!(ctx.messages().is_empty());
36///
37/// let mut message = ctx.assistant_message()?;
38/// message.delta("Hello")?;
39/// message.end()?;
40/// # Ok::<(), ag_ui::server::Error>(())
41/// ```
42#[derive(Debug)]
43pub struct RunContext<S> {
44    input: RunAgentInput,
45    state: RunState<S>,
46    sink: EventSink,
47    next_message: u64,
48    next_tool_call: u64,
49    next_subagent: u64,
50}
51
52impl<S: AgentState> RunContext<S> {
53    /// Builds a context and the receiving half of its event stream.
54    ///
55    /// This is the harness for unit-testing an [`Agent`](crate::server::Agent) without
56    /// the run driver: call the agent's body, then assert on
57    /// [`EventReceiver::drain`]. Nothing emits `RUN_STARTED` here — that is the
58    /// driver's job, and skipping it lets a test exercise one method in
59    /// isolation.
60    pub fn new(input: RunAgentInput) -> Result<(Self, EventReceiver)> {
61        let (tx, rx) = mpsc::unbounded();
62        let sink = EventSink::new(tx, TransformerChain::new(), CancellationToken::new());
63        let state = decode_state(&input.state)?;
64        Ok((Self::from_parts(input, state, sink), EventReceiver::new(rx)))
65    }
66
67    /// Assembles a context from an already-decoded state.
68    ///
69    /// The run driver decodes first so that a state that does not fit `S` is
70    /// reported through the sink it still owns, as a `RUN_ERROR`.
71    pub(crate) fn from_parts(input: RunAgentInput, state: S, sink: EventSink) -> Self {
72        Self {
73            input,
74            state: RunState::new(state),
75            sink,
76            next_message: 0,
77            next_tool_call: 0,
78            next_subagent: 0,
79        }
80    }
81
82    /// The typed state, as of the last publish.
83    pub fn state(&self) -> &S {
84        self.state.get()
85    }
86
87    /// The typed state, mutably. Nothing is emitted until you call
88    /// [`publish_state`](Self::publish_state).
89    pub fn state_mut(&mut self) -> &mut S {
90        self.state.get_mut()
91    }
92
93    /// Replaces the state and publishes the change.
94    ///
95    /// The first publish of a run is a `STATE_SNAPSHOT`; later ones are a
96    /// `STATE_DELTA` unless the patch would be no smaller than the snapshot.
97    /// See [`StateManager`](crate::server::StateManager).
98    pub fn set_state(&mut self, state: &S) -> Result<()> {
99        self.state.replace(&mut self.sink, state)
100    }
101
102    /// Mutates the state in place and publishes the change.
103    ///
104    /// ```
105    /// # use ag_ui::RunAgentInput;
106    /// # use ag_ui::server::RunContext;
107    /// # use serde::{Deserialize, Serialize};
108    /// #[derive(Default, Serialize, Deserialize)]
109    /// struct Draft { revision: u32 }
110    ///
111    /// # let (mut ctx, _events) = RunContext::<Draft>::new(RunAgentInput::new("t", "r"))?;
112    /// ctx.update_state(|draft| draft.revision += 1)?;
113    /// assert_eq!(ctx.state().revision, 1);
114    /// # Ok::<(), ag_ui::server::Error>(())
115    /// ```
116    pub fn update_state(&mut self, update: impl FnOnce(&mut S)) -> Result<()> {
117        update(self.state.get_mut());
118        self.publish_state()
119    }
120
121    /// Publishes whatever [`state_mut`](Self::state_mut) left behind.
122    ///
123    /// A no-op when nothing changed since the last publish.
124    pub fn publish_state(&mut self) -> Result<()> {
125        self.state.publish(&mut self.sink)
126    }
127}
128
129impl<S> RunContext<S> {
130    /// The whole request, for anything the accessors do not cover.
131    pub fn input(&self) -> &RunAgentInput {
132        &self.input
133    }
134
135    /// The conversation this run belongs to.
136    pub fn thread_id(&self) -> &ThreadId {
137        &self.input.thread_id
138    }
139
140    /// This run's id.
141    pub fn run_id(&self) -> &RunId {
142        &self.input.run_id
143    }
144
145    /// The run that spawned this one, for nested agents.
146    pub fn parent_run_id(&self) -> Option<&RunId> {
147        self.input.parent_run_id.as_ref()
148    }
149
150    /// Conversation history, oldest first.
151    pub fn messages(&self) -> &[Message] {
152        &self.input.messages
153    }
154
155    /// What the user said last, as text.
156    ///
157    /// The turn an agent is almost always answering. Non-text parts of a
158    /// multimodal message are dropped — see
159    /// [`UserContent::to_text`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/message/enum.UserContent.html#method.to_text); reach into
160    /// [`RunContext::messages`] directly if the images matter.
161    ///
162    /// `None` when the history holds no user message at all, which is distinct
163    /// from a user who sent an empty one.
164    ///
165    /// ```
166    /// # use ag_ui::{RunAgentInput, Message};
167    /// # use ag_ui::server::RunContext;
168    /// let mut input = RunAgentInput::new("thread-1", "run-1");
169    /// input.messages = vec![Message::user("msg-1", "add milk")];
170    /// let (ctx, _events) = RunContext::<()>::new(input)?;
171    ///
172    /// assert_eq!(ctx.last_user_text().as_deref(), Some("add milk"));
173    /// # Ok::<(), ag_ui::server::Error>(())
174    /// ```
175    pub fn last_user_text(&self) -> Option<String> {
176        self.input
177            .messages
178            .iter()
179            .rev()
180            .find_map(|message| match message {
181                Message::User(user) => Some(user.content.to_text()),
182                _ => None,
183            })
184    }
185
186    /// Tools the client is offering for this run.
187    pub fn tools(&self) -> &[Tool] {
188        &self.input.tools
189    }
190
191    /// One offered tool by name.
192    pub fn tool(&self, name: &str) -> Option<&Tool> {
193        self.input.tools.iter().find(|tool| tool.name == name)
194    }
195
196    /// Ambient context entries.
197    pub fn context(&self) -> &[Context] {
198        &self.input.context
199    }
200
201    /// Arbitrary passthrough properties, opaque to the protocol.
202    pub fn forwarded_props(&self) -> &Value {
203        &self.input.forwarded_props
204    }
205
206    /// Answers to the interrupts a previous run paused on.
207    ///
208    /// Empty unless this request resumes a paused run.
209    pub fn resume(&self) -> &[ResumeEntry] {
210        self.input.resume.as_deref().unwrap_or_default()
211    }
212
213    /// The answer to one interrupt, by its [`Interrupt::id`].
214    ///
215    /// [`Interrupt::id`]: https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/outcome/struct.Interrupt.html#structfield.id
216    pub fn resume_for(&self, interrupt_id: &str) -> Option<&ResumeEntry> {
217        self.resume()
218            .iter()
219            .find(|entry| entry.interrupt_id == interrupt_id)
220    }
221
222    /// Whether this request resumes a paused run.
223    pub fn is_resume(&self) -> bool {
224        self.input.is_resume()
225    }
226
227    /// Whether the run has been cancelled.
228    pub fn is_cancelled(&self) -> bool {
229        self.sink.cancel_token().is_cancelled()
230    }
231
232    /// [`Error::Cancelled`] once the run is cancelled, for use with `?`.
233    pub fn check_cancelled(&self) -> Result<()> {
234        if self.is_cancelled() {
235            return Err(Error::Cancelled);
236        }
237        Ok(())
238    }
239
240    /// A handle a transport can trip on client disconnect.
241    pub fn cancel_token(&self) -> CancellationToken {
242        self.sink.cancel_token().clone()
243    }
244
245    /// Resolves once the run is cancelled.
246    pub fn cancelled(&self) -> Cancelled {
247        self.sink.cancel_token().cancelled()
248    }
249
250    /// Races `future` against cancellation, returning `None` if cancellation
251    /// won.
252    ///
253    /// The way to make a long model call interruptible:
254    ///
255    /// ```
256    /// # use ag_ui::RunAgentInput;
257    /// # use ag_ui::server::{Error, RunContext};
258    /// # let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
259    /// # rt.block_on(async {
260    /// # let (ctx, _events) = RunContext::<()>::new(RunAgentInput::new("t", "r"))?;
261    /// let answer = ctx
262    ///     .until_cancelled(async { "the model's reply" })
263    ///     .await
264    ///     .ok_or(Error::Cancelled)?;
265    /// assert_eq!(answer, "the model's reply");
266    /// # Ok::<(), Error>(())
267    /// # })?;
268    /// # Ok::<(), Error>(())
269    /// ```
270    /// Deliberately not an `async fn`: that would capture `&self` in the
271    /// returned future, and a future holding a borrow of the run context is
272    /// only `Send` if the context is `Sync` — which it is not, since a
273    /// transformer only has to be `Send`.
274    pub fn until_cancelled<F: Future>(&self, future: F) -> impl Future<Output = Option<F::Output>> {
275        let cancelled = self.cancelled();
276        async move {
277            futures_util::pin_mut!(future, cancelled);
278            match select(future, cancelled).await {
279                Either::Left((output, _)) => Some(output),
280                Either::Right(((), _)) => None,
281            }
282        }
283    }
284
285    /// Emits an event as-is — the escape hatch under the typed emitters.
286    ///
287    /// Everything the handles emit goes through here, so a raw event is
288    /// transformed and verified like any other.
289    pub fn emit(&mut self, event: Event) -> Result<()> {
290        self.sink.emit(event)
291    }
292
293    /// A fresh message id, unique within the run.
294    ///
295    /// Derived from the run id and a counter rather than a UUID: the protocol
296    /// asks for opaque strings, this crate takes no `uuid` dependency, and a
297    /// deterministic id makes a recorded stream diffable. Pass your own id to
298    /// [`message_with_id`](Self::message_with_id) when you need one.
299    pub fn new_message_id(&mut self) -> MessageId {
300        self.next_message += 1;
301        MessageId::new(format!("{}-msg-{}", self.id_prefix(), self.next_message))
302    }
303
304    /// A fresh tool call id, unique within the run.
305    pub fn new_tool_call_id(&mut self) -> ToolCallId {
306        self.next_tool_call += 1;
307        ToolCallId::new(format!("{}-call-{}", self.id_prefix(), self.next_tool_call))
308    }
309
310    /// A fresh subagent invocation id, unique within the run.
311    ///
312    /// Derived like the others, from the run id and a counter — so, like
313    /// message ids, it is unique across runs only while run ids are. A
314    /// resuming run that continues a *suspended* subagent should reuse the
315    /// suspended id instead — see [`subagent_with`](Self::subagent_with).
316    pub fn new_subagent_run_id(&mut self) -> SubagentRunId {
317        self.next_subagent += 1;
318        SubagentRunId::new(format!("{}-sub-{}", self.id_prefix(), self.next_subagent))
319    }
320
321    /// The subagent everything emitted right now is attributed to — `None`
322    /// outside any [`subagent`](Self::subagent) scope.
323    pub fn subagent_run_id(&self) -> Option<&SubagentRunId> {
324        self.sink.attribution()
325    }
326
327    /// Replaces the attribution scope and returns the previous one. Only the
328    /// subagent handle calls this, on the way in and on the way out.
329    pub(crate) fn set_attribution(
330        &mut self,
331        attribution: Option<SubagentRunId>,
332    ) -> Option<SubagentRunId> {
333        self.sink.set_attribution(attribution)
334    }
335
336    /// Announces a subagent under a fresh id and scopes everything emitted
337    /// through the returned handle to it — `SUBAGENT_STARTED` now,
338    /// `SUBAGENT_FINISHED` when the handle drops.
339    ///
340    /// `name` is the subagent's reusable type or name, for display; the id is
341    /// this invocation's alone. See [`SubagentHandle`].
342    pub fn subagent(&mut self, name: impl Into<String>) -> Result<SubagentHandle<'_, S>> {
343        let id = self.new_subagent_run_id();
344        self.subagent_with(SubagentStartedEvent::new(id, name))
345    }
346
347    /// The same, from an announcement you built — for a description, an
348    /// explicit id, or the agents-as-tools links.
349    ///
350    /// A `parent_subagent_run_id` left absent is filled from the enclosing
351    /// scope, so nesting needs no help. An explicit id is how a resuming run
352    /// continues a subagent that suspended: announce the id the suspended
353    /// invocation had, and a client transitions its group from waiting back
354    /// to running rather than drawing a second one.
355    ///
356    /// ```
357    /// # use ag_ui::{Event, EventType, RunAgentInput, SubagentStartedEvent};
358    /// # use ag_ui::server::RunContext;
359    /// # let (mut ctx, mut events) = RunContext::<()>::new(RunAgentInput::new("t", "r"))?;
360    /// let mut call = ctx.tool_call("task")?;
361    /// call.args(r#"{"brief":"find sources"}"#)?;
362    /// let (call_id, result_id) = (call.id().clone(), call.result_message_id().clone());
363    /// call.end()?;                                   // the client sees the call close…
364    ///
365    /// let announce = SubagentStartedEvent::new("researcher-7", "researcher")
366    ///     .with_parent_tool_call(call_id.clone());
367    /// let mut researcher = ctx.subagent_with(announce)?;
368    /// researcher.say("Three sources found.")?;       // …then the subagent it spawned…
369    /// researcher.finish()?;
370    ///
371    /// ctx.emit(Event::tool_call_result(result_id, call_id, "3 sources"))?;  // …then its result
372    /// let types: Vec<_> = events.drain().iter().map(Event::event_type).collect();
373    /// assert_eq!(types[3], EventType::SubagentStarted);
374    /// assert_eq!(types[7], EventType::SubagentFinished);
375    /// assert_eq!(types[8], EventType::ToolCallResult);
376    /// # Ok::<(), ag_ui::server::Error>(())
377    /// ```
378    pub fn subagent_with(
379        &mut self,
380        started: SubagentStartedEvent,
381    ) -> Result<SubagentHandle<'_, S>> {
382        SubagentHandle::start(self, started)
383    }
384
385    fn id_prefix(&self) -> &str {
386        if self.input.run_id.is_empty() {
387            "run"
388        } else {
389            self.input.run_id.as_str()
390        }
391    }
392
393    /// Opens an assistant message under a fresh id — `TEXT_MESSAGE_START`.
394    pub fn assistant_message(&mut self) -> Result<MessageHandle<'_, S>> {
395        self.message(TextMessageRole::Assistant)
396    }
397
398    /// Opens a message with the given role under a fresh id.
399    pub fn message(&mut self, role: TextMessageRole) -> Result<MessageHandle<'_, S>> {
400        let id = self.new_message_id();
401        self.message_with_id(id, role)
402    }
403
404    /// Opens a message under an id you choose.
405    pub fn message_with_id(
406        &mut self,
407        id: impl Into<MessageId>,
408        role: TextMessageRole,
409    ) -> Result<MessageHandle<'_, S>> {
410        // Two disjoint field borrows, not a borrow of the context: the handle
411        // reaches the state without being able to open a second block.
412        MessageHandle::start(&mut self.sink, &mut self.state, id.into(), role)
413    }
414
415    /// Emits a whole assistant message — start, content, end — and returns its
416    /// id.
417    pub fn say(&mut self, text: impl Into<String>) -> Result<MessageId> {
418        let mut message = self.message(TextMessageRole::Assistant)?;
419        message.delta(text)?;
420        let id = message.id().clone();
421        message.end()?;
422        Ok(id)
423    }
424
425    /// Opens a reasoning block under a fresh id — `REASONING_START`.
426    pub fn reasoning(&mut self) -> Result<ReasoningHandle<'_, S>> {
427        let id = self.new_message_id();
428        self.reasoning_with_id(id)
429    }
430
431    /// Opens a reasoning block under an id you choose.
432    pub fn reasoning_with_id(
433        &mut self,
434        id: impl Into<MessageId>,
435    ) -> Result<ReasoningHandle<'_, S>> {
436        ReasoningHandle::start(&mut self.sink, &mut self.state, id.into())
437    }
438
439    /// Emits a whole reasoning block in one call and returns its id.
440    pub fn think(&mut self, text: impl Into<String>) -> Result<MessageId> {
441        let mut reasoning = self.reasoning()?;
442        reasoning.delta(text)?;
443        let id = reasoning.id().clone();
444        reasoning.end()?;
445        Ok(id)
446    }
447
448    /// Opens a call to `name` under a fresh id — `TOOL_CALL_START`.
449    pub fn tool_call(&mut self, name: &str) -> Result<ToolCallHandle<'_, S>> {
450        let id = self.new_tool_call_id();
451        self.tool_call_with_id(id, name)
452    }
453
454    /// Opens a call to `name` under an id you choose.
455    pub fn tool_call_with_id(
456        &mut self,
457        id: impl Into<ToolCallId>,
458        name: &str,
459    ) -> Result<ToolCallHandle<'_, S>> {
460        let id = id.into();
461        let result_message_id = self.new_message_id();
462        ToolCallHandle::start(
463            &mut self.sink,
464            &mut self.state,
465            id,
466            name,
467            None,
468            result_message_id,
469        )
470    }
471
472    /// Opens a named step — `STEP_STARTED`.
473    ///
474    /// The returned guard dereferences to this context, and emits
475    /// `STEP_FINISHED` when it drops.
476    pub fn step(&mut self, name: impl Into<StepName>) -> Result<StepGuard<'_, S>> {
477        StepGuard::start(self, name.into())
478    }
479
480    /// Whether a terminal event has already gone out.
481    pub(crate) fn is_terminated(&self) -> bool {
482        self.sink.is_terminated()
483    }
484
485    /// Recovers the sink so the run driver can emit the terminal event through
486    /// the same transformers and the same verifier the agent used.
487    pub(crate) fn into_sink(self) -> EventSink {
488        self.sink
489    }
490}
491
492/// Reads `RunAgentInput::state` as `S`.
493///
494/// An absent state — JSON `null`, or the empty object clients send for "no
495/// state yet" — becomes `S::default()` rather than a deserialization error, so
496/// a stateless agent (`State = ()`) works against every client.
497pub(crate) fn decode_state<S: AgentState>(value: &Value) -> Result<S> {
498    let empty = value.is_null() || value.as_object().is_some_and(serde_json::Map::is_empty);
499    if empty {
500        return Ok(S::default());
501    }
502    Ok(serde_json::from_value(value.clone())?)
503}
504
505#[cfg(test)]
506mod tests {
507    use super::*;
508    use serde::{Deserialize, Serialize};
509    use serde_json::json;
510
511    #[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
512    struct Counter {
513        clicks: u32,
514    }
515
516    fn context<S: AgentState>(input: RunAgentInput) -> (RunContext<S>, EventReceiver) {
517        RunContext::new(input).expect("state should decode")
518    }
519
520    #[test]
521    fn ids_are_derived_from_the_run_id() {
522        let (mut ctx, _events) = context::<()>(RunAgentInput::new("t", "run-7"));
523        assert_eq!(ctx.new_message_id().as_str(), "run-7-msg-1");
524        assert_eq!(ctx.new_message_id().as_str(), "run-7-msg-2");
525        assert_eq!(ctx.new_tool_call_id().as_str(), "run-7-call-1");
526    }
527
528    #[test]
529    fn the_last_user_turn_is_the_one_returned_and_absence_is_not_emptiness() {
530        let (ctx, _events) = context::<()>(RunAgentInput::new("t", "r"));
531        assert_eq!(ctx.last_user_text(), None, "no user message is not \"\"");
532
533        let mut input = RunAgentInput::new("t", "r");
534        input.messages = vec![
535            Message::user("m-1", "add milk"),
536            Message::assistant("m-2", "milk it is"),
537            Message::user("m-3", "and bread"),
538        ];
539        let (ctx, _events) = context::<()>(input);
540        assert_eq!(ctx.last_user_text().as_deref(), Some("and bread"));
541    }
542
543    #[test]
544    fn an_empty_state_object_decodes_to_the_default() {
545        let mut input = RunAgentInput::new("t", "r");
546        input.state = json!({});
547        let (ctx, _events) = context::<Counter>(input);
548        assert_eq!(ctx.state(), &Counter::default());
549    }
550
551    #[test]
552    fn typed_state_comes_from_the_input() {
553        let mut input = RunAgentInput::new("t", "r");
554        input.state = json!({"clicks": 3});
555        let (ctx, _events) = context::<Counter>(input);
556        assert_eq!(ctx.state().clicks, 3);
557    }
558
559    #[test]
560    fn a_state_that_does_not_fit_is_an_error() {
561        let mut input = RunAgentInput::new("t", "r");
562        input.state = json!({"clicks": "three"});
563        let error = RunContext::<Counter>::new(input).expect_err("should not decode");
564        assert!(matches!(error, Error::Json(_)), "{error}");
565    }
566
567    #[test]
568    fn resume_entries_are_addressable_by_interrupt_id() {
569        let mut input = RunAgentInput::new("t", "r");
570        input.resume = Some(vec![ResumeEntry::resolved("i-1", json!(true))]);
571        let (ctx, _events) = context::<()>(input);
572        assert!(ctx.is_resume());
573        assert_eq!(ctx.resume().len(), 1);
574        assert!(ctx.resume_for("i-1").is_some());
575        assert!(ctx.resume_for("i-2").is_none());
576    }
577
578    #[test]
579    fn a_cancelled_run_fails_every_emit() {
580        let (mut ctx, _events) = context::<()>(RunAgentInput::new("t", "r"));
581        ctx.cancel_token().cancel();
582        assert!(ctx.is_cancelled());
583        let error = ctx.say("too late").expect_err("emit should fail");
584        assert!(error.is_cancelled(), "{error}");
585    }
586
587    #[test]
588    fn a_dropped_receiver_disconnects_the_run() {
589        let (mut ctx, events) = context::<()>(RunAgentInput::new("t", "r"));
590        drop(events);
591        let error = ctx.say("nobody home").expect_err("emit should fail");
592        assert!(error.is_disconnected(), "{error}");
593    }
594}
ag_ui/server/context.rs

ag_ui/server/
context.rs
