fix(self-driving): stop a 401 stranding the GitHub gate

src/lib/agent/runner/shared/__tests__/refresh-access-token-if-needed.test.ts → src/lib/__tests__/session-token.test.ts

@@ -1,4 +1,4 @@
-import { refreshAccessTokenIfNeeded } from '../authenticate';
+import { refreshAccessTokenIfNeeded } from '@lib/session-token';
 import { refreshAccessToken } from '@utils/oauth';
 import { OAuthError } from '@utils/oauth-errors';
 import { isGrantRevoked, resetAuthSessionState } from '@lib/auth-session-state';
@@ -120,7 +120,7 @@ describe('refreshAccessTokenIfNeeded', () => {
     mockedRefresh.mockRejectedValueOnce(new Error('network down'));
     const session = sessionWith(aging());
 
-    await expect(refreshAccessTokenIfNeeded(session)).resolves.toBeUndefined();
+    await expect(refreshAccessTokenIfNeeded(session)).resolves.toBe(false);
     expect(session.credentials!.accessToken).toBe('pha_old');
     expect(setAccessToken).not.toHaveBeenCalled();
   });
@@ -140,4 +140,32 @@ describe('refreshAccessTokenIfNeeded', () => {
 
     expect(isGrantRevoked()).toBe(false);
   });
+
+  // The Self-driving GitHub gate calls this after a 401: the token was
+  // rejected, so its stated expiry proves nothing.
+  it('refreshes a token that still looks fresh when forced', async () => {
+    mockedRefresh.mockResolvedValueOnce({
+      access_token: 'pha_new',
+      expires_in: 3600,
+      token_type: 'Bearer',
+      scope: 'project:read',
+    });
+    const session = sessionWith(
+      aging({ expiresAt: Date.now() + 6 * 60 * 60 * 1000 }),
+    );
+
+    await expect(
+      refreshAccessTokenIfNeeded(session, { force: true }),
+    ).resolves.toBe(true);
+    expect(session.credentials!.accessToken).toBe('pha_new');
+  });
+
+  it('still needs a refresh token when forced', async () => {
+    await expect(
+      refreshAccessTokenIfNeeded(sessionWith({ accessToken: 'pha_ci_key' }), {
+        force: true,
+      }),
+    ).resolves.toBe(false);
+    expect(mockedRefresh).not.toHaveBeenCalled();
+  });
 });

src/lib/agent/runner/shared/authenticate.ts

@@ -10,12 +10,9 @@
  * back rather than fetching again.
  */
 
-import type { Credentials, WizardSession } from '@lib/wizard-session';
+import type { WizardSession } from '@lib/wizard-session';
 import type { ProgramId } from '@lib/programs/program-registry';
 import { getOrAskForProjectData } from '@utils/setup-utils';
-import { refreshAccessToken } from '@utils/oauth';
-import { OAuthError } from '@utils/oauth-errors';
-import { markGrantRevoked } from '@lib/auth-session-state';
 import { analytics, groupsFromUser } from '@utils/analytics';
 import { getUI } from '@ui';
 import { logToFile } from '@utils/debug';
@@ -74,57 +71,3 @@ export async function authenticate(
   if (user) analytics.identifyUser(user);
   analytics.setGroups(groupsFromUser(user, host.apiHost));
 }
-
-// Below this remaining lifetime a run risks outliving its token; just-minted and 7-day tokens skip.
-// Only a second agent run in one invocation can be this old — see self-driving's chained phases.
-const REFRESH_WHEN_REMAINING_MS = 50 * 60 * 1000;
-
-/**
- * Grants the token endpoint refuses permanently. A dead grant means the login
- * is gone, not that the network blipped, so only these mark the session.
- */
-const DEAD_GRANT_CODES = new Set(['invalid_grant', 'invalid_client']);
-
-// Best-effort pre-run refresh: no refresh token or a failed grant keeps the existing token.
-export async function refreshAccessTokenIfNeeded(
-  session: WizardSession,
-): Promise<void> {
-  const credentials = session.credentials;
-  if (!credentials?.refreshToken) return;
-
-  // No expiry means we cannot tell how much life is left, so leave it alone —
-  // refreshing every run would spend a rotation for nothing.
-  if (credentials.expiresAt === undefined) return;
-  if (credentials.expiresAt - Date.now() >= REFRESH_WHEN_REMAINING_MS) return;
-
-  try {
-    const token = await refreshAccessToken(
-      credentials.refreshToken,
-      session.baseUrl,
-      credentials.oauthClientId,
-    );
-    // Replaced, not mutated: readers hold this object, and a new one keeps the
-    // store and the (possibly shallow-copied) session explicitly in step.
-    const refreshed: Credentials = {
-      ...credentials,
-      accessToken: token.access_token,
-      // Rotation: keep the returned refresh token or the old one stops working.
-      refreshToken: token.refresh_token ?? credentials.refreshToken,
-      expiresAt: Date.now() + token.expires_in * 1000,
-    };
-    session.credentials = refreshed;
-    getUI().setAccessToken(refreshed);
-  } catch (error) {
-    // A dead grant is recorded but not thrown: the current token may still have
-    // minutes of life, and failing here would break runs that would have worked.
-    // If a 401 does follow, the auth-error screen can finally name the cause.
-    if (error instanceof OAuthError && DEAD_GRANT_CODES.has(error.code)) {
-      markGrantRevoked();
-      analytics.wizardCapture('auth session expired', { reason: error.code });
-    }
-    logToFile(
-      '[oauth] pre-run token refresh failed, continuing with the existing token:',
-      error instanceof Error ? error.message : error,
-    );
-  }
-}

