diff --git a/changelog.d/projection-mid-turn-images.fixed.md b/changelog.d/projection-mid-turn-images.fixed.md new file mode 100644 index 000000000..c0d454281 --- /dev/null +++ b/changelog.d/projection-mid-turn-images.fixed.md @@ -0,0 +1 @@ +- The physical-window projection (`physicalWindowTokens`) prices what a mid-turn message brings into a continuation round: an image or a generated image at its `tokenEstimate`, or 1,600 tokens without a usable one, and a document, audio or video at 1,000, the flat estimates context-manager's store uses for images and other media. It priced only such a message's text, so a round that carried photos or audio could cross the window the projection exists to keep a request inside. diff --git a/src/framework.ts b/src/framework.ts index 6e7f08a2f..e088ba7c6 100644 --- a/src/framework.ts +++ b/src/framework.ts @@ -336,11 +336,17 @@ const isConversationalInjection = (metadata?: MessageMetadata): boolean => { /** * Rough token estimate for the blocks a continuation round is about to * append to the live wire (issue #92 physical-window projection): the - * spilled tool-result strings at chars/4, a flat per-image cost for native - * image blocks the wire path preserves (provider vision tokens are - * resolution-dependent; 1600 ≈ a max-size Anthropic tile), mid-turn injected - * text at chars/4, and a small per-block envelope allowance. Deliberately a - * ceiling-ish heuristic — the cost of overestimating is one early recompile; + * spilled tool-result strings at chars/4, a flat per-image cost for the + * native image blocks a tool result carries (provider vision tokens are + * resolution-dependent; 1600 ≈ one image at Anthropic's largest size), a + * mid-turn message's text at chars/4 and its other blocks about as + * context-manager's store prices them (injectedBlockTokens), and a small + * per-block envelope allowance. On membrane's native tool path a mid-turn + * message rides into the next round with its stored content as it is + * (provideToolResults' injectedMessages); its XML tool path drops injected + * messages, so there the estimate errs high. Deliberately a + * ceiling-ish heuristic, but for a mid-turn message's media at the store's + * flat prices — the cost of overestimating is one early recompile; * the cost of underestimating is the 400 this projection exists to prevent. */ function estimateAppendedRoundTokens( @@ -359,15 +365,52 @@ function estimateAppendedRoundTokens( } } } + let injectedTokens = 0; for (const inj of injections) { for (const block of inj.content) { if (block.type === 'text' && typeof (block as { text?: unknown }).text === 'string') { chars += (block as { text: string }).text.length; + } else { + injectedTokens += injectedBlockTokens(block); } } } const blockCount = toolResults.length + injections.length; - return Math.ceil(chars / 4) + images * 1600 + blockCount * 20; + return Math.ceil(chars / 4) + images * 1600 + injectedTokens + blockCount * 20; +} + +/** + * A mid-turn message's non-text block, priced about as context-manager's store + * prices it (MessageStore's raw estimate at 0.14.1): an image at its + * `tokenEstimate` stamp, or 1600 without one, and a document, audio or video + * at 1000. These are the store's flat estimates, not ceilings: a PDF's pages + * can cost more than 1000. Of the writers in agent-framework, channel + * messages and MCPL bring images, MCPL push events and channel messages + * bring audio, and none writes a document or a video; a host module's + * addMessage can write any of them. Two departures err high. A generated + * image, which the store prices at 0, is priced as an image, since + * membrane's Anthropic formatter sends it as one. A stamp that isn't a + * finite number, 0 or more, is priced at 1600, where the store takes it as + * given: a NaN would make the whole projection NaN, and the round would + * never restart. Thinking and other blocks are priced at 0: they come from + * provider responses, stored as the agent's own messages, and nothing in + * agent-framework writes them under another participant, though a host + * module's addMessage could. + */ +function injectedBlockTokens(block: ContentBlock): number { + switch (block.type) { + case 'image': + case 'generated_image': { + const stamp = (block as { tokenEstimate?: unknown }).tokenEstimate; + return typeof stamp === 'number' && Number.isFinite(stamp) && stamp >= 0 ? stamp : 1600; + } + case 'document': + case 'audio': + case 'video': + return 1000; + default: + return 0; + } } /** diff --git a/test/physical-window-restart.test.ts b/test/physical-window-restart.test.ts index a513eff83..796e49182 100644 --- a/test/physical-window-restart.test.ts +++ b/test/physical-window-restart.test.ts @@ -18,7 +18,7 @@ import assert from 'node:assert'; import { mkdtempSync, rmSync } from 'node:fs'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; -import type { NormalizedRequest, NormalizedResponse, YieldingStream } from '@animalabs/membrane'; +import type { ContentBlock, NormalizedRequest, NormalizedResponse, YieldingStream } from '@animalabs/membrane'; import type { EventResponse, Module, @@ -76,6 +76,62 @@ class BlobToolModule implements Module { } } +/** A 1x1 PNG: a real image, so the message is one a store keeps as sent. */ +const PNG_1PX = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII='; +const PHOTOS_TEXT = 'Here are the photos.'; +const image = (): ContentBlock => ({ type: 'image', source: { type: 'base64', data: PNG_1PX, mediaType: 'image/png' } }) as ContentBlock; + +/** + * A tool whose first call brings in a message with the given blocks beside + * a line of text while it runs, as someone sending photos mid-turn does. + * The framework holds the message, and the next round carries it as a + * mid-turn injection. Later calls bring in nothing. + */ +class PhotosDuringToolModule implements Module { + readonly name = 'canned'; + framework: AgentFramework | undefined = undefined; + private messages = 0; + private brought = false; + private readonly blocks: ContentBlock[]; + constructor(blocks: number | ContentBlock[]) { + this.blocks = typeof blocks === 'number' ? Array.from({ length: blocks }, image) : blocks; + } + async start(): Promise {} + async stop(): Promise {} + getTools(): ToolDefinition[] { + return [{ name: 'fetch', description: 'small', inputSchema: { type: 'object', properties: {} } }]; + } + async handleToolCall(): Promise { + if (!this.brought) { + this.brought = true; + const framework = this.framework!; + framework.pushEvent({ + type: 'external-message', + source: 'test', + content: [{ type: 'text', text: PHOTOS_TEXT }, ...this.blocks], + metadata: {}, + } as unknown as ProcessEvent); + // The call returns once the framework holds this message for the + // round's boundary. + const held = () => (framework as unknown as { deferredMessages: Array<{ content: ContentBlock[] }> }).deferredMessages + .some((m) => m.content.some((b) => b.type === 'text' && b.text === PHOTOS_TEXT)); + for (const end = Date.now() + 5_000; !held();) { + if (Date.now() > end) throw new Error('the message was never held for the boundary'); + await new Promise((r) => setTimeout(r, 5)); + } + } + return { success: true, data: { ok: true } }; + } + async onProcess(event: ProcessEvent): Promise { + if (event.type !== 'external-message') return {}; + // Only the first message, the turn's own, asks for a turn. + return { + addMessages: [{ participant: 'User', content: (event as { content: unknown }).content as never }], + ...(this.messages++ === 0 ? { requestInference: true } : {}), + }; + } +} + /** Tool round whose usage reports a near-cap REAL prefix: only 5k fresh * input (far under maxStreamTokens) but 195k cache-read — the projection * must count cached tokens or it misses the window entirely. */ @@ -268,6 +324,7 @@ describe('physical-window mid-turn restart (issue #92)', () => { modules: [module], syncIntervalMs: 0, }); + if ('framework' in module) (module as { framework?: AgentFramework }).framework = framework; const traces: TraceEvent[] = []; framework.onTrace((e) => traces.push(e)); framework.start(); @@ -500,6 +557,135 @@ describe('physical-window mid-turn restart (issue #92)', () => { } }); + /** The messages of a request that carry the photos message's text. */ + const photosIn = (request: { messages: Array<{ content: ContentBlock[] }> }) => + request.messages.filter((m) => m.content.some((b) => b.type === 'text' && b.text.includes(PHOTOS_TEXT))); + + for (const images of [0, 3]) { + it(`counts a mid-turn message's images toward a restart: ${images} image${images === 1 ? '' : 's'}`, async () => { + // A 163k prompt, 2k of output and the 32k reserve make 197k. The round + // appends a short tool result and a message, and each image in the + // message counts 1,600 tokens, as a tool result's does: with none the + // projection fits the 200k window, and three cross it. + const t = await windowTurn([ + [ + round('call_1', { inputTokens: 20_000, cacheReadTokens: 143_000, outputTokens: 2_000 }), + round(null, { inputTokens: 21_000, cacheReadTokens: 306_000, outputTokens: 2_500 }), + ], + afterRestart, + ], new PhotosDuringToolModule(images)); + try { + assert.ok(t.done, 'the turn should finish'); + if (images === 0) { + assert.strictEqual(t.restarts.length, 0, 'control: without images the round fits'); + const injected = t.membrane.lastStream!.receivedToolResultOptions + .flatMap((o) => o?.injectedMessages ?? []) + .flatMap((m) => m.content as Array<{ type: string; text?: string }>); + assert.ok(injected.some((b) => b.text === PHOTOS_TEXT), 'control: the message rode into the next round'); + } else { + assert.strictEqual(t.restarts.length, 1, "the message's images cross the window"); + assert.strictEqual(t.restarts[0]!.reason, 'physical_window'); + // The round's short texts and envelopes come to under 100 tokens, + // which this doesn't pin: the images' 1,600 each is what it checks. + const appended = t.restarts[0]!.inputTokens! - (163_000 + 2_000 + 32_000); + assert.ok(appended >= images * 1_600 && appended < images * 1_600 + 100, `the round's estimate counts each image at 1,600 tokens: ${appended}`); + assert.ok(storedTexts(t.framework).includes('Answered after a restart'), "the answer is the restarted stream's"); + assert.strictEqual(t.membrane.calls.length, 2, 'the first stream, then the restart'); + const carried = photosIn(t.membrane.calls[1]!); + assert.strictEqual(carried.length, 1, 'the restarted request carries the photos message once'); + assert.strictEqual(carried[0]!.content.filter((b) => b.type === 'image').length, images, 'with its images'); + } + } finally { + await t.close(); + } + }); + } + + /** One round on the same figures as the image test, with a message of + * the given blocks: whether it restarts, what the round added, and + * whether the message reached the next round or the restarted request. */ + async function photosRound(blocks: ContentBlock[]): Promise<{ restarted: boolean; appended: number | undefined; carried: boolean }> { + const t = await windowTurn([ + [ + round('call_1', { inputTokens: 20_000, cacheReadTokens: 143_000, outputTokens: 2_000 }), + round(null, { inputTokens: 21_000, cacheReadTokens: 306_000, outputTokens: 2_500 }), + ], + afterRestart, + ], new PhotosDuringToolModule(blocks)); + try { + assert.ok(t.done, 'the turn should finish'); + const restart = t.restarts[0]; + const injected = t.membrane.lastStream!.receivedToolResultOptions + .flatMap((o) => o?.injectedMessages ?? []) + .some((m) => (m.content as ContentBlock[]).some((b) => b.type === 'text' && b.text === PHOTOS_TEXT)); + const carried = restart !== undefined ? photosIn(t.membrane.calls[1]!).length === 1 : injected; + return { restarted: restart !== undefined, appended: restart ? restart.inputTokens! - (163_000 + 2_000 + 32_000) : undefined, carried }; + } finally { + await t.close(); + } + } + + it("prices a mid-turn generated image as an image", async () => { + const generated = { type: 'generated_image', data: PNG_1PX, mimeType: 'image/png' } as ContentBlock; + const { restarted, appended } = await photosRound([generated, generated, generated]); + assert.ok(restarted, 'three generated images cross the window, as three images do'); + assert.ok(appended! >= 3 * 1_600 && appended! < 3 * 1_600 + 100, `each at 1,600 tokens: ${appended}`); + }); + + it("prices a mid-turn image at its tokenEstimate stamp", async () => { + const stamped = { ...image(), tokenEstimate: 100 } as ContentBlock; + const { restarted, carried } = await photosRound([stamped, stamped, stamped]); + assert.ok(!restarted, 'three images stamped at 100 tokens fit, where three unstamped cross'); + assert.ok(carried, 'the message rode into the next round'); + }); + + for (const stamp of [Number.NaN, -1_600]) { + it(`prices a mid-turn image whose stamp is ${stamp} at 1,600, as an unstamped one`, async () => { + const stamped = { ...image(), tokenEstimate: stamp } as ContentBlock; + const { restarted, appended } = await photosRound([stamped, stamped, stamped]); + assert.ok(restarted, 'three such images cross the window, as three unstamped do'); + assert.ok(appended! >= 3 * 1_600 && appended! < 3 * 1_600 + 100, `each at 1,600 tokens: ${appended}`); + }); + } + + for (const kind of ['document', 'audio', 'video'] as const) { + it(`prices a mid-turn ${kind} at 1,000 tokens, as the store does`, async () => { + const mediaType = { document: 'application/pdf', audio: 'audio/wav', video: 'video/mp4' }[kind]; + const block = { type: kind, source: { type: 'base64', data: 'AAAA', mediaType } } as ContentBlock; + const { restarted, appended } = await photosRound([block, block, block, block]); + assert.ok(restarted, `four of them cross the window`); + assert.ok(appended! >= 4 * 1_000 && appended! < 4 * 1_000 + 100, `each at 1,000 tokens: ${appended}`); + }); + } + + it("carries a mid-turn message's images past a call that reports no prompt", async () => { + // Call 1 reports a 160k prompt and 2k of output. During its tool a + // message brings in three photos: 160k + 2k + the round (the photos' + // 4,800 and under 100 for its texts and envelopes) + the 32k reserve is + // under 199k, so the round fits and goes on, and what it appended joins + // the growth the next boundary carries. Call 2 reports 2k of output but + // no prompt, so the second boundary projects call 1's figures plus that + // growth, past the window. Without the photos in the growth it would + // project about 196.1k and dispatch. + const t = await windowTurn([ + [ + round('call_1', { inputTokens: 20_000, cacheReadTokens: 140_000, outputTokens: 2_000 }), + round('call_2', { inputTokens: 20_000, cacheReadTokens: 140_000, outputTokens: 4_000 }), + round(null, { inputTokens: 21_000, cacheReadTokens: 300_000, outputTokens: 4_500 }), + ], + afterRestart, + ], new PhotosDuringToolModule(3)); + try { + assert.ok(t.done, 'the turn should finish'); + assert.strictEqual(t.restarts.length, 1, 'the second boundary restarts, the first fits'); + assert.strictEqual(t.membrane.calls.length, 2, 'the first stream, then the restart'); + const carried = t.restarts[0]!.inputTokens! - (160_000 + 2_000 + 2_000 + 32_000); + assert.ok(carried >= 3 * 1_600 && carried < 3 * 1_600 + 100, `the growth carried the photos at 1,600 each: ${carried}`); + } finally { + await t.close(); + } + }); + it('leaves behavior unchanged when physicalWindowTokens is unset', async () => { const h = await runTurn(undefined, new MockMembrane()); try {