nomadcode/services/core/internal/scheduler/jobs_test.go

package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"testing"
	"time"

	"github.com/riverqueue/river"

	"github.com/nomadcode/nomadcode-core/internal/agent"
	"github.com/nomadcode/nomadcode-core/internal/model"
	"github.com/nomadcode/nomadcode-core/internal/notification"
	"github.com/nomadcode/nomadcode-core/internal/storage"
	"github.com/nomadcode/nomadcode-core/internal/workflow"
)

func TestRunAgentTaskCompletesFromMessage(t *testing.T) {
	worker := &TaskWorker{
		Agent: fakeAgentClient{
			result: agent.SendMessageResult{
				Message: &agent.Message{
					Kind:      "message",
					Role:      "agent",
					MessageID: "msg-1",
					Parts:     []agent.Part{{Kind: "text", Text: "done"}},
				},
				Raw: json.RawMessage(`{"kind":"message","messageId":"msg-1"}`),
			},
		},
	}

	raw, message, err := worker.runAgentTask(context.Background(), storage.Task{
		ID:      "task-1",
		Title:   "Fix issue",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"please fix it"}`),
	})
	if err != nil {
		t.Fatalf("runAgentTask returned error: %v", err)
	}
	if message != "done" {
		t.Fatalf("unexpected message: %q", message)
	}

	var output map[string]any
	if err := json.Unmarshal(raw, &output); err != nil {
		t.Fatalf("decode output: %v", err)
	}
	if output["mode"] != "a2a" || output["message"] != "done" || output["message_id"] != "msg-1" {
		t.Fatalf("unexpected output: %#v", output)
	}
}

func TestRunAgentTaskRejectsNonTerminalTask(t *testing.T) {
	worker := &TaskWorker{
		Agent: fakeAgentClient{
			result: agent.SendMessageResult{
				Task: &agent.Task{
					Kind:   "task",
					ID:     "remote-1",
					Status: agent.TaskStatus{State: agent.TaskStateWorking},
				},
			},
		},
	}

	_, _, err := worker.runAgentTask(context.Background(), storage.Task{
		ID:      "task-1",
		Title:   "Fix issue",
		Source:  "manual",
		Payload: json.RawMessage(`{}`),
	})
	if err == nil || !strings.Contains(err.Error(), "non-terminal") {
		t.Fatalf("expected non-terminal state error, got %v", err)
	}
}

func TestBuildAgentInputUsesPromptPayload(t *testing.T) {
	input := buildAgentInput(storage.Task{
		ID:      "task-1",
		Title:   "Fallback title",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"do the thing","instructions":"be concise"}`),
	})

	if !input.Blocking {
		t.Fatal("expected blocking A2A call")
	}
	if !strings.Contains(input.Text, "be concise") || !strings.Contains(input.Text, "do the thing") {
		t.Fatalf("unexpected text: %q", input.Text)
	}
	if input.Metadata["nomadcode_task_id"] != "task-1" {
		t.Fatalf("unexpected metadata: %#v", input.Metadata)
	}
}

func TestBuildGenerateInputIncludesTaskMetadata(t *testing.T) {
	input := buildGenerateInput(storage.Task{
		ID:      "task-1",
		Title:   "Fallback title",
		Source:  "plane",
		Payload: json.RawMessage(`{"prompt":"do the thing","instructions":"be concise"}`),
	})

	if input.Input != "do the thing" {
		t.Fatalf("unexpected input: %q", input.Input)
	}
	if input.Instructions != "be concise" {
		t.Fatalf("unexpected instructions: %q", input.Instructions)
	}
	if input.Metadata["task_id"] != "task-1" || input.Metadata["source"] != "plane" {
		t.Fatalf("unexpected metadata: %#v", input.Metadata)
	}
}

type fakeAgentClient struct {
	result agent.SendMessageResult
	err    error
}

func (c fakeAgentClient) SendMessage(context.Context, agent.SendMessageInput) (agent.SendMessageResult, error) {
	return c.result, c.err
}

