From 38a68152acfce16a385194eb0d0622efe3fd3c09 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Fri, 31 Jul 2026 11:22:30 -0700
Subject: [PATCH] fix(ui): clear the shared ID-JAG leg-1 columns on an auth
 type switch

ID-JAG and OAuth token exchange both persist leg 1 in token_exchange_endpoint,
audience and subject_token_type, but the edit form treated all three as owned
by token exchange alone. Converting an ID-JAG server to another auth type left
a stale endpoint and audience behind, and switching token exchange to ID-JAG
wiped the values the new mode was about to reuse, which the save path then
rejected as a half-configured server.

The three shared columns now clear when a server leaves the pair rather than
when it leaves one member of it, and the token-exchange-only clearing narrows
to token_exchange_profile, the one field that mode alone owns
---
 .../_components/mcp_server_edit.test.tsx      | 80 +++++++++++++++++++
 .../_components/mcp_server_edit.tsx           |  7 +-
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
index b660c4bdb76..8a23c445173 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
@@ -432,6 +432,86 @@ describe("MCPServerEdit (auth type switch)", () => {
     expect(payload.registration_url).toBeNull();
   });
 
+  it("preserves the shared leg-1 endpoint and audience when switching token exchange to ID-JAG", async () => {
+    vi.mocked(networking.updateMCPServer).mockResolvedValue({
+      ...interactiveOAuthServer,
+      auth_type: "oauth2_id_jag",
+    });
+
+    render(
+      <MCPServerEdit
+        mcpServer={{
+          ...interactiveOAuthServer,
+          auth_type: "oauth2_token_exchange",
+          token_exchange_endpoint: "https://idp.example.com/oauth2/token",
+          audience: "api://existing-resource",
+          subject_token_type: "urn:ietf:params:oauth:token-type:saml2",
+        }}
+        accessToken="access-token"
+        onCancel={vi.fn()}
+        onSuccess={vi.fn()}
+        availableAccessGroups={[]}
+      />,
+    );
+
+    await selectAntOption("Authentication", "ID-JAG (Okta Cross App Access)");
+
+    const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
+    await act(async () => {
+      fireEvent.click(saveButtons[0]);
+    });
+
+    await waitFor(() => {
+      expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
+    });
+
+    const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
+    expect(payload.auth_type).toBe("oauth2_id_jag");
+    expect(payload.token_exchange_endpoint).toBe("https://idp.example.com/oauth2/token");
+    expect(payload.audience).toBe("api://existing-resource");
+    expect(payload.subject_token_type).toBe("urn:ietf:params:oauth:token-type:saml2");
+  });
+
+  it("clears the shared leg-1 endpoint and audience when switching ID-JAG to an auth type that uses neither", async () => {
+    vi.mocked(networking.updateMCPServer).mockResolvedValue({
+      ...interactiveOAuthServer,
+      auth_type: "none",
+    });
+
+    render(
+      <MCPServerEdit
+        mcpServer={{
+          ...interactiveOAuthServer,
+          auth_type: "oauth2_id_jag",
+          token_exchange_endpoint: "https://idp.example.com/oauth2/token",
+          audience: "api://existing-resource",
+          subject_token_type: "urn:ietf:params:oauth:token-type:saml2",
+        }}
+        accessToken="access-token"
+        onCancel={vi.fn()}
+        onSuccess={vi.fn()}
+        availableAccessGroups={[]}
+      />,
+    );
+
+    await selectAntOption("Authentication", "None");
+
+    const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
+    await act(async () => {
+      fireEvent.click(saveButtons[0]);
+    });
+
+    await waitFor(() => {
+      expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
+    });
+
+    const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
+    expect(payload.auth_type).toBe("none");
+    expect(payload.token_exchange_endpoint).toBeNull();
+    expect(payload.audience).toBeNull();
+    expect(payload.subject_token_type).toBeNull();
+  });
+
   it("keeps oauth2 endpoint overrides when the auth type is unchanged", async () => {
     vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer });
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
index acecec105eb..e9311fce67b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
@@ -70,6 +70,7 @@ const AUTH_TYPES_REQUIRING_CREDENTIALS = [
   AUTH_TYPE.TRUE_PASSTHROUGH,
   AUTH_TYPE.OAUTH_DELEGATE,
 ];
+const AUTH_TYPES_SHARING_TOKEN_EXCHANGE_COLUMNS = [AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE, AUTH_TYPE.OAUTH2_ID_JAG];
 export const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
 
 const MCPServerEdit: React.FC<MCPServerEditProps> = ({
@@ -858,9 +859,13 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
         ...(mcpServer.auth_type === AUTH_TYPE.OAUTH2 && restValues.auth_type !== AUTH_TYPE.OAUTH2
           ? { issuer: null, authorization_url: null, token_url: null, registration_url: null }
           : {}),
+        ...(AUTH_TYPES_SHARING_TOKEN_EXCHANGE_COLUMNS.includes(mcpServer.auth_type ?? "") &&
+        !AUTH_TYPES_SHARING_TOKEN_EXCHANGE_COLUMNS.includes(restValues.auth_type ?? "")
+          ? { token_exchange_endpoint: null, audience: null, subject_token_type: null }
+          : {}),
         ...(mcpServer.auth_type === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE &&
         restValues.auth_type !== AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE
-          ? { token_exchange_endpoint: null, audience: null, subject_token_type: null, token_exchange_profile: null }
+          ? { token_exchange_profile: null }
           : {}),
         server_id: mcpServer.server_id,
         mcp_info: {