From 6a458ef29e150b42e1f55ff8e01722728b11b218 Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <hi@obviy.us>
Date: Sun, 15 Mar 2026 12:13:23 +0530
Subject: [PATCH] fix: harden compaction timeout follow-ups

---
 ...bedded-runner.compaction-safety-timeout.test.ts | 14 ++++++++++++++
 .../compaction-safety-timeout.ts                   |  7 ++++++-
 src/agents/pi-embedded-runner/run/attempt.ts       |  6 +++++-
 .../run/compaction-timeout.test.ts                 | 10 ++++++++++
 .../pi-embedded-runner/run/compaction-timeout.ts   |  7 +++++++
 5 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts b/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
index a44359c78da..1898373cfc9 100644
--- a/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
+++ b/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
@@ -76,6 +76,20 @@ describe("compactWithSafetyTimeout", () => {
     expect(onCancel).toHaveBeenCalledTimes(1);
     expect(vi.getTimerCount()).toBe(0);
   });
+
+  it("ignores onCancel errors and still rejects with the timeout", async () => {
+    vi.useFakeTimers();
+    const compactPromise = compactWithSafetyTimeout(() => new Promise<never>(() => {}), 30, {
+      onCancel: () => {
+        throw new Error("abortCompaction failed");
+      },
+    });
+    const timeoutAssertion = expect(compactPromise).rejects.toThrow("Compaction timed out");
+
+    await vi.advanceTimersByTimeAsync(30);
+    await timeoutAssertion;
+    expect(vi.getTimerCount()).toBe(0);
+  });
 });
 
 describe("resolveCompactionTimeoutMs", () => {
diff --git a/src/agents/pi-embedded-runner/compaction-safety-timeout.ts b/src/agents/pi-embedded-runner/compaction-safety-timeout.ts
index f50a112300a..bd15368ee2a 100644
--- a/src/agents/pi-embedded-runner/compaction-safety-timeout.ts
+++ b/src/agents/pi-embedded-runner/compaction-safety-timeout.ts
@@ -37,7 +37,12 @@ export async function compactWithSafetyTimeout<T>(
       return;
     }
     canceled = true;
-    opts?.onCancel?.();
+    try {
+      opts?.onCancel?.();
+    } catch {
+      // Best-effort cancellation hook. Keep the timeout/abort path intact even
+      // if the underlying compaction cancel operation throws.
+    }
   };
 
   return await withTimeout(
diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
index 105797c865a..ef5a63cdcd1 100644
--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -132,6 +132,7 @@ import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.
 import { waitForCompactionRetryWithAggregateTimeout } from "./compaction-retry-aggregate-timeout.js";
 import {
   resolveRunTimeoutDuringCompaction,
+  resolveRunTimeoutWithCompactionGraceMs,
   selectCompactionTimeoutSnapshot,
   shouldFlagCompactionTimeout,
 } from "./compaction-timeout.js";
@@ -1708,7 +1709,10 @@ export async function runEmbeddedAttempt(
     const sessionLock = await acquireSessionWriteLock({
       sessionFile: params.sessionFile,
       maxHoldMs: resolveSessionLockMaxHoldFromTimeout({
-        timeoutMs: params.timeoutMs,
+        timeoutMs: resolveRunTimeoutWithCompactionGraceMs({
+          runTimeoutMs: params.timeoutMs,
+          compactionTimeoutMs: resolveCompactionTimeoutMs(params.config),
+        }),
       }),
     });
 
diff --git a/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts b/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
index 5da781c615d..3853e0ebd25 100644
--- a/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
+++ b/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
@@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest";
 import { castAgentMessage } from "../../test-helpers/agent-message-fixtures.js";
 import {
   resolveRunTimeoutDuringCompaction,
+  resolveRunTimeoutWithCompactionGraceMs,
   selectCompactionTimeoutSnapshot,
   shouldFlagCompactionTimeout,
 } from "./compaction-timeout.js";
@@ -62,6 +63,15 @@ describe("compaction-timeout helpers", () => {
     ).toBe("abort");
   });
 
+  it("adds one compaction grace window to the run timeout budget", () => {
+    expect(
+      resolveRunTimeoutWithCompactionGraceMs({
+        runTimeoutMs: 600_000,
+        compactionTimeoutMs: 900_000,
+      }),
+    ).toBe(1_500_000);
+  });
+
   it("uses pre-compaction snapshot when compaction timeout occurs", () => {
     const pre = [castAgentMessage({ role: "assistant", content: "pre" })] as const;
     const current = [castAgentMessage({ role: "assistant", content: "current" })] as const;
diff --git a/src/agents/pi-embedded-runner/run/compaction-timeout.ts b/src/agents/pi-embedded-runner/run/compaction-timeout.ts
index 11a88455c96..97e1dfff4e3 100644
--- a/src/agents/pi-embedded-runner/run/compaction-timeout.ts
+++ b/src/agents/pi-embedded-runner/run/compaction-timeout.ts
@@ -24,6 +24,13 @@ export function resolveRunTimeoutDuringCompaction(params: {
   return params.graceAlreadyUsed ? "abort" : "extend";
 }
 
+export function resolveRunTimeoutWithCompactionGraceMs(params: {
+  runTimeoutMs: number;
+  compactionTimeoutMs: number;
+}): number {
+  return params.runTimeoutMs + params.compactionTimeoutMs;
+}
+
 export type SnapshotSelectionParams = {
   timedOutDuringCompaction: boolean;
   preCompactionSnapshot: AgentMessage[] | null;