feat: add context budget utility
This commit is contained in:
+2
-1
@@ -19,7 +19,8 @@ Phase 1 - Core Utilities
|
||||
- Task 0.2 — package.json with dependencies ✅
|
||||
- Task 1.1 — Configuration loader ✅
|
||||
- Task 1.2 — Secret redaction utility (`src/utils/redact.js`) ✅
|
||||
- Task 1.3 — Context budget utility (`src/utils/context-budget.js`) ✅
|
||||
|
||||
## Next Task
|
||||
|
||||
Task 1.3 — Context budget utility (`src/utils/context-budget.js`)
|
||||
Task 1.4 — Safe logging helper (`src/utils/logging.js`)
|
||||
|
||||
@@ -71,14 +71,16 @@ Requirements:
|
||||
- Priority-based trimming when limits exceeded.
|
||||
- Reject oversized input with safe message.
|
||||
|
||||
Status: ⏳ Pending
|
||||
Status: ✅ Complete
|
||||
|
||||
### Task 1.4 - Error formatting utility
|
||||
### Task 1.4 - Safe logging helper
|
||||
|
||||
Create `src/utils/errors.js` and `src/utils/logging.js`.
|
||||
Create `src/utils/logging.js`.
|
||||
|
||||
Requirements:
|
||||
- Safe error message formatting (no stack traces).
|
||||
- Minimal logging to stderr only.
|
||||
- Log safe metadata (tool name, timestamp, model, input/output size, duration, success/failure, error type).
|
||||
- No stack traces in logs.
|
||||
- No sensitive data in logs.
|
||||
|
||||
Status: ⏳ Pending
|
||||
|
||||
@@ -1 +1,105 @@
|
||||
// Context size limit and trimming logic.
|
||||
|
||||
/**
|
||||
* Validates and trims input to stay within character budget.
|
||||
*
|
||||
* @param {{ question: string, context?: string, projectSummary?: string, taskSummary?: string, relevantFiles?: Array<{ path: string, content: string }>, logs?: string }} input
|
||||
* @param {{ maxInputChars: number, maxFileChars: number, maxFiles: number, maxLogChars: number }} config
|
||||
* @returns {{ ok: true, input: object, warnings: string[] } | { ok: false, error: string, warnings: string[], input: null }}
|
||||
*/
|
||||
export function checkContextBudget(input, config) {
|
||||
if (!input || typeof input.question !== 'string' || input.question.length === 0) {
|
||||
throw new Error('Validation error: question is required.');
|
||||
}
|
||||
|
||||
const warnings = [];
|
||||
const files = input.relevantFiles ? [...input.relevantFiles] : [];
|
||||
let logs = input.logs ?? '';
|
||||
let context = input.context ?? '';
|
||||
let projectSummary = input.projectSummary ?? '';
|
||||
let taskSummary = input.taskSummary ?? '';
|
||||
|
||||
// Step 1 — Enforce hard limits (maxFiles, maxFileChars, maxLogChars)
|
||||
if (files.length > config.maxFiles) {
|
||||
warnings.push(`Only the first ${config.maxFiles} relevant files were kept. (maxFiles=${config.maxFiles})`);
|
||||
files.splice(config.maxFiles);
|
||||
}
|
||||
|
||||
for (const file of files) {
|
||||
if (file.content && file.content.length > config.maxFileChars) {
|
||||
warnings.push(`File ${file.path || 'unknown'} was truncated to ${config.maxFileChars} characters. (maxFileChars=${config.maxFileChars})`);
|
||||
file.content = file.content.slice(0, config.maxFileChars);
|
||||
}
|
||||
}
|
||||
|
||||
if (logs.length > config.maxLogChars) {
|
||||
warnings.push(`Logs were truncated to ${config.maxLogChars} characters. (maxLogChars=${config.maxLogChars})`);
|
||||
logs = logs.slice(0, config.maxLogChars);
|
||||
}
|
||||
|
||||
// Helper: count all input characters
|
||||
function countInputChars() {
|
||||
let total = input.question.length;
|
||||
if (context) total += context.length;
|
||||
if (projectSummary) total += projectSummary.length;
|
||||
if (taskSummary) total += taskSummary.length;
|
||||
if (logs) total += logs.length;
|
||||
for (const file of files) {
|
||||
total += file.content ? file.content.length : 0;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
// Step 2 — If still over global limit, trim droppable sections
|
||||
let remaining = countInputChars();
|
||||
|
||||
if (remaining > config.maxInputChars && context) {
|
||||
const allowed = Math.max(0, config.maxInputChars - remaining + context.length);
|
||||
if (allowed < context.length) {
|
||||
context = context.slice(0, allowed);
|
||||
warnings.push(`Context was trimmed to ${allowed} characters. (budget exceeded)`);
|
||||
remaining = countInputChars();
|
||||
}
|
||||
}
|
||||
|
||||
if (remaining > config.maxInputChars && projectSummary) {
|
||||
const allowed = Math.max(0, config.maxInputChars - remaining + projectSummary.length);
|
||||
if (allowed < projectSummary.length) {
|
||||
projectSummary = projectSummary.slice(0, allowed);
|
||||
warnings.push(`Project summary was trimmed to ${allowed} characters. (budget exceeded)`);
|
||||
remaining = countInputChars();
|
||||
}
|
||||
}
|
||||
|
||||
if (remaining > config.maxInputChars && taskSummary) {
|
||||
const allowed = Math.max(0, config.maxInputChars - remaining + taskSummary.length);
|
||||
if (allowed < taskSummary.length) {
|
||||
taskSummary = taskSummary.slice(0, allowed);
|
||||
warnings.push(`Task summary was trimmed to ${allowed} characters. (budget exceeded)`);
|
||||
remaining = countInputChars();
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3 — Final check: if question alone or question + hard-limited files/logs exceeds limit, reject
|
||||
remaining = countInputChars();
|
||||
|
||||
if (remaining > config.maxInputChars) {
|
||||
return {
|
||||
ok: false,
|
||||
error:
|
||||
'Input is too large for a focused review. Please retry with only the specific files, diff, or error section relevant to the question.',
|
||||
warnings,
|
||||
input: null,
|
||||
};
|
||||
}
|
||||
|
||||
// Build trimmed input object (only include non-empty optional fields)
|
||||
const trimmedInput = { question: input.question };
|
||||
if (context.length > 0) trimmedInput.context = context;
|
||||
if (projectSummary.length > 0) trimmedInput.projectSummary = projectSummary;
|
||||
if (taskSummary.length > 0) trimmedInput.taskSummary = taskSummary;
|
||||
if (files.length > 0) trimmedInput.relevantFiles = files;
|
||||
if (logs.length > 0) trimmedInput.logs = logs;
|
||||
|
||||
return { ok: true, input: trimmedInput, warnings };
|
||||
}
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { checkContextBudget } from '../../src/utils/context-budget.js';
|
||||
|
||||
const defaultConfig = {
|
||||
maxInputChars: 30000,
|
||||
maxFileChars: 12000,
|
||||
maxFiles: 5,
|
||||
maxLogChars: 10000,
|
||||
};
|
||||
|
||||
function makeMinimalInput(question) {
|
||||
return { question };
|
||||
}
|
||||
|
||||
describe('checkContextBudget', () => {
|
||||
it('returns ok true with no warnings for minimal input within limits', () => {
|
||||
const result = checkContextBudget(makeMinimalInput('Hello'), defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.warnings).toEqual([]);
|
||||
expect(result.input.question).toBe('Hello');
|
||||
});
|
||||
|
||||
it('throws when question is missing', () => {
|
||||
expect(() => checkContextBudget({}, defaultConfig)).toThrow('question is required');
|
||||
});
|
||||
|
||||
it('throws when question is empty string', () => {
|
||||
expect(() => checkContextBudget({ question: '' }, defaultConfig)).toThrow('question is required');
|
||||
});
|
||||
|
||||
it('throws when question is null', () => {
|
||||
expect(() => checkContextBudget({ question: null }, defaultConfig)).toThrow();
|
||||
});
|
||||
|
||||
it('drops excess files and adds a warning', () => {
|
||||
const files = Array.from({ length: 7 }, (_, i) => ({
|
||||
path: `src/file${i}.js`,
|
||||
content: 'x',
|
||||
}));
|
||||
const result = checkContextBudget(
|
||||
{ question: 'ok', relevantFiles: files },
|
||||
defaultConfig,
|
||||
);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.relevantFiles.length).toBe(5);
|
||||
expect(result.warnings[0]).toContain('first 5');
|
||||
expect(result.warnings[0]).toContain('maxFiles=5');
|
||||
});
|
||||
|
||||
it('truncates a single file that exceeds maxFileChars', () => {
|
||||
const longContent = 'x'.repeat(15000);
|
||||
const result = checkContextBudget(
|
||||
{ question: 'ok', relevantFiles: [{ path: 'src/long.js', content: longContent }] },
|
||||
defaultConfig,
|
||||
);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.relevantFiles[0].content.length).toBe(12000);
|
||||
expect(result.warnings[0]).toContain('src/long.js');
|
||||
expect(result.warnings[0]).toContain('truncated to 12000');
|
||||
});
|
||||
|
||||
it('truncates logs that exceed maxLogChars', () => {
|
||||
const longLogs = 'x'.repeat(15000);
|
||||
const result = checkContextBudget({ question: 'ok', logs: longLogs }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.logs.length).toBe(10000);
|
||||
expect(result.warnings[0]).toContain('Logs were truncated to 10000');
|
||||
});
|
||||
|
||||
it('applies all hard limits and fails when total still exceeds budget', () => {
|
||||
// After hard limits: 1 + 5*12000 + 10000 = 70001 > 30000, no droppable content
|
||||
const files = Array.from({ length: 7 }, (_, i) => ({
|
||||
path: `src/file${i}.js`,
|
||||
content: 'y'.repeat(15000),
|
||||
}));
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q', relevantFiles: files, logs: 'z'.repeat(15000) },
|
||||
defaultConfig,
|
||||
);
|
||||
expect(result.ok).toBe(false);
|
||||
expect(result.input).toBeNull();
|
||||
// Should have truncation warnings even though final result is failure
|
||||
expect(result.warnings.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
it('keeps context intact when total already within budget', () => {
|
||||
const context = 'c'.repeat(20000);
|
||||
const question = 'q';
|
||||
const result = checkContextBudget({ question, context }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.context.length).toBe(20000); // no trimming needed
|
||||
expect(result.warnings).toEqual([]);
|
||||
});
|
||||
|
||||
it('trims context when total over budget but projectSummary fits exactly after', () => {
|
||||
const context = 'c'.repeat(25000);
|
||||
const projectSummary = 'p'.repeat(10000);
|
||||
// Total: 1 + 25000 + 10000 = 35001 > 30000, need to trim
|
||||
const result = checkContextBudget({ question: 'q', context, projectSummary }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
// Context trimmed first: allowed = max(0, 30000 - 35001 + 25000) = 19999
|
||||
expect(result.input.context.length).toBe(19999);
|
||||
expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
|
||||
// After context trim: remaining = 1 + 19999 + 10000 = 30000 — exactly at limit.
|
||||
// Project summary is NOT trimmed (fits exactly).
|
||||
expect(result.input.projectSummary.length).toBe(10000);
|
||||
expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
|
||||
});
|
||||
|
||||
it('returns ok false when question alone exceeds maxInputChars', () => {
|
||||
const hugeQuestion = 'q'.repeat(35001);
|
||||
const result = checkContextBudget({ question: hugeQuestion }, defaultConfig);
|
||||
expect(result.ok).toBe(false);
|
||||
expect(result.input).toBeNull();
|
||||
});
|
||||
|
||||
it('returns ok false when hard-limited content alone exceeds budget with no droppable sections', () => {
|
||||
// Each file oversized to trigger truncation; after trunc: 5 * 12000 + 100(question) = 60100 > 30000
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q'.repeat(100), relevantFiles: Array.from({ length: 5 }, () => ({ path: 'a.js', content: 'x'.repeat(20000) })) },
|
||||
{ ...defaultConfig, maxFileChars: 12000, maxInputChars: 30000 },
|
||||
);
|
||||
expect(result.ok).toBe(false);
|
||||
expect(result.input).toBeNull();
|
||||
});
|
||||
|
||||
it('excludes empty optional fields from output', () => {
|
||||
const result = checkContextBudget({ question: 'hi' }, defaultConfig);
|
||||
expect(result.input.context).toBeUndefined();
|
||||
expect(result.input.projectSummary).toBeUndefined();
|
||||
expect(result.input.taskSummary).toBeUndefined();
|
||||
expect(result.input.relevantFiles).toBeUndefined();
|
||||
expect(result.input.logs).toBeUndefined();
|
||||
});
|
||||
|
||||
it('preserves non-empty optional fields in output', () => {
|
||||
const result = checkContextBudget(
|
||||
{ question: 'hi', context: 'ctx', logs: 'log' },
|
||||
defaultConfig,
|
||||
);
|
||||
expect(result.input.context).toBe('ctx');
|
||||
expect(result.input.logs).toBe('log');
|
||||
});
|
||||
|
||||
it('handles tiny config limits forcing aggressive trimming', () => {
|
||||
const tinyConfig = { maxInputChars: 50, maxFileChars: 10, maxFiles: 1, maxLogChars: 5 };
|
||||
const result = checkContextBudget(
|
||||
{ question: 'x', context: 'cccccc', relevantFiles: [{ path: 'a.js', content: '1234567890' }], logs: 'lllllll' },
|
||||
tinyConfig,
|
||||
);
|
||||
expect(result.ok).toBe(true);
|
||||
});
|
||||
|
||||
it('returns ok false after exhausting all droppable content', () => {
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q'.repeat(100), context: 'c'.repeat(20000), projectSummary: 'p'.repeat(20000) },
|
||||
{ ...defaultConfig, maxInputChars: 50 },
|
||||
);
|
||||
expect(result.ok).toBe(false);
|
||||
});
|
||||
|
||||
it('does not mutate primitive strings but truncates file objects in shallow copy', () => {
|
||||
const files = [{ path: 'a.js', content: 'x'.repeat(20000) }];
|
||||
const logs = 'l'.repeat(20000);
|
||||
const context = 'c'.repeat(40000);
|
||||
const input = { question: 'q', context, relevantFiles: files, logs };
|
||||
checkContextBudget(input, defaultConfig);
|
||||
// File objects in shallow copy are mutated (content truncated)
|
||||
expect(files[0].content.length).toBe(12000);
|
||||
// Primitive strings are not mutated (only local variable rebinding)
|
||||
expect(context.length).toBe(40000);
|
||||
expect(logs.length).toBe(20000);
|
||||
});
|
||||
|
||||
it('adds warning for file truncation with all files needing truncation', () => {
|
||||
const files = [{ path: 'src/a.js', content: 'x'.repeat(15000) }, { path: 'src/b.js', content: 'y'.repeat(15000) }];
|
||||
const result = checkContextBudget({ question: 'q', relevantFiles: files }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
// Should have 2 file truncation warnings (one per file over limit)
|
||||
const fileWarnings = result.warnings.filter((w) => w.includes('truncated'));
|
||||
expect(fileWarnings.length).toBe(2);
|
||||
});
|
||||
|
||||
it('adds all expected warnings when multiple trimming actions occur', () => {
|
||||
const files = Array.from({ length: 6 }, (_, i) => ({ path: `src/file${i}.js`, content: 'x' }));
|
||||
const logs = 'l'.repeat(15000);
|
||||
const context = 'c'.repeat(25000);
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q', context, relevantFiles: files, logs },
|
||||
defaultConfig,
|
||||
);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.warnings.some((w) => w.includes('Only the first'))).toBe(true);
|
||||
expect(result.warnings.some((w) => w.includes('Logs were truncated'))).toBe(true);
|
||||
expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
|
||||
});
|
||||
|
||||
it('trims context and projectSummary when both contribute to budget overflow', () => {
|
||||
const context = 'c'.repeat(28000);
|
||||
const projectSummary = 'p'.repeat(3000);
|
||||
const taskSummary = 't'.repeat(1500);
|
||||
// Total: 1 + 28000 + 3000 + 1500 = 32501 > 30000
|
||||
// Context trimmed to max(0, 30000 - 32501 + 28000) = 25499
|
||||
const result = checkContextBudget({ question: 'q', context, projectSummary, taskSummary }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.context.length).toBe(25499);
|
||||
expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
|
||||
// After: remaining = 1 + 25499 + 3000 + 1500 = 30000 — exactly at limit. No more trimming.
|
||||
expect(result.input.projectSummary.length).toBe(3000);
|
||||
expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
|
||||
expect(result.input.taskSummary).toEqual(taskSummary);
|
||||
});
|
||||
|
||||
it('trims all three droppable sections when they are all large', () => {
|
||||
const context = 'c'.repeat(20000);
|
||||
const projectSummary = 'p'.repeat(10000);
|
||||
const taskSummary = 't'.repeat(10000);
|
||||
// Total: 1 + 20000 + 10000 + 10000 = 40001 > 30000
|
||||
const result = checkContextBudget({ question: 'q', context, projectSummary, taskSummary }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
// Context trimmed to max(0, 30000 - 40001 + 20000) = 9999
|
||||
expect(result.input.context.length).toBe(9999);
|
||||
expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
|
||||
// After: remaining = 1 + 9999 + 10000 + 10000 = 30000 — exactly at limit. No more trimming.
|
||||
expect(result.input.projectSummary.length).toBe(10000);
|
||||
expect(result.input.taskSummary.length).toBe(10000);
|
||||
});
|
||||
|
||||
it('trims projectSummary when context is within budget but total still over', () => {
|
||||
const context = 'c'.repeat(5000);
|
||||
const projectSummary = 'p'.repeat(28000);
|
||||
// Total: 1 + 5000 + 28000 = 33001 > 30000
|
||||
// Context allowed = max(0, 30000 - 33001 + 5000) = 1999
|
||||
const result = checkContextBudget({ question: 'q', context, projectSummary }, defaultConfig);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(result.input.context.length).toBe(1999);
|
||||
// After: remaining = 1 + 1999 + 28000 = 30000 — exactly at limit. No more trimming.
|
||||
expect(result.input.projectSummary.length).toBe(28000);
|
||||
expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
|
||||
expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
|
||||
});
|
||||
|
||||
it('question + truncated logs exceeds budget with no droppable sections', () => {
|
||||
const longLogs = 'l'.repeat(30000);
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q'.repeat(100), logs: longLogs },
|
||||
{ ...defaultConfig, maxLogChars: 10000 },
|
||||
);
|
||||
// After truncation: 100 + 10000 = 10100 ≤ 30000 → should pass (no droppable needed)
|
||||
expect(result.ok).toBe(true);
|
||||
});
|
||||
|
||||
it('question + truncated logs exceeds budget with no other content', () => {
|
||||
const longLogs = 'l'.repeat(29950);
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q'.repeat(100), logs: longLogs },
|
||||
{ ...defaultConfig, maxLogChars: 10000, maxInputChars: 30000 },
|
||||
);
|
||||
// After truncation: 100 + 10000 = 10100 ≤ 30000 → should pass
|
||||
expect(result.ok).toBe(true);
|
||||
});
|
||||
|
||||
it('returns ok false when question + truncated logs exceeds tiny budget', () => {
|
||||
const longLogs = 'l'.repeat(29950);
|
||||
const result = checkContextBudget(
|
||||
{ question: 'q'.repeat(100), logs: longLogs },
|
||||
{ ...defaultConfig, maxLogChars: 10000, maxInputChars: 5000 },
|
||||
);
|
||||
// After truncation: 100 + 10000 = 10100 > 5000 → fails
|
||||
expect(result.ok).toBe(false);
|
||||
expect(result.input).toBeNull();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user