Avoid mid-word splits in replace text transform

.changeset/replace-prefix-holdback.md

@@ -0,0 +1,5 @@
+---
+'@livekit/agents': patch
+---
+
+Avoid splitting streamed replacement output mid-word when no replacement key prefix is pending.

agents/src/voice/transcription/text_transforms.test.ts

@@ -16,6 +16,17 @@ function streamText(text: string, chunkSize: number): ReadableStream<string> {
   });
 }
 
+function streamChunks(chunks: string[]): ReadableStream<string> {
+  return new ReadableStream<string>({
+    start(controller) {
+      for (const chunk of chunks) {
+        controller.enqueue(chunk);
+      }
+      controller.close();
+    },
+  });
+}
+
 async function collect(stream: ReadableStream<string>): Promise<string> {
   const reader = stream.getReader();
   let result = '';
@@ -31,6 +42,21 @@ async function collect(stream: ReadableStream<string>): Promise<string> {
   return result;
 }
 
+async function collectChunks(stream: ReadableStream<string>): Promise<string[]> {
+  const reader = stream.getReader();
+  const result: string[] = [];
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      result.push(value);
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  return result;
+}
+
 describe('textTransforms.replace', () => {
   for (const chunkSize of [1, 2, 5, 11, 50]) {
     it(`replaces across chunk size ${chunkSize}`, async () => {
@@ -70,6 +96,117 @@ describe('textTransforms.replace', () => {
     ).toBe(String.raw`a \1 \n \t here`);
   });
 
+  it('flushes non-prefix text immediately', async () => {
+    const transform = replace({ LiveKit: 'Lyve Kit' });
+    const chunks = await collectChunks(transform(streamText('you connect.', 100)));
+    expect(chunks).toEqual(['you connect.']);
+  });
+
+  it('preserves no-match source chunk topology', async () => {
+    const transform = replace({ LiveKit: 'Lyve Kit' });
+    expect(await collectChunks(transform(streamChunks(['you connect.'])))).toEqual([
+      'you connect.',
+    ]);
+    expect(await collectChunks(transform(streamChunks(['you con', 'nect.'])))).toEqual([
+      'you con',
+      'nect.',
+    ]);
+  });
+
+  it('holds only potential prefix', async () => {
+    const transform = replace({ LiveKit: 'Lyve Kit' });
+    const chunks = await collectChunks(transform(streamText('visit Live', 100)));
+    expect(chunks[0]).toBe('visit ');
+    expect(chunks.join('')).toBe('visit Live');
+  });
+
+  it('completes or rejects a held prefix with the next source chunk', async () => {
+    const transform = replace({ LiveKit: 'Lyve Kit' });
+    expect(await collectChunks(transform(streamChunks(['visit Live', 'Kit now'])))).toEqual([
+      'visit ',
+      'Lyve Kit now',
+    ]);
+    expect(await collectChunks(transform(streamChunks(['visit Live', 'ly now'])))).toEqual([
+      'visit ',
+      'Lively now',
+    ]);
+  });
+
+  it('holds whitespace prefixes and flushes unresolved source only at EOF', async () => {
+    const whitespaceTransform = replace({ '\nLive': ' START' });
+    expect(await collectChunks(whitespaceTransform(streamChunks(['ready\n', 'Live!'])))).toEqual([
+      'ready',
+      ' START!',
+    ]);
+
+    const unresolvedTransform = replace({ LiveKit: 'Lyve Kit' });
+    expect(await collectChunks(unresolvedTransform(streamChunks(['visit Live'])))).toEqual([
+      'visit ',
+      'Live',
+    ]);
+  });
+
+  it('prefers longest overlapping key', async () => {
+    const transform = replace({ a: 'X', ab: 'Y' });
+    expect(await collect(transform(streamText('ab', 100)))).toBe('Y');
+  });
+
+  it('retains the source split-overlap fallback policy', async () => {
+    const transform = replace({ a: 'X', ab: 'Y' });
+    expect(await collectChunks(transform(streamChunks(['a', 'b'])))).toEqual(['X', 'b']);
+  });
+
+  it('does not cascade', async () => {
+    const transform = replace({ a: 'b', b: 'c' });
+    expect(await collect(transform(streamText('a', 100)))).toBe('b');
+  });
+
+  it('does not cascade replacement output across source chunks', async () => {
+    const transform = replace({ a: 'b', bc: 'X' });
+    expect(await collectChunks(transform(streamChunks(['a', 'c'])))).toEqual(['b', 'c']);
+  });
+
+  it('replaces Unicode regex case equivalents', async () => {
+    const transform = replace({ S: 'ess', Σ: 'sigma' });
+    expect(await collectChunks(transform(streamChunks(['Aſ', 'BςC'])))).toEqual([
+      'Aess',
+      'BsigmaC',
+    ]);
+  });
+
+  it('uses source order for case-equivalent keys of equal length', async () => {
+    const transform = replace({ S: 'first', s: 'second' });
+    expect(await collect(transform(streamChunks(['s'])))).toBe('first');
+  });
+
+  it('handles UTF-16 splits without normalizing combining marks', async () => {
+    const transform = replace({ '😀x': 'smile', é: 'e' });
+    expect(
+      await collectChunks(transform(streamChunks(['go \ud83d', '\ude00x; ', 'e\u0301']))),
+    ).toEqual(['go ', 'smile; ', 'e\u0301']);
+  });
+
+  it('propagates source errors without flushing an unresolved prefix', async () => {
+    const sourceError = new Error('source failed');
+    let sourceController: ReadableStreamDefaultController<string> | undefined;
+    const source = new ReadableStream<string>({
+      start(controller) {
+        sourceController = controller;
+      },
+    });
+    const reader = replace({ LiveKit: 'Lyve Kit' })(source).getReader();
+    if (!sourceController) throw new Error('source controller was not initialized');
+
+    sourceController.enqueue('already emitted ');
+    expect(await reader.read()).toEqual({ done: false, value: 'already emitted ' });
+    sourceController.enqueue('Live');
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    sourceController.error(sourceError);
+
+    await expect(reader.read()).rejects.toBe(sourceError);
+    reader.releaseLock();
+  });
+
   it('applies callable and built-in transforms', async () => {
     expect(
       await collect(

agents/src/voice/transcription/text_transforms.ts

@@ -193,34 +193,68 @@ export function replace(
   options: { caseSensitive?: boolean } = {},
 ): (text: ReadableStream<string>) => ReadableStream<string> {
   const entries = Object.entries(replacements);
-  const tailLen = entries.length > 0 ? Math.max(...entries.map(([old]) => old.length)) - 1 : 0;
-  const flags = options.caseSensitive ? 'gu' : 'giu';
+  const sortedEntries = [...entries].sort(([a], [b]) => b.length - a.length);
+  const flags = options.caseSensitive ? 'u' : 'iu';
+  const entryPatterns = sortedEntries.map(
+    ([old, replacement]) => [new RegExp(`^(?:${escapeRegex(old)})$`, flags), replacement] as const,
+  );
+  const pattern =
+    entries.length > 0
+      ? new RegExp(sortedEntries.map(([old]) => escapeRegex(old)).join('|'), `g${flags}`)
+      : null;
+  const maxPrefix = entries.length > 0 ? Math.max(...entries.map(([old]) => old.length - 1)) : 0;
+  const prefixes = new Set<string>();
+  for (const [old] of entries) {
+    for (let length = 1; length < old.length; length += 1) {
+      prefixes.add(old.slice(0, length));
+    }
+  }
+  const holdbackPattern =
+    prefixes.size > 0
+      ? new RegExp(`(?:${Array.from(prefixes).map(escapeRegex).join('|')})$`, flags)
+      : null;
+
+  const apply = (value: string): { output: string; lastMatchEnd: number } => {
+    if (!pattern) return { output: value, lastMatchEnd: 0 };
+    let lastMatchEnd = 0;
+    const output = value.replace(pattern, (match: string, offset: number) => {
+      const entry = entryPatterns.find(([entryPattern]) => entryPattern.test(match));
+      if (!entry) {
+        throw new Error(`Unable to resolve replacement for matched text: ${match}`);
+      }
+      const replacement = entry[1];
+      lastMatchEnd = offset + match.length;
+      return replacement;
+    });
+    return { output, lastMatchEnd };
+  };
+
+  const holdback = (value: string): number => {
+    if (!holdbackPattern) return 0;
+    const match = holdbackPattern.exec(value.slice(-maxPrefix));
+    return match ? match[0].length : 0;
+  };
 
   return (text: ReadableStream<string>) =>
     streamFromAsyncIterable(
       (async function* () {
-        let buffer = '';
+        let sourceBuffer = '';
 
         for await (const chunk of readStream(text)) {
-          buffer += chunk;
-          if (buffer.length <= tailLen) {
-            continue;
+          const source = sourceBuffer + chunk;
+          const applied = apply(source);
+          const heldLength = holdback(source);
+          const heldStart = source.length - heldLength;
+          const retainSource = heldLength > 0 && heldStart >= applied.lastMatchEnd;
+          sourceBuffer = retainSource ? source.slice(heldStart) : '';
+          const emitted = retainSource ? applied.output.slice(0, -heldLength) : applied.output;
+          if (emitted) {
+            yield emitted;
           }
-
-          for (const [old, replacement] of entries) {
-            buffer = buffer.replace(new RegExp(escapeRegex(old), flags), () => replacement);
-          }
-
-          const flushTo = buffer.length - tailLen;
-          yield buffer.slice(0, flushTo);
-          buffer = buffer.slice(flushTo);
         }
 
-        if (buffer) {
-          for (const [old, replacement] of entries) {
-            buffer = buffer.replace(new RegExp(escapeRegex(old), flags), () => replacement);
-          }
-          yield buffer;
+        if (sourceBuffer) {
+          yield sourceBuffer;
         }
       })(),
     );