108 lines
3.6 KiB
Rust
108 lines
3.6 KiB
Rust
//! Projection socket coverage: snapshots replace state on reconnect and live
|
|
//! deltas patch rows without client-side event folding.
|
|
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use axum::extract::ws::{Message, WebSocketUpgrade};
|
|
use axum::extract::State;
|
|
use axum::response::Response;
|
|
use axum::routing::any;
|
|
use axum::Router;
|
|
use maestro_protocol::{
|
|
ReducedState, StructuredEvent, TranscriptDelta, TranscriptRow, TranscriptRowKind,
|
|
TranscriptSnapshot,
|
|
};
|
|
use maestro_tui::session_task::SessionState;
|
|
use maestro_tui::structured_task::StructuredTask;
|
|
use serde_json::json;
|
|
|
|
async fn structured_handler(
|
|
State(connects): State<Arc<AtomicUsize>>,
|
|
ws: WebSocketUpgrade,
|
|
) -> Response {
|
|
ws.on_upgrade(move |mut socket| async move {
|
|
let attempt = connects.fetch_add(1, Ordering::SeqCst);
|
|
let snapshot = if attempt == 0 {
|
|
TranscriptSnapshot {
|
|
cursor: 1,
|
|
rows: vec![row("old", json!({"content":{"type":"text","text":"old"}}))],
|
|
}
|
|
} else {
|
|
TranscriptSnapshot {
|
|
cursor: 3,
|
|
rows: vec![row("live", json!({"content":{"type":"text","text":"one"}}))],
|
|
}
|
|
};
|
|
send(
|
|
&mut socket,
|
|
StructuredEvent::TranscriptSnapshot { snapshot },
|
|
)
|
|
.await;
|
|
if attempt == 0 {
|
|
let _ = socket.close().await;
|
|
return;
|
|
}
|
|
send(
|
|
&mut socket,
|
|
StructuredEvent::ProjectionUpdate {
|
|
transcript: Some(TranscriptDelta::Patch {
|
|
row: row("live", json!({"content":{"type":"text","text":"two"}})),
|
|
}),
|
|
state: ReducedState {
|
|
model: Some("test-model".into()),
|
|
..ReducedState::default()
|
|
},
|
|
},
|
|
)
|
|
.await;
|
|
tokio::time::sleep(Duration::from_secs(3)).await;
|
|
})
|
|
}
|
|
|
|
fn row(id: &str, data: serde_json::Value) -> TranscriptRow {
|
|
TranscriptRow {
|
|
id: id.into(),
|
|
kind: TranscriptRowKind::AgentText,
|
|
data,
|
|
}
|
|
}
|
|
|
|
async fn send(socket: &mut axum::extract::ws::WebSocket, event: StructuredEvent) {
|
|
socket
|
|
.send(Message::Text(serde_json::to_string(&event).unwrap()))
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn reconnect_snapshot_replaces_then_delta_patches_projection() {
|
|
let connects = Arc::new(AtomicUsize::new(0));
|
|
let app = Router::new()
|
|
.route("/api/structured/x", any(structured_handler))
|
|
.with_state(connects.clone());
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
|
|
|
let handle = StructuredTask::spawn(format!("ws://{addr}/api/structured/x"));
|
|
for _ in 0..250 {
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
let state = handle.state.lock().unwrap().clone();
|
|
if state.reduced.model.as_deref() == Some("test-model") {
|
|
assert_eq!(state.rows.len(), 1);
|
|
assert_eq!(state.rows[0].id, "live");
|
|
assert_eq!(state.rows[0].data["content"]["text"], "two");
|
|
assert!(connects.load(Ordering::SeqCst) >= 2);
|
|
handle.end();
|
|
assert_eq!(handle.connection(), SessionState::Ended);
|
|
return;
|
|
}
|
|
}
|
|
panic!(
|
|
"structured projection never reached reconnect delta; connection={:?}",
|
|
handle.connection()
|
|
);
|
|
}
|