//! Service quiescence vs explicit operator Stop (review 7, R2 / workstream B).
//!
//! The same teardown call (`SessionRegistry::shutdown_all`, the one
//! `serve_listener` invokes) must behave differently by intent:
//! - an ordinary service bounce releases processes but preserves resumable
//! session/work/wait state, including a persisted waiting deadline;
//! - an explicit operator Stop remains persisted across restart.
//!
//! These tests run the real `serve_listener` and the real turn coordinator
//! (`trigger_session_turn`); nothing calls `SessionRuntime::execute` to bypass
//! restored session state.
#![forbid(unsafe_code)]
use std::{
net::SocketAddr,
path::PathBuf,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Duration,
};
use anyhow::Result;
use klbr_core::models::{LlmClient, ModelConfig, ModelsConfig};
use klbr_daemon::{
daemon::{parse_session_id, serve_listener, trigger_session_turn},
ingress::{admit_input, SourceRef},
session_registry::{ServiceTaskGroup, SessionRegistry, SessionState},
};
use klbr_ipc::v2::V2ServerMsg;
use klbr_runtime::{
behavior::Release,
protocol::OperationId,
store::Store,
work::{ConsequenceMetadata, Impact, Provenance, RunCause, WorkItem, WorkStatus},
worker::WorkerConfig,
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
sync::{broadcast, mpsc, watch},
time::timeout,
};
fn repo_root() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.parent()
.expect("manifest has parent repo root")
.to_path_buf()
}
fn publish_test_release(temp: &tempfile::TempDir) -> Result<Release> {
Release::publish(&repo_root().join("behavior"), &temp.path().join("releases"))
}
fn v2_dummy_tx() -> broadcast::Sender<V2ServerMsg> {
broadcast::channel(64).0
}
#[derive(Clone, Copy)]
enum MockMode {
/// normal completion: tool call that sends and then waits on a deadline
Complete,
/// accept the request and never answer: keeps the runner honestly in-flight
Hang,
}
async fn spawn_mock_llm_server(mode: MockMode) -> Result<SocketAddr> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
tokio::spawn(async move {
loop {
let (mut socket, _) = match listener.accept().await {
Ok(conn) => conn,
Err(_) => break,
};
tokio::spawn(async move {
let mut buf = [0_u8; 8192];
let _ = socket.read(&mut buf).await;
match mode {
MockMode::Hang => {
// hold the connection open without any provider event
futures::future::pending::<()>().await;
}
MockMode::Complete => {
let chunk1 = serde_json::json!({
"choices": [{"delta": {"role": "assistant", "content": "working."}}]
});
let args_str = serde_json::to_string(&serde_json::json!({
"code": "from klbr import local, runtime\nawait local.send('quiesce gate output')\nawait runtime.wait(seconds=600.0, reason='quiesce gate wait')"
}))
.unwrap();
let chunk2 = serde_json::json!({
"choices": [{"delta": {"tool_calls": [{
"index": 0,
"id": "call_quiesce_1",
"type": "function",
"function": {
"name": "python",
"arguments": args_str,
}
}]}}]
});
let body = format!(
"data: {}\n\ndata: {}\n\ndata: [DONE]\n\n",
serde_json::to_string(&chunk1).unwrap(),
serde_json::to_string(&chunk2).unwrap()
);
let resp = format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n{}",
body
);
let _ = socket.write_all(resp.as_bytes()).await;
let _ = socket.shutdown().await;
}
}
});
}
});
Ok(addr)
}
fn llm_at(addr: SocketAddr) -> LlmClient {
LlmClient::new(ModelsConfig {
llm: ModelConfig {
url: format!("http://{addr}"),
model: "mock-model".into(),
..Default::default()
},
..Default::default()
})
}
struct TestService {
pub store: Store,
pub registry: Arc<SessionRegistry>,
pub shutdown_tx: watch::Sender<bool>,
pub handle: tokio::task::JoinHandle<()>,
}
async fn start_service(
temp: &tempfile::TempDir,
llm: Option<LlmClient>,
db_path: &std::path::Path,
) -> Result<TestService> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let store = Store::open(db_path)?;
let worker_config = WorkerConfig::source_tree(&repo_root());
let registry = Arc::new(SessionRegistry::new(store.clone(), worker_config));
let (v2_event_tx, _) = broadcast::channel::<V2ServerMsg>(1024);
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let initial_rel = publish_test_release(temp).ok();
let serve_registry = registry.clone();
let handle = tokio::spawn(async move {
let _ = serve_listener(
listener,
None,
v2_event_tx,
"test_soul".to_string(),
serve_registry,
llm,
initial_rel,
shutdown_rx,
)
.await;
});
Ok(TestService {
store,
registry,
shutdown_tx,
handle,
})
}
impl TestService {
/// The owned stop-and-join boundary: signal the listener loop (which calls
/// `shutdown_all` itself), apply the boundary again to prove idempotence,
/// then await the listener task. Returns the registry and its store.
async fn quiesce_and_join(self) -> Result<(Arc<SessionRegistry>, Store)> {
let _ = self.shutdown_tx.send(true);
self.registry.shutdown_all().await;
timeout(Duration::from_secs(30), self.handle).await??;
Ok((self.registry, self.store))
}
}
async fn wait_until(
dur: Duration,
mut check: impl FnMut() -> futures::future::BoxFuture<'static, bool>,
) -> Result<()> {
timeout(dur, async {
loop {
if check().await {
return Ok::<_, anyhow::Error>(());
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await??;
Ok(())
}
/// A waiting session with a deadline survives service quiescence, is restored
/// on reopen, and its deadline fires exactly once through the production
/// scheduler query.
#[tokio::test]
async fn service_quiesce_preserves_waiting_deadline_and_resumes() -> Result<()> {
let temp = tempfile::tempdir()?;
let db_path = temp.path().join("quiesce.sqlite");
let mock_addr = spawn_mock_llm_server(MockMode::Complete).await?;
let llm = llm_at(mock_addr);
let service = start_service(&temp, Some(llm.clone()), &db_path).await?;
let id = parse_session_id("review7_quiesce_wait");
let source = SourceRef::from_v2_input(
id.as_str(),
"operator",
"hello",
Some("quiesce-wait".into()),
);
admit_input(&service.registry, &id, source).await?;
trigger_session_turn(
service.registry.clone(),
llm,
id.clone(),
publish_test_release(&temp)?,
v2_dummy_tx(),
)
.await?;
// the turn ran through the real coordinator and persisted a wait deadline
{
let store = service.store.clone();
let id_watch = id.clone();
wait_until(Duration::from_secs(30), move || {
let store = store.clone();
let id_watch = id_watch.clone();
Box::pin(async move {
matches!(
store.session_state(&id_watch).await.ok().flatten(),
Some((ref st, Some(_))) if st == "waiting"
)
})
})
.await?;
}
let (_, deadline) = service
.store
.session_state(&id)
.await?
.expect("session row exists");
assert!(deadline.is_some(), "runtime.wait must persist a deadline");
// the identical teardown serve_listener performs, applied twice
let (registry, store) = service.quiesce_and_join().await?;
assert_eq!(
store.session_state(&id).await?,
Some(("waiting".into(), deadline)),
"service quiesce must preserve continuation state and deadline"
);
// no new work may be admitted by a quiesced owner
if let Some(entry) = registry.get(&id).await {
let op = OperationId::fresh();
let (h, _reg) = futures::future::AbortHandle::new_pair();
assert!(
!entry.try_start_runner(op, h).await?,
"quiesced session must refuse new runners"
);
}
drop(registry);
drop(store);
// reopen through a fresh registry (no provider needed to inspect state)
let service2 = start_service(&temp, None, &db_path).await?;
let entry2 = service2
.registry
.get_or_open(&id, &publish_test_release(&temp)?)
.await?;
assert_eq!(
entry2.state().await,
SessionState::Waiting {
receipt: None,
wake_at_ms: deadline,
reason: None
},
"restart must restore the waiting session with its deadline"
);
// the deadline fires exactly once through the production scheduler query
let now = deadline.expect("deadline persisted") + 1;
let woken = service2.store.wake_due(now).await?;
assert_eq!(
woken,
vec![id.to_string()],
"the preserved deadline must wake the session exactly once"
);
assert!(
service2.store.wake_due(now + 1).await?.is_empty(),
"a fired deadline must not wake twice"
);
let wake_work: Vec<WorkItem> = service2
.store
.pending_work_items(&id)
.await?
.into_iter()
.filter(|w| matches!(w.cause, RunCause::DeadlineWake { .. }))
.collect();
assert_eq!(
wake_work.len(),
1,
"the wake must produce exactly one runnable deadline cause"
);
service2.registry.shutdown_all().await;
Ok(())
}
/// An explicit operator Stop persists through the same service teardown and a
/// full restart; the restarted service does not silently revive it.
#[tokio::test]
async fn explicit_operator_stop_survives_service_quiesce_and_restart() -> Result<()> {
let temp = tempfile::tempdir()?;
let db_path = temp.path().join("stop.sqlite");
let mock_addr = spawn_mock_llm_server(MockMode::Complete).await?;
let llm = llm_at(mock_addr);
let service = start_service(&temp, Some(llm.clone()), &db_path).await?;
let id = parse_session_id("review7_operator_stop");
let source = SourceRef::from_v2_input(id.as_str(), "operator", "hi", Some("stop-gate".into()));
admit_input(&service.registry, &id, source).await?;
trigger_session_turn(
service.registry.clone(),
llm,
id.clone(),
publish_test_release(&temp)?,
v2_dummy_tx(),
)
.await?;
// explicit operator Stop
service.registry.stop(&id).await?;
assert_eq!(
service.store.session_state(&id).await?,
Some(("stopped".into(), None)),
"operator Stop persists immediately"
);
// the identical teardown the service runs on shutdown
let (registry, store) = service.quiesce_and_join().await?;
assert_eq!(
store.session_state(&id).await?,
Some(("stopped".into(), None)),
"service quiesce must not erase an explicit operator Stop"
);
if let Some(entry) = registry.get(&id).await {
assert_eq!(entry.state().await, SessionState::Stopped);
}
drop(registry);
drop(store);
// restart: the stopped session is restored as Stopped and refuses work
let service2 = start_service(&temp, None, &db_path).await?;
let entry2 = service2
.registry
.get_or_open(&id, &publish_test_release(&temp)?)
.await?;
assert_eq!(entry2.state().await, SessionState::Stopped);
assert!(!entry2.can_admit().await);
let op = OperationId::fresh();
let (h, _reg) = futures::future::AbortHandle::new_pair();
assert!(
entry2.try_start_runner(op, h).await.is_err(),
"stopped session must refuse runners"
);
let late_input =
SourceRef::from_v2_input(id.as_str(), "operator", "no", Some("late-1".to_string()));
assert!(
admit_input(&service2.registry, &id, late_input)
.await
.is_err(),
"admission must refuse a session restored as Stopped"
);
service2.registry.shutdown_all().await;
Ok(())
}
/// Stop stays terminal for every automatic path, but not for the operator: Reset is
/// the only way back, and it opens a fresh runtime rather than reviving the
/// shut-down one.
#[tokio::test]
async fn operator_reset_reopens_a_stopped_session() -> Result<()> {
let temp = tempfile::tempdir()?;
let db_path = temp.path().join("reopen.sqlite");
let service = start_service(&temp, None, &db_path).await?;
let id = parse_session_id("review7_operator_reopen");
let release = publish_test_release(&temp)?;
let source =
SourceRef::from_v2_input(id.as_str(), "operator", "hi", Some("reopen-gate".into()));
admit_input(&service.registry, &id, source).await?;
service.registry.get_or_open(&id, &release).await?;
service.registry.stop(&id).await?;
assert_eq!(
service.store.session_state(&id).await?,
Some(("stopped".into(), None))
);
// Without this, the session is reachable by nothing for the life of the
// database: no input is admitted, so no runner starts, so nothing clears it.
let reopened = service.registry.reopen(&id, &release).await?;
assert_eq!(reopened.state().await, SessionState::Ready);
assert!(reopened.can_admit().await);
assert_eq!(
service.store.session_state(&id).await?,
Some(("ready".into(), None))
);
let back = SourceRef::from_v2_input(id.as_str(), "operator", "back", Some("reopen-2".into()));
admit_input(&service.registry, &id, back).await?;
service.registry.shutdown_all().await;
Ok(())
}
/// After quiescence completes, no previous owner may still hold or claim work:
/// the in-flight attempt is joined, its durable claim is released, and the
/// work is claimable again by a fresh attempt.
#[tokio::test]
async fn no_old_owner_claims_or_delivers_after_quiesce() -> Result<()> {
let temp = tempfile::tempdir()?;
let db_path = temp.path().join("owner.sqlite");
let mock_addr = spawn_mock_llm_server(MockMode::Hang).await?;
let llm = llm_at(mock_addr);
let service = start_service(&temp, Some(llm.clone()), &db_path).await?;
let id = parse_session_id("review7_owner_quiesce");
let release = publish_test_release(&temp)?;
let _entry = service.registry.get_or_open(&id, &release).await?;
// a runnable maintenance work item the coordinator will claim
let work = WorkItem::new(
"review7-owner-work",
"long in-flight maintenance",
RunCause::MaintenanceReady {
item_id: "review7-owner-work".into(),
purpose: "in-flight settle check".into(),
},
ConsequenceMetadata::new(Impact::Normal),
Provenance::agent("service-lifecycle-fixture", 1000),
1000,
)
.with_session(id.to_string());
assert!(service.store.insert_work_item(&work).await?);
trigger_session_turn(
service.registry.clone(),
llm,
id.clone(),
release,
v2_dummy_tx(),
)
.await?;
// wait until the attempt durably owns the work item
let (holder_tx, holder_rx) = watch::channel(None::<OperationId>);
{
let store = service.store.clone();
let work_id = work.id.clone();
wait_until(Duration::from_secs(30), move || {
let store = store.clone();
let holder_tx = holder_tx.clone();
let work_id = work_id.clone();
Box::pin(async move {
if let Ok(Some(item)) = store.get_work_item(&work_id).await {
if let WorkStatus::Claimed { attempt_id, .. } = item.status {
let _ = holder_tx.send(Some(attempt_id));
return true;
}
}
false
})
})
.await?;
}
let attempt = holder_rx.borrow().clone().expect("attempt identity");
assert!(matches!(
service.registry.get(&id).await.unwrap().state().await,
SessionState::Executing { .. }
));
// quiesce aborts and joins the hung owner instead of sleeping
timeout(Duration::from_secs(30), service.registry.shutdown_all())
.await
.expect("shutdown_all must not hang on an in-flight attempt");
let item = service
.store
.get_work_item(&work.id)
.await?
.expect("work item retained");
assert_eq!(
item.status,
WorkStatus::Pending,
"the interrupted attempt's durable claim must be released, not orphaned"
);
// the old owner cannot deliver: no output may land after quiesce
let events = service.store.history(&id, 0, 1024).await?;
assert!(
!events
.iter()
.any(|e| e.kind == "local.sent" && e.payload["text"] == "quiesce gate output"),
"no output may be delivered by an owner that was quiesced"
);
// a fresh attempt can claim the preserved work
let fresh = OperationId::fresh();
let claim = service
.store
.claim_work_item(&work.id, &fresh, 3000)
.await?;
assert!(
claim.is_claimed(),
"preserved work must be claimable by a new owner after quiesce"
);
drop(attempt);
Ok(())
}
/// Component check without a provider: one quiesce pass over mixed intents
/// preserves each session's truthful state (waiting deadline vs explicit Stop)
/// and is stable across repeated teardowns.
#[tokio::test]
async fn quiesce_all_preserves_mixed_intents_idempotently() -> Result<()> {
let temp = tempfile::tempdir()?;
let store = Store::open(temp.path().join("mixed.sqlite"))?;
let registry = Arc::new(SessionRegistry::new(
store.clone(),
WorkerConfig::source_tree(&repo_root()),
));
let waiting = parse_session_id("review7_mixed_wait");
let waiting_entry = registry
.get_or_open(&waiting, &publish_test_release(&temp)?)
.await?;
waiting_entry
.transition_to_executing(
waiting_entry.runtime.hook_epoch().await,
OperationId::fresh(),
)
.await?;
waiting_entry
.transition_to_waiting_with_store(&store, None, Some(90_000), Some("pending check".into()))
.await?;
let stopped = parse_session_id("review7_mixed_stop");
registry
.get_or_open(&stopped, &publish_test_release(&temp)?)
.await?;
registry.stop(&stopped).await?;
// the same call serve_listener performs, applied repeatedly
registry.shutdown_all().await;
registry.shutdown_all().await;
assert_eq!(
store.session_state(&waiting).await?,
Some(("waiting".into(), Some(90_000))),
"waiting deadline must survive repeated service teardowns"
);
assert_eq!(
store.session_state(&stopped).await?,
Some(("stopped".into(), None)),
"explicit Stop must survive repeated service teardowns"
);
// a fresh registry restores both intents faithfully
let registry2 = Arc::new(SessionRegistry::new(
store.clone(),
WorkerConfig::source_tree(&repo_root()),
));
let waiting_restored = registry2
.get_or_open(&waiting, &publish_test_release(&temp)?)
.await?;
assert!(matches!(
waiting_restored.state().await,
SessionState::Waiting {
wake_at_ms: Some(90_000),
..
}
));
let stopped_restored = registry2
.get_or_open(&stopped, &publish_test_release(&temp)?)
.await?;
assert_eq!(stopped_restored.state().await, SessionState::Stopped);
registry2.shutdown_all().await;
Ok(())
}
/// Service-level tasks (driver, recovery scan, client handlers) tracked by the
/// registry are joined by the quiesce boundary: nothing tracked survives it.
#[tokio::test]
async fn tracked_service_tasks_are_joined_by_quiesce() -> Result<()> {
let temp = tempfile::tempdir()?;
let store = Store::open(temp.path().join("tracked.sqlite"))?;
let registry = Arc::new(SessionRegistry::new(
store,
WorkerConfig::source_tree(&repo_root()),
));
let (tx, mut rx) = mpsc::unbounded_channel::<u64>();
let looper = tokio::spawn(async move {
let mut n = 0_u64;
loop {
tokio::time::sleep(Duration::from_millis(10)).await;
n += 1;
if tx.send(n).is_err() {
break;
}
}
});
registry.track_service_task(ServiceTaskGroup::Driver, looper);
// ticks flow before quiesce
assert!(
timeout(Duration::from_secs(5), rx.recv()).await?.is_some(),
"the tracked coordinator must be ticking before quiesce"
);
registry.shutdown_all().await;
// drain whatever was buffered before the abort took effect
while rx.try_recv().is_ok() {}
let ticked = timeout(Duration::from_millis(300), rx.recv()).await;
assert!(
matches!(ticked, Ok(None) | Err(_)),
"no tracked coordinator may tick after quiesce completes"
);
Ok(())
}
/// A runner task must not outlive its retirement, abort, or service quiescence.
#[tokio::test]
async fn no_runner_task_outlives_retirement_abort_or_quiescence() -> Result<()> {
let temp = tempfile::tempdir()?;
let store = Store::open(temp.path().join("runner_leak.sqlite"))?;
let initial = publish_test_release(&temp)?;
let registry = Arc::new(SessionRegistry::new(
store.clone(),
WorkerConfig::source_tree(&repo_root()),
));
let id = parse_session_id("runner_leak_test");
let entry = registry.get_or_open(&id, &initial).await?;
// 1. Retirement must join the runner task
let op1 = OperationId::fresh();
let (abort_h1, _abort_reg1) = futures::future::AbortHandle::new_pair();
assert!(entry.try_start_runner(op1.clone(), abort_h1).await?);
let running1 = Arc::new(AtomicBool::new(true));
let r1_clone = running1.clone();
let task1 = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
r1_clone.store(false, Ordering::SeqCst);
});
assert!(entry.attach_runner_task(&op1, task1).await);
assert!(entry.retire_runner(&op1).await);
assert!(
!running1.load(Ordering::SeqCst),
"runner 1 must not outlive retire_runner"
);
// 2. Abort must join the runner task
let op2 = OperationId::fresh();
let (abort_h2, abort_reg2) = futures::future::AbortHandle::new_pair();
assert!(entry.try_start_runner(op2.clone(), abort_h2).await?);
let running2 = Arc::new(AtomicBool::new(true));
let r2_clone = running2.clone();
let task2 = tokio::spawn(async move {
let _ = futures::future::Abortable::new(
async {
tokio::time::sleep(Duration::from_millis(500)).await;
},
abort_reg2,
)
.await;
r2_clone.store(false, Ordering::SeqCst);
});
assert!(entry.attach_runner_task(&op2, task2).await);
entry.abort_runner().await;
assert!(
!running2.load(Ordering::SeqCst),
"runner 2 must not outlive abort_runner"
);
// 3. Quiesce must join retired/retiring runner tasks
let op3 = OperationId::fresh();
let (abort_h3, _abort_reg3) = futures::future::AbortHandle::new_pair();
assert!(entry.try_start_runner(op3.clone(), abort_h3).await?);
let running3 = Arc::new(AtomicBool::new(true));
let r3_clone = running3.clone();
let e_clone = entry.clone();
let op3_clone = op3.clone();
let task3 = tokio::spawn(async move {
// simulate runner task retiring from inside the task
tokio::time::sleep(Duration::from_millis(20)).await;
e_clone.retire_runner(&op3_clone).await;
// simulate delayed cleanup / thread preemption after retirement
tokio::time::sleep(Duration::from_millis(100)).await;
r3_clone.store(false, Ordering::SeqCst);
});
assert!(entry.attach_runner_task(&op3, task3).await);
// wait until retire_runner ran inside the task
tokio::time::sleep(Duration::from_millis(40)).await;
entry.quiesce(&store).await?;
assert!(
!running3.load(Ordering::SeqCst),
"runner 3 must not outlive quiesce"
);
Ok(())
}