kandev/apps/backend/internal/task/repository/sqlite/session_test.go

package sqlite

//revive:disable:file-length-limit // SQLite session regression coverage is intentionally scenario-heavy.

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/jmoiron/sqlx"

	"github.com/kandev/kandev/internal/db"
	"github.com/kandev/kandev/internal/db/dialect"
	"github.com/kandev/kandev/internal/task/models"
	"github.com/kandev/kandev/internal/testutil"
	"github.com/stretchr/testify/require"
)

func newRepoForSessionTests(t *testing.T) *Repository {
	t.Helper()
	dbPath := filepath.Join(t.TempDir(), "session-test.db")
	dbConn, err := db.OpenSQLite(dbPath)
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	sqlxDB := sqlx.NewDb(dbConn, "sqlite3")
	repo, err := NewWithDB(sqlxDB, sqlxDB, nil)
	if err != nil {
		t.Fatalf("new repo: %v", err)
	}
	t.Cleanup(func() { _ = sqlxDB.Close() })
	return repo
}

func TestTaskSessionWorkspacePathUsesCurrentEnvironmentRoot(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	const (
		taskID    = "task-workspace-root"
		sessionID = "session-workspace-root"
		envID     = "env-workspace-root"
	)

	if err := repo.CreateTask(ctx, &models.Task{ID: taskID, Title: "Workspace root"}); err != nil {
		t.Fatalf("CreateTask: %v", err)
	}
	if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
		ID:            envID,
		TaskID:        taskID,
		ExecutorType:  string(models.ExecutorTypeWorktree),
		Status:        models.TaskEnvironmentStatusReady,
		WorkspacePath: "/task-root/kandev",
		Repos: []*models.TaskEnvironmentRepo{{
			RepositoryID: "repo-root",
			WorktreeID:   "worktree-primary",
			WorktreePath: "/task-root/kandev",
		}},
	}); err != nil {
		t.Fatalf("CreateTaskEnvironment: %v", err)
	}
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID:                sessionID,
		TaskID:            taskID,
		TaskEnvironmentID: envID,
		WorkspacePath:     "/task-root/kandev",
	}); err != nil {
		t.Fatalf("CreateTaskSession: %v", err)
	}
	env, err := repo.GetTaskEnvironment(ctx, envID)
	if err != nil {
		t.Fatalf("GetTaskEnvironment: %v", err)
	}
	env.WorkspacePath = "/task-root"
	if err := repo.UpdateTaskEnvironment(ctx, env); err != nil {
		t.Fatalf("UpdateTaskEnvironment: %v", err)
	}

	got, err := repo.GetTaskSession(ctx, sessionID)
	if err != nil {
		t.Fatalf("GetTaskSession: %v", err)
	}
	if got.WorkspacePath != "/task-root" {
		t.Fatalf("GetTaskSession WorkspacePath = %q, want %q", got.WorkspacePath, "/task-root")
	}
	if len(got.Worktrees) != 1 || got.Worktrees[0].WorktreePath != "/task-root/kandev" {
		t.Fatalf("GetTaskSession primary worktree = %+v, want repository path", got.Worktrees)
	}

	listed, err := repo.ListTaskSessions(ctx, taskID)
	if err != nil {
		t.Fatalf("ListTaskSessions: %v", err)
	}
	if len(listed) != 1 || listed[0].WorkspacePath != "/task-root" {
		t.Fatalf("ListTaskSessions workspace = %+v, want %q", listed, "/task-root")
	}
}

func TestTaskSessionWorkspacePathFallsBackWithoutEnvironment(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	if err := repo.CreateTask(ctx, &models.Task{ID: "task-legacy-workspace", Title: "Legacy workspace"}); err != nil {
		t.Fatalf("CreateTask: %v", err)
	}
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID:            "session-legacy-workspace",
		TaskID:        "task-legacy-workspace",
		WorkspacePath: "/legacy/repository",
	}); err != nil {
		t.Fatalf("CreateTaskSession: %v", err)
	}

	got, err := repo.GetTaskSession(ctx, "session-legacy-workspace")
	if err != nil {
		t.Fatalf("GetTaskSession: %v", err)
	}
	if got.WorkspacePath != "/legacy/repository" {
		t.Fatalf("WorkspacePath = %q, want legacy fallback", got.WorkspacePath)
	}
}

func TestCreateOfficeTaskSessionMarksOnlyTheFirstConcurrentSessionAsOrigin(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	const taskID = "task-office-origin-race"
	if err := repo.CreateTask(ctx, &models.Task{ID: taskID, Title: "Office origin race"}); err != nil {
		t.Fatalf("CreateTask: %v", err)
	}

	sessions := []*models.TaskSession{
		{ID: "office-origin-agent-1", TaskID: taskID, AgentProfileID: "agent-1", State: models.TaskSessionStateCreated},
		{ID: "office-origin-agent-2", TaskID: taskID, AgentProfileID: "agent-2", State: models.TaskSessionStateCreated},
	}
	errs := make([]error, len(sessions))
	var wg sync.WaitGroup
	for i, session := range sessions {
		wg.Add(1)
		go func(i int, session *models.TaskSession) {
			defer wg.Done()
			errs[i] = repo.CreateOfficeTaskSession(ctx, session)
		}(i, session)
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			t.Fatalf("CreateOfficeTaskSession(%d): %v", i, err)
		}
	}

	created, err := repo.ListTaskSessions(ctx, taskID)
	if err != nil {
		t.Fatalf("ListTaskSessions: %v", err)
	}
	if len(created) != len(sessions) {
		t.Fatalf("created sessions = %d, want %d", len(created), len(sessions))
	}
	originCount := 0
	for _, session := range created {
		if models.IsOriginalTaskSession(session.Metadata) {
			originCount++
		}
	}
	if originCount != 1 {
		t.Fatalf("origin-marked sessions = %d, want exactly one", originCount)
	}
}

func TestPostgresCreateOfficeTaskSessionMarksOnlyTheFirstConcurrentSessionAsOrigin(t *testing.T) {
	db := openIsolatedPostgresMultiConn(t, testutil.PostgresDSNFromEnv(t), 2)
	repo, err := NewWithDB(db, db, nil)
	if err != nil {
		t.Fatalf("init postgres schema: %v", err)
	}
	ctx := context.Background()
	const taskID = "task-office-origin-race-postgres"
	now := time.Now().UTC()
	if _, err := db.Exec(db.Rebind(`
		INSERT INTO tasks (id, workspace_id, title, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?)
	`), taskID, "", "Office origin race (Postgres)", now, now); err != nil {
		t.Fatalf("seed task: %v", err)
	}

	sessions := []*models.TaskSession{
		{ID: "office-origin-postgres-agent-1", TaskID: taskID, AgentProfileID: "agent-1", State: models.TaskSessionStateCreated},
		{ID: "office-origin-postgres-agent-2", TaskID: taskID, AgentProfileID: "agent-2", State: models.TaskSessionStateCreated},
	}
	errs := make([]error, len(sessions))
	var wg sync.WaitGroup
	for i, session := range sessions {
		wg.Add(1)
		go func(i int, session *models.TaskSession) {
			defer wg.Done()
			errs[i] = repo.CreateOfficeTaskSession(ctx, session)
		}(i, session)
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			t.Fatalf("CreateOfficeTaskSession(%d): %v", i, err)
		}
	}

	created, err := repo.ListTaskSessions(ctx, taskID)
	if err != nil {
		t.Fatalf("ListTaskSessions: %v", err)
	}
	if len(created) != len(sessions) {
		t.Fatalf("created sessions = %d, want %d", len(created), len(sessions))
	}
	originCount := 0
	for _, session := range created {
		if models.IsOriginalTaskSession(session.Metadata) {
			originCount++
		}
	}
	if originCount != 1 {
		t.Fatalf("origin-marked sessions = %d, want exactly one", originCount)
	}
}

// seedForMsgTest seeds task, session, and turn rows so that all FK constraints
// on task_session_messages are satisfied. Returns the turn ID for use in inserts.
func seedForMsgTest(t *testing.T, repo *Repository, taskID, sessionID, turnID string) {
	t.Helper()
	now := time.Now().UTC()
	_, err := repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO tasks (id, workspace_id, title, created_at, updated_at)
		VALUES (?, '', 'test task', ?, ?)
	`), taskID, now, now)
	if err != nil {
		t.Fatalf("seed task %s: %v", taskID, err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO task_sessions
			(id, task_id, started_at, updated_at)
		VALUES (?, ?, ?, ?)
	`), sessionID, taskID, now, now)
	if err != nil {
		t.Fatalf("seed session %s: %v", sessionID, err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO task_session_turns
			(id, task_session_id, task_id, started_at, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?)
	`), turnID, sessionID, taskID, now, now, now)
	if err != nil {
		t.Fatalf("seed turn %s: %v", turnID, err)
	}
}

// insertAgentMsg inserts a message row directly into the DB under the given
// session and turn. authorType must be 'agent' or 'user'.
func insertAgentMsg(t *testing.T, repo *Repository, id, sessionID, turnID, authorType, content string, ts time.Time) {
	t.Helper()
	_, err := repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_session_messages
			(id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at)
		VALUES (?, ?, '', ?, ?, '', ?, 0, 'message', '{}', ?)
	`), id, sessionID, turnID, authorType, content, ts)
	if err != nil {
		t.Fatalf("insert message %s: %v", id, err)
	}
}

func TestRenameTaskSession(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	if err := repo.RenameTaskSession(ctx, "missing-session", "reviewer"); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("RenameTaskSession error = %v, want ErrTaskSessionNotFound", err)
	}

	seedForMsgTest(t, repo, "task-rename", "session-rename", "turn-rename")
	if err := repo.RenameTaskSession(ctx, "session-rename", "reviewer"); err != nil {
		t.Fatalf("RenameTaskSession: %v", err)
	}
	session, err := repo.GetTaskSession(ctx, "session-rename")
	if err != nil {
		t.Fatalf("GetTaskSession after rename: %v", err)
	}
	if session.Name != "reviewer" {
		t.Fatalf("session.Name = %q, want %q", session.Name, "reviewer")
	}
	if got := session.ToAPI()["name"]; got != "reviewer" {
		t.Fatalf(`ToAPI()["name"] = %v, want "reviewer"`, got)
	}

	// Clearing the name falls back to the derived tab title on the frontend.
	if err := repo.RenameTaskSession(ctx, "session-rename", ""); err != nil {
		t.Fatalf("RenameTaskSession clear: %v", err)
	}
	session, err = repo.GetTaskSession(ctx, "session-rename")
	if err != nil {
		t.Fatalf("GetTaskSession after clear: %v", err)
	}
	if session.Name != "" {
		t.Fatalf("session.Name = %q, want empty after clear", session.Name)
	}
	if _, ok := session.ToAPI()["name"]; ok {
		t.Fatalf("ToAPI() should omit name when empty")
	}

	// Name survives CreateTaskSession round-trips and list scans.
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID: "session-named", TaskID: "task-rename", Name: "verifier",
	}); err != nil {
		t.Fatalf("CreateTaskSession with name: %v", err)
	}
	sessions, err := repo.ListTaskSessions(ctx, "task-rename")
	if err != nil {
		t.Fatalf("ListTaskSessionsByTaskID: %v", err)
	}
	found := false
	for _, s := range sessions {
		if s.ID == "session-named" {
			found = true
			if s.Name != "verifier" {
				t.Fatalf("listed session Name = %q, want %q", s.Name, "verifier")
			}
		}
	}
	if !found {
		t.Fatalf("session-named not returned by ListTaskSessions")
	}
}

