fix(queue): auto-merge compatible messages at full queues (#2656)

This commit is contained in:
Clément Fiere
2026-08-15 13:51:43 +02:00
committed by GitHub
parent dfa100df13
commit 488abb4265
17 changed files with 2365 additions and 23 deletions
@@ -237,6 +237,12 @@ func (h *QueueHandlers) wsQueueMessage(ctx context.Context, msg *ws.Message) (*w
fieldMax: status.Max,
})
}
if errors.Is(err, messagequeue.ErrTaskInactive) {
// The task was archived or deleted between the caller's
// authorization and the queue admission; do not queue a message
// that would be orphaned behind the task's purge.
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "Task is no longer active", nil)
}
if errors.Is(err, errQueuedAttachmentRollback) {
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Failed to roll back queued attachment", nil)
}
@@ -91,6 +91,51 @@ func TestWsQueueMessage_AutoMergeBlocksLaterAdmissionUntilClaimFinalizes(t *test
require.Equal(t, "first\n\nsecond\n\nthird", status.Entries[0].Content)
}
func TestWsQueueMessage_AutoMergeFoldsAdmissionAtFullQueue(t *testing.T) {
handlers, service := setupQueueHandlers(t)
// Fill the queue with auto-merge disabled so the session actually reaches
// capacity, then enable it: the next compatible message must fold into the
// tail instead of surfacing "queue full".
service.SetMaxPerSession(1)
service.SetAutoMergeEnabled(false)
_, err := service.QueueMessage(context.Background(), "session", "task", "first", "", messagequeue.QueuedByUser, false, nil)
require.NoError(t, err)
service.SetAutoMergeEnabled(true)
response, err := handlers.wsQueueMessage(context.Background(), createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{
"session_id": "session",
"task_id": "task",
"content": "second",
}))
require.NoError(t, err)
require.Equal(t, ws.MessageTypeResponse, response.Type)
var queued messagequeue.QueuedMessage
require.NoError(t, json.Unmarshal(response.Payload, &queued))
require.Equal(t, "first\n\nsecond", queued.Content)
status := service.GetStatus(context.Background(), "session")
require.Equal(t, 1, status.Count)
require.Equal(t, "first\n\nsecond", status.Entries[0].Content)
}
func TestWsQueueMessage_AutoMergeAtFullQueueRejectsIncompatible(t *testing.T) {
handlers, service := setupQueueHandlers(t)
service.SetMaxPerSession(1)
service.SetAutoMergeEnabled(false)
_, err := service.QueueMessage(context.Background(), "session", "task", "first", "", messagequeue.QueuedByUser, false, nil)
require.NoError(t, err)
service.SetAutoMergeEnabled(true)
response, err := handlers.wsQueueMessage(context.Background(), createTestMessage(t, ws.ActionMessageQueueAdd, map[string]interface{}{
"session_id": "session",
"task_id": "task",
"content": "second",
"model": "other-model",
}))
require.NoError(t, err)
require.Equal(t, ws.MessageTypeError, response.Type)
require.Equal(t, messagequeue.QueueFullErrorCode, parseError(t, response).Code)
}
type blockingQueueAttachmentClaimer struct {
started chan struct{}
release chan struct{}
@@ -122,6 +122,14 @@ type Repository interface {
// successful skip and returns false without mutating either row.
AutoMergeIntoAbove(ctx context.Context, sessionID, sourceID string) (*QueuedMessage, bool, error)
// AutoMergeCandidateIntoAbove folds a not-yet-admitted candidate message
// into the session's tail entry when compatible, without inserting it. Used
// by admission at full capacity: the fold itself is the admission, so a
// message that would merge anyway is accepted instead of rejected with
// ErrQueueFull. A missing tail or incompatible pair is a successful skip
// and returns false without mutating any row.
AutoMergeCandidateIntoAbove(ctx context.Context, candidate *QueuedMessage) (*QueuedMessage, bool, error)
// ReorderEntries atomically rewrites the FIFO positions of the session's
// pending entries to match orderedIDs. The submitted list must be an exact
// permutation of the currently visible pending (not reserved-in-flight)
@@ -17,6 +17,10 @@ type automaticMergeRepository interface {
AutoMergeIntoAbove(context.Context, string, string) (*QueuedMessage, bool, error)
}
type automaticMergeCandidateRepository interface {
AutoMergeCandidateIntoAbove(context.Context, *QueuedMessage) (*QueuedMessage, bool, error)
}
type autoMergeRepositoryFactory struct {
name string
new func(*testing.T) Repository
@@ -236,6 +240,106 @@ func TestRepository_AutoMergeIntoAbove_EmptyContentHasNoExtraSeparator(t *testin
}
}
func TestRepository_AutoMergeCandidateIntoAbove_FoldsIntoTail(t *testing.T) {
for _, tt := range autoMergeRepositoryFactories() {
t.Run(tt.name, func(t *testing.T) {
repo := tt.new(t)
autoRepo, ok := repo.(automaticMergeCandidateRepository)
if !ok {
t.Fatalf("repository %T does not implement candidate tail merge", repo)
}
insertAutoMergeEntry(t, repo, defaultAutoMergeEntry("candidate", "first"))
tail := insertAutoMergeEntry(t, repo, defaultAutoMergeEntry("candidate", "second"))
candidate := defaultAutoMergeEntry("candidate", "third")
merged, didMerge, err := autoRepo.AutoMergeCandidateIntoAbove(context.Background(), &candidate)
if err != nil || !didMerge || merged == nil {
t.Fatalf("candidate merge: result=%+v merged=%v err=%v", merged, didMerge, err)
}
if merged.ID != tail.ID {
t.Fatalf("survivor = %s, want tail %s", merged.ID, tail.ID)
}
if merged.Content != "second\n\nthird" {
t.Errorf("content = %q, want %q", merged.Content, "second\n\nthird")
}
entries, err := repo.ListBySession(context.Background(), "candidate")
if err != nil {
t.Fatalf("list: %v", err)
}
if len(entries) != 2 || entries[1].ID != tail.ID {
t.Fatalf("entries after candidate fold = %+v, want two rows with folded tail", entries)
}
})
}
}
func TestRepository_AutoMergeCandidateIntoAbove_IncompatibleSkips(t *testing.T) {
for _, tt := range autoMergeRepositoryFactories() {
t.Run(tt.name, func(t *testing.T) {
repo := tt.new(t)
autoRepo, ok := repo.(automaticMergeCandidateRepository)
if !ok {
t.Fatalf("repository %T does not implement candidate tail merge", repo)
}
insertAutoMergeEntry(t, repo, defaultAutoMergeEntry("candidate-skip", "first"))
candidate := defaultAutoMergeEntry("candidate-skip", "other-model")
candidate.Model = "other-model"
before := autoQueueSnapshot(t, repo, "candidate-skip")
merged, didMerge, err := autoRepo.AutoMergeCandidateIntoAbove(context.Background(), &candidate)
if err != nil || didMerge || merged != nil {
t.Fatalf("incompatible candidate: result=%+v merged=%v err=%v", merged, didMerge, err)
}
if after := autoQueueSnapshot(t, repo, "candidate-skip"); after != before {
t.Fatalf("queue mutated on skip\nbefore=%s\nafter=%s", before, after)
}
})
}
}
func TestRepository_AutoMergeCandidateIntoAbove_EmptyQueueSkips(t *testing.T) {
for _, tt := range autoMergeRepositoryFactories() {
t.Run(tt.name, func(t *testing.T) {
repo := tt.new(t)
autoRepo, ok := repo.(automaticMergeCandidateRepository)
if !ok {
t.Fatalf("repository %T does not implement candidate tail merge", repo)
}
candidate := defaultAutoMergeEntry("candidate-empty", "only")
merged, didMerge, err := autoRepo.AutoMergeCandidateIntoAbove(context.Background(), &candidate)
if err != nil || didMerge || merged != nil {
t.Fatalf("empty queue: result=%+v merged=%v err=%v", merged, didMerge, err)
}
})
}
}
func TestSQLiteRepository_AutoMergeCandidateIntoAbove_WriteFailureRollsBack(t *testing.T) {
repo := newTestSQLiteRepo(t)
db := repo.(*sqliteRepository).db
insertAutoMergeEntry(t, repo, defaultAutoMergeEntry("candidate-rollback", "first"))
insertAutoMergeEntry(t, repo, defaultAutoMergeEntry("candidate-rollback", "second"))
before := autoQueueSnapshot(t, repo, "candidate-rollback")
if _, err := db.Exec(`
CREATE TRIGGER fail_auto_merge_candidate_update
BEFORE UPDATE ON queued_messages
BEGIN
SELECT RAISE(ABORT, 'forced automatic candidate merge failure');
END;
`); err != nil {
t.Fatalf("create failure trigger: %v", err)
}
candidate := defaultAutoMergeEntry("candidate-rollback", "third")
merged, didMerge, err := repo.(automaticMergeCandidateRepository).AutoMergeCandidateIntoAbove(context.Background(), &candidate)
if err == nil || didMerge || merged != nil {
t.Fatalf("forced write failure: result=%+v merged=%v err=%v", merged, didMerge, err)
}
if after := autoQueueSnapshot(t, repo, "candidate-rollback"); after != before {
t.Fatalf("write failure was not atomic\nbefore=%s\nafter=%s", before, after)
}
}
func TestSQLiteRepository_AutoMergeIntoAbove_WriteFailureRollsBack(t *testing.T) {
repo := newTestSQLiteRepo(t)
testAutoMergeWriteRollback(t, repo, `
@@ -349,6 +453,20 @@ func autoMergeSkipCases() []autoMergeSkipCase {
target.QueuedBy, source.QueuedBy = QueuedByAgent, QueuedByAgent
target.Metadata, source.Metadata = senderMetadata("one"), senderMetadata("two")
}},
{name: "agent session mismatch", mutate: func(target, source *QueuedMessage) {
// Different sessions of the same sender task are different agents
// and execution contexts; auto-merge must keep their prompts
// separate even though manual merge would allow the fold.
target.QueuedBy, source.QueuedBy = QueuedByAgent, QueuedByAgent
target.Metadata = map[string]interface{}{
"sender_task_id": "sender-task", "sender_session_id": "session-one",
"sender_session_name": "reviewer-one", "sender_task_title": "Sender task",
}
source.Metadata = map[string]interface{}{
"sender_task_id": "sender-task", "sender_session_id": "session-two",
"sender_session_name": "reviewer-two", "sender_task_title": "Sender task",
}
}},
{name: "workflow source", mutate: func(target, source *QueuedMessage) {
target.QueuedBy, source.QueuedBy = QueuedByWorkflow, QueuedByWorkflow
}},
@@ -698,6 +698,33 @@ func (r *memoryRepository) AutoMergeIntoAbove(_ context.Context, sessionID, sour
return cloneQueuedMessage(target), true, nil
}
// AutoMergeCandidateIntoAbove folds a not-yet-admitted candidate into the
// session's tail entry when compatible, without inserting it. Incompatibility
// and missing tails are successful skips.
func (r *memoryRepository) AutoMergeCandidateIntoAbove(_ context.Context, candidate *QueuedMessage) (*QueuedMessage, bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
list := r.entries[candidate.SessionID]
var target *QueuedMessage
for _, message := range list {
if target == nil || message.Position > target.Position {
target = message
}
}
if target == nil {
return nil, false, nil
}
values, compatible := buildAutoMergedEntry(target, candidate)
if !compatible {
return nil, false, nil
}
target.Content = values.content
target.Attachments = values.attachments
target.Metadata = values.metadata
return cloneQueuedMessage(target), true, nil
}
// ReorderEntries rewrites the session's visible pending order to match
// orderedIDs, mirroring the sqlite repository's semantics: reserved in-flight
// rows keep their place in the sequence, visible rows appear in the submitted
@@ -2,10 +2,16 @@ package messagequeue
import (
"context"
"database/sql"
"errors"
"fmt"
"strings"
"testing"
"time"
"github.com/google/uuid"
_ "github.com/jackc/pgx/v5/stdlib"
"github.com/jmoiron/sqlx"
"github.com/kandev/kandev/internal/entityrefs"
"github.com/kandev/kandev/internal/testutil"
@@ -30,6 +36,695 @@ func newTestPostgresRepo(t *testing.T) Repository {
// TestPostgresRepository_MergeDrainOrdering_DrainWins asserts the drain-wins
// ordering of the merge/drain race: after the source drains, a merge reports
// ErrEntryNotFound and the target is untouched, so no content is lost.
// pgBackendPID returns the backend PID of the repository's single connection,
// used to scope pg_locks barrier polls to this test's worker: transactionid
// waits are server-global, so an unrelated waiter from a concurrently running
// package could otherwise false-trigger the barrier.
func pgBackendPID(t *testing.T, db *sqlx.DB) int {
t.Helper()
var pid int
if err := db.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil {
t.Fatalf("pg_backend_pid: %v", err)
}
return pid
}
// waitForWaitingLocks polls pg_locks for the expected number of NOT-granted
// transactionid waits owned by the given backend PID, scoping the barrier to
// this test's worker. The poller must be the held transaction's connection
// (each test handle has a single connection busy inside its lock tx).
func waitForWaitingLocks(t *testing.T, poller interface {
QueryRowContext(context.Context, string, ...any) *sql.Row
}, backendPID, want int, what string) {
t.Helper()
deadline := time.Now().Add(10 * time.Second)
for {
var waiting int
if err := poller.QueryRowContext(context.Background(), `
SELECT count(*) FROM pg_locks
WHERE NOT granted AND locktype = 'transactionid' AND pid = $1
`, backendPID).Scan(&waiting); err != nil {
t.Fatalf("query pg_locks: %v", err)
}
if waiting >= want {
return
}
if time.Now().After(deadline) {
t.Fatalf("%s never reached %d waiting lock(s)", what, want)
}
time.Sleep(10 * time.Millisecond)
}
}
// newTestPostgresRepoPair opens one isolated Postgres schema (skipping unless
// KANDEV_TEST_POSTGRES_DSN is set) and constructs two repository instances
// over it, plus an extra single-connection DB handle on the same schema. Two
// instances simulate two backend processes sharing one queue database: their
// per-session admission locks are process-local, so cross-instance races are
// guarded only by the repository's affected-row checks. The extra handle gives
// a test a third connection — each worker pool is capped at one connection, so
// a second concurrent operation on the same worker needs its own.
func newTestPostgresRepoPair(t *testing.T) (Repository, Repository, *sqlx.DB) {
t.Helper()
dsn := testutil.PostgresDSNFromEnv(t)
schema := "kandev_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
// Create the shared schema exactly once; each handle below only selects it
// via search_path so both repository instances see the same tables.
setupDB, err := sqlx.Open("pgx", dsn)
if err != nil {
t.Fatalf("open postgres: %v", err)
}
if _, err := setupDB.Exec("CREATE SCHEMA " + schema); err != nil {
_ = setupDB.Close()
t.Fatalf("create postgres schema %s: %v", schema, err)
}
t.Cleanup(func() {
_, _ = setupDB.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
_ = setupDB.Close()
})
open := func() *sqlx.DB {
db, err := sqlx.Open("pgx", dsn)
if err != nil {
t.Fatalf("open postgres: %v", err)
}
db.SetMaxOpenConns(1)
t.Cleanup(func() { _ = db.Close() })
if _, err := db.Exec("SET search_path TO " + schema); err != nil {
t.Fatalf("set postgres search_path %s: %v", schema, err)
}
return db
}
dbA, dbB, dbC := open(), open(), open()
repoA, err := NewSQLiteRepository(dbA, dbA)
if err != nil {
t.Fatalf("NewSQLiteRepository(postgres A): %v", err)
}
repoB, err := NewSQLiteRepository(dbB, dbB)
if err != nil {
t.Fatalf("NewSQLiteRepository(postgres B): %v", err)
}
return repoA, repoB, dbC
}
// TestPostgresRepository_AutoMergeCandidateIntoAbove_ConcurrentAdmissionCannotLoseContent
// proves the compare-and-swap guard in AutoMergeCandidateIntoAbove: two
// repository instances (two backend processes) that both read the same tail
// must not both fold successfully, silently dropping one accepted message.
//
// The dangerous interleaving is forced deterministically with a row lock.
// Instance A holds a FOR UPDATE lock on the tail while instance B runs the
// fold: B reads the original tail, then its CAS UPDATE blocks on A's lock.
// Only once the poll observes B waiting is B's read snapshot guaranteed to
// predate A's competing fold. A then commits its own fold, B's UPDATE
// re-evaluates its CAS condition against the new row version, affects zero
// rows, and reports a skip — the winner's content survives intact and the
// loser surfaces a failed admission instead of overwriting it.
func TestPostgresRepository_AutoMergeCandidateIntoAbove_ConcurrentAdmissionCannotLoseContent(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
tail := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("race", "first"))
dbA := repoA.(*sqliteRepository).db
// Instance A holds a row lock on the tail, simulating a competing fold
// that is about to commit.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `SELECT id FROM queued_messages WHERE id = $1 FOR UPDATE`, tail.ID); err != nil {
t.Fatalf("lock tail row: %v", err)
}
// Instance B starts a fold while A holds the lock: B reads the original
// tail, then its CAS UPDATE blocks on A's row lock. The PID is captured
// before the goroutine (single-connection pool).
foldPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
foldDone := make(chan error, 1)
candidate := defaultAutoMergeEntry("race", "second")
go func() {
merged, didMerge, err := repoB.(automaticMergeCandidateRepository).AutoMergeCandidateIntoAbove(ctx, &candidate)
if err != nil {
foldDone <- err
return
}
if didMerge || merged != nil {
foldDone <- fmt.Errorf("concurrent fold succeeded (didMerge=%v merged=%+v), want stale-write skip", didMerge, merged)
return
}
foldDone <- nil
}()
// Wait until B's UPDATE is queued on the row lock. Only then is B's read
// snapshot guaranteed to predate A's competing fold below.
waitForWaitingLocks(t, lockTx, foldPID, 1, "instance B fold on the tail row lock")
// Instance A's competing fold commits first.
if _, err := lockTx.ExecContext(ctx, `
UPDATE queued_messages SET content = $1 WHERE id = $2
`, "first\n\ncompeting", tail.ID); err != nil {
t.Fatalf("competing fold update: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit competing fold: %v", err)
}
if err := <-foldDone; err != nil {
t.Fatalf("instance B fold: %v", err)
}
// The winner's content is intact and the loser's candidate is absent:
// no silent content loss, and the loser's admission failed loudly.
entries, err := repoA.ListBySession(ctx, "race")
if err != nil {
t.Fatalf("list queue: %v", err)
}
if len(entries) != 1 || entries[0].ID != tail.ID {
t.Fatalf("entries after concurrent admission = %+v, want one surviving tail", entries)
}
if entries[0].Content != "first\n\ncompeting" {
t.Fatalf("tail content = %q, want %q (loser must not overwrite the winner)", entries[0].Content, "first\n\ncompeting")
}
}
// TestPostgresRepository_AutoMergeCandidateIntoAbove_RaceWithReorderSkips
// proves the position component of the CAS guard: a concurrent reorder that
// only changes positions must invalidate an in-flight fold, otherwise the fold
// lands on a row that is no longer the tail (e.g. it became the head) and
// drains in the wrong order. The interleaving is forced deterministically with
// a row lock, mirroring the concurrent-admission test.
func TestPostgresRepository_AutoMergeCandidateIntoAbove_RaceWithReorderSkips(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
first := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("race-reorder", "first"))
second := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("race-reorder", "second"))
dbA := repoA.(*sqliteRepository).db
// Instance A holds a row lock on the tail (the "second" entry), simulating
// a concurrent reorder that is about to commit.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `SELECT id FROM queued_messages WHERE id = $1 FOR UPDATE`, second.ID); err != nil {
t.Fatalf("lock tail row: %v", err)
}
// Instance B starts a fold while A holds the lock: B scans the tail
// ("second", position 2), then its CAS UPDATE blocks on A's row lock.
foldPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
foldDone := make(chan error, 1)
candidate := defaultAutoMergeEntry("race-reorder", "third")
go func() {
merged, didMerge, err := repoB.(automaticMergeCandidateRepository).AutoMergeCandidateIntoAbove(ctx, &candidate)
if err != nil {
foldDone <- err
return
}
if didMerge || merged != nil {
foldDone <- fmt.Errorf("fold survived a concurrent reorder (didMerge=%v merged=%+v), want stale-tail skip", didMerge, merged)
return
}
foldDone <- nil
}()
// Wait until B's UPDATE is queued on the row lock, guaranteeing B's scan
// predates the reorder below.
waitForWaitingLocks(t, lockTx, foldPID, 1, "instance B fold on the tail row lock")
// The concurrent reorder swaps only positions: "second" becomes the head.
if _, err := lockTx.ExecContext(ctx, `UPDATE queued_messages SET position = 1 WHERE id = $1`, second.ID); err != nil {
t.Fatalf("reorder second: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `UPDATE queued_messages SET position = 2 WHERE id = $1`, first.ID); err != nil {
t.Fatalf("reorder first: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit reorder: %v", err)
}
if err := <-foldDone; err != nil {
t.Fatalf("instance B fold: %v", err)
}
// The reordered queue is intact and the fold content is absent: the tail
// is no longer the tail, so the fold must not have landed on it.
entries, err := repoA.ListBySession(ctx, "race-reorder")
if err != nil {
t.Fatalf("list queue: %v", err)
}
if len(entries) != 2 || entries[0].ID != second.ID || entries[1].ID != first.ID {
t.Fatalf("entries after reorder = %+v, want second first", entries)
}
if entries[0].Content != "second" {
t.Fatalf("head content = %q, want %q (fold must not land on a non-tail row)", entries[0].Content, "second")
}
}
// TestPostgresRepository_AutoMergeCandidateIntoAbove_RaceWithInsertAbove proves
// the per-session cross-process lock: a fold must never land on a row that
// stopped being the tail. The interleaving is forced deterministically — the
// competing backend holds the queue_session_locks row, the fold blocks on it
// before scanning, and only after the competing drain+insert commits does the
// fold run — so the fold always scans the current tail, never a stale one.
func TestPostgresRepository_AutoMergeCandidateIntoAbove_RaceWithInsertAbove(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
first := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("race-tail", "first"))
second := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("race-tail", "second"))
dbA := repoA.(*sqliteRepository).db
// Instance A holds the per-session cross-process lock, simulating a
// concurrent drain+insert that is about to commit.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('race-tail')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'race-tail' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
// Instance B's fold blocks on the session lock before it can scan.
foldPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
foldDone := make(chan error, 1)
candidate := defaultAutoMergeEntry("race-tail", "third")
go func() {
merged, didMerge, err := repoB.(automaticMergeCandidateRepository).AutoMergeCandidateIntoAbove(ctx, &candidate)
if err != nil {
foldDone <- err
return
}
if !didMerge || merged == nil {
foldDone <- fmt.Errorf("fold skipped (didMerge=%v merged=%+v), want fold onto the new tail", didMerge, merged)
return
}
if merged.ID == second.ID {
foldDone <- errors.New("fold landed on the stale tail")
return
}
foldDone <- nil
}()
// Wait until B's fold is queued on the session lock row.
waitForWaitingLocks(t, lockTx, foldPID, 1, "instance B fold on the session lock")
// The competing backend drains the head and inserts a new row above the
// old tail, then commits before the fold can run.
if _, err := lockTx.ExecContext(ctx, `DELETE FROM queued_messages WHERE id = $1`, first.ID); err != nil {
t.Fatalf("drain head: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queued_messages
(id, session_id, task_id, position, content, model, plan_mode,
attachments_json, metadata_json, queued_at, queued_by)
VALUES ($1, 'race-tail', 'task', 3, 'drained-in', 'model', 0, '[]', '{}', now(), 'user')
`, uuid.NewString()); err != nil {
t.Fatalf("insert above tail: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit competing drain+insert: %v", err)
}
if err := <-foldDone; err != nil {
t.Fatalf("instance B fold: %v", err)
}
// The fold landed on the current tail (the newly inserted row), never on
// the stale "second" row, and the queue order is intact.
entries, err := repoA.ListBySession(ctx, "race-tail")
if err != nil {
t.Fatalf("list queue: %v", err)
}
if len(entries) != 2 {
t.Fatalf("entries = %+v, want two rows", entries)
}
if entries[0].ID != second.ID || entries[0].Content != "second" {
t.Fatalf("stale tail = %s %q, want untouched second entry", entries[0].ID, entries[0].Content)
}
if entries[1].Content != "drained-in\n\nthird" {
t.Fatalf("new tail content = %q, want candidate folded into the current tail", entries[1].Content)
}
}
// TestPostgresRepository_UpdateContentAndMetadata_RaceWithMerge proves the
// per-session lock also covers content edits: a merge's scan-to-write window
// must not silently overwrite a concurrent UpdateContentAndMetadata. The
// interleaving is forced deterministically — the competing backend holds the
// session lock (simulating a merge mid-transaction), the edit blocks on it,
// the merge commits, and only then does the edit apply on top.
func TestPostgresRepository_UpdateContentAndMetadata_RaceWithMerge(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
target := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("merge-update", "first"))
source := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("merge-update", "second"))
dbA := repoA.(*sqliteRepository).db
// Instance A holds the per-session lock, simulating a merge that has
// scanned its target/source and is about to write.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('merge-update')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'merge-update' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
// Instance B's content edit blocks on the session lock until the merge
// commits.
editPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
updateDone := make(chan error, 1)
go func() {
err := repoB.UpdateContentAndMetadata(ctx, "merge-update", target.ID, "edited", nil,
map[string]interface{}{"edited": true}, "user")
updateDone <- err
}()
// Wait until B's edit is queued on the session lock row.
waitForWaitingLocks(t, lockTx, editPID, 1, "instance B edit on the session lock")
// The merge commits its stale-snapshot write (fold target + delete
// source) while the edit is still waiting.
if _, err := lockTx.ExecContext(ctx, `
UPDATE queued_messages SET content = $1 WHERE id = $2 AND session_id = 'merge-update'
`, "first\n\nsecond", target.ID); err != nil {
t.Fatalf("merge target write: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
DELETE FROM queued_messages WHERE id = $1 AND session_id = 'merge-update'
`, source.ID); err != nil {
t.Fatalf("merge source delete: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit merge: %v", err)
}
if err := <-updateDone; err != nil {
t.Fatalf("instance B edit: %v", err)
}
// The edit applied after the merge and survives: nothing was silently
// overwritten by the merge's stale write.
entries, err := repoA.ListBySession(ctx, "merge-update")
if err != nil {
t.Fatalf("list queue: %v", err)
}
if len(entries) != 1 || entries[0].ID != target.ID {
t.Fatalf("entries = %+v, want the merged target only", entries)
}
if entries[0].Content != "edited" {
t.Fatalf("target content = %q, want %q (concurrent edit must survive the merge)", entries[0].Content, "edited")
}
if edited := entries[0].Metadata["edited"]; edited != true {
t.Fatalf("target metadata = %+v, want the edit's metadata applied", entries[0].Metadata)
}
}
// TestPostgresRepository_TransferSession_RaceWithSourceInsert proves the real
// TransferSession holds the SOURCE session lock too: a source-side admission
// must serialize with the transfer instead of interleaving with its READ
// COMMITTED move. The test drives the actual TransferSession method and forces
// the interleaving deterministically:
//
// - session ids are chosen so the transfer's stable-sorted acquisition locks
// the source first and the destination second ("a-source" < "b-dest");
// - the competing backend holds the DESTINATION lock, so TransferSession
// acquires the source lock and then blocks on the destination;
// - a source-side admission is started only after the transfer is blocked,
// and must itself block on the source lock;
// - releasing the destination lets the transfer commit, and only then does
// the admission land on the (now empty) source session.
//
// If TransferSession regressed to destination-only locking, the admission
// would never block on the source and the second barrier would time out.
func TestPostgresRepository_TransferSession_RaceWithSourceInsert(t *testing.T) {
repoA, repoB, extraDB := newTestPostgresRepoPair(t)
ctx := context.Background()
const (
oldSession = "a-source"
newSession = "b-dest"
)
insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry(oldSession, "first"))
dbA := repoA.(*sqliteRepository).db
// The source-side admission runs on a THIRD connection: repoB's pool has
// one connection, already held by the blocked transfer.
repoC, err := NewSQLiteRepository(extraDB, extraDB)
if err != nil {
t.Fatalf("NewSQLiteRepository(extra): %v", err)
}
// Capture worker PIDs BEFORE starting the operations: the pools are
// single-connection, so a post-start lookup could block on the busy
// worker or observe a different connection.
transferPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
insertPID := pgBackendPID(t, extraDB)
// Instance A holds the DESTINATION lock. TransferSession locks the source
// first (stable sorted order), so it will block here on the destination.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('b-dest')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'b-dest' FOR UPDATE
`); err != nil {
t.Fatalf("lock destination session: %v", err)
}
// The real transfer starts: it acquires the source lock, then blocks on
// the destination.
transferDone := make(chan error, 1)
go func() {
transferDone <- repoB.TransferSession(ctx, oldSession, newSession)
}()
// Transfer blocked on the destination (holding the source lock).
waitForWaitingLocks(t, lockTx, transferPID, 1, "transfer on the destination")
// The source-side admission (own connection) must now block on the source
// lock held by the transfer.
insertDone := make(chan error, 1)
admission := defaultAutoMergeEntry(oldSession, "inserted")
go func() {
insertDone <- repoC.Insert(ctx, &admission, 0)
}()
waitForWaitingLocks(t, lockTx, insertPID, 1, "source admission on the source lock")
// Release the destination: the transfer commits its move before the
// admission proceeds.
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-transferDone; err != nil {
t.Fatalf("transfer: %v", err)
}
if err := <-insertDone; err != nil {
t.Fatalf("source admission: %v", err)
}
// The transfer moved every pre-existing row; the admission that waited
// landed after it, on the source session it targeted. Nothing is lost or
// stranded mid-move.
oldEntries, err := repoA.ListBySession(ctx, oldSession)
if err != nil {
t.Fatalf("list old: %v", err)
}
if len(oldEntries) != 1 || oldEntries[0].Content != "inserted" {
t.Fatalf("old entries = %+v, want the post-transfer admission only", oldEntries)
}
newEntries, err := repoA.ListBySession(ctx, newSession)
if err != nil {
t.Fatalf("list new: %v", err)
}
if len(newEntries) != 1 || newEntries[0].Content != "first" {
t.Fatalf("new entries = %+v, want the transferred row", newEntries)
}
}
// TestPostgresRepository_AcknowledgeByID_RaceWithRestore proves the ack is
// cross-process serialized on the session lock: an autocommit DELETE could
// otherwise remove a row a concurrent backend just restored/reinserted,
// losing durable retry state. The competing backend holds the session lock,
// the ack blocks on it, and only after release does the ack delete the row.
func TestPostgresRepository_AcknowledgeByID_RaceWithSessionLock(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
entry := insertAutoMergeEntry(t, repoA, defaultAutoMergeEntry("ack", "first"))
dbA := repoA.(*sqliteRepository).db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('ack')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'ack' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
ackPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
ackDone := make(chan error, 1)
go func() {
ackDone <- repoB.AcknowledgeByID(ctx, "ack", entry.ID)
}()
waitForWaitingLocks(t, lockTx, ackPID, 1, "ack on the session lock")
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-ackDone; err != nil {
t.Fatalf("ack: %v", err)
}
entries, err := repoA.ListBySession(ctx, "ack")
if err != nil {
t.Fatalf("list after ack: %v", err)
}
if len(entries) != 0 {
t.Fatalf("entries after ack = %+v, want removed", entries)
}
}
// TestPostgresRepository_SetPendingMove_RaceWithTransferLock proves pending
// moves serialize on the session lock so a transfer cannot orphan one.
func TestPostgresRepository_SetPendingMove_RaceWithSessionLock(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
dbA := repoA.(*sqliteRepository).db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('pm')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'pm' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
setPID := pgBackendPID(t, repoB.(*sqliteRepository).db)
setDone := make(chan error, 1)
go func() {
setDone <- repoB.SetPendingMove(ctx, "pm", &PendingMove{
TaskID: "t", WorkflowID: "w", WorkflowStepID: "s", Position: 1, Actor: "test",
})
}()
waitForWaitingLocks(t, lockTx, setPID, 1, "set pending move on the session lock")
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-setDone; err != nil {
t.Fatalf("set pending move: %v", err)
}
move, err := repoA.GetPendingMove(ctx, "pm")
if err != nil || move == nil || move.TaskID != "t" {
t.Fatalf("pending move = %+v err=%v, want the set move", move, err)
}
}
// TestPostgresRepository_TakePendingMove_TwoTakersSerialized proves two
// backends cannot both return the same pending move: the second taker blocks
// on the session lock until the first commits, then sees no move.
func TestPostgresRepository_TakePendingMove_TwoTakersSerialized(t *testing.T) {
repoA, repoB, _ := newTestPostgresRepoPair(t)
ctx := context.Background()
if err := repoA.SetPendingMove(ctx, "pm-take", &PendingMove{
TaskID: "t", WorkflowID: "w", WorkflowStepID: "s", Position: 1, Actor: "test",
}); err != nil {
t.Fatalf("seed pending move: %v", err)
}
dbA := repoA.(*sqliteRepository).db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('pm-take')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'pm-take' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
takePID := pgBackendPID(t, repoB.(*sqliteRepository).db)
takeDone := make(chan error, 1)
go func() {
move, err := repoB.TakePendingMove(ctx, "pm-take")
if err != nil {
takeDone <- err
return
}
if move == nil || move.TaskID != "t" {
takeDone <- fmt.Errorf("take returned %+v, want the seeded move", move)
return
}
takeDone <- nil
}()
waitForWaitingLocks(t, lockTx, takePID, 1, "take pending move on the session lock")
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-takeDone; err != nil {
t.Fatalf("take pending move: %v", err)
}
if move, err := repoA.GetPendingMove(ctx, "pm-take"); err != nil || move != nil {
t.Fatalf("pending move after take = %+v err=%v, want gone", move, err)
}
}
func TestPostgresRepository_MergeDrainOrdering_DrainWins(t *testing.T) {
repo := newTestPostgresRepo(t)
ctx := context.Background()
@@ -6,12 +6,13 @@ import (
"encoding/json"
"errors"
"fmt"
"sort"
"sync"
"time"
"github.com/google/uuid"
"github.com/jmoiron/sqlx"
"github.com/kandev/kandev/internal/db"
internaldb "github.com/kandev/kandev/internal/db"
)
// sqliteRepository persists queued messages and pending moves.
@@ -21,6 +22,14 @@ type sqliteRepository struct {
mu sync.Mutex
sessionLocks map[string]*sync.Mutex
// tasksTablePresent is whether the owning tasks table exists, resolved at
// construction. The queue repository's isolated tests create only queue
// tables; production databases always have the task schema. It is checked
// OUTSIDE any transaction because a failed statement on PostgreSQL aborts
// the whole transaction — the guard must never issue its UPDATE against a
// missing table inside a tx.
tasksTablePresent bool
}
// NewSQLiteRepository creates a SQLite-backed Repository. The supplied writer
@@ -30,6 +39,17 @@ func NewSQLiteRepository(writer, reader *sqlx.DB) (Repository, error) {
if err := r.initSchema(); err != nil {
return nil, fmt.Errorf("messagequeue: init schema: %w", err)
}
var present bool
var err error
if writer.DriverName() == "pgx" {
err = writer.Get(&present, `SELECT to_regclass('tasks') IS NOT NULL`)
} else {
err = writer.Get(&present, `SELECT EXISTS (SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'tasks')`)
}
if err != nil {
return nil, fmt.Errorf("messagequeue: resolve tasks table presence: %w", err)
}
r.tasksTablePresent = present
return r, nil
}
@@ -39,6 +59,70 @@ func NewSQLiteRepository(writer, reader *sqlx.DB) (Repository, error) {
// affected-row checks remain the authoritative guard across processes (e.g.
// multiple backends sharing a Postgres queue); the lock makes in-process
// merge-wins/drain-wins ordering deterministic.
//
// lockSessionTx is the cross-process counterpart: every mutating method takes
// the per-session queue_session_locks row inside its transaction, so tail
// scans and tail changes cannot interleave between backend instances. It is a
// no-op on SQLite (single writer; FOR UPDATE is ignored).
func (r *sqliteRepository) lockSessionTx(ctx context.Context, tx *sqlx.Tx, sessionID string) error {
return lockSessionTxIn(ctx, tx, r.db, sessionID)
}
// guardActiveTaskTx rejects queue admissions whose owning task is not live
// (archived or deleted). It takes the task-row lock, so admission serializes
// with task lifecycle cleanup in the global task-row -> session-lock order and
// a post-delete/post-archive admission cannot leave a queue row that survives
// the task's purge. SQLite's single writer is the serialization. When the
// owning tasks table is absent (the queue repository's isolated tests create
// only queue tables) the guard is skipped — the presence check happens at
// construction, never inside the transaction, because a failed statement
// would abort the whole PostgreSQL transaction.
func (r *sqliteRepository) guardActiveTaskTx(ctx context.Context, tx *sqlx.Tx, taskID string) error {
if !r.tasksTablePresent {
return nil
}
res, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE tasks SET updated_at = updated_at
WHERE id = ? AND archived_at IS NULL
`), taskID)
if err != nil {
return fmt.Errorf("guard active task for queue admission: %w", err)
}
affected, err := res.RowsAffected()
if err != nil {
return fmt.Errorf("guard active task rows affected: %w", err)
}
if affected == 0 {
return ErrTaskInactive
}
return nil
}
// lockSessionTxIn takes the per-session cross-process lock inside an existing
// transaction (see lockSessionTx). It is the shared core used by the
// repository methods and by PurgeTaskInTransaction, which runs inside the task
// repository's archive transaction.
func lockSessionTxIn(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, sessionID string) error {
if _, err := tx.ExecContext(ctx, db.Rebind(`
INSERT INTO queue_session_locks (session_id) VALUES (?)
ON CONFLICT(session_id) DO NOTHING
`), sessionID); err != nil {
return fmt.Errorf("ensure queue session lock row: %w", err)
}
if db.DriverName() != "pgx" {
// SQLite has a single writer: the INSERT above already holds the
// database write lock for this transaction, serializing it against
// every other writer. FOR UPDATE is not valid SQLite syntax.
return nil
}
var one int
if err := tx.GetContext(ctx, &one, db.Rebind(`
SELECT 1 FROM queue_session_locks WHERE session_id = ? FOR UPDATE
`), sessionID); err != nil {
return fmt.Errorf("acquire queue session lock: %w", err)
}
return nil
}
func (r *sqliteRepository) withSessionLock(sessionID string) func() {
r.mu.Lock()
lock := r.sessionLocks[sessionID]
@@ -84,6 +168,15 @@ func (r *sqliteRepository) initSchema() error {
queued_at TIMESTAMP NOT NULL,
actor TEXT NOT NULL DEFAULT ''
);
-- Per-session cross-process mutex. Every queue mutation takes this row
-- inside its transaction, so a tail scan and a tail change (insert,
-- drain, reorder, transfer) cannot interleave between backend instances
-- sharing one queue database. SQLite treats FOR UPDATE as a no-op and
-- serializes writes at the DB level anyway; Postgres enforces the row lock.
CREATE TABLE IF NOT EXISTS queue_session_locks (
session_id TEXT PRIMARY KEY
);
`)
if err != nil {
return err
@@ -91,7 +184,7 @@ func (r *sqliteRepository) initSchema() error {
// Existing installations may have the pre-audit shape; fresh installs
// already get the column from CREATE TABLE above, so this replays as a
// duplicate-column error there.
if _, alterErr := r.db.Exec(`ALTER TABLE pending_moves ADD COLUMN actor TEXT NOT NULL DEFAULT ''`); alterErr != nil && !db.IsDuplicateColumnError(alterErr) {
if _, alterErr := r.db.Exec(`ALTER TABLE pending_moves ADD COLUMN actor TEXT NOT NULL DEFAULT ''`); alterErr != nil && !internaldb.IsDuplicateColumnError(alterErr) {
return alterErr
}
return nil
@@ -104,6 +197,12 @@ func (r *sqliteRepository) Insert(ctx context.Context, msg *QueuedMessage, maxPe
return fmt.Errorf("begin insert tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.guardActiveTaskTx(ctx, tx, msg.TaskID); err != nil {
return err
}
if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil {
return err
}
if maxPerSession > 0 {
var count int
@@ -156,6 +255,12 @@ func (r *sqliteRepository) Restore(ctx context.Context, msg *QueuedMessage, maxP
return fmt.Errorf("begin restore tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.guardActiveTaskTx(ctx, tx, msg.TaskID); err != nil {
return err
}
if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil {
return err
}
if maxPerSession > 0 {
var count int
@@ -199,6 +304,12 @@ func (r *sqliteRepository) AppendOrInsertTail(ctx context.Context, sessionID, ta
return nil, false, fmt.Errorf("begin append tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.guardActiveTaskTx(ctx, tx, taskID); err != nil {
return nil, false, err
}
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, false, err
}
tail, err := r.scanTail(ctx, tx, sessionID)
if err != nil {
@@ -275,6 +386,12 @@ func (r *sqliteRepository) InsertOrReplaceByCoalesceKey(ctx context.Context, msg
return nil, false, fmt.Errorf("begin coalesce tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.guardActiveTaskTx(ctx, tx, msg.TaskID); err != nil {
return nil, false, err
}
if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil {
return nil, false, err
}
existing, err := r.findCoalesced(ctx, tx, msg.SessionID, msg.QueuedBy, coalesceKey)
if err != nil {
@@ -314,6 +431,11 @@ func (r *sqliteRepository) InsertOrReplaceLifecycleByCoalesceKey(ctx context.Con
}
defer func() { _ = tx.Rollback() }()
// Global lock order: task row first, then per-session queue locks.
// ArchiveTaskIfActive locks the task row and then (via
// PurgeTaskInTransaction) the affected session locks; taking the session
// lock here before the task guard would invert that order and deadlock
// the two on Postgres (each waiting on the other's held lock).
guard, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE tasks SET updated_at = updated_at
WHERE id = ? AND archived_at IS NULL
@@ -328,6 +450,9 @@ func (r *sqliteRepository) InsertOrReplaceLifecycleByCoalesceKey(ctx context.Con
if rows == 0 {
return nil, false, ErrTaskInactive
}
if err := r.lockSessionTx(ctx, tx, msg.SessionID); err != nil {
return nil, false, err
}
expectedGeneration, ok := lifecycleGenerationFromMetadata(msg.Metadata)
if !ok {
return nil, false, ErrLifecycleCancelled
@@ -388,7 +513,7 @@ func (r *sqliteRepository) PurgeTask(ctx context.Context, taskID string) (int, e
return 0, fmt.Errorf("begin lifecycle queue purge: %w", err)
}
defer func() { _ = tx.Rollback() }()
removed, err := PurgeTaskInTransaction(ctx, tx, r.db, taskID)
removed, err := PurgeTaskInTransaction(ctx, tx, r.db, taskID, nil)
if err != nil {
return 0, err
}
@@ -413,10 +538,75 @@ func lifecycleGenerationInTx(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, task
return generation, nil
}
// purgeQueueRowSessions lists the distinct sessions holding queued rows for
// the task, checking rows.Err() after iteration.
func purgeQueueRowSessions(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskID string) ([]string, error) {
rows, err := tx.QueryxContext(ctx, db.Rebind(`
SELECT DISTINCT session_id FROM queued_messages WHERE task_id = ?
`), taskID)
if err != nil {
return nil, fmt.Errorf("list purge sessions: %w", err)
}
var sessions []string
for rows.Next() {
var sessionID string
if err := rows.Scan(&sessionID); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("scan purge session: %w", err)
}
sessions = append(sessions, sessionID)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("iterate purge sessions: %w", err)
}
if err := rows.Close(); err != nil {
return nil, fmt.Errorf("close purge sessions: %w", err)
}
return sessions, nil
}
// PurgeTaskInTransaction lets the task repository make archive/delete and
// durable queue invalidation one SQLite transaction. It is backend-internal:
// user queue handlers must keep using ownership-checked deletion methods.
func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskID string) (int, error) {
//
// taskSessions is the task's authoritative session set, discovered by the
// caller from its own task_sessions schema — the queue repository never
// reaches across schemas, and a failed cross-schema query would abort the
// caller's transaction on PostgreSQL. The standalone PurgeTask passes nil and
// locks only the sessions that currently hold queue rows.
func PurgeTaskInTransaction(ctx context.Context, tx *sqlx.Tx, db *sqlx.DB, taskID string, taskSessions []string) (int, error) {
// Serialize with per-session tail operations: a purge that races a fold
// or insert could otherwise delete a row an admission just accepted, or
// admit into a queue being purged. Lock the AUTHORITATIVE session set —
// the union of sessions currently holding queue rows and the task's
// sessions (taskSessions, discovered by the caller from its own
// task_sessions schema): a session that is empty at discovery time could
// otherwise admit a task row during the purge and have it survive the
// DELETE snapshot. Locks are taken in sorted order — the same ordering
// every multi-session lock uses — so concurrent purges and queue
// mutations cannot deadlock.
rowSessions, err := purgeQueueRowSessions(ctx, tx, db, taskID)
if err != nil {
return 0, err
}
seen := make(map[string]struct{}, len(rowSessions)+len(taskSessions))
for _, sessionID := range rowSessions {
seen[sessionID] = struct{}{}
}
for _, sessionID := range taskSessions {
seen[sessionID] = struct{}{}
}
ordered := make([]string, 0, len(seen))
for sessionID := range seen {
ordered = append(ordered, sessionID)
}
sort.Strings(ordered)
for _, sessionID := range ordered {
if err := lockSessionTxIn(ctx, tx, db, sessionID); err != nil {
return 0, err
}
}
result, err := tx.ExecContext(ctx, db.Rebind(`DELETE FROM queued_messages WHERE task_id = ?`), taskID)
if err != nil {
return 0, fmt.Errorf("purge queued task entries: %w", err)
@@ -616,6 +806,9 @@ func (r *sqliteRepository) TakeHead(ctx context.Context, sessionID string) (*Que
return nil, fmt.Errorf("begin take tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
row := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT id, session_id, task_id, position, content, model, plan_mode,
@@ -665,6 +858,9 @@ func (r *sqliteRepository) ReserveHead(ctx context.Context, sessionID string) (*
return nil, fmt.Errorf("begin reserve tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
row := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT id, session_id, task_id, position, content, model, plan_mode,
@@ -734,7 +930,18 @@ func (r *sqliteRepository) AcknowledgeByID(ctx context.Context, sessionID, entry
unlock := r.withSessionLock(sessionID)
defer unlock()
res, err := r.db.ExecContext(ctx, r.db.Rebind(`
// The ack must serialize with restore/replace on the same session across
// processes: an autocommit DELETE could otherwise delete a row that a
// concurrent backend just restored/reinserted, losing durable retry state.
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return fmt.Errorf("begin acknowledge tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
res, err := tx.ExecContext(ctx, r.db.Rebind(`
DELETE FROM queued_messages WHERE id = ? AND session_id = ?
`), entryID, sessionID)
if err != nil {
@@ -747,6 +954,9 @@ func (r *sqliteRepository) AcknowledgeByID(ctx context.Context, sessionID, entry
if affected == 0 {
return ErrEntryNotFound
}
if err := tx.Commit(); err != nil {
return err
}
return nil
}
@@ -768,6 +978,9 @@ func (r *sqliteRepository) TakeByID(ctx context.Context, sessionID, entryID stri
return nil, fmt.Errorf("begin take tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
row := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT id, session_id, task_id, position, content, model, plan_mode,
@@ -817,6 +1030,9 @@ func (r *sqliteRepository) ClaimSendNow(ctx context.Context, sessionID string, e
return nil, fmt.Errorf("begin send-now claim tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
ordered, storedByID, err := r.listOrderedStoredSessionEntries(ctx, tx, sessionID)
if err != nil {
@@ -928,6 +1144,15 @@ func (r *sqliteRepository) beginSendNowClaimTx(
unlock()
return nil, "", nil, fmt.Errorf("begin send-now %s tx: %w", action, err)
}
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
// Roll back the started transaction: callers register
// `defer tx.Rollback()` only after this function succeeds, so an
// abandoned tx would keep the pooled connection inside an open
// transaction.
_ = tx.Rollback()
unlock()
return nil, "", nil, err
}
return tx, sessionID, unlock, nil
}
@@ -1238,6 +1463,13 @@ func (r *sqliteRepository) UpdateContentAndMetadata(ctx context.Context, session
return fmt.Errorf("begin update queued tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
// Content edits serialize with merges/folds on the same session: without
// the cross-process session lock a merge's scan-to-write window could
// silently overwrite a concurrent edit (the merge's affected-row check
// only detects deletion, not stale content).
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
var metadataJSON string
query := `SELECT metadata_json FROM queued_messages WHERE id = ? AND session_id = ? AND queued_by = ? AND queued_by NOT IN (?, ?, ?)`
@@ -1289,6 +1521,9 @@ func (r *sqliteRepository) MergeIntoAbove(ctx context.Context, sessionID, source
return nil, fmt.Errorf("begin merge tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
source, err := readMergeSource(ctx, r, tx, sessionID, sourceID)
if err != nil {
@@ -1334,6 +1569,9 @@ func (r *sqliteRepository) AutoMergeIntoAbove(ctx context.Context, sessionID, so
return nil, false, fmt.Errorf("begin automatic merge tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, false, err
}
source, err := readMergeSource(ctx, r, tx, sessionID, sourceID)
if errors.Is(err, ErrEntryNotFound) {
@@ -1368,6 +1606,99 @@ func (r *sqliteRepository) AutoMergeIntoAbove(ctx context.Context, sessionID, so
return target, true, nil
}
// AutoMergeCandidateIntoAbove folds a not-yet-admitted candidate into the
// session's tail entry when compatible. Unlike AutoMergeIntoAbove there is no
// source row to delete: admission at full capacity could not insert one, so
// the fold is the admission. Missing or incompatible tails are successful
// skips and leave storage unchanged.
func (r *sqliteRepository) AutoMergeCandidateIntoAbove(ctx context.Context, candidate *QueuedMessage) (*QueuedMessage, bool, error) {
unlock := r.withSessionLock(candidate.SessionID)
defer unlock()
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return nil, false, fmt.Errorf("begin automatic candidate merge tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
// The full-queue fold is its own admission and runs in its own
// transaction after the guarded insert failed, so it must re-guard the
// task row: an archive/delete can commit between the failed insert and
// the fold, and the fold must not accept a message the purge will then
// silently delete. Task row first, then the session lock.
if err := r.guardActiveTaskTx(ctx, tx, candidate.TaskID); err != nil {
return nil, false, err
}
if err := r.lockSessionTx(ctx, tx, candidate.SessionID); err != nil {
return nil, false, err
}
target, storedContent, storedAttachmentsJSON, storedMetadataJSON, err := r.scanTailWithRawJSON(ctx, tx, candidate.SessionID)
if err != nil {
return nil, false, err
}
if target == nil {
return nil, false, nil
}
// The tail scan already captured the exact stored bytes for the
// compare-and-swap below. The CAS compares raw storage bytes, not
// re-marshalled structs: metadata values round-trip through JSON as maps
// whose key order differs from the struct field order used at write time,
// so re-marshalling would never match. Comparing the raw strings keeps the
// guard exact.
values, compatible := buildAutoMergedEntry(target, candidate)
if !compatible {
return nil, false, nil
}
attachmentsJSON, err := marshalAttachments(values.attachments)
if err != nil {
return nil, false, err
}
metadataJSON, err := marshalMetadata(values.metadata)
if err != nil {
return nil, false, err
}
// The UPDATE is compare-and-swapped against the snapshot read in this
// transaction. Across backend processes the session mutex is process-local,
// so two instances can read the same tail and both attempt a fold; the
// CAS makes the loser's UPDATE affect zero rows and roll back instead of
// silently overwriting the winner's accepted message. Position is part of
// the CAS because a concurrent reorder only changes positions: without it
// a fold could land on a row that is no longer the tail (e.g. it became
// the head), violating FIFO drain order.
res, err := tx.ExecContext(ctx, r.db.Rebind(`
UPDATE queued_messages
SET content = ?, attachments_json = ?, metadata_json = ?
WHERE id = ? AND session_id = ?
AND position = ?
AND content = ?
AND attachments_json = ?
AND metadata_json = ?
`), values.content, attachmentsJSON, metadataJSON,
target.ID, candidate.SessionID, target.Position,
storedContent, storedAttachmentsJSON, storedMetadataJSON)
if err != nil {
return nil, false, fmt.Errorf("update automatic candidate merge target: %w", err)
}
affected, err := res.RowsAffected()
if err != nil {
return nil, false, fmt.Errorf("update automatic candidate merge rows affected: %w", err)
}
if affected == 0 {
// The tail changed between our read and write — a concurrent fold
// committed first or the row was drained. Roll back and report the
// skip so the caller rejects (or retries) rather than overwriting the
// concurrently accepted message.
return nil, false, nil
}
if err := tx.Commit(); err != nil {
return nil, false, err
}
target.Content = values.content
target.Attachments = values.attachments
target.Metadata = values.metadata
return target, true, nil
}
// ReorderEntries atomically rewrites the FIFO positions of the session's
// visible pending entries to match orderedIDs. Reserved in-flight lifecycle
// rows keep their place in the sequence; visible rows are interleaved in the
@@ -1390,6 +1721,9 @@ func (r *sqliteRepository) ReorderEntries(ctx context.Context, sessionID string,
return fmt.Errorf("begin reorder tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
rows, err := tx.QueryxContext(ctx, r.db.Rebind(`
SELECT id, session_id, task_id, position, content, model, plan_mode,
@@ -1573,6 +1907,9 @@ func (r *sqliteRepository) DeleteByID(ctx context.Context, sessionID, entryID st
return fmt.Errorf("begin delete queued tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
var metadataJSON string
err = tx.GetContext(ctx, &metadataJSON, r.db.Rebind(`
@@ -1619,6 +1956,9 @@ func (r *sqliteRepository) DeleteAllBySession(ctx context.Context, sessionID str
return 0, fmt.Errorf("begin delete all queued tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return 0, err
}
candidates, err := cancellationCandidates(ctx, tx, sessionID)
if err != nil {
@@ -1697,11 +2037,36 @@ func isReservedMetadataJSON(metadataJSON string) (bool, error) {
// TransferSession moves all entries (and any pending move) from one session to another.
func (r *sqliteRepository) TransferSession(ctx context.Context, oldSessionID, newSessionID string) error {
// The transfer moves rows out of the source and into the destination, so
// it must hold BOTH sessions' locks: a concurrent source-side insert
// (holding the source lock) could otherwise commit a row the transfer's
// READ COMMITTED UPDATE missed, orphaning it on the old session. Acquire
// in stable sorted order — in-process and cross-process alike — so
// concurrent transfers in opposite directions cannot deadlock.
first, second := oldSessionID, newSessionID
if first > second {
first, second = second, first
}
unlockFirst := r.withSessionLock(first)
defer unlockFirst()
if first != second {
unlockSecond := r.withSessionLock(second)
defer unlockSecond()
}
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return fmt.Errorf("begin transfer tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, first); err != nil {
return err
}
if first != second {
if err := r.lockSessionTx(ctx, tx, second); err != nil {
return err
}
}
// Shift positions on the destination so transferred entries land at the tail
// without colliding on the (session_id, position) implicit ordering.
@@ -1737,6 +2102,9 @@ func (r *sqliteRepository) ReplaceSession(ctx context.Context, sessionID string,
return fmt.Errorf("begin replace session tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM queued_messages WHERE session_id = ?`), sessionID); err != nil {
return fmt.Errorf("clear queued messages: %w", err)
@@ -1792,7 +2160,19 @@ func (r *sqliteRepository) SetPendingMove(ctx context.Context, sessionID string,
if move.QueuedAt.IsZero() {
move.QueuedAt = time.Now().UTC()
}
_, err := r.db.ExecContext(ctx, r.db.Rebind(`
// The pending move is session-scoped and TransferSession moves it with
// the session's queue, so it must serialize on the same per-session lock:
// an unlocked upsert could land on the old session after a transfer
// committed, orphaning the move.
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return fmt.Errorf("begin set pending move tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, r.db.Rebind(`
INSERT INTO pending_moves (id, session_id, task_id, workflow_id, workflow_step_id, step_position, queued_at, actor)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(session_id) DO UPDATE SET
@@ -1804,11 +2184,10 @@ func (r *sqliteRepository) SetPendingMove(ctx context.Context, sessionID string,
actor = excluded.actor
`),
uuid.New().String(), sessionID, move.TaskID, move.WorkflowID, move.WorkflowStepID, move.Position, move.QueuedAt, move.Actor,
)
if err != nil {
); err != nil {
return fmt.Errorf("upsert pending move: %w", err)
}
return nil
return tx.Commit()
}
// GetPendingMove returns the deferred workflow move for a session, or nil when absent.
@@ -1845,6 +2224,11 @@ func (r *sqliteRepository) TakePendingMove(ctx context.Context, sessionID string
return nil, fmt.Errorf("begin take pending tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
// Serialize on the per-session lock so two backend instances cannot both
// read and delete the same pending move, or race a transfer that moves it.
if err := r.lockSessionTx(ctx, tx, sessionID); err != nil {
return nil, err
}
var (
taskID, workflowID, workflowStepID string
@@ -1898,6 +2282,31 @@ func (r *sqliteRepository) scanTail(ctx context.Context, tx *sqlx.Tx, sessionID
return msg, nil
}
// scanTailWithRawJSON reads the highest-position entry for a session plus its
// exact stored content/attachments_json/metadata_json bytes in one query, so
// the CAS condition in AutoMergeCandidateIntoAbove never needs a second
// round-trip for the same row. Returns nil, "", "", "", nil when the queue is
// empty.
func (r *sqliteRepository) scanTailWithRawJSON(ctx context.Context, tx *sqlx.Tx, sessionID string) (*QueuedMessage, string, string, string, error) {
row := tx.QueryRowxContext(ctx, r.db.Rebind(`
SELECT id, session_id, task_id, position, content, model, plan_mode,
attachments_json, metadata_json, queued_at, queued_by,
content, attachments_json, metadata_json
FROM queued_messages
WHERE session_id = ?
ORDER BY position DESC
LIMIT 1
`), sessionID)
msg, rawContent, rawAttachments, rawMetadata, err := scanQueuedRowWithRawJSON(row)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, "", "", "", nil
}
return nil, "", "", "", fmt.Errorf("scan tail with raw bytes: %w", err)
}
return msg, rawContent, rawAttachments, rawMetadata, nil
}
// findCoalesced locates the entry matching the session, owner, and coalesce key inside a transaction.
func (r *sqliteRepository) findCoalesced(ctx context.Context, tx *sqlx.Tx, sessionID, queuedBy, coalesceKey string) (*QueuedMessage, error) {
rows, err := tx.QueryxContext(ctx, r.db.Rebind(`
@@ -1961,6 +2370,43 @@ func scanQueuedRowWithMetadataJSON(
return &msg, metaJSON, nil
}
// scanQueuedRowWithRawJSON scans a queue row plus its exact stored
// content/attachments_json/metadata_json strings. The raw strings back the
// compare-and-swap in AutoMergeCandidateIntoAbove: metadata values round-trip
// through JSON as maps whose key order differs from the struct field order
// used at write time, so only raw bytes compare exactly. The caller's SELECT
// must project those three columns a second time after the regular eleven.
func scanQueuedRowWithRawJSON(
scanner interface{ Scan(dest ...any) error },
) (*QueuedMessage, string, string, string, error) {
var (
msg QueuedMessage
planModeInt int
attachmentsJSON, metaJSON string
rawContent string
rawAttachments, rawMeta string
)
if err := scanner.Scan(
&msg.ID, &msg.SessionID, &msg.TaskID, &msg.Position, &msg.Content, &msg.Model,
&planModeInt, &attachmentsJSON, &metaJSON, &msg.QueuedAt, &msg.QueuedBy,
&rawContent, &rawAttachments, &rawMeta,
); err != nil {
return nil, "", "", "", err
}
msg.PlanMode = planModeInt != 0
if attachmentsJSON != "" && attachmentsJSON != "[]" {
if err := json.Unmarshal([]byte(attachmentsJSON), &msg.Attachments); err != nil {
return nil, "", "", "", fmt.Errorf("unmarshal attachments: %w", err)
}
}
if metaJSON != "" && metaJSON != "{}" {
if err := json.Unmarshal([]byte(metaJSON), &msg.Metadata); err != nil {
return nil, "", "", "", fmt.Errorf("unmarshal metadata: %w", err)
}
}
return &msg, rawContent, rawAttachments, rawMeta, nil
}
// marshalAttachments serializes attachments for storage.
func marshalAttachments(att []MessageAttachment) (string, error) {
if len(att) == 0 {
@@ -179,7 +179,62 @@ func (s *Service) queueMessageWithMetadataAdmission(ctx context.Context, session
admittedCtx, sessionID, taskID, content, model, userID, planMode, attachments, metadata, maxPerSession,
)
if err != nil {
return err
if errors.Is(err, ErrQueueFull) && autoMergeEnabled && afterInsert == nil {
// A full queue must still accept a message that would fold into
// the tail: admission-time auto-merge runs after a successful
// insert, so at capacity it could never fire. Fold the
// candidate directly — the fold is the admission. The
// afterInsert hook is excluded because it claims staged
// attachments against a persisted source row, which this path
// never creates.
candidate := &QueuedMessage{
SessionID: sessionID,
TaskID: taskID,
Content: content,
Model: model,
PlanMode: planMode,
Attachments: attachments,
Metadata: copyMessageMetadata(metadata, 0),
QueuedBy: userID,
}
merged, didMerge, mergeErr := s.repo.AutoMergeCandidateIntoAbove(admittedCtx, candidate)
switch {
case mergeErr != nil:
if errors.Is(mergeErr, ErrTaskInactive) {
// The task was archived or deleted while the admission
// was in flight; surface the inactive-task contract
// instead of a misleading queue-full rejection.
return mergeErr
}
s.logger.Error("automatic merge into full queue failed; preserving queue full rejection",
zap.String("session_id", sessionID),
zap.Error(mergeErr))
return err
case didMerge && merged != nil:
s.logger.Info("automatically merged queued entry into full tail",
zap.String("session_id", sessionID),
zap.String("surviving_entry_id", merged.ID))
queued = merged
return nil
default:
// The fold skipped because the tail is absent or changed —
// a concurrent drain freed capacity (or a concurrent fold or
// drain moved the tail) between the failed insert and the
// fold scan. Retry the ordinary insert once under the held
// admission lock: it succeeds now that capacity is free, and
// still returns ErrQueueFull when another writer keeps the
// cap occupied. Without this retry a stale ErrQueueFull
// would drop a message that is admissible at fold time.
source, err = s.insertQueueMessageWithMetadata(
admittedCtx, sessionID, taskID, content, model, userID, planMode, attachments, metadata, maxPerSession,
)
if err != nil {
return err
}
}
} else {
return err
}
}
if afterInsert != nil {
if err := afterInsert(admittedCtx, source); err != nil {
@@ -3,6 +3,7 @@ package messagequeue
import (
"context"
"errors"
"sync"
"testing"
"time"
@@ -69,19 +70,193 @@ func TestService_AutoMergeChainsThreeAdmissions(t *testing.T) {
}
}
func TestService_AutoMergeDoesNotBypassCapacity(t *testing.T) {
func TestService_AutoMergeKeepsDifferentSenderSessionsSeparate(t *testing.T) {
svc := newAutoMergeTestService(t, 10)
sender := func(session string) map[string]interface{} {
return map[string]interface{}{
"sender_task_id": "sender-task", "sender_session_id": session,
"sender_session_name": "reviewer-" + session, "sender_task_title": "Sender task",
}
}
first, err := svc.QueueMessageWithMetadata(context.Background(), "session", "task", "first", "", QueuedByAgent, false, nil, sender("one"))
if err != nil {
t.Fatalf("queue first: %v", err)
}
// A different session of the same sender task is a different agent and
// execution context; its prompt must stay a separate queue entry.
second, err := svc.QueueMessageWithMetadata(context.Background(), "session", "task", "second", "", QueuedByAgent, false, nil, sender("two"))
if err != nil {
t.Fatalf("queue second: %v", err)
}
if second.ID == first.ID {
t.Fatal("different sender sessions must stay separate, got a fold")
}
status := svc.GetStatus(context.Background(), "session")
if status.Count != 2 {
t.Fatalf("status = %+v, want two separate entries", status)
}
}
func TestService_AutoMergeFoldsCompatibleAdmissionAtCapacity(t *testing.T) {
svc := newAutoMergeTestService(t, 1)
first, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil)
if err != nil {
t.Fatalf("queue first: %v", err)
}
// A compatible message must fold into the tail instead of being rejected:
// the fold is the admission, so the queue never reports full for a message
// that would merge anyway.
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "", QueuedByUser, false, nil)
if err != nil {
t.Fatalf("compatible admission at capacity = %v, want fold", err)
}
if second.ID != first.ID || second.Content != "first\n\nsecond" {
t.Fatalf("second admission = %+v, want folded survivor %s", second, first.ID)
}
status := svc.GetStatus(context.Background(), "session")
if status.Count != 1 || status.Entries[0].ID != first.ID || status.Entries[0].Content != "first\n\nsecond" {
t.Fatalf("queue after full-queue fold: %+v", status)
}
}
func TestService_AutoMergeAtFullQueueRejectsIncompatibleAdmission(t *testing.T) {
svc := newAutoMergeTestService(t, 1)
if _, err := svc.QueueMessage(context.Background(), "session", "task", "first", "model-a", QueuedByUser, false, nil); err != nil {
t.Fatalf("queue first: %v", err)
}
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "model-b", QueuedByUser, false, nil)
if !errors.Is(err, ErrQueueFull) || second != nil {
t.Fatalf("incompatible full admission = %+v, err=%v, want queue full", second, err)
}
}
func TestService_AutoMergeDisabledAtFullQueueRejects(t *testing.T) {
svc := newAutoMergeTestService(t, 1)
svc.SetAutoMergeEnabled(false)
if _, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil); err != nil {
t.Fatalf("queue first: %v", err)
}
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "", QueuedByUser, false, nil)
if !errors.Is(err, ErrQueueFull) || second != nil {
t.Fatalf("second admission = %+v, err=%v, want queue full", second, err)
t.Fatalf("disabled full admission = %+v, err=%v, want queue full", second, err)
}
}
func TestService_AutoMergeAtFullQueueSkipsAfterInsertHook(t *testing.T) {
svc := newAutoMergeTestService(t, 1)
if _, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil); err != nil {
t.Fatalf("queue first: %v", err)
}
ran := false
second, err := svc.QueueMessageWithMetadataAfterInsert(
context.Background(), "session", "task", "second", "", QueuedByUser, false, nil, nil,
func(context.Context, *QueuedMessage) error { ran = true; return nil },
)
if !errors.Is(err, ErrQueueFull) || second != nil {
t.Fatalf("after-insert full admission = %+v, err=%v, want queue full", second, err)
}
if ran {
t.Fatal("after-insert hook ran for a rejected full-queue admission")
}
}
// autoMergeCandidateErrorRepository wraps a repository and fails every
// full-queue candidate fold with the configured error.
type autoMergeCandidateErrorRepository struct {
Repository
err error
}
func (r *autoMergeCandidateErrorRepository) AutoMergeCandidateIntoAbove(context.Context, *QueuedMessage) (*QueuedMessage, bool, error) {
return nil, false, r.err
}
// failNextInsertRepository fails the next N Insert calls with ErrQueueFull,
// simulating a cross-process admission that observed a stale full count while
// the underlying queue already drained.
type failNextInsertRepository struct {
Repository
mu sync.Mutex
failNext int
}
func (r *failNextInsertRepository) Insert(ctx context.Context, msg *QueuedMessage, maxPerSession int) error {
r.mu.Lock()
if r.failNext > 0 {
r.failNext--
r.mu.Unlock()
return ErrQueueFull
}
r.mu.Unlock()
return r.Repository.Insert(ctx, msg, maxPerSession)
}
func TestService_AutoMergeFullQueueRetriesInsertAfterFoldSkip(t *testing.T) {
repo := &failNextInsertRepository{Repository: NewMemoryRepository()}
svc := newAutoMergeTestServiceWithRepository(t, repo, 1)
if _, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil); err != nil {
t.Fatalf("queue first: %v", err)
}
// A concurrent drain frees capacity between the failed insert and the
// fold scan: the fold sees an empty queue and skips, and admission must
// retry the ordinary insert instead of returning the stale ErrQueueFull.
if _, ok := svc.TakeQueued(context.Background(), "session"); !ok {
t.Fatal("drain first entry")
}
repo.mu.Lock()
repo.failNext = 1
repo.mu.Unlock()
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "", QueuedByUser, false, nil)
if err != nil {
t.Fatalf("admission after fold skip = %v, want insert retry to succeed", err)
}
if second == nil || second.Content != "second" {
t.Fatalf("queued entry = %+v, want the retried message", second)
}
status := svc.GetStatus(context.Background(), "session")
if status.Count != 1 || status.Entries[0].Content != "second" {
t.Fatalf("status = %+v, want the retried message queued", status)
}
}
func TestService_AutoMergeFullQueueCandidateMergeErrorDegradesToQueueFull(t *testing.T) {
wantErr := errors.New("automatic candidate merge unavailable")
repo := &autoMergeCandidateErrorRepository{Repository: NewMemoryRepository(), err: wantErr}
svc := newAutoMergeTestServiceWithRepository(t, repo, 1)
first, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil)
if err != nil {
t.Fatalf("queue first: %v", err)
}
// The full-queue fold itself failed: admission must degrade to the
// original ErrQueueFull (never accept a message whose fold did not
// happen, and never surface the repository error).
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "", QueuedByUser, false, nil)
if !errors.Is(err, ErrQueueFull) || second != nil {
t.Fatalf("candidate merge error admission = %+v, err=%v, want queue full", second, err)
}
status := svc.GetStatus(context.Background(), "session")
if status.Count != 1 || status.Entries[0].ID != first.ID || status.Entries[0].Content != "first" {
t.Fatalf("queue changed after full admission: %+v", status)
t.Fatalf("queue changed after degraded full admission: %+v", status)
}
}
func TestService_AutoMergeFullQueueTaskInactiveSurfacesContract(t *testing.T) {
repo := &autoMergeCandidateErrorRepository{Repository: NewMemoryRepository(), err: ErrTaskInactive}
svc := newAutoMergeTestServiceWithRepository(t, repo, 1)
if _, err := svc.QueueMessage(context.Background(), "session", "task", "first", "", QueuedByUser, false, nil); err != nil {
t.Fatalf("queue first: %v", err)
}
// The full-queue fold's task guard failed (task archived/deleted between
// the failed insert and the fold): admission must surface ErrTaskInactive,
// not the stale ErrQueueFull.
second, err := svc.QueueMessage(context.Background(), "session", "task", "second", "", QueuedByUser, false, nil)
if !errors.Is(err, ErrTaskInactive) || second != nil {
t.Fatalf("task-inactive admission = %+v, err=%v, want ErrTaskInactive", second, err)
}
status := svc.GetStatus(context.Background(), "session")
if status.Count != 1 || status.Entries[0].Content != "first" {
t.Fatalf("queue changed after rejected admission: %+v", status)
}
}
@@ -215,6 +215,13 @@ func (r *Repository) CreateTaskWithWorkflowStepAdmission(
}
defer func() { _ = tx.Rollback() }()
// Same order as createTask and the workspace cascade: workspace row
// first, then workflow-step locks. Without it an admission holding a step
// lock while the cascade holds the workspace and waits for that step
// deadlocks on Postgres.
if err := r.lockWorkspaceRowStdTx(ctx, tx, task.WorkspaceID); err != nil {
return err
}
if err := r.lockWorkflowStepsForAdmission(ctx, tx, targetStepID, feederStepID); err != nil {
return err
}
@@ -283,6 +290,17 @@ func (r *Repository) createTask(ctx context.Context, task *models.Task, targetSt
}
defer func() { _ = tx.Rollback() }()
// Serialize with the workspace delete cascade: the cascade locks the
// workspace row before inventorying its tasks, so a task created here
// either commits before the cascade's inventory (and is purged with the
// rest) or blocks until the cascade finishes, when the workspace is gone
// and the insert fails its foreign key. The workspace lock is taken
// before any workflow-step lock so the creation/admission paths share one
// order with the cascade.
if err := r.lockWorkspaceRowStdTx(ctx, tx, task.WorkspaceID); err != nil {
return err
}
if err := r.ensureWorkflowStepCapacity(ctx, tx, targetStepID, limit); err != nil {
return err
}
@@ -563,6 +581,24 @@ func (r *Repository) updateTaskWithWorkflowStepAdmission(
}
defer func() { _ = tx.Rollback() }()
// Workspace row before workflow-step lock, matching createTask and the
// workspace cascade: the update path must not hold a step lock while the
// cascade holds the workspace and waits for that step (Postgres
// deadlock). The task's workspace is read from its row so the caller's
// model cannot bypass the ordering.
var workspaceID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(`SELECT workspace_id FROM tasks WHERE id = ?`), task.ID).Scan(&workspaceID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
// Preserve the ErrTaskNotFound sentinel callers relied on before
// the workspace read was introduced (a task deleted concurrently
// with a move is reachable).
return false, fmt.Errorf("%w: %s", ErrTaskNotFound, task.ID)
}
return false, fmt.Errorf("read task workspace for admission: %w", err)
}
if err := r.lockWorkspaceRowStdTx(ctx, tx, workspaceID); err != nil {
return false, err
}
if err := lockWorkflowStepForCapacity(ctx, tx, r.db.DriverName(), r.db.Rebind, targetStepID); err != nil {
return false, err
}
@@ -1012,6 +1048,20 @@ func (r *Repository) DeleteTask(ctx context.Context, id string) error {
return err
}
defer func() { _ = tx.Rollback() }()
// Serialize with session/worktree creation FIRST (task-row lock), then
// capture the authoritative session set: a concurrent CreateTaskSession
// holds the same task-row barrier, so every session committed before this
// lock is visible to the capture and anything after blocks until the task
// row is gone. Capturing before the lock could use a stale set (a session
// created mid-flight would never be purged). The session capture must also
// precede the task-row DELETE because task_sessions cascades on deletion.
if err := r.lockTaskRowInTx(ctx, tx, id); err != nil {
return err
}
sessions, err := r.taskQueueSessionsInTx(ctx, tx, id)
if err != nil {
return err
}
result, err := tx.ExecContext(ctx, r.db.Rebind(`DELETE FROM tasks WHERE id = ?`), id)
if err != nil {
return err
@@ -1021,7 +1071,7 @@ func (r *Repository) DeleteTask(ctx context.Context, id string) error {
if rows == 0 {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
if err := r.purgeTaskQueueInTx(ctx, tx, id); err != nil {
if err := r.purgeTaskQueueInTx(ctx, tx, id, sessions); err != nil {
return err
}
if err := tx.Commit(); err != nil {
@@ -1747,7 +1797,11 @@ func (r *Repository) ArchiveTask(ctx context.Context, id string) error {
if rows == 0 {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
if err := r.purgeTaskQueueInTx(ctx, tx, id); err != nil {
sessions, err := r.taskQueueSessionsInTx(ctx, tx, id)
if err != nil {
return err
}
if err := r.purgeTaskQueueInTx(ctx, tx, id, sessions); err != nil {
return err
}
if err := tx.Commit(); err != nil {
@@ -1786,7 +1840,11 @@ func (r *Repository) ArchiveTaskIfActive(ctx context.Context, id, cascadeID stri
if rows == 0 {
return false, tx.Commit()
}
if err := r.purgeTaskQueueInTx(ctx, tx, id); err != nil {
sessions, err := r.taskQueueSessionsInTx(ctx, tx, id)
if err != nil {
return false, err
}
if err := r.purgeTaskQueueInTx(ctx, tx, id, sessions); err != nil {
return false, err
}
if err := tx.Commit(); err != nil {
@@ -1796,8 +1854,37 @@ func (r *Repository) ArchiveTaskIfActive(ctx context.Context, id, cascadeID stri
return true, nil
}
func (r *Repository) purgeTaskQueueInTx(ctx context.Context, tx *sqlx.Tx, taskID string) error {
_, err := messagequeue.PurgeTaskInTransaction(ctx, tx, r.db, taskID)
// taskQueueSessionsInTx returns the task's authoritative session set (its
// task_sessions rows). Callers must capture it BEFORE any statement that
// deletes the task row: task_sessions cascades on task deletion, so a
// post-delete discovery returns nothing and a concurrent admission to an
// empty session could survive the purge.
func (r *Repository) taskQueueSessionsInTx(ctx context.Context, tx *sqlx.Tx, taskID string) ([]string, error) {
var sessions []string
rows, err := tx.QueryxContext(ctx, r.db.Rebind(`SELECT id FROM task_sessions WHERE task_id = ?`), taskID)
if err != nil {
return nil, fmt.Errorf("list task purge sessions: %w", err)
}
for rows.Next() {
var sessionID string
if err := rows.Scan(&sessionID); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("scan task purge session: %w", err)
}
sessions = append(sessions, sessionID)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("iterate task purge sessions: %w", err)
}
if err := rows.Close(); err != nil {
return nil, fmt.Errorf("close task purge sessions: %w", err)
}
return sessions, nil
}
func (r *Repository) purgeTaskQueueInTx(ctx context.Context, tx *sqlx.Tx, taskID string, sessions []string) error {
_, err := messagequeue.PurgeTaskInTransaction(ctx, tx, r.db, taskID, sessions)
if internaldb.IsMissingTableError(err) {
return nil
}
@@ -2,6 +2,8 @@ package sqlite
import (
"context"
"database/sql"
"errors"
"fmt"
"github.com/jmoiron/sqlx"
@@ -47,3 +49,61 @@ func (r *Repository) taskCleanupBarrierLocked(ctx context.Context, tx *sqlx.Tx,
}
return nil
}
// lockTaskRowInTx takes the task row lock (Postgres SELECT FOR UPDATE; SQLite
// serializes via its single writer). It is the raw row lock behind the
// creation barrier, used by lifecycle paths that must serialize against
// session/worktree creation but should not be rejected by an in-flight
// cleanup job.
func (r *Repository) lockTaskRowInTx(ctx context.Context, tx *sqlx.Tx, taskID string) error {
if !dialect.IsPostgres(r.db.DriverName()) {
return nil
}
var lockedTaskID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT id FROM tasks WHERE id = ? FOR UPDATE`,
), taskID).Scan(&lockedTaskID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
// On PostgreSQL a missing task surfaces here (before any
// RowsAffected check); keep the ErrTaskNotFound classification.
return fmt.Errorf("%w: %s", ErrTaskNotFound, taskID)
}
return fmt.Errorf("lock task row: %w", err)
}
return nil
}
// lockWorkspaceRowInTx takes the workspace row lock (Postgres SELECT FOR
// UPDATE; SQLite serializes via its single writer). Task creation and the
// workspace delete cascade take it so a task created after a cascade's
// inventory cannot escape the queue purge.
func (r *Repository) lockWorkspaceRowInTx(ctx context.Context, tx *sqlx.Tx, workspaceID string) error {
return r.lockWorkspaceRowStdTx(ctx, tx, workspaceID)
}
// lockWorkspaceRowStdTx is the stdlib-transaction variant used by the task
// creation/admission paths (which begin database/sql transactions). An empty
// workspace id (config tasks) skips the lock. A NON-EMPTY workspace that no
// longer has a row is rejected: a creator waiting behind a workspace delete
// cascade resumes with sql.ErrNoRows after the cascade commits, and inserting
// then would orphan the task (tasks.workspace_id is not a foreign key).
func (r *Repository) lockWorkspaceRowStdTx(ctx context.Context, tx interface {
QueryRowContext(context.Context, string, ...any) *sql.Row
}, workspaceID string) error {
if !dialect.IsPostgres(r.db.DriverName()) {
return nil
}
if workspaceID == "" {
return nil
}
var lockedWorkspaceID string
if err := tx.QueryRowContext(ctx, r.db.Rebind(
`SELECT id FROM workspaces WHERE id = ? FOR UPDATE`,
), workspaceID).Scan(&lockedWorkspaceID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return fmt.Errorf("%w: %s", repoerrors.ErrWorkspaceNotFound, workspaceID)
}
return fmt.Errorf("lock workspace row: %w", err)
}
return nil
}
@@ -0,0 +1,542 @@
package sqlite
import (
"context"
"database/sql"
"errors"
"fmt"
"strings"
"testing"
"time"
"github.com/google/uuid"
_ "github.com/jackc/pgx/v5/stdlib"
"github.com/jmoiron/sqlx"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/task/models"
"github.com/kandev/kandev/internal/testutil"
)
// pgBackendPID returns the backend PID of the repository's single connection,
// used to scope pg_locks barrier polls to this test's worker: transactionid
// waits are server-global, so an unrelated waiter from a concurrently running
// package could otherwise false-trigger the barrier.
func pgBackendPID(t *testing.T, db *sqlx.DB) int {
t.Helper()
var pid int
if err := db.QueryRow(`SELECT pg_backend_pid()`).Scan(&pid); err != nil {
t.Fatalf("pg_backend_pid: %v", err)
}
return pid
}
// waitForWaitingLocks polls pg_locks for the expected number of NOT-granted
// transactionid waits owned by the given backend PID, scoping the barrier to
// this test's worker. The poller must be the held transaction's connection
// (each test handle has a single connection busy inside its lock tx).
func waitForWaitingLocks(t *testing.T, poller interface {
QueryRowContext(context.Context, string, ...any) *sql.Row
}, backendPID, want int, what string) {
t.Helper()
deadline := time.Now().Add(10 * time.Second)
for {
var waiting int
if err := poller.QueryRowContext(context.Background(), `
SELECT count(*) FROM pg_locks
WHERE NOT granted AND locktype = 'transactionid' AND pid = $1
`, backendPID).Scan(&waiting); err != nil {
t.Fatalf("query pg_locks: %v", err)
}
if waiting >= want {
return
}
if time.Now().After(deadline) {
t.Fatalf("%s never reached %d waiting lock(s)", what, want)
}
time.Sleep(10 * time.Millisecond)
}
}
// newTaskPostgresRepoPair opens one isolated Postgres schema and constructs
// two task repositories over it (plus the queue tables the purge touches,
// which production creates on the same database), plus an extra
// single-connection DB handle on the same schema. Two instances simulate two
// backend processes sharing a queue database; the extra handle lets a test
// hold a queue_session_locks row on a separate connection while one of the
// worker pools is busy in its lock transaction.
func newTaskPostgresRepoPair(t *testing.T) (*Repository, *Repository, *sqlx.DB) {
t.Helper()
dsn := testutil.PostgresDSNFromEnv(t)
schema := "kandev_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
setupDB, err := sqlx.Open("pgx", dsn)
if err != nil {
t.Fatalf("open postgres: %v", err)
}
if _, err := setupDB.Exec("CREATE SCHEMA " + schema); err != nil {
_ = setupDB.Close()
t.Fatalf("create postgres schema %s: %v", schema, err)
}
t.Cleanup(func() {
_, _ = setupDB.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
_ = setupDB.Close()
})
open := func() *Repository {
db, err := sqlx.Open("pgx", dsn)
if err != nil {
t.Fatalf("open postgres: %v", err)
}
db.SetMaxOpenConns(1)
t.Cleanup(func() { _ = db.Close() })
if _, err := db.Exec("SET search_path TO " + schema); err != nil {
t.Fatalf("set postgres search_path %s: %v", schema, err)
}
// Task schema first: the queue repository resolves tasks-table
// presence at construction, so it must exist by then for the
// admission liveness guard to engage.
repo, err := NewWithDB(db, db, nil)
if err != nil {
t.Fatalf("init task schema: %v", err)
}
if _, err := messagequeue.NewSQLiteRepository(db, db); err != nil {
t.Fatalf("init queue schema: %v", err)
}
return repo
}
repoA, repoB := open(), open()
dbC := mustOpenTaskTestDB(t, dsn, schema)
return repoA, repoB, dbC
}
func mustOpenTaskTestDB(t *testing.T, dsn, schema string) *sqlx.DB {
t.Helper()
db, err := sqlx.Open("pgx", dsn)
if err != nil {
t.Fatalf("open postgres: %v", err)
}
db.SetMaxOpenConns(1)
t.Cleanup(func() { _ = db.Close() })
if _, err := db.Exec("SET search_path TO " + schema); err != nil {
t.Fatalf("set postgres search_path %s: %v", schema, err)
}
return db
}
// seedTaskWithSession creates a workspace, a task in it, and one session.
func seedTaskWithSession(t *testing.T, repo *Repository, taskID, workspaceID, sessionID string) {
t.Helper()
ctx := context.Background()
if err := repo.CreateWorkspace(ctx, &models.Workspace{ID: workspaceID, Name: "Purge race"}); err != nil {
t.Fatalf("create workspace: %v", err)
}
if err := repo.CreateTask(ctx, &models.Task{
ID: taskID, WorkspaceID: workspaceID, WorkflowID: "wf-" + taskID,
WorkflowStepID: "step-" + taskID, Title: taskID, Priority: "medium",
}); err != nil {
t.Fatalf("create task: %v", err)
}
if err := repo.CreateTaskSession(ctx, &models.TaskSession{ID: sessionID, TaskID: taskID, State: models.TaskSessionStateIdle}); err != nil {
t.Fatalf("create session: %v", err)
}
}
// TestPostgresRepository_DeleteTask_LocksEmptySessionDuringPurge proves
// DeleteTask captures the task's session set BEFORE deleting the task row
// (task_sessions cascades on deletion): the purge must lock the session even
// when its queue is empty, so a concurrent admission cannot survive the
// purge. The competing backend holds the session lock; DeleteTask blocks on
// it, proving the empty session was still locked.
func TestPostgresRepository_DeleteTask_LocksEmptySessionDuringPurge(t *testing.T) {
repoA, repoB, _ := newTaskPostgresRepoPair(t)
ctx := context.Background()
const (
taskID = "task-del-race"
sessID = "sess-del-race"
)
seedTaskWithSession(t, repoA, taskID, "ws-del-race", sessID)
dbA := repoA.db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('sess-del-race')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'sess-del-race' FOR UPDATE
`); err != nil {
t.Fatalf("lock session: %v", err)
}
delPID := pgBackendPID(t, repoB.db)
delDone := make(chan error, 1)
go func() {
delDone <- repoB.DeleteTask(ctx, taskID)
}()
waitForWaitingLocks(t, lockTx, delPID, 1, "DeleteTask on the empty session's lock")
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-delDone; err != nil {
t.Fatalf("delete task: %v", err)
}
var count int
if err := dbA.GetContext(ctx, &count, `SELECT count(*) FROM tasks WHERE id = 'task-del-race'`); err != nil {
t.Fatalf("count tasks: %v", err)
}
if count != 0 {
t.Fatalf("task survived delete: count=%d", count)
}
}
// TestPostgresRepository_WorkspaceCascade_GuardsTaskRowsBeforeSessionLocks
// proves the cascade establishes the global task-row -> session-lock order:
// it must block on a held task row BEFORE taking any queue session lock,
// otherwise lifecycle admission (task row first, then session lock) and the
// cascade deadlock on Postgres.
func TestPostgresRepository_WorkspaceCascade_GuardsTaskRowsBeforeSessionLocks(t *testing.T) {
repoA, repoB, _ := newTaskPostgresRepoPair(t)
ctx := context.Background()
const taskID = "task-cascade-race"
seedTaskWithSession(t, repoA, taskID, "ws-cascade-race", "sess-cascade-race")
dbA := repoA.db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
SELECT 1 FROM tasks WHERE id = 'task-cascade-race' FOR UPDATE
`); err != nil {
t.Fatalf("lock task row: %v", err)
}
cascadePID := pgBackendPID(t, repoB.db)
cascadeDone := make(chan error, 1)
go func() {
_, _, err := repoB.deleteWorkspaceCascade(ctx, "ws-cascade-race", nil, nil)
cascadeDone <- err
}()
waitForWaitingLocks(t, lockTx, cascadePID, 1, "cascade on the task row guard (lock order)")
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit lock tx: %v", err)
}
if err := <-cascadeDone; err != nil {
t.Fatalf("workspace cascade: %v", err)
}
var count int
if err := dbA.GetContext(ctx, &count, `SELECT count(*) FROM workspaces WHERE id = 'ws-cascade-race'`); err != nil {
t.Fatalf("count workspaces: %v", err)
}
if count != 0 {
t.Fatalf("workspace survived cascade: count=%d", count)
}
}
// TestPostgresRepository_QueueAdmissionRejectedAfterTaskDelete proves the
// ordinary queue admission task-liveness guard: after DeleteTask commits, a
// stale admission targeting the deleted task's session is rejected with
// ErrTaskInactive and no queue row survives.
func TestPostgresRepository_QueueAdmissionRejectedAfterTaskDelete(t *testing.T) {
repoA, _, _ := newTaskPostgresRepoPair(t)
ctx := context.Background()
seedTaskWithSession(t, repoA, "task-post-del", "ws-post-del", "sess-post-del")
queueRepo, err := messagequeue.NewSQLiteRepository(repoA.db, repoA.db)
if err != nil {
t.Fatalf("init queue repo: %v", err)
}
if err := repoA.DeleteTask(ctx, "task-post-del"); err != nil {
t.Fatalf("delete task: %v", err)
}
err = queueRepo.Insert(ctx, &messagequeue.QueuedMessage{
SessionID: "sess-post-del", TaskID: "task-post-del", Content: "orphan", QueuedBy: messagequeue.QueuedByUser,
}, 0)
if !errors.Is(err, messagequeue.ErrTaskInactive) {
t.Fatalf("post-delete admission err = %v, want ErrTaskInactive", err)
}
count, err := queueRepo.CountBySession(ctx, "sess-post-del")
if err != nil {
t.Fatalf("count: %v", err)
}
if count != 0 {
t.Fatalf("queue survived task delete: count=%d", count)
}
}
// TestPostgresRepository_AutoMergeCandidateIntoAbove_RejectedAfterArchive
// proves the full-queue fold re-guards the task row in its own transaction:
// an archive committing between the failed insert and the fold must make the
// fold fail with ErrTaskInactive instead of accepting a message the purge
// will silently delete. The competing backend holds the task row (simulating
// the archive mid-flight), the fold blocks on the guard, the archive commits,
// and the fold is rejected.
func TestPostgresRepository_AutoMergeCandidateIntoAbove_RejectedAfterArchive(t *testing.T) {
repoA, repoB, _ := newTaskPostgresRepoPair(t)
ctx := context.Background()
const (
taskID = "task-fold-arch-race"
sessID = "sess-fold-arch-race"
)
seedTaskWithSession(t, repoA, taskID, "ws-fold-arch-race", sessID)
queueRepoA, err := messagequeue.NewSQLiteRepository(repoA.db, repoA.db)
if err != nil {
t.Fatalf("init queue repo A: %v", err)
}
if err := queueRepoA.Insert(ctx, &messagequeue.QueuedMessage{
SessionID: sessID, TaskID: taskID, Content: "first", QueuedBy: messagequeue.QueuedByUser,
}, 0); err != nil {
t.Fatalf("seed tail: %v", err)
}
queueRepoB, err := messagequeue.NewSQLiteRepository(repoB.db, repoB.db)
if err != nil {
t.Fatalf("init queue repo B: %v", err)
}
// Simulate an archive that holds the task row while it purges.
dbA := repoA.db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
SELECT id FROM tasks WHERE id = 'task-fold-arch-race' FOR UPDATE
`); err != nil {
t.Fatalf("lock task row: %v", err)
}
foldPID := pgBackendPID(t, repoB.db)
foldDone := make(chan error, 1)
candidate := &messagequeue.QueuedMessage{
SessionID: sessID, TaskID: taskID, Content: "second", QueuedBy: messagequeue.QueuedByUser,
}
go func() {
merged, didMerge, err := queueRepoB.AutoMergeCandidateIntoAbove(ctx, candidate)
if err != nil {
foldDone <- err
return
}
foldDone <- fmt.Errorf("fold succeeded (didMerge=%v merged=%+v), want ErrTaskInactive after archive", didMerge, merged)
}()
waitForWaitingLocks(t, lockTx, foldPID, 1, "fold on the task row guard")
// The archive commits (task archived) while the fold waits on the guard.
if _, err := lockTx.ExecContext(ctx, `
UPDATE tasks SET archived_at = now(), updated_at = now() WHERE id = 'task-fold-arch-race'
`); err != nil {
t.Fatalf("archive task: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit archive: %v", err)
}
if err := <-foldDone; !errors.Is(err, messagequeue.ErrTaskInactive) {
t.Fatalf("fold after archive err = %v, want ErrTaskInactive", err)
}
}
// TestPostgresRepository_DeleteTask_MissingTaskReturnsErrTaskNotFound proves
// the task-row lock preserves the ErrTaskNotFound classification: on
// PostgreSQL a missing task surfaces as sql.ErrNoRows from the FOR UPDATE
// before any RowsAffected check.
func TestPostgresRepository_DeleteTask_MissingTaskReturnsErrTaskNotFound(t *testing.T) {
repoA, _, _ := newTaskPostgresRepoPair(t)
ctx := context.Background()
if err := repoA.DeleteTask(ctx, "no-such-task"); !errors.Is(err, ErrTaskNotFound) {
t.Fatalf("DeleteTask missing task err = %v, want ErrTaskNotFound", err)
}
}
// TestPostgresRepository_DeleteTask_SerializesWithSessionCreation proves
// DeleteTask takes the task-row creation barrier BEFORE capturing the session
// set: a session created mid-flight (while the task row is locked) must be
// included in the purge, otherwise its queue survives task deletion. The
// competing backend holds the task row, DeleteTask blocks on it, a new
// session is committed, and the subsequent purge must block on that new
// session's lock (proving it was captured).
func TestPostgresRepository_DeleteTask_SerializesWithSessionCreation(t *testing.T) {
repoA, repoB, dbC := newTaskPostgresRepoPair(t)
ctx := context.Background()
const (
taskID = "task-del-sess-race"
sessID = "sess-del-sess-race"
)
seedTaskWithSession(t, repoA, taskID, "ws-del-sess-race", sessID)
dbA := repoA.db
// Simulate CreateTaskSession holding the task-row creation barrier.
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
SELECT id FROM tasks WHERE id = 'task-del-sess-race' FOR UPDATE
`); err != nil {
t.Fatalf("lock task row: %v", err)
}
delPID := pgBackendPID(t, repoB.db)
delDone := make(chan error, 1)
go func() {
delDone <- repoB.DeleteTask(ctx, taskID)
}()
waitForWaitingLocks(t, lockTx, delPID, 1, "DeleteTask on the task row barrier")
// Hold the NEW session's queue lock on a SEPARATE connection BEFORE
// releasing the task row: DeleteTask must not be able to pass the session
// lock in the gap between the two acquisitions.
lockTx2, err := dbC.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin second lock tx: %v", err)
}
defer func() { _ = lockTx2.Rollback() }()
if _, err := lockTx2.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('sess-new')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx2.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'sess-new' FOR UPDATE
`); err != nil {
t.Fatalf("lock new session: %v", err)
}
// A new session commits while DeleteTask waits on the task row.
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO task_sessions (id, task_id, started_at, updated_at)
VALUES ('sess-new', 'task-del-sess-race', now(), now())
`); err != nil {
t.Fatalf("insert session mid-delete: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit session creation: %v", err)
}
// With the fix the purge captured the new session (capture happens after
// the barrier) and blocks on the held lock; without the fix the stale
// capture misses it and DeleteTask completes, so the barrier times out.
waitForWaitingLocks(t, lockTx2, delPID, 1, "DeleteTask purge on the mid-flight session")
if err := lockTx2.Commit(); err != nil {
t.Fatalf("commit second lock tx: %v", err)
}
if err := <-delDone; err != nil {
t.Fatalf("delete task: %v", err)
}
var count int
if err := dbA.GetContext(ctx, &count, `SELECT count(*) FROM tasks WHERE id = 'task-del-sess-race'`); err != nil {
t.Fatalf("count tasks: %v", err)
}
if count != 0 {
t.Fatalf("task survived delete: count=%d", count)
}
}
// TestPostgresRepository_WorkspaceCascade_SerializesWithTaskCreation proves
// the cascade locks the workspace row BEFORE inventorying its tasks: a task
// created mid-cascade is either visible to the inventory (and purged) or
// blocks until the cascade commits. The competing backend holds the workspace
// row, the cascade blocks on it, a new task+session commits, and the purge
// must then block on that new session's lock (proving the inventory saw it).
func TestPostgresRepository_WorkspaceCascade_SerializesWithTaskCreation(t *testing.T) {
repoA, repoB, dbC := newTaskPostgresRepoPair(t)
ctx := context.Background()
const (
wsID = "ws-cascade-sess-race"
taskID = "task-cascade-sess-race"
)
seedTaskWithSession(t, repoA, taskID, wsID, "sess-cascade-sess-race")
dbA := repoA.db
lockTx, err := dbA.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin lock tx: %v", err)
}
defer func() { _ = lockTx.Rollback() }()
if _, err := lockTx.ExecContext(ctx, `
SELECT id FROM workspaces WHERE id = 'ws-cascade-sess-race' FOR UPDATE
`); err != nil {
t.Fatalf("lock workspace row: %v", err)
}
cascadePID := pgBackendPID(t, repoB.db)
cascadeDone := make(chan error, 1)
go func() {
_, _, err := repoB.deleteWorkspaceCascade(ctx, wsID, nil, nil)
cascadeDone <- err
}()
waitForWaitingLocks(t, lockTx, cascadePID, 1, "cascade on the workspace row")
// Hold the new task's session queue lock on a SEPARATE connection BEFORE
// releasing the workspace row: the cascade must not be able to pass the
// session lock in the gap between the two acquisitions.
lockTx2, err := dbC.BeginTxx(ctx, nil)
if err != nil {
t.Fatalf("begin second lock tx: %v", err)
}
defer func() { _ = lockTx2.Rollback() }()
if _, err := lockTx2.ExecContext(ctx, `
INSERT INTO queue_session_locks (session_id) VALUES ('sess-cascade-new')
ON CONFLICT(session_id) DO NOTHING
`); err != nil {
t.Fatalf("ensure session lock row: %v", err)
}
if _, err := lockTx2.ExecContext(ctx, `
SELECT 1 FROM queue_session_locks WHERE session_id = 'sess-cascade-new' FOR UPDATE
`); err != nil {
t.Fatalf("lock new session: %v", err)
}
// A new task with a session commits while the cascade waits on the
// workspace row.
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO tasks (id, workspace_id, workflow_id, workflow_step_id, title, created_at, updated_at)
VALUES ('task-cascade-new', 'ws-cascade-sess-race', 'wf', 'step', 'new', now(), now())
`); err != nil {
t.Fatalf("insert task mid-cascade: %v", err)
}
if _, err := lockTx.ExecContext(ctx, `
INSERT INTO task_sessions (id, task_id, started_at, updated_at)
VALUES ('sess-cascade-new', 'task-cascade-new', now(), now())
`); err != nil {
t.Fatalf("insert session mid-cascade: %v", err)
}
if err := lockTx.Commit(); err != nil {
t.Fatalf("commit task creation: %v", err)
}
// With the fix the cascade's inventory (taken after the workspace lock)
// includes the new task and its purge blocks on the held lock; without the
// fix the stale inventory misses it and the cascade completes, so the
// barrier times out.
waitForWaitingLocks(t, lockTx2, cascadePID, 1, "cascade purge on the mid-cascade task session")
if err := lockTx2.Commit(); err != nil {
t.Fatalf("commit second lock tx: %v", err)
}
if err := <-cascadeDone; err != nil {
t.Fatalf("workspace cascade: %v", err)
}
var count int
if err := dbA.GetContext(ctx, &count, `SELECT count(*) FROM workspaces WHERE id = 'ws-cascade-sess-race'`); err != nil {
t.Fatalf("count workspaces: %v", err)
}
if count != 0 {
t.Fatalf("workspace survived cascade: count=%d", count)
}
}
@@ -26,6 +26,11 @@ func TestPostgresUpdateTaskWithWorkflowStepAdmission_ConcurrentLastSlot(t *testi
t.Fatalf("init postgres schema: %v", err)
}
ctx := context.Background()
// The creation/admission paths lock the workspace row before the
// workflow-step locks; the tasks below reference it, so it must exist.
if err := repo.CreateWorkspace(ctx, &models.Workspace{ID: "postgres-wip-workspace", Name: "Postgres WIP workspace"}); err != nil {
t.Fatalf("seed workspace: %v", err)
}
if _, err := db.Exec(db.Rebind(`
INSERT INTO workflow_steps (id, workflow_id, name, position)
VALUES (?, ?, ?, ?)
@@ -5,6 +5,7 @@ import (
"database/sql"
"errors"
"fmt"
"sort"
"time"
"github.com/google/uuid"
@@ -189,6 +190,15 @@ func (r *Repository) deleteWorkspaceCascade(
}
defer func() { _ = tx.Rollback() }()
// Lock the workspace row BEFORE inventorying its tasks: task creation
// takes the same lock, so a task created after this point either commits
// before the inventory (and is purged with the rest) or blocks until the
// cascade commits, when the workspace is gone and its insert fails. An
// unlocked inventory could miss a task created mid-cascade and leave its
// queued messages orphaned after the task rows are deleted.
if err := r.lockWorkspaceRowInTx(ctx, tx, id); err != nil {
return nil, nil, err
}
if expectedName != nil {
if err := r.confirmWorkspaceNameForCascadeDelete(ctx, tx, id, *expectedName); err != nil {
return nil, nil, err
@@ -202,6 +212,23 @@ func (r *Repository) deleteWorkspaceCascade(
if err != nil {
return nil, nil, err
}
// Establish the global lock order task-row -> queue-session before
// purging the queues: lifecycle admission takes the task row first and
// then the session lock, so taking session locks first here would invert
// the order and deadlock the two on Postgres. Guard every affected task
// row in stable sorted id order, WITHOUT reordering the inventory returned
// to the caller (listWorkspaceCascadeDeleteTasks orders created_at ASC,
// id ASC and the service consumes that order).
lockIDs := make([]string, len(tasks))
for i, task := range tasks {
lockIDs[i] = task.ID
}
sort.Strings(lockIDs)
for _, taskID := range lockIDs {
if err := r.lockTaskRowInTx(ctx, tx, taskID); err != nil {
return nil, nil, fmt.Errorf("guard cascade task row %s: %w", taskID, err)
}
}
if err := r.purgeWorkspaceTaskQueuesInTx(ctx, tx, tasks); err != nil {
return nil, nil, err
}
@@ -255,7 +282,13 @@ func (r *Repository) deleteWorkspaceCascade(
func (r *Repository) purgeWorkspaceTaskQueuesInTx(ctx context.Context, tx *sqlx.Tx, tasks []*models.Task) error {
for _, task := range tasks {
if err := r.purgeTaskQueueInTx(ctx, tx, task.ID); err != nil {
// The tasks still exist at this point (deletion happens later in the
// cascade), so the authoritative session set is discoverable now.
sessions, err := r.taskQueueSessionsInTx(ctx, tx, task.ID)
if err != nil {
return fmt.Errorf("task queue sessions for cascade task %s: %w", task.ID, err)
}
if err := r.purgeTaskQueueInTx(ctx, tx, task.ID, sessions); err != nil {
return fmt.Errorf("purge task queue for workspace cascade task %s: %w", task.ID, err)
}
}
@@ -5,12 +5,18 @@ import { SessionPage } from "../../pages/session-page";
const FULL_QUEUE_SIZE = 10;
export type FullQueueSeed = {
taskId: string;
sessionId: string;
session: SessionPage;
};
export async function seedFullQueueTask(
page: Page,
apiClient: ApiClient,
seedData: SeedData,
title: string,
): Promise<SessionPage> {
): Promise<FullQueueSeed> {
const task = await apiClient.createTask(seedData.workspaceId, title, {
description: "Queue scrolling regression",
workflow_id: seedData.workflowId,
@@ -34,7 +40,7 @@ export async function seedFullQueueTask(
await page.goto(`/t/${task.id}`);
const session = new SessionPage(page);
await session.waitForLoad();
return session;
return { taskId: task.id, sessionId, session };
}
export async function expectFullQueueScrolls(session: SessionPage): Promise<void> {
+30 -1
View File
@@ -225,7 +225,7 @@ test.describe("Task session queue", () => {
apiClient,
seedData,
}) => {
const session = await seedFullQueueTask(
const { session } = await seedFullQueueTask(
testPage,
apiClient,
seedData,
@@ -271,6 +271,35 @@ test.describe("Task session queue", () => {
await expect(entries.nth(2)).toHaveText("separate fourth");
});
test("automatic merge accepts a compatible message into a full queue", async ({
testPage,
apiClient,
seedData,
}) => {
const { taskId, sessionId, session } = await seedFullQueueTask(
testPage,
apiClient,
seedData,
"Full queue auto merge",
);
// The queue is at capacity (10 separate rows, auto-merge off per spec
// isolation). Enabling automatic merge must fold the next compatible
// message into the tail instead of rejecting it as "queue full".
await requestMessageQueueSettings(apiClient, "PATCH", { auto_merge_enabled: true });
await apiClient.queueMessage(taskId, sessionId, "folded while full");
const chat = session.activeChat();
const chip = chat.getByTestId("queue-chip");
await expect(chip).toBeVisible({ timeout: 10_000 });
await chip.click();
const panel = chat.getByTestId("queued-ghost-list");
const entries = panel.getByTestId("queue-entry-text");
await expect(entries).toHaveCount(10);
await expect(entries.last()).toContainText("Queued item 10");
await expect(entries.last()).toContainText("folded while full");
});
test("queue message via submit button on task session page", async ({
testPage,
apiClient,
@@ -48,7 +48,12 @@ test("mobile full queue stays usable while removing and clearing messages", asyn
apiClient,
seedData,
}) => {
const session = await seedFullQueueTask(testPage, apiClient, seedData, "Mobile queue management");
const { session } = await seedFullQueueTask(
testPage,
apiClient,
seedData,
"Mobile queue management",
);
await expectFullQueueScrolls(session);