1use arrow::array::RecordBatch;
11use chrono::Utc;
12use futures::StreamExt;
13use serde::{Deserialize, Serialize};
14use std::sync::Arc;
15use tokio::sync::Mutex;
16
17pub const UPDATES_SUBJECT: &str = "nusy.code.graph.updates";
19pub const CONFLICT_SUBJECT: &str = "nusy.code.graph.conflict";
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct CodeGraphUpdate {
25 pub update_type: UpdateType,
27 pub node_id: String,
29 pub body_hash: String,
31 pub new_body: Option<String>,
33 pub timestamp_ms: i64,
35 pub agent: String,
37 pub rationale: String,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
43pub enum UpdateType {
44 NodeUpdated,
45 NodeCreated,
46 NodeDeleted,
47 EdgeCreated,
48 EdgeDeleted,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct ConflictEvent {
54 pub node_id: String,
55 pub local_hash: String,
56 pub remote_hash: String,
57 pub remote_agent: String,
58 pub resolution: String,
59 pub timestamp_ms: i64,
60}
61
62pub struct CodeGraphPublisher {
66 client: async_nats::Client,
67 agent_name: String,
68}
69
70impl CodeGraphPublisher {
71 pub async fn new(
73 nats_url: &str,
74 agent_name: &str,
75 ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
76 let client = async_nats::connect(nats_url).await?;
77 Ok(Self {
78 client,
79 agent_name: agent_name.to_string(),
80 })
81 }
82
83 pub fn client_ref(&self) -> &async_nats::Client {
85 &self.client
86 }
87
88 pub async fn publish(&self, update: &CodeGraphUpdate) -> Result<(), async_nats::PublishError> {
90 let data = serde_json::to_vec(update).unwrap_or_default();
91 self.client.publish(UPDATES_SUBJECT, data.into()).await
92 }
93
94 pub fn make_update(
96 &self,
97 node_id: &str,
98 new_body: &str,
99 body_hash: &str,
100 rationale: &str,
101 ) -> CodeGraphUpdate {
102 CodeGraphUpdate {
103 update_type: UpdateType::NodeUpdated,
104 node_id: node_id.to_string(),
105 body_hash: body_hash.to_string(),
106 new_body: Some(new_body.to_string()),
107 timestamp_ms: Utc::now().timestamp_millis(),
108 agent: self.agent_name.clone(),
109 rationale: rationale.to_string(),
110 }
111 }
112}
113
114pub type SharedGraphState = Arc<Mutex<SyncableGraph>>;
118
119pub struct SyncableGraph {
121 pub nodes: RecordBatch,
122 pub edges: RecordBatch,
123 pub body_hashes: std::collections::HashMap<String, String>,
125}
126
127pub async fn subscribe_and_apply(nats_url: &str, graph: SharedGraphState, own_agent: &str) {
129 let client = match async_nats::connect(nats_url).await {
130 Ok(c) => c,
131 Err(e) => {
132 tracing::error!(error = %e, "failed to connect for graph sync");
133 return;
134 }
135 };
136
137 let mut subscriber = match client.subscribe(UPDATES_SUBJECT).await {
138 Ok(s) => s,
139 Err(e) => {
140 tracing::error!(error = %e, "failed to subscribe to graph updates");
141 return;
142 }
143 };
144
145 let own_agent = own_agent.to_string();
146 tracing::info!(subject = UPDATES_SUBJECT, "graph sync subscriber active");
147
148 while let Some(msg) = subscriber.next().await {
149 let update: CodeGraphUpdate = match serde_json::from_slice(&msg.payload) {
150 Ok(u) => u,
151 Err(e) => {
152 tracing::warn!(error = %e, "invalid graph update payload");
153 continue;
154 }
155 };
156
157 if update.agent == own_agent {
159 tracing::debug!(node = %update.node_id, "skipping own update");
160 continue;
161 }
162
163 tracing::info!(
164 node = %update.node_id,
165 agent = %update.agent,
166 update_type = ?update.update_type,
167 "received graph update"
168 );
169
170 let mut graph = graph.lock().await;
171
172 if let Some(local_hash) = graph.body_hashes.get(&update.node_id)
174 && *local_hash != update.body_hash
175 && update.update_type == UpdateType::NodeUpdated
176 {
177 let conflict = ConflictEvent {
178 node_id: update.node_id.clone(),
179 local_hash: local_hash.clone(),
180 remote_hash: update.body_hash.clone(),
181 remote_agent: update.agent.clone(),
182 resolution: "last_write_wins".to_string(),
183 timestamp_ms: Utc::now().timestamp_millis(),
184 };
185 tracing::warn!(
186 node = %conflict.node_id,
187 local = %conflict.local_hash,
188 remote = %conflict.remote_hash,
189 "code graph conflict detected — last write wins"
190 );
191 if let Ok(data) = serde_json::to_vec(&conflict) {
192 let _ = client.publish(CONFLICT_SUBJECT, data.into()).await;
193 }
194 }
195
196 if update.update_type == UpdateType::NodeUpdated
198 && let Some(ref new_body) = update.new_body
199 {
200 let node_update = crate::mcp_tools::NodeUpdate {
201 body: Some(new_body.clone()),
202 signature: None,
203 docstring: None,
204 body_hash: None,
205 loc: None,
206 cyclomatic_complexity: None,
207 coverage_pct: None,
208 };
209 match crate::mcp_tools::codegraph_update_object(
210 &graph.nodes,
211 &update.node_id,
212 &node_update,
213 ) {
214 Ok(new_batch) => {
215 graph.nodes = new_batch;
216 graph
217 .body_hashes
218 .insert(update.node_id.clone(), update.body_hash.clone());
219 tracing::info!(node = %update.node_id, "applied remote update");
220 }
221 Err(e) => {
222 tracing::warn!(
223 node = %update.node_id,
224 error = %e,
225 "failed to apply remote update"
226 );
227 }
228 }
229 }
230 }
231}
232
233#[cfg(test)]
234mod tests {
235 use super::*;
236
237 #[test]
238 fn test_update_serialization_roundtrip() {
239 let update = CodeGraphUpdate {
240 update_type: UpdateType::NodeUpdated,
241 node_id: "rust_fn:crates/foo/src/lib.rs::bar".to_string(),
242 body_hash: "abc123".to_string(),
243 new_body: Some("fn bar() { 42 }".to_string()),
244 timestamp_ms: 1700000000000,
245 agent: "Mini".to_string(),
246 rationale: "fix typo".to_string(),
247 };
248 let json = serde_json::to_vec(&update).unwrap();
249 let back: CodeGraphUpdate = serde_json::from_slice(&json).unwrap();
250 assert_eq!(back.node_id, update.node_id);
251 assert_eq!(back.agent, "Mini");
252 assert_eq!(back.update_type, UpdateType::NodeUpdated);
253 }
254
255 #[test]
256 fn test_conflict_event_serialization() {
257 let conflict = ConflictEvent {
258 node_id: "rust_fn:foo::bar".to_string(),
259 local_hash: "aaa".to_string(),
260 remote_hash: "bbb".to_string(),
261 remote_agent: "DGX".to_string(),
262 resolution: "last_write_wins".to_string(),
263 timestamp_ms: 1700000000000,
264 };
265 let json = serde_json::to_vec(&conflict).unwrap();
266 let back: ConflictEvent = serde_json::from_slice(&json).unwrap();
267 assert_eq!(back.node_id, "rust_fn:foo::bar");
268 assert_eq!(back.resolution, "last_write_wins");
269 }
270
271 #[test]
272 fn test_make_update() {
273 let rt = tokio::runtime::Runtime::new().unwrap();
274 rt.block_on(async {
275 let update = CodeGraphUpdate {
277 update_type: UpdateType::NodeUpdated,
278 node_id: "test::node".to_string(),
279 body_hash: "hash123".to_string(),
280 new_body: Some("new body".to_string()),
281 timestamp_ms: chrono::Utc::now().timestamp_millis(),
282 agent: "TestAgent".to_string(),
283 rationale: "test".to_string(),
284 };
285 assert_eq!(update.agent, "TestAgent");
286 assert!(update.new_body.is_some());
287 });
288 }
289}