From 42649d835d0b78c2d3b48350cb160bf3f9e3c78b Mon Sep 17 00:00:00 2001 From: digua Date: Wed, 10 Jun 2026 20:51:07 +0800 Subject: [PATCH] fix(ai): count replayed tool results in compression token accounting --- packages/node-runtime/src/ai/agent/history.ts | 4 +- packages/node-runtime/src/ai/chats.ts | 24 ++- .../compression/__tests__/compressor.test.ts | 159 ++++++++++++++++++ .../src/ai/compression/compressor.ts | 57 +++++-- 4 files changed, 227 insertions(+), 17 deletions(-) create mode 100644 packages/node-runtime/src/ai/compression/__tests__/compressor.test.ts diff --git a/packages/node-runtime/src/ai/agent/history.ts b/packages/node-runtime/src/ai/agent/history.ts index 9046ea32..1461efdf 100644 --- a/packages/node-runtime/src/ai/agent/history.ts +++ b/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 -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 ( diff --git a/packages/node-runtime/src/ai/chats.ts b/packages/node-runtime/src/ai/chats.ts index 744ec75d..ca8b0993 100644 --- a/packages/node-runtime/src/ai/chats.ts +++ b/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 { diff --git a/packages/node-runtime/src/ai/compression/__tests__/compressor.test.ts b/packages/node-runtime/src/ai/compression/__tests__/compressor.test.ts new file mode 100644 index 00000000..15c85a92 --- /dev/null +++ b/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) + } + }) +}) diff --git a/packages/node-runtime/src/ai/compression/compressor.ts b/packages/node-runtime/src/ai/compression/compressor.ts index fe4da93f..2c06d66a 100644 --- a/packages/node-runtime/src/ai/compression/compressor.ts +++ b/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 => 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( + 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')