fix(projection): a mid-turn message's media count toward the physical-window restart

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.

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;
+  }
 }
 
 /**

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<void> {}
+  async stop(): Promise<void> {}
+  getTools(): ToolDefinition[] {
+    return [{ name: 'fetch', description: 'small', inputSchema: { type: 'object', properties: {} } }];
+  }
+  async handleToolCall(): Promise<ToolResult> {
+    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<EventResponse> {
+    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 {