query-orchestration/internal/bot/service_turn_test.go

// Package bot — Milestone 4 service-layer concurrency tests for AddTurn.
//
// Plan §11 specifies three concurrency invariants:
//
//  1. Ten concurrent creates for the same next ordinal yield exactly one
//     persisted in_flight (or completed once FastAPI returns) row. Other
//     callers receive 409 responses, never 500.
//  2. Concurrent calls that both observe a grace-expired row are
//     idempotent on abandonment. At most one new turn proceeds; losing
//     callers receive a documented 409.
//  3. A stale tx2 callback after abandonment or retry affects zero rows
//     because both attempt_id and status='in_flight' reject it
//     (turn_superseded).
//
// Compile contract: this file references bot.Service.AddTurn,
// bot.AddTurnInput, bot.AddTurnResult, bot.ErrTurnInFlight,
// bot.ErrPriorTurnInFlight, bot.ErrOrdinalOutOfRange, bot.ErrTurnSuperseded,
// and the M4-only sqlc queries via the existing GetDBQueries() chain.
// Backend-eng adds service_turn.go and the package compiles. Until
// then this file fails to compile — the failing state plan §10 M4
// endorses.
package bot_test

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"queryorchestration/internal/bot"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/chatbot"
	"queryorchestration/internal/test"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// turnTestCfg returns a chatbot config with a short request timeout so
// timeout-sensitive scenarios do not block the test for 60s.
func turnTestCfg(baseURL string) chatbot.ChatbotConfig {
	return chatbot.ChatbotConfig{
		ServiceURL:            baseURL,
		APIKey:                "test-api-key",
		RequestTimeoutSeconds: 1,
		MaxTurnChars:          2000,
		RecentTurns:           5,
	}
}

// resetForTurnTests does the same idempotency dance the M2/M3 helpers
// do but lives in this file so the bot package's tests do not pull in
// queryapi_test helpers.
func resetForTurnTests(t *testing.T, ctx context.Context, dbCfg *serviceconfig.BaseConfig, clientID, name string) {
	t.Helper()
	pool := dbCfg.GetDBPool()
	stmts := []string{
		`DELETE FROM bot_session_documents WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
		`DELETE FROM bot_session_folders WHERE folder_id IN (SELECT id FROM folders WHERE clientId = $1)`,
		`DELETE FROM bot_sessions WHERE client_id = $1`,
		`DELETE FROM documents WHERE clientId = $1`,
		`DELETE FROM folders WHERE clientId = $1`,
		`DELETE FROM clients WHERE clientId = $1`,
	}
	for _, stmt := range stmts {
		_, err := pool.Exec(ctx, stmt, clientID)
		require.NoErrorf(t, err, "reset stmt failed: %s", stmt)
	}
	require.NoError(t, dbCfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Name:     name,
		Clientid: clientID,
	}))
}

// fastAPIServerForTurnTests returns a minimal FastAPI test server whose
// /agent/chat handler emits a deterministic 200 success body. The chosen
// answer text echoes the supplied prefix so concurrent callers can
// distinguish their results.
func fastAPIServerForTurnTests(t *testing.T, answerText string) (*httptest.Server, *int32) {
	t.Helper()
	var hits int32
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"status":"ok","service":"aaria-chatbot","version":"0.1.0"}`))
	})
	mux.HandleFunc("/agent/chat", func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&hits, 1)
		w.WriteHeader(http.StatusOK)
		// Build the response body inline so we do not need a separate
		// helper. The test only checks status + persisted state, not
		// the wire body shape (that is covered by M3 tests).
		body := []byte(`{"request_id":"any","status":"success","answer":{"text":"` + answerText + `","result_type":"summary","confidence":0.9}}`)
		_, _ = w.Write(body)
	})
	server := httptest.NewServer(mux)
	t.Cleanup(server.Close)
	return server, &hits
}

