From efd562f0d50a095617025bc77f250cd8b3187bfb Mon Sep 17 00:00:00 2001 From: ABE HSU Date: Sat, 25 Jul 2026 20:37:53 +0000 Subject: [PATCH] fix: preserve chat state across MCP tool OAuth redirect Connecting an MCP tool that requires OAuth does a full-page redirect, and the OAuth callback always lands the user back on the site root. Only the pending tool id was restored, so the active chat url, selected model, draft prompt, other enabled tools, selected skills/filters, and the web-search / image-generation / code-interpreter toggles were all lost. initiateOAuthRedirect() now snapshots the chat config to sessionStorage before redirecting, and new-chat initialization restores it and navigates back to the originating route. The return path is validated to be a same-origin relative path to avoid open redirects. Selected models ride the existing one-shot sessionStorage.selectedModels channel (consumed by new-chat init) rather than the snapshot. This fixes the case where a user manually picks a non-default model in a new chat and enables a tool: the reactive persistence gate only fires for existing chats, so without this the model fell back to default after the redirect, which in turn triggered the model-change reactive that wiped the restored tool selection. The legacy single-tool restore path is kept as a fallback for redirects started by an older client still in flight. Addresses #18271. --- src/lib/apis/configs/index.ts | 66 ++++++++++- src/lib/apis/configs/oauth-chat-state.test.ts | 112 ++++++++++++++++++ src/lib/components/chat/Chat.svelte | 75 ++++++++++-- src/lib/components/chat/MessageInput.svelte | 12 +- .../chat/MessageInput/IntegrationsMenu.svelte | 24 +++- 5 files changed, 270 insertions(+), 19 deletions(-) create mode 100644 src/lib/apis/configs/oauth-chat-state.test.ts diff --git a/src/lib/apis/configs/index.ts b/src/lib/apis/configs/index.ts index edb80028ed..c2843a5650 100644 --- a/src/lib/apis/configs/index.ts +++ b/src/lib/apis/configs/index.ts @@ -565,13 +565,69 @@ export const getOAuthClientAuthorizationUrl = (clientId: string, type: null | st return `${WEBUI_BASE_URL}/oauth/clients/${oauthClientId}/authorize`; }; -export const initiateOAuthRedirect = (tool: { - id: string; - serverId: string; - authType?: string | null; -}) => { +// Restore reads this to reject snapshots written by an older, incompatible client. +export const PENDING_OAUTH_CHAT_STATE_VERSION = 1; + +export type PendingOAuthChatState = { + selectedModels?: string[]; + selectedToolIds?: string[]; + selectedSkillIds?: string[]; + selectedFilterIds?: string[]; + webSearchEnabled?: boolean; + imageGenerationEnabled?: boolean; + codeInterpreterEnabled?: boolean; + prompt?: string; + returnTo?: string; +}; + +// Rejects absolute and protocol-relative ("//host") paths to prevent open redirects. +export const isSafeOAuthReturnPath = (path: unknown): path is string => { + return typeof path === 'string' && path.startsWith('/') && !path.startsWith('//'); +}; + +export const initiateOAuthRedirect = ( + tool: { + id: string; + serverId: string; + authType?: string | null; + }, + chatState: PendingOAuthChatState = {} +) => { sessionStorage.setItem('pendingOAuthToolId', tool.id); sessionStorage.setItem('oauthRedirectInProgressToolId', tool.id); + + // The OAuth callback always redirects to the site root, discarding the active + // chat. Snapshot the chat config here so it survives the round-trip. + try { + // Models ride the existing one-shot sessionStorage.selectedModels channel + // (consumed by new-chat init) rather than the snapshot, so restoring them + // does not re-trigger the model-change reactive that wipes tool selection. + const models = chatState.selectedModels?.filter( + (id): id is string => typeof id === 'string' && id !== '' + ); + if (models?.length) { + sessionStorage.setItem('selectedModels', JSON.stringify(models)); + } + + sessionStorage.setItem( + 'pendingOAuthChatState', + JSON.stringify({ + version: PENDING_OAUTH_CHAT_STATE_VERSION, + pendingOAuthToolId: tool.id, + selectedToolIds: chatState.selectedToolIds ?? [], + selectedSkillIds: chatState.selectedSkillIds ?? [], + selectedFilterIds: chatState.selectedFilterIds ?? [], + webSearchEnabled: chatState.webSearchEnabled ?? false, + imageGenerationEnabled: chatState.imageGenerationEnabled ?? false, + codeInterpreterEnabled: chatState.codeInterpreterEnabled ?? false, + prompt: chatState.prompt ?? '', + returnTo: chatState.returnTo ?? '' + }) + ); + } catch (err) { + console.error('Failed to persist chat state before OAuth redirect:', err); + } + const authUrl = getOAuthClientAuthorizationUrl(tool.serverId, tool.authType ?? 'mcp'); window.open(authUrl, '_self', 'noopener'); }; diff --git a/src/lib/apis/configs/oauth-chat-state.test.ts b/src/lib/apis/configs/oauth-chat-state.test.ts new file mode 100644 index 0000000000..f12b21cf39 --- /dev/null +++ b/src/lib/apis/configs/oauth-chat-state.test.ts @@ -0,0 +1,112 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { + initiateOAuthRedirect, + isSafeOAuthReturnPath, + PENDING_OAUTH_CHAT_STATE_VERSION +} from './index'; + +describe('isSafeOAuthReturnPath', () => { + it('accepts same-origin relative paths', () => { + expect(isSafeOAuthReturnPath('/')).toBe(true); + expect(isSafeOAuthReturnPath('/c/abc123')).toBe(true); + expect(isSafeOAuthReturnPath('/?tools=a,b&web-search=true')).toBe(true); + }); + + it('rejects protocol-relative paths', () => { + expect(isSafeOAuthReturnPath('//evil.com')).toBe(false); + expect(isSafeOAuthReturnPath('//evil.com/c/1')).toBe(false); + }); + + it('rejects absolute URLs and non-rooted values', () => { + expect(isSafeOAuthReturnPath('https://evil.com')).toBe(false); + expect(isSafeOAuthReturnPath('c/abc123')).toBe(false); + expect(isSafeOAuthReturnPath('')).toBe(false); + }); + + it('rejects non-string values', () => { + expect(isSafeOAuthReturnPath(undefined)).toBe(false); + expect(isSafeOAuthReturnPath(null)).toBe(false); + expect(isSafeOAuthReturnPath(42)).toBe(false); + expect(isSafeOAuthReturnPath({})).toBe(false); + }); +}); + +describe('initiateOAuthRedirect snapshot', () => { + const store = new Map(); + + beforeEach(() => { + store.clear(); + vi.stubGlobal('sessionStorage', { + getItem: (k: string) => (store.has(k) ? store.get(k)! : null), + setItem: (k: string, v: string) => void store.set(k, v), + removeItem: (k: string) => void store.delete(k) + }); + vi.stubGlobal('window', { open: vi.fn() }); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + const tool = { id: 'server:mcp', serverId: 'mcp', authType: 'mcp' as const }; + + it('persists a versioned snapshot with the provided chat state', () => { + initiateOAuthRedirect(tool, { + selectedToolIds: ['t1', 't2'], + selectedSkillIds: ['s1'], + selectedFilterIds: ['f1'], + webSearchEnabled: true, + imageGenerationEnabled: false, + codeInterpreterEnabled: true, + prompt: 'draft text', + returnTo: '/c/xyz' + }); + + const snap = JSON.parse(store.get('pendingOAuthChatState')!); + expect(snap.version).toBe(PENDING_OAUTH_CHAT_STATE_VERSION); + expect(snap.pendingOAuthToolId).toBe('server:mcp'); + expect(snap.selectedToolIds).toEqual(['t1', 't2']); + expect(snap.selectedSkillIds).toEqual(['s1']); + expect(snap.selectedFilterIds).toEqual(['f1']); + expect(snap.webSearchEnabled).toBe(true); + expect(snap.imageGenerationEnabled).toBe(false); + expect(snap.codeInterpreterEnabled).toBe(true); + expect(snap.prompt).toBe('draft text'); + expect(snap.returnTo).toBe('/c/xyz'); + }); + + it('mirrors selected models to the one-shot selectedModels key, not the snapshot', () => { + initiateOAuthRedirect(tool, { selectedModels: ['opus-4-8'], selectedToolIds: ['t1'] }); + + expect(JSON.parse(store.get('selectedModels')!)).toEqual(['opus-4-8']); + const snap = JSON.parse(store.get('pendingOAuthChatState')!); + expect('selectedModels' in snap).toBe(false); + }); + + it('does not write selectedModels for empty or placeholder model selection', () => { + initiateOAuthRedirect(tool, { selectedModels: [''] }); + expect(store.has('selectedModels')).toBe(false); + + initiateOAuthRedirect(tool, { selectedModels: [] }); + expect(store.has('selectedModels')).toBe(false); + }); + + it('defaults every field when chat state is omitted', () => { + initiateOAuthRedirect(tool); + const snap = JSON.parse(store.get('pendingOAuthChatState')!); + expect(snap.selectedToolIds).toEqual([]); + expect(snap.selectedSkillIds).toEqual([]); + expect(snap.selectedFilterIds).toEqual([]); + expect(snap.webSearchEnabled).toBe(false); + expect(snap.imageGenerationEnabled).toBe(false); + expect(snap.codeInterpreterEnabled).toBe(false); + expect(snap.prompt).toBe(''); + expect(snap.returnTo).toBe(''); + expect(store.has('selectedModels')).toBe(false); + }); + + it('keeps the legacy pendingOAuthToolId key for backward compatibility', () => { + initiateOAuthRedirect(tool); + expect(store.get('pendingOAuthToolId')).toBe('server:mcp'); + }); +}); diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index 8652df4630..2339801a51 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -97,7 +97,11 @@ import { uploadFile } from '$lib/apis/files'; import { createOpenAITextStream } from '$lib/apis/streaming'; import { getFunctions } from '$lib/apis/functions'; - import { initiateOAuthRedirect } from '$lib/apis/configs'; + import { + initiateOAuthRedirect, + isSafeOAuthReturnPath, + PENDING_OAUTH_CHAT_STATE_VERSION + } from '$lib/apis/configs'; import { updateFolderById } from '$lib/apis/folders'; import Banner from '../common/Banner.svelte'; @@ -704,7 +708,17 @@ saveSessionSelectedModels(); await tick(); - initiateOAuthRedirect(nextTool); + initiateOAuthRedirect(nextTool, { + selectedModels: selectedModelIds, + selectedToolIds, + selectedSkillIds, + selectedFilterIds, + webSearchEnabled, + imageGenerationEnabled, + codeInterpreterEnabled, + prompt, + returnTo: window.location.pathname + window.location.search + }); }; const resetInput = async () => { @@ -1866,12 +1880,59 @@ .filter((id) => id); } - // Restore tool selection after OAuth redirect - const pendingToolId = sessionStorage.getItem('pendingOAuthToolId'); - if (pendingToolId) { + // Restore chat state captured before an OAuth redirect. The snapshot form + // supersedes the legacy single-tool form; the latter is kept as a fallback + // for redirects started by an older client still in flight. + const oauthSnapshotRaw = sessionStorage.getItem('pendingOAuthChatState'); + if (oauthSnapshotRaw) { + sessionStorage.removeItem('pendingOAuthChatState'); sessionStorage.removeItem('pendingOAuthToolId'); - if (!selectedToolIds.includes(pendingToolId)) { - selectedToolIds = [...selectedToolIds, pendingToolId]; + try { + const snap = JSON.parse(oauthSnapshotRaw); + if (snap?.version === PENDING_OAUTH_CHAT_STATE_VERSION) { + if (Array.isArray(snap.selectedToolIds)) { + selectedToolIds = [...snap.selectedToolIds]; + } + if (typeof snap.pendingOAuthToolId === 'string' && snap.pendingOAuthToolId) { + if (!selectedToolIds.includes(snap.pendingOAuthToolId)) { + selectedToolIds = [...selectedToolIds, snap.pendingOAuthToolId]; + } + } + if (Array.isArray(snap.selectedSkillIds)) { + selectedSkillIds = [...snap.selectedSkillIds]; + } + if (Array.isArray(snap.selectedFilterIds)) { + selectedFilterIds = [...snap.selectedFilterIds]; + } + if (typeof snap.webSearchEnabled === 'boolean') { + webSearchEnabled = snap.webSearchEnabled; + } + if (typeof snap.imageGenerationEnabled === 'boolean') { + imageGenerationEnabled = snap.imageGenerationEnabled; + } + if (typeof snap.codeInterpreterEnabled === 'boolean') { + codeInterpreterEnabled = snap.codeInterpreterEnabled; + } + if (typeof snap.prompt === 'string' && snap.prompt) { + messageInput?.setText(snap.prompt); + } + if ( + isSafeOAuthReturnPath(snap.returnTo) && + snap.returnTo !== $page.url.pathname + $page.url.search + ) { + goto(snap.returnTo); + } + } + } catch (e) { + console.error('Failed to restore OAuth chat snapshot:', e); + } + } else { + const pendingToolId = sessionStorage.getItem('pendingOAuthToolId'); + if (pendingToolId) { + sessionStorage.removeItem('pendingOAuthToolId'); + if (!selectedToolIds.includes(pendingToolId)) { + selectedToolIds = [...selectedToolIds, pendingToolId]; + } } } diff --git a/src/lib/components/chat/MessageInput.svelte b/src/lib/components/chat/MessageInput.svelte index 883bdbad08..f492f48622 100644 --- a/src/lib/components/chat/MessageInput.svelte +++ b/src/lib/components/chat/MessageInput.svelte @@ -2210,7 +2210,17 @@