| 1 | import { ToolMessage } from '@langchain/core/messages' |
| 2 | import { describe, expect, it, vi } from 'vitest' |
| 3 | import { |
| 4 | createThinkingToolInputRecoveryMiddleware, |
| 5 | isThinkingToolInputError |
| 6 | } from '../../../src/main/thinking/tool-error-recovery' |
| 7 | |
| 8 | type ToolHandler = (request: Record<string, unknown>) => Promise<ToolMessage> |
| 9 | |
| 10 | describe('thinking tool input recovery', () => { |
| 11 | it('recognizes workflow tool schema validation errors', () => { |
| 12 | expect( |
| 13 | isThinkingToolInputError( |
| 14 | new Error( |
| 15 | 'Error invoking tool update_context_document: Received tool input did not match expected schema' |
| 16 | ) |
| 17 | ) |
| 18 | ).toBe(true) |
| 19 | expect(isThinkingToolInputError(new Error('disk write failed'))).toBe(false) |
| 20 | }) |
| 21 | |
| 22 | it('returns the first schema error to the model so it can correct the arguments', async () => { |
| 23 | const middleware = createThinkingToolInputRecoveryMiddleware() |
| 24 | const wrapToolCall = middleware.wrapToolCall as unknown as ( |
| 25 | request: Record<string, unknown>, |
| 26 | handler: ToolHandler |
| 27 | ) => Promise<ToolMessage> |
| 28 | const parsingError = new Error('Received tool input did not match expected schema') |
| 29 | const handler = vi.fn<ToolHandler>().mockRejectedValue(parsingError) |
| 30 | const request = { |
| 31 | toolCall: { |
| 32 | name: 'update_context_document', |
| 33 | id: 'call-1' |
| 34 | } |
| 35 | } |
| 36 | |
| 37 | const result = await wrapToolCall(request, handler) |
| 38 | |
| 39 | expect(ToolMessage.isInstance(result)).toBe(true) |
| 40 | expect(result.status).toBe('error') |
| 41 | expect(result.tool_call_id).toBe('call-1') |
| 42 | expect(String(result.content)).toContain('match the declared schema') |
| 43 | await expect(wrapToolCall(request, handler)).rejects.toBe(parsingError) |
| 44 | }) |
| 45 | |
| 46 | it('does not swallow non-schema errors or errors from unrelated tools', async () => { |
| 47 | const middleware = createThinkingToolInputRecoveryMiddleware() |
| 48 | const wrapToolCall = middleware.wrapToolCall as unknown as ( |
| 49 | request: Record<string, unknown>, |
| 50 | handler: ToolHandler |
| 51 | ) => Promise<ToolMessage> |
| 52 | const runtimeError = new Error('disk write failed') |
| 53 | const schemaError = new Error('Received tool input did not match expected schema') |
| 54 | |
| 55 | await expect( |
| 56 | wrapToolCall( |
| 57 | { toolCall: { name: 'update_context_document', id: 'call-1' } }, |
| 58 | vi.fn<ToolHandler>().mockRejectedValue(runtimeError) |
| 59 | ) |
| 60 | ).rejects.toBe(runtimeError) |
| 61 | await expect( |
| 62 | wrapToolCall( |
| 63 | { toolCall: { name: 'read_file', id: 'call-2' } }, |
| 64 | vi.fn<ToolHandler>().mockRejectedValue(schemaError) |
| 65 | ) |
| 66 | ).rejects.toBe(schemaError) |
| 67 | }) |
| 68 | }) |
| 69 |