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.

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')}`
   )

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 }