1use std::pin::Pin;
2use std::sync::{Arc, RwLock};
3
4use tokio::sync::mpsc;
5use tokio_stream::{wrappers::ReceiverStream, Stream};
6use tonic::{Request, Response, Status};
7
8use crate::event_bus::{Event, EventBus};
9use crate::proto::graph_service_server::GraphService;
10use crate::proto::{
11 Edge, GetGraphRequest, GetGraphResponse, GetNodeRequest, GetNodeResponse, GraphEvent, Node,
12 SearchNodesRequest, SearchNodesResponse, ShortestPathRequest, ShortestPathResponse,
13 WatchGraphRequest,
14};
15use crate::state::{GraphEdge, GraphNode, GraphState};
16
17pub struct GraphServiceImpl {
18 state: Arc<RwLock<GraphState>>,
19 bus: Arc<EventBus>,
20}
21
22impl GraphServiceImpl {
23 pub fn new(state: Arc<RwLock<GraphState>>, bus: Arc<EventBus>) -> Self {
24 Self { state, bus }
25 }
26
27 pub fn add_node(&self, node: GraphNode) {
28 let ev = Event::NodeAdded {
29 id: node.id.clone(),
30 label: node.label.clone(),
31 source_file: node.source_file.clone(),
32 };
33 self.state.write().unwrap().add_node(node);
34 self.bus.emit(ev);
35 }
36
37 pub fn add_edge(&self, edge: GraphEdge) {
38 let ev = Event::EdgeAdded {
39 source: edge.source.clone(),
40 target: edge.target.clone(),
41 relation: edge.relation.clone(),
42 };
43 self.state.write().unwrap().add_edge(edge);
44 self.bus.emit(ev);
45 }
46}
47
48fn node_to_proto(n: &GraphNode) -> Node {
49 Node {
50 id: n.id.clone(),
51 label: n.label.clone(),
52 source_file: n.source_file.clone(),
53 source_location: n.source_location.clone(),
54 community: n.community,
55 }
56}
57
58fn edge_to_proto(e: &GraphEdge) -> Edge {
59 Edge {
60 source: e.source.clone(),
61 target: e.target.clone(),
62 relation: e.relation.clone(),
63 confidence: e.confidence.clone(),
64 }
65}
66
67fn event_to_proto(ev: Event) -> Option<GraphEvent> {
68 use crate::proto::graph_event::Event as PEvent;
69 let inner = match ev {
70 Event::NodeAdded {
71 id,
72 label,
73 source_file,
74 } => PEvent::NodeAdded(Node {
75 id,
76 label,
77 source_file,
78 source_location: String::new(),
79 community: 0,
80 }),
81 Event::EdgeAdded {
82 source,
83 target,
84 relation,
85 } => PEvent::EdgeAdded(Edge {
86 source,
87 target,
88 relation,
89 confidence: String::new(),
90 }),
91 Event::NodeRemoved { id } => PEvent::NodeRemoved(id),
92 Event::GraphReset => PEvent::GraphReset("reset".to_string()),
93 };
94 Some(GraphEvent { event: Some(inner) })
95}
96
97pub(crate) type BoxStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
98
99#[tonic::async_trait]
100impl GraphService for GraphServiceImpl {
101 async fn get_graph(
102 &self,
103 _req: Request<GetGraphRequest>,
104 ) -> Result<Response<GetGraphResponse>, Status> {
105 let state = self.state.read().unwrap();
106 let nodes = state.get_nodes().iter().map(node_to_proto).collect();
107 let edges = state.get_edges().iter().map(edge_to_proto).collect();
108 Ok(Response::new(GetGraphResponse { nodes, edges }))
109 }
110
111 async fn get_node(
112 &self,
113 req: Request<GetNodeRequest>,
114 ) -> Result<Response<GetNodeResponse>, Status> {
115 let id = req.into_inner().id;
116 let state = self.state.read().unwrap();
117 match state.get_node(&id) {
118 Some(n) => Ok(Response::new(GetNodeResponse {
119 node: Some(node_to_proto(n)),
120 })),
121 None => Err(Status::not_found(format!("node {} not found", id))),
122 }
123 }
124
125 async fn search_nodes(
126 &self,
127 req: Request<SearchNodesRequest>,
128 ) -> Result<Response<SearchNodesResponse>, Status> {
129 let inner = req.into_inner();
130 let limit = if inner.limit <= 0 {
131 20
132 } else {
133 inner.limit as usize
134 };
135 let state = self.state.read().unwrap();
136 let nodes = state
137 .search_nodes(&inner.query, limit)
138 .iter()
139 .map(node_to_proto)
140 .collect();
141 Ok(Response::new(SearchNodesResponse { nodes }))
142 }
143
144 async fn shortest_path(
145 &self,
146 req: Request<ShortestPathRequest>,
147 ) -> Result<Response<ShortestPathResponse>, Status> {
148 let inner = req.into_inner();
149 let state = self.state.read().unwrap();
150 match state.shortest_path(&inner.source, &inner.target) {
151 Some(path) => Ok(Response::new(ShortestPathResponse {
152 node_ids: path,
153 found: true,
154 })),
155 None => Ok(Response::new(ShortestPathResponse {
156 node_ids: vec![],
157 found: false,
158 })),
159 }
160 }
161
162 type WatchGraphStream = BoxStream<GraphEvent>;
163
164 async fn watch_graph(
165 &self,
166 _req: Request<WatchGraphRequest>,
167 ) -> Result<Response<Self::WatchGraphStream>, Status> {
168 let (tx, rx) = mpsc::channel(128);
169 let mut sub = self.bus.subscribe();
170 tokio::spawn(async move {
171 while let Ok(ev) = sub.recv().await {
172 if let Some(proto_ev) = event_to_proto(ev) {
173 if tx.send(Ok(proto_ev)).await.is_err() {
174 break;
175 }
176 }
177 }
178 });
179 Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
180 }
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186 use crate::proto::{GetGraphRequest, GetNodeRequest, SearchNodesRequest, ShortestPathRequest};
187 use tokio_stream::StreamExt;
188
189 fn make_service() -> GraphServiceImpl {
190 let state = Arc::new(RwLock::new(GraphState::new()));
191 let bus = Arc::new(EventBus::new());
192 GraphServiceImpl::new(state, bus)
193 }
194
195 fn node(id: &str, label: &str) -> GraphNode {
196 GraphNode {
197 id: id.to_string(),
198 label: label.to_string(),
199 source_file: "test.rs".to_string(),
200 source_location: "1:1".to_string(),
201 community: 0,
202 }
203 }
204
205 fn edge(src: &str, tgt: &str) -> GraphEdge {
206 GraphEdge {
207 source: src.to_string(),
208 target: tgt.to_string(),
209 relation: "calls".to_string(),
210 confidence: "1.0".to_string(),
211 }
212 }
213
214 #[tokio::test]
215 async fn test_get_graph_empty() {
216 let svc = make_service();
217 let resp = svc
218 .get_graph(Request::new(GetGraphRequest {}))
219 .await
220 .unwrap();
221 let inner = resp.into_inner();
222 assert!(inner.nodes.is_empty());
223 assert!(inner.edges.is_empty());
224 }
225
226 #[tokio::test]
227 async fn test_get_graph_with_nodes_and_edges() {
228 let svc = make_service();
229 svc.add_node(node("a", "Alpha"));
230 svc.add_node(node("b", "Beta"));
231 svc.add_edge(edge("a", "b"));
232
233 let resp = svc
234 .get_graph(Request::new(GetGraphRequest {}))
235 .await
236 .unwrap();
237 let inner = resp.into_inner();
238 assert_eq!(inner.nodes.len(), 2);
239 assert_eq!(inner.edges.len(), 1);
240 }
241
242 #[tokio::test]
243 async fn test_get_node_found() {
244 let svc = make_service();
245 svc.add_node(node("n1", "Foo"));
246
247 let resp = svc
248 .get_node(Request::new(GetNodeRequest {
249 id: "n1".to_string(),
250 }))
251 .await
252 .unwrap();
253 assert_eq!(resp.into_inner().node.unwrap().label, "Foo");
254 }
255
256 #[tokio::test]
257 async fn test_get_node_not_found_returns_status_not_found() {
258 let svc = make_service();
259 let err = svc
260 .get_node(Request::new(GetNodeRequest {
261 id: "missing".to_string(),
262 }))
263 .await
264 .unwrap_err();
265 assert_eq!(err.code(), tonic::Code::NotFound);
266 }
267
268 #[tokio::test]
269 async fn test_search_nodes_matches_substring() {
270 let svc = make_service();
271 svc.add_node(node("1", "AuthService"));
272 svc.add_node(node("2", "AuthController"));
273 svc.add_node(node("3", "UserService"));
274
275 let resp = svc
276 .search_nodes(Request::new(SearchNodesRequest {
277 query: "auth".to_string(),
278 limit: 10,
279 }))
280 .await
281 .unwrap();
282 assert_eq!(resp.into_inner().nodes.len(), 2);
283 }
284
285 #[tokio::test]
286 async fn test_search_nodes_limit_respected() {
287 let svc = make_service();
288 for i in 0..5 {
289 svc.add_node(node(&i.to_string(), &format!("Service{}", i)));
290 }
291 let resp = svc
292 .search_nodes(Request::new(SearchNodesRequest {
293 query: "service".to_string(),
294 limit: 2,
295 }))
296 .await
297 .unwrap();
298 assert_eq!(resp.into_inner().nodes.len(), 2);
299 }
300
301 #[tokio::test]
302 async fn test_search_nodes_zero_limit_defaults_to_20() {
303 let svc = make_service();
304 for i in 0..25 {
305 svc.add_node(node(&i.to_string(), &format!("Node{}", i)));
306 }
307 let resp = svc
308 .search_nodes(Request::new(SearchNodesRequest {
309 query: "node".to_string(),
310 limit: 0,
311 }))
312 .await
313 .unwrap();
314 assert_eq!(resp.into_inner().nodes.len(), 20);
315 }
316
317 #[tokio::test]
318 async fn test_shortest_path_found() {
319 let svc = make_service();
320 svc.add_node(node("a", "A"));
321 svc.add_node(node("b", "B"));
322 svc.add_node(node("c", "C"));
323 svc.add_edge(edge("a", "b"));
324 svc.add_edge(edge("b", "c"));
325
326 let resp = svc
327 .shortest_path(Request::new(ShortestPathRequest {
328 source: "a".to_string(),
329 target: "c".to_string(),
330 }))
331 .await
332 .unwrap();
333 let inner = resp.into_inner();
334 assert!(inner.found);
335 assert_eq!(inner.node_ids, vec!["a", "b", "c"]);
336 }
337
338 #[tokio::test]
339 async fn test_shortest_path_not_found() {
340 let svc = make_service();
341 svc.add_node(node("x", "X"));
342 svc.add_node(node("y", "Y"));
343
344 let resp = svc
345 .shortest_path(Request::new(ShortestPathRequest {
346 source: "x".to_string(),
347 target: "y".to_string(),
348 }))
349 .await
350 .unwrap();
351 let inner = resp.into_inner();
352 assert!(!inner.found);
353 assert!(inner.node_ids.is_empty());
354 }
355
356 #[tokio::test]
357 async fn test_watch_graph_receives_node_added_event() {
358 let svc = make_service();
359 let resp = svc
360 .watch_graph(Request::new(crate::proto::WatchGraphRequest {}))
361 .await
362 .unwrap();
363 let mut stream = resp.into_inner();
364
365 svc.add_node(node("w1", "Watcher"));
366
367 let event = tokio::time::timeout(std::time::Duration::from_millis(200), stream.next())
368 .await
369 .expect("timed out waiting for event")
370 .expect("stream ended")
371 .unwrap();
372
373 assert!(matches!(
374 event.event,
375 Some(crate::proto::graph_event::Event::NodeAdded(_))
376 ));
377 }
378}