🐛 fix(conversation): preserve mention runtime context (#13223)

* 🐛 fix(conversation): preserve mention context on retry

* 🐛 fix(runtime): preserve initial payload for mention context

*  feat(store): expose Zustand stores on window.__LOBE_STORES in dev

Made-with: Cursor
This commit is contained in:
Innei
2026-03-24 19:50:26 +08:00
committed by GitHub
parent 72ba8c8923
commit 995d5ea354
28 changed files with 256 additions and 9 deletions

View File

@@ -1,3 +1,4 @@
import { AgentManagementIdentifier } from '@lobechat/builtin-tool-agent-management';
import { act } from '@testing-library/react';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
@@ -738,6 +739,77 @@ describe('Generation Actions', () => {
);
});
it('should restore mention-based initialContext when regenerating a user message', async () => {
const { useChatStore } = await import('@/store/chat');
vi.mocked(useChatStore.getState).mockReturnValue({
messagesMap: {},
operations: {},
messageLoadingIds: [],
cancelOperations: mockCancelOperations,
cancelOperation: mockCancelOperation,
deleteMessage: mockDeleteMessage,
switchMessageBranch: mockSwitchMessageBranch,
startOperation: mockStartOperation,
completeOperation: mockCompleteOperation,
failOperation: mockFailOperation,
internal_execAgentRuntime: mockInternalExecAgentRuntime,
} as any);
const context: ConversationContext = {
agentId: 'session-1',
topicId: 'topic-1',
threadId: null,
};
const store = createStore({ context });
act(() => {
store.setState({
displayMessages: [
{
id: 'msg-1',
role: 'user',
content: '<mention name="Agent A" id="agent-a" /> hello',
editorData: {
root: {
type: 'root',
children: [
{
type: 'paragraph',
children: [
{
type: 'mention',
label: 'Agent A',
metadata: { id: 'agent-a', type: 'agent' },
},
{ type: 'text', text: ' hello' },
],
},
],
},
},
},
],
} as any);
});
await act(async () => {
await store.getState().regenerateUserMessage('msg-1');
});
expect(mockInternalExecAgentRuntime).toHaveBeenCalledWith(
expect.objectContaining({
initialContext: {
initialContext: {
mentionedAgents: [{ id: 'agent-a', name: 'Agent A' }],
selectedTools: [{ identifier: AgentManagementIdentifier, name: 'Agent Management' }],
},
phase: 'init',
},
}),
);
});
it('should not regenerate if message is already loading', async () => {
// Mock messageLoadingIds to include the target message
const { useChatStore } = await import('@/store/chat');

View File

@@ -1,11 +1,46 @@
import { AgentManagementIdentifier } from '@lobechat/builtin-tool-agent-management';
import { type StateCreator } from 'zustand';
import { MESSAGE_CANCEL_FLAT } from '@/const/index';
import { useChatStore } from '@/store/chat';
import {
parseMentionedAgentsFromEditorData,
parseSelectedSkillsFromEditorData,
parseSelectedToolsFromEditorData,
} from '@/store/chat/slices/aiChat/actions/commandBus';
import { INPUT_LOADING_OPERATION_TYPES } from '@/store/chat/slices/operation/types';
import { type Store as ConversationStore } from '../../action';
const buildRetryInitialContext = (editorData: Record<string, any> | null | undefined) => {
const normalizedEditorData = editorData ?? undefined;
const selectedSkills = parseSelectedSkillsFromEditorData(normalizedEditorData);
const selectedTools = parseSelectedToolsFromEditorData(normalizedEditorData);
const mentionedAgents = parseMentionedAgentsFromEditorData(normalizedEditorData);
const effectiveSelectedTools =
mentionedAgents.length > 0 &&
!selectedTools.some((tool) => tool.identifier === AgentManagementIdentifier)
? [...selectedTools, { identifier: AgentManagementIdentifier, name: 'Agent Management' }]
: selectedTools;
const hasInitialContext =
effectiveSelectedTools.length > 0 || selectedSkills.length > 0 || mentionedAgents.length > 0;
if (!hasInitialContext) return undefined;
return {
initialContext: {
...(selectedSkills.length > 0 ? { selectedSkills } : undefined),
...(effectiveSelectedTools.length > 0
? { selectedTools: effectiveSelectedTools }
: undefined),
...(mentionedAgents.length > 0 ? { mentionedAgents } : undefined),
},
phase: 'init' as const,
};
};
/**
* Generation Actions
*
@@ -287,6 +322,7 @@ export const generationSlice: StateCreator<
const currentIndex = displayMessages.findIndex((c) => c.id === messageId);
const item = displayMessages[currentIndex];
if (!item) return;
const initialContext = buildRetryInitialContext(item.editorData);
// Get context messages up to and including the target message
const contextMessages = displayMessages.slice(0, currentIndex + 1);
@@ -320,6 +356,7 @@ export const generationSlice: StateCreator<
// Execute agent runtime with full context from ConversationStore
await chatStore.internal_execAgentRuntime({
context,
initialContext,
messages: contextMessages,
parentMessageId: messageId,
parentMessageType: 'user',

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type AgentStoreState } from './initialState';
import { initialState } from './initialState';
@@ -58,4 +59,6 @@ const devtools = createDevtools('agent');
export const useAgentStore = createWithEqualityFn<AgentStore>()(devtools(createStore), shallow);
expose('agent', useAgentStore);
export const getAgentStoreState = () => useAgentStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { type ChatGroupAction } from './action';
import { chatGroupAction } from './action';
import { type ChatGroupState } from './initialState';
@@ -24,4 +25,6 @@ export const useAgentGroupStore = createWithEqualityFn<ChatGroupStore>()(
shallow,
);
expose('agentGroup', useAgentGroupStore);
export const getChatGroupStoreState = () => useAgentGroupStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type AIProviderStoreState } from './initialState';
import { initialState } from './initialState';
@@ -34,4 +35,6 @@ const devtools = createDevtools('aiInfra');
export const useAiInfraStore = createWithEqualityFn<AiInfraStore>()(devtools(createStore), shallow);
expose('aiInfra', useAiInfraStore);
export const getAiInfraStoreState = () => useAiInfraStore.getState();

View File

@@ -899,6 +899,59 @@ describe('StreamingExecutor actions', () => {
);
});
it('should preserve default model/provider payload when initialContext is provided', () => {
act(() => {
useChatStore.setState({ internal_execAgentRuntime: realExecAgentRuntime });
});
const { result } = renderHook(() => useChatStore());
const userMessage = {
id: TEST_IDS.USER_MESSAGE_ID,
role: 'user',
content: TEST_CONTENT.USER_MESSAGE,
sessionId: TEST_IDS.SESSION_ID,
topicId: TEST_IDS.TOPIC_ID,
} as UIChatMessage;
vi.spyOn(agentConfigResolver, 'resolveAgentConfig').mockReturnValue({
agentConfig: createMockAgentConfig({
model: 'claude-sonnet-4-6',
provider: 'lobehub',
}),
chatConfig: createMockChatConfig(),
isBuiltinAgent: false,
plugins: [],
});
vi.spyOn(toolEngineering, 'createAgentToolsEngine').mockReturnValue({
generateToolsDetailed: vi.fn().mockReturnValue({
enabledManifests: [],
enabledToolIds: [],
tools: [],
}),
} as any);
const { context } = result.current.internal_createAgentState({
messages: [userMessage],
parentMessageId: userMessage.id,
agentId: TEST_IDS.SESSION_ID,
topicId: TEST_IDS.TOPIC_ID,
initialContext: {
phase: 'init',
initialContext: {
selectedTools: [{ identifier: 'lobe-notebook', name: 'Notebook' }],
},
},
});
expect(context.payload).toEqual(
expect.objectContaining({
model: 'claude-sonnet-4-6',
parentMessageId: TEST_IDS.USER_MESSAGE_ID,
provider: 'lobehub',
}),
);
});
it('should pass merged resolvedAgentConfig to chatService when selectedTools are provided', async () => {
act(() => {
useChatStore.setState({ internal_execAgentRuntime: realExecAgentRuntime });

View File

@@ -287,19 +287,29 @@ export class StreamingExecutorActionImpl {
}
: undefined;
const defaultPayload = {
model: agentConfigData.model,
parentMessageId,
provider: agentConfigData.provider,
};
const existingPayload =
initialContext?.payload && typeof initialContext.payload === 'object'
? (initialContext.payload as Record<string, unknown>)
: undefined;
// Create initial context or use provided context
const context: AgentRuntimeContext = initialContext
? {
...initialContext,
payload: {
...defaultPayload,
...existingPayload,
},
initialContext: mergedRuntimeInitialContext,
}
: {
phase: 'init',
payload: {
model: agentConfigData.model,
provider: agentConfigData.provider,
parentMessageId,
},
payload: defaultPayload,
session: {
sessionId: agentId,
messageCount: messages.length,

View File

@@ -5,6 +5,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type ChatStoreState } from './initialState';
import { initialState } from './initialState';
@@ -76,8 +77,6 @@ export const useChatStore = createWithEqualityFn<ChatStore>()(
shallow,
);
if (typeof window !== 'undefined') {
window.__CHAT_STORE__ = useChatStore;
}
expose('chat', useChatStore);
export const getChatStoreState = () => useChatStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type AssistantAction } from './slices/assistant/action';
import { createAssistantSlice } from './slices/assistant/action';
@@ -69,4 +70,6 @@ export const useDiscoverStore = createWithEqualityFn<DiscoverStore>()(
shallow,
);
expose('discover', useDiscoverStore);
export const getDiscoverStoreState = () => useDiscoverStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type DocumentAction } from './slices/document';
import { createDocumentSlice } from './slices/document';
@@ -40,4 +41,6 @@ export const useDocumentStore = createWithEqualityFn<DocumentStore>()(
shallow,
);
expose('document', useDocumentStore);
export const getDocumentStoreState = () => useDocumentStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type ElectronAppAction } from './actions/app';
import { createElectronAppSlice } from './actions/app';
@@ -63,4 +64,6 @@ export const useElectronStore = createWithEqualityFn<ElectronStore>()(
shallow,
);
expose('electron', useElectronStore);
export const getElectronStoreState = () => useElectronStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import type { StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { type EvalStoreState, initialState } from './initialState';
import { type BenchmarkAction, createBenchmarkSlice } from './slices/benchmark/action';
import { createDatasetSlice, type DatasetAction } from './slices/dataset/action';
@@ -26,3 +27,5 @@ const createStore: StateCreator<EvalStore, [['zustand/devtools', never]]> = (set
const devtools = createDevtools('eval');
export const useEvalStore = createWithEqualityFn<EvalStore>()(devtools(createStore), shallow);
expose('eval', useEvalStore);

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type FilesStoreState } from './initialState';
import { initialState } from './initialState';
@@ -62,4 +63,6 @@ const devtools = createDevtools('file');
export const useFileStore = createWithEqualityFn<FileStore>()(devtools(createStore), shallow);
expose('file', useFileStore);
export const getFileStoreState = () => useFileStore.getState();

View File

@@ -4,6 +4,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type GlobalGeneralAction } from './actions/general';
import { generalActionSlice } from './actions/general';
@@ -38,3 +39,5 @@ export const useGlobalStore = createWithEqualityFn<GlobalStore>()(
subscribeWithSelector(devtools(createStore)),
shallow,
);
expose('global', useGlobalStore);

View File

@@ -6,6 +6,7 @@ import { type StateCreator } from 'zustand/vanilla';
import { isDev } from '@/utils/env';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type HomeStoreState } from './initialState';
import { initialState } from './initialState';
@@ -62,4 +63,6 @@ export const useHomeStore = createWithEqualityFn<HomeStore>()(
shallow,
);
expose('home', useHomeStore);
export const getHomeStoreState = () => useHomeStore.getState();

View File

@@ -4,6 +4,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type ImageStoreState } from './initialState';
import { initialState } from './initialState';
@@ -52,4 +53,6 @@ export const useImageStore = createWithEqualityFn<ImageStore>()(
shallow,
);
expose('image', useImageStore);
export const getImageStoreState = () => useImageStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type KnowledgeBaseStoreState } from './initialState';
import { initialState } from './initialState';
@@ -46,3 +47,5 @@ export const useKnowledgeBaseStore = createWithEqualityFn<KnowledgeBaseStore>()(
devtools(createStore),
shallow,
);
expose('library', useKnowledgeBaseStore);

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type MentionAction } from './action';
import { createMentionSlice } from './action';
@@ -22,4 +23,6 @@ const devtools = createDevtools('mention');
export const useMentionStore = createWithEqualityFn<MentionStore>()(devtools(createStore), shallow);
expose('mention', useMentionStore);
export const getMentionStoreState = () => useMentionStore.getState();

View File

@@ -0,0 +1,12 @@
import { isDev } from '@/utils/env';
/**
* In development, registers the store on `window.__LOBE_STORES[name]` as a getter that returns
* the current snapshot from `store.getState()`.
*/
export function expose<T>(name: string, store: { getState: () => T }): void {
if (!isDev || typeof window === 'undefined') return;
window.__LOBE_STORES ??= {};
window.__LOBE_STORES[name] = () => store.getState();
}

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type NotebookAction } from './action';
import { createNotebookAction } from './action';
@@ -25,4 +26,6 @@ export const useNotebookStore = createWithEqualityFn<NotebookStore>()(
shallow,
);
expose('notebook', useNotebookStore);
export const getNotebookStoreState = () => useNotebookStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type PageState } from './initialState';
import { initialState } from './initialState';
@@ -38,4 +39,6 @@ const devtools = createDevtools('page');
export const usePageStore = createWithEqualityFn<PageStore>()(devtools(createStore), shallow);
expose('page', usePageStore);
export const getPageStoreState = () => usePageStore.getState();

