fix(grok-cli): accept full auth.json object in import-token endpoint

src/lib/oauth/providers/grok-cli.ts

@@ -57,35 +57,61 @@ function parseJwtPayload(token: string): {
  *   - Raw JWT string (no refresh_token available)
  *   - The entire auth.json object: { "https://auth.x.ai::...": { "key": "eyJ...", "refresh_token": "..." } }
  */
-function extractTokenAndRefresh(input: any): { accessToken: string; refreshToken: string | null } {
-  if (typeof input === "string") return { accessToken: input, refreshToken: null };
+function extractTokenAndRefresh(input: unknown): {
+  accessToken: string;
+  refreshToken: string | null;
+  rawAuthJson: Record<string, unknown> | null;
+} {
+  // Direct JWT string
+  if (typeof input === "string")
+    return { accessToken: input, refreshToken: null, rawAuthJson: null };
 
   if (input && typeof input === "object") {
-    // auth.json format: first entry's "key" and "refresh_token" fields
-    const keys = Object.keys(input);
-    if (keys.length > 0 && input[keys[0]]?.key) {
-      return {
-        accessToken: input[keys[0]].key,
-        refreshToken: input[keys[0]].refresh_token || null,
-      };
+    const obj = input as Record<string, unknown>;
+
+    // The route handler wraps the token: { accessToken: <token> }.
+    // Unwrap once before checking the inner value.
+    const inner =
+      typeof obj.accessToken === "object" && obj.accessToken !== null
+        ? (obj.accessToken as Record<string, unknown>)
+        : obj;
+
+    // auth.json format: { "https://auth.x.ai::...": { key: "eyJ...", refresh_token: "..." } }
+    if (inner && typeof inner === "object") {
+      const innerKeys = Object.keys(inner);
+      for (const k of innerKeys) {
+        const entry = inner[k];
+        if (entry && typeof entry === "object" && "key" in entry) {
+          const e = entry as Record<string, unknown>;
+          if (typeof e.key === "string" && e.key.startsWith("eyJ")) {
+            return {
+              accessToken: e.key,
+              refreshToken: typeof e.refresh_token === "string" ? e.refresh_token : null,
+              rawAuthJson: inner as Record<string, unknown>,
+            };
+          }
+        }
+      }
     }
-    // Already has accessToken
-    if (input.accessToken) {
+
+    // Raw JWT passed as { accessToken: "eyJ..." }
+    if (typeof obj.accessToken === "string" && obj.accessToken.length > 0) {
       return {
-        accessToken: input.accessToken,
-        refreshToken: input.refreshToken || null,
+        accessToken: obj.accessToken,
+        refreshToken: typeof obj.refreshToken === "string" ? obj.refreshToken : null,
+        rawAuthJson: null,
       };
     }
   }
 
-  return { accessToken: "", refreshToken: null };
+  return { accessToken: "", refreshToken: null, rawAuthJson: null };
 }
 
 export const grokCli = {
   config: GROK_CLI_CONFIG,
   flowType: "import_token",
-  mapTokens: (token: any) => {
-    const { accessToken, refreshToken } = extractTokenAndRefresh(token);
+  mapTokens: (token: unknown) => {
+    const { accessToken, refreshToken, rawAuthJson } = extractTokenAndRefresh(token);
     const { email, authInfo } = parseJwtPayload(accessToken);
 
     return {
@@ -98,6 +124,7 @@ export const grokCli = {
         teamId: authInfo?.team_id || null,
         tier: authInfo?.tier || 1,
         principalType: authInfo?.principal_type || "User",
+        rawAuthJson: rawAuthJson || undefined,
       },
     };
   },

src/shared/validation/schemas/auth.ts

@@ -133,7 +133,7 @@ export const oauthPollSchema = z.object({
 
 /** Import a raw API token (e.g. WINDSURF_API_KEY) without going through the browser OAuth flow. */
 export const oauthImportTokenSchema = z.object({
-  token: z.string().trim().min(1, "Token is required"),
+  token: z.union([z.string().trim().min(1, "Token is required"), z.record(z.string(), z.unknown())]),
   connectionId: z.string().optional(),
 });
 

tests/unit/grok-cli-oauth.test.ts

@@ -65,3 +65,55 @@ test("Grok Build OAuth Provider - mapTokens from object with accessToken", () =>
   const result = grokCli.mapTokens(input, null);
   assert.equal(result.accessToken, "direct-token");
 });
+
+test("Grok Build OAuth Provider - mapTokens from route-wrapped auth.json", () => {
+  // The route handler wraps the token: { accessToken: <token> }.
+  // This simulates what the import-token endpoint passes to mapTokens.
+  const authJson = {
+    "https://auth.x.ai::b1a00492-073a-47ea-816f-4c329264a828": {
+      key: "eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.eyJlbWFpbCI6InRlc3RAZXhhbXBsZS5jb20ifQ.signature",
+      refresh_token: "test-refresh-token-wrapped",
+      expires_at: "2026-12-31T00:00:00Z",
+    },
+  };
+  const wrapped = { accessToken: authJson };
+  const result = grokCli.mapTokens(wrapped, null);
+
+  assert.ok(
+    result.accessToken.startsWith("eyJ"),
+    "accessToken should be JWT from wrapped auth.json"
+  );
+  assert.equal(result.refreshToken, "test-refresh-token-wrapped");
+  assert.equal(result.email, "test@example.com");
+  assert.ok(result.providerSpecificData?.rawAuthJson, "rawAuthJson should be populated");
+  assert.deepEqual(
+    result.providerSpecificData?.rawAuthJson,
+    authJson,
+    "rawAuthJson should equal the original auth.json"
+  );
+});
+
+test("Grok Build OAuth Provider - mapTokens from direct auth.json has rawAuthJson", () => {
+  const authJson = {
+    "https://auth.x.ai::clientId": {
+      key: "eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.eyJlbWFpbCI6InRlc3RAZXhhbXBsZS5jb20ifQ.signature",
+      refresh_token: "direct-refresh",
+    },
+  };
+  const result = grokCli.mapTokens(authJson, null);
+
+  assert.ok(result.accessToken.startsWith("eyJ"));
+  assert.equal(result.refreshToken, "direct-refresh");
+  assert.deepEqual(result.providerSpecificData?.rawAuthJson, authJson);
+});
+
+test("Grok Build OAuth Provider - mapTokens from raw JWT has no rawAuthJson", () => {
+  const payload = { sub: "12345", email: "test@example.com" };
+  const payloadBase64 = Buffer.from(JSON.stringify(payload)).toString("base64url");
+  const mockJwt = `eyJhbGciOiJFUzI1NiJ9.${payloadBase64}.signature`;
+  const result = grokCli.mapTokens(mockJwt, null);
+
+  assert.equal(result.accessToken, mockJwt);
+  assert.equal(result.refreshToken, null);
+  assert.equal(result.providerSpecificData?.rawAuthJson, undefined);
+});