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

package scheduler

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

	"github.com/riverqueue/river"

	"github.com/nomadcode/nomadcode-core/internal/agent"
	"github.com/nomadcode/nomadcode-core/internal/authoring"
	"github.com/nomadcode/nomadcode-core/internal/model"
	"github.com/nomadcode/nomadcode-core/internal/notification"
	"github.com/nomadcode/nomadcode-core/internal/projectsync"
	"github.com/nomadcode/nomadcode-core/internal/roadmapsync"
	"github.com/nomadcode/nomadcode-core/internal/storage"
	"github.com/nomadcode/nomadcode-core/internal/workflow"
	"github.com/nomadcode/nomadcode-core/internal/workitem"
)

type authoringPushError struct {
	err error
}

func (e authoringPushError) Error() string { return e.err.Error() }
func (e authoringPushError) Unwrap() error { return e.err }

type authoringIdentityError struct {
	err error
}

func (e authoringIdentityError) Error() string { return e.err.Error() }
func (e authoringIdentityError) Unwrap() error { return e.err }

type TaskJobArgs struct {
	TaskID string `json:"task_id"`
}

func (TaskJobArgs) Kind() string {
	return "task_run"
}

// TaskLifecycle handles task state transitions and canonical metadata rules.
//
// Responsibility boundaries:
//   - Workflow service owns create/list/get/enqueue API semantics.
//   - Lifecycle owns all canonical status writes and metadata merge.
//   - Scheduler worker owns invoking StartTask, running execution, then invoking CompleteTask or FailTask.
//   - Provider projection failure must not roll back canonical task state.
type TaskLifecycle interface {
	StartTask(ctx context.Context, id string) (storage.Task, error)
	CompleteTask(ctx context.Context, id string, result json.RawMessage) (storage.Task, error)
	FailTask(ctx context.Context, id string, message string) (storage.Task, error)
	FailTaskWithMetadata(ctx context.Context, id string, input workflow.FailureInput) (storage.Task, error)
	MergeTaskMetadata(ctx context.Context, id string, updates map[string]any) (storage.Task, error)
}

type TaskWorker struct {
	river.WorkerDefaults[TaskJobArgs]

	Lifecycle      TaskLifecycle
	Notifications  *notification.Service
	Agent          agent.Client
	Model          model.Client
	SlotUpdater    WorkspaceSlotStateUpdater
	RunTimeout     time.Duration
	Logger         *slog.Logger
	IdentityWriter authoring.IdentityWriter
}

