@vellumai/assistant 0.5.15 → 0.6.0

package/src/__tests__/context-token-estimator.test.ts

@@ -9,6 +9,27 @@ import {
 } from "../context/token-estimator.js";
 import type { Message } from "../providers/types.js";
 
+/** Build a minimal valid PNG header with the given dimensions, returned as base64. */
+function makePngBase64(width: number, height: number): string {
+  const header = Buffer.alloc(24);
+  header[0] = 0x89;
+  header[1] = 0x50;
+  header[2] = 0x4e;
+  header[3] = 0x47;
+  header[4] = 0x0d;
+  header[5] = 0x0a;
+  header[6] = 0x1a;
+  header[7] = 0x0a;
+  header.writeUInt32BE(13, 8);
+  header[12] = 0x49;
+  header[13] = 0x48;
+  header[14] = 0x44;
+  header[15] = 0x52;
+  header.writeUInt32BE(width, 16);
+  header.writeUInt32BE(height, 20);
+  return header.toString("base64");
+}
+
 describe("token estimator", () => {
   test("estimates text tokens from character length", () => {
     expect(estimateTextTokens("")).toBe(0);
@@ -48,7 +69,7 @@ describe("token estimator", () => {
           data: "a".repeat(100),
         },
       }),
-    ).toBeGreaterThan(500);
+    ).toBeGreaterThan(0);
   });
 
   test("estimates message and prompt totals", () => {
@@ -264,9 +285,11 @@ describe("token estimator", () => {
       { providerName: "anthropic" },
     );
 
-    // 1920x1080 scaled to fit 1568x1568: scale = 1568/1920 = 0.8167
+    // 1920x1080 scaled to fit 1568px bounding box: dimScale = 1568/1920 = 0.8167
     // scaledWidth = round(1920 * 0.8167) = 1568, scaledHeight = round(1080 * 0.8167) = 882
-    // tokens = ceil(1568 * 882 / 750) = ceil(1843.968) = ~1844
+    // pixels = 1568 * 882 = 1,382,976 > 1,200,000 → mpScale = sqrt(1200000/1382976) = 0.9315
+    // scaledWidth = round(1568 * 0.9315) = 1461, scaledHeight = round(882 * 0.9315) = 822
+    // tokens = ceil(1461 * 822 / 750) = ceil(1601.26) = ~1,602
     // With IMAGE_BLOCK_OVERHEAD_TOKENS and media_type overhead, still well under 5000
     expect(anthropicTokens).toBeLessThan(5_000);
 
@@ -299,13 +322,10 @@ describe("token estimator", () => {
       { providerName: "anthropic" },
     );
 
-    // Should fall back to ANTHROPIC_IMAGE_MAX_TOKENS (~3,277)
-    // The total will include IMAGE_BLOCK_OVERHEAD_TOKENS + media_type overhead,
-    // but the max is applied at the outer Math.max(IMAGE_BLOCK_TOKENS, ...) level
-    // ANTHROPIC_IMAGE_MAX_TOKENS = ceil(1568*1568/750) = 3277
-    // Total = max(1024, 16 + ceil(9/4) + 3277) = max(1024, 3296) = 3296
-    expect(tokens).toBeGreaterThanOrEqual(3_277);
-    expect(tokens).toBeLessThan(4_000);
+    // Should fall back to ANTHROPIC_IMAGE_MAX_TOKENS (1,600)
+    // Total = 16 (block overhead) + ceil(9/4) (media_type) + 1600 = 1619
+    expect(tokens).toBeGreaterThanOrEqual(1_600);
+    expect(tokens).toBeLessThan(2_000);
   });
 
   test("Anthropic image tokens are the same for same-dimension images regardless of payload size", () => {
@@ -356,4 +376,149 @@ describe("token estimator", () => {
     // For Anthropic, same dimensions should produce the same estimate
     expect(largeTokens).toBe(smallTokens);
   });
+
+  test("applies megapixel cap for square images on Anthropic", () => {
+    // Build a minimal valid PNG header encoding 2000x2000 dimensions.
+    const pngHeader = Buffer.alloc(24);
+    // PNG signature
+    pngHeader[0] = 0x89;
+    pngHeader[1] = 0x50;
+    pngHeader[2] = 0x4e;
+    pngHeader[3] = 0x47;
+    pngHeader[4] = 0x0d;
+    pngHeader[5] = 0x0a;
+    pngHeader[6] = 0x1a;
+    pngHeader[7] = 0x0a;
+    // IHDR chunk length (13 bytes)
+    pngHeader.writeUInt32BE(13, 8);
+    // "IHDR"
+    pngHeader[12] = 0x49;
+    pngHeader[13] = 0x48;
+    pngHeader[14] = 0x44;
+    pngHeader[15] = 0x52;
+    // Width: 2000
+    pngHeader.writeUInt32BE(2000, 16);
+    // Height: 2000
+    pngHeader.writeUInt32BE(2000, 20);
+
+    const base64Data = pngHeader.toString("base64");
+
+    const tokens = estimateContentBlockTokens(
+      {
+        type: "image",
+        source: { type: "base64", media_type: "image/png", data: base64Data },
+      },
+      { providerName: "anthropic" },
+    );
+
+    // 2000x2000 → dimScale = 1568/2000 = 0.784 → 1568x1568 = 2,458,624 pixels
+    // 2,458,624 > 1,200,000 → mpScale = sqrt(1200000/2458624) ≈ 0.6987
+    // scaledWidth = round(1568 * 0.6987) = 1096, scaledHeight = round(1568 * 0.6987) = 1096
+    // tokens = ceil(1096 * 1096 / 750) = ceil(1601.6) ≈ 1602
+    // Without megapixel cap would have been ceil(1568 * 1568 / 750) ≈ 3277
+    expect(tokens).toBeLessThanOrEqual(1_700);
+  });
+
+  test("small Anthropic images are not inflated to 1024 tokens", () => {
+    // 200x200 image: ceil(200*200/750) = ceil(53.33) = 54 tokens
+    const tokens = estimateContentBlockTokens(
+      {
+        type: "image",
+        source: {
+          type: "base64",
+          media_type: "image/png",
+          data: makePngBase64(200, 200),
+        },
+      },
+      { providerName: "anthropic" },
+    );
+
+    // 54 (dimension-based) + 16 (block overhead) + 3 (media type) = 73
+    expect(tokens).toBeLessThan(100);
+    expect(tokens).toBeGreaterThan(50);
+  });
+
+  test("thumbnail Anthropic images estimate accurately", () => {
+    // 150x150 image: ceil(150*150/750) = ceil(30) = 30 tokens
+    const tokens = estimateContentBlockTokens(
+      {
+        type: "image",
+        source: {
+          type: "base64",
+          media_type: "image/png",
+          data: makePngBase64(150, 150),
+        },
+      },
+      { providerName: "anthropic" },
+    );
+
+    // 30 + 16 + 3 = 49
+    expect(tokens).toBeLessThan(70);
+    expect(tokens).toBeGreaterThan(30);
+  });
+
+  test("many small Anthropic images do not trigger phantom token inflation", () => {
+    const messages: Message[] = [
+      {
+        role: "user",
+        content: Array.from({ length: 100 }, () => ({
+          type: "image" as const,
+          source: {
+            type: "base64" as const,
+            media_type: "image/png",
+            data: makePngBase64(200, 200),
+          },
+        })),
+      },
+    ];
+
+    const total = estimateMessagesTokens(messages, { providerName: "anthropic" });
+
+    // Each image: ~73 tokens. 100 images + message overhead ≈ 7,304
+    // Old behavior: 100 * ~1,043 = ~104,300 (14x overestimate)
+    expect(total).toBeLessThan(10_000);
+  });
+
+  test("matches Anthropic's published table for common aspect ratios", () => {
+    // These are the max sizes that should NOT be further scaled (at or below the megapixel cap).
+    // 1:1 → 1092x1092 (~1,590 tokens)
+    const squareTokens = estimateContentBlockTokens(
+      {
+        type: "image",
+        source: {
+          type: "base64",
+          media_type: "image/png",
+          data: makePngBase64(1092, 1092),
+        },
+      },
+      { providerName: "anthropic" },
+    );
+
+    // 1:2 → 784x1568 (~1,639 tokens)
+    const tallTokens = estimateContentBlockTokens(
+      {
+        type: "image",
+        source: {
+          type: "base64",
+          media_type: "image/png",
+          data: makePngBase64(784, 1568),
+        },
+      },
+      { providerName: "anthropic" },
+    );
+
+    // 1092x1092 = 1,192,464 pixels < 1,200,000 → no megapixel scaling needed.
+    // tokens = ceil(1092 * 1092 / 750) = ceil(1589.95) ≈ 1590
+    // With overhead: 16 + 3 + 1590 = 1609
+    expect(squareTokens).toBeGreaterThan(1_400);
+    expect(squareTokens).toBeLessThan(1_800);
+
+    // 784*1568 = 1,229,312 > 1,200,000 → slight scaling applies
+    // mpScale = sqrt(1200000/1229312) ≈ 0.9881
+    // scaledWidth = round(784 * 0.9881) = 775, scaledHeight = round(1568 * 0.9881) = 1549
+    // tokens = ceil(775 * 1549 / 750) = ceil(1600.6) ≈ 1601
+    // With overhead: 16 + 3 + 1601 = 1620
+    expect(tallTokens).toBeGreaterThan(1_400);
+    expect(tallTokens).toBeLessThan(1_800);
+  });
 });

package/src/__tests__/conversation-agent-loop-overflow.test.ts

@@ -457,6 +457,15 @@ function makeCtx(
       assistantMessageChannel: "vellum" as const,
     }),
 
+    graphMemory: {
+      onCompacted: () => {},
+      prepareMemory: async () => ({ runMessages: [], injectedTokens: 0, latencyMs: 0, mode: "none" as const }),
+      reinjectCachedMemory: (messages: Message[]) => ({
+        runMessages: messages,
+        injectedTokens: 0,
+      }),
+    } as unknown as AgentLoopConversationContext["graphMemory"],
+
     ...overrides,
   } as AgentLoopConversationContext;
 }

package/src/__tests__/conversation-agent-loop.test.ts

@@ -445,6 +445,15 @@ function makeCtx(
       assistantMessageChannel: "vellum" as const,
     }),
 
+    graphMemory: {
+      onCompacted: () => {},
+      prepareMemory: async () => ({ runMessages: [], injectedTokens: 0, latencyMs: 0, mode: "none" as const }),
+      reinjectCachedMemory: (messages: Message[]) => ({
+        runMessages: messages,
+        injectedTokens: 0,
+      }),
+    } as unknown as AgentLoopConversationContext["graphMemory"],
+
     ...overrides,
   } as AgentLoopConversationContext;
 }

package/src/__tests__/conversation-attachments.test.ts

@@ -1,5 +1,5 @@
 import { existsSync } from "node:fs";
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 mock.module("../util/logger.js", () => ({
   getLogger: () =>
@@ -48,14 +48,10 @@ mock.module("../permissions/types.js", () => ({
 import type { AssistantAttachmentDraft } from "../daemon/assistant-attachments.js";
 import { getFilePathForAttachment } from "../memory/attachments-store.js";
 import { addMessage, createConversation } from "../memory/conversation-crud.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 
 initializeDb();
 
-afterAll(() => {
-  resetDb();
-});
-
 function resetTables() {
   const db = getDb();
   db.run("DELETE FROM message_attachments");

package/src/__tests__/conversation-attention-store.test.ts

@@ -1,4 +1,4 @@
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 mock.module("../util/logger.js", () => ({
   getLogger: () =>
@@ -17,7 +17,7 @@ import {
   projectAssistantMessage,
   recordConversationSeenSignal,
 } from "../memory/conversation-attention-store.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 import {
   conversationAssistantAttentionState,
   conversationAttentionEvents,
@@ -66,10 +66,6 @@ function insertAssistantMessage(
     .run();
 }
 
-afterAll(() => {
-  resetDb();
-});
-
 describe("conversation-attention-store", () => {
   beforeEach(() => {
     clearTables();

package/src/__tests__/conversation-clear-safety.test.ts

@@ -9,7 +9,7 @@
  * - lifecycle_events contains `conversations_clear_all` audit entry after clear
  */
 
-import { afterAll, describe, expect, mock, test } from "bun:test";
+import { describe, expect, mock, test } from "bun:test";
 
 mock.module("../util/logger.js", () => ({
   getLogger: () =>
@@ -31,17 +31,13 @@ import {
   createConversation,
   getConversation,
 } from "../memory/conversation-crud.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 import { enforcePolicy, getPolicy } from "../runtime/auth/route-policy.js";
 import type { AuthContext, Scope } from "../runtime/auth/types.js";
 import { conversationManagementRouteDefinitions } from "../runtime/routes/conversation-management-routes.js";
 
 initializeDb();
 
-afterAll(() => {
-  resetDb();
-});
-
 /** Build a synthetic AuthContext for testing. */
 function buildAuthContext(overrides?: {
   principalType?: AuthContext["principalType"];

package/src/__tests__/conversation-delete-schedule-cleanup.test.ts

@@ -3,7 +3,7 @@
  * job also deletes the schedule, preventing orphaned scheduled automations.
  */
 
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 mock.module("../util/logger.js", () => ({
   getLogger: () =>
@@ -23,16 +23,12 @@ import {
   createConversation,
   getConversation,
 } from "../memory/conversation-crud.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 import { conversationManagementRouteDefinitions } from "../runtime/routes/conversation-management-routes.js";
 import { createSchedule, getSchedule } from "../schedule/schedule-store.js";
 
 initializeDb();
 
-afterAll(() => {
-  resetDb();
-});
-
 function getRawDb(): Database {
   return (getDb() as unknown as { $client: Database }).$client;
 }
@@ -72,9 +68,8 @@ describe("DELETE /conversations/:id — schedule cleanup", () => {
   beforeEach(() => {
     getRawDb().run("DELETE FROM cron_runs");
     getRawDb().run("DELETE FROM cron_jobs");
-    getRawDb().run("DELETE FROM memory_item_sources");
+    getRawDb().run("DELETE FROM memory_graph_nodes");
     getRawDb().run("DELETE FROM memory_segments");
-    getRawDb().run("DELETE FROM memory_items");
     getRawDb().run("DELETE FROM memory_summaries");
     getRawDb().run("DELETE FROM memory_embeddings");
     getRawDb().run("DELETE FROM memory_jobs");
@@ -297,9 +292,8 @@ describe("POST /conversations/:id/wipe — schedule cleanup", () => {
   beforeEach(() => {
     getRawDb().run("DELETE FROM cron_runs");
     getRawDb().run("DELETE FROM cron_jobs");
-    getRawDb().run("DELETE FROM memory_item_sources");
+    getRawDb().run("DELETE FROM memory_graph_nodes");
     getRawDb().run("DELETE FROM memory_segments");
-    getRawDb().run("DELETE FROM memory_items");
     getRawDb().run("DELETE FROM memory_summaries");
     getRawDb().run("DELETE FROM memory_embeddings");
     getRawDb().run("DELETE FROM memory_jobs");

package/src/__tests__/conversation-disk-view-integration.test.ts

@@ -13,7 +13,7 @@ import {
   rmSync,
 } from "node:fs";
 import { join } from "node:path";
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 // ---------------------------------------------------------------------------
 // Mocks — must come before any imports that depend on them
@@ -59,14 +59,10 @@ import {
   getConversationDirPath,
   syncMessageToDisk,
 } from "../memory/conversation-disk-view.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 
 initializeDb();
 
-afterAll(() => {
-  resetDb();
-});
-
 function resetTables() {
   const db = getDb();
   db.run("DELETE FROM message_attachments");

package/src/__tests__/conversation-disk-view.test.ts

@@ -8,7 +8,7 @@ import {
 } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 // ---------------------------------------------------------------------------
 // Mocks — must come before any imports that depend on them
@@ -61,14 +61,10 @@ import {
   syncMessageToDisk,
   updateMetaFile,
 } from "../memory/conversation-disk-view.js";
-import { getDb, initializeDb, rawRun, resetDb } from "../memory/db.js";
+import { getDb, initializeDb, rawRun } from "../memory/db.js";
 
 initializeDb();
 
-afterAll(() => {
-  resetDb();
-});
-
 function resetTables() {
   const db = getDb();
   db.run("DELETE FROM message_attachments");

package/src/__tests__/conversation-error.test.ts

@@ -90,7 +90,6 @@ describe("classifyConversationError", () => {
       "rate limit exceeded",
       "Rate-limit hit",
       "too many requests",
-      "overloaded",
     ];
 
     for (const msg of cases) {
@@ -98,12 +97,44 @@ describe("classifyConversationError", () => {
         const result = classifyConversationError(new Error(msg), baseCtx);
         expect(result.code).toBe("PROVIDER_RATE_LIMIT");
         expect(result.retryable).toBe(true);
-        expect(result.userMessage).toContain("busy");
+        expect(result.userMessage).toContain("rate limited");
         expect(result.errorCategory).toBe("rate_limit");
       });
     }
   });
 
+  describe("provider overloaded errors", () => {
+    it('classifies "overloaded" as PROVIDER_OVERLOADED', () => {
+      const result = classifyConversationError(
+        new Error("overloaded"),
+        baseCtx,
+      );
+      expect(result.code).toBe("PROVIDER_OVERLOADED");
+      expect(result.retryable).toBe(true);
+      expect(result.userMessage).toContain("overloaded");
+      expect(result.errorCategory).toBe("provider_overloaded");
+    });
+
+    it("classifies Anthropic overloaded_error (no statusCode) as PROVIDER_OVERLOADED", () => {
+      const err = new ProviderError(
+        'Anthropic API error (undefined): {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}',
+        "anthropic",
+      );
+      const result = classifyConversationError(err, baseCtx);
+      expect(result.code).toBe("PROVIDER_OVERLOADED");
+      expect(result.retryable).toBe(true);
+      expect(result.errorCategory).toBe("provider_overloaded");
+    });
+
+    it("classifies ProviderError with 529 as PROVIDER_OVERLOADED", () => {
+      const err = new ProviderError("Overloaded", "anthropic", 529);
+      const result = classifyConversationError(err, baseCtx);
+      expect(result.code).toBe("PROVIDER_OVERLOADED");
+      expect(result.retryable).toBe(true);
+      expect(result.errorCategory).toBe("provider_overloaded");
+    });
+  });
+
   describe("provider API errors", () => {
     const cases = [
       "HTTP 500 Internal Server Error",

package/src/__tests__/conversation-fork-crud.test.ts

@@ -1,6 +1,6 @@
 import { existsSync, readFileSync } from "node:fs";
 import { join } from "node:path";
-import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+import { beforeEach, describe, expect, mock, test } from "bun:test";
 
 import { eq, like } from "drizzle-orm";
 
@@ -39,7 +39,7 @@ import {
   PRIVATE_CONVERSATION_FORK_ERROR,
 } from "../memory/conversation-crud.js";
 import { getConversationDirPath } from "../memory/conversation-disk-view.js";
-import { getDb, initializeDb, resetDb } from "../memory/db.js";
+import { getDb, initializeDb } from "../memory/db.js";
 import { getRequestLogsByMessageId } from "../memory/llm-request-log-store.js";
 import {
   channelInboundEvents,
@@ -71,10 +71,6 @@ function parseMetadata(metadata: string | null): unknown {
   return metadata == null ? null : JSON.parse(metadata);
 }
 
-afterAll(() => {
-  resetDb();
-});
-
 describe("forkConversation", () => {
   beforeEach(() => {
     resetTables();

package/src/__tests__/conversation-history-web-search.test.ts

@@ -782,6 +782,11 @@ describe("web_search_tool_result structural guard", () => {
     // Extracts tool results from persisted message content for work-item
     // display. web_search_tool_result blocks are not relevant here.
     "runtime/routes/work-items-routes.ts",
+
+    // Media token counting iterates tool_result.contentBlocks for nested
+    // image/file blocks. web_search_tool_result has opaque content with no
+    // contentBlocks property, so it cannot contain nested media.
+    "daemon/context-overflow-reducer.ts",
   ]);
 
   /**

package/src/__tests__/conversation-load-history-repair.test.ts

@@ -106,7 +106,7 @@ function makeConversation(): Conversation {
       stopReason: "end_turn",
     }),
   };
-  return new Conversation(
+  const conv = new Conversation(
     "conv-1",
     provider,
     "system prompt",
@@ -114,6 +114,10 @@ function makeConversation(): Conversation {
     () => {},
     "/tmp",
   );
+  // Default to guardian trust so history repair tests load all messages.
+  // Tests that exercise untrusted-actor filtering override this explicitly.
+  conv.setTrustContext({ trustClass: "guardian", sourceChannel: "vellum" });
+  return conv;
 }
 
 describe("loadFromDb history repair", () => {

package/src/__tests__/conversation-media-retry.test.ts

@@ -18,6 +18,15 @@ function makeImageBlock(
   };
 }
 
+function makeImageBlockWithSize(
+  dataLength: number,
+): Extract<ContentBlock, { type: "image" }> {
+  return {
+    type: "image",
+    source: { type: "base64", media_type: "image/png", data: "A".repeat(dataLength) },
+  };
+}
+
 function makeUserMessage(...blocks: ContentBlock[]): Message {
   return { role: "user", content: blocks };
 }
@@ -88,4 +97,86 @@ describe("stripMediaPayloadsForRetry", () => {
     const latestBlocks = result.messages[3].content;
     expect(latestBlocks.filter((b) => b.type === "image").length).toBe(1);
   });
+
+  // ---------------------------------------------------------------------------
+  // Budget-aware media retention
+  // ---------------------------------------------------------------------------
+
+  test("budget-aware: keeps images that fit within token budget", () => {
+    // Non-Anthropic estimation: estimateTextTokens(base64Data) + overhead (~19 tokens).
+    // Data length 4000 → 1000 data tokens + 19 overhead ≈ 1019 tokens/image.
+    // Budget of 3500 allows 3 images (3 * 1019 = 3057 <= 3500) but not 4.
+    const images = Array.from({ length: 5 }, () => makeImageBlockWithSize(4000));
+    const messages: Message[] = [
+      makeUserMessage({ type: "text", text: "describe these" }, ...images),
+    ];
+
+    const result = stripMediaPayloadsForRetry(messages, {
+      mediaTokenBudget: 3500,
+      providerName: "mock",
+    });
+    expect(result.modified).toBe(true);
+
+    const content = result.messages[0].content;
+    const keptImages = content.filter((b) => b.type === "image");
+    const stubs = content.filter(
+      (b) => b.type === "text" && (b as { text: string }).text.includes("Image omitted"),
+    );
+    expect(keptImages.length).toBe(3);
+    expect(stubs.length).toBe(2);
+  });
+
+  test("budget-aware: keeps all images when budget is generous", () => {
+    const images = Array.from({ length: 3 }, () => makeImageBlockWithSize(100));
+    const messages: Message[] = [
+      makeUserMessage({ type: "text", text: "describe these" }, ...images),
+    ];
+
+    const result = stripMediaPayloadsForRetry(messages, {
+      mediaTokenBudget: 100_000,
+      providerName: "mock",
+    });
+    expect(result.modified).toBe(false);
+    expect(result.replacedBlocks).toBe(0);
+
+    const content = result.messages[0].content;
+    const keptImages = content.filter((b) => b.type === "image");
+    expect(keptImages.length).toBe(3);
+  });
+
+  test("budget-aware: stubs all when budget is zero", () => {
+    const images = Array.from({ length: 3 }, () => makeImageBlockWithSize(100));
+    const messages: Message[] = [
+      makeUserMessage({ type: "text", text: "describe these" }, ...images),
+    ];
+
+    const result = stripMediaPayloadsForRetry(messages, {
+      mediaTokenBudget: 0,
+      providerName: "mock",
+    });
+    expect(result.modified).toBe(true);
+    expect(result.replacedBlocks).toBe(3);
+
+    const content = result.messages[0].content;
+    const keptImages = content.filter((b) => b.type === "image");
+    expect(keptImages.length).toBe(0);
+  });
+
+  test("no options falls back to hardcoded limit", () => {
+    const images = Array.from({ length: 5 }, () => makeImageBlockWithSize(100));
+    const messages: Message[] = [
+      makeUserMessage({ type: "text", text: "describe these" }, ...images),
+    ];
+
+    const result = stripMediaPayloadsForRetry(messages);
+    expect(result.modified).toBe(true);
+
+    const content = result.messages[0].content;
+    const keptImages = content.filter((b) => b.type === "image");
+    const stubs = content.filter(
+      (b) => b.type === "text" && (b as { text: string }).text.includes("Image omitted"),
+    );
+    expect(keptImages.length).toBe(3);
+    expect(stubs.length).toBe(2);
+  });
 });

package/src/__tests__/conversation-runtime-assembly.test.ts

@@ -1458,7 +1458,10 @@ describe("injectNowScratchpad", () => {
     const messageWithMemory: Message = {
       role: "user",
       content: [
-        { type: "text", text: "<memory_context __injected>\nrecalled notes\n</memory_context>" },
+        {
+          type: "text",
+          text: "<memory_context __injected>\nrecalled notes\n</memory_context>",
+        },
         { type: "text", text: "What should I work on?" },
       ],
     };
@@ -1635,9 +1638,9 @@ describe("applyRuntimeInjections with nowScratchpad", () => {
       (result[0].content[0] as { type: "text"; text: string }).text,
     ).toContain("<NOW.md");
     // Original text is last
-    expect(
-      (result[0].content[1] as { type: "text"; text: string }).text,
-    ).toBe("What should I do?");
+    expect((result[0].content[1] as { type: "text"; text: string }).text).toBe(
+      "What should I do?",
+    );
   });
 
   test("does not inject when nowScratchpad is null", () => {