opencode-episodic-memory 0.2.0 → 0.3.0

package/src/indexer.ts

@@ -1,10 +1,10 @@
 // Incremental, idempotent indexer. Watermark = session.time_updated; a session
 // is re-embedded only when the source changed since we last indexed it.
 import type { Database } from "bun:sqlite";
-import { getTranscriptChecked, listSessions, type SourceSession } from "./reader";
+import { getTranscriptChecked, listSessions, transcriptHasMarker, type SourceSession } from "./reader";
 import { parseTranscript, exchangeText } from "./parser";
 import { embed } from "./embed";
-import { getIndexedSession, replaceSessionChunks } from "./store";
+import type { IndexStore } from "./store";
 
 export interface SyncResult {
   scanned: number;
@@ -17,20 +17,33 @@ export interface SyncResult {
 
 export async function syncSession(
   source: Database,
-  index: Database,
+  index: IndexStore,
   s: SourceSession,
   force = false
 ): Promise<"indexed" | "fresh" | "excluded" | "empty"> {
-  const prior = getIndexedSession(index, s.id);
+  const removeIfRemoteExcluded = async (): Promise<boolean> => {
+    if (!index.remote || !transcriptHasMarker(source, s.id)) return false;
+    await index.removeSession(s.id);
+    return true;
+  };
+  // Remote freshness must never preserve metadata for a newly excluded session.
+  // Local mode intentionally keeps its established cheap freshness-first path.
+  const checked = index.remote ? getTranscriptChecked(source, s.id) : undefined;
+  if (checked?.excluded) {
+    await index.removeSession(s.id);
+    return "excluded";
+  }
+  const prior = await index.getIndexedSession(s.id);
+  if (await removeIfRemoteExcluded()) return "excluded";
   if (!force && prior && prior.source_time_updated >= s.time_updated) return "fresh";
 
   // Authoritative opt-out gate lives inside getTranscriptChecked (raw-blob
   // scan before any read); parseTranscript's own parsed-text check is a
   // harmless redundant fast path for the non-excluded branch.
-  const checked = getTranscriptChecked(source, s.id);
-  const { exchanges, excluded } = checked.excluded
+  const transcript = checked ?? getTranscriptChecked(source, s.id);
+  const { exchanges, excluded } = transcript.excluded
     ? { exchanges: [], excluded: true }
-    : parseTranscript(checked.messages);
+    : parseTranscript(transcript.messages);
   const meta = {
     id: s.id, project_id: s.project_id, parent_id: s.parent_id,
     title: s.title, directory: s.directory,
@@ -38,28 +51,37 @@ export async function syncSession(
   };
 
   if (excluded) {
-    replaceSessionChunks(index, meta, [], "excluded");
+    // A remote index is an opt-in upload boundary: unlike the local index's
+    // useful excluded tombstone, it must retain no metadata for marked chats.
+    if (index.remote) await index.removeSession(s.id);
+    else await index.replaceSessionChunks(meta, [], "excluded");
     return "excluded";
   }
   if (exchanges.length === 0) {
-    replaceSessionChunks(index, meta, [], "empty");
+    if (await removeIfRemoteExcluded()) return "excluded";
+    await index.replaceSessionChunks(meta, [], "empty");
     return "empty";
   }
 
   const date = new Date(s.time_created).toISOString().slice(0, 10);
   const texts = exchanges.map((e) => exchangeText(s.title, date, e));
   const vectors = await embed(texts);
-  replaceSessionChunks(
-    index,
+  // Embedding can take long enough for the source conversation to change. Run
+  // the cheap authoritative raw-marker check again immediately before a remote
+  // upload so a marker added during embedding never exports the prepared data.
+  if (await removeIfRemoteExcluded()) return "excluded";
+  await index.replaceSessionChunks(
     meta,
-    exchanges.map((e, i) => ({ seq: i, time_created: e.time, text: texts[i], embedding: vectors[i] }))
+    exchanges.map((e, i) => ({
+      seq: i, time_created: e.time, text: texts[i], embedding: vectors[i], anchor_message_id: e.anchorMessageId,
+    }))
   );
   return "indexed";
 }
 
 export async function syncAll(
   source: Database,
-  index: Database,
+  index: IndexStore,
   opts: { force?: boolean; onProgress?: (done: number, total: number, title: string) => void } = {}
 ): Promise<SyncResult> {
   const sessions = listSessions(source);
@@ -76,7 +98,7 @@ export async function syncAll(
 
   // Prune index rows whose session no longer exists in the source DB;
   // otherwise their stale (possibly wrong-dims) chunks linger forever.
-  result.pruned = pruneOrphans(source, index, sessions);
+  result.pruned = await pruneOrphans(source, index, sessions);
 
   return result;
 }
@@ -85,15 +107,6 @@ export async function syncAll(
 // source DB. Extracted so the plugin's full-reindex path can call it without
 // re-running the whole sync. Pass already-fetched sessions to avoid a redundant
 // query in syncAll; omitted, it re-reads the source.
-export function pruneOrphans(source: Database, index: Database, knownSource?: SourceSession[]): number {
-  const sourceIds = new Set((knownSource ?? listSessions(source)).map((s) => s.id));
-  const indexedIds = index.prepare<{ id: string }, []>("SELECT id FROM sessions").all();
-  let pruned = 0;
-  for (const { id } of indexedIds) {
-    if (sourceIds.has(id)) continue;
-    index.run("DELETE FROM chunks WHERE session_id = ?", [id]);
-    index.run("DELETE FROM sessions WHERE id = ?", [id]);
-    pruned++;
-  }
-  return pruned;
+export async function pruneOrphans(source: Database, index: IndexStore, knownSource?: SourceSession[]): Promise<number> {
+  return index.pruneOrphans((knownSource ?? listSessions(source)).map((s) => s.id));
 }

package/src/parser.ts

@@ -23,6 +23,7 @@ export function hasExcludeMarker(messages: SourceMessage[]): boolean {
 }
 
 export interface Exchange {
+  anchorMessageId: string;
   user: string;
   assistant: string;
   tools: string[];
@@ -60,7 +61,7 @@ export function parseTranscript(messages: SourceMessage[]): {
     if (m.role === "user") {
       const text = textOf(m.parts);
       if (!text) continue; // e.g. pure tool-result turns
-      current = { user: text, assistant: "", tools: [], time: m.timeCreated };
+      current = { anchorMessageId: m.id, user: text, assistant: "", tools: [], time: m.timeCreated };
       exchanges.push(current);
     } else if (m.role === "assistant" && current) {
       const text = textOf(m.parts);
@@ -73,7 +74,7 @@ export function parseTranscript(messages: SourceMessage[]): {
   return { exchanges: exchanges.filter((e) => e.user || e.assistant), excluded: false };
 }
 
-// Text stored per chunk (also displayed by episodic_read). Capped at 4000 chars
+// Text stored per chunk (also displayed by episodic_read_session). Capped at 4000 chars
 // to keep storage sane; the embedding step (embed.ts) further truncates to 2000
 // chars where retrieval quality peaks. The head of an exchange carries the
 // most signal.

package/src/reader.ts

@@ -42,11 +42,20 @@ const MessageRowSchema = z.object({
   time_created: z.number(),
   data: z.string(),
 });
+type SourceMessageRow = z.infer<typeof MessageRowSchema>;
+
+const AnchorRowSchema = z.object({
+  id: z.string(),
+  time_created: z.number(),
+});
 
 const PartRowSchema = z.object({
   message_id: z.string(),
   data: z.string(),
 });
+type SourcePartRow = z.infer<typeof PartRowSchema>;
+
+const PartCountSchema = z.object({ message_id: z.string(), n: z.number() });
 
 // Aggregate row for the raw marker scan (structural: throw on drift).
 const MarkerCountSchema = z.object({ n: z.number() });
@@ -70,8 +79,12 @@ export interface SourceMessage {
   role: string;
   timeCreated: number;
   parts: SourcePart[];
+  contextPartsOmitted?: number;
 }
 
+const MAX_CONTEXT_PART_BYTES = 8_192;
+const MAX_CONTEXT_PARTS_PER_MESSAGE = 20;
+
 export function openSource(path: string = sourceDbPath()): Database {
   return new Database(path, { readonly: true });
 }
@@ -140,30 +153,90 @@ function getTranscript(db: Database, sessionId: string): SourceMessage[] {
       )
       .all(sessionId)
   );
+  return materializeMessages(db, sessionId, messages);
+}
 
-  const parts = PartRowSchema.array().parse(
-    db
-      .prepare(
-        `SELECT message_id, data FROM part
-         WHERE session_id = ? ORDER BY time_created, id`
-      )
-      .all(sessionId)
-  );
+// Parse parts only for the supplied message rows. Full transcript reads pass
+// every row; bounded context reads pass just their selected SQL window.
+function materializeMessages(
+  db: Database,
+  sessionId: string,
+  messages: SourceMessageRow[],
+  contextLimits?: { maxPartBytes: number; maxPartsPerMessage: number }
+): SourceMessage[] {
+  if (messages.length === 0) return [];
+  const ids = messages.map((message) => message.id);
+  const placeholders = ids.map(() => "?").join(", ");
+  const materialized = contextLimits
+    ? boundedParts(db, sessionId, ids, placeholders, contextLimits)
+    : {
+      rows: PartRowSchema.array().parse(
+        db
+          .prepare(
+            `SELECT message_id, data FROM part
+             WHERE session_id = ? AND message_id IN (${placeholders}) ORDER BY time_created, id`
+          )
+          .all(sessionId, ...ids)
+      ),
+      omittedByMessage: new Map<string, number>(),
+    };
 
   const partsByMsg = new Map<string, SourcePart[]>();
-  for (const p of parts) {
+  for (const p of materialized.rows) {
     const d = PartDataSchema.parse(safeJsonParse(p.data));
     let list = partsByMsg.get(p.message_id);
     if (!list) partsByMsg.set(p.message_id, (list = []));
     list.push(d);
   }
 
-  return messages.map((m) => ({
-    id: m.id,
-    role: MessageDataSchema.parse(safeJsonParse(m.data)).role,
-    timeCreated: m.time_created,
-    parts: partsByMsg.get(m.id) ?? [],
-  }));
+  return messages.map((m) => {
+    const omitted = materialized.omittedByMessage.get(m.id) ?? 0;
+    return {
+      id: m.id,
+      role: MessageDataSchema.parse(safeJsonParse(m.data)).role,
+      timeCreated: m.time_created,
+      parts: partsByMsg.get(m.id) ?? [],
+      ...(omitted > 0 ? { contextPartsOmitted: omitted } : {}),
+    };
+  });
+}
+
+// Context-only part fetch: SQL excludes oversized raw blobs before they cross
+// into JS, ranks remaining parts per selected message, and records omissions.
+// Full transcript reads continue through the unbounded path above.
+function boundedParts(
+  db: Database,
+  sessionId: string,
+  ids: string[],
+  placeholders: string,
+  limits: { maxPartBytes: number; maxPartsPerMessage: number }
+): { rows: SourcePartRow[]; omittedByMessage: Map<string, number> } {
+  const counts = PartCountSchema.array().parse(
+    db.prepare(
+      `SELECT message_id, COUNT(*) AS n FROM part
+       WHERE session_id = ? AND message_id IN (${placeholders}) GROUP BY message_id`
+    ).all(sessionId, ...ids)
+  );
+  const rows = PartRowSchema.array().parse(
+    db.prepare(
+      `WITH ranked AS (
+         SELECT message_id, data, time_created, id,
+                ROW_NUMBER() OVER (PARTITION BY message_id ORDER BY time_created, id) AS part_rank
+         FROM part
+         WHERE session_id = ? AND message_id IN (${placeholders})
+           AND (data IS NULL OR length(CAST(data AS BLOB)) <= ?)
+       )
+       SELECT message_id, data FROM ranked WHERE part_rank <= ? ORDER BY time_created, id`
+    ).all(sessionId, ...ids, limits.maxPartBytes, limits.maxPartsPerMessage)
+  );
+  const retained = new Map<string, number>();
+  for (const row of rows) retained.set(row.message_id, (retained.get(row.message_id) ?? 0) + 1);
+  const omittedByMessage = new Map<string, number>();
+  for (const count of counts) {
+    const omitted = count.n - (retained.get(count.message_id) ?? 0);
+    if (omitted > 0) omittedByMessage.set(count.message_id, omitted);
+  }
+  return { rows, omittedByMessage };
 }
 
 // Discriminated result: excluded conversations never yield a transcript.
@@ -174,9 +247,88 @@ export type CheckedTranscript =
 // The single privacy-gated entry point for reading a transcript. Runs the
 // AUTHORITATIVE raw-blob exclusion check (transcriptHasMarker) BEFORE reading,
 // so the opt-out marker cannot be bypassed by a caller forgetting to check.
-// All production call sites (CLI read, plugin episodic_read, indexer) use this;
+// All production call sites (CLI read, plugin episodic_read_session, indexer) use this;
 // the raw getTranscript is module-internal.
 export function getTranscriptChecked(db: Database, sessionId: string): CheckedTranscript {
   if (transcriptHasMarker(db, sessionId)) return { excluded: true };
   return { excluded: false, messages: getTranscript(db, sessionId) };
 }
+
+export const MAX_CONTEXT_MESSAGES = 20;
+
+export type TranscriptContext =
+  | { ok: true; session: SourceSession; messages: SourceMessage[]; anchorIndex: number; sliceStart: number; total: number }
+  | { ok: false; reason: "unknown_session" | "excluded" | "invalid_anchor" | "invalid_bounds" };
+
+// Read a small, chronological live-source window around an indexed user-message
+// anchor. This deliberately has no indexed fallback: an index may outlive its
+// source transcript, but it cannot safely reconstruct source message context.
+export function getTranscriptContext(
+  db: Database,
+  sessionId: string,
+  anchorMessageId: string,
+  before: number = 3,
+  after: number = 3
+): TranscriptContext {
+  if (!isContextBound(before) || !isContextBound(after)) return { ok: false, reason: "invalid_bounds" };
+  return readSnapshot(db, () => {
+    // Keep the whole-session raw scan first: privacy is session-wide, while
+    // every subsequent query stays bounded to the requested context window.
+    if (transcriptHasMarker(db, sessionId)) return { ok: false, reason: "excluded" };
+    const session = getSession(db, sessionId);
+    if (!session) return { ok: false, reason: "unknown_session" };
+    const anchorRow = db.prepare("SELECT id, time_created FROM message WHERE session_id = ? AND id = ?").get(sessionId, anchorMessageId);
+    if (anchorRow === null || anchorRow === undefined) return { ok: false, reason: "invalid_anchor" };
+    const anchor = AnchorRowSchema.parse(anchorRow);
+    const total = MarkerCountSchema.parse(
+      db.prepare("SELECT COUNT(*) AS n FROM message WHERE session_id = ?").get(sessionId)
+    ).n;
+    const anchorIndex = MarkerCountSchema.parse(
+      db.prepare(
+        `SELECT COUNT(*) AS n FROM message
+         WHERE session_id = ? AND (time_created < ? OR (time_created = ? AND id < ?))`
+      ).get(sessionId, anchor.time_created, anchor.time_created, anchor.id)
+    ).n;
+    const sliceStart = Math.max(0, anchorIndex - before);
+    const sliceEnd = Math.min(total, anchorIndex + after + 1);
+    const sliceLength = sliceEnd - sliceStart;
+    const rows = MessageRowSchema.array().parse(
+      db.prepare(
+        `SELECT id, time_created, data FROM message
+         WHERE session_id = ? ORDER BY time_created, id LIMIT ? OFFSET ?`
+      ).all(sessionId, sliceLength, sliceStart)
+    );
+    return {
+      ok: true,
+      session,
+      messages: materializeMessages(db, sessionId, rows, {
+        maxPartBytes: MAX_CONTEXT_PART_BYTES,
+        maxPartsPerMessage: MAX_CONTEXT_PARTS_PER_MESSAGE,
+      }),
+      anchorIndex,
+      sliceStart,
+      total,
+    };
+  });
+}
+
+// BEGIN is deferred, so this remains a read transaction against the readonly
+// source DB while pinning all context queries to one SQLite snapshot.
+function readSnapshot<T>(db: Database, read: () => T): T {
+  let active = false;
+  try {
+    db.run("BEGIN");
+    active = true;
+    const result = read();
+    db.run("COMMIT");
+    active = false;
+    return result;
+  } catch (error) {
+    if (active) db.run("ROLLBACK");
+    throw error;
+  }
+}
+
+function isContextBound(value: number): boolean {
+  return Number.isInteger(value) && value >= 0 && value <= MAX_CONTEXT_MESSAGES;
+}