diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 05f972a1e67..fcf34e0fff8 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -2543,7 +2543,7 @@ describe('Session', () => { mockToolRegistry.getTool.mockReturnValue(tool); mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); mockClient.extMethod = vi.fn().mockResolvedValue({ - messages: ['please also check tests'], + messages: [' please also check tests '], }); mockChat.sendMessageStream = vi .fn() @@ -2576,14 +2576,821 @@ describe('Session', () => { ); const secondCall = vi.mocked(mockChat.sendMessageStream).mock.calls[1]; const midTurnPart = { - text: '\n[User message received during tool execution]: please also check tests', + text: '\n[User message received during tool execution]: please also check tests ', }; expect(secondCall?.[1].message).toEqual( expect.arrayContaining([midTurnPart]), ); expect( mockChatRecordingService.recordMidTurnUserMessage, - ).toHaveBeenCalledWith([midTurnPart], 'please also check tests'); + ).toHaveBeenCalledWith([midTurnPart], ' please also check tests '); + }); + + it('injects drained structured mid-turn user messages with images', async () => { + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [ + { type: 'text', text: 'please inspect this image' }, + { + type: 'image', + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + { + type: 'audio', + mimeType: 'audio/wav', + data: 'UklGRgAAAA==', + }, + { + type: 'image', + mimeType: 'text/html', + data: '', + }, + { + type: 'audio', + mimeType: 'text/plain', + data: 'not-audio', + }, + { + type: 'video', + mimeType: 'video/mp4', + data: 'not-supported', + }, + ], + displayText: 'please inspect this image', + }, + ], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const midTurnParts: Part[] = [ + { + text: '\n[User message received during tool execution]: please inspect this image', + }, + { + inlineData: { + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + }, + { + inlineData: { + mimeType: 'audio/wav', + data: 'UklGRgAAAA==', + }, + }, + ]; + const secondCall = vi.mocked(mockChat.sendMessageStream).mock.calls[1]; + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining(midTurnParts), + ); + expect(secondCall?.[1].message).not.toEqual( + expect.arrayContaining([ + { + inlineData: { + mimeType: 'text/html', + data: '', + }, + }, + ]), + ); + expect(secondCall?.[1].message).not.toEqual( + expect.arrayContaining([ + { + inlineData: { + mimeType: 'text/plain', + data: 'not-audio', + }, + }, + ]), + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith(midTurnParts, 'please inspect this image'); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Unknown ContentBlock type: video', + ); + }); + + it('keeps later structured mid-turn messages when one resolution fails', async () => { + const clampSpy = vi + .spyOn(core, 'clampInlineMediaPart') + .mockImplementation(() => { + throw new Error('image decode failed'); + }); + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [ + { + type: 'image', + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + ], + displayText: 'please inspect this image', + }, + { + content: [{ type: 'text', text: 'safe follow-up' }], + displayText: 'safe follow-up', + }, + ], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + try { + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const fallbackPart = { + text: '\n[User message received during tool execution]: please inspect this image', + }; + const attachmentFailurePart = { + text: '[Attachment could not be processed]', + }; + const followUpPart = { + text: '\n[User message received during tool execution]: safe follow-up', + }; + const secondCall = vi.mocked(mockChat.sendMessageStream).mock + .calls[1]; + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining([ + fallbackPart, + attachmentFailurePart, + followUpPart, + ]), + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith( + [fallbackPart, attachmentFailurePart], + 'please inspect this image', + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith([followUpPart], 'safe follow-up'); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Failed to resolve mid-turn message: image decode failed', + ); + } finally { + clampSpy.mockRestore(); + } + }); + + it('adds a fallback marker when audio resolution fails', async () => { + const clampSpy = vi + .spyOn(core, 'clampInlineMediaPart') + .mockImplementation(() => { + throw new Error('audio decode failed'); + }); + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [ + { + type: 'audio', + mimeType: 'audio/wav', + data: 'UklGRgAAAA==', + }, + ], + displayText: 'please listen to this audio', + }, + ], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + try { + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const fallbackPart = { + text: '\n[User message received during tool execution]: please listen to this audio', + }; + const attachmentFailurePart = { + text: '[Attachment could not be processed]', + }; + const secondCall = vi.mocked(mockChat.sendMessageStream).mock + .calls[1]; + + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining([fallbackPart, attachmentFailurePart]), + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith( + [fallbackPart, attachmentFailurePart], + 'please listen to this audio', + ); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Failed to resolve mid-turn message: audio decode failed', + ); + } finally { + clampSpy.mockRestore(); + } + }); + + it('caps structured mid-turn drain items', async () => { + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: Array.from({ length: 12 }, (_value, index) => ({ + content: [{ type: 'text', text: `mid-turn ${index}` }], + displayText: `mid-turn ${index}`, + })), + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const secondCall = vi.mocked(mockChat.sendMessageStream).mock.calls[1]; + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining([ + { + text: '\n[User message received during tool execution]: mid-turn 0', + }, + { + text: '\n[User message received during tool execution]: mid-turn 9', + }, + ]), + ); + expect(secondCall?.[1].message).not.toEqual( + expect.arrayContaining([ + { + text: '\n[User message received during tool execution]: mid-turn 10', + }, + ]), + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledTimes(10); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Mid-turn drain response had 12 item(s); processing first 10', + ); + }); + + it('stops draining mid-turn messages when structured resolution is aborted', async () => { + let promptSignalAborted = false; + const clampSpy = vi + .spyOn(core, 'clampInlineMediaPart') + .mockImplementation(() => { + const pendingPrompt = ( + session as unknown as { pendingPrompt: AbortController | null } + ).pendingPrompt; + pendingPrompt?.abort(); + promptSignalAborted = pendingPrompt?.signal.aborted ?? false; + const abortError = new Error('aborted'); + abortError.name = 'AbortError'; + throw abortError; + }); + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [{ type: 'text', text: 'already queued' }], + displayText: 'already queued', + }, + { + content: [ + { + type: 'image', + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + ], + displayText: 'inspect this image', + }, + { + content: [{ type: 'text', text: 'should not be processed' }], + displayText: 'should not be processed', + }, + ], + }); + mockChat.sendMessageStream = vi.fn().mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ); + + try { + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const retainedMidTurnPart = { + text: '\n[User message received during tool execution]: already queued', + }; + const abortedMidTurnPart = { + text: '\n[User message received during tool execution]: inspect this image', + }; + const skippedMidTurnPart = { + text: '\n[User message received during tool execution]: should not be processed', + }; + const preservedMessage = vi.mocked(mockChat.addHistory).mock + .calls[0]?.[0] as Content | undefined; + + expect(promptSignalAborted).toBe(true); + expect(clampSpy).toHaveBeenCalledTimes(1); + expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1); + expect(preservedMessage?.parts).toEqual( + expect.arrayContaining([retainedMidTurnPart]), + ); + expect(preservedMessage?.parts).not.toEqual( + expect.arrayContaining([abortedMidTurnPart]), + ); + expect(preservedMessage?.parts).not.toEqual( + expect.arrayContaining([skippedMidTurnPart]), + ); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith([retainedMidTurnPart], 'already queued'); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).not.toHaveBeenCalledWith( + [skippedMidTurnPart], + 'should not be processed', + ); + } finally { + clampSpy.mockRestore(); + } + }); + + it('logs unrecognized mid-turn drain response fields', async () => { + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + payload: ['safe follow-up'], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + "Mid-turn drain response had no recognized 'items' or 'messages' field; keys: payload", + ); + }); + + it('rejects mid-turn resource links and keeps valid messages in the same batch', async () => { + const readManyFilesSpy = vi + .spyOn(core, 'readManyFiles') + .mockResolvedValue({ + contentParts: 'secret file', + files: [], + }); + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [ + { type: 'text', text: 'mixed safe follow-up' }, + { + type: 'resource_link', + uri: 'file:///etc/passwd', + name: 'passwd', + }, + ], + displayText: 'mixed safe follow-up', + }, + { + content: [ + { + type: 'resource_link', + uri: 'file:///etc/passwd', + name: 'passwd', + }, + ], + displayText: 'secret file', + }, + { + content: [{ type: 'text', text: 'safe follow-up' }], + displayText: 'safe follow-up', + }, + ], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + try { + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const mixedMidTurnPart = { + text: '\n[User message received during tool execution]: mixed safe follow-up', + }; + const midTurnPart = { + text: '\n[User message received during tool execution]: safe follow-up', + }; + const secondCall = vi.mocked(mockChat.sendMessageStream).mock + .calls[1]; + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining([mixedMidTurnPart, midTurnPart]), + ); + expect(readManyFilesSpy).not.toHaveBeenCalled(); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith([mixedMidTurnPart], 'mixed safe follow-up'); + expect( + mockChatRecordingService.recordMidTurnUserMessage, + ).toHaveBeenCalledWith([midTurnPart], 'safe follow-up'); + } finally { + readManyFilesSpy.mockRestore(); + } + }); + + it('accepts valid mid-turn embedded resources and drops invalid ones', async () => { + const executeSpy = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + const tool = { + name: 'read_file', + kind: core.Kind.Read, + build: vi.fn().mockReturnValue({ + params: { path: '/tmp/test.txt' }, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue('Read file'), + toolLocations: vi.fn().mockReturnValue([]), + execute: executeSpy, + }), + }; + + mockToolRegistry.getTool.mockReturnValue(tool); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockClient.extMethod = vi.fn().mockResolvedValue({ + items: [ + { + content: [ + { + type: 'resource', + resource: { + uri: 'file:///notes.txt', + text: 'note contents', + }, + }, + ], + displayText: 'read embedded notes', + }, + { + content: [ + { + type: 'resource', + resource: { + uri: 'file:///image.png', + mimeType: 'image/png', + blob: 'iVBORw0KGgo=', + }, + }, + ], + displayText: 'read embedded image', + }, + { + content: [ + { + type: 'resource', + resource: { + uri: 'file:///invalid.txt', + }, + }, + ], + displayText: 'invalid resource', + }, + { + content: [ + { + type: 'resource', + resource: { + uri: 'file:///huge.txt', + text: 'x'.repeat(100_001), + }, + }, + ], + displayText: 'huge resource', + }, + ], + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { path: '/tmp/test.txt' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + debugLoggerWarnSpy.mockClear(); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read file' }], + }); + + const secondCall = vi.mocked(mockChat.sendMessageStream).mock.calls[1]; + expect(secondCall?.[1].message).toEqual( + expect.arrayContaining([ + { + text: '\n[User message received during tool execution]: @file:///notes.txt', + }, + { + text: 'File: file:///notes.txt\nnote contents', + }, + { + text: '\n[User message received during tool execution]: @file:///image.png', + }, + { + inlineData: { + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + }, + ]), + ); + expect(secondCall?.[1].message).not.toEqual( + expect.arrayContaining([ + { + text: '\n[User message received during tool execution]: invalid resource', + }, + { + text: '\n[User message received during tool execution]: huge resource', + }, + ]), + ); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Dropped 1 invalid mid-turn content block(s): "invalid resource"', + ); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + 'Dropped 1 invalid mid-turn content block(s): "huge resource"', + ); }); it('latches mid-turn drain off after a permanent (-32601) error', async () => { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index d1e3b799f69..0b6f4f125ba 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -123,6 +123,7 @@ import type { import type { LoadedSettings } from '../../config/settings.js'; import { z } from 'zod'; import { normalizePartList } from '../../utils/nonInteractiveHelpers.js'; +import { prefixMidTurnUserMessageParts } from '../../utils/midTurnUserMessage.js'; import { handleSlashCommand, getAvailableCommands, @@ -184,12 +185,205 @@ const ASK_USER_QUESTION_CANCEL_SKIP_MESSAGE = // means the client silently drops unknown methods; without a deadline the // await would wedge the prompt turn forever. const MID_TURN_QUEUE_DRAIN_TIMEOUT_MS = 2_000; +const MID_TURN_QUEUE_RESOLVE_TIMEOUT_MS = 10_000; +const MAX_MID_TURN_DRAIN_ITEMS = 10; +const MID_TURN_ATTACHMENT_PROCESSING_FAILURE_TEXT = + '[Attachment could not be processed]'; +const MAX_MID_TURN_RESOURCE_TEXT_LENGTH = 100_000; // Latch the drain off only after this many consecutive timeouts: one slow // answer must not permanently disable mid-turn messages for a // conforming-but-busy client, while a client that never answers stops // costing a stall per tool batch after a few batches. const MID_TURN_QUEUE_DRAIN_MAX_TIMEOUT_STRIKES = 3; +type DrainedMidTurnMessage = + | { kind: 'text'; message: string } + | { kind: 'structured'; content: ContentBlock[]; displayText: string }; + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object'; +} + +function isContentBlock(value: unknown): value is ContentBlock { + if (!isRecord(value) || typeof value['type'] !== 'string') return false; + + switch (value['type']) { + case 'text': + return typeof value['text'] === 'string'; + case 'image': + return ( + typeof value['mimeType'] === 'string' && + value['mimeType'].startsWith('image/') && + typeof value['data'] === 'string' + ); + case 'audio': + return ( + typeof value['mimeType'] === 'string' && + value['mimeType'].startsWith('audio/') && + typeof value['data'] === 'string' + ); + case 'resource_link': + return false; + case 'resource': + return isEmbeddedResourceResource(value['resource']); + default: + debugLogger.warn(`Unknown ContentBlock type: ${value['type']}`); + return false; + } +} + +async function withTimeoutSignal( + parentSignal: AbortSignal, + timeoutMs: number, + fn: (signal: AbortSignal) => Promise, +): Promise { + const signal = AbortSignal.any([ + parentSignal, + AbortSignal.timeout(timeoutMs), + ]); + + const toAbortError = () => + signal.reason instanceof Error + ? signal.reason + : new Error('Mid-turn message resolution aborted'); + + if (signal.aborted) throw toAbortError(); + + let rejectOnAbort: (() => void) | undefined; + const abortPromise = new Promise((_, reject) => { + rejectOnAbort = () => reject(toAbortError()); + signal.addEventListener('abort', rejectOnAbort, { once: true }); + if (signal.aborted) rejectOnAbort(); + }); + + try { + return await Promise.race([fn(signal), abortPromise]); + } finally { + if (rejectOnAbort) signal.removeEventListener('abort', rejectOnAbort); + } +} + +function isEmbeddedResourceResource( + value: unknown, +): value is EmbeddedResourceResource { + if (!isRecord(value) || typeof value['uri'] !== 'string') return false; + if (typeof value['text'] === 'string') { + return value['text'].length <= MAX_MID_TURN_RESOURCE_TEXT_LENGTH; + } + return typeof value['blob'] === 'string'; +} + +function hasInlineMediaContentBlock(content: ContentBlock[]): boolean { + return content.some((part) => part.type === 'image' || part.type === 'audio'); +} + +function capMidTurnDrainItems(items: T[], fieldName: string): T[] { + if (items.length <= MAX_MID_TURN_DRAIN_ITEMS) return items; + + debugLogger.warn( + `Mid-turn drain response had ${items.length} ${fieldName}; processing first ${MAX_MID_TURN_DRAIN_ITEMS}`, + ); + return items.slice(0, MAX_MID_TURN_DRAIN_ITEMS); +} + +function getMidTurnItemDisplayTextForLog(displayText: unknown): string { + if (typeof displayText !== 'string' || displayText.trim().length === 0) { + return '(no display text)'; + } + return JSON.stringify(displayText.trim().slice(0, 120)); +} + +function getValidMidTurnContentBlocks( + content: unknown, + displayText: unknown, +): ContentBlock[] { + if (!Array.isArray(content)) { + debugLogger.warn( + `Dropped invalid mid-turn item: ${getMidTurnItemDisplayTextForLog( + displayText, + )}`, + ); + return []; + } + + const validBlocks = content.filter(isContentBlock); + const invalidBlockCount = content.length - validBlocks.length; + if (invalidBlockCount > 0) { + debugLogger.warn( + `Dropped ${invalidBlockCount} invalid mid-turn content block(s): ${getMidTurnItemDisplayTextForLog( + displayText, + )}`, + ); + } + + return validBlocks; +} + +function getStructuredMidTurnDisplayText( + content: ContentBlock[], + displayText: unknown, +): string { + if (typeof displayText === 'string' && displayText.trim().length > 0) { + return displayText.trim(); + } + + const text = content + .filter( + (part): part is Extract => + part.type === 'text', + ) + .map((part) => part.text) + .join('\n') + .trim(); + + return text || '[User message with attachments]'; +} + +function parseMidTurnDrainResponse(response: unknown): DrainedMidTurnMessage[] { + if (!isRecord(response)) return []; + + if (Array.isArray(response['items'])) { + return capMidTurnDrainItems(response['items'], 'item(s)').flatMap( + (item): DrainedMidTurnMessage[] => { + if (!isRecord(item)) { + return []; + } + const content = getValidMidTurnContentBlocks( + item['content'], + item['displayText'], + ); + if (content.length === 0) return []; + return [ + { + kind: 'structured', + content, + displayText: getStructuredMidTurnDisplayText( + content, + item['displayText'], + ), + }, + ]; + }, + ); + } + + if (!Array.isArray(response['messages'])) { + debugLogger.warn( + `Mid-turn drain response had no recognized 'items' or 'messages' field; keys: ${Object.keys( + response, + ).join(', ')}`, + ); + return []; + } + + return capMidTurnDrainItems(response['messages'], 'message(s)') + .filter( + (message): message is string => + typeof message === 'string' && message.trim().length > 0, + ) + .map((message) => ({ kind: 'text', message })); +} + class MidTurnDrainTimeoutError extends Error { constructor() { super( @@ -1334,6 +1528,7 @@ export class Session implements SessionContext { if (toolRun.stopAfterUserQuestionCancel) { await this.#preserveCancelledAskUserQuestionToolRun( toolRun, + pendingSend.signal, ); return { stopReason: 'end_turn' }; } @@ -1341,7 +1536,9 @@ export class Session implements SessionContext { role: 'user', parts: [ ...toolRun.parts, - ...(await this.#drainMidTurnUserMessages()), + ...(await this.#drainMidTurnUserMessages( + pendingSend.signal, + )), ], }; } @@ -1603,14 +1800,17 @@ export class Session implements SessionContext { functionCalls, ); if (toolRun.stopAfterUserQuestionCancel) { - await this.#preserveCancelledAskUserQuestionToolRun(toolRun); + await this.#preserveCancelledAskUserQuestionToolRun( + toolRun, + pendingSend.signal, + ); return { stopReason: 'end_turn' }; } nextMessage = { role: 'user', parts: [ ...toolRun.parts, - ...(await this.#drainMidTurnUserMessages()), + ...(await this.#drainMidTurnUserMessages(pendingSend.signal)), ], }; } @@ -1775,11 +1975,15 @@ export class Session implements SessionContext { async #preserveCancelledAskUserQuestionToolRun( toolRun: RunToolResult, + abortSignal: AbortSignal, ): Promise { this.#preserveUnsentMessageHistory( { role: 'user', - parts: [...toolRun.parts, ...(await this.#drainMidTurnUserMessages())], + parts: [ + ...toolRun.parts, + ...(await this.#drainMidTurnUserMessages(abortSignal)), + ], }, true, ); @@ -1892,7 +2096,7 @@ export class Session implements SessionContext { }); } - async #drainMidTurnUserMessages(): Promise { + async #drainMidTurnUserMessages(abortSignal: AbortSignal): Promise { if (this.midTurnDrainUnavailable) return []; let drainPromise: ReturnType | undefined; @@ -1914,28 +2118,49 @@ export class Session implements SessionContext { clearTimeout(timeoutHandle); } this.midTurnDrainTimeoutStrikes = 0; - // A client may legally resolve with `result: null` (passed through - // unwrapped by the ACP SDK); guard the object access so that doesn't - // throw a TypeError and get misclassified as a transient drain error. - const messages = - response && - typeof response === 'object' && - Array.isArray(response['messages']) - ? response['messages'].filter( - (message): message is string => - typeof message === 'string' && message.trim().length > 0, - ) - : []; - - return messages.map((message) => { - const part = { - text: `\n[User message received during tool execution]: ${message}`, - }; + const drainedMessages = parseMidTurnDrainResponse(response); + const drainedParts: Part[] = []; + for (const message of drainedMessages) { + const displayText = + message.kind === 'text' ? message.message : message.displayText; + let rawParts: Part[]; + try { + rawParts = + message.kind === 'text' + ? [{ text: message.message }] + : await withTimeoutSignal( + abortSignal, + MID_TURN_QUEUE_RESOLVE_TIMEOUT_MS, + (signal) => this.#resolvePrompt(message.content, signal), + ); + } catch (messageError) { + if (abortSignal.aborted) return drainedParts; + const errorMessage = this.#formatError(messageError); + debugLogger.warn( + `Failed to resolve mid-turn message: ${errorMessage}`, + ); + rawParts = [ + { + text: displayText, + }, + ]; + if ( + message.kind === 'structured' && + hasInlineMediaContentBlock(message.content) + ) { + rawParts.push({ + text: MID_TURN_ATTACHMENT_PROCESSING_FAILURE_TEXT, + }); + } + } + const parts = prefixMidTurnUserMessageParts(rawParts, displayText); this.config .getChatRecordingService() - ?.recordMidTurnUserMessage([part], message); - return part; - }); + ?.recordMidTurnUserMessage(parts, displayText); + drainedParts.push(...parts); + } + + return drainedParts; } catch (error) { // The ACP SDK rejects with the raw JSON-RPC error object // (`{ code, message, data }`), which is not an `Error` instance, so @@ -2196,6 +2421,7 @@ export class Session implements SessionContext { if (toolRun.stopAfterUserQuestionCancel) { await this.#preserveCancelledAskUserQuestionToolRun( toolRun, + ac.signal, ); return; } @@ -2203,7 +2429,7 @@ export class Session implements SessionContext { role: 'user', parts: [ ...toolRun.parts, - ...(await this.#drainMidTurnUserMessages()), + ...(await this.#drainMidTurnUserMessages(ac.signal)), ], }; } @@ -2506,7 +2732,10 @@ export class Session implements SessionContext { functionCalls, ); if (toolRun.stopAfterUserQuestionCancel) { - await this.#preserveCancelledAskUserQuestionToolRun(toolRun); + await this.#preserveCancelledAskUserQuestionToolRun( + toolRun, + ac.signal, + ); await this.#emitBackgroundNotificationEndTurn('end_turn'); return; } @@ -2514,7 +2743,7 @@ export class Session implements SessionContext { role: 'user', parts: [ ...toolRun.parts, - ...(await this.#drainMidTurnUserMessages()), + ...(await this.#drainMidTurnUserMessages(ac.signal)), ], }; } diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index 2bb5a578d27..c497db2bc1d 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -22,22 +22,35 @@ import type { } from '../types.js'; import { ToolCallStatus } from '../types.js'; -interface HandleAtCommandParams { +export interface ResolveAtCommandParams { query: string; config: Config; onDebugMessage: (message: string) => void; messageId: number; signal: AbortSignal; +} + +interface HandleAtCommandParams extends ResolveAtCommandParams { addItem?: (item: HistoryItemWithoutId, baseTimestamp: number) => number; } -interface HandleAtCommandResult { +export interface HandleAtCommandResult { processedQuery: PartListUnion | null; shouldProceed: boolean; toolDisplays?: IndividualToolCallDisplay[]; filesRead?: string[]; } +export interface AtCommandRecording { + filesRead: string[]; + status: 'success' | 'error'; + message?: string; +} + +export interface ResolveAtCommandResult extends HandleAtCommandResult { + recording?: AtCommandRecording; +} + interface AtCommandPart { type: 'text' | 'atPath'; content: string; @@ -126,30 +139,18 @@ function parseAllAtCommands(query: string): AtCommandPart[] { * @returns An object indicating whether the main hook should proceed with an * LLM call and the processed query parts (including file content). */ -export async function handleAtCommand({ +export async function resolveAtCommandQuery({ query, config, onDebugMessage, messageId: userMessageTimestamp, signal, - addItem, -}: HandleAtCommandParams): Promise { +}: ResolveAtCommandParams): Promise { const commandParts = parseAllAtCommands(query); const atPathCommandParts = commandParts.filter( (part) => part.type === 'atPath', ); - const addToolGroup = (result: HandleAtCommandResult): void => { - if (!addItem) return; - if (result.toolDisplays && result.toolDisplays.length > 0) { - const toolGroupItem: HistoryItemToolGroup = { - type: 'tool_group', - tools: result.toolDisplays, - }; - addItem(toolGroupItem, userMessageTimestamp); - } - }; - if (atPathCommandParts.length === 0) { return { processedQuery: [{ text: query }], shouldProceed: true }; } @@ -395,15 +396,13 @@ export async function handleAtCommand({ filesRead: contentLabelsForDisplay, }; - const chatRecorder = config.getChatRecordingService?.(); - chatRecorder?.recordAtCommand({ - filesRead: contentLabelsForDisplay, - status: 'success', - userText: query, - }); - - addToolGroup(processedResult); - return processedResult; + return { + ...processedResult, + recording: { + filesRead: contentLabelsForDisplay, + status: 'success', + }, + }; } catch (error: unknown) { const errorToolCallDisplay: IndividualToolCallDisplay = { callId: `client-read-${userMessageTimestamp}`, @@ -413,24 +412,53 @@ export async function handleAtCommand({ resultDisplay: `Error reading files (${contentLabelsForDisplay.join(', ')}): ${getErrorMessage(error)}`, confirmationDetails: undefined, }; - const chatRecorder = config.getChatRecordingService?.(); const errorMessage = typeof errorToolCallDisplay.resultDisplay === 'string' ? errorToolCallDisplay.resultDisplay : undefined; - chatRecorder?.recordAtCommand({ - filesRead: contentLabelsForDisplay, - status: 'error', - message: errorMessage, - userText: query, - }); - const result = { + return { processedQuery: null, shouldProceed: false, toolDisplays: [errorToolCallDisplay], filesRead: contentLabelsForDisplay, + recording: { + filesRead: contentLabelsForDisplay, + status: 'error', + message: errorMessage, + }, }; - addToolGroup(result); - return result; } } + +export async function handleAtCommand( + params: HandleAtCommandParams, +): Promise { + const result = await resolveAtCommandQuery(params); + + if (result.recording) { + const chatRecorder = params.config.getChatRecordingService?.(); + chatRecorder?.recordAtCommand({ + filesRead: result.recording.filesRead, + status: result.recording.status, + ...(result.recording.message + ? { message: result.recording.message } + : {}), + userText: params.query, + }); + } + + if (params.addItem && result.toolDisplays && result.toolDisplays.length > 0) { + const toolGroupItem: HistoryItemToolGroup = { + type: 'tool_group', + tools: result.toolDisplays, + }; + params.addItem(toolGroupItem, params.messageId); + } + + return { + processedQuery: result.processedQuery, + shouldProceed: result.shouldProceed, + toolDisplays: result.toolDisplays, + filesRead: result.filesRead, + }; +} diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 922bfca5f76..aeabc18d77c 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -35,7 +35,7 @@ import { import type { Part, PartListUnion } from '@google/genai'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import type { HistoryItem, SlashCommandProcessorResult } from '../types.js'; -import { MessageType, StreamingState } from '../types.js'; +import { MessageType, StreamingState, ToolCallStatus } from '../types.js'; import type { LoadedSettings } from '../../config/settings.js'; import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; @@ -720,7 +720,7 @@ describe('useGeminiStream', () => { text: `\n[User message received during tool execution]: ${queuedPrompt}`, }; expect(recordMidTurnUserMessage).toHaveBeenCalledWith( - expectedMidTurnMessage, + [expectedMidTurnMessage], queuedPrompt, ); const queuedPromptAddItemIndex = mockAddItem.mock.calls.findIndex( @@ -746,6 +746,731 @@ describe('useGeminiStream', () => { ); }); + it('resolves mid-turn @ image messages before submitting tool results', async () => { + const queuedPrompt = 'inspect @/tmp/screenshot.png'; + const resolvedImagePart: Part = { + inlineData: { + mimeType: 'image/png', + data: 'iVBORw0KGgo=', + }, + }; + const resolvedTextPart: Part = { text: 'inspect @/tmp/screenshot.png' }; + const recordMidTurnUserMessage = vi.fn(); + const recordAtCommand = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordAtCommand, + recordMidTurnUserMessage, + }); + const resolveAtCommandQuerySpy = vi + .spyOn(atCommandProcessor, 'resolveAtCommandQuery') + .mockResolvedValue({ + processedQuery: [resolvedTextPart, resolvedImagePart], + shouldProceed: true, + recording: { + filesRead: ['/tmp/screenshot.png'], + status: 'success', + }, + }); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-image', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(completedToolCalls); + } + }); + + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + const expectedMidTurnParts: Part[] = [ + { + text: `\n[User message received during tool execution]: ${resolvedTextPart.text}`, + }, + resolvedImagePart, + ]; + expect(resolveAtCommandQuerySpy).toHaveBeenCalledWith( + expect.objectContaining({ + query: queuedPrompt, + config: mockConfig, + onDebugMessage: mockOnDebugMessage, + signal: expect.any(AbortSignal), + }), + ); + expect(recordMidTurnUserMessage).toHaveBeenCalledWith( + expectedMidTurnParts, + queuedPrompt, + ); + expect(recordAtCommand).toHaveBeenCalledWith({ + filesRead: ['/tmp/screenshot.png'], + status: 'success', + userText: queuedPrompt, + }); + expect(handleAtCommandSpy).not.toHaveBeenCalled(); + expect(mockSendMessageStream).toHaveBeenCalledWith( + [...toolCallResponseParts, ...expectedMidTurnParts], + expect.any(AbortSignal), + 'prompt-id-midturn-image', + { type: SendMessageType.ToolResult }, + ); + }); + + it('skips mid-turn @ injection when resolution should not proceed', async () => { + const queuedPrompt = 'inspect @/tmp/missing.png'; + const recordMidTurnUserMessage = vi.fn(); + const recordAtCommand = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordAtCommand, + recordMidTurnUserMessage, + }); + const toolDisplays = [ + { + callId: 'client-read-midturn-at-error', + name: 'Read File(s)', + description: 'Error attempting to read files', + status: ToolCallStatus.Error, + resultDisplay: 'Error reading files (/tmp/missing.png): not found', + confirmationDetails: undefined, + }, + ]; + const resolveAtCommandQuerySpy = vi + .spyOn(atCommandProcessor, 'resolveAtCommandQuery') + .mockResolvedValue({ + processedQuery: null, + shouldProceed: false, + toolDisplays, + recording: { + filesRead: ['/tmp/missing.png'], + status: 'error', + message: 'Error reading files (/tmp/missing.png): not found', + }, + }); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-at-error', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(completedToolCalls); + } + }); + + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + expect(resolveAtCommandQuerySpy).toHaveBeenCalledWith( + expect.objectContaining({ + query: queuedPrompt, + config: mockConfig, + onDebugMessage: mockOnDebugMessage, + signal: expect.any(AbortSignal), + }), + ); + expect(recordAtCommand).toHaveBeenCalledWith({ + filesRead: ['/tmp/missing.png'], + status: 'error', + message: 'Error reading files (/tmp/missing.png): not found', + userText: queuedPrompt, + }); + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'tool_group', + tools: toolDisplays, + }, + expect.any(Number), + ); + expect(recordMidTurnUserMessage).not.toHaveBeenCalled(); + expect(mockSendMessageStream).toHaveBeenCalledWith( + toolCallResponseParts, + expect.any(AbortSignal), + 'prompt-id-midturn-at-error', + { type: SendMessageType.ToolResult }, + ); + }); + + it('warns and skips mid-turn @ injection when resolution fails', async () => { + const queuedPrompt = 'inspect @/tmp/unreadable.png'; + const recordMidTurnUserMessage = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordMidTurnUserMessage, + }); + vi.spyOn(atCommandProcessor, 'resolveAtCommandQuery').mockRejectedValue( + new Error('permission denied'), + ); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-at-throw', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(completedToolCalls); + } + }); + + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + expect(mockAddItem).toHaveBeenCalledWith( + { + type: MessageType.WARNING, + text: 'Could not attach file: permission denied', + }, + expect.any(Number), + ); + expect(recordMidTurnUserMessage).not.toHaveBeenCalled(); + expect(mockSendMessageStream).toHaveBeenCalledWith( + toolCallResponseParts, + expect.any(AbortSignal), + 'prompt-id-midturn-at-throw', + { type: SendMessageType.ToolResult }, + ); + }); + + it('times out stalled mid-turn @ resolution before submitting tool results', async () => { + vi.useFakeTimers(); + + const queuedPrompt = 'inspect @/tmp/slow.png'; + const recordMidTurnUserMessage = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordMidTurnUserMessage, + }); + let resolveSignal: AbortSignal | undefined; + let rejectResolve: ((error: Error) => void) | undefined; + vi.spyOn(atCommandProcessor, 'resolveAtCommandQuery').mockImplementation( + ({ signal }) => { + resolveSignal = signal; + return new Promise((_, reject) => { + rejectResolve = reject; + signal.addEventListener( + 'abort', + () => + reject( + signal.reason instanceof Error + ? signal.reason + : new Error('aborted'), + ), + { once: true }, + ); + }); + }, + ); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-timeout', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + let completePromise: Promise | undefined; + await act(async () => { + if (capturedOnComplete) { + completePromise = capturedOnComplete(completedToolCalls); + } + }); + + try { + expect(resolveSignal).toBeDefined(); + await act(async () => { + await vi.advanceTimersByTimeAsync(10_000); + }); + expect(resolveSignal?.aborted).toBe(true); + await act(async () => { + await completePromise; + }); + } finally { + if (!resolveSignal?.aborted) { + rejectResolve?.(new Error('cleanup')); + await completePromise; + } + } + + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it('skips mid-turn @ fallback side effects when cancelled during resolution', async () => { + const queuedPrompt = 'inspect @/tmp/cancelled.png'; + const recordMidTurnUserMessage = vi.fn(); + mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ + recordMidTurnUserMessage, + }); + let resolveSignal: AbortSignal | undefined; + vi.spyOn(atCommandProcessor, 'resolveAtCommandQuery').mockImplementation( + ({ signal }) => { + resolveSignal = signal; + return new Promise(() => {}); + }, + ); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-cancel', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + completedToolCalls, + mockScheduleToolCalls, + mockMarkToolsAsSubmitted, + ]; + }); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + let completePromise: Promise | undefined; + act(() => { + if (capturedOnComplete) { + completePromise = capturedOnComplete(completedToolCalls); + } + }); + + await waitFor(() => { + expect(resolveSignal).toBeDefined(); + }); + + await act(async () => { + result.current.cancelOngoingRequest(); + await completePromise; + }); + + expect(resolveSignal?.aborted).toBe(true); + expect(recordMidTurnUserMessage).not.toHaveBeenCalled(); + expect(mockAddItem).not.toHaveBeenCalledWith( + { type: MessageType.NOTIFICATION, text: queuedPrompt }, + expect.any(Number), + ); + expect(mockSendMessageStream).not.toHaveBeenCalled(); + }); + + it('does not show mid-turn @ fallback warnings after cancellation and timeout overlap', async () => { + vi.useFakeTimers(); + + const queuedPrompt = 'inspect @/tmp/cancelled-slow.png'; + vi.spyOn(atCommandProcessor, 'resolveAtCommandQuery').mockImplementation( + () => new Promise(() => {}), + ); + const toolCallResponseParts: Part[] = [ + { + functionResponse: { + id: 'call1', + name: 'testTool', + response: { result: 'ok' }, + }, + }, + ]; + const completedToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'call1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-midturn-cancel-timeout', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call1', + responseParts: toolCallResponseParts, + errorType: undefined, + }, + tool: { + displayName: 'MockTool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as TrackedCompletedToolCall, + ]; + const midTurnDrainRef = { + current: vi.fn().mockReturnValue([queuedPrompt]), + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + completedToolCalls, + mockScheduleToolCalls, + mockMarkToolsAsSubmitted, + ]; + }); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + midTurnDrainRef, + ), + ); + + let completePromise: Promise | undefined; + act(() => { + if (capturedOnComplete) { + completePromise = capturedOnComplete(completedToolCalls); + } + }); + expect(completePromise).toBeDefined(); + + act(() => { + result.current.cancelOngoingRequest(); + vi.advanceTimersByTime(10_000); + }); + await act(async () => { + await completePromise; + }); + + expect(mockAddItem).not.toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.WARNING, + text: expect.stringContaining('Could not attach file:'), + }), + expect.any(Number), + ); + expect(mockSendMessageStream).not.toHaveBeenCalled(); + }); + it('handles mid-turn drain when chat recording is not configured', async () => { const queuedPrompt = 'save the logs locally first'; mockConfig.getChatRecordingService = vi.fn().mockReturnValue(undefined); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index c278d8d9e5b..ff6e1acc10f 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -72,9 +72,13 @@ import { isSlashCommand, } from '../utils/commandUtils.js'; import { useShellCommandProcessor } from './shellCommandProcessor.js'; -import { handleAtCommand } from './atCommandProcessor.js'; +import { + handleAtCommand, + resolveAtCommandQuery, +} from './atCommandProcessor.js'; import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; import { useStateAndRef } from './useStateAndRef.js'; +import { prefixMidTurnUserMessageParts } from '../../utils/midTurnUserMessage.js'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import { useReactToolScheduler, @@ -95,6 +99,9 @@ import { recordGoalStatusItem } from '../utils/restoreGoal.js'; import process from 'node:process'; const debugLogger = createDebugLogger('GEMINI_STREAM'); +const MID_TURN_AT_COMMAND_RESOLVE_TIMEOUT_MS = 10_000; +const MID_TURN_AT_COMMAND_RESOLVE_TIMEOUT_MESSAGE = + 'Mid-turn @ command resolution timed out'; /** * Pull the assistant's most recent visible text from the UI history. Used as @@ -119,6 +126,35 @@ function stripLeadingBlankLines(text: string): string { return text.replace(/^(?:[ \t]*\r?\n)+/, ''); } +async function resolveWithAbort( + signal: AbortSignal, + run: () => Promise, +): Promise { + let onAbort: (() => void) | undefined; + const abortPromise = new Promise((_, reject) => { + onAbort = () => { + reject( + signal.reason instanceof Error + ? signal.reason + : new Error('Mid-turn @ command resolution aborted'), + ); + }; + if (signal.aborted) { + onAbort(); + return; + } + signal.addEventListener('abort', onAbort, { once: true }); + }); + + try { + return await Promise.race([run(), abortPromise]); + } finally { + if (onAbort) { + signal.removeEventListener('abort', onAbort); + } + } +} + /** * Flatten `functionResponse` parts into a compact string for the summarizer. * The summarizer itself truncates to 300 chars per field, so we just join @@ -359,13 +395,11 @@ export const useGeminiStream = ( useLayoutEffect(() => { historyRef.current = history; }, [history]); - // In-flight tool-use-summary aborters. Each batch gets its own AbortController - // because the captured turn controller is replaced when submitQuery starts - // the next turn, and the summary call outlives the current turn (that's the - // whole point — it overlaps with the next turn's streaming). cancelOngoingRequest - // aborts all in-flight summaries so Ctrl+C during the next turn also kills - // this turn's stale summary work. - const summaryAbortRefsRef = useRef>(new Set()); + // In-flight auxiliary work. Some work is batch-scoped rather than turn-scoped: + // summaries intentionally outlive the turn, and mid-turn @ resolution may run + // before submitQuery installs the next turn controller. + // cancelOngoingRequest aborts these controllers so Ctrl+C still cancels them. + const auxiliaryAbortRefsRef = useRef>(new Set()); const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] = useStateAndRef(null); // Streamed model reasoning for the current turn. Rendered (height-limited) @@ -645,12 +679,12 @@ export const useGeminiStream = ( turnCancelledRef.current = true; isSubmittingQueryRef.current = false; abortControllerRef.current?.abort(); - // Cancel any in-flight tool-use-summary generations so their Promise.then - // doesn't addItem a stale label after the user cancelled. - for (const ac of summaryAbortRefsRef.current) { + // Cancel any in-flight auxiliary work so its Promise.then doesn't add + // stale content after the user cancelled. + for (const ac of auxiliaryAbortRefsRef.current) { ac.abort(); } - summaryAbortRefsRef.current.clear(); + auxiliaryAbortRefsRef.current.clear(); // Report cancellation to arena status reporter (if in arena mode). // This is needed because cancellation during tool execution won't @@ -2357,7 +2391,7 @@ export const useGeminiStream = ( // resolve time (which covers both Ctrl+C on the next turn and // mid-flight cancellation of this batch via turnCancelledRef). const summaryAbort = new AbortController(); - summaryAbortRefsRef.current.add(summaryAbort); + auxiliaryAbortRefsRef.current.add(summaryAbort); // Capture the first callId so we can locate "our" tool_group at // resolve time. If a newer tool_group has been added since we @@ -2372,7 +2406,7 @@ export const useGeminiStream = ( lastAssistantText, }) .then((summary) => { - summaryAbortRefsRef.current.delete(summaryAbort); + auxiliaryAbortRefsRef.current.delete(summaryAbort); const cancelled = turnCancelledRef.current || abortControllerRef.current?.signal.aborted || @@ -2409,7 +2443,7 @@ export const useGeminiStream = ( } }) .catch(() => { - summaryAbortRefsRef.current.delete(summaryAbort); + auxiliaryAbortRefsRef.current.delete(summaryAbort); }); } } @@ -2427,16 +2461,128 @@ export const useGeminiStream = ( ? [] : (midTurnDrainRef?.current?.() ?? []); if (drained.length > 0) { - for (const msg of drained) { - const midTurnUserMessage = { - text: `\n[User message received during tool execution]: ${msg}`, - }; - responsesToSend.push(midTurnUserMessage); - config - .getChatRecordingService() - ?.recordMidTurnUserMessage(midTurnUserMessage, msg); - addItem({ type: MessageType.NOTIFICATION, text: msg }, Date.now()); + const midTurnTimestamp = Date.now(); + const midTurnAbort = + abortControllerRef.current ?? new AbortController(); + const shouldTrackMidTurnAbort = !abortControllerRef.current; + if (shouldTrackMidTurnAbort) { + auxiliaryAbortRefsRef.current.add(midTurnAbort); } + try { + for (let index = 0; index < drained.length; index += 1) { + if (midTurnAbort.signal.aborted) { + break; + } + const msg = drained[index]; + let resolvedMidTurnQuery: PartListUnion = [{ text: msg }]; + if (isAtCommand(msg)) { + const atCommandTimeout = new AbortController(); + const atCommandSignal = AbortSignal.any([ + midTurnAbort.signal, + atCommandTimeout.signal, + ]); + const atCommandTimeoutId = setTimeout(() => { + atCommandTimeout.abort( + new Error(MID_TURN_AT_COMMAND_RESOLVE_TIMEOUT_MESSAGE), + ); + }, MID_TURN_AT_COMMAND_RESOLVE_TIMEOUT_MS); + try { + const atCommandResult = await resolveWithAbort( + atCommandSignal, + () => + resolveAtCommandQuery({ + query: msg, + config, + onDebugMessage, + messageId: midTurnTimestamp + index, + signal: atCommandSignal, + }), + ); + const shouldSkipMidTurnMessage = + !atCommandResult.shouldProceed && + (atCommandResult.toolDisplays?.length ?? 0) > 0; + if ( + atCommandResult.shouldProceed && + atCommandResult.processedQuery !== null + ) { + resolvedMidTurnQuery = atCommandResult.processedQuery; + } else if (atCommandResult.toolDisplays?.length) { + addItem( + { type: 'tool_group', tools: atCommandResult.toolDisplays }, + midTurnTimestamp + index, + ); + } + if (atCommandResult.recording) { + config.getChatRecordingService()?.recordAtCommand?.({ + filesRead: atCommandResult.recording.filesRead, + status: atCommandResult.recording.status, + ...(atCommandResult.recording.message + ? { message: atCommandResult.recording.message } + : {}), + userText: msg, + }); + } + if (shouldSkipMidTurnMessage) { + continue; + } + } catch (error) { + const errorMessage = getErrorMessage(error); + onDebugMessage( + `Failed to resolve mid-turn @ command: ${errorMessage}`, + ); + if (!midTurnAbort.signal.aborted) { + addItem( + { + type: MessageType.WARNING, + text: `Could not attach file: ${errorMessage}`, + }, + Date.now(), + ); + } + continue; + } finally { + clearTimeout(atCommandTimeoutId); + } + if (midTurnAbort.signal.aborted) { + break; + } + } + + const midTurnUserMessageParts = prefixMidTurnUserMessageParts( + resolvedMidTurnQuery, + msg, + ); + const formatCheck = checkImageFormatsSupport( + midTurnUserMessageParts, + ); + if (formatCheck.hasUnsupportedFormats) { + addItem( + { + type: MessageType.INFO, + text: getUnsupportedImageFormatWarning(), + }, + Date.now(), + ); + } + responsesToSend.push(...midTurnUserMessageParts); + config + .getChatRecordingService() + ?.recordMidTurnUserMessage(midTurnUserMessageParts, msg); + addItem({ type: MessageType.NOTIFICATION, text: msg }, Date.now()); + } + } finally { + if (shouldTrackMidTurnAbort) { + auxiliaryAbortRefsRef.current.delete(midTurnAbort); + midTurnAbort.abort(); + } + } + } + + if ( + turnCancelledRef.current || + abortControllerRef.current?.signal.aborted + ) { + return; } submitQuery(responsesToSend, SendMessageType.ToolResult, prompt_ids[0]); @@ -2451,6 +2597,7 @@ export const useGeminiStream = ( midTurnDrainRef, addItem, dualOutput, + onDebugMessage, ], ); diff --git a/packages/cli/src/utils/midTurnUserMessage.test.ts b/packages/cli/src/utils/midTurnUserMessage.test.ts new file mode 100644 index 00000000000..ffb5ab91363 --- /dev/null +++ b/packages/cli/src/utils/midTurnUserMessage.test.ts @@ -0,0 +1,40 @@ +/** + * @license + * Copyright 2026 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import type { Part } from '@google/genai'; +import { + MID_TURN_USER_MESSAGE_PREFIX, + prefixMidTurnUserMessageParts, +} from './midTurnUserMessage.js'; + +describe('prefixMidTurnUserMessageParts', () => { + it('returns a text-only part when parts normalize to empty', () => { + expect(prefixMidTurnUserMessageParts([], 'fallback')).toEqual([ + { text: `${MID_TURN_USER_MESSAGE_PREFIX}fallback` }, + ]); + }); + + it('prepends the prefix to the first text part', () => { + const parts: Part[] = [{ text: 'hello' }, { text: 'world' }]; + + expect(prefixMidTurnUserMessageParts(parts, 'hello')).toEqual([ + { text: `${MID_TURN_USER_MESSAGE_PREFIX}hello` }, + { text: 'world' }, + ]); + }); + + it('prepends a text prefix before non-text first parts', () => { + const imagePart: Part = { + inlineData: { mimeType: 'image/png', data: 'abc' }, + }; + + expect(prefixMidTurnUserMessageParts([imagePart], 'inspect this')).toEqual([ + { text: `${MID_TURN_USER_MESSAGE_PREFIX}inspect this` }, + imagePart, + ]); + }); +}); diff --git a/packages/cli/src/utils/midTurnUserMessage.ts b/packages/cli/src/utils/midTurnUserMessage.ts new file mode 100644 index 00000000000..097caeda934 --- /dev/null +++ b/packages/cli/src/utils/midTurnUserMessage.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2026 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Part, PartListUnion } from '@google/genai'; +import { normalizePartList } from './nonInteractiveHelpers.js'; + +export const MID_TURN_USER_MESSAGE_PREFIX = + '\n[User message received during tool execution]: '; + +export function prefixMidTurnUserMessageParts( + parts: PartListUnion, + displayText: string, +): Part[] { + const partArray = normalizePartList(parts); + if (partArray.length === 0) { + return [{ text: `${MID_TURN_USER_MESSAGE_PREFIX}${displayText}` }]; + } + + const [firstPart, ...rest] = partArray; + if ('text' in firstPart && typeof firstPart.text === 'string') { + return [ + { + ...firstPart, + text: `${MID_TURN_USER_MESSAGE_PREFIX}${firstPart.text}`, + }, + ...rest, + ]; + } + + return [ + { text: `${MID_TURN_USER_MESSAGE_PREFIX}${displayText}` }, + ...partArray, + ]; +} diff --git a/packages/desktop/packages/server-core/src/sessions/SessionManager.ts b/packages/desktop/packages/server-core/src/sessions/SessionManager.ts index 370b7eaa41e..a529195973e 100644 --- a/packages/desktop/packages/server-core/src/sessions/SessionManager.ts +++ b/packages/desktop/packages/server-core/src/sessions/SessionManager.ts @@ -305,6 +305,19 @@ export const AGENT_FLAGS = { defaultModesEnabled: true, } as const +function canOfferMidTurnAttachments( + attachments?: FileAttachment[], +): boolean { + if (!attachments?.length) { + return true + } + + return attachments.every((attachment) => { + if (!attachment.mimeType?.startsWith('image/')) return false + return typeof attachment.base64 === 'string' && attachment.base64.length > 0 + }) +} + const MAX_ADMIN_REMEMBER_MINUTES = 60 const MAX_ANNOTATIONS_PER_MESSAGE = 200 const MAX_ANNOTATION_JSON_BYTES = 32 * 1024 @@ -4729,6 +4742,68 @@ export class SessionManager implements ISessionManager { } } + private createMidTurnMessagesDrainedCallback( + managed: ManagedSession, + ): (messageIds: string[]) => void { + return (messageIds: string[]) => { + const drainedEntries: Array<{ + messageId?: string + optimisticMessageId?: string + }> = [] + for (const messageId of messageIds) { + const index = managed.messageQueue.findIndex((entry) => { + if (!entry.midTurnPending) return false + if ( + entry.messageId === messageId || + entry.optimisticMessageId === messageId + ) { + return true + } + if (!entry.messageId && !entry.optimisticMessageId) { + return entry.message === messageId + } + return false + }) + if (index >= 0) { + const [entry] = managed.messageQueue.splice(index, 1) + drainedEntries.push({ + messageId: entry.messageId, + optimisticMessageId: entry.optimisticMessageId, + }) + } + } + if (drainedEntries.length < messageIds.length) { + sessionLog.warn( + `Mid-turn drain acknowledgement matched ${drainedEntries.length}/${messageIds.length} entries for session ${managed.id}`, + ) + } + if (drainedEntries.length > 0) { + sessionLog.info( + `Acknowledged ${drainedEntries.length} mid-turn queued message(s) for session ${managed.id}`, + ) + for (const entry of drainedEntries) { + if (!entry.messageId) continue + const existingMessage = managed.messages.find( + (m) => m.id === entry.messageId, + ) + if (!existingMessage) continue + existingMessage.isQueued = false + this.sendEvent( + { + type: 'user_message', + sessionId: managed.id, + message: existingMessage, + status: 'accepted', + optimisticMessageId: entry.optimisticMessageId, + }, + managed.workspace.id, + ) + } + this.persistSession(managed) + } + } + } + private async createAgentForManagedSession( managed: ManagedSession, ): Promise { @@ -4978,48 +5053,8 @@ export class SessionManager implements ISessionManager { }) } - const onMidTurnMessagesDrained = (messages: string[]) => { - const drainedEntries: Array<{ - messageId?: string - optimisticMessageId?: string - }> = [] - for (const message of messages) { - const index = managed.messageQueue.findIndex( - (entry) => entry.midTurnPending && entry.message === message, - ) - if (index >= 0) { - const [entry] = managed.messageQueue.splice(index, 1) - drainedEntries.push({ - messageId: entry.messageId, - optimisticMessageId: entry.optimisticMessageId, - }) - } - } - if (drainedEntries.length > 0) { - sessionLog.info( - `Acknowledged ${drainedEntries.length} mid-turn queued message(s) for session ${managed.id}`, - ) - for (const entry of drainedEntries) { - if (!entry.messageId) continue - const existingMessage = managed.messages.find( - (m) => m.id === entry.messageId, - ) - if (!existingMessage) continue - existingMessage.isQueued = false - this.sendEvent( - { - type: 'user_message', - sessionId: managed.id, - message: existingMessage, - status: 'accepted', - optimisticMessageId: entry.optimisticMessageId, - }, - managed.workspace.id, - ) - } - this.persistSession(managed) - } - } + const onMidTurnMessagesDrained = + this.createMidTurnMessagesDrainedCallback(managed) // ============================================================ // Construct backend via factory @@ -8097,15 +8132,6 @@ export class SessionManager implements ISessionManager { // injection point is available, keep the message queued for the next turn. if (managed.isProcessing) { const agent = managed.agent - const canInjectMidTurn = - !attachments?.length && - !storedAttachments?.length && - (agent?.enqueueMidTurnMessage?.(message) ?? false) - - sessionLog.info( - `Session ${sessionId} ${canInjectMidTurn ? 'queued message for mid-turn injection' : 'queued message for next turn'}`, - ) - // Create user message for UI const userMessage: Message = { id: generateMessageId(), @@ -8116,6 +8142,21 @@ export class SessionManager implements ISessionManager { textElements: options?.textElements, isQueued: true, } + const hasStoredAttachmentsWithoutLivePayload = + (storedAttachments?.length ?? 0) > 0 && !attachments?.length + const canInjectMidTurn = + !hasStoredAttachmentsWithoutLivePayload && + canOfferMidTurnAttachments(attachments) && + (agent?.enqueueMidTurnMessage?.(message, attachments, { + messageId: userMessage.id, + optimisticMessageId: options?.optimisticMessageId, + }) ?? + false) + + sessionLog.info( + `Session ${sessionId} ${canInjectMidTurn ? 'queued message for mid-turn injection' : 'queued message for next turn'}`, + ) + managed.messages.push(userMessage) // Always show the message as queued while the current turn is still diff --git a/packages/desktop/packages/server-core/src/sessions/qwen-native-history.test.ts b/packages/desktop/packages/server-core/src/sessions/qwen-native-history.test.ts index 890dc9a8894..1fa466f53a3 100644 --- a/packages/desktop/packages/server-core/src/sessions/qwen-native-history.test.ts +++ b/packages/desktop/packages/server-core/src/sessions/qwen-native-history.test.ts @@ -10,7 +10,7 @@ import { join } from 'node:path'; import { tmpdir } from 'node:os'; import type { AgentBackend } from '@craft-agent/shared/agent/backend'; import type { Workspace } from '@craft-agent/shared/config'; -import { RPC_CHANNELS } from '@craft-agent/shared/protocol'; +import { RPC_CHANNELS, type FileAttachment } from '@craft-agent/shared/protocol'; import { loadSession, saveSession, @@ -1390,6 +1390,644 @@ describe('Qwen native history loading', () => { expect(persisted?.messages[0]?.attachments).toEqual([attachment]); }); + it('offers live image attachments to Qwen mid-turn injection', async () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-live-visual'; + const timestamp = Date.now(); + const liveAttachment: FileAttachment = { + type: 'image', + path: join(workspaceRoot, 'subscription.jpg'), + name: 'subscription.jpg', + mimeType: 'image/jpeg', + base64: 'base64-image', + size: 1024, + }; + const storedAttachment: NonNullable[number] = { + id: 'attachment-1', + type: 'image', + name: 'subscription.jpg', + mimeType: 'image/jpeg', + size: 1024, + storedPath: join( + workspaceRoot, + 'sessions', + sessionId, + 'attachments', + 'subscription.jpg', + ), + thumbnailBase64: 'data:image/jpeg;base64,thumb', + }; + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + const enqueueCalls: Array<{ + message: string; + attachments?: FileAttachment[]; + metadata?: { messageId?: string; optimisticMessageId?: string }; + }> = []; + managed.agent = { + enqueueMidTurnMessage: ( + message: string, + attachments?: FileAttachment[], + metadata?: { messageId?: string; optimisticMessageId?: string }, + ) => { + enqueueCalls.push({ message, attachments, metadata }); + return true; + }, + destroy: () => {}, + dispose: () => {}, + } as unknown as AgentBackend; + const manager = new SessionManager(); + ( + manager as unknown as { sessions: Map } + ).sessions.set(sessionId, managed); + + await manager.sendMessage( + sessionId, + '这个是什么图片', + [liveAttachment], + [storedAttachment], + { optimisticMessageId: 'optimistic-1' }, + ); + + expect(enqueueCalls).toHaveLength(1); + expect(enqueueCalls[0]).toEqual({ + message: '这个是什么图片', + attachments: [liveAttachment], + metadata: { + messageId: expect.any(String), + optimisticMessageId: 'optimistic-1', + }, + }); + expect(managed.messageQueue[0]?.midTurnPending).toBe(true); + expect(managed.messageQueue[0]?.attachments).toEqual([liveAttachment]); + expect(managed.messageQueue[0]?.storedAttachments).toEqual([ + storedAttachment, + ]); + }); + + it('acknowledges Qwen mid-turn queued messages by messageId', () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-messageid-ack'; + const timestamp = Date.now(); + const storedAttachment: NonNullable[number] = { + id: 'attachment-1', + type: 'image', + name: 'subscription.jpg', + mimeType: 'image/jpeg', + size: 1024, + storedPath: join( + workspaceRoot, + 'sessions', + sessionId, + 'attachments', + 'subscription.jpg', + ), + thumbnailBase64: 'data:image/jpeg;base64,thumb', + }; + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + managed.messages.push({ + id: 'message-1', + role: 'user', + content: '这个是什么图片', + timestamp, + isQueued: true, + attachments: [storedAttachment], + }); + managed.messageQueue.push({ + message: '这个是什么图片', + storedAttachments: [storedAttachment], + messageId: 'message-1', + optimisticMessageId: 'optimistic-1', + midTurnPending: true, + }); + + const events: unknown[] = []; + const manager = new SessionManager(); + manager.setEventSink((_channel, _target, event) => { + events.push(event); + }); + const onMidTurnMessagesDrained = ( + manager as unknown as { + createMidTurnMessagesDrainedCallback: ( + managedSession: unknown, + ) => (messageIds: string[]) => void; + } + ).createMidTurnMessagesDrainedCallback(managed); + + onMidTurnMessagesDrained(['message-1']); + + expect(managed.messageQueue).toHaveLength(0); + expect(managed.messages[0]?.isQueued).toBe(false); + expect(events).toContainEqual( + expect.objectContaining({ + type: 'user_message', + sessionId, + status: 'accepted', + optimisticMessageId: 'optimistic-1', + workspaceId: workspace.id, + message: expect.objectContaining({ + id: 'message-1', + isQueued: false, + }), + }), + ); + }); + + it('does not text-match identified Qwen mid-turn queued messages', () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-text-collision'; + const timestamp = Date.now(); + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + managed.messageQueue.push( + { + message: 'duplicate response', + messageId: 'message-with-id', + midTurnPending: true, + }, + { + message: 'duplicate response', + midTurnPending: true, + }, + ); + + const manager = new SessionManager(); + const onMidTurnMessagesDrained = ( + manager as unknown as { + createMidTurnMessagesDrainedCallback: ( + managedSession: unknown, + ) => (messageIds: string[]) => void; + } + ).createMidTurnMessagesDrainedCallback(managed); + + onMidTurnMessagesDrained(['duplicate response']); + + expect(managed.messageQueue).toHaveLength(1); + expect(managed.messageQueue[0]?.messageId).toBe('message-with-id'); + }); + + it('acknowledges metadata-free Qwen mid-turn queued messages by empty text', () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-empty-text-ack'; + const timestamp = Date.now(); + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + managed.messageQueue.push( + { + message: '', + midTurnPending: true, + }, + { + message: '', + midTurnPending: true, + }, + ); + + const manager = new SessionManager(); + const onMidTurnMessagesDrained = ( + manager as unknown as { + createMidTurnMessagesDrainedCallback: ( + managedSession: unknown, + ) => (messageIds: string[]) => void; + } + ).createMidTurnMessagesDrainedCallback(managed); + + onMidTurnMessagesDrained(['', '']); + + expect(managed.messageQueue).toHaveLength(0); + }); + + it('warns when Qwen mid-turn drain acknowledgements do not match', () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-unmatched-ack'; + const timestamp = Date.now(); + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + managed.messageQueue.push({ + message: 'queued response', + messageId: 'message-with-id', + midTurnPending: true, + }); + + const warnings: unknown[][] = []; + const originalWarn = logger.warn; + logger.warn = (...args: unknown[]) => { + warnings.push(args); + }; + try { + const manager = new SessionManager(); + const onMidTurnMessagesDrained = ( + manager as unknown as { + createMidTurnMessagesDrainedCallback: ( + managedSession: unknown, + ) => (messageIds: string[]) => void; + } + ).createMidTurnMessagesDrainedCallback(managed); + + onMidTurnMessagesDrained(['missing-message-id']); + + expect(managed.messageQueue).toHaveLength(1); + expect(warnings).toContainEqual([ + '[session]', + `Mid-turn drain acknowledgement matched 0/1 entries for session ${sessionId}`, + ]); + } finally { + logger.warn = originalWarn; + } + }); + + it('offers plain text follow-ups to Qwen mid-turn injection after visual messages', async () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-text-after-visual'; + const timestamp = Date.now(); + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + const enqueueCalls: Array<{ + message: string; + attachments?: FileAttachment[]; + metadata?: { messageId?: string; optimisticMessageId?: string }; + }> = []; + managed.agent = { + enqueueMidTurnMessage: ( + message: string, + attachments?: FileAttachment[], + metadata?: { messageId?: string; optimisticMessageId?: string }, + ) => { + enqueueCalls.push({ message, attachments, metadata }); + return true; + }, + destroy: () => {}, + dispose: () => {}, + } as unknown as AgentBackend; + const manager = new SessionManager(); + ( + manager as unknown as { sessions: Map } + ).sessions.set(sessionId, managed); + + await manager.sendMessage( + sessionId, + 'also summarize the visible text', + undefined, + undefined, + { optimisticMessageId: 'optimistic-text' }, + ); + + expect(enqueueCalls).toHaveLength(1); + expect(enqueueCalls[0]).toEqual({ + message: 'also summarize the visible text', + attachments: undefined, + metadata: { + messageId: expect.any(String), + optimisticMessageId: 'optimistic-text', + }, + }); + expect(managed.messageQueue[0]?.midTurnPending).toBe(true); + expect(managed.messageQueue[0]?.storedAttachments).toBeUndefined(); + }); + + it('keeps stored-only image attachments queued for the next turn', async () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-stored-only-visual'; + const timestamp = Date.now(); + const storedAttachment: NonNullable[number] = { + id: 'attachment-1', + type: 'image', + name: 'subscription.jpg', + mimeType: 'image/jpeg', + size: 1024, + storedPath: join( + workspaceRoot, + 'sessions', + sessionId, + 'attachments', + 'subscription.jpg', + ), + thumbnailBase64: 'data:image/jpeg;base64,thumb', + }; + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + const enqueueCalls: string[] = []; + managed.agent = { + enqueueMidTurnMessage: (message: string) => { + enqueueCalls.push(message); + return true; + }, + destroy: () => {}, + dispose: () => {}, + } as unknown as AgentBackend; + const manager = new SessionManager(); + ( + manager as unknown as { sessions: Map } + ).sessions.set(sessionId, managed); + + await manager.sendMessage( + sessionId, + '这个是什么图片', + undefined, + [storedAttachment], + { optimisticMessageId: 'optimistic-1' }, + ); + + expect(enqueueCalls).toEqual([]); + expect(managed.messageQueue[0]?.midTurnPending).toBe(false); + expect(managed.messageQueue[0]?.attachments).toBeUndefined(); + expect(managed.messageQueue[0]?.storedAttachments).toEqual([ + storedAttachment, + ]); + }); + + it('keeps image attachments queued when live base64 is unavailable', async () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-stored-visual'; + const timestamp = Date.now(); + const liveAttachment: FileAttachment = { + type: 'image', + path: join(workspaceRoot, 'subscription.jpg'), + name: 'subscription.jpg', + mimeType: 'image/jpeg', + size: 1024, + }; + const storedAttachment: NonNullable[number] = { + id: 'attachment-1', + type: 'image', + name: 'subscription.jpg', + mimeType: 'image/jpeg', + size: 1024, + storedPath: join( + workspaceRoot, + 'sessions', + sessionId, + 'attachments', + 'subscription.jpg', + ), + thumbnailBase64: 'data:image/jpeg;base64,thumb', + }; + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + const enqueueCalls: string[] = []; + managed.agent = { + enqueueMidTurnMessage: (message: string) => { + enqueueCalls.push(message); + return true; + }, + destroy: () => {}, + dispose: () => {}, + } as unknown as AgentBackend; + const manager = new SessionManager(); + ( + manager as unknown as { sessions: Map } + ).sessions.set(sessionId, managed); + + await manager.sendMessage( + sessionId, + '这个是什么图片', + [liveAttachment], + [storedAttachment], + { optimisticMessageId: 'optimistic-1' }, + ); + + expect(enqueueCalls).toEqual([]); + expect(managed.messageQueue[0]?.midTurnPending).toBe(false); + expect(managed.messageQueue[0]?.attachments).toEqual([liveAttachment]); + expect(managed.messageQueue[0]?.storedAttachments).toEqual([ + storedAttachment, + ]); + }); + + it('keeps non-image attachments queued for the next turn', async () => { + const workspaceRoot = mkdtempSync( + join(tmpdir(), 'craft-managed-workspace-'), + ); + tempRoots.push(workspaceRoot); + + const sessionId = '260602-qwen-midturn-pdf'; + const timestamp = Date.now(); + const liveAttachment: FileAttachment = { + type: 'pdf', + path: join(workspaceRoot, 'report.pdf'), + name: 'report.pdf', + mimeType: 'application/pdf', + base64: 'base64-pdf', + size: 1024, + }; + const workspace: Workspace = { + id: 'workspace-qwen', + name: 'qwen-code', + slug: 'qwen-code', + rootPath: workspaceRoot, + createdAt: timestamp, + }; + const managed = createManagedSession( + { + id: sessionId, + sdkSessionId: sessionId, + sdkCwd: workspaceRoot, + workingDirectory: workspaceRoot, + name: 'existing qwen title', + llmConnection: 'qwen-code', + lastMessageAt: timestamp, + }, + workspace, + { isProcessing: true, messagesLoaded: true }, + ); + const enqueueCalls: string[] = []; + managed.agent = { + enqueueMidTurnMessage: (message: string) => { + enqueueCalls.push(message); + return true; + }, + destroy: () => {}, + dispose: () => {}, + } as unknown as AgentBackend; + const manager = new SessionManager(); + ( + manager as unknown as { sessions: Map } + ).sessions.set(sessionId, managed); + + await manager.sendMessage(sessionId, 'read this pdf', [liveAttachment]); + + expect(enqueueCalls).toEqual([]); + expect(managed.messageQueue[0]?.midTurnPending).toBe(false); + expect(managed.messageQueue[0]?.attachments).toEqual([liveAttachment]); + }); + it('merges local visual overlays into provider-loaded Qwen messages', async () => { const workspaceRoot = mkdtempSync( join(tmpdir(), 'craft-managed-workspace-'), diff --git a/packages/desktop/packages/shared/src/agent/__tests__/qwen-agent-slash-history.test.ts b/packages/desktop/packages/shared/src/agent/__tests__/qwen-agent-slash-history.test.ts index c077e87736a..96c3293d5d0 100644 --- a/packages/desktop/packages/shared/src/agent/__tests__/qwen-agent-slash-history.test.ts +++ b/packages/desktop/packages/shared/src/agent/__tests__/qwen-agent-slash-history.test.ts @@ -10,6 +10,7 @@ import { join, resolve } from 'node:path'; import { tmpdir } from 'node:os'; import type { AgentEvent, Message } from '@craft-agent/core/types'; import { QwenAgent } from '../qwen-agent.ts'; +import type { FileAttachment } from '../../utils/files.ts'; type QwenAgentConfig = ConstructorParameters[0]; @@ -39,6 +40,8 @@ type QwenHistoryInternals = { type QwenPromptBlock = { type: string; text?: string; + data?: string; + mimeType?: string; resource?: { uri?: string; mimeType?: string | null; @@ -48,7 +51,15 @@ type QwenPromptBlock = { }; type QwenPromptInternals = { - buildPromptBlocks: (message: string) => QwenPromptBlock[]; + buildPromptBlocks: ( + message: string, + attachments?: FileAttachment[], + options?: { includeContext?: boolean }, + ) => QwenPromptBlock[]; +}; + +type QwenDebugInternals = { + onDebug?: (message: string) => void; }; type QwenAvailableCommandsInternals = { @@ -211,6 +222,36 @@ describe('QwenAgent slash command history', () => { expect(blocks).toEqual([{ type: 'text', text: 'hello' }]); }); + it('logs attachments skipped while building prompt blocks', () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const agent = createAgent(cwd); + const debugMessages: string[] = []; + (agent as unknown as QwenDebugInternals).onDebug = (message) => { + debugMessages.push(message); + }; + + const attachment: FileAttachment = { + type: 'unknown', + path: '', + name: 'empty.bin', + mimeType: 'application/octet-stream', + size: 0, + }; + const blocks = (agent as unknown as QwenPromptInternals).buildPromptBlocks( + 'hello', + [attachment], + ); + + expect(blocks).toEqual([{ type: 'text', text: 'hello' }]); + expect(debugMessages).toContain( + '[QwenAgent] Skipping attachment empty.bin while building prompt blocks: no readable content', + ); + + agent.destroy(); + }); + it('drains queued mid-turn messages through the ACP extension handler', async () => { const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); tempRoots.push(cwd); @@ -221,7 +262,11 @@ describe('QwenAgent slash command history', () => { internals.qwenSessionId = 'sdk-session-qwen'; internals._isProcessing = true; - expect(agent.enqueueMidTurnMessage('please also inspect tests')).toBe(true); + expect( + agent.enqueueMidTurnMessage('please also inspect tests', undefined, { + messageId: 'queued-1', + }), + ).toBe(true); await expect( internals.handleExtMethod('craft/drainMidTurnQueue', { @@ -235,26 +280,367 @@ describe('QwenAgent slash command history', () => { ).resolves.toEqual({ messages: ['please also inspect tests'], }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith(['queued-1']); + + expect( + agent.enqueueMidTurnMessage('and summarize findings', undefined, { + messageId: 'queued-2', + }), + ).toBe(true); + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + messages: ['and summarize findings'], + }); + expect(onMidTurnMessagesDrained).toHaveBeenLastCalledWith(['queued-2']); + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ messages: [] }); + + agent.destroy(); + }); + + it('acknowledges drained mid-turn messages without metadata by text', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock(() => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + expect(agent.enqueueMidTurnMessage('legacy queued message')).toBe(true); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + messages: ['legacy queued message'], + }); expect(onMidTurnMessagesDrained).toHaveBeenCalledWith([ - 'please also inspect tests', + 'legacy queued message', ]); - expect(agent.enqueueMidTurnMessage('and summarize findings')).toBe(true); + agent.destroy(); + }); + + it('acknowledges metadata-free image-only mid-turn messages by empty text', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock((_messageIds: string[]) => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + const attachment: FileAttachment = { + type: 'image', + path: join(cwd, 'screenshot.png'), + name: 'screenshot.png', + mimeType: 'image/png', + base64: 'iVBORw0KGgo=', + size: 8, + }; + expect(agent.enqueueMidTurnMessage('', [attachment])).toBe(true); + expect(agent.enqueueMidTurnMessage('', [attachment])).toBe(true); + await expect( internals.handleExtMethod('craft/drainMidTurnQueue', { sessionId: 'sdk-session-qwen', }), ).resolves.toEqual({ - messages: ['and summarize findings'], + items: [ + { + content: [ + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: '[User message with attachments]', + }, + { + content: [ + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: '[User message with attachments]', + }, + ], }); - expect(onMidTurnMessagesDrained).toHaveBeenLastCalledWith([ - 'and summarize findings', + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith(['', '']); + + agent.destroy(); + }); + + it('rejects empty mid-turn messages without attachments', () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const agent = createAgent(cwd); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals._isProcessing = true; + + expect(agent.enqueueMidTurnMessage('')).toBe(false); + expect(agent.enqueueMidTurnMessage(' ')).toBe(false); + + agent.destroy(); + }); + + it('drains queued mid-turn image attachments as ACP content blocks', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock(() => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + const attachment: FileAttachment = { + type: 'image', + path: join(cwd, 'screenshot.png'), + name: 'screenshot.png', + mimeType: 'image/png', + base64: 'iVBORw0KGgo=', + size: 8, + }; + expect( + agent.enqueueMidTurnMessage('please inspect this image', [attachment], { + messageId: 'queued-image', + }), + ).toBe(true); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + items: [ + { + content: [ + { type: 'text', text: 'please inspect this image' }, + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: 'please inspect this image', + }, + ], + }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith(['queued-image']); + + agent.destroy(); + }); + + it('retries and falls back when mid-turn attachment messages fail to build', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock(() => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + const promptInternals = agent as unknown as QwenPromptInternals; + const originalBuildPromptBlocks = + promptInternals.buildPromptBlocks.bind(agent); + promptInternals.buildPromptBlocks = (message, attachments, options) => { + if (message === 'bad image') { + throw new Error('image decode failed'); + } + return originalBuildPromptBlocks(message, attachments, options); + }; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + const attachment: FileAttachment = { + type: 'image', + path: join(cwd, 'screenshot.png'), + name: 'screenshot.png', + mimeType: 'image/png', + base64: 'iVBORw0KGgo=', + size: 8, + }; + expect( + agent.enqueueMidTurnMessage('bad image', [attachment], { + messageId: 'bad-image', + }), + ).toBe(true); + expect( + agent.enqueueMidTurnMessage('good image', [attachment], { + messageId: 'good-image', + }), + ).toBe(true); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + items: [ + { + content: [ + { type: 'text', text: 'good image' }, + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: 'good image', + }, + ], + }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith(['good-image']); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ items: [] }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledTimes(1); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + items: [ + { + content: [ + { type: 'text', text: 'bad image' }, + { + type: 'text', + text: '[Attachment could not be processed]', + }, + ], + displayText: 'bad image', + }, + ], + }); + expect(onMidTurnMessagesDrained).toHaveBeenLastCalledWith(['bad-image']); + expect(onMidTurnMessagesDrained).toHaveBeenCalledTimes(2); + + agent.destroy(); + }); + + it('acknowledges image-only mid-turn messages by optimistic id', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock(() => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + const attachment: FileAttachment = { + type: 'image', + path: join(cwd, 'screenshot.png'), + name: 'screenshot.png', + mimeType: 'image/png', + base64: 'iVBORw0KGgo=', + size: 8, + }; + expect( + agent.enqueueMidTurnMessage('', [attachment], { + optimisticMessageId: 'optimistic-image', + }), + ).toBe(true); + + await expect( + internals.handleExtMethod('craft/drainMidTurnQueue', { + sessionId: 'sdk-session-qwen', + }), + ).resolves.toEqual({ + items: [ + { + content: [ + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: '[User message with attachments]', + }, + ], + }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith([ + 'optimistic-image', ]); + + agent.destroy(); + }); + + it('drains mixed text and image mid-turn messages as ACP items', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'qwen-cwd-')); + tempRoots.push(cwd); + + const onMidTurnMessagesDrained = mock(() => {}); + const agent = createAgent(cwd, undefined, onMidTurnMessagesDrained); + const internals = agent as unknown as QwenAvailableCommandsInternals; + internals.qwenSessionId = 'sdk-session-qwen'; + internals._isProcessing = true; + + const attachment: FileAttachment = { + type: 'image', + path: join(cwd, 'screenshot.png'), + name: 'screenshot.png', + mimeType: 'image/png', + base64: 'iVBORw0KGgo=', + size: 8, + }; + expect( + agent.enqueueMidTurnMessage('first text only', undefined, { + messageId: 'queued-text', + }), + ).toBe(true); + expect( + agent.enqueueMidTurnMessage('then inspect image', [attachment], { + messageId: 'queued-image', + }), + ).toBe(true); + await expect( internals.handleExtMethod('craft/drainMidTurnQueue', { sessionId: 'sdk-session-qwen', }), - ).resolves.toEqual({ messages: [] }); + ).resolves.toEqual({ + items: [ + { + content: [{ type: 'text', text: 'first text only' }], + displayText: 'first text only', + }, + { + content: [ + { type: 'text', text: 'then inspect image' }, + { + type: 'image', + data: 'iVBORw0KGgo=', + mimeType: 'image/png', + }, + ], + displayText: 'then inspect image', + }, + ], + }); + expect(onMidTurnMessagesDrained).toHaveBeenCalledWith([ + 'queued-text', + 'queued-image', + ]); agent.destroy(); }); diff --git a/packages/desktop/packages/shared/src/agent/backend/types.ts b/packages/desktop/packages/shared/src/agent/backend/types.ts index a74479ef17f..ba224ccb1c1 100644 --- a/packages/desktop/packages/shared/src/agent/backend/types.ts +++ b/packages/desktop/packages/shared/src/agent/backend/types.ts @@ -58,6 +58,11 @@ import type { QwenSkillSetEnabledResult, } from '../../protocol/dto.ts'; +export interface MidTurnMessageMetadata { + messageId?: string; + optimisticMessageId?: string; +} + /** * Provider identifier for AI backends. * @deprecated Use ModelProvider from config/models.ts instead @@ -269,7 +274,7 @@ export interface CoreBackendConfig { * Callback when ACP has consumed queued mid-turn user messages. * Hosts can use this as an acknowledgement to remove replay fallbacks. */ - onMidTurnMessagesDrained?: (messages: string[]) => void; + onMidTurnMessagesDrained?: (messageIds: string[]) => void; /** Callback when a backend reports its live model list for the session. */ onAvailableModelsUpdate?: ( @@ -473,7 +478,11 @@ export interface AgentBackend { * the message was drained, so backends can accept candidates before knowing * whether this turn will actually produce a tool boundary. */ - enqueueMidTurnMessage?(message: string): boolean; + enqueueMidTurnMessage?( + message: string, + attachments?: FileAttachment[], + metadata?: MidTurnMessageMetadata, + ): boolean; /** * Run a simple text completion using the backend's auth infrastructure. diff --git a/packages/desktop/packages/shared/src/agent/qwen-agent.ts b/packages/desktop/packages/shared/src/agent/qwen-agent.ts index f33862d9439..af480061bb5 100644 --- a/packages/desktop/packages/shared/src/agent/qwen-agent.ts +++ b/packages/desktop/packages/shared/src/agent/qwen-agent.ts @@ -56,6 +56,7 @@ import type { BackendRewindResult, ChatOptions, BackendHostRuntimeContext, + MidTurnMessageMetadata, PermissionRequestType, SdkMcpServerConfig, } from './backend/types.ts'; @@ -102,6 +103,9 @@ type JsonRecord = Record; const QWEN_RESPONSE_INTERRUPTED_MESSAGE = 'Response interrupted'; const QWEN_TOOL_RESULT_MISSING_MESSAGE = 'Tool result was not recorded.'; +const MAX_MID_TURN_CONTENT_BUILD_FAILURES = 3; +const MID_TURN_ATTACHMENT_PROCESSING_FAILURE_TEXT = + '[Attachment could not be processed]'; function getErrorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); @@ -1624,6 +1628,12 @@ function permissionTypeForKind( } } +interface QueuedMidTurnMessage extends MidTurnMessageMetadata { + message: string; + attachments?: FileAttachment[]; + buildFailureCount?: number; +} + export class QwenAgent extends BaseAgent { protected backendName = 'Qwen Code'; @@ -1671,7 +1681,7 @@ export class QwenAgent extends BaseAgent { private toolNames = new Map(); private toolInputs = new Map>(); private activeParentToolUseIds = new Set(); - private midTurnMessageQueue: string[] = []; + private midTurnMessageQueue: QueuedMidTurnMessage[] = []; constructor(config: BackendConfig) { super(config, config.model || ''); @@ -1937,11 +1947,26 @@ export class QwenAgent extends BaseAgent { } } - enqueueMidTurnMessage(message: string): boolean { + enqueueMidTurnMessage( + message: string, + attachments?: FileAttachment[], + metadata?: MidTurnMessageMetadata, + ): boolean { const trimmed = message.trim(); - if (!trimmed || !this._isProcessing || this.abortReason) return false; + if ( + (!trimmed && !attachments?.length) || + !this._isProcessing || + this.abortReason + ) { + return false; + } - this.midTurnMessageQueue.push(trimmed); + this.midTurnMessageQueue.push({ + message: trimmed, + attachments, + messageId: metadata?.messageId, + optimisticMessageId: metadata?.optimisticMessageId, + }); this.debug( `Queued mid-turn user message for Qwen ACP injection (${this.midTurnMessageQueue.length} pending)`, ); @@ -2788,14 +2813,76 @@ export class QwenAgent extends BaseAgent { return {}; } - const messages = this.midTurnMessageQueue.splice(0); - if (messages.length > 0) { + const entries = this.midTurnMessageQueue.splice(0); + + const hasAttachments = entries.some( + (entry) => entry.attachments && entry.attachments.length > 0, + ); + if (!hasAttachments) { + if (entries.length > 0) { + this.debug( + `Drained ${entries.length} mid-turn user message(s) to Qwen ACP`, + ); + this.config.onMidTurnMessagesDrained?.( + entries.map( + (entry) => + entry.messageId ?? entry.optimisticMessageId ?? entry.message, + ), + ); + } + return { messages: entries.map((entry) => entry.message) }; + } + + const items: Array<{ content: ContentBlock[]; displayText: string }> = []; + const messageIds: string[] = []; + const failedEntries: QueuedMidTurnMessage[] = []; + for (const entry of entries) { + const displayText = entry.message || '[User message with attachments]'; + try { + items.push({ + content: this.buildPromptBlocks(entry.message, entry.attachments, { + includeContext: false, + }), + displayText, + }); + messageIds.push( + entry.messageId ?? entry.optimisticMessageId ?? entry.message, + ); + } catch (error) { + const buildFailureCount = (entry.buildFailureCount ?? 0) + 1; + this.debug( + `Failed to build mid-turn content blocks (${buildFailureCount}/${MAX_MID_TURN_CONTENT_BUILD_FAILURES}): ${getErrorMessage(error)}`, + ); + if (buildFailureCount >= MAX_MID_TURN_CONTENT_BUILD_FAILURES) { + items.push({ + content: [ + { type: 'text', text: displayText }, + { type: 'text', text: MID_TURN_ATTACHMENT_PROCESSING_FAILURE_TEXT }, + ], + displayText, + }); + messageIds.push( + entry.messageId ?? entry.optimisticMessageId ?? entry.message, + ); + } else { + failedEntries.push({ ...entry, buildFailureCount }); + } + } + } + + if (failedEntries.length > 0) { + this.midTurnMessageQueue.unshift(...failedEntries); + } + if (messageIds.length > 0) { this.debug( - `Drained ${messages.length} mid-turn user message(s) to Qwen ACP`, + `Drained ${messageIds.length} mid-turn user message(s) to Qwen ACP`, ); - this.config.onMidTurnMessagesDrained?.(messages); + this.config.onMidTurnMessagesDrained?.(messageIds); } - return { messages }; + + return { + items, + }; } private getAcpConnection(): ClientSideConnection { @@ -3447,13 +3534,15 @@ export class QwenAgent extends BaseAgent { private buildPromptBlocks( message: string, attachments?: FileAttachment[], + options?: { includeContext?: boolean }, ): ContentBlock[] { - if (isSlashCommandPrompt(message, attachments)) { + const includeContext = options?.includeContext ?? true; + if (includeContext && isSlashCommandPrompt(message, attachments)) { return [{ type: 'text', text: message.trim() }]; } const textParts: string[] = []; - const context = INCLUDE_CRAFT_CONTEXT_IN_QWEN_PROMPTS + const context = includeContext && INCLUDE_CRAFT_CONTEXT_IN_QWEN_PROMPTS ? this.buildCraftContext() : ''; @@ -3471,17 +3560,22 @@ export class QwenAgent extends BaseAgent { textParts.push( `[Attached text: ${attachment.name}]\n${attachment.text}`, ); + } else { + this.debug( + `Skipping attachment ${attachment.name} while building prompt blocks: no readable content`, + ); } } textParts.push(message); const text = textParts.filter(Boolean).join('\n\n'); - const blocks: ContentBlock[] = [ - { + const blocks: ContentBlock[] = []; + if (text || context) { + blocks.push({ type: 'text', text: context ? `${text}\n\n` : text, - }, - ]; + }); + } if (context) { blocks.push({