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

@@ -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: