From 60a12ae0f1b4e18610beffb16399821777cb0a26 Mon Sep 17 00:00:00 2001 From: Amp Date: Fri, 28 Aug 2026 13:14:47 +0000 Subject: [PATCH] feat(zettel): add tenant-safe graph store Amp-Thread-ID: https://ampcode.com/threads/T-01a047e7-d402-72a8-ad22-eb5d3d3df6b1 Co-authored-by: Aditya Balakrishnan --- apps/zettel/src/notes/store.test.ts | 144 ++++++++++++++++-- apps/zettel/src/notes/store.ts | 222 ++++++++++++++++++++++++++++ 2 files changed, 355 insertions(+), 11 deletions(-) diff --git a/apps/zettel/src/notes/store.test.ts b/apps/zettel/src/notes/store.test.ts index eaf7eb4..2d1ea05 100644 --- a/apps/zettel/src/notes/store.test.ts +++ b/apps/zettel/src/notes/store.test.ts @@ -1,13 +1,10 @@ -import path from "node:path"; -import os from "node:os"; -import { mkdirSync } from "node:fs"; +import { describe, it, expect, vi } from "vitest"; -// Force a fresh test directory for the database before importing store.js -const testDir = path.join(os.tmpdir(), "agentx-zettel-test-" + Date.now()); -mkdirSync(testDir, { recursive: true }); -process.env.ZETTEL_DIR = testDir; - -import { describe, it, expect } from "vitest"; +vi.hoisted(() => { + process.env.ZETTEL_DIR = `/tmp/agentx-zettel-test-${process.pid}-${Math.random()}`; + delete process.env.TURSO_DATABASE_URL; + delete process.env.TURSO_AUTH_TOKEN; +}); import { writeNote, listNotes, @@ -17,6 +14,9 @@ import { backlinksOf, updateNote, deleteNote, + replaceNoteGraph, + traverseGraphStore, + client, } from "./store.js"; describe("Multi-tenant Notes Database Isolation", () => { @@ -90,8 +90,8 @@ describe("Multi-tenant Notes Database Isolation", () => { const noteA1 = await writeNote(userA, { content: "Note A1" }); const noteA2 = await writeNote(userA, { content: "Note A2" }); - const noteB1 = await writeNote(userB, { content: "Note B1" }); - const noteB2 = await writeNote(userB, { content: "Note B2" }); + await writeNote(userB, { content: "Note B1" }); + await writeNote(userB, { content: "Note B2" }); // Link A1 to A2 await addLink(userA, noteA1.id, noteA2.id); @@ -199,3 +199,125 @@ describe("Multi-tenant Notes Database Isolation", () => { expect(read).not.toBeNull(); }); }); + +describe("tenant graph store", () => { + const userA = "graph-user-A"; + const userB = "graph-user-B"; + + it("stores same-named entities independently for each tenant", async () => { + const noteA = await writeNote(userA, { content: "A graph note" }); + const noteB = await writeNote(userB, { content: "B graph note" }); + + await replaceNoteGraph(userA, noteA.id, { + entities: [{ name: "Mercury", type: "planet", description: "A planet" }], + relations: [], + }); + await replaceNoteGraph(userB, noteB.id, { + entities: [{ name: "Mercury", type: "element", description: "A metal" }], + relations: [], + }); + + const rows = await client.execute( + "SELECT user_id, type, description FROM entities WHERE name = 'Mercury' ORDER BY user_id", + ); + expect(rows.rows).toEqual([ + { user_id: userA, type: "planet", description: "A planet" }, + { user_id: userB, type: "element", description: "A metal" }, + ]); + }); + + it("enforces note ownership and atomically replaces relations", async () => { + const note = await writeNote(userA, { content: "Owned graph note" }); + const firstGraph = { + entities: [ + { name: "A", type: "concept", description: "first" }, + { name: "B", type: "concept", description: "second" }, + ], + relations: [{ source: "A", target: "B", relationship: "leads to" }], + }; + + await expect(replaceNoteGraph(userB, note.id, firstGraph)).rejects.toThrow("Note not found"); + await replaceNoteGraph(userA, note.id, firstGraph); + await replaceNoteGraph(userA, note.id, firstGraph); + expect((await traverseGraphStore(userA, "A", 1)).relations).toHaveLength(1); + + await replaceNoteGraph(userA, note.id, { + entities: [ + { name: "A", type: "concept", description: "updated" }, + { name: "C", type: "concept", description: "third" }, + ], + relations: [{ source: "A", target: "C", relationship: "replaces" }], + }); + expect(await traverseGraphStore(userA, "A", 1)).toEqual({ + entities: ["A", "C"], + relations: [{ source: "A", target: "C", relationship: "replaces", noteId: note.id }], + }); + }); + + it("traverses deterministically with tenant isolation, deduplication, and bounded depth", async () => { + const traversalUser = "graph-traversal-user"; + const noteA1 = await writeNote(traversalUser, { content: "Graph one" }); + const noteA2 = await writeNote(traversalUser, { content: "Graph two" }); + const noteB = await writeNote(userB, { content: "Other tenant graph" }); + const entities = ["A", "B", "C", "D"].map((name) => ({ + name, + type: "concept", + description: name, + })); + + await replaceNoteGraph(traversalUser, noteA1.id, { + entities, + relations: [ + { source: "B", target: "C", relationship: "next" }, + { source: "A", target: "B", relationship: "next" }, + ], + }); + await replaceNoteGraph(traversalUser, noteA2.id, { + entities, + relations: [ + { source: "C", target: "D", relationship: "next" }, + { source: "A", target: "B", relationship: "next" }, + ], + }); + await replaceNoteGraph(userB, noteB.id, { + entities: [...entities, { name: "SECRET", type: "concept", description: "hidden" }], + relations: [{ source: "A", target: "SECRET", relationship: "private" }], + }); + + expect(await traverseGraphStore(traversalUser, "A", 2)).toEqual({ + entities: ["A", "B", "C"], + relations: [ + { source: "A", target: "B", relationship: "next", noteId: noteA1.id }, + { source: "A", target: "B", relationship: "next", noteId: noteA2.id }, + { source: "B", target: "C", relationship: "next", noteId: noteA1.id }, + ], + }); + expect(await traverseGraphStore(traversalUser, "A", 0)).toEqual({ + entities: ["A"], + relations: [], + }); + expect(await traverseGraphStore(traversalUser, "missing", 3)).toEqual({ + entities: [], + relations: [], + }); + await expect(traverseGraphStore(traversalUser, "A", -1)).rejects.toThrow("depth"); + await expect(traverseGraphStore(traversalUser, "A", 101)).rejects.toThrow("depth"); + }); + + it("removes note-owned relations when a note is deleted", async () => { + const note = await writeNote(userA, { content: "Disposable graph" }); + await replaceNoteGraph(userA, note.id, { + entities: [ + { name: "Delete", type: "concept", description: "source" }, + { name: "Me", type: "concept", description: "target" }, + ], + relations: [{ source: "Delete", target: "Me", relationship: "owns" }], + }); + + await deleteNote(userA, note.id); + expect(await traverseGraphStore(userA, "Delete", 2)).toEqual({ + entities: ["Delete"], + relations: [], + }); + }); +}); diff --git a/apps/zettel/src/notes/store.ts b/apps/zettel/src/notes/store.ts index 033be3f..d0f5576 100644 --- a/apps/zettel/src/notes/store.ts +++ b/apps/zettel/src/notes/store.ts @@ -55,6 +55,7 @@ export const client = createClient({ async function initDb(): Promise { try { await client.execute("PRAGMA busy_timeout = 5000"); + await client.execute("PRAGMA foreign_keys = ON"); } catch { // Ignore if not supported (e.g., remote HTTP database) } @@ -179,6 +180,76 @@ async function initDb(): Promise { FOREIGN KEY(user_id) REFERENCES "user"(id) ON DELETE CASCADE ) `); + + await client.execute(` + CREATE UNIQUE INDEX IF NOT EXISTS notes_user_id_id_idx ON notes(user_id, id) + `); + + const legacyRelationColumns = await client.execute("PRAGMA table_info(entity_relations)"); + if ( + legacyRelationColumns.rows.length > 0 && + !legacyRelationColumns.rows.some((row) => row.name === "user_id") + ) { + await client.batch([ + "ALTER TABLE entity_relations RENAME TO legacy_entity_relations", + "ALTER TABLE entities RENAME TO legacy_entities", + `CREATE TABLE entities ( + user_id TEXT NOT NULL, name TEXT NOT NULL, type TEXT NOT NULL, description TEXT NOT NULL, + PRIMARY KEY (user_id, name) + )`, + `CREATE TABLE entity_relations ( + user_id TEXT NOT NULL, note_id TEXT NOT NULL, source TEXT NOT NULL, + target TEXT NOT NULL, relationship TEXT NOT NULL, + PRIMARY KEY (user_id, note_id, source, target, relationship), + FOREIGN KEY (note_id, user_id) REFERENCES notes(id, user_id) ON DELETE CASCADE, + FOREIGN KEY (user_id, source) REFERENCES entities(user_id, name), + FOREIGN KEY (user_id, target) REFERENCES entities(user_id, name) + )`, + `INSERT OR IGNORE INTO entities (user_id, name, type, description) + SELECT DISTINCT n.user_id, endpoints.name, + COALESCE(e.type, 'unknown'), COALESCE(e.description, '') + FROM ( + SELECT note_id, source AS name FROM legacy_entity_relations + UNION SELECT note_id, target AS name FROM legacy_entity_relations + ) endpoints + JOIN notes n ON n.id = endpoints.note_id + LEFT JOIN legacy_entities e ON e.name = endpoints.name`, + `INSERT OR IGNORE INTO entity_relations (user_id, note_id, source, target, relationship) + SELECT n.user_id, r.note_id, r.source, r.target, r.relationship + FROM legacy_entity_relations r JOIN notes n ON n.id = r.note_id`, + "DROP TABLE legacy_entity_relations", + "DROP TABLE legacy_entities", + ]); + } + + await client.execute(` + CREATE TABLE IF NOT EXISTS entities ( + user_id TEXT NOT NULL, + name TEXT NOT NULL, + type TEXT NOT NULL, + description TEXT NOT NULL, + PRIMARY KEY (user_id, name) + ) + `); + + await client.execute(` + CREATE TABLE IF NOT EXISTS entity_relations ( + user_id TEXT NOT NULL, + note_id TEXT NOT NULL, + source TEXT NOT NULL, + target TEXT NOT NULL, + relationship TEXT NOT NULL, + PRIMARY KEY (user_id, note_id, source, target, relationship), + FOREIGN KEY (note_id, user_id) REFERENCES notes(id, user_id) ON DELETE CASCADE, + FOREIGN KEY (user_id, source) REFERENCES entities(user_id, name), + FOREIGN KEY (user_id, target) REFERENCES entities(user_id, name) + ) + `); + + await client.execute(` + CREATE INDEX IF NOT EXISTS entity_relations_traversal_idx + ON entity_relations(user_id, source, target) + `); } // ── Markdown Parser for Migration ────────────────────────────────────────────── @@ -201,6 +272,89 @@ export interface WriteNoteInput { source?: NoteSource; } +export interface GraphEntityInput { + name: string; + type: string; + description: string; +} + +export interface GraphRelationInput { + source: string; + target: string; + relationship: string; +} + +export interface NoteGraphInput { + entities: GraphEntityInput[]; + relations: GraphRelationInput[]; +} + +/** Atomically replace the graph relations owned by one note. */ +export async function replaceNoteGraph( + userId: string, + noteId: string, + graph: NoteGraphInput, +): Promise { + await ensureDb(); + + const ownedNote = await client.execute({ + sql: "SELECT 1 FROM notes WHERE id = ? AND user_id = ?", + args: [noteId, userId], + }); + if (ownedNote.rows.length === 0) throw new Error("Note not found"); + + const entities = [...graph.entities] + .map((entity) => ({ + name: entity.name.trim(), + type: entity.type.trim(), + description: entity.description.trim(), + })) + .sort((a, b) => a.name.localeCompare(b.name)); + const names = new Set(); + for (const entity of entities) { + if (!entity.name || !entity.type) throw new Error("Entity name and type are required"); + if (names.has(entity.name)) throw new Error(`Duplicate entity: ${entity.name}`); + names.add(entity.name); + } + + const relations = [...graph.relations] + .map((relation) => ({ + source: relation.source.trim(), + target: relation.target.trim(), + relationship: relation.relationship.trim(), + })) + .sort( + (a, b) => + a.source.localeCompare(b.source) || + a.target.localeCompare(b.target) || + a.relationship.localeCompare(b.relationship), + ); + for (const relation of relations) { + if (!relation.relationship) throw new Error("Relationship is required"); + if (!names.has(relation.source) || !names.has(relation.target)) { + throw new Error("Every relation endpoint must be included in entities"); + } + } + + await client.batch([ + ...entities.map((entity) => ({ + sql: `INSERT INTO entities (user_id, name, type, description) VALUES (?, ?, ?, ?) + ON CONFLICT(user_id, name) DO UPDATE SET + type = excluded.type, description = excluded.description`, + args: [userId, entity.name, entity.type, entity.description], + })), + { + sql: "DELETE FROM entity_relations WHERE user_id = ? AND note_id = ?", + args: [userId, noteId], + }, + ...relations.map((relation) => ({ + sql: `INSERT OR IGNORE INTO entity_relations + (user_id, note_id, source, target, relationship) VALUES (?, ?, ?, ?, ?)`, + args: [userId, noteId, relation.source, relation.target, relation.relationship], + })), + ]); +} + /** Create a new atomic note in the database. */ export async function writeNote(userId: string, input: WriteNoteInput): Promise { await ensureDb(); @@ -630,6 +784,10 @@ export async function deleteNote(userId: string, id: string): Promise { // Clean up in transaction await client.batch([ + { + sql: "DELETE FROM entity_relations WHERE user_id = ? AND note_id = ?", + args: [userId, id], + }, { sql: "DELETE FROM note_tags WHERE note_id = ?", args: [id], @@ -644,3 +802,67 @@ export async function deleteNote(userId: string, id: string): Promise { }, ]); } + +export interface GraphTraversalResult { + entities: string[]; + relations: Array<{ source: string; target: string; relationship: string; noteId: string }>; +} + +/** Traverse both directions from an entity, with depth limited to 0..100 edges. */ +export async function traverseGraphStore( + userId: string, + entityName: string, + depth = 2, +): Promise { + await ensureDb(); + if (!Number.isInteger(depth) || depth < 0 || depth > 100) { + throw new Error("Graph traversal depth must be an integer between 0 and 100"); + } + + const start = entityName.trim(); + if (!start) return { entities: [], relations: [] }; + const exists = await client.execute({ + sql: "SELECT 1 FROM entities WHERE user_id = ? AND name = ?", + args: [userId, start], + }); + if (exists.rows.length === 0) return { entities: [], relations: [] }; + + const entities = new Set([start]); + const relations = new Map< + string, + { source: string; target: string; relationship: string; noteId: string } + >(); + let frontier = [start]; + + for (let level = 0; level < depth && frontier.length > 0; level += 1) { + const placeholders = frontier.map(() => "?").join(", "); + const result = await client.execute({ + sql: `SELECT source, target, relationship, note_id + FROM entity_relations + WHERE user_id = ? + AND (source IN (${placeholders}) OR target IN (${placeholders})) + ORDER BY source, target, relationship, note_id`, + args: [userId, ...frontier, ...frontier], + }); + const next = new Set(); + for (const row of result.rows) { + const relation = { + source: row.source as string, + target: row.target as string, + relationship: row.relationship as string, + noteId: row.note_id as string, + }; + const key = JSON.stringify(relation); + if (!relations.has(key)) relations.set(key, relation); + for (const name of [relation.source, relation.target]) { + if (!entities.has(name)) { + entities.add(name); + next.add(name); + } + } + } + frontier = [...next].sort((a, b) => a.localeCompare(b)); + } + + return { entities: [...entities], relations: [...relations.values()] }; +}