func (c fakeAgentClient) GetTask(context.Context, agent.GetTaskInput) (agent.Task, error) {
	return agent.Task{}, nil
}

func (c fakeAgentClient) CancelTask(context.Context, agent.CancelTaskInput) (agent.Task, error) {
	return agent.Task{}, nil
}

type fakeTaskLifecycle struct {
	started        []string
	completed      []string
	failed         []string
	failMessages   []string
	failInputs     []workflow.FailureInput
	mergedMetadata []map[string]any

	task storage.Task
}

func (f *fakeTaskLifecycle) StartTask(ctx context.Context, id string) (storage.Task, error) {
	if f.task.Status != "running" && f.task.Status != "pending" && f.task.Status != "queued" && f.task.Status != "failed" && f.task.Status != "" {
		return storage.Task{}, errors.New("invalid transition to running")
	}
	f.started = append(f.started, id)
	f.task.Status = "running"
	attempt := len(f.started)
	// Merge attempt into existing metadata rather than replacing it, so that
	// checkout metadata (used for authoring task detection) is preserved.
	var existing map[string]any
	if len(f.task.Metadata) > 0 && string(f.task.Metadata) != "null" {
		_ = json.Unmarshal(f.task.Metadata, &existing)
	}
	if existing == nil {
		existing = make(map[string]any)
	}
	existing["attempt"] = attempt
	f.task.Metadata, _ = json.Marshal(existing)
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) CompleteTask(ctx context.Context, id string, result json.RawMessage) (storage.Task, error) {
	if f.task.Status != "running" {
		return storage.Task{}, errors.New("invalid transition to completed")
	}
	f.completed = append(f.completed, id)
	f.task.Status = "completed"
	attempt := len(f.started)
	f.task.Metadata = json.RawMessage(fmt.Sprintf(`{"attempt":%d}`, attempt))
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) FailTaskWithMetadata(ctx context.Context, id string, input workflow.FailureInput) (storage.Task, error) {
	if f.task.Status != "running" {
		return storage.Task{}, errors.New("invalid transition to failed")
	}
	f.failed = append(f.failed, id)
	f.failMessages = append(f.failMessages, input.Message)
	f.failInputs = append(f.failInputs, input)
	f.task.Status = "failed"
	attempt := len(f.started)
	f.task.Metadata = json.RawMessage(fmt.Sprintf(`{"attempt":%d,"status_reason":"%s"}`, attempt, input.Message))
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
	return f.FailTaskWithMetadata(ctx, id, workflow.FailureInput{
		Message: message,
		Type:    workflow.FailureTypeExecution,
	})
}

func (f *fakeTaskLifecycle) MergeTaskMetadata(_ context.Context, _ string, updates map[string]any) (storage.Task, error) {
	f.mergedMetadata = append(f.mergedMetadata, updates)
	return f.task, nil
}

func TestWorkCompletesThroughLifecycle(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 1 || fakeLifecycle.completed[0] != "task-1" {
		t.Fatalf("expected completed to have task-1, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 0 {
		t.Fatalf("expected no failed tasks, got %v", fakeLifecycle.failed)
	}
}

func TestWorkFailsThroughLifecycle(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
		Agent: fakeAgentClient{
			err: errors.New("agent failed"),
		},
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return error")
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected no completed tasks, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 1 || fakeLifecycle.failed[0] != "task-1" {
		t.Fatalf("expected failed to have task-1, got %v", fakeLifecycle.failed)
	}
}

type blockingModelClient struct{}

func (m blockingModelClient) Generate(ctx context.Context, input model.GenerateInput) (model.GenerateResult, error) {
	<-ctx.Done()
	return model.GenerateResult{}, ctx.Err()
}

func TestWorkMarksTimeoutFailure(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:  fakeLifecycle,
		Model:      blockingModelClient{},
		RunTimeout: 1 * time.Millisecond,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return timeout error")
	}
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("expected context.DeadlineExceeded error, got: %v", err)
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected no completed tasks, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 1 || fakeLifecycle.failed[0] != "task-1" {
		t.Fatalf("expected failed to have task-1, got %v", fakeLifecycle.failed)
	}
	if len(fakeLifecycle.failMessages) != 1 || !strings.Contains(fakeLifecycle.failMessages[0], "timeout") {
		t.Fatalf("expected fail message containing timeout, got %v", fakeLifecycle.failMessages)
	}
	if len(fakeLifecycle.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
	}
	if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeTimeout {
		t.Errorf("expected failure type %q, got %q", workflow.FailureTypeTimeout, fakeLifecycle.failInputs[0].Type)
	}
}

