@contextableai/openclaw-memory-graphiti 0.1.2 → 0.2.0

package/index.ts

@@ -13,9 +13,7 @@
 import type { OpenClawPluginApi } from "openclaw/plugin-sdk";
 import { Type } from "@sinclair/typebox";
 import { readFileSync } from "node:fs";
-import { readdir, readFile, stat } from "node:fs/promises";
-import { join, dirname, basename, resolve } from "node:path";
-import { homedir } from "node:os";
+import { join, dirname } from "node:path";
 import { fileURLToPath } from "node:url";
 
 import { graphitiMemoryConfigSchema } from "./config.js";
@@ -35,6 +33,7 @@ import {
   formatDualResults,
   deduplicateSessionResults,
 } from "./search.js";
+import { registerCommands } from "./cli.js";
 
 // ============================================================================
 // Session helpers
@@ -74,6 +73,14 @@ const memoryGraphitiPlugin = {
     // Track current session ID — updated from hook event context
     let currentSessionId: string | undefined;
 
+    // Track most recent ZedToken from SpiceDB writes for causal consistency.
+    // Reads use at_least_as_fresh(token) after own writes, minimize_latency otherwise.
+    let lastWriteToken: string | undefined;
+
+    // Map tracking UUIDs to resolvedUuid promises so memory_forget can translate
+    // a tracking UUID (from memory_store) to the real server-side UUID.
+    const pendingResolutions = new Map<string, Promise<string>>();
+
     api.logger.info(
       `memory-graphiti: registered (graphiti: ${cfg.graphiti.endpoint}, spicedb: ${cfg.spicedb.endpoint})`,
     );
@@ -97,16 +104,28 @@ const memoryGraphitiPlugin = {
               { description: "Memory scope: 'session' (current session only), 'long-term' (persistent), or 'all' (both). Default: 'all'" },
             ),
           ),
