refactor(adk): introduce AgentHandler interface for runtime agent customization

adk/chatmodel_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 
+	"github.com/cloudwego/eino/components/model"
 	"github.com/cloudwego/eino/components/tool"
 	"github.com/cloudwego/eino/compose"
 	mockModel "github.com/cloudwego/eino/internal/mock/components/model"
@@ -456,8 +457,8 @@ func TestParallelReturnDirectlyToolCall(t *testing.T) {
 					&myTool{name: "tool3", desc: "tool3", waitTime: 100 * time.Millisecond},
 				},
 			},
-			ReturnDirectly: map[string]bool{
-				"tool1": true,
+			ReturnDirectly: map[string]struct{}{
+				"tool1": {},
 			},
 		},
 	})
@@ -1183,3 +1184,116 @@ func (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string,
 func (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) {
 	return schema.StreamReaderFromArray([]string{s.result}), nil
 }
+
+func TestGetComposeOptions(t *testing.T) {
+	t.Run("WithChatModelOptions", func(t *testing.T) {
+		ctx := context.Background()
+		ctrl := gomock.NewController(t)
+		cm := mockModel.NewMockToolCallingChatModel(ctrl)
+
+		var capturedTemperature float32
+		cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+				options := model.GetCommonOptions(&model.Options{}, opts...)
+				if options.Temperature != nil {
+					capturedTemperature = *options.Temperature
+				}
+				return schema.AssistantMessage("response", nil), nil
+			}).Times(1)
+
+		agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
+			Name:        "TestAgent",
+			Description: "Test agent",
+			Model:       cm,
+		})
+		assert.NoError(t, err)
+
+		temp := float32(0.7)
+		iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}},
+			WithChatModelOptions([]model.Option{model.WithTemperature(temp)}))
+		for {
+			_, ok := iter.Next()
+			if !ok {
+				break
+			}
+		}
+
+		assert.Equal(t, temp, capturedTemperature, "Temperature should be passed through WithChatModelOptions")
+	})
+
+	t.Run("WithToolOptions", func(t *testing.T) {
+		ctx := context.Background()
+		ctrl := gomock.NewController(t)
+		cm := mockModel.NewMockToolCallingChatModel(ctrl)
+
+		var toolOptionsCaptured bool
+		testTool := &toolOptionCapturingTool{
+			name: "test_tool",
+			onRun: func(opts []tool.Option) {
+				if len(opts) > 0 {
+					toolOptionsCaptured = true
+				}
+			},
+		}
+		info, _ := testTool.Info(ctx)
+
+		cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes()
+		cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			Return(schema.AssistantMessage("Using tool", []schema.ToolCall{
+				{ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}},
+			}), nil).Times(1)
+		cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
+			Return(schema.AssistantMessage("done", nil), nil).Times(1)
+
+		agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
+			Name:        "TestAgent",
+			Description: "Test agent",
+			Model:       cm,
+			ToolsConfig: ToolsConfig{
+				ToolsNodeConfig: compose.ToolsNodeConfig{
+					Tools: []tool.BaseTool{testTool},
+				},
+			},
+		})
+		assert.NoError(t, err)
+
+		iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}},
+			WithToolOptions([]tool.Option{testToolOption("test_value")}))
+		for {
+			_, ok := iter.Next()
+			if !ok {
+				break
+			}
+		}
+
+		assert.True(t, toolOptionsCaptured, "Tool options should be passed through WithToolOptions")
+	})
+
+
+}
+
+type toolOptionCapturingTool struct {
+	name  string
+	onRun func(opts []tool.Option)
+}
+
+func (t *toolOptionCapturingTool) Info(_ context.Context) (*schema.ToolInfo, error) {
+	return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil
+}
+
+func (t *toolOptionCapturingTool) InvokableRun(_ context.Context, _ string, opts ...tool.Option) (string, error) {
+	if t.onRun != nil {
+		t.onRun(opts)
+	}
+	return t.name + " result", nil
+}
+
+type testToolOptions struct {
+	value string
+}
+
+func testToolOption(value string) tool.Option {
+	return tool.WrapImplSpecificOptFn(func(o *testToolOptions) {
+		o.value = value
+	})
+}

adk/handler.go

