From 3f3e94a604c1ee1bde6a1d650ca616d60cae6ef6 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 9 May 2026 20:03:37 +0100
Subject: [PATCH] test: tighten history image prune assertions

---
 .../run/history-image-prune.test.ts           | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/agents/pi-embedded-runner/run/history-image-prune.test.ts b/src/agents/pi-embedded-runner/run/history-image-prune.test.ts
index 27c9fd0b77e..1149007efa9 100644
--- a/src/agents/pi-embedded-runner/run/history-image-prune.test.ts
+++ b/src/agents/pi-embedded-runner/run/history-image-prune.test.ts
@@ -19,6 +19,19 @@ function expectArrayMessageContent(
   return message.content as Array<{ type: string; text?: string; data?: string }>;
 }
 
+function expectContentBlock(
+  block: { type: string; text?: string; data?: string } | undefined,
+  expected: { type: string; text?: string; data?: string },
+) {
+  expect(block?.type).toBe(expected.type);
+  if ("text" in expected) {
+    expect(block?.text).toBe(expected.text);
+  }
+  if ("data" in expected) {
+    expect(block?.data).toBe(expected.data);
+  }
+}
+
 function expectPrunedImageMessage(
   messages: AgentMessage[],
   errorMessage: string,
@@ -26,7 +39,7 @@ function expectPrunedImageMessage(
   const pruned = expectPrunedMessages(messages);
   const content = expectArrayMessageContent(pruned[0], errorMessage);
   expect(content).toHaveLength(2);
-  expect(content[1]).toMatchObject({ type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
+  expectContentBlock(content[1], { type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
   return content;
 }
 
@@ -45,7 +58,7 @@ function expectImageMessagePreserved(messages: AgentMessage[], errorMessage: str
 
   expect(pruned).toBeNull();
   const content = expectArrayMessageContent(messages[0], errorMessage);
-  expect(content[1]).toMatchObject({ type: "image", data: "abc" });
+  expectContentBlock(content[1], { type: "image", data: "abc" });
 }
 
 function oldEnoughTail(): AgentMessage[] {
@@ -117,13 +130,13 @@ describe("pruneProcessedHistoryImages", () => {
         PRUNED_HISTORY_MEDIA_REFERENCE_MARKER,
       ].join("\n"),
     );
-    expect(content[1]).toMatchObject({ type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
+    expectContentBlock(content[1], { type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
     const originalContent = expectArrayMessageContent(
       messages[0],
       "expected original user content",
     );
     expect(originalContent[0]?.text).toContain("[media attached: media://inbound/old.png]");
-    expect(originalContent[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(originalContent[1], { type: "image", data: "abc" });
   });
 
   it("scrubs old media attachment markers from string content without image blocks", () => {
@@ -207,7 +220,7 @@ describe("pruneProcessedHistoryImages", () => {
     expect(pruned).toBeNull();
     const content = expectArrayMessageContent(messages[0], "expected user array content");
     expect(content[0]?.text).toBe("recent [media attached: media://inbound/current.png]");
-    expect(content[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(content[1], { type: "image", data: "abc" });
   });
 
   it("does not count multiple assistant messages from one tool loop as separate turns", () => {
@@ -249,7 +262,7 @@ describe("pruneProcessedHistoryImages", () => {
     expect(pruned).toBeNull();
     const content = expectArrayMessageContent(messages[0], "expected user array content");
     expect(content).toHaveLength(2);
-    expect(content[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(content[1], { type: "image", data: "abc" });
   });
 
   it("prunes image blocks from toolResult messages older than 3 assistant turns", () => {
@@ -292,16 +305,16 @@ describe("pruneProcessedHistoryImages", () => {
     const pruned = expectPrunedMessages(messages);
 
     const oldContent = expectArrayMessageContent(pruned[0], "expected old user content");
-    expect(oldContent[1]).toMatchObject({ type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
+    expectContentBlock(oldContent[1], { type: "text", text: PRUNED_HISTORY_IMAGE_MARKER });
 
     const recentContent = expectArrayMessageContent(pruned[6], "expected recent user content");
-    expect(recentContent[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(recentContent[1], { type: "image", data: "abc" });
 
     const originalOldContent = expectArrayMessageContent(
       messages[0],
       "expected original old user content",
     );
-    expect(originalOldContent[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(originalOldContent[1], { type: "image", data: "abc" });
   });
 
   it("does not change messages when no assistant turn exists", () => {
@@ -359,7 +372,7 @@ describe("installHistoryImagePruneContextTransform", () => {
       "expected replay user array content",
     );
     expect(replayContent[0]?.text).toBe(`stale ${PRUNED_HISTORY_MEDIA_REFERENCE_MARKER}`);
-    expect(replayContent[1]).toMatchObject({
+    expectContentBlock(replayContent[1], {
       type: "text",
       text: PRUNED_HISTORY_IMAGE_MARKER,
     });
@@ -368,7 +381,7 @@ describe("installHistoryImagePruneContextTransform", () => {
       "expected original transformed content",
     );
     expect(originalContent[0]?.text).toContain("media://inbound/old.png");
-    expect(originalContent[1]).toMatchObject({ type: "image", data: "abc" });
+    expectContentBlock(originalContent[1], { type: "image", data: "abc" });
 
     restore();
     expect(agent.transformContext).toBe(originalTransformContext);