From 4b6951cfa8d3690a9da5f2092ca41d051594f556 Mon Sep 17 00:00:00 2001 From: toki Date: Sun, 24 May 2026 19:36:43 +0900 Subject: [PATCH] =?UTF-8?q?feat:=20notification=20=EB=B0=8F=20workflow=20?= =?UTF-8?q?=EA=B8=B0=EB=8A=A5=20=EA=B0=9C=EC=84=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - notification 서비스 테스트 추가 - workflow lifecycle 관리 구현 - config, scheduler, workflow 서비스 개선 --- services/core/Makefile | 1 + services/core/cmd/server/main.go | 4 +- services/core/internal/config/config.go | 82 ++-- services/core/internal/config/config_test.go | 13 + services/core/internal/notification/model.go | 17 + .../core/internal/notification/service.go | 17 +- .../internal/notification/service_test.go | 114 ++++++ services/core/internal/scheduler/jobs.go | 81 +++- services/core/internal/scheduler/jobs_test.go | 358 ++++++++++++++++++ services/core/internal/scheduler/river.go | 7 +- services/core/internal/workflow/lifecycle.go | 213 +++++++++++ services/core/internal/workflow/service.go | 5 +- .../core/internal/workflow/service_test.go | 102 +++++ 13 files changed, 947 insertions(+), 67 deletions(-) create mode 100644 services/core/internal/notification/service_test.go create mode 100644 services/core/internal/workflow/lifecycle.go diff --git a/services/core/Makefile b/services/core/Makefile index ae0de94..63d3b64 100644 --- a/services/core/Makefile +++ b/services/core/Makefile @@ -14,6 +14,7 @@ export MODEL_TIMEOUT_SEC ?= 300 export A2A_EDGE_URL ?= export A2A_AGENT_URL ?= export A2A_TIMEOUT_SEC ?= 300 +export WORKFLOW_TASK_TIMEOUT_SEC ?= 300 export GOOSE_DRIVER ?= postgres export GOOSE_DBSTRING ?= $(DATABASE_URL) export OUTPUT ?= .build/nomadcode-core diff --git a/services/core/cmd/server/main.go b/services/core/cmd/server/main.go index e669f5e..3fe9a64 100644 --- a/services/core/cmd/server/main.go +++ b/services/core/cmd/server/main.go @@ -72,7 +72,9 @@ func run(logger *slog.Logger) error { logger.Info("a2a edge input enabled", "url", cfg.A2AEdgeURL) } - taskScheduler, err := scheduler.New(pool, store, notificationService, agentClient, modelClient, logger) + lifecycle := workflow.NewLifecycle(store, logger) + + taskScheduler, err := scheduler.New(pool, lifecycle, notificationService, agentClient, modelClient, time.Duration(cfg.WorkflowTaskTimeoutSec)*time.Second, logger) if err != nil { return err } diff --git a/services/core/internal/config/config.go b/services/core/internal/config/config.go index 626d3a2..39d92ed 100644 --- a/services/core/internal/config/config.go +++ b/services/core/internal/config/config.go @@ -6,50 +6,52 @@ import ( ) type Config struct { - AppEnv string - HTTPAddr string - DatabaseURL string - RedisURL string - RedisKeyPrefix string - AuthUsername string - AuthPassword string - ModelBaseURL string - ModelAPIKey string - ModelName string - ModelContextSize int - ModelTimeoutSec int - A2AEdgeURL string - A2AAgentURL string - A2AToken string - A2ATimeoutSec int - MattermostBaseURL string - MattermostToken string - PlaneBaseURL string - PlaneToken string + AppEnv string + HTTPAddr string + DatabaseURL string + RedisURL string + RedisKeyPrefix string + AuthUsername string + AuthPassword string + ModelBaseURL string + ModelAPIKey string + ModelName string + ModelContextSize int + ModelTimeoutSec int + A2AEdgeURL string + A2AAgentURL string + A2AToken string + A2ATimeoutSec int + MattermostBaseURL string + MattermostToken string + PlaneBaseURL string + PlaneToken string + WorkflowTaskTimeoutSec int } func Load() Config { return Config{ - AppEnv: getEnv("APP_ENV", "local"), - HTTPAddr: getEnv("HTTP_ADDR", ":8080"), - DatabaseURL: os.Getenv("DATABASE_URL"), - RedisURL: os.Getenv("REDIS_URL"), - RedisKeyPrefix: os.Getenv("REDIS_KEY_PREFIX"), - AuthUsername: getEnv("AUTH_USERNAME", "nomadcode"), - AuthPassword: os.Getenv("AUTH_PASSWORD"), - ModelBaseURL: getEnv("MODEL_BASE_URL", "http://192.168.0.91:11434"), - ModelAPIKey: getEnv("MODEL_API_KEY", "ollama"), - ModelName: getEnv("MODEL_NAME", "qwen3.6:35b-a3b-bf16"), - ModelContextSize: getEnvInt("MODEL_CONTEXT_SIZE", 262144), - ModelTimeoutSec: getEnvInt("MODEL_TIMEOUT_SEC", 300), - A2AEdgeURL: firstEnv("A2A_EDGE_URL", "A2A_AGENT_URL"), - A2AAgentURL: firstEnv("A2A_AGENT_URL", "A2A_EDGE_URL"), - A2AToken: os.Getenv("A2A_TOKEN"), - A2ATimeoutSec: getEnvInt("A2A_TIMEOUT_SEC", 300), - MattermostBaseURL: os.Getenv("MATTERMOST_BASE_URL"), - MattermostToken: os.Getenv("MATTERMOST_TOKEN"), - PlaneBaseURL: os.Getenv("PLANE_BASE_URL"), - PlaneToken: os.Getenv("PLANE_TOKEN"), + AppEnv: getEnv("APP_ENV", "local"), + HTTPAddr: getEnv("HTTP_ADDR", ":8080"), + DatabaseURL: os.Getenv("DATABASE_URL"), + RedisURL: os.Getenv("REDIS_URL"), + RedisKeyPrefix: os.Getenv("REDIS_KEY_PREFIX"), + AuthUsername: getEnv("AUTH_USERNAME", "nomadcode"), + AuthPassword: os.Getenv("AUTH_PASSWORD"), + ModelBaseURL: getEnv("MODEL_BASE_URL", "http://192.168.0.91:11434"), + ModelAPIKey: getEnv("MODEL_API_KEY", "ollama"), + ModelName: getEnv("MODEL_NAME", "qwen3.6:35b-a3b-bf16"), + ModelContextSize: getEnvInt("MODEL_CONTEXT_SIZE", 262144), + ModelTimeoutSec: getEnvInt("MODEL_TIMEOUT_SEC", 300), + A2AEdgeURL: firstEnv("A2A_EDGE_URL", "A2A_AGENT_URL"), + A2AAgentURL: firstEnv("A2A_AGENT_URL", "A2A_EDGE_URL"), + A2AToken: os.Getenv("A2A_TOKEN"), + A2ATimeoutSec: getEnvInt("A2A_TIMEOUT_SEC", 300), + MattermostBaseURL: os.Getenv("MATTERMOST_BASE_URL"), + MattermostToken: os.Getenv("MATTERMOST_TOKEN"), + PlaneBaseURL: os.Getenv("PLANE_BASE_URL"), + PlaneToken: os.Getenv("PLANE_TOKEN"), + WorkflowTaskTimeoutSec: getEnvInt("WORKFLOW_TASK_TIMEOUT_SEC", 300), } } diff --git a/services/core/internal/config/config_test.go b/services/core/internal/config/config_test.go index 3983073..a7c6c0a 100644 --- a/services/core/internal/config/config_test.go +++ b/services/core/internal/config/config_test.go @@ -51,3 +51,16 @@ func TestLoadModelAndA2ADefaults(t *testing.T) { t.Fatalf("A2ATimeoutSec: got %d", cfg.A2ATimeoutSec) } } + +func TestConfigLoadsWorkflowTaskTimeout(t *testing.T) { + cfg := Load() + if cfg.WorkflowTaskTimeoutSec != 300 { + t.Fatalf("WorkflowTaskTimeoutSec default: got %d, want 300", cfg.WorkflowTaskTimeoutSec) + } + + t.Setenv("WORKFLOW_TASK_TIMEOUT_SEC", "120") + cfgOverride := Load() + if cfgOverride.WorkflowTaskTimeoutSec != 120 { + t.Fatalf("WorkflowTaskTimeoutSec override: got %d, want 120", cfgOverride.WorkflowTaskTimeoutSec) + } +} diff --git a/services/core/internal/notification/model.go b/services/core/internal/notification/model.go index 13051e7..9a14e7f 100644 --- a/services/core/internal/notification/model.go +++ b/services/core/internal/notification/model.go @@ -6,3 +6,20 @@ type TaskNotification struct { Status string Message string } + +type TaskEventType string + +const ( + TaskEventRunning TaskEventType = "task.running" + TaskEventCompleted TaskEventType = "task.completed" + TaskEventFailed TaskEventType = "task.failed" + TaskEventCanceled TaskEventType = "task.canceled" +) + +type TaskEvent struct { + Type TaskEventType + TaskID string + Title string + Status string + Message string +} diff --git a/services/core/internal/notification/service.go b/services/core/internal/notification/service.go index 690d39a..03ab971 100644 --- a/services/core/internal/notification/service.go +++ b/services/core/internal/notification/service.go @@ -20,9 +20,24 @@ func NewService(mattermostClient *mattermost.Client, logger *slog.Logger) *Servi } func (s *Service) NotifyTaskCompleted(ctx context.Context, input TaskNotification) error { + return s.NotifyTaskEvent(ctx, TaskEvent{ + Type: TaskEventCompleted, + TaskID: input.TaskID, + Title: input.Title, + Status: input.Status, + Message: input.Message, + }) +} + +func (s *Service) NotifyTaskEvent(ctx context.Context, input TaskEvent) error { if s.logger != nil { - s.logger.Info("task completed notification requested", "task_id", input.TaskID) + s.logger.Info("task event notification requested", "type", input.Type, "task_id", input.TaskID) } + + if input.Type != TaskEventCompleted { + return nil + } + if s.mattermost == nil { return nil } diff --git a/services/core/internal/notification/service_test.go b/services/core/internal/notification/service_test.go new file mode 100644 index 0000000..4404320 --- /dev/null +++ b/services/core/internal/notification/service_test.go @@ -0,0 +1,114 @@ +package notification_test + +import ( + "context" + "log/slog" + "testing" + + "github.com/nomadcode/nomadcode-core/internal/notification" +) + +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 +} + +func TestNotifyTaskCompletedUsesTaskEventCompleted(t *testing.T) { + spy := &spyHandler{} + logger := slog.New(spy) + + service := notification.NewService(nil, logger) + + err := service.NotifyTaskCompleted(context.Background(), notification.TaskNotification{ + TaskID: "task-123", + Title: "Test Task", + Status: "completed", + Message: "Finished successfully", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(spy.records) == 0 { + t.Fatal("expected log records, got none") + } + + found := false + for _, rec := range spy.records { + if rec.msg == "task event notification requested" { + found = true + if rec.args["type"] != notification.TaskEventCompleted { + t.Errorf("expected event type %s, got %v", notification.TaskEventCompleted, rec.args["type"]) + } + if rec.args["task_id"] != "task-123" { + t.Errorf("expected task_id task-123, got %v", rec.args["task_id"]) + } + } + } + if !found { + t.Error("expected 'task event notification requested' log message, not found") + } +} + +func TestNotifyTaskEventAllowsFailedWithoutMattermost(t *testing.T) { + spy := &spyHandler{} + logger := slog.New(spy) + + service := notification.NewService(nil, logger) + + err := service.NotifyTaskEvent(context.Background(), notification.TaskEvent{ + Type: notification.TaskEventFailed, + TaskID: "task-456", + Title: "Failed Task", + Status: "failed", + Message: "Something went wrong", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + found := false + for _, rec := range spy.records { + if rec.msg == "task event notification requested" { + found = true + if rec.args["type"] != notification.TaskEventFailed { + t.Errorf("expected event type %s, got %v", notification.TaskEventFailed, rec.args["type"]) + } + if rec.args["task_id"] != "task-456" { + t.Errorf("expected task_id task-456, got %v", rec.args["task_id"]) + } + } + } + if !found { + t.Error("expected 'task event notification requested' log message, not found") + } +} diff --git a/services/core/internal/scheduler/jobs.go b/services/core/internal/scheduler/jobs.go index 9f1b5bd..65c1034 100644 --- a/services/core/internal/scheduler/jobs.go +++ b/services/core/internal/scheduler/jobs.go @@ -3,6 +3,7 @@ package scheduler import ( "context" "encoding/json" + "errors" "fmt" "log/slog" "strings" @@ -14,7 +15,6 @@ import ( "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 { @@ -25,13 +25,20 @@ func (TaskJobArgs) Kind() string { return "task_run" } +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) +} + type TaskWorker struct { river.WorkerDefaults[TaskJobArgs] - Store *storage.Store + Lifecycle TaskLifecycle Notifications *notification.Service Agent agent.Client Model model.Client + RunTimeout time.Duration Logger *slog.Logger } @@ -41,34 +48,44 @@ func (w *TaskWorker) Work(ctx context.Context, job *river.Job[TaskJobArgs]) erro w.Logger.Info("task job started", "task_id", taskID) } - task, err := w.Store.UpdateStatus(ctx, taskID, string(workflow.StatusRunning)) + task, err := w.Lifecycle.StartTask(ctx, taskID) if err != nil { return err } - result, message, err := w.runTask(ctx, task) + 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.Store.CompleteTask(ctx, taskID, result) + task, err = w.Lifecycle.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) - } - } + 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) @@ -291,18 +308,42 @@ func partsText(parts []agent.Part) string { 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.Store == nil { + if err == nil || w.Lifecycle == 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) + msg := err.Error() + if errors.Is(err, context.DeadlineExceeded) { + msg = "timeout" + } + + task, failErr := w.Lifecycle.FailTask(ctx, taskID, msg) + 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, + }) } diff --git a/services/core/internal/scheduler/jobs_test.go b/services/core/internal/scheduler/jobs_test.go index 5842087..258b6ad 100644 --- a/services/core/internal/scheduler/jobs_test.go +++ b/services/core/internal/scheduler/jobs_test.go @@ -3,10 +3,17 @@ package scheduler import ( "context" "encoding/json" + "errors" + "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" ) @@ -106,3 +113,354 @@ func (c fakeAgentClient) GetTask(context.Context, agent.GetTaskInput) (agent.Tas 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 + + task storage.Task +} + +func (f *fakeTaskLifecycle) StartTask(ctx context.Context, id string) (storage.Task, error) { + if 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" + 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" + return f.task, nil +} + +func (f *fakeTaskLifecycle) FailTask(ctx context.Context, id string, message string) (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, message) + f.task.Status = "failed" + 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 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) + } +} + +func TestWorkRetriesAfterFailureState(t *testing.T) { + fakeLifecycle := &fakeTaskLifecycle{ + task: storage.Task{ + ID: "task-1", + Status: "pending", + Title: "Test task", + Source: "test", + Payload: json.RawMessage(`{}`), + }, + } + + worker := &TaskWorker{ + Lifecycle: fakeLifecycle, + 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) + } + + 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)) + } +} + +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 +} + +func TestWorkEmitsRunningAndCompletedEvents(t *testing.T) { + spy := &spyHandler{} + logger := slog.New(spy) + notifService := notification.NewService(nil, logger) + + 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 + var eventTypes []notification.TaskEventType + for _, rec := range spy.records { + if rec.msg == "task event notification requested" { + if tp, ok := rec.args["type"].(notification.TaskEventType); ok { + eventTypes = append(eventTypes, tp) + } + } + } + + if len(eventTypes) != 2 { + t.Fatalf("expected 2 notification events, got %d: %v", len(eventTypes), eventTypes) + } + + if eventTypes[0] != notification.TaskEventRunning { + t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, eventTypes[0]) + } + + if eventTypes[1] != notification.TaskEventCompleted { + t.Errorf("expected second event to be %s, got %s", notification.TaskEventCompleted, eventTypes[1]) + } +} + +func TestWorkEmitsFailedEvent(t *testing.T) { + spy := &spyHandler{} + logger := slog.New(spy) + notifService := notification.NewService(nil, logger) + + 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 + var eventTypes []notification.TaskEventType + for _, rec := range spy.records { + if rec.msg == "task event notification requested" { + if tp, ok := rec.args["type"].(notification.TaskEventType); ok { + eventTypes = append(eventTypes, tp) + } + } + } + + // Running + Failed + if len(eventTypes) != 2 { + t.Fatalf("expected 2 notification events, got %d: %v", len(eventTypes), eventTypes) + } + + if eventTypes[0] != notification.TaskEventRunning { + t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, eventTypes[0]) + } + + if eventTypes[1] != notification.TaskEventFailed { + t.Errorf("expected second event to be %s, got %s", notification.TaskEventFailed, eventTypes[1]) + } +} diff --git a/services/core/internal/scheduler/river.go b/services/core/internal/scheduler/river.go index 7b0f732..deb47a3 100644 --- a/services/core/internal/scheduler/river.go +++ b/services/core/internal/scheduler/river.go @@ -3,6 +3,7 @@ package scheduler import ( "context" "log/slog" + "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -13,7 +14,6 @@ import ( "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" ) type Client struct { @@ -22,13 +22,14 @@ type Client struct { logger *slog.Logger } -func New(pool *pgxpool.Pool, store *storage.Store, notifications *notification.Service, agentClient agent.Client, modelClient model.Client, logger *slog.Logger) (*Client, error) { +func New(pool *pgxpool.Pool, lifecycle TaskLifecycle, notifications *notification.Service, agentClient agent.Client, modelClient model.Client, runTimeout time.Duration, logger *slog.Logger) (*Client, error) { workers := river.NewWorkers() river.AddWorker(workers, &TaskWorker{ - Store: store, + Lifecycle: lifecycle, Notifications: notifications, Agent: agentClient, Model: modelClient, + RunTimeout: runTimeout, Logger: logger, }) diff --git a/services/core/internal/workflow/lifecycle.go b/services/core/internal/workflow/lifecycle.go new file mode 100644 index 0000000..a93b95e --- /dev/null +++ b/services/core/internal/workflow/lifecycle.go @@ -0,0 +1,213 @@ +package workflow + +import ( + "context" + "encoding/json" + "log/slog" + "strings" + "time" + + "github.com/nomadcode/nomadcode-core/internal/storage" +) + +type Lifecycle struct { + store *storage.Store + logger *slog.Logger +} + +func NewLifecycle(store *storage.Store, logger *slog.Logger) *Lifecycle { + return &Lifecycle{ + store: store, + logger: logger, + } +} + +func canTransition(from, to TaskStatus) bool { + switch from { + case StatusPending, StatusQueued, StatusFailed: + return to == StatusRunning + case StatusRunning: + return to == StatusCompleted || to == StatusFailed || to == StatusCanceled + default: + return false + } +} + +func nextTaskAttempt(metadata json.RawMessage) (int, error) { + if len(metadata) == 0 || string(metadata) == "null" { + return 1, nil + } + var currentMap map[string]any + if err := json.Unmarshal(metadata, ¤tMap); err != nil { + return 0, ErrInvalidTaskInput + } + attempt := 1 + if currentMap != nil { + if val, ok := currentMap[MetadataKeyAttempt]; ok { + switch v := val.(type) { + case float64: + attempt = int(v) + 1 + case int: + attempt = v + 1 + case int64: + attempt = int(v) + 1 + } + } + } + return attempt, nil +} + +func mergeTaskMetadata(current json.RawMessage, updates map[string]any) (json.RawMessage, error) { + var currentMap map[string]any + if len(current) > 0 && string(current) != "null" { + if err := json.Unmarshal(current, ¤tMap); err != nil { + return nil, ErrInvalidTaskInput + } + } else { + currentMap = make(map[string]any) + } + + for k, v := range updates { + if v == nil { + delete(currentMap, k) + } else { + currentMap[k] = v + } + } + + return json.Marshal(currentMap) +} + +func (l *Lifecycle) StartTask(ctx context.Context, id string) (storage.Task, error) { + task, err := l.store.GetTask(ctx, id) + if err != nil { + return storage.Task{}, err + } + + if !canTransition(TaskStatus(task.Status), StatusRunning) { + return storage.Task{}, ErrInvalidTaskTransition + } + + attempt, err := nextTaskAttempt(task.Metadata) + if err != nil { + return storage.Task{}, err + } + + updates := map[string]any{ + MetadataKeyAttempt: attempt, + MetadataKeyAgentRunState: "running", + MetadataKeyLastHeartbeat: time.Now().UTC().Format(time.RFC3339), + } + + newMeta, err := mergeTaskMetadata(task.Metadata, updates) + if err != nil { + return storage.Task{}, err + } + + task, err = l.store.UpdateMetadata(ctx, id, newMeta) + if err != nil { + return storage.Task{}, err + } + + return l.store.UpdateStatus(ctx, id, string(StatusRunning)) +} + +func (l *Lifecycle) CompleteTask(ctx context.Context, id string, result json.RawMessage) (storage.Task, error) { + task, err := l.store.GetTask(ctx, id) + if err != nil { + return storage.Task{}, err + } + + if !canTransition(TaskStatus(task.Status), StatusCompleted) { + return storage.Task{}, ErrInvalidTaskTransition + } + + updates := map[string]any{ + MetadataKeyAgentRunState: "completed", + MetadataKeyWaitType: nil, + } + + var summary string + if len(result) > 0 && string(result) != "null" { + var resMap map[string]interface{} + if err := json.Unmarshal(result, &resMap); err == nil { + if sVal, ok := resMap["summary"].(string); ok { + summary = strings.TrimSpace(sVal) + } + } + } + + if summary != "" { + updates[MetadataKeyStatusReason] = summary + } + + newMeta, err := mergeTaskMetadata(task.Metadata, updates) + if err != nil { + return storage.Task{}, err + } + + task, err = l.store.UpdateMetadata(ctx, id, newMeta) + if err != nil { + return storage.Task{}, err + } + + return l.store.CompleteTask(ctx, id, result) +} + +func (l *Lifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) { + task, err := l.store.GetTask(ctx, id) + if err != nil { + return storage.Task{}, err + } + + if !canTransition(TaskStatus(task.Status), StatusFailed) { + return storage.Task{}, ErrInvalidTaskTransition + } + + updates := map[string]any{ + MetadataKeyAgentRunState: "failed", + MetadataKeyStatusReason: message, + } + + newMeta, err := mergeTaskMetadata(task.Metadata, updates) + if err != nil { + return storage.Task{}, err + } + + task, err = l.store.UpdateMetadata(ctx, id, newMeta) + if err != nil { + return storage.Task{}, err + } + + return l.store.FailTask(ctx, id, message) +} + +func (l *Lifecycle) CancelTask(ctx context.Context, id string, message string) (storage.Task, error) { + task, err := l.store.GetTask(ctx, id) + if err != nil { + return storage.Task{}, err + } + + if !canTransition(TaskStatus(task.Status), StatusCanceled) { + return storage.Task{}, ErrInvalidTaskTransition + } + + updates := map[string]any{ + MetadataKeyAgentRunState: "canceled", + } + if message != "" { + updates[MetadataKeyStatusReason] = message + } + + newMeta, err := mergeTaskMetadata(task.Metadata, updates) + if err != nil { + return storage.Task{}, err + } + + task, err = l.store.UpdateMetadata(ctx, id, newMeta) + if err != nil { + return storage.Task{}, err + } + + return l.store.UpdateStatus(ctx, id, string(StatusCanceled)) +} diff --git a/services/core/internal/workflow/service.go b/services/core/internal/workflow/service.go index 7561b73..ca42a7d 100644 --- a/services/core/internal/workflow/service.go +++ b/services/core/internal/workflow/service.go @@ -11,8 +11,9 @@ import ( ) var ( - ErrInvalidTaskInput = errors.New("invalid task input") - ErrTaskCannotBeEnqueued = errors.New("task cannot be enqueued in current status") + ErrInvalidTaskInput = errors.New("invalid task input") + ErrTaskCannotBeEnqueued = errors.New("task cannot be enqueued in current status") + ErrInvalidTaskTransition = errors.New("invalid task transition") ) type TaskEnqueuer interface { diff --git a/services/core/internal/workflow/service_test.go b/services/core/internal/workflow/service_test.go index e7b54b8..58daa21 100644 --- a/services/core/internal/workflow/service_test.go +++ b/services/core/internal/workflow/service_test.go @@ -71,3 +71,105 @@ func TestNormalizeTaskMetadataRejectsInvalidJSON(t *testing.T) { t.Fatal("expected invalid metadata error") } } + +func TestLifecycleStartTaskAllowsQueuedAndMarksRunning(t *testing.T) { + tests := []struct { + from TaskStatus + to TaskStatus + wanted bool + }{ + {StatusPending, StatusRunning, true}, + {StatusQueued, StatusRunning, true}, + } + + for _, tt := range tests { + got := canTransition(tt.from, tt.to) + if got != tt.wanted { + t.Errorf("canTransition(%s, %s) = %v; want %v", tt.from, tt.to, got, tt.wanted) + } + } +} + +func TestCanTransitionRejectsInvalidLifecycleMove(t *testing.T) { + tests := []struct { + from TaskStatus + to TaskStatus + wanted bool + }{ + {StatusCompleted, StatusRunning, false}, + {StatusPending, StatusCompleted, false}, + {StatusCanceled, StatusRunning, false}, + {StatusRunning, StatusQueued, false}, + } + + for _, tt := range tests { + got := canTransition(tt.from, tt.to) + if got != tt.wanted { + t.Errorf("canTransition(%s, %s) = %v; want %v", tt.from, tt.to, got, tt.wanted) + } + } +} + +func TestMergeTaskMetadataPreservesExistingKeys(t *testing.T) { + current := json.RawMessage(`{"plan_ref":"plan-123","attempt":2}`) + updates := map[string]any{ + "attempt": 3, + "agent_run_state": "running", + } + merged, err := mergeTaskMetadata(current, updates) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var res map[string]any + if err := json.Unmarshal(merged, &res); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + + if res["plan_ref"] != "plan-123" { + t.Errorf("plan_ref was modified: %v", res["plan_ref"]) + } + if res["attempt"] != float64(3) && res["attempt"] != 3 { + t.Errorf("attempt was not updated: %v", res["attempt"]) + } + if res["agent_run_state"] != "running" { + t.Errorf("agent_run_state was not set: %v", res["agent_run_state"]) + } +} + +func TestLifecycleMetadataStartIncrementsAttempt(t *testing.T) { + current1 := json.RawMessage(`{"plan_ref":"plan-123"}`) + attempt1, err := nextTaskAttempt(current1) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if attempt1 != 1 { + t.Errorf("expected attempt 1, got: %d", attempt1) + } + + current2 := json.RawMessage(`{"attempt":2}`) + attempt2, err := nextTaskAttempt(current2) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if attempt2 != 3 { + t.Errorf("expected attempt 3, got: %d", attempt2) + } +} + +func TestLifecycleMetadataRejectsInvalidExistingMetadata(t *testing.T) { + current := json.RawMessage(`{invalid_json}`) + _, err := nextTaskAttempt(current) + if err == nil { + t.Fatal("expected error on invalid metadata, got nil") + } + if err != ErrInvalidTaskInput { + t.Errorf("expected ErrInvalidTaskInput, got: %v", err) + } +} + +func TestCanTransitionAllowsRetryFromFailedToRunning(t *testing.T) { + if !canTransition(StatusFailed, StatusRunning) { + t.Error("expected canTransition(StatusFailed, StatusRunning) to be true, got false") + } +}