diff --git a/apps/backend/internal/mcp/handlers/config_handlers_test.go b/apps/backend/internal/mcp/handlers/config_handlers_test.go index 0035c32e43..29bd5aeef1 100644 --- a/apps/backend/internal/mcp/handlers/config_handlers_test.go +++ b/apps/backend/internal/mcp/handlers/config_handlers_test.go @@ -597,6 +597,74 @@ func TestHandleArchiveTask_MergedPRRunAcceptsBoundTarget(t *testing.T) { assert.NotNil(t, archived.ArchivedAt) } +// TestHandleArchiveTask_MergedPRRunAcceptsRefreshedTargetAfterResume is the +// regression guard for the reuse_thread + github_pr_merged defect: once a +// resumed continuation task's metadata is refreshed for a NEW firing (as +// orchestrator.refreshAutomationContinuationMetadata now does), the guard +// must bind to the refreshed target, not the stale one from the first +// firing. It exercises the metadata mutation through repo.SetTaskMetadataKey +// — the exact concurrent-key-safe primitive, on the exact same *sqliterepo. +// Repository type, that orchestrator.refreshAutomationContinuationMetadata +// calls in production (rather than the unrelated Service.UpdateTaskMetadata +// full-row path) — so this proves the guard accepts a binding written the +// same way production writes it, not just whatever happens to be in the DB. +func TestHandleArchiveTask_MergedPRRunAcceptsRefreshedTargetAfterResume(t *testing.T) { + svc, repo := newTestTaskService(t) + ctx := context.Background() + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ID: "ws-resume", Name: "Resume"})) + require.NoError(t, repo.CreateWorkflow(ctx, &models.Workflow{ID: "wf-resume", WorkspaceID: "ws-resume", Name: "Board"})) + + firstTarget := &models.Task{ + ID: "first-target", WorkspaceID: "ws-resume", WorkflowID: "wf-resume", + Title: "First merge target", State: v1.TaskStateTODO, + } + require.NoError(t, repo.CreateTask(ctx, firstTarget)) + secondTarget := &models.Task{ + ID: "second-target", WorkspaceID: "ws-resume", WorkflowID: "wf-resume", + Title: "Second merge target", State: v1.TaskStateTODO, + } + require.NoError(t, repo.CreateTask(ctx, secondTarget)) + caller := &models.Task{ + ID: "automation-run", WorkspaceID: "ws-resume", WorkflowID: "wf-resume", + Title: "Automation run", State: v1.TaskStateTODO, + Origin: models.TaskOriginAutomationRun, + Metadata: map[string]interface{}{ + "trigger_type": "github_pr_merged", + models.MetaKeyAutomationTargetTaskID: firstTarget.ID, + }, + } + require.NoError(t, repo.CreateTask(ctx, caller)) + + h := &Handlers{taskSvc: svc, logger: testLogger(t).WithFields()} + + // A second firing resumes the same continuation task and refreshes its + // binding — this is what orchestrator.refreshAutomationContinuationMetadata + // does on the reuse path, one key at a time via the same primitive. + require.NoError(t, repo.SetTaskMetadataKey(ctx, caller.ID, "trigger_type", "github_pr_merged")) + require.NoError(t, repo.SetTaskMetadataKey(ctx, caller.ID, models.MetaKeyAutomationTargetTaskID, secondTarget.ID)) + + // The second merge's target is now accepted... + secondMsg := makeWSMessage(t, ws.ActionMCPArchiveTask, map[string]string{ + "task_id": secondTarget.ID, "caller_task_id": caller.ID, + }) + resp, err := h.handleArchiveTask(ctx, secondMsg) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, resp.Type) + archived, err := svc.GetTask(ctx, secondTarget.ID) + require.NoError(t, err) + assert.NotNil(t, archived.ArchivedAt, "the refreshed target must be archivable") + + // ...and the stale first target is correctly refused against the now-current + // binding while it is still unarchived, so this assertion specifically tests + // the target guard rather than ordinary archive validation. + staleMsg := makeWSMessage(t, ws.ActionMCPArchiveTask, map[string]string{ + "task_id": firstTarget.ID, "caller_task_id": caller.ID, + }) + resp, err = h.handleArchiveTask(ctx, staleMsg) + require.NoError(t, err) + assertWSError(t, resp, ws.ErrorCodeValidation) +} + func TestHandleArchiveTask_MergedPRRunRejectsMissingBinding(t *testing.T) { svc, repo := newTestTaskService(t) ctx := context.Background() diff --git a/apps/backend/internal/orchestrator/event_handlers_automation.go b/apps/backend/internal/orchestrator/event_handlers_automation.go index 001bc5ce76..6bf111b366 100644 --- a/apps/backend/internal/orchestrator/event_handlers_automation.go +++ b/apps/backend/internal/orchestrator/event_handlers_automation.go @@ -340,7 +340,7 @@ func (s *Service) createAutomationTask(ctx context.Context, evt *automation.Auto zap.String("trigger_type", string(evt.TriggerType))) if continuationSession != nil { - s.dispatchAutomationContinuation(ctx, a, task, continuationSession, prompt, evt.RunID, action, reason) + s.dispatchAutomationContinuation(ctx, a, task, continuationSession, prompt, metadata, evt.RunID, action, reason) return } @@ -424,6 +424,130 @@ func (s *Service) prepareAutomationTask( return task, nil, action, reason, nil } +// automationContinuationMetadataSnapshot records the values that a firing +// replaces. It lets a failed dispatch restore the previous authorization +// target instead of leaving a stopped firing's target on the shared task. +type automationContinuationMetadataSnapshot struct { + values map[string]interface{} + present map[string]bool +} + +// refreshAutomationContinuationMetadata is kept as the small direct helper +// used by tests and legacy callers. The run-aware path uses the snapshot form +// so it can restore the task when dispatch or exact binding fails. +func (s *Service) refreshAutomationContinuationMetadata(ctx context.Context, task *models.Task, metadata map[string]interface{}) error { + _, err := s.refreshAutomationContinuationMetadataWithSnapshot(ctx, task, metadata) + return err +} + +// refreshAutomationContinuationMetadataWithSnapshot merges this firing's +// freshly-computed metadata onto a resumed continuation task and persists it. +// The writes happen only after DispatchRun has admitted the firing and holds +// its per-automation lock. This prevents a stopped or deleted run from +// changing the shared task's archive authorization. +// +// Each key is patched independently through the concurrent-key-safe +// SetTaskMetadataKey primitive rather than a full-row UpdateTask. A full-row +// write built from an in-memory snapshot could silently resurrect a key that +// another lifecycle operation just removed. Deferred-launch ownership is +// server-managed and is skipped here. +func (s *Service) refreshAutomationContinuationMetadataWithSnapshot( + ctx context.Context, + task *models.Task, + metadata map[string]interface{}, +) (*automationContinuationMetadataSnapshot, error) { + keys := orderedAutomationContinuationMetadataKeys(metadata) + snapshot := &automationContinuationMetadataSnapshot{ + values: make(map[string]interface{}, len(keys)), + present: make(map[string]bool, len(keys)), + } + for _, key := range keys { + value, ok := task.Metadata[key] + snapshot.values[key] = value + snapshot.present[key] = ok + } + + written := make([]string, 0, len(keys)) + for _, key := range keys { + if err := s.repo.SetTaskMetadataKey(ctx, task.ID, key, metadata[key]); err != nil { + s.logger.Error("failed to refresh automation continuation task metadata", + zap.String("task_id", task.ID), zap.String("metadata_key", key), zap.Error(err)) + if restoreErr := s.restoreAutomationContinuationMetadata(ctx, task, snapshot, written); restoreErr != nil { + s.logger.Warn("failed to restore automation continuation task metadata after refresh failure", + zap.String("task_id", task.ID), zap.Error(restoreErr)) + } + return nil, fmt.Errorf("refresh automation continuation metadata key %q: %w", key, err) + } + written = append(written, key) + } + if len(written) == 0 { + return snapshot, nil + } + if err := s.reloadAutomationContinuationTask(ctx, task); err != nil { + if restoreErr := s.restoreAutomationContinuationMetadata(ctx, task, snapshot, written); restoreErr != nil { + s.logger.Warn("failed to restore automation continuation task metadata after reload failure", + zap.String("task_id", task.ID), zap.Error(restoreErr)) + } + return nil, fmt.Errorf("reload refreshed automation continuation task: %w", err) + } + s.publishTaskUpdated(ctx, task) + return snapshot, nil +} + +func orderedAutomationContinuationMetadataKeys(metadata map[string]interface{}) []string { + keys := make([]string, 0, len(metadata)) + for key := range metadata { + if key == models.MetaKeyDeferredLaunch || key == models.MetaKeyAutomationTargetTaskID { + continue + } + keys = append(keys, key) + } + sort.Strings(keys) + if _, ok := metadata[models.MetaKeyAutomationTargetTaskID]; ok { + keys = append([]string{models.MetaKeyAutomationTargetTaskID}, keys...) + } + return keys +} + +func (s *Service) restoreAutomationContinuationMetadata( + ctx context.Context, + task *models.Task, + snapshot *automationContinuationMetadataSnapshot, + keys []string, +) error { + if snapshot == nil || len(keys) == 0 { + return nil + } + for _, key := range keys { + var err error + if snapshot.present[key] { + err = s.repo.SetTaskMetadataKey(ctx, task.ID, key, snapshot.values[key]) + } else { + _, err = s.repo.RemoveTaskMetadataKey(ctx, task.ID, key) + } + if err != nil { + return fmt.Errorf("restore automation continuation metadata key %q: %w", key, err) + } + } + if err := s.reloadAutomationContinuationTask(ctx, task); err != nil { + return err + } + s.publishTaskUpdated(ctx, task) + return nil +} + +func (s *Service) reloadAutomationContinuationTask(ctx context.Context, task *models.Task) error { + refreshed, err := s.repo.GetTask(ctx, task.ID) + if err != nil { + return err + } + if refreshed == nil { + return fmt.Errorf("task %s was not found", task.ID) + } + *task = *refreshed + return nil +} + func (s *Service) findAutomationContinuation(ctx context.Context, a *automation.Automation, evt *automation.AutomationTriggeredEvent) (*models.Task, *models.TaskSession, string) { if a.ContinuationTaskID == "" { return nil, nil, "no continuation task exists; created the first task" @@ -563,6 +687,7 @@ func (s *Service) dispatchAutomationRun( action automation.ThreadAction, reason, operation string, dispatch func() (automation.RunDispatch, error), + onFailure func(), ) bool { if runID == "" { return false @@ -574,6 +699,9 @@ func (s *Service) dispatchAutomationRun( if err := dispatcher.DispatchRun(ctx, runID, action, reason, dispatch); err == nil { return true } else { + if onFailure != nil { + onFailure() + } s.logger.Error("failed to dispatch automation run", zap.String("operation", operation), zap.String("automation_id", automationID), zap.String("task_id", taskID), zap.String("session_id", sessionID), zap.Error(err)) @@ -618,14 +746,37 @@ func (s *Service) bindAutomationRun( return false } -func (s *Service) dispatchAutomationContinuation(ctx context.Context, a *automation.Automation, task *models.Task, session *models.TaskSession, prompt, runID string, action automation.ThreadAction, reason string) { - if s.dispatchAutomationRun(ctx, a.ID, task.ID, session.ID, runID, action, reason, "continuation", func() (automation.RunDispatch, error) { - return s.promptAutomationContinuation(ctx, task, session, prompt) - }) { +func (s *Service) dispatchAutomationContinuation(ctx context.Context, a *automation.Automation, task *models.Task, session *models.TaskSession, prompt string, metadata map[string]interface{}, runID string, action automation.ThreadAction, reason string) { + var snapshot *automationContinuationMetadataSnapshot + restore := func() { + if snapshot == nil { + return + } + keys := orderedAutomationContinuationMetadataKeys(metadata) + if err := s.restoreAutomationContinuationMetadata(ctx, task, snapshot, keys); err != nil { + s.logger.Warn("failed to restore automation continuation task metadata", + zap.String("task_id", task.ID), zap.Error(err)) + } + snapshot = nil + } + dispatch := func() (automation.RunDispatch, error) { + var err error + snapshot, err = s.refreshAutomationContinuationMetadataWithSnapshot(ctx, task, metadata) + if err != nil { + return automation.RunDispatch{}, err + } + result, err := s.promptAutomationContinuation(ctx, task, session, prompt) + if err != nil { + restore() + return automation.RunDispatch{}, err + } + return result, nil + } + if s.dispatchAutomationRun(ctx, a.ID, task.ID, session.ID, runID, action, reason, "continuation", dispatch, restore) { return } - dispatch, err := s.promptAutomationContinuation(ctx, task, session, prompt) + dispatchResult, err := dispatch() if err != nil { s.logger.Error("failed to dispatch automation continuation", zap.String("automation_id", a.ID), zap.String("task_id", task.ID), zap.String("session_id", session.ID), zap.Error(err)) @@ -634,7 +785,8 @@ func (s *Service) dispatchAutomationContinuation(ctx context.Context, a *automat } return } - if !s.bindAutomationRun(ctx, runID, dispatch.TaskID, dispatch.SessionID, dispatch.TurnID, action, reason, "continuation") { + if !s.bindAutomationRun(ctx, runID, dispatchResult.TaskID, dispatchResult.SessionID, dispatchResult.TurnID, action, reason, "continuation") { + restore() return } } @@ -693,7 +845,7 @@ func (s *Service) autoStartAutomationTask(ctx context.Context, a *automation.Aut func (s *Service) autoStartAutomationTaskForRun(ctx context.Context, a *automation.Automation, task *models.Task, workflowStepID, runID string, action automation.ThreadAction, reason string) { if s.dispatchAutomationRun(ctx, a.ID, task.ID, "", runID, action, reason, "auto-start", func() (automation.RunDispatch, error) { return s.startAutomationTask(ctx, a, task, workflowStepID) - }) { + }, nil) { return } diff --git a/apps/backend/internal/orchestrator/event_handlers_automation_test.go b/apps/backend/internal/orchestrator/event_handlers_automation_test.go index 45efbf2c5a..4b13d15781 100644 --- a/apps/backend/internal/orchestrator/event_handlers_automation_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_automation_test.go @@ -331,6 +331,24 @@ func (s *stubAutomationService) RunWorkspaceInUse(_ context.Context, taskID stri return s.inUse[taskID], nil } +// rejectingAutomationService models a run that is stopped after admission, +// but before the service-owned dispatcher can acquire its lock. The +// continuation task must not receive the firing's target in that case. +type rejectingAutomationService struct { + *stubAutomationService + dispatchErr error +} + +func (s *rejectingAutomationService) DispatchRun( + context.Context, + string, + automation.ThreadAction, + string, + func() (automation.RunDispatch, error), +) error { + return s.dispatchErr +} + // seedAutomationTask writes a task exactly as the automation path now creates // one: ordinary and persistent, tagged only by its origin. func seedAutomationTask(t *testing.T, repo *sqliterepo.Repository, taskID, origin string, ephemeral bool) { @@ -419,6 +437,184 @@ func TestCreateAutomationTask_BindsMergedPRTarget(t *testing.T) { require.Equal(t, "target-task-1", creator.got.Metadata[models.MetaKeyAutomationTargetTaskID]) } +// TestRefreshAutomationContinuationMetadataOnResume is the regression +// guard for the reuse_thread + github_pr_merged defect: a resumed task keeps +// whatever automation_target_task_id its FIRST firing stamped, so every merge +// after the first is refused by validateAutomationArchiveTarget forever. The +// per-firing metadata computed for the SECOND firing must be merged onto the +// resumed task after dispatch admission (and persisted), not discarded at the +// reuse early return. It also proves the write goes through the per-key +// SetTaskMetadataKey primitive rather than a full-row UpdateTask (the persisted +// deferred-launch value survives even though the firing's own metadata carries +// a conflicting value for that key — the skip branch is exercised, not +// vacuously true). +func TestRefreshAutomationContinuationMetadataOnResume(t *testing.T) { + repo := setupTestRepo(t) + now := time.Now().UTC() + ctx := context.Background() + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ + ID: "ws-1", Name: "Test", CreatedAt: now, UpdatedAt: now, + })) + require.NoError(t, repo.CreateTask(ctx, &models.Task{ + ID: "continuation-task", + WorkspaceID: "ws-1", + Origin: models.TaskOriginAutomationRun, + Metadata: map[string]interface{}{ + models.MetaKeyAutomationTaskMode: string(automation.TaskModeAutomationRun), + models.MetaKeyAutomationRepositoryMode: string(automation.RepositoryModeNone), + models.MetaKeyAutomationTargetTaskID: "target-task-1", + models.MetaKeyDeferredLaunch: "must-survive", + }, + CreatedAt: now, + UpdatedAt: now, + })) + require.NoError(t, repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "continuation-session", TaskID: "continuation-task", + State: models.TaskSessionStateWaitingForInput, StartedAt: now, UpdatedAt: now, + })) + + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + a := &automation.Automation{ + ID: "a-merged", WorkspaceID: "ws-1", ContinuationTaskID: "continuation-task", + ContinuationPolicy: automation.ContinuationPolicyReuseThread, + TaskMode: automation.TaskModeAutomationRun, + RepositoryMode: automation.RepositoryModeNone, + } + evt := &automation.AutomationTriggeredEvent{ + TriggerType: automation.TriggerTypeGitHubPRMerged, + TriggerData: json.RawMessage(`{"task_id":"target-task-2"}`), + } + metadata := map[string]interface{}{ + "trigger_id": "trg-2", + "trigger_type": string(automation.TriggerTypeGitHubPRMerged), + models.MetaKeyAutomationTargetTaskID: "target-task-2", + // Deliberately present (with a different value than the task already + // carries) so the skip branch in refreshAutomationContinuationMetadata + // actually executes. Without this key in the input map, the assertion + // below would pass even with no skip at all — a merge trivially + // preserves a key the source map never mentions. + models.MetaKeyDeferredLaunch: "attacker-supplied-value", + } + + task, session, action, reason, err := svc.prepareAutomationTask(ctx, a, evt, "title", "prompt", metadata) + require.NoError(t, err) + require.Equal(t, automation.ThreadActionResumed, action, "reason=%q", reason) + require.NotNil(t, session) + require.Equal(t, "target-task-1", task.Metadata[models.MetaKeyAutomationTargetTaskID], + "preparing a resumed task must not mutate the target before dispatch admission") + require.NoError(t, svc.refreshAutomationContinuationMetadata(ctx, task, metadata)) + require.Equal(t, "target-task-2", task.Metadata[models.MetaKeyAutomationTargetTaskID], + "the admitted firing must carry the SECOND firing's target, not the first") + + persisted, err := repo.GetTask(ctx, "continuation-task") + require.NoError(t, err) + require.Equal(t, "target-task-2", persisted.Metadata[models.MetaKeyAutomationTargetTaskID], + "the refreshed binding must be persisted, or the guard reads the stale value back out on the next lookup") + require.Equal(t, "must-survive", persisted.Metadata[models.MetaKeyDeferredLaunch], + "deferred launch ownership is server-managed and must not be clobbered by a metadata refresh, even when the firing's own metadata map carries a value for it") + require.Equal(t, "must-survive", task.Metadata[models.MetaKeyDeferredLaunch], + "the in-memory task must also keep the original value, not the firing's attempted overwrite") +} + +func TestPrepareAutomationTask_DefersRefreshUntilDispatchAdmission(t *testing.T) { + repo := setupTestRepo(t) + now := time.Now().UTC() + ctx := context.Background() + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ + ID: "ws-stop", Name: "Test", CreatedAt: now, UpdatedAt: now, + })) + require.NoError(t, repo.CreateTask(ctx, &models.Task{ + ID: "continuation-stop-task", + WorkspaceID: "ws-stop", + Origin: models.TaskOriginAutomationRun, + Metadata: map[string]interface{}{ + models.MetaKeyAutomationTaskMode: string(automation.TaskModeAutomationRun), + models.MetaKeyAutomationRepositoryMode: string(automation.RepositoryModeNone), + models.MetaKeyAutomationTargetTaskID: "target-before-stop", + }, + CreatedAt: now, + UpdatedAt: now, + })) + require.NoError(t, repo.CreateTaskSession(ctx, &models.TaskSession{ + ID: "continuation-stop-session", TaskID: "continuation-stop-task", + State: models.TaskSessionStateWaitingForInput, StartedAt: now, UpdatedAt: now, + })) + + autoSvc := &rejectingAutomationService{ + stubAutomationService: &stubAutomationService{automation: &automation.Automation{ + ID: "a-stop", WorkspaceID: "ws-stop", ContinuationTaskID: "continuation-stop-task", + ContinuationPolicy: automation.ContinuationPolicyReuseThread, + TaskMode: automation.TaskModeAutomationRun, + RepositoryMode: automation.RepositoryModeNone, + }}, + dispatchErr: errors.New("run stopped before dispatch"), + } + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + svc.SetAutomationService(autoSvc) + evt := &automation.AutomationTriggeredEvent{ + RunID: "run-stop", + TriggerType: automation.TriggerTypeGitHubPRMerged, + TriggerData: json.RawMessage(`{"task_id":"target-after-stop"}`), + } + metadata := map[string]interface{}{ + "trigger_id": "trigger-stop", + "trigger_type": string(automation.TriggerTypeGitHubPRMerged), + models.MetaKeyAutomationTargetTaskID: "target-after-stop", + } + + task, session, action, reason, err := svc.prepareAutomationTask(ctx, autoSvc.automation, evt, "title", "prompt", metadata) + require.NoError(t, err) + require.Equal(t, automation.ThreadActionResumed, action, "reason=%q", reason) + require.NotNil(t, session) + require.Equal(t, "target-before-stop", task.Metadata[models.MetaKeyAutomationTargetTaskID], + "preparing a resumed task must not authorize a run before dispatch admission") + + svc.dispatchAutomationContinuation(ctx, autoSvc.automation, task, session, "prompt", metadata, evt.RunID, action, reason) + persisted, err := repo.GetTask(ctx, task.ID) + require.NoError(t, err) + require.Equal(t, "target-before-stop", persisted.Metadata[models.MetaKeyAutomationTargetTaskID], + "a run rejected before dispatch must not leave its target on the continuation task") +} + +func TestRefreshAutomationContinuationMetadataPublishesTaskUpdated(t *testing.T) { + repo := setupTestRepo(t) + now := time.Now().UTC() + ctx := context.Background() + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ + ID: "ws-publish", Name: "Test", CreatedAt: now, UpdatedAt: now, + })) + task := &models.Task{ + ID: "continuation-publish-task", WorkspaceID: "ws-publish", + Origin: models.TaskOriginAutomationRun, + Metadata: map[string]interface{}{ + models.MetaKeyAutomationTargetTaskID: "target-before-publish", + }, + CreatedAt: now, UpdatedAt: now, + } + require.NoError(t, repo.CreateTask(ctx, task)) + + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + publisher := &recordingTaskUpdatedPublisher{} + svc.SetTaskEventPublisher(publisher) + require.NoError(t, svc.refreshAutomationContinuationMetadata(ctx, task, map[string]interface{}{ + models.MetaKeyAutomationTargetTaskID: "target-after-publish", + })) + + require.Equal(t, []string{task.ID}, publisher.updatedTaskIDs, + "metadata changes must publish the refreshed task to live subscribers") +} + +func TestOrderedAutomationContinuationMetadataKeysPutsTargetFirst(t *testing.T) { + keys := orderedAutomationContinuationMetadataKeys(map[string]interface{}{ + "zeta": true, + models.MetaKeyDeferredLaunch: "must-survive", + models.MetaKeyAutomationTargetTaskID: "target", + "alpha": true, + }) + + require.Equal(t, []string{models.MetaKeyAutomationTargetTaskID, "alpha", "zeta"}, keys) +} + func TestRecordFailedRun_MergedPRDoesNotConsumeDedupKey(t *testing.T) { autoSvc := &stubAutomationService{} svc := &Service{automationService: autoSvc}