func (w *TaskWorker) Work(ctx context.Context, job *river.Job[TaskJobArgs]) error {
	taskID := job.Args.TaskID
	if w.Logger != nil {
		w.Logger.Info("task job started", "task_id", taskID)
	}

	task, err := w.Lifecycle.StartTask(ctx, taskID)
	if err != nil {
		return err
	}

	// Record authoring_run_state=in_progress before execution so a stale
	// detection layer can compare authoring_run_updated_at against a threshold
	// without waiting for the run to complete or fail.
	if _, isAuthoring := authoring.TaskCheckoutMetadata(task); isAuthoring {
		if _, mergeErr := w.Lifecycle.MergeTaskMetadata(ctx, taskID, map[string]any{
			workflow.MetadataKeyAuthoringRunState:     "in_progress",
			workflow.MetadataKeyAuthoringRunUpdatedAt: time.Now().UTC().Format(time.RFC3339),
		}); mergeErr != nil && w.Logger != nil {
			w.Logger.Warn("authoring in_progress metadata merge failed", "task_id", taskID, "error", mergeErr)
		}
	}

	attempt, reason, occurredAt := parseTaskEventContext(task, notification.TaskEventRunning)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventRunning,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})

	runCtx := ctx
	if w.RunTimeout > 0 {
		// Plane-origin authoring tasks rely on MODEL_TIMEOUT_SEC for their
		// execution bound; applying a generic worker deadline would prematurely
		// cancel IOP Responses calls that run for 10+ minutes.
		if _, isAuthoring := authoring.TaskCheckoutMetadata(task); !isAuthoring {
			var cancel context.CancelFunc
			runCtx, cancel = context.WithTimeout(ctx, w.RunTimeout)
			defer cancel()
		}
	}

	result, message, err := w.runTask(runCtx, task)
	if err != nil {
		w.markFailedWithTask(taskID, task, err)
		return err
	}

	// Plane-origin authoring task의 경우, bridge 성공만으로는 task를 완료하지 않고
	// develop match를 대기하는 상태(running)를 유지한다.
	_, isAuthoring := authoring.BuildAuthoringGenerateInput(task)
	if isAuthoring {
		var resMap map[string]any
		if err := json.Unmarshal(result, &resMap); err == nil {
			if state, ok := resMap[workflow.MetadataKeyAuthoringRunState].(string); ok && state == "in_progress" {
				updates := map[string]any{
					workflow.MetadataKeyAuthoringRunState:     "in_progress",
					workflow.MetadataKeyStatusReason:          message,
					workflow.MetadataKeyWaitType:              "develop_match",
					workflow.MetadataKeyAuthoringRunUpdatedAt: time.Now().UTC().Format(time.RFC3339),
				}
				for k, v := range resMap {
					updates[k] = v
				}

				if _, mergeErr := w.Lifecycle.MergeTaskMetadata(ctx, taskID, updates); mergeErr != nil && w.Logger != nil {
					w.Logger.Warn("authoring running metadata merge failed", "task_id", taskID, "error", mergeErr)
				}

				if w.Logger != nil {
					w.Logger.Info("task job remains running for develop match", "task_id", taskID)
				}
				return nil
			}
		}
	}

	task, err = w.Lifecycle.CompleteTask(ctx, taskID, result)
	if err != nil {
		w.markFailed(taskID, err)
		return err
	}

	attempt, reason, occurredAt = parseTaskEventContext(task, notification.TaskEventCompleted)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventCompleted,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Message:    message,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})

	if w.Logger != nil {
		w.Logger.Info("task job completed", "task_id", taskID)
	}
	return nil
}

func (w *TaskWorker) Timeout(job *river.Job[TaskJobArgs]) time.Duration {
	return 15 * time.Minute
}

