ctx-mcp 0.1.0

package/src/mcp/tools.ts

@@ -0,0 +1,522 @@
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp";
+
+import {
+  traceAddEdgeInputSchema,
+  traceAddEdgeOutputSchema,
+  traceAddNodeInputSchema,
+  traceAddNodeOutputSchema,
+  traceAttachArtifactInputSchema,
+  traceAttachArtifactOutputSchema,
+  traceExplainDecisionInputSchema,
+  traceExplainDecisionOutputSchema,
+  traceRiskCheckInputSchema,
+  traceRiskCheckOutputSchema,
+  traceSimilarityInputSchema,
+  traceSimilarityOutputSchema,
+  traceFindPathsInputSchema,
+  traceFindPathsOutputSchema,
+  traceFinishInputSchema,
+  traceFinishOutputSchema,
+  traceGetSubgraphInputSchema,
+  traceGetSubgraphOutputSchema,
+  traceQueryInputSchema,
+  traceQueryOutputSchema,
+  traceStartInputSchema,
+  traceStartOutputSchema,
+} from "../schemas.js";
+import { newId } from "../util/ids.js";
+import { sha256 } from "../util/hashing.js";
+import { redactSecrets } from "../util/redact.js";
+import {
+  finishTrace,
+  getArtifactsForNodes,
+  getTrace,
+  getTraceEdges,
+  getTraceNodes,
+  insertArtifact,
+  insertEdge,
+  insertNode,
+  insertTrace,
+  searchNodes,
+  searchNodesForTraces,
+  type ArtifactRow,
+  type NodeRow,
+  type TraceRow,
+} from "../storage/queries.js";
+import type { StorageEngine } from "../storage/engine.js";
+import { buildGraph, extractSubgraph, findPaths } from "../graph/graph.js";
+import { explainDecision } from "../explain/explain.js";
+import {
+  assessDecisionRisk,
+  findSimilarDecisions,
+  invalidateTraceContext,
+} from "../similarity/similarity.js";
+
+const ok = <T>(payload: T) => payload;
+
+const error = (code: string, message: string, details?: unknown) => ({
+  ok: false,
+  error: { code, message, details },
+});
+
+function toNodeOutput(node: NodeRow) {
+  return {
+    ...node,
+    data: JSON.parse(node.data_json),
+    metadata: JSON.parse(node.metadata_json),
+  };
+}
+
+function toTraceOutput(trace: TraceRow) {
+  return {
+    ...trace,
+    tags: JSON.parse(trace.tags_json),
+    metadata: JSON.parse(trace.metadata_json),
+    outcome: trace.outcome_json ? JSON.parse(trace.outcome_json) : null,
+  };
+}
+
+function toArtifactOutput(artifact: ArtifactRow) {
+  return {
+    ...artifact,
+    metadata: JSON.parse(artifact.metadata_json),
+  };
+}
+
+function ensureTrace(engine: StorageEngine, traceId: string): TraceRow {
+  const trace = getTrace(engine, traceId);
+  if (!trace) {
+    throw new Error(`Trace ${traceId} not found`);
+  }
+  return trace;
+}
+
+export function registerTools(server: McpServer, engine: StorageEngine): void {
+  server.tool(
+    "trace.start",
+    traceStartInputSchema,
+    async (input) => {
+      const payload = traceStartInputSchema.parse(input);
+      const traceId = payload.trace_id ?? newId();
+
+      const row: TraceRow = {
+        trace_id: traceId,
+        workflow_id: payload.workflow_id ?? null,
+        actor: payload.actor ?? null,
+        intent: payload.intent,
+        tags_json: JSON.stringify(payload.tags ?? []),
+        metadata_json: JSON.stringify(payload.metadata ?? {}),
+        started_at: payload.started_at ?? Date.now(),
+        finished_at: null,
+        status: null,
+        outcome_json: null,
+      };
+
+      try {
+        engine.transaction(() => {
+          insertTrace(engine, row);
+        });
+      } catch (err) {
+        return error("trace_start_failed", "Failed to start trace", (err as Error).message);
+      }
+
+      return ok(traceStartOutputSchema.parse({ ok: true, trace_id: traceId }));
+    }
+  );
+
+  server.tool(
+    "trace.add_node",
+    traceAddNodeInputSchema,
+    async (input) => {
+      const payload = traceAddNodeInputSchema.parse(input);
+      const nodeId = payload.node_id ?? newId();
+
+      try {
+        engine.transaction(() => {
+          ensureTrace(engine, payload.trace_id);
+          insertNode(engine, {
+            node_id: nodeId,
+            trace_id: payload.trace_id,
+            type: payload.type,
+            summary: payload.summary,
+            data_json: JSON.stringify(payload.data ?? {}),
+            confidence: payload.confidence ?? null,
+            metadata_json: JSON.stringify(payload.metadata ?? {}),
+            created_at: payload.created_at ?? Date.now(),
+          });
+        });
+        invalidateTraceContext(engine, payload.trace_id);
+      } catch (err) {
+        return error("trace_add_node_failed", "Failed to add node", (err as Error).message);
+      }
+
+      return ok(traceAddNodeOutputSchema.parse({ ok: true, node_id: nodeId }));
+    }
+  );
+
+  server.tool(
+    "trace.add_edge",
+    traceAddEdgeInputSchema,
+    async (input) => {
+      const payload = traceAddEdgeInputSchema.parse(input);
+      const edgeId = payload.edge_id ?? newId();
+
+      try {
+        engine.transaction(() => {
+          ensureTrace(engine, payload.trace_id);
+          insertEdge(engine, {
+            edge_id: edgeId,
+            trace_id: payload.trace_id,
+            from_node_id: payload.from_node_id,
+            to_node_id: payload.to_node_id,
+            relation_type: payload.relation_type,
+            data_json: JSON.stringify(payload.data ?? {}),
+            created_at: payload.created_at ?? Date.now(),
+          });
+        });
+        invalidateTraceContext(engine, payload.trace_id);
+      } catch (err) {
+        return error("trace_add_edge_failed", "Failed to add edge", (err as Error).message);
+      }
+
+      return ok(traceAddEdgeOutputSchema.parse({ ok: true, edge_id: edgeId }));
+    }
+  );
+
+  server.tool(
+    "trace.attach_artifact",
+    traceAttachArtifactInputSchema,
+    async (input) => {
+      const payload = traceAttachArtifactInputSchema.parse(input);
+      const artifactId = payload.artifact_id ?? newId();
+      const redactResult = redactSecrets(payload.content);
+      const redactionLevel = redactResult.wasRedacted
+        ? "scrubbed"
+        : payload.redaction_level ?? "internal";
+
+      try {
+        engine.transaction(() => {
+          ensureTrace(engine, payload.trace_id);
+          insertArtifact(engine, {
+            artifact_id: artifactId,
+            trace_id: payload.trace_id,
+            node_id: payload.node_id ?? null,
+            artifact_type: payload.artifact_type,
+            content: redactResult.redacted,
+            redaction_level: redactionLevel,
+            metadata_json: JSON.stringify(payload.metadata ?? {}),
+            sha256: sha256(redactResult.redacted),
+            created_at: payload.created_at ?? Date.now(),
+          });
+        });
+      } catch (err) {
+        return error(
+          "trace_attach_artifact_failed",
+          "Failed to attach artifact",
+          (err as Error).message
+        );
+      }
+
+      return ok(
+        traceAttachArtifactOutputSchema.parse({
+          ok: true,
+          artifact_id: artifactId,
+          sha256: sha256(redactResult.redacted),
+          redaction_level: redactionLevel,
+        })
+      );
+    }
+  );
+
+  server.tool(
+    "trace.finish",
+    traceFinishInputSchema,
+    async (input) => {
+      const payload = traceFinishInputSchema.parse(input);
+      const finishedAt = payload.finished_at ?? Date.now();
+
+      try {
+        engine.transaction(() => {
+          ensureTrace(engine, payload.trace_id);
+          finishTrace(
+            engine,
+            payload.trace_id,
+            finishedAt,
+            payload.status ?? null,
+            payload.outcome ? JSON.stringify(payload.outcome) : null
+          );
+        });
+      } catch (err) {
+        return error("trace_finish_failed", "Failed to finish trace", (err as Error).message);
+      }
+
+      return ok(traceFinishOutputSchema.parse({ ok: true, trace_id: payload.trace_id }));
+    }
+  );
+
+  server.tool(
+    "trace.query",
+    traceQueryInputSchema,
+    async (input) => {
+      const payload = traceQueryInputSchema.parse(input);
+      const limit = payload.limit ?? 50;
+      const offset = payload.offset ?? 0;
+      const hasTags = Boolean(payload.tags && payload.tags.length > 0);
+
+      let nodes: NodeRow[] = [];
+      try {
+        if (hasTags) {
+          if (payload.trace_id) {
+            const trace = getTrace(engine, payload.trace_id);
+            const traceTags = trace ? toTraceOutput(trace).tags : [];
+            if (!payload.tags?.every((tag) => traceTags.includes(tag))) {
+              nodes = [];
+            } else {
+              nodes = searchNodes(engine, {
+                traceId: payload.trace_id,
+                type: payload.type,
+                text: payload.text,
+                limit,
+                offset,
+              });
+            }
+          } else {
+            const traceRows = engine.queryAll<TraceRow>("SELECT * FROM traces;");
+            const matchingTraceIds = traceRows
+              .map((row) => toTraceOutput(row))
+              .filter((row) => payload.tags?.every((tag) => row.tags.includes(tag)))
+              .map((row) => row.trace_id);
+
+            nodes = searchNodesForTraces(engine, {
+              traceIds: matchingTraceIds,
+              type: payload.type,
+              text: payload.text,
+              limit,
+              offset,
+            });
+          }
+        } else {
+          nodes = searchNodes(engine, {
+            traceId: payload.trace_id,
+            type: payload.type,
+            text: payload.text,
+            limit,
+            offset,
+          });
+        }
+      } catch (err) {
+        return error("trace_query_failed", "Failed to query trace", (err as Error).message);
+      }
+
+      return ok(traceQueryOutputSchema.parse({ ok: true, nodes: nodes.map(toNodeOutput) }));
+    }
+  );
+
+  server.tool(
+    "trace.get_subgraph",
+    traceGetSubgraphInputSchema,
+    async (input) => {
+      const payload = traceGetSubgraphInputSchema.parse(input);
+
+      try {
+        ensureTrace(engine, payload.trace_id);
+        const nodes = getTraceNodes(engine, payload.trace_id);
+        const edges = getTraceEdges(engine, payload.trace_id);
+        const graph = buildGraph(nodes, edges);
+        const subgraph = extractSubgraph(graph, {
+          center: payload.center,
+          depth: payload.depth ?? 2,
+          direction: payload.direction ?? "both",
+          maxNodes: payload.max_nodes ?? 250,
+        });
+
+        return ok(
+          traceGetSubgraphOutputSchema.parse({
+            ok: true,
+            nodes: subgraph.nodes.map(toNodeOutput),
+            edges: subgraph.edges,
+          })
+        );
+      } catch (err) {
+        return error("trace_get_subgraph_failed", "Failed to get subgraph", (err as Error).message);
+      }
+    }
+  );
+
+  server.tool(
+    "trace.find_paths",
+    traceFindPathsInputSchema,
+    async (input) => {
+      const payload = traceFindPathsInputSchema.parse(input);
+
+      try {
+        ensureTrace(engine, payload.trace_id);
+        const nodes = getTraceNodes(engine, payload.trace_id);
+        const edges = getTraceEdges(engine, payload.trace_id);
+        const graph = buildGraph(nodes, edges);
+        const paths = findPaths(graph, {
+          from: payload.from,
+          to: payload.to,
+          maxDepth: payload.max_depth ?? 4,
+          maxPaths: payload.max_paths ?? 10,
+          allowRelations: payload.allow_relations
+            ? new Set(payload.allow_relations)
+            : undefined,
+        });
+
+        return ok(traceFindPathsOutputSchema.parse({ ok: true, paths }));
+      } catch (err) {
+        return error("trace_find_paths_failed", "Failed to find paths", (err as Error).message);
+      }
+    }
+  );
+
+  server.tool(
+    "trace.explain_decision",
+    traceExplainDecisionInputSchema,
+    async (input) => {
+      const payload = traceExplainDecisionInputSchema.parse(input);
+
+      try {
+        ensureTrace(engine, payload.trace_id);
+        const nodes = getTraceNodes(engine, payload.trace_id);
+        const edges = getTraceEdges(engine, payload.trace_id);
+        const artifacts = getArtifactsForNodes(
+          engine,
+          nodes.map((node) => node.node_id)
+        );
+        const explanation = explainDecision(
+          payload.decision_node_id,
+          nodes,
+          edges,
+          artifacts,
+          payload.depth ?? 4
+        );
+
+        return ok(
+          traceExplainDecisionOutputSchema.parse({
+            ok: true,
+            explanation: {
+              ...explanation,
+              decision_node: toNodeOutput(explanation.decision_node),
+              assumptions: explanation.assumptions.map(toNodeOutput),
+              evidence: explanation.evidence.map(toNodeOutput),
+              alternatives: explanation.alternatives.map(toNodeOutput),
+              checks: explanation.checks.map(toNodeOutput),
+              outcome: explanation.outcome ? toNodeOutput(explanation.outcome) : null,
+              artifacts: explanation.artifacts.map(toArtifactOutput),
+            },
+          })
+        );
+      } catch (err) {
+        return error(
+          "trace_explain_decision_failed",
+          "Failed to explain decision",
+          (err as Error).message
+        );
+      }
+    }
+  );
+
+  server.tool(
+    "trace.similarity",
+    traceSimilarityInputSchema,
+    async (input) => {
+      const payload = traceSimilarityInputSchema.parse(input);
+
+      try {
+        const matches = findSimilarDecisions(engine, payload.decision_node_id, {
+          scope: payload.scope ?? "all",
+          limit: payload.limit ?? 10,
+          depth: payload.depth ?? 4,
+          maxTraces: payload.max_traces ?? 200,
+        });
+
+        return ok(
+          traceSimilarityOutputSchema.parse({
+            ok: true,
+            matches: matches.map((match) => ({
+              node: toNodeOutput(match.node),
+              trace: toTraceOutput(match.trace),
+              score: match.score,
+              reasons: match.reasons,
+            })),
+          })
+        );
+      } catch (err) {
+        return error(
+          "trace_similarity_failed",
+          "Failed to compute similarity",
+          (err as Error).message
+        );
+      }
+    }
+  );
+
+  server.tool(
+    "trace.risk_check",
+    traceRiskCheckInputSchema,
+    async (input) => {
+      const payload = traceRiskCheckInputSchema.parse(input);
+
+      try {
+        const assessment = assessDecisionRisk(engine, payload.decision_node_id, {
+          scope: payload.scope ?? "all",
+          limit: payload.limit ?? 10,
+          depth: payload.depth ?? 4,
+          threshold: payload.threshold ?? 0.6,
+          maxTraces: payload.max_traces ?? 200,
+        });
+
+        return ok(
+          traceRiskCheckOutputSchema.parse({
+            ok: true,
+            assessment: {
+              decision: toNodeOutput(assessment.decision),
+              novelty: assessment.novelty,
+              similar_failures: assessment.similar_failures.map((match) => ({
+                node: toNodeOutput(match.node),
+                trace: toTraceOutput(match.trace),
+                score: match.score,
+                reasons: match.reasons,
+              })),
+              signature_failed_match: assessment.signature_failed_match,
+              warnings: assessment.warnings,
+            },
+          })
+        );
+      } catch (err) {
+        return error(
+          "trace_risk_check_failed",
+          "Failed to assess decision risk",
+          (err as Error).message
+        );
+      }
+    }
+  );
+}
+
+export const ToolSchemas = {
+  traceStartInputSchema,
+  traceStartOutputSchema,
+  traceAddNodeInputSchema,
+  traceAddNodeOutputSchema,
+  traceAddEdgeInputSchema,
+  traceAddEdgeOutputSchema,
+  traceAttachArtifactInputSchema,
+  traceAttachArtifactOutputSchema,
+  traceFinishInputSchema,
+  traceFinishOutputSchema,
+  traceQueryInputSchema,
+  traceQueryOutputSchema,
+  traceGetSubgraphInputSchema,
+  traceGetSubgraphOutputSchema,
+  traceFindPathsInputSchema,
+  traceFindPathsOutputSchema,
+  traceExplainDecisionInputSchema,
+  traceExplainDecisionOutputSchema,
+  traceSimilarityInputSchema,
+  traceSimilarityOutputSchema,
+  traceRiskCheckInputSchema,
+  traceRiskCheckOutputSchema,
+};