stratagate-dsh 0.2.72 → 0.2.73

package/dist/plugin.js

@@ -1905,6 +1905,29 @@ CREATE TABLE IF NOT EXISTS external_memory_import_jobs (
   PRIMARY KEY (namespace, id),
   FOREIGN KEY (namespace) REFERENCES memory_spaces(namespace) ON DELETE CASCADE
 ) STRICT;
+
+CREATE TABLE IF NOT EXISTS raw_message_fts_meta (
+  namespace TEXT NOT NULL,
+  message_id TEXT NOT NULL,
+  tokens TEXT NOT NULL,
+  fts_rowid INTEGER NOT NULL,
+  PRIMARY KEY (namespace, message_id),
+  FOREIGN KEY (namespace, message_id) REFERENCES messages(namespace, id) ON DELETE CASCADE
+) STRICT;
+
+CREATE TABLE IF NOT EXISTS raw_message_fts_state (
+  namespace TEXT PRIMARY KEY,
+  backfill_complete INTEGER NOT NULL DEFAULT 0,
+  FOREIGN KEY (namespace) REFERENCES memory_spaces(namespace) ON DELETE CASCADE
+) STRICT;
+`;
+var RAW_MESSAGE_FTS_SCHEMA = `
+CREATE VIRTUAL TABLE IF NOT EXISTS raw_message_fts USING fts5(
+  namespace UNINDEXED,
+  message_id UNINDEXED,
+  tokens,
+  tokenize = 'unicode61'
+);
 `;
 var THREAD_INDEXES = `
 CREATE INDEX IF NOT EXISTS messages_thread_idx ON messages(namespace, thread_id, position);
@@ -1922,8 +1945,15 @@ function nonEmptyNamespace(namespace) {
   if (!normalized) throw new TypeError("Storage namespace must not be empty");
   return normalized;
 }