// TestUpdateTaskSessionLastReadMessageID verifies the read-cursor setter
// rejects a missing session, persists the message id (round-tripping through
// both GetTaskSession and ListTaskSessions, which scan via the single-row and
// multi-row helpers respectively), that ToAPI omits an empty cursor, and that
// the cursor never regresses when a stale/out-of-order messageID is applied
// after a newer one already landed.
func TestUpdateTaskSessionLastReadMessageID(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "missing-session", "msg-1"); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("UpdateTaskSessionLastReadMessageID error = %v, want ErrTaskSessionNotFound", err)
	}

	seedForMsgTest(t, repo, "task-read", "session-read", "turn-read")
	now := time.Now().UTC()
	insertAgentMsg(t, repo, "msg-1", "session-read", "turn-read", "user", "hi", now)
	insertAgentMsg(t, repo, "msg-2", "session-read", "turn-read", "agent", "hello", now.Add(time.Second))

	session, err := repo.GetTaskSession(ctx, "session-read")
	if err != nil {
		t.Fatalf("GetTaskSession before mark-read: %v", err)
	}
	if _, ok := session.ToAPI()["last_read_message_id"]; ok {
		t.Fatalf("ToAPI() should omit last_read_message_id when empty")
	}

	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "session-read", "msg-1"); err != nil {
		t.Fatalf("UpdateTaskSessionLastReadMessageID: %v", err)
	}
	session, err = repo.GetTaskSession(ctx, "session-read")
	if err != nil {
		t.Fatalf("GetTaskSession after mark-read: %v", err)
	}
	if session.LastReadMessageID != "msg-1" {
		t.Fatalf("session.LastReadMessageID = %q, want %q", session.LastReadMessageID, "msg-1")
	}
	if got := session.ToAPI()["last_read_message_id"]; got != "msg-1" {
		t.Fatalf(`ToAPI()["last_read_message_id"] = %v, want "msg-1"`, got)
	}

	// Round-trips through the multi-row scan path (ListTaskSessions) too.
	sessions, err := repo.ListTaskSessions(ctx, "task-read")
	if err != nil {
		t.Fatalf("ListTaskSessions: %v", err)
	}
	if len(sessions) != 1 || sessions[0].LastReadMessageID != "msg-1" {
		t.Fatalf("ListTaskSessions did not carry LastReadMessageID: %#v", sessions)
	}

	// Advancing to a newer message overwrites the cursor.
	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "session-read", "msg-2"); err != nil {
		t.Fatalf("UpdateTaskSessionLastReadMessageID advance: %v", err)
	}
	session, err = repo.GetTaskSession(ctx, "session-read")
	if err != nil {
		t.Fatalf("GetTaskSession after advance: %v", err)
	}
	if session.LastReadMessageID != "msg-2" {
		t.Fatalf("session.LastReadMessageID = %q, want %q", session.LastReadMessageID, "msg-2")
	}

	// A delayed/retried request for the older message (out-of-order arrival)
	// must not regress the cursor — it's a silent no-op, not an error.
	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "session-read", "msg-1"); err != nil {
		t.Fatalf("UpdateTaskSessionLastReadMessageID stale update: %v", err)
	}
	session, err = repo.GetTaskSession(ctx, "session-read")
	if err != nil {
		t.Fatalf("GetTaskSession after stale update: %v", err)
	}
	if session.LastReadMessageID != "msg-2" {
		t.Fatalf("session.LastReadMessageID = %q, want unchanged %q after stale update", session.LastReadMessageID, "msg-2")
	}
}

// TestUpdateTaskSessionLastReadMessageIDTiebreaksEqualTimestampsByID verifies
// the monotonic guard's deterministic tiebreaker: when two messages share the
// exact same created_at (e.g. a burst of messages persisted in the same
// instant), ordering falls back to comparing message id, and a "stale" id in
// that ordering is still rejected rather than silently accepted just because
// the timestamps tie.
func TestUpdateTaskSessionLastReadMessageIDTiebreaksEqualTimestampsByID(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-tie", "session-tie", "turn-tie")
	sameInstant := time.Now().UTC()
	insertAgentMsg(t, repo, "msg-a", "session-tie", "turn-tie", "user", "a", sameInstant)
	insertAgentMsg(t, repo, "msg-b", "session-tie", "turn-tie", "agent", "b", sameInstant)

	// "msg-b" > "msg-a" lexically, so it's the tiebreak winner at this
	// shared timestamp — advancing to it must succeed.
	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "session-tie", "msg-b"); err != nil {
		t.Fatalf("UpdateTaskSessionLastReadMessageID to tiebreak winner: %v", err)
	}

	// Falling back to the tiebreak loser at the same timestamp must be
	// rejected as stale, exactly like a strictly-older timestamp would be.
	if err := repo.UpdateTaskSessionLastReadMessageID(ctx, "session-tie", "msg-a"); err != nil {
		t.Fatalf("UpdateTaskSessionLastReadMessageID stale tiebreak: %v", err)
	}
	session, err := repo.GetTaskSession(ctx, "session-tie")
	if err != nil {
		t.Fatalf("GetTaskSession: %v", err)
	}
	if session.LastReadMessageID != "msg-b" {
		t.Fatalf("session.LastReadMessageID = %q, want unchanged %q after stale tiebreak", session.LastReadMessageID, "msg-b")
	}
}

// TestTaskSessionNotFoundErrorsAreTyped verifies that GetTaskSession,
// UpdateTaskSession, UpdateTaskSessionState, and UpdateTaskSessionBaseCommit
// all return ErrTaskSessionNotFound for a missing session, that
// GetTaskSessionByTaskAndAgent translates not-found to (nil, nil), and that
// GetPrimarySessionByTaskID returns ErrNoPrimarySession, while operations on
// an existing session succeed.
func TestTaskSessionNotFoundErrorsAreTyped(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	if _, err := repo.GetTaskSession(ctx, "missing-session"); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("GetTaskSession error = %v, want ErrTaskSessionNotFound", err)
	}
	if err := repo.UpdateTaskSession(ctx, &models.TaskSession{ID: "missing-session"}); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("UpdateTaskSession error = %v, want ErrTaskSessionNotFound", err)
	}
	if err := repo.UpdateTaskSessionState(ctx, "missing-session", models.TaskSessionStateCompleted, ""); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("UpdateTaskSessionState error = %v, want ErrTaskSessionNotFound", err)
	}
	if err := repo.UpdateTaskSessionBaseCommit(ctx, "missing-session", "abc123"); !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Fatalf("UpdateTaskSessionBaseCommit error = %v, want ErrTaskSessionNotFound", err)
	}

	session, err := repo.GetTaskSessionByTaskAndAgent(ctx, "task-missing", "agent-missing")
	if err != nil {
		t.Fatalf("GetTaskSessionByTaskAndAgent should translate not found to nil, nil: %v", err)
	}
	if session != nil {
		t.Fatalf("GetTaskSessionByTaskAndAgent session = %#v, want nil", session)
	}
	if _, err := repo.GetPrimarySessionByTaskID(ctx, "task-missing"); !errors.Is(err, ErrNoPrimarySession) {
		t.Fatalf("GetPrimarySessionByTaskID error = %v, want ErrNoPrimarySession", err)
	}

	seedForMsgTest(t, repo, "task-found", "session-found", "turn-found")
	if _, err := repo.GetTaskSession(ctx, "session-found"); err != nil {
		t.Fatalf("GetTaskSession existing row: %v", err)
	}
	if err := repo.UpdateTaskSessionState(ctx, "session-found", models.TaskSessionStateCompleted, ""); err != nil {
		t.Fatalf("UpdateTaskSessionState existing row: %v", err)
	}
}

func TestClaimPromptableTaskSessionIfActive(t *testing.T) {
	ctx := context.Background()

	t.Run("claims an active promptable session and persists running", func(t *testing.T) {
		repo := newRepoForSessionTests(t)
		seedForMsgTest(t, repo, "task-active", "session-active", "turn-active")
		require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-active", models.TaskSessionStateWaitingForInput, ""))

		claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "session-active")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionClaimed, claim.Status)
		require.Equal(t, models.TaskSessionStateWaitingForInput, claim.PreviousState)

		persisted, err := repo.GetTaskSession(ctx, "session-active")
		require.NoError(t, err)
		require.Equal(t, models.TaskSessionStateRunning, persisted.State)
	})

	t.Run("clears completion timestamp when reclaiming a completed session", func(t *testing.T) {
		repo := newRepoForSessionTests(t)
		seedForMsgTest(t, repo, "task-completed", "session-completed", "turn-completed")
		require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-completed", models.TaskSessionStateCompleted, ""))

		before, err := repo.GetTaskSession(ctx, "session-completed")
		require.NoError(t, err)
		require.NotNil(t, before.CompletedAt)

		claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "session-completed")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionClaimed, claim.Status)

		persisted, err := repo.GetTaskSession(ctx, "session-completed")
		require.NoError(t, err)
		require.Equal(t, models.TaskSessionStateRunning, persisted.State)
		require.Nil(t, persisted.CompletedAt)
	})

	t.Run("does not claim an archived task session", func(t *testing.T) {
		repo := newRepoForSessionTests(t)
		seedForMsgTest(t, repo, "task-archived", "session-archived", "turn-archived")
		require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-archived", models.TaskSessionStateIdle, ""))
		require.NoError(t, repo.ArchiveTask(ctx, "task-archived"))

		claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "session-archived")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionInactive, claim.Status)

		persisted, err := repo.GetTaskSession(ctx, "session-archived")
		require.NoError(t, err)
		require.Equal(t, models.TaskSessionStateIdle, persisted.State)
	})

	t.Run("does not claim a deleted or missing task session", func(t *testing.T) {
		repo := newRepoForSessionTests(t)

		claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "missing-session")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionInactive, claim.Status)

		seedForMsgTest(t, repo, "task-deleted", "session-deleted", "turn-deleted")
		require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-deleted", models.TaskSessionStateCompleted, ""))
		require.NoError(t, repo.DeleteTask(ctx, "task-deleted"))

		claim, err = repo.ClaimPromptableTaskSessionIfActive(ctx, "session-deleted")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionInactive, claim.Status)
	})

	t.Run("does not claim a non-promptable session", func(t *testing.T) {
		repo := newRepoForSessionTests(t)
		seedForMsgTest(t, repo, "task-running", "session-running", "turn-running")
		require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-running", models.TaskSessionStateRunning, ""))

		claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "session-running")
		require.NoError(t, err)
		require.Equal(t, models.PromptableTaskSessionBusy, claim.Status)

		persisted, err := repo.GetTaskSession(ctx, "session-running")
		require.NoError(t, err)
		require.Equal(t, models.TaskSessionStateRunning, persisted.State)
	})
}

func TestClaimPromptableTaskSessionIfActive_ZeroRowAfterConcurrentBusyTransition(t *testing.T) {
	ctx := context.Background()
	repo := newRepoForSessionTests(t)
	seedForMsgTest(t, repo, "task-race", "session-race", "turn-race")
	require.NoError(t, repo.UpdateTaskSessionState(ctx, "session-race", models.TaskSessionStateWaitingForInput, ""))

	// Simulate the competing writer winning after ClaimPromptableTaskSessionIfActive
	// has read WAITING_FOR_INPUT but before its guarded UPDATE can take ownership.
	// RAISE(IGNORE) makes the outer UPDATE report zero rows deterministically.
	_, err := repo.db.Exec(`
		CREATE TRIGGER claim_race_before_running
		BEFORE UPDATE OF state ON task_sessions
		WHEN OLD.id = 'session-race'
		  AND OLD.state = 'WAITING_FOR_INPUT'
		  AND NEW.state = 'RUNNING'
		BEGIN
			UPDATE task_sessions SET state = 'RUNNING' WHERE id = OLD.id;
			SELECT RAISE(IGNORE);
		END
	`)
	require.NoError(t, err)

	claim, err := repo.ClaimPromptableTaskSessionIfActive(ctx, "session-race")
	require.NoError(t, err)
	require.Equal(t, models.PromptableTaskSessionBusy, claim.Status,
		"a zero-row guarded update must distinguish a concurrent active session from an inactive task")

	persisted, err := repo.GetTaskSession(ctx, "session-race")
	require.NoError(t, err)
	require.Equal(t, models.TaskSessionStateRunning, persisted.State)
}

func TestClassifyPromptableTaskSessionClaimPropagatesScanError(t *testing.T) {
	ctx := context.Background()
	repo := newRepoForSessionTests(t)
	seedForMsgTest(t, repo, "task-scan-error", "session-scan-error", "turn-scan-error")

	tx, err := repo.db.BeginTxx(ctx, nil)
	require.NoError(t, err)
	t.Cleanup(func() { _ = tx.Rollback() })
	_, err = tx.ExecContext(ctx, `ALTER TABLE tasks RENAME TO tasks_unavailable`)
	require.NoError(t, err)

	_, err = repo.classifyPromptableTaskSessionClaim(ctx, tx, "session-scan-error")
	require.Error(t, err)
}

