From 3dfc833671e5b620982e4d2e40cabe37da17dd75 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Sun, 7 Jun 2026 09:00:01 +0200 Subject: [PATCH] Replace supervisor with agentrun worker service Move agent-run orchestration from the legacy supervisor path into the new agentrun worker/service package and wire it through coredata, server, policies, and GraphQL resolvers. This consolidates run lifecycle handling around lease-aware workers and aligns API surface with the new agent-run domain model so reviewers can follow one coherent execution path. Signed-off-by: Bryan Frimin --- contrib/claude/agent.md | 4 +- pkg/agent/checkpoint.go | 6 +- pkg/agentrun/export_test.go | 24 + .../handler.go} | 105 +- pkg/agentrun/internal_test.go | 111 +++ pkg/agentrun/service.go | 114 +++ pkg/agentrun/service_test.go | 95 ++ .../worker.go} | 63 +- pkg/agentruntest/agent_run_supervisor_test.go | 906 ------------------ pkg/agentruntest/agentruntest.go | 247 ----- pkg/coredata/agent_run.go | 140 ++- pkg/coredata/migrations/20260607T060000Z.sql | 17 + pkg/coredata/pg_checkpointer_test.go | 13 +- pkg/coredata/pg_test_helpers_test.go | 202 ++++ pkg/probo/actions.go | 4 + pkg/probo/policies.go | 2 + pkg/probod/probod.go | 4 + pkg/server/api/api.go | 3 + .../api/console/v1/agent_run_resolvers.go | 82 ++ pkg/server/api/console/v1/base_resolvers.go | 10 + .../api/console/v1/graphql/agent_run.graphql | 60 ++ .../console/v1/graphql/organization.graphql | 8 + pkg/server/api/console/v1/graphql_handler.go | 3 + .../api/console/v1/organization_resolvers.go | 30 + pkg/server/api/console/v1/resolver.go | 4 + pkg/server/api/console/v1/types/agent_run.go | 76 ++ pkg/server/server.go | 3 + 27 files changed, 1076 insertions(+), 1260 deletions(-) create mode 100644 pkg/agentrun/export_test.go rename pkg/{probo/agent_run_handler.go => agentrun/handler.go} (76%) create mode 100644 pkg/agentrun/internal_test.go create mode 100644 pkg/agentrun/service.go create mode 100644 pkg/agentrun/service_test.go rename pkg/{probo/agent_run_supervisor.go => agentrun/worker.go} (54%) delete mode 100644 pkg/agentruntest/agent_run_supervisor_test.go delete mode 100644 pkg/agentruntest/agentruntest.go create mode 100644 pkg/coredata/migrations/20260607T060000Z.sql create mode 100644 pkg/coredata/pg_test_helpers_test.go create mode 100644 pkg/server/api/console/v1/agent_run_resolvers.go create mode 100644 pkg/server/api/console/v1/graphql/agent_run.graphql create mode 100644 pkg/server/api/console/v1/types/agent_run.go diff --git a/contrib/claude/agent.md b/contrib/claude/agent.md index 0f294b9c7..235520c58 100644 --- a/contrib/claude/agent.md +++ b/contrib/claude/agent.md @@ -62,7 +62,7 @@ the cancellation so they complete naturally; the cancel is only observed at the next safe boundary. Use `agent.ErrSuspendForCheckpoint` as the cancel cause when the -intent is graceful suspend — supervisors that distinguish a +intent is graceful suspend — workers that distinguish a graceful-stop request from infrastructure-level causes (lease loss, heartbeat failure) inspect `context.Cause(ctx)` to dispatch. @@ -77,7 +77,7 @@ Implications: it must derive its own with `context.WithTimeout(ctx, ...)` inside the tool body. -The supervisor (`pkg/probo/agent_run_handler.go`) maps a SIGTERM-driven +The agent run worker (`pkg/agentrun/handler.go`) maps a SIGTERM-driven shutdown broadcast onto a per-run `cancelRun(agent.ErrSuspendForCheckpoint)`, so the same contract drives both the public Go API and the worker infrastructure path. diff --git a/pkg/agent/checkpoint.go b/pkg/agent/checkpoint.go index 840d3d1b3..54700160b 100644 --- a/pkg/agent/checkpoint.go +++ b/pkg/agent/checkpoint.go @@ -63,9 +63,9 @@ type ( Result ToolResult } - // Checkpointer is supervisor-internal. Implementations may use raw + // Checkpointer is worker-internal. Implementations may use raw // run IDs because public API/service methods perform tenant scoping and - // authorization before a run reaches the supervisor. + // authorization before a run reaches the worker. Checkpointer interface { Save(ctx context.Context, runID string, cp *Checkpoint) error Load(ctx context.Context, runID string) (*Checkpoint, error) @@ -76,7 +76,7 @@ type ( } SuspendedError struct { - RunID string // Set when the outer loop has a store+runID (supervisor-managed). + RunID string // Set when the outer loop has a store+runID (worker-managed). Checkpoint *Checkpoint // Set when returning from an inner agent-as-tool (no store). } ) diff --git a/pkg/agentrun/export_test.go b/pkg/agentrun/export_test.go new file mode 100644 index 000000000..4ecef5e04 --- /dev/null +++ b/pkg/agentrun/export_test.go @@ -0,0 +1,24 @@ +// 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 + +// ShutdownBroadcast returns the channel that closes once the worker has +// broadcast graceful shutdown to all in-flight runs. It is compiled only +// in test builds so external tests can synchronize tool release with +// shutdown propagation without leaking a test-only method into the +// worker's public API. +func (w *Worker) ShutdownBroadcast() <-chan struct{} { + return w.handler.shutdownCh +} diff --git a/pkg/probo/agent_run_handler.go b/pkg/agentrun/handler.go similarity index 76% rename from pkg/probo/agent_run_handler.go rename to pkg/agentrun/handler.go index e5c361adb..8bd79eacc 100644 --- a/pkg/probo/agent_run_handler.go +++ b/pkg/agentrun/handler.go @@ -12,7 +12,7 @@ // OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR // PERFORMANCE OF THIS SOFTWARE. -package probo +package agentrun import ( "context" @@ -31,30 +31,28 @@ import ( "go.probo.inc/probo/pkg/llm" ) -type agentRunHandler struct { +type handler struct { pg *pg.Client store *coredata.PGCheckpointer registry agent.AgentRegistry logger *log.Logger leaseDuration time.Duration - workerID string shutdownCh chan struct{} shutdownOnce sync.Once } var ( - _ worker.Handler[coredata.AgentRun] = (*agentRunHandler)(nil) - _ worker.StaleRecoverer = (*agentRunHandler)(nil) + _ worker.Handler[coredata.AgentRun] = (*handler)(nil) + _ worker.StaleRecoverer = (*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. -func (h *agentRunHandler) Claim(ctx context.Context) (coredata.AgentRun, error) { +func (h *handler) Claim(ctx context.Context) (coredata.AgentRun, error) { var ( run = coredata.AgentRun{} now = time.Now() - leaseOwner = h.workerID leaseExpiresAt = now.Add(h.leaseDuration) ) @@ -67,8 +65,8 @@ func (h *agentRunHandler) Claim(ctx context.Context) (coredata.AgentRun, error) run.Status = coredata.AgentRunStatusRunning run.StartedAt = &now - run.LeaseOwner = &leaseOwner run.LeaseExpiresAt = &leaseExpiresAt + run.LeaseGeneration++ run.UpdatedAt = now if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil { @@ -99,9 +97,10 @@ func (h *agentRunHandler) Claim(ctx context.Context) (coredata.AgentRun, error) // nil is returned for both successful runs and graceful exits // (lease loss, infrastructure suspension) where the row state is // already consistent. -func (h *agentRunHandler) Process(ctx context.Context, run coredata.AgentRun) error { +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) @@ -117,14 +116,14 @@ func (h *agentRunHandler) Process(ctx context.Context, run coredata.AgentRun) er heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx) defer cancelHeartbeat() - go h.heartbeatLease(heartbeatCtx, run.ID.String(), cancelRun) + go h.heartbeatLease(heartbeatCtx, run.ID.String(), leaseGeneration, cancelRun) - return h.executeRun(runCtx, &run) + return h.executeRun(runCtx, &run, leaseGeneration) } // RecoverStale resets agent runs whose worker lease has expired back to -// PENDING so a fresh supervisor can pick them up on the next cycle. -func (h *agentRunHandler) RecoverStale(ctx context.Context) error { +// 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 { @@ -141,13 +140,14 @@ func (h *agentRunHandler) RecoverStale(ctx context.Context) error { // in-flight Process forwarder goroutines observe the close and propagate // it to their per-run agent stop channels, letting agents checkpoint at // the next turn boundary before Process returns. -func (h *agentRunHandler) signalShutdown() { +func (h *handler) signalShutdown() { h.shutdownOnce.Do(func() { close(h.shutdownCh) }) } -func (h *agentRunHandler) heartbeatLease( +func (h *handler) heartbeatLease( ctx context.Context, runID string, + leaseGeneration int64, cancelRun context.CancelCauseFunc, ) { ticker := time.NewTicker(h.leaseDuration / 3) @@ -163,13 +163,19 @@ func (h *agentRunHandler) heartbeatLease( if err := h.pg.WithConn( ctx, func(ctx context.Context, conn pg.Querier) error { - rowsAffected, err := coredata.HeartbeatAgentRunLease(ctx, conn, runID, h.workerID, expiresAt) + rowsAffected, err := coredata.HeartbeatAgentRunLease( + ctx, + conn, + runID, + leaseGeneration, + expiresAt, + ) if err != nil { return err } if rowsAffected == 0 { - return ErrAgentRunLeaseLost + return ErrLeaseLost } return nil @@ -177,10 +183,10 @@ func (h *agentRunHandler) heartbeatLease( ); err != nil { h.logger.ErrorCtx(ctx, "cannot heartbeat agent run lease", log.Error(err)) - if errors.Is(err, ErrAgentRunLeaseLost) { - cancelRun(ErrAgentRunLeaseLost) + if errors.Is(err, ErrLeaseLost) { + cancelRun(ErrLeaseLost) } else { - cancelRun(fmt.Errorf("%w: %w", ErrAgentRunHeartbeatFailed, err)) + cancelRun(fmt.Errorf("%w: %w", ErrHeartbeatFailed, err)) } return @@ -190,21 +196,21 @@ func (h *agentRunHandler) heartbeatLease( } const ( - // agentRunErrorMessageMaxLen caps the error string persisted to the + // errorMessageMaxLen caps the error string persisted to the // agent_runs.error_message column. Raw tool or LLM errors can embed // URLs with credentials, response snippets containing PII, or partial // records from failed DB lookups; the full context is logged while // only a truncated summary is stored for caller-visible state. - agentRunErrorMessageMaxLen = 512 + errorMessageMaxLen = 512 ) -func sanitizeAgentRunError(err error) string { +func sanitizeError(err error) string { msg := err.Error() - if len(msg) <= agentRunErrorMessageMaxLen { + if len(msg) <= errorMessageMaxLen { return msg } - cut := agentRunErrorMessageMaxLen + cut := errorMessageMaxLen for cut > 0 && !utf8.RuneStart(msg[cut]) { cut-- } @@ -212,8 +218,29 @@ func sanitizeAgentRunError(err error) string { return msg[:cut] + "…" } -func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun) error { +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 { runID := run.ID.String() + checkpointer := leasedCheckpointer{ + store: h.store, + leaseGeneration: leaseGeneration, + } var ( result *agent.Result @@ -222,7 +249,7 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun if run.Checkpoint != nil { h.logger.InfoCtx(ctx, "resuming agent run", log.String("run_id", runID)) - result, runErr = agent.Restore(ctx, h.store, runID, h.registry) + result, runErr = agent.Restore(ctx, checkpointer, runID, h.registry) } else { h.logger.InfoCtx(ctx, "starting agent run", log.String("run_id", runID)) @@ -237,7 +264,7 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun result, runErr = a.Run( ctx, inputMsgs, - agent.WithCheckpointer(h.store, runID), + agent.WithCheckpointer(checkpointer, runID), ) } } @@ -246,7 +273,7 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun // 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, ErrAgentRunLeaseLost) || errors.Is(cause, ErrAgentRunHeartbeatFailed) { + 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", @@ -277,7 +304,6 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun now := time.Now() run.UpdatedAt = now run.StartedAt = nil - run.LeaseOwner = nil run.LeaseExpiresAt = nil if runErr == nil { @@ -304,7 +330,7 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun log.String("run_id", runID), log.Error(runErr), ) - msg := sanitizeAgentRunError(runErr) + msg := sanitizeError(runErr) run.ErrorMessage = &msg } @@ -313,10 +339,15 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun if err := h.pg.WithTx( commitCtx, func(ctx context.Context, tx pg.Tx) error { - if err := run.Update(ctx, tx, coredata.NewNoScope()); err != nil { + rowsAffected, err := coredata.CommitAgentRunResult(ctx, tx, run, leaseGeneration) + if err != nil { return err } + if rowsAffected == 0 { + return ErrLeaseLost + } + if run.Status == coredata.AgentRunStatusCompleted { if err := run.ClearCheckpoint(ctx, tx, coredata.NewNoScope()); err != nil { return err @@ -326,6 +357,16 @@ func (h *agentRunHandler) executeRun(ctx context.Context, run *coredata.AgentRun 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) } diff --git a/pkg/agentrun/internal_test.go b/pkg/agentrun/internal_test.go new file mode 100644 index 000000000..bca126159 --- /dev/null +++ b/pkg/agentrun/internal_test.go @@ -0,0 +1,111 @@ +// 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" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestSanitizeError(t *testing.T) { + t.Parallel() + + t.Run( + "short message unchanged", + func(t *testing.T) { + t.Parallel() + + err := errors.New("short") + assert.Equal(t, "short", sanitizeError(err)) + }, + ) + + t.Run( + "boundary length unchanged", + func(t *testing.T) { + t.Parallel() + + msg := strings.Repeat("a", errorMessageMaxLen) + assert.Equal(t, msg, sanitizeError(errors.New(msg))) + }, + ) + + t.Run( + "long utf8 message is rune safe and suffixed", + func(t *testing.T) { + t.Parallel() + + msg := strings.Repeat("é", errorMessageMaxLen) + sanitized := sanitizeError(errors.New(msg)) + + assert.True(t, strings.HasSuffix(sanitized, "…")) + assert.True(t, len(sanitized) <= errorMessageMaxLen+len("…")) + assert.True(t, strings.HasPrefix(msg, strings.TrimSuffix(sanitized, "…"))) + }, + ) +} + +func TestWorkerOptions(t *testing.T) { + t.Parallel() + + t.Run( + "interval updates only when positive", + func(t *testing.T) { + t.Parallel() + + cfg := workerConfig{interval: 3 * time.Second} + + WithWorkerInterval(0)(&cfg) + assert.Equal(t, 3*time.Second, cfg.interval) + + WithWorkerInterval(7 * time.Second)(&cfg) + assert.Equal(t, 7*time.Second, cfg.interval) + }, + ) + + t.Run( + "lease duration updates only when positive", + func(t *testing.T) { + t.Parallel() + + cfg := workerConfig{leaseDuration: 5 * time.Second} + + WithWorkerLeaseDuration(-1)(&cfg) + assert.Equal(t, 5*time.Second, cfg.leaseDuration) + + WithWorkerLeaseDuration(12 * time.Second)(&cfg) + assert.Equal(t, 12*time.Second, cfg.leaseDuration) + }, + ) + + t.Run( + "max concurrency updates only when positive", + func(t *testing.T) { + t.Parallel() + + cfg := workerConfig{maxConcurrency: 2} + + WithWorkerMaxConcurrency(0)(&cfg) + assert.Equal(t, 2, cfg.maxConcurrency) + + WithWorkerMaxConcurrency(9)(&cfg) + assert.Equal(t, 9, cfg.maxConcurrency) + }, + ) +} diff --git a/pkg/agentrun/service.go b/pkg/agentrun/service.go new file mode 100644 index 000000000..55d7f615e --- /dev/null +++ b/pkg/agentrun/service.go @@ -0,0 +1,114 @@ +// 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 ( + "context" + "fmt" + + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" +) + +type Service struct { + pg *pg.Client +} + +func NewService(pgClient *pg.Client) *Service { + return &Service{pg: pgClient} +} + +func (s *Service) Get( + ctx context.Context, + scope coredata.Scoper, + agentRunID gid.GID, +) (*coredata.AgentRun, error) { + run := &coredata.AgentRun{} + + err := s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + if err := run.LoadByID(ctx, conn, scope, agentRunID); err != nil { + return fmt.Errorf("cannot load agent run: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return run, nil +} + +func (s *Service) ListForOrganizationID( + ctx context.Context, + scope coredata.Scoper, + organizationID gid.GID, + cursor *page.Cursor[coredata.AgentRunOrderField], +) (*page.Page[*coredata.AgentRun, coredata.AgentRunOrderField], error) { + var runs coredata.AgentRuns + + err := s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + organization := &coredata.Organization{} + if err := organization.LoadByID(ctx, conn, scope, organizationID); err != nil { + return fmt.Errorf("cannot load organization: %w", err) + } + + if err := runs.LoadByOrganizationID(ctx, conn, scope, organization.ID, cursor); err != nil { + return fmt.Errorf("cannot load agent runs: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return page.NewPage(runs, cursor), nil +} + +func (s *Service) CountForOrganizationID( + ctx context.Context, + scope coredata.Scoper, + organizationID gid.GID, +) (int, error) { + var count int + + err := s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) (err error) { + runs := &coredata.AgentRuns{} + + count, err = runs.CountByOrganizationID(ctx, conn, scope, organizationID) + if err != nil { + return fmt.Errorf("cannot count agent runs: %w", err) + } + + return nil + }, + ) + if err != nil { + return 0, err + } + + return count, nil +} diff --git a/pkg/agentrun/service_test.go b/pkg/agentrun/service_test.go new file mode 100644 index 000000000..81d7a56cb --- /dev/null +++ b/pkg/agentrun/service_test.go @@ -0,0 +1,95 @@ +// 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_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.probo.inc/probo/pkg/agentrun" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" +) + +func TestService_Get(t *testing.T) { + client := pgClient(t) + svc := agentrun.NewService(client) + + run := insertPendingRun( + t, + client, + "service-get-agent", + nil, + ) + + got, err := svc.Get(context.Background(), coredata.NewNoScope(), run.ID) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, run.ID, got.ID) + + missingID := gid.New(run.ID.TenantID(), coredata.AgentRunEntityType) + _, err = svc.Get(context.Background(), coredata.NewNoScope(), missingID) + require.Error(t, err) + assert.ErrorIs(t, err, coredata.ErrResourceNotFound) +} + +func TestService_ListForOrganizationID(t *testing.T) { + client := pgClient(t) + svc := agentrun.NewService(client) + + orgID := insertTestOrganization(t, client) + + runA := insertPendingRunInOrg(t, client, orgID, "service-list-agent-a", nil) + runB := insertPendingRunInOrg(t, client, orgID, "service-list-agent-b", nil) + + cursor := page.NewCursor( + 10, + nil, + page.Head, + page.OrderBy[coredata.AgentRunOrderField]{ + Field: coredata.AgentRunOrderFieldCreatedAt, + Direction: page.OrderDirectionDesc, + }, + ) + + got, err := svc.ListForOrganizationID(context.Background(), coredata.NewNoScope(), orgID, cursor) + require.NoError(t, err) + require.NotNil(t, got) + + ids := make(map[gid.GID]bool) + for _, run := range got.Data { + ids[run.ID] = true + } + + assert.True(t, ids[runA.ID]) + assert.True(t, ids[runB.ID]) +} + +func TestService_CountForOrganizationID(t *testing.T) { + client := pgClient(t) + svc := agentrun.NewService(client) + + orgID := insertTestOrganization(t, client) + _ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-a", nil) + _ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-b", nil) + _ = insertPendingRunInOrg(t, client, orgID, "service-count-agent-c", nil) + + count, err := svc.CountForOrganizationID(context.Background(), coredata.NewNoScope(), orgID) + require.NoError(t, err) + assert.Equal(t, 3, count) +} diff --git a/pkg/probo/agent_run_supervisor.go b/pkg/agentrun/worker.go similarity index 54% rename from pkg/probo/agent_run_supervisor.go rename to pkg/agentrun/worker.go index ca09558ae..4d78dc0d2 100644 --- a/pkg/probo/agent_run_supervisor.go +++ b/pkg/agentrun/worker.go @@ -12,14 +12,13 @@ // OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR // PERFORMANCE OF THIS SOFTWARE. -package probo +package agentrun import ( "context" "errors" "time" - "go.gearno.de/crypto/uuid" "go.gearno.de/kit/log" "go.gearno.de/kit/pg" "go.gearno.de/kit/worker" @@ -28,14 +27,14 @@ import ( ) type ( - AgentRunSupervisor struct { - handler *agentRunHandler - worker *worker.Worker[coredata.AgentRun] + Worker struct { + handler *handler + kitWorker *worker.Worker[coredata.AgentRun] } - AgentRunSupervisorOption func(*agentRunSupervisorConfig) + WorkerOption func(*workerConfig) - agentRunSupervisorConfig struct { + workerConfig struct { interval time.Duration leaseDuration time.Duration maxConcurrency int @@ -43,42 +42,42 @@ type ( ) var ( - ErrAgentRunHeartbeatFailed = errors.New("agent run heartbeat failed") - ErrAgentRunLeaseLost = errors.New("agent run lease lost") + ErrHeartbeatFailed = errors.New("agent run heartbeat failed") + ErrLeaseLost = errors.New("agent run lease lost") ) -func WithAgentRunSupervisorInterval(d time.Duration) AgentRunSupervisorOption { - return func(c *agentRunSupervisorConfig) { +func WithWorkerInterval(d time.Duration) WorkerOption { + return func(c *workerConfig) { if d > 0 { c.interval = d } } } -func WithAgentRunSupervisorLeaseDuration(d time.Duration) AgentRunSupervisorOption { - return func(c *agentRunSupervisorConfig) { +func WithWorkerLeaseDuration(d time.Duration) WorkerOption { + return func(c *workerConfig) { if d > 0 { c.leaseDuration = d } } } -func WithAgentRunSupervisorMaxConcurrency(n int) AgentRunSupervisorOption { - return func(c *agentRunSupervisorConfig) { +func WithWorkerMaxConcurrency(n int) WorkerOption { + return func(c *workerConfig) { if n > 0 { c.maxConcurrency = n } } } -func NewAgentRunSupervisor( +func NewWorker( pgClient *pg.Client, store *coredata.PGCheckpointer, registry agent.AgentRegistry, logger *log.Logger, - opts ...AgentRunSupervisorOption, -) *AgentRunSupervisor { - cfg := agentRunSupervisorConfig{ + opts ...WorkerOption, +) *Worker { + cfg := workerConfig{ interval: 10 * time.Second, leaseDuration: 5 * time.Minute, maxConcurrency: 5, @@ -88,45 +87,35 @@ func NewAgentRunSupervisor( opt(&cfg) } - h := &agentRunHandler{ + h := &handler{ pg: pgClient, store: store, registry: registry, logger: logger, leaseDuration: cfg.leaseDuration, - workerID: uuid.MustNewV4().String(), shutdownCh: make(chan struct{}), } w := worker.New( - "agent-run-supervisor", + "agent-run-worker", h, logger, worker.WithInterval(cfg.interval), worker.WithMaxConcurrency(cfg.maxConcurrency), ) - return &AgentRunSupervisor{handler: h, worker: w} + return &Worker{handler: h, kitWorker: w} } -// Run starts the supervisor loop. It blocks until ctx is cancelled, then +// 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 worker.Run returning. -func (s *AgentRunSupervisor) Run(ctx context.Context) error { - context.AfterFunc(ctx, s.handler.signalShutdown) - return s.worker.Run(ctx) -} - -// ShutdownBroadcastForTests returns a channel that closes once the -// supervisor has broadcast graceful shutdown to all in-flight runs. -// Exposed for tests that need to synchronize tool release with shutdown -// propagation; not a stable API and not part of the supervisor's public -// operational contract. -func (s *AgentRunSupervisor) ShutdownBroadcastForTests() <-chan struct{} { - return s.handler.shutdownCh +// 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) } diff --git a/pkg/agentruntest/agent_run_supervisor_test.go b/pkg/agentruntest/agent_run_supervisor_test.go deleted file mode 100644 index edd5d0da5..000000000 --- a/pkg/agentruntest/agent_run_supervisor_test.go +++ /dev/null @@ -1,906 +0,0 @@ -// 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 agentruntest_test - -import ( - "context" - "errors" - "fmt" - "os" - "os/exec" - "os/signal" - "path/filepath" - "strconv" - "sync" - "syscall" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "go.gearno.de/kit/log" - "go.gearno.de/kit/pg" - "go.probo.inc/probo/pkg/agent" - "go.probo.inc/probo/pkg/agentruntest" - "go.probo.inc/probo/pkg/coredata" - "go.probo.inc/probo/pkg/llm" - "go.probo.inc/probo/pkg/probo" -) - -// --------------------------------------------------------------------------- -// Test helpers -// --------------------------------------------------------------------------- - -func testLogger() *log.Logger { - return log.NewLogger(log.WithFormat(log.FormatPretty)) -} - -// --------------------------------------------------------------------------- -// Mock LLM provider -// --------------------------------------------------------------------------- - -type mockProvider struct { - mu sync.Mutex - responses []*llm.ChatCompletionResponse - calls int -} - -func (m *mockProvider) ChatCompletion(_ context.Context, _ *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) { - m.mu.Lock() - defer m.mu.Unlock() - - if m.calls >= len(m.responses) { - return nil, errors.New("no more mock responses") - } - - resp := m.responses[m.calls] - m.calls++ - - return resp, nil -} - -func (m *mockProvider) ChatCompletionStream(_ context.Context, _ *llm.ChatCompletionRequest) (llm.ChatCompletionStream, error) { - return nil, errors.New("not implemented") -} - -func newTestClient(provider llm.Provider) *llm.Client { - return llm.NewClient(provider, "test") -} - -func stopResponse(text string) *llm.ChatCompletionResponse { - return &llm.ChatCompletionResponse{ - Model: "test-model", - Message: llm.Message{ - Role: llm.RoleAssistant, - Parts: []llm.Part{llm.TextPart{Text: text}}, - }, - Usage: llm.Usage{InputTokens: 10, OutputTokens: 5}, - FinishReason: llm.FinishReasonStop, - } -} - -func toolCallResponse(toolCalls ...llm.ToolCall) *llm.ChatCompletionResponse { - return &llm.ChatCompletionResponse{ - Model: "test-model", - Message: llm.Message{ - Role: llm.RoleAssistant, - ToolCalls: toolCalls, - }, - Usage: llm.Usage{InputTokens: 10, OutputTokens: 5}, - FinishReason: llm.FinishReasonToolCalls, - } -} - -// --------------------------------------------------------------------------- -// Simple agent registry -// --------------------------------------------------------------------------- - -type simpleRegistry struct { - agents map[string]*agent.Agent -} - -func (r *simpleRegistry) Agent(name string) (*agent.Agent, error) { - a, ok := r.agents[name] - if !ok { - return nil, fmt.Errorf("agent %q not found", name) - } - - return a, nil -} - -// --------------------------------------------------------------------------- -// Test 3: Supervisor picks up a PENDING run and completes it -// --------------------------------------------------------------------------- - -// Supervisor tests are intentionally sequential. The supervisor claims -// runs cross-tenant via LoadNextPendingForUpdateSkipLocked; running two -// supervisors against the same test database would steal each other's -// runs. If a per-tenant claim filter is ever added, these can go back -// to t.Parallel(). -func TestAgentRunSupervisor_PicksUpAndCompletes(t *testing.T) { - client := agentruntest.PGClient(t) - store := coredata.NewPGCheckpointer(client) - - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - stopResponse("Done."), - }, - } - - ag := agent.New( - "echo-agent", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithInstructions("Reply with done."), - ) - - registry := &simpleRegistry{ - agents: map[string]*agent.Agent{"echo-agent": ag}, - } - - run := agentruntest.InsertPendingRun( - t, - client, - "echo-agent", - []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, - ) - - supervisor := probo.NewAgentRunSupervisor( - client, - store, - registry, - testLogger(), - probo.WithAgentRunSupervisorInterval(500*time.Millisecond), - probo.WithAgentRunSupervisorLeaseDuration(30*time.Second), - ) - - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel() - - go func() { _ = supervisor.Run(ctx) }() - - // Poll until the run is completed. - require.Eventually( - t, - func() bool { - r, err := agentruntest.TryLoadAgentRun(client, run.ID) - return err == nil && r.Status == coredata.AgentRunStatusCompleted - }, - 10*time.Second, - 200*time.Millisecond, - "run should reach COMPLETED status", - ) - - completed := agentruntest.LoadAgentRun(t, client, run.ID) - assert.Equal(t, coredata.AgentRunStatusCompleted, completed.Status) - assert.NotNil(t, completed.Result) - assert.Nil(t, completed.Checkpoint, "checkpoint should be cleared after completion") - assert.Nil(t, completed.ErrorMessage) -} - -// --------------------------------------------------------------------------- -// Test 4: Supervisor stop/resume cycle with checkpoint -// --------------------------------------------------------------------------- - -func TestAgentRunSupervisor_StopAndResume(t *testing.T) { - client := agentruntest.PGClient(t) - store := coredata.NewPGCheckpointer(client) - - // The tool blocks until signaled, giving us time to trigger graceful - // shutdown via the supervisor context while the agent is mid-turn. - toolReady := make(chan struct{}) - toolRelease := make(chan struct{}) - - slowTool := agent.FunctionTool[struct{}]( - "slow_work", - "Does slow work", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - close(toolReady) - <-toolRelease - - return agent.ToolResult{Content: "work done"}, nil - }, - ) - - // Provider sequence: - // Call 1: request tool call (first execution) - // Call 2: final stop response (after restoration) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - // First execution: LLM asks to call the tool. - toolCallResponse(llm.ToolCall{ - ID: "tc_1", - Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, - }), - // After resume: the incremental checkpoint saved after tool completion - // means restore continues with these messages; LLM returns final answer. - stopResponse("All done after resume."), - }, - } - - ag := agent.New( - "worker-agent", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(slowTool), - ) - - registry := &simpleRegistry{ - agents: map[string]*agent.Agent{"worker-agent": ag}, - } - - run := agentruntest.InsertPendingRun( - t, - client, - "worker-agent", - []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, - ) - - supervisor := probo.NewAgentRunSupervisor( - client, - store, - registry, - testLogger(), - probo.WithAgentRunSupervisorInterval(500*time.Millisecond), - probo.WithAgentRunSupervisorLeaseDuration(30*time.Second), - ) - - // --- Phase 1: Start and let the supervisor pick up the run --- - - ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel1() - - go func() { _ = supervisor.Run(ctx1) }() - - // Wait for the tool to start executing — this confirms the supervisor - // claimed the run and the agent called the tool. - select { - case <-toolReady: - case <-ctx1.Done(): - t.Fatal("timed out waiting for tool to start") - } - - // The run should now be RUNNING. - running := agentruntest.LoadAgentRun(t, client, run.ID) - assert.Equal(t, coredata.AgentRunStatusRunning, running.Status) - - // Trigger graceful shutdown of the supervisor: context.AfterFunc - // registered in Run() fires signalShutdown, closing the shutdown - // broadcast channel; the per-run forwarder goroutine closes the - // agent's stop channel. - cancel1() - - // Wait for the shutdown broadcast to be observed (the AfterFunc - // goroutine closes it) before releasing the tool. This is - // deterministic: no wall-clock sleep. The per-run forwarder - // goroutine observes the same close synchronously and closes the - // agent stop channel. - select { - case <-supervisor.ShutdownBroadcastForTests(): - case <-time.After(5 * time.Second): - t.Fatal("timed out waiting for supervisor shutdown broadcast") - } - - // Now release the tool. When the coreLoop resumes control at the - // next turn boundary it observes the closed stop channel, saves - // the suspension checkpoint, and returns SuspendedError. - close(toolRelease) - - // Wait for the checkpoint to appear. The supervisor leaves the row - // in RUNNING because SuspendedError triggers the "leaving for stale - // recovery" path. - require.Eventually( - t, - func() bool { - r, err := agentruntest.TryLoadAgentRun(client, run.ID) - return err == nil && r.Checkpoint != nil - }, - 10*time.Second, - 200*time.Millisecond, - "checkpoint should be saved after stop", - ) - - // Verify checkpoint content. - cp, err := store.Load(context.Background(), run.ID.String()) - require.NoError(t, err) - require.NotNil(t, cp, "checkpoint must exist after suspension") - assert.Equal(t, agent.AgentStatusSuspended, cp.Status) - assert.Equal(t, "worker-agent", cp.AgentName) - assert.True(t, len(cp.Messages) > 0, "checkpoint should contain messages") - - // --- Phase 2: Simulate resume by resetting to PENDING --- - - err = client.WithConn( - context.Background(), - func(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec( - ctx, - `UPDATE agent_runs - SET status = 'PENDING', - started_at = NULL, - lease_owner = NULL, - lease_expires_at = NULL, - updated_at = now() - WHERE id = $1`, - run.ID.String(), - ) - - return err - }, - ) - require.NoError(t, err) - - // Start a fresh supervisor to pick up the resumed run. - supervisor2 := probo.NewAgentRunSupervisor( - client, - store, - registry, - testLogger(), - probo.WithAgentRunSupervisorInterval(500*time.Millisecond), - probo.WithAgentRunSupervisorLeaseDuration(30*time.Second), - ) - - ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) - defer cancel2() - - go func() { _ = supervisor2.Run(ctx2) }() - - // The resumed run should load the checkpoint, call Restore, get the - // second LLM response (stopResponse), and complete. - require.Eventually( - t, - func() bool { - r, err := agentruntest.TryLoadAgentRun(client, run.ID) - return err == nil && r.Status == coredata.AgentRunStatusCompleted - }, - 10*time.Second, - 200*time.Millisecond, - "run should reach COMPLETED after resume", - ) - - completed := agentruntest.LoadAgentRun(t, client, run.ID) - assert.Equal(t, coredata.AgentRunStatusCompleted, completed.Status) - assert.NotNil(t, completed.Result) - assert.Nil(t, completed.Checkpoint, "checkpoint should be cleared after completion") - assert.Nil(t, completed.ErrorMessage) -} - -// --------------------------------------------------------------------------- -// Test 5: SIGTERM battle test — realistic multi-turn security audit -// with parallel tool calls, long-running operations, thinking turns, -// and multiple kill/resume cycles. -// -// Simulated workflow (10 tool-call turns + 1 final response): -// -// Turn 0: [think] scan_repos (single, 800ms) -// Turn 1: [think] fetch_config ×3 (parallel, 300-500ms) -// Turn 2: [think] analyze (single, 1000ms) -// Turn 3: [think] check ×3 (parallel, 400-800ms) -// Turn 4: [think] deep_analysis (single, 1500ms — long running) -// Turn 5: [think] generate ×2 (parallel, 500-600ms) -// Turn 6: [think] cve_lookup (single, 700ms) -// Turn 7: [think] compile (single, 600ms) -// Turn 8: [think] validate + format (parallel, 300-400ms) -// Turn 9: [think] publish (single, 400ms) -// Turn 10: final response — "Security audit complete..." -// -// SIGTERM is sent 3 times at different points, each time interrupting -// during tool execution (sometimes single, sometimes parallel). -// After each kill the checkpoint is verified to show progressive -// accumulation. A final in-process resume runs the remaining turns -// to completion. -// --------------------------------------------------------------------------- - -// workInput is the shared parameter type for all battle-test tools. -type workInput struct { - Task string `json:"task"` - DurationMs int `json:"duration_ms"` -} - -// battleTestResponses returns the full LLM response sequence for a -// simulated security-audit agent. Each tool-call turn includes -// thinking text so the checkpoint messages are realistic. -func battleTestResponses() []*llm.ChatCompletionResponse { - tc := func(id, name, task string, ms int) llm.ToolCall { - return llm.ToolCall{ - ID: id, - Function: llm.FunctionCall{ - Name: name, - Arguments: fmt.Sprintf(`{"task":%q,"duration_ms":%d}`, task, ms), - }, - } - } - - think := func(text string, calls ...llm.ToolCall) *llm.ChatCompletionResponse { - return &llm.ChatCompletionResponse{ - Model: "test-model", - Message: llm.Message{ - Role: llm.RoleAssistant, - Parts: []llm.Part{llm.TextPart{Text: text}}, - ToolCalls: calls, - }, - Usage: llm.Usage{InputTokens: 50, OutputTokens: 30}, - FinishReason: llm.FinishReasonToolCalls, - } - } - - return []*llm.ChatCompletionResponse{ - // Turn 0 — single long scan - think( - "I'll begin the security audit by scanning all repositories to identify codebases, dependency manifests, and access-control configurations.", - tc("tc_0_1", "scan", "scan_repos", 800), - ), - - // Turn 1 — 3 parallel fetches - think( - "Found 3 repositories: api-gateway, auth-service, data-pipeline. Fetching their configurations in parallel to save time.", - tc("tc_1_1", "fetch", "fetch_api_config", 300), - tc("tc_1_2", "fetch", "fetch_auth_config", 400), - tc("tc_1_3", "fetch", "fetch_data_config", 500), - ), - - // Turn 2 — single analysis - think( - "All configurations retrieved. Running a comprehensive vulnerability analysis against the OWASP Top-10 checklist.", - tc("tc_2_1", "analyze", "analyze_configs", 1000), - ), - - // Turn 3 — 3 parallel security checks - think( - "Analysis flagged several areas of concern. Running dependency audit, secret scanning, and IAM permission checks in parallel.", - tc("tc_3_1", "check", "check_dependencies", 600), - tc("tc_3_2", "check", "check_secrets", 400), - tc("tc_3_3", "check", "check_permissions", 800), - ), - - // Turn 4 — single very long deep-dive - think( - "Multiple issues found: 3 outdated dependencies with known CVEs, 2 overly permissive IAM roles. Performing a deep analysis on the critical findings to determine exploitability and blast radius.", - tc("tc_4_1", "analyze", "deep_analysis", 1500), - ), - - // Turn 5 — 2 parallel report sections - think( - "Deep analysis complete. auth-service uses deprecated TLS 1.1 and data-pipeline stores PII unencrypted. Generating the executive summary and detailed findings sections in parallel.", - tc("tc_5_1", "generate", "generate_summary", 500), - tc("tc_5_2", "generate", "generate_findings", 600), - ), - - // Turn 6 — single CVE lookup - think( - "Report sections drafted. Cross-referencing all findings against the NVD and GitHub Advisory databases for known CVE identifiers.", - tc("tc_6_1", "lookup", "cve_lookup", 700), - ), - - // Turn 7 — single compile - think( - "CVE-2026-1234 matches the auth-service TLS vulnerability (CVSS 9.1). Compiling all sections, references, and remediation steps into the final report.", - tc("tc_7_1", "compile", "compile_report", 600), - ), - - // Turn 8 — 2 parallel validation + formatting - think( - "Draft report assembled (12 pages). Running structural validation and PDF formatting concurrently.", - tc("tc_8_1", "validate", "validate_report", 400), - tc("tc_8_2", "format", "format_pdf", 300), - ), - - // Turn 9 — single publish - think( - "Validation passed, PDF formatted. Publishing the finalized audit report to the internal portal.", - tc("tc_9_1", "publish", "publish_report", 400), - ), - - // Turn 10 — final text response - { - Model: "test-model", - Message: llm.Message{ - Role: llm.RoleAssistant, - Parts: []llm.Part{llm.TextPart{Text: "Security audit complete.\n\nFindings:\n- 3 critical (auth-service TLS 1.1, unencrypted PII, CVE-2026-1234)\n- 5 medium (outdated deps, permissive IAM)\n- 4 low (missing rate-limiting, verbose logging)\n\nFull report: https://audits.internal/report-2026-04"}}, - }, - Usage: llm.Usage{InputTokens: 50, OutputTokens: 30}, - FinishReason: llm.FinishReasonStop, - }, - } -} - -// makeBattleTools creates the tool set for the battle test. Every tool -// shares the same handler that sleeps for the requested duration and -// records progress to a shared file. -func makeBattleTools(progressFile string) []agent.Tool { - var mu sync.Mutex - - handler := func(_ context.Context, input workInput) (agent.ToolResult, error) { - // Simulate real work. - time.Sleep(time.Duration(input.DurationMs) * time.Millisecond) - - // Record completion — written AFTER the sleep so the parent's - // step count reflects truly-finished work. - mu.Lock() - - f, err := os.OpenFile(progressFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - mu.Unlock() - return agent.ToolResult{}, err - } - - _, _ = fmt.Fprintln(f, input.Task) - _ = f.Close() - mu.Unlock() - - return agent.ToolResult{Content: fmt.Sprintf("completed: %s", input.Task)}, nil - } - - names := []struct{ name, desc string }{ - {"scan", "Scan repositories for audit targets"}, - {"fetch", "Fetch configuration or source files"}, - {"analyze", "Run vulnerability analysis"}, - {"check", "Execute a specific security check"}, - {"generate", "Generate a report section"}, - {"lookup", "Query external vulnerability databases"}, - {"compile", "Compile report sections into final document"}, - {"validate", "Validate report structure"}, - {"format", "Apply output formatting"}, - {"publish", "Publish report to internal portal"}, - } - - tools := make([]agent.Tool, len(names)) - for i, n := range names { - tools[i] = agent.FunctionTool[workInput](n.name, n.desc, handler) - } - - return tools -} - -func TestAgentRunSupervisor_SIGTERM(t *testing.T) { - // ---- Subprocess mode ---- - if os.Getenv("TEST_SIGTERM_SUBPROCESS") == "1" { - runSIGTERMSubprocess() - return - } - - // ---- Parent mode ---- - client := agentruntest.PGClient(t) - store := coredata.NewPGCheckpointer(client) - - run := agentruntest.InsertPendingRun( - t, - client, - "battle-agent", - []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "Run a full security audit on all repositories."}}}}, - ) - - progressFile := filepath.Join(t.TempDir(), "progress") - - // ---- Helpers ---- - - countSteps := func() int { - data, err := os.ReadFile(progressFile) - if err != nil { - return 0 - } - - n := 0 - - for _, b := range data { - if b == '\n' { - n++ - } - } - - return n - } - - startSubprocess := func(skipResponses int) *exec.Cmd { - cmd := exec.Command( - os.Args[0], - "-test.run=^TestAgentRunSupervisor_SIGTERM$", - "-test.v", - ) - - cmd.Env = append(os.Environ(), - "TEST_SIGTERM_SUBPROCESS=1", - "TEST_SIGTERM_PROGRESS_FILE="+progressFile, - "TEST_SIGTERM_SKIP_RESPONSES="+strconv.Itoa(skipResponses), - ) - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - require.NoError(t, cmd.Start()) - - return cmd - } - - killAndWait := func(cmd *exec.Cmd) { - require.NoError(t, cmd.Process.Signal(syscall.SIGTERM)) - - err := cmd.Wait() - if err == nil { - return - } - - exitErr, ok := errors.AsType[*exec.ExitError](err) - if !ok { - t.Fatalf("subprocess error: %v", err) - } - - ws, ok := exitErr.Sys().(syscall.WaitStatus) - if !ok { - t.Fatalf("subprocess exited with unexpected wait status: %v", exitErr) - } - - if ws.Signaled() && ws.Signal() == syscall.SIGTERM { - t.Logf("subprocess terminated by SIGTERM") - return - } - - t.Fatalf("subprocess exited unexpectedly (signaled=%v signal=%v exit=%d): %v", - ws.Signaled(), ws.Signal(), ws.ExitStatus(), exitErr) - } - - resetToPending := func() { - err := client.WithConn( - context.Background(), - func(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec(ctx, ` - UPDATE agent_runs - SET status = 'PENDING', - started_at = NULL, - lease_owner = NULL, - lease_expires_at = NULL, - updated_at = now() - WHERE id = $1`, - run.ID.String(), - ) - - return err - }, - ) - require.NoError(t, err) - } - - waitForSteps := func(target int) { - require.Eventually( - t, - func() bool { return countSteps() >= target }, - 30*time.Second, - 100*time.Millisecond, - fmt.Sprintf("expected at least %d completed tool executions", target), - ) - } - - verifyCheckpoint := func(phase int) *agent.Checkpoint { - cp, err := store.Load(context.Background(), run.ID.String()) - require.NoError(t, err) - require.NotNil(t, cp, "phase %d: checkpoint must exist", phase) - assert.Equal(t, agent.AgentStatusSuspended, cp.Status, "phase %d", phase) - assert.Equal(t, "battle-agent", cp.AgentName, "phase %d", phase) - assert.Greater(t, len(cp.Messages), 1, "phase %d: checkpoint should have messages", phase) - assert.Greater(t, cp.Turns, 0, "phase %d: checkpoint should have turns", phase) - t.Logf( - " checkpoint: %d messages, %d turns, usage=%+v", - len(cp.Messages), cp.Turns, cp.Usage, - ) - - return cp - } - - // ============================================================ - // Phase 1: SIGTERM during the parallel fetch (turn 1) - // Steps so far: turn0=1(scan) + turn1=3(fetch×3) = 4 - // ============================================================ - t.Log("=== Phase 1: SIGTERM after scan + parallel fetch (4 steps) ===") - - cmd1 := startSubprocess(0) - - waitForSteps(4) - killAndWait(cmd1) - - steps1 := countSteps() - t.Logf(" %d tool executions completed", steps1) - require.GreaterOrEqual(t, steps1, 4) - - cp1 := verifyCheckpoint(1) - require.GreaterOrEqual(t, cp1.Turns, 2, "should have completed at least turns 0-1") - - resetToPending() - - // ============================================================ - // Phase 2: SIGTERM during the parallel security checks (turn 3) - // New steps: turn2=1(analyze) + turn3=3(check×3) = 4 - // ============================================================ - t.Log("=== Phase 2: SIGTERM after analyze + parallel checks (4 more steps) ===") - - cmd2 := startSubprocess(cp1.Turns) - - waitForSteps(steps1 + 4) - killAndWait(cmd2) - - steps2 := countSteps() - t.Logf(" %d tool executions completed (total)", steps2) - require.GreaterOrEqual(t, steps2, steps1+4) - - cp2 := verifyCheckpoint(2) - assert.Greater(t, cp2.Turns, cp1.Turns, "turns should grow") - assert.Greater(t, len(cp2.Messages), len(cp1.Messages), "messages should grow") - - resetToPending() - - // ============================================================ - // Phase 3: SIGTERM during deep_analysis (turn 4, long-running) - // or after generate ×2 (turn 5) - // New steps: turn4=1(deep) + turn5=2(generate×2) = 3 - // ============================================================ - t.Log("=== Phase 3: SIGTERM during long-running deep analysis (3 more steps) ===") - - cmd3 := startSubprocess(cp2.Turns) - - waitForSteps(steps2 + 3) - killAndWait(cmd3) - - steps3 := countSteps() - t.Logf(" %d tool executions completed (total)", steps3) - require.GreaterOrEqual(t, steps3, steps2+3) - - cp3 := verifyCheckpoint(3) - assert.Greater(t, cp3.Turns, cp2.Turns, "turns should grow again") - assert.Greater(t, len(cp3.Messages), len(cp2.Messages), "messages should grow again") - assert.Greater(t, cp3.Usage.InputTokens, 0, "usage should accumulate") - assert.Greater(t, cp3.Usage.OutputTokens, 0, "usage should accumulate") - - t.Logf( - " after 3 SIGTERM cycles: %d steps, %d turns, %d messages, usage=%+v", - steps3, cp3.Turns, len(cp3.Messages), cp3.Usage, - ) - - resetToPending() - - // ============================================================ - // Phase 4: final in-process resume — run remaining turns to - // completion (lookup, compile, validate+format, publish, done) - // ============================================================ - t.Log("=== Phase 4: in-process resume to completion ===") - - remaining := battleTestResponses()[cp3.Turns:] - t.Logf(" %d LLM responses remaining (turns %d–10)", len(remaining), cp3.Turns) - - tools := makeBattleTools(progressFile) - - resumeAgent := agent.New( - "battle-agent", - newTestClient(&mockProvider{responses: remaining}), - agent.WithModel("test-model"), - agent.WithTools(tools...), - agent.WithMaxTurns(25), - ) - - supervisor := probo.NewAgentRunSupervisor( - client, - store, - &simpleRegistry{agents: map[string]*agent.Agent{"battle-agent": resumeAgent}}, - testLogger(), - probo.WithAgentRunSupervisorInterval(500*time.Millisecond), - probo.WithAgentRunSupervisorLeaseDuration(30*time.Second), - ) - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - go func() { _ = supervisor.Run(ctx) }() - - require.Eventually( - t, - func() bool { - r, err := agentruntest.TryLoadAgentRun(client, run.ID) - return err == nil && r.Status == coredata.AgentRunStatusCompleted - }, - 25*time.Second, - 200*time.Millisecond, - "run should complete after final resume", - ) - - stepsFinal := countSteps() - final := agentruntest.LoadAgentRun(t, client, run.ID) - assert.Equal(t, coredata.AgentRunStatusCompleted, final.Status) - assert.NotNil(t, final.Result) - assert.Nil(t, final.Checkpoint, "checkpoint should be cleared") - assert.Nil(t, final.ErrorMessage) - assert.Contains(t, string(final.Result), "Security audit complete") - - t.Logf( - " battle test done: %d total tool executions across 3 SIGTERM cycles + final resume", - stepsFinal, - ) -} - -// runSIGTERMSubprocess is the child-process entry point. It sets up a -// supervisor with the full security-audit agent (10 distinct tools, -// thinking text, parallel calls, varying durations) and handles -// SIGTERM via signal.NotifyContext — identical to production probod. -func runSIGTERMSubprocess() { - progressFile := os.Getenv("TEST_SIGTERM_PROGRESS_FILE") - skip, _ := strconv.Atoi(os.Getenv("TEST_SIGTERM_SKIP_RESPONSES")) - - addr := os.Getenv("PROBO_TEST_PG_ADDR") - if addr == "" { - addr = "localhost:5432" - } - - user := os.Getenv("PROBO_TEST_PG_USER") - if user == "" { - user = "probod" - } - - password := os.Getenv("PROBO_TEST_PG_PASSWORD") - if password == "" { - password = "probod" - } - - database := os.Getenv("PROBO_TEST_PG_DATABASE") - if database == "" { - database = "probod_test" - } - - pgClient, err := pg.NewClient( - pg.WithAddr(addr), - pg.WithUser(user), - pg.WithPassword(password), - pg.WithDatabase(database), - pg.WithPoolSize(5), - ) - if err != nil { - fmt.Fprintf(os.Stderr, "subprocess: cannot create pg client: %v\n", err) - os.Exit(1) - } - defer pgClient.Close() - - store := coredata.NewPGCheckpointer(pgClient) - tools := makeBattleTools(progressFile) - - responses := battleTestResponses() - if skip > 0 && skip < len(responses) { - responses = responses[skip:] - } - - ag := agent.New( - "battle-agent", - newTestClient(&mockProvider{responses: responses}), - agent.WithModel("test-model"), - agent.WithTools(tools...), - agent.WithMaxTurns(25), - ) - - supervisor := probo.NewAgentRunSupervisor( - pgClient, - store, - &simpleRegistry{agents: map[string]*agent.Agent{"battle-agent": ag}}, - log.NewLogger(log.WithFormat(log.FormatPretty)), - probo.WithAgentRunSupervisorInterval(500*time.Millisecond), - probo.WithAgentRunSupervisorLeaseDuration(5*time.Second), - ) - - ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM) - defer stop() - - // signal.NotifyContext sets the cancellation cause to the signal, - // so context.Cause(ctx) returns syscall.SIGTERM rather than - // context.Canceled. Treat any ctx-cancellation outcome as graceful. - err = supervisor.Run(ctx) - if err != nil && ctx.Err() == nil { - fmt.Fprintf(os.Stderr, "subprocess: supervisor error: %v\n", err) - os.Exit(1) - } - - os.Exit(0) -} diff --git a/pkg/agentruntest/agentruntest.go b/pkg/agentruntest/agentruntest.go deleted file mode 100644 index e8bca26e3..000000000 --- a/pkg/agentruntest/agentruntest.go +++ /dev/null @@ -1,247 +0,0 @@ -// 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 agentruntest provides shared test helpers for agent run -// integration tests that require a PostgreSQL database. -package agentruntest - -import ( - "context" - "encoding/json" - "fmt" - "os" - "sync" - "testing" - "time" - - "github.com/stretchr/testify/require" - "go.gearno.de/kit/pg" - "go.probo.inc/probo/pkg/coredata" - "go.probo.inc/probo/pkg/gid" - "go.probo.inc/probo/pkg/llm" -) - -var ( - sharedPGClient *pg.Client - pgOnce sync.Once - pgInitErr error - ensureTableOnce sync.Once - ensureTableErr error -) - -// PGClient returns a shared pg.Client connected to the test database. -// Skips the test if the database is not reachable. -func PGClient(t *testing.T) *pg.Client { - t.Helper() - - pgOnce.Do(func() { - addr := os.Getenv("PROBO_TEST_PG_ADDR") - if addr == "" { - addr = "localhost:5432" - } - - user := os.Getenv("PROBO_TEST_PG_USER") - if user == "" { - user = "probod" - } - - password := os.Getenv("PROBO_TEST_PG_PASSWORD") - if password == "" { - password = "probod" - } - - database := os.Getenv("PROBO_TEST_PG_DATABASE") - if database == "" { - database = "probod_test" - } - - sharedPGClient, pgInitErr = pg.NewClient( - pg.WithAddr(addr), - pg.WithUser(user), - pg.WithPassword(password), - pg.WithDatabase(database), - pg.WithPoolSize(5), - ) - if pgInitErr != nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - pgInitErr = sharedPGClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec(ctx, "SELECT 1") - return err - }) - }) - - if pgInitErr != nil { - t.Skipf("cannot connect to test database: %v", pgInitErr) - } - - EnsureAgentRunsTable(t, sharedPGClient) - - return sharedPGClient -} - -// EnsureAgentRunsTable creates the agent_runs table against the test -// database using the embedded migration, if the table is not already -// present. If the table exists with a stale schema (e.g. missing the -// FK added later), drop it manually or let the production migration -// runner apply the current version — this helper does not rewrite an -// existing table to avoid racing concurrent test processes that share -// the same database. -func EnsureAgentRunsTable(t *testing.T, client *pg.Client) { - t.Helper() - - ensureTableOnce.Do(func() { - ctx := context.Background() - ensureTableErr = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { - var exists bool - if err := conn.QueryRow( - ctx, - `SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'agent_runs')`, - ).Scan(&exists); err != nil { - return fmt.Errorf("cannot check agent_runs existence: %w", err) - } - - if exists { - return nil - } - - ddl, err := coredata.Migrations.ReadFile("migrations/20260424T173529Z.sql") - if err != nil { - return fmt.Errorf("cannot read agent_runs migration: %w", err) - } - - if _, err := conn.Exec(ctx, string(ddl)); err != nil { - return fmt.Errorf("cannot apply agent_runs migration: %w", err) - } - - return nil - }) - }) - require.NoError(t, ensureTableErr, "cannot ensure agent_runs table") -} - -// CleanupAgentRun deletes an agent run by ID. Safe to call from -// t.Cleanup. -func CleanupAgentRun(client *pg.Client, id gid.GID) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - _ = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec(ctx, "DELETE FROM agent_runs WHERE id = $1", id.String()) - return err - }) -} - -// InsertPendingRun inserts a PENDING agent run and registers cleanup. -// A placeholder organization row is created first so the agent_runs FK -// on organization_id is satisfied. -func InsertPendingRun( - t *testing.T, - client *pg.Client, - agentName string, - inputMessages []llm.Message, -) coredata.AgentRun { - t.Helper() - - tenantID := gid.NewTenantID() - orgID := gid.New(tenantID, 1) - runID := gid.New(tenantID, 2) - - inputJSON, err := json.Marshal(inputMessages) - require.NoError(t, err) - - now := time.Now() - - run := coredata.AgentRun{ - ID: runID, - OrganizationID: orgID, - StartAgentName: agentName, - Status: coredata.AgentRunStatusPending, - InputMessages: inputJSON, - CreatedAt: now, - UpdatedAt: now, - } - - err = client.WithTx( - context.Background(), - func(ctx context.Context, tx pg.Tx) error { - if _, err := tx.Exec( - ctx, - `INSERT INTO organizations (id, tenant_id, name, created_at, updated_at) VALUES ($1, $2, $3, $4, $5)`, - orgID.String(), tenantID.String(), "test-org-"+orgID.String(), now, now, - ); err != nil { - return fmt.Errorf("cannot insert placeholder organization: %w", err) - } - - return run.Insert(ctx, tx, coredata.NewScope(tenantID)) - }, - ) - require.NoError(t, err) - - t.Cleanup(func() { - cleanupOrganization(client, orgID) - }) - - return run -} - -// cleanupOrganization deletes the test organization row; the agent_runs -// FK has ON DELETE CASCADE so the associated run is removed too. -func cleanupOrganization(client *pg.Client, id gid.GID) { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - _ = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec(ctx, "DELETE FROM organizations WHERE id = $1", id.String()) - return err - }) -} - -// LoadAgentRun loads an agent run by ID, failing the test on error. -func LoadAgentRun(t *testing.T, client *pg.Client, id gid.GID) coredata.AgentRun { - t.Helper() - - var run coredata.AgentRun - - err := client.WithConn( - context.Background(), - func(ctx context.Context, conn pg.Querier) error { - return run.LoadByID(ctx, conn, coredata.NewNoScope(), id) - }, - ) - if err != nil { - t.Fatalf("cannot load agent run %s: %v", id, err) - } - - return run -} - -// TryLoadAgentRun is a non-fatal variant safe for use inside -// require.Eventually callbacks (which recover panics). -func TryLoadAgentRun(client *pg.Client, id gid.GID) (coredata.AgentRun, error) { - var run coredata.AgentRun - - err := client.WithConn( - context.Background(), - func(ctx context.Context, conn pg.Querier) error { - return run.LoadByID(ctx, conn, coredata.NewNoScope(), id) - }, - ) - - return run, err -} diff --git a/pkg/coredata/agent_run.go b/pkg/coredata/agent_run.go index 763b758d5..c74913698 100644 --- a/pkg/coredata/agent_run.go +++ b/pkg/coredata/agent_run.go @@ -35,19 +35,19 @@ 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"` - LeaseOwner *string `db:"lease_owner"` - LeaseExpiresAt *time.Time `db:"lease_expires_at"` - 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"` + LeaseExpiresAt *time.Time `db:"lease_expires_at"` + LeaseGeneration int64 `db:"lease_generation"` + CreatedAt time.Time `db:"created_at"` + UpdatedAt time.Time `db:"updated_at"` } AgentRuns []*AgentRun @@ -178,8 +178,8 @@ SELECT result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at FROM @@ -231,8 +231,8 @@ SELECT result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at FROM @@ -286,8 +286,8 @@ SELECT result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at FROM @@ -383,8 +383,8 @@ RETURNING result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at; ` @@ -432,8 +432,8 @@ SET result = @result, error_message = @error_message, started_at = @started_at, - lease_owner = @lease_owner, lease_expires_at = @lease_expires_at, + lease_generation = @lease_generation, updated_at = @updated_at WHERE %s @@ -448,8 +448,8 @@ RETURNING result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at; ` @@ -462,8 +462,8 @@ RETURNING "result": e.Result, "error_message": e.ErrorMessage, "started_at": e.StartedAt, - "lease_owner": e.LeaseOwner, "lease_expires_at": e.LeaseExpiresAt, + "lease_generation": e.LeaseGeneration, "updated_at": e.UpdatedAt, } maps.Copy(args, scope.SQLArguments()) @@ -516,6 +516,46 @@ WHERE return nil } +func CommitAgentRunResult( + ctx context.Context, + tx pg.Tx, + e *AgentRun, + leaseGeneration int64, +) (int64, error) { + q := ` +UPDATE agent_runs +SET + status = @status, + 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; +` + + 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, + } + + tag, err := tx.Exec(ctx, q, args) + if err != nil { + return 0, fmt.Errorf("cannot commit agent run result: %w", err) + } + + return tag.RowsAffected(), nil +} + func (e *AgentRun) LoadNextPendingForUpdateSkipLocked( ctx context.Context, tx pg.Tx, @@ -531,8 +571,8 @@ SELECT result, error_message, started_at, - lease_owner, lease_expires_at, + lease_generation, created_at, updated_at FROM @@ -566,7 +606,7 @@ FOR UPDATE SKIP LOCKED; // 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 supervisor auto-resumes +// 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 := ` @@ -574,7 +614,6 @@ UPDATE agent_runs SET status = 'PENDING', started_at = NULL, - lease_owner = NULL, lease_expires_at = NULL, updated_at = now() WHERE @@ -597,7 +636,7 @@ func HeartbeatAgentRunLease( ctx context.Context, conn pg.Querier, runID string, - leaseOwner string, + leaseGeneration int64, expiresAt time.Time, ) (int64, error) { q := ` @@ -608,13 +647,13 @@ SET WHERE id = @id AND status = 'RUNNING' - AND lease_owner = @lease_owner; + AND lease_generation = @lease_generation; ` args := pgx.StrictNamedArgs{ "id": runID, - "lease_owner": leaseOwner, "lease_expires_at": expiresAt, + "lease_generation": leaseGeneration, } tag, err := conn.Exec(ctx, q, args) @@ -699,6 +738,55 @@ 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/20260607T060000Z.sql b/pkg/coredata/migrations/20260607T060000Z.sql new file mode 100644 index 000000000..1688b4a6e --- /dev/null +++ b/pkg/coredata/migrations/20260607T060000Z.sql @@ -0,0 +1,17 @@ +-- 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. + +ALTER TABLE agent_runs + ADD COLUMN lease_generation BIGINT NOT NULL DEFAULT 0, + DROP COLUMN lease_owner; diff --git a/pkg/coredata/pg_checkpointer_test.go b/pkg/coredata/pg_checkpointer_test.go index 2dd962698..fcdf3c169 100644 --- a/pkg/coredata/pg_checkpointer_test.go +++ b/pkg/coredata/pg_checkpointer_test.go @@ -21,7 +21,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.probo.inc/probo/pkg/agent" - "go.probo.inc/probo/pkg/agentruntest" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/llm" @@ -30,7 +29,7 @@ import ( func TestPGCheckpointer(t *testing.T) { t.Parallel() - client := agentruntest.PGClient(t) + client := pgClient(t) store := coredata.NewPGCheckpointer(client) t.Run( @@ -39,7 +38,7 @@ func TestPGCheckpointer(t *testing.T) { t.Parallel() ctx := context.Background() - run := agentruntest.InsertPendingRun( + run := insertPendingRun( t, client, "test-agent", @@ -58,7 +57,7 @@ func TestPGCheckpointer(t *testing.T) { t.Parallel() ctx := context.Background() - run := agentruntest.InsertPendingRun( + run := insertPendingRun( t, client, "test-agent", @@ -100,7 +99,7 @@ func TestPGCheckpointer(t *testing.T) { t.Parallel() ctx := context.Background() - run := agentruntest.InsertPendingRun( + run := insertPendingRun( t, client, "test-agent", @@ -143,7 +142,7 @@ func TestPGCheckpointer(t *testing.T) { t.Parallel() ctx := context.Background() - run := agentruntest.InsertPendingRun( + run := insertPendingRun( t, client, "test-agent", @@ -202,7 +201,7 @@ func TestPGCheckpointer(t *testing.T) { t.Parallel() ctx := context.Background() - run := agentruntest.InsertPendingRun( + run := insertPendingRun( t, client, "test-agent", diff --git a/pkg/coredata/pg_test_helpers_test.go b/pkg/coredata/pg_test_helpers_test.go new file mode 100644 index 000000000..405cd2ffd --- /dev/null +++ b/pkg/coredata/pg_test_helpers_test.go @@ -0,0 +1,202 @@ +// 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 coredata_test + +import ( + "context" + "encoding/json" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/llm" +) + +var ( + sharedPGClientCoredata *pg.Client + pgOnceCoredata sync.Once + pgInitErrCoredata error + ensureTableOnceCoredata sync.Once + ensureTableErrCoredata error +) + +func pgClient(t *testing.T) *pg.Client { + t.Helper() + + pgOnceCoredata.Do(func() { + addr := os.Getenv("PROBO_TEST_PG_ADDR") + if addr == "" { + addr = "localhost:5432" + } + + user := os.Getenv("PROBO_TEST_PG_USER") + if user == "" { + user = "probod" + } + + password := os.Getenv("PROBO_TEST_PG_PASSWORD") + if password == "" { + password = "probod" + } + + database := os.Getenv("PROBO_TEST_PG_DATABASE") + if database == "" { + database = "probod_test" + } + + sharedPGClientCoredata, pgInitErrCoredata = pg.NewClient( + pg.WithAddr(addr), + pg.WithUser(user), + pg.WithPassword(password), + pg.WithDatabase(database), + pg.WithPoolSize(5), + ) + if pgInitErrCoredata != nil { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + pgInitErrCoredata = sharedPGClientCoredata.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + _, err := conn.Exec(ctx, "SELECT 1") + return err + }) + }) + + if pgInitErrCoredata != nil { + t.Skipf("cannot connect to test database: %v", pgInitErrCoredata) + } + + ensureAgentRunsTable(t, sharedPGClientCoredata) + + return sharedPGClientCoredata +} + +func ensureAgentRunsTable(t *testing.T, client *pg.Client) { + t.Helper() + + ensureTableOnceCoredata.Do(func() { + ctx := context.Background() + ensureTableErrCoredata = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + var exists bool + if err := conn.QueryRow( + ctx, + `SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'agent_runs')`, + ).Scan(&exists); err != nil { + return fmt.Errorf("cannot check agent_runs existence: %w", err) + } + + if !exists { + ddl, err := coredata.Migrations.ReadFile("migrations/20260424T173529Z.sql") + if err != nil { + return fmt.Errorf("cannot read agent_runs base migration: %w", err) + } + + if _, err := conn.Exec(ctx, string(ddl)); err != nil { + return fmt.Errorf("cannot apply agent_runs base migration: %w", err) + } + } + + var hasLeaseGeneration bool + if err := conn.QueryRow( + ctx, + `SELECT EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'agent_runs' + AND column_name = 'lease_generation' + )`, + ).Scan(&hasLeaseGeneration); err != nil { + return fmt.Errorf("cannot check lease_generation column: %w", err) + } + + if !hasLeaseGeneration { + ddl, err := coredata.Migrations.ReadFile("migrations/20260607T060000Z.sql") + if err != nil { + return fmt.Errorf("cannot read agent_runs lease generation migration: %w", err) + } + + if _, err := conn.Exec(ctx, string(ddl)); err != nil { + return fmt.Errorf("cannot apply agent_runs lease generation migration: %w", err) + } + } + + return nil + }) + }) + + require.NoError(t, ensureTableErrCoredata, "cannot ensure agent_runs table") +} + +func insertPendingRun( + t *testing.T, + client *pg.Client, + agentName string, + inputMessages []llm.Message, +) coredata.AgentRun { + t.Helper() + + tenantID := gid.NewTenantID() + orgID := gid.New(tenantID, coredata.OrganizationEntityType) + runID := gid.New(tenantID, coredata.AgentRunEntityType) + inputJSON, err := json.Marshal(inputMessages) + require.NoError(t, err) + + now := time.Now() + run := coredata.AgentRun{ + ID: runID, + OrganizationID: orgID, + StartAgentName: agentName, + Status: coredata.AgentRunStatusPending, + InputMessages: inputJSON, + CreatedAt: now, + UpdatedAt: now, + } + + err = client.WithTx( + context.Background(), + func(ctx context.Context, tx pg.Tx) error { + if _, err := tx.Exec( + ctx, + `INSERT INTO organizations (id, tenant_id, name, created_at, updated_at) VALUES ($1, $2, $3, $4, $5)`, + orgID.String(), tenantID.String(), "test-org-"+orgID.String(), now, now, + ); err != nil { + return fmt.Errorf("cannot insert placeholder organization: %w", err) + } + + return run.Insert(ctx, tx, coredata.NewScope(tenantID)) + }, + ) + require.NoError(t, err) + + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _ = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + _, err := conn.Exec(ctx, "DELETE FROM organizations WHERE id = $1", orgID.String()) + return err + }) + }) + + return run +} diff --git a/pkg/probo/actions.go b/pkg/probo/actions.go index e5530151b..b46231777 100644 --- a/pkg/probo/actions.go +++ b/pkg/probo/actions.go @@ -133,6 +133,10 @@ const ( ActionThirdPartyRiskAssessmentCreate = "core:thirdParty-risk-assessment:create" ActionThirdPartyRiskAssessmentList = "core:thirdParty-risk-assessment:list" + // AgentRun actions + ActionAgentRunGet = "core:agent-run:get" + ActionAgentRunList = "core:agent-run:list" + // Framework actions ActionFrameworkGet = "core:framework:get" ActionFrameworkList = "core:framework:list" diff --git a/pkg/probo/policies.go b/pkg/probo/policies.go index e81cdc390..4842a0335 100644 --- a/pkg/probo/policies.go +++ b/pkg/probo/policies.go @@ -56,6 +56,7 @@ var ViewerPolicy = policy.NewPolicy( ActionThirdPartyDataPrivacyAgreementGet, ActionThirdPartyRiskAssessmentList, ActionThirdPartyRelationList, + ActionAgentRunGet, ActionAgentRunList, ActionFrameworkGet, ActionFrameworkList, ActionControlGet, ActionControlList, ActionMeasureGet, ActionMeasureList, @@ -144,6 +145,7 @@ var AuditorPolicy = policy.NewPolicy( ActionThirdPartyDataPrivacyAgreementGet, ActionThirdPartyRiskAssessmentList, ActionThirdPartyRelationList, + ActionAgentRunGet, ActionAgentRunList, ActionFrameworkGet, ActionFrameworkList, ActionControlGet, ActionControlList, ActionMeasureGet, ActionMeasureList, diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go index 2d291f5af..58746c0bc 100644 --- a/pkg/probod/probod.go +++ b/pkg/probod/probod.go @@ -43,6 +43,7 @@ import ( "go.opentelemetry.io/otel/trace" "go.probo.inc/probo/packages/emails" "go.probo.inc/probo/pkg/accessreview" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/awsconfig" "go.probo.inc/probo/pkg/baseurl" "go.probo.inc/probo/pkg/certmanager" @@ -593,6 +594,8 @@ func (impl *Implm) Run( l.Named("access-review"), ) + agentRunService := agentrun.NewService(pgClient) + thirdPartyService := thirdparty.NewService(pgClient, fileService, thirdPartyVetter) riskManagementService := riskmanagement.NewService(pgClient) @@ -606,6 +609,7 @@ func (impl *Implm) Run( Trust: trustService, ESign: esignService, AccessReview: accessReviewService, + AgentRun: agentRunService, Mailman: mailmanService, CookieBanner: cookieBannerService, Geoloc: geolocService, diff --git a/pkg/server/api/api.go b/pkg/server/api/api.go index f19b45453..c37244270 100644 --- a/pkg/server/api/api.go +++ b/pkg/server/api/api.go @@ -26,6 +26,7 @@ import ( "go.gearno.de/kit/httpserver" "go.gearno.de/kit/log" "go.probo.inc/probo/pkg/accessreview" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/baseurl" "go.probo.inc/probo/pkg/connector" "go.probo.inc/probo/pkg/connector/provider" @@ -60,6 +61,7 @@ type ( Trust *trust.Service ESign *esign.Service AccessReview *accessreview.Service + AgentRun *agentrun.Service Slack *slack.Service Mailman *mailman.Service CookieBanner *cookiebanner.Service @@ -189,6 +191,7 @@ func NewServer(cfg Config) (*Server, error) { cfg.IAM, cfg.ESign, cfg.AccessReview, + cfg.AgentRun, cfg.Mailman, cfg.CookieBanner, cfg.Cookie, diff --git a/pkg/server/api/console/v1/agent_run_resolvers.go b/pkg/server/api/console/v1/agent_run_resolvers.go new file mode 100644 index 000000000..0754336a1 --- /dev/null +++ b/pkg/server/api/console/v1/agent_run_resolvers.go @@ -0,0 +1,82 @@ +package console_v1 + +// This file will be automatically regenerated based on the schema, any resolver +// implementations +// will be copied through when generating and any unknown code will be moved to the end. +// Code generated by github.com/99designs/gqlgen version v0.17.90 + +import ( + "context" + "errors" + "fmt" + + "github.com/vikstrous/dataloadgen" + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/probo" + "go.probo.inc/probo/pkg/server/api/console/v1/dataloader" + "go.probo.inc/probo/pkg/server/api/console/v1/schema" + "go.probo.inc/probo/pkg/server/api/console/v1/types" + "go.probo.inc/probo/pkg/server/gqlutils" +) + +// Organization is the resolver for the organization field. +func (r *agentRunResolver) Organization(ctx context.Context, obj *types.AgentRun) (*types.Organization, error) { + if _, err := r.authorize(ctx, obj.ID, probo.ActionOrganizationGet); err != nil { + return nil, err + } + + loaders := dataloader.FromContext(ctx) + + organization, err := loaders.Organization.Load(ctx, obj.Organization.ID) + if err != nil { + if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) { + return nil, gqlutils.NotFound(ctx, err) + } + + r.logger.ErrorCtx(ctx, "cannot load organization", log.Error(err)) + + return nil, gqlutils.Internal(ctx) + } + + return types.NewOrganization(organization), nil +} + +// Permission is the resolver for the permission field. +func (r *agentRunResolver) Permission(ctx context.Context, obj *types.AgentRun, action string) (bool, error) { + return r.Resolver.Permission(ctx, obj, action) +} + +// 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) + if err != nil { + return 0, err + } + + switch obj.Resolver.(type) { + case *organizationResolver: + count, err := r.agentRun.CountForOrganizationID(ctx, scope, obj.ParentID) + if err != nil { + r.logger.ErrorCtx(ctx, "cannot count agent runs", log.Error(err)) + return 0, gqlutils.Internal(ctx) + } + + return count, nil + } + + r.logger.ErrorCtx(ctx, "unsupported resolver for agent run connection", log.String("resolver", fmt.Sprintf("%T", obj.Resolver))) + + return 0, gqlutils.Internal(ctx) +} + +// AgentRun returns schema.AgentRunResolver implementation. +func (r *Resolver) AgentRun() schema.AgentRunResolver { return &agentRunResolver{r} } + +// AgentRunConnection returns schema.AgentRunConnectionResolver implementation. +func (r *Resolver) AgentRunConnection() schema.AgentRunConnectionResolver { + return &agentRunConnectionResolver{r} +} + +type agentRunResolver struct{ *Resolver } +type agentRunConnectionResolver struct{ *Resolver } diff --git a/pkg/server/api/console/v1/base_resolvers.go b/pkg/server/api/console/v1/base_resolvers.go index a150f0075..f3cd8f19a 100644 --- a/pkg/server/api/console/v1/base_resolvers.go +++ b/pkg/server/api/console/v1/base_resolvers.go @@ -369,6 +369,16 @@ func (r *queryResolver) Node(ctx context.Context, id gid.GID) (types.Node, error return types.NewWebhookSubscription(wc), nil } + case coredata.AgentRunEntityType: + action = probo.ActionAgentRunGet + loadNode = func(ctx context.Context, scope *coredata.Scope, id gid.GID) (types.Node, error) { + run, err := r.agentRun.Get(ctx, scope, id) + if err != nil { + return nil, err + } + + return types.NewAgentRun(run), nil + } case coredata.AccessReviewCampaignEntityType: action = probo.ActionAccessReviewCampaignGet loadNode = func(ctx context.Context, scope *coredata.Scope, id gid.GID) (types.Node, error) { diff --git a/pkg/server/api/console/v1/graphql/agent_run.graphql b/pkg/server/api/console/v1/graphql/agent_run.graphql new file mode 100644 index 000000000..3df26f44f --- /dev/null +++ b/pkg/server/api/console/v1/graphql/agent_run.graphql @@ -0,0 +1,60 @@ +enum AgentRunStatus + @goModel(model: "go.probo.inc/probo/pkg/coredata.AgentRunStatus") { + PENDING + @goEnum(value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusPending") + RUNNING + @goEnum(value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusRunning") + SUSPENDED + @goEnum(value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusSuspended") + AWAITING_APPROVAL + @goEnum( + value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusAwaitingApproval" + ) + COMPLETED + @goEnum(value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusCompleted") + FAILED + @goEnum(value: "go.probo.inc/probo/pkg/coredata.AgentRunStatusFailed") +} + +enum AgentRunOrderField + @goModel(model: "go.probo.inc/probo/pkg/coredata.AgentRunOrderField") { + CREATED_AT + @goEnum( + value: "go.probo.inc/probo/pkg/coredata.AgentRunOrderFieldCreatedAt" + ) +} + +input AgentRunOrder + @goModel( + model: "go.probo.inc/probo/pkg/server/api/console/v1/types.AgentRunOrderBy" + ) { + direction: OrderDirection! + field: AgentRunOrderField! +} + +type AgentRun implements Node { + id: ID! + organization: Organization! @goField(forceResolver: true) + agentName: String! + status: AgentRunStatus! + errorMessage: String + startedAt: Datetime + createdAt: Datetime! + updatedAt: Datetime! + + permission(action: String!): Boolean! @goField(forceResolver: true) +} + +type AgentRunConnection + @goModel( + model: "go.probo.inc/probo/pkg/server/api/console/v1/types.AgentRunConnection" + ) { + totalCount: Int! @goField(forceResolver: true) + edges: [AgentRunEdge!]! + pageInfo: PageInfo! +} + +type AgentRunEdge { + cursor: CursorKey! + node: AgentRun! +} diff --git a/pkg/server/api/console/v1/graphql/organization.graphql b/pkg/server/api/console/v1/graphql/organization.graphql index 2ab7003a7..0c245eb9e 100644 --- a/pkg/server/api/console/v1/graphql/organization.graphql +++ b/pkg/server/api/console/v1/graphql/organization.graphql @@ -311,6 +311,14 @@ type Organization implements Node { orderBy: TaskOrder ): TaskConnection! @goField(forceResolver: true) + agentRuns( + first: Int + after: CursorKey + last: Int + before: CursorKey + orderBy: AgentRunOrder + ): AgentRunConnection! @goField(forceResolver: true) + trustCenter: TrustCenter @goField(forceResolver: true) customDomain: CustomDomain @goField(forceResolver: true) trustCenterFiles( diff --git a/pkg/server/api/console/v1/graphql_handler.go b/pkg/server/api/console/v1/graphql_handler.go index 5cd7eeae1..a4eca6abc 100644 --- a/pkg/server/api/console/v1/graphql_handler.go +++ b/pkg/server/api/console/v1/graphql_handler.go @@ -19,6 +19,7 @@ import ( "go.gearno.de/kit/log" "go.probo.inc/probo/pkg/accessreview" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/connector" "go.probo.inc/probo/pkg/connector/provider" "go.probo.inc/probo/pkg/cookiebanner" @@ -39,6 +40,7 @@ func NewGraphQLHandler( proboSvc *probo.Service, esignSvc *esign.Service, accessReviewSvc *accessreview.Service, + agentRunSvc *agentrun.Service, mailmanSvc *mailman.Service, cookieBannerSvc *cookiebanner.Service, connectorRegistry *connector.ConnectorRegistry, @@ -56,6 +58,7 @@ func NewGraphQLHandler( iam: iamSvc, esign: esignSvc, accessReview: accessReviewSvc, + agentRun: agentRunSvc, mailman: mailmanSvc, cookieBanner: cookieBannerSvc, connectorRegistry: connectorRegistry, diff --git a/pkg/server/api/console/v1/organization_resolvers.go b/pkg/server/api/console/v1/organization_resolvers.go index cb5107fb1..f411cb936 100644 --- a/pkg/server/api/console/v1/organization_resolvers.go +++ b/pkg/server/api/console/v1/organization_resolvers.go @@ -1173,6 +1173,36 @@ func (r *organizationResolver) Tasks(ctx context.Context, obj *types.Organizatio return types.NewTaskConnection(page, r, obj.ID), nil } +// AgentRuns is the resolver for the agentRuns field. +func (r *organizationResolver) AgentRuns(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.AgentRunOrderBy) (*types.AgentRunConnection, error) { + scope, err := r.authorize(ctx, obj.ID, probo.ActionAgentRunList) + if err != nil { + return nil, err + } + + pageOrderBy := page.OrderBy[coredata.AgentRunOrderField]{ + Field: coredata.AgentRunOrderFieldCreatedAt, + Direction: page.OrderDirectionDesc, + } + + if orderBy != nil { + pageOrderBy = page.OrderBy[coredata.AgentRunOrderField]{ + Field: orderBy.Field, + Direction: orderBy.Direction, + } + } + + cursor := types.NewCursor(first, after, last, before, pageOrderBy) + + page, err := r.agentRun.ListForOrganizationID(ctx, scope, obj.ID, cursor) + if err != nil { + r.logger.ErrorCtx(ctx, "cannot list organization agent runs", log.Error(err)) + return nil, gqlutils.Internal(ctx) + } + + return types.NewAgentRunConnection(page, r, obj.ID), nil +} + // TrustCenter is the resolver for the trustCenter field. func (r *organizationResolver) TrustCenter(ctx context.Context, obj *types.Organization) (*types.TrustCenter, error) { scope, err := r.authorize(ctx, obj.ID, probo.ActionTrustCenterGet) diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go index 7208c5eda..e76d576bf 100644 --- a/pkg/server/api/console/v1/resolver.go +++ b/pkg/server/api/console/v1/resolver.go @@ -27,6 +27,7 @@ import ( "go.gearno.de/kit/httpserver" "go.gearno.de/kit/log" "go.probo.inc/probo/pkg/accessreview" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/baseurl" "go.probo.inc/probo/pkg/connector" "go.probo.inc/probo/pkg/connector/provider" @@ -55,6 +56,7 @@ type ( iam *iam.Service esign *esign.Service accessReview *accessreview.Service + agentRun *agentrun.Service mailman *mailman.Service cookieBanner *cookiebanner.Service connectorRegistry *connector.ConnectorRegistry @@ -72,6 +74,7 @@ func NewMux( iamSvc *iam.Service, esignSvc *esign.Service, accessReviewSvc *accessreview.Service, + agentRunSvc *agentrun.Service, mailmanSvc *mailman.Service, cookieBannerSvc *cookiebanner.Service, cookieConfig securecookie.Config, @@ -92,6 +95,7 @@ func NewMux( proboSvc, esignSvc, accessReviewSvc, + agentRunSvc, mailmanSvc, cookieBannerSvc, connectorRegistry, diff --git a/pkg/server/api/console/v1/types/agent_run.go b/pkg/server/api/console/v1/types/agent_run.go new file mode 100644 index 000000000..868bf53e8 --- /dev/null +++ b/pkg/server/api/console/v1/types/agent_run.go @@ -0,0 +1,76 @@ +// 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 types + +import ( + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" +) + +type ( + AgentRunOrderBy OrderBy[coredata.AgentRunOrderField] + + AgentRunConnection struct { + TotalCount int + Edges []*AgentRunEdge + PageInfo PageInfo + + Resolver any + ParentID gid.GID + } +) + +func NewAgentRunConnection( + p *page.Page[*coredata.AgentRun, coredata.AgentRunOrderField], + parentType any, + parentID gid.GID, +) *AgentRunConnection { + var edges = make([]*AgentRunEdge, len(p.Data)) + + for i := range edges { + edges[i] = NewAgentRunEdge(p.Data[i], p.Cursor.OrderBy.Field) + } + + return &AgentRunConnection{ + Edges: edges, + PageInfo: *NewPageInfo(p), + + Resolver: parentType, + ParentID: parentID, + } +} + +func NewAgentRunEdge(run *coredata.AgentRun, orderBy coredata.AgentRunOrderField) *AgentRunEdge { + return &AgentRunEdge{ + Cursor: run.CursorKey(orderBy), + Node: NewAgentRun(run), + } +} + +func NewAgentRun(run *coredata.AgentRun) *AgentRun { + return &AgentRun{ + ID: run.ID, + Organization: &Organization{ + ID: run.OrganizationID, + }, + AgentName: run.StartAgentName, + Status: run.Status, + ErrorMessage: run.ErrorMessage, + StartedAt: run.StartedAt, + CreatedAt: run.CreatedAt, + UpdatedAt: run.UpdatedAt, + } +} diff --git a/pkg/server/server.go b/pkg/server/server.go index 380e5d6e5..23d32386e 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -25,6 +25,7 @@ import ( "go.gearno.de/kit/log" "go.gearno.de/x/ref" "go.probo.inc/probo/pkg/accessreview" + "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/baseurl" "go.probo.inc/probo/pkg/connector" "go.probo.inc/probo/pkg/connector/provider" @@ -59,6 +60,7 @@ type Config struct { Trust *trust.Service ESign *esign.Service AccessReview *accessreview.Service + AgentRun *agentrun.Service Slack *slack.Service Mailman *mailman.Service CookieBanner *cookiebanner.Service @@ -97,6 +99,7 @@ func NewServer(cfg Config) (*Server, error) { Trust: cfg.Trust, ESign: cfg.ESign, AccessReview: cfg.AccessReview, + AgentRun: cfg.AgentRun, Slack: cfg.Slack, Mailman: cfg.Mailman, CookieBanner: cfg.CookieBanner,