func (w *TaskWorker) runTask(ctx context.Context, task storage.Task) (json.RawMessage, string, error) {
	// Plane-origin authoring tasks always route through the IOP OpenAI-compatible
	// HTTP bridge even when an A2A agent client is configured, because IOP runs
	// workspace-bound agent authoring rather than a generic A2A task.
	generateInput, isAuthoring, authoringInputErr := authoring.BuildAuthoringGenerateInputStrict(task)
	if isAuthoring {
		if authoringInputErr != nil {
			return nil, "", authoringInputErr
		}
		if w.Model == nil {
			return nil, "", fmt.Errorf("model client is required for authoring tasks")
		}
		var progressMode string
		var progressReason string
		generateInput.OnProgress = func(p model.GenerateProgress) {
			progressMode = p.Mode
			progressReason = p.Reason
			updates := map[string]any{
				workflow.MetadataKeyAuthoringProgressMode:   p.Mode,
				workflow.MetadataKeyAuthoringProgressReason: p.Reason,
				workflow.MetadataKeyAuthoringRunUpdatedAt:   p.LastEventTime.Format(time.RFC3339),
			}
			if _, err := w.Lifecycle.MergeTaskMetadata(ctx, task.ID, updates); err != nil && w.Logger != nil {
				w.Logger.Warn("authoring progress metadata merge failed", "task_id", task.ID, "error", err)
			}
		}

		generated, err := w.Model.Generate(ctx, generateInput)
		if err != nil {
			return nil, "", err
		}

		if w.IdentityWriter == nil {
			return nil, "", authoringIdentityError{err: errors.New("identity writer is required for authoring tasks")}
		}
		if generateInput.WorkspaceMetadata == nil {
			return nil, "", authoringIdentityError{err: errors.New("workspace metadata is required for authoring tasks")}
		}

		_, err = w.IdentityWriter.EnsureProviderIdentity(ctx, authoring.EnsureProviderIdentityInput{
			WorkspacePath: generateInput.WorkspaceMetadata.Path,
			Identity: roadmapsync.Identity{
				Provider:   workitem.ProviderID(generateInput.WorkspaceMetadata.Provider),
				Tenant:     generateInput.WorkspaceMetadata.Tenant,
				Project:    generateInput.WorkspaceMetadata.Project,
				WorkItemID: generateInput.WorkspaceMetadata.WorkItemID,
			},
		})
		if err != nil {
			return nil, "", authoringIdentityError{err: err}
		}

		// Push-only second stage: send the authored content to the workspace
		// agent so it can push to the remote develop branch before we gate on
		// the develop match confirmation.
		pushInput, _, pushInputErr := authoring.BuildAuthoringPushInputStrict(task)
		if pushInputErr != nil {
			return nil, "", pushInputErr
		}
		pushGenerated, pushErr := w.Model.Generate(ctx, pushInput)
		if pushErr != nil {
			return nil, "", authoringPushError{pushErr}
		}

		decision := authoring.DecideAuthoringResult(authoring.DecisionInput{
			BridgeSuccess: true,
		})
		result := map[string]any{
			"message":                  generated.Text,
			"mode":                     "authoring_run",
			"model":                    generated.Model,
			"response_id":              generated.ID,
			"push_response_id":         pushGenerated.ID,
			"push_request_state":       "succeeded",
			"authoring_run_state":      decision.State,
			"authoring_run_updated_at": time.Now().UTC().Format(time.RFC3339),
			workflow.MetadataKeyAuthoringProgressMode:   progressMode,
			workflow.MetadataKeyAuthoringProgressReason: progressReason,
			"usage": map[string]int{
				"input_tokens":  generated.Usage.InputTokens,
				"output_tokens": generated.Usage.OutputTokens,
				"total_tokens":  generated.Usage.TotalTokens,
			},
		}
		raw, err := json.Marshal(result)
		if err != nil {
			return nil, "", err
		}
		return raw, decision.Reason, nil
	}

	if w.Agent != nil {
		return w.runAgentTask(ctx, task)
	}

	if w.Model == nil {
		select {
		case <-time.After(500 * time.Millisecond):
		case <-ctx.Done():
			return nil, "", ctx.Err()
		}
		return json.RawMessage(`{"message":"dummy task completed","mode":"dummy"}`), "dummy task completed", nil
	}

	generated, err := w.Model.Generate(ctx, buildGenerateInput(task))
	if err != nil {
		return nil, "", err
	}

	result := map[string]any{
		"message":     generated.Text,
		"mode":        "openai_responses",
		"model":       generated.Model,
		"response_id": generated.ID,
		"usage": map[string]int{
			"input_tokens":  generated.Usage.InputTokens,
			"output_tokens": generated.Usage.OutputTokens,
			"total_tokens":  generated.Usage.TotalTokens,
		},
	}
	raw, err := json.Marshal(result)
	if err != nil {
		return nil, "", err
	}
	return raw, "model task completed", nil
}

func (w *TaskWorker) runAgentTask(ctx context.Context, task storage.Task) (json.RawMessage, string, error) {
	result, err := w.Agent.SendMessage(ctx, buildAgentInput(task))
	if err != nil {
		return nil, "", err
	}

	output := map[string]any{
		"mode": "a2a",
	}
	if len(result.Raw) > 0 {
		output["raw"] = json.RawMessage(result.Raw)
	}

	message := "a2a task completed"
	if result.Message != nil {
		text := messageText(*result.Message)
		if text != "" {
			message = text
			output["message"] = text
		}
		output["message_id"] = result.Message.MessageID
		raw, err := json.Marshal(output)
		if err != nil {
			return nil, "", err
		}
		return raw, message, nil
	}

	if result.Task == nil {
		return nil, "", fmt.Errorf("a2a result does not include message or task")
	}

	state := result.Task.Status.State
	output["a2a_task_id"] = result.Task.ID
	output["a2a_state"] = state

	text := taskOutputText(*result.Task)
	if text != "" {
		message = text
		output["message"] = text
	}

	switch state {
	case agent.TaskStateCompleted:
	case agent.TaskStateFailed, agent.TaskStateCanceled, agent.TaskStateRejected, agent.TaskStateInputNeeded, agent.TaskStateAuthRequired:
		if text == "" {
			text = fmt.Sprintf("a2a task ended with state %s", state)
		}
		return nil, "", fmt.Errorf("%s", text)
	default:
		return nil, "", fmt.Errorf("a2a blocking call returned non-terminal task state %s", state)
	}

	raw, err := json.Marshal(output)
	if err != nil {
		return nil, "", err
	}
	return raw, message, nil
}

