test: cover Codex usage backfill execution

Co-authored-by: multica-agent <github@multica.ai>
2026-07-17 23:29:01 +02:00 · 2026-06-18 12:19:31 +08:00
parent 099dda2051
commit cbb7adc0ea
3 changed files with 425 additions and 0 deletions
--- a/2
+++ b/2
@@ -20,6 +20,7 @@ RUN cd server && CGO_ENABLED=0 go build -ldflags "-s -w -X main.version=${VERSIO
 RUN cd server && CGO_ENABLED=0 go build -ldflags "-s -w -X main.version=${VERSION} -X main.commit=${COMMIT} -X main.date=${DATE}" -o bin/multica ./cmd/multica
 RUN cd server && CGO_ENABLED=0 go build -ldflags "-s -w" -o bin/migrate ./cmd/migrate
 RUN cd server && CGO_ENABLED=0 go build -ldflags "-s -w" -o bin/backfill_task_usage_hourly ./cmd/backfill_task_usage_hourly
+RUN cd server && CGO_ENABLED=0 go build -ldflags "-s -w" -o bin/backfill_codex_usage_cache ./cmd/backfill_codex_usage_cache

 # --- Runtime stage ---
 FROM alpine:3.21
@@ -32,6 +33,7 @@ COPY --from=builder /src/server/bin/server .
 COPY --from=builder /src/server/bin/multica .
 COPY --from=builder /src/server/bin/migrate .
 COPY --from=builder /src/server/bin/backfill_task_usage_hourly .
+COPY --from=builder /src/server/bin/backfill_codex_usage_cache .
 COPY server/migrations/ ./migrations/
 COPY docker/entrypoint.sh .
 RUN sed -i 's/\r$//' entrypoint.sh && chmod +x entrypoint.sh
