rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rebuild AgentRunSupervisor on go.gearno.de/kit/worker

Browse Source 
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:
Aurélien Sibiril
2026-04-24 16:09:39 +02:00
parent fb88318c54
commit 6887294c9e
5 changed files with 340 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
pkg/agentruntest/agent_run_supervisor_test.go
View File

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

42
pkg/coredata/agent_run.go
View File

@@ -41,7 +41,6 @@ type (
InputMessages json.RawMessage `db:"input_messages"`
Result json.RawMessage `db:"result"`
ErrorMessage *string `db:"error_message"`
StopRequested bool `db:"stop_requested"`
StartedAt *time.Time `db:"started_at"`
LeaseOwner *string `db:"lease_owner"`
LeaseExpiresAt *time.Time `db:"lease_expires_at"`
@@ -111,7 +110,6 @@ SELECT
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -164,7 +162,6 @@ SELECT
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -219,7 +216,6 @@ SELECT
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -317,7 +313,6 @@ RETURNING
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -362,7 +357,6 @@ SET
status = @status,
result = @result,
error_message = @error_message,
stop_requested = @stop_requested,
started_at = @started_at,
lease_owner = @lease_owner,
lease_expires_at = @lease_expires_at,
@@ -379,7 +373,6 @@ RETURNING
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -394,7 +387,6 @@ RETURNING
"status": e.Status,
"result": e.Result,
"error_message": e.ErrorMessage,
"stop_requested": e.StopRequested,
"started_at": e.StartedAt,
"lease_owner": e.LeaseOwner,
"lease_expires_at": e.LeaseExpiresAt,
@@ -481,7 +473,6 @@ RETURNING
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -526,7 +517,6 @@ SELECT
input_messages,
result,
error_message,
stop_requested,
started_at,
lease_owner,
lease_expires_at,
@@ -536,7 +526,6 @@ FROM
agent_runs
WHERE
status = 'PENDING'
AND stop_requested = false
ORDER BY created_at ASC
LIMIT 1
FOR UPDATE SKIP LOCKED;
@@ -564,8 +553,7 @@ FOR UPDATE SKIP LOCKED;
// The worker refreshes lease_expires_at from a separate heartbeat goroutine,
// so a long LLM or tool call is not considered stale while the process is alive.
// Stale recovery returns rows to PENDING so the supervisor auto-resumes
// from checkpoint when one exists. User-requested suspension remains
// SUSPENDED and is resumed only by the Resume service method.
// from checkpoint when one exists.
func ResetStaleAgentRuns(ctx context.Context, conn pg.Querier) error {
q := `
UPDATE agent_runs
@@ -574,7 +562,6 @@ SET
started_at = NULL,
lease_owner = NULL,
lease_expires_at = NULL,
stop_requested = false,
updated_at = now()
WHERE
status = 'RUNNING'
@@ -624,33 +611,6 @@ WHERE
return tag.RowsAffected(), nil
}
// LoadRunningStopRequestedIDs returns IDs of running agent runs with
// stop_requested = true.
func LoadRunningStopRequestedIDs(ctx context.Context, conn pg.Querier) ([]string, error) {
q := `SELECT id FROM agent_runs WHERE status = 'RUNNING' AND stop_requested = true;`
rows, err := conn.Query(ctx, q)
if err != nil {
return nil, fmt.Errorf("cannot query stop-requested agent runs: %w", err)
}
defer rows.Close()
var ids []string
for rows.Next() {
var id string
if err := rows.Scan(&id); err != nil {
return nil, fmt.Errorf("cannot scan stop-requested agent run ID: %w", err)
}
ids = append(ids, id)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("cannot iterate stop-requested agent runs: %w", err)
}
return ids, nil
}
// PGCheckpointer implements agent.Checkpointer backed by the
// agent_runs table checkpoint column. It is supervisor-internal and
// intentionally uses raw run IDs with no tenant scope; public service/API

1
pkg/coredata/migrations/20260424T120000Z.sql
View File

@@ -22,7 +22,6 @@ CREATE TABLE agent_runs (
input_messages JSONB NOT NULL,
result JSONB,
error_message TEXT,
stop_requested BOOLEAN NOT NULL DEFAULT FALSE,
started_at TIMESTAMPTZ,
lease_owner TEXT,
lease_expires_at TIMESTAMPTZ,

284
pkg/probo/agent_run_handler.go Normal file
View File

@@ -0,0 +1,284 @@
// 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 probo
import (
"context"
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/llm"
)
type agentRunHandler struct {
pg *pg.Client
store *coredata.PGCheckpointer
registry agent.AgentRegistry
logger *log.Logger
leaseDuration time.Duration
workerID string
shutdownCh chan struct{}
shutdownOnce sync.Once
}
var (
_ worker.Handler[coredata.AgentRun] = (*agentRunHandler)(nil)
_ worker.StaleRecoverer = (*agentRunHandler)(nil)
)
// Claim loads the next pending agent run, marks it RUNNING with a lease
// owned by this worker, and returns the row. When no work is available it
// returns worker.ErrNoTask so the kit can back off until the next tick.
func (h *agentRunHandler) Claim(ctx context.Context) (coredata.AgentRun, error) {
var (
run = coredata.AgentRun{}
now = time.Now()
leaseOwner = h.workerID
leaseExpiresAt = now.Add(h.leaseDuration)
)
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
run.Status = coredata.AgentRunStatusRunning
run.StartedAt = &now
run.LeaseOwner = &leaseOwner
run.LeaseExpiresAt = &leaseExpiresAt
run.UpdatedAt = now
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update agent run: %w", err)
}
return nil
},
); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.AgentRun{}, worker.ErrNoTask
}
return coredata.AgentRun{}, err
}
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.
func (h *agentRunHandler) Process(ctx context.Context, run coredata.AgentRun) error {
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)
case <-forwarderDone:
}
}()
heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx)
defer cancelHeartbeat()
go h.heartbeatLease(heartbeatCtx, run.ID.String(), cancelRun)
runCtx = agent.WithStopSignal(runCtx, stopCh)
h.executeRun(runCtx, &run)
return nil
}
// RecoverStale resets agent runs whose worker lease has expired back to
// PENDING so a fresh supervisor can pick them up on the next cycle.
func (h *agentRunHandler) RecoverStale(ctx context.Context) error {
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleAgentRuns(ctx, conn)
},
); err != nil {
return fmt.Errorf("cannot reset stale agent runs: %w", err)
}
return nil
}
// signalShutdown closes the handler-level shutdown broadcast channel. All
// in-flight Process forwarder goroutines observe the close and propagate
// it to their per-run agent stop channels, letting agents checkpoint at
// the next turn boundary before Process returns.
func (h *agentRunHandler) signalShutdown() {
h.shutdownOnce.Do(func() { close(h.shutdownCh) })
}
func (h *agentRunHandler) heartbeatLease(
ctx context.Context,
runID string,
cancelRun context.CancelCauseFunc,
) {
ticker := time.NewTicker(h.leaseDuration / 3)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
expiresAt := time.Now().Add(h.leaseDuration)
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
rowsAffected, err := coredata.HeartbeatAgentRunLease(ctx, conn, runID, h.workerID, expiresAt)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrAgentRunLeaseLost
}
return nil
},
); err != nil {
h.logger.ErrorCtx(ctx, "cannot heartbeat agent run lease", log.Error(err))
if errors.Is(err, ErrAgentRunLeaseLost) {
cancelRun(ErrAgentRunLeaseLost)
} else {
cancelRun(fmt.Errorf("%w: %w", ErrAgentRunHeartbeatFailed, err))
}
return
}
}
}
}
func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun) {
runID := run.ID.String()
var (
result *agent.Result
runErr error
)
if run.Checkpoint != nil {
h.logger.InfoCtx(ctx, "resuming agent run", log.String("run_id", runID))
result, runErr = agent.Restore(ctx, h.store, runID, h.registry)
} else {
h.logger.InfoCtx(ctx, "starting agent run", log.String("run_id", runID))
a, err := h.registry.Agent(run.StartAgentName)
if err != nil {
runErr = fmt.Errorf("cannot resolve agent %q: %w", run.StartAgentName, err)
} else {
var inputMsgs []llm.Message
if err := json.Unmarshal(run.InputMessages, &inputMsgs); err != nil {
runErr = fmt.Errorf("cannot unmarshal input messages: %w", err)
} else {
result, runErr = a.RunWithOpts(
ctx,
inputMsgs,
agent.WithCheckpointer(h.store, runID),
)
}
}
}
// Heartbeat loss: another worker may have taken over. Do not commit
// any status — stale recovery will handle the row.
if cause := context.Cause(ctx); errors.Is(cause, ErrAgentRunLeaseLost) || errors.Is(cause, ErrAgentRunHeartbeatFailed) {
h.logger.WarnCtx(
context.WithoutCancel(ctx),
"agent run stopped after heartbeat failure; leaving status for stale recovery",
log.String("run_id", runID),
log.Error(cause),
)
return
}
// Infrastructure-triggered suspension (graceful shutdown): leave the
// row as RUNNING so stale recovery resets it to PENDING on restart.
// The checkpoint was already saved by coreLoop before returning
// SuspendedError, so Restore will pick up where it left off.
if runErr != nil {
if _, ok := errors.AsType[*agent.SuspendedError](runErr); ok {
h.logger.InfoCtx(
context.WithoutCancel(ctx),
"agent run suspended by infrastructure; leaving for stale recovery",
log.String("run_id", runID),
)
return
}
}
now := time.Now()
run.UpdatedAt = now
run.StartedAt = nil
run.LeaseOwner = nil
run.LeaseExpiresAt = nil
switch {
case runErr == nil:
run.Status = coredata.AgentRunStatusCompleted
if result != nil {
data, err := json.Marshal(result)
if err != nil {
h.logger.ErrorCtx(ctx, "cannot marshal agent run result", log.Error(err))
} else {
run.Result = data
}
}
default:
run.Status = coredata.AgentRunStatusFailed
msg := runErr.Error()
run.ErrorMessage = &msg
}
commitCtx := context.WithoutCancel(ctx)
if err := h.pg.WithTx(
commitCtx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return err
}
if run.Status == coredata.AgentRunStatusCompleted {
if err := run.ClearCheckpoint(ctx, tx, coredata.NewNoScope()); err != nil {
return err
}
}
return nil
},
); err != nil {
h.logger.ErrorCtx(commitCtx, "cannot commit agent run status", log.Error(err))
}
}

380
pkg/probo/agent_run_supervisor.go
View File

@@ -16,42 +16,30 @@ package probo
import (
"context"
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/crypto/uuid"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/llm"
)
type (
AgentRunSupervisor struct {
pg *pg.Client
store *coredata.PGCheckpointer
registry agent.AgentRegistry
logger *log.Logger
handler *agentRunHandler
worker *worker.Worker[coredata.AgentRun]
}
AgentRunSupervisorOption func(*agentRunSupervisorConfig)
agentRunSupervisorConfig struct {
interval time.Duration
leaseDuration time.Duration
maxConcurrency int
workerID string
mu sync.Mutex
running map[string]*runHandle // runID -> handle
}
// runHandle wraps a stop channel with a sync.Once to prevent double-close panics.
runHandle struct {
stopCh chan struct{}
once sync.Once
}
AgentRunSupervisorOption func(*AgentRunSupervisor)
)
var (
@@ -59,30 +47,26 @@ var (
ErrAgentRunLeaseLost = errors.New("agent run lease lost")
)
func (h *runHandle) stop() {
h.once.Do(func() { close(h.stopCh) })
}
func WithAgentRunSupervisorInterval(d time.Duration) AgentRunSupervisorOption {
return func(s *AgentRunSupervisor) {
return func(c *agentRunSupervisorConfig) {
if d > 0 {
s.interval = d
c.interval = d
}
}
}
func WithAgentRunSupervisorLeaseDuration(d time.Duration) AgentRunSupervisorOption {
return func(s *AgentRunSupervisor) {
return func(c *agentRunSupervisorConfig) {
if d > 0 {
s.leaseDuration = d
c.leaseDuration = d
}
}
}
func WithAgentRunSupervisorMaxConcurrency(n int) AgentRunSupervisorOption {
return func(s *AgentRunSupervisor) {
return func(c *agentRunSupervisorConfig) {
if n > 0 {
s.maxConcurrency = n
c.maxConcurrency = n
}
}
}
@@ -94,323 +78,43 @@ func NewAgentRunSupervisor(
logger *log.Logger,
opts ...AgentRunSupervisorOption,
) *AgentRunSupervisor {
s := &AgentRunSupervisor{
pg: pgClient,
store: store,
registry: registry,
logger: logger,
cfg := agentRunSupervisorConfig{
interval: 10 * time.Second,
leaseDuration: 5 * time.Minute,
maxConcurrency: 5,
workerID: uuid.MustNewV4().String(),
running: make(map[string]*runHandle),
}
for _, opt := range opts {
opt(s)
opt(&cfg)
}
return s
h := &agentRunHandler{
pg: pgClient,
store: store,
registry: registry,
logger: logger,
leaseDuration: cfg.leaseDuration,
workerID: uuid.MustNewV4().String(),
shutdownCh: make(chan struct{}),
}
w := worker.New(
"agent-run-supervisor",
h,
logger,
worker.WithInterval(cfg.interval),
worker.WithMaxConcurrency(cfg.maxConcurrency),
)
return &AgentRunSupervisor{handler: h, worker: w}
}
// Run starts the supervisor loop. It blocks until ctx is cancelled, then
// closes the shutdown broadcast channel so in-flight Process calls can
// checkpoint and exit, and waits for all of them to drain before
// returning.
func (s *AgentRunSupervisor) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, s.maxConcurrency)
ticker = time.NewTicker(s.interval)
)
defer ticker.Stop()
defer wg.Wait()
defer s.stopAll()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
s.recoverStaleRuns(nonCancelableCtx)
s.checkStopRequests(nonCancelableCtx)
for {
if err := s.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
s.logger.ErrorCtx(nonCancelableCtx, "cannot claim agent run", log.Error(err))
}
break
}
}
}
}
}
func (s *AgentRunSupervisor) processNext(
ctx context.Context,
sem chan struct{},
wg *sync.WaitGroup,
) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
run = coredata.AgentRun{}
now = time.Now()
leaseOwner = s.workerID
leaseExpiresAt = now.Add(s.leaseDuration)
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := s.pg.WithTx(
nonCancelableCtx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
run.Status = coredata.AgentRunStatusRunning
run.StartedAt = &now
run.LeaseOwner = &leaseOwner
run.LeaseExpiresAt = &leaseExpiresAt
run.UpdatedAt = now
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update agent run: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
}
handle := &runHandle{stopCh: make(chan struct{})}
s.mu.Lock()
s.running[run.ID.String()] = handle
s.mu.Unlock()
wg.Add(1)
go func(run coredata.AgentRun) {
defer wg.Done()
defer func() { <-sem }()
defer func() {
s.mu.Lock()
delete(s.running, run.ID.String())
s.mu.Unlock()
}()
runCtx, cancelRun := context.WithCancelCause(nonCancelableCtx)
defer cancelRun(nil)
heartbeatCtx, cancelHeartbeat := context.WithCancel(nonCancelableCtx)
defer cancelHeartbeat()
go s.heartbeatLease(heartbeatCtx, run.ID.String(), cancelRun)
runCtx = agent.WithStopSignal(runCtx, handle.stopCh)
s.executeRun(runCtx, &run)
}(run)
return nil
}
func (s *AgentRunSupervisor) heartbeatLease(
ctx context.Context,
runID string,
cancelRun context.CancelCauseFunc,
) {
ticker := time.NewTicker(s.leaseDuration / 3)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
expiresAt := time.Now().Add(s.leaseDuration)
if err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
rowsAffected, err := coredata.HeartbeatAgentRunLease(ctx, conn, runID, s.workerID, expiresAt)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrAgentRunLeaseLost
}
return nil
},
); err != nil {
s.logger.ErrorCtx(ctx, "cannot heartbeat agent run lease", log.Error(err))
if errors.Is(err, ErrAgentRunLeaseLost) {
cancelRun(ErrAgentRunLeaseLost)
} else {
cancelRun(fmt.Errorf("%w: %w", ErrAgentRunHeartbeatFailed, err))
}
return
}
}
}
}
func (s *AgentRunSupervisor) executeRun(ctx context.Context, run *coredata.AgentRun) {
runID := run.ID.String()
var (
result *agent.Result
runErr error
)
if run.Checkpoint != nil {
// Resume from checkpoint.
s.logger.InfoCtx(ctx, "resuming agent run", log.String("run_id", runID))
result, runErr = agent.Restore(ctx, s.store, runID, s.registry)
} else {
// Start fresh.
s.logger.InfoCtx(ctx, "starting agent run", log.String("run_id", runID))
a, err := s.registry.Agent(run.StartAgentName)
if err != nil {
runErr = fmt.Errorf("cannot resolve agent %q: %w", run.StartAgentName, err)
} else {
var inputMsgs []llm.Message
if err := json.Unmarshal(run.InputMessages, &inputMsgs); err != nil {
runErr = fmt.Errorf("cannot unmarshal input messages: %w", err)
} else {
result, runErr = a.RunWithOpts(
ctx,
inputMsgs,
agent.WithCheckpointer(s.store, runID),
)
}
}
}
// Heartbeat loss: another worker may have taken over. Do not commit
// any status — stale recovery will handle the row.
if cause := context.Cause(ctx); errors.Is(cause, ErrAgentRunLeaseLost) || errors.Is(cause, ErrAgentRunHeartbeatFailed) {
s.logger.WarnCtx(
context.WithoutCancel(ctx),
"agent run stopped after heartbeat failure; leaving status for stale recovery",
log.String("run_id", runID),
log.Error(cause),
)
return
}
// Infrastructure-triggered suspension (graceful shutdown): leave the
// row as RUNNING so stale recovery resets it to PENDING on restart.
// The checkpoint was already saved by coreLoop before returning
// SuspendedError, so Restore will pick up where it left off.
if runErr != nil {
if _, ok := errors.AsType[*agent.SuspendedError](runErr); ok {
s.logger.InfoCtx(
context.WithoutCancel(ctx),
"agent run suspended by infrastructure; leaving for stale recovery",
log.String("run_id", runID),
)
return
}
}
// Update run status based on outcome.
now := time.Now()
run.UpdatedAt = now
run.StartedAt = nil
run.LeaseOwner = nil
run.LeaseExpiresAt = nil
switch {
case runErr == nil:
run.Status = coredata.AgentRunStatusCompleted
if result != nil {
data, err := json.Marshal(result)
if err != nil {
s.logger.ErrorCtx(ctx, "cannot marshal agent run result", log.Error(err))
} else {
run.Result = data
}
}
run.StopRequested = false
default:
run.Status = coredata.AgentRunStatusFailed
msg := runErr.Error()
run.ErrorMessage = &msg
}
commitCtx := context.WithoutCancel(ctx)
if err := s.pg.WithTx(
commitCtx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return err
}
if run.Status == coredata.AgentRunStatusCompleted {
if err := run.ClearCheckpoint(ctx, tx, coredata.NewNoScope()); err != nil {
return err
}
}
return nil
},
); err != nil {
s.logger.ErrorCtx(commitCtx, "cannot commit agent run status", log.Error(err))
}
}
func (s *AgentRunSupervisor) checkStopRequests(ctx context.Context) {
var ids []string
if err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
var err error
ids, err = coredata.LoadRunningStopRequestedIDs(ctx, conn)
return err
},
); err != nil {
s.logger.ErrorCtx(ctx, "cannot check stop requests", log.Error(err))
return
}
for _, id := range ids {
s.mu.Lock()
if h, ok := s.running[id]; ok {
h.stop()
}
s.mu.Unlock()
}
}
func (s *AgentRunSupervisor) stopAll() {
s.mu.Lock()
defer s.mu.Unlock()
for _, h := range s.running {
h.stop()
}
}
func (s *AgentRunSupervisor) recoverStaleRuns(ctx context.Context) {
if err := s.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := coredata.ResetStaleAgentRuns(ctx, conn); err != nil {
return fmt.Errorf("cannot reset stale agent runs: %w", err)
}
return nil
},
); err != nil {
s.logger.ErrorCtx(ctx, "cannot recover stale agent runs", log.Error(err))
}
stop := context.AfterFunc(ctx, s.handler.signalShutdown)
defer stop()
return s.worker.Run(ctx)
}