View File

@@ -8,6 +8,7 @@ import { createContext } from 'zustand-utils';
import { type IFeatureFlagsState } from '@/config/featureFlags';
import { DEFAULT_FEATURE_FLAGS, mapFeatureFlagsEnvToState } from '@/config/featureFlags';
import { createDevtools } from '@/store/middleware/createDevtools';
import { expose } from '@/store/middleware/expose';
import { type GlobalServerConfig } from '@/types/serverConfig';
import { merge } from '@/utils/merge';
@@ -73,6 +74,8 @@ export const createServerConfigStore = (initState?: Partial<ServerConfigStore>)
if (typeof window !== 'undefined') {
window.global_serverConfigStore = store;
}
expose('serverConfig', store);
}
return store;

View File

@@ -6,6 +6,7 @@ import { type StateCreator } from 'zustand/vanilla';
import { isDev } from '@/utils/env';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type SessionStoreState } from './initialState';
import { initialState } from './initialState';
@@ -49,4 +50,6 @@ export const useSessionStore = createWithEqualityFn<SessionStore>()(
shallow,
);
expose('session', useSessionStore);
export const getSessionStoreState = () => useSessionStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { initialState, type ToolStoreState } from './initialState';
import { type AgentSkillsAction, createAgentSkillsSlice } from './slices/agentSkills';
@@ -60,4 +61,6 @@ const devtools = createDevtools('tools');
export const useToolStore = createWithEqualityFn<ToolStore>()(devtools(createStore), shallow);
expose('tool', useToolStore);
export const getToolStoreState = () => useToolStore.getState();

