From 7f071a6a8eaa90f92b847e47c7a0cf2422523d52 Mon Sep 17 00:00:00 2001
From: Ted Li <tl2493@columbia.edu>
Date: Fri, 10 Apr 2026 12:07:57 -0700
Subject: [PATCH] Agents: persist bootstrap marker after clean sessions_yield

---
 ...t.spawn-workspace.bootstrap-marker.test.ts | 78 +++++++++++++++++++
 .../run/attempt.thread-helpers.ts             | 19 +++++
 src/agents/pi-embedded-runner/run/attempt.ts  | 15 ++--
 3 files changed, 106 insertions(+), 6 deletions(-)
 create mode 100644 src/agents/pi-embedded-runner/run/attempt.spawn-workspace.bootstrap-marker.test.ts

diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.bootstrap-marker.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.bootstrap-marker.test.ts
new file mode 100644
index 00000000000..a6a59d64f3f
--- /dev/null
+++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.bootstrap-marker.test.ts
@@ -0,0 +1,78 @@
+import { describe, expect, it } from "vitest";
+import { shouldPersistCompletedBootstrapTurn } from "./attempt.thread-helpers.js";
+
+describe("runEmbeddedAttempt bootstrap completion marker", () => {
+  it("keeps marker persistence enabled for clean sessions_yield exits", () => {
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: true,
+        promptError: undefined,
+        aborted: false,
+        yieldAborted: true,
+        timedOutDuringCompaction: false,
+        compactionOccurredThisAttempt: false,
+      }),
+    ).toBe(true);
+  });
+
+  it("skips marker persistence when recording is disabled", () => {
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: false,
+        promptError: undefined,
+        aborted: false,
+        yieldAborted: false,
+        timedOutDuringCompaction: false,
+        compactionOccurredThisAttempt: false,
+      }),
+    ).toBe(false);
+  });
+
+  it("skips marker persistence when the attempt aborted", () => {
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: true,
+        promptError: undefined,
+        aborted: true,
+        yieldAborted: false,
+        timedOutDuringCompaction: false,
+        compactionOccurredThisAttempt: false,
+      }),
+    ).toBe(false);
+  });
+
+  it("skips marker persistence for prompt errors and compaction-side outcomes", () => {
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: true,
+        promptError: new Error("prompt failed"),
+        aborted: false,
+        yieldAborted: false,
+        timedOutDuringCompaction: false,
+        compactionOccurredThisAttempt: false,
+      }),
+    ).toBe(false);
+
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: true,
+        promptError: undefined,
+        aborted: false,
+        yieldAborted: false,
+        timedOutDuringCompaction: true,
+        compactionOccurredThisAttempt: false,
+      }),
+    ).toBe(false);
+
+    expect(
+      shouldPersistCompletedBootstrapTurn({
+        shouldRecordCompletedBootstrapTurn: true,
+        promptError: undefined,
+        aborted: false,
+        yieldAborted: false,
+        timedOutDuringCompaction: false,
+        compactionOccurredThisAttempt: true,
+      }),
+    ).toBe(false);
+  });
+});
diff --git a/src/agents/pi-embedded-runner/run/attempt.thread-helpers.ts b/src/agents/pi-embedded-runner/run/attempt.thread-helpers.ts
index ef48ce83ec3..8fafed3a557 100644
--- a/src/agents/pi-embedded-runner/run/attempt.thread-helpers.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.thread-helpers.ts
@@ -88,3 +88,22 @@ export function appendAttemptCacheTtlIfNeeded(params: {
   });
   return true;
 }
+
+export function shouldPersistCompletedBootstrapTurn(params: {
+  shouldRecordCompletedBootstrapTurn: boolean;
+  promptError: unknown;
+  aborted: boolean;
+  yieldAborted: boolean;
+  timedOutDuringCompaction: boolean;
+  compactionOccurredThisAttempt: boolean;
+}): boolean {
+  if (!params.shouldRecordCompletedBootstrapTurn || params.promptError || params.aborted) {
+    return false;
+  }
+  if (params.timedOutDuringCompaction || params.compactionOccurredThisAttempt) {
+    return false;
+  }
+  // Intentionally allow clean sessions_yield exits here so continuation-skip
+  // can treat the next relay/user turn as an existing bootstrap session.
+  return true;
+}
diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
index ddb011762c4..6e1d0bc8c54 100644
--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -200,6 +200,7 @@ import {
   appendAttemptCacheTtlIfNeeded,
   composeSystemPromptWithHookContext,
   resolveAttemptSpawnWorkspaceDir,
+  shouldPersistCompletedBootstrapTurn,
   shouldUseOpenAIWebSocketTransport,
 } from "./attempt.thread-helpers.js";
 import {
@@ -2169,12 +2170,14 @@ export async function runEmbeddedAttempt(
         }
 
         if (
-          shouldRecordCompletedBootstrapTurn &&
-          !promptError &&
-          !aborted &&
-          !yieldAborted &&
-          !timedOutDuringCompaction &&
-          !compactionOccurredThisAttempt
+          shouldPersistCompletedBootstrapTurn({
+            shouldRecordCompletedBootstrapTurn,
+            promptError,
+            aborted,
+            yieldAborted,
+            timedOutDuringCompaction,
+            compactionOccurredThisAttempt,
+          })
         ) {
           try {
             sessionManager.appendCustomEntry(FULL_BOOTSTRAP_COMPLETED_CUSTOM_TYPE, {