// seedSessionDirect inserts a bot_sessions row via raw SQL. Used by the
// concurrency tests that need a session id without exercising the M2
// CreateSession path.
func seedSessionDirect(t *testing.T, ctx context.Context, dbCfg *serviceconfig.BaseConfig, clientID, actor string) uuid.UUID {
	t.Helper()
	var id uuid.UUID
	err := dbCfg.GetDBPool().QueryRow(ctx, `
		INSERT INTO bot_sessions (client_id, created_by) VALUES ($1, $2) RETURNING id
	`, clientID, actor).Scan(&id)
	require.NoError(t, err)
	return id
}

// seedTerminalTurn inserts a bot_turns row in a terminal status
// (errored, abandoned, session_deleted). The CHECK constraint
// bot_turns_status_fields_consistent requires `error` to be non-null
// for these statuses; this helper supplies a placeholder JSON value so
// callers do not have to repeat the boilerplate.
func seedTerminalTurn(t *testing.T, ctx context.Context, dbCfg *serviceconfig.BaseConfig,
	sessionID uuid.UUID, ordinal int, status, prompt string) {
	t.Helper()
	_, err := dbCfg.GetDBPool().Exec(ctx, `
		INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, error, attempt_started_at, created_at)
		VALUES ($1, $2, $3, $4, $5, $6::jsonb, NOW() - INTERVAL '1 hour', NOW() - INTERVAL '1 hour')
	`, sessionID, ordinal, prompt, status, uuid.New(), `{"code":"x","message":"y"}`)
	require.NoError(t, err)
}

// TestAddTurn_TenConcurrentCreates_OneWinner: ten goroutines POST
// ordinal 1 simultaneously to a fresh session. The partial unique
// index `bot_turns_one_inflight_per_session` and the FOR UPDATE lock
// in LockBotSessionForOwner serialize the writes so exactly one row
// persists (in flight or completed). The other nine callers receive a
// 409-class typed error; nobody gets 500.
func TestAddTurn_TenConcurrentCreates_OneWinner(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CONC_10"
		actor    = "concurrent@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-conc-10")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	server, _ := fastAPIServerForTurnTests(t, "concurrent answer")
	cfg := turnTestCfg(server.URL)

	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	const N = 10
	type result struct {
		err error
	}
	results := make([]result, N)

	var wg sync.WaitGroup
	wg.Add(N)
	for i := 0; i < N; i++ {
		idx := i
		go func() {
			defer wg.Done()
			_, err := svc.AddTurn(ctx, bot.AddTurnInput{
				SessionID: sessionID,
				ClientID:  clientID,
				Actor:     actor,
				Ordinal:   1,
				Prompt:    "concurrent prompt",
			})
			results[idx] = result{err: err}
		}()
	}
	wg.Wait()

	// Exactly one bot_turns row persists for this session.
	var rowCount int
	err = dbCfg.GetDBPool().QueryRow(ctx, `
		SELECT COUNT(*) FROM bot_turns WHERE session_id = $1
	`, sessionID).Scan(&rowCount)
	require.NoError(t, err)
	assert.Equal(t, 1, rowCount,
		"plan §11: 10 concurrent ordinal-1 creates -> exactly 1 persisted row")

	// Exactly one caller succeeded; the other nine got typed 409 errors.
	wins := 0
	losses := 0
	for i, r := range results {
		switch {
		case r.err == nil:
			wins++
		case errors.Is(r.err, bot.ErrTurnInFlight),
			errors.Is(r.err, bot.ErrPriorTurnInFlight),
			errors.Is(r.err, bot.ErrOrdinalOutOfRange):
			losses++
		default:
			t.Errorf("call %d returned non-409 error: %v", i, r.err)
		}
	}
	assert.Equal(t, 1, wins, "exactly one goroutine must succeed")
	assert.Equal(t, N-1, losses, "the other nine must receive typed 409 errors, not 500")
}

