fix: distinguish archived automation runs from cancelled ones (#1860)

2026-07-24 10:50:36 +02:00
parent 93c6987bf5
commit 62919efbdc
12 changed files with 576 additions and 63 deletions
@@ -2,9 +2,16 @@ package sqlite

 import (
 	"context"
+	"fmt"
+	"net/url"
+	"strings"
+	"sync"
 	"testing"
 	"time"

+	"github.com/google/uuid"
+	"github.com/jmoiron/sqlx"
+
 	"github.com/kandev/kandev/internal/agentruntime"
 	"github.com/kandev/kandev/internal/task/models"
 	"github.com/kandev/kandev/internal/testutil"
@@ -314,3 +321,124 @@ func TestPostgresTaskEnvironmentReposMultiBranchMigration(t *testing.T) {
 		t.Fatal("expected duplicate env/repo/branch insert to fail")
 	}
 }
+
+// openIsolatedPostgresMultiConn is like testutil.OpenIsolatedPostgres but
+// supports a real multi-connection pool. OpenIsolatedPostgres scopes its
+// isolated schema via a session-level `SET search_path` issued on one
+// connection — fine for its single-connection pool, but a second pooled
+// connection never sees that SET and falls back to the default "public"
+// schema. Baking search_path into the DSN's libpq `options` param instead
+// makes every new connection resolve unqualified table names against the
+// isolated schema, so the pool can be sized for genuine concurrency tests.
+func openIsolatedPostgresMultiConn(t *testing.T, dsn string, maxConns int) *sqlx.DB {
+	t.Helper()
+	schema := "kandev_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
+
+	setup, err := sqlx.Open("pgx", dsn)
+	if err != nil {
+		t.Fatalf("open postgres (schema setup): %v", err)
+	}
+	if _, err := setup.Exec("CREATE SCHEMA " + schema); err != nil {
+		_ = setup.Close()
+		t.Fatalf("create postgres schema %s: %v", schema, err)
+	}
+	_ = setup.Close()
+	// Register schema-drop cleanup immediately: any t.Fatalf between here and
+	// the end of this function (e.g. sqlx.Open below failing) must still not
+	// leak the schema on the Postgres test instance.
+	t.Cleanup(func() {
+		if cleanup, cerr := sqlx.Open("pgx", dsn); cerr == nil {
+			_, _ = cleanup.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
+			_ = cleanup.Close()
+		}
+	})
+
+	// dsn is either a URL ("postgres://user:pass@host/db?sslmode=...") or a
+	// libpq keyword/value string ("host=... port=... sslmode=disable" —
+	// what CI's KANDEV_TEST_POSTGRES_DSN uses). Appending "?options=..." to
+	// the latter corrupts the last keyword's value instead of adding a new
+	// one, so the two forms need different separators.
+	var scopedDSN string
+	if strings.Contains(dsn, "://") {
+		sep := "?"
+		if strings.Contains(dsn, "?") {
+			sep = "&"
+		}
+		scopedDSN = dsn + sep + "options=" + url.QueryEscape("-c search_path="+schema)
+	} else {
+		scopedDSN = dsn + " options='-c search_path=" + schema + "'"
+	}
+	db, err := sqlx.Open("pgx", scopedDSN)
+	if err != nil {
+		t.Fatalf("open postgres (scoped, %d conns): %v", maxConns, err)
+	}
+	db.SetMaxOpenConns(maxConns)
+	db.SetMaxIdleConns(maxConns)
+	t.Cleanup(func() { _ = db.Close() })
+	return db
+}
+
+// TestPostgresSetSessionPrimary_ConcurrentPromotionsLeaveExactlyOnePrimary
+// exercises the race Greptile flagged on PR #1860: SetSessionPrimary demotes
+// every other session for a task and promotes the target in one
+// transaction, plus a `SELECT ... FOR UPDATE` row lock on the owning task
+// so concurrent promotions on separate Postgres connections serialize
+// instead of interleaving their demote/promote pairs. Uses
+// openIsolatedPostgresMultiConn (not testutil.OpenIsolatedPostgres, which
+// caps the pool at one connection) because genuine cross-connection
+// concurrency is the whole point — a one-connection pool would trivially
+// serialize even the pre-fix code and prove nothing.
+func TestPostgresSetSessionPrimary_ConcurrentPromotionsLeaveExactlyOnePrimary(t *testing.T) {
+	const concurrency = 8
+	db := openIsolatedPostgresMultiConn(t, testutil.PostgresDSNFromEnv(t), concurrency)
+	repo, err := NewWithDB(db, db, nil)
+	if err != nil {
+		t.Fatalf("init postgres schema: %v", err)
+	}
+	ctx := context.Background()
+
+	now := time.Now().UTC()
+	if _, err := db.Exec(db.Rebind(`
+		INSERT INTO tasks (id, workspace_id, title, created_at, updated_at)
+		VALUES (?, ?, ?, ?, ?)
+	`), "task-primary-race", "ws-primary-race", "Task primary race", now, now); err != nil {
+		t.Fatalf("seed task: %v", err)
+	}
+
+	sessionIDs := make([]string, concurrency)
+	for i := range sessionIDs {
+		sessionIDs[i] = fmt.Sprintf("session-primary-race-%d", i)
+		if err := repo.CreateTaskSession(ctx, &models.TaskSession{
+			ID: sessionIDs[i], TaskID: "task-primary-race", State: models.TaskSessionStateRunning,
+		}); err != nil {
+			t.Fatalf("CreateTaskSession(%s): %v", sessionIDs[i], err)
+		}
+	}
+
+	var wg sync.WaitGroup
+	errs := make([]error, concurrency)
+	for i, sessionID := range sessionIDs {
+		wg.Add(1)
+		go func(i int, sessionID string) {
+			defer wg.Done()
+			errs[i] = repo.SetSessionPrimary(ctx, sessionID)
+		}(i, sessionID)
+	}
+	wg.Wait()
+
+	for i, err := range errs {
+		if err != nil {
+			t.Fatalf("SetSessionPrimary(%s) concurrent call failed: %v", sessionIDs[i], err)
+		}
+	}
+
+	var primaryCount int
+	if err := db.Get(&primaryCount, db.Rebind(
+		`SELECT COUNT(*) FROM task_sessions WHERE task_id = ? AND is_primary = 1`,
+	), "task-primary-race"); err != nil {
+		t.Fatalf("count primary sessions: %v", err)
+	}
+	if primaryCount != 1 {
+		t.Errorf("expected exactly 1 primary session after %d concurrent promotions, got %d — FOR UPDATE lock not serializing across connections", concurrency, primaryCount)
+	}
+}
@@ -1914,13 +1914,32 @@ func (r *Repository) GetPrimarySessionInfoByTaskIDs(ctx context.Context, taskIDs
 	return result, rows.Err()
 }

-// SetSessionPrimary marks a session as primary and clears primary flag on other sessions for the same task
+// SetSessionPrimary marks a session as primary and clears the primary flag
+// on every other session for the same task, atomically. Both writes go
+// through a single transaction so a concurrent caller can't observe (or
+// write) a half-applied state — e.g. two sessions racing on is_primary=1,
+// or a reader seeing zero primary sessions mid-swap.
+//
+// On SQLite the writer pool is a single connection (see
+// internal/db.NewSQLiteDB), so only one transaction can hold it at a time —
+// the transaction alone fully serializes concurrent callers. On Postgres,
+// separate connections could otherwise run two of these transactions truly
+// concurrently, each seeing zero primary sessions and both promoting; to
+// close that window we take an exclusive row lock on the owning task
+// (`SELECT ... FOR UPDATE`) before touching its sessions, so a second
+// concurrent promotion for the same task blocks until the first commits.
 func (r *Repository) SetSessionPrimary(ctx context.Context, sessionID string) error {
 	now := time.Now().UTC()

+	tx, err := r.db.BeginTxx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = tx.Rollback() }()
+
 	// First, get the task_id for this session
 	var taskID string
-	err := r.db.QueryRowContext(ctx, r.db.Rebind(`SELECT task_id FROM task_sessions WHERE id = ?`), sessionID).Scan(&taskID)
+	err = tx.QueryRowContext(ctx, r.db.Rebind(`SELECT task_id FROM task_sessions WHERE id = ?`), sessionID).Scan(&taskID)
 	if err == sql.ErrNoRows {
 		return fmt.Errorf("session not found: %s", sessionID)
 	}
@@ -1928,8 +1947,19 @@ func (r *Repository) SetSessionPrimary(ctx context.Context, sessionID string) er
 		return err
 	}

