From c14bacb157686de493da9cac12788545f6fa3028 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Mon, 8 Jun 2026 14:51:12 +0200 Subject: [PATCH] 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 --- pkg/agent/approval.go | 53 ++++ pkg/agentrun/errors.go | 35 +++ pkg/agentrun/handler.go | 254 +++++------------- pkg/agentrun/service.go | 64 +++++ pkg/agentrun/worker.go | 27 +- pkg/coredata/agent_run.go | 238 ++++++---------- pkg/coredata/migrations/20260608T090000Z.sql | 25 ++ .../api/console/v1/agent_run_resolvers.go | 44 ++- .../api/console/v1/graphql/agent_run.graphql | 21 ++ 9 files changed, 398 insertions(+), 363 deletions(-) create mode 100644 pkg/agentrun/errors.go create mode 100644 pkg/coredata/migrations/20260608T090000Z.sql diff --git a/pkg/agent/approval.go b/pkg/agent/approval.go index 0271da553..396cc7205 100644 --- a/pkg/agent/approval.go +++ b/pkg/agent/approval.go @@ -16,10 +16,19 @@ package agent import ( "context" + "encoding/json" + "errors" + "fmt" "go.probo.inc/probo/pkg/llm" ) +// ErrApprovalDecisionsMismatch is returned by MergeApprovalDecisions when +// the supplied decisions do not cover exactly the checkpoint's pending +// approvals. A missing decision would resume as an implicit denial, so +// partial submissions are rejected. +var ErrApprovalDecisionsMismatch = errors.New("approval decisions do not match the checkpoint's pending approvals") + type ( ApprovalConfig struct { ToolNames []string @@ -38,6 +47,50 @@ type ( } ) +// MergeApprovalDecisions decodes an awaiting-approval checkpoint, records +// the human decisions into its ApprovalInput, and returns the re-encoded +// checkpoint ready to be persisted. decisions is keyed by pending +// tool-call ID and must cover exactly the checkpoint's pending approvals; +// otherwise ErrApprovalDecisionsMismatch is returned. +func MergeApprovalDecisions( + raw json.RawMessage, + decisions map[string]ApprovalResult, +) (json.RawMessage, error) { + var cp Checkpoint + if err := json.Unmarshal(raw, &cp); err != nil { + return nil, fmt.Errorf("cannot unmarshal checkpoint: %w", err) + } + + pending := make(map[string]struct{}, len(cp.PendingApprovals)) + for _, toolCall := range cp.PendingApprovals { + pending[toolCall.ID] = struct{}{} + } + + if len(decisions) != len(pending) { + return nil, ErrApprovalDecisionsMismatch + } + + for id := range decisions { + if _, ok := pending[id]; !ok { + return nil, ErrApprovalDecisionsMismatch + } + } + + if cp.ApprovalInput == nil { + cp.ApprovalInput = make(map[string]ApprovalResult, len(decisions)) + } + for id, decision := range decisions { + cp.ApprovalInput[id] = decision + } + + data, err := json.Marshal(&cp) + if err != nil { + return nil, fmt.Errorf("cannot marshal checkpoint: %w", err) + } + + return data, nil +} + func buildToolNameSet(names []string) map[string]struct{} { if len(names) == 0 { return nil diff --git a/pkg/agentrun/errors.go b/pkg/agentrun/errors.go new file mode 100644 index 000000000..2485fd4f7 --- /dev/null +++ b/pkg/agentrun/errors.go @@ -0,0 +1,35 @@ +// 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 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") +) diff --git a/pkg/agentrun/handler.go b/pkg/agentrun/handler.go index 86b41f26f..fb3c9fc6d 100644 --- a/pkg/agentrun/handler.go +++ b/pkg/agentrun/handler.go @@ -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 +} diff --git a/pkg/agentrun/service.go b/pkg/agentrun/service.go index 55d7f615e..78cdde662 100644 --- a/pkg/agentrun/service.go +++ b/pkg/agentrun/service.go @@ -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, diff --git a/pkg/agentrun/worker.go b/pkg/agentrun/worker.go index 4d78dc0d2..9683741b2 100644 --- a/pkg/agentrun/worker.go +++ b/pkg/agentrun/worker.go @@ -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( diff --git a/pkg/coredata/agent_run.go b/pkg/coredata/agent_run.go index c74913698..ae64a7146 100644 --- a/pkg/coredata/agent_run.go +++ b/pkg/coredata/agent_run.go @@ -35,19 +35,17 @@ type ( AgentRunStatus string AgentRun struct { - ID gid.GID `db:"id"` - OrganizationID gid.GID `db:"organization_id"` - StartAgentName string `db:"start_agent_name"` - Status AgentRunStatus `db:"status"` - Checkpoint json.RawMessage `db:"checkpoint"` - InputMessages json.RawMessage `db:"input_messages"` - Result json.RawMessage `db:"result"` - ErrorMessage *string `db:"error_message"` - StartedAt *time.Time `db:"started_at"` - LeaseExpiresAt *time.Time `db:"lease_expires_at"` - LeaseGeneration int64 `db:"lease_generation"` - CreatedAt time.Time `db:"created_at"` - UpdatedAt time.Time `db:"updated_at"` + ID gid.GID `db:"id"` + OrganizationID gid.GID `db:"organization_id"` + StartAgentName string `db:"start_agent_name"` + Status AgentRunStatus `db:"status"` + Checkpoint json.RawMessage `db:"checkpoint"` + InputMessages json.RawMessage `db:"input_messages"` + Result json.RawMessage `db:"result"` + ErrorMessage *string `db:"error_message"` + StartedAt *time.Time `db:"started_at"` + CreatedAt time.Time `db:"created_at"` + UpdatedAt time.Time `db:"updated_at"` } AgentRuns []*AgentRun @@ -178,8 +176,6 @@ SELECT result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at FROM @@ -231,8 +227,6 @@ SELECT result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at FROM @@ -286,8 +280,6 @@ SELECT result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at FROM @@ -383,8 +375,6 @@ RETURNING result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at; ` @@ -432,8 +422,6 @@ SET result = @result, error_message = @error_message, started_at = @started_at, - lease_expires_at = @lease_expires_at, - lease_generation = @lease_generation, updated_at = @updated_at WHERE %s @@ -448,8 +436,6 @@ RETURNING result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at; ` @@ -457,14 +443,12 @@ RETURNING q = fmt.Sprintf(q, scope.SQLFragment()) args := pgx.StrictNamedArgs{ - "id": e.ID.String(), - "status": e.Status, - "result": e.Result, - "error_message": e.ErrorMessage, - "started_at": e.StartedAt, - "lease_expires_at": e.LeaseExpiresAt, - "lease_generation": e.LeaseGeneration, - "updated_at": e.UpdatedAt, + "id": e.ID.String(), + "status": e.Status, + "result": e.Result, + "error_message": e.ErrorMessage, + "started_at": e.StartedAt, + "updated_at": e.UpdatedAt, } maps.Copy(args, scope.SQLArguments()) @@ -516,11 +500,17 @@ WHERE return nil } +// CommitAgentRunResult writes the terminal or resting state of a run +// (COMPLETED, FAILED, PENDING on graceful suspend, or AWAITING_APPROVAL) +// guarded on the row still being RUNNING. The guard is a lightweight +// safety net: it discards a commit for a run a human moved out of +// RUNNING manually (the only way a run leaves RUNNING out from under an +// active worker now that lease-based recovery is gone). Returns the +// number of rows affected (0 when the guard rejected the write). func CommitAgentRunResult( ctx context.Context, tx pg.Tx, e *AgentRun, - leaseGeneration int64, ) (int64, error) { q := ` UPDATE agent_runs @@ -529,23 +519,19 @@ SET result = @result, error_message = @error_message, started_at = @started_at, - lease_expires_at = @lease_expires_at, updated_at = @updated_at WHERE id = @id - AND status = 'RUNNING' - AND lease_generation = @lease_generation; + AND status = 'RUNNING'; ` args := pgx.StrictNamedArgs{ - "id": e.ID.String(), - "status": e.Status, - "result": e.Result, - "error_message": e.ErrorMessage, - "started_at": e.StartedAt, - "lease_expires_at": e.LeaseExpiresAt, - "updated_at": e.UpdatedAt, - "lease_generation": leaseGeneration, + "id": e.ID.String(), + "status": e.Status, + "result": e.Result, + "error_message": e.ErrorMessage, + "started_at": e.StartedAt, + "updated_at": e.UpdatedAt, } tag, err := tx.Exec(ctx, q, args) @@ -556,6 +542,56 @@ WHERE return tag.RowsAffected(), nil } +// RequeueForApprovalResume persists the run's checkpoint and returns it +// to PENDING so a worker resumes from the approval boundary. The caller +// populates e.Checkpoint (with the approval decisions merged in), +// e.Status, e.StartedAt, and e.UpdatedAt before calling. The write is +// guarded on the row still being AWAITING_APPROVAL; ErrResourceNotFound +// is returned when the guard rejects it. This is the one path that writes +// the checkpoint alongside a status change — Update deliberately omits the +// checkpoint column. +func (e *AgentRun) RequeueForApprovalResume( + ctx context.Context, + tx pg.Tx, + scope Scoper, +) error { + q := ` +UPDATE agent_runs +SET + checkpoint = @checkpoint, + status = @status, + started_at = @started_at, + updated_at = @updated_at +WHERE + %s + AND id = @id + AND status = @expected_status; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "id": e.ID.String(), + "checkpoint": e.Checkpoint, + "status": e.Status, + "started_at": e.StartedAt, + "updated_at": e.UpdatedAt, + "expected_status": AgentRunStatusAwaitingApproval, + } + maps.Copy(args, scope.SQLArguments()) + + result, err := tx.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot requeue agent run for approval resume: %w", err) + } + + if result.RowsAffected() == 0 { + return ErrResourceNotFound + } + + return nil +} + func (e *AgentRun) LoadNextPendingForUpdateSkipLocked( ctx context.Context, tx pg.Tx, @@ -571,8 +607,6 @@ SELECT result, error_message, started_at, - lease_expires_at, - lease_generation, created_at, updated_at FROM @@ -603,67 +637,6 @@ FOR UPDATE SKIP LOCKED; return nil } -// ResetStaleAgentRuns resets agent runs whose worker lease has expired. -// 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 worker auto-resumes -// from checkpoint when one exists. -func ResetStaleAgentRuns(ctx context.Context, conn pg.Querier) error { - q := ` -UPDATE agent_runs -SET - status = 'PENDING', - started_at = NULL, - lease_expires_at = NULL, - updated_at = now() -WHERE - status = 'RUNNING' - AND lease_expires_at IS NOT NULL - AND lease_expires_at < now(); -` - - _, err := conn.Exec(ctx, q) - if err != nil { - return fmt.Errorf("cannot reset stale agent runs: %w", err) - } - - return nil -} - -// HeartbeatAgentRunLease refreshes the lease for a running agent run. -// Returns the number of rows affected (0 if lease was lost). -func HeartbeatAgentRunLease( - ctx context.Context, - conn pg.Querier, - runID string, - leaseGeneration int64, - expiresAt time.Time, -) (int64, error) { - q := ` -UPDATE agent_runs -SET - lease_expires_at = @lease_expires_at, - updated_at = now() -WHERE - id = @id - AND status = 'RUNNING' - AND lease_generation = @lease_generation; -` - - args := pgx.StrictNamedArgs{ - "id": runID, - "lease_expires_at": expiresAt, - "lease_generation": leaseGeneration, - } - - tag, err := conn.Exec(ctx, q, args) - if err != nil { - return 0, fmt.Errorf("cannot heartbeat agent run lease: %w", err) - } - - return tag.RowsAffected(), nil -} - // PGCheckpointer implements agent.Checkpointer backed by the // agent_runs table checkpoint column. The runID is validated as a GID // up front so malformed identifiers fail closed; rows are then scoped @@ -738,55 +711,6 @@ WHERE ) } -func (s *PGCheckpointer) SaveForLease( - ctx context.Context, - runID string, - cp *agent.Checkpoint, - leaseGeneration int64, -) error { - if _, err := gid.ParseGID(runID); err != nil { - return fmt.Errorf("cannot parse agent run id: %w", err) - } - - data, err := s.marshalAgentCheckpoint(cp) - if err != nil { - return err - } - - return s.pg.WithConn( - ctx, - func(ctx context.Context, conn pg.Querier) error { - q := ` -UPDATE agent_runs -SET - checkpoint = @checkpoint, - updated_at = now() -WHERE - id = @id - AND status = 'RUNNING' - AND lease_generation = @lease_generation; -` - - args := pgx.StrictNamedArgs{ - "id": runID, - "checkpoint": json.RawMessage(data), - "lease_generation": leaseGeneration, - } - - tag, err := conn.Exec(ctx, q, args) - if err != nil { - return fmt.Errorf("cannot save checkpoint: %w", err) - } - - if tag.RowsAffected() == 0 { - return fmt.Errorf("cannot save checkpoint: lease lost") - } - - return nil - }, - ) -} - func (s *PGCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) { if _, err := gid.ParseGID(runID); err != nil { return nil, fmt.Errorf("cannot parse agent run id: %w", err) diff --git a/pkg/coredata/migrations/20260608T090000Z.sql b/pkg/coredata/migrations/20260608T090000Z.sql new file mode 100644 index 000000000..a7294a489 --- /dev/null +++ b/pkg/coredata/migrations/20260608T090000Z.sql @@ -0,0 +1,25 @@ +-- 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. + +-- Drop the agent-run worker lease. Recovery from a crashed worker is now +-- manual (inspect logs, move the row out of RUNNING by hand). Known stops +-- transition explicitly: graceful shutdown returns the row to PENDING and +-- approval pauses park it in AWAITING_APPROVAL, so nothing relies on a +-- timeout to requeue. Re-introduce leasing here when the system matures. + +DROP INDEX IF EXISTS idx_agent_runs_running_lease; + +ALTER TABLE agent_runs + DROP COLUMN lease_expires_at, + DROP COLUMN lease_generation; diff --git a/pkg/server/api/console/v1/agent_run_resolvers.go b/pkg/server/api/console/v1/agent_run_resolvers.go index 0754336a1..e1aaefa4c 100644 --- a/pkg/server/api/console/v1/agent_run_resolvers.go +++ b/pkg/server/api/console/v1/agent_run_resolvers.go @@ -12,6 +12,8 @@ import ( "github.com/vikstrous/dataloadgen" "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/probo" "go.probo.inc/probo/pkg/server/api/console/v1/dataloader" @@ -49,7 +51,7 @@ func (r *agentRunResolver) Permission(ctx context.Context, obj *types.AgentRun, // TotalCount is the resolver for the totalCount field. func (r *agentRunConnectionResolver) TotalCount(ctx context.Context, obj *types.AgentRunConnection) (int, error) { - scope, err := r.authorize(ctx, obj.ParentID, probo.ActionAgentRunList) + scope, err := r.authorize(ctx, obj.ParentID, agentrun.ActionAgentRunList) if err != nil { return 0, err } @@ -70,6 +72,46 @@ func (r *agentRunConnectionResolver) TotalCount(ctx context.Context, obj *types. return 0, gqlutils.Internal(ctx) } +// SubmitAgentRunApproval is the resolver for the submitAgentRunApproval field. +func (r *mutationResolver) SubmitAgentRunApproval(ctx context.Context, input types.SubmitAgentRunApprovalInput) (*types.SubmitAgentRunApprovalPayload, error) { + scope, err := r.authorize(ctx, input.AgentRunID, agentrun.ActionAgentRunApprove) + if err != nil { + return nil, err + } + + decisions := make(map[string]agent.ApprovalResult, len(input.Decisions)) + for _, decision := range input.Decisions { + message := "" + if decision.Reason != nil { + message = *decision.Reason + } + + decisions[decision.ToolCallID] = agent.ApprovalResult{ + Approved: decision.Approved, + Message: message, + } + } + + run, err := r.agentRun.SubmitApproval(ctx, scope, input.AgentRunID, decisions) + if err != nil { + switch { + case errors.Is(err, agentrun.ErrAgentRunNotFound): + return nil, gqlutils.NotFound(ctx, err) + case errors.Is(err, agentrun.ErrNotAwaitingApproval): + return nil, gqlutils.Conflictf(ctx, "agent run is not awaiting approval") + case errors.Is(err, agentrun.ErrApprovalDecisionsMismatch): + return nil, gqlutils.Invalidf(ctx, "approval decisions must match the run's pending approvals") + default: + r.logger.ErrorCtx(ctx, "cannot submit agent run approval", log.Error(err)) + return nil, gqlutils.Internal(ctx) + } + } + + return &types.SubmitAgentRunApprovalPayload{ + AgentRun: types.NewAgentRun(run), + }, nil +} + // AgentRun returns schema.AgentRunResolver implementation. func (r *Resolver) AgentRun() schema.AgentRunResolver { return &agentRunResolver{r} } diff --git a/pkg/server/api/console/v1/graphql/agent_run.graphql b/pkg/server/api/console/v1/graphql/agent_run.graphql index 3df26f44f..0056c6f05 100644 --- a/pkg/server/api/console/v1/graphql/agent_run.graphql +++ b/pkg/server/api/console/v1/graphql/agent_run.graphql @@ -58,3 +58,24 @@ type AgentRunEdge { cursor: CursorKey! node: AgentRun! } + +extend type Mutation { + submitAgentRunApproval( + input: SubmitAgentRunApprovalInput! + ): SubmitAgentRunApprovalPayload! +} + +input AgentRunApprovalDecisionInput { + toolCallId: String! + approved: Boolean! + reason: String +} + +input SubmitAgentRunApprovalInput { + agentRunId: ID! + decisions: [AgentRunApprovalDecisionInput!]! +} + +type SubmitAgentRunApprovalPayload { + agentRun: AgentRun! +}