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! +}