use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock, PoisonError};
use std::time::Duration;
use bevy_ecs::prelude::Resource;
use leviath_core::interaction::{InteractionRequest, InteractionResponse};
use tokio::sync::{Notify, oneshot};
use crate::dynamic_interaction::InteractionBackend;
struct PendingEntry {
agent_id: String,
request: InteractionRequest,
responder: oneshot::Sender<InteractionResponse>,
}
#[derive(Clone, Default, Resource)]
pub struct InteractionHub {
pending: Arc<Mutex<HashMap<String, PendingEntry>>>,
wake: Arc<OnceLock<Arc<Notify>>>,
timeout_secs: Arc<AtomicU64>,
}
pub const DEFAULT_INTERACTION_TIMEOUT_SECS: u64 = 3600;
impl InteractionHub {
pub fn new() -> Self {
Self::default()
}
pub fn attach_wake(&self, wake: Arc<Notify>) {
let _ = self.wake.set(wake);
}
pub fn set_timeout_secs(&self, secs: u64) {
self.timeout_secs.store(secs, Ordering::Relaxed);
}
fn timeout(&self) -> Option<Duration> {
match self.timeout_secs.load(Ordering::Relaxed) {
0 => None,
secs => Some(Duration::from_secs(secs)),
}
}
fn nudge(&self) {
if let Some(wake) = self.wake.get() {
wake.notify_one();
}
}
async fn submit(&self, agent_id: &str, request: InteractionRequest) -> InteractionResponse {
let id = request.id.clone();
let (responder, rx) = oneshot::channel();
self.pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(
id.clone(),
PendingEntry {
agent_id: agent_id.to_string(),
request,
responder,
},
);
self.nudge();
let Some(deadline) = self.timeout() else {
return crate::tool_bridge::off_lane(rx)
.await
.unwrap_or_else(|_| InteractionResponse::text(id, ""));
};
let mut rx = rx;
match crate::tool_bridge::off_lane(tokio::time::timeout(deadline, &mut rx)).await {
Ok(answered) => answered.unwrap_or_else(|_| InteractionResponse::text(id, "")),
Err(_elapsed) => self.expire(agent_id, &id, &mut rx),
}
}
fn expire(
&self,
agent_id: &str,
id: &str,
rx: &mut oneshot::Receiver<InteractionResponse>,
) -> InteractionResponse {
self.pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(id);
if let Ok(answered) = rx.try_recv() {
return answered;
}
tracing::warn!(
agent = %agent_id,
request = %id,
"no answer within the interaction timeout - resolving it as unanswered"
);
self.nudge();
InteractionResponse::text(id, "")
}
pub fn pending(&self) -> Vec<(String, InteractionRequest)> {
self.pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.values()
.map(|e| (e.agent_id.clone(), e.request.clone()))
.collect()
}
pub fn answer(&self, response: InteractionResponse) -> bool {
let entry = self
.pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&response.request_id);
match entry {
Some(entry) => {
let _ = entry.responder.send(response);
self.nudge();
true
}
None => false,
}
}
pub fn cancel(&self, request_id: &str) -> bool {
let removed = self
.pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(request_id)
.is_some();
if removed {
self.nudge();
}
removed
}
pub fn cancel_for_agent(&self, agent_id: &str) -> usize {
let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
let before = pending.len();
pending.retain(|_, entry| entry.agent_id != agent_id);
let removed = before - pending.len();
drop(pending);
if removed > 0 {
self.nudge();
}
removed
}
pub fn backend_for(&self, agent_id: impl Into<String>) -> HubInteractionBackend {
HubInteractionBackend {
hub: self.clone(),
agent_id: agent_id.into(),
}
}
}
#[derive(Clone)]
pub struct HubInteractionBackend {
hub: InteractionHub,
agent_id: String,
}
#[async_trait::async_trait]
impl InteractionBackend for HubInteractionBackend {
async fn ask(&self, request: InteractionRequest) -> InteractionResponse {
self.hub.submit(&self.agent_id, request).await
}
}
#[cfg(test)]
#[cfg(test)]
#[path = "interaction_hub_tests.rs"]
mod tests;
pub struct PromptLane<T> {
pub hub: InteractionHub,
pub outcomes: tokio::sync::mpsc::UnboundedSender<T>,
pub wake: std::sync::Arc<tokio::sync::Notify>,
}