@lobehub/lobehub 2.0.0-next.84 → 2.0.0-next.86

package/src/store/chat/slices/aiChat/actions/conversationControl.ts

@@ -2,18 +2,13 @@
 // Disable the auto sort key eslint rule to make the code more logic and readable
 import { type AgentRuntimeContext } from '@lobechat/agent-runtime';
 import { MESSAGE_CANCEL_FLAT } from '@lobechat/const';
-import { produce } from 'immer';
 import { StateCreator } from 'zustand/vanilla';
 
 import { ChatStore } from '@/store/chat/store';
-import { setNamespace } from '@/utils/storeDebug';
 
 import { displayMessageSelectors } from '../../../selectors';
 import { messageMapKey } from '../../../utils/messageMapKey';
 import { dbMessageSelectors } from '../../message/selectors';
-import { MainSendMessageOperation } from '../initialState';
-
-const n = setNamespace('ai');
 
 /**
  * Actions for controlling conversation operations like cancellation and error handling
@@ -47,22 +42,6 @@ export interface ConversationControlAction {
    * Reject tool intervention and continue
    */
   rejectAndContinueToolCalling: (messageId: string, reason?: string) => Promise<void>;
-  /**
-   * Toggle sendMessage operation state
-   */
-  internal_toggleSendMessageOperation: (
-    key: string | { sessionId: string; topicId?: string | null },
-    loading: boolean,
-    cancelReason?: string,
-  ) => AbortController | undefined;
-  /**
-   * Update sendMessage operation metadata
-   */
-  internal_updateSendMessageOperation: (
-    key: string | { sessionId: string; topicId?: string | null },
-    value: Partial<MainSendMessageOperation> | null,
-    actionName?: any,
-  ) => void;
 }
 
 export const conversationControl: StateCreator<
@@ -72,13 +51,18 @@ export const conversationControl: StateCreator<
   ConversationControlAction
 > = (set, get) => ({
   stopGenerateMessage: () => {
-    const { chatLoadingIdsAbortController, internal_toggleChatLoading } = get();
-
-    if (!chatLoadingIdsAbortController) return;
-
-    chatLoadingIdsAbortController.abort(MESSAGE_CANCEL_FLAT);
-
-    internal_toggleChatLoading(false, undefined, n('stopGenerateMessage') as string);
+    const { activeId, activeTopicId, cancelOperations } = get();
+
+    // Cancel all running execAgentRuntime operations in the current context
+    cancelOperations(
+      {
+        type: 'execAgentRuntime',
+        status: 'running',
+        sessionId: activeId,
+        topicId: activeTopicId,
+      },
+      MESSAGE_CANCEL_FLAT,
+    );
   },
 
   cancelSendMessageInServer: (topicId?: string) => {
@@ -86,66 +70,48 @@ export const conversationControl: StateCreator<
 
     // Determine which operation to cancel
     const targetTopicId = topicId ?? activeTopicId;
-    const operationKey = messageMapKey(activeId, targetTopicId);
+    const contextKey = messageMapKey(activeId, targetTopicId);
 
-    // Cancel the specific operation
-    get().internal_toggleSendMessageOperation(
-      operationKey,
-      false,
-      'User cancelled sendMessage operation',
-    );
+    // Cancel operations in the operation system
+    const operationIds = get().operationsByContext[contextKey] || [];
 
-    // Only clear creating message state if it's the active session
-    if (operationKey === messageMapKey(activeId, activeTopicId)) {
-      const editorTempState = get().mainSendMessageOperations[operationKey]?.inputEditorTempState;
+    operationIds.forEach((opId) => {
+      const operation = get().operations[opId];
+      if (operation && operation.type === 'sendMessage' && operation.status === 'running') {
+        get().cancelOperation(opId, 'User cancelled');
+      }
+    });
 
-      if (editorTempState) get().mainInputEditor?.setJSONState(editorTempState);
+    // Restore editor state if it's the active session
+    if (contextKey === messageMapKey(activeId, activeTopicId)) {
+      // Find the latest sendMessage operation with editor state
+      for (const opId of [...operationIds].reverse()) {
+        const op = get().operations[opId];
+        if (op && op.type === 'sendMessage' && op.metadata.inputEditorTempState) {
+          get().mainInputEditor?.setJSONState(op.metadata.inputEditorTempState);
+          break;
+        }
+      }
     }
   },
 
   clearSendMessageError: () => {
-    get().internal_updateSendMessageOperation(
-      { sessionId: get().activeId, topicId: get().activeTopicId },
-      null,
-      'clearSendMessageError',
-    );
+    const { activeId, activeTopicId } = get();
+    const contextKey = messageMapKey(activeId, activeTopicId);
+    const operationIds = get().operationsByContext[contextKey] || [];
+
+    // Clear error message from all sendMessage operations in current context
+    operationIds.forEach((opId) => {
+      const op = get().operations[opId];
+      if (op && op.type === 'sendMessage' && op.metadata.inputSendErrorMsg) {
+        get().updateOperationMetadata(opId, { inputSendErrorMsg: undefined });
+      }
+    });
   },
 
   switchMessageBranch: async (messageId, branchIndex) => {
     await get().optimisticUpdateMessageMetadata(messageId, { activeBranchIndex: branchIndex });
   },
-
-  internal_toggleSendMessageOperation: (key, loading: boolean, cancelReason?: string) => {
-    if (loading) {
-      const abortController = new AbortController();
-
-      get().internal_updateSendMessageOperation(
-        key,
-        { isLoading: true, abortController },
-        n('toggleSendMessageOperation(start)', { key }),
-      );
-
-      return abortController;
-    } else {
-      const operationKey =
-        typeof key === 'string' ? key : messageMapKey(key.sessionId, key.topicId);
-
-      const operation = get().mainSendMessageOperations[operationKey];
-
-      // If cancelReason is provided, abort the operation first
-      if (cancelReason && operation?.isLoading) {
-        operation.abortController?.abort(cancelReason);
-      }
-
-      get().internal_updateSendMessageOperation(
-        key,
-        { isLoading: false, abortController: null },
-        n('toggleSendMessageOperation(stop)', { key, cancelReason }),
-      );
-
-      return undefined;
-    }
-  },
   approveToolCalling: async (toolMessageId) => {
     const { activeThreadId, internal_execAgentRuntime } = get();
 
@@ -251,27 +217,4 @@ export const conversationControl: StateCreator<
       console.error('[rejectAndContinueToolCalling] Error executing agent runtime:', error);
     }
   },
-
-  internal_updateSendMessageOperation: (key, value, actionName) => {
-    const operationKey = typeof key === 'string' ? key : messageMapKey(key.sessionId, key.topicId);
-
-    set(
-      produce((draft) => {
-        if (!draft.mainSendMessageOperations[operationKey])
-          draft.mainSendMessageOperations[operationKey] = value;
-        else {
-          if (value === null) {
-            delete draft.mainSendMessageOperations[operationKey];
-          } else {
-            draft.mainSendMessageOperations[operationKey] = {
-              ...draft.mainSendMessageOperations[operationKey],
-              ...value,
-            };
-          }
-        }
-      }),
-      false,
-      actionName ?? n('updateSendMessageOperation', { operationKey, value }),
-    );
-  },
 });

package/src/store/chat/slices/aiChat/actions/conversationLifecycle.ts

@@ -134,17 +134,31 @@ export const conversationLifecycle: StateCreator<
     });
     get().internal_toggleMessageLoading(true, tempId);
 
-    const operationKey = messageMapKey(activeId, activeTopicId);
+    // Create operation for send message
+    const { operationId, abortController } = get().startOperation({
+      type: 'sendMessage',
+      context: {
+        sessionId: activeId,
+        topicId: activeTopicId,
+        threadId: activeThreadId,
+        messageId: tempId,
+      },
+      label: 'Send Message',
+      metadata: {
+        // Mark this as main window operation (not thread)
+        inThread: false,
+      },
+    });
 
-    // Start tracking sendMessage operation with AbortController
-    const abortController = get().internal_toggleSendMessageOperation(operationKey, true)!;
+    // Associate temp message with operation
+    get().associateMessageWithOperation(tempId, operationId);
 
+    // Store editor state in operation metadata for cancel restoration
     const jsonState = mainInputEditor?.getJSONState();
-    get().internal_updateSendMessageOperation(
-      operationKey,
-      { inputSendErrorMsg: undefined, inputEditorTempState: jsonState },
-      'creatingMessage/start',
-    );
+    get().updateOperationMetadata(operationId, {
+      inputEditorTempState: jsonState,
+      inputSendErrorMsg: undefined,
+    });
 
     let data: SendMessageServerResponse | undefined;
     try {
@@ -172,6 +186,9 @@ export const conversationLifecycle: StateCreator<
       if (data?.topics) {
         get().internal_dispatchTopic({ type: 'updateTopics', value: data.topics });
         topicId = data.topicId;
+
+        // Record the created topicId in metadata (not context)
+        get().updateOperationMetadata(operationId, { createdTopicId: data.topicId });
       }
 
       get().replaceMessages(data.messages, {
@@ -184,33 +201,36 @@ export const conversationLifecycle: StateCreator<
         await get().switchTopic(data.topicId, true);
       }
     } catch (e) {
+      // Fail operation on error
+      get().failOperation(operationId, {
+        type: e instanceof Error ? e.name : 'unknown_error',
+        message: e instanceof Error ? e.message : 'Unknown error',
+      });
+
       if (e instanceof TRPCClientError) {
         const isAbort = e.message.includes('aborted') || e.name === 'AbortError';
         // Check if error is due to cancellation
         if (!isAbort) {
-          get().internal_updateSendMessageOperation(operationKey, { inputSendErrorMsg: e.message });
+          get().updateOperationMetadata(operationId, { inputSendErrorMsg: e.message });
           get().mainInputEditor?.setJSONState(jsonState);
         }
       }
     } finally {
-      // Stop tracking sendMessage operation
-      get().internal_toggleSendMessageOperation(operationKey, false);
-    }
-
-    // remove temporally message
-    if (data?.isCreateNewTopic) {
-      get().internal_dispatchMessage(
-        { type: 'deleteMessages', ids: [tempId, tempAssistantId] },
-        { topicId: activeTopicId, sessionId: activeId },
-      );
+      // 创建了新topic 或者 用户 cancel 了消息（或者失败了），此时无 data
+      if (data?.isCreateNewTopic || !data) {
+        get().internal_dispatchMessage(
+          { type: 'deleteMessages', ids: [tempId, tempAssistantId] },
+          { operationId },
+        );
+      }
     }
 
     get().internal_toggleMessageLoading(false, tempId);
-    get().internal_updateSendMessageOperation(
-      operationKey,
-      { inputEditorTempState: null },
-      'creatingMessage/finished',
-    );
+
+    // Clear editor temp state after message created
+    if (data) {
+      get().updateOperationMetadata(operationId, { inputEditorTempState: null });
+    }
 
     if (!data) return;
 
@@ -251,6 +271,7 @@ export const conversationLifecycle: StateCreator<
         parentMessageType: 'assistant',
         sessionId: activeId,
         topicId: data.topicId ?? activeTopicId,
+        parentOperationId: operationId, // Pass as parent operation
         ragQuery: get().internal_shouldUseRAG() ? message : undefined,
         threadId: activeThreadId,
         skipCreateFirstMessage: true,
@@ -267,8 +288,16 @@ export const conversationLifecycle: StateCreator<
       if (userFiles.length > 0) {
         await getAgentStoreState().addFilesToAgent(userFiles, false);
       }
+
+      // Complete operation on success
+      get().completeOperation(operationId);
     } catch (e) {
       console.error(e);
+      // Fail operation on error
+      get().failOperation(operationId, {
+        type: e instanceof Error ? e.name : 'unknown_error',
+        message: e instanceof Error ? e.message : 'AI generation failed',
+      });
     } finally {
       if (data.topicId) get().internal_updateTopicLoading(data.topicId, false);
     }
@@ -288,16 +317,15 @@ export const conversationLifecycle: StateCreator<
 
     if (contextMessages.length <= 0) return;
 
-    try {
-      const { internal_execAgentRuntime, activeThreadId } = get();
+    const { internal_execAgentRuntime, activeThreadId, activeId, activeTopicId } = get();
 
-      // Mark message as regenerating
-      set(
-        { regeneratingIds: [...get().regeneratingIds, id] },
-        false,
-        'regenerateUserMessage/start',
-      );
+    // Create regenerate operation
+    const { operationId } = get().startOperation({
+      type: 'regenerate',
+      context: { sessionId: activeId, topicId: activeTopicId, messageId: id },
+    });
 
+    try {
       const traceId = params?.traceId ?? dbMessageSelectors.getTraceIdByDbMessageId(id)(get());
 
       // 切一个新的激活分支
@@ -307,23 +335,25 @@ export const conversationLifecycle: StateCreator<
         messages: contextMessages,
         parentMessageId: id,
         parentMessageType: 'user',
-        sessionId: get().activeId,
-        topicId: get().activeTopicId,
+        sessionId: activeId,
+        topicId: activeTopicId,
         traceId,
         ragQuery: get().internal_shouldUseRAG() ? item.content : undefined,
         threadId: activeThreadId,
+        parentOperationId: operationId,
       });
 
       // trace the regenerate message
       if (!params?.skipTrace)
         get().internal_traceMessage(id, { eventType: TraceEventType.RegenerateMessage });
-    } finally {
-      // Remove message from regenerating state
-      set(
-        { regeneratingIds: get().regeneratingIds.filter((msgId) => msgId !== id) },
-        false,
-        'regenerateUserMessage/end',
-      );
+
+      get().completeOperation(operationId);
+    } catch (error) {
+      get().failOperation(operationId, {
+        type: 'RegenerateError',
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
     }
   },
 
@@ -350,30 +380,33 @@ export const conversationLifecycle: StateCreator<
     const message = dbMessageSelectors.getDbMessageById(id)(get());
     if (!message) return;
 
-    try {
-      // Mark message as continuing
-      set(
-        { continuingIds: [...get().continuingIds, messageId] },
-        false,
-        'continueGenerationMessage/start',
-      );
+    const { activeId, activeTopicId } = get();
 
+    // Create continue operation
+    const { operationId } = get().startOperation({
+      type: 'continue',
+      context: { sessionId: activeId, topicId: activeTopicId, messageId },
+    });
+
+    try {
       const chats = displayMessageSelectors.mainAIChatsWithHistoryConfig(get());
 
       await get().internal_execAgentRuntime({
         messages: chats,
         parentMessageId: id,
         parentMessageType: message.role as 'assistant' | 'tool' | 'user',
-        sessionId: get().activeId,
-        topicId: get().activeTopicId,
+        sessionId: activeId,
+        topicId: activeTopicId,
+        parentOperationId: operationId,
       });
-    } finally {
-      // Remove message from continuing state
-      set(
-        { continuingIds: get().continuingIds.filter((msgId) => msgId !== messageId) },
-        false,
-        'continueGenerationMessage/end',
-      );
+
+      get().completeOperation(operationId);
+    } catch (error) {
+      get().failOperation(operationId, {
+        type: 'ContinueError',
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
     }
   },
 

package/src/store/chat/slices/aiChat/actions/generateAIGroupChat.ts

@@ -1,5 +1,6 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */
 // Disable the auto sort key eslint rule to make the code more logic and readable
+import { LOADING_FLAT } from '@lobechat/const';
 import {
   GroupMemberInfo,
   buildGroupChatSystemPrompt,
@@ -14,7 +15,6 @@ import {
 import { produce } from 'immer';
 import { StateCreator } from 'zustand/vanilla';
 
-import { LOADING_FLAT } from '@/const/message';
 import { DEFAULT_CHAT_GROUP_CHAT_CONFIG } from '@/const/settings';
 import { ChatStore } from '@/store/chat/store';
 import { messageMapKey } from '@/store/chat/utils/messageMapKey';
@@ -602,8 +602,6 @@ export const chatAiGroupChat: StateCreator<
         refreshMessages,
         activeTopicId,
         internal_dispatchMessage,
-        internal_toggleChatLoading,
-        triggerToolCalls,
       } = get();
 
       try {
@@ -710,30 +708,14 @@ export const chatAiGroupChat: StateCreator<
         const messagesForAPI = [systemMessage, ...messagesWithAuthors];
 
         if (assistantId) {
-          const { isFunctionCall } = await internal_fetchAIChatMessage({
-            messages: messagesForAPI,
+          await internal_fetchAIChatMessage({
             messageId: assistantId,
+            messages: messagesForAPI,
             model: agentModel,
             provider: agentProvider,
-            params: {
-              traceId: `group-${groupId}-agent-${agentId}`,
-              agentConfig: agentData,
-            },
+            agentConfig: agentData,
+            traceId: `group-${groupId}-agent-${agentId}`,
           });
-
-          // Handle tool calling in group chat like single chat
-          if (isFunctionCall) {
-            get().internal_toggleMessageInToolsCalling(true, assistantId);
-            await refreshMessages();
-            await triggerToolCalls(assistantId, {
-              threadId: undefined,
-              inPortalThread: false,
-            });
-            // Change: if an agent message is a tool call, make the same agent speak again
-            // instead of asking supervisor for a decision.
-            await get().internal_processAgentMessage(groupId, agentId, targetId, instruction);
-            return;
-          }
         }
 
         await refreshMessages();
@@ -770,8 +752,6 @@ export const chatAiGroupChat: StateCreator<
             },
           });
         }
-      } finally {
-        internal_toggleChatLoading(false, undefined, n('processAgentMessage(end)'));
       }
     },