Skip to main content

piw/
client.rs

1//! Reconnecting client used by local Unix-socket and remote WebSocket modes.
2
3use crate::layout::{layout_graph, GraphLayout};
4use crate::protocol::{
5    apply_patch, encode_request, parse_server_message, ClientRequest, PageKind, ServerMessage,
6    TargetPatch, PROTOCOL_ID,
7};
8use crate::state::types::{
9    as_artifact_ref, ArtifactRef, DefinitionSnapshot, Manifest, RunState, StepRecord,
10};
11use anyhow::{Context, Result};
12use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
13use futures_util::{SinkExt, StreamExt};
14use serde_json::{json, Value};
15use sha2::{Digest, Sha256};
16use std::collections::{HashMap, HashSet};
17use std::path::{Path, PathBuf};
18use std::sync::{Arc, Mutex};
19use std::thread::JoinHandle;
20use std::time::Duration;
21use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
22#[cfg(windows)]
23use tokio::net::windows::named_pipe::ClientOptions;
24#[cfg(unix)]
25use tokio::net::UnixStream;
26use tokio::sync::mpsc;
27use tokio_tungstenite::tungstenite::Message;
28
29pub struct RemoteView {
30    pub revision: u64,
31    pub graph_revision: u64,
32    generation: u64,
33    pub manifest: Manifest,
34    pub state: RunState,
35    pub graph_steps: Vec<StepRecord>,
36    pub taken_transitions: Vec<String>,
37    pub graph_cursor: u64,
38    pub step_start: u64,
39    pub step_total: u64,
40    pub snapshot: Option<DefinitionSnapshot>,
41    pub graph_layout: Option<GraphLayout>,
42    pub events: Vec<Value>,
43    pub trace_start: u64,
44    pub trace_total: u64,
45    pub session_binding: Option<Value>,
46    pub session_entries: Vec<Value>,
47    pub session_entry_start: u64,
48    pub session_entry_total: u64,
49    pub session_events: Vec<Value>,
50    pub session_event_start: u64,
51    pub session_event_total: u64,
52    pub session_events_malformed: bool,
53    pub session_events_torn_tail: bool,
54    pub session_capture: Option<Value>,
55    pub session_replay_checkpoint: Option<Value>,
56    pub settings_scopes: Vec<Value>,
57    pub settings_start: u64,
58    pub settings_total: u64,
59    pub follow_up_queue: Option<Value>,
60    pub follow_up_start: u64,
61    pub follow_up_total: u64,
62    pub update_start: u64,
63    pub update_total: u64,
64    pub live: bool,
65    pub possibly_interrupted: bool,
66}
67
68fn decode_view(
69    revision: u64,
70    generation: u64,
71    raw: &Value,
72    definition_content: Option<&str>,
73    graph_history_content: Option<&str>,
74) -> Option<RemoteView> {
75    let graph_revision = raw
76        .get("graphRevision")
77        .and_then(Value::as_u64)
78        .unwrap_or(revision);
79    let manifest: Manifest = serde_json::from_value(raw.get("manifest")?.clone()).ok()?;
80    let state: RunState = serde_json::from_value(raw.get("state")?.clone()).ok()?;
81    let graph_history = graph_history_value(raw, graph_history_content);
82    let graph_steps = graph_history
83        .as_ref()
84        .and_then(|value| value.get("steps"))
85        .and_then(|value| serde_json::from_value(value.clone()).ok())
86        .or_else(|| {
87            raw.get("graphSteps")
88                .and_then(|value| serde_json::from_value(value.clone()).ok())
89        })
90        .unwrap_or_else(|| state.steps.clone());
91    let taken_transitions = graph_history
92        .as_ref()
93        .and_then(|value| value.get("transitions"))
94        .and_then(|value| serde_json::from_value(value.clone()).ok())
95        .or_else(|| {
96            raw.get("takenTransitions")
97                .and_then(|value| serde_json::from_value(value.clone()).ok())
98        })
99        .unwrap_or_default();
100    let graph_cursor = raw
101        .get("graphCursor")
102        .and_then(Value::as_u64)
103        .unwrap_or_else(|| state.steps.len().saturating_sub(1) as u64);
104    let step_start = raw.get("stepStart").and_then(Value::as_u64).unwrap_or(0);
105    let step_total = raw
106        .get("stepTotal")
107        .and_then(Value::as_u64)
108        .unwrap_or(state.steps.len() as u64);
109    let snapshot: Option<DefinitionSnapshot> = workflow_definition_value(raw, definition_content)
110        .and_then(|value| serde_json::from_value(value).ok());
111    let graph_layout = raw
112        .get("graphScene")
113        .and_then(|value| serde_json::from_value(value.clone()).ok())
114        .or_else(|| snapshot.as_ref().map(layout_graph));
115    let events = page_items(raw, "/tracePage/items");
116    let trace_start = pointer_u64(raw, "/tracePage/start");
117    let trace_total = pointer_u64(raw, "/tracePage/total").max(events.len() as u64);
118    let session_binding = raw
119        .pointer("/session/binding")
120        .cloned()
121        .filter(|v| !v.is_null());
122    let session_entries = page_items(raw, "/session/entryPage/items");
123    let session_entry_start = pointer_u64(raw, "/session/entryPage/start");
124    let session_entry_total =
125        pointer_u64(raw, "/session/entryPage/total").max(session_entries.len() as u64);
126    let session_events = page_items(raw, "/session/eventPage/items");
127    let session_event_start = pointer_u64(raw, "/session/eventPage/start");
128    let session_event_total =
129        pointer_u64(raw, "/session/eventPage/total").max(session_events.len() as u64);
130    let settings_scopes = raw
131        .get("settingsScopes")
132        .and_then(Value::as_array)
133        .cloned()
134        .unwrap_or_default();
135    let follow_up_queue = raw.get("followUpQueue").cloned().filter(|v| !v.is_null());
136    let update_total = raw
137        .get("updateTotal")
138        .and_then(Value::as_u64)
139        .unwrap_or_else(|| {
140            state
141                .updates
142                .as_ref()
143                .map_or(0, |updates| updates.len() as u64)
144        });
145    Some(RemoteView {
146        revision,
147        graph_revision,
148        generation,
149        manifest,
150        state,
151        graph_steps,
152        taken_transitions,
153        graph_cursor,
154        step_start,
155        step_total,
156        snapshot,
157        graph_layout,
158        events,
159        trace_start,
160        trace_total,
161        session_binding,
162        session_entries,
163        session_entry_start,
164        session_entry_total,
165        session_events,
166        session_event_start,
167        session_event_total,
168        session_events_malformed: raw
169            .pointer("/session/integrity/malformed")
170            .and_then(Value::as_bool)
171            .unwrap_or(false),
172        session_events_torn_tail: raw
173            .pointer("/session/integrity/tornTail")
174            .and_then(Value::as_bool)
175            .unwrap_or(false),
176        session_capture: raw
177            .pointer("/session/capture")
178            .cloned()
179            .filter(|v| !v.is_null()),
180        session_replay_checkpoint: raw
181            .pointer("/session/replayCheckpoint")
182            .cloned()
183            .filter(|v| !v.is_null()),
184        settings_start: raw
185            .get("settingsStart")
186            .and_then(Value::as_u64)
187            .unwrap_or(0),
188        settings_total: raw
189            .get("settingsTotal")
190            .and_then(Value::as_u64)
191            .unwrap_or(settings_scopes.len() as u64),
192        settings_scopes,
193        follow_up_start: raw
194            .get("followUpStart")
195            .and_then(Value::as_u64)
196            .unwrap_or(0),
197        follow_up_total: raw
198            .get("followUpTotal")
199            .and_then(Value::as_u64)
200            .unwrap_or(0),
201        follow_up_queue,
202        update_start: raw.get("updateStart").and_then(Value::as_u64).unwrap_or(0),
203        update_total,
204        live: raw.get("live").and_then(Value::as_bool).unwrap_or(false),
205        possibly_interrupted: raw
206            .get("possiblyInterrupted")
207            .and_then(Value::as_bool)
208            .unwrap_or(false),
209    })
210}
211
212fn page_items(raw: &Value, pointer: &str) -> Vec<Value> {
213    raw.pointer(pointer)
214        .and_then(Value::as_array)
215        .cloned()
216        .unwrap_or_default()
217}
218
219fn pointer_u64(raw: &Value, pointer: &str) -> u64 {
220    raw.pointer(pointer).and_then(Value::as_u64).unwrap_or(0)
221}
222
223fn workflow_definition_artifact(raw: &Value) -> Option<ArtifactRef> {
224    let workflow = raw.get("workflow")?;
225    as_artifact_ref(workflow).or_else(|| workflow.get("content").and_then(as_artifact_ref))
226}
227
228fn workflow_definition_value(raw: &Value, content: Option<&str>) -> Option<Value> {
229    let workflow = raw.get("workflow")?;
230    if workflow_definition_artifact(raw).is_some() {
231        return serde_json::from_str(content?).ok();
232    }
233    Some(workflow.clone())
234}
235
236fn graph_history_artifact(raw: &Value) -> Option<ArtifactRef> {
237    raw.get("graphHistory").and_then(as_artifact_ref)
238}
239
240fn graph_history_value(raw: &Value, content: Option<&str>) -> Option<Value> {
241    let history = raw.get("graphHistory")?;
242    if graph_history_artifact(raw).is_some() {
243        return serde_json::from_str(content?).ok();
244    }
245    Some(history.clone())
246}
247
248fn workflow_definition_content(
249    state: &mut Shared,
250    run_id: &str,
251    raw: &Value,
252) -> (Option<String>, bool) {
253    let Some(artifact) = workflow_definition_artifact(raw) else {
254        return (None, false);
255    };
256    referenced_json_content(state, run_id, artifact, "workflow definition")
257}
258
259fn graph_history_content(state: &mut Shared, run_id: &str, raw: &Value) -> (Option<String>, bool) {
260    let Some(artifact) = graph_history_artifact(raw) else {
261        return (None, false);
262    };
263    referenced_json_content(state, run_id, artifact, "workflow graph history")
264}
265
266fn referenced_json_content(
267    state: &mut Shared,
268    run_id: &str,
269    artifact: ArtifactRef,
270    label: &str,
271) -> (Option<String>, bool) {
272    let key = (run_id.to_string(), artifact.path.clone());
273    match state.artifacts.get(&key).cloned() {
274        Some(ArtifactEntry::Ready(content)) => {
275            let valid = artifact.media_type == "application/json"
276                && content.len() as u64 == artifact.bytes
277                && hex_sha256(content.as_bytes()) == artifact.sha256;
278            if valid {
279                return (Some(content), false);
280            }
281            let error = format!("{label} content does not match its reference");
282            state
283                .artifacts
284                .insert(key, ArtifactEntry::Error(error.clone()));
285            state.error = Some(error);
286            (None, false)
287        }
288        Some(ArtifactEntry::Error(error)) => {
289            state.error = Some(format!("{label} content is unavailable: {error}"));
290            (None, false)
291        }
292        Some(ArtifactEntry::Loading(_)) => (None, false),
293        None => {
294            state
295                .artifacts
296                .insert(key.clone(), ArtifactEntry::Loading(Vec::new()));
297            state.content_requests.insert(key, 0);
298            (None, true)
299        }
300    }
301}
302
303fn apply_target_patches(view: &mut Value, targets: &[TargetPatch]) -> Result<()> {
304    let mut next = view.clone();
305    for target in targets {
306        if target.target_key.ends_with(":tail") {
307            let pointer = if target.target_type == "timeline" {
308                if target.target_key.starts_with("session:") {
309                    "/session/eventPage"
310                } else {
311                    "/tracePage"
312                }
313            } else if target.target_key.starts_with("entries:") {
314                "/session/entryPage"
315            } else {
316                "/session/eventPage"
317            };
318            if next.pointer(pointer).is_none_or(|page| !is_tail_page(page)) {
319                continue;
320            }
321        }
322        let document = match target.target_type.as_str() {
323            "timeline" if target.target_key.starts_with("session:") => {
324                next.pointer_mut("/session/eventPage")
325            }
326            "timeline" => next.pointer_mut("/tracePage"),
327            "conversation" if target.target_key.starts_with("entries:") => {
328                next.pointer_mut("/session/entryPage")
329            }
330            "conversation" if target.target_key.starts_with("events:") => {
331                next.pointer_mut("/session/eventPage")
332            }
333            "conversation" => next.pointer_mut("/session"),
334            "summary" | "graph" | "replay" | "inspector" => Some(&mut next),
335            _ => None,
336        }
337        .with_context(|| format!("projection target is not loaded: {}", target.target_key))?;
338        apply_patch(document, &target.patch).map_err(anyhow::Error::msg)?;
339    }
340    *view = next;
341    Ok(())
342}
343
344fn is_tail_page(page: &Value) -> bool {
345    let Some(start) = page.get("start").and_then(Value::as_u64) else {
346        return false;
347    };
348    let Some(total) = page.get("total").and_then(Value::as_u64) else {
349        return false;
350    };
351    let Some(items) = page.get("items").and_then(Value::as_array) else {
352        return false;
353    };
354    start.saturating_add(items.len() as u64) == total
355}
356
357#[derive(Debug, Clone)]
358enum ArtifactEntry {
359    Loading(Vec<u8>),
360    Ready(String),
361    Error(String),
362}
363
364struct RunListAssembly {
365    revision: String,
366    total: u64,
367    items: Vec<Value>,
368}
369
370#[derive(Default)]
371struct Shared {
372    connected: bool,
373    connecting: bool,
374    reconnect_attempt: u32,
375    error: Option<String>,
376    summaries: Vec<Value>,
377    run_list: Option<RunListAssembly>,
378    run_list_request: Option<(String, u64)>,
379    raw_views: HashMap<String, (u64, u64, u64, Value)>,
380    next_view_generation: u64,
381    watched: HashSet<String>,
382    page_requests: HashMap<(String, PageKind), u64>,
383    content_requests: HashMap<(String, String), u64>,
384    artifacts: HashMap<(String, String), ArtifactEntry>,
385}
386
387#[derive(Clone)]
388enum Endpoint {
389    WebSocket(String),
390    Local(PathBuf),
391}
392
393fn discard_run_state(shared: &mut Shared, run_id: &str) {
394    shared.raw_views.remove(run_id);
395    shared
396        .page_requests
397        .retain(|(candidate, _), _| candidate != run_id);
398    shared
399        .content_requests
400        .retain(|(candidate, _), _| candidate != run_id);
401    shared
402        .artifacts
403        .retain(|(candidate, _), _| candidate != run_id);
404}
405
406pub struct RemoteRuns {
407    shared: Arc<Mutex<Shared>>,
408    wake: Option<mpsc::UnboundedSender<()>>,
409    worker: Option<JoinHandle<()>>,
410    decoded: HashMap<String, RemoteView>,
411}
412
413impl RemoteRuns {
414    pub fn connect(url: &str) -> Result<Self> {
415        Self::start(Endpoint::WebSocket(url.to_string()))
416    }
417
418    pub fn connect_local(path: &Path) -> Result<Self> {
419        Self::start(Endpoint::Local(path.to_path_buf()))
420    }
421
422    fn start(endpoint: Endpoint) -> Result<Self> {
423        let shared = Arc::new(Mutex::new(Shared {
424            connecting: true,
425            ..Shared::default()
426        }));
427        let (wake_tx, wake_rx) = mpsc::unbounded_channel();
428        let task_shared = Arc::clone(&shared);
429        let worker = std::thread::spawn(move || {
430            let runtime = tokio::runtime::Builder::new_current_thread()
431                .enable_all()
432                .build()
433                .expect("tokio runtime");
434            runtime.block_on(run_reconnecting(endpoint, task_shared, wake_rx));
435        });
436        let _ = wake_tx.send(());
437        Ok(Self {
438            shared,
439            wake: Some(wake_tx),
440            worker: Some(worker),
441            decoded: HashMap::new(),
442        })
443    }
444
445    pub fn connected(&self) -> bool {
446        self.shared.lock().unwrap().connected
447    }
448
449    pub fn status_label(&self) -> &'static str {
450        let shared = self.shared.lock().unwrap();
451        if shared.connected {
452            "connected"
453        } else if shared.reconnect_attempt > 0 {
454            "reconnecting"
455        } else if shared.connecting {
456            "connecting"
457        } else {
458            "disconnected"
459        }
460    }
461
462    pub fn error(&self) -> Option<String> {
463        self.shared.lock().unwrap().error.clone()
464    }
465
466    pub fn summaries(&self) -> Vec<Value> {
467        self.shared.lock().unwrap().summaries.clone()
468    }
469
470    pub fn watch(&mut self, run_id: &str) {
471        let mut shared = self.shared.lock().unwrap();
472        let old: Vec<String> = shared.watched.drain().collect();
473        for old_id in old {
474            if old_id != run_id {
475                discard_run_state(&mut shared, &old_id);
476                self.decoded.remove(&old_id);
477            }
478        }
479        shared.watched.insert(run_id.to_string());
480        drop(shared);
481        self.wake();
482    }
483
484    pub fn request_page(&self, run_id: &str, kind: PageKind, cursor: u64) {
485        self.shared
486            .lock()
487            .unwrap()
488            .page_requests
489            .insert((run_id.to_string(), kind), cursor);
490        self.wake();
491    }
492
493    pub fn request_artifact(&self, run_id: &str, path: &str) {
494        let key = (run_id.to_string(), path.to_string());
495        let mut shared = self.shared.lock().unwrap();
496        if shared.artifacts.contains_key(&key) {
497            return;
498        }
499        shared
500            .artifacts
501            .insert(key.clone(), ArtifactEntry::Loading(Vec::new()));
502        shared.content_requests.insert(key, 0);
503        drop(shared);
504        self.wake();
505    }
506
507    pub fn artifact_snapshot(
508        &self,
509        run_id: &str,
510    ) -> HashMap<String, std::result::Result<String, String>> {
511        self.shared
512            .lock()
513            .unwrap()
514            .artifacts
515            .iter()
516            .filter(|((candidate, _), _)| candidate == run_id)
517            .filter_map(|((_, path), entry)| {
518                let value = match entry {
519                    ArtifactEntry::Loading(_) => return None,
520                    ArtifactEntry::Ready(content) => Ok(content.clone()),
521                    ArtifactEntry::Error(error) => Err(error.clone()),
522                };
523                Some((path.clone(), value))
524            })
525            .collect()
526    }
527
528    pub fn view(&mut self, run_id: &str) -> Option<&RemoteView> {
529        let raw = self.shared.lock().unwrap().raw_views.get(run_id).cloned();
530        let Some((_, revision, generation, raw)) = raw else {
531            self.decoded.remove(run_id);
532            return None;
533        };
534        let (definition_content, graph_history_content, requested) = {
535            let mut shared = self.shared.lock().unwrap();
536            let (definition, definition_requested) =
537                workflow_definition_content(&mut shared, run_id, &raw);
538            let (graph_history, graph_requested) = graph_history_content(&mut shared, run_id, &raw);
539            (
540                definition,
541                graph_history,
542                definition_requested || graph_requested,
543            )
544        };
545        if requested {
546            self.wake();
547        }
548        let stale = self
549            .decoded
550            .get(run_id)
551            .is_none_or(|view| view.generation != generation);
552        if stale {
553            if let Some(decoded) = decode_view(
554                revision,
555                generation,
556                &raw,
557                definition_content.as_deref(),
558                graph_history_content.as_deref(),
559            ) {
560                self.decoded.insert(run_id.to_string(), decoded);
561            } else {
562                self.decoded.remove(run_id);
563            }
564        }
565        self.decoded.get(run_id)
566    }
567
568    fn wake(&self) {
569        if let Some(wake) = &self.wake {
570            let _ = wake.send(());
571        }
572    }
573}
574
575impl Drop for RemoteRuns {
576    fn drop(&mut self) {
577        self.wake.take();
578        if let Some(worker) = self.worker.take() {
579            let _ = worker.join();
580        }
581    }
582}
583
584async fn run_reconnecting(
585    endpoint: Endpoint,
586    shared: Arc<Mutex<Shared>>,
587    mut wake: mpsc::UnboundedReceiver<()>,
588) {
589    let mut attempt = 0u32;
590    loop {
591        {
592            let mut state = shared.lock().unwrap();
593            state.connected = false;
594            state.connecting = true;
595            state.reconnect_attempt = attempt;
596        }
597        let result = match &endpoint {
598            Endpoint::WebSocket(url) => run_websocket(url, Arc::clone(&shared), &mut wake).await,
599            Endpoint::Local(path) => run_local(path, Arc::clone(&shared), &mut wake).await,
600        };
601        if wake.is_closed() {
602            return;
603        }
604        {
605            let mut state = shared.lock().unwrap();
606            state.connected = false;
607            state.connecting = false;
608            state.error = Some(match result {
609                Ok(()) => "connection closed".to_string(),
610                Err(error) => format!("{error:#}"),
611            });
612        }
613        attempt = attempt.saturating_add(1);
614        let delay = (250u64.saturating_mul(1u64 << attempt.min(5))).min(10_000);
615        tokio::select! {
616            _ = tokio::time::sleep(Duration::from_millis(delay)) => {}
617            message = wake.recv() => if message.is_none() { return; }
618        }
619    }
620}
621
622async fn run_websocket(
623    url: &str,
624    shared: Arc<Mutex<Shared>>,
625    wake: &mut mpsc::UnboundedReceiver<()>,
626) -> Result<()> {
627    let (socket, _) = tokio_tungstenite::connect_async(url)
628        .await
629        .with_context(|| format!("connecting to {url}"))?;
630    let (mut sink, mut stream) = socket.split();
631    let mut sent = HashSet::new();
632    let mut counter = 0u64;
633    loop {
634        for request in reconcile_requests(&shared, &mut sent, &mut counter) {
635            let text = encode_request(&request).map_err(anyhow::Error::msg)?;
636            sink.send(Message::Text(text.into())).await?;
637        }
638        tokio::select! {
639            message = stream.next() => {
640                let Some(message) = message else { return Ok(()) };
641                match message? {
642                    Message::Text(text) => handle_server_message(text.as_ref(), &shared)?,
643                    Message::Close(_) => return Ok(()),
644                    Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
645                    Message::Pong(_) => {}
646                    Message::Binary(_) | Message::Frame(_) => anyhow::bail!("server frame must be text"),
647                }
648            }
649            message = wake.recv() => {
650                if message.is_none() { return Ok(()); }
651            }
652            _ = tokio::time::sleep(Duration::from_millis(250)) => {}
653        }
654    }
655}
656
657async fn run_local(
658    path: &Path,
659    shared: Arc<Mutex<Shared>>,
660    wake: &mut mpsc::UnboundedReceiver<()>,
661) -> Result<()> {
662    #[cfg(unix)]
663    {
664        let socket = UnixStream::connect(path)
665            .await
666            .with_context(|| format!("connecting to workflow host {}", path.display()))?;
667        run_local_connection(socket, shared, wake).await
668    }
669    #[cfg(windows)]
670    {
671        let socket = ClientOptions::new()
672            .open(path)
673            .with_context(|| format!("connecting to workflow host {}", path.display()))?;
674        run_local_connection(socket, shared, wake).await
675    }
676    #[cfg(not(any(unix, windows)))]
677    anyhow::bail!("local workflow host transport is not supported on this platform");
678}
679
680async fn run_local_connection<S>(
681    socket: S,
682    shared: Arc<Mutex<Shared>>,
683    wake: &mut mpsc::UnboundedReceiver<()>,
684) -> Result<()>
685where
686    S: AsyncRead + AsyncWrite + Unpin,
687{
688    let (read, mut write) = tokio::io::split(socket);
689    let mut lines = BufReader::new(read).lines();
690    let mut sent = HashSet::new();
691    let mut counter = 0u64;
692    loop {
693        for request in reconcile_requests(&shared, &mut sent, &mut counter) {
694            let text = encode_request(&request).map_err(anyhow::Error::msg)?;
695            write.write_all(text.as_bytes()).await?;
696            write.write_all(b"\n").await?;
697            write.flush().await?;
698        }
699        tokio::select! {
700            line = lines.next_line() => {
701                let Some(line) = line? else { return Ok(()) };
702                handle_server_message(&line, &shared)?;
703            }
704            message = wake.recv() => {
705                if message.is_none() { return Ok(()); }
706            }
707            _ = tokio::time::sleep(Duration::from_millis(250)) => {}
708        }
709    }
710}
711
712fn reconcile_requests(
713    shared: &Arc<Mutex<Shared>>,
714    sent: &mut HashSet<String>,
715    counter: &mut u64,
716) -> Vec<ClientRequest> {
717    let (watched, pages, contents, revisions, run_list_request) = {
718        let state = shared.lock().unwrap();
719        (
720            state.watched.clone(),
721            state.page_requests.clone(),
722            state.content_requests.clone(),
723            state
724                .raw_views
725                .iter()
726                .map(|(run_id, (_, revision, _, _))| (run_id.clone(), *revision))
727                .collect::<HashMap<_, _>>(),
728            state.run_list_request.clone(),
729        )
730    };
731    let mut requests = Vec::new();
732    if !sent.contains("runs") {
733        requests.push(request(
734            counter,
735            "view.runs.watch",
736            None,
737            json!({"subscriptionId":"runs"}),
738        ));
739        sent.insert("runs".to_string());
740    }
741    sent.retain(|key| {
742        !key.starts_with("runs-page:")
743            || run_list_request
744                .as_ref()
745                .is_some_and(|(revision, cursor)| key == &format!("runs-page:{revision}:{cursor}"))
746    });
747    if let Some((revision, cursor)) = run_list_request {
748        let key = format!("runs-page:{revision}:{cursor}");
749        if sent.insert(key) {
750            requests.push(request(
751                counter,
752                "view.runs.page",
753                None,
754                json!({"revision":revision,"cursor":cursor}),
755            ));
756        }
757    }
758    let removals: Vec<String> = sent
759        .iter()
760        .filter_map(|id| id.strip_prefix("run:").map(str::to_string))
761        .filter(|run_id| !watched.contains(run_id))
762        .collect();
763    for run_id in removals {
764        requests.push(request(
765            counter,
766            "view.run.unwatch",
767            Some(&run_id),
768            json!({"subscriptionId":format!("run:{run_id}")}),
769        ));
770        sent.remove(&format!("run:{run_id}"));
771    }
772    for run_id in watched {
773        let subscription = format!("run:{run_id}");
774        if sent.insert(subscription.clone()) {
775            requests.push(request(
776                counter,
777                "view.run.watch",
778                Some(&run_id),
779                json!({"subscriptionId":subscription}),
780            ));
781        }
782    }
783    let desired_page_keys = pages
784        .iter()
785        .map(|((run_id, kind), cursor)| {
786            let revision = revisions.get(run_id).copied().unwrap_or(0);
787            format!("page:{run_id}:{}:{cursor}:{revision}", page_name(*kind))
788        })
789        .collect::<HashSet<_>>();
790    sent.retain(|key| !key.starts_with("page:") || desired_page_keys.contains(key));
791    for ((run_id, kind), cursor) in pages {
792        let revision = revisions.get(&run_id).copied().unwrap_or(0);
793        let key = format!("page:{run_id}:{}:{cursor}:{revision}", page_name(kind));
794        if sent.insert(key) {
795            requests.push(request(
796                counter,
797                "view.page",
798                Some(&run_id),
799                json!({"kind":page_name(kind),"cursor":cursor}),
800            ));
801        }
802    }
803    let desired_content_keys = contents
804        .iter()
805        .map(|((run_id, path), offset)| format!("content:{run_id}:{path}:{offset}"))
806        .collect::<HashSet<_>>();
807    sent.retain(|key| !key.starts_with("content:") || desired_content_keys.contains(key));
808    for ((run_id, path), offset) in contents {
809        let key = format!("content:{run_id}:{path}:{offset}");
810        if sent.insert(key) {
811            requests.push(request(
812                counter,
813                "view.content",
814                Some(&run_id),
815                json!({"path":path,"offset":offset}),
816            ));
817        }
818    }
819    requests
820}
821
822fn request(
823    counter: &mut u64,
824    operation: &str,
825    run_id: Option<&str>,
826    payload: Value,
827) -> ClientRequest {
828    *counter = counter.wrapping_add(1);
829    let id = format!("piw-{}-{counter}", std::process::id());
830    ClientRequest {
831        schema: PROTOCOL_ID.to_string(),
832        message_type: "request".to_string(),
833        request_id: id.clone(),
834        client_id: format!("piw-{}", std::process::id()),
835        operation: operation.to_string(),
836        idempotency_key: id,
837        run_id: run_id.map(str::to_string),
838        expected_revision: None,
839        payload,
840    }
841}
842
843fn handle_server_message(text: &str, shared: &Arc<Mutex<Shared>>) -> Result<()> {
844    match parse_server_message(text).map_err(anyhow::Error::msg)? {
845        ServerMessage::Hello(hello) => {
846            anyhow::ensure!(
847                hello.package_version == env!("CARGO_PKG_VERSION"),
848                "workflow client version mismatch: host {}, piw {}",
849                hello.package_version,
850                env!("CARGO_PKG_VERSION")
851            );
852            let mut state = shared.lock().unwrap();
853            state.connected = true;
854            state.connecting = false;
855            state.reconnect_attempt = 0;
856            state.error = None;
857        }
858        ServerMessage::Event(event) => match event.event.as_str() {
859            "runs" => {
860                start_run_list(&mut shared.lock().unwrap(), &event.payload)?;
861            }
862            "run_snapshot" => {
863                let Some(run_id) = event.run_id else {
864                    return Ok(());
865                };
866                store_snapshot(
867                    &mut shared.lock().unwrap(),
868                    run_id,
869                    event.revision.unwrap_or(0),
870                    event.payload,
871                );
872            }
873            "run_patch" => {
874                let Some(run_id) = event.run_id else {
875                    return Ok(());
876                };
877                let targets: Vec<TargetPatch> = serde_json::from_value(event.payload)?;
878                let mut state = shared.lock().unwrap();
879                if let Some((event_revision, view_revision, generation, view)) =
880                    state.raw_views.get_mut(&run_id)
881                {
882                    if event.revision == Some(*event_revision + 1)
883                        && apply_target_patches(view, &targets).is_ok()
884                    {
885                        *event_revision += 1;
886                        *view_revision = view
887                            .get("revision")
888                            .and_then(Value::as_u64)
889                            .unwrap_or(view_revision.saturating_add(1));
890                        *generation = generation.wrapping_add(1);
891                    }
892                }
893            }
894            _ => {}
895        },
896        ServerMessage::Response(response) => {
897            if response.outcome == "unavailable"
898                || response.outcome == "rejected"
899                || response.outcome == "notFound"
900            {
901                shared.lock().unwrap().error = Some(
902                    response
903                        .error
904                        .unwrap_or_else(|| format!("workflow request was {}", response.outcome)),
905                );
906            } else if let Some(receipt) = response.receipt {
907                match receipt.get("schema").and_then(Value::as_str) {
908                    Some("pi-workflows.run-list-page.v1") => {
909                        if response.outcome == "accepted" {
910                            merge_run_list(&mut shared.lock().unwrap(), &receipt)?;
911                        } else {
912                            shared.lock().unwrap().run_list_request = None;
913                        }
914                    }
915                    Some("pi-workflows.run-page.v1") => {
916                        if response.outcome == "accepted" {
917                            merge_page(&mut shared.lock().unwrap(), &receipt)?;
918                        }
919                    }
920                    Some("pi-workflows.content-chunk.v1") => {
921                        merge_content(&mut shared.lock().unwrap(), &receipt)?;
922                    }
923                    _ => {}
924                }
925            }
926        }
927    }
928    Ok(())
929}
930
931fn start_run_list(state: &mut Shared, page: &Value) -> Result<()> {
932    let revision = page
933        .get("revision")
934        .and_then(Value::as_str)
935        .context("run list page has no revision")?
936        .to_string();
937    let start = page
938        .get("start")
939        .and_then(Value::as_u64)
940        .context("run list page has no start")?;
941    let total = page
942        .get("total")
943        .and_then(Value::as_u64)
944        .context("run list page has no total")?;
945    let items = page
946        .get("items")
947        .and_then(Value::as_array)
948        .cloned()
949        .context("run list page items are not an array")?;
950    anyhow::ensure!(start == 0, "run list snapshot does not start at zero");
951    anyhow::ensure!(
952        items.len() as u64 <= total,
953        "run list page exceeds its total"
954    );
955    if items.len() as u64 == total {
956        state.summaries = items;
957        state.run_list = None;
958        state.run_list_request = None;
959    } else {
960        anyhow::ensure!(!items.is_empty(), "run list page made no progress");
961        let cursor = items.len() as u64;
962        state.run_list = Some(RunListAssembly {
963            revision: revision.clone(),
964            total,
965            items,
966        });
967        state.run_list_request = Some((revision, cursor));
968    }
969    Ok(())
970}
971
972fn merge_run_list(state: &mut Shared, page: &Value) -> Result<()> {
973    let revision = page
974        .get("revision")
975        .and_then(Value::as_str)
976        .context("run list page has no revision")?;
977    let start = page
978        .get("start")
979        .and_then(Value::as_u64)
980        .context("run list page has no start")?;
981    let total = page
982        .get("total")
983        .and_then(Value::as_u64)
984        .context("run list page has no total")?;
985    let items = page
986        .get("items")
987        .and_then(Value::as_array)
988        .cloned()
989        .context("run list page items are not an array")?;
990    let Some(assembly) = state.run_list.as_mut() else {
991        return Ok(());
992    };
993    if assembly.revision != revision {
994        return Ok(());
995    }
996    let expected = assembly.items.len() as u64;
997    anyhow::ensure!(
998        start == expected,
999        "run list page does not continue the snapshot"
1000    );
1001    anyhow::ensure!(
1002        total == assembly.total,
1003        "run list total changed while paging"
1004    );
1005    anyhow::ensure!(!items.is_empty(), "run list page made no progress");
1006    assembly.items.extend(items);
1007    anyhow::ensure!(
1008        assembly.items.len() as u64 <= assembly.total,
1009        "run list page exceeds its total"
1010    );
1011    if assembly.items.len() as u64 == assembly.total {
1012        state.summaries = std::mem::take(&mut assembly.items);
1013        state.run_list = None;
1014        state.run_list_request = None;
1015    } else {
1016        state.run_list_request = Some((revision.to_string(), assembly.items.len() as u64));
1017    }
1018    Ok(())
1019}
1020
1021fn merge_content(state: &mut Shared, receipt: &Value) -> Result<()> {
1022    let run_id = receipt
1023        .get("runId")
1024        .and_then(Value::as_str)
1025        .context("content chunk has no runId")?;
1026    let path = receipt
1027        .get("path")
1028        .and_then(Value::as_str)
1029        .context("content chunk has no path")?;
1030    let offset = receipt
1031        .get("offset")
1032        .and_then(Value::as_u64)
1033        .context("content chunk has no offset")?;
1034    let next_offset = receipt
1035        .get("nextOffset")
1036        .and_then(Value::as_u64)
1037        .context("content chunk has no nextOffset")?;
1038    let total = receipt
1039        .get("bytes")
1040        .and_then(Value::as_u64)
1041        .context("content chunk has no byte total")?;
1042    let sha256 = receipt
1043        .get("sha256")
1044        .and_then(Value::as_str)
1045        .context("content chunk has no digest")?;
1046    let complete = receipt
1047        .get("complete")
1048        .and_then(Value::as_bool)
1049        .context("content chunk has no completion marker")?;
1050    let data = BASE64
1051        .decode(
1052            receipt
1053                .get("data")
1054                .and_then(Value::as_str)
1055                .context("content chunk has no data")?,
1056        )
1057        .context("content chunk is not valid base64")?;
1058    let key = (run_id.to_string(), path.to_string());
1059    let Some(entry) = state.artifacts.remove(&key) else {
1060        return Ok(());
1061    };
1062    let ArtifactEntry::Loading(mut bytes) = entry else {
1063        state.artifacts.insert(key, entry);
1064        return Ok(());
1065    };
1066    if bytes.len() as u64 != offset || offset.saturating_add(data.len() as u64) != next_offset {
1067        state.artifacts.insert(
1068            key.clone(),
1069            ArtifactEntry::Error("workflow content chunk offset is invalid".to_string()),
1070        );
1071        state.content_requests.remove(&key);
1072        bump_view_generation(state, run_id);
1073        return Ok(());
1074    }
1075    bytes.extend_from_slice(&data);
1076    if complete {
1077        let digest = hex_sha256(&bytes);
1078        if bytes.len() as u64 != total || next_offset != total || digest != sha256 {
1079            state.artifacts.insert(
1080                key.clone(),
1081                ArtifactEntry::Error("workflow content digest does not match".to_string()),
1082            );
1083        } else {
1084            match String::from_utf8(bytes) {
1085                Ok(content) => {
1086                    state
1087                        .artifacts
1088                        .insert(key.clone(), ArtifactEntry::Ready(content));
1089                }
1090                Err(_) => {
1091                    state.artifacts.insert(
1092                        key.clone(),
1093                        ArtifactEntry::Error("workflow content is not UTF-8".to_string()),
1094                    );
1095                }
1096            }
1097        }
1098        state.content_requests.remove(&key);
1099        bump_view_generation(state, run_id);
1100    } else {
1101        state
1102            .artifacts
1103            .insert(key.clone(), ArtifactEntry::Loading(bytes));
1104        state.content_requests.insert(key, next_offset);
1105    }
1106    Ok(())
1107}
1108
1109fn hex_sha256(bytes: &[u8]) -> String {
1110    Sha256::digest(bytes)
1111        .iter()
1112        .map(|byte| format!("{byte:02x}"))
1113        .collect()
1114}
1115
1116fn merge_page(state: &mut Shared, receipt: &Value) -> Result<()> {
1117    let run_id = receipt
1118        .get("runId")
1119        .and_then(Value::as_str)
1120        .context("run page has no runId")?;
1121    let kind = receipt
1122        .get("kind")
1123        .and_then(Value::as_str)
1124        .context("run page has no kind")?;
1125    let start = receipt.get("start").and_then(Value::as_u64).unwrap_or(0);
1126    let total = receipt.get("total").and_then(Value::as_u64).unwrap_or(0);
1127    let items = receipt
1128        .get("items")
1129        .and_then(Value::as_array)
1130        .cloned()
1131        .context("run page items are not an array")?;
1132    let page_kind = page_kind(kind).context("run page kind is invalid")?;
1133    let cursor = receipt
1134        .get("cursor")
1135        .and_then(Value::as_u64)
1136        .context("run page has no cursor")?;
1137    let revision = receipt
1138        .get("revision")
1139        .and_then(Value::as_u64)
1140        .context("run page has no revision")?;
1141    if state.page_requests.get(&(run_id.to_string(), page_kind)) != Some(&cursor) {
1142        return Ok(());
1143    }
1144    let Some((_, view_revision, generation, view)) = state.raw_views.get_mut(run_id) else {
1145        return Ok(());
1146    };
1147    if *view_revision != revision {
1148        return Ok(());
1149    }
1150
1151    match kind {
1152        "steps" => {
1153            view["graphSteps"] = receipt
1154                .get("graphSteps")
1155                .cloned()
1156                .unwrap_or_else(|| Value::Array(items.clone()));
1157            view["stepStart"] = json!(start);
1158            view["stepTotal"] = json!(total);
1159            view["graphCursor"] = receipt.get("graphCursor").cloned().unwrap_or(json!(0));
1160            view["takenTransitions"] = receipt
1161                .get("takenTransitions")
1162                .cloned()
1163                .unwrap_or_else(|| json!([]));
1164            if let Some(workflow_state) = view.get_mut("state").and_then(Value::as_object_mut) {
1165                workflow_state.insert("steps".to_string(), Value::Array(items));
1166            }
1167        }
1168        "trace" | "trace_at_step" => {
1169            view["tracePage"] = json!({"start":start,"total":total,"items":items});
1170        }
1171        "session_entries" | "session_events" => {
1172            let session = view
1173                .get_mut("session")
1174                .and_then(Value::as_object_mut)
1175                .context("run view has no session object")?;
1176            session.insert(
1177                if kind == "session_entries" {
1178                    "entryPage".to_string()
1179                } else {
1180                    "eventPage".to_string()
1181                },
1182                json!({"start":start,"total":total,"items":items}),
1183            );
1184            if kind == "session_events" {
1185                session.insert(
1186                    "replayCheckpoint".to_string(),
1187                    receipt
1188                        .get("replayCheckpoint")
1189                        .cloned()
1190                        .unwrap_or(Value::Null),
1191                );
1192            }
1193        }
1194        "settings" => {
1195            view["settingsScopes"] = Value::Array(items);
1196            view["settingsStart"] = json!(start);
1197            view["settingsTotal"] = json!(total);
1198        }
1199        "follow_ups" => {
1200            if !view.get("followUpQueue").is_some_and(Value::is_object) {
1201                view["followUpQueue"] = json!({});
1202            }
1203            if let Some(queue) = view.get_mut("followUpQueue").and_then(Value::as_object_mut) {
1204                queue.insert("items".to_string(), Value::Array(items));
1205            }
1206            view["followUpStart"] = json!(start);
1207            view["followUpTotal"] = json!(total);
1208        }
1209        "updates" => {
1210            view["updates"] = Value::Array(items.clone());
1211            view["updateStart"] = json!(start);
1212            view["updateTotal"] = json!(total);
1213            if let Some(workflow_state) = view.get_mut("state").and_then(Value::as_object_mut) {
1214                workflow_state.insert("updates".to_string(), Value::Array(items));
1215            }
1216        }
1217        _ => anyhow::bail!("unsupported run page kind {kind}"),
1218    }
1219    *generation = generation.wrapping_add(1);
1220    Ok(())
1221}
1222
1223fn bump_view_generation(state: &mut Shared, run_id: &str) {
1224    if let Some((_, _, generation, _)) = state.raw_views.get_mut(run_id) {
1225        *generation = generation.wrapping_add(1);
1226    }
1227}
1228
1229fn store_snapshot(state: &mut Shared, run_id: String, event_revision: u64, value: Value) {
1230    state.next_view_generation = state.next_view_generation.wrapping_add(1);
1231    let generation = state.next_view_generation;
1232    let view_revision = value
1233        .get("revision")
1234        .and_then(Value::as_u64)
1235        .unwrap_or(event_revision);
1236    state
1237        .raw_views
1238        .insert(run_id, (event_revision, view_revision, generation, value));
1239}
1240
1241fn page_kind(value: &str) -> Option<PageKind> {
1242    match value {
1243        "steps" => Some(PageKind::Steps),
1244        "trace" => Some(PageKind::Trace),
1245        "trace_at_step" => Some(PageKind::TraceAtStep),
1246        "session_entries" => Some(PageKind::SessionEntries),
1247        "session_events" => Some(PageKind::SessionEvents),
1248        "settings" => Some(PageKind::Settings),
1249        "follow_ups" => Some(PageKind::FollowUps),
1250        "updates" => Some(PageKind::Updates),
1251        _ => None,
1252    }
1253}
1254
1255fn page_name(kind: PageKind) -> &'static str {
1256    match kind {
1257        PageKind::Steps => "steps",
1258        PageKind::Trace => "trace",
1259        PageKind::TraceAtStep => "trace_at_step",
1260        PageKind::SessionEntries => "session_entries",
1261        PageKind::SessionEvents => "session_events",
1262        PageKind::Settings => "settings",
1263        PageKind::FollowUps => "follow_ups",
1264        PageKind::Updates => "updates",
1265    }
1266}
1267
1268#[cfg(test)]
1269mod tests {
1270    use super::*;
1271    use crate::protocol::PatchOp;
1272
1273    #[test]
1274    fn missing_run_watch_sets_a_visible_client_error() {
1275        let shared = Arc::new(Mutex::new(Shared::default()));
1276        let message = crate::protocol::canonical_json(&json!({
1277            "schema": PROTOCOL_ID,
1278            "type": "response",
1279            "requestId": "watch-missing",
1280            "outcome": "notFound",
1281            "error": "Workflow run not found"
1282        }))
1283        .unwrap();
1284        handle_server_message(&message, &shared).unwrap();
1285        assert_eq!(
1286            shared.lock().unwrap().error.as_deref(),
1287            Some("Workflow run not found")
1288        );
1289    }
1290
1291    #[test]
1292    fn decodes_the_host_owned_run_view_contract() {
1293        let source = json!({"kind":"file","path":"/tmp/smoke.workflow.ts","hash":"abc"});
1294        let raw = json!({
1295            "manifest": {
1296                "schema":"pi-workflows.run-manifest.v1",
1297                "runId":"run-1",
1298                "workflowName":"smoke",
1299                "workflowSource":source,
1300                "startedAt":"2026-01-01T00:00:00.000Z",
1301                "finishedAt":"2026-01-01T00:00:01.000Z",
1302                "status":"completed",
1303                "traceSchema":"pi-workflows.trace-event.v1",
1304                "paths":{"workflow":"host","state":"host","trace":"host"}
1305            },
1306            "state": {
1307                "schema":"pi-workflows.run-state.v1",
1308                "traceSeq":1,
1309                "runId":"run-1",
1310                "workflowName":"smoke",
1311                "workflowSource":source,
1312                "startedAt":"2026-01-01T00:00:00.000Z",
1313                "finishedAt":"2026-01-01T00:00:01.000Z",
1314                "updatedAt":"2026-01-01T00:00:01.000Z",
1315                "status":"completed",
1316                "input":{},
1317                "outputs":{},
1318                "results":{},
1319                "steps":[],
1320                "updates":[{"seq":299}]
1321            },
1322            "workflow": {
1323                "schema":"pi-workflows.definition-snapshot.v1",
1324                "name":"smoke",
1325                "startAt":"done",
1326                "nodes":{"done":{"nodeType":"compute"}},
1327                "edges":[]
1328            },
1329            "graphSteps":[],
1330            "takenTransitions":[],
1331            "graphCursor":0,
1332            "stepStart":0,
1333            "stepTotal":0,
1334            "tracePage":{"start":0,"total":0,"items":[]},
1335            "session":{
1336                "binding":null,
1337                "entryPage":{"start":0,"total":0,"items":[]},
1338                "eventPage":{"start":0,"total":0,"items":[]},
1339                "capture":null,
1340                "integrity":null
1341            },
1342            "settingsScopes":[],
1343            "settingsStart":0,
1344            "settingsTotal":0,
1345            "followUpQueue":null,
1346            "followUpStart":0,
1347            "followUpTotal":0,
1348            "updateStart":299,
1349            "updateTotal":300,
1350            "live":false,
1351            "possiblyInterrupted":false
1352        });
1353
1354        let view = decode_view(4, 1, &raw, None, None).expect("host view should decode");
1355        assert_eq!(view.manifest.workflow_name, "smoke");
1356        assert_eq!(view.state.status, crate::state::types::RunStatus::Completed);
1357        assert_eq!(view.update_total, 300);
1358    }
1359
1360    #[test]
1361    fn large_workflow_definitions_load_before_layout() {
1362        let nodes = (0..300)
1363            .map(|index| (format!("node-{index}"), json!({"nodeType":"compute"})))
1364            .collect::<serde_json::Map<_, _>>();
1365        let full = json!({
1366            "schema":"pi-workflows.definition-snapshot.v1",
1367            "name":"large",
1368            "startAt":"node-0",
1369            "nodes":nodes,
1370            "edges":[]
1371        });
1372        let content = serde_json::to_string(&full).unwrap();
1373        let path = "artifacts/sha256/definition.json";
1374        let raw = json!({
1375            "workflow": {
1376                "schema":"pi-workflows.definition-snapshot.v1",
1377                "name":"large",
1378                "startAt":"node-0",
1379                "nodes":{"node-0":{"nodeType":"compute"}},
1380                "edges":[],
1381                "content": {
1382                    "$artifact": {
1383                        "path":path,
1384                        "mediaType":"application/json",
1385                        "bytes":content.len(),
1386                        "sha256":hex_sha256(content.as_bytes()),
1387                        "opaque":true
1388                    }
1389                }
1390            }
1391        });
1392        let mut state = Shared::default();
1393        let (missing, requested) = workflow_definition_content(&mut state, "run-1", &raw);
1394        assert!(missing.is_none());
1395        assert!(requested);
1396        assert_eq!(
1397            state
1398                .content_requests
1399                .get(&("run-1".to_string(), path.to_string())),
1400            Some(&0)
1401        );
1402        state.artifacts.insert(
1403            ("run-1".to_string(), path.to_string()),
1404            ArtifactEntry::Ready(content),
1405        );
1406        let (loaded, requested) = workflow_definition_content(&mut state, "run-1", &raw);
1407        assert!(!requested);
1408        let snapshot: DefinitionSnapshot =
1409            serde_json::from_value(workflow_definition_value(&raw, loaded.as_deref()).unwrap())
1410                .unwrap();
1411        assert_eq!(snapshot.nodes.len(), 300);
1412        assert_eq!(layout_graph(&snapshot).rank_of_node.len(), 300);
1413
1414        let mut invalid = raw;
1415        invalid["workflow"]["content"]["$artifact"]["sha256"] = json!("0".repeat(64));
1416        let (loaded, requested) = workflow_definition_content(&mut state, "run-1", &invalid);
1417        assert!(loaded.is_none());
1418        assert!(!requested);
1419        assert!(state
1420            .error
1421            .as_deref()
1422            .is_some_and(|error| error.contains("does not match")));
1423    }
1424
1425    #[test]
1426    fn complete_graph_history_uses_the_shared_content_loader() {
1427        let history = json!({
1428            "steps":[],
1429            "transitions":(0..300).map(|index| format!("node-{index}->node-{}", index + 1)).collect::<Vec<_>>()
1430        });
1431        let content = serde_json::to_string(&history).unwrap();
1432        let path = "artifacts/sha256/graph-history.json";
1433        let raw = json!({
1434            "graphSteps":[],
1435            "takenTransitions":["node-0->node-1"],
1436            "graphHistory": {
1437                "$artifact": {
1438                    "path":path,
1439                    "mediaType":"application/json",
1440                    "bytes":content.len(),
1441                    "sha256":hex_sha256(content.as_bytes()),
1442                    "opaque":true
1443                }
1444            }
1445        });
1446        let mut state = Shared::default();
1447        let (missing, requested) = graph_history_content(&mut state, "run-1", &raw);
1448        assert!(missing.is_none());
1449        assert!(requested);
1450        state.artifacts.insert(
1451            ("run-1".to_string(), path.to_string()),
1452            ArtifactEntry::Ready(content),
1453        );
1454        let (loaded, requested) = graph_history_content(&mut state, "run-1", &raw);
1455        assert!(!requested);
1456        let complete = graph_history_value(&raw, loaded.as_deref()).unwrap();
1457        assert_eq!(complete["transitions"].as_array().map(Vec::len), Some(300));
1458    }
1459
1460    #[test]
1461    fn a_run_page_updates_only_its_selected_window() {
1462        let mut state = Shared::default();
1463        state.raw_views.insert(
1464            "run-1".to_string(),
1465            (
1466                3,
1467                3,
1468                1,
1469                json!({
1470                    "state":{"steps":[],"updates":[]},
1471                    "session":{"entryPage":{"items":[]},"eventPage":{"items":[]}},
1472                    "settingsScopes":[{"change":1}],
1473                    "settingsStart":0,
1474                    "settingsTotal":1
1475                }),
1476            ),
1477        );
1478        state
1479            .page_requests
1480            .insert(("run-1".to_string(), PageKind::Settings), 44);
1481        state
1482            .page_requests
1483            .insert(("run-1".to_string(), PageKind::Steps), 0);
1484        state
1485            .page_requests
1486            .insert(("run-1".to_string(), PageKind::FollowUps), 2);
1487        state
1488            .page_requests
1489            .insert(("run-1".to_string(), PageKind::SessionEvents), 256);
1490        merge_page(
1491            &mut state,
1492            &json!({
1493                "schema":"pi-workflows.run-page.v1",
1494                "runId":"run-1",
1495                "revision":3,
1496                "kind":"settings",
1497                "cursor":44,
1498                "start":44,
1499                "total":300,
1500                "items":[{"change":299}]
1501            }),
1502        )
1503        .unwrap();
1504        merge_page(
1505            &mut state,
1506            &json!({
1507                "schema":"pi-workflows.run-page.v1",
1508                "runId":"run-1",
1509                "revision":3,
1510                "kind":"steps",
1511                "cursor":0,
1512                "start":0,
1513                "total":1,
1514                "items":[{"step":0}],
1515                "graphSteps":[{"node":"one"}],
1516                "graphCursor":0,
1517                "takenTransitions":[]
1518            }),
1519        )
1520        .unwrap();
1521        merge_page(
1522            &mut state,
1523            &json!({
1524                "schema":"pi-workflows.run-page.v1",
1525                "runId":"run-1",
1526                "revision":3,
1527                "kind":"follow_ups",
1528                "cursor":2,
1529                "start":2,
1530                "total":3,
1531                "items":[{"followUpId":"follow-3"}]
1532            }),
1533        )
1534        .unwrap();
1535        merge_page(
1536            &mut state,
1537            &json!({
1538                "schema":"pi-workflows.run-page.v1",
1539                "runId":"run-1",
1540                "revision":3,
1541                "kind":"session_events",
1542                "cursor":256,
1543                "start":256,
1544                "total":300,
1545                "items":[{"seq":257}],
1546                "replayCheckpoint":{"throughSeq":256}
1547            }),
1548        )
1549        .unwrap();
1550        let (_, _, generation, view) = state.raw_views.get("run-1").unwrap();
1551        assert_eq!(*generation, 5);
1552        assert_eq!(view["settingsStart"], 44);
1553        assert_eq!(view["settingsScopes"], json!([{"change":299}]));
1554        assert_eq!(view["graphSteps"], json!([{"node":"one"}]));
1555        assert_eq!(view["state"]["steps"], json!([{"step":0}]));
1556        assert_eq!(
1557            view["followUpQueue"]["items"],
1558            json!([{"followUpId":"follow-3"}])
1559        );
1560        assert_eq!(view["session"]["replayCheckpoint"]["throughSeq"], 256);
1561    }
1562
1563    #[test]
1564    fn stale_run_pages_cannot_replace_the_requested_window() {
1565        let mut state = Shared::default();
1566        state.raw_views.insert(
1567            "run-1".to_string(),
1568            (
1569                1,
1570                3,
1571                1,
1572                json!({"state":{"steps":[]},"graphSteps":[],"takenTransitions":[]}),
1573            ),
1574        );
1575        state
1576            .page_requests
1577            .insert(("run-1".to_string(), PageKind::Steps), 20);
1578        let receipt = |cursor: u64, revision: u64, item: u64| {
1579            json!({
1580                "schema":"pi-workflows.run-page.v1",
1581                "runId":"run-1",
1582                "revision":revision,
1583                "kind":"steps",
1584                "cursor":cursor,
1585                "start":cursor,
1586                "total":100,
1587                "items":[{"step":item}],
1588                "graphSteps":[],
1589                "graphCursor":cursor,
1590                "takenTransitions":[]
1591            })
1592        };
1593        merge_page(&mut state, &receipt(10, 3, 10)).unwrap();
1594        merge_page(&mut state, &receipt(20, 2, 2)).unwrap();
1595        assert_eq!(state.raw_views["run-1"].2, 1);
1596        merge_page(&mut state, &receipt(20, 3, 20)).unwrap();
1597        assert_eq!(state.raw_views["run-1"].2, 2);
1598        assert_eq!(
1599            state.raw_views["run-1"].3["state"]["steps"],
1600            json!([{"step":20}])
1601        );
1602    }
1603
1604    #[test]
1605    fn paged_run_lists_publish_only_complete_matching_revisions() {
1606        let mut state = Shared::default();
1607        start_run_list(
1608            &mut state,
1609            &json!({
1610                "schema":"pi-workflows.run-list-page.v1",
1611                "revision":"3:10:20",
1612                "start":0,
1613                "total":3,
1614                "items":[{"runId":"run-1"},{"runId":"run-2"}]
1615            }),
1616        )
1617        .unwrap();
1618        assert!(state.summaries.is_empty());
1619        assert_eq!(state.run_list_request, Some(("3:10:20".to_string(), 2)));
1620        merge_run_list(
1621            &mut state,
1622            &json!({
1623                "schema":"pi-workflows.run-list-page.v1",
1624                "revision":"stale",
1625                "start":2,
1626                "total":3,
1627                "items":[{"runId":"stale"}]
1628            }),
1629        )
1630        .unwrap();
1631        assert!(state.summaries.is_empty());
1632        merge_run_list(
1633            &mut state,
1634            &json!({
1635                "schema":"pi-workflows.run-list-page.v1",
1636                "revision":"3:10:20",
1637                "start":2,
1638                "total":3,
1639                "items":[{"runId":"run-3"}]
1640            }),
1641        )
1642        .unwrap();
1643        assert_eq!(state.summaries.len(), 3);
1644        assert!(state.run_list_request.is_none());
1645    }
1646
1647    #[test]
1648    fn switching_runs_discards_old_pages_content_and_artifacts() {
1649        let mut state = Shared::default();
1650        state
1651            .raw_views
1652            .insert("old".to_string(), (1, 1, 1, json!({})));
1653        state
1654            .page_requests
1655            .insert(("old".to_string(), PageKind::Steps), 10);
1656        state
1657            .content_requests
1658            .insert(("old".to_string(), "old.json".to_string()), 0);
1659        state.artifacts.insert(
1660            ("old".to_string(), "old.json".to_string()),
1661            ArtifactEntry::Ready("old".to_string()),
1662        );
1663        state
1664            .page_requests
1665            .insert(("current".to_string(), PageKind::Steps), 20);
1666
1667        discard_run_state(&mut state, "old");
1668
1669        assert!(!state.raw_views.contains_key("old"));
1670        assert!(state
1671            .page_requests
1672            .keys()
1673            .all(|(run_id, _)| run_id != "old"));
1674        assert!(state
1675            .content_requests
1676            .keys()
1677            .all(|(run_id, _)| run_id != "old"));
1678        assert!(state.artifacts.keys().all(|(run_id, _)| run_id != "old"));
1679        assert!(state
1680            .page_requests
1681            .contains_key(&("current".to_string(), PageKind::Steps)));
1682    }
1683
1684    #[test]
1685    fn page_requests_can_return_to_an_earlier_window() {
1686        let shared = Arc::new(Mutex::new(Shared::default()));
1687        shared
1688            .lock()
1689            .unwrap()
1690            .raw_views
1691            .insert("run-1".to_string(), (7, 7, 1, json!({})));
1692        let mut sent = HashSet::new();
1693        let mut counter = 0;
1694        let request_at = |cursor: u64,
1695                          shared: &Arc<Mutex<Shared>>,
1696                          sent: &mut HashSet<String>,
1697                          counter: &mut u64| {
1698            shared
1699                .lock()
1700                .unwrap()
1701                .page_requests
1702                .insert(("run-1".to_string(), PageKind::Steps), cursor);
1703            reconcile_requests(shared, sent, counter)
1704                .into_iter()
1705                .filter(|request| request.operation == "view.page")
1706                .count()
1707        };
1708        assert_eq!(request_at(10, &shared, &mut sent, &mut counter), 1);
1709        assert_eq!(request_at(20, &shared, &mut sent, &mut counter), 1);
1710        assert_eq!(request_at(10, &shared, &mut sent, &mut counter), 1);
1711    }
1712
1713    #[test]
1714    fn content_chunks_are_reassembled_and_verified() {
1715        let content = br#"{"complete":true}"#.to_vec();
1716        let path = "artifacts/sha256/content.json";
1717        let key = ("run-1".to_string(), path.to_string());
1718        let mut state = Shared::default();
1719        state
1720            .raw_views
1721            .insert("run-1".to_string(), (1, 1, 1, json!({"workflow":null})));
1722        state
1723            .artifacts
1724            .insert(key.clone(), ArtifactEntry::Loading(Vec::new()));
1725        state.content_requests.insert(key.clone(), 0);
1726        merge_content(
1727            &mut state,
1728            &json!({
1729                "schema":"pi-workflows.content-chunk.v1",
1730                "runId":"run-1",
1731                "path":path,
1732                "mediaType":"application/json",
1733                "bytes":content.len(),
1734                "sha256":hex_sha256(&content),
1735                "offset":0,
1736                "nextOffset":content.len(),
1737                "complete":true,
1738                "data":BASE64.encode(&content)
1739            }),
1740        )
1741        .unwrap();
1742        assert!(
1743            matches!(state.artifacts.get(&key), Some(ArtifactEntry::Ready(value)) if value == "{\"complete\":true}")
1744        );
1745        assert!(!state.content_requests.contains_key(&key));
1746        assert_eq!(state.raw_views["run-1"].2, 2);
1747    }
1748
1749    #[test]
1750    fn a_bad_patch_keeps_the_last_good_view() {
1751        let mut view = json!({"presentationRevision":1});
1752        let before = view.clone();
1753        let result = apply_target_patches(
1754            &mut view,
1755            &[TargetPatch {
1756                target_type: "graph".to_string(),
1757                target_key: String::new(),
1758                patch: vec![PatchOp::Replace {
1759                    path: "/missing/value".to_string(),
1760                    value: json!(2),
1761                }],
1762            }],
1763        );
1764        assert!(result.is_err());
1765        assert_eq!(view, before);
1766    }
1767}