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 @@