From b7f59a75ea7d87e71a80bbc21fe749dd454cf4b1 Mon Sep 17 00:00:00 2001 From: digua Date: Fri, 7 Aug 2026 15:22:09 +0800 Subject: [PATCH] fix(ai): bound bootstrap context compression Compress oversized browser history in model-sized stages so switching to a smaller model can recover. Bound the recent-history fallback when summarization fails without deleting persisted messages. --- .../ai-runtime/src/context-compression.ts | 137 ++++++++++++++++-- packages/ai-runtime/src/runtime.test.ts | 76 +++++++++- 2 files changed, 199 insertions(+), 14 deletions(-) diff --git a/packages/ai-runtime/src/context-compression.ts b/packages/ai-runtime/src/context-compression.ts index ae228fef..4495c9c5 100644 --- a/packages/ai-runtime/src/context-compression.ts +++ b/packages/ai-runtime/src/context-compression.ts @@ -9,6 +9,7 @@ import type { RunAgentInput, RuntimeContextSummary, RuntimeMessage } from './typ const DEFAULT_CONTEXT_WINDOW = 8_192 const COMPRESSION_THRESHOLD_PERCENT = 75 const BUFFER_SIZE_PERCENT = 20 +const COMPRESSION_SAFETY_PERCENT = 5 const MAX_SUMMARY_TOKENS = 4_096 const textEncoder = new TextEncoder() @@ -40,20 +41,32 @@ export async function prepareContextMessages( const bufferTokenBudget = Math.floor(contextWindow * (BUFFER_SIZE_PERCENT / 100)) const { bufferMessages, messagesToCompress } = splitMessages(recentMessages, bufferTokenBudget) if (messagesToCompress.length === 0) return effectiveMessages + const fallbackMessages = (): RuntimeMessage[] => + fitMessagesWithinTokenBudget(effectiveMessages, thresholdTokens, input.systemPrompt, input.userMessage) input.onEvent({ type: 'context-compression-start' }) let compressed = false try { - const history = formatCompressionHistory(summary, messagesToCompress) const maxOutputTokens = Math.max(64, Math.min(MAX_SUMMARY_TOKENS, Math.floor(contextWindow * 0.1))) - const result = await generateText({ - model: input.model.model, - prompt: COMPRESSION_PROMPT.replace('{history}', history), - maxOutputTokens, - abortSignal: input.signal, - }) - const content = result.text.trim() - if (!content) return effectiveMessages + const safetyTokens = Math.max(1, Math.floor(contextWindow * (COMPRESSION_SAFETY_PERCENT / 100))) + const promptTokenBudget = Math.max(1, contextWindow - maxOutputTokens - safetyTokens) + let content = summary?.content ?? null + let offset = 0 + + while (offset < messagesToCompress.length) { + const stage = createCompressionStage(content, messagesToCompress.slice(offset), promptTokenBudget) + if (stage.messageCount === 0) return fallbackMessages() + const result = await generateText({ + model: input.model.model, + prompt: stage.prompt, + maxOutputTokens, + abortSignal: input.signal, + }) + content = result.text.trim() + if (!content) return fallbackMessages() + offset += stage.messageCount + } + if (!content) return fallbackMessages() const boundaryMessageId = messagesToCompress.at(-1)!.id const nextSummary = await input.repository.saveContextSummary({ @@ -66,7 +79,7 @@ export async function prepareContextMessages( return [toSummaryMessage(nextSummary), ...bufferMessages] } catch (error) { if (input.signal.aborted) throw error - return effectiveMessages + return fallbackMessages() } finally { input.onEvent({ type: 'context-compression-finish', compressed }) } @@ -131,9 +144,109 @@ function estimateTextTokens(text: string): number { return Math.ceil(textEncoder.encode(text).byteLength / 3) } -function formatCompressionHistory(summary: RuntimeContextSummary | null, messages: RuntimeMessage[]): string { +function createCompressionStage( + previousSummary: string | null, + messages: RuntimeMessage[], + promptTokenBudget: number +): { prompt: string; messageCount: number } { + const maxPromptBytes = promptTokenBudget * 3 + const minimumMessageBytes = textEncoder.encode('[ASSISTANT]\n.').byteLength + const summary = fitSummaryWithinPromptBudget(previousSummary, maxPromptBytes - minimumMessageBytes) + const basePrompt = COMPRESSION_PROMPT.replace('{history}', formatCompressionHistory(summary, [])) + let usedBytes = textEncoder.encode(basePrompt).byteLength + if (usedBytes >= maxPromptBytes) return { prompt: basePrompt, messageCount: 0 } + + const fragments: string[] = [] + for (const message of messages) { + const prefix = `${fragments.length === 0 ? '' : '\n\n'}[${message.role.toUpperCase()}]\n` + const fragment = `${prefix}${message.content}` + const fragmentBytes = textEncoder.encode(fragment).byteLength + if (usedBytes + fragmentBytes <= maxPromptBytes) { + fragments.push(fragment) + usedBytes += fragmentBytes + continue + } + if (fragments.length > 0) break + + const marker = '\n[TRUNCATED TO FIT MODEL CONTEXT]' + const prefixBytes = textEncoder.encode(prefix).byteLength + const markerBytes = textEncoder.encode(marker).byteLength + const contentByteBudget = maxPromptBytes - usedBytes - prefixBytes + if (contentByteBudget <= 0) break + const canIncludeMarker = contentByteBudget > markerBytes + const boundedContent = truncateTextToByteBudget( + message.content, + contentByteBudget - (canIncludeMarker ? markerBytes : 0) + ) + fragments.push(`${prefix}${boundedContent}${canIncludeMarker ? marker : ''}`) + break + } + + return { prompt: `${basePrompt}${fragments.join('')}`, messageCount: fragments.length } +} + +function fitSummaryWithinPromptBudget(summary: string | null, maxPromptBytes: number): string | null { + if (!summary) return null + const promptBytes = (content: string): number => + textEncoder.encode(COMPRESSION_PROMPT.replace('{history}', formatCompressionHistory(content, []))).byteLength + if (promptBytes(summary) <= maxPromptBytes) return summary + + let low = 0 + let high = summary.length + while (low < high) { + const middle = Math.ceil((low + high) / 2) + if (promptBytes(summary.slice(0, middle)) <= maxPromptBytes) low = middle + else high = middle - 1 + } + if (low > 0 && /[\uD800-\uDBFF]/.test(summary[low - 1])) low -= 1 + return summary.slice(0, low) +} + +function fitMessagesWithinTokenBudget( + messages: RuntimeMessage[], + tokenBudget: number, + systemPrompt: string, + userMessage: string +): RuntimeMessage[] { + let remainingTokens = Math.max(0, tokenBudget - estimateTextTokens(systemPrompt) - estimateTextTokens(userMessage)) + const result: RuntimeMessage[] = [] + + for (let index = messages.length - 1; index >= 0; index--) { + const message = messages[index] + const messageTokens = estimateMessageTokens(message) + if (messageTokens <= remainingTokens) { + result.push(message) + remainingTokens -= messageTokens + continue + } + if (result.length === 0 && remainingTokens > 4) { + const content = truncateTextToByteBudget(message.content, (remainingTokens - 4) * 3) + if (content) result.push({ ...message, content }) + } + break + } + + return result.reverse() +} + +function truncateTextToByteBudget(text: string, maxBytes: number): string { + if (maxBytes <= 0) return '' + if (textEncoder.encode(text).byteLength <= maxBytes) return text + + let low = 0 + let high = text.length + while (low < high) { + const middle = Math.ceil((low + high) / 2) + if (textEncoder.encode(text.slice(0, middle)).byteLength <= maxBytes) low = middle + else high = middle - 1 + } + if (low > 0 && /[\uD800-\uDBFF]/.test(text[low - 1])) low -= 1 + return text.slice(0, low) +} + +function formatCompressionHistory(summary: string | null, messages: RuntimeMessage[]): string { const sections: string[] = [] - if (summary) sections.push(`[PREVIOUS SUMMARY]\n${summary.content}`) + if (summary) sections.push(`[PREVIOUS SUMMARY]\n${summary}`) sections.push( `[MESSAGES]\n${messages.map((message) => `[${message.role.toUpperCase()}]\n${message.content}`).join('\n\n')}` ) diff --git a/packages/ai-runtime/src/runtime.test.ts b/packages/ai-runtime/src/runtime.test.ts index 8fb900ab..58eb1cfb 100644 --- a/packages/ai-runtime/src/runtime.test.ts +++ b/packages/ai-runtime/src/runtime.test.ts @@ -373,6 +373,76 @@ describe('runAgent', () => { ) }) + it('recovers when existing browser history is larger than the model context window', async () => { + const maxPromptBytes = 2_048 * 3 + const promptBytes = (prompt: unknown): number => new TextEncoder().encode(JSON.stringify(prompt)).byteLength + let summaryIndex = 0 + const model = new MockLanguageModelV3({ + doGenerate: async (options) => { + if (promptBytes(options.prompt) > maxPromptBytes) { + throw new Error('context length exceeded') + } + summaryIndex += 1 + return { + content: [{ type: 'text', text: `BOUNDED SUMMARY ${summaryIndex}` }], + finishReason: { unified: 'stop', raw: 'stop' }, + usage: EMPTY_USAGE, + warnings: [], + } + }, + doStream: async (options) => { + if (promptBytes(options.prompt) > maxPromptBytes) { + throw new Error('context length exceeded') + } + return { + stream: simulateReadableStream({ + chunks: [ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 'text-1' }, + { type: 'text-delta', id: 'text-1', delta: '恢复成功' }, + { type: 'text-end', id: 'text-1' }, + { type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage: EMPTY_USAGE }, + ], + }), + } + }, + }) + const repository = new MemoryRepository() + for (let index = 0; index < 48; index++) { + await repository.appendMessage({ + conversationId: 'conversation-1', + role: index % 2 === 0 ? 'user' : 'assistant', + content: `oversized-marker-${index} ${'history '.repeat(40)}`, + createdAt: index + 1, + }) + } + const originalMessageIds = repository.messages.map((message) => message.id) + + const result = await runAgent({ + conversationId: 'conversation-1', + sessionId: 'session-1', + systemPrompt: 'You are a test assistant.', + userMessage: 'continue', + model: { model, contextWindow: 2_048 }, + repository, + tools: new FakeToolExecutor(), + signal: new AbortController().signal, + onEvent: () => undefined, + }) + + assert.equal(result.message.content, '恢复成功') + assert.ok(model.doGenerateCalls.length > 1) + assert.equal( + model.doGenerateCalls.every((call) => promptBytes(call.prompt) <= maxPromptBytes), + true + ) + assert.equal(promptBytes(model.doStreamCalls[0]?.prompt) <= maxPromptBytes, true) + assert.deepEqual( + repository.messages.slice(0, originalMessageIds.length).map((message) => message.id), + originalMessageIds + ) + }) + it('merges an existing browser summary into the next automatic compression', async () => { const model = new MockLanguageModelV3({ doGenerate: { @@ -438,7 +508,7 @@ describe('runAgent', () => { assert.doesNotMatch(modelPrompt, /PREVIOUS SUMMARY/) }) - it('continues with the original browser history when automatic compression fails', async () => { + it('continues with bounded recent browser history when automatic compression fails', async () => { const model = new MockLanguageModelV3({ doGenerate: async () => { throw new Error('compression failed') @@ -480,7 +550,9 @@ describe('runAgent', () => { assert.equal(result.message.content, 'fallback answer') assert.equal(repository.contextSummary, null) - assert.match(JSON.stringify(model.doStreamCalls[0]?.prompt), /fallback-marker-0/) + assert.doesNotMatch(JSON.stringify(model.doStreamCalls[0]?.prompt), /fallback-marker-0/) + assert.match(JSON.stringify(model.doStreamCalls[0]?.prompt), /fallback-marker-7/) + assert.equal(repository.messages.length, 9) assert.deepEqual( events.find((event) => event.type === 'context-compression-finish'), { type: 'context-compression-finish', compressed: false }