refactor: dedupe embedded runner helpers

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<SubscribeEmbeddedPiSessionFn>(
-    (_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<SubscribeEmbeddedPiSessionFn>(() =>
+    createSubscriptionMock(),
   );
   const acquireSessionWriteLockMock = vi.fn<AcquireSessionWriteLockFn>(async (_params) => ({
     release: async () => {},
@@ -655,32 +658,6 @@ export type MutableSession = {
   steer: (text: string) => Promise<void>;
 };
 
-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,

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<Api>;
+    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) {

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