diff --git a/packages/sdk/src/plugins/result-eviction/plugin.ts b/packages/sdk/src/plugins/result-eviction/plugin.ts index ffe62a5..225f8bf 100644 --- a/packages/sdk/src/plugins/result-eviction/plugin.ts +++ b/packages/sdk/src/plugins/result-eviction/plugin.ts @@ -6,6 +6,7 @@ * head + tail + file path is returned instead. */ +import { type ChatMessageContentItem, contentToString, type ToolResultContent } from '~/core/llm/llm-log-types.js' import { truncateByTokens } from '~/core/llm/tokens.js' import { definePlugin } from '~/core/plugins/plugin-builder.js' @@ -23,6 +24,21 @@ export interface EvictionAgentConfig { const DEFAULT_MAX_TOKENS = 20_000 +const isImage = (item: ChatMessageContentItem): boolean => item.type === 'image_url' + +/** + * Rebuild a multi-part result with its text collapsed to `preview` and every + * image part preserved in order. + * + * Images are kept rather than truncated because there is no partial image — the + * alternative to keeping it is dropping the observation entirely. Text is the + * part that both dominates the payload and survives being summarised. + */ +const withTextReplaced = (content: ChatMessageContentItem[], preview: string): ChatMessageContentItem[] => [ + { type: 'text', text: preview }, + ...content.filter(isImage), +] + export const resultEvictionPlugin = definePlugin('result-eviction') .agentConfig() .hook('afterToolCall', async (ctx) => { @@ -34,31 +50,30 @@ export const resultEvictionPlugin = definePlugin('result-eviction') } const { content } = ctx.result - - // Only evict string content - if (typeof content !== 'string') { - return null - } - const maxTokens = agentConfig?.config?.maxTokens ?? DEFAULT_MAX_TOKENS - const truncation = truncateByTokens(content, maxTokens) + // Multi-part results (e.g. a page snapshot returning JSON + a screenshot) + // used to bail out here, so no threshold could ever evict them however + // large the text grew. Measure the text, evict the text, keep the images. + const text = contentToString(content) + const truncation = truncateByTokens(text, maxTokens) if (!truncation) { return null } - // Write full content to file via FileStore + // Write full text to file via FileStore const fileName = `${ctx.toolCall.id}.txt` const filePath = `.results/${fileName}` - await ctx.files.session.write(filePath, content) + await ctx.files.session.write(filePath, text) - const truncatedContent = `${truncation.content}\n\n[Full output saved to: ${filePath}]` + const preview = `${truncation.content}\n\n[Full output saved to: ${filePath}]` + const evicted: ToolResultContent = typeof content === 'string' ? preview : withTextReplaced(content, preview) return { action: 'modify', result: { isError: false, - content: truncatedContent, + content: evicted, }, } }) diff --git a/packages/sdk/src/plugins/result-eviction/result-eviction.integration.test.ts b/packages/sdk/src/plugins/result-eviction/result-eviction.integration.test.ts index d13317a..04106cf 100644 --- a/packages/sdk/src/plugins/result-eviction/result-eviction.integration.test.ts +++ b/packages/sdk/src/plugins/result-eviction/result-eviction.integration.test.ts @@ -19,6 +19,18 @@ const bigOutputPlugin = definePlugin('big-output') input: z.object({ size: z.number() }), execute: async (input) => Ok('x'.repeat(input.size)), }), + // Mirrors a page-snapshot style result: a large JSON/text block plus an + // image part. These used to bypass eviction entirely. + createTool({ + name: 'generate_snapshot', + description: 'Generate a text + image result of a given text size', + input: z.object({ size: z.number() }), + execute: async (input) => + Ok([ + { type: 'text' as const, text: 'x'.repeat(input.size) }, + { type: 'image_url' as const, imageUrl: { url: 'data:image/png;base64,AAAA' } }, + ]), + }), ]) .build() @@ -100,6 +112,82 @@ describe('result-eviction plugin', () => { await harness.shutdown() }) + it('large multi-part output → text evicted, image part kept', async () => { + let toolContent: unknown = null + const harness = createEvictionHarness({ + presets: [createEvictionPreset()], + mockHandler: (request) => { + const callCount = harness.llmProvider.getCallCount() + if (callCount === 1) { + return { + content: null, + toolCalls: [{ id: ToolCallId('tc1'), name: 'generate_snapshot', input: { size: 250_000 } }], + finishReason: 'stop', + metrics: MockLLMProvider.defaultMetrics(), + } + } + for (const msg of request.messages) { + if (msg.role === 'tool' && Array.isArray(msg.content)) { + toolContent = msg.content + } + } + return { content: 'Done', toolCalls: [], finishReason: 'stop', metrics: MockLLMProvider.defaultMetrics() } + }, + }) + + const session = await harness.createSession('test') + await session.sendAndWaitForIdle('Generate large snapshot') + + expect(Array.isArray(toolContent)).toBe(true) + const parts = toolContent as Array<{ type: string; text?: string }> + + // Text collapsed to a preview that names the spill file... + const textParts = parts.filter((p) => p.type === 'text') + expect(textParts).toHaveLength(1) + expect(textParts[0].text).toContain('truncated') + expect(textParts[0].text).toContain('.results/') + expect(textParts[0].text!.length).toBeLessThan(250_000) + + // ...and the image survives, because there is no partial image: the only + // alternative to keeping it is losing the observation. + expect(parts.filter((p) => p.type === 'image_url')).toHaveLength(1) + + await harness.shutdown() + }) + + it('small multi-part output → returned unchanged', async () => { + let toolContent: unknown = null + const harness = createEvictionHarness({ + presets: [createEvictionPreset()], + mockHandler: (request) => { + const callCount = harness.llmProvider.getCallCount() + if (callCount === 1) { + return { + content: null, + toolCalls: [{ id: ToolCallId('tc1'), name: 'generate_snapshot', input: { size: 100 } }], + finishReason: 'stop', + metrics: MockLLMProvider.defaultMetrics(), + } + } + for (const msg of request.messages) { + if (msg.role === 'tool' && Array.isArray(msg.content)) { + toolContent = msg.content + } + } + return { content: 'Done', toolCalls: [], finishReason: 'stop', metrics: MockLLMProvider.defaultMetrics() } + }, + }) + + const session = await harness.createSession('test') + await session.sendAndWaitForIdle('Generate small snapshot') + + const parts = toolContent as Array<{ type: string; text?: string }> + expect(parts.filter((p) => p.type === 'text')[0].text).toBe('x'.repeat(100)) + expect(parts.filter((p) => p.type === 'image_url')).toHaveLength(1) + + await harness.shutdown() + }) + it('full output saved to .results/.txt', async () => { const harness = createEvictionHarness({ presets: [createEvictionPreset()],