func TestWorkRetriesAfterFailureState(t *testing.T) {
	sink := &spySink{}
	notifService := notification.NewService(nil, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Status:  "pending",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Agent: fakeAgentClient{
			err: errors.New("first attempt agent failure"),
		},
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected first Work to return error")
	}

	if fakeLifecycle.task.Status != "failed" {
		t.Fatalf("expected task status to be failed, got: %s", fakeLifecycle.task.Status)
	}

	if len(fakeLifecycle.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
	}
	if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeExecution {
		t.Errorf("expected failure type %q, got %q", workflow.FailureTypeExecution, fakeLifecycle.failInputs[0].Type)
	}

	worker.Agent = nil

	err = worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("expected second Work to succeed, got error: %v", err)
	}

	if fakeLifecycle.task.Status != "completed" {
		t.Fatalf("expected task status to be completed, got: %s", fakeLifecycle.task.Status)
	}

	if len(fakeLifecycle.started) != 2 {
		t.Fatalf("expected StartTask to be called 2 times, got: %d", len(fakeLifecycle.started))
	}
	if len(fakeLifecycle.failed) != 1 {
		t.Fatalf("expected FailTask to be called 1 time, got: %d", len(fakeLifecycle.failed))
	}
	if len(fakeLifecycle.completed) != 1 {
		t.Fatalf("expected CompleteTask to be called 1 time, got: %d", len(fakeLifecycle.completed))
	}

	// Verify the emitted notification event sequence and context fields (preventing stale reason leak)
	if len(sink.events) != 4 {
		t.Fatalf("expected 4 notification events fanned out, got %d: %+v", len(sink.events), sink.events)
	}

	// First attempt Running event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("ev0: expected type Running, got %s", ev0.Type)
	}
	if ev0.Attempt != 1 {
		t.Errorf("ev0: expected attempt 1, got %d", ev0.Attempt)
	}
	if ev0.Reason != "" {
		t.Errorf("ev0: expected empty reason (no stale failure leak), got %q", ev0.Reason)
	}

	// First attempt Failed event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventFailed {
		t.Errorf("ev1: expected type Failed, got %s", ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("ev1: expected attempt 1, got %d", ev1.Attempt)
	}
	if ev1.Reason != "first attempt agent failure" {
		t.Errorf("ev1: expected reason 'first attempt agent failure', got %q", ev1.Reason)
	}

	// Second attempt Running event
	ev2 := sink.events[2]
	if ev2.Type != notification.TaskEventRunning {
		t.Errorf("ev2: expected type Running, got %s", ev2.Type)
	}
	if ev2.Attempt != 2 {
		t.Errorf("ev2: expected attempt 2, got %d", ev2.Attempt)
	}
	if ev2.Reason != "" {
		t.Errorf("ev2: expected empty reason (prevented stale failed reason leak), got %q", ev2.Reason)
	}

	// Second attempt Completed event
	ev3 := sink.events[3]
	if ev3.Type != notification.TaskEventCompleted {
		t.Errorf("ev3: expected type Completed, got %s", ev3.Type)
	}
	if ev3.Attempt != 2 {
		t.Errorf("ev3: expected attempt 2, got %d", ev3.Attempt)
	}
	if ev3.Reason != "" {
		t.Errorf("ev3: expected empty reason (prevented stale failed reason leak), got %q", ev3.Reason)
	}
}

