multica/server/internal/handler/chat_draft_restore_test.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
)

// withDraftRestoreParams sets both chi URL params in one route context —
// chaining withURLParam would replace the context and drop the first param.
func withDraftRestoreParams(req *http.Request, sessionID, restoreID string) *http.Request {
	rctx := chi.NewRouteContext()
	rctx.URLParams.Add("sessionId", sessionID)
	rctx.URLParams.Add("restoreId", restoreID)
	return req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))
}

// seedDraftRestore inserts a chat_draft_restore row the way the deferred
// finalize tx does: id is the deleted user message's id, attachments already
// detached (chat_message_id NULL).
func seedDraftRestore(t *testing.T, sessionID, content string, attachmentIDs []string) string {
	t.Helper()
	restoreID := uuid.NewString()
	ids := "{}"
	if len(attachmentIDs) > 0 {
		ids = "{"
		for i, id := range attachmentIDs {
			if i > 0 {
				ids += ","
			}
			ids += id
		}
		ids += "}"
	}
	if _, err := testPool.Exec(context.Background(), `
		INSERT INTO chat_draft_restore (id, chat_session_id, task_id, content, attachment_ids)
		VALUES ($1, $2, $3, $4, $5::uuid[])
	`, restoreID, sessionID, uuid.NewString(), content, ids); err != nil {
		t.Fatalf("seed draft restore: %v", err)
	}
	t.Cleanup(func() {
		testPool.Exec(context.Background(), `DELETE FROM chat_draft_restore WHERE id = $1`, restoreID)
	})
	return restoreID
}

func seedDetachedChatAttachment(t *testing.T, sessionID string) string {
	t.Helper()
	var attachmentID string
	if err := testPool.QueryRow(context.Background(), `
		INSERT INTO attachment (
			workspace_id, chat_session_id, uploader_type, uploader_id,
			filename, url, content_type, size_bytes
		)
		VALUES ($1, $2, 'member', $3, 'notes.txt', 'https://files.test/notes.txt', 'text/plain', 12)
		RETURNING id
	`, testWorkspaceID, sessionID, testUserID).Scan(&attachmentID); err != nil {
		t.Fatalf("seed detached attachment: %v", err)
	}
	t.Cleanup(func() {
		testPool.Exec(context.Background(), `DELETE FROM attachment WHERE id = $1`, attachmentID)
	})
	return attachmentID
}

// The recovery path the chat:cancel_finalized broadcast cannot guarantee: a
// creator whose client missed the event fetches the pending restore, gets the
// prompt content plus re-bindable attachments, and consumes it idempotently.
func TestChatDraftRestores_CreatorFetchesAndConsumes(t *testing.T) {
	agentID := createHandlerTestAgent(t, "DraftRestoreAgent", []byte("[]"))
	sessionID := createHandlerTestChatSession(t, agentID)
	attachmentID := seedDetachedChatAttachment(t, sessionID)
	restoreID := seedDraftRestore(t, sessionID, "run the thing", []string{attachmentID})

	listReq := withURLParam(newRequest(http.MethodGet, "/api/chat/sessions/"+sessionID+"/draft-restores", nil), "sessionId", sessionID)
	listReq = withChatTestWorkspaceCtx(t, listReq)
	listW := httptest.NewRecorder()
	testHandler.ListChatDraftRestores(listW, listReq)
	if listW.Code != http.StatusOK {
		t.Fatalf("list: expected 200, got %d: %s", listW.Code, listW.Body.String())
	}
	var resp ChatDraftRestoresResponse
	if err := json.Unmarshal(listW.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode list: %v", err)
	}
	if len(resp.Restores) != 1 {
		t.Fatalf("restores = %d, want 1", len(resp.Restores))
	}
	got := resp.Restores[0]
	if got.ID != restoreID {
		t.Errorf("id = %q, want %q", got.ID, restoreID)
	}
	if got.Content != "run the thing" {
		t.Errorf("content = %q, want %q", got.Content, "run the thing")
	}
	if len(got.Attachments) != 1 || got.Attachments[0].ID != attachmentID {
		t.Fatalf("attachments = %+v, want the seeded attachment", got.Attachments)
	}
	if got.Attachments[0].Filename != "notes.txt" {
		t.Errorf("attachment filename = %q, want notes.txt", got.Attachments[0].Filename)
	}

	// Consume twice: both must be 204 (idempotent), and the row must be gone.
	for i := 0; i < 2; i++ {
		delReq := withDraftRestoreParams(newRequest(http.MethodDelete, "/api/chat/sessions/"+sessionID+"/draft-restores/"+restoreID, nil), sessionID, restoreID)
		delReq = withChatTestWorkspaceCtx(t, delReq)
		delW := httptest.NewRecorder()
		testHandler.ConsumeChatDraftRestore(delW, delReq)
		if delW.Code != http.StatusNoContent {
			t.Fatalf("consume call %d: expected 204, got %d: %s", i+1, delW.Code, delW.Body.String())
		}
	}

	listW = httptest.NewRecorder()
	testHandler.ListChatDraftRestores(listW, withChatTestWorkspaceCtx(t, withURLParam(newRequest(http.MethodGet, "/api/chat/sessions/"+sessionID+"/draft-restores", nil), "sessionId", sessionID)))
	if listW.Code != http.StatusOK {
		t.Fatalf("relist: expected 200, got %d", listW.Code)
	}
	resp = ChatDraftRestoresResponse{}
	if err := json.Unmarshal(listW.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode relist: %v", err)
	}
	if len(resp.Restores) != 0 {
		t.Errorf("restores after consume = %d, want 0", len(resp.Restores))
	}
}