+	// Serialize concurrent promotions for the same task across Postgres
+	// connections. SQLite has no FOR UPDATE / row-level locking and doesn't
+	// need it — the single-connection writer pool already serializes here.
+	if dialect.IsPostgres(r.db.DriverName()) {
+		var lockedTaskID string
+		err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT id FROM tasks WHERE id = ? FOR UPDATE`), taskID).Scan(&lockedTaskID)
+		if err != nil && err != sql.ErrNoRows {
+			return err
+		}
+	}
+
 	// Clear primary flag on all sessions for this task
-	_, err = r.db.ExecContext(ctx, r.db.Rebind(`
+	_, err = tx.ExecContext(ctx, r.db.Rebind(`
 		UPDATE task_sessions SET is_primary = 0, updated_at = ? WHERE task_id = ?
 	`), now, taskID)
 	if err != nil {
@@ -1937,7 +1967,7 @@ func (r *Repository) SetSessionPrimary(ctx context.Context, sessionID string) er
 	}

 	// Set primary flag on the specified session
-	result, err := r.db.ExecContext(ctx, r.db.Rebind(`
+	result, err := tx.ExecContext(ctx, r.db.Rebind(`
 		UPDATE task_sessions SET is_primary = 1, updated_at = ? WHERE id = ?
 	`), now, sessionID)
 	if err != nil {
@@ -1948,5 +1978,5 @@ func (r *Repository) SetSessionPrimary(ctx context.Context, sessionID string) er
 	if rows == 0 {
 		return fmt.Errorf("session not found: %s", sessionID)
 	}
-	return nil
+	return tx.Commit()
 }