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
2026-06-07 09:00:14 +02:00
parent 3dfc833671
commit 0a1b47607b
7 changed files with 2117 additions and 1 deletions

View File

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

View File

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

View File

@@ -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"`)
},
)
}

View File

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

View File

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