use async_trait::async_trait;
use everruns_core::error::{AgentLoopError, Result};
use everruns_core::platform_store::PlatformMessage;
use everruns_core::session::Session;
use everruns_core::typed_id::{AgentId, HarnessId, SessionId};
use everruns_local::LocalSessionRunner;
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use tokio::sync::mpsc;
pub type WakeSender = mpsc::UnboundedSender<String>;
pub type WakeReceiver = mpsc::UnboundedReceiver<String>;
fn wake_routes() -> &'static Mutex<HashMap<SessionId, WakeSender>> {
static ROUTES: OnceLock<Mutex<HashMap<SessionId, WakeSender>>> = OnceLock::new();
ROUTES.get_or_init(|| Mutex::new(HashMap::new()))
}
pub struct WakeRunner {
session_id: SessionId,
wake_tx: WakeSender,
}
impl WakeRunner {
pub fn new(session_id: SessionId, wake_tx: WakeSender) -> Self {
wake_routes()
.lock()
.unwrap()
.insert(session_id, wake_tx.clone());
Self {
session_id,
wake_tx,
}
}
}
impl Drop for WakeRunner {
fn drop(&mut self) {
let mut routes = wake_routes().lock().unwrap();
if routes
.get(&self.session_id)
.is_some_and(|current| current.same_channel(&self.wake_tx))
{
routes.remove(&self.session_id);
}
}
}
pub fn frame_wake_prompt(message: &str) -> String {
format!(
"[automatic] A background task you started has finished:\n\n{message}\n\nThis is not a \
user message. Read the run's result if useful (its `result_path`/`log_path` are session \
files you can read), then continue the work it was for or report the result."
)
}
#[async_trait]
impl LocalSessionRunner for WakeRunner {
async fn send_message(&self, session_id: SessionId, content: &str) -> Result<()> {
let sender = wake_routes().lock().unwrap().get(&session_id).cloned();
let Some(sender) = sender else {
return Err(AgentLoopError::tool(format!(
"session {session_id} is not active; wake remains retryable"
)));
};
sender.send(content.to_string()).map_err(|_| {
let mut routes = wake_routes().lock().unwrap();
if routes
.get(&session_id)
.is_some_and(|current| current.same_channel(&sender))
{
routes.remove(&session_id);
}
AgentLoopError::tool(format!(
"session {session_id} wake receiver is closed; wake remains retryable"
))
})
}
async fn create_session(
&self,
_harness_id: HarnessId,
_agent_id: Option<AgentId>,
_title: Option<&str>,
_locale: Option<&str>,
_parent_session_id: Option<SessionId>,
) -> Result<Session> {
Err(unsupported("create_session"))
}
async fn list_sessions(
&self,
_limit: Option<usize>,
_agent_id: Option<AgentId>,
) -> Result<Vec<Session>> {
Ok(Vec::new())
}
async fn get_session(&self, _session_id: SessionId) -> Result<Option<Session>> {
Ok(None)
}
async fn get_messages(
&self,
_session_id: SessionId,
_limit: Option<usize>,
) -> Result<Vec<PlatformMessage>> {
Ok(Vec::new())
}
async fn get_session_status(&self, _session_id: SessionId) -> Result<Option<String>> {
Ok(None)
}
}
fn unsupported(op: &str) -> AgentLoopError {
AgentLoopError::tool(format!(
"the local wake runner does not support '{op}' (background completions only)"
))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn routes_wakes_to_the_target_session() {
let session_a = SessionId::from_seed(910_001);
let session_b = SessionId::from_seed(910_002);
let (tx_a, mut rx_a) = mpsc::unbounded_channel();
let (tx_b, mut rx_b) = mpsc::unbounded_channel();
let runner_a = WakeRunner::new(session_a, tx_a);
let _runner_b = WakeRunner::new(session_b, tx_b);
runner_a
.send_message(session_b, "for b")
.await
.expect("route to active session b");
assert_eq!(rx_b.recv().await.as_deref(), Some("for b"));
assert!(rx_a.try_recv().is_err());
}
#[tokio::test]
async fn inactive_session_wakes_remain_retryable() {
let active = SessionId::from_seed(910_003);
let inactive = SessionId::from_seed(910_004);
let (tx, _rx) = mpsc::unbounded_channel();
let runner = WakeRunner::new(active, tx);
let error = runner
.send_message(inactive, "later")
.await
.expect_err("inactive session must reject delivery");
assert!(error.to_string().contains("not active"));
}
}