fix: prevent stalls while processing large diffs

src/generateCommitMessageFromGitDiff.ts

@@ -18,7 +18,11 @@ import {
 } from './utils/errors';
 import { GenerateCommitMessageErrorEnum } from './utils/generateCommitMessageErrors';
 import { mergeDiffs } from './utils/mergeDiffs';
-import { tokenCount } from './utils/tokenCount';
+import {
+  splitByTokenLimit,
+  tokenCount,
+  tokenCountAsync
+} from './utils/tokenCount';
 
 const config = getConfig();
 const MAX_TOKENS_INPUT = config.OCO_TOKENS_MAX_INPUT;
@@ -158,7 +162,7 @@ export const generateCommitMessageByDiff = async (
       INIT_MESSAGES_PROMPT_LENGTH -
       MAX_TOKENS_OUTPUT;
 
-    if (tokenCount(diff) >= MAX_REQUEST_TOKENS) {
+    if ((await tokenCountAsync(diff)) >= MAX_REQUEST_TOKENS) {
       const commitMessagePromises = await getCommitMsgsPromisesFromFileDiffs(
         diff,
         MAX_REQUEST_TOKENS,
@@ -215,28 +219,28 @@ export const generateCommitMessageByDiff = async (
   }
 };
 
-function getMessagesPromisesByChangesInFile(
+async function getMessagesPromisesByChangesInFile(
   fileDiff: string,
   separator: string,
   maxChangeLength: number,
   fullGitMojiSpec: boolean,
   context: string
-) {
+): Promise<Array<Promise<string | null | undefined>>> {
   const hunkHeaderSeparator = '@@ ';
   const [fileHeader, ...fileDiffByLines] = fileDiff.split(hunkHeaderSeparator);
 
   // merge multiple line-diffs into 1 to save tokens
-  const mergedChanges = mergeDiffs(
+  const mergedChanges = await mergeDiffs(
     fileDiffByLines.map((line) => hunkHeaderSeparator + line),
     maxChangeLength
   );
 
   const lineDiffsWithHeader = [] as string[];
   for (const change of mergedChanges) {
     const totalChange = fileHeader + change;
-    if (tokenCount(totalChange) > maxChangeLength) {
+    if ((await tokenCountAsync(totalChange)) > maxChangeLength) {
       // If the totalChange is too large, split it into smaller pieces
-      const splitChanges = splitDiff(totalChange, maxChangeLength);
+      const splitChanges = await splitDiff(totalChange, maxChangeLength);
       lineDiffsWithHeader.push(...splitChanges);
     } else {
       lineDiffsWithHeader.push(totalChange);
@@ -259,40 +263,12 @@ function getMessagesPromisesByChangesInFile(
   return commitMsgsFromFileLineDiffs;
 }
 
-function splitDiff(diff: string, maxChangeLength: number) {
-  const lines = diff.split('\n');
-  const splitDiffs = [] as string[];
-  let currentDiff = '';
-
+async function splitDiff(diff: string, maxChangeLength: number) {
   if (maxChangeLength <= 0) {
     throw new Error(GenerateCommitMessageErrorEnum.outputTokensTooHigh);
   }
 
-  for (let line of lines) {
-    // If a single line exceeds maxChangeLength, split it into multiple lines
-    while (tokenCount(line) > maxChangeLength) {
-      const subLine = line.substring(0, maxChangeLength);
-      line = line.substring(maxChangeLength);
-      splitDiffs.push(subLine);
-    }
-
-    // Check the tokenCount of the currentDiff and the line separately
-    if (tokenCount(currentDiff) + tokenCount('\n' + line) > maxChangeLength) {
-      // If adding the next line would exceed the maxChangeLength, start a new diff
-      splitDiffs.push(currentDiff);
-      currentDiff = line;
-    } else {
-      // Otherwise, add the line to the current diff
-      currentDiff += '\n' + line;
-    }
-  }
-
-  // Add the last diff
-  if (currentDiff) {
-    splitDiffs.push(currentDiff);
-  }
-
-  return splitDiffs;
+  return splitByTokenLimit(diff, maxChangeLength);
 }
 
 export const getCommitMsgsPromisesFromFileDiffs = async (
@@ -306,14 +282,14 @@ export const getCommitMsgsPromisesFromFileDiffs = async (
   const diffByFiles = diff.split(separator).slice(1);
 
   // merge multiple files-diffs into 1 prompt to save tokens
-  const mergedFilesDiffs = mergeDiffs(diffByFiles, maxDiffLength);
+  const mergedFilesDiffs = await mergeDiffs(diffByFiles, maxDiffLength);
 
   const commitMessagePromises = [] as Promise<string | null | undefined>[];
 
   for (const fileDiff of mergedFilesDiffs) {
-    if (tokenCount(fileDiff) >= maxDiffLength) {
+    if ((await tokenCountAsync(fileDiff)) >= maxDiffLength) {
       // if file-diff is bigger than gpt context — split fileDiff into lineDiff
-      const messagesPromises = getMessagesPromisesByChangesInFile(
+      const messagesPromises = await getMessagesPromisesByChangesInFile(
         fileDiff,
         separator,
         maxDiffLength,

src/utils/mergeDiffs.ts

@@ -1,14 +1,27 @@
-import { tokenCount } from './tokenCount';
+import { tokenCountAsync } from './tokenCount';
+
+export async function mergeDiffs(
+  arr: string[],
+  maxStringLength: number
+): Promise<string[]> {
+  if (!arr.length) return [];
 
-export function mergeDiffs(arr: string[], maxStringLength: number): string[] {
   const mergedArr: string[] = [];
   let currentItem: string = arr[0];
+  let currentItemTokens = await tokenCountAsync(currentItem);
+
   for (const item of arr.slice(1)) {
-    if (tokenCount(currentItem + item) <= maxStringLength) {
+    const itemTokens = await tokenCountAsync(item);
+
+    // Adding independently counted chunks is conservative at a BPE boundary
+    // and avoids repeatedly tokenizing an ever-growing merged diff.
+    if (currentItemTokens + itemTokens <= maxStringLength) {
       currentItem += item;
+      currentItemTokens += itemTokens;
     } else {
       mergedArr.push(currentItem);
       currentItem = item;
+      currentItemTokens = itemTokens;
     }
   }
 

src/utils/tokenCount.ts

@@ -1,13 +1,129 @@
 import cl100k_base from '@dqbd/tiktoken/encoders/cl100k_base.json';
 import { Tiktoken } from '@dqbd/tiktoken/lite';
 
-export function tokenCount(content: string): number {
-  const encoding = new Tiktoken(
+const TOKENIZER_CHUNK_LENGTH = 8_000;
+
+interface CountedTextChunk {
+  content: string;
+  tokens: number;
+}
+
+let encoding: Tiktoken | undefined;
+
+const getEncoding = (): Tiktoken => {
+  // OpenCommit is a short-lived CLI. Reusing the WASM encoder avoids paying its
+  // initialization cost for every file, line, and growing merged diff.
+  encoding ??= new Tiktoken(
     cl100k_base.bpe_ranks,
     cl100k_base.special_tokens,
     cl100k_base.pat_str
   );
-  const tokens = encoding.encode(content);
-  encoding.free();
-  return tokens.length;
+
+  return encoding;
+};
+
+const getSafeSliceEnd = (content: string, start: number, length: number) => {
+  let end = Math.min(start + length, content.length);
+
+  // Do not split a UTF-16 surrogate pair between tokenizer chunks.
+  if (
+    end < content.length &&
+    end > start &&
+    /[\uD800-\uDBFF]/.test(content[end - 1]) &&
+    /[\uDC00-\uDFFF]/.test(content[end])
+  ) {
+    end -= 1;
+  }
+
+  return end;
+};
+
+const getTextChunks = (content: string): string[] => {
+  const chunks: string[] = [];
+
+  for (let start = 0; start < content.length; ) {
+    const end = getSafeSliceEnd(content, start, TOKENIZER_CHUNK_LENGTH);
+    chunks.push(content.slice(start, end));
+    start = end;
+  }
+
+  return chunks;
+};
+
+const countTextChunk = (content: string): number =>
+  getEncoding().encode(content).length;
+
+const yieldToEventLoop = () =>
+  new Promise<void>((resolve) => setImmediate(resolve));
+
+export function tokenCount(content: string): number {
+  return getTextChunks(content).reduce(
+    (total, chunk) => total + countTextChunk(chunk),
+    0
+  );
+}
+
+export async function tokenCountAsync(content: string): Promise<number> {
+  let total = 0;
+
+  for (const chunk of getTextChunks(content)) {
+    total += countTextChunk(chunk);
+    await yieldToEventLoop();
+  }
+
+  return total;
+}
+
+const getBoundedTokenChunks = async (
+  content: string,
+  maxTokens: number
+): Promise<CountedTextChunk[]> => {
+  const tokens = countTextChunk(content);
+  await yieldToEventLoop();
+
+  if (tokens <= maxTokens) return [{ content, tokens }];
+
+  const middle = getSafeSliceEnd(content, 0, Math.floor(content.length / 2));
+
+  // A single Unicode code point cannot be divided without corrupting it. This
+  // is only reachable with impractically tiny token limits.
+  if (middle === 0 || middle === content.length) return [{ content, tokens }];
+
+  return [
+    ...(await getBoundedTokenChunks(content.slice(0, middle), maxTokens)),
+    ...(await getBoundedTokenChunks(content.slice(middle), maxTokens))
+  ];
+};
+
+export async function splitByTokenLimit(
+  content: string,
+  maxTokens: number
+): Promise<string[]> {
+  if (maxTokens <= 0) throw new Error('maxTokens must be greater than zero');
+  if (!content) return [];
+
+  const countedChunks: CountedTextChunk[] = [];
+
+  for (const chunk of getTextChunks(content)) {
+    countedChunks.push(...(await getBoundedTokenChunks(chunk, maxTokens)));
+  }
+
+  const mergedChunks: string[] = [];
+  let currentContent = '';
+  let currentTokens = 0;
+
+  for (const chunk of countedChunks) {
+    if (currentContent && currentTokens + chunk.tokens > maxTokens) {
+      mergedChunks.push(currentContent);
+      currentContent = '';
+      currentTokens = 0;
+    }
+
+    currentContent += chunk.content;
+    currentTokens += chunk.tokens;
+  }
+
+  if (currentContent) mergedChunks.push(currentContent);
+
+  return mergedChunks;
 }

test/e2e/cliBehavior.test.ts

@@ -115,6 +115,38 @@ it('cli flow passes --fgm through to the full GitMoji prompt', async () => {
   }
 });
 
+it('cli flow processes a large minified diff without stalling', async () => {
+  const { gitDir, cleanup } = await prepareEnvironment();
+  const server = await startMockOpenAiServer(
+    'fix(diff): process a large minified file'
+  );
+
+  try {
+    await prepareRepo(
+      gitDir,
+      {
+        'content.json': `{"content":"${'a'.repeat(100_000)}"}`
+      },
+      { stage: true }
+    );
+
+    const oco = await runCli(['--yes'], {
+      cwd: gitDir,
+      env: getMockOpenAiEnv(server.baseUrl, {
+        OCO_TOKENS_MAX_INPUT: '4096',
+        OCO_TOKENS_MAX_OUTPUT: '500'
+      })
+    });
+
+    expect(await waitForExit(oco, 30_000)).toBe(0);
+    await assertHeadCommit(gitDir, 'fix(diff): process a large minified file');
+    expect(server.requestBodies.length).toBeGreaterThan(1);
+  } finally {
+    await server.cleanup();
+    await cleanup();
+  }
+});
+
 it('cli flow allows editing the generated commit message before committing', async () => {
   const { gitDir, cleanup } = await prepareEnvironment();
   const server = await startMockOpenAiServer(

test/unit/mergeDiffs.test.ts

@@ -0,0 +1,18 @@
+import { mergeDiffs } from '../../src/utils/mergeDiffs';
+
+describe('mergeDiffs', () => {
+  it('returns no chunks for an empty diff list', async () => {
+    await expect(mergeDiffs([], 10)).resolves.toEqual([]);
+  });
+
+  it('merges adjacent diffs without repeatedly counting the merged text', async () => {
+    await expect(mergeDiffs(['a', 'b', 'c'], 2)).resolves.toEqual(['ab', 'c']);
+  });
+
+  it('preserves every diff when splitting groups', async () => {
+    const diffs = ['first change\n', 'second change\n', 'third change\n'];
+    const chunks = await mergeDiffs(diffs, 3);
+
+    expect(chunks.join('')).toBe(diffs.join(''));
+  });
+});

test/unit/tokenCount.test.ts

@@ -0,0 +1,41 @@
+import {
+  splitByTokenLimit,
+  tokenCount,
+  tokenCountAsync
+} from '../../src/utils/tokenCount';
+
+describe('tokenCount', () => {
+  it('counts short text exactly', () => {
+    expect(tokenCount('hello world')).toBe(2);
+  });
+
+  it('yields to the event loop while counting a long minified line', async () => {
+    let timerRan = false;
+    const timer = setTimeout(() => {
+      timerRan = true;
+    }, 0);
+
+    const count = await tokenCountAsync('a'.repeat(24_000));
+    clearTimeout(timer);
+
+    expect(count).toBe(3_000);
+    expect(timerRan).toBe(true);
+  });
+});
+
+describe('splitByTokenLimit', () => {
+  it('preserves long single-line content in bounded chunks', async () => {
+    const content = `${'a'.repeat(20_000)}${'🙂'.repeat(500)}`;
+    const chunks = await splitByTokenLimit(content, 500);
+
+    expect(chunks.length).toBeGreaterThan(1);
+    expect(chunks.join('')).toBe(content);
+    expect(chunks.every((chunk) => tokenCount(chunk) <= 500)).toBe(true);
+  });
+
+  it('rejects non-positive token limits', async () => {
+    await expect(splitByTokenLimit('content', 0)).rejects.toThrow(
+      'maxTokens must be greater than zero'
+    );
+  });
+});