pi-optchat 0.10.2 → 0.11.1

package/src/durable/optchat.ts

@@ -0,0 +1,274 @@
+import {
+  AssistantEntry, CompactionTask, defineExtension, defineTask, defineTool, GenerationTask, hook, InboxDoc, LiveDoc, ResetEntry, ROOT_CONVERSATION_ID,
+  section, ToolResultEntry, UserEntry, type Conversation, type ConversationId, type EntryId, type EntryRecord, type GenerationHooks, type Harness,
+  type HarnessOptions, type ToolExecutionResult,
+} from '@earendil-works/pi-durable';
+import type { Static, TSchema } from 'typebox';
+import {
+  allowSearch, boundedMessage, cachePayload, completedExchange, CONTINUITY, decode, isImage, isMime, logMessage, memoryTools, refreshBudget, refreshedRun,
+  refreshPoint, SEARCH_DOC, searchTool, textContent, TURN_REFRESHED, VIEW_DOC, withView, type AnyMessage, type ImageContent, type Memory,
+} from '../core/index.ts';
+
+export type Context = Parameters<Harness['commit']>[1];
+type Message = Parameters<GenerationHooks['beforeRequest']>[0]['messages'][number];
+type Models = Exclude<HarnessOptions['models'], () => unknown>;
+
+/** The entries that are the chat: other kinds (system prompt changes, resets, compactions, other extensions' own) stay out of memory. */
+const CHAT = new Set([UserEntry.kind, AssistantEntry.kind, ToolResultEntry.kind]);
+const chat = (entry: EntryRecord) => CHAT.has(entry.kind) ? entry.model ?? [] : [];
+/** What memory logs of an entry: its chat, or a reset's note (`Conversation.reset(text)`, a tool's handoff) marked as one. Pi sends
+ * the note in the context it starts, which is logged only once it ends, so the model never sees it twice. */
+const remembered = (entry: EntryRecord) => !ResetEntry.is(entry) ? chat(entry)
+  : (entry.model ?? []).map(note => ({ role: 'user' as const, content: `[Handoff note] ${textContent(note.content)}`, timestamp: note.timestamp }));
+/** Each message OptChat logs from a Pi entry carries the entry's id, so the log itself says how far it has read Pi's transcript. */
+const RECEIPT = 'pi:';
+/** The answer that ends a run, whether it succeeded, failed or was aborted. */
+const endsRun = (m: Message) => m.role === 'assistant' && m.stopReason !== 'toolUse';
+const encoder = new TextEncoder();
+
+/** The images memory names, as Pi keeps them: a user's and a tool's, but not zoom's, which are kept already. */
+const keptImages = (message: AnyMessage) => message.role === 'user' || message.role === 'toolResult' && !('toolName' in message && message.toolName === 'zoom')
+  ? Array.isArray(message.content) ? message.content.filter(isImage) : [] : [];
+/** Each distinct image by its base64, named as Pi's OptChat names it: the first 16 hex digits of the SHA-256 of its base64. Web
+ * Crypto is asynchronous and core's ImageRef is not, so images are named before core reads the names. */
+async function imageNames(images: readonly ImageContent[]) {
+  const distinct = new Map(images.map(image => [image.data, image]));
+  return new Map(await Promise.all([...distinct.values()].map(async image => {
+    const hash = new Uint8Array(await crypto.subtle.digest('SHA-256', encoder.encode(image.data)), 0, 8);
+    return [image.data, { image, name: Array.from(hash, byte => byte.toString(16).padStart(2, '0')).join('') }] as const;
+  })));
+}
+
+/** One of core's memory tools as a Pi Durable tool, its arguments typed by its own schema. */
+const durableTool = <P extends TSchema>(tool: {
+  name: string; description: string; parameters: P;
+  execute: (callId: string, args: Static<P>) => Promise<{ content: ToolExecutionResult['output'] }>;
+}) => defineTool({
+  name: tool.name, description: tool.description, parameters: tool.parameters, replay: 'safe',
+  execute: async (args, api) => ({ output: (await tool.execute(api.callId, args)).content }),
+});
+
+export interface OptChatOptions {
+  memory: Memory;
+  /** The conversation OptChat remembers; Pi Durable's root by default. Other conversations that select the extension are left alone. */
+  conversation?: ConversationId;
+  /** The memory part of the system prompt: VIEW_DOC (plus CONTINUITY and SEARCH_DOC as set) by default. */
+  prompt?: string;
+  /** Offers the `search` tool. Off by default, as in Pi. */
+  search?: boolean;
+  /** The previous exchange is sent in full after the view when it is at most this many bytes; 0 never sends it. */
+  previousExchangeBytes?: number;
+  /** When a run's messages since the last refresh pass this share of the model's window, they move into memory and the run goes
+   * on with a fresh view, as Pi's Turn refresh threshold (#146). 0.3 by default; 0 is off. */
+  refreshAt?: number;
+}
+
+/** OptChat's memory for one Pi Durable conversation. Install `extension` in the registry Harness.open gets, then `attach` the opened
+ * Harness. Each request is built as OptChat builds it in Pi: the view, the previous exchange and the current context, so the model
+ * never sees Pi's older transcript. When the conversation goes idle, OptChat resets Pi's context and a background task logs what the
+ * reset ended into memory and builds its summaries; Pi Durable runs that task, and wakes the Durable Object for it after an eviction.
+ * Turn Pi's own compaction off (`settings.compaction.enabled: false`); the extension also declines any that starts. */
+export function optchat({ memory, conversation: id = ROOT_CONVERSATION_ID, prompt, search = false, previousExchangeBytes = 16_000, refreshAt = 0.3 }: OptChatOptions) {
+  const attached = attachment('OptChat'), opened = attached.get;
+  /** The Pi entry logging resumes after: messages before the context's head are logged, the rest are sent as they are. On open it is
+   * the one before the newest entry in memory, which is read again: a restart may have cut it off between its records (an answer
+   * with tool calls makes several), so the first `saved` records it makes, which memory has, are skipped. The newest is the highest
+   * id, not the last in the log: a rebuild orders turns by their dates, which a clock can put out of order. */
+  let logged = 0, saved = 0;
+  const newest = memory.root.reduce((max, e) => e.receipt?.startsWith(RECEIPT) ? Math.max(max, Number(e.receipt.slice(RECEIPT.length))) : max, 0);
+  if (newest) { logged = newest - 1; saved = memory.root.filter(e => e.receipt === `${RECEIPT}${newest}`).length; }
+  const resumed = new Proxy(memory, {
+    get: (target, key) => key !== 'append' ? Reflect.get(target, key) : (...args: Parameters<Memory['append']>) => saved > 0 ? void saved-- : target.append(...args),
+  });
+  let logging = Promise.resolve();
+  /** What a context's requests share, fixed when its first one is built so they read the same cached prefix. */
+  let turn: { head: EntryId | undefined; view: string; previous: Message[] } | undefined;
+
+  const handle = async (ctx: Context): Promise<Conversation> => {
+    const conversation = await opened().harness.conversation(id, ctx);
+    if (!conversation) throw new Error(`No conversation ${id}.`);
+    return conversation;
+  };
+  /** Logs Pi's entries before `head` that memory doesn't have yet, one caller at a time. A call that fails (aborted, say) leaves
+   * the next one to carry on from the last entry it logged. */
+  const logUpTo = (conversation: Conversation, head: EntryId, ctx: Context) => {
+    const run = logging.then(async () => {
+      let cursor: Parameters<Conversation['entries']>[2];
+      do {
+        if (logged + 1 > head - 1) return;
+        const page = await conversation.entries({ order: 'ascending', minEntryId: logged + 1 as EntryId, maxEntryId: head - 1 as EntryId }, 256, cursor, ctx);
+        for (const entry of page.items) {
+          const messages = remembered(entry), names = await imageNames(messages.flatMap(keptImages));
+          // Kept before the message is logged, so the image its text names is there for zoom; one that can't be kept is only named.
+          // An image sent again, or read again after a restart, is written again under its name.
+          for (const { image, name } of names.values()) if (isMime(image.mimeType)) {
+            try { await memory.store.putImage(name, image.mimeType, decode(image.data)); }
+            catch (error) { console.error(`OptChat could not keep an image: ${error instanceof Error ? error.message : String(error)}`); }
+          }
+          // Once a save failed (also one queued before, failing during the awaits above), memory refuses appends: requests keep the
+          // view it has, and a restart logs the rest from Pi's transcript.
+          if (memory.saveError) return;
+          for (const message of messages) logMessage(resumed, message, { receipt: `${RECEIPT}${entry.id}`, imageRef: image => `[image ${names.get(image.data)?.name}]` });
+          logged = entry.id;
+        }
+        cursor = page.next;
+      } while (cursor);
+    });
+    logging = run.catch(() => {});
+    return run;
+  };
+  /** The last run before `head`, from Pi's transcript, if it was answered and is short enough to send whole. It is read from the
+   * entries, not the context's messages, which leave failed and aborted answers out and so would join a run to the one before.
+   * A text-only answer ends a run here even when steering followed it in the same Pi run. */
+  const previousBefore = async (conversation: Conversation, head: EntryId, ctx: Context) => {
+    if (!previousExchangeBytes) return [];
+    const [last] = (await conversation.entries({ order: 'descending', maxEntryId: head - 1 as EntryId }, 1, undefined, ctx)).items;
+    if (!last) return [];
+    const messages = (await conversation.context(ctx, { at: last.id })).entries.flatMap(chat);
+    const start = messages.slice(0, messages.findLastIndex(m => m.role === 'user')).findLastIndex(endsRun) + 1;
+    const run = messages.slice(start), names = await imageNames(run.filter(m => m.role === 'user').flatMap(keptImages));
+    const exchange = completedExchange(run, image => `[image ${names.get(image.data)?.name}]`);
+    const bytes = exchange.reduce((sum, m) => sum + encoder.encode(textContent(m.content)).length, 0);
+    return bytes <= previousExchangeBytes ? exchange : [];
+  };
+  const beforeRequest: GenerationHooks['beforeRequest'] = async ({ messages }, api, ctx) => {
+    if (api.conversationId !== id) return undefined;
+    const conversation = await handle(ctx);
+    const { head, entries } = await conversation.context(ctx);
+    if (!turn || turn.head !== head?.id) {
+      if (head) await logUpTo(conversation, head.id, ctx);
+      await memory.settle(ctx.abortSignal);
+      turn = { head: head?.id, view: memory.render(), previous: head ? await previousBefore(conversation, head.id, ctx) : [] };
+    }
+    const at = messages.findIndex(m => m.role !== 'system');
+    if (at < 0) return undefined;
+    // A tool result over 30K characters is sent as its head and tail, as in Pi.
+    const run = messages.slice(at).map(boundedMessage), shared = turn;
+    // `run` is the whole context, which can hold more than one run: inputs queued with no idle reset between them, or a reset's note
+    // first. The run being answered starts at its task, the last of the user messages after the newest finished answer; a refresh
+    // moves what comes before it into memory with the steps it cuts, and a steer before the cut too, as in Pi.
+    let task = run.findLastIndex(endsRun) + 1;
+    while (run[task + 1]?.role === 'user') task++;
+    const own = run.slice(task);
+    /** The run's answers and their entries. Pi Durable's messages are copies of its entries', so an answer is found by its first
+     * tool call's id, or its timestamp without one. */
+    const answer = (m: Message) => m.role === 'assistant' ? m.content.find(c => c.type === 'toolCall')?.id ?? m.timestamp : undefined;
+    const steps = entries.flatMap(entry => {
+      const sent = AssistantEntry.is(entry) && entry.model?.[0];
+      const i = sent ? own.findIndex(m => answer(m) === answer(sent)) : -1;
+      return i > 0 ? [{ entry, i }] : [];
+    });
+    /** Logs Pi's entries before `end`, rebuilding the view if that logged anything. False once a save failed: memory then lacks
+     * part of the run, so it is sent whole. */
+    const logThrough = async (end: EntryId) => {
+      const count = memory.root.length;
+      await logUpTo(conversation, end, ctx);
+      if (memory.saveError) return false;
+      if (memory.root.length !== count) { await memory.settle(ctx.abortSignal); shared.view = memory.render(); }
+      return true;
+    };
+    // Once a refresh has logged part of this context, an earlier run in it that has ended goes into memory whole, as the reset it
+    // missed would have done: a run queued behind a long one doesn't send it again.
+    const ended = task > 0 && logged > (head?.id ?? 0) && run[task - 1].role === 'assistant'
+      ? entries.find(entry => AssistantEntry.is(entry) && entry.model?.[0] && answer(entry.model[0]) === answer(run[task - 1])) : undefined;
+    const whole = ended && await logThrough(ended.id + 1 as EntryId) ? own : run;
+    // A refresh logs the run up to the answer it cuts at, so once any of its answers is in the log, the run is sent from the first
+    // answer after it, after a restart too; a restart in the middle of that logging leaves the rest of it to do.
+    let keep = 0;
+    const kept = steps.some(s => s.entry.id <= logged) ? steps.find(s => s.entry.id > logged) : undefined;
+    if (kept && await logThrough(kept.entry.id)) keep = kept.i;
+    const { model } = await conversation.agent(ctx);
+    const refresh = !memory.saveError && refreshPoint(own, keep, refreshBudget(model && api.models.getModel(model.provider, model.modelId)?.contextWindow, refreshAt));
+    const cut = refresh && steps.find(s => s.i === refresh.cut);
+    if (cut && await logThrough(cut.entry.id)) keep = cut.i;
+    const sent = keep ? refreshedRun(own, keep, TURN_REFRESHED) : whole;
+    return { messages: withView([...messages.slice(0, at), ...turn.previous, ...sent], turn.view) };
+  };
+
+  const Remember = defineTask<Record<string, never>, { phase: 'remember' }, null>({
+    name: 'optchat.remember', version: 1,
+    initial: () => ({ phase: 'remember' }),
+    phases: {
+      // Logs what the newest reset ended and waits until the next turn's view is built, or every summary owed has failed. A phase
+      // cut off by an eviction runs again from the top: logging skips what is in memory, and memory rebuilds what it owes on open.
+      remember: async (_task, runtime, ctx) => {
+        const conversation = await handle(ctx), { head } = await conversation.context(ctx);
+        if (head) await logUpTo(conversation, head.id, ctx);
+        await memory.settle(runtime.signal);
+        await runtime.commit(() => ({ status: 'terminal', outcome: { status: 'completed', result: null } }), ctx);
+      },
+    },
+    abort: (_task, runtime, ctx) => runtime.commit(() => ({ status: 'terminal', outcome: { status: 'aborted' } }), ctx),
+  });
+
+  /** Once the conversation is idle with a chat in its context, resets the context and starts remembering it, in one commit. */
+  const resetIfIdle = async () => {
+    const { harness, context } = opened();
+    const busy = async () => !!(await harness.snapshot(LiveDoc, id, context))?.run || !!(await harness.snapshot(InboxDoc, id, context))?.items.length;
+    if (await busy() || !(await (await handle(context)).context(context)).entries.some(e => chat(e).length)) return;
+    await harness.commit(async tx => {
+      // Checked again where it is atomic: an input that came in meanwhile keeps its context.
+      if ((await tx.doc(LiveDoc, id)).run || (await tx.doc(InboxDoc, id)).items.length) return;
+      await tx.appendEntry(id, { kind: ResetEntry.kind, head: 'self' });
+      await tx.createTask(Remember, {}, { ownership: { kind: 'conversation' }, conversationId: id, background: true });
+    }, context);
+  };
+  /** Checks after commits, one check at a time; commits during a check ask for one more, however many there were. */
+  let checking = false, again = false;
+  const schedule = async () => {
+    again = true;
+    if (checking) return;
+    checking = true;
+    while (again) {
+      again = false;
+      await resetIfIdle().catch(error => console.error(`OptChat could not reset the context: ${error instanceof Error ? error.message : String(error)}`));
+    }
+    checking = false;
+  };
+
+  const [zoom, date] = memoryTools(() => memory);
+  const tools = [durableTool(zoom), durableTool(date), ...search ? [durableTool(searchTool(() => memory))] : []];
+  const text = prompt ?? allowSearch(`${VIEW_DOC}${previousExchangeBytes ? CONTINUITY : ''}${search ? SEARCH_DOC : ''}`, search);
+
+  return {
+    extension: defineExtension({
+      name: 'optchat',
+      tools,
+      sections: [section('optchat', () => text, { tag: false })],
+      hooks: [hook(GenerationTask, { beforeRequest }), hook(CompactionTask, { beforeCompact: (_compaction, api) => api.conversationId === id ? { decline: true } : undefined })],
+      tasks: [Remember],
+    }),
+    /** Starts following the conversation, and resets it now if a run ended before a restart. Returns a function that stops. */
+    attach(harness: Harness, ctx: Context) {
+      attached.set(harness, ctx);
+      const stop = harness.subscribeCommits(() => void schedule());
+      void schedule();
+      return stop;
+    },
+  };
+}
+
+/** The Harness and context an extension's `attach` was given; `get` throws before it is attached, or after `clear`. */
+export function attachment(name: string) {
+  let at: { harness: Harness; context: Context } | undefined;
+  return {
+    set: (harness: Harness, context: Context) => { at = { harness, context }; },
+    clear: () => { at = undefined; },
+    get: () => { if (!at) throw new Error(`${name} is not attached to a Harness yet.`); return at; },
+  };
+}
+
+/** Pi Durable's `Models` with OptChat's cache marks on Anthropic requests that carry a view: in blocks of 4 lines, so a view that
+ * grows at its end is read from the cache (pi-ai's own marks miss it). A request that sets its own `onPayload` is left alone, as
+ * the compactor's are. A Proxy, since `createModels()` returns a class instance whose methods a spread would lose. */
+export function withCacheMarks<M extends Models>(models: M): M {
+  return new Proxy(models, {
+    get(target, key) {
+      const value: unknown = Reflect.get(target, key, target);
+      if (key !== 'streamSimple') return typeof value === 'function' ? value.bind(target) : value;
+      const stream: Models['streamSimple'] = (model, request, options) => target.streamSimple(model, request,
+        options?.onPayload || model.api !== 'anthropic-messages' ? options : { ...options, onPayload: payload => cachePayload(payload) });
+      return stream;
+    },
+  });
+}

package/src/durable/store.ts

@@ -0,0 +1,162 @@
+import type { SqliteDatabase, SqliteExecutor, SqliteValue } from '@earendil-works/pi-durable/storage/sqlite';
+import { isMime, type Entry, type ImageStore, type ImportMode, type Mime, type Part, type Saved, type Store, type Summary } from '../core/index.ts';
+
+/** A Durable Object's SQLite refuses a row over 2 MB. A message's text is stored as JSON in parts of at most this many UTF-16 units,
+ * so a part is at most ~1.5 MB of UTF-8 however the text is written, and a huge pasted email still fits. */
+const PART = 512 * 1024;
+export interface SqliteStoreOptions {
+  /** Where images are kept instead of `optchat_images`, such as an R2 bucket that keeps large binary data out of the database
+   * (its row size limit and storage cost). */
+  images?: ImageStore;
+  /** The generation to open, as `optchatImport` does for the one it builds. Default: the live one, whose load also deletes the
+   * rows of generations that are neither live nor being imported. */
+  generation?: number;
+}
+
+/** Log, tree and view rows carry their generation `g`: an import builds a new one beside the live one. `optchat_state` names the
+ * live one, the pending import's (`target`, NULL when none) and the last one handed out; the planned log is in `optchat_plan`.
+ * Images are shared. */
+const TABLES = `CREATE TABLE IF NOT EXISTS optchat_log (g INTEGER NOT NULL, i INTEGER NOT NULL, part INTEGER NOT NULL, meta TEXT, text TEXT NOT NULL, PRIMARY KEY (g, i, part));
+CREATE TABLE IF NOT EXISTS optchat_tree (g INTEGER NOT NULL, l INTEGER NOT NULL, i INTEGER NOT NULL, size INTEGER NOT NULL, text TEXT NOT NULL, PRIMARY KEY (g, l, i));
+CREATE TABLE IF NOT EXISTS optchat_view (g INTEGER PRIMARY KEY, json TEXT NOT NULL);
+CREATE TABLE IF NOT EXISTS optchat_images (name TEXT PRIMARY KEY, mime TEXT NOT NULL, data BLOB NOT NULL);
+CREATE TABLE IF NOT EXISTS optchat_state (id INTEGER PRIMARY KEY CHECK (id = 0), live INTEGER NOT NULL, last INTEGER NOT NULL DEFAULT 0, target INTEGER, mode TEXT, total INTEGER);
+CREATE TABLE IF NOT EXISTS optchat_plan (i INTEGER NOT NULL, part INTEGER NOT NULL, meta TEXT, text TEXT NOT NULL, PRIMARY KEY (i, part));
+INSERT OR IGNORE INTO optchat_state (id, live) VALUES (0, 0);`;
+
+/** Splits `text` every PART units, never between the two halves of a surrogate pair. */
+function parts(text: string) {
+  const result: string[] = [];
+  for (let at = 0; at < text.length;) {
+    let end = Math.min(text.length, at + PART);
+    if (end < text.length && /[\ud800-\udbff]/.test(text[end - 1])) end--;
+    result.push(text.slice(at, end)); at = end;
+  }
+  return result;
+}
+/** An entry as rows of `(part, meta, text)`: its text as JSON in parts, its metadata on the first. */
+const rows = ({ text, ...meta }: Entry) => parts(JSON.stringify(text)).map((chunk, part) => [part, part ? null : JSON.stringify(meta), chunk] as const);
+
+/** A synchronous query on the same database, for the texts memory reads while it builds a view, zooms or searches. In a Durable
+ * Object: `(sql, ...params) => ctx.storage.sql.exec(sql, ...params)`; on Node, node:sqlite's `(sql, ...params) => db.prepare(sql).iterate(...params)`. */
+export type Read = (sql: string, ...params: SqliteValue[]) => Iterable<Record<string, unknown>>;
+/** Texts are stored as JSON, which keeps an unpaired surrogate that SQLite's TEXT would replace. A text that isn't there throws. */
+function texts(rows: Iterable<Record<string, unknown>>, what: string) {
+  const found = [...rows].map(row => typeof row.text === 'string' ? row.text : '');
+  if (!found.length) throw new Error(`No ${what} in this store.`);
+  const text: unknown = JSON.parse(found.join(''));
+  if (typeof text !== 'string') throw new Error(`Invalid ${what} in this store.`);
+  return text;
+}
+
+/** OptChat's memory in any Chord `SqliteDatabase`, such as Node's (`NodeSqliteDatabase`) or a Durable Object's. Give it the
+ * database Pi Durable's storage uses, so both share one queue: in a Durable Object, the `DurableObjectSqliteDatabase` over
+ * `ctx.storage` (direct `ctx.storage.sql` writes could be rolled back with a Pi transaction). `read` reads texts from the same
+ * database synchronously; it only reads rows whose writes have finished, so it needs no queue. Its tables are `optchat_*`, beside
+ * Pi's `pi_*`. Memory loads only metadata and node sizes; texts stay here. Messages are split so no row passes a Durable Object's
+ * 2 MB limit; images go to `images` when given, otherwise to `optchat_images`, one row each. */
+export class SqliteStore implements Store {
+  #g = 0;
+  constructor(private readonly db: SqliteDatabase, private readonly read: Read, private readonly options: SqliteStoreOptions = {}) {}
+  async load(): Promise<Saved> {
+    await this.db.exec(TABLES);
+    if (this.options.generation !== undefined) this.#g = this.options.generation;
+    else {
+      const state = await this.db.get<{ live: number; target: number | null }>('SELECT live, target FROM optchat_state');
+      this.#g = state?.live ?? 0;
+      // An old memory may still write to its generation until the object restarts after a switch, so old rows go only here.
+      await this.db.transaction(async tx => {
+        for (const table of ['log', 'tree', 'view']) await tx.run(`DELETE FROM optchat_${table} WHERE g NOT IN (?, ?)`, this.#g, state?.target ?? this.#g);
+      });
+    }
+    const log = await this.db.all<{ meta: string }>('SELECT meta FROM optchat_log WHERE g = ? AND part = 0 ORDER BY i', this.#g);
+    const tree = await this.db.all<{ l: number; i: number; size: number }>('SELECT l, i, size FROM optchat_tree WHERE g = ?', this.#g);
+    const view = await this.db.get<{ json: string }>('SELECT json FROM optchat_view WHERE g = ?', this.#g);
+    return { entries: log.map(row => JSON.parse(row.meta)), tree, view: view?.json };
+  }
+  text(i: number) { return texts(this.read('SELECT text FROM optchat_log WHERE g = ? AND i = ? ORDER BY part', this.#g, i), `message ${i}`); }
+  summary({ l, i }: Part) { return texts(this.read('SELECT text FROM optchat_tree WHERE g = ? AND l = ? AND i = ?', this.#g, l, i), `summary ${l}/${i}`); }
+  /** The primary key refuses a second message at an index, so a store that was not the last to append writes nothing. */
+  async append(entry: Entry) {
+    try {
+      await this.db.transaction(async tx => {
+        for (const row of rows(entry)) await tx.run('INSERT INTO optchat_log (g, i, part, meta, text) VALUES (?, ?, ?, ?, ?)', this.#g, entry.i, ...row);
+      });
+    } catch (error) {
+      if (await this.db.get('SELECT 1 FROM optchat_log WHERE g = ? AND i = ? AND part = 0', this.#g, entry.i))
+        throw new Error(`Someone else already saved message ${entry.i}; nothing was written.`, { cause: error });
+      throw error;
+    }
+  }
+  async appendNode(node: Summary) {
+    await this.db.run('INSERT OR REPLACE INTO optchat_tree (g, l, i, size, text) VALUES (?, ?, ?, ?, ?)', this.#g, node.l, node.i, node.size, JSON.stringify(node.text));
+  }
+  async saveView(view: string) {
+    await this.db.run('INSERT OR REPLACE INTO optchat_view (g, json) VALUES (?, ?)', this.#g, view);
+  }
+  async putImage(name: string, mimeType: Mime, data: Uint8Array) {
+    if (this.options.images) return this.options.images.putImage(name, mimeType, data);
+    await this.db.run('INSERT OR REPLACE INTO optchat_images (name, mime, data) VALUES (?, ?, ?)', name, mimeType, data);
+  }
+  async image(name: string) {
+    if (this.options.images) return this.options.images.image(name);
+    const row = await this.db.get<{ mime: string; data: Uint8Array }>('SELECT mime, data FROM optchat_images WHERE name = ?', name);
+    return row && isMime(row.mime) ? { mimeType: row.mime, data: row.data } : undefined;
+  }
+}
+
+/** A pending import: the generation it builds, and its plan's mode and length. */
+export interface ImportJob { generation: number; mode: ImportMode; total: number }
+
+/** The rows `optchatImport` keeps: which generation is live, the pending import and its plan. */
+export function importTables(db: SqliteDatabase, read: Read) {
+  const state = async (tx: SqliteExecutor = db) => {
+    await tx.exec(TABLES);
+    const row = await tx.get<{ live: number; target: number | null; mode: ImportMode | null; total: number | null }>('SELECT live, target, mode, total FROM optchat_state');
+    const job: ImportJob | undefined = row?.target != null && row.mode && row.total != null ? { generation: row.target, mode: row.mode, total: row.total } : undefined;
+    return { live: row?.live ?? 0, job };
+  };
+  return {
+    state: () => state(),
+    /** Saves the plan, `entries` in order, made from generation `from`, and makes it the pending import of a new generation. A
+     * generation is never handed out twice, so a late call about a discarded one can't touch the next. Append's new generation starts with the live one's first `kept` messages, the nodes
+     * over them, and its view. Refused unless `from` is still live: an import went live since. */
+    stage: (from: number, mode: ImportMode, entries: Iterable<Entry>, kept: number) => db.transaction(async tx => {
+      const { live, job } = await state(tx);
+      if (job) throw new Error('An import is pending: let it finish, or discard it.');
+      if (live !== from) throw new Error('An import went live: restart before the next one.');
+      const generation = ((await tx.get<{ last: number }>('SELECT last FROM optchat_state'))?.last ?? live) + 1;
+      await tx.run('DELETE FROM optchat_plan');
+      let total = 0;
+      for (const entry of entries) { total++; for (const row of rows(entry)) await tx.run('INSERT INTO optchat_plan (i, part, meta, text) VALUES (?, ?, ?, ?)', entry.i, ...row); }
+      if (kept) {
+        // Nodes only over the messages copied: a message still being written isn't, and a node past the log would refuse to open.
+        // The view keeps its cached prefix; one that covers more than was copied is rebuilt.
+        await tx.run('INSERT INTO optchat_log (g, i, part, meta, text) SELECT ?, i, part, meta, text FROM optchat_log WHERE g = ? AND i < ?', generation, live, kept);
+        await tx.run(`INSERT INTO optchat_tree (g, l, i, size, text) SELECT ?, l, i, size, text FROM optchat_tree
+          WHERE g = ? AND ((i + 1) << l) <= (SELECT COUNT(*) FROM optchat_log WHERE g = ? AND part = 0)`, generation, live, generation);
+        await tx.run('INSERT INTO optchat_view (g, json) SELECT ?, json FROM optchat_view WHERE g = ?', generation, live);
+      }
+      await tx.run('UPDATE optchat_state SET last = ?, target = ?, mode = ?, total = ?', generation, generation, mode, total);
+      return { generation, mode, total };
+    }),
+    /** Planned message `i`, as saved; validated where it is used. */
+    planned(i: number): unknown {
+      const [first] = read('SELECT meta FROM optchat_plan WHERE i = ? AND part = 0', i);
+      if (typeof first?.meta !== 'string') throw new Error(`No planned message ${i} in this store.`);
+      return { ...JSON.parse(first.meta), text: texts(read('SELECT text FROM optchat_plan WHERE i = ? ORDER BY part', i), `planned message ${i}`) };
+    },
+    /** Makes `generation` live and drops its plan, in one transaction; refused if its import was discarded meanwhile. */
+    activate: (generation: number) => db.transaction(async tx => {
+      if ((await state(tx)).job?.generation !== generation) throw new Error('The import was discarded.');
+      await tx.run('UPDATE optchat_state SET live = target, target = NULL, mode = NULL, total = NULL');
+      await tx.run('DELETE FROM optchat_plan');
+    }),
+    /** Drops the pending import if it is still `generation`'s; its rows go when the live store next loads. */
+    discard: (generation: number) => db.transaction(async tx => {
+      if ((await state(tx)).job?.generation !== generation) return;
+      await tx.run('UPDATE optchat_state SET target = NULL, mode = NULL, total = NULL');
+      await tx.run('DELETE FROM optchat_plan');
+    }),
+  };
+}

package/src/handoff.ts

@@ -1,7 +1,8 @@
 import type { ModelRegistry } from '@earendil-works/pi-coding-agent';
 import type { AgentMessage } from '@earendil-works/pi-agent-core';
-import type { AssistantMessage } from '@earendil-works/pi-ai';
-import { reasoningFor, retryTransient, type ModelChoice } from './compactor.ts';
+import { isRetryableAssistantError, type AssistantMessage } from '@earendil-works/pi-ai';
+import { reasoningFor, type ModelChoice } from './compactor.ts';
+import { retryTransient } from './core/compactor.ts';
 import { HANDOFF } from './prompts.ts';
 import { undelivered, type RunInfo } from './runs.ts';
 import { textContent } from './transcript.ts';
@@ -53,7 +54,7 @@ export function createHandoffSummarizer(registry: ModelRegistry, choice: () => M
       const reply = await retryTransient(() => registry.streamSimple(model, {
         systemPrompt: HANDOFF,
         messages: [{ role: 'user', timestamp: Date.now(), content: prefix + transcript.subarray(offset, end).toString('utf8') }],
-      }, { reasoning: reasoningFor(model, selected.thinking), maxTokens, signal: AbortSignal.any([AbortSignal.timeout(300_000), ...signal ? [signal] : []]) }).result(), usage, retryDelays, signal);
+      }, { reasoning: reasoningFor(model, selected.thinking), maxTokens, signal: AbortSignal.any([AbortSignal.timeout(300_000), ...signal ? [signal] : []]) }).result(), usage, isRetryableAssistantError, retryDelays, signal);
       summary = handoffText(reply);
       offset = end;
     }

package/src/images.ts

@@ -1,7 +1,8 @@
 import { createHash } from 'node:crypto';
 import type { ImageContent } from '@earendil-works/pi-ai';
 import { resizeImage } from '@earendil-works/pi-coding-agent';
-import { isMime, type Store } from './store.ts';
+import { isImage } from './core/images.ts';
+import { isMime, type Store } from './core/store.ts';
 
 /** Kept images fit 2048 px on the long side and ~1.5 MB (2 MB as base64), re-encoded as JPEG 85 only when they don't already. */
 export const IMAGE_LIMITS = { maxWidth: 2048, maxHeight: 2048, maxBytes: 2 * 1024 * 1024, jpegQuality: 85 };
@@ -9,10 +10,6 @@ export const IMAGE_LIMITS = { maxWidth: 2048, maxHeight: 2048, maxBytes: 2 * 102
 const hash = (image: ImageContent) => createHash('sha256').update(image.data).digest('hex').slice(0, 16);
 /** How a message's image reads in memory: named by the image as the message holds it, before any shrinking, so it always gets the same name. */
 export const imageRef = (image: ImageContent) => `[image ${hash(image)}]`;
-const REF = /\[image ([0-9a-f]{16})\]/g;
-
-export const isImage = (part: unknown): part is ImageContent => typeof part === 'object' && part !== null
-  && 'type' in part && part.type === 'image' && 'data' in part && typeof part.data === 'string' && 'mimeType' in part && typeof part.mimeType === 'string';
 
 /** Keeps each image of a message's content in the store, once per distinct image, shrunk to IMAGE_LIMITS. */
 export async function saveImages(store: Store, content: unknown) {
@@ -25,10 +22,3 @@ export async function saveImages(store: Store, content: unknown) {
     if (isMime(kept.mimeType)) await store.putImage(name, kept.mimeType, Buffer.from(kept.data, 'base64'));
   }
 }
-
-/** The kept images a text refers to, in order, each once; ones never kept are skipped. */
-export async function loadImages(store: Store, text: string): Promise<ImageContent[]> {
-  const names = [...new Set(Array.from(text.matchAll(REF), match => match[1]))];
-  return (await Promise.all(names.map(name => store.image(name))))
-    .flatMap(kept => kept ? [{ type: 'image' as const, data: Buffer.from(kept.data).toString('base64'), mimeType: kept.mimeType }] : []);
-}

package/src/import/job.ts

@@ -1,13 +1,15 @@
 import { existsSync, readFileSync, rmSync, cpSync, mkdirSync } from 'node:fs';
 import { join, relative, sep } from 'node:path';
 import { randomUUID } from 'node:crypto';
-import { Memory, bytes, isEntry, type Entry, type Compressor } from '../memory.ts';
+import { Memory } from '../memory.ts';
+import { isEntry, type Entry, type Compressor } from '../core/memory.ts';
+import { deduplicate as dedupe, planImport, runImport as build, type ImportMode } from '../core/import.ts';
 import { atomicWrite } from '../store.ts';
-import { record } from '../cache.ts';
+import { record } from '../core/cache.ts';
 import { copyKey, type ImportedEntry } from './sources.ts';
 import { DEFAULT_SETTINGS } from '../settings.ts';
 
-export type ImportMode = 'append' | 'rebuild';
+export type { ImportMode };
 export interface ImportJob {
   id: string; mode: ImportMode; previous: string; target: string; created: string;
   added: number; skipped: number; total: number; inputBytes: number;
@@ -39,41 +41,14 @@ export function pendingImport(dir: string): ImportJob | undefined {
   return { id: j.id, mode: j.mode === 'append' ? 'append' : 'rebuild', previous: j.previous, target: `memories/${j.id}`,
     created: j.created, added: Number(j.added), skipped: Number(j.skipped), total: Number(j.total), inputBytes: Number(j.inputBytes) };
 }
-export function deduplicate(existing: readonly Entry[], incoming: readonly ImportedEntry[]) {
-  const receipts = new Set(existing.map(e => e.receipt).filter(Boolean));
-  const copies = new Set(existing.map(copyKey).filter(Boolean));
-  const added: ImportedEntry[] = []; let skipped = 0;
-  for (const entry of incoming) {
-    if (!entry.receipt?.startsWith('import:')) throw new Error('Imported entry has no stable source identity.');
-    const copy = copyKey(entry);
-    if (receipts.has(entry.receipt) || copy && copies.has(copy)) { skipped++; continue; }
-    receipts.add(entry.receipt); if (copy) copies.add(copy); added.push(entry);
-  }
-  return { added, skipped };
-}
-/** Imported conversations remain contiguous; existing native history is grouped by user turn. */
-export function chronological(entries: readonly ImportedEntry[]) {
-  const groups = new Map<string, { date: string; entries: ImportedEntry[] }>();
-  let native = 0;
-  for (const entry of entries) {
-    if (!entry.origin && entry.kind === 'user') native++;
-    const id = entry.origin ? JSON.stringify([entry.origin.source, entry.origin.conversation]) : `native:${native}`;
-    const group = groups.get(id) ?? { date: entry.date, entries: [] };
-    group.date = group.date < entry.date ? group.date : entry.date;
-    group.entries.push(entry); groups.set(id, group);
-  }
-  return [...groups.values()].sort((a, b) => a.date.localeCompare(b.date)).flatMap(g => g.entries);
-}
+/** Pi also drops a Claude message copied into a forked session (`copyKey`). */
+export const deduplicate = (existing: readonly Entry[], incoming: readonly ImportedEntry[]) => dedupe(existing, incoming, copyKey);
 /** Call only while holding the profile lock, after closing its old Memory worker. */
 export function prepareImport(dir: string, old: Memory, incoming: readonly ImportedEntry[], mode: ImportMode): ImportJob | undefined {
   if (pendingImport(dir)) throw new Error('Resume or discard the pending import first.');
-  const { added, skipped } = deduplicate(old.root, incoming);
-  if (!added.length) return undefined;
-  const id = randomUUID(), target = `memories/${id}`, path = generationPath(dir, target);
-  // Append copies existing summaries without changing their indices. Rebuild regenerates every node.
-  const all = mode === 'rebuild' ? chronological([...old.root, ...added]) : [...old.root, ...chronological(added)];
-  const entries: Entry[] = all.map((e, i) => ({ ...e, i, size: bytes(`${e.kind}: ${e.text}`) }));
-  const kept = mode === 'append' ? old.root.length : 0;
+  const plan = planImport(old.entries(), incoming, mode, copyKey);
+  if (!plan) return undefined;
+  const { entries, kept } = plan, id = randomUUID(), target = `memories/${id}`, path = generationPath(dir, target);
   for (const sub of ['main', 'tree']) mkdirSync(join(path, sub), { recursive: true, mode: 0o700 });
   if (kept) atomicWrite(join(path, 'main', '000-import.jsonl'), entries.slice(0, kept).map(e => JSON.stringify(e)).join('\n') + '\n');
   atomicWrite(join(path, STAGED), entries.map(e => JSON.stringify(e)).join('\n') + '\n');
@@ -83,8 +58,7 @@ export function prepareImport(dir: string, old: Memory, incoming: readonly Impor
     if (existsSync(join(old.directory, 'view.json'))) cpSync(join(old.directory, 'view.json'), join(path, 'view.json'));
   }
   const job: ImportJob = { id, mode, previous: relative(dir, old.directory).split(sep).join('/') || '.', target,
-    created: new Date().toISOString(), added: added.length, skipped, total: entries.length,
-    inputBytes: (mode === 'rebuild' ? entries : added).reduce((n, e) => n + bytes(e.text), 0) };
+    created: new Date().toISOString(), added: plan.added, skipped: plan.skipped, total: entries.length, inputBytes: plan.inputBytes };
   atomicWrite(pendingFile(dir), JSON.stringify(job, null, 2));
   return job;
 }
@@ -104,22 +78,16 @@ export async function runImport(dir: string, compress: Compressor, signal: Abort
   if (memoryDirectory(dir) === path) { finish(); return job; }
   if (!existsSync(join(path, STAGED))) throw new Error('Import staging data is missing; original memory remains intact.');
   const plan: unknown[] = readFileSync(join(path, STAGED), 'utf8').split('\n').filter(Boolean).map(line => JSON.parse(line));
-  // Every line must parse, and the plan must be whole and in order, so a damaged plan can't drop messages silently.
+  // Every line must parse, and the plan must be whole and in order, so a damaged plan is refused before anything is written.
   if (!plan.every(isEntry) || plan.length !== job.total || plan.some((e, i) => e.i !== i)) throw new Error('Import staging data is invalid; original memory remains intact.');
-  const memory = await Memory.open(path, compress, () => {}, undefined, jobs);
-  const report = () => progress({ messages: memory.root.length - memory.pending, total: plan.length, summaries: memory.tree.size, error: memory.lastError });
+  let memory: Memory | undefined;
+  const report = () => memory && progress({ messages: memory.root.length - memory.pending, total: plan.length, summaries: memory.tree.size, error: memory.lastError });
   const timer = setInterval(report, 500);
   try {
-    if (memory.root.length > plan.length || memory.root.some((e, i) => e.text !== plan[i].text)) throw new Error('Import staging data does not match its log; original memory remains intact.');
-    report();
-    // Recipe §10: imported messages are compressed like any other. Each is logged once its node can start, as a burst of live
-    // messages is, so up to `jobs` summaries run at once (the "Summaries at once" setting). Unlike a turn, an import waits through failures (its progress shows the error).
-    for (const e of plan.slice(memory.root.length)) { await memory.settle(signal, 'ahead'); memory.append(e.kind, e.text, e.date, e.receipt, e.origin); }
-    await memory.settle(signal, 'tree'); signal.throwIfAborted();
-    // Close all writers before the single atomic pointer swap. The previous generation stays intact.
-    await memory.close();
-    atomicWrite(join(dir, 'imports', `${job.id}.json`), JSON.stringify({ ...job, completed: new Date().toISOString() }, null, 2));
-    atomicWrite(join(dir, 'active-memory.json'), JSON.stringify(job.target));
+    await build({ entry: i => plan[i], total: job.total, open: async () => { memory = await Memory.open(path, compress, { warn: () => {}, jobs }); report(); return memory; }, activate: () => {
+      atomicWrite(join(dir, 'imports', `${job.id}.json`), JSON.stringify({ ...job, completed: new Date().toISOString() }, null, 2));
+      atomicWrite(join(dir, 'active-memory.json'), JSON.stringify(job.target));
+    } }, signal);
     finish(); report(); return job;
-  } finally { clearInterval(timer); await memory.close(); }
+  } finally { clearInterval(timer); await memory?.close(); } // Also when `open` fails after opening: a throwing progress callback.
 }

package/src/import/sources.ts

@@ -7,11 +7,13 @@ import { execFile } from 'node:child_process';
 import { promisify } from 'node:util';
 import { createHash } from 'node:crypto';
 import { getAgentDir, parseSkillBlock } from '@earendil-works/pi-coding-agent';
-import { record } from '../cache.ts';
-import { bytes, type Entry, type Kind, type Origin } from '../memory.ts';
+import { record } from '../core/cache.ts';
+import { bytes, type Kind, type Origin } from '../core/memory.ts';
+import type { ImportedEntry } from '../core/import.ts';
+export type { ImportedEntry };
 
-export type Source = Origin['source'];
-export type ImportedEntry = Omit<Entry, 'i' | 'size'>;
+/** The histories Pi's OptChat imports, as their entries' `origin.source`. */
+export type Source = 'claude' | 'claude-memory' | 'codex' | 'pi' | 'chatgpt';
 export interface Conversation {
   source: Source; id: string; file: string; title: string; project: string; date: string; size: number;
   exported?: Record<string, unknown>;