From 01d34faa09ecc2d84447fe09bead604663a21bf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Fri, 29 May 2026 11:06:33 +0100 Subject: [PATCH] fix: backfill task_sessions cost columns for legacy DBs (#1145) --- .../task/repository/sqlite/base_migrations.go | 19 +++++ .../task/repository/sqlite/session_test.go | 78 +++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/apps/backend/internal/task/repository/sqlite/base_migrations.go b/apps/backend/internal/task/repository/sqlite/base_migrations.go index 9a1f0a2d85..1b46c732ca 100644 --- a/apps/backend/internal/task/repository/sqlite/base_migrations.go +++ b/apps/backend/internal/task/repository/sqlite/base_migrations.go @@ -26,6 +26,20 @@ func (r *Repository) migrateTaskSessions() error { return nil } +// migrateSessionsAddCostColumns backfills the per-session cost/token columns on +// task_sessions. These are otherwise only introduced inside two gated +// CREATE-TABLE statements — the fresh create (a no-op once the table exists) and +// the agent_execution_id-drop rebuild (gated on that column still being present). +// A DB that no longer contains the rebuild trigger columns would never gain the +// cost columns, breaking the office cost subscriber's IncrementTaskSessionUsage +// with "no such column: tokens_in". These additive ALTERs are idempotent — the +// MigrateLogger swallows "duplicate column name" on DBs that already have them. +func (r *Repository) migrateSessionsAddCostColumns() { + r.migrate.Apply("task_sessions.cost_subcents", `ALTER TABLE task_sessions ADD COLUMN cost_subcents INTEGER NOT NULL DEFAULT 0`) + r.migrate.Apply("task_sessions.tokens_in", `ALTER TABLE task_sessions ADD COLUMN tokens_in INTEGER NOT NULL DEFAULT 0`) + r.migrate.Apply("task_sessions.tokens_out", `ALTER TABLE task_sessions ADD COLUMN tokens_out INTEGER NOT NULL DEFAULT 0`) +} + // runMigrations applies idempotent ALTER TABLE migrations for schema evolution. func (r *Repository) runMigrations() error { r.migrate.Apply("executors_running.last_message_uuid", `ALTER TABLE executors_running ADD COLUMN last_message_uuid TEXT DEFAULT ''`) @@ -66,6 +80,11 @@ func (r *Repository) runMigrations() error { r.migrate.Apply("task_sessions.workspace_path", `ALTER TABLE task_sessions ADD COLUMN workspace_path TEXT DEFAULT ''`) r.migrate.Apply("repositories.copy_files", `ALTER TABLE repositories ADD COLUMN copy_files TEXT DEFAULT ''`) + // Backfill the per-session cost/token columns. Runs after the gated + // task_sessions rebuilds above so it repairs legacy DBs whose schema can no + // longer trigger a rebuild (see migrateSessionsAddCostColumns). + r.migrateSessionsAddCostColumns() + // Office task extensions - net-new columns on existing main tables. // Idempotent ALTERs; main upgrades pick them up at first boot. // The transient in-branch columns (requires_approval, diff --git a/apps/backend/internal/task/repository/sqlite/session_test.go b/apps/backend/internal/task/repository/sqlite/session_test.go index 3d916e6c41..29bbc9bb2a 100644 --- a/apps/backend/internal/task/repository/sqlite/session_test.go +++ b/apps/backend/internal/task/repository/sqlite/session_test.go @@ -192,6 +192,84 @@ func TestIncrementTaskSessionUsage_EmptySessionIDNoOp(t *testing.T) { } } +// rebuildSessionsWithoutCostColumns drops and recreates task_sessions with the +// post-migration schema MINUS the cost/token columns (cost_subcents, tokens_in, +// tokens_out) and without the agent_execution_id / workflow_step_id trigger +// columns. This reproduces a legacy DB that can never gain the cost columns: the +// gated CREATE-TABLE rebuilds won't fire (their trigger columns are absent) and +// the fresh-create is a no-op because the table already exists. +func rebuildSessionsWithoutCostColumns(t *testing.T, repo *Repository) { + t.Helper() + if _, err := repo.db.Exec(`PRAGMA foreign_keys=OFF`); err != nil { + t.Fatalf("disable fk: %v", err) + } + defer func() { _, _ = repo.db.Exec(`PRAGMA foreign_keys=ON`) }() + stmts := []string{ + `DROP TABLE task_sessions`, + `CREATE TABLE task_sessions ( + id TEXT PRIMARY KEY, + task_id TEXT NOT NULL, + agent_profile_id TEXT, + executor_id TEXT DEFAULT '', + executor_profile_id TEXT DEFAULT '', + environment_id TEXT DEFAULT '', + repository_id TEXT DEFAULT '', + base_branch TEXT DEFAULT '', + agent_profile_snapshot TEXT DEFAULT '{}', + executor_snapshot TEXT DEFAULT '{}', + environment_snapshot TEXT DEFAULT '{}', + repository_snapshot TEXT DEFAULT '{}', + state TEXT NOT NULL DEFAULT 'CREATED', + error_message TEXT DEFAULT '', + metadata TEXT DEFAULT '{}', + started_at TIMESTAMP NOT NULL, + completed_at TIMESTAMP, + updated_at TIMESTAMP NOT NULL, + is_primary INTEGER DEFAULT 0, + is_passthrough INTEGER DEFAULT 0, + review_status TEXT DEFAULT '', + base_commit_sha TEXT DEFAULT '', + task_environment_id TEXT DEFAULT '', + FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE + )`, + } + for _, stmt := range stmts { + if _, err := repo.db.Exec(stmt); err != nil { + t.Fatalf("rebuild task_sessions: %v", err) + } + } +} + +// TestMigrateSessionsAddCostColumns_BackfillsLegacySchema reproduces the office +// cost subscriber failure ("no such column: tokens_in"): a task_sessions table +// that predates the cost columns and no longer contains the rebuild trigger +// columns can never gain them, so IncrementTaskSessionUsage fails. The additive +// migration must backfill the columns idempotently. +func TestMigrateSessionsAddCostColumns_BackfillsLegacySchema(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + + rebuildSessionsWithoutCostColumns(t, repo) + seedForMsgTest(t, repo, "task-mig", "sess-mig", "turn-mig") + + // Precondition: this is the reported bug on a legacy schema. + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 1, 2, 3); err == nil { + t.Fatal("expected missing-column error before backfill") + } + + repo.migrateSessionsAddCostColumns() + + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 1, 2, 3); err != nil { + t.Fatalf("IncrementTaskSessionUsage after backfill: %v", err) + } + + // Idempotent: a second pass over a table that already has the columns is a no-op. + repo.migrateSessionsAddCostColumns() + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 10, 20, 30); err != nil { + t.Fatalf("IncrementTaskSessionUsage after second pass: %v", err) + } +} + // seedRepoLink wires up a workspace, repository, task, task_repositories link // row, and a task_session row in the given state. Used to exercise the join // in CountActiveTaskSessionsByRepository.