From 10685fca9f978584d5bbd9fab27e34e6dcfc60b2 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 14 Feb 2026 13:18:50 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=97=82=EF=B8=8F=20refactor:=20Artifacts?= =?UTF-8?q?=20via=20Model=20Specs=20&=20Scope=20Badge=20Persistence=20by?= =?UTF-8?q?=20Spec=20Context=20(#11796)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🔧 refactor: Simplify MCP selection logic in useMCPSelect hook - Removed redundant useEffect for setting ephemeral agent when MCP values change. - Integrated ephemeral agent update directly into the MCP value change handler, improving code clarity and reducing unnecessary re-renders. - Updated dependencies in the effect hook to ensure proper state management. Why Effect 2 Was Added (PR #9528) PR #9528 was a refactor that migrated MCP state from useLocalStorage hooks to Jotai atomWithStorage. Before that PR, useLocalStorage handled bidirectional sync between localStorage and Recoil in one abstraction. After the migration, the two useEffect hooks were introduced to bridge Jotai ↔ Recoil: - Effect 1 (Recoil → Jotai): When ephemeralAgent.mcp changes externally, update the Jotai atom (which drives the UI dropdown) - Effect 2 (Jotai → Recoil): When mcpValues changes, push it back to ephemeralAgent.mcp (which is read at submission time) Effect 2 was needed because in that PR's design, setMCPValues only wrote to Jotai — it never touched Recoil. Effect 2 was the bridge to propagate user selections into the ephemeral agent. Why Removing It Is Correct All user-initiated MCP changes go through setMCPValues. The callers are in useMCPServerManager: toggleServerSelection, batchToggleServers, OAuth success callbacks, and access revocation. Our change puts the Recoil write directly in that callback, so all these paths are covered. All external changes go through Recoil, handled by Effect 1 (kept). Model spec application (applyModelSpecEphemeralAgent), agent template application after submission, and BadgeRowContext initialization all write directly to ephemeralAgentByConvoId. Effect 1 watches ephemeralAgent?.mcp and syncs those into the Jotai atom for the UI. There is no code path where mcpValues changes without going through setMCPValues or Effect 1. The only other source is atomWithStorage's getOnInit reading from localStorage on mount — that's just restoring persisted state and is harmless (overwritten by Effect 1 if the ephemeral agent has values). Additional Benefits - Eliminates the race condition. Effect 2 fired on mount with Jotai's stale default ([]), overwriting ephemeralAgent.mcp that had been set by a model spec. Our change prevents that because the imperative sync only fires on explicit user action. - Eliminates infinite loop risk. The old bidirectional two-effect approach relied on isEqual/JSON.stringify checks to break cycles. The new unidirectional-reactive (Effect 1) + imperative (setMCPValues) approach has no such risk. - Effect 1's enhancements are preserved. The mcp_clear sentinel handling and configuredServers filtering (both added after PR #9528) continue to work correctly. * ✨ feat: Add artifacts support to model specifications and ephemeral agents - Introduced `artifacts` property in the model specification and ephemeral agent types, allowing for string or boolean values. - Updated `applyModelSpecEphemeralAgent` to handle artifacts, defaulting to 'default' if true or an empty string if not specified. - Enhanced localStorage handling to store artifacts alongside other agent properties, improving state management for ephemeral agents. * 🔧 refactor: Update BadgeRowContext to improve localStorage handling - Modified the logic to only apply values from localStorage that were actually stored, preventing unnecessary overrides of the ephemeral agent. - Simplified the setting of ephemeral agent values by directly using initialValues, enhancing code clarity and maintainability. * 🔧 refactor: Enhance ephemeral agent handling in BadgeRowContext and model spec application - Updated BadgeRowContext to apply localStorage values only for tools not already set in ephemeralAgent, improving state management. - Modified useApplyModelSpecEffects to reset the ephemeral agent when no spec is provided but specs are configured, ensuring localStorage defaults are applied correctly. - Streamlined the logic for applying model spec properties, enhancing clarity and maintainability. * refactor: Isolate spec and non-spec tool/MCP state with environment-keyed storage Spec tool state (badges, MCP) and non-spec user preferences previously shared conversation-keyed localStorage, causing cross-pollination when switching between spec and non-spec models. This introduces environment-keyed storage so each context maintains independent persisted state. Key changes: - Spec active: no localStorage persistence — admin config always applied fresh - Non-spec (with specs configured): tool/MCP state persisted to __defaults__ key - No specs configured: zero behavior change (conversation-keyed storage) - Per-conversation isolation preserved for existing conversations - Dual-write on user interaction updates both conversation and environment keys - Remove mcp_clear sentinel in favor of null ephemeral agent reset * refactor: Enhance ephemeral agent initialization and MCP handling in BadgeRowContext and useMCPSelect - Updated BadgeRowContext to clarify the handling of localStorage values for ephemeral agents, ensuring proper initialization based on conversation state. - Improved useMCPSelect tests to accurately reflect behavior when setting empty MCP values, ensuring the visual selection clears as expected. - Introduced environment-keyed storage logic to maintain independent state for spec and non-spec contexts, enhancing user experience during context switching. * test: Add comprehensive tests for useToolToggle and applyModelSpecEphemeralAgent hooks - Introduced unit tests for the useToolToggle hook, covering dual-write behavior in non-spec mode and per-conversation isolation. - Added tests for applyModelSpecEphemeralAgent, ensuring correct application of model specifications and user overrides from localStorage. - Enhanced test coverage for ephemeral agent state management during conversation transitions, validating expected behaviors for both new and existing conversations. --- client/src/Providers/BadgeRowContext.tsx | 137 +++++--- client/src/components/Chat/Input/BadgeRow.tsx | 8 +- client/src/components/Chat/Input/ChatForm.tsx | 1 + .../src/components/Chat/Input/MCPSelect.tsx | 8 +- .../src/components/Chat/Input/MCPSubMenu.tsx | 6 +- client/src/components/MCP/MCPConfigDialog.tsx | 3 + .../MCP/ServerInitializationSection.tsx | 4 +- .../hooks/Agents/useApplyModelSpecAgents.ts | 13 + .../hooks/MCP/__tests__/useMCPSelect.test.tsx | 252 +++++++++++++- client/src/hooks/MCP/useMCPSelect.ts | 60 ++-- client/src/hooks/MCP/useMCPServerManager.ts | 6 +- .../Plugins/__tests__/useToolToggle.test.tsx | 328 ++++++++++++++++++ client/src/hooks/Plugins/useToolToggle.ts | 18 +- .../applyModelSpecEphemeralAgent.test.ts | 274 +++++++++++++++ client/src/utils/endpoints.ts | 46 ++- packages/data-provider/src/config.ts | 2 + packages/data-provider/src/models.ts | 2 + packages/data-provider/src/types.ts | 1 + 18 files changed, 1084 insertions(+), 85 deletions(-) create mode 100644 client/src/hooks/Plugins/__tests__/useToolToggle.test.tsx create mode 100644 client/src/utils/__tests__/applyModelSpecEphemeralAgent.test.ts diff --git a/client/src/Providers/BadgeRowContext.tsx b/client/src/Providers/BadgeRowContext.tsx index 40df795aba..dce1c38a78 100644 --- a/client/src/Providers/BadgeRowContext.tsx +++ b/client/src/Providers/BadgeRowContext.tsx @@ -1,4 +1,4 @@ -import React, { createContext, useContext, useEffect, useRef } from 'react'; +import React, { createContext, useContext, useEffect, useMemo, useRef } from 'react'; import { useSetRecoilState } from 'recoil'; import { Tools, Constants, LocalStorageKeys, AgentCapabilities } from 'librechat-data-provider'; import type { TAgentsEndpoint } from 'librechat-data-provider'; @@ -9,11 +9,13 @@ import { useCodeApiKeyForm, useToolToggle, } from '~/hooks'; -import { getTimestampedValue, setTimestamp } from '~/utils/timestamps'; +import { getTimestampedValue } from '~/utils/timestamps'; +import { useGetStartupConfig } from '~/data-provider'; import { ephemeralAgentByConvoId } from '~/store'; interface BadgeRowContextType { conversationId?: string | null; + storageContextKey?: string; agentsConfig?: TAgentsEndpoint | null; webSearch: ReturnType; artifacts: ReturnType; @@ -38,34 +40,70 @@ interface BadgeRowProviderProps { children: React.ReactNode; isSubmitting?: boolean; conversationId?: string | null; + specName?: string | null; } export default function BadgeRowProvider({ children, isSubmitting, conversationId, + specName, }: BadgeRowProviderProps) { - const lastKeyRef = useRef(''); + const lastContextKeyRef = useRef(''); const hasInitializedRef = useRef(false); const { agentsConfig } = useGetAgentsConfig(); + const { data: startupConfig } = useGetStartupConfig(); const key = conversationId ?? Constants.NEW_CONVO; + const hasModelSpecs = (startupConfig?.modelSpecs?.list?.length ?? 0) > 0; + + /** + * Compute the storage context key for non-spec persistence: + * - `__defaults__`: specs configured but none active → shared defaults key + * - undefined: spec active (no persistence) or no specs configured (original behavior) + * + * When a spec is active, tool/MCP state is NOT persisted — the admin's spec + * configuration is always applied fresh. Only non-spec user preferences persist. + */ + const storageContextKey = useMemo(() => { + if (!specName && hasModelSpecs) { + return Constants.spec_defaults_key as string; + } + return undefined; + }, [specName, hasModelSpecs]); + + /** + * Compute the storage suffix for reading localStorage defaults: + * - New conversations read from environment key (spec or non-spec defaults) + * - Existing conversations read from conversation key (per-conversation state) + */ + const isNewConvo = key === Constants.NEW_CONVO; + const storageSuffix = isNewConvo && storageContextKey ? storageContextKey : key; const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(key)); - /** Initialize ephemeralAgent from localStorage on mount and when conversation changes */ + /** Initialize ephemeralAgent from localStorage on mount and when conversation/spec changes. + * Skipped when a spec is active — applyModelSpecEphemeralAgent handles both new conversations + * (pure spec values) and existing conversations (spec values + localStorage overrides). */ useEffect(() => { if (isSubmitting) { return; } - // Check if this is a new conversation or the first load - if (!hasInitializedRef.current || lastKeyRef.current !== key) { + if (specName) { + // Spec active: applyModelSpecEphemeralAgent handles all state (spec base + localStorage + // overrides for existing conversations). Reset init flag so switching back to non-spec + // triggers a fresh re-init. + hasInitializedRef.current = false; + return; + } + // Check if this is a new conversation/spec or the first load + if (!hasInitializedRef.current || lastContextKeyRef.current !== storageSuffix) { hasInitializedRef.current = true; - lastKeyRef.current = key; + lastContextKeyRef.current = storageSuffix; - const codeToggleKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${key}`; - const webSearchToggleKey = `${LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_}${key}`; - const fileSearchToggleKey = `${LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_}${key}`; - const artifactsToggleKey = `${LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_}${key}`; + const codeToggleKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${storageSuffix}`; + const webSearchToggleKey = `${LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_}${storageSuffix}`; + const fileSearchToggleKey = `${LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_}${storageSuffix}`; + const artifactsToggleKey = `${LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_}${storageSuffix}`; const codeToggleValue = getTimestampedValue(codeToggleKey); const webSearchToggleValue = getTimestampedValue(webSearchToggleKey); @@ -106,39 +144,53 @@ export default function BadgeRowProvider({ } } - /** - * Always set values for all tools (use defaults if not in `localStorage`) - * If `ephemeralAgent` is `null`, create a new object with just our tool values - */ - const finalValues = { - [Tools.execute_code]: initialValues[Tools.execute_code] ?? false, - [Tools.web_search]: initialValues[Tools.web_search] ?? false, - [Tools.file_search]: initialValues[Tools.file_search] ?? false, - [AgentCapabilities.artifacts]: initialValues[AgentCapabilities.artifacts] ?? false, - }; + const hasOverrides = Object.keys(initialValues).length > 0; - setEphemeralAgent((prev) => ({ - ...(prev || {}), - ...finalValues, - })); - - Object.entries(finalValues).forEach(([toolKey, value]) => { - if (value !== false) { - let storageKey = artifactsToggleKey; - if (toolKey === Tools.execute_code) { - storageKey = codeToggleKey; - } else if (toolKey === Tools.web_search) { - storageKey = webSearchToggleKey; - } else if (toolKey === Tools.file_search) { - storageKey = fileSearchToggleKey; + /** Read persisted MCP values from localStorage */ + let mcpOverrides: string[] | null = null; + const mcpStorageKey = `${LocalStorageKeys.LAST_MCP_}${storageSuffix}`; + const mcpRaw = localStorage.getItem(mcpStorageKey); + if (mcpRaw !== null) { + try { + const parsed = JSON.parse(mcpRaw); + if (Array.isArray(parsed) && parsed.length > 0) { + mcpOverrides = parsed; } - // Store the value and set timestamp for existing values - localStorage.setItem(storageKey, JSON.stringify(value)); - setTimestamp(storageKey); + } catch (e) { + console.error('Failed to parse MCP values:', e); } + } + + setEphemeralAgent((prev) => { + if (prev == null) { + /** ephemeralAgent is null — use localStorage defaults */ + if (hasOverrides || mcpOverrides) { + const result = { ...initialValues }; + if (mcpOverrides) { + result.mcp = mcpOverrides; + } + return result; + } + return prev; + } + /** ephemeralAgent already has values (from prior state). + * Only fill in undefined keys from localStorage. */ + let changed = false; + const result = { ...prev }; + for (const [toolKey, value] of Object.entries(initialValues)) { + if (result[toolKey] === undefined) { + result[toolKey] = value; + changed = true; + } + } + if (mcpOverrides && result.mcp === undefined) { + result.mcp = mcpOverrides; + changed = true; + } + return changed ? result : prev; }); } - }, [key, isSubmitting, setEphemeralAgent]); + }, [storageSuffix, specName, isSubmitting, setEphemeralAgent]); /** CodeInterpreter hooks */ const codeApiKeyForm = useCodeApiKeyForm({}); @@ -146,6 +198,7 @@ export default function BadgeRowProvider({ const codeInterpreter = useToolToggle({ conversationId, + storageContextKey, setIsDialogOpen: setCodeDialogOpen, toolKey: Tools.execute_code, localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, @@ -161,6 +214,7 @@ export default function BadgeRowProvider({ const webSearch = useToolToggle({ conversationId, + storageContextKey, toolKey: Tools.web_search, localStorageKey: LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, setIsDialogOpen: setWebSearchDialogOpen, @@ -173,6 +227,7 @@ export default function BadgeRowProvider({ /** FileSearch hook */ const fileSearch = useToolToggle({ conversationId, + storageContextKey, toolKey: Tools.file_search, localStorageKey: LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_, isAuthenticated: true, @@ -181,12 +236,13 @@ export default function BadgeRowProvider({ /** Artifacts hook - using a custom key since it's not a Tool but a capability */ const artifacts = useToolToggle({ conversationId, + storageContextKey, toolKey: AgentCapabilities.artifacts, localStorageKey: LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_, isAuthenticated: true, }); - const mcpServerManager = useMCPServerManager({ conversationId }); + const mcpServerManager = useMCPServerManager({ conversationId, storageContextKey }); const value: BadgeRowContextType = { webSearch, @@ -194,6 +250,7 @@ export default function BadgeRowProvider({ fileSearch, agentsConfig, conversationId, + storageContextKey, codeApiKeyForm, codeInterpreter, searchApiKeyForm, diff --git a/client/src/components/Chat/Input/BadgeRow.tsx b/client/src/components/Chat/Input/BadgeRow.tsx index 5036dcd5e4..6fea6b0d58 100644 --- a/client/src/components/Chat/Input/BadgeRow.tsx +++ b/client/src/components/Chat/Input/BadgeRow.tsx @@ -28,6 +28,7 @@ interface BadgeRowProps { onChange: (badges: Pick[]) => void; onToggle?: (badgeId: string, currentActive: boolean) => void; conversationId?: string | null; + specName?: string | null; isSubmitting?: boolean; isInChat: boolean; } @@ -142,6 +143,7 @@ const dragReducer = (state: DragState, action: DragAction): DragState => { function BadgeRow({ showEphemeralBadges, conversationId, + specName, isSubmitting, onChange, onToggle, @@ -320,7 +322,11 @@ function BadgeRow({ }, [dragState.draggedBadge, handleMouseMove, handleMouseUp]); return ( - +
{showEphemeralBadges === true && } {tempBadges.map((badge, index) => ( diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx index f8f0fbb40b..45277e5b9c 100644 --- a/client/src/components/Chat/Input/ChatForm.tsx +++ b/client/src/components/Chat/Input/ChatForm.tsx @@ -325,6 +325,7 @@ const ChatForm = memo(({ index = 0 }: { index?: number }) => { } isSubmitting={isSubmitting} conversationId={conversationId} + specName={conversation?.spec} onChange={setBadges} isInChat={ Array.isArray(conversation?.messages) && conversation.messages.length >= 1 diff --git a/client/src/components/Chat/Input/MCPSelect.tsx b/client/src/components/Chat/Input/MCPSelect.tsx index 278e603db0..a5356f5094 100644 --- a/client/src/components/Chat/Input/MCPSelect.tsx +++ b/client/src/components/Chat/Input/MCPSelect.tsx @@ -11,7 +11,7 @@ import { useHasAccess } from '~/hooks'; import { cn } from '~/utils'; function MCPSelectContent() { - const { conversationId, mcpServerManager } = useBadgeRowContext(); + const { conversationId, storageContextKey, mcpServerManager } = useBadgeRowContext(); const { localize, isPinned, @@ -128,7 +128,11 @@ function MCPSelectContent() { {configDialogProps && ( - + )} ); diff --git a/client/src/components/Chat/Input/MCPSubMenu.tsx b/client/src/components/Chat/Input/MCPSubMenu.tsx index ca547ca1f7..b0b8fad1bb 100644 --- a/client/src/components/Chat/Input/MCPSubMenu.tsx +++ b/client/src/components/Chat/Input/MCPSubMenu.tsx @@ -15,7 +15,7 @@ interface MCPSubMenuProps { const MCPSubMenu = React.forwardRef( ({ placeholder, ...props }, ref) => { const localize = useLocalize(); - const { mcpServerManager } = useBadgeRowContext(); + const { storageContextKey, mcpServerManager } = useBadgeRowContext(); const { isPinned, mcpValues, @@ -106,7 +106,9 @@ const MCPSubMenu = React.forwardRef(
- {configDialogProps && } + {configDialogProps && ( + + )} ); }, diff --git a/client/src/components/MCP/MCPConfigDialog.tsx b/client/src/components/MCP/MCPConfigDialog.tsx index a3727971e9..f1079c2799 100644 --- a/client/src/components/MCP/MCPConfigDialog.tsx +++ b/client/src/components/MCP/MCPConfigDialog.tsx @@ -24,6 +24,7 @@ interface MCPConfigDialogProps { serverName: string; serverStatus?: MCPServerStatus; conversationId?: string | null; + storageContextKey?: string; } export default function MCPConfigDialog({ @@ -36,6 +37,7 @@ export default function MCPConfigDialog({ serverName, serverStatus, conversationId, + storageContextKey, }: MCPConfigDialogProps) { const localize = useLocalize(); @@ -167,6 +169,7 @@ export default function MCPConfigDialog({ 0} /> diff --git a/client/src/components/MCP/ServerInitializationSection.tsx b/client/src/components/MCP/ServerInitializationSection.tsx index b5f71335d7..c080866b3d 100644 --- a/client/src/components/MCP/ServerInitializationSection.tsx +++ b/client/src/components/MCP/ServerInitializationSection.tsx @@ -9,12 +9,14 @@ interface ServerInitializationSectionProps { requiresOAuth: boolean; hasCustomUserVars?: boolean; conversationId?: string | null; + storageContextKey?: string; } export default function ServerInitializationSection({ serverName, requiresOAuth, conversationId, + storageContextKey, sidePanel = false, hasCustomUserVars = false, }: ServerInitializationSectionProps) { @@ -28,7 +30,7 @@ export default function ServerInitializationSection({ initializeServer, availableMCPServers, revokeOAuthForServer, - } = useMCPServerManager({ conversationId }); + } = useMCPServerManager({ conversationId, storageContextKey }); const { connectionStatus } = useMCPConnectionStatus({ enabled: !!availableMCPServers && availableMCPServers.length > 0, diff --git a/client/src/hooks/Agents/useApplyModelSpecAgents.ts b/client/src/hooks/Agents/useApplyModelSpecAgents.ts index 94d62a058a..2c677f85ca 100644 --- a/client/src/hooks/Agents/useApplyModelSpecAgents.ts +++ b/client/src/hooks/Agents/useApplyModelSpecAgents.ts @@ -1,4 +1,5 @@ import { useCallback } from 'react'; +import { Constants } from 'librechat-data-provider'; import type { TStartupConfig, TSubmission } from 'librechat-data-provider'; import { useUpdateEphemeralAgent, useApplyNewAgentTemplate } from '~/store/agents'; import { getModelSpec, applyModelSpecEphemeralAgent } from '~/utils'; @@ -6,6 +7,10 @@ import { getModelSpec, applyModelSpecEphemeralAgent } from '~/utils'; /** * Hook that applies a model spec from a preset to an ephemeral agent. * This is used when initializing a new conversation with a preset that has a spec. + * + * When a spec is provided, its tool settings are applied to the ephemeral agent. + * When no spec is provided but specs are configured, the ephemeral agent is reset + * to null so BadgeRowContext can apply localStorage defaults (non-spec experience). */ export function useApplyModelSpecEffects() { const updateEphemeralAgent = useUpdateEphemeralAgent(); @@ -20,6 +25,11 @@ export function useApplyModelSpecEffects() { startupConfig?: TStartupConfig; }) => { if (specName == null || !specName) { + if (startupConfig?.modelSpecs?.list?.length) { + /** Specs are configured but none selected — reset ephemeral agent to null + * so BadgeRowContext fills all values (tool toggles + MCP) from localStorage. */ + updateEphemeralAgent((convoId ?? Constants.NEW_CONVO) || Constants.NEW_CONVO, null); + } return; } @@ -80,6 +90,9 @@ export function useApplyAgentTemplate() { web_search: ephemeralAgent?.web_search ?? modelSpec.webSearch ?? false, file_search: ephemeralAgent?.file_search ?? modelSpec.fileSearch ?? false, execute_code: ephemeralAgent?.execute_code ?? modelSpec.executeCode ?? false, + artifacts: + ephemeralAgent?.artifacts ?? + (modelSpec.artifacts === true ? 'default' : modelSpec.artifacts || ''), }; mergedAgent.mcp = [...new Set(mergedAgent.mcp)]; diff --git a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx index 26595b611c..783f525b9c 100644 --- a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx +++ b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx @@ -415,7 +415,7 @@ describe('useMCPSelect', () => { }); }); - it('should handle empty ephemeralAgent.mcp array correctly', async () => { + it('should clear mcpValues when ephemeralAgent.mcp is set to empty array', async () => { // Create a shared wrapper const { Wrapper, servers } = createWrapper(['initial-value']); @@ -437,19 +437,21 @@ describe('useMCPSelect', () => { expect(result.current.mcpHook.mcpValues).toEqual(['initial-value']); }); - // Try to set empty array externally + // Set empty array externally (e.g., spec with no MCP servers) act(() => { result.current.setEphemeralAgent({ mcp: [], }); }); - // Values should remain unchanged since empty mcp array doesn't trigger update - // (due to the condition: ephemeralAgent?.mcp && ephemeralAgent.mcp.length > 0) - expect(result.current.mcpHook.mcpValues).toEqual(['initial-value']); + // Jotai atom should be cleared — an explicit empty mcp array means + // the spec (or reset) has no MCP servers, so the visual selection must clear + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual([]); + }); }); - it('should handle ephemeralAgent with clear mcp value', async () => { + it('should handle ephemeralAgent being reset to null', async () => { // Create a shared wrapper const { Wrapper, servers } = createWrapper(['server1', 'server2']); @@ -471,16 +473,15 @@ describe('useMCPSelect', () => { expect(result.current.mcpHook.mcpValues).toEqual(['server1', 'server2']); }); - // Set ephemeralAgent with clear value + // Reset ephemeralAgent to null (simulating non-spec reset) act(() => { - result.current.setEphemeralAgent({ - mcp: [Constants.mcp_clear as string], - }); + result.current.setEphemeralAgent(null); }); - // mcpValues should be cleared + // mcpValues should remain unchanged since null ephemeral agent + // doesn't trigger the sync effect (mcps.length === 0) await waitFor(() => { - expect(result.current.mcpHook.mcpValues).toEqual([]); + expect(result.current.mcpHook.mcpValues).toEqual(['server1', 'server2']); }); }); @@ -590,6 +591,233 @@ describe('useMCPSelect', () => { }); }); + describe('Environment-Keyed Storage (storageContextKey)', () => { + it('should use storageContextKey as atom key for new conversations', async () => { + const { Wrapper, servers } = createWrapper(['server1', 'server2']); + const storageContextKey = '__defaults__'; + + // Hook A: new conversation with storageContextKey + const { result: resultA } = renderHook( + () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + { wrapper: Wrapper }, + ); + + act(() => { + resultA.current.setMCPValues(['server1']); + }); + + await waitFor(() => { + expect(resultA.current.mcpValues).toEqual(['server1']); + }); + + // Hook B: new conversation WITHOUT storageContextKey (different environment) + const { result: resultB } = renderHook( + () => useMCPSelect({ conversationId: null, servers }), + { wrapper: Wrapper }, + ); + + // Should NOT see server1 since it's a different atom (NEW_CONVO vs __defaults__) + expect(resultB.current.mcpValues).toEqual([]); + }); + + it('should use conversationId as atom key for existing conversations even with storageContextKey', async () => { + const conversationId = 'existing-convo-123'; + const { Wrapper, servers } = createWrapper(['server1', 'server2']); + const storageContextKey = '__defaults__'; + + const { result } = renderHook( + () => useMCPSelect({ conversationId, storageContextKey, servers }), + { wrapper: Wrapper }, + ); + + act(() => { + result.current.setMCPValues(['server1', 'server2']); + }); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual(['server1', 'server2']); + }); + + // Verify timestamp was written to the conversation key, not the environment key + const convoKey = `${LocalStorageKeys.LAST_MCP_}${conversationId}`; + expect(setTimestamp).toHaveBeenCalledWith(convoKey); + }); + + it('should dual-write to environment key when storageContextKey is provided', async () => { + const { Wrapper, servers } = createWrapper(['server1', 'server2']); + const storageContextKey = '__defaults__'; + + const { result } = renderHook( + () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + { wrapper: Wrapper }, + ); + + act(() => { + result.current.setMCPValues(['server1', 'server2']); + }); + + await waitFor(() => { + // Verify dual-write to environment key + const envKey = `${LocalStorageKeys.LAST_MCP_}${storageContextKey}`; + expect(localStorage.getItem(envKey)).toEqual(JSON.stringify(['server1', 'server2'])); + expect(setTimestamp).toHaveBeenCalledWith(envKey); + }); + }); + + it('should NOT dual-write when storageContextKey is undefined', async () => { + const conversationId = 'convo-no-specs'; + const { Wrapper, servers } = createWrapper(['server1']); + + const { result } = renderHook(() => useMCPSelect({ conversationId, servers }), { + wrapper: Wrapper, + }); + + act(() => { + result.current.setMCPValues(['server1']); + }); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual(['server1']); + }); + + // Only the conversation-keyed timestamp should be set, no environment key + const envKey = `${LocalStorageKeys.LAST_MCP_}__defaults__`; + expect(localStorage.getItem(envKey)).toBeNull(); + }); + + it('should isolate per-conversation state from environment defaults', async () => { + const { Wrapper, servers } = createWrapper(['server1', 'server2', 'server3']); + const storageContextKey = '__defaults__'; + + // Set environment defaults via new conversation + const { result: newConvoResult } = renderHook( + () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + { wrapper: Wrapper }, + ); + + act(() => { + newConvoResult.current.setMCPValues(['server1', 'server2']); + }); + + await waitFor(() => { + expect(newConvoResult.current.mcpValues).toEqual(['server1', 'server2']); + }); + + // Existing conversation should have its own isolated state + const { result: existingResult } = renderHook( + () => useMCPSelect({ conversationId: 'existing-convo', storageContextKey, servers }), + { wrapper: Wrapper }, + ); + + // Should start empty (its own atom), not inherit from defaults + expect(existingResult.current.mcpValues).toEqual([]); + + // Set different value for existing conversation + act(() => { + existingResult.current.setMCPValues(['server3']); + }); + + await waitFor(() => { + expect(existingResult.current.mcpValues).toEqual(['server3']); + }); + + // New conversation defaults should be unchanged + expect(newConvoResult.current.mcpValues).toEqual(['server1', 'server2']); + }); + }); + + describe('Spec/Non-Spec Context Switching', () => { + it('should clear MCP when ephemeral agent switches to empty mcp (spec with no MCP)', async () => { + const { Wrapper, servers } = createWrapper(['server1', 'server2']); + const storageContextKey = '__defaults__'; + + const TestComponent = ({ ctxKey }: { ctxKey?: string }) => { + const mcpHook = useMCPSelect({ conversationId: null, storageContextKey: ctxKey, servers }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, setEphemeralAgent }; + }; + + // Start in non-spec context with some servers selected + const { result } = renderHook(() => TestComponent({ ctxKey: storageContextKey }), { + wrapper: Wrapper, + }); + + act(() => { + result.current.mcpHook.setMCPValues(['server1', 'server2']); + }); + + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual(['server1', 'server2']); + }); + + // Simulate switching to a spec with no MCP — ephemeral agent gets mcp: [] + act(() => { + result.current.setEphemeralAgent({ mcp: [] }); + }); + + // MCP values should clear since the spec explicitly has no MCP servers + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual([]); + }); + }); + + it('should handle ephemeral agent with spec MCP servers syncing to Jotai atom', async () => { + const { Wrapper, servers } = createWrapper(['spec-server1', 'spec-server2']); + + const TestComponent = () => { + const mcpHook = useMCPSelect({ conversationId: null, servers }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + // Simulate spec application setting ephemeral agent MCP + act(() => { + result.current.setEphemeralAgent({ + mcp: ['spec-server1', 'spec-server2'], + execute_code: true, + }); + }); + + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual(['spec-server1', 'spec-server2']); + }); + }); + + it('should handle null ephemeral agent reset (non-spec with specs configured)', async () => { + const { Wrapper, servers } = createWrapper(['server1', 'server2']); + + const TestComponent = () => { + const mcpHook = useMCPSelect({ servers }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + // Set values from a spec + act(() => { + result.current.setEphemeralAgent({ mcp: ['server1', 'server2'] }); + }); + + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual(['server1', 'server2']); + }); + + // Reset ephemeral agent to null (switching to non-spec) + act(() => { + result.current.setEphemeralAgent(null); + }); + + // mcpValues should remain unchanged — null ephemeral agent doesn't trigger sync + // (BadgeRowContext will fill from localStorage defaults separately) + await waitFor(() => { + expect(result.current.mcpHook.mcpValues).toEqual(['server1', 'server2']); + }); + }); + }); + describe('Memory Leak Prevention', () => { it('should not leak memory on repeated updates', async () => { const values = Array.from({ length: 100 }, (_, i) => `value-${i}`); diff --git a/client/src/hooks/MCP/useMCPSelect.ts b/client/src/hooks/MCP/useMCPSelect.ts index e0118b8be1..b15786f678 100644 --- a/client/src/hooks/MCP/useMCPSelect.ts +++ b/client/src/hooks/MCP/useMCPSelect.ts @@ -9,9 +9,11 @@ import { MCPServerDefinition } from './useMCPServerManager'; export function useMCPSelect({ conversationId, + storageContextKey, servers, }: { conversationId?: string | null; + storageContextKey?: string; servers: MCPServerDefinition[]; }) { const key = conversationId ?? Constants.NEW_CONVO; @@ -19,47 +21,61 @@ export function useMCPSelect({ return new Set(servers?.map((s) => s.serverName)); }, [servers]); + /** + * For new conversations, key the MCP atom by environment (spec or defaults) + * so switching between spec ↔ non-spec gives each its own atom. + * For existing conversations, key by conversation ID for per-conversation isolation. + */ + const isNewConvo = key === Constants.NEW_CONVO; + const mcpAtomKey = isNewConvo && storageContextKey ? storageContextKey : key; + const [isPinned, setIsPinned] = useAtom(mcpPinnedAtom); - const [mcpValues, setMCPValuesRaw] = useAtom(mcpValuesAtomFamily(key)); + const [mcpValues, setMCPValuesRaw] = useAtom(mcpValuesAtomFamily(mcpAtomKey)); const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key)); - // Sync Jotai state with ephemeral agent state + // Sync ephemeral agent MCP → Jotai atom (strip unconfigured servers) useEffect(() => { - const mcps = ephemeralAgent?.mcp ?? []; - if (mcps.length === 1 && mcps[0] === Constants.mcp_clear) { - setMCPValuesRaw([]); - } else if (mcps.length > 0 && configuredServers.size > 0) { - // Strip out servers that are not available in the startup config + const mcps = ephemeralAgent?.mcp; + if (Array.isArray(mcps) && mcps.length > 0 && configuredServers.size > 0) { const activeMcps = mcps.filter((mcp) => configuredServers.has(mcp)); - setMCPValuesRaw(activeMcps); - } - }, [ephemeralAgent?.mcp, setMCPValuesRaw, configuredServers]); - - useEffect(() => { - setEphemeralAgent((prev) => { - if (!isEqual(prev?.mcp, mcpValues)) { - return { ...(prev ?? {}), mcp: mcpValues }; + if (!isEqual(activeMcps, mcpValues)) { + setMCPValuesRaw(activeMcps); } - return prev; - }); - }, [mcpValues, setEphemeralAgent]); + } else if (Array.isArray(mcps) && mcps.length === 0 && mcpValues.length > 0) { + // Ephemeral agent explicitly has empty MCP (e.g., spec with no MCP servers) — clear atom + setMCPValuesRaw([]); + } + }, [ephemeralAgent?.mcp, setMCPValuesRaw, configuredServers, mcpValues]); + // Write timestamp when MCP values change useEffect(() => { - const mcpStorageKey = `${LocalStorageKeys.LAST_MCP_}${key}`; + const mcpStorageKey = `${LocalStorageKeys.LAST_MCP_}${mcpAtomKey}`; if (mcpValues.length > 0) { setTimestamp(mcpStorageKey); } - }, [mcpValues, key]); + }, [mcpValues, mcpAtomKey]); - /** Stable memoized setter */ + /** Stable memoized setter with dual-write to environment key */ const setMCPValues = useCallback( (value: string[]) => { if (!Array.isArray(value)) { return; } setMCPValuesRaw(value); + setEphemeralAgent((prev) => { + if (!isEqual(prev?.mcp, value)) { + return { ...(prev ?? {}), mcp: value }; + } + return prev; + }); + // Dual-write to environment key for new conversation defaults + if (storageContextKey) { + const envKey = `${LocalStorageKeys.LAST_MCP_}${storageContextKey}`; + localStorage.setItem(envKey, JSON.stringify(value)); + setTimestamp(envKey); + } }, - [setMCPValuesRaw], + [setMCPValuesRaw, setEphemeralAgent, storageContextKey], ); return { diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index bb5214be7c..fcca040af2 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -28,7 +28,10 @@ export interface MCPServerDefinition { // The init states (isInitializing, isCancellable, etc.) are stored in the global Jotai atom type PollIntervals = Record; -export function useMCPServerManager({ conversationId }: { conversationId?: string | null } = {}) { +export function useMCPServerManager({ + conversationId, + storageContextKey, +}: { conversationId?: string | null; storageContextKey?: string } = {}) { const localize = useLocalize(); const queryClient = useQueryClient(); const { showToast } = useToastContext(); @@ -73,6 +76,7 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin const { mcpValues, setMCPValues, isPinned, setIsPinned } = useMCPSelect({ conversationId, + storageContextKey, servers: selectableServers, }); const mcpValuesRef = useRef(mcpValues); diff --git a/client/src/hooks/Plugins/__tests__/useToolToggle.test.tsx b/client/src/hooks/Plugins/__tests__/useToolToggle.test.tsx new file mode 100644 index 0000000000..f617db2249 --- /dev/null +++ b/client/src/hooks/Plugins/__tests__/useToolToggle.test.tsx @@ -0,0 +1,328 @@ +import React from 'react'; +import { renderHook, act, waitFor } from '@testing-library/react'; +import { LocalStorageKeys, Tools } from 'librechat-data-provider'; +import { RecoilRoot, useRecoilValue, useSetRecoilState } from 'recoil'; +import { ephemeralAgentByConvoId } from '~/store'; +import { useToolToggle } from '../useToolToggle'; + +/** + * Tests for useToolToggle — the hook responsible for toggling tool badges + * (code execution, web search, file search, artifacts) and persisting state. + * + * Desired behaviors: + * - User toggles persist to per-conversation localStorage + * - In non-spec mode with specs configured (storageContextKey = '__defaults__'), + * toggles ALSO persist to the defaults key so future new conversations inherit them + * - In spec mode (storageContextKey = undefined), toggles only persist per-conversation + * - The hook reflects the current ephemeral agent state + */ + +// Mock data-provider auth query +jest.mock('~/data-provider', () => ({ + useVerifyAgentToolAuth: jest.fn().mockReturnValue({ + data: { authenticated: true }, + }), +})); + +// Mock timestamps (track calls without actual localStorage timestamp logic) +jest.mock('~/utils/timestamps', () => ({ + setTimestamp: jest.fn(), +})); + +// Mock useLocalStorageAlt (isPinned state — not relevant to our behavior tests) +jest.mock('~/hooks/useLocalStorageAlt', () => jest.fn(() => [false, jest.fn()])); + +const Wrapper: React.FC<{ children: React.ReactNode }> = ({ children }) => ( + {children} +); + +describe('useToolToggle', () => { + beforeEach(() => { + jest.clearAllMocks(); + localStorage.clear(); + }); + + // ─── Dual-Write Behavior ─────────────────────────────────────────── + + describe('non-spec mode: dual-write to defaults key', () => { + const storageContextKey = '__defaults__'; + + it('should write to both conversation key and defaults key when user toggles a tool', () => { + const conversationId = 'convo-123'; + const { result } = renderHook( + () => + useToolToggle({ + conversationId, + storageContextKey, + toolKey: Tools.execute_code, + localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, + isAuthenticated: true, + }), + { wrapper: Wrapper }, + ); + + act(() => { + result.current.handleChange({ value: true }); + }); + + // Conversation key: per-conversation persistence + const convoKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${conversationId}`; + // Defaults key: persists for future new conversations + const defaultsKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${storageContextKey}`; + + // Sync effect writes to conversation key + expect(localStorage.getItem(convoKey)).toBe(JSON.stringify(true)); + // handleChange dual-writes to defaults key + expect(localStorage.getItem(defaultsKey)).toBe(JSON.stringify(true)); + }); + + it('should persist false values to defaults key when user disables a tool', () => { + const { result } = renderHook( + () => + useToolToggle({ + conversationId: 'convo-456', + storageContextKey, + toolKey: Tools.web_search, + localStorageKey: LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, + isAuthenticated: true, + }), + { wrapper: Wrapper }, + ); + + // Enable then disable + act(() => { + result.current.handleChange({ value: true }); + }); + act(() => { + result.current.handleChange({ value: false }); + }); + + const defaultsKey = `${LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_}${storageContextKey}`; + expect(localStorage.getItem(defaultsKey)).toBe(JSON.stringify(false)); + }); + }); + + describe('spec mode: no dual-write', () => { + it('should only write to conversation key, not to any defaults key', () => { + const conversationId = 'spec-convo-789'; + const { result } = renderHook( + () => + useToolToggle({ + conversationId, + storageContextKey: undefined, // spec mode + toolKey: Tools.execute_code, + localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, + isAuthenticated: true, + }), + { wrapper: Wrapper }, + ); + + act(() => { + result.current.handleChange({ value: true }); + }); + + // Conversation key should have the value + const convoKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${conversationId}`; + expect(localStorage.getItem(convoKey)).toBe(JSON.stringify(true)); + + // Defaults key should NOT have a value + const defaultsKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}__defaults__`; + expect(localStorage.getItem(defaultsKey)).toBeNull(); + }); + }); + + // ─── Per-Conversation Isolation ──────────────────────────────────── + + describe('per-conversation isolation', () => { + it('should maintain separate toggle state per conversation', () => { + const TestComponent = ({ conversationId }: { conversationId: string }) => { + const toggle = useToolToggle({ + conversationId, + toolKey: Tools.execute_code, + localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, + isAuthenticated: true, + }); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(conversationId)); + return { toggle, ephemeralAgent }; + }; + + // Conversation A: enable code + const { result: resultA } = renderHook(() => TestComponent({ conversationId: 'convo-A' }), { + wrapper: Wrapper, + }); + + act(() => { + resultA.current.toggle.handleChange({ value: true }); + }); + + // Conversation B: disable code + const { result: resultB } = renderHook(() => TestComponent({ conversationId: 'convo-B' }), { + wrapper: Wrapper, + }); + + act(() => { + resultB.current.toggle.handleChange({ value: false }); + }); + + // Each conversation has its own value in localStorage + expect(localStorage.getItem(`${LocalStorageKeys.LAST_CODE_TOGGLE_}convo-A`)).toBe('true'); + expect(localStorage.getItem(`${LocalStorageKeys.LAST_CODE_TOGGLE_}convo-B`)).toBe('false'); + }); + }); + + // ─── Ephemeral Agent Sync ────────────────────────────────────────── + + describe('ephemeral agent reflects toggle state', () => { + it('should update ephemeral agent when user toggles a tool', async () => { + const conversationId = 'convo-sync-test'; + const TestComponent = () => { + const toggle = useToolToggle({ + conversationId, + toolKey: Tools.execute_code, + localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, + isAuthenticated: true, + }); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(conversationId)); + return { toggle, ephemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.toggle.handleChange({ value: true }); + }); + + await waitFor(() => { + expect(result.current.ephemeralAgent?.execute_code).toBe(true); + }); + + act(() => { + result.current.toggle.handleChange({ value: false }); + }); + + await waitFor(() => { + expect(result.current.ephemeralAgent?.execute_code).toBe(false); + }); + }); + + it('should reflect external ephemeral agent changes in toolValue', async () => { + const conversationId = 'convo-external'; + const TestComponent = () => { + const toggle = useToolToggle({ + conversationId, + toolKey: Tools.web_search, + localStorageKey: LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, + isAuthenticated: true, + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(conversationId)); + return { toggle, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + // External update (e.g., from applyModelSpecEphemeralAgent) + act(() => { + result.current.setEphemeralAgent({ web_search: true, execute_code: false }); + }); + + await waitFor(() => { + expect(result.current.toggle.toolValue).toBe(true); + expect(result.current.toggle.isToolEnabled).toBe(true); + }); + }); + + it('should sync externally-set ephemeral agent values to localStorage', async () => { + const conversationId = 'convo-sync-ls'; + const TestComponent = () => { + const toggle = useToolToggle({ + conversationId, + toolKey: Tools.file_search, + localStorageKey: LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_, + isAuthenticated: true, + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(conversationId)); + return { toggle, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + // Simulate applyModelSpecEphemeralAgent setting a value + act(() => { + result.current.setEphemeralAgent({ file_search: true }); + }); + + // The sync effect should write to conversation-keyed localStorage + await waitFor(() => { + const storageKey = `${LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_}${conversationId}`; + expect(localStorage.getItem(storageKey)).toBe(JSON.stringify(true)); + }); + }); + }); + + // ─── isToolEnabled computation ───────────────────────────────────── + + describe('isToolEnabled computation', () => { + it('should return false when tool is not set', () => { + const { result } = renderHook( + () => + useToolToggle({ + conversationId: 'convo-1', + toolKey: Tools.execute_code, + localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_, + isAuthenticated: true, + }), + { wrapper: Wrapper }, + ); + + expect(result.current.isToolEnabled).toBe(false); + }); + + it('should treat non-empty string as enabled (artifacts)', async () => { + const conversationId = 'convo-artifacts'; + const TestComponent = () => { + const toggle = useToolToggle({ + conversationId, + toolKey: 'artifacts', + localStorageKey: LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_, + isAuthenticated: true, + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(conversationId)); + return { toggle, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ artifacts: 'default' }); + }); + + await waitFor(() => { + expect(result.current.toggle.isToolEnabled).toBe(true); + }); + }); + + it('should treat empty string as disabled (artifacts off)', async () => { + const conversationId = 'convo-no-artifacts'; + const TestComponent = () => { + const toggle = useToolToggle({ + conversationId, + toolKey: 'artifacts', + localStorageKey: LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_, + isAuthenticated: true, + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(conversationId)); + return { toggle, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ artifacts: '' }); + }); + + await waitFor(() => { + expect(result.current.toggle.isToolEnabled).toBe(false); + }); + }); + }); +}); diff --git a/client/src/hooks/Plugins/useToolToggle.ts b/client/src/hooks/Plugins/useToolToggle.ts index 3b12e87d51..d8026cad1c 100644 --- a/client/src/hooks/Plugins/useToolToggle.ts +++ b/client/src/hooks/Plugins/useToolToggle.ts @@ -13,6 +13,7 @@ type ToolValue = boolean | string; interface UseToolToggleOptions { conversationId?: string | null; + storageContextKey?: string; toolKey: string; localStorageKey: LocalStorageKeys; isAuthenticated?: boolean; @@ -26,6 +27,7 @@ interface UseToolToggleOptions { export function useToolToggle({ conversationId, + storageContextKey, toolKey: _toolKey, localStorageKey, isAuthenticated: externalIsAuthenticated, @@ -93,8 +95,22 @@ export function useToolToggle({ ...(prev || {}), [toolKey]: value, })); + + // Dual-write to environment key for new conversation defaults + if (storageContextKey) { + const envKey = `${localStorageKey}${storageContextKey}`; + localStorage.setItem(envKey, JSON.stringify(value)); + setTimestamp(envKey); + } }, - [setIsDialogOpen, isAuthenticated, setEphemeralAgent, toolKey], + [ + setIsDialogOpen, + isAuthenticated, + setEphemeralAgent, + toolKey, + storageContextKey, + localStorageKey, + ], ); const debouncedChange = useMemo( diff --git a/client/src/utils/__tests__/applyModelSpecEphemeralAgent.test.ts b/client/src/utils/__tests__/applyModelSpecEphemeralAgent.test.ts new file mode 100644 index 0000000000..44bfbb82f7 --- /dev/null +++ b/client/src/utils/__tests__/applyModelSpecEphemeralAgent.test.ts @@ -0,0 +1,274 @@ +import { Constants, LocalStorageKeys } from 'librechat-data-provider'; +import type { TModelSpec, TEphemeralAgent } from 'librechat-data-provider'; +import { applyModelSpecEphemeralAgent } from '../endpoints'; +import { setTimestamp } from '../timestamps'; + +/** + * Tests for applyModelSpecEphemeralAgent — the function responsible for + * constructing the ephemeral agent state when navigating to a spec conversation. + * + * Desired behaviors: + * - New conversations always get the admin's exact spec configuration + * - Existing conversations merge per-conversation localStorage overrides on top of spec + * - Cleared localStorage for existing conversations falls back to fresh spec config + */ + +const createModelSpec = (overrides: Partial = {}): TModelSpec => + ({ + name: 'test-spec', + label: 'Test Spec', + preset: { endpoint: 'agents' }, + mcpServers: ['spec-server1'], + webSearch: true, + executeCode: true, + fileSearch: false, + artifacts: true, + ...overrides, + }) as TModelSpec; + +/** Write a value + fresh timestamp to localStorage (simulates a user toggle) */ +function writeToolToggle(storagePrefix: string, convoId: string, value: unknown): void { + const key = `${storagePrefix}${convoId}`; + localStorage.setItem(key, JSON.stringify(value)); + setTimestamp(key); +} + +describe('applyModelSpecEphemeralAgent', () => { + let updateEphemeralAgent: jest.Mock; + + beforeEach(() => { + localStorage.clear(); + updateEphemeralAgent = jest.fn(); + }); + + // ─── New Conversations ───────────────────────────────────────────── + + describe('new conversations always get fresh admin spec config', () => { + it('should apply exactly the admin-configured tools and MCP servers', () => { + const modelSpec = createModelSpec({ + mcpServers: ['clickhouse', 'github'], + executeCode: true, + webSearch: false, + fileSearch: true, + artifacts: true, + }); + + applyModelSpecEphemeralAgent({ + convoId: null, + modelSpec, + updateEphemeralAgent, + }); + + expect(updateEphemeralAgent).toHaveBeenCalledWith(Constants.NEW_CONVO, { + mcp: ['clickhouse', 'github'], + execute_code: true, + web_search: false, + file_search: true, + artifacts: 'default', + }); + }); + + it('should not read from localStorage even if stale values exist', () => { + // Simulate stale localStorage from a previous session + writeToolToggle(LocalStorageKeys.LAST_CODE_TOGGLE_, Constants.NEW_CONVO, false); + writeToolToggle(LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, Constants.NEW_CONVO, true); + localStorage.setItem( + `${LocalStorageKeys.LAST_MCP_}${Constants.NEW_CONVO}`, + JSON.stringify(['stale-server']), + ); + + const modelSpec = createModelSpec({ executeCode: true, webSearch: false, mcpServers: [] }); + + applyModelSpecEphemeralAgent({ + convoId: null, + modelSpec, + updateEphemeralAgent, + }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + // Should be spec values, NOT localStorage values + expect(agent.execute_code).toBe(true); + expect(agent.web_search).toBe(false); + expect(agent.mcp).toEqual([]); + }); + + it('should handle spec with no MCP servers', () => { + const modelSpec = createModelSpec({ mcpServers: undefined }); + + applyModelSpecEphemeralAgent({ convoId: null, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.mcp).toEqual([]); + }); + + it('should map artifacts: true to "default" string', () => { + const modelSpec = createModelSpec({ artifacts: true }); + + applyModelSpecEphemeralAgent({ convoId: null, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.artifacts).toBe('default'); + }); + + it('should pass through artifacts string value directly', () => { + const modelSpec = createModelSpec({ artifacts: 'custom-renderer' as any }); + + applyModelSpecEphemeralAgent({ convoId: null, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.artifacts).toBe('custom-renderer'); + }); + }); + + // ─── Existing Conversations: Per-Conversation Persistence ────────── + + describe('existing conversations merge user overrides from localStorage', () => { + const convoId = 'convo-abc-123'; + + it('should preserve user tool modifications across navigation', () => { + // User previously toggled off code execution and enabled file search + writeToolToggle(LocalStorageKeys.LAST_CODE_TOGGLE_, convoId, false); + writeToolToggle(LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_, convoId, true); + + const modelSpec = createModelSpec({ + executeCode: true, + fileSearch: false, + webSearch: true, + }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.execute_code).toBe(false); // user override + expect(agent.file_search).toBe(true); // user override + expect(agent.web_search).toBe(true); // not overridden, spec value + }); + + it('should preserve user-added MCP servers across navigation', () => { + // Spec has clickhouse, user also added github during the conversation + localStorage.setItem( + `${LocalStorageKeys.LAST_MCP_}${convoId}`, + JSON.stringify(['clickhouse', 'github']), + ); + + const modelSpec = createModelSpec({ mcpServers: ['clickhouse'] }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.mcp).toEqual(['clickhouse', 'github']); + }); + + it('should preserve user-removed MCP servers (empty array) across navigation', () => { + // User removed all MCP servers during the conversation + localStorage.setItem(`${LocalStorageKeys.LAST_MCP_}${convoId}`, JSON.stringify([])); + + const modelSpec = createModelSpec({ mcpServers: ['clickhouse', 'github'] }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.mcp).toEqual([]); + }); + + it('should only override keys that exist in localStorage, leaving the rest as spec defaults', () => { + // User only changed artifacts, nothing else + writeToolToggle(LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_, convoId, ''); + + const modelSpec = createModelSpec({ + executeCode: true, + webSearch: true, + fileSearch: false, + artifacts: true, + mcpServers: ['server1'], + }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + expect(agent.execute_code).toBe(true); // spec default (not in localStorage) + expect(agent.web_search).toBe(true); // spec default + expect(agent.file_search).toBe(false); // spec default + expect(agent.artifacts).toBe(''); // user override + expect(agent.mcp).toEqual(['server1']); // spec default (not in localStorage) + }); + }); + + // ─── Existing Conversations: Cleared localStorage ────────────────── + + describe('existing conversations with cleared localStorage get fresh spec config', () => { + const convoId = 'convo-cleared-456'; + + it('should fall back to pure spec values when localStorage is empty', () => { + // localStorage.clear() was already called in beforeEach + + const modelSpec = createModelSpec({ + executeCode: true, + webSearch: false, + fileSearch: true, + artifacts: true, + mcpServers: ['server1', 'server2'], + }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + expect(updateEphemeralAgent).toHaveBeenCalledWith(convoId, { + mcp: ['server1', 'server2'], + execute_code: true, + web_search: false, + file_search: true, + artifacts: 'default', + }); + }); + + it('should fall back to spec values when timestamps have expired (>2 days)', () => { + // Write values with expired timestamps (3 days old) + const expiredTimestamp = (Date.now() - 3 * 24 * 60 * 60 * 1000).toString(); + const codeKey = `${LocalStorageKeys.LAST_CODE_TOGGLE_}${convoId}`; + localStorage.setItem(codeKey, JSON.stringify(false)); + localStorage.setItem(`${codeKey}_TIMESTAMP`, expiredTimestamp); + + const modelSpec = createModelSpec({ executeCode: true }); + + applyModelSpecEphemeralAgent({ convoId, modelSpec, updateEphemeralAgent }); + + const agent = updateEphemeralAgent.mock.calls[0][1] as TEphemeralAgent; + // Expired override should be ignored — spec value wins + expect(agent.execute_code).toBe(true); + }); + }); + + // ─── Guard Clauses ───────────────────────────────────────────────── + + describe('guard clauses', () => { + it('should not call updateEphemeralAgent when modelSpec is undefined', () => { + applyModelSpecEphemeralAgent({ + convoId: 'convo-1', + modelSpec: undefined, + updateEphemeralAgent, + }); + + expect(updateEphemeralAgent).not.toHaveBeenCalled(); + }); + + it('should not throw when updateEphemeralAgent is undefined', () => { + expect(() => + applyModelSpecEphemeralAgent({ + convoId: 'convo-1', + modelSpec: createModelSpec(), + updateEphemeralAgent: undefined, + }), + ).not.toThrow(); + }); + + it('should use NEW_CONVO key when convoId is empty string', () => { + applyModelSpecEphemeralAgent({ + convoId: '', + modelSpec: createModelSpec(), + updateEphemeralAgent, + }); + + expect(updateEphemeralAgent).toHaveBeenCalledWith(Constants.NEW_CONVO, expect.any(Object)); + }); + }); +}); diff --git a/client/src/utils/endpoints.ts b/client/src/utils/endpoints.ts index eb9e60386f..33aa7a8525 100644 --- a/client/src/utils/endpoints.ts +++ b/client/src/utils/endpoints.ts @@ -11,6 +11,7 @@ import { } from 'librechat-data-provider'; import type * as t from 'librechat-data-provider'; import type { LocalizeFunction, IconsRecord } from '~/common'; +import { getTimestampedValue } from './timestamps'; /** * Clears model for non-ephemeral agent conversations. @@ -219,12 +220,51 @@ export function applyModelSpecEphemeralAgent({ if (!modelSpec || !updateEphemeralAgent) { return; } - updateEphemeralAgent((convoId ?? Constants.NEW_CONVO) || Constants.NEW_CONVO, { - mcp: modelSpec.mcpServers ?? [Constants.mcp_clear as string], + const key = (convoId ?? Constants.NEW_CONVO) || Constants.NEW_CONVO; + const agent: t.TEphemeralAgent = { + mcp: modelSpec.mcpServers ?? [], web_search: modelSpec.webSearch ?? false, file_search: modelSpec.fileSearch ?? false, execute_code: modelSpec.executeCode ?? false, - }); + artifacts: modelSpec.artifacts === true ? 'default' : modelSpec.artifacts || '', + }; + + // For existing conversations, layer per-conversation localStorage overrides + // on top of spec defaults so user modifications persist across navigation. + // If localStorage is empty (e.g., cleared), spec values stand alone. + if (key !== Constants.NEW_CONVO) { + const toolStorageMap: Array<[keyof t.TEphemeralAgent, string]> = [ + ['execute_code', LocalStorageKeys.LAST_CODE_TOGGLE_], + ['web_search', LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_], + ['file_search', LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_], + ['artifacts', LocalStorageKeys.LAST_ARTIFACTS_TOGGLE_], + ]; + + for (const [toolKey, storagePrefix] of toolStorageMap) { + const raw = getTimestampedValue(`${storagePrefix}${key}`); + if (raw !== null) { + try { + agent[toolKey] = JSON.parse(raw) as never; + } catch { + // ignore parse errors + } + } + } + + const mcpRaw = localStorage.getItem(`${LocalStorageKeys.LAST_MCP_}${key}`); + if (mcpRaw !== null) { + try { + const parsed = JSON.parse(mcpRaw); + if (Array.isArray(parsed)) { + agent.mcp = parsed; + } + } catch { + // ignore parse errors + } + } + } + + updateEphemeralAgent(key, agent); } /** diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index f6567e8da9..02174b6496 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1758,6 +1758,8 @@ export enum Constants { mcp_all = 'sys__all__sys', /** Unique value to indicate clearing MCP servers from UI state. For frontend use only. */ mcp_clear = 'sys__clear__sys', + /** Key suffix for non-spec user default tool storage */ + spec_defaults_key = '__defaults__', /** * Unique value to indicate the MCP tool was added to an agent. * This helps inform the UI if the mcp server was previously added. diff --git a/packages/data-provider/src/models.ts b/packages/data-provider/src/models.ts index 3c3c197660..c2dbe2cf77 100644 --- a/packages/data-provider/src/models.ts +++ b/packages/data-provider/src/models.ts @@ -35,6 +35,7 @@ export type TModelSpec = { webSearch?: boolean; fileSearch?: boolean; executeCode?: boolean; + artifacts?: string | boolean; mcpServers?: string[]; }; @@ -54,6 +55,7 @@ export const tModelSpecSchema = z.object({ webSearch: z.boolean().optional(), fileSearch: z.boolean().optional(), executeCode: z.boolean().optional(), + artifacts: z.union([z.string(), z.boolean()]).optional(), mcpServers: z.array(z.string()).optional(), }); diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 1198f97b80..a7782a3bc6 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -99,6 +99,7 @@ export type TEphemeralAgent = { web_search?: boolean; file_search?: boolean; execute_code?: boolean; + artifacts?: string; }; export type TPayload = Partial &