multica/server/internal/handler/task_terminal_wakeup_test.go

package handler

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

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

type terminalWakeupRecorder struct {
	calls []struct{ runtimeID, taskID string }
}

func (r *terminalWakeupRecorder) NotifyTaskAvailable(runtimeID, taskID string) {
	r.calls = append(r.calls, struct{ runtimeID, taskID string }{runtimeID, taskID})
}

func failTaskViaHandler(t *testing.T, taskID string) *httptest.ResponseRecorder {
	t.Helper()
	w := httptest.NewRecorder()
	req := newDaemonTokenRequest("POST", "/api/daemon/tasks/"+taskID+"/fail",
		map[string]any{"error": "agent failed", "failure_reason": "agent_error"},
		testWorkspaceID, "legit-daemon")
	rctx := chi.NewRouteContext()
	rctx.URLParams.Add("taskId", taskID)
	req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))
	testHandler.FailTask(w, req)
	return w
}

// TestTerminalTransitionsNotifyRuntime verifies the server-side half of the
// queued-successor fix. Completion wakes only after reconciliation, while fail
// and cancel wake after their terminal side effects, so an old daemon can retry
// its normal atomic claim path without waiting for the 30-second fallback.
func TestTerminalTransitionsNotifyRuntime(t *testing.T) {
	if testHandler == nil || testPool == nil {
		t.Skip("database not available")
	}
	ctx := context.Background()

	var agentID, runtimeID string
	if err := testPool.QueryRow(ctx,
		`SELECT id, runtime_id FROM agent WHERE workspace_id = $1 AND runtime_id IS NOT NULL LIMIT 1`,
		testWorkspaceID).Scan(&agentID, &runtimeID); err != nil {
		t.Fatalf("setup: get agent: %v", err)
	}

	var issueID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO issue (workspace_id, title, status, priority, creator_id, creator_type, number, position, assignee_type, assignee_id)
		VALUES ($1, 'terminal wakeup fixture', 'in_progress', 'none', $2, 'member', 999099, 0, 'agent', $3)
		RETURNING id
	`, testWorkspaceID, testUserID, agentID).Scan(&issueID); err != nil {
		t.Fatalf("setup: create issue: %v", err)
	}
	t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM issue WHERE id = $1`, issueID) })

	var taskID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, created_at, started_at)
		VALUES ($1, $2, $3, 'running', 0, now() - interval '2 minutes', now() - interval '1 minute')
		RETURNING id
	`, agentID, runtimeID, issueID).Scan(&taskID); err != nil {
		t.Fatalf("setup: running task: %v", err)
	}

	recorder := &terminalWakeupRecorder{}
	previous := testHandler.TaskService.Wakeup
	testHandler.TaskService.Wakeup = recorder
	t.Cleanup(func() { testHandler.TaskService.Wakeup = previous })

	if w := completeTaskViaHandler(t, taskID, "done"); w.Code != http.StatusOK {
		t.Fatalf("CompleteTask: expected 200, got %d: %s", w.Code, w.Body.String())
	}
	if got := len(recorder.calls); got != 1 {
		t.Fatalf("expected 1 terminal runtime wakeup, got %d", got)
	}
	if recorder.calls[0].runtimeID != runtimeID {
		t.Fatalf("wakeup runtime = %q, want %q", recorder.calls[0].runtimeID, runtimeID)
	}
	if recorder.calls[0].taskID != "" {
		t.Fatalf("terminal wakeup must omit completed task id, got %q", recorder.calls[0].taskID)
	}

	// Failure uses the daemon callback handler, matching production. It must
	// send the same runtime-only hint after FailTask commits.
	recorder.calls = nil
	var failedTaskID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, created_at, started_at)
		VALUES ($1, $2, $3, 'running', 0, now() - interval '2 minutes', now() - interval '1 minute')
		RETURNING id
	`, agentID, runtimeID, issueID).Scan(&failedTaskID); err != nil {
		t.Fatalf("setup: failed task: %v", err)
	}
	if w := failTaskViaHandler(t, failedTaskID); w.Code != http.StatusOK {
		t.Fatalf("FailTask: expected 200, got %d: %s", w.Code, w.Body.String())
	}
	if got := len(recorder.calls); got != 1 || recorder.calls[0].runtimeID != runtimeID || recorder.calls[0].taskID != "" {
		t.Fatalf("unexpected failure wakeups: %#v", recorder.calls)
	}

	// Cancellation is finalized in TaskService because both issue and chat
	// handlers share that path.
	recorder.calls = nil
	var cancelledTaskID string
	if err := testPool.QueryRow(ctx, `
		INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, created_at, started_at)
		VALUES ($1, $2, $3, 'running', 0, now() - interval '2 minutes', now() - interval '1 minute')
		RETURNING id
	`, agentID, runtimeID, issueID).Scan(&cancelledTaskID); err != nil {
		t.Fatalf("setup: cancelled task: %v", err)
	}
	cancelled, err := testHandler.TaskService.CancelTask(ctx, parseUUID(cancelledTaskID))
	if err != nil {
		t.Fatalf("CancelTask: %v", err)
	}
	if cancelled.Status != "cancelled" {
		t.Fatalf("CancelTask status = %q, want cancelled", cancelled.Status)
	}
	if got := len(recorder.calls); got != 1 || recorder.calls[0].runtimeID != runtimeID || recorder.calls[0].taskID != "" {
		t.Fatalf("unexpected cancellation wakeups: %#v", recorder.calls)
	}
}