@office-ai/aioncli-core 0.2.3 → 0.18.4

package/dist/src/core/turn.test.js

@@ -4,8 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { Turn, GeminiEventType, } from './turn.js';
+import { Turn, GeminiEventType } from './turn.js';
 import { reportError } from '../utils/errorReporting.js';
+import { InvalidStreamError, StreamEventType } from './geminiChat.js';
 const mockSendMessageStream = vi.fn();
 const mockGetHistory = vi.fn();
 const mockMaybeIncludeSchemaDepthContext = vi.fn();
@@ -24,6 +25,7 @@ vi.mock('@google/genai', async (importOriginal) => {
 vi.mock('../utils/errorReporting', () => ({
     reportError: vi.fn(),
 }));
+// Use the actual implementation from partUtils now that it's provided.
 vi.mock('../utils/generateContentResponseUtilities', () => ({
     getResponseText: (resp) => resp.candidates?.[0]?.content?.parts?.map((part) => part.text).join('') ||
         undefined,
@@ -55,22 +57,25 @@ describe('Turn', () => {
         it('should yield content events for text parts', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [{ content: { parts: [{ text: 'Hello' }] } }],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [{ content: { parts: [{ text: 'Hello' }] } }],
+                    },
                 };
                 yield {
-                    candidates: [{ content: { parts: [{ text: ' world' }] } }],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [{ content: { parts: [{ text: ' world' }] } }],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
             const reqParts = [{ text: 'Hi' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 events.push(event);
             }
-            expect(mockSendMessageStream).toHaveBeenCalledWith({
-                message: reqParts,
-                config: { abortSignal: expect.any(AbortSignal) },
-            }, 'prompt-id-1');
+            expect(mockSendMessageStream).toHaveBeenCalledWith({ model: 'gemini' }, reqParts, 'prompt-id-1', expect.any(AbortSignal));
             expect(events).toEqual([
                 { type: GeminiEventType.Content, value: 'Hello' },
                 { type: GeminiEventType.Content, value: ' world' },
@@ -80,21 +85,28 @@ describe('Turn', () => {
         it('should yield tool_call_request events for function calls', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    functionCalls: [
-                        {
-                            id: 'fc1',
-                            name: 'tool1',
-                            args: { arg1: 'val1' },
-                            isClientInitiated: false,
-                        },
-                        { name: 'tool2', args: { arg2: 'val2' }, isClientInitiated: false }, // No ID
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        functionCalls: [
+                            {
+                                id: 'fc1',
+                                name: 'tool1',
+                                args: { arg1: 'val1' },
+                                isClientInitiated: false,
+                            },
+                            {
+                                name: 'tool2',
+                                args: { arg2: 'val2' },
+                                isClientInitiated: false,
+                            }, // No ID
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
             const reqParts = [{ text: 'Use tools' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 events.push(event);
             }
             expect(events.length).toBe(2);
@@ -122,23 +134,29 @@ describe('Turn', () => {
             const abortController = new AbortController();
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [{ content: { parts: [{ text: 'First part' }] } }],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [{ content: { parts: [{ text: 'First part' }] } }],
+                    },
                 };
                 abortController.abort();
                 yield {
-                    candidates: [
-                        {
-                            content: {
-                                parts: [{ text: 'Second part - should not be processed' }],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: {
+                                    parts: [{ text: 'Second part - should not be processed' }],
+                                },
                             },
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
             const reqParts = [{ text: 'Test abort' }];
-            for await (const event of turn.run(reqParts, abortController.signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, abortController.signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
@@ -147,6 +165,18 @@ describe('Turn', () => {
             ]);
             expect(turn.getDebugResponses().length).toBe(1);
         });
+        it('should yield InvalidStream event if sendMessageStream throws InvalidStreamError', async () => {
+            const error = new InvalidStreamError('Test invalid stream', 'NO_FINISH_REASON');
+            mockSendMessageStream.mockRejectedValue(error);
+            const reqParts = [{ text: 'Trigger invalid stream' }];
+            const events = [];
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([{ type: GeminiEventType.InvalidStream }]);
+            expect(turn.getDebugResponses().length).toBe(0);
+            expect(reportError).not.toHaveBeenCalled(); // Should not report as error
+        });
         it('should yield Error event and report if sendMessageStream throws', async () => {
             const error = new Error('API Error');
             mockSendMessageStream.mockRejectedValue(error);
@@ -157,7 +187,7 @@ describe('Turn', () => {
             mockGetHistory.mockReturnValue(historyContent);
             mockMaybeIncludeSchemaDepthContext.mockResolvedValue(undefined);
             const events = [];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 events.push(event);
             }
             expect(events.length).toBe(1);
@@ -167,183 +197,406 @@ describe('Turn', () => {
                 error: { message: 'API Error', status: undefined },
             });
             expect(turn.getDebugResponses().length).toBe(0);
-            expect(reportError).toHaveBeenCalledWith(error, 'Error when talking to Gemini API', [...historyContent, reqParts], 'Turn.run-sendMessageStream');
+            expect(reportError).toHaveBeenCalledWith(error, 'Error when talking to Gemini API', [...historyContent, { role: 'user', parts: reqParts }], 'Turn.run-sendMessageStream');
         });
         it('should handle function calls with undefined name or args', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    functionCalls: [
-                        { id: 'fc1', name: undefined, args: { arg1: 'val1' } },
-                        { id: 'fc2', name: 'tool2', args: undefined },
-                        { id: 'fc3', name: undefined, args: undefined },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [],
+                        functionCalls: [
+                            // Add `id` back to the mock to match what the code expects
+                            { id: 'fc1', name: undefined, args: { arg1: 'val1' } },
+                            { id: 'fc2', name: 'tool2', args: undefined },
+                            { id: 'fc3', name: undefined, args: undefined },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Test undefined tool parts' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'Test undefined tool parts' }], new AbortController().signal)) {
                 events.push(event);
             }
             expect(events.length).toBe(3);
+            // Assertions for each specific tool call event
             const event1 = events[0];
-            expect(event1.type).toBe(GeminiEventType.ToolCallRequest);
-            expect(event1.value).toEqual(expect.objectContaining({
+            expect(event1.value).toMatchObject({
                 callId: 'fc1',
                 name: 'undefined_tool_name',
                 args: { arg1: 'val1' },
-                isClientInitiated: false,
-            }));
-            expect(turn.pendingToolCalls[0]).toEqual(event1.value);
+            });
             const event2 = events[1];
-            expect(event2.type).toBe(GeminiEventType.ToolCallRequest);
-            expect(event2.value).toEqual(expect.objectContaining({
+            expect(event2.value).toMatchObject({
                 callId: 'fc2',
                 name: 'tool2',
                 args: {},
-                isClientInitiated: false,
-            }));
-            expect(turn.pendingToolCalls[1]).toEqual(event2.value);
+            });
             const event3 = events[2];
-            expect(event3.type).toBe(GeminiEventType.ToolCallRequest);
-            expect(event3.value).toEqual(expect.objectContaining({
+            expect(event3.value).toMatchObject({
                 callId: 'fc3',
                 name: 'undefined_tool_name',
                 args: {},
-                isClientInitiated: false,
-            }));
-            expect(turn.pendingToolCalls[2]).toEqual(event3.value);
-            expect(turn.getDebugResponses().length).toBe(1);
+            });
         });
-        it('should yield finished event when response has finish reason', async () => {
+        it.each([
+            {
+                description: 'should yield finished event when response has finish reason',
+                contentText: 'Partial response',
+                finishReason: 'STOP',
+                usageMetadata: {
+                    promptTokenCount: 17,
+                    candidatesTokenCount: 50,
+                    cachedContentTokenCount: 10,
+                    thoughtsTokenCount: 5,
+                    toolUsePromptTokenCount: 2,
+                },
+            },
+            {
+                description: 'should yield finished event for MAX_TOKENS finish reason',
+                contentText: 'This is a long response that was cut off...',
+                finishReason: 'MAX_TOKENS',
+                usageMetadata: undefined,
+            },
+            {
+                description: 'should yield finished event for SAFETY finish reason',
+                contentText: 'Content blocked',
+                finishReason: 'SAFETY',
+                usageMetadata: undefined,
+            },
+        ])('$description', async ({ contentText, finishReason, usageMetadata }) => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Partial response' }] },
-                            finishReason: 'STOP',
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: contentText }] },
+                                finishReason,
+                            },
+                        ],
+                        usageMetadata,
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Test finish reason' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'Test' }], new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
-                { type: GeminiEventType.Content, value: 'Partial response' },
-                { type: GeminiEventType.Finished, value: 'STOP' },
+                { type: GeminiEventType.Content, value: contentText },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: finishReason, usageMetadata },
+                },
             ]);
         });
-        it('should yield finished event for MAX_TOKENS finish reason', async () => {
+        it('should yield finished event with undefined reason when there is no finish reason', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: {
-                                parts: [
-                                    { text: 'This is a long response that was cut off...' },
-                                ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: {
+                                    parts: [{ text: 'Response without finish reason' }],
+                                },
+                                // No finishReason property
                             },
-                            finishReason: 'MAX_TOKENS',
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Generate long text' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            const reqParts = [{ text: 'Test no finish reason' }];
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
                 {
                     type: GeminiEventType.Content,
-                    value: 'This is a long response that was cut off...',
+                    value: 'Response without finish reason',
                 },
-                { type: GeminiEventType.Finished, value: 'MAX_TOKENS' },
             ]);
         });
-        it('should yield finished event for SAFETY finish reason', async () => {
+        it('should handle multiple responses with different finish reasons', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Content blocked' }] },
-                            finishReason: 'SAFETY',
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'First part' }] },
+                                // No finish reason on first response
+                            },
+                        ],
+                    },
+                };
+                yield {
+                    value: {
+                        type: StreamEventType.CHUNK,
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Second part' }] },
+                                finishReason: 'OTHER',
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Test safety' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            const reqParts = [{ text: 'Test multiple responses' }];
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
-                { type: GeminiEventType.Content, value: 'Content blocked' },
-                { type: GeminiEventType.Finished, value: 'SAFETY' },
+                { type: GeminiEventType.Content, value: 'First part' },
+                { type: GeminiEventType.Content, value: 'Second part' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'OTHER', usageMetadata: undefined },
+                },
             ]);
         });
-        it('should not yield finished event when there is no finish reason', async () => {
+        it('should yield citation and finished events when response has citationMetadata', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Response without finish reason' }] },
-                            // No finishReason property
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Source 1 Title',
+                                        },
+                                    ],
+                                },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Test no finish reason' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'Test citations' }], new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
+                { type: GeminiEventType.Content, value: 'Some text.' },
                 {
-                    type: GeminiEventType.Content,
-                    value: 'Response without finish reason',
+                    type: GeminiEventType.Citation,
+                    value: 'Citations:\n(Source 1 Title) https://example.com/source1',
+                },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'STOP', usageMetadata: undefined },
                 },
             ]);
-            // No Finished event should be emitted
         });