func checkoutTaskMeta(slotPath, sourceBranch string) json.RawMessage {
	raw, _ := json.Marshal(map[string]any{
		"checkout": map[string]any{
			"slot_path":     slotPath,
			"source_branch": sourceBranch,
			"slot_index":    0,
			"slot_id":       1,
		},
	})
	return raw
}

func TestRunTaskRoutesPlaneOriginToAuthoringBridge(t *testing.T) {
	var capturedInput model.GenerateInput

	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			capturedInput = input
			return model.GenerateResult{
				ID:    "resp-1",
				Model: "m",
				Text:  "authoring done",
			}, nil
		},
	}

	worker := &TaskWorker{Model: fakeModel}

	task := storage.Task{
		ID:       "task-plane-1",
		Title:    "NOMAD-42: Create milestone",
		Source:   "plane",
		Payload:  json.RawMessage(`{"prompt":"Author milestone for NOMAD-42","instructions":"Use roadmap skill"}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	raw, msg, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if msg == "" {
		t.Fatal("expected non-empty message")
	}
	if raw == nil {
		t.Fatal("expected non-nil result")
	}

	if capturedInput.WorkspaceMetadata == nil {
		t.Fatal("expected WorkspaceMetadata to be set for Plane-origin authoring task")
	}
	if capturedInput.WorkspaceMetadata.Path != "/home/user/workspace/nomadcode/slots/000" {
		t.Errorf("unexpected workspace path: %q", capturedInput.WorkspaceMetadata.Path)
	}
	if capturedInput.WorkspaceMetadata.SourceBranch != "develop" {
		t.Errorf("unexpected source branch: %q", capturedInput.WorkspaceMetadata.SourceBranch)
	}
	if capturedInput.WorkspaceMetadata.Provider != "plane" {
		t.Errorf("unexpected provider: %q", capturedInput.WorkspaceMetadata.Provider)
	}
	if capturedInput.WorkspaceMetadata.WorkItemID != "task-plane-1" {
		t.Errorf("unexpected work_item_id: %q", capturedInput.WorkspaceMetadata.WorkItemID)
	}
}

func TestRunTaskAuthoringPreemptsA2AClient(t *testing.T) {
	var modelCalled bool
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			modelCalled = true
			if input.WorkspaceMetadata == nil {
				t.Error("expected WorkspaceMetadata for authoring task routed through model")
			}
			return model.GenerateResult{Text: "done", Model: "m"}, nil
		},
	}
	// A2A client is set — authoring task must still go to model bridge.
	worker := &TaskWorker{
		Model: fakeModel,
		Agent: fakeAgentClient{err: errors.New("a2a should not be called")},
	}

	task := storage.Task{
		ID:       "task-preempt",
		Source:   "plane",
		Payload:  json.RawMessage(`{}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	raw, _, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if !modelCalled {
		t.Error("expected model client to be called for authoring task, not A2A")
	}
	var result map[string]any
	_ = json.Unmarshal(raw, &result)
	if result["mode"] != "authoring_run" {
		t.Errorf("expected mode=authoring_run, got %#v", result["mode"])
	}
}

func TestRunTaskGenericModelHasNoWorkspaceMetadata(t *testing.T) {
	var capturedInput model.GenerateInput

	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			capturedInput = input
			return model.GenerateResult{Text: "done"}, nil
		},
	}

	worker := &TaskWorker{Model: fakeModel}

	task := storage.Task{
		ID:      "task-manual-1",
		Title:   "Generic task",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"do something"}`),
	}

	if _, _, err := worker.runTask(context.Background(), task); err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if capturedInput.WorkspaceMetadata != nil {
		t.Errorf("expected no WorkspaceMetadata for generic task, got %#v", capturedInput.WorkspaceMetadata)
	}
}

type fakeModelClient struct {
	generate func(context.Context, model.GenerateInput) (model.GenerateResult, error)
}

func (f fakeModelClient) Generate(ctx context.Context, input model.GenerateInput) (model.GenerateResult, error) {
	return f.generate(ctx, input)
}

func TestRunTaskAuthoringSuccessRecordsState(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "milestone authored", Model: "m"}, nil
		},
	}
	worker := &TaskWorker{Model: fakeModel}

	task := storage.Task{
		ID:       "task-authoring-ok",
		Title:    "Author milestone",
		Source:   "plane",
		Payload:  json.RawMessage(`{}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	raw, _, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}

	var result map[string]any
	if err := json.Unmarshal(raw, &result); err != nil {
		t.Fatalf("unmarshal result: %v", err)
	}
	if result["authoring_run_state"] != "in_progress" {
		t.Errorf("expected authoring_run_state=in_progress, got %#v", result["authoring_run_state"])
	}
	if result["authoring_run_updated_at"] == nil || result["authoring_run_updated_at"] == "" {
		t.Error("expected authoring_run_updated_at to be set")
	}
	if result["mode"] != "authoring_run" {
		t.Errorf("expected mode=authoring_run, got %#v", result["mode"])
	}
}

