Drop agent-run lease and add approval resume
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user