test: cover history image injection

src/agents/pi-embedded-runner/run/attempt.test.ts

@@ -0,0 +1,55 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import type { ImageContent } from "@mariozechner/pi-ai";
+import { describe, expect, it } from "vitest";
+
+import { injectHistoryImagesIntoMessages } from "./attempt.js";
+
+describe("injectHistoryImagesIntoMessages", () => {
+  const image: ImageContent = { type: "image", data: "abc", mimeType: "image/png" };
+
+  it("injects history images and converts string content", () => {
+    const messages: AgentMessage[] = [
+      {
+        role: "user",
+        content: "See /tmp/photo.png",
+      } as AgentMessage,
+    ];
+
+    const didMutate = injectHistoryImagesIntoMessages(messages, new Map([[0, [image]]]));
+
+    expect(didMutate).toBe(true);
+    expect(Array.isArray(messages[0]?.content)).toBe(true);
+    const content = messages[0]?.content as Array<{ type: string; text?: string; data?: string }>;
+    expect(content).toHaveLength(2);
+    expect(content[0]?.type).toBe("text");
+    expect(content[1]).toMatchObject({ type: "image", data: "abc" });
+  });
+
+  it("avoids duplicating existing image content", () => {
+    const messages: AgentMessage[] = [
+      {
+        role: "user",
+        content: [{ type: "text", text: "See /tmp/photo.png" }, { ...image }],
+      } as AgentMessage,
+    ];
+
+    const didMutate = injectHistoryImagesIntoMessages(messages, new Map([[0, [image]]]));
+
+    expect(didMutate).toBe(false);
+    expect((messages[0]?.content as unknown[]).length).toBe(2);
+  });
+
+  it("ignores non-user messages and out-of-range indices", () => {
+    const messages: AgentMessage[] = [
+      {
+        role: "assistant",
+        content: "noop",
+      } as AgentMessage,
+    ];
+
+    const didMutate = injectHistoryImagesIntoMessages(messages, new Map([[1, [image]]]));
+
+    expect(didMutate).toBe(false);
+    expect(messages[0]?.content).toBe("noop");
+  });
+});

src/agents/pi-embedded-runner/run/attempt.ts

@@ -77,6 +77,50 @@ import { MAX_IMAGE_BYTES } from "../../../media/constants.js";
 import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js";
 import { detectAndLoadPromptImages } from "./images.js";
 
+export function injectHistoryImagesIntoMessages(
+  messages: AgentMessage[],
+  historyImagesByIndex: Map<number, ImageContent[]>,
+): boolean {
+  if (historyImagesByIndex.size === 0) return false;
+  let didMutate = false;
+
+  for (const [msgIndex, images] of historyImagesByIndex) {
+    // Bounds check: ensure index is valid before accessing
+    if (msgIndex < 0 || msgIndex >= messages.length) continue;
+    const msg = messages[msgIndex];
+    if (msg && msg.role === "user") {
+      // Convert string content to array format if needed
+      if (typeof msg.content === "string") {
+        msg.content = [{ type: "text", text: msg.content }];
+        didMutate = true;
+      }
+      if (Array.isArray(msg.content)) {
+        // Check for existing image content to avoid duplicates across turns
+        const existingImageData = new Set(
+          msg.content
+            .filter(
+              (c): c is ImageContent =>
+                c != null &&
+                typeof c === "object" &&
+                c.type === "image" &&
+                typeof c.data === "string",
+            )
+            .map((c) => c.data),
+        );
+        for (const img of images) {
+          // Only add if this image isn't already in the message
+          if (!existingImageData.has(img.data)) {
+            msg.content.push(img);
+            didMutate = true;
+          }
+        }
+      }
+    }
+  }
+
+  return didMutate;
+}
+
 export async function runEmbeddedAttempt(
   params: EmbeddedRunAttemptParams,
 ): Promise<EmbeddedRunAttemptResult> {
@@ -626,46 +670,14 @@ export async function runEmbeddedAttempt(
 
           // Inject history images into their original message positions.
           // This ensures the model sees images in context (e.g., "compare to the first image").
-          if (imageResult.historyImagesByIndex.size > 0) {
+          const didMutate = injectHistoryImagesIntoMessages(
-            let didMutate = false;
+            activeSession.messages,
-            for (const [msgIndex, images] of imageResult.historyImagesByIndex) {
+            imageResult.historyImagesByIndex,
-              // Bounds check: ensure index is valid before accessing
-              if (msgIndex < 0 || msgIndex >= activeSession.messages.length) continue;
-              const msg = activeSession.messages[msgIndex];
-              if (msg && msg.role === "user") {
-                // Convert string content to array format if needed
-                if (typeof msg.content === "string") {
-                  msg.content = [{ type: "text", text: msg.content }];
-                  didMutate = true;
-                }
-                if (Array.isArray(msg.content)) {
-                  // Check for existing image content to avoid duplicates across turns
-                  const existingImageData = new Set(
-                    msg.content
-                      .filter(
-                        (c): c is ImageContent =>
-                          c != null &&
-                          typeof c === "object" &&
-                          c.type === "image" &&
-                          typeof c.data === "string",
-                      )
-                      .map((c) => c.data),
           );
-                  for (const img of images) {
-                    // Only add if this image isn't already in the message
-                    if (!existingImageData.has(img.data)) {
-                      msg.content.push(img);
-                      didMutate = true;
-                    }
-                  }
-                }
-              }
-            }
           if (didMutate) {
             // Persist message mutations (e.g., injected history images) so we don't re-scan/reload.
             activeSession.agent.replaceMessages(activeSession.messages);
           }
-          }
 
           cacheTrace?.recordStage("prompt:images", {
             prompt: effectivePrompt,