From d74aed1e07d32a16ff3be770cd3e4e7f0bcf8a1a Mon Sep 17 00:00:00 2001 From: Multica Eve Date: Thu, 16 Jul 2026 11:21:20 +0800 Subject: [PATCH] MUL-4302 Remove legacy attribution invariant exemption (#5504) * fix(attribution): remove legacy invariant exemption Co-authored-by: multica-agent * test(attribution): enforce strict legacy row rejection Co-authored-by: multica-agent --------- Co-authored-by: Eve Co-authored-by: multica-agent --- .../attribution_constraint_migration_test.go | 156 ++++++++++++++++++ .../service/attribution_stamp_test.go | 41 ++--- ...ask_attribution_strict_constraint.down.sql | 2 + ..._task_attribution_strict_constraint.up.sql | 23 +++ ...bution_strict_constraint_validate.down.sql | 15 ++ ...ribution_strict_constraint_validate.up.sql | 9 + ...tribution_legacy_exemption_remove.down.sql | 18 ++ ...attribution_legacy_exemption_remove.up.sql | 10 ++ 8 files changed, 246 insertions(+), 28 deletions(-) create mode 100644 server/internal/migrations/attribution_constraint_migration_test.go create mode 100644 server/migrations/197_agent_task_attribution_strict_constraint.down.sql create mode 100644 server/migrations/197_agent_task_attribution_strict_constraint.up.sql create mode 100644 server/migrations/198_agent_task_attribution_strict_constraint_validate.down.sql create mode 100644 server/migrations/198_agent_task_attribution_strict_constraint_validate.up.sql create mode 100644 server/migrations/199_agent_task_attribution_legacy_exemption_remove.down.sql create mode 100644 server/migrations/199_agent_task_attribution_legacy_exemption_remove.up.sql diff --git a/server/internal/migrations/attribution_constraint_migration_test.go b/server/internal/migrations/attribution_constraint_migration_test.go new file mode 100644 index 000000000..3113de4aa --- /dev/null +++ b/server/internal/migrations/attribution_constraint_migration_test.go @@ -0,0 +1,156 @@ +package migrations + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" +) + +const ( + originatorID = "00000000-0000-0000-0000-000000000001" + otherUserID = "00000000-0000-0000-0000-000000000002" +) + +func TestAttributionStrictConstraintMigrations(t *testing.T) { + dbURL := os.Getenv("DATABASE_URL") + if dbURL == "" { + t.Skip("integration test requires Postgres at DATABASE_URL") + } + + ctx := context.Background() + pool, err := pgxpool.New(ctx, dbURL) + if err != nil { + t.Fatalf("connect to Postgres: %v", err) + } + defer pool.Close() + + conn, err := pool.Acquire(ctx) + if err != nil { + t.Fatalf("acquire Postgres connection: %v", err) + } + defer conn.Release() + + if _, err := conn.Exec(ctx, ` + CREATE TEMP TABLE agent_task_queue ( + id UUID PRIMARY KEY, + originator_source TEXT NULL, + originator_user_id UUID NULL, + accountable_user_id UUID NULL + ) + `); err != nil { + t.Fatalf("create temporary queue table: %v", err) + } + + applyMigrationFile(t, ctx, conn.Conn(), "190_agent_task_attribution_invariant_check.up.sql") + + // Migration 190 allows this pre-backfill shape solely because its source is + // NULL. The strict shadow must refuse to validate until the row is repaired. + if _, err := conn.Exec(ctx, ` + INSERT INTO agent_task_queue (id, originator_user_id) + VALUES ('00000000-0000-0000-0000-000000000010', $1) + `, originatorID); err != nil { + t.Fatalf("insert transitional legacy row: %v", err) + } + + applyMigrationFile(t, ctx, conn.Conn(), "197_agent_task_attribution_strict_constraint.up.sql") + validateSQL := readMigrationFile(t, "198_agent_task_attribution_strict_constraint_validate.up.sql") + if _, err := conn.Exec(ctx, validateSQL); !isCheckViolation(err) { + t.Fatalf("validate strict constraint with legacy row: got %v, want check violation", err) + } + + if _, err := conn.Exec(ctx, ` + UPDATE agent_task_queue + SET accountable_user_id = originator_user_id, + originator_source = 'backfill' + WHERE id = '00000000-0000-0000-0000-000000000010' + `); err != nil { + t.Fatalf("backfill transitional row: %v", err) + } + + applyMigrationFile(t, ctx, conn.Conn(), "198_agent_task_attribution_strict_constraint_validate.up.sql") + applyMigrationFile(t, ctx, conn.Conn(), "199_agent_task_attribution_legacy_exemption_remove.up.sql") + + var validated bool + var definition string + if err := conn.QueryRow(ctx, ` + SELECT convalidated, pg_get_constraintdef(oid) + FROM pg_constraint + WHERE conrelid = 'agent_task_queue'::regclass + AND conname = 'agent_task_queue_accountable_matches_originator' + `).Scan(&validated, &definition); err != nil { + t.Fatalf("read final constraint: %v", err) + } + if !validated { + t.Fatal("final attribution constraint is not validated") + } + if strings.Contains(definition, "originator_source") { + t.Fatalf("final constraint still contains legacy exemption: %s", definition) + } + + assertInsertCheckViolation(t, ctx, conn.Conn(), ` + INSERT INTO agent_task_queue (id, originator_source, originator_user_id) + VALUES ('00000000-0000-0000-0000-000000000011', NULL, $1) + `, originatorID) + assertInsertCheckViolation(t, ctx, conn.Conn(), ` + INSERT INTO agent_task_queue ( + id, originator_source, originator_user_id, accountable_user_id + ) VALUES ('00000000-0000-0000-0000-000000000012', NULL, $1, $2) + `, originatorID, otherUserID) + + if _, err := conn.Exec(ctx, ` + INSERT INTO agent_task_queue ( + id, originator_source, originator_user_id, accountable_user_id + ) VALUES ('00000000-0000-0000-0000-000000000013', NULL, $1, $1) + `, originatorID); err != nil { + t.Fatalf("insert matching attribution under strict constraint: %v", err) + } + + applyMigrationFile(t, ctx, conn.Conn(), "199_agent_task_attribution_legacy_exemption_remove.down.sql") + applyMigrationFile(t, ctx, conn.Conn(), "198_agent_task_attribution_strict_constraint_validate.down.sql") + applyMigrationFile(t, ctx, conn.Conn(), "197_agent_task_attribution_strict_constraint.down.sql") + + if _, err := conn.Exec(ctx, ` + INSERT INTO agent_task_queue (id, originator_user_id) + VALUES ('00000000-0000-0000-0000-000000000014', $1) + `, originatorID); err != nil { + t.Fatalf("insert legacy row after rollback restored exemption: %v", err) + } +} + +func applyMigrationFile(t *testing.T, ctx context.Context, conn interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) +}, name string) { + t.Helper() + if _, err := conn.Exec(ctx, readMigrationFile(t, name)); err != nil { + t.Fatalf("apply migration %s: %v", name, err) + } +} + +func readMigrationFile(t *testing.T, name string) string { + t.Helper() + contents, err := os.ReadFile(filepath.Join(realMigrationsDir(t), name)) + if err != nil { + t.Fatalf("read migration %s: %v", name, err) + } + return string(contents) +} + +func assertInsertCheckViolation(t *testing.T, ctx context.Context, conn interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) +}, sql string, args ...any) { + t.Helper() + if _, err := conn.Exec(ctx, sql, args...); !isCheckViolation(err) { + t.Fatalf("insert under strict constraint: got %v, want check violation", err) + } +} + +func isCheckViolation(err error) bool { + var pgErr *pgconn.PgError + return errors.As(err, &pgErr) && pgErr.Code == "23514" +} diff --git a/server/internal/service/attribution_stamp_test.go b/server/internal/service/attribution_stamp_test.go index 998bd8576..3b1522896 100644 --- a/server/internal/service/attribution_stamp_test.go +++ b/server/internal/service/attribution_stamp_test.go @@ -314,9 +314,9 @@ func TestAttributionForMergedComment_HonorsFailClosedPolicy(t *testing.T) { // real enqueue / coalesce path stamps it) that sets originator_user_id but leaves // accountable_user_id NULL — or different — is rejected at the database, so a future // code path that bypasses finalizeAttribution fails loudly instead of silently -// mis-attributing an audited run (the #5192 comment-merge bug class). The legacy-row -// exemption (originator_source IS NULL) is covered by -// TestAttributionInvariantCheck_ExemptsLegacyRows. +// mis-attributing an audited run (the #5192 comment-merge bug class). The strict +// post-backfill handling of source-NULL rows is covered by +// TestAttributionInvariantCheck_RejectsUnbackfilledLegacyRows. func TestAttributionInvariantCheck_RejectsBypass(t *testing.T) { pool := newResolveOriginatorPool(t) ctx := context.Background() @@ -355,36 +355,21 @@ func TestAttributionInvariantCheck_RejectsBypass(t *testing.T) { t.Cleanup(func() { pool.Exec(context.Background(), `DELETE FROM agent_task_queue WHERE id = $1`, okTaskID) }) } -// TestAttributionInvariantCheck_ExemptsLegacyRows is the upgrade-safety test Elon -// required (MUL-4302, Option A): a pre-migration row — originator set, accountable -// NULL, and originator_source NULL (its column did not exist yet) — must survive the -// NOT VALID CHECK and, crucially, must still accept a later status UPDATE by the new -// backend (claim / complete / cancel), even though that UPDATE does not touch the -// attribution columns. This reproduces the cross-deployment stale-task scenario the -// two-phase rollout protects; without the `originator_source IS NULL` exemption the -// UPDATE would fail the CHECK on the whole updated row. -func TestAttributionInvariantCheck_ExemptsLegacyRows(t *testing.T) { +// TestAttributionInvariantCheck_RejectsUnbackfilledLegacyRows verifies the second +// phase of the two-phase rollout (MUL-4302). Once the out-of-band backfill is complete, +// originator_source=NULL no longer exempts a row from the one-way invariant. A stale +// writer or missed backfill that tries to persist originator set with accountable NULL +// must fail loudly instead of recreating the legacy shape. +func TestAttributionInvariantCheck_RejectsUnbackfilledLegacyRows(t *testing.T) { pool := newResolveOriginatorPool(t) ctx := context.Background() _, userA, agentID, issueID := seedAttributionFixture(t, pool) - // A legacy-shaped row: originator set, accountable + source NULL. The exemption - // admits it (as it admits the real pre-migration rows the migration cannot rewrite). - var legacyID string - if err := pool.QueryRow(ctx, ` + if _, err := pool.Exec(ctx, ` INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, originator_user_id) - VALUES ($1, (SELECT runtime_id FROM agent WHERE id = $1), $2, 'queued', 0, $3) RETURNING id`, - agentID, issueID, userA).Scan(&legacyID); err != nil { - t.Fatalf("legacy-shaped row must be admitted by the exemption, got %v", err) - } - t.Cleanup(func() { pool.Exec(context.Background(), `DELETE FROM agent_task_queue WHERE id = $1`, legacyID) }) - - // The status transitions the new backend performs on such a stale task must all - // succeed — this is exactly what a bare invariant CHECK would have broken. - for _, status := range []string{"running", "completed", "cancelled"} { - if _, err := pool.Exec(ctx, `UPDATE agent_task_queue SET status = $2 WHERE id = $1`, legacyID, status); err != nil { - t.Fatalf("status UPDATE to %q on a legacy row must succeed, got %v", status, err) - } + VALUES ($1, (SELECT runtime_id FROM agent WHERE id = $1), $2, 'queued', 0, $3)`, + agentID, issueID, userA); err == nil { + t.Fatal("expected the strict CHECK to reject an unbackfilled legacy row, but insert succeeded") } } diff --git a/server/migrations/197_agent_task_attribution_strict_constraint.down.sql b/server/migrations/197_agent_task_attribution_strict_constraint.down.sql new file mode 100644 index 000000000..3a5ed0a59 --- /dev/null +++ b/server/migrations/197_agent_task_attribution_strict_constraint.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE agent_task_queue + DROP CONSTRAINT IF EXISTS agent_task_queue_accountable_matches_originator_strict; diff --git a/server/migrations/197_agent_task_attribution_strict_constraint.up.sql b/server/migrations/197_agent_task_attribution_strict_constraint.up.sql new file mode 100644 index 000000000..39a8b95c8 --- /dev/null +++ b/server/migrations/197_agent_task_attribution_strict_constraint.up.sql @@ -0,0 +1,23 @@ +-- Human Attribution rollout, phase 2: install the strict one-way invariant +-- alongside migration 190's transitional constraint (MUL-4302). +-- +-- Migration 190 temporarily exempted rows with originator_source IS NULL so +-- active legacy tasks and stale writers could survive the rolling deployment. +-- Those rows have now been backfilled out of band. Install the strict form +-- under a temporary name first so it starts protecting new writes immediately +-- while the existing constraint remains in place. +-- +-- Keep this ADD separate from migration 198's full-table VALIDATE. PostgreSQL +-- holds the ACCESS EXCLUSIVE lock taken here until the statement transaction +-- ends; combining both steps would unnecessarily block queue traffic for the +-- duration of the validation scan. +ALTER TABLE agent_task_queue + ADD CONSTRAINT agent_task_queue_accountable_matches_originator_strict + CHECK ( + originator_user_id IS NULL + OR ( + accountable_user_id IS NOT NULL + AND accountable_user_id = originator_user_id + ) + ) + NOT VALID; diff --git a/server/migrations/198_agent_task_attribution_strict_constraint_validate.down.sql b/server/migrations/198_agent_task_attribution_strict_constraint_validate.down.sql new file mode 100644 index 000000000..2fa73d10f --- /dev/null +++ b/server/migrations/198_agent_task_attribution_strict_constraint_validate.down.sql @@ -0,0 +1,15 @@ +-- PostgreSQL cannot mark a validated constraint NOT VALID again. Recreate the +-- shadow constraint to restore the state immediately after migration 197. +ALTER TABLE agent_task_queue + DROP CONSTRAINT IF EXISTS agent_task_queue_accountable_matches_originator_strict; + +ALTER TABLE agent_task_queue + ADD CONSTRAINT agent_task_queue_accountable_matches_originator_strict + CHECK ( + originator_user_id IS NULL + OR ( + accountable_user_id IS NOT NULL + AND accountable_user_id = originator_user_id + ) + ) + NOT VALID; diff --git a/server/migrations/198_agent_task_attribution_strict_constraint_validate.up.sql b/server/migrations/198_agent_task_attribution_strict_constraint_validate.up.sql new file mode 100644 index 000000000..3b6d1e5eb --- /dev/null +++ b/server/migrations/198_agent_task_attribution_strict_constraint_validate.up.sql @@ -0,0 +1,9 @@ +-- Validate the strict shadow constraint only after the attribution backfill has +-- completed. This takes a SHARE UPDATE EXCLUSIVE lock while scanning the table, +-- which permits normal INSERT/UPDATE/DELETE traffic to continue. +-- +-- Any remaining row where originator_user_id is set but accountable_user_id is +-- NULL or different causes this migration to fail closed. The validated shadow +-- replaces migration 190's transitional constraint in migration 199. +ALTER TABLE agent_task_queue + VALIDATE CONSTRAINT agent_task_queue_accountable_matches_originator_strict; diff --git a/server/migrations/199_agent_task_attribution_legacy_exemption_remove.down.sql b/server/migrations/199_agent_task_attribution_legacy_exemption_remove.down.sql new file mode 100644 index 000000000..554a7d9d0 --- /dev/null +++ b/server/migrations/199_agent_task_attribution_legacy_exemption_remove.down.sql @@ -0,0 +1,18 @@ +-- Restore migration 190's transitional constraint without changing data. Keep +-- the validated strict constraint under its shadow name until the transitional +-- constraint is installed, so rollback never leaves the table unprotected. +ALTER TABLE agent_task_queue + RENAME CONSTRAINT agent_task_queue_accountable_matches_originator + TO agent_task_queue_accountable_matches_originator_strict; + +ALTER TABLE agent_task_queue + ADD CONSTRAINT agent_task_queue_accountable_matches_originator + CHECK ( + originator_source IS NULL + OR originator_user_id IS NULL + OR ( + accountable_user_id IS NOT NULL + AND accountable_user_id = originator_user_id + ) + ) + NOT VALID; diff --git a/server/migrations/199_agent_task_attribution_legacy_exemption_remove.up.sql b/server/migrations/199_agent_task_attribution_legacy_exemption_remove.up.sql new file mode 100644 index 000000000..032038101 --- /dev/null +++ b/server/migrations/199_agent_task_attribution_legacy_exemption_remove.up.sql @@ -0,0 +1,10 @@ +-- The strict shadow constraint is validated, so migration 190's transitional +-- originator_source IS NULL exemption can now be removed. Both statements are +-- sent in one migration query and therefore commit atomically: readers never +-- observe the table without either constraint in place. +ALTER TABLE agent_task_queue + DROP CONSTRAINT agent_task_queue_accountable_matches_originator; + +ALTER TABLE agent_task_queue + RENAME CONSTRAINT agent_task_queue_accountable_matches_originator_strict + TO agent_task_queue_accountable_matches_originator;