Skip to main content

yoagent_state/
graph.rs

1use crate::{ArtifactRef, NodeId, StateError, StateOp};
2use serde::{Deserialize, Serialize};
3use serde_json::{Map, Value as JsonValue};
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct Node {
8    pub id: NodeId,
9    pub kind: String,
10    pub props: JsonValue,
11    pub stale: bool,
12    pub tombstoned: bool,
13    pub artifacts: Vec<ArtifactRef>,
14}
15
16impl Node {
17    pub fn new(id: NodeId, kind: impl Into<String>, props: JsonValue) -> Self {
18        Self {
19            id,
20            kind: kind.into(),
21            props,
22            stale: false,
23            tombstoned: false,
24            artifacts: Vec::new(),
25        }
26    }
27}
28
29#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30pub struct Relation {
31    pub from: NodeId,
32    pub rel: String,
33    pub to: NodeId,
34    pub props: JsonValue,
35}
36
37impl Relation {
38    pub fn new(from: NodeId, rel: impl Into<String>, to: NodeId, props: JsonValue) -> Self {
39        Self {
40            from,
41            rel: rel.into(),
42            to,
43            props,
44        }
45    }
46}
47
48#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
49pub struct Graph {
50    pub nodes: HashMap<NodeId, Node>,
51    pub relations: Vec<Relation>,
52    pub version: u64,
53}
54
55pub type GraphSnapshot = Graph;
56
57impl Graph {
58    pub fn apply_ops(&mut self, ops: &[StateOp]) -> Result<(), StateError> {
59        for op in ops {
60            self.apply_op(op)?;
61            self.version += 1;
62        }
63        Ok(())
64    }
65
66    pub fn apply_op(&mut self, op: &StateOp) -> Result<(), StateError> {
67        match op {
68            StateOp::CreateNode { id, kind, props } => {
69                self.nodes.insert(
70                    id.clone(),
71                    Node::new(id.clone(), kind.clone(), props.clone()),
72                );
73            }
74            StateOp::UpdateNode { id, props } => {
75                let node = self
76                    .nodes
77                    .get_mut(id)
78                    .ok_or_else(|| StateError::NodeNotFound(id.clone()))?;
79                merge_json(&mut node.props, props.clone());
80            }
81            StateOp::TombstoneNode { id, reason } => {
82                let node = self
83                    .nodes
84                    .get_mut(id)
85                    .ok_or_else(|| StateError::NodeNotFound(id.clone()))?;
86                node.tombstoned = true;
87                merge_json(
88                    &mut node.props,
89                    serde_json::json!({ "tombstone_reason": reason }),
90                );
91            }
92            StateOp::CreateRelation {
93                from,
94                rel,
95                to,
96                props,
97            } => {
98                self.relations.push(Relation::new(
99                    from.clone(),
100                    rel.clone(),
101                    to.clone(),
102                    props.clone(),
103                ));
104            }
105            StateOp::DeleteRelation { from, rel, to } => {
106                self.relations
107                    .retain(|r| &r.from != from || &r.rel != rel || &r.to != to);
108            }
109            StateOp::MarkStale { id, reason } => {
110                let node = self
111                    .nodes
112                    .get_mut(id)
113                    .ok_or_else(|| StateError::NodeNotFound(id.clone()))?;
114                node.stale = true;
115                merge_json(
116                    &mut node.props,
117                    serde_json::json!({ "stale_reason": reason }),
118                );
119            }
120            StateOp::AttachArtifact { id, artifact } => {
121                let node = self
122                    .nodes
123                    .get_mut(id)
124                    .ok_or_else(|| StateError::NodeNotFound(id.clone()))?;
125                node.artifacts.push(artifact.clone());
126            }
127        }
128
129        Ok(())
130    }
131
132    pub fn get_node(&self, id: &NodeId) -> Option<&Node> {
133        self.nodes.get(id)
134    }
135
136    pub fn outgoing(&self, id: &NodeId, rel: Option<&str>) -> Vec<Relation> {
137        self.relations
138            .iter()
139            .filter(|r| &r.from == id && rel.is_none_or(|expected| r.rel == expected))
140            .cloned()
141            .collect()
142    }
143
144    pub fn incoming(&self, id: &NodeId, rel: Option<&str>) -> Vec<Relation> {
145        self.relations
146            .iter()
147            .filter(|r| &r.to == id && rel.is_none_or(|expected| r.rel == expected))
148            .cloned()
149            .collect()
150    }
151
152    pub fn related(&self, id: &NodeId) -> Vec<Relation> {
153        self.relations
154            .iter()
155            .filter(|r| &r.from == id || &r.to == id)
156            .cloned()
157            .collect()
158    }
159}
160
161fn merge_json(target: &mut JsonValue, patch: JsonValue) {
162    match (target, patch) {
163        (JsonValue::Object(target), JsonValue::Object(patch)) => {
164            for (key, value) in patch {
165                merge_json(target.entry(key).or_insert(JsonValue::Null), value);
166            }
167        }
168        (slot, value) => *slot = value,
169    }
170}
171
172pub fn props() -> JsonValue {
173    JsonValue::Object(Map::new())
174}