func buildGenerateInput(task storage.Task) model.GenerateInput {
	input := strings.TrimSpace(task.Title)
	instructions := ""

	if len(task.Payload) > 0 && string(task.Payload) != "null" {
		if payloadInput, payloadInstructions, ok := parsePayloadPrompt(task.Payload); ok {
			if payloadInput != "" {
				input = payloadInput
			}
			instructions = payloadInstructions
		} else if string(task.Payload) != "{}" {
			input = fmt.Sprintf("Task title: %s\nPayload:\n%s", task.Title, string(task.Payload))
		}
	}

	return model.GenerateInput{
		Input:        input,
		Instructions: instructions,
		Metadata: map[string]string{
			"task_id": task.ID,
			"source":  task.Source,
		},
	}
}

func buildAgentInput(task storage.Task) agent.SendMessageInput {
	generateInput := buildGenerateInput(task)
	text := strings.TrimSpace(generateInput.Input)
	if generateInput.Instructions != "" {
		text = fmt.Sprintf("Instructions:\n%s\n\nTask:\n%s", generateInput.Instructions, text)
	}
	if text == "" {
		text = task.Title
	}

	return agent.SendMessageInput{
		Text:                text,
		AcceptedOutputModes: []string{"text/plain", "application/json"},
		Blocking:            true,
		Metadata: map[string]any{
			"nomadcode_task_id": task.ID,
			"source":            task.Source,
			"title":             task.Title,
		},
	}
}

func parsePayloadPrompt(payload json.RawMessage) (input string, instructions string, ok bool) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(payload, &fields); err != nil {
		return "", "", false
	}
	ok = true
	input = firstStringField(fields, "prompt", "message", "input")
	instructions = firstStringField(fields, "instructions", "system")
	return input, instructions, ok
}

func firstStringField(fields map[string]json.RawMessage, names ...string) string {
	for _, name := range names {
		raw, exists := fields[name]
		if !exists {
			continue
		}
		var value string
		if err := json.Unmarshal(raw, &value); err == nil {
			if value = strings.TrimSpace(value); value != "" {
				return value
			}
		}
	}
	return ""
}

func taskOutputText(task agent.Task) string {
	if task.Status.Message != nil {
		if text := messageText(*task.Status.Message); text != "" {
			return text
		}
	}

	for _, artifact := range task.Artifacts {
		if text := partsText(artifact.Parts); text != "" {
			return text
		}
	}

	for i := len(task.History) - 1; i >= 0; i-- {
		if task.History[i].Role == "agent" || task.History[i].Role == "assistant" {
			if text := messageText(task.History[i]); text != "" {
				return text
			}
		}
	}
	for i := len(task.History) - 1; i >= 0; i-- {
		if text := messageText(task.History[i]); text != "" {
			return text
		}
	}

	return ""
}

func messageText(message agent.Message) string {
	return partsText(message.Parts)
}

func partsText(parts []agent.Part) string {
	texts := make([]string, 0, len(parts))
	for _, part := range parts {
		text := strings.TrimSpace(part.Text)
		if text != "" {
			texts = append(texts, text)
		}
	}
	return strings.Join(texts, "\n")
}

func (w *TaskWorker) notifyTaskEvent(ctx context.Context, event notification.TaskEvent) {
	if w.Notifications == nil {
		return
	}
	err := w.Notifications.NotifyTaskEvent(ctx, event)
	if err != nil && w.Logger != nil {
		w.Logger.Warn("task notification failed", "task_id", event.TaskID, "type", event.Type, "error", err)
	}
}

