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

package scheduler

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"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"
)

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

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

type TaskWorker struct {
	river.WorkerDefaults[TaskJobArgs]

	Store         *storage.Store
	Notifications *notification.Service
	Agent         agent.Client
	Model         model.Client
	Logger        *slog.Logger
}

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.Store.UpdateStatus(ctx, taskID, string(workflow.StatusRunning))
	if err != nil {
		return err
	}

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

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

	if w.Notifications != nil {
		err = w.Notifications.NotifyTaskCompleted(ctx, notification.TaskNotification{
			TaskID:  task.ID,
			Title:   task.Title,
			Status:  task.Status,
			Message: message,
		})
		if err != nil && w.Logger != nil {
			w.Logger.Warn("task notification failed", "task_id", taskID, "error", err)
		}
	}

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

func (w *TaskWorker) runTask(ctx context.Context, task storage.Task) (json.RawMessage, string, error) {
	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) markFailed(taskID string, err error) {
	if err == nil || w.Store == nil {
		return
	}

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

	_, failErr := w.Store.FailTask(ctx, taskID, err.Error())
	if failErr != nil && w.Logger != nil {
		w.Logger.Error("failed to mark task failed", "task_id", taskID, "error", failErr)
		return
	}

}