+function encodeRawSearchToken(token) {
+  return `t${Buffer.from(token, "utf8").toString("hex")}`;
+}
+function encodedRawSearchTokens(content) {
+  return searchTokens(content).map(encodeRawSearchToken).join(" ");
+}
 var SqliteStorage = class {
   database;
+  rawSearchFtsAvailable = false;
   closed = false;
   constructor(options) {
     if (!options.filename.trim()) throw new TypeError("SQLite filename must not be empty");
@@ -2213,21 +2243,46 @@ var SqliteStorage = class {
     };
     return { snapshot: cloneSnapshot(normalizeSnapshot(snapshot)), revision: space.revision };
   }
-  async save(namespace, snapshot, expectedRevision) {
+  async save(namespace, snapshot, expectedRevision, rawMessageIndexDelta) {
     this.assertOpen();
     assertValidSnapshot(snapshot);
     if (!Number.isSafeInteger(expectedRevision) || expectedRevision < 0) {
       throw new TypeError("expectedRevision must be a non-negative integer");
     }
     const key = nonEmptyNamespace(namespace);
-    return this.immediateTransaction(() => this.persistSnapshot(key, snapshot, expectedRevision));
+    return this.immediateTransaction(() => this.persistSnapshot(key, snapshot, expectedRevision, rawMessageIndexDelta));
   }
   async close() {
     if (this.closed) return;
     this.database.close();
     this.closed = true;
   }
-  persistSnapshot(namespace, snapshot, expectedRevision) {
+  searchRawMessageIds(namespace, tokens, limit, threadId, includeUnthreaded = false) {
+    this.assertOpen();
+    if (!this.rawSearchFtsAvailable || tokens.length === 0) return null;
+    const key = nonEmptyNamespace(namespace);
+    const boundedLimit = Math.max(1, Math.min(5e3, Math.floor(limit)));
+    const match = tokens.map(encodeRawSearchToken).filter(Boolean).map((token) => `"${token}"`).join(" OR ");
+    if (!match) return [];
+    try {
+      const rows = this.database.prepare(`
+        SELECT f.message_id
+        FROM raw_message_fts AS f
+        INNER JOIN messages AS m
+          ON m.namespace = f.namespace AND m.id = f.message_id
+        WHERE raw_message_fts MATCH ?
+          AND f.namespace = ?
+          ${threadId === void 0 ? "" : includeUnthreaded ? "AND (m.thread_id = ? OR m.thread_id IS NULL)" : "AND m.thread_id = ?"}
+        ORDER BY rank
+        LIMIT ?
+      `).all(...threadId === void 0 ? [match, key, boundedLimit] : [match, key, threadId, boundedLimit]);
+      return rows.map(({ message_id }) => message_id);
+    } catch {
+      this.rawSearchFtsAvailable = false;
+      return null;
+    }
+  }
+  persistSnapshot(namespace, snapshot, expectedRevision, rawMessageIndexDelta) {
     const current = this.database.prepare("SELECT revision FROM memory_spaces WHERE namespace = ?").get(namespace);
     const actualRevision = current?.revision ?? null;
     if ((actualRevision ?? 0) !== expectedRevision || actualRevision === null && expectedRevision !== 0) {
@@ -2346,6 +2401,14 @@ var SqliteStorage = class {
     };
     insertMessages(snapshot.openTail, null);
     for (const block of snapshot.blocks) insertMessages(block.l5Raw, block.id);
+    this.syncRawSearchIndex(namespace, rawMessageIndexDelta ?? {
+      upsert: snapshot.blocks.flatMap(({ l5Raw }) => l5Raw),
+      deleteIds: []
+    });
+    if (rawMessageIndexDelta) {
+      const deleteMessage = this.database.prepare("DELETE FROM messages WHERE namespace = ? AND id = ?");
+      for (const messageId of rawMessageIndexDelta.deleteIds) deleteMessage.run(namespace, messageId);
+    }
     const insertEvent = this.database.prepare(`
       INSERT INTO events (
         namespace, id, position, title, summary, narrative, tags_json, quotes_json, source_block_id,
@@ -2620,6 +2683,7 @@ var SqliteStorage = class {
       this.immediateTransaction(() => {
         this.database.exec(SCHEMA);
         this.database.exec(THREAD_INDEXES);
+        this.enableRawSearchFts();
         this.database.exec(`PRAGMA user_version = ${STRATAGATE_STORAGE_SCHEMA_VERSION}`);
       });
     } else if (version >= 1 && version < STRATAGATE_STORAGE_SCHEMA_VERSION) {
@@ -2687,15 +2751,95 @@ var SqliteStorage = class {
           this.database.exec("ALTER TABLE messages ADD COLUMN thread_id TEXT");
         }
         this.database.exec(THREAD_INDEXES);
+        this.enableRawSearchFts();
         this.database.prepare("UPDATE memory_spaces SET schema_version = ? WHERE schema_version < ?").run(STRATAGATE_STORAGE_SCHEMA_VERSION, STRATAGATE_STORAGE_SCHEMA_VERSION);
         this.database.exec(`PRAGMA user_version = ${STRATAGATE_STORAGE_SCHEMA_VERSION}`);
       });
     } else if (version === STRATAGATE_STORAGE_SCHEMA_VERSION) {
       this.database.exec(SCHEMA);
       this.database.exec(THREAD_INDEXES);
+      this.enableRawSearchFts();
     }
     this.assertSchemaVersion();
   }
+  enableRawSearchFts() {
+    try {
+      this.database.exec(RAW_MESSAGE_FTS_SCHEMA);
+      this.rawSearchFtsAvailable = true;
+      const namespaces = this.database.prepare("SELECT namespace FROM memory_spaces").all();
+      for (const { namespace } of namespaces) {
+        const state = this.database.prepare(
+          "SELECT backfill_complete FROM raw_message_fts_state WHERE namespace = ?"
+        ).get(namespace);
+        const messageCount = this.database.prepare(
+          "SELECT COUNT(*) AS count FROM messages WHERE namespace = ? AND block_id IS NOT NULL"
+        ).get(namespace);
+        const indexedCount = this.database.prepare(
+          "SELECT COUNT(*) AS count FROM raw_message_fts_meta WHERE namespace = ?"
+        ).get(namespace);
+        if (state?.backfill_complete === 1 && messageCount.count === indexedCount.count) continue;
+        const messages = state?.backfill_complete === 1 ? this.database.prepare(`
+            SELECT m.id, m.content
+            FROM messages AS m
+            LEFT JOIN raw_message_fts_meta AS i
+              ON i.namespace = m.namespace AND i.message_id = m.id
+            WHERE m.namespace = ? AND m.block_id IS NOT NULL AND i.message_id IS NULL
+          `).all(namespace) : this.database.prepare(
+          "SELECT id, content FROM messages WHERE namespace = ? AND block_id IS NOT NULL"
+        ).all(namespace);
+        this.syncRawSearchIndex(namespace, { upsert: messages, deleteIds: [] });
+        this.database.prepare(`
+          INSERT INTO raw_message_fts_state (namespace, backfill_complete)
+          VALUES (?, 1)
+          ON CONFLICT (namespace) DO UPDATE SET backfill_complete = 1
+        `).run(namespace);
+      }
+    } catch {
+      this.rawSearchFtsAvailable = false;
+    }
+  }
+  syncRawSearchIndex(namespace, delta) {
+    if (!this.rawSearchFtsAvailable) return;
+    const findMeta = this.database.prepare(`
+      SELECT message_id, tokens, fts_rowid
+      FROM raw_message_fts_meta WHERE namespace = ? AND message_id = ?
+    `);
+    const deleteFts = this.database.prepare("DELETE FROM raw_message_fts WHERE rowid = ?");
+    const deleteMeta = this.database.prepare("DELETE FROM raw_message_fts_meta WHERE namespace = ? AND message_id = ?");
+    for (const messageId of delta.deleteIds) {
+      const row = findMeta.get(namespace, messageId);
+      if (!row) continue;
+      deleteFts.run(row.fts_rowid);
+      deleteMeta.run(namespace, messageId);
+    }
+    const insertFts = this.database.prepare(
+      "INSERT INTO raw_message_fts (namespace, message_id, tokens) VALUES (?, ?, ?)"
+    );
+    const insertMeta = this.database.prepare(`
+      INSERT INTO raw_message_fts_meta (namespace, message_id, tokens, fts_rowid)
+      VALUES (?, ?, ?, ?)
+      ON CONFLICT (namespace, message_id) DO UPDATE SET tokens = excluded.tokens, fts_rowid = excluded.fts_rowid
+    `);
+    for (const message of delta.upsert) {
+      const messageId = message.id;
+      const tokens = encodedRawSearchTokens(message.content);
+      const previous = findMeta.get(namespace, messageId);
+      if (previous?.tokens === tokens) continue;
+      if (previous) deleteFts.run(previous.fts_rowid);
+      if (!tokens) {
+        if (previous) deleteMeta.run(namespace, messageId);
+        continue;
+      }
+      insertFts.run(namespace, messageId, tokens);
+      const row = this.database.prepare("SELECT last_insert_rowid() AS rowid").get();
+      insertMeta.run(namespace, messageId, tokens, row.rowid);
+    }
+    this.database.prepare(`
+      INSERT INTO raw_message_fts_state (namespace, backfill_complete)
+      VALUES (?, 1)
+      ON CONFLICT (namespace) DO UPDATE SET backfill_complete = 1
+    `).run(namespace);
+  }
   assertSchemaVersion() {
     const version = this.userVersion();
     if (version !== STRATAGATE_STORAGE_SCHEMA_VERSION) {
@@ -2799,6 +2943,9 @@ var StrataGate = class _StrataGate {
   successfulModelResponses = [];
   ingestionReceipts = /* @__PURE__ */ new Map();
   externalMemoryImportJobs = /* @__PURE__ */ new Map();
+  rawMessageLookup = /* @__PURE__ */ new Map();
+  pendingRawUpserts = /* @__PURE__ */ new Map();
+  pendingRawDeletes = /* @__PURE__ */ new Set();
   currentTurn = 0;
   storage;
   namespace;
@@ -3643,6 +3790,7 @@ var StrataGate = class _StrataGate {
         if (receipt.eventIds.length === 0 && receipt.elementIds.length === 0) this.usageReceipts.delete(receiptId);
       }
       this.blocks.splice(sourceIndex, 1);
+      this.markRawMessagesDeleted(sourceMessageIds);
       return {
         sourceBlockId: id,
         removedEventIds: [...importedEventIds],
@@ -3894,24 +4042,40 @@ var StrataGate = class _StrataGate {
     if (!element) throw new Error(`Unknown element: ${id}`);
     return elementViewAt(element, at);
   }
-  searchRawMemory(query, limit = 6) {
+  searchRawMemory(query, limit = 6, options = {}) {
     const tokens = searchTokens(query);
     if (tokens.length === 0) return [];
-    const hits = [];
-    for (const block of this.blocks) {
-      for (const [index, message] of block.l5Raw.entries()) {
-        const messageTokens = new Set(searchTokens(message.content));
-        if (!tokens.some((token) => messageTokens.has(token))) continue;
-        hits.push({
-          blockId: block.id,
-          turnRange: [block.startTurn, block.endTurn],
-          message,
-          nearby: block.l5Raw.slice(Math.max(0, index - 1), index + 2)
-        });
-        if (hits.length >= limit) return hits;
-      }
+    const boundedLimit = Math.max(1, Math.floor(limit));
+    let candidateIds;
+    try {
+      candidateIds = this.storage?.searchRawMessageIds?.(
+        this.namespace ?? "",
+        tokens,
+        Math.min(5e3, Math.max(100, boundedLimit * 20)),
+        options.threadId,
+        options.includeUnthreaded
+      );
+    } catch {
+      candidateIds = null;
     }
-    return hits;
+    const candidates = candidateIds === null || candidateIds === void 0 ? this.blocks.flatMap((block) => block.l5Raw.map((message, index) => ({
+      id: message.id,
+      block,
+      index,
+      message
+    }))).filter(({ message }) => options.threadId === void 0 || message.threadId === options.threadId || options.includeUnthreaded === true && message.threadId === void 0) : [...new Set(candidateIds)].flatMap((id) => {
+      const item = this.rawMessageLookup.get(id);
+      if (!item) return [];
+      if (options.threadId !== void 0 && item.message.threadId !== options.threadId && !(options.includeUnthreaded === true && item.message.threadId === void 0)) return [];
+      return [{ id, ...item }];
+    });
+    const ranked = bm25Rank(candidates, query, ({ message }) => searchTokens(message.content)).sort((left, right) => right.score - left.score || right.item.message.createdAt.localeCompare(left.item.message.createdAt) || right.item.block.sequence - left.item.block.sequence || right.item.index - left.item.index || left.item.id.localeCompare(right.item.id));
+    return ranked.slice(0, boundedLimit).map(({ item }) => ({
+      blockId: item.block.id,
+      turnRange: [item.block.startTurn, item.block.endTurn],
+      message: item.message,
+      nearby: item.block.l5Raw.slice(Math.max(0, item.index - 1), item.index + 2)
+    }));
   }
   /**
    * Return decayed block views. Passing a threadId limits the result to that
@@ -4083,6 +4247,8 @@ var StrataGate = class _StrataGate {
       lastLiftedBy: null
     };
     this.blocks.push(block);
+    this.indexRawBlock(block);
+    this.markRawMessagesForUpsert(block.l5Raw);
     return block;
   }
   addEventInMemory(input) {
@@ -4247,6 +4413,28 @@ var StrataGate = class _StrataGate {
     this.elementProjectionJobs.set(job.id, job);
     return job;
   }
+  indexRawBlock(block) {
+    for (const [index, message] of block.l5Raw.entries()) {
+      this.rawMessageLookup.set(message.id, { block, index, message });
+    }
+  }
+  markRawMessagesForUpsert(messages) {
+    for (const message of messages) {
+      this.pendingRawDeletes.delete(message.id);
+      this.pendingRawUpserts.set(message.id, structuredClone(message));
+    }
+  }
+  markRawMessagesDeleted(ids) {
+    for (const id of ids) {
+      this.pendingRawUpserts.delete(id);
+      this.pendingRawDeletes.add(id);
+      this.rawMessageLookup.delete(id);
+    }
+  }
+  rebuildRawMessageLookup() {
+    this.rawMessageLookup.clear();
+    for (const block of this.blocks) this.indexRawBlock(block);
+  }
   threadOpenTail(threadId) {
     return this.openTail.filter((message) => message.threadId === threadId);
   }
@@ -4325,6 +4513,8 @@ var StrataGate = class _StrataGate {
       const remaining = this.openTail.filter((message) => !sealedIds.has(message.id));
       this.openTail.splice(0, this.openTail.length, ...remaining);
       this.blocks.push(block);
+      this.indexRawBlock(block);
+      this.markRawMessagesForUpsert(block.l5Raw);
       const updatedAt = toUtc8Iso(this.now());
       this.summaryJobs.set(block.id, {
         blockId: block.id,
@@ -4558,6 +4748,8 @@ var StrataGate = class _StrataGate {
     }
     const before = this.exportSnapshot();
     const beforeRevision = this.revision;
+    const beforeRawUpserts = new Map(this.pendingRawUpserts);
+    const beforeRawDeletes = new Set(this.pendingRawDeletes);
     try {
       const result = await mutation();
       await this.persist();
@@ -4565,6 +4757,10 @@ var StrataGate = class _StrataGate {
     } catch (error) {
       this.restoreSnapshot(before);
       this.revision = beforeRevision;
+      this.pendingRawUpserts.clear();
+      for (const [id, message] of beforeRawUpserts) this.pendingRawUpserts.set(id, message);
+      this.pendingRawDeletes.clear();
+      for (const id of beforeRawDeletes) this.pendingRawDeletes.add(id);
       throw error;
     } finally {
       release();
@@ -4572,7 +4768,13 @@ var StrataGate = class _StrataGate {
   }
   async persist() {
     if (!this.storage || !this.namespace) return;
-    this.revision = await this.storage.save(this.namespace, this.exportSnapshot(), this.revision);
+    const delta = {
+      upsert: [...this.pendingRawUpserts.values()],
+      deleteIds: [...this.pendingRawDeletes]
+    };
+    this.revision = await this.storage.save(this.namespace, this.exportSnapshot(), this.revision, delta);
+    this.pendingRawUpserts.clear();
+    this.pendingRawDeletes.clear();
   }
   restoreSnapshot(snapshot) {
     const normalized = normalizeSnapshot(snapshot);
@@ -4584,6 +4786,9 @@ var StrataGate = class _StrataGate {
     this.currentTurn = copy.currentTurn;
     this.openTail.splice(0, this.openTail.length, ...copy.openTail);
     this.blocks.splice(0, this.blocks.length, ...copy.blocks);
+    this.rebuildRawMessageLookup();
+    this.pendingRawUpserts.clear();
+    this.pendingRawDeletes.clear();
     this.summaryJobs.clear();
     for (const job of copy.summaryJobs) this.summaryJobs.set(job.blockId, job);
     this.events.splice(0, this.events.length, ...copy.events);
@@ -6118,7 +6323,7 @@ var StrataGateRuntime = class {
     await this.flush();
     const memory = await this.space(session);
     const threadId = String(session.id);
-    const results = memory.searchRawMemory(query, scope === "namespace" ? limit : Number.MAX_SAFE_INTEGER).filter((result) => scope === "namespace" || result.message.threadId === threadId || result.message.threadId === void 0).slice(0, limit);
+    const results = memory.searchRawMemory(query, limit, scope === "namespace" ? {} : { threadId, includeUnthreaded: true }).slice(0, limit);
     const blockTitles = new Map(memory.listBlocks().map((block) => [block.id, blockCitationTitle(block)]));
     return this.batch(session, results.map((result, index) => ({
       ref: `raw:${result.blockId}:${result.message.id}:${index}`,