From 63cb4fbf3bb5310e51ea843694043b9e86cd7d1f Mon Sep 17 00:00:00 2001 From: Faisal Amir Date: Tue, 29 Jul 2025 09:50:07 +0700 Subject: [PATCH] fix: assistant with last used and fix metadata (#5955) * fix: assistant with last used and fix metadata * chore: revert instruction and desc * chore: fix current assistant state * chore: updae metadata message assistant * chore: update test case --- web-app/src/constants/localStorage.ts | 1 + .../src/hooks/__tests__/useAppState.test.ts | 3 +- web-app/src/hooks/__tests__/useChat.test.ts | 37 ++++++++--- web-app/src/hooks/useAppState.ts | 8 ++- web-app/src/hooks/useAssistant.ts | 65 +++++++++++++++++-- web-app/src/hooks/useChat.ts | 11 +++- web-app/src/hooks/useMessages.ts | 13 ++-- web-app/src/providers/DataProvider.tsx | 3 +- 8 files changed, 115 insertions(+), 26 deletions(-) diff --git a/web-app/src/constants/localStorage.ts b/web-app/src/constants/localStorage.ts index 6c968c324..cef01b0fd 100644 --- a/web-app/src/constants/localStorage.ts +++ b/web-app/src/constants/localStorage.ts @@ -18,4 +18,5 @@ export const localStorageKey = { toolAvailability: 'tool-availability', mcpGlobalPermissions: 'mcp-global-permissions', lastUsedModel: 'last-used-model', + lastUsedAssistant: 'last-used-assistant', } diff --git a/web-app/src/hooks/__tests__/useAppState.test.ts b/web-app/src/hooks/__tests__/useAppState.test.ts index 6d6567774..a677b6564 100644 --- a/web-app/src/hooks/__tests__/useAppState.test.ts +++ b/web-app/src/hooks/__tests__/useAppState.test.ts @@ -11,7 +11,8 @@ vi.mock('../useAssistant', () => ({ })), { getState: vi.fn(() => ({ - currentAssistant: { id: 'test-assistant', name: 'Test Assistant' } + currentAssistant: { id: 'test-assistant', name: 'Test Assistant' }, + assistants: [{ id: 'test-assistant', name: 'Test Assistant' }] })) } ) diff --git a/web-app/src/hooks/__tests__/useChat.test.ts b/web-app/src/hooks/__tests__/useChat.test.ts index ea54aba13..67f86b5a3 100644 --- a/web-app/src/hooks/__tests__/useChat.test.ts +++ b/web-app/src/hooks/__tests__/useChat.test.ts @@ -11,19 +11,31 @@ vi.mock('../usePrompt', () => ({ })) vi.mock('../useAppState', () => ({ - useAppState: vi.fn(() => ({ - tools: [], - updateTokenSpeed: vi.fn(), - resetTokenSpeed: vi.fn(), - updateTools: vi.fn(), - updateStreamingContent: vi.fn(), - updateLoadingModel: vi.fn(), - setAbortController: vi.fn(), - })), + useAppState: Object.assign( + vi.fn(() => ({ + tools: [], + updateTokenSpeed: vi.fn(), + resetTokenSpeed: vi.fn(), + updateTools: vi.fn(), + updateStreamingContent: vi.fn(), + updateLoadingModel: vi.fn(), + setAbortController: vi.fn(), + })), + { + getState: vi.fn(() => ({ + tokenSpeed: { tokensPerSecond: 10 }, + })) + } + ), })) vi.mock('../useAssistant', () => ({ useAssistant: vi.fn(() => ({ + assistants: [{ + id: 'test-assistant', + instructions: 'test instructions', + parameters: { stream: true }, + }], currentAssistant: { id: 'test-assistant', instructions: 'test instructions', @@ -88,6 +100,12 @@ vi.mock('../useModelContextApproval', () => ({ })), })) +vi.mock('../useModelLoad', () => ({ + useModelLoad: vi.fn(() => ({ + setModelLoadError: vi.fn(), + })), +})) + vi.mock('@tanstack/react-router', () => ({ useRouter: vi.fn(() => ({ navigate: vi.fn(), @@ -96,6 +114,7 @@ vi.mock('@tanstack/react-router', () => ({ vi.mock('@/lib/completion', () => ({ emptyThreadContent: { thread_id: 'test-thread', content: '' }, + extractToolCall: vi.fn(), newUserThreadContent: vi.fn(() => ({ thread_id: 'test-thread', content: 'user message' })), newAssistantThreadContent: vi.fn(() => ({ thread_id: 'test-thread', content: 'assistant message' })), sendCompletion: vi.fn(), diff --git a/web-app/src/hooks/useAppState.ts b/web-app/src/hooks/useAppState.ts index dc29f7f8a..eade55af5 100644 --- a/web-app/src/hooks/useAppState.ts +++ b/web-app/src/hooks/useAppState.ts @@ -35,6 +35,12 @@ export const useAppState = create()((set) => ({ tokenSpeed: undefined, currentToolCall: undefined, updateStreamingContent: (content: ThreadMessage | undefined) => { + const assistants = useAssistant.getState().assistants + const currentAssistant = useAssistant.getState().currentAssistant + + const selectedAssistant = + assistants.find((a) => a.id === currentAssistant.id) || assistants[0] + set(() => ({ streamingContent: content ? { @@ -42,7 +48,7 @@ export const useAppState = create()((set) => ({ created_at: content.created_at || Date.now(), metadata: { ...content.metadata, - assistant: useAssistant.getState().currentAssistant, + assistant: selectedAssistant, }, } : undefined, diff --git a/web-app/src/hooks/useAssistant.ts b/web-app/src/hooks/useAssistant.ts index 1813a9c7b..51cb1f97f 100644 --- a/web-app/src/hooks/useAssistant.ts +++ b/web-app/src/hooks/useAssistant.ts @@ -1,6 +1,7 @@ import { createAssistant, deleteAssistant } from '@/services/assistants' import { Assistant as CoreAssistant } from '@janhq/core' import { create } from 'zustand' +import { localStorageKey } from '@/constants/localStorage' interface AssistantState { assistants: Assistant[] @@ -8,8 +9,29 @@ interface AssistantState { addAssistant: (assistant: Assistant) => void updateAssistant: (assistant: Assistant) => void deleteAssistant: (id: string) => void - setCurrentAssistant: (assistant: Assistant) => void + setCurrentAssistant: (assistant: Assistant, saveToStorage?: boolean) => void setAssistants: (assistants: Assistant[]) => void + getLastUsedAssistant: () => string | null + setLastUsedAssistant: (assistantId: string) => void + initializeWithLastUsed: () => void +} + +// Helper functions for localStorage +const getLastUsedAssistantId = (): string | null => { + try { + return localStorage.getItem(localStorageKey.lastUsedAssistant) + } catch (error) { + console.debug('Failed to get last used assistant from localStorage:', error) + return null + } +} + +const setLastUsedAssistantId = (assistantId: string) => { + try { + localStorage.setItem(localStorageKey.lastUsedAssistant, assistantId) + } catch (error) { + console.debug('Failed to set last used assistant in localStorage:', error) + } } export const defaultAssistant: Assistant = { @@ -51,17 +73,52 @@ export const useAssistant = create()((set, get) => ({ }) }, deleteAssistant: (id) => { + const state = get() deleteAssistant( - get().assistants.find((e) => e.id === id) as unknown as CoreAssistant + state.assistants.find((e) => e.id === id) as unknown as CoreAssistant ).catch((error) => { console.error('Failed to delete assistant:', error) }) - set({ assistants: get().assistants.filter((a) => a.id !== id) }) + + // Check if we're deleting the current assistant + const wasCurrentAssistant = state.currentAssistant.id === id + + set({ assistants: state.assistants.filter((a) => a.id !== id) }) + + // If the deleted assistant was current, fallback to default and update localStorage + if (wasCurrentAssistant) { + set({ currentAssistant: defaultAssistant }) + setLastUsedAssistantId(defaultAssistant.id) + } }, - setCurrentAssistant: (assistant) => { + setCurrentAssistant: (assistant, saveToStorage = true) => { set({ currentAssistant: assistant }) + if (saveToStorage) { + setLastUsedAssistantId(assistant.id) + } }, setAssistants: (assistants) => { set({ assistants }) }, + getLastUsedAssistant: () => { + return getLastUsedAssistantId() + }, + setLastUsedAssistant: (assistantId) => { + setLastUsedAssistantId(assistantId) + }, + initializeWithLastUsed: () => { + const lastUsedId = getLastUsedAssistantId() + if (lastUsedId) { + const lastUsedAssistant = get().assistants.find( + (a) => a.id === lastUsedId + ) + if (lastUsedAssistant) { + set({ currentAssistant: lastUsedAssistant }) + } else { + // Fallback to default if last used assistant was deleted + set({ currentAssistant: defaultAssistant }) + setLastUsedAssistantId(defaultAssistant.id) + } + } + }, })) diff --git a/web-app/src/hooks/useChat.ts b/web-app/src/hooks/useChat.ts index e52129145..11e12bec1 100644 --- a/web-app/src/hooks/useChat.ts +++ b/web-app/src/hooks/useChat.ts @@ -45,7 +45,7 @@ export const useChat = () => { updateLoadingModel, setAbortController, } = useAppState() - const { currentAssistant } = useAssistant() + const { assistants, currentAssistant } = useAssistant() const { updateProvider } = useModelProvider() const { approvedTools, showApprovalModal, allowAllMCPPermissions } = @@ -74,6 +74,9 @@ export const useChat = () => { return provider?.provider || selectedProvider }, [provider, selectedProvider]) + const selectedAssistant = + assistants.find((a) => a.id === currentAssistant.id) || assistants[0] + useEffect(() => { function setTools() { getTools().then((data: MCPTool[]) => { @@ -92,6 +95,7 @@ export const useChat = () => { const getCurrentThread = useCallback(async () => { let currentThread = retrieveThread() + if (!currentThread) { currentThread = await createThread( { @@ -99,7 +103,7 @@ export const useChat = () => { provider: selectedProvider, }, prompt, - currentAssistant + selectedAssistant ) router.navigate({ to: route.threadsDetail, @@ -114,7 +118,7 @@ export const useChat = () => { router, selectedModel?.id, selectedProvider, - currentAssistant, + selectedAssistant, ]) const restartModel = useCallback( @@ -402,6 +406,7 @@ export const useChat = () => { accumulatedText, { tokenSpeed: useAppState.getState().tokenSpeed, + assistant: currentAssistant, } ) diff --git a/web-app/src/hooks/useMessages.ts b/web-app/src/hooks/useMessages.ts index 3841a50b2..bead31641 100644 --- a/web-app/src/hooks/useMessages.ts +++ b/web-app/src/hooks/useMessages.ts @@ -28,19 +28,18 @@ export const useMessages = create()((set, get) => ({ })) }, addMessage: (message) => { + const assistants = useAssistant.getState().assistants const currentAssistant = useAssistant.getState().currentAssistant + + const selectedAssistant = + assistants.find((a) => a.id === currentAssistant.id) || assistants[0] + const newMessage = { ...message, created_at: message.created_at || Date.now(), metadata: { ...message.metadata, - assistant: { - id: currentAssistant?.id || '', - name: currentAssistant?.name || '', - avatar: currentAssistant?.avatar || '', - instructions: currentAssistant?.instructions || '', - parameters: currentAssistant?.parameters || '', - }, + assistant: selectedAssistant, }, } createMessage(newMessage).then((createdMessage) => { diff --git a/web-app/src/providers/DataProvider.tsx b/web-app/src/providers/DataProvider.tsx index 4141b4ec3..9fdb802ea 100644 --- a/web-app/src/providers/DataProvider.tsx +++ b/web-app/src/providers/DataProvider.tsx @@ -24,7 +24,7 @@ export function DataProvider() { const { setMessages } = useMessages() const { checkForUpdate } = useAppUpdater() const { setServers } = useMCPServers() - const { setAssistants } = useAssistant() + const { setAssistants, initializeWithLastUsed } = useAssistant() const { setThreads } = useThreads() const navigate = useNavigate() @@ -37,6 +37,7 @@ export function DataProvider() { // Only update assistants if we have valid data if (data && Array.isArray(data) && data.length > 0) { setAssistants(data as unknown as Assistant[]) + initializeWithLastUsed() } }) .catch((error) => {