feat(adk): add BeforeFinalAnswer hook to ChatModelAgentMiddleware

.gitignore

@@ -47,10 +47,14 @@ output/*
 
 # Reports (generated analysis files)
 reports/
+/todos
 
 .DS_Store
-*.log
+*.log*
+.claude
 CLAUDE.md
+*.jsonl
+*.txt
 
 # Specs directories
 */specs

adk/agent_tool.go

@@ -103,14 +103,34 @@ func NewAgentTool(_ context.Context, agent Agent, options ...AgentToolOption) to
 	}
 }
 
-type agentTool struct {
-	agent Agent
+// NewTypedAgentTool creates a new agent tool that wraps a TypedAgent as a tool.BaseTool.
+func NewTypedAgentTool[M messageType](_ context.Context, agent TypedAgent[M], options ...AgentToolOption) tool.BaseTool {
+	opts := &AgentToolOptions{}
+	for _, opt := range options {
+		opt(opts)
+	}
+
+	return &typedAgentTool[M]{
+		agent:                  agent,
+		fullChatHistoryAsInput: opts.fullChatHistoryAsInput,
+		inputSchema:            opts.agentInputSchema,
+	}
+}
+
+type typedAgentTool[M messageType] struct {
+	agent TypedAgent[M]
 
 	fullChatHistoryAsInput bool
 	inputSchema            *schema.ParamsOneOf
 }
 
-func (at *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
+type agentTool = typedAgentTool[*schema.Message]
+
+type agentToolRequest struct {
+	Request string `json:"request"`
+}
+
+func (at *typedAgentTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) {
 	name := at.agent.Name(ctx)
 	if name == "" {
 		return nil, errors.New("agent tool requires a non-empty Name")
@@ -119,7 +139,6 @@ func (at *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
 	if desc == "" {
 		return nil, errors.New("agent tool requires a non-empty Description")
 	}
-
 	param := at.inputSchema
 	if param == nil {
 		param = defaultAgentToolParam
@@ -132,57 +151,60 @@ func (at *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
 	}, nil
 }
 
-func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
+func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
 	gen, enableStreaming := getEmitGeneratorAndEnableStreaming(opts)
 	var ms *bridgeStore
-	var iter *AsyncIterator[*AgentEvent]
+	var iter *AsyncIterator[*TypedAgentEvent[M]]
 	var err error
 
 	wasInterrupted, hasState, state := tool.GetInterruptState[[]byte](ctx)
 	if !wasInterrupted {
 		ms = newBridgeStore()
-		var input []Message
+
+		var input []M
 		if at.fullChatHistoryAsInput {
-			input, err = getReactChatHistory(ctx, at.agent.Name(ctx))
-			if err != nil {
-				return "", err
+			var zero M
+			if _, ok := any(zero).(*schema.Message); !ok {
+				return "", fmt.Errorf("fullChatHistoryAsInput is only supported for *schema.Message agents")
 			}
+			msgInput, histErr := getReactChatHistory(ctx, at.agent.Name(ctx))
+			if histErr != nil {
+				return "", histErr
+			}
+			input = any(msgInput).([]M)
 		} else {
 			if at.inputSchema == nil {
-				// default input schema
-				type request struct {
-					Request string `json:"request"`
-				}
-
-				req := &request{}
+				req := &agentToolRequest{}
 				err = sonic.UnmarshalString(argumentsInJSON, req)
 				if err != nil {
 					return "", err
 				}
 				argumentsInJSON = req.Request
 			}
-			input = []Message{
-				schema.UserMessage(argumentsInJSON),
-			}
+			input = newTypedUserMessages[M](argumentsInJSON)
 		}
 
-		iter = newInvokableAgentToolRunner(at.agent, ms, enableStreaming).Run(ctx, input,
-			append(getOptionsByAgentName(at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...)
+		runner := newTypedInvokableAgentToolRunner[M](at.agent, ms, enableStreaming)
+		iter = runner.Run(ctx, input,
+			append(extractAndDeriveCancelCtx(ctx, at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...)
 	} else {
 		if !hasState {
 			return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx))
 		}
 
-		ms = newResumeBridgeStore(state)
+		ms = newResumeBridgeStore(bridgeCheckpointID, state)
 
-		iter, err = newInvokableAgentToolRunner(at.agent, ms, enableStreaming).
-			Resume(ctx, bridgeCheckpointID, append(getOptionsByAgentName(at.agent.Name(ctx), opts), withSharedParentSession())...)
+		agentOpts := extractAndDeriveCancelCtx(ctx, at.agent.Name(ctx), opts)
+		agentOpts = append(agentOpts, withSharedParentSession())
+
+		runner := newTypedInvokableAgentToolRunner[M](at.agent, ms, enableStreaming)
+		iter, err = runner.Resume(ctx, bridgeCheckpointID, agentOpts...)
 		if err != nil {
 			return "", err
 		}
 	}
 
-	var lastEvent *AgentEvent
+	var lastEvent *TypedAgentEvent[M]
 	for {
 		event, ok := iter.Next()
 		if !ok {
@@ -208,9 +230,13 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 					rp = append(rp, event.RunPath...)
 					event.RunPath = rp
 				}
-				tmp := copyAgentEvent(event)
-				gen.Send(event)
-				event = tmp
+				if msgEvent, ok := any(event).(*AgentEvent); ok {
+					tmp := copyTypedAgentEvent(msgEvent)
+					gen.Send(msgEvent)
+					event = any(tmp).(*TypedAgentEvent[M])
+				} else {
+					return "", fmt.Errorf("cross-message-type agent tools are not supported: cannot use an AgenticMessage agent as a tool of a Message agent")
+				}
 			}
 		}
 
@@ -241,7 +267,7 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 			if err != nil {
 				return "", err
 			}
-			ret = msg.Content
+			ret = extractTextContent(msg)
 		}
 	}
 
@@ -281,6 +307,18 @@ func getOptionsByAgentName(agentName string, opts []tool.Option) []AgentRunOptio
 	return ret
 }
 
+func extractAndDeriveCancelCtx(ctx context.Context, agentName string, opts []tool.Option) []AgentRunOption {
+	agentOpts := getOptionsByAgentName(agentName, opts)
+	baseOpts := getCommonOptions(nil, agentOpts...)
+	if baseOpts.cancelCtx != nil {
+		childCtx := baseOpts.cancelCtx.deriveChild(ctx)
+		agentOpts = append(agentOpts, WrapImplSpecificOptFn(func(o *options) {
+			o.cancelCtx = childCtx
+		}))
+	}
+	return agentOpts
+}
+
 func getEmitGeneratorAndEnableStreaming(opts []tool.Option) (*AsyncGenerator[*AgentEvent], bool) {
 	o := tool.GetImplSpecificOptions[agentToolOptions](nil, opts...)
 	if o == nil {
@@ -293,8 +331,11 @@ func getEmitGeneratorAndEnableStreaming(opts []tool.Option) (*AsyncGenerator[*Ag
 func getReactChatHistory(ctx context.Context, destAgentName string) ([]Message, error) {
 	var messages []Message
 	err := compose.ProcessState(ctx, func(ctx context.Context, st *State) error {
+		if len(st.Messages) == 0 {
+			return nil
+		}
 		messages = make([]Message, len(st.Messages)-1)
-		copy(messages, st.Messages[:len(st.Messages)-1]) // remove the last assistant message, which is the tool call message
+		copy(messages, st.Messages[:len(st.Messages)-1])
 		return nil
 	})
 	if err != nil {
@@ -324,8 +365,20 @@ func getReactChatHistory(ctx context.Context, destAgentName string) ([]Message,
 	return history, nil
 }
 
-func newInvokableAgentToolRunner(agent Agent, store compose.CheckPointStore, enableStreaming bool) *Runner {
-	return &Runner{
+func newTypedUserMessages[M messageType](text string) []M {
+	var zero M
+	switch any(zero).(type) {
+	case *schema.Message:
+		return any([]Message{schema.UserMessage(text)}).([]M)
+	case *schema.AgenticMessage:
+		return any([]*schema.AgenticMessage{schema.UserAgenticMessage(text)}).([]M)
+	default:
+		return nil
+	}
+}
+
+func newTypedInvokableAgentToolRunner[M messageType](agent TypedAgent[M], store compose.CheckPointStore, enableStreaming bool) *TypedRunner[M] {
+	return &TypedRunner[M]{
 		a:               agent,
 		enableStreaming: enableStreaming,
 		store:           store,

adk/agent_tool_test.go

@@ -21,16 +21,36 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/cloudwego/eino/components/model"
 	"github.com/cloudwego/eino/components/tool"
 	"github.com/cloudwego/eino/compose"
 	"github.com/cloudwego/eino/schema"
 )
 
+type mockChatModelForAttack struct {
+	generateFn func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)
+}
+
+func (m *mockChatModelForAttack) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+	return m.generateFn(ctx, input, opts...)
+}
+
+func (m *mockChatModelForAttack) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+	result, err := m.generateFn(ctx, input, opts...)
+	if err != nil {
+		return nil, err
+	}
+	r, w := schema.Pipe[*schema.Message](1)
+	go func() { defer w.Close(); w.Send(result, nil) }()
+	return r, nil
+}
+
 // mockAgent implements the Agent interface for testing
 type mockAgentForTool struct {
 	name        string
@@ -1146,3 +1166,76 @@ func TestInvokableAgentTool_ErrorCases(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "", out2)
 }
+
+func TestCrossTypeAgentToolGracefulError(t *testing.T) {
+	ctx := context.Background()
+
+	innerModel := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg("inner result"), nil
+		},
+	}
+
+	innerAgent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "AgenticInner",
+		Description: "An agentic agent used as a tool",
+		Model:       innerModel,
+	})
+	require.NoError(t, err)
+
+	agenticAgentTool := NewTypedAgentTool(ctx, TypedAgent[*schema.AgenticMessage](innerAgent))
+
+	var outerCallCount int32
+	outerModel := &mockChatModelForAttack{
+		generateFn: func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+			count := atomic.AddInt32(&outerCallCount, 1)
+			if count == 1 {
+				return &schema.Message{
+					Role: schema.Assistant,
+					ToolCalls: []schema.ToolCall{
+						{ID: "c1", Function: schema.FunctionCall{Name: "AgenticInner", Arguments: `{"request":"test"}`}},
+					},
+				}, nil
+			}
+			return schema.AssistantMessage("done", nil), nil
+		},
+	}
+
+	outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{
+		Name:        "OuterMessageAgent",
+		Description: "A Message agent using an AgenticMessage sub-agent tool",
+		Model:       outerModel,
+		ToolsConfig: ToolsConfig{
+			ToolsNodeConfig: compose.ToolsNodeConfig{
+				Tools: []tool.BaseTool{agenticAgentTool},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	runner := NewRunner(ctx, RunnerConfig{Agent: outerAgent, EnableStreaming: true})
+	iter := runner.Query(ctx, "test cross-type")
+
+	var capturedErr error
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if event.Err != nil {
+			capturedErr = event.Err
+			t.Logf("Cross-type error message: %v", event.Err)
+		}
+	}
+
+	if capturedErr == nil {
+		t.Log("DESIGN CONCERN: Cross-type agent tool (AgenticMessage sub-agent in Message agent) " +
+			"only errors at event forwarding time when streaming is enabled. " +
+			"The error check happens in the gen.Send path, which is only exercised " +
+			"when the outer agent actually calls the tool AND streaming is enabled. " +
+			"Without streaming, the tool result is returned as a string, so no type mismatch occurs.")
+	} else {
+		assert.Contains(t, capturedErr.Error(), "cross-message-type",
+			"Error should mention cross-message-type incompatibility")
+	}
+}