Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 67 additions & 0 deletions src/openai/format.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,73 @@ function runWithinBudget<T>(fn: () => T, message: string): T {
);
})();

// Guards the truncation loop against collapsing the user slot to the empty
// string for token-dense, space-sparse input. The word-count formula computes
// `end = words.length - Math.ceil((tokens - MAX)/2)`; when the overage-per-2 is
// greater than or equal to the word count, `Array.prototype.slice(0, end)`
// clamps to `[]` and the user message becomes `""`. Compact JSON (the default
// output of `JSON.stringify`, with no inter-token whitespace) hits this path
// because a large blob encodes to thousands of tokens but splits into only a
// handful of space-delimited words.
(function testTokenDenseSpaceSparseInputDoesNotCollapseToEmpty() {
const jsonError = JSON.stringify({
error: 'InternalServerError',
code: 500,
message: 'DB pool exhausted',
stack: Array.from({ length: 120 }, (_, i) => ({
file: `/app/src/modules/mod${i}/service${i}.ts`,
line: i * 50,
col: i,
fn: `handleRequest${i}`,
err: `ECONNREFUSED:${i}`,
})),
});
const text = [
`What did you do yesterday?\nI hit a 500 error and captured the response:\n${jsonError}`,
'What will you do today?\nDebug the API error.',
'Any blockers?\nDB pool issue.',
].join('\n\n');
assert(encode(text).length > MODEL_MAX_TOKENS, 'compact-JSON input is over the token budget before truncation');
assert(
text.split(' ').length < Math.ceil((encode(text).length - MODEL_MAX_TOKENS) / 2),
'compact-JSON input is in the single-step collapse region (C >= W)',
);

const messages = runWithinBudget(
() => textToCompletionPrompt(ModelType.Vague, text),
'compact-JSON input does not infinite-loop',
);

const userContent = userContentOf(messages);
assert(userContent.length > 0, 'compact-JSON input does not collapse the user slot to the empty string');
const finalTokens = encode(userContent).length;
assert(
finalTokens <= MODEL_MAX_TOKENS,
`compact-JSON user content fits the token budget (final=${finalTokens}, budget=${MODEL_MAX_TOKENS})`,
);
})();

// Guards the single-word over-budget edge case: when the entire over-budget
// input is one whitespace-free token, the keep-count clamps to 1 and
// `slice(0, 1)` returns the same word, so the `next >= truncatedTextTokenLength`
// guard must break out of the loop rather than spinning forever. The result
// stays non-empty (the whole word is preserved since it cannot be split), even
// though it remains over the model's token budget.
(function testSingleOverBudgetWordTerminatesAndStaysNonEmpty() {
const text = 'A'.repeat(20000); // ~2500 tokens, single whitespace-free "word"
assert(encode(text).length > MODEL_MAX_TOKENS, 'single-word input is over the token budget before truncation');
assert(text.split(' ').length === 1, 'single-word input has exactly one space-delimited word');

const messages = runWithinBudget(
() => textToCompletionPrompt(ModelType.Vague, text),
'single over-budget word does not infinite-loop',
);

const userContent = userContentOf(messages);
assert(userContent.length > 0, 'single over-budget word does not collapse the user slot to the empty string');
assertEqual(userContent, text, 'single over-budget word is preserved verbatim (cannot be sub-word split)');
})();

// Guards against over-truncation: input already within the budget must be
// passed through verbatim.
(function testUnderlimitInputPassedThroughUnchanged() {
Expand Down
12 changes: 7 additions & 5 deletions src/openai/format.ts
Original file line number Diff line number Diff line change
Expand Up @@ -81,11 +81,13 @@ export function textToCompletionPrompt(

while (truncatedTextTokenLength > MODEL_MAX_TOKENS) {
const words = truncatedText.split(' ');
truncatedText = words
.slice(0, words.length - Math.ceil((truncatedTextTokenLength - MODEL_MAX_TOKENS) / 2))
.join(' ');

truncatedTextTokenLength = encode(truncatedText).length;
let keep = words.length - Math.ceil((truncatedTextTokenLength - MODEL_MAX_TOKENS) / 2);
if (keep < 1) keep = 1;
if (keep > words.length) keep = words.length;
truncatedText = words.slice(0, keep).join(' ');
const next = encode(truncatedText).length;
if (next >= truncatedTextTokenLength) break;

@augmentcode augmentcode Bot Sep 25, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For a single whitespace-free value whose token count exceeds MODEL_MAX_TOKENS, next equals the original count, so this break returns that still-oversized value. getTrainingExample serializes it directly into the fine-tuning JSONL, so this reintroduces an invalid over-budget training example (and similarly lets scoring requests exceed the function's documented limit).

Severity: medium

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checked src/openai/format.ts:82-91 and the PR's own testSingleOverBudgetWordTerminatesAndStaysNonEmpty (src/openai/format.test.ts:154-167), which deliberately asserts an unsplittable over-budget single word is preserved verbatim and stays over MODEL_MAX_TOKENS (the comment states "the whole word is preserved since it cannot be split, even though it remains over the model's token budget").

I attempted to fix this by hard-truncating to the token budget once word-splitting can no longer reduce the count (sub-word cut via decode(encode(text).slice(0, MODEL_MAX_TOKENS))). That brings the over-budget single word under MODEL_MAX_TOKENS, but it directly fails the PR's assertEqual(userContent, text, '... preserved verbatim ...') contract, so landing it would rewrite that test and change the PR's stated purpose. For a single whitespace-free over-budget token, "non-empty" and "under budget" are mutually exclusive under word-based truncation; getting under budget requires sub-word splitting, which the PR deliberately rejects.

Leaving the code unchanged pending a human decision: (a) keep the current behavior — preserve non-empty verbatim and accept that an unsplittable over-budget input exceeds MODEL_MAX_TOKENS in getTrainingExample/scoring, or (b) change the PR's purpose to hard-truncate unsplittable over-budget input to the token budget (which rewrites testSingleOverBudgetWordTerminatesAndStaysNonEmpty).

truncatedTextTokenLength = next;
}

return [
Expand Down
Loading