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

Drop agent-run lease and add approval resume

Browse Source 
The worker leaned on a lease plus a heartbeat goroutine and a stale
recovery sweep to reclaim runs from crashed workers. That machinery
raced with long LLM and tool calls and conflated graceful stops with
failures. Remove the lease columns, heartbeat, and stale recovery, and
rely on FOR UPDATE SKIP LOCKED for single-claim plus explicit state
transitions: a graceful suspend returns the run to PENDING and a crash
now leaves it RUNNING for manual recovery.

Treat an approval interruption as a known stop that parks the run in
AWAITING_APPROVAL, and add SubmitApproval to merge human decisions into
the checkpoint and requeue the run to PENDING. The decisions must cover
exactly the pending approvals, since a missing one would resume as an
implicit denial. Expose this through the submitAgentRunApproval
mutation.

Signed-off-by: Bryan Frimin <bryan@probo.com>
This commit is contained in:
Bryan Frimin
2026-06-08 14:51:12 +02:00
parent 98a8d90391
commit c14bacb157
9 changed files with 398 additions and 363 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
pkg/agentrun/errors.go Normal file
View File

@@ -0,0 +1,35 @@
// 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 agentrun
import (
"errors"
)
var (
// ErrAgentRunNotFound is returned when the target agent run does not
// exist. It wraps coredata.ErrResourceNotFound so callers depend on the
// agentrun API rather than the underlying data layer.
ErrAgentRunNotFound = errors.New("agent run not found")
// ErrNotAwaitingApproval is returned when an approval decision is
// submitted for a run that is not currently parked in AWAITING_APPROVAL.
ErrNotAwaitingApproval = errors.New("agent run is not awaiting approval")
// ErrApprovalDecisionsMismatch is returned when the submitted decisions
// do not cover exactly the run's pending approvals. It shields callers
// from the agent package's internal mismatch error.
ErrApprovalDecisionsMismatch = errors.New("approval decisions do not match the run's pending approvals")
)

254
pkg/agentrun/handler.go
View File

