Update agent-run tests for lease-free worker

Drop the lease, heartbeat, and stale-recovery test coverage that no
longer applies, and add coverage for the new behaviour: graceful
suspend returning a run to PENDING, an approval interruption parking it
in AWAITING_APPROVAL, and SubmitApproval merging decisions and
requeuing the run. Exercise the submitAgentRunApproval mutation end to
end, including the authorization paths for the relocated policies.

Signed-off-by: Bryan Frimin <bryan@probo.com>

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")
+}

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

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

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)

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) {
+// TestWorker_ReclaimedRunDoesNotClobberWinner simulates the residual
-	client := test.PGClient(t)
+// manual-recovery risk now that leasing is gone: a human moves a still
-
+// in-flight run back to PENDING (resetRunToPending) while worker A is
-	toolReady := make(chan struct{})
+// blocked in a tool. Worker B then claims and finishes it. When worker A
-	toolRelease := make(chan struct{})
+// finally returns, its commit must be discarded because the row is no
-
+// longer RUNNING. The CommitAgentRunResult `status = 'RUNNING'` guard is
-	var readyOnce sync.Once
+// the only fence protecting the winner's result.
-
-	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)
-}
-
 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),
 	)