saving

ts/packages/mcp/src/index.ts

@@ -0,0 +1,2 @@
+export { createMcpServer, run } from "./server.js";
+export { SocketManager, type SocketManagerConfig } from "./socket-manager.js";

ts/packages/mcp/src/server.ts

@@ -0,0 +1,174 @@
+/**
+ * TrustGraph MCP server.
+ *
+ * Exposes TrustGraph capabilities as MCP tools for AI assistants.
+ * Communicates with the TrustGraph gateway via WebSocket.
+ *
+ * Python reference: trustgraph-mcp/trustgraph/mcp_server/mcp.py
+ */
+
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { z } from "zod";
+import { SocketManager } from "./socket-manager.js";
+
+export function createMcpServer(config: {
+  gatewayUrl: string;
+  token?: string;
+  flowId?: string;
+}) {
+  const server = new McpServer({
+    name: "trustgraph",
+    version: "0.1.0",
+  });
+
+  const socket = new SocketManager({
+    gatewayUrl: config.gatewayUrl,
+    token: config.token,
+  });
+
+  const flowId = config.flowId ?? "default";
+
+  // --- Text Completion ---
+  server.tool(
+    "text_completion",
+    "Run a text completion using the configured LLM",
+    {
+      system: z.string().describe("System prompt"),
+      prompt: z.string().describe("User prompt"),
+    },
+    async ({ system, prompt }) => {
+      const resp = await socket.request("text-completion", { system, prompt }, { flowId }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: String(resp.response ?? resp) }] };
+    },
+  );
+
+  // --- Graph RAG ---
+  server.tool(
+    "graph_rag",
+    "Query the knowledge graph using RAG",
+    {
+      query: z.string().describe("Natural language query"),
+      entity_limit: z.number().optional().describe("Max entities to retrieve"),
+      triple_limit: z.number().optional().describe("Max triples per entity"),
+    },
+    async ({ query, entity_limit, triple_limit }) => {
+      const resp = await socket.request(
+        "graph-rag",
+        { query, entity_limit, triple_limit },
+        { flowId },
+      ) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: String(resp.response ?? resp) }] };
+    },
+  );
+
+  // --- Agent ---
+  server.tool(
+    "agent",
+    "Ask the TrustGraph agent a question",
+    {
+      question: z.string().describe("Question for the agent"),
+    },
+    async ({ question }) => {
+      const resp = await socket.request("agent", { question }, { flowId }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: String(resp.answer ?? resp) }] };
+    },
+  );
+
+  // --- Embeddings ---
+  server.tool(
+    "embeddings",
+    "Generate text embeddings",
+    {
+      text: z.array(z.string()).describe("Texts to embed"),
+    },
+    async ({ text }) => {
+      const resp = await socket.request("embeddings", { text }, { flowId }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
+    },
+  );
+
+  // --- Triples Query ---
+  server.tool(
+    "triples_query",
+    "Query the knowledge graph for triples matching a pattern",
+    {
+      s: z.string().optional().describe("Subject IRI"),
+      p: z.string().optional().describe("Predicate IRI"),
+      o: z.string().optional().describe("Object IRI or literal"),
+      limit: z.number().optional().describe("Max results"),
+    },
+    async ({ s, p, o, limit }) => {
+      const request: Record<string, unknown> = { limit };
+      if (s) request.s = { type: "IRI", iri: s };
+      if (p) request.p = { type: "IRI", iri: p };
+      if (o) request.o = { type: "IRI", iri: o };
+
+      const resp = await socket.request("triples-query", request, { flowId }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
+    },
+  );
+
+  // --- Graph Embeddings Query ---
+  server.tool(
+    "graph_embeddings_query",
+    "Find entities similar to a text query using vector embeddings",
+    {
+      query: z.string().describe("Text to find similar entities for"),
+      limit: z.number().optional().describe("Max results"),
+    },
+    async ({ query, limit }) => {
+      // First embed the query, then search
+      const embResp = await socket.request("embeddings", { text: [query] }, { flowId }) as { vectors: number[][] };
+      const resp = await socket.request(
+        "graph-embeddings-query",
+        { vectors: embResp.vectors, limit: limit ?? 10 },
+        { flowId },
+      ) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
+    },
+  );
+
+  // --- Config ---
+  server.tool(
+    "get_config",
+    "Get configuration values",
+    {
+      keys: z.array(z.string()).describe("Config keys to retrieve"),
+    },
+    async ({ keys }) => {
+      const resp = await socket.request("config", { operation: "get", keys }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
+    },
+  );
+
+  server.tool(
+    "put_config",
+    "Set configuration values",
+    {
+      values: z.record(z.unknown()).describe("Key-value pairs to set"),
+    },
+    async ({ values }) => {
+      const resp = await socket.request("config", { operation: "put", values }) as Record<string, unknown>;
+      return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
+    },
+  );
+
+  return { server, socket };
+}
+
+export async function run(): Promise<void> {
+  const { server, socket } = createMcpServer({
+    gatewayUrl: process.env.GATEWAY_URL ?? "ws://localhost:8088/api/v1/socket",
+    token: process.env.GATEWAY_SECRET,
+    flowId: process.env.FLOW_ID ?? "default",
+  });
+
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+
+  process.on("SIGINT", async () => {
+    await socket.close();
+    process.exit(0);
+  });
+}

ts/packages/mcp/src/socket-manager.ts

@@ -0,0 +1,147 @@
+/**
+ * WebSocket manager for communicating with the TrustGraph gateway.
+ *
+ * Maintains a persistent connection per user and handles request/response
+ * correlation via UUIDs.
+ *
+ * Python reference: trustgraph-mcp/trustgraph/mcp_server/tg_socket.py
+ */
+
+import WebSocket from "ws";
+import { randomUUID } from "node:crypto";
+
+export interface SocketManagerConfig {
+  gatewayUrl: string;
+  token?: string;
+}
+
+interface PendingRequest {
+  resolve: (value: unknown) => void;
+  reject: (error: Error) => void;
+  responses: unknown[];
+  streaming: boolean;
+  onChunk?: (chunk: unknown) => void;
+}
+
+export class SocketManager {
+  private ws: WebSocket | null = null;
+  private pending = new Map<string, PendingRequest>();
+  private connected = false;
+
+  constructor(private readonly config: SocketManagerConfig) {}
+
+  async connect(): Promise<void> {
+    if (this.connected) return;
+
+    const url = new URL(this.config.gatewayUrl);
+    if (this.config.token) {
+      url.searchParams.set("token", this.config.token);
+    }
+
+    return new Promise((resolve, reject) => {
+      this.ws = new WebSocket(url.toString());
+
+      this.ws.on("open", () => {
+        this.connected = true;
+        resolve();
+      });
+
+      this.ws.on("error", (err) => {
+        if (!this.connected) reject(err);
+        else console.error("[SocketManager] WebSocket error:", err);
+      });
+
+      this.ws.on("message", (data) => {
+        try {
+          const msg = JSON.parse(data.toString());
+          const { id, response, error, complete } = msg;
+
+          const req = this.pending.get(id);
+          if (!req) return;
+
+          if (error) {
+            req.reject(new Error(`${error.type}: ${error.message}`));
+            this.pending.delete(id);
+            return;
+          }
+
+          if (req.streaming && req.onChunk) {
+            req.onChunk(response);
+          }
+
+          req.responses.push(response);
+
+          if (complete) {
+            req.resolve(req.streaming ? req.responses : response);
+            this.pending.delete(id);
+          }
+        } catch (err) {
+          console.error("[SocketManager] Failed to parse message:", err);
+        }
+      });
+
+      this.ws.on("close", () => {
+        this.connected = false;
+        // Reject all pending requests
+        for (const [id, req] of this.pending) {
+          req.reject(new Error("WebSocket closed"));
+        }
+        this.pending.clear();
+      });
+    });
+  }
+
+  async request(
+    service: string,
+    requestData: Record<string, unknown>,
+    options?: {
+      flowId?: string;
+      timeoutMs?: number;
+      onChunk?: (chunk: unknown) => void;
+    },
+  ): Promise<unknown> {
+    await this.connect();
+    if (!this.ws) throw new Error("Not connected");
+
+    const id = randomUUID();
+    const timeoutMs = options?.timeoutMs ?? 300_000;
+
+    const msg = {
+      id,
+      service,
+      flow: options?.flowId ?? "default",
+      request: requestData,
+    };
+
+    return new Promise((resolve, reject) => {
+      const timer = setTimeout(() => {
+        this.pending.delete(id);
+        reject(new Error(`Request timed out after ${timeoutMs}ms`));
+      }, timeoutMs);
+
+      this.pending.set(id, {
+        resolve: (value) => {
+          clearTimeout(timer);
+          resolve(value);
+        },
+        reject: (err) => {
+          clearTimeout(timer);
+          reject(err);
+        },
+        responses: [],
+        streaming: !!options?.onChunk,
+        onChunk: options?.onChunk,
+      });
+
+      this.ws!.send(JSON.stringify(msg));
+    });
+  }
+
+  async close(): Promise<void> {
+    if (this.ws) {
+      this.ws.close();
+      this.ws = null;
+      this.connected = false;
+    }
+  }
+}