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 205ab1a9cd..664270e172 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 @@ -217,7 +218,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) { @@ -230,7 +231,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