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> {
181 self.question(None, request).await
182 }
183
184 pub async fn ask_from(
186 &self,
187 from: SessionId,
188 request: Request,
189 ) -> Result<AgentEvent, SessionClosed> {
190 self.question(Some(from), request).await
191 }
192
193 async fn question(
194 &self,
195 from: Option<SessionId>,
196 request: Request,
197 ) -> Result<AgentEvent, SessionClosed> {
198 let (reply_tx, reply_rx) = oneshot::channel();
199 self.inbox
200 .send(Message::Ask {
201 from,
202 request,
203 reply: reply_tx,
204 })
205 .map_err(|_| SessionClosed)?;
206 reply_rx.await.map_err(|_| SessionClosed)
207 }
208
209 pub fn state(&self) -> MemberState {
212 *self.state.borrow()
213 }
214
215 pub fn state_rx(&self) -> watch::Receiver<MemberState> {
219 self.state.subscribe()
220 }
221 pub fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
224 self.events.subscribe()
225 }
226}
227
228#[derive(Debug, Clone, Copy, PartialEq, Eq)]
230pub struct SessionClosed;
231
232impl fmt::Display for SessionClosed {
233 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
234 f.write_str("session actor has shut down")
235 }
236}
237
238impl std::error::Error for SessionClosed {}
239
240enum Message {
242 Tell {
245 from: Option<SessionId>,
246 request: Request,
247 },
248 Ask {
250 from: Option<SessionId>,
251 request: Request,
252 reply: oneshot::Sender<AgentEvent>,
253 },
254}
255
256pub struct SessionActor;
261
262impl SessionActor {
263 pub fn spawn(agent: Agent) -> SessionHandle {
264 let (inbox_tx, inbox_rx) = mpsc::unbounded_channel();
265 let (events_tx, _) = broadcast::channel(EVENT_BUFFER);
266 let events = events_tx.clone();
267 let (state_tx, _state_rx) = watch::channel(MemberState::Idle);
268 tokio::spawn(serve(agent, inbox_rx, events_tx, state_tx.clone()));
269 SessionHandle {
270 inbox: inbox_tx,
271 events,
272 state: state_tx,
273 }
274 }
275}
276
277async fn serve(
279 mut agent: Agent,
280 mut inbox: mpsc::UnboundedReceiver<Message>,
281 events: broadcast::Sender<AgentEvent>,
282 state: watch::Sender<MemberState>,
283) {
284 let mut deferred: VecDeque<Message> = VecDeque::new();
288
289 loop {
290 let message = match deferred.pop_front() {
291 Some(message) => message,
292 None => match inbox.recv().await {
293 Some(message) => message,
294 None => break, },
296 };
297
298 dispatch(&mut agent, message, &mut inbox, &events, &mut deferred, &state).await;
299 }
300}
301
302fn tag(sender: &SessionId, content: &str) -> String {
305 format!("[message from {sender}]\n{content}")
306}
307
308async fn gate_inbound(hooks: &HooksSet, message: &mut Message) -> Option<String> {
312 let (from, request) = match message {
313 Message::Tell { from, request } | Message::Ask { from, request, .. } => (from, request),
314 };
315 if !crate::protocol::is_inbound(request) {
316 return None;
317 }
318 hooks.before_inbound(from.as_ref(), request).await
319}
320
321async fn dispatch(
324 agent: &mut Agent,
325 mut message: Message,
326 inbox: &mut mpsc::UnboundedReceiver<Message>,
327 events: &broadcast::Sender<AgentEvent>,
328 deferred: &mut VecDeque<Message>,
329 state: &watch::Sender<MemberState>,
330) {
331 if let Some(reason) = gate_inbound(agent.hooks(), &mut message).await {
334 if let Message::Ask { reply, .. } = message {
335 let _ = reply.send(AgentEvent::Error {
336 message: format!("inbound blocked: {reason}"),
337 });
338 }
339 return;
340 }
341
342 let (from, request, reply) = match message {
343 Message::Tell { from, request } => (from, request, None),
344 Message::Ask { from, request, reply } => (from, request, Some(reply)),
345 };
346 let sender = from.unwrap_or_else(SessionId::user);
347
348 let event = match request {
349 Request::Submit { text } => match run(agent, &text, inbox, events, deferred, state).await {
350 Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
351 Err(e) => AgentEvent::Error {
352 message: e.to_string(),
353 },
354 },
355 Request::Notify { content } => match agent.notify(tag(&sender, &content)) {
358 Ok(_) => {
359 let _ = events.send(AgentEvent::MessageReceived {
360 from: sender.clone(),
361 content,
362 });
363 AgentEvent::Ack
364 }
365 Err(e) => AgentEvent::Error {
366 message: e.to_string(),
367 },
368 },
369 Request::Interrupt { content } => {
370 agent.steer(AgentMessage::user_text(tag(&sender, &content)));
371 let _ = events.send(AgentEvent::MessageReceived {
372 from: sender.clone(),
373 content,
374 });
375 AgentEvent::Ack
376 }
377 Request::Wake { content } => {
381 let _ = events.send(AgentEvent::MessageReceived {
382 from: sender.clone(),
383 content: content.clone(),
384 });
385 match run(agent, &tag(&sender, &content), inbox, events, deferred, state).await {
386 Ok(stop_reason) => AgentEvent::Stopped { stop_reason },
387 Err(e) => AgentEvent::Error {
388 message: e.to_string(),
389 },
390 }
391 }
392 Request::Cancel => AgentEvent::Ack,
396 Request::SetModel { model } => match agent.set_model(model) {
397 Ok(()) => AgentEvent::Ack,
398 Err(e) => AgentEvent::Error {
399 message: e.to_string(),
400 },
401 },
402 Request::SetEffort { effort } => match agent.set_effort(effort) {
403 Ok(()) => AgentEvent::Ack,
404 Err(e) => AgentEvent::Error {
405 message: e.to_string(),
406 },
407 },
408 Request::Compact { instructions } => match agent.compact(instructions.as_deref()).await {
409 Ok(CompactOutcome::NothingToDo) => AgentEvent::Ack,
410 Ok(CompactOutcome::Done {
411 summarized, kept, ..
412 }) => AgentEvent::Compaction { summarized, kept },
413 Err(e) => AgentEvent::Error { message: e },
414 },
415 Request::SideAsk { text } => match agent.side_ask(&text).await {
420 Ok(answer) => AgentEvent::SideAnswer {
421 text: answer.text,
422 usage: answer.usage,
423 },
424 Err(e) => AgentEvent::Error {
425 message: e.to_string(),
426 },
427 },
428 Request::SetPlanMode { on } => {
431 agent.set_plan_mode(on);
432 AgentEvent::Ack
433 }
434 Request::GetHistory => AgentEvent::History {
435 messages: agent.messages().to_vec(),
436 },
437 Request::Status => AgentEvent::Status {
438 last_assistant_text: crate::message::last_assistant_text(agent.messages()),
439 },
440 Request::Define { .. } => AgentEvent::Error {
443 message: "agent definition is served, not in-process".into(),
444 },
445 Request::ListSessions | Request::ListModels | Request::Unknown => AgentEvent::Ack,
446 };
447
448 if let Some(reply) = reply {
449 let _ = reply.send(event);
450 }
451}
452
453async fn run(
456 agent: &mut Agent,
457 text: &str,
458 inbox: &mut mpsc::UnboundedReceiver<Message>,
459 events: &broadcast::Sender<AgentEvent>,
460 deferred: &mut VecDeque<Message>,
461 state: &watch::Sender<MemberState>,
462) -> Result<StopReason, LoopError> {
463 state.send_replace(MemberState::Running);
470 let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::<AgentEvent>();
471 let steer = agent.steer_sender();
474 let follow_up = agent.follow_up_sender();
475 let cancel = agent.cancel_token();
476 let hooks = agent.hooks().clone();
477 let side = agent.side_handle();
479
480 let mut run = Box::pin(agent.run(text, sink_tx));
481 let result = loop {
482 tokio::select! {
488 biased;
489 Some(mut message) = inbox.recv() => {
490 if let Some(reason) = gate_inbound(&hooks, &mut message).await {
491 if let Message::Ask { reply, .. } = message {
492 let _ = reply.send(AgentEvent::Error {
493 message: format!("inbound blocked: {reason}"),
494 });
495 }
496 continue;
497 }
498 match message {
499 Message::Tell {
500 request: Request::Cancel,
501 ..
502 } => cancel.cancel(),
503 Message::Tell {
504 request: Request::Notify { content } | Request::Interrupt { content },
505 from,
506 } => {
507 let from = from.unwrap_or_else(SessionId::user);
508 let _ = steer.send(AgentMessage::user_text(tag(&from, &content)));
509 let _ = events.send(AgentEvent::MessageReceived {
510 from,
511 content,
512 });
513 }
514 Message::Tell {
515 request: Request::Wake { content },
516 from,
517 } => {
518 let from = from.unwrap_or_else(SessionId::user);
519 let _ = follow_up.send(AgentMessage::user_text(tag(&from, &content)));
520 let _ = events.send(AgentEvent::MessageReceived {
521 from,
522 content,
523 });
524 }
525 Message::Ask {
530 request: Request::GetHistory,
531 reply,
532 ..
533 } => {
534 let _ = reply.send(AgentEvent::History { messages: side.history() });
535 }
536 Message::Ask {
537 request: Request::Status,
538 reply,
539 ..
540 } => {
541 let _ = reply.send(AgentEvent::Status {
542 last_assistant_text: side.last_assistant_text(),
543 });
544 }
545 Message::Ask {
546 request: Request::SideAsk { text },
547 reply,
548 ..
549 } => {
550 let side = side.clone();
553 tokio::spawn(async move {
554 let event = match side.side_ask(&text).await {
555 Ok(answer) => AgentEvent::SideAnswer {
556 text: answer.text,
557 usage: answer.usage,
558 },
559 Err(e) => AgentEvent::Error {
560 message: e.to_string(),
561 },
562 };
563 let _ = reply.send(event);
564 });
565 }
566 Message::Ask {
571 from,
572 request:
573 req @ (Request::SetModel { .. }
574 | Request::SetEffort { .. }
575 | Request::SetPlanMode { .. }),
576 reply,
577 } => {
578 let _ = reply.send(AgentEvent::Ack);
579 deferred.push_back(Message::Tell { from, request: req });
580 }
581 other => deferred.push_back(other),
582 }
583 },
584 result = &mut run => break result,
585 Some(event) = sink_rx.recv() => {
586 let _ = events.send(event);
587 }
588 }
589 };
590
591 while let Ok(event) = sink_rx.try_recv() {
594 let _ = events.send(event);
595 }
596
597 if let Err(e) = &result {
601 let _ = events.send(AgentEvent::Error {
602 message: e.to_string(),
603 });
604 }
605 match &result {
609 Ok(StopReason::Error) | Err(_) => {
610 state.send_replace(MemberState::Failed);
611 }
612 Ok(_) => {
613 state.send_replace(MemberState::Done);
614 }
615 }
616 result
617}
618
619#[cfg(test)]
620mod tests {
621 use super::*;
622 use crate::agent::{Agent, AgentConfig};
623 use crate::compaction::CompactionPolicy;
624 use crate::event::LlmStreamEvent;
625 use crate::hooks::HooksSet;
626 use crate::message::StopReason;
627 use crate::streamfn::{LlmOpts, LlmStream, StreamFn};
628 use crate::tool::{ToolContext, ToolOutput, TypedTool, erased};
629 use futures::StreamExt as _;
630 use serde::Deserialize;
631 use serde_json::json;
632 use std::sync::{Arc, Mutex};
633 use std::time::Duration;
634
635 #[derive(Clone, Default)]
636 struct Recorder {
637 script: Arc<Mutex<VecDeque<Vec<LlmStreamEvent>>>>,
638 calls: Arc<Mutex<Vec<Vec<AgentMessage>>>>,
639 models: Arc<Mutex<Vec<String>>>,
640 }
641
642 impl Recorder {
643 fn push(&self, events: Vec<LlmStreamEvent>) {
644 self.script.lock().unwrap().push_back(events);
645 }
646 fn calls(&self) -> Vec<Vec<AgentMessage>> {
647 self.calls.lock().unwrap().clone()
648 }
649 fn models(&self) -> Vec<String> {
650 self.models.lock().unwrap().clone()
651 }
652 }
653
654 fn fake_stream_fn(rec: &Recorder) -> StreamFn {
655 let rec = rec.clone();
656 Arc::new(
657 move |ctx: &[AgentMessage], _system, _tools, opts: &LlmOpts| {
658 rec.calls.lock().unwrap().push(ctx.to_vec());
659 rec.models.lock().unwrap().push(opts.model.clone());
660 let events = rec.script.lock().unwrap().pop_front().unwrap_or_default();
661 Box::pin(futures::stream::iter(events)) as LlmStream
662 },
663 )
664 }
665
666 fn agent_config(stream_fn: StreamFn, tools: Vec<crate::tool::Tool>) -> AgentConfig {
667 AgentConfig {
668 system: "sys".into(),
669 tools,
670 llm: LlmOpts {
671 model: "m1".into(),
672 ..LlmOpts::default()
673 },
674 stream_fn,
675 hooks: HooksSet::default(),
676 session: None,
677 context: Vec::new(),
678 working_dir: std::path::PathBuf::new(),
679 max_turns: crate::loop_::DEFAULT_MAX_TURNS,
680 parallel_tools: true,
681 compaction: CompactionPolicy::default(),
682 plan_mode: crate::hooks::PlanModeHandle::new(),
683 }
684 }
685
686 async fn wait_for_end(rx: &mut broadcast::Receiver<AgentEvent>) {
688 loop {
689 let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
690 .await
691 .expect("an event within the timeout")
692 .expect("the outbox stays open");
693 if matches!(event, AgentEvent::AgentEnd) {
694 return;
695 }
696 }
697 }
698
699 #[tokio::test]
700 async fn submit_streams_events_to_a_subscriber() {
701 let rec = Recorder::default();
702 rec.push(vec![
703 LlmStreamEvent::TextDelta("hello".into()),
704 LlmStreamEvent::Done {
705 stop_reason: StopReason::Stop,
706 usage: None,
707 },
708 ]);
709 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
710 let mut rx = handle.subscribe();
711 handle.send(Request::Submit { text: "hi".into() }).unwrap();
712
713 let mut saw_text = false;
714 loop {
715 let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
716 .await
717 .expect("an event")
718 .expect("open");
719 match event {
720 AgentEvent::MessageUpdate { message } if message.as_text() == "hello" => {
721 saw_text = true;
722 }
723 AgentEvent::AgentEnd => break,
724 _ => {}
725 }
726 }
727 assert!(saw_text, "the streamed text reached the subscriber");
728 }
729
730 #[tokio::test]
731 async fn idle_cancel_does_not_poison_the_next_run() {
732 let rec = Recorder::default();
733 for text in ["one", "two"] {
734 rec.push(vec![
735 LlmStreamEvent::TextDelta(text.into()),
736 LlmStreamEvent::Done {
737 stop_reason: StopReason::Stop,
738 usage: None,
739 },
740 ]);
741 }
742 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
743 let mut rx = handle.subscribe();
744
745 handle.send(Request::Submit { text: "a".into() }).unwrap();
746 wait_for_end(&mut rx).await;
747
748 handle.send(Request::Cancel).unwrap();
750 handle.send(Request::Submit { text: "b".into() }).unwrap();
751
752 let mut saw_two = false;
753 loop {
754 let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
755 .await
756 .expect("an event")
757 .expect("open");
758 match event {
759 AgentEvent::MessageUpdate { message } if message.as_text() == "two" => {
760 saw_two = true;
761 }
762 AgentEvent::AgentEnd => break,
763 _ => {}
764 }
765 }
766 assert!(saw_two, "the second run completed despite the idle cancel");
767 }
768
769 #[tokio::test]
770 async fn cancel_aborts_an_in_flight_run() {
771 let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
773 let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
774 Box::pin(head.chain(futures::stream::pending())) as LlmStream
775 });
776 let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
777 let mut rx = handle.subscribe();
778 handle.send(Request::Submit { text: "hi".into() }).unwrap();
779
780 loop {
781 let event = tokio::time::timeout(Duration::from_secs(5), rx.recv())
782 .await
783 .expect("an event")
784 .expect("open");
785 if matches!(event, AgentEvent::MessageUpdate { .. }) {
786 handle.send(Request::Cancel).unwrap();
787 break;
788 }
789 }
790 wait_for_end(&mut rx).await;
792 }
793
794 #[tokio::test]
795 async fn state_is_running_during_a_turn_then_done_after() {
796 let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
799 let head = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
800 Box::pin(head.chain(futures::stream::pending())) as LlmStream
801 });
802 let handle = SessionActor::spawn(Agent::new(agent_config(stream_fn, vec![])));
803 assert_eq!(handle.state(), MemberState::Idle, "a virgin session is idle");
804
805 let mut state_rx = handle.state_rx();
807 handle.send(Request::Submit { text: "hi".into() }).unwrap();
808 tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
809 .await
810 .expect("a state change")
811 .expect("open");
812 assert_eq!(*state_rx.borrow_and_update(), MemberState::Running);
813 assert_eq!(handle.state(), MemberState::Running, "the run is in flight");
814
815 handle.send(Request::Cancel).unwrap();
816 tokio::time::timeout(Duration::from_secs(5), state_rx.changed())
817 .await
818 .expect("a state change")
819 .expect("open");
820 assert_eq!(
821 *state_rx.borrow_and_update(),
822 MemberState::Done,
823 "an abort ends Done, not Failed"
824 );
825 assert_eq!(handle.state(), MemberState::Done);
826 }
827
828 #[tokio::test]
829 async fn a_fatal_stream_error_ends_failed() {
830 let rec = Recorder::default();
831 rec.push(vec![LlmStreamEvent::Error {
832 message: "boom".into(),
833 fatal: true,
834 }]);
835 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
836 let mut rx = handle.subscribe();
837 handle.send(Request::Submit { text: "hi".into() }).unwrap();
838 wait_for_end(&mut rx).await;
839 assert_eq!(handle.state(), MemberState::Failed);
840 }
841
842 #[derive(Deserialize, schemars::JsonSchema)]
843 struct GateArgs {
844 text: String,
845 }
846
847 struct GateTool {
850 entered: mpsc::UnboundedSender<()>,
851 release: Arc<tokio::sync::Notify>,
852 }
853
854 #[async_trait::async_trait]
855 impl TypedTool for GateTool {
856 type Args = GateArgs;
857 fn name(&self) -> &str {
858 "gate"
859 }
860 fn description(&self) -> &str {
861 "blocks until released"
862 }
863 async fn execute(&self, args: GateArgs, _ctx: &ToolContext) -> ToolOutput {
864 self.entered.send(()).ok();
865 self.release.notified().await;
866 ToolOutput {
867 output: format!("released:{}", args.text),
868 is_error: false,
869 diff: None,
870 path: None,
871 }
872 }
873 }
874
875 #[tokio::test]
876 async fn config_request_during_a_run_is_deferred_to_after_it() {
877 let rec = Recorder::default();
880 rec.push(vec![
881 LlmStreamEvent::ToolCall {
882 id: "c1".into(),
883 name: "gate".into(),
884 arguments: json!({ "text": "hi" }),
885 },
886 LlmStreamEvent::Done {
887 stop_reason: StopReason::ToolUse,
888 usage: None,
889 },
890 ]);
891 for text in ["turn two", "run two"] {
892 rec.push(vec![
893 LlmStreamEvent::TextDelta(text.into()),
894 LlmStreamEvent::Done {
895 stop_reason: StopReason::Stop,
896 usage: None,
897 },
898 ]);
899 }
900
901 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
902 let release = Arc::new(tokio::sync::Notify::new());
903 let gate = erased(GateTool {
904 entered: entered_tx,
905 release: release.clone(),
906 });
907
908 let handle =
909 SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
910 let mut rx = handle.subscribe();
911 handle.send(Request::Submit { text: "a".into() }).unwrap();
912
913 entered_rx.recv().await.unwrap(); handle
915 .send(Request::SetModel { model: "m2".into() })
916 .unwrap(); release.notify_one();
918 wait_for_end(&mut rx).await; handle.send(Request::Submit { text: "b".into() }).unwrap();
921 wait_for_end(&mut rx).await; assert_eq!(
924 rec.models(),
925 ["m1", "m1", "m2"],
926 "the swap waited for the run, then applied to the next one"
927 );
928 }
929
930 #[tokio::test]
931 async fn ask_history_returns_the_conversation() {
932 let rec = Recorder::default();
933 rec.push(vec![
934 LlmStreamEvent::TextDelta("hi back".into()),
935 LlmStreamEvent::Done {
936 stop_reason: StopReason::Stop,
937 usage: None,
938 },
939 ]);
940 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
941 let mut rx = handle.subscribe();
942 handle
943 .send(Request::Submit {
944 text: "hello".into(),
945 })
946 .unwrap();
947 wait_for_end(&mut rx).await;
948
949 let reply = handle.ask(Request::GetHistory).await.unwrap();
950 let AgentEvent::History { messages } = reply else {
951 panic!("expected a History reply, got {reply:?}");
952 };
953 assert!(
954 messages.iter().any(|m| m.as_text() == "hello"),
955 "the read-back carries the user turn: {messages:?}"
956 );
957 }
958
959 #[tokio::test]
960 async fn ask_set_model_replies_ack_and_takes_effect() {
961 let rec = Recorder::default();
962 rec.push(vec![
963 LlmStreamEvent::TextDelta("x".into()),
964 LlmStreamEvent::Done {
965 stop_reason: StopReason::Stop,
966 usage: None,
967 },
968 ]);
969 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
970 let mut rx = handle.subscribe();
971
972 let reply = handle
973 .ask(Request::SetModel { model: "m2".into() })
974 .await
975 .unwrap();
976 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
977
978 handle.send(Request::Submit { text: "go".into() }).unwrap();
979 wait_for_end(&mut rx).await;
980 assert_eq!(rec.models(), ["m2"], "the swap applied to the next run");
981 }
982
983 #[tokio::test]
984 async fn define_on_a_local_session_is_an_error() {
985 let rec = Recorder::default();
986 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
987 let reply = handle
988 .ask(Request::Define {
989 name: Some("w1".into()),
990 model: None,
991 role: None,
992 tools: None,
993 base_url: None,
994 api_key: None,
995 read_only: false,
996 effort: None,
997 })
998 .await
999 .unwrap();
1000 assert!(
1001 matches!(&reply, AgentEvent::Error { message } if message.contains("served")),
1002 "{reply:?}"
1003 );
1004 }
1005
1006 #[tokio::test]
1007 async fn ask_compact_with_nothing_to_do_replies_ack() {
1008 let handle = SessionActor::spawn(Agent::new(agent_config(
1009 fake_stream_fn(&Recorder::default()),
1010 vec![],
1011 )));
1012 let reply = handle
1013 .ask(Request::Compact { instructions: None })
1014 .await
1015 .unwrap();
1016 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1017 }
1018
1019 #[tokio::test]
1020 async fn ask_submit_replies_stopped_with_the_reason() {
1021 let rec = Recorder::default();
1022 rec.push(vec![
1023 LlmStreamEvent::TextDelta("done".into()),
1024 LlmStreamEvent::Done {
1025 stop_reason: StopReason::Stop,
1026 usage: None,
1027 },
1028 ]);
1029 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1030 let reply = handle
1031 .ask(Request::Submit { text: "hi".into() })
1032 .await
1033 .unwrap();
1034 assert!(
1035 matches!(
1036 reply,
1037 AgentEvent::Stopped {
1038 stop_reason: StopReason::Stop
1039 }
1040 ),
1041 "{reply:?}"
1042 );
1043 }
1044
1045 #[tokio::test]
1046 async fn notify_appends_without_starting_a_turn() {
1047 let rec = Recorder::default();
1048 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1049 let mut rx = handle.subscribe();
1050
1051 let reply = handle
1052 .ask(Request::Notify {
1053 content: "ping".into(),
1054 })
1055 .await
1056 .unwrap();
1057 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1058
1059 let surfaced = tokio::time::timeout(Duration::from_secs(1), rx.recv())
1061 .await
1062 .expect("an event")
1063 .expect("open");
1064 assert!(
1065 matches!(&surfaced, AgentEvent::MessageReceived { content, .. } if content == "ping"),
1066 "{surfaced:?}"
1067 );
1068
1069 let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1071 panic!("expected History");
1072 };
1073 assert!(
1074 messages
1075 .iter()
1076 .any(|m| m.as_text() == "[message from user]\nping"),
1077 "{messages:?}"
1078 );
1079
1080 assert!(rec.calls().is_empty(), "notify must not run the loop");
1082 }
1083
1084 #[tokio::test]
1085 async fn wake_starts_a_run_when_idle() {
1086 let rec = Recorder::default();
1087 rec.push(vec![
1088 LlmStreamEvent::TextDelta("acted".into()),
1089 LlmStreamEvent::Done {
1090 stop_reason: StopReason::Stop,
1091 usage: None,
1092 },
1093 ]);
1094 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1095 let mut rx = handle.subscribe();
1096
1097 let reply = handle
1098 .ask(Request::Wake {
1099 content: "do the task".into(),
1100 })
1101 .await
1102 .unwrap();
1103 assert!(
1104 matches!(
1105 reply,
1106 AgentEvent::Stopped {
1107 stop_reason: StopReason::Stop
1108 }
1109 ),
1110 "{reply:?}"
1111 );
1112 wait_for_end(&mut rx).await;
1113
1114 let calls = rec.calls();
1115 assert_eq!(calls.len(), 1, "Wake ran the loop once");
1116 assert!(
1117 calls[0]
1118 .iter()
1119 .any(|m| m.as_text() == "[message from user]\ndo the task"),
1120 "the tagged task is the prompt: {:?}",
1121 calls[0]
1122 );
1123 }
1124
1125 struct InboundPolicy;
1126
1127 #[async_trait::async_trait]
1128 impl crate::hooks::Hooks for InboundPolicy {
1129 async fn before_inbound(
1130 &self,
1131 _from: Option<&SessionId>,
1132 request: &mut Request,
1133 ) -> Option<String> {
1134 let Request::Notify { content } = request else {
1135 return None;
1136 };
1137 if content.starts_with("drop:") {
1138 return Some("policy".into());
1139 }
1140 *content = format!("{content}!");
1141 None
1142 }
1143 }
1144
1145 #[tokio::test]
1146 async fn before_inbound_drops_and_rewrites_a_notify() {
1147 let rec = Recorder::default();
1148 let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
1149 cfg.hooks = HooksSet::one(Arc::new(InboundPolicy));
1150 let handle = SessionActor::spawn(Agent::new(cfg));
1151
1152 let reply = handle
1154 .ask(Request::Notify {
1155 content: "drop: me".into(),
1156 })
1157 .await
1158 .unwrap();
1159 assert!(
1160 matches!(&reply, AgentEvent::Error { message } if message.contains("policy")),
1161 "{reply:?}"
1162 );
1163
1164 let reply = handle
1166 .ask(Request::Notify {
1167 content: "keep me".into(),
1168 })
1169 .await
1170 .unwrap();
1171 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1172
1173 let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1174 panic!("expected History");
1175 };
1176 let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
1177 assert!(
1178 !texts.iter().any(|t| t == "drop: me"),
1179 "the dropped message never landed: {texts:?}"
1180 );
1181 assert!(
1182 texts.iter().any(|t| t == "[message from user]\nkeep me!"),
1183 "the rewrite landed: {texts:?}"
1184 );
1185 }
1186
1187 #[tokio::test]
1188 async fn interrupt_reaches_the_next_turn_while_a_tool_is_blocked() {
1189 let rec = Recorder::default();
1190 rec.push(vec![
1191 LlmStreamEvent::ToolCall {
1192 id: "c1".into(),
1193 name: "gate".into(),
1194 arguments: json!({ "text": "hi" }),
1195 },
1196 LlmStreamEvent::Done {
1197 stop_reason: StopReason::ToolUse,
1198 usage: None,
1199 },
1200 ]);
1201 rec.push(vec![
1202 LlmStreamEvent::TextDelta("after interrupt".into()),
1203 LlmStreamEvent::Done {
1204 stop_reason: StopReason::Stop,
1205 usage: None,
1206 },
1207 ]);
1208
1209 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1210 let release = Arc::new(tokio::sync::Notify::new());
1211 let gate = erased(GateTool {
1212 entered: entered_tx,
1213 release: release.clone(),
1214 });
1215
1216 let handle =
1217 SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1218 let mut rx = handle.subscribe();
1219 handle.send(Request::Submit { text: "hi".into() }).unwrap();
1220
1221 entered_rx.recv().await.unwrap(); handle
1223 .send(Request::Interrupt {
1224 content: "mid-run interrupt".into(),
1225 })
1226 .unwrap();
1227 release.notify_one();
1228
1229 wait_for_end(&mut rx).await;
1230
1231 let calls = rec.calls();
1232 assert_eq!(calls.len(), 2, "two stream calls: the tool turn and the next");
1233 assert!(
1234 calls[1]
1235 .iter()
1236 .any(|m| matches!(m, AgentMessage::User { .. })
1237 && m.as_text() == "[message from user]\nmid-run interrupt"),
1238 "the interrupt drained into the next turn: {:?}",
1239 calls[1]
1240 );
1241 }
1242
1243 #[tokio::test]
1244 async fn send_from_attributes_the_sender() {
1245 let rec = Recorder::default();
1246 let handle = SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![])));
1247 let mut rx = handle.subscribe();
1248
1249 handle
1250 .send_from(
1251 SessionId::agent("b"),
1252 Request::Notify {
1253 content: "ping".into(),
1254 },
1255 )
1256 .unwrap();
1257
1258 let event = tokio::time::timeout(Duration::from_secs(1), rx.recv())
1259 .await
1260 .expect("an event")
1261 .expect("open");
1262 assert!(
1263 matches!(&event, AgentEvent::MessageReceived { from, content }
1264 if from.as_str() == "agent:b" && content == "ping"),
1265 "{event:?}"
1266 );
1267 }
1268
1269 struct SenderPolicy;
1270
1271 #[async_trait::async_trait]
1272 impl crate::hooks::Hooks for SenderPolicy {
1273 async fn before_inbound(
1274 &self,
1275 from: Option<&SessionId>,
1276 _request: &mut Request,
1277 ) -> Option<String> {
1278 (from.map(SessionId::as_str) == Some("agent:evil")).then(|| "blocked sender".into())
1279 }
1280 }
1281
1282 #[tokio::test]
1283 async fn before_inbound_sees_the_sender() {
1284 let rec = Recorder::default();
1285 let mut cfg = agent_config(fake_stream_fn(&rec), vec![]);
1286 cfg.hooks = HooksSet::one(Arc::new(SenderPolicy));
1287 let handle = SessionActor::spawn(Agent::new(cfg));
1288
1289 let reply = handle
1291 .ask_from(
1292 SessionId::agent("evil"),
1293 Request::Notify {
1294 content: "let me in".into(),
1295 },
1296 )
1297 .await
1298 .unwrap();
1299 assert!(
1300 matches!(&reply, AgentEvent::Error { message } if message.contains("blocked sender")),
1301 "{reply:?}"
1302 );
1303
1304 let reply = handle
1306 .ask_from(
1307 SessionId::agent("friend"),
1308 Request::Notify {
1309 content: "hello".into(),
1310 },
1311 )
1312 .await
1313 .unwrap();
1314 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1315
1316 let AgentEvent::History { messages } = handle.ask(Request::GetHistory).await.unwrap() else {
1317 panic!("expected History");
1318 };
1319 let texts: Vec<String> = messages.iter().map(|m| m.as_text()).collect();
1320 assert!(
1321 texts
1322 .iter()
1323 .any(|t| t == "[message from agent:friend]\nhello"),
1324 "{texts:?}"
1325 );
1326 assert!(!texts.iter().any(|t| t.contains("let me in")), "{texts:?}");
1327 }
1328
1329 #[tokio::test]
1330 async fn side_ask_is_answered_mid_run() {
1331 let rec = Recorder::default();
1334 rec.push(vec![
1335 LlmStreamEvent::ToolCall {
1336 id: "c1".into(),
1337 name: "gate".into(),
1338 arguments: json!({ "text": "hi" }),
1339 },
1340 LlmStreamEvent::Done {
1341 stop_reason: StopReason::ToolUse,
1342 usage: None,
1343 },
1344 ]);
1345 rec.push(vec![
1347 LlmStreamEvent::TextDelta("aside".into()),
1348 LlmStreamEvent::Done {
1349 stop_reason: StopReason::Stop,
1350 usage: None,
1351 },
1352 ]);
1353 rec.push(vec![
1355 LlmStreamEvent::TextDelta("done".into()),
1356 LlmStreamEvent::Done {
1357 stop_reason: StopReason::Stop,
1358 usage: None,
1359 },
1360 ]);
1361
1362 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1363 let release = Arc::new(tokio::sync::Notify::new());
1364 let gate = erased(GateTool {
1365 entered: entered_tx,
1366 release: release.clone(),
1367 });
1368 let handle =
1369 SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1370 let mut rx = handle.subscribe();
1371 handle.send(Request::Submit { text: "a".into() }).unwrap();
1372
1373 entered_rx.recv().await.unwrap(); let reply = tokio::time::timeout(
1377 Duration::from_secs(5),
1378 handle.ask(Request::SideAsk { text: "why?".into() }),
1379 )
1380 .await
1381 .expect("the side ask must not wait for the run")
1382 .unwrap();
1383 assert!(
1384 matches!(&reply, AgentEvent::SideAnswer { text, .. } if text == "aside"),
1385 "{reply:?}"
1386 );
1387
1388 release.notify_one();
1389 wait_for_end(&mut rx).await;
1390 }
1391
1392 #[tokio::test]
1393 async fn get_history_is_answered_mid_run_from_the_snapshot() {
1394 let rec = Recorder::default();
1395 rec.push(vec![
1396 LlmStreamEvent::ToolCall {
1397 id: "c1".into(),
1398 name: "gate".into(),
1399 arguments: json!({ "text": "hi" }),
1400 },
1401 LlmStreamEvent::Done {
1402 stop_reason: StopReason::ToolUse,
1403 usage: None,
1404 },
1405 ]);
1406 rec.push(vec![
1407 LlmStreamEvent::TextDelta("done".into()),
1408 LlmStreamEvent::Done {
1409 stop_reason: StopReason::Stop,
1410 usage: None,
1411 },
1412 ]);
1413
1414 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1415 let release = Arc::new(tokio::sync::Notify::new());
1416 let gate = erased(GateTool {
1417 entered: entered_tx,
1418 release: release.clone(),
1419 });
1420 let handle =
1421 SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1422 let mut rx = handle.subscribe();
1423 handle.send(Request::Submit { text: "a".into() }).unwrap();
1424
1425 entered_rx.recv().await.unwrap();
1426
1427 let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::GetHistory))
1428 .await
1429 .expect("history must not wait for the run")
1430 .unwrap();
1431 let AgentEvent::History { messages } = reply else {
1432 panic!("expected History, got {reply:?}");
1433 };
1434 assert!(
1435 messages.iter().any(|m| m.as_text() == "a"),
1436 "the mid-run snapshot carries the submit: {messages:?}"
1437 );
1438
1439 release.notify_one();
1440 wait_for_end(&mut rx).await;
1441 }
1442
1443 #[tokio::test]
1444 async fn a_deferred_command_is_acked_before_the_run_ends() {
1445 let rec = Recorder::default();
1446 rec.push(vec![
1447 LlmStreamEvent::ToolCall {
1448 id: "c1".into(),
1449 name: "gate".into(),
1450 arguments: json!({ "text": "hi" }),
1451 },
1452 LlmStreamEvent::Done {
1453 stop_reason: StopReason::ToolUse,
1454 usage: None,
1455 },
1456 ]);
1457 for text in ["turn two", "run two"] {
1458 rec.push(vec![
1459 LlmStreamEvent::TextDelta(text.into()),
1460 LlmStreamEvent::Done {
1461 stop_reason: StopReason::Stop,
1462 usage: None,
1463 },
1464 ]);
1465 }
1466
1467 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1468 let release = Arc::new(tokio::sync::Notify::new());
1469 let gate = erased(GateTool {
1470 entered: entered_tx,
1471 release: release.clone(),
1472 });
1473 let handle =
1474 SessionActor::spawn(Agent::new(agent_config(fake_stream_fn(&rec), vec![gate])));
1475 let mut rx = handle.subscribe();
1476 handle.send(Request::Submit { text: "a".into() }).unwrap();
1477
1478 entered_rx.recv().await.unwrap();
1479
1480 let reply = tokio::time::timeout(
1482 Duration::from_secs(5),
1483 handle.ask(Request::SetModel { model: "m2".into() }),
1484 )
1485 .await
1486 .expect("a deferred command must be acked, not hang")
1487 .unwrap();
1488 assert!(matches!(reply, AgentEvent::Ack), "{reply:?}");
1489
1490 release.notify_one();
1491 wait_for_end(&mut rx).await; handle.send(Request::Submit { text: "b".into() }).unwrap();
1494 wait_for_end(&mut rx).await; assert_eq!(
1496 rec.models(),
1497 ["m1", "m1", "m2"],
1498 "the swap applied after the run, not before"
1499 );
1500 }
1501
1502 #[tokio::test]
1503 async fn ask_status_is_answered_mid_run_from_the_snapshot() {
1504 let rec = Recorder::default();
1507 rec.push(vec![
1508 LlmStreamEvent::ToolCall {
1509 id: "c1".into(),
1510 name: "gate".into(),
1511 arguments: json!({ "text": "hi" }),
1512 },
1513 LlmStreamEvent::Done {
1514 stop_reason: StopReason::ToolUse,
1515 usage: None,
1516 },
1517 ]);
1518 rec.push(vec![
1519 LlmStreamEvent::TextDelta("done".into()),
1520 LlmStreamEvent::Done {
1521 stop_reason: StopReason::Stop,
1522 usage: None,
1523 },
1524 ]);
1525
1526 let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
1527 let release = Arc::new(tokio::sync::Notify::new());
1528 let gate = erased(GateTool {
1529 entered: entered_tx,
1530 release: release.clone(),
1531 });
1532 let mut cfg = agent_config(fake_stream_fn(&rec), vec![gate]);
1534 cfg.context = vec![AgentMessage::Assistant {
1535 content: vec![crate::message::ContentBlock::Text {
1536 text: "seeded work".into(),
1537 }],
1538 stop_reason: StopReason::Stop,
1539 usage: None,
1540 model: None,
1541 }];
1542 let handle = SessionActor::spawn(Agent::new(cfg));
1543 let mut rx = handle.subscribe();
1544 handle.send(Request::Submit { text: "a".into() }).unwrap();
1545
1546 entered_rx.recv().await.unwrap(); let reply = tokio::time::timeout(Duration::from_secs(5), handle.ask(Request::Status))
1549 .await
1550 .expect("status must not wait for the run")
1551 .unwrap();
1552 let AgentEvent::Status { last_assistant_text } = reply else {
1553 panic!("expected Status, got {reply:?}");
1554 };
1555 assert_eq!(last_assistant_text.as_deref(), Some("seeded work"));
1556
1557 release.notify_one();
1558 wait_for_end(&mut rx).await;
1559 }
1560}