@springbrand/agent-runtime 0.2.0-alpha.68 → 0.2.0-alpha.70

package/src/pi/session/migration.ts

@@ -0,0 +1,646 @@
+import { isEqual } from "lodash-es";
+import {
+  decodePiToolRecoveryMilestone,
+  replayPiToolRecovery,
+} from "../turn/tool-recovery";
+/** Versioned SQL data migration, called before the Runtime opens its transcript. */
+import {
+  branchTip,
+  entryLabel,
+  insertEntry,
+  insertUsage,
+  setValue,
+  deleteValue,
+  sessionName,
+  sessionEntryToContextMessages,
+  value,
+  type Entry,
+  type JsonValue,
+  type Write,
+} from "@earendil-works/pi-agent-core/harness/session";
+import {
+  createToolResultMessage,
+  type AgentToolResult,
+  type AgentMessage,
+} from "@earendil-works/pi-agent-core";
+import type { Usage, ToolResultMessage } from "@earendil-works/pi-ai";
+import { CloudflareSessionStorage } from "./storage";
+
+type LegacyRecord = Record<string, unknown> & {
+  id: string;
+  parentId: string | null;
+  type: string;
+  timestamp: string;
+};
+export interface PiSessionMigrationReport {
+  sourceEntries: number;
+  entries: number;
+  tip: string | null;
+  alreadyMigrated: boolean;
+}
+
+function object(input: unknown): Record<string, unknown> {
+  if (!input || typeof input !== "object" || Array.isArray(input))
+    throw new Error("Invalid legacy object");
+  return input as Record<string, unknown>;
+}
+function text(input: unknown): string {
+  if (typeof input !== "string") throw new Error("Invalid legacy string");
+  return input;
+}
+function time(input: unknown): number {
+  const timestamp = Date.parse(text(input));
+  if (!Number.isSafeInteger(timestamp) || timestamp < 0)
+    throw new Error("Invalid legacy timestamp");
+  return timestamp;
+}
+function finite(input: unknown): number {
+  if (typeof input !== "number" || !Number.isFinite(input) || input < 0)
+    throw new Error("Invalid legacy number");
+  return input;
+}
+function message(input: unknown): AgentMessage {
+  const record = object(input);
+  if (
+    ![
+      "user",
+      "assistant",
+      "toolResult",
+      "custom",
+      "bashExecution",
+      "branchSummary",
+      "compactionSummary",
+    ].includes(text(record.role))
+  )
+    throw new Error("Unknown legacy message role");
+  finite(record.timestamp);
+  if (record.role === "assistant" && record.stopReason === "pending")
+    throw new Error("Unsettled legacy assistant message");
+  return record as unknown as AgentMessage;
+}
+function hasTable(sql: SqlStorage, name: string): boolean {
+  return (
+    sql
+      .exec(
+        "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?",
+        name,
+      )
+      .toArray().length > 0
+  );
+}
+
+function assertSameToolResult(
+  left: ToolResultMessage,
+  right: ToolResultMessage,
+): void {
+  for (const key of [
+    "toolCallId",
+    "toolName",
+    "content",
+    "details",
+    "usage",
+    "addedToolNames",
+    "isError",
+  ] as const) {
+    if (!isEqual(left[key], right[key]))
+      throw new Error(`Conflicting legacy Tool result: ${left.toolCallId}`);
+  }
+}
+
+/**
+ * Called synchronously during DO startup, before admission or recovery can execute.
+ * Durable pending work keeps its IDs and state; no old executor runs across this boundary.
+ * A dry run exercises the real import then rolls back DDL, records, and the format switch together.
+ */
+export function migrateLegacyPiSession(
+  storage: DurableObjectStorage,
+  options: { dryRun: boolean },
+): PiSessionMigrationReport {
+  const sql = storage.sql;
+  let report!: PiSessionMigrationReport;
+  const rollbackDryRun = new Error("rollback migration dry run");
+  try {
+    storage.transactionSync(() => {
+      if (hasTable(sql, "pi_storage_state")) {
+        const store = new CloudflareSessionStorage(storage);
+        report = {
+          sourceEntries: 0,
+          entries: store.entries().length,
+          tip: store.branch().at(-1)?.id ?? null,
+          alreadyMigrated: true,
+        };
+        return;
+      }
+      if (
+        !hasTable(sql, "pi_session_entries") ||
+        !hasTable(sql, "pi_session_metadata")
+      )
+        throw new Error("Missing legacy session snapshot");
+      const metadata = sql
+        .exec<{
+          id: string;
+          created_at: string;
+        }>("SELECT id, created_at FROM pi_session_metadata WHERE singleton = 1")
+        .one();
+      const createdAt = time(metadata.created_at);
+      const rows = sql
+        .exec<{
+          seq: number;
+          id: string;
+          body: string;
+        }>("SELECT seq, id, body FROM pi_session_entries ORDER BY seq")
+        .toArray();
+      const records = new Map<string, LegacyRecord>();
+      for (const row of rows) {
+        const raw = object(JSON.parse(row.body));
+        if (raw.id !== row.id || !row.id || records.has(row.id))
+          throw new Error("Invalid or duplicate legacy entry ID");
+        if (
+          raw.parentId !== null &&
+          (typeof raw.parentId !== "string" || !records.has(raw.parentId))
+        )
+          throw new Error(`Missing or forward legacy parent: ${row.id}`);
+        time(raw.timestamp);
+        text(raw.type);
+        records.set(row.id, raw as LegacyRecord);
+      }
+      const converted = new Map<string, Entry>();
+      const resolved = new Map<string, string | null>();
+      const values: Write[] = [];
+      const resolve = (id: unknown): string | null => {
+        if (id === null) return null;
+        const key = text(id);
+        if (converted.has(key)) return key;
+        if (!resolved.has(key))
+          throw new Error(`Unknown legacy reference: ${key}`);
+        return resolved.get(key)!;
+      };
+      for (const [index, row] of rows.entries()) {
+        const raw = records.get(row.id)!;
+        const parentId = resolve(raw.parentId);
+        const timestamp = time(raw.timestamp);
+        const base = { id: raw.id, parentId, timestamp, seq: index + 1 };
+        let entry: Entry | undefined;
+        switch (raw.type) {
+          case "message":
+            entry = { ...base, type: "message", message: message(raw.message) };
+            break;
+          case "custom_message":
+            entry = {
+              ...base,
+              type: "message",
+              message: {
+                role: "custom",
+                customType: text(raw.customType),
+                content: raw.content,
+                details: raw.details,
+                display: raw.display,
+                timestamp,
+              } as AgentMessage,
+            };
+            break;
+          case "custom":
+            entry = {
+              ...base,
+              type: "custom",
+              customType: text(raw.customType),
+              ...(raw.data === undefined
+                ? {}
+                : { data: raw.data as JsonValue }),
+            };
+            break;
+          case "branch_summary":
+            entry = {
+              ...base,
+              type: "branch_summary",
+              fromId: raw.fromId === "root" ? null : resolve(raw.fromId),
+              summary: text(raw.summary),
+              fromHook: raw.fromHook === true,
+              ...(raw.details === undefined
+                ? {}
+                : { details: raw.details as JsonValue }),
+              ...(raw.usage === undefined ? {} : { usage: raw.usage as Usage }),
+            };
+            break;
+          case "compaction": {
+            let retainedTail: AgentMessage[];
+            if (Array.isArray(raw.retainedTail))
+              retainedTail = raw.retainedTail.map(message);
+            else {
+              const kept = text(raw.firstKeptEntryId);
+              const tail: Entry[] = [];
+              let cursor = raw.parentId;
+              let found = false;
+              while (cursor !== null) {
+                const ancestor = records.get(cursor);
+                if (!ancestor)
+                  throw new Error(`Missing compaction ancestor: ${cursor}`);
+                const convertedEntry = converted.get(cursor);
+                if (convertedEntry) tail.unshift(convertedEntry);
+                if (cursor === kept) {
+                  found = true;
+                  break;
+                }
+                cursor = ancestor.parentId;
+              }
+              if (!found)
+                throw new Error(
+                  `Compaction retained tail is not on its historical branch: ${raw.id}`,
+                );
+              retainedTail = tail.flatMap(sessionEntryToContextMessages);
+            }
+            entry = {
+              ...base,
+              type: "compaction",
+              summary: text(raw.summary),
+              retainedTail,
+              tokensBefore: finite(raw.tokensBefore),
+              fromHook: raw.fromHook === true,
+              ...(raw.details === undefined
+                ? {}
+                : { details: raw.details as JsonValue }),
+              ...(raw.usage === undefined ? {} : { usage: raw.usage as Usage }),
+            };
+            break;
+          }
+          case "leaf":
+            resolved.set(raw.id, resolve(raw.targetId));
+            break;
+          case "label": {
+            const target = resolve(raw.targetId);
+            if (!target)
+              throw new Error(`Label has no content target: ${raw.id}`);
+            values.push(
+              raw.label === undefined
+                ? deleteValue(entryLabel(target))
+                : setValue(entryLabel(target), text(raw.label)),
+            );
+            resolved.set(raw.id, parentId);
+            break;
+          }
+          case "session_info":
+            if (raw.name !== undefined)
+              values.push(setValue(sessionName, text(raw.name)));
+            resolved.set(raw.id, parentId);
+            break;
+          case "model_change":
+          case "thinking_level_change":
+          case "active_tools_change": {
+            // Historical settings are values, not conversation entries or a second execution lane.
+            const setting =
+              raw.type === "model_change"
+                ? { provider: text(raw.provider), modelId: text(raw.modelId) }
+                : raw.type === "thinking_level_change"
+                  ? text(raw.thinkingLevel)
+                  : raw.activeToolNames;
+            if (setting === undefined)
+              throw new Error(`Missing legacy configuration: ${raw.id}`);
+            values.push(
+              setValue(value(`springbrand.history.${raw.type}`, raw.id), {
+                at: parentId,
+                timestamp,
+                setting,
+              }),
+            );
+            resolved.set(raw.id, parentId);
+            break;
+          }
+          default:
+            throw new Error(`Unsupported legacy entry type: ${raw.type}`);
+        }
+        if (entry) converted.set(raw.id, entry);
+      }
+      const tip = rows.length ? resolve(rows.at(-1)!.id) : null;
+      if (hasTable(sql, "pi_session_message_submission")) {
+        for (const link of sql
+          .exec<{
+            message_id: string;
+            submission_id: string;
+          }>(
+            "SELECT message_id, submission_id FROM pi_session_message_submission",
+          )
+          .toArray()) {
+          if (converted.get(link.message_id)?.type !== "message")
+            throw new Error(`Dangling message association: ${link.message_id}`);
+          const persistedSubmission =
+            hasTable(sql, "pi_submissions") &&
+            sql
+              .exec(
+                "SELECT 1 FROM pi_submissions WHERE submission_id = ?",
+                link.submission_id,
+              )
+              .toArray().length > 0;
+          if (!persistedSubmission) {
+            let ancestor = converted.get(link.message_id);
+            while (
+              ancestor &&
+              !(ancestor.type === "message" && ancestor.message.role === "user")
+            )
+              ancestor = ancestor.parentId
+                ? converted.get(ancestor.parentId)
+                : undefined;
+            if (!ancestor || link.submission_id !== `import:${ancestor.id}`)
+              throw new Error(
+                `Dangling submission association: ${link.submission_id}`,
+              );
+          }
+        }
+      }
+      const messageOwners = new Map(
+        hasTable(sql, "pi_session_message_submission")
+          ? sql
+              .exec<{ message_id: string; submission_id: string }>(
+                "SELECT message_id, submission_id FROM pi_session_message_submission",
+              )
+              .toArray()
+              .map((row) => [row.message_id, row.submission_id] as const)
+          : [],
+      );
+      const results = new Map<string, ToolResultMessage>();
+      for (const entry of converted.values()) {
+        if (entry.type !== "message" || entry.message.role !== "toolResult")
+          continue;
+        const result = entry.message;
+        let ancestor = entry.parentId
+          ? converted.get(entry.parentId)
+          : undefined;
+        let found = false;
+        while (ancestor) {
+          if (
+            ancestor.type === "message" &&
+            ancestor.message.role === "assistant"
+          ) {
+            const call = ancestor.message.content.find(
+              (block) =>
+                block.type === "toolCall" && block.id === result.toolCallId,
+            );
+            if (call?.type === "toolCall") {
+              if (call.name !== result.toolName)
+                throw new Error("Mismatched legacy Tool name");
+              found = true;
+              break;
+            }
+          }
+          ancestor = ancestor.parentId
+            ? converted.get(ancestor.parentId)
+            : undefined;
+        }
+        if (!found)
+          throw new Error(`Orphan legacy Tool result: ${result.toolCallId}`);
+        const owner = messageOwners.get(entry.id);
+        if (owner) {
+          const key = `${owner}:${result.toolCallId}`;
+          const previous = results.get(key);
+          if (previous) assertSameToolResult(previous, result);
+          results.set(key, result);
+        }
+      }
+      if (hasTable(sql, "pi_pending_steers")) {
+        for (const steer of sql
+          .exec<{
+            submission_id: string;
+            message_id: string;
+            canonical_json: string;
+          }>(
+            "SELECT submission_id, message_id, canonical_json FROM pi_pending_steers",
+          )
+          .toArray()) {
+          if (
+            !hasTable(sql, "pi_submissions") ||
+            !sql
+              .exec(
+                "SELECT 1 FROM pi_submissions WHERE submission_id = ?",
+                steer.submission_id,
+              )
+              .toArray().length
+          ) {
+            throw new Error("Dangling pending steer submission");
+          }
+          const pending = message(JSON.parse(steer.canonical_json));
+          if (pending.role !== "user" || converted.has(steer.message_id)) {
+            throw new Error("Invalid or already committed pending steer");
+          }
+        }
+      }
+      // Business records keep their stable IDs and payloads. Validate ownership before switching.
+      if (hasTable(sql, "pi_submissions")) {
+        const submissions = sql
+          .exec<{
+            submission_id: string;
+            status: string;
+            queued_input_json: string | null;
+            queued_ui_message_json: string | null;
+            user_message_id: string | null;
+            regenerate_message_id: string | null;
+          }>("SELECT * FROM pi_submissions")
+          .toArray();
+        const knownSubmissions = new Set(
+          submissions.map((row) => row.submission_id),
+        );
+        for (const submission of submissions) {
+          if (submission.status !== "pending" || !submission.queued_input_json)
+            continue;
+          const queued = message(JSON.parse(submission.queued_input_json));
+          if (queued.role !== "user" || !submission.user_message_id) {
+            throw new Error("Invalid queued user message");
+          }
+          if (
+            submission.queued_ui_message_json &&
+            object(JSON.parse(submission.queued_ui_message_json)).role !==
+              "user"
+          ) {
+            throw new Error("Invalid queued UI user message");
+          }
+          if (submission.regenerate_message_id) {
+            const target = converted.get(submission.regenerate_message_id);
+            if (
+              target?.type !== "message" ||
+              target.message.role !== "user" ||
+              !isEqual(
+                typeof target.message.content === "string"
+                  ? [{ type: "text", text: target.message.content }]
+                  : target.message.content,
+                typeof queued.content === "string"
+                  ? [{ type: "text", text: queued.content }]
+                  : queued.content,
+              )
+            ) {
+              throw new Error("Invalid queued regeneration target");
+            }
+          } else if (converted.has(submission.user_message_id)) {
+            throw new Error("Queued user message is already committed");
+          }
+        }
+        const toolInputs = new Set<string>();
+        if (hasTable(sql, "pi_recovery_milestones")) {
+          const bySubmission = new Map<string, string[]>();
+          for (const row of sql
+            .exec<{
+              submission_id: string;
+              body: string;
+            }>(
+              "SELECT submission_id, body FROM pi_recovery_milestones ORDER BY submission_id, seq",
+            )
+            .toArray()) {
+            if (!knownSubmissions.has(row.submission_id))
+              throw new Error("Dangling recovery submission");
+            const record = decodePiToolRecoveryMilestone(row.body);
+            if (!record || record.turnId !== row.submission_id)
+              throw new Error("Invalid legacy recovery milestone");
+            const bodies = bySubmission.get(row.submission_id) ?? [];
+            bodies.push(row.body);
+            bySubmission.set(row.submission_id, bodies);
+            if (record.type === "tool-result" || record.type === "tool-error") {
+              for (const result of replayPiToolRecovery([row.body])
+                .canonicalToolResults) {
+                const key = `${row.submission_id}:${result.toolCallId}`;
+                const previous = results.get(key);
+                if (previous) assertSameToolResult(previous, result);
+                results.set(key, result);
+              }
+            }
+          }
+          for (const [submissionId, bodies] of bySubmission) {
+            const recovered = replayPiToolRecovery(bodies);
+            for (const call of Object.values(recovered.toolCalls)) {
+              if (call.input !== undefined || call.retry !== undefined)
+                toolInputs.add(`${submissionId}:${call.toolCallId}`);
+              if (call.result) {
+                const key = `${submissionId}:${call.toolCallId}`;
+                const previous = results.get(key);
+                if (previous) assertSameToolResult(previous, call.result);
+                results.set(key, call.result);
+              }
+            }
+          }
+        }
+        if (hasTable(sql, "pi_session_message_submission")) {
+          for (const row of sql
+            .exec<{
+              message_id: string;
+              submission_id: string;
+            }>(
+              "SELECT message_id, submission_id FROM pi_session_message_submission",
+            )
+            .toArray()) {
+            const entry = converted.get(row.message_id);
+            if (
+              entry?.type === "message" &&
+              entry.message.role === "assistant"
+            ) {
+              for (const block of entry.message.content)
+                if (block.type === "toolCall")
+                  toolInputs.add(`${row.submission_id}:${block.id}`);
+            }
+          }
+        }
+        for (const table of [
+          "pi_tool_settlements",
+          "pi_approvals",
+          "pi_tool_interactions",
+        ]) {
+          if (!hasTable(sql, table)) continue;
+          for (const row of sql
+            .exec<{
+              submission_id: string;
+              tool_call_id: string;
+            }>(`SELECT submission_id, tool_call_id FROM ${table}`)
+            .toArray()) {
+            if (
+              !knownSubmissions.has(row.submission_id) ||
+              !toolInputs.has(`${row.submission_id}:${row.tool_call_id}`)
+            )
+              throw new Error(`Dangling Tool association in ${table}`);
+          }
+        }
+      }
+      if (hasTable(sql, "pi_tool_settlements")) {
+        for (const row of sql
+          .exec<{
+            submission_id: string;
+            tool_call_id: string;
+            tool_name: string;
+            result: string;
+            is_error: number;
+            created_at: number;
+          }>(
+            "SELECT submission_id, tool_call_id, tool_name, result, is_error, created_at FROM pi_tool_settlements",
+          )
+          .toArray()) {
+          const expected = createToolResultMessage({
+            toolCallId: row.tool_call_id,
+            toolName: row.tool_name,
+            result: JSON.parse(row.result) as AgentToolResult<unknown>,
+            isError: row.is_error === 1,
+            timestamp: row.created_at,
+          });
+          const recorded = results.get(
+            `${row.submission_id}:${row.tool_call_id}`,
+          );
+          if (recorded) assertSameToolResult(recorded, expected);
+        }
+      }
+      // The old table remains an inert backup. Failure rolls back DDL, data and the version together.
+      sql.exec(
+        "ALTER TABLE pi_session_entries RENAME TO pi_legacy_entries_backup",
+      );
+      const store = new CloudflareSessionStorage(storage);
+      for (const entry of converted.values()) {
+        const { seq: _seq, timestamp, ...next } = entry;
+        const writes: Write[] = [insertEntry(next)];
+        const usage =
+          entry.type === "message" &&
+          (entry.message.role === "assistant" ||
+            entry.message.role === "toolResult")
+            ? entry.message.usage
+            : entry.type === "compaction" || entry.type === "branch_summary"
+              ? entry.usage
+              : undefined;
+        if (usage) {
+          let usageId = `migration:usage:${entry.id}`;
+          while (records.has(usageId)) usageId = `migration:${usageId}`;
+          writes.push(
+            insertUsage({
+              id: usageId,
+              entryId: entry.id,
+              usage,
+              adjustment: false,
+            }),
+          );
+        }
+        const result = store.commitSync(writes);
+        // Import timestamps belong to the historical snapshot, not the time of the migration.
+        sql.exec(
+          "UPDATE pi_entries SET timestamp = ?, body = json_set(body, '$.timestamp', ?) WHERE id = ?",
+          timestamp,
+          timestamp,
+          entry.id,
+        );
+        entry.seq = result.firstSeq;
+      }
+      const parents = new Set(
+        [...converted.values()].map((entry) => entry.parentId),
+      );
+      for (const entry of converted.values()) {
+        if (entry.id !== tip && !parents.has(entry.id))
+          values.push(setValue(branchTip(`history:${entry.id}`), entry.id));
+      }
+      store.commitSync([...values, setValue(branchTip("main"), tip)]);
+      sql.exec(
+        "UPDATE pi_storage_state SET session_id = ?, created_at = ? WHERE singleton = 1",
+        metadata.id,
+        createdAt,
+      );
+      report = {
+        sourceEntries: rows.length,
+        entries: converted.size,
+        tip,
+        alreadyMigrated: false,
+      };
+      if (options.dryRun) throw rollbackDryRun;
+    });
+  } catch (error) {
+    if (error !== rollbackDryRun) throw error;
+  }
+  return report;
+}