View File

@@ -4,6 +4,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type UserState } from './initialState';
import { initialState } from './initialState';
@@ -55,4 +56,6 @@ export const useUserStore = createWithEqualityFn<UserStore>()(
shallow,
);
expose('user', useUserStore);
export const getUserStoreState = () => useUserStore.getState();

View File

@@ -3,6 +3,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { flattenActions } from '../utils/flattenActions';
import { type UserMemoryStoreState } from './initialState';
import { initialState } from './initialState';
@@ -67,4 +68,6 @@ export const useUserMemoryStore = createWithEqualityFn<UserMemoryStore>()(
shallow,
);
expose('userMemory', useUserMemoryStore);
export const getUserMemoryStoreState = () => useUserMemoryStore.getState();

View File

@@ -4,6 +4,7 @@ import { createWithEqualityFn } from 'zustand/traditional';
import { type StateCreator } from 'zustand/vanilla';
import { createDevtools } from '../middleware/createDevtools';
import { expose } from '../middleware/expose';
import { initialState, type VideoStoreState } from './initialState';
import { createCreateVideoSlice, type CreateVideoAction } from './slices/createVideo/action';
import {
@@ -46,4 +47,6 @@ export const useVideoStore = createWithEqualityFn<VideoStore>()(
shallow,
);
expose('video', useVideoStore);
export const getVideoStoreState = () => useVideoStore.getState();

View File

@@ -18,9 +18,10 @@ declare module 'styled-components' {
declare global {
interface Window {
__CHAT_STORE__?: any;
__DEBUG_PROXY__: boolean | undefined;
__editor?: IEditor;
/** Dev-only: Zustand store snapshots via `getState()` keyed by store name */
__LOBE_STORES?: Record<string, () => unknown>;
__SERVER_CONFIG__: SPAServerConfig | undefined;
lobeEnv?: {
darwinMajorVersion?: number;