-        it('should handle multiple responses with different finish reasons', async () => {
+        it('should yield a single citation event for multiple citations in one response', async () => {
+            const mockResponseStream = (async function* () {
+                yield {
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source2',
+                                            title: 'Title2',
+                                        },
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Title1',
+                                        },
+                                    ],
+                                },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    },
+                };
+            })();
+            mockSendMessageStream.mockResolvedValue(mockResponseStream);
+            const events = [];
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'test' }], new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([
+                { type: GeminiEventType.Content, value: 'Some text.' },
+                {
+                    type: GeminiEventType.Citation,
+                    value: 'Citations:\n(Title1) https://example.com/source1\n(Title2) https://example.com/source2',
+                },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'STOP', usageMetadata: undefined },
+                },
+            ]);
+        });
+        it('should not yield citation event if there is no finish reason', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'First part' }] },
-                            // No finish reason on first response
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Source 1 Title',
+                                        },
+                                    ],
+                                },
+                                // No finishReason
+                            },
+                        ],
+                    },
                 };
+            })();
+            mockSendMessageStream.mockResolvedValue(mockResponseStream);
+            const events = [];
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'test' }], new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([
+                { type: GeminiEventType.Content, value: 'Some text.' },
+            ]);
+            // No Citation event (but we do get a Finished event with undefined reason)
+            expect(events.some((e) => e.type === GeminiEventType.Citation)).toBe(false);
+        });
+        it('should ignore citations without a URI', async () => {
+            const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Second part' }] },
-                            finishReason: 'OTHER',
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Good Source',
+                                        },
+                                        {
+                                            // uri is undefined
+                                            title: 'Bad Source',
+                                        },
+                                    ],
+                                },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const events = [];
-            const reqParts = [{ text: 'Test multiple responses' }];
-            for await (const event of turn.run(reqParts, new AbortController().signal)) {
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'test' }], new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
-                { type: GeminiEventType.Content, value: 'First part' },
-                { type: GeminiEventType.Content, value: 'Second part' },
-                { type: GeminiEventType.Finished, value: 'OTHER' },
+                { type: GeminiEventType.Content, value: 'Some text.' },
+                {
+                    type: GeminiEventType.Citation,
+                    value: 'Citations:\n(Good Source) https://example.com/source1',
+                },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'STOP', usageMetadata: undefined },
+                },
+            ]);
+        });
+        it('should not crash when cancelled request has malformed error', async () => {
+            const abortController = new AbortController();
+            const errorToThrow = {
+                response: {
+                    data: undefined, // Malformed error data
+                },
+            };
+            mockSendMessageStream.mockImplementation(async () => {
+                abortController.abort();
+                throw errorToThrow;
+            });
+            const events = [];
+            const reqParts = [{ text: 'Test malformed error handling' }];
+            for await (const event of turn.run({ model: 'gemini' }, reqParts, abortController.signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([{ type: GeminiEventType.UserCancelled }]);
+            expect(reportError).not.toHaveBeenCalled();
+        });
+        it('should yield a Retry event when it receives one from the chat stream', async () => {
+            const mockResponseStream = (async function* () {
+                yield { type: StreamEventType.RETRY };
+                yield {
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [{ content: { parts: [{ text: 'Success' }] } }],
+                    },
+                };
+            })();
+            mockSendMessageStream.mockResolvedValue(mockResponseStream);
+            const events = [];
+            for await (const event of turn.run({ model: 'gemini' }, [], new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([
+                { type: GeminiEventType.Retry },
+                { type: GeminiEventType.Content, value: 'Success' },
             ]);
         });
+        it.each([
+            {
+                description: 'should yield content events with traceId',
+                part: { text: 'Hello' },
+                responseId: 'trace-123',
+                expectedEvent: {
+                    type: GeminiEventType.Content,
+                    value: 'Hello',
+                    traceId: 'trace-123',
+                },
+            },
+            {
+                description: 'should yield thought events with traceId',
+                part: { text: '[Thought: thinking]', thought: 'thinking' },
+                responseId: 'trace-456',
+                expectedEvent: {
+                    type: GeminiEventType.Thought,
+                    value: { subject: '', description: '[Thought: thinking]' },
+                    traceId: 'trace-456',
+                },
+            },
+        ])('$description', async ({ part, responseId, expectedEvent }) => {
+            const mockResponseStream = (async function* () {
+                yield {
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [{ content: { parts: [part] } }],
+                        responseId,
+                    },
+                };
+            })();
+            mockSendMessageStream.mockResolvedValue(mockResponseStream);
+            const events = [];
+            for await (const event of turn.run({ model: 'gemini' }, [{ text: 'Hi' }], new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([expectedEvent]);
+        });
     });
     describe('getDebugResponses', () => {
         it('should return collected debug responses', async () => {
@@ -354,12 +607,12 @@ describe('Turn', () => {
                 functionCalls: [{ name: 'debugTool' }],
             };
             const mockResponseStream = (async function* () {
-                yield resp1;
-                yield resp2;
+                yield { type: StreamEventType.CHUNK, value: resp1 };
+                yield { type: StreamEventType.CHUNK, value: resp2 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
             const reqParts = [{ text: 'Hi' }];
-            for await (const _ of turn.run(reqParts, new AbortController().signal)) {
+            for await (const _ of turn.run({ model: 'gemini' }, reqParts, new AbortController().signal)) {
                 // consume stream
             }
             expect(turn.getDebugResponses()).toEqual([resp1, resp2]);