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