include remaining workspace changes

src/server/transformers/gemini/generate-content/stream.test.ts

@@ -5,8 +5,10 @@ import { extractResponseMetadata, serializeGeminiAggregateResponse } from './out
 import {
   applyJsonPayloadToAggregate,
   applySsePayloadsToAggregate,
+  consumeUpstreamSseBuffer,
   geminiGenerateContentStream,
   parseGeminiStreamPayload,
+  serializeUpstreamJsonPayload,
   serializeAggregateSsePayload,
 } from './stream.js';
 import { extractGeminiUsage } from './usage.js';
@@ -169,6 +171,56 @@ describe('geminiGenerateContentStream', () => {
     });
   });
 
+  it('returns raw upstream sse blocks while aggregating tool-calling state', () => {
+    const state = createGeminiGenerateContentAggregateState();
+    const firstBlock = 'data: {"promptFeedback":{"blockReason":"BLOCK_REASON_UNSPECIFIED"},"candidates":[{"content":{"parts":[{"functionCall":{"id":"tool-1","name":"lookup","args":{"q":"cat"}},"thoughtSignature":"sig-tool-1"}]}}]}\r\n\r\n';
+    const secondBlock = 'data: {"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":4,"totalTokenCount":14}}\r\n\r\n';
+    const doneBlock = 'data: [DONE]\r\n\r\n';
+    const result = consumeUpstreamSseBuffer(
+      state,
+      `${firstBlock}${secondBlock}${doneBlock}data: {"responseId":"partial"`,
+    );
+
+    expect(result.lines).toEqual([firstBlock, secondBlock, doneBlock]);
+    expect(result.events).toEqual([
+      {
+        promptFeedback: { blockReason: 'BLOCK_REASON_UNSPECIFIED' },
+        candidates: [
+          {
+            content: {
+              parts: [{ functionCall: { id: 'tool-1', name: 'lookup', args: { q: 'cat' } }, thoughtSignature: 'sig-tool-1' }],
+            },
+          },
+        ],
+      },
+      {
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'answer' }],
+            },
+            finishReason: 'STOP',
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 10,
+          candidatesTokenCount: 4,
+          totalTokenCount: 14,
+        },
+      },
+    ]);
+    expect(result.rest).toBe('data: {"responseId":"partial"');
+    expect(extractGeminiUsage(result.state)).toEqual({
+      promptTokens: 10,
+      completionTokens: 4,
+      totalTokens: 14,
+      cachedTokens: 0,
+      cacheReadTokens: 0,
+      cacheCreationTokens: 0,
+      reasoningTokens: 0,
+    });
+  });
+
   it('normalizes object payloads through the JSON-array path', () => {
     const state = createGeminiGenerateContentAggregateState();
 
@@ -203,4 +255,47 @@ describe('geminiGenerateContentStream', () => {
       ],
     });
   });
+
+  it('keeps json stream chunks raw while still aggregating usage', () => {
+    const state = createGeminiGenerateContentAggregateState();
+    const payload = [
+      {
+        promptFeedback: { blockReason: 'BLOCK_REASON_UNSPECIFIED' },
+        candidates: [
+          {
+            content: {
+              parts: [{ functionCall: { id: 'tool-1', name: 'lookup', args: { q: 'cat' } } }],
+            },
+          },
+        ],
+      },
+      {
+        serverContent: { modelTurn: { parts: [{ text: 'tool result received' }] } },
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'final answer' }],
+            },
+            finishReason: 'STOP',
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 8,
+          candidatesTokenCount: 5,
+          totalTokenCount: 13,
+        },
+      },
+    ];
+
+    expect(serializeUpstreamJsonPayload(state, payload, true)).toEqual(payload);
+    expect(extractGeminiUsage(state)).toEqual({
+      promptTokens: 8,
+      completionTokens: 5,
+      totalTokens: 13,
+      cachedTokens: 0,
+      cacheReadTokens: 0,
+      cacheCreationTokens: 0,
+      reasoningTokens: 0,
+    });
+  });
 });

src/server/transformers/gemini/generate-content/stream.ts

@@ -7,6 +7,7 @@ import { serializeGeminiAggregateResponse } from './outbound.js';
 
 type ParsedSsePayloads = {
   events: unknown[];
+  lines: string[];
   rest: string;
 };
 
@@ -15,6 +16,7 @@ type GeminiGenerateContentStreamFormat = 'sse' | 'json';
 type ParsedGeminiStreamPayload = {
   format: GeminiGenerateContentStreamFormat;
   events: unknown[];
+  lines?: string[];
   rest: string;
 };
 
@@ -27,19 +29,24 @@ function serializeSsePayload(payload: unknown): string {
 }
 
 function parseSsePayloads(buffer: string): ParsedSsePayloads {
-  const normalized = buffer.replace(/\r\n/g, '\n');
   const events: unknown[] = [];
-  let rest = normalized;
+  const lines: string[] = [];
+  let cursor = 0;
 
-  while (true) {
-    const boundary = rest.indexOf('\n\n');
-    if (boundary < 0) break;
+  while (cursor < buffer.length) {
+    const boundaryMatch = /\r?\n\r?\n/.exec(buffer.slice(cursor));
+    if (!boundaryMatch || typeof boundaryMatch.index !== 'number') break;
+
+    const boundary = cursor + boundaryMatch.index;
+    const block = buffer.slice(cursor, boundary);
+    const rawBlock = buffer.slice(cursor, boundary + boundaryMatch[0].length);
+    cursor = boundary + boundaryMatch[0].length;
 
-    const block = rest.slice(0, boundary);
-    rest = rest.slice(boundary + 2);
     if (!block.trim()) continue;
+    lines.push(rawBlock);
 
     const data = block
+      .replace(/\r\n/g, '\n')
       .split('\n')
       .filter((line) => line.startsWith('data:'))
       .map((line) => line.slice(5).trimStart())
@@ -55,7 +62,11 @@ function parseSsePayloads(buffer: string): ParsedSsePayloads {
     }
   }
 
-  return { events, rest };
+  return {
+    events,
+    lines,
+    rest: buffer.slice(cursor),
+  };
 }
 
 function isRecord(value: unknown): value is Record<string, unknown> {
@@ -77,6 +88,7 @@ function parseGeminiStreamPayload(
     return {
       format: 'sse',
       events: parsed.events,
+      lines: parsed.lines,
       rest: parsed.rest,
     };
   }
@@ -123,8 +135,7 @@ function applyParsedPayloadToAggregate(
   }
 
   return {
-    events: parsed.events,
-    rest: parsed.rest,
+    ...parsed,
     state,
   };
 }
@@ -151,7 +162,7 @@ function consumeUpstreamSseBuffer(
   const applied = applySsePayloadsToAggregate(state, buffer);
   return {
     ...applied,
-    lines: applied.events.map((event) => serializeAggregateSsePayload(event)),
+    lines: applied.lines ?? [],
   };
 }
 
@@ -182,10 +193,11 @@ function serializeUpstreamJsonPayload(
   streamAction = false,
 ): unknown {
   if (streamAction) {
-    return parseJsonArrayPayload(payload).map((event) => {
+    const events = parseJsonArrayPayload(payload);
+    for (const event of events) {
       applyGeminiGenerateContentAggregate(state, event);
-      return serializeAggregateJsonPayload(event);
-    });
+    }
+    return payload;
   }
 
   applyJsonPayloadToAggregate(state, payload);