From 2a7ee29ef36cee6108eab41cef3948927594b7f3 Mon Sep 17 00:00:00 2001
From: Amp <amp@ampcode.com>
Date: Fri, 28 Aug 2026 13:40:08 +0000
Subject: [PATCH] feat(zettel): index graphs on note capture

Amp-Thread-ID: https://ampcode.com/threads/T-01a047e7-d402-72a8-ad22-eb5d3d3df6b1
Co-authored-by: Aditya Balakrishnan <adibalak1925@gmail.com>
---
 apps/zettel/src/tools/notes.test.ts | 73 ++++++++++++++++++++++++++---
 apps/zettel/src/tools/notes.ts      | 14 ++++++
 2 files changed, 81 insertions(+), 6 deletions(-)

diff --git a/apps/zettel/src/tools/notes.test.ts b/apps/zettel/src/tools/notes.test.ts
index c54d432..2aaf457 100644
--- a/apps/zettel/src/tools/notes.test.ts
+++ b/apps/zettel/src/tools/notes.test.ts
@@ -1,8 +1,15 @@
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
-const { traverseGraphStore } = vi.hoisted(() => ({
-  traverseGraphStore: vi.fn(),
-}));
+const { createDefaultGraphModelProvider, extractAndReplaceNoteGraph, traverseGraphStore, writeNote } =
+  vi.hoisted(() => ({
+    createDefaultGraphModelProvider: vi.fn(),
+    extractAndReplaceNoteGraph: vi.fn(),
+    traverseGraphStore: vi.fn(),
+    writeNote: vi.fn(),
+  }));
+
+vi.mock("../notes/graph-extraction.js", () => ({ extractAndReplaceNoteGraph }));
+vi.mock("../notes/graph-model-provider.js", () => ({ createDefaultGraphModelProvider }));
 
 vi.mock("../notes/store.js", () => ({
   addLink: vi.fn(),
@@ -10,10 +17,64 @@ vi.mock("../notes/store.js", () => ({
   readNote: vi.fn(),
   searchNotes: vi.fn(),
   traverseGraphStore,
-  writeNote: vi.fn(),
+  writeNote,
 }));
 
-import { traverseGraph, traverseGraphSchema } from "./notes.js";
+import { createNote, traverseGraph, traverseGraphSchema } from "./notes.js";
+
+describe("createNote graph indexing", () => {
+  beforeEach(() => {
+    vi.stubEnv("GROQ_API_KEY", "test-key");
+    createDefaultGraphModelProvider.mockReset();
+    extractAndReplaceNoteGraph.mockReset();
+    writeNote.mockReset();
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  it("indexes a successfully created note for the owning tenant", async () => {
+    const provider = { extractGraph: vi.fn() };
+    createDefaultGraphModelProvider.mockReturnValue(provider);
+    writeNote.mockResolvedValue({ id: "note-1" });
+    extractAndReplaceNoteGraph.mockResolvedValue({ entities: [], relations: [] });
+
+    await expect(createNote({ content: "Ada designed an engine.", userId: "tenant-a" })).resolves.toEqual({
+      success: true,
+      id: "note-1",
+    });
+    expect(extractAndReplaceNoteGraph).toHaveBeenCalledWith(
+      "tenant-a",
+      "note-1",
+      "Ada designed an engine.",
+      provider,
+    );
+  });
+
+  it("keeps the captured note when graph extraction fails", async () => {
+    createDefaultGraphModelProvider.mockReturnValue({ extractGraph: vi.fn() });
+    writeNote.mockResolvedValue({ id: "note-2" });
+    extractAndReplaceNoteGraph.mockRejectedValue(new Error("model unavailable"));
+
+    await expect(createNote({ content: "A durable note." })).resolves.toEqual({
+      success: true,
+      id: "note-2",
+    });
+  });
+
+  it("does not call the model when graph extraction is not configured", async () => {
+    vi.stubEnv("GROQ_API_KEY", "");
+    writeNote.mockResolvedValue({ id: "note-3" });
+
+    await expect(createNote({ content: "An offline note." })).resolves.toEqual({
+      success: true,
+      id: "note-3",
+    });
+    expect(createDefaultGraphModelProvider).not.toHaveBeenCalled();
+    expect(extractAndReplaceNoteGraph).not.toHaveBeenCalled();
+  });
+});
 
 describe("traverseGraph tool", () => {
   beforeEach(() => {
diff --git a/apps/zettel/src/tools/notes.ts b/apps/zettel/src/tools/notes.ts
index 92bf3cb..f95e32a 100644
--- a/apps/zettel/src/tools/notes.ts
+++ b/apps/zettel/src/tools/notes.ts
@@ -16,6 +16,8 @@ import {
   backlinksOf,
   traverseGraphStore,
 } from "../notes/store.js";
+import { extractAndReplaceNoteGraph } from "../notes/graph-extraction.js";
+import { createDefaultGraphModelProvider } from "../notes/graph-model-provider.js";
 
 // ── createNote ────────────────────────────────────────────────────────────────
 
@@ -35,6 +37,18 @@ export async function createNote(args: CreateNoteInput & { userId?: string }) {
   const userId = args.userId ?? "default";
   try {
     const note = await writeNote(userId, { content, title, tags, source });
+    if (process.env.GROQ_API_KEY) {
+      try {
+        await extractAndReplaceNoteGraph(
+          userId,
+          note.id,
+          content,
+          createDefaultGraphModelProvider(),
+        );
+      } catch (err) {
+        console.error("[zettel] Graph indexing failed:", err);
+      }
+    }
     return { success: true, id: note.id };
   } catch (err) {
     return { success: false, error: err instanceof Error ? err.message : String(err) };