--- a/docs/codex-usage-cache-backfill.md
+++ b/docs/codex-usage-cache-backfill.md
@@ -0,0 +1,91 @@
+# Codex Usage Cache Backfill
+
+This runbook describes the one-time hosted data repair for Codex usage rows
+created before cached input was normalized at ingestion time.
+
+Do not run this as an automatic database migration. The write step needs an
+operator-selected cutoff, dry-run review, and an explicit execute command.
+
+## When To Run
+
+Run this only after the backend image containing `backfill_codex_usage_cache`
+has been deployed, and only for databases that need historical Codex usage
+correction.
+
+Use the actual hosted deployment time of PR #4083 as `--cutoff`. Do not use the
+PR merge time unless it is also the real production cutover time.
+
+## Execution Model
+
+Run the command from the released backend image as a one-time operator job, such
+as a Kubernetes Job with the normal backend database secret and network access.
+Override the command to execute `./backfill_codex_usage_cache`.
+
+The command defaults to dry-run. It mutates data only when `--execute` is passed.
+
+## Dry Run
+
+First run:
+
+```bash
+./backfill_codex_usage_cache --cutoff <RFC3339_DEPLOY_TIME>
+```
+
+Optionally limit scope while validating:
+
+```bash
+./backfill_codex_usage_cache --cutoff <RFC3339_DEPLOY_TIME> --workspace-id <workspace-uuid>
+```
+
+Review the per-workspace/date output:
+
+- `rows`
+- `input_before`
+- `input_after`
+- `input_tokens_removed`
+- `clamped_rows`
+
+Proceed only if the totals match the expected overcount shape.
+
+## Execute
+
+After dry-run review:
+
+```bash
+./backfill_codex_usage_cache --cutoff <RFC3339_DEPLOY_TIME> --execute
+```
+
+For large datasets, throttle writes:
+
+```bash
+./backfill_codex_usage_cache \
+  --cutoff <RFC3339_DEPLOY_TIME> \
+  --execute \
+  --batch-size 500 \
+  --sleep-between-batches 1s
+```
+
+By default, execution rebuilds affected hourly rollups by calling
+`rollup_task_usage_hourly_window(...)` for the database update window. Leave
+`--rebuild-rollup=true` unless an operator intentionally plans a separate rollup
+rebuild.
+
+## Verification
+
+After execution, run the dry-run command again with the same cutoff and scope.
+Eligible rows should be zero.
+
+Then verify Usage / Runtime dashboard periods that were previously inflated.
+
+## Safety Boundaries
+
+The command updates only rows that match all of these conditions:
+
+- `provider = 'codex'`
+- `cache_read_tokens > 0`
+- `input_tokens > 0`
+- `COALESCE(updated_at, created_at) < --cutoff`
+- optional `--workspace-id` match
+
+Rows without persisted `cache_read_tokens` are intentionally ignored because the
+current database cannot accurately reconstruct cached input for them.
--- a/server/cmd/backfill_codex_usage_cache/integration_test.go
+++ b/server/cmd/backfill_codex_usage_cache/integration_test.go
@@ -0,0 +1,332 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"sort"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/multica-ai/multica/server/internal/migrations"
+)
+
+func TestExecuteBackfillUpdatesOnlyEligibleCodexRowsAndRebuildsRollup(t *testing.T) {
+	adminURL := os.Getenv("DATABASE_URL")
+	if adminURL == "" {
+		adminURL = "postgres://multica:multica@localhost:5432/multica?sslmode=disable"
+	}
+
+	ctx := context.Background()
+	if !testDatabaseReachable(ctx, adminURL) {
+		t.Skip("integration test requires Postgres at DATABASE_URL")
+	}
+
+	tmpDB := fmt.Sprintf("multica_codex_usage_backfill_%d", time.Now().UnixNano())
+	if err := testCreateDatabase(ctx, adminURL, tmpDB); err != nil {
+		t.Fatalf("create temp database %s: %v", tmpDB, err)
+	}
+	t.Cleanup(func() {
+		if err := testDropDatabase(context.Background(), adminURL, tmpDB); err != nil {
+			t.Logf("drop temp database %s: %v", tmpDB, err)
+		}
+	})
+
+	pool, err := pgxpool.New(ctx, testReplaceDatabase(adminURL, tmpDB))
+	if err != nil {
+		t.Fatalf("connect to temp database: %v", err)
+	}
+	t.Cleanup(pool.Close)
+
+	if err := testApplyMigrationsUpTo(ctx, pool, "102_task_usage_hourly_pipeline"); err != nil {
+		t.Fatalf("apply migrations to 102: %v", err)
+	}
+
+	cutoff := time.Date(2026, 6, 13, 8, 0, 0, 0, time.UTC)
+	createdAt := cutoff.Add(-48 * time.Hour)
+	oldUpdatedAt := cutoff.Add(-24 * time.Hour)
+	newUpdatedAt := cutoff.Add(24 * time.Hour)
+
+	workspaceID, _, _, targetTaskID := seedBackfillTaskUsageFixture(t, ctx, pool, "target")
+	_, _, _, otherWorkspaceTaskID := seedBackfillTaskUsageFixture(t, ctx, pool, "other")
+
+	targetID := seedTaskUsageRow(t, ctx, pool, targetTaskID, "codex", "gpt-5-codex", 1000, 50, 300, createdAt, oldUpdatedAt)
+	postCutoffID := seedTaskUsageRow(t, ctx, pool, targetTaskID, "codex", "gpt-5-codex-new", 900, 50, 200, createdAt, newUpdatedAt)
+	nonCodexID := seedTaskUsageRow(t, ctx, pool, targetTaskID, "claude", "claude-test", 800, 50, 200, createdAt, oldUpdatedAt)
+	zeroCacheID := seedTaskUsageRow(t, ctx, pool, targetTaskID, "codex", "gpt-5-codex-no-cache", 700, 50, 0, createdAt, oldUpdatedAt)
+	otherWorkspaceID := seedTaskUsageRow(t, ctx, pool, otherWorkspaceTaskID, "codex", "gpt-5-codex-other-workspace", 600, 50, 100, createdAt, oldUpdatedAt)
+
+	cfg := config{
+		cutoff:      cutoff,
+		workspaceID: workspaceID,
+		batchSize:   10,
+	}
+
+	summary, total, err := loadDryRunSummary(ctx, pool, cfg)
+	if err != nil {
+		t.Fatalf("load dry-run summary: %v", err)
+	}
+	if len(summary) != 1 {
+		t.Fatalf("summary rows = %d, want 1: %#v", len(summary), summary)
+	}
+	if total.Rows != 1 || total.InputBefore != 1000 || total.InputAfter != 700 || total.Overcount != 300 {
+		t.Fatalf("unexpected dry-run total: %+v", total)
+	}
+	assertTaskUsageInput(t, ctx, pool, targetID, 1000)
+
+	updateStartedAt, err := databaseClock(ctx, pool)
+	if err != nil {
+		t.Fatalf("read update start clock: %v", err)
+	}
+	updatedRows, removedTokens, err := executeBackfill(ctx, pool, cfg)
+	if err != nil {
+		t.Fatalf("execute backfill: %v", err)
+	}
+	if updatedRows != 1 || removedTokens != 300 {
+		t.Fatalf("updatedRows=%d removedTokens=%d, want 1/300", updatedRows, removedTokens)
+	}
+	updateFinishedAt, err := databaseClock(ctx, pool)
+	if err != nil {
+		t.Fatalf("read update finish clock: %v", err)
+	}
+	rollupFrom, rollupTo := rollupWindow(updateStartedAt, updateFinishedAt)
+	var rollupRows int64
+	if err := pool.QueryRow(ctx,
+		`SELECT rollup_task_usage_hourly_window($1::timestamptz, $2::timestamptz)`,
+		rollupFrom, rollupTo,
+	).Scan(&rollupRows); err != nil {
+		t.Fatalf("rebuild hourly rollup: %v", err)
+	}
+	if rollupRows == 0 {
+		t.Fatalf("expected hourly rollup rows to be touched")
+	}
+
+	assertTaskUsageInput(t, ctx, pool, targetID, 700)
+	assertTaskUsageInput(t, ctx, pool, postCutoffID, 900)
+	assertTaskUsageInput(t, ctx, pool, nonCodexID, 800)
+	assertTaskUsageInput(t, ctx, pool, zeroCacheID, 700)
+	assertTaskUsageInput(t, ctx, pool, otherWorkspaceID, 600)
+
+	var targetUpdatedAt time.Time
+	if err := pool.QueryRow(ctx, `SELECT updated_at FROM task_usage WHERE id = $1`, targetID).Scan(&targetUpdatedAt); err != nil {
+		t.Fatalf("read target updated_at: %v", err)
+	}
+	if !targetUpdatedAt.After(cutoff) {
+		t.Fatalf("target updated_at = %s, want after cutoff %s", targetUpdatedAt, cutoff)
+	}
+
+	updatedRows, removedTokens, err = executeBackfill(ctx, pool, cfg)
+	if err != nil {
+		t.Fatalf("execute backfill second time: %v", err)
+	}
+	if updatedRows != 0 || removedTokens != 0 {
+		t.Fatalf("second execution updatedRows=%d removedTokens=%d, want 0/0", updatedRows, removedTokens)
+	}
+
+	var hourlyInput, hourlyCache int64
+	if err := pool.QueryRow(ctx, `
+		SELECT COALESCE(SUM(input_tokens), 0)::bigint,
+		       COALESCE(SUM(cache_read_tokens), 0)::bigint
+		  FROM task_usage_hourly
+		 WHERE workspace_id = $1
+		   AND provider = 'codex'
+		   AND model = 'gpt-5-codex'
+	`, workspaceID).Scan(&hourlyInput, &hourlyCache); err != nil {
+		t.Fatalf("read hourly rollup: %v", err)
+	}
+	if hourlyInput != 700 || hourlyCache != 300 {
+		t.Fatalf("hourly rollup input/cache = %d/%d, want 700/300", hourlyInput, hourlyCache)
+	}
+}
+
+func seedTaskUsageRow(
+	t *testing.T,
+	ctx context.Context,
+	pool *pgxpool.Pool,
+	taskID, provider, model string,
+	input, output, cacheRead int64,
+	createdAt, updatedAt time.Time,
+) string {
+	t.Helper()
+
+	var id string
+	if err := pool.QueryRow(ctx, `
+		INSERT INTO task_usage (
+			task_id, provider, model,
+			input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
+			created_at, updated_at
+		) VALUES ($1, $2, $3, $4, $5, $6, 0, $7, $8)
+		RETURNING id
+	`, taskID, provider, model, input, output, cacheRead, createdAt, updatedAt).Scan(&id); err != nil {
+		t.Fatalf("seed task_usage row %s/%s: %v", provider, model, err)
+	}
+	return id
+}
+
+func assertTaskUsageInput(t *testing.T, ctx context.Context, pool *pgxpool.Pool, id string, want int64) {
+	t.Helper()
+	var got int64
+	if err := pool.QueryRow(ctx, `SELECT input_tokens FROM task_usage WHERE id = $1`, id).Scan(&got); err != nil {
+		t.Fatalf("read task_usage %s: %v", id, err)
+	}
+	if got != want {
+		t.Fatalf("task_usage %s input_tokens = %d, want %d", id, got, want)
+	}
+}
+
+func seedBackfillTaskUsageFixture(t *testing.T, ctx context.Context, pool *pgxpool.Pool, suffix string) (string, string, string, string) {
+	t.Helper()
+
+	var wsID, runtimeID, agentID, taskID string
+	if err := pool.QueryRow(ctx, `
+		INSERT INTO workspace (name, slug)
+		VALUES ($1, $2)
+		RETURNING id
+	`, "codex backfill "+suffix, "codex-backfill-"+suffix).Scan(&wsID); err != nil {
+		t.Fatalf("seed workspace: %v", err)
+	}
+	if err := pool.QueryRow(ctx, `
+		INSERT INTO agent_runtime (
+			workspace_id, daemon_id, name, runtime_mode, provider, status,
+			device_info, metadata, last_seen_at
+		)
+		VALUES ($1, NULL, $2, 'cloud', 'codex', 'online',
+		        '{}'::jsonb, '{}'::jsonb, now())
+		RETURNING id
+	`, wsID, "codex backfill runtime "+suffix).Scan(&runtimeID); err != nil {
+		t.Fatalf("seed runtime: %v", err)
+	}
+	if err := pool.QueryRow(ctx, `
+		INSERT INTO agent (
+			workspace_id, name, description, runtime_mode, runtime_config,
+			runtime_id, visibility, max_concurrent_tasks
+		)
+		VALUES ($1, $2, '', 'cloud', '{}'::jsonb, $3, 'workspace', 1)
+		RETURNING id
+	`, wsID, "codex backfill agent "+suffix, runtimeID).Scan(&agentID); err != nil {
+		t.Fatalf("seed agent: %v", err)
+	}
+	if err := pool.QueryRow(ctx, `
+		INSERT INTO agent_task_queue (
+			agent_id, runtime_id, status, payload
+		)
+		VALUES ($1, $2, 'queued', '{}'::jsonb)
+		RETURNING id
+	`, agentID, runtimeID).Scan(&taskID); err != nil {
+		altErr := pool.QueryRow(ctx, `
+			INSERT INTO agent_task_queue (agent_id, runtime_id)
+			VALUES ($1, $2)
+			RETURNING id
+		`, agentID, runtimeID).Scan(&taskID)
+		if altErr != nil {
+			t.Fatalf("seed agent_task_queue: %v / %v", err, altErr)
+		}
+	}
+	return wsID, runtimeID, agentID, taskID
+}
+
+func testResolveMigrationsDir() (string, error) {
+	_, here, _, ok := runtime.Caller(0)
+	if !ok {
+		return "", fmt.Errorf("runtime.Caller(0) failed")
+	}
+	dir := filepath.Join(filepath.Dir(here), "..", "..", "migrations")
+	dir = filepath.Clean(dir)
+	if info, err := os.Stat(dir); err != nil || !info.IsDir() {
+		return "", fmt.Errorf("migrations dir not at %s", dir)
+	}
+	return dir, nil
+}
+
+func testDatabaseReachable(ctx context.Context, url string) bool {
+	pool, err := pgxpool.New(ctx, url)
+	if err != nil {
+		return false
+	}
+	defer pool.Close()
+	return pool.Ping(ctx) == nil
+}
+
+func testCreateDatabase(ctx context.Context, adminURL, name string) error {
+	conn, err := pgx.Connect(ctx, adminURL)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	if _, err := conn.Exec(ctx, fmt.Sprintf(`CREATE DATABASE %q`, name)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func testDropDatabase(ctx context.Context, adminURL, name string) error {
+	conn, err := pgx.Connect(ctx, adminURL)
+	if err != nil {
+		return err
+	}
+	defer conn.Close(ctx)
+	if _, err := conn.Exec(ctx, fmt.Sprintf(`DROP DATABASE IF EXISTS %q WITH (FORCE)`, name)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func testReplaceDatabase(url, name string) string {
+	idx := strings.LastIndex(url, "/")
+	if idx < 0 {
+		return url
+	}
+	rest := url[idx+1:]
+	q := strings.Index(rest, "?")
+	if q < 0 {
+		return url[:idx+1] + name
+	}
+	return url[:idx+1] + name + rest[q:]
+}
+
+func testApplyMigrationsUpTo(ctx context.Context, pool *pgxpool.Pool, lastVersion string) error {
+	dir, err := testResolveMigrationsDir()
+	if err != nil {
+		return err
+	}
+	files, err := filepath.Glob(filepath.Join(dir, "*.up.sql"))
+	if err != nil {
+		return err
+	}
+	sort.Strings(files)
+
+	if _, err := pool.Exec(ctx, `
+		CREATE TABLE IF NOT EXISTS schema_migrations (
+			version TEXT PRIMARY KEY,
+			applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
+		)
+	`); err != nil {
+		return err
+	}
+
+	for _, f := range files {
+		v := migrations.ExtractVersion(f)
+		sql, err := os.ReadFile(f)
+		if err != nil {
+			return err
+		}
+		if _, err := pool.Exec(ctx, string(sql)); err != nil {
+			return fmt.Errorf("apply %s: %w", v, err)
+		}
+		if _, err := pool.Exec(ctx,
+			`INSERT INTO schema_migrations (version) VALUES ($1) ON CONFLICT DO NOTHING`,
+			v); err != nil {
+			return err
+		}
+		if v == lastVersion {
+			return nil
+		}
+	}
+	return fmt.Errorf("migration %q not found", lastVersion)
+}