basant307/AI_Governance_Project
045
1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';8import type {9 Content,10 GenerateContentConfig,11 GenerateContentResponse,12 Part,13} from '@google/genai';14import { ApiError } from '@google/genai';15import { AuthType, type ContentGenerator } from '../core/contentGenerator.js';16import {17 GeminiChat,18 InvalidStreamError,19 redactStructuredOutputArgsForRecording,20 StreamEventType,21 type StreamEvent,22} from './geminiChat.js';23import { RETRYABLE_STREAM_TRANSPORT_CODES } from './stream-transport-retry.js';24import { classifyRetryError } from '../utils/retryErrorClassification.js';25import { StreamContentError } from './openaiContentGenerator/pipeline.js';26import type { Config } from '../config/config.js';27import { setSimulate429 } from '../utils/testUtils.js';28import { uiTelemetryService } from '../telemetry/uiTelemetry.js';29import { CompressionStatus, type ChatCompressionInfo } from './turn.js';30import {31 ChatCompressionService,32 MAX_CONSECUTIVE_FAILURES,33} from '../services/chatCompressionService.js';34import {35 estimateContentTokens,36 estimatePromptTokens,37} from '../services/tokenEstimation.js';38import { SYSTEM_REMINDER_OPEN } from '../utils/environmentContext.js';39import { SessionStartSource } from '../hooks/types.js';40 41// Mock fs module to prevent actual file system operations during tests42const mockFileSystem = new Map<string, string>();43 44vi.mock('node:fs', () => {45 const fsModule = {46 mkdirSync: vi.fn(),47 writeFileSync: vi.fn((path: string, data: string) => {48 mockFileSystem.set(path, data);49 }),50 readFileSync: vi.fn((path: string) => {51 if (mockFileSystem.has(path)) {52 return mockFileSystem.get(path);53 }54 throw Object.assign(new Error('ENOENT: no such file or directory'), {55 code: 'ENOENT',56 });57 }),58 existsSync: vi.fn((path: string) => mockFileSystem.has(path)),59 appendFileSync: vi.fn(),60 };61 62 return {63 default: fsModule,64 ...fsModule,65 };66});67 68// Add mock for the retry utility69const { mockRetryWithBackoff } = vi.hoisted(() => ({70 mockRetryWithBackoff: vi.fn(),71}));72 73vi.mock('../utils/retry.js', async (importOriginal) => {74 const actual = await importOriginal<typeof import('../utils/retry.js')>();75 return {76 ...actual,77 retryWithBackoff: mockRetryWithBackoff,78 };79});80 81const { mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({82 mockLogContentRetry: vi.fn(),83 mockLogContentRetryFailure: vi.fn(),84}));85 86vi.mock('../telemetry/loggers.js', () => ({87 logContentRetry: mockLogContentRetry,88 logContentRetryFailure: mockLogContentRetryFailure,89 // Real ChatCompressionService.compress() calls logChatCompression on90 // every attempt; the R3.4 integration test exercises that path, so the91 // mock has to expose it (no-op).92 logChatCompression: vi.fn(),93}));94 95vi.mock('../telemetry/uiTelemetry.js', () => ({96 uiTelemetryService: {97 setLastPromptTokenCount: vi.fn(),98 },99}));100 101const { mockAcquireSleepInhibitor, mockSleepInhibitorRelease } = vi.hoisted(102 () => ({103 mockAcquireSleepInhibitor: vi.fn(),104 mockSleepInhibitorRelease: vi.fn(),105 }),106);107 108vi.mock('../services/sleepInhibitor.js', () => ({109 acquireSleepInhibitor: mockAcquireSleepInhibitor,110}));111 112const { mockDebugLoggerWarn } = vi.hoisted(() => ({113 mockDebugLoggerWarn: vi.fn(),114}));115 116vi.mock('../utils/debugLogger.js', async (importOriginal) => {117 const actual =118 await importOriginal<typeof import('../utils/debugLogger.js')>();119 return {120 ...actual,121 createDebugLogger: () => ({122 debug: vi.fn(),123 info: vi.fn(),124 warn: mockDebugLoggerWarn,125 error: vi.fn(),126 }),127 };128});129 130describe('GeminiChat', async () => {131 let mockContentGenerator: ContentGenerator;132 let chat: GeminiChat;133 let mockConfig: Config;134 const config: GenerateContentConfig = {};135 136 beforeEach(() => {137 vi.clearAllMocks();138 mockAcquireSleepInhibitor.mockReturnValue({139 release: mockSleepInhibitorRelease,140 });141 vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();142 mockContentGenerator = {143 generateContent: vi.fn(),144 generateContentStream: vi.fn(),145 countTokens: vi.fn(),146 embedContent: vi.fn(),147 batchEmbedContents: vi.fn(),148 useSummarizedThinking: vi.fn().mockReturnValue(false),149 } as unknown as ContentGenerator;150 151 // Default mock implementation for tests that don't care about retry logic152 mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());153 mockConfig = {154 getSessionId: () => 'test-session-id',155 getTelemetryLogPromptsEnabled: () => true,156 getUsageStatisticsEnabled: () => true,157 getDebugMode: () => false,158 getContentGeneratorConfig: vi.fn().mockReturnValue({159 authType: 'gemini', // Ensure this is set for fallback tests160 model: 'test-model',161 }),162 getModel: vi.fn().mockReturnValue('gemini-pro'),163 setModel: vi.fn(),164 getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),165 getTargetDir: vi.fn().mockReturnValue('/test/project/root'),166 getCliVersion: vi.fn().mockReturnValue('1.0.0'),167 storage: {168 getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),169 },170 getToolRegistry: vi.fn().mockReturnValue({171 getTool: vi.fn(),172 }),173 getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),174 getBaseLlmClient: vi.fn().mockReturnValue(undefined),175 getModelFallbacks: vi.fn().mockReturnValue([]),176 getChatCompression: vi.fn().mockReturnValue(undefined),177 getAutoCompactThreshold: vi.fn().mockReturnValue(undefined),178 getHookSystem: vi.fn().mockReturnValue(undefined),179 getDebugLogger: vi180 .fn()181 .mockReturnValue({ debug: vi.fn(), warn: vi.fn(), info: vi.fn() }),182 getApprovalMode: vi.fn().mockReturnValue('default'),183 getFileReadCache: vi.fn().mockReturnValue({ clear: vi.fn() }),184 } as unknown as Config;185 186 // Disable 429 simulation for tests187 setSimulate429(false);188 // Reset history for each test by creating a new instance189 chat = new GeminiChat(190 mockConfig,191 config,192 [],193 undefined,194 uiTelemetryService,195 );196 });197 198 afterEach(() => {199 vi.restoreAllMocks();200 vi.resetAllMocks();201 });202 203 /**204 * Helper: consume a stream and expect it to throw InvalidStreamError205 * after all transient retries exhaust. Uses fake timers to skip delays.206 * Must be called within a vi.useFakeTimers() / vi.useRealTimers() block.207 */208 async function expectStreamExhaustion(209 stream: AsyncGenerator<StreamEvent>,210 ): Promise<void> {211 const collecting = (async () => {212 for await (const _ of stream) {213 /* consume */214 }215 })();216 // Get assertion promise first (don't await), then advance timers.217 const resultPromise = (async () => {218 await expect(collecting).rejects.toThrow(InvalidStreamError);219 })();220 await vi.advanceTimersByTimeAsync(0);221 await vi.advanceTimersByTimeAsync(35_000);222 await resultPromise;223 }224 225 async function collectStreamWithFakeTimers(226 stream: AsyncGenerator<StreamEvent>,227 advanceByMs: number = 10_000,228 ): Promise<StreamEvent[]> {229 const events: StreamEvent[] = [];230 const collecting = (async () => {231 for await (const event of stream) {232 events.push(event);233 }234 return events;235 })();236 await vi.advanceTimersByTimeAsync(0);237 await vi.advanceTimersByTimeAsync(advanceByMs);238 return collecting;239 }240 241 describe('system instruction helpers', () => {242 it('replaces prior session-start context instead of appending indefinitely', () => {243 const isolatedChat = new GeminiChat(244 mockConfig,245 {},246 [],247 undefined,248 uiTelemetryService,249 );250 isolatedChat.setSystemInstruction('Base instruction');251 252 isolatedChat.setSessionStartContext('Ctx1');253 isolatedChat.setSessionStartContext('Ctx2');254 255 expect(isolatedChat['generationConfig'].systemInstruction).toBe(256 'Base instruction\n\n<qwen:session-start-context hidden="true">\nSessionStart additional context:\nCtx2\n</qwen:session-start-context>',257 );258 });259 260 it('preserves existing system prompt suffixes when replacing session-start context', () => {261 const isolatedChat = new GeminiChat(262 mockConfig,263 {},264 [],265 undefined,266 uiTelemetryService,267 );268 isolatedChat.setSystemInstruction(269 'Base instruction\n\n---\n\nUser memory\n\n---\n\nAppended rule',270 );271 272 isolatedChat.setSessionStartContext('Ctx1');273 isolatedChat.setSessionStartContext('Ctx2');274 275 expect(isolatedChat['generationConfig'].systemInstruction).toBe(276 'Base instruction\n\n---\n\nUser memory\n\n---\n\nAppended rule\n\n<qwen:session-start-context hidden="true">\nSessionStart additional context:\nCtx2\n</qwen:session-start-context>',277 );278 });279 280 it('preserves non-string systemInstruction content when applying session-start context', () => {281 const isolatedChat = new GeminiChat(282 mockConfig,283 {284 systemInstruction: {285 role: 'system',286 parts: [{ text: 'Base content instruction' }],287 },288 },289 [],290 undefined,291 uiTelemetryService,292 );293 294 isolatedChat.setSessionStartContext('Ctx1');295 isolatedChat.setSessionStartContext('Ctx2');296 297 expect(isolatedChat['generationConfig'].systemInstruction).toBe(298 'Base content instruction\n\n<qwen:session-start-context hidden="true">\nSessionStart additional context:\nCtx2\n</qwen:session-start-context>',299 );300 });301 302 it('applies session-start context synchronously via applySessionStartContext', () => {303 const isolatedChat = new GeminiChat(304 mockConfig,305 {},306 [],307 undefined,308 uiTelemetryService,309 );310 isolatedChat.setSystemInstruction('Base instruction');311 312 isolatedChat.applySessionStartContext(313 ' Sync ctx ',314 SessionStartSource.Startup,315 );316 317 expect(isolatedChat['generationConfig'].systemInstruction).toBe(318 'Base instruction\n\n<qwen:session-start-context hidden="true">\nSessionStart additional context:\nSync ctx\n</qwen:session-start-context>',319 );320 });321 322 it('does not strip legitimate content that only resembles the old plain-text marker', () => {323 const isolatedChat = new GeminiChat(324 mockConfig,325 {},326 [],327 undefined,328 uiTelemetryService,329 );330 isolatedChat.setSystemInstruction(331 'Base instruction\n\n---\n\nSessionStart additional context:\nLegitimate content',332 );333 334 isolatedChat.setSessionStartContext('Ctx1');335 336 expect(isolatedChat['generationConfig'].systemInstruction).toContain(337 'Legitimate content',338 );339 expect(isolatedChat['generationConfig'].systemInstruction).toContain(340 '<qwen:session-start-context hidden="true">\nSessionStart additional context:\nCtx1\n</qwen:session-start-context>',341 );342 });343 });344 345 describe('sendMessageStream', () => {346 it('releases the sleep inhibitor after the stream is consumed', async () => {347 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(348 (async function* () {349 yield {350 candidates: [351 {352 content: { role: 'model', parts: [{ text: 'done' }] },353 finishReason: 'STOP',354 },355 ],356 } as unknown as GenerateContentResponse;357 })(),358 );359 360 const stream = await chat.sendMessageStream(361 'test-model',362 { message: 'test message' },363 'prompt-id-sleep-inhibitor',364 );365 for await (const _ of stream) {366 /* consume stream */367 }368 369 expect(mockAcquireSleepInhibitor).toHaveBeenCalledWith(370 mockConfig,371 'Qwen Code is streaming a model response',372 );373 expect(mockSleepInhibitorRelease).toHaveBeenCalledTimes(1);374 });375 376 it('increments the user-content push counter once per surviving send', async () => {377 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(378 (async function* () {379 yield {380 candidates: [381 {382 content: { role: 'model', parts: [{ text: 'done' }] },383 finishReason: 'STOP',384 },385 ],386 } as unknown as GenerateContentResponse;387 })(),388 );389 390 const before = chat.getUserContentPushCount();391 const stream = await chat.sendMessageStream(392 'test-model',393 { message: 'hello' },394 'prompt-id-push-count',395 );396 for await (const _ of stream) {397 /* consume stream */398 }399 400 // The user content landed exactly once, so the counter advanced by one —401 // this is the signal the Retry strip/restore in client.ts gates on.402 expect(chat.getUserContentPushCount()).toBe(before + 1);403 });404 405 it('releases the sleep inhibitor when the stream errors', async () => {406 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(407 (async function* () {408 yield {409 candidates: [410 {411 content: { role: 'model', parts: [{ text: 'partial' }] },412 },413 ],414 } as unknown as GenerateContentResponse;415 throw new Error('stream aborted');416 })(),417 );418 419 const stream = await chat.sendMessageStream(420 'test-model',421 { message: 'fail' },422 'prompt-id-stream-error',423 );424 425 await expect(426 (async () => {427 for await (const _ of stream) {428 /* consume stream */429 }430 })(),431 ).rejects.toThrow('stream aborted');432 433 expect(mockSleepInhibitorRelease).toHaveBeenCalledTimes(1);434 });435 436 it('should succeed if a tool call is followed by an empty part', async () => {437 // 1. Mock a stream that contains a tool call, then an invalid (empty) part.438 const streamWithToolCall = (async function* () {439 yield {440 candidates: [441 {442 content: {443 role: 'model',444 parts: [{ functionCall: { name: 'test_tool', args: {} } }],445 },446 },447 ],448 } as unknown as GenerateContentResponse;449 // This second chunk is invalid according to isValidResponse450 yield {451 candidates: [452 {453 content: {454 role: 'model',455 parts: [{ text: '' }],456 },457 },458 ],459 } as unknown as GenerateContentResponse;460 })();461 462 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(463 streamWithToolCall,464 );465 466 // 2. Action & Assert: The stream processing should complete without throwing an error467 // because the presence of a tool call makes the empty final chunk acceptable.468 const stream = await chat.sendMessageStream(469 'test-model',470 { message: 'test message' },471 'prompt-id-tool-call-empty-end',472 );473 await expect(474 (async () => {475 for await (const _ of stream) {476 /* consume stream */477 }478 })(),479 ).resolves.not.toThrow();480 481 // 3. Verify history was recorded correctly482 const history = chat.getHistory();483 expect(history.length).toBe(2); // user turn + model turn484 const modelTurn = history[1]!;485 expect(modelTurn?.parts?.length).toBe(1); // The empty part is discarded486 expect(modelTurn?.parts![0]!.functionCall).toBeDefined();487 });488 489 it('should fail if the stream ends with an empty part and has no finishReason', async () => {490 vi.useFakeTimers();491 try {492 const streamWithNoFinish = (async function* () {493 yield {494 candidates: [495 {496 content: {497 role: 'model',498 parts: [{ text: 'Initial content...' }],499 },500 },501 ],502 } as unknown as GenerateContentResponse;503 yield {504 candidates: [505 {506 content: {507 role: 'model',508 parts: [{ text: '' }],509 },510 },511 ],512 } as unknown as GenerateContentResponse;513 })();514 515 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(516 streamWithNoFinish,517 );518 519 const stream = await chat.sendMessageStream(520 'test-model',521 { message: 'test message' },522 'prompt-id-no-finish-empty-end',523 );524 await expectStreamExhaustion(stream);525 } finally {526 vi.useRealTimers();527 }528 });529 530 it('should succeed if the stream ends with an invalid part but has a finishReason and contained a valid part', async () => {531 // 1. Mock a stream that sends a valid chunk, then an invalid one, but has a finish reason.532 const streamWithInvalidEnd = (async function* () {533 yield {534 candidates: [535 {536 content: {537 role: 'model',538 parts: [{ text: 'Initial valid content...' }],539 },540 },541 ],542 } as unknown as GenerateContentResponse;543 // This second chunk is invalid, but the response has a finishReason.544 yield {545 candidates: [546 {547 content: {548 role: 'model',549 parts: [{ text: '' }], // Invalid part550 },551 finishReason: 'STOP',552 },553 ],554 } as unknown as GenerateContentResponse;555 })();556 557 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(558 streamWithInvalidEnd,559 );560 561 // 2. Action & Assert: The stream should complete without throwing an error.562 const stream = await chat.sendMessageStream(563 'test-model',564 { message: 'test message' },565 'prompt-id-valid-then-invalid-end',566 );567 await expect(568 (async () => {569 for await (const _ of stream) {570 /* consume stream */571 }572 })(),573 ).resolves.not.toThrow();574 575 // 3. Verify history was recorded correctly with only the valid part.576 const history = chat.getHistory();577 expect(history.length).toBe(2); // user turn + model turn578 const modelTurn = history[1]!;579 expect(modelTurn?.parts?.length).toBe(1);580 expect(modelTurn?.parts![0]!.text).toBe('Initial valid content...');581 });582 583 it('should consolidate subsequent text chunks after receiving an empty text chunk', async () => {584 // 1. Mock the API to return a stream where one chunk is just an empty text part.585 const multiChunkStream = (async function* () {586 yield {587 candidates: [588 { content: { role: 'model', parts: [{ text: 'Hello' }] } },589 ],590 } as unknown as GenerateContentResponse;591 // FIX: The original test used { text: '' }, which is invalid.592 // A chunk can be empty but still valid. This chunk is now removed593 // as the important part is consolidating what comes after.594 yield {595 candidates: [596 {597 content: { role: 'model', parts: [{ text: ' World!' }] },598 finishReason: 'STOP',599 },600 ],601 } as unknown as GenerateContentResponse;602 })();603 604 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(605 multiChunkStream,606 );607 608 // 2. Action: Send a message and consume the stream.609 const stream = await chat.sendMessageStream(610 'test-model',611 { message: 'test message' },612 'prompt-id-empty-chunk-consolidation',613 );614 for await (const _ of stream) {615 // Consume the stream616 }617 618 // 3. Assert: Check that the final history was correctly consolidated.619 const history = chat.getHistory();620 expect(history.length).toBe(2);621 const modelTurn = history[1]!;622 expect(modelTurn?.parts?.length).toBe(1);623 expect(modelTurn?.parts![0]!.text).toBe('Hello World!');624 });625 626 it('should consolidate adjacent text parts that arrive in separate stream chunks', async () => {627 // 1. Mock the API to return a stream of multiple, adjacent text chunks.628 const multiChunkStream = (async function* () {629 yield {630 candidates: [631 { content: { role: 'model', parts: [{ text: 'This is the ' }] } },632 ],633 } as unknown as GenerateContentResponse;634 yield {635 candidates: [636 { content: { role: 'model', parts: [{ text: 'first part.' }] } },637 ],638 } as unknown as GenerateContentResponse;639 // This function call should break the consolidation.640 yield {641 candidates: [642 {643 content: {644 role: 'model',645 parts: [{ functionCall: { name: 'do_stuff', args: {} } }],646 },647 },648 ],649 } as unknown as GenerateContentResponse;650 yield {651 candidates: [652 {653 content: {654 role: 'model',655 parts: [{ text: 'This is the second part.' }],656 },657 },658 ],659 } as unknown as GenerateContentResponse;660 })();661 662 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(663 multiChunkStream,664 );665 666 // 2. Action: Send a message and consume the stream.667 const stream = await chat.sendMessageStream(668 'test-model',669 { message: 'test message' },670 'prompt-id-multi-chunk',671 );672 for await (const _ of stream) {673 // Consume the stream to trigger history recording.674 }675 676 // 3. Assert: Check that the final history was correctly consolidated.677 const history = chat.getHistory();678 679 // The history should contain the user's turn and ONE consolidated model turn.680 expect(history.length).toBe(2);681 682 const modelTurn = history[1]!;683 expect(modelTurn.role).toBe('model');684 685 // The model turn should have 3 distinct parts: the merged text, the function call, and the final text.686 expect(modelTurn?.parts?.length).toBe(3);687 expect(modelTurn?.parts![0]!.text).toBe('This is the first part.');688 expect(modelTurn.parts![1]!.functionCall).toBeDefined();689 expect(modelTurn.parts![2]!.text).toBe('This is the second part.');690 });691 it('should preserve text parts that stream in the same chunk as a thought', async () => {692 // 1. Mock the API to return a single chunk containing both a thought and visible text.693 const mixedContentStream = (async function* () {694 yield {695 candidates: [696 {697 content: {698 role: 'model',699 parts: [700 { thought: 'This is a thought.' },701 { text: 'This is the visible text that should not be lost.' },702 ],703 },704 finishReason: 'STOP',705 },706 ],707 } as unknown as GenerateContentResponse;708 })();709 710 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(711 mixedContentStream,712 );713 714 // 2. Action: Send a message and fully consume the stream to trigger history recording.715 const stream = await chat.sendMessageStream(716 'test-model',717 { message: 'test message' },718 'prompt-id-mixed-chunk',719 );720 for await (const _ of stream) {721 // This loop consumes the stream.722 }723 724 // 3. Assert: Check the final state of the history.725 const history = chat.getHistory();726 727 // The history should contain two turns: the user's message and the model's response.728 expect(history.length).toBe(2);729 730 const modelTurn = history[1]!;731 expect(modelTurn.role).toBe('model');732 733 // CRUCIAL ASSERTION:734 // The buggy code would fail here, resulting in parts.length being 0.735 // The corrected code will pass, preserving the single visible text part.736 expect(modelTurn?.parts?.length).toBe(1);737 expect(modelTurn?.parts![0]!.text).toBe(738 'This is the visible text that should not be lost.',739 );740 });741 742 it('synthesizes a functionResponse for a dangling tool_use before sending', async () => {743 // End-to-end: when sendMessageStream is invoked on a chat whose744 // history carries a dangling `model[functionCall]` (typical state745 // after a Ctrl+Y race or a crash-resume on a partial-tool_use746 // turn), the inline repair pass closes the pair against the747 // just-pushed user content so the wire payload doesn't 400 with748 // "tool_use_id ... corresponding tool_use".749 chat.setHistory([750 { role: 'user', parts: [{ text: 'first message' }] },751 {752 role: 'model',753 parts: [754 {755 functionCall: {756 id: 'call_dangling_for_send',757 name: 'read_file',758 args: { path: '/tmp/x' },759 },760 },761 ],762 },763 ]);764 765 const ackStream = (async function* () {766 yield {767 candidates: [768 {769 content: { role: 'model', parts: [{ text: 'ok' }] },770 finishReason: 'STOP',771 },772 ],773 } as unknown as GenerateContentResponse;774 })();775 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(776 ackStream,777 );778 779 const stream = await chat.sendMessageStream(780 'test-model',781 { message: 'next user prompt after a stream-error-mid-tool_use' },782 'prompt-send-repair',783 );784 for await (const _ of stream) {785 /* drain */786 }787 788 const history = chat.getHistory();789 // The dangling fc should now be followed by a user turn that790 // carries both the user-supplied text AND the synthetic fr that791 // closes the pair.792 const userTurn = history[2]!;793 expect(userTurn.role).toBe('user');794 const fr = userTurn.parts!.find((p) => p.functionResponse);795 expect(fr?.functionResponse?.id).toBe('call_dangling_for_send');796 expect(fr?.functionResponse?.name).toBe('read_file');797 expect(798 (fr?.functionResponse?.response as { error?: string })?.error,799 ).toMatch(/interrupted/i);800 // The user's own text part is still present.801 expect(802 userTurn.parts!.some(803 (p) =>804 p.text === 'next user prompt after a stream-error-mid-tool_use',805 ),806 ).toBe(true);807 // tool_result block must come BEFORE the text — Anthropic-808 // compatible backends reject a user message whose first content809 // block isn't the tool_result answering the immediately preceding810 // tool_use. Mirrors upstream Claude Code's `hoistToolResults`.811 expect(userTurn.parts![0]!.functionResponse?.id).toBe(812 'call_dangling_for_send',813 );814 });815 816 it('does NOT synthesize when the user supplies a matching tool_result', async () => {817 // Retry-of-ToolResult case (lastPrompt is a functionResponse Part818 // array): the user-supplied tool_result must close the pair before819 // the inline repair pass sees it, so no synthetic error is820 // injected. Otherwise the wire payload would carry two821 // functionResponse parts for the same callId — the real one and a822 // bogus synthetic.823 chat.setHistory([824 { role: 'user', parts: [{ text: 'do the read' }] },825 {826 role: 'model',827 parts: [828 {829 functionCall: {830 id: 'call_retry_real_fr',831 name: 'read_file',832 args: { path: '/tmp/y' },833 },834 },835 ],836 },837 ]);838 839 const ackStream = (async function* () {840 yield {841 candidates: [842 {843 content: { role: 'model', parts: [{ text: 'ack' }] },844 finishReason: 'STOP',845 },846 ],847 } as unknown as GenerateContentResponse;848 })();849 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(850 ackStream,851 );852 853 const stream = await chat.sendMessageStream(854 'test-model',855 {856 message: {857 functionResponse: {858 id: 'call_retry_real_fr',859 name: 'read_file',860 response: { output: 'real-tool-output' },861 },862 },863 },864 'prompt-retry-real-fr',865 );866 for await (const _ of stream) {867 /* drain */868 }869 870 const userTurn = chat.getHistory()[2]!;871 const frParts = userTurn.parts!.filter((p) => p.functionResponse);872 // Exactly ONE functionResponse — the real one. No synthetic.873 expect(frParts.length).toBe(1);874 expect(frParts[0]!.functionResponse?.id).toBe('call_retry_real_fr');875 expect(876 (frParts[0]!.functionResponse?.response as { output?: string })?.output,877 ).toBe('real-tool-output');878 });879 880 it('should throw an error when a tool call is followed by an empty stream response', async () => {881 vi.useFakeTimers();882 try {883 // 1. Setup: A history where the model has just made a function call.884 const initialHistory: Content[] = [885 {886 role: 'user',887 parts: [{ text: 'Find a good Italian restaurant for me.' }],888 },889 {890 role: 'model',891 parts: [892 {893 functionCall: {894 name: 'find_restaurant',895 args: { cuisine: 'Italian' },896 },897 },898 ],899 },900 ];901 chat.setHistory(initialHistory);902 903 // 2. Mock the API to return an empty/thought-only stream.904 const emptyStreamResponse = (async function* () {905 yield {906 candidates: [907 {908 content: { role: 'model', parts: [{ thought: true }] },909 finishReason: 'STOP',910 },911 ],912 } as unknown as GenerateContentResponse;913 })();914 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(915 emptyStreamResponse,916 );917 918 // 3. Action: Send the function response back to the model and consume the stream.919 const stream = await chat.sendMessageStream(920 'test-model',921 {922 message: {923 functionResponse: {924 name: 'find_restaurant',925 response: { name: 'Vesuvio' },926 },927 },928 },929 'prompt-id-stream-1',930 );931 932 // 4. Assert: The stream processing should throw an InvalidStreamError.933 await expectStreamExhaustion(stream);934 } finally {935 vi.useRealTimers();936 }937 });938 939 it('should succeed when there is a tool call without finish reason', async () => {940 // Setup: Stream with tool call but no finish reason941 const streamWithToolCall = (async function* () {942 yield {943 candidates: [944 {945 content: {946 role: 'model',947 parts: [948 {949 functionCall: {950 name: 'test_function',951 args: { param: 'value' },952 },953 },954 ],955 },956 // No finishReason957 },958 ],959 } as unknown as GenerateContentResponse;960 })();961 962 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(963 streamWithToolCall,964 );965 966 const stream = await chat.sendMessageStream(967 'test-model',968 { message: 'test' },969 'prompt-id-1',970 );971 972 // Should not throw an error973 await expect(974 (async () => {975 for await (const _ of stream) {976 // consume stream977 }978 })(),979 ).resolves.not.toThrow();980 });981 982 it('persists partial assistant turn when stream throws after a tool_use chunk', async () => {983 // Weak-network scenario: Anthropic-compatible providers emit the984 // `functionCall` part on `content_block_stop`; the SSE may then drop985 // before `message_stop`. The yielded chunk is enough for `Turn.run`986 // to queue a `ToolCallRequest`, the tool scheduler will eventually987 // submit a `functionResponse` user turn — without a matching988 // tool_use in history, the next request body shows989 // `user → user[tool_result]` and DeepSeek/Anthropic rejects with990 // "tool_use_id ... must have a corresponding tool_use block in the991 // previous message". `processStreamResponse` must persist the992 // partial model turn before re-throwing so the pairing is intact.993 mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());994 const networkError = new Error('SSE connection reset by peer');995 const streamThatThrowsAfterToolCall = (async function* () {996 yield {997 candidates: [998 {999 content: {1000 role: 'model',1001 parts: [1002 {1003 functionCall: {1004 id: 'call_00_CeJrKJB0PSmXUZTCWHET7332',1005 name: 'read_file',1006 args: { path: '/tmp/x.txt' },1007 },1008 },1009 ],1010 },1011 },1012 ],1013 } as unknown as GenerateContentResponse;1014 throw networkError;1015 })();1016 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(1017 streamThatThrowsAfterToolCall,1018 );1019 1020 const stream = await chat.sendMessageStream(1021 'test-model',1022 { message: 'open /tmp/x.txt please' },1023 'prompt-weak-network-tool',1024 );1025 await expect(1026 (async () => {1027 for await (const _ of stream) {1028 /* drain */1029 }1030 })(),1031 ).rejects.toBe(networkError);1032 1033 const history = chat.getHistory();1034 expect(history.length).toBe(2);1035 expect(history[0]!.role).toBe('user');1036 const modelTurn = history[1]!;1037 expect(modelTurn.role).toBe('model');1038 expect(modelTurn.parts).toBeDefined();1039 const functionCallPart = modelTurn.parts!.find((p) => p.functionCall);1040 expect(functionCallPart?.functionCall?.id).toBe(1041 'call_00_CeJrKJB0PSmXUZTCWHET7332',1042 );1043 expect(functionCallPart?.functionCall?.name).toBe('read_file');1044 });1045 1046 it('preserves thinking parts alongside tool_use when stream throws mid-tool', async () => {1047 // Covers reasoning-mode providers (DeepSeek, Claude 4.6+) where the1048 // assistant turn carries both a thinking block and a tool_use. The1049 // partial-history push must keep the thinking part so DeepSeek's1050 // `injectThinkingOnToolUseTurns` converter pass sees an existing1051 // block on the replayed turn and does not pre-pend a synthetic one1052 // (which would discard the model's original reasoning text).1053 mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());1054 const networkError = new Error('SSE timeout');1055 const streamWithThinkingAndTool = (async function* () {1056 yield {1057 candidates: [1058 {1059 content: {1060 role: 'model',1061 parts: [{ text: 'planning the read', thought: true }],1062 },1063 },1064 ],1065 } as unknown as GenerateContentResponse;1066 yield {1067 candidates: [1068 {1069 content: {1070 role: 'model',1071 parts: [1072 {1073 functionCall: {1074 id: 'call_thinking_tool_use',1075 name: 'read_file',1076 args: { path: '/tmp/a.txt' },1077 },1078 },1079 ],1080 },1081 },1082 ],1083 } as unknown as GenerateContentResponse;1084 throw networkError;1085 })();1086 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(1087 streamWithThinkingAndTool,1088 );1089 1090 const stream = await chat.sendMessageStream(1091 'test-model',1092 { message: 'read /tmp/a.txt' },1093 'prompt-thinking-tool-weak-network',1094 );1095 await expect(1096 (async () => {1097 for await (const _ of stream) {1098 /* drain */1099 }1100 })(),1101 ).rejects.toBe(networkError);1102 1103 const history = chat.getHistory();1104 expect(history.length).toBe(2);1105 const modelTurn = history[1]!;1106 expect(modelTurn.role).toBe('model');1107 const parts = modelTurn.parts!;1108 // The thinking part must come before the functionCall — Anthropic1109 // requires thinking blocks first in the assistant content array.1110 expect(parts[0]!.thought).toBe(true);1111 expect(parts[0]!.text).toBe('planning the read');1112 const functionCallPart = parts.find((p) => p.functionCall);1113 expect(functionCallPart?.functionCall?.id).toBe('call_thinking_tool_use');1114 });1115 1116 it('does NOT persist partial assistant turn when stream throws before any tool_use chunk', async () => {1117 // Plain-text partial responses are deliberately dropped on stream1118 // error: the Retry path pops the trailing user prompt and re-issues1119 // it, so a stale partial-text model turn between them would bias1120 // the retry or surface as duplicate output. Only tool_use turns1121 // need the partial-history bridge to preserve the tool_use →1122 // tool_result invariant — text alone has no such invariant.1123 mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());1124 const networkError = new Error('connection reset');1125 const streamThatThrowsAfterText = (async function* () {1126 yield {1127 candidates: [1128 {1129 content: {1130 role: 'model',1131 parts: [{ text: 'partial reply that will be lost' }],1132 },1133 },1134 ],1135 } as unknown as GenerateContentResponse;1136 throw networkError;1137 })();1138 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(1139 streamThatThrowsAfterText,1140 );1141 1142 const stream = await chat.sendMessageStream(1143 'test-model',1144 { message: 'hello' },1145 'prompt-weak-network-text',1146 );1147 await expect(1148 (async () => {1149 for await (const _ of stream) {1150 /* drain */1151 }1152 })(),1153 ).rejects.toBe(networkError);1154 1155 const history = chat.getHistory();1156 // Only the user turn is in history — the partial-text model turn is1157 // intentionally not persisted.1158 expect(history.length).toBe(1);1159 expect(history[0]!.role).toBe('user');1160 });1161 1162 it('should throw InvalidStreamError when no tool call and no finish reason', async () => {1163 vi.useFakeTimers();1164 try {1165 // Setup: Stream with text but no finish reason and no tool call1166 const streamWithoutFinishReason = (async function* () {1167 yield {1168 candidates: [1169 {1170 content: {1171 role: 'model',1172 parts: [{ text: 'some response' }],1173 },1174 // No finishReason1175 },1176 ],1177 } as unknown as GenerateContentResponse;1178 })();1179 1180 vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(1181 streamWithoutFinishReason,1182 );1183 1184 const stream = await chat.sendMessageStream(1185 'test-model',1186 { message: 'test' },1187 'prompt-id-1',1188 );1189 await expectStreamExhaustion(stream);1190 } finally {1191 vi.useRealTimers();1192 }1193 });1194 1195 it('should throw InvalidStreamError when there is finish reason but truly empty response (no text, no thought)', async () => {1196 vi.useFakeTimers();1197 try {1198 // Setup: Stream with finish reason but completely empty parts1199 const streamWithEmptyResponse = (async function* () {1200 yield {