Skip to main content

nusy_codegraph/
nats_sync.rs

1//! NATS-based code graph synchronization (EX-3184).
2//!
3//! When an agent modifies a CodeNode via `code_replace`, the change is published
4//! to `nusy.code.graph.updates`. Other agents subscribed to that subject receive
5//! the update within milliseconds and apply it to their local graph.
6//!
7//! This is multi-writer (unlike kanban's single-writer pattern). Conflicts are
8//! detected and logged to `nusy.code.graph.conflict` for the Captain to review.
9
10use arrow::array::RecordBatch;
11use chrono::Utc;
12use futures::StreamExt;
13use serde::{Deserialize, Serialize};
14use std::sync::Arc;
15use tokio::sync::Mutex;
16
17/// Subject for code graph update events.
18pub const UPDATES_SUBJECT: &str = "nusy.code.graph.updates";
19/// Subject for conflict notifications.
20pub const CONFLICT_SUBJECT: &str = "nusy.code.graph.conflict";
21
22/// A code graph mutation event, published over NATS.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct CodeGraphUpdate {
25    /// Type of mutation.
26    pub update_type: UpdateType,
27    /// Node ID that was modified (e.g., `rust_fn:crates/foo/src/lib.rs::bar`).
28    pub node_id: String,
29    /// SHA-256 hash of the new body (first 16 hex chars).
30    pub body_hash: String,
31    /// New body content. None for deletes.
32    pub new_body: Option<String>,
33    /// Unix timestamp (milliseconds).
34    pub timestamp_ms: i64,
35    /// Agent that made the change.
36    pub agent: String,
37    /// Why this change was made.
38    pub rationale: String,
39}
40
41/// Type of graph mutation.
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
43pub enum UpdateType {
44    NodeUpdated,
45    NodeCreated,
46    NodeDeleted,
47    EdgeCreated,
48    EdgeDeleted,
49}
50
51/// A conflict detected during sync (two agents edited the same node).
52#[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
62// ─── Publisher ──────────────────────────────────────────────────────────────
63
64/// Publishes code graph updates to NATS.
65pub struct CodeGraphPublisher {
66    client: async_nats::Client,
67    agent_name: String,
68}
69
70impl CodeGraphPublisher {
71    /// Connect to NATS and create a publisher.
72    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    /// Get a reference to the NATS client (for fire-and-forget from sync handlers).
84    pub fn client_ref(&self) -> &async_nats::Client {
85        &self.client
86    }
87
88    /// Publish a code graph update.
89    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    /// Create an update event for a node replacement.
95    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
114// ─── Subscriber ─────────────────────────────────────────────────────────────
115
116/// Shared graph state that the subscriber can update.
117pub type SharedGraphState = Arc<Mutex<SyncableGraph>>;
118
119/// The subset of graph state needed for sync operations.
120pub struct SyncableGraph {
121    pub nodes: RecordBatch,
122    pub edges: RecordBatch,
123    /// Track body hashes for conflict detection.
124    pub body_hashes: std::collections::HashMap<String, String>,
125}
126
127/// Subscribe to code graph updates and apply them locally.
128pub 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        // Skip our own updates to avoid echo.
158        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        // Conflict detection: check if we have a different version.
173        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        // Apply the update (last write wins).
197        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            // Can't actually connect without NATS, but test the struct
276            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}