diff --git a/internal/session/todo.go b/internal/session/todo.go index d4b4de0e..cbb90fcf 100644 --- a/internal/session/todo.go +++ b/internal/session/todo.go @@ -9,7 +9,7 @@ import ( ) // CurrentTodoVersion 表示当前 Todo 结构版本。 -const CurrentTodoVersion = 2 +const CurrentTodoVersion = 3 // TodoStatus 表示 Todo 项的状态枚举。 type TodoStatus string @@ -260,10 +260,14 @@ func (s *Session) UpdateTodo(id string, patch TodoPatch, expectedRevision int64) // ClaimTodo 用于 SubAgent 领取 Todo，并设置 owner 与执行中状态。 func (s *Session) ClaimTodo(id string, ownerType string, ownerID string, expectedRevision int64) error { status := TodoStatusInProgress + failureReason := "" + nextRetryAt := time.Time{} patch := TodoPatch{ - Status: &status, - OwnerType: &ownerType, - OwnerID: &ownerID, + Status: &status, + OwnerType: &ownerType, + OwnerID: &ownerID, + FailureReason: &failureReason, + NextRetryAt: &nextRetryAt, } return s.UpdateTodo(id, patch, expectedRevision) } @@ -271,9 +275,15 @@ func (s *Session) ClaimTodo(id string, ownerType string, ownerID string, expecte // CompleteTodo 将 Todo 标记为完成并写入产物列表。 func (s *Session) CompleteTodo(id string, artifacts []string, expectedRevision int64) error { status := TodoStatusCompleted + failureReason := "" + retryCount := 0 + nextRetryAt := time.Time{} patch := TodoPatch{ - Status: &status, - Artifacts: &artifacts, + Status: &status, + Artifacts: &artifacts, + FailureReason: &failureReason, + RetryCount: &retryCount, + NextRetryAt: &nextRetryAt, } return s.UpdateTodo(id, patch, expectedRevision) } @@ -281,9 +291,11 @@ func (s *Session) CompleteTodo(id string, artifacts []string, expectedRevision i // FailTodo 将 Todo 标记为失败并记录失败原因。 func (s *Session) FailTodo(id string, reason string, expectedRevision int64) error { status := TodoStatusFailed + nextRetryAt := time.Time{} patch := TodoPatch{ Status: &status, FailureReason: &reason, + NextRetryAt: &nextRetryAt, } return s.UpdateTodo(id, patch, expectedRevision) } @@ -395,6 +407,15 @@ func normalizeTodoItem(item TodoItem) (TodoItem, error) { item.Acceptance = normalizeTodoTextList(item.Acceptance) item.Artifacts = normalizeTodoTextList(item.Artifacts) item.FailureReason = strings.TrimSpace(item.FailureReason) + if item.RetryCount < 0 { + item.RetryCount = 0 + } + if item.RetryLimit < 0 { + item.RetryLimit = 0 + } + if !item.NextRetryAt.IsZero() { + item.NextRetryAt = item.NextRetryAt.UTC() + } if item.Status == "" { item.Status = TodoStatusPending } @@ -413,9 +434,12 @@ func normalizeTodoItem(item TodoItem) (TodoItem, error) { return TodoItem{}, fmt.Errorf("session: invalid todo owner_type %q", item.OwnerType) } - if item.Status != TodoStatusFailed { + if item.Status != TodoStatusFailed && item.Status != TodoStatusPending && item.Status != TodoStatusBlocked { item.FailureReason = "" } + if item.Status != TodoStatusPending && item.Status != TodoStatusFailed { + item.NextRetryAt = time.Time{} + } return item, nil } diff --git a/internal/session/todo_test.go b/internal/session/todo_test.go index d669467c..57c4ccce 100644 --- a/internal/session/todo_test.go +++ b/internal/session/todo_test.go @@ -4,6 +4,7 @@ import ( "errors" "strings" "testing" + "time" ) func TestTodoStatusValidAndTransition(t *testing.T) { @@ -194,9 +195,12 @@ func TestSessionClaimCompleteFail(t *testing.T) { t.Fatalf("AddTodo(base) error = %v", err) } if err := session.AddTodo(TodoItem{ - ID: "task", - Content: "task", - Dependencies: []string{"base"}, + ID: "task", + Content: "task", + Dependencies: []string{"base"}, + FailureReason: "previous failure", + RetryCount: 2, + NextRetryAt: time.Now().Add(2 * time.Minute), }); err != nil { t.Fatalf("AddTodo(task) error = %v", err) } @@ -208,6 +212,17 @@ func TestSessionClaimCompleteFail(t *testing.T) { if task.OwnerType != TodoOwnerTypeSubAgent || task.OwnerID != "worker-1" { t.Fatalf("unexpected owner after claim: %+v", task) } + if task.FailureReason != "" || !task.NextRetryAt.IsZero() { + t.Fatalf("claim should clear failure/next_retry_at, got %+v", task) + } + retryCount := 5 + if err := session.UpdateTodo("task", TodoPatch{RetryCount: &retryCount}, task.Revision); err != nil { + t.Fatalf("UpdateTodo(retry_count) error = %v", err) + } + task, _ = session.FindTodo("task") + if task.RetryCount != 5 { + t.Fatalf("retry_count = %d, want 5", task.RetryCount) + } if err := session.FailTodo("task", "compile failed", task.Revision); err != nil { t.Fatalf("FailTodo error = %v", err) @@ -216,6 +231,9 @@ func TestSessionClaimCompleteFail(t *testing.T) { if task.Status != TodoStatusFailed || task.FailureReason != "compile failed" { t.Fatalf("unexpected fail state: %+v", task) } + if !task.NextRetryAt.IsZero() { + t.Fatalf("FailTodo should clear next_retry_at, got %+v", task) + } if err := session.SetTodoStatus("task", TodoStatusInProgress, task.Revision); err == nil || !errors.Is(err, ErrInvalidTransition) { t.Fatalf("terminal transition error = %v, want invalid transition", err) @@ -356,6 +374,24 @@ func TestTodoInternalHelpers(t *testing.T) { if got := normalizeTodoDependencies([]string{" a ", "a", "b"}); len(got) != 2 { t.Fatalf("normalizeTodoDependencies unexpected result: %+v", got) } + + normalized, err := normalizeTodoItem(TodoItem{ + ID: "n1", + Content: "content", + Status: TodoStatusBlocked, + FailureReason: " retry failed ", + RetryCount: -3, + RetryLimit: -2, + }) + if err != nil { + t.Fatalf("normalizeTodoItem error = %v", err) + } + if normalized.FailureReason != "retry failed" { + t.Fatalf("blocked todo should keep failure reason, got %q", normalized.FailureReason) + } + if normalized.RetryCount != 0 || normalized.RetryLimit != 0 { + t.Fatalf("negative retry fields should be normalized to 0, got count=%d limit=%d", normalized.RetryCount, normalized.RetryLimit) + } } func TestApplyTodoPatchCoverage(t *testing.T) { @@ -381,6 +417,9 @@ func TestApplyTodoPatchCoverage(t *testing.T) { acceptance := []string{"ok"} artifacts := []string{"a.txt"} reason := "boom" + retryCount := 1 + retryLimit := 3 + nextRetryAt := time.Now().Add(3 * time.Minute).UTC() status := TodoStatusInProgress patch := TodoPatch{ Content: &content, @@ -391,6 +430,9 @@ func TestApplyTodoPatchCoverage(t *testing.T) { Acceptance: &acceptance, Artifacts: &artifacts, FailureReason: &reason, + RetryCount: &retryCount, + RetryLimit: &retryLimit, + NextRetryAt: &nextRetryAt, Status: &status, } @@ -407,6 +449,26 @@ func TestApplyTodoPatchCoverage(t *testing.T) { if next.FailureReason != "" { t.Fatalf("non-failed status should clear failure reason, got %q", next.FailureReason) } + if next.RetryCount != 1 || next.RetryLimit != 3 { + t.Fatalf("retry fields not applied, got %+v", next) + } + if !next.NextRetryAt.IsZero() { + t.Fatalf("in_progress should clear next_retry_at, got %+v", next) + } + + blocked := TodoStatusBlocked + blockedReason := "retry later" + blockedPatch := TodoPatch{ + Status: &blocked, + FailureReason: &blockedReason, + } + blockedNext, err := applyTodoPatch(base, blockedPatch) + if err != nil { + t.Fatalf("applyTodoPatch(blocked) error = %v", err) + } + if blockedNext.FailureReason != blockedReason { + t.Fatalf("blocked status should preserve failure reason, got %q", blockedNext.FailureReason) + } invalidStatus := TodoStatus("paused") if _, err := applyTodoPatch(base, TodoPatch{Status: &invalidStatus}); err == nil { diff --git a/internal/subagent/scheduler.go b/internal/subagent/scheduler.go new file mode 100644 index 00000000..fd116ba6 --- /dev/null +++ b/internal/subagent/scheduler.go @@ -0,0 +1,699 @@ +package subagent + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "time" + + agentsession "neo-code/internal/session" +) + +// Scheduler 负责按 Todo DAG 依赖关系驱动子代理任务执行与状态回写。 +type Scheduler struct { + store TodoStore + factory Factory + cfg SchedulerConfig +} + +type runningTask struct { + id string + attempt int +} + +type taskOutcome struct { + id string + attempt int + result Result + err error +} + +type schedulerState struct { + running map[string]runningTask + readySince map[string]time.Time + started map[string]int + outcomes chan taskOutcome +} + +// NewScheduler 创建 Task DAG 调度器，并校验核心依赖是否可用。 +func NewScheduler(store TodoStore, factory Factory, cfg SchedulerConfig) (*Scheduler, error) { + if store == nil { + return nil, errorsf("scheduler todo store is nil") + } + if factory == nil { + return nil, errorsf("scheduler factory is nil") + } + return &Scheduler{ + store: store, + factory: factory, + cfg: cfg.normalize(), + }, nil +} + +// Run 执行一次调度轮次，直到所有任务终态、上下文取消或出现不可恢复错误。 +func (s *Scheduler) Run(ctx context.Context) (ScheduleResult, error) { + if err := ctx.Err(); err != nil { + return ScheduleResult{}, err + } + + initial := s.store.ListTodos() + graph, err := buildTaskGraph(initial) + if err != nil { + return ScheduleResult{}, err + } + + now := s.cfg.Clock() + result := ScheduleResult{ + StartedAt: now, + Total: len(graph.order), + Retried: make(map[string]int), + } + state := newSchedulerState(s.cfg.MaxConcurrency) + finalize := func(current ScheduleResult) ScheduleResult { + current.EndedAt = s.cfg.Clock() + current.BlockedLeft = collectBlockedLeft(graph.order, s.store.ListTodos(), state.running) + s.emit(SchedulerEvent{ + Type: SchedulerEventFinished, + QueueSize: len(current.BlockedLeft), + Running: len(state.running), + At: current.EndedAt, + }) + return current + } + + for { + if err := ctx.Err(); err != nil { + s.cancelRunningTodos(state, err) + return finalize(result), err + } + + snapshot := mapTodosByID(s.store.ListTodos()) + ready, err := s.collectReadyTasks(snapshot, graph, state) + if err != nil { + s.cancelRunningTodos(state, err) + return finalize(result), err + } + + s.pruneReadySince(state, ready) + s.sortReadyTasks(ready, state.readySince) + + started, err := s.startReadyTasks(ctx, ready, state) + if err != nil { + s.cancelRunningTodos(state, err) + return finalize(result), err + } + if started > 0 { + continue + } + + if len(state.running) == 0 { + if !hasSchedulablePotential(graph.order, snapshot) { + return finalize(result), nil + } + if err := waitWithContext(ctx, s.nextPollDelay(snapshot)); err != nil { + s.cancelRunningTodos(state, err) + return finalize(result), err + } + continue + } + + if err := s.handleOneOutcome(ctx, state, &result); err != nil { + s.cancelRunningTodos(state, err) + return finalize(result), err + } + } +} + +// newSchedulerState 初始化调度运行态，避免循环中重复分配映射与通道。 +func newSchedulerState(maxConcurrency int) *schedulerState { + buffer := maxConcurrency + if buffer < 1 { + buffer = 1 + } + return &schedulerState{ + running: make(map[string]runningTask, maxConcurrency), + readySince: make(map[string]time.Time, 32), + started: make(map[string]int, 32), + outcomes: make(chan taskOutcome, buffer), + } +} + +// collectReadyTasks 基于依赖关系、重试窗口与当前状态筛选可执行任务。 +func (s *Scheduler) collectReadyTasks( + snapshot map[string]agentsession.TodoItem, + graph *taskGraph, + state *schedulerState, +) ([]agentsession.TodoItem, error) { + now := s.cfg.Clock() + ready := make([]agentsession.TodoItem, 0, len(graph.order)) + + for _, id := range graph.order { + item, ok := snapshot[id] + if !ok || item.Status.IsTerminal() { + continue + } + if _, running := state.running[id]; running { + continue + } + + depsSatisfied := dependenciesCompleted(item, snapshot) + if !depsSatisfied { + if err := s.ensureBlocked(item, "dependency_unmet"); err != nil { + return nil, err + } + continue + } + + item, ok, err := s.ensureReadyStatus(item) + if err != nil { + return nil, err + } + if !ok { + continue + } + + if !item.NextRetryAt.IsZero() && now.Before(item.NextRetryAt) { + s.emit(SchedulerEvent{ + Type: SchedulerEventBlocked, + TaskID: item.ID, + Reason: "retry_backoff", + QueueSize: len(ready), + Running: len(state.running), + At: now, + }) + continue + } + ready = append(ready, item.Clone()) + if _, exists := state.readySince[item.ID]; !exists { + state.readySince[item.ID] = now + } + } + return ready, nil +} + +// ensureBlocked 将未满足依赖的任务收敛到 blocked 状态并发出可观测事件。 +func (s *Scheduler) ensureBlocked(item agentsession.TodoItem, reason string) error { + if item.Status != agentsession.TodoStatusBlocked { + status := agentsession.TodoStatusBlocked + patch := agentsession.TodoPatch{Status: &status} + if err := s.store.UpdateTodo(item.ID, patch, item.Revision); err != nil && !isRevisionConflict(err) { + return fmt.Errorf("subagent: mark todo blocked: %w", err) + } + } + + s.emit(SchedulerEvent{ + Type: SchedulerEventBlocked, + TaskID: item.ID, + Reason: reason, + Running: 0, + At: s.cfg.Clock(), + }) + return nil +} + +// ensureReadyStatus 处理 blocked 到 pending 的解锁与可执行状态判定。 +func (s *Scheduler) ensureReadyStatus(item agentsession.TodoItem) (agentsession.TodoItem, bool, error) { + switch item.Status { + case agentsession.TodoStatusPending, agentsession.TodoStatusInProgress: + return item, true, nil + case agentsession.TodoStatusBlocked: + if !item.NextRetryAt.IsZero() && s.cfg.Clock().Before(item.NextRetryAt) { + return item, false, nil + } + status := agentsession.TodoStatusPending + zeroRetry := time.Time{} + patch := agentsession.TodoPatch{ + Status: &status, + NextRetryAt: &zeroRetry, + } + if err := s.store.UpdateTodo(item.ID, patch, item.Revision); err != nil { + if isRevisionConflict(err) { + return item, false, nil + } + return item, false, fmt.Errorf("subagent: unlock blocked todo: %w", err) + } + next, ok := s.store.FindTodo(item.ID) + if !ok { + return item, false, nil + } + return next, true, nil + default: + return item, false, nil + } +} + +// startReadyTasks 在并发上限内领取并启动可执行任务。 +func (s *Scheduler) startReadyTasks(ctx context.Context, ready []agentsession.TodoItem, state *schedulerState) (int, error) { + if len(ready) == 0 { + return 0, nil + } + started := 0 + for _, item := range ready { + if len(state.running) >= s.cfg.MaxConcurrency { + break + } + if _, exists := state.running[item.ID]; exists { + continue + } + attempt := item.RetryCount + 1 + workerID := fmt.Sprintf("%s-%s", s.cfg.WorkerIDPrefix, item.ID) + if err := s.store.ClaimTodo(item.ID, agentsession.TodoOwnerTypeSubAgent, workerID, item.Revision); err != nil { + if isRevisionConflict(err) { + continue + } + return started, fmt.Errorf("subagent: claim todo %q: %w", item.ID, err) + } + + role := s.cfg.RoleSelector(item) + budget := s.cfg.BudgetSelector(item, s.cfg.DefaultBudget).normalize(s.cfg.DefaultBudget) + capability := s.cfg.Capabilities(item).normalize() + task := Task{ + ID: item.ID, + Goal: strings.TrimSpace(item.Content), + ExpectedOutput: strings.Join(item.Acceptance, "\n"), + } + + state.running[item.ID] = runningTask{id: item.ID, attempt: attempt} + state.started[item.ID] = attempt + started++ + + s.emit(SchedulerEvent{ + Type: SchedulerEventQueued, + TaskID: item.ID, + Attempt: attempt, + QueueSize: len(ready), + Running: len(state.running), + At: s.cfg.Clock(), + }) + s.emit(SchedulerEvent{ + Type: SchedulerEventClaimed, + TaskID: item.ID, + Attempt: attempt, + Running: len(state.running), + At: s.cfg.Clock(), + }) + s.emit(SchedulerEvent{ + Type: SchedulerEventRunning, + TaskID: item.ID, + Attempt: attempt, + Running: len(state.running), + At: s.cfg.Clock(), + }) + + go s.executeTaskAsync(ctx, state.outcomes, taskOutcome{ + id: item.ID, + attempt: attempt, + }, scheduleTaskInput{ + Task: task, + Role: role, + Budget: budget, + Capability: capability, + }) + } + return started, nil +} + +// executeTaskAsync 在独立 goroutine 中执行任务，并把结果回传给调度主循环。 +func (s *Scheduler) executeTaskAsync( + parent context.Context, + out chan<- taskOutcome, + base taskOutcome, + input scheduleTaskInput, +) { + execCtx := parent + cancel := func() {} + if input.Budget.Timeout > 0 { + execCtx, cancel = context.WithTimeout(parent, input.Budget.Timeout) + } + defer cancel() + + result, err := executeTaskWithFactory(execCtx, s.factory, input) + base.result = result + base.err = err + select { + case out <- base: + default: + } +} + +// handleOneOutcome 消费单个任务执行结果并完成成功/失败/重试回写。 +func (s *Scheduler) handleOneOutcome(ctx context.Context, state *schedulerState, summary *ScheduleResult) error { + select { + case <-ctx.Done(): + return ctx.Err() + case outcome := <-state.outcomes: + delete(state.running, outcome.id) + delete(state.readySince, outcome.id) + return s.applyOutcome(outcome, state, summary) + case <-time.After(s.cfg.PollInterval): + return nil + } +} + +// applyOutcome 根据执行结果写回 Todo 状态，并更新调度统计与事件。 +func (s *Scheduler) applyOutcome(outcome taskOutcome, state *schedulerState, summary *ScheduleResult) error { + current, ok := s.store.FindTodo(outcome.id) + if !ok { + return nil + } + if current.Status.IsTerminal() { + return nil + } + + if outcome.err == nil && outcome.result.State == StateSucceeded { + artifacts := append([]string(nil), outcome.result.Output.Artifacts...) + if err := s.store.CompleteTodo(current.ID, artifacts, current.Revision); err != nil { + if isRevisionConflict(err) { + return nil + } + return fmt.Errorf("subagent: complete todo %q: %w", current.ID, err) + } + appendUniqueString(&summary.Succeeded, current.ID) + s.emit(SchedulerEvent{ + Type: SchedulerEventCompleted, + TaskID: current.ID, + Attempt: outcome.attempt, + Running: len(state.running), + At: s.cfg.Clock(), + }) + return nil + } + + return s.handleTaskFailure(current, outcome, state, summary) +} + +// handleTaskFailure 处理失败回写，按重试预算决定重排或终态失败。 +func (s *Scheduler) handleTaskFailure( + current agentsession.TodoItem, + outcome taskOutcome, + state *schedulerState, + summary *ScheduleResult, +) error { + reason := resolveTaskFailureReason(outcome) + retryLimit := current.RetryLimit + if retryLimit <= 0 { + retryLimit = s.cfg.MaxRetries + } + nextRetryCount := current.RetryCount + 1 + + if nextRetryCount <= retryLimit { + backoff := s.cfg.Backoff(nextRetryCount) + nextRetryAt := s.cfg.Clock().Add(backoff) + status := agentsession.TodoStatusBlocked + patch := agentsession.TodoPatch{ + Status: &status, + FailureReason: &reason, + RetryCount: &nextRetryCount, + RetryLimit: &retryLimit, + NextRetryAt: &nextRetryAt, + } + if err := s.store.UpdateTodo(current.ID, patch, current.Revision); err != nil { + if isRevisionConflict(err) { + return nil + } + return fmt.Errorf("subagent: schedule retry for todo %q: %w", current.ID, err) + } + + summary.Retried[current.ID] = nextRetryCount + s.emit(SchedulerEvent{ + Type: SchedulerEventRetryScheduled, + TaskID: current.ID, + Attempt: nextRetryCount, + Reason: reason, + Running: len(state.running), + At: s.cfg.Clock(), + }) + return nil + } + + status := agentsession.TodoStatusFailed + zeroRetryAt := time.Time{} + patch := agentsession.TodoPatch{ + Status: &status, + FailureReason: &reason, + RetryCount: &nextRetryCount, + RetryLimit: &retryLimit, + NextRetryAt: &zeroRetryAt, + } + if err := s.store.UpdateTodo(current.ID, patch, current.Revision); err != nil { + if isRevisionConflict(err) { + return nil + } + return fmt.Errorf("subagent: fail todo %q: %w", current.ID, err) + } + + appendUniqueString(&summary.Failed, current.ID) + s.emit(SchedulerEvent{ + Type: SchedulerEventFailed, + TaskID: current.ID, + Attempt: nextRetryCount, + Reason: reason, + Running: len(state.running), + At: s.cfg.Clock(), + }) + return nil +} + +// cancelRunningTodos 在调度中断时把仍在执行的任务统一回写为 canceled。 +func (s *Scheduler) cancelRunningTodos(state *schedulerState, cause error) { + if state == nil || len(state.running) == 0 { + return + } + reason := "scheduler canceled" + if cause != nil && strings.TrimSpace(cause.Error()) != "" { + reason = strings.TrimSpace(cause.Error()) + } + status := agentsession.TodoStatusCanceled + for id := range state.running { + item, ok := s.store.FindTodo(id) + if !ok || item.Status.IsTerminal() { + continue + } + patch := agentsession.TodoPatch{ + Status: &status, + FailureReason: &reason, + } + if err := s.store.UpdateTodo(item.ID, patch, item.Revision); err != nil && !isRevisionConflict(err) { + continue + } + s.emit(SchedulerEvent{ + Type: SchedulerEventCanceled, + TaskID: item.ID, + Reason: reason, + Running: len(state.running), + At: s.cfg.Clock(), + }) + } +} + +// sortReadyTasks 按优先级与等待时间排序，兼顾高优先级与公平性。 +func (s *Scheduler) sortReadyTasks(ready []agentsession.TodoItem, readySince map[string]time.Time) { + now := s.cfg.Clock() + agingWindow := 5 * s.cfg.PollInterval + if agingWindow <= 0 { + agingWindow = time.Second + } + sort.SliceStable(ready, func(i, j int) bool { + left := ready[i] + right := ready[j] + lp := effectivePriority(left, readySince[left.ID], now, agingWindow) + rp := effectivePriority(right, readySince[right.ID], now, agingWindow) + if lp != rp { + return lp > rp + } + if !left.CreatedAt.Equal(right.CreatedAt) { + return left.CreatedAt.Before(right.CreatedAt) + } + return left.ID < right.ID + }) +} + +// pruneReadySince 清理当前不可调度任务的 ready 时间戳，避免无界增长。 +func (s *Scheduler) pruneReadySince(state *schedulerState, ready []agentsession.TodoItem) { + if len(state.readySince) == 0 { + return + } + allowed := make(map[string]struct{}, len(ready)) + for _, item := range ready { + allowed[item.ID] = struct{}{} + } + for id := range state.readySince { + if _, ok := allowed[id]; !ok { + delete(state.readySince, id) + } + } +} + +// nextPollDelay 计算下一次轮询等待时间，优先对齐最近重试窗口。 +func (s *Scheduler) nextPollDelay(snapshot map[string]agentsession.TodoItem) time.Duration { + now := s.cfg.Clock() + minDelay := s.cfg.PollInterval + if minDelay <= 0 { + minDelay = 200 * time.Millisecond + } + for _, item := range snapshot { + if item.Status.IsTerminal() || item.NextRetryAt.IsZero() { + continue + } + if !item.NextRetryAt.After(now) { + return 0 + } + delay := item.NextRetryAt.Sub(now) + if delay < minDelay { + minDelay = delay + } + } + return minDelay +} + +// emit 发射调度事件，统一补齐时间戳并调用观察器。 +func (s *Scheduler) emit(event SchedulerEvent) { + if event.At.IsZero() { + event.At = s.cfg.Clock() + } + s.cfg.Observer(event) +} + +// mapTodosByID 将 todo 列表转为 ID 索引映射，便于依赖与状态查询。 +func mapTodosByID(items []agentsession.TodoItem) map[string]agentsession.TodoItem { + result := make(map[string]agentsession.TodoItem, len(items)) + for _, item := range items { + result[item.ID] = item + } + return result +} + +// dependenciesCompleted 判断任务依赖是否全部处于 completed 状态。 +func dependenciesCompleted(item agentsession.TodoItem, byID map[string]agentsession.TodoItem) bool { + for _, depID := range item.Dependencies { + dependency, ok := byID[depID] + if !ok || dependency.Status != agentsession.TodoStatusCompleted { + return false + } + } + return true +} + +// hasSchedulablePotential 判断当前非终态任务是否仍可能通过调度推进到可执行状态。 +func hasSchedulablePotential(order []string, byID map[string]agentsession.TodoItem) bool { + memo := make(map[string]bool, len(byID)) + visiting := make(map[string]bool, len(byID)) + + var satisfiable func(id string) bool + satisfiable = func(id string) bool { + item, ok := byID[id] + if !ok { + return false + } + if item.Status == agentsession.TodoStatusCompleted { + return true + } + if item.Status == agentsession.TodoStatusFailed || item.Status == agentsession.TodoStatusCanceled { + return false + } + if known, ok := memo[id]; ok { + return known + } + if visiting[id] { + return false + } + + visiting[id] = true + defer delete(visiting, id) + for _, dependencyID := range item.Dependencies { + if !satisfiable(dependencyID) { + memo[id] = false + return false + } + } + memo[id] = true + return true + } + + for _, id := range order { + item, ok := byID[id] + if !ok || item.Status.IsTerminal() { + continue + } + if satisfiable(id) { + return true + } + } + return false +} + +// collectBlockedLeft 汇总结束时仍处于非终态的任务 ID。 +func collectBlockedLeft(order []string, items []agentsession.TodoItem, running map[string]runningTask) []string { + byID := mapTodosByID(items) + left := make([]string, 0) + for _, id := range order { + if _, ok := running[id]; ok { + left = append(left, id) + continue + } + item, ok := byID[id] + if !ok || item.Status.IsTerminal() { + continue + } + left = append(left, id) + } + return left +} + +// effectivePriority 基于原始优先级与等待时长计算动态调度优先级。 +func effectivePriority(item agentsession.TodoItem, readySince time.Time, now time.Time, agingWindow time.Duration) int { + priority := item.Priority + if readySince.IsZero() || agingWindow <= 0 { + return priority + } + if now.Before(readySince) { + return priority + } + ageBoost := int(now.Sub(readySince) / agingWindow) + return priority + ageBoost +} + +// appendUniqueString 追加去重字符串，避免调度统计结果重复计数。 +func appendUniqueString(dst *[]string, value string) { + for _, current := range *dst { + if current == value { + return + } + } + *dst = append(*dst, value) +} + +// resolveTaskFailureReason 统一提取失败原因，保证回写文本稳定可读。 +func resolveTaskFailureReason(outcome taskOutcome) string { + if outcome.err != nil { + if text := strings.TrimSpace(outcome.err.Error()); text != "" { + return text + } + } + if text := strings.TrimSpace(outcome.result.Error); text != "" { + return text + } + return "subagent task execution failed" +} + +// isRevisionConflict 判断错误是否为 revision 竞争冲突，便于调度层重试。 +func isRevisionConflict(err error) bool { + return errors.Is(err, agentsession.ErrRevisionConflict) +} + +// waitWithContext 在保留可取消语义的前提下等待指定时长。 +func waitWithContext(ctx context.Context, wait time.Duration) error { + if wait <= 0 { + return nil + } + timer := time.NewTimer(wait) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} diff --git a/internal/subagent/scheduler_executor.go b/internal/subagent/scheduler_executor.go new file mode 100644 index 00000000..807fd6a5 --- /dev/null +++ b/internal/subagent/scheduler_executor.go @@ -0,0 +1,73 @@ +package subagent + +import ( + "context" + "errors" + "strings" +) + +// scheduleTaskInput 描述调度器发起的单任务执行请求。 +type scheduleTaskInput struct { + Task Task + Role Role + Budget Budget + Capability Capability +} + +// executeTaskWithFactory 通过 WorkerFactory 执行单任务并返回标准化结果。 +func executeTaskWithFactory(ctx context.Context, factory Factory, input scheduleTaskInput) (Result, error) { + worker, err := factory.Create(input.Role) + if err != nil { + return Result{}, err + } + if err := worker.Start(input.Task, input.Budget, input.Capability); err != nil { + return Result{}, err + } + + for { + stepResult, stepErr := worker.Step(ctx) + if stepErr != nil { + if errors.Is(stepErr, context.Canceled) || errors.Is(stepErr, context.DeadlineExceeded) { + _ = worker.Stop(StopReasonCanceled) + result, resultErr := worker.Result() + if resultErr == nil { + return result, stepErr + } + return Result{ + Role: input.Role, + TaskID: input.Task.ID, + State: StateCanceled, + StopReason: StopReasonCanceled, + Error: strings.TrimSpace(stepErr.Error()), + }, stepErr + } + + result, resultErr := worker.Result() + if resultErr == nil { + return result, stepErr + } + return Result{ + Role: input.Role, + TaskID: input.Task.ID, + State: StateFailed, + StopReason: StopReasonError, + Error: strings.TrimSpace(stepErr.Error()), + }, stepErr + } + if !stepResult.Done { + continue + } + + result, resultErr := worker.Result() + if resultErr != nil { + return Result{}, resultErr + } + if result.State == StateSucceeded { + return result, nil + } + if strings.TrimSpace(result.Error) != "" { + return result, errors.New(result.Error) + } + return result, errorsf("worker finished with state=%s stop_reason=%s", result.State, result.StopReason) + } +} diff --git a/internal/subagent/scheduler_graph.go b/internal/subagent/scheduler_graph.go new file mode 100644 index 00000000..bd38d344 --- /dev/null +++ b/internal/subagent/scheduler_graph.go @@ -0,0 +1,104 @@ +package subagent + +import ( + "fmt" + "slices" + "strings" + + agentsession "neo-code/internal/session" +) + +type taskNode struct { + todo agentsession.TodoItem + dependents []string +} + +// taskGraph 保存调度时的 Todo DAG 快照，供就绪判定和依赖推进复用。 +type taskGraph struct { + nodes map[string]*taskNode + order []string +} + +// buildTaskGraph 从当前 Todo 列表构建 DAG，并校验依赖合法性和无环约束。 +func buildTaskGraph(items []agentsession.TodoItem) (*taskGraph, error) { + graph := &taskGraph{ + nodes: make(map[string]*taskNode, len(items)), + order: make([]string, 0, len(items)), + } + for _, item := range items { + id := strings.TrimSpace(item.ID) + if id == "" { + return nil, errorsf("scheduler graph contains empty todo id") + } + if _, exists := graph.nodes[id]; exists { + return nil, errorsf("scheduler graph contains duplicate todo id %q", id) + } + graph.nodes[id] = &taskNode{todo: item.Clone()} + graph.order = append(graph.order, id) + } + for _, id := range graph.order { + node := graph.nodes[id] + for _, dep := range node.todo.Dependencies { + dependencyID := strings.TrimSpace(dep) + dependency, ok := graph.nodes[dependencyID] + if !ok { + return nil, errorsf("scheduler graph todo %q references unknown dependency %q", id, dep) + } + dependency.dependents = append(dependency.dependents, id) + } + } + if err := detectTaskCycle(graph); err != nil { + return nil, err + } + return graph, nil +} + +// detectTaskCycle 使用 Kahn 算法检测图中环路，避免调度阶段死锁。 +func detectTaskCycle(graph *taskGraph) error { + inDegree := make(map[string]int, len(graph.nodes)) + for _, id := range graph.order { + inDegree[id] = 0 + } + for _, id := range graph.order { + node := graph.nodes[id] + for _, dep := range node.todo.Dependencies { + inDegree[id]++ + if strings.TrimSpace(dep) == id { + return fmt.Errorf("%w: %q", agentsession.ErrCyclicDependency, id) + } + } + } + + queue := make([]string, 0, len(graph.order)) + for _, id := range graph.order { + if inDegree[id] == 0 { + queue = append(queue, id) + } + } + + visited := 0 + for len(queue) > 0 { + current := queue[0] + queue = queue[1:] + visited++ + + for _, dependentID := range graph.nodes[current].dependents { + inDegree[dependentID]-- + if inDegree[dependentID] == 0 { + queue = append(queue, dependentID) + } + } + } + if visited == len(graph.order) { + return nil + } + + cycleNodes := make([]string, 0) + for _, id := range graph.order { + if inDegree[id] > 0 { + cycleNodes = append(cycleNodes, id) + } + } + slices.Sort(cycleNodes) + return fmt.Errorf("%w: %s", agentsession.ErrCyclicDependency, strings.Join(cycleNodes, ", ")) +} diff --git a/internal/subagent/scheduler_test.go b/internal/subagent/scheduler_test.go new file mode 100644 index 00000000..82ad411a --- /dev/null +++ b/internal/subagent/scheduler_test.go @@ -0,0 +1,987 @@ +package subagent + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + agentsession "neo-code/internal/session" +) + +type schedulerStore struct { + mu sync.Mutex + session *agentsession.Session + claimConflicts map[string]int + updateConflicts map[string]int +} + +type schedulerStoreWithClaimError struct { + *schedulerStore + claimErrors map[string]error +} + +func newSchedulerStore(t *testing.T, items []agentsession.TodoItem) *schedulerStore { + t.Helper() + session := agentsession.New("scheduler") + if err := session.ReplaceTodos(items); err != nil { + t.Fatalf("ReplaceTodos() error = %v", err) + } + return &schedulerStore{ + session: &session, + claimConflicts: make(map[string]int), + updateConflicts: make(map[string]int), + } +} + +func (s *schedulerStore) ListTodos() []agentsession.TodoItem { + s.mu.Lock() + defer s.mu.Unlock() + return s.session.ListTodos() +} + +func (s *schedulerStore) FindTodo(id string) (agentsession.TodoItem, bool) { + s.mu.Lock() + defer s.mu.Unlock() + return s.session.FindTodo(id) +} + +func (s *schedulerStore) UpdateTodo(id string, patch agentsession.TodoPatch, expectedRevision int64) error { + s.mu.Lock() + defer s.mu.Unlock() + if n := s.updateConflicts[id]; n > 0 { + s.updateConflicts[id] = n - 1 + return fmt.Errorf("%w: injected update conflict", agentsession.ErrRevisionConflict) + } + return s.session.UpdateTodo(id, patch, expectedRevision) +} + +func (s *schedulerStore) ClaimTodo(id string, ownerType string, ownerID string, expectedRevision int64) error { + s.mu.Lock() + defer s.mu.Unlock() + if n := s.claimConflicts[id]; n > 0 { + s.claimConflicts[id] = n - 1 + return fmt.Errorf("%w: injected claim conflict", agentsession.ErrRevisionConflict) + } + return s.session.ClaimTodo(id, ownerType, ownerID, expectedRevision) +} + +func (s *schedulerStore) CompleteTodo(id string, artifacts []string, expectedRevision int64) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.session.CompleteTodo(id, artifacts, expectedRevision) +} + +func (s *schedulerStore) FailTodo(id string, reason string, expectedRevision int64) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.session.FailTodo(id, reason, expectedRevision) +} + +func (s *schedulerStoreWithClaimError) ClaimTodo(id string, ownerType string, ownerID string, expectedRevision int64) error { + s.mu.Lock() + if err, ok := s.claimErrors[id]; ok && err != nil { + delete(s.claimErrors, id) + s.mu.Unlock() + return err + } + s.mu.Unlock() + return s.schedulerStore.ClaimTodo(id, ownerType, ownerID, expectedRevision) +} + +type scriptedFactory struct { + mu sync.Mutex + attempts map[string]int + run func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) +} + +func newScriptedFactory( + run func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error), +) *scriptedFactory { + return &scriptedFactory{ + attempts: make(map[string]int), + run: run, + } +} + +func (f *scriptedFactory) Create(role Role) (WorkerRuntime, error) { + policy, err := DefaultRolePolicy(role) + if err != nil { + return nil, err + } + engine := EngineFunc(func(ctx context.Context, input StepInput) (StepOutput, error) { + f.mu.Lock() + f.attempts[input.Task.ID]++ + attempt := f.attempts[input.Task.ID] + f.mu.Unlock() + return f.run(ctx, input.Task.ID, attempt, input) + }) + return NewWorker(role, policy, engine) +} + +func successStep(taskID string) StepOutput { + return StepOutput{ + Done: true, + Output: Output{ + Summary: "done: " + taskID, + Findings: []string{"ok"}, + Patches: []string{"none"}, + Risks: []string{"low"}, + NextActions: []string{"continue"}, + Artifacts: []string{taskID + ".artifact"}, + }, + } +} + +func TestSchedulerConfigNormalize(t *testing.T) { + t.Parallel() + + cfg := (SchedulerConfig{}).normalize() + if cfg.MaxConcurrency != 2 { + t.Fatalf("MaxConcurrency = %d, want 2", cfg.MaxConcurrency) + } + if cfg.DefaultRole != RoleCoder { + t.Fatalf("DefaultRole = %q, want %q", cfg.DefaultRole, RoleCoder) + } + if cfg.DefaultBudget.MaxSteps <= 0 || cfg.DefaultBudget.Timeout <= 0 { + t.Fatalf("DefaultBudget not normalized: %+v", cfg.DefaultBudget) + } + if cfg.MaxRetries != 0 { + t.Fatalf("MaxRetries = %d, want 0", cfg.MaxRetries) + } + if cfg.PollInterval <= 0 { + t.Fatalf("PollInterval should be positive") + } + if cfg.WorkerIDPrefix == "" { + t.Fatalf("WorkerIDPrefix should not be empty") + } + if cfg.Clock == nil || cfg.RoleSelector == nil || cfg.BudgetSelector == nil || cfg.Backoff == nil || cfg.Observer == nil { + t.Fatalf("normalize() should set defaults") + } + + role := cfg.RoleSelector(agentsession.TodoItem{OwnerID: string(RoleReviewer)}) + if role != cfg.DefaultRole { + t.Fatalf("RoleSelector() = %q, want default %q", role, cfg.DefaultRole) + } + role = cfg.RoleSelector(agentsession.TodoItem{OwnerID: "unknown"}) + if role != cfg.DefaultRole { + t.Fatalf("RoleSelector fallback = %q, want %q", role, cfg.DefaultRole) + } + if got := cfg.Backoff(3); got != 3*time.Second { + t.Fatalf("Backoff(3) = %v, want 3s", got) + } + if got := cfg.BudgetSelector(agentsession.TodoItem{}, cfg.DefaultBudget); got != cfg.DefaultBudget { + t.Fatalf("BudgetSelector() = %+v, want %+v", got, cfg.DefaultBudget) + } +} + +func TestNewSchedulerValidationErrors(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{{ID: "a", Content: "a"}}) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = taskID + _ = attempt + _ = input + return successStep("a"), nil + }) + + if _, err := NewScheduler(nil, factory, SchedulerConfig{}); err == nil { + t.Fatalf("NewScheduler(nil, factory) should fail") + } + if _, err := NewScheduler(store, nil, SchedulerConfig{}); err == nil { + t.Fatalf("NewScheduler(store, nil) should fail") + } +} + +func TestSchedulerRunEarlyErrorPaths(t *testing.T) { + t.Parallel() + + t.Run("context canceled before run", func(t *testing.T) { + t.Parallel() + store := newSchedulerStore(t, []agentsession.TodoItem{{ID: "a", Content: "a"}}) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = taskID + _ = attempt + _ = input + return successStep("a"), nil + }) + scheduler, err := NewScheduler(store, factory, SchedulerConfig{}) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := scheduler.Run(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("Run() error = %v, want context canceled", err) + } + }) + + t.Run("invalid graph", func(t *testing.T) { + t.Parallel() + store := newSchedulerStore(t, []agentsession.TodoItem{{ID: "a", Content: "a"}}) + store.session.Todos = append(store.session.Todos, agentsession.TodoItem{ID: "", Content: "bad"}) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = taskID + _ = attempt + _ = input + return successStep("a"), nil + }) + scheduler, err := NewScheduler(store, factory, SchedulerConfig{}) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + _, err = scheduler.Run(context.Background()) + if err == nil || !strings.Contains(err.Error(), "empty todo id") { + t.Fatalf("Run() error = %v, want empty todo id", err) + } + }) +} + +func TestBuildTaskGraphValidation(t *testing.T) { + t.Parallel() + + t.Run("duplicate id", func(t *testing.T) { + t.Parallel() + _, err := buildTaskGraph([]agentsession.TodoItem{ + {ID: "a", Content: "a"}, + {ID: "a", Content: "dup"}, + }) + if err == nil || !strings.Contains(err.Error(), "duplicate") { + t.Fatalf("buildTaskGraph() error = %v, want duplicate", err) + } + }) + + t.Run("unknown dependency", func(t *testing.T) { + t.Parallel() + _, err := buildTaskGraph([]agentsession.TodoItem{ + {ID: "a", Content: "a", Dependencies: []string{"x"}}, + }) + if err == nil || !strings.Contains(err.Error(), "unknown dependency") { + t.Fatalf("buildTaskGraph() error = %v, want unknown dependency", err) + } + }) + + t.Run("cycle", func(t *testing.T) { + t.Parallel() + _, err := buildTaskGraph([]agentsession.TodoItem{ + {ID: "a", Content: "a", Dependencies: []string{"b"}}, + {ID: "b", Content: "b", Dependencies: []string{"a"}}, + }) + if err == nil || !errors.Is(err, agentsession.ErrCyclicDependency) { + t.Fatalf("buildTaskGraph() error = %v, want cyclic dependency", err) + } + }) +} + +func TestSchedulerRunDependencyUnlock(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "a", Content: "task-a", Priority: 2}, + {ID: "b", Content: "task-b", Dependencies: []string{"a"}, Priority: 1}, + }) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = attempt + _ = input + return successStep(taskID), nil + }) + + var ( + mu sync.Mutex + events []SchedulerEvent + ) + cfg := SchedulerConfig{ + MaxConcurrency: 2, + PollInterval: 2 * time.Millisecond, + Observer: func(event SchedulerEvent) { + mu.Lock() + events = append(events, event) + mu.Unlock() + }, + } + scheduler, err := NewScheduler(store, factory, cfg) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + result, err := scheduler.Run(ctx) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + + if len(result.Succeeded) != 2 { + t.Fatalf("Succeeded = %v, want 2 tasks", result.Succeeded) + } + if len(result.Failed) != 0 { + t.Fatalf("Failed = %v, want empty", result.Failed) + } + if len(result.BlockedLeft) != 0 { + t.Fatalf("BlockedLeft = %v, want empty", result.BlockedLeft) + } + + a, ok := store.FindTodo("a") + if !ok || a.Status != agentsession.TodoStatusCompleted { + t.Fatalf("todo a status = %+v, want completed", a) + } + b, ok := store.FindTodo("b") + if !ok || b.Status != agentsession.TodoStatusCompleted { + t.Fatalf("todo b status = %+v, want completed", b) + } + if len(b.Artifacts) == 0 || b.Artifacts[0] != "b.artifact" { + t.Fatalf("todo b artifacts = %v, want b.artifact", b.Artifacts) + } + + mu.Lock() + defer mu.Unlock() + if !hasEvent(events, SchedulerEventBlocked, "b") { + t.Fatalf("expected blocked event for b, events=%+v", events) + } + if !hasEvent(events, SchedulerEventCompleted, "a") || !hasEvent(events, SchedulerEventCompleted, "b") { + t.Fatalf("expected completed events for a/b, events=%+v", events) + } +} + +func TestSchedulerRunConcurrencyLimit(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "t1", Content: "t1"}, + {ID: "t2", Content: "t2"}, + {ID: "t3", Content: "t3"}, + }) + + var running int32 + var maxRunning int32 + started := make(chan string, 8) + release := make(chan struct{}) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = attempt + _ = input + nowRunning := atomic.AddInt32(&running, 1) + for { + prev := atomic.LoadInt32(&maxRunning) + if nowRunning <= prev || atomic.CompareAndSwapInt32(&maxRunning, prev, nowRunning) { + break + } + } + started <- taskID + select { + case <-release: + case <-ctx.Done(): + atomic.AddInt32(&running, -1) + return StepOutput{}, ctx.Err() + } + atomic.AddInt32(&running, -1) + return successStep(taskID), nil + }) + + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 2, + PollInterval: 2 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + done := make(chan error, 1) + go func() { + _, runErr := scheduler.Run(ctx) + done <- runErr + }() + + timeout := time.After(500 * time.Millisecond) + count := 0 + for count < 2 { + select { + case <-started: + count++ + case <-timeout: + t.Fatalf("timeout waiting first two tasks to start") + } + } + close(release) + + select { + case runErr := <-done: + if runErr != nil { + t.Fatalf("Run() error = %v", runErr) + } + case <-time.After(2 * time.Second): + t.Fatalf("timeout waiting scheduler complete") + } + + if got := atomic.LoadInt32(&maxRunning); got != 2 { + t.Fatalf("max concurrency = %d, want 2", got) + } +} + +func TestSchedulerRunRetryAndGiveUp(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "r1", Content: "retry once then pass", RetryLimit: 1}, + {ID: "r2", Content: "always fail", RetryLimit: 1}, + }) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = input + switch taskID { + case "r1": + if attempt == 1 { + return StepOutput{}, errors.New("r1 first fail") + } + return successStep(taskID), nil + case "r2": + return StepOutput{}, errors.New("r2 always fail") + default: + return successStep(taskID), nil + } + }) + + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 2, + MaxRetries: 2, + PollInterval: 2 * time.Millisecond, + Backoff: func(attempt int) time.Duration { + _ = attempt + return 0 + }, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + result, err := scheduler.Run(ctx) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + + if got := result.Retried["r1"]; got != 1 { + t.Fatalf("Retried[r1] = %d, want 1", got) + } + if got := result.Retried["r2"]; got != 1 { + t.Fatalf("Retried[r2] = %d, want 1", got) + } + if !contains(result.Failed, "r2") { + t.Fatalf("Failed = %v, want r2", result.Failed) + } + + r1, _ := store.FindTodo("r1") + if r1.Status != agentsession.TodoStatusCompleted { + t.Fatalf("r1 status = %q, want completed", r1.Status) + } + r2, _ := store.FindTodo("r2") + if r2.Status != agentsession.TodoStatusFailed { + t.Fatalf("r2 status = %q, want failed", r2.Status) + } + if r2.RetryCount != 2 { + t.Fatalf("r2 retry_count = %d, want 2", r2.RetryCount) + } + if !strings.Contains(r2.FailureReason, "always fail") { + t.Fatalf("r2 failure_reason = %q, want contains always fail", r2.FailureReason) + } +} + +func TestSchedulerRunRevisionConflict(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "a", Content: "task-a"}, + }) + store.claimConflicts["a"] = 1 + store.updateConflicts["a"] = 1 + + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = input + _ = attempt + return successStep(taskID), nil + }) + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 1, + PollInterval: 2 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + result, err := scheduler.Run(ctx) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !contains(result.Succeeded, "a") { + t.Fatalf("Succeeded = %v, want a", result.Succeeded) + } +} + +func TestSchedulerRunStopsOnDependencyDeadEnd(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "root", Content: "root", Status: agentsession.TodoStatusFailed}, + {ID: "child", Content: "child", Dependencies: []string{"root"}, Status: agentsession.TodoStatusPending}, + }) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = taskID + _ = attempt + _ = input + return successStep(taskID), nil + }) + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 1, + PollInterval: 2 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result, err := scheduler.Run(ctx) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !contains(result.BlockedLeft, "child") { + t.Fatalf("BlockedLeft = %v, want child", result.BlockedLeft) + } +} + +func TestSchedulerRunCancellationWriteback(t *testing.T) { + t.Parallel() + + store := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "t1", Content: "task-1"}, + }) + started := make(chan struct{}, 1) + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = taskID + _ = attempt + _ = input + select { + case started <- struct{}{}: + default: + } + <-ctx.Done() + return StepOutput{}, ctx.Err() + }) + + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 1, + PollInterval: 2 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, runErr := scheduler.Run(ctx) + done <- runErr + }() + + select { + case <-started: + case <-time.After(time.Second): + t.Fatalf("timeout waiting task start") + } + cancel() + + select { + case runErr := <-done: + if !errors.Is(runErr, context.Canceled) { + t.Fatalf("Run() error = %v, want context canceled", runErr) + } + case <-time.After(time.Second): + t.Fatalf("timeout waiting scheduler stop") + } + + item, _ := store.FindTodo("t1") + if item.Status != agentsession.TodoStatusCanceled { + t.Fatalf("todo status = %q, want canceled", item.Status) + } +} + +func TestSchedulerRunClaimErrorCancelsRunningTodo(t *testing.T) { + t.Parallel() + + baseStore := newSchedulerStore(t, []agentsession.TodoItem{ + {ID: "a", Content: "task-a"}, + {ID: "b", Content: "task-b"}, + }) + store := &schedulerStoreWithClaimError{ + schedulerStore: baseStore, + claimErrors: map[string]error{ + "b": errors.New("injected claim failure"), + }, + } + factory := newScriptedFactory(func(ctx context.Context, taskID string, attempt int, input StepInput) (StepOutput, error) { + _ = ctx + _ = taskID + _ = attempt + _ = input + return successStep(taskID), nil + }) + scheduler, err := NewScheduler(store, factory, SchedulerConfig{ + MaxConcurrency: 2, + PollInterval: 2 * time.Millisecond, + }) + if err != nil { + t.Fatalf("NewScheduler() error = %v", err) + } + + _, err = scheduler.Run(context.Background()) + if err == nil || !strings.Contains(err.Error(), "claim todo") { + t.Fatalf("Run() error = %v, want claim todo error", err) + } + + item, ok := store.FindTodo("a") + if !ok { + t.Fatalf("FindTodo(a) not found") + } + if item.Status != agentsession.TodoStatusCanceled { + t.Fatalf("todo a status = %q, want canceled", item.Status) + } +} + +type fakeFactory struct { + create func(role Role) (WorkerRuntime, error) +} + +func (f fakeFactory) Create(role Role) (WorkerRuntime, error) { + return f.create(role) +} + +type fakeWorker struct { + startErr error + stepResult StepResult + stepErr error + stepFunc func(ctx context.Context) (StepResult, error) + result Result + resultErr error + state State +} + +func (w *fakeWorker) Start(task Task, budget Budget, capability Capability) error { + _ = task + _ = budget + _ = capability + if w.startErr != nil { + return w.startErr + } + w.state = StateRunning + return nil +} + +func (w *fakeWorker) Step(ctx context.Context) (StepResult, error) { + if w.stepFunc != nil { + return w.stepFunc(ctx) + } + _ = ctx + return w.stepResult, w.stepErr +} + +func (w *fakeWorker) Stop(reason StopReason) error { + if reason == StopReasonCanceled { + w.state = StateCanceled + } + return nil +} + +func (w *fakeWorker) Result() (Result, error) { + return w.result, w.resultErr +} + +func (w *fakeWorker) State() State { + return w.state +} + +func (w *fakeWorker) Policy() RolePolicy { + return RolePolicy{} +} + +func TestExecuteTaskWithFactoryBranches(t *testing.T) { + t.Parallel() + + t.Run("factory create failed", func(t *testing.T) { + t.Parallel() + _, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return nil, errors.New("create failed") + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "create failed") { + t.Fatalf("executeTaskWithFactory() error = %v, want create failed", err) + } + }) + + t.Run("start failed", func(t *testing.T) { + t.Parallel() + _, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{startErr: errors.New("start failed")}, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "start failed") { + t.Fatalf("executeTaskWithFactory() error = %v, want start failed", err) + } + }) + + t.Run("step canceled with result fallback", func(t *testing.T) { + t.Parallel() + result, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepErr: context.Canceled, + resultErr: errors.New("no result"), + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("error = %v, want context canceled", err) + } + if result.State != StateCanceled || result.StopReason != StopReasonCanceled { + t.Fatalf("result = %+v, want canceled fallback", result) + } + }) + + t.Run("step failed with result", func(t *testing.T) { + t.Parallel() + result, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepErr: errors.New("boom"), + result: Result{ + State: StateFailed, + Error: "boom", + }, + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("error = %v, want boom", err) + } + if result.State != StateFailed { + t.Fatalf("result state = %q, want failed", result.State) + } + }) + + t.Run("done with failed state and no error message", func(t *testing.T) { + t.Parallel() + _, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepResult: StepResult{Done: true}, + result: Result{ + State: StateFailed, + StopReason: StopReasonError, + }, + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "state=failed") { + t.Fatalf("error = %v, want fallback state error", err) + } + }) + + t.Run("step error fallback without result", func(t *testing.T) { + t.Parallel() + result, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepErr: errors.New("step failed"), + resultErr: errors.New("result unavailable"), + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "step failed") { + t.Fatalf("error = %v, want step failed", err) + } + if result.State != StateFailed || result.StopReason != StopReasonError { + t.Fatalf("result = %+v, want failed fallback", result) + } + }) + + t.Run("step retry then success", func(t *testing.T) { + t.Parallel() + result, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + calls := 0 + return &fakeWorker{ + stepFunc: func(ctx context.Context) (StepResult, error) { + _ = ctx + calls++ + if calls == 1 { + return StepResult{Done: false}, nil + } + return StepResult{Done: true}, nil + }, + result: Result{ + State: StateSucceeded, + StopReason: StopReasonCompleted, + }, + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err != nil { + t.Fatalf("error = %v, want nil", err) + } + if result.State != StateSucceeded { + t.Fatalf("result state = %q, want succeeded", result.State) + } + }) + + t.Run("done but result failed with explicit message", func(t *testing.T) { + t.Parallel() + _, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepResult: StepResult{Done: true}, + result: Result{ + State: StateFailed, + Error: "explicit fail", + }, + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "explicit fail") { + t.Fatalf("error = %v, want explicit fail", err) + } + }) + + t.Run("done but result unavailable", func(t *testing.T) { + t.Parallel() + _, err := executeTaskWithFactory(context.Background(), fakeFactory{ + create: func(role Role) (WorkerRuntime, error) { + _ = role + return &fakeWorker{ + stepResult: StepResult{Done: true}, + resultErr: errors.New("result unavailable"), + }, nil + }, + }, scheduleTaskInput{Role: RoleCoder, Task: Task{ID: "t", Goal: "g"}}) + if err == nil || !strings.Contains(err.Error(), "result unavailable") { + t.Fatalf("error = %v, want result unavailable", err) + } + }) +} + +func TestSchedulerHelpersCoverage(t *testing.T) { + t.Parallel() + + items := []agentsession.TodoItem{ + {ID: "a", Content: "a", Status: agentsession.TodoStatusCompleted}, + {ID: "b", Content: "b", Dependencies: []string{"a"}, Status: agentsession.TodoStatusPending}, + } + byID := mapTodosByID(items) + if !dependenciesCompleted(byID["b"], byID) { + t.Fatalf("dependenciesCompleted should be true") + } + + waited := effectivePriority(agentsession.TodoItem{Priority: 1}, time.Now().Add(-20*time.Second), time.Now(), 5*time.Second) + if waited <= 1 { + t.Fatalf("effectivePriority should include aging boost, got %d", waited) + } + + list := []string{"a"} + appendUniqueString(&list, "a") + appendUniqueString(&list, "b") + if len(list) != 2 || list[1] != "b" { + t.Fatalf("appendUniqueString unexpected result: %v", list) + } + + if !isRevisionConflict(fmt.Errorf("%w: x", agentsession.ErrRevisionConflict)) { + t.Fatalf("isRevisionConflict should match wrapped error") + } + if isRevisionConflict(errors.New("other")) { + t.Fatalf("isRevisionConflict should reject unrelated errors") + } + + if err := waitWithContext(context.Background(), 0); err != nil { + t.Fatalf("waitWithContext(0) error = %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := waitWithContext(ctx, time.Second); !errors.Is(err, context.Canceled) { + t.Fatalf("waitWithContext canceled error = %v", err) + } + + left := collectBlockedLeft([]string{"a", "b", "c"}, []agentsession.TodoItem{ + {ID: "a", Content: "a", Status: agentsession.TodoStatusCompleted}, + {ID: "b", Content: "b", Status: agentsession.TodoStatusBlocked}, + }, map[string]runningTask{ + "c": {id: "c"}, + }) + if len(left) != 2 || left[0] != "b" || left[1] != "c" { + t.Fatalf("collectBlockedLeft() = %v, want [b c]", left) + } + + outcome := taskOutcome{err: errors.New(" boom ")} + if got := resolveTaskFailureReason(outcome); got != "boom" { + t.Fatalf("resolveTaskFailureReason(err) = %q, want boom", got) + } + outcome = taskOutcome{result: Result{Error: " failed "}} + if got := resolveTaskFailureReason(outcome); got != "failed" { + t.Fatalf("resolveTaskFailureReason(result) = %q, want failed", got) + } + if got := resolveTaskFailureReason(taskOutcome{}); got == "" { + t.Fatalf("resolveTaskFailureReason() should return fallback") + } + + if got := defaultRetryBackoff(0); got != 0 { + t.Fatalf("defaultRetryBackoff(0) = %v, want 0", got) + } + cfg := (SchedulerConfig{MaxRetries: -1}).normalize() + if cfg.MaxRetries != 0 { + t.Fatalf("normalize MaxRetries = %d, want 0", cfg.MaxRetries) + } + if _, err := buildTaskGraph([]agentsession.TodoItem{{ID: " ", Content: "bad"}}); err == nil { + t.Fatalf("buildTaskGraph should reject empty id") + } +} + +func hasEvent(events []SchedulerEvent, eventType SchedulerEventType, taskID string) bool { + for _, event := range events { + if event.Type == eventType && event.TaskID == taskID { + return true + } + } + return false +} + +func contains(items []string, target string) bool { + for _, item := range items { + if item == target { + return true + } + } + return false +} diff --git a/internal/subagent/scheduler_types.go b/internal/subagent/scheduler_types.go new file mode 100644 index 00000000..7f229784 --- /dev/null +++ b/internal/subagent/scheduler_types.go @@ -0,0 +1,159 @@ +package subagent + +import ( + "time" + + agentsession "neo-code/internal/session" +) + +// TodoStore 定义调度器读写 Todo 的最小依赖，便于复用 runtime/session 不同实现。 +type TodoStore interface { + ListTodos() []agentsession.TodoItem + FindTodo(id string) (agentsession.TodoItem, bool) + UpdateTodo(id string, patch agentsession.TodoPatch, expectedRevision int64) error + ClaimTodo(id string, ownerType string, ownerID string, expectedRevision int64) error + CompleteTodo(id string, artifacts []string, expectedRevision int64) error + FailTodo(id string, reason string, expectedRevision int64) error +} + +// RoleSelector 根据 Todo 节点选择执行角色。 +type RoleSelector func(todo agentsession.TodoItem) Role + +// BudgetSelector 根据 Todo 节点生成本任务预算。 +type BudgetSelector func(todo agentsession.TodoItem, defaults Budget) Budget + +// CapabilitySelector 根据 Todo 节点生成本任务能力边界。 +type CapabilitySelector func(todo agentsession.TodoItem) Capability + +// RetryBackoff 计算第 N 次重试应等待的退避时长。 +type RetryBackoff func(attempt int) time.Duration + +// SchedulerEventType 描述调度器可观测事件类型。 +type SchedulerEventType string + +const ( + // SchedulerEventQueued 表示任务进入可调度队列。 + SchedulerEventQueued SchedulerEventType = "queued" + // SchedulerEventClaimed 表示任务已被领取并进入执行中。 + SchedulerEventClaimed SchedulerEventType = "claimed" + // SchedulerEventRunning 表示任务开始执行。 + SchedulerEventRunning SchedulerEventType = "running" + // SchedulerEventCompleted 表示任务执行成功并回写完成。 + SchedulerEventCompleted SchedulerEventType = "completed" + // SchedulerEventRetryScheduled 表示任务失败后已安排重试。 + SchedulerEventRetryScheduled SchedulerEventType = "retry_scheduled" + // SchedulerEventFailed 表示任务执行失败且不再重试。 + SchedulerEventFailed SchedulerEventType = "failed" + // SchedulerEventCanceled 表示任务因外部取消被回写为 canceled。 + SchedulerEventCanceled SchedulerEventType = "canceled" + // SchedulerEventBlocked 表示任务因依赖未满足处于阻塞态。 + SchedulerEventBlocked SchedulerEventType = "blocked" + // SchedulerEventFinished 表示调度器本轮结束。 + SchedulerEventFinished SchedulerEventType = "finished" +) + +// SchedulerEvent 描述单次调度决策事件，供日志或 runtime 事件桥接使用。 +type SchedulerEvent struct { + Type SchedulerEventType + TaskID string + Attempt int + Running int + QueueSize int + Reason string + At time.Time +} + +// SchedulerObserver 用于消费调度器事件。 +type SchedulerObserver func(event SchedulerEvent) + +// ScheduleResult 汇总一次调度执行的结果统计。 +type ScheduleResult struct { + StartedAt time.Time + EndedAt time.Time + Total int + Succeeded []string + Failed []string + Retried map[string]int + BlockedLeft []string +} + +// SchedulerConfig 描述调度器策略参数。 +type SchedulerConfig struct { + MaxConcurrency int + DefaultRole Role + DefaultBudget Budget + MaxRetries int + WorkerIDPrefix string + PollInterval time.Duration + Clock func() time.Time + RoleSelector RoleSelector + BudgetSelector BudgetSelector + Backoff RetryBackoff + Capabilities CapabilitySelector + Observer SchedulerObserver +} + +// normalize 返回带默认值的配置副本，保证调度器执行阶段不出现隐式零值。 +func (c SchedulerConfig) normalize() SchedulerConfig { + out := c + if out.MaxConcurrency <= 0 { + out.MaxConcurrency = 2 + } + if !out.DefaultRole.Valid() { + out.DefaultRole = RoleCoder + } + out.DefaultBudget = out.DefaultBudget.normalize(Budget{MaxSteps: 6, Timeout: 30 * time.Second}) + if out.MaxRetries < 0 { + out.MaxRetries = 0 + } + if out.PollInterval <= 0 { + out.PollInterval = 200 * time.Millisecond + } + if out.WorkerIDPrefix == "" { + out.WorkerIDPrefix = "subagent" + } + if out.Clock == nil { + out.Clock = time.Now + } + if out.RoleSelector == nil { + out.RoleSelector = defaultRoleSelector(out.DefaultRole) + } + if out.BudgetSelector == nil { + out.BudgetSelector = defaultBudgetSelector + } + if out.Backoff == nil { + out.Backoff = defaultRetryBackoff + } + if out.Capabilities == nil { + out.Capabilities = func(todo agentsession.TodoItem) Capability { + _ = todo + return Capability{} + } + } + if out.Observer == nil { + out.Observer = func(SchedulerEvent) {} + } + return out +} + +// defaultRoleSelector 生成默认角色选择器：仅使用调度器可信默认角色，避免读取任务可变元数据。 +func defaultRoleSelector(defaultRole Role) RoleSelector { + return func(todo agentsession.TodoItem) Role { + _ = todo + return defaultRole + } +} + +// defaultBudgetSelector 对每个任务统一返回标准预算，可被配置覆盖。 +func defaultBudgetSelector(todo agentsession.TodoItem, defaults Budget) Budget { + _ = todo + return defaults +} + +// defaultRetryBackoff 提供线性重试退避，避免瞬时失败造成热循环。 +func defaultRetryBackoff(attempt int) time.Duration { + if attempt <= 0 { + return 0 + } + return time.Duration(attempt) * time.Second +} diff --git a/internal/tui/infra/clipboard_common.go b/internal/tui/infra/clipboard_common.go index 3079eb72..42bec744 100644 --- a/internal/tui/infra/clipboard_common.go +++ b/internal/tui/infra/clipboard_common.go @@ -1,15 +1,21 @@ package infra import ( + "errors" "os" + "strings" ) +var errClipboardImageUnsupported = errors.New("clipboard image is not supported on this platform") + func SaveImageToTempFile(data []byte, prefix string) (string, error) { pattern := "image-*.png" if cleaned := sanitizeTempPrefix(prefix); cleaned != "" { pattern = cleaned + "-*.png" } - f, err := os.CreateTemp("", pattern) + + tempDir := strings.TrimSpace(os.Getenv("TMPDIR")) + f, err := os.CreateTemp(tempDir, pattern) if err != nil { return "", err } diff --git a/internal/tui/infra/clipboard_fallback.go b/internal/tui/infra/clipboard_fallback.go index ab19067f..1ff456f5 100644 --- a/internal/tui/infra/clipboard_fallback.go +++ b/internal/tui/infra/clipboard_fallback.go @@ -3,13 +3,9 @@ package infra import ( - "errors" - clipboardtext "github.com/atotto/clipboard" ) -var errClipboardImageUnsupported = errors.New("clipboard image is not supported on this platform") - func CopyText(text string) error { return clipboardtext.WriteAll(text) } diff --git a/internal/tui/infra/clipboard_native.go b/internal/tui/infra/clipboard_native.go index 604098b2..3f8f15e0 100644 --- a/internal/tui/infra/clipboard_native.go +++ b/internal/tui/infra/clipboard_native.go @@ -40,7 +40,7 @@ func ReadClipboardImage() ([]byte, error) { } data := clipboard.Read(clipboard.FmtImage) if len(data) == 0 { - return nil, nil + return nil, errClipboardImageUnsupported } return data, nil } diff --git a/internal/tui/infra/image.go b/internal/tui/infra/image.go index f7eb6504..3ade0574 100644 --- a/internal/tui/infra/image.go +++ b/internal/tui/infra/image.go @@ -34,7 +34,7 @@ func DetectImageMimeType(path string) string { } detected := mime.TypeByExtension(ext) - if detected != "" { + if strings.HasPrefix(detected, "image/") { return detected } @@ -50,17 +50,15 @@ func DetectImageMimeType(path string) string { if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF { return "image/jpeg" } - if len(data) >= 12 { - if string(data[0:4]) == "GIF8" { - return "image/gif" - } - if string(data[0:4]) == "RIFF" && string(data[8:12]) == "WEBP" { - return "image/webp" - } - } + } + if len(data) >= 4 && string(data[0:4]) == "GIF8" { + return "image/gif" + } + if len(data) >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WEBP" { + return "image/webp" } - return "" + return detected } func IsSupportedImageFormat(path string) bool { diff --git a/internal/tui/infra/infra_test.go b/internal/tui/infra/infra_test.go index 9e39bb45..4dfc1b9b 100644 --- a/internal/tui/infra/infra_test.go +++ b/internal/tui/infra/infra_test.go @@ -403,11 +403,11 @@ func TestClipboardFallbackFunctions(t *testing.T) { t.Fatalf("expected clipboard text or an error, got empty success result") } data, err := ReadClipboardImage() - if err != errClipboardImageUnsupported { - t.Fatalf("expected unsupported image error, got %v", err) + if err != nil && err != errClipboardImageUnsupported { + t.Fatalf("expected nil or unsupported image error, got %v", err) } - if data != nil { - t.Fatalf("expected nil image data on unsupported platform") + if err == nil && len(data) == 0 && data != nil { + t.Fatalf("expected nil for empty clipboard image state, got zero-length slice") } }