1use std::collections::{HashMap, HashSet, VecDeque};
9use std::path::PathBuf;
10use std::sync::{Arc, Mutex as StdMutex};
11
12use chrono::{DateTime, Utc};
13use tokio::sync::{broadcast, mpsc, Mutex as AsyncMutex, RwLock as TokioRwLock};
14use tokio_util::sync::CancellationToken;
15
16use bamboo_agent_core::tools::ToolExecutor;
17use bamboo_agent_core::{AgentEvent, Message, Session};
18use bamboo_domain::reasoning::ReasoningEffort;
19use bamboo_engine::execution::runner_state::AgentRunner;
20use bamboo_engine::execution::{
21 create_event_forwarder, get_or_create_event_sender, spawn_session_execution,
22 try_reserve_runner, SessionExecutionArgs,
23};
24use bamboo_engine::{AuxiliaryModelConfig, SessionRepository};
25use bamboo_llm::{Config, ProviderRegistry};
26
27use super::approvals::{self, ParkedAsk, RespondAndResumeOutcome, Responder};
28use super::platform::{CallbackQuery, InboundMessage, OutboundMessage, Platform, ReplyCtx};
29use super::render;
30
31#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct SessionKey {
35 pub platform: String,
36 pub chat_id: String,
37 pub user_id: String,
38}
39
40impl SessionKey {
41 pub fn as_string(&self) -> String {
42 format!("{}:{}:{}", self.platform, self.chat_id, self.user_id)
43 }
44}
45
46const DEDUP_CAPACITY: usize = 10_000;
55
56struct BoundedSeenSet {
62 set: HashSet<String>,
63 order: VecDeque<String>,
64 capacity: usize,
65}
66
67impl BoundedSeenSet {
68 fn new(capacity: usize) -> Self {
69 Self {
70 set: HashSet::new(),
71 order: VecDeque::new(),
72 capacity: capacity.max(1),
73 }
74 }
75
76 fn insert(&mut self, key: String) -> bool {
80 if !self.set.insert(key.clone()) {
81 return false;
82 }
83 self.order.push_back(key);
84 while self.order.len() > self.capacity {
85 if let Some(oldest) = self.order.pop_front() {
86 self.set.remove(&oldest);
87 }
88 }
89 true
90 }
91
92 #[cfg(test)]
93 fn len(&self) -> usize {
94 self.set.len()
95 }
96}
97
98#[derive(Clone)]
102pub struct ConnectContext {
103 pub agent: Arc<bamboo_engine::Agent>,
104 pub tools: Arc<dyn ToolExecutor>,
105 pub session_repo: SessionRepository,
106 pub agent_runners: Arc<TokioRwLock<HashMap<String, AgentRunner>>>,
107 pub session_event_senders: Arc<TokioRwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
108 pub account_feed_inbox: Option<bamboo_engine::execution::AccountFeedInbox>,
109 pub app_data_dir: Option<PathBuf>,
110 pub config: Arc<tokio::sync::RwLock<Config>>,
111 pub provider_registry: Arc<ProviderRegistry>,
112 pub permission_checker: Arc<dyn bamboo_tools::permission::PermissionChecker>,
118}
119
120#[derive(Default)]
127struct ChatState {
128 busy: bool,
129 queue: VecDeque<(Arc<dyn Platform>, InboundMessage)>,
130 cancel_token: Option<CancellationToken>,
131 pending_ask: Option<ParkedAsk>,
134 ask_resolution: Option<mpsc::Sender<AskResolution>>,
142}
143
144#[derive(Debug, Clone)]
146enum AskResolution {
147 Answer(String),
150 Invalidated,
154}
155
156type ResolvedConnectRunConfig = bamboo_engine::resolved_defaults::ResolvedDefaultRunConfig;
165
166fn resolve_connect_run_config(
167 config_snapshot: &Config,
168 provider_registry: &Arc<ProviderRegistry>,
169) -> ResolvedConnectRunConfig {
170 bamboo_engine::resolved_defaults::resolve_default_run_config(config_snapshot, provider_registry)
171}
172
173fn create_connect_session(
188 key: &str,
189 model: &str,
190 system_prompt: &str,
191 base_system_prompt: &str,
192 workspace_path: Option<&str>,
193 reasoning_effort: Option<ReasoningEffort>,
194) -> Session {
195 let session_id = uuid::Uuid::new_v4().to_string();
196 let mut session = Session::new(session_id.clone(), model.to_string());
197 session.title = format!("Connect: {key}");
198 session
199 .metadata
200 .insert("created_by_connect_key".to_string(), key.to_string());
201 session.metadata.insert(
202 "base_system_prompt".to_string(),
203 base_system_prompt.to_string(),
204 );
205 if let Some(path) = workspace_path {
206 session.set_workspace_path_meta(path);
207 bamboo_tools::tools::workspace_state::ensure_session_workspace(
208 &session_id,
209 Some(PathBuf::from(path)),
210 );
211 }
212 if let Some(effort) = reasoning_effort {
213 session.set_reasoning_effort_meta(effort.as_str());
214 }
215 session.add_message(Message::system(system_prompt.to_string()));
216 bamboo_engine::runner::refresh_prompt_snapshot(&mut session);
217 session
218 .agent_runtime_state
219 .get_or_insert_with(bamboo_domain::AgentRuntimeState::default)
220 .no_human_approver = false;
221 session
222}
223
224fn strip_command_suffix(text: &str) -> &str {
227 text.split('@').next().unwrap_or(text)
228}
229
230async fn reply_text(platform: &Arc<dyn Platform>, ctx: &ReplyCtx, text: impl Into<String>) {
231 if let Err(error) = platform.reply(ctx, OutboundMessage::text(text)).await {
232 tracing::warn!("connect: failed to send reply: {error}");
233 }
234}
235
236pub struct ConnectBridge {
240 ctx: ConnectContext,
241 session_map: TokioRwLock<HashMap<String, String>>,
244 map_path: Option<PathBuf>,
245 chat_state: AsyncMutex<HashMap<String, ChatState>>,
246 seen_message_ids: StdMutex<BoundedSeenSet>,
253 process_start: DateTime<Utc>,
254 responder: Arc<dyn Responder>,
257}
258
259impl ConnectBridge {
260 pub fn new(ctx: ConnectContext, map_path: Option<PathBuf>) -> Self {
264 let responder = Arc::new(approvals::EngineResponder::new(ctx.clone()));
265 Self::with_responder(ctx, map_path, responder)
266 }
267
268 pub fn with_responder(
272 ctx: ConnectContext,
273 map_path: Option<PathBuf>,
274 responder: Arc<dyn Responder>,
275 ) -> Self {
276 Self {
277 ctx,
278 session_map: TokioRwLock::new(HashMap::new()),
279 map_path,
280 chat_state: AsyncMutex::new(HashMap::new()),
281 seen_message_ids: StdMutex::new(BoundedSeenSet::new(DEDUP_CAPACITY)),
282 process_start: Utc::now(),
283 responder,
284 }
285 }
286
287 pub async fn load_session_map(&self) {
292 let Some(path) = &self.map_path else {
293 return;
294 };
295 match tokio::fs::read(path).await {
296 Ok(bytes) => match serde_json::from_slice::<HashMap<String, String>>(&bytes) {
297 Ok(map) => *self.session_map.write().await = map,
298 Err(error) => {
299 tracing::warn!(
300 "connect: session map at {path:?} is corrupt, starting empty: {error}"
301 );
302 }
303 },
304 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
305 Err(error) => {
306 tracing::warn!("connect: failed to read session map at {path:?}: {error}");
307 }
308 }
309 }
310
311 pub async fn session_id_for_key(&self, key: &str) -> Option<String> {
312 self.session_map.read().await.get(key).cloned()
313 }
314
315 async fn set_session_id_for_key(&self, key: &str, session_id: &str) {
316 {
317 let mut map = self.session_map.write().await;
318 map.insert(key.to_string(), session_id.to_string());
319 }
320 self.persist_session_map().await;
321 }
322
323 async fn rotate_session(&self, key: &str) {
328 self.invalidate_pending_ask(key).await;
329 {
330 let mut map = self.session_map.write().await;
331 map.remove(key);
332 }
333 self.persist_session_map().await;
334 }
335
336 async fn has_pending_ask(&self, key: &str) -> bool {
338 self.chat_state
339 .lock()
340 .await
341 .get(key)
342 .is_some_and(|state| state.pending_ask.is_some())
343 }
344
345 async fn try_resolve_pending_ask(
353 &self,
354 key: &str,
355 resolve: impl FnOnce(&ParkedAsk) -> Option<String>,
356 ) -> Option<(String, mpsc::Sender<AskResolution>)> {
357 let mut guard = self.chat_state.lock().await;
358 let state = guard.get_mut(key)?;
359 let ask_ref = state.pending_ask.as_ref()?;
360 let answer = resolve(ask_ref)?;
361 let sender = state.ask_resolution.take()?;
362 state.pending_ask = None;
363 Some((answer, sender))
364 }
365
366 async fn invalidate_pending_ask(&self, key: &str) {
369 let sender = {
370 let mut guard = self.chat_state.lock().await;
371 match guard.get_mut(key) {
372 Some(state) => {
373 state.pending_ask = None;
374 state.ask_resolution.take()
375 }
376 None => None,
377 }
378 };
379 if let Some(sender) = sender {
380 let _ = sender.send(AskResolution::Invalidated).await;
381 }
382 }
383
384 async fn clear_pending_ask(&self, key: &str) {
388 let mut guard = self.chat_state.lock().await;
389 if let Some(state) = guard.get_mut(key) {
390 state.pending_ask = None;
391 state.ask_resolution = None;
392 }
393 }
394
395 async fn persist_session_map(&self) {
396 let Some(path) = &self.map_path else {
397 return;
398 };
399 let snapshot = self.session_map.read().await.clone();
400 let json = match serde_json::to_vec_pretty(&snapshot) {
401 Ok(json) => json,
402 Err(error) => {
403 tracing::warn!("connect: failed to serialize session map: {error}");
404 return;
405 }
406 };
407 if let Err(error) = atomic_write(path, &json).await {
408 tracing::warn!("connect: failed to persist session map at {path:?}: {error}");
409 }
410 }
411
412 async fn set_cancel_token(&self, key: &str, token: CancellationToken) {
413 let mut guard = self.chat_state.lock().await;
414 guard.entry(key.to_string()).or_default().cancel_token = Some(token);
415 }
416
417 async fn clear_cancel_token(&self, key: &str) {
418 let mut guard = self.chat_state.lock().await;
419 if let Some(state) = guard.get_mut(key) {
420 state.cancel_token = None;
421 }
422 }
423
424 pub async fn handle_inbound(
436 self: Arc<Self>,
437 platform: Arc<dyn Platform>,
438 allow_from: Vec<String>,
439 msg: InboundMessage,
440 ) {
441 if !allow_from.iter().any(|allowed| allowed == &msg.user_id) {
442 tracing::warn!(
443 platform = %msg.platform,
444 chat_id = %msg.chat_id,
445 user_id = %msg.user_id,
446 "connect: rejected inbound message — user not in allow_from"
447 );
448 return;
449 }
450
451 if msg.sent_at < self.process_start {
452 tracing::debug!(
453 platform = %msg.platform,
454 message_id = %msg.message_id,
455 "connect: dropping message older than process start"
456 );
457 return;
458 }
459
460 let dedup_key = format!("{}:{}", msg.platform, msg.message_id);
461 {
462 let mut seen = self.seen_message_ids.lock().unwrap();
463 if !seen.insert(dedup_key) {
464 tracing::debug!(
465 platform = %msg.platform,
466 message_id = %msg.message_id,
467 "connect: dropping duplicate message_id"
468 );
469 return;
470 }
471 }
472
473 let key = SessionKey {
474 platform: msg.platform.clone(),
475 chat_id: msg.chat_id.clone(),
476 user_id: msg.user_id.clone(),
477 }
478 .as_string();
479
480 let command = strip_command_suffix(msg.text.trim());
481 if command.eq_ignore_ascii_case("/stop") {
482 self.handle_stop(&key, &platform, &msg.reply_ctx).await;
483 return;
484 }
485 if command.eq_ignore_ascii_case("/status") {
486 self.handle_status(&key, &platform, &msg.reply_ctx).await;
487 return;
488 }
489
490 if let Some((answer, sender)) = self
497 .try_resolve_pending_ask(&key, |ask| approvals::match_text_answer(ask, &msg.text))
498 .await
499 {
500 let _ = sender.send(AskResolution::Answer(answer)).await;
501 return;
502 }
503
504 if command.eq_ignore_ascii_case("/new") && self.has_pending_ask(&key).await {
509 self.rotate_session(&key).await;
510 reply_text(&platform, &msg.reply_ctx, "Started a new session.").await;
511 return;
512 }
513
514 let mut guard = self.chat_state.lock().await;
515 let state = guard.entry(key.clone()).or_default();
516 if state.busy {
517 state.queue.push_back((platform, msg));
518 return;
519 }
520 state.busy = true;
521 drop(guard);
522
523 let bridge = self.clone();
524 tokio::spawn(async move {
525 bridge.drain_chat(key, platform, msg).await;
526 });
527 }
528
529 async fn drain_chat(
533 self: Arc<Self>,
534 key: String,
535 mut platform: Arc<dyn Platform>,
536 mut msg: InboundMessage,
537 ) {
538 loop {
539 self.process_one(&key, platform.clone(), msg).await;
540
541 let next = {
542 let mut guard = self.chat_state.lock().await;
543 match guard.get_mut(&key) {
544 Some(state) => match state.queue.pop_front() {
545 Some(item) => Some(item),
546 None => {
547 state.busy = false;
548 None
549 }
550 },
551 None => None,
552 }
553 };
554
555 match next {
556 Some((p, m)) => {
557 platform = p;
558 msg = m;
559 }
560 None => break,
561 }
562 }
563 }
564
565 pub async fn handle_callback(
573 self: Arc<Self>,
574 platform: Arc<dyn Platform>,
575 allow_from: Vec<String>,
576 callback: CallbackQuery,
577 ) {
578 if !allow_from
579 .iter()
580 .any(|allowed| allowed == &callback.user_id)
581 {
582 tracing::warn!(
583 platform = %callback.platform,
584 chat_id = %callback.chat_id,
585 user_id = %callback.user_id,
586 "connect: rejected callback query — user not in allow_from"
587 );
588 let _ = platform
589 .answer_callback(&callback.callback_query_id, None)
590 .await;
591 return;
592 }
593
594 let key = SessionKey {
595 platform: callback.platform.clone(),
596 chat_id: callback.chat_id.clone(),
597 user_id: callback.user_id.clone(),
598 }
599 .as_string();
600
601 let resolved = self
602 .try_resolve_pending_ask(&key, |ask| {
603 approvals::match_callback_data(ask, &callback.data)
604 })
605 .await;
606
607 match resolved {
608 Some((answer, sender)) => {
609 let _ = platform
610 .answer_callback(&callback.callback_query_id, None)
611 .await;
612 let _ = sender.send(AskResolution::Answer(answer)).await;
613 }
614 None => {
615 tracing::debug!(
616 platform = %callback.platform,
617 chat_id = %callback.chat_id,
618 "connect: dropping stale/forged callback_data"
619 );
620 let _ = platform
621 .answer_callback(
622 &callback.callback_query_id,
623 Some("This action has expired."),
624 )
625 .await;
626 }
627 }
628 }
629
630 async fn process_one(&self, key: &str, platform: Arc<dyn Platform>, msg: InboundMessage) {
631 let command = strip_command_suffix(msg.text.trim());
632 if command.eq_ignore_ascii_case("/new") {
633 self.rotate_session(key).await;
634 reply_text(&platform, &msg.reply_ctx, "Started a new session.").await;
635 return;
636 }
637
638 let text = msg.text.trim();
639 if text.is_empty() {
640 return;
641 }
642
643 self.run_prompt(key, platform, &msg.reply_ctx, text).await;
644 }
645
646 async fn handle_stop(&self, key: &str, platform: &Arc<dyn Platform>, reply_ctx: &ReplyCtx) {
647 let token = {
648 self.chat_state
649 .lock()
650 .await
651 .get(key)
652 .and_then(|state| state.cancel_token.clone())
653 };
654 let had_pending_ask = self.has_pending_ask(key).await;
659 if had_pending_ask {
660 self.invalidate_pending_ask(key).await;
661 }
662 match (token, had_pending_ask) {
663 (Some(token), _) => {
664 token.cancel();
665 reply_text(platform, reply_ctx, "Stopping the current run…").await;
666 }
667 (None, true) => {
668 reply_text(
669 platform,
670 reply_ctx,
671 "Stopped — the pending question was cancelled.",
672 )
673 .await;
674 }
675 (None, false) => {
676 reply_text(platform, reply_ctx, "Nothing is running.").await;
677 }
678 }
679 }
680
681 async fn handle_status(&self, key: &str, platform: &Arc<dyn Platform>, reply_ctx: &ReplyCtx) {
682 let session_id = self.session_id_for_key(key).await;
683 let busy = {
684 self.chat_state
685 .lock()
686 .await
687 .get(key)
688 .map(|state| state.busy)
689 .unwrap_or(false)
690 };
691 let text = match session_id {
692 Some(id) => format!(
693 "Session: {id}\nStatus: {}",
694 if busy { "busy" } else { "idle" }
695 ),
696 None => "No session yet. Send a message to start one.".to_string(),
697 };
698 reply_text(platform, reply_ctx, text).await;
699 }
700
701 async fn create_and_register_session(
702 &self,
703 key: &str,
704 resolved: &ResolvedConnectRunConfig,
705 ) -> Session {
706 let model = resolved.model_roster.model.clone().unwrap_or_default();
707 let session = create_connect_session(
708 key,
709 &model,
710 &resolved.system_prompt,
711 &resolved.base_system_prompt,
712 resolved.workspace_path.as_deref(),
713 resolved.reasoning_effort,
714 );
715 self.set_session_id_for_key(key, &session.id).await;
716 session
717 }
718
719 async fn run_prompt(
727 &self,
728 key: &str,
729 platform: Arc<dyn Platform>,
730 reply_ctx: &ReplyCtx,
731 text: &str,
732 ) {
733 let config_snapshot = self.ctx.config.read().await.clone();
734 let resolved = resolve_connect_run_config(&config_snapshot, &self.ctx.provider_registry);
735
736 if resolved
737 .model_roster
738 .model
739 .as_deref()
740 .unwrap_or("")
741 .trim()
742 .is_empty()
743 {
744 reply_text(
745 &platform,
746 reply_ctx,
747 "No model is configured for this bamboo instance; cannot run your request.",
748 )
749 .await;
750 return;
751 }
752
753 let existing_id = self.session_id_for_key(key).await;
754 let mut session = match existing_id {
755 Some(id) => match self.ctx.session_repo.load_merged(&id).await {
756 Some(session) => session,
757 None => self.create_and_register_session(key, &resolved).await,
758 },
759 None => self.create_and_register_session(key, &resolved).await,
760 };
761
762 session.add_message(Message::user(text.to_string()));
763 let session_id = session.id.clone();
764 self.ctx.session_repo.save_and_cache(&mut session).await;
765
766 let session_tx =
767 get_or_create_event_sender(&self.ctx.session_event_senders, &session_id).await;
768 let rx = session_tx.subscribe();
769
770 let Some(reservation) = try_reserve_runner(
771 &self.ctx.agent_runners,
772 &self.ctx.session_event_senders,
773 &session_id,
774 &session_tx,
775 )
776 .await
777 else {
778 reply_text(
779 &platform,
780 reply_ctx,
781 "This session is already running elsewhere; please wait for it to finish.",
782 )
783 .await;
784 return;
785 };
786
787 self.set_cancel_token(key, reservation.cancel_token.clone())
788 .await;
789
790 let (mpsc_tx, _forwarder_handle) = create_event_forwarder(
791 session_id.clone(),
792 session_tx.clone(),
793 self.ctx.agent_runners.clone(),
794 self.ctx.account_feed_inbox.clone(),
795 );
796
797 let aux_fast_model = resolved.model_roster.fast_model();
800 let aux_fast_provider = resolved.model_roster.fast_model_provider();
801 let aux_background_model = resolved.model_roster.background_model();
802 let aux_background_provider = resolved.model_roster.background_model_provider();
803 let aux_summarization_model = resolved.model_roster.summarization_model();
804 let aux_summarization_provider = resolved.model_roster.summarization_model_provider();
805 let auxiliary_model_resolver = Arc::new(move || AuxiliaryModelConfig {
806 fast_model_name: aux_fast_model.clone(),
807 fast_model_provider: aux_fast_provider.clone(),
808 background_model_name: aux_background_model.clone(),
809 planning_model_name: None,
810 search_model_name: None,
811 summarization_model_name: aux_summarization_model.clone(),
812 background_model_provider: aux_background_provider.clone(),
813 summarization_model_provider: aux_summarization_provider.clone(),
814 });
815
816 spawn_session_execution(SessionExecutionArgs {
817 agent: self.ctx.agent.clone(),
818 session_id: session_id.clone(),
819 session,
820 tools_override: Some(self.ctx.tools.clone()),
821 provider_override: None,
822 model_roster: resolved.model_roster.clone(),
823 reasoning_effort: resolved.reasoning_effort,
824 reasoning_effort_source: "connect".to_string(),
825 auxiliary_model_resolver: Some(auxiliary_model_resolver),
826 disabled_filter_resolver: None,
827 disabled_tools: None,
828 disabled_skill_ids: None,
829 selected_skill_ids: None,
830 selected_skill_mode: None,
831 cancel_token: reservation.cancel_token,
832 mpsc_tx,
833 image_fallback: None,
834 gold_config: resolved.gold_config.clone(),
835 guardian_config: None,
839 guardian_spawner: None,
840 bash_resume_hook: None,
841 bash_completion_sink: None,
842 app_data_dir: self.ctx.app_data_dir.clone(),
843 run_budget: None,
846 runners: self.ctx.agent_runners.clone(),
847 sessions_cache: self.ctx.session_repo.cache().clone(),
848 on_complete: None,
849 child_completion_handler: None,
852 });
853
854 self.render_until_settled(key, platform, reply_ctx.clone(), &session_id, rx)
855 .await;
856
857 self.clear_cancel_token(key).await;
858 }
859
860 async fn render_until_settled(
876 &self,
877 key: &str,
878 platform: Arc<dyn Platform>,
879 reply_ctx: ReplyCtx,
880 session_id: &str,
881 mut rx: broadcast::Receiver<AgentEvent>,
882 ) {
883 let mut stream_state: Option<Box<render::StreamState>> = None;
884 loop {
885 match render::stream_execution(
886 platform.clone(),
887 reply_ctx.clone(),
888 rx,
889 stream_state.take(),
890 )
891 .await
892 {
893 render::RunOutcome::Terminal => return,
894 render::RunOutcome::Paused {
895 ask,
896 stream_state: paused_state,
897 } => {
898 stream_state = paused_state;
899 let caps = platform.capabilities();
900 let parked =
901 ParkedAsk::new(approvals::new_nonce(), session_id.to_string(), &ask);
902
903 if let Err(error) =
904 approvals::render_ask(&platform, &reply_ctx, &parked, caps.buttons).await
905 {
906 tracing::warn!("connect: failed to render pending ask: {error}");
907 }
908
909 let (ask_tx, mut ask_rx) = mpsc::channel(1);
910 {
911 let mut guard = self.chat_state.lock().await;
912 let state = guard.entry(key.to_string()).or_default();
913 state.pending_ask = Some(parked);
914 state.ask_resolution = Some(ask_tx);
915 }
916
917 match ask_rx.recv().await {
918 Some(AskResolution::Answer(answer)) => {
919 match self.responder.respond_and_resume(session_id, answer).await {
920 Ok(RespondAndResumeOutcome::Resumed(new_rx)) => {
921 rx = new_rx;
922 continue;
923 }
924 Ok(RespondAndResumeOutcome::NotResumed(reason)) => {
925 reply_text(&platform, &reply_ctx, format!("({reason})")).await;
926 return;
927 }
928 Err(error) => {
929 reply_text(
930 &platform,
931 &reply_ctx,
932 format!("Failed to record your answer: {error}"),
933 )
934 .await;
935 return;
936 }
937 }
938 }
939 Some(AskResolution::Invalidated) | None => {
940 self.clear_pending_ask(key).await;
946 return;
947 }
948 }
949 }
950 }
951 }
952 }
953}
954
955async fn atomic_write(path: &std::path::Path, bytes: &[u8]) -> std::io::Result<()> {
960 if let Some(parent) = path.parent() {
961 tokio::fs::create_dir_all(parent).await?;
962 }
963 let tmp = path.with_extension(format!("tmp.{}", uuid::Uuid::new_v4()));
964 {
965 let mut file = tokio::fs::File::create(&tmp).await?;
966 tokio::io::AsyncWriteExt::write_all(&mut file, bytes).await?;
967 file.sync_all().await?;
968 }
969 tokio::fs::rename(&tmp, path).await?;
970 if let Some(parent) = path.parent() {
971 if let Ok(dir) = tokio::fs::File::open(parent).await {
972 let _ = dir.sync_all().await;
973 }
974 }
975 Ok(())
976}
977
978#[cfg(test)]
979mod tests {
980 use super::*;
981 use crate::app_state::AppState;
982 use crate::tools::ToolSurface;
983 use std::time::Duration;
984 use tokio::sync::Mutex as TokioMutex;
985
986 struct FakePlatform {
992 label: String,
993 capabilities: super::super::platform::Capabilities,
994 sent: TokioMutex<Vec<String>>,
995 sent_messages: TokioMutex<Vec<OutboundMessage>>,
996 edits: TokioMutex<Vec<String>>,
997 answered_callbacks: TokioMutex<Vec<(String, Option<String>)>>,
998 }
999
1000 impl FakePlatform {
1001 fn new(label: &str) -> Arc<Self> {
1002 Self::with_capabilities(label, Default::default())
1003 }
1004
1005 fn with_capabilities(
1006 label: &str,
1007 capabilities: super::super::platform::Capabilities,
1008 ) -> Arc<Self> {
1009 Arc::new(Self {
1010 label: label.to_string(),
1011 capabilities,
1012 sent: TokioMutex::new(Vec::new()),
1013 sent_messages: TokioMutex::new(Vec::new()),
1014 edits: TokioMutex::new(Vec::new()),
1015 answered_callbacks: TokioMutex::new(Vec::new()),
1016 })
1017 }
1018
1019 async fn sent_texts(&self) -> Vec<String> {
1020 self.sent.lock().await.clone()
1021 }
1022 }
1023
1024 #[async_trait::async_trait]
1025 impl Platform for FakePlatform {
1026 fn name(&self) -> &str {
1027 &self.label
1028 }
1029 fn capabilities(&self) -> super::super::platform::Capabilities {
1030 self.capabilities
1031 }
1032 async fn start(
1033 &self,
1034 _inbound: tokio::sync::mpsc::Sender<super::super::platform::Inbound>,
1035 ) -> super::super::platform::PlatformResult<()> {
1036 Ok(())
1037 }
1038 async fn reply(
1039 &self,
1040 _ctx: &ReplyCtx,
1041 msg: OutboundMessage,
1042 ) -> super::super::platform::PlatformResult<super::super::platform::MessageRef> {
1043 self.sent.lock().await.push(msg.text.clone());
1044 self.sent_messages.lock().await.push(msg);
1045 Ok(super::super::platform::MessageRef(serde_json::Value::Null))
1046 }
1047 async fn edit(
1048 &self,
1049 _msg_ref: &super::super::platform::MessageRef,
1050 new: OutboundMessage,
1051 ) -> super::super::platform::PlatformResult<()> {
1052 self.edits.lock().await.push(new.text);
1053 Ok(())
1054 }
1055 async fn answer_callback(
1056 &self,
1057 callback_query_id: &str,
1058 text: Option<&str>,
1059 ) -> super::super::platform::PlatformResult<()> {
1060 self.answered_callbacks
1061 .lock()
1062 .await
1063 .push((callback_query_id.to_string(), text.map(str::to_string)));
1064 Ok(())
1065 }
1066 async fn stop(&self) -> super::super::platform::PlatformResult<()> {
1067 Ok(())
1068 }
1069 }
1070
1071 struct FakeResponder {
1076 calls: TokioMutex<Vec<(String, String)>>,
1077 resume_sender: broadcast::Sender<AgentEvent>,
1078 fail_with: Option<String>,
1079 }
1080
1081 impl FakeResponder {
1082 fn new(resume_sender: broadcast::Sender<AgentEvent>) -> Arc<Self> {
1083 Arc::new(Self {
1084 calls: TokioMutex::new(Vec::new()),
1085 resume_sender,
1086 fail_with: None,
1087 })
1088 }
1089
1090 fn failing(resume_sender: broadcast::Sender<AgentEvent>, reason: &str) -> Arc<Self> {
1091 Arc::new(Self {
1092 calls: TokioMutex::new(Vec::new()),
1093 resume_sender,
1094 fail_with: Some(reason.to_string()),
1095 })
1096 }
1097 }
1098
1099 #[async_trait::async_trait]
1100 impl Responder for FakeResponder {
1101 async fn respond_and_resume(
1102 &self,
1103 session_id: &str,
1104 answer: String,
1105 ) -> Result<RespondAndResumeOutcome, super::super::approvals::ResponderError> {
1106 self.calls
1107 .lock()
1108 .await
1109 .push((session_id.to_string(), answer));
1110 if let Some(reason) = &self.fail_with {
1111 return Err(super::super::approvals::ResponderError::Other(
1112 reason.clone(),
1113 ));
1114 }
1115 Ok(RespondAndResumeOutcome::Resumed(
1116 self.resume_sender.subscribe(),
1117 ))
1118 }
1119 }
1120
1121 async fn wait_for_parked_ask(bridge: &ConnectBridge, key: &str) -> ParkedAsk {
1126 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
1127 loop {
1128 if let Some(ask) = bridge
1129 .chat_state
1130 .lock()
1131 .await
1132 .get(key)
1133 .and_then(|state| state.pending_ask.clone())
1134 {
1135 return ask;
1136 }
1137 assert!(
1138 tokio::time::Instant::now() < deadline,
1139 "ask was never parked for {key}"
1140 );
1141 tokio::time::sleep(Duration::from_millis(10)).await;
1142 }
1143 }
1144
1145 async fn wait_for_responder_calls(responder: &FakeResponder, count: usize) {
1152 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
1153 loop {
1154 if responder.calls.lock().await.len() >= count {
1155 return;
1156 }
1157 assert!(
1158 tokio::time::Instant::now() < deadline,
1159 "responder never reached {count} call(s)"
1160 );
1161 tokio::time::sleep(Duration::from_millis(10)).await;
1162 }
1163 }
1164
1165 async fn test_context() -> (ConnectContext, tempfile::TempDir) {
1166 let dir = tempfile::tempdir().unwrap();
1167 let state = AppState::new(dir.path().to_path_buf())
1176 .await
1177 .expect("app state");
1178 {
1179 let mut cfg = state.config.write().await;
1180 cfg.provider = "openai".to_string();
1181 cfg.providers.openai = Some(bamboo_config::OpenAIConfig {
1182 api_key: String::new(),
1183 api_key_from_env: false,
1184 api_key_encrypted: None,
1185 base_url: None,
1186 model: Some("gpt-4o-mini".to_string()),
1187 fast_model: None,
1188 vision_model: None,
1189 reasoning_effort: None,
1190 responses_only_models: Vec::new(),
1191 request_overrides: None,
1192 extra: Default::default(),
1193 });
1194 }
1195 let ctx = ConnectContext {
1196 agent: state.agent.clone(),
1197 tools: state.tools_for(ToolSurface::Root),
1198 session_repo: state.session_repo.clone(),
1199 agent_runners: state.agent_runners.clone(),
1200 session_event_senders: state.session_event_senders.clone(),
1201 account_feed_inbox: None,
1202 app_data_dir: Some(state.app_data_dir.clone()),
1203 config: state.config.clone(),
1204 provider_registry: state.provider_registry.clone(),
1205 permission_checker: state.permission_checker.clone(),
1206 };
1207 (ctx, dir)
1208 }
1209
1210 fn inbound(chat_id: &str, user_id: &str, message_id: &str, text: &str) -> InboundMessage {
1211 InboundMessage {
1212 platform: "fake".to_string(),
1213 chat_id: chat_id.to_string(),
1214 user_id: user_id.to_string(),
1215 message_id: message_id.to_string(),
1216 sent_at: Utc::now(),
1217 text: text.to_string(),
1218 reply_ctx: ReplyCtx(serde_json::json!({ "chat_id": chat_id })),
1219 }
1220 }
1221
1222 fn key_for(chat_id: &str, user_id: &str) -> String {
1223 SessionKey {
1224 platform: "fake".to_string(),
1225 chat_id: chat_id.to_string(),
1226 user_id: user_id.to_string(),
1227 }
1228 .as_string()
1229 }
1230
1231 #[test]
1232 fn session_key_formats_as_platform_chat_user() {
1233 let key = SessionKey {
1234 platform: "telegram".to_string(),
1235 chat_id: "42".to_string(),
1236 user_id: "7".to_string(),
1237 };
1238 assert_eq!(key.as_string(), "telegram:42:7");
1239 }
1240
1241 #[test]
1244 fn bounded_seen_set_evicts_the_oldest_entry_once_over_capacity() {
1245 let mut set = BoundedSeenSet::new(2);
1246 assert!(set.insert("a".to_string()));
1247 assert!(set.insert("b".to_string()));
1248 assert_eq!(set.len(), 2);
1249
1250 assert!(set.insert("c".to_string()));
1253 assert_eq!(set.len(), 2);
1254 assert!(!set.insert("b".to_string()), "b must still be tracked");
1255 assert!(!set.insert("c".to_string()), "c must still be tracked");
1256
1257 assert!(set.insert("a".to_string()));
1260 assert_eq!(set.len(), 2);
1261 }
1262
1263 #[test]
1264 fn bounded_seen_set_still_dedups_within_capacity() {
1265 let mut set = BoundedSeenSet::new(10);
1266 assert!(set.insert("x".to_string()));
1267 assert!(!set.insert("x".to_string()), "duplicate must be rejected");
1268 assert_eq!(set.len(), 1);
1269 }
1270
1271 #[test]
1272 fn bounded_seen_set_never_grows_past_capacity() {
1273 let mut set = BoundedSeenSet::new(50);
1274 for i in 0..5_000 {
1275 set.insert(format!("msg-{i}"));
1276 }
1277 assert_eq!(set.len(), 50);
1278 }
1279
1280 #[tokio::test]
1281 async fn allow_from_denies_users_not_in_the_list() {
1282 let (ctx, _dir) = test_context().await;
1283 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1284 let platform = FakePlatform::new("fake");
1285
1286 ConnectBridge::handle_inbound(
1287 bridge.clone(),
1288 platform.clone(),
1289 vec!["allowed-user".to_string()],
1290 inbound("chat1", "someone-else", "1", "hello"),
1291 )
1292 .await;
1293
1294 tokio::time::sleep(Duration::from_millis(50)).await;
1295 assert!(platform.sent_texts().await.is_empty());
1296 assert!(bridge
1297 .session_id_for_key(&key_for("chat1", "someone-else"))
1298 .await
1299 .is_none());
1300 }
1301
1302 #[tokio::test]
1303 async fn dedup_drops_repeated_message_ids() {
1304 let (ctx, _dir) = test_context().await;
1305 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1306 let platform = FakePlatform::new("fake");
1307 let allow = vec!["u1".to_string()];
1308
1309 ConnectBridge::handle_inbound(
1313 bridge.clone(),
1314 platform.clone(),
1315 allow.clone(),
1316 inbound("chat1", "u1", "dup-1", "/status"),
1317 )
1318 .await;
1319 ConnectBridge::handle_inbound(
1320 bridge.clone(),
1321 platform.clone(),
1322 allow,
1323 inbound("chat1", "u1", "dup-1", "/status"),
1324 )
1325 .await;
1326
1327 let sent = platform.sent_texts().await;
1328 assert_eq!(
1329 sent.len(),
1330 1,
1331 "duplicate message_id must be dropped: {sent:?}"
1332 );
1333 }
1334
1335 #[tokio::test]
1336 async fn older_than_process_start_messages_are_dropped() {
1337 let (ctx, _dir) = test_context().await;
1338 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1339 let platform = FakePlatform::new("fake");
1340 let mut msg = inbound("chat1", "u1", "1", "/status");
1341 msg.sent_at = bridge.process_start - chrono::Duration::seconds(5);
1342
1343 ConnectBridge::handle_inbound(
1344 bridge.clone(),
1345 platform.clone(),
1346 vec!["u1".to_string()],
1347 msg,
1348 )
1349 .await;
1350
1351 assert!(platform.sent_texts().await.is_empty());
1352 }
1353
1354 #[tokio::test]
1355 async fn status_command_reports_idle_with_no_session_yet() {
1356 let (ctx, _dir) = test_context().await;
1357 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1358 let platform = FakePlatform::new("fake");
1359
1360 ConnectBridge::handle_inbound(
1361 bridge.clone(),
1362 platform.clone(),
1363 vec!["u1".to_string()],
1364 inbound("chat1", "u1", "1", "/status"),
1365 )
1366 .await;
1367
1368 let sent = platform.sent_texts().await;
1369 assert_eq!(sent.len(), 1);
1370 assert!(sent[0].contains("No session yet"), "got: {:?}", sent[0]);
1371 }
1372
1373 #[tokio::test]
1374 async fn stop_with_nothing_running_replies_nothing_running() {
1375 let (ctx, _dir) = test_context().await;
1376 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1377 let platform = FakePlatform::new("fake");
1378
1379 ConnectBridge::handle_inbound(
1380 bridge.clone(),
1381 platform.clone(),
1382 vec!["u1".to_string()],
1383 inbound("chat1", "u1", "1", "/stop"),
1384 )
1385 .await;
1386
1387 assert_eq!(
1388 platform.sent_texts().await,
1389 vec!["Nothing is running.".to_string()]
1390 );
1391 }
1392
1393 #[tokio::test]
1394 async fn prompt_creates_a_session_and_maps_it_to_the_chat_key() {
1395 let (ctx, _dir) = test_context().await;
1396 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1397 let platform = FakePlatform::new("fake");
1398 let key = key_for("chat1", "u1");
1399
1400 ConnectBridge::handle_inbound(
1401 bridge.clone(),
1402 platform.clone(),
1403 vec!["u1".to_string()],
1404 inbound("chat1", "u1", "1", "hello there"),
1405 )
1406 .await;
1407
1408 let deadline = tokio::time::Instant::now() + Duration::from_secs(15);
1409 loop {
1410 if bridge.session_id_for_key(&key).await.is_some() {
1411 break;
1412 }
1413 assert!(
1414 tokio::time::Instant::now() < deadline,
1415 "session was never created for the chat key"
1416 );
1417 tokio::time::sleep(Duration::from_millis(20)).await;
1418 }
1419 }
1420
1421 #[tokio::test]
1422 async fn new_command_rotates_the_session_mapping() {
1423 let (ctx, _dir) = test_context().await;
1424 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1425 let platform = FakePlatform::new("fake");
1426 let key = key_for("chat1", "u1");
1427
1428 bridge
1429 .set_session_id_for_key(&key, "pre-existing-session")
1430 .await;
1431 assert_eq!(
1432 bridge.session_id_for_key(&key).await.as_deref(),
1433 Some("pre-existing-session")
1434 );
1435
1436 ConnectBridge::handle_inbound(
1437 bridge.clone(),
1438 platform.clone(),
1439 vec!["u1".to_string()],
1440 inbound("chat1", "u1", "1", "/new"),
1441 )
1442 .await;
1443
1444 let deadline = tokio::time::Instant::now() + Duration::from_secs(15);
1445 loop {
1446 if bridge.session_id_for_key(&key).await.is_none() {
1447 break;
1448 }
1449 assert!(
1450 tokio::time::Instant::now() < deadline,
1451 "session mapping was never rotated away"
1452 );
1453 tokio::time::sleep(Duration::from_millis(20)).await;
1454 }
1455
1456 let sent = platform.sent_texts().await;
1457 assert!(sent.iter().any(|t| t == "Started a new session."));
1458 }
1459
1460 #[tokio::test]
1461 async fn busy_queue_drains_a_second_message_after_the_first_finishes() {
1462 let (ctx, _dir) = test_context().await;
1463 let bridge = Arc::new(ConnectBridge::new(ctx, None));
1464 let platform = FakePlatform::new("fake");
1465 let allow = vec!["u1".to_string()];
1466 let key = key_for("chat1", "u1");
1467
1468 ConnectBridge::handle_inbound(
1472 bridge.clone(),
1473 platform.clone(),
1474 allow.clone(),
1475 inbound("chat1", "u1", "1", "first"),
1476 )
1477 .await;
1478 ConnectBridge::handle_inbound(
1479 bridge.clone(),
1480 platform.clone(),
1481 allow,
1482 inbound("chat1", "u1", "2", "second"),
1483 )
1484 .await;
1485
1486 let deadline = tokio::time::Instant::now() + Duration::from_secs(15);
1487 loop {
1488 let idle = {
1489 !bridge
1490 .chat_state
1491 .lock()
1492 .await
1493 .get(&key)
1494 .map(|state| state.busy)
1495 .unwrap_or(true)
1496 };
1497 if idle {
1498 break;
1499 }
1500 assert!(
1501 tokio::time::Instant::now() < deadline,
1502 "chat never drained back to idle"
1503 );
1504 tokio::time::sleep(Duration::from_millis(20)).await;
1505 }
1506
1507 assert!(bridge.session_id_for_key(&key).await.is_some());
1508 }
1509
1510 #[tokio::test]
1511 async fn session_map_persists_and_reloads_across_bridge_instances() {
1512 let (ctx, dir) = test_context().await;
1513 let map_path = dir.path().join("connect_sessions.json");
1514 let bridge = ConnectBridge::new(ctx.clone(), Some(map_path.clone()));
1515 bridge.set_session_id_for_key("k1", "sess-1").await;
1516
1517 let bridge2 = ConnectBridge::new(ctx, Some(map_path));
1518 bridge2.load_session_map().await;
1519
1520 assert_eq!(
1521 bridge2.session_id_for_key("k1").await.as_deref(),
1522 Some("sess-1")
1523 );
1524 }
1525
1526 fn buttons_and_edit_capabilities() -> super::super::platform::Capabilities {
1529 super::super::platform::Capabilities {
1530 buttons: true,
1531 edit_message: true,
1532 images: false,
1533 files: false,
1534 }
1535 }
1536
1537 fn need_clarification_event(
1538 question: &str,
1539 options: Vec<&str>,
1540 allow_custom: bool,
1541 ) -> AgentEvent {
1542 AgentEvent::NeedClarification {
1543 question: question.to_string(),
1544 options: Some(options.into_iter().map(str::to_string).collect()),
1545 tool_call_id: Some("call-1".to_string()),
1546 tool_name: Some("request_permissions".to_string()),
1547 allow_custom,
1548 }
1549 }
1550
1551 #[tokio::test]
1552 async fn paused_run_renders_buttons_with_nonce_and_resolves_via_callback() {
1553 let (ctx, _dir) = test_context().await;
1554 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1555 let responder = FakeResponder::new(resume_tx.clone());
1556 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1557 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1558 let key = key_for("chat1", "u1");
1559 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1560
1561 let (tx, rx) = broadcast::channel(16);
1562 tx.send(need_clarification_event(
1563 "Approve?",
1564 vec!["Approve", "Deny"],
1565 false,
1566 ))
1567 .unwrap();
1568
1569 let render_handle = {
1570 let bridge = bridge.clone();
1571 let platform = platform.clone();
1572 let reply_ctx = reply_ctx.clone();
1573 let key = key.clone();
1574 tokio::spawn(async move {
1575 bridge
1576 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1577 .await;
1578 })
1579 };
1580
1581 let parked = wait_for_parked_ask(&bridge, &key).await;
1582 assert_eq!(
1583 parked.options,
1584 vec!["Approve".to_string(), "Deny".to_string()]
1585 );
1586
1587 let sent = platform.sent_messages.lock().await.clone();
1590 let ask_message = sent
1591 .iter()
1592 .find(|message| message.buttons.is_some())
1593 .expect("expected a buttoned ask message");
1594 let buttons = ask_message.buttons.as_ref().unwrap();
1595 assert_eq!(buttons.len(), 2);
1596 assert_eq!(buttons[0][0].callback_data, format!("{}:0", parked.nonce));
1597 assert_eq!(buttons[1][0].callback_data, format!("{}:1", parked.nonce));
1598
1599 let callback = CallbackQuery {
1600 platform: "fake".to_string(),
1601 chat_id: "chat1".to_string(),
1602 user_id: "u1".to_string(),
1603 callback_query_id: "cbq-1".to_string(),
1604 data: format!("{}:0", parked.nonce),
1605 reply_ctx: reply_ctx.clone(),
1606 };
1607 ConnectBridge::handle_callback(
1608 bridge.clone(),
1609 platform.clone(),
1610 vec!["u1".to_string()],
1611 callback,
1612 )
1613 .await;
1614
1615 wait_for_responder_calls(&responder, 1).await;
1620
1621 resume_tx
1622 .send(AgentEvent::Complete {
1623 usage: bamboo_agent_core::TokenUsage {
1624 prompt_tokens: 1,
1625 completion_tokens: 1,
1626 total_tokens: 2,
1627 },
1628 })
1629 .unwrap();
1630
1631 tokio::time::timeout(Duration::from_secs(5), render_handle)
1632 .await
1633 .expect("render task must finish")
1634 .unwrap();
1635
1636 assert_eq!(
1637 responder.calls.lock().await.as_slice(),
1638 &[("sess-1".to_string(), "Approve".to_string())]
1639 );
1640 assert_eq!(
1641 platform.answered_callbacks.lock().await.as_slice(),
1642 &[("cbq-1".to_string(), None)]
1643 );
1644 assert!(!bridge.has_pending_ask(&key).await);
1645 }
1646
1647 #[tokio::test]
1648 async fn stale_callback_nonce_is_dropped_and_acked_without_resolving() {
1649 let (ctx, _dir) = test_context().await;
1650 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1651 let responder = FakeResponder::new(resume_tx.clone());
1652 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1653 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1654 let key = key_for("chat1", "u1");
1655 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1656
1657 let (tx, rx) = broadcast::channel(16);
1658 tx.send(need_clarification_event(
1659 "Approve?",
1660 vec!["Approve", "Deny"],
1661 false,
1662 ))
1663 .unwrap();
1664
1665 let render_handle = {
1666 let bridge = bridge.clone();
1667 let platform = platform.clone();
1668 let reply_ctx = reply_ctx.clone();
1669 let key = key.clone();
1670 tokio::spawn(async move {
1671 bridge
1672 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1673 .await;
1674 })
1675 };
1676
1677 wait_for_parked_ask(&bridge, &key).await;
1678
1679 let stale_callback = CallbackQuery {
1680 platform: "fake".to_string(),
1681 chat_id: "chat1".to_string(),
1682 user_id: "u1".to_string(),
1683 callback_query_id: "cbq-stale".to_string(),
1684 data: "totally-wrong-nonce:0".to_string(),
1685 reply_ctx: reply_ctx.clone(),
1686 };
1687 ConnectBridge::handle_callback(
1688 bridge.clone(),
1689 platform.clone(),
1690 vec!["u1".to_string()],
1691 stale_callback,
1692 )
1693 .await;
1694
1695 let acked = platform.answered_callbacks.lock().await.clone();
1698 assert_eq!(acked.len(), 1);
1699 assert_eq!(acked[0].0, "cbq-stale");
1700 assert!(acked[0].1.is_some());
1701 assert!(responder.calls.lock().await.is_empty());
1702 assert!(bridge.has_pending_ask(&key).await);
1704
1705 bridge.invalidate_pending_ask(&key).await;
1706 tokio::time::timeout(Duration::from_secs(5), render_handle)
1707 .await
1708 .expect("render task must finish")
1709 .unwrap();
1710 }
1711
1712 #[tokio::test]
1713 async fn text_answer_resolves_an_open_question() {
1714 let (ctx, _dir) = test_context().await;
1715 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1716 let responder = FakeResponder::new(resume_tx.clone());
1717 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1718 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1719 let key = key_for("chat1", "u1");
1720 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1721
1722 let (tx, rx) = broadcast::channel(16);
1723 tx.send(need_clarification_event(
1724 "Anything else?",
1725 vec!["OK", "Need changes"],
1726 true,
1727 ))
1728 .unwrap();
1729
1730 let render_handle = {
1731 let bridge = bridge.clone();
1732 let platform = platform.clone();
1733 let reply_ctx = reply_ctx.clone();
1734 let key = key.clone();
1735 tokio::spawn(async move {
1736 bridge
1737 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1738 .await;
1739 })
1740 };
1741
1742 wait_for_parked_ask(&bridge, &key).await;
1743
1744 ConnectBridge::handle_inbound(
1745 bridge.clone(),
1746 platform.clone(),
1747 vec!["u1".to_string()],
1748 inbound("chat1", "u1", "answer-1", "please also add tests"),
1749 )
1750 .await;
1751
1752 wait_for_responder_calls(&responder, 1).await;
1757
1758 resume_tx
1759 .send(AgentEvent::Complete {
1760 usage: bamboo_agent_core::TokenUsage {
1761 prompt_tokens: 1,
1762 completion_tokens: 1,
1763 total_tokens: 2,
1764 },
1765 })
1766 .unwrap();
1767
1768 tokio::time::timeout(Duration::from_secs(5), render_handle)
1769 .await
1770 .expect("render task must finish")
1771 .unwrap();
1772
1773 assert_eq!(
1774 responder.calls.lock().await.as_slice(),
1775 &[("sess-1".to_string(), "please also add tests".to_string())]
1776 );
1777 }
1778
1779 #[tokio::test]
1780 async fn binary_ask_keyword_mapping_resolves_via_text() {
1781 let (ctx, _dir) = test_context().await;
1782 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1783 let responder = FakeResponder::new(resume_tx.clone());
1784 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1785 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1786 let key = key_for("chat1", "u1");
1787 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1788
1789 let (tx, rx) = broadcast::channel(16);
1790 tx.send(need_clarification_event(
1791 "Approve?",
1792 vec!["Approve", "Deny"],
1793 false,
1794 ))
1795 .unwrap();
1796
1797 let render_handle = {
1798 let bridge = bridge.clone();
1799 let platform = platform.clone();
1800 let reply_ctx = reply_ctx.clone();
1801 let key = key.clone();
1802 tokio::spawn(async move {
1803 bridge
1804 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1805 .await;
1806 })
1807 };
1808
1809 wait_for_parked_ask(&bridge, &key).await;
1810
1811 ConnectBridge::handle_inbound(
1814 bridge.clone(),
1815 platform.clone(),
1816 vec!["u1".to_string()],
1817 inbound("chat1", "u1", "answer-1", "允许"),
1818 )
1819 .await;
1820
1821 wait_for_responder_calls(&responder, 1).await;
1826
1827 resume_tx
1828 .send(AgentEvent::Complete {
1829 usage: bamboo_agent_core::TokenUsage {
1830 prompt_tokens: 1,
1831 completion_tokens: 1,
1832 total_tokens: 2,
1833 },
1834 })
1835 .unwrap();
1836
1837 tokio::time::timeout(Duration::from_secs(5), render_handle)
1838 .await
1839 .expect("render task must finish")
1840 .unwrap();
1841
1842 assert_eq!(
1843 responder.calls.lock().await.as_slice(),
1844 &[("sess-1".to_string(), "Approve".to_string())]
1845 );
1846 }
1847
1848 #[tokio::test]
1849 async fn new_command_invalidates_a_parked_ask_instead_of_answering_it() {
1850 let (ctx, _dir) = test_context().await;
1851 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1852 let responder = FakeResponder::new(resume_tx.clone());
1853 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1854 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1855 let key = key_for("chat1", "u1");
1856 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1857
1858 bridge.set_session_id_for_key(&key, "sess-1").await;
1859
1860 let (tx, rx) = broadcast::channel(16);
1861 tx.send(need_clarification_event(
1862 "Approve?",
1863 vec!["Approve", "Deny"],
1864 false,
1865 ))
1866 .unwrap();
1867
1868 let render_handle = {
1869 let bridge = bridge.clone();
1870 let platform = platform.clone();
1871 let reply_ctx = reply_ctx.clone();
1872 let key = key.clone();
1873 tokio::spawn(async move {
1874 bridge
1875 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1876 .await;
1877 })
1878 };
1879
1880 wait_for_parked_ask(&bridge, &key).await;
1881
1882 ConnectBridge::handle_inbound(
1883 bridge.clone(),
1884 platform.clone(),
1885 vec!["u1".to_string()],
1886 inbound("chat1", "u1", "new-1", "/new"),
1887 )
1888 .await;
1889
1890 tokio::time::timeout(Duration::from_secs(5), render_handle)
1891 .await
1892 .expect("render task must finish")
1893 .unwrap();
1894
1895 assert!(responder.calls.lock().await.is_empty());
1896 assert!(!bridge.has_pending_ask(&key).await);
1897 assert!(bridge.session_id_for_key(&key).await.is_none());
1899 }
1900
1901 #[tokio::test]
1902 async fn respond_error_reports_to_the_chat_without_hanging() {
1903 let (ctx, _dir) = test_context().await;
1904 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1905 let responder = FakeResponder::failing(resume_tx.clone(), "boom");
1906 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1907 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1908 let key = key_for("chat1", "u1");
1909 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1910
1911 let (tx, rx) = broadcast::channel(16);
1912 tx.send(need_clarification_event(
1913 "Approve?",
1914 vec!["Approve", "Deny"],
1915 false,
1916 ))
1917 .unwrap();
1918
1919 let render_handle = {
1920 let bridge = bridge.clone();
1921 let platform = platform.clone();
1922 let reply_ctx = reply_ctx.clone();
1923 let key = key.clone();
1924 tokio::spawn(async move {
1925 bridge
1926 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1927 .await;
1928 })
1929 };
1930
1931 let parked = wait_for_parked_ask(&bridge, &key).await;
1932 let callback = CallbackQuery {
1933 platform: "fake".to_string(),
1934 chat_id: "chat1".to_string(),
1935 user_id: "u1".to_string(),
1936 callback_query_id: "cbq-1".to_string(),
1937 data: format!("{}:0", parked.nonce),
1938 reply_ctx: reply_ctx.clone(),
1939 };
1940 ConnectBridge::handle_callback(
1941 bridge.clone(),
1942 platform.clone(),
1943 vec!["u1".to_string()],
1944 callback,
1945 )
1946 .await;
1947
1948 tokio::time::timeout(Duration::from_secs(5), render_handle)
1949 .await
1950 .expect("render task must finish even when the responder errors")
1951 .unwrap();
1952
1953 let sent = platform.sent_texts().await;
1954 assert!(
1955 sent.iter()
1956 .any(|text| text.contains("Failed to record your answer")),
1957 "expected an error report, got: {sent:?}"
1958 );
1959 }
1960
1961 #[tokio::test]
1966 async fn run_that_pauses_twice_keeps_a_single_status_message() {
1967 let (ctx, _dir) = test_context().await;
1968 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
1969 let responder = FakeResponder::new(resume_tx.clone());
1970 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
1971 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
1972 let key = key_for("chat1", "u1");
1973 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
1974
1975 let (tx, rx) = broadcast::channel(16);
1977 tx.send(AgentEvent::Token {
1978 content: "segment one ".to_string(),
1979 })
1980 .unwrap();
1981 tx.send(need_clarification_event(
1982 "First?",
1983 vec!["Approve", "Deny"],
1984 false,
1985 ))
1986 .unwrap();
1987
1988 let render_handle = {
1989 let bridge = bridge.clone();
1990 let platform = platform.clone();
1991 let reply_ctx = reply_ctx.clone();
1992 let key = key.clone();
1993 tokio::spawn(async move {
1994 bridge
1995 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
1996 .await;
1997 })
1998 };
1999
2000 let first_ask = wait_for_parked_ask(&bridge, &key).await;
2001 ConnectBridge::handle_inbound(
2002 bridge.clone(),
2003 platform.clone(),
2004 vec!["u1".to_string()],
2005 inbound("chat1", "u1", "a1", "1"),
2006 )
2007 .await;
2008 wait_for_responder_calls(&responder, 1).await;
2009
2010 resume_tx
2012 .send(AgentEvent::Token {
2013 content: "segment two ".to_string(),
2014 })
2015 .unwrap();
2016 resume_tx
2017 .send(need_clarification_event(
2018 "Second?",
2019 vec!["Approve", "Deny"],
2020 false,
2021 ))
2022 .unwrap();
2023
2024 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
2027 loop {
2028 let parked = {
2029 bridge
2030 .chat_state
2031 .lock()
2032 .await
2033 .get(&key)
2034 .and_then(|state| state.pending_ask.clone())
2035 };
2036 if let Some(ask) = parked {
2037 if ask.nonce != first_ask.nonce {
2038 break;
2039 }
2040 }
2041 assert!(
2042 tokio::time::Instant::now() < deadline,
2043 "second ask was never parked"
2044 );
2045 tokio::time::sleep(Duration::from_millis(10)).await;
2046 }
2047
2048 ConnectBridge::handle_inbound(
2049 bridge.clone(),
2050 platform.clone(),
2051 vec!["u1".to_string()],
2052 inbound("chat1", "u1", "a2", "1"),
2053 )
2054 .await;
2055 wait_for_responder_calls(&responder, 2).await;
2056
2057 resume_tx
2059 .send(AgentEvent::Token {
2060 content: "segment three".to_string(),
2061 })
2062 .unwrap();
2063 resume_tx
2064 .send(AgentEvent::Complete {
2065 usage: bamboo_agent_core::TokenUsage {
2066 prompt_tokens: 1,
2067 completion_tokens: 1,
2068 total_tokens: 2,
2069 },
2070 })
2071 .unwrap();
2072
2073 tokio::time::timeout(Duration::from_secs(5), render_handle)
2074 .await
2075 .expect("render task must finish")
2076 .unwrap();
2077
2078 let sent = platform.sent_texts().await;
2081 let status_count = sent.iter().filter(|text| text.contains('⏳')).count();
2082 assert_eq!(
2083 status_count, 1,
2084 "expected exactly one status bubble, got: {sent:?}"
2085 );
2086 assert_eq!(sent.len(), 3, "status + 2 asks expected, got: {sent:?}");
2087
2088 let edits = platform.edits.lock().await;
2091 let last = edits.last().expect("expected a final edit");
2092 assert!(last.starts_with('✅'), "final edit not a success: {last}");
2093 assert!(last.contains("segment one"));
2094 assert!(last.contains("segment two"));
2095 assert!(last.contains("segment three"));
2096 }
2097
2098 #[tokio::test]
2103 async fn stop_while_paused_cancels_the_pending_question() {
2104 let (ctx, _dir) = test_context().await;
2105 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
2106 let responder = FakeResponder::new(resume_tx.clone());
2107 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
2108 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
2109 let key = key_for("chat1", "u1");
2110 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
2111
2112 let (tx, rx) = broadcast::channel(16);
2113 tx.send(need_clarification_event(
2114 "Approve?",
2115 vec!["Approve", "Deny"],
2116 false,
2117 ))
2118 .unwrap();
2119
2120 let render_handle = {
2121 let bridge = bridge.clone();
2122 let platform = platform.clone();
2123 let reply_ctx = reply_ctx.clone();
2124 let key = key.clone();
2125 tokio::spawn(async move {
2126 bridge
2127 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
2128 .await;
2129 })
2130 };
2131
2132 wait_for_parked_ask(&bridge, &key).await;
2133
2134 ConnectBridge::handle_inbound(
2135 bridge.clone(),
2136 platform.clone(),
2137 vec!["u1".to_string()],
2138 inbound("chat1", "u1", "stop-1", "/stop"),
2139 )
2140 .await;
2141
2142 tokio::time::timeout(Duration::from_secs(5), render_handle)
2143 .await
2144 .expect("render task must finish after /stop invalidates the ask")
2145 .unwrap();
2146
2147 let sent = platform.sent_texts().await;
2148 assert!(
2149 sent.iter()
2150 .any(|text| text == "Stopped — the pending question was cancelled."),
2151 "expected the pending-question-cancelled reply, got: {sent:?}"
2152 );
2153 assert!(!bridge.has_pending_ask(&key).await);
2154 assert!(responder.calls.lock().await.is_empty());
2155 }
2156
2157 #[tokio::test]
2162 async fn stop_while_paused_with_live_token_cancels_both() {
2163 let (ctx, _dir) = test_context().await;
2164 let resume_tx = broadcast::channel::<AgentEvent>(16).0;
2165 let responder = FakeResponder::new(resume_tx.clone());
2166 let bridge = Arc::new(ConnectBridge::with_responder(ctx, None, responder.clone()));
2167 let platform = FakePlatform::with_capabilities("fake", buttons_and_edit_capabilities());
2168 let key = key_for("chat1", "u1");
2169 let reply_ctx = ReplyCtx(serde_json::json!({ "chat_id": "chat1" }));
2170
2171 let token = CancellationToken::new();
2172 bridge.set_cancel_token(&key, token.clone()).await;
2173
2174 let (tx, rx) = broadcast::channel(16);
2175 tx.send(need_clarification_event(
2176 "Approve?",
2177 vec!["Approve", "Deny"],
2178 false,
2179 ))
2180 .unwrap();
2181
2182 let render_handle = {
2183 let bridge = bridge.clone();
2184 let platform = platform.clone();
2185 let reply_ctx = reply_ctx.clone();
2186 let key = key.clone();
2187 tokio::spawn(async move {
2188 bridge
2189 .render_until_settled(&key, platform, reply_ctx, "sess-1", rx)
2190 .await;
2191 })
2192 };
2193
2194 wait_for_parked_ask(&bridge, &key).await;
2195
2196 ConnectBridge::handle_inbound(
2197 bridge.clone(),
2198 platform.clone(),
2199 vec!["u1".to_string()],
2200 inbound("chat1", "u1", "stop-1", "/stop"),
2201 )
2202 .await;
2203
2204 tokio::time::timeout(Duration::from_secs(5), render_handle)
2205 .await
2206 .expect("render task must finish after /stop")
2207 .unwrap();
2208
2209 assert!(token.is_cancelled(), "cancel token must be cancelled");
2210 assert!(!bridge.has_pending_ask(&key).await);
2211 let sent = platform.sent_texts().await;
2212 assert!(
2213 sent.iter().any(|text| text == "Stopping the current run…"),
2214 "expected the stopping reply, got: {sent:?}"
2215 );
2216 }
2217}