From 7ed56925975102377bbfc5817380c03776b8dc75 Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Thu, 19 Mar 2026 20:42:49 -0700
Subject: [PATCH] feat: add exchangeLoginCode for SSO code exchange, update
 LoginPage

- Add exchangeLoginCode() to networking.tsx for SSO callback
- LoginPage detects ?code= from SSO redirect, exchanges for JWT
- Keep backwards compat for legacy ?token= flow
---
 .../src/app/login/LoginPage.tsx               | 19 ++++++++++++---
 .../src/components/networking.tsx             | 24 +++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx
index dc9aae5d532..5a9d420456c 100644
--- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx
+++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx
@@ -3,7 +3,7 @@
 import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin";
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
-import { getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
+import { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
 import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
 import { isJwtExpired } from "@/utils/jwtUtils";
 import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
@@ -43,12 +43,25 @@ function LoginPageContent() {
       return;
     }
 
-    // Cross-origin SSO: worker redirected back with token in URL
+    // Cross-origin SSO: worker redirected back with a single-use code.
+    // Exchange it for the JWT via the worker's /v3/login/exchange endpoint.
     const params = new URLSearchParams(window.location.search);
+    const ssoCode = params.get("code");
+    if (ssoCode) {
+      const workerUrl = localStorage.getItem("litellm_worker_url");
+      exchangeLoginCode(ssoCode, workerUrl).then(() => {
+        params.delete("code");
+        const cleanSearch = params.toString();
+        window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""));
+        router.replace("/ui/?login=success");
+      });
+      return;
+    }
+
+    // Backwards compat: handle direct token in URL (legacy flow)
     const urlToken = params.get("token");
     if (urlToken && !isJwtExpired(urlToken)) {
       document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
-      // Strip token from URL to keep it out of browser history
       params.delete("token");
       const cleanSearch = params.toString();
       window.history.replaceState(
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index dbae46389b0..548a695d5f4 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -9100,6 +9100,30 @@ export const loginCall = async (username: string, password: string, useV3?: bool
   return data;
 };
 
+/**
+ * Exchange a single-use login code for a JWT token.
+ * Used by the SSO callback when the worker redirects back with ?code=.
+ */
+export const exchangeLoginCode = async (code: string, workerBaseUrl?: string | null): Promise<string> => {
+  const base = workerBaseUrl || getProxyBaseUrl();
+  const response = await fetch(`${base}/v3/login/exchange`, {
+    method: "POST",
+    body: JSON.stringify({ code }),
+    headers: { "Content-Type": "application/json" },
+  });
+
+  if (!response.ok) {
+    const errorData = await response.json();
+    throw new Error(deriveErrorMessage(errorData));
+  }
+
+  const data = await response.json();
+  if (data.token) {
+    document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
+  }
+  return data.token;
+};
+
 export const getUiSettings = async () => {
   const proxyBaseUrl = getProxyBaseUrl();
   const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_settings` : `/get/ui_settings`;