@@ -32,28 +32,25 @@ import (
)
type handler struct {
pg *pg.Client
store *coredata.PGCheckpointer
registry agent.AgentRegistry
logger *log.Logger
leaseDuration time.Duration
shutdownCh chan struct{}
shutdownOnce sync.Once
pg *pg.Client
store *coredata.PGCheckpointer
registry agent.AgentRegistry
logger *log.Logger
shutdownCh chan struct{}
shutdownOnce sync.Once
}
var (
_ worker.Handler[coredata.AgentRun] = (*handler)(nil)
_ worker.StaleRecoverer = (*handler)(nil)
)
var _ worker.Handler[coredata.AgentRun] = (*handler)(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.
// Claim loads the next pending agent run and marks it RUNNING. When no
// work is available it returns worker.ErrNoTask so the kit backs off
// until the next tick. The FOR UPDATE SKIP LOCKED select guarantees only
// one worker claims a given row; there is no lease, so a worker that
// crashes mid-run leaves the row RUNNING for manual recovery.
func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) {
var (
run = coredata.AgentRun{}
now = time.Now()
leaseExpiresAt = now.Add(h.leaseDuration)
run = coredata.AgentRun{}
now = time.Now()
)
if err := h.pg.WithTx(
@@ -65,8 +62,6 @@ func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) {
run.Status = coredata.AgentRunStatusRunning
run.StartedAt = &now
run.LeaseExpiresAt = &leaseExpiresAt
run.LeaseGeneration++
run.UpdatedAt = now
if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil {
@@ -86,23 +81,20 @@ func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) {
return run, nil
}
// 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.
// Process executes a single agent run. It spawns 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.
// nil is returned for both successful runs and graceful exits
// (lease loss, infrastructure suspension) where the row state is
// already consistent.
// The returned error mirrors the run outcome so the worker kit's task
// metrics and OTel span status reflect actual agent failures. nil is
// returned for successful runs and for known stops (graceful suspend,
// awaiting approval) where the row was already committed to a resumable
// state.
func (h *handler) Process(ctx context.Context, run coredata.AgentRun) error {
runCtx, cancelRun := context.WithCancelCause(ctx)
defer cancelRun(nil)
leaseGeneration := run.LeaseGeneration
forwarderDone := make(chan struct{})
defer close(forwarderDone)
@@ -114,27 +106,7 @@ func (h *handler) Process(ctx context.Context, run coredata.AgentRun) error {
}
}()
heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx)
defer cancelHeartbeat()
go h.heartbeatLease(heartbeatCtx, run.ID.String(), leaseGeneration, cancelRun)
return h.executeRun(runCtx, &run, leaseGeneration)
}
// RecoverStale resets agent runs whose worker lease has expired back to
// PENDING so a fresh worker can pick them up on the next cycle.
func (h *handler) 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
return h.executeRun(runCtx, &run)
}
// signalShutdown closes the handler-level shutdown broadcast channel. All
@@ -145,57 +117,6 @@ func (h *handler) signalShutdown() {
h.shutdownOnce.Do(func() { close(h.shutdownCh) })
}
func (h *handler) heartbeatLease(
ctx context.Context,
runID string,
leaseGeneration int64,
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,
leaseGeneration,
expiresAt,
)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrLeaseLost
}
return nil
},
); err != nil {
h.logger.ErrorCtx(ctx, "cannot heartbeat agent run lease", log.Error(err))
if errors.Is(err, ErrLeaseLost) {
cancelRun(ErrLeaseLost)
} else {
cancelRun(fmt.Errorf("%w: %w", ErrHeartbeatFailed, err))
}
return
}
}
}
}
const (
// errorMessageMaxLen caps the error string persisted to the
// agent_runs.error_message column. Raw tool or LLM errors can embed
@@ -219,29 +140,8 @@ func sanitizeError(err error) string {
return msg[:cut] + "…"
}
type leasedCheckpointer struct {
store *coredata.PGCheckpointer
leaseGeneration int64
}
func (s leasedCheckpointer) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error {
return s.store.SaveForLease(ctx, runID, cp, s.leaseGeneration)
}
func (s leasedCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) {
return s.store.Load(ctx, runID)
}
func (h *handler) executeRun(
ctx context.Context,
run *coredata.AgentRun,
leaseGeneration int64,
) error {
func (h *handler) executeRun(ctx context.Context, run *coredata.AgentRun) error {
runID := run.ID.String()
checkpointer := leasedCheckpointer{
store: h.store,
leaseGeneration: leaseGeneration,
}
var (
result *agent.Result
@@ -250,7 +150,7 @@ func (h *handler) executeRun(
if run.Checkpoint != nil {
h.logger.InfoCtx(ctx, "resuming agent run", log.String("run_id", runID))
result, runErr = agent.Restore(ctx, checkpointer, runID, h.registry)
result, runErr = agent.Restore(ctx, h.store, runID, h.registry)
} else {
h.logger.InfoCtx(ctx, "starting agent run", log.String("run_id", runID))
@@ -265,55 +165,25 @@ func (h *handler) executeRun(
result, runErr = a.Run(
ctx,
inputMsgs,
agent.WithCheckpointer(checkpointer, runID),
agent.WithCheckpointer(h.store, runID),
)
}
}
}
// Heartbeat loss: another worker may have taken over. Do not commit
// any status — stale recovery will handle the row. Surface the cause
// so the worker kit logs and traces a failure for this attempt.
if cause := context.Cause(ctx); errors.Is(cause, ErrLeaseLost) || errors.Is(cause, ErrHeartbeatFailed) {
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 cause
}
// 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. This
// is not a failure from the worker kit's perspective.
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 nil
}
}
now := time.Now()
run.UpdatedAt = now
run.StartedAt = nil
run.LeaseExpiresAt = nil
run.Result = nil
run.ErrorMessage = nil
if runErr == 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))
runErr = fmt.Errorf("cannot marshal agent run result: %w", err)
} else {
run.Result = data
@@ -321,18 +191,34 @@ func (h *handler) executeRun(
}
}
// Known stops are not failures: the agent loop already saved a
// checkpoint before returning. Graceful suspend returns the run to
// PENDING so any worker resumes it from the checkpoint; an approval
// interruption parks it in AWAITING_APPROVAL until an approval
// decision requeues it. Anything else is a genuine failure.
if runErr != nil {
run.Status = coredata.AgentRunStatusFailed
run.Result = nil
switch {
case isType[*agent.SuspendedError](runErr):
run.Status = coredata.AgentRunStatusPending
runErr = nil
h.logger.ErrorCtx(
context.WithoutCancel(ctx),
"agent run failed",
log.String("run_id", runID),
log.Error(runErr),
)
msg := sanitizeError(runErr)
run.ErrorMessage = &msg
case isType[*agent.InterruptedError](runErr):
run.Status = coredata.AgentRunStatusAwaitingApproval
runErr = nil
default:
run.Status = coredata.AgentRunStatusFailed
run.Result = nil
h.logger.ErrorCtx(
context.WithoutCancel(ctx),
"agent run failed",
log.String("run_id", runID),
log.Error(runErr),
)
msg := sanitizeError(runErr)
run.ErrorMessage = &msg
}
}
commitCtx := context.WithoutCancel(ctx)
@@ -340,13 +226,19 @@ func (h *handler) executeRun(
if err := h.pg.WithTx(
commitCtx,
func(ctx context.Context, tx pg.Tx) error {
rowsAffected, err := coredata.CommitAgentRunResult(ctx, tx, run, leaseGeneration)
rowsAffected, err := coredata.CommitAgentRunResult(ctx, tx, run)
if err != nil {
return err
}
if rowsAffected == 0 {
return ErrLeaseLost
h.logger.WarnCtx(
ctx,
"agent run no longer RUNNING at commit; discarding result",
log.String("run_id", runID),
)
return nil
}
if run.Status == coredata.AgentRunStatusCompleted {
@@ -358,16 +250,6 @@ func (h *handler) executeRun(
return nil
},
); err != nil {
if errors.Is(err, ErrLeaseLost) {
h.logger.WarnCtx(
commitCtx,
"agent run lost lease before commit; discarding stale completion",
log.String("run_id", runID),
)
return nil
}
h.logger.ErrorCtx(commitCtx, "cannot commit agent run status", log.Error(err))
return fmt.Errorf("cannot commit agent run status: %w", err)
@@ -375,3 +257,9 @@ func (h *handler) executeRun(
return runErr
}
func isType[T error](err error) bool {
_, ok := errors.AsType[T](err)
return ok
}

64
pkg/agentrun/service.go
View File

@@ -16,9 +16,12 @@ package agentrun
import (
"context"
"errors"
"fmt"
"time"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/agent"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/page"
@@ -86,6 +89,67 @@ func (s *Service) ListForOrganizationID(
return page.NewPage(runs, cursor), nil
}
// SubmitApproval records human approval decisions for a run parked in
// AWAITING_APPROVAL and requeues it to PENDING so a worker resumes it.
// decisions is keyed by pending tool-call ID and must cover exactly the
// run's pending approvals (a missing decision would be treated as an
// implicit denial on resume, so partial submissions are rejected). The
// refreshed run is returned.
func (s *Service) SubmitApproval(
ctx context.Context,
scope coredata.Scoper,
agentRunID gid.GID,
decisions map[string]agent.ApprovalResult,
) (*coredata.AgentRun, error) {
run := &coredata.AgentRun{}
err := s.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := run.LoadByIDForUpdate(ctx, tx, scope, agentRunID); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return ErrAgentRunNotFound
}
return fmt.Errorf("cannot load agent run: %w", err)
}
if run.Status != coredata.AgentRunStatusAwaitingApproval {
return ErrNotAwaitingApproval
}
if run.Checkpoint == nil {
return fmt.Errorf("agent run %s has no checkpoint", agentRunID)
}
checkpoint, err := agent.MergeApprovalDecisions(run.Checkpoint, decisions)
if err != nil {
if errors.Is(err, agent.ErrApprovalDecisionsMismatch) {
return ErrApprovalDecisionsMismatch
}
return fmt.Errorf("cannot merge approval decisions: %w", err)
}
run.Checkpoint = checkpoint
run.Status = coredata.AgentRunStatusPending
run.StartedAt = nil
run.UpdatedAt = time.Now()
if err := run.RequeueForApprovalResume(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot requeue agent run for approval resume: %w", err)
}
return nil
},
)
if err != nil {
return nil, err
}
return run, nil
}
func (s *Service) CountForOrganizationID(
ctx context.Context,
scope coredata.Scoper,

27
pkg/agentrun/worker.go
View File

@@ -16,7 +16,6 @@ package agentrun
import (
"context"
"errors"
"time"
"go.gearno.de/kit/log"
@@ -36,16 +35,10 @@ type (
workerConfig struct {
interval time.Duration
leaseDuration time.Duration
maxConcurrency int
}
)
var (
ErrHeartbeatFailed = errors.New("agent run heartbeat failed")
ErrLeaseLost = errors.New("agent run lease lost")
)
func WithWorkerInterval(d time.Duration) WorkerOption {
return func(c *workerConfig) {
if d > 0 {
@@ -54,14 +47,6 @@ func WithWorkerInterval(d time.Duration) WorkerOption {
}
}
func WithWorkerLeaseDuration(d time.Duration) WorkerOption {
return func(c *workerConfig) {
if d > 0 {
c.leaseDuration = d
}
}
}
func WithWorkerMaxConcurrency(n int) WorkerOption {
return func(c *workerConfig) {
if n > 0 {
@@ -79,7 +64,6 @@ func NewWorker(
) *Worker {
cfg := workerConfig{
interval: 10 * time.Second,
leaseDuration: 5 * time.Minute,
maxConcurrency: 5,
}
@@ -88,12 +72,11 @@ func NewWorker(
}
h := &handler{
pg: pgClient,
store: store,
registry: registry,
logger: logger,
leaseDuration: cfg.leaseDuration,
shutdownCh: make(chan struct{}),
pg: pgClient,
store: store,
registry: registry,
logger: logger,
shutdownCh: make(chan struct{}),
}
w := worker.New(