From 6996c0f330cfa03b5889ee2a02ddfd1597d1bfcb Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Wed, 21 Jan 2026 18:40:09 +0000 Subject: [PATCH] test: cover history image injection --- .../pi-embedded-runner/run/attempt.test.ts | 55 ++++++++++++ src/agents/pi-embedded-runner/run/attempt.ts | 90 +++++++++++-------- 2 files changed, 106 insertions(+), 39 deletions(-) create mode 100644 src/agents/pi-embedded-runner/run/attempt.test.ts diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts new file mode 100644 index 000000000..d87cabd1d --- /dev/null +++ b/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"); + }); +}); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index b54ff4087..6a1bd3978 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/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, +): 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 { @@ -626,45 +670,13 @@ 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) { - let didMutate = false; - for (const [msgIndex, images] of 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); - } + const didMutate = injectHistoryImagesIntoMessages( + activeSession.messages, + imageResult.historyImagesByIndex, + ); + 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", {