rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replace supervisor with agentrun worker service

Browse Source 
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 <bryan@probo.com>
This commit is contained in:
Bryan Frimin
2026-06-07 09:00:01 +02:00
parent 400800fd41
commit 3dfc833671
27 changed files with 1076 additions and 1260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
pkg/coredata/agent_run.go
View File

@@ -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)

17
pkg/coredata/migrations/20260607T060000Z.sql Normal file
View File

@@ -0,0 +1,17 @@
-- Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
--
-- Permission to use, copy, modify, and/or distribute this software for any
-- purpose with or without fee is hereby granted, provided that the above
-- copyright notice and this permission notice appear in all copies.
--
-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-- PERFORMANCE OF THIS SOFTWARE.
ALTER TABLE agent_runs
ADD COLUMN lease_generation BIGINT NOT NULL DEFAULT 0,
DROP COLUMN lease_owner;

13
pkg/coredata/pg_checkpointer_test.go
View File

@@ -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",

202
pkg/coredata/pg_test_helpers_test.go Normal file
View File

@@ -0,0 +1,202 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package 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
}