From 5a8654969ecc59c23f02fd764888322ff9166809 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@probo.com>
Date: Sun, 7 Jun 2026 12:15:00 +0200
Subject: [PATCH] Make agent suspend tests release on the suspend signal

The single- and multi-level suspend/restore tests gated their slow
leaf tool on a manual release channel closed 50ms after cancel().
That sleep was a guess at how long the suspend signal takes to reach
the running sub-agent, so the post-tool turn boundary could observe
the release before cancellation and complete the run instead of
checkpointing, making the assertions timing-dependent.

Expose the per-run suspend signal through SuspendSignalFrom in an
export_test shim and have the leaf tools block on it directly. The
tool now returns only once the graceful-suspend signal has actually
propagated to its agent, so suspension is observed deterministically
without sleeps or release channels.

Signed-off-by: Bryan Frimin <bryan@probo.com>
---
 pkg/agent/agent_tool_test.go | 28 ++++++++++++++++++----------
 pkg/agent/export_test.go     | 21 +++++++++++++++++++++
 pkg/agent/run.go             |  9 ++++++---
 3 files changed, 45 insertions(+), 13 deletions(-)
 create mode 100644 pkg/agent/export_test.go

diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go
index 32bf6e632..463608a57 100644
--- a/pkg/agent/agent_tool_test.go
+++ b/pkg/agent/agent_tool_test.go
@@ -924,16 +924,22 @@ func TestAgentTool_Execute_SuspendAndRestoreSingleLevel(t *testing.T) {
 	store := newMemoryCheckpointer()
 
 	toolReady := make(chan struct{})
-	toolRelease := make(chan struct{})
 
 	var readyOnce sync.Once
 
 	slowTool := agent.FunctionTool[struct{}](
 		"slow_inner_work",
 		"Slow inner work",
-		func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+		func(ctx context.Context, _ struct{}) (agent.ToolResult, error) {
 			readyOnce.Do(func() { close(toolReady) })
-			<-toolRelease
+
+			// Release only once the graceful-suspend signal has
+			// reached this inner agent, so the post-tool turn
+			// boundary deterministically observes cancellation and
+			// checkpoints instead of completing the run.
+			if sig := agent.SuspendSignalFrom(ctx); sig != nil {
+				<-sig.Done()
+			}
 
 			return agent.ToolResult{Content: "inner tool done"}, nil
 		},
@@ -996,8 +1002,6 @@ func TestAgentTool_Execute_SuspendAndRestoreSingleLevel(t *testing.T) {
 	}
 
 	cancel()
-	time.Sleep(50 * time.Millisecond)
-	close(toolRelease)
 
 	select {
 	case err := <-errCh:
@@ -1043,16 +1047,22 @@ func TestAgentTool_Execute_SuspendAndRestoreMultiLevel(t *testing.T) {
 	store := newMemoryCheckpointer()
 
 	toolReady := make(chan struct{})
-	toolRelease := make(chan struct{})
 
 	var readyOnce sync.Once
 
 	slowTool := agent.FunctionTool[struct{}](
 		"slow_grandchild_work",
 		"Slow grandchild work",
-		func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+		func(ctx context.Context, _ struct{}) (agent.ToolResult, error) {
 			readyOnce.Do(func() { close(toolReady) })
-			<-toolRelease
+
+			// Release only once the graceful-suspend signal has
+			// propagated down to this grandchild agent, so the
+			// post-tool turn boundary deterministically observes
+			// cancellation and checkpoints instead of completing.
+			if sig := agent.SuspendSignalFrom(ctx); sig != nil {
+				<-sig.Done()
+			}
 
 			return agent.ToolResult{Content: "grandchild tool done"}, nil
 		},
@@ -1134,8 +1144,6 @@ func TestAgentTool_Execute_SuspendAndRestoreMultiLevel(t *testing.T) {
 	}
 
 	cancel()
-	time.Sleep(50 * time.Millisecond)
-	close(toolRelease)
 
 	select {
 	case err := <-errCh:
diff --git a/pkg/agent/export_test.go b/pkg/agent/export_test.go
new file mode 100644
index 000000000..d2da0ab5d
--- /dev/null
+++ b/pkg/agent/export_test.go
@@ -0,0 +1,21 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package agent
+
+// SuspendSignalFrom exposes the per-run graceful-suspend signal carried in
+// ctx to external tests. Tests use it to block in-flight leaf tools until
+// the suspend signal has actually reached the running (sub-)agent, making
+// suspend/restore assertions deterministic instead of timing-dependent.
+var SuspendSignalFrom = suspendSignalFrom
diff --git a/pkg/agent/run.go b/pkg/agent/run.go
index 2080afecc..3cfc7fe63 100644
--- a/pkg/agent/run.go
+++ b/pkg/agent/run.go
@@ -109,9 +109,12 @@ func withSuspendableToolContext(ctx context.Context, tool Tool) (context.Context
 	}
 
 	execCtx, cancel := context.WithCancelCause(ctx)
-	stop := context.AfterFunc(signal, func() {
-		cancel(context.Cause(signal))
-	})
+	stop := context.AfterFunc(
+		signal,
+		func() {
+			cancel(context.Cause(signal))
+		},
+	)
 
 	return execCtx, func() {
 		stop()