Skip to content

Commit 88bc08c

Browse files
vincentkocvincentkoc
committed
refactor(memory): canonicalize agent database tables
1 parent acdcf37 commit 88bc08c

29 files changed

Lines changed: 522 additions & 229 deletions

‎extensions/memory-core/src/memory/index.test.ts‎

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,13 @@ describe("memory index", () => {
341341
};
342342
}
343343
).db;
344-
for (const table of ["files", "chunks", "embedding_cache", "chunks_fts", "chunks_vec"]) {
344+
for (const table of [
345+
"memory_index_sources",
346+
"memory_index_chunks",
347+
"memory_embedding_cache",
348+
"memory_index_chunks_fts",
349+
"memory_index_chunks_vec",
350+
]) {
345351
const existingTable = db
346352
.prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?")
347353
.get(table);
@@ -452,18 +458,18 @@ describe("memory index", () => {
452458
run: (...params: unknown[]) => void;
453459
};
454460
};
455-
const metaRow = db.prepare("SELECT value FROM meta WHERE key = ?").get("memory_index_meta_v1");
461+
const metaRow = db
462+
.prepare("SELECT value FROM memory_index_meta WHERE key = ?")
463+
.get("memory_index_meta_v1");
456464
const meta = JSON.parse(metaRow?.value ?? "{}") as MemoryIndexMeta;
457-
db.prepare("UPDATE meta SET value = ? WHERE key = ?").run(
465+
db.prepare("UPDATE memory_index_meta SET value = ? WHERE key = ?").run(
458466
JSON.stringify({ ...meta, model, providerKey }),
459467
"memory_index_meta_v1",
460468
);
461-
db.prepare("UPDATE chunks SET model = ?").run(model);
462-
db.prepare("UPDATE embedding_cache SET model = ?, provider_key = ? WHERE provider = ?").run(
463-
model,
464-
providerKey,
465-
identityAliasFixture.provider,
466-
);
469+
db.prepare("UPDATE memory_index_chunks SET model = ?").run(model);
470+
db.prepare(
471+
"UPDATE memory_embedding_cache SET model = ?, provider_key = ? WHERE provider = ?",
472+
).run(model, providerKey, identityAliasFixture.provider);
467473
}
468474

469475
async function expectHybridKeywordSearchFindsMemory(cfg: TestCfg) {
@@ -634,7 +640,7 @@ describe("memory index", () => {
634640
db: { prepare: (sql: string) => { get: (...args: unknown[]) => unknown } };
635641
}
636642
).db
637-
.prepare("SELECT embedding FROM chunks WHERE path LIKE ? AND source = ?")
643+
.prepare("SELECT embedding FROM memory_index_chunks WHERE path LIKE ? AND source = ?")
638644
.get("%2026-01-13.md", "memory") as { embedding: string } | undefined;
639645

640646
expect(betaRow).toBeDefined();
@@ -1009,7 +1015,7 @@ describe("memory index", () => {
10091015
nextManager as unknown as {
10101016
db: { exec: (sql: string) => void };
10111017
}
1012-
).db.exec(`DELETE FROM meta WHERE key = 'memory_index_meta_v1'`);
1018+
).db.exec(`DELETE FROM memory_index_meta WHERE key = 'memory_index_meta_v1'`);
10131019
expect(nextManager.status().custom?.indexIdentity).toEqual({
10141020
status: "missing",
10151021
reason: "index metadata is missing",
@@ -1086,7 +1092,7 @@ describe("memory index", () => {
10861092
};
10871093
}
10881094
).db;
1089-
db.exec(`DELETE FROM meta WHERE key = 'memory_index_meta_v1'`);
1095+
db.exec(`DELETE FROM memory_index_meta WHERE key = 'memory_index_meta_v1'`);
10901096

10911097
await nextManager.sync({ reason: "test" });
10921098

@@ -1095,7 +1101,7 @@ describe("memory index", () => {
10951101
status: "missing",
10961102
reason: "index metadata is missing",
10971103
});
1098-
const row = db.prepare("SELECT model FROM chunks LIMIT 1").get();
1104+
const row = db.prepare("SELECT model FROM memory_index_chunks LIMIT 1").get();
10991105
expect(row?.model).toBe("semantic-embed");
11001106
} finally {
11011107
await nextManager.close?.();
@@ -1197,7 +1203,7 @@ describe("memory index", () => {
11971203
nextManager as unknown as {
11981204
db: { exec: (sql: string) => void };
11991205
}
1200-
).db.exec(`DELETE FROM meta WHERE key = 'memory_index_meta_v1'`);
1206+
).db.exec(`DELETE FROM memory_index_meta WHERE key = 'memory_index_meta_v1'`);
12011207

12021208
const status = nextManager.status();
12031209

@@ -1627,7 +1633,10 @@ describe("memory index", () => {
16271633
const originalPrepare = db.prepare.bind(db);
16281634
let ftsSelects = 0;
16291635
const prepareSpy = vi.spyOn(db, "prepare").mockImplementation((sql: string) => {
1630-
if (sql.includes("FROM chunks_fts") && sql.includes("WHERE chunks_fts MATCH ?")) {
1636+
if (
1637+
sql.includes("FROM memory_index_chunks_fts") &&
1638+
sql.includes("WHERE memory_index_chunks_fts MATCH ?")
1639+
) {
16311640
ftsSelects += 1;
16321641
}
16331642
return originalPrepare(sql);

‎extensions/memory-core/src/memory/manager-db.test.ts‎

Lines changed: 59 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,10 @@ import {
1515
} from "./manager-db.js";
1616
import { acquireMemoryReindexLock } from "./manager-reindex-lock.js";
1717

18-
function ensureTestMemorySchema(db: DatabaseSync): void {
18+
function ensureTestMemorySchema(db: DatabaseSync, cacheEnabled = true): void {
1919
ensureMemoryIndexSchema({
2020
db,
21-
embeddingCacheTable: "embedding_cache",
22-
cacheEnabled: true,
23-
ftsTable: "chunks_fts",
21+
cacheEnabled,
2422
ftsEnabled: false,
2523
});
2624
}
@@ -48,8 +46,10 @@ describe("memory manager database publication", () => {
4846
try {
4947
ensureTestMemorySchema(targetDb);
5048
ensureTestMemorySchema(sourceDb);
51-
targetDb.exec("CREATE TABLE chunks_vec (id TEXT PRIMARY KEY, embedding BLOB)");
52-
targetDb.prepare("INSERT INTO chunks_vec (id, embedding) VALUES (?, ?)").run("stale", "[]");
49+
targetDb.exec("CREATE TABLE memory_index_chunks_vec (id TEXT PRIMARY KEY, embedding BLOB)");
50+
targetDb
51+
.prepare("INSERT INTO memory_index_chunks_vec (id, embedding) VALUES (?, ?)")
52+
.run("stale", "[]");
5353
sourceDb.close();
5454

5555
await publishMemoryDatabaseTables({
@@ -61,7 +61,9 @@ describe("memory manager database publication", () => {
6161

6262
expect(
6363
targetDb
64-
.prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'chunks_vec'")
64+
.prepare(
65+
"SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'memory_index_chunks_vec'",
66+
)
6567
.get(),
6668
).toBeUndefined();
6769
} finally {
@@ -85,13 +87,13 @@ describe("memory manager database publication", () => {
8587
return;
8688
}
8789
sourceDb.exec(`
88-
CREATE VIRTUAL TABLE chunks_vec USING vec0(
90+
CREATE VIRTUAL TABLE memory_index_chunks_vec USING vec0(
8991
id TEXT PRIMARY KEY,
9092
embedding FLOAT[3]
9193
)
9294
`);
9395
sourceDb
94-
.prepare("INSERT INTO chunks_vec (id, embedding) VALUES (?, ?)")
96+
.prepare("INSERT INTO memory_index_chunks_vec (id, embedding) VALUES (?, ?)")
9597
.run("vector", JSON.stringify([0, 1, 0]));
9698
sourceDb.close();
9799

@@ -103,7 +105,9 @@ describe("memory manager database publication", () => {
103105
vectorExtensionPath: sourceVector.extensionPath,
104106
});
105107

106-
expect(targetDb.prepare("SELECT id FROM chunks_vec").all()).toEqual([{ id: "vector" }]);
108+
expect(targetDb.prepare("SELECT id FROM memory_index_chunks_vec").all()).toEqual([
109+
{ id: "vector" },
110+
]);
107111
} finally {
108112
try {
109113
sourceDb.close();
@@ -122,16 +126,22 @@ describe("memory manager database publication", () => {
122126
ensureTestMemorySchema(targetDb);
123127
ensureTestMemorySchema(sourceDb);
124128
targetDb
125-
.prepare("INSERT INTO files (path, source, hash, mtime, size) VALUES (?, ?, ?, ?, ?)")
129+
.prepare(
130+
"INSERT INTO memory_index_sources (path, source, hash, mtime, size) VALUES (?, ?, ?, ?, ?)",
131+
)
126132
.run("memory.md", "memory", "published", 1, 1);
127133
sourceDb
128-
.prepare("INSERT INTO files (path, source, hash, mtime, size) VALUES (?, ?, ?, ?, ?)")
134+
.prepare(
135+
"INSERT INTO memory_index_sources (path, source, hash, mtime, size) VALUES (?, ?, ?, ?, ?)",
136+
)
129137
.run("memory.md", "memory", "shadow", 1, 1);
130138
const expectedRevision = readMemoryDatabaseRevision(targetDb);
131139
sourceDb.close();
132140

133141
concurrentDb = new DatabaseSync(targetPath);
134-
concurrentDb.prepare("UPDATE files SET hash = ? WHERE path = ?").run("newer", "memory.md");
142+
concurrentDb
143+
.prepare("UPDATE memory_index_sources SET hash = ? WHERE path = ?")
144+
.run("newer", "memory.md");
135145
concurrentDb.close();
136146
concurrentDb = undefined;
137147

@@ -144,7 +154,7 @@ describe("memory manager database publication", () => {
144154
}),
145155
).rejects.toThrow(/changed while full reindex was building/);
146156
expect(
147-
targetDb.prepare("SELECT hash FROM files WHERE path = ?").get("memory.md"),
157+
targetDb.prepare("SELECT hash FROM memory_index_sources WHERE path = ?").get("memory.md"),
148158
).toEqual({ hash: "newer" });
149159
} finally {
150160
try {
@@ -157,6 +167,41 @@ describe("memory manager database publication", () => {
157167
}
158168
});
159169

170+
it("preserves the live embedding cache when the shadow index has caching disabled", async () => {
171+
const targetPath = path.join(fixtureRoot, "target.sqlite");
172+
const sourcePath = path.join(fixtureRoot, "source.sqlite");
173+
const targetDb = new DatabaseSync(targetPath);
174+
const sourceDb = new DatabaseSync(sourcePath);
175+
try {
176+
ensureTestMemorySchema(targetDb);
177+
ensureTestMemorySchema(sourceDb, false);
178+
targetDb
179+
.prepare(
180+
`INSERT INTO memory_embedding_cache (
181+
provider, model, provider_key, hash, embedding, dims, updated_at
182+
) VALUES (?, ?, ?, ?, ?, ?, ?)`,
183+
)
184+
.run("test", "model", "key", "hash", "[]", 0, 1);
185+
sourceDb.close();
186+
187+
await publishMemoryDatabaseTables({
188+
targetDb,
189+
sourcePath,
190+
metaKey: "memory_index_meta",
191+
expectedRevision: readMemoryDatabaseRevision(targetDb),
192+
});
193+
194+
expect(targetDb.prepare("SELECT hash FROM memory_embedding_cache").all()).toEqual([
195+
{ hash: "hash" },
196+
]);
197+
} finally {
198+
try {
199+
sourceDb.close();
200+
} catch {}
201+
targetDb.close();
202+
}
203+
});
204+
160205
it("removes aged orphan shadows but preserves young and locked shadows", async () => {
161206
const databasePath = path.join(fixtureRoot, "agent.sqlite");
162207
const database = new DatabaseSync(databasePath);

‎extensions/memory-core/src/memory/manager-db.ts‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ export function readMemoryDatabaseRevision(db: DatabaseSync): number {
9191

9292
function replaceVirtualTable(params: {
9393
db: DatabaseSync;
94-
tableName: "chunks_fts" | "chunks_vec";
94+
tableName: "memory_index_chunks_fts" | "memory_index_chunks_vec";
9595
columns: string;
9696
ignoreDropErrorWhenSourceMissing?: boolean;
9797
}): void {
@@ -126,7 +126,7 @@ export async function publishMemoryDatabaseTables(params: {
126126
params.targetDb.prepare(`ATTACH DATABASE ? AS ${MEMORY_REINDEX_SCHEMA}`).run(params.sourcePath);
127127
try {
128128
if (
129-
tableExists(params.targetDb, MEMORY_REINDEX_SCHEMA, "chunks_vec") &&
129+
tableExists(params.targetDb, MEMORY_REINDEX_SCHEMA, "memory_index_chunks_vec") &&
130130
!hasSqliteVecExtension(params.targetDb)
131131
) {
132132
const loaded = await loadSqliteVecExtension({
@@ -148,47 +148,49 @@ export async function publishMemoryDatabaseTables(params: {
148148
`(expected revision ${params.expectedRevision}, found ${liveRevision}); retry the full reindex.`,
149149
);
150150
}
151-
params.targetDb.prepare("DELETE FROM main.meta WHERE key = ?").run(params.metaKey);
151+
params.targetDb
152+
.prepare("DELETE FROM main.memory_index_meta WHERE key = ?")
153+
.run(params.metaKey);
152154
params.targetDb
153155
.prepare(
154-
`INSERT INTO main.meta (key, value)
155-
SELECT key, value FROM ${MEMORY_REINDEX_SCHEMA}.meta WHERE key = ?`,
156+
`INSERT INTO main.memory_index_meta (key, value)
157+
SELECT key, value FROM ${MEMORY_REINDEX_SCHEMA}.memory_index_meta WHERE key = ?`,
156158
)
157159
.run(params.metaKey);
158160

159161
params.targetDb.exec(`
160-
DELETE FROM main.files;
161-
INSERT INTO main.files (path, source, hash, mtime, size)
162-
SELECT path, source, hash, mtime, size FROM ${MEMORY_REINDEX_SCHEMA}.files;
162+
DELETE FROM main.memory_index_sources;
163+
INSERT INTO main.memory_index_sources (path, source, hash, mtime, size)
164+
SELECT path, source, hash, mtime, size FROM ${MEMORY_REINDEX_SCHEMA}.memory_index_sources;
163165
164-
DELETE FROM main.chunks;
165-
INSERT INTO main.chunks (
166+
DELETE FROM main.memory_index_chunks;
167+
INSERT INTO main.memory_index_chunks (
166168
id, path, source, start_line, end_line, hash, model, text, embedding, updated_at
167169
)
168170
SELECT
169171
id, path, source, start_line, end_line, hash, model, text, embedding, updated_at
170-
FROM ${MEMORY_REINDEX_SCHEMA}.chunks;
172+
FROM ${MEMORY_REINDEX_SCHEMA}.memory_index_chunks;
171173
`);
172174

173-
if (tableExists(params.targetDb, MEMORY_REINDEX_SCHEMA, "embedding_cache")) {
175+
if (tableExists(params.targetDb, MEMORY_REINDEX_SCHEMA, "memory_embedding_cache")) {
174176
params.targetDb.exec(`
175-
DELETE FROM main.embedding_cache;
176-
INSERT INTO main.embedding_cache (
177+
DELETE FROM main.memory_embedding_cache;
178+
INSERT INTO main.memory_embedding_cache (
177179
provider, model, provider_key, hash, embedding, dims, updated_at
178180
)
179181
SELECT provider, model, provider_key, hash, embedding, dims, updated_at
180-
FROM ${MEMORY_REINDEX_SCHEMA}.embedding_cache;
182+
FROM ${MEMORY_REINDEX_SCHEMA}.memory_embedding_cache;
181183
`);
182184
}
183185

184186
replaceVirtualTable({
185187
db: params.targetDb,
186-
tableName: "chunks_fts",
188+
tableName: "memory_index_chunks_fts",
187189
columns: "text, id, path, source, model, start_line, end_line",
188190
});
189191
replaceVirtualTable({
190192
db: params.targetDb,
191-
tableName: "chunks_vec",
193+
tableName: "memory_index_chunks_vec",
192194
columns: "id, embedding",
193195
// A vector-disabled connection may not have sqlite-vec loaded and cannot
194196
// drop an old virtual table. Missing vector metadata forces a strict

‎extensions/memory-core/src/memory/manager-embedding-cache.test.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,7 @@ describe("memory embedding cache", () => {
1717
const db = new DatabaseSync(":memory:");
1818
ensureMemoryIndexSchema({
1919
db,
20-
embeddingCacheTable: "embedding_cache",
2120
cacheEnabled: true,
22-
ftsTable: "chunks_fts",
2321
ftsEnabled: false,
2422
ftsTokenizer: "unicode61",
2523
});

‎extensions/memory-core/src/memory/manager-embedding-cache.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ export function loadMemoryEmbeddingCache(params: {
3636
return new Map();
3737
}
3838

39-
const tableName = params.tableName ?? "embedding_cache";
39+
const tableName = params.tableName ?? "memory_embedding_cache";
4040
const out = new Map<string, number[]>();
4141
const batchSize = 400;
4242
for (const identity of params.providerIdentities) {
@@ -73,7 +73,7 @@ export function upsertMemoryEmbeddingCache(params: {
7373
if (!params.enabled || !provider || !params.providerKey || params.entries.length === 0) {
7474
return;
7575
}
76-
const tableName = params.tableName ?? "embedding_cache";
76+
const tableName = params.tableName ?? "memory_embedding_cache";
7777
const now = params.now ?? Date.now();
7878
const stmt = params.db.prepare(
7979
`INSERT INTO ${tableName} (provider, model, provider_key, hash, embedding, dims, updated_at)\n` +

0 commit comments

Comments
 (0)