From 0a1b47607b8f8fa4a99a4f40341168d93eb07229 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Sun, 7 Jun 2026 09:00:14 +0200 Subject: [PATCH] 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 --- pkg/agent/agent_tool.go | 4 + pkg/agent/agent_tool_test.go | 360 ++++++++++++ pkg/agent/restore_test.go | 236 ++++++++ pkg/agent/run.go | 41 +- pkg/agent/tool.go | 11 + pkg/agentrun/helpers_test.go | 452 +++++++++++++++ pkg/agentrun/worker_test.go | 1014 ++++++++++++++++++++++++++++++++++ 7 files changed, 2117 insertions(+), 1 deletion(-) create mode 100644 pkg/agentrun/helpers_test.go create mode 100644 pkg/agentrun/worker_test.go diff --git a/pkg/agent/agent_tool.go b/pkg/agent/agent_tool.go index 440228284..922f9b45c 100644 --- a/pkg/agent/agent_tool.go +++ b/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, diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go index dd4b4e57f..32bf6e632 100644 --- a/pkg/agent/agent_tool_test.go +++ b/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) } diff --git a/pkg/agent/restore_test.go b/pkg/agent/restore_test.go index 60d4dc9ad..6cf23be1a 100644 --- a/pkg/agent/restore_test.go +++ b/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"`) + }, + ) } diff --git a/pkg/agent/run.go b/pkg/agent/run.go index c5c09b0a7..2080afecc 100644 --- a/pkg/agent/run.go +++ b/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 diff --git a/pkg/agent/tool.go b/pkg/agent/tool.go index 890723a44..3d3be2c80 100644 --- a/pkg/agent/tool.go +++ b/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. diff --git a/pkg/agentrun/helpers_test.go b/pkg/agentrun/helpers_test.go new file mode 100644 index 000000000..c9168ba05 --- /dev/null +++ b/pkg/agentrun/helpers_test.go @@ -0,0 +1,452 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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) +} diff --git a/pkg/agentrun/worker_test.go b/pkg/agentrun/worker_test.go new file mode 100644 index 000000000..8a2256a45 --- /dev/null +++ b/pkg/agentrun/worker_test.go @@ -0,0 +1,1014 @@ +// Copyright (c) 2026 Probo Inc . +// +// 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 ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "os/signal" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/agentrun" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/llm" +) + +func TestWorker_PicksUpAndCompletes(t *testing.T) { + client := pgClient(t) + ag := newDummyAgent( + "echo-agent", + []*llm.ChatCompletionResponse{ + stopResponse("Done."), + }, + ) + + run := insertPendingRun( + t, + client, + "echo-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"echo-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.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) +} + +func TestWorker_StopAndResume(t *testing.T) { + client := pgClient(t) + store := coredata.NewPGCheckpointer(client) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + slowTool := agent.FunctionTool[struct{}]( + "slow_work", + "Does slow work", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + close(toolReady) + <-toolRelease + return agent.ToolResult{Content: "work done"}, nil + }, + ) + + ag := newDummyAgent( + "worker-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, + }), + stopResponse("All done after resume."), + }, + slowTool, + ) + + run := insertPendingRun( + t, + client, + "worker-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, + ) + + ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() + + select { + case <-toolReady: + case <-ctx1.Done(): + t.Fatal("timed out waiting for tool to start") + } + + cancel1() + + select { + case <-runWorker.ShutdownBroadcast(): + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for worker shutdown broadcast") + } + + close(toolRelease) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Checkpoint != nil + }, + 10*time.Second, + 200*time.Millisecond, + ) + + 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) + + runWorker2 := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, + ) + + ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() + + 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) +} + +// 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 +// resumed run continues in that branch and completes. +func TestWorker_StopAndResumeAcrossHandoff(t *testing.T) { + client := pgClient(t) + store := coredata.NewPGCheckpointer(client) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + slowTool := agent.FunctionTool[struct{}]( + "slow_work", + "Does slow work", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + close(toolReady) + <-toolRelease + return agent.ToolResult{Content: "child work done"}, nil + }, + ) + + childAgent := newDummyAgent( + "child-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_child", + Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, + }), + stopResponse("child done"), + }, + slowTool, + ) + + rootProvider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_root", + Function: llm.FunctionCall{Name: "transfer_to_child_agent", Arguments: `{}`}, + }), + }, + } + + rootAgent := agent.New( + "root-agent", + newTestClient(rootProvider), + agent.WithModel("test-model"), + agent.WithHandoffs(childAgent), + ) + + registry := &simpleRegistry{ + agents: map[string]*agent.Agent{ + "root-agent": rootAgent, + "child-agent": childAgent, + }, + } + + run := insertPendingRun( + t, + client, + "root-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, + ) + + runWorker := newTestWorker(client, registry) + + ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() + + select { + case <-toolReady: + case <-ctx1.Done(): + t.Fatal("timed out waiting for child agent tool to start") + } + + cancel1() + + select { + case <-runWorker.ShutdownBroadcast(): + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for worker shutdown broadcast") + } + + close(toolRelease) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Checkpoint != nil + }, + 10*time.Second, + 200*time.Millisecond, + ) + + cp, err := store.Load(context.Background(), run.ID.String()) + require.NoError(t, err) + require.NotNil(t, cp) + assert.Equal(t, agent.AgentStatusSuspended, cp.Status) + assert.Equal( + t, + "child-agent", + cp.AgentName, + "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) + defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() + + 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) +} + +func TestWorker_StopAndResumeNestedSubAgent(t *testing.T) { + client := pgClient(t) + store := coredata.NewPGCheckpointer(client) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + var readyOnce sync.Once + + 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: "inner work done"}, nil + }, + ) + + innerAgent := newDummyAgent( + "inner-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse( + llm.ToolCall{ + ID: "tc_inner", + Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, + }, + ), + stopResponse("inner done"), + }, + slowTool, + ) + + outerAgent := newDummyAgent( + "outer-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse( + llm.ToolCall{ + ID: "tc_outer", + Function: llm.FunctionCall{Name: "call_inner", Arguments: `{"input":"delegate"}`}, + }, + ), + stopResponse("outer done"), + }, + innerAgent.AsTool("call_inner", "Call inner"), + ) + + registry := &simpleRegistry{ + agents: map[string]*agent.Agent{ + "outer-agent": outerAgent, + "inner-agent": innerAgent, + }, + } + + run := insertPendingRun( + t, + client, + "outer-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, + ) + + runWorker := newTestWorker(client, registry) + + ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() + + select { + case <-toolReady: + case <-ctx1.Done(): + t.Fatal("timed out waiting for nested sub-agent tool to start") + } + + cancel1() + + select { + case <-runWorker.ShutdownBroadcast(): + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for worker shutdown broadcast") + } + + close(toolRelease) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Checkpoint != nil + }, + 10*time.Second, + 200*time.Millisecond, + ) + + cp, err := store.Load(context.Background(), run.ID.String()) + 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 for outer tool call") + require.NotNil(t, innerCP) + 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) + defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() + + 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) +} + +func TestWorker_StopAndResumeNestedSubAgentMultiLevel(t *testing.T) { + client := pgClient(t) + store := coredata.NewPGCheckpointer(client) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + var readyOnce sync.Once + + 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: "grandchild work done"}, nil + }, + ) + + grandchildAgent := newDummyAgent( + "grandchild-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse( + llm.ToolCall{ + ID: "tc_grandchild", + Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, + }, + ), + stopResponse("grandchild done"), + }, + slowTool, + ) + + childAgent := newDummyAgent( + "child-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse( + llm.ToolCall{ + ID: "tc_child", + Function: llm.FunctionCall{Name: "call_grandchild", Arguments: `{"input":"delegate deeper"}`}, + }, + ), + stopResponse("child done"), + }, + grandchildAgent.AsTool("call_grandchild", "Call grandchild"), + ) + + outerAgent := newDummyAgent( + "outer-agent", + []*llm.ChatCompletionResponse{ + toolCallResponse( + llm.ToolCall{ + ID: "tc_outer", + Function: llm.FunctionCall{Name: "call_child", Arguments: `{"input":"delegate"}`}, + }, + ), + stopResponse("outer done"), + }, + childAgent.AsTool("call_child", "Call child"), + ) + + registry := &simpleRegistry{ + agents: map[string]*agent.Agent{ + "outer-agent": outerAgent, + "child-agent": childAgent, + "grandchild-agent": grandchildAgent, + }, + } + + run := insertPendingRun( + t, + client, + "outer-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, + ) + + runWorker := newTestWorker(client, registry) + + ctx1, cancel1 := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel1() + go func() { _ = runWorker.Run(ctx1) }() + + select { + case <-toolReady: + case <-ctx1.Done(): + t.Fatal("timed out waiting for grandchild tool to start") + } + + cancel1() + + select { + case <-runWorker.ShutdownBroadcast(): + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for worker shutdown broadcast") + } + + close(toolRelease) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Checkpoint != nil + }, + 10*time.Second, + 200*time.Millisecond, + ) + + cp, err := store.Load(context.Background(), run.ID.String()) + 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) + require.NotNil(t, grandchildCP) + 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) + defer cancel2() + go func() { _ = runWorker2.Run(ctx2) }() + + 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) +} + +func TestWorker_HeartbeatLeaseLostLeavesRunForRecovery(t *testing.T) { + client := pgClient(t) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + var readyOnce sync.Once + + 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 := pgClient(t) + + toolReady := make(chan struct{}) + toolRelease := make(chan struct{}) + + slowTool := agent.FunctionTool[struct{}]( + "slow_work", + "Does slow work", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + close(toolReady) + <-toolRelease + return agent.ToolResult{Content: "work done"}, nil + }, + ) + + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "slow_work", Arguments: `{}`}, + }), + stopResponse("winner result"), + stopResponse("stale result"), + }, + } + + ag := agent.New( + "worker-agent", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(slowTool), + ) + + run := insertPendingRun( + t, + client, + "worker-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "do work"}}}}, + ) + + runWorkerA := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, + agentrun.WithWorkerLeaseDuration(5*time.Second), + agentrun.WithWorkerMaxConcurrency(1), + ) + + ctxA, cancelA := context.WithTimeout(context.Background(), 30*time.Second) + defer cancelA() + go func() { _ = runWorkerA.Run(ctxA) }() + + select { + case <-toolReady: + case <-ctxA.Done(): + t.Fatal("timed out waiting for first worker tool call") + } + + resetRunToPending(t, client, run.ID) + + runWorkerB := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"worker-agent": ag}}, + agentrun.WithWorkerLeaseDuration(5*time.Second), + agentrun.WithWorkerMaxConcurrency(1), + ) + + ctxB, cancelB := context.WithTimeout(context.Background(), 30*time.Second) + defer cancelB() + go func() { _ = runWorkerB.Run(ctxB) }() + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + return err == nil && r.Status == coredata.AgentRunStatusCompleted + }, + 15*time.Second, + 200*time.Millisecond, + ) + + winner := loadAgentRun(t, client, run.ID) + winnerResult := append(json.RawMessage(nil), winner.Result...) + require.NotNil(t, winnerResult) + + close(toolRelease) + + require.Eventually( + t, + func() bool { + provider.mu.Lock() + defer provider.mu.Unlock() + return provider.calls >= 3 + }, + 15*time.Second, + 200*time.Millisecond, + ) + + require.Eventually( + t, + func() bool { + r, err := tryLoadAgentRun(client, run.ID) + if err != nil { + return false + } + + return r.Status == coredata.AgentRunStatusCompleted && string(r.Result) == string(winnerResult) + }, + 10*time.Second, + 200*time.Millisecond, + ) +} + +func TestWorker_UnknownAgentFails(t *testing.T) { + client := pgClient(t) + + run := insertPendingRun( + t, + client, + "missing-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{}}, + ) + + 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.AgentRunStatusFailed + }, + 10*time.Second, + 200*time.Millisecond, + ) + + failed := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusFailed, failed.Status) + require.NotNil(t, failed.ErrorMessage) + assert.Contains(t, *failed.ErrorMessage, "cannot resolve agent") +} + +func TestWorker_InvalidInputMessagesFails(t *testing.T) { + client := pgClient(t) + ag := newDummyAgent( + "worker-agent", + []*llm.ChatCompletionResponse{ + stopResponse("Done."), + }, + ) + + run := insertPendingRun( + t, + client, + "worker-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "go"}}}}, + ) + + overwriteRunInputMessagesRaw(t, client, run.ID, `"invalid-json"`) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"worker-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.AgentRunStatusFailed + }, + 10*time.Second, + 200*time.Millisecond, + ) + + failed := loadAgentRun(t, client, run.ID) + assert.Equal(t, coredata.AgentRunStatusFailed, failed.Status) + require.NotNil(t, failed.ErrorMessage) + assert.Contains(t, *failed.ErrorMessage, "cannot unmarshal input messages") +} + +func TestWorker_SIGTERM(t *testing.T) { + if os.Getenv("TEST_SIGTERM_SUBPROCESS") == "1" { + runSIGTERMSubprocess(t) + return + } + + cmd := exec.Command(os.Args[0], "-test.run=^TestWorker_SIGTERM$") + cmd.Env = append(os.Environ(), "TEST_SIGTERM_SUBPROCESS=1") + + stdout, err := cmd.StdoutPipe() + require.NoError(t, err) + cmd.Stderr = cmd.Stdout + + require.NoError(t, cmd.Start()) + + ready := make(chan struct{}) + scanDone := make(chan struct{}) + var linesMu sync.Mutex + var lines []string + snapshotLines := func() string { + linesMu.Lock() + defer linesMu.Unlock() + return strings.Join(lines, "\n") + } + go func() { + defer close(scanDone) + + scanner := bufio.NewScanner(stdout) + for scanner.Scan() { + line := scanner.Text() + linesMu.Lock() + lines = append(lines, line) + linesMu.Unlock() + if line == "READY" { + close(ready) + } + } + }() + + select { + case <-ready: + case <-time.After(20 * time.Second): + _ = cmd.Process.Kill() + t.Fatalf("subprocess did not become ready for SIGTERM\n%s", snapshotLines()) + } + + require.NoError(t, cmd.Process.Signal(syscall.SIGTERM)) + if err := cmd.Wait(); err != nil { + t.Fatalf("subprocess failed: %v\n%s", err, snapshotLines()) + } + <-scanDone +} + +func runSIGTERMSubprocess(t *testing.T) { + client := pgClient(t) + + workStarted := make(chan struct{}) + + ag := newDummyAgent( + "battle-agent", + battleTestResponses(), + makeBattleTools(workStarted)..., + ) + + run := insertPendingRun( + t, + client, + "battle-agent", + []llm.Message{{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: "start"}}}}, + ) + + runWorker := newTestWorker( + client, + &simpleRegistry{agents: map[string]*agent.Agent{"battle-agent": ag}}, + agentrun.WithWorkerInterval(150*time.Millisecond), + agentrun.WithWorkerMaxConcurrency(1), + ) + + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM) + defer stop() + + go func() { + _ = runWorker.Run(ctx) + }() + + select { + case <-workStarted: + case <-time.After(15 * time.Second): + t.Fatal("tool did not start before SIGTERM") + } + + fmt.Fprintln(os.Stdout, "READY") + + select { + case <-runWorker.ShutdownBroadcast(): + case <-time.After(15 * time.Second): + t.Fatal("worker did not broadcast shutdown after SIGTERM") + } + + time.Sleep(300 * time.Millisecond) + + // The in-flight run may checkpoint or be recovered later depending on + // timing, but it must still be queryable after graceful shutdown. + _, err := tryLoadAgentRun(client, run.ID) + require.NoError(t, err) +} + +type workInput struct { + Step string `json:"step"` +} + +func makeBattleTools(workStarted chan<- struct{}) []agent.Tool { + return []agent.Tool{ + agent.FunctionTool[workInput]( + "do_work", + "Performs interruptible work", + func(ctx context.Context, _ workInput) (agent.ToolResult, error) { + close(workStarted) + <-ctx.Done() + return agent.ToolResult{Content: "interrupted"}, ctx.Err() + }, + ), + } +} + +func battleTestResponses() []*llm.ChatCompletionResponse { + return []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_battle_1", + Function: llm.FunctionCall{ + Name: "do_work", + Arguments: `{"step":"one"}`, + }, + }), + stopResponse("done"), + } +}