feat: notification 및 workflow 기능 개선

- notification 서비스 테스트 추가
- workflow lifecycle 관리 구현
- config, scheduler, workflow 서비스 개선
This commit is contained in:
toki 2026-05-24 19:36:43 +09:00
parent 02a4ced8d7
commit 4b6951cfa8
13 changed files with 947 additions and 67 deletions

View file

@ -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

View file

@ -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
}

View file

@ -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),
}
}

View file

@ -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)
}
}

View file

@ -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
}

View file

@ -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
}

View file

@ -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")
}
}

View file

@ -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,
})
}

View file

@ -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])
}
}

View file

@ -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,
})

View file

@ -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, &currentMap); 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, &currentMap); 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))
}

View file

@ -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 {

View file

@ -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")
}
}