feat(adk): add retries for agent tools

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 {

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()