use std::collections::VecDeque;
use std::fmt;
use tokio::sync::{broadcast, mpsc, oneshot, watch};
use crate::agent::Agent;
use crate::compaction::CompactOutcome;
use crate::event::AgentEvent;
use crate::hooks::HooksSet;
use crate::loop_::LoopError;
use crate::message::{AgentMessage, StopReason};
use crate::protocol::{MemberState, Request, SessionId};
pub const EVENT_BUFFER: usize = 1024;
#[derive(Clone)]
pub struct SessionHandle {
inbox: mpsc::UnboundedSender<Message>,
events: broadcast::Sender<AgentEvent>,
state: watch::Sender<MemberState>,
}
impl SessionHandle {
pub fn send(&self, request: Request) -> Result<(), SessionClosed> {
self.tell(None, request)
}
pub fn send_from(&self, from: SessionId, request: Request) -> Result<(), SessionClosed> {
self.tell(Some(from), request)
}
fn tell(&self, from: Option<SessionId>, request: Request) -> Result<(), SessionClosed> {
self.inbox
.send(Message::Tell { from, request })
.map_err(|_| SessionClosed)
}
pub async fn ask(&self, request: Request) -> Result<AgentEvent, SessionClosed> {
self.question(None, request).await
}
pub async fn ask_from(
&self,
from: SessionId,
request: Request,
) -> Result<AgentEvent, SessionClosed> {
self.question(Some(from), request).await
}
async fn question(
&self,
from: Option<SessionId>,
request: Request,
) -> Result<AgentEvent, SessionClosed> {
let (reply_tx, reply_rx) = oneshot::channel();
self.inbox
.send(Message::Ask {
from,
request,
reply: reply_tx,
})
.map_err(|_| SessionClosed)?;
reply_rx.await.map_err(|_| SessionClosed)
}
pub fn state(&self) -> MemberState {
*self.state.borrow()
}
pub fn state_rx(&self) -> watch::Receiver<MemberState> {
self.state.subscribe()
}
pub fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
self.events.subscribe()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SessionClosed;
impl fmt::Display for SessionClosed {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("session actor has shut down")
}
}
impl std::error::Error for SessionClosed {}
enum Message {
Tell {
from: Option<SessionId>,
request: Request,
},
Ask {
from: Option<SessionId>,
request: Request,
reply: oneshot::Sender<AgentEvent>,
},
}
pub struct SessionActor;
impl SessionActor {
pub fn spawn(agent: Agent) -> SessionHandle {
let (inbox_tx, inbox_rx) = mpsc::unbounded_channel();
let (events_tx, _) = broadcast::channel(EVENT_BUFFER);
let events = events_tx.clone();
let (state_tx, _state_rx) = watch::channel(MemberState::Idle);
tokio::spawn(serve(agent, inbox_rx, events_tx, state_tx.clone()));
SessionHandle {
inbox: inbox_tx,
events,
state: state_tx,
}
}
}
async fn serve(
mut agent: Agent,
mut inbox: mpsc::UnboundedReceiver<Message>,
events: broadcast::Sender<AgentEvent>,
state: watch::Sender<MemberState>,
) {
let mut deferred: VecDeque<Message> = VecDeque::new();
loop {
let message = match deferred.pop_front() {
Some(message) => message,
None => match inbox.recv().await {
Some(message) => message,
None => break, },
};
dispatch(&mut agent, message, &mut inbox, &events, &mut deferred, &state).await;
}
}
fn tag(sender: &SessionId, content: &str) -> String {
format!("[message from {sender}]\n{content}")
}
async fn gate_inbound(hooks: &HooksSet, message: &mut Message) -> Option<String> {
let (from, request) = match message {
Message::Tell { from, request } | Message::Ask { from, request, .. } => (from, request),
};
if !crate::protocol::is_inbound(request) {
return None;
}
hooks.before_inbound(from.as_ref(), request).await
}
async fn dispatch(
agent: &mut Agent,
mut message: Message,
inbox: &mut mpsc::UnboundedReceiver<Message>,
events: &broadcast::Sender<AgentEvent>,
deferred: &mut VecDeque<Message>,
state: &watch::Sender<MemberState>,
) {
if let Some(reason) = gate_inbound(agent.hooks(), &mut message).await {
if let Message::Ask { reply, .. } = message {
let _ = reply.send(AgentEvent::Error {
message: format!("inbound blocked: {reason}"),
});
}
return;
}
let (from, request, reply) = match message {
Message::Tell { from, request } => (from, request, None),
Message::Ask { from, request, reply } => (from, request, Some(reply)),
};
let sender = from.unwrap_or_else(SessionId::user);
let event = match request {
Request::Submit { text } => match run(agent, &text, inbox, events, deferred, state).await {
Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
},
Request::Notify { content } => match agent.notify(tag(&sender, &content)) {
Ok(_) => {
let _ = events.send(AgentEvent::MessageReceived {
from: sender.clone(),
content,
});
AgentEvent::Ack
}
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
},
Request::Interrupt { content } => {
agent.steer(AgentMessage::user_text(tag(&sender, &content)));
let _ = events.send(AgentEvent::MessageReceived {
from: sender.clone(),
content,
});
AgentEvent::Ack
}
Request::Wake { content } => {
let _ = events.send(AgentEvent::MessageReceived {
from: sender.clone(),
content: content.clone(),
});
match run(agent, &tag(&sender, &content), inbox, events, deferred, state).await {
Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
}
}
Request::Cancel => AgentEvent::Ack,
Request::SetModel { model } => match agent.set_model(model) {
Ok(()) => AgentEvent::Ack,
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
},
Request::SetEffort { effort } => match agent.set_effort(effort) {
Ok(()) => AgentEvent::Ack,
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
},
Request::Compact { instructions } => match agent.compact(instructions.as_deref()).await {
Ok(CompactOutcome::NothingToDo) => AgentEvent::Ack,
Ok(CompactOutcome::Done {
summarized, kept, ..
}) => AgentEvent::Compaction { summarized, kept },
Err(e) => AgentEvent::Error { message: e },
},
Request::SideAsk { text } => match agent.side_ask(&text).await {
Ok(answer) => AgentEvent::SideAnswer {
text: answer.text,
usage: answer.usage,
},
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
},
Request::SetPlanMode { on } => {
agent.set_plan_mode(on);
AgentEvent::Ack
}
Request::GetHistory => AgentEvent::History {
messages: agent.messages().to_vec(),
},
Request::Status => AgentEvent::Status {
last_assistant_text: crate::message::last_assistant_text(agent.messages()),
},
Request::Define { .. } => AgentEvent::Error {
message: "agent definition is served, not in-process".into(),
},
Request::ListSessions | Request::ListModels | Request::Unknown => AgentEvent::Ack,
};
if let Some(reply) = reply {
let _ = reply.send(event);
}
}
async fn run(
agent: &mut Agent,
text: &str,
inbox: &mut mpsc::UnboundedReceiver<Message>,
events: &broadcast::Sender<AgentEvent>,
deferred: &mut VecDeque<Message>,
state: &watch::Sender<MemberState>,
) -> Result<StopReason, LoopError> {
state.send_replace(MemberState::Running);
let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::<AgentEvent>();
let steer = agent.steer_sender();
let follow_up = agent.follow_up_sender();
let cancel = agent.cancel_token();
let hooks = agent.hooks().clone();
let side = agent.side_handle();
let mut run = Box::pin(agent.run(text, sink_tx));
let result = loop {
tokio::select! {
biased;
Some(mut message) = inbox.recv() => {
if let Some(reason) = gate_inbound(&hooks, &mut message).await {
if let Message::Ask { reply, .. } = message {
let _ = reply.send(AgentEvent::Error {
message: format!("inbound blocked: {reason}"),
});
}
continue;
}
match message {
Message::Tell {
request: Request::Cancel,
..
} => cancel.cancel(),
Message::Tell {
request: Request::Notify { content } | Request::Interrupt { content },
from,
} => {
let from = from.unwrap_or_else(SessionId::user);
let _ = steer.send(AgentMessage::user_text(tag(&from, &content)));
let _ = events.send(AgentEvent::MessageReceived {
from,
content,
});
}
Message::Tell {
request: Request::Wake { content },
from,
} => {
let from = from.unwrap_or_else(SessionId::user);
let _ = follow_up.send(AgentMessage::user_text(tag(&from, &content)));
let _ = events.send(AgentEvent::MessageReceived {
from,
content,
});
}
Message::Ask {
request: Request::GetHistory,
reply,
..
} => {
let _ = reply.send(AgentEvent::History { messages: side.history() });
}
Message::Ask {
request: Request::Status,
reply,
..
} => {
let _ = reply.send(AgentEvent::Status {
last_assistant_text: side.last_assistant_text(),
});
}
Message::Ask {
request: Request::SideAsk { text },
reply,
..
} => {
let side = side.clone();
tokio::spawn(async move {
let event = match side.side_ask(&text).await {
Ok(answer) => AgentEvent::SideAnswer {
text: answer.text,
usage: answer.usage,
},
Err(e) => AgentEvent::Error {
message: e.to_string(),
},
};
let _ = reply.send(event);
});
}
Message::Ask {
from,
request:
req @ (Request::SetModel { .. }
| Request::SetEffort { .. }
| Request::SetPlanMode { .. }),
reply,
} => {
let _ = reply.send(AgentEvent::Ack);
deferred.push_back(Message::Tell { from, request: req });
}
other => deferred.push_back(other),
}
},
result = &mut run => break result,
Some(event) = sink_rx.recv() => {
let _ = events.send(event);
}
}
};
while let Ok(event) = sink_rx.try_recv() {
let _ = events.send(event);
}
if let Err(e) = &result {
let _ = events.send(AgentEvent::Error {
message: e.to_string(),
});
}
match &result {
Ok(StopReason::Error) | Err(_) => {
state.send_replace(MemberState::Failed);
}
Ok(_) => {
state.send_replace(MemberState::Done);
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::{Agent, AgentConfig};
use crate::compaction::CompactionPolicy;
use crate::event::LlmStreamEvent;
use crate::hooks::HooksSet;
use crate::message::StopReason;
use crate::streamfn::{LlmOpts, LlmStream, StreamFn};
use crate::tool::{ToolContext, ToolOutput, TypedTool, erased};
use futures::StreamExt as _;
use serde::Deserialize;
use serde_json::json;
use std::sync::{Arc, Mutex};
use std::time::Duration;
#[derive(Clone, Default)]
struct Recorder {
script: Arc<Mutex<VecDeque<Vec<LlmStreamEvent>>>>,
calls: Arc<Mutex<Vec<Vec<AgentMessage>>>>,
models: Arc<Mutex<Vec<String>>>,
}
impl Recorder {
fn push(&self, events: Vec<LlmStreamEvent>) {
self.script.lock().unwrap().push_back(events);
}
fn calls(&self) -> Vec<Vec<AgentMessage>> {
self.calls.lock().unwrap().clone()
}
fn models(&self) -> Vec<String> {
self.models.lock().unwrap().clone()
}
}
fn fake_stream_fn(rec: &Recorder) -> StreamFn {
let rec = rec.clone();
Arc::new(
move |ctx: &[AgentMessage], _system, _tools, opts: &LlmOpts| {
rec.calls.lock().unwrap().push(ctx.to_vec());
rec.models.lock().unwrap().push(opts.model.clone());
let events = rec.script.lock().unwrap().pop_front().unwrap_or_default();
Box::pin(futures::stream::iter(events)) as LlmStream
},
)
}
fn agent_config(stream_fn: StreamFn, tools: Vec<crate::tool::Tool>) -> AgentConfig {
AgentConfig {
system: "sys".into(),
tools,
llm: LlmOpts {
model: "m1".into(),
..LlmOpts::default()
},
stream_fn,
hooks: HooksSet::default(),
session: None,
context: Vec::new(),
working_dir: std::path::PathBuf::new(),
max_turns: crate::loop_::DEFAULT_MAX_TURNS,
parallel_tools: true,
compaction: CompactionPolicy::default(),
plan_mode: crate::hooks::PlanModeHandle::new(),
}
}
async fn wait_for_end(rx: &mut broadcast::Receiver<AgentEvent>) {
loop {
let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("an event within the timeout")
.expect("the outbox stays open");
if matches!(event, AgentEvent::AgentEnd) {
return;
}
}
}
#[tokio::test]
async fn submit_streams_events_to_a_subscriber() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("hello".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "hi".into() }).unwrap();
let mut saw_text = false;
loop {
let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("an event")
.expect("open");
match event {
AgentEvent::MessageUpdate { message } if message.as_text() == "hello" => {
saw_text = true;
}
AgentEvent::AgentEnd => break,
_ => {}
}
}
assert!(saw_text, "the streamed text reached the subscriber");
}
#[tokio::test]
async fn idle_cancel_does_not_poison_the_next_run() {
let rec = Recorder::default();
for text in ["one", "two"] {
rec.push(vec![
LlmStreamEvent::TextDelta(text.into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
wait_for_end(&mut rx).await;
handle.send(Request::Cancel).unwrap();
handle.send(Request::Submit { text: "b".into() }).unwrap();
let mut saw_two = false;
loop {
let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("an event")
.expect("open");
match event {
AgentEvent::MessageUpdate { message } if message.as_text() == "two" => {
saw_two = true;
}
AgentEvent::AgentEnd => break,
_ => {}
}
}
assert!(saw_two, "the second run completed despite the idle cancel");
}
#[tokio::test]
async fn cancel_aborts_an_in_flight_run() {
let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
Box::pin(head.chain(futures::stream::pending())) as LlmStream
});
let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "hi".into() }).unwrap();
loop {
let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("an event")
.expect("open");
if matches!(event, AgentEvent::MessageUpdate { .. }) {
handle.send(Request::Cancel).unwrap();
break;
}
}
wait_for_end(&mut rx).await;
}
#[tokio::test]
async fn state_is_running_during_a_turn_then_done_after() {
let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
Box::pin(head.chain(futures::stream::pending())) as LlmStream
});
let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
assert_eq!(handle.state(), MemberState::Idle, "a virgin session is idle");
let mut state_rx = handle.state_rx();
handle.send(Request::Submit { text: "hi".into() }).unwrap();
tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
.await
.expect("a state change")
.expect("open");
assert_eq!(*state_rx.borrow_and_update(), MemberState::Running);
assert_eq!(handle.state(), MemberState::Running, "the run is in flight");
handle.send(Request::Cancel).unwrap();
tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
.await
.expect("a state change")
.expect("open");
assert_eq!(
*state_rx.borrow_and_update(),
MemberState::Done,
"an abort ends Done, not Failed"
);
assert_eq!(handle.state(), MemberState::Done);
}
#[tokio::test]
async fn a_fatal_stream_error_ends_failed() {
let rec = Recorder::default();
rec.push(vec![LlmStreamEvent::Error {
message: "boom".into(),
fatal: true,
}]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "hi".into() }).unwrap();
wait_for_end(&mut rx).await;
assert_eq!(handle.state(), MemberState::Failed);
}
#[derive(Deserialize, schemars::JsonSchema)]
struct GateArgs {
text: String,
}
struct GateTool {
entered: mpsc::UnboundedSender<()>,
release: Arc<tokio::sync::Notify>,
}
#[async_trait::async_trait]
impl TypedTool for GateTool {
type Args = GateArgs;
fn name(&self) -> &str {
"gate"
}
fn description(&self) -> &str {
"blocks until released"
}
async fn execute(&self, args: GateArgs, _ctx: &ToolContext) -> ToolOutput {
self.entered.send(()).ok();
self.release.notified().await;
ToolOutput {
output: format!("released:{}", args.text),
is_error: false,
diff: None,
path: None,
}
}
}
#[tokio::test]
async fn config_request_during_a_run_is_deferred_to_after_it() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
for text in ["turn two", "run two"] {
rec.push(vec![
LlmStreamEvent::TextDelta(text.into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let handle =
SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
entered_rx.recv().await.unwrap(); handle
.send(Request::SetModel { model: "m2".into() })
.unwrap(); release.notify_one();
wait_for_end(&mut rx).await;
handle.send(Request::Submit { text: "b".into() }).unwrap();
wait_for_end(&mut rx).await;
assert_eq!(
rec.models(),
["m1", "m1", "m2"],
"the swap waited for the run, then applied to the next one"
);
}
#[tokio::test]
async fn ask_history_returns_the_conversation() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("hi back".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
handle
.send(Request::Submit {
text: "hello".into(),
})
.unwrap();
wait_for_end(&mut rx).await;
let reply = handle.ask(Request::GetHistory).await.unwrap();
let AgentEvent::History { messages } = reply else {
panic!("expected a History reply, got {reply:?}");
};
assert!(
messages.iter().any(|m| m.as_text() == "hello"),
"the read-back carries the user turn: {messages:?}"
);
}
#[tokio::test]
async fn ask_set_model_replies_ack_and_takes_effect() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("x".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
let reply = handle
.ask(Request::SetModel { model: "m2".into() })
.await
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
handle.send(Request::Submit { text: "go".into() }).unwrap();
wait_for_end(&mut rx).await;
assert_eq!(rec.models(), ["m2"], "the swap applied to the next run");
}
#[tokio::test]
async fn define_on_a_local_session_is_an_error() {
let rec = Recorder::default();
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let reply = handle
.ask(Request::Define {
name: Some("w1".into()),
model: None,
role: None,
tools: None,
base_url: None,
api_key: None,
read_only: false,
effort: None,
})
.await
.unwrap();
assert!(
matches!(&reply, AgentEvent::Error { message } if message.contains("served")),
"{reply:?}"
);
}
#[tokio::test]
async fn ask_compact_with_nothing_to_do_replies_ack() {
let handle = SessionActor::spawn(Agent::new(agent_config(
fake_stream_fn(&Recorder::default()),
vec![],
)));
let reply = handle
.ask(Request::Compact { instructions: None })
.await
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
}
#[tokio::test]
async fn ask_submit_replies_stopped_with_the_reason() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("done".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let reply = handle
.ask(Request::Submit { text: "hi".into() })
.await
.unwrap();
assert!(
matches!(
reply,
AgentEvent::Stopped {
stop_reason: StopReason::Stop
}
),
"{reply:?}"
);
}
#[tokio::test]
async fn notify_appends_without_starting_a_turn() {
let rec = Recorder::default();
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
let reply = handle
.ask(Request::Notify {
content: "ping".into(),
})
.await
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
let surfaced = tokio::time::timeout(Duration::from_secs(1), rx.recv())
.await
.expect("an event")
.expect("open");
assert!(
matches!(&surfaced, AgentEvent::MessageReceived { content, .. } if content == "ping"),
"{surfaced:?}"
);
let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
panic!("expected History");
};
assert!(
messages
.iter()
.any(|m| m.as_text() == "[message from user]\nping"),
"{messages:?}"
);
assert!(rec.calls().is_empty(), "notify must not run the loop");
}
#[tokio::test]
async fn wake_starts_a_run_when_idle() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("acted".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
let reply = handle
.ask(Request::Wake {
content: "do the task".into(),
})
.await
.unwrap();
assert!(
matches!(
reply,
AgentEvent::Stopped {
stop_reason: StopReason::Stop
}
),
"{reply:?}"
);
wait_for_end(&mut rx).await;
let calls = rec.calls();
assert_eq!(calls.len(), 1, "Wake ran the loop once");
assert!(
calls[0]
.iter()
.any(|m| m.as_text() == "[message from user]\ndo the task"),
"the tagged task is the prompt: {:?}",
calls[0]
);
}
struct InboundPolicy;
#[async_trait::async_trait]
impl crate::hooks::Hooks for InboundPolicy {
async fn before_inbound(
&self,
_from: Option<&SessionId>,
request: &mut Request,
) -> Option<String> {
let Request::Notify { content } = request else {
return None;
};
if content.starts_with("drop:") {
return Some("policy".into());
}
*content = format!("{content}!");
None
}
}
#[tokio::test]
async fn before_inbound_drops_and_rewrites_a_notify() {
let rec = Recorder::default();
let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
cfg.hooks = HooksSet::one(Arc::new(InboundPolicy));
let handle = SessionActor::spawn(Agent::new(cfg));
let reply = handle
.ask(Request::Notify {
content: "drop: me".into(),
})
.await
.unwrap();
assert!(
matches!(&reply, AgentEvent::Error { message } if message.contains("policy")),
"{reply:?}"
);
let reply = handle
.ask(Request::Notify {
content: "keep me".into(),
})
.await
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
panic!("expected History");
};
let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
assert!(
!texts.iter().any(|t| t == "drop: me"),
"the dropped message never landed: {texts:?}"
);
assert!(
texts.iter().any(|t| t == "[message from user]\nkeep me!"),
"the rewrite landed: {texts:?}"
);
}
#[tokio::test]
async fn interrupt_reaches_the_next_turn_while_a_tool_is_blocked() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("after interrupt".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let handle =
SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "hi".into() }).unwrap();
entered_rx.recv().await.unwrap(); handle
.send(Request::Interrupt {
content: "mid-run interrupt".into(),
})
.unwrap();
release.notify_one();
wait_for_end(&mut rx).await;
let calls = rec.calls();
assert_eq!(calls.len(), 2, "two stream calls: the tool turn and the next");
assert!(
calls[1]
.iter()
.any(|m| matches!(m, AgentMessage::User { .. })
&& m.as_text() == "[message from user]\nmid-run interrupt"),
"the interrupt drained into the next turn: {:?}",
calls[1]
);
}
#[tokio::test]
async fn send_from_attributes_the_sender() {
let rec = Recorder::default();
let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
let mut rx = handle.subscribe();
handle
.send_from(
SessionId::agent("b"),
Request::Notify {
content: "ping".into(),
},
)
.unwrap();
let event = tokio::time::timeout(Duration::from_secs(1), rx.recv())
.await
.expect("an event")
.expect("open");
assert!(
matches!(&event, AgentEvent::MessageReceived { from, content }
if from.as_str() == "agent:b" && content == "ping"),
"{event:?}"
);
}
struct SenderPolicy;
#[async_trait::async_trait]
impl crate::hooks::Hooks for SenderPolicy {
async fn before_inbound(
&self,
from: Option<&SessionId>,
_request: &mut Request,
) -> Option<String> {
(from.map(SessionId::as_str) == Some("agent:evil")).then(|| "blocked sender".into())
}
}
#[tokio::test]
async fn before_inbound_sees_the_sender() {
let rec = Recorder::default();
let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
cfg.hooks = HooksSet::one(Arc::new(SenderPolicy));
let handle = SessionActor::spawn(Agent::new(cfg));
let reply = handle
.ask_from(
SessionId::agent("evil"),
Request::Notify {
content: "let me in".into(),
},
)
.await
.unwrap();
assert!(
matches!(&reply, AgentEvent::Error { message } if message.contains("blocked sender")),
"{reply:?}"
);
let reply = handle
.ask_from(
SessionId::agent("friend"),
Request::Notify {
content: "hello".into(),
},
)
.await
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
panic!("expected History");
};
let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
assert!(
texts
.iter()
.any(|t| t == "[message from agent:friend]\nhello"),
"{texts:?}"
);
assert!(!texts.iter().any(|t| t.contains("let me in")), "{texts:?}");
}
#[tokio::test]
async fn side_ask_is_answered_mid_run() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("aside".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("done".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let handle =
SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
entered_rx.recv().await.unwrap();
let reply = tokio::time::timeout(
Duration::from_secs(5),
handle.ask(Request::SideAsk { text: "why?".into() }),
)
.await
.expect("the side ask must not wait for the run")
.unwrap();
assert!(
matches!(&reply, AgentEvent::SideAnswer { text, .. } if text == "aside"),
"{reply:?}"
);
release.notify_one();
wait_for_end(&mut rx).await;
}
#[tokio::test]
async fn get_history_is_answered_mid_run_from_the_snapshot() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("done".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let handle =
SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
entered_rx.recv().await.unwrap();
let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::GetHistory))
.await
.expect("history must not wait for the run")
.unwrap();
let AgentEvent::History { messages } = reply else {
panic!("expected History, got {reply:?}");
};
assert!(
messages.iter().any(|m| m.as_text() == "a"),
"the mid-run snapshot carries the submit: {messages:?}"
);
release.notify_one();
wait_for_end(&mut rx).await;
}
#[tokio::test]
async fn a_deferred_command_is_acked_before_the_run_ends() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
for text in ["turn two", "run two"] {
rec.push(vec![
LlmStreamEvent::TextDelta(text.into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let handle =
SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
entered_rx.recv().await.unwrap();
let reply = tokio::time::timeout(
Duration::from_secs(5),
handle.ask(Request::SetModel { model: "m2".into() }),
)
.await
.expect("a deferred command must be acked, not hang")
.unwrap();
assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
release.notify_one();
wait_for_end(&mut rx).await;
handle.send(Request::Submit { text: "b".into() }).unwrap();
wait_for_end(&mut rx).await; assert_eq!(
rec.models(),
["m1", "m1", "m2"],
"the swap applied after the run, not before"
);
}
#[tokio::test]
async fn ask_status_is_answered_mid_run_from_the_snapshot() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("done".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let mut cfg = agent_config(fake_stream_fn(&rec), vec![gate]);
cfg.context = vec![AgentMessage::Assistant {
content: vec![crate::message::ContentBlock::Text {
text: "seeded work".into(),
}],
stop_reason: StopReason::Stop,
usage: None,
model: None,
}];
let handle = SessionActor::spawn(Agent::new(cfg));
let mut rx = handle.subscribe();
handle.send(Request::Submit { text: "a".into() }).unwrap();
entered_rx.recv().await.unwrap();
let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::Status))
.await
.expect("status must not wait for the run")
.unwrap();
let AgentEvent::Status { last_assistant_text } = reply else {
panic!("expected Status, got {reply:?}");
};
assert_eq!(last_assistant_text.as_deref(), Some("seeded work"));
release.notify_one();
wait_for_end(&mut rx).await;
}
}