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
This commit is contained in:
Faisal Amir 2025-07-29 09:50:07 +07:00 committed by GitHub
parent 160d158152
commit 63cb4fbf3b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 115 additions and 26 deletions

View File

@ -18,4 +18,5 @@ export const localStorageKey = {
toolAvailability: 'tool-availability',
mcpGlobalPermissions: 'mcp-global-permissions',
lastUsedModel: 'last-used-model',
lastUsedAssistant: 'last-used-assistant',
}

View File

@ -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' }]
}))
}
)

View File

@ -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(),

View File

@ -35,6 +35,12 @@ export const useAppState = create<AppState>()((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<AppState>()((set) => ({
created_at: content.created_at || Date.now(),
metadata: {
...content.metadata,
assistant: useAssistant.getState().currentAssistant,
assistant: selectedAssistant,
},
}
: undefined,

View File

@ -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<AssistantState>()((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)
}
}
},
}))

View File

@ -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,
}
)

View File

@ -28,19 +28,18 @@ export const useMessages = create<MessageState>()((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) => {

View File

@ -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) => {