From d801c88888118e9cab435335f1e383ba627352a9 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 10 May 2026 13:11:52 +0900 Subject: [PATCH] fix(team-mode): preserve member context on retry --- .../fallback-retry-handler.test.ts | 38 ++++++++++++++----- .../fallback-retry-handler.ts | 2 + src/features/background-agent/manager.ts | 1 + src/features/background-agent/spawner.ts | 1 + src/features/background-agent/types.ts | 1 + src/features/team-mode/team-runtime/create.ts | 8 +++- 6 files changed, 40 insertions(+), 11 deletions(-) diff --git a/src/features/background-agent/fallback-retry-handler.test.ts b/src/features/background-agent/fallback-retry-handler.test.ts index a9c4f0cd1..f51ce2059 100644 --- a/src/features/background-agent/fallback-retry-handler.test.ts +++ b/src/features/background-agent/fallback-retry-handler.test.ts @@ -2,7 +2,7 @@ import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test" const sharedLogMock = mock(() => {}) const readConnectedProvidersCacheMock = mock(() => null) -const readProviderModelsCacheMock = mock(() => null) +const readProviderModelsCacheMock = mock((): { connected: string[] } | null => null) const shouldRetryErrorMock = mock(() => true) const getNextFallbackMock = mock((chain: Array<{ model: string }>, attempt: number) => chain[attempt]) const hasMoreFallbacksMock = mock((chain: Array<{ model: string }>, attempt: number) => attempt < chain.length) @@ -88,7 +88,7 @@ function createMockConcurrencyManager(): ConcurrencyManager { acquire: mock(async () => {}), getQueueLength: mock(() => 0), getActiveCount: mock(() => 0), - } as unknown as ConcurrencyManager + } as never } function createMockClient(): { @@ -101,7 +101,7 @@ function createMockClient(): { session: { abort: abortMock, }, - } as unknown as OpencodeClient, + } as never, abortMock, } } @@ -133,9 +133,9 @@ describe("tryFallbackRetry", () => { }) beforeEach(() => { - ;(shouldRetryError as any).mockImplementation(() => true) - ;(selectFallbackProvider as any).mockImplementation((providers: string[]) => providers[0]) - ;(readProviderModelsCache as any).mockReturnValue(null) + shouldRetryError.mockImplementation(() => true) + selectFallbackProvider.mockImplementation((providers: string[]) => providers[0]) + readProviderModelsCache.mockReturnValue(null) }) describe("#given retryable error with fallback chain", () => { @@ -260,6 +260,21 @@ describe("tryFallbackRetry", () => { expect(args.processKey).toHaveBeenCalledWith(key) }) + test("preserves team identity and session callback in retry input", async () => { + const onSessionCreated = mock(async () => {}) + const args = createDefaultArgs({ + teamRunId: "team-run-1", + onSessionCreated, + }) + + await tryFallbackRetry(args) + + const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` + const retryInput = args.queuesByKey.get(key)?.[0]?.input + expect(retryInput?.teamRunId).toBe("team-run-1") + expect(retryInput?.onSessionCreated).toBe(onSessionCreated) + }) + test("finalizes the failed attempt, creates a new pending attempt, and enqueues its explicit attemptID", async () => { const args = createDefaultArgs({ status: "running", @@ -308,13 +323,16 @@ describe("tryFallbackRetry", () => { const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` const queue = args.queuesByKey.get(key) expect(queue).toBeDefined() - expect((queue?.[0] as QueueItem & { attemptID?: string })?.attemptID).toBe(nextAttempt?.attemptId) + const queuedAttemptID = queue?.[0]?.attemptID + expect(queuedAttemptID).toBeDefined() + expect(nextAttempt?.attemptId).toBeDefined() + expect(queuedAttemptID).toBe(nextAttempt?.attemptId ?? "") }) }) describe("#given non-retryable error", () => { test("returns false when shouldRetryError returns false", async () => { - ;(shouldRetryError as any).mockImplementation(() => false) + shouldRetryError.mockImplementation(() => false) const args = createDefaultArgs() const result = await tryFallbackRetry(args) @@ -415,8 +433,8 @@ describe("tryFallbackRetry", () => { describe("#given disconnected fallback providers with connected preferred provider", () => { test("keeps fallback entry and selects connected preferred provider", async () => { - ;(readProviderModelsCache as any).mockReturnValueOnce({ connected: ["provider-a"] }) - ;(selectFallbackProvider as any).mockImplementationOnce( + readProviderModelsCache.mockReturnValueOnce({ connected: ["provider-a"] }) + selectFallbackProvider.mockImplementationOnce( (_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b", ) diff --git a/src/features/background-agent/fallback-retry-handler.ts b/src/features/background-agent/fallback-retry-handler.ts index f5f31a6ad..d61a839b1 100644 --- a/src/features/background-agent/fallback-retry-handler.ts +++ b/src/features/background-agent/fallback-retry-handler.ts @@ -170,10 +170,12 @@ export async function tryFallbackRetry(args: { parentModel: task.parentModel, parentAgent: task.parentAgent, parentTools: task.parentTools, + teamRunId: task.teamRunId, model: nextModel, fallbackChain: task.fallbackChain, category: task.category, isUnstableAgent: task.isUnstableAgent, + onSessionCreated: task.onSessionCreated, } if (previousSessionID) { diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index e717e806f..1b0d4acdf 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -424,6 +424,7 @@ export class BackgroundManager { fallbackChain: input.fallbackChain, attemptCount: 0, category: input.category, + onSessionCreated: input.onSessionCreated, } const firstAttempt = startAttempt(task, input.model) diff --git a/src/features/background-agent/spawner.ts b/src/features/background-agent/spawner.ts index 2e49537b7..bfbe675da 100644 --- a/src/features/background-agent/spawner.ts +++ b/src/features/background-agent/spawner.ts @@ -65,6 +65,7 @@ export function createTask(input: LaunchInput): BackgroundTask { parentModel: input.parentModel, parentAgent: input.parentAgent, model: input.model, + onSessionCreated: input.onSessionCreated, } } diff --git a/src/features/background-agent/types.ts b/src/features/background-agent/types.ts index d72ab4a5e..76c242a41 100644 --- a/src/features/background-agent/types.ts +++ b/src/features/background-agent/types.ts @@ -77,6 +77,7 @@ export interface BackgroundTask { isUnstableAgent?: boolean /** Category used for this task (e.g., 'quick', 'visual-engineering') */ category?: string + onSessionCreated?: (sessionId: string) => void | Promise /** Pending retry notification details for the next spawned retry session */ retryNotification?: { previousSessionID?: string diff --git a/src/features/team-mode/team-runtime/create.ts b/src/features/team-mode/team-runtime/create.ts index 5bfdf9e01..7671b03ed 100644 --- a/src/features/team-mode/team-runtime/create.ts +++ b/src/features/team-mode/team-runtime/create.ts @@ -199,12 +199,18 @@ export async function createTeamRun( skillContent: resolvedMember.systemContent, category: member.kind === "category" ? member.category : undefined, sessionPermission: QUESTION_DENIED_SESSION_PERMISSION, - onSessionCreated: (sessionId) => { + onSessionCreated: async (sessionId) => { registerTeamSession(sessionId, { teamRunId: runtimeState.teamRunId, memberName: member.name, role: member.name === spec.leadAgentId ? "lead" : "member", }) + runtimeState = await transitionRuntimeState(runtimeState.teamRunId, (currentState) => ({ + ...currentState, + members: currentState.members.map((currentMember, currentIndex) => currentIndex === memberIndex + ? { ...currentMember, sessionId, status: "running" } + : currentMember), + }), config) }, }) resource.taskId = task.id