From ea1e933b29a2d6a1d88752d3565add35237489fb Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 18 Apr 2026 21:24:53 +0100
Subject: [PATCH] refactor: share sessions spawn attachment checks

---
 .../run/attempt.tool-call-normalization.ts    | 28 ++---------------
 src/agents/tool-call-id.ts                    | 22 ++------------
 src/agents/tool-call-shared.ts                | 30 +++++++++++++++++++
 3 files changed, 34 insertions(+), 46 deletions(-)

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 307223e59ea..544973df384 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
@@ -2,15 +2,13 @@ import type { AgentMessage, StreamFn } from "@mariozechner/pi-agent-core";
 import { streamSimple } from "@mariozechner/pi-ai";
 import { normalizeLowercaseStringOrEmpty } from "../../../shared/string-coerce.js";
 import { validateAnthropicTurns, validateGeminiTurns } from "../../pi-embedded-helpers.js";
-import {
-  isRedactedSessionsSpawnAttachment,
-  sanitizeToolUseResultPairing,
-} from "../../session-transcript-repair.js";
+import { sanitizeToolUseResultPairing } from "../../session-transcript-repair.js";
 import {
   extractToolCallsFromAssistant,
   sanitizeToolCallIdsForCloudCodeAssist,
   type ToolCallIdMode,
 } from "../../tool-call-id.js";
+import { hasUnredactedSessionsSpawnAttachments } from "../../tool-call-shared.js";
 import { normalizeToolName } from "../../tool-policy.js";
 import { shouldAllowProviderOwnedThinkingReplay } from "../../transcript-policy.js";
 import type { TranscriptPolicy } from "../../transcript-policy.js";
@@ -255,28 +253,6 @@ function isThinkingLikeReplayBlock(block: unknown): boolean {
   return type === "thinking" || type === "redacted_thinking";
 }
 
-function hasUnredactedSessionsSpawnAttachments(block: ReplayToolCallBlock): boolean {
-  const rawName = typeof block.name === "string" ? block.name.trim() : "";
-  if (normalizeLowercaseStringOrEmpty(rawName) !== "sessions_spawn") {
-    return false;
-  }
-  for (const payload of [block.arguments, block.input]) {
-    if (!payload || typeof payload !== "object") {
-      continue;
-    }
-    const attachments = (payload as { attachments?: unknown }).attachments;
-    if (!Array.isArray(attachments)) {
-      continue;
-    }
-    for (const attachment of attachments) {
-      if (!isRedactedSessionsSpawnAttachment(attachment)) {
-        return true;
-      }
-    }
-  }
-  return false;
-}
-
 function isReplaySafeThinkingTurn(content: unknown[], allowedToolNames?: Set<string>): boolean {
   const seenToolCallIds = new Set<string>();
   for (const block of content) {
diff --git a/src/agents/tool-call-id.ts b/src/agents/tool-call-id.ts
index 1d0b9f87578..54a06ff2142 100644
--- a/src/agents/tool-call-id.ts
+++ b/src/agents/tool-call-id.ts
@@ -1,9 +1,8 @@
 import { createHash } from "node:crypto";
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
-import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js";
 import {
+  hasUnredactedSessionsSpawnAttachments,
   isAllowedToolCallName,
-  isRedactedSessionsSpawnAttachment,
   normalizeAllowedToolNames,
 } from "./tool-call-shared.js";
 
@@ -118,24 +117,7 @@ function toolCallNeedsReplayMutation(block: ReplaySafeToolCallBlock): boolean {
   if (rawName && rawName !== trimmedName) {
     return true;
   }
-  if (normalizeLowercaseStringOrEmpty(trimmedName) !== "sessions_spawn") {
-    return false;
-  }
-  for (const payload of [block.arguments, block.input]) {
-    if (!payload || typeof payload !== "object") {
-      continue;
-    }
-    const attachments = (payload as { attachments?: unknown }).attachments;
-    if (!Array.isArray(attachments)) {
-      continue;
-    }
-    for (const attachment of attachments) {
-      if (!isRedactedSessionsSpawnAttachment(attachment)) {
-        return true;
-      }
-    }
-  }
-  return false;
+  return hasUnredactedSessionsSpawnAttachments(block);
 }
 
 function isReplaySafeThinkingAssistantMessage(
diff --git a/src/agents/tool-call-shared.ts b/src/agents/tool-call-shared.ts
index 4a45c595334..a1de09a9c29 100644
--- a/src/agents/tool-call-shared.ts
+++ b/src/agents/tool-call-shared.ts
@@ -65,3 +65,33 @@ export function isRedactedSessionsSpawnAttachment(item: unknown): boolean {
   }
   return true;
 }
+
+export type SessionsSpawnAttachmentToolCallBlock = {
+  name?: unknown;
+  input?: unknown;
+  arguments?: unknown;
+};
+
+export function hasUnredactedSessionsSpawnAttachments(
+  block: SessionsSpawnAttachmentToolCallBlock,
+): boolean {
+  const rawName = typeof block.name === "string" ? block.name.trim() : "";
+  if (normalizeLowercaseStringOrEmpty(rawName) !== "sessions_spawn") {
+    return false;
+  }
+  for (const payload of [block.arguments, block.input]) {
+    if (!payload || typeof payload !== "object") {
+      continue;
+    }
+    const attachments = (payload as { attachments?: unknown }).attachments;
+    if (!Array.isArray(attachments)) {
+      continue;
+    }
+    for (const attachment of attachments) {
+      if (!isRedactedSessionsSpawnAttachment(attachment)) {
+        return true;
+      }
+    }
+  }
+  return false;
+}