// TestAddTurn_GraceExpiredIdempotent: two callers concurrently observe
// the same expired in_flight row. One wins; the other receives a
// documented 409 (turn_in_flight or prior_turn_in_flight depending on
// observed ordering); no double-abandon, no leaked in_flight.
func TestAddTurn_GraceExpiredIdempotent(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CONC_GR"
		actor    = "grace@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-conc-gr")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	// Seed expired in_flight at ordinal 1. Default grace = 2 *
	// RequestTimeoutSeconds = 2s; 600s is well outside.
	expiredAt := time.Now().UTC().Add(-600 * time.Second)
	_, err := dbCfg.GetDBPool().Exec(ctx, `
		INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, attempt_started_at, created_at)
		VALUES ($1, 1, 'stuck', 'in_flight', $2, $3, $3)
	`, sessionID, uuid.New(), expiredAt)
	require.NoError(t, err)

	server, _ := fastAPIServerForTurnTests(t, "after-grace answer")
	cfg := turnTestCfg(server.URL)

	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	type result struct {
		err error
	}
	results := make([]result, 2)
	var wg sync.WaitGroup
	wg.Add(2)
	for i := 0; i < 2; i++ {
		idx := i
		go func() {
			defer wg.Done()
			_, err := svc.AddTurn(ctx, bot.AddTurnInput{
				SessionID: sessionID,
				ClientID:  clientID,
				Actor:     actor,
				Ordinal:   2,
				Prompt:    "after grace prompt",
			})
			results[idx] = result{err: err}
		}()
	}
	wg.Wait()

	// Exactly one ordinal-2 row persists (in_flight or completed).
	var ord2Count int
	err = dbCfg.GetDBPool().QueryRow(ctx, `
		SELECT COUNT(*) FROM bot_turns WHERE session_id = $1 AND ordinal = 2
	`, sessionID).Scan(&ord2Count)
	require.NoError(t, err)
	assert.Equal(t, 1, ord2Count,
		"plan §11: at most one new turn proceeds")

	// Original ordinal 1 must be exactly one row, status=abandoned.
	var ord1Status string
	var ord1Count int
	err = dbCfg.GetDBPool().QueryRow(ctx, `
		SELECT COUNT(*), MAX(status) FROM bot_turns WHERE session_id = $1 AND ordinal = 1
	`, sessionID).Scan(&ord1Count, &ord1Status)
	require.NoError(t, err)
	assert.Equal(t, 1, ord1Count, "no double-abandon — ordinal 1 has exactly one row")
	assert.Equal(t, "abandoned", ord1Status, "expired in_flight must be abandoned")

	// Exactly one caller succeeded; the other got a typed 409.
	wins := 0
	losses := 0
	for _, r := range results {
		switch {
		case r.err == nil:
			wins++
		case errors.Is(r.err, bot.ErrTurnInFlight),
			errors.Is(r.err, bot.ErrPriorTurnInFlight),
			errors.Is(r.err, bot.ErrOrdinalOutOfRange):
			losses++
		default:
			t.Errorf("unexpected error type: %v", r.err)
		}
	}
	assert.Equal(t, 1, wins)
	assert.Equal(t, 1, losses)
}

