Files
kandev/apps/backend/internal/task/repository/sqlite/session.go
T
2026-09-21 14:07:49 +01:00

4451 lines
165 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Package sqlite provides SQLite-based repository implementations.
package sqlite
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/google/uuid"
"github.com/jmoiron/sqlx"
agentdto "github.com/kandev/kandev/internal/agent/dto"
"github.com/kandev/kandev/internal/agentctl/tracing"
"github.com/kandev/kandev/internal/db"
"github.com/kandev/kandev/internal/db/dialect"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/task/models"
"github.com/kandev/kandev/internal/task/recoveryclaim"
)
type taskSessionExecutor interface {
ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
}
// Turn operations
// CreateTurn creates a new turn
func (r *Repository) CreateTurn(ctx context.Context, turn *models.Turn) error {
stampTurnDefaults(turn)
return r.insertTurnWithSessionLock(ctx, turn)
}
// CreateTurnWithStepStamp is documented on the TurnRepository interface. It
// reads the task's current step inside a transaction that takes the same
// readTaskStepInTx lock a step move takes, so the read and the turn insert
// are serialized against concurrent movers of the same task row rather than
// racing a plain unlocked GetTask against a later, separate insert. A
// failure to open a transaction or read the step degrades to a plain,
// unstamped insert — see the spec's failure-modes table: turn creation must
// never fail because telemetry could not be resolved.
func (r *Repository) CreateTurnWithStepStamp(ctx context.Context, turn *models.Turn) (bool, error) {
stampTurnDefaults(turn)
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return false, r.insertTurnWithSessionLock(ctx, turn)
}
committed := false
defer func() {
if !committed {
_ = tx.Rollback()
}
}()
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), turn.TaskSessionID); err != nil {
return false, err
}
_, stepID, found, stepErr := r.readTaskStepInTx(ctx, tx, turn.TaskID)
if stepErr != nil {
_ = tx.Rollback()
committed = true
return false, r.insertTurnWithSessionLock(ctx, turn)
}
stamped := false
if found && stepID != "" {
if turn.Metadata == nil {
turn.Metadata = map[string]interface{}{}
}
turn.Metadata[models.TurnMetaKeyWorkflowStepIDAtStart] = stepID
stamped = true
}
if err := r.insertTurnRow(ctx, tx, turn); err != nil {
return false, err
}
if err := tx.Commit(); err != nil {
return false, err
}
committed = true
return stamped, nil
}
// stampTurnDefaults fills in the ID/timestamp defaults CreateTurn and
// CreateTurnWithStepStamp both need before inserting.
func stampTurnDefaults(turn *models.Turn) {
if turn.ID == "" {
turn.ID = uuid.New().String()
}
now := time.Now().UTC()
if turn.StartedAt.IsZero() {
turn.StartedAt = now
}
if turn.CreatedAt.IsZero() {
turn.CreatedAt = now
}
turn.UpdatedAt = now
}
// insertTurnRow inserts turn's row via execer, which is either r.db (a plain,
// non-transactional insert) or a *sql.Tx (participating in the caller's
// transaction). A new turn is session activity, so the same statement batch
// also refreshes task_sessions.updated_at: the session reconciliation sweep
// measures event silence from that row (and the newest message), and a turn
// dispatched without a user message (auto-start, workflow on_enter) must
// reset the clock even though no message row exists yet.
func (r *Repository) insertTurnRow(ctx context.Context, execer taskSessionExecutor, turn *models.Turn) error {
metadataJSON := "{}"
if turn.Metadata != nil {
metadataBytes, err := json.Marshal(turn.Metadata)
if err != nil {
return fmt.Errorf("failed to serialize turn metadata: %w", err)
}
metadataJSON = string(metadataBytes)
}
if _, err := execer.ExecContext(ctx, r.db.Rebind(`
INSERT INTO task_session_turns (id, task_session_id, task_id, execution_profile_id, route_generation, started_at, completed_at, metadata, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`), turn.ID, turn.TaskSessionID, turn.TaskID, turn.ExecutionProfileID, turn.RouteGeneration, turn.StartedAt, turn.CompletedAt, metadataJSON, turn.CreatedAt, turn.UpdatedAt); err != nil {
return err
}
_, err := execer.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET updated_at = ? WHERE id = ?
`), turn.UpdatedAt, turn.TaskSessionID)
return err
}
// insertTurnWithSessionLock serializes successor-turn creation with every
// current-turn clarification decision on every database. The transaction is
// also required for SQLite: inserting a turn and refreshing its parent
// session's activity clock must commit or roll back as one unit.
func (r *Repository) insertTurnWithSessionLock(ctx context.Context, turn *models.Turn) error {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return fmt.Errorf("begin turn creation: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), turn.TaskSessionID); err != nil {
return err
}
if err := r.insertTurnRow(ctx, tx, turn); err != nil {
return err
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("commit turn creation: %w", err)
}
return nil
}
// DeleteTurnIfUnreferenced removes a rejected pre-dispatch turn only while it
// has no messages. The message guard preserves an ambiguously accepted prompt.
func (r *Repository) DeleteTurnIfUnreferenced(
ctx context.Context,
sessionID, turnID string,
) (bool, error) {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return false, fmt.Errorf("begin turn rollback: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), sessionID); err != nil {
return false, err
}
result, err := tx.ExecContext(ctx, r.db.Rebind(`
DELETE FROM task_session_turns
WHERE id = ?
AND task_session_id = ?
AND NOT EXISTS (
SELECT 1
FROM task_session_messages
WHERE turn_id = task_session_turns.id
)
`), turnID, sessionID)
if err != nil {
return false, fmt.Errorf("delete unreferenced turn: %w", err)
}
deleted, err := result.RowsAffected()
if err != nil {
return false, fmt.Errorf("inspect unreferenced turn deletion: %w", err)
}
if err := tx.Commit(); err != nil {
return false, fmt.Errorf("commit turn rollback: %w", err)
}
return deleted == 1, nil
}
type turnScanner interface {
Scan(dest ...interface{}) error
}
func scanTurn(scanner turnScanner) (*models.Turn, error) {
turn := &models.Turn{}
var metadataJSON string
var completedAt sql.NullTime
err := scanner.Scan(&turn.ID, &turn.TaskSessionID, &turn.TaskID, &turn.ExecutionProfileID, &turn.RouteGeneration, &turn.StartedAt, &completedAt, &metadataJSON, &turn.CreatedAt, &turn.UpdatedAt)
if err != nil {
return nil, err
}
if completedAt.Valid {
turn.CompletedAt = &completedAt.Time
}
if metadataJSON != "" && metadataJSON != "{}" {
if err := json.Unmarshal([]byte(metadataJSON), &turn.Metadata); err != nil {
return nil, fmt.Errorf("failed to deserialize turn metadata: %w", err)
}
}
return turn, nil
}
func scanTurnRow(row *sql.Row) (*models.Turn, error) {
return scanTurn(row)
}
// GetTurn retrieves a turn by ID
func (r *Repository) GetTurn(ctx context.Context, id string) (*models.Turn, error) {
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT id, task_session_id, task_id, execution_profile_id, route_generation, started_at, completed_at, metadata, created_at, updated_at
FROM task_session_turns WHERE id = ?
`), id)
return scanTurnRow(row)
}
// GetActiveTurnBySessionID gets the currently active (non-completed) turn for
// a session. This is not the same query as "current turn": it keeps its own
// completed_at IS NULL filter alongside currentTurnAuthority's predicate and
// ordering (D10) so AbandonOpenTurns' re-bury loop keeps working.
func (r *Repository) GetActiveTurnBySessionID(ctx context.Context, sessionID string) (*models.Turn, error) {
predicate, orderBy := currentTurnAuthority(r.ro.DriverName(), "turn_row")
query := fmt.Sprintf(`
SELECT id, task_session_id, task_id, execution_profile_id, route_generation, started_at, completed_at, metadata, created_at, updated_at
FROM task_session_turns turn_row
WHERE turn_row.task_session_id = ?
AND turn_row.completed_at IS NULL
AND %s
ORDER BY %s
LIMIT 1
`, predicate, orderBy)
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), sessionID)
return scanTurnRow(row)
}
// UpdateTurn updates an existing turn
func (r *Repository) UpdateTurn(ctx context.Context, turn *models.Turn) error {
metadataJSON, err := serializeTurnMetadata(turn.Metadata)
if err != nil {
return err
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return fmt.Errorf("begin turn update: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), turn.TaskSessionID); err != nil {
return err
}
updatedAt := r.nowUTC()
result, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_session_turns
SET completed_at = ?, metadata = ?, execution_profile_id = ?, route_generation = ?, updated_at = ?
WHERE id = ? AND task_session_id = ? AND updated_at = ?
`), turn.CompletedAt, metadataJSON, turn.ExecutionProfileID, turn.RouteGeneration, updatedAt, turn.ID, turn.TaskSessionID, turn.UpdatedAt)
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("inspect turn update: %w", err)
}
if affected != 1 {
return fmt.Errorf("update turn %s: stale metadata snapshot", turn.ID)
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("commit turn update: %w", err)
}
turn.UpdatedAt = updatedAt
return nil
}
// UpdateActiveTurnMetadata applies a narrow metadata patch without copying a
// caller's stale snapshot over unrelated fields or a concurrently completed turn.
func (r *Repository) UpdateActiveTurnMetadata(
ctx context.Context,
sessionID, turnID string,
updates map[string]interface{},
removeKeys []string,
) (bool, map[string]interface{}, time.Time, error) {
return r.patchTurnMetadata(ctx, sessionID, turnID, updates, removeKeys, true)
}
// PatchTurnMetadata merges metadata into an active or completed turn.
func (r *Repository) PatchTurnMetadata(
ctx context.Context,
sessionID, turnID string,
updates map[string]interface{},
) (bool, time.Time, error) {
updated, _, updatedAt, err := r.patchTurnMetadata(ctx, sessionID, turnID, updates, nil, false)
return updated, updatedAt, err
}
// ClearTurnPromptDispatchMetadata removes durable recovery state only after
// turn.started publication succeeds. It intentionally accepts a completed
// turn because provider output can settle a fast turn while publication is in
// progress; clearing metadata must never reopen or otherwise alter completion.
func (r *Repository) ClearTurnPromptDispatchMetadata(
ctx context.Context,
sessionID, turnID string,
) (bool, map[string]interface{}, time.Time, error) {
return r.patchTurnMetadata(ctx, sessionID, turnID, nil, []string{
models.TurnMetaKeyPromptDispatchPending,
models.TurnMetaKeyPromptDispatchAttempted,
models.TurnMetaKeyPromptDispatchClarificationPendingID,
models.TurnMetaKeyPromptDispatchClarificationTurnID,
models.TurnMetaKeyPromptDispatchClarificationMessageIDs,
models.TurnMetaKeyPromptDispatchStartEventPending,
}, false)
}
func (r *Repository) patchTurnMetadata(
ctx context.Context,
sessionID, turnID string,
updates map[string]interface{},
removeKeys []string,
activeOnly bool,
) (bool, map[string]interface{}, time.Time, error) {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return false, nil, time.Time{}, fmt.Errorf("begin active turn metadata update: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), sessionID); err != nil {
return false, nil, time.Time{}, err
}
activeClause := ""
if activeOnly {
activeClause = " AND completed_at IS NULL"
}
selectQuery := fmt.Sprintf(`
SELECT metadata
FROM task_session_turns
WHERE id = ? AND task_session_id = ?%s
`, activeClause)
var metadataJSON string
err = tx.QueryRowContext(ctx, r.db.Rebind(selectQuery), turnID, sessionID).Scan(&metadataJSON)
if errors.Is(err, sql.ErrNoRows) {
return false, nil, time.Time{}, nil
}
if err != nil {
return false, nil, time.Time{}, fmt.Errorf("read active turn metadata: %w", err)
}
metadata, metadataJSON, err := applyTurnMetadataPatch(metadataJSON, updates, removeKeys)
if err != nil {
return false, nil, time.Time{}, err
}
updatedAt := r.nowUTC()
updateQuery := fmt.Sprintf(`
UPDATE task_session_turns
SET metadata = ?, updated_at = ?
WHERE id = ? AND task_session_id = ?%s
`, activeClause)
result, err := tx.ExecContext(ctx, r.db.Rebind(updateQuery), metadataJSON, updatedAt, turnID, sessionID)
if err != nil {
return false, nil, time.Time{}, err
}
affected, err := result.RowsAffected()
if err != nil {
return false, nil, time.Time{}, fmt.Errorf("inspect active turn metadata update: %w", err)
}
if err := tx.Commit(); err != nil {
return false, nil, time.Time{}, fmt.Errorf("commit active turn metadata update: %w", err)
}
if affected != 1 {
return false, nil, time.Time{}, nil
}
return true, metadata, updatedAt, nil
}
func applyTurnMetadataPatch(
metadataJSON string,
updates map[string]interface{},
removeKeys []string,
) (map[string]interface{}, string, error) {
metadata := make(map[string]interface{})
if metadataJSON != "" && metadataJSON != "{}" {
if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil {
return nil, "", fmt.Errorf("deserialize turn metadata patch: %w", err)
}
}
for key, value := range updates {
metadata[key] = value
}
for _, key := range removeKeys {
delete(metadata, key)
}
serialized, err := serializeTurnMetadata(metadata)
return metadata, serialized, err
}
func serializeTurnMetadata(metadata map[string]interface{}) (string, error) {
if metadata == nil {
return "{}", nil
}
metadataBytes, err := json.Marshal(metadata)
if err != nil {
return "", fmt.Errorf("failed to serialize turn metadata: %w", err)
}
return string(metadataBytes), nil
}
// CompleteTurn marks a turn as completed with the current time
func (r *Repository) CompleteTurn(ctx context.Context, id string) error {
now := time.Now().UTC()
_, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_session_turns
SET completed_at = ?, updated_at = ?
WHERE id = ?
`), now, now, id)
return err
}
// AbandonTurn marks a turn as completed with completed_at = started_at, giving it
// zero duration. Used when a turn was orphaned by an interruption (backend
// restart, agent crash) and the previous "running" window was not real work —
// recording `now` would inflate analytics and the UI's last-turn duration with
// hours of dead time.
func (r *Repository) AbandonTurn(ctx context.Context, id string) error {
now := time.Now().UTC()
_, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_session_turns
SET completed_at = started_at, updated_at = ?
WHERE id = ? AND completed_at IS NULL
`), now, id)
return err
}
// ListTurnsBySession returns all turns for a session ordered by start time
func (r *Repository) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) {
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListTurnsBySession")
defer span.End()
query := fmt.Sprintf(`
SELECT id, task_session_id, task_id, execution_profile_id, route_generation, started_at, completed_at, metadata, created_at, updated_at
FROM task_session_turns turn_row
WHERE turn_row.task_session_id = ? AND %s
ORDER BY turn_row.started_at ASC, turn_row.created_at ASC, turn_row.id ASC
`, turnHistoryPredicate(r.ro.DriverName(), "turn_row"))
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), sessionID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var result []*models.Turn
for rows.Next() {
turn, err := scanTurn(rows)
if err != nil {
return nil, err
}
result = append(result, turn)
}
if err := rows.Err(); err != nil {
return nil, err
}
return result, nil
}
// taskSessionSelectCols is the column list used by every SELECT that scans
// into a *models.TaskSession via scanTaskSession / scanTaskSessionRow. Centralised
// so columns can be added in one place. Order MUST match the scan helpers.
//
// agent_execution_id and container_id come from executors_running (the single
// source of truth for active execution state) via LEFT JOIN. The effective
// workspace_path comes from the linked task environment when available so a
// promoted multi-repo task is correct after a session refresh; all other
// columns come from task_sessions (aliased ts).
//
// ADR 0005: agent_profile_id is the single column for both kanban (FK to a
// shallow profile) and office (FK to a per-workspace rich profile) sessions —
// the two column names that used to live here have collapsed into one.
const taskSessionSelectCols = `ts.id, ts.task_id, ts.queue_incarnation_id,
COALESCE(er.agent_execution_id, ''), COALESCE(er.container_id, ''),
ts.agent_profile_id, ts.execution_profile_id, ts.route_generation, ts.route_state, ts.route_reason, ts.downstream_acp_session_id,
ts.executor_id, ts.executor_profile_id, ts.environment_id,
ts.repository_id, ts.base_branch, ts.base_commit_sha,
COALESCE(NULLIF(te.workspace_path, ''), ts.workspace_path),
ts.agent_profile_snapshot, ts.executor_snapshot, ts.environment_snapshot, ts.repository_snapshot,
ts.state, ts.error_message, ts.metadata, ts.started_at, ts.completed_at, ts.updated_at,
ts.is_primary, ts.review_status, ts.is_passthrough, ts.task_environment_id, ts.name, ts.last_read_message_id,
ts.cost_subcents, ts.tokens_in, ts.tokens_cached_in, ts.tokens_out`
// taskSessionFromClause is the FROM clause that pairs with taskSessionSelectCols.
// Always reference task_sessions as `ts` and executors_running as `er` in WHERE/ORDER.
const taskSessionFromClause = `FROM task_sessions ts
LEFT JOIN executors_running er ON er.session_id = ts.id
LEFT JOIN task_environments te ON te.id = ts.task_environment_id`
// Task Session operations
// CreateTaskSession creates a new agent session
func (r *Repository) CreateTaskSession(ctx context.Context, session *models.TaskSession) error {
// The task cleanup barrier serializes session creation against archive/
// delete preparation (ADR-2026-08-08): PostgreSQL takes a row lock on the
// task, SQLite serializes the writer transaction. A creation admitted
// after cleanup inventory was captured would be missed by the snapshot.
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.taskCleanupBarrierLocked(ctx, tx, session.TaskID); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
return tx.Commit()
}
// CreateTaskSessionWithInitialRuntimeSeed claims the task's launch-only
// runtime seed while creating the session. The task lock / SQLite writer
// transaction makes the session count, seed read, session insert, and seed
// removal one serialized operation, so a concurrent launch or replacement
// session cannot inherit the seed a second time.
func (r *Repository) CreateTaskSessionWithInitialRuntimeSeed(ctx context.Context, session *models.TaskSession) error {
return r.createTaskSessionWithInitialRuntimeSeed(ctx, session, nil)
}
// CreateTaskSessionWithInitialRuntimeSeedAndWorkflowRoute atomically claims
// the launch-only runtime seed, inserts the session, and records a prepared
// workflow route. The route destination is therefore durable before any
// external lifecycle work can promote or launch the session.
func (r *Repository) CreateTaskSessionWithInitialRuntimeSeedAndWorkflowRoute(
ctx context.Context,
session *models.TaskSession,
route *models.WorkflowSessionRoute,
) error {
return r.createTaskSessionWithInitialRuntimeSeed(ctx, session, route)
}
func (r *Repository) createTaskSessionWithInitialRuntimeSeed(
ctx context.Context,
session *models.TaskSession,
route *models.WorkflowSessionRoute,
) error {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.taskCleanupBarrierLocked(ctx, tx, session.TaskID); err != nil {
return err
}
if err := r.applyInitialRuntimeSeedTx(ctx, tx, session); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.persistWorkflowSessionRouteTx(ctx, tx, session.TaskID, session.ID, route); err != nil {
return err
}
return tx.Commit()
}
// CreateTaskSessionWithWorkflowSessionRoute inserts a session and its
// prepared route in one transaction for repositories that do not need a
// launch-only seed or workspace election.
func (r *Repository) CreateTaskSessionWithWorkflowSessionRoute(
ctx context.Context,
session *models.TaskSession,
route *models.WorkflowSessionRoute,
) error {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.taskCleanupBarrierLocked(ctx, tx, session.TaskID); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.persistWorkflowSessionRouteTx(ctx, tx, session.TaskID, session.ID, route); err != nil {
return err
}
return tx.Commit()
}
// CreateTaskSessionWithWorkspaceBinding atomically elects the first session
// allowed to materialize a task workspace, or attaches a later session to its
// ready canonical environment. The transaction deliberately happens before a
// session row is committed: a preparing or unsafe workspace therefore leaves
// no orphan session for a caller to clean up.
//
// The candidate is persisted only when this call wins the election. It may be
// a worktree environment without a workspace path because that path is not
// known until the elected launch has completed preparation.
//
//nolint:cyclop // The state cases are the durable workspace binding state machine.
func (r *Repository) CreateTaskSessionWithWorkspaceBinding(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
) error {
return r.createTaskSessionWithWorkspaceBinding(ctx, session, candidate, nil)
}
// CreateTaskSessionWithWorkspaceBindingAndWorkflowRoute is the route-aware
// variant used by explicit workflow starts. It keeps workspace election,
// session insertion, and prepared-route recording in one transaction.
func (r *Repository) CreateTaskSessionWithWorkspaceBindingAndWorkflowRoute(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
route *models.WorkflowSessionRoute,
) error {
return r.createTaskSessionWithWorkspaceBinding(ctx, session, candidate, route)
}
//nolint:cyclop // The state cases are the durable workspace binding state machine.
func (r *Repository) createTaskSessionWithWorkspaceBinding(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
route *models.WorkflowSessionRoute,
) error {
if candidate == nil {
return fmt.Errorf("workspace binding candidate is required")
}
if candidate.TaskID != session.TaskID {
return fmt.Errorf("workspace binding task mismatch")
}
if session.ID == "" {
session.ID = uuid.New().String()
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.taskCleanupBarrierLocked(ctx, tx, session.TaskID); err != nil {
return err
}
var envID, status, materializationSessionID string
err = tx.QueryRowContext(ctx, r.db.Rebind(`
SELECT id, status, materialization_session_id FROM task_environments WHERE task_id = ?
`), session.TaskID).Scan(&envID, &status, &materializationSessionID)
switch {
case errors.Is(err, sql.ErrNoRows):
if candidate.ID == "" {
candidate.ID = uuid.New().String()
}
candidate.Status = models.TaskEnvironmentStatusCreating
candidate.MaterializationSessionID = session.ID
candidate.CreatedAt = r.nowUTC()
candidate.UpdatedAt = candidate.CreatedAt
if err := r.insertCreatingWorkspaceEnvironment(ctx, tx, candidate); err != nil {
return err
}
session.TaskEnvironmentID = candidate.ID
case err != nil:
return fmt.Errorf("load workspace binding: %w", err)
case models.TaskEnvironmentStatus(status) == models.TaskEnvironmentStatusCreating:
abandoned, err := r.failAbandonedWorkspaceMaterialization(ctx, tx, envID, materializationSessionID)
if err != nil {
return err
}
if abandoned {
if err := tx.Commit(); err != nil {
return fmt.Errorf("persist abandoned workspace materialization failure: %w", err)
}
return fmt.Errorf("%w: the initial workspace materialization did not complete", models.ErrWorkspaceReuseUnsafe)
}
return fmt.Errorf("%w: retry after the initial workspace launch completes", models.ErrWorkspacePreparing)
case models.TaskEnvironmentStatus(status) == models.TaskEnvironmentStatusReady || models.TaskEnvironmentStatus(status) == models.TaskEnvironmentStatusStopped:
session.TaskEnvironmentID = envID
default:
return fmt.Errorf("%w: existing task environment is not attachable", models.ErrWorkspaceReuseUnsafe)
}
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, session.TaskEnvironmentID); err != nil {
return err
}
if err := r.applyInitialRuntimeSeedTx(ctx, tx, session); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.persistWorkflowSessionRouteTx(ctx, tx, session.TaskID, session.ID, route); err != nil {
return err
}
return tx.Commit()
}
// CreateTaskSessionWithSharedGroupWorkspaceBinding atomically elects the
// shared group's first materializer. The elected environment is published to
// the group while it is still creating, so another member can never create a
// second physical workspace: it receives workspace_preparing until the elected
// owner finalizes the canonical environment.
func (r *Repository) CreateTaskSessionWithSharedGroupWorkspaceBinding(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
groupID string,
) error {
return r.createTaskSessionWithSharedGroupWorkspaceBinding(ctx, session, candidate, groupID, nil)
}
// CreateTaskSessionWithSharedGroupWorkspaceBindingAndWorkflowRoute is the
// route-aware shared-workspace variant used by explicit workflow starts.
func (r *Repository) CreateTaskSessionWithSharedGroupWorkspaceBindingAndWorkflowRoute(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
groupID string,
route *models.WorkflowSessionRoute,
) error {
return r.createTaskSessionWithSharedGroupWorkspaceBinding(ctx, session, candidate, groupID, route)
}
func (r *Repository) createTaskSessionWithSharedGroupWorkspaceBinding(
ctx context.Context,
session *models.TaskSession,
candidate *models.TaskEnvironment,
groupID string,
route *models.WorkflowSessionRoute,
) error {
if candidate == nil || groupID == "" {
return fmt.Errorf("shared workspace binding requires a candidate and group")
}
if candidate.TaskID != session.TaskID {
return fmt.Errorf("shared workspace binding task mismatch")
}
if session.ID == "" {
session.ID = uuid.New().String()
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.taskCleanupBarrierLocked(ctx, tx, session.TaskID); err != nil {
return err
}
if candidate.ID == "" {
candidate.ID = uuid.New().String()
}
candidate.Status = models.TaskEnvironmentStatusCreating
candidate.MaterializationSessionID = session.ID
candidate.CreatedAt = r.nowUTC()
candidate.UpdatedAt = candidate.CreatedAt
res, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_workspace_groups
SET materialized_environment_id = ?, updated_at = ?
WHERE id = ?
AND materialized_environment_id = ''
AND EXISTS (
SELECT 1 FROM task_workspace_group_members
WHERE workspace_group_id = ? AND task_id = ? AND released_at IS NULL
)
`), candidate.ID, candidate.UpdatedAt, groupID, groupID, session.TaskID)
if err != nil {
return fmt.Errorf("elect shared workspace materializer: %w", err)
}
won, err := res.RowsAffected()
if err != nil {
return fmt.Errorf("inspect shared workspace election: %w", err)
}
if won == 0 {
return r.bindReadySharedGroupEnvironment(ctx, tx, session, groupID, route)
}
if err := r.insertCreatingWorkspaceEnvironment(ctx, tx, candidate); err != nil {
return err
}
session.TaskEnvironmentID = candidate.ID
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, session.TaskEnvironmentID); err != nil {
return err
}
if err := r.applyInitialRuntimeSeedTx(ctx, tx, session); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.persistWorkflowSessionRouteTx(ctx, tx, session.TaskID, session.ID, route); err != nil {
return err
}
return tx.Commit()
}
func (r *Repository) bindReadySharedGroupEnvironment(
ctx context.Context,
tx *sqlx.Tx,
session *models.TaskSession,
groupID string,
route *models.WorkflowSessionRoute,
) error {
var environmentID, status string
err := tx.QueryRowContext(ctx, r.db.Rebind(`
SELECT COALESCE(e.id, ''), COALESCE(e.status, '')
FROM task_workspace_groups g
JOIN task_workspace_group_members m
ON m.workspace_group_id = g.id AND m.task_id = ? AND m.released_at IS NULL
LEFT JOIN task_environments e ON e.id = g.materialized_environment_id
WHERE g.id = ?
`), session.TaskID, groupID).Scan(&environmentID, &status)
if errors.Is(err, sql.ErrNoRows) || environmentID == "" {
return fmt.Errorf("%w: shared workspace group is unavailable", models.ErrWorkspaceReuseUnsafe)
}
if err != nil {
return fmt.Errorf("load shared workspace environment: %w", err)
}
switch models.TaskEnvironmentStatus(status) {
case models.TaskEnvironmentStatusReady, models.TaskEnvironmentStatusStopped:
session.TaskEnvironmentID = environmentID
case models.TaskEnvironmentStatusCreating:
return fmt.Errorf("%w: retry after the shared workspace launch completes", models.ErrWorkspacePreparing)
default:
return fmt.Errorf("%w: shared workspace is not attachable", models.ErrWorkspaceReuseUnsafe)
}
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, session.TaskEnvironmentID); err != nil {
return err
}
if err := r.applyInitialRuntimeSeedTx(ctx, tx, session); err != nil {
return err
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.persistWorkflowSessionRouteTx(ctx, tx, session.TaskID, session.ID, route); err != nil {
return err
}
return tx.Commit()
}
func (r *Repository) insertCreatingWorkspaceEnvironment(ctx context.Context, tx *sqlx.Tx, candidate *models.TaskEnvironment) error {
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
INSERT INTO task_environments (
id, task_id, executor_type, executor_id, executor_profile_id,
control_port, status, materialization_session_id, workspace_path,
container_id, sandbox_id, task_dir_name, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`), candidate.ID, candidate.TaskID, candidate.ExecutorType, candidate.ExecutorID,
candidate.ExecutorProfileID, candidate.ControlPort, string(candidate.Status),
candidate.MaterializationSessionID, candidate.WorkspacePath, candidate.ContainerID,
candidate.SandboxID, candidate.TaskDirName, candidate.CreatedAt, candidate.UpdatedAt); err != nil {
return fmt.Errorf("create workspace binding: %w", err)
}
return nil
}
func (r *Repository) failAbandonedWorkspaceMaterialization(
ctx context.Context,
tx *sqlx.Tx,
environmentID string,
ownerID string,
) (bool, error) {
var ownerState models.TaskSessionState
ownerErr := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT state FROM task_sessions WHERE id = ?`), ownerID).Scan(&ownerState)
abandoned := ownerID == "" || errors.Is(ownerErr, sql.ErrNoRows) || isTerminalWorkspaceMaterializerState(ownerState)
if ownerErr != nil && !errors.Is(ownerErr, sql.ErrNoRows) {
return false, fmt.Errorf("load workspace materialization owner: %w", ownerErr)
}
if !abandoned {
return false, nil
}
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, environmentID); err != nil {
return false, err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_environments
SET status = ?, materialization_session_id = '', updated_at = ?
WHERE id = ? AND status = ?
`), string(models.TaskEnvironmentStatusFailed), r.nowUTC(), environmentID, string(models.TaskEnvironmentStatusCreating)); err != nil {
return false, fmt.Errorf("fail abandoned workspace materialization: %w", err)
}
return true, nil
}
func isTerminalWorkspaceMaterializerState(state models.TaskSessionState) bool {
switch state {
case models.TaskSessionStateCompleted, models.TaskSessionStateFailed, models.TaskSessionStateCancelled:
return true
default:
return false
}
}
// applyInitialRuntimeSeedTx preserves CreateTaskSessionWithInitialRuntimeSeed's
// one-time seed semantics for the workspace-binding creation path.
func (r *Repository) applyInitialRuntimeSeedTx(ctx context.Context, tx *sqlx.Tx, session *models.TaskSession) error {
initialRuntimeConfig, hasInitialRuntimeConfig, initialRuntimeConfigProfileID, hasInitialRuntimeSeedKey, err := r.loadInitialSessionRuntimeSeedTx(ctx, tx, session.TaskID)
if err != nil {
return err
}
var sessionCount int
if err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT COUNT(*) FROM task_sessions WHERE task_id = ?`), session.TaskID).Scan(&sessionCount); err != nil {
return fmt.Errorf("check task sessions before workspace binding: %w", err)
}
if sessionCount == 0 {
if session.Metadata == nil {
session.Metadata = make(map[string]interface{})
}
session.Metadata[models.SessionMetaKeyOrigin] = models.SessionOriginTaskInitial
if hasInitialRuntimeConfig && initialRuntimeConfigProfileID == session.AgentProfileID {
session.Metadata[models.SessionMetaKeyRuntimeConfigOverrides] = initialRuntimeConfig
}
} else if models.IsOriginalTaskSession(session.Metadata) {
delete(session.Metadata, models.SessionMetaKeyOrigin)
}
if hasInitialRuntimeSeedKey {
if _, err := r.removeTaskMetadataKeyWithExecutor(ctx, tx, session.TaskID, models.MetaKeyInitialSessionRuntimeConfig); err != nil {
return fmt.Errorf("consume initial runtime seed: %w", err)
}
if _, err := r.removeTaskMetadataKeyWithExecutor(ctx, tx, session.TaskID, models.MetaKeyInitialSessionRuntimeConfigProfileID); err != nil {
return fmt.Errorf("consume initial runtime seed profile: %w", err)
}
}
return nil
}
func (r *Repository) loadInitialSessionRuntimeSeedTx(
ctx context.Context,
tx *sqlx.Tx,
taskID string,
) (models.SessionRuntimeConfig, bool, string, bool, error) {
var metadataJSON sql.NullString
err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT metadata FROM tasks WHERE id = ?`,
), taskID).Scan(&metadataJSON)
if errors.Is(err, sql.ErrNoRows) {
return models.SessionRuntimeConfig{}, false, "", false, fmt.Errorf("task not found: %s", taskID)
}
if err != nil {
return models.SessionRuntimeConfig{}, false, "", false, fmt.Errorf("load task metadata for initial runtime seed: %w", err)
}
metadata := make(map[string]interface{})
raw := strings.TrimSpace(metadataJSON.String)
if metadataJSON.Valid && raw != "" && raw != "null" && raw != "{}" {
if err := json.Unmarshal([]byte(raw), &metadata); err != nil {
return models.SessionRuntimeConfig{}, false, "", false, fmt.Errorf("decode task metadata for initial runtime seed: %w", err)
}
}
seed, ok := models.LoadInitialSessionRuntimeConfig(metadata)
profileID := models.LoadInitialSessionRuntimeConfigProfileID(metadata)
_, hasSeedKey := metadata[models.MetaKeyInitialSessionRuntimeConfig]
return seed, ok, profileID, hasSeedKey, nil
}
// CreateOfficeTaskSession creates an Office session and atomically marks it as
// the task's initial session when no earlier session exists. The task row lock
// serializes callers across PostgreSQL connections; SQLite's single writer
// connection serializes the transaction. Within that same transaction it also
// enforces office-session uniqueness for the (task_id, agent_profile_id) pair:
// if a live row already exists for the pair, the insert is refused with
// ErrOfficeSessionRaceConflict rather than creating a second live session.
func (r *Repository) CreateOfficeTaskSession(ctx context.Context, session *models.TaskSession) error {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if dialect.IsPostgres(r.db.DriverName()) {
var lockedTaskID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT id FROM tasks WHERE id = ? FOR UPDATE`,
), session.TaskID).Scan(&lockedTaskID); err != nil {
return fmt.Errorf("lock task for office session: %w", err)
}
}
var sessionCount int
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT COUNT(*) FROM task_sessions WHERE task_id = ?`,
), session.TaskID).Scan(&sessionCount); err != nil {
return fmt.Errorf("check task sessions before office session: %w", err)
}
if sessionCount == 0 {
if session.Metadata == nil {
session.Metadata = make(map[string]interface{})
}
session.Metadata[models.SessionMetaKeyOrigin] = models.SessionOriginTaskInitial
}
if session.AgentProfileID != "" {
var liveCount int
if err := tx.QueryRowContext(ctx, r.db.Rebind(`
SELECT COUNT(*) FROM task_sessions
WHERE task_id = ? AND agent_profile_id = ?
AND state NOT IN ('COMPLETED', 'FAILED', 'CANCELLED')
`), session.TaskID, session.AgentProfileID).Scan(&liveCount); err != nil {
return fmt.Errorf("check live office session for pair: %w", err)
}
if liveCount > 0 {
return ErrOfficeSessionRaceConflict
}
}
if err := r.createTaskSession(ctx, tx, session); err != nil {
return err
}
return tx.Commit()
}
// createTaskSession inserts a task session through the supplied database handle.
func (r *Repository) createTaskSession(ctx context.Context, exec taskSessionExecutor, session *models.TaskSession) error {
if session.ID == "" {
session.ID = uuid.New().String()
}
if session.QueueIncarnationID == "" {
session.QueueIncarnationID = uuid.New().String()
}
now := time.Now().UTC()
// Only default StartedAt / UpdatedAt when the caller hasn't supplied
// one. The test harness backdates StartedAt so completed sessions
// have a non-zero duration (e.g. "Agent worked for 30s"); blowing
// those values away here defeats that.
if session.StartedAt.IsZero() {
session.StartedAt = now
}
if session.UpdatedAt.IsZero() {
session.UpdatedAt = now
}
if session.State == "" {
session.State = models.TaskSessionStateCreated
}
if tx, ok := exec.(*sqlx.Tx); ok {
if err := r.verifyTaskRunnerResolutionTx(ctx, tx, session); err != nil {
return err
}
if session.TaskEnvironmentID != "" {
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, session.TaskEnvironmentID); err != nil {
return err
}
}
if err := r.prepareWorkflowInitialSessionTx(ctx, tx, session); err != nil {
return err
}
}
metadataJSON, err := json.Marshal(session.Metadata)
if err != nil {
return fmt.Errorf("failed to serialize agent session metadata: %w", err)
}
agentProfileSnapshotJSON, err := json.Marshal(session.AgentProfileSnapshot)
if err != nil {
return fmt.Errorf("failed to serialize agent profile snapshot: %w", err)
}
executorSnapshotJSON, err := json.Marshal(session.ExecutorSnapshot)
if err != nil {
return fmt.Errorf("failed to serialize executor snapshot: %w", err)
}
environmentSnapshotJSON, err := json.Marshal(session.EnvironmentSnapshot)
if err != nil {
return fmt.Errorf("failed to serialize environment snapshot: %w", err)
}
repositorySnapshotJSON, err := json.Marshal(session.RepositorySnapshot)
if err != nil {
return fmt.Errorf("failed to serialize repository snapshot: %w", err)
}
// agent_profile_id is NULL-able and stored as NULL when empty. No unique
// index currently constrains (task_id, agent_profile_id) — see
// ErrOfficeSessionRaceConflict's doc comment (errors.go). Per ADR 0005,
// kanban and office share this column, and every row (kanban included)
// carries a non-NULL value in practice — nothing here scopes NULL to
// kanban specifically.
var agentProfileID interface{}
if session.AgentProfileID != "" {
agentProfileID = session.AgentProfileID
}
_, err = exec.ExecContext(ctx, r.db.Rebind(`
INSERT INTO task_sessions (
id, task_id, queue_incarnation_id, agent_profile_id, execution_profile_id, route_generation, route_state, route_reason, downstream_acp_session_id,
executor_id, executor_profile_id, environment_id,
repository_id, base_branch, base_commit_sha, workspace_path,
agent_profile_snapshot, executor_snapshot, environment_snapshot, repository_snapshot,
state, error_message, metadata, started_at, completed_at, updated_at,
is_primary, review_status, is_passthrough, task_environment_id, name
) VALUES (
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
?
)
`), session.ID, session.TaskID, session.QueueIncarnationID, agentProfileID,
session.ExecutionProfileID, session.RouteGeneration, session.RouteState, session.RouteReason, session.DownstreamACPSessionID,
session.ExecutorID, session.ExecutorProfileID, session.EnvironmentID, session.RepositoryID, session.BaseBranch, session.BaseCommitSHA, session.WorkspacePath,
string(agentProfileSnapshotJSON), string(executorSnapshotJSON), string(environmentSnapshotJSON), string(repositorySnapshotJSON),
string(session.State), session.ErrorMessage, string(metadataJSON),
session.StartedAt, session.CompletedAt, session.UpdatedAt,
dialect.BoolToInt(session.IsPrimary), session.ReviewStatus,
dialect.BoolToInt(session.IsPassthrough), session.TaskEnvironmentID, session.Name)
return err
}
// verifyTaskRunnerResolutionTx rejects a session built from a task snapshot
// that is older than the task-row lock it now holds. A runner switch and
// session creation therefore have one total order: the switch commits first
// and this check asks the caller to retry with the new runner, or this
// transaction inserts the session first and the switch is rejected by its
// mutability gate.
func (r *Repository) verifyTaskRunnerResolutionTx(ctx context.Context, tx *sqlx.Tx, session *models.TaskSession) error {
if session == nil || !session.TaskRunnerResolvedFromTask {
return nil
}
var metadataJSON string
err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT metadata FROM tasks WHERE id = ?`), session.TaskID).Scan(&metadataJSON)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("%w: task %s no longer exists", models.ErrTaskRunnerChanged, session.TaskID)
}
if err != nil {
return fmt.Errorf("verify task runner resolution: %w", err)
}
metadata := make(map[string]interface{})
if strings.TrimSpace(metadataJSON) != "" && strings.TrimSpace(metadataJSON) != "{}" {
if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil {
return fmt.Errorf("verify task runner resolution metadata: %w", err)
}
}
currentProfileID, _ := metadata[models.MetaKeyExecutorProfileID].(string)
if currentProfileID != session.TaskRunnerProfileAtResolution || currentProfileID != session.ExecutorProfileID {
return models.ErrTaskRunnerChanged
}
return nil
}
func (r *Repository) persistWorkflowSessionRouteTx(
ctx context.Context,
exec taskSessionExecutor,
taskID, destinationID string,
route *models.WorkflowSessionRoute,
) error {
if route == nil {
return nil
}
prepared := *route
if prepared.DestinationID == "" {
prepared.DestinationID = destinationID
}
if prepared.DestinationID != destinationID {
return fmt.Errorf("workflow route destination does not match session")
}
if prepared.Phase == "" {
prepared.Phase = "prepared"
}
return r.setTaskMetadataKeyWithExecutor(ctx, exec, taskID, models.MetaKeyWorkflowSessionRoute, prepared, time.Now().UTC())
}
// setTaskMetadataKeyWithExecutor is SetTaskMetadataKey's tx-capable sibling,
// following the same shape as removeTaskMetadataKeyWithExecutor: every
// existing single-key metadata writer executes on the shared handle and
// none accepts a transaction, but a caller inside a serialized transaction
// (workflow session route persistence, the runner switch) needs its write to
// land only if that transaction commits. updatedAt is supplied by the
// caller, rather than sampled here, so the value written to the row and the
// value the caller carries forward (into a returned task or event) are the
// same instant rather than two independent clock reads either side of the
// write.
func (r *Repository) setTaskMetadataKeyWithExecutor(
ctx context.Context,
exec taskSessionExecutor,
taskID, key string,
value interface{},
updatedAt time.Time,
) error {
payload, err := json.Marshal(value)
if err != nil {
return err
}
if updatedAt.IsZero() {
updatedAt = time.Now().UTC()
}
var query string
if dialect.IsPostgres(r.db.DriverName()) {
query = `UPDATE tasks SET metadata = jsonb_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END, ARRAY[?]::text[], ?::jsonb, true)::text, updated_at = ? WHERE id = ?`
} else {
query = `UPDATE tasks SET metadata = json_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?, json(?)), updated_at = ? WHERE id = ?`
}
path := key
if !dialect.IsPostgres(r.db.DriverName()) {
path = jsonPath(key)
}
result, err := exec.ExecContext(ctx, r.db.Rebind(query), path, string(payload), updatedAt, taskID)
if err != nil {
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows == 0 {
return fmt.Errorf("task not found: %s", taskID)
}
return nil
}
func (r *Repository) commitWorkflowSessionRouteTx(
ctx context.Context,
tx *sqlx.Tx,
taskID, sessionID string,
route *models.WorkflowSessionRoute,
now time.Time,
) error {
if route == nil {
return nil
}
var metadataJSON sql.NullString
if err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT metadata FROM tasks WHERE id = ?`), taskID).Scan(&metadataJSON); err != nil {
return err
}
metadata := make(map[string]interface{})
if metadataJSON.Valid && metadataJSON.String != "" && metadataJSON.String != "null" {
if err := json.Unmarshal([]byte(metadataJSON.String), &metadata); err != nil {
return fmt.Errorf("decode workflow session route: %w", err)
}
}
recorded, ok := models.LoadWorkflowSessionRoute(metadata)
if !ok || recorded.OperationID != route.OperationID || recorded.DestinationID != sessionID {
return fmt.Errorf("workflow session route was superseded before promotion")
}
committed := *route
committed.DestinationID = sessionID
committed.Phase = "committed"
return r.setTaskMetadataKeyWithExecutor(ctx, tx, taskID, models.MetaKeyWorkflowSessionRoute, committed, now)
}
// prepareWorkflowInitialSessionTx elects the first task session and records
// its immutable workflow-routing snapshot in the same transaction as the
// session insert. Later callers cannot replace the snapshot or leave a stale
// origin marker on a session created after the first one.
func (r *Repository) prepareWorkflowInitialSessionTx(ctx context.Context, tx *sqlx.Tx, session *models.TaskSession) error {
var sessionCount int
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT COUNT(*) FROM task_sessions WHERE task_id = ?`,
), session.TaskID).Scan(&sessionCount); err != nil {
return fmt.Errorf("check task sessions before workflow snapshot: %w", err)
}
if sessionCount == 0 {
if session.Metadata == nil {
session.Metadata = make(map[string]interface{})
}
session.Metadata[models.SessionMetaKeyOrigin] = models.SessionOriginTaskInitial
_, err := r.setTaskMetadataKeyIfAbsentWithoutTimestampWithExecutor(ctx, tx, session.TaskID, models.MetaKeyWorkflowInitialSession, models.WorkflowInitialSessionSnapshot{
SessionID: session.ID,
AgentProfileID: session.AgentProfileID,
})
return err
}
if models.IsOriginalTaskSession(session.Metadata) {
delete(session.Metadata, models.SessionMetaKeyOrigin)
}
return nil
}
func (r *Repository) setTaskMetadataKeyIfAbsentWithoutTimestampWithExecutor(ctx context.Context, exec taskSessionExecutor, taskID, key string, value interface{}) (bool, error) {
return r.setTaskMetadataKeyIfAbsentWithExecutorOptions(ctx, exec, taskID, key, value, false)
}
func (r *Repository) setTaskMetadataKeyIfAbsentWithExecutorOptions(ctx context.Context, exec taskSessionExecutor, taskID, key string, value interface{}, updateTaskTimestamp bool) (bool, error) {
payload, err := json.Marshal(value)
if err != nil {
return false, err
}
args := []interface{}{}
var query string
if dialect.IsPostgres(r.db.DriverName()) {
query = `UPDATE tasks
SET metadata = jsonb_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END, ARRAY[?]::text[], ?::jsonb, true)::text`
if updateTaskTimestamp {
query += `, updated_at = ?`
}
query += ` WHERE id = ? AND jsonb_extract_path(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END, ?) IS NULL`
} else {
query = `UPDATE tasks
SET metadata = json_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?, json(?))`
if updateTaskTimestamp {
query += `, updated_at = ?`
}
query += ` WHERE id = ? AND json_type(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?) IS NULL`
}
path := key
if !dialect.IsPostgres(r.db.DriverName()) {
path = jsonPath(key)
}
args = append(args, path, string(payload))
if updateTaskTimestamp {
args = append(args, time.Now().UTC())
}
args = append(args, taskID, path)
result, err := exec.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
return rows > 0, err
}
// unmarshalSessionJSON deserializes a JSON string into dest, skipping empty/placeholder values.
func unmarshalSessionJSON(jsonStr string, dest interface{}, fieldDesc string) error {
if jsonStr == "" || jsonStr == "{}" {
return nil
}
if err := json.Unmarshal([]byte(jsonStr), dest); err != nil {
return fmt.Errorf("failed to deserialize %s: %w", fieldDesc, err)
}
return nil
}
func (r *Repository) scanTaskSession(ctx context.Context, row *sql.Row, noRowsErr string) (*models.TaskSession, error) {
session := &models.TaskSession{}
var state string
var metadataJSON string
var agentProfileSnapshotJSON string
var executorSnapshotJSON string
var environmentSnapshotJSON string
var repositorySnapshotJSON string
var completedAt sql.NullTime
var isPrimary int
var isPassthrough int
var reviewStatus sql.NullString
// agent_profile_id is nullable (kanban / quick-chat rows store NULL); decode
// via NullString so the empty case maps to "" on the model.
var agentProfileID sql.NullString
var name sql.NullString
var lastReadMessageID sql.NullString
err := row.Scan(
&session.ID, &session.TaskID, &session.QueueIncarnationID,
&session.AgentExecutionID, &session.ContainerID, &agentProfileID,
&session.ExecutionProfileID, &session.RouteGeneration, &session.RouteState, &session.RouteReason, &session.DownstreamACPSessionID,
&session.ExecutorID, &session.ExecutorProfileID, &session.EnvironmentID,
&session.RepositoryID, &session.BaseBranch, &session.BaseCommitSHA, &session.WorkspacePath,
&agentProfileSnapshotJSON, &executorSnapshotJSON, &environmentSnapshotJSON, &repositorySnapshotJSON,
&state, &session.ErrorMessage, &metadataJSON, &session.StartedAt, &completedAt, &session.UpdatedAt,
&isPrimary, &reviewStatus, &isPassthrough, &session.TaskEnvironmentID, &name, &lastReadMessageID,
&session.CostSubcents, &session.TokensIn, &session.TokensCachedIn, &session.TokensOut,
)
if err == sql.ErrNoRows {
return nil, fmt.Errorf("%w: %s", models.ErrTaskSessionNotFound, noRowsErr)
}
if err != nil {
return nil, err
}
session.State = models.TaskSessionState(state)
session.IsPrimary = isPrimary == 1
session.IsPassthrough = isPassthrough == 1
if reviewStatus.Valid {
session.ReviewStatus = models.ReviewStatus(reviewStatus.String)
}
if agentProfileID.Valid {
session.AgentProfileID = agentProfileID.String
}
if name.Valid {
session.Name = name.String
}
if lastReadMessageID.Valid {
session.LastReadMessageID = lastReadMessageID.String
}
if completedAt.Valid {
session.CompletedAt = &completedAt.Time
}
if err := unmarshalSessionJSON(metadataJSON, &session.Metadata, "agent session metadata"); err != nil {
return nil, err
}
if err := unmarshalSessionJSON(agentProfileSnapshotJSON, &session.AgentProfileSnapshot, "agent profile snapshot"); err != nil {
return nil, err
}
if err := unmarshalSessionJSON(executorSnapshotJSON, &session.ExecutorSnapshot, "executor snapshot"); err != nil {
return nil, err
}
if err := unmarshalSessionJSON(environmentSnapshotJSON, &session.EnvironmentSnapshot, "environment snapshot"); err != nil {
return nil, err
}
if err := unmarshalSessionJSON(repositorySnapshotJSON, &session.RepositorySnapshot, "repository snapshot"); err != nil {
return nil, err
}
worktrees, err := r.ListTaskSessionWorktrees(ctx, session.ID)
if err != nil {
return nil, fmt.Errorf("failed to load session worktrees: %w", err)
}
session.Worktrees = worktrees
if err := r.hydrateDynamicRoutePolicy(ctx, session); err != nil {
return nil, err
}
return session, nil
}
// hydrateDynamicRoutePolicy projects the durable route decision into session
// responses. The route row remains the source of truth, so this projection
// survives a backend restart without adding a second mutable policy store to
// task_sessions.
func (r *Repository) hydrateDynamicRoutePolicy(ctx context.Context, session *models.TaskSession) error {
if session == nil || session.RouteGeneration <= 0 {
return nil
}
var policyStateJSON string
err := r.ro.QueryRowxContext(ctx, r.ro.Rebind(
`SELECT policy_state_json FROM dynamic_route_states WHERE session_id = ?`,
), session.ID).Scan(&policyStateJSON)
if errors.Is(err, sql.ErrNoRows) {
return nil
}
if err != nil {
return fmt.Errorf("failed to load dynamic route policy projection: %w", err)
}
var projection struct {
FailureCode string `json:"failure_code"`
FailureClass string `json:"failure_class"`
CatalogueVersion string `json:"catalogue_version"`
RetryOrdinal int64 `json:"retry_ordinal"`
Deadline *time.Time `json:"deadline"`
PendingOutcome string `json:"pending_outcome"`
}
if policyStateJSON == "" {
return nil
}
if err := json.Unmarshal([]byte(policyStateJSON), &projection); err != nil {
return fmt.Errorf("failed to decode dynamic route policy projection: %w", err)
}
session.RouteErrorCode = projection.FailureCode
session.RouteErrorClass = projection.FailureClass
session.RouteCatalogueVersion = projection.CatalogueVersion
session.RouteRetryOrdinal = projection.RetryOrdinal
session.RouteDeadline = projection.Deadline
session.RoutePendingOutcome = projection.PendingOutcome
return nil
}
// GetTaskSession retrieves an agent session by ID
func (r *Repository) GetTaskSession(ctx context.Context, id string) (*models.TaskSession, error) {
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.id = ?`,
), id)
return r.scanTaskSession(ctx, row, fmt.Sprintf("agent session not found: %s", id))
}
// ClaimPromptableTaskSessionIfActive atomically claims a ready session for a
// prompt while its task is still active. It intentionally performs no agent
// I/O; callers dispatch only after this bounded database claim commits.
func (r *Repository) ClaimPromptableTaskSessionIfActive(ctx context.Context, id string) (models.PromptableTaskSessionClaim, error) {
return r.claimPromptableTaskSessionIfActive(ctx, id, "", "")
}
func (r *Repository) ClaimPromptableTaskSessionIfActiveForIdentity(
ctx context.Context,
taskID, id, incarnationID string,
) (models.PromptableTaskSessionClaim, error) {
if taskID == "" || id == "" || incarnationID == "" {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionInactive}, nil
}
return r.claimPromptableTaskSessionIfActive(ctx, id, taskID, incarnationID)
}
func (r *Repository) claimPromptableTaskSessionIfActive(
ctx context.Context,
id, taskID, incarnationID string,
) (models.PromptableTaskSessionClaim, error) {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
defer func() { _ = tx.Rollback() }()
var state models.TaskSessionState
var active bool
err = tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT ts.state, t.archived_at IS NULL
FROM task_sessions ts JOIN tasks t ON t.id = ts.task_id
WHERE ts.id = ?
AND (? = '' OR ts.task_id = ?)
AND (? = '' OR ts.queue_incarnation_id = ?)
`), id, taskID, taskID, incarnationID, incarnationID).Scan(&state, &active)
if errors.Is(err, sql.ErrNoRows) {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionInactive}, nil
}
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
if !active {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionInactive}, nil
}
if !isPromptableSessionState(state) {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionBusy}, nil
}
result, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET state = ?, completed_at = NULL, updated_at = ?
WHERE id = ? AND state = ?
AND (? = '' OR task_id = ?)
AND (? = '' OR queue_incarnation_id = ?)
AND EXISTS (SELECT 1 FROM tasks WHERE tasks.id = task_sessions.task_id AND tasks.archived_at IS NULL)
`), models.TaskSessionStateRunning, time.Now().UTC(), id, state,
taskID, taskID, incarnationID, incarnationID)
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
changed, err := result.RowsAffected()
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
if changed == 0 {
claim, err := r.classifyPromptableTaskSessionClaimForIdentity(ctx, tx, id, taskID, incarnationID)
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
if err := tx.Commit(); err != nil {
return models.PromptableTaskSessionClaim{}, err
}
return claim, nil
}
if err := tx.Commit(); err != nil {
return models.PromptableTaskSessionClaim{}, err
}
return models.PromptableTaskSessionClaim{
Status: models.PromptableTaskSessionClaimed, PreviousState: state,
}, nil
}
// classifyPromptableTaskSessionClaim distinguishes a concurrent session-state
// transition from task/session removal after a guarded claim updates no rows.
// It must run in the claim transaction so the result describes the same
// ownership window as the failed UPDATE.
func (r *Repository) classifyPromptableTaskSessionClaim(
ctx context.Context,
tx *sqlx.Tx,
id string,
) (models.PromptableTaskSessionClaim, error) {
return r.classifyPromptableTaskSessionClaimForIdentity(ctx, tx, id, "", "")
}
func (r *Repository) classifyPromptableTaskSessionClaimForIdentity(
ctx context.Context,
tx *sqlx.Tx,
id, taskID, incarnationID string,
) (models.PromptableTaskSessionClaim, error) {
var state models.TaskSessionState
var active bool
err := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT ts.state, t.archived_at IS NULL
FROM task_sessions ts JOIN tasks t ON t.id = ts.task_id
WHERE ts.id = ?
AND (? = '' OR ts.task_id = ?)
AND (? = '' OR ts.queue_incarnation_id = ?)
`), id, taskID, taskID, incarnationID, incarnationID).Scan(&state, &active)
if errors.Is(err, sql.ErrNoRows) {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionInactive}, nil
}
if err != nil {
return models.PromptableTaskSessionClaim{}, err
}
if !active {
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionInactive}, nil
}
return models.PromptableTaskSessionClaim{Status: models.PromptableTaskSessionBusy}, nil
}
func isPromptableSessionState(state models.TaskSessionState) bool {
return state == models.TaskSessionStateWaitingForInput ||
state == models.TaskSessionStateIdle ||
state == models.TaskSessionStateCompleted
}
// GetTaskSessionByTaskID retrieves the most recent agent session for a task
func (r *Repository) GetTaskSessionByTaskID(ctx context.Context, taskID string) (*models.TaskSession, error) {
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.task_id = ? ORDER BY ts.started_at DESC LIMIT 1`,
), taskID)
return r.scanTaskSession(ctx, row, fmt.Sprintf("agent session not found for task: %s", taskID))
}
// GetActiveTaskSessionByTaskID retrieves the active (running/waiting) agent session for a task
func (r *Repository) GetActiveTaskSessionByTaskID(ctx context.Context, taskID string) (*models.TaskSession, error) {
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+`
WHERE ts.task_id = ? AND ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
ORDER BY ts.started_at DESC LIMIT 1`,
), taskID)
return r.scanTaskSession(ctx, row, fmt.Sprintf("no active agent session for task: %s", taskID))
}
// GetTaskSessionByTaskAndAgent retrieves the office task session for the given
// (task_id, agent_profile_id) pair. This pair is NOT unique at the schema
// level — the office creation guard (CreateOfficeTaskSession) constrains it
// only at the moment of insert, so legacy duplicate-pair data and rows
// created outside that guard can still leave multiple rows for the same
// pair. The ORDER BY therefore prefers a live row over a terminal one
// regardless of which is newer — a terminal row created after a live row
// (e.g. a stale duplicate resolving while the real session is still running)
// must not shadow the live one — then falls back to started_at DESC, then id
// DESC as a total tiebreak. Returns nil, nil when no session exists for the
// pair.
func (r *Repository) GetTaskSessionByTaskAndAgent(ctx context.Context, taskID, agentInstanceID string) (*models.TaskSession, error) {
if taskID == "" || agentInstanceID == "" {
return nil, nil
}
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+`
WHERE ts.task_id = ? AND ts.agent_profile_id = ?
ORDER BY
CASE WHEN ts.state IN ('COMPLETED', 'FAILED', 'CANCELLED') THEN 1 ELSE 0 END,
ts.started_at DESC,
ts.id DESC
LIMIT 1`,
), taskID, agentInstanceID)
session, err := r.scanTaskSession(ctx, row, "task_sessions: no matching row")
if errors.Is(err, models.ErrTaskSessionNotFound) {
return nil, nil
}
return session, err
}
// ListNonTerminalSessionsByAgentInstance returns every office task_session row
// for the given agent_profile_id whose state is NOT terminal
// (CREATED / STARTING / RUNNING / IDLE / WAITING_FOR_INPUT). Used by the
// agent-instance deletion cascade in office, which must terminate all of an
// agent's live sessions across every task.
func (r *Repository) ListNonTerminalSessionsByAgentInstance(ctx context.Context, agentInstanceID string) ([]*models.TaskSession, error) {
if agentInstanceID == "" {
return nil, nil
}
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+`
WHERE ts.agent_profile_id = ?
AND ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'IDLE', 'WAITING_FOR_INPUT')`,
), agentInstanceID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
return r.scanTaskSessions(ctx, rows)
}
// UpdateTaskSession updates an existing agent session.
// Note: metadata is NOT updated here to prevent clobbering concurrent writes
// from UpdateSessionMetadata callers (e.g. setSessionPlanMode, persistPrepareResult).
// Use UpdateSessionMetadata for metadata changes.
func (r *Repository) UpdateTaskSession(ctx context.Context, session *models.TaskSession) error {
session.UpdatedAt = time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.updateTaskSession(ctx, tx, session); err != nil {
return err
}
return tx.Commit()
}
// UpdateTaskSessionIfCurrentState persists a full session row only while the
// stored state still matches expected. Runtime launch/resume paths use this to
// retain their non-state fields without reviving a concurrently cancelled
// session from a stale snapshot.
func (r *Repository) UpdateTaskSessionIfCurrentState(
ctx context.Context,
session *models.TaskSession,
expected models.TaskSessionState,
) (bool, error) {
session.UpdatedAt = time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return false, err
}
defer func() { _ = tx.Rollback() }()
changed, err := r.updateTaskSessionWithStateGuard(ctx, tx, session, &expected)
if err != nil || !changed {
return changed, err
}
if err := tx.Commit(); err != nil {
return false, err
}
return true, nil
}
// UpdateTaskSessionIfCurrentStateRemovingMetadataKeys persists a full session
// row and removes provider-owned metadata atomically while the stored state
// still matches expected. JSON removal preserves unrelated concurrent keys.
func (r *Repository) UpdateTaskSessionIfCurrentStateRemovingMetadataKeys(
ctx context.Context,
session *models.TaskSession,
expected models.TaskSessionState,
keys []string,
) (bool, error) {
session.UpdatedAt = time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return false, err
}
defer func() { _ = tx.Rollback() }()
changed, err := r.updateTaskSessionWithStateGuard(ctx, tx, session, &expected)
if err != nil || !changed {
return changed, err
}
if err := r.removeSessionMetadataKeys(ctx, tx, session.ID, keys, session.UpdatedAt); err != nil {
return false, err
}
if err := tx.Commit(); err != nil {
return false, err
}
return true, nil
}
func (r *Repository) removeSessionMetadataKeys(
ctx context.Context,
exec taskSessionExecutor,
sessionID string,
keys []string,
updatedAt time.Time,
) error {
if len(keys) == 0 {
return nil
}
driver := r.db.DriverName()
var query string
args := make([]interface{}, 0, len(keys)+2)
if dialect.IsPostgres(driver) {
expr := "(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END)"
for range keys {
expr = fmt.Sprintf("(%s #- ARRAY[?]::text[])", expr)
}
for _, key := range keys {
args = append(args, key)
}
query = fmt.Sprintf("UPDATE task_sessions SET metadata = %s::text, updated_at = ? WHERE id = ?", expr)
} else {
paths := make([]string, len(keys))
for i, key := range keys {
paths[i] = "?"
args = append(args, "$."+key)
}
query = fmt.Sprintf(
"UPDATE task_sessions SET metadata = json_remove(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, %s), updated_at = ? WHERE id = ?",
strings.Join(paths, ", "),
)
}
args = append(args, updatedAt, sessionID)
result, err := exec.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("agent session not found: %s", sessionID)
}
return nil
}
// UpdateTaskSessionWithMetadata updates the session row and metadata column in
// one transaction so callers cannot observe a partially-applied update.
func (r *Repository) UpdateTaskSessionWithMetadata(
ctx context.Context,
session *models.TaskSession,
metadata map[string]interface{},
) error {
session.UpdatedAt = time.Now().UTC()
metadataJSON, err := marshalSessionMetadata(metadata)
if err != nil {
return err
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.updateTaskSession(ctx, tx, session); err != nil {
return err
}
if err := r.updateSessionMetadataJSON(ctx, tx, session.ID, metadataJSON, session.UpdatedAt); err != nil {
return err
}
return tx.Commit()
}
func (r *Repository) updateTaskSession(
ctx context.Context,
exec taskSessionExecutor,
session *models.TaskSession,
) error {
changed, err := r.updateTaskSessionWithStateGuard(ctx, exec, session, nil)
if err != nil {
return err
}
if !changed {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, session.ID)
}
return nil
}
// updateTaskSessionWithStateGuard writes a full session row while an optional
// expected state guard still matches the stored row.
func (r *Repository) updateTaskSessionWithStateGuard(
ctx context.Context,
exec taskSessionExecutor,
session *models.TaskSession,
expected *models.TaskSessionState,
) (bool, error) {
if tx, ok := exec.(*sqlx.Tx); ok {
if err := r.ensureTaskSessionEnvironmentAvailableTx(ctx, tx, session.ID, session.TaskEnvironmentID); err != nil {
return false, err
}
}
agentProfileSnapshotJSON, err := json.Marshal(session.AgentProfileSnapshot)
if err != nil {
return false, fmt.Errorf("failed to serialize agent profile snapshot: %w", err)
}
executorSnapshotJSON, err := json.Marshal(session.ExecutorSnapshot)
if err != nil {
return false, fmt.Errorf("failed to serialize executor snapshot: %w", err)
}
environmentSnapshotJSON, err := json.Marshal(session.EnvironmentSnapshot)
if err != nil {
return false, fmt.Errorf("failed to serialize environment snapshot: %w", err)
}
repositorySnapshotJSON, err := json.Marshal(session.RepositorySnapshot)
if err != nil {
return false, fmt.Errorf("failed to serialize repository snapshot: %w", err)
}
// metadata is NOT written here — callers wanting to change it must use
// UpdateSessionMetadata or SetSessionMetadataKey. A full-row write here
// would clobber metadata set via those side-channel paths since the
// caller's in-memory copy may be stale.
// agent_profile_id is stored as NULL when empty. No unique index
// constrains (task_id, agent_profile_id) at this UPDATE path — office
// session-uniqueness enforcement lives only in CreateOfficeTaskSession's
// in-transaction guard (see ErrOfficeSessionRaceConflict's doc comment).
var agentProfileID interface{}
if session.AgentProfileID != "" {
agentProfileID = session.AgentProfileID
}
query := `
UPDATE task_sessions SET
agent_profile_id = ?, execution_profile_id = ?, route_generation = ?, route_state = ?, route_reason = ?, downstream_acp_session_id = ?,
executor_id = ?, executor_profile_id = ?, environment_id = ?,
repository_id = ?, base_branch = ?, base_commit_sha = ?, workspace_path = ?,
agent_profile_snapshot = ?, executor_snapshot = ?, environment_snapshot = ?, repository_snapshot = ?,
state = ?, error_message = ?, completed_at = ?, updated_at = ?,
is_primary = ?, review_status = ?, is_passthrough = ?, task_environment_id = ?
WHERE id = ?`
args := []interface{}{agentProfileID, session.ExecutionProfileID, session.RouteGeneration, session.RouteState, session.RouteReason, session.DownstreamACPSessionID,
session.ExecutorID, session.ExecutorProfileID, session.EnvironmentID,
session.RepositoryID, session.BaseBranch, session.BaseCommitSHA, session.WorkspacePath,
string(agentProfileSnapshotJSON), string(executorSnapshotJSON), string(environmentSnapshotJSON), string(repositorySnapshotJSON),
string(session.State), session.ErrorMessage, session.CompletedAt, session.UpdatedAt,
dialect.BoolToInt(session.IsPrimary), session.ReviewStatus,
dialect.BoolToInt(session.IsPassthrough), session.TaskEnvironmentID,
session.ID}
if expected != nil {
query += " AND state = ?"
args = append(args, string(*expected))
}
result, err := exec.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, err
}
return rows > 0, nil
}
// ensureTaskSessionEnvironmentAvailableTx protects full-row session updates
// that can attach a session to an environment. The current environment and a
// newly supplied environment are both checked because a stale session object
// can otherwise move an environment pointer through an active recovery claim.
func (r *Repository) ensureTaskSessionEnvironmentAvailableTx(
ctx context.Context,
tx *sqlx.Tx,
sessionID, nextEnvironmentID string,
) error {
var current sql.NullString
if err := tx.QueryRowContext(ctx, r.db.Rebind(`
SELECT task_environment_id FROM task_sessions WHERE id = ?
`), sessionID).Scan(&current); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
return err
}
if current.Valid && current.String != "" {
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, current.String); err != nil {
return err
}
}
if nextEnvironmentID != "" && (!current.Valid || current.String != nextEnvironmentID) {
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, nextEnvironmentID); err != nil {
return err
}
}
return nil
}
// RenameTaskSession updates just the user-supplied name of an agent session.
// name is deliberately excluded from the full-row updateTaskSession write (like
// metadata) so concurrent session updates can't clobber a rename.
func (r *Repository) RenameTaskSession(ctx context.Context, id, name string) error {
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET name = ?, updated_at = ? WHERE id = ?
`), name, time.Now().UTC(), id)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, id)
}
return nil
}
// UpdateTaskSessionAgentProfileSnapshot updates only the profile snapshot.
// Runtime capability events may carry a stale full session row, so keeping
// this write narrow prevents them from overwriting terminal state.
func (r *Repository) UpdateTaskSessionAgentProfileSnapshot(
ctx context.Context,
id string,
snapshot map[string]interface{},
) error {
snapshotJSON, err := json.Marshal(snapshot)
if err != nil {
return fmt.Errorf("failed to serialize agent profile snapshot: %w", err)
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET agent_profile_snapshot = ?, updated_at = ?
WHERE id = ?
`), string(snapshotJSON), time.Now().UTC(), id)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, id)
}
return nil
}
// UpdateTaskSessionState updates just the state and error message of an agent session
func (r *Repository) UpdateTaskSessionState(ctx context.Context, id string, status models.TaskSessionState, errorMessage string) error {
now := time.Now().UTC()
completedAt := completedAtForTaskSessionState(status, now)
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET state = ?, error_message = ?, completed_at = ?, updated_at = ? WHERE id = ?
`), string(status), errorMessage, completedAt, now, id)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, id)
}
return nil
}
// UpdateTaskSessionStateIfCurrent transitions a session only when its state
// still matches the caller's observation. The returned timestamp belongs to
// the committed write and remains authoritative even if a later read fails.
func (r *Repository) UpdateTaskSessionStateIfCurrent(
ctx context.Context,
id string,
expected, status models.TaskSessionState,
errorMessage string,
) (bool, time.Time, error) {
now := time.Now().UTC()
completedAt := completedAtForTaskSessionState(status, now)
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE id = ? AND state = ?
`), string(status), errorMessage, completedAt, now, id, string(expected))
if err != nil {
return false, time.Time{}, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, time.Time{}, err
}
return rows > 0, now, nil
}
func (r *Repository) UpdateTaskSessionStateIfCurrentIdentity(
ctx context.Context,
taskID, id, incarnationID string,
expected, status models.TaskSessionState,
errorMessage string,
) (bool, time.Time, error) {
now := time.Now().UTC()
completedAt := completedAtForTaskSessionState(status, now)
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE id = ? AND task_id = ? AND queue_incarnation_id = ? AND state = ?
`), string(status), errorMessage, completedAt, now,
id, taskID, incarnationID, string(expected))
if err != nil {
return false, time.Time{}, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, time.Time{}, err
}
return rows > 0, now, nil
}
// UpdateTaskSessionDynamicRouteIfCurrent changes only the route projection
// while the session generation and projected route state still match the
// caller's observation. It avoids writing a stale full session row after an
// asynchronous launch or recovery callback.
func (r *Repository) UpdateTaskSessionDynamicRouteIfCurrent(
ctx context.Context,
id string,
expectedGeneration int64,
expectedRouteState, routeState, routeReason string,
) (bool, time.Time, error) {
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET route_state = ?, route_reason = ?, updated_at = ?
WHERE id = ? AND route_generation = ? AND route_state = ?
`), routeState, routeReason, now, id, expectedGeneration, expectedRouteState)
if err != nil {
return false, time.Time{}, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, time.Time{}, err
}
return rows > 0, now, nil
}
// CancelActiveTaskSession atomically transitions one active session, including
// a parked IDLE session, to CANCELLED. A false result means the row exists in
// a non-active state, or it was concurrently changed before this conditional
// write; callers re-read to distinguish those cases from a missing row. The
// returned timestamp belongs
// to the committed cancellation, so accepting callers never need a fallible
// post-write read before scheduling teardown.
func (r *Repository) CancelActiveTaskSession(ctx context.Context, id, reason string) (bool, time.Time, error) {
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE id = ?
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT', 'IDLE')
`), string(models.TaskSessionStateCancelled), reason, now, now, id)
if err != nil {
return false, time.Time{}, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, time.Time{}, err
}
return rows > 0, now, nil
}
func completedAtForTaskSessionState(status models.TaskSessionState, now time.Time) *time.Time {
if status == models.TaskSessionStateCompleted ||
status == models.TaskSessionStateFailed ||
status == models.TaskSessionStateCancelled {
return &now
}
return nil
}
// CancelActiveTaskSessionsByTaskID transitions every active session of a task
// (CREATED/STARTING/RUNNING/WAITING_FOR_INPUT) to CANCELLED, returning the
// full row of each session actually transitioned. The transition is a pure
// DB state change and does not require a live agent execution, making it
// the authoritative way to finalize a task's sessions independent of
// agent-process teardown.
//
// The UPDATE and the row selection happen in a single atomic statement via
// RETURNING, so the returned rows are exactly what this call changed — a
// session created or transitioned to active state concurrently, after this
// statement starts, is simply outside its snapshot; the task_id + state
// predicate is evaluated once per row as it commits, so no session matching
// it at commit time is missed by a separate pre-update snapshot. Callers use
// the returned sessions to publish a matching session.state_changed event
// per session — without this, clients that cache session state
// independently of the task (e.g. an Office task list's "is running"
// indicator) never learn the session left its active state and spin
// forever after the owning task is archived.
//
// The RETURNING clause carries every field publishSessionsCancelled needs
// to build its event payload directly, so callers never fall back to a
// separate post-commit read (e.g. GetTaskSession) to assemble the event —
// closing the read-after-write gap where a session could commit CANCELLED
// but its event never gets published because that follow-up read failed or
// timed out. Returned sessions therefore carry only the fields the
// RETURNING clause selects (ID, TaskID, AgentProfileID,
// AgentProfileSnapshot, IsPassthrough, Name, ReviewStatus, Metadata,
// TaskEnvironmentID, State, UpdatedAt, IsPrimary) — every other models.TaskSession
// field is left at its zero value, and callers must not rely on fields
// outside this list being populated.
func (r *Repository) CancelActiveTaskSessionsByTaskID(ctx context.Context, taskID, reason string) ([]*models.TaskSession, error) {
now := time.Now().UTC()
// Detach from ctx via WithoutCancel: this write must survive a client
// disconnect (see the doc comment above) so the CANCELLED transition and
// its returned rows are never lost to a caller that hung up mid-request.
// Bound it with a timeout so a locked SQLite UPDATE can't block forever
// on a request-independent DB stall just because the deadline was
// dropped.
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel()
rows, err := r.db.QueryContext(writeCtx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE task_id = ?
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
RETURNING id, agent_profile_id, agent_profile_snapshot, is_passthrough, name,
review_status, metadata, task_environment_id, state, updated_at, is_primary
`), string(models.TaskSessionStateCancelled), reason, now, now, taskID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var sessions []*models.TaskSession
for rows.Next() {
session, err := scanCancelledTaskSessionRow(rows, taskID)
if err != nil {
return nil, err
}
sessions = append(sessions, session)
}
return sessions, rows.Err()
}
// ListStaleRunningSessionsOnUnarchivedTasks returns every STARTING/RUNNING
// session of an unarchived task whose updated_at is older than staleBefore.
// It is the candidate list for the orphan-session reconciliation sweep (see
// service.runOrphanedSessionReconciliation): a session still holding one of
// those states past the launch grace window, with no live in-memory execution
// backing it, can never reach a terminal state on its own — its actor died
// with the process. The staleBefore cutoff is applied in SQL so the sweep
// never even loads fresh rows that may belong to an in-flight launch.
//
// Not part of the SessionRepository interface: it exists for that one sweep,
// which reaches it through the narrow orphanedSessionRepository capability —
// widening the interface would force every test double in the tree to grow
// methods this sweep never exercises through them.
func (r *Repository) ListStaleRunningSessionsOnUnarchivedTasks(ctx context.Context, staleBefore time.Time) ([]*models.TaskSession, error) {
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+`
JOIN tasks t ON t.id = ts.task_id
WHERE ts.state IN ('STARTING', 'RUNNING')
AND ts.updated_at < ?
AND t.archived_at IS NULL
ORDER BY ts.updated_at ASC
`), staleBefore)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
sessions, err := r.scanTaskSessions(ctx, rows)
if err != nil {
return nil, err
}
return r.loadWorktreesBatch(ctx, sessions)
}
// CancelRunningTaskSessionByID transitions a single stale STARTING/RUNNING
// session to CANCELLED, returning the transitioned row, or nil when the
// session no longer matches (already terminal, raced to another state, or
// refreshed since the sweep read it). Like CancelActiveTaskSessionsByTaskID
// it is a pure DB state change that requires no live agent execution, and it
// is session-scoped so the orphan sweep can terminalize one session without
// cancelling healthy sibling sessions of the same task.
//
// staleBefore re-asserts the sweep's staleness cutoff at write time, closing
// the read-then-write race against an in-flight launch: a launch CAS-writes
// its session to STARTING (bumping updated_at) before it registers an
// execution in the in-memory store, so between the sweep's liveness check and
// this UPDATE the row can pass from "no live execution" to "launch in
// progress". A row refreshed since the candidate read no longer satisfies
// updated_at < staleBefore, the UPDATE matches nothing, and the next tick
// re-evaluates a fresh row that the grace window protects until its
// execution registers.
//
// The UPDATE and the row selection happen in one atomic RETURNING statement;
// the returned row carries only the fields that clause selects (same set as
// CancelActiveTaskSessionsByTaskID, minus TaskID — callers pass the task ID
// they already know). The write detaches from ctx like
// CancelActiveTaskSessionsByTaskID: once the sweep has decided this session
// is orphaned, the terminal transition must not be lost to a caller-context
// cancellation, and the 10s bound keeps a locked SQLite writer from stalling
// the sweep pass. A failed write simply retries on the next sweep tick.
func (r *Repository) CancelRunningTaskSessionByID(ctx context.Context, sessionID, reason string, staleBefore time.Time) (*models.TaskSession, error) {
now := time.Now().UTC()
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel()
rows, err := r.db.QueryContext(writeCtx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE id = ?
AND state IN ('STARTING', 'RUNNING')
AND updated_at < ?
RETURNING id, agent_profile_id, agent_profile_snapshot, is_passthrough, name,
review_status, metadata, task_environment_id, state, updated_at, is_primary
`), string(models.TaskSessionStateCancelled), reason, now, now, sessionID, staleBefore)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
if !rows.Next() {
return nil, rows.Err()
}
// The scanner backfills TaskID from the parameter because RETURNING
// does not carry it; sessionID is not the task ID, so set it from the
// candidate row the sweep already read instead.
session, err := scanCancelledTaskSessionRow(rows, "")
if err != nil {
return nil, err
}
return session, rows.Err()
}
// RecoverTaskSessionByCandidate returns one execution-less session to
// WAITING_FOR_INPUT when its state, activity clock, and current turn still
// match the reconciliation snapshot. The optional staleBefore cutoff is used
// by the restart pass to preserve its launch grace window; a zero cutoff is
// used by the active-task stall pass.
//
// The write is session-scoped and compare-and-set guarded. A newer message,
// state transition, successor turn, or refreshed launch causes the UPDATE to
// match no row, leaving the current owner untouched for the next sweep.
//
//nolint:cyclop,funlen // The SQL predicate mirrors the complete recovery CAS contract.
func (r *Repository) RecoverTaskSessionByCandidate(
ctx context.Context,
candidate models.ActiveSessionRecoveryCandidate,
staleBefore time.Time,
) (*models.TaskSession, error) {
if candidate.SessionID == "" || candidate.TaskID == "" {
return nil, nil
}
switch candidate.ExpectedState {
case models.TaskSessionStateCreated,
models.TaskSessionStateStarting,
models.TaskSessionStateRunning,
models.TaskSessionStateWaitingForInput:
default:
return nil, nil
}
// PostgreSQL timestamp columns retain microsecond precision. Use the same
// precision in the JSON settlement snapshot so a retry compares the
// stored generation exactly on both database dialects.
now := time.Now().UTC().Truncate(time.Microsecond)
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel()
// Preserve a durable recovery token and immutable effect snapshot in the
// same compare-and-set write as the WAITING_FOR_INPUT transition. A later
// focus request can therefore distinguish this recovered conversation from
// an ordinary idle session even when executors_running has no row, and a
// partial post-commit settlement can retry without inspecting a successor.
settlementPending := candidate.ExpectedTurnID != "" ||
candidate.ExpectedState == models.TaskSessionStateCreated ||
candidate.ExpectedState == models.TaskSessionStateStarting ||
candidate.ExpectedState == models.TaskSessionStateRunning
setClause := "state = ?, error_message = ?, completed_at = ?, updated_at = ?"
query := `
UPDATE task_sessions
SET `
args := make([]interface{}, 0, 16)
if settlementPending {
token := interruptedRecoveryToken(candidate)
settlement := models.InterruptedRecoverySettlement{
Token: token,
ExpectedState: candidate.ExpectedState,
RecoveredUpdatedAt: now,
ExpectedTurnID: candidate.ExpectedTurnID,
ExpectedExecutorID: candidate.ExpectedExecutorID,
ExpectedExecutorAgentExecutionID: candidate.ExpectedExecutorAgentExecutionID,
ExpectedExecutorUpdatedAt: candidate.ExpectedExecutorUpdatedAt,
}
pendingJSON, err := json.Marshal(token)
if err != nil {
return nil, err
}
settlementJSON, err := json.Marshal(settlement)
if err != nil {
return nil, err
}
if dialect.IsPostgres(r.db.DriverName()) {
base := postgresMetadataObject
setClause = "metadata = jsonb_set(jsonb_set(" + base + ", ARRAY[?]::text[], ?::jsonb, true), ARRAY[?]::text[], ?::jsonb, true)::text, " + setClause
args = append(args,
models.SessionMetaKeyInterruptedRecoveryPending, string(pendingJSON),
models.SessionMetaKeyRecoverySettlementPending, string(settlementJSON),
)
} else {
base := sqliteMetadataObject
setClause = "metadata = json_set(json_set(" + base + ", ?, json(?)), ?, json(?)), " + setClause
args = append(args,
jsonPath(models.SessionMetaKeyInterruptedRecoveryPending), string(pendingJSON),
jsonPath(models.SessionMetaKeyRecoverySettlementPending), string(settlementJSON),
)
}
}
query += setClause + `
WHERE id = ?
AND task_id = ?
AND EXISTS (
SELECT 1 FROM tasks recoverable_task
WHERE recoverable_task.id = task_sessions.task_id
AND recoverable_task.archived_at IS NULL
)
AND state = ?
AND updated_at = ?
`
args = append(args,
string(models.TaskSessionStateWaitingForInput), "", nil, now,
candidate.SessionID, candidate.TaskID, string(candidate.ExpectedState), candidate.ExpectedUpdatedAt,
)
if !staleBefore.IsZero() {
query += " AND updated_at < ?\n"
args = append(args, staleBefore)
}
if !candidate.ExpectedLastEventAt.IsZero() {
query += `
AND NOT EXISTS (
SELECT 1 FROM task_session_messages newer_message
WHERE newer_message.task_session_id = task_sessions.id
AND newer_message.updated_at > ?
)
`
args = append(args, candidate.ExpectedLastEventAt)
}
if candidate.ExpectedTurnID == "" {
query += `
AND NOT EXISTS (
SELECT 1 FROM task_session_turns active_turn
WHERE active_turn.task_session_id = task_sessions.id
AND active_turn.completed_at IS NULL
)
`
} else {
query += `
AND EXISTS (
SELECT 1 FROM task_session_turns active_turn
WHERE active_turn.id = ?
AND active_turn.task_session_id = task_sessions.id
AND active_turn.completed_at IS NULL
)
`
args = append(args, candidate.ExpectedTurnID)
}
query += `
RETURNING id, agent_profile_id, agent_profile_snapshot, is_passthrough, name,
review_status, metadata, task_environment_id, state, updated_at, is_primary
`
rows, err := r.db.QueryContext(writeCtx, r.db.Rebind(query), args...)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
if !rows.Next() {
return nil, rows.Err()
}
session, err := scanCancelledTaskSessionRow(rows, candidate.TaskID)
if err != nil {
return nil, err
}
return session, rows.Err()
}
func interruptedRecoveryToken(candidate models.ActiveSessionRecoveryCandidate) string {
return fmt.Sprintf("%s:%s:%s:%s", candidate.SessionID,
candidate.ExpectedState,
candidate.ExpectedUpdatedAt.UTC().Format(time.RFC3339Nano),
candidate.ExpectedTurnID)
}
// CancelActiveTaskSessionsByIDs is documented on the SessionRepository
// interface. It shares CancelActiveTaskSessionsByTaskID's atomic
// UPDATE ... RETURNING shape and detached-but-bounded write context, so the
// returned rows are exactly the sessions this call transitioned, and a
// client disconnect mid-write cannot lose the committed cancellation.
// Chunking keeps the placeholder count below the host-parameter limit; the
// active set of one task is far below it, so multi-chunk calls are a
// correctness backstop, not an expected path.
func (r *Repository) CancelActiveTaskSessionsByIDs(ctx context.Context, taskID string, sessionIDs []string, reason string) ([]*models.TaskSession, error) {
if len(sessionIDs) == 0 {
return nil, nil
}
now := time.Now().UTC()
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel()
var sessions []*models.TaskSession
for _, chunk := range chunkIDs(sessionIDs, sqliteMaxHostParams) {
placeholders, args := buildInPlaceholders(chunk)
args = append([]interface{}{string(models.TaskSessionStateCancelled), reason, now, now, taskID}, args...)
rows, err := r.db.QueryContext(writeCtx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE task_id = ?
AND id IN (`+placeholders+`)
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
RETURNING id, agent_profile_id, agent_profile_snapshot, is_passthrough, name,
review_status, metadata, task_environment_id, state, updated_at, is_primary
`), args...)
if err != nil {
return nil, err
}
for rows.Next() {
session, err := scanCancelledTaskSessionRow(rows, taskID)
if err != nil {
_ = rows.Close()
return nil, err
}
sessions = append(sessions, session)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, err
}
_ = rows.Close()
}
return sessions, nil
}
// CancelActiveTaskSessionsByCandidates cancels only candidates whose session
// activity and current-turn identity are unchanged since classification. The
// compare-and-set predicates run in the same UPDATE that transitions the row,
// so a resumed session, a newer message, or a successor turn wins the race and
// remains active for the next reconciliation pass.
func (r *Repository) CancelActiveTaskSessionsByCandidates(
ctx context.Context,
taskID string,
candidates []models.ActiveSessionCancellationCandidate,
reason string,
) ([]*models.TaskSession, error) {
if len(candidates) == 0 {
return nil, nil
}
now := time.Now().UTC()
writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel()
// Each candidate contributes at most five bind parameters. Keep a margin
// below SQLite's host-parameter limit and use the same bounded-chunk shape
// as the ID-scoped cancellation path.
chunkSize := sqliteMaxHostParams / 5
if chunkSize < 1 {
chunkSize = 1
}
var sessions []*models.TaskSession
for start := 0; start < len(candidates); start += chunkSize {
end := start + chunkSize
if end > len(candidates) {
end = len(candidates)
}
chunk := candidates[start:end]
predicates := make([]string, 0, len(chunk))
predicateArgs := make([]interface{}, 0, len(chunk)*5)
for _, candidate := range chunk {
predicate, args := activeSessionCancellationCandidatePredicate(candidate)
predicates = append(predicates, predicate)
predicateArgs = append(predicateArgs, args...)
}
args := []interface{}{string(models.TaskSessionStateCancelled), reason, now, now, taskID}
args = append(args, predicateArgs...)
rows, err := r.db.QueryContext(writeCtx, r.db.Rebind(`
UPDATE task_sessions
SET state = ?, error_message = ?, completed_at = ?, updated_at = ?
WHERE task_id = ?
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
AND (`+strings.Join(predicates, " OR ")+`)
RETURNING id, agent_profile_id, agent_profile_snapshot, is_passthrough, name,
review_status, metadata, task_environment_id, state, updated_at, is_primary
`), args...)
if err != nil {
return nil, err
}
for rows.Next() {
session, err := scanCancelledTaskSessionRow(rows, taskID)
if err != nil {
_ = rows.Close()
return nil, err
}
sessions = append(sessions, session)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, err
}
_ = rows.Close()
}
return sessions, nil
}
func activeSessionCancellationCandidatePredicate(
candidate models.ActiveSessionCancellationCandidate,
) (string, []interface{}) {
predicate := "(id = ? AND updated_at = ?"
args := []interface{}{candidate.SessionID, candidate.ExpectedUpdatedAt}
if !candidate.ExpectedLastEventAt.IsZero() {
predicate += `
AND NOT EXISTS (
SELECT 1 FROM task_session_messages newer_message
WHERE newer_message.task_session_id = task_sessions.id
AND newer_message.updated_at > ?
)`
args = append(args, candidate.ExpectedLastEventAt)
}
if candidate.ExpectedTurnID == "" {
predicate += `
AND NOT EXISTS (
SELECT 1 FROM task_session_turns active_turn
WHERE active_turn.task_session_id = task_sessions.id
AND active_turn.completed_at IS NULL
)`
} else {
predicate += `
AND EXISTS (
SELECT 1 FROM task_session_turns active_turn
WHERE active_turn.id = ?
AND active_turn.task_session_id = task_sessions.id
AND active_turn.completed_at IS NULL
)`
args = append(args, candidate.ExpectedTurnID)
}
return predicate + ")", args
}
// scanCancelledTaskSessionRow scans one row produced by
// CancelActiveTaskSessionsByTaskID's UPDATE ... RETURNING into a
// *models.TaskSession, mirroring scanTaskSessionRow's JSON-unmarshal and
// int-to-bool/nullable-string conventions but for the narrower RETURNING
// column set (id, agent_profile_id, agent_profile_snapshot, is_passthrough,
// name, review_status, metadata, task_environment_id, state, updated_at,
// is_primary). taskID backfills TaskID, which RETURNING cannot supply since
// it's a query parameter, not a returned column.
func scanCancelledTaskSessionRow(rows *sql.Rows, taskID string) (*models.TaskSession, error) {
session := &models.TaskSession{TaskID: taskID}
var state string
var metadataJSON string
var agentProfileSnapshotJSON string
var isPassthrough int
var isPrimary int
var reviewStatus sql.NullString
var agentProfileID sql.NullString
var name sql.NullString
if err := rows.Scan(
&session.ID, &agentProfileID, &agentProfileSnapshotJSON, &isPassthrough, &name,
&reviewStatus, &metadataJSON, &session.TaskEnvironmentID, &state, &session.UpdatedAt,
&isPrimary,
); err != nil {
return nil, err
}
session.State = models.TaskSessionState(state)
session.IsPassthrough = isPassthrough == 1
session.IsPrimary = isPrimary == 1
if reviewStatus.Valid {
session.ReviewStatus = models.ReviewStatus(reviewStatus.String)
}
if agentProfileID.Valid {
session.AgentProfileID = agentProfileID.String
}
if name.Valid {
session.Name = name.String
}
if err := unmarshalSessionJSON(metadataJSON, &session.Metadata, "agent session metadata"); err != nil {
return nil, err
}
if err := unmarshalSessionJSON(agentProfileSnapshotJSON, &session.AgentProfileSnapshot, "agent profile snapshot"); err != nil {
return nil, err
}
return session, nil
}
// UpdateSessionMetadata updates only the metadata column of a session,
// avoiding a full-row overwrite that could clobber concurrent field updates.
func (r *Repository) UpdateSessionMetadata(ctx context.Context, sessionID string, metadata map[string]interface{}) error {
metadataJSON, err := marshalSessionMetadata(metadata)
if err != nil {
return err
}
now := time.Now().UTC()
return r.updateSessionMetadataJSON(ctx, r.db, sessionID, metadataJSON, now)
}
func marshalSessionMetadata(metadata map[string]interface{}) (string, error) {
metadataJSON, err := json.Marshal(metadata)
if err != nil {
return "", fmt.Errorf("failed to serialize metadata: %w", err)
}
return string(metadataJSON), nil
}
func (r *Repository) updateSessionMetadataJSON(
ctx context.Context,
exec taskSessionExecutor,
sessionID string,
metadataJSON string,
updatedAt time.Time,
) error {
result, err := exec.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET metadata = ?, updated_at = ? WHERE id = ?
`), metadataJSON, updatedAt, sessionID)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("agent session not found: %s", sessionID)
}
return nil
}
// SetSessionMetadataKey atomically sets a single key in the session's metadata
// using the active database dialect. Unlike UpdateSessionMetadata (which does
// a full replacement), this preserves all other metadata keys.
func (r *Repository) SetSessionMetadataKey(ctx context.Context, sessionID, key string, value interface{}) error {
valueJSON, err := json.Marshal(value)
if err != nil {
return fmt.Errorf("failed to serialize metadata value: %w", err)
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(metadataKeyUpdateQuery("task_sessions", r.db.DriverName())), metadataKeyUpdateArgs(r.db.DriverName(), key, string(valueJSON), r.nowUTC(), sessionID)...)
if err != nil {
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows == 0 {
return fmt.Errorf("agent session not found: %s", sessionID)
}
return nil
}
// SetSessionMetadataKeyIfState atomically sets one metadata key only while the
// session remains in expectedState. Runtime recovery uses this to keep a
// follow-up marker from being attached to a session that was stopped between
// its guarded state transition and metadata persistence.
func (r *Repository) SetSessionMetadataKeyIfState(
ctx context.Context,
sessionID, key string,
value interface{},
expectedState models.TaskSessionState,
) (bool, error) {
valueJSON, err := json.Marshal(value)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
query := metadataKeyUpdateQuery("task_sessions", r.db.DriverName()) + " AND state = ?"
args := metadataKeyUpdateArgs(r.db.DriverName(), key, string(valueJSON), r.nowUTC(), sessionID)
args = append(args, string(expectedState))
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, err
}
return rows > 0, nil
}
// UpdateSessionContextWindow stores a context-window sample and atomically
// increments the session's inferred compaction count when the new used-token
// value is lower than the previous persisted sample.
func (r *Repository) UpdateSessionContextWindow(
ctx context.Context,
sessionID string,
contextWindow map[string]interface{},
) (int64, error) {
windowJSON, err := json.Marshal(contextWindow)
if err != nil {
return 0, fmt.Errorf("failed to serialize context window: %w", err)
}
used, err := contextWindowUsed(contextWindow)
if err != nil {
return 0, err
}
now := time.Now().UTC()
var count int64
row := r.db.QueryRowxContext(ctx, r.db.Rebind(updateSessionContextWindowQuery(r.db.DriverName())), string(windowJSON), used, now, sessionID)
if err := row.Scan(&count); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return 0, fmt.Errorf("agent session not found: %s", sessionID)
}
return 0, err
}
return count, nil
}
func contextWindowUsed(contextWindow map[string]interface{}) (int64, error) {
value, ok := contextWindow["used"]
if !ok {
return 0, errors.New("context window is missing used tokens")
}
switch typed := value.(type) {
case int:
return int64(typed), nil
case int32:
return int64(typed), nil
case int64:
return typed, nil
case float64:
return int64(typed), nil
case json.Number:
used, err := typed.Int64()
if err != nil {
return 0, fmt.Errorf("invalid context window used tokens: %w", err)
}
return used, nil
default:
return 0, fmt.Errorf("invalid context window used tokens: %T", value)
}
}
//nolint:dupword // nested JSON setter calls are intentional in the atomic SQL.
func updateSessionContextWindowQuery(driver string) string {
if dialect.IsPostgres(driver) {
base := "CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END"
count := "GREATEST(COALESCE(NULLIF((" + base + ") #>> '{context_compaction_count}', '')::bigint, 0), 0)"
previousUsed := "(" + base + ") #>> '{context_window,used}'"
return `
UPDATE task_sessions
SET metadata = jsonb_set(
jsonb_set(` + base + `, '{context_window}', ?::jsonb, true),
'{context_compaction_count}',
to_jsonb(CASE
WHEN jsonb_typeof((` + base + `) #> '{context_window,used}') = 'number'
AND (` + previousUsed + `)::bigint > ?
THEN ` + count + ` + 1
ELSE ` + count + `
END),
true
)::text,
updated_at = ?
WHERE id = ?
RETURNING ((metadata::jsonb) #>> '{context_compaction_count}')::bigint
`
}
base := "CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END"
count := "MAX(COALESCE(CAST(json_extract(" + base + ", '$.context_compaction_count') AS INTEGER), 0), 0)"
previousUsed := "json_extract(" + base + ", '$.context_window.used')"
return `
UPDATE task_sessions
SET metadata = json_set(
json_set(` + base + `, '$.context_window', json(?)),
'$.context_compaction_count',
CASE
WHEN json_type(` + base + `, '$.context_window.used') IN ('integer', 'real')
AND CAST(` + previousUsed + ` AS INTEGER) > ?
THEN ` + count + ` + 1
ELSE ` + count + `
END
),
updated_at = ?
WHERE id = ?
RETURNING CAST(json_extract(metadata, '$.context_compaction_count') AS INTEGER)
`
}
// SetSessionMetadataKeyIfAbsent atomically writes a metadata key only when it
// does not already exist. The returned bool reports whether this call stored
// the value.
func (r *Repository) SetSessionMetadataKeyIfAbsent(
ctx context.Context,
sessionID string,
key string,
value interface{},
) (bool, error) {
valueJSON, err := json.Marshal(value)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
now := time.Now().UTC()
driver := r.db.DriverName()
path := key
if !dialect.IsPostgres(driver) {
path = "$." + key
}
query := setSessionMetadataKeyIfAbsentQuery(driver)
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), path, string(valueJSON), now, sessionID, path)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
// SetSessionMetadataKeyIfAbsentOrDifferentStep atomically writes a metadata
// key when it is absent or its stored value has a different step_id. The
// returned bool reports whether this call stored the value.
func (r *Repository) SetSessionMetadataKeyIfAbsentOrDifferentStep(
ctx context.Context,
sessionID, key, stepID string,
value interface{},
) (bool, error) {
valueJSON, err := json.Marshal(value)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
now := time.Now().UTC()
driver := r.db.DriverName()
path := key
stepPath := key
if !dialect.IsPostgres(driver) {
path = "$." + key
stepPath = path + ".step_id"
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(setSessionMetadataKeyIfAbsentOrDifferentStepQuery(driver)),
path, string(valueJSON), now, sessionID, stepPath, stepID)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
// SetSessionMetadataKeyIfAbsentOrDifferentStepIfTaskAtStep claims a pending
// completion signal only while the session's task remains on the turn's
// launch step. The task row is locked before the session metadata write so a
// concurrent workflow move cannot turn a stale signal into a valid successor
// signal.
func (r *Repository) SetSessionMetadataKeyIfAbsentOrDifferentStepIfTaskAtStep(
ctx context.Context,
taskID, sessionID, key, stepID string,
value interface{},
) (bool, error) {
valueJSON, err := json.Marshal(value)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
tx, err := r.db.BeginTx(ctx, nil)
if err != nil {
return false, err
}
committed := false
defer func() {
if !committed {
_ = tx.Rollback()
}
}()
_, currentStepID, found, err := r.readTaskStepInTx(ctx, tx, taskID)
if err != nil {
return false, err
}
if !found || currentStepID != stepID {
return false, nil
}
var sessionTaskID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT task_id FROM task_sessions WHERE id = ?`,
), sessionID).Scan(&sessionTaskID); err != nil {
if err == sql.ErrNoRows {
return false, nil
}
return false, err
}
if sessionTaskID != taskID {
return false, nil
}
now := time.Now().UTC()
driver := r.db.DriverName()
path := key
stepPath := key
if !dialect.IsPostgres(driver) {
path = "$." + key
stepPath = path + ".step_id"
}
result, err := tx.ExecContext(ctx, r.db.Rebind(setSessionMetadataKeyIfAbsentOrDifferentStepQuery(driver)),
path, string(valueJSON), now, sessionID, stepPath, stepID)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
if err != nil {
return false, err
}
if err := tx.Commit(); err != nil {
return false, err
}
committed = true
return rows > 0, nil
}
// SetSessionMetadataKeyIfAbsentIfState atomically claims a metadata key only
// while the session remains in expectedState. It is used when a terminal
// transition owns a one-time side effect that must not be emitted by a stale
// launch callback.
func (r *Repository) SetSessionMetadataKeyIfAbsentIfState(
ctx context.Context,
sessionID, key string,
value interface{},
expectedState models.TaskSessionState,
) (bool, error) {
valueJSON, err := json.Marshal(value)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
now := time.Now().UTC()
driver := r.db.DriverName()
path := key
if !dialect.IsPostgres(driver) {
path = "$." + key
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(setSessionMetadataKeyIfAbsentIfStateQuery(driver)), path, string(valueJSON), now, sessionID, path, expectedState)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
// RemoveSessionMetadataKeyIfState removes a claimed metadata key only while
// the session remains in expectedState. It releases one-time side-effect
// claims when their downstream write fails, allowing a later retry.
func (r *Repository) RemoveSessionMetadataKeyIfState(
ctx context.Context,
sessionID, key string,
expectedState models.TaskSessionState,
) (bool, error) {
now := time.Now().UTC()
driver := r.db.DriverName()
path := key
if !dialect.IsPostgres(driver) {
path = "$." + key
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(removeSessionMetadataKeyIfStateQuery(driver)), path, now, sessionID, path, expectedState)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
// RemoveSessionMetadataKeyIfStamp removes a metadata object only when its
// nested stamp still equals expectedStamp. The comparison and key removal are
// one statement so a successful recovery cannot erase a newer session error.
func (r *Repository) RemoveSessionMetadataKeyIfStamp(
ctx context.Context,
sessionID, key, expectedStamp string,
) (bool, error) {
if strings.TrimSpace(expectedStamp) == "" {
return false, nil
}
now := time.Now().UTC()
driver := r.db.DriverName()
var query string
var args []interface{}
if dialect.IsPostgres(driver) {
query = `
UPDATE task_sessions
SET metadata = (CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END #- ARRAY[?]::text[])::text,
updated_at = ?
WHERE id = ?
AND jsonb_extract_path_text(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END, ?, 'stamp') = ?
`
args = []interface{}{key, now, sessionID, key, expectedStamp}
} else {
path := "$" + "." + key
query = `
UPDATE task_sessions
SET metadata = json_remove(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?),
updated_at = ?
WHERE id = ?
AND json_extract(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?) = ?
`
args = []interface{}{path, now, sessionID, path + ".stamp", expectedStamp}
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
return rows > 0, err
}
// RemoveSessionMetadataKeyIfJSONValue removes one metadata key only when its
// complete JSON value still equals expectedValue. Recovery settlement uses
// this compare-and-set boundary so a delayed effect cannot erase a newer
// recovery snapshot written after a successor launch.
func (r *Repository) RemoveSessionMetadataKeyIfJSONValue(
ctx context.Context,
sessionID, key string,
expectedValue interface{},
) (bool, error) {
payload, err := json.Marshal(expectedValue)
if err != nil {
return false, fmt.Errorf("failed to serialize expected session metadata: %w", err)
}
now := time.Now().UTC()
driver := r.db.DriverName()
var query string
var args []interface{}
if dialect.IsPostgres(driver) {
query = `
UPDATE task_sessions
SET metadata = (
CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END
#- ARRAY[?]::text[]
)::text, updated_at = ?
WHERE id = ?
AND (CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END -> ?) = ?::jsonb
`
args = []interface{}{key, now, sessionID, key, string(payload)}
} else {
path := jsonPath(key)
query = `
UPDATE task_sessions
SET metadata = json_remove(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?), updated_at = ?
WHERE id = ?
AND json_extract(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?) = json(?)
`
args = []interface{}{path, now, sessionID, path, string(payload)}
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), args...)
if err != nil {
return false, err
}
rows, err := result.RowsAffected()
return rows > 0, err
}
func setSessionMetadataKeyIfAbsentQuery(driver string) string {
if dialect.IsPostgres(driver) {
return `
UPDATE task_sessions
SET metadata = jsonb_set(
CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END,
ARRAY[?]::text[],
?::jsonb,
true
)::text,
updated_at = ?
WHERE id = ?
AND jsonb_extract_path(
CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END,
?
) IS NULL
`
}
return `
UPDATE task_sessions
SET metadata = json_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?, json(?)),
updated_at = ?
WHERE id = ?
AND json_type(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?) IS NULL
`
}
func setSessionMetadataKeyIfAbsentOrDifferentStepQuery(driver string) string {
if dialect.IsPostgres(driver) {
base := "CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END"
return `
UPDATE task_sessions
SET metadata = jsonb_set(
` + base + `,
ARRAY[?]::text[],
?::jsonb,
true
)::text,
updated_at = ?
WHERE id = ?
AND jsonb_extract_path_text(` + base + `, ?, 'step_id') IS DISTINCT FROM ?
`
}
base := "CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END"
return `
UPDATE task_sessions
SET metadata = json_set(` + base + `, ?, json(?)),
updated_at = ?
WHERE id = ?
AND (
json_extract(` + base + `, ?) IS NOT ?
)
`
}
func setSessionMetadataKeyIfAbsentIfStateQuery(driver string) string {
return setSessionMetadataKeyIfAbsentQuery(driver) + " AND state = ?"
}
func removeSessionMetadataKeyIfStateQuery(driver string) string {
if dialect.IsPostgres(driver) {
return `
UPDATE task_sessions
SET metadata = (
CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END
#- ARRAY[?]::text[]
)::text,
updated_at = ?
WHERE id = ?
AND jsonb_extract_path(
CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}'::jsonb ELSE metadata::jsonb END,
?
) IS NOT NULL
AND state = ?
`
}
return `
UPDATE task_sessions
SET metadata = json_remove(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?),
updated_at = ?
WHERE id = ?
AND json_type(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, ?) IS NOT NULL
AND state = ?
`
}
// SetSessionACPSessionID mirrors the agent's ACP session id into the session's
// "acp" metadata map as a single atomic UPDATE. json_patch merges the sub-key,
// so keys session_info already wrote (title, ...) survive without a
// read-modify-write round trip. The write only happens while the session's
// executors_running row still holds acpSessionID as its resume token — i.e.
// the CAS that stored the token hasn't been superseded by a rotated execution
// — and is skipped when the stored id is already current, so a no-op never
// touches updated_at. Returns whether a row was written.
func (r *Repository) SetSessionACPSessionID(ctx context.Context, sessionID, acpSessionID string) (bool, error) {
patch, err := json.Marshal(map[string]interface{}{"acp": map[string]string{"session_id": acpSessionID}})
if err != nil {
return false, fmt.Errorf("failed to serialize metadata patch: %w", err)
}
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET metadata = json_patch(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, json(?)),
updated_at = ?
WHERE id = ?
AND json_extract(metadata, '$.acp.session_id') IS NOT ?
AND EXISTS (
SELECT 1 FROM executors_running er
WHERE er.session_id = task_sessions.id AND er.resume_token = ?
)
`), string(patch), now, sessionID, acpSessionID, acpSessionID)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
func (r *Repository) DismissLastAgentError(ctx context.Context, sessionID string, expected models.LastAgentError, dismissedAt time.Time) (bool, error) {
next := expected
next.DismissedAt = &dismissedAt
valueJSON, err := json.Marshal(next)
if err != nil {
return false, fmt.Errorf("failed to serialize metadata value: %w", err)
}
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET metadata = json_set(CASE WHEN metadata IS NULL OR metadata = 'null' OR metadata = '' THEN '{}' ELSE metadata END, '$.last_agent_error', json(?)),
updated_at = ?
WHERE id = ?
AND json_extract(metadata, '$.last_agent_error.message') = ?
AND (
json_extract(metadata, '$.last_agent_error.occurred_at') = ?
OR julianday(json_extract(metadata, '$.last_agent_error.occurred_at')) = julianday(?)
)
`), string(valueJSON), now, sessionID, expected.Message, expected.OccurredAt.UTC().Format(time.RFC3339Nano), expected.OccurredAt.UTC().Format(time.RFC3339Nano))
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
return rows > 0, nil
}
// GetLastAgentMessage returns the content of the most recent agent message in a session.
func (r *Repository) GetLastAgentMessage(ctx context.Context, sessionID string) (string, error) {
var content string
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT content FROM task_session_messages
WHERE task_session_id = ? AND author_type = 'agent' AND type = 'message'
ORDER BY created_at DESC LIMIT 1
`), sessionID).Scan(&content)
if err != nil {
return "", err
}
return content, nil
}
// IncrementTaskSessionUsageTx adds the given deltas to the cumulative
// tokens / cost columns on task_sessions, including cached input tokens
// (tokens_cached_in mirrors office_cost_events.tokens_cached_in and is kept
// separate from tokens_in because it is priced differently). Surfaced on
// models.TaskSession and dto.TaskSessionDTO. internal/task/usage's writer is
// the sole production caller (docs/specs/task-cost-ledger/spec.md AC-10,
// AC-21) — it inserts a task_usage_events row and increments this rollup in
// one transaction (insertUsageEventAndRollup in this package), so this
// method is executed against tx when non-nil (falling back to r.db, the
// shared writer connection, when tx is nil).
func (r *Repository) IncrementTaskSessionUsageTx(
ctx context.Context, tx *sqlx.Tx, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64,
) error {
if sessionID == "" {
return nil
}
var exec interface {
ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
Rebind(query string) string
}
if tx != nil {
exec = tx
} else {
exec = r.db
}
_, err := exec.ExecContext(ctx, exec.Rebind(`
UPDATE task_sessions
SET tokens_in = COALESCE(tokens_in, 0) + ?,
tokens_cached_in = COALESCE(tokens_cached_in, 0) + ?,
tokens_out = COALESCE(tokens_out, 0) + ?,
cost_subcents = COALESCE(cost_subcents, 0) + ?
WHERE id = ?
`), tokensIn, tokensCachedIn, tokensOut, costSubcents, sessionID)
return err
}
// UpdateTaskSessionBaseCommit updates the base_commit_sha for a session.
// This is called after agent launch to capture the HEAD commit at session start.
func (r *Repository) UpdateTaskSessionBaseCommit(ctx context.Context, id string, baseCommitSHA string) error {
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET base_commit_sha = ?, updated_at = ? WHERE id = ?
`), baseCommitSHA, now, id)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, id)
}
return nil
}
// UpdateTaskSessionLastReadMessageID persists messageID as the session's
// read cursor — deliberately a narrow single-column write (like
// RenameTaskSession/UpdateTaskSessionBaseCommit) so it never collides with a
// concurrent metadata or full-row write, and vice versa.
//
// The write is monotonic: it only advances the cursor when messageID is not
// older than the currently persisted one, comparing each message's
// (created_at, id) — created_at as the primary ordering (matches every
// other transcript ordering in this codebase, e.g. ListMessages' `ORDER BY
// created_at ASC`), with id as a deterministic tiebreaker for messages
// created in the same instant. Portable across SQLite and Postgres — no
// dialect branching needed, unlike the rowid-based approach this replaced
// (SQLite's rowid pseudo-column doesn't exist on Postgres). This guards
// against out-of-order delivery — e.g. two overlapping mark-read requests
// where the response for an older message is processed after a newer one —
// silently regressing last_read_message_id and resurrecting the "New"
// divider over transcript the user already read. A cursor referencing a
// message that no longer exists (deleted) is treated as having no rank, so
// the guard never wedges: NOT EXISTS is satisfied and the update proceeds.
func (r *Repository) UpdateTaskSessionLastReadMessageID(ctx context.Context, id, messageID string) error {
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET last_read_message_id = ?, updated_at = ?
WHERE id = ?
AND (
last_read_message_id IS NULL OR last_read_message_id = ''
OR NOT EXISTS (
SELECT 1 FROM task_session_messages cur, task_session_messages incoming
WHERE cur.id = task_sessions.last_read_message_id
AND incoming.id = ?
AND (
cur.created_at > incoming.created_at
OR (cur.created_at = incoming.created_at AND cur.id > incoming.id)
)
)
)
`), messageID, now, id, messageID)
if err != nil {
return err
}
rows, _ := result.RowsAffected()
if rows > 0 {
return nil
}
// No row was updated — either the session doesn't exist, or the guard
// above correctly rejected a stale/out-of-order cursor (messageID is not
// newer than what's already persisted). Distinguish the two so a stale
// update is a silent no-op rather than a spurious not-found error.
var exists int
err = r.db.QueryRowContext(ctx, r.db.Rebind(`SELECT 1 FROM task_sessions WHERE id = ?`), id).Scan(&exists)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, id)
}
return err
}
return nil
}
// ResetTaskSessionBasesForRepository rewrites base_branch and clears
// base_commit_sha on every task_session belonging to (taskID, repositoryID).
// Used by service.UpdateRepositoryBaseBranch after the changes-panel
// "Compare against" picker changes the recorded base.
//
// Clearing base_commit_sha matters because git_handlers.computeGitCommits and
// computeCumulativeDiff read it first and only fall back to the live
// agentctl GetGitStatus().BaseCommit (which now resolves against the new
// base_branch via Phase 1) when the column is empty. Without this reset the
// task-card stat counts would refresh against the new base while the
// commits panel + cumulative diff would keep filtering against the OLD
// commit snapshot — visible to users as "Commits section disappeared
// after I changed the base".
//
// Returns the number of sessions touched so callers can log / no-op when
// the task has no sessions yet.
func (r *Repository) ResetTaskSessionBasesForRepository(ctx context.Context, taskID, repositoryID, baseBranch string) (int64, error) {
if taskID == "" {
return 0, fmt.Errorf("task_id is required")
}
if repositoryID == "" {
return 0, fmt.Errorf("repository_id is required")
}
now := time.Now().UTC()
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET base_branch = ?, base_commit_sha = '', updated_at = ?
WHERE task_id = ? AND repository_id = ?
`), baseBranch, now, taskID, repositoryID)
if err != nil {
return 0, err
}
rows, _ := result.RowsAffected()
return rows, nil
}
// ListTaskSessions returns all agent sessions for a task
func (r *Repository) ListTaskSessions(ctx context.Context, taskID string) ([]*models.TaskSession, error) {
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListTaskSessions")
defer span.End()
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.task_id = ? ORDER BY ts.started_at DESC`,
), taskID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
sessions, err := r.scanTaskSessions(ctx, rows)
if err != nil {
return nil, err
}
return r.loadWorktreesBatch(ctx, sessions)
}
// ListTaskSessionsByTaskEnvironment returns every session bound to an
// environment, including inherited sessions whose task differs from the
// environment owner's task. The environment binding, rather than task row
// ownership, is the source-selection boundary for shared workspaces.
func (r *Repository) ListTaskSessionsByTaskEnvironment(ctx context.Context, environmentID string) ([]*models.TaskSession, error) {
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListTaskSessionsByTaskEnvironment")
defer span.End()
if environmentID == "" {
return []*models.TaskSession{}, nil
}
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.task_environment_id = ? ORDER BY ts.started_at DESC`,
), environmentID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
sessions, err := r.scanTaskSessions(ctx, rows)
if err != nil {
return nil, err
}
return r.loadWorktreesBatch(ctx, sessions)
}
// GetTaskSessionWorkspacePathsByTaskEnvironment returns only the raw workspace
// paths recorded on session rows bound to an environment. It deliberately does
// not use the effective environment projection because an empty or historical
// session path must remain distinguishable from the canonical path.
func (r *Repository) GetTaskSessionWorkspacePathsByTaskEnvironment(ctx context.Context, environmentID string) (map[string]string, error) {
paths := make(map[string]string)
if environmentID == "" {
return paths, nil
}
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
SELECT id, workspace_path
FROM task_sessions
WHERE task_environment_id = ?
`), environmentID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
for rows.Next() {
var sessionID string
var workspacePath sql.NullString
if err := rows.Scan(&sessionID, &workspacePath); err != nil {
return nil, err
}
paths[sessionID] = workspacePath.String
}
if err := rows.Err(); err != nil {
return nil, err
}
return paths, nil
}
// ListActiveTaskSessions returns all active agent sessions across all tasks
func (r *Repository) ListActiveTaskSessions(ctx context.Context) ([]*models.TaskSession, error) {
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListActiveTaskSessions")
defer span.End()
rows, err := r.ro.QueryContext(ctx,
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT') ORDER BY ts.started_at DESC`,
)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
sessions, err := r.scanTaskSessions(ctx, rows)
if err != nil {
return nil, err
}
return r.loadWorktreesBatch(ctx, sessions)
}
// ListLiveWorkspaceSessions returns every session across all tasks whose
// state is one of the five live states (CREATED, STARTING, RUNNING, IDLE,
// WAITING_FOR_INPUT), with each session's effective workspace_path
// (taskSessionSelectCols already prefers the linked task environment's path
// over the possibly-stale session column).
// It exists only for the orphan-reap ownership check: ListActiveTaskSessions
// omits IDLE deliberately for its own unrelated callers, and does not batch
// worktrees, which this check does not need.
func (r *Repository) ListLiveWorkspaceSessions(ctx context.Context) ([]*models.TaskSession, error) {
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListLiveWorkspaceSessions")
defer span.End()
rows, err := r.ro.QueryContext(ctx,
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+
` WHERE ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'IDLE', 'WAITING_FOR_INPUT')`,
)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
return r.scanTaskSessions(ctx, rows)
}
// ListActiveTaskSessionsByTaskID returns all active agent sessions for a specific task
func (r *Repository) ListActiveTaskSessionsByTaskID(ctx context.Context, taskID string) ([]*models.TaskSession, error) {
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.task_id = ? AND ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT') ORDER BY ts.started_at DESC`,
), taskID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
sessions, err := r.scanTaskSessions(ctx, rows)
if err != nil {
return nil, err
}
return r.loadWorktreesBatch(ctx, sessions)
}
// loadWorktreesBatch loads worktrees for multiple sessions in a single query.
func (r *Repository) loadWorktreesBatch(ctx context.Context, sessions []*models.TaskSession) ([]*models.TaskSession, error) {
if len(sessions) == 0 {
return sessions, nil
}
sessionIDs := make([]string, len(sessions))
for i, s := range sessions {
sessionIDs[i] = s.ID
}
worktreeMap, err := r.ListWorktreesBySessionIDs(ctx, sessionIDs)
if err != nil {
return nil, fmt.Errorf("failed to batch-load session worktrees: %w", err)
}
for _, session := range sessions {
session.Worktrees = worktreeMap[session.ID]
}
return sessions, nil
}
// BatchGetSessionsByTaskIDs returns all task_sessions for the given task IDs,
// grouped by task ID and ordered by started_at DESC within each task. The
// returned sessions carry their associated worktrees (loaded in one extra
// query). Used by callers that need primary session, session count, and
// session info for many tasks in one round trip (e.g. the workspace
// task-list endpoint) to avoid issuing one GetSession per task.
//
// Returns an empty map for an empty input. Chunks input larger than
// sqliteMaxHostParams to stay well below SQLite's compile-time placeholder
// limit.
func (r *Repository) BatchGetSessionsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]*models.TaskSession, error) {
result := make(map[string][]*models.TaskSession, len(taskIDs))
if len(taskIDs) == 0 {
return result, nil
}
for _, chunk := range chunkIDs(taskIDs, sqliteMaxHostParams) {
placeholders, args := buildInPlaceholders(chunk)
query := `SELECT ` + taskSessionSelectCols + ` ` + taskSessionFromClause +
` WHERE ts.task_id IN (` + placeholders + `) ORDER BY ts.task_id, ts.started_at DESC`
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, err
}
sessions, scanErr := r.scanTaskSessions(ctx, rows)
_ = rows.Close()
if scanErr != nil {
return nil, scanErr
}
sessions, err = r.loadWorktreesBatch(ctx, sessions)
if err != nil {
return nil, err
}
for _, s := range sessions {
result[s.TaskID] = append(result[s.TaskID], s)
}
}
return result, nil
}
func (r *Repository) HasActiveTaskSessionsByAgentProfile(ctx context.Context, agentProfileID string) (bool, error) {
var exists int
// Exclude ephemeral tasks (quick chat, config chat) - they shouldn't block profile deletion.
// Automation runs are excluded only where they are parked: a finished run
// rests in WAITING_FOR_INPUT so it stays answerable, and counting those would
// let one nightly report block its agent profile from ever being deleted. A
// run that is still working is using the profile now and blocks like any
// other task — see GetActiveTaskInfoByAgentProfile for the same split.
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT 1 FROM task_sessions ts
JOIN tasks t ON ts.task_id = t.id
WHERE ts.agent_profile_id = ?
AND ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
AND t.is_ephemeral = 0
AND NOT (COALESCE(t.origin, '') = '`+models.TaskOriginAutomationRun+`'
AND ts.state = 'WAITING_FOR_INPUT')
LIMIT 1
`), agentProfileID).Scan(&exists)
if err == sql.ErrNoRows {
return false, nil
}
return err == nil, err
}
func (r *Repository) GetActiveTaskInfoByAgentProfile(ctx context.Context, agentProfileID string) ([]agentdto.ActiveTaskInfo, error) {
// This list is "what is blocking this deletion", and automation runs sit on
// both sides of that question.
//
// A run that is CREATED, STARTING or RUNNING is using the profile right now.
// Deleting it out from under live work is the same failure as for any other
// task, so it blocks — the row names a task the user cannot find on a board,
// but the alternative is a silent kill.
//
// A run parked in WAITING_FOR_INPUT is different. That is where every
// finished automation run comes to rest, by design, so counting those would
// let one nightly report block its profile's deletion forever. Those are
// excluded, which makes them non-resumable if the profile goes: replying to
// an old run afterwards fails. That is the accepted trade — see
// docs/specs/office/requirements/automations-settings.md.
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
SELECT DISTINCT t.id, t.title, t.is_ephemeral
FROM task_sessions ts
JOIN tasks t ON t.id = ts.task_id
WHERE ts.agent_profile_id = ? AND ts.state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
AND NOT (COALESCE(t.origin, '') = '`+models.TaskOriginAutomationRun+`'
AND ts.state = 'WAITING_FOR_INPUT')
`), agentProfileID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var result []agentdto.ActiveTaskInfo
for rows.Next() {
var info agentdto.ActiveTaskInfo
if err := rows.Scan(&info.TaskID, &info.TaskTitle, &info.IsEphemeral); err != nil {
return nil, err
}
result = append(result, info)
}
return result, rows.Err()
}
func (r *Repository) HasActiveTaskSessionsByEnvironment(ctx context.Context, environmentID string) (bool, error) {
var exists int
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT 1 FROM task_sessions
WHERE environment_id = ? AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
LIMIT 1
`), environmentID).Scan(&exists)
if err == sql.ErrNoRows {
return false, nil
}
return err == nil, err
}
func (r *Repository) HasActiveTaskSessionsByTaskEnvironmentExcludingTask(ctx context.Context, taskEnvironmentID, taskID string) (bool, error) {
var exists int
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT 1 FROM task_sessions
WHERE task_environment_id = ?
AND task_id != ?
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
LIMIT 1
`), taskEnvironmentID, taskID).Scan(&exists)
if err == sql.ErrNoRows {
return false, nil
}
return err == nil, err
}
func (r *Repository) FindActiveTaskSessionTaskIDByTaskEnvironmentExcludingTask(ctx context.Context, taskEnvironmentID, taskID string) (string, error) {
var borrowerTaskID string
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT task_id FROM task_sessions
WHERE task_environment_id = ?
AND task_id != ?
AND state IN ('CREATED', 'STARTING', 'RUNNING', 'WAITING_FOR_INPUT')
ORDER BY updated_at DESC
LIMIT 1
`), taskEnvironmentID, taskID).Scan(&borrowerTaskID)
if err == sql.ErrNoRows {
return "", nil
}
return borrowerTaskID, err
}
func (r *Repository) HasActiveTaskSessionsByRepository(ctx context.Context, repositoryID string) (bool, error) {
var exists int
// Only sessions of live (non-archived) tasks count; archived tasks never
// block repository deletion. IDLE sessions are resumable and must
// preserve their repository rows.
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT 1
FROM task_sessions s
INNER JOIN task_repositories tr ON tr.task_id = s.task_id
INNER JOIN tasks t ON t.id = s.task_id
WHERE s.state IN ('CREATED', 'STARTING', 'RUNNING', 'IDLE', 'WAITING_FOR_INPUT')
AND tr.repository_id = ?
AND t.archived_at IS NULL
LIMIT 1
`), repositoryID).Scan(&exists)
if err == sql.ErrNoRows {
return false, nil
}
return err == nil, err
}
func (r *Repository) CountActiveTaskSessionsByRepository(ctx context.Context, repositoryID string) (int, error) {
var count int
// Counts only sessions of live (non-archived) tasks, including resumable
// IDLE sessions, matching HasActiveTaskSessionsByRepository.
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
SELECT COUNT(*)
FROM task_sessions s
INNER JOIN task_repositories tr ON tr.task_id = s.task_id
INNER JOIN tasks t ON t.id = s.task_id
WHERE s.state IN ('CREATED', 'STARTING', 'RUNNING', 'IDLE', 'WAITING_FOR_INPUT')
AND tr.repository_id = ?
AND t.archived_at IS NULL
`), repositoryID).Scan(&count)
if err != nil {
return 0, err
}
return count, nil
}
// DeleteEphemeralTasksByAgentProfile deletes all ephemeral tasks (and their sessions)
// that are using the specified agent profile. This is used during profile deletion
// to clean up transient quick chat / config chat tasks.
func (r *Repository) DeleteEphemeralTasksByAgentProfile(ctx context.Context, agentProfileID string) (int64, error) {
// Delete tasks that are ephemeral and have sessions using this profile.
// CASCADE will handle session deletion.
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
DELETE FROM tasks
WHERE is_ephemeral = 1
AND id IN (
SELECT DISTINCT task_id FROM task_sessions WHERE agent_profile_id = ?
)
`), agentProfileID)
if err != nil {
return 0, err
}
return result.RowsAffected()
}
// scanTaskSessions is a helper to scan multiple agent session rows
func (r *Repository) scanTaskSessions(ctx context.Context, rows *sql.Rows) ([]*models.TaskSession, error) {
var result []*models.TaskSession
for rows.Next() {
session, err := scanTaskSessionRow(rows)
if err != nil {
return nil, err
}
result = append(result, session)
}
return result, rows.Err()
}
// scanTaskSessionRow scans a single row into a TaskSession, applying all field mappings and JSON unmarshalling.
func scanTaskSessionRow(rows *sql.Rows) (*models.TaskSession, error) {
session := &models.TaskSession{}
var state string
var metadataJSON string
var agentProfileSnapshotJSON string
var executorSnapshotJSON string
var environmentSnapshotJSON string
var repositorySnapshotJSON string
var completedAt sql.NullTime
var isPrimary int
var isPassthrough int
var reviewStatus sql.NullString
var agentProfileID sql.NullString
var name sql.NullString
var lastReadMessageID sql.NullString
err := rows.Scan(
&session.ID, &session.TaskID, &session.QueueIncarnationID,
&session.AgentExecutionID, &session.ContainerID, &agentProfileID,
&session.ExecutionProfileID, &session.RouteGeneration, &session.RouteState, &session.RouteReason, &session.DownstreamACPSessionID,
&session.ExecutorID, &session.ExecutorProfileID, &session.EnvironmentID,
&session.RepositoryID, &session.BaseBranch, &session.BaseCommitSHA, &session.WorkspacePath,
&agentProfileSnapshotJSON, &executorSnapshotJSON, &environmentSnapshotJSON, &repositorySnapshotJSON,
&state, &session.ErrorMessage, &metadataJSON, &session.StartedAt, &completedAt, &session.UpdatedAt,
&isPrimary, &reviewStatus, &isPassthrough, &session.TaskEnvironmentID, &name, &lastReadMessageID,
&session.CostSubcents, &session.TokensIn, &session.TokensCachedIn, &session.TokensOut,
)
if err != nil {
return nil, err
}
session.State = models.TaskSessionState(state)
session.IsPrimary = isPrimary == 1
session.IsPassthrough = isPassthrough == 1
if reviewStatus.Valid {
session.ReviewStatus = models.ReviewStatus(reviewStatus.String)
}
if agentProfileID.Valid {
session.AgentProfileID = agentProfileID.String
}
if name.Valid {
session.Name = name.String
}
if lastReadMessageID.Valid {
session.LastReadMessageID = lastReadMessageID.String
}
if completedAt.Valid {
session.CompletedAt = &completedAt.Time
}
if err := unmarshalSessionSnapshots(session, metadataJSON, agentProfileSnapshotJSON,
executorSnapshotJSON, environmentSnapshotJSON, repositorySnapshotJSON); err != nil {
return nil, err
}
return session, nil
}
// unmarshalSessionSnapshots deserializes all JSON snapshot fields into the session struct.
func unmarshalSessionSnapshots(
session *models.TaskSession,
metadataJSON, agentProfileSnapshotJSON, executorSnapshotJSON, environmentSnapshotJSON, repositorySnapshotJSON string,
) error {
if err := unmarshalSessionJSON(metadataJSON, &session.Metadata, "agent session metadata"); err != nil {
return err
}
if err := unmarshalSessionJSON(agentProfileSnapshotJSON, &session.AgentProfileSnapshot, "agent profile snapshot"); err != nil {
return err
}
if err := unmarshalSessionJSON(executorSnapshotJSON, &session.ExecutorSnapshot, "executor snapshot"); err != nil {
return err
}
if err := unmarshalSessionJSON(environmentSnapshotJSON, &session.EnvironmentSnapshot, "environment snapshot"); err != nil {
return err
}
return unmarshalSessionJSON(repositorySnapshotJSON, &session.RepositorySnapshot, "repository snapshot")
}
// queueSessionLockTablePresent reports whether queue mutations can participate
// in this task transaction. Some repository unit tests intentionally omit the
// message queue schema; production databases always include it.
func (r *Repository) queueSessionLockTablePresent(ctx context.Context) (bool, error) {
return db.TableExistsContext(ctx, r.db, "queue_session_locks")
}
// DeleteTaskSession deletes the exact session incarnation and its pending queue
// rows. The session identity fences delayed deletion from removing a replacement
// that reuses the same textual session ID.
func (r *Repository) DeleteTaskSession(ctx context.Context, session *models.TaskSession) error {
_, err := r.DeleteTaskSessionWithAttachments(ctx, session)
return err
}
// DeleteTaskSessionWithAttachments also returns claimed attachment descriptors
// whose registry rows were removed, allowing the service layer to remove their
// private bytes after the transaction commits.
func (r *Repository) DeleteTaskSessionWithAttachments(
ctx context.Context,
session *models.TaskSession,
) ([]*models.TaskMessageAttachment, error) {
if session == nil || session.ID == "" || session.TaskID == "" || session.QueueIncarnationID == "" {
return nil, errors.New("task session identity is required")
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return nil, err
}
defer func() { _ = tx.Rollback() }()
if err := r.confirmTaskSessionDeletionIdentityTx(ctx, tx, session); err != nil {
return nil, err
}
var environmentID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT COALESCE(task_environment_id, '') FROM task_sessions WHERE id = ?`), session.ID).Scan(&environmentID); err != nil {
return nil, fmt.Errorf("load deleted session environment: %w", err)
}
if err := recoveryclaim.EnsureAvailableTx(ctx, r.db, tx, environmentID); err != nil {
return nil, err
}
deletedAttachments, err := r.purgeTaskSessionStateTx(ctx, tx, session)
if err != nil {
return nil, err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_workflow_session_bindings
SET session_id = NULL, updated_at = ?
WHERE session_id = ?
`), r.nowUTC(), session.ID); err != nil {
if !db.IsMissingTableError(err) {
return nil, fmt.Errorf("clear workflow session binding: %w", err)
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
r.notifyTaskSessionQueuePurged(ctx, session.TaskID, session.ID)
return deletedAttachments, nil
}
// Task Session Worktree operations
func (r *Repository) confirmTaskSessionDeletionIdentityTx(
ctx context.Context,
tx *sqlx.Tx,
session *models.TaskSession,
) error {
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE tasks SET updated_at = updated_at WHERE id = ?
`), session.TaskID); err != nil {
return fmt.Errorf("lock deleted session task: %w", err)
}
var found string
err := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT id
FROM task_sessions
WHERE id = ? AND task_id = ? AND queue_incarnation_id = ?
`), session.ID, session.TaskID, session.QueueIncarnationID).Scan(&found)
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("agent session not found: %s", session.ID)
}
if err != nil {
return fmt.Errorf("confirm deleted session identity: %w", err)
}
return nil
}
func (r *Repository) purgeTaskSessionStateTx(
ctx context.Context,
tx *sqlx.Tx,
session *models.TaskSession,
) ([]*models.TaskMessageAttachment, error) {
identity := messagequeue.QueueSessionIdentity{
TaskID: session.TaskID,
SessionID: session.ID,
SessionIncarnationID: session.QueueIncarnationID,
}
if err := messagequeue.DeleteSessionInTransaction(ctx, tx, r.db, identity); err != nil {
if !db.IsMissingTableError(err) {
return nil, fmt.Errorf("purge deleted session queue: %w", err)
}
}
deletedAttachments, err := r.deleteUnreferencedSessionAttachmentClaimsTx(ctx, tx, session.TaskID, session.ID)
if err != nil {
return nil, fmt.Errorf("purge deleted session attachment claims: %w", err)
}
result, err := tx.ExecContext(ctx, r.db.Rebind(`
DELETE FROM task_sessions
WHERE id = ? AND task_id = ? AND queue_incarnation_id = ?
`), session.ID, session.TaskID, session.QueueIncarnationID)
if err != nil {
return nil, err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return nil, fmt.Errorf("agent session not found: %s", session.ID)
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM task_session_prompt_seq WHERE task_session_id = ?`), session.ID); err != nil {
if !db.IsMissingTableError(err) {
return nil, fmt.Errorf("purge prompt history for session %s: %w", session.ID, err)
}
}
return deletedAttachments, nil
}
//
// Sessions reference worktrees only through task_sessions.task_environment_id;
// the physical worktree records live on task_environment_repos (owned by the
// task environment, not the session). All session-scoped queries below join
// through that link.
// envRepoSelectCols is the SELECT projection for a task_environment_repos row
// in session-scoped worktree queries.
const envRepoSelectCols = `
ter.id, ter.task_environment_id, ter.repository_id,
COALESCE(ter.branch_slug, ''), COALESCE(ter.worktree_id, ''),
COALESCE(ter.worktree_path, ''), COALESCE(ter.worktree_branch, ''),
ter.position, COALESCE(ter.error_message, ''), ter.status,
ter.created_at, ter.updated_at, ter.merged_at, ter.deleted_at`
// scanEnvRepoRow scans one task_environment_repos row into a
// models.TaskEnvironmentRepo.
func scanEnvRepoRow(scanner rowScanner) (*models.TaskEnvironmentRepo, error) {
row := &models.TaskEnvironmentRepo{}
var mergedAt, deletedAt sql.NullTime
if err := scanner.Scan(
&row.ID, &row.TaskEnvironmentID, &row.RepositoryID, &row.BranchSlug,
&row.WorktreeID, &row.WorktreePath, &row.WorktreeBranch, &row.Position,
&row.ErrorMessage, &row.Status, &row.CreatedAt, &row.UpdatedAt,
&mergedAt, &deletedAt,
); err != nil {
return nil, err
}
if mergedAt.Valid {
t := mergedAt.Time
row.MergedAt = &t
}
if deletedAt.Valid {
t := deletedAt.Time
row.DeletedAt = &t
}
return row, nil
}
// rowScanner abstracts *sql.Row and *sql.Rows so env-repo rows scan through
// one helper.
type rowScanner interface {
Scan(dest ...interface{}) error
}
// ListTaskSessionWorktrees returns the active environment-repository rows of
// the session's task environment. Sessions sharing an environment observe the
// same worktrees.
func (r *Repository) ListTaskSessionWorktrees(ctx context.Context, sessionID string) ([]*models.TaskEnvironmentRepo, error) {
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
SELECT `+envRepoSelectCols+`
FROM task_environment_repos ter
INNER JOIN task_sessions ts ON ts.task_environment_id = ter.task_environment_id
WHERE ts.id = ?
AND ter.deleted_at IS NULL
AND ter.status = 'active'
ORDER BY ter.position ASC, ter.created_at ASC
`), sessionID)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var repos []*models.TaskEnvironmentRepo
for rows.Next() {
row, err := scanEnvRepoRow(rows)
if err != nil {
return nil, err
}
repos = append(repos, row)
}
return repos, rows.Err()
}
// UpdateTaskSessionWorktreeBranch updates the cached worktree_branch for all
// repository rows of the session's task environment. Called when a branch
// switch or rename is detected in the live workspace so downstream consumers
// (PR watch reconciliation, branch listings) see the current branch rather
// than the value captured at worktree creation.
func (r *Repository) UpdateTaskSessionWorktreeBranch(ctx context.Context, sessionID, branch string) error {
now := time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.ensureTaskSessionEnvironmentAvailableTx(ctx, tx, sessionID, ""); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_environment_repos SET worktree_branch = ?, updated_at = ?
WHERE task_environment_id = (SELECT task_environment_id FROM task_sessions WHERE id = ?)
AND deleted_at IS NULL
AND status = 'active'
`), branch, now, sessionID); err != nil {
return err
}
return tx.Commit()
}
// UpdateTaskSessionWorktreeBranchByRepository updates the cached worktree_branch
// for one repository row in the session's environment. Use this for repo-scoped
// live git operations in multi-repo tasks so sibling repositories keep their
// branch snapshots.
func (r *Repository) UpdateTaskSessionWorktreeBranchByRepository(ctx context.Context, sessionID, repositoryID, branch string) error {
now := time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.ensureTaskSessionEnvironmentAvailableTx(ctx, tx, sessionID, ""); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_environment_repos
SET worktree_branch = ?, updated_at = ?
WHERE task_environment_id = (SELECT task_environment_id FROM task_sessions WHERE id = ?)
AND repository_id = ?
AND deleted_at IS NULL
AND status = 'active'
`), branch, now, sessionID, repositoryID); err != nil {
return err
}
return tx.Commit()
}
// UpdateTaskSessionWorktreeBranchByWorktree updates exactly one worktree row.
// This is the repository-scoped variant needed when a task attaches multiple
// branches from the same repository.
func (r *Repository) UpdateTaskSessionWorktreeBranchByWorktree(ctx context.Context, sessionID, worktreeID, branch string) error {
now := time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return err
}
defer func() { _ = tx.Rollback() }()
if err := r.ensureTaskSessionEnvironmentAvailableTx(ctx, tx, sessionID, ""); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_environment_repos
SET worktree_branch = ?, updated_at = ?
WHERE task_environment_id = (SELECT task_environment_id FROM task_sessions WHERE id = ?)
AND worktree_id = ?
AND deleted_at IS NULL
AND status = 'active'
`), branch, now, sessionID, worktreeID); err != nil {
return err
}
return tx.Commit()
}
// ListWorktreesBySessionIDs returns the active environment-repository rows for
// the given session IDs, grouped by session ID. This eliminates N+1 queries
// when loading worktrees for multiple sessions. Chunks input above
// sqliteMaxHostParams (500) because callers like loadWorktreesBatch — invoked
// from BatchGetSessionsByTaskIDs — can pass
// `chunk_size_tasks × avg_sessions_per_task` IDs, which crosses SQLite's
// SQLITE_MAX_VARIABLE_NUMBER (999 on older builds) at modest task volumes.
func (r *Repository) ListWorktreesBySessionIDs(ctx context.Context, sessionIDs []string) (map[string][]*models.TaskEnvironmentRepo, error) {
result := make(map[string][]*models.TaskEnvironmentRepo, len(sessionIDs))
if len(sessionIDs) == 0 {
return result, nil
}
for _, chunk := range chunkIDs(sessionIDs, sqliteMaxHostParams) {
if err := r.appendWorktreesForSessionChunk(ctx, chunk, result); err != nil {
return nil, err
}
}
return result, nil
}
// appendWorktreesForSessionChunk runs the worktree SELECT for one chunk of
// session IDs and merges the rows into result. Extracted from
// ListWorktreesBySessionIDs so the public method stays inside the funlen cap
// after the chunking loop was added.
func (r *Repository) appendWorktreesForSessionChunk(
ctx context.Context,
sessionIDs []string,
result map[string][]*models.TaskEnvironmentRepo,
) error {
placeholders, args := buildInPlaceholders(sessionIDs)
query := `SELECT ts.id AS session_id, ` + envRepoSelectCols + `
FROM task_environment_repos ter
INNER JOIN task_sessions ts ON ts.task_environment_id = ter.task_environment_id
WHERE ts.id IN (` + placeholders + `)
AND ter.deleted_at IS NULL
AND ter.status = 'active'
ORDER BY ter.position ASC, ter.created_at ASC`
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return err
}
defer func() { _ = rows.Close() }()
for rows.Next() {
var sessionID string
var mergedAt, deletedAt sql.NullTime
row := &models.TaskEnvironmentRepo{}
if err := rows.Scan(&sessionID, &row.ID, &row.TaskEnvironmentID, &row.RepositoryID,
&row.BranchSlug, &row.WorktreeID, &row.WorktreePath, &row.WorktreeBranch,
&row.Position, &row.ErrorMessage, &row.Status, &row.CreatedAt, &row.UpdatedAt,
&mergedAt, &deletedAt); err != nil {
return err
}
if mergedAt.Valid {
t := mergedAt.Time
row.MergedAt = &t
}
if deletedAt.Valid {
t := deletedAt.Time
row.DeletedAt = &t
}
result[sessionID] = append(result[sessionID], row)
}
return rows.Err()
}
// GetPrimarySessionByTaskID retrieves the primary session for a task.
// Returns ErrNoPrimarySession (wrapped) when the task has no primary session
// row; callers should use errors.Is to distinguish this from real DB errors.
func (r *Repository) GetPrimarySessionByTaskID(ctx context.Context, taskID string) (*models.TaskSession, error) {
row := r.ro.QueryRowContext(ctx, r.ro.Rebind(
`SELECT `+taskSessionSelectCols+` `+taskSessionFromClause+` WHERE ts.task_id = ? AND ts.is_primary = 1 LIMIT 1`,
), taskID)
session, err := r.scanTaskSession(ctx, row, "task_sessions: no primary session row")
if errors.Is(err, models.ErrTaskSessionNotFound) {
return nil, fmt.Errorf("%w: %s", ErrNoPrimarySession, taskID)
}
return session, err
}
// GetPrimarySessionIDsByTaskIDs returns a map of task ID to primary session ID for the given task IDs.
// Tasks without a primary session are not included in the result.
func (r *Repository) GetPrimarySessionIDsByTaskIDs(ctx context.Context, taskIDs []string) (map[string]string, error) {
if len(taskIDs) == 0 {
return make(map[string]string), nil
}
// Build placeholders for IN clause
placeholders := make([]string, len(taskIDs))
args := make([]interface{}, len(taskIDs))
for i, id := range taskIDs {
placeholders[i] = "?"
args[i] = id
}
query := fmt.Sprintf(`
SELECT id, task_id FROM task_sessions
WHERE task_id IN (%s) AND is_primary = 1
`, strings.Join(placeholders, ","))
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
result := make(map[string]string)
for rows.Next() {
var sessionID, taskID string
if err := rows.Scan(&sessionID, &taskID); err != nil {
return nil, err
}
result[taskID] = sessionID
}
return result, rows.Err()
}
// GetSessionCountsByTaskIDs returns a map of task ID to session count for the given task IDs.
func (r *Repository) GetSessionCountsByTaskIDs(ctx context.Context, taskIDs []string) (map[string]int, error) {
if len(taskIDs) == 0 {
return make(map[string]int), nil
}
// Build placeholders for IN clause
placeholders := make([]string, len(taskIDs))
args := make([]interface{}, len(taskIDs))
for i, id := range taskIDs {
placeholders[i] = "?"
args[i] = id
}
query := fmt.Sprintf(`
SELECT task_id, COUNT(*) as count FROM task_sessions
WHERE task_id IN (%s)
GROUP BY task_id
`, strings.Join(placeholders, ","))
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
result := make(map[string]int)
for rows.Next() {
var taskID string
var count int
if err := rows.Scan(&taskID, &count); err != nil {
return nil, err
}
result[taskID] = count
}
return result, rows.Err()
}
// GetPrimarySessionInfoByTaskIDs returns a map of task ID to primary session for the given task IDs.
// Returns review_status, executor info, agent profile snapshot, and repository snapshot.
func (r *Repository) GetPrimarySessionInfoByTaskIDs(ctx context.Context, taskIDs []string) (map[string]*models.TaskSession, error) {
if len(taskIDs) == 0 {
return make(map[string]*models.TaskSession), nil
}
// Build placeholders for IN clause
placeholders := make([]string, len(taskIDs))
args := make([]interface{}, len(taskIDs))
for i, id := range taskIDs {
placeholders[i] = "?"
args[i] = id
}
query := fmt.Sprintf(`
SELECT ts.id, ts.task_id, ts.review_status, ts.executor_id, ts.executor_profile_id, ts.state,
ts.agent_profile_snapshot, ts.repository_snapshot,
e.type, e.name
FROM task_sessions ts
LEFT JOIN executors e ON e.id = ts.executor_id
WHERE ts.task_id IN (%s) AND ts.is_primary = 1
`, strings.Join(placeholders, ","))
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
result := make(map[string]*models.TaskSession)
for rows.Next() {
var sessionID string
var taskID string
var reviewStatus sql.NullString
var executorID sql.NullString
var executorProfileID sql.NullString
var sessionState sql.NullString
var agentProfileSnapshotJSON sql.NullString
var repositorySnapshotJSON sql.NullString
var executorType sql.NullString
var executorName sql.NullString
if err := rows.Scan(&sessionID, &taskID, &reviewStatus, &executorID, &executorProfileID, &sessionState, &agentProfileSnapshotJSON, &repositorySnapshotJSON, &executorType, &executorName); err != nil {
return nil, err
}
session := &models.TaskSession{
ID: sessionID,
TaskID: taskID,
}
if sessionState.Valid {
session.State = models.TaskSessionState(sessionState.String)
}
if reviewStatus.Valid {
session.ReviewStatus = models.ReviewStatus(reviewStatus.String)
}
if executorID.Valid {
session.ExecutorID = executorID.String
}
if executorProfileID.Valid {
session.ExecutorProfileID = executorProfileID.String
}
if executorType.Valid || executorName.Valid {
session.ExecutorSnapshot = make(map[string]interface{}, 2)
if executorType.Valid {
session.ExecutorSnapshot["executor_type"] = executorType.String
}
if executorName.Valid {
session.ExecutorSnapshot["executor_name"] = executorName.String
}
}
if agentProfileSnapshotJSON.Valid && agentProfileSnapshotJSON.String != "" {
if err := json.Unmarshal([]byte(agentProfileSnapshotJSON.String), &session.AgentProfileSnapshot); err != nil {
return nil, fmt.Errorf("failed to unmarshal agent profile snapshot for task %s: %w", taskID, err)
}
}
if repositorySnapshotJSON.Valid && repositorySnapshotJSON.String != "" {
if err := json.Unmarshal([]byte(repositorySnapshotJSON.String), &session.RepositorySnapshot); err != nil {
return nil, fmt.Errorf("failed to unmarshal repository snapshot for task %s: %w", taskID, err)
}
}
result[taskID] = session
}
return result, rows.Err()
}
// SetSessionPrimary marks a session as primary and clears the primary flag
// on every other session for the same task, atomically. Both writes go
// through a single transaction so a concurrent caller can't observe (or
// write) a half-applied state — e.g. two sessions racing on is_primary=1,
// or a reader seeing zero primary sessions mid-swap.
//
// On SQLite the writer pool is a single connection (see
// internal/db.NewSQLiteDB), so only one transaction can hold it at a time —
// the transaction alone fully serializes concurrent callers. On Postgres,
// separate connections could otherwise run two of these transactions truly
// concurrently, each seeing zero primary sessions and both promoting; to
// close that window we take an exclusive row lock on the owning task
// (`SELECT ... FOR UPDATE`) before touching its sessions, so a second
// concurrent promotion for the same task blocks until the first commits.
func (r *Repository) SetSessionPrimary(ctx context.Context, sessionID string) error {
_, err := r.setSessionPrimary(ctx, sessionID, false, nil)
return err
}
// SetSessionPrimaryIfNonterminal marks a session primary only while it remains
// nonterminal. It is used by workflow profile switching so a completed agent
// cannot be promoted from a stale lookup and have its ACP conversation resumed.
func (r *Repository) SetSessionPrimaryIfNonterminal(ctx context.Context, sessionID string) (bool, error) {
return r.setSessionPrimary(ctx, sessionID, true, nil)
}
// SetSessionPrimaryWithWorkflowSessionRouteIfNonterminal promotes a route
// destination and commits its prepared route in the same database transaction.
// The operation and destination checks fence a delayed workflow entry from
// promoting or committing over a newer route.
func (r *Repository) SetSessionPrimaryWithWorkflowSessionRouteIfNonterminal(
ctx context.Context,
sessionID string,
route models.WorkflowSessionRoute,
) (bool, error) {
return r.setSessionPrimary(ctx, sessionID, true, &route)
}
//nolint:cyclop // Primary promotion preserves one transaction boundary for both legacy and route-aware callers.
func (r *Repository) setSessionPrimary(
ctx context.Context,
sessionID string,
requireNonterminal bool,
route *models.WorkflowSessionRoute,
) (bool, error) {
now := time.Now().UTC()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return false, err
}
defer func() { _ = tx.Rollback() }()
// First, get the task_id for this session. Do not lock the target row here:
// every primary promotion must take the owning task lock first so concurrent
// promotions keep one lock order.
var taskID string
query := `SELECT task_id FROM task_sessions WHERE id = ?`
err = tx.QueryRowContext(ctx, r.db.Rebind(query), sessionID).Scan(&taskID)
if err == sql.ErrNoRows {
return primarySessionNotPromoted(sessionID, requireNonterminal)
}
if err != nil {
return false, err
}
// Serialize concurrent promotions for the same task across Postgres
// connections. SQLite has no FOR UPDATE / row-level locking and doesn't
// need it — the single-connection writer pool already serializes here.
if dialect.IsPostgres(r.db.DriverName()) {
var lockedTaskID string
err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT id FROM tasks WHERE id = ? FOR UPDATE`), taskID).Scan(&lockedTaskID)
if err != nil && err != sql.ErrNoRows {
return false, err
}
}
// Once the task lock is held, lock and validate the target row before
// promoting it. This serializes the nonterminal check with a concurrent
// state transition without reversing the task -> session lock order above.
if requireNonterminal {
valid, err := r.lockNonterminalPrimarySession(ctx, tx, sessionID)
if err != nil {
return false, err
}
if !valid {
return false, nil
}
}
if route != nil {
if err := r.commitWorkflowSessionRouteTx(ctx, tx, taskID, sessionID, route, now); err != nil {
return false, err
}
}
// Clear primary flag on all sessions for this task
_, err = tx.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions SET is_primary = 0, updated_at = ? WHERE task_id = ?
`), now, taskID)
if err != nil {
return false, err
}
// Set primary flag on the specified session
promoteQuery := `UPDATE task_sessions SET is_primary = 1, updated_at = ? WHERE id = ?`
if requireNonterminal {
promoteQuery += ` AND state IN ('CREATED', 'STARTING', 'RUNNING', 'IDLE', 'WAITING_FOR_INPUT')`
}
result, err := tx.ExecContext(ctx, r.db.Rebind(promoteQuery), now, sessionID)
if err != nil {
return false, err
}
rows, _ := result.RowsAffected()
if rows == 0 {
return primarySessionNotPromoted(sessionID, requireNonterminal)
}
if err := tx.Commit(); err != nil {
return false, err
}
return true, nil
}