@contextableai/openclaw-memory-graphiti 0.1.2 → 0.2.1

package/graphiti.ts

@@ -6,7 +6,8 @@
  * session ID tracking, and SSE response parsing.
  *
  * Wraps core tools: add_memory, search_nodes, search_memory_facts,
- * get_episodes, delete_episode, get_status.
+ * get_episodes, delete_episode, get_status, get_entity_edge,
+ * delete_entity_edge, clear_graph.
  */
 
 import { randomUUID } from "node:crypto";
@@ -49,6 +50,8 @@ export type GraphitiFact = {
 
 export type AddEpisodeResult = {
   episode_uuid: string;
+  /** Resolves to the real server-side UUID once Graphiti finishes processing. */
+  resolvedUuid: Promise<string>;
 };
 
 type JsonRpcRequest = {
@@ -74,8 +77,17 @@ export class GraphitiClient {
   private sessionId: string | null = null;
   private initPromise: Promise<void> | null = null;
 
+  /** Polling interval (ms) for UUID resolution after addEpisode. */
+  uuidPollIntervalMs = 3000;
+  /** Max polling attempts for UUID resolution (total wait = interval * attempts). */
+  uuidPollMaxAttempts = 30;
+
   constructor(private readonly endpoint: string) {}
 
+  private get mcpUrl(): string {
+    return `${this.endpoint}/mcp`;
+  }
+
   // --------------------------------------------------------------------------
   // MCP Session Lifecycle
   // --------------------------------------------------------------------------
@@ -100,7 +112,7 @@ export class GraphitiClient {
       },
     };
 
-    const response = await fetch(`${this.endpoint}/mcp`, {
+    const response = await fetch(this.mcpUrl, {
       method: "POST",
       headers: {
         "Content-Type": "application/json",
@@ -128,7 +140,7 @@ export class GraphitiClient {
     if (this.sessionId) {
       headers["mcp-session-id"] = this.sessionId;
     }
-    await fetch(`${this.endpoint}/mcp`, {
+    await fetch(this.mcpUrl, {
       method: "POST",
       headers,
       body: JSON.stringify({ jsonrpc: "2.0", method: "notifications/initialized" }),
@@ -138,7 +150,7 @@ export class GraphitiClient {
   async close(): Promise<void> {
     if (this.sessionId) {
       try {
-        await fetch(`${this.endpoint}/mcp`, {
+        await fetch(this.mcpUrl, {
           method: "DELETE",
           headers: { "mcp-session-id": this.sessionId },
         });
@@ -172,7 +184,7 @@ export class GraphitiClient {
       headers["mcp-session-id"] = this.sessionId;
     }
 
-    const response = await fetch(`${this.endpoint}/mcp`, {
+    const response = await fetch(this.mcpUrl, {
       method: "POST",
       headers,
       body: JSON.stringify(request),
@@ -255,7 +267,7 @@ export class GraphitiClient {
     // Note: The Graphiti MCP server's uuid parameter is for re-processing
     // existing episodes, NOT for setting the UUID of new ones. We generate
     // a client-side tracking UUID instead (it won't match the server-side UUID).
-    const trackingUuid = randomUUID();
+    const trackingUuid = `tmp-${randomUUID()}`;
 
     // Prepend custom extraction instructions to episode_body since the MCP
     // server doesn't support custom_extraction_instructions as a parameter.
@@ -281,7 +293,39 @@ export class GraphitiClient {
     }
 
     await this.callTool("add_memory", args);
-    return { episode_uuid: trackingUuid };
+
+    // Graphiti's add_memory queues the episode for async LLM processing and
+    // returns only a "queued" message — no UUID. We poll getEpisodes in the
+    // background to discover the real server-side UUID by name match.
+    let resolvedUuid: Promise<string>;
+    if (params.group_id) {
+      resolvedUuid = this.resolveEpisodeUuid(params.name, params.group_id);
+      resolvedUuid.catch(() => {}); // Prevent unhandled rejection if caller ignores
+    } else {
+      resolvedUuid = Promise.resolve(trackingUuid);
+    }
+
+    return { episode_uuid: trackingUuid, resolvedUuid };
+  }
+
+  private async resolveEpisodeUuid(name: string, groupId: string): Promise<string> {
+    const totalTimeoutSec = (this.uuidPollMaxAttempts * this.uuidPollIntervalMs) / 1000;
+    for (let i = 0; i < this.uuidPollMaxAttempts; i++) {
+      await new Promise((r) => setTimeout(r, this.uuidPollIntervalMs));
+      try {
+        const episodes = await this.getEpisodes(groupId, 50);
+        const match = episodes.find((ep) => ep.name === name);
+        if (match) return match.uuid;
+      } catch (err) {
+        // Log transient errors to aid debugging
+        console.warn(
+          `[graphiti] UUID poll ${i + 1}/${this.uuidPollMaxAttempts} for "${name}" in "${groupId}" failed: ${err instanceof Error ? err.message : String(err)}`,
+        );
+      }
+    }
+    throw new Error(
+      `Failed to resolve episode UUID for "${name}" in group "${groupId}" after ${totalTimeoutSec}s — episode not yet visible in get_episodes (Graphiti LLM processing may still be running)`,
+    );
   }
 
   async getEpisodes(groupId: string, lastN: number): Promise<GraphitiEpisode[]> {
@@ -305,6 +349,7 @@ export class GraphitiClient {
     group_id?: string;
     group_ids?: string[];
     limit?: number;
+    entity_types?: string[];
   }): Promise<GraphitiNode[]> {
     const args: Record<string, unknown> = {
       query: params.query,
@@ -316,6 +361,9 @@ export class GraphitiClient {
     if (params.limit !== undefined) {
       args.max_nodes = params.limit;
     }
+    if (params.entity_types && params.entity_types.length > 0) {
+      args.entity_types = params.entity_types;
+    }
 
     const result = await this.callTool("search_nodes", args);
     return parseToolResult<GraphitiNode[]>(result, "nodes");
@@ -326,6 +374,7 @@ export class GraphitiClient {
     group_id?: string;
     group_ids?: string[];
     limit?: number;
+    center_node_uuid?: string;
     /** @deprecated No longer supported by the Graphiti MCP server */
     created_after?: string;
   }): Promise<GraphitiFact[]> {
@@ -339,10 +388,34 @@ export class GraphitiClient {
     if (params.limit !== undefined) {
       args.max_facts = params.limit;
     }
+    if (params.center_node_uuid) {
+      args.center_node_uuid = params.center_node_uuid;
+    }
 
     const result = await this.callTool("search_memory_facts", args);
     return parseToolResult<GraphitiFact[]>(result, "facts");
   }
+
+  // --------------------------------------------------------------------------
+  // Entity Edge Operations
+  // --------------------------------------------------------------------------
+
+  async getEntityEdge(uuid: string): Promise<GraphitiFact> {
+    const result = await this.callTool("get_entity_edge", { uuid });
+    return parseJsonResult<GraphitiFact>(result);
+  }
+
+  async deleteEntityEdge(uuid: string): Promise<void> {
+    await this.callTool("delete_entity_edge", { uuid });
+  }
+
+  async clearGraph(groupIds?: string[]): Promise<void> {
+    const args: Record<string, unknown> = {};
+    if (groupIds && groupIds.length > 0) {
+      args.group_ids = groupIds;
+    }
+    await this.callTool("clear_graph", args);
+  }
 }
 
 // ============================================================================