fix: stabilize kiro desktop authorize callback flow

background/kiro/desktop-authorize-runner.js

@@ -271,7 +271,7 @@
       return result;
     }
 
-    function waitForResolved(expectedState = '', timeoutMs = 120000) {
+    function waitForResolved(expectedState = '', timeoutMs = DEFAULT_KIRO_PAGE_LOAD_TIMEOUT_MS) {
       const stateKey = cleanString(expectedState);
       const immediate = consumeResolved(stateKey);
       if (immediate) {
@@ -289,7 +289,7 @@
             pendingSessions.set(stateKey, nextSession);
           }
           reject(new Error('等待桌面授权回调超时。'));
-        }, Math.max(1000, Math.floor(Number(timeoutMs) || 120000)));
+        }, Math.max(1000, Math.floor(Number(timeoutMs) || DEFAULT_KIRO_PAGE_LOAD_TIMEOUT_MS)));
         session.waiters.push({
           resolve: (result) => {
             clearTimeout(timer);
@@ -414,6 +414,52 @@
       }));
     }
 
+    function isMissingTabError(error) {
+      return /No tab with id/i.test(getErrorMessage(error));
+    }
+
+    async function finalizeDesktopAuthorizeCallback(currentState = {}, runtimeState = {}, resolvedCallback = {}, nodeId = '') {
+      if (resolvedCallback?.error) {
+        throw new Error(`桌面授权回调失败：${resolvedCallback.error}`);
+      }
+
+      const authorizationCode = cleanString(resolvedCallback?.code);
+      if (!authorizationCode) {
+        throw new Error('桌面授权回调缺少 authorization code。');
+      }
+
+      const tokenResult = await desktopClientApi.exchangeDesktopAuthorizationCode({
+        region: runtimeState.desktopAuth?.region || DEFAULT_REGION,
+        clientId: runtimeState.desktopAuth?.clientId,
+        clientSecret: runtimeState.desktopAuth?.clientSecret,
+        redirectUri: runtimeState.desktopAuth?.redirectUri,
+        code: authorizationCode,
+        codeVerifier: runtimeState.desktopAuth?.codeVerifier,
+      }, fetchImpl);
+      const payload = await applyRuntimeState(currentState, {
+        session: {
+          currentStage: 'upload',
+          pageState: 'callback_captured',
+          pageUrl: resolvedCallback.url,
+          lastError: '',
+        },
+        desktopAuth: {
+          authorizationCode,
+          accessToken: tokenResult.accessToken,
+          refreshToken: tokenResult.refreshToken,
+          status: 'authorized',
+          authorizedAt: Date.now(),
+        },
+        upload: {
+          status: 'ready_to_upload',
+          error: '',
+        },
+      });
+      await log('步骤 8：桌面授权回调已捕获，Token 换取成功。', 'ok', nodeId);
+      await completeNodeFromBackground(nodeId, payload);
+      return payload;
+    }
+
     async function ensureDesktopAuthorizeTab(state = {}, options = {}) {
       const runtimeState = readKiroRuntime(state);
       let tabId = Number.isInteger(runtimeState.session?.desktopTabId)
@@ -534,6 +580,7 @@
     }
 
     async function executeDesktopAction(tabId, action, payload = {}, options = {}) {
+      const timeoutBudget = resolveTimeoutBudget(options);
       const result = await sendToContentScriptResilient(KIRO_DESKTOP_SOURCE_ID, {
         type: 'EXECUTE_KIRO_DESKTOP_AUTHORIZE_ACTION',
         step: options.step || 0,
@@ -543,9 +590,12 @@
           ...payload,
         },
       }, {
-        timeoutMs: 30000,
+        timeoutMs: timeoutBudget.getRemainingMs(1000),
         retryDelayMs: 700,
-        onRetryableError: buildDesktopRetryRecovery(tabId, options),
+        onRetryableError: buildDesktopRetryRecovery(tabId, {
+          ...options,
+          timeoutBudget,
+        }),
         logMessage: options.logMessage || '正在执行 Kiro 桌面授权动作...',
       });
       if (result?.error) {
@@ -816,7 +866,7 @@
 
     async function executeKiroCompleteDesktopAuthorize(state = {}) {
       const nodeId = String(state?.nodeId || 'kiro-complete-desktop-authorize').trim();
-      const currentState = await getExecutionState(state);
+      let currentState = await getExecutionState(state);
       let runtimeState = readKiroRuntime(currentState);
       const desktopState = cleanString(runtimeState.desktopAuth?.state);
       try {
@@ -839,71 +889,105 @@
           tabId: runtimeState.session?.desktopTabId,
         });
 
-        const deadline = Date.now() + 120000;
+        const timeoutBudget = createTimeoutBudget(DEFAULT_KIRO_PAGE_LOAD_TIMEOUT_MS);
+        const deadline = Date.now() + timeoutBudget.totalTimeoutMs;
+        let awaitingCallbackAfterConsent = false;
+        const updateLoopState = async (patch = {}) => {
+          const runtimePatch = mergeRuntimePatch(currentState, patch);
+          await setState(runtimePatch);
+          currentState = {
+            ...currentState,
+            ...runtimePatch,
+          };
+          runtimeState = readKiroRuntime(currentState);
+          return runtimePatch;
+        };
+
         while (Date.now() < deadline) {
           throwIfStopped();
 
           const resolvedCallback = callbackTracker.consumeResolved(desktopState);
           if (resolvedCallback) {
-            if (resolvedCallback.error) {
+            await finalizeDesktopAuthorizeCallback(currentState, runtimeState, resolvedCallback, nodeId);
-              throw new Error(`桌面授权回调失败：${resolvedCallback.error}`);
-            }
-            const tokenResult = await desktopClientApi.exchangeDesktopAuthorizationCode({
-              region: runtimeState.desktopAuth.region || DEFAULT_REGION,
-              clientId: runtimeState.desktopAuth.clientId,
-              clientSecret: runtimeState.desktopAuth.clientSecret,
-              redirectUri: runtimeState.desktopAuth.redirectUri,
-              code: resolvedCallback.code,
-              codeVerifier: runtimeState.desktopAuth.codeVerifier,
-            }, fetchImpl);
-            const payload = await applyRuntimeState(currentState, {
-              session: {
-                currentStage: 'upload',
-                pageState: 'callback_captured',
-                pageUrl: resolvedCallback.url,
-                lastError: '',
-              },
-              desktopAuth: {
-                authorizationCode: resolvedCallback.code,
-                accessToken: tokenResult.accessToken,
-                refreshToken: tokenResult.refreshToken,
-                status: 'authorized',
-                authorizedAt: Date.now(),
-              },
-              upload: {
-                status: 'ready_to_upload',
-                error: '',
-              },
-            });
-            await log('步骤 8：桌面授权回调已捕获，Token 换取成功。', 'ok', nodeId);
-            await completeNodeFromBackground(nodeId, payload);
             return;
           }
 
-          const tabId = await activateDesktopAuthorizeTab(currentState, {
+          if (awaitingCallbackAfterConsent) {
+            const waitedCallback = await callbackTracker.waitForResolved(
+              desktopState,
+              Math.min(timeoutBudget.getRemainingMs(1000), 1500)
+            ).catch(() => null);
+            if (waitedCallback) {
+              await finalizeDesktopAuthorizeCallback(currentState, runtimeState, waitedCallback, nodeId);
+              return;
+            }
+          }
+
+          let tabId = null;
+          if (awaitingCallbackAfterConsent) {
+            tabId = await getTabId(KIRO_DESKTOP_SOURCE_ID).catch(() => null);
+            if (!Number.isInteger(tabId)) {
+              await sleepWithStop(1000);
+              continue;
+            }
+
+            const trackedTab = await chrome?.tabs?.get?.(tabId).catch(() => null);
+            if (!trackedTab) {
+              await sleepWithStop(1000);
+              continue;
+            }
+
+            const trackedCallback = parseDesktopCallbackUrl(
+              trackedTab.url,
+              desktopState,
+              runtimeState.desktopAuth?.redirectPort
+            );
+            if (trackedCallback) {
+              await finalizeDesktopAuthorizeCallback(currentState, runtimeState, trackedCallback, nodeId);
+              return;
+            }
+
+            if (String(trackedTab.status || '') !== 'complete') {
+              await sleepWithStop(1000);
+              continue;
+            }
+          } else {
+            tabId = await activateDesktopAuthorizeTab(currentState, {
               missingUrlMessage: '缺少桌面授权地址，请先执行步骤 7。',
               openFailedMessage: '无法恢复桌面授权页，请重新执行步骤 7。',
             });
+          }
 
-          const pageState = await getDesktopAuthorizePageState(tabId, {
+          let pageState = null;
+          try {
+            pageState = await getDesktopAuthorizePageState(tabId, {
               step: 8,
+              timeoutBudget,
               injectLogMessage: '步骤 8：Kiro 桌面授权页内容脚本未就绪，正在等待页面恢复...',
               readyLogMessage: '步骤 8：正在读取 Kiro 桌面授权页当前状态...',
             });
+          } catch (error) {
+            if (awaitingCallbackAfterConsent && isMissingTabError(error)) {
+              await sleepWithStop(1000);
+              continue;
+            }
+            throw error;
+          }
 
-          await setState(mergeRuntimePatch(currentState, {
+          await updateLoopState({
             session: {
               pageState: pageState?.state || '',
               pageUrl: pageState?.url || '',
               lastError: '',
             },
-          }));
+          });
 
           if (pageState.state === 'relogin_email') {
             const email = cleanString(runtimeState.register?.email || currentState?.email);
             await log(`步骤 8：桌面授权页要求重新输入邮箱，正在填写 ${email}...`, 'info', nodeId);
             await executeDesktopAction(tabId, 'submit-email', { email }, {
               step: 8,
+              timeoutBudget,
               logMessage: '步骤 8：正在向桌面授权页提交邮箱...',
             });
             await sleepWithStop(1200);
@@ -915,6 +999,7 @@
             await log('步骤 8：桌面授权页要求重新输入密码，正在填写密码...', 'info', nodeId);
             await executeDesktopAction(tabId, 'submit-password', { password }, {
               step: 8,
+              timeoutBudget,
               logMessage: '步骤 8：正在向桌面授权页提交密码...',
             });
             await sleepWithStop(1200);
@@ -922,14 +1007,13 @@
           }
 
           if (pageState.state === 'otp_page') {
-            runtimeState = readKiroRuntime(currentState);
             if (!runtimeState.desktopAuth?.otpRequestedAt) {
-              await setState(mergeRuntimePatch(currentState, {
+              await updateLoopState({
                 desktopAuth: {
                   otpRequestedAt: Date.now(),
                   status: 'waiting_otp',
                 },
-              }));
+              });
             }
             const codeResult = await pollDesktopOtpCode(8, currentState, nodeId);
             const code = cleanString(codeResult?.code);
@@ -939,6 +1023,7 @@
             await log(`步骤 8：已获取桌面授权验证码 ${code}，正在提交...`, 'info', nodeId);
             await executeDesktopAction(tabId, 'submit-otp', { code }, {
               step: 8,
+              timeoutBudget,
               logMessage: '步骤 8：正在向桌面授权页提交验证码...',
             });
             await sleepWithStop(1200);
@@ -949,8 +1034,10 @@
             await log('步骤 8：正在确认 Kiro 桌面授权访问...', 'info', nodeId);
             await executeDesktopAction(tabId, 'confirm-consent', {}, {
               step: 8,
+              timeoutBudget,
               logMessage: '步骤 8：正在确认桌面授权访问...',
             });
+            awaitingCallbackAfterConsent = true;
             await sleepWithStop(1200);
             continue;
           }
@@ -958,59 +1045,20 @@
           if (pageState.state === 'callback_page') {
             const parsedCallback = parseDesktopCallbackUrl(pageState.url, desktopState, runtimeState.desktopAuth?.redirectPort);
             if (parsedCallback) {
-              callbackTracker.registerPending({
+              await finalizeDesktopAuthorizeCallback(currentState, runtimeState, parsedCallback, nodeId);
-                expectedState: desktopState,
+              return;
-                redirectPort: runtimeState.desktopAuth?.redirectPort,
-                tabId,
-              });
-              if (parsedCallback.code) {
-                await setState(mergeRuntimePatch(currentState, {
-                  session: {
-                    pageState: 'callback_page',
-                    pageUrl: parsedCallback.url,
-                  },
-                }));
-              }
             }
           }
 
           await sleepWithStop(1000);
         }
 
-        const lastResult = await callbackTracker.waitForResolved(desktopState, 2000).catch(() => null);
+        const lastResult = await callbackTracker.waitForResolved(
-        if (lastResult?.error) {
+          desktopState,
-          throw new Error(`桌面授权回调失败：${lastResult.error}`);
+          Math.min(timeoutBudget.getRemainingMs(1000), 2000)
-        }
+        ).catch(() => null);
-        if (lastResult?.code) {
+        if (lastResult) {
-          const tokenResult = await desktopClientApi.exchangeDesktopAuthorizationCode({
+          await finalizeDesktopAuthorizeCallback(currentState, runtimeState, lastResult, nodeId);
-            region: runtimeState.desktopAuth.region || DEFAULT_REGION,
-            clientId: runtimeState.desktopAuth.clientId,
-            clientSecret: runtimeState.desktopAuth.clientSecret,
-            redirectUri: runtimeState.desktopAuth.redirectUri,
-            code: lastResult.code,
-            codeVerifier: runtimeState.desktopAuth.codeVerifier,
-          }, fetchImpl);
-          const payload = await applyRuntimeState(currentState, {
-            session: {
-              currentStage: 'upload',
-              pageState: 'callback_captured',
-              pageUrl: lastResult.url,
-              lastError: '',
-            },
-            desktopAuth: {
-              authorizationCode: lastResult.code,
-              accessToken: tokenResult.accessToken,
-              refreshToken: tokenResult.refreshToken,
-              status: 'authorized',
-              authorizedAt: Date.now(),
-            },
-            upload: {
-              status: 'ready_to_upload',
-              error: '',
-            },
-          });
-          await log('步骤 8：桌面授权回调已捕获，Token 换取成功。', 'ok', nodeId);
-          await completeNodeFromBackground(nodeId, payload);
           return;
         }
 

tests/background-kiro-desktop-authorize-runner-module.test.js

@@ -11,6 +11,31 @@ function loadDesktopAuthorizeRunnerApi() {
   return globalScope.MultiPageBackgroundKiroDesktopAuthorizeRunner;
 }
 
+function createDesktopAuthorizeState(overrides = {}) {
+  return {
+    kiroRuntime: {
+      session: {
+        desktopTabId: 91,
+      },
+      register: {
+        email: 'kiro-user@example.com',
+      },
+      desktopAuth: {
+        region: 'us-east-1',
+        clientId: 'desktop-client-id',
+        clientSecret: 'desktop-client-secret',
+        state: 'desktop-state-001',
+        codeVerifier: 'desktop-code-verifier',
+        redirectUri: 'http://127.0.0.1:43121/oauth/callback',
+        redirectPort: 43121,
+        authorizeUrl: 'https://example.com/authorize',
+      },
+      upload: {},
+    },
+    ...overrides,
+  };
+}
+
 test('kiro desktop authorize runner exposes a factory and callback parser', () => {
   const api = loadDesktopAuthorizeRunnerApi();
   assert.equal(typeof api?.createKiroDesktopAuthorizeRunner, 'function');
@@ -54,4 +79,183 @@ test('kiro desktop authorize runner uses a shared 3-minute page-load timeout bud
   assert.match(source, /resolveTimeoutBudget/);
   assert.match(source, /timeoutBudget\.getRemainingMs\(1000\)/);
   assert.match(source, /onRetryableError: buildDesktopRetryRecovery\(tabId, \{\s*\.\.\.options,\s*timeoutBudget,/);
+  assert.match(source, /awaitingCallbackAfterConsent/);
+  assert.match(source, /finalizeDesktopAuthorizeCallback/);
+});
+
+test('executeKiroCompleteDesktopAuthorize finishes from callback page without waiting for tracker replay', async () => {
+  const api = loadDesktopAuthorizeRunnerApi();
+  let currentState = createDesktopAuthorizeState();
+  let completedPayload = null;
+
+  const runner = api.createKiroDesktopAuthorizeRunner({
+    addLog: async () => {},
+    chrome: {
+      tabs: {
+        get: async (tabId) => ({
+          id: tabId,
+          status: 'complete',
+          url: 'http://127.0.0.1:43121/oauth/callback?code=auth-code-001&state=desktop-state-001',
+        }),
+        remove: async () => {},
+        update: async () => ({}),
+      },
+      webNavigation: {
+        onBeforeNavigate: { addListener: () => {} },
+        onCommitted: { addListener: () => {} },
+      },
+      webRequest: {
+        onBeforeRequest: { addListener: () => {} },
+      },
+    },
+    completeNodeFromBackground: async (_nodeId, payload) => {
+      completedPayload = payload;
+    },
+    ensureContentScriptReadyOnTab: async () => {},
+    fetchImpl: async () => ({
+      ok: true,
+      text: async () => JSON.stringify({
+        accessToken: 'access-token-001',
+        refreshToken: 'refresh-token-001',
+      }),
+    }),
+    getState: async () => currentState,
+    getTabId: async () => 91,
+    isTabAlive: async () => true,
+    registerTab: async () => {},
+    sendToContentScriptResilient: async (_source, message) => {
+      if (message.type === 'GET_KIRO_DESKTOP_AUTHORIZE_STATE') {
+        return {
+          state: 'callback_page',
+          url: 'http://127.0.0.1:43121/oauth/callback?code=auth-code-001&state=desktop-state-001',
+        };
+      }
+      throw new Error(`Unexpected message: ${message.type}`);
+    },
+    setState: async (patch) => {
+      currentState = { ...currentState, ...patch };
+    },
+    sleepWithStop: async () => {},
+    throwIfStopped: () => {},
+    waitForTabStableComplete: async () => ({ status: 'complete' }),
+  });
+
+  await runner.executeKiroCompleteDesktopAuthorize(currentState);
+
+  assert.equal(completedPayload?.kiroRuntime?.desktopAuth?.authorizationCode, 'auth-code-001');
+  assert.equal(completedPayload?.kiroRuntime?.desktopAuth?.refreshToken, 'refresh-token-001');
+});
+
+test('executeKiroCompleteDesktopAuthorize waits for callback after consent even if original tab disappears', async () => {
+  const api = loadDesktopAuthorizeRunnerApi();
+  let currentState = createDesktopAuthorizeState({
+    kiroRuntime: {
+      session: {
+        desktopTabId: 91,
+      },
+      register: {
+        email: 'kiro-user@example.com',
+      },
+      desktopAuth: {
+        region: 'us-east-1',
+        clientId: 'desktop-client-id',
+        clientSecret: 'desktop-client-secret',
+        state: 'desktop-state-002',
+        codeVerifier: 'desktop-code-verifier',
+        redirectUri: 'http://127.0.0.1:43121/oauth/callback',
+        redirectPort: 43121,
+        authorizeUrl: 'https://example.com/authorize',
+      },
+      upload: {},
+    },
+  });
+  let completedPayload = null;
+  let tabAlive = true;
+  let navigationListener = null;
+
+  const runner = api.createKiroDesktopAuthorizeRunner({
+    addLog: async () => {},
+    chrome: {
+      tabs: {
+        get: async (tabId) => {
+          if (!tabAlive) {
+            throw new Error(`No tab with id: ${tabId}`);
+          }
+          return {
+            id: tabId,
+            status: 'complete',
+            url: 'https://example.com/consent',
+          };
+        },
+        remove: async () => {},
+        update: async (tabId) => {
+          if (!tabAlive) {
+            throw new Error(`No tab with id: ${tabId}`);
+          }
+          return { id: tabId };
+        },
+      },
+      webNavigation: {
+        onBeforeNavigate: { addListener: () => {} },
+        onCommitted: {
+          addListener: (listener) => {
+            navigationListener = listener;
+          },
+        },
+      },
+      webRequest: {
+        onBeforeRequest: { addListener: () => {} },
+      },
+    },
+    completeNodeFromBackground: async (_nodeId, payload) => {
+      completedPayload = payload;
+    },
+    ensureContentScriptReadyOnTab: async () => {},
+    fetchImpl: async () => ({
+      ok: true,
+      text: async () => JSON.stringify({
+        accessToken: 'access-token-002',
+        refreshToken: 'refresh-token-002',
+      }),
+    }),
+    getState: async () => currentState,
+    getTabId: async () => 91,
+    isTabAlive: async () => true,
+    registerTab: async () => {},
+    sendToContentScriptResilient: async (_source, message) => {
+      if (message.type === 'GET_KIRO_DESKTOP_AUTHORIZE_STATE') {
+        return {
+          state: 'consent_page',
+          url: 'https://example.com/consent',
+        };
+      }
+      if (message.type === 'EXECUTE_KIRO_DESKTOP_AUTHORIZE_ACTION') {
+        assert.equal(message.payload?.action, 'confirm-consent');
+        tabAlive = false;
+        setTimeout(() => {
+          navigationListener?.({
+            url: 'http://127.0.0.1:43121/oauth/callback?code=auth-code-002&state=desktop-state-002',
+            tabId: 91,
+          });
+        }, 50);
+        return {
+          submitted: true,
+          state: 'consent_submitted',
+          url: 'https://example.com/consent',
+        };
+      }
+      throw new Error(`Unexpected message: ${message.type}`);
+    },
+    setState: async (patch) => {
+      currentState = { ...currentState, ...patch };
+    },
+    sleepWithStop: async () => {},
+    throwIfStopped: () => {},
+    waitForTabStableComplete: async () => ({ status: 'complete' }),
+  });
+
+  await runner.executeKiroCompleteDesktopAuthorize(currentState);
+
+  assert.equal(completedPayload?.kiroRuntime?.desktopAuth?.authorizationCode, 'auth-code-002');
+  assert.equal(completedPayload?.kiroRuntime?.desktopAuth?.refreshToken, 'refresh-token-002');
 });