fix: user setup uses org token on provision (#495)

Author: v3nant

src/api/ci-token.client.ts

@@ -17,6 +17,7 @@ export type IamAccessOrgTokensInput =
 
 export interface ProvisionCITokenResponse {
   refresh_token: string;
+  access_token: string;
 }
 
 export interface ProvisionCITokenOptions {
@@ -139,5 +140,5 @@ export async function provisionCIToken(options: ProvisionCITokenOptions = {}): P
     throw new Error('Either orgId or previousToken is required to provision a CI token');
   }
   const result = await getOrgAccessTokens(input);
-  return { refresh_token: result.refreshToken };
+  return { access_token: result.accessToken, refresh_token: result.refreshToken };
 }

src/api/user-setup.client.ts

@@ -35,11 +35,11 @@ function extractErrorCode(errors: ReadonlyArray<GraphQLFormattedError>): ApiErro
   return code;
 }
 
-export async function getUserSetupStatus(options?: { preferOAuth?: boolean }): Promise<{
+export async function getUserSetupStatus(options?: { preferOAuth?: boolean; orgAccessToken?: string }): Promise<{
   isComplete: boolean;
   orgId?: number | null;
 }> {
-  const tokenProvider = getTokenProvider(options?.preferOAuth);
+  const tokenProvider = getTokenProvider(options?.preferOAuth, options?.orgAccessToken);
   const client = createApollo(getGraphqlUrl(), tokenProvider);
   const res = await client.query<UserSetupStatusResponse>({ query: userSetupStatusQuery });
 
@@ -70,11 +70,11 @@ export async function getUserSetupStatus(options?: { preferOAuth?: boolean }): P
   return { isComplete: status.isComplete, orgId: status.orgId ?? undefined };
 }
 
-export async function completeUserSetup(options?: { preferOAuth?: boolean }): Promise<{
+export async function completeUserSetup(options?: { preferOAuth?: boolean; orgAccessToken?: string }): Promise<{
   isComplete: boolean;
   orgId?: number | null;
 }> {
-  const tokenProvider = getTokenProvider(options?.preferOAuth);
+  const tokenProvider = getTokenProvider(options?.preferOAuth, options?.orgAccessToken);
   const client = createApollo(getGraphqlUrl(), tokenProvider);
   const res = await client.mutate<CompleteUserSetupResponse>({ mutation: completeUserSetupMutation });
 
@@ -105,7 +105,7 @@ export async function completeUserSetup(options?: { preferOAuth?: boolean }): Pr
   return { isComplete: true, orgId: result.orgId ?? undefined };
 }
 
-export async function ensureUserSetup(options?: { preferOAuth?: boolean }): Promise<number> {
+export async function ensureUserSetup(options?: { preferOAuth?: boolean; orgAccessToken?: string }): Promise<number> {
   const status = await withRetries('user-setup-status', () => getUserSetupStatus(options), {
     attempts: USER_SETUP_MAX_ATTEMPTS,
     baseDelayMs: USER_SETUP_RETRY_DELAY_MS,

src/commands/auth/provision-ci-token.ts

@@ -45,6 +45,15 @@ export default class AuthProvisionCiToken extends Command {
 
     try {
       const result = await provisionCIToken({ orgId });
+      try {
+        await ensureUserSetup({ orgAccessToken: result.access_token });
+      } catch (error) {
+        track('CLI CI Token Provision Failed', () => ({
+          command: 'auth provision-ci-token',
+          error: `user_setup_failed:${getErrorMessage(error)}`,
+        }));
+        this.error(`User Org setup failed. ${getErrorMessage(error)}`);
+      }
       const refreshToken = result.refresh_token;
       saveCIToken(refreshToken);
       this.log('CI token provisioned and saved locally.');

src/service/auth.svc.ts

@@ -24,7 +24,11 @@ export const AUTH_ERROR_MESSAGES = {
 
 export async function getTokenForScanWithSource(
   preferOAuth?: boolean,
+  orgAccessToken?: string,
 ): Promise<{ token: string; source: TokenSource }> {
+  if (orgAccessToken) {
+    return { token: orgAccessToken, source: 'ci' };
+  }
   if (preferOAuth) {
     const token = await requireAccessToken();
     return { token, source: 'oauth' };
@@ -94,9 +98,9 @@ export async function getAccessToken(): Promise<string | undefined> {
   return refreshed.access_token;
 }
 
-export function getTokenProvider(preferOAuth?: boolean): TokenProvider {
+export function getTokenProvider(preferOAuth?: boolean, orgAccessToken?: string): TokenProvider {
   return async (_forceRefresh?: boolean): Promise<string> => {
-    const { token } = await getTokenForScanWithSource(preferOAuth);
+    const { token } = await getTokenForScanWithSource(preferOAuth, orgAccessToken);
     return token;
   };
 }

test/api/ci-token.client.test.ts

@@ -107,7 +107,7 @@ describe('ci-token.client', () => {
     );
 
     const result = await provisionCIToken({ orgId: 42 });
-    expect(result).toEqual({ refresh_token: 'ci-refresh-token-123' });
+    expect(result).toEqual({ refresh_token: 'ci-refresh-token-123', access_token: 'new-access' });
 
     const calls = fetchMock.getCalls();
     expect(calls).toHaveLength(1);