Skip to main content

graphwalker_mcp/
lib.rs

1mod dto;
2
3use graphwalker_service as service;
4use rmcp::{
5    handler::server::wrapper::Parameters,
6    model::{CallToolResult, Implementation, ServerCapabilities, ServerInfo},
7    tool, tool_handler, tool_router, ServerHandler,
8};
9use serde::Serialize;
10use serde_json::json;
11
12use dto::*;
13
14/// MCP protocol adapter for GraphWalker model authoring and execution.
15#[derive(Clone, Default)]
16pub struct GraphWalkerMcp {
17    drafts: service::DraftRegistry,
18    executions: service::ExecutionRegistry,
19}
20
21fn patch<T>(value: McpPatch<T>) -> service::FieldPatch<T> {
22    match value {
23        McpPatch::Keep => service::FieldPatch::Keep,
24        McpPatch::Set(value) => service::FieldPatch::Set(value),
25        McpPatch::Clear => service::FieldPatch::Clear,
26    }
27}
28
29fn structured<T: Serialize>(value: T) -> CallToolResult {
30    match serde_json::to_value(value) {
31        Ok(value) => CallToolResult::structured(value),
32        Err(error) => adapter_error(
33            "internal",
34            format!("Could not serialize tool result: {error}"),
35        ),
36    }
37}
38
39fn adapter_error(code: impl Into<String>, message: impl Into<String>) -> CallToolResult {
40    let error = ToolErrorOutput {
41        code: code.into(),
42        message: message.into(),
43    };
44    CallToolResult::structured_error(
45        serde_json::to_value(error).unwrap_or_else(
46            |_| json!({ "code": "internal", "message": "Internal MCP adapter error" }),
47        ),
48    )
49}
50
51fn service_error(error: service::ServiceError) -> CallToolResult {
52    let code = serde_json::to_value(error.code)
53        .ok()
54        .and_then(|value| value.as_str().map(str::to_owned))
55        .unwrap_or_else(|| "internal".to_string());
56    adapter_error(code, error.message)
57}
58
59fn result<T: Serialize>(value: Result<T, service::ServiceError>) -> CallToolResult {
60    match value {
61        Ok(value) => structured(value),
62        Err(error) => service_error(error),
63    }
64}
65
66fn vertex_output(vertex: graphwalker_io::json::JsonVertex) -> VertexOutput {
67    VertexOutput {
68        id: vertex.id,
69        name: vertex.name,
70        shared_state: vertex.shared_state,
71        actions: vertex.actions,
72        requirements: vertex.requirements,
73        properties: vertex.properties,
74    }
75}
76
77fn edge_output(edge: graphwalker_io::json::JsonEdge) -> Result<EdgeOutput, service::ServiceError> {
78    let target_vertex_id = edge.target_vertex_id.ok_or_else(|| {
79        service::ServiceError::new(
80            service::ServiceErrorCode::Internal,
81            "The service returned an edge without a target vertex",
82        )
83    })?;
84    Ok(EdgeOutput {
85        id: edge.id,
86        name: edge.name,
87        source_vertex_id: edge.source_vertex_id,
88        target_vertex_id,
89        guard: edge.guard,
90        actions: edge.actions,
91        requirements: edge.requirements,
92        properties: edge.properties,
93        weight: edge.weight,
94        dependency: edge.dependency,
95    })
96}
97
98fn statistics_output(value: service::ExecutionStatistics) -> StatisticsOutput {
99    StatisticsOutput {
100        total_vertices: value.total_vertices,
101        total_edges: value.total_edges,
102        visited_vertices: value.visited_vertices,
103        visited_edges: value.visited_edges,
104        unvisited_vertices: value.unvisited_vertices,
105        unvisited_edges: value.unvisited_edges,
106        vertex_coverage: value.vertex_coverage,
107        edge_coverage: value.edge_coverage,
108    }
109}
110
111#[tool_router]
112impl GraphWalkerMcp {
113    #[tool(
114        description = "Check that the local GraphWalker MCP server is running",
115        annotations(read_only_hint = true, open_world_hint = false)
116    )]
117    fn health(&self) -> rmcp::Json<HealthOutput> {
118        rmcp::Json(HealthOutput {
119            status: "ok",
120            server_version: env!("CARGO_PKG_VERSION"),
121        })
122    }
123
124    #[tool(
125        description = "Create an empty, process-local GraphWalker model draft",
126        output_schema = rmcp::handler::server::tool::schema_for_type::<CreateModelOutput>(),
127        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
128    )]
129    fn create_model(&self, Parameters(input): Parameters<CreateModelInput>) -> CallToolResult {
130        result(
131            self.drafts
132                .create_model(service::CreateModel {
133                    model_id: input.model_id,
134                    name: input.name,
135                    generator: input.generator,
136                    actions: input.actions,
137                    requirements: input.requirements,
138                    properties: input.properties,
139                })
140                .map(|created| CreateModelOutput {
141                    draft_id: created.draft_id.to_string(),
142                    model_id: created.model_id,
143                    revision: created.revision,
144                }),
145        )
146    }
147
148    #[tool(
149        description = "Add one vertex to a GraphWalker model draft and increment its revision",
150        output_schema = rmcp::handler::server::tool::schema_for_type::<VertexMutationOutput>(),
151        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
152    )]
153    fn add_vertex(&self, Parameters(input): Parameters<AddVertexInput>) -> CallToolResult {
154        result(
155            self.drafts
156                .add_vertex(service::AddVertex {
157                    draft_id: input.draft_id.as_str().into(),
158                    id: input.id,
159                    name: input.name,
160                    shared_state: input.shared_state,
161                    actions: input.actions,
162                    requirements: input.requirements,
163                    properties: input.properties,
164                    expected_revision: input.expected_revision,
165                })
166                .map(|added| VertexMutationOutput {
167                    vertex: vertex_output(added.vertex),
168                    revision: added.revision,
169                }),
170        )
171    }
172
173    #[tool(
174        description = "Add one edge to a GraphWalker model draft and increment its revision",
175        output_schema = rmcp::handler::server::tool::schema_for_type::<EdgeMutationOutput>(),
176        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
177    )]
178    fn add_edge(&self, Parameters(input): Parameters<AddEdgeInput>) -> CallToolResult {
179        let added = self
180            .drafts
181            .add_edge(service::AddEdge {
182                draft_id: input.draft_id.as_str().into(),
183                id: input.id,
184                name: input.name,
185                source_vertex_id: input.source_vertex_id,
186                target_vertex_id: Some(input.target_vertex_id),
187                guard: input.guard,
188                actions: input.actions,
189                requirements: input.requirements,
190                properties: input.properties,
191                weight: input.weight,
192                dependency: input.dependency,
193                expected_revision: input.expected_revision,
194            })
195            .and_then(|added| {
196                Ok(EdgeMutationOutput {
197                    edge: edge_output(added.edge)?,
198                    revision: added.revision,
199                })
200            });
201        result(added)
202    }
203
204    #[tool(
205        description = "Patch GraphWalker model metadata in a draft and increment its revision",
206        output_schema = rmcp::handler::server::tool::schema_for_type::<ModelMutationOutput>(),
207        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
208    )]
209    fn update_model(&self, Parameters(input): Parameters<UpdateModelInput>) -> CallToolResult {
210        result(
211            self.drafts
212                .update_model(service::UpdateModel {
213                    draft_id: input.draft_id.as_str().into(),
214                    name: patch(input.name),
215                    generator: patch(input.generator),
216                    start_element_id: patch(input.start_element_id),
217                    actions: patch(input.actions),
218                    requirements: patch(input.requirements),
219                    properties: patch(input.properties),
220                    predefined_path_edge_ids: patch(input.predefined_path_edge_ids),
221                    expected_revision: input.expected_revision,
222                })
223                .and_then(|updated| {
224                    serde_json::to_value(updated.model)
225                        .map(|model| ModelMutationOutput {
226                            model,
227                            revision: updated.revision,
228                        })
229                        .map_err(|error| {
230                            service::ServiceError::new(
231                                service::ServiceErrorCode::Internal,
232                                error.to_string(),
233                            )
234                        })
235                }),
236        )
237    }
238
239    #[tool(
240        description = "Patch one vertex in a GraphWalker model draft and increment its revision",
241        output_schema = rmcp::handler::server::tool::schema_for_type::<VertexMutationOutput>(),
242        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
243    )]
244    fn update_vertex(&self, Parameters(input): Parameters<UpdateVertexInput>) -> CallToolResult {
245        result(
246            self.drafts
247                .update_vertex(service::UpdateVertex {
248                    draft_id: input.draft_id.as_str().into(),
249                    vertex_id: input.vertex_id,
250                    name: patch(input.name),
251                    shared_state: patch(input.shared_state),
252                    actions: patch(input.actions),
253                    requirements: patch(input.requirements),
254                    properties: patch(input.properties),
255                    expected_revision: input.expected_revision,
256                })
257                .map(|updated| VertexMutationOutput {
258                    vertex: vertex_output(updated.vertex),
259                    revision: updated.revision,
260                }),
261        )
262    }
263
264    #[tool(
265        description = "Patch one edge in a GraphWalker model draft and increment its revision",
266        output_schema = rmcp::handler::server::tool::schema_for_type::<EdgeMutationOutput>(),
267        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
268    )]
269    fn update_edge(&self, Parameters(input): Parameters<UpdateEdgeInput>) -> CallToolResult {
270        let updated = self
271            .drafts
272            .update_edge(service::UpdateEdge {
273                draft_id: input.draft_id.as_str().into(),
274                edge_id: input.edge_id,
275                name: patch(input.name),
276                source_vertex_id: patch(input.source_vertex_id),
277                target_vertex_id: patch(input.target_vertex_id),
278                guard: patch(input.guard),
279                actions: patch(input.actions),
280                requirements: patch(input.requirements),
281                properties: patch(input.properties),
282                weight: patch(input.weight),
283                dependency: patch(input.dependency),
284                expected_revision: input.expected_revision,
285            })
286            .and_then(|updated| {
287                Ok(EdgeMutationOutput {
288                    edge: edge_output(updated.edge)?,
289                    revision: updated.revision,
290                })
291            });
292        result(updated)
293    }
294
295    #[tool(
296        description = "Remove an element from a GraphWalker draft; cascade and reference cleanup require explicit opt-in",
297        output_schema = rmcp::handler::server::tool::schema_for_type::<RemoveElementOutput>(),
298        annotations(read_only_hint = false, destructive_hint = true, open_world_hint = false)
299    )]
300    fn remove_element(&self, Parameters(input): Parameters<RemoveElementInput>) -> CallToolResult {
301        result(
302            self.drafts
303                .remove_element(service::RemoveElement {
304                    draft_id: input.draft_id.as_str().into(),
305                    element_id: input.element_id,
306                    cascade: input.cascade,
307                    cleanup_references: input.cleanup_references,
308                    expected_revision: input.expected_revision,
309                })
310                .map(|removed| RemoveElementOutput {
311                    removed_ids: removed.removed_ids,
312                    revision: removed.revision,
313                }),
314        )
315    }
316
317    #[tool(
318        description = "Export a GraphWalker model draft as a canonical JSON object without changing it",
319        output_schema = rmcp::handler::server::tool::schema_for_type::<ExportModelOutput>(),
320        annotations(read_only_hint = true, open_world_hint = false)
321    )]
322    fn export_model(&self, Parameters(input): Parameters<DraftInput>) -> CallToolResult {
323        result(
324            self.drafts
325                .export_model(&input.draft_id.as_str().into())
326                .map(|exported| ExportModelOutput {
327                    model: exported.model,
328                    revision: exported.revision,
329                }),
330        )
331    }
332
333    #[tool(
334        description = "Permanently discard a process-local GraphWalker model draft",
335        output_schema = rmcp::handler::server::tool::schema_for_type::<DiscardModelOutput>(),
336        annotations(read_only_hint = false, destructive_hint = true, idempotent_hint = false, open_world_hint = false)
337    )]
338    fn discard_model(&self, Parameters(input): Parameters<DraftInput>) -> CallToolResult {
339        result(
340            self.drafts
341                .discard(&input.draft_id.as_str().into())
342                .map(|discarded| DiscardModelOutput {
343                    discarded: discarded.discarded,
344                }),
345        )
346    }
347
348    #[tool(
349        description = "Validate exactly one inline GraphWalker JSON model or process-local draft without changing it",
350        output_schema = rmcp::handler::server::tool::schema_for_type::<ValidateModelOutput>(),
351        annotations(read_only_hint = true, open_world_hint = false)
352    )]
353    fn validate_model(&self, Parameters(input): Parameters<ValidateModelInput>) -> CallToolResult {
354        match (input.model, input.draft_id) {
355            (Some(model), None) => result(service::validate_model(&model).map(|validation| {
356                ValidateModelOutput {
357                    valid: validation.valid,
358                    issues: validation
359                        .issues
360                        .into_iter()
361                        .map(|issue| issue.message)
362                        .collect(),
363                    revision: None,
364                }
365            })),
366            (None, Some(draft_id)) => result(self.drafts.validate(&draft_id.as_str().into()).map(
367                |validation| {
368                    ValidateModelOutput {
369                        valid: validation.valid,
370                        issues: validation
371                            .issues
372                            .into_iter()
373                            .map(|issue| issue.message)
374                            .collect(),
375                        revision: Some(validation.revision),
376                    }
377                },
378            )),
379            _ => adapter_error(
380                "invalid_input",
381                "Provide exactly one of 'model' or 'draft_id'",
382            ),
383        }
384    }
385
386    #[tool(
387        description = "Convert an inline GraphML/yEd document into canonical GraphWalker JSON without storing it",
388        output_schema = rmcp::handler::server::tool::schema_for_type::<ConvertGraphmlOutput>(),
389        annotations(read_only_hint = true, open_world_hint = false)
390    )]
391    fn convert_graphml(
392        &self,
393        Parameters(input): Parameters<ConvertGraphmlInput>,
394    ) -> CallToolResult {
395        result(
396            service::convert_graphml(&input.graphml).map(|converted| ConvertGraphmlOutput {
397                model: converted.model,
398            }),
399        )
400    }
401
402    #[tool(
403        description = "Start isolated GraphWalker execution state from exactly one inline model or draft snapshot",
404        output_schema = rmcp::handler::server::tool::schema_for_type::<StartExecutionOutput>(),
405        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
406    )]
407    fn start_execution(
408        &self,
409        Parameters(input): Parameters<StartExecutionInput>,
410    ) -> CallToolResult {
411        let (model, source_revision) = match (input.model, input.draft_id) {
412            (Some(model), None) if input.revision.is_none() => {
413                let validation = match service::validate_model(&model) {
414                    Ok(validation) => validation,
415                    Err(error) => return service_error(error),
416                };
417                if !validation.valid {
418                    return adapter_error(
419                        "invalid_model",
420                        validation
421                            .issues
422                            .into_iter()
423                            .map(|issue| issue.message)
424                            .collect::<Vec<_>>()
425                            .join("; "),
426                    );
427                }
428                (model, None)
429            }
430            (None, Some(draft_id)) => {
431                let snapshot = match self
432                    .drafts
433                    .snapshot(&draft_id.as_str().into(), input.revision)
434                {
435                    Ok(snapshot) => snapshot,
436                    Err(error) => return service_error(error),
437                };
438                let validation = match self.drafts.validate(&draft_id.as_str().into()) {
439                    Ok(validation) => validation,
440                    Err(error) => return service_error(error),
441                };
442                if validation.revision != snapshot.revision {
443                    return service_error(service::ServiceError::new(
444                        service::ServiceErrorCode::RevisionConflict,
445                        format!(
446                            "Draft changed from revision {} to {} while starting execution",
447                            snapshot.revision, validation.revision
448                        ),
449                    ));
450                }
451                if !validation.valid {
452                    return adapter_error(
453                        "invalid_model",
454                        validation
455                            .issues
456                            .into_iter()
457                            .map(|issue| issue.message)
458                            .collect::<Vec<_>>()
459                            .join("; "),
460                    );
461                }
462                (snapshot.model, Some(snapshot.revision))
463            }
464            (Some(_), None) => {
465                return adapter_error("invalid_input", "'revision' is only valid with 'draft_id'")
466            }
467            _ => {
468                return adapter_error(
469                    "invalid_input",
470                    "Provide exactly one of 'model' or 'draft_id'",
471                )
472            }
473        };
474        result(
475            self.executions
476                .start(service::StartExecution {
477                    model,
478                    seed: input.seed,
479                    global_data: input.global_data,
480                })
481                .map(|started| StartExecutionOutput {
482                    execution_id: started.execution_id.to_string(),
483                    seed: started.seed,
484                    source_revision,
485                }),
486        )
487    }
488
489    #[tool(
490        description = "Advance a GraphWalker execution by at most one element",
491        output_schema = rmcp::handler::server::tool::schema_for_type::<NextStepOutput>(),
492        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
493    )]
494    fn next_step(&self, Parameters(input): Parameters<ExecutionInput>) -> CallToolResult {
495        result(
496            self.executions
497                .next_step(&input.execution_id.as_str().into())
498                .map(|step| NextStepOutput {
499                    completed: step.completed,
500                    element: step.element.map(|element| StepElementOutput {
501                        id: element.id,
502                        name: element.name,
503                        model_id: element.model_id,
504                        kind: match element.kind {
505                            service::ElementKind::Edge => "edge",
506                            service::ElementKind::Vertex => "vertex",
507                        }
508                        .to_string(),
509                        data: element.data,
510                        visited_count: element.visited_count,
511                        total_count: element.total_count,
512                        stop_condition_fulfillment: element.stop_condition_fulfillment,
513                    }),
514                }),
515        )
516    }
517
518    #[tool(
519        description = "Inspect GraphWalker execution data and coverage without advancing it",
520        output_schema = rmcp::handler::server::tool::schema_for_type::<ExecutionStatusOutput>(),
521        annotations(read_only_hint = true, open_world_hint = false)
522    )]
523    fn execution_status(
524        &self,
525        Parameters(input): Parameters<ExecutionStatusInput>,
526    ) -> CallToolResult {
527        let execution_id: service::ExecutionId = input.execution_id.as_str().into();
528        let status = match self.executions.status(&execution_id) {
529            Ok(status) => status,
530            Err(error) => return service_error(error),
531        };
532        let statistics = match self.executions.statistics(&execution_id) {
533            Ok(statistics) => statistics,
534            Err(error) => return service_error(error),
535        };
536        let elements = if input.include_elements {
537            match self.executions.elements(&execution_id) {
538                Ok(elements) => Some(
539                    elements
540                        .into_iter()
541                        .map(|element| ElementStatusOutput {
542                            model_id: element.model_id,
543                            element_id: element.element_id,
544                            visited_count: element.visited_count,
545                        })
546                        .collect(),
547                ),
548                Err(error) => return service_error(error),
549            }
550        } else {
551            None
552        };
553        structured(ExecutionStatusOutput {
554            has_next: status.has_next,
555            data: status.data,
556            statistics: statistics_output(statistics),
557            elements,
558        })
559    }
560
561    #[tool(
562        description = "Execute a data script in the current GraphWalker execution context",
563        output_schema = rmcp::handler::server::tool::schema_for_type::<SetExecutionDataOutput>(),
564        annotations(read_only_hint = false, destructive_hint = false, open_world_hint = false)
565    )]
566    fn set_execution_data(
567        &self,
568        Parameters(input): Parameters<SetExecutionDataInput>,
569    ) -> CallToolResult {
570        result(
571            self.executions
572                .set_data(&input.execution_id.as_str().into(), input.script)
573                .map(|updated| SetExecutionDataOutput { data: updated.data }),
574        )
575    }
576
577    #[tool(
578        description = "Restart a GraphWalker execution using its original model, seed, and global data",
579        output_schema = rmcp::handler::server::tool::schema_for_type::<RestartExecutionOutput>(),
580        annotations(read_only_hint = false, destructive_hint = true, idempotent_hint = true, open_world_hint = false)
581    )]
582    fn restart_execution(&self, Parameters(input): Parameters<ExecutionInput>) -> CallToolResult {
583        result(
584            self.executions
585                .restart(&input.execution_id.as_str().into())
586                .map(|restarted| RestartExecutionOutput {
587                    restarted: true,
588                    seed: restarted.seed,
589                }),
590        )
591    }
592
593    #[tool(
594        description = "Close and permanently release one process-local GraphWalker execution",
595        output_schema = rmcp::handler::server::tool::schema_for_type::<CloseExecutionOutput>(),
596        annotations(read_only_hint = false, destructive_hint = true, idempotent_hint = false, open_world_hint = false)
597    )]
598    fn close_execution(&self, Parameters(input): Parameters<ExecutionInput>) -> CallToolResult {
599        result(
600            self.executions
601                .close(&input.execution_id.as_str().into())
602                .map(|()| CloseExecutionOutput { closed: true }),
603        )
604    }
605}
606
607#[tool_handler]
608impl ServerHandler for GraphWalkerMcp {
609    fn get_info(&self) -> ServerInfo {
610        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
611            .with_server_info(Implementation::new(
612                env!("CARGO_PKG_NAME"),
613                env!("CARGO_PKG_VERSION"),
614            ))
615            .with_instructions(
616                "Build, revise, export, validate, and execute GraphWalker model-based tests. Draft and execution IDs are process-local; export models before disconnecting if they must persist.",
617            )
618    }
619}