// markFailedWithTask is markFailed with the originating task available so
// authoring-specific failure metadata can be appended before the lifecycle call.
func (w *TaskWorker) markFailedWithTask(taskID string, task storage.Task, err error) {
	if err == nil || w.Lifecycle == nil {
		return
	}

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

	msg := err.Error()
	failureType := workflow.FailureTypeExecution
	if errors.Is(err, context.DeadlineExceeded) {
		msg = "timeout"
		failureType = workflow.FailureTypeTimeout
	}

	failInput := workflow.FailureInput{Message: msg, Type: failureType}

	var terminalSlotState projectsync.SlotState
	var releaseSlot bool
	if _, isAuthoring := authoring.TaskCheckoutMetadata(task); isAuthoring {
		var pushErr authoringPushError
		var identErr authoringIdentityError
		decisionInput := authoring.DecisionInput{BridgeSuccess: false}
		if errors.As(err, &pushErr) {
			decisionInput = authoring.DecisionInput{BridgeSuccess: true, PushFailed: true}
		} else if errors.As(err, &identErr) {
			innerErr := identErr.Unwrap()
			if errors.Is(innerErr, roadmapsync.ErrInvalidIdentity) || errors.Is(innerErr, authoring.ErrProviderIdentityMissing) {
				decisionInput = authoring.DecisionInput{BridgeSuccess: true, IdentityMissing: true}
			} else {
				decisionInput = authoring.DecisionInput{BridgeSuccess: true, IdentityWriteFailed: true}
			}
		}
		decision := authoring.DecideAuthoringResult(decisionInput)
		failInput.ExtraMetadata = map[string]any{
			workflow.MetadataKeyAuthoringRunState:        decision.State,
			workflow.MetadataKeyAuthoringRunUpdatedAt:    time.Now().UTC().Format(time.RFC3339),
			workflow.MetadataKeyAuthoringFailureType:     string(failureType),
			workflow.MetadataKeyAuthoringFailureCategory: decision.FailureCategory,
		}
		if state, ok := authoring.SlotStateForAuthoringOutcome(decision); ok {
			terminalSlotState = state
			releaseSlot = true
		}
	}

	failed, failErr := w.Lifecycle.FailTaskWithMetadata(ctx, taskID, failInput)
	if failErr != nil {
		if w.Logger != nil {
			w.Logger.Error("failed to mark task failed", "task_id", taskID, "error", failErr)
		}
		return
	}
	if releaseSlot {
		applyTaskSlotTerminalState(ctx, w.SlotUpdater, task, terminalSlotState, w.Logger)
	}

	attempt, reason, occurredAt := parseTaskEventContext(failed, notification.TaskEventFailed)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventFailed,
		TaskID:     failed.ID,
		Title:      failed.Title,
		Status:     failed.Status,
		Message:    msg,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})
}

func (w *TaskWorker) markFailed(taskID string, err error) {
	if err == nil || w.Lifecycle == nil {
		return
	}

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

	msg := err.Error()
	failureType := workflow.FailureTypeExecution
	if errors.Is(err, context.DeadlineExceeded) {
		msg = "timeout"
		failureType = workflow.FailureTypeTimeout
	}

	task, failErr := w.Lifecycle.FailTaskWithMetadata(ctx, taskID, workflow.FailureInput{
		Message: msg,
		Type:    failureType,
	})
	if failErr != nil {
		if w.Logger != nil {
			w.Logger.Error("failed to mark task failed", "task_id", taskID, "error", failErr)
		}
		return
	}

	attempt, reason, occurredAt := parseTaskEventContext(task, notification.TaskEventFailed)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventFailed,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Message:    msg,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})
}

func parseTaskEventContext(task storage.Task, eventType notification.TaskEventType) (int, string, time.Time) {
	attempt := 0
	reason := ""
	occurredAt := task.UpdatedAt

	if len(task.Metadata) > 0 && string(task.Metadata) != "null" {
		var meta map[string]any
		if err := json.Unmarshal(task.Metadata, &meta); err == nil {
			if val, ok := meta[workflow.MetadataKeyAttempt]; ok {
				switch v := val.(type) {
				case float64:
					attempt = int(v)
				case int:
					attempt = v
				case int64:
					attempt = int(v)
				}
			}
			if eventType == notification.TaskEventFailed {
				if val, ok := meta[workflow.MetadataKeyStatusReason]; ok {
					if str, ok := val.(string); ok {
						reason = str
					}
				}
			}
		}
	}
	return attempt, reason, occurredAt
}