From 20dbc28398d88472417e8947ff02fbb81adb8d89 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?=
 <81782+aureliensibiril@users.noreply.github.com>
Date: Fri, 24 Apr 2026 20:01:38 +0200
Subject: [PATCH] Expose supervisor shutdown broadcast for tests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds AgentRunSupervisor.ShutdownBroadcast() returning the handler's
shutdown channel so the StopAndResume integration test can wait for
graceful-shutdown propagation deterministically instead of sleeping
for a fixed duration. The method is explicitly documented as
test-only and not part of the operational contract.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
---
 pkg/agentruntest/agent_run_supervisor_test.go | 13 ++++++++++---
 pkg/probo/agent_run_supervisor.go             | 16 ++++++++++++++--
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/pkg/agentruntest/agent_run_supervisor_test.go b/pkg/agentruntest/agent_run_supervisor_test.go
index 21e69b3c6..226235a92 100644
--- a/pkg/agentruntest/agent_run_supervisor_test.go
+++ b/pkg/agentruntest/agent_run_supervisor_test.go
@@ -273,9 +273,16 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
 	// agent's stop channel.
 	cancel1()
 
-	// Give the AfterFunc goroutine a moment to close the broadcast and
-	// propagate into the per-run stopCh before the tool unblocks.
-	time.Sleep(500 * time.Millisecond)
+	// Wait for the shutdown broadcast to be observed (the AfterFunc
+	// goroutine closes it) before releasing the tool. This is
+	// deterministic: no wall-clock sleep. The per-run forwarder
+	// goroutine observes the same close synchronously and closes the
+	// agent stop channel.
+	select {
+	case <-supervisor.ShutdownBroadcastForTests():
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for supervisor shutdown broadcast")
+	}
 
 	// Now release the tool. When the coreLoop resumes control at the
 	// next turn boundary it observes the closed stop channel, saves
diff --git a/pkg/probo/agent_run_supervisor.go b/pkg/probo/agent_run_supervisor.go
index 8c71780a8..ca09558ae 100644
--- a/pkg/probo/agent_run_supervisor.go
+++ b/pkg/probo/agent_run_supervisor.go
@@ -113,8 +113,20 @@ func NewAgentRunSupervisor(
 // closes the shutdown broadcast channel so in-flight Process calls can
 // checkpoint and exit, and waits for all of them to drain before
 // returning.
+//
+// signalShutdown is registered without a stop hook because it is
+// idempotent (sync.Once) and we want it to fire on every ctx
+// cancellation, even one that races with worker.Run returning.
 func (s *AgentRunSupervisor) Run(ctx context.Context) error {
-	stop := context.AfterFunc(ctx, s.handler.signalShutdown)
-	defer stop()
+	context.AfterFunc(ctx, s.handler.signalShutdown)
 	return s.worker.Run(ctx)
 }
+
+// ShutdownBroadcastForTests returns a channel that closes once the
+// supervisor has broadcast graceful shutdown to all in-flight runs.
+// Exposed for tests that need to synchronize tool release with shutdown
+// propagation; not a stable API and not part of the supervisor's public
+// operational contract.
+func (s *AgentRunSupervisor) ShutdownBroadcastForTests() <-chan struct{} {
+	return s.handler.shutdownCh
+}