//! Integration tests for klbr-daemon runtime-v2:
//! - WebSocket submit/ack
//! - Commit failure with zero executions
//! - Retry after lost ack deduplicates with same ID
//! - Changed payload under same key returns Conflict
//! - Operator stop cancels in-flight runner
//! - Ping / Pong and Subscription replay
#![forbid(unsafe_code)]
use std::{path::PathBuf, sync::Arc, time::Duration};
use anyhow::Result;
use futures::{SinkExt, StreamExt};
use klbr_core::models::{LlmClient, ModelConfig, ModelsConfig};
use klbr_daemon::{
daemon::{parse_session_id, serve_listener},
ingress::{admit_input, SourceRef},
session_registry::{SessionRegistry, SessionState},
};
use klbr_ipc::v2::{Lane, V2ClientMsg, V2ServerMsg};
use klbr_runtime::{
behavior::Release, frontier::Frontier, protocol::OperationId, store::Store,
worker::WorkerConfig,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::{
net::TcpListener,
sync::{broadcast, watch},
time::timeout,
};
use tokio_tungstenite::{connect_async, tungstenite::Message};
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"))
}
async fn open_store_with_retry(path: &std::path::Path) -> Result<Store> {
for _ in 0..40 {
if let Ok(s) = Store::open(path) {
return Ok(s);
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
Store::open(path)
}
struct TestDaemon {
pub addr: std::net::SocketAddr,
pub store: Store,
pub registry: Arc<SessionRegistry>,
#[allow(dead_code)]
pub v2_event_tx: broadcast::Sender<V2ServerMsg>,
pub shutdown_tx: watch::Sender<bool>,
pub handle: tokio::task::JoinHandle<()>,
}
impl TestDaemon {
pub async fn stop_and_join(self) -> Result<()> {
let _ = self.shutdown_tx.send(true);
self.registry.shutdown_all().await;
let _ = self.handle.await;
Ok(())
}
}
async fn start_test_daemon_with_release(
temp: &tempfile::TempDir,
llm: Option<klbr_core::models::LlmClient>,
policy: Option<klbr_daemon::observation::RecoveryObservationPolicy>,
periodic_interval_ms: Option<usize>,
initial_release: Option<Release>,
) -> Result<TestDaemon> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
let db_path = temp.path().join("daemon_test.sqlite");
let store = open_store_with_retry(&db_path).await?;
let worker_config = WorkerConfig::source_tree(&repo_root());
let mut reg = SessionRegistry::new(store.clone(), worker_config);
if let Some(pol) = policy {
reg = reg.with_recovery_policy(pol);
}
let registry = Arc::new(reg);
let (v2_event_tx, _) = broadcast::channel::<V2ServerMsg>(1024);
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let reg_clone = registry.clone();
let v2_tx_clone = v2_event_tx.clone();
let initial_rel = initial_release.or_else(|| publish_test_release(temp).ok());
let _ = periodic_interval_ms;
let handle = tokio::spawn(async move {
let _ = serve_listener(
listener,
None,
v2_tx_clone,
"test_soul".to_string(),
reg_clone,
llm,
initial_rel,
shutdown_rx,
)
.await;
});
Ok(TestDaemon {
addr,
store,
registry,
v2_event_tx,
shutdown_tx,
handle,
})
}
async fn start_test_daemon_with_options(
temp: &tempfile::TempDir,
llm: Option<klbr_core::models::LlmClient>,
policy: Option<klbr_daemon::observation::RecoveryObservationPolicy>,
periodic_interval_ms: Option<usize>,
) -> Result<TestDaemon> {
start_test_daemon_with_release(temp, llm, policy, periodic_interval_ms, None).await
}
async fn start_test_daemon_with_llm(
temp: &tempfile::TempDir,
llm: Option<klbr_core::models::LlmClient>,
) -> Result<TestDaemon> {
start_test_daemon_with_options(temp, llm, None, None).await
}
async fn start_test_daemon(temp: &tempfile::TempDir) -> Result<TestDaemon> {
start_test_daemon_with_llm(temp, None).await
}
/// Reads the next ADMISSION frame, stepping over the daemon's no-model refusal.
///
/// The admission tests in this file run a service with NO model on purpose: their subject is dedupe
/// keys and conflict handling, not turn execution. A service with no model now refuses to the client
/// (`no_model_configured`) instead of accepting an input that nothing will ever answer, so that frame
/// arrives immediately after every Ack. These tests step over it rather than asserting that a
/// model-less service be silent; the refusal itself is asserted by
/// `test_input_to_a_service_with_no_model_is_refused`, and with a model configured no such frame is
/// sent at all.
async fn next_admission_frame<S, E>(read: &mut S) -> Result<V2ServerMsg>
where
S: futures::Stream<Item = std::result::Result<Message, E>> + Unpin,
E: std::error::Error + Send + Sync + 'static,
{
loop {
let frame = timeout(Duration::from_secs(5), read.next())
.await?
.expect("a response is expected")?;
let Message::Text(text) = frame else {
continue;
};
let parsed: V2ServerMsg = serde_json::from_str(&text)?;
match &parsed {
V2ServerMsg::Error { code, .. } if code == "no_model_configured" => continue,
_ => return Ok(parsed),
}
}
}
/// A replay read to its end: the range it announced, and the sequences it carried.
type Replay = (i64, i64, Option<i64>, Option<i64>);
async fn collect_replay<S, E>(read: &mut S) -> Result<(Replay, Vec<i64>)>
where
S: futures::Stream<Item = std::result::Result<Message, E>> + Unpin,
E: std::error::Error + Send + Sync + 'static,
{
let mut records = Vec::new();
let mut started = None;
loop {
let frame = timeout(Duration::from_secs(5), read.next())
.await?
.expect("a response is expected")?;
let Message::Text(text) = frame else { continue };
match serde_json::from_str::<V2ServerMsg>(&text)? {
V2ServerMsg::SyncStarted {
from_seq,
to_seq,
earliest_seq,
fold_seq,
..
} => {
started = Some((from_seq, to_seq, earliest_seq, fold_seq));
}
V2ServerMsg::Event { record } => records.push(record.seq),
V2ServerMsg::SyncCompleted { .. } => break,
_ => {}
}
}
Ok((
started.expect("a replay says what range it is sending"),
records,
))
}
/// A service with no model must not accept an input in silence.
///
/// The input is admitted - it is durable, and that is not in question - and the client is then told,
/// in the protocol, that nothing will answer it. Before this the trigger was skipped and the client
/// heard nothing at all, which is indistinguishable from a slow turn.
#[tokio::test]
async fn test_input_to_a_service_with_no_model_is_refused() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let (ws_stream, _) = connect_async(format!("ws://{}", daemon.addr)).await?;
let (mut write, mut read) = ws_stream.split();
write
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::SendInput {
session_id: "test_sess_no_model".to_string(),
source: "operator".to_string(),
text: "an input nothing can answer".to_string(),
client_request_id: Some("req-no-model-1".to_string()),
},
)?))
.await?;
let admitted = next_admission_frame(&mut read).await?;
let seq = match admitted {
V2ServerMsg::Ack { seq, duplicate, .. } => {
assert_eq!(duplicate, Some(false), "the input must be admitted first");
seq
}
other => panic!("expected the input to be admitted, got {other:?}"),
};
// The refusal, typed, instead of silence. Read raw: the helper would step over this frame.
let refused = timeout(Duration::from_secs(5), read.next())
.await?
.expect("a refusal is expected")?;
let Message::Text(text) = refused else {
panic!()
};
match serde_json::from_str::<V2ServerMsg>(&text)? {
V2ServerMsg::Error { code, message } => {
assert_eq!(code, "no_model_configured");
assert!(
message.contains("stored but nothing will answer it"),
"the refusal must say the input is stored and unanswered: {message}"
);
}
other => panic!("expected a typed refusal, got {other:?}"),
}
// Nothing ran: the admitted input is still pending, not consumed by a turn that cannot happen.
let s_id = parse_session_id("test_sess_no_model");
assert_eq!(daemon.store.pending(&s_id, 10).await?, vec![seq]);
daemon.stop_and_join().await?;
Ok(())
}
#[tokio::test]
async fn operator_controls_and_status_use_the_v2_socket() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let session_id = "operator-control".to_string();
let release = publish_test_release(&temp)?;
daemon
.registry
.get_or_open(&parse_session_id(&session_id), &release)
.await?;
let (mut socket, _) = connect_async(format!("ws://{}", daemon.addr)).await?;
socket
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::SessionStatus {
session_id: session_id.clone(),
},
)?))
.await?;
let Message::Text(status) = timeout(Duration::from_secs(5), socket.next())
.await?
.expect("status response")?
else {
panic!("status response must be text")
};
assert!(matches!(
serde_json::from_str::<V2ServerMsg>(&status)?,
V2ServerMsg::SessionStatus { .. }
));
for (action, request) in [
(
"reset_kernel",
V2ClientMsg::ResetKernel {
session_id: session_id.clone(),
},
),
(
"stop",
V2ClientMsg::Stop {
session_id: session_id.clone(),
},
),
] {
socket
.send(Message::Text(serde_json::to_string(&request)?))
.await?;
let Message::Text(reply) = timeout(Duration::from_secs(5), socket.next())
.await?
.expect("operator response")?
else {
panic!("operator response must be text")
};
match serde_json::from_str::<V2ServerMsg>(&reply)? {
V2ServerMsg::OperatorAck {
action: received,
success,
..
} => {
assert_eq!(received, action);
assert!(success, "{action} failed");
}
other => panic!("expected operator acknowledgement, got {other:?}"),
}
}
daemon.stop_and_join().await?;
Ok(())
}
#[tokio::test]
async fn offline_model_failure_cannot_recursively_originate_repairs() -> Result<()> {
let temp = tempfile::tempdir()?;
let unused = TcpListener::bind("127.0.0.1:0").await?;
let unavailable_url = format!("http://{}/v1", unused.local_addr()?);
drop(unused);
let mut models = ModelsConfig::default();
models.llm.url = unavailable_url;
models.llm.model = "offline-test".into();
let daemon = start_test_daemon_with_llm(&temp, Some(LlmClient::new(models))).await?;
let session_id = "offline-model".to_string();
let (mut socket, _) = connect_async(format!("ws://{}", daemon.addr)).await?;
socket
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::SendInput {
session_id: session_id.clone(),
source: "operator".into(),
text: "exercise one bounded failure".into(),
client_request_id: Some("offline-once".into()),
},
)?))
.await?;
let sid = parse_session_id(&session_id);
let failures = timeout(Duration::from_secs(5), async {
loop {
let history = daemon.store.history(&sid, 0, 100).await?;
let failures = history
.iter()
.filter(|event| {
event.kind == "turn.failed"
&& event.payload.get("type").and_then(|value| value.as_str())
== Some("turn_failed")
})
.count();
if failures >= 2 && daemon.store.all_pending_work_items().await?.is_empty() {
return Ok::<_, anyhow::Error>(failures);
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await??;
assert_eq!(
failures, 2,
"one input failure plus one bounded repair failure"
);
let input_work = daemon
.store
.get_work_item("work-input-1")
.await?
.expect("input work item");
assert!(
matches!(
input_work.status,
klbr_runtime::work::WorkStatus::Failed { .. }
),
"provider failure must not be overwritten as consumed_by_turn"
);
tokio::time::sleep(Duration::from_millis(250)).await;
let later = daemon
.store
.history(&sid, 0, 100)
.await?
.iter()
.filter(|event| {
event.kind == "turn.failed"
&& event.payload.get("type").and_then(|value| value.as_str()) == Some("turn_failed")
})
.count();
assert_eq!(
later, failures,
"offline model failure recurred without new input"
);
daemon.stop_and_join().await?;
Ok(())
}
/// 1. WebSocket submit/ack:
/// - Client connects over WebSocket.
/// - Submits SendInput with client_request_id.
/// - Server admits and returns Ack { seq, client_request_id, duplicate: false }.
/// - Store durably records the pending inbox item.
#[tokio::test]
async fn test_websocket_submit_and_ack() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
let session_id = "test_sess_ack";
let send_msg = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".to_string(),
text: "hello v2 daemon".to_string(),
client_request_id: Some("req-ack-1".to_string()),
};
write
.send(Message::Text(serde_json::to_string(&send_msg)?))
.await?;
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.expect("websocket message expected")?;
let Message::Text(text) = resp else {
panic!("expected text response");
};
let server_msg: V2ServerMsg = serde_json::from_str(&text)?;
match server_msg {
V2ServerMsg::Ack {
session_id: ack_sess,
seq,
client_request_id,
duplicate,
} => {
assert_eq!(ack_sess, session_id);
assert!(seq > 0, "seq must be positive");
assert_eq!(client_request_id.as_deref(), Some("req-ack-1"));
assert_eq!(duplicate, Some(false));
// Verify store pending inbox item
let s_id = parse_session_id(session_id);
let pending = daemon.store.pending(&s_id, 10).await?;
assert_eq!(pending, vec![seq]);
}
other => panic!("expected Ack message, got: {other:?}"),
}
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 2. Commit failure with zero executions:
/// - Stop/drain the session in registry.
/// - Submit input via WebSocket.
/// - Admission is rejected before execution.
/// - Zero events committed, zero inbox entries created.
#[tokio::test]
async fn test_commit_failure_zero_executions() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let release = publish_test_release(&temp)?;
let session_id = "test_sess_stopped";
let s_id = parse_session_id(session_id);
// Initialize session and then stop it
let _entry = daemon.registry.get_or_open(&s_id, &release).await?;
daemon.registry.stop(&s_id).await?;
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
let send_msg = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".to_string(),
text: "should not be admitted".to_string(),
client_request_id: Some("req-fail-1".to_string()),
};
write
.send(Message::Text(serde_json::to_string(&send_msg)?))
.await?;
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.expect("websocket message expected")?;
let Message::Text(text) = resp else {
panic!("expected text response");
};
let server_msg: V2ServerMsg = serde_json::from_str(&text)?;
match server_msg {
V2ServerMsg::Error { code, message } => {
assert!(
code == "admission_rejected" || code == "admission_error",
"code was {code}"
);
assert!(
message.contains("cannot admit input") || message.contains("draining or stopped"),
"message was {message}"
);
}
other => panic!("expected Error message, got: {other:?}"),
}
// Zero inbox items
let pending = daemon.store.pending(&s_id, 10).await.unwrap_or_default();
assert!(
pending.is_empty(),
"pending inbox must be empty on commit failure"
);
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 3. Retry after lost ack deduplicates with same ID:
/// - Submit SendInput with client_request_id.
/// - Server admits and returns Ack { seq: N, duplicate: false }.
/// - Client resends identical request (simulating lost ack / retry).
/// - Server returns Ack { seq: N, duplicate: true } with identical seq!
/// - Exactly one event and one inbox item in store.
#[tokio::test]
async fn test_retry_after_lost_ack_deduplicates_same_id() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
let session_id = "test_sess_dedupe";
let send_msg = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".to_string(),
text: "idempotent payload".to_string(),
client_request_id: Some("req-dedupe-key-42".to_string()),
};
// First attempt
write
.send(Message::Text(serde_json::to_string(&send_msg)?))
.await?;
let resp1 = timeout(Duration::from_secs(5), read.next())
.await?
.expect("first response expected")?;
let Message::Text(text1) = resp1 else {
panic!()
};
let ack1: V2ServerMsg = serde_json::from_str(&text1)?;
let first_seq = match ack1 {
V2ServerMsg::Ack { seq, duplicate, .. } => {
assert_eq!(duplicate, Some(false));
seq
}
other => panic!("expected Ack, got {other:?}"),
};
// Second attempt (retry after lost ack)
write
.send(Message::Text(serde_json::to_string(&send_msg)?))
.await?;
let ack2 = next_admission_frame(&mut read).await?;
match ack2 {
V2ServerMsg::Ack { seq, duplicate, .. } => {
assert_eq!(seq, first_seq, "retry must return identical seq number");
assert_eq!(duplicate, Some(true), "retry must indicate duplicate: true");
}
other => panic!("expected duplicate Ack, got {other:?}"),
}
// Verify exactly 1 pending item in inbox
let s_id = parse_session_id(session_id);
let pending = daemon.store.pending(&s_id, 10).await?;
assert_eq!(pending, vec![first_seq]);
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 4. Changed payload under same key returns Conflict:
/// - Submit SendInput with client_request_id and text "alpha".
/// - Server admits it.
/// - Submit SendInput with same client_request_id but text "beta".
/// - Server returns V2ServerMsg::Error { code: "conflict", .. }.
#[tokio::test]
async fn test_changed_payload_under_same_key_returns_conflict() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
let session_id = "test_sess_conflict";
let msg1 = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".to_string(),
text: "initial content".to_string(),
client_request_id: Some("req-conflict-key-1".to_string()),
};
write
.send(Message::Text(serde_json::to_string(&msg1)?))
.await?;
let resp1 = timeout(Duration::from_secs(5), read.next())
.await?
.expect("first response expected")?;
let Message::Text(text1) = resp1 else {
panic!()
};
let ack1: V2ServerMsg = serde_json::from_str(&text1)?;
assert!(matches!(ack1, V2ServerMsg::Ack { .. }));
// Send different text under same key
let msg2 = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".to_string(),
text: "tampered / different content".to_string(),
client_request_id: Some("req-conflict-key-1".to_string()),
};
write
.send(Message::Text(serde_json::to_string(&msg2)?))
.await?;
let err2 = next_admission_frame(&mut read).await?;
match err2 {
V2ServerMsg::Error { code, message } => {
assert_eq!(code, "conflict", "expected code == conflict");
assert!(
message.contains("source key reused with different content"),
"message was {message}"
);
}
other => panic!("expected conflict error, got {other:?}"),
}
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 5. Operator stop cancels in-flight runner:
/// - Simulate an active in-flight runner with an abort handle.
/// - Call registry.stop(&session_id).
/// - Runner is immediately cancelled/aborted.
/// - Session transitions to Stopped in registry and Store.
/// - Subsequent admission calls are rejected.
#[tokio::test]
async fn test_operator_stop_cancels_inflight_runner() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let release = publish_test_release(&temp)?;
let session_id = "test_sess_cancel_runner";
let s_id = parse_session_id(session_id);
let entry = daemon.registry.get_or_open(&s_id, &release).await?;
// Simulate in-flight runner
let (abort_handle, abort_reg) = futures::future::AbortHandle::new_pair();
let op_id = OperationId::fresh();
entry.register_runner(abort_handle).await;
entry.transition_to_executing(1, op_id).await?;
let runner_completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
let comp_clone = runner_completed.clone();
let runner_task = tokio::spawn(async move {
let long_work = async {
tokio::time::sleep(Duration::from_secs(30)).await;
comp_clone.store(true, std::sync::atomic::Ordering::SeqCst);
};
let _ = futures::future::Abortable::new(long_work, abort_reg).await;
});
// Operator issues stop
daemon.registry.stop(&s_id).await?;
// Wait briefly for runner task
let _ = timeout(Duration::from_secs(2), runner_task).await?;
// Runner must NOT have completed
assert!(
!runner_completed.load(std::sync::atomic::Ordering::SeqCst),
"In-flight runner must be aborted before completing work"
);
// Verify typestate in registry is Stopped
assert_eq!(entry.state().await, SessionState::Stopped);
// Verify Store state is "stopped"
let store_state = daemon.store.session_state(&s_id).await?;
assert_eq!(
store_state.map(|(s, _)| s),
Some("stopped".to_string()),
"Store state must be stopped"
);
// Subsequent admission rejected
let sref = SourceRef::from_v2_input(session_id, "operator", "post stop", None);
let res = admit_input(&daemon.registry, &s_id, sref).await;
assert!(res.is_err(), "admit_input must fail on stopped session");
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 6. Ping / Pong and Subscription replay:
/// - Client sends Ping { timestamp: 12345 } -> Pong { timestamp: 12345 }.
/// - Client sends Subscribe with after_seq -> receives SyncStarted, Event, SyncCompleted.
#[tokio::test]
async fn test_ping_pong_and_subscribe_replay() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
// 1. Test Ping / Pong
let ping = V2ClientMsg::Ping {
timestamp: Some(12345),
};
write
.send(Message::Text(serde_json::to_string(&ping)?))
.await?;
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.expect("response expected")?;
let Message::Text(text) = resp else { panic!() };
let pong: V2ServerMsg = serde_json::from_str(&text)?;
assert_eq!(
pong,
V2ServerMsg::Pong {
timestamp: Some(12345)
}
);
// 2. Pre-populate store with an admitted event
let session_id = "test_sess_replay";
let s_id = parse_session_id(session_id);
let sref = SourceRef::from_v2_input(session_id, "user", "event 1", Some("req-rep-1".into()));
let outcome = admit_input(&daemon.registry, &s_id, sref).await?;
let seq1 = outcome.seq().unwrap();
// 3. Send Subscribe with after_seq = 0
let sub = V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(0),
limit: None,
};
write
.send(Message::Text(serde_json::to_string(&sub)?))
.await?;
// Receive SyncStarted
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.unwrap()?;
let Message::Text(text) = resp else { panic!() };
let sync_start: V2ServerMsg = serde_json::from_str(&text)?;
assert!(matches!(sync_start, V2ServerMsg::SyncStarted { .. }));
// Receive Event
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.unwrap()?;
let Message::Text(text) = resp else { panic!() };
let ev: V2ServerMsg = serde_json::from_str(&text)?;
match ev {
V2ServerMsg::Event { record } => {
assert_eq!(record.seq, seq1);
assert_eq!(record.session_id, session_id);
}
other => panic!("expected Event, got {other:?}"),
}
// Receive SyncCompleted
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.unwrap()?;
let Message::Text(text) = resp else { panic!() };
let sync_comp: V2ServerMsg = serde_json::from_str(&text)?;
assert!(matches!(sync_comp, V2ServerMsg::SyncCompleted { .. }));
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
/// 7. A window and the history behind it:
/// - a fresh subscribe is a window, bounded by the session's fold and by the `limit` it asks for;
/// - the daemon names the floor it can still serve and the record its fold drained through, with the
/// range rather than at the end of it;
/// - a history request returns the page behind a loaded window, oldest first;
/// - a client that is catching up (`after_seq > 0`) still gets every record it missed.
#[tokio::test]
async fn windowed_replay_reports_the_fold_and_pages_the_history_behind_it() -> Result<()> {
let temp = tempfile::tempdir()?;
let daemon = start_test_daemon(&temp).await?;
let ws_url = format!("ws://{}", daemon.addr);
let session_id = "test_sess_window";
let s_id = parse_session_id(session_id);
let store = daemon.registry.store();
// Thirty records: enough for a window to be a window and for a fold to sit inside the session.
store.ensure_session(&s_id).await?;
for index in 0..30 {
store
.append(
&s_id,
"scratchpad",
serde_json::json!({ "text": format!("record {index}") }),
)
.await?;
}
// A window of ten, in a session of thirty records and no fold yet: the newest ten, and the floor the
// daemon can still serve is the first record it has.
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
write
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(0),
limit: Some(10),
},
)?))
.await?;
let (started, records) = collect_replay(&mut read).await?;
let (from_seq, to_seq, earliest, fold) = started;
assert_eq!(to_seq, 30, "the window ends at the newest record");
assert_eq!(from_seq, 21, "ten records, ending at the newest");
assert_eq!(records, (21..=30).collect::<Vec<_>>());
assert_eq!(
earliest,
Some(1),
"the daemon can still serve the whole session"
);
assert_eq!(
fold, None,
"nothing has been folded, so nothing bounds the window but its size"
);
// A fold: everything through record 25 is carried by the summary from now on, so a window opens above
// it even when the client asks for far more records than the fold leaves.
store
.install_frontier(
&s_id,
&Frontier {
epoch: 1,
drained_through_seq: 25,
summary: "records 1 to 25: the background".to_string(),
derived_by: "test".to_string(),
created_ms: 0,
source_children: Vec::new(),
actor_testimony: None,
actor_testimony_source: None,
},
)
.await?;
write
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(0),
limit: Some(100),
},
)?))
.await?;
let (started, records) = collect_replay(&mut read).await?;
let (from_seq, to_seq, _earliest, fold) = started;
assert_eq!(
from_seq, 26,
"the window opens after the fold, not before it"
);
// The fold is itself a record: the record that says the prefix is summarized is part of the window,
// so the transcript can say why it starts where it does.
assert_eq!(to_seq, 31, "the fold's own record is the newest event");
assert_eq!(records, (26..=31).collect::<Vec<_>>());
assert_eq!(
fold,
Some(25),
"the daemon names the record its fold drained through"
);
// The history behind the window, one page at a time, oldest first.
write
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::History {
session_id: session_id.to_string(),
before_seq: 26,
limit: 5,
},
)?))
.await?;
let (started, records) = collect_replay(&mut read).await?;
let (from_seq, to_seq, earliest, fold) = started;
assert_eq!((from_seq, to_seq), (21, 25));
assert_eq!(records, (21..=25).collect::<Vec<_>>());
assert_eq!(earliest, Some(1));
assert_eq!(fold, Some(25), "a page carries the session's fold with it");
// Catching up is not choosing what to hold: a client with a cursor gets every record it missed.
write
.send(Message::Text(serde_json::to_string(
&V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(26),
limit: Some(1),
},
)?))
.await?;
let (started, records) = collect_replay(&mut read).await?;
let (from_seq, _to_seq, _earliest, _fold) = started;
assert_eq!(
from_seq, 27,
"the gap starts where the client's cursor ends"
);
assert_eq!(
records,
(27..=31).collect::<Vec<_>>(),
"and the limit does not cap a catch-up"
);
let _ = daemon.shutdown_tx.send(true);
Ok(())
}
async fn spawn_mock_llm_server(
should_fail: Arc<std::sync::atomic::AtomicBool>,
) -> Result<std::net::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,
};
let fail = should_fail.load(std::sync::atomic::Ordering::SeqCst);
tokio::spawn(async move {
let mut buf = [0u8; 4096];
let _ = socket.read(&mut buf).await;
if fail {
let resp = "HTTP/1.1 500 Internal Server Error\r\ncontent-type: application/json\r\nconnection: close\r\n\r\n{\"error\":\"model exploded\"}";
let _ = socket.write_all(resp.as_bytes()).await;
} else {
let chunk1 = serde_json::json!({
"choices": [{"delta": {"role": "assistant", "content": "I am thinking in scratchpad."}}]
});
let args_str = serde_json::to_string(&serde_json::json!({
"code": "from klbr import local, runtime\nawait local.send('hello from integration gate')\nawait runtime.wait(seconds=5.0, reason='gate wait')"
})).unwrap();
let chunk2 = serde_json::json!({
"choices": [{"delta": {"tool_calls": [{
"index": 0,
"id": "call_gate_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;
}
});
}
});
Ok(addr)
}
/// 7. Integration Gate: end-to-end multi-lane test with fake provider and WebSocket client
/// Crosses all 4 workstreams:
/// - Workstream A: Daemon admission, WebSocket protocol, trigger_session_turn wiring
/// - Workstream B: Behavior policy active in workbench cell, host verification
/// - Workstream C: Provider stream parsing, scratchpad lane, tool call execution, local.send, runtime.wait
/// - Workstream D: Multi-lane WebSocket broadcast, sparse cursor replay on reconnect
#[tokio::test]
async fn test_end_to_end_v2_cross_workstream_scenario() -> Result<()> {
let temp = tempfile::tempdir()?;
let should_fail = Arc::new(std::sync::atomic::AtomicBool::new(false));
let mock_addr = spawn_mock_llm_server(should_fail.clone()).await?;
let models_cfg = ModelsConfig {
llm: ModelConfig {
url: format!("http://{}", mock_addr),
model: "mock-model".into(),
..Default::default()
},
..Default::default()
};
let llm = LlmClient::new(models_cfg);
let daemon = start_test_daemon_with_llm(&temp, Some(llm)).await?;
let session_id = "gate_test_session";
let s_id = parse_session_id(session_id);
let ws_url = format!("ws://{}", daemon.addr);
let (ws_stream, _) = connect_async(&ws_url).await?;
let (mut write, mut read) = ws_stream.split();
// 1. Subscribe to session
let sub = V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(0),
limit: None,
};
write
.send(Message::Text(serde_json::to_string(&sub)?))
.await?;
// Drain initial sync messages
let mut initial_sync_completed = false;
while let Ok(Some(Ok(Message::Text(text)))) =
timeout(Duration::from_millis(500), read.next()).await
{
if let Ok(V2ServerMsg::SyncCompleted { .. }) = serde_json::from_str(&text) {
initial_sync_completed = true;
break;
}
}
assert!(initial_sync_completed, "Initial sync should complete");
// 2. Submit identified operator input
let send_input = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".into(),
text: "Run integration gate scenario".into(),
client_request_id: Some("req-gate-1".into()),
};
write
.send(Message::Text(serde_json::to_string(&send_input)?))
.await?;
// Receive Ack
let resp = timeout(Duration::from_secs(5), read.next())
.await?
.unwrap()?;
let Message::Text(ack_text) = resp else {
panic!("expected text ack")
};
let ack: V2ServerMsg = serde_json::from_str(&ack_text)?;
assert!(
matches!(ack, V2ServerMsg::Ack { ref client_request_id, duplicate: Some(false), .. } if client_request_id.as_deref() == Some("req-gate-1")),
"Expected clean Ack for req-gate-1, got {ack:?}"
);
// 3. Collect LiveChunks across lanes
let mut saw_scratchpad = false;
let mut saw_speech = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(12);
while tokio::time::Instant::now() < deadline {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if let Ok(Some(Ok(Message::Text(text)))) = timeout(remaining, read.next()).await {
println!("RECV WS MSG: {}", text);
if let Ok(V2ServerMsg::LiveChunk { lane, delta, .. }) =
serde_json::from_str::<V2ServerMsg>(&text)
{
if lane == Lane::Scratchpad && delta.contains("thinking in scratchpad") {
saw_scratchpad = true;
}
if lane == Lane::Speech && delta.contains("hello from integration gate") {
saw_speech = true;
}
}
}
if saw_scratchpad && saw_speech {
break;
}
}
assert!(
saw_scratchpad,
"Must receive Scratchpad LiveChunk with model thoughts"
);
assert!(
saw_speech,
"Must receive Speech LiveChunk with confirmed local.send text"
);
// 4. Invariant checks on Store and SessionRegistry
// Pending inputs must be acknowledged (wait up to 3s for turn loop to commit acknowledgment)
let ack_deadline = tokio::time::Instant::now() + Duration::from_secs(3);
let mut pending_empty = false;
while tokio::time::Instant::now() < ack_deadline {
if daemon.store.pending(&s_id, 10).await?.is_empty() {
pending_empty = true;
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(
pending_empty,
"Admitted input must be acknowledged after turn"
);
// Confirmed speech must exist in store history
let hist = daemon.store.history(&s_id, 0, 50).await?;
println!("STORE HISTORY LEN: {}", hist.len());
for ev in &hist {
println!(" EVENT: kind={}, payload={}", ev.kind, ev.payload);
}
let has_speech_event = hist.iter().any(|ev| {
ev.kind == "local.sent"
&& ev.payload.get("text").and_then(|t| t.as_str())
== Some("hello from integration gate")
});
assert!(
has_speech_event,
"Store must contain durable local.sent event"
);
// Session typestate must be Waiting with reason 'gate wait'
let entry = daemon
.registry
.get_or_open(&s_id, &publish_test_release(&temp)?)
.await?;
let wait_deadline = tokio::time::Instant::now() + Duration::from_secs(3);
let mut is_waiting = false;
while tokio::time::Instant::now() < wait_deadline {
if let SessionState::Waiting { ref reason, .. } = entry.state().await {
assert_eq!(reason.as_deref(), Some("gate wait"));
is_waiting = true;
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(
is_waiting,
"Expected session to be Waiting with reason 'gate wait'"
);
// 5. Reconnect fresh client and verify cursor replay
drop(write);
drop(read);
let (ws_stream_reconnect, _) = connect_async(&ws_url).await?;
let (mut write_re, mut read_re) = ws_stream_reconnect.split();
let sub_re = V2ClientMsg::Subscribe {
session_id: session_id.to_string(),
after_seq: Some(0),
limit: None,
};
write_re
.send(Message::Text(serde_json::to_string(&sub_re)?))
.await?;
let mut replayed_events = Vec::new();
let mut reconnect_sync_completed = false;
while let Ok(Some(Ok(Message::Text(text)))) =
timeout(Duration::from_secs(5), read_re.next()).await
{
if let Ok(msg) = serde_json::from_str::<V2ServerMsg>(&text) {
match msg {
V2ServerMsg::Event { record } => replayed_events.push(record),
V2ServerMsg::SyncCompleted { .. } => {
reconnect_sync_completed = true;
break;
}
_ => {}
}
}
}
assert!(reconnect_sync_completed, "Reconnect sync must complete");
assert!(
replayed_events.len() >= 2,
"Replay must deliver admitted input and confirmed speech events"
);
// Verify monotonic ordering of replayed events
for window in replayed_events.windows(2) {
assert!(
window[0].seq < window[1].seq,
"Replayed sequence must be monotonically increasing"
);
}
// 6. Test failed model request does not manufacture success
should_fail.store(true, std::sync::atomic::Ordering::SeqCst);
let fail_input = V2ClientMsg::SendInput {
session_id: session_id.to_string(),
source: "operator".into(),
text: "This input encounters 500 error".into(),
client_request_id: Some("req-gate-fail".into()),
};
write_re
.send(Message::Text(serde_json::to_string(&fail_input)?))
.await?;
// Ack is admitted
let resp = timeout(Duration::from_secs(5), read_re.next())
.await?
.unwrap()?;
let Message::Text(ack_fail_text) = resp else {
panic!()
};
let ack_fail: V2ServerMsg = serde_json::from_str(&ack_fail_text)?;
assert!(matches!(ack_fail, V2ServerMsg::Ack { .. }));
// Wait for turn runner to process and fail
tokio::time::sleep(Duration::from_millis(500)).await;
// Verify no new local.sent event was fabricated
let hist_after_fail = daemon.store.history(&s_id, 0, 50).await?;
let speech_count = hist_after_fail
.iter()
.filter(|ev| ev.kind == "local.sent")
.count();
assert_eq!(
speech_count, 1,
"Failed model turn must not emit or commit speech events"
);
let _ = daemon.shutdown_tx.send(true);
Ok(())
}