board_watch/
fake.rs

1//! The awkward agent: a backend that produces the shapes real providers do.
2//!
3//! This is not the application. It is the fixture the application is aimed at,
4//! and it lives here because a client dogfood is only worth as much as the
5//! streams it is pointed at. `task-board`'s agent is well-behaved — bracketed
6//! messages, whole tool arguments, one call at a time — so it proves the happy
7//! path and nothing else. Real streams are worse:
8//!
9//! - text arrives as `TEXT_MESSAGE_CHUNK` with the id **only on the first one**;
10//! - tool arguments are split at arbitrary byte offsets, including between a
11//!   backslash and the character it escapes;
12//! - a model calls two tools at once and their events interleave;
13//! - a producer in another language sends something the protocol forbids.
14//!
15//! Every scenario is chosen by the first word of the user's message, so a
16//! transcript names what it exercised. Nothing here is timed and nothing is
17//! random: the same message produces the same bytes every run.
18//!
19//! # Why so much of it is raw `emit`
20//!
21//! `ag-ui-server`'s typestate handles bracket what they open, which is exactly
22//! what a chunk event is defined not to do — there is no `ctx.text_chunk()`,
23//! and two overlapping [`ToolCallHandle`](ag_ui::server::ToolCallHandle)s are a
24//! borrow-check error by design. Producing provider-shaped output therefore
25//! means dropping to [`RunContext::emit`], which is the documented escape
26//! hatch. See the report: this is a finding, not a complaint about the design.
27
28use ag_ui::axum::RouterExt;
29use ag_ui::server::{Agent, CancellationToken, Error, Result, RunContext};
30use ag_ui::{Event, Interrupt, MessageId, ResumeStatus, RunOutcome, TextMessageRole, ToolCallId};
31use ag_ui_a2ui::constants::RENDER_A2UI_TOOL_NAME;
32use ag_ui_a2ui::message::Component;
33use ag_ui_a2ui::toolkit::envelope::wrap_as_operations_envelope;
34use ag_ui_a2ui::toolkit::ops::{Intent, SurfaceSpec, assemble_ops};
35use axum::Router;
36use axum::response::{IntoResponse, Response};
37use axum::routing::{get, post};
38use serde_json::json;
39use tokio::sync::mpsc::UnboundedSender;
40
41use crate::board::{Board, Task};
42
43/// Where the awkward agent is mounted.
44pub const ROUTE: &str = "/agent";
45
46/// Where the hand-framed, deliberately illegal streams are mounted.
47///
48/// `/raw/{scenario}` — see [`raw_script`].
49pub const RAW_ROUTE: &str = "/raw/{scenario}";
50
51/// The interrupt ids the `approve` scenario pauses on.
52pub const BUDGET: &str = "approve-budget";
53/// The second of them. Two, because answering one per request never terminates.
54pub const DATE: &str = "confirm-date";
55
56/// An agent that answers in the shapes a provider adapter produces.
57#[derive(Debug, Default)]
58pub struct Awkward {
59    /// Set by tests: reports, as the run's future is dropped, whether the run
60    /// had been cancelled. A client that stops reading has to reach the agent,
61    /// and nothing observable from the client side proves that it did.
62    exits: Option<UnboundedSender<bool>>,
63}
64
65impl Awkward {
66    /// The agent as the CLI serves it.
67    pub fn new() -> Self {
68        Self::default()
69    }
70
71    /// The agent, reporting every run's cancellation state as it exits.
72    pub fn reporting(exits: UnboundedSender<bool>) -> Self {
73        Self { exits: Some(exits) }
74    }
75}
76
77/// Reports whether the run was cancelled on *every* way out, including the one
78/// where the agent's future is simply dropped mid-await.
79struct ExitGuard {
80    token: CancellationToken,
81    report: UnboundedSender<bool>,
82}
83
84impl Drop for ExitGuard {
85    fn drop(&mut self) {
86        let _ = self.report.send(self.token.is_cancelled());
87    }
88}
89
90impl Agent for Awkward {
91    type State = Board;
92
93    async fn run(&self, ctx: &mut RunContext<Board>) -> Result<RunOutcome> {
94        let said = ctx.last_user_text().unwrap_or_default();
95        let scenario = said.split_whitespace().next().unwrap_or("").to_lowercase();
96
97        let _guard = self.exits.clone().map(|report| ExitGuard {
98            token: ctx.cancel_token(),
99            report,
100        });
101
102        match scenario.as_str() {
103            "chunks" => chunked_text(ctx),
104            "call" => chunked_call(ctx),
105            "parallel" => parallel_calls(ctx),
106            "mixed" => mixed_stream(ctx),
107            "busy" => return busy(ctx),
108            "approve" => return Ok(approvals(ctx)),
109            "slow" => return slow(ctx).await,
110            "fail" => Err(Error::agent("the model refused, and said so at length")),
111            _ => board_turn(ctx),
112        }?;
113
114        Ok(RunOutcome::Success)
115    }
116}
117
118/// Text as `TEXT_MESSAGE_CHUNK`, the id on the first chunk only.
119///
120/// The last three fragments split a ZWJ emoji sequence between its parts: every
121/// fragment is valid UTF-8 on its own — a `String` cannot be otherwise — but
122/// the *grapheme* is only whole once they are rejoined.
123fn chunked_text(ctx: &mut RunContext<Board>) -> Result<()> {
124    let id = ctx.new_message_id();
125    let fragments = [
126        "Chunked text arrives in frag",
127        "ments, and the client rejoins ",
128        "them — emoji included: 👩",
129        "\u{200d}",
130        "💻.",
131    ];
132
133    for (index, fragment) in fragments.iter().enumerate() {
134        // The id rides on the first chunk and nothing else. A client that does
135        // not remember it drops everything after this line.
136        let carried = (index == 0).then(|| id.clone());
137        ctx.emit(Event::text_message_chunk(
138            carried,
139            Some((*fragment).to_owned()),
140        ))?;
141    }
142    Ok(())
143}
144
145/// One tool call as `TOOL_CALL_CHUNK`, arguments split at hostile offsets.
146///
147/// The split after `line\` is the case every provider adapter gets wrong once:
148/// the backslash and the `n` it escapes arrive in different events, so anything
149/// that parses a fragment on its own sees invalid JSON.
150fn chunked_call(ctx: &mut RunContext<Board>) -> Result<()> {
151    let id = ctx.new_tool_call_id();
152    let fragments = [
153        r#"{"no"#,
154        r#"te":"line\"#,
155        r#"nbreak","ti"#,
156        r#"tle":"ship "#,
157        r#"the SDK","depth":3}"#,
158    ];
159
160    for (index, fragment) in fragments.iter().enumerate() {
161        let first = index == 0;
162        ctx.emit(Event::tool_call_chunk(
163            first.then(|| id.clone()),
164            first.then(|| "add_task".to_owned()),
165            Some((*fragment).to_owned()),
166        ))?;
167    }
168
169    // A result still needs an explicit event; only the call itself chunks.
170    let message_id = ctx.new_message_id();
171    ctx.emit(Event::tool_call_result(
172        message_id,
173        id,
174        r#"{"id":1,"title":"ship the SDK"}"#,
175    ))
176}
177
178/// Two calls in flight at once, their events interleaved by id.
179///
180/// Legal on the wire and legal for the applier, and *unwritable* with the
181/// typestate handles: two open [`ToolCallHandle`](ag_ui::server::ToolCallHandle)s
182/// do not compile. Hence the raw emits.
183fn parallel_calls(ctx: &mut RunContext<Board>) -> Result<()> {
184    let (first, second) = (ctx.new_tool_call_id(), ctx.new_tool_call_id());
185    let (first_result, second_result) = (ctx.new_message_id(), ctx.new_message_id());
186
187    ctx.emit(Event::tool_call_start(first.clone(), "add_task"))?;
188    ctx.emit(Event::tool_call_start(second.clone(), "add_task"))?;
189    ctx.emit(Event::tool_call_args(first.clone(), r#"{"title":"#))?;
190    ctx.emit(Event::tool_call_args(second.clone(), r#"{"title":"#))?;
191    ctx.emit(Event::tool_call_args(first.clone(), r#""write it down"}"#))?;
192    ctx.emit(Event::tool_call_args(second.clone(), r#""read it back"}"#))?;
193    ctx.emit(Event::tool_call_end(first.clone()))?;
194    ctx.emit(Event::tool_call_end(second.clone()))?;
195    ctx.emit(Event::tool_call_result(
196        first_result,
197        first,
198        r#"{"id":1,"title":"write it down"}"#,
199    ))?;
200    ctx.emit(Event::tool_call_result(
201        second_result,
202        second,
203        r#"{"id":2,"title":"read it back"}"#,
204    ))?;
205
206    ctx.update_state(|board| {
207        board.tasks = vec![task(1, "write it down"), task(2, "read it back")];
208    })
209}
210
211/// Reasoning, then text, then a call — all as chunks, with no explicit end
212/// anywhere.
213///
214/// Each stream is closed only by the next one starting, and the last by the end
215/// of the run. A client that waits for an explicit terminator hangs on the
216/// final message forever.
217fn mixed_stream(ctx: &mut RunContext<Board>) -> Result<()> {
218    let thought = ctx.new_message_id();
219    ctx.emit(Event::reasoning_message_chunk(
220        Some(thought),
221        Some("three streams, no brackets".to_owned()),
222    ))?;
223
224    let text = ctx.new_message_id();
225    ctx.emit(Event::text_message_chunk(
226        Some(text),
227        Some("Reading the board".to_owned()),
228    ))?;
229    ctx.emit(Event::text_message_chunk(
230        None,
231        Some(", then adding to it.".to_owned()),
232    ))?;
233
234    let call = ctx.new_tool_call_id();
235    ctx.emit(Event::tool_call_chunk(
236        Some(call),
237        Some("add_task".to_owned()),
238        Some(r#"{"title":"unbracketed"}"#.to_owned()),
239    ))
240}
241
242/// Does work *and then* pauses, in one run.
243///
244/// The interaction the separate scenarios miss: two calls in flight, state
245/// published, and only then a decision the agent needs a human for. What it
246/// exercises on the client side is that a run which already grew the
247/// conversation can still pause — and that resuming carries the tool messages
248/// the first half produced, rather than starting from the user's turn.
249fn busy(ctx: &mut RunContext<Board>) -> Result<RunOutcome> {
250    if ctx.resume_for(BUDGET).is_none() {
251        parallel_calls(ctx)?;
252        ctx.say("Two added. The third needs sign-off.")?;
253        return Ok(RunOutcome::interrupt(vec![Interrupt {
254            id: BUDGET.to_owned(),
255            reason: "tool_approval".to_owned(),
256            message: Some("Add the third task too?".to_owned()),
257            ..Default::default()
258        }]));
259    }
260
261    let approved = ctx
262        .resume_for(BUDGET)
263        .is_some_and(|entry| entry.status == ResumeStatus::Resolved);
264    if approved {
265        let id = ctx.new_tool_call_id();
266        let result = ctx.new_message_id();
267        ctx.emit(Event::tool_call_start(id.clone(), "add_task"))?;
268        ctx.emit(Event::tool_call_args(
269            id.clone(),
270            r#"{"title":"sign it off"}"#,
271        ))?;
272        ctx.emit(Event::tool_call_end(id.clone()))?;
273        ctx.emit(Event::tool_call_result(
274            result,
275            id,
276            r#"{"id":3,"title":"sign it off"}"#,
277        ))?;
278        ctx.update_state(|board| board.tasks.push(task(3, "sign it off")))?;
279        ctx.say("Three on the board.")?;
280    } else {
281        ctx.say("Left it at two.")?;
282    }
283    Ok(RunOutcome::Success)
284}
285
286/// Pauses on two decisions at once.
287///
288/// Two, not one, because answering them one request at a time never terminates
289/// — the agent only ever sees the answers the current request carries.
290fn approvals(ctx: &mut RunContext<Board>) -> RunOutcome {
291    let pending: Vec<Interrupt> = [(BUDGET, "Approve the budget?"), (DATE, "Confirm the date?")]
292        .into_iter()
293        .filter(|(id, _)| ctx.resume_for(id).is_none())
294        .map(|(id, question)| Interrupt {
295            id: id.to_owned(),
296            reason: "tool_approval".to_owned(),
297            message: Some(question.to_owned()),
298            ..Default::default()
299        })
300        .collect();
301
302    if !pending.is_empty() {
303        return RunOutcome::interrupt(pending);
304    }
305
306    let declined: Vec<&str> = [BUDGET, DATE]
307        .into_iter()
308        .filter(|id| {
309            ctx.resume_for(id)
310                .is_some_and(|entry| entry.status == ResumeStatus::Cancelled)
311        })
312        .collect();
313
314    let _ = match declined.len() {
315        0 => ctx.say("Both approved. Booked."),
316        2 => ctx.say("Both declined. Nothing booked."),
317        _ => ctx.say(format!(
318            "Declined: {}. Nothing booked.",
319            declined.join(", ")
320        )),
321    };
322    RunOutcome::Success
323}
324
325/// Says one thing, then waits forever — an agent thirty seconds into a model
326/// call, which is when a user actually hits stop.
327async fn slow(ctx: &mut RunContext<Board>) -> Result<RunOutcome> {
328    ctx.say("working on it, this will take a while")?;
329    std::future::pending::<()>().await;
330    Ok(RunOutcome::Success)
331}
332
333/// The well-behaved turn: reasoning, a bracketed message, state, and a surface.
334///
335/// Here so that pointing the watcher at this server shows the same shape as
336/// pointing it at `task-board`, and the awkward scenarios read as departures
337/// from something rather than as the only thing on offer.
338fn board_turn(ctx: &mut RunContext<Board>) -> Result<()> {
339    let mut step = ctx.step("turn")?;
340    step.think("nothing unusual about this one")?;
341
342    let mut message = step.assistant_message()?;
343    for word in "Two tasks on the board. ".split_inclusive(' ') {
344        message.delta(word)?;
345    }
346    // A handle can reach the run state mid-message, so the reply can quote the
347    // board it is about to publish without closing first.
348    message.state_mut().tasks = vec![task(1, "draft the agenda"), task(2, "book the room")];
349    message.delta(message.state().summary())?;
350    message.end()?;
351
352    step.publish_state()?;
353    surface(&mut step)
354}
355
356/// Ships the board as an A2UI surface in a tool result envelope.
357fn surface(ctx: &mut RunContext<Board>) -> Result<()> {
358    let spec = SurfaceSpec::new("board-watch")
359        .with_components(vec![
360            Component::new("root", "Column").with("children", json!(["heading", "list"])),
361            Component::new("heading", "Text")
362                .with("text", json!({"path": "/title"}))
363                .with("variant", json!("h2")),
364            Component::new("list", "List")
365                .with("children", json!({"componentId": "row", "path": "/tasks"})),
366            Component::new("row", "Text").with("text", json!({"path": "line"})),
367        ])
368        .with_data_model(json!({
369            "title": "Watched board",
370            "tasks": ctx
371                .state()
372                .tasks
373                .iter()
374                .map(|task| json!({"line": task.line()}))
375                .collect::<Vec<_>>(),
376        }));
377
378    let envelope =
379        wrap_as_operations_envelope(&assemble_ops(Intent::Create, &spec)).map_err(Error::agent)?;
380    let mut call = ctx.tool_call(RENDER_A2UI_TOOL_NAME)?;
381    call.args_json(&json!({"surfaceId": "board-watch"}))?;
382    call.result(envelope)?;
383    Ok(())
384}
385
386fn task(id: u32, title: &str) -> Task {
387    Task {
388        id,
389        title: title.to_owned(),
390        estimate_minutes: None,
391        done: false,
392    }
393}
394
395/// The whole fake backend: the agent, the raw endpoint, and a health check.
396pub fn router(agent: Awkward) -> Router {
397    Router::new()
398        .route("/health", get(|| async { "ok" }))
399        .route(RAW_ROUTE, post(raw).get(raw))
400        .route_agui(ROUTE, agent)
401}
402
403/// A hand-framed `text/event-stream`, bypassing `ag-ui-server` entirely.
404///
405/// The server's own verifier will not emit a malformed stream — that is what it
406/// is for — so the only way to hand the *client's* verifier something to reject
407/// is to frame the bytes here, the way a producer in another language does.
408/// [`SseFormatter`](ag_ui::SseFormatter) is the same encoder the real
409/// endpoint uses; only the ordering is wrong.
410async fn raw(axum::extract::Path(scenario): axum::extract::Path<String>) -> Response {
411    use ag_ui::SseFormatter;
412
413    let formatter = SseFormatter::new();
414    let mut body = String::new();
415    for event in raw_script(&scenario) {
416        match formatter.encode_to_string(&event) {
417            Ok(frame) => body.push_str(&frame),
418            Err(error) => {
419                return (
420                    axum::http::StatusCode::INTERNAL_SERVER_ERROR,
421                    error.to_string(),
422                )
423                    .into_response();
424            }
425        }
426    }
427
428    (
429        [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
430        body,
431    )
432        .into_response()
433}
434
435/// The scripted illegal streams, by name.
436///
437/// - `unbracketed` — content for a message that was never started.
438/// - `truncated` — a run that stops without saying how.
439/// - `orphan-result` — a result for a call nobody made.
440pub fn raw_script(scenario: &str) -> Vec<Event> {
441    let started = Event::run_started("raw", "raw-1");
442    match scenario {
443        "unbracketed" => vec![
444            started,
445            Event::text_message_content(MessageId::new("ghost"), "text nobody opened"),
446            Event::run_finished_success("raw", "raw-1"),
447        ],
448        "truncated" => vec![
449            started,
450            Event::text_message_start(MessageId::new("cut"), TextMessageRole::Assistant),
451            Event::text_message_content(MessageId::new("cut"), "half a sen"),
452        ],
453        "orphan-result" => vec![
454            started,
455            Event::tool_call_result(
456                MessageId::new("answer"),
457                ToolCallId::new("never-called"),
458                "{}",
459            ),
460            Event::run_finished_success("raw", "raw-1"),
461        ],
462        _ => vec![started, Event::run_finished_success("raw", "raw-1")],
463    }
464}
board_watch/fake.rs

board_watch/
fake.rs
