fix: respect user-chosen agent on workflow step transitions (#941)

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 != "" {

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

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)

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"

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

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;

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,