363 lines
9.1 KiB
Go
363 lines
9.1 KiB
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/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"
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
|
|
type TaskWorker struct {
|
|
river.WorkerDefaults[TaskJobArgs]
|
|
|
|
Lifecycle TaskLifecycle
|
|
Notifications *notification.Service
|
|
Agent agent.Client
|
|
Model model.Client
|
|
RunTimeout time.Duration
|
|
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.Lifecycle.StartTask(ctx, taskID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
w.notifyTaskEvent(ctx, notification.TaskEvent{
|
|
Type: notification.TaskEventRunning,
|
|
TaskID: task.ID,
|
|
Title: task.Title,
|
|
Status: task.Status,
|
|
})
|
|
|
|
runCtx := ctx
|
|
if w.RunTimeout > 0 {
|
|
var cancel context.CancelFunc
|
|
runCtx, cancel = context.WithTimeout(ctx, w.RunTimeout)
|
|
defer cancel()
|
|
}
|
|
|
|
result, message, err := w.runTask(runCtx, task)
|
|
if err != nil {
|
|
w.markFailed(taskID, err)
|
|
return err
|
|
}
|
|
|
|
task, err = w.Lifecycle.CompleteTask(ctx, taskID, result)
|
|
if err != nil {
|
|
w.markFailed(taskID, err)
|
|
return err
|
|
}
|
|
|
|
w.notifyTaskEvent(ctx, notification.TaskEvent{
|
|
Type: notification.TaskEventCompleted,
|
|
TaskID: task.ID,
|
|
Title: task.Title,
|
|
Status: task.Status,
|
|
Message: message,
|
|
})
|
|
|
|
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) 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)
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
w.notifyTaskEvent(ctx, notification.TaskEvent{
|
|
Type: notification.TaskEventFailed,
|
|
TaskID: task.ID,
|
|
Title: task.Title,
|
|
Status: task.Status,
|
|
Message: msg,
|
|
})
|
|
}
|