diff --git a/e2e/console/agent_run_test.go b/e2e/console/agent_run_test.go index c33aacbc8..8ac911b86 100644 --- a/e2e/console/agent_run_test.go +++ b/e2e/console/agent_run_test.go @@ -16,14 +16,17 @@ package console_test import ( "context" + "encoding/json" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.probo.inc/probo/e2e/internal/testutil" + "go.probo.inc/probo/pkg/agent" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/llm" ) // agentRunSeed describes the agent run row inserted directly into the test @@ -36,6 +39,7 @@ type agentRunSeed struct { errorMessage *string startedAt *time.Time createdAt time.Time + checkpoint []byte } func seedAgentRun(t *testing.T, organizationID gid.GID, seed agentRunSeed) gid.GID { @@ -59,12 +63,17 @@ func seedAgentRun(t *testing.T, organizationID gid.GID, seed agentRunSeed) gid.G id := gid.New(organizationID.TenantID(), coredata.AgentRunEntityType) + var checkpoint any + if len(seed.checkpoint) > 0 { + checkpoint = string(seed.checkpoint) + } + _, err := conn.Exec(ctx, ` INSERT INTO agent_runs ( id, tenant_id, organization_id, start_agent_name, status, - input_messages, error_message, started_at, created_at, updated_at + input_messages, checkpoint, error_message, started_at, created_at, updated_at ) VALUES ( - $1, $2, $3, $4, $5, $6::jsonb, $7, $8, $9, $9 + $1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8, $9, $10, $10 ) `, id, @@ -73,6 +82,7 @@ func seedAgentRun(t *testing.T, organizationID gid.GID, seed agentRunSeed) gid.G seed.agentName, seed.status, "[]", + checkpoint, seed.errorMessage, seed.startedAt, seed.createdAt, @@ -356,7 +366,7 @@ func TestAgentRun_Get(t *testing.T) { createdAt updatedAt organization { id } - permission(action: "core:agent-run:get") + permission(action: "agentrun:agent-run:get") } } } @@ -486,3 +496,176 @@ func TestAgentRun_TenantIsolation(t *testing.T) { assert.Empty(t, result.Node.AgentRuns.Edges) }) } + +// awaitingApprovalCheckpoint builds the JSON checkpoint a worker persists +// when a run pauses for approval, carrying the pending tool-call IDs the +// approval mutation must reconcile against. +func awaitingApprovalCheckpoint(t *testing.T, toolCallIDs ...string) []byte { + t.Helper() + + approvals := make([]llm.ToolCall, len(toolCallIDs)) + for i, id := range toolCallIDs { + approvals[i] = llm.ToolCall{ + ID: id, + Function: llm.FunctionCall{Name: "danger", Arguments: "{}"}, + } + } + + cp := agent.Checkpoint{ + Status: agent.AgentStatusAwaitingApproval, + AgentName: "approval-agent", + Messages: []llm.Message{ + {Role: llm.RoleAssistant, ToolCalls: approvals}, + }, + PendingToolCalls: approvals, + PendingApprovals: approvals, + } + + data, err := json.Marshal(&cp) + require.NoError(t, err, "cannot marshal approval checkpoint") + + return data +} + +const submitAgentRunApprovalMutation = ` + mutation($input: SubmitAgentRunApprovalInput!) { + submitAgentRunApproval(input: $input) { + agentRun { + id + status + } + } + } +` + +type submitAgentRunApprovalResult struct { + SubmitAgentRunApproval struct { + AgentRun struct { + ID string `json:"id"` + Status string `json:"status"` + } `json:"agentRun"` + } `json:"submitAgentRunApproval"` +} + +func TestAgentRun_SubmitApproval(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "approval-agent", + status: coredata.AgentRunStatusAwaitingApproval, + checkpoint: awaitingApprovalCheckpoint(t, "tc_1"), + }) + + var result submitAgentRunApprovalResult + + err := owner.Execute(submitAgentRunApprovalMutation, map[string]any{ + "input": map[string]any{ + "agentRunId": runID.String(), + "decisions": []map[string]any{ + {"toolCallId": "tc_1", "approved": true}, + }, + }, + }, &result) + require.NoError(t, err) + + // A submitted decision requeues the run so a worker resumes it. + assert.Equal(t, runID.String(), result.SubmitAgentRunApproval.AgentRun.ID) + assert.Equal(t, "PENDING", result.SubmitAgentRunApproval.AgentRun.Status) +} + +func TestAgentRun_SubmitApproval_NotAwaiting(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "approval-agent", + status: coredata.AgentRunStatusCompleted, + }) + + var result submitAgentRunApprovalResult + + err := owner.Execute(submitAgentRunApprovalMutation, map[string]any{ + "input": map[string]any{ + "agentRunId": runID.String(), + "decisions": []map[string]any{ + {"toolCallId": "tc_1", "approved": true}, + }, + }, + }, &result) + testutil.RequireErrorCode(t, err, "CONFLICT") +} + +func TestAgentRun_SubmitApproval_IncompleteDecisions(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + // Two pending approvals, but only one decision is supplied. + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "approval-agent", + status: coredata.AgentRunStatusAwaitingApproval, + checkpoint: awaitingApprovalCheckpoint(t, "tc_1", "tc_2"), + }) + + var result submitAgentRunApprovalResult + + err := owner.Execute(submitAgentRunApprovalMutation, map[string]any{ + "input": map[string]any{ + "agentRunId": runID.String(), + "decisions": []map[string]any{ + {"toolCallId": "tc_1", "approved": true}, + }, + }, + }, &result) + testutil.RequireErrorCode(t, err, "INVALID") +} + +func TestAgentRun_SubmitApproval_RBAC(t *testing.T) { + t.Parallel() + owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, owner.GetOrganizationID(), agentRunSeed{ + agentName: "approval-agent", + status: coredata.AgentRunStatusAwaitingApproval, + checkpoint: awaitingApprovalCheckpoint(t, "tc_1"), + }) + + viewer := testutil.NewClientInOrg(t, testutil.RoleViewer, owner) + + var result submitAgentRunApprovalResult + + err := viewer.Execute(submitAgentRunApprovalMutation, map[string]any{ + "input": map[string]any{ + "agentRunId": runID.String(), + "decisions": []map[string]any{ + {"toolCallId": "tc_1", "approved": true}, + }, + }, + }, &result) + testutil.RequireForbiddenError(t, err, "viewer should not be able to approve agent runs") +} + +func TestAgentRun_SubmitApproval_TenantIsolation(t *testing.T) { + t.Parallel() + + org1Owner := testutil.NewClient(t, testutil.RoleOwner) + org2Owner := testutil.NewClient(t, testutil.RoleOwner) + + runID := seedAgentRun(t, org1Owner.GetOrganizationID(), agentRunSeed{ + agentName: "approval-agent", + status: coredata.AgentRunStatusAwaitingApproval, + checkpoint: awaitingApprovalCheckpoint(t, "tc_1"), + }) + + var result submitAgentRunApprovalResult + + err := org2Owner.Execute(submitAgentRunApprovalMutation, map[string]any{ + "input": map[string]any{ + "agentRunId": runID.String(), + "decisions": []map[string]any{ + {"toolCallId": "tc_1", "approved": true}, + }, + }, + }, &result) + testutil.RequireForbiddenError(t, err, "other org should not be able to approve the run") +} diff --git a/pkg/agentrun/helpers_test.go b/pkg/agentrun/helpers_test.go index ef47ad665..3eace0288 100644 --- a/pkg/agentrun/helpers_test.go +++ b/pkg/agentrun/helpers_test.go @@ -93,7 +93,6 @@ func newTestWorker( baseOpts := []agentrun.WorkerOption{ agentrun.WithWorkerInterval(250 * time.Millisecond), - agentrun.WithWorkerLeaseDuration(30 * time.Second), } baseOpts = append(baseOpts, opts...) @@ -276,7 +275,6 @@ func resetRunToPending(t *testing.T, client *pg.Client, runID gid.GID) { `UPDATE agent_runs SET status = 'PENDING', started_at = NULL, - lease_expires_at = NULL, updated_at = now() WHERE id = $1`, runID.String(), @@ -314,24 +312,3 @@ func overwriteRunInputMessagesRaw( ) 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) -} diff --git a/pkg/agentrun/internal_test.go b/pkg/agentrun/internal_test.go index bca126159..639d2b719 100644 --- a/pkg/agentrun/internal_test.go +++ b/pkg/agentrun/internal_test.go @@ -79,21 +79,6 @@ func TestWorkerOptions(t *testing.T) { }, ) - 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) { diff --git a/pkg/agentrun/service_test.go b/pkg/agentrun/service_test.go index 9e0ccacd5..c1ff729c5 100644 --- a/pkg/agentrun/service_test.go +++ b/pkg/agentrun/service_test.go @@ -21,6 +21,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.probo.inc/probo/internal/test" + "go.probo.inc/probo/pkg/agent" "go.probo.inc/probo/pkg/agentrun" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" @@ -81,6 +82,23 @@ func TestService_ListForOrganizationID(t *testing.T) { assert.True(t, ids[runB.ID]) } +func TestService_SubmitApproval_NotAwaitingApproval(t *testing.T) { + client := test.PGClient(t) + svc := agentrun.NewService(client) + + // A freshly inserted run is PENDING, not AWAITING_APPROVAL. + run := insertPendingRun(t, client, "service-approval-agent", nil) + + _, err := svc.SubmitApproval( + context.Background(), + coredata.NewNoScope(), + run.ID, + map[string]agent.ApprovalResult{"tc_x": {Approved: true}}, + ) + require.Error(t, err) + assert.ErrorIs(t, err, agentrun.ErrNotAwaitingApproval) +} + func TestService_CountForOrganizationID(t *testing.T) { client := test.PGClient(t) svc := agentrun.NewService(client) diff --git a/pkg/agentrun/worker_test.go b/pkg/agentrun/worker_test.go index 8a8509f33..c4daac34d 100644 --- a/pkg/agentrun/worker_test.go +++ b/pkg/agentrun/worker_test.go @@ -24,6 +24,7 @@ import ( "os/signal" "strings" "sync" + "sync/atomic" "syscall" "testing" "time" @@ -143,23 +144,38 @@ func TestWorker_StopAndResume(t *testing.T) { close(toolRelease) + // Graceful shutdown must commit the run back to PENDING (with its + // checkpoint intact) so another worker resumes it. Nothing relies on + // a lease timeout to requeue it. require.Eventually( t, func() bool { r, err := tryLoadAgentRun(client, run.ID) - return err == nil && r.Checkpoint != nil + return err == nil && + r.Status == coredata.AgentRunStatusPending && + r.Checkpoint != nil }, 10*time.Second, 200*time.Millisecond, ) + suspended := loadAgentRun(t, client, run.ID) + assert.Equal( + t, + coredata.AgentRunStatusPending, + suspended.Status, + "graceful shutdown must requeue the run as PENDING without manual recovery", + ) + assert.Nil(t, suspended.Result) + assert.Nil(t, suspended.ErrorMessage) + cp, err := store.Load(context.Background(), run.ID.String()) require.NoError(t, err) require.NotNil(t, cp) assert.Equal(t, agent.AgentStatusSuspended, cp.Status) - resetRunToPending(t, client, run.ID) - + // No manual reset: the run is already PENDING from the graceful + // shutdown, so a fresh worker must pick it up and resume on its own. runWorker2 := newTestWorker( client, &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, @@ -187,6 +203,270 @@ func TestWorker_StopAndResume(t *testing.T) { assert.Nil(t, completed.ErrorMessage) } +// TestWorker_AwaitsApprovalDoesNotFail covers the regression where a tool +// call requiring approval surfaced as InterruptedError and was committed +// as FAILED. The run must instead park in AWAITING_APPROVAL with its +// checkpoint (and the pending approvals) preserved, and must not be +// re-claimed while it rests. +func TestWorker_AwaitsApprovalDoesNotFail(t *testing.T) { + client := test.PGClient(t) + store := coredata.NewPGCheckpointer(client) + + dangerTool := agent.FunctionTool[struct{}]( + "danger", + "Performs a dangerous action", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + return agent.ToolResult{Content: "must not run before approval"}, nil + }, + ) + + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_danger", + Function: llm.FunctionCall{Name: "danger", Arguments: `{}`}, + }), + }, + } + + ag := agent.New( + "approval-agent", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(dangerTool), + agent.WithApproval(agent.ApprovalConfig{ToolNames: []string{"danger"}}), + ) + + run := insertPendingRun( + t, + client, + "approval-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do the dangerous thing"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"approval-agent": ag}}, + ) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + go func() { _ = runWorker.Run(ctx) }() + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusAwaitingApproval + }, + 10*time.Second, + 200*time.Millisecond, + ) + + awaiting := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusAwaitingApproval, awaiting.Status) + assert.Nil(t, awaiting.Result) + assert.Nil(t, awaiting.ErrorMessage) + assert.NotNil(t, awaiting.Checkpoint) + + cp, err := store.Load(context.Background(), run.ID.String()) + require.NoError(t, err) + require.NotNil(t, cp) + assert.Equal(t, agent.AgentStatusAwaitingApproval, cp.Status) + require.Len(t, cp.PendingApprovals, 1) + assert.Equal(t, "danger", cp.PendingApprovals[0].Function.Name) + + // The run must stay parked: only one mock response exists, so a + // re-claim would error with "no more mock responses" and flip it to + // FAILED. Confirm it holds AWAITING_APPROVAL. + time.Sleep(time.Second) + + stillAwaiting := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusAwaitingApproval, stillAwaiting.Status) +} + +// TestWorker_ApprovalApprovedResumesAndCompletes is the full happy-path +// approval cycle: the run parks in AWAITING_APPROVAL, a decision approves +// the pending tool call via the service, and the same worker resumes from +// the checkpoint, executes the approved tool, and completes. +func TestWorker_ApprovalApprovedResumesAndCompletes(t *testing.T) { + client := test.PGClient(t) + + var executed atomic.Bool + + dangerTool := agent.FunctionTool[struct{}]( + "danger", + "Performs a dangerous action", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + executed.Store(true) + + return agent.ToolResult{Content: "danger executed"}, nil + }, + ) + + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_danger", + Function: llm.FunctionCall{Name: "danger", Arguments: `{}`}, + }), + stopResponse("all done"), + }, + } + + ag := agent.New( + "approval-agent", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(dangerTool), + agent.WithApproval(agent.ApprovalConfig{ToolNames: []string{"danger"}}), + ) + + run := insertPendingRun( + t, + client, + "approval-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"approval-agent": ag}}, + ) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + go func() { _ = runWorker.Run(ctx) }() + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusAwaitingApproval + }, + 10*time.Second, + 200*time.Millisecond, + ) + + svc := agentrun.NewService(client) + _, err := svc.SubmitApproval( + context.Background(), + coredata.NewNoScope(), + run.ID, + map[string]agent.ApprovalResult{"tc_danger": {Approved: true}}, + ) + require.NoError(t, err) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusCompleted + }, + 10*time.Second, + 200*time.Millisecond, + ) + + completed := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusCompleted, completed.Status) + assert.NotNil(t, completed.Result) + assert.Nil(t, completed.Checkpoint) + assert.Nil(t, completed.ErrorMessage) + assert.True(t, executed.Load(), "approved tool must execute on resume") +} + +// TestWorker_ApprovalDeniedResumesAndCompletes covers the denial path: the +// run resumes without executing the gated tool and completes, with the +// denial fed back to the model as the tool result. +func TestWorker_ApprovalDeniedResumesAndCompletes(t *testing.T) { + client := test.PGClient(t) + + var executed atomic.Bool + + dangerTool := agent.FunctionTool[struct{}]( + "danger", + "Performs a dangerous action", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + executed.Store(true) + + return agent.ToolResult{Content: "danger executed"}, nil + }, + ) + + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_danger", + Function: llm.FunctionCall{Name: "danger", Arguments: `{}`}, + }), + stopResponse("acknowledged the denial"), + }, + } + + ag := agent.New( + "approval-agent", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(dangerTool), + agent.WithApproval(agent.ApprovalConfig{ToolNames: []string{"danger"}}), + ) + + run := insertPendingRun( + t, + client, + "approval-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"approval-agent": ag}}, + ) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + go func() { _ = runWorker.Run(ctx) }() + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusAwaitingApproval + }, + 10*time.Second, + 200*time.Millisecond, + ) + + svc := agentrun.NewService(client) + _, err := svc.SubmitApproval( + context.Background(), + coredata.NewNoScope(), + run.ID, + map[string]agent.ApprovalResult{"tc_danger": {Approved: false, Message: "denied by reviewer"}}, + ) + require.NoError(t, err) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusCompleted + }, + 10*time.Second, + 200*time.Millisecond, + ) + + completed := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusCompleted, completed.Status) + assert.NotNil(t, completed.Result) + assert.Nil(t, completed.Checkpoint) + assert.Nil(t, completed.ErrorMessage) + assert.False(t, executed.Load(), "denied tool must not execute on resume") +} + // TestWorker_StopAndResumeAcrossHandoff exercises tree suspension where the // active branch is a handed-off child agent. The checkpoint must record the // child as active, and restore must resolve it from the registry so the @@ -278,7 +558,9 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { t, func() bool { r, err := tryLoadAgentRun(client, run.ID) - return err == nil && r.Checkpoint != nil + return err == nil && + r.Status == coredata.AgentRunStatusPending && + r.Checkpoint != nil }, 10*time.Second, 200*time.Millisecond, @@ -295,8 +577,6 @@ func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { "checkpoint must record the handed-off child as the active agent", ) - resetRunToPending(t, client, run.ID) - runWorker2 := newTestWorker(client, registry) ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) @@ -410,7 +690,9 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { t, func() bool { r, err := tryLoadAgentRun(client, run.ID) - return err == nil && r.Checkpoint != nil + return err == nil && + r.Status == coredata.AgentRunStatusPending && + r.Checkpoint != nil }, 10*time.Second, 200*time.Millisecond, @@ -428,8 +710,6 @@ func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { assert.Equal(t, "inner-agent", innerCP.AgentName) assert.Equal(t, agent.AgentStatusSuspended, innerCP.Status) - resetRunToPending(t, client, run.ID) - runWorker2 := newTestWorker(client, registry) ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) @@ -558,7 +838,9 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { t, func() bool { r, err := tryLoadAgentRun(client, run.ID) - return err == nil && r.Checkpoint != nil + return err == nil && + r.Status == coredata.AgentRunStatusPending && + r.Checkpoint != nil }, 10*time.Second, 200*time.Millisecond, @@ -580,8 +862,6 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { assert.Equal(t, "grandchild-agent", grandchildCP.AgentName) assert.Equal(t, agent.AgentStatusSuspended, grandchildCP.Status) - resetRunToPending(t, client, run.ID) - runWorker2 := newTestWorker(client, registry) ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) @@ -606,98 +886,13 @@ func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { assert.Nil(t, completed.ErrorMessage) } -func TestWorker_HeartbeatLeaseLostLeavesRunForRecovery(t *testing.T) { - client := test.PGClient(t) - - toolReady := make(chan struct{}) - toolRelease := make(chan struct{}) - - var readyOnce sync.Once - - slowTool := agent.FunctionTool[struct{}]( - "slow_work", - "Does slow work", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - readyOnce.Do(func() { close(toolReady) }) - <-toolRelease - - return agent.ToolResult{Content: "work done"}, nil - }, - ) - - ag := newDummyAgent( - "worker-agent", - []*llm.ChatCompletionResponse{ - toolCallResponse( - llm.ToolCall{ - ID: "tc_heartbeat", - Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, - }, - ), - stopResponse("done"), - }, - slowTool, - ) - - run := insertPendingRun( - t, - client, - "worker-agent", - []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, - ) - - leaseDuration := 300 * time.Millisecond - runWorker := newTestWorker( - client, - &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, - agentrun.WithWorkerInterval(100*time.Millisecond), - agentrun.WithWorkerLeaseDuration(leaseDuration), - agentrun.WithWorkerMaxConcurrency(1), - ) - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - go func() { - _ = runWorker.Run(ctx) - }() - - select { - case <-toolReady: - case <-ctx.Done(): - t.Fatal("timed out waiting for tool to start") - } - - // Simulate another worker takeover by changing the lease generation. - bumpRunLeaseGeneration(t, client, run.ID) - - // Keep the tool blocked long enough for the heartbeat goroutine to - // observe rowsAffected=0 and cancel this run with ErrLeaseLost. - time.Sleep(2 * leaseDuration) - close(toolRelease) - - require.Eventually( - t, - func() bool { - r, err := tryLoadAgentRun(client, run.ID) - if err != nil { - return false - } - - return r.Status == coredata.AgentRunStatusRunning - }, - 10*time.Second, - 100*time.Millisecond, - ) - - current := loadAgentRun(t, client, run.ID) - assert.Equal(t, coredata.AgentRunStatusRunning, current.Status) - assert.Nil(t, current.Result) - assert.Nil(t, current.ErrorMessage) - assert.NotNil(t, current.LeaseExpiresAt) - assert.Equal(t, int64(2), current.LeaseGeneration) -} - +// TestWorker_ReclaimedRunDoesNotClobberWinner simulates the residual +// manual-recovery risk now that leasing is gone: a human moves a still +// in-flight run back to PENDING (resetRunToPending) while worker A is +// blocked in a tool. Worker B then claims and finishes it. When worker A +// finally returns, its commit must be discarded because the row is no +// longer RUNNING. The CommitAgentRunResult `status = 'RUNNING'` guard is +// the only fence protecting the winner's result. func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { client := test.PGClient(t) @@ -743,7 +938,6 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { runWorkerA := newTestWorker( client, &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, - agentrun.WithWorkerLeaseDuration(5*time.Second), agentrun.WithWorkerMaxConcurrency(1), ) @@ -763,7 +957,6 @@ func TestWorker_ReclaimedRunDoesNotClobberWinner(t *testing.T) { runWorkerB := newTestWorker( client, &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, - agentrun.WithWorkerLeaseDuration(5*time.Second), agentrun.WithWorkerMaxConcurrency(1), )