1use crate::layout::{layout_graph, GraphLayout};
6use crate::protocol::{
7 apply_patch, ClientMessage, PageKind, ServerMessage, TargetPatch, PROTOCOL_ID,
8};
9use crate::state::types::{DefinitionSnapshot, Manifest, RunState, StepRecord};
10use anyhow::{Context, Result};
11use futures_util::{SinkExt, StreamExt};
12use serde_json::Value;
13use std::collections::{HashMap, HashSet};
14use std::sync::{Arc, Mutex};
15use std::thread::JoinHandle;
16use std::time::Duration;
17use tokio::sync::mpsc;
18use tokio_tungstenite::tungstenite::Message;
19
20pub struct RemoteView {
21 pub revision: u64,
22 pub graph_revision: u64,
23 generation: u64,
24 pub manifest: Manifest,
25 pub state: RunState,
26 pub graph_steps: Vec<StepRecord>,
27 pub taken_transitions: Vec<String>,
28 pub graph_cursor: u64,
29 pub step_start: u64,
30 pub step_total: u64,
31 pub snapshot: Option<DefinitionSnapshot>,
32 pub graph_layout: Option<GraphLayout>,
33 pub events: Vec<Value>,
34 pub trace_start: u64,
35 pub trace_total: u64,
36 pub session_binding: Option<Value>,
37 pub session_entries: Vec<Value>,
38 pub session_entry_start: u64,
39 pub session_entry_total: u64,
40 pub session_events: Vec<Value>,
41 pub session_event_start: u64,
42 pub session_event_total: u64,
43 pub session_events_malformed: bool,
44 pub session_events_torn_tail: bool,
45 pub session_capture: Option<Value>,
46 pub session_replay_checkpoint: Option<Value>,
47 pub settings_scopes: Vec<Value>,
48 pub settings_start: u64,
49 pub settings_total: u64,
50 pub follow_up_queue: Option<Value>,
51 pub follow_up_start: u64,
52 pub follow_up_total: u64,
53 pub update_start: u64,
54 pub update_total: u64,
55 pub live: bool,
56 pub possibly_interrupted: bool,
57}
58
59fn decode_view(revision: u64, generation: u64, raw: &Value) -> Option<RemoteView> {
60 let graph_revision = raw
61 .get("graphRevision")
62 .and_then(Value::as_u64)
63 .unwrap_or(revision);
64 let manifest: Manifest = serde_json::from_value(raw.get("manifest")?.clone()).ok()?;
65 let state: RunState = serde_json::from_value(raw.get("state")?.clone()).ok()?;
66 let graph_steps = raw
67 .get("graphSteps")
68 .and_then(|value| serde_json::from_value(value.clone()).ok())
69 .unwrap_or_else(|| state.steps.clone());
70 let taken_transitions = raw
71 .get("takenTransitions")
72 .and_then(|value| serde_json::from_value(value.clone()).ok())
73 .unwrap_or_default();
74 let graph_cursor = raw
75 .get("graphCursor")
76 .and_then(Value::as_u64)
77 .unwrap_or_else(|| state.steps.len().saturating_sub(1) as u64);
78 let step_start = raw.get("stepStart").and_then(Value::as_u64).unwrap_or(0);
79 let step_total = raw
80 .get("stepTotal")
81 .and_then(Value::as_u64)
82 .unwrap_or(state.steps.len() as u64);
83 let snapshot: Option<DefinitionSnapshot> = raw
84 .get("workflow")
85 .and_then(|value| serde_json::from_value(value.clone()).ok());
86 let graph_layout = raw
87 .get("graphScene")
88 .and_then(|value| serde_json::from_value(value.clone()).ok())
89 .or_else(|| snapshot.as_ref().map(layout_graph));
90 let events = raw
91 .pointer("/tracePage/items")
92 .and_then(Value::as_array)
93 .cloned()
94 .unwrap_or_default();
95 let trace_start = raw
96 .pointer("/tracePage/start")
97 .and_then(Value::as_u64)
98 .unwrap_or(0);
99 let trace_total = raw
100 .pointer("/tracePage/total")
101 .and_then(Value::as_u64)
102 .unwrap_or(events.len() as u64);
103 let session_binding = raw.pointer("/session/binding").cloned();
104 let session_entries = raw
105 .pointer("/session/entryPage/items")
106 .and_then(Value::as_array)
107 .cloned()
108 .unwrap_or_default();
109 let session_entry_start = raw
110 .pointer("/session/entryPage/start")
111 .and_then(Value::as_u64)
112 .unwrap_or(0);
113 let session_entry_total = raw
114 .pointer("/session/entryPage/total")
115 .and_then(Value::as_u64)
116 .unwrap_or(session_entries.len() as u64);
117 let session_events = raw
118 .pointer("/session/eventPage/items")
119 .and_then(Value::as_array)
120 .cloned()
121 .unwrap_or_default();
122 let session_event_start = raw
123 .pointer("/session/eventPage/start")
124 .and_then(Value::as_u64)
125 .unwrap_or(0);
126 let session_event_total = raw
127 .pointer("/session/eventPage/total")
128 .and_then(Value::as_u64)
129 .unwrap_or(session_events.len() as u64);
130 let session_events_malformed = raw
131 .pointer("/session/eventsMalformed")
132 .and_then(Value::as_bool)
133 .unwrap_or(false);
134 let session_events_torn_tail = raw
135 .pointer("/session/eventsTornTail")
136 .and_then(Value::as_bool)
137 .unwrap_or(false);
138 let session_capture = raw.pointer("/session/capture").cloned();
139 let session_replay_checkpoint = raw
140 .pointer("/session/replayCheckpoint")
141 .cloned()
142 .filter(|value| !value.is_null());
143 let settings_scopes = raw
144 .get("settingsScopes")
145 .and_then(Value::as_array)
146 .cloned()
147 .unwrap_or_default();
148 let settings_start = raw
149 .get("settingsStart")
150 .and_then(Value::as_u64)
151 .unwrap_or(0);
152 let settings_total = raw
153 .get("settingsTotal")
154 .and_then(Value::as_u64)
155 .unwrap_or(settings_scopes.len() as u64);
156 let follow_up_queue = raw
157 .get("followUpQueue")
158 .cloned()
159 .filter(|value| !value.is_null());
160 let follow_up_start = raw
161 .get("followUpStart")
162 .and_then(Value::as_u64)
163 .unwrap_or(0);
164 let follow_up_total = raw
165 .get("followUpTotal")
166 .and_then(Value::as_u64)
167 .unwrap_or_else(|| {
168 follow_up_queue
169 .as_ref()
170 .and_then(|queue| queue.get("items"))
171 .and_then(Value::as_array)
172 .map_or(0, |items| items.len() as u64)
173 });
174 let update_start = raw.get("updateStart").and_then(Value::as_u64).unwrap_or(0);
175 let update_total = raw
176 .get("updateTotal")
177 .and_then(Value::as_u64)
178 .unwrap_or_else(|| {
179 state
180 .updates
181 .as_ref()
182 .map_or(0, |updates| updates.len() as u64)
183 });
184 Some(RemoteView {
185 revision,
186 graph_revision,
187 generation,
188 manifest,
189 state,
190 graph_steps,
191 taken_transitions,
192 graph_cursor,
193 step_start,
194 step_total,
195 snapshot,
196 graph_layout,
197 events,
198 trace_start,
199 trace_total,
200 session_binding,
201 session_entries,
202 session_entry_start,
203 session_entry_total,
204 session_events,
205 session_event_start,
206 session_event_total,
207 session_events_malformed,
208 session_events_torn_tail,
209 session_capture,
210 session_replay_checkpoint,
211 settings_scopes,
212 settings_start,
213 settings_total,
214 follow_up_queue,
215 follow_up_start,
216 follow_up_total,
217 update_start,
218 update_total,
219 live: raw.get("live").and_then(Value::as_bool).unwrap_or(false),
220 possibly_interrupted: raw
221 .get("possiblyInterrupted")
222 .and_then(Value::as_bool)
223 .unwrap_or(false),
224 })
225}
226
227fn apply_target_patches(view: &mut Value, targets: &[TargetPatch]) -> Result<(), String> {
228 let mut next = view.clone();
229 for target in targets {
230 if target.target_key.ends_with(":tail") {
231 let pointer = if target.target_type == "timeline" {
232 if target.target_key.starts_with("session:") {
233 "/session/eventPage"
234 } else {
235 "/tracePage"
236 }
237 } else if target.target_key.starts_with("entries:") {
238 "/session/entryPage"
239 } else {
240 "/session/eventPage"
241 };
242 if next.pointer(pointer).is_none_or(|page| !is_tail_page(page)) {
243 continue;
244 }
245 }
246 let document = match target.target_type.as_str() {
247 "timeline" if target.target_key.starts_with("session:") => {
248 next.pointer_mut("/session/eventPage")
249 }
250 "timeline" => next.pointer_mut("/tracePage"),
251 "conversation" if target.target_key.starts_with("entries:") => {
252 next.pointer_mut("/session/entryPage")
253 }
254 "conversation" if target.target_key.starts_with("events:") => {
255 next.pointer_mut("/session/eventPage")
256 }
257 "conversation" => next.pointer_mut("/session"),
258 "summary" | "graph" | "replay" | "inspector" => Some(&mut next),
259 _ => None,
260 }
261 .ok_or_else(|| format!("projection target is not loaded: {}", target.target_key))?;
262 apply_patch(document, &target.patch)?;
263 }
264 *view = next;
265 Ok(())
266}
267
268fn session_reload_cursor(targets: &[TargetPatch], view: &Value) -> Option<u64> {
269 targets
270 .iter()
271 .any(|target| target.target_type == "timeline" && target.target_key == "session:reload")
272 .then(|| {
273 view.pointer("/session/eventPage/total")
274 .and_then(Value::as_u64)
275 })
276 .flatten()
277 .and_then(|total| total.checked_sub(1))
278}
279
280fn step_reload_cursor(targets: &[TargetPatch], view: &Value) -> Option<u64> {
281 targets
282 .iter()
283 .any(|target| target.target_type == "replay" && target.target_key == "steps:reload")
284 .then(|| view.get("stepTotal").and_then(Value::as_u64))
285 .flatten()
286 .and_then(|total| total.checked_sub(1))
287}
288
289fn accept_page_response(
290 desired: &mut HashMap<(String, PageKind), u64>,
291 submitted: &mut HashMap<(String, PageKind), u64>,
292 key: &(String, PageKind),
293 cursor: u64,
294) -> bool {
295 if submitted.get(key) == Some(&cursor) {
296 submitted.remove(key);
297 }
298 match desired.get(key) {
299 Some(desired_cursor) if *desired_cursor == cursor => {
300 desired.remove(key);
301 true
302 }
303 Some(_) => false,
304 None => true,
305 }
306}
307
308fn is_tail_page(page: &Value) -> bool {
309 let Some(start) = page.get("start").and_then(Value::as_u64) else {
310 return false;
311 };
312 let Some(total) = page.get("total").and_then(Value::as_u64) else {
313 return false;
314 };
315 let Some(items) = page.get("items").and_then(Value::as_array) else {
316 return false;
317 };
318 start.saturating_add(items.len() as u64) == total
319}
320
321#[derive(Debug, Clone)]
322enum ArtifactEntry {
323 Loading,
324 Ready(String),
325 Error(String),
326}
327
328#[derive(Default)]
329struct Shared {
330 connected: bool,
331 connecting: bool,
332 reconnect_attempt: u32,
333 error: Option<String>,
334 summaries: Vec<Value>,
335 raw_views: HashMap<String, (u64, u64, Value)>,
336 next_view_generation: u64,
337 watched: HashSet<String>,
338 page_requests: HashMap<(String, PageKind), u64>,
339 artifacts: HashMap<(String, String), ArtifactEntry>,
340}
341
342pub struct RemoteRuns {
343 shared: Arc<Mutex<Shared>>,
344 wake: Option<mpsc::UnboundedSender<()>>,
345 worker: Option<JoinHandle<()>>,
346 decoded: HashMap<String, RemoteView>,
347}
348
349impl RemoteRuns {
350 pub fn connect(url: &str) -> Result<Self> {
351 let shared = Arc::new(Mutex::new(Shared {
352 connecting: true,
353 ..Shared::default()
354 }));
355 let (wake_tx, wake_rx) = mpsc::unbounded_channel();
356 let task_shared = Arc::clone(&shared);
357 let url = url.to_string();
358 let worker = std::thread::spawn(move || {
359 let runtime = tokio::runtime::Builder::new_current_thread()
360 .enable_all()
361 .build()
362 .expect("tokio runtime");
363 runtime.block_on(run_reconnecting(&url, task_shared, wake_rx));
364 });
365 let _ = wake_tx.send(());
366 Ok(Self {
367 shared,
368 wake: Some(wake_tx),
369 worker: Some(worker),
370 decoded: HashMap::new(),
371 })
372 }
373
374 pub fn connected(&self) -> bool {
375 self.shared.lock().unwrap().connected
376 }
377
378 pub fn status_label(&self) -> &'static str {
379 let shared = self.shared.lock().unwrap();
380 if shared.connected {
381 "connected"
382 } else if shared.reconnect_attempt > 0 {
383 "reconnecting"
384 } else if shared.connecting {
385 "connecting"
386 } else {
387 "disconnected"
388 }
389 }
390
391 pub fn error(&self) -> Option<String> {
392 self.shared.lock().unwrap().error.clone()
393 }
394
395 pub fn summaries(&self) -> Vec<Value> {
396 self.shared.lock().unwrap().summaries.clone()
397 }
398
399 pub fn watch(&mut self, run_id: &str) {
400 let previous: Vec<String> = {
401 let mut shared = self.shared.lock().unwrap();
402 if shared.watched.len() == 1 && shared.watched.contains(run_id) {
403 return;
404 }
405 let previous: Vec<String> = shared.watched.drain().collect();
406 for old in &previous {
407 shared.raw_views.remove(old);
408 }
409 shared
410 .page_requests
411 .retain(|(candidate, _), _| candidate == run_id);
412 shared
413 .artifacts
414 .retain(|(candidate, _), _| candidate == run_id);
415 shared.watched.insert(run_id.to_string());
416 previous
417 };
418 for old in previous {
419 self.decoded.remove(&old);
420 }
421 self.wake();
422 }
423
424 pub fn request_page(&self, run_id: &str, kind: PageKind, cursor: u64) {
425 self.shared
426 .lock()
427 .unwrap()
428 .page_requests
429 .insert((run_id.to_string(), kind), cursor);
430 self.wake();
431 }
432
433 pub fn request_artifact(&self, run_id: &str, path: &str) {
434 let inserted = {
435 let mut shared = self.shared.lock().unwrap();
436 let key = (run_id.to_string(), path.to_string());
437 if let std::collections::hash_map::Entry::Vacant(entry) = shared.artifacts.entry(key) {
438 entry.insert(ArtifactEntry::Loading);
439 true
440 } else {
441 false
442 }
443 };
444 if inserted {
445 self.wake();
446 }
447 }
448
449 pub fn artifact_content(&self, run_id: &str, path: &str) -> Option<Result<String, String>> {
450 match self
451 .shared
452 .lock()
453 .unwrap()
454 .artifacts
455 .get(&(run_id.to_string(), path.to_string()))
456 .cloned()?
457 {
458 ArtifactEntry::Loading => None,
459 ArtifactEntry::Ready(content) => Some(Ok(content)),
460 ArtifactEntry::Error(error) => Some(Err(error)),
461 }
462 }
463
464 pub fn artifact_snapshot(&self, run_id: &str) -> HashMap<String, Result<String, String>> {
465 self.shared
466 .lock()
467 .unwrap()
468 .artifacts
469 .iter()
470 .filter_map(|((candidate_run, path), entry)| {
471 if candidate_run != run_id {
472 return None;
473 }
474 match entry {
475 ArtifactEntry::Loading => None,
476 ArtifactEntry::Ready(content) => Some((path.clone(), Ok(content.clone()))),
477 ArtifactEntry::Error(error) => Some((path.clone(), Err(error.clone()))),
478 }
479 })
480 .collect()
481 }
482
483 pub fn view(&mut self, run_id: &str) -> Option<&RemoteView> {
484 let raw = {
485 let shared = self.shared.lock().unwrap();
486 let (revision, generation, raw) = shared.raw_views.get(run_id)?;
487 let cached = self.decoded.get(run_id);
488 if cached
489 .is_some_and(|view| view.revision == *revision && view.generation == *generation)
490 {
491 None
492 } else {
493 Some((*revision, *generation, raw.clone()))
494 }
495 };
496 if let Some((revision, generation, raw)) = raw {
497 if let Some(view) = decode_view(revision, generation, &raw) {
498 self.decoded.insert(run_id.to_string(), view);
499 }
500 }
501 self.decoded.get(run_id)
502 }
503
504 fn wake(&self) {
505 if let Some(wake) = &self.wake {
506 let _ = wake.send(());
507 }
508 }
509}
510
511impl Drop for RemoteRuns {
512 fn drop(&mut self) {
513 self.wake.take();
514 if let Some(worker) = self.worker.take() {
515 let _ = worker.join();
516 }
517 }
518}
519
520async fn run_reconnecting(
521 url: &str,
522 shared: Arc<Mutex<Shared>>,
523 mut wake: mpsc::UnboundedReceiver<()>,
524) {
525 let mut attempt = 0u32;
526 loop {
527 {
528 let mut shared = shared.lock().unwrap();
529 shared.connected = false;
530 shared.connecting = true;
531 shared.reconnect_attempt = attempt;
532 }
533 let connection = tokio::select! {
534 connection = tokio_tungstenite::connect_async(url) => connection,
535 message = wake.recv() => {
536 if message.is_none() {
537 return;
538 }
539 continue;
540 }
541 };
542 match connection {
543 Ok((socket, _)) => {
544 attempt = 0;
545 let result = run_socket(socket, Arc::clone(&shared), &mut wake).await;
546 let mut state = shared.lock().unwrap();
547 state.connected = false;
548 state.connecting = false;
549 if state.error.is_none() {
550 state.error = Some(match result {
551 Ok(()) => "connection closed".to_string(),
552 Err(error) => format!("{error:#}"),
553 });
554 }
555 }
556 Err(error) => {
557 let mut state = shared.lock().unwrap();
558 state.connected = false;
559 state.connecting = false;
560 state.error = Some(format!("connecting to {url}: {error}"));
561 }
562 }
563 attempt = attempt.saturating_add(1);
564 {
565 let mut state = shared.lock().unwrap();
566 state.reconnect_attempt = attempt;
567 }
568 let base_ms = (250u64.saturating_mul(1u64 << attempt.min(5))).min(10_000);
569 let jitter_ms = (u64::from(attempt).wrapping_mul(137)) % 251;
570 tokio::select! {
571 _ = tokio::time::sleep(Duration::from_millis(base_ms + jitter_ms)) => {}
572 message = wake.recv() => {
573 if message.is_none() {
574 return;
575 }
576 }
577 }
578 if wake.is_closed() {
579 return;
580 }
581 }
582}
583
584async fn run_socket(
585 socket: tokio_tungstenite::WebSocketStream<
586 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
587 >,
588 shared: Arc<Mutex<Shared>>,
589 wake: &mut mpsc::UnboundedReceiver<()>,
590) -> Result<()> {
591 let (mut sink, mut reads) = socket.split();
592 let mut hello_received = false;
593 let mut subscribed = HashSet::new();
594 let mut submitted_artifacts = HashSet::new();
595 let mut submitted_pages = HashMap::new();
596
597 loop {
598 tokio::select! {
599 wake_message = wake.recv() => {
600 if wake_message.is_none() {
601 return Ok(());
602 }
603 if hello_received {
604 reconcile_desired(
605 &mut sink,
606 &shared,
607 &mut subscribed,
608 &mut submitted_artifacts,
609 &mut submitted_pages,
610 ).await?;
611 }
612 }
613 incoming = reads.next() => {
614 let Some(incoming) = incoming else { return Ok(()) };
615 let text = match incoming? {
616 Message::Text(text) => text,
617 Message::Close(_) => return Ok(()),
618 _ => continue,
619 };
620 let Ok(message) = serde_json::from_str::<ServerMessage>(&text) else {
621 continue;
622 };
623 let mut resubscribe = None;
624 match message {
625 ServerMessage::Hello { protocol } => {
626 if protocol != PROTOCOL_ID {
627 anyhow::bail!("unsupported protocol {protocol}");
628 }
629 {
630 let mut state = shared.lock().unwrap();
631 state.connected = true;
632 state.connecting = false;
633 state.reconnect_attempt = 0;
634 state.error = None;
635 }
636 hello_received = true;
637 send_message(&mut sink, &ClientMessage::WatchRuns).await?;
638 reconcile_desired(
639 &mut sink,
640 &shared,
641 &mut subscribed,
642 &mut submitted_artifacts,
643 &mut submitted_pages,
644 ).await?;
645 }
646 ServerMessage::Runs { runs } => {
647 shared.lock().unwrap().summaries = runs;
648 }
649 ServerMessage::RunSnapshot { run_id, revision, view } => {
650 let mut state = shared.lock().unwrap();
651 if state.watched.contains(&run_id) {
652 state.next_view_generation = state.next_view_generation.wrapping_add(1);
653 let generation = state.next_view_generation;
654 state.raw_views.insert(run_id, (revision, generation, view));
655 }
656 }
657 ServerMessage::RunPatch { run_id, revision, targets } => {
658 let mut state = shared.lock().unwrap();
659 let mut step_cursor = None;
660 let mut session_cursor = None;
661 match state.raw_views.get_mut(&run_id) {
662 Some((current, _, _)) if revision == *current => {}
663 Some((current, generation, view)) if revision == *current + 1 => {
664 if apply_target_patches(view, &targets).is_ok() {
665 *current = revision;
666 *generation = (*generation).wrapping_add(1);
667 step_cursor = step_reload_cursor(&targets, view);
668 session_cursor = session_reload_cursor(&targets, view);
669 } else {
670 resubscribe = Some(run_id.clone());
671 }
672 }
673 Some(_) => resubscribe = Some(run_id.clone()),
674 None => {}
675 }
676 if let Some(cursor) = step_cursor {
677 state
678 .page_requests
679 .insert((run_id.clone(), PageKind::Steps), cursor);
680 }
681 if let Some(cursor) = session_cursor {
682 state
683 .page_requests
684 .insert((run_id.clone(), PageKind::SessionEvents), cursor);
685 state
686 .page_requests
687 .insert((run_id, PageKind::SessionEntries), cursor);
688 }
689 }
690 ServerMessage::RunPage {
691 run_id,
692 revision,
693 kind,
694 cursor,
695 start,
696 total,
697 items,
698 graph_cursor,
699 graph_steps,
700 taken_transitions,
701 replay_checkpoint,
702 } => {
703 let page_key = (run_id.clone(), kind);
704 let mut state = shared.lock().unwrap();
705 let accepted = accept_page_response(
706 &mut state.page_requests,
707 &mut submitted_pages,
708 &page_key,
709 cursor,
710 );
711 if accepted {
712 if let Some((current, generation, view)) =
713 state.raw_views.get_mut(&run_id)
714 {
715 if revision != *current {
716 resubscribe = Some(run_id);
717 } else {
718 let pointer = match kind {
719 PageKind::Steps => "/state/steps",
720 PageKind::Trace | PageKind::TraceAtStep => "/tracePage",
721 PageKind::SessionEntries => "/session/entryPage",
722 PageKind::SessionEvents => "/session/eventPage",
723 PageKind::Settings => "/settingsScopes",
724 PageKind::FollowUps => "/followUpQueue/items",
725 PageKind::Updates => "/state/updates",
726 };
727 if let Some(page) = view.pointer_mut(pointer) {
728 match kind {
729 PageKind::Steps => {
730 *page = Value::Array(items);
731 view["stepStart"] = serde_json::json!(start);
732 view["stepTotal"] = serde_json::json!(total);
733 if let Some(cursor) = graph_cursor {
734 view["graphCursor"] = serde_json::json!(cursor);
735 }
736 if let Some(steps) = graph_steps {
737 view["graphSteps"] = Value::Array(steps);
738 }
739 if let Some(transitions) = taken_transitions {
740 view["takenTransitions"] =
741 serde_json::json!(transitions);
742 }
743 }
744 PageKind::Settings => {
745 *page = Value::Array(items);
746 view["settingsStart"] = serde_json::json!(start);
747 view["settingsTotal"] = serde_json::json!(total);
748 }
749 PageKind::FollowUps => {
750 *page = Value::Array(items);
751 view["followUpStart"] = serde_json::json!(start);
752 view["followUpTotal"] = serde_json::json!(total);
753 }
754 PageKind::Updates => {
755 *page = Value::Array(items);
756 view["updateStart"] = serde_json::json!(start);
757 view["updateTotal"] = serde_json::json!(total);
758 }
759 PageKind::Trace
760 | PageKind::TraceAtStep
761 | PageKind::SessionEntries
762 | PageKind::SessionEvents => {
763 *page = serde_json::json!({
764 "presentationRevision": revision,
765 "start": start,
766 "total": total,
767 "items": items,
768 });
769 }
770 }
771 if kind == PageKind::SessionEvents {
772 if let Some(session) = view
773 .get_mut("session")
774 .and_then(Value::as_object_mut)
775 {
776 session.insert(
777 "replayCheckpoint".to_string(),
778 replay_checkpoint.unwrap_or(Value::Null),
779 );
780 }
781 }
782 *generation = (*generation).wrapping_add(1);
783 } else {
784 resubscribe = Some(run_id);
785 }
786 }
787 }
788 }
789 }
790 ServerMessage::Artifact { run_id, path, content } => {
791 let key = (run_id, path);
792 submitted_artifacts.remove(&key);
793 shared
794 .lock()
795 .unwrap()
796 .artifacts
797 .insert(key, ArtifactEntry::Ready(content));
798 }
799 ServerMessage::Error { message, run_id } => {
800 let mut state = shared.lock().unwrap();
801 if let Some(run_id) = run_id {
802 if let Some(key) = submitted_artifacts
803 .iter()
804 .find(|(candidate_run, _)| candidate_run == &run_id)
805 .cloned()
806 {
807 submitted_artifacts.remove(&key);
808 state.artifacts.insert(key, ArtifactEntry::Error(message));
809 } else {
810 state.error = Some(message);
811 }
812 } else {
813 state.error = Some(message);
814 }
815 }
816 }
817 if hello_received {
818 reconcile_desired(
819 &mut sink,
820 &shared,
821 &mut subscribed,
822 &mut submitted_artifacts,
823 &mut submitted_pages,
824 ).await?;
825 }
826 if let Some(run_id) = resubscribe {
827 send_message(
828 &mut sink,
829 &ClientMessage::WatchRun {
830 run_id,
831 revision: None,
832 step_cursor: None,
833 trace_cursor: None,
834 session_entry_cursor: None,
835 session_event_cursor: None,
836 },
837 )
838 .await?;
839 }
840 }
841 }
842 }
843}
844
845async fn reconcile_desired(
846 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
847 shared: &Arc<Mutex<Shared>>,
848 subscribed: &mut HashSet<String>,
849 submitted_artifacts: &mut HashSet<(String, String)>,
850 submitted_pages: &mut HashMap<(String, PageKind), u64>,
851) -> Result<()> {
852 let (desired, artifacts, pages) = {
853 let state = shared.lock().unwrap();
854 let desired = state.watched.clone();
855 let artifacts = state
856 .artifacts
857 .iter()
858 .filter_map(|(key, entry)| {
859 matches!(entry, ArtifactEntry::Loading).then_some(key.clone())
860 })
861 .collect::<Vec<_>>();
862 let pages = state
863 .page_requests
864 .iter()
865 .map(|(key, cursor)| (key.clone(), *cursor))
866 .collect::<Vec<_>>();
867 (desired, artifacts, pages)
868 };
869 let removals: Vec<String> = subscribed.difference(&desired).cloned().collect();
870 let additions: Vec<String> = desired.difference(subscribed).cloned().collect();
871 for run_id in removals {
872 send_message(
873 sink,
874 &ClientMessage::UnwatchRun {
875 run_id: run_id.clone(),
876 },
877 )
878 .await?;
879 subscribed.remove(&run_id);
880 submitted_pages.retain(|(candidate, _), _| candidate != &run_id);
881 submitted_artifacts.retain(|(candidate, _)| candidate != &run_id);
882 }
883 for run_id in additions {
884 let (revision, step_cursor, trace_cursor, session_entry_cursor, session_event_cursor) = {
885 let state = shared.lock().unwrap();
886 state.raw_views.get(&run_id).map_or(
887 (None, None, None, None, None),
888 |(revision, _, view)| {
889 (
890 Some(*revision),
891 view.get("stepStart").and_then(Value::as_u64),
892 view.pointer("/tracePage/start").and_then(Value::as_u64),
893 view.pointer("/session/entryPage/start")
894 .and_then(Value::as_u64),
895 view.pointer("/session/eventPage/start")
896 .and_then(Value::as_u64),
897 )
898 },
899 )
900 };
901 send_message(
902 sink,
903 &ClientMessage::WatchRun {
904 run_id: run_id.clone(),
905 revision,
906 step_cursor,
907 trace_cursor,
908 session_entry_cursor,
909 session_event_cursor,
910 },
911 )
912 .await?;
913 subscribed.insert(run_id);
914 }
915 for ((run_id, kind), cursor) in pages {
916 let key = (run_id.clone(), kind);
917 if submitted_pages.get(&key) != Some(&cursor) {
918 submitted_pages.insert(key, cursor);
919 send_message(
920 sink,
921 &ClientMessage::FetchPage {
922 run_id,
923 kind,
924 cursor,
925 },
926 )
927 .await?;
928 }
929 }
930 if submitted_artifacts.is_empty() {
931 if let Some((run_id, path)) = artifacts.into_iter().next() {
932 let key = (run_id.clone(), path.clone());
933 submitted_artifacts.insert(key);
934 send_message(sink, &ClientMessage::FetchArtifact { run_id, path }).await?;
935 }
936 }
937 Ok(())
938}
939
940async fn send_message(
941 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
942 message: &ClientMessage,
943) -> Result<()> {
944 let text = serde_json::to_string(message).context("encoding client message")?;
945 sink.send(Message::Text(text.into())).await?;
946 Ok(())
947}
948
949#[cfg(test)]
950mod tests {
951 use super::*;
952 use crate::protocol::PatchOp;
953 use serde_json::json;
954
955 fn entry_tail_target() -> TargetPatch {
956 TargetPatch {
957 target_type: "conversation".to_string(),
958 target_key: "entries:tail".to_string(),
959 patch: vec![
960 PatchOp::Replace {
961 path: "/presentationRevision".to_string(),
962 value: json!(2),
963 },
964 PatchOp::Remove {
965 path: "/items/0".to_string(),
966 },
967 PatchOp::Append {
968 path: "/items".to_string(),
969 value: vec![json!({"seq": 3})],
970 },
971 PatchOp::Replace {
972 path: "/start".to_string(),
973 value: json!(1),
974 },
975 PatchOp::Replace {
976 path: "/total".to_string(),
977 value: json!(3),
978 },
979 ],
980 }
981 }
982
983 #[test]
984 fn applies_tail_patches_only_to_the_loaded_tail_page() {
985 let mut tail = json!({
986 "session": {
987 "entryPage": {
988 "presentationRevision": 1,
989 "start": 0,
990 "total": 2,
991 "items": [{"seq": 1}, {"seq": 2}]
992 }
993 }
994 });
995 apply_target_patches(&mut tail, &[entry_tail_target()]).unwrap();
996 assert_eq!(tail.pointer("/session/entryPage/start"), Some(&json!(1)));
997 assert_eq!(
998 tail.pointer("/session/entryPage/items/1/seq"),
999 Some(&json!(3))
1000 );
1001
1002 let mut middle = json!({
1003 "session": {
1004 "entryPage": {
1005 "presentationRevision": 1,
1006 "start": 0,
1007 "total": 5,
1008 "items": [{"seq": 1}, {"seq": 2}]
1009 }
1010 }
1011 });
1012 let before = middle.clone();
1013 apply_target_patches(&mut middle, &[entry_tail_target()]).unwrap();
1014 assert_eq!(middle, before);
1015 }
1016
1017 #[test]
1018 fn a_session_reload_delta_requests_the_new_aligned_page() {
1019 let targets = vec![TargetPatch {
1020 target_type: "timeline".to_string(),
1021 target_key: "session:reload".to_string(),
1022 patch: vec![PatchOp::Replace {
1023 path: "/total".to_string(),
1024 value: json!(257),
1025 }],
1026 }];
1027 let view = json!({"session": {"eventPage": {"total": 257}}});
1028 assert_eq!(session_reload_cursor(&targets, &view), Some(256));
1029 }
1030
1031 #[test]
1032 fn a_step_reload_delta_requests_the_latest_step_page() {
1033 let targets = vec![TargetPatch {
1034 target_type: "replay".to_string(),
1035 target_key: "steps:reload".to_string(),
1036 patch: vec![PatchOp::Replace {
1037 path: "/stepTotal".to_string(),
1038 value: json!(12),
1039 }],
1040 }];
1041 assert_eq!(
1042 step_reload_cursor(&targets, &json!({"stepTotal": 12})),
1043 Some(11)
1044 );
1045 assert_eq!(step_reload_cursor(&targets, &json!({"stepTotal": 0})), None);
1046 }
1047
1048 #[test]
1049 fn an_older_response_cannot_discard_a_newer_page_request() {
1050 let key = ("run-1".to_string(), PageKind::SessionEvents);
1051 let mut desired = HashMap::from([(key.clone(), 900)]);
1052 let mut submitted = HashMap::from([(key.clone(), 900)]);
1053
1054 assert!(!accept_page_response(
1055 &mut desired,
1056 &mut submitted,
1057 &key,
1058 800,
1059 ));
1060 assert_eq!(desired.get(&key), Some(&900));
1061 assert_eq!(submitted.get(&key), Some(&900));
1062
1063 assert!(accept_page_response(
1064 &mut desired,
1065 &mut submitted,
1066 &key,
1067 900,
1068 ));
1069 assert!(!desired.contains_key(&key));
1070 assert!(!submitted.contains_key(&key));
1071 }
1072
1073 #[test]
1074 fn rejects_a_patch_without_mutating_the_last_good_view() {
1075 let mut view = json!({"presentationRevision": 1});
1076 let before = view.clone();
1077 let result = apply_target_patches(
1078 &mut view,
1079 &[TargetPatch {
1080 target_type: "graph".to_string(),
1081 target_key: String::new(),
1082 patch: vec![PatchOp::Replace {
1083 path: "/missing/value".to_string(),
1084 value: json!(2),
1085 }],
1086 }],
1087 );
1088 assert!(result.is_err());
1089 assert_eq!(view, before);
1090 }
1091}