1use crate::protocol::{ClientMessage, PageKind, PatchOp, ServerMessage, TargetPatch, PROTOCOL_ID};
4use crate::source::{ProjectionUpdate, RefreshOutcome, RunSource};
5use crate::state::reader::ViewerDeltaRead;
6use anyhow::{Context, Result};
7use futures_util::{SinkExt, StreamExt};
8use std::collections::HashMap;
9use std::path::PathBuf;
10use std::sync::Arc;
11use tokio::net::{TcpListener, TcpStream};
12use tokio::sync::{broadcast, Mutex};
13use tokio_tungstenite::tungstenite::Message;
14
15#[derive(Clone, Debug)]
16enum Update {
17 Runs(Vec<serde_json::Value>),
18 Delta {
19 run_id: String,
20 revision: u64,
21 targets: Vec<TargetPatch>,
22 },
23 SnapshotRequired {
24 run_id: String,
25 },
26}
27
28pub struct ServeOptions {
29 pub database_path: PathBuf,
30 pub bind: String,
31}
32
33pub async fn serve(options: ServeOptions) -> Result<()> {
34 let listener = TcpListener::bind(&options.bind)
35 .await
36 .with_context(|| format!("binding {}", options.bind))?;
37 eprintln!(
38 "piw serve: watching {} on ws://{}/ws",
39 options.database_path.display(),
40 listener.local_addr()?
41 );
42 serve_on(listener, options.database_path).await
43}
44
45pub async fn serve_on(listener: TcpListener, database_path: PathBuf) -> Result<()> {
46 let local = listener.local_addr()?;
47 if !local.ip().is_loopback() {
48 anyhow::bail!(
49 "refusing to serve on non-loopback address {local}: the live replay \
50 protocol is unauthenticated; bind to 127.0.0.1 and use an SSH tunnel \
51 for remote access"
52 );
53 }
54 let source = Arc::new(Mutex::new(RunSource::new(&database_path)?));
55 let (updates_tx, _) = broadcast::channel::<Update>(256);
56
57 {
58 let source = Arc::clone(&source);
59 let updates_tx = updates_tx.clone();
60 tokio::spawn(async move {
61 loop {
62 tokio::time::sleep(std::time::Duration::from_millis(250)).await;
63 let Ok(outcome) = run_blocking(&source, RunSource::refresh_all).await else {
64 continue;
65 };
66 broadcast_outcome(&updates_tx, &source, outcome).await;
67 }
68 });
69 }
70
71 loop {
72 let (stream, _addr) = listener.accept().await?;
73 let source = Arc::clone(&source);
74 let updates_rx = updates_tx.subscribe();
75 tokio::spawn(async move {
76 let _ = handle_connection(stream, source, updates_rx).await;
77 });
78 }
79}
80
81async fn broadcast_outcome(
82 updates_tx: &broadcast::Sender<Update>,
83 source: &Arc<Mutex<RunSource>>,
84 outcome: RefreshOutcome,
85) {
86 for ProjectionUpdate { run_id, delta } in outcome.updates {
87 let targets = delta
88 .targets
89 .into_iter()
90 .map(|target| TargetPatch {
91 target_type: target.target_type,
92 target_key: target.target_key,
93 patch: target.patch,
94 })
95 .collect::<Vec<_>>();
96 if targets_are_direct(&targets) {
97 let _ = updates_tx.send(Update::Delta {
98 run_id,
99 revision: delta.revision,
100 targets,
101 });
102 } else {
103 let _ = updates_tx.send(Update::SnapshotRequired { run_id });
104 }
105 }
106 for run_id in outcome.snapshots_required {
107 let _ = updates_tx.send(Update::SnapshotRequired { run_id });
108 }
109 if outcome.listing_changed {
110 let runs = source.lock().await.summaries();
111 let _ = updates_tx.send(Update::Runs(runs));
112 }
113}
114
115fn page_cursors_for_run(
116 page_cursors: &HashMap<(String, PageKind), u64>,
117 run_id: &str,
118) -> Vec<(PageKind, u64)> {
119 let mut cursors = page_cursors
120 .iter()
121 .filter_map(|((candidate, kind), cursor)| (candidate == run_id).then_some((*kind, *cursor)))
122 .collect::<Vec<_>>();
123 cursors.sort_unstable();
124 cursors
125}
126
127fn targets_are_direct(targets: &[TargetPatch]) -> bool {
128 !targets.is_empty()
129 && targets.iter().all(|target| {
130 target.patch.iter().any(|operation| {
131 !matches!(
132 operation,
133 PatchOp::Replace { path, .. }
134 if path == "/presentationRevision" || path == "/graphRevision"
135 )
136 })
137 })
138}
139
140async fn run_blocking<T, F>(source: &Arc<Mutex<RunSource>>, operation: F) -> Result<T>
141where
142 T: Send + 'static,
143 F: FnOnce(&mut RunSource) -> T + Send + 'static,
144{
145 let source = Arc::clone(source);
146 Ok(tokio::task::spawn_blocking(move || {
147 let mut source = source.blocking_lock();
148 operation(&mut source)
149 })
150 .await?)
151}
152
153async fn send(
154 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
155 message: &ServerMessage,
156) -> Result<()> {
157 let text = serde_json::to_string(message)?;
158 sink.send(Message::Text(text.into())).await?;
159 Ok(())
160}
161
162#[allow(clippy::result_large_err)]
163fn reject_browser_origins(
164 request: &tokio_tungstenite::tungstenite::handshake::server::Request,
165 response: tokio_tungstenite::tungstenite::handshake::server::Response,
166) -> Result<
167 tokio_tungstenite::tungstenite::handshake::server::Response,
168 tokio_tungstenite::tungstenite::handshake::server::ErrorResponse,
169> {
170 if request.headers().contains_key("origin") {
171 let mut rejection = tokio_tungstenite::tungstenite::handshake::server::ErrorResponse::new(
172 Some("browser origins are not allowed".to_string()),
173 );
174 *rejection.status_mut() = tokio_tungstenite::tungstenite::http::StatusCode::FORBIDDEN;
175 return Err(rejection);
176 }
177 Ok(response)
178}
179
180async fn handle_connection(
181 stream: TcpStream,
182 source: Arc<Mutex<RunSource>>,
183 mut updates_rx: broadcast::Receiver<Update>,
184) -> Result<()> {
185 let ws = tokio_tungstenite::accept_hdr_async(stream, reject_browser_origins).await?;
186 let (mut sink, mut reads) = ws.split();
187 send(
188 &mut sink,
189 &ServerMessage::Hello {
190 protocol: PROTOCOL_ID.to_string(),
191 },
192 )
193 .await?;
194
195 let mut watching_runs = false;
196 let mut watched: HashMap<String, u64> = HashMap::new();
197 let mut page_cursors: HashMap<(String, PageKind), u64> = HashMap::new();
198
199 let session_result: Result<()> = async {
200 loop {
201 tokio::select! {
202 incoming = reads.next() => {
203 let Some(incoming) = incoming else { break };
204 let message = match incoming {
205 Ok(Message::Text(text)) => text,
206 Ok(Message::Close(_)) => break,
207 Ok(_) => continue,
208 Err(_) => break,
209 };
210 let Ok(request) = serde_json::from_str::<ClientMessage>(&message) else {
211 continue;
212 };
213 match request {
214 ClientMessage::WatchRuns => {
215 watching_runs = true;
216 let runs = source.lock().await.summaries();
217 send(&mut sink, &ServerMessage::Runs { runs }).await?;
218 }
219 ClientMessage::WatchRun {
220 run_id,
221 revision,
222 step_cursor,
223 trace_cursor,
224 session_entry_cursor,
225 session_event_cursor,
226 } => {
227 let first_watch = !watched.contains_key(&run_id);
228 let request_run_id = run_id.clone();
229 let result = run_blocking(&source, move |source| -> Result<_> {
230 if first_watch {
231 source.watch(&request_run_id)?;
232 }
233 let result = (|| -> Result<_> {
234 let resume = revision
235 .map(|cursor| source.deltas_after(&request_run_id, cursor))
236 .transpose()?;
237 let snapshot = source
238 .get(&request_run_id)
239 .map(|entry| (entry.revision, entry.view()));
240 Ok((resume, snapshot))
241 })();
242 if first_watch && result.is_err() {
243 source.unwatch(&request_run_id);
244 }
245 result
246 }).await?;
247 if first_watch && result.is_ok() {
248 watched.insert(run_id.clone(), revision.unwrap_or(0));
249 }
250 let available = match result {
251 Ok((Some(ViewerDeltaRead::Deltas { deltas, current_revision }), snapshot))
252 if revision.is_some() => {
253 let mut sent = revision.unwrap_or(0);
254 for delta in deltas {
255 let targets = delta.targets.into_iter().map(|target| TargetPatch {
256 target_type: target.target_type,
257 target_key: target.target_key,
258 patch: target.patch,
259 }).collect::<Vec<_>>();
260 if !targets_are_direct(&targets) {
261 sent = 0;
262 break;
263 }
264 send(&mut sink, &ServerMessage::RunPatch {
265 run_id: run_id.clone(),
266 revision: delta.revision,
267 targets,
268 }).await?;
269 sent = delta.revision;
270 }
271 watched.insert(run_id.clone(), sent.max(current_revision));
272 if sent < current_revision {
273 if let Some((snapshot_revision, view)) = snapshot {
274 watched.insert(run_id.clone(), snapshot_revision);
275 send(&mut sink, &ServerMessage::RunSnapshot {
276 run_id: run_id.clone(),
277 revision: snapshot_revision,
278 view,
279 }).await?;
280 }
281 }
282 true
283 }
284 Ok((_, Some((snapshot_revision, view)))) => {
285 watched.insert(run_id.clone(), snapshot_revision);
286 send(&mut sink, &ServerMessage::RunSnapshot {
287 run_id: run_id.clone(),
288 revision: snapshot_revision,
289 view,
290 }).await?;
291 true
292 }
293 Ok((_, None)) | Err(_) => {
294 send(&mut sink, &ServerMessage::Error {
295 message: format!("run {run_id} is unavailable"),
296 run_id: Some(run_id.clone()),
297 }).await?;
298 false
299 }
300 };
301 if available {
302 for (kind, cursor) in [
303 (PageKind::Steps, step_cursor),
304 (PageKind::Trace, trace_cursor),
305 (PageKind::SessionEntries, session_entry_cursor),
306 (PageKind::SessionEvents, session_event_cursor),
307 ] {
308 if let Some(cursor) = cursor {
309 page_cursors.insert((run_id.clone(), kind), cursor);
310 send_projection_page(
311 &mut sink,
312 &source,
313 run_id.clone(),
314 kind,
315 cursor,
316 )
317 .await?;
318 }
319 }
320 }
321 }
322 ClientMessage::UnwatchRun { run_id } => {
323 if watched.remove(&run_id).is_some() {
324 page_cursors.retain(|(candidate, _), _| candidate != &run_id);
325 let remove_id = run_id.clone();
326 let _ = run_blocking(&source, move |source| source.unwatch(&remove_id)).await;
327 }
328 }
329 ClientMessage::FetchPage { run_id, kind, cursor } => {
330 if watched.contains_key(&run_id) {
331 page_cursors.insert((run_id.clone(), kind), cursor);
332 send_projection_page(&mut sink, &source, run_id, kind, cursor).await?;
333 }
334 }
335 ClientMessage::FetchArtifact { run_id, path } => {
336 send(&mut sink, &ServerMessage::Error {
337 message: format!("artifact {path} not available; values are stored in SQLite"),
338 run_id: Some(run_id),
339 }).await?;
340 }
341 }
342 }
343 update = updates_rx.recv() => {
344 match update {
345 Ok(Update::Runs(runs)) if watching_runs => {
346 send(&mut sink, &ServerMessage::Runs { runs }).await?;
347 }
348 Ok(Update::Runs(_)) => {}
349 Ok(Update::Delta { run_id, revision, targets }) => {
350 let Some(&last) = watched.get(&run_id) else { continue };
351 if revision == last + 1 {
352 watched.insert(run_id.clone(), revision);
353 send(&mut sink, &ServerMessage::RunPatch { run_id, revision, targets }).await?;
354 } else if revision > last {
355 send_snapshot(
356 &mut sink,
357 &source,
358 &mut watched,
359 &page_cursors,
360 run_id,
361 )
362 .await?;
363 }
364 }
365 Ok(Update::SnapshotRequired { run_id }) => {
366 if watched.contains_key(&run_id) {
367 send_snapshot(
368 &mut sink,
369 &source,
370 &mut watched,
371 &page_cursors,
372 run_id,
373 )
374 .await?;
375 }
376 }
377 Err(broadcast::error::RecvError::Lagged(_)) => {
378 let run_ids: Vec<String> = watched.keys().cloned().collect();
379 for run_id in run_ids {
380 send_snapshot(
381 &mut sink,
382 &source,
383 &mut watched,
384 &page_cursors,
385 run_id,
386 )
387 .await?;
388 }
389 }
390 Err(broadcast::error::RecvError::Closed) => break,
391 }
392 }
393 }
394 }
395 Ok(())
396 }.await;
397
398 let watched_ids: Vec<String> = watched.into_keys().collect();
399 let _ = run_blocking(&source, move |source| {
400 for run_id in watched_ids {
401 source.unwatch(&run_id);
402 }
403 })
404 .await;
405 session_result
406}
407
408async fn send_projection_page(
409 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
410 source: &Arc<Mutex<RunSource>>,
411 run_id: String,
412 kind: PageKind,
413 cursor: u64,
414) -> Result<()> {
415 let request_id = run_id.clone();
416 let page = run_blocking(source, move |source| source.page(&request_id, kind, cursor)).await?;
417 match page {
418 Ok((revision, page)) => {
419 let graph_steps = page
420 .graph_steps
421 .map(|steps| {
422 steps
423 .into_iter()
424 .map(serde_json::to_value)
425 .collect::<Result<Vec<_>, _>>()
426 })
427 .transpose()?;
428 send(
429 sink,
430 &ServerMessage::RunPage {
431 run_id,
432 revision,
433 kind,
434 cursor,
435 start: page.start,
436 total: page.total,
437 items: page.items,
438 graph_cursor: page.graph_cursor,
439 graph_steps,
440 taken_transitions: page.taken_transitions,
441 replay_checkpoint: page.replay_checkpoint,
442 },
443 )
444 .await?;
445 }
446 Err(_) => {
447 send(
448 sink,
449 &ServerMessage::Error {
450 message: "run page is unavailable".to_string(),
451 run_id: Some(run_id),
452 },
453 )
454 .await?;
455 }
456 }
457 Ok(())
458}
459
460async fn send_snapshot(
461 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
462 source: &Arc<Mutex<RunSource>>,
463 watched: &mut HashMap<String, u64>,
464 page_cursors: &HashMap<(String, PageKind), u64>,
465 run_id: String,
466) -> Result<()> {
467 let snapshot_id = run_id.clone();
468 let snapshot = run_blocking(source, move |source| {
469 source
470 .get(&snapshot_id)
471 .map(|entry| (entry.revision, entry.view()))
472 })
473 .await?;
474 if let Some((revision, view)) = snapshot {
475 watched.insert(run_id.clone(), revision);
476 send(
477 sink,
478 &ServerMessage::RunSnapshot {
479 run_id: run_id.clone(),
480 revision,
481 view,
482 },
483 )
484 .await?;
485 for (kind, cursor) in page_cursors_for_run(page_cursors, &run_id) {
486 send_projection_page(sink, source, run_id.clone(), kind, cursor).await?;
487 }
488 }
489 Ok(())
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495 use serde_json::json;
496
497 #[test]
498 fn recovery_keeps_each_run_page_cursor() {
499 let cursors = HashMap::from([
500 (("run-1".to_string(), PageKind::Trace), 40),
501 (("run-1".to_string(), PageKind::SessionEvents), 80),
502 (("run-2".to_string(), PageKind::Trace), 120),
503 ]);
504 assert_eq!(
505 page_cursors_for_run(&cursors, "run-1"),
506 vec![(PageKind::Trace, 40), (PageKind::SessionEvents, 80)]
507 );
508 }
509
510 #[test]
511 fn sends_only_complete_direct_target_patches() {
512 let revision_only = TargetPatch {
513 target_type: "graph".to_string(),
514 target_key: String::new(),
515 patch: vec![PatchOp::Replace {
516 path: "/presentationRevision".to_string(),
517 value: json!(2),
518 }],
519 };
520 assert!(!targets_are_direct(&[revision_only]));
521
522 let tail = TargetPatch {
523 target_type: "conversation".to_string(),
524 target_key: "entries:tail".to_string(),
525 patch: vec![
526 PatchOp::Replace {
527 path: "/presentationRevision".to_string(),
528 value: json!(2),
529 },
530 PatchOp::Append {
531 path: "/items".to_string(),
532 value: vec![json!({"seq": 1})],
533 },
534 ],
535 };
536 assert!(targets_are_direct(&[tail]));
537 }
538}