From 8aad43681b5047beb858bfab3ee4dbfd5ceaecc3 Mon Sep 17 00:00:00 2001
From: Gustavo Madeira Santana <gumadeiras@gmail.com>
Date: Sat, 14 Feb 2026 17:36:41 -0500
Subject: [PATCH] refactor(agent): extract tested compaction safety timeout

---
 ...d-runner.compaction-safety-timeout.test.ts | 45 +++++++++++++++++++
 src/agents/pi-embedded-runner/compact.ts      | 18 ++------
 .../compaction-safety-timeout.ts              | 10 +++++
 src/node-host/with-timeout.ts                 |  1 +
 4 files changed, 59 insertions(+), 15 deletions(-)
 create mode 100644 src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
 create mode 100644 src/agents/pi-embedded-runner/compaction-safety-timeout.ts

diff --git a/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts b/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
new file mode 100644
index 00000000000..31906dd733e
--- /dev/null
+++ b/src/agents/pi-embedded-runner.compaction-safety-timeout.test.ts
@@ -0,0 +1,45 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  compactWithSafetyTimeout,
+  EMBEDDED_COMPACTION_TIMEOUT_MS,
+} from "./pi-embedded-runner/compaction-safety-timeout.js";
+
+describe("compactWithSafetyTimeout", () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("rejects with timeout when compaction never settles", async () => {
+    vi.useFakeTimers();
+    const compactPromise = compactWithSafetyTimeout(() => new Promise<never>(() => {}));
+    const timeoutAssertion = expect(compactPromise).rejects.toThrow("Compaction timed out");
+
+    await vi.advanceTimersByTimeAsync(EMBEDDED_COMPACTION_TIMEOUT_MS);
+    await timeoutAssertion;
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("returns result and clears timer when compaction settles first", async () => {
+    vi.useFakeTimers();
+    const compactPromise = compactWithSafetyTimeout(
+      () => new Promise<string>((resolve) => setTimeout(() => resolve("ok"), 10)),
+      30,
+    );
+
+    await vi.advanceTimersByTimeAsync(10);
+    await expect(compactPromise).resolves.toBe("ok");
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("preserves compaction errors and clears timer", async () => {
+    vi.useFakeTimers();
+    const error = new Error("provider exploded");
+
+    await expect(
+      compactWithSafetyTimeout(async () => {
+        throw error;
+      }, 30),
+    ).rejects.toBe(error);
+    expect(vi.getTimerCount()).toBe(0);
+  });
+});
diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts
index 6a52f8eca09..408808312d5 100644
--- a/src/agents/pi-embedded-runner/compact.ts
+++ b/src/agents/pi-embedded-runner/compact.ts
@@ -57,6 +57,7 @@ import {
   type SkillSnapshot,
 } from "../skills.js";
 import { resolveTranscriptPolicy } from "../transcript-policy.js";
+import { compactWithSafetyTimeout } from "./compaction-safety-timeout.js";
 import { buildEmbeddedExtensionPaths } from "./extensions.js";
 import {
   logToolSchemasForGoogle,
@@ -632,22 +633,9 @@ export async function compactEmbeddedPiSessionDirect(
         }
 
         const compactStartedAt = Date.now();
-        const COMPACT_TIMEOUT_MS = 300_000; // 5 minutes safety timeout
-        let compactTimer: ReturnType<typeof setTimeout> | undefined;
-        const result = await Promise.race([
+        const result = await compactWithSafetyTimeout(() =>
           session.compact(params.customInstructions),
-          new Promise<never>((_, reject) => {
-            compactTimer = setTimeout(
-              () => reject(new Error("Compaction timed out")),
-              COMPACT_TIMEOUT_MS,
-            );
-            compactTimer.unref?.();
-          }),
-        ]).finally(() => {
-          if (compactTimer) {
-            clearTimeout(compactTimer);
-          }
-        });
+        );
         // Estimate tokens after compaction by summing token estimates for remaining messages
         let tokensAfter: number | undefined;
         try {
diff --git a/src/agents/pi-embedded-runner/compaction-safety-timeout.ts b/src/agents/pi-embedded-runner/compaction-safety-timeout.ts
new file mode 100644
index 00000000000..689aa9a931f
--- /dev/null
+++ b/src/agents/pi-embedded-runner/compaction-safety-timeout.ts
@@ -0,0 +1,10 @@
+import { withTimeout } from "../../node-host/with-timeout.js";
+
+export const EMBEDDED_COMPACTION_TIMEOUT_MS = 300_000;
+
+export async function compactWithSafetyTimeout<T>(
+  compact: () => Promise<T>,
+  timeoutMs: number = EMBEDDED_COMPACTION_TIMEOUT_MS,
+): Promise<T> {
+  return await withTimeout(() => compact(), timeoutMs, "Compaction");
+}
diff --git a/src/node-host/with-timeout.ts b/src/node-host/with-timeout.ts
index 07ea1415493..1acf525a79e 100644
--- a/src/node-host/with-timeout.ts
+++ b/src/node-host/with-timeout.ts
@@ -14,6 +14,7 @@ export async function withTimeout<T>(
   const abortCtrl = new AbortController();
   const timeoutError = new Error(`${label ?? "request"} timed out`);
   const timer = setTimeout(() => abortCtrl.abort(timeoutError), resolved);
+  timer.unref?.();
 
   let abortListener: (() => void) | undefined;
   const abortPromise: Promise<never> = abortCtrl.signal.aborted