multica/server/internal/service/task_cancel_finalize_test.go

package service

import (
	"context"
	"fmt"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/multica-ai/multica/server/internal/events"
	"github.com/multica-ai/multica/server/internal/util"
	db "github.com/multica-ai/multica/server/pkg/db/generated"
	"github.com/multica-ai/multica/server/pkg/protocol"
)

func newCancelFinalizePool(t *testing.T) *pgxpool.Pool {
	t.Helper()

	dbURL := os.Getenv("DATABASE_URL")
	if dbURL == "" {
		dbURL = "postgres://multica:multica@localhost:5432/multica?sslmode=disable"
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	pool, err := pgxpool.New(ctx, dbURL)
	if err != nil {
		t.Skipf("database unavailable: %v", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		t.Skipf("database unreachable: %v", err)
	}
	t.Cleanup(pool.Close)
	return pool
}

type cancelFinalizeFixture struct {
	pool          *pgxpool.Pool
	workspaceID   string
	agentID       string
	chatSessionID string
	taskID        string
	userMessageID string
}

// createCancelFinalizeFixture seeds a chat task in the given status with one
// triggering user chat_message bound to it. started marks the task as picked
// up by a daemon (started_at set).
func createCancelFinalizeFixture(t *testing.T, ctx context.Context, pool *pgxpool.Pool, status string, started bool) cancelFinalizeFixture {
	t.Helper()

	suffix := time.Now().UnixNano()
	email := fmt.Sprintf("cancel-finalize-%d@multica.ai", suffix)
	slug := fmt.Sprintf("cancel-finalize-%d", suffix)

	var userID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO "user" (name, email)
		VALUES ($1, $2)
		RETURNING id
	`, "Cancel Finalize Test", email).Scan(&userID); err != nil {
		t.Fatalf("create user: %v", err)
	}

	var workspaceID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO workspace (name, slug, description, issue_prefix)
		VALUES ($1, $2, $3, $4)
		RETURNING id
	`, "Cancel Finalize Test", slug, "temporary cancel finalize test workspace", "CFT").Scan(&workspaceID); err != nil {
		t.Fatalf("create workspace: %v", err)
	}
	if _, err := pool.Exec(ctx, `
		INSERT INTO member (workspace_id, user_id, role)
		VALUES ($1, $2, 'owner')
	`, workspaceID, userID); err != nil {
		t.Fatalf("create member: %v", err)
	}

	var runtimeID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO agent_runtime (
			workspace_id, daemon_id, name, runtime_mode, provider,
			status, device_info, metadata, last_seen_at, visibility, owner_id
		)
		VALUES ($1, NULL, $2, 'cloud', 'cancel_finalize_test', 'online', 'test runtime', '{}'::jsonb, now(), 'private', $3)
		RETURNING id
	`, workspaceID, "Cancel Finalize Runtime", userID).Scan(&runtimeID); err != nil {
		t.Fatalf("create runtime: %v", err)
	}

	var agentID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO agent (
			workspace_id, name, description, runtime_mode, runtime_config,
			runtime_id, visibility, max_concurrent_tasks, owner_id
		)
		VALUES ($1, $2, '', 'cloud', '{}'::jsonb, $3, 'private', 1, $4)
		RETURNING id
	`, workspaceID, "Cancel Finalize Agent", runtimeID, userID).Scan(&agentID); err != nil {
		t.Fatalf("create agent: %v", err)
	}

	var chatSessionID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO chat_session (workspace_id, agent_id, creator_id, title)
		VALUES ($1, $2, $3, 'cancel finalize test')
		RETURNING id
	`, workspaceID, agentID, userID).Scan(&chatSessionID); err != nil {
		t.Fatalf("create chat session: %v", err)
	}

	startedAt := "NULL"
	if started {
		startedAt = "now()"
	}
	var taskID string
	if err := pool.QueryRow(ctx, fmt.Sprintf(`
		INSERT INTO agent_task_queue (agent_id, chat_session_id, status, priority, context, runtime_id, started_at)
		VALUES ($1, $2, $3, 0, '{}'::jsonb, $4, %s)
		RETURNING id
	`, startedAt), agentID, chatSessionID, status, runtimeID).Scan(&taskID); err != nil {
		t.Fatalf("create task: %v", err)
	}

	var userMessageID string
	if err := pool.QueryRow(ctx, `
		INSERT INTO chat_message (chat_session_id, role, content, task_id)
		VALUES ($1, 'user', 'run the thing', $2)
		RETURNING id
	`, chatSessionID, taskID).Scan(&userMessageID); err != nil {
		t.Fatalf("create user chat message: %v", err)
	}

	t.Cleanup(func() {
		cleanupCtx := context.Background()
		pool.Exec(cleanupCtx, `DELETE FROM task_message WHERE task_id = $1`, taskID)
		pool.Exec(cleanupCtx, `DELETE FROM agent_task_queue WHERE agent_id = $1`, agentID)
		pool.Exec(cleanupCtx, `DELETE FROM chat_message WHERE chat_session_id = $1`, chatSessionID)
		pool.Exec(cleanupCtx, `DELETE FROM chat_session WHERE id = $1`, chatSessionID)
		pool.Exec(cleanupCtx, `DELETE FROM agent WHERE id = $1`, agentID)
		pool.Exec(cleanupCtx, `DELETE FROM agent_runtime WHERE id = $1`, runtimeID)
		pool.Exec(cleanupCtx, `DELETE FROM member WHERE workspace_id = $1 AND user_id = $2`, workspaceID, userID)
		pool.Exec(cleanupCtx, `DELETE FROM workspace WHERE id = $1`, workspaceID)
		pool.Exec(cleanupCtx, `DELETE FROM "user" WHERE id = $1`, userID)
	})

	return cancelFinalizeFixture{
		pool:          pool,
		workspaceID:   workspaceID,
		agentID:       agentID,
		chatSessionID: chatSessionID,
		taskID:        taskID,
		userMessageID: userMessageID,
	}
}

func (f cancelFinalizeFixture) insertTranscriptRow(t *testing.T, ctx context.Context) {
	t.Helper()
	if _, err := f.pool.Exec(ctx, `
		INSERT INTO task_message (task_id, seq, type, content)
		VALUES ($1, 1, 'text', 'partial output')
	`, f.taskID); err != nil {
		t.Fatalf("insert task message: %v", err)
	}
}

func (f cancelFinalizeFixture) chatFinalizeDeferredAt(t *testing.T, ctx context.Context) *time.Time {
	t.Helper()
	var deferredAt *time.Time
	if err := f.pool.QueryRow(ctx, `
		SELECT chat_finalize_deferred_at FROM agent_task_queue WHERE id = $1
	`, f.taskID).Scan(&deferredAt); err != nil {
		t.Fatalf("read chat_finalize_deferred_at: %v", err)
	}
	return deferredAt
}

func (f cancelFinalizeFixture) userMessageExists(t *testing.T, ctx context.Context) bool {
	t.Helper()
	var count int
	if err := f.pool.QueryRow(ctx, `
		SELECT count(*) FROM chat_message WHERE id = $1
	`, f.userMessageID).Scan(&count); err != nil {
		t.Fatalf("count user message: %v", err)
	}
	return count > 0
}

// attachUserMessageAttachment binds an attachment row to the fixture's user
// message, mirroring a prompt sent with a file.
func (f cancelFinalizeFixture) attachUserMessageAttachment(t *testing.T, ctx context.Context) string {
	t.Helper()
	var attachmentID string
	if err := f.pool.QueryRow(ctx, `
		INSERT INTO attachment (
			workspace_id, chat_session_id, chat_message_id,
			uploader_type, uploader_id, filename, url, content_type, size_bytes
		)
		SELECT $1, $2, $3, 'member', creator_id, 'notes.txt', 'https://files.test/notes.txt', 'text/plain', 12
		FROM chat_session WHERE id = $2
		RETURNING id
	`, f.workspaceID, f.chatSessionID, f.userMessageID).Scan(&attachmentID); err != nil {
		t.Fatalf("create attachment: %v", err)
	}
	return attachmentID
}

func (f cancelFinalizeFixture) draftRestores(t *testing.T, ctx context.Context) []db.ChatDraftRestore {
	t.Helper()
	rows, err := db.New(f.pool).ListChatDraftRestoresBySession(ctx, util.MustParseUUID(f.chatSessionID))
	if err != nil {
		t.Fatalf("list draft restores: %v", err)
	}
	return rows
}

func (f cancelFinalizeFixture) assistantMessages(t *testing.T, ctx context.Context) []string {
	t.Helper()
	rows, err := f.pool.Query(ctx, `
		SELECT content FROM chat_message WHERE task_id = $1 AND role = 'assistant' ORDER BY created_at
	`, f.taskID)
	if err != nil {
		t.Fatalf("list assistant messages: %v", err)
	}
	defer rows.Close()
	var contents []string
	for rows.Next() {
		var c string
		if err := rows.Scan(&c); err != nil {
			t.Fatalf("scan assistant message: %v", err)
		}
		contents = append(contents, c)
	}
	return contents
}

// eventRecorder captures chat:cancel_finalized events published on the bus.
type eventRecorder struct {
	mu     sync.Mutex
	events []events.Event
}

func recordCancelFinalizedEvents(bus *events.Bus) *eventRecorder {
	rec := &eventRecorder{}
	bus.Subscribe(protocol.EventChatCancelFinalized, func(e events.Event) {
		rec.mu.Lock()
		defer rec.mu.Unlock()
		rec.events = append(rec.events, e)
	})
	return rec
}

func (r *eventRecorder) snapshot() []events.Event {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]events.Event(nil), r.events...)
}

// Cancelled before any daemon started the task: the transcript can never gain
// rows, so the empty judgment is final and the draft restore stays synchronous.
func TestCancelTask_NotStarted_RestoresDraftSynchronously(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "queued", false)
	svc := NewTaskService(db.New(pool), pool, nil, events.New())

	result, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true})
	if err != nil {
		t.Fatalf("cancel task: %v", err)
	}
	if result.CancelledChatMessage == nil {
		t.Fatal("expected a synchronous CancelledChatMessage restore result")
	}
	if !result.CancelledChatMessage.RestoreToInput {
		t.Error("expected RestoreToInput=true")
	}
	if result.CancelledChatMessage.Content != "run the thing" {
		t.Errorf("restored content = %q, want %q", result.CancelledChatMessage.Content, "run the thing")
	}
	if f.userMessageExists(t, ctx) {
		t.Error("user chat message should have been deleted")
	}
	if got := f.chatFinalizeDeferredAt(t, ctx); got != nil {
		t.Errorf("chat_finalize_deferred_at should stay NULL, got %v", got)
	}
}

// Cancelled after start with transcript rows already persisted: "non-empty"
// is monotonic (late rows only append), so the Stopped. outcome stays
// synchronous.
func TestCancelTask_StartedNonEmptyTranscript_StopsSynchronously(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	f.insertTranscriptRow(t, ctx)
	svc := NewTaskService(db.New(pool), pool, nil, events.New())

	result, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true})
	if err != nil {
		t.Fatalf("cancel task: %v", err)
	}
	if result.CancelledChatMessage != nil {
		t.Fatalf("expected no restore result, got %+v", result.CancelledChatMessage)
	}
	if got := f.assistantMessages(t, ctx); len(got) != 1 || got[0] != "Stopped." {
		t.Errorf("assistant messages = %v, want [Stopped.]", got)
	}
	if !f.userMessageExists(t, ctx) {
		t.Error("user chat message should be preserved")
	}
	if got := f.chatFinalizeDeferredAt(t, ctx); got != nil {
		t.Errorf("chat_finalize_deferred_at should stay NULL, got %v", got)
	}
}

// Cancelled after start with an empty transcript: the daemon may still be
// flushing, so no judgment is made — the task is marked deferred instead.
func TestCancelTask_StartedEmptyTranscript_DefersJudgment(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	svc := NewTaskService(db.New(pool), pool, nil, events.New())

	result, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true})
	if err != nil {
		t.Fatalf("cancel task: %v", err)
	}
	if result.CancelledChatMessage != nil {
		t.Fatalf("expected no synchronous restore result, got %+v", result.CancelledChatMessage)
	}
	if !f.userMessageExists(t, ctx) {
		t.Error("user chat message must not be deleted while deferred")
	}
	if got := f.assistantMessages(t, ctx); len(got) != 0 {
		t.Errorf("no assistant message should exist while deferred, got %v", got)
	}
	if got := f.chatFinalizeDeferredAt(t, ctx); got == nil {
		t.Error("chat_finalize_deferred_at should be set")
	}
}

// Deferred finalize with a transcript that stayed empty: the user message is
// deleted, a durable chat_draft_restore row is persisted in the same tx, and
// a content-free restored-outcome event is broadcast as an invalidation hint.
func TestFinalizeDeferredCancelledChat_StillEmpty_RestoresDraft(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	bus := events.New()
	rec := recordCancelFinalizedEvents(bus)
	svc := NewTaskService(db.New(pool), pool, nil, bus)

	if _, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true}); err != nil {
		t.Fatalf("cancel task: %v", err)
	}

	svc.FinalizeDeferredCancelledChat(ctx, util.MustParseUUID(f.taskID))

	if f.userMessageExists(t, ctx) {
		t.Error("user chat message should have been deleted")
	}
	if got := f.chatFinalizeDeferredAt(t, ctx); got != nil {
		t.Errorf("chat_finalize_deferred_at should be cleared, got %v", got)
	}
	restores := f.draftRestores(t, ctx)
	if len(restores) != 1 {
		t.Fatalf("expected 1 chat_draft_restore row, got %d", len(restores))
	}
	if got := util.UUIDToString(restores[0].ID); got != f.userMessageID {
		t.Errorf("restore id = %q, want deleted message id %q", got, f.userMessageID)
	}
	if restores[0].Content != "run the thing" {
		t.Errorf("restore content = %q, want %q", restores[0].Content, "run the thing")
	}
	if got := util.UUIDToString(restores[0].TaskID); got != f.taskID {
		t.Errorf("restore task_id = %q, want %q", got, f.taskID)
	}
	evts := rec.snapshot()
	if len(evts) != 1 {
		t.Fatalf("expected 1 chat:cancel_finalized event, got %d", len(evts))
	}
	payload, ok := evts[0].Payload.(protocol.ChatCancelFinalizedPayload)
	if !ok {
		t.Fatalf("payload type = %T", evts[0].Payload)
	}
	if payload.Outcome != protocol.ChatCancelOutcomeRestored {
		t.Errorf("outcome = %q, want restored", payload.Outcome)
	}
	if payload.MessageID != f.userMessageID {
		t.Errorf("message_id = %q, want %q", payload.MessageID, f.userMessageID)
	}
	if payload.Content != "" {
		t.Errorf("restored event must not carry the prompt content, got %q", payload.Content)
	}
	if payload.ChatSessionID != f.chatSessionID {
		t.Errorf("chat_session_id = %q, want %q", payload.ChatSessionID, f.chatSessionID)
	}
}

// The recoverability guarantee behind the restored outcome: the draft-restore
// row commits with the settlement tx, so a client that never receives the
// chat:cancel_finalized broadcast (offline, or the server dies between commit
// and publish — simulated here by a bus with no subscribers) still finds the
// restore, with its detached attachment ids, on the next fetch.
func TestFinalizeDeferredCancelledChat_BroadcastLost_RestoreIsRecoverable(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	attachmentID := f.attachUserMessageAttachment(t, ctx)
	queries := db.New(pool)
	svc := NewTaskService(queries, pool, nil, events.New())

	if _, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true}); err != nil {
		t.Fatalf("cancel task: %v", err)
	}

	svc.FinalizeDeferredCancelledChat(ctx, util.MustParseUUID(f.taskID))

	restores, err := queries.ListChatDraftRestoresBySession(ctx, util.MustParseUUID(f.chatSessionID))
	if err != nil {
		t.Fatalf("list draft restores: %v", err)
	}
	if len(restores) != 1 {
		t.Fatalf("expected 1 recoverable draft restore, got %d", len(restores))
	}
	if restores[0].Content != "run the thing" {
		t.Errorf("restore content = %q, want %q", restores[0].Content, "run the thing")
	}
	if len(restores[0].AttachmentIds) != 1 || util.UUIDToString(restores[0].AttachmentIds[0]) != attachmentID {
		t.Errorf("restore attachment_ids = %v, want [%s]", restores[0].AttachmentIds, attachmentID)
	}
	// The attachment row itself must have survived the message delete
	// (detached, so the restored draft can re-bind it on re-send).
	var chatMessageID *string
	if err := pool.QueryRow(ctx, `
		SELECT chat_message_id FROM attachment WHERE id = $1
	`, attachmentID).Scan(&chatMessageID); err != nil {
		t.Fatalf("attachment row should survive: %v", err)
	}
	if chatMessageID != nil {
		t.Errorf("attachment should be detached, still bound to %v", *chatMessageID)
	}

	// Consume is idempotent: first delete claims the row, second is a no-op.
	for i, want := range []int64{1, 0} {
		n, err := queries.DeleteChatDraftRestore(ctx, db.DeleteChatDraftRestoreParams{
			ID:            restores[0].ID,
			ChatSessionID: util.MustParseUUID(f.chatSessionID),
		})
		if err != nil {
			t.Fatalf("consume draft restore (call %d): %v", i+1, err)
		}
		if n != want {
			t.Errorf("consume rows (call %d) = %d, want %d", i+1, n, want)
		}
	}
}

// Deferred finalize after transcript rows landed late: the judgment flips to
// non-empty, a Stopped. row is written, and a stopped-outcome event is
// broadcast.
func TestFinalizeDeferredCancelledChat_RowsLanded_WritesStopped(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	bus := events.New()
	rec := recordCancelFinalizedEvents(bus)
	svc := NewTaskService(db.New(pool), pool, nil, bus)

	if _, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true}); err != nil {
		t.Fatalf("cancel task: %v", err)
	}
	// The daemon's late flush lands after the cancel commit.
	f.insertTranscriptRow(t, ctx)

	svc.FinalizeDeferredCancelledChat(ctx, util.MustParseUUID(f.taskID))

	if !f.userMessageExists(t, ctx) {
		t.Error("user chat message should be preserved")
	}
	if got := f.assistantMessages(t, ctx); len(got) != 1 || got[0] != "Stopped." {
		t.Errorf("assistant messages = %v, want [Stopped.]", got)
	}
	if got := f.chatFinalizeDeferredAt(t, ctx); got != nil {
		t.Errorf("chat_finalize_deferred_at should be cleared, got %v", got)
	}
	evts := rec.snapshot()
	if len(evts) != 1 {
		t.Fatalf("expected 1 chat:cancel_finalized event, got %d", len(evts))
	}
	payload, ok := evts[0].Payload.(protocol.ChatCancelFinalizedPayload)
	if !ok {
		t.Fatalf("payload type = %T", evts[0].Payload)
	}
	if payload.Outcome != protocol.ChatCancelOutcomeStopped {
		t.Errorf("outcome = %q, want stopped", payload.Outcome)
	}
	if payload.Content != "Stopped." {
		t.Errorf("content = %q, want Stopped.", payload.Content)
	}
	if payload.MessageID == "" {
		t.Error("message_id should carry the new assistant row id")
	}
}

// The marker is an atomic claim: a second finalize call must be a no-op.
func TestFinalizeDeferredCancelledChat_SecondCallIsNoop(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	bus := events.New()
	rec := recordCancelFinalizedEvents(bus)
	svc := NewTaskService(db.New(pool), pool, nil, bus)

	if _, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true}); err != nil {
		t.Fatalf("cancel task: %v", err)
	}
	f.insertTranscriptRow(t, ctx)

	svc.FinalizeDeferredCancelledChat(ctx, util.MustParseUUID(f.taskID))
	svc.FinalizeDeferredCancelledChat(ctx, util.MustParseUUID(f.taskID))

	if got := f.assistantMessages(t, ctx); len(got) != 1 {
		t.Errorf("assistant messages = %v, want exactly one Stopped.", got)
	}
	if got := len(rec.snapshot()); got != 1 {
		t.Errorf("events = %d, want 1", got)
	}
}

// Sweeper query: only markers older than the grace period are returned.
func TestListChatFinalizeDeferredExpired_HonorsGrace(t *testing.T) {
	ctx := context.Background()
	pool := newCancelFinalizePool(t)
	f := createCancelFinalizeFixture(t, ctx, pool, "running", true)
	queries := db.New(pool)
	svc := NewTaskService(queries, pool, nil, events.New())

	if _, err := svc.CancelTaskWithResult(ctx, util.MustParseUUID(f.taskID), CancelTaskOptions{ClientSupportsDraftRestore: true}); err != nil {
		t.Fatalf("cancel task: %v", err)
	}

	fresh, err := queries.ListChatFinalizeDeferredExpired(ctx, db.ListChatFinalizeDeferredExpiredParams{
		GraceSecs:  60,
		MaxPerTick: 100,
	})
	if err != nil {
		t.Fatalf("list expired: %v", err)
	}
	for _, row := range fresh {
		if util.UUIDToString(row.ID) == f.taskID {
			t.Fatal("fresh marker must not be returned within the grace period")
		}
	}

	if _, err := pool.Exec(ctx, `
		UPDATE agent_task_queue SET chat_finalize_deferred_at = now() - interval '2 minutes' WHERE id = $1
	`, f.taskID); err != nil {
		t.Fatalf("backdate marker: %v", err)
	}

	expired, err := queries.ListChatFinalizeDeferredExpired(ctx, db.ListChatFinalizeDeferredExpiredParams{
		GraceSecs:  60,
		MaxPerTick: 100,
	})
	if err != nil {
		t.Fatalf("list expired: %v", err)
	}
	found := false
	for _, row := range expired {
		if util.UUIDToString(row.ID) == f.taskID {
			found = true
		}
	}
	if !found {
		t.Error("backdated marker should be returned once past the grace period")
	}
}