1use 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
125pub const EVENT_BUFFER: usize = 1024;
129
130#[derive(Clone)]
135pub struct SessionHandle {
136 inbox: mpsc::UnboundedSender<Message>,
137 events: broadcast::Sender<AgentEvent>,
138 state: watch::Sender<MemberState>,
143}
144
145impl SessionHandle {
146 pub fn send(&self, request: Request) -> Result<(), SessionClosed> {
153 self.tell(None, request)
154 }
155
156 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 pub async fn ask(&self, request: Request) -> Result<AgentEvent, SessionClosed> {
182 self.question(None, request).await
183 }
184
185 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 pub fn state(&self) -> MemberState {
213 *self.state.borrow()
214 }
215
216 pub fn state_rx(&self) -> watch::Receiver<MemberState> {
220 self.state.subscribe()
221 }
222 pub fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
225 self.events.subscribe()
226 }
227}
228
229#[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
241enum Message {
243 Tell {
246 from: Option<SessionId>,
247 request: Request,
248 },
249 Ask {
251 from: Option<SessionId>,
252 request: Request,
253 reply: oneshot::Sender<AgentEvent>,
254 },
255}
256
257pub 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
278async fn serve(
280 mut agent: Agent,
281 mut inbox: mpsc::UnboundedReceiver<Message>,
282 events: broadcast::Sender<AgentEvent>,
283 state: watch::Sender<MemberState>,
284) {
285 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, },
297 };
298
299 dispatch(&mut agent, message, &mut inbox, &events, &mut deferred, &state).await;
300 }
301}
302
303fn tag(sender: &SessionId, content: &str) -> String {
306 format!("[message from {sender}]\n{content}")
307}
308
309async 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
322async 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 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 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 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 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 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 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 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
454async 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.send_replace(MemberState::Running);
471 let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::<AgentEvent>();
472 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 let side = agent.side_handle();
480
481 let mut run = Box::pin(agent.run(text, sink_tx));
482 let result = loop {
483 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 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 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 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 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 while let Ok(event) = sink_rx.try_recv() {
634 let _ = events.send(event);
635 }
636
637 if let Err(e) = &result {
641 let _ = events.send(AgentEvent::Error {
642 message: e.to_string(),
643 });
644 }
645 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 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 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 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 wait_for_end(&mut rx).await;
832 }
833
834 #[tokio::test]
835 async fn ask_cancel_aborts_an_in_flight_run() {
836 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 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 wait_for_end(&mut rx).await;
865 }
866
867 #[tokio::test]
868 async fn state_is_running_during_a_turn_then_done_after() {
869 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 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 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 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(); handle
988 .send(Request::SetModel { model: "m2".into() })
989 .unwrap(); release.notify_one();
991 wait_for_end(&mut rx).await; handle.send(Request::Submit { text: "b".into() }).unwrap();
994 wait_for_end(&mut rx).await; 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 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 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 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 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 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(); 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 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 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 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 rec.push(vec![
1420 LlmStreamEvent::TextDelta("aside".into()),
1421 LlmStreamEvent::Done {
1422 stop_reason: StopReason::Stop,
1423 usage: None,
1424 },
1425 ]);
1426 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(); 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 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; handle.send(Request::Submit { text: "b".into() }).unwrap();
1567 wait_for_end(&mut rx).await; 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 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 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(); 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}