diff --git a/pkg/agentruntest/agent_run_supervisor_test.go b/pkg/agentruntest/agent_run_supervisor_test.go index 736edc7d5..21e69b3c6 100644 --- a/pkg/agentruntest/agent_run_supervisor_test.go +++ b/pkg/agentruntest/agent_run_supervisor_test.go @@ -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, diff --git a/pkg/coredata/agent_run.go b/pkg/coredata/agent_run.go index 073e9c908..2009640c5 100644 --- a/pkg/coredata/agent_run.go +++ b/pkg/coredata/agent_run.go @@ -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 diff --git a/pkg/coredata/migrations/20260424T120000Z.sql b/pkg/coredata/migrations/20260424T120000Z.sql index f3e6a1c6b..aa15ac008 100644 --- a/pkg/coredata/migrations/20260424T120000Z.sql +++ b/pkg/coredata/migrations/20260424T120000Z.sql @@ -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, diff --git a/pkg/probo/agent_run_handler.go b/pkg/probo/agent_run_handler.go new file mode 100644 index 000000000..bcaee782a --- /dev/null +++ b/pkg/probo/agent_run_handler.go @@ -0,0 +1,284 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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)) + } +} diff --git a/pkg/probo/agent_run_supervisor.go b/pkg/probo/agent_run_supervisor.go index ac1883f68..8c71780a8 100644 --- a/pkg/probo/agent_run_supervisor.go +++ b/pkg/probo/agent_run_supervisor.go @@ -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) }