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>

pkg/agent/agent_tool.go

@@ -41,6 +41,8 @@ type (
 
 var (
 	agentToolParamsSchema = mustJSONSchemaFor[agentToolParams]()
+
+	_ SuspendableTool = (*agentTool)(nil)
 )
 
 func agentToolDepth(ctx context.Context) int {
@@ -62,6 +64,8 @@ func newAgentTool(agent *Agent, name, description string) *agentTool {
 
 func (t *agentTool) Name() string { return t.toolName }
 
+func (t *agentTool) Suspendable() {}
+
 func (t *agentTool) Definition() llm.Tool {
 	return llm.Tool{
 		Name:        t.toolName,

pkg/agent/agent_tool_test.go

@@ -17,7 +17,10 @@ package agent_test
 import (
 	"context"
 	"encoding/json"
+	"sync"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -915,6 +918,362 @@ func TestAgentTool_Execute_DepthLimit(t *testing.T) {
 	)
 }
 
+func TestAgentTool_Execute_SuspendAndRestoreSingleLevel(t *testing.T) {
+	t.Parallel()
+
+	store := newMemoryCheckpointer()
+
+	toolReady := make(chan struct{})
+	toolRelease := make(chan struct{})
+
+	var readyOnce sync.Once
+
+	slowTool := agent.FunctionTool[struct{}](
+		"slow_inner_work",
+		"Slow inner work",
+		func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+			readyOnce.Do(func() { close(toolReady) })
+			<-toolRelease
+
+			return agent.ToolResult{Content: "inner tool done"}, nil
+		},
+	)
+
+	innerProvider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_inner",
+					Function: llm.FunctionCall{Name: "slow_inner_work", Arguments: `{}`},
+				},
+			),
+			stopResponse("inner completed"),
+		},
+	}
+
+	innerAgent := agent.New(
+		"inner-agent",
+		newTestClient(innerProvider),
+		agent.WithModel("test-model"),
+		agent.WithTools(slowTool),
+	)
+
+	outerProvider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_outer",
+					Function: llm.FunctionCall{Name: "call_inner", Arguments: `{"input":"delegate"}`},
+				},
+			),
+			stopResponse("outer completed"),
+		},
+	}
+
+	outerAgent := agent.New(
+		"outer-agent",
+		newTestClient(outerProvider),
+		agent.WithModel("test-model"),
+		agent.WithTools(innerAgent.AsTool("call_inner", "Call inner")),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	errCh := make(chan error, 1)
+
+	go func() {
+		_, err := outerAgent.Run(
+			ctx,
+			[]llm.Message{userMessage("go")},
+			agent.WithCheckpointer(store, "run-single-level"),
+		)
+		errCh <- err
+	}()
+
+	select {
+	case <-toolReady:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for inner leaf tool to start")
+	}
+
+	cancel()
+	time.Sleep(50 * time.Millisecond)
+	close(toolRelease)
+
+	select {
+	case err := <-errCh:
+		var se *agent.SuspendedError
+		require.ErrorAs(t, err, &se)
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for run to suspend")
+	}
+
+	cp, err := store.Load(context.Background(), "run-single-level")
+	require.NoError(t, err)
+	require.NotNil(t, cp)
+	assert.Equal(t, agent.AgentStatusSuspended, cp.Status)
+	assert.Equal(t, "outer-agent", cp.AgentName)
+
+	innerCP, ok := cp.InnerCheckpoints["tc_outer"]
+	require.True(t, ok, "expected nested checkpoint keyed by outer tool call")
+	require.NotNil(t, innerCP)
+	assert.Equal(t, "inner-agent", innerCP.AgentName)
+	assert.Equal(t, agent.AgentStatusSuspended, innerCP.Status)
+
+	registry := &simpleRegistry{
+		agents: map[string]*agent.Agent{
+			"outer-agent": outerAgent,
+			"inner-agent": innerAgent,
+		},
+	}
+
+	result, err := agent.Restore(
+		context.Background(),
+		store,
+		"run-single-level",
+		registry,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, "outer completed", result.FinalMessage().Text())
+	assert.Equal(t, "outer-agent", result.LastAgent.Name())
+}
+
+func TestAgentTool_Execute_SuspendAndRestoreMultiLevel(t *testing.T) {
+	t.Parallel()
+
+	store := newMemoryCheckpointer()
+
+	toolReady := make(chan struct{})
+	toolRelease := make(chan struct{})
+
+	var readyOnce sync.Once
+
+	slowTool := agent.FunctionTool[struct{}](
+		"slow_grandchild_work",
+		"Slow grandchild work",
+		func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+			readyOnce.Do(func() { close(toolReady) })
+			<-toolRelease
+
+			return agent.ToolResult{Content: "grandchild tool done"}, nil
+		},
+	)
+
+	grandchildProvider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_grandchild",
+					Function: llm.FunctionCall{Name: "slow_grandchild_work", Arguments: `{}`},
+				},
+			),
+			stopResponse("grandchild completed"),
+		},
+	}
+
+	grandchildAgent := agent.New(
+		"grandchild-agent",
+		newTestClient(grandchildProvider),
+		agent.WithModel("test-model"),
+		agent.WithTools(slowTool),
+	)
+
+	childProvider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_child",
+					Function: llm.FunctionCall{Name: "call_grandchild", Arguments: `{"input":"delegate deeper"}`},
+				},
+			),
+			stopResponse("child completed"),
+		},
+	}
+
+	childAgent := agent.New(
+		"child-agent",
+		newTestClient(childProvider),
+		agent.WithModel("test-model"),
+		agent.WithTools(grandchildAgent.AsTool("call_grandchild", "Call grandchild")),
+	)
+
+	outerProvider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_outer",
+					Function: llm.FunctionCall{Name: "call_child", Arguments: `{"input":"delegate"}`},
+				},
+			),
+			stopResponse("outer completed"),
+		},
+	}
+
+	outerAgent := agent.New(
+		"outer-agent",
+		newTestClient(outerProvider),
+		agent.WithModel("test-model"),
+		agent.WithTools(childAgent.AsTool("call_child", "Call child")),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	errCh := make(chan error, 1)
+
+	go func() {
+		_, err := outerAgent.Run(
+			ctx,
+			[]llm.Message{userMessage("go")},
+			agent.WithCheckpointer(store, "run-multi-level"),
+		)
+		errCh <- err
+	}()
+
+	select {
+	case <-toolReady:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for grandchild leaf tool to start")
+	}
+
+	cancel()
+	time.Sleep(50 * time.Millisecond)
+	close(toolRelease)
+
+	select {
+	case err := <-errCh:
+		var se *agent.SuspendedError
+		require.ErrorAs(t, err, &se)
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for multi-level run to suspend")
+	}
+
+	cp, err := store.Load(context.Background(), "run-multi-level")
+	require.NoError(t, err)
+	require.NotNil(t, cp)
+	assert.Equal(t, "outer-agent", cp.AgentName)
+
+	childCP, ok := cp.InnerCheckpoints["tc_outer"]
+	require.True(t, ok)
+	require.NotNil(t, childCP)
+	assert.Equal(t, "child-agent", childCP.AgentName)
+
+	grandchildCP, ok := childCP.InnerCheckpoints["tc_child"]
+	require.True(t, ok, "child checkpoint keys: %v", childCP.InnerCheckpoints)
+	require.NotNil(t, grandchildCP)
+	assert.Equal(t, "grandchild-agent", grandchildCP.AgentName)
+	assert.Equal(t, agent.AgentStatusSuspended, grandchildCP.Status)
+
+	registry := &simpleRegistry{
+		agents: map[string]*agent.Agent{
+			"outer-agent":      outerAgent,
+			"child-agent":      childAgent,
+			"grandchild-agent": grandchildAgent,
+		},
+	}
+
+	result, err := agent.Restore(
+		context.Background(),
+		store,
+		"run-multi-level",
+		registry,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, "outer completed", result.FinalMessage().Text())
+	assert.Equal(t, "outer-agent", result.LastAgent.Name())
+}
+
+func TestAgentTool_Execute_LeafToolsRemainDetachedOnSuspend(t *testing.T) {
+	t.Parallel()
+
+	var leafCtxCanceled atomic.Bool
+
+	leafStarted := make(chan struct{})
+	leafRelease := make(chan struct{})
+
+	leafTool := agent.FunctionTool[struct{}](
+		"slow_leaf",
+		"Slow leaf tool",
+		func(ctx context.Context, _ struct{}) (agent.ToolResult, error) {
+			close(leafStarted)
+
+			select {
+			case <-ctx.Done():
+				leafCtxCanceled.Store(true)
+				return agent.ToolResult{Content: "leaf cancelled", IsError: true}, nil
+			case <-leafRelease:
+				if ctx.Err() != nil {
+					leafCtxCanceled.Store(true)
+				}
+				return agent.ToolResult{Content: "leaf completed"}, nil
+			}
+		},
+	)
+
+	provider := &mockProvider{
+		responses: []*llm.ChatCompletionResponse{
+			toolCallResponse(
+				llm.ToolCall{
+					ID:       "tc_leaf",
+					Function: llm.FunctionCall{Name: "slow_leaf", Arguments: `{}`},
+				},
+			),
+			stopResponse("done"),
+		},
+	}
+
+	ag := agent.New(
+		"leaf-agent",
+		newTestClient(provider),
+		agent.WithModel("test-model"),
+		agent.WithTools(leafTool),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	type runResult struct {
+		result *agent.Result
+		err    error
+	}
+	runDone := make(chan runResult, 1)
+
+	go func() {
+		result, err := ag.Run(ctx, []llm.Message{userMessage("go")})
+		runDone <- runResult{result: result, err: err}
+	}()
+
+	select {
+	case <-leafStarted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for leaf tool to start")
+	}
+
+	cancel()
+
+	select {
+	case outcome := <-runDone:
+		t.Fatalf(
+			"run returned before leaf tool release: err=%v result=%v",
+			outcome.err,
+			outcome.result,
+		)
+	case <-time.After(250 * time.Millisecond):
+	}
+
+	close(leafRelease)
+
+	select {
+	case outcome := <-runDone:
+		var se *agent.SuspendedError
+		require.ErrorAs(t, outcome.err, &se)
+		assert.Nil(t, outcome.result)
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for run completion after leaf release")
+	}
+
+	assert.False(
+		t,
+		leafCtxCanceled.Load(),
+		"leaf tool ctx should remain detached from suspend cancellation",
+	)
+}
+
 func TestAgentTool_InterfaceSatisfaction(t *testing.T) {
 	t.Parallel()
 
@@ -928,4 +1287,5 @@ func TestAgentTool_InterfaceSatisfaction(t *testing.T) {
 
 	assert.Implements(t, (*agent.Tool)(nil), tool)
 	assert.Implements(t, (*agent.ToolDescriptor)(nil), tool)
+	assert.Implements(t, (*agent.SuspendableTool)(nil), tool)
 }

pkg/agent/restore_test.go

@@ -16,6 +16,7 @@ package agent_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"testing"
@@ -74,6 +75,20 @@ func (r *simpleRegistry) Agent(name string) (*agent.Agent, error) {
 	return a, nil
 }
 
+type saveFailCheckpointer struct {
+	cp *agent.Checkpoint
+}
+
+func (s *saveFailCheckpointer) Save(_ context.Context, _ string, _ *agent.Checkpoint) error {
+	return errors.New("save exploded")
+}
+
+func (s *saveFailCheckpointer) Load(_ context.Context, _ string) (*agent.Checkpoint, error) {
+	clone := *s.cp
+
+	return &clone, nil
+}
+
 func TestRestore(t *testing.T) {
 	t.Parallel()
 
@@ -560,4 +575,225 @@ func TestRestore(t *testing.T) {
 			assert.Equal(t, "Completed after resume.", result.FinalMessage().Text())
 		},
 	)
+
+	t.Run(
+		"nested suspended restore keeps progress when inner agent missing",
+		func(t *testing.T) {
+			t.Parallel()
+
+			outerAgent := agent.New(
+				"outer-agent",
+				newTestClient(&mockProvider{}),
+				agent.WithModel("test-model"),
+			)
+
+			store := newMemoryCheckpointer()
+			err := store.Save(context.Background(), "run-nested-missing-inner", &agent.Checkpoint{
+				Status:    agent.AgentStatusSuspended,
+				AgentName: "outer-agent",
+				Messages: []llm.Message{
+					{
+						Role:  llm.RoleUser,
+						Parts: []llm.Part{llm.TextPart{Text: "continue"}},
+					},
+				},
+				AllToolCalls: []llm.ToolCall{
+					{
+						ID: "tc_missing",
+						Function: llm.FunctionCall{
+							Name:      "call_inner",
+							Arguments: `{"input":"go"}`,
+						},
+					},
+				},
+				InnerCheckpoints: map[string]*agent.Checkpoint{
+					"tc_missing": {
+						Status:    agent.AgentStatusSuspended,
+						AgentName: "inner-agent",
+					},
+				},
+			})
+			require.NoError(t, err)
+
+			registry := &simpleRegistry{
+				agents: map[string]*agent.Agent{
+					"outer-agent": outerAgent,
+				},
+			}
+
+			_, err = agent.Restore(
+				context.Background(),
+				store,
+				"run-nested-missing-inner",
+				registry,
+			)
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), `cannot resolve inner agent "inner-agent"`)
+
+			cp, loadErr := store.Load(context.Background(), "run-nested-missing-inner")
+			require.NoError(t, loadErr)
+			require.NotNil(t, cp)
+			require.Contains(t, cp.InnerCheckpoints, "tc_missing")
+		},
+	)
+
+	t.Run(
+		"nested suspended restore returns suspended when inner stays suspended",
+		func(t *testing.T) {
+			t.Parallel()
+
+			outerAgent := agent.New(
+				"outer-agent",
+				newTestClient(&mockProvider{}),
+				agent.WithModel("test-model"),
+			)
+			innerAgent := agent.New(
+				"inner-agent",
+				newTestClient(&mockProvider{}),
+				agent.WithModel("test-model"),
+			)
+
+			store := newMemoryCheckpointer()
+			err := store.Save(context.Background(), "run-nested-still-suspended", &agent.Checkpoint{
+				Status:    agent.AgentStatusSuspended,
+				AgentName: "outer-agent",
+				Messages: []llm.Message{
+					{
+						Role:  llm.RoleUser,
+						Parts: []llm.Part{llm.TextPart{Text: "continue"}},
+					},
+				},
+				AllToolCalls: []llm.ToolCall{
+					{
+						ID: "tc_inner",
+						Function: llm.FunctionCall{
+							Name:      "call_inner",
+							Arguments: `{"input":"go"}`,
+						},
+					},
+				},
+				InnerCheckpoints: map[string]*agent.Checkpoint{
+					"tc_inner": {
+						Status:    agent.AgentStatusSuspended,
+						AgentName: "inner-agent",
+					},
+				},
+			})
+			require.NoError(t, err)
+
+			registry := &simpleRegistry{
+				agents: map[string]*agent.Agent{
+					"outer-agent": outerAgent,
+					"inner-agent": innerAgent,
+				},
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+
+			_, err = agent.Restore(
+				ctx,
+				store,
+				"run-nested-still-suspended",
+				registry,
+			)
+
+			var se *agent.SuspendedError
+			require.ErrorAs(t, err, &se)
+			require.NotNil(t, se.Checkpoint)
+			require.Contains(t, se.Checkpoint.InnerCheckpoints, "tc_inner")
+		},
+	)
+
+	t.Run(
+		"nested suspended restore joins save failure with restore error",
+		func(t *testing.T) {
+			t.Parallel()
+
+			outerAgent := agent.New(
+				"outer-agent",
+				newTestClient(&mockProvider{}),
+				agent.WithModel("test-model"),
+			)
+
+			store := &saveFailCheckpointer{
+				cp: &agent.Checkpoint{
+					Status:    agent.AgentStatusSuspended,
+					AgentName: "outer-agent",
+					AllToolCalls: []llm.ToolCall{
+						{
+							ID: "tc_missing",
+							Function: llm.FunctionCall{
+								Name:      "call_inner",
+								Arguments: `{"input":"go"}`,
+							},
+						},
+					},
+					InnerCheckpoints: map[string]*agent.Checkpoint{
+						"tc_missing": {
+							Status:    agent.AgentStatusSuspended,
+							AgentName: "inner-agent",
+						},
+					},
+				},
+			}
+
+			registry := &simpleRegistry{
+				agents: map[string]*agent.Agent{
+					"outer-agent": outerAgent,
+				},
+			}
+
+			_, err := agent.Restore(
+				context.Background(),
+				store,
+				"run-nested-save-fail",
+				registry,
+			)
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), `cannot resolve inner agent "inner-agent"`)
+			assert.Contains(t, err.Error(), "cannot save nested restore progress")
+		},
+	)
+
+	t.Run(
+		"nested awaiting approval with unknown inner agent returns error",
+		func(t *testing.T) {
+			t.Parallel()
+
+			outerAgent := agent.New(
+				"outer-agent",
+				newTestClient(&mockProvider{}),
+				agent.WithModel("test-model"),
+			)
+
+			store := newMemoryCheckpointer()
+			err := store.Save(context.Background(), "run-awaiting-missing-inner", &agent.Checkpoint{
+				Status:    agent.AgentStatusAwaitingApproval,
+				AgentName: "outer-agent",
+				InnerCheckpoints: map[string]*agent.Checkpoint{
+					"tc_inner": {
+						Status:    agent.AgentStatusAwaitingApproval,
+						AgentName: "inner-agent",
+					},
+				},
+			})
+			require.NoError(t, err)
+
+			registry := &simpleRegistry{
+				agents: map[string]*agent.Agent{
+					"outer-agent": outerAgent,
+				},
+			}
+
+			_, err = agent.Restore(
+				context.Background(),
+				store,
+				"run-awaiting-missing-inner",
+				registry,
+			)
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), `cannot resolve inner agent "inner-agent"`)
+		},
+	)
 }

pkg/agent/run.go

@@ -75,6 +75,8 @@ type (
 		result ToolResult
 		err    error
 	}
+
+	suspendSignalKey struct{}
 )
 
 func WithCheckpointer(cp Checkpointer, runID string) RunOption {
@@ -86,6 +88,37 @@ func WithCheckpointer(cp Checkpointer, runID string) RunOption {
 
 func noopEvent(_ context.Context, _ StreamEvent) {}
 
+func withSuspendSignal(ctx context.Context, signal context.Context) context.Context {
+	return context.WithValue(ctx, suspendSignalKey{}, signal)
+}
+
+func suspendSignalFrom(ctx context.Context) context.Context {
+	signal, _ := ctx.Value(suspendSignalKey{}).(context.Context)
+
+	return signal
+}
+
+func withSuspendableToolContext(ctx context.Context, tool Tool) (context.Context, func()) {
+	if _, ok := tool.(SuspendableTool); !ok {
+		return ctx, func() {}
+	}
+
+	signal := suspendSignalFrom(ctx)
+	if signal == nil {
+		return ctx, func() {}
+	}
+
+	execCtx, cancel := context.WithCancelCause(ctx)
+	stop := context.AfterFunc(signal, func() {
+		cancel(context.Cause(signal))
+	})
+
+	return execCtx, func() {
+		stop()
+		cancel(nil)
+	}
+}
+
 func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
 	resp, err := agent.client.ChatCompletion(ctx, req)
 	if err == nil {
@@ -287,6 +320,7 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag
 	// hooks, save) carries through to completion once a checkpoint is
 	// requested.
 	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
+	ctx = withSuspendSignal(ctx, outerCtx)
 
 	s := &loopState{
 		agent:         startAgent,
@@ -1151,7 +1185,10 @@ func executeSingleTool(
 		log.String("tool", tool.Name()),
 	)
 
-	result, err := tool.Execute(toolCtx, tc.Function.Arguments)
+	execCtx, cleanupExecCtx := withSuspendableToolContext(toolCtx, tool)
+	defer cleanupExecCtx()
+
+	result, err := tool.Execute(execCtx, tc.Function.Arguments)
 	if err != nil {
 		if _, ok := errors.AsType[*InterruptedError](err); ok {
 			toolSpan.SetAttributes(attribute.Bool("tool.interrupted", true))
@@ -1306,6 +1343,7 @@ func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInpu
 
 func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
 	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
+	ctx = withSuspendSignal(ctx, outerCtx)
 
 	if interrupted.outerState != nil {
 		return resumeNested(outerCtx, interrupted, input, ro)
@@ -1470,6 +1508,7 @@ func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input Re
 
 func resumeNested(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
 	outerCtx, ctx := ctx, context.WithoutCancel(ctx)
+	ctx = withSuspendSignal(ctx, outerCtx)
 
 	outer := interrupted.outerState
 	logger := outer.agent.logger

pkg/agent/tool.go

@@ -39,6 +39,17 @@ type (
 		ToolDescriptor
 		Execute(ctx context.Context, arguments string) (ToolResult, error)
 	}
+
+	// SuspendableTool marks tools that can safely receive the run's
+	// graceful-suspend signal and checkpoint their own progress.
+	//
+	// Leaf tools should generally not implement this interface: they are
+	// expected to run on a detached context so in-flight side effects are
+	// not aborted during shutdown.
+	SuspendableTool interface {
+		Tool
+		Suspendable()
+	}
 )
 
 // ResultJSON marshals v to JSON and returns a successful ToolResult.

pkg/agentrun/helpers_test.go

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