fix(team-mode): preserve member context on retry

This commit is contained in:
YeonGyu-Kim
2026-05-10 13:11:52 +09:00
parent b36389ef2c
commit d801c88888
6 changed files with 40 additions and 11 deletions
@@ -2,7 +2,7 @@ import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test"
const sharedLogMock = mock(() => {}) const sharedLogMock = mock(() => {})
const readConnectedProvidersCacheMock = mock(() => null) const readConnectedProvidersCacheMock = mock(() => null)
const readProviderModelsCacheMock = mock(() => null) const readProviderModelsCacheMock = mock((): { connected: string[] } | null => null)
const shouldRetryErrorMock = mock(() => true) const shouldRetryErrorMock = mock(() => true)
const getNextFallbackMock = mock((chain: Array<{ model: string }>, attempt: number) => chain[attempt]) const getNextFallbackMock = mock((chain: Array<{ model: string }>, attempt: number) => chain[attempt])
const hasMoreFallbacksMock = mock((chain: Array<{ model: string }>, attempt: number) => attempt < chain.length) const hasMoreFallbacksMock = mock((chain: Array<{ model: string }>, attempt: number) => attempt < chain.length)
@@ -88,7 +88,7 @@ function createMockConcurrencyManager(): ConcurrencyManager {
acquire: mock(async () => {}), acquire: mock(async () => {}),
getQueueLength: mock(() => 0), getQueueLength: mock(() => 0),
getActiveCount: mock(() => 0), getActiveCount: mock(() => 0),
} as unknown as ConcurrencyManager } as never
} }
function createMockClient(): { function createMockClient(): {
@@ -101,7 +101,7 @@ function createMockClient(): {
session: { session: {
abort: abortMock, abort: abortMock,
}, },
} as unknown as OpencodeClient, } as never,
abortMock, abortMock,
} }
} }
@@ -133,9 +133,9 @@ describe("tryFallbackRetry", () => {
}) })
beforeEach(() => { beforeEach(() => {
;(shouldRetryError as any).mockImplementation(() => true) shouldRetryError.mockImplementation(() => true)
;(selectFallbackProvider as any).mockImplementation((providers: string[]) => providers[0]) selectFallbackProvider.mockImplementation((providers: string[]) => providers[0])
;(readProviderModelsCache as any).mockReturnValue(null) readProviderModelsCache.mockReturnValue(null)
}) })
describe("#given retryable error with fallback chain", () => { describe("#given retryable error with fallback chain", () => {
@@ -260,6 +260,21 @@ describe("tryFallbackRetry", () => {
expect(args.processKey).toHaveBeenCalledWith(key) 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 () => { test("finalizes the failed attempt, creates a new pending attempt, and enqueues its explicit attemptID", async () => {
const args = createDefaultArgs({ const args = createDefaultArgs({
status: "running", status: "running",
@@ -308,13 +323,16 @@ describe("tryFallbackRetry", () => {
const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` const key = `${args.task.model!.providerID}/${args.task.model!.modelID}`
const queue = args.queuesByKey.get(key) const queue = args.queuesByKey.get(key)
expect(queue).toBeDefined() 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", () => { describe("#given non-retryable error", () => {
test("returns false when shouldRetryError returns false", async () => { test("returns false when shouldRetryError returns false", async () => {
;(shouldRetryError as any).mockImplementation(() => false) shouldRetryError.mockImplementation(() => false)
const args = createDefaultArgs() const args = createDefaultArgs()
const result = await tryFallbackRetry(args) const result = await tryFallbackRetry(args)
@@ -415,8 +433,8 @@ describe("tryFallbackRetry", () => {
describe("#given disconnected fallback providers with connected preferred provider", () => { describe("#given disconnected fallback providers with connected preferred provider", () => {
test("keeps fallback entry and selects connected preferred provider", async () => { test("keeps fallback entry and selects connected preferred provider", async () => {
;(readProviderModelsCache as any).mockReturnValueOnce({ connected: ["provider-a"] }) readProviderModelsCache.mockReturnValueOnce({ connected: ["provider-a"] })
;(selectFallbackProvider as any).mockImplementationOnce( selectFallbackProvider.mockImplementationOnce(
(_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b", (_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b",
) )
@@ -170,10 +170,12 @@ export async function tryFallbackRetry(args: {
parentModel: task.parentModel, parentModel: task.parentModel,
parentAgent: task.parentAgent, parentAgent: task.parentAgent,
parentTools: task.parentTools, parentTools: task.parentTools,
teamRunId: task.teamRunId,
model: nextModel, model: nextModel,
fallbackChain: task.fallbackChain, fallbackChain: task.fallbackChain,
category: task.category, category: task.category,
isUnstableAgent: task.isUnstableAgent, isUnstableAgent: task.isUnstableAgent,
onSessionCreated: task.onSessionCreated,
} }
if (previousSessionID) { if (previousSessionID) {
+1
View File
@@ -424,6 +424,7 @@ export class BackgroundManager {
fallbackChain: input.fallbackChain, fallbackChain: input.fallbackChain,
attemptCount: 0, attemptCount: 0,
category: input.category, category: input.category,
onSessionCreated: input.onSessionCreated,
} }
const firstAttempt = startAttempt(task, input.model) const firstAttempt = startAttempt(task, input.model)
+1
View File
@@ -65,6 +65,7 @@ export function createTask(input: LaunchInput): BackgroundTask {
parentModel: input.parentModel, parentModel: input.parentModel,
parentAgent: input.parentAgent, parentAgent: input.parentAgent,
model: input.model, model: input.model,
onSessionCreated: input.onSessionCreated,
} }
} }
+1
View File
@@ -77,6 +77,7 @@ export interface BackgroundTask {
isUnstableAgent?: boolean isUnstableAgent?: boolean
/** Category used for this task (e.g., 'quick', 'visual-engineering') */ /** Category used for this task (e.g., 'quick', 'visual-engineering') */
category?: string category?: string
onSessionCreated?: (sessionId: string) => void | Promise<void>
/** Pending retry notification details for the next spawned retry session */ /** Pending retry notification details for the next spawned retry session */
retryNotification?: { retryNotification?: {
previousSessionID?: string previousSessionID?: string
@@ -199,12 +199,18 @@ export async function createTeamRun(
skillContent: resolvedMember.systemContent, skillContent: resolvedMember.systemContent,
category: member.kind === "category" ? member.category : undefined, category: member.kind === "category" ? member.category : undefined,
sessionPermission: QUESTION_DENIED_SESSION_PERMISSION, sessionPermission: QUESTION_DENIED_SESSION_PERMISSION,
onSessionCreated: (sessionId) => { onSessionCreated: async (sessionId) => {
registerTeamSession(sessionId, { registerTeamSession(sessionId, {
teamRunId: runtimeState.teamRunId, teamRunId: runtimeState.teamRunId,
memberName: member.name, memberName: member.name,
role: member.name === spec.leadAgentId ? "lead" : "member", 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 resource.taskId = task.id