diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow.go b/apps/backend/internal/orchestrator/event_handlers_workflow.go index 058e8a9b5e..42991933a5 100644 --- a/apps/backend/internal/orchestrator/event_handlers_workflow.go +++ b/apps/backend/internal/orchestrator/event_handlers_workflow.go @@ -467,6 +467,30 @@ func (s *Service) resolveStepAgentProfile(ctx context.Context, step *wfmodels.Wo return "" } +// sessionWasWorkflowSwitched reports whether the session was spawned by a +// previous workflow step's agent_profile override (createNewSessionForStep) +// rather than by the user. Used to decide whether transitioning to a step +// with no override should revert to the task default. +func sessionWasWorkflowSwitched(session *models.TaskSession) bool { + if session == nil || session.Metadata == nil { + return false + } + v, _ := session.Metadata[models.SessionMetaKeyCreatedBy].(string) + return v == models.SessionCreatedByWorkflowSwitch +} + +// tagSessionAsWorkflowSwitched marks a session's metadata so it's recognised +// as workflow-spawned by sessionWasWorkflowSwitched. Used by every code path +// that ends up with a session whose agent_profile_id was decided by the +// workflow step override rather than by the user. Uses the atomic +// SetSessionMetadataKey (json_set) so other metadata keys are preserved. +func (s *Service) tagSessionAsWorkflowSwitched(ctx context.Context, sessionID string) { + if err := s.repo.SetSessionMetadataKey(ctx, sessionID, models.SessionMetaKeyCreatedBy, models.SessionCreatedByWorkflowSwitch); err != nil { + s.logger.Warn("failed to persist workflow-switch tag", + zap.String("session_id", sessionID), zap.Error(err)) + } +} + // switchSessionForStep activates a session for the new agent profile. // If an existing session on this task already uses the target profile it is // reused (re-promoted to primary, brought out of COMPLETED if it had been @@ -558,6 +582,12 @@ func (s *Service) reuseSessionForStep(ctx context.Context, taskID string, curren s.reviveReusedSession(ctx, existing) } + // Note: do not stamp created_by here. Reused sessions keep whatever tag + // they had when first created — workflow-spawned ones (createNewSessionForStep, + // StartCreatedSession, startTask) already carry the workflow_switch tag, + // and user-created ones must stay untagged so a later plain step doesn't + // silently revert their explicit profile choice to the task default. + if err := s.SetPrimarySession(ctx, existing.ID); err != nil { s.logger.Warn("failed to set reused session as primary", zap.String("session_id", existing.ID), zap.Error(err)) @@ -640,6 +670,12 @@ func (s *Service) createNewSessionForStep(ctx context.Context, taskID string, cu return nil, fmt.Errorf("failed to get new session: %w", err) } + // Tag the session as workflow-spawned so a later transition to a plain + // step (no agent_profile override) knows it's safe to revert to the task + // default. User-created sessions intentionally lack this tag — reverting + // them would override an explicit user choice. + s.tagSessionAsWorkflowSwitched(ctx, newSession.ID) + // Inherit the task environment from the old session — the workspace is shared // across sessions within the same task, so the new session can reuse the // existing agentctl connection and workspace files. @@ -733,9 +769,13 @@ func (s *Service) maybySwitchSessionForProfile( } effectiveProfile := s.resolveStepAgentProfile(ctx, step) - // When the step has no override, fall back to the task's original agent profile - // so that moving to a step without an agent_profile reverts to the default. - if effectiveProfile == "" { + // When the step has no override, fall back to the task's original agent + // profile so that moving to a step without an agent_profile reverts to the + // default. Only do this when the current session was itself spawned by a + // previous step's override — a user-chosen session (e.g. "New Agent" + // button) must not be silently replaced by the task default, which would + // override the user's explicit choice. + if effectiveProfile == "" && sessionWasWorkflowSwitched(session) { task, err := s.repo.GetTask(ctx, taskID) if err == nil && task.Metadata != nil { if pid, ok := task.Metadata[models.MetaKeyAgentProfileID].(string); ok && pid != "" { diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go b/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go index 5265c755cd..4e52a16c45 100644 --- a/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_workflow_profile_test.go @@ -731,6 +731,180 @@ func TestProcessOnEnter_ProfileSwitch(t *testing.T) { t.Errorf("expected 1 session, got %d", len(sessions)) } }) + + // The user created the task with profile-a, then manually added a new + // session with profile-b ("New Agent" button). When the workflow + // transitions to a step with no agent_profile_id override, the user's + // explicit choice (profile-b) must win — we must NOT silently switch + // back to the task's original profile-a just because that's what + // task.Metadata[agent_profile_id] still says. + t.Run("keeps user-chosen session when step has no override", func(t *testing.T) { + repo := setupTestRepo(t) + now := time.Now().UTC() + + ws := &models.Workspace{ID: "ws1", Name: "Test", CreatedAt: now, UpdatedAt: now} + _ = repo.CreateWorkspace(ctx, ws) + wf := &models.Workflow{ID: "wf1", WorkspaceID: "ws1", Name: "WF", CreatedAt: now, UpdatedAt: now} + _ = repo.CreateWorkflow(ctx, wf) + // Task was created with profile-a as the default agent. + task := &models.Task{ + ID: "t1", WorkflowID: "wf1", WorkflowStepID: "step1", + Title: "Test", Description: "desc", State: v1.TaskStateInProgress, + Metadata: map[string]interface{}{models.MetaKeyAgentProfileID: "profile-a"}, + CreatedAt: now, UpdatedAt: now, + } + _ = repo.CreateTask(ctx, task) + + // User clicked "New Agent" and started a profile-b session — it has + // no created_by metadata tag because it was user-chosen, not spawned + // by a workflow step override. + session := &models.TaskSession{ + ID: "s1", + TaskID: "t1", + AgentProfileID: "profile-b", + ExecutorID: "exec-local", + ExecutorProfileID: "ep1", + State: models.TaskSessionStateWaitingForInput, + IsPrimary: true, + StartedAt: now, + UpdatedAt: now, + } + _ = repo.CreateTaskSession(ctx, session) + + sg := newMockStepGetter() + step := &wfmodels.WorkflowStep{ + ID: "step1", + WorkflowID: "wf1", + Name: "Review", + // No AgentProfileID — step has no override. + } + sg.steps["step1"] = step + + svc := createTestService(repo, sg, newMockTaskRepo()) + svc.processOnEnter(ctx, "t1", session, step, "desc") + + // Critical: no new profile-a session should be spawned, and the + // user-chosen profile-b session must NOT be marked COMPLETED. + sessions, err := repo.ListTaskSessions(ctx, "t1") + if err != nil { + t.Fatalf("failed to list sessions: %v", err) + } + if len(sessions) != 1 { + t.Errorf("expected 1 session (no respawn), got %d", len(sessions)) + } + updated, err := repo.GetTaskSession(ctx, "s1") + if err != nil { + t.Fatalf("failed to get session: %v", err) + } + if updated.State == models.TaskSessionStateCompleted { + t.Error("user-chosen session must not be completed when step has no override") + } + if updated.AgentProfileID != "profile-b" { + t.Errorf("primary session must remain profile-b, got %q", updated.AgentProfileID) + } + }) + + // Inverse of the above: when the current session was spawned by a + // previous step's agent_profile override, transitioning to a plain step + // SHOULD revert to the task default. + t.Run("reverts to task default when current session was workflow-spawned", func(t *testing.T) { + repo := setupTestRepo(t) + now := time.Now().UTC() + + ws := &models.Workspace{ID: "ws1", Name: "Test", CreatedAt: now, UpdatedAt: now} + _ = repo.CreateWorkspace(ctx, ws) + wf := &models.Workflow{ID: "wf1", WorkspaceID: "ws1", Name: "WF", CreatedAt: now, UpdatedAt: now} + _ = repo.CreateWorkflow(ctx, wf) + task := &models.Task{ + ID: "t1", WorkflowID: "wf1", WorkflowStepID: "step2", + Title: "Test", Description: "desc", State: v1.TaskStateInProgress, + Metadata: map[string]interface{}{models.MetaKeyAgentProfileID: "profile-a"}, + CreatedAt: now, UpdatedAt: now, + } + _ = repo.CreateTask(ctx, task) + + // Session was spawned by createNewSessionForStep — tagged accordingly. + session := &models.TaskSession{ + ID: "s1", + TaskID: "t1", + AgentProfileID: "profile-b", + ExecutorID: "exec-local", + ExecutorProfileID: "ep1", + State: models.TaskSessionStateWaitingForInput, + IsPrimary: true, + Metadata: map[string]interface{}{models.SessionMetaKeyCreatedBy: models.SessionCreatedByWorkflowSwitch}, + StartedAt: now, + UpdatedAt: now, + } + _ = repo.CreateTaskSession(ctx, session) + + taskRepo := newMockTaskRepo() + taskRepo.tasks["t1"] = &v1.Task{ + ID: "t1", WorkspaceID: "ws1", WorkflowID: "wf1", + Title: "Test", Description: "desc", State: v1.TaskStateInProgress, + } + + sg := newMockStepGetter() + step := &wfmodels.WorkflowStep{ + ID: "step2", + WorkflowID: "wf1", + Name: "Done", + // No AgentProfileID — plain step. + } + sg.steps["step2"] = step + + agentMgr := &mockAgentManager{repoForExecutionLookup: repo} + log := testLogger() + exec := executor.NewExecutor(agentMgr, repo, log, executor.ExecutorConfig{}) + sched := scheduler.NewScheduler(queue.NewTaskQueue(100), exec, taskRepo, log, scheduler.SchedulerConfig{}) + svc := &Service{ + logger: log, + repo: repo, + workflowStepGetter: sg, + taskRepo: taskRepo, + agentManager: agentMgr, + messageQueue: messagequeue.NewServiceMemory(log), + executor: exec, + scheduler: sched, + } + + svc.processOnEnter(ctx, "t1", session, step, "desc") + + oldSession, err := repo.GetTaskSession(ctx, "s1") + if err != nil { + t.Fatalf("failed to get old session: %v", err) + } + if oldSession.State != models.TaskSessionStateCompleted { + t.Errorf("workflow-spawned session should be completed when reverting to default, got %s", oldSession.State) + } + + sessions, err := repo.ListTaskSessions(ctx, "t1") + if err != nil { + t.Fatalf("failed to list sessions: %v", err) + } + if len(sessions) != 2 { + t.Fatalf("expected exactly 2 sessions (old + reverted), got %d", len(sessions)) + } + var newSession *models.TaskSession + profileACount := 0 + for _, s := range sessions { + if s.AgentProfileID == "profile-a" { + profileACount++ + } + if s.ID != "s1" { + newSession = s + } + } + if profileACount != 1 { + t.Fatalf("expected exactly one profile-a session after revert, got %d", profileACount) + } + if newSession == nil { + t.Fatal("expected a new session for task default profile") + } + if newSession.AgentProfileID != "profile-a" { + t.Errorf("expected reverted profile-a, got %q", newSession.AgentProfileID) + } + }) } func TestSwitchSessionForStep_PreservesOldSessionOnFailure(t *testing.T) { diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index f6adad737a..cddb33372d 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/apps/backend/internal/orchestrator/task_operations.go @@ -257,6 +257,12 @@ func (s *Service) StartCreatedSession(ctx context.Context, taskID, sessionID, ag zap.String("old_profile", session.AgentProfileID), zap.String("new_profile", effectiveProfileID)) session.AgentProfileID = effectiveProfileID + // Tag as workflow-spawned so a later transition to a plain step can + // correctly revert to the task default. Without this, the in-place + // profile mutation here is invisible to maybySwitchSessionForProfile, + // which would treat the session as user-chosen and leave it on the + // override forever. + s.tagSessionAsWorkflowSwitched(ctx, sessionID) // Re-resolve the agent profile snapshot so the tab shows the correct agent logo/name. // Set a minimal snapshot first so stale data is never persisted if resolution fails. session.AgentProfileSnapshot = map[string]interface{}{"id": effectiveProfileID} @@ -485,7 +491,9 @@ func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID s // Resolve the workflow step's agent profile override. // The frontend may pass the workspace default profile, but the step may // require a different agent (e.g., Codex on "In Progress", Auggie on "Review"). + callerProfileID := agentProfileID agentProfileID = s.resolveEffectiveAgentProfile(ctx, taskID, workflowStepID, agentProfileID) + overrideApplied := agentProfileID != callerProfileID // Fetch the task from the repository to get complete task info task, err := s.scheduler.GetTask(ctx, taskID) @@ -515,6 +523,14 @@ func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID s return nil, err } + // When the workflow step overrode the caller's profile, tag the session + // so maybySwitchSessionForProfile knows it can revert to the task default + // on a later plain-step transition. Without this tag, the session looks + // indistinguishable from a user-chosen one and stays on the override. + if overrideApplied { + s.tagSessionAsWorkflowSwitched(ctx, sessionID) + } + effectivePrompt, planModeActive := s.applyWorkflowAndPlanMode(ctx, effectivePrompt, task.ID, sessionID, workflowStepID, planMode, task.IsEphemeral) // Inject config context for config-mode sessions (dedicated settings chat) @@ -592,6 +608,15 @@ func (s *Service) prepareSessionForStart( if ensureErr != nil { return "", ensureErr } + // The office/assignee path bypasses StartCreatedSession's override + // mutation: the session is created directly with the assignee + // profile (which falls back to the step's agent_profile_id via the + // runner projection). If that profile differs from the caller's, + // the assignment was workflow-driven — tag so a later transition + // to a plain step knows it can revert to the task default. + if agentProfileID != "" && session.AgentProfileID != "" && session.AgentProfileID != agentProfileID { + s.tagSessionAsWorkflowSwitched(ctx, session.ID) + } return session.ID, nil } return s.executor.PrepareSession(ctx, task, agentProfileID, executorID, executorProfileID, workflowStepID) diff --git a/apps/backend/internal/task/models/models.go b/apps/backend/internal/task/models/models.go index f0fb341662..1aa1bc16b3 100644 --- a/apps/backend/internal/task/models/models.go +++ b/apps/backend/internal/task/models/models.go @@ -45,6 +45,16 @@ const ( MetaKeyWorkspacePath = "workspace_path" ) +// TaskSession.Metadata key that records how the session came into existence. +// Read by maybySwitchSessionForProfile to decide whether transitioning to a +// step with no agent_profile override should revert to the task default: +// workflow-spawned sessions (created_by=workflow_switch) -> yes, revert. +// user-created sessions (no created_by tag) -> no, preserve. +const ( + SessionMetaKeyCreatedBy = "created_by" + SessionCreatedByWorkflowSwitch = "workflow_switch" +) + // Task origin values for the Origin field. const ( TaskOriginManual = "manual" diff --git a/apps/backend/internal/task/repository/sqlite/session.go b/apps/backend/internal/task/repository/sqlite/session.go index af728ca311..b84db3bd59 100644 --- a/apps/backend/internal/task/repository/sqlite/session.go +++ b/apps/backend/internal/task/repository/sqlite/session.go @@ -443,6 +443,11 @@ func (r *Repository) UpdateTaskSession(ctx context.Context, session *models.Task return fmt.Errorf("failed to serialize repository snapshot: %w", err) } + // metadata is NOT written here — callers wanting to change it must use + // UpdateSessionMetadata or SetSessionMetadataKey. A full-row write here + // would clobber metadata set via those side-channel paths since the + // caller's in-memory copy may be stale. + // agent_profile_id is stored as NULL when empty so the partial unique // index over (task_id, agent_profile_id) ignores kanban / quick-chat rows. var agentProfileID interface{} diff --git a/apps/web/e2e/helpers/api-client.ts b/apps/web/e2e/helpers/api-client.ts index 3096697f84..436a77cc7e 100644 --- a/apps/web/e2e/helpers/api-client.ts +++ b/apps/web/e2e/helpers/api-client.ts @@ -1011,6 +1011,7 @@ export class ApiClient { title: string; primary_session_id?: string | null; state?: string; + workflow_step_id?: string; repositories?: Array<{ id: string; task_id: string; diff --git a/apps/web/e2e/tests/workflow/workflow-agent-switch.spec.ts b/apps/web/e2e/tests/workflow/workflow-agent-switch.spec.ts index 932da92844..b9b49c0571 100644 --- a/apps/web/e2e/tests/workflow/workflow-agent-switch.spec.ts +++ b/apps/web/e2e/tests/workflow/workflow-agent-switch.spec.ts @@ -710,6 +710,119 @@ test.describe("Workflow agent profile switching", () => { await expect(session.primaryStarInTab("Profile A")).toBeVisible({ timeout: 30_000 }); }); + /** + * A user manually adds a session with a different agent (the "New Agent" + * button), and the active step has on_turn_complete: move_to_next with the + * next step having no agent override. The user's explicit profileB choice + * must win — we must NOT silently respawn a profileA session after each + * completed turn. + */ + test("user-added agent session is not respawned to task default on turn complete", async ({ + apiClient, + seedData, + }) => { + test.setTimeout(60_000); + const { profileA, profileB } = await createProfiles(apiClient); + + // Step1: no override, auto_start + on_turn_complete → next. + // Step2: no override (this is the trigger condition for the bug). + const workflow = await apiClient.createWorkflow(seedData.workspaceId, "Respawn Regression"); + const step1 = await apiClient.createWorkflowStep(workflow.id, "Step1", 0, { + is_start_step: true, + }); + const step2 = await apiClient.createWorkflowStep(workflow.id, "Step2", 1); + await apiClient.createWorkflowStep(workflow.id, "Done", 2); + + // step1 has auto_start + on_turn_complete:move_to_next but NO startup + // prompt — a prompt would race the user's New Agent launch and could + // advance the task to step2 before the user-added session even exists, + // letting this test pass without exercising the bug path. + await apiClient.updateWorkflowStep(step1.id, { + events: { + on_enter: [{ type: "auto_start_agent" }], + on_turn_complete: [{ type: "move_to_next" }], + }, + }); + // step2 deliberately has no agent_profile_id and no on_enter. + + // Task is created with profileA as the default agent. + const task = await apiClient.createTaskWithAgent( + seedData.workspaceId, + "Respawn Regression Task", + profileA.id, + { + workflow_id: workflow.id, + workflow_step_id: step1.id, + repository_ids: [seedData.repositoryId], + }, + ); + + // Wait for the auto-started profileA session to be ready AND confirm the + // task is still on step1 — without a startup prompt the only way to + // advance is via the user-added session's turn-complete below. + await expect + .poll( + async () => { + const { sessions } = await apiClient.listTaskSessions(task.id); + return sessions.find((s) => s.agent_profile_id === profileA.id)?.state ?? ""; + }, + { timeout: 30_000, message: "Waiting for profileA session to be ready" }, + ) + .toBe("WAITING_FOR_INPUT"); + await expect + .poll(async () => (await apiClient.getTask(task.id)).workflow_step_id, { + timeout: 5_000, + message: "Task should still be on step1 before New Agent", + }) + .toBe(step1.id); + + // User clicks "New Agent" — launch a profileB session on the same task. + await apiClient.launchSession({ + task_id: task.id, + agent_profile_id: profileB.id, + prompt: 'e2e:message("hello from B")', + intent: "start", + }); + + // Wait for the on_turn_complete handler to advance the task to Step2 — + // a deterministic signal that the handler ran (and either spawned a new + // session, the bug, or correctly did not). + await expect + .poll( + async () => { + const t = await apiClient.getTask(task.id); + return t.workflow_step_id; + }, + { timeout: 30_000, message: "Waiting for task to advance to step2" }, + ) + .toBe(step2.id); + + // The handler runs in a goroutine after the step-id write, so a buggy + // respawn could land a few moments after the advance is visible. Sample + // the session count repeatedly across a short window and fail fast if + // it ever exceeds 2. `expect.poll(...).toBe(2)` would return on the + // first matching sample and miss a delayed respawn — we need to prove + // the count *stays* at 2. + const stabilityWindowMs = 3_000; + const stabilityStart = Date.now(); + while (Date.now() - stabilityStart < stabilityWindowMs) { + const { sessions: stableSessions } = await apiClient.listTaskSessions(task.id); + expect(stableSessions.length, "no extra session should spawn within stability window").toBe( + 2, + ); + await new Promise((r) => setTimeout(r, 250)); + } + + // Critical assertion: no extra profileA session was spawned after the + // user's profileB turn completed. The task must still have exactly two + // sessions (the original profileA and the user-added profileB). + const { sessions } = await apiClient.listTaskSessions(task.id); + const profileASessions = sessions.filter((s) => s.agent_profile_id === profileA.id); + const profileBSessions = sessions.filter((s) => s.agent_profile_id === profileB.id); + expect(profileBSessions.length).toBe(1); + expect(profileASessions.length).toBe(1); + }); + test("reset context checkbox is disabled when step has agent profile override", async ({ testPage, apiClient,