Fix SQLite session status migration

Author: stainlu

src/store/sqlite.test.ts

@@ -41,10 +41,22 @@ function columns(db: Database.Database, table: string): Set<string> {
 
 describe("SqliteStore — fresh-DB schema", () => {
   it("creates every table and every current column on a brand new file", () => {
-    const store = new SqliteStore(newDbPath());
+    const path = newDbPath();
+    const store = new SqliteStore(path);
+    const agent = store.agents.create({
+      model: "m",
+      tools: [],
+      instructions: "",
+      permissionPolicy: { type: "always_allow" },
+      callableAgents: [],
+      maxSubagentDepth: 0,
+    });
+    const session = store.sessions.create({ agentId: agent.agentId });
+    store.sessions.beginRun(session.sessionId);
+    expect(store.sessions.get(session.sessionId)?.status).toBe("starting");
     // Read via a second handle so we can snapshot the schema without
     // going through the ORM-shaped methods.
-    const probe = new Database(newDbPath(), { readonly: true });
+    const probe = new Database(path, { readonly: true });
     const agents = columns(probe, "agents");
     const sessions = columns(probe, "sessions");
     const envs = columns(probe, "environments");
@@ -150,6 +162,94 @@ describe("SqliteStore — additive migrations on pre-existing DBs", () => {
     expect(row.environment_id).toBeNull(); // nullable column
   });
 
+  it("upgrades the sessions status constraint in place and preserves session_container foreign keys", () => {
+    const path = newDbPath();
+    const seed = new Database(path);
+    seed.pragma("foreign_keys = ON");
+    seed.exec(`
+      CREATE TABLE sessions (
+        session_id TEXT PRIMARY KEY,
+        agent_id TEXT NOT NULL,
+        environment_id TEXT,
+        status TEXT NOT NULL CHECK (status IN ('idle','running','failed')),
+        ephemeral INTEGER NOT NULL DEFAULT 0,
+        remaining_subagent_depth INTEGER NOT NULL DEFAULT 0,
+        turns INTEGER NOT NULL DEFAULT 0,
+        tokens_in INTEGER NOT NULL DEFAULT 0,
+        tokens_out INTEGER NOT NULL DEFAULT 0,
+        cost_usd REAL NOT NULL DEFAULT 0,
+        error TEXT,
+        created_at INTEGER NOT NULL,
+        last_event_at INTEGER,
+        vault_id TEXT,
+        parent_session_id TEXT,
+        user_id TEXT
+      );
+      CREATE TABLE session_containers (
+        session_id TEXT PRIMARY KEY REFERENCES sessions(session_id) ON DELETE CASCADE,
+        agent_id TEXT NOT NULL,
+        container_id TEXT NOT NULL,
+        container_name TEXT NOT NULL,
+        container_port INTEGER NOT NULL,
+        gateway_token TEXT NOT NULL,
+        claimed_at INTEGER NOT NULL,
+        boot_ms INTEGER,
+        pool_source TEXT
+      );
+      CREATE INDEX idx_session_containers_container ON session_containers(container_id);
+    `);
+    seed.prepare(
+      `INSERT INTO sessions (
+        session_id, agent_id, environment_id, status, ephemeral,
+        remaining_subagent_depth, turns, tokens_in, tokens_out, cost_usd,
+        error, created_at, last_event_at, vault_id, parent_session_id, user_id
+      ) VALUES (
+        ?, ?, NULL, 'idle', 0,
+        0, 0, 0, 0, 0,
+        NULL, ?, NULL, NULL, NULL, NULL
+      )`,
+    ).run("ses_old", "agt_old", 1000);
+    seed.prepare(
+      `INSERT INTO session_containers (
+        session_id, agent_id, container_id, container_name, container_port,
+        gateway_token, claimed_at, boot_ms, pool_source
+      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+    ).run(
+      "ses_old",
+      "agt_old",
+      "cid_old",
+      "openclaw-agt-old",
+      18789,
+      "tok_old",
+      1001,
+      250,
+      "warm",
+    );
+    seed.close();
+
+    const store = new SqliteStore(path);
+    store.sessions.beginRun("ses_old");
+    expect(store.sessions.get("ses_old")?.status).toBe("starting");
+    expect(store.sessions.delete("ses_old")).toBe(true);
+    store.close();
+
+    const probe = new Database(path, { readonly: true });
+    const fk = probe.pragma("foreign_key_list(session_containers)") as Array<{
+      table: string;
+    }>;
+    const remainingContainers = probe
+      .prepare(`SELECT COUNT(*) as count FROM session_containers`)
+      .get() as { count: number };
+    const sessionsSql = probe
+      .prepare(`SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'sessions'`)
+      .get() as { sql: string };
+    probe.close();
+
+    expect(fk[0]?.table).toBe("sessions");
+    expect(remainingContainers.count).toBe(0);
+    expect(sessionsSql.sql).toContain("'starting'");
+  });
+
   it("adds every agents migration column when opening a v1-era DB", () => {
     // Pre-migration agents table: only the fields that existed before
     // Items 12-14, 17, 19 landed. Minimum viable shape.

src/store/sqlite.ts

@@ -40,6 +40,7 @@ import type {
 } from "./types.js";
 
 const nanoid = customAlphabet("abcdefghijklmnopqrstuvwxyz0123456789", 12);
+const SESSION_STATUS_VALUES_SQL = `'idle', 'starting', 'running', 'failed'`;
 
 // ---------- Row shapes ----------
 
@@ -158,6 +159,46 @@ function rowToSession(r: SessionRow): Session {
   };
 }
 
+function createSessionsTableSql(tableName: string): string {
+  return `
+CREATE TABLE ${tableName} (
+  session_id TEXT PRIMARY KEY,
+  agent_id TEXT NOT NULL,
+  environment_id TEXT,
+  status TEXT NOT NULL CHECK (status IN (${SESSION_STATUS_VALUES_SQL})),
+  ephemeral INTEGER NOT NULL DEFAULT 0,
+  remaining_subagent_depth INTEGER NOT NULL DEFAULT 0,
+  turns INTEGER NOT NULL DEFAULT 0,
+  tokens_in INTEGER NOT NULL DEFAULT 0,
+  tokens_out INTEGER NOT NULL DEFAULT 0,
+  cost_usd REAL NOT NULL DEFAULT 0,
+  error TEXT,
+  created_at INTEGER NOT NULL,
+  last_event_at INTEGER,
+  vault_id TEXT,
+  parent_session_id TEXT,
+  user_id TEXT
+)`;
+}
+
+function createSessionContainersTableSql(
+  tableName: string,
+  sessionsTableName: string,
+): string {
+  return `
+CREATE TABLE ${tableName} (
+  session_id TEXT PRIMARY KEY REFERENCES ${sessionsTableName}(session_id) ON DELETE CASCADE,
+  agent_id TEXT NOT NULL,
+  container_id TEXT NOT NULL,
+  container_name TEXT NOT NULL,
+  container_port INTEGER NOT NULL,
+  gateway_token TEXT NOT NULL,
+  claimed_at INTEGER NOT NULL,
+  boot_ms INTEGER,
+  pool_source TEXT
+)`;
+}
+
 // ---------- Schema bootstrap ----------
 
 // Applied once per database. Idempotent — every CREATE uses IF NOT EXISTS.
@@ -226,7 +267,8 @@ CREATE TABLE IF NOT EXISTS environments (
 CREATE TABLE IF NOT EXISTS sessions (
   session_id TEXT PRIMARY KEY,
   agent_id TEXT NOT NULL,
-  status TEXT NOT NULL CHECK (status IN ('idle', 'running', 'failed')),
+  environment_id TEXT,
+  status TEXT NOT NULL CHECK (status IN (${SESSION_STATUS_VALUES_SQL})),
   ephemeral INTEGER NOT NULL DEFAULT 0,
   remaining_subagent_depth INTEGER NOT NULL DEFAULT 0,
   turns INTEGER NOT NULL DEFAULT 0,
@@ -237,7 +279,8 @@ CREATE TABLE IF NOT EXISTS sessions (
   created_at INTEGER NOT NULL,
   last_event_at INTEGER,
   vault_id TEXT,
-  parent_session_id TEXT
+  parent_session_id TEXT,
+  user_id TEXT
 );
 
 CREATE INDEX IF NOT EXISTS idx_sessions_agent_id ON sessions(agent_id);
@@ -1794,6 +1837,7 @@ export class SqliteStore implements Store {
     if (scCols.length > 0 && !scCols.some((c) => c.name === "pool_source")) {
       this.db.exec("ALTER TABLE session_containers ADD COLUMN pool_source TEXT");
     }
+    this.migrateSessionStatusConstraint();
 
     const envCols = this.db.pragma("table_info(environments)") as Array<{
       name: string;
@@ -1845,6 +1889,86 @@ export class SqliteStore implements Store {
     this.migratedVaultCredentials = vaultsStore.migratePlaintextCredentials();
   }
 
+  private migrateSessionStatusConstraint(): void {
+    const row = this.db
+      .prepare(
+        `SELECT sql
+           FROM sqlite_master
+          WHERE type = 'table' AND name = 'sessions'`,
+      )
+      .get() as { sql: string | null } | undefined;
+    const sessionsSql = row?.sql ?? "";
+    if (sessionsSql.includes("'starting'")) return;
+
+    const fkWasEnabled =
+      (this.db.pragma("foreign_keys", { simple: true }) as number) === 1;
+    if (fkWasEnabled) {
+      this.db.pragma("foreign_keys = OFF");
+    }
+    try {
+      const migrate = this.db.transaction(() => {
+        this.db.exec(createSessionsTableSql("sessions_new"));
+        this.db.exec(`
+          INSERT INTO sessions_new (
+            session_id, agent_id, environment_id, status, ephemeral,
+            remaining_subagent_depth, turns, tokens_in, tokens_out, cost_usd,
+            error, created_at, last_event_at, vault_id, parent_session_id, user_id
+          )
+          SELECT
+            session_id, agent_id, environment_id, status, ephemeral,
+            remaining_subagent_depth, turns, tokens_in, tokens_out, cost_usd,
+            error, created_at, last_event_at, vault_id, parent_session_id, user_id
+          FROM sessions
+        `);
+        this.db.exec(
+          createSessionContainersTableSql(
+            "session_containers_new",
+            "sessions_new",
+          ),
+        );
+        this.db.exec(`
+          INSERT INTO session_containers_new (
+            session_id, agent_id, container_id, container_name, container_port,
+            gateway_token, claimed_at, boot_ms, pool_source
+          )
+          SELECT
+            session_id, agent_id, container_id, container_name, container_port,
+            gateway_token, claimed_at, boot_ms, pool_source
+          FROM session_containers
+        `);
+        this.db.exec("DROP TABLE session_containers");
+        this.db.exec("DROP TABLE sessions");
+        this.db.exec("ALTER TABLE sessions_new RENAME TO sessions");
+        this.db.exec(
+          "ALTER TABLE session_containers_new RENAME TO session_containers",
+        );
+        this.db.exec(
+          "CREATE INDEX IF NOT EXISTS idx_sessions_agent_id ON sessions(agent_id)",
+        );
+        this.db.exec(
+          "CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id)",
+        );
+        this.db.exec(
+          "CREATE INDEX IF NOT EXISTS idx_session_containers_container ON session_containers(container_id)",
+        );
+      });
+      migrate();
+    } finally {
+      if (fkWasEnabled) {
+        this.db.pragma("foreign_keys = ON");
+      }
+    }
+
+    const fkViolations = this.db.pragma(
+      "foreign_key_check",
+    ) as Array<Record<string, unknown>>;
+    if (fkViolations.length > 0) {
+      throw new Error(
+        `session status migration left foreign key violations: ${JSON.stringify(fkViolations)}`,
+      );
+    }
+  }
+
   close(): void {
     if (this.closed) return;
     this.db.close();