func TestWorkRecordsAuthoringInProgressBeforeExecution(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "done", Model: "m"}, nil
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-auth-ip",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
		Model:     fakeModel,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-ip"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	// MergeTaskMetadata must have been called at least once with in_progress.
	found := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyAuthoringRunState] == "in_progress" {
			found = true
			if m[workflow.MetadataKeyAuthoringRunUpdatedAt] == nil {
				t.Error("expected authoring_run_updated_at to be set with in_progress")
			}
		}
	}
	if !found {
		t.Errorf("expected MergeTaskMetadata to be called with authoring_run_state=in_progress, got %v", fakeLifecycle.mergedMetadata)
	}
}

func TestWorkDoesNotRecordAuthoringInProgressForGenericTask(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:     "task-generic-ip",
			Source: "manual",
			Status: "pending",
		},
	}

	worker := &TaskWorker{Lifecycle: fakeLifecycle}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-generic-ip"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyAuthoringRunState] != nil {
			t.Errorf("expected no authoring metadata for generic task, got %v", m)
		}
	}
}

func TestMarkFailedWithTaskAuthoringRecordsState(t *testing.T) {
	lc := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-authoring-fail",
			Title:    "Author milestone",
			Source:   "plane",
			Status:   "running",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		},
	}
	lc.task.Metadata = checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop")

	worker := &TaskWorker{Lifecycle: lc}

	task := storage.Task{
		ID:       "task-authoring-fail",
		Source:   "plane",
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	worker.markFailedWithTask("task-authoring-fail", task, errors.New("IOP bridge failed"))

	if len(lc.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(lc.failInputs))
	}
	extra := lc.failInputs[0].ExtraMetadata
	if extra == nil {
		t.Fatal("expected ExtraMetadata to be set for authoring task failure")
	}
	if extra[workflow.MetadataKeyAuthoringRunState] != "failed" {
		t.Errorf("expected authoring_run_state=failed, got %#v", extra[workflow.MetadataKeyAuthoringRunState])
	}
	if extra[workflow.MetadataKeyAuthoringRunUpdatedAt] == nil {
		t.Error("expected authoring_run_updated_at to be set")
	}
	if extra[workflow.MetadataKeyAuthoringFailureCategory] != "bridge_failed" {
		t.Errorf("expected authoring_failure_category=bridge_failed, got %#v", extra[workflow.MetadataKeyAuthoringFailureCategory])
	}
}

func TestMarkFailedWithTaskGenericNoAuthoringState(t *testing.T) {
	lc := &fakeTaskLifecycle{
		task: storage.Task{
			ID:     "task-generic-fail",
			Source: "manual",
			Status: "running",
		},
	}
	worker := &TaskWorker{Lifecycle: lc}

	task := storage.Task{
		ID:     "task-generic-fail",
		Source: "manual",
	}

	worker.markFailedWithTask("task-generic-fail", task, errors.New("generic error"))

	if len(lc.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(lc.failInputs))
	}
	if lc.failInputs[0].ExtraMetadata != nil {
		t.Errorf("expected no ExtraMetadata for generic task, got %#v", lc.failInputs[0].ExtraMetadata)
	}
}

