From 53747733b3fddb7c08fba03134d393631047373e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?=
 <81782+aureliensibiril@users.noreply.github.com>
Date: Mon, 27 Apr 2026 09:51:56 +0200
Subject: [PATCH] Map supervisor shutdown onto run ctx cancellation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Process now has the per-run forwarder goroutine call
cancelRun(ErrSuspendForCheckpoint) when h.shutdownCh closes,
rather than closing a separate stopCh and embedding it via
agent.WithStopSignal. The agent loop's new ctx-cancel = graceful
suspend contract covers the rest. h.shutdownCh and signalShutdown
stay as the supervisor-level broadcast (still observable through
ShutdownBroadcastForTests).

The lease-loss path keeps its existing cancelRun call; under the
new contract that triggers a best-effort save before executeRun
detects ErrAgentRunLeaseLost and skips the row commit, which is
race-safe because Worker B can only claim the row after stale
recovery — by then our save has long landed.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
---
 pkg/probo/agent_run_handler.go | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/pkg/probo/agent_run_handler.go b/pkg/probo/agent_run_handler.go
index 33db6cb12..adc6fd10f 100644
--- a/pkg/probo/agent_run_handler.go
+++ b/pkg/probo/agent_run_handler.go
@@ -45,6 +45,13 @@ type agentRunHandler struct {
 var (
 	_ worker.Handler[coredata.AgentRun] = (*agentRunHandler)(nil)
 	_ worker.StaleRecoverer             = (*agentRunHandler)(nil)
+
+	// ErrSuspendForCheckpoint is the cancel cause used when the
+	// supervisor asks an in-flight run to gracefully suspend so it can
+	// checkpoint and exit. The agent loop sees ctx.Err() at its next
+	// turn boundary and returns *SuspendedError; executeRun treats
+	// that outcome as a graceful exit (no row-status commit).
+	ErrSuspendForCheckpoint = errors.New("agent run: graceful suspend requested")
 )
 
 // Claim loads the next pending agent run, marks it RUNNING with a lease
@@ -87,10 +94,11 @@ func (h *agentRunHandler) Claim(ctx context.Context) (coredata.AgentRun, error)
 	return run, nil
 }
 
-// Process executes a single agent run. It spawns a heartbeat goroutine to
-// renew the lease while the run is active, and a forwarder goroutine that
-// bridges the handler-level shutdown signal to the run's agent stop
-// channel so the agent checkpoints cleanly at the next turn boundary.
+// Process executes a single agent run. It spawns a heartbeat goroutine
+// that renews the lease while the run is active, and a forwarder
+// goroutine that converts the handler-level shutdown broadcast into a
+// per-run ctx cancellation so the agent loop checkpoints cleanly at
+// its next turn boundary.
 //
 // The returned error mirrors the run outcome so the worker kit's
 // task metrics and OTel span status reflect actual agent failures.
@@ -101,13 +109,12 @@ func (h *agentRunHandler) Process(ctx context.Context, run coredata.AgentRun) er
 	runCtx, cancelRun := context.WithCancelCause(ctx)
 	defer cancelRun(nil)
 
-	stopCh := make(chan struct{})
 	forwarderDone := make(chan struct{})
 	defer close(forwarderDone)
 	go func() {
 		select {
 		case <-h.shutdownCh:
-			close(stopCh)
+			cancelRun(ErrSuspendForCheckpoint)
 		case <-forwarderDone:
 		}
 	}()
@@ -116,7 +123,6 @@ func (h *agentRunHandler) Process(ctx context.Context, run coredata.AgentRun) er
 	defer cancelHeartbeat()
 	go h.heartbeatLease(heartbeatCtx, run.ID.String(), cancelRun)
 
-	runCtx = agent.WithStopSignal(runCtx, stopCh)
 	return h.executeRun(runCtx, &run)
 }