src/lib/agent/runner/shared/bootstrap.ts

@@ -10,7 +10,8 @@
 import type { WizardSession } from '@lib/wizard-session';
 import { analytics } from '@utils/analytics';
 import { getUI } from '@ui';
-import { authenticate, refreshAccessTokenIfNeeded } from './authenticate';
+import { authenticate } from './authenticate';
+import { refreshAccessTokenIfNeeded } from '@lib/session-token';
 import { maybeStampAiSdkDetected } from '@lib/programs/posthog-integration/detect';
 import { createTriageLLMProvider } from '@lib/agent/triage-provider';
 import { gatewayAuth } from '@lib/gateway-session';

src/lib/api.ts

@@ -235,58 +235,60 @@ const IntegrationsResponseSchema = z.object({
 });
 
 /**
- * Check whether the project already has a Slack integration connected.
- * Requires the `integration:read` scope. Throws on failure — callers
- * (including the SlackConnectScreen poll) decide how to degrade and
- * are responsible for capturing the error exactly once.
+ * Check whether the project already has an integration of `kind` connected.
+ * Requires the `integration:read` scope. Throws an `ApiError` on failure —
+ * the connect-screen polls decide how to degrade and are responsible for
+ * capturing the error exactly once. The status matters to them: the
+ * Self-driving GitHub gate reads a 401 as an expired login, not as an
+ * integration that is missing.
  */