type logRecord struct {
	msg  string
	args map[string]any
}

type spyHandler struct {
	records []logRecord
}

func (h *spyHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return true
}

func (h *spyHandler) Handle(ctx context.Context, r slog.Record) error {
	args := make(map[string]any)
	r.Attrs(func(a slog.Attr) bool {
		args[a.Key] = a.Value.Any()
		return true
	})
	h.records = append(h.records, logRecord{
		msg:  r.Message,
		args: args,
	})
	return nil
}

func (h *spyHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return h
}

func (h *spyHandler) WithGroup(name string) slog.Handler {
	return h
}

type spySink struct {
	events []notification.TaskEvent
}

func (s *spySink) HandleTaskEvent(ctx context.Context, event notification.TaskEvent) error {
	s.events = append(s.events, event)
	return nil
}

func TestWorkEmitsRunningAndCompletedEvents(t *testing.T) {
	spy := &spyHandler{}
	logger := slog.New(spy)
	sink := &spySink{}
	notifService := notification.NewService(logger, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Success Task",
			Source:  "test",
			Status:  "pending",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Logger:        logger,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	// Verify events from sink
	if len(sink.events) != 2 {
		t.Fatalf("expected 2 notification events from sink, got %d", len(sink.events))
	}

	// Running Event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, ev0.Type)
	}
	if ev0.Attempt != 1 {
		t.Errorf("expected running attempt 1, got %d", ev0.Attempt)
	}
	if ev0.OccurredAt.IsZero() {
		t.Error("expected running occurred_at to be non-zero")
	}

	// Completed Event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventCompleted {
		t.Errorf("expected second event to be %s, got %s", notification.TaskEventCompleted, ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("expected completed attempt 1, got %d", ev1.Attempt)
	}
	if ev1.OccurredAt.IsZero() {
		t.Error("expected completed occurred_at to be non-zero")
	}
}

func TestWorkEmitsFailedEvent(t *testing.T) {
	spy := &spyHandler{}
	logger := slog.New(spy)
	sink := &spySink{}
	notifService := notification.NewService(logger, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-2",
			Title:   "Failure Task",
			Source:  "test",
			Status:  "pending",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Agent: fakeAgentClient{
			err: errors.New("agent failed"),
		},
		Logger: logger,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-2"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return error")
	}

	// Verify events from sink
	if len(sink.events) != 2 {
		t.Fatalf("expected 2 notification events from sink, got %d", len(sink.events))
	}

	// Running Event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, ev0.Type)
	}

	// Failed Event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventFailed {
		t.Errorf("expected second event to be %s, got %s", notification.TaskEventFailed, ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("expected failed attempt 1, got %d", ev1.Attempt)
	}
	if ev1.Reason != "agent failed" {
		t.Errorf("expected failed reason %q, got %q", "agent failed", ev1.Reason)
	}
	if ev1.OccurredAt.IsZero() {
		t.Error("expected failed occurred_at to be non-zero")
	}
}

func TestWorkKeepsAuthoringTaskRunningUntilDevelopMatch(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "milestone authored", Model: "m"}, nil
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-auth-wait",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
		Model:     fakeModel,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-wait"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected CompletedTask to NOT be called, but was called: %v", fakeLifecycle.completed)
	}

	foundWaitType := false
	foundRunState := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyWaitType] == "develop_match" {
			foundWaitType = true
		}
		if m[workflow.MetadataKeyAuthoringRunState] == "in_progress" {
			foundRunState = true
		}
	}
	if !foundWaitType {
		t.Errorf("expected MergeTaskMetadata to be called with wait_type=develop_match, got %v", fakeLifecycle.mergedMetadata)
	}
	if !foundRunState {
		t.Errorf("expected MergeTaskMetadata to be called with authoring_run_state=in_progress, got %v", fakeLifecycle.mergedMetadata)
	}
}