diff --git a/src/renderer/src/services/messageStreaming/BlockManager.ts b/src/renderer/src/services/messageStreaming/BlockManager.ts index fb0d65913a..c07c67f804 100644 --- a/src/renderer/src/services/messageStreaming/BlockManager.ts +++ b/src/renderer/src/services/messageStreaming/BlockManager.ts @@ -121,7 +121,8 @@ export class BlockManager { newMessagesActions.upsertBlockReference({ messageId: this.deps.assistantMsgId, blockId: newBlock.id, - status: newBlock.status + status: newBlock.status, + blockType: newBlock.type }) ) diff --git a/src/renderer/src/store/newMessage.ts b/src/renderer/src/store/newMessage.ts index f5b171f07d..82b81adb8b 100644 --- a/src/renderer/src/store/newMessage.ts +++ b/src/renderer/src/store/newMessage.ts @@ -2,7 +2,7 @@ import { loggerService } from '@logger' import { createEntityAdapter, createSlice, EntityState, PayloadAction } from '@reduxjs/toolkit' // Separate type-only imports from value imports import type { Message } from '@renderer/types/newMessage' -import { AssistantMessageStatus, MessageBlockStatus } from '@renderer/types/newMessage' +import { AssistantMessageStatus, MessageBlockStatus, MessageBlockType } from '@renderer/types/newMessage' const logger = loggerService.withContext('newMessage') @@ -50,6 +50,7 @@ interface UpsertBlockReferencePayload { messageId: string blockId: string status?: MessageBlockStatus + blockType?: MessageBlockType } // Payload for removing a single message @@ -218,7 +219,7 @@ export const messagesSlice = createSlice({ messagesAdapter.removeMany(state, messageIds) }, upsertBlockReference(state, action: PayloadAction) { - const { messageId, blockId, status } = action.payload + const { messageId, blockId, status, blockType } = action.payload const messageToUpdate = state.entities[messageId] if (!messageToUpdate) { @@ -231,7 +232,11 @@ export const messagesSlice = createSlice({ // Update Block ID const currentBlocks = messageToUpdate.blocks || [] if (!currentBlocks.includes(blockId)) { - changes.blocks = [...currentBlocks, blockId] + if (blockType === MessageBlockType.THINKING) { + changes.blocks = [blockId, ...currentBlocks] + } else { + changes.blocks = [...currentBlocks, blockId] + } } // Update Message Status based on Block Status