fix(ai): count replayed tool results in compression token accounting

packages/node-runtime/src/ai/agent/history.ts

@@ -16,7 +16,7 @@ import type { SimpleHistoryMessage } from './types'
 import type { ContentBlock } from '../chats'
 
 type ToolBlock = Extract<ContentBlock, { type: 'tool' }>
-type ReplayableToolBlock = ToolBlock & { tool: ToolBlock['tool'] & { toolCallId: string; result: string } }
+export type ReplayableToolBlock = ToolBlock & { tool: ToolBlock['tool'] & { toolCallId: string; result: string } }
 
 function createEmptyPiUsage(): PiUsage {
   return {
@@ -29,7 +29,7 @@ function createEmptyPiUsage(): PiUsage {
   }
 }
 
-function isReplayableToolBlock(block: ContentBlock): block is ReplayableToolBlock {
+export function isReplayableToolBlock(block: ContentBlock): block is ReplayableToolBlock {
   if (block.type !== 'tool') return false
   const { toolCallId, result, status } = block.tool
   return (

packages/node-runtime/src/ai/chats.ts

@@ -944,16 +944,32 @@ export class AIChatManager {
   getMessagesAfterSummary(
     aiChatId: string,
     summaryTimestamp: number
-  ): Array<{ role: AIMessageRole; content: string; timestamp: number }> {
+  ): Array<{ role: AIMessageRole; content: string; timestamp: number; contentBlocks?: ContentBlock[] }> {
     return this.getActivePathRows(aiChatId)
       .filter((row) => row.timestamp > summaryTimestamp && (row.role === 'user' || row.role === 'assistant'))
-      .map((row) => ({ role: row.role as AIMessageRole, content: row.content, timestamp: row.timestamp }))
+      .map((row) => this.toCompressionMessage(row))
   }
 
-  getAllUserAssistantMessages(aiChatId: string): Array<{ role: AIMessageRole; content: string; timestamp: number }> {
+  getAllUserAssistantMessages(
+    aiChatId: string
+  ): Array<{ role: AIMessageRole; content: string; timestamp: number; contentBlocks?: ContentBlock[] }> {
     return this.getActivePathRows(aiChatId)
       .filter((row) => row.role === 'user' || row.role === 'assistant')
-      .map((row) => ({ role: row.role as AIMessageRole, content: row.content, timestamp: row.timestamp }))
+      .map((row) => this.toCompressionMessage(row))
+  }
+
+  private toCompressionMessage(row: AIMessageRow): {
+    role: AIMessageRole
+    content: string
+    timestamp: number
+    contentBlocks?: ContentBlock[]
+  } {
+    return {
+      role: row.role as AIMessageRole,
+      content: row.content,
+      timestamp: row.timestamp,
+      contentBlocks: row.contentBlocks ? JSON.parse(row.contentBlocks) : undefined,
+    }
   }
 
   getMessageCountAfterSummary(aiChatId: string): number {

packages/node-runtime/src/ai/compression/__tests__/compressor.test.ts

@@ -0,0 +1,159 @@
+import { describe, it } from 'node:test'
+import assert from 'node:assert/strict'
+import { mkdtempSync, rmSync } from 'node:fs'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { AIChatManager } from '../../chats'
+import type { ContentBlock } from '../../chats'
+import { checkAndCompress } from '../compressor'
+import type { CompressionConfig, CompressionLlmAdapter } from '../types'
+
+const sqliteNativeBinding = process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING
+
+function createTempDir(): string {
+  return mkdtempSync(join(tmpdir(), 'chatlab-compression-'))
+}
+
+function createManager(dir: string): AIChatManager {
+  return sqliteNativeBinding ? new AIChatManager(dir, { nativeBinding: sqliteNativeBinding }) : new AIChatManager(dir)
+}
+
+function cleanup(dir: string): void {
+  try {
+    rmSync(dir, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 })
+  } catch {
+    // Windows can hold SQLite WAL handles briefly after close; temp cleanup is best-effort.
+  }
+}
+
+const CONFIG: CompressionConfig = {
+  enabled: true,
+  tokenThresholdPercent: 75,
+  bufferSizePercent: 20,
+}
+
+function createAdapter(captured: { prompt: string | null }): CompressionLlmAdapter {
+  return {
+    contextWindow: 1000,
+    compress: async (prompt: string) => {
+      captured.prompt = prompt
+      return 'COMPRESSED SUMMARY'
+    },
+  }
+}
+
+/** A tool result large enough that replaying it dominates the context window. */
+function bigToolResult(marker: string): string {
+  return `${marker} ` + Array.from({ length: 700 }, (_, i) => `record${i} value${i * 3}`).join(' ')
+}
+
+function toolBlock(name: string, result: string | undefined): ContentBlock {
+  return {
+    type: 'tool',
+    tool: {
+      name,
+      displayName: name,
+      status: 'done',
+      params: { query: 'stats' },
+      ...(result !== undefined ? { toolCallId: `call_${name}`, result } : {}),
+    },
+  }
+}
+
+/**
+ * Seeds a conversation whose plain message text is tiny, but whose persisted
+ * tool results (replayed into the LLM context each turn) are large.
+ */
+function seedToolHeavyChat(manager: AIChatManager, withToolResults: boolean): string {
+  const chat = manager.createAIChat('session-1', 'Compression', 'general_cn')
+  const result1 = withToolResults ? bigToolResult('TOOL_DATA_ALPHA') : undefined
+  const result2 = withToolResults ? bigToolResult('TOOL_DATA_BETA') : undefined
+
+  manager.addMessage(chat.id, 'user', 'show me stats')
+  manager.addMessage(chat.id, 'assistant', 'here are the stats', undefined, undefined, [toolBlock('query_a', result1)])
+  manager.addMessage(chat.id, 'user', 'and more details')
+  manager.addMessage(chat.id, 'assistant', 'more stats', undefined, undefined, [toolBlock('query_b', result2)])
+  manager.addMessage(chat.id, 'user', 'thanks')
+  manager.addMessage(chat.id, 'assistant', 'you are welcome')
+  return chat.id
+}
+
+describe('checkAndCompress tool result token accounting', () => {
+  it('counts replayed tool results toward the compression threshold', async () => {
+    const dir = createTempDir()
+    const manager = createManager(dir)
+    try {
+      const chatId = seedToolHeavyChat(manager, true)
+      const captured: { prompt: string | null } = { prompt: null }
+
+      const result = await checkAndCompress(chatId, CONFIG, 'system', createAdapter(captured), manager)
+
+      assert.equal(result.compressed, true)
+      assert.equal(result.reason, 'success')
+      assert.ok(result.tokensBefore! > result.tokensAfter!)
+      assert.ok(manager.getLatestSummary(chatId), 'summary message should be persisted')
+
+      // Tool results must be part of the compression input so the summary
+      // can preserve their key data points.
+      assert.ok(captured.prompt!.includes('TOOL_DATA_ALPHA'))
+      assert.ok(captured.prompt!.includes('TOOL_DATA_BETA'))
+      assert.ok(captured.prompt!.includes('[Tool result: query_a]'))
+    } finally {
+      manager.close()
+      cleanup(dir)
+    }
+  })
+
+  it('does not count tool blocks without a persisted result (not replayed)', async () => {
+    const dir = createTempDir()
+    const manager = createManager(dir)
+    try {
+      const chatId = seedToolHeavyChat(manager, false)
+      const captured: { prompt: string | null } = { prompt: null }
+
+      const result = await checkAndCompress(chatId, CONFIG, 'system', createAdapter(captured), manager)
+
+      assert.equal(result.compressed, false)
+      assert.equal(result.reason, 'skipped_below_threshold')
+      assert.equal(captured.prompt, null)
+      assert.equal(manager.getLatestSummary(chatId), null)
+    } finally {
+      manager.close()
+      cleanup(dir)
+    }
+  })
+
+  it('counts tool results on the progressive path (after an existing summary)', async () => {
+    const dir = createTempDir()
+    const manager = createManager(dir)
+    try {
+      const chat = manager.createAIChat('session-1', 'Compression', 'general_cn')
+      manager.addSummaryMessage(chat.id, 'old summary of earlier topics', {
+        bufferBoundaryTimestamp: Math.floor(Date.now() / 1000) - 100,
+        compressedMessageCount: 10,
+      })
+      manager.addMessage(chat.id, 'user', 'show me stats')
+      manager.addMessage(chat.id, 'assistant', 'here are the stats', undefined, undefined, [
+        toolBlock('query_a', bigToolResult('TOOL_DATA_GAMMA')),
+      ])
+      manager.addMessage(chat.id, 'user', 'and more details')
+      manager.addMessage(chat.id, 'assistant', 'more stats', undefined, undefined, [
+        toolBlock('query_b', bigToolResult('TOOL_DATA_DELTA')),
+      ])
+      manager.addMessage(chat.id, 'user', 'thanks')
+      manager.addMessage(chat.id, 'assistant', 'you are welcome')
+
+      const captured: { prompt: string | null } = { prompt: null }
+      const result = await checkAndCompress(chat.id, CONFIG, 'system', createAdapter(captured), manager)
+
+      assert.equal(result.compressed, true)
+      assert.equal(result.reason, 'success')
+      assert.ok(captured.prompt!.includes('[PREVIOUS SUMMARY'))
+      assert.ok(captured.prompt!.includes('old summary of earlier topics'))
+      assert.ok(captured.prompt!.includes('TOOL_DATA_GAMMA'))
+    } finally {
+      manager.close()
+      cleanup(dir)
+    }
+  })
+})

packages/node-runtime/src/ai/compression/compressor.ts

@@ -5,10 +5,20 @@
  * 通过 AIChatManager 操作对话数据。
  */
 
+import { truncateToolResultText } from '@openchatlab/core'
+
 import { countTokens, countMessagesTokens } from '../tokenizer'
+import { isReplayableToolBlock } from '../agent/history'
 import type { AIChatManager, ContentBlock, AIMessageRole } from '../chats'
 import type { CompressionConfig, CompressionResult, CompressionLogger, CompressionLlmAdapter } from './types'
 
+interface CompressibleMessage {
+  role: string
+  content: string
+  timestamp: number
+  contentBlocks?: ContentBlock[]
+}
+
 const DEFAULT_CONTEXT_WINDOW = 128000
 
 const INITIAL_COMPRESSION_PROMPT = `You are a context compression assistant. Compress the conversation below into a structured summary.
@@ -70,7 +80,7 @@ export async function checkAndCompress(
 
     const summary = convManager.getLatestSummary(aiChatId)
 
-    let messages: Array<{ role: AIMessageRole; content: string; timestamp: number }>
+    let messages: Array<{ role: AIMessageRole; content: string; timestamp: number; contentBlocks?: ContentBlock[] }>
     if (summary) {
       const metaBlock = summary.contentBlocks?.find(
         (b): b is Extract<ContentBlock, { type: 'summary_meta' }> => b.type === 'summary_meta'
@@ -87,6 +97,11 @@ export async function checkAndCompress(
     }
     for (const msg of messages) {
       historyForTokenCount.push({ role: msg.role, content: msg.content })
+      // Persisted tool results are replayed as toolCall/toolResult pairs each
+      // turn (see agent/history.ts), so they occupy real context and must be counted.
+      for (const toolText of replayedToolResultTexts(msg.contentBlocks)) {
+        historyForTokenCount.push({ role: 'tool', content: toolText })
+      }
     }
 
     const currentTokens = countMessagesTokens(historyForTokenCount, systemPrompt)
@@ -133,10 +148,13 @@ export async function checkAndCompress(
       compressedMessageCount: messagesToCompress.length,
     })
 
-    const afterTokenCount: Array<{ role: string; content: string }> = [
-      { role: 'assistant', content: summaryText },
-      ...bufferMessages.map((m) => ({ role: m.role, content: m.content })),
-    ]
+    const afterTokenCount: Array<{ role: string; content: string }> = [{ role: 'assistant', content: summaryText }]
+    for (const m of bufferMessages) {
+      afterTokenCount.push({ role: m.role, content: m.content })
+      for (const toolText of replayedToolResultTexts(m.contentBlocks)) {
+        afterTokenCount.push({ role: 'tool', content: toolText })
+      }
+    }
     const tokensAfter = countMessagesTokens(afterTokenCount, systemPrompt)
 
     if (tokensAfter >= thresholdTokens) {
@@ -186,18 +204,32 @@ export async function manualCompress(
 
 // ==================== Internal Helpers ====================
 
-function splitMessagesForCompression(
-  messages: Array<{ role: string; content: string; timestamp: number }>,
+/** Tool result texts that history replay will inject back into the LLM context. */
+function replayedToolResultTexts(blocks?: ContentBlock[]): string[] {
+  if (!blocks) return []
+  return blocks.filter(isReplayableToolBlock).map((block) => truncateToolResultText(block.tool.result))
+}
+
+function countMessageTokensWithTools(msg: CompressibleMessage): number {
+  let tokens = countTokens(msg.content) + 4
+  for (const toolText of replayedToolResultTexts(msg.contentBlocks)) {
+    tokens += countTokens(toolText) + 4
+  }
+  return tokens
+}
+
+function splitMessagesForCompression<T extends CompressibleMessage>(
+  messages: T[],
   bufferTokenBudget: number
 ): {
-  bufferMessages: Array<{ role: string; content: string; timestamp: number }>
-  messagesToCompress: Array<{ role: string; content: string; timestamp: number }>
+  bufferMessages: T[]
+  messagesToCompress: T[]
 } {
   let bufferTokens = 0
   let splitIndex = messages.length
 
   for (let i = messages.length - 1; i >= 0; i--) {
-    const msgTokens = countTokens(messages[i].content) + 4
+    const msgTokens = countMessageTokensWithTools(messages[i])
     if (bufferTokens + msgTokens > bufferTokenBudget) {
       splitIndex = i + 1
       break
@@ -215,7 +247,7 @@ function splitMessagesForCompression(
 }
 
 function buildCompressionInput(
-  messagesToCompress: Array<{ role: string; content: string }>,
+  messagesToCompress: Array<{ role: string; content: string; contentBlocks?: ContentBlock[] }>,
   existingSummary: { content: string } | null
 ): string {
   const parts: string[] = []
@@ -228,6 +260,9 @@ function buildCompressionInput(
   for (const msg of messagesToCompress) {
     const roleLabel = msg.role === 'user' ? 'User' : 'Assistant'
     parts.push(`${roleLabel}: ${msg.content}`)
+    for (const block of (msg.contentBlocks ?? []).filter(isReplayableToolBlock)) {
+      parts.push(`[Tool result: ${block.tool.name}]\n${truncateToolResultText(block.tool.result)}`)
+    }
   }
 
   return parts.join('\n\n')