+          entity_types: Type.Optional(
+            Type.Array(Type.String(), {
+              description: "Filter by entity type (e.g., 'Preference', 'Organization', 'Procedure')",
+            }),
+          ),
+          center_node_uuid: Type.Optional(
+            Type.String({
+              description: "UUID of an entity node to center the fact search around",
+            }),
+          ),
         }),
         async execute(_toolCallId, params) {
-          const { query, limit = 10, scope = "all" } = params as {
+          const { query, limit = 10, scope = "all", entity_types, center_node_uuid } = params as {
             query: string;
             limit?: number;
             scope?: "session" | "long-term" | "all";
+            entity_types?: string[];
+            center_node_uuid?: string;
           };
 
           // 1. Get authorized groups for current subject
-          const authorizedGroups = await lookupAuthorizedGroups(spicedb, currentSubject);
+          const authorizedGroups = await lookupAuthorizedGroups(spicedb, currentSubject, lastWriteToken);
 
           if (authorizedGroups.length === 0) {
             return {
@@ -145,12 +164,13 @@ const memoryGraphitiPlugin = {
           }
 
           // 3. Parallel search across groups
+          const searchOpts = { entityTypes: entity_types, centerNodeUuid: center_node_uuid };
           const [longTermResults, rawSessionResults] = await Promise.all([
             longTermGroups.length > 0
-              ? searchAuthorizedMemories(graphiti, { query, groupIds: longTermGroups, limit })
+              ? searchAuthorizedMemories(graphiti, { query, groupIds: longTermGroups, limit, ...searchOpts })
               : Promise.resolve([]),
             sessionGroups.length > 0
-              ? searchAuthorizedMemories(graphiti, { query, groupIds: sessionGroups, limit })
+              ? searchAuthorizedMemories(graphiti, { query, groupIds: sessionGroups, limit, ...searchOpts })
               : Promise.resolve([]),
           ]);
 
@@ -247,13 +267,14 @@ const memoryGraphitiPlugin = {
 
           if (isOwnSession) {
             try {
-              await ensureGroupMembership(spicedb, targetGroupId, currentSubject);
+              const token = await ensureGroupMembership(spicedb, targetGroupId, currentSubject);
+              if (token) lastWriteToken = token;
             } catch {
               api.logger.warn(`memory-graphiti: failed to ensure membership in ${targetGroupId}`);
             }
           } else {
             // All other groups (non-session AND foreign session) require write permission
-            const allowed = await canWriteToGroup(spicedb, currentSubject, targetGroupId);
+            const allowed = await canWriteToGroup(spicedb, currentSubject, targetGroupId, lastWriteToken);
             if (!allowed) {
               return {
                 content: [
@@ -276,19 +297,35 @@ const memoryGraphitiPlugin = {
             custom_extraction_instructions: cfg.customInstructions,
           });
 
-          const fragmentId = result.episode_uuid;
-
-          // 2. Write authorization relationships in SpiceDB
+          // 2. Write authorization relationships in SpiceDB (background).
+          // Graphiti processes episodes asynchronously — the real UUID isn't
+          // available immediately. resolvedUuid polls in the background and
+          // writes SpiceDB relationships once the real UUID is known, so the
+          // tool response isn't blocked.
           const involvedSubjects: Subject[] = involves.map((id) => ({
             type: "person" as const,
             id,
           }));
 
-          await writeFragmentRelationships(spicedb, {
-            fragmentId,
-            groupId: targetGroupId,
-            sharedBy: currentSubject,
-            involves: involvedSubjects,
+          // Chain UUID resolution → SpiceDB write, and store the promise so
+          // memory_forget can await both before checking permissions.
+          const deferredWrite = result.resolvedUuid
+            .then(async (realUuid) => {
+              const writeToken = await writeFragmentRelationships(spicedb, {
+                fragmentId: realUuid,
+                groupId: targetGroupId,
+                sharedBy: currentSubject,
+                involves: involvedSubjects,
+              });
+              if (writeToken) lastWriteToken = writeToken;
+              return realUuid;
+            });
+
+          pendingResolutions.set(result.episode_uuid, deferredWrite);
+          deferredWrite.catch((err) => {
+            api.logger.warn(
+              `memory-graphiti: deferred SpiceDB write failed for memory_store: ${err}`,
+            );
           });
 
           return {
@@ -300,7 +337,7 @@ const memoryGraphitiPlugin = {
             ],
             details: {
               action: "created",
-              episodeId: fragmentId,
+              episodeId: result.episode_uuid,
               groupId: targetGroupId,
               longTerm,
               involves,
@@ -315,15 +352,71 @@ const memoryGraphitiPlugin = {
       {
         name: "memory_forget",
         label: "Memory Forget",
-        description: "Delete a memory episode. Requires delete permission.",
+        description:
+          "Delete a memory episode or fact. Provide either episode_id or fact_id (not both). Requires delete/write permission.",
         parameters: Type.Object({
-          episode_id: Type.String({ description: "Episode UUID to delete" }),
+          episode_id: Type.Optional(Type.String({ description: "Episode UUID to delete" })),
+          fact_id: Type.Optional(Type.String({ description: "Fact (entity edge) UUID to delete" })),
         }),
         async execute(_toolCallId, params) {
-          const { episode_id } = params as { episode_id: string };
+          const { episode_id, fact_id } = params as { episode_id?: string; fact_id?: string };
+
+          if (!episode_id && !fact_id) {
+            return {
+              content: [{ type: "text", text: "Either episode_id or fact_id must be provided." }],
+              details: { action: "error" },
+            };
+          }
+
+          // --- Fact deletion ---
+          if (fact_id) {
+            // 1. Fetch fact to get group_id for authorization
+            let fact: Awaited<ReturnType<typeof graphiti.getEntityEdge>>;
+            try {
+              fact = await graphiti.getEntityEdge(fact_id);
+            } catch {
+              return {
+                content: [{ type: "text", text: `Fact ${fact_id} not found.` }],
+                details: { action: "error", factId: fact_id },
+              };
+            }
+
+            // 2. Check write permission on the fact's group
+            const allowed = await canWriteToGroup(spicedb, currentSubject, fact.group_id, lastWriteToken);
+            if (!allowed) {
+              return {
+                content: [{ type: "text", text: `Permission denied: cannot delete fact ${fact_id}` }],
+                details: { action: "denied", factId: fact_id },
+              };
+            }
+
+            // 3. Delete fact from Graphiti
+            await graphiti.deleteEntityEdge(fact_id);
+
+            return {
+              content: [{ type: "text", text: `Fact ${fact_id} forgotten.` }],
+              details: { action: "deleted", factId: fact_id },
+            };
+          }
+
+          // --- Episode deletion (existing flow) ---
+
+          // Resolve tracking UUID → real server-side UUID if this came
+          // from a recent memory_store call. Awaits the background
+          // resolution so permission checks use the correct UUID.
+          let effectiveId = episode_id!;
+          const pending = pendingResolutions.get(episode_id!);
+          if (pending) {
+            try {
+              effectiveId = await pending;
+            } catch {
+              // Resolution failed — try with original UUID
+            }
+            pendingResolutions.delete(episode_id!);
+          }
 
           // 1. Check delete permission
-          const allowed = await canDeleteFragment(spicedb, currentSubject, episode_id);
+          const allowed = await canDeleteFragment(spicedb, currentSubject, effectiveId, lastWriteToken);
           if (!allowed) {
             return {
               content: [
@@ -337,15 +430,12 @@ const memoryGraphitiPlugin = {
           }
 
           // 2. Delete from Graphiti
-          await graphiti.deleteEpisode(episode_id);
+          await graphiti.deleteEpisode(effectiveId);
 
           // 3. Clean up SpiceDB relationships (best-effort)
           try {
-            await deleteFragmentRelationships(spicedb, episode_id, {
-              fragmentId: episode_id,
-              groupId: cfg.graphiti.defaultGroupId,
-              sharedBy: currentSubject,
-            });
+            const deleteToken = await deleteFragmentRelationships(spicedb, effectiveId);
+            if (deleteToken) lastWriteToken = deleteToken;
           } catch {
             api.logger.warn(
               `memory-graphiti: failed to clean up SpiceDB relationships for ${episode_id}`,
@@ -410,298 +500,13 @@ const memoryGraphitiPlugin = {
         const mem = program
           .command("graphiti-mem")
           .description("Graphiti + SpiceDB memory plugin commands");
-
-        mem
-          .command("search")
-          .description("Search memories with authorization")
-          .argument("<query>", "Search query")
-          .option("--limit <n>", "Max results", "10")
-          .option("--scope <scope>", "Memory scope: session, long-term, all", "all")
-          .action(async (query: string, opts: { limit: string; scope: string }) => {
-            const authorizedGroups = await lookupAuthorizedGroups(spicedb, currentSubject);
-            if (authorizedGroups.length === 0) {
-              console.log("No accessible memory groups.");
-              return;
-            }
-
-            console.log(`Searching ${authorizedGroups.length} authorized groups...`);
-            const results = await searchAuthorizedMemories(graphiti, {
-              query,
-              groupIds: authorizedGroups,
-              limit: parseInt(opts.limit),
-            });
-
-            if (results.length === 0) {
-              console.log("No results found.");
-              return;
-            }
-
-            console.log(JSON.stringify(results, null, 2));
-          });
-
-        mem
-          .command("episodes")
-          .description("List recent episodes")
-          .option("--last <n>", "Number of episodes", "10")
-          .option("--group <id>", "Group ID", cfg.graphiti.defaultGroupId)
-          .action(async (opts: { last: string; group: string }) => {
-            const episodes = await graphiti.getEpisodes(opts.group, parseInt(opts.last));
-            console.log(JSON.stringify(episodes, null, 2));
-          });
-
-        mem
-          .command("status")
-          .description("Check SpiceDB + Graphiti health")
-          .action(async () => {
-            const graphitiOk = await graphiti.healthCheck();
-            let spicedbOk = false;
-            try {
-              await spicedb.readSchema();
-              spicedbOk = true;
-            } catch {
-              // unreachable
-            }
-
-            console.log(`Graphiti MCP: ${graphitiOk ? "OK" : "UNREACHABLE"} (${cfg.graphiti.endpoint})`);
-            console.log(`SpiceDB:      ${spicedbOk ? "OK" : "UNREACHABLE"} (${cfg.spicedb.endpoint})`);
-          });
-
-        mem
-          .command("schema-write")
-          .description("Write/update SpiceDB authorization schema")
-          .action(async () => {
-            const schemaPath = join(dirname(fileURLToPath(import.meta.url)), "schema.zed");
-            const schema = readFileSync(schemaPath, "utf-8");
-            await spicedb.writeSchema(schema);
-            console.log("SpiceDB schema written successfully.");
-          });
-
-        mem
-          .command("groups")
-          .description("List authorized groups for current subject")
-          .action(async () => {
-            const groups = await lookupAuthorizedGroups(spicedb, currentSubject);
-            if (groups.length === 0) {
-              console.log("No authorized groups.");
-              return;
-            }
-            console.log(`Authorized groups for ${currentSubject.type}:${currentSubject.id}:`);
-            for (const g of groups) {
-              console.log(`  - ${g}`);
-            }
-          });
-
-        mem
-          .command("add-member")
-          .description("Add a subject to a group")
-          .argument("<group-id>", "Group ID")
-          .argument("<subject-id>", "Subject ID")
-          .option("--type <type>", "Subject type (agent|person)", "person")
-          .action(async (groupId: string, subjectId: string, opts: { type: string }) => {
-            const subjectType = opts.type === "agent" ? "agent" : "person";
-            await ensureGroupMembership(spicedb, groupId, {
-              type: subjectType as "agent" | "person",
-              id: subjectId,
-            });
-            console.log(`Added ${subjectType}:${subjectId} to group:${groupId}`);
-          });
-
-        mem
-          .command("import")
-          .description("Import workspace markdown files (and optionally session transcripts) into Graphiti")
-          .option("--workspace <path>", "Workspace directory", join(homedir(), ".openclaw", "workspace"))
-          .option("--include-sessions", "Also import session JSONL transcripts", false)
-          .option("--sessions-only", "Only import session transcripts (skip workspace files)", false)
-          .option("--session-dir <path>", "Session transcripts directory", join(homedir(), ".openclaw", "agents", "main", "sessions"))
-          .option("--group <id>", "Target group for workspace files", cfg.graphiti.defaultGroupId)
-          .option("--dry-run", "List files without importing", false)
-          .action(async (opts: {
-            workspace: string;
-            includeSessions: boolean;
-            sessionsOnly: boolean;
-            sessionDir: string;
-            group: string;
-            dryRun: boolean;
-          }) => {
-            const workspacePath = resolve(opts.workspace);
-            const targetGroup = opts.group;
-            const importSessions = opts.includeSessions || opts.sessionsOnly;
-            const importWorkspace = !opts.sessionsOnly;
-
-            // Discover workspace markdown files
-            let mdFiles: string[] = [];
-            try {
-              const entries = await readdir(workspacePath);
-              mdFiles = entries.filter((f) => f.endsWith(".md")).sort();
-            } catch {
-              console.error(`Cannot read workspace directory: ${workspacePath}`);
-              return;
-            }
-
-            // Also check for memory/ subdirectory
-            try {
-              const memDir = join(workspacePath, "memory");
-              const memEntries = await readdir(memDir);
-              for (const f of memEntries) {
-                if (f.endsWith(".md")) {
-                  mdFiles.push(join("memory", f));
-                }
-              }
-            } catch {
-              // No memory/ subdirectory — that's fine
-            }
-
-            if (mdFiles.length === 0) {
-              console.log("No markdown files found in workspace.");
-              return;
-            }
-
-            console.log(`Found ${mdFiles.length} workspace file(s) in ${workspacePath}:`);
-            for (const f of mdFiles) {
-              const filePath = join(workspacePath, f);
-              const info = await stat(filePath);
-              console.log(`  ${f} (${info.size} bytes)`);
-            }
-
-            if (opts.dryRun) {
-              console.log("\n[dry-run] No files imported.");
-              if (importSessions) {
-                const sessionPath = resolve(opts.sessionDir);
-                try {
-                  const sessions = (await readdir(sessionPath)).filter((f) => f.endsWith(".jsonl"));
-                  console.log(`\nFound ${sessions.length} session transcript(s) in ${sessionPath}:`);
-                  for (const f of sessions) {
-                    const info = await stat(join(sessionPath, f));
-                    console.log(`  ${f} (${info.size} bytes)`);
-                  }
-                } catch {
-                  console.log(`\nCannot read session directory: ${sessionPath}`);
-                }
-              }
-              return;
-            }
-
-            // Import workspace files
-            if (importWorkspace) {
-              console.log(`\nImporting workspace files to group: ${targetGroup}`);
-              let imported = 0;
-              for (const f of mdFiles) {
-                const filePath = join(workspacePath, f);
-                const content = await readFile(filePath, "utf-8");
-                if (!content.trim()) {
-                  console.log(`  Skipping ${f} (empty)`);
-                  continue;
-                }
-                try {
-                  const result = await graphiti.addEpisode({
-                    name: f,
-                    episode_body: content,
-                    source_description: `Imported from OpenClaw workspace: ${f}`,
-                    group_id: targetGroup,
-                    source: "text",
-                  });
-                  await writeFragmentRelationships(spicedb, {
-                    fragmentId: result.episode_uuid,
-                    groupId: targetGroup,
-                    sharedBy: currentSubject,
-                  });
-                  console.log(`  Imported ${f} (${content.length} bytes) → episode ${result.episode_uuid}`);
-                  imported++;
-                } catch (err) {
-                  console.error(`  Failed to import ${f}: ${err instanceof Error ? err.message : String(err)}`);
-                }
-              }
-              console.log(`\nWorkspace import complete: ${imported}/${mdFiles.length} files.`);
-            }
-
-            // Import session transcripts
-            if (importSessions) {
-              const sessionPath = resolve(opts.sessionDir);
-              let jsonlFiles: string[] = [];
-              try {
-                jsonlFiles = (await readdir(sessionPath)).filter((f) => f.endsWith(".jsonl")).sort();
-              } catch {
-                console.error(`\nCannot read session directory: ${sessionPath}`);
-                return;
-              }
-
-              if (jsonlFiles.length === 0) {
-                console.log("\nNo session transcripts found.");
-                return;
-              }
-
-              console.log(`\nImporting ${jsonlFiles.length} session transcript(s)...`);
-              let sessionsImported = 0;
-              for (const f of jsonlFiles) {
-                const sessionId = basename(f, ".jsonl");
-                const sessionGroup = sessionGroupId(sessionId);
-                const filePath = join(sessionPath, f);
-                const raw = await readFile(filePath, "utf-8");
-                const lines = raw.split("\n").filter(Boolean);
-
-                // Extract user/assistant message text from JSONL
-                const conversationLines: string[] = [];
-                for (const line of lines) {
-                  try {
-                    const entry = JSON.parse(line) as Record<string, unknown>;
-                    // OpenClaw JSONL format: {"type":"message","message":{"role":"user","content":[...]}}
-                    const msg = (entry.type === "message" && entry.message && typeof entry.message === "object")
-                      ? entry.message as Record<string, unknown>
-                      : entry;
-                    const role = msg.role as string | undefined;
-                    if (role !== "user" && role !== "assistant") continue;
-                    const content = msg.content;
-                    let text = "";
-                    if (typeof content === "string") {
-                      text = content;
-                    } else if (Array.isArray(content)) {
-                      text = content
-                        .filter((b: unknown) =>
-                          typeof b === "object" && b !== null &&
-                          (b as Record<string, unknown>).type === "text" &&
-                          typeof (b as Record<string, unknown>).text === "string",
-                        )
-                        .map((b: unknown) => (b as Record<string, unknown>).text as string)
-                        .join("\n");
-                    }
-                    if (text && text.length >= 5 && !text.includes("<relevant-memories>") && !text.includes("<memory-tools>")) {
-                      const roleLabel = role === "user" ? "User" : "Assistant";
-                      conversationLines.push(`${roleLabel}: ${text}`);
-                    }
-                  } catch {
-                    // Skip malformed JSONL lines
-                  }
-                }
-
-                if (conversationLines.length === 0) {
-                  console.log(`  Skipping ${f} (no user/assistant messages)`);
-                  continue;
-                }
-
-                try {
-                  await ensureGroupMembership(spicedb, sessionGroup, currentSubject);
-                  const episodeBody = conversationLines.join("\n");
-                  const result = await graphiti.addEpisode({
-                    name: `session_${sessionId}`,
-                    episode_body: episodeBody,
-                    source_description: `Imported session transcript: ${sessionId}`,
-                    group_id: sessionGroup,
-                    source: "message",
-                  });
-                  await writeFragmentRelationships(spicedb, {
-                    fragmentId: result.episode_uuid,
-                    groupId: sessionGroup,
-                    sharedBy: currentSubject,
-                  });
-                  console.log(`  Imported ${f} (${conversationLines.length} messages) → episode ${result.episode_uuid} [group: ${sessionGroup}]`);
-                  sessionsImported++;
-                } catch (err) {
-                  console.error(`  Failed to import ${f}: ${err instanceof Error ? err.message : String(err)}`);
-                }
-              }
-              console.log(`\nSession import complete: ${sessionsImported}/${jsonlFiles.length} transcripts.`);
-            }
-          });
+        registerCommands(mem, {
+          graphiti,
+          spicedb,
+          cfg,
+          currentSubject,
+          getLastWriteToken: () => lastWriteToken,
+        });
       },
       { commands: ["graphiti-mem"] },
     );
@@ -711,10 +516,10 @@ const memoryGraphitiPlugin = {
     // ========================================================================
 
     if (cfg.autoRecall) {
-      api.on("before_agent_start", async (event) => {
-        // Track session ID from event context
-        if (event.ctx?.sessionKey) {
-          currentSessionId = event.ctx.sessionKey as string;
+      api.on("before_agent_start", async (event, ctx) => {
+        // Track session ID from hook context
+        if (ctx?.sessionKey) {
+          currentSessionId = ctx.sessionKey;
         }
 
         if (!event.prompt || event.prompt.length < 5) {
@@ -722,7 +527,7 @@ const memoryGraphitiPlugin = {
         }
 
         try {
-          const authorizedGroups = await lookupAuthorizedGroups(spicedb, currentSubject);
+          const authorizedGroups = await lookupAuthorizedGroups(spicedb, currentSubject, lastWriteToken);
           if (authorizedGroups.length === 0) {
             return;
           }
@@ -787,10 +592,10 @@ const memoryGraphitiPlugin = {
     }
 
     if (cfg.autoCapture) {
-      api.on("agent_end", async (event) => {
-        // Track session ID from event context
-        if (event.ctx?.sessionKey) {
-          currentSessionId = event.ctx.sessionKey as string;
+      api.on("agent_end", async (event, ctx) => {
+        // Track session ID from hook context
+        if (ctx?.sessionKey) {
+          currentSessionId = ctx.sessionKey;
         }
 
         if (!event.success || !event.messages || event.messages.length === 0) {
@@ -863,12 +668,13 @@ const memoryGraphitiPlugin = {
 
           if (isOwnSession) {
             try {
-              await ensureGroupMembership(spicedb, targetGroupId, currentSubject);
+              const token = await ensureGroupMembership(spicedb, targetGroupId, currentSubject);
+              if (token) lastWriteToken = token;
             } catch {
               // Best-effort
             }
           } else {
-            const allowed = await canWriteToGroup(spicedb, currentSubject, targetGroupId);
+            const allowed = await canWriteToGroup(spicedb, currentSubject, targetGroupId, lastWriteToken);
             if (!allowed) {
               api.logger.warn(`memory-graphiti: auto-capture denied for group ${targetGroupId}`);
               return;
@@ -883,11 +689,21 @@ const memoryGraphitiPlugin = {
             custom_extraction_instructions: cfg.customInstructions,
           });
 
-          await writeFragmentRelationships(spicedb, {
-            fragmentId: result.episode_uuid,
-            groupId: targetGroupId,
-            sharedBy: currentSubject,
-          });
+          // SpiceDB writes use the real UUID once Graphiti finishes processing
+          result.resolvedUuid
+            .then(async (realUuid) => {
+              const writeToken = await writeFragmentRelationships(spicedb, {
+                fragmentId: realUuid,
+                groupId: targetGroupId,
+                sharedBy: currentSubject,
+              });
+              if (writeToken) lastWriteToken = writeToken;
+            })
+            .catch((err) => {
+              api.logger.warn(
+                `memory-graphiti: deferred SpiceDB write (auto-capture) failed: ${err}`,
+              );
+            });
 
           api.logger.info(
             `memory-graphiti: auto-captured ${conversationLines.length} messages as batch episode to ${targetGroupId}`,
@@ -927,11 +743,12 @@ const memoryGraphitiPlugin = {
         // Ensure current subject is a member of the default group
         if (spicedbOk) {
           try {
-            await ensureGroupMembership(
+            const token = await ensureGroupMembership(
               spicedb,
               cfg.graphiti.defaultGroupId,
               currentSubject,
             );
+            if (token) lastWriteToken = token;
           } catch {
             api.logger.warn("memory-graphiti: failed to ensure default group membership");
           }