// TestAddTurn_StaleTx2CallbackZeroRows: a stale tx2 callback (whose
// attempt_id no longer matches the persisted row) must affect zero
// rows and the service must return ErrTurnSuperseded. Plan §11.
//
// We exercise this by seeding an in_flight row with attempt_id A,
// simulating that another caller has already abandoned + retried
// (status now = in_flight with attempt_id B), then directly invoking
// the service-level CompleteBotTurn helper that the AddTurn tx2 path
// uses. The helper must use compare-and-set on attempt_id and report
// 0 rows affected -> ErrTurnSuperseded.
func TestAddTurn_StaleTx2CallbackZeroRows(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CONC_STALE"
		actor    = "stale@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-conc-stale")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	// Simulated state: ordinal 1 was originally in_flight with
	// attempt_id A; another caller has since reset it and the row now
	// holds attempt_id B in status='in_flight'. The stale tx2
	// callback from attempt A will try to complete the turn.
	staleAttemptID := uuid.New() // A — what the stale callback thinks
	currentAttemptID := uuid.New()
	_, err := dbCfg.GetDBPool().Exec(ctx, `
		INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id)
		VALUES ($1, 1, 'p', 'in_flight', $2)
	`, sessionID, currentAttemptID)
	require.NoError(t, err)

	server, _ := fastAPIServerForTurnTests(t, "stale answer")
	cfg := turnTestCfg(server.URL)
	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	// CompleteBotTurnWithAttempt runs the same compare-and-set the
	// AddTurn tx2 path uses. Affected rows = 0 because attempt_id
	// does not match.
	err = svc.CompleteBotTurnWithAttempt(ctx, bot.CompleteTurnInput{
		SessionID:  sessionID,
		Ordinal:    1,
		AttemptID:  staleAttemptID,
		Completion: "answer from stale call",
		LatencyMs:  100,
		TokensIn:   5,
		TokensOut:  10,
	})
	require.Error(t, err)
	assert.True(t, errors.Is(err, bot.ErrTurnSuperseded),
		"plan §11: stale tx2 callback must return ErrTurnSuperseded; got %v", err)

	// Row must remain in_flight with the current (non-stale) attempt_id.
	var status string
	var attemptID uuid.UUID
	err = dbCfg.GetDBPool().QueryRow(ctx, `
		SELECT status, attempt_id FROM bot_turns WHERE session_id = $1 AND ordinal = 1
	`, sessionID).Scan(&status, &attemptID)
	require.NoError(t, err)
	assert.Equal(t, "in_flight", status,
		"row must remain in_flight after stale callback")
	assert.Equal(t, currentAttemptID, attemptID,
		"row's attempt_id must remain the current value (B), not the stale (A)")
}

// ---------- classifyExistingTurn branch coverage ----------
//
// The next three tests exercise the errored/abandoned arms of
// service_turn.go classifyExistingTurn that the public HTTP handler
// cannot reach because FindTargetOrdinalForPrompt pre-resolves the
// target ordinal. They call svc.AddTurn directly with crafted DB state
// so coverage on these defensive arms does not depend on integration
// flakes elsewhere in the suite.

// TestAddTurn_ExistingErroredOrdinalMismatch_OutOfRange covers the
// errored/abandoned -> input.Ordinal != session.LastTerminalOrdinal arm.
//
// Setup: ord 1 errored ("X"), ord 2 completed -> LastTerminalOrdinal = 2.
// AddTurn(Ordinal=1) lands on the errored row at ord 1, but 1 != 2 so
// classifyExistingTurn returns ErrOrdinalOutOfRange before touching
// the prompt-match check.
func TestAddTurn_ExistingErroredOrdinalMismatch_OutOfRange(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CET_ORD"
		actor    = "ord-mismatch@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-cet-ord")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	seedTerminalTurn(t, ctx, dbCfg, sessionID, 1, "errored", "X")
	_, err := dbCfg.GetDBPool().Exec(ctx, `
		INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, completion, attempt_started_at, created_at)
		VALUES ($1, 2, 'second', 'completed', $2, 'a', NOW(), NOW())
	`, sessionID, uuid.New())
	require.NoError(t, err)

	server, _ := fastAPIServerForTurnTests(t, "unused")
	cfg := turnTestCfg(server.URL)
	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	_, err = svc.AddTurn(ctx, bot.AddTurnInput{
		SessionID: sessionID,
		ClientID:  clientID,
		Actor:     actor,
		Ordinal:   1,
		Prompt:    "X",
	})
	require.Error(t, err)
	assert.Truef(t, errors.Is(err, bot.ErrOrdinalOutOfRange),
		"expected ErrOrdinalOutOfRange when ord != LastTerminalOrdinal in errored arm; got %v", err)
}

// TestAddTurn_ExistingErroredPromptMismatch covers the
// errored/abandoned -> existing.Prompt != input.Prompt arm.
//
// Setup: ord 1 errored ("X"), no other turns -> LastTerminalOrdinal = 1.
// AddTurn(Ordinal=1, Prompt="Y") clears the ord check (1 == 1) then
// trips the prompt-match check ("X" != "Y").
func TestAddTurn_ExistingErroredPromptMismatch(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CET_PRM"
		actor    = "prompt-mismatch@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-cet-prm")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	seedTerminalTurn(t, ctx, dbCfg, sessionID, 1, "errored", "X")

	server, _ := fastAPIServerForTurnTests(t, "unused")
	cfg := turnTestCfg(server.URL)
	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	_, err = svc.AddTurn(ctx, bot.AddTurnInput{
		SessionID: sessionID,
		ClientID:  clientID,
		Actor:     actor,
		Ordinal:   1,
		Prompt:    "Y",
	})
	require.Error(t, err)
	assert.Truef(t, errors.Is(err, bot.ErrPromptMismatch),
		"expected ErrPromptMismatch when existing.Prompt != input.Prompt; got %v", err)
}

// TestAddTurn_RetryRejectedByInflightIndex
// covers the errored/abandoned -> ResetBotTurnForRetry unique-violation
// arm (isUniqueViolation -> ErrPriorTurnInFlight).
//
// Setup: ord 1 errored ("X"), ord 2 in_flight (fresh, so the grace-
// abandon pass leaves it alone). LastTerminalOrdinal = 1 because in_flight
// rows are excluded from the MAX. AddTurn(Ordinal=1, Prompt="X") clears
// the ord and prompt checks, then ResetBotTurnForRetry tries to flip ord
// 1 to in_flight — the partial unique index
// bot_turns_one_inflight_per_session rejects the second in_flight row
// for this session, surfacing SQLSTATE 23505.
func TestAddTurn_RetryRejectedByInflightIndex(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	ctx := t.Context()
	dbCfg := &serviceconfig.BaseConfig{}
	test.CreateDB(t, dbCfg)

	const (
		clientID = "BOT_M4_CET_PIF"
		actor    = "prior-inflight@example.com"
	)
	resetForTurnTests(t, ctx, dbCfg, clientID, "bot-m4-cet-pif")
	sessionID := seedSessionDirect(t, ctx, dbCfg, clientID, actor)

	seedTerminalTurn(t, ctx, dbCfg, sessionID, 1, "errored", "X")
	_, err := dbCfg.GetDBPool().Exec(ctx, `
		INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, attempt_started_at, created_at)
		VALUES ($1, 2, 'second', 'in_flight', $2, NOW(), NOW())
	`, sessionID, uuid.New())
	require.NoError(t, err)

	server, _ := fastAPIServerForTurnTests(t, "unused")
	cfg := turnTestCfg(server.URL)
	svc, err := bot.New(cfg, dbCfg)
	require.NoError(t, err)

	_, err = svc.AddTurn(ctx, bot.AddTurnInput{
		SessionID: sessionID,
		ClientID:  clientID,
		Actor:     actor,
		Ordinal:   1,
		Prompt:    "X",
	})
	require.Error(t, err)
	assert.Truef(t, errors.Is(err, bot.ErrPriorTurnInFlight),
		"expected ErrPriorTurnInFlight when reset trips partial unique index; got %v", err)

	// Sanity: the errored row at ord 1 must not have been flipped.
	var ord1Status string
	err = dbCfg.GetDBPool().QueryRow(ctx, `
		SELECT status FROM bot_turns WHERE session_id = $1 AND ordinal = 1
	`, sessionID).Scan(&ord1Status)
	require.NoError(t, err)
	assert.Equal(t, "errored", ord1Status,
		"reset must roll back when the unique index rejects it; ord 1 stays errored")
}