mirror of
https://github.com/CherryHQ/cherry-studio.git
synced 2026-02-04 01:51:07 +08:00
fix/selection-not-copy (#8276)
* fix/selection-not-copy * fix: abort not copy * test: add actionUtil test Update onStream callback to accept content parameter and set copyable content during streaming instead of only at completion. Removes redundant content tracking in ActionUtils and properly passes content to callbacks as it's received. Adds test directory structure for components. * refactor: improve streaming content handling and fix error states Refine content handling in ActionUtils by tracking content internally rather than passing it through onStream callbacks. Properly set error status in message updates and ensure content is finalized on errors. Fix typo in test filename. * fix: show error * chore: remove unuse * chore: remove console
This commit is contained in:
SuYao
GitHub
parent
d0c375aa0a
commit
92513024b5
@ -8,8 +8,8 @@ import { cancelThrottledBlockUpdate, throttledBlockUpdate } from '@renderer/stor
|
||||
import { Assistant, Topic } from '@renderer/types'
|
||||
import { Chunk, ChunkType } from '@renderer/types/chunk'
|
||||
import { AssistantMessageStatus, MessageBlockStatus } from '@renderer/types/newMessage'
|
||||
import { isAbortError } from '@renderer/utils/error'
|
||||
import { createMainTextBlock, createThinkingBlock } from '@renderer/utils/messageUtils/create'
|
||||
import { formatErrorMessage, isAbortError } from '@renderer/utils/error'
|
||||
import { createErrorBlock, createMainTextBlock, createThinkingBlock } from '@renderer/utils/messageUtils/create'
|
||||
|
||||
const logger = loggerService.withContext('ActionUtils')
|
||||
|
||||
@ -38,7 +38,7 @@ export const processMessages = async (
|
||||
|
||||
let textBlockId: string | null = null
|
||||
let thinkingBlockId: string | null = null
|
||||
const textBlockContent: string = ''
|
||||
let textBlockContent: string = ''
|
||||
|
||||
const assistantMessage = getAssistantMessage({
|
||||
assistant,
|
||||
@ -133,6 +133,7 @@ export const processMessages = async (
|
||||
throttledBlockUpdate(textBlockId, { content: chunk.text })
|
||||
}
|
||||
onStream()
|
||||
textBlockContent = chunk.text
|
||||
}
|
||||
break
|
||||
case ChunkType.TEXT_COMPLETE:
|
||||
@ -146,6 +147,7 @@ export const processMessages = async (
|
||||
})
|
||||
)
|
||||
onFinish(chunk.text)
|
||||
textBlockContent = chunk.text
|
||||
textBlockId = null
|
||||
}
|
||||
}
|
||||
@ -159,7 +161,6 @@ export const processMessages = async (
|
||||
updates: { status: AssistantMessageStatus.SUCCESS }
|
||||
})
|
||||
)
|
||||
onFinish(textBlockContent)
|
||||
}
|
||||
break
|
||||
case ChunkType.ERROR:
|
||||
@ -175,12 +176,36 @@ export const processMessages = async (
|
||||
})
|
||||
)
|
||||
}
|
||||
const isErrorTypeAbort = isAbortError(chunk.error)
|
||||
let pauseErrorLanguagePlaceholder = ''
|
||||
if (isErrorTypeAbort) {
|
||||
pauseErrorLanguagePlaceholder = 'pause_placeholder'
|
||||
}
|
||||
const serializableError = {
|
||||
name: chunk.error.name,
|
||||
message: pauseErrorLanguagePlaceholder || chunk.error.message || formatErrorMessage(chunk.error),
|
||||
originalMessage: chunk.error.message,
|
||||
stack: chunk.error.stack,
|
||||
status: chunk.error.status || chunk.error.code,
|
||||
requestId: chunk.error.request_id
|
||||
}
|
||||
const errorBlock = createErrorBlock(assistantMessage.id, serializableError, {
|
||||
status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR
|
||||
})
|
||||
store.dispatch(
|
||||
newMessagesActions.updateMessage({
|
||||
topicId: topic.id,
|
||||
messageId: assistantMessage.id,
|
||||
updates: { blockInstruction: { id: errorBlock.id } }
|
||||
})
|
||||
)
|
||||
store.dispatch(upsertOneBlock(errorBlock))
|
||||
store.dispatch(
|
||||
newMessagesActions.updateMessage({
|
||||
topicId: topic.id,
|
||||
messageId: assistantMessage.id,
|
||||
updates: {
|
||||
status: isAbortError(chunk.error) ? AssistantMessageStatus.PAUSED : AssistantMessageStatus.SUCCESS
|
||||
status: isAbortError(chunk.error) ? AssistantMessageStatus.PAUSED : AssistantMessageStatus.ERROR
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
@ -0,0 +1,524 @@
|
||||
import type { Assistant, Topic } from '@renderer/types'
|
||||
import { ChunkType } from '@renderer/types/chunk'
|
||||
import { AssistantMessageStatus, MessageBlockStatus } from '@renderer/types/newMessage'
|
||||
import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
|
||||
|
||||
import { processMessages } from '../ActionUtils'
|
||||
|
||||
// Mock all dependencies
|
||||
vi.mock('@renderer/services/ApiService', () => ({
|
||||
fetchChatCompletion: vi.fn()
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/services/MessagesService', () => ({
|
||||
getUserMessage: vi.fn(),
|
||||
getAssistantMessage: vi.fn()
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/store', () => ({
|
||||
default: {
|
||||
dispatch: vi.fn()
|
||||
}
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/store/messageBlock', () => ({
|
||||
updateOneBlock: vi.fn(),
|
||||
upsertManyBlocks: vi.fn(),
|
||||
upsertOneBlock: vi.fn()
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/store/newMessage', () => ({
|
||||
newMessagesActions: {
|
||||
addMessage: vi.fn(),
|
||||
updateMessage: vi.fn()
|
||||
}
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/store/thunk/messageThunk', () => ({
|
||||
cancelThrottledBlockUpdate: vi.fn(),
|
||||
throttledBlockUpdate: vi.fn()
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/utils/error', () => ({
|
||||
isAbortError: vi.fn(),
|
||||
formatErrorMessage: vi.fn()
|
||||
}))
|
||||
|
||||
vi.mock('@renderer/utils/messageUtils/create', () => ({
|
||||
createMainTextBlock: vi.fn(),
|
||||
createThinkingBlock: vi.fn(),
|
||||
createErrorBlock: vi.fn()
|
||||
}))
|
||||
|
||||
// Import mocked modules
|
||||
import { fetchChatCompletion } from '@renderer/services/ApiService'
|
||||
import { getAssistantMessage, getUserMessage } from '@renderer/services/MessagesService'
|
||||
import store from '@renderer/store'
|
||||
import { updateOneBlock } from '@renderer/store/messageBlock'
|
||||
import { newMessagesActions } from '@renderer/store/newMessage'
|
||||
import { cancelThrottledBlockUpdate, throttledBlockUpdate } from '@renderer/store/thunk/messageThunk'
|
||||
import { formatErrorMessage, isAbortError } from '@renderer/utils/error'
|
||||
import { createErrorBlock, createMainTextBlock, createThinkingBlock } from '@renderer/utils/messageUtils/create'
|
||||
|
||||
describe('processMessages', () => {
|
||||
let mockAssistant: Assistant
|
||||
let mockTopic: Topic
|
||||
let mockSetAskId: Mock
|
||||
let mockOnStream: Mock
|
||||
let mockOnFinish: Mock
|
||||
let mockOnError: Mock
|
||||
|
||||
beforeEach(() => {
|
||||
// Setup mock data
|
||||
mockAssistant = {
|
||||
id: 'assistant-1',
|
||||
name: 'Test Assistant',
|
||||
model: {
|
||||
id: 'model-1',
|
||||
name: 'test model',
|
||||
provider: 'test provider',
|
||||
group: 'test group'
|
||||
},
|
||||
prompt: '',
|
||||
topics: [],
|
||||
type: 'assistant'
|
||||
} as Assistant
|
||||
|
||||
mockTopic = {
|
||||
id: 'topic-1',
|
||||
name: 'Test Topic'
|
||||
} as Topic
|
||||
|
||||
// Setup mock callbacks
|
||||
mockSetAskId = vi.fn()
|
||||
mockOnStream = vi.fn()
|
||||
mockOnFinish = vi.fn()
|
||||
mockOnError = vi.fn()
|
||||
|
||||
// Reset all mocks
|
||||
vi.clearAllMocks()
|
||||
|
||||
// Setup default mock implementations
|
||||
vi.mocked(getUserMessage).mockReturnValue({
|
||||
message: { id: 'user-message-1', role: 'user', content: 'test prompt' },
|
||||
blocks: []
|
||||
} as any)
|
||||
|
||||
vi.mocked(getAssistantMessage).mockReturnValue({
|
||||
id: 'assistant-message-1',
|
||||
role: 'assistant',
|
||||
content: ''
|
||||
} as any)
|
||||
|
||||
vi.mocked(createThinkingBlock).mockReturnValue({
|
||||
id: 'thinking-block-1',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
|
||||
vi.mocked(createMainTextBlock).mockReturnValue({
|
||||
id: 'text-block-1',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
|
||||
vi.mocked(createErrorBlock).mockReturnValue({
|
||||
id: 'error-block-1',
|
||||
content: '',
|
||||
status: MessageBlockStatus.ERROR
|
||||
} as any)
|
||||
|
||||
vi.mocked(isAbortError).mockReturnValue(false)
|
||||
vi.mocked(formatErrorMessage).mockReturnValue('Formatted error message')
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
describe('normal complete stream with thinking flow', () => {
|
||||
it('should process a complete stream with thinking and text blocks', async () => {
|
||||
// Mock chunk stream for normal flow
|
||||
const mockChunks = [
|
||||
{ type: ChunkType.THINKING_START },
|
||||
{ type: ChunkType.THINKING_DELTA, text: 'I need to think about this...', thinking_millsec: 1000 },
|
||||
{
|
||||
type: ChunkType.THINKING_DELTA,
|
||||
text: 'I need to think about this... Let me consider the options.',
|
||||
thinking_millsec: 2000
|
||||
},
|
||||
{
|
||||
type: ChunkType.THINKING_COMPLETE,
|
||||
text: 'I need to think about this... Let me consider the options. Now I have a solution.',
|
||||
thinking_millsec: 3000
|
||||
},
|
||||
{ type: ChunkType.TEXT_START },
|
||||
{ type: ChunkType.TEXT_DELTA, text: 'Here is' },
|
||||
{ type: ChunkType.TEXT_DELTA, text: 'Here is my' },
|
||||
{ type: ChunkType.TEXT_DELTA, text: 'Here is my answer' },
|
||||
{ type: ChunkType.TEXT_COMPLETE, text: 'Here is my answer to your question.' },
|
||||
{ type: ChunkType.BLOCK_COMPLETE }
|
||||
]
|
||||
|
||||
vi.mocked(fetchChatCompletion).mockImplementation(async ({ onChunkReceived }: any) => {
|
||||
for (const chunk of mockChunks) {
|
||||
await onChunkReceived(chunk)
|
||||
}
|
||||
})
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify setAskId was called
|
||||
expect(mockSetAskId).toHaveBeenCalledWith('user-message-1')
|
||||
|
||||
// Verify store dispatches for user message
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
newMessagesActions.addMessage({
|
||||
topicId: 'topic-1',
|
||||
message: expect.objectContaining({ id: 'user-message-1' })
|
||||
})
|
||||
)
|
||||
|
||||
// Verify store dispatches for assistant message
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
newMessagesActions.addMessage({
|
||||
topicId: 'topic-1',
|
||||
message: expect.objectContaining({ id: 'assistant-message-1' })
|
||||
})
|
||||
)
|
||||
|
||||
// Verify thinking block creation and updates
|
||||
expect(createThinkingBlock).toHaveBeenCalledWith('assistant-message-1', '', {
|
||||
status: MessageBlockStatus.STREAMING
|
||||
})
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('thinking-block-1', {
|
||||
content: 'I need to think about this...',
|
||||
thinking_millsec: 1000
|
||||
})
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('thinking-block-1', {
|
||||
content: 'I need to think about this... Let me consider the options.',
|
||||
thinking_millsec: 2000
|
||||
})
|
||||
|
||||
// Verify thinking block completion
|
||||
expect(cancelThrottledBlockUpdate).toHaveBeenCalledWith('thinking-block-1')
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
updateOneBlock({
|
||||
id: 'thinking-block-1',
|
||||
changes: {
|
||||
content: 'I need to think about this... Let me consider the options. Now I have a solution.',
|
||||
status: MessageBlockStatus.SUCCESS,
|
||||
thinking_millsec: 3000
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
// Verify text block creation and updates
|
||||
expect(createMainTextBlock).toHaveBeenCalledWith('assistant-message-1', '', {
|
||||
status: MessageBlockStatus.STREAMING
|
||||
})
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('text-block-1', { content: 'Here is' })
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('text-block-1', { content: 'Here is my' })
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('text-block-1', { content: 'Here is my answer' })
|
||||
|
||||
// Verify text block completion
|
||||
expect(cancelThrottledBlockUpdate).toHaveBeenCalledWith('text-block-1')
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
updateOneBlock({
|
||||
id: 'text-block-1',
|
||||
changes: {
|
||||
content: 'Here is my answer to your question.',
|
||||
status: MessageBlockStatus.SUCCESS
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
// Verify callbacks
|
||||
expect(mockOnStream).toHaveBeenCalledWith()
|
||||
|
||||
// Verify final message status update
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
newMessagesActions.updateMessage({
|
||||
topicId: 'topic-1',
|
||||
messageId: 'assistant-message-1',
|
||||
updates: { status: AssistantMessageStatus.SUCCESS }
|
||||
})
|
||||
)
|
||||
|
||||
expect(mockOnFinish).toHaveBeenCalledWith('Here is my answer to your question.')
|
||||
// Verify no errors
|
||||
expect(mockOnError).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe('stream with exceptions', () => {
|
||||
it('should handle error chunks properly', async () => {
|
||||
const mockError = new Error('Stream processing error')
|
||||
const mockChunks = [
|
||||
{ type: ChunkType.TEXT_START },
|
||||
{ type: ChunkType.TEXT_DELTA, text: 'Partial response' },
|
||||
{ type: ChunkType.ERROR, error: mockError }
|
||||
]
|
||||
|
||||
vi.mocked(fetchChatCompletion).mockImplementation(async ({ onChunkReceived }: any) => {
|
||||
for (const chunk of mockChunks) {
|
||||
await onChunkReceived(chunk)
|
||||
}
|
||||
})
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify text block was created and updated
|
||||
expect(createMainTextBlock).toHaveBeenCalled()
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('text-block-1', { content: 'Partial response' })
|
||||
expect(mockOnStream).toHaveBeenCalledWith()
|
||||
|
||||
// Verify error handling
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
updateOneBlock({
|
||||
id: 'text-block-1',
|
||||
changes: {
|
||||
status: MessageBlockStatus.ERROR
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
// Verify error block creation
|
||||
expect(createErrorBlock).toHaveBeenCalledWith(
|
||||
'assistant-message-1',
|
||||
expect.objectContaining({
|
||||
name: 'Error',
|
||||
message: 'Stream processing error'
|
||||
}),
|
||||
{ status: MessageBlockStatus.ERROR }
|
||||
)
|
||||
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
newMessagesActions.updateMessage({
|
||||
topicId: 'topic-1',
|
||||
messageId: 'assistant-message-1',
|
||||
updates: {
|
||||
status: AssistantMessageStatus.ERROR
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
// Verify onFinish is called with text content accumulated so far
|
||||
expect(mockOnFinish).toHaveBeenCalledWith('Partial response')
|
||||
expect(mockOnError).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle fetchChatCompletion errors', async () => {
|
||||
const mockError = new Error('API Error')
|
||||
vi.mocked(fetchChatCompletion).mockRejectedValue(mockError)
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify error callback is called
|
||||
expect(mockOnError).toHaveBeenCalledWith(mockError)
|
||||
})
|
||||
})
|
||||
|
||||
describe('actively aborted stream', () => {
|
||||
it('should handle aborted streams properly', async () => {
|
||||
const mockAbortError = new Error('AbortError')
|
||||
vi.mocked(isAbortError).mockReturnValue(true)
|
||||
|
||||
const mockChunks = [
|
||||
{ type: ChunkType.THINKING_START },
|
||||
{ type: ChunkType.THINKING_DELTA, text: 'Starting to think...', thinking_millsec: 1000 },
|
||||
{ type: ChunkType.TEXT_START },
|
||||
{ type: ChunkType.TEXT_DELTA, text: 'Partial' },
|
||||
{ type: ChunkType.ERROR, error: mockAbortError }
|
||||
]
|
||||
|
||||
vi.mocked(fetchChatCompletion).mockImplementation(async ({ onChunkReceived }: any) => {
|
||||
for (const chunk of mockChunks) {
|
||||
await onChunkReceived(chunk)
|
||||
}
|
||||
})
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify both blocks were created
|
||||
expect(createThinkingBlock).toHaveBeenCalled()
|
||||
expect(createMainTextBlock).toHaveBeenCalled()
|
||||
|
||||
// Verify partial updates were made
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('thinking-block-1', {
|
||||
content: 'Starting to think...',
|
||||
thinking_millsec: 1000
|
||||
})
|
||||
expect(throttledBlockUpdate).toHaveBeenCalledWith('text-block-1', { content: 'Partial' })
|
||||
|
||||
// Verify abort handling - should set status to PAUSED
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
updateOneBlock({
|
||||
id: 'text-block-1',
|
||||
changes: {
|
||||
status: MessageBlockStatus.PAUSED
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
expect(store.dispatch).toHaveBeenCalledWith(
|
||||
newMessagesActions.updateMessage({
|
||||
topicId: 'topic-1',
|
||||
messageId: 'assistant-message-1',
|
||||
updates: {
|
||||
status: AssistantMessageStatus.PAUSED
|
||||
}
|
||||
})
|
||||
)
|
||||
|
||||
// Verify error block creation for abort
|
||||
expect(createErrorBlock).toHaveBeenCalledWith(
|
||||
'assistant-message-1',
|
||||
expect.objectContaining({
|
||||
name: 'Error',
|
||||
message: 'pause_placeholder'
|
||||
}),
|
||||
{ status: MessageBlockStatus.PAUSED }
|
||||
)
|
||||
|
||||
// Verify callbacks
|
||||
expect(mockOnStream).toHaveBeenCalledWith()
|
||||
expect(mockOnFinish).toHaveBeenCalledWith('Partial')
|
||||
expect(mockOnError).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle aborted fetchChatCompletion gracefully', async () => {
|
||||
const mockAbortError = new Error('AbortError')
|
||||
vi.mocked(isAbortError).mockReturnValue(true)
|
||||
vi.mocked(fetchChatCompletion).mockRejectedValue(mockAbortError)
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify that abort errors are handled gracefully (no error callback)
|
||||
expect(mockOnError).not.toHaveBeenCalled()
|
||||
expect(mockOnFinish).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('should handle missing assistant or topic', async () => {
|
||||
await processMessages(
|
||||
null as any,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Should return early without making any calls
|
||||
expect(fetchChatCompletion).not.toHaveBeenCalled()
|
||||
expect(mockSetAskId).not.toHaveBeenCalled()
|
||||
expect(mockOnStream).not.toHaveBeenCalled()
|
||||
expect(mockOnFinish).not.toHaveBeenCalled()
|
||||
expect(mockOnError).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle multiple text/thinking blocks correctly', async () => {
|
||||
const mockChunks = [
|
||||
{ type: ChunkType.THINKING_START },
|
||||
{ type: ChunkType.THINKING_COMPLETE, text: 'First thinking', thinking_millsec: 1000 },
|
||||
{ type: ChunkType.TEXT_START },
|
||||
{ type: ChunkType.TEXT_COMPLETE, text: 'First text' },
|
||||
{ type: ChunkType.THINKING_START },
|
||||
{ type: ChunkType.THINKING_COMPLETE, text: 'Second thinking', thinking_millsec: 2000 },
|
||||
{ type: ChunkType.TEXT_START },
|
||||
{ type: ChunkType.TEXT_COMPLETE, text: 'Second text' }
|
||||
]
|
||||
|
||||
vi.mocked(createThinkingBlock)
|
||||
.mockReturnValueOnce({
|
||||
id: 'thinking-block-1',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
.mockReturnValueOnce({
|
||||
id: 'thinking-block-2',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
|
||||
vi.mocked(createMainTextBlock)
|
||||
.mockReturnValueOnce({
|
||||
id: 'text-block-1',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
.mockReturnValueOnce({
|
||||
id: 'text-block-2',
|
||||
content: '',
|
||||
status: MessageBlockStatus.STREAMING
|
||||
} as any)
|
||||
|
||||
vi.mocked(fetchChatCompletion).mockImplementation(async ({ onChunkReceived }: any) => {
|
||||
for (const chunk of mockChunks) {
|
||||
await onChunkReceived(chunk)
|
||||
}
|
||||
})
|
||||
|
||||
await processMessages(
|
||||
mockAssistant,
|
||||
mockTopic,
|
||||
'test prompt',
|
||||
mockSetAskId,
|
||||
mockOnStream,
|
||||
mockOnFinish,
|
||||
mockOnError
|
||||
)
|
||||
|
||||
// Verify both thinking blocks were created and completed
|
||||
expect(createThinkingBlock).toHaveBeenCalledTimes(2)
|
||||
expect(createMainTextBlock).toHaveBeenCalledTimes(2)
|
||||
|
||||
// Verify onFinish was called for both text completions
|
||||
expect(mockOnFinish).toHaveBeenCalledWith('First text')
|
||||
expect(mockOnFinish).toHaveBeenCalledWith('Second text')
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Reference in New Issue
Block a user