Make agent-tool subtrees suspendable
Propagate graceful-suspend signals through detached run contexts and let only opt-in suspendable tools re-attach cancellation, so AsTool sub-agents can checkpoint and restore across nested trees while leaf tools keep running detached. Add focused agent and worker tests for single and multi-level suspend/ restore flows, plus heartbeat lease-loss and nested-restore error paths to harden functional behavior under failure conditions. Signed-off-by: Bryan Frimin <bryan@probo.com>
This commit is contained in:
Bryan Frimin
452
pkg/agentrun/helpers_test.go
Normal file
452
pkg/agentrun/helpers_test.go
Normal file
@@ -0,0 +1,452 @@
|
||||
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
|
||||
//
|
||||
// Permission to use, copy, modify, and/or distribute this software for any
|
||||
// purpose with or without fee is hereby granted, provided that the above
|
||||
// copyright notice and this permission notice appear in all copies.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
|
||||
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
|
||||
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
|
||||
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
|
||||
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
|
||||
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
// PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
package agentrun_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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/agentrun"
|
||||
"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
|
||||
)
|
||||
|
||||
func testLogger() *log.Logger {
|
||||
return log.NewLogger(log.WithFormat(log.FormatPretty))
|
||||
}
|
||||
|
||||
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 newDummyAgent(name string, responses []*llm.ChatCompletionResponse, tools ...agent.Tool) *agent.Agent {
|
||||
provider := &mockProvider{
|
||||
responses: responses,
|
||||
}
|
||||
|
||||
opts := []agent.Option{
|
||||
agent.WithModel("test-model"),
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
opts = append(opts, agent.WithTools(tools...))
|
||||
}
|
||||
|
||||
return agent.New(
|
||||
name,
|
||||
newTestClient(provider),
|
||||
opts...,
|
||||
)
|
||||
}
|
||||
|
||||
func newTestWorker(
|
||||
client *pg.Client,
|
||||
registry agent.AgentRegistry,
|
||||
opts ...agentrun.WorkerOption,
|
||||
) *agentrun.Worker {
|
||||
store := coredata.NewPGCheckpointer(client)
|
||||
|
||||
baseOpts := []agentrun.WorkerOption{
|
||||
agentrun.WithWorkerInterval(250 * time.Millisecond),
|
||||
agentrun.WithWorkerLeaseDuration(30 * time.Second),
|
||||
}
|
||||
|
||||
baseOpts = append(baseOpts, opts...)
|
||||
|
||||
return agentrun.NewWorker(
|
||||
client,
|
||||
store,
|
||||
registry,
|
||||
testLogger(),
|
||||
baseOpts...,
|
||||
)
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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 {
|
||||
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, ensureTableErr, "cannot ensure agent_runs table")
|
||||
}
|
||||
|
||||
func insertTestOrganization(t *testing.T, client *pg.Client) gid.GID {
|
||||
t.Helper()
|
||||
|
||||
tenantID := gid.NewTenantID()
|
||||
orgID := gid.New(tenantID, coredata.OrganizationEntityType)
|
||||
now := time.Now()
|
||||
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
_, err := conn.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,
|
||||
)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
t.Cleanup(func() {
|
||||
cleanupOrganization(client, orgID)
|
||||
})
|
||||
|
||||
return orgID
|
||||
}
|
||||
|
||||
func insertPendingRun(
|
||||
t *testing.T,
|
||||
client *pg.Client,
|
||||
agentName string,
|
||||
inputMessages []llm.Message,
|
||||
) coredata.AgentRun {
|
||||
t.Helper()
|
||||
|
||||
orgID := insertTestOrganization(t, client)
|
||||
|
||||
return insertPendingRunInOrg(t, client, orgID, agentName, inputMessages)
|
||||
}
|
||||
|
||||
func insertPendingRunInOrg(
|
||||
t *testing.T,
|
||||
client *pg.Client,
|
||||
organizationID gid.GID,
|
||||
agentName string,
|
||||
inputMessages []llm.Message,
|
||||
) coredata.AgentRun {
|
||||
t.Helper()
|
||||
|
||||
runID := gid.New(organizationID.TenantID(), coredata.AgentRunEntityType)
|
||||
|
||||
inputJSON, err := json.Marshal(inputMessages)
|
||||
require.NoError(t, err)
|
||||
|
||||
now := time.Now()
|
||||
|
||||
run := coredata.AgentRun{
|
||||
ID: runID,
|
||||
OrganizationID: organizationID,
|
||||
StartAgentName: agentName,
|
||||
Status: coredata.AgentRunStatusPending,
|
||||
InputMessages: inputJSON,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
err = client.WithTx(
|
||||
context.Background(),
|
||||
func(ctx context.Context, tx pg.Tx) error {
|
||||
return run.Insert(ctx, tx, coredata.NewScope(organizationID.TenantID()))
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
return run
|
||||
}
|
||||
|
||||
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
|
||||
})
|
||||
}
|
||||
|
||||
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)
|
||||
},
|
||||
)
|
||||
require.NoError(t, err, "cannot load agent run %s", id)
|
||||
|
||||
return run
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
func resetRunToPending(t *testing.T, client *pg.Client, runID gid.GID) {
|
||||
t.Helper()
|
||||
|
||||
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_expires_at = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1`,
|
||||
runID.String(),
|
||||
)
|
||||
return err
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func overwriteRunInputMessagesRaw(
|
||||
t *testing.T,
|
||||
client *pg.Client,
|
||||
runID gid.GID,
|
||||
rawJSON string,
|
||||
) {
|
||||
t.Helper()
|
||||
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
_, err := conn.Exec(
|
||||
ctx,
|
||||
`UPDATE agent_runs
|
||||
SET input_messages = $2::jsonb,
|
||||
updated_at = now()
|
||||
WHERE id = $1`,
|
||||
runID.String(),
|
||||
rawJSON,
|
||||
)
|
||||
|
||||
return err
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func bumpRunLeaseGeneration(t *testing.T, client *pg.Client, runID gid.GID) {
|
||||
t.Helper()
|
||||
|
||||
err := client.WithConn(
|
||||
context.Background(),
|
||||
func(ctx context.Context, conn pg.Querier) error {
|
||||
_, err := conn.Exec(
|
||||
ctx,
|
||||
`UPDATE agent_runs
|
||||
SET lease_generation = lease_generation + 1,
|
||||
updated_at = now()
|
||||
WHERE id = $1`,
|
||||
runID.String(),
|
||||
)
|
||||
return err
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
Reference in New Issue
Block a user