use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};
use parking_lot::Mutex;
use tokio::sync::{broadcast, mpsc};
use tokio_util::sync::CancellationToken;
use zeph_common::SessionId;
use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
use crate::agent::Agent;
use crate::channel::{ChannelMessage, LoopbackChannel, LoopbackEvent, LoopbackHandle};
const SESSION_ACTOR_STACK_SIZE: usize = 8 * 1024 * 1024;
const LOOPBACK_CHANNEL_CAPACITY: usize = 8;
#[derive(Debug)]
pub enum SessionCommand {
Prompt {
text: String,
},
Cancel,
Shutdown,
}
#[derive(Debug, Clone, serde::Serialize)]
#[serde(tag = "type", content = "data", rename_all = "snake_case")]
pub enum SessionOutput {
Token(String),
ToolCall {
tool_name: String,
tool_call_id: String,
},
ToolResult {
tool_name: String,
display: String,
},
TurnComplete,
Error(String),
}
const OUTPUT_CHANNEL_CAPACITY: usize = 256;
pub struct SessionActor;
impl SessionActor {
#[must_use]
pub fn spawn<F>(
supervisor: &TaskSupervisor,
registry: &Arc<LiveSessionRegistry>,
session_id: &SessionId,
build_agent: F,
mailbox_capacity: usize,
) -> (SessionActorHandle, BlockingHandle<()>)
where
F: FnOnce(LoopbackChannel) -> Agent<LoopbackChannel> + Send + 'static,
{
let (cmd_tx, cmd_rx) = mpsc::channel(mailbox_capacity.max(1));
let (tx_out, _first_subscriber) = broadcast::channel(OUTPUT_CHANNEL_CAPACITY);
let tx_out_for_actor = tx_out.clone();
let id_str = session_id.as_str().to_owned();
let session_cancel = CancellationToken::new();
let session_cancel_for_thread = session_cancel.clone();
let session_cancel_for_coordinator = session_cancel.clone();
let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
let thread_name = format!("serve-session-{id_str}");
let thread_session_id = id_str.clone();
let spawn_result = std::thread::Builder::new()
.name(thread_name)
.stack_size(SESSION_ACTOR_STACK_SIZE)
.spawn(move || {
let rt = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(rt) => rt,
Err(e) => {
tracing::error!(
session_id = %thread_session_id,
error = %e,
"failed to build session actor tokio runtime"
);
let _ = done_tx.send(());
return;
}
};
let (channel, handle) = LoopbackChannel::pair(LOOPBACK_CHANNEL_CAPACITY);
let agent = build_agent(channel);
let local = tokio::task::LocalSet::new();
rt.block_on(local.run_until(Self::drive(
agent,
handle,
cmd_rx,
tx_out_for_actor,
session_cancel_for_thread,
)));
let _ = done_tx.send(());
});
if let Err(e) = spawn_result {
tracing::error!(error = %e, "failed to spawn dedicated session actor thread");
}
let name: Arc<str> = Arc::from(format!("serve.session.{id_str}"));
let supervisor_cancel = supervisor.cancellation_token();
let registry_for_coordinator = Arc::clone(registry);
let session_id_for_coordinator = session_id.clone();
let tx_for_coordinator = cmd_tx.clone();
let blocking_handle = supervisor.spawn_oneshot(name, move || {
Self::coordinate(
done_rx,
supervisor_cancel,
session_cancel_for_coordinator,
registry_for_coordinator,
session_id_for_coordinator,
tx_for_coordinator,
)
});
(
SessionActorHandle {
tx: cmd_tx,
tx_out,
last_active: Instant::now(),
cancel: session_cancel,
},
blocking_handle,
)
}
async fn coordinate(
done_rx: tokio::sync::oneshot::Receiver<()>,
supervisor_cancel: CancellationToken,
session_cancel: CancellationToken,
registry: Arc<LiveSessionRegistry>,
session_id: SessionId,
tx: mpsc::Sender<SessionCommand>,
) {
let mut done_rx = done_rx;
tokio::select! {
_ = &mut done_rx => {}
() = supervisor_cancel.cancelled() => {
session_cancel.cancel();
let _ = done_rx.await;
}
}
registry.remove_if_current(&session_id, &tx);
}
async fn drive(
mut agent: Agent<LoopbackChannel>,
handle: LoopbackHandle,
mut cmd_rx: mpsc::Receiver<SessionCommand>,
tx_out: broadcast::Sender<SessionOutput>,
cancel: CancellationToken,
) {
let LoopbackHandle {
input_tx,
mut output_rx,
cancel_signal,
} = handle;
let mut input_tx = Some(input_tx);
let mut agent_run = std::pin::pin!(agent.run());
let mut agent_done = false;
while !agent_done {
tokio::select! {
biased;
result = &mut agent_run => {
agent_done = true;
if let Err(e) = result {
tracing::warn!(error = %e, "session actor: agent run ended with error");
let _ = tx_out.send(SessionOutput::Error(e.to_string()));
}
}
Some(event) = output_rx.recv() => {
if let Some(output) = translate_loopback_event(event) {
let _ = tx_out.send(output);
}
}
() = cancel.cancelled(), if input_tx.is_some() => {
tracing::info!("session actor: supervisor shutdown, flushing and exiting");
input_tx = None;
}
cmd = cmd_rx.recv() => {
match cmd {
Some(SessionCommand::Prompt { text }) => {
if let Some(tx) = &input_tx {
let msg = ChannelMessage {
text,
attachments: Vec::new(),
is_guest_context: false,
is_from_bot: false,
};
tracing::debug!("session actor: forwarding prompt to agent channel");
let _ = tx.send(msg).await;
tracing::debug!("session actor: prompt forwarded");
}
}
Some(SessionCommand::Cancel) => cancel_signal.notify_one(),
Some(SessionCommand::Shutdown) | None => {
input_tx = None;
}
}
}
}
}
while let Ok(event) = output_rx.try_recv() {
if let Some(output) = translate_loopback_event(event) {
let _ = tx_out.send(output);
}
}
}
}
fn translate_loopback_event(event: LoopbackEvent) -> Option<SessionOutput> {
match event {
LoopbackEvent::Chunk(text) | LoopbackEvent::FullMessage(text) => {
Some(SessionOutput::Token(text))
}
LoopbackEvent::Flush => Some(SessionOutput::TurnComplete),
LoopbackEvent::ToolStart(ev) => Some(SessionOutput::ToolCall {
tool_name: ev.tool_name.to_string(),
tool_call_id: ev.tool_call_id,
}),
LoopbackEvent::ToolOutput(ev) => Some(SessionOutput::ToolResult {
tool_name: ev.tool_name.to_string(),
display: ev.display,
}),
_ => None,
}
}
#[derive(Clone)]
pub struct SessionActorHandle {
pub tx: mpsc::Sender<SessionCommand>,
pub tx_out: broadcast::Sender<SessionOutput>,
pub last_active: Instant,
pub cancel: CancellationToken,
}
#[derive(Default)]
pub struct LiveSessionRegistry {
sessions: Mutex<HashMap<SessionId, SessionActorHandle>>,
reactivation_lock: tokio::sync::Mutex<()>,
}
impl LiveSessionRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn get(&self, id: &SessionId) -> Option<SessionActorHandle> {
let mut sessions = self.sessions.lock();
let handle = sessions.get_mut(id)?;
handle.last_active = Instant::now();
Some(handle.clone())
}
pub fn insert(&self, id: SessionId, handle: SessionActorHandle) {
self.sessions.lock().insert(id, handle);
}
#[tracing::instrument(
name = "core.serve.registry.get_or_reactivate",
skip_all,
level = "debug",
fields(session_id = id.as_str())
)]
pub async fn get_or_reactivate<F, Fut>(
&self,
id: &SessionId,
reactivate: F,
) -> Option<SessionActorHandle>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Option<SessionActorHandle>>,
{
if let Some(handle) = self.get(id) {
return Some(handle);
}
let _guard = self.reactivation_lock.lock().await;
if let Some(handle) = self.get(id) {
return Some(handle);
}
reactivate().await
}
pub fn remove_if_current(
&self,
id: &SessionId,
tx: &mpsc::Sender<SessionCommand>,
) -> Option<SessionActorHandle> {
let mut sessions = self.sessions.lock();
if sessions.get(id).is_some_and(|h| h.tx.same_channel(tx)) {
sessions.remove(id)
} else {
None
}
}
pub fn remove(&self, id: &SessionId) -> Option<SessionActorHandle> {
self.sessions.lock().remove(id)
}
#[must_use]
pub fn idle_candidates(&self, ttl: Duration) -> Vec<SessionId> {
let sessions = self.sessions.lock();
let now = Instant::now();
sessions
.iter()
.filter(|(_, handle)| {
handle.tx_out.receiver_count() == 0 && now.duration_since(handle.last_active) >= ttl
})
.map(|(id, _)| id.clone())
.collect()
}
#[must_use]
pub fn ids(&self) -> Vec<SessionId> {
self.sessions.lock().keys().cloned().collect()
}
#[must_use]
pub fn len(&self) -> usize {
self.sessions.lock().len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
fn make_handle() -> SessionActorHandle {
let (tx, _rx) = mpsc::channel(4);
let (tx_out, _sub) = broadcast::channel(4);
SessionActorHandle {
tx,
tx_out,
last_active: Instant::now(),
cancel: CancellationToken::new(),
}
}
#[test]
fn registry_insert_and_get_round_trips() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
registry.insert(id.clone(), make_handle());
assert!(registry.get(&id).is_some());
assert_eq!(registry.len(), 1);
}
#[test]
fn registry_get_missing_returns_none() {
let registry = LiveSessionRegistry::new();
assert!(registry.get(&SessionId::new("nope")).is_none());
}
#[test]
fn registry_remove_drops_entry() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
registry.insert(id.clone(), make_handle());
assert!(registry.remove(&id).is_some());
assert!(registry.get(&id).is_none());
assert!(registry.is_empty());
}
#[test]
fn registry_remove_if_current_drops_matching_entry() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
let handle = make_handle();
let tx = handle.tx.clone();
registry.insert(id.clone(), handle);
assert!(registry.remove_if_current(&id, &tx).is_some());
assert!(registry.get(&id).is_none());
}
#[test]
fn registry_remove_if_current_ignores_stale_tx_after_reactivation() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
let stale_handle = make_handle();
let stale_tx = stale_handle.tx.clone();
registry.insert(id.clone(), stale_handle);
registry.insert(id.clone(), make_handle());
assert!(registry.remove_if_current(&id, &stale_tx).is_none());
assert!(
registry.get(&id).is_some(),
"a stale coordinator must never evict a concurrently-reactivated entry"
);
}
#[tokio::test]
async fn get_or_reactivate_serializes_concurrent_reactivation_for_the_same_id() {
use std::sync::atomic::{AtomicUsize, Ordering};
let registry = Arc::new(LiveSessionRegistry::new());
let id = SessionId::new("race-test");
let reactivate_calls = Arc::new(AtomicUsize::new(0));
let mut tasks = Vec::new();
for _ in 0..8 {
let registry = Arc::clone(®istry);
let id = id.clone();
let reactivate_calls = Arc::clone(&reactivate_calls);
tasks.push(tokio::spawn(async move {
registry
.get_or_reactivate(&id, || {
let registry = Arc::clone(®istry);
let id = id.clone();
let reactivate_calls = Arc::clone(&reactivate_calls);
async move {
tokio::task::yield_now().await;
reactivate_calls.fetch_add(1, Ordering::SeqCst);
let handle = make_handle();
registry.insert(id, handle.clone());
Some(handle)
}
})
.await
}));
}
for task in tasks {
assert!(
task.await.unwrap().is_some(),
"every concurrent caller must resolve to a live handle, win or lose the race"
);
}
assert_eq!(
reactivate_calls.load(Ordering::SeqCst),
1,
"exactly one concurrent caller may run the reactivation closure — a second run means \
two SessionActors would have been spawned over the same durable log (N1)"
);
}
#[test]
fn registry_ids_lists_all_tracked_sessions() {
let registry = LiveSessionRegistry::new();
assert!(registry.ids().is_empty());
registry.insert(SessionId::new("s1"), make_handle());
registry.insert(SessionId::new("s2"), make_handle());
let mut ids: Vec<String> = registry
.ids()
.into_iter()
.map(|id| id.as_str().to_owned())
.collect();
ids.sort();
assert_eq!(ids, vec!["s1".to_owned(), "s2".to_owned()]);
}
#[test]
fn registry_idle_candidates_requires_no_subscribers_and_expired_ttl() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
let mut handle = make_handle();
handle.last_active = Instant::now();
registry.insert(id.clone(), handle);
assert!(registry.idle_candidates(Duration::from_hours(1)).is_empty());
}
#[test]
fn registry_idle_candidates_skips_sessions_with_active_subscribers() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
let mut handle = make_handle();
handle.last_active = Instant::now()
.checked_sub(Duration::from_secs(9999))
.unwrap();
let _subscriber = handle.tx_out.subscribe();
registry.insert(id, handle);
assert!(registry.idle_candidates(Duration::from_secs(1)).is_empty());
}
#[test]
fn registry_idle_candidates_returns_expired_unattached_sessions() {
let registry = LiveSessionRegistry::new();
let id = SessionId::new("s1");
let mut handle = make_handle();
handle.last_active = Instant::now()
.checked_sub(Duration::from_secs(9999))
.unwrap();
registry.insert(id.clone(), handle);
let candidates = registry.idle_candidates(Duration::from_secs(1));
assert_eq!(candidates, vec![id]);
}
#[tokio::test]
async fn session_actor_drive_shuts_down_cleanly_on_command() {
use crate::agent::Agent;
use crate::agent::agent_tests::{MockToolExecutor, create_test_registry, mock_provider};
let (channel, handle) = LoopbackChannel::pair(8);
let provider = mock_provider(vec!["ok".to_owned()]);
let registry = create_test_registry();
let executor = MockToolExecutor::no_tools();
let agent: Agent<LoopbackChannel> =
Agent::new(provider, channel, registry, None, 5, executor);
let (cmd_tx, cmd_rx) = mpsc::channel::<SessionCommand>(4);
let (tx_out, _sub) = broadcast::channel::<SessionOutput>(16);
cmd_tx
.send(SessionCommand::Prompt {
text: "hello".to_owned(),
})
.await
.unwrap();
cmd_tx.send(SessionCommand::Shutdown).await.unwrap();
drop(cmd_tx);
tokio::time::timeout(
Duration::from_secs(10),
Box::pin(SessionActor::drive(
agent,
handle,
cmd_rx,
tx_out,
CancellationToken::new(),
)),
)
.await
.expect("drive must finish within the timeout");
}
#[tokio::test]
async fn session_actor_spawn_runs_on_dedicated_thread_and_shuts_down() {
use crate::agent::Agent;
use crate::agent::agent_tests::{MockToolExecutor, create_test_registry, mock_provider};
let supervisor = TaskSupervisor::new(CancellationToken::new());
let registry = Arc::new(LiveSessionRegistry::new());
let session_id = SessionId::new("spawn-test");
let (handle, blocking_handle) = SessionActor::spawn(
&supervisor,
®istry,
&session_id,
move |channel| {
let provider = mock_provider(vec!["ok".to_owned()]);
let registry = create_test_registry();
let executor = MockToolExecutor::no_tools();
let agent: Agent<LoopbackChannel> =
Agent::new(provider, channel, registry, None, 5, executor);
agent
},
4,
);
handle
.tx
.send(SessionCommand::Prompt {
text: "hello".to_owned(),
})
.await
.unwrap();
handle.tx.send(SessionCommand::Shutdown).await.unwrap();
tokio::time::timeout(Duration::from_secs(10), blocking_handle.join())
.await
.expect("session actor must finish within the timeout")
.expect("session actor task must not panic or be aborted");
assert!(
registry.get(&session_id).is_none(),
"registry entry must be reaped once the session actor's coordinator completes"
);
}
#[tokio::test]
async fn session_actor_handle_cancel_shuts_down_without_supervisor_shutdown() {
use crate::agent::Agent;
use crate::agent::agent_tests::{MockToolExecutor, create_test_registry, mock_provider};
let supervisor = TaskSupervisor::new(CancellationToken::new());
let registry = Arc::new(LiveSessionRegistry::new());
let session_id = SessionId::new("cancel-test");
let (handle, blocking_handle) = SessionActor::spawn(
&supervisor,
®istry,
&session_id,
move |channel| {
let provider = mock_provider(vec!["ok".to_owned()]);
let registry = create_test_registry();
let executor = MockToolExecutor::no_tools();
Agent::new(provider, channel, registry, None, 5, executor)
},
4,
);
handle.cancel.cancel();
tokio::time::timeout(Duration::from_secs(10), blocking_handle.join())
.await
.expect("session actor must finish within the timeout after its own token cancels")
.expect("session actor task must not panic or be aborted");
}
fn sample_rss(sys: &mut sysinfo::System, pid: sysinfo::Pid) -> u64 {
sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[pid]), true);
sys.process(pid).map_or(0, sysinfo::Process::memory)
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "spawns up to 100 real OS threads; run explicitly to verify NFR-P7, e.g. \
`cargo nextest run -p zeph-core -E 'test(nfr_p7)' --run-ignored ignored-only \
--no-capture`"]
async fn nfr_p7_real_agent_idle_session_memory_floor() {
use crate::agent::Agent;
use crate::agent::agent_tests::{MockToolExecutor, create_test_registry, mock_provider};
let supervisor = TaskSupervisor::new(CancellationToken::new());
let registry = Arc::new(LiveSessionRegistry::new());
let pid = sysinfo::get_current_pid().expect("current pid must be resolvable");
let mut sys = sysinfo::System::new();
let skill_registry = Arc::new(parking_lot::RwLock::new(create_test_registry()));
let mut actors = Vec::new();
let mut checkpoints: Vec<(usize, u64)> = Vec::new();
for target in [10usize, 25, 50, 100] {
while actors.len() < target {
let session_id = SessionId::new(format!("nfr-p7-{}", actors.len()));
let shared_registry = Arc::clone(&skill_registry);
let (handle, blocking) = SessionActor::spawn(
&supervisor,
®istry,
&session_id,
move |channel| {
let provider = mock_provider(vec!["ok".to_owned()]);
let embedding_provider = provider.clone();
let executor = MockToolExecutor::no_tools();
Agent::new_with_registry_arc(
provider,
embedding_provider,
channel,
shared_registry,
None,
5,
executor,
)
},
4,
);
actors.push((handle, blocking));
}
tokio::time::sleep(Duration::from_millis(200)).await;
checkpoints.push((target, sample_rss(&mut sys, pid)));
}
for (handle, _) in &actors {
let _ = handle.tx.send(SessionCommand::Shutdown).await;
}
for (_, blocking) in actors {
let _ = tokio::time::timeout(Duration::from_secs(10), blocking.join()).await;
}
let (n_prev, rss_prev) = checkpoints[checkpoints.len() - 2];
let (n_last, rss_last) = checkpoints[checkpoints.len() - 1];
let marginal_per_session = rss_last.saturating_sub(rss_prev) / (n_last - n_prev) as u64;
eprintln!("NFR-P7 real-Agent memory floor checkpoints: {checkpoints:?}");
eprintln!(
"NFR-P7 real-Agent per-session marginal floor (housing + Agent owned state, \
synthetic mock-backed lower bound): {marginal_per_session} bytes ({} KiB)",
marginal_per_session / 1024
);
assert!(
marginal_per_session < 1_048_576,
"NFR-P7 composite budget (informational threshold, see nfr.md's #5840 rationale) \
exceeded: measured {marginal_per_session} bytes/session >= 1 MiB \
(checkpoints: {checkpoints:?})"
);
}
}