diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts index 22660fc64d6..9aa30e1db83 100644 --- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts +++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts @@ -76,6 +76,32 @@ type AttemptSpawnWorkspaceHoisted = { sessionManager: SessionManagerMocks; }; +export function createSubscriptionMock(): SubscriptionMock { + return { + assistantTexts: [] as string[], + toolMetas: [] as Array<{ toolName: string; meta?: string }>, + unsubscribe: () => {}, + setTerminalLifecycleMeta: () => {}, + waitForCompactionRetry: async () => {}, + getMessagingToolSentTexts: () => [] as string[], + getMessagingToolSentMediaUrls: () => [] as string[], + getMessagingToolSentTargets: () => [] as MessagingToolSend[], + getSuccessfulCronAdds: () => 0, + getReplayState: () => ({ + replayInvalid: false, + hadPotentialSideEffects: false, + }), + didSendViaMessagingTool: () => false, + didSendDeterministicApprovalPrompt: () => false, + getLastToolError: () => undefined, + getUsageTotals: () => undefined, + getCompactionCount: () => 0, + getItemLifecycle: () => ({ startedCount: 0, completedCount: 0, activeCount: 0 }), + isCompacting: () => false, + isCompactionInFlight: () => false, + }; +} + const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => { const spawnSubagentDirectMock = vi.fn(); const createAgentSessionMock = vi.fn(); @@ -87,31 +113,8 @@ const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => { const installToolResultContextGuardMock = vi.fn(() => () => {}); const flushPendingToolResultsAfterIdleMock = vi.fn(async () => {}); const releaseWsSessionMock = vi.fn(() => {}); - const subscribeEmbeddedPiSessionMock = vi.fn( - (_params) => - ({ - assistantTexts: [] as string[], - toolMetas: [] as Array<{ toolName: string; meta?: string }>, - unsubscribe: () => {}, - setTerminalLifecycleMeta: () => {}, - waitForCompactionRetry: async () => {}, - getMessagingToolSentTexts: () => [] as string[], - getMessagingToolSentMediaUrls: () => [] as string[], - getMessagingToolSentTargets: () => [] as MessagingToolSend[], - getSuccessfulCronAdds: () => 0, - getReplayState: () => ({ - replayInvalid: false, - hadPotentialSideEffects: false, - }), - didSendViaMessagingTool: () => false, - didSendDeterministicApprovalPrompt: () => false, - getLastToolError: () => undefined, - getUsageTotals: () => undefined, - getCompactionCount: () => 0, - getItemLifecycle: () => ({ startedCount: 0, completedCount: 0, activeCount: 0 }), - isCompacting: () => false, - isCompactionInFlight: () => false, - }) satisfies SubscriptionMock, + const subscribeEmbeddedPiSessionMock = vi.fn(() => + createSubscriptionMock(), ); const acquireSessionWriteLockMock = vi.fn(async (_params) => ({ release: async () => {}, @@ -655,32 +658,6 @@ export type MutableSession = { steer: (text: string) => Promise; }; -export function createSubscriptionMock(): SubscriptionMock { - return { - assistantTexts: [] as string[], - toolMetas: [] as Array<{ toolName: string; meta?: string }>, - unsubscribe: () => {}, - setTerminalLifecycleMeta: () => {}, - waitForCompactionRetry: async () => {}, - getMessagingToolSentTexts: () => [] as string[], - getMessagingToolSentMediaUrls: () => [] as string[], - getMessagingToolSentTargets: () => [] as MessagingToolSend[], - getSuccessfulCronAdds: () => 0, - getReplayState: () => ({ - replayInvalid: false, - hadPotentialSideEffects: false, - }), - didSendViaMessagingTool: () => false, - didSendDeterministicApprovalPrompt: () => false, - getLastToolError: () => undefined, - getUsageTotals: () => undefined, - getCompactionCount: () => 0, - getItemLifecycle: () => ({ startedCount: 0, completedCount: 0, activeCount: 0 }), - isCompacting: () => false, - isCompactionInFlight: () => false, - }; -} - type SessionPromptOverride = ( session: MutableSession, prompt: string, diff --git a/src/agents/pi-embedded-runner/run/auth-controller.ts b/src/agents/pi-embedded-runner/run/auth-controller.ts index ceaed0b0e3c..10f92b3a590 100644 --- a/src/agents/pi-embedded-runner/run/auth-controller.ts +++ b/src/agents/pi-embedded-runner/run/auth-controller.ts @@ -115,6 +115,31 @@ export function createEmbeddedRunAuthController(params: { const nextRuntimeAuthGeneration = () => (params.getRuntimeAuthState()?.generation ?? 0) + 1; + const prepareRuntimeAuthForModel = async (prepareParams: { + runtimeModel: Model; + apiKey: string; + authMode: string; + profileId?: string; + }) => + prepareProviderRuntimeAuth({ + provider: prepareParams.runtimeModel.provider, + config: params.config, + workspaceDir: params.workspaceDir, + env: process.env, + context: { + config: params.config, + agentDir: params.agentDir, + workspaceDir: params.workspaceDir, + env: process.env, + provider: prepareParams.runtimeModel.provider, + modelId: params.getModelId(), + model: prepareParams.runtimeModel, + apiKey: prepareParams.apiKey, + authMode: prepareParams.authMode, + profileId: prepareParams.profileId, + }, + }); + const clearRuntimeAuthRefreshTimer = () => { const runtimeAuthState = params.getRuntimeAuthState(); if (!runtimeAuthState?.refreshTimer) { @@ -152,23 +177,11 @@ export function createEmbeddedRunAuthController(params: { } const runtimeModel = params.getRuntimeModel(); params.log.debug(`Refreshing runtime auth for ${runtimeModel.provider} (${reason})...`); - const preparedAuth = await prepareProviderRuntimeAuth({ - provider: runtimeModel.provider, - config: params.config, - workspaceDir: params.workspaceDir, - env: process.env, - context: { - config: params.config, - agentDir: params.agentDir, - workspaceDir: params.workspaceDir, - env: process.env, - provider: runtimeModel.provider, - modelId: params.getModelId(), - model: runtimeModel, - apiKey: sourceApiKey, - authMode: currentRuntimeAuthState?.authMode ?? "unknown", - profileId: currentRuntimeAuthState?.profileId, - }, + const preparedAuth = await prepareRuntimeAuthForModel({ + runtimeModel, + apiKey: sourceApiKey, + authMode: currentRuntimeAuthState?.authMode ?? "unknown", + profileId: currentRuntimeAuthState?.profileId, }); if (!preparedAuth?.apiKey) { throw new Error( @@ -360,23 +373,11 @@ export function createEmbeddedRunAuthController(params: { } let runtimeAuthHandled = false; const runtimeModel = params.getRuntimeModel(); - const preparedAuth = await prepareProviderRuntimeAuth({ - provider: runtimeModel.provider, - config: params.config, - workspaceDir: params.workspaceDir, - env: process.env, - context: { - config: params.config, - agentDir: params.agentDir, - workspaceDir: params.workspaceDir, - env: process.env, - provider: runtimeModel.provider, - modelId: params.getModelId(), - model: runtimeModel, - apiKey: apiKeyInfo.apiKey, - authMode: apiKeyInfo.mode, - profileId: apiKeyInfo.profileId, - }, + const preparedAuth = await prepareRuntimeAuthForModel({ + runtimeModel, + apiKey: apiKeyInfo.apiKey, + authMode: apiKeyInfo.mode, + profileId: apiKeyInfo.profileId, }); applyPreparedRuntimeRequestOverrides({ runtimeModel, preparedAuth: preparedAuth ?? {} }); if (preparedAuth?.apiKey) { diff --git a/src/agents/pi-embedded-runner/run/incomplete-turn.ts b/src/agents/pi-embedded-runner/run/incomplete-turn.ts index 9341af36100..7da969a6398 100644 --- a/src/agents/pi-embedded-runner/run/incomplete-turn.ts +++ b/src/agents/pi-embedded-runner/run/incomplete-turn.ts @@ -275,6 +275,22 @@ function isEmptyResponseAssistantTurn(params: { return true; } +function shouldSkipPlanningOnlyRetry(params: { + aborted: boolean; + timedOut: boolean; + attempt: IncompleteTurnAttempt; +}): boolean { + return Boolean( + params.aborted || + params.timedOut || + params.attempt.clientToolCall || + params.attempt.yieldDetected || + params.attempt.didSendDeterministicApprovalPrompt || + params.attempt.lastToolError || + params.attempt.replayMetadata.hadPotentialSideEffects, + ); +} + export function resolveReasoningOnlyRetryInstruction(params: { provider?: string; modelId?: string; @@ -282,15 +298,7 @@ export function resolveReasoningOnlyRetryInstruction(params: { timedOut: boolean; attempt: IncompleteTurnAttempt; }): string | null { - if ( - params.aborted || - params.timedOut || - params.attempt.clientToolCall || - params.attempt.yieldDetected || - params.attempt.didSendDeterministicApprovalPrompt || - params.attempt.lastToolError || - params.attempt.replayMetadata.hadPotentialSideEffects - ) { + if (shouldSkipPlanningOnlyRetry(params)) { return null; } @@ -325,15 +333,7 @@ export function resolveEmptyResponseRetryInstruction(params: { timedOut: boolean; attempt: IncompleteTurnAttempt; }): string | null { - if ( - params.aborted || - params.timedOut || - params.attempt.clientToolCall || - params.attempt.yieldDetected || - params.attempt.didSendDeterministicApprovalPrompt || - params.attempt.lastToolError || - params.attempt.replayMetadata.hadPotentialSideEffects - ) { + if (shouldSkipPlanningOnlyRetry(params)) { return null; }