@machina.ai/cell-cli-core 1.2.2-rc1 → 1.4.0-rc1

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

@@ -6,6 +6,7 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { Turn, GeminiEventType } from './turn.js';
 import { reportError } from '../utils/errorReporting.js';
+import { 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,10 +57,16 @@ 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);
@@ -80,15 +88,22 @@ 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);
@@ -122,17 +137,23 @@ 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);
@@ -172,84 +193,103 @@ describe('Turn', () => {
         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([{ 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 () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Partial response' }] },
-                            finishReason: 'STOP',
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Partial response' }] },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                        usageMetadata: {
+                            promptTokenCount: 17,
+                            candidatesTokenCount: 50,
+                            cachedContentTokenCount: 10,
+                            thoughtsTokenCount: 5,
+                            toolUsePromptTokenCount: 2,
                         },
-                    ],
+                    },
                 };
             })();
             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([{ text: 'Test finish reason' }], new AbortController().signal)) {
                 events.push(event);
             }
             expect(events).toEqual([
                 { type: GeminiEventType.Content, value: 'Partial response' },
-                { type: GeminiEventType.Finished, value: 'STOP' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: {
+                        reason: 'STOP',
+                        usageMetadata: {
+                            promptTokenCount: 17,
+                            candidatesTokenCount: 50,
+                            cachedContentTokenCount: 10,
+                            thoughtsTokenCount: 5,
+                            toolUsePromptTokenCount: 2,
+                        },
+                    },
+                },
             ]);
         });
         it('should yield finished event for MAX_TOKENS 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: 'This is a long response that was cut off...' },
+                                    ],
+                                },
+                                finishReason: 'MAX_TOKENS',
                             },
-                            finishReason: 'MAX_TOKENS',
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -263,18 +303,24 @@ describe('Turn', () => {
                     type: GeminiEventType.Content,
                     value: 'This is a long response that was cut off...',
                 },
-                { type: GeminiEventType.Finished, value: 'MAX_TOKENS' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'MAX_TOKENS', usageMetadata: undefined },
+                },
             ]);
         });
         it('should yield finished event for SAFETY finish reason', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Content blocked' }] },
-                            finishReason: 'SAFETY',
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Content blocked' }] },
+                                finishReason: 'SAFETY',
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -285,18 +331,26 @@ describe('Turn', () => {
             }
             expect(events).toEqual([
                 { type: GeminiEventType.Content, value: 'Content blocked' },
-                { type: GeminiEventType.Finished, value: 'SAFETY' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'SAFETY', usageMetadata: undefined },
+                },
             ]);
         });
-        it('should not yield finished event when there is no 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: 'Response without finish reason' }] },
-                            // No finishReason property
-                        },
-                    ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: {
+                                    parts: [{ text: 'Response without finish reason' }],
+                                },
+                                // No finishReason property
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -311,25 +365,30 @@ describe('Turn', () => {
                     value: 'Response without finish reason',
                 },
             ]);
-            // No Finished event should be emitted
         });
         it('should handle multiple responses with different finish reasons', 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: 'First part' }] },
+                                // No finish reason on first response
+                            },
+                        ],
+                    },
                 };
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Second part' }] },
-                            finishReason: 'OTHER',
-                        },
-                    ],
+                    value: {
+                        type: StreamEventType.CHUNK,
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Second part' }] },
+                                finishReason: 'OTHER',
+                            },
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -341,26 +400,32 @@ describe('Turn', () => {
             expect(events).toEqual([
                 { type: GeminiEventType.Content, value: 'First part' },
                 { type: GeminiEventType.Content, value: 'Second part' },
-                { type: GeminiEventType.Finished, value: 'OTHER' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'OTHER', usageMetadata: undefined },
+                },
             ]);
         });
         it('should yield citation and finished events when response has citationMetadata', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Some text.' }] },
-                            citationMetadata: {
-                                citations: [
-                                    {
-                                        uri: 'https://example.com/source1',
-                                        title: 'Source 1 Title',
-                                    },
-                                ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Source 1 Title',
+                                        },
+                                    ],
+                                },
+                                finishReason: 'STOP',
                             },
-                            finishReason: 'STOP',
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -374,30 +439,36 @@ describe('Turn', () => {
                     type: GeminiEventType.Citation,
                     value: 'Citations:\n(Source 1 Title) https://example.com/source1',
                 },
-                { type: GeminiEventType.Finished, value: 'STOP' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'STOP', usageMetadata: undefined },
+                },
             ]);
         });
         it('should yield a single citation event for multiple citations in one response', async () => {
             const mockResponseStream = (async function* () {
                 yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Some text.' }] },
-                            citationMetadata: {
-                                citations: [
-                                    {
-                                        uri: 'https://example.com/source2',
-                                        title: 'Title2',
-                                    },
-                                    {
-                                        uri: 'https://example.com/source1',
-                                        title: 'Title1',
-                                    },
-                                ],
+                    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',
                             },
-                            finishReason: 'STOP',
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -411,26 +482,32 @@ describe('Turn', () => {
                     type: GeminiEventType.Citation,
                     value: 'Citations:\n(Title1) https://example.com/source1\n(Title2) https://example.com/source2',
                 },
-                { type: GeminiEventType.Finished, value: 'STOP' },
+                {
+                    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: 'Some text.' }] },
-                            citationMetadata: {
-                                citations: [
-                                    {
-                                        uri: 'https://example.com/source1',
-                                        title: 'Source 1 Title',
-                                    },
-                                ],
+                    type: StreamEventType.CHUNK,
+                    value: {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Some text.' }] },
+                                citationMetadata: {
+                                    citations: [
+                                        {
+                                            uri: 'https://example.com/source1',
+                                            title: 'Source 1 Title',
+                                        },
+                                    ],
+                                },
+                                // No finishReason
                             },
-                            // No finishReason
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -441,30 +518,33 @@ describe('Turn', () => {
             expect(events).toEqual([
                 { type: GeminiEventType.Content, value: 'Some text.' },
             ]);
-            // No Citation or Finished event
+            // 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: 'Some text.' }] },
-                            citationMetadata: {
-                                citations: [
-                                    {
-                                        uri: 'https://example.com/source1',
-                                        title: 'Good Source',
-                                    },
-                                    {
-                                        // uri is undefined
-                                        title: 'Bad Source',
-                                    },
-                                ],
+                    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',
                             },
-                            finishReason: 'STOP',
-                        },
-                    ],
+                        ],
+                    },
                 };
             })();
             mockSendMessageStream.mockResolvedValue(mockResponseStream);
@@ -478,7 +558,10 @@ describe('Turn', () => {
                     type: GeminiEventType.Citation,
                     value: 'Citations:\n(Good Source) https://example.com/source1',
                 },
-                { type: GeminiEventType.Finished, value: 'STOP' },
+                {
+                    type: GeminiEventType.Finished,
+                    value: { reason: 'STOP', usageMetadata: undefined },
+                },
             ]);
         });
         it('should not crash when cancelled request has malformed error', async () => {
@@ -500,6 +583,26 @@ describe('Turn', () => {
             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([], new AbortController().signal)) {
+                events.push(event);
+            }
+            expect(events).toEqual([
+                { type: GeminiEventType.Retry },
+                { type: GeminiEventType.Content, value: 'Success' },
+            ]);
+        });
     });
     describe('getDebugResponses', () => {
         it('should return collected debug responses', async () => {
@@ -510,8 +613,8 @@ 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' }];