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..e58b74dd7f 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); + if (fallback === messages) return fallback; + const preferredBytes = options.preferredProjectionBytes; + if ( + preferredBytes === undefined + || preferredBytes <= maxBytes + || !options.preferredProjectionFits + ) return fallback; + + try { + return buildPreferredNewestSingletonProjection( + messages, + preferredBytes, + options, + options.preferredProjectionFits, + ) ?? fallback; + } catch { + // Exact provider tokenization is optional. Any failure retains the + // independently valid worst-case byte projection. + return fallback; + } +} + +/** + * Build a minimal preferred projection around the newest singleton read. + * Older rounds become durable ledger evidence instead of consuming headroom. + */ +function buildPreferredNewestSingletonProjection( + messages: ChatMessage[], + preferredBytes: number, + options: ToolLoopContextOptions, + fits: (candidate: ChatMessage[]) => boolean, +): ChatMessage[] | null { + const { rounds, otherCount } = collectToolRounds(messages); + const originalRound = 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; + + 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. */ +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..91a8e70bd3 --- /dev/null +++ b/test/ai/gateway-tool-loop-openai-budget.test.ts @@ -0,0 +1,455 @@ +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_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', + 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 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: [{ + 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 durableSnapshot = structuredClone(durableMessages); + const mutatingToolNames = new Set(['brain_stage_ingestion_proposal_page']); + 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); + 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 => { + 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_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: 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({ + budgets, + system, + tools, + modelMessages: toModelMessages(providerMessages), + })).toBe(true); + expect(durableMessages).toEqual(durableSnapshot); + }, 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('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 = [ + { + 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); + } + }); +});