@@ -0,0 +1,139 @@
+/*
+ * Copyright 2025 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package adk
+
+import (
+	"context"
+	"reflect"
+
+	"github.com/cloudwego/eino/components/model"
+	"github.com/cloudwego/eino/components/tool"
+)
+
+// AgentContext contains runtime information passed to handlers before each agent run.
+// Handlers can modify Instruction, Tools, and ReturnDirectly to customize agent behavior.
+type AgentContext struct {
+	Instruction    string
+	Tools          []tool.BaseTool
+	ReturnDirectly map[string]struct{}
+}
+
+// HandlerMiddleware defines the interface for customizing agent behavior.
+//
+// Why HandlerMiddleware instead of AgentMiddleware?
+//
+// AgentMiddleware is a struct type, which has inherent limitations:
+//   - Struct types are closed: users cannot add new methods to extend functionality
+//   - The framework only recognizes AgentMiddleware's fixed fields, so even if users
+//     embed AgentMiddleware in a custom struct and add methods, the framework cannot
+//     call those methods (config.Middlewares is []AgentMiddleware, not a user type)
+//   - Callbacks in AgentMiddleware only return error, cannot return modified context
+//
+// HandlerMiddleware is an interface type, which is open for extension:
+//   - Users can implement custom handlers with arbitrary internal state and methods
+//   - All methods return (context.Context, ..., error), allowing context propagation
+//   - Configuration is centralized in struct fields rather than scattered in closures
+//
+// HandlerMiddleware vs AgentMiddleware:
+//   - Use AgentMiddleware for simple, static additions (extra instruction/tools)
+//   - Use HandlerMiddleware for dynamic behavior, context modification, or call wrapping
+//   - AgentMiddleware is kept for backward compatibility with existing users
+//   - Both can be used together; middlewares are applied first, then handlers
+//
+// Use *BaseHandlerMiddleware as an embedded struct to provide default no-op
+// implementations for all methods.
+type HandlerMiddleware interface {
+	// BeforeAgent is called before each agent run, allowing modification of
+	// the agent's instruction and tools configuration.
+	BeforeAgent(ctx context.Context, runCtx *AgentContext) (context.Context, *AgentContext, error)
+
+	// BeforeModelRewriteHistory is called before each model invocation.
+	// The returned messages are persisted to the agent's internal state and passed to the model.
+	// The returned context is propagated to the model call and subsequent handlers.
+	BeforeModelRewriteHistory(ctx context.Context, messages []Message) (context.Context, []Message, error)
+
+	// AfterModelRewriteHistory is called after each model invocation.
+	// The input messages include the model's response as the last message.
+	// The returned messages are persisted to the agent's internal state.
+	AfterModelRewriteHistory(ctx context.Context, messages []Message) (context.Context, []Message, error)
+
+	// WrapTool wraps a tool with custom behavior.
+	// Return the input tool unchanged if no wrapping is needed.
+	// Called at construction time (or after BeforeAgent if tools are modified dynamically).
+	WrapTool(ctx context.Context, t tool.BaseTool) (tool.BaseTool, error)
+
+	// WrapModel wraps a chat model with custom behavior.
+	// Return the input model unchanged if no wrapping is needed.
+	// Called once when the agent is built, not per-call.
+	WrapModel(ctx context.Context, m model.BaseChatModel) (model.BaseChatModel, error)
+}
+
+// BaseHandlerMiddleware provides default no-op implementations for HandlerMiddleware.
+// Embed *BaseHandlerMiddleware in custom handlers to only override the methods you need.
+type BaseHandlerMiddleware struct{}
+
+func (b *BaseHandlerMiddleware) WrapTool(_ context.Context, t tool.BaseTool) (tool.BaseTool, error) {
+	return t, nil
+}
+
+func (b *BaseHandlerMiddleware) WrapModel(_ context.Context, m model.BaseChatModel) (model.BaseChatModel, error) {
+	return m, nil
+}
+
+func (b *BaseHandlerMiddleware) BeforeAgent(ctx context.Context, runCtx *AgentContext) (context.Context, *AgentContext, error) {
+	return ctx, runCtx, nil
+}
+
+func (b *BaseHandlerMiddleware) BeforeModelRewriteHistory(ctx context.Context, messages []Message) (context.Context, []Message, error) {
+	return ctx, messages, nil
+}
+
+func (b *BaseHandlerMiddleware) AfterModelRewriteHistory(ctx context.Context, messages []Message) (context.Context, []Message, error) {
+	return ctx, messages, nil
+}
+
+type handlerInfo struct {
+	handler                      HandlerMiddleware
+	hasBeforeAgent               bool
+	hasBeforeModelRewriteHistory bool
+	hasAfterModelRewriteHistory  bool
+	hasWrapTool                  bool
+	hasWrapModel                 bool
+}
+
+var baseHandlerMiddlewareType = reflect.TypeOf(&BaseHandlerMiddleware{})
+
+func isMethodOverridden(handler HandlerMiddleware, methodName string) bool {
+	handlerType := reflect.TypeOf(handler)
+	handlerMethod, ok1 := handlerType.MethodByName(methodName)
+	baseMethod, ok2 := baseHandlerMiddlewareType.MethodByName(methodName)
+	if !ok1 || !ok2 {
+		return true
+	}
+	return handlerMethod.Func.Pointer() != baseMethod.Func.Pointer()
+}
+
+func newHandlerInfo(h HandlerMiddleware) handlerInfo {
+	return handlerInfo{
+		handler:                      h,
+		hasBeforeAgent:               isMethodOverridden(h, "BeforeAgent"),
+		hasBeforeModelRewriteHistory: isMethodOverridden(h, "BeforeModelRewriteHistory"),
+		hasAfterModelRewriteHistory:  isMethodOverridden(h, "AfterModelRewriteHistory"),
+		hasWrapTool:                  isMethodOverridden(h, "WrapTool"),
+		hasWrapModel:                 isMethodOverridden(h, "WrapModel"),
+	}
+}