toki
6ab8dd8f63
- Add authoring adapter for openai integration - Move 02+01_iop_authoring_bridge to archive - Update scheduler jobs and workflow lifecycle - Enhance openai client and model definitions
499 lines
14 KiB
Go
499 lines
14 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/authoring"
|
|
"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)
|
|
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
|
|
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
|
|
}
|
|
|
|
// 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.BuildAuthoringGenerateInput(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 {
|
|
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
|
|
}
|
|
|
|
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) 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 := authoring.BuildAuthoringGenerateInput(task)
|
|
if isAuthoring {
|
|
if w.Model == nil {
|
|
return nil, "", fmt.Errorf("model client is required for authoring tasks")
|
|
}
|
|
generated, err := w.Model.Generate(ctx, generateInput)
|
|
if err != nil {
|
|
return nil, "", err
|
|
}
|
|
result := map[string]any{
|
|
"message": generated.Text,
|
|
"mode": "authoring_run",
|
|
"model": generated.Model,
|
|
"response_id": generated.ID,
|
|
"authoring_run_state": "succeeded",
|
|
"authoring_run_updated_at": time.Now().UTC().Format(time.RFC3339),
|
|
"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, "authoring run completed", 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}
|
|
|
|
if _, isAuthoring := authoring.BuildAuthoringGenerateInput(task); isAuthoring {
|
|
failInput.ExtraMetadata = map[string]any{
|
|
workflow.MetadataKeyAuthoringRunState: "failed",
|
|
workflow.MetadataKeyAuthoringRunUpdatedAt: time.Now().UTC().Format(time.RFC3339),
|
|
workflow.MetadataKeyAuthoringFailureType: string(failureType),
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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
|
|
}
|