rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make agent-tool subtrees suspendable

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

236
pkg/agent/restore_test.go
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"`)
},
)
}