func TestSetSessionMetadataKeyIfAbsentSQLiteIsWriteOnce(t *testing.T) {
	repo := newRepoForSessionTests(t)
	seedForMsgTest(t, repo, "task-baseline", "session-baseline", "turn-baseline")
	ctx := context.Background()

	stored, err := repo.SetSessionMetadataKeyIfAbsent(ctx, "session-baseline", "baseline", map[string]string{"effort": "high"})
	if err != nil {
		t.Fatalf("first SetSessionMetadataKeyIfAbsent: %v", err)
	}
	if !stored {
		t.Fatal("first SetSessionMetadataKeyIfAbsent should store")
	}
	stored, err = repo.SetSessionMetadataKeyIfAbsent(ctx, "session-baseline", "baseline", map[string]string{"effort": "low"})
	if err != nil {
		t.Fatalf("second SetSessionMetadataKeyIfAbsent: %v", err)
	}
	if stored {
		t.Fatal("second SetSessionMetadataKeyIfAbsent should not overwrite")
	}

	session, err := repo.GetTaskSession(ctx, "session-baseline")
	if err != nil {
		t.Fatalf("GetTaskSession: %v", err)
	}
	baseline, ok := session.Metadata["baseline"].(map[string]interface{})
	if !ok || baseline["effort"] != "high" {
		t.Fatalf("baseline = %#v, want effort=high", session.Metadata["baseline"])
	}
}

func TestSetSessionMetadataKeyIfAbsentOrDifferentStepSQLiteReplacesOnlyStaleStep(t *testing.T) {
	repo := newRepoForSessionTests(t)
	seedForMsgTest(t, repo, "task-step-claim", "session-step-claim", "turn-step-claim")
	ctx := context.Background()

	first := models.PendingStepCompletionSignal{StepID: "step-1", Summary: "first"}
	stored, err := repo.SetSessionMetadataKeyIfAbsentOrDifferentStep(
		ctx, "session-step-claim", models.SessionMetaKeyPendingStepCompletion, "step-1", first)
	require.NoError(t, err)
	require.True(t, stored, "an empty signal bag should be claimed")

	second := models.PendingStepCompletionSignal{StepID: "step-2", Summary: "second"}
	stored, err = repo.SetSessionMetadataKeyIfAbsentOrDifferentStep(
		ctx, "session-step-claim", models.SessionMetaKeyPendingStepCompletion, "step-2", second)
	require.NoError(t, err)
	require.True(t, stored, "a signal from an older step should be replaced")

	third := models.PendingStepCompletionSignal{StepID: "step-2", Summary: "third"}
	stored, err = repo.SetSessionMetadataKeyIfAbsentOrDifferentStep(
		ctx, "session-step-claim", models.SessionMetaKeyPendingStepCompletion, "step-2", third)
	require.NoError(t, err)
	require.False(t, stored, "a signal for the current step should keep the first payload")

	session, err := repo.GetTaskSession(ctx, "session-step-claim")
	require.NoError(t, err)
	signal, ok := models.LoadPendingStepSignal(session.Metadata)
	require.True(t, ok)
	require.Equal(t, "step-2", signal.StepID)
	require.Equal(t, "second", signal.Summary)
}

func TestUpdateSessionContextWindowSQLiteCountsStrictUsageDrops(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-context-count", "session-context-count", "turn-context-count")
	require.NoError(t, repo.SetSessionMetadataKey(ctx, "session-context-count", "unrelated", "kept"))

	count, err := repo.UpdateSessionContextWindow(ctx, "session-context-count", map[string]interface{}{
		"size": int64(200000), "used": int64(120000), "remaining": int64(80000),
	})
	require.NoError(t, err)
	require.Equal(t, int64(0), count)

	count, err = repo.UpdateSessionContextWindow(ctx, "session-context-count", map[string]interface{}{
		"size": int64(200000), "used": int64(120000), "remaining": int64(80000),
	})
	require.NoError(t, err)
	require.Equal(t, int64(0), count)

	count, err = repo.UpdateSessionContextWindow(ctx, "session-context-count", map[string]interface{}{
		"size": int64(200000), "used": int64(80000), "remaining": int64(120000),
	})
	require.NoError(t, err)
	require.Equal(t, int64(1), count)

	count, err = repo.UpdateSessionContextWindow(ctx, "session-context-count", map[string]interface{}{
		"size": int64(200000), "used": int64(80000), "remaining": int64(120000),
	})
	require.NoError(t, err)
	require.Equal(t, int64(1), count)

	count, err = repo.UpdateSessionContextWindow(ctx, "session-context-count", map[string]interface{}{
		"size": int64(200000), "used": int64(100000), "remaining": int64(100000),
	})
	require.NoError(t, err)
	require.Equal(t, int64(1), count)

	session, err := repo.GetTaskSession(ctx, "session-context-count")
	require.NoError(t, err)
	require.Equal(t, "kept", session.Metadata["unrelated"])
	require.Equal(t, float64(1), session.Metadata[models.SessionMetaKeyContextCompactionCount])
	window, ok := session.Metadata[models.SessionMetaKeyContextWindow].(map[string]interface{})
	require.True(t, ok)
	require.Equal(t, float64(100000), window["used"])
}

func TestSetSessionMetadataKeyIfAbsentQueryUsesPostgresJSONB(t *testing.T) {
	query := setSessionMetadataKeyIfAbsentQuery(dialect.PGX)
	if strings.Contains(query, "json_set") || strings.Contains(query, "json_type") || strings.Contains(query, "json(?)") {
		t.Fatalf("postgres write-once query uses SQLite JSON functions: %s", query)
	}
	if !strings.Contains(query, "jsonb_set") || !strings.Contains(query, "jsonb_extract_path") {
		t.Fatalf("postgres write-once query must use JSONB set/existence operations: %s", query)
	}
}

func TestSetSessionMetadataKeyIfAbsentOrDifferentStepQueryUsesPostgresJSONB(t *testing.T) {
	query := setSessionMetadataKeyIfAbsentOrDifferentStepQuery(dialect.PGX)
	if strings.Contains(query, "json_set") || strings.Contains(query, "json_extract") || strings.Contains(query, "json(?)") {
		t.Fatalf("postgres step-aware claim query uses SQLite JSON functions: %s", query)
	}
	if !strings.Contains(query, "jsonb_set") || !strings.Contains(query, "jsonb_extract_path_text") || !strings.Contains(query, "IS DISTINCT FROM") {
		t.Fatalf("postgres step-aware claim query must use atomic JSONB comparison: %s", query)
	}
}

func TestListTaskSessionWorktreesFiltersInactiveRows(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-worktrees", "session-worktrees", "turn-worktrees")
	if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
		ID: "env-worktrees", TaskID: "task-worktrees", ExecutorType: "worktree",
		WorkspacePath: "/tmp", Status: models.TaskEnvironmentStatusReady,
	}); err != nil {
		t.Fatalf("CreateTaskEnvironment: %v", err)
	}
	if _, err := repo.db.Exec(repo.db.Rebind(
		`UPDATE task_sessions SET task_environment_id = ? WHERE id = ?`),
		"env-worktrees", "session-worktrees"); err != nil {
		t.Fatalf("link session to env: %v", err)
	}
	worktrees := []*models.TaskEnvironmentRepo{
		{
			ID: "wt-active", TaskEnvironmentID: "env-worktrees",
			WorktreeID: "worktree-active", RepositoryID: "repo-1", BranchSlug: "main",
		},
		{
			ID: "wt-status-deleted", TaskEnvironmentID: "env-worktrees",
			WorktreeID: "worktree-status-deleted", RepositoryID: "repo-1", BranchSlug: "deleted-status",
		},
		{
			ID: "wt-timestamp-deleted", TaskEnvironmentID: "env-worktrees",
			WorktreeID: "worktree-timestamp-deleted", RepositoryID: "repo-1", BranchSlug: "deleted-at",
		},
	}
	for _, wt := range worktrees {
		if err := repo.CreateTaskEnvironmentRepo(ctx, wt); err != nil {
			t.Fatalf("CreateTaskEnvironmentRepo(%s): %v", wt.ID, err)
		}
	}
	now := time.Now().UTC()
	if _, err := repo.db.Exec(repo.db.Rebind(`
		UPDATE task_environment_repos
		SET status = 'deleted', updated_at = ?
		WHERE id = ?
	`), now, "wt-status-deleted"); err != nil {
		t.Fatalf("mark status deleted: %v", err)
	}
	if _, err := repo.db.Exec(repo.db.Rebind(`
		UPDATE task_environment_repos
		SET deleted_at = ?, updated_at = ?
		WHERE id = ?
	`), now, now, "wt-timestamp-deleted"); err != nil {
		t.Fatalf("mark timestamp deleted: %v", err)
	}

	listed, err := repo.ListTaskSessionWorktrees(ctx, "session-worktrees")
	if err != nil {
		t.Fatalf("ListTaskSessionWorktrees: %v", err)
	}
	if len(listed) != 1 || listed[0].ID != "wt-active" {
		t.Fatalf("ListTaskSessionWorktrees = %+v, want only wt-active", listed)
	}
	batched, err := repo.ListWorktreesBySessionIDs(ctx, []string{"session-worktrees"})
	if err != nil {
		t.Fatalf("ListWorktreesBySessionIDs: %v", err)
	}
	rows := batched["session-worktrees"]
	if len(rows) != 1 || rows[0].ID != "wt-active" {
		t.Fatalf("ListWorktreesBySessionIDs = %+v, want only wt-active", rows)
	}
}

func TestUpdateTaskSessionWorktreeBranchByRepositoryScopesUpdate(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-worktrees", "session-worktrees", "turn-worktrees")
	linkSessionToEnvForTest(t, repo, "session-worktrees", "env-worktrees")

	worktrees := []*models.TaskEnvironmentRepo{
		{
			ID:                "wt-repo-1",
			TaskEnvironmentID: "env-worktrees",
			WorktreeID:        "worktree-repo-1",
			RepositoryID:      "repo-1",
			WorktreeBranch:    "feature/old-one",
		},
		{
			ID:                "wt-repo-2",
			TaskEnvironmentID: "env-worktrees",
			WorktreeID:        "worktree-repo-2",
			RepositoryID:      "repo-2",
			WorktreeBranch:    "feature/old-two",
		},
	}
	for _, wt := range worktrees {
		if err := repo.CreateTaskEnvironmentRepo(ctx, wt); err != nil {
			t.Fatalf("CreateTaskEnvironmentRepo(%s): %v", wt.ID, err)
		}
	}

	if err := repo.UpdateTaskSessionWorktreeBranchByRepository(ctx, "session-worktrees", "repo-1", "feature/new-one"); err != nil {
		t.Fatalf("UpdateTaskSessionWorktreeBranchByRepository: %v", err)
	}

	listed, err := repo.ListTaskSessionWorktrees(ctx, "session-worktrees")
	if err != nil {
		t.Fatalf("ListTaskSessionWorktrees: %v", err)
	}
	branches := map[string]string{}
	for _, wt := range listed {
		branches[wt.RepositoryID] = wt.WorktreeBranch
	}
	if branches["repo-1"] != "feature/new-one" {
		t.Fatalf("repo-1 branch = %q, want feature/new-one", branches["repo-1"])
	}
	if branches["repo-2"] != "feature/old-two" {
		t.Fatalf("repo-2 branch = %q, want feature/old-two", branches["repo-2"])
	}
}

