From 1ae214adf6f455bfdf07b342e70601dca497d4fd Mon Sep 17 00:00:00 2001 From: Josh Lehman Date: Fri, 7 Aug 2026 23:20:43 -0700 Subject: [PATCH 1/2] lore-b74.15: retain exact OpenAI singleton reads --- src/core/ai/gateway.ts | 19 +- src/core/ai/tool-loop-context.ts | 147 +++++++++- .../gateway-tool-loop-openai-budget.test.ts | 264 ++++++++++++++++++ 3 files changed, 426 insertions(+), 4 deletions(-) create mode 100644 test/ai/gateway-tool-loop-openai-budget.test.ts diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index aaac52a25c..475e3b5f75 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -56,7 +56,11 @@ import { AIConfigError, AITransientError, normalizeAIError } from './errors.ts'; import { runGuardrails, hasGuardrails, type GuardrailHook } from '../guardrails.ts'; import { loadConfig } from '../config.ts'; import { buildGatewayConfig } from './build-gateway-config.ts'; -import { compactToolLoopMessages, resolveToolLoopMessageBudget } from './tool-loop-context.ts'; +import { + compactToolLoopMessages, + openAiToolLoopRequestFits, + resolveToolLoopMessageBudgets, +} from './tool-loop-context.ts'; import { assertProposalToolTurnPersistable } from '../minions/agent-job-proposals.ts'; // ---- Gateway-wide AI-HTTP timeout (v0.42.20.0, #1762/#1775) ---- @@ -3359,7 +3363,7 @@ export async function toolLoop(opts: ToolLoopOpts): Promise { const maxTokens = opts.maxTokens ?? defaultMaxOutputTokens(opts.model ?? getChatModel()); const handlers = opts.toolHandlers; const model = opts.model ?? getChatModel(); - const messageBudget = resolveToolLoopMessageBudget({ + const messageBudgets = resolveToolLoopMessageBudgets({ model, maxOutputTokens: maxTokens, system: opts.system, @@ -3394,12 +3398,21 @@ export async function toolLoop(opts: ToolLoopOpts): Promise { opts.onHeartbeat?.('turn_start', { turn_idx: turnIdx }); const balancedMessages = repairToolPairing(messages); - const providerMessages = compactToolLoopMessages(balancedMessages, messageBudget, { + const providerMessages = compactToolLoopMessages(balancedMessages, messageBudgets.byteSafeBytes, { mutatingToolNames: new Set( opts.tools .filter(tool => opts.toolHandlers.get(tool.name)?.mutating !== false) .map(tool => tool.name), ), + preferredProjectionBytes: messageBudgets.preferredProjectionBytes, + preferredProjectionFits: messageBudgets.openAiTokenLimits + ? candidate => openAiToolLoopRequestFits({ + budgets: messageBudgets, + system: opts.system, + tools: opts.tools, + modelMessages: toModelMessages(candidate), + }) + : undefined, }); if (providerMessages !== balancedMessages) { opts.onHeartbeat?.('context_compacted', { diff --git a/src/core/ai/tool-loop-context.ts b/src/core/ai/tool-loop-context.ts index de4c78b7a3..f673a042de 100644 --- a/src/core/ai/tool-loop-context.ts +++ b/src/core/ai/tool-loop-context.ts @@ -86,6 +86,27 @@ function isOpenAiModel(model: string): boolean { export interface ToolLoopContextOptions { /** Tools whose effects must retain a distinct identity whenever context compacts. */ mutatingToolNames?: ReadonlySet; + /** + * Larger candidate ceiling for providers with an exact request-token check. + * This is never an acceptance boundary by itself. + */ + preferredProjectionBytes?: number; + /** Provider-owned proof that the complete preferred request fits safely. */ + preferredProjectionFits?: (candidate: ChatMessage[]) => boolean; +} + +/** Provider-specific limits used to construct and verify tool-loop projections. */ +export interface ToolLoopMessageBudgets { + /** Worst-case UTF-8 limit that is safe even at one token per byte. */ + byteSafeBytes: number; + /** Candidate ceiling whose result still requires exact provider validation. */ + preferredProjectionBytes: number; + /** OpenAI total-token limits used to verify a transformed provider request. */ + openAiTokenLimits: { + targetTotalTokens: number; + hardTotalTokens: number; + maxOutputTokens: number; + } | null; } /** Raised when no valid, evidence-preserving provider projection can fit. */ @@ -104,6 +125,17 @@ export function resolveToolLoopMessageBudget(args: { tools: ChatToolDef[]; contextWindowTokens?: number; }): number { + return resolveToolLoopMessageBudgets(args).byteSafeBytes; +} + +/** Resolve preferred and worst-case-safe provider message budgets. */ +export function resolveToolLoopMessageBudgets(args: { + model: string; + maxOutputTokens: number; + system?: string; + tools: ChatToolDef[]; + contextWindowTokens?: number; +}): ToolLoopMessageBudgets { let contextTokens = args.contextWindowTokens; if (contextTokens === undefined) { try { @@ -139,7 +171,43 @@ export function resolveToolLoopMessageBudget(args: { 0, hardInputTokens - openAiStaticTokens - OPENAI_PROTOCOL_TOKEN_RESERVE, ) * CONSERVATIVE_BYTES_PER_TOKEN; - return Math.max(0, Math.min(targetBudget, byteSafeBudget)); + return { + byteSafeBytes: Math.max(0, Math.min(targetBudget, byteSafeBudget)), + preferredProjectionBytes: openAiStaticTokens === null + ? Math.max(0, Math.min(targetBudget, byteSafeBudget)) + : Math.max(0, targetBudget), + openAiTokenLimits: openAiStaticTokens === null + ? null + : { + targetTotalTokens: targetTokens, + hardTotalTokens: contextTokens, + maxOutputTokens: args.maxOutputTokens, + }, + }; +} + +/** Verify the complete transformed OpenAI request against both token limits. */ +export function openAiToolLoopRequestFits(args: { + budgets: ToolLoopMessageBudgets, + system?: string; + tools: ChatToolDef[]; + modelMessages: unknown[]; +}): boolean { + const limits = args.budgets.openAiTokenLimits; + if (!limits) return false; + // Count the provider-facing message shape, not gbrain's durable ChatMessage + // shape. The enclosing object includes the request's JSON keys and + // separators; the fixed reserve covers SDK/provider framing not represented + // here. + const requestTokens = countOpenAiTokens(safeJson({ + system: args.system ?? '', + tools: args.tools, + messages: args.modelMessages, + })) + + OPENAI_PROTOCOL_TOKEN_RESERVE + + limits.maxOutputTokens; + return requestTokens <= limits.targetTotalTokens + && requestTokens <= limits.hardTotalTokens; } /** @@ -150,6 +218,83 @@ export function compactToolLoopMessages( messages: ChatMessage[], maxBytes: number, options: ToolLoopContextOptions = {}, +): ChatMessage[] { + const fallback = compactToolLoopMessagesToByteBudget(messages, maxBytes, options); + const preferredBytes = options.preferredProjectionBytes; + if ( + preferredBytes === undefined + || preferredBytes <= maxBytes + || !options.preferredProjectionFits + ) return fallback; + + const preferred = restorePreferredNewestSingletonResult( + messages, + fallback, + options, + ); + if (!preferred || jsonBytes(preferred) > preferredBytes) return fallback; + try { + return options.preferredProjectionFits(preferred) ? preferred : fallback; + } catch { + // Exact provider tokenization is optional. Any failure retains the + // independently valid worst-case byte projection. + return fallback; + } +} + +/** + * Restore only the newest singleton read result onto an already-safe fallback. + * Tool identity comes from the paired assistant call, never the result label. + */ +function restorePreferredNewestSingletonResult( + originalMessages: ChatMessage[], + fallback: ChatMessage[], + options: ToolLoopContextOptions, +): ChatMessage[] | null { + const originalRound = collectToolRounds(originalMessages).rounds.at(-1); + if (!originalRound || originalRound.evidence.length !== 1) return null; + + const evidence = originalRound.evidence[0]!; + if ( + evidence.failed + || isMutationSensitive(evidence.toolName, options.mutatingToolNames) + ) return null; + + const originalResult = toolResultBlocks(originalRound.result).find(block => ( + block.toolCallId === evidence.toolCallId + )); + if (!originalResult || originalResult.toolName !== evidence.toolName) return null; + + // The byte fallback must still retain this same newest balanced round. Do + // not inject an exact result into a summary or an unrelated historical call. + const fallbackRound = collectToolRounds(fallback).rounds.at(-1); + if ( + !fallbackRound + || fallbackRound.evidence.length !== 1 + || fallbackRound.evidence[0]?.toolCallId !== evidence.toolCallId + || fallbackRound.evidence[0]?.toolName !== evidence.toolName + ) return null; + + const fallbackResultIndex = fallback.indexOf(fallbackRound.result); + if (fallbackResultIndex < 0) return null; + const restoredResult: ChatMessage = { + ...fallbackRound.result, + content: mapBlocks(fallbackRound.result, block => ( + block.type === 'tool-result' && block.toolCallId === evidence.toolCallId + ? { ...block, output: originalResult.output } + : block + )), + }; + const preferred = [...fallback]; + preferred[fallbackResultIndex] = restoredResult; + return preferred; +} + +/** Build one evidence-preserving projection under a UTF-8 byte ceiling. */ +function compactToolLoopMessagesToByteBudget( + messages: ChatMessage[], + maxBytes: number, + options: ToolLoopContextOptions, ): ChatMessage[] { if (jsonBytes(messages) <= maxBytes) return messages; diff --git a/test/ai/gateway-tool-loop-openai-budget.test.ts b/test/ai/gateway-tool-loop-openai-budget.test.ts new file mode 100644 index 0000000000..1cf821c58f --- /dev/null +++ b/test/ai/gateway-tool-loop-openai-budget.test.ts @@ -0,0 +1,264 @@ +import { afterEach, describe, expect, it } from 'bun:test'; +import { + __setChatTransportForTests, + configureGateway, + resetGateway, + toModelMessages, + toolLoop, +} from '../../src/core/ai/gateway.ts'; +import type { + ChatBlock, + ChatMessage, + ChatToolDef, +} from '../../src/core/ai/gateway.ts'; +import { + compactToolLoopMessages, + openAiToolLoopRequestFits, + resolveToolLoopMessageBudgets, +} from '../../src/core/ai/tool-loop-context.ts'; + +/** Find a persisted tool result in the provider-facing prompt. */ +function resultOutput(messages: ChatMessage[], toolCallId: string): unknown { + for (const message of messages) { + if (typeof message.content === 'string') continue; + const result = message.content.find(block => ( + block.type === 'tool-result' && block.toolCallId === toolCallId + )); + if (result?.type === 'tool-result') return result.output; + } + throw new Error(`Missing tool result ${toolCallId}`); +} + +describe('OpenAI tool-loop context budgeting', () => { + afterEach(() => { + __setChatTransportForTests(null); + resetGateway(); + }); + + it('retains a production-sized latest read when exact OpenAI tokens fit', async () => { + configureGateway({ + chat_model: 'openai:gpt-5.6-terra', + embedding_model: 'openai:text-embedding-3-large', + embedding_dimensions: 1536, + expansion_model: 'openai:gpt-5.6-luna', + env: { OPENAI_API_KEY: 'stub' }, + }); + const system = 'Published ingestion instructions.\n'.repeat(1_200); + const tools: ChatToolDef[] = [{ + name: 'brain_get_page', + description: 'Read one canonical brain page. '.repeat(24), + inputSchema: { type: 'object', properties: { slug: { type: 'string' } } }, + }]; + const exactOutput = { + id: 6_537, + slug: 'companies/signalcore', + source_id: 'martian', + content_hash: 'a'.repeat(64), + timeline: 't'.repeat(52_663), + compiled_truth: 'b'.repeat(9_091), + }; + expect(Buffer.byteLength(JSON.stringify(exactOutput), 'utf8')).toBe(61_933); + const durableMessages: ChatMessage[] = [ + { role: 'user', content: 'p'.repeat(116_659) }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'canonical-baseline', + toolName: 'brain_get_page', + input: { slug: 'companies/signalcore', source_id: 'martian' }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'canonical-baseline', + toolName: 'brain_get_page', + output: exactOutput, + }], + }, + ]; + const budgets = resolveToolLoopMessageBudgets({ + model: 'openai:gpt-5.6-terra', + maxOutputTokens: 32_768, + contextWindowTokens: 200_000, + system, + tools, + }); + expect(Buffer.byteLength(JSON.stringify(durableMessages), 'utf8')) + .toBeGreaterThan(budgets.byteSafeBytes); + + let providerMessages: ChatMessage[] = []; + __setChatTransportForTests(async options => { + providerMessages = structuredClone(options.messages); + return { + text: 'safe to stage', + blocks: [{ type: 'text', text: 'safe to stage' }] as ChatBlock[], + stopReason: 'end', + usage: { + input_tokens: 75_000, + output_tokens: 4, + cache_read_tokens: 0, + cache_creation_tokens: 0, + }, + model: 'openai:gpt-5.6-terra', + providerId: 'openai', + }; + }); + + await toolLoop({ + model: 'openai:gpt-5.6-terra', + system, + initialMessages: [], + tools, + toolHandlers: new Map([['brain_get_page', { + idempotent: true, + mutating: false, + async execute() { return null; }, + }]]), + maxTokens: 32_768, + contextWindowTokens: 200_000, + replayState: { + priorMessages: durableMessages, + priorTools: new Map(), + nextTurnIdx: 1, + nextMessageIdx: 3, + }, + }); + + expect(resultOutput(providerMessages, 'canonical-baseline')).toEqual(exactOutput); + expect(Buffer.byteLength(JSON.stringify(providerMessages), 'utf8')) + .toBeLessThanOrEqual(budgets.preferredProjectionBytes); + expect(openAiToolLoopRequestFits({ + budgets, + system, + tools, + modelMessages: toModelMessages(providerMessages), + })).toBe(true); + expect(durableMessages[2]?.content).toEqual([expect.objectContaining({ output: exactOutput })]); + }, 15_000); + + it('rejects a preferred dense-Unicode result when exact tokens exceed the target', () => { + const model = 'openai:gpt-5.6-terra'; + const system = 'Keep exact source text only when it fits.'; + const tools: ChatToolDef[] = [{ + name: 'brain_get_page', + description: 'Read one page.', + inputSchema: { type: 'object', properties: { slug: { type: 'string' } } }, + }]; + const budgets = resolveToolLoopMessageBudgets({ + model, + maxOutputTokens: 1_000, + contextWindowTokens: 10_000, + system, + tools, + }); + const exactOutput = { body: '🧠'.repeat(2_500) }; + const messages: ChatMessage[] = [ + { role: 'user', content: 'Read the latest page.' }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'unicode-read', + toolName: 'brain_get_page', + input: { slug: 'notes/unicode' }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'unicode-read', + toolName: 'brain_get_page', + output: exactOutput, + }], + }, + ]; + const durableSnapshot = structuredClone(messages); + + const compacted = compactToolLoopMessages(messages, budgets.byteSafeBytes, { + mutatingToolNames: new Set(), + preferredProjectionBytes: budgets.preferredProjectionBytes, + preferredProjectionFits: candidate => openAiToolLoopRequestFits({ + budgets, + system, + tools, + modelMessages: toModelMessages(candidate), + }), + }); + + expect(resultOutput(compacted, 'unicode-read')).not.toEqual(exactOutput); + expect(JSON.stringify(compacted)).toContain('working_context_projection'); + expect(Buffer.byteLength(JSON.stringify(compacted), 'utf8')) + .toBeLessThanOrEqual(budgets.byteSafeBytes); + expect(messages).toEqual(durableSnapshot); + }); + + it('never restores mutation, mismatched-name, or failed singleton outputs', () => { + const scenarios = [ + { + id: 'mutation', + callName: 'put_page', + resultName: 'put_page', + isError: false, + mutating: new Set(['put_page']), + }, + { + id: 'mismatch', + callName: 'brain_get_page', + resultName: 'query', + isError: false, + mutating: new Set(), + }, + { + id: 'failed', + callName: 'brain_get_page', + resultName: 'brain_get_page', + isError: true, + mutating: new Set(), + }, + ]; + + for (const scenario of scenarios) { + const exactOutput = { body: `FORBIDDEN_${scenario.id}_${'x'.repeat(4_000)}` }; + const messages: ChatMessage[] = [ + { role: 'user', content: 'Continue safely.' }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: scenario.id, + toolName: scenario.callName, + input: { slug: `notes/${scenario.id}`, content: 'private'.repeat(500) }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: scenario.id, + toolName: scenario.resultName, + output: exactOutput, + ...(scenario.isError ? { isError: true } : {}), + }], + }, + ]; + let preferredChecks = 0; + + const compacted = compactToolLoopMessages(messages, 1_000, { + mutatingToolNames: scenario.mutating, + preferredProjectionBytes: 20_000, + preferredProjectionFits: () => { + preferredChecks++; + return true; + }, + }); + + expect(resultOutput(compacted, scenario.id)).not.toEqual(exactOutput); + expect(JSON.stringify(compacted)).not.toContain(`FORBIDDEN_${scenario.id}_`); + expect(preferredChecks).toBe(0); + } + }); +}); From 19ee2590baef01ad44136c1fc448f8a4538df3af Mon Sep 17 00:00:00 2001 From: Josh Lehman Date: Fri, 7 Aug 2026 23:41:19 -0700 Subject: [PATCH 2/2] fix(lore-b74.15): preserve exact latest read context --- src/core/ai/tool-loop-context.ts | 72 +++--- .../gateway-tool-loop-openai-budget.test.ts | 215 +++++++++++++++++- 2 files changed, 239 insertions(+), 48 deletions(-) diff --git a/src/core/ai/tool-loop-context.ts b/src/core/ai/tool-loop-context.ts index f673a042de..e58b74dd7f 100644 --- a/src/core/ai/tool-loop-context.ts +++ b/src/core/ai/tool-loop-context.ts @@ -220,6 +220,7 @@ export function compactToolLoopMessages( options: ToolLoopContextOptions = {}, ): ChatMessage[] { const fallback = compactToolLoopMessagesToByteBudget(messages, maxBytes, options); + if (fallback === messages) return fallback; const preferredBytes = options.preferredProjectionBytes; if ( preferredBytes === undefined @@ -227,14 +228,13 @@ export function compactToolLoopMessages( || !options.preferredProjectionFits ) return fallback; - const preferred = restorePreferredNewestSingletonResult( - messages, - fallback, - options, - ); - if (!preferred || jsonBytes(preferred) > preferredBytes) return fallback; try { - return options.preferredProjectionFits(preferred) ? preferred : fallback; + return buildPreferredNewestSingletonProjection( + messages, + preferredBytes, + options, + options.preferredProjectionFits, + ) ?? fallback; } catch { // Exact provider tokenization is optional. Any failure retains the // independently valid worst-case byte projection. @@ -243,15 +243,17 @@ export function compactToolLoopMessages( } /** - * Restore only the newest singleton read result onto an already-safe fallback. - * Tool identity comes from the paired assistant call, never the result label. + * Build a minimal preferred projection around the newest singleton read. + * Older rounds become durable ledger evidence instead of consuming headroom. */ -function restorePreferredNewestSingletonResult( - originalMessages: ChatMessage[], - fallback: ChatMessage[], +function buildPreferredNewestSingletonProjection( + messages: ChatMessage[], + preferredBytes: number, options: ToolLoopContextOptions, + fits: (candidate: ChatMessage[]) => boolean, ): ChatMessage[] | null { - const originalRound = collectToolRounds(originalMessages).rounds.at(-1); + const { rounds, otherCount } = collectToolRounds(messages); + const originalRound = rounds.at(-1); if (!originalRound || originalRound.evidence.length !== 1) return null; const evidence = originalRound.evidence[0]!; @@ -265,29 +267,27 @@ function restorePreferredNewestSingletonResult( )); if (!originalResult || originalResult.toolName !== evidence.toolName) return null; - // The byte fallback must still retain this same newest balanced round. Do - // not inject an exact result into a summary or an unrelated historical call. - const fallbackRound = collectToolRounds(fallback).rounds.at(-1); - if ( - !fallbackRound - || fallbackRound.evidence.length !== 1 - || fallbackRound.evidence[0]?.toolCallId !== evidence.toolCallId - || fallbackRound.evidence[0]?.toolName !== evidence.toolName - ) return null; - - const fallbackResultIndex = fallback.indexOf(fallbackRound.result); - if (fallbackResultIndex < 0) return null; - const restoredResult: ChatMessage = { - ...fallbackRound.result, - content: mapBlocks(fallbackRound.result, block => ( - block.type === 'tool-result' && block.toolCallId === evidence.toolCallId - ? { ...block, output: originalResult.output } - : block - )), - }; - const preferred = [...fallback]; - preferred[fallbackResultIndex] = restoredResult; - return preferred; + const task = buildTaskAnchor(messages); + const summary = buildLedgerSummary(rounds.slice(0, -1), otherCount, options); + for (const perPayload of PAYLOAD_LIMITS) { + const compacted = compactRound(originalRound, perPayload, options); + const exactResult: ChatMessage = { + ...compacted.result, + content: mapBlocks(compacted.result, block => ( + block.type === 'tool-result' && block.toolCallId === evidence.toolCallId + ? { ...block, output: originalResult.output } + : block + )), + }; + const preferred = [ + task, + ...(summary ? [summary] : []), + compacted.assistant, + exactResult, + ]; + if (jsonBytes(preferred) <= preferredBytes && fits(preferred)) return preferred; + } + return null; } /** Build one evidence-preserving projection under a UTF-8 byte ceiling. */ diff --git a/test/ai/gateway-tool-loop-openai-budget.test.ts b/test/ai/gateway-tool-loop-openai-budget.test.ts index 1cf821c58f..91a8e70bd3 100644 --- a/test/ai/gateway-tool-loop-openai-budget.test.ts +++ b/test/ai/gateway-tool-loop-openai-budget.test.ts @@ -44,11 +44,23 @@ describe('OpenAI tool-loop context budgeting', () => { env: { OPENAI_API_KEY: 'stub' }, }); const system = 'Published ingestion instructions.\n'.repeat(1_200); - const tools: ChatToolDef[] = [{ - name: 'brain_get_page', - description: 'Read one canonical brain page. '.repeat(24), - inputSchema: { type: 'object', properties: { slug: { type: 'string' } } }, - }]; + const tools: ChatToolDef[] = [ + { + name: 'brain_search', + description: 'Search the brain for relevant evidence. '.repeat(12), + inputSchema: { type: 'object', properties: { query: { type: 'string' } } }, + }, + { + name: 'brain_stage_ingestion_proposal_page', + description: 'Stage one reviewed ingestion proposal page. '.repeat(12), + inputSchema: { type: 'object', properties: { slug: { type: 'string' } } }, + }, + { + name: 'brain_get_page', + description: 'Read one canonical brain page. '.repeat(24), + inputSchema: { type: 'object', properties: { slug: { type: 'string' } } }, + }, + ]; const exactOutput = { id: 6_537, slug: 'companies/signalcore', @@ -58,8 +70,66 @@ describe('OpenAI tool-loop context budgeting', () => { compiled_truth: 'b'.repeat(9_091), }; expect(Buffer.byteLength(JSON.stringify(exactOutput), 'utf8')).toBe(61_933); + const oldSearchBody = `OLD_SEARCH_RAW_${'s'.repeat(44_000)}`; + const oldStageBody = `OLD_STAGE_RAW_${'m'.repeat(8_000)}`; const durableMessages: ChatMessage[] = [ { role: 'user', content: 'p'.repeat(116_659) }, + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'search-company', + toolName: 'brain_search', + input: { query: 'SignalCore company baseline' }, + }, + { + type: 'tool-call', + toolCallId: 'search-people', + toolName: 'brain_search', + input: { query: 'SignalCore meeting attendees' }, + }, + ], + }, + { + role: 'user', + content: [ + { + type: 'tool-result', + toolCallId: 'search-company', + toolName: 'brain_search', + output: { matches: oldSearchBody }, + }, + { + type: 'tool-result', + toolCallId: 'search-people', + toolName: 'brain_search', + output: { matches: oldSearchBody }, + }, + ], + }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'stage-page-fragment', + toolName: 'brain_stage_ingestion_proposal_page', + input: { + slug: 'companies/signalcore', + source_id: 'martian', + content: oldStageBody, + }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'stage-page-fragment', + toolName: 'brain_stage_ingestion_proposal_page', + output: { staged: true, fragment: 1, total_fragments: 13 }, + }], + }, { role: 'assistant', content: [{ @@ -79,6 +149,8 @@ describe('OpenAI tool-loop context budgeting', () => { }], }, ]; + const durableSnapshot = structuredClone(durableMessages); + const mutatingToolNames = new Set(['brain_stage_ingestion_proposal_page']); const budgets = resolveToolLoopMessageBudgets({ model: 'openai:gpt-5.6-terra', maxOutputTokens: 32_768, @@ -88,6 +160,14 @@ describe('OpenAI tool-loop context budgeting', () => { }); expect(Buffer.byteLength(JSON.stringify(durableMessages), 'utf8')) .toBeGreaterThan(budgets.byteSafeBytes); + const byteSafeOnly = compactToolLoopMessages( + durableMessages, + budgets.byteSafeBytes, + { mutatingToolNames }, + ); + expect(resultOutput(byteSafeOnly, 'canonical-baseline')).not.toEqual(exactOutput); + expect(JSON.stringify(resultOutput(byteSafeOnly, 'canonical-baseline'))) + .toContain('working_context_projection'); let providerMessages: ChatMessage[] = []; __setChatTransportForTests(async options => { @@ -112,22 +192,39 @@ describe('OpenAI tool-loop context budgeting', () => { system, initialMessages: [], tools, - toolHandlers: new Map([['brain_get_page', { - idempotent: true, - mutating: false, - async execute() { return null; }, - }]]), + toolHandlers: new Map([ + ['brain_search', { + idempotent: true, + mutating: false, + async execute() { return null; }, + }], + ['brain_stage_ingestion_proposal_page', { + idempotent: false, + mutating: true, + async execute() { return null; }, + }], + ['brain_get_page', { + idempotent: true, + mutating: false, + async execute() { return null; }, + }], + ]), maxTokens: 32_768, contextWindowTokens: 200_000, replayState: { priorMessages: durableMessages, priorTools: new Map(), nextTurnIdx: 1, - nextMessageIdx: 3, + nextMessageIdx: durableMessages.length, }, }); expect(resultOutput(providerMessages, 'canonical-baseline')).toEqual(exactOutput); + expect(providerMessages).toHaveLength(4); + expect(JSON.stringify(providerMessages)).toContain('Distinct mutation evidence'); + expect(JSON.stringify(providerMessages)).toContain('call_id=stage-page-fragment'); + expect(JSON.stringify(providerMessages)).not.toContain('OLD_SEARCH_RAW_'); + expect(JSON.stringify(providerMessages)).not.toContain('OLD_STAGE_RAW_'); expect(Buffer.byteLength(JSON.stringify(providerMessages), 'utf8')) .toBeLessThanOrEqual(budgets.preferredProjectionBytes); expect(openAiToolLoopRequestFits({ @@ -136,7 +233,7 @@ describe('OpenAI tool-loop context budgeting', () => { tools, modelMessages: toModelMessages(providerMessages), })).toBe(true); - expect(durableMessages[2]?.content).toEqual([expect.objectContaining({ output: exactOutput })]); + expect(durableMessages).toEqual(durableSnapshot); }, 15_000); it('rejects a preferred dense-Unicode result when exact tokens exceed the target', () => { @@ -196,6 +293,100 @@ describe('OpenAI tool-loop context budgeting', () => { expect(messages).toEqual(durableSnapshot); }); + it('tries smaller input projections until the exact-result candidate fits', () => { + const exactOutput = { body: `EXACT_RESULT_${'r'.repeat(4_000)}` }; + const messages: ChatMessage[] = [ + { role: 'user', content: 'Read the page and continue.' }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'large-input-read', + toolName: 'brain_get_page', + input: { slug: 'notes/large-input', filter: '🧠'.repeat(2_500) }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'large-input-read', + toolName: 'brain_get_page', + output: exactOutput, + }], + }, + ]; + let preferredChecks = 0; + + const compacted = compactToolLoopMessages(messages, 1_000, { + mutatingToolNames: new Set(), + preferredProjectionBytes: 20_000, + preferredProjectionFits: candidate => { + preferredChecks++; + return JSON.stringify(candidate).includes('working_context_projection'); + }, + }); + + expect(preferredChecks).toBe(2); + expect(resultOutput(compacted, 'large-input-read')).toEqual(exactOutput); + expect(JSON.stringify(compacted)).toContain('working_context_projection'); + }); + + it('leaves an already byte-safe transcript completely unchanged', () => { + const messages: ChatMessage[] = [ + { role: 'user', content: 'Compare both pages.' }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'first-read', + toolName: 'brain_get_page', + input: { slug: 'notes/first' }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'first-read', + toolName: 'brain_get_page', + output: { body: 'first exact result' }, + }], + }, + { + role: 'assistant', + content: [{ + type: 'tool-call', + toolCallId: 'second-read', + toolName: 'brain_get_page', + input: { slug: 'notes/second' }, + }], + }, + { + role: 'user', + content: [{ + type: 'tool-result', + toolCallId: 'second-read', + toolName: 'brain_get_page', + output: { body: 'second exact result' }, + }], + }, + ]; + let preferredChecks = 0; + + const compacted = compactToolLoopMessages(messages, 10_000, { + mutatingToolNames: new Set(), + preferredProjectionBytes: 20_000, + preferredProjectionFits: () => { + preferredChecks++; + return true; + }, + }); + + expect(compacted).toBe(messages); + expect(preferredChecks).toBe(0); + }); + it('never restores mutation, mismatched-name, or failed singleton outputs', () => { const scenarios = [ {