-export async function fetchSlackConnected(
+async function fetchIntegrationConnected(
+  kind: string,
   accessToken: string,
   projectId: number,
   baseUrl: string,
   signal?: AbortSignal,
 ): Promise<boolean> {
-  const response = await axios.get(
-    `${baseUrl}/api/projects/${projectId}/integrations/`,
-    {
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        'User-Agent': WIZARD_USER_AGENT,
+  try {
+    const response = await axios.get(
+      `${baseUrl}/api/projects/${projectId}/integrations/`,
+      {
+        headers: {
+          Authorization: `Bearer ${accessToken}`,
+          'User-Agent': WIZARD_USER_AGENT,
+        },
+        signal,
+        // Bound the request so one stalled socket can't wedge the poll. The
+        // Self-driving GitHub gate can't be skipped, so a hang here would
+        // strand the user the same way a rejected token does. A timeout raises
+        // ECONNABORTED with no response, which the poll reads as a non-auth
+        // blip and retries.
+        timeout: 10_000,
       },
-      signal,
-    },
-  );
-  const parsed = IntegrationsResponseSchema.safeParse(response.data);
-  if (!parsed.success) return false;
-  return parsed.data.results.some((i) => i.kind === 'slack');
+    );
+    const parsed = IntegrationsResponseSchema.safeParse(response.data);
+    if (!parsed.success) return false;
+    return parsed.data.results.some((i) => i.kind === kind);
+  } catch (error) {
+    throw handleApiError(error, `check the ${kind} connection`);
+  }
 }
 
-/**
- * Check whether the project already has a GitHub App integration connected.
- * Requires the `integration:read` scope. Throws on failure — callers (the
- * SelfDrivingGitHubScreen poll) decide how to degrade and are responsible for
- * capturing the error exactly once.
- */
-export async function fetchGithubConnected(
+export const fetchSlackConnected = (
   accessToken: string,
   projectId: number,
   baseUrl: string,
   signal?: AbortSignal,
-): Promise<boolean> {
-  const response = await axios.get(
-    `${baseUrl}/api/projects/${projectId}/integrations/`,
-    {
-      headers: {
-        Authorization: `Bearer ${accessToken}`,
-        'User-Agent': WIZARD_USER_AGENT,
-      },
-      signal,
-    },
-  );
-  const parsed = IntegrationsResponseSchema.safeParse(response.data);
-  if (!parsed.success) return false;
-  return parsed.data.results.some((i) => i.kind === 'github');
-}
+): Promise<boolean> =>
+  fetchIntegrationConnected('slack', accessToken, projectId, baseUrl, signal);
+
+export const fetchGithubConnected = (
+  accessToken: string,
+  projectId: number,
+  baseUrl: string,
+  signal?: AbortSignal,
+): Promise<boolean> =>
+  fetchIntegrationConnected('github', accessToken, projectId, baseUrl, signal);
 
 export function handleApiError(error: unknown, operation: string): ApiError {
   if (axios.isAxiosError(error)) {

src/lib/session-token.ts

@@ -0,0 +1,84 @@
+/**
+ * The run's access token, kept fresh.
+ *
+ * A leaf module on purpose. Two callers far apart need it: the agent bootstrap
+ * refreshes right before it mints a gateway token, and the Self-driving GitHub
+ * gate refreshes before it polls the PostHog API — a screen that runs *before*
+ * the bootstrap and would otherwise poll with whatever token the earlier
+ * integration phase left behind.
+ */
+
+import type { Credentials, WizardSession } from '@lib/wizard-session';
+import { refreshAccessToken } from '@utils/oauth';
+import { OAuthError } from '@utils/oauth-errors';
+import { markGrantRevoked } from '@lib/auth-session-state';
+import { analytics } from '@utils/analytics';
+import { getUI } from '@ui';
+import { logToFile } from '@utils/debug';
+
+// Below this remaining lifetime a run risks outliving its token; just-minted and 7-day tokens skip.
+// Only a second agent run in one invocation can be this old — see self-driving's chained phases.
+const REFRESH_WHEN_REMAINING_MS = 50 * 60 * 1000;
+
+/**
+ * Grants the token endpoint refuses permanently. A dead grant means the login
+ * is gone, not that the network blipped, so only these mark the session.
+ */
+const DEAD_GRANT_CODES = new Set(['invalid_grant', 'invalid_client']);
+
+/**
+ * Best-effort refresh: no refresh token or a failed grant keeps the existing
+ * token. Returns true only when the token was replaced.
+ *
+ * `force` skips the remaining-lifetime check. Pass it when the server already
+ * rejected the token, because then the expiry says nothing useful — a token can
+ * be revoked long before it runs out.
+ */
+export async function refreshAccessTokenIfNeeded(
+  session: WizardSession,
+  options: { force?: boolean } = {},
+): Promise<boolean> {
+  const credentials = session.credentials;
+  if (!credentials?.refreshToken) return false;
+
+  if (!options.force) {
+    // No expiry means we cannot tell how much life is left, so leave it alone —
+    // refreshing every run would spend a rotation for nothing.
+    if (credentials.expiresAt === undefined) return false;
+    if (credentials.expiresAt - Date.now() >= REFRESH_WHEN_REMAINING_MS)
+      return false;
+  }
+
+  try {
+    const token = await refreshAccessToken(
+      credentials.refreshToken,
+      session.baseUrl,
+      credentials.oauthClientId,
+    );
+    // Replaced, not mutated: readers hold this object, and a new one keeps the
+    // store and the (possibly shallow-copied) session explicitly in step.
+    const refreshed: Credentials = {
+      ...credentials,
+      accessToken: token.access_token,
+      // Rotation: keep the returned refresh token or the old one stops working.
+      refreshToken: token.refresh_token ?? credentials.refreshToken,
+      expiresAt: Date.now() + token.expires_in * 1000,
+    };
+    session.credentials = refreshed;
+    getUI().setAccessToken(refreshed);
+    return true;
+  } catch (error) {
+    // A dead grant is recorded but not thrown: the current token may still have
+    // minutes of life, and failing here would break runs that would have worked.
+    // If a 401 does follow, the auth-error screen can finally name the cause.
+    if (error instanceof OAuthError && DEAD_GRANT_CODES.has(error.code)) {
+      markGrantRevoked();
+      analytics.wizardCapture('auth session expired', { reason: error.code });
+    }
+    logToFile(
+      '[oauth] token refresh failed, continuing with the existing token:',
+      error instanceof Error ? error.message : error,
+    );
+    return false;
+  }
+}