feat: workflow scheduler and lifecycle improvements

This commit is contained in:
toki 2026-06-01 11:03:02 +09:00
parent 6f383f1213
commit cbcbd63e0d
5 changed files with 165 additions and 6 deletions

View file

@ -15,6 +15,7 @@ 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 {
@ -36,6 +37,7 @@ 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 {
@ -334,11 +336,16 @@ func (w *TaskWorker) markFailed(taskID string, err error) {
defer cancel()
msg := err.Error()
failureType := workflow.FailureTypeExecution
if errors.Is(err, context.DeadlineExceeded) {
msg = "timeout"
failureType = workflow.FailureTypeTimeout
}
task, failErr := w.Lifecycle.FailTask(ctx, taskID, msg)
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)

View file

@ -15,6 +15,7 @@ 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"
)
func TestRunAgentTaskCompletesFromMessage(t *testing.T) {
@ -119,6 +120,7 @@ type fakeTaskLifecycle struct {
completed []string
failed []string
failMessages []string
failInputs []workflow.FailureInput
task storage.Task
}
@ -141,16 +143,24 @@ func (f *fakeTaskLifecycle) CompleteTask(ctx context.Context, id string, result
return f.task, nil
}
func (f *fakeTaskLifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
func (f *fakeTaskLifecycle) FailTaskWithMetadata(ctx context.Context, id string, input workflow.FailureInput) (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.failMessages = append(f.failMessages, input.Message)
f.failInputs = append(f.failInputs, input)
f.task.Status = "failed"
return f.task, nil
}
func (f *fakeTaskLifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
return f.FailTaskWithMetadata(ctx, id, workflow.FailureInput{
Message: message,
Type: workflow.FailureTypeExecution,
})
}
func TestWorkCompletesThroughLifecycle(t *testing.T) {
fakeLifecycle := &fakeTaskLifecycle{
task: storage.Task{
@ -266,6 +276,12 @@ func TestWorkMarksTimeoutFailure(t *testing.T) {
if len(fakeLifecycle.failMessages) != 1 || !strings.Contains(fakeLifecycle.failMessages[0], "timeout") {
t.Fatalf("expected fail message containing timeout, got %v", fakeLifecycle.failMessages)
}
if len(fakeLifecycle.failInputs) != 1 {
t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
}
if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeTimeout {
t.Errorf("expected failure type %q, got %q", workflow.FailureTypeTimeout, fakeLifecycle.failInputs[0].Type)
}
}
func TestWorkRetriesAfterFailureState(t *testing.T) {
@ -299,6 +315,13 @@ func TestWorkRetriesAfterFailureState(t *testing.T) {
t.Fatalf("expected task status to be failed, got: %s", fakeLifecycle.task.Status)
}
if len(fakeLifecycle.failInputs) != 1 {
t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
}
if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeExecution {
t.Errorf("expected failure type %q, got %q", workflow.FailureTypeExecution, fakeLifecycle.failInputs[0].Type)
}
worker.Agent = nil
err = worker.Work(context.Background(), job)

View file

@ -206,7 +206,12 @@ func (l *Lifecycle) CompleteTask(ctx context.Context, id string, result json.Raw
return l.store.CompleteTask(ctx, id, result)
}
func (l *Lifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
type FailureInput struct {
Message string `json:"message"`
Type FailureType `json:"type"`
}
func (l *Lifecycle) FailTaskWithMetadata(ctx context.Context, id string, input FailureInput) (storage.Task, error) {
task, err := l.store.GetTask(ctx, id)
if err != nil {
return storage.Task{}, err
@ -216,9 +221,22 @@ func (l *Lifecycle) FailTask(ctx context.Context, id string, message string) (st
return storage.Task{}, ErrInvalidTaskTransition
}
msg := strings.TrimSpace(input.Message)
if msg == "" {
msg = "unknown failure"
}
failType := input.Type
if failType == "" {
failType = FailureTypeExecution
}
updates := map[string]any{
MetadataKeyAgentRunState: "failed",
MetadataKeyStatusReason: message,
MetadataKeyStatusReason: msg,
MetadataKeyWaitType: nil,
MetadataKeyFailureType: string(failType),
MetadataKeyFailedAt: time.Now().UTC().Format(time.RFC3339),
}
newMeta, err := mergeTaskMetadata(task.Metadata, updates)
@ -231,9 +249,17 @@ func (l *Lifecycle) FailTask(ctx context.Context, id string, message string) (st
return storage.Task{}, err
}
return l.store.FailTask(ctx, id, message)
return l.store.FailTask(ctx, id, msg)
}
func (l *Lifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
return l.FailTaskWithMetadata(ctx, id, FailureInput{
Message: message,
Type: FailureTypeExecution,
})
}
func (l *Lifecycle) CancelTask(ctx context.Context, id string, message string) (storage.Task, error) {
task, err := l.store.GetTask(ctx, id)
if err != nil {

View file

@ -28,12 +28,23 @@ type ExternalRefInput struct {
Metadata json.RawMessage `json:"metadata"`
}
type FailureType string
const (
FailureTypeTimeout FailureType = "timeout"
FailureTypeExecution FailureType = "execution"
FailureTypeEnqueue FailureType = "enqueue"
)
const (
MetadataKeyAgentRunState = "agent_run_state"
MetadataKeyAgentPhase = "agent_phase"
MetadataKeyWaitType = "wait_type"
MetadataKeyStatusReason = "status_reason"
MetadataKeyFailureType = "failure_type"
MetadataKeyFailedAt = "failed_at"
MetadataKeyLastHeartbeat = "last_heartbeat_at"
MetadataKeyPlanRef = "plan_ref"
MetadataKeyAttempt = "attempt"
)

View file

@ -498,6 +498,98 @@ func TestLifecycleStartCompleteFailTransitionsWithStore(t *testing.T) {
})
}
func TestLifecycleFailTaskStoresFailureMetadataAndError(t *testing.T) {
ctx := context.Background()
store := newFakeTaskStore()
store.tasks["task-fail-meta"] = storage.Task{
ID: "task-fail-meta",
Status: string(StatusRunning),
Metadata: json.RawMessage(`{"attempt":1,"agent_run_state":"running"}`),
}
l := &Lifecycle{store: store}
// Fail with FailTaskWithMetadata
task, err := l.FailTaskWithMetadata(ctx, "task-fail-meta", FailureInput{
Message: "timeout occurred",
Type: FailureTypeTimeout,
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if task.Status != string(StatusFailed) {
t.Errorf("expected status failed, got: %s", task.Status)
}
if task.Error == nil || *task.Error != "timeout occurred" {
t.Errorf("expected error message 'timeout occurred', got: %v", task.Error)
}
var meta map[string]any
if err := json.Unmarshal(task.Metadata, &meta); err != nil {
t.Fatalf("failed to unmarshal metadata: %v", err)
}
if meta[MetadataKeyAgentRunState] != "failed" {
t.Errorf("expected agent_run_state failed, got: %v", meta[MetadataKeyAgentRunState])
}
if meta[MetadataKeyStatusReason] != "timeout occurred" {
t.Errorf("expected status_reason 'timeout occurred', got: %v", meta[MetadataKeyStatusReason])
}
if meta[MetadataKeyFailureType] != string(FailureTypeTimeout) {
t.Errorf("expected failure_type 'timeout', got: %v", meta[MetadataKeyFailureType])
}
if _, ok := meta[MetadataKeyFailedAt]; !ok {
t.Errorf("expected failed_at to be set in metadata")
}
// Test trimming and empty/default behavior
store.tasks["task-fail-empty"] = storage.Task{
ID: "task-fail-empty",
Status: string(StatusRunning),
Metadata: json.RawMessage(`{}`),
}
task2, err := l.FailTaskWithMetadata(ctx, "task-fail-empty", FailureInput{
Message: " ",
Type: "",
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
var meta2 map[string]any
if err := json.Unmarshal(task2.Metadata, &meta2); err != nil {
t.Fatalf("failed to unmarshal metadata: %v", err)
}
if meta2[MetadataKeyStatusReason] != "unknown failure" {
t.Errorf("expected status_reason fallback 'unknown failure', got: %v", meta2[MetadataKeyStatusReason])
}
if meta2[MetadataKeyFailureType] != string(FailureTypeExecution) {
t.Errorf("expected default failure_type 'execution', got: %v", meta2[MetadataKeyFailureType])
}
}
func TestLifecycleFailTaskClearsWaitType(t *testing.T) {
ctx := context.Background()
store := newFakeTaskStore()
store.tasks["task-wait"] = storage.Task{
ID: "task-wait",
Status: string(StatusRunning),
Metadata: json.RawMessage(`{"attempt":1,"agent_run_state":"running","wait_type":"user_input"}`),
}
l := &Lifecycle{store: store}
task, err := l.FailTask(ctx, "task-wait", "failed by agent")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
var meta map[string]any
if err := json.Unmarshal(task.Metadata, &meta); err != nil {
t.Fatalf("failed to unmarshal metadata: %v", err)
}
if val, ok := meta[MetadataKeyWaitType]; ok && val != nil {
t.Errorf("expected wait_type to be cleared, got: %v", val)
}
}
type fakeTaskEnqueuer struct {
err error
calls []string