From 85c1c59c5f887eb2e69aac7acaee63c3eed43737 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Mon, 20 Apr 2026 14:53:31 +0100
Subject: [PATCH] refactor: share message content block visitor

---
 .../run/attempt.tool-call-normalization.ts    | 19 +++++--------------
 src/plugin-sdk/provider-stream-shared.ts      | 17 ++++-------------
 src/shared/message-content-blocks.ts          | 18 ++++++++++++++++++
 3 files changed, 27 insertions(+), 27 deletions(-)
 create mode 100644 src/shared/message-content-blocks.ts

diff --git a/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts b/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
index 953f68a99da..747f71be6d8 100644
--- a/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
@@ -1,5 +1,6 @@
 import type { AgentMessage, StreamFn } from "@mariozechner/pi-agent-core";
 import { streamSimple } from "@mariozechner/pi-ai";
+import { visitObjectContentBlocks } from "../../../shared/message-content-blocks.js";
 import { normalizeLowercaseStringOrEmpty } from "../../../shared/string-coerce.js";
 import { validateAnthropicTurns, validateGeminiTurns } from "../../pi-embedded-helpers.js";
 import { sanitizeToolUseResultPairing } from "../../session-transcript-repair.js";
@@ -586,20 +587,10 @@ function trimWhitespaceFromToolCallNamesInMessage(
   message: unknown,
   allowedToolNames?: Set<string>,
 ): void {
-  if (!message || typeof message !== "object") {
-    return;
-  }
-  const content = (message as { content?: unknown }).content;
-  if (!Array.isArray(content)) {
-    return;
-  }
-  for (const block of content) {
-    if (!block || typeof block !== "object") {
-      continue;
-    }
+  visitObjectContentBlocks(message, (block) => {
     const typedBlock = block as { type?: unknown; name?: unknown; id?: unknown };
     if (!isToolCallBlockType(typedBlock.type)) {
-      continue;
+      return;
     }
     const rawId = typeof typedBlock.id === "string" ? typedBlock.id : undefined;
     if (typeof typedBlock.name === "string") {
@@ -607,13 +598,13 @@ function trimWhitespaceFromToolCallNamesInMessage(
       if (normalized !== typedBlock.name) {
         typedBlock.name = normalized;
       }
-      continue;
+      return;
     }
     const inferred = inferToolNameFromToolCallId(rawId, allowedToolNames);
     if (inferred) {
       typedBlock.name = inferred;
     }
-  }
+  });
   normalizeToolCallIdsInMessage(message);
 }
 
diff --git a/src/plugin-sdk/provider-stream-shared.ts b/src/plugin-sdk/provider-stream-shared.ts
index 8c24a32673f..b8e538142d4 100644
--- a/src/plugin-sdk/provider-stream-shared.ts
+++ b/src/plugin-sdk/provider-stream-shared.ts
@@ -1,6 +1,7 @@
 import type { StreamFn } from "@mariozechner/pi-agent-core";
 import { streamSimple } from "@mariozechner/pi-ai";
 import { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js";
+import { visitObjectContentBlocks } from "../shared/message-content-blocks.js";
 import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js";
 import type { ProviderWrapStreamFnContext } from "./plugin-entry.js";
 
@@ -64,25 +65,15 @@ export function decodeHtmlEntitiesInObject(value: unknown): unknown {
 }
 
 function decodeToolCallArgumentsHtmlEntitiesInMessage(message: unknown): void {
-  if (!message || typeof message !== "object") {
-    return;
-  }
-  const content = (message as { content?: unknown }).content;
-  if (!Array.isArray(content)) {
-    return;
-  }
-  for (const block of content) {
-    if (!block || typeof block !== "object") {
-      continue;
-    }
+  visitObjectContentBlocks(message, (block) => {
     const typedBlock = block as { type?: unknown; arguments?: unknown };
     if (typedBlock.type !== "toolCall" || !typedBlock.arguments) {
-      continue;
+      return;
     }
     if (typeof typedBlock.arguments === "object") {
       typedBlock.arguments = decodeHtmlEntitiesInObject(typedBlock.arguments);
     }
-  }
+  });
 }
 
 function wrapStreamDecodeToolCallArgumentHtmlEntities(
diff --git a/src/shared/message-content-blocks.ts b/src/shared/message-content-blocks.ts
new file mode 100644
index 00000000000..4a0224e6ff4
--- /dev/null
+++ b/src/shared/message-content-blocks.ts
@@ -0,0 +1,18 @@
+export function visitObjectContentBlocks(
+  message: unknown,
+  visitor: (block: Record<string, unknown>) => void,
+): void {
+  if (!message || typeof message !== "object") {
+    return;
+  }
+  const content = (message as { content?: unknown }).content;
+  if (!Array.isArray(content)) {
+    return;
+  }
+  for (const block of content) {
+    if (!block || typeof block !== "object") {
+      continue;
+    }
+    visitor(block as Record<string, unknown>);
+  }
+}