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>
105 lines
2.6 KiB
Go
105 lines
2.6 KiB
Go
// 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 (
|
|
"context"
|
|
"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"
|
|
)
|
|
|
|
type (
|
|
Worker struct {
|
|
handler *handler
|
|
kitWorker *worker.Worker[coredata.AgentRun]
|
|
}
|
|
|
|
WorkerOption func(*workerConfig)
|
|
|
|
workerConfig struct {
|
|
interval time.Duration
|
|
maxConcurrency int
|
|
}
|
|
)
|
|
|
|
func WithWorkerInterval(d time.Duration) WorkerOption {
|
|
return func(c *workerConfig) {
|
|
if d > 0 {
|
|
c.interval = d
|
|
}
|
|
}
|
|
}
|
|
|
|
func WithWorkerMaxConcurrency(n int) WorkerOption {
|
|
return func(c *workerConfig) {
|
|
if n > 0 {
|
|
c.maxConcurrency = n
|
|
}
|
|
}
|
|
}
|
|
|
|
func NewWorker(
|
|
pgClient *pg.Client,
|
|
store *coredata.PGCheckpointer,
|
|
registry agent.AgentRegistry,
|
|
logger *log.Logger,
|
|
opts ...WorkerOption,
|
|
) *Worker {
|
|
cfg := workerConfig{
|
|
interval: 10 * time.Second,
|
|
maxConcurrency: 5,
|
|
}
|
|
|
|
for _, opt := range opts {
|
|
opt(&cfg)
|
|
}
|
|
|
|
h := &handler{
|
|
pg: pgClient,
|
|
store: store,
|
|
registry: registry,
|
|
logger: logger,
|
|
shutdownCh: make(chan struct{}),
|
|
}
|
|
|
|
w := worker.New(
|
|
"agent-run-worker",
|
|
h,
|
|
logger,
|
|
worker.WithInterval(cfg.interval),
|
|
worker.WithMaxConcurrency(cfg.maxConcurrency),
|
|
)
|
|
|
|
return &Worker{handler: h, kitWorker: w}
|
|
}
|
|
|
|
// Run starts the worker 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.
|
|
//
|
|
// signalShutdown is registered without a stop hook because it is
|
|
// idempotent (sync.Once) and we want it to fire on every ctx
|
|
// cancellation, even one that races with kitWorker.Run returning.
|
|
func (w *Worker) Run(ctx context.Context) error {
|
|
context.AfterFunc(ctx, w.handler.signalShutdown)
|
|
return w.kitWorker.Run(ctx)
|
|
}
|