func TestUpdateTaskSessionWorktreeBranchByWorktreeScopesRepeatedRepository(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-repeated-repo", "session-repeated-repo", "turn-repeated-repo")
	if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
		ID: "env-repeated-repo", TaskID: "task-repeated-repo", ExecutorType: "worktree",
		WorkspacePath: "/tmp", Status: models.TaskEnvironmentStatusReady,
	}); err != nil {
		t.Fatalf("CreateTaskEnvironment: %v", err)
	}
	if _, err := repo.db.Exec(repo.db.Rebind(
		`UPDATE task_sessions SET task_environment_id = ? WHERE id = ?`),
		"env-repeated-repo", "session-repeated-repo"); err != nil {
		t.Fatalf("link session to env: %v", err)
	}
	for _, wt := range []*models.TaskEnvironmentRepo{
		{ID: "wt-repeated-one", TaskEnvironmentID: "env-repeated-repo", WorktreeID: "worktree-repeated-one", RepositoryID: "repo-repeated", BranchSlug: "one", WorktreeBranch: "feature/one", Position: 0},
		{ID: "wt-repeated-two", TaskEnvironmentID: "env-repeated-repo", WorktreeID: "worktree-repeated-two", RepositoryID: "repo-repeated", BranchSlug: "two", WorktreeBranch: "feature/two", Position: 1},
	} {
		require.NoError(t, repo.CreateTaskEnvironmentRepo(ctx, wt))
	}
	require.NoError(t, repo.UpdateTaskSessionWorktreeBranchByWorktree(ctx, "session-repeated-repo", "worktree-repeated-two", "feature/two-renamed"))
	worktrees, err := repo.ListTaskSessionWorktrees(ctx, "session-repeated-repo")
	require.NoError(t, err)
	require.Equal(t, "feature/one", worktrees[0].WorktreeBranch)
	require.Equal(t, "feature/two-renamed", worktrees[1].WorktreeBranch)
}

// TestGetLastAgentMessage_NoMessages verifies that a session with no messages
// returns an empty string and sql.ErrNoRows.
// linkSessionToEnvForTest creates the environment row a worktree test
// references and points the session at it.
func linkSessionToEnvForTest(t *testing.T, repo *Repository, sessionID, envID string) {
	t.Helper()
	ctx := context.Background()
	if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
		ID: envID, TaskID: "task-worktrees", ExecutorType: "worktree",
		WorkspacePath: "/tmp", Status: models.TaskEnvironmentStatusReady,
	}); err != nil {
		t.Fatalf("CreateTaskEnvironment: %v", err)
	}
	if _, err := repo.db.Exec(repo.db.Rebind(
		`UPDATE task_sessions SET task_environment_id = ? WHERE id = ?`),
		envID, sessionID); err != nil {
		t.Fatalf("link session to env: %v", err)
	}
}

func TestGetLastAgentMessage_NoMessages(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	msg, err := repo.GetLastAgentMessage(ctx, "sess-empty")
	if err != sql.ErrNoRows {
		t.Errorf("expected sql.ErrNoRows, got %v", err)
	}
	if msg != "" {
		t.Errorf("expected empty string, got %q", msg)
	}
}

// TestGetLastAgentMessage_MessagesAllEmptyContent verifies that when the agent
// message has empty content the function returns "" without error (content
// column allows empty string).
func TestGetLastAgentMessage_MessagesAllEmptyContent(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-ec", "sess-ec", "turn-ec")
	insertAgentMsg(t, repo, "msg-ec-1", "sess-ec", "turn-ec", "agent", "", time.Now().UTC())

	msg, err := repo.GetLastAgentMessage(ctx, "sess-ec")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if msg != "" {
		t.Errorf("expected empty string for empty-content message, got %q", msg)
	}
}

// TestGetLastAgentMessage_ReturnsLatestAgentMessage verifies that the most
// recent agent message is returned, and that user messages are ignored.
func TestGetLastAgentMessage_ReturnsLatestAgentMessage(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-1", "sess-1", "turn-1")

	base := time.Now().UTC()
	// User message — must be ignored by GetLastAgentMessage.
	insertAgentMsg(t, repo, "msg-u-1", "sess-1", "turn-1", "user", "user question", base)
	// First agent message.
	insertAgentMsg(t, repo, "msg-a-1", "sess-1", "turn-1", "agent", "first agent reply", base.Add(time.Second))
	// Second (latest) agent message — this must be returned.
	insertAgentMsg(t, repo, "msg-a-2", "sess-1", "turn-1", "agent", "second agent reply", base.Add(2*time.Second))

	msg, err := repo.GetLastAgentMessage(ctx, "sess-1")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if msg != "second agent reply" {
		t.Errorf("expected 'second agent reply', got %q", msg)
	}
}

// TestGetLastAgentMessage_SessionDoesNotExist verifies that looking up a
// session that has no messages returns an empty string and sql.ErrNoRows.
func TestGetLastAgentMessage_SessionDoesNotExist(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	msg, err := repo.GetLastAgentMessage(ctx, "sess-nonexistent")
	if err != sql.ErrNoRows {
		t.Errorf("expected sql.ErrNoRows, got %v", err)
	}
	if msg != "" {
		t.Errorf("expected empty string for non-existent session, got %q", msg)
	}
}

// TestIncrementTaskSessionUsage_AccumulatesAcrossCalls confirms multiple
// calls compound onto the same row. The DB-only columns are seeded via
// the migration's CREATE TABLE defaults (zero) and bumped via the
// UPDATE in the helper.
func TestIncrementTaskSessionUsage_AccumulatesAcrossCalls(t *testing.T) {
	repo := newRepoForSessionTests(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 {
		t.Fatalf("first increment: %v", err)
	}
	if err := repo.IncrementTaskSessionUsage(ctx, "sess-usage", 10, 20, 5); err != nil {
		t.Fatalf("second increment: %v", err)
	}

	var tokensIn, tokensOut, costSubcents int64
	err := repo.ro.QueryRowx(repo.ro.Rebind(
		`SELECT tokens_in, tokens_out, cost_subcents FROM task_sessions WHERE id = ?`),
		"sess-usage").Scan(&tokensIn, &tokensOut, &costSubcents)
	if err != nil {
		t.Fatalf("read row: %v", err)
	}
	if tokensIn != 110 || tokensOut != 220 || costSubcents != 55 {
		t.Errorf("totals = (%d,%d,%d), want (110,220,55)", tokensIn, tokensOut, costSubcents)
	}
}

// TestIncrementTaskSessionUsage_UnknownSessionNoError tolerates a
// 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 {
		t.Errorf("expected no error for unknown session, got %v", err)
	}
}

// TestIncrementTaskSessionUsage_EmptySessionIDNoOp guards against the
// 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 {
		t.Errorf("empty session id should be a no-op, got %v", err)
	}
}

// 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.
func seedRepoLink(t *testing.T, repo *Repository, workspaceID, repositoryID, taskID, sessionID, state string) {
	t.Helper()
	now := time.Now().UTC()
	_, err := repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO workspaces (id, name, created_at, updated_at)
		VALUES (?, 'ws', ?, ?)
	`), workspaceID, now, now)
	if err != nil {
		t.Fatalf("seed workspace: %v", err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO repositories (id, workspace_id, name, created_at, updated_at)
		VALUES (?, ?, 'repo', ?, ?)
	`), repositoryID, workspaceID, now, now)
	if err != nil {
		t.Fatalf("seed repository: %v", err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT OR IGNORE INTO tasks (id, workspace_id, title, created_at, updated_at)
		VALUES (?, ?, 'test task', ?, ?)
	`), taskID, workspaceID, now, now)
	if err != nil {
		t.Fatalf("seed task: %v", err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_repositories (id, task_id, repository_id, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?)
	`), "tr-"+taskID+"-"+repositoryID, taskID, repositoryID, now, now)
	if err != nil {
		t.Fatalf("seed task_repositories: %v", err)
	}
	_, err = repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_sessions (id, task_id, state, started_at, updated_at)
		VALUES (?, ?, ?, ?, ?)
	`), sessionID, taskID, state, now, now)
	if err != nil {
		t.Fatalf("seed task_session: %v", err)
	}
}

// TestCountActiveTaskSessionsByRepository_NoSessions verifies the count is
// zero when no sessions reference the repository at all.
func TestCountActiveTaskSessionsByRepository_NoSessions(t *testing.T) {
	repo := newRepoForSessionTests(t)
	count, err := repo.CountActiveTaskSessionsByRepository(context.Background(), "repo-empty")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if count != 0 {
		t.Errorf("expected 0, got %d", count)
	}
}

// TestCountActiveTaskSessionsByRepository_CountsActiveOnly verifies the join
// counts sessions in active or resumable states (CREATED, STARTING, RUNNING,
// IDLE, WAITING_FOR_INPUT) and excludes sessions in terminal states.
func TestCountActiveTaskSessionsByRepository_CountsActiveOnly(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	// Two active sessions across two tasks linked to the repo.
	seedRepoLink(t, repo, "ws-a", "repo-a", "task-a1", "sess-a1", "RUNNING")
	seedRepoLink(t, repo, "ws-a", "repo-a", "task-a2", "sess-a2", "WAITING_FOR_INPUT")
	seedRepoLink(t, repo, "ws-a", "repo-a", "task-a4", "sess-a4", "IDLE")
	// Terminal-state session linked to the repo — must NOT count.
	seedRepoLink(t, repo, "ws-a", "repo-a", "task-a3", "sess-a3", "COMPLETED")
	// Active session linked to a different repo — must NOT count.
	seedRepoLink(t, repo, "ws-a", "repo-b", "task-b1", "sess-b1", "RUNNING")

	count, err := repo.CountActiveTaskSessionsByRepository(ctx, "repo-a")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if count != 3 {
		t.Errorf("expected 3 active or resumable sessions, got %d", count)
	}
}

// TestCountActiveTaskSessionsByRepository_RequiresJoinRow verifies that a
// session whose task is NOT linked via task_repositories is not counted, even
// if the session is active. This guards against accidentally widening the
// query to use task_sessions.repository_id.
func TestCountActiveTaskSessionsByRepository_RequiresJoinRow(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	// Seed workspace, repo, and a task with the link.
	seedRepoLink(t, repo, "ws-j", "repo-j", "task-j1", "sess-j1", "RUNNING")

	// Seed a second task in the same workspace with an active session, but
	// without inserting a task_repositories row pointing at repo-j.
	now := time.Now().UTC()
	if _, err := repo.db.Exec(repo.db.Rebind(`
		INSERT INTO tasks (id, workspace_id, title, created_at, updated_at)
		VALUES ('task-j2', 'ws-j', 'orphan', ?, ?)
	`), now, now); err != nil {
		t.Fatalf("seed orphan task: %v", err)
	}
	if _, err := repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_sessions (id, task_id, state, started_at, updated_at)
		VALUES ('sess-j2', 'task-j2', 'RUNNING', ?, ?)
	`), now, now); err != nil {
		t.Fatalf("seed orphan session: %v", err)
	}

	count, err := repo.CountActiveTaskSessionsByRepository(ctx, "repo-j")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if count != 1 {
		t.Errorf("expected only the linked session to be counted, got %d", count)
	}
}

// archiveTask marks a seeded task as archived so the repo-delete guard tests
// can exercise the archived_at exclusion.
func archiveTask(t *testing.T, repo *Repository, taskID string) {
	t.Helper()
	if _, err := repo.db.Exec(repo.db.Rebind(
		`UPDATE tasks SET archived_at = ? WHERE id = ?`), time.Now().UTC(), taskID); err != nil {
		t.Fatalf("archive task %s: %v", taskID, err)
	}
}

// TestCountActiveTaskSessionsByRepository_ExcludesArchivedTasks verifies that an
// active session belonging to an archived task is not counted, so archived tasks
// never block repository deletion, while a live task's active session still is.
func TestCountActiveTaskSessionsByRepository_ExcludesArchivedTasks(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	// Active session on an archived task — must NOT count.
	seedRepoLink(t, repo, "ws-x", "repo-x", "task-x1", "sess-x1", "WAITING_FOR_INPUT")
	archiveTask(t, repo, "task-x1")

	count, err := repo.CountActiveTaskSessionsByRepository(ctx, "repo-x")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if count != 0 {
		t.Errorf("archived task must not block deletion, got %d active sessions", count)
	}

	// A live (non-archived) task with an active session on the same repo still
	// counts — pins that the archived_at filter did not over-broaden.
	seedRepoLink(t, repo, "ws-x", "repo-x", "task-x2", "sess-x2", "RUNNING")
	count, err = repo.CountActiveTaskSessionsByRepository(ctx, "repo-x")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if count != 1 {
		t.Errorf("live task session must still count, got %d", count)
	}
}

// TestHasActiveTaskSessionsByRepository_ExcludesArchivedTasks verifies the
// boolean delete guard mirrors the count: an archived task's active session
// does not report the repository as in use, but a live task's does.
func TestHasActiveTaskSessionsByRepository_ExcludesArchivedTasks(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedRepoLink(t, repo, "ws-h", "repo-h", "task-h1", "sess-h1", "WAITING_FOR_INPUT")
	archiveTask(t, repo, "task-h1")

	active, err := repo.HasActiveTaskSessionsByRepository(ctx, "repo-h")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if active {
		t.Error("archived task must not mark repository as having active sessions")
	}

	// Add a live task on the same repo — now it must report active.
	seedRepoLink(t, repo, "ws-h", "repo-h", "task-h2", "sess-h2", "RUNNING")
	active, err = repo.HasActiveTaskSessionsByRepository(ctx, "repo-h")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !active {
		t.Error("live task session must mark repository as active")
	}
}

// insertSession inserts a task_session row in the given state directly, for
// seeding multiple sessions on a single task (seedRepoLink can only create one
// session per task because its task_repositories PK is keyed on task+repo).
func insertSession(t *testing.T, repo *Repository, sessionID, taskID, state string) {
	t.Helper()
	now := time.Now().UTC()
	if _, err := repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_sessions (id, task_id, state, started_at, updated_at)
		VALUES (?, ?, ?, ?, ?)
	`), sessionID, taskID, state, now, now); err != nil {
		t.Fatalf("seed session %s: %v", sessionID, err)
	}
}

func sessionState(t *testing.T, repo *Repository, sessionID string) string {
	t.Helper()
	var state string
	if err := repo.ro.QueryRowx(repo.ro.Rebind(
		`SELECT state FROM task_sessions WHERE id = ?`), sessionID).Scan(&state); err != nil {
		t.Fatalf("read state for %s: %v", sessionID, err)
	}
	return state
}

func TestCancelActiveTaskSessionIsTerminalSafe(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-cas", "session-running", "turn-cas")
	if err := repo.UpdateTaskSessionState(ctx, "session-running", models.TaskSessionStateRunning, ""); err != nil {
		t.Fatalf("seed running state: %v", err)
	}
	insertSession(t, repo, "session-completed", "task-cas", string(models.TaskSessionStateCompleted))

	changed, cancelledAt, err := repo.CancelActiveTaskSession(ctx, "session-running", "coordinator stop")
	if err != nil {
		t.Fatalf("cancel running session: %v", err)
	}
	if !changed {
		t.Fatal("running session was not cancelled")
	}
	if got := sessionState(t, repo, "session-running"); got != string(models.TaskSessionStateCancelled) {
		t.Fatalf("running session state = %q, want CANCELLED", got)
	}
	cancelled, err := repo.GetTaskSession(ctx, "session-running")
	if err != nil {
		t.Fatalf("read cancelled session: %v", err)
	}
	if !cancelled.UpdatedAt.Equal(cancelledAt) {
		t.Fatalf("cancel timestamp = %s, stored updated_at = %s", cancelledAt, cancelled.UpdatedAt)
	}

	changed, _, err = repo.CancelActiveTaskSession(ctx, "session-completed", "coordinator stop")
	if err != nil {
		t.Fatalf("cancel completed session: %v", err)
	}
	if changed {
		t.Fatal("completed session reported a cancellation")
	}
	if got := sessionState(t, repo, "session-completed"); got != string(models.TaskSessionStateCompleted) {
		t.Fatalf("completed session state = %q, want COMPLETED", got)
	}
}

func TestUpdateTaskSessionStateIfCurrentRejectsStaleActiveWriter(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-state-cas", "session-state-cas", "turn-state-cas")
	if err := repo.UpdateTaskSessionState(ctx, "session-state-cas", models.TaskSessionStateRunning, ""); err != nil {
		t.Fatalf("seed running state: %v", err)
	}
	if changed, _, err := repo.CancelActiveTaskSession(ctx, "session-state-cas", "coordinator stop"); err != nil || !changed {
		t.Fatalf("cancel session: changed=%v err=%v", changed, err)
	}

	changed, _, err := repo.UpdateTaskSessionStateIfCurrent(
		ctx,
		"session-state-cas",
		models.TaskSessionStateRunning,
		models.TaskSessionStateWaitingForInput,
		"",
	)
	if err != nil {
		t.Fatalf("stale conditional update: %v", err)
	}
	if changed {
		t.Fatal("stale RUNNING writer changed a CANCELLED session")
	}
	if got := sessionState(t, repo, "session-state-cas"); got != string(models.TaskSessionStateCancelled) {
		t.Fatalf("session state = %q, want CANCELLED", got)
	}
}

func TestUpdateTaskSessionIfCurrentStateRejectsStaleFullRowWriter(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-full-row-cas", "session-full-row-cas", "turn-full-row-cas")
	require.NoError(t, repo.UpdateTaskSessionState(
		ctx, "session-full-row-cas", models.TaskSessionStateRunning, "",
	))
	stale, err := repo.GetTaskSession(ctx, "session-full-row-cas")
	require.NoError(t, err)
	require.Equal(t, models.TaskSessionStateRunning, stale.State)
	changed, _, err := repo.CancelActiveTaskSession(ctx, stale.ID, "stopped by parent task via MCP")
	require.NoError(t, err)
	require.True(t, changed)

	stale.State = models.TaskSessionStateStarting
	stale.ExecutorID = "late-executor"
	changed, err = repo.UpdateTaskSessionIfCurrentState(
		ctx, stale, models.TaskSessionStateRunning,
	)
	require.NoError(t, err)
	require.False(t, changed)
	stored, err := repo.GetTaskSession(ctx, stale.ID)
	require.NoError(t, err)
	require.Equal(t, models.TaskSessionStateCancelled, stored.State)
	require.Empty(t, stored.ExecutorID)
}

func TestUpdateTaskSessionWithMetadataRejectsInvalidMetadataBeforeStateWrite(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-atomic", "sess-atomic", "turn-atomic")
	if err := repo.UpdateSessionMetadata(ctx, "sess-atomic", map[string]interface{}{"keep": "yes"}); err != nil {
		t.Fatalf("seed metadata: %v", err)
	}
	session, err := repo.GetTaskSession(ctx, "sess-atomic")
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	session.State = models.TaskSessionStateFailed

	err = repo.UpdateTaskSessionWithMetadata(ctx, session, map[string]interface{}{"bad": func() {}})
	if err == nil {
		t.Fatal("expected invalid metadata error")
	}
	if got := sessionState(t, repo, "sess-atomic"); got == string(models.TaskSessionStateFailed) {
		t.Fatalf("state was partially updated to %q", got)
	}
	got, err := repo.GetTaskSession(ctx, "sess-atomic")
	if err != nil {
		t.Fatalf("get session after failed update: %v", err)
	}
	if got.Metadata["keep"] != "yes" {
		t.Fatalf("metadata keep = %v, want yes", got.Metadata["keep"])
	}
}

func TestUpdateTaskSessionIfCurrentStateRemovingMetadataKeys(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-remove-metadata", "sess-remove-metadata", "turn-remove-metadata")
	require.NoError(t, repo.UpdateSessionMetadata(ctx, "sess-remove-metadata", map[string]interface{}{
		"provider_state": "stale",
		"keep":           "newer",
	}))
	session, err := repo.GetTaskSession(ctx, "sess-remove-metadata")
	require.NoError(t, err)
	changed, err := repo.UpdateTaskSessionIfCurrentStateRemovingMetadataKeys(
		ctx,
		session,
		models.TaskSessionStateCreated,
		[]string{"provider_state"},
	)
	require.NoError(t, err)
	require.True(t, changed)
	stored, err := repo.GetTaskSession(ctx, "sess-remove-metadata")
	require.NoError(t, err)
	require.NotContains(t, stored.Metadata, "provider_state")
	require.Equal(t, "newer", stored.Metadata["keep"])
}

func TestUpdateTaskSessionIfCurrentStateRemovingMetadataKeysStateMismatch(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedForMsgTest(t, repo, "task-remove-metadata-mismatch", "sess-remove-metadata-mismatch", "turn-remove-metadata-mismatch")
	require.NoError(t, repo.UpdateSessionMetadata(ctx, "sess-remove-metadata-mismatch", map[string]interface{}{
		"provider_state": "stale",
		"keep":           "untouched",
	}))
	session, err := repo.GetTaskSession(ctx, "sess-remove-metadata-mismatch")
	require.NoError(t, err)
	changed, err := repo.UpdateTaskSessionIfCurrentStateRemovingMetadataKeys(
		ctx,
		session,
		models.TaskSessionStateRunning,
		[]string{"provider_state"},
	)
	require.NoError(t, err)
	require.False(t, changed)
	stored, err := repo.GetTaskSession(ctx, "sess-remove-metadata-mismatch")
	require.NoError(t, err)
	require.Equal(t, "stale", stored.Metadata["provider_state"])
	require.Equal(t, "untouched", stored.Metadata["keep"])
}

func TestDismissLastAgentErrorDoesNotOverwriteNewerError(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-error", "sess-error", "turn-error")
	oldErr := models.LastAgentError{
		Message:    "old error",
		OccurredAt: time.Date(2026, 6, 14, 10, 0, 0, 0, time.UTC),
	}
	newErr := models.LastAgentError{
		Message:    "new error",
		OccurredAt: time.Date(2026, 6, 14, 10, 5, 0, 0, time.UTC),
	}
	if err := repo.SetSessionMetadataKey(ctx, "sess-error", models.SessionMetaKeyLastAgentError, oldErr); err != nil {
		t.Fatalf("seed old error: %v", err)
	}
	if err := repo.SetSessionMetadataKey(ctx, "sess-error", models.SessionMetaKeyLastAgentError, newErr); err != nil {
		t.Fatalf("seed new error: %v", err)
	}

	updated, err := repo.DismissLastAgentError(ctx, "sess-error", oldErr, time.Now().UTC())
	if err != nil {
		t.Fatalf("dismiss stale error: %v", err)
	}
	if updated {
		t.Fatalf("expected stale dismiss to be ignored")
	}
	session, err := repo.GetTaskSession(ctx, "sess-error")
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	got, ok := models.LoadLastAgentError(session.Metadata)
	if !ok {
		t.Fatalf("expected last agent error metadata")
	}
	if got.Message != newErr.Message || got.IsDismissed() {
		t.Fatalf("last agent error = %#v, want undismissed newer error", got)
	}
}

func TestDismissLastAgentErrorMatchesEquivalentTimestampText(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	seedForMsgTest(t, repo, "task-error", "sess-error", "turn-error")
	occurredAt, err := time.Parse(time.RFC3339Nano, "2026-06-14T12:00:00.310Z")
	if err != nil {
		t.Fatalf("parse occurred_at: %v", err)
	}
	lastErr := models.LastAgentError{
		Message:    "peer disconnected before response",
		OccurredAt: occurredAt,
	}
	if err := repo.SetSessionMetadataKey(ctx, "sess-error", models.SessionMetaKeyLastAgentError, map[string]any{
		"message":     lastErr.Message,
		"occurred_at": "2026-06-14T12:00:00.310Z",
	}); err != nil {
		t.Fatalf("seed last agent error: %v", err)
	}

	updated, err := repo.DismissLastAgentError(ctx, "sess-error", lastErr, time.Now().UTC())
	if err != nil {
		t.Fatalf("dismiss last agent error: %v", err)
	}
	if !updated {
		t.Fatalf("expected equivalent timestamp text to match")
	}
	session, err := repo.GetTaskSession(ctx, "sess-error")
	if err != nil {
		t.Fatalf("get session: %v", err)
	}
	got, ok := models.LoadLastAgentError(session.Metadata)
	if !ok {
		t.Fatalf("expected last agent error metadata")
	}
	if !got.IsDismissed() {
		t.Fatalf("last agent error = %#v, want dismissed", got)
	}
}

// Nothing used to retire a stored agent failure, so installations carry records
// whose errors the agent recovered from long ago — each one still painting a red
// icon on the task list.
func TestClearRecoveredAgentErrorsBackfill(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	occurredAt := time.Date(2026, 6, 14, 10, 0, 0, 0, time.UTC)
	seedStale := func(t *testing.T, taskID, sessionID, turnID string) {
		t.Helper()
		seedForMsgTest(t, repo, taskID, sessionID, turnID)
		if err := repo.SetSessionMetadataKey(ctx, sessionID, models.SessionMetaKeyLastAgentError,
			models.LastAgentError{Message: "agent crashed", OccurredAt: occurredAt}); err != nil {
			t.Fatalf("seed last agent error on %s: %v", sessionID, err)
		}
	}

	// Recovered: an ordinary agent message lands after the failure.
	seedStale(t, "task-recovered", "sess-recovered", "turn-recovered")
	insertAgentMsg(t, repo, "msg-recovered", "sess-recovered", "turn-recovered",
		"agent", "back on track", occurredAt.Add(time.Hour))

	// Still failing: nothing successful has happened since.
	seedStale(t, "task-current", "sess-current", "turn-current")

	// Recovery that predates the failure must not count.
	seedStale(t, "task-older", "sess-older", "turn-older")
	insertAgentMsg(t, repo, "msg-older", "sess-older", "turn-older",
		"agent", "earlier output", occurredAt.Add(-time.Hour))

	// A user message after the failure is not the agent producing good output.
	seedStale(t, "task-user", "sess-user", "turn-user")
	insertAgentMsg(t, repo, "msg-user", "sess-user", "turn-user",
		"user", "are you stuck?", occurredAt.Add(time.Hour))

	// These statements scan every session before filtering, so a row the
	// repository wrote as "no metadata" must not abort them. migrate.Apply
	// swallows SQL errors, so the proof is that the recovered session below is
	// still cleared with these rows present — an aborted statement would leave
	// it untouched.
	for index, blank := range []string{"", "null"} {
		id := fmt.Sprintf("sess-blank-%d", index)
		seedForMsgTest(t, repo, "task-blank-"+strconv.Itoa(index), id, "turn-blank-"+strconv.Itoa(index))
		if _, err := repo.db.Exec(repo.db.Rebind(
			`UPDATE task_sessions SET metadata = ? WHERE id = ?`), blank, id); err != nil {
			t.Fatalf("seed %q metadata: %v", blank, err)
		}
	}

	// The cached summary keeps the icon up until it is cleared too.
	if _, err := repo.db.Exec(repo.db.Rebind(`
		INSERT INTO task_status_summaries (task_id, workspace_id, revision, summary, updated_at)
		VALUES (?, '', 4, ?, ?)
	`), "task-recovered",
		`{"active_error":{"session_id":"sess-recovered","stamp":"s","preview":"agent crashed"},"pending_action":"permission"}`,
		time.Now().UTC()); err != nil {
		t.Fatalf("seed status summary: %v", err)
	}

	if err := repo.clearRecoveredAgentErrors(); err != nil {
		t.Fatalf("clearRecoveredAgentErrors: %v", err)
	}

	hasError := func(t *testing.T, sessionID string) bool {
		t.Helper()
		session, err := repo.GetTaskSession(ctx, sessionID)
		if err != nil {
			t.Fatalf("get session %s: %v", sessionID, err)
		}
		_, ok := models.LoadLastAgentError(session.Metadata)
		return ok
	}

	if hasError(t, "sess-recovered") {
		t.Fatalf("a failure the agent recovered from must be cleared")
	}
	for _, sessionID := range []string{"sess-current", "sess-older", "sess-user"} {
		if !hasError(t, sessionID) {
			t.Fatalf("%s has no successful agent work after the failure; it must stay", sessionID)
		}
	}

	var summary string
	if err := repo.db.QueryRow(repo.db.Rebind(
		`SELECT summary FROM task_status_summaries WHERE task_id = ?`), "task-recovered",
	).Scan(&summary); err != nil {
		t.Fatalf("read status summary: %v", err)
	}
	if strings.Contains(summary, "active_error") {
		t.Fatalf("summary = %s, want the cached error cleared", summary)
	}
	if !strings.Contains(summary, "permission") {
		t.Fatalf("summary = %s, want the rest of the projection preserved", summary)
	}
}

func sessionCancellationMetadata(t *testing.T, repo *Repository, sessionID string) (string, sql.NullTime, time.Time) {
	t.Helper()
	var errorMessage string
	var completedAt sql.NullTime
	var updatedAt time.Time
	if err := repo.ro.QueryRowx(repo.ro.Rebind(
		`SELECT error_message, completed_at, updated_at FROM task_sessions WHERE id = ?`),
		sessionID,
	).Scan(&errorMessage, &completedAt, &updatedAt); err != nil {
		t.Fatalf("read cancellation metadata for %s: %v", sessionID, err)
	}
	return errorMessage, completedAt, updatedAt
}

func assertReapedSession(t *testing.T, repo *Repository, sessionID string, reapedAfter time.Time) {
	t.Helper()
	if got := sessionState(t, repo, sessionID); got != "CANCELLED" {
		t.Errorf("%s = %q, want CANCELLED", sessionID, got)
	}
	errorMessage, completedAt, updatedAt := sessionCancellationMetadata(t, repo, sessionID)
	if errorMessage != "task archived" {
		t.Errorf("%s error_message = %q, want task archived", sessionID, errorMessage)
	}
	if !completedAt.Valid {
		t.Errorf("%s completed_at should be set", sessionID)
	}
	if updatedAt.Before(reapedAfter) {
		t.Errorf("%s updated_at = %s, want >= %s", sessionID, updatedAt, reapedAfter)
	}
}

// TestCancelActiveTaskSessionsByTaskID verifies the archive reaper transitions
// only the target task's still-active sessions to CANCELLED, leaves terminal
// sessions and other tasks untouched, and reports exactly the sessions it
// changed. It also confirms the repo-delete guard reports the repository as
// free afterward — the end-to-end purpose of the reap.
func TestCancelActiveTaskSessionsByTaskID(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	// Target task: one active session via the link helper, plus a second active
	// session, two pre-run sessions, and an already-terminal session inserted directly.
	seedRepoLink(t, repo, "ws-r", "repo-r", "task-r", "sess-r1", "WAITING_FOR_INPUT")
	insertSession(t, repo, "sess-r2", "task-r", "RUNNING")
	insertSession(t, repo, "sess-r3", "task-r", "CREATED")
	insertSession(t, repo, "sess-r4", "task-r", "STARTING")
	insertSession(t, repo, "sess-r5", "task-r", "COMPLETED")
	// A different task on a different repo — must be untouched.
	seedRepoLink(t, repo, "ws-r", "repo-other", "task-other", "sess-o1", "RUNNING")

	reapedAfter := time.Now().UTC()
	sessions, err := repo.CancelActiveTaskSessionsByTaskID(ctx, "task-r", "task archived")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	wantIDs := []string{"sess-r1", "sess-r2", "sess-r3", "sess-r4"}
	var gotIDs []string
	for _, sess := range sessions {
		gotIDs = append(gotIDs, sess.ID)
		if sess.TaskID != "task-r" {
			t.Errorf("session %s TaskID = %q, want task-r", sess.ID, sess.TaskID)
		}
		if sess.State != models.TaskSessionStateCancelled {
			t.Errorf("session %s State = %q, want CANCELLED", sess.ID, sess.State)
		}
		if sess.UpdatedAt.Before(reapedAfter) {
			t.Errorf("session %s UpdatedAt = %s, want >= %s", sess.ID, sess.UpdatedAt, reapedAfter)
		}
	}
	if !sameStringSet(gotIDs, wantIDs) {
		t.Errorf("cancelled ids = %v, want %v", gotIDs, wantIDs)
	}
	for _, sessionID := range wantIDs {
		assertReapedSession(t, repo, sessionID, reapedAfter)
	}
	if got := sessionState(t, repo, "sess-r5"); got != "COMPLETED" {
		t.Errorf("sess-r5 (terminal) = %q, want unchanged COMPLETED", got)
	}
	if got := sessionState(t, repo, "sess-o1"); got != "RUNNING" {
		t.Errorf("sess-o1 (other task) = %q, want unchanged RUNNING", got)
	}

	// End-to-end: the repository that only the reaped task referenced is now
	// reported as free by the delete guard.
	active, err := repo.HasActiveTaskSessionsByRepository(ctx, "repo-r")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if active {
		t.Error("repo-r should have no active sessions after reaping its task")
	}

	// Idempotent: a second call changes nothing.
	sessions, err = repo.CancelActiveTaskSessionsByTaskID(ctx, "task-r", "task archived")
	if err != nil {
		t.Fatalf("unexpected error on second call: %v", err)
	}
	if len(sessions) != 0 {
		t.Errorf("expected 0 sessions on idempotent re-run, got %v", sessions)
	}
}

// sameStringSet reports whether got and want contain the same strings,
// ignoring order and duplicates count-for-count — used to compare the set of
// cancelled session IDs against an expected set regardless of return order.
func sameStringSet(got, want []string) bool {
	if len(got) != len(want) {
		return false
	}
	set := make(map[string]int, len(want))
	for _, s := range want {
		set[s]++
	}
	for _, s := range got {
		set[s]--
	}
	for _, count := range set {
		if count != 0 {
			return false
		}
	}
	return true
}

// --- Turn lifecycle -------------------------------------------------------
//
// task_session_turns is the analytics/duration record behind the UI's
// last-turn readout. AbandonTurn's zero-duration close is the subtle part:
// it must collapse an orphaned turn to started_at rather than to now.

func seedSessionForTurns(t *testing.T, repo *Repository, taskID, sessionID string) {
	t.Helper()
	ctx := context.Background()
	if err := repo.CreateTask(ctx, &models.Task{ID: taskID, Title: taskID}); err != nil {
		t.Fatalf("CreateTask(%s): %v", taskID, err)
	}
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID: sessionID, TaskID: taskID, State: models.TaskSessionStateRunning,
	}); err != nil {
		t.Fatalf("CreateTaskSession(%s): %v", sessionID, err)
	}
}

func TestCreateTurnRoundTripsEveryFieldAndDefaultsTimestamps(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedSessionForTurns(t, repo, "task-turn-full", "session-turn-full")

	completedAt := time.Date(2026, 5, 5, 10, 30, 0, 0, time.UTC)
	want := &models.Turn{
		ID:            "turn-full",
		TaskSessionID: "session-turn-full",
		TaskID:        "task-turn-full",
		StartedAt:     time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC),
		CompletedAt:   &completedAt,
		Metadata:      map[string]interface{}{"tokens": float64(1234), "model": "opus"},
		CreatedAt:     time.Date(2026, 5, 5, 9, 59, 0, 0, time.UTC),
	}
	if err := repo.CreateTurn(ctx, want); err != nil {
		t.Fatalf("CreateTurn: %v", err)
	}

	got, err := repo.GetTurn(ctx, "turn-full")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	if got.ID != want.ID || got.TaskSessionID != want.TaskSessionID || got.TaskID != want.TaskID {
		t.Errorf("identity = %q/%q/%q, want %q/%q/%q",
			got.ID, got.TaskSessionID, got.TaskID, want.ID, want.TaskSessionID, want.TaskID)
	}
	assertTimeEqual(t, "StartedAt", got.StartedAt, want.StartedAt)
	assertTimeEqual(t, "CreatedAt", got.CreatedAt, want.CreatedAt)
	assertTimeEqual(t, "UpdatedAt", got.UpdatedAt, want.UpdatedAt)
	if got.CompletedAt == nil {
		t.Fatal("CompletedAt = nil, want the supplied value")
	}
	assertTimeEqual(t, "CompletedAt", *got.CompletedAt, completedAt)
	assertJSONMapEqual(t, "Metadata", got.Metadata, want.Metadata)

	// An id-less, timestamp-less turn gets both stamped.
	before := time.Now().UTC().Add(-time.Second)
	bare := &models.Turn{TaskSessionID: "session-turn-full", TaskID: "task-turn-full"}
	if err := repo.CreateTurn(ctx, bare); err != nil {
		t.Fatalf("CreateTurn(bare): %v", err)
	}
	if bare.ID == "" {
		t.Error("ID left empty; a UUID should have been generated")
	}
	if bare.StartedAt.Before(before) || bare.CreatedAt.Before(before) {
		t.Errorf("StartedAt=%v CreatedAt=%v, want fresh stamps after %v", bare.StartedAt, bare.CreatedAt, before)
	}
	gotBare, err := repo.GetTurn(ctx, bare.ID)
	if err != nil {
		t.Fatalf("GetTurn(bare): %v", err)
	}
	if gotBare.CompletedAt != nil {
		t.Errorf("CompletedAt = %v, want nil for an open turn", *gotBare.CompletedAt)
	}
	if gotBare.Metadata != nil {
		t.Errorf("Metadata = %#v, want nil (the {} sentinel decodes to nil)", gotBare.Metadata)
	}
}

func TestGetTurnAndActiveTurnReturnErrNoRowsWhenAbsent(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedSessionForTurns(t, repo, "task-turn-none", "session-turn-none")

	if _, err := repo.GetTurn(ctx, "turn-missing"); !errors.Is(err, sql.ErrNoRows) {
		t.Errorf("GetTurn error = %v, want sql.ErrNoRows", err)
	}
	if _, err := repo.GetActiveTurnBySessionID(ctx, "session-turn-none"); !errors.Is(err, sql.ErrNoRows) {
		t.Errorf("GetActiveTurnBySessionID error = %v, want sql.ErrNoRows", err)
	}
	turns, err := repo.ListTurnsBySession(ctx, "session-turn-none")
	if err != nil {
		t.Fatalf("ListTurnsBySession: %v", err)
	}
	if len(turns) != 0 {
		t.Errorf("ListTurnsBySession = %+v, want empty", turns)
	}
}

func TestGetActiveTurnBySessionIDPicksNewestOpenTurn(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedSessionForTurns(t, repo, "task-turn-active", "session-turn-active")
	seedSessionForTurns(t, repo, "task-turn-other", "session-turn-other")

	base := time.Date(2026, 6, 1, 8, 0, 0, 0, time.UTC)
	closedAt := base.Add(30 * time.Minute)
	turns := []*models.Turn{
		{ID: "turn-closed", TaskSessionID: "session-turn-active", TaskID: "task-turn-active",
			StartedAt: base, CompletedAt: &closedAt},
		{ID: "turn-open-old", TaskSessionID: "session-turn-active", TaskID: "task-turn-active",
			StartedAt: base.Add(time.Hour)},
		{ID: "turn-open-new", TaskSessionID: "session-turn-active", TaskID: "task-turn-active",
			StartedAt: base.Add(2 * time.Hour)},
		{ID: "turn-foreign", TaskSessionID: "session-turn-other", TaskID: "task-turn-other",
			StartedAt: base.Add(3 * time.Hour)},
	}
	for _, turn := range turns {
		if err := repo.CreateTurn(ctx, turn); err != nil {
			t.Fatalf("CreateTurn(%s): %v", turn.ID, err)
		}
	}

	active, err := repo.GetActiveTurnBySessionID(ctx, "session-turn-active")
	if err != nil {
		t.Fatalf("GetActiveTurnBySessionID: %v", err)
	}
	if active.ID != "turn-open-new" {
		t.Errorf("GetActiveTurnBySessionID = %q, want turn-open-new (newest open turn)", active.ID)
	}

	listed, err := repo.ListTurnsBySession(ctx, "session-turn-active")
	if err != nil {
		t.Fatalf("ListTurnsBySession: %v", err)
	}
	wantOrder := []string{"turn-closed", "turn-open-old", "turn-open-new"}
	if len(listed) != len(wantOrder) {
		t.Fatalf("ListTurnsBySession returned %d turns, want %d (session-scoped)", len(listed), len(wantOrder))
	}
	for i := range wantOrder {
		if listed[i].ID != wantOrder[i] {
			t.Fatalf("ListTurnsBySession order = %q at index %d, want %q (started_at ASC)", listed[i].ID, i, wantOrder[i])
		}
	}
	if listed[0].CompletedAt == nil {
		t.Error("turn-closed CompletedAt = nil, want it populated in the list read")
	}
	if listed[1].CompletedAt != nil {
		t.Errorf("turn-open-old CompletedAt = %v, want nil", *listed[1].CompletedAt)
	}
}

func TestUpdateTurnWritesCompletionAndMetadata(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedSessionForTurns(t, repo, "task-turn-update", "session-turn-update")

	turn := &models.Turn{
		ID: "turn-update", TaskSessionID: "session-turn-update", TaskID: "task-turn-update",
		StartedAt: time.Date(2026, 6, 2, 8, 0, 0, 0, time.UTC),
		Metadata:  map[string]interface{}{"stage": "start"},
	}
	if err := repo.CreateTurn(ctx, turn); err != nil {
		t.Fatalf("CreateTurn: %v", err)
	}

	completedAt := time.Date(2026, 6, 2, 8, 45, 0, 0, time.UTC)
	turn.CompletedAt = &completedAt
	turn.Metadata = map[string]interface{}{"stage": "done", "tokens": float64(42)}
	if err := repo.UpdateTurn(ctx, turn); err != nil {
		t.Fatalf("UpdateTurn: %v", err)
	}

	got, err := repo.GetTurn(ctx, "turn-update")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	if got.CompletedAt == nil {
		t.Fatal("CompletedAt = nil after UpdateTurn")
	}
	assertTimeEqual(t, "CompletedAt", *got.CompletedAt, completedAt)
	assertJSONMapEqual(t, "Metadata", got.Metadata, turn.Metadata)
	assertTimeEqual(t, "StartedAt", got.StartedAt, turn.StartedAt)
	if !got.UpdatedAt.After(got.StartedAt) {
		t.Errorf("UpdatedAt = %v, want it bumped past StartedAt %v", got.UpdatedAt, got.StartedAt)
	}

	// Clearing Metadata writes the {} sentinel back, which reads as nil.
	turn.Metadata = nil
	if err := repo.UpdateTurn(ctx, turn); err != nil {
		t.Fatalf("UpdateTurn(clear metadata): %v", err)
	}
	got, err = repo.GetTurn(ctx, "turn-update")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	if got.Metadata != nil {
		t.Errorf("Metadata = %#v, want nil after being cleared", got.Metadata)
	}
}

func TestCompleteTurnStampsNowAndAbandonTurnCollapsesToStartedAt(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()
	seedSessionForTurns(t, repo, "task-turn-close", "session-turn-close")

	startedAt := time.Date(2026, 6, 3, 8, 0, 0, 0, time.UTC)
	for _, id := range []string{"turn-complete", "turn-abandon", "turn-already-closed"} {
		if err := repo.CreateTurn(ctx, &models.Turn{
			ID: id, TaskSessionID: "session-turn-close", TaskID: "task-turn-close", StartedAt: startedAt,
		}); err != nil {
			t.Fatalf("CreateTurn(%s): %v", id, err)
		}
	}

	if err := repo.CompleteTurn(ctx, "turn-complete"); err != nil {
		t.Fatalf("CompleteTurn: %v", err)
	}
	completed, err := repo.GetTurn(ctx, "turn-complete")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	if completed.CompletedAt == nil {
		t.Fatal("CompletedAt = nil after CompleteTurn")
	}
	if !completed.CompletedAt.After(startedAt) {
		t.Errorf("CompletedAt = %v, want a wall-clock stamp after StartedAt %v", *completed.CompletedAt, startedAt)
	}

	// AbandonTurn gives the orphaned turn zero duration rather than hours of
	// dead time.
	if err := repo.AbandonTurn(ctx, "turn-abandon"); err != nil {
		t.Fatalf("AbandonTurn: %v", err)
	}
	abandoned, err := repo.GetTurn(ctx, "turn-abandon")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	if abandoned.CompletedAt == nil {
		t.Fatal("CompletedAt = nil after AbandonTurn")
	}
	assertTimeEqual(t, "abandoned CompletedAt", *abandoned.CompletedAt, startedAt)

	// AbandonTurn must not rewrite a turn that is already closed.
	closedAt := startedAt.Add(90 * time.Minute)
	alreadyClosed, err := repo.GetTurn(ctx, "turn-already-closed")
	if err != nil {
		t.Fatalf("GetTurn(turn-already-closed): %v", err)
	}
	alreadyClosed.CompletedAt = &closedAt
	if err := repo.UpdateTurn(ctx, alreadyClosed); err != nil {
		t.Fatalf("UpdateTurn: %v", err)
	}
	if err := repo.AbandonTurn(ctx, "turn-already-closed"); err != nil {
		t.Fatalf("AbandonTurn(already closed): %v", err)
	}
	stillClosed, err := repo.GetTurn(ctx, "turn-already-closed")
	if err != nil {
		t.Fatalf("GetTurn: %v", err)
	}
	assertTimeEqual(t, "already-closed CompletedAt", *stillClosed.CompletedAt, closedAt)
}

// --- Batch and active-session reads --------------------------------------
//
// These feed the kanban card summaries and the environment-sharing guards.
// They share one fixture: three tasks with a mix of primary/secondary and
// active/terminal sessions on two task environments.

type sessionBatchFixture struct {
	repo *Repository
}

func seedSessionBatchFixture(t *testing.T) *sessionBatchFixture {
	t.Helper()
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	for _, taskID := range []string{"task-batch-1", "task-batch-2", "task-batch-3"} {
		if err := repo.CreateTask(ctx, &models.Task{ID: taskID, Title: taskID}); err != nil {
			t.Fatalf("CreateTask(%s): %v", taskID, err)
		}
	}
	// One environment per task (task_environments.task_id is UNIQUE); the
	// "shared" one is owned by task-batch-1 and borrowed by task-batch-2.
	for envID, ownerTaskID := range map[string]string{
		"env-batch-shared": "task-batch-1",
		"env-batch-solo":   "task-batch-3",
	} {
		if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
			ID: envID, TaskID: ownerTaskID, ExecutorType: string(models.ExecutorTypeLocal),
			Status: models.TaskEnvironmentStatusReady,
		}); err != nil {
			t.Fatalf("CreateTaskEnvironment(%s): %v", envID, err)
		}
	}

	sessions := []*models.TaskSession{
		{ID: "session-b1-primary", TaskID: "task-batch-1", IsPrimary: true,
			State: models.TaskSessionStateRunning, TaskEnvironmentID: "env-batch-shared",
			EnvironmentID: "env-batch-shared", ReviewStatus: models.ReviewStatusPending},
		{ID: "session-b1-secondary", TaskID: "task-batch-1",
			State: models.TaskSessionStateCompleted, TaskEnvironmentID: "env-batch-shared"},
		{ID: "session-b2-primary", TaskID: "task-batch-2", IsPrimary: true,
			State: models.TaskSessionStateWaitingForInput, TaskEnvironmentID: "env-batch-shared"},
		// task-batch-3 has only a non-primary, terminal session.
		{ID: "session-b3-secondary", TaskID: "task-batch-3",
			State: models.TaskSessionStateCompleted, TaskEnvironmentID: "env-batch-solo"},
	}
	for _, session := range sessions {
		if err := repo.CreateTaskSession(ctx, session); err != nil {
			t.Fatalf("CreateTaskSession(%s): %v", session.ID, err)
		}
	}
	return &sessionBatchFixture{repo: repo}
}

func TestGetPrimarySessionIDsByTaskIDsOmitsTasksWithoutAPrimary(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()

	got, err := fixture.repo.GetPrimarySessionIDsByTaskIDs(ctx,
		[]string{"task-batch-1", "task-batch-2", "task-batch-3", "task-batch-missing"})
	if err != nil {
		t.Fatalf("GetPrimarySessionIDsByTaskIDs: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("map = %v, want exactly the two tasks with a primary session", got)
	}
	if got["task-batch-1"] != "session-b1-primary" {
		t.Errorf("task-batch-1 = %q, want session-b1-primary", got["task-batch-1"])
	}
	if got["task-batch-2"] != "session-b2-primary" {
		t.Errorf("task-batch-2 = %q, want session-b2-primary", got["task-batch-2"])
	}

	empty, err := fixture.repo.GetPrimarySessionIDsByTaskIDs(ctx, nil)
	if err != nil {
		t.Fatalf("GetPrimarySessionIDsByTaskIDs(nil): %v", err)
	}
	if empty == nil || len(empty) != 0 {
		t.Errorf("nil input = %v, want a non-nil empty map", empty)
	}
}

func TestGetSessionCountsByTaskIDsCountsEverySessionState(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()

	got, err := fixture.repo.GetSessionCountsByTaskIDs(ctx,
		[]string{"task-batch-1", "task-batch-2", "task-batch-3", "task-batch-missing"})
	if err != nil {
		t.Fatalf("GetSessionCountsByTaskIDs: %v", err)
	}
	want := map[string]int{"task-batch-1": 2, "task-batch-2": 1, "task-batch-3": 1}
	if len(got) != len(want) {
		t.Fatalf("map = %v, want %v (tasks with no sessions are omitted)", got, want)
	}
	for taskID, count := range want {
		if got[taskID] != count {
			t.Errorf("count[%s] = %d, want %d", taskID, got[taskID], count)
		}
	}

	empty, err := fixture.repo.GetSessionCountsByTaskIDs(ctx, nil)
	if err != nil {
		t.Fatalf("GetSessionCountsByTaskIDs(nil): %v", err)
	}
	if empty == nil || len(empty) != 0 {
		t.Errorf("nil input = %v, want a non-nil empty map", empty)
	}
}

func TestGetPrimarySessionInfoByTaskIDsProjectsStateAndExecutor(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()

	got, err := fixture.repo.GetPrimarySessionInfoByTaskIDs(ctx,
		[]string{"task-batch-1", "task-batch-2", "task-batch-3"})
	if err != nil {
		t.Fatalf("GetPrimarySessionInfoByTaskIDs: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("map has %d entries (%v), want 2 — task-batch-3 has no primary", len(got), got)
	}
	first := got["task-batch-1"]
	if first == nil {
		t.Fatal("task-batch-1 missing from the result")
	}
	if first.ID != "session-b1-primary" || first.TaskID != "task-batch-1" {
		t.Errorf("identity = %q/%q, want session-b1-primary/task-batch-1", first.ID, first.TaskID)
	}
	if first.State != models.TaskSessionStateRunning {
		t.Errorf("State = %q, want RUNNING", first.State)
	}
	if first.ReviewStatus != models.ReviewStatusPending {
		t.Errorf("ReviewStatus = %q, want %q", first.ReviewStatus, models.ReviewStatusPending)
	}
	second := got["task-batch-2"]
	if second == nil || second.State != models.TaskSessionStateWaitingForInput {
		t.Errorf("task-batch-2 = %+v, want a WAITING_FOR_INPUT primary session", second)
	}

	empty, err := fixture.repo.GetPrimarySessionInfoByTaskIDs(ctx, nil)
	if err != nil {
		t.Fatalf("GetPrimarySessionInfoByTaskIDs(nil): %v", err)
	}
	if empty == nil || len(empty) != 0 {
		t.Errorf("nil input = %v, want a non-nil empty map", empty)
	}
}

func TestListActiveTaskSessionsFiltersTerminalStates(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()

	all, err := fixture.repo.ListActiveTaskSessions(ctx)
	if err != nil {
		t.Fatalf("ListActiveTaskSessions: %v", err)
	}
	activeIDs := make(map[string]bool, len(all))
	for _, session := range all {
		activeIDs[session.ID] = true
	}
	if len(all) != 2 {
		t.Fatalf("ListActiveTaskSessions returned %d sessions (%v), want 2", len(all), activeIDs)
	}
	if !activeIDs["session-b1-primary"] || !activeIDs["session-b2-primary"] {
		t.Errorf("active sessions = %v, want the RUNNING and WAITING_FOR_INPUT rows", activeIDs)
	}
	if activeIDs["session-b1-secondary"] || activeIDs["session-b3-secondary"] {
		t.Errorf("active sessions = %v, must exclude COMPLETED rows", activeIDs)
	}

	scoped, err := fixture.repo.ListActiveTaskSessionsByTaskID(ctx, "task-batch-1")
	if err != nil {
		t.Fatalf("ListActiveTaskSessionsByTaskID: %v", err)
	}
	if len(scoped) != 1 || scoped[0].ID != "session-b1-primary" {
		t.Errorf("task-scoped active sessions = %+v, want only session-b1-primary", scoped)
	}
	none, err := fixture.repo.ListActiveTaskSessionsByTaskID(ctx, "task-batch-3")
	if err != nil {
		t.Fatalf("ListActiveTaskSessionsByTaskID(terminal-only): %v", err)
	}
	if len(none) != 0 {
		t.Errorf("terminal-only task returned %+v, want empty", none)
	}
}

func TestTaskEnvironmentSharingGuardsSeeOnlyActiveForeignSessions(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()
	repo := fixture.repo

	hasActive, err := repo.HasActiveTaskSessionsByEnvironment(ctx, "env-batch-shared")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByEnvironment: %v", err)
	}
	if !hasActive {
		t.Error("HasActiveTaskSessionsByEnvironment = false, want true for the RUNNING session")
	}
	hasActive, err = repo.HasActiveTaskSessionsByEnvironment(ctx, "env-batch-solo")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByEnvironment(solo): %v", err)
	}
	if hasActive {
		t.Error("HasActiveTaskSessionsByEnvironment(solo) = true, want false — its only session is COMPLETED")
	}

	// task-batch-2's active session borrows env-batch-shared, so the guard
	// fires for task-batch-1 but not for task-batch-2 itself.
	borrowed, err := repo.HasActiveTaskSessionsByTaskEnvironmentExcludingTask(ctx, "env-batch-shared", "task-batch-1")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByTaskEnvironmentExcludingTask: %v", err)
	}
	if !borrowed {
		t.Error("guard = false, want true — task-batch-2 holds an active session on the shared environment")
	}
	borrowerID, err := repo.FindActiveTaskSessionTaskIDByTaskEnvironmentExcludingTask(ctx, "env-batch-shared", "task-batch-1")
	if err != nil {
		t.Fatalf("FindActiveTaskSessionTaskIDByTaskEnvironmentExcludingTask: %v", err)
	}
	if borrowerID != "task-batch-2" {
		t.Errorf("borrower = %q, want task-batch-2", borrowerID)
	}

	borrowed, err = repo.HasActiveTaskSessionsByTaskEnvironmentExcludingTask(ctx, "env-batch-shared", "task-batch-2")
	if err != nil {
		t.Fatalf("guard excluding task-batch-2: %v", err)
	}
	if !borrowed {
		t.Error("guard = false, want true — task-batch-1 also holds an active session there")
	}

	// No active foreign session at all: both guards report the empty answer
	// rather than an error.
	borrowed, err = repo.HasActiveTaskSessionsByTaskEnvironmentExcludingTask(ctx, "env-batch-solo", "task-batch-3")
	if err != nil {
		t.Fatalf("guard on the solo environment: %v", err)
	}
	if borrowed {
		t.Error("guard = true on the solo environment, want false")
	}
	borrowerID, err = repo.FindActiveTaskSessionTaskIDByTaskEnvironmentExcludingTask(ctx, "env-batch-solo", "task-batch-3")
	if err != nil {
		t.Fatalf("borrower lookup on the solo environment: %v", err)
	}
	if borrowerID != "" {
		t.Errorf("borrower = %q, want an empty string when nothing matches", borrowerID)
	}
}

func TestGetTaskSessionByTaskIDReturnsNewestWhileActiveVariantFiltersState(t *testing.T) {
	fixture := seedSessionBatchFixture(t)
	ctx := context.Background()
	repo := fixture.repo

	// GetTaskSessionByTaskID is newest-by-started_at, not primary-first, so
	// pin the two task-batch-1 sessions with the terminal one newest.
	stamps := map[string]time.Time{
		"session-b1-primary":   time.Date(2026, 7, 1, 8, 0, 0, 0, time.UTC),
		"session-b1-secondary": time.Date(2026, 7, 1, 9, 0, 0, 0, time.UTC),
	}
	for sessionID, startedAt := range stamps {
		if _, err := repo.db.ExecContext(ctx, repo.db.Rebind(
			`UPDATE task_sessions SET started_at = ? WHERE id = ?`), startedAt, sessionID); err != nil {
			t.Fatalf("pin started_at for %s: %v", sessionID, err)
		}
	}

	got, err := repo.GetTaskSessionByTaskID(ctx, "task-batch-1")
	if err != nil {
		t.Fatalf("GetTaskSessionByTaskID: %v", err)
	}
	if got.ID != "session-b1-secondary" {
		t.Errorf("GetTaskSessionByTaskID = %q, want session-b1-secondary (newest by started_at)", got.ID)
	}

	// The active-only variant skips the newer COMPLETED row.
	active, err := repo.GetActiveTaskSessionByTaskID(ctx, "task-batch-1")
	if err != nil {
		t.Fatalf("GetActiveTaskSessionByTaskID: %v", err)
	}
	if active.ID != "session-b1-primary" {
		t.Errorf("GetActiveTaskSessionByTaskID = %q, want session-b1-primary (the only active row)", active.ID)
	}
	if active.State != models.TaskSessionStateRunning {
		t.Errorf("State = %q, want RUNNING", active.State)
	}

	// task-batch-3's only session is COMPLETED: found by the plain read,
	// reported as not-found by the active-only read.
	got, err = repo.GetTaskSessionByTaskID(ctx, "task-batch-3")
	if err != nil {
		t.Fatalf("GetTaskSessionByTaskID(task-batch-3): %v", err)
	}
	if got.ID != "session-b3-secondary" {
		t.Errorf("GetTaskSessionByTaskID = %q, want session-b3-secondary", got.ID)
	}
	_, err = repo.GetActiveTaskSessionByTaskID(ctx, "task-batch-3")
	if !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Errorf("GetActiveTaskSessionByTaskID(terminal-only) error = %v, want ErrTaskSessionNotFound", err)
	}
	_, err = repo.GetTaskSessionByTaskID(ctx, "task-batch-missing")
	if !errors.Is(err, models.ErrTaskSessionNotFound) {
		t.Errorf("GetTaskSessionByTaskID(missing) error = %v, want ErrTaskSessionNotFound", err)
	}
}

// TestHasActiveTaskSessionsByEnvironmentKeysOffEnvironmentIDNotTaskEnvironmentID
// discriminates the two columns. The shared batch fixture sets environment_id
// and task_environment_id to the same value, so on its own it cannot tell
// which column the guard reads — swapping them there would go unnoticed.
func TestHasActiveTaskSessionsByEnvironmentKeysOffEnvironmentIDNotTaskEnvironmentID(t *testing.T) {
	repo := newRepoForSessionTests(t)
	ctx := context.Background()

	if err := repo.CreateTask(ctx, &models.Task{ID: "task-envcol", Title: "env column"}); err != nil {
		t.Fatalf("CreateTask: %v", err)
	}
	if err := repo.CreateTaskEnvironment(ctx, &models.TaskEnvironment{
		ID: "env-col-taskenv", TaskID: "task-envcol",
		ExecutorType: string(models.ExecutorTypeLocal), Status: models.TaskEnvironmentStatusReady,
	}); err != nil {
		t.Fatalf("CreateTaskEnvironment: %v", err)
	}
	// Active, but only task_environment_id points at env-col-taskenv;
	// environment_id is left empty.
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID: "session-envcol-taskenv-only", TaskID: "task-envcol",
		State: models.TaskSessionStateRunning, TaskEnvironmentID: "env-col-taskenv",
	}); err != nil {
		t.Fatalf("CreateTaskSession(task-env only): %v", err)
	}

	hasActive, err := repo.HasActiveTaskSessionsByEnvironment(ctx, "env-col-taskenv")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByEnvironment: %v", err)
	}
	if hasActive {
		t.Error("HasActiveTaskSessionsByEnvironment = true for a value only present in task_environment_id; " +
			"the guard must key off environment_id")
	}
	// The task-environment guard does see the same row, which is what makes
	// the negative above meaningful rather than an empty-table artifact.
	borrowed, err := repo.HasActiveTaskSessionsByTaskEnvironmentExcludingTask(ctx, "env-col-taskenv", "task-other")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByTaskEnvironmentExcludingTask: %v", err)
	}
	if !borrowed {
		t.Fatal("task-environment guard = false; the fixture row is not visible, so the negative above proves nothing")
	}

	// Now a session whose environment_id carries the value: the guard fires.
	if err := repo.CreateTaskSession(ctx, &models.TaskSession{
		ID: "session-envcol-envid", TaskID: "task-envcol",
		State: models.TaskSessionStateRunning, EnvironmentID: "env-col-envid",
	}); err != nil {
		t.Fatalf("CreateTaskSession(environment_id): %v", err)
	}
	hasActive, err = repo.HasActiveTaskSessionsByEnvironment(ctx, "env-col-envid")
	if err != nil {
		t.Fatalf("HasActiveTaskSessionsByEnvironment(environment_id): %v", err)
	}
	if !hasActive {
		t.Error("HasActiveTaskSessionsByEnvironment = false for a value present in environment_id, want true")
	}
}