fix: preserve UTF-8 in logged ReadableStream byte chunks (#2701)

Co-authored-by: lforst <8118419+lforst@users.noreply.github.com>

Author: DiadkoShmek

.changeset/fix-logged-byte-stream-utf8.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix: preserve UTF-8 in logged ReadableStream byte chunks

js/src/functions/stream.test.ts

@@ -1,4 +1,4 @@
-import { expect, test } from "vitest";
+import { expect, test, vi } from "vitest";
 import {
   BraintrustStream,
   BraintrustStreamChunk,
@@ -103,3 +103,64 @@ test("final value passthrough with abort", async () => {
     "This operation was aborted",
   );
 });
+
+async function collectFinalValue(
+  chunks: (string | Uint8Array | BraintrustStreamChunk)[],
+) {
+  const final = vi.fn();
+  const error = vi.fn();
+  const emitted: BraintrustStreamChunk[] = [];
+  const input = new ReadableStream({
+    start(controller) {
+      for (const chunk of chunks) controller.enqueue(chunk);
+      controller.close();
+    },
+  });
+  await input
+    .pipeThrough(createFinalValuePassThroughStream(final, error))
+    .pipeTo(
+      new WritableStream({
+        write(chunk) {
+          emitted.push(chunk);
+        },
+      }),
+    );
+  return { final, error, emitted };
+}
+
+test.each(["é", "€", "🧠", "Пам’ять 🧠 漢字"])(
+  "final value and forwarded text preserve every UTF-8 split: %s",
+  async (text) => {
+    const bytes = new TextEncoder().encode(text);
+    for (let split = 1; split < bytes.length; split++) {
+      const { final, error, emitted } = await collectFinalValue([
+        bytes.subarray(0, split),
+        new Uint8Array(),
+        bytes.subarray(split),
+      ]);
+      expect(final).toHaveBeenCalledExactlyOnceWith(text);
+      expect(error).not.toHaveBeenCalled();
+      expect(emitted.map((chunk) => chunk.data).join("")).toBe(text);
+    }
+  },
+);
+
+test.each([
+  { chunks: [Uint8Array.of(0xf0, 0x9f)], expected: "�" },
+  { chunks: [Uint8Array.of(0xf0), Uint8Array.of(0x9f)], expected: "�" },
+  {
+    chunks: [
+      Uint8Array.of(0xc3),
+      { type: "text_delta", data: "middle" } as const,
+      Uint8Array.of(0xa9),
+    ],
+    expected: "�middle�",
+  },
+])(
+  "final value replaces incomplete byte runs: $expected",
+  async ({ chunks, expected }) => {
+    const { final, emitted } = await collectFinalValue(chunks);
+    expect(final).toHaveBeenCalledExactlyOnceWith(expected);
+    expect(emitted.map((chunk) => chunk.data).join("")).toBe(expected);
+  },
+);

js/src/functions/stream.ts

@@ -317,20 +317,34 @@ export function createFinalValuePassThroughStream<
   const textChunks: string[] = [];
   const jsonChunks: string[] = [];
   const reasoningChunks: string[] = [];
+  const flushDecoder = (
+    controller: TransformStreamDefaultController<BraintrustStreamChunk>,
+  ) => {
+    const tail = decoder.decode();
+    if (tail) {
+      textChunks.push(tail);
+      controller.enqueue({ type: "text_delta", data: tail });
+    }
+  };
   const transformStream = new TransformStream<T, BraintrustStreamChunk>({
     transform(chunk, controller) {
+      // A non-byte chunk ends the current byte run, so flush it in order.
+      if (!(chunk instanceof Uint8Array)) flushDecoder(controller);
       if (typeof chunk === "string") {
         textChunks.push(chunk);
         controller.enqueue({
           type: "text_delta",
           data: chunk,
         });
       } else if (chunk instanceof Uint8Array) {
-        textChunks.push(decoder.decode(chunk));
-        controller.enqueue({
-          type: "text_delta",
-          data: decoder.decode(chunk),
-        });
+        const text = decoder.decode(chunk, { stream: true });
+        textChunks.push(text);
+        if (text) {
+          controller.enqueue({
+            type: "text_delta",
+            data: text,
+          });
+        }
       } else if (braintrustStreamChunkSchema.safeParse(chunk).success) {
         const chunkType = chunk.type;
         switch (chunkType) {
@@ -361,6 +375,7 @@ export function createFinalValuePassThroughStream<
       }
     },
     flush(controller) {
+      flushDecoder(controller);
       if (jsonChunks.length > 0) {
         // If we received both text and json deltas in the same stream, we
         // only return the json delta

js/src/logger.test.ts

@@ -456,6 +456,45 @@ test("verify MemoryBackgroundLogger intercepts logs", async () => {
   _exportsForTestingOnly.clearTestBackgroundLogger(); // can go back to normal
 });
 
+test.each([
+  ["single chunk", false],
+  ["split bytes", true],
+])(
+  "logs a Unicode byte stream without changing its text: %s",
+  async (_, split) => {
+    _exportsForTestingOnly.simulateLoginForTests();
+    const memoryLogger = _exportsForTestingOnly.useTestBackgroundLogger();
+    const logger = initLogger({
+      projectName: "test",
+      projectId: "test-project-id",
+    });
+
+    const text = "Пам’ять 🧠 漢字";
+    const bytes = new TextEncoder().encode(text);
+    const output = new ReadableStream<Uint8Array>({
+      start(controller) {
+        if (split) {
+          for (const byte of bytes) controller.enqueue(Uint8Array.of(byte));
+        } else {
+          controller.enqueue(bytes);
+        }
+        controller.close();
+      },
+    });
+    const span = logger.startSpan({ name: "byte-stream" });
+    span.log({ output });
+    span.end();
+    await memoryLogger.flush();
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const events = (await memoryLogger.drain()) as any[];
+    expect(events).toHaveLength(1);
+    expect(events[0].output).toBe(text);
+
+    _exportsForTestingOnly.clearTestBackgroundLogger();
+  },
+);
+
 test("init validation", () => {
   expect(() => init({})).toThrow(
     "Must specify at least one of project or projectId",