// Draft restores hold the creator's private prompt: another workspace member
// must not be able to read or consume them.
func TestChatDraftRestores_NonCreatorForbidden(t *testing.T) {
	agentID := createHandlerTestAgent(t, "DraftRestoreOtherAgent", []byte("[]"))
	sessionID := createHandlerTestChatSession(t, agentID)
	restoreID := seedDraftRestore(t, sessionID, "secret prompt", nil)

	var otherUserID string
	if err := testPool.QueryRow(context.Background(), `
		INSERT INTO "user" (name, email)
		VALUES ('Draft Restore Other', $1)
		RETURNING id
	`, fmt.Sprintf("draft-restore-other-%d@multica.ai", time.Now().UnixNano())).Scan(&otherUserID); err != nil {
		t.Fatalf("create other user: %v", err)
	}
	if _, err := testPool.Exec(context.Background(), `
		INSERT INTO member (workspace_id, user_id, role)
		VALUES ($1, $2, 'member')
	`, testWorkspaceID, otherUserID); err != nil {
		t.Fatalf("create other member: %v", err)
	}
	t.Cleanup(func() {
		testPool.Exec(context.Background(), `DELETE FROM member WHERE workspace_id = $1 AND user_id = $2`, testWorkspaceID, otherUserID)
		testPool.Exec(context.Background(), `DELETE FROM "user" WHERE id = $1`, otherUserID)
	})

	listReq := withURLParam(newRequest(http.MethodGet, "/api/chat/sessions/"+sessionID+"/draft-restores", nil), "sessionId", sessionID)
	listReq = withChatTestWorkspaceCtx(t, listReq)
	listReq.Header.Set("X-User-ID", otherUserID)
	listW := httptest.NewRecorder()
	testHandler.ListChatDraftRestores(listW, listReq)
	if listW.Code != http.StatusForbidden {
		t.Fatalf("list as non-creator: expected 403, got %d: %s", listW.Code, listW.Body.String())
	}

	delReq := withDraftRestoreParams(newRequest(http.MethodDelete, "/api/chat/sessions/"+sessionID+"/draft-restores/"+restoreID, nil), sessionID, restoreID)
	delReq = withChatTestWorkspaceCtx(t, delReq)
	delReq.Header.Set("X-User-ID", otherUserID)
	delW := httptest.NewRecorder()
	testHandler.ConsumeChatDraftRestore(delW, delReq)
	if delW.Code != http.StatusForbidden {
		t.Fatalf("consume as non-creator: expected 403, got %d: %s", delW.Code, delW.Body.String())
	}

	var count int
	if err := testPool.QueryRow(context.Background(), `SELECT count(*) FROM chat_draft_restore WHERE id = $1`, restoreID).Scan(&count); err != nil {
		t.Fatalf("count restore: %v", err)
	}
	if count != 1 {
		t.Errorf("restore row must survive a non-creator consume attempt, count = %d", count)
	}
}

// chat_draft_restore has no chat_session FK (MUL-3515), so nothing but
// DeleteChatSession's own transaction prunes an unconsumed restore.
func TestDeleteChatSession_PrunesDraftRestores(t *testing.T) {
	agentID := createHandlerTestAgent(t, "DraftRestorePruneAgent", []byte("[]"))
	sessionID := createHandlerTestChatSession(t, agentID)
	seedDraftRestore(t, sessionID, "unconsumed prompt", nil)

	req := withURLParam(newRequest(http.MethodDelete, "/api/chat/sessions/"+sessionID, nil), "sessionId", sessionID)
	req = withChatTestWorkspaceCtx(t, req)
	req.Header.Set("X-User-ID", testUserID)
	w := httptest.NewRecorder()
	testHandler.DeleteChatSession(w, req)
	if w.Code != http.StatusNoContent {
		t.Fatalf("DeleteChatSession: expected 204, got %d: %s", w.Code, w.Body.String())
	}

	var count int
	if err := testPool.QueryRow(context.Background(),
		`SELECT count(*) FROM chat_draft_restore WHERE chat_session_id = $1`, sessionID).Scan(&count); err != nil {
		t.Fatalf("count restores: %v", err)
	}
	if count != 0 {
		t.Errorf("draft restores must be pruned with their session, count = %d", count)
	}
}

func countDraftRestores(t *testing.T, sessionID string) int {
	t.Helper()
	var count int
	if err := testPool.QueryRow(context.Background(),
		`SELECT count(*) FROM chat_draft_restore WHERE chat_session_id = $1`, sessionID).Scan(&count); err != nil {
		t.Fatalf("count restores: %v", err)
	}
	return count
}

// DeleteChatSession is not the only way a chat_session dies: it also cascades
// from agent (migration 033). Since MUL-5559 a runtime delete only hard-deletes
// the runtime's SYSTEM agents, so that is the cascade whose draft restores would
// otherwise be stranded (chat_draft_restore has no FK).
func TestDeleteAgentRuntime_PrunesDraftRestoresOfSystemAgentSessions(t *testing.T) {
	if testHandler == nil {
		t.Skip("database not available")
	}

	runtimeID := seedIsolatedRuntime(t, "Runtime With Draft Restore")
	systemAgent := seedAgentOnRuntime(t, runtimeID, "System Agent With Draft Restore", false)
	if _, err := testPool.Exec(context.Background(),
		`UPDATE agent SET kind = 'system', system_key = 'draft_restore_probe' WHERE id = $1`,
		systemAgent); err != nil {
		t.Fatalf("make agent a system agent: %v", err)
	}
	sessionID := createHandlerTestChatSession(t, systemAgent)
	seedDraftRestore(t, sessionID, "prompt stranded by the agent cascade", nil)

	w := httptest.NewRecorder()
	req := withURLParam(newRequest(http.MethodDelete, "/api/runtimes/"+runtimeID, nil), "runtimeId", runtimeID)
	testHandler.DeleteAgentRuntime(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("DeleteAgentRuntime: expected 200, got %d: %s", w.Code, w.Body.String())
	}

	if agentExists(t, systemAgent) {
		t.Fatal("system agent should have been deleted — the cascade under test never ran")
	}
	if n := countDraftRestores(t, sessionID); n != 0 {
		t.Errorf("draft restores of a cascade-deleted session must be pruned, count = %d", n)
	}
}

// The mirror case, and the regression for the cleanup scope: an ARCHIVED USER
// agent is no longer deleted with its runtime, so its session and the pending
// restore holding the user's prompt must both survive. Pruning by "archived
// agents on this runtime" — what the old teardown did — would throw away the
// draft of an agent that is still there.
func TestDeleteAgentRuntime_KeepsDraftRestoresOfUnboundAgentSessions(t *testing.T) {
	if testHandler == nil {
		t.Skip("database not available")
	}

	runtimeID := seedIsolatedRuntime(t, "Runtime With Surviving Draft Restore")
	archivedAgent := seedAgentOnRuntime(t, runtimeID, "Archived Agent With Surviving Draft", true)
	sessionID := createHandlerTestChatSession(t, archivedAgent)
	seedDraftRestore(t, sessionID, "prompt that must survive the machine being retired", nil)

	w := httptest.NewRecorder()
	req := withURLParam(newRequest(http.MethodDelete, "/api/runtimes/"+runtimeID, nil), "runtimeId", runtimeID)
	testHandler.DeleteAgentRuntime(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("DeleteAgentRuntime: expected 200, got %d: %s", w.Code, w.Body.String())
	}

	if !agentExists(t, archivedAgent) {
		t.Fatal("archived user agent must survive its runtime as an unbound agent")
	}
	if n := countDraftRestores(t, sessionID); n != 1 {
		t.Errorf("draft restore of a surviving agent must be kept, count = %d", n)
	}
}

// The workspace cascade, same class. DeleteWorkspace runs outside a transaction
// (its atomicity comes from being one multi-CTE statement), so the prune has to
// live inside that statement — this is what guards it.
func TestDeleteWorkspace_PrunesDraftRestoresOfCascadedSessions(t *testing.T) {
	if testHandler == nil {
		t.Skip("database not available")
	}
	ctx := context.Background()

	const slug = "handler-tests-delete-draft-restores"
	_, _ = testPool.Exec(ctx, `DELETE FROM workspace WHERE slug = $1`, slug)

	var wsID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO workspace (name, slug, description)
		VALUES ($1, $2, '')
		RETURNING id
	`, "Handler Test Draft Restore Cascade", slug).Scan(&wsID); err != nil {
		t.Fatalf("create workspace: %v", err)
	}
	t.Cleanup(func() {
		testPool.Exec(context.Background(), `DELETE FROM workspace WHERE id = $1`, wsID)
	})
	if _, err := testPool.Exec(ctx, `
		INSERT INTO member (workspace_id, user_id, role)
		VALUES ($1, $2, 'owner')
	`, wsID, testUserID); err != nil {
		t.Fatalf("create owner member: %v", err)
	}

	var runtimeID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO agent_runtime (
			workspace_id, daemon_id, name, runtime_mode, provider, status, device_info, metadata, last_seen_at
		)
		VALUES ($1, NULL, 'Draft Restore Cascade Runtime', 'cloud', 'isolated_test', 'online', '', '{}'::jsonb, now())
		RETURNING id
	`, wsID).Scan(&runtimeID); err != nil {
		t.Fatalf("create runtime: %v", err)
	}

	var agentID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO agent (
			workspace_id, name, description, runtime_mode, runtime_config,
			runtime_id, visibility, max_concurrent_tasks, owner_id
		)
		VALUES ($1, 'Draft Restore Cascade Agent', '', 'cloud', '{}'::jsonb, $2, 'workspace', 1, $3)
		RETURNING id
	`, wsID, runtimeID, testUserID).Scan(&agentID); err != nil {
		t.Fatalf("create agent: %v", err)
	}

	var sessionID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO chat_session (workspace_id, agent_id, creator_id, title, status)
		VALUES ($1, $2, $3, 'Draft Restore Cascade Session', 'active')
		RETURNING id
	`, wsID, agentID, testUserID).Scan(&sessionID); err != nil {
		t.Fatalf("create chat session: %v", err)
	}
	seedDraftRestore(t, sessionID, "prompt stranded by the workspace cascade", nil)

	w := httptest.NewRecorder()
	req := withURLParam(newRequest(http.MethodDelete, "/api/workspaces/"+wsID, nil), "id", wsID)
	testHandler.DeleteWorkspace(w, req)
	if w.Code != http.StatusNoContent {
		t.Fatalf("DeleteWorkspace: expected 204, got %d: %s", w.Code, w.Body.String())
	}

	if n := countDraftRestores(t, sessionID); n != 0 {
		t.Errorf("draft restores of a workspace-cascaded session must be pruned, count = %d", n)
	}
}