1use std::collections::HashMap;
22use std::sync::atomic::{AtomicU64, Ordering};
23use std::sync::{Arc, Mutex, OnceLock, PoisonError};
24use std::time::Duration;
25
26use bevy_ecs::prelude::Resource;
27use leviath_core::interaction::{InteractionRequest, InteractionResponse};
28use tokio::sync::{Notify, oneshot};
29
30use crate::dynamic_interaction::InteractionBackend;
31
32struct PendingEntry {
34 agent_id: String,
36 request: InteractionRequest,
38 responder: oneshot::Sender<InteractionResponse>,
40}
41
42#[derive(Clone, Default, Resource)]
47pub struct InteractionHub {
48 pending: Arc<Mutex<HashMap<String, PendingEntry>>>,
49 wake: Arc<OnceLock<Arc<Notify>>>,
54 timeout_secs: Arc<AtomicU64>,
58}
59
60pub const DEFAULT_INTERACTION_TIMEOUT_SECS: u64 = 3600;
66
67impl InteractionHub {
68 pub fn new() -> Self {
70 Self::default()
71 }
72
73 pub fn attach_wake(&self, wake: Arc<Notify>) {
76 let _ = self.wake.set(wake);
77 }
78
79 pub fn set_timeout_secs(&self, secs: u64) {
85 self.timeout_secs.store(secs, Ordering::Relaxed);
86 }
87
88 fn timeout(&self) -> Option<Duration> {
90 match self.timeout_secs.load(Ordering::Relaxed) {
91 0 => None,
92 secs => Some(Duration::from_secs(secs)),
93 }
94 }
95
96 fn nudge(&self) {
98 if let Some(wake) = self.wake.get() {
99 wake.notify_one();
100 }
101 }
102
103 async fn submit(&self, agent_id: &str, request: InteractionRequest) -> InteractionResponse {
113 let id = request.id.clone();
114 let (responder, rx) = oneshot::channel();
115 self.pending
116 .lock()
117 .unwrap_or_else(PoisonError::into_inner)
118 .insert(
119 id.clone(),
120 PendingEntry {
121 agent_id: agent_id.to_string(),
122 request,
123 responder,
124 },
125 );
126 self.nudge();
129 let Some(deadline) = self.timeout() else {
137 return crate::tool_bridge::off_lane(rx)
138 .await
139 .unwrap_or_else(|_| InteractionResponse::text(id, ""));
140 };
141 let mut rx = rx;
145 match crate::tool_bridge::off_lane(tokio::time::timeout(deadline, &mut rx)).await {
146 Ok(answered) => answered.unwrap_or_else(|_| InteractionResponse::text(id, "")),
147 Err(_elapsed) => self.expire(agent_id, &id, &mut rx),
148 }
149 }
150
151 fn expire(
159 &self,
160 agent_id: &str,
161 id: &str,
162 rx: &mut oneshot::Receiver<InteractionResponse>,
163 ) -> InteractionResponse {
164 self.pending
165 .lock()
166 .unwrap_or_else(PoisonError::into_inner)
167 .remove(id);
168 if let Ok(answered) = rx.try_recv() {
169 return answered;
170 }
171 tracing::warn!(
172 agent = %agent_id,
173 request = %id,
174 "no answer within the interaction timeout - resolving it as unanswered"
175 );
176 self.nudge();
179 InteractionResponse::text(id, "")
180 }
181
182 pub fn pending(&self) -> Vec<(String, InteractionRequest)> {
185 self.pending
186 .lock()
187 .unwrap_or_else(PoisonError::into_inner)
188 .values()
189 .map(|e| (e.agent_id.clone(), e.request.clone()))
190 .collect()
191 }
192
193 pub fn answer(&self, response: InteractionResponse) -> bool {
196 let entry = self
197 .pending
198 .lock()
199 .unwrap_or_else(PoisonError::into_inner)
200 .remove(&response.request_id);
201 match entry {
202 Some(entry) => {
203 let _ = entry.responder.send(response);
206 self.nudge();
209 true
210 }
211 None => false,
212 }
213 }
214
215 pub fn cancel(&self, request_id: &str) -> bool {
218 let removed = self
220 .pending
221 .lock()
222 .unwrap_or_else(PoisonError::into_inner)
223 .remove(request_id)
224 .is_some();
225 if removed {
226 self.nudge();
227 }
228 removed
229 }
230
231 pub fn cancel_for_agent(&self, agent_id: &str) -> usize {
240 let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
242 let before = pending.len();
243 pending.retain(|_, entry| entry.agent_id != agent_id);
244 let removed = before - pending.len();
245 drop(pending);
246 if removed > 0 {
247 self.nudge();
248 }
249 removed
250 }
251
252 pub fn backend_for(&self, agent_id: impl Into<String>) -> HubInteractionBackend {
254 HubInteractionBackend {
255 hub: self.clone(),
256 agent_id: agent_id.into(),
257 }
258 }
259}
260
261#[derive(Clone)]
264pub struct HubInteractionBackend {
265 hub: InteractionHub,
266 agent_id: String,
267}
268
269#[async_trait::async_trait]
270impl InteractionBackend for HubInteractionBackend {
271 async fn ask(&self, request: InteractionRequest) -> InteractionResponse {
272 self.hub.submit(&self.agent_id, request).await
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279
280 fn req(id: &str) -> InteractionRequest {
281 InteractionRequest::free_text(id, "prompt?", "stage", true)
282 }
283
284 async fn settle() {
288 for _ in 0..8 {
289 tokio::task::yield_now().await;
290 }
291 }
292
293 #[test]
294 fn a_poisoned_registry_still_serves_every_other_agent() {
295 let hub = InteractionHub::new();
300 let prev = std::panic::take_hook();
301 std::panic::set_hook(Box::new(|_| {})); let poisoned = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
303 let _guard = hub.pending.lock().expect("fresh lock");
304 panic!("a panic while holding the interaction registry");
305 }));
306 std::panic::set_hook(prev);
307 assert!(poisoned.is_err());
308 assert!(hub.pending.is_poisoned(), "the lock really is poisoned");
309
310 assert!(hub.pending().is_empty());
311 assert!(!hub.cancel("nope"));
312 assert!(!hub.answer(InteractionResponse::text("nope", "x")));
313 }
314
315 #[tokio::test]
316 async fn ask_is_answered_through_the_hub() {
317 let hub = InteractionHub::new();
318 let backend = hub.backend_for("agent-a");
319 let asking = tokio::spawn(async move { backend.ask(req("q1")).await });
320
321 settle().await;
322 let pending = hub.pending();
323 assert_eq!(pending.len(), 1);
324 assert_eq!(pending[0].0, "agent-a");
325 assert_eq!(pending[0].1.id, "q1");
326
327 assert!(hub.answer(InteractionResponse::text("q1", "hello")));
328 let response = asking.await.unwrap();
329 assert_eq!(response.value.as_deref(), Some("hello"));
330 assert!(hub.pending().is_empty());
332 }
333
334 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
341 async fn a_batch_waiting_on_a_prompt_does_not_hold_the_tool_lane() {
342 use crate::tool_bridge::{ToolJob, ToolLane, ToolLaneStats};
343 use bevy_ecs::entity::Entity;
344
345 let hub = InteractionHub::new();
346 let (job_tx, job_rx) = tokio::sync::mpsc::unbounded_channel();
347 let (result_tx, mut results) = tokio::sync::mpsc::unbounded_channel();
348 let stats = Arc::new(ToolLaneStats::new(1));
349 let lane = ToolLane::new(
350 tokio::runtime::Handle::current(),
351 result_tx,
352 Arc::new(Notify::new()),
353 1,
354 stats.clone(),
355 );
356 let serving = lane.serve(job_rx);
357 let submit = |entity: u32, exec: crate::tool_bridge::BoxedToolExec| {
358 stats.enqueued();
359 job_tx
360 .send(ToolJob {
361 entity: Entity::from_raw_u32(entity).expect("a small index is a valid id"),
362 exec,
363 cancel: crate::cancel::CancelToken::new(),
364 })
365 .expect("the lane is serving");
366 };
367
368 let asking = hub.backend_for("agent-a");
370 submit(
371 1,
372 Box::new(move || {
373 Box::pin(async move {
374 let response = asking.ask(req("q1")).await;
375 vec![("q1".to_string(), response.value.unwrap_or_default())]
376 })
377 }),
378 );
379 wait_for_prompt(&hub).await;
380 assert_eq!(stats.parked(), 1, "the asker stepped off the lane");
381
382 let answering = hub.clone();
384 submit(
385 2,
386 Box::new(move || {
387 Box::pin(async move {
388 answering.answer(InteractionResponse::text("q1", "hello"));
389 vec![("answered".to_string(), "ok".to_string())]
390 })
391 }),
392 );
393
394 let mut answers = Vec::new();
395 for _ in 0..2 {
396 let outcome = tokio::time::timeout(std::time::Duration::from_secs(30), results.recv())
397 .await
398 .expect("both batches finished")
399 .expect("an outcome arrived");
400 answers.extend(outcome.results);
401 }
402 answers.sort();
403 assert_eq!(
404 answers,
405 vec![
406 ("answered".to_string(), "ok".to_string()),
407 ("q1".to_string(), "hello".to_string()),
408 ],
409 "the asker got its answer from the batch behind it"
410 );
411
412 drop(job_tx);
413 tokio::time::timeout(std::time::Duration::from_secs(30), serving)
414 .await
415 .expect("the lane drained")
416 .expect("the lane task ended");
417 }
418
419 async fn wait_for_prompt(hub: &InteractionHub) {
423 tokio::time::timeout(std::time::Duration::from_secs(30), async {
424 while hub.pending().is_empty() {
425 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
426 }
427 })
428 .await
429 .expect("the prompt was raised");
430 }
431
432 #[tokio::test]
433 async fn answer_unknown_request_is_false() {
434 let hub = InteractionHub::new();
435 assert!(!hub.answer(InteractionResponse::text("nope", "x")));
436 }
437
438 #[tokio::test]
439 async fn submit_and_answer_nudge_the_attached_wake() {
440 let hub = InteractionHub::new();
441 let wake = Arc::new(Notify::new());
442 hub.attach_wake(wake.clone());
443 hub.attach_wake(Arc::new(Notify::new()));
445
446 let backend = hub.backend_for("agent-a");
447 let asking = tokio::spawn(async move { backend.ask(req("q1")).await });
448 settle().await;
449
450 wake.notified().await;
452
453 assert!(hub.answer(InteractionResponse::text("q1", "hi")));
455 wake.notified().await;
456 assert_eq!(asking.await.unwrap().value.as_deref(), Some("hi"));
457 }
458
459 #[tokio::test]
460 async fn cancel_nudges_the_attached_wake() {
461 let hub = InteractionHub::new();
462 let wake = Arc::new(Notify::new());
463 hub.attach_wake(wake.clone());
464
465 let backend = hub.backend_for("agent-a");
466 let asking = tokio::spawn(async move { backend.ask(req("q2")).await });
467 settle().await;
468 wake.notified().await; assert!(hub.cancel("q2"));
471 wake.notified().await; let _ = asking.await.unwrap();
473 }
474
475 #[tokio::test]
476 async fn cancel_wakes_submit_with_neutral_response() {
477 let hub = InteractionHub::new();
478 let backend = hub.backend_for("agent-a");
479 let asking = tokio::spawn(async move { backend.ask(req("q2")).await });
480
481 settle().await;
482 assert!(hub.cancel("q2"));
483 let response = asking.await.unwrap();
484 assert_eq!(response.request_id, "q2");
485 assert_eq!(response.value.as_deref(), Some("")); assert!(!hub.cancel("q2"));
489 }
490
491 #[tokio::test(start_paused = true)]
494 async fn a_prompt_nobody_answers_is_released_when_the_deadline_passes() {
495 let hub = InteractionHub::new();
499 hub.set_timeout_secs(60);
500 let backend = hub.backend_for("agent-a");
501 let asking = tokio::spawn(async move { backend.ask(req("q1")).await });
502
503 settle().await;
504 assert_eq!(hub.pending().len(), 1, "the prompt is open while it waits");
505
506 let response = asking.await.unwrap();
508 assert_eq!(response.request_id, "q1");
509 assert_eq!(response.value.as_deref(), Some(""));
511 assert_eq!(response.approved, None);
512 assert!(
513 hub.pending().is_empty(),
514 "the expired request is off the open list, so the agent leaves Waiting"
515 );
516 }
517
518 #[tokio::test(start_paused = true)]
519 async fn a_deadline_changes_nothing_for_a_prompt_that_is_answered() {
520 let hub = InteractionHub::new();
523 hub.set_timeout_secs(3600);
524
525 let answered_backend = hub.backend_for("agent-a");
526 let answered = tokio::spawn(async move { answered_backend.ask(req("q1")).await });
527 let cancelled_backend = hub.backend_for("agent-b");
528 let cancelled = tokio::spawn(async move { cancelled_backend.ask(req("q2")).await });
529 settle().await;
530
531 assert!(hub.answer(InteractionResponse::text("q1", "yes, go on")));
532 assert_eq!(answered.await.unwrap().value.as_deref(), Some("yes, go on"));
533
534 assert!(hub.cancel("q2"));
535 assert_eq!(cancelled.await.unwrap().value.as_deref(), Some(""));
536 }
537
538 #[tokio::test(start_paused = true)]
539 async fn a_zero_deadline_waits_for_a_person_however_long_it_takes() {
540 let hub = InteractionHub::new();
543 hub.set_timeout_secs(0);
544 let backend = hub.backend_for("agent-a");
545 let asking = tokio::spawn(async move { backend.ask(req("q1")).await });
546
547 settle().await;
548 tokio::time::advance(Duration::from_secs(86_400)).await;
549 assert_eq!(hub.pending().len(), 1, "a day later, still waiting");
550
551 assert!(hub.answer(InteractionResponse::text("q1", "here I am")));
552 assert_eq!(asking.await.unwrap().value.as_deref(), Some("here I am"));
553 }
554
555 #[tokio::test(start_paused = true)]
556 async fn the_deadline_denies_rather_than_approves() {
557 let hub = InteractionHub::new();
561 hub.set_timeout_secs(30);
562 let backend = hub.backend_for("agent-a");
563 let asking = tokio::spawn(async move {
564 backend
565 .ask(InteractionRequest::tool_approval(
566 "t1",
567 "shell",
568 serde_json::json!({"command": "rm -rf /"}),
569 "implement",
570 ))
571 .await
572 });
573
574 let response = asking.await.unwrap();
575 assert!(!leviath_core::interaction::response_approved(&response));
576 }
577
578 #[tokio::test]
579 async fn an_answer_that_lands_as_the_deadline_passes_still_wins() {
580 let hub = InteractionHub::new();
584 let (responder, mut rx) = oneshot::channel();
585 responder
586 .send(InteractionResponse::text("q1", "approved by hand"))
587 .expect("the receiver is still alive");
588
589 let response = hub.expire("agent-a", "q1", &mut rx);
590 assert_eq!(response.value.as_deref(), Some("approved by hand"));
591 }
592}