Skip to main content

piw/
protocol.rs

1//! Wire types for the live replay protocol (`pi-workflows.replay.v1`), plus
2//! the JSON Patch subset (with the `append` extension op) used for view
3//! synchronization. See `docs/live-replay-protocol.md`.
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8pub const PROTOCOL_ID: &str = "pi-workflows.replay.v1";
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(tag = "type", rename_all = "snake_case")]
12pub enum ClientMessage {
13    WatchRuns,
14    WatchRun {
15        #[serde(rename = "runId")]
16        run_id: String,
17        #[serde(skip_serializing_if = "Option::is_none")]
18        revision: Option<u64>,
19        #[serde(rename = "stepCursor", skip_serializing_if = "Option::is_none")]
20        step_cursor: Option<u64>,
21        #[serde(rename = "traceCursor", skip_serializing_if = "Option::is_none")]
22        trace_cursor: Option<u64>,
23        #[serde(rename = "sessionEntryCursor", skip_serializing_if = "Option::is_none")]
24        session_entry_cursor: Option<u64>,
25        #[serde(rename = "sessionEventCursor", skip_serializing_if = "Option::is_none")]
26        session_event_cursor: Option<u64>,
27    },
28    UnwatchRun {
29        #[serde(rename = "runId")]
30        run_id: String,
31    },
32    FetchPage {
33        #[serde(rename = "runId")]
34        run_id: String,
35        kind: PageKind,
36        cursor: u64,
37    },
38    FetchArtifact {
39        #[serde(rename = "runId")]
40        run_id: String,
41        path: String,
42    },
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
46#[serde(tag = "type", rename_all = "snake_case")]
47pub enum ServerMessage {
48    Hello {
49        protocol: String,
50    },
51    Runs {
52        runs: Vec<Value>,
53    },
54    RunSnapshot {
55        #[serde(rename = "runId")]
56        run_id: String,
57        revision: u64,
58        view: Value,
59    },
60    RunPatch {
61        #[serde(rename = "runId")]
62        run_id: String,
63        revision: u64,
64        targets: Vec<TargetPatch>,
65    },
66    RunPage {
67        #[serde(rename = "runId")]
68        run_id: String,
69        revision: u64,
70        kind: PageKind,
71        cursor: u64,
72        start: u64,
73        total: u64,
74        items: Vec<Value>,
75        #[serde(rename = "graphCursor", skip_serializing_if = "Option::is_none")]
76        graph_cursor: Option<u64>,
77        #[serde(rename = "graphSteps", skip_serializing_if = "Option::is_none")]
78        graph_steps: Option<Vec<Value>>,
79        #[serde(rename = "takenTransitions", skip_serializing_if = "Option::is_none")]
80        taken_transitions: Option<Vec<String>>,
81        #[serde(rename = "replayCheckpoint", skip_serializing_if = "Option::is_none")]
82        replay_checkpoint: Option<Value>,
83    },
84    Artifact {
85        #[serde(rename = "runId")]
86        run_id: String,
87        path: String,
88        content: String,
89    },
90    Error {
91        message: String,
92        #[serde(rename = "runId", skip_serializing_if = "Option::is_none")]
93        run_id: Option<String>,
94    },
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
98#[serde(rename_all = "snake_case")]
99pub enum PageKind {
100    Steps,
101    Trace,
102    TraceAtStep,
103    SessionEntries,
104    SessionEvents,
105    Settings,
106    FollowUps,
107    Updates,
108}
109
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct TargetPatch {
112    #[serde(rename = "targetType")]
113    pub target_type: String,
114    #[serde(rename = "targetKey")]
115    pub target_key: String,
116    pub patch: Vec<PatchOp>,
117}
118
119/// RFC 6902 ops we use, plus `append` (add a batch of items to an array).
120#[derive(Debug, Clone, Serialize, Deserialize)]
121#[serde(tag = "op", rename_all = "snake_case")]
122pub enum PatchOp {
123    Add { path: String, value: Value },
124    Replace { path: String, value: Value },
125    Remove { path: String },
126    Append { path: String, value: Vec<Value> },
127}
128
129/// Apply a patch to a view document in place. Returns an error message on
130/// the first op that does not apply (the caller should resnapshot).
131pub fn apply_patch(target: &mut Value, patch: &[PatchOp]) -> Result<(), String> {
132    for op in patch {
133        match op {
134            PatchOp::Add { path, value } => {
135                set_path(target, path, value.clone(), false)?;
136            }
137            PatchOp::Replace { path, value } => {
138                set_path(target, path, value.clone(), true)?;
139            }
140            PatchOp::Remove { path } => {
141                remove_path(target, path)?;
142            }
143            PatchOp::Append { path, value } => {
144                let array = resolve_path(target, path)?
145                    .as_array_mut()
146                    .ok_or_else(|| format!("append target {path} is not an array"))?;
147                array.extend(value.iter().cloned());
148            }
149        }
150    }
151    Ok(())
152}
153
154fn unescape_token(token: &str) -> String {
155    token.replace("~1", "/").replace("~0", "~")
156}
157
158fn resolve_path<'a>(target: &'a mut Value, path: &str) -> Result<&'a mut Value, String> {
159    let mut current = target;
160    for token in path.split('/').skip(1) {
161        let token = unescape_token(token);
162        current = match current {
163            Value::Object(object) => object
164                .get_mut(&token)
165                .ok_or_else(|| format!("missing key {token} in {path}"))?,
166            Value::Array(array) => {
167                let index: usize = token
168                    .parse()
169                    .map_err(|_| format!("bad array index {token} in {path}"))?;
170                array
171                    .get_mut(index)
172                    .ok_or_else(|| format!("index {index} out of bounds in {path}"))?
173            }
174            _ => return Err(format!("cannot traverse into scalar at {token} in {path}")),
175        };
176    }
177    Ok(current)
178}
179
180fn set_path(target: &mut Value, path: &str, value: Value, replace: bool) -> Result<(), String> {
181    if path.is_empty() {
182        *target = value;
183        return Ok(());
184    }
185    let Some((parent_path, key)) = path.rsplit_once('/') else {
186        return Err(format!("bad path {path}"));
187    };
188    let parent = resolve_path(target, parent_path)?;
189    let key = unescape_token(key);
190    match parent {
191        Value::Object(object) => {
192            if replace && !object.contains_key(&key) {
193                return Err(format!("missing key {key} in {path}"));
194            }
195            object.insert(key, value);
196            Ok(())
197        }
198        Value::Array(array) => {
199            if !replace && key == "-" {
200                array.push(value);
201                return Ok(());
202            }
203            let index: usize = key
204                .parse()
205                .map_err(|_| format!("bad array index {key} in {path}"))?;
206            if replace {
207                let member = array
208                    .get_mut(index)
209                    .ok_or_else(|| format!("index {index} out of bounds in {path}"))?;
210                *member = value;
211                return Ok(());
212            }
213            if index > array.len() {
214                return Err(format!("index {index} out of bounds in {path}"));
215            }
216            array.insert(index, value);
217            Ok(())
218        }
219        _ => Err(format!("cannot set {key} on scalar in {path}")),
220    }
221}
222
223fn remove_path(target: &mut Value, path: &str) -> Result<(), String> {
224    let Some((parent_path, key)) = path.rsplit_once('/') else {
225        return Err(format!("bad path {path}"));
226    };
227    let parent = resolve_path(target, parent_path)?;
228    let key = unescape_token(key);
229    match parent {
230        Value::Object(object) => {
231            object
232                .remove(&key)
233                .ok_or_else(|| format!("missing key {key} in {path}"))?;
234            Ok(())
235        }
236        Value::Array(array) => {
237            let index: usize = key
238                .parse()
239                .map_err(|_| format!("bad array index {key} in {path}"))?;
240            if index >= array.len() {
241                return Err(format!("index {index} out of bounds in {path}"));
242            }
243            array.remove(index);
244            Ok(())
245        }
246        _ => Err(format!("cannot remove {key} from scalar in {path}")),
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use serde_json::json;
254
255    #[test]
256    fn applies_replace_append_and_remove() {
257        let mut view = json!({ "state": { "status": "running" }, "events": [1] });
258        apply_patch(
259            &mut view,
260            &[
261                PatchOp::Replace {
262                    path: "/state/status".into(),
263                    value: json!("completed"),
264                },
265                PatchOp::Append {
266                    path: "/events".into(),
267                    value: vec![json!(2), json!(3)],
268                },
269                PatchOp::Add {
270                    path: "/session".into(),
271                    value: json!({ "binding": null, "entries": [] }),
272                },
273            ],
274        )
275        .unwrap();
276        assert_eq!(
277            view,
278            json!({
279                "state": { "status": "completed" },
280                "events": [1, 2, 3],
281                "session": { "binding": null, "entries": [] }
282            })
283        );
284        apply_patch(
285            &mut view,
286            &[PatchOp::Remove {
287                path: "/session".into(),
288            }],
289        )
290        .unwrap();
291        assert!(view.get("session").is_none());
292    }
293
294    #[test]
295    fn array_add_inserts_but_replace_overwrites() {
296        let mut view = json!({ "events": [1, 2] });
297        apply_patch(
298            &mut view,
299            &[
300                PatchOp::Add {
301                    path: "/events/1".into(),
302                    value: json!(3),
303                },
304                PatchOp::Replace {
305                    path: "/events/0".into(),
306                    value: json!(4),
307                },
308            ],
309        )
310        .unwrap();
311        assert_eq!(view, json!({ "events": [4, 3, 2] }));
312    }
313
314    #[test]
315    fn escapes_json_pointer_tokens() {
316        let mut view = json!({ "a/b": { "c~d": 1 } });
317        apply_patch(
318            &mut view,
319            &[PatchOp::Replace {
320                path: "/a~1b/c~0d".into(),
321                value: json!(2),
322            }],
323        )
324        .unwrap();
325        assert_eq!(view, json!({ "a/b": { "c~d": 2 } }));
326    }
327
328    #[test]
329    fn gap_or_missing_path_errors() {
330        let mut view = json!({ "events": [] });
331        assert!(apply_patch(
332            &mut view,
333            &[PatchOp::Replace {
334                path: "/missing/deep".into(),
335                value: json!(1),
336            }],
337        )
338        .is_err());
339    }
340}