diff --git a/apps/backend/internal/backendapp/storage.go b/apps/backend/internal/backendapp/storage.go index c2aaa5fc13..6fff18b40f 100644 --- a/apps/backend/internal/backendapp/storage.go +++ b/apps/backend/internal/backendapp/storage.go @@ -69,6 +69,13 @@ func provideRepositories(ctx context.Context, cfg *config.Config, log *logger.Lo return nil, nil, nil, fmt.Errorf("office repo: %w", err) } cleanups = append(cleanups, officeCleanup) + // Must run after office.Provide succeeds: it reconciles task_sessions + // against office's office_cost_events ledger, and construction order is + // what guarantees that table exists now - see the doc comment on + // sqlite.Repository.BackfillSessionTokensCachedIn. + if err := backfillSessionCachedTokens(ctx, taskRepoImpl, log); err != nil { + return nil, nil, nil, err + } terminalRepoImpl, err := terminalrepo.NewWithDB(writer, reader, log) if err != nil { return nil, nil, nil, fmt.Errorf("terminal repo: %w", err) @@ -121,6 +128,30 @@ func provideRepositories(ctx context.Context, cfg *config.Config, log *logger.Lo return pool, repos, cleanups, nil } +// sessionCachedTokenBackfiller is the narrow slice of *sqlite.Repository +// (internal/task/repository/sqlite) this composition root needs, defined +// locally so this file doesn't have to import that subpackage just for one +// call. See sqlite.Repository.BackfillSessionTokensCachedIn's doc comment +// for why this call lives here rather than inside the task repository's own +// runMigrations(): it reconciles task_sessions against the office +// repository's office_cost_events ledger, and only this composition root +// knows both repositories have finished initializing. +type sessionCachedTokenBackfiller interface { + BackfillSessionTokensCachedIn(ctx context.Context) (int64, error) +} + +// backfillSessionCachedTokens triggers the task repository's cached-token +// reconciliation and logs how many task_sessions rows it reconciled. +// Extracted out of provideRepositories to keep it within the funlen limit. +func backfillSessionCachedTokens(ctx context.Context, taskRepo sessionCachedTokenBackfiller, log *logger.Logger) error { + affected, err := taskRepo.BackfillSessionTokensCachedIn(ctx) + if err != nil { + return err + } + log.Info("reconciled task_sessions.tokens_cached_in against office_cost_events", zap.Int64("rows", affected)) + return nil +} + // supportRepositorySet groups the lighter-weight support repositories // that share a common (writer, reader) wire-up pattern. type supportRepositorySet struct { diff --git a/apps/backend/internal/backendapp/storage_cached_tokens_test.go b/apps/backend/internal/backendapp/storage_cached_tokens_test.go new file mode 100644 index 0000000000..5aa4e42dd4 --- /dev/null +++ b/apps/backend/internal/backendapp/storage_cached_tokens_test.go @@ -0,0 +1,93 @@ +package backendapp + +import ( + "context" + "testing" + "time" + + "github.com/kandev/kandev/internal/common/config" +) + +// TestProvideRepositoriesBackfillsSessionCachedTokens is the composition-level +// regression test for moving task_sessions.tokens_cached_in's backfill +// trigger out of the task repository's own runMigrations() and into this +// package (see the "PR Fixup" plan section on carlosflorencio's PR #2521 +// architecture comment). Neither repository's own package tests can see the +// other's construction order, so only a test that boots the real repository +// graph through provideRepositories can prove the backfill still actually +// runs — a regression here would silently resurrect the original card's bug +// (task_sessions.tokens_cached_in drifting from the office_cost_events +// ledger) without any single-package test failing. +func TestProvideRepositoriesBackfillsSessionCachedTokens(t *testing.T) { + ctx := context.Background() + cfg := &config.Config{ + HomeDir: t.TempDir(), + Database: config.DatabaseConfig{Driver: "sqlite"}, + } + log := newTestLogger() + + // First boot: creates the schema (both repositories' tables, including + // office_cost_events) against a real on-disk SQLite file. Then seed a + // drifted task_sessions row directly - bypassing IncrementTaskSessionUsage + // - the same way a missed live increment would leave one, plus a ledger + // row establishing the true value the next boot's backfill must recover. + pool, _, cleanups, err := provideRepositories(ctx, cfg, log, "test") + if err != nil { + t.Fatalf("provideRepositories (first boot): %v", err) + } + writer := pool.Writer() + now := time.Now().UTC() + if _, err := writer.Exec(writer.Rebind( + `INSERT INTO workspaces (id, name, created_at, updated_at) VALUES (?, 'test workspace', ?, ?)`), + "ws-cached-boot", now, now); err != nil { + t.Fatalf("seed workspace: %v", err) + } + if _, err := writer.Exec(writer.Rebind( + `INSERT INTO tasks (id, workspace_id, title, created_at, updated_at) VALUES (?, ?, 'test', ?, ?)`), + "task-cached-boot", "ws-cached-boot", now, now); err != nil { + t.Fatalf("seed task: %v", err) + } + if _, err := writer.Exec(writer.Rebind( + `INSERT INTO task_sessions (id, task_id, started_at, updated_at, tokens_cached_in) VALUES (?, ?, ?, ?, ?)`), + "sess-cached-boot", "task-cached-boot", now, now, 2_000); err != nil { + t.Fatalf("seed drifted session: %v", err) + } + if _, err := writer.Exec(writer.Rebind( + `INSERT INTO office_cost_events (id, session_id, tokens_cached_in, occurred_at, created_at) VALUES (?, ?, ?, ?, ?)`), + "cost-cached-boot", "sess-cached-boot", 98_805_109, now, now); err != nil { + t.Fatalf("seed ledger row: %v", err) + } + for i := len(cleanups) - 1; i >= 0; i-- { + if cleanups[i] != nil { + _ = cleanups[i]() + } + } + + // Second boot against the same on-disk DB (same cfg.HomeDir, so the same + // kandev.db file) simulates a restart. The composition-level backfill + // call in provideRepositories must reconcile the drifted row above. + pool2, _, cleanups2, err := provideRepositories(ctx, cfg, log, "test") + if err != nil { + t.Fatalf("provideRepositories (second boot): %v", err) + } + t.Cleanup(func() { + for i := len(cleanups2) - 1; i >= 0; i-- { + if cleanups2[i] != nil { + _ = cleanups2[i]() + } + } + }) + + var got int64 + writer2 := pool2.Writer() + if err := writer2.QueryRowx(writer2.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-cached-boot", + ).Scan(&got); err != nil { + t.Fatalf("read row after second boot: %v", err) + } + if got != 98_805_109 { + t.Errorf("tokens_cached_in after reboot = %d, want 98805109 (the drifted 2000 must be "+ + "reconciled against the office_cost_events ledger by the composition-level backfill "+ + "call in provideRepositories, after office.Provide succeeds)", got) + } +} diff --git a/apps/backend/internal/office/repository/sqlite/base.go b/apps/backend/internal/office/repository/sqlite/base.go index b784431e29..1f11c51065 100644 --- a/apps/backend/internal/office/repository/sqlite/base.go +++ b/apps/backend/internal/office/repository/sqlite/base.go @@ -279,6 +279,13 @@ func (r *Repository) createCostTables() error { CREATE INDEX IF NOT EXISTS idx_office_cost_agent ON office_cost_events(agent_profile_id); CREATE INDEX IF NOT EXISTS idx_office_cost_occurred ON office_cost_events(occurred_at DESC); CREATE INDEX IF NOT EXISTS idx_office_cost_task ON office_cost_events(task_id); + -- Indexes the join column internal/task/repository/sqlite's + -- BackfillSessionTokensCachedIn correlates task_sessions against. Lives + -- here (not in the task repository) because it indexes a table this + -- repository owns. Without it that correlated subquery falls back to a + -- full table scan per task_sessions row - measured at 76.72s at a modest + -- 4,000 sessions / 80,000 events versus 0.17s indexed. + CREATE INDEX IF NOT EXISTS idx_office_cost_events_session_id ON office_cost_events(session_id); CREATE TABLE IF NOT EXISTS office_budget_policies ( id TEXT PRIMARY KEY, diff --git a/apps/backend/internal/office/repository/sqlite/costs_test.go b/apps/backend/internal/office/repository/sqlite/costs_test.go index 766a7e4274..41de5fa575 100644 --- a/apps/backend/internal/office/repository/sqlite/costs_test.go +++ b/apps/backend/internal/office/repository/sqlite/costs_test.go @@ -35,6 +35,25 @@ func seedTasksTable(t *testing.T, repo *sqlite.Repository, taskID, workspaceID s } } +// TestCreateCostTables_IndexesSessionID guards the boot-time cost of +// internal/task/repository/sqlite's BackfillSessionTokensCachedIn, which +// correlates task_sessions against office_cost_events on session_id. Without +// this index that correlated subquery falls back to a full table scan per +// task_sessions row (a measured ~450x slowdown at 4,000 sessions / 80,000 +// events - 76.72s vs 0.17s). The index lives here, not in the task +// repository, because it indexes a table this repository owns. +func TestCreateCostTables_IndexesSessionID(t *testing.T) { + repo := newTestRepo(t) + + var name string + err := repo.ReaderDB().QueryRow( + `SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'office_cost_events' AND sql LIKE '%session_id%'`, + ).Scan(&name) + if err != nil { + t.Fatalf("expected an index on office_cost_events(session_id) to exist after repo init, got: %v", err) + } +} + func TestCostEvent_CreateAndList(t *testing.T) { repo := newTestRepo(t) ctx := context.Background() diff --git a/apps/backend/internal/office/service/event_subscribers.go b/apps/backend/internal/office/service/event_subscribers.go index c741b67253..8b5b84a586 100644 --- a/apps/backend/internal/office/service/event_subscribers.go +++ b/apps/backend/internal/office/service/event_subscribers.go @@ -534,9 +534,9 @@ func (s *Service) tryPostStartFallback( // normalize the model id and look up pricing. On miss the row // records cost_subcents=0 with estimated=true. // -// After insert the session totals (tokens_in / tokens_out / cost_subcents) -// are incremented on task_sessions, and any applicable budget policy is -// evaluated. Estimated rows count toward budget totals at face value. +// After insert the session totals (tokens_in / tokens_cached_in / tokens_out / +// cost_subcents) are incremented on task_sessions, and any applicable budget +// policy is evaluated. Estimated rows count toward budget totals at face value. func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error { data, err := decodeEventData[PromptUsageData](event) if err != nil || data.TaskID == "" || data.SessionID == "" { @@ -549,6 +549,7 @@ func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error costSubcents, estimated := s.resolveCostForUsage(ctx, *data) provider := resolveProvider(*data) + tokensCachedIn := data.Usage.CachedReadTokens + data.Usage.CachedWriteTokens costEvent := &models.CostEvent{ SessionID: data.SessionID, @@ -558,7 +559,7 @@ func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error Model: data.Model, Provider: provider, TokensIn: data.Usage.InputTokens, - TokensCachedIn: data.Usage.CachedReadTokens + data.Usage.CachedWriteTokens, + TokensCachedIn: tokensCachedIn, TokensOut: data.Usage.OutputTokens, CostSubcents: costSubcents, Estimated: estimated, @@ -570,7 +571,7 @@ func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error s.incrementSessionUsageTotals( ctx, data.SessionID, - data.Usage.InputTokens, data.Usage.OutputTokens, costSubcents, + data.Usage.InputTokens, tokensCachedIn, data.Usage.OutputTokens, costSubcents, ) if fields.WorkspaceID != "" { @@ -652,13 +653,13 @@ func providerFromCLI(cli string) string { } func (s *Service) incrementSessionUsageTotals( - ctx context.Context, sessionID string, tokensIn, tokensOut, costSubcents int64, + ctx context.Context, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, ) { if s.sessionUsageWriter == nil || sessionID == "" { return } if err := s.sessionUsageWriter.IncrementTaskSessionUsage( - ctx, sessionID, tokensIn, tokensOut, costSubcents, + ctx, sessionID, tokensIn, tokensCachedIn, tokensOut, costSubcents, ); err != nil { s.logger.Warn("increment task_session usage failed", zap.String("session_id", sessionID), zap.Error(err)) diff --git a/apps/backend/internal/office/service/session_usage_reconcile_test.go b/apps/backend/internal/office/service/session_usage_reconcile_test.go new file mode 100644 index 0000000000..f7414b0580 --- /dev/null +++ b/apps/backend/internal/office/service/session_usage_reconcile_test.go @@ -0,0 +1,120 @@ +package service_test + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/events/bus" +) + +// fakeSessionUsageWriter accumulates the deltas handlePromptUsage sends to +// the task-session rollup, mirroring what the real SQLite repo's +// IncrementTaskSessionUsage does with UPDATE ... SET x = x + ?. +type fakeSessionUsageWriter struct { + tokensIn int64 + tokensCachedIn int64 + tokensOut int64 + costSubcents int64 +} + +func (f *fakeSessionUsageWriter) IncrementTaskSessionUsage( + _ context.Context, _ string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, +) error { + f.tokensIn += tokensIn + f.tokensCachedIn += tokensCachedIn + f.tokensOut += tokensOut + f.costSubcents += costSubcents + return nil +} + +// TestPromptUsage_RollupReconcilesWithCostLedger asserts the actual +// invariant this rollup exists to maintain: task_sessions' cumulative +// totals must equal the sum of the office_cost_events rows they were +// derived from. Before the fix, tokensCachedIn silently never reaches the +// writer even though every event row carries it, which is exactly what let +// task_sessions.tokens_in under-report cache-heavy sessions by orders of +// magnitude in production. +func TestPromptUsage_RollupReconcilesWithCostLedger(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + writer := &fakeSessionUsageWriter{} + svc.SetSessionUsageWriter(writer) + + createTestAgent(t, svc, "ws-1", "worker-reconcile") + svc.ExecSQL(t, `INSERT INTO tasks ( + id, workspace_id, project_id, title, created_at, updated_at + ) VALUES ( + 'task-reconcile', 'ws-1', 'project-1', 'Reconcile task', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP + )`) + setTestTaskAssignee(t, svc, "task-reconcile", "worker-reconcile") + + const sessionID = "session-reconcile-1" + + // Cache-heavy shape mirroring the reported bug: cached read tokens + // dominate input tokens by several orders of magnitude. + usages := []map[string]interface{}{ + { + "input_tokens": 100, + "cached_read_tokens": 50_000_000, + "cached_write_tokens": 1_000_000, + "output_tokens": 200, + }, + { + "input_tokens": 10, + "cached_read_tokens": 2_000, + "cached_write_tokens": 0, + "output_tokens": 5, + }, + } + for _, usage := range usages { + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-reconcile", + "session_id": sessionID, + "agent_id": "claude-acp", + "model": "default", + "usage": usage, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject(sessionID), event); err != nil { + t.Fatalf("publish prompt usage: %v", err) + } + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil { + t.Fatalf("list costs: %v", err) + } + + var wantIn, wantCachedIn, wantOut, wantCost int64 + var ledgerRows int + for _, c := range costs { + if c.SessionID != sessionID { + continue + } + ledgerRows++ + wantIn += c.TokensIn + wantCachedIn += c.TokensCachedIn + wantOut += c.TokensOut + wantCost += c.CostSubcents + } + if ledgerRows != len(usages) { + t.Fatalf("ledger rows for session = %d, want %d", ledgerRows, len(usages)) + } + if wantCachedIn == 0 { + t.Fatal("test setup bug: ledger recorded zero cached tokens, reconciliation would be vacuous") + } + + if writer.tokensIn != wantIn { + t.Errorf("rollup tokens_in = %d, want %d (sum of office_cost_events.tokens_in)", writer.tokensIn, wantIn) + } + if writer.tokensCachedIn != wantCachedIn { + t.Errorf("rollup tokens_cached_in = %d, want %d (sum of office_cost_events.tokens_cached_in)", + writer.tokensCachedIn, wantCachedIn) + } + if writer.tokensOut != wantOut { + t.Errorf("rollup tokens_out = %d, want %d (sum of office_cost_events.tokens_out)", writer.tokensOut, wantOut) + } + if writer.costSubcents != wantCost { + t.Errorf("rollup cost_subcents = %d, want %d (sum of office_cost_events.cost_subcents)", writer.costSubcents, wantCost) + } +} diff --git a/apps/backend/internal/office/shared/interfaces.go b/apps/backend/internal/office/shared/interfaces.go index 6f550b9e61..6081c72b88 100644 --- a/apps/backend/internal/office/shared/interfaces.go +++ b/apps/backend/internal/office/shared/interfaces.go @@ -154,5 +154,5 @@ type PricingLookup interface { // task_sessions when a cost event lands. Implemented by the task repo. type SessionUsageWriter interface { IncrementTaskSessionUsage(ctx context.Context, sessionID string, - tokensIn, tokensOut, costSubcents int64) error + tokensIn, tokensCachedIn, tokensOut, costSubcents int64) error } diff --git a/apps/backend/internal/task/repository/sqlite/base_migrations.go b/apps/backend/internal/task/repository/sqlite/base_migrations.go index 1cf3301d25..b923cb5aad 100644 --- a/apps/backend/internal/task/repository/sqlite/base_migrations.go +++ b/apps/backend/internal/task/repository/sqlite/base_migrations.go @@ -2,6 +2,7 @@ package sqlite import ( + "context" "database/sql" "errors" "fmt" @@ -39,6 +40,84 @@ 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`) + // BIGINT, not INTEGER: office_cost_events.tokens_cached_in routinely + // accumulates well past int4's 2,147,483,647 ceiling over a long-running + // session (the reported bug measured up to 98,805,109 on one already- + // completed task). SQLite's INTEGER is 64-bit regardless, but on Postgres + // INTEGER is int4 - an overflowing session would abort the single + // multi-column UPDATE in IncrementTaskSessionUsage and the single + // table-wide UPDATE in BackfillSessionTokensCachedIn, silently taking + // tokens_in/tokens_out/cost_subcents down with it for that session. + r.migrate.Apply("task_sessions.tokens_cached_in", `ALTER TABLE task_sessions ADD COLUMN tokens_cached_in BIGINT NOT NULL DEFAULT 0`) +} + +// BackfillSessionTokensCachedIn recomputes task_sessions.tokens_cached_in from +// the office_cost_events ledger. The rollup is purely derived from the ledger, +// so recomputing from it is restoring data, not inventing it. +// +// Deliberately NOT called from runMigrations(): office_cost_events is owned and +// created by internal/office/repository/sqlite, a separate repository that +// shares this database but initializes independently, and on a fresh boot the +// task repository's own migrations run first (internal/backendapp/storage.go). +// An earlier version of this method guarded on the ledger table's existence to +// tolerate that ordering from inside runMigrations(); that guard, and the table +// existence it checked, both went away in favor of the caller in +// internal/backendapp/storage.go, which only invokes this after office.Provide +// has already succeeded — construction order guarantees office_cost_events +// exists by then, so no guard is needed here. This keeps the task repository +// from having to know office's schema/table-existence details, and lets the +// office_cost_events(session_id) index it depends on live with the table it +// indexes (internal/office/repository/sqlite/base.go createCostTables). +// +// Deliberately unconditional per session it touches (assignment, not +// increment, so it is idempotent across replays) rather than gated on +// "already nonzero": the rollup can go out of sync with the ledger without +// erroring — a live IncrementTaskSessionUsage call is a best-effort +// UPDATE ... WHERE id = ? that silently matches zero rows if the session row +// doesn't exist yet (see the "missing row" case in +// TestIncrementTaskSessionUsage_UnknownSessionNoError), and +// event_subscribers.go only logs a Warn when the rollup write fails while the +// ledger insert has already committed. A "skip if nonzero" guard would make +// exactly that drift permanent instead of self-healing it on the next boot. +// (If the ledger for a session was deleted separately from the session row +// itself - see the two-transaction delete in workspace_deletion.go - this +// recompute correctly zeroes that session's tokens_cached_in, since an empty +// ledger sums to zero; the rollup's only contract is to equal the ledger sum.) +// +// The WHERE clause below is a boot-cost optimization, not a correctness gate: +// it restricts the write to sessions that either have ledger rows (need a +// possible recompute) or already hold a nonzero value (need a possible +// zeroing, per the paragraph above). Every row it excludes has +// tokens_cached_in = 0 AND no ledger rows, so the unconditional +// COALESCE(NULL, 0) = 0 this statement would otherwise compute for it is +// already what the row holds — excluding it is a provable no-op, never a +// skip of a row that needs correcting. This runs unconditionally on every +// boot (no run-once tracking - see MigrateLogger.Apply, which this method +// deliberately does not use so a real failure propagates instead of being +// swallowed as a Warn), so office_cost_events(session_id) must already be +// indexed: an unindexed correlated subquery here is quadratic in +// (sessions x ledger rows), measured at 76.72s at a modest 4,000 sessions / +// 80,000 events versus 0.17s indexed. +// +// Returns the number of task_sessions rows the WHERE clause matched, so a +// caller can log it for boot-time observability; not otherwise used for +// correctness. +func (r *Repository) BackfillSessionTokensCachedIn(ctx context.Context) (int64, error) { + result, err := r.db.ExecContext(ctx, ` + UPDATE task_sessions + SET tokens_cached_in = COALESCE( + (SELECT SUM(e.tokens_cached_in) FROM office_cost_events e + WHERE e.session_id = task_sessions.id), 0) + WHERE EXISTS (SELECT 1 FROM office_cost_events e WHERE e.session_id = task_sessions.id) + OR tokens_cached_in <> 0`) + if err != nil { + return 0, fmt.Errorf("backfill task_sessions.tokens_cached_in: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return 0, fmt.Errorf("backfill task_sessions.tokens_cached_in: rows affected: %w", err) + } + return affected, nil } // runMigrations applies idempotent ALTER TABLE migrations for schema evolution. @@ -154,6 +233,10 @@ func (r *Repository) runMigrations() error { // task_sessions rebuilds above so it repairs legacy DBs whose schema can no // longer trigger a rebuild (see migrateSessionsAddCostColumns). r.migrateSessionsAddCostColumns() + // BackfillSessionTokensCachedIn is deliberately NOT called here - see its + // doc comment. It runs from internal/backendapp/storage.go, after both + // this repository and the office repository (which owns office_cost_events) + // have finished initializing. // Office task extensions - net-new columns on existing main tables. // Idempotent ALTERs; main upgrades pick them up at first boot. @@ -191,9 +274,9 @@ func (r *Repository) runMigrations() error { // Office session cost tracking extensions are declared in // initSessionWorktreeSchema's CREATE TABLE (cost_subcents, tokens_in, - // tokens_out). task_sessions.agent_profile_id existed on main as - // NOT NULL; migrateSessionsRemoveAgentExecutionID rebuilds the table - // with the column nullable and the cost columns added. + // tokens_cached_in, tokens_out). task_sessions.agent_profile_id existed + // on main as NOT NULL; migrateSessionsRemoveAgentExecutionID rebuilds the + // table with the column nullable and the cost columns added. r.migrate.Apply("workflows.is_system", `ALTER TABLE workflows ADD COLUMN is_system INTEGER DEFAULT 0`) @@ -631,6 +714,7 @@ func (r *Repository) migrateSessionsRemoveAgentExecutionID() error { task_environment_id TEXT DEFAULT '', cost_subcents INTEGER NOT NULL DEFAULT 0, tokens_in INTEGER NOT NULL DEFAULT 0, + tokens_cached_in BIGINT NOT NULL DEFAULT 0, tokens_out INTEGER NOT NULL DEFAULT 0, FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE )`, @@ -641,7 +725,7 @@ func (r *Repository) migrateSessionsRemoveAgentExecutionID() error { state, error_message, metadata, started_at, completed_at, updated_at, is_primary, is_passthrough, review_status, COALESCE(base_commit_sha, ''), COALESCE(task_environment_id, ''), - 0, 0, 0 + 0, 0, 0, 0 FROM task_sessions`, `DROP TABLE task_sessions`, `ALTER TABLE task_sessions_new RENAME TO task_sessions`, diff --git a/apps/backend/internal/task/repository/sqlite/base_schema.go b/apps/backend/internal/task/repository/sqlite/base_schema.go index a353ae81a0..4dabf04ca2 100644 --- a/apps/backend/internal/task/repository/sqlite/base_schema.go +++ b/apps/backend/internal/task/repository/sqlite/base_schema.go @@ -701,6 +701,7 @@ const sessionWorktreeSchemaDDL = ` task_environment_id TEXT DEFAULT '', cost_subcents INTEGER NOT NULL DEFAULT 0, tokens_in INTEGER NOT NULL DEFAULT 0, + tokens_cached_in BIGINT NOT NULL DEFAULT 0, tokens_out INTEGER NOT NULL DEFAULT 0, FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE ); diff --git a/apps/backend/internal/task/repository/sqlite/session.go b/apps/backend/internal/task/repository/sqlite/session.go index e500c2e2f1..f9bac52478 100644 --- a/apps/backend/internal/task/repository/sqlite/session.go +++ b/apps/backend/internal/task/repository/sqlite/session.go @@ -1447,24 +1447,27 @@ func (r *Repository) GetLastAgentMessage(ctx context.Context, sessionID string) } // IncrementTaskSessionUsage adds the given deltas to the cumulative -// tokens / cost columns on task_sessions. Used by the office cost -// subscriber after a cost event lands so the per-session totals stay -// in sync without re-summing office_cost_events. The model + DTO +// tokens / cost columns on task_sessions, including cached input tokens +// (tokens_cached_in mirrors office_cost_events.tokens_cached_in and is kept +// separate from tokens_in because it is priced differently). Used by the +// office cost subscriber after a cost event lands so the per-session totals +// stay in sync without re-summing office_cost_events. The model + DTO // don't surface these columns yet (DB-only per the office-costs // wedge); the cost explorer follow-up will expose them. func (r *Repository) IncrementTaskSessionUsage( - ctx context.Context, sessionID string, tokensIn, tokensOut, costSubcents int64, + ctx context.Context, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, ) error { if sessionID == "" { return nil } _, err := r.db.ExecContext(ctx, r.db.Rebind(` UPDATE task_sessions - SET tokens_in = COALESCE(tokens_in, 0) + ?, - tokens_out = COALESCE(tokens_out, 0) + ?, - cost_subcents = COALESCE(cost_subcents, 0) + ? + SET tokens_in = COALESCE(tokens_in, 0) + ?, + tokens_cached_in = COALESCE(tokens_cached_in, 0) + ?, + tokens_out = COALESCE(tokens_out, 0) + ?, + cost_subcents = COALESCE(cost_subcents, 0) + ? WHERE id = ? - `), tokensIn, tokensOut, costSubcents, sessionID) + `), tokensIn, tokensCachedIn, tokensOut, costSubcents, sessionID) return err } diff --git a/apps/backend/internal/task/repository/sqlite/session_test.go b/apps/backend/internal/task/repository/sqlite/session_test.go index 79ada6c6d2..434de374eb 100644 --- a/apps/backend/internal/task/repository/sqlite/session_test.go +++ b/apps/backend/internal/task/repository/sqlite/session_test.go @@ -985,10 +985,10 @@ func TestIncrementTaskSessionUsage_AccumulatesAcrossCalls(t *testing.T) { ctx := context.Background() seedForMsgTest(t, repo, "task-usage", "sess-usage", "turn-usage") - if err := repo.IncrementTaskSessionUsage(ctx, "sess-usage", 100, 200, 50); err != nil { + if err := repo.IncrementTaskSessionUsage(ctx, "sess-usage", 100, 0, 200, 50); err != nil { t.Fatalf("first increment: %v", err) } - if err := repo.IncrementTaskSessionUsage(ctx, "sess-usage", 10, 20, 5); err != nil { + if err := repo.IncrementTaskSessionUsage(ctx, "sess-usage", 10, 0, 20, 5); err != nil { t.Fatalf("second increment: %v", err) } @@ -1008,7 +1008,7 @@ func TestIncrementTaskSessionUsage_AccumulatesAcrossCalls(t *testing.T) { // missing row (subscriber may race against session creation). func TestIncrementTaskSessionUsage_UnknownSessionNoError(t *testing.T) { repo := newRepoForSessionTests(t) - if err := repo.IncrementTaskSessionUsage(context.Background(), "no-such", 1, 2, 3); err != nil { + if err := repo.IncrementTaskSessionUsage(context.Background(), "no-such", 1, 1, 2, 3); err != nil { t.Errorf("expected no error for unknown session, got %v", err) } } @@ -1017,7 +1017,7 @@ func TestIncrementTaskSessionUsage_UnknownSessionNoError(t *testing.T) { // orchestrator publishing a usage event before SessionID is set. func TestIncrementTaskSessionUsage_EmptySessionIDNoOp(t *testing.T) { repo := newRepoForSessionTests(t) - if err := repo.IncrementTaskSessionUsage(context.Background(), "", 1, 2, 3); err != nil { + if err := repo.IncrementTaskSessionUsage(context.Background(), "", 1, 1, 2, 3); err != nil { t.Errorf("empty session id should be a no-op, got %v", err) } } @@ -1083,19 +1083,19 @@ func TestMigrateSessionsAddCostColumns_BackfillsLegacySchema(t *testing.T) { 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 { + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 1, 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 { + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 1, 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 { + if err := repo.IncrementTaskSessionUsage(ctx, "sess-mig", 10, 10, 20, 30); err != nil { t.Fatalf("IncrementTaskSessionUsage after second pass: %v", err) } } diff --git a/apps/backend/internal/task/repository/sqlite/session_usage_postgres_test.go b/apps/backend/internal/task/repository/sqlite/session_usage_postgres_test.go new file mode 100644 index 0000000000..bc79c631f7 --- /dev/null +++ b/apps/backend/internal/task/repository/sqlite/session_usage_postgres_test.go @@ -0,0 +1,163 @@ +package sqlite + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/kandev/kandev/internal/testutil" +) + +// The cached-token rollup (BackfillSessionTokensCachedIn, +// IncrementTaskSessionUsage) is dialect-sensitive in a way the SQLite tests in +// session_usage_test.go and session_test.go cannot exercise: SQLite's INTEGER +// is 64-bit, so a column declared INTEGER never overflows there. On Postgres, +// INTEGER is int4 (max 2,147,483,647), and office_cost_events.tokens_cached_in +// routinely accumulates well past that for a single long-running session (the +// reported bug measured up to 98,805,109 on one already-completed task). +// Without this file the type of the new column was never exercised against +// Postgres at all. +// +// Skips unless KANDEV_TEST_POSTGRES_DSN is set. + +func seedPostgresTaskSession(t *testing.T, repo *Repository, taskID, sessionID string) { + t.Helper() + now := time.Now().UTC() + if _, err := repo.db.Exec(repo.db.Rebind(` + INSERT INTO tasks (id, workspace_id, title, created_at, updated_at) + VALUES (?, '', 'test task', ?, ?) + `), taskID, now, now); err != nil { + t.Fatalf("seed task %s: %v", taskID, err) + } + if _, err := repo.db.Exec(repo.db.Rebind(` + INSERT INTO task_sessions (id, task_id, started_at, updated_at) + VALUES (?, ?, ?, ?) + `), sessionID, taskID, now, now); err != nil { + t.Fatalf("seed session %s: %v", sessionID, err) + } +} + +// TestPostgresBackfillSessionTokensCachedIn_SumsLedger is the Postgres +// counterpart to TestBackfillSessionTokensCachedIn_SumsLedger: proves the +// correlated-subquery UPDATE and its CREATE INDEX IF NOT EXISTS both work +// against a real Postgres instance, not just SQLite. +func TestPostgresBackfillSessionTokensCachedIn_SumsLedger(t *testing.T) { + db := testutil.OpenIsolatedPostgres(t, testutil.PostgresDSNFromEnv(t)) + repo, err := NewWithDB(db, db, nil) + if err != nil { + t.Fatalf("init postgres schema: %v", err) + } + seedPostgresTaskSession(t, repo, "task-backfill-pg", "sess-backfill-pg") + createOfficeCostEventsTable(t, repo) + insertOfficeCostEvent(t, repo, "cost-1", "sess-backfill-pg", 98_805_109) + insertOfficeCostEvent(t, repo, "cost-2", "sess-backfill-pg", 1_194_891) + insertOfficeCostEvent(t, repo, "cost-other", "sess-other-pg", 555) + + if _, err := repo.BackfillSessionTokensCachedIn(context.Background()); err != nil { + t.Fatalf("backfill: %v", err) + } + + var got int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-backfill-pg", + ).Scan(&got); err != nil { + t.Fatalf("read row: %v", err) + } + if got != 100_000_000 { + t.Errorf("tokens_cached_in = %d, want 100000000 (sum of office_cost_events for this session)", got) + } +} + +// TestPostgresBackfillSessionTokensCachedIn_HandlesValuesBeyondInt32Range is +// the regression test for F8: task_sessions.tokens_cached_in was declared +// INTEGER (int4 on Postgres, max 2,147,483,647). Because the backfill is a +// single UPDATE across the whole task_sessions table, one session whose +// ledger sum exceeds the int4 ceiling made the entire statement fail with +// "integer out of range" — swallowed to a Warn by MigrateLogger.Apply, so the +// self-healing recompute silently stopped working for every session, not just +// the overflowing one. Fails on an INTEGER column, passes once the column is +// BIGINT. +func TestPostgresBackfillSessionTokensCachedIn_HandlesValuesBeyondInt32Range(t *testing.T) { + db := testutil.OpenIsolatedPostgres(t, testutil.PostgresDSNFromEnv(t)) + repo, err := NewWithDB(db, db, nil) + if err != nil { + t.Fatalf("init postgres schema: %v", err) + } + seedPostgresTaskSession(t, repo, "task-overflow-pg", "sess-overflow-pg") + seedPostgresTaskSession(t, repo, "task-normal-pg", "sess-normal-pg") + createOfficeCostEventsTable(t, repo) + + // 200 events of 14,000,000 cached tokens each = 2,800,000,000, well past + // the int32 ceiling of 2,147,483,647 - the shape the card's own live-instance + // measurement (98,805,109 on a single completed task) extrapolates to over + // a long-running session. + const perEvent = int64(14_000_000) + const eventCount = 200 + const wantOverflow = perEvent * eventCount + for i := 0; i < eventCount; i++ { + insertOfficeCostEvent(t, repo, fmt.Sprintf("cost-overflow-%d", i), "sess-overflow-pg", perEvent) + } + // A normal-sized session must still backfill correctly in the same pass - + // the whole point of F8 is that one overflowing row previously took the + // entire UPDATE down with it. + insertOfficeCostEvent(t, repo, "cost-normal", "sess-normal-pg", 1_000) + + if _, err := repo.BackfillSessionTokensCachedIn(context.Background()); err != nil { + t.Fatalf("backfill must not error on a ledger sum beyond int32 range: %v", err) + } + + var gotOverflow, gotNormal int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-overflow-pg", + ).Scan(&gotOverflow); err != nil { + t.Fatalf("read overflow row: %v", err) + } + if gotOverflow != wantOverflow { + t.Errorf("tokens_cached_in = %d, want %d (sum beyond int32 range must survive the backfill)", + gotOverflow, wantOverflow) + } + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-normal-pg", + ).Scan(&gotNormal); err != nil { + t.Fatalf("read normal row: %v", err) + } + if gotNormal != 1_000 { + t.Errorf("tokens_cached_in for the non-overflowing session = %d, want 1000 "+ + "(a sibling session's overflow must not abort the whole backfill statement)", gotNormal) + } +} + +// TestPostgresIncrementTaskSessionUsage_AccumulatesValuesBeyondInt32Range +// covers F8's other failure mode: IncrementTaskSessionUsage sets tokens_in, +// tokens_cached_in, tokens_out and cost_subcents in one UPDATE. On an INTEGER +// column that UPDATE fails once tokens_cached_in crosses the int32 ceiling, +// silently stopping tokens_in/tokens_out/cost_subcents from updating too - +// columns the card states reconcile correctly today. Fails on an INTEGER +// column, passes once the column is BIGINT. +func TestPostgresIncrementTaskSessionUsage_AccumulatesValuesBeyondInt32Range(t *testing.T) { + db := testutil.OpenIsolatedPostgres(t, testutil.PostgresDSNFromEnv(t)) + repo, err := NewWithDB(db, db, nil) + if err != nil { + t.Fatalf("init postgres schema: %v", err) + } + ctx := context.Background() + seedPostgresTaskSession(t, repo, "task-inc-overflow-pg", "sess-inc-overflow-pg") + + const beyondInt32 = int64(2_800_000_000) + if err := repo.IncrementTaskSessionUsage(ctx, "sess-inc-overflow-pg", 100, beyondInt32, 200, 50); err != nil { + t.Fatalf("IncrementTaskSessionUsage must not error on a cached-token delta beyond int32 range: %v", err) + } + + var tokensIn, tokensCachedIn, tokensOut, costSubcents int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_in, tokens_cached_in, tokens_out, cost_subcents FROM task_sessions WHERE id = ?`), + "sess-inc-overflow-pg").Scan(&tokensIn, &tokensCachedIn, &tokensOut, &costSubcents); err != nil { + t.Fatalf("read row: %v", err) + } + if tokensIn != 100 || tokensCachedIn != beyondInt32 || tokensOut != 200 || costSubcents != 50 { + t.Errorf("totals = (%d,%d,%d,%d), want (100,%d,200,50) - a cached-token delta beyond int32 "+ + "range must not roll back tokens_in/tokens_out/cost_subcents in the same statement", + tokensIn, tokensCachedIn, tokensOut, costSubcents, beyondInt32) + } +} diff --git a/apps/backend/internal/task/repository/sqlite/session_usage_test.go b/apps/backend/internal/task/repository/sqlite/session_usage_test.go new file mode 100644 index 0000000000..67d5524b7e --- /dev/null +++ b/apps/backend/internal/task/repository/sqlite/session_usage_test.go @@ -0,0 +1,235 @@ +package sqlite + +import ( + "context" + "testing" + "time" +) + +// TestIncrementTaskSessionUsage_AccumulatesCachedInAcrossCalls confirms +// tokens_cached_in compounds the same way tokens_in/tokens_out/cost_subcents +// already do. Before the fix this column didn't exist and the parameter +// couldn't be threaded at all. +func TestIncrementTaskSessionUsage_AccumulatesCachedInAcrossCalls(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedForMsgTest(t, repo, "task-cached-usage", "sess-cached-usage", "turn-cached-usage") + + if err := repo.IncrementTaskSessionUsage(ctx, "sess-cached-usage", 100, 50_000_000, 200, 50); err != nil { + t.Fatalf("first increment: %v", err) + } + if err := repo.IncrementTaskSessionUsage(ctx, "sess-cached-usage", 10, 2_000, 20, 5); err != nil { + t.Fatalf("second increment: %v", err) + } + + var tokensIn, tokensCachedIn, tokensOut, costSubcents int64 + err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_in, tokens_cached_in, tokens_out, cost_subcents FROM task_sessions WHERE id = ?`), + "sess-cached-usage").Scan(&tokensIn, &tokensCachedIn, &tokensOut, &costSubcents) + if err != nil { + t.Fatalf("read row: %v", err) + } + if tokensIn != 110 || tokensCachedIn != 50_002_000 || tokensOut != 220 || costSubcents != 55 { + t.Errorf("totals = (%d,%d,%d,%d), want (110,50002000,220,55)", + tokensIn, tokensCachedIn, tokensOut, costSubcents) + } +} + +// createOfficeCostEventsTable creates a minimal standalone copy of +// office_cost_events (owned by internal/office/repository/sqlite) so this +// package's migration tests can exercise the guarded backfill without +// importing the office package (which would create an import cycle back +// into internal/task). Column set mirrors office/repository/sqlite/base.go +// createCostTables. +func createOfficeCostEventsTable(t *testing.T, repo *Repository) { + t.Helper() + if _, err := repo.db.Exec(` + CREATE TABLE IF NOT EXISTS office_cost_events ( + id TEXT PRIMARY KEY, + session_id TEXT DEFAULT '', + task_id TEXT DEFAULT '', + agent_profile_id TEXT DEFAULT '', + project_id TEXT DEFAULT '', + model TEXT DEFAULT '', + provider TEXT DEFAULT '', + tokens_in INTEGER DEFAULT 0, + tokens_cached_in INTEGER DEFAULT 0, + tokens_out INTEGER DEFAULT 0, + cost_subcents INTEGER NOT NULL DEFAULT 0, + estimated INTEGER NOT NULL DEFAULT 0, + occurred_at TIMESTAMP NOT NULL, + created_at TIMESTAMP NOT NULL + )`); err != nil { + t.Fatalf("create office_cost_events: %v", err) + } +} + +func insertOfficeCostEvent(t *testing.T, repo *Repository, id, sessionID string, tokensCachedIn int64) { + t.Helper() + now := time.Now().UTC() + _, err := repo.db.Exec(repo.db.Rebind(` + INSERT INTO office_cost_events (id, session_id, tokens_cached_in, occurred_at, created_at) + VALUES (?, ?, ?, ?, ?) + `), id, sessionID, tokensCachedIn, now, now) + if err != nil { + t.Fatalf("insert office_cost_event %s: %v", id, err) + } +} + +// TestBackfillSessionTokensCachedIn_SumsLedger proves the backfill invariant +// directly: task_sessions.tokens_cached_in must equal the sum of +// office_cost_events.tokens_cached_in for that session. This is the +// reconciliation the card asked for, exercised at the SQL layer. +func TestBackfillSessionTokensCachedIn_SumsLedger(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedForMsgTest(t, repo, "task-backfill", "sess-backfill", "turn-backfill") + createOfficeCostEventsTable(t, repo) + insertOfficeCostEvent(t, repo, "cost-1", "sess-backfill", 98_805_109) + insertOfficeCostEvent(t, repo, "cost-2", "sess-backfill", 1_194_891) + // A row for a different session must not leak into this session's total. + insertOfficeCostEvent(t, repo, "cost-other", "sess-other", 555) + + if _, err := repo.BackfillSessionTokensCachedIn(ctx); err != nil { + t.Fatalf("backfill: %v", err) + } + + var got int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-backfill", + ).Scan(&got); err != nil { + t.Fatalf("read row: %v", err) + } + if got != 100_000_000 { + t.Errorf("tokens_cached_in = %d, want 100000000 (sum of office_cost_events for this session)", got) + } + + // Idempotent: re-running against an unchanged ledger must not + // double-count the already-correct total. + if _, err := repo.BackfillSessionTokensCachedIn(ctx); err != nil { + t.Fatalf("second backfill: %v", err) + } + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-backfill", + ).Scan(&got); err != nil { + t.Fatalf("read row after second pass: %v", err) + } + if got != 100_000_000 { + t.Errorf("tokens_cached_in after second backfill pass = %d, want unchanged 100000000", got) + } +} + +// TestBackfillSessionTokensCachedIn_CorrectsDriftedRollup covers the case an +// unconditional recompute exists to repair: a session whose rollup went out +// of sync with the ledger without ever erroring. IncrementTaskSessionUsage +// is a best-effort `UPDATE ... WHERE id = ?` that silently matches zero rows +// if the session row doesn't exist yet (see +// TestIncrementTaskSessionUsage_UnknownSessionNoError), and the office +// subscriber only logs a Warn when that write fails while the ledger insert +// has already committed (event_subscribers.go handlePromptUsage). A backfill +// gated on "already nonzero" would treat this drifted-but-nonzero row as +// already correct and skip it forever; the unconditional recompute must +// instead overwrite it with the true ledger sum. +func TestBackfillSessionTokensCachedIn_CorrectsDriftedRollup(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedForMsgTest(t, repo, "task-drifted", "sess-drifted", "turn-drifted") + createOfficeCostEventsTable(t, repo) + insertOfficeCostEvent(t, repo, "cost-1", "sess-drifted", 98_805_109) + insertOfficeCostEvent(t, repo, "cost-2", "sess-drifted", 1_194_891) + + // Simulate the drift directly: the rollup already holds a nonzero but + // wrong value, as it would after a missed increment. + if err := repo.IncrementTaskSessionUsage(ctx, "sess-drifted", 0, 2_000, 0, 0); err != nil { + t.Fatalf("seed drifted rollup: %v", err) + } + + var before int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-drifted", + ).Scan(&before); err != nil { + t.Fatalf("read row before backfill: %v", err) + } + if before != 2_000 { + t.Fatalf("test setup bug: rollup = %d, want the drifted seed value 2000", before) + } + + if _, err := repo.BackfillSessionTokensCachedIn(ctx); err != nil { + t.Fatalf("backfill: %v", err) + } + + var got int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-drifted", + ).Scan(&got); err != nil { + t.Fatalf("read row after backfill: %v", err) + } + if got != 100_000_000 { + t.Errorf("tokens_cached_in after backfill = %d, want 100000000 (the drifted 2000 must be "+ + "overwritten with the full ledger sum, not skipped as already-nonzero)", got) + } +} + +// TestBackfillSessionTokensCachedIn_ZeroesRollupWithNoMatchingLedgerRows pins +// the behaviour when the ledger table exists but has no rows for a session +// whose rollup is already nonzero (e.g. the ledger rows for that session were +// deleted separately from the session itself - see workspace_deletion.go, +// which deletes office_cost_events before the task-side workspace in two +// separate transactions). The unconditional recompute is a COALESCE(...,0), +// so it deliberately zeroes tokens_cached_in in this case rather than leaving +// the stale value: the rollup is supposed to always equal the ledger sum for +// that session, and an empty ledger sums to zero. +func TestBackfillSessionTokensCachedIn_ZeroesRollupWithNoMatchingLedgerRows(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedForMsgTest(t, repo, "task-orphaned", "sess-orphaned", "turn-orphaned") + createOfficeCostEventsTable(t, repo) + + // Rollup already holds a nonzero value, but no office_cost_events rows + // exist for this session (simulates the ledger having been deleted). + if err := repo.IncrementTaskSessionUsage(ctx, "sess-orphaned", 0, 12_345, 0, 0); err != nil { + t.Fatalf("seed rollup: %v", err) + } + + if _, err := repo.BackfillSessionTokensCachedIn(ctx); err != nil { + t.Fatalf("backfill: %v", err) + } + + var got int64 + if err := repo.ro.QueryRowx(repo.ro.Rebind( + `SELECT tokens_cached_in FROM task_sessions WHERE id = ?`), "sess-orphaned", + ).Scan(&got); err != nil { + t.Fatalf("read row: %v", err) + } + if got != 0 { + t.Errorf("tokens_cached_in = %d, want 0 (no matching ledger rows sums to zero)", got) + } +} + +// TestBackfillSessionTokensCachedIn_SkipsAlreadyCorrectUntouchedSessions pins +// the boot-cost optimization: a session with no office_cost_events rows and a +// tokens_cached_in already at its default of 0 must be excluded by the +// UPDATE's WHERE clause, not just happen to compute the same value. Proven +// via the reported affected-row count (matches SQL UPDATE semantics on both +// SQLite and Postgres: rows matched by WHERE, not rows whose value changed), +// comparing a session the WHERE clause must match (has ledger rows) against +// one it must not (no ledger rows, untouched default). +func TestBackfillSessionTokensCachedIn_SkipsAlreadyCorrectUntouchedSessions(t *testing.T) { + repo := newRepoForSessionTests(t) + ctx := context.Background() + seedForMsgTest(t, repo, "task-untouched", "sess-untouched", "turn-untouched") + seedForMsgTest(t, repo, "task-with-ledger", "sess-with-ledger", "turn-with-ledger") + createOfficeCostEventsTable(t, repo) + insertOfficeCostEvent(t, repo, "cost-1", "sess-with-ledger", 500) + // sess-untouched: no office_cost_events rows, tokens_cached_in stays at + // its untouched default of 0. + + affected, err := repo.BackfillSessionTokensCachedIn(ctx) + if err != nil { + t.Fatalf("backfill: %v", err) + } + if affected != 1 { + t.Errorf("rows affected = %d, want 1 (only sess-with-ledger should match the WHERE clause; "+ + "sess-untouched has no ledger rows and was already 0)", affected) + } +} diff --git a/docs/specs/office/costs.md b/docs/specs/office/costs.md index 19f26d32b5..1a26eac02f 100644 --- a/docs/specs/office/costs.md +++ b/docs/specs/office/costs.md @@ -68,7 +68,7 @@ Multiple policies can apply to the same scope (e.g. monthly + total). ### `TaskSession` additions -`cost_subcents`, `tokens_in`, `tokens_out` incrementally updated as cost events arrive, providing quick per-session totals without scanning `office_cost_events`. +`cost_subcents`, `tokens_in`, `tokens_out`, `tokens_cached_in` incrementally updated as cost events arrive, providing quick per-session totals without scanning `office_cost_events`. `tokens_cached_in` is DB-only today (no model/DTO/API/web reader) — it exists so the rollup reconciles column-for-column against the ledger; a future cost-explorer surface will expose it. ### Per-agent budget @@ -187,7 +187,7 @@ There is no per-field permission model. Conformance tests should assert that cos - `office_cost_events` rows (full history, never trimmed). Disk-runner rows keyed by `(session_id, provider_event_id)` survive re-ingestion without duplicating; the wire-side `estimated=true` rows for codex are deleted once the matching aggregate row lands. - `office_budget_policies` rows. -- `TaskSession.cost_subcents` / `tokens_in` / `tokens_out` running totals, kept in sync with `office_cost_events` so per-session totals are correct without a re-scan. +- `TaskSession.cost_subcents` / `tokens_in` / `tokens_out` / `tokens_cached_in` running totals, kept in sync with `office_cost_events` so per-session totals are correct without a re-scan. - Per-agent `budget_monthly_cents` stored on the agent instance row. - Per-model pricing overrides stored in workspace settings. - The on-disk `models.dev` cache at `/cache/models-dev.json`. Recovery on next boot: the in-memory pricing map is empty until first query; queries fall back to the on-disk file when the background refresh has not yet completed.