diff --git a/adk/agent_tool.go b/adk/agent_tool.go index c660feea3..d3af71976 100644 --- a/adk/agent_tool.go +++ b/adk/agent_tool.go @@ -21,11 +21,13 @@ import ( "context" "errors" "fmt" + "time" "github.com/bytedance/sonic" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/internal/core" "github.com/cloudwego/eino/schema" ) @@ -42,10 +44,28 @@ var ( type AgentToolOptions struct { fullChatHistoryAsInput bool agentInputSchema *schema.ParamsOneOf + retryConfig *AgentToolRetryConfig } type AgentToolOption func(*AgentToolOptions) +// AgentToolRetryConfig configures retries for an agent wrapped as a tool. +// +// MaxRetries is the number of additional attempts after the initial attempt. +// Retries start a fresh child-agent run and therefore should only be enabled +// for agents whose side effects are safe to repeat. Forwarded internal events +// from a failed attempt remain observable. +type AgentToolRetryConfig struct { + MaxRetries int + // IsRetryable decides whether an invocation error should be retried. + // If nil, all errors except cancellation, deadline, and interrupt signals + // are retryable. + IsRetryable func(ctx context.Context, err error) bool + // Backoff returns the delay before a retry. attempt is one-based. + // If nil, retries start immediately. + Backoff func(ctx context.Context, attempt int) time.Duration +} + // WithFullChatHistoryAsInput enables using the full chat history as input. func WithFullChatHistoryAsInput() AgentToolOption { return func(options *AgentToolOptions) { @@ -60,6 +80,21 @@ func WithAgentInputSchema(schema *schema.ParamsOneOf) AgentToolOption { } } +// WithAgentToolRetry enables whole-invocation retries for an agent tool. +// +// A nil config disables retries. Interrupt/resume invocations are never +// retried because they must preserve their existing checkpoint state. +func WithAgentToolRetry(config *AgentToolRetryConfig) AgentToolOption { + return func(options *AgentToolOptions) { + if config == nil { + options.retryConfig = nil + return + } + cloned := *config + options.retryConfig = &cloned + } +} + func withAgentToolEnableStreaming(enabled bool) tool.Option { return tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) { opt.enableStreaming = enabled @@ -100,6 +135,7 @@ func NewAgentTool(_ context.Context, agent Agent, options ...AgentToolOption) to agent: agent, fullChatHistoryAsInput: opts.fullChatHistoryAsInput, inputSchema: opts.agentInputSchema, + retryConfig: opts.retryConfig, } } @@ -114,6 +150,7 @@ func NewTypedAgentTool[M MessageType](_ context.Context, agent TypedAgent[M], op agent: agent, fullChatHistoryAsInput: opts.fullChatHistoryAsInput, inputSchema: opts.agentInputSchema, + retryConfig: opts.retryConfig, } } @@ -122,6 +159,7 @@ type typedAgentTool[M MessageType] struct { fullChatHistoryAsInput bool inputSchema *schema.ParamsOneOf + retryConfig *AgentToolRetryConfig } type agentTool = typedAgentTool[*schema.Message] @@ -152,6 +190,36 @@ func (at *typedAgentTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) } func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + if at.retryConfig == nil { + return at.invokeOnce(ctx, argumentsInJSON, opts...) + } + if at.retryConfig.MaxRetries < 0 { + return "", errors.New("agent tool retry MaxRetries must be non-negative") + } + + wasInterrupted, _, _ := tool.GetInterruptState[[]byte](ctx) + if wasInterrupted { + return at.invokeOnce(ctx, argumentsInJSON, opts...) + } + + for attempt := 0; ; attempt++ { + result, err := at.invokeOnce(ctx, argumentsInJSON, opts...) + if err == nil { + return result, nil + } + if ctxErr := ctx.Err(); ctxErr != nil { + return "", ctxErr + } + if attempt >= at.retryConfig.MaxRetries || !at.shouldRetry(ctx, err) { + return "", err + } + if err = waitAgentToolRetry(ctx, at.retryConfig, attempt+1); err != nil { + return "", err + } + } +} + +func (at *typedAgentTool[M]) invokeOnce(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { if cancelCtx := getCancelContext(ctx); cancelCtx != nil { cancelCtx.markCheckpointAwareDescendant() } @@ -279,6 +347,45 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s return ret, nil } +func (at *typedAgentTool[M]) shouldRetry(ctx context.Context, err error) bool { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + + var cancelErr *CancelError + if errors.As(err, &cancelErr) { + return false + } + var interruptSignal *core.InterruptSignal + if errors.As(err, &interruptSignal) { + return false + } + var interruptProvider core.InterruptContextsProvider + if errors.As(err, &interruptProvider) { + return false + } + return at.retryConfig.IsRetryable == nil || at.retryConfig.IsRetryable(ctx, err) +} + +func waitAgentToolRetry(ctx context.Context, config *AgentToolRetryConfig, attempt int) error { + if config.Backoff == nil { + return nil + } + delay := config.Backoff(ctx, attempt) + if delay <= 0 { + return nil + } + + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + // agentToolOptions is a wrapper structure used to convert AgentRunOption slices to tool.Option. // It stores the agent name and corresponding run options for tool-specific processing. type agentToolOptions struct { diff --git a/adk/agent_tool_test.go b/adk/agent_tool_test.go index 785ad995a..cd304a9ad 100644 --- a/adk/agent_tool_test.go +++ b/adk/agent_tool_test.go @@ -18,11 +18,13 @@ package adk import ( "context" + "errors" "fmt" "strings" "sync" "sync/atomic" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -763,6 +765,130 @@ func TestNestedAgentTool_RunPath(t *testing.T) { } } +func TestAgentToolRetry(t *testing.T) { + t.Run("retries retryable error", func(t *testing.T) { + retryErr := errors.New("retry") + agent := &retryAgentToolAgent{ + failures: 2, + err: retryErr, + } + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: 2, + })).(tool.InvokableTool) + + result, err := agentTool.InvokableRun(context.Background(), `{"request":"hello"}`) + require.NoError(t, err) + require.Equal(t, "success", result) + require.Equal(t, int32(3), agent.attempts.Load()) + }) + + t.Run("predicate rejects retry", func(t *testing.T) { + terminalErr := errors.New("terminal") + agent := &retryAgentToolAgent{failures: 2, err: terminalErr} + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: 3, + IsRetryable: func(_ context.Context, err error) bool { + return !errors.Is(err, terminalErr) + }, + })).(tool.InvokableTool) + + _, err := agentTool.InvokableRun(context.Background(), `{"request":"hello"}`) + require.ErrorIs(t, err, terminalErr) + require.Equal(t, int32(1), agent.attempts.Load()) + }) + + t.Run("returns last error when exhausted", func(t *testing.T) { + retryErr := errors.New("retry") + agent := &retryAgentToolAgent{failures: 3, err: retryErr} + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: 1, + })).(tool.InvokableTool) + + _, err := agentTool.InvokableRun(context.Background(), `{"request":"hello"}`) + require.ErrorIs(t, err, retryErr) + require.Equal(t, int32(2), agent.attempts.Load()) + }) + + t.Run("backoff observes context cancellation", func(t *testing.T) { + retryErr := errors.New("retry") + agent := &retryAgentToolAgent{failures: 2, err: retryErr} + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: 2, + Backoff: func(_ context.Context, _ int) time.Duration { + return time.Minute + }, + })).(tool.InvokableTool) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := agentTool.InvokableRun(ctx, `{"request":"hello"}`) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, int32(1), agent.attempts.Load()) + }) + + t.Run("interrupt is not retried", func(t *testing.T) { + agent := &retryAgentToolAgent{interrupt: true} + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: 2, + })).(tool.InvokableTool) + + _, err := agentTool.InvokableRun(context.Background(), `{"request":"hello"}`) + require.Error(t, err) + require.Equal(t, int32(1), agent.attempts.Load()) + }) + + t.Run("rejects negative max retries", func(t *testing.T) { + agent := &retryAgentToolAgent{} + agentTool := NewAgentTool(context.Background(), agent, WithAgentToolRetry(&AgentToolRetryConfig{ + MaxRetries: -1, + })).(tool.InvokableTool) + + _, err := agentTool.InvokableRun(context.Background(), `{"request":"hello"}`) + require.ErrorContains(t, err, "must be non-negative") + require.Zero(t, agent.attempts.Load()) + }) +} + +type retryAgentToolAgent struct { + attempts atomic.Int32 + failures int32 + err error + interrupt bool +} + +func (a *retryAgentToolAgent) Name(context.Context) string { + return "retry_agent" +} + +func (a *retryAgentToolAgent) Description(context.Context) string { + return "retry agent" +} + +func (a *retryAgentToolAgent) Run(ctx context.Context, _ *AgentInput, + _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + attempt := a.attempts.Add(1) + go func() { + defer gen.Close() + if a.interrupt { + gen.Send(&AgentEvent{Err: tool.Interrupt(ctx, "pause")}) + return + } + if attempt <= a.failures { + gen.Send(&AgentEvent{Err: a.err}) + return + } + gen.Send(&AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("success", nil), + }, + }, + }) + }() + return iter +} + func TestNestedAgentTool_NoInternalEventsWhenDisabled(t *testing.T) { ctx := context.Background()