Skip to main content

wcode_harness/
actor.rs

1//! In-process session actor — one mailbox, one stream.
2//!
3//! [`SessionActor`] turns an [`Agent`] into a task with an inbox and a
4//! subscribable outbox, so a caller drives the session entirely through
5//! [`Request`]s and observes it through [`AgentEvent`]s — never by holding the
6//! `Agent` itself:
7//!
8//! ```text
9//!   caller  (REPL / TUI / future peer session — all identical here)
10//!     │  send(Request)                          ▲  subscribe() → recv()
11//!     ▼                                         │  (own cursor per subscriber)
12//!   inbox: mpsc::Unbounded<Request>       outbox: broadcast<AgentEvent>
13//!     │                                         ▲  events.send(ev)
14//!     └───────────► SessionActor task ◄─────────┘
15//!                   owns the Agent; serves one Request at a time;
16//!                   `deferred: VecDeque<Request>` parks mid-run commands
17//! ```
18//!
19//! This is the seam the whole protocol design hangs on. A REPL, a TUI, and
20//! (later) a peer session are all just callers that hold a [`SessionHandle`];
21//! they differ only in how they render the stream and which address they hold,
22//! not in the mechanism they speak. Today the inbox and outbox are in-process
23//! channels; the shapes are chosen so that a socket ([`Frame`]s over NDJSON,
24//! stage S2) can replace the transport without changing this API.
25//!
26//! ## Flow
27//!
28//! `serve()` — no run in flight: pull ONE Request, service it, repeat.
29//!
30//! ```text
31//!   deferred non-empty? ──yes──► take from deferred (FIFO, order kept)
32//!        │ no
33//!        ▼
34//!   inbox.recv().await ──► match Request:
35//!        Submit    {..} ─► run(...)              → enters the in-run map below
36//!        Notify    {..} ─► agent.notify(tagged)  → append, no turn
37//!        Interrupt {..} ─► agent.steer(tagged)   → queued for the next turn
38//!        Wake      {..} ─► run(tagged)           → a turn even when idle
39//!        Cancel         ─► (no-op — nothing is running)
40//!        SetModel  {..} ─► agent.set_model(..)
41//!        SetEffort {..} ─► agent.set_effort(..)
42//!        Compact   {..} ─► agent.compact(..).await
43//!        GetHistory     ─► reply with the context
44//!        Usage          ─► reply with folded stats (+ $ when priced)
45//!        Unknown        ─► (ignored)
46//! ```
47//!
48//! `run()` — the Agent is borrowed for the whole run, so it is reached only
49//! through clones taken *before* the borrow.
50//!
51//! ```text
52//!   steer     = agent.steer_sender()
53//!   follow_up = agent.follow_up_sender()
54//!   cancel    = agent.cancel_token()
55//!
56//!   loop { tokio::select! { biased; ... } }   ← biased: inbox polled FIRST
57//!        ├─ Some(req) = inbox.recv() ─► Cancel   ─► cancel.cancel()         ┐ to the
58//!        │                              Steer    ─► steer.send(msg)         │ live run
59//!        │                              FollowUp ─► follow_up.send(msg)     │
60//!        │                              other    ─► deferred.push_back(req) ┘ waits
61//!        ├─ result  = &mut run  ──────► run finished → break
62//!        └─ Some(ev) = sink_rx.recv() ► events.send(ev)  (fan-out to subscribers)
63//!   }
64//!   then: drain leftover sink_rx → events.send ; on Err(run) → Error event
65//! ```
66//!
67//! [`Frame`]: crate::protocol::Frame
68//!
69//! ## Delivery rules
70//!
71//! A run holds the agent's `&mut` for its whole life, so the actor splits the
72//! inbox by what each request needs:
73//!
74//! | request | while a run is in flight |
75//! |---------|--------------------------|
76//! | `Cancel` | delivered at once — the token is cancelled; the loop aborts |
77//! | `Interrupt` / `Wake` | forwarded to the run's channels at once |
78//! | `SetModel` / `SetEffort` / `Compact` | deferred, applied when the run ends |
79//! | another `Submit` | deferred, becomes the next run |
80//!
81//! The interactive path therefore never waits on a run. The deferred set is
82//! exactly the requests that need `&mut Agent` — you cannot swap the model or
83//! compact the context out from under a run — and deferring them is also the
84//! documented meaning: `set_model`/`set_effort` "take effect on the next run".
85//!
86//! One caveat lives in the kernel, not here — the actor hands a `Cancel` over at
87//! once, but the *loop* is what acts on it, and its latency depends on where the
88//! run is:
89//!
90//! ```text
91//! where the run is when Cancel arrives          when it stops
92//! ---------------------------------------------  -----------------------
93//! model streaming     (loop races the token)     immediately
94//! between turns       (checked at turn start)    immediately
95//! inside a tool call  (no race on the future)    when the tool returns *
96//! ```
97//!
98//! `*` unless the tool watches `ToolContext::cancel` itself (e.g. `bash` kills
99//! its child) — then it can stop early.
100//!
101//! ## Requests and replies
102//!
103//! `send` is fire-and-forget; `ask` awaits a reply. The reply is an
104//! [`AgentEvent`] reply variant — [`AgentEvent::History`] for `GetHistory`,
105//! [`AgentEvent::Ack`]/[`AgentEvent::Error`] for `SetModel`/`SetEffort`/
106//! [`AgentEvent::Compaction`] for `Compact`, [`AgentEvent::Usage`] for `Usage`
107//! — correlated in-process by a
108//! one-shot channel. Replies are never fanned out on the event stream; the wire
109//! (S2) will carry the same variants correlated by the envelope's `reply_to`.
110//!
111//! Scope so far: every request variant that maps onto an existing [`Agent`]
112//! method is serviced. Still to come (S1b): pointing the CLI at a handle.
113
114use std::collections::VecDeque;
115use std::fmt;
116
117use tokio::sync::{broadcast, mpsc, oneshot, watch};
118
119use crate::agent::Agent;
120use crate::compaction::CompactOutcome;
121use crate::event::AgentEvent;
122use crate::hooks::HooksSet;
123use crate::loop_::LoopError;
124use crate::message::{AgentMessage, StopReason};
125use crate::protocol::{MemberState, Request, SessionId};
126use crate::streamfn::LlmOpts;
127
128/// Default outbox buffer: how far a subscriber may lag before it starts losing
129/// events (`broadcast` semantics — the slowest reader drops, it never blocks the
130/// run).
131pub const EVENT_BUFFER: usize = 1024;
132
133/// Handle to a running [`SessionActor`].
134///
135/// Cloning gives another caller the same mailbox and stream. When the last
136/// handle drops, the session shuts down and its [`Agent`] is dropped.
137#[derive(Clone)]
138pub struct SessionHandle {
139    inbox: mpsc::UnboundedSender<Message>,
140    events: broadcast::Sender<AgentEvent>,
141    /// The run-state cell — the ONE server-side liveness fact (bug 2). A `watch`
142    /// (not a bare `AtomicU8`) so a registry/server task can WAKE on a transition
143    /// and re-push the roster; `send_replace` keeps the latest value even with no
144    /// receiver yet. `Clone` rides freely (the sender is shared).
145    state: watch::Sender<MemberState>,
146}
147
148impl SessionHandle {
149    /// Submit a request to the session's inbox.
150    ///
151    /// Fails only if the session has shut down (every handle was dropped and the
152    /// actor task has ended). A `Submit` runs to completion before the next
153    /// non-interrupt request is serviced, but `Interrupt`/`Wake`/`Cancel` sent
154    /// during a run are handled immediately (see the actor's delivery rules).
155    pub fn send(&self, request: Request) -> Result<(), SessionClosed> {
156        self.tell(None, request)
157    }
158
159    /// Like [`SessionHandle::send`], but attributed to a sender address so the
160    /// recipient can attribute and policy-check the message (a peer session
161    /// passes its own `"agent:<id>"` address).
162    pub fn send_from(&self, from: SessionId, request: Request) -> Result<(), SessionClosed> {
163        self.tell(Some(from), request)
164    }
165
166    fn tell(&self, from: Option<SessionId>, request: Request) -> Result<(), SessionClosed> {
167        self.inbox
168            .send(Message::Tell { from, request })
169            .map_err(|_| SessionClosed)
170    }
171
172    /// Send a request and await its reply.
173    ///
174    /// `ask` is for requests that expect a reply: `GetHistory` (→
175    /// [`AgentEvent::History`]), `SetModel`/`SetEffort` (→ [`AgentEvent::Ack`]
176    /// or [`AgentEvent::Error`]), `Compact` (→ [`AgentEvent::Compaction`],
177    /// [`AgentEvent::Ack`], or [`AgentEvent::Error`]). It also works for
178    /// `Submit`, replying when the run completes. Interrupts
179    /// (`Interrupt`/`Wake`/`Cancel`) want no reply — use `send`; an `ask` on one
180    /// still lands mid-run (the actor answers it `Ack`), so a server-side `cancel`
181    /// is never deferred.
182    ///
183    /// Fails only if the session has shut down.
184    pub async fn ask(&self, request: Request) -> Result<AgentEvent, SessionClosed> {
185        self.question(None, request).await
186    }
187
188    /// Like [`SessionHandle::ask`], attributed to a sender address.
189    pub async fn ask_from(
190        &self,
191        from: SessionId,
192        request: Request,
193    ) -> Result<AgentEvent, SessionClosed> {
194        self.question(Some(from), request).await
195    }
196
197    async fn question(
198        &self,
199        from: Option<SessionId>,
200        request: Request,
201    ) -> Result<AgentEvent, SessionClosed> {
202        let (reply_tx, reply_rx) = oneshot::channel();
203        self.inbox
204            .send(Message::Ask {
205                from,
206                request,
207                reply: reply_tx,
208            })
209            .map_err(|_| SessionClosed)?;
210        reply_rx.await.map_err(|_| SessionClosed)
211    }
212
213    /// The session's current run-state. Read-only, lock-free (a `watch` borrow),
214    /// so a `peers` listing / roster read over N members is cheap.
215    pub fn state(&self) -> MemberState {
216        *self.state.borrow()
217    }
218
219    /// A live view of the run-state, so a registry/server task can wake on a
220    /// transition instead of polling. Seeded with the current value (like
221    /// `Registry::subscribe`).
222    pub fn state_rx(&self) -> watch::Receiver<MemberState> {
223        self.state.subscribe()
224    }
225    /// Subscribe to the session's event stream. Each subscriber gets its own
226    /// receiver; a lagging subscriber drops events rather than stalling the run.
227    pub fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
228        self.events.subscribe()
229    }
230}
231
232/// Returned when a [`SessionHandle::send`] targets a session that has ended.
233#[derive(Debug, Clone, Copy, PartialEq, Eq)]
234pub struct SessionClosed;
235
236impl fmt::Display for SessionClosed {
237    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
238        f.write_str("session actor has shut down")
239    }
240}
241
242impl std::error::Error for SessionClosed {}
243
244/// A mailbox item: fire-and-forget, or a request that expects a reply.
245enum Message {
246    /// No reply wanted. `from` is the sender's address — `None` for a local
247    /// caller (the human), `Some(peer)` when another session sends.
248    Tell {
249        from: Option<SessionId>,
250        request: Request,
251    },
252    /// The reply to the request is sent here when it is serviced.
253    Ask {
254        from: Option<SessionId>,
255        request: Request,
256        reply: oneshot::Sender<AgentEvent>,
257    },
258}
259
260/// Spawns the actor task that owns an [`Agent`].
261///
262/// The `Agent` moves into a background task; all interaction returns through the
263/// [`SessionHandle`]. The task exits when the last handle drops.
264pub struct SessionActor;
265
266impl SessionActor {
267    pub fn spawn(agent: Agent) -> SessionHandle {
268        let (inbox_tx, inbox_rx) = mpsc::unbounded_channel();
269        let (events_tx, _) = broadcast::channel(EVENT_BUFFER);
270        let events = events_tx.clone();
271        let (state_tx, _state_rx) = watch::channel(MemberState::Idle);
272        tokio::spawn(serve(agent, inbox_rx, events_tx, state_tx.clone()));
273        SessionHandle {
274            inbox: inbox_tx,
275            events,
276            state: state_tx,
277        }
278    }
279}
280
281/// The actor loop: pull one request at a time, service it, repeat.
282async fn serve(
283    mut agent: Agent,
284    mut inbox: mpsc::UnboundedReceiver<Message>,
285    events: broadcast::Sender<AgentEvent>,
286    state: watch::Sender<MemberState>,
287) {
288    // Requests that arrive *during* a run but cannot be applied to it (a
289    // second `Submit`, a model swap) wait here and are serviced once the run
290    // ends, preserving their order.
291    let mut deferred: VecDeque<Message> = VecDeque::new();
292
293    loop {
294        let message = match deferred.pop_front() {
295            Some(message) => message,
296            None => match inbox.recv().await {
297                Some(message) => message,
298                None => break, // every handle dropped: shut down
299            },
300        };
301
302        dispatch(&mut agent, message, &mut inbox, &events, &mut deferred, &state).await;
303    }
304}
305
306/// Prefix an inbound peer message with its sender, so the model sees who sent
307/// it and knows the address to reply to (§13.13).
308fn tag(sender: &SessionId, content: &str) -> String {
309    format!("[message from {sender}]\n{content}")
310}
311
312/// Run the inbound policy on an A2A message (the mirror of
313/// `before_tool_call`). Returns `Some(reason)` if the message was dropped; the
314/// request may have been rewritten in place by the hook either way.
315async fn gate_inbound(hooks: &HooksSet, message: &mut Message) -> Option<String> {
316    let (from, request) = match message {
317        Message::Tell { from, request } | Message::Ask { from, request, .. } => (from, request),
318    };
319    if !crate::protocol::is_inbound(request) {
320        return None;
321    }
322    hooks.before_inbound(from.as_ref(), request).await
323}
324
325/// Build the `Request::Usage` reply: fold the transcript, then price it with
326/// the provider THIS session bills against — the one body both arms share, so
327/// an idle and a mid-run report can never disagree about what a report is.
328///
329/// Contract:
330/// - the fold is the pure [`crate::stats::session_stats`] over `messages` (the
331///   idle arm passes the live `Agent::messages()`, the mid-run arm the
332///   turn-boundary snapshot);
333/// - `cost` comes from [`crate::pricing::cost_of`] on the SAME slice with
334///   `llm.base_url.as_deref()` / `&llm.model`, so a model switched mid-session
335///   is priced per message and the turns that cannot be priced are COUNTED
336///   (`Cost::unpriced_turns`), not silently dropped;
337/// - an unknown provider or model yields `cost: None` inside an otherwise
338///   valid report — a missing price is never an [`AgentEvent::Error`];
339/// - empty `messages` ⇒ a `SessionStats::default()` in the reply, i.e.
340///   "(no usage reported)" downstream: a fresh session answers, it does not fail.
341///
342/// Does NOT promise: a `$` for a non-zen provider, tier rates, or a per-model
343/// `base_url` (see `pricing`'s module doc). Panics on nothing.
344fn usage_reply(messages: &[AgentMessage], llm: &LlmOpts) -> AgentEvent {
345    let mut stats = crate::stats::session_stats(messages);
346    stats.cost = crate::pricing::cost_of(messages, llm.base_url.as_deref(), &llm.model);
347    AgentEvent::Usage { stats }
348}
349
350/// Service one mailbox item: run the request, then answer an `Ask` with the
351/// reply it produced (a `Tell` drops the reply on the floor).
352async fn dispatch(
353    agent: &mut Agent,
354    mut message: Message,
355    inbox: &mut mpsc::UnboundedReceiver<Message>,
356    events: &broadcast::Sender<AgentEvent>,
357    deferred: &mut VecDeque<Message>,
358    state: &watch::Sender<MemberState>,
359) {
360    // Inbound peer-message policy (S4-1): a drop skips delivery entirely; a
361    // rewrite (in place) is what gets serviced.
362    if let Some(reason) = gate_inbound(agent.hooks(), &mut message).await {
363        if let Message::Ask { reply, .. } = message {
364            let _ = reply.send(AgentEvent::Error {
365                message: format!("inbound blocked: {reason}"),
366            });
367        }
368        return;
369    }
370
371    let (from, request, reply) = match message {
372        Message::Tell { from, request } => (from, request, None),
373        Message::Ask { from, request, reply } => (from, request, Some(reply)),
374    };
375    let sender = from.unwrap_or_else(SessionId::user);
376
377    let event = match request {
378        Request::Submit { text } => match run(agent, &text, inbox, events, deferred, state).await {
379            Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
380            Err(e) => AgentEvent::Error {
381                message: e.to_string(),
382            },
383        },
384        // A2A delivery verbs (S4-1). `Notify` appends with no turn;
385        // `Interrupt`/`Wake` are the canonical interrupt/continue forms.
386        Request::Notify { content } => match agent.notify(tag(&sender, &content)) {
387            Ok(_) => {
388                let _ = events.send(AgentEvent::MessageReceived {
389                    from: sender.clone(),
390                    content,
391                });
392                AgentEvent::Ack
393            }
394            Err(e) => AgentEvent::Error {
395                message: e.to_string(),
396            },
397        },
398        Request::Interrupt { content } => {
399            agent.steer(AgentMessage::user_text(tag(&sender, &content)));
400            let _ = events.send(AgentEvent::MessageReceived {
401                from: sender.clone(),
402                content,
403            });
404            AgentEvent::Ack
405        }
406        // A `Wake` runs even when idle: start a turn with the (tagged) content,
407        // like `Submit` — otherwise a worker handed a task would only queue a
408        // follow-up and never act on it.
409        Request::Wake { content } => {
410            let _ = events.send(AgentEvent::MessageReceived {
411                from: sender.clone(),
412                content: content.clone(),
413            });
414            match run(agent, &tag(&sender, &content), inbox, events, deferred, state).await {
415                Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
416                Err(e) => AgentEvent::Error {
417                    message: e.to_string(),
418                },
419            }
420        }
421        // An idle cancel is a no-op (see the delivery rules): the token is
422        // minted fresh at the end of each run, so cancelling an idle session
423        // would otherwise abort the next run on its first poll.
424        Request::Cancel => AgentEvent::Ack,
425        Request::SetModel { model } => match agent.set_model(model) {
426            Ok(()) => AgentEvent::Ack,
427            Err(e) => AgentEvent::Error {
428                message: e.to_string(),
429            },
430        },
431        Request::SetEffort { effort } => match agent.set_effort(effort) {
432            Ok(()) => AgentEvent::Ack,
433            Err(e) => AgentEvent::Error {
434                message: e.to_string(),
435            },
436        },
437        Request::Compact { instructions } => match agent.compact(instructions.as_deref()).await {
438            Ok(CompactOutcome::NothingToDo) => AgentEvent::Ack,
439            Ok(CompactOutcome::Done {
440                summarized, kept, ..
441            }) => AgentEvent::Compaction { summarized, kept },
442            Err(e) => AgentEvent::Error { message: e },
443        },
444        // `/btw`: a tool-free side call. `side_ask` takes `&self`, so no borrow
445        // conflict with the run. A `SideAsk` arriving mid-run is deferred to a
446        // turn boundary like any other inbox request (see the run loop's
447        // `deferred.push_back`) — queued, not refused.
448        Request::SideAsk { text } => match agent.side_ask(&text).await {
449            Ok(answer) => AgentEvent::SideAnswer {
450                text: answer.text,
451                usage: answer.usage,
452            },
453            Err(e) => AgentEvent::Error {
454                message: e.to_string(),
455            },
456        },
457        // Toggle plan mode: flips the shared handle (the `PlanModeHooks` reads
458        // the same one) and recomposes the system prompt. Infallible.
459        Request::SetPlanMode { on } => {
460            agent.set_plan_mode(on);
461            AgentEvent::Ack
462        }
463        Request::GetHistory => AgentEvent::History {
464            messages: agent.messages().to_vec(),
465        },
466        Request::Status => AgentEvent::Status {
467            last_assistant_text: crate::message::last_assistant_text(agent.messages()),
468        },
469        Request::Usage => usage_reply(agent.messages(), agent.llm()),
470        // A served peer is defined by a socket server's injected handler; an
471        // in-process session has no factory, so it can only refuse.
472        Request::Define { .. } => AgentEvent::Error {
473            message: "agent definition is served, not in-process".into(),
474        },
475        Request::ListSessions | Request::ListModels | Request::Unknown => AgentEvent::Ack,
476    };
477
478    if let Some(reply) = reply {
479        let _ = reply.send(event);
480    }
481}
482
483/// Drive one `Agent::run`, forwarding its events to the outbox while still
484/// servicing the inbox so `Interrupt`/`Wake`/`Cancel` reach the in-flight run.
485async fn run(
486    agent: &mut Agent,
487    text: &str,
488    inbox: &mut mpsc::UnboundedReceiver<Message>,
489    events: &broadcast::Sender<AgentEvent>,
490    deferred: &mut VecDeque<Message>,
491    state: &watch::Sender<MemberState>,
492) -> Result<StopReason, LoopError> {
493    // State transitions live at the run boundary — this `run()` is the in-flight
494    // window. The event stream's `AgentStart`/`AgentEnd` are the OTHER place this
495    // could be derived; the actor-owned cell is authoritative because it is set
496    // synchronously before the first await, so a reader never misses `Running`.
497    // A `Cancel` ends on `Ok(StopReason::Aborted)` -> `Done`: an abort is not a
498    // failure, matching the TUI.
499    state.send_replace(MemberState::Running);
500    let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::<AgentEvent>();
501    // Sender/token clones taken *before* `run` borrows the agent: steering and
502    // cancelling the in-flight run must not need the `&mut`.
503    let steer = agent.steer_sender();
504    let follow_up = agent.follow_up_sender();
505    let cancel = agent.cancel_token();
506    let hooks = agent.hooks().clone();
507    // Read-only handle for mid-run `GetHistory`/`SideAsk` (see the arms below).
508    let side = agent.side_handle();
509
510    let mut run = Box::pin(agent.run(text, sink_tx));
511    let result = loop {
512        // Biased: service the inbox *before* advancing the run. An `Interrupt` must
513        // be in the agent's steering channel before the run reaches its next
514        // turn-boundary drain, or it is silently a turn late. Polling the run
515        // first would let it cross that boundary while the request still sits
516        // here.
517        tokio::select! {
518            biased;
519            Some(mut message) = inbox.recv() => {
520                if let Some(reason) = gate_inbound(&hooks, &mut message).await {
521                    if let Message::Ask { reply, .. } = message {
522                        let _ = reply.send(AgentEvent::Error {
523                            message: format!("inbound blocked: {reason}"),
524                        });
525                    }
526                    continue;
527                }
528                match message {
529                    // The interrupt verbs (Cancel / Notify-Interrupt / Wake) want NO reply —
530                    // but the server delivers EVERY request via `ask` (`server.rs`), so handle
531                    // them for BOTH `Tell` and `Ask`, or an `Ask{Cancel}` falls through to
532                    // `deferred` and is queued until the run ENDS: the interrupt never lands.
533                    // An `Ask` gets its `Ack`; a `Tell` needs none.
534                    Message::Tell {
535                        request: Request::Cancel,
536                        ..
537                    } => cancel.cancel(),
538                    Message::Ask {
539                        request: Request::Cancel,
540                        reply,
541                        ..
542                    } => {
543                        cancel.cancel();
544                        let _ = reply.send(AgentEvent::Ack);
545                    }
546                    Message::Tell {
547                        request: Request::Notify { content } | Request::Interrupt { content },
548                        from,
549                    } => {
550                        let from = from.unwrap_or_else(SessionId::user);
551                        let _ = steer.send(AgentMessage::user_text(tag(&from, &content)));
552                        let _ = events.send(AgentEvent::MessageReceived {
553                            from,
554                            content,
555                        });
556                    }
557                    Message::Ask {
558                        request: Request::Notify { content } | Request::Interrupt { content },
559                        from,
560                        reply,
561                    } => {
562                        let from = from.unwrap_or_else(SessionId::user);
563                        let _ = steer.send(AgentMessage::user_text(tag(&from, &content)));
564                        let _ = events.send(AgentEvent::MessageReceived {
565                            from,
566                            content,
567                        });
568                        let _ = reply.send(AgentEvent::Ack);
569                    }
570                    Message::Tell {
571                        request: Request::Wake { content },
572                        from,
573                    } => {
574                        let from = from.unwrap_or_else(SessionId::user);
575                        let _ = follow_up.send(AgentMessage::user_text(tag(&from, &content)));
576                        let _ = events.send(AgentEvent::MessageReceived {
577                            from,
578                            content,
579                        });
580                    }
581                    Message::Ask {
582                        request: Request::Wake { content },
583                        from,
584                        reply,
585                    } => {
586                        let from = from.unwrap_or_else(SessionId::user);
587                        let _ = follow_up.send(AgentMessage::user_text(tag(&from, &content)));
588                        let _ = events.send(AgentEvent::MessageReceived {
589                            from,
590                            content,
591                        });
592                        let _ = reply.send(AgentEvent::Ack);
593                    }
594                    // Read-only queries are answered mid-run from the context
595                    // snapshot `run_loop` republishes each turn — a `/btw` or
596                    // `/usage` never waits for the whole run. `side` was cloned
597                    // before the run borrowed the agent.
598                    Message::Ask {
599                        request: Request::GetHistory,
600                        reply,
601                        ..
602                    } => {
603                        let _ = reply.send(AgentEvent::History { messages: side.history() });
604                    }
605                    Message::Ask {
606                        request: Request::Status,
607                        reply,
608                        ..
609                    } => {
610                        let _ = reply.send(AgentEvent::Status {
611                            last_assistant_text: side.last_assistant_text(),
612                        });
613                    }
614                    // The mid-run twin of the idle arm: answered AT ONCE from the
615                    // turn-boundary snapshot, so a `/usage` never waits for the run
616                    // to finish. `side.history()` clones the snapshot; the fold needs
617                    // the slice and the reply owns the result. This arm must NOT fall
618                    // through to `deferred` — a read that queues behind a run is the
619                    // bug `GetHistory` already fixed.
620                    Message::Ask {
621                        request: Request::Usage,
622                        reply,
623                        ..
624                    } => {
625                        let _ = reply.send(usage_reply(&side.history(), side.llm()));
626                    }
627                    Message::Ask {
628                        request: Request::SideAsk { text },
629                        reply,
630                        ..
631                    } => {
632                        // A side ask is itself a model call: spawn it so it runs
633                        // alongside the run instead of stalling the inbox.
634                        let side = side.clone();
635                        tokio::spawn(async move {
636                            let event = match side.side_ask(&text).await {
637                                Ok(answer) => AgentEvent::SideAnswer {
638                                    text: answer.text,
639                                    usage: answer.usage,
640                                },
641                                Err(e) => AgentEvent::Error {
642                                    message: e.to_string(),
643                                },
644                            };
645                            let _ = reply.send(event);
646                        });
647                    }
648                    // A mutation cannot apply mid-run (it needs `&mut Agent`), but
649                    // its contract is "next run": acknowledge now and re-queue it
650                    // as a `Tell` so it is applied silently at the run boundary —
651                    // instead of leaving the ask to time out (item-41 ASK_TIMEOUT).
652                    Message::Ask {
653                        from,
654                        request:
655                            req @ (Request::SetModel { .. }
656                            | Request::SetEffort { .. }
657                            | Request::SetPlanMode { .. }),
658                        reply,
659                    } => {
660                        let _ = reply.send(AgentEvent::Ack);
661                        deferred.push_back(Message::Tell { from, request: req });
662                    }
663                    other => deferred.push_back(other),
664                }
665            },
666            result = &mut run => break result,
667            Some(event) = sink_rx.recv() => {
668                let _ = events.send(event);
669            }
670        }
671    };
672
673    // The run dropped its sink on completion; forward whatever is still buffered
674    // so the last events (e.g. `AgentEnd`) are never stranded.
675    while let Ok(event) = sink_rx.try_recv() {
676        let _ = events.send(event);
677    }
678
679    // A session-write error inside the run is a `LoopError`; the run's own
680    // `AgentEvent::Error` covers stream failures, so surface this one the same
681    // way rather than dropping it silently.
682    if let Err(e) = &result {
683        let _ = events.send(AgentEvent::Error {
684            message: e.to_string(),
685        });
686    }
687    // Final transition: a normal stop is `Done`; a run-failure — the loop's
688    // `StopReason::Error` or a `LoopError` — is `Failed` (sticky until the next
689    // run sets `Running`).
690    match &result {
691        Ok(StopReason::Error) | Err(_) => {
692            state.send_replace(MemberState::Failed);
693        }
694        Ok(_) => {
695            state.send_replace(MemberState::Done);
696        }
697    }
698    result
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704    use crate::agent::{Agent, AgentConfig};
705    use crate::compaction::CompactionPolicy;
706    use crate::event::LlmStreamEvent;
707    use crate::hooks::HooksSet;
708    use crate::message::StopReason;
709    use crate::streamfn::{LlmOpts, LlmStream, StreamFn};
710    use crate::tool::{ToolContext, ToolOutput, TypedTool, erased};
711    use futures::StreamExt as _;
712    use serde::Deserialize;
713    use serde_json::json;
714    use std::sync::{Arc, Mutex};
715    use std::time::Duration;
716
717    #[derive(Clone, Default)]
718    struct Recorder {
719        script: Arc<Mutex<VecDeque<Vec<LlmStreamEvent>>>>,
720        calls: Arc<Mutex<Vec<Vec<AgentMessage>>>>,
721        models: Arc<Mutex<Vec<String>>>,
722    }
723
724    impl Recorder {
725        fn push(&self, events: Vec<LlmStreamEvent>) {
726            self.script.lock().unwrap().push_back(events);
727        }
728        fn calls(&self) -> Vec<Vec<AgentMessage>> {
729            self.calls.lock().unwrap().clone()
730        }
731        fn models(&self) -> Vec<String> {
732            self.models.lock().unwrap().clone()
733        }
734    }
735
736    fn fake_stream_fn(rec: &Recorder) -> StreamFn {
737        let rec = rec.clone();
738        Arc::new(
739            move |ctx: &[AgentMessage], _system, _tools, opts: &LlmOpts| {
740                rec.calls.lock().unwrap().push(ctx.to_vec());
741                rec.models.lock().unwrap().push(opts.model.clone());
742                let events = rec.script.lock().unwrap().pop_front().unwrap_or_default();
743                Box::pin(futures::stream::iter(events)) as LlmStream
744            },
745        )
746    }
747
748    fn agent_config(stream_fn: StreamFn, tools: Vec<crate::tool::Tool>) -> AgentConfig {
749        AgentConfig {
750            system: "sys".into(),
751            tools,
752            llm: LlmOpts {
753                model: "m1".into(),
754                ..LlmOpts::default()
755            },
756            stream_fn,
757            hooks: HooksSet::default(),
758            session: None,
759            context: Vec::new(),
760            working_dir: std::path::PathBuf::new(),
761            max_turns: crate::loop_::DEFAULT_MAX_TURNS,
762            parallel_tools: true,
763            compaction: CompactionPolicy::default(),
764            plan_mode: crate::hooks::PlanModeHandle::new(),
765        }
766    }
767
768    /// Block until the run ends, panicking if it stalls.
769    async fn wait_for_end(rx: &mut broadcast::Receiver<AgentEvent>) {
770        loop {
771            let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
772                .await
773                .expect("an event within the timeout")
774                .expect("the outbox stays open");
775            if matches!(event, AgentEvent::AgentEnd) {
776                return;
777            }
778        }
779    }
780
781    #[tokio::test]
782    async fn submit_streams_events_to_a_subscriber() {
783        let rec = Recorder::default();
784        rec.push(vec![
785            LlmStreamEvent::TextDelta("hello".into()),
786            LlmStreamEvent::Done {
787                stop_reason: StopReason::Stop,
788                usage: None,
789            },
790        ]);
791        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
792        let mut rx = handle.subscribe();
793        handle.send(Request::Submit { text: "hi".into() }).unwrap();
794
795        let mut saw_text = false;
796        loop {
797            let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
798                .await
799                .expect("an event")
800                .expect("open");
801            match event {
802                AgentEvent::MessageUpdate { message } if message.as_text() == "hello" => {
803                    saw_text = true;
804                }
805                AgentEvent::AgentEnd => break,
806                _ => {}
807            }
808        }
809        assert!(saw_text, "the streamed text reached the subscriber");
810    }
811
812    #[tokio::test]
813    async fn idle_cancel_does_not_poison_the_next_run() {
814        let rec = Recorder::default();
815        for text in ["one", "two"] {
816            rec.push(vec![
817                LlmStreamEvent::TextDelta(text.into()),
818                LlmStreamEvent::Done {
819                    stop_reason: StopReason::Stop,
820                    usage: None,
821                },
822            ]);
823        }
824        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
825        let mut rx = handle.subscribe();
826
827        handle.send(Request::Submit { text: "a".into() }).unwrap();
828        wait_for_end(&mut rx).await;
829
830        // No run in flight: this must be swallowed, not arm a kill for run two.
831        handle.send(Request::Cancel).unwrap();
832        handle.send(Request::Submit { text: "b".into() }).unwrap();
833
834        let mut saw_two = false;
835        loop {
836            let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
837                .await
838                .expect("an event")
839                .expect("open");
840            match event {
841                AgentEvent::MessageUpdate { message } if message.as_text() == "two" => {
842                    saw_two = true;
843                }
844                AgentEvent::AgentEnd => break,
845                _ => {}
846            }
847        }
848        assert!(saw_two, "the second run completed despite the idle cancel");
849    }
850
851    #[tokio::test]
852    async fn cancel_aborts_an_in_flight_run() {
853        // One delta, then the stream hangs: only a working cancel ends the run.
854        let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
855            let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
856            Box::pin(head.chain(futures::stream::pending())) as LlmStream
857        });
858        let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
859        let mut rx = handle.subscribe();
860        handle.send(Request::Submit { text: "hi".into() }).unwrap();
861
862        loop {
863            let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
864                .await
865                .expect("an event")
866                .expect("open");
867            if matches!(event, AgentEvent::MessageUpdate { .. }) {
868                handle.send(Request::Cancel).unwrap();
869                break;
870            }
871        }
872        // Without a live cancel this would time out on the pending stream.
873        wait_for_end(&mut rx).await;
874    }
875
876    #[tokio::test]
877    async fn ask_cancel_aborts_an_in_flight_run() {
878        // The server delivers EVERY request via `ask`; an `Ask{Cancel}` must abort the run
879        // mid-flight, NOT be deferred to its end (the bug this fixes). One delta, then a
880        // hang: only a live cancel ends the run.
881        let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
882            let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
883            Box::pin(head.chain(futures::stream::pending())) as LlmStream
884        });
885        let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
886        let mut rx = handle.subscribe();
887        handle.send(Request::Submit { text: "hi".into() }).unwrap();
888
889        loop {
890            let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
891                .await
892                .expect("an event")
893                .expect("open");
894            if matches!(event, AgentEvent::MessageUpdate { .. }) {
895                // ASK for the cancel: deferred, this hangs and the 2 s timeout fires (the
896                // pre-fix failure); handled, it replies `Ack` at once.
897                let reply = tokio::time::timeout(Duration::from_secs(2), handle.ask(Request::Cancel))
898                    .await
899                    .expect("the cancel ask is answered mid-run, not deferred")
900                    .unwrap();
901                assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
902                break;
903            }
904        }
905        // Without a live cancel this would time out on the pending stream.
906        wait_for_end(&mut rx).await;
907    }
908
909    #[tokio::test]
910    async fn state_is_running_during_a_turn_then_done_after() {
911        // A stream that emits one delta then hangs: the run is observably in
912        // flight until a cancel ends it.
913        let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
914            let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
915            Box::pin(head.chain(futures::stream::pending())) as LlmStream
916        });
917        let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
918        assert_eq!(handle.state(), MemberState::Idle, "a virgin session is idle");
919
920        // A watcher wakes on the transition rather than polling.
921        let mut state_rx = handle.state_rx();
922        handle.send(Request::Submit { text: "hi".into() }).unwrap();
923        tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
924            .await
925            .expect("a state change")
926            .expect("open");
927        assert_eq!(*state_rx.borrow_and_update(), MemberState::Running);
928        assert_eq!(handle.state(), MemberState::Running, "the run is in flight");
929
930        handle.send(Request::Cancel).unwrap();
931        tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
932            .await
933            .expect("a state change")
934            .expect("open");
935        assert_eq!(
936            *state_rx.borrow_and_update(),
937            MemberState::Done,
938            "an abort ends Done, not Failed"
939        );
940        assert_eq!(handle.state(), MemberState::Done);
941    }
942
943    #[tokio::test]
944    async fn a_fatal_stream_error_ends_failed() {
945        let rec = Recorder::default();
946        rec.push(vec![LlmStreamEvent::Error {
947            message: "boom".into(),
948            fatal: true,
949        }]);
950        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
951        let mut rx = handle.subscribe();
952        handle.send(Request::Submit { text: "hi".into() }).unwrap();
953        wait_for_end(&mut rx).await;
954        assert_eq!(handle.state(), MemberState::Failed);
955    }
956
957    #[derive(Deserialize, schemars::JsonSchema)]
958    struct GateArgs {
959        text: String,
960    }
961
962    /// Blocks until released, so the actor can be steered while a run is inside
963    /// a tool call.
964    struct GateTool {
965        entered: mpsc::UnboundedSender<()>,
966        release: Arc<tokio::sync::Notify>,
967    }
968
969    #[async_trait::async_trait]
970    impl TypedTool for GateTool {
971        type Args = GateArgs;
972        fn name(&self) -> &str {
973            "gate"
974        }
975        fn description(&self) -> &str {
976            "blocks until released"
977        }
978        async fn execute(&self, args: GateArgs, _ctx: &ToolContext) -> ToolOutput {
979            self.entered.send(()).ok();
980            self.release.notified().await;
981            ToolOutput {
982                output: format!("released:{}", args.text),
983                is_error: false,
984                diff: None,
985                path: None,
986            }
987        }
988    }
989
990    #[tokio::test]
991    async fn config_request_during_a_run_is_deferred_to_after_it() {
992        // Run one: turn 1 blocks in the gate, turn 2 finishes — both on the
993        // run's own model. Run two is a fresh request and must see the swap.
994        let rec = Recorder::default();
995        rec.push(vec![
996            LlmStreamEvent::ToolCall {
997                id: "c1".into(),
998                name: "gate".into(),
999                arguments: json!({ "text": "hi" }),
1000            },
1001            LlmStreamEvent::Done {
1002                stop_reason: StopReason::ToolUse,
1003                usage: None,
1004            },
1005        ]);
1006        for text in ["turn two", "run two"] {
1007            rec.push(vec![
1008                LlmStreamEvent::TextDelta(text.into()),
1009                LlmStreamEvent::Done {
1010                    stop_reason: StopReason::Stop,
1011                    usage: None,
1012                },
1013            ]);
1014        }
1015
1016        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1017        let release = Arc::new(tokio::sync::Notify::new());
1018        let gate = erased(GateTool {
1019            entered: entered_tx,
1020            release: release.clone(),
1021        });
1022
1023        let handle =
1024            SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1025        let mut rx = handle.subscribe();
1026        handle.send(Request::Submit { text: "a".into() }).unwrap();
1027
1028        entered_rx.recv().await.unwrap(); // run one is inside the tool
1029        handle
1030            .send(Request::SetModel { model: "m2".into() })
1031            .unwrap(); // must not take effect until run one is done
1032        release.notify_one();
1033        wait_for_end(&mut rx).await; // run one ends on m1
1034
1035        handle.send(Request::Submit { text: "b".into() }).unwrap();
1036        wait_for_end(&mut rx).await; // run two
1037
1038        assert_eq!(
1039            rec.models(),
1040            ["m1", "m1", "m2"],
1041            "the swap waited for the run, then applied to the next one"
1042        );
1043    }
1044
1045    #[tokio::test]
1046    async fn ask_history_returns_the_conversation() {
1047        let rec = Recorder::default();
1048        rec.push(vec![
1049            LlmStreamEvent::TextDelta("hi back".into()),
1050            LlmStreamEvent::Done {
1051                stop_reason: StopReason::Stop,
1052                usage: None,
1053            },
1054        ]);
1055        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1056        let mut rx = handle.subscribe();
1057        handle
1058            .send(Request::Submit {
1059                text: "hello".into(),
1060            })
1061            .unwrap();
1062        wait_for_end(&mut rx).await;
1063
1064        let reply = handle.ask(Request::GetHistory).await.unwrap();
1065        let AgentEvent::History { messages } = reply else {
1066            panic!("expected a History reply, got {reply:?}");
1067        };
1068        assert!(
1069            messages.iter().any(|m| m.as_text() == "hello"),
1070            "the read-back carries the user turn: {messages:?}"
1071        );
1072    }
1073
1074    #[tokio::test]
1075    async fn ask_set_model_replies_ack_and_takes_effect() {
1076        let rec = Recorder::default();
1077        rec.push(vec![
1078            LlmStreamEvent::TextDelta("x".into()),
1079            LlmStreamEvent::Done {
1080                stop_reason: StopReason::Stop,
1081                usage: None,
1082            },
1083        ]);
1084        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1085        let mut rx = handle.subscribe();
1086
1087        let reply = handle
1088            .ask(Request::SetModel { model: "m2".into() })
1089            .await
1090            .unwrap();
1091        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1092
1093        handle.send(Request::Submit { text: "go".into() }).unwrap();
1094        wait_for_end(&mut rx).await;
1095        assert_eq!(rec.models(), ["m2"], "the swap applied to the next run");
1096    }
1097
1098    #[tokio::test]
1099    async fn define_on_a_local_session_is_an_error() {
1100        let rec = Recorder::default();
1101        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1102        let reply = handle
1103            .ask(Request::Define {
1104                name: Some("w1".into()),
1105                model: None,
1106                role: None,
1107                tools: None,
1108                base_url: None,
1109                api_key: None,
1110                read_only: false,
1111                effort: None,
1112            })
1113            .await
1114            .unwrap();
1115        assert!(
1116            matches!(&reply, AgentEvent::Error { message } if message.contains("served")),
1117            "{reply:?}"
1118        );
1119    }
1120
1121    #[tokio::test]
1122    async fn ask_compact_with_nothing_to_do_replies_ack() {
1123        let handle = SessionActor::spawn(Agent::new(agent_config(
1124            fake_stream_fn(&Recorder::default()),
1125            vec![],
1126        )));
1127        let reply = handle
1128            .ask(Request::Compact { instructions: None })
1129            .await
1130            .unwrap();
1131        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1132    }
1133
1134    #[tokio::test]
1135    async fn ask_submit_replies_stopped_with_the_reason() {
1136        let rec = Recorder::default();
1137        rec.push(vec![
1138            LlmStreamEvent::TextDelta("done".into()),
1139            LlmStreamEvent::Done {
1140                stop_reason: StopReason::Stop,
1141                usage: None,
1142            },
1143        ]);
1144        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1145        let reply = handle
1146            .ask(Request::Submit { text: "hi".into() })
1147            .await
1148            .unwrap();
1149        assert!(
1150            matches!(
1151                reply,
1152                AgentEvent::Stopped {
1153                    stop_reason: StopReason::Stop
1154                }
1155            ),
1156            "{reply:?}"
1157        );
1158    }
1159
1160    #[tokio::test]
1161    async fn notify_appends_without_starting_a_turn() {
1162        let rec = Recorder::default();
1163        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1164        let mut rx = handle.subscribe();
1165
1166        let reply = handle
1167            .ask(Request::Notify {
1168                content: "ping".into(),
1169            })
1170            .await
1171            .unwrap();
1172        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1173
1174        // The message was surfaced to subscribers...
1175        let surfaced = tokio::time::timeout(Duration::from_secs(1), rx.recv())
1176            .await
1177            .expect("an event")
1178            .expect("open");
1179        assert!(
1180            matches!(&surfaced, AgentEvent::MessageReceived { content, .. } if content == "ping"),
1181            "{surfaced:?}"
1182        );
1183
1184        // ...the conversation carries it...
1185        let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1186            panic!("expected History");
1187        };
1188        assert!(
1189            messages
1190                .iter()
1191                .any(|m| m.as_text() == "[message from user]\nping"),
1192            "{messages:?}"
1193        );
1194
1195        // ...and no turn ran (the loop was never entered).
1196        assert!(rec.calls().is_empty(), "notify must not run the loop");
1197    }
1198
1199    #[tokio::test]
1200    async fn wake_starts_a_run_when_idle() {
1201        let rec = Recorder::default();
1202        rec.push(vec![
1203            LlmStreamEvent::TextDelta("acted".into()),
1204            LlmStreamEvent::Done {
1205                stop_reason: StopReason::Stop,
1206                usage: None,
1207            },
1208        ]);
1209        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1210        let mut rx = handle.subscribe();
1211
1212        let reply = handle
1213            .ask(Request::Wake {
1214                content: "do the task".into(),
1215            })
1216            .await
1217            .unwrap();
1218        assert!(
1219            matches!(
1220                reply,
1221                AgentEvent::Stopped {
1222                    stop_reason: StopReason::Stop
1223                }
1224            ),
1225            "{reply:?}"
1226        );
1227        wait_for_end(&mut rx).await;
1228
1229        let calls = rec.calls();
1230        assert_eq!(calls.len(), 1, "Wake ran the loop once");
1231        assert!(
1232            calls[0]
1233                .iter()
1234                .any(|m| m.as_text() == "[message from user]\ndo the task"),
1235            "the tagged task is the prompt: {:?}",
1236            calls[0]
1237        );
1238    }
1239
1240    struct InboundPolicy;
1241
1242    #[async_trait::async_trait]
1243    impl crate::hooks::Hooks for InboundPolicy {
1244        async fn before_inbound(
1245            &self,
1246            _from: Option<&SessionId>,
1247            request: &mut Request,
1248        ) -> Option<String> {
1249            let Request::Notify { content } = request else {
1250                return None;
1251            };
1252            if content.starts_with("drop:") {
1253                return Some("policy".into());
1254            }
1255            *content = format!("{content}!");
1256            None
1257        }
1258    }
1259
1260    #[tokio::test]
1261    async fn before_inbound_drops_and_rewrites_a_notify() {
1262        let rec = Recorder::default();
1263        let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
1264        cfg.hooks = HooksSet::one(Arc::new(InboundPolicy));
1265        let handle = SessionActor::spawn(Agent::new(cfg));
1266
1267        // Dropped: an `ask` is answered with an Error, and nothing lands.
1268        let reply = handle
1269            .ask(Request::Notify {
1270                content: "drop: me".into(),
1271            })
1272            .await
1273            .unwrap();
1274        assert!(
1275            matches!(&reply, AgentEvent::Error { message } if message.contains("policy")),
1276            "{reply:?}"
1277        );
1278
1279        // Rewritten: the hook's edit is what is recorded.
1280        let reply = handle
1281            .ask(Request::Notify {
1282                content: "keep me".into(),
1283            })
1284            .await
1285            .unwrap();
1286        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1287
1288        let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1289            panic!("expected History");
1290        };
1291        let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
1292        assert!(
1293            !texts.iter().any(|t| t == "drop: me"),
1294            "the dropped message never landed: {texts:?}"
1295        );
1296        assert!(
1297            texts.iter().any(|t| t == "[message from user]\nkeep me!"),
1298            "the rewrite landed: {texts:?}"
1299        );
1300    }
1301
1302    #[tokio::test]
1303    async fn interrupt_reaches_the_next_turn_while_a_tool_is_blocked() {
1304        let rec = Recorder::default();
1305        rec.push(vec![
1306            LlmStreamEvent::ToolCall {
1307                id: "c1".into(),
1308                name: "gate".into(),
1309                arguments: json!({ "text": "hi" }),
1310            },
1311            LlmStreamEvent::Done {
1312                stop_reason: StopReason::ToolUse,
1313                usage: None,
1314            },
1315        ]);
1316        rec.push(vec![
1317            LlmStreamEvent::TextDelta("after interrupt".into()),
1318            LlmStreamEvent::Done {
1319                stop_reason: StopReason::Stop,
1320                usage: None,
1321            },
1322        ]);
1323
1324        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1325        let release = Arc::new(tokio::sync::Notify::new());
1326        let gate = erased(GateTool {
1327            entered: entered_tx,
1328            release: release.clone(),
1329        });
1330
1331        let handle =
1332            SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1333        let mut rx = handle.subscribe();
1334        handle.send(Request::Submit { text: "hi".into() }).unwrap();
1335
1336        entered_rx.recv().await.unwrap(); // the tool is now blocked
1337        handle
1338            .send(Request::Interrupt {
1339                content: "mid-run interrupt".into(),
1340            })
1341            .unwrap();
1342        release.notify_one();
1343
1344        wait_for_end(&mut rx).await;
1345
1346        let calls = rec.calls();
1347        assert_eq!(calls.len(), 2, "two stream calls: the tool turn and the next");
1348        assert!(
1349            calls[1]
1350                .iter()
1351                .any(|m| matches!(m, AgentMessage::User { .. })
1352                    && m.as_text() == "[message from user]\nmid-run interrupt"),
1353            "the interrupt drained into the next turn: {:?}",
1354            calls[1]
1355        );
1356    }
1357
1358    #[tokio::test]
1359    async fn send_from_attributes_the_sender() {
1360        let rec = Recorder::default();
1361        let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1362        let mut rx = handle.subscribe();
1363
1364        handle
1365            .send_from(
1366                SessionId::agent("b"),
1367                Request::Notify {
1368                    content: "ping".into(),
1369                },
1370            )
1371            .unwrap();
1372
1373        let event = tokio::time::timeout(Duration::from_secs(1), rx.recv())
1374            .await
1375            .expect("an event")
1376            .expect("open");
1377        assert!(
1378            matches!(&event, AgentEvent::MessageReceived { from, content }
1379                if from.as_str() == "agent:b" && content == "ping"),
1380            "{event:?}"
1381        );
1382    }
1383
1384    struct SenderPolicy;
1385
1386    #[async_trait::async_trait]
1387    impl crate::hooks::Hooks for SenderPolicy {
1388        async fn before_inbound(
1389            &self,
1390            from: Option<&SessionId>,
1391            _request: &mut Request,
1392        ) -> Option<String> {
1393            (from.map(SessionId::as_str) == Some("agent:evil")).then(|| "blocked sender".into())
1394        }
1395    }
1396
1397    #[tokio::test]
1398    async fn before_inbound_sees_the_sender() {
1399        let rec = Recorder::default();
1400        let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
1401        cfg.hooks = HooksSet::one(Arc::new(SenderPolicy));
1402        let handle = SessionActor::spawn(Agent::new(cfg));
1403
1404        // A hostile peer is dropped...
1405        let reply = handle
1406            .ask_from(
1407                SessionId::agent("evil"),
1408                Request::Notify {
1409                    content: "let me in".into(),
1410                },
1411            )
1412            .await
1413            .unwrap();
1414        assert!(
1415            matches!(&reply, AgentEvent::Error { message } if message.contains("blocked sender")),
1416            "{reply:?}"
1417        );
1418
1419        // ...a friendly one is not.
1420        let reply = handle
1421            .ask_from(
1422                SessionId::agent("friend"),
1423                Request::Notify {
1424                    content: "hello".into(),
1425                },
1426            )
1427            .await
1428            .unwrap();
1429        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1430
1431        let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1432            panic!("expected History");
1433        };
1434        let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
1435        assert!(
1436            texts
1437                .iter()
1438                .any(|t| t == "[message from agent:friend]\nhello"),
1439            "{texts:?}"
1440        );
1441        assert!(!texts.iter().any(|t| t.contains("let me in")), "{texts:?}");
1442    }
1443
1444    #[tokio::test]
1445    async fn side_ask_is_answered_mid_run() {
1446        // Run one blocks in the gate (turn 1), so the actor is mid-run when the
1447        // side ask arrives; it must be answered without waiting for the run.
1448        let rec = Recorder::default();
1449        rec.push(vec![
1450            LlmStreamEvent::ToolCall {
1451                id: "c1".into(),
1452                name: "gate".into(),
1453                arguments: json!({ "text": "hi" }),
1454            },
1455            LlmStreamEvent::Done {
1456                stop_reason: StopReason::ToolUse,
1457                usage: None,
1458            },
1459        ]);
1460        // The side ask's own model call (empty tool list) pops this.
1461        rec.push(vec![
1462            LlmStreamEvent::TextDelta("aside".into()),
1463            LlmStreamEvent::Done {
1464                stop_reason: StopReason::Stop,
1465                usage: None,
1466            },
1467        ]);
1468        // Run one, turn 2, after the gate releases.
1469        rec.push(vec![
1470            LlmStreamEvent::TextDelta("done".into()),
1471            LlmStreamEvent::Done {
1472                stop_reason: StopReason::Stop,
1473                usage: None,
1474            },
1475        ]);
1476
1477        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1478        let release = Arc::new(tokio::sync::Notify::new());
1479        let gate = erased(GateTool {
1480            entered: entered_tx,
1481            release: release.clone(),
1482        });
1483        let handle =
1484            SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1485        let mut rx = handle.subscribe();
1486        handle.send(Request::Submit { text: "a".into() }).unwrap();
1487
1488        entered_rx.recv().await.unwrap(); // run one is inside the tool, still in flight
1489
1490        // Answered now, not after the run: `ask` returns before we release.
1491        let reply = tokio::time::timeout(
1492            Duration::from_secs(5),
1493            handle.ask(Request::SideAsk { text: "why?".into() }),
1494        )
1495        .await
1496        .expect("the side ask must not wait for the run")
1497        .unwrap();
1498        assert!(
1499            matches!(&reply, AgentEvent::SideAnswer { text, .. } if text == "aside"),
1500            "{reply:?}"
1501        );
1502
1503        release.notify_one();
1504        wait_for_end(&mut rx).await;
1505    }
1506
1507    #[tokio::test]
1508    async fn get_history_is_answered_mid_run_from_the_snapshot() {
1509        let rec = Recorder::default();
1510        rec.push(vec![
1511            LlmStreamEvent::ToolCall {
1512                id: "c1".into(),
1513                name: "gate".into(),
1514                arguments: json!({ "text": "hi" }),
1515            },
1516            LlmStreamEvent::Done {
1517                stop_reason: StopReason::ToolUse,
1518                usage: None,
1519            },
1520        ]);
1521        rec.push(vec![
1522            LlmStreamEvent::TextDelta("done".into()),
1523            LlmStreamEvent::Done {
1524                stop_reason: StopReason::Stop,
1525                usage: None,
1526            },
1527        ]);
1528
1529        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1530        let release = Arc::new(tokio::sync::Notify::new());
1531        let gate = erased(GateTool {
1532            entered: entered_tx,
1533            release: release.clone(),
1534        });
1535        let handle =
1536            SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1537        let mut rx = handle.subscribe();
1538        handle.send(Request::Submit { text: "a".into() }).unwrap();
1539
1540        entered_rx.recv().await.unwrap();
1541
1542        let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::GetHistory))
1543            .await
1544            .expect("history must not wait for the run")
1545            .unwrap();
1546        let AgentEvent::History { messages } = reply else {
1547            panic!("expected History, got {reply:?}");
1548        };
1549        assert!(
1550            messages.iter().any(|m| m.as_text() == "a"),
1551            "the mid-run snapshot carries the submit: {messages:?}"
1552        );
1553
1554        release.notify_one();
1555        wait_for_end(&mut rx).await;
1556    }
1557
1558    #[tokio::test]
1559    async fn a_deferred_command_is_acked_before_the_run_ends() {
1560        let rec = Recorder::default();
1561        rec.push(vec![
1562            LlmStreamEvent::ToolCall {
1563                id: "c1".into(),
1564                name: "gate".into(),
1565                arguments: json!({ "text": "hi" }),
1566            },
1567            LlmStreamEvent::Done {
1568                stop_reason: StopReason::ToolUse,
1569                usage: None,
1570            },
1571        ]);
1572        for text in ["turn two", "run two"] {
1573            rec.push(vec![
1574                LlmStreamEvent::TextDelta(text.into()),
1575                LlmStreamEvent::Done {
1576                    stop_reason: StopReason::Stop,
1577                    usage: None,
1578                },
1579            ]);
1580        }
1581
1582        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1583        let release = Arc::new(tokio::sync::Notify::new());
1584        let gate = erased(GateTool {
1585            entered: entered_tx,
1586            release: release.clone(),
1587        });
1588        let handle =
1589            SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1590        let mut rx = handle.subscribe();
1591        handle.send(Request::Submit { text: "a".into() }).unwrap();
1592
1593        entered_rx.recv().await.unwrap();
1594
1595        // Acked now (mid-run), not left to time out; applied at the run boundary.
1596        let reply = tokio::time::timeout(
1597            Duration::from_secs(5),
1598            handle.ask(Request::SetModel { model: "m2".into() }),
1599        )
1600        .await
1601        .expect("a deferred command must be acked, not hang")
1602        .unwrap();
1603        assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1604
1605        release.notify_one();
1606        wait_for_end(&mut rx).await; // run one ends on m1
1607
1608        handle.send(Request::Submit { text: "b".into() }).unwrap();
1609        wait_for_end(&mut rx).await; // run two
1610        assert_eq!(
1611            rec.models(),
1612            ["m1", "m1", "m2"],
1613            "the swap applied after the run, not before"
1614        );
1615    }
1616
1617    #[tokio::test]
1618    async fn usage_reply_is_answered_mid_run_without_waiting() {
1619        // Cloned from `ask_status_is_answered_mid_run_from_the_snapshot`: same
1620        // GateTool + release choreography, one assertion changed — the mid-run
1621        // ask answers at once with folded, PRICED stats.
1622        let rec = Recorder::default();
1623        rec.push(vec![
1624            LlmStreamEvent::ToolCall {
1625                id: "c1".into(),
1626                name: "gate".into(),
1627                arguments: json!({ "text": "hi" }),
1628            },
1629            LlmStreamEvent::Done {
1630                stop_reason: StopReason::ToolUse,
1631                usage: None,
1632            },
1633        ]);
1634        rec.push(vec![
1635            LlmStreamEvent::TextDelta("done".into()),
1636            LlmStreamEvent::Done {
1637                stop_reason: StopReason::Stop,
1638                usage: None,
1639            },
1640        ]);
1641
1642        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1643        let release = Arc::new(tokio::sync::Notify::new());
1644        let gate = erased(GateTool {
1645            entered: entered_tx,
1646            release: release.clone(),
1647        });
1648        // Seed a priced assistant (zen base + deepseek-v4.1-flash) so the mid-run
1649        // reply folds ONE turn and reports a cost.
1650        let mut cfg = agent_config(fake_stream_fn(&rec), vec![gate]);
1651        cfg.llm.base_url = Some("https://opencode.ai/zen/go/v1".into());
1652        cfg.llm.model = "deepseek-v4.1-flash".into();
1653        cfg.context = vec![AgentMessage::Assistant {
1654            content: vec![crate::message::ContentBlock::Text {
1655                text: "seeded work".into(),
1656            }],
1657            stop_reason: StopReason::Stop,
1658            usage: Some(crate::message::Usage {
1659                input_tokens: 1_000,
1660                output_tokens: 500,
1661                cache_read_tokens: Some(400),
1662                cache_write_tokens: None,
1663            }),
1664            model: Some("deepseek-v4.1-flash".into()),
1665        }];
1666        let handle = SessionActor::spawn(Agent::new(cfg));
1667        let mut rx = handle.subscribe();
1668        handle.send(Request::Submit { text: "a".into() }).unwrap();
1669
1670        entered_rx.recv().await.unwrap(); // run is inside the tool
1671
1672        let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::Usage))
1673            .await
1674            .expect("usage must not wait for the run")
1675            .unwrap();
1676        let AgentEvent::Usage { stats } = reply else {
1677            panic!("expected Usage, got {reply:?}");
1678        };
1679        assert_eq!(stats.turns, 1);
1680        assert!(stats.cost.is_some(), "a priced model reports a cost mid-run");
1681
1682        release.notify_one();
1683        wait_for_end(&mut rx).await;
1684    }
1685
1686    #[tokio::test]
1687    async fn an_idle_usage_reply_folds_the_live_transcript() {
1688        // The idle arm reads the FRESH `Agent::messages()`, not the snapshot.
1689        let rec = Recorder::default();
1690        let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
1691        cfg.llm.base_url = None; // unpriced: a valid report with cost: None
1692        cfg.context = vec![AgentMessage::Assistant {
1693            content: vec![crate::message::ContentBlock::Text {
1694                text: "done".into(),
1695            }],
1696            stop_reason: StopReason::Stop,
1697            usage: Some(crate::message::Usage {
1698                input_tokens: 10,
1699                output_tokens: 20,
1700                cache_read_tokens: None,
1701                cache_write_tokens: None,
1702            }),
1703            model: None,
1704        }];
1705        let handle = SessionActor::spawn(Agent::new(cfg));
1706
1707        let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::Usage))
1708            .await
1709            .expect("an idle usage read answers")
1710            .unwrap();
1711        let AgentEvent::Usage { stats } = reply else {
1712            panic!("expected Usage, got {reply:?}");
1713        };
1714        assert_eq!(stats.turns, 1);
1715        assert_eq!(stats.input_tokens, 10);
1716        assert_eq!(stats.cost, None, "no base_url ⇒ no price, never $0");
1717    }
1718
1719    #[tokio::test]
1720    async fn ask_status_is_answered_mid_run_from_the_snapshot() {
1721        // A lean `Status` (no transcript): the member's last assistant text,
1722        // answered while a run is in flight from the turn-boundary snapshot.
1723        let rec = Recorder::default();
1724        rec.push(vec![
1725            LlmStreamEvent::ToolCall {
1726                id: "c1".into(),
1727                name: "gate".into(),
1728                arguments: json!({ "text": "hi" }),
1729            },
1730            LlmStreamEvent::Done {
1731                stop_reason: StopReason::ToolUse,
1732                usage: None,
1733            },
1734        ]);
1735        rec.push(vec![
1736            LlmStreamEvent::TextDelta("done".into()),
1737            LlmStreamEvent::Done {
1738                stop_reason: StopReason::Stop,
1739                usage: None,
1740            },
1741        ]);
1742
1743        let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1744        let release = Arc::new(tokio::sync::Notify::new());
1745        let gate = erased(GateTool {
1746            entered: entered_tx,
1747            release: release.clone(),
1748        });
1749        // Seed an assistant so the mid-run snapshot has a "last work".
1750        let mut cfg = agent_config(fake_stream_fn(&rec), vec![gate]);
1751        cfg.context = vec![AgentMessage::Assistant {
1752            content: vec![crate::message::ContentBlock::Text {
1753                text: "seeded work".into(),
1754            }],
1755            stop_reason: StopReason::Stop,
1756            usage: None,
1757            model: None,
1758        }];
1759        let handle = SessionActor::spawn(Agent::new(cfg));
1760        let mut rx = handle.subscribe();
1761        handle.send(Request::Submit { text: "a".into() }).unwrap();
1762
1763        entered_rx.recv().await.unwrap(); // run is inside the tool
1764
1765        let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::Status))
1766            .await
1767            .expect("status must not wait for the run")
1768            .unwrap();
1769        let AgentEvent::Status { last_assistant_text } = reply else {
1770            panic!("expected Status, got {reply:?}");
1771        };
1772        assert_eq!(last_assistant_text.as_deref(), Some("seeded work"));
1773
1774        release.notify_one();
1775        wait_for_end(&mut rx).await;
1776    }
1777}