import type { DynamicToolUIPart, ToolUIPart, UIMessage } from 'ai' import { describe, expect, it } from 'vitest' import { getParallelApprovalIdsToReject, isManualApprovalRequested, prepareMessagesForAPI, } from './message-utils' import { createAssistantMessageWithUpdateNotebookTool } from './test-fixtures' const makeApprovalPart = (id: string, isAutomatic = false): DynamicToolUIPart => ({ type: 'dynamic-tool', toolName: 'test_tool', toolCallId: id, state: 'approval-requested', input: {}, approval: { id, ...(isAutomatic ? { isAutomatic: true } : {}) }, }) as DynamicToolUIPart const makeResultPart = (id: string): DynamicToolUIPart => ({ type: 'dynamic-tool', toolName: 'test_tool', toolCallId: id, state: 'output-available', input: {}, output: {}, }) describe('isManualApprovalRequested', () => { it('returns true for a human approval-requested tool part', () => { expect(isManualApprovalRequested(makeApprovalPart('a1'))).toBe(true) }) it('returns false for an automatic approval', () => { expect(isManualApprovalRequested(makeApprovalPart('a1', true))).toBe(false) }) it('returns false for a tool result part', () => { expect(isManualApprovalRequested(makeResultPart('r1'))).toBe(false) }) it('returns false for a content part with no state or approval', () => { expect(isManualApprovalRequested({ type: 'text', text: 'hello' })).toBe(false) }) }) describe('getParallelApprovalIdsToReject', () => { it('returns [] for empty messages', () => { expect(getParallelApprovalIdsToReject([])).toEqual([]) }) it('returns [] when there are no assistant messages', () => { const messages: UIMessage[] = [{ id: '1', role: 'user', parts: [] }] expect(getParallelApprovalIdsToReject(messages)).toEqual([]) }) it('returns [] when last assistant message has no pending approvals', () => { const messages: UIMessage[] = [{ id: '1', role: 'assistant', parts: [makeResultPart('r1')] }] expect(getParallelApprovalIdsToReject(messages)).toEqual([]) }) it('returns [] when there is only one pending approval', () => { const messages: UIMessage[] = [{ id: '1', role: 'assistant', parts: [makeApprovalPart('a1')] }] expect(getParallelApprovalIdsToReject(messages)).toEqual([]) }) it('returns all but the first id when there are multiple pending approvals', () => { const messages: UIMessage[] = [ { id: '1', role: 'assistant', parts: [makeApprovalPart('a1'), makeApprovalPart('a2'), makeApprovalPart('a3')], }, ] expect(getParallelApprovalIdsToReject(messages)).toEqual(['a2', 'a3']) }) it('only inspects the last assistant message', () => { const messages: UIMessage[] = [ { id: '1', role: 'assistant', parts: [makeApprovalPart('old1'), makeApprovalPart('old2')], }, { id: '2', role: 'user', parts: [] }, { id: '3', role: 'assistant', parts: [makeApprovalPart('new1')] }, ] expect(getParallelApprovalIdsToReject(messages)).toEqual([]) }) it('ignores non-approval tool parts', () => { const messages: UIMessage[] = [ { id: '1', role: 'assistant', parts: [makeResultPart('r1'), makeApprovalPart('a1'), makeApprovalPart('a2')], }, ] expect(getParallelApprovalIdsToReject(messages)).toEqual(['a2']) }) it('ignores automatic approvals when picking extras to reject', () => { const messages: UIMessage[] = [ { id: '1', role: 'assistant', parts: [ makeApprovalPart('auto', true), makeApprovalPart('manual-1'), makeApprovalPart('manual-2'), ], }, ] expect(getParallelApprovalIdsToReject(messages)).toEqual(['manual-2']) }) }) describe('prepareMessagesForAPI', () => { it('should limit messages to last 7 entries', () => { const messages: UIMessage[] = Array.from({ length: 10 }, (_, i) => ({ id: `msg-${i}`, role: 'user', parts: [{ type: 'text', text: `Message ${i}` }], })) const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(7) expect(result[0].parts[0]).toEqual({ type: 'text', text: 'Message 3' }) expect(result[6].parts[0]).toEqual({ type: 'text', text: 'Message 9' }) }) it('should remove results property from assistant messages', () => { const messages = [ { id: 'msg-1', role: 'assistant', parts: [{ type: 'text', text: 'Response' }], results: { data: 'some data' }, }, ] as Array const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(1) expect(result[0]).not.toHaveProperty('results') expect(result[0].parts[0]).toEqual({ type: 'text', text: 'Response' }) }) it('should preserve messages without results', () => { const messages: UIMessage[] = [ { id: 'msg-1', role: 'user', parts: [{ type: 'text', text: 'Question' }], }, { id: 'msg-2', role: 'assistant', parts: [{ type: 'text', text: 'Answer' }], }, ] const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(2) expect(result[0]).toEqual(messages[0]) expect(result[1]).toEqual(messages[1]) }) it('should handle empty array', () => { const messages: UIMessage[] = [] const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(0) expect(result).toEqual([]) }) it('should handle array with fewer than 7 messages', () => { const messages: UIMessage[] = [ { id: 'msg-1', role: 'user', parts: [{ type: 'text', text: 'Message 1' }] }, { id: 'msg-2', role: 'assistant', parts: [{ type: 'text', text: 'Message 2' }] }, { id: 'msg-3', role: 'user', parts: [{ type: 'text', text: 'Message 3' }] }, ] const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(3) expect(result).toEqual(messages) }) it('should handle array with exactly 7 messages', () => { const messages: UIMessage[] = Array.from({ length: 7 }, (_, i) => ({ id: `msg-${i}`, role: i % 2 === 0 ? 'user' : 'assistant', parts: [{ type: 'text', text: `Message ${i}` }], })) const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(7) expect(result).toEqual(messages) }) it('should only remove results from assistant messages, not user messages', () => { const messages = [ { id: 'msg-1', role: 'user', parts: [{ type: 'text', text: 'Question' }], results: { data: 'user data' }, }, { id: 'msg-2', role: 'assistant', parts: [{ type: 'text', text: 'Answer' }], results: { data: 'assistant data' }, }, ] as Array const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(2) // User message keeps results (not removed by the function) expect((result[0] as any).results).toEqual({ data: 'user data' }) // Assistant message has results removed expect(result[1]).not.toHaveProperty('results') }) it('should handle mixed messages with and without results', () => { const messages = [ { id: 'msg-1', role: 'assistant', parts: [{ type: 'text', text: 'First' }], results: { data: 'data1' }, }, { id: 'msg-2', role: 'user', parts: [{ type: 'text', text: 'Second' }], }, { id: 'msg-3', role: 'assistant', parts: [{ type: 'text', text: 'Third' }], }, { id: 'msg-4', role: 'assistant', parts: [{ type: 'text', text: 'Fourth' }], results: { data: 'data2' }, }, ] as Array const result = prepareMessagesForAPI(messages) expect(result).toHaveLength(4) expect(result[0]).not.toHaveProperty('results') expect(result[1]).toEqual(messages[1]) expect(result[2]).toEqual(messages[2]) expect(result[3]).not.toHaveProperty('results') }) it('strips update_notebook previous_content before re-uploading to the API', () => { const messages = [createAssistantMessageWithUpdateNotebookTool()] const result = prepareMessagesForAPI(messages) expect((result[0].parts[0] as ToolUIPart).output).toEqual({ id: 'notebook-1', name: 'Signup funnel', }) }) it('does not mutate the original message parts when stripping previous_content', () => { const messages = [createAssistantMessageWithUpdateNotebookTool()] const originalParts = messages[0].parts prepareMessagesForAPI(messages) expect(messages[0].parts).toBe(originalParts) expect((originalParts[0] as ToolUIPart).output).toHaveProperty('previous_content') }) it('leaves an update_notebook output without previous_content unchanged', () => { const messages = [ createAssistantMessageWithUpdateNotebookTool({ id: 'notebook-1', name: 'Signup funnel' }), ] const result = prepareMessagesForAPI(messages) expect((result[0].parts[0] as ToolUIPart).output).toEqual({ id: 'notebook-1', name: 'Signup funnel', }) }) })