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:
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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"`)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user