Rebuild AgentRunSupervisor on go.gearno.de/kit/worker
The supervisor was a hand-rolled polling, semaphore, and wait-group loop predating the project's adoption of the shared worker kit. Two sibling workers in pkg/probo already use the kit, and go-worker.md documents it as the project convention. This commit introduces agentRunHandler, which implements worker.Handler[coredata.AgentRun] and worker.StaleRecoverer, and reduces AgentRunSupervisor to a thin wrapper that owns the handler plus a worker.Worker and bridges ctx cancellation into a handler- level shutdown broadcast via context.AfterFunc. The agent stop channel is now closed by a per-Process forwarder goroutine when the broadcast fires, so in-flight runs checkpoint at the next turn boundary and drain through wg.Wait before Run returns. The stop_requested column, struct field, supporting SQL, and the LoadRunningStopRequestedIDs function are removed end-to-end. None of it was ever wired to an external surface; it existed purely to let the supervisor find runs the operator wanted to halt. With the kit handling the polling cadence and the AfterFunc bridging shutdown, per-row flagging is dead weight. The supervisor's public API (NewAgentRunSupervisor, Run, the With* option helpers, and the error sentinels) stays intact so probod.go needs no change. The integration test now triggers stop by cancelling the supervisor context, which is the actual production path through SIGTERM rather than a synthetic DB flag. Prometheus counters and OTel spans labelled worker="agent-run-supervisor" come for free. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
This commit is contained in:
@@ -180,7 +180,6 @@ func TestAgentRunSupervisor_PicksUpAndCompletes(t *testing.T) {
|
||||
assert.NotNil(t, completed.Result)
|
||||
assert.Nil(t, completed.Checkpoint, "checkpoint should be cleared after completion")
|
||||
assert.Nil(t, completed.ErrorMessage)
|
||||
assert.False(t, completed.StopRequested)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -191,7 +190,8 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
client := agentruntest.PGClient(t)
|
||||
store := coredata.NewPGCheckpointer(client)
|
||||
|
||||
// The tool blocks until signaled, giving us time to set stop_requested.
|
||||
// The tool blocks until signaled, giving us time to trigger graceful
|
||||
// shutdown via the supervisor context while the agent is mid-turn.
|
||||
toolReady := make(chan struct{})
|
||||
toolRelease := make(chan struct{})
|
||||
|
||||
@@ -267,29 +267,19 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
running := agentruntest.LoadAgentRun(t, client, run.ID)
|
||||
assert.Equal(t, coredata.AgentRunStatusRunning, running.Status)
|
||||
|
||||
// Set stop_requested in the database WHILE the tool is still blocked.
|
||||
// The supervisor polls for this on each tick and signals the run's
|
||||
// stop channel.
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
_, err := conn.Exec(
|
||||
ctx,
|
||||
"UPDATE agent_runs SET stop_requested = true WHERE id = $1",
|
||||
run.ID.String(),
|
||||
)
|
||||
return err
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
// Trigger graceful shutdown of the supervisor: context.AfterFunc
|
||||
// registered in Run() fires signalShutdown, closing the shutdown
|
||||
// broadcast channel; the per-run forwarder goroutine closes the
|
||||
// agent's stop channel.
|
||||
cancel1()
|
||||
|
||||
// Give the supervisor at least one tick to poll stop requests and
|
||||
// close the run's stop channel before the tool finishes.
|
||||
time.Sleep(1 * time.Second)
|
||||
// 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)
|
||||
|
||||
// Now release the tool. After completion the coreLoop saves an
|
||||
// incremental checkpoint and checks the stop signal at the next
|
||||
// turn boundary — it should already be closed.
|
||||
// Now release the tool. When the coreLoop resumes control at the
|
||||
// next turn boundary it observes the closed stop channel, saves
|
||||
// the suspension checkpoint, and returns SuspendedError.
|
||||
close(toolRelease)
|
||||
|
||||
// Wait for the checkpoint to appear. The supervisor leaves the row
|
||||
@@ -306,10 +296,6 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
"checkpoint should be saved after stop",
|
||||
)
|
||||
|
||||
// Stop the first supervisor.
|
||||
cancel1()
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Verify checkpoint content.
|
||||
cp, err := store.Load(context.Background(), run.ID.String())
|
||||
require.NoError(t, err)
|
||||
@@ -327,7 +313,6 @@ func TestAgentRunSupervisor_StopAndResume(t *testing.T) {
|
||||
ctx,
|
||||
`UPDATE agent_runs
|
||||
SET status = 'PENDING',
|
||||
stop_requested = false,
|
||||
started_at = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
@@ -633,7 +618,6 @@ func TestAgentRunSupervisor_SIGTERM(t *testing.T) {
|
||||
_, err := conn.Exec(ctx, `
|
||||
UPDATE agent_runs
|
||||
SET status = 'PENDING',
|
||||
stop_requested = false,
|
||||
started_at = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
|
||||
Reference in New Issue
Block a user