@springbrand/agent-runtime 0.2.0-alpha.67 → 0.2.0-alpha.69

package/src/pi/session/storage.ts

@@ -1,353 +1,565 @@
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
 import {
-  SessionError,
-  uuidv7,
-  type AgentMessage,
+  prepareStorageCommit,
+  validateCommittedWrites,
+  resolveListReadOptions,
+  value,
+  branchTip,
+  setValue,
+  insertEntry,
+  insertUsage,
   type CompactionEntry,
+  type Storage,
+  type Write,
+  type CommitResult,
+  type Entry,
+  type EntryStructure,
+  type EntryScan,
+  type StorageBranchScan,
+  type UsageRow,
+  type UsageScan,
+  type Value,
+  type ValueList,
+  type StoredValue,
+  type ListElement,
+  type ListReadOptions,
+  type SessionStats,
   type SessionMetadata,
-  type SessionTreeEntry,
-} from "@earendil-works/pi-agent-core";
-import type { SqlTaggedTemplate } from "agents/chat";
+} from "@earendil-works/pi-agent-core/harness/session";
+import type { Usage } from "@earendil-works/pi-ai";
+import type { Context } from "@earendil-works/pi-agent-core/harness/context";
 
-type EntryRow = {
-  id: string;
-  body: string;
-};
+/** Pi's Storage contract, with every physical commit confined to one DO transaction. */
+export class CloudflareSessionStorage implements Storage {
+  private closed = false;
+  private readonly sql: SqlStorage;
+  readonly metadata: SessionMetadata;
 
-type LeafRow = {
-  id: string;
-  type: string;
-  target_id: string | null;
-};
-
-function errorCause(error: unknown): Error {
-  return error instanceof Error ? error : new Error(String(error));
-}
-
-function parseEntry(row: EntryRow): SessionTreeEntry {
-  try {
-    return JSON.parse(row.body) as SessionTreeEntry;
-  } catch (error) {
-    throw new SessionError(
-      "invalid_entry",
-      `Invalid SQLite session entry ${row.id}`,
-      errorCause(error),
+  constructor(private readonly storage: DurableObjectStorage) {
+    this.sql = storage.sql;
+    const hasVersion =
+      this.sql
+        .exec(
+          "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'pi_storage_state'",
+        )
+        .toArray().length > 0;
+    const hasLegacy =
+      this.sql
+        .exec(
+          "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'pi_session_entries'",
+        )
+        .toArray().length > 0;
+    if (!hasVersion && hasLegacy) {
+      throw new Error(
+        "Pi session requires an explicit legacy-to-0.85 migration",
+      );
+    }
+    if (hasVersion) {
+      const version = this.sql
+        .exec<{
+          storage_version: number;
+        }>("SELECT storage_version FROM pi_storage_state WHERE singleton = 1")
+        .one().storage_version;
+      if (version !== 1)
+        throw new Error(
+          `Unsupported Cloudflare Pi storage version: ${version}`,
+        );
+    }
+    this.sql.exec(`
+      CREATE TABLE IF NOT EXISTS pi_session_message_submission (message_id TEXT PRIMARY KEY, submission_id TEXT NOT NULL);
+      CREATE TABLE IF NOT EXISTS pi_entries (
+        id TEXT PRIMARY KEY, seq INTEGER NOT NULL UNIQUE, parent_id TEXT,
+        type TEXT NOT NULL, custom_type TEXT, timestamp INTEGER NOT NULL, body TEXT NOT NULL
+      );
+      CREATE TABLE IF NOT EXISTS pi_values (
+        namespace TEXT NOT NULL, key TEXT NOT NULL, seq INTEGER NOT NULL, body TEXT NOT NULL,
+        PRIMARY KEY(namespace, key)
+      );
+      CREATE TABLE IF NOT EXISTS pi_lists (
+        namespace TEXT NOT NULL, key TEXT NOT NULL, seq INTEGER PRIMARY KEY, body TEXT NOT NULL
+      );
+      CREATE INDEX IF NOT EXISTS pi_lists_address ON pi_lists(namespace, key, seq);
+      CREATE TABLE IF NOT EXISTS pi_usage (id TEXT PRIMARY KEY, seq INTEGER NOT NULL UNIQUE, body TEXT NOT NULL);
+      CREATE TABLE IF NOT EXISTS pi_storage_state (singleton INTEGER PRIMARY KEY CHECK(singleton = 1), next_seq INTEGER NOT NULL, storage_version INTEGER NOT NULL, session_id TEXT NOT NULL, created_at INTEGER NOT NULL, stats TEXT NOT NULL);
+    `);
+    this.sql.exec(
+      "INSERT OR IGNORE INTO pi_storage_state VALUES (1, 1, 1, ?, ?, ?)",
+      crypto.randomUUID(),
+      Date.now(),
+      JSON.stringify({
+        messageCount: 0,
+        usage: {
+          input: 0,
+          output: 0,
+          cacheRead: 0,
+          cacheWrite: 0,
+          totalTokens: 0,
+          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+        },
+      }),
     );
+    const row = this.sql
+      .exec<{
+        storage_version: number;
+        session_id: string;
+        created_at: number;
+      }>(
+        "SELECT storage_version, session_id, created_at FROM pi_storage_state WHERE singleton = 1",
+      )
+      .one();
+    this.metadata = {
+      id: row.session_id,
+      createdAt: row.created_at,
+      storageVersion: row.storage_version,
+    };
   }
-}
 
-function sessionEntryTargetId(entry: SessionTreeEntry): string | null {
-  return entry.type === "leaf" || entry.type === "label"
-    ? entry.targetId
-    : null;
-}
+  private assertOpen(context?: Context): void {
+    if (this.closed) throw new Error("Session storage is closed");
+    context?.abortSignal?.throwIfAborted();
+  }
 
-function getEntrySync(
-  sql: SqlTaggedTemplate,
-  id: string,
-): SessionTreeEntry | undefined {
-  const row = sql<EntryRow>`
-    SELECT id, body
-    FROM pi_session_entries
-    WHERE id = ${id}
-  `[0];
-  return row ? parseEntry(row) : undefined;
-}
+  /** Synchronous boundary also used by the host's admission/settlement transaction. */
+  commitSync(writes: Write[], context?: Context): CommitResult {
+    this.assertOpen(context);
+    return this.storage.transactionSync(() => {
+      const firstSeq = this.sql
+        .exec<{
+          next_seq: number;
+        }>("SELECT next_seq FROM pi_storage_state WHERE singleton = 1")
+        .one().next_seq;
+      const prepared = prepareStorageCommit(writes, firstSeq, Date.now());
+      validateCommittedWrites(prepared.writes, firstSeq, {
+        hasEntryId: (id) =>
+          this.sql.exec("SELECT 1 FROM pi_entries WHERE id = ?", id).toArray()
+            .length > 0,
+        hasEntryOrUsageId: (id) =>
+          this.sql
+            .exec(
+              "SELECT 1 FROM pi_entries WHERE id = ? UNION ALL SELECT 1 FROM pi_usage WHERE id = ?",
+              id,
+              id,
+            )
+            .toArray().length > 0,
+      });
+      const stats = this.statsSync();
+      for (const write of prepared.writes) {
+        switch (write.kind) {
+          case "entry": {
+            const { kind: _kind, ...entry } = write;
+            if (entry.type === "message") stats.messageCount += 1;
+            this.sql.exec(
+              "INSERT INTO pi_entries VALUES (?, ?, ?, ?, ?, ?, ?)",
+              entry.id,
+              entry.seq,
+              entry.parentId,
+              entry.type,
+              entry.customType ?? null,
+              entry.timestamp,
+              JSON.stringify(entry),
+            );
+            break;
+          }
+          case "usage": {
+            const { kind: _kind, ...row } = write;
+            for (const key of [
+              "input",
+              "output",
+              "cacheRead",
+              "cacheWrite",
+              "totalTokens",
+              "cacheWrite1h",
+              "reasoning",
+            ] as const) {
+              if (row.usage[key] !== undefined)
+                stats.usage[key] = (stats.usage[key] ?? 0) + row.usage[key];
+            }
+            for (const key of [
+              "input",
+              "output",
+              "cacheRead",
+              "cacheWrite",
+              "total",
+            ] as const)
+              stats.usage.cost[key] += row.usage.cost[key];
+            this.sql.exec(
+              "INSERT INTO pi_usage VALUES (?, ?, ?)",
+              row.id,
+              row.seq,
+              JSON.stringify(row),
+            );
+            break;
+          }
+          case "value":
+            if (write.op === "delete")
+              this.sql.exec(
+                "DELETE FROM pi_values WHERE namespace = ? AND key = ?",
+                write.namespace,
+                write.key,
+              );
+            else
+              this.sql.exec(
+                "INSERT INTO pi_values VALUES (?, ?, ?, ?) ON CONFLICT(namespace, key) DO UPDATE SET seq = excluded.seq, body = excluded.body",
+                write.namespace,
+                write.key,
+                write.seq,
+                JSON.stringify(write.value),
+              );
+            break;
+          case "list":
+            if (write.op === "delete")
+              this.sql.exec(
+                "DELETE FROM pi_lists WHERE namespace = ? AND key = ?",
+                write.namespace,
+                write.key,
+              );
+            else
+              this.sql.exec(
+                "INSERT INTO pi_lists VALUES (?, ?, ?, ?)",
+                write.namespace,
+                write.key,
+                write.seq,
+                JSON.stringify(write.value),
+              );
+            break;
+        }
+      }
+      this.sql.exec(
+        "UPDATE pi_storage_state SET next_seq = ?, stats = ? WHERE singleton = 1",
+        firstSeq + writes.length,
+        JSON.stringify(stats),
+      );
+      return { ...prepared.result, stats };
+    });
+  }
 
-function getLeafIdSync(sql: SqlTaggedTemplate): string | null {
-  const row = sql<LeafRow>`
-    SELECT id, type, target_id
-    FROM pi_session_entries
-    ORDER BY seq DESC
-    LIMIT 1
-  `[0];
-  if (!row) return null;
+  async commit(writes: Write[], context: Context): Promise<CommitResult> {
+    return this.commitSync(writes, context);
+  }
 
-  const leafId = row.type === "leaf" ? row.target_id : row.id;
-  if (leafId !== null && !getEntrySync(sql, leafId)) {
-    throw new SessionError(
-      "invalid_session",
-      `Entry ${leafId} not found`,
-    );
+  async getEntries(
+    ids: string[],
+    context: Context,
+  ): Promise<Map<string, Entry>> {
+    return this.getEntriesSync(ids, context);
   }
-  return leafId;
-}
 
-function appendEntrySync(
-  sql: SqlTaggedTemplate,
-  entry: SessionTreeEntry,
-): boolean {
-  const before = sql<{ count: number }>`
-    SELECT COUNT(*) AS count
-    FROM pi_session_entries
-    WHERE id = ${entry.id}
-  `[0]?.count ?? 0;
-  if (before > 0) return false;
+  getEntriesSync(ids: string[], context?: Context): Map<string, Entry> {
+    this.assertOpen(context);
+    const entries = new Map<string, Entry>();
+    for (const id of ids) {
+      const row = this.sql
+        .exec<{ body: string }>("SELECT body FROM pi_entries WHERE id = ?", id)
+        .toArray()[0];
+      if (row) entries.set(id, JSON.parse(row.body));
+    }
+    return entries;
+  }
 
-  sql`
-    INSERT INTO pi_session_entries (
-      id, type, parent_id, target_id, body
-    ) VALUES (
-      ${entry.id},
-      ${entry.type},
-      ${entry.parentId},
-      ${sessionEntryTargetId(entry)},
-      ${JSON.stringify(entry)}
-    )
-  `;
-  return true;
-}
+  async getValue<T>(
+    address: Value<T>,
+    context: Context,
+  ): Promise<StoredValue<T> | undefined> {
+    return this.getValueSync(address, context);
+  }
 
-/**
- * 给 Durable Object 的 `transactionSync` 调用方用的同步写原语。
- *
- * 必须是同步的:admission 走在 DO 的同步事务里,任何 await 都会破坏那次事务的原子性。
- */
-export function appendSessionMessageSync(
-  sql: SqlTaggedTemplate,
-  input: {
-    id: string;
-    message: AgentMessage;
-    submissionId?: string;
-    createdAt?: number;
-  },
-): boolean {
-  const timestamp = new Date(
-    input.createdAt ??
-      ("timestamp" in input.message &&
-        typeof input.message.timestamp === "number"
-        ? input.message.timestamp
-        : Date.now()),
-  ).toISOString();
-  const inserted = appendEntrySync(sql, {
-    type: "message",
-    id: input.id,
-    parentId: getLeafIdSync(sql),
-    timestamp,
-    message: input.message,
-  });
-  if (inserted && input.submissionId) {
-    sql`
-      INSERT INTO pi_session_message_submission (
-        message_id, submission_id
-      ) VALUES (
-        ${input.id}, ${input.submissionId}
+  getValueSync<T>(
+    address: Value<T>,
+    context?: Context,
+  ): StoredValue<T> | undefined {
+    this.assertOpen(context);
+    const row = this.sql
+      .exec<{
+        seq: number;
+        body: string;
+      }>(
+        "SELECT seq, body FROM pi_values WHERE namespace = ? AND key = ?",
+        address.namespace,
+        address.key,
       )
-    `;
+      .toArray()[0];
+    return row
+      ? { address, seq: row.seq, value: JSON.parse(row.body) }
+      : undefined;
   }
-  return inserted;
-}
 
-export function moveSessionLeafSync(
-  sql: SqlTaggedTemplate,
-  targetId: string | null,
-): void {
-  if (targetId !== null && !getEntrySync(sql, targetId)) {
-    throw new SessionError("not_found", `Entry ${targetId} not found`);
-  }
-  const id = uuidv7();
-  if (!appendEntrySync(sql, {
-    type: "leaf",
-    id,
-    parentId: getLeafIdSync(sql),
-    timestamp: new Date().toISOString(),
-    targetId,
-  })) {
-    throw new SessionError(
-      "storage",
-      `Session entry ${id} already exists`,
-    );
+  async scanValues<T>(
+    prefix: Value<T>,
+    context: Context,
+  ): Promise<StoredValue<T>[]> {
+    this.assertOpen(context);
+    return this.sql
+      .exec<{ key: string; seq: number; body: string }>(
+        "SELECT key, seq, body FROM pi_values WHERE namespace = ? AND substr(key, 1, length(?)) = ? ORDER BY key COLLATE BINARY",
+        prefix.namespace,
+        prefix.key,
+        prefix.key,
+      )
+      .toArray()
+      .map((row) => ({
+        address: value<T>(prefix.namespace, row.key),
+        seq: row.seq,
+        value: JSON.parse(row.body),
+      }));
   }
-}
 
-export function clearSessionSync(sql: SqlTaggedTemplate): void {
-  sql`DELETE FROM pi_session_message_submission`;
-  sql`DELETE FROM pi_session_entries`;
-}
+  async readList<T>(
+    address: ValueList<T>,
+    options: ListReadOptions | undefined,
+    context: Context,
+  ): Promise<ListElement<T>[]> {
+    this.assertOpen(context);
+    const resolved = resolveListReadOptions(options);
+    const descending = resolved.order === "desc";
+    const comparison = descending ? "<" : ">";
+    return this.sql
+      .exec<{ seq: number; body: string }>(
+        `SELECT seq, body FROM pi_lists WHERE namespace = ? AND key = ? AND (? IS NULL OR seq ${comparison} ?) ORDER BY seq ${descending ? "DESC" : "ASC"} LIMIT ?`,
+        address.namespace,
+        address.key,
+        resolved.cursor?.seq ?? null,
+        resolved.cursor?.seq ?? null,
+        resolved.limit,
+      )
+      .toArray()
+      .map((row) => ({ seq: row.seq, value: JSON.parse(row.body) }));
+  }
 
-/**
- * 回合事实源:这个 Durable Object 自己拥有的会话记录流。
- *
- * ## 为什么不再 implements pi 的 SessionStorage
- *
- * 我们只用到 branch / entry / appendCompaction 三件事,却曾为此实现一个十方法接口,
- * 并把「一个回合到底长什么样」的定义权交给了 pi 的 **harness 层** —— 而 `SessionStorage`
- * 住在 `dist/harness/types.d.ts`,是 pi 给本地 CLI 用的那一层,不是 pi 内核的必需品
- * (`AgentOptions` 里根本没有 session 字段)。它也是 pi 发版最勤的一层:
- * 0.81.0 已经把 `getPathToRoot` 改名成 `getPathToRootOrCompaction` 并新增两个必填方法,
- * `[Unreleased]` 又要把整个 `SessionStorage` 换成 `SessionRepository` + caller-owned
- * `SessionStore`。**实现别人的接口 = 别人改语义时编译器不会告诉我们。**
- *
- * 现在语义由我们定、写在我们的测试里(见 storage.test.ts 的差分对拍);pi 只作为
- * 纯函数被调用(`prepareCompaction` / `compact` / `buildSessionContext`),
- * 那些改签名编译器一定会报。
- *
- * ## 会话模型(读代码前先看这三条,否则会误改)
- *
- * 1. **leaf 是派生的,不是存的**:seq 最大的那一行;若它是 `leaf` 标记条目,则跟随它的
- *    `targetId`。所以**追加任何 entry 都会自动成为新 leaf**,不需要额外「移动」。
- * 2. **回退 leaf 靠追加标记**:`moveSessionLeafSync` append 一条 `leaf` 类型条目指向目标,
- *    表本身始终 append-only,从不改写历史行。
- * 3. **`leaf` 标记永远不会成为任何条目的 parent**:新条目的 parentId 取自
- *    `getLeafIdSync()`,那已经是解析过标记之后的真实条目 id。所以按 parentId 上溯时
- *    不会走进标记里。
- */
-export class PiTranscriptStore {
-  constructor(
-    private readonly sql: SqlTaggedTemplate,
-    metadata?: SessionMetadata,
-  ) {
-    this.sql`
-      CREATE TABLE IF NOT EXISTS pi_session_metadata (
-        singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
-        id TEXT NOT NULL,
-        created_at TEXT NOT NULL
+  private branchRows(query: StorageBranchScan, structure: boolean) {
+    // Stops follow the requested traversal direction; filters must not hide a stop entry.
+    const oldestFirst = query.order === "oldestFirst";
+    const rows = this.sql
+      .exec<{
+        id: string;
+        parentId: string | null;
+        seq: number;
+        timestamp: number;
+        type: Entry["type"];
+        customType: string | null;
+        body?: string;
+      }>(
+        `
+      WITH RECURSIVE path AS (
+        SELECT id, parent_id, seq, timestamp, type, custom_type${structure ? "" : ", body"} FROM pi_entries WHERE id = ?
+        UNION ALL
+        SELECT e.id, e.parent_id, e.seq, e.timestamp, e.type, e.custom_type${structure ? "" : ", e.body"}
+        FROM pi_entries e JOIN path p ON e.id = p.parent_id
+      ) SELECT id, parent_id AS parentId, seq, timestamp, type, custom_type AS customType${structure ? "" : ", body"}
+        FROM path ORDER BY seq ${oldestFirst ? "ASC" : "DESC"}
+    `,
+        query.start,
       )
-    `;
-    this.sql`
-      CREATE TABLE IF NOT EXISTS pi_session_entries (
-        seq INTEGER PRIMARY KEY AUTOINCREMENT,
-        id TEXT NOT NULL UNIQUE,
-        type TEXT NOT NULL,
-        parent_id TEXT,
-        target_id TEXT,
-        body TEXT NOT NULL
+      .toArray();
+    if (!rows.length) throw new Error(`Unknown branch start: ${query.start}`);
+    const root = oldestFirst ? rows[0]! : rows.at(-1)!;
+    if (root.parentId !== null)
+      throw new Error("Corrupt branch: missing parent");
+    const stop = rows.findIndex(
+      (row) => row.id === query.stopAtId || row.type === query.stopAtType,
+    );
+    return (stop < 0 ? rows : rows.slice(0, stop + 1))
+      .filter(
+        (row) =>
+          (query.type === undefined || query.type === row.type) &&
+          (query.customType === undefined ||
+            query.customType === row.customType) &&
+          (query.cursor === undefined ||
+            (oldestFirst
+              ? row.seq > query.cursor.seq
+              : row.seq < query.cursor.seq)),
       )
-    `;
-    this.sql`
-      CREATE INDEX IF NOT EXISTS pi_session_entries_type
-      ON pi_session_entries(type)
-    `;
-    this.sql`
-      CREATE INDEX IF NOT EXISTS pi_session_entries_target
-      ON pi_session_entries(type, target_id, seq)
-    `;
-    this.sql`
-      CREATE TABLE IF NOT EXISTS pi_session_message_submission (
-        message_id TEXT PRIMARY KEY,
-        submission_id TEXT NOT NULL
+      .slice(
+        0,
+        query.limit === undefined ? undefined : Math.max(0, query.limit),
+      );
+  }
+
+  async scanBranch(
+    query: StorageBranchScan,
+    context: Context,
+  ): Promise<Entry[]> {
+    return this.scanBranchSync(query, context);
+  }
+
+  scanBranchSync(query: StorageBranchScan, context?: Context): Entry[] {
+    this.assertOpen(context);
+    return this.branchRows(query, false).map((row) => JSON.parse(row.body!));
+  }
+
+  async scanBranchStructure(
+    query: StorageBranchScan,
+    context: Context,
+  ): Promise<EntryStructure[]> {
+    this.assertOpen(context);
+    return this.branchRows(query, true).map(
+      ({ customType, body: _body, ...row }) => ({
+        ...row,
+        ...(customType === null ? {} : { customType }),
+      }),
+    );
+  }
+
+  async scanEntries(query: EntryScan, context: Context): Promise<Entry[]> {
+    this.assertOpen(context);
+    return this.sql
+      .exec<{ body: string }>(
+        `
+      SELECT body FROM pi_entries WHERE (? IS NULL OR type = ?) AND (? IS NULL OR custom_type = ?)
+      AND (? IS NULL OR seq >= ?) AND (? IS NULL OR seq <= ?) ORDER BY seq ${query.order === "desc" ? "DESC" : "ASC"} LIMIT ?`,
+        query.type ?? null,
+        query.type ?? null,
+        query.customType ?? null,
+        query.customType ?? null,
+        query.fromSeq ?? null,
+        query.fromSeq ?? null,
+        query.toSeq ?? null,
+        query.toSeq ?? null,
+        query.limit === undefined ? -1 : Math.max(0, Math.trunc(query.limit)),
       )
-    `;
-    this.sql`
-      CREATE INDEX IF NOT EXISTS pi_session_message_submission_id
-      ON pi_session_message_submission(submission_id)
-    `;
-    const initial = metadata ?? {
-      id: uuidv7(),
-      createdAt: new Date().toISOString(),
-    };
-    this.sql`
-      INSERT OR IGNORE INTO pi_session_metadata (
-        singleton, id, created_at
-      ) VALUES (
-        1, ${initial.id}, ${initial.createdAt}
+      .toArray()
+      .map((row) => JSON.parse(row.body));
+  }
+
+  async scanUsage(query: UsageScan, context: Context): Promise<UsageRow[]> {
+    this.assertOpen(context);
+    return this.sql
+      .exec<{ body: string }>(
+        `
+      SELECT body FROM pi_usage WHERE (? IS NULL OR seq >= ?) AND (? IS NULL OR seq <= ?)
+      ORDER BY seq ${query.order === "desc" ? "DESC" : "ASC"} LIMIT ?`,
+        query.fromSeq ?? null,
+        query.fromSeq ?? null,
+        query.toSeq ?? null,
+        query.toSeq ?? null,
+        query.limit === undefined ? -1 : Math.max(0, Math.trunc(query.limit)),
       )
-    `;
+      .toArray()
+      .map((row) => JSON.parse(row.body));
   }
 
-  /**
-   * 从当前 leaf 上溯到 root 的完整路径。
-   *
-   * 用一条递归 CTE 取全路径。旧实现是 while 循环里每个祖先打一次 SELECT ——
-   * canonicalMessages / storedMessages 各调一次,等于每次刷 UI 都是 O(消息数) 次 SQL。
-   * JOIN 落在 `id` 上(UNIQUE,有索引),不需要为 parent_id 另建索引。
-   */
-  branch(): SessionTreeEntry[] {
-    const leafId = getLeafIdSync(this.sql);
-    if (leafId === null) return [];
+  private statsSync(): SessionStats {
+    return JSON.parse(
+      this.sql
+        .exec<{
+          stats: string;
+        }>("SELECT stats FROM pi_storage_state WHERE singleton = 1")
+        .one().stats,
+    );
+  }
 
-    const rows = this.sql<EntryRow>`
-      WITH RECURSIVE path(id, body, parent_id) AS (
-        SELECT id, body, parent_id
-        FROM pi_session_entries
-        WHERE id = ${leafId}
-        UNION ALL
-        SELECT e.id, e.body, e.parent_id
-        FROM pi_session_entries e
-        JOIN path p ON e.id = p.parent_id
-      )
-      SELECT id, body FROM path
-    `;
-    if (rows.length === 0) {
-      throw new SessionError("not_found", `Entry ${leafId} not found`);
-    }
-    // CTE 从 leaf 往上产出,这里翻成 root→leaf 与旧的 unshift 顺序一致。
-    const path = rows.map(parseEntry).reverse();
-    // 断链检测:旧实现在父条目缺失时抛 invalid_session,而 CTE 遇到断点只会静默停止
-    // 递归。顶端条目的 parentId 必须为空,否则说明中间有一条祖先不见了。
-    const root = path[0];
-    if (root?.parentId) {
-      throw new SessionError(
-        "invalid_session",
-        `Entry ${root.parentId} not found`,
-      );
-    }
-    return path;
+  async getStats(context: Context): Promise<SessionStats> {
+    this.assertOpen(context);
+    return this.statsSync();
   }
 
-  entry(id: string): SessionTreeEntry | undefined {
-    return getEntrySync(this.sql, id);
+  /** Host commits use this single branch; off-branch entries remain available by ID. */
+  branch(): Entry[] {
+    const tip = this.getValueSync(branchTip("main"))?.value;
+    return tip ? this.scanBranchSync({ start: tip, order: "oldestFirst" }) : [];
   }
 
-  /**
-   * 追加一条压缩条目,整对象写入。
-   *
-   * 旧路径是把 assembly 侧已经造好的完整 CompactionEntry 拆成 5 个位置参数喂给
-   * pi 的 `Session.appendCompaction` —— pi 0.81 把它加到 7 个参数(usage / retainedTail)
-   * 之后,多出来的两个会被静默丢掉,编译器不报。整对象写入让这类失败模式在结构上
-   * 不存在:调用方给什么就落什么。
-   *
-   * id / parentId / timestamp 由这里补齐,语义与 pi 的 `appendTypedEntry` 一致
-   * (parentId 取当前 leaf;追加后本条自动成为新 leaf)。
-   */
-  appendCompaction(
-    entry: Omit<CompactionEntry, "id" | "parentId" | "timestamp">,
-  ): string {
-    const id = this.createEntryId();
-    const appended = appendEntrySync(this.sql, {
-      ...entry,
-      id,
-      parentId: getLeafIdSync(this.sql),
-      timestamp: new Date().toISOString(),
+  entry(id: string): Entry | undefined {
+    return this.getEntriesSync([id]).get(id);
+  }
+
+  appendMessage(input: {
+    id: string;
+    message: AgentMessage;
+    submissionId?: string;
+  }): boolean {
+    return this.storage.transactionSync(() => {
+      if (this.entry(input.id)) return false;
+      const usage =
+        input.message.role === "assistant" ||
+        input.message.role === "toolResult"
+          ? input.message.usage
+          : undefined;
+      this.commitSync([
+        insertEntry({
+          id: input.id,
+          parentId: this.getValueSync(branchTip("main"))?.value ?? null,
+          type: "message",
+          message: input.message,
+        }),
+        setValue(branchTip("main"), input.id),
+        ...(usage
+          ? [
+              insertUsage({
+                id: crypto.randomUUID(),
+                entryId: input.id,
+                usage,
+                adjustment: false,
+              }),
+            ]
+          : []),
+      ]);
+      if (input.submissionId)
+        this.sql.exec(
+          "INSERT INTO pi_session_message_submission VALUES (?, ?)",
+          input.id,
+          input.submissionId,
+        );
+      return true;
     });
-    if (!appended) {
-      throw new SessionError("storage", `Session entry ${id} already exists`);
-    }
+  }
+
+  moveBranchTip(id: string | null): void {
+    if (id !== null && !this.entry(id))
+      throw new Error(`Unknown branch target: ${id}`);
+    this.commitSync([setValue(branchTip("main"), id)]);
+  }
+
+  recordCompaction(
+    entry: Omit<CompactionEntry, "id" | "parentId" | "seq" | "timestamp"> | undefined,
+    usage: Usage | undefined = entry?.usage,
+  ): string {
+    const id = crypto.randomUUID();
+    const writes: Write[] = [];
+    if (entry) writes.push(
+      insertEntry({
+        ...entry, id,
+        parentId: this.getValueSync(branchTip("main"))?.value ?? null,
+      }),
+      setValue(branchTip("main"), id),
+    );
+    if (usage) writes.push(insertUsage({
+      id: entry ? crypto.randomUUID() : id,
+      ...(entry ? { entryId: id } : {}),
+      usage,
+      adjustment: false,
+    }));
+    this.commitSync(writes);
     return id;
   }
 
-  /**
-   * 按类型取条目。生产路径无消费者,但压缩相关测试靠它断言「压缩条目真的落库了」,
-   * 是那条行为契约的观测口。
-   */
-  findEntries<TType extends SessionTreeEntry["type"]>(
-    type: TType,
-  ): Array<Extract<SessionTreeEntry, { type: TType }>> {
-    return this.sql<EntryRow>`
-      SELECT id, body
-      FROM pi_session_entries
-      WHERE type = ${type}
-      ORDER BY seq ASC
-    `.map(parseEntry) as Array<Extract<SessionTreeEntry, { type: TType }>>;
+  entries(): Entry[] {
+    return this.sql
+      .exec<{ body: string }>("SELECT body FROM pi_entries ORDER BY seq")
+      .toArray()
+      .map((row) => JSON.parse(row.body));
   }
 
-  /** 全量条目(含 leaf 标记),按写入顺序。仅供测试观测整棵树用。 */
-  entries(): SessionTreeEntry[] {
-    return this.sql<EntryRow>`
-      SELECT id, body
-      FROM pi_session_entries
-      ORDER BY seq ASC
-    `.map(parseEntry);
+  clear(): void {
+    this.storage.transactionSync(() => {
+      this.sql.exec(
+        "DELETE FROM pi_session_message_submission; DELETE FROM pi_entries; DELETE FROM pi_values; DELETE FROM pi_lists; DELETE FROM pi_usage;",
+      );
+      const usage = {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+      };
+      this.sql.exec(
+        "UPDATE pi_storage_state SET stats = ? WHERE singleton = 1",
+        JSON.stringify({ messageCount: 0, usage }),
+      );
+    });
   }
 
-  /**
-   * 短 id,与 pi 的 `createEntryId` 同规则(uuidv7 尾 8 位 + 碰撞重试)。
-   * 必须保持一致:压缩条目的 firstKeptEntryId 引用的就是这种 id。
-   */
-  private createEntryId(): string {
-    for (let attempt = 0; attempt < 100; attempt += 1) {
-      const id = uuidv7().slice(-8);
-      if (!getEntrySync(this.sql, id)) return id;
-    }
-    return uuidv7();
+  async close(_context?: Context): Promise<void> {
+    this.closed = true;
   }
 }