feat(adk): integrate AgenticMessage into ADK (#920)

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,41 +151,41 @@ 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,
+		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 {
@@ -178,14 +197,14 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 		agentOpts := extractAndDeriveCancelCtx(ctx, at.agent.Name(ctx), opts)
 		agentOpts = append(agentOpts, withSharedParentSession())
 
-		iter, err = newInvokableAgentToolRunner(at.agent, ms, enableStreaming).
-			Resume(ctx, bridgeCheckpointID, agentOpts...)
+		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 {
@@ -211,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")
+				}
 			}
 		}
 
@@ -244,7 +267,7 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 			if err != nil {
 				return "", err
 			}
-			ret = msg.Content
+			ret = extractTextContent(msg)
 		}
 	}
 
@@ -308,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 {
@@ -339,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,9 +21,11 @@ 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"
@@ -31,6 +33,24 @@ import (
 	"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")
+	}
+}

adk/agentic_callback_integration_test.go

@@ -0,0 +1,268 @@
+/*
+ * Copyright 2026 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"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/cloudwego/eino/callbacks"
+	"github.com/cloudwego/eino/components/model"
+	"github.com/cloudwego/eino/schema"
+)
+
+type agenticCallbackRecorder struct {
+	mu             sync.Mutex
+	onStartCalled  bool
+	onEndCalled    bool
+	runInfo        *callbacks.RunInfo
+	inputReceived  *TypedAgentCallbackInput[*schema.AgenticMessage]
+	eventsReceived []*TypedAgentEvent[*schema.AgenticMessage]
+	eventsDone     chan struct{}
+	closeOnce      sync.Once
+}
+
+func (r *agenticCallbackRecorder) getOnStartCalled() bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.onStartCalled
+}
+
+func (r *agenticCallbackRecorder) getOnEndCalled() bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.onEndCalled
+}
+
+func (r *agenticCallbackRecorder) getEventsReceived() []*TypedAgentEvent[*schema.AgenticMessage] {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	result := make([]*TypedAgentEvent[*schema.AgenticMessage], len(r.eventsReceived))
+	copy(result, r.eventsReceived)
+	return result
+}
+
+func newAgenticRecordingHandler(recorder *agenticCallbackRecorder) callbacks.Handler {
+	recorder.eventsDone = make(chan struct{})
+	return callbacks.NewHandlerBuilder().
+		OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
+			if info.Component != ComponentOfAgenticAgent {
+				return ctx
+			}
+			recorder.mu.Lock()
+			defer recorder.mu.Unlock()
+			recorder.onStartCalled = true
+			recorder.runInfo = info
+			if agentInput := ConvTypedCallbackInput[*schema.AgenticMessage](input); agentInput != nil {
+				recorder.inputReceived = agentInput
+			}
+			return ctx
+		}).
+		OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+			if info.Component != ComponentOfAgenticAgent {
+				return ctx
+			}
+			recorder.mu.Lock()
+			recorder.onEndCalled = true
+			recorder.runInfo = info
+			recorder.mu.Unlock()
+
+			if agentOutput := ConvTypedCallbackOutput[*schema.AgenticMessage](output); agentOutput != nil {
+				if agentOutput.Events != nil {
+					go func() {
+						defer recorder.closeOnce.Do(func() { close(recorder.eventsDone) })
+						for {
+							event, ok := agentOutput.Events.Next()
+							if !ok {
+								break
+							}
+							recorder.mu.Lock()
+							recorder.eventsReceived = append(recorder.eventsReceived, event)
+							recorder.mu.Unlock()
+						}
+					}()
+					return ctx
+				}
+			}
+			recorder.closeOnce.Do(func() { close(recorder.eventsDone) })
+			return ctx
+		}).
+		Build()
+}
+
+func TestAgenticCallback(t *testing.T) {
+	ctx := context.Background()
+
+	expectedContent := "This is the test response content"
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg(expectedContent), nil
+		},
+	}
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "TestChatAgent",
+		Description: "Test chat agent",
+		Instruction: "You are a test agent",
+		Model:       m,
+	})
+	require.NoError(t, err)
+
+	recorder := &agenticCallbackRecorder{}
+	handler := newAgenticRecordingHandler(recorder)
+
+	var agentEvents []*TypedAgentEvent[*schema.AgenticMessage]
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent})
+	iter := runner.Query(ctx, "hello", WithCallbacks(handler))
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		agentEvents = append(agentEvents, event)
+	}
+
+	<-recorder.eventsDone
+	assertAgenticEventRoleFields(t, agentEvents)
+
+	t.Run("OnStart_Invocation", func(t *testing.T) {
+		assert.True(t, recorder.getOnStartCalled(), "OnStart should be called")
+		require.NotNil(t, recorder.inputReceived, "Input should be received")
+		require.NotNil(t, recorder.inputReceived.Input, "AgentInput should be set")
+		assert.Len(t, recorder.inputReceived.Input.Messages, 1)
+	})
+
+	t.Run("OnEnd_Invocation", func(t *testing.T) {
+		assert.True(t, recorder.getOnEndCalled(), "OnEnd should be called")
+		assert.Len(t, recorder.getEventsReceived(), 1)
+	})
+
+	t.Run("RunInfo_Fields", func(t *testing.T) {
+		require.NotNil(t, recorder.runInfo)
+		assert.Equal(t, "TestChatAgent", recorder.runInfo.Name)
+		assert.Equal(t, ComponentOfAgenticAgent, recorder.runInfo.Component)
+	})
+
+	t.Run("Events_MatchAgentOutput", func(t *testing.T) {
+		require.NotEmpty(t, agentEvents, "Agent should emit events")
+		received := recorder.getEventsReceived()
+		require.NotEmpty(t, received, "Callback should receive events")
+
+		require.Len(t, received, 1, "Callback should receive exactly 1 event")
+		require.NotNil(t, received[0].Output)
+		require.NotNil(t, received[0].Output.MessageOutput)
+		require.NotNil(t, received[0].Output.MessageOutput.Message)
+		assert.Equal(t, expectedContent, agenticTextContent(received[0].Output.MessageOutput.Message))
+	})
+}
+
+func TestAgenticCallbackMultipleHandlers(t *testing.T) {
+	ctx := context.Background()
+
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg("test response"), nil
+		},
+	}
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "TestAgent",
+		Description: "Test agent",
+		Instruction: "You are a test agent",
+		Model:       m,
+	})
+	require.NoError(t, err)
+
+	recorder1 := &agenticCallbackRecorder{}
+	recorder2 := &agenticCallbackRecorder{}
+	handler1 := newAgenticRecordingHandler(recorder1)
+	handler2 := newAgenticRecordingHandler(recorder2)
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent})
+	iter := runner.Query(ctx, "hello", WithCallbacks(handler1, handler2))
+	for {
+		_, ok := iter.Next()
+		if !ok {
+			break
+		}
+	}
+
+	<-recorder1.eventsDone
+	<-recorder2.eventsDone
+
+	assert.True(t, recorder1.getOnStartCalled(), "Handler1 OnStart should be called")
+	assert.True(t, recorder2.getOnStartCalled(), "Handler2 OnStart should be called")
+	assert.True(t, recorder1.getOnEndCalled(), "Handler1 OnEnd should be called")
+	assert.True(t, recorder2.getOnEndCalled(), "Handler2 OnEnd should be called")
+
+	assert.NotEmpty(t, recorder1.getEventsReceived(), "Handler1 should receive events")
+	assert.NotEmpty(t, recorder2.getEventsReceived(), "Handler2 should receive events")
+}
+
+func TestCoverage_WrapAgenticIterWithOnEnd(t *testing.T) {
+	ctx := context.Background()
+
+	var onEndCalled bool
+	handler := callbacks.NewHandlerBuilder().
+		OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
+			return ctx
+		}).
+		OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
+			if info.Component == ComponentOfAgenticAgent {
+				onEndCalled = true
+			}
+			return ctx
+		}).
+		Build()
+
+	ctx = initAgenticCallbacks(ctx, "test-agent", "ChatModel",
+		WithCallbacks(handler))
+
+	cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{
+		Input: &TypedAgentInput[*schema.AgenticMessage]{
+			Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")},
+		},
+	}
+	ctx = callbacks.OnStart(ctx, cbInput)
+
+	origIter, origGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+	go func() {
+		defer origGen.Close()
+		origGen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+			Output: &TypedAgentOutput[*schema.AgenticMessage]{
+				MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
+					Message: agenticMsg("done"),
+				},
+			},
+		})
+	}()
+
+	wrappedIter := wrapAgenticIterWithOnEnd(ctx, origIter)
+
+	for {
+		_, ok := wrappedIter.Next()
+		if !ok {
+			break
+		}
+	}
+
+	assert.True(t, onEndCalled, "OnEnd callback should have been called")
+}

adk/agentic_integration_test.go

@@ -0,0 +1,665 @@
+/*
+ * Copyright 2026 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"
+	"encoding/json"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/eino-contrib/jsonschema"
+	"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"
+)
+
+func agenticMsg(text string) *schema.AgenticMessage {
+	return &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeAssistant,
+		ContentBlocks: []*schema.ContentBlock{
+			schema.NewContentBlock(&schema.AssistantGenText{Text: text}),
+		},
+	}
+}
+
+func agenticTextContent(msg *schema.AgenticMessage) string {
+	for _, b := range msg.ContentBlocks {
+		if b.AssistantGenText != nil {
+			return b.AssistantGenText.Text
+		}
+	}
+	return ""
+}
+
+func TestAgenticIntegration_ChatModelSingleShot(t *testing.T) {
+	ctx := context.Background()
+
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg("Handled internally with tool result: 42"), nil
+		},
+	}
+
+	dummyTool := newSlowTool("calculator", 0, "42")
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "ToolCallAgent",
+		Description: "Agent with tools for agentic model",
+		Instruction: "You are a calculator.",
+		Model:       m,
+		ToolsConfig: ToolsConfig{
+			ToolsNodeConfig: compose.ToolsNodeConfig{
+				Tools: []tool.BaseTool{dummyTool},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent: agent,
+	})
+
+	iter := runner.Query(ctx, "What is 6*7?")
+
+	var events []*TypedAgentEvent[*schema.AgenticMessage]
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		events = append(events, event)
+	}
+
+	require.Len(t, events, 1)
+	assertAgenticEventRoleFields(t, events)
+	lastEvent := events[len(events)-1]
+	require.Nil(t, lastEvent.Err)
+	require.NotNil(t, lastEvent.Output)
+	require.NotNil(t, lastEvent.Output.MessageOutput)
+	assert.Equal(t, "Handled internally with tool result: 42",
+		agenticTextContent(lastEvent.Output.MessageOutput.Message))
+}
+
+func TestAgenticIntegration_ChatModelToolsPassedViaOptions(t *testing.T) {
+	ctx := context.Background()
+
+	var receivedTools []*schema.ToolInfo
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			o := model.GetCommonOptions(&model.Options{}, opts...)
+			receivedTools = o.Tools
+			return agenticMsg("done"), nil
+		},
+	}
+
+	dummyTool := newSlowTool("my_tool", 0, "result")
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "ToolOptAgent",
+		Description: "Agent verifying tools are passed via options",
+		Model:       m,
+		ToolsConfig: ToolsConfig{
+			ToolsNodeConfig: compose.ToolsNodeConfig{
+				Tools: []tool.BaseTool{dummyTool},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent: agent,
+	})
+	iter := runner.Query(ctx, "test tools")
+	for {
+		_, ok := iter.Next()
+		if !ok {
+			break
+		}
+	}
+
+	require.NotNil(t, receivedTools, "tools should be passed via model.Options")
+	require.Len(t, receivedTools, 1)
+	assert.Equal(t, "my_tool", receivedTools[0].Name)
+}
+
+func TestAgenticIntegration_StreamingWithRunner(t *testing.T) {
+	ctx := context.Background()
+
+	chunk1 := &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeAssistant,
+		ContentBlocks: []*schema.ContentBlock{
+			schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}),
+		},
+	}
+	chunk2 := &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeAssistant,
+		ContentBlocks: []*schema.ContentBlock{
+			schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}),
+		},
+	}
+
+	m := &mockAgenticModel{
+		streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) {
+			r, w := schema.Pipe[*schema.AgenticMessage](2)
+			go func() {
+				defer w.Close()
+				w.Send(chunk1, nil)
+				w.Send(chunk2, nil)
+			}()
+			return r, nil
+		},
+	}
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "StreamRunner",
+		Description: "Streaming runner agent",
+		Model:       m,
+	})
+	require.NoError(t, err)
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent:           agent,
+		EnableStreaming: true,
+	})
+
+	iter := runner.Query(ctx, "stream me")
+
+	event, ok := iter.Next()
+	require.True(t, ok)
+	assert.Nil(t, event.Err)
+	require.NotNil(t, event.Output)
+	require.NotNil(t, event.Output.MessageOutput)
+
+	if event.Output.MessageOutput.IsStreaming {
+		require.NotNil(t, event.Output.MessageOutput.MessageStream)
+		var chunks []*schema.AgenticMessage
+		for {
+			chunk, err := event.Output.MessageOutput.MessageStream.Recv()
+			if err != nil {
+				break
+			}
+			chunks = append(chunks, chunk)
+		}
+		assert.Equal(t, 2, len(chunks))
+	} else {
+		assert.NotNil(t, event.Output.MessageOutput.Message)
+	}
+
+	_, ok = iter.Next()
+	assert.False(t, ok)
+}
+
+func TestAgenticIntegration_CancelDuringExecution(t *testing.T) {
+	ctx := context.Background()
+
+	modelStarted := make(chan struct{}, 1)
+	modelBlocked := make(chan struct{})
+
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			select {
+			case modelStarted <- struct{}{}:
+			default:
+			}
+			select {
+			case <-modelBlocked:
+				return agenticMsg("should not reach"), nil
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			}
+		},
+	}
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "CancelAgent",
+		Description: "cancel test",
+		Model:       m,
+	})
+	require.NoError(t, err)
+
+	cancelCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent: agent,
+	})
+	iter := runner.Run(cancelCtx, []*schema.AgenticMessage{
+		schema.UserAgenticMessage("Hi"),
+	})
+
+	<-modelStarted
+	cancel()
+
+	var capturedErr error
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if event.Err != nil {
+			capturedErr = event.Err
+		}
+	}
+	require.Error(t, capturedErr, "should propagate cancel error")
+	assert.ErrorIs(t, capturedErr, context.Canceled)
+}
+
+func TestAgenticIntegration_CancelWithTimeout(t *testing.T) {
+	ctx := context.Background()
+
+	sa := &myAgenticAgent{
+		name: "slow-agent",
+		runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer generator.Close()
+				select {
+				case <-time.After(10 * time.Second):
+					generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+						Output: &TypedAgentOutput[*schema.AgenticMessage]{
+							MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
+								Message: agenticMsg("slow response"),
+							},
+						},
+					})
+				case <-ctx.Done():
+					generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+						Err: ctx.Err(),
+					})
+				}
+			}()
+			return iter
+		},
+	}
+
+	timeoutCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+	defer cancel()
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent: sa,
+	})
+	iter := runner.Run(timeoutCtx, []*schema.AgenticMessage{
+		schema.UserAgenticMessage("slow request"),
+	})
+
+	var capturedErr error
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if event.Err != nil {
+			capturedErr = event.Err
+		}
+	}
+
+	require.Error(t, capturedErr, "should get timeout/cancel error")
+	assert.ErrorIs(t, capturedErr, context.DeadlineExceeded)
+}
+func TestAgenticIntegration_AgentTool(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 tool result"), nil
+		},
+	}
+
+	innerAgent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "InnerAgent",
+		Description: "An agent used as a tool",
+		Model:       innerModel,
+	})
+	require.NoError(t, err)
+
+	agentTool := NewTypedAgentTool(ctx, TypedAgent[*schema.AgenticMessage](innerAgent))
+	require.NotNil(t, agentTool)
+
+	info, err := agentTool.Info(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, "InnerAgent", info.Name)
+	assert.Equal(t, "An agent used as a tool", info.Desc)
+
+	outerModel := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg("outer response after inner tool"), nil
+		},
+	}
+
+	outerAgent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "OuterAgent",
+		Description: "Outer agent with agent tool",
+		Model:       outerModel,
+		ToolsConfig: ToolsConfig{
+			ToolsNodeConfig: compose.ToolsNodeConfig{
+				Tools: []tool.BaseTool{agentTool},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent: outerAgent,
+	})
+	iter := runner.Query(ctx, "delegate to inner")
+
+	var events []*TypedAgentEvent[*schema.AgenticMessage]
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		events = append(events, event)
+	}
+
+	require.NotEmpty(t, events)
+	assertAgenticEventRoleFields(t, events)
+	lastEvent := events[len(events)-1]
+	assert.Nil(t, lastEvent.Err)
+	assert.NotNil(t, lastEvent.Output)
+}
+func TestAgenticIntegration_InterruptEventFormation(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("simple interrupt", func(t *testing.T) {
+		agent := &myAgenticAgent{
+			name: "int-agent",
+			runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+				iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+				go func() {
+					defer generator.Close()
+					intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "approval needed")
+					intEvent.Action.Interrupted.Data = "approval data"
+					generator.Send(intEvent)
+				}()
+				return iter
+			},
+		}
+
+		runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+			Agent: agent,
+		})
+		iter := runner.Query(ctx, "interrupt test")
+
+		var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			if event.Action != nil && event.Action.Interrupted != nil {
+				interruptEvent = event
+			}
+		}
+
+		require.NotNil(t, interruptEvent)
+		assert.Equal(t, "approval data", interruptEvent.Action.Interrupted.Data)
+		require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
+		assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID)
+		assert.Equal(t, "approval needed", interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
+		assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause)
+	})
+
+	t.Run("stateful interrupt", func(t *testing.T) {
+		agent := &myAgenticAgent{
+			name: "st-agent",
+			runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+				iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+				go func() {
+					defer generator.Close()
+					intEvent := TypedStatefulInterrupt[*schema.AgenticMessage](ctx, "state interrupt", "my-state")
+					intEvent.Action.Interrupted.Data = "stateful data"
+					generator.Send(intEvent)
+				}()
+				return iter
+			},
+		}
+
+		runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+			Agent: agent,
+		})
+		iter := runner.Query(ctx, "stateful test")
+
+		var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			if event.Action != nil && event.Action.Interrupted != nil {
+				interruptEvent = event
+			}
+		}
+
+		require.NotNil(t, interruptEvent)
+		assert.Equal(t, "stateful data", interruptEvent.Action.Interrupted.Data)
+		require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
+		assert.Equal(t, "state interrupt", interruptEvent.Action.Interrupted.InterruptContexts[0].Info)
+	})
+}
+func TestAgenticIntegration_CheckpointInterruptResume(t *testing.T) {
+	ctx := context.Background()
+
+	var resumeCalled int32
+	agent := &myAgenticAgent{
+		name: "ckpt-agent",
+		runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer generator.Close()
+				generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+					AgentName: "ckpt-agent",
+					Output: &TypedAgentOutput[*schema.AgenticMessage]{
+						MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
+							Message: agenticMsg("before interrupt"),
+						},
+					},
+				})
+				intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "need approval")
+				intEvent.Action.Interrupted.Data = "approval data"
+				generator.Send(intEvent)
+			}()
+			return iter
+		},
+		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			atomic.StoreInt32(&resumeCalled, 1)
+			iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer generator.Close()
+				generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+					AgentName: "ckpt-agent",
+					Output: &TypedAgentOutput[*schema.AgenticMessage]{
+						MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
+							Message: agenticMsg("after resume"),
+						},
+					},
+				})
+			}()
+			return iter
+		},
+	}
+
+	store := newMyStore()
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent:           agent,
+		CheckPointStore: store,
+	})
+
+	iter := runner.Query(ctx, "checkpoint test", WithCheckPointID("ckpt-1"))
+
+	var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
+	var preInterruptOutputs []string
+	for {
+		event, ok := iter.Next()
+		if !ok {
+			break
+		}
+		require.Nil(t, event.Err)
+		if event.Action != nil && event.Action.Interrupted != nil {
+			interruptEvent = event
+		}
+		if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {
+			preInterruptOutputs = append(preInterruptOutputs, agenticTextContent(event.Output.MessageOutput.Message))
+		}
+	}
+
+	require.NotNil(t, interruptEvent, "should receive interrupt event")
+	assert.Contains(t, preInterruptOutputs, "before interrupt")
+	require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts)
+
+	interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
+	require.NotEmpty(t, interruptID)
+
+	resumeIter, err := runner.ResumeWithParams(ctx, "ckpt-1", &ResumeParams{
+		Targets: map[string]any{
+			interruptID: nil,
+		},
+	})
+	require.NoError(t, err)
+
+	var postResumeOutputs []string
+	for {
+		event, ok := resumeIter.Next()
+		if !ok {
+			break
+		}
+		if event.Err != nil {
+			t.Fatalf("unexpected error during resume: %v", event.Err)
+		}
+		if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil {
+			postResumeOutputs = append(postResumeOutputs, agenticTextContent(event.Output.MessageOutput.Message))
+		}
+	}
+
+	assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled), "resume function should have been called")
+	assert.Contains(t, postResumeOutputs, "after resume")
+}
+
+func TestAgenticIntegration_CheckpointWithMCPListToolsResult(t *testing.T) {
+	ctx := context.Background()
+
+	inputSchemaJSON := `{
+		"type": "object",
+		"properties": {
+			"query": {"type": "string", "description": "search query"},
+			"limit": {"type": "integer", "description": "max results"}
+		},
+		"required": ["query"]
+	}`
+	var inputSchema jsonschema.Schema
+	require.NoError(t, json.Unmarshal([]byte(inputSchemaJSON), &inputSchema))
+
+	mcpMsg := &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeAssistant,
+		ContentBlocks: []*schema.ContentBlock{
+			{
+				Type: schema.ContentBlockTypeMCPListToolsResult,
+				MCPListToolsResult: &schema.MCPListToolsResult{
+					ServerLabel: "test-server",
+					Tools: []*schema.MCPListToolsItem{
+						{
+							Name:        "search",
+							Description: "search the web",
+							InputSchema: &inputSchema,
+						},
+					},
+				},
+			},
+			schema.NewContentBlock(&schema.AssistantGenText{Text: "here are tools"}),
+		},
+	}
+
+	var resumeCalled int32
+	agent := &myAgenticAgent{
+		name: "mcp-agent",
+		runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer gen.Close()
+				gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+					AgentName: "mcp-agent",
+					Output: &TypedAgentOutput[*schema.AgenticMessage]{
+						MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: mcpMsg},
+					},
+				})
+				gen.Send(TypedInterrupt[*schema.AgenticMessage](ctx, "approve tools"))
+			}()
+			return iter
+		},
+		resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			atomic.StoreInt32(&resumeCalled, 1)
+			iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer gen.Close()
+				gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+					AgentName: "mcp-agent",
+					Output: &TypedAgentOutput[*schema.AgenticMessage]{
+						MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("tools approved")},
+					},
+				})
+			}()
+			return iter
+		},
+	}
+
+	store := newMyStore()
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent:           agent,
+		CheckPointStore: store,
+	})
+
+	iter := runner.Query(ctx, "list tools", WithCheckPointID("mcp-1"))
+	var interruptEvent *TypedAgentEvent[*schema.AgenticMessage]
+	for {
+		ev, ok := iter.Next()
+		if !ok {
+			break
+		}
+		require.Nil(t, ev.Err)
+		if ev.Action != nil && ev.Action.Interrupted != nil {
+			interruptEvent = ev
+		}
+	}
+	require.NotNil(t, interruptEvent)
+	interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID
+
+	resumeIter, err := runner.ResumeWithParams(ctx, "mcp-1", &ResumeParams{
+		Targets: map[string]any{interruptID: nil},
+	})
+	require.NoError(t, err)
+
+	var outputs []string
+	for {
+		ev, ok := resumeIter.Next()
+		if !ok {
+			break
+		}
+		require.Nil(t, ev.Err)
+		if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.Message != nil {
+			outputs = append(outputs, agenticTextContent(ev.Output.MessageOutput.Message))
+		}
+	}
+
+	assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled))
+	assert.Contains(t, outputs, "tools approved")
+}

adk/callback.go

@@ -43,18 +43,18 @@ type AgentCallbackOutput struct {
 	Events *AsyncIterator[*AgentEvent]
 }
 
-func copyEventIterator(iter *AsyncIterator[*AgentEvent], n int) []*AsyncIterator[*AgentEvent] {
+func copyTypedEventIterator[M messageType](iter *AsyncIterator[*TypedAgentEvent[M]], n int) []*AsyncIterator[*TypedAgentEvent[M]] {
 	if n <= 0 {
 		return nil
 	}
 	if n == 1 {
-		return []*AsyncIterator[*AgentEvent]{iter}
+		return []*AsyncIterator[*TypedAgentEvent[M]]{iter}
 	}
 
-	iterators := make([]*AsyncIterator[*AgentEvent], n)
-	generators := make([]*AsyncGenerator[*AgentEvent], n)
+	iterators := make([]*AsyncIterator[*TypedAgentEvent[M]], n)
+	generators := make([]*AsyncGenerator[*TypedAgentEvent[M]], n)
 	for i := 0; i < n; i++ {
-		iterators[i], generators[i] = NewAsyncIteratorPair[*AgentEvent]()
+		iterators[i], generators[i] = NewAsyncIteratorPair[*TypedAgentEvent[M]]()
 	}
 
 	go func() {
@@ -70,7 +70,7 @@ func copyEventIterator(iter *AsyncIterator[*AgentEvent], n int) []*AsyncIterator
 				break
 			}
 			for i := 0; i < n-1; i++ {
-				generators[i].Send(copyAgentEvent(event))
+				generators[i].Send(copyTypedAgentEvent(event))
 			}
 			generators[n-1].Send(event)
 		}
@@ -87,7 +87,7 @@ func copyAgentCallbackOutput(out *AgentCallbackOutput, n int) []*AgentCallbackOu
 		}
 		return result
 	}
-	iters := copyEventIterator(out.Events, n)
+	iters := copyTypedEventIterator(out.Events, n)
 	result := make([]*AgentCallbackOutput, n)
 	for i, iter := range iters {
 		result[i] = &AgentCallbackOutput{Events: iter}
@@ -133,3 +133,70 @@ func getAgentType(agent Agent) string {
 	}
 	return ""
 }
+
+// TypedAgentCallbackInput represents the input passed to typed agent callbacks during OnStart.
+// Use ConvTypedCallbackInput to safely convert from callbacks.CallbackInput.
+type TypedAgentCallbackInput[M messageType] struct {
+	// Input contains the agent input for a new run. Nil when resuming.
+	Input *TypedAgentInput[M]
+	// ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs.
+	ResumeInfo *ResumeInfo
+}
+
+// TypedAgentCallbackOutput represents the output passed to typed agent callbacks during OnEnd.
+// Use ConvTypedCallbackOutput to safely convert from callbacks.CallbackOutput.
+//
+// Important: The Events iterator should be consumed asynchronously to avoid blocking
+// the agent execution. Each callback handler receives an independent copy of the iterator.
+type TypedAgentCallbackOutput[M messageType] struct {
+	// Events provides a stream of agent events. Each handler receives its own copy.
+	Events *AsyncIterator[*TypedAgentEvent[M]]
+}
+
+// ConvTypedCallbackInput converts a callbacks.CallbackInput to *TypedAgentCallbackInput[M].
+// Returns nil if the input is not of the expected type.
+func ConvTypedCallbackInput[M messageType](input callbacks.CallbackInput) *TypedAgentCallbackInput[M] {
+	if v, ok := input.(*TypedAgentCallbackInput[M]); ok {
+		return v
+	}
+	return nil
+}
+
+// ConvTypedCallbackOutput converts a callbacks.CallbackOutput to *TypedAgentCallbackOutput[M].
+// Returns nil if the output is not of the expected type.
+func ConvTypedCallbackOutput[M messageType](output callbacks.CallbackOutput) *TypedAgentCallbackOutput[M] {
+	if v, ok := output.(*TypedAgentCallbackOutput[M]); ok {
+		return v
+	}
+	return nil
+}
+
+func copyTypedCallbackOutput[M messageType](out *TypedAgentCallbackOutput[M], n int) []*TypedAgentCallbackOutput[M] {
+	if out == nil || out.Events == nil {
+		result := make([]*TypedAgentCallbackOutput[M], n)
+		for i := 0; i < n; i++ {
+			result[i] = out
+		}
+		return result
+	}
+	iters := copyTypedEventIterator(out.Events, n)
+	result := make([]*TypedAgentCallbackOutput[M], n)
+	for i, iter := range iters {
+		result[i] = &TypedAgentCallbackOutput[M]{Events: iter}
+	}
+	return result
+}
+
+func initAgenticCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context {
+	ri := &callbacks.RunInfo{
+		Name:      agentName,
+		Type:      agentType,
+		Component: ComponentOfAgenticAgent,
+	}
+
+	o := getCommonOptions(nil, opts...)
+	if len(o.handlers) == 0 {
+		return icb.ReuseHandlers(ctx, ri)
+	}
+	return icb.AppendHandlers(ctx, ri, o.handlers...)
+}

adk/callback_test.go

@@ -22,12 +22,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/cloudwego/eino/callbacks"
 	"github.com/cloudwego/eino/schema"
 )
 
-func TestCopyEventIterator(t *testing.T) {
+func TestCopyTypedEventIterator(t *testing.T) {
 	t.Run("n=0 returns nil", func(t *testing.T) {
 		iter, gen := NewAsyncIteratorPair[*AgentEvent]()
 		go func() {
@@ -35,7 +36,7 @@ func TestCopyEventIterator(t *testing.T) {
 			gen.Close()
 		}()
 
-		result := copyEventIterator(iter, 0)
+		result := copyTypedEventIterator(iter, 0)
 		assert.Nil(t, result)
 	})
 
@@ -46,7 +47,7 @@ func TestCopyEventIterator(t *testing.T) {
 			gen.Close()
 		}()
 
-		result := copyEventIterator(iter, 1)
+		result := copyTypedEventIterator(iter, 1)
 		assert.Len(t, result, 1)
 		assert.Equal(t, iter, result[0])
 	})
@@ -66,7 +67,7 @@ func TestCopyEventIterator(t *testing.T) {
 		}()
 
 		n := 3
-		copies := copyEventIterator(iter, n)
+		copies := copyTypedEventIterator(iter, n)
 		assert.Len(t, copies, n)
 
 		var wg sync.WaitGroup
@@ -127,7 +128,7 @@ func TestCopyAgentCallbackOutput(t *testing.T) {
 		assert.Len(t, result, 2)
 
 		for i, r := range result {
-			assert.NotNil(t, r, "result[%d] should not be nil", i)
+			require.NotNil(t, r, "result[%d] should not be nil", i)
 			assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i)
 		}
 	})
@@ -234,3 +235,154 @@ func TestWithMultipleCallbacksOption(t *testing.T) {
 
 	assert.Len(t, opts.handlers, 2)
 }
+
+func TestCopyTypedEventIteratorAgentic(t *testing.T) {
+	t.Run("n=0 returns nil", func(t *testing.T) {
+		iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+		go func() {
+			gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
+			gen.Close()
+		}()
+
+		result := copyTypedEventIterator(iter, 0)
+		assert.Nil(t, result)
+	})
+
+	t.Run("n=1 returns original iterator", func(t *testing.T) {
+		iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+		go func() {
+			gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
+			gen.Close()
+		}()
+
+		result := copyTypedEventIterator(iter, 1)
+		assert.Len(t, result, 1)
+		assert.Equal(t, iter, result[0])
+	})
+
+	t.Run("n>1 creates n independent copies", func(t *testing.T) {
+		iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+		events := []*TypedAgentEvent[*schema.AgenticMessage]{
+			{AgentName: "agent1", Output: &TypedAgentOutput[*schema.AgenticMessage]{
+				MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg1")},
+			}},
+			{AgentName: "agent2", Output: &TypedAgentOutput[*schema.AgenticMessage]{
+				MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg2")},
+			}},
+		}
+
+		go func() {
+			for _, e := range events {
+				gen.Send(e)
+			}
+			gen.Close()
+		}()
+
+		n := 3
+		copies := copyTypedEventIterator(iter, n)
+		assert.Len(t, copies, n)
+
+		var wg sync.WaitGroup
+		receivedEvents := make([][]*TypedAgentEvent[*schema.AgenticMessage], n)
+
+		for i := 0; i < n; i++ {
+			wg.Add(1)
+			go func(idx int) {
+				defer wg.Done()
+				for {
+					event, ok := copies[idx].Next()
+					if !ok {
+						break
+					}
+					receivedEvents[idx] = append(receivedEvents[idx], event)
+				}
+			}(i)
+		}
+
+		wg.Wait()
+
+		for i := 0; i < n; i++ {
+			assert.Len(t, receivedEvents[i], len(events), "iterator %d should receive all events", i)
+			for j, e := range receivedEvents[i] {
+				assert.Equal(t, events[j].AgentName, e.AgentName)
+			}
+		}
+	})
+}
+
+func TestCopyTypedCallbackOutput(t *testing.T) {
+	t.Run("nil output", func(t *testing.T) {
+		result := copyTypedCallbackOutput[*schema.AgenticMessage](nil, 3)
+		assert.Len(t, result, 3)
+		for _, r := range result {
+			assert.Nil(t, r)
+		}
+	})
+
+	t.Run("output with nil Events", func(t *testing.T) {
+		out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: nil}
+		result := copyTypedCallbackOutput(out, 3)
+		assert.Len(t, result, 3)
+		for _, r := range result {
+			assert.Equal(t, out, r)
+		}
+	})
+
+	t.Run("valid output with events", func(t *testing.T) {
+		iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+		go func() {
+			gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"})
+			gen.Close()
+		}()
+
+		out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter}
+		result := copyTypedCallbackOutput(out, 2)
+		assert.Len(t, result, 2)
+
+		for i, r := range result {
+			require.NotNil(t, r, "result[%d] should not be nil", i)
+			assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i)
+		}
+	})
+}
+
+func TestConvTypedCallbackInput(t *testing.T) {
+	t.Run("valid TypedAgentCallbackInput", func(t *testing.T) {
+		input := &TypedAgentCallbackInput[*schema.AgenticMessage]{
+			Input: &TypedAgentInput[*schema.AgenticMessage]{
+				Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")},
+			},
+		}
+		result := ConvTypedCallbackInput[*schema.AgenticMessage](input)
+		assert.Equal(t, input, result)
+	})
+
+	t.Run("invalid type returns nil", func(t *testing.T) {
+		result := ConvTypedCallbackInput[*schema.AgenticMessage]("invalid")
+		assert.Nil(t, result)
+	})
+
+	t.Run("nil returns nil", func(t *testing.T) {
+		result := ConvTypedCallbackInput[*schema.AgenticMessage](nil)
+		assert.Nil(t, result)
+	})
+}
+
+func TestConvTypedCallbackOutput(t *testing.T) {
+	t.Run("valid TypedAgentCallbackOutput", func(t *testing.T) {
+		iter, _ := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+		output := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter}
+		result := ConvTypedCallbackOutput[*schema.AgenticMessage](output)
+		assert.Equal(t, output, result)
+	})
+
+	t.Run("invalid type returns nil", func(t *testing.T) {
+		result := ConvTypedCallbackOutput[*schema.AgenticMessage]("invalid")
+		assert.Nil(t, result)
+	})
+
+	t.Run("nil returns nil", func(t *testing.T) {
+		result := ConvTypedCallbackOutput[*schema.AgenticMessage](nil)
+		assert.Nil(t, result)
+	})
+}

adk/cancel.go

@@ -812,11 +812,11 @@ func (cc *cancelContext) buildCancelFunc() AgentCancelFunc {
 // were passed through unconverted, markDone would transition stateCancelling→stateDone
 // before the Runner goroutine could call createAndMarkCancelHandled, causing it
 // to fail the CAS.
-func wrapIterWithCancelCtx(iter *AsyncIterator[*AgentEvent], cancelCtx *cancelContext) *AsyncIterator[*AgentEvent] {
+func wrapIterWithCancelCtx[M messageType](iter *AsyncIterator[*TypedAgentEvent[M]], cancelCtx *cancelContext) *AsyncIterator[*TypedAgentEvent[M]] {
 	if cancelCtx == nil {
 		return iter
 	}
-	it, gen := NewAsyncIteratorPair[*AgentEvent]()
+	it, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
 	go func() {
 		defer cancelCtx.markDone()
 		defer gen.Close()
@@ -831,7 +831,7 @@ func wrapIterWithCancelCtx(iter *AsyncIterator[*AgentEvent], cancelCtx *cancelCo
 					cancelErr, ok := cancelCtx.createAndMarkCancelHandled()
 					if ok {
 						cancelErr.interruptSignal = event.Action.internalInterrupted
-						gen.Send(&AgentEvent{Err: cancelErr})
+						gen.Send(&TypedAgentEvent[M]{Err: cancelErr})
 					}
 					return
 				}
@@ -843,13 +843,13 @@ func wrapIterWithCancelCtx(iter *AsyncIterator[*AgentEvent], cancelCtx *cancelCo
 	return it
 }
 
-// cancelMonitoredModel wraps a model with cancel monitoring.
+// typedCancelMonitoredModel wraps a model with cancel monitoring.
 // Generate: pure delegate to the inner model (CancelAfterChatModel is handled
 // by a dedicated node after the ChatModel in the compose graph).
 // Stream: pipes chunks through a goroutine that selects on immediateChan for
 // CancelImmediate abort.
-type cancelMonitoredModel struct {
-	inner         model.BaseChatModel
+type typedCancelMonitoredModel[M messageType] struct {
+	inner         model.BaseModel[M]
 	cancelContext *cancelContext
 }
 
@@ -858,11 +858,11 @@ type recvResult[T any] struct {
 	err  error
 }
 
-func (m *cancelMonitoredModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (m *typedCancelMonitoredModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	return m.inner.Generate(ctx, input, opts...)
 }
 
-func (m *cancelMonitoredModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+func (m *typedCancelMonitoredModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 	stream, err := m.inner.Stream(ctx, input, opts...)
 	if err != nil {
 		return nil, err

adk/cancel_edge_test.go

@@ -1336,7 +1336,7 @@ func TestWithCancel_CancelImmediate_StreamableToolAborted(t *testing.T) {
 	tcm := &toolCallStreamModel{}
 	st := &slowStreamingTool{
 		name:          "slow_tool",
-		chunkInterval: 200 * time.Millisecond,
+		chunkInterval: 100 * time.Millisecond,
 		chunks:        []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"},
 		started:       make(chan struct{}, 1),
 	}
@@ -1366,7 +1366,7 @@ func TestWithCancel_CancelImmediate_StreamableToolAborted(t *testing.T) {
 		t.Fatal("tool did not start streaming")
 	}
 	// Let a few chunks through, then cancel mid-stream
-	time.Sleep(300 * time.Millisecond)
+	time.Sleep(500 * time.Millisecond)
 
 	handle, _ := cancelFn()
 	cancelErr := handle.Wait()

adk/cancel_test.go

@@ -2368,7 +2368,7 @@ func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) {
 			cancelCtx: cc,
 		}
 
-		ctx := withChatModelAgentExecCtx(context.Background(), execCtx)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), execCtx)
 
 		assert.NotPanics(t, func() {
 			err := SendEvent(ctx, &AgentEvent{AgentName: "test"})

adk/chatmodel.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"math"
 	"runtime/debug"
+	"strings"
 	"sync"
 	"sync/atomic"
 
@@ -38,14 +39,15 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-var _ ResumableAgent = &ChatModelAgent{}
+var _ ResumableAgent = &TypedChatModelAgent[*schema.Message]{}
+var _ TypedResumableAgent[*schema.AgenticMessage] = &TypedChatModelAgent[*schema.AgenticMessage]{}
 
-type chatModelAgentExecCtx struct {
+type typedChatModelAgentExecCtx[M messageType] struct {
 	runtimeReturnDirectly map[string]bool
-	generator             *AsyncGenerator[*AgentEvent]
+	generator             *AsyncGenerator[*TypedAgentEvent[M]]
 	cancelCtx             *cancelContext
 
-	failoverLastSuccessModel model.BaseChatModel
+	failoverLastSuccessModel model.BaseModel[M]
 
 	// suppressEventSend prevents eventSenderModel from emitting AgentEvents for the current
 	// Generate call. Set to true before each rejected retry attempt and reset to false after.
@@ -54,7 +56,7 @@ type chatModelAgentExecCtx struct {
 	retryVerdictSignal *retryVerdictSignal
 }
 
-func (e *chatModelAgentExecCtx) send(event *AgentEvent) {
+func (e *typedChatModelAgentExecCtx[M]) send(event *TypedAgentEvent[M]) {
 	if e == nil || e.generator == nil {
 		return
 	}
@@ -64,15 +66,17 @@ func (e *chatModelAgentExecCtx) send(event *AgentEvent) {
 	e.generator.trySend(event)
 }
 
-type chatModelAgentExecCtxKey struct{}
+type chatModelAgentExecCtx = typedChatModelAgentExecCtx[*schema.Message]
 
-func withChatModelAgentExecCtx(ctx context.Context, execCtx *chatModelAgentExecCtx) context.Context {
-	return context.WithValue(ctx, chatModelAgentExecCtxKey{}, execCtx)
+type typedChatModelAgentExecCtxKey[M messageType] struct{}
+
+func withTypedChatModelAgentExecCtx[M messageType](ctx context.Context, execCtx *typedChatModelAgentExecCtx[M]) context.Context {
+	return context.WithValue(ctx, typedChatModelAgentExecCtxKey[M]{}, execCtx)
 }
 
-func getChatModelAgentExecCtx(ctx context.Context) *chatModelAgentExecCtx {
-	if v := ctx.Value(chatModelAgentExecCtxKey{}); v != nil {
-		return v.(*chatModelAgentExecCtx)
+func getTypedChatModelAgentExecCtx[M messageType](ctx context.Context) *typedChatModelAgentExecCtx[M] {
+	if v := ctx.Value(typedChatModelAgentExecCtxKey[M]{}); v != nil {
+		return v.(*typedChatModelAgentExecCtx[M])
 	}
 	return nil
 }
@@ -137,8 +141,14 @@ type ToolsConfig struct {
 	EmitInternalEvents bool
 }
 
+// TypedGenModelInput transforms the agent's system instruction and user input into model input
+// messages ([]M). This is the primary customization point for controlling what the model sees.
+// The default implementation prepends a system message (if instruction is non-empty),
+// followed by the user's input messages.
+type TypedGenModelInput[M messageType] func(ctx context.Context, instruction string, input *TypedAgentInput[M]) ([]M, error)
+
 // GenModelInput transforms agent instructions and input into a format suitable for the model.
-type GenModelInput func(ctx context.Context, instruction string, input *AgentInput) ([]Message, error)
+type GenModelInput = TypedGenModelInput[*schema.Message]
 
 func defaultGenModelInput(ctx context.Context, instruction string, input *AgentInput) ([]Message, error) {
 	msgs := make([]Message, 0, len(input.Messages)+1)
@@ -168,13 +178,35 @@ func defaultGenModelInput(ctx context.Context, instruction string, input *AgentI
 	return msgs, nil
 }
 
-// ChatModelAgentState represents the state of a chat model agent during conversation.
-// This is the primary state type for both ChatModelAgentMiddleware and AgentMiddleware callbacks.
-type ChatModelAgentState struct {
-	// Messages contains all messages in the current conversation session.
-	Messages []Message
+func newDefaultGenModelInput[M messageType]() TypedGenModelInput[M] {
+	var zero M
+	switch any(zero).(type) {
+	case *schema.Message:
+		return any(GenModelInput(defaultGenModelInput)).(TypedGenModelInput[M])
+	case *schema.AgenticMessage:
+		return any(TypedGenModelInput[*schema.AgenticMessage](func(_ context.Context, instruction string, input *TypedAgentInput[*schema.AgenticMessage]) ([]*schema.AgenticMessage, error) {
+			msgs := make([]*schema.AgenticMessage, 0, len(input.Messages)+1)
+			if instruction != "" {
+				msgs = append(msgs, schema.SystemAgenticMessage(instruction))
+			}
+			msgs = append(msgs, input.Messages...)
+			return msgs, nil
+		})).(TypedGenModelInput[M])
+	default:
+		panic("unreachable: unknown messageType")
+	}
 }
 
+// TypedChatModelAgentState represents the state of a chat model agent during conversation.
+// This is the primary state type for both TypedChatModelAgentMiddleware and AgentMiddleware callbacks.
+type TypedChatModelAgentState[M messageType] struct {
+	// Messages contains all messages in the current conversation session.
+	Messages []M
+}
+
+// ChatModelAgentState is the default state type using *schema.Message.
+type ChatModelAgentState = TypedChatModelAgentState[*schema.Message]
+
 // AgentMiddleware provides hooks to customize agent behavior at various stages of execution.
 //
 // Limitations of AgentMiddleware (struct-based):
@@ -207,7 +239,8 @@ type AgentMiddleware struct {
 	WrapToolCall compose.ToolMiddleware
 }
 
-type ChatModelAgentConfig struct {
+// TypedChatModelAgentConfig is the generic configuration for ChatModelAgent.
+type TypedChatModelAgentConfig[M messageType] struct {
 	// Name of the agent. Better be unique across all agents.
 	// Optional. If empty, the agent can still run standalone but cannot be used as
 	// a sub-agent tool via NewAgentTool (which requires a non-empty Name).
@@ -227,13 +260,13 @@ type ChatModelAgentConfig struct {
 	// Model is the chat model used by the agent.
 	// If your ChatModelAgent uses any tools, this model must support the model.WithTools
 	// call option, as that's how ChatModelAgent configures the model with tool information.
-	Model model.BaseChatModel
+	Model model.BaseModel[M]
 
 	ToolsConfig ToolsConfig
 
 	// GenModelInput transforms instructions and input messages into the model's input format.
 	// Optional. Defaults to defaultGenModelInput which combines instruction and messages.
-	GenModelInput GenModelInput
+	GenModelInput TypedGenModelInput[M]
 
 	// Exit defines the tool used to terminate the agent process.
 	// Optional. If nil, no Exit Action will be generated.
@@ -347,7 +380,7 @@ type ChatModelAgentConfig struct {
 	//     passed to ChatModel, NOT the actual tools available for execution. Use this for
 	//     dynamic tool filtering/selection based on conversation context. The modification
 	//     is scoped to this model request only.
-	Handlers []ChatModelAgentMiddleware
+	Handlers []TypedChatModelAgentMiddleware[M]
 
 	// ModelRetryConfig configures retry behavior for the ChatModel.
 	// When set, the agent will automatically retry failed ChatModel calls
@@ -363,42 +396,52 @@ type ChatModelAgentConfig struct {
 	ModelFailoverConfig *ModelFailoverConfig
 }
 
-type ChatModelAgent struct {
+type ChatModelAgentConfig = TypedChatModelAgentConfig[*schema.Message]
+
+// TypedChatModelAgent is a chat model-backed agent parameterized by message type.
+//
+// For M = *schema.Message, the full ReAct loop (model → tool calls → model) is used.
+// For M = *schema.AgenticMessage, a single-shot chain is used since agentic models
+// handle tool calling internally. Cancel monitoring and retry on the model stream
+// are not yet supported for agentic models.
+type TypedChatModelAgent[M messageType] struct {
 	name        string
 	description string
 	instruction string
 
-	model       model.BaseChatModel
+	model       model.BaseModel[M]
 	toolsConfig ToolsConfig
 
-	genModelInput GenModelInput
+	genModelInput TypedGenModelInput[M]
 
 	outputKey     string
 	maxIterations int
 
-	subAgents   []Agent
-	parentAgent Agent
+	subAgents   []TypedAgent[M]
+	parentAgent TypedAgent[M]
 
 	disallowTransferToParent bool
 
 	exit tool.BaseTool
 
-	handlers    []ChatModelAgentMiddleware
+	handlers    []TypedChatModelAgentMiddleware[M]
 	middlewares []AgentMiddleware
 
 	modelRetryConfig    *ModelRetryConfig
 	modelFailoverConfig *ModelFailoverConfig
 
 	once   sync.Once
-	run    runFunc
+	run    typedRunFunc[M]
 	frozen uint32
 	exeCtx *execContext
 }
 
-// runParams holds the parameters for a runFunc invocation.
-type runParams struct {
-	input          *AgentInput
-	generator      *AsyncGenerator[*AgentEvent]
+type ChatModelAgent = TypedChatModelAgent[*schema.Message]
+
+// typedRunParams holds the parameters for a typedRunFunc invocation.
+type typedRunParams[M messageType] struct {
+	input          *TypedAgentInput[M]
+	generator      *AsyncGenerator[*TypedAgentEvent[M]]
 	store          *bridgeStore
 	instruction    string
 	returnDirectly map[string]bool
@@ -407,10 +450,15 @@ type runParams struct {
 	composeOpts    []compose.Option
 }
 
-type runFunc func(ctx context.Context, p *runParams)
+type typedRunFunc[M messageType] func(ctx context.Context, p *typedRunParams[M])
 
-// NewChatModelAgent constructs a chat model-backed agent with the provided config.
+// NewChatModelAgent creates a new ChatModelAgent with the given config.
 func NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*ChatModelAgent, error) {
+	return NewTypedChatModelAgent[*schema.Message](ctx, config)
+}
+
+// NewTypedChatModelAgent creates a new TypedChatModelAgent with the given config.
+func NewTypedChatModelAgent[M messageType](ctx context.Context, config *TypedChatModelAgentConfig[M]) (*TypedChatModelAgent[M], error) {
 	if config.ModelFailoverConfig != nil {
 		if config.ModelFailoverConfig.GetFailoverModel == nil {
 			return nil, errors.New("ModelFailoverConfig.GetFailoverModel is required when ModelFailoverConfig is set")
@@ -426,9 +474,11 @@ func NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*Chat
 		return nil, errors.New("agent 'Model' is required")
 	}
 
-	genInput := defaultGenModelInput
+	var genInput TypedGenModelInput[M]
 	if config.GenModelInput != nil {
 		genInput = config.GenModelInput
+	} else {
+		genInput = newDefaultGenModelInput[M]()
 	}
 
 	tc := config.ToolsConfig
@@ -455,7 +505,7 @@ func NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*Chat
 		EnhancedStreamable: cancelToolHandler.WrapEnhancedStreamableToolCall,
 	})
 
-	return &ChatModelAgent{
+	return &TypedChatModelAgent[M]{
 		name:                config.Name,
 		description:         config.Description,
 		instruction:         config.Instruction,
@@ -580,15 +630,15 @@ func (tta transferToAgent) InvokableRun(ctx context.Context, argumentsInJSON str
 	return transferToAgentToolOutput(params.AgentName), nil
 }
 
-func (a *ChatModelAgent) Name(_ context.Context) string {
+func (a *TypedChatModelAgent[M]) Name(_ context.Context) string {
 	return a.name
 }
 
-func (a *ChatModelAgent) Description(_ context.Context) string {
+func (a *TypedChatModelAgent[M]) Description(_ context.Context) string {
 	return a.description
 }
 
-func (a *ChatModelAgent) GetType() string {
+func (a *TypedChatModelAgent[M]) GetType() string {
 	return "ChatModel"
 }
 
@@ -597,7 +647,7 @@ func (a *ChatModelAgent) GetType() string {
 // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
 // to be more effective empirically. Consider using ChatModelAgent with AgentTool
 // or DeepAgent instead for most multi-agent scenarios.
-func (a *ChatModelAgent) OnSetSubAgents(_ context.Context, subAgents []Agent) error {
+func (a *TypedChatModelAgent[M]) OnSetSubAgents(_ context.Context, subAgents []TypedAgent[M]) error {
 	if atomic.LoadUint32(&a.frozen) == 1 {
 		return errors.New("agent has been frozen after run")
 	}
@@ -615,7 +665,7 @@ func (a *ChatModelAgent) OnSetSubAgents(_ context.Context, subAgents []Agent) er
 // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
 // to be more effective empirically. Consider using ChatModelAgent with AgentTool
 // or DeepAgent instead for most multi-agent scenarios.
-func (a *ChatModelAgent) OnSetAsSubAgent(_ context.Context, parent Agent) error {
+func (a *TypedChatModelAgent[M]) OnSetAsSubAgent(_ context.Context, parent TypedAgent[M]) error {
 	if atomic.LoadUint32(&a.frozen) == 1 {
 		return errors.New("agent has been frozen after run")
 	}
@@ -633,7 +683,7 @@ func (a *ChatModelAgent) OnSetAsSubAgent(_ context.Context, parent Agent) error
 // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
 // to be more effective empirically. Consider using ChatModelAgent with AgentTool
 // or DeepAgent instead for most multi-agent scenarios.
-func (a *ChatModelAgent) OnDisallowTransferToParent(_ context.Context) error {
+func (a *TypedChatModelAgent[M]) OnDisallowTransferToParent(_ context.Context) error {
 	if atomic.LoadUint32(&a.frozen) == 1 {
 		return errors.New("agent has been frozen after run")
 	}
@@ -652,24 +702,41 @@ func init() {
 	schema.RegisterName[*ChatModelAgentInterruptInfo]("_eino_adk_chat_model_agent_interrupt_info")
 }
 
-func setOutputToSession(ctx context.Context, msg Message, msgStream MessageStream, outputKey string) error {
-	if msg != nil {
-		AddSessionValue(ctx, outputKey, msg.Content)
+func extractTextContent[M messageType](msg M) string {
+	switch v := any(msg).(type) {
+	case *schema.Message:
+		return v.Content
+	case *schema.AgenticMessage:
+		var texts []string
+		for _, block := range v.ContentBlocks {
+			if block != nil && block.Type == schema.ContentBlockTypeAssistantGenText && block.AssistantGenText != nil {
+				texts = append(texts, block.AssistantGenText.Text)
+			}
+		}
+		return strings.Join(texts, "\n")
+	default:
+		return ""
+	}
+}
+
+func setOutputToSession[M messageType](ctx context.Context, msg M, msgStream *schema.StreamReader[M], outputKey string) error {
+	if !isNilMessage(msg) {
+		AddSessionValue(ctx, outputKey, extractTextContent(msg))
 		return nil
 	}
 
-	concatenated, err := schema.ConcatMessageStream(msgStream)
+	concatenated, err := concatMessageStream(msgStream)
 	if err != nil {
 		return err
 	}
 
-	AddSessionValue(ctx, outputKey, concatenated.Content)
+	AddSessionValue(ctx, outputKey, extractTextContent(concatenated))
 	return nil
 }
 
-func errFunc(err error) runFunc {
-	return func(ctx context.Context, p *runParams) {
-		p.generator.Send(&AgentEvent{Err: err})
+func typedErrFunc[M messageType](err error) typedRunFunc[M] {
+	return func(ctx context.Context, p *typedRunParams[M]) {
+		p.generator.Send(&TypedAgentEvent[M]{Err: err})
 	}
 }
 
@@ -693,7 +760,7 @@ type execContext struct {
 	toolUpdated  bool // whether needs to pass a compose.WithToolList option to ToolsNode due to tool list change
 }
 
-func (a *ChatModelAgent) applyBeforeAgent(ctx context.Context, ec *execContext) (context.Context, *execContext, error) {
+func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execContext) (context.Context, *execContext, error) {
 	runCtx := &ChatModelAgentContext{
 		Instruction:    ec.instruction,
 		Tools:          cloneSlice(ec.unwrappedTools),
@@ -731,7 +798,7 @@ func (a *ChatModelAgent) applyBeforeAgent(ctx context.Context, ec *execContext)
 	return ctx, runtimeEC, nil
 }
 
-func (a *ChatModelAgent) prepareExecContext(ctx context.Context) (*execContext, error) {
+func (a *TypedChatModelAgent[M]) prepareExecContext(ctx context.Context) (*execContext, error) {
 	instruction := a.instruction
 	toolsNodeConf := a.toolsConfig.ToolsNodeConfig
 	toolsNodeConf.Tools = cloneSlice(a.toolsConfig.Tools)
@@ -790,39 +857,39 @@ func (a *ChatModelAgent) prepareExecContext(ctx context.Context) (*execContext,
 // handleRunFuncError is the common error handler for buildNoToolsRunFunc and buildReActRunFunc.
 // It handles compose interrupts (both cancel-triggered and business)
 // and generic errors, sending the appropriate event to the generator.
-func (a *ChatModelAgent) handleRunFuncError(
+func (a *TypedChatModelAgent[M]) handleRunFuncError(
 	ctx context.Context,
 	err error,
 	cancelCtx *cancelContext,
 	cancelCtxOwned bool,
 	store *bridgeStore,
-	generator *AsyncGenerator[*AgentEvent],
+	generator *AsyncGenerator[*TypedAgentEvent[M]],
 ) {
 	info, ok := compose.ExtractInterruptInfo(err)
 	if ok {
 		if cancelCtx != nil {
-			// Note: there is a benign TOCTOU window here. Between shouldCancel()
-			// returning false and markDone() executing, a concurrent cancel could
-			// transition stateRunning→stateCancelling. markDone() then does
-			// stateCancelling→stateDone, and the cancel func receives
-			// ErrExecutionEnded (execution finished before cancel took effect).
 			if !cancelCtx.shouldCancel() {
+				// Note: there is a benign TOCTOU window here. Between shouldCancel()
+				// returning false and markDone() executing, a concurrent cancel could
+				// transition stateRunning→stateCancelling. markDone() then does
+				// stateCancelling→stateDone, and the cancel func receives
+				// ErrExecutionEnded (execution finished before cancel took effect).
 				cancelCtx.markDone()
 			}
 		}
 
 		data, existed, sErr := store.Get(ctx, bridgeCheckpointID)
 		if sErr != nil {
-			generator.Send(&AgentEvent{AgentName: a.name, Err: fmt.Errorf("failed to get interrupt info: %w", sErr)})
+			generator.Send(&TypedAgentEvent[M]{AgentName: a.name, Err: fmt.Errorf("failed to get interrupt info: %w", sErr)})
 			return
 		}
 		if !existed {
-			generator.Send(&AgentEvent{AgentName: a.name, Err: fmt.Errorf("interrupt occurred but checkpoint data is missing")})
+			generator.Send(&TypedAgentEvent[M]{AgentName: a.name, Err: fmt.Errorf("interrupt occurred but checkpoint data is missing")})
 			return
 		}
 
 		is := FromInterruptContexts(info.InterruptContexts)
-		event := CompositeInterrupt(ctx, info, data, is)
+		event := TypedCompositeInterrupt[M](ctx, info, data, is)
 		event.Action.Interrupted.Data = &ChatModelAgentInterruptInfo{
 			Info: info,
 			Data: data,
@@ -835,20 +902,30 @@ func (a *ChatModelAgent) handleRunFuncError(
 	if cancelCtxOwned && cancelCtx != nil {
 		cancelCtx.markDone()
 	}
-	generator.Send(&AgentEvent{Err: err})
+	generator.Send(&TypedAgentEvent[M]{Err: err})
 }
 
-func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
-	type noToolsInput struct {
-		input       *AgentInput
-		instruction string
-	}
+type typedNoToolsInput[M messageType] struct {
+	input       *TypedAgentInput[M]
+	instruction string
+}
 
-	return func(ctx context.Context, p *runParams) {
+func appendModelToChain[I, O any, M messageType](chain *compose.Chain[I, O], m model.BaseModel[M]) {
+	var zero M
+	switch any(zero).(type) {
+	case *schema.Message:
+		chain.AppendChatModel(any(m).(model.BaseChatModel))
+	case *schema.AgenticMessage:
+		chain.AppendAgenticModel(any(m).(model.AgenticModel))
+	}
+}
+
+func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRunFunc[M], error) {
+	return func(ctx context.Context, p *typedRunParams[M]) {
 		cancelCtx := p.cancelCtx
 		ctx = withCancelContext(ctx, cancelCtx)
 
-		wrappedModel := buildModelWrappers(a.model, &modelWrapperConfig{
+		wrappedModel := buildModelWrappers(a.model, &typedModelWrapperConfig[M]{
 			handlers:       a.handlers,
 			middlewares:    a.middlewares,
 			retryConfig:    a.modelRetryConfig,
@@ -856,22 +933,26 @@ func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
 			cancelContext:  cancelCtx,
 		})
 
-		chain := compose.NewChain[noToolsInput, Message](
-			compose.WithGenLocalState(func(ctx context.Context) (state *State) {
-				return &State{}
-			})).
-			AppendLambda(compose.InvokableLambda(func(ctx context.Context, in noToolsInput) ([]Message, error) {
-				messages, err := a.genModelInput(ctx, in.instruction, in.input)
-				if err != nil {
-					return nil, err
-				}
-				_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
-					st.Messages = append(st.Messages, messages...)
-					return nil
-				})
-				return messages, nil
-			})).
-			AppendChatModel(wrappedModel)
+		chain := compose.NewChain[typedNoToolsInput[M], M](
+			compose.WithGenLocalState(func(ctx context.Context) (state *typedState[M]) {
+				return &typedState[M]{}
+			}))
+
+		chain.AppendLambda(compose.InvokableLambda(func(ctx context.Context, in typedNoToolsInput[M]) ([]M, error) {
+			messages, err := a.genModelInput(ctx, in.instruction, in.input)
+			if err != nil {
+				return nil, err
+			}
+			if err := compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
+				st.Messages = append(st.Messages, messages...)
+				return nil
+			}); err != nil {
+				return nil, err
+			}
+			return messages, nil
+		}))
+
+		appendModelToChain(chain, wrappedModel)
 
 		var compileOptions []compose.GraphCompileOption
 		compileOptions = append(compileOptions,
@@ -887,11 +968,11 @@ func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
 
 		r, err := chain.Compile(ctx, compileOptions...)
 		if err != nil {
-			p.generator.Send(&AgentEvent{Err: err})
+			p.generator.Send(&TypedAgentEvent[M]{Err: err})
 			return
 		}
 
-		ctx = withChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{
+		ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[M]{
 			generator:                p.generator,
 			cancelCtx:                cancelCtx,
 			failoverLastSuccessModel: a.model,
@@ -904,15 +985,15 @@ func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
 				if !ok {
 					return
 				}
-				p.generator.Send(&AgentEvent{Err: cancelErr})
+				p.generator.Send(&TypedAgentEvent[M]{Err: cancelErr})
 				return
 			}
 		}
 
-		in := noToolsInput{input: p.input, instruction: p.instruction}
+		in := typedNoToolsInput[M]{input: p.input, instruction: p.instruction}
 
-		var msg Message
-		var msgStream MessageStream
+		var msg M
+		var msgStream *schema.StreamReader[M]
 		if p.input.EnableStreaming {
 			msgStream, err = r.Stream(ctx, in, p.composeOpts...)
 		} else {
@@ -923,7 +1004,7 @@ func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
 			if a.outputKey != "" {
 				err = setOutputToSession(ctx, msg, msgStream, a.outputKey)
 				if err != nil {
-					p.generator.Send(&AgentEvent{Err: err})
+					p.generator.Send(&TypedAgentEvent[M]{Err: err})
 				}
 			} else if msgStream != nil {
 				msgStream.Close()
@@ -932,15 +1013,37 @@ func (a *ChatModelAgent) buildNoToolsRunFunc(_ context.Context) runFunc {
 		}
 
 		a.handleRunFuncError(ctx, err, cancelCtx, p.cancelCtxOwned, p.store, p.generator)
+	}, nil
+}
+
+func (a *TypedChatModelAgent[M]) buildReActRunFunc(ctx context.Context, bc *execContext) (typedRunFunc[M], error) {
+	var zero M
+	switch any(zero).(type) {
+	case *schema.Message:
+		return a.buildMessageReActRunFunc(ctx, bc)
+	case *schema.AgenticMessage:
+		// single-shot: agentic models handle tool calling internally
+		return a.buildAgenticReActRunFunc(ctx, bc)
+	default:
+		return nil, fmt.Errorf("unsupported message type %T for ReAct run mode", zero)
 	}
 }
 
-func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (runFunc, error) {
-	conf := &reactConfig{
-		model:       a.model,
+type reactRunInput struct {
+	input       *AgentInput
+	instruction string
+}
+
+func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(ctx context.Context, bc *execContext) (typedRunFunc[M], error) {
+	// safe: only called when M = *schema.Message (guarded by type switch in buildReActRunFunc)
+	msgModel := any(a.model).(model.BaseChatModel)
+	msgHandlers := any(a.handlers).([]ChatModelAgentMiddleware)
+	genModelInputFn := any(a.genModelInput).(GenModelInput)
+	msgConf := &reactConfig{
+		model:       msgModel,
 		toolsConfig: &bc.toolsNodeConf,
 		modelWrapperConf: &modelWrapperConfig{
-			handlers:       a.handlers,
+			handlers:       msgHandlers,
 			middlewares:    a.middlewares,
 			retryConfig:    a.modelRetryConfig,
 			failoverConfig: a.modelFailoverConfig,
@@ -951,29 +1054,25 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 		maxIterations:       a.maxIterations,
 	}
 
-	type reactRunInput struct {
-		input       *AgentInput
-		instruction string
-	}
-
-	return func(ctx context.Context, p *runParams) {
-		cancelCtx := p.cancelCtx
-		conf.cancelCtx = cancelCtx
-		if conf.modelWrapperConf != nil {
-			conf.modelWrapperConf.cancelContext = cancelCtx
+	return func(ctx context.Context, p *typedRunParams[M]) {
+		mp := any(p).(*typedRunParams[*schema.Message])
+		cancelCtx := mp.cancelCtx
+		msgConf.cancelCtx = cancelCtx
+		if msgConf.modelWrapperConf != nil {
+			msgConf.modelWrapperConf.cancelContext = cancelCtx
 		}
 		ctx = withCancelContext(ctx, cancelCtx)
 
-		g, err := newReact(ctx, conf)
+		g, err := newReact(ctx, msgConf)
 		if err != nil {
-			p.generator.Send(&AgentEvent{Err: err})
+			mp.generator.Send(&AgentEvent{Err: err})
 			return
 		}
 
 		chain := compose.NewChain[reactRunInput, Message]().
 			AppendLambda(
 				compose.InvokableLambda(func(ctx context.Context, in reactRunInput) (*reactInput, error) {
-					messages, genErr := a.genModelInput(ctx, in.instruction, in.input)
+					messages, genErr := genModelInputFn(ctx, in.instruction, in.input)
 					if genErr != nil {
 						return nil, genErr
 					}
@@ -987,7 +1086,7 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 		var compileOptions []compose.GraphCompileOption
 		compileOptions = append(compileOptions,
 			compose.WithGraphName(a.name),
-			compose.WithCheckPointStore(p.store),
+			compose.WithCheckPointStore(mp.store),
 			compose.WithSerializer(&gobSerializer{}),
 			compose.WithMaxRunSteps(math.MaxInt))
 
@@ -999,15 +1098,15 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 
 		runnable, err_ := chain.Compile(ctx, compileOptions...)
 		if err_ != nil {
-			p.generator.Send(&AgentEvent{Err: err_})
+			mp.generator.Send(&AgentEvent{Err: err_})
 			return
 		}
 
-		ctx = withChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{
-			runtimeReturnDirectly:    p.returnDirectly,
-			generator:                p.generator,
+		ctx = withTypedChatModelAgentExecCtx[*schema.Message](ctx, &chatModelAgentExecCtx{
+			runtimeReturnDirectly:    mp.returnDirectly,
+			generator:                mp.generator,
 			cancelCtx:                cancelCtx,
-			failoverLastSuccessModel: a.model,
+			failoverLastSuccessModel: msgModel,
 		})
 
 		// Pre-execution cancel check
@@ -1017,28 +1116,149 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 				if !ok {
 					return
 				}
-				p.generator.Send(&AgentEvent{Err: cancelErr})
+				mp.generator.Send(&AgentEvent{Err: cancelErr})
 				return
 			}
 		}
 
 		in := reactRunInput{
-			input:       p.input,
-			instruction: p.instruction,
+			input:       mp.input,
+			instruction: mp.instruction,
 		}
 
 		var runOpts []compose.Option
-		runOpts = append(runOpts, p.composeOpts...)
+		runOpts = append(runOpts, mp.composeOpts...)
 		if a.toolsConfig.EmitInternalEvents {
-			runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEventGenerator(p.generator))))
+			runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEventGenerator(mp.generator))))
 		}
-		if p.input.EnableStreaming {
+		if mp.input.EnableStreaming {
 			runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEnableStreaming(true))))
 		}
 
 		var msg Message
 		var msgStream MessageStream
-		if p.input.EnableStreaming {
+		if mp.input.EnableStreaming {
+			msgStream, err_ = runnable.Stream(ctx, in, runOpts...)
+		} else {
+			msg, err_ = runnable.Invoke(ctx, in, runOpts...)
+		}
+
+		if err_ == nil {
+			if a.outputKey != "" {
+				err_ = setOutputToSession[*schema.Message](ctx, msg, msgStream, a.outputKey)
+				if err_ != nil {
+					mp.generator.Send(&AgentEvent{Err: err_})
+				}
+			} else if msgStream != nil {
+				msgStream.Close()
+			}
+
+			return
+		}
+
+		a.handleRunFuncError(ctx, err_, cancelCtx, mp.cancelCtxOwned, mp.store, p.generator)
+	}, nil
+}
+
+type agenticReactRunInput struct {
+	input       *TypedAgentInput[*schema.AgenticMessage]
+	instruction string
+}
+
+func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(ctx context.Context, bc *execContext) (typedRunFunc[M], error) {
+	agenticModel := any(a.model).(model.AgenticModel)
+	agenticHandlers := any(a.handlers).([]TypedChatModelAgentMiddleware[*schema.AgenticMessage])
+	genModelInputFn := any(a.genModelInput).(TypedGenModelInput[*schema.AgenticMessage])
+	agenticConf := &agenticReactConfig{
+		model:       agenticModel,
+		toolsConfig: &bc.toolsNodeConf,
+		modelWrapperConf: &typedModelWrapperConfig[*schema.AgenticMessage]{
+			handlers:    agenticHandlers,
+			middlewares: a.middlewares,
+			retryConfig: a.modelRetryConfig,
+			toolInfos:   bc.toolInfos,
+		},
+		toolsReturnDirectly: bc.returnDirectly,
+		agentName:           a.name,
+		maxIterations:       a.maxIterations,
+	}
+
+	return func(ctx context.Context, p *typedRunParams[M]) {
+		ap := any(p).(*typedRunParams[*schema.AgenticMessage])
+		cancelCtx := ap.cancelCtx
+		agenticConf.cancelCtx = cancelCtx
+		if agenticConf.modelWrapperConf != nil {
+			agenticConf.modelWrapperConf.cancelContext = cancelCtx
+		}
+		ctx = withCancelContext(ctx, cancelCtx)
+
+		g, err := newAgenticReact(ctx, agenticConf)
+		if err != nil {
+			ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err})
+			return
+		}
+
+		chain := compose.NewChain[agenticReactRunInput, *schema.AgenticMessage]().
+			AppendLambda(
+				compose.InvokableLambda(func(ctx context.Context, in agenticReactRunInput) (*agenticReactInput, error) {
+					messages, genErr := genModelInputFn(ctx, in.instruction, in.input)
+					if genErr != nil {
+						return nil, genErr
+					}
+					return &agenticReactInput{
+						Messages: messages,
+					}, nil
+				}),
+			).
+			AppendGraph(g, compose.WithNodeName("ReAct"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt)))
+
+		var compileOptions []compose.GraphCompileOption
+		compileOptions = append(compileOptions,
+			compose.WithGraphName(a.name),
+			compose.WithCheckPointStore(ap.store),
+			compose.WithSerializer(&gobSerializer{}),
+			compose.WithMaxRunSteps(math.MaxInt))
+
+		if cancelCtx != nil {
+			var interrupt func(...compose.GraphInterruptOption)
+			ctx, interrupt = compose.WithGraphInterrupt(ctx)
+			cancelCtx.setGraphInterruptFunc(cancelCtx.wrapGraphInterruptWithGracePeriod(interrupt))
+		}
+
+		runnable, err_ := chain.Compile(ctx, compileOptions...)
+		if err_ != nil {
+			ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err_})
+			return
+		}
+
+		ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[*schema.AgenticMessage]{
+			runtimeReturnDirectly: ap.returnDirectly,
+			generator:             ap.generator,
+		})
+
+		// Pre-execution cancel check
+		if cancelCtx != nil && cancelCtx.shouldCancel() {
+			if cancelCtx.getMode() == CancelImmediate || atomic.LoadInt32(&cancelCtx.escalated) == 1 {
+				cancelErr, ok := cancelCtx.createAndMarkCancelHandled()
+				if !ok {
+					return
+				}
+				ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: cancelErr})
+				return
+			}
+		}
+
+		in := agenticReactRunInput{input: ap.input, instruction: ap.instruction}
+
+		var runOpts []compose.Option
+		runOpts = append(runOpts, ap.composeOpts...)
+		if ap.input.EnableStreaming {
+			runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEnableStreaming(true))))
+		}
+
+		var msg *schema.AgenticMessage
+		var msgStream *schema.StreamReader[*schema.AgenticMessage]
+		if ap.input.EnableStreaming {
 			msgStream, err_ = runnable.Stream(ctx, in, runOpts...)
 		} else {
 			msg, err_ = runnable.Invoke(ctx, in, runOpts...)
@@ -1048,7 +1268,7 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 			if a.outputKey != "" {
 				err_ = setOutputToSession(ctx, msg, msgStream, a.outputKey)
 				if err_ != nil {
-					p.generator.Send(&AgentEvent{Err: err_})
+					ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err_})
 				}
 			} else if msgStream != nil {
 				msgStream.Close()
@@ -1057,28 +1277,35 @@ func (a *ChatModelAgent) buildReActRunFunc(_ context.Context, bc *execContext) (
 			return
 		}
 
-		a.handleRunFuncError(ctx, err_, cancelCtx, p.cancelCtxOwned, p.store, p.generator)
+		a.handleRunFuncError(ctx, err_, cancelCtx, ap.cancelCtxOwned, ap.store, p.generator)
 	}, nil
 }
 
-func (a *ChatModelAgent) buildRunFunc(ctx context.Context) runFunc {
+func (a *TypedChatModelAgent[M]) buildRunFunc(ctx context.Context) typedRunFunc[M] {
 	a.once.Do(func() {
 		ec, err := a.prepareExecContext(ctx)
 		if err != nil {
-			a.run = errFunc(err)
+			a.run = typedErrFunc[M](err)
 			return
 		}
 
 		a.exeCtx = ec
 
 		if len(ec.toolsNodeConf.Tools) == 0 {
-			a.run = a.buildNoToolsRunFunc(ctx)
+			var run typedRunFunc[M]
+			run, err = a.buildNoToolsRunFunc(ctx)
+			if err != nil {
+				a.run = typedErrFunc[M](err)
+				return
+			}
+			a.run = run
 			return
 		}
 
-		run, err := a.buildReActRunFunc(ctx, ec)
+		var run typedRunFunc[M]
+		run, err = a.buildReActRunFunc(ctx, ec)
 		if err != nil {
-			a.run = errFunc(err)
+			a.run = typedErrFunc[M](err)
 			return
 		}
 		a.run = run
@@ -1089,7 +1316,7 @@ func (a *ChatModelAgent) buildRunFunc(ctx context.Context) runFunc {
 	return a.run
 }
 
-func (a *ChatModelAgent) getRunFunc(ctx context.Context) (context.Context, runFunc, *execContext, error) {
+func (a *TypedChatModelAgent[M]) getRunFunc(ctx context.Context) (context.Context, typedRunFunc[M], *execContext, error) {
 	defaultRun := a.buildRunFunc(ctx)
 	bc := a.exeCtx
 
@@ -1116,9 +1343,12 @@ func (a *ChatModelAgent) getRunFunc(ctx context.Context) (context.Context, runFu
 		return ctx, defaultRun, runtimeBC, nil
 	}
 
-	var tempRun runFunc
+	var tempRun typedRunFunc[M]
 	if len(runtimeBC.toolsNodeConf.Tools) == 0 {
-		tempRun = a.buildNoToolsRunFunc(ctx)
+		tempRun, err = a.buildNoToolsRunFunc(ctx)
+		if err != nil {
+			return ctx, nil, nil, err
+		}
 	} else {
 		tempRun, err = a.buildReActRunFunc(ctx, runtimeBC)
 		if err != nil {
@@ -1129,8 +1359,8 @@ func (a *ChatModelAgent) getRunFunc(ctx context.Context) (context.Context, runFu
 	return ctx, tempRun, runtimeBC, nil
 }
 
-func (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
-	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
+func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
 
 	o := getCommonOptions(nil, opts...)
 	cancelCtx := o.cancelCtx
@@ -1145,7 +1375,7 @@ func (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...Age
 			if cancelCtxOwned && cancelCtx != nil {
 				defer cancelCtx.markDone()
 			}
-			generator.Send(&AgentEvent{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)})
+			generator.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)})
 			generator.Close()
 		}()
 		return iterator
@@ -1166,7 +1396,7 @@ func (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...Age
 			panicErr := recover()
 			if panicErr != nil {
 				e := safe.NewPanicErr(panicErr, debug.Stack())
-				generator.Send(&AgentEvent{Err: e})
+				generator.Send(&TypedAgentEvent[M]{Err: e})
 			}
 
 			generator.Close()
@@ -1182,7 +1412,7 @@ func (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...Age
 			returnDirectly = bc.returnDirectly
 		}
 
-		run(ctx, &runParams{
+		run(ctx, &typedRunParams[M]{
 			input:          input,
 			generator:      generator,
 			store:          newBridgeStore(),
@@ -1200,8 +1430,8 @@ func (a *ChatModelAgent) Run(ctx context.Context, input *AgentInput, opts ...Age
 	return iterator
 }
 
-func (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
-	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
+func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
 
 	o := getCommonOptions(nil, opts...)
 	cancelCtx := o.cancelCtx
@@ -1216,7 +1446,7 @@ func (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...A
 			if cancelCtxOwned && cancelCtx != nil {
 				defer cancelCtx.markDone()
 			}
-			generator.Send(&AgentEvent{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)})
+			generator.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)})
 			generator.Close()
 		}()
 		return iterator
@@ -1255,7 +1485,7 @@ func (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...A
 	stateByte, err = preprocessComposeCheckpoint(stateByte)
 	if err != nil {
 		go func() {
-			generator.Send(&AgentEvent{Err: err})
+			generator.Send(&TypedAgentEvent[M]{Err: err})
 			generator.Close()
 		}()
 		return iterator
@@ -1287,7 +1517,7 @@ func (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...A
 			panicErr := recover()
 			if panicErr != nil {
 				e := safe.NewPanicErr(panicErr, debug.Stack())
-				generator.Send(&AgentEvent{Err: e})
+				generator.Send(&TypedAgentEvent[M]{Err: e})
 			}
 
 			generator.Close()
@@ -1303,8 +1533,8 @@ func (a *ChatModelAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...A
 			returnDirectly = bc.returnDirectly
 		}
 
-		run(ctx, &runParams{
-			input:          &AgentInput{EnableStreaming: info.EnableStreaming},
+		run(ctx, &typedRunParams[M]{
+			input:          &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming},
 			generator:      generator,
 			store:          newResumeBridgeStore(bridgeCheckpointID, stateByte),
 			instruction:    instruction,

adk/deterministic_transfer.go

@@ -250,7 +250,7 @@ func handleFlowAgentEvents(ctx context.Context, iter *AsyncIterator[*AgentEvent]
 		}
 
 		if parentSession != nil && (event.Action == nil || event.Action.Interrupted == nil) {
-			copied := copyAgentEvent(event)
+			copied := copyTypedAgentEvent(event)
 			setAutomaticClose(copied)
 			setAutomaticClose(event)
 			parentSession.addEvent(copied)

adk/failover_chatmodel.go

@@ -31,21 +31,13 @@ import (
 
 type failoverCurrentModelKey struct{}
 
-type failoverCurrentModel struct {
-	model model.BaseChatModel
+func typedSetFailoverCurrentModel[M messageType](ctx context.Context, currentModel model.BaseModel[M]) context.Context {
+	return context.WithValue(ctx, failoverCurrentModelKey{}, currentModel)
 }
 
-func setFailoverCurrentModel(ctx context.Context, currentModel model.BaseChatModel) context.Context {
-	return context.WithValue(ctx, failoverCurrentModelKey{}, &failoverCurrentModel{
-		model: currentModel,
-	})
-}
-
-func getFailoverCurrentModel(ctx context.Context) *failoverCurrentModel {
-	if fm, ok := ctx.Value(failoverCurrentModelKey{}).(*failoverCurrentModel); ok {
-		return fm
-	}
-	return nil
+func typedGetFailoverCurrentModel[M messageType](ctx context.Context) (model.BaseModel[M], bool) {
+	m, ok := ctx.Value(failoverCurrentModelKey{}).(model.BaseModel[M])
+	return m, ok
 }
 
 type failoverHasMoreAttemptsKey struct{}
@@ -64,30 +56,30 @@ func getFailoverHasMoreAttempts(ctx context.Context) bool {
 	return v
 }
 
-type failoverProxyModel struct {
+type typedFailoverProxyModel[M messageType] struct {
 }
 
-func (m *failoverProxyModel) prepareCallbacks(ctx context.Context) (context.Context, model.BaseChatModel, error) {
-	current := getFailoverCurrentModel(ctx)
-	if current == nil || current.model == nil {
+func (m *typedFailoverProxyModel[M]) prepareCallbacks(ctx context.Context) (context.Context, model.BaseModel[M], error) {
+	target, ok := typedGetFailoverCurrentModel[M](ctx)
+	if !ok {
 		return nil, nil, errors.New("failover current model not found in context")
 	}
 
-	typ, _ := components.GetType(current.model)
+	typ, _ := components.GetType(target)
 	ctx = callbacks.EnsureRunInfo(ctx, typ, components.ComponentOfChatModel)
 
-	target := current.model
 	if !components.IsCallbacksEnabled(target) {
-		target = (&callbackInjectionModelWrapper{}).WrapModel(target)
+		target = typedCallbackInjectionModelWrapper[M]{}.wrapModel(target)
 	}
 
 	return ctx, target, nil
 }
 
-func (m *failoverProxyModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (m *typedFailoverProxyModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	nCtx, target, err := m.prepareCallbacks(ctx)
 	if err != nil {
-		return nil, err
+		var zero M
+		return zero, err
 	}
 
 	ctx = callbacks.OnStart(ctx, input)
@@ -103,7 +95,7 @@ func (m *failoverProxyModel) Generate(ctx context.Context, input []*schema.Messa
 	return result, nil
 }
 
-func (m *failoverProxyModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+func (m *typedFailoverProxyModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 	nCtx, target, err := m.prepareCallbacks(ctx)
 	if err != nil {
 		return nil, err
@@ -121,14 +113,16 @@ func (m *failoverProxyModel) Stream(ctx context.Context, input []*schema.Message
 	return wrappedStream, nil
 }
 
-func (m *failoverProxyModel) IsCallbacksEnabled() bool {
+func (m *typedFailoverProxyModel[M]) IsCallbacksEnabled() bool {
 	return true
 }
 
-func (m *failoverProxyModel) GetType() string {
+func (m *typedFailoverProxyModel[M]) GetType() string {
 	return "FailoverProxyModel"
 }
 
+type failoverProxyModel = typedFailoverProxyModel[*schema.Message]
+
 // FailoverContext contains context information during failover process.
 type FailoverContext struct {
 	// FailoverAttempt is the current failover attempt number, starting from 1.
@@ -199,32 +193,35 @@ type ModelFailoverConfig struct {
 		failoverModel model.BaseChatModel, failoverModelInputMessages []*schema.Message, failoverErr error)
 }
 
-func getLastSuccessModel(ctx context.Context) model.BaseChatModel {
-	if execCtx := getChatModelAgentExecCtx(ctx); execCtx != nil {
-		return execCtx.failoverLastSuccessModel
+func typedGetFailoverLastSuccessModel[M messageType](ctx context.Context) model.BaseModel[M] {
+	execCtx := getTypedChatModelAgentExecCtx[M](ctx)
+	if execCtx == nil {
+		return nil
 	}
-	return nil
+	return execCtx.failoverLastSuccessModel
 }
 
-func setLastSuccessModel(ctx context.Context, m model.BaseChatModel) {
-	if execCtx := getChatModelAgentExecCtx(ctx); execCtx != nil {
+func typedSetFailoverLastSuccessModel[M messageType](ctx context.Context, m model.BaseModel[M]) {
+	if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil {
 		execCtx.failoverLastSuccessModel = m
 	}
 }
 
-type failoverModelWrapper struct {
+type typedFailoverModelWrapper[M messageType] struct {
 	config *ModelFailoverConfig
-	inner  model.BaseChatModel
+	inner  model.BaseModel[M]
 }
 
-func newFailoverModelWrapper(inner model.BaseChatModel, config *ModelFailoverConfig) *failoverModelWrapper {
-	return &failoverModelWrapper{
+type failoverModelWrapper = typedFailoverModelWrapper[*schema.Message]
+
+func newTypedFailoverModelWrapper[M messageType](inner model.BaseModel[M], config *ModelFailoverConfig) *typedFailoverModelWrapper[M] {
+	return &typedFailoverModelWrapper[M]{
 		config: config,
 		inner:  inner,
 	}
 }
 
-func (f *failoverModelWrapper) needFailover(ctx context.Context, outputMessage *schema.Message, outputErr error) bool {
+func (f *typedFailoverModelWrapper[M]) needFailover(ctx context.Context, outputMessage M, outputErr error) bool {
 	if ctx.Err() != nil {
 		return false
 	}
@@ -236,25 +233,51 @@ func (f *failoverModelWrapper) needFailover(ctx context.Context, outputMessage *
 	}
 
 	// ShouldFailover is validated at agent construction; nil here indicates a programmer error.
-	return f.config.ShouldFailover(ctx, outputMessage, outputErr)
+	schemaMsg, _ := any(outputMessage).(*schema.Message)
+	return f.config.ShouldFailover(ctx, schemaMsg, outputErr)
 }
 
-func (f *failoverModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (f *typedFailoverModelWrapper[M]) getFailoverModel(ctx context.Context, failoverCtx *FailoverContext) (model.BaseModel[M], []M, error) {
+	chatModel, msgs, err := f.config.GetFailoverModel(ctx, failoverCtx)
+	if err != nil {
+		return nil, nil, err
+	}
+	if chatModel == nil {
+		return nil, nil, nil
+	}
+
+	typedModel, ok := any(chatModel).(model.BaseModel[M])
+	if !ok {
+		return nil, nil, fmt.Errorf("failover GetFailoverModel returned model of type %T, expected model.BaseModel[%T]", chatModel, *new(M))
+	}
+
+	var typedMsgs []M
+	if msgs != nil {
+		if m, ok := any(msgs).([]M); ok {
+			typedMsgs = m
+		}
+	}
+
+	return typedModel, typedMsgs, nil
+}
+
+func (f *typedFailoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	// Defensive: GetFailoverModel is validated non-nil at agent construction.
 	if f.config.GetFailoverModel == nil {
 		return f.inner.Generate(ctx, input, opts...)
 	}
 
-	var lastOutputMessage *schema.Message
+	var lastOutputMessage M
 	var lastErr error
 
 	// Try lastSuccessModel first if available.
-	if lastSuccess := getLastSuccessModel(ctx); lastSuccess != nil {
+	if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil {
 		if err := ctx.Err(); err != nil {
-			return nil, err
+			var zero M
+			return zero, err
 		}
 
-		modelCtx := setFailoverCurrentModel(ctx, lastSuccess)
+		modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess)
 		modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0)
 		result, err := f.inner.Generate(modelCtx, input, opts...)
 		if err == nil {
@@ -273,36 +296,41 @@ func (f *failoverModelWrapper) Generate(ctx context.Context, input []*schema.Mes
 
 	for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ {
 		if err := ctx.Err(); err != nil {
-			return nil, err
+			var zero M
+			return zero, err
 		}
 
+		inputMsgs, _ := any(input).([]*schema.Message)
+		lastOutputMsg, _ := any(lastOutputMessage).(*schema.Message)
 		failoverCtx := &FailoverContext{
 			FailoverAttempt:   attempt,
-			InputMessages:     input,
-			LastOutputMessage: lastOutputMessage,
+			InputMessages:     inputMsgs,
+			LastOutputMessage: lastOutputMsg,
 			LastErr:           lastErr,
 		}
 
-		currentModel, currentInput, err := f.config.GetFailoverModel(ctx, failoverCtx)
+		currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx)
 		if err != nil {
-			return nil, err
+			var zero M
+			return zero, err
 		}
 		if currentModel == nil {
-			return nil, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt)
+			var zero M
+			return zero, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt)
 		}
 
 		if currentInput == nil {
 			currentInput = input
 		}
 
-		modelCtx := setFailoverCurrentModel(ctx, currentModel)
+		modelCtx := typedSetFailoverCurrentModel(ctx, currentModel)
 		modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries)
 		result, err := f.inner.Generate(modelCtx, currentInput, opts...)
 		lastOutputMessage = result
 		lastErr = err
 
 		if err == nil {
-			setLastSuccessModel(ctx, currentModel)
+			typedSetFailoverLastSuccessModel[M](ctx, currentModel)
 			return result, nil
 		}
 
@@ -318,28 +346,29 @@ func (f *failoverModelWrapper) Generate(ctx context.Context, input []*schema.Mes
 	return lastOutputMessage, lastErr
 }
 
-func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (
-	*schema.StreamReader[*schema.Message], error) {
+func (f *typedFailoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (
+	*schema.StreamReader[M], error) {
 	// Defensive: GetFailoverModel is validated non-nil at agent construction.
 	if f.config.GetFailoverModel == nil {
 		return f.inner.Stream(ctx, input, opts...)
 	}
 
-	var lastOutputMessage *schema.Message
+	var lastOutputMessage M
 	var lastErr error
 
 	// Try lastSuccessModel first if available.
-	if lastSuccess := getLastSuccessModel(ctx); lastSuccess != nil {
+	if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil {
 		if err := ctx.Err(); err != nil {
 			return nil, err
 		}
 
-		modelCtx := setFailoverCurrentModel(ctx, lastSuccess)
+		modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess)
 		modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0)
 		stream, err := f.inner.Stream(modelCtx, input, opts...)
 		if err != nil {
 			lastErr = err
-			if !f.needFailover(ctx, nil, err) {
+			var zero M
+			if !f.needFailover(ctx, zero, err) {
 				return nil, err
 			}
 			log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", err)
@@ -348,7 +377,7 @@ func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Messa
 			checkCopy := copies[0]
 			returnCopy := copies[1]
 
-			outMsg, streamErr := consumeStream(checkCopy)
+			outMsg, streamErr := typedConsumeStream(checkCopy)
 			if streamErr != nil {
 				lastOutputMessage = outMsg
 				lastErr = streamErr
@@ -369,14 +398,16 @@ func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Messa
 			return nil, err
 		}
 
+		inputMsgs2, _ := any(input).([]*schema.Message)
+		lastOutputMsg2, _ := any(lastOutputMessage).(*schema.Message)
 		failoverCtx := &FailoverContext{
 			FailoverAttempt:   attempt,
-			InputMessages:     input,
-			LastOutputMessage: lastOutputMessage,
+			InputMessages:     inputMsgs2,
+			LastOutputMessage: lastOutputMsg2,
 			LastErr:           lastErr,
 		}
 
-		currentModel, currentInput, err := f.config.GetFailoverModel(ctx, failoverCtx)
+		currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx)
 		if err != nil {
 			return nil, err
 		}
@@ -388,14 +419,15 @@ func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Messa
 			currentInput = input
 		}
 
-		modelCtx := setFailoverCurrentModel(ctx, currentModel)
+		modelCtx := typedSetFailoverCurrentModel(ctx, currentModel)
 		modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries)
 		stream, err := f.inner.Stream(modelCtx, currentInput, opts...)
 		if err != nil {
 			lastErr = err
-			lastOutputMessage = nil
+			var zero M
+			lastOutputMessage = zero
 
-			if !f.needFailover(ctx, nil, err) {
+			if !f.needFailover(ctx, zero, err) {
 				return nil, err
 			}
 
@@ -425,7 +457,7 @@ func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Messa
 		checkCopy := copies[0]
 		returnCopy := copies[1]
 
-		outMsg, streamErr := consumeStream(checkCopy)
+		outMsg, streamErr := typedConsumeStream(checkCopy)
 		if streamErr != nil {
 			lastOutputMessage = outMsg
 			lastErr = streamErr
@@ -441,32 +473,61 @@ func (f *failoverModelWrapper) Stream(ctx context.Context, input []*schema.Messa
 			continue
 		}
 
-		setLastSuccessModel(ctx, currentModel)
+		typedSetFailoverLastSuccessModel[M](ctx, currentModel)
 		return returnCopy, nil
 	}
 
 	return nil, lastErr
 }
 
-func consumeStream(stream *schema.StreamReader[*schema.Message]) (*schema.Message, error) {
+func typedConsumeStream[M messageType](stream *schema.StreamReader[M]) (M, error) {
+	var zero M
 	defer stream.Close()
-	chunks := make([]*schema.Message, 0)
-	for {
-		chunk, err := stream.Recv()
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			// ignore concat error
-			msg, _ := schema.ConcatMessages(chunks)
-			return msg, err
-		}
 
-		chunks = append(chunks, chunk)
+	switch s := any(stream).(type) {
+	case *schema.StreamReader[*schema.Message]:
+		chunks := make([]*schema.Message, 0)
+		for {
+			chunk, err := s.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				msg, _ := schema.ConcatMessages(chunks)
+				if msg != nil {
+					return any(msg).(M), err
+				}
+				return zero, err
+			}
+			chunks = append(chunks, chunk)
+		}
+		msg, _ := schema.ConcatMessages(chunks)
+		if msg != nil {
+			return any(msg).(M), nil
+		}
+		return zero, nil
+	case *schema.StreamReader[*schema.AgenticMessage]:
+		chunks := make([]*schema.AgenticMessage, 0)
+		for {
+			chunk, err := s.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				msg, _ := schema.ConcatAgenticMessages(chunks)
+				if msg != nil {
+					return any(msg).(M), err
+				}
+				return zero, err
+			}
+			chunks = append(chunks, chunk)
+		}
+		msg, _ := schema.ConcatAgenticMessages(chunks)
+		if msg != nil {
+			return any(msg).(M), nil
+		}
+		return zero, nil
+	default:
+		panic("unreachable: unknown messageType")
 	}
-
-	// Stream completed successfully (EOF). ConcatMessages error is not a stream error,
-	// so ignore it to avoid incorrectly triggering failover.
-	msg, _ := schema.ConcatMessages(chunks)
-	return msg, nil
 }

adk/failover_chatmodel_test.go

@@ -104,19 +104,21 @@ func TestFailoverCurrentModelContext(t *testing.T) {
 				return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
 			},
 		}
-		ctx = setFailoverCurrentModel(ctx, m)
-		got := getFailoverCurrentModel(ctx)
-		require.NotNil(t, got)
-		require.Same(t, m, got.model)
+		ctx = typedSetFailoverCurrentModel[*schema.Message](ctx, m)
+		got, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
+		require.True(t, ok)
+		require.Same(t, m, got)
 	})
 
 	t.Run("wrong type", func(t *testing.T) {
 		ctx := context.WithValue(context.Background(), failoverCurrentModelKey{}, "bad")
-		require.Nil(t, getFailoverCurrentModel(ctx))
+		_, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
+		require.False(t, ok)
 	})
 
 	t.Run("missing", func(t *testing.T) {
-		require.Nil(t, getFailoverCurrentModel(context.Background()))
+		_, ok := typedGetFailoverCurrentModel[*schema.Message](context.Background())
+		require.False(t, ok)
 	})
 }
 
@@ -145,7 +147,7 @@ func TestFailoverProxyModel(t *testing.T) {
 				return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("routed", nil)}), nil
 			},
 		}
-		ctx := setFailoverCurrentModel(context.Background(), target)
+		ctx := typedSetFailoverCurrentModel[*schema.Message](context.Background(), target)
 		p := &failoverProxyModel{}
 		msg, err := p.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
 		require.NoError(t, err)
@@ -167,7 +169,7 @@ func TestFailoverModelWrapper_Generate(t *testing.T) {
 				return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("inner", nil)}), nil
 			},
 		}
-		w := newFailoverModelWrapper(inner, &ModelFailoverConfig{
+		w := newTypedFailoverModelWrapper[*schema.Message](inner, &ModelFailoverConfig{
 			MaxRetries:       2,
 			ShouldFailover:   func(context.Context, *schema.Message, error) bool { return true },
 			GetFailoverModel: nil,
@@ -217,8 +219,8 @@ func TestFailoverModelWrapper_Generate(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		msg, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -253,8 +255,8 @@ func TestFailoverModelWrapper_Generate(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -285,7 +287,7 @@ func TestFailoverModelWrapper_Generate(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(inner, cfg)
+		w := newTypedFailoverModelWrapper[*schema.Message](inner, cfg)
 		_, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
 		require.ErrorIs(t, err, wantErr)
 		require.Equal(t, int32(0), atomic.LoadInt32(&called))
@@ -300,7 +302,7 @@ func TestFailoverModelWrapper_Generate(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
 		msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
 		require.Nil(t, msg)
 		require.Error(t, err)
@@ -339,8 +341,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{in})
@@ -392,8 +394,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -453,8 +455,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -491,8 +493,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -524,8 +526,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -563,8 +565,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -583,7 +585,7 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
 		sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
 		require.Nil(t, sr)
 		require.Error(t, err)
@@ -612,7 +614,7 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(inner, cfg)
+		w := newTypedFailoverModelWrapper[*schema.Message](inner, cfg)
 		sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
 		require.Nil(t, sr)
 		require.ErrorIs(t, err, wantErr)
@@ -665,8 +667,8 @@ func TestFailoverModelWrapper_Stream(t *testing.T) {
 			},
 		}
 
-		w := newFailoverModelWrapper(&failoverProxyModel{}, cfg)
-		baseCtx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		w := newTypedFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
+		baseCtx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		ctx, cancel := context.WithCancel(baseCtx)

adk/flow.go

@@ -261,7 +261,7 @@ func genMsg(entry *HistoryEntry, agentName string) (Message, error) {
 	return msg, nil
 }
 
-func (ai *AgentInput) deepCopy() *AgentInput {
+func deepCopyAgentInput(ai *AgentInput) *AgentInput {
 	copied := &AgentInput{
 		Messages:        make([]Message, len(ai.Messages)),
 		EnableStreaming: ai.EnableStreaming,
@@ -273,7 +273,7 @@ func (ai *AgentInput) deepCopy() *AgentInput {
 }
 
 func (a *flowAgent) genAgentInput(ctx context.Context, runCtx *runContext, skipTransferMessages bool) (*AgentInput, error) {
-	input := runCtx.RootInput.deepCopy()
+	input := deepCopyAgentInput(runCtx.RootInput)
 
 	events := runCtx.Session.getEvents()
 	historyEntries := make([]*HistoryEntry, 0)
@@ -521,7 +521,7 @@ func (a *flowAgent) run(
 			// copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time
 			// interrupt action are not added to session, because ALL information contained in it
 			// is either presented to end-user, or made available to agents through other means
-			copied := copyAgentEvent(event)
+			copied := copyTypedAgentEvent(event)
 			setAutomaticClose(copied)
 			setAutomaticClose(event)
 			runCtx.Session.addEvent(copied)
@@ -532,7 +532,7 @@ func (a *flowAgent) run(
 		if exactRunPathMatch(runCtx.RunPath, event.RunPath) {
 			lastAction = event.Action
 		}
-		copied := copyAgentEvent(event)
+		copied := copyTypedAgentEvent(event)
 		setAutomaticClose(copied)
 		setAutomaticClose(event)
 		cbGen.Send(copied)
@@ -604,10 +604,206 @@ func wrapIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*AgentEvent]) *A
 			if !ok {
 				break
 			}
-			copied := copyAgentEvent(event)
+			copied := copyTypedAgentEvent(event)
 			cbGen.Send(copied)
 			outGen.Send(event)
 		}
 	}()
 	return outIter
 }
+
+// ---------------------------------------------------------------------------
+// Typed wrapper for the agentic path (TypedAgent[*schema.AgenticMessage]).
+//
+// typedFlowAgent is a minimal wrapper used exclusively by TypedRunner and
+// AgentTool to execute a TypedAgent[*schema.AgenticMessage]. It handles
+// callbacks, event recording, and run-path tracking. Transfer, sub-agent
+// orchestration, and history rewriting are handled solely by the concrete
+// flowAgent (the *schema.Message path).
+// ---------------------------------------------------------------------------
+
+type typedFlowAgent[M messageType] struct {
+	TypedAgent[M]
+
+	checkPointStore compose.CheckPointStore
+}
+
+func toTypedFlowAgent[M messageType](agent TypedAgent[M]) *typedFlowAgent[M] {
+	if fa, ok := agent.(*typedFlowAgent[M]); ok {
+		return fa
+	}
+	return &typedFlowAgent[M]{TypedAgent: agent}
+}
+
+func getTypedAgentType[M messageType](agent TypedAgent[M]) string {
+	if msgAgent, ok := any(agent).(Agent); ok {
+		return getAgentType(msgAgent)
+	}
+	if typer, ok := any(agent).(interface{ GetType() string }); ok {
+		return typer.GetType()
+	}
+	return ""
+}
+
+func (a *typedFlowAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	agentName := a.Name(ctx)
+
+	var runCtx *runContext
+	ctx, runCtx = initTypedRunCtx(ctx, agentName, input)
+	ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName)
+
+	o := getCommonOptions(nil, opts...)
+	cancelCtx := o.cancelCtx
+
+	ctxForSubAgents := ctx
+
+	agentType := getTypedAgentType(a.TypedAgent)
+	ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
+	cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{Input: any(input).(*TypedAgentInput[*schema.AgenticMessage])}
+	ctx = callbacks.OnStart(ctx, cbInput)
+
+	aIter := a.TypedAgent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...)
+
+	iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+
+	go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...)
+
+	return wrapIterWithCancelCtx(iterator, cancelCtx)
+}
+
+func (a *typedFlowAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	agentName := a.Name(ctx)
+
+	ctx, info = buildResumeInfo(ctx, agentName, info)
+
+	ctxForSubAgents := ctx
+
+	o := getCommonOptions(nil, opts...)
+	cancelCtx := o.cancelCtx
+
+	agentType := getTypedAgentType(a.TypedAgent)
+	ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...)
+	cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{ResumeInfo: info}
+	ctx = callbacks.OnStart(ctx, cbInput)
+
+	if info.WasInterrupted {
+		if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
+			aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...)
+
+			iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+			go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...)
+			return wrapIterWithCancelCtx(iterator, cancelCtx)
+		}
+
+		if cancelCtx != nil {
+			cancelCtx.markDone()
+		}
+		return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+
+			"but is not a ResumableAgent", agentName))
+	}
+
+	_, err := getNextResumeAgent(ctx, info)
+	if err != nil {
+		if cancelCtx != nil {
+			cancelCtx.markDone()
+		}
+		return typedErrorIterWithOnEnd[M](ctx, err)
+	}
+
+	if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok {
+		ctx = withCancelContext(ctx, cancelCtx)
+		innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...)
+		return wrapIterWithCancelCtx(typedWrapIterWithOnEnd[M](ctx, innerIter), cancelCtx)
+	}
+	return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf(
+		"failed to resume agent: agent '%s' (type %T) does not implement ResumableAgent interface. "+
+			"To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.TypedAgent))
+}
+
+func (a *typedFlowAgent[M]) run(
+	ctx context.Context,
+	_ context.Context,
+	runCtx *runContext,
+	aIter *AsyncIterator[*TypedAgentEvent[M]],
+	generator *AsyncGenerator[*TypedAgentEvent[M]],
+	_ ...AgentRunOption) {
+
+	agenticCbIter, agenticCbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+	cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: agenticCbIter}
+	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)
+
+	defer func() {
+		panicErr := recover()
+		if panicErr != nil {
+			e := safe.NewPanicErr(panicErr, debug.Stack())
+			generator.Send(&TypedAgentEvent[M]{Err: e})
+		}
+
+		agenticCbGen.Close()
+		generator.Close()
+	}()
+
+	for {
+		event, ok := aIter.Next()
+		if !ok {
+			break
+		}
+
+		if len(event.RunPath) == 0 {
+			event.AgentName = a.Name(ctx)
+			event.RunPath = runCtx.RunPath
+		}
+		if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) {
+			copied := copyTypedAgentEvent(event)
+			typedSetAutomaticClose(copied)
+			typedSetAutomaticClose(event)
+			addTypedEvent(runCtx.Session, copied)
+		}
+
+		agenticCopied := copyTypedAgentEvent(event)
+		typedSetAutomaticClose(agenticCopied)
+		typedSetAutomaticClose(event)
+		agenticCbGen.Send(any(agenticCopied).(*TypedAgentEvent[*schema.AgenticMessage]))
+		generator.Send(event)
+	}
+}
+
+func wrapAgenticIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+	cbIter, cbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+	cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: cbIter}
+	icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false)
+
+	outIter, outGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+	go func() {
+		defer func() {
+			cbGen.Close()
+			outGen.Close()
+		}()
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			copied := copyTypedAgentEvent(event)
+			cbGen.Send(copied)
+			outGen.Send(event)
+		}
+	}()
+	return outIter
+}
+
+func genAgenticErrorIter(err error) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+	iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+	gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err})
+	gen.Close()
+	return iter
+}
+
+func typedWrapIterWithOnEnd[M messageType](ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[M]]) *AsyncIterator[*TypedAgentEvent[M]] {
+	agenticIter := any(iter).(*AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]])
+	return any(wrapAgenticIterWithOnEnd(ctx, agenticIter)).(*AsyncIterator[*TypedAgentEvent[M]])
+}
+
+func typedErrorIterWithOnEnd[M messageType](ctx context.Context, err error) *AsyncIterator[*TypedAgentEvent[M]] {
+	return typedWrapIterWithOnEnd[M](ctx, typedErrorIter[M](err))
+}

adk/handler.go

@@ -96,12 +96,12 @@ type ChatModelAgentContext struct {
 	ReturnDirectly map[string]bool
 }
 
-// ChatModelAgentMiddleware defines the interface for customizing ChatModelAgent behavior.
+// TypedChatModelAgentMiddleware defines the interface for customizing TypedChatModelAgent behavior.
 //
-// IMPORTANT: This interface is specifically designed for ChatModelAgent and agents built
+// IMPORTANT: This interface is specifically designed for TypedChatModelAgent and agents built
 // on top of it (e.g., DeepAgent).
 //
-// Why ChatModelAgentMiddleware instead of AgentMiddleware?
+// Why TypedChatModelAgentMiddleware instead of AgentMiddleware?
 //
 // AgentMiddleware is a struct type, which has inherent limitations:
 //   - Struct types are closed: users cannot add new methods to extend functionality
@@ -110,22 +110,22 @@ type ChatModelAgentContext struct {
 //     call those methods (config.Middlewares is []AgentMiddleware, not a user type)
 //   - Callbacks in AgentMiddleware only return error, cannot return modified context
 //
-// ChatModelAgentMiddleware is an interface type, which is open for extension:
+// TypedChatModelAgentMiddleware is an interface type, which is open for extension:
 //   - Users can implement custom handlers with arbitrary internal state and methods
 //   - Hook methods return (context.Context, ..., error) for direct context propagation
 //   - Wrapper methods (WrapToolCall, WrapModel) enable context propagation through the
 //     wrapped endpoint chain: wrappers can pass modified context to the next wrapper
 //   - Configuration is centralized in struct fields rather than scattered in closures
 //
-// ChatModelAgentMiddleware vs AgentMiddleware:
+// TypedChatModelAgentMiddleware vs AgentMiddleware:
 //   - Use AgentMiddleware for simple, static additions (extra instruction/tools)
-//   - Use ChatModelAgentMiddleware for dynamic behavior, context modification, or call wrapping
+//   - Use TypedChatModelAgentMiddleware for dynamic behavior, context modification, or call wrapping
 //   - AgentMiddleware is kept for backward compatibility with existing users
 //   - Both can be used together; see AgentMiddleware documentation for execution order
 //
-// Use *BaseChatModelAgentMiddleware as an embedded struct to provide default no-op
+// Use *TypedBaseChatModelAgentMiddleware as an embedded struct to provide default no-op
 // implementations for all methods.
-type ChatModelAgentMiddleware interface {
+type TypedChatModelAgentMiddleware[M messageType] interface {
 	// BeforeAgent is called before each agent run, allowing modification of
 	// the agent's instruction and tools configuration.
 	BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error)
@@ -139,7 +139,7 @@ type ChatModelAgentMiddleware interface {
 	//
 	// The ModelContext struct provides read-only access to:
 	//   - Tools: the current tool list that will be sent to the model
-	BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)
+	BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *ModelContext) (context.Context, *TypedChatModelAgentState[M], error)
 
 	// AfterModelRewriteState is called after each model invocation.
 	// The input state includes the model's response as the last message.
@@ -150,7 +150,7 @@ type ChatModelAgentMiddleware interface {
 	//
 	// The ModelContext struct provides read-only access to:
 	//   - Tools: the current tool list that was sent to the model
-	AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error)
+	AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *ModelContext) (context.Context, *TypedChatModelAgentState[M], error)
 
 	// AfterToolCallsRewriteState is called after all concurrent tool calls in an iteration complete.
 	// The input state includes all messages up to and including the tool call results.
@@ -158,7 +158,7 @@ type ChatModelAgentMiddleware interface {
 	//
 	// The ToolCallsContext provides metadata about the tool calls that just completed,
 	// derived from the assistant message's ToolCalls field.
-	AfterToolCallsRewriteState(ctx context.Context, state *ChatModelAgentState, tc *ToolCallsContext) (context.Context, *ChatModelAgentState, error)
+	AfterToolCallsRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], tc *ToolCallsContext) (context.Context, *TypedChatModelAgentState[M], error)
 
 	// WrapInvokableToolCall wraps a tool's synchronous execution with custom behavior.
 	// Return the input endpoint unchanged and nil error if no wrapping is needed.
@@ -212,15 +212,21 @@ type ChatModelAgentMiddleware interface {
 	// Return the input model unchanged and nil error if no wrapping is needed.
 	//
 	// This method is called at request time when the model is about to be invoked.
-	// Note: The parameter is BaseChatModel (not ToolCallingChatModel) because wrappers
+	// Note: The parameter is model.BaseModel[M] (not ToolCallingChatModel) because wrappers
 	// only need to intercept Generate/Stream calls. Tool binding (WithTools) is handled
 	// separately by the framework and does not flow through user wrappers.
 	//
 	// The mc parameter contains the current tool configuration:
 	//   - Tools: The tool infos that will be sent to the model
-	WrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error)
+	WrapModel(ctx context.Context, m model.BaseModel[M], mc *ModelContext) (model.BaseModel[M], error)
 }
 
+// ChatModelAgentMiddleware is the default middleware type using *schema.Message.
+// See TypedChatModelAgentMiddleware for full documentation.
+type ChatModelAgentMiddleware = TypedChatModelAgentMiddleware[*schema.Message]
+
+type TypedBaseChatModelAgentMiddleware[M messageType] struct{}
+
 // BaseChatModelAgentMiddleware provides default no-op implementations for ChatModelAgentMiddleware.
 // Embed *BaseChatModelAgentMiddleware in custom handlers to only override the methods you need.
 //
@@ -235,44 +241,58 @@ type ChatModelAgentMiddleware interface {
 //		// custom logic
 //		return ctx, state, nil
 //	}
-type BaseChatModelAgentMiddleware struct{}
+type BaseChatModelAgentMiddleware = TypedBaseChatModelAgentMiddleware[*schema.Message]
 
-func (b *BaseChatModelAgentMiddleware) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) {
 	return endpoint, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) {
 	return endpoint, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) {
 	return endpoint, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) {
 	return endpoint, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) WrapModel(_ context.Context, m model.BaseChatModel, _ *ModelContext) (model.BaseChatModel, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) WrapModel(_ context.Context, m model.BaseModel[M], _ *ModelContext) (model.BaseModel[M], error) {
 	return m, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) {
 	return ctx, runCtx, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *ModelContext) (context.Context, *TypedChatModelAgentState[M], error) {
 	return ctx, state, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *ModelContext) (context.Context, *TypedChatModelAgentState[M], error) {
 	return ctx, state, nil
 }
 
-func (b *BaseChatModelAgentMiddleware) AfterToolCallsRewriteState(ctx context.Context, state *ChatModelAgentState, tc *ToolCallsContext) (context.Context, *ChatModelAgentState, error) {
+func (b *TypedBaseChatModelAgentMiddleware[M]) AfterToolCallsRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], tc *ToolCallsContext) (context.Context, *TypedChatModelAgentState[M], error) {
 	return ctx, state, nil
 }
 
+func processTypedState(ctx context.Context, fn func(extra map[string]any) map[string]any) error {
+	runCtx := getRunCtx(ctx)
+	if runCtx != nil && runCtx.AgenticRootInput != nil {
+		return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.AgenticMessage]) error {
+			st.Extra = fn(st.Extra)
+			return nil
+		})
+	}
+	return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.Message]) error {
+		st.Extra = fn(st.Extra)
+		return nil
+	})
+}
+
 // SetRunLocalValue sets a key-value pair that persists for the duration of the current agent Run() invocation.
 // The value is scoped to this specific execution and is not shared across different Run() calls or agent instances.
 //
@@ -287,12 +307,12 @@ func SetRunLocalValue(ctx context.Context, key string, value any) error {
 		return err
 	}
 
-	err := compose.ProcessState(ctx, func(_ context.Context, st *State) error {
-		if st.Extra == nil {
-			st.Extra = make(map[string]any)
+	err := processTypedState(ctx, func(extra map[string]any) map[string]any {
+		if extra == nil {
+			extra = make(map[string]any)
 		}
-		st.Extra[key] = value
-		return nil
+		extra[key] = value
+		return extra
 	})
 	if err != nil {
 		return fmt.Errorf("SetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
@@ -313,11 +333,11 @@ func SetRunLocalValue(ctx context.Context, key string, value any) error {
 func GetRunLocalValue(ctx context.Context, key string) (any, bool, error) {
 	var val any
 	var found bool
-	err := compose.ProcessState(ctx, func(_ context.Context, st *State) error {
-		if st.Extra != nil {
-			val, found = st.Extra[key]
+	err := processTypedState(ctx, func(extra map[string]any) map[string]any {
+		if extra != nil {
+			val, found = extra[key]
 		}
-		return nil
+		return extra
 	})
 	if err != nil {
 		return nil, false, fmt.Errorf("GetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
@@ -330,11 +350,11 @@ func GetRunLocalValue(ctx context.Context, key string) (any, bool, error) {
 // This function can only be called from within a ChatModelAgentMiddleware during agent execution.
 // Returns an error if called outside of an agent execution context.
 func DeleteRunLocalValue(ctx context.Context, key string) error {
-	err := compose.ProcessState(ctx, func(_ context.Context, st *State) error {
-		if st.Extra != nil {
-			delete(st.Extra, key)
+	err := processTypedState(ctx, func(extra map[string]any) map[string]any {
+		if extra != nil {
+			delete(extra, key)
 		}
-		return nil
+		return extra
 	})
 	if err != nil {
 		return fmt.Errorf("DeleteRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err)
@@ -342,6 +362,21 @@ func DeleteRunLocalValue(ctx context.Context, key string) error {
 	return nil
 }
 
+// TypedSendEvent sends a custom TypedAgentEvent to the event stream during agent execution.
+// This allows TypedChatModelAgentMiddleware implementations to emit custom events that will be
+// received by the caller iterating over the agent's event stream.
+//
+// This function can only be called from within a TypedChatModelAgentMiddleware during agent execution.
+// Returns an error if called outside of an agent execution context.
+func TypedSendEvent[M messageType](ctx context.Context, event *TypedAgentEvent[M]) error {
+	execCtx := getTypedChatModelAgentExecCtx[M](ctx)
+	if execCtx == nil || execCtx.generator == nil {
+		return fmt.Errorf("TypedSendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context")
+	}
+	execCtx.send(event)
+	return nil
+}
+
 // SendEvent sends a custom AgentEvent to the event stream during agent execution.
 // This allows ChatModelAgentMiddleware implementations to emit custom events that will be
 // received by the caller iterating over the agent's event stream.
@@ -349,12 +384,7 @@ func DeleteRunLocalValue(ctx context.Context, key string) error {
 // This function can only be called from within a ChatModelAgentMiddleware during agent execution.
 // Returns an error if called outside of an agent execution context.
 func SendEvent(ctx context.Context, event *AgentEvent) error {
-	execCtx := getChatModelAgentExecCtx(ctx)
-	if execCtx == nil || execCtx.generator == nil {
-		return fmt.Errorf("SendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context")
-	}
-	execCtx.send(event)
-	return nil
+	return TypedSendEvent[*schema.Message](ctx, event)
 }
 
 // checkGobEncodability probes whether the value can be gob-encoded as part of

adk/instruction.go

@@ -45,7 +45,7 @@ When transferring: OUTPUT ONLY THE FUNCTION CALL`
 	agentDescriptionTplChinese = "\n- Agent 名字: %s\n  Agent 描述: %s"
 )
 
-func genTransferToAgentInstruction(ctx context.Context, agents []Agent) string {
+func genTransferToAgentInstruction[M messageType](ctx context.Context, agents []TypedAgent[M]) string {
 	tpl := internal.SelectPrompt(internal.I18nPrompts{
 		English: agentDescriptionTpl,
 		Chinese: agentDescriptionTplChinese,

adk/interface.go

@@ -32,36 +32,80 @@ import (
 // Use this to filter callback events to only agent-related events.
 const ComponentOfAgent components.Component = "Agent"
 
+// ComponentOfAgenticAgent is the component type identifier for ADK agents
+// that use *schema.AgenticMessage in callbacks.
+const ComponentOfAgenticAgent components.Component = "AgenticAgent"
+
+// messageType is the sealed type constraint for message types used in ADK.
+// Only *schema.Message and *schema.AgenticMessage satisfy this constraint.
+// External packages cannot add new types to this union; all generic functions
+// in ADK use exhaustive type switches on these two types.
+type messageType interface {
+	*schema.Message | *schema.AgenticMessage
+}
+
 type Message = *schema.Message
 type MessageStream = *schema.StreamReader[Message]
 
-type MessageVariant struct {
+type AgenticMessage = *schema.AgenticMessage
+type AgenticMessageStream = *schema.StreamReader[AgenticMessage]
+
+// isNilMessage checks whether a generic message value is nil.
+// Direct `msg == nil` does not compile for generic pointer types in Go;
+// the canonical workaround is to compare through the `any` interface.
+func isNilMessage[M messageType](msg M) bool {
+	var zero M
+	return any(msg) == any(zero)
+}
+
+// TypedMessageVariant represents a message output from an agent event.
+// It carries either a complete message or a streaming reader, along with
+// metadata describing the event's origin.
+//
+// Role and ToolName are only meaningful for *schema.Message events. For
+// *schema.AgenticMessage events (created via EventFromAgenticMessage), these
+// fields are always zero-valued because AgenticMessage carries tool results as
+// ContentBlocks within the message itself and does not support agent transfer.
+//
+// For *schema.Message events, Role and ToolName exist independently of the inner
+// Message because in streaming mode (IsStreaming=true, Message=nil), the message
+// has not materialized yet and the consumer needs metadata without consuming the stream.
+type TypedMessageVariant[M messageType] struct {
 	IsStreaming bool
 
-	Message       Message
-	MessageStream MessageStream
-	// message role: Assistant or Tool
+	Message       M
+	MessageStream *schema.StreamReader[M]
+
+	// Role indicates the origin of this event within the agent's ReAct loop.
+	// Only meaningful for *schema.Message events:
+	//   - schema.Assistant: the event carries model output (generation or stream).
+	//   - schema.Tool: the event carries a tool execution result.
+	// Always zero-valued for *schema.AgenticMessage events; use AgenticRole instead.
 	Role schema.RoleType
-	// only used when Role is Tool
+
+	// AgenticRole indicates the role of the agentic message (assistant, user, system).
+	// Only meaningful for *schema.AgenticMessage events.
+	// In streaming mode, this is available before consuming the stream.
+	// Always zero-valued for *schema.Message events; use Role instead.
+	AgenticRole schema.AgenticRoleType
+
+	// ToolName is the name of the tool that produced this event.
+	// Only meaningful for *schema.Message events: non-empty when Role == schema.Tool.
+	// In streaming mode, this is the only way to identify the source tool before
+	// the stream is consumed.
+	// Always empty for *schema.AgenticMessage events.
 	ToolName string
 }
 
-// EventFromMessage wraps a message or stream into an AgentEvent with role metadata.
-func EventFromMessage(msg Message, msgStream MessageStream,
-	role schema.RoleType, toolName string) *AgentEvent {
-	return &AgentEvent{
-		Output: &AgentOutput{
-			MessageOutput: &MessageVariant{
-				IsStreaming:   msgStream != nil,
-				Message:       msg,
-				MessageStream: msgStream,
-				Role:          role,
-				ToolName:      toolName,
-			},
-		},
+func (mv *TypedMessageVariant[M]) GetMessage() (M, error) {
+	if mv.IsStreaming {
+		return concatMessageStream(mv.MessageStream)
 	}
+	return mv.Message, nil
 }
 
+type MessageVariant = TypedMessageVariant[*schema.Message]
+
 type messageVariantSerialization struct {
 	IsStreaming   bool
 	Message       Message
@@ -70,7 +114,36 @@ type messageVariantSerialization struct {
 	ToolName      string
 }
 
-func (mv *MessageVariant) GobEncode() ([]byte, error) {
+type agenticMessageVariantSerialization struct {
+	IsStreaming   bool
+	Message       *schema.AgenticMessage
+	MessageStream *schema.AgenticMessage
+	Role          schema.RoleType
+	AgenticRole   schema.AgenticRoleType
+	ToolName      string
+}
+
+func (mv *TypedMessageVariant[M]) GobEncode() ([]byte, error) {
+	if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok {
+		return gobEncodeMessageVariant(mvMsg)
+	}
+	if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok {
+		return gobEncodeAgenticMessageVariant(mvAgentic)
+	}
+	return nil, fmt.Errorf("gob encoding not supported for this message type")
+}
+
+func (mv *TypedMessageVariant[M]) GobDecode(b []byte) error {
+	if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok {
+		return gobDecodeMessageVariant(mvMsg, b)
+	}
+	if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok {
+		return gobDecodeAgenticMessageVariant(mvAgentic, b)
+	}
+	return fmt.Errorf("gob decoding not supported for this message type")
+}
+
+func gobEncodeMessageVariant(mv *TypedMessageVariant[*schema.Message]) ([]byte, error) {
 	s := &messageVariantSerialization{
 		IsStreaming: mv.IsStreaming,
 		Message:     mv.Message,
@@ -103,7 +176,7 @@ func (mv *MessageVariant) GobEncode() ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
-func (mv *MessageVariant) GobDecode(b []byte) error {
+func gobDecodeMessageVariant(mv *TypedMessageVariant[*schema.Message], b []byte) error {
 	s := &messageVariantSerialization{}
 	err := gob.NewDecoder(bytes.NewReader(b)).Decode(s)
 	if err != nil {
@@ -119,19 +192,120 @@ func (mv *MessageVariant) GobDecode(b []byte) error {
 	return nil
 }
 
-func (mv *MessageVariant) GetMessage() (Message, error) {
-	var message Message
-	if mv.IsStreaming {
-		var err error
-		message, err = schema.ConcatMessageStream(mv.MessageStream)
-		if err != nil {
-			return nil, err
-		}
-	} else {
-		message = mv.Message
+func gobEncodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage]) ([]byte, error) {
+	s := &agenticMessageVariantSerialization{
+		IsStreaming: mv.IsStreaming,
+		Message:     mv.Message,
+		Role:        mv.Role,
+		AgenticRole: mv.AgenticRole,
+		ToolName:    mv.ToolName,
 	}
+	if mv.IsStreaming {
+		var messages []*schema.AgenticMessage
+		for {
+			frame, err := mv.MessageStream.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				return nil, fmt.Errorf("error receiving agentic message stream: %w", err)
+			}
+			messages = append(messages, frame)
+		}
+		m, err := schema.ConcatAgenticMessages(messages)
+		if err != nil {
+			return nil, fmt.Errorf("failed to encode agentic message: cannot concat message stream: %w", err)
+		}
+		s.MessageStream = m
+	}
+	buf := &bytes.Buffer{}
+	err := gob.NewEncoder(buf).Encode(s)
+	if err != nil {
+		return nil, fmt.Errorf("failed to gob encode agentic message variant: %w", err)
+	}
+	return buf.Bytes(), nil
+}
 
-	return message, nil
+func gobDecodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage], b []byte) error {
+	s := &agenticMessageVariantSerialization{}
+	err := gob.NewDecoder(bytes.NewReader(b)).Decode(s)
+	if err != nil {
+		return fmt.Errorf("failed to decode agentic message variant: %w", err)
+	}
+	mv.IsStreaming = s.IsStreaming
+	mv.Message = s.Message
+	mv.Role = s.Role
+	mv.AgenticRole = s.AgenticRole
+	mv.ToolName = s.ToolName
+	if s.MessageStream != nil {
+		mv.MessageStream = schema.StreamReaderFromArray([]*schema.AgenticMessage{s.MessageStream})
+	}
+	return nil
+}
+
+// typedEventFromMessage creates a TypedAgentEvent containing the given message and optional stream.
+func typedEventFromMessage[M messageType](msg M, msgStream *schema.StreamReader[M],
+	role schema.RoleType, toolName string) *TypedAgentEvent[M] {
+	return &TypedAgentEvent[M]{
+		Output: &TypedAgentOutput[M]{
+			MessageOutput: &TypedMessageVariant[M]{
+				IsStreaming:   msgStream != nil,
+				Message:       msg,
+				MessageStream: msgStream,
+				Role:          role,
+				ToolName:      toolName,
+			},
+		},
+	}
+}
+
+// typedModelOutputEvent creates a model-output event for the generic path.
+// For *schema.Message, Role is set to schema.Assistant.
+// For *schema.AgenticMessage, AgenticRole is set to schema.AgenticRoleTypeAssistant.
+func typedModelOutputEvent[M messageType](msg M, msgStream *schema.StreamReader[M]) *TypedAgentEvent[M] {
+	var role schema.RoleType
+	var agenticRole schema.AgenticRoleType
+	var zero M
+	if _, ok := any(zero).(*schema.Message); ok {
+		role = schema.Assistant
+	} else {
+		agenticRole = schema.AgenticRoleTypeAssistant
+	}
+	event := typedEventFromMessage(msg, msgStream, role, "")
+	event.Output.MessageOutput.AgenticRole = agenticRole
+	return event
+}
+
+// EventFromMessage creates an AgentEvent containing the given message and optional stream.
+//
+// role identifies the origin of this event:
+//   - schema.Assistant: model output (generation or stream).
+//   - schema.Tool: tool execution result; toolName must be non-empty.
+//
+// For *schema.AgenticMessage events, use EventFromAgenticMessage instead.
+func EventFromMessage(msg Message, msgStream *schema.StreamReader[Message],
+	role schema.RoleType, toolName string) *AgentEvent {
+	return typedEventFromMessage(msg, msgStream, role, toolName)
+}
+
+// EventFromAgenticMessage creates a TypedAgentEvent for the AgenticMessage path.
+// Unlike EventFromMessage, it does not require role or toolName parameters because
+// AgenticMessage carries tool results as ContentBlocks within the message itself,
+// and does not support agent transfer.
+//
+// agenticRole identifies the role of the message (e.g. schema.AgenticRoleTypeAssistant).
+// In streaming mode, the role is available on the event before consuming the stream.
+func EventFromAgenticMessage(msg AgenticMessage, msgStream AgenticMessageStream, agenticRole schema.AgenticRoleType) *TypedAgentEvent[AgenticMessage] {
+	return &TypedAgentEvent[AgenticMessage]{
+		Output: &TypedAgentOutput[AgenticMessage]{
+			MessageOutput: &TypedMessageVariant[AgenticMessage]{
+				IsStreaming:   msgStream != nil,
+				Message:       msg,
+				MessageStream: msgStream,
+				AgenticRole:   agenticRole,
+			},
+		},
+	}
 }
 
 // TransferToAgentAction represents a transfer-to-agent action.
@@ -143,12 +317,14 @@ type TransferToAgentAction struct {
 	DestAgentName string
 }
 
-type AgentOutput struct {
-	MessageOutput *MessageVariant
+type TypedAgentOutput[M messageType] struct {
+	MessageOutput *TypedMessageVariant[M]
 
 	CustomizedOutput any
 }
 
+type AgentOutput = TypedAgentOutput[*schema.Message]
+
 // NewTransferToAgentAction creates an action to transfer to the specified agent.
 //
 // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
@@ -238,8 +414,9 @@ type runStepSerialization struct {
 	AgentName string
 }
 
-// AgentEvent CheckpointSchema: persisted via serialization.RunCtx (gob).
-type AgentEvent struct {
+// TypedAgentEvent represents a single event emitted during agent execution.
+// CheckpointSchema: persisted via serialization.RunCtx (gob).
+type TypedAgentEvent[M messageType] struct {
 	AgentName string
 
 	// RunPath represents the execution path from root agent to the current event source.
@@ -250,20 +427,30 @@ type AgentEvent struct {
 	// AgentTool or DeepAgent, RunPath is trivial. Consider those patterns instead.
 	RunPath []RunStep
 
-	Output *AgentOutput
+	Output *TypedAgentOutput[M]
 
 	Action *AgentAction
 
 	Err error
 }
 
-type AgentInput struct {
-	Messages        []Message
+// AgentEvent is the default event type using *schema.Message.
+type AgentEvent = TypedAgentEvent[*schema.Message]
+
+type TypedAgentInput[M messageType] struct {
+	Messages        []M
 	EnableStreaming bool
 }
 
-//go:generate  mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk -source interface.go
-type Agent interface {
+type AgentInput = TypedAgentInput[*schema.Message]
+
+// TypedAgent is the base agent interface parameterized by message type.
+//
+// For M = *schema.Message, the full ADK feature set is supported (multi-agent
+// orchestration, cancel monitoring, retry, flowAgent).
+// For M = *schema.AgenticMessage, single-agent execution works but cancel
+// monitoring on the model stream and retry are not yet wired.
+type TypedAgent[M messageType] interface {
 	Name(ctx context.Context) string
 	Description(ctx context.Context) string
 
@@ -273,9 +460,12 @@ type Agent interface {
 	// the MessageStream MUST be exclusive and safe to be received directly.
 	// NOTE: it's recommended to use SetAutomaticClose() on the MessageStream of AgentEvents emitted by AsyncIterator,
 	// so that even the events are not processed, the MessageStream can still be closed.
-	Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent]
+	Run(ctx context.Context, input *TypedAgentInput[M], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]]
 }
 
+//go:generate  mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk github.com/cloudwego/eino/adk Agent,ResumableAgent
+type Agent = TypedAgent[*schema.Message]
+
 // OnSubAgents is the interface for agents that support sub-agent registration and transfer.
 //
 // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven
@@ -288,8 +478,42 @@ type OnSubAgents interface {
 	OnDisallowTransferToParent(ctx context.Context) error
 }
 
-type ResumableAgent interface {
-	Agent
+type TypedResumableAgent[M messageType] interface {
+	TypedAgent[M]
 
-	Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent]
+	Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]]
+}
+
+type ResumableAgent = TypedResumableAgent[*schema.Message]
+
+func concatMessageStream[M messageType](stream *schema.StreamReader[M]) (M, error) {
+	var zero M
+	switch s := any(stream).(type) {
+	case *schema.StreamReader[*schema.Message]:
+		result, err := schema.ConcatMessageStream(s)
+		if err != nil {
+			return zero, err
+		}
+		return any(result).(M), nil
+	case *schema.StreamReader[*schema.AgenticMessage]:
+		defer s.Close()
+		var msgs []*schema.AgenticMessage
+		for {
+			frame, err := s.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				return zero, err
+			}
+			msgs = append(msgs, frame)
+		}
+		result, err := schema.ConcatAgenticMessages(msgs)
+		if err != nil {
+			return zero, err
+		}
+		return any(result).(M), nil
+	default:
+		panic("unreachable: unknown messageType")
+	}
 }

adk/interrupt.go

@@ -54,11 +54,9 @@ type InterruptInfo struct {
 	InterruptContexts []*InterruptCtx
 }
 
-// Interrupt creates a basic interrupt action.
-// This is used when an agent needs to pause its execution to request external input or intervention,
-// but does not need to save any internal state to be restored upon resumption.
-// The `info` parameter is user-facing data that describes the reason for the interrupt.
-func Interrupt(ctx context.Context, info any) *AgentEvent {
+// TypedInterrupt creates a typed interrupt event that pauses execution to request external input.
+// It is the generic counterpart of Interrupt; see Interrupt for full documentation.
+func TypedInterrupt[M messageType](ctx context.Context, info any) *TypedAgentEvent[M] {
 	var rp []RunStep
 	rCtx := getRunCtx(ctx)
 	if rCtx != nil {
@@ -68,12 +66,47 @@ func Interrupt(ctx context.Context, info any) *AgentEvent {
 	is, err := core.Interrupt(ctx, info, nil, nil,
 		core.WithLayerPayload(rp))
 	if err != nil {
-		return &AgentEvent{Err: err}
+		return &TypedAgentEvent[M]{Err: err}
 	}
 
 	contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
 
-	return &AgentEvent{
+	return &TypedAgentEvent[M]{
+		Action: &AgentAction{
+			Interrupted: &InterruptInfo{
+				InterruptContexts: contexts,
+			},
+			internalInterrupted: is,
+		},
+	}
+}
+
+// Interrupt creates a basic interrupt action.
+// This is used when an agent needs to pause its execution to request external input or intervention,
+// but does not need to save any internal state to be restored upon resumption.
+// The `info` parameter is user-facing data that describes the reason for the interrupt.
+func Interrupt(ctx context.Context, info any) *AgentEvent {
+	return TypedInterrupt[*schema.Message](ctx, info)
+}
+
+// TypedStatefulInterrupt creates a typed interrupt event that also saves the agent's internal state.
+// It is the generic counterpart of StatefulInterrupt; see StatefulInterrupt for full documentation.
+func TypedStatefulInterrupt[M messageType](ctx context.Context, info any, state any) *TypedAgentEvent[M] {
+	var rp []RunStep
+	rCtx := getRunCtx(ctx)
+	if rCtx != nil {
+		rp = rCtx.RunPath
+	}
+
+	is, err := core.Interrupt(ctx, info, state, nil,
+		core.WithLayerPayload(rp))
+	if err != nil {
+		return &TypedAgentEvent[M]{Err: err}
+	}
+
+	contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
+
+	return &TypedAgentEvent[M]{
 		Action: &AgentAction{
 			Interrupted: &InterruptInfo{
 				InterruptContexts: contexts,
@@ -88,38 +121,13 @@ func Interrupt(ctx context.Context, info any) *AgentEvent {
 // The `info` parameter is user-facing data describing the interrupt.
 // The `state` parameter is the agent's internal state object, which will be serialized and stored.
 func StatefulInterrupt(ctx context.Context, info any, state any) *AgentEvent {
-	var rp []RunStep
-	rCtx := getRunCtx(ctx)
-	if rCtx != nil {
-		rp = rCtx.RunPath
-	}
-
-	is, err := core.Interrupt(ctx, info, state, nil,
-		core.WithLayerPayload(rp))
-	if err != nil {
-		return &AgentEvent{Err: err}
-	}
-
-	contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
-
-	return &AgentEvent{
-		Action: &AgentAction{
-			Interrupted: &InterruptInfo{
-				InterruptContexts: contexts,
-			},
-			internalInterrupted: is,
-		},
-	}
+	return TypedStatefulInterrupt[*schema.Message](ctx, info, state)
 }
 
-// CompositeInterrupt creates an interrupt action for a workflow agent.
-// It combines the interrupts from one or more of its sub-agents into a single, cohesive interrupt.
-// This is used by workflow agents (like Sequential, Parallel, or Loop) to propagate interrupts from their children.
-// The `info` parameter is user-facing data describing the workflow's own reason for interrupting.
-// The `state` parameter is the workflow agent's own state (e.g., the index of the sub-agent that was interrupted).
-// The `subInterruptSignals` is a variadic list of the InterruptSignal objects from the interrupted sub-agents.
-func CompositeInterrupt(ctx context.Context, info any, state any,
-	subInterruptSignals ...*InterruptSignal) *AgentEvent {
+// TypedCompositeInterrupt creates a typed interrupt event that aggregates sub-interrupt signals.
+// It is the generic counterpart of CompositeInterrupt; see CompositeInterrupt for full documentation.
+func TypedCompositeInterrupt[M messageType](ctx context.Context, info any, state any,
+	subInterruptSignals ...*InterruptSignal) *TypedAgentEvent[M] {
 	var rp []RunStep
 	rCtx := getRunCtx(ctx)
 	if rCtx != nil {
@@ -129,12 +137,12 @@ func CompositeInterrupt(ctx context.Context, info any, state any,
 	is, err := core.Interrupt(ctx, info, state, subInterruptSignals,
 		core.WithLayerPayload(rp))
 	if err != nil {
-		return &AgentEvent{Err: err}
+		return &TypedAgentEvent[M]{Err: err}
 	}
 
 	contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes)
 
-	return &AgentEvent{
+	return &TypedAgentEvent[M]{
 		Action: &AgentAction{
 			Interrupted: &InterruptInfo{
 				InterruptContexts: contexts,
@@ -144,6 +152,12 @@ func CompositeInterrupt(ctx context.Context, info any, state any,
 	}
 }
 
+// CompositeInterrupt creates an interrupt event that aggregates sub-interrupt signals.
+func CompositeInterrupt(ctx context.Context, info any, state any,
+	subInterruptSignals ...*InterruptSignal) *AgentEvent {
+	return TypedCompositeInterrupt[*schema.Message](ctx, info, state, subInterruptSignals...)
+}
+
 // Address represents the unique, hierarchical address of a component within an execution.
 // It is a slice of AddressSegments, where each segment represents one level of nesting.
 // This is a type alias for core.Address. See the core package for more details.
@@ -202,9 +216,9 @@ type serialization struct {
 	InterruptID2State   map[string]core.InterruptState
 }
 
-func (r *Runner) loadCheckPoint(ctx context.Context, checkpointID string) (
+func runnerLoadCheckPointImpl(store CheckPointStore, ctx context.Context, checkpointID string) (
 	context.Context, *runContext, *ResumeInfo, error) {
-	data, existed, err := r.store.Get(ctx, checkpointID)
+	data, existed, err := store.Get(ctx, checkpointID)
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("failed to get checkpoint from store: %w", err)
 	}
@@ -266,13 +280,15 @@ func preprocessADKCheckpoint(data []byte) []byte {
 		[]byte(lenPrefixedCompatName))
 }
 
-func (r *Runner) saveCheckPoint(
+func runnerSaveCheckPointImpl(
+	enableStreaming bool,
+	store CheckPointStore,
 	ctx context.Context,
 	key string,
 	info *InterruptInfo,
 	is *core.InterruptSignal,
 ) error {
-	if r.store == nil {
+	if store == nil {
 		return nil
 	}
 
@@ -286,12 +302,12 @@ func (r *Runner) saveCheckPoint(
 		Info:                info,
 		InterruptID2Address: id2Addr,
 		InterruptID2State:   id2State,
-		EnableStreaming:     r.enableStreaming,
+		EnableStreaming:     enableStreaming,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to encode checkpoint: %w", err)
 	}
-	return r.store.Set(ctx, key, buf.Bytes())
+	return store.Set(ctx, key, buf.Bytes())
 }
 
 const bridgeCheckpointID = "adk_react_mock_key"

adk/react.go

@@ -31,14 +31,8 @@ import (
 // ErrExceedMaxIterations indicates the agent reached the maximum iterations limit.
 var ErrExceedMaxIterations = errors.New("exceeds max iterations")
 
-// State holds agent runtime state including messages and user-extensible storage.
-//
-// Deprecated: This type will be unexported in v1.0.0. Use ChatModelAgentState
-// in HandlerMiddleware and AgentMiddleware callbacks instead. Direct use of
-// compose.ProcessState[*State] is discouraged and will stop working in v1.0.0;
-// use the handler APIs instead.
-type State struct {
-	Messages []Message
+type typedState[M messageType] struct {
+	Messages []M
 	Extra    map[string]any
 
 	// Internal fields below - do not access directly.
@@ -48,10 +42,18 @@ type State struct {
 	ToolGenActions           map[string]*AgentAction
 	AgentName                string
 	RemainingIterations      int
-	ReturnDirectlyEvent      *AgentEvent
+	ReturnDirectlyEvent      *TypedAgentEvent[M]
 	RetryAttempt             int
 }
 
+// State is the internal state of the ChatModelAgent.
+//
+// Deprecated: State is exported only for checkpoint backward compatibility.
+// Do not use it directly.
+type State = typedState[*schema.Message]
+
+type agenticState = typedState[*schema.AgenticMessage]
+
 const (
 	stateGobNameV07 = "_eino_adk_react_state"
 
@@ -77,50 +79,57 @@ func init() {
 	schema.RegisterName[*State](stateGobNameV07)
 	schema.RegisterName[*stateV080](stateGobNameV080)
 
-	// the following two lines of registration mainly for backward compatibility
-	// when decoding checkpoints created by v0.8.0 - v0.8.3
+	schema.RegisterName[*typedState[*schema.AgenticMessage]]("_eino_adk_agentic_state")
+	schema.RegisterName[*TypedAgentEvent[*schema.AgenticMessage]]("_eino_adk_agentic_event")
+
+	// backward compatibility when decoding checkpoints created by v0.8.0 - v0.8.3
 	gob.Register(&AgentEvent{})
 	gob.Register(int(0))
+
+	schema.RegisterName[*TypedAgentInput[*schema.AgenticMessage]]("_eino_adk_agentic_agent_input")
+	schema.RegisterName[*typedAgentEventWrapper[*schema.AgenticMessage]]("_eino_adk_agentic_event_wrapper")
+	schema.RegisterName[*[]*typedAgentEventWrapper[*schema.AgenticMessage]]("_eino_adk_agentic_event_wrapper_slice")
 	schema.RegisterName[*reactInput]("_eino_adk_react_input")
+	schema.RegisterName[*agenticReactInput]("_eino_adk_agentic_react_input")
 }
 
-func (s *State) getReturnDirectlyEvent() *AgentEvent {
+func (s *typedState[M]) getReturnDirectlyEvent() *TypedAgentEvent[M] {
 	return s.ReturnDirectlyEvent
 }
 
-func (s *State) setReturnDirectlyEvent(event *AgentEvent) {
+func (s *typedState[M]) setReturnDirectlyEvent(event *TypedAgentEvent[M]) {
 	s.ReturnDirectlyEvent = event
 }
 
-func (s *State) getRetryAttempt() int {
+func (s *typedState[M]) getRetryAttempt() int {
 	return s.RetryAttempt
 }
 
-func (s *State) setRetryAttempt(attempt int) {
+func (s *typedState[M]) setRetryAttempt(attempt int) {
 	s.RetryAttempt = attempt
 }
 
-func (s *State) getReturnDirectlyToolCallID() string {
+func (s *typedState[M]) getReturnDirectlyToolCallID() string {
 	return s.ReturnDirectlyToolCallID
 }
 
-func (s *State) setReturnDirectlyToolCallID(id string) {
+func (s *typedState[M]) setReturnDirectlyToolCallID(id string) {
 	s.ReturnDirectlyToolCallID = id
 	s.HasReturnDirectly = id != ""
 }
 
-func (s *State) getToolGenActions() map[string]*AgentAction {
+func (s *typedState[M]) getToolGenActions() map[string]*AgentAction {
 	return s.ToolGenActions
 }
 
-func (s *State) setToolGenAction(key string, action *AgentAction) {
+func (s *typedState[M]) setToolGenAction(key string, action *AgentAction) {
 	if s.ToolGenActions == nil {
 		s.ToolGenActions = make(map[string]*AgentAction)
 	}
 	s.ToolGenActions[key] = action
 }
 
-func (s *State) popToolGenAction(key string) *AgentAction {
+func (s *typedState[M]) popToolGenAction(key string) *AgentAction {
 	if s.ToolGenActions == nil {
 		return nil
 	}
@@ -129,15 +138,15 @@ func (s *State) popToolGenAction(key string) *AgentAction {
 	return action
 }
 
-func (s *State) getRemainingIterations() int {
+func (s *typedState[M]) getRemainingIterations() int {
 	return s.RemainingIterations
 }
 
-func (s *State) setRemainingIterations(iterations int) {
+func (s *typedState[M]) setRemainingIterations(iterations int) {
 	s.RemainingIterations = iterations
 }
 
-func (s *State) decrementRemainingIterations() {
+func (s *typedState[M]) decrementRemainingIterations() {
 	current := s.getRemainingIterations()
 	s.RemainingIterations = current - 1
 }
@@ -241,13 +250,11 @@ type reactInput struct {
 	Messages []Message
 }
 
-type reactConfig struct {
-	// model is the chat model used by the react graph.
-	// Tools are configured via model.WithTools call option, not the WithTools method.
-	model model.BaseChatModel
+type typedReactConfig[M messageType] struct {
+	model model.BaseModel[M]
 
 	toolsConfig      *compose.ToolsNodeConfig
-	modelWrapperConf *modelWrapperConfig
+	modelWrapperConf *typedModelWrapperConfig[M]
 
 	toolsReturnDirectly map[string]bool
 
@@ -258,6 +265,8 @@ type reactConfig struct {
 	cancelCtx *cancelContext
 }
 
+type reactConfig = typedReactConfig[*schema.Message]
+
 func genToolInfos(ctx context.Context, config *compose.ToolsNodeConfig) ([]*schema.ToolInfo, error) {
 	toolInfos := make([]*schema.ToolInfo, 0, len(config.Tools))
 	for _, t := range config.Tools {
@@ -360,7 +369,7 @@ func newReact(ctx context.Context, config *reactConfig) (reactGraph, error) {
 	toolPreHandle := func(ctx context.Context, _ Message, st *State) (Message, error) {
 		input := st.Messages[len(st.Messages)-1]
 		returnDirectly := config.toolsReturnDirectly
-		if execCtx := getChatModelAgentExecCtx(ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 {
+		if execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 {
 			returnDirectly = execCtx.runtimeReturnDirectly
 		}
 		if len(returnDirectly) > 0 {
@@ -375,7 +384,7 @@ func newReact(ctx context.Context, config *reactConfig) (reactGraph, error) {
 	}
 	toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.Message], st *State) (*schema.StreamReader[[]*schema.Message], error) {
 		if event := st.getReturnDirectlyEvent(); event != nil {
-			getChatModelAgentExecCtx(ctx).send(event)
+			getTypedChatModelAgentExecCtx[*schema.Message](ctx).send(event)
 			st.setReturnDirectlyEvent(nil)
 		}
 		return out, nil
@@ -501,3 +510,218 @@ func newReact(ctx context.Context, config *reactConfig) (reactGraph, error) {
 
 	return g, nil
 }
+
+type agenticReactInput struct {
+	Messages []*schema.AgenticMessage
+}
+
+type agenticReactConfig = typedReactConfig[*schema.AgenticMessage]
+
+type agenticReactGraph = *compose.Graph[*agenticReactInput, *schema.AgenticMessage]
+
+func getAgenticReturnDirectlyToolCallID(ctx context.Context) (string, bool) {
+	var toolCallID string
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error {
+		toolCallID = st.getReturnDirectlyToolCallID()
+		return nil
+	})
+	return toolCallID, toolCallID != ""
+}
+
+func genAgenticReactState(config *agenticReactConfig) func(ctx context.Context) *agenticState {
+	return func(ctx context.Context) *agenticState {
+		st := &agenticState{
+			AgentName: config.agentName,
+		}
+		maxIter := 20
+		if config.maxIterations > 0 {
+			maxIter = config.maxIterations
+		}
+		st.setRemainingIterations(maxIter)
+		return st
+	}
+}
+
+func agenticMessageHasToolCalls(msg *schema.AgenticMessage) bool {
+	if msg == nil {
+		return false
+	}
+	for _, block := range msg.ContentBlocks {
+		if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil {
+			return true
+		}
+	}
+	return false
+}
+
+func newAgenticReact(ctx context.Context, config *agenticReactConfig) (agenticReactGraph, error) {
+	const (
+		initNode_                      = "Init"
+		chatModel_                     = "ChatModel"
+		cancelCheckNode_               = "CancelCheck"
+		toolNode_                      = "ToolNode"
+		afterToolCallsNode_            = "AfterToolCalls"
+		afterToolCallsCancelCheckNode_ = "AfterToolCallsCancelCheck"
+	)
+
+	cancelCtx := config.cancelCtx
+	g := compose.NewGraph[*agenticReactInput, *schema.AgenticMessage](
+		compose.WithGenLocalState(genAgenticReactState(config)))
+	_ = g.AddLambdaNode(initNode_, compose.InvokableLambda(func(ctx context.Context, input *agenticReactInput) ([]*schema.AgenticMessage, error) {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error {
+			st.Messages = append(st.Messages, input.Messages...)
+			return nil
+		})
+		return input.Messages, nil
+	}), compose.WithNodeName(initNode_))
+
+	var wrappedModel model.AgenticModel = config.model
+	if config.modelWrapperConf != nil {
+		wrappedModel = buildModelWrappers(config.model, config.modelWrapperConf)
+	}
+
+	toolsNode, err := compose.NewAgenticToolsNode(ctx, config.toolsConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	_ = g.AddAgenticModelNode(chatModel_, wrappedModel, compose.WithStatePreHandler(
+		func(ctx context.Context, input []*schema.AgenticMessage, st *agenticState) ([]*schema.AgenticMessage, error) {
+			if st.getRemainingIterations() <= 0 {
+				return nil, ErrExceedMaxIterations
+			}
+			st.decrementRemainingIterations()
+			return input, nil
+		}), compose.WithNodeName(chatModel_))
+
+	_ = g.AddLambdaNode(cancelCheckNode_, compose.InvokableLambda(func(ctx context.Context, msg *schema.AgenticMessage) (*schema.AgenticMessage, error) {
+		if cancelCtx != nil && cancelCtx.shouldCancel() {
+			if cancelCtx.getMode()&CancelAfterChatModel != 0 {
+				return nil, compose.StatefulInterrupt(ctx, "CancelAfterChatModel", msg)
+			}
+		}
+		wasInterrupted, hasState, state := compose.GetInterruptState[*schema.AgenticMessage](ctx)
+		if wasInterrupted && hasState {
+			msg = state
+		}
+		return msg, nil
+	}), compose.WithNodeName(cancelCheckNode_))
+
+	toolPreHandle := func(ctx context.Context, _ *schema.AgenticMessage, st *agenticState) (*schema.AgenticMessage, error) {
+		input := st.Messages[len(st.Messages)-1]
+		returnDirectly := config.toolsReturnDirectly
+		if execCtx := getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 {
+			returnDirectly = execCtx.runtimeReturnDirectly
+		}
+		if len(returnDirectly) > 0 {
+			for _, block := range input.ContentBlocks {
+				if block == nil || block.Type != schema.ContentBlockTypeFunctionToolCall || block.FunctionToolCall == nil {
+					continue
+				}
+				if _, ok := returnDirectly[block.FunctionToolCall.Name]; ok {
+					st.setReturnDirectlyToolCallID(block.FunctionToolCall.CallID)
+				}
+			}
+		}
+		return input, nil
+	}
+	toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.AgenticMessage], st *agenticState) (*schema.StreamReader[[]*schema.AgenticMessage], error) {
+		if event := st.getReturnDirectlyEvent(); event != nil {
+			getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx).send(event)
+			st.setReturnDirectlyEvent(nil)
+		}
+		return out, nil
+	}
+	_ = g.AddAgenticToolsNode(toolNode_, toolsNode,
+		compose.WithStatePreHandler(toolPreHandle),
+		compose.WithStreamStatePostHandler(toolPostHandle),
+		compose.WithNodeName(toolNode_))
+
+	afterToolCalls := func(ctx context.Context, toolResults []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error {
+			st.Messages = append(st.Messages, toolResults...)
+			return nil
+		})
+		return toolResults, nil
+	}
+	_ = g.AddLambdaNode(afterToolCallsNode_, compose.InvokableLambda(afterToolCalls),
+		compose.WithNodeName(afterToolCallsNode_))
+
+	afterToolCallsCancelCheck := func(ctx context.Context, toolResults []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) {
+		if cancelCtx != nil && cancelCtx.shouldCancel() {
+			if cancelCtx.getMode()&CancelAfterToolCalls != 0 {
+				return nil, compose.Interrupt(ctx, "CancelAfterToolCalls")
+			}
+		}
+		return toolResults, nil
+	}
+	_ = g.AddLambdaNode(afterToolCallsCancelCheckNode_, compose.InvokableLambda(afterToolCallsCancelCheck),
+		compose.WithNodeName(afterToolCallsCancelCheckNode_))
+
+	_ = g.AddEdge(compose.START, initNode_)
+	_ = g.AddEdge(initNode_, chatModel_)
+
+	toolCallCheck := func(ctx context.Context, sMsg *schema.StreamReader[*schema.AgenticMessage]) (string, error) {
+		defer sMsg.Close()
+		for {
+			chunk, err_ := sMsg.Recv()
+			if err_ != nil {
+				if err_ == io.EOF {
+					return compose.END, nil
+				}
+				return "", err_
+			}
+			if agenticMessageHasToolCalls(chunk) {
+				return cancelCheckNode_, nil
+			}
+		}
+	}
+	branch := compose.NewStreamGraphBranch(toolCallCheck, map[string]bool{compose.END: true, cancelCheckNode_: true})
+	_ = g.AddBranch(chatModel_, branch)
+
+	_ = g.AddEdge(cancelCheckNode_, toolNode_)
+	_ = g.AddEdge(toolNode_, afterToolCallsNode_)
+	_ = g.AddEdge(afterToolCallsNode_, afterToolCallsCancelCheckNode_)
+
+	if len(config.toolsReturnDirectly) > 0 {
+		const (
+			toolNodeToEndConverter = "ToolNodeToEndConverter"
+		)
+
+		cvt := func(ctx context.Context, toolResults []*schema.AgenticMessage) (*schema.AgenticMessage, error) {
+			id, _ := getAgenticReturnDirectlyToolCallID(ctx)
+			for _, msg := range toolResults {
+				if msg == nil {
+					continue
+				}
+				for _, block := range msg.ContentBlocks {
+					if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult &&
+						block.FunctionToolResult != nil && block.FunctionToolResult.CallID == id {
+						return msg, nil
+					}
+				}
+			}
+			return nil, errors.New("return directly tool call result not found")
+		}
+
+		_ = g.AddLambdaNode(toolNodeToEndConverter, compose.InvokableLambda(cvt),
+			compose.WithNodeName(toolNodeToEndConverter))
+		_ = g.AddEdge(toolNodeToEndConverter, compose.END)
+
+		checkReturnDirect := func(ctx context.Context, toolResults []*schema.AgenticMessage) (string, error) {
+			_, ok := getAgenticReturnDirectlyToolCallID(ctx)
+			if ok {
+				return toolNodeToEndConverter, nil
+			}
+			return chatModel_, nil
+		}
+
+		returnDirectBranch := compose.NewGraphBranch(checkReturnDirect,
+			map[string]bool{toolNodeToEndConverter: true, chatModel_: true})
+		_ = g.AddBranch(afterToolCallsCancelCheckNode_, returnDirectBranch)
+	} else {
+		_ = g.AddEdge(afterToolCallsCancelCheckNode_, chatModel_)
+	}
+
+	return g, nil
+}

adk/react_test.go

@@ -29,6 +29,7 @@ import (
 
 	"github.com/bytedance/sonic"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
 	"github.com/cloudwego/eino/components/model"
@@ -642,3 +643,30 @@ func randStrForTest() string {
 	}
 	return string(b)
 }
+
+func TestReactHistory_EmptyMessages(t *testing.T) {
+	g := compose.NewGraph[string, []Message](compose.WithGenLocalState(func(ctx context.Context) (state *State) {
+		return &State{
+			Messages: []Message{},
+		}
+	}))
+	require.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output []Message, err error) {
+		return getReactChatHistory(ctx, "DestAgent")
+	})))
+	require.NoError(t, g.AddEdge(compose.START, "1"))
+	require.NoError(t, g.AddEdge("1", compose.END))
+
+	ctx := context.Background()
+	ctx, _ = initRunCtx(ctx, "MyAgent", nil)
+	runner, err := g.Compile(ctx)
+	require.NoError(t, err)
+
+	require.NotPanics(t, func() {
+		result, err := runner.Invoke(ctx, "")
+		if err != nil {
+			t.Logf("Got error (acceptable): %v", err)
+			return
+		}
+		t.Logf("Got %d messages", len(result))
+	}, "BUG: getReactChatHistory should not panic with empty Messages slice")
+}

adk/retry_chatmodel.go

@@ -264,7 +264,7 @@ func genErrWrapper(ctx context.Context, maxRetries, attempt int, isRetryAbleFunc
 	}
 }
 
-func consumeStreamForError(stream *schema.StreamReader[*schema.Message]) error {
+func consumeStreamForError[M any](stream *schema.StreamReader[M]) error {
 	defer stream.Close()
 	for {
 		_, err := stream.Recv()
@@ -292,23 +292,27 @@ type retryVerdict struct {
 // This is used inside the model wrapper chain, positioned between eventSenderModelWrapper
 // and stateModelWrapper, so that retry only affects the inner chain (event sending, user wrappers,
 // callback injection) without re-running state management (BeforeModelRewriteState/AfterModelRewriteState).
-type retryModelWrapper struct {
-	inner  model.BaseChatModel
+type typedRetryModelWrapper[M messageType] struct {
+	inner  model.BaseModel[M]
 	config *ModelRetryConfig
 }
 
-func newRetryModelWrapper(inner model.BaseChatModel, config *ModelRetryConfig) *retryModelWrapper {
-	return &retryModelWrapper{inner: inner, config: config}
+func newTypedRetryModelWrapper[M messageType](inner model.BaseModel[M], config *ModelRetryConfig) *typedRetryModelWrapper[M] {
+	return &typedRetryModelWrapper[M]{inner: inner, config: config}
 }
 
-func (r *retryModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (r *typedRetryModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	if r.config.ShouldRetry != nil {
-		return r.generateWithShouldRetry(ctx, input, opts...)
+		// ShouldRetry is *schema.Message-specific (RetryContext.OutputMessage is *schema.Message).
+		msgR, _ := any(r).(*typedRetryModelWrapper[*schema.Message])
+		msgInput, _ := any(input).([]Message)
+		out, err := generateWithShouldRetry(msgR, ctx, msgInput, opts...)
+		return any(out).(M), err
 	}
 	return r.generateLegacy(ctx, input, opts...)
 }
 
-func (r *retryModelWrapper) generateLegacy(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []M, opts ...model.Option) (zero M, _ error) {
 	isRetryAble := r.config.IsRetryAble
 	if isRetryAble == nil {
 		isRetryAble = defaultIsRetryAble
@@ -325,37 +329,36 @@ func (r *retryModelWrapper) generateLegacy(ctx context.Context, input []*schema.
 			return out, nil
 		}
 
-		// Never retry interrupt errors (e.g. cancel safe-point interrupts).
 		if _, ok := compose.ExtractInterruptInfo(err); ok {
-			return nil, err
+			return zero, err
 		}
 
 		if errors.Is(err, ErrStreamCanceled) {
-			return nil, err
+			return zero, err
 		}
 
 		if !isRetryAble(ctx, err) {
-			return nil, err
+			return zero, err
 		}
 
 		lastErr = err
 		if attempt < r.config.MaxRetries {
 			if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil {
-				return nil, err
+				return zero, err
 			}
 		}
 	}
 
-	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
+	return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
 }
 
-func (r *retryModelWrapper) generateWithShouldRetry(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func generateWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
 	backoffFunc := r.config.BackoffFunc
 	if backoffFunc == nil {
 		backoffFunc = defaultBackoff
 	}
 
-	execCtx := getChatModelAgentExecCtx(ctx)
+	execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 
 	currentInput := input
 	currentOpts := opts
@@ -431,7 +434,7 @@ func (r *retryModelWrapper) generateWithShouldRetry(ctx context.Context, input [
 			break
 		}
 
-		r.applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
+		applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 
 		delay := decision.Backoff
 		if delay == 0 {
@@ -446,7 +449,7 @@ func (r *retryModelWrapper) generateWithShouldRetry(ctx context.Context, input [
 	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
 }
 
-func (r *retryModelWrapper) contextAwareSleep(ctx context.Context, delay time.Duration) error {
+func (r *typedRetryModelWrapper[M]) contextAwareSleep(ctx context.Context, delay time.Duration) error {
 	if delay <= 0 {
 		return nil
 	}
@@ -481,7 +484,7 @@ func consumeStreamForMessage(stream *schema.StreamReader[*schema.Message]) (*sch
 	}
 }
 
-func (r *retryModelWrapper) streamWithShouldRetry(ctx context.Context, input []*schema.Message, opts ...model.Option) (
+func streamWithShouldRetry(r *typedRetryModelWrapper[*schema.Message], ctx context.Context, input []*schema.Message, opts ...model.Option) (
 	*schema.StreamReader[*schema.Message], error) {
 
 	backoffFunc := r.config.BackoffFunc
@@ -496,7 +499,7 @@ func (r *retryModelWrapper) streamWithShouldRetry(ctx context.Context, input []*
 		})
 	}()
 
-	execCtx := getChatModelAgentExecCtx(ctx)
+	execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 
 	currentInput := input
 	currentOpts := opts
@@ -568,7 +571,7 @@ func (r *retryModelWrapper) streamWithShouldRetry(ctx context.Context, input []*
 
 			lastErr = err
 			if attempt < r.config.MaxRetries {
-				r.applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
+				applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 				delay := decision.Backoff
 				if delay == 0 {
 					delay = backoffFunc(ctx, attempt+1)
@@ -638,7 +641,7 @@ func (r *retryModelWrapper) streamWithShouldRetry(ctx context.Context, input []*
 		lastErr = verdictErr
 
 		if attempt < r.config.MaxRetries {
-			r.applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
+			applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
 			delay := decision.Backoff
 			if delay == 0 {
 				delay = backoffFunc(ctx, attempt+1)
@@ -652,7 +655,7 @@ func (r *retryModelWrapper) streamWithShouldRetry(ctx context.Context, input []*
 	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
 }
 
-func (r *retryModelWrapper) applyDecisionForRetry(currentInput *[]*schema.Message, currentOpts *[]model.Option, ctx context.Context, decision *RetryDecision) {
+func applyDecisionForRetry(currentInput *[]*schema.Message, currentOpts *[]model.Option, ctx context.Context, decision *RetryDecision) {
 	if decision.ModifiedInputMessages != nil {
 		*currentInput = decision.ModifiedInputMessages
 		if decision.PersistModifiedInputMessages {
@@ -671,17 +674,24 @@ func (r *retryModelWrapper) applyDecisionForRetry(currentInput *[]*schema.Messag
 	}
 }
 
-func (r *retryModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (
-	*schema.StreamReader[*schema.Message], error) {
+func (r *typedRetryModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (
+	*schema.StreamReader[M], error) {
 
 	if r.config.ShouldRetry != nil {
-		return r.streamWithShouldRetry(ctx, input, opts...)
+		// ShouldRetry is *schema.Message-specific (RetryContext.OutputMessage is *schema.Message).
+		msgR, _ := any(r).(*typedRetryModelWrapper[*schema.Message])
+		msgInput, _ := any(input).([]Message)
+		sr, err := streamWithShouldRetry(msgR, ctx, msgInput, opts...)
+		if err != nil {
+			return nil, err
+		}
+		return any(sr).(*schema.StreamReader[M]), nil
 	}
 	return r.streamLegacy(ctx, input, opts...)
 }
 
-func (r *retryModelWrapper) streamLegacy(ctx context.Context, input []*schema.Message, opts ...model.Option) (
-	*schema.StreamReader[*schema.Message], error) {
+func (r *typedRetryModelWrapper[M]) streamLegacy(ctx context.Context, input []M, opts ...model.Option) (
+	*schema.StreamReader[M], error) {
 
 	isRetryAble := r.config.IsRetryAble
 	if isRetryAble == nil {
@@ -693,7 +703,7 @@ func (r *retryModelWrapper) streamLegacy(ctx context.Context, input []*schema.Me
 	}
 
 	defer func() {
-		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 			st.setRetryAttempt(0)
 			return nil
 		})
@@ -701,7 +711,7 @@ func (r *retryModelWrapper) streamLegacy(ctx context.Context, input []*schema.Me
 
 	var lastErr error
 	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
-		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 			st.setRetryAttempt(attempt)
 			return nil
 		})
@@ -730,7 +740,7 @@ func (r *retryModelWrapper) streamLegacy(ctx context.Context, input []*schema.Me
 		checkCopy := copies[0]
 		returnCopy := copies[1]
 
-		streamErr := consumeStreamForError(checkCopy)
+		streamErr := consumeStreamForError[M](checkCopy)
 		if streamErr == nil {
 			return returnCopy, nil
 		}

adk/runctx.go

@@ -20,10 +20,14 @@ import (
 	"bytes"
 	"context"
 	"encoding/gob"
+	"errors"
 	"fmt"
+	"io"
 	"sort"
 	"sync"
 	"time"
+
+	"github.com/cloudwego/eino/schema"
 )
 
 // runSession CheckpointSchema: persisted via serialization.RunCtx (gob).
@@ -34,6 +38,11 @@ type runSession struct {
 	Events     []*agentEventWrapper
 	LaneEvents *laneEvents
 	mtx        sync.Mutex
+
+	// TypedEvents stores *[]*typedAgentEventWrapper[M] for M != *schema.Message.
+	// For M = *schema.Message, the existing Events field is used instead.
+	// The any type is required because Go does not support generic fields in non-generic structs.
+	TypedEvents any
 }
 
 // laneEvents CheckpointSchema: persisted via serialization.RunCtx (gob).
@@ -60,6 +69,105 @@ type agentEventWrapper struct {
 	StreamErr error
 }
 
+type typedAgentEventWrapper[M messageType] struct {
+	event               *TypedAgentEvent[M]
+	mu                  sync.Mutex
+	concatenatedMessage M
+	TS                  int64
+	StreamErr           error
+}
+
+// typedAgentEventWrapperForGob is a gob-serializable representation of typedAgentEventWrapper.
+// We encode the event and TS separately to avoid the sync.Mutex and non-exported fields.
+type typedAgentEventWrapperForGob[M messageType] struct {
+	Event *TypedAgentEvent[M]
+	TS    int64
+}
+
+func (e *typedAgentEventWrapper[M]) GobEncode() ([]byte, error) {
+	if e.event != nil && e.event.Output != nil && e.event.Output.MessageOutput != nil && e.event.Output.MessageOutput.IsStreaming {
+		// Materialize the stream before encoding.
+		if isNilMessage(e.concatenatedMessage) && e.StreamErr == nil {
+			e.consumeStream()
+		}
+	}
+
+	buf := &bytes.Buffer{}
+	err := gob.NewEncoder(buf).Encode(&typedAgentEventWrapperForGob[M]{
+		Event: e.event,
+		TS:    e.TS,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to gob encode generic agent event wrapper: %w", err)
+	}
+	return buf.Bytes(), nil
+}
+
+func (e *typedAgentEventWrapper[M]) GobDecode(b []byte) error {
+	g := &typedAgentEventWrapperForGob[M]{}
+	if err := gob.NewDecoder(bytes.NewReader(b)).Decode(g); err != nil {
+		return fmt.Errorf("failed to gob decode generic agent event wrapper: %w", err)
+	}
+	e.event = g.Event
+	e.TS = g.TS
+	return nil
+}
+
+// consumeStream drains the typed message stream, setting concatenatedMessage on success
+// or StreamErr on failure. The stream is replaced with a materialized version safe for
+// gob encoding.
+//
+// NOTE: This method parallels agentEventWrapper.consumeStream in utils.go. The two
+// implementations exist because agentEventWrapper is non-generic (uses *schema.Message
+// directly) while typedAgentEventWrapper[M] is generic. They cannot be unified without
+// making the non-generic wrapper generic, which would cascade through the entire
+// non-generic event storage layer.
+func (e *typedAgentEventWrapper[M]) consumeStream() {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if !isNilMessage(e.concatenatedMessage) {
+		return
+	}
+
+	s := e.event.Output.MessageOutput.MessageStream
+	var msgs []M
+
+	defer s.Close()
+	for {
+		msg, err := s.Recv()
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			e.StreamErr = err
+			e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs)
+			return
+		}
+		msgs = append(msgs, msg)
+	}
+
+	if len(msgs) == 0 {
+		e.StreamErr = errors.New("no messages in typedAgentEventWrapper.MessageStream")
+		e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs)
+		return
+	}
+
+	if len(msgs) == 1 {
+		e.concatenatedMessage = msgs[0]
+	} else {
+		var err error
+		e.concatenatedMessage, err = concatMessageStream(schema.StreamReaderFromArray(msgs))
+		if err != nil {
+			e.StreamErr = err
+			e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs)
+			return
+		}
+	}
+
+	e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray([]M{e.concatenatedMessage})
+}
+
 type otherAgentEventWrapperForEncode agentEventWrapper
 
 func (a *agentEventWrapper) GobEncode() ([]byte, error) {
@@ -184,6 +292,71 @@ func (rs *runSession) getEvents() []*agentEventWrapper {
 	return finalEvents
 }
 
+func addTypedEvent[M messageType](session *runSession, event *TypedAgentEvent[M]) {
+	var zero M
+	if _, ok := any(zero).(*schema.Message); ok {
+		session.addEvent(any(event).(*AgentEvent))
+		return
+	}
+	session.mtx.Lock()
+	defer session.mtx.Unlock()
+	wrapper := &typedAgentEventWrapper[M]{event: event, TS: time.Now().UnixNano()}
+	store, _ := session.TypedEvents.(*[]*typedAgentEventWrapper[M])
+	if store == nil {
+		s := make([]*typedAgentEventWrapper[M], 0)
+		store = &s
+		session.TypedEvents = store
+	}
+	*store = append(*store, wrapper)
+}
+
+func getTypedEvents[M messageType](session *runSession) []*typedAgentEventWrapper[M] {
+	var zero M
+	if _, ok := any(zero).(*schema.Message); ok {
+		events := session.getEvents()
+		result := make([]*typedAgentEventWrapper[M], 0, len(events))
+		for _, e := range events {
+			w := &typedAgentEventWrapper[M]{
+				event:     any(e.AgentEvent).(*TypedAgentEvent[M]),
+				TS:        e.TS,
+				StreamErr: e.StreamErr,
+			}
+			if e.concatenatedMessage != nil {
+				w.concatenatedMessage = any(e.concatenatedMessage).(M)
+			}
+			result = append(result, w)
+		}
+		return result
+	}
+
+	session.mtx.Lock()
+	defer session.mtx.Unlock()
+
+	store, _ := session.TypedEvents.(*[]*typedAgentEventWrapper[M])
+	if store == nil {
+		if len(session.Events) == 0 {
+			return nil
+		}
+		result := make([]*typedAgentEventWrapper[M], 0, len(session.Events))
+		for _, e := range session.Events {
+			w := &typedAgentEventWrapper[M]{
+				event:     any(e.AgentEvent).(*TypedAgentEvent[M]),
+				TS:        e.TS,
+				StreamErr: e.StreamErr,
+			}
+			if e.concatenatedMessage != nil {
+				w.concatenatedMessage = any(e.concatenatedMessage).(M)
+			}
+			result = append(result, w)
+		}
+		return result
+	}
+
+	result := make([]*typedAgentEventWrapper[M], len(*store))
+	copy(result, *store)
+	return result
+}
+
 func (rs *runSession) getValues() map[string]any {
 	rs.valuesMtx.Lock()
 	values := make(map[string]any, len(rs.Values))
@@ -221,6 +394,8 @@ type runContext struct {
 	RootInput *AgentInput
 	RunPath   []RunStep
 
+	AgenticRootInput any
+
 	Session *runSession
 }
 
@@ -230,9 +405,10 @@ func (rc *runContext) isRoot() bool {
 
 func (rc *runContext) deepCopy() *runContext {
 	copied := &runContext{
-		RootInput: rc.RootInput,
-		RunPath:   make([]RunStep, len(rc.RunPath)),
-		Session:   rc.Session,
+		RootInput:        rc.RootInput,
+		AgenticRootInput: rc.AgenticRootInput,
+		RunPath:          make([]RunStep, len(rc.RunPath)),
+		Session:          rc.Session,
 	}
 
 	copy(copied.RunPath, rc.RunPath)
@@ -270,6 +446,27 @@ func initRunCtx(ctx context.Context, agentName string, input *AgentInput) (conte
 	return setRunCtx(ctx, runCtx), runCtx
 }
 
+func initTypedRunCtx[M messageType](ctx context.Context, agentName string, input *TypedAgentInput[M]) (context.Context, *runContext) {
+	runCtx := getRunCtx(ctx)
+	if runCtx != nil {
+		runCtx = runCtx.deepCopy()
+	} else {
+		runCtx = &runContext{Session: newRunSession()}
+	}
+
+	runCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName})
+	if runCtx.isRoot() && input != nil {
+		var zero M
+		if _, ok := any(zero).(*schema.Message); ok {
+			runCtx.RootInput = any(input).(*AgentInput)
+		} else {
+			runCtx.AgenticRootInput = input
+		}
+	}
+
+	return setRunCtx(ctx, runCtx), runCtx
+}
+
 func joinRunCtxs(parentCtx context.Context, childCtxs ...context.Context) {
 	switch len(childCtxs) {
 	case 0:
@@ -384,7 +581,7 @@ func ClearRunCtx(ctx context.Context) context.Context {
 	return context.WithValue(ctx, runCtxKey{}, nil)
 }
 
-func ctxWithNewRunCtx(ctx context.Context, input *AgentInput, sharedParentSession bool) context.Context {
+func ctxWithNewTypedRunCtx[M messageType](ctx context.Context, input *TypedAgentInput[M], sharedParentSession bool) context.Context {
 	var session *runSession
 	if sharedParentSession {
 		if parentSession := getSession(ctx); parentSession != nil {
@@ -397,7 +594,14 @@ func ctxWithNewRunCtx(ctx context.Context, input *AgentInput, sharedParentSessio
 	if session == nil {
 		session = newRunSession()
 	}
-	return setRunCtx(ctx, &runContext{Session: session, RootInput: input})
+	var zero M
+	rc := &runContext{Session: session}
+	if _, ok := any(zero).(*schema.Message); ok {
+		rc.RootInput = any(input).(*AgentInput)
+	} else {
+		rc.AgenticRootInput = input
+	}
+	return setRunCtx(ctx, rc)
 }
 
 func getSession(ctx context.Context) *runSession {

adk/runner.go

@@ -28,29 +28,53 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-// Runner is the primary entry point for executing an Agent.
-// It manages the agent's lifecycle, including starting, resuming, and checkpointing.
-type Runner struct {
-	// a is the agent to be executed.
-	a Agent
-	// enableStreaming dictates whether the execution should be in streaming mode.
-	enableStreaming bool
-	// store is the checkpoint store used to persist agent state upon interruption.
-	// If nil, checkpointing is disabled.
-	store CheckPointStore
+func errorIterator[M messageType](err error) *AsyncIterator[*TypedAgentEvent[M]] {
+	iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+	gen.Send(&TypedAgentEvent[M]{Err: err})
+	gen.Close()
+	return iter
 }
 
+func newUserMessage[M messageType](query string) (M, error) {
+	var zero M
+	switch any(zero).(type) {
+	case *schema.Message:
+		return any(schema.UserMessage(query)).(M), nil
+	case *schema.AgenticMessage:
+		return any(schema.UserAgenticMessage(query)).(M), nil
+	default:
+		return zero, fmt.Errorf("unsupported message type %T", zero)
+	}
+}
+
+// TypedRunner is the primary entry point for executing an Agent.
+// It manages the agent's lifecycle, including starting, resuming, and checkpointing.
+//
+// Execution always goes through the flowAgent pipeline, which handles
+// multi-agent orchestration, callbacks, agent naming, run paths, and cancellation.
+type TypedRunner[M messageType] struct {
+	a               TypedAgent[M]
+	enableStreaming bool
+	store           CheckPointStore
+}
+
+// Runner is the default runner type using *schema.Message.
+type Runner = TypedRunner[*schema.Message]
+
 type CheckPointStore = core.CheckPointStore
 
 type CheckPointDeleter = core.CheckPointDeleter
 
-type RunnerConfig struct {
-	Agent           Agent
+type TypedRunnerConfig[M messageType] struct {
+	Agent           TypedAgent[M]
 	EnableStreaming bool
 
 	CheckPointStore CheckPointStore
 }
 
+// RunnerConfig is the default runner config type using *schema.Message.
+type RunnerConfig = TypedRunnerConfig[*schema.Message]
+
 // ResumeParams contains all parameters needed to resume an execution.
 // This struct provides an extensible way to pass resume parameters without
 // requiring breaking changes to method signatures.
@@ -61,52 +85,33 @@ type ResumeParams struct {
 	// Future extensible fields can be added here without breaking changes
 }
 
-// NewRunner creates a Runner that executes an Agent with optional streaming
-// and checkpoint persistence.
+// NewRunner creates a new Runner with the given config.
 func NewRunner(_ context.Context, conf RunnerConfig) *Runner {
-	return &Runner{
+	return NewTypedRunner[*schema.Message](conf)
+}
+
+// NewTypedRunner creates a new TypedRunner with the given config.
+func NewTypedRunner[M messageType](conf TypedRunnerConfig[M]) *TypedRunner[M] {
+	return &TypedRunner[M]{
 		enableStreaming: conf.EnableStreaming,
 		a:               conf.Agent,
 		store:           conf.CheckPointStore,
 	}
 }
 
-// Run starts a new execution of the agent with a given set of messages.
-// It returns an iterator that yields agent events as they occur.
-// If the Runner was configured with a CheckPointStore, it will automatically save the agent's state
-// upon interruption.
-func (r *Runner) Run(ctx context.Context, messages []Message,
-	opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
-	o := getCommonOptions(nil, opts...)
-
-	fa := toFlowAgent(ctx, r.a)
-
-	input := &AgentInput{
-		Messages:        messages,
-		EnableStreaming: r.enableStreaming,
-	}
-
-	ctx = ctxWithNewRunCtx(ctx, input, o.sharedParentSession)
-
-	AddSessionValues(ctx, o.sessionValues)
-
-	iter := fa.Run(ctx, input, opts...)
-
-	if r.store == nil && o.cancelCtx == nil {
-		return iter
-	}
-
-	niter, gen := NewAsyncIteratorPair[*AgentEvent]()
-
-	go r.handleIter(ctx, iter, gen, o.checkPointID, o.cancelCtx)
-	return niter
+func (r *TypedRunner[M]) Run(ctx context.Context, messages []M,
+	opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	return typedRunnerRunImpl(r.a, r.enableStreaming, r.store, ctx, messages, opts...)
 }
 
 // Query is a convenience method that starts a new execution with a single user query string.
-func (r *Runner) Query(ctx context.Context,
-	query string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
-
-	return r.Run(ctx, []Message{schema.UserMessage(query)}, opts...)
+func (r *TypedRunner[M]) Query(ctx context.Context,
+	query string, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	msgs, err := newUserMessage[M](query)
+	if err != nil {
+		return errorIterator[M](err)
+	}
+	return r.Run(ctx, []M{msgs}, opts...)
 }
 
 // Resume continues an interrupted execution from a checkpoint, using an "Implicit Resume All" strategy.
@@ -116,8 +121,8 @@ func (r *Runner) Query(ctx context.Context,
 // When using this method, all interrupted agents will receive `isResumeFlow = false` when they
 // call `GetResumeContext`, as no specific agent was targeted. This is suitable for the "Simple Confirmation"
 // pattern where an agent only needs to know `wasInterrupted` is true to continue.
-func (r *Runner) Resume(ctx context.Context, checkPointID string, opts ...AgentRunOption) (
-	*AsyncIterator[*AgentEvent], error) {
+func (r *TypedRunner[M]) Resume(ctx context.Context, checkPointID string, opts ...AgentRunOption) (
+	*AsyncIterator[*TypedAgentEvent[M]], error) {
 	return r.resumeInternal(ctx, checkPointID, nil, opts...)
 }
 
@@ -139,17 +144,71 @@ func (r *Runner) Resume(ctx context.Context, checkPointID string, opts ...AgentR
 //     execution. They act as conduits, allowing the resume signal to flow to their children. They will
 //     naturally re-interrupt if one of their interrupted children re-interrupts, as they receive the
 //     new `CompositeInterrupt` signal from them.
-func (r *Runner) ResumeWithParams(ctx context.Context, checkPointID string, params *ResumeParams, opts ...AgentRunOption) (*AsyncIterator[*AgentEvent], error) {
+func (r *TypedRunner[M]) ResumeWithParams(ctx context.Context, checkPointID string, params *ResumeParams, opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) {
 	return r.resumeInternal(ctx, checkPointID, params.Targets, opts...)
 }
 
-func (r *Runner) resumeInternal(ctx context.Context, checkPointID string, resumeData map[string]any,
-	opts ...AgentRunOption) (*AsyncIterator[*AgentEvent], error) {
-	if r.store == nil {
+func (r *TypedRunner[M]) resumeInternal(ctx context.Context, checkPointID string, resumeData map[string]any,
+	opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) {
+	return typedRunnerResumeInternalImpl(r.a, r.enableStreaming, r.store, ctx, checkPointID, resumeData, opts...)
+}
+
+func typedRunnerRunImpl[M messageType](a TypedAgent[M], enableStreaming bool, store CheckPointStore, ctx context.Context, messages []M, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] {
+	o := getCommonOptions(nil, opts...)
+
+	input := &TypedAgentInput[M]{
+		Messages:        messages,
+		EnableStreaming: enableStreaming,
+	}
+
+	var zero M
+	if _, ok := any(zero).(*schema.Message); ok {
+		concreteAgent, _ := any(a).(Agent)
+		fa := toFlowAgent(ctx, concreteAgent)
+		if store != nil {
+			fa.checkPointStore = store
+		}
+		concreteInput := any(input).(*AgentInput)
+		ctx = ctxWithNewTypedRunCtx(ctx, input, o.sharedParentSession)
+		AddSessionValues(ctx, o.sessionValues)
+
+		iter := fa.Run(ctx, concreteInput, opts...)
+
+		if store == nil && o.cancelCtx == nil {
+			return any(iter).(*AsyncIterator[*TypedAgentEvent[M]])
+		}
+
+		niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+		go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(iter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, o.checkPointID, o.cancelCtx)
+		return niter
+	}
+
+	fa := toTypedFlowAgent(a)
+	if store != nil {
+		fa.checkPointStore = store
+	}
+
+	ctx = ctxWithNewTypedRunCtx(ctx, input, o.sharedParentSession)
+	AddSessionValues(ctx, o.sessionValues)
+
+	iter := fa.Run(ctx, input, opts...)
+
+	if store == nil && o.cancelCtx == nil {
+		return iter
+	}
+
+	niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+	go typedRunnerHandleIterImpl(enableStreaming, store, ctx, iter, gen, o.checkPointID, o.cancelCtx)
+	return niter
+}
+
+func typedRunnerResumeInternalImpl[M messageType](a TypedAgent[M], enableStreaming bool, store CheckPointStore, ctx context.Context, checkPointID string, resumeData map[string]any, //nolint:revive // argument-limit
+	opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) {
+	if store == nil {
 		return nil, fmt.Errorf("failed to resume: store is nil")
 	}
 
-	ctx, runCtx, resumeInfo, err := r.loadCheckPoint(ctx, checkPointID)
+	ctx, runCtx, resumeInfo, err := runnerLoadCheckPointImpl(store, ctx, checkPointID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to load from checkpoint: %w", err)
 	}
@@ -170,34 +229,46 @@ func (r *Runner) resumeInternal(ctx context.Context, checkPointID string, resume
 	}
 
 	ctx = setRunCtx(ctx, runCtx)
-
 	AddSessionValues(ctx, o.sessionValues)
 
 	if len(resumeData) > 0 {
 		ctx = core.BatchResumeWithData(ctx, resumeData)
 	}
 
-	fa := toFlowAgent(ctx, r.a)
+	var zero M
+	if _, ok := any(zero).(*schema.Message); ok {
+		concreteAgent, _ := any(a).(Agent)
+		fa := toFlowAgent(ctx, concreteAgent)
+		ra, ok := Agent(fa).(ResumableAgent)
+		if !ok {
+			return nil, fmt.Errorf("agent %T does not support resume", a)
+		}
+		aIter := ra.Resume(ctx, resumeInfo, opts...)
 
-	aIter := fa.Resume(ctx, resumeInfo, opts...)
-
-	if r.store == nil && o.cancelCtx == nil {
-		return aIter, nil
+		niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+		go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(aIter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, &checkPointID, o.cancelCtx)
+		return niter, nil
 	}
 
-	niter, gen := NewAsyncIteratorPair[*AgentEvent]()
+	fa := toTypedFlowAgent(a)
+	ra, ok := TypedAgent[M](fa).(TypedResumableAgent[M])
+	if !ok {
+		return nil, fmt.Errorf("agent %T does not support resume", a)
+	}
+	aIter := ra.Resume(ctx, resumeInfo, opts...)
 
-	go r.handleIter(ctx, aIter, gen, &checkPointID, o.cancelCtx)
+	niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+	go typedRunnerHandleIterImpl(enableStreaming, store, ctx, aIter, gen, &checkPointID, o.cancelCtx)
 	return niter, nil
 }
 
-func (r *Runner) handleIter(ctx context.Context, aIter *AsyncIterator[*AgentEvent],
-	gen *AsyncGenerator[*AgentEvent], checkPointID *string, cancelCtx *cancelContext) {
+func typedRunnerHandleIterImpl[M messageType](enableStreaming bool, store CheckPointStore, ctx context.Context, aIter *AsyncIterator[*TypedAgentEvent[M]], //nolint:revive // argument-limit
+	gen *AsyncGenerator[*TypedAgentEvent[M]], checkPointID *string, cancelCtx *cancelContext) {
 	defer func() {
 		panicErr := recover()
 		if panicErr != nil {
 			e := safe.NewPanicErr(panicErr, debug.Stack())
-			gen.Send(&AgentEvent{Err: e})
+			gen.Send(&TypedAgentEvent[M]{Err: e})
 		}
 
 		gen.Close()
@@ -220,9 +291,9 @@ func (r *Runner) handleIter(ctx context.Context, aIter *AsyncIterator[*AgentEven
 				}
 				if cancelErr.interruptSignal != nil && checkPointID != nil {
 					cancelErr.InterruptContexts = core.ToInterruptContexts(cancelErr.interruptSignal, allowedAddressSegmentTypes)
-					err := r.saveCheckPoint(ctx, *checkPointID, &InterruptInfo{}, cancelErr.interruptSignal)
+					err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{}, cancelErr.interruptSignal)
 					if err != nil {
-						gen.Send(&AgentEvent{Err: fmt.Errorf("failed to save checkpoint on cancel: %w", err)})
+						gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint on cancel: %w", err)})
 					}
 				}
 				gen.Send(event)
@@ -232,14 +303,11 @@ func (r *Runner) handleIter(ctx context.Context, aIter *AsyncIterator[*AgentEven
 
 		if event.Action != nil && event.Action.internalInterrupted != nil {
 			if interruptSignal != nil {
-				// even if multiple interrupt happens, they should be merged into one
-				// action by CompositeInterrupt, so here in Runner we must assume at most
-				// one interrupt action happens
 				panic("multiple interrupt actions should not happen in Runner")
 			}
 			interruptSignal = event.Action.internalInterrupted
 			interruptContexts := core.ToInterruptContexts(interruptSignal, allowedAddressSegmentTypes)
-			event = &AgentEvent{
+			event = &TypedAgentEvent[M]{
 				AgentName: event.AgentName,
 				RunPath:   event.RunPath,
 				Output:    event.Output,
@@ -254,12 +322,11 @@ func (r *Runner) handleIter(ctx context.Context, aIter *AsyncIterator[*AgentEven
 			legacyData = event.Action.Interrupted.Data
 
 			if checkPointID != nil {
-				// save checkpoint first before sending interrupt event, so when end-user receives interrupt event, they can resume from this checkpoint
-				err := r.saveCheckPoint(ctx, *checkPointID, &InterruptInfo{
+				err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{
 					Data: legacyData,
 				}, interruptSignal)
 				if err != nil {
-					gen.Send(&AgentEvent{Err: fmt.Errorf("failed to save checkpoint: %w", err)})
+					gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint: %w", err)})
 				}
 			}
 		}

adk/runner_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/cloudwego/eino/schema"
 )
@@ -261,3 +262,50 @@ func TestRunner_Query_WithStreaming(t *testing.T) {
 	_, ok = iterator.Next()
 	assert.False(t, ok)
 }
+
+func TestResumeWithMissingCheckpoint(t *testing.T) {
+	ctx := context.Background()
+
+	agent := &myAgenticAgent{
+		name: "resume-agent",
+		runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] {
+			iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]()
+			go func() {
+				defer gen.Close()
+				gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{
+					Output: &TypedAgentOutput[*schema.AgenticMessage]{
+						MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{
+							Message: agenticMsg("ok"),
+						},
+					},
+				})
+			}()
+			return iter
+		},
+	}
+
+	store := newMyStore()
+	runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{
+		Agent:           agent,
+		CheckPointStore: store,
+	})
+
+	require.NotPanics(t, func() {
+		iter, err := runner.ResumeWithParams(ctx, "nonexistent-checkpoint", &ResumeParams{
+			Targets: map[string]any{"fake-id": nil},
+		})
+		if err != nil {
+			t.Logf("Got expected error: %v", err)
+			return
+		}
+		for {
+			event, ok := iter.Next()
+			if !ok {
+				break
+			}
+			if event.Err != nil {
+				t.Logf("Got error event: %v", event.Err)
+			}
+		}
+	}, "ResumeWithParams with nonexistent checkpoint should not panic")
+}

adk/turn_loop.go

@@ -356,12 +356,12 @@ func (s *preemptSignal) drainAll() {
 }
 
 // TurnLoopConfig is the configuration for creating a TurnLoop.
-type TurnLoopConfig[T any] struct {
+type TurnLoopConfig[T any, M messageType] struct {
 	// GenInput receives the TurnLoop instance and all buffered items, and decides what to process.
 	// It returns which items to consume now vs keep for later turns.
 	// The loop parameter allows calling Push() or Stop() directly from within the callback.
 	// Required.
-	GenInput func(ctx context.Context, loop *TurnLoop[T], items []T) (*GenInputResult[T], error)
+	GenInput func(ctx context.Context, loop *TurnLoop[T, M], items []T) (*GenInputResult[T, M], error)
 
 	// GenResume is called at most once during Run(). When CheckpointID is
 	// configured, Run() queries Store for the checkpoint:
@@ -378,7 +378,7 @@ type TurnLoopConfig[T any] struct {
 	// It returns a GenResumeResult describing how to resume the interrupted agent
 	// turn (optional ResumeParams) and how to manipulate the buffer
 	// (Consumed/Remaining) before continuing.
-	GenResume func(ctx context.Context, loop *TurnLoop[T], canceledItems, unhandledItems, newItems []T) (*GenResumeResult[T], error)
+	GenResume func(ctx context.Context, loop *TurnLoop[T, M], canceledItems, unhandledItems, newItems []T) (*GenResumeResult[T, M], error)
 
 	// PrepareAgent returns an Agent configured to handle the consumed items.
 	// This callback should set up the agent with appropriate system prompt,
@@ -386,7 +386,7 @@ type TurnLoopConfig[T any] struct {
 	// Called once per turn with the items that GenInput decided to consume.
 	// The loop parameter allows calling Push() or Stop() directly from within the callback.
 	// Required.
-	PrepareAgent func(ctx context.Context, loop *TurnLoop[T], consumed []T) (Agent, error)
+	PrepareAgent func(ctx context.Context, loop *TurnLoop[T, M], consumed []T) (TypedAgent[M], error)
 
 	// OnAgentEvents is called to handle events emitted by the agent.
 	// The TurnContext provides per-turn info and control:
@@ -405,7 +405,7 @@ type TurnLoopConfig[T any] struct {
 	//
 	// Optional. If not provided, events are drained and the first error
 	// (including CancelError from Stop) is returned as ExitReason.
-	OnAgentEvents func(ctx context.Context, tc *TurnContext[T], events *AsyncIterator[*AgentEvent]) error
+	OnAgentEvents func(ctx context.Context, tc *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error
 
 	// Store is the checkpoint store for persistence and resume. Optional.
 	// When set together with CheckpointID, enables automatic checkpoint-based resume.
@@ -430,7 +430,7 @@ type TurnLoopConfig[T any] struct {
 }
 
 // GenInputResult contains the result of GenInput processing.
-type GenInputResult[T any] struct {
+type GenInputResult[T any, M messageType] struct {
 	// RunCtx, if non-nil, overrides the context for this turn's execution
 	// (PrepareAgent, agent run, OnAgentEvents).
 	//
@@ -444,7 +444,7 @@ type GenInputResult[T any] struct {
 	RunCtx context.Context
 
 	// Input is the agent input to execute
-	Input *AgentInput
+	Input *TypedAgentInput[M]
 
 	// RunOpts are the options for this agent run.
 	// Note: do not pass WithCheckPointID here; the TurnLoop automatically
@@ -464,7 +464,7 @@ type GenInputResult[T any] struct {
 }
 
 // GenResumeResult contains the result of GenResume processing.
-type GenResumeResult[T any] struct {
+type GenResumeResult[T any, M messageType] struct {
 	// RunCtx, if non-nil, overrides the context for this resumed turn's execution
 	// (PrepareAgent, agent resume, OnAgentEvents).
 	RunCtx context.Context
@@ -489,9 +489,9 @@ type GenResumeResult[T any] struct {
 	Remaining []T
 }
 
-type turnRunSpec[T any] struct {
+type turnRunSpec[T any, M messageType] struct {
 	runCtx       context.Context
-	input        *AgentInput
+	input        *TypedAgentInput[M]
 	runOpts      []AgentRunOption
 	resumeParams *ResumeParams
 	isResume     bool
@@ -499,18 +499,18 @@ type turnRunSpec[T any] struct {
 	resumeBytes  []byte
 }
 
-type turnPlan[T any] struct {
+type turnPlan[T any, M messageType] struct {
 	turnCtx   context.Context
 	remaining []T
-	spec      *turnRunSpec[T]
+	spec      *turnRunSpec[T, M]
 }
 
-func (l *TurnLoop[T]) planTurn(
+func (l *TurnLoop[T, M]) planTurn(
 	ctx context.Context,
 	isResume bool,
 	items []T,
 	pr *turnLoopPendingResume[T],
-) (*turnPlan[T], error) {
+) (*turnPlan[T, M], error) {
 	if !isResume {
 		result, err := l.config.GenInput(ctx, l, items)
 		if err != nil {
@@ -526,10 +526,10 @@ func (l *TurnLoop[T]) planTurn(
 		if result.RunCtx != nil {
 			turnCtx = result.RunCtx
 		}
-		return &turnPlan[T]{
+		return &turnPlan[T, M]{
 			turnCtx:   turnCtx,
 			remaining: result.Remaining,
-			spec: &turnRunSpec[T]{
+			spec: &turnRunSpec[T, M]{
 				runCtx:   result.RunCtx,
 				input:    result.Input,
 				runOpts:  result.RunOpts,
@@ -554,10 +554,10 @@ func (l *TurnLoop[T]) planTurn(
 	if resumeResult.RunCtx != nil {
 		turnCtx = resumeResult.RunCtx
 	}
-	return &turnPlan[T]{
+	return &turnPlan[T, M]{
 		turnCtx:   turnCtx,
 		remaining: resumeResult.Remaining,
-		spec: &turnRunSpec[T]{
+		spec: &turnRunSpec[T, M]{
 			runCtx:       resumeResult.RunCtx,
 			runOpts:      resumeResult.RunOpts,
 			resumeParams: resumeResult.ResumeParams,
@@ -570,7 +570,7 @@ func (l *TurnLoop[T]) planTurn(
 
 // TurnLoopExitState is returned when TurnLoop exits, containing the exit reason
 // and any items that were not processed.
-type TurnLoopExitState[T any] struct {
+type TurnLoopExitState[T any, M messageType] struct {
 	// ExitReason indicates why the loop exited.
 	// nil means clean exit (Stop() was called without cancel options, or the
 	// agent completed normally before Stop took effect).
@@ -621,9 +621,9 @@ type TurnLoopExitState[T any] struct {
 }
 
 // TurnContext provides per-turn context to the OnAgentEvents callback.
-type TurnContext[T any] struct {
+type TurnContext[T any, M messageType] struct {
 	// Loop is the TurnLoop instance, allowing Push() or Stop() calls.
-	Loop *TurnLoop[T]
+	Loop *TurnLoop[T, M]
 
 	// Consumed contains items that triggered this agent execution.
 	Consumed []T
@@ -672,8 +672,8 @@ type TurnContext[T any] struct {
 //   - Wait: blocks until Run is called AND the loop exits. If Run is never
 //     called, Wait blocks forever (this is a programming error, analogous
 //     to reading from a channel that nobody writes to).
-type TurnLoop[T any] struct {
-	config TurnLoopConfig[T]
+type TurnLoop[T any, M messageType] struct {
+	config TurnLoopConfig[T, M]
 
 	buffer *turnBuffer[T]
 
@@ -682,7 +682,7 @@ type TurnLoop[T any] struct {
 
 	done chan struct{}
 
-	result *TurnLoopExitState[T]
+	result *TurnLoopExitState[T, M]
 
 	stopOnce sync.Once
 
@@ -702,14 +702,14 @@ type TurnLoop[T any] struct {
 
 	loadCheckpointID string
 
-	onAgentEvents func(ctx context.Context, tc *TurnContext[T], events *AsyncIterator[*AgentEvent]) error
+	onAgentEvents func(ctx context.Context, tc *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error
 
 	lateMu     sync.Mutex
 	lateItems  []T
 	lateSealed bool
 }
 
-func (l *TurnLoop[T]) appendLate(item T) {
+func (l *TurnLoop[T, M]) appendLate(item T) {
 	l.lateMu.Lock()
 	defer l.lateMu.Unlock()
 	if l.lateSealed {
@@ -744,7 +744,7 @@ func unmarshalTurnLoopCheckpoint[T any](data []byte) (*turnLoopCheckpoint[T], er
 	return &c, nil
 }
 
-func (l *TurnLoop[T]) saveTurnLoopCheckpoint(ctx context.Context, checkPointID string, c *turnLoopCheckpoint[T]) error {
+func (l *TurnLoop[T, M]) saveTurnLoopCheckpoint(ctx context.Context, checkPointID string, c *turnLoopCheckpoint[T]) error {
 	if l.config.Store == nil {
 		return errors.New("checkpoint store is nil")
 	}
@@ -755,7 +755,7 @@ func (l *TurnLoop[T]) saveTurnLoopCheckpoint(ctx context.Context, checkPointID s
 	return l.config.Store.Set(ctx, checkPointID, data)
 }
 
-func (l *TurnLoop[T]) deleteTurnLoopCheckpoint(ctx context.Context, checkPointID string) error {
+func (l *TurnLoop[T, M]) deleteTurnLoopCheckpoint(ctx context.Context, checkPointID string) error {
 	if l.config.Store == nil {
 		return nil
 	}
@@ -765,7 +765,7 @@ func (l *TurnLoop[T]) deleteTurnLoopCheckpoint(ctx context.Context, checkPointID
 	return nil
 }
 
-func (l *TurnLoop[T]) tryLoadCheckpoint(ctx context.Context) error {
+func (l *TurnLoop[T, M]) tryLoadCheckpoint(ctx context.Context) error {
 	checkPointID := l.config.CheckpointID
 	if checkPointID == "" || l.config.Store == nil {
 		return nil
@@ -973,15 +973,15 @@ func UntilIdleFor(duration time.Duration) StopOption {
 	}
 }
 
-type pushConfig[T any] struct {
+type pushConfig[T any, M messageType] struct {
 	preempt         bool
 	preemptDelay    time.Duration
 	agentCancelOpts []AgentCancelOption
-	pushStrategy    func(context.Context, *TurnContext[T]) []PushOption[T]
+	pushStrategy    func(context.Context, *TurnContext[T, M]) []PushOption[T, M]
 }
 
 // PushOption is an option for Push().
-type PushOption[T any] func(*pushConfig[T])
+type PushOption[T any, M messageType] func(*pushConfig[T, M])
 
 // WithPreempt signals that the current agent turn should be cancelled at the
 // specified safePoint after pushing the new item. The loop cancels the current
@@ -1001,11 +1001,11 @@ type PushOption[T any] func(*pushConfig[T])
 // passed to the same Push call, the last one wins.
 //
 // safePoint must not be zero; passing SafePoint(0) panics.
-func WithPreempt[T any](safePoint SafePoint) PushOption[T] {
+func WithPreempt[T any, M messageType](safePoint SafePoint) PushOption[T, M] {
 	if safePoint == 0 {
 		panic("adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint")
 	}
-	return func(cfg *pushConfig[T]) {
+	return func(cfg *pushConfig[T, M]) {
 		cfg.preempt = true
 		cfg.agentCancelOpts = []AgentCancelOption{
 			WithAgentCancelMode(safePoint.toCancelMode()),
@@ -1019,11 +1019,11 @@ func WithPreempt[T any](safePoint SafePoint) PushOption[T] {
 // also receive the cancel signal and be torn down.
 //
 // safePoint must not be zero; passing SafePoint(0) panics.
-func WithPreemptTimeout[T any](safePoint SafePoint, timeout time.Duration) PushOption[T] {
+func WithPreemptTimeout[T any, M messageType](safePoint SafePoint, timeout time.Duration) PushOption[T, M] {
 	if safePoint == 0 {
 		panic("adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint")
 	}
-	return func(cfg *pushConfig[T]) {
+	return func(cfg *pushConfig[T, M]) {
 		cfg.preempt = true
 		cfg.agentCancelOpts = []AgentCancelOption{
 			WithAgentCancelMode(safePoint.toCancelMode()),
@@ -1038,8 +1038,8 @@ func WithPreemptTimeout[T any](safePoint SafePoint, timeout time.Duration) PushO
 // immediately, but the preemption signal will be delayed by the specified
 // duration. This allows the current agent to continue processing for a grace
 // period before being preempted.
-func WithPreemptDelay[T any](delay time.Duration) PushOption[T] {
-	return func(cfg *pushConfig[T]) {
+func WithPreemptDelay[T any, M messageType](delay time.Duration) PushOption[T, M] {
+	return func(cfg *pushConfig[T, M]) {
 		cfg.preemptDelay = delay
 	}
 }
@@ -1054,22 +1054,22 @@ func WithPreemptDelay[T any](delay time.Duration) PushOption[T] {
 //
 // Example: preempt only if the current turn is processing low-priority items:
 //
-//	loop.Push(urgentItem, WithPushStrategy(func(ctx context.Context, tc *TurnContext[MyItem]) []PushOption[MyItem] {
+//	loop.Push(urgentItem, WithPushStrategy(func(ctx context.Context, tc *TurnContext[MyItem, *schema.Message]) []PushOption[MyItem, *schema.Message] {
 //	    if tc == nil {
 //	        return nil // between turns, plain push
 //	    }
 //	    if isLowPriority(tc.Consumed) {
-//	        return []PushOption[MyItem]{WithPreempt[MyItem](AnySafePoint)}
+//	        return []PushOption[MyItem, *schema.Message]{WithPreempt[MyItem, *schema.Message](AnySafePoint)}
 //	    }
 //	    return nil // don't preempt high-priority work
 //	}))
-func WithPushStrategy[T any](fn func(ctx context.Context, tc *TurnContext[T]) []PushOption[T]) PushOption[T] {
-	return func(cfg *pushConfig[T]) {
+func WithPushStrategy[T any, M messageType](fn func(ctx context.Context, tc *TurnContext[T, M]) []PushOption[T, M]) PushOption[T, M] {
+	return func(cfg *pushConfig[T, M]) {
 		cfg.pushStrategy = fn
 	}
 }
 
-func defaultTurnLoopOnAgentEvents[T any](_ context.Context, _ *TurnContext[T], events *AsyncIterator[*AgentEvent]) error {
+func defaultTurnLoopOnAgentEvents[T any, M messageType](_ context.Context, _ *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error {
 	for {
 		event, ok := events.Next()
 		if !ok {
@@ -1088,7 +1088,7 @@ func defaultTurnLoopOnAgentEvents[T any](_ context.Context, _ *TurnContext[T], e
 // Call Run to start the processing goroutine.
 //
 // NewTurnLoop panics if GenInput or PrepareAgent is nil.
-func NewTurnLoop[T any](cfg TurnLoopConfig[T]) *TurnLoop[T] {
+func NewTurnLoop[T any, M messageType](cfg TurnLoopConfig[T, M]) *TurnLoop[T, M] {
 	if cfg.GenInput == nil {
 		panic("adk: NewTurnLoop: GenInput is required")
 	}
@@ -1096,7 +1096,7 @@ func NewTurnLoop[T any](cfg TurnLoopConfig[T]) *TurnLoop[T] {
 		panic("adk: NewTurnLoop: PrepareAgent is required")
 	}
 
-	l := &TurnLoop[T]{
+	l := &TurnLoop[T, M]{
 		config:     cfg,
 		buffer:     newTurnBuffer[T](),
 		done:       make(chan struct{}),
@@ -1106,12 +1106,12 @@ func NewTurnLoop[T any](cfg TurnLoopConfig[T]) *TurnLoop[T] {
 	if cfg.OnAgentEvents != nil {
 		l.onAgentEvents = cfg.OnAgentEvents
 	} else {
-		l.onAgentEvents = defaultTurnLoopOnAgentEvents[T]
+		l.onAgentEvents = defaultTurnLoopOnAgentEvents[T, M]
 	}
 	return l
 }
 
-func (l *TurnLoop[T]) start(ctx context.Context) {
+func (l *TurnLoop[T, M]) start(ctx context.Context) {
 	l.runOnce.Do(func() {
 		atomic.StoreInt32(&l.started, 1)
 		go l.run(ctx)
@@ -1126,7 +1126,7 @@ func (l *TurnLoop[T]) start(ctx context.Context) {
 // Otherwise it starts fresh with whatever items were Push()-ed.
 //
 // Calling Run more than once is a no-op: only the first call starts the loop.
-func (l *TurnLoop[T]) Run(ctx context.Context) {
+func (l *TurnLoop[T, M]) Run(ctx context.Context) {
 	l.start(ctx)
 }
 
@@ -1156,8 +1156,8 @@ func (l *TurnLoop[T]) Run(ctx context.Context) {
 // the preemption signal.
 // Push returns immediately after the item is buffered, and a goroutine is spawned
 // to signal preemption after the delay.
-func (l *TurnLoop[T]) Push(item T, opts ...PushOption[T]) (bool, <-chan struct{}) {
-	cfg := &pushConfig[T]{}
+func (l *TurnLoop[T, M]) Push(item T, opts ...PushOption[T, M]) (bool, <-chan struct{}) {
+	cfg := &pushConfig[T, M]{}
 	for _, opt := range opts {
 		opt(cfg)
 	}
@@ -1173,19 +1173,19 @@ func (l *TurnLoop[T]) Push(item T, opts ...PushOption[T]) (bool, <-chan struct{}
 // then calls the strategy callback with a guaranteed-stable TurnContext. If the
 // strategy returns preempt options, the hold is kept and a preempt is requested;
 // otherwise the hold is released and the item is buffered as a plain push.
-func (l *TurnLoop[T]) pushWithStrategy(item T, cfg *pushConfig[T]) (bool, <-chan struct{}) {
+func (l *TurnLoop[T, M]) pushWithStrategy(item T, cfg *pushConfig[T, M]) (bool, <-chan struct{}) {
 	strategy := cfg.pushStrategy
 
 	runCtx, tcAny := l.preemptSig.holdAndGetTurn()
 	if runCtx == nil {
 		runCtx = context.Background()
 	}
-	var tc *TurnContext[T]
+	var tc *TurnContext[T, M]
 	if tcAny != nil {
-		tc = tcAny.(*TurnContext[T])
+		tc = tcAny.(*TurnContext[T, M])
 	}
 	realOpts := strategy(runCtx, tc)
-	cfg = &pushConfig[T]{}
+	cfg = &pushConfig[T, M]{}
 	for _, opt := range realOpts {
 		opt(cfg)
 	}
@@ -1235,7 +1235,7 @@ func (l *TurnLoop[T]) pushWithStrategy(item T, cfg *pushConfig[T]) (bool, <-chan
 	return true, ack
 }
 
-func (l *TurnLoop[T]) pushWithConfig(item T, cfg *pushConfig[T]) (bool, <-chan struct{}) {
+func (l *TurnLoop[T, M]) pushWithConfig(item T, cfg *pushConfig[T, M]) (bool, <-chan struct{}) {
 	if atomic.LoadInt32(&l.stopped) != 0 {
 		l.appendLate(item)
 		return false, nil
@@ -1304,7 +1304,7 @@ func (l *TurnLoop[T]) pushWithConfig(item T, cfg *pushConfig[T]) (bool, <-chan s
 // all cancel-related options (WithImmediate, WithGraceful, WithGracefulTimeout)
 // degrade to "exit the loop on entering the next iteration" — the current
 // agent turn runs to completion before the loop exits.
-func (l *TurnLoop[T]) Stop(opts ...StopOption) {
+func (l *TurnLoop[T, M]) Stop(opts ...StopOption) {
 	cfg := &stopConfig{}
 	for _, opt := range opts {
 		opt(cfg)
@@ -1327,7 +1327,7 @@ func (l *TurnLoop[T]) Stop(opts ...StopOption) {
 	l.commitStop()
 }
 
-func (l *TurnLoop[T]) commitStop() {
+func (l *TurnLoop[T, M]) commitStop() {
 	l.stopOnce.Do(func() {
 		l.stopSig.closeDone()
 		atomic.StoreInt32(&l.stopped, 1)
@@ -1341,12 +1341,12 @@ func (l *TurnLoop[T]) commitStop() {
 //
 // Wait blocks until Run is called AND the loop exits. If Run is
 // never called, Wait blocks forever.
-func (l *TurnLoop[T]) Wait() *TurnLoopExitState[T] {
+func (l *TurnLoop[T, M]) Wait() *TurnLoopExitState[T, M] {
 	<-l.done
 	return l.result
 }
 
-func (l *TurnLoop[T]) run(ctx context.Context) {
+func (l *TurnLoop[T, M]) run(ctx context.Context) {
 	defer l.cleanup(ctx)
 
 	if err := l.tryLoadCheckpoint(ctx); err != nil {
@@ -1502,7 +1502,7 @@ func (l *TurnLoop[T]) run(ctx context.Context) {
 	}
 }
 
-func (l *TurnLoop[T]) setupBridgeStore(spec *turnRunSpec[T], runOpts []AgentRunOption) ([]AgentRunOption, *bridgeStore, error) {
+func (l *TurnLoop[T, M]) setupBridgeStore(spec *turnRunSpec[T, M], runOpts []AgentRunOption) ([]AgentRunOption, *bridgeStore, error) {
 	store := l.config.Store
 	if store == nil && spec.isResume {
 		return nil, nil, fmt.Errorf("failed to resume agent: checkpoint store is nil")
@@ -1530,7 +1530,7 @@ func (l *TurnLoop[T]) setupBridgeStore(spec *turnRunSpec[T], runOpts []AgentRunO
 // On the first preempt whose cancel actually contributed (i.e. the cancel options
 // were accepted before the CancelError was finalized), preemptDone is closed to
 // wake runAgentAndHandleEvents's select.
-func (l *TurnLoop[T]) watchPreemptSignal(done <-chan struct{}, agentCancelFunc AgentCancelFunc, preemptDone chan struct{}) {
+func (l *TurnLoop[T, M]) watchPreemptSignal(done <-chan struct{}, agentCancelFunc AgentCancelFunc, preemptDone chan struct{}) {
 	var lastGen uint64
 	for {
 		select {
@@ -1573,7 +1573,7 @@ func (l *TurnLoop[T]) watchPreemptSignal(done <-chan struct{}, agentCancelFunc A
 // On the first cancel that actually contributed (i.e. the cancel was accepted
 // before the CancelError was finalized), stoppedDone is closed to wake
 // runAgentAndHandleEvents's select.
-func (l *TurnLoop[T]) watchStopSignal(done <-chan struct{}, agentCancelFunc AgentCancelFunc, stoppedDone chan struct{}) {
+func (l *TurnLoop[T, M]) watchStopSignal(done <-chan struct{}, agentCancelFunc AgentCancelFunc, stoppedDone chan struct{}) {
 	var lastGen uint64
 	stoppedClosed := false
 
@@ -1612,12 +1612,12 @@ func (l *TurnLoop[T]) watchStopSignal(done <-chan struct{}, agentCancelFunc Agen
 	}
 }
 
-func (l *TurnLoop[T]) runAgentAndHandleEvents(
+func (l *TurnLoop[T, M]) runAgentAndHandleEvents(
 	ctx context.Context,
-	agent Agent,
-	spec *turnRunSpec[T],
+	agent TypedAgent[M],
+	spec *turnRunSpec[T, M],
 ) error {
-	var iter *AsyncIterator[*AgentEvent]
+	var iter *AsyncIterator[*TypedAgentEvent[M]]
 
 	runOpts, ms, err := l.setupBridgeStore(spec, spec.runOpts)
 	if err != nil {
@@ -1631,7 +1631,7 @@ func (l *TurnLoop[T]) runAgentAndHandleEvents(
 	if spec.input != nil {
 		enableStreaming = spec.input.EnableStreaming
 	}
-	runner := NewRunner(ctx, RunnerConfig{
+	runner := NewTypedRunner[M](TypedRunnerConfig[M]{
 		EnableStreaming: enableStreaming,
 		Agent:           agent,
 		CheckPointStore: ms,
@@ -1640,7 +1640,7 @@ func (l *TurnLoop[T]) runAgentAndHandleEvents(
 	preemptDone := make(chan struct{})
 	stoppedDone := make(chan struct{})
 
-	tc := &TurnContext[T]{
+	tc := &TurnContext[T, M]{
 		Loop:      l,
 		Consumed:  spec.consumed,
 		Preempted: preemptDone,
@@ -1743,7 +1743,7 @@ func (l *TurnLoop[T]) runAgentAndHandleEvents(
 	}
 }
 
-func (l *TurnLoop[T]) cleanup(ctx context.Context) {
+func (l *TurnLoop[T, M]) cleanup(ctx context.Context) {
 	atomic.StoreInt32(&l.stopped, 1)
 
 	unhandled := l.buffer.TakeAll()
@@ -1777,7 +1777,7 @@ func (l *TurnLoop[T]) cleanup(ctx context.Context) {
 	var takeLateOnce sync.Once
 	var takeLateResult []T
 
-	l.result = &TurnLoopExitState[T]{
+	l.result = &TurnLoopExitState[T, M]{
 		ExitReason:          l.runErr,
 		UnhandledItems:      unhandled,
 		CanceledItems:       l.canceledItems,

adk/utils.go

@@ -102,8 +102,7 @@ func GenTransferMessages(_ context.Context, destAgentName string) (Message, Mess
 	return assistantMessage, toolMessage
 }
 
-// set automatic close for event's message stream
-func setAutomaticClose(e *AgentEvent) {
+func typedSetAutomaticClose[M messageType](e *TypedAgentEvent[M]) {
 	if e.Output == nil || e.Output.MessageOutput == nil || !e.Output.MessageOutput.IsStreaming {
 		return
 	}
@@ -111,10 +110,41 @@ func setAutomaticClose(e *AgentEvent) {
 	e.Output.MessageOutput.MessageStream.SetAutomaticClose()
 }
 
+// set automatic close for event's message stream
+func setAutomaticClose(e *AgentEvent) {
+	typedSetAutomaticClose(e)
+}
+
 // getMessageFromWrappedEvent extracts the message from an AgentEvent.
 // If the stream contains an error chunk, this function returns (nil, err) and
 // sets StreamErr to prevent re-consumption. The nil message ensures that
 // failed stream responses are not included in subsequent agents' context windows.
+func getMessageFromTypedWrappedEvent[M messageType](e *typedAgentEventWrapper[M]) (M, error) {
+	var zero M
+	if e.event.Output == nil || e.event.Output.MessageOutput == nil {
+		return zero, nil
+	}
+
+	if !e.event.Output.MessageOutput.IsStreaming {
+		return e.event.Output.MessageOutput.Message, nil
+	}
+
+	if e.StreamErr != nil {
+		return zero, e.StreamErr
+	}
+
+	if !isNilMessage(e.concatenatedMessage) {
+		return e.concatenatedMessage, nil
+	}
+
+	e.consumeStream()
+
+	if e.StreamErr != nil {
+		return zero, e.StreamErr
+	}
+	return e.concatenatedMessage, nil
+}
+
 func getMessageFromWrappedEvent(e *agentEventWrapper) (Message, error) {
 	if e.AgentEvent.Output == nil || e.AgentEvent.Output.MessageOutput == nil {
 		return nil, nil
@@ -194,21 +224,21 @@ func (e *agentEventWrapper) consumeStream() {
 	e.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray([]Message{e.concatenatedMessage})
 }
 
-// copyAgentEvent copies an AgentEvent.
+// copyTypedAgentEvent copies a TypedAgentEvent.
 // If the MessageVariant is streaming, the MessageStream will be copied.
 // RunPath will be deep copied.
-// The result of Copy will be a new AgentEvent that is:
-// - safe to set fields of AgentEvent
+// The result of Copy will be a new TypedAgentEvent that is:
+// - safe to set fields of TypedAgentEvent
 // - safe to extend RunPath
 // - safe to receive from MessageStream
-// NOTE: even if the AgentEvent is copied, it's still not recommended to modify
+// NOTE: even if the event is copied, it's still not recommended to modify
 // the Message itself or Chunks of the MessageStream, as they are not copied.
 // NOTE: if you have CustomizedOutput or CustomizedAction, they are NOT copied.
-func copyAgentEvent(ae *AgentEvent) *AgentEvent {
+func copyTypedAgentEvent[M messageType](ae *TypedAgentEvent[M]) *TypedAgentEvent[M] {
 	rp := make([]RunStep, len(ae.RunPath))
 	copy(rp, ae.RunPath)
 
-	copied := &AgentEvent{
+	copied := &TypedAgentEvent[M]{
 		AgentName: ae.AgentName,
 		RunPath:   rp,
 		Action:    ae.Action,
@@ -219,7 +249,7 @@ func copyAgentEvent(ae *AgentEvent) *AgentEvent {
 		return copied
 	}
 
-	copied.Output = &AgentOutput{
+	copied.Output = &TypedAgentOutput[M]{
 		CustomizedOutput: ae.Output.CustomizedOutput,
 	}
 
@@ -228,7 +258,7 @@ func copyAgentEvent(ae *AgentEvent) *AgentEvent {
 		return copied
 	}
 
-	copied.Output.MessageOutput = &MessageVariant{
+	copied.Output.MessageOutput = &TypedMessageVariant[M]{
 		IsStreaming: mv.IsStreaming,
 		Role:        mv.Role,
 		ToolName:    mv.ToolName,
@@ -244,11 +274,11 @@ func copyAgentEvent(ae *AgentEvent) *AgentEvent {
 	return copied
 }
 
-// GetMessage extracts the Message from an AgentEvent. For streaming output,
-// it duplicates the stream and concatenates it into a single Message.
-func GetMessage(e *AgentEvent) (Message, *AgentEvent, error) {
+// TypedGetMessage extracts the message from a TypedAgentEvent, concatenating a stream if present.
+func TypedGetMessage[M messageType](e *TypedAgentEvent[M]) (M, *TypedAgentEvent[M], error) {
+	var zero M
 	if e.Output == nil || e.Output.MessageOutput == nil {
-		return nil, e, nil
+		return zero, e, nil
 	}
 
 	msgOutput := e.Output.MessageOutput
@@ -256,7 +286,7 @@ func GetMessage(e *AgentEvent) (Message, *AgentEvent, error) {
 		ss := msgOutput.MessageStream.Copy(2)
 		e.Output.MessageOutput.MessageStream = ss[0]
 
-		msg, err := schema.ConcatMessageStream(ss[1])
+		msg, err := concatMessageStream(ss[1])
 
 		return msg, e, err
 	}
@@ -264,9 +294,19 @@ func GetMessage(e *AgentEvent) (Message, *AgentEvent, error) {
 	return msgOutput.Message, e, nil
 }
 
-func genErrorIter(err error) *AsyncIterator[*AgentEvent] {
-	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
-	generator.Send(&AgentEvent{Err: err})
+// GetMessage extracts the Message from an AgentEvent. For streaming output,
+// it duplicates the stream and concatenates it into a single Message.
+func GetMessage(e *AgentEvent) (Message, *AgentEvent, error) {
+	return TypedGetMessage(e)
+}
+
+func typedErrorIter[M messageType](err error) *AsyncIterator[*TypedAgentEvent[M]] {
+	iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]()
+	generator.Send(&TypedAgentEvent[M]{Err: err})
 	generator.Close()
 	return iterator
 }
+
+func genErrorIter(err error) *AsyncIterator[*AgentEvent] {
+	return typedErrorIter[*schema.Message](err)
+}

adk/workflow_test.go

@@ -1021,7 +1021,7 @@ func TestWorkflowAgentUnsupportedMode(t *testing.T) {
 		name:        "UnsupportedModeAgent",
 		description: "Agent with unsupported mode",
 		subAgents:   []*flowAgent{},
-		mode:        workflowAgentMode(999), // Invalid mode
+		mode:        workflowAgentMode(999),
 	}
 
 	// Run the agent and expect error

adk/wrappers.go

@@ -32,11 +32,11 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-type generateEndpoint func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)
-type streamEndpoint func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)
+type typedGenerateEndpoint[M messageType] func(ctx context.Context, input []M, opts ...model.Option) (M, error)
+type typedStreamEndpoint[M messageType] func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error)
 
-type modelWrapperConfig struct {
-	handlers       []ChatModelAgentMiddleware
+type typedModelWrapperConfig[M messageType] struct {
+	handlers       []TypedChatModelAgentMiddleware[M]
 	middlewares    []AgentMiddleware
 	retryConfig    *ModelRetryConfig
 	failoverConfig *ModelFailoverConfig
@@ -44,19 +44,24 @@ type modelWrapperConfig struct {
 	cancelContext  *cancelContext
 }
 
-func buildModelWrappers(m model.BaseChatModel, config *modelWrapperConfig) model.BaseChatModel {
-	var wrapped model.BaseChatModel = m
+type modelWrapperConfig = typedModelWrapperConfig[*schema.Message]
+
+func buildModelWrappers[M messageType](m model.BaseModel[M], config *typedModelWrapperConfig[M]) model.BaseModel[M] {
+	return buildModelWrappersImpl(m, config)
+}
+
+func buildModelWrappersImpl[M messageType](m model.BaseModel[M], config *typedModelWrapperConfig[M]) model.BaseModel[M] {
+	var wrapped model.BaseModel[M] = m
 
-	// failoverProxyModel must be the innermost wrapper to read the selected failover model from context.
 	if config.failoverConfig != nil {
-		wrapped = &failoverProxyModel{}
+		wrapped = &typedFailoverProxyModel[M]{}
 	}
 
 	if !components.IsCallbacksEnabled(wrapped) {
-		wrapped = (&callbackInjectionModelWrapper{}).WrapModel(wrapped)
+		wrapped = typedCallbackInjectionModelWrapper[M]{}.wrapModel(wrapped)
 	}
 
-	wrapped = &stateModelWrapper{
+	wrapped = &typedStateModelWrapper[M]{
 		inner:               wrapped,
 		original:            m,
 		handlers:            config.handlers,
@@ -70,28 +75,29 @@ func buildModelWrappers(m model.BaseChatModel, config *modelWrapperConfig) model
 	return wrapped
 }
 
-type callbackInjectionModelWrapper struct{}
+type typedCallbackInjectionModelWrapper[M messageType] struct{}
 
-func (w *callbackInjectionModelWrapper) WrapModel(m model.BaseChatModel) model.BaseChatModel {
-	return &callbackInjectedModel{inner: m}
+func (w typedCallbackInjectionModelWrapper[M]) wrapModel(m model.BaseModel[M]) model.BaseModel[M] {
+	return &typedCallbackInjectedModel[M]{inner: m}
 }
 
-type callbackInjectedModel struct {
-	inner model.BaseChatModel
+type typedCallbackInjectedModel[M messageType] struct {
+	inner model.BaseModel[M]
 }
 
-func (m *callbackInjectedModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (m *typedCallbackInjectedModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	ctx = callbacks.OnStart(ctx, input)
 	result, err := m.inner.Generate(ctx, input, opts...)
 	if err != nil {
 		callbacks.OnError(ctx, err)
-		return nil, err
+		var zero M
+		return zero, err
 	}
 	callbacks.OnEnd(ctx, result)
 	return result, nil
 }
 
-func (m *callbackInjectedModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+func (m *typedCallbackInjectedModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 	ctx = callbacks.OnStart(ctx, input)
 	result, err := m.inner.Stream(ctx, input, opts...)
 	if err != nil {
@@ -102,7 +108,7 @@ func (m *callbackInjectedModel) Stream(ctx context.Context, input []*schema.Mess
 	return wrappedStream, nil
 }
 
-func handlersToToolMiddlewares(handlers []ChatModelAgentMiddleware) []compose.ToolMiddleware {
+func handlersToToolMiddlewares[M messageType](handlers []TypedChatModelAgentMiddleware[M]) []compose.ToolMiddleware {
 	var middlewares []compose.ToolMiddleware
 	// Forward iteration: compose.wrapToolCall applies middlewares in reverse order
 	// (len-1 down to 0), so keeping the original handler order here means
@@ -249,25 +255,21 @@ func handlersToToolMiddlewares(handlers []ChatModelAgentMiddleware) []compose.To
 	return middlewares
 }
 
-type eventSenderModelWrapper struct {
-	*BaseChatModelAgentMiddleware
+type typedEventSenderModelWrapper[M messageType] struct {
+	*TypedBaseChatModelAgentMiddleware[M]
 }
 
-// NewEventSenderModelWrapper returns a ChatModelAgentMiddleware that sends model response events.
-// By default, the framework applies this wrapper after all user middlewares, so events contain
-// modified messages. To send events with original (unmodified) output, pass this as a Handler
-// after the modifying middleware (placing it innermost in the wrapper chain).
-// When detected in Handlers, the framework skips the default event sender to avoid duplicates.
+// NewEventSenderModelWrapper creates a ChatModelAgentMiddleware that sends model output as agent events.
 func NewEventSenderModelWrapper() ChatModelAgentMiddleware {
-	return &eventSenderModelWrapper{
-		BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{},
+	return &typedEventSenderModelWrapper[*schema.Message]{
+		TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[*schema.Message]{},
 	}
 }
 
-func (w *eventSenderModelWrapper) WrapModel(_ context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error) {
+func (w *typedEventSenderModelWrapper[M]) WrapModel(_ context.Context, m model.BaseModel[M], mc *ModelContext) (model.BaseModel[M], error) {
 	inner := m
 	if mc != nil && mc.cancelContext != nil {
-		inner = &cancelMonitoredModel{
+		inner = &typedCancelMonitoredModel[M]{
 			inner:         inner,
 			cancelContext: mc.cancelContext,
 		}
@@ -280,43 +282,44 @@ func (w *eventSenderModelWrapper) WrapModel(_ context.Context, m model.BaseChatM
 	if mc != nil {
 		failoverConfig = mc.ModelFailoverConfig
 	}
-	return &eventSenderModel{inner: inner, modelRetryConfig: retryConfig, modelFailoverConfig: failoverConfig}, nil
+	return &typedEventSenderModel[M]{inner: inner, modelRetryConfig: retryConfig, modelFailoverConfig: failoverConfig}, nil
 }
 
-type eventSenderModel struct {
-	inner               model.BaseChatModel
+type typedEventSenderModel[M messageType] struct {
+	inner               model.BaseModel[M]
 	modelRetryConfig    *ModelRetryConfig
 	modelFailoverConfig *ModelFailoverConfig
 }
 
-func (m *eventSenderModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (m *typedEventSenderModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	result, err := m.inner.Generate(ctx, input, opts...)
 	if err != nil {
-		return nil, err
+		var zero M
+		return zero, err
 	}
 
-	execCtx := getChatModelAgentExecCtx(ctx)
+	execCtx := getTypedChatModelAgentExecCtx[M](ctx)
 	if execCtx != nil && execCtx.suppressEventSend {
 		return result, nil
 	}
 	if execCtx == nil || execCtx.generator == nil {
-		return nil, errors.New("generator is nil when sending event in Generate: ensure agent state is properly initialized")
+		var zero M
+		return zero, errors.New("generator is nil when sending event in Generate: ensure agent state is properly initialized")
 	}
 
-	msgCopy := *result
-	event := EventFromMessage(&msgCopy, nil, schema.Assistant, "")
+	event := typedModelOutputEvent(copyMessage(result), nil)
 	execCtx.send(event)
 
 	return result, nil
 }
 
-func (m *eventSenderModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+func (m *typedEventSenderModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 	result, err := m.inner.Stream(ctx, input, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	execCtx := getChatModelAgentExecCtx(ctx)
+	execCtx := getTypedChatModelAgentExecCtx[M](ctx)
 	if execCtx == nil || execCtx.generator == nil {
 		result.Close()
 		return nil, errors.New("generator is nil when sending event in Stream: ensure agent state is properly initialized")
@@ -327,11 +330,12 @@ func (m *eventSenderModel) Stream(ctx context.Context, input []*schema.Message,
 	eventStream := streams[0]
 	if convertOpts := m.buildStreamConvertOptions(ctx); len(convertOpts) > 0 {
 		eventStream = schema.StreamReaderWithConvert(streams[0],
-			func(msg *schema.Message) (*schema.Message, error) { return msg, nil },
+			func(msg M) (M, error) { return msg, nil },
 			convertOpts...)
 	}
 
-	event := EventFromMessage(nil, eventStream, schema.Assistant, "")
+	var zero M
+	event := typedModelOutputEvent[M](zero, eventStream)
 	execCtx.send(event)
 
 	return streams[1], nil
@@ -356,9 +360,9 @@ func (m *eventSenderModel) Stream(ctx context.Context, input []*schema.Message,
 // This prevents a goroutine leak when a mid-stream error is followed by EOF: errWrapper fires
 // first (caching the verdict), and onEOF reuses the cached value instead of blocking on a
 // drained channel.
-func (m *eventSenderModel) buildStreamConvertOptions(ctx context.Context) []schema.ConvertOption {
+func (m *typedEventSenderModel[M]) buildStreamConvertOptions(ctx context.Context) []schema.ConvertOption {
 	var retryAttempt int
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		retryAttempt = st.getRetryAttempt()
 		return nil
 	})
@@ -377,7 +381,7 @@ func (m *eventSenderModel) buildStreamConvertOptions(ctx context.Context) []sche
 	var retryWrapper func(error) error
 	if m.modelRetryConfig != nil {
 		if m.modelRetryConfig.ShouldRetry != nil {
-			execCtx := getChatModelAgentExecCtx(ctx)
+			execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 			signal := (*retryVerdictSignal)(nil)
 			if execCtx != nil {
 				signal = execCtx.retryVerdictSignal
@@ -462,11 +466,24 @@ func (m *eventSenderModel) buildStreamConvertOptions(ctx context.Context) []sche
 	return opts
 }
 
-func popToolGenAction(ctx context.Context, toolName string) *AgentAction {
+func copyMessage[M messageType](msg M) M {
+	switch v := any(msg).(type) {
+	case *schema.Message:
+		cp := *v
+		return any(&cp).(M)
+	case *schema.AgenticMessage:
+		cp := *v
+		return any(&cp).(M)
+	default:
+		return msg
+	}
+}
+
+func typedPopToolGenAction[M messageType](ctx context.Context, toolName string) *AgentAction {
 	toolCallID := compose.GetToolCallID(ctx)
 
 	var action *AgentAction
-	_ = compose.ProcessState(ctx, func(ctx context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(ctx context.Context, st *typedState[M]) error {
 		if len(toolCallID) > 0 {
 			if a := st.popToolGenAction(toolCallID); a != nil {
 				action = a
@@ -484,10 +501,23 @@ func popToolGenAction(ctx context.Context, toolName string) *AgentAction {
 	return action
 }
 
+func popToolGenAction(ctx context.Context, toolName string) *AgentAction {
+	return typedPopToolGenAction[*schema.Message](ctx, toolName)
+}
+
 type eventSenderToolWrapper struct {
 	*BaseChatModelAgentMiddleware
 }
 
+func (*eventSenderToolWrapper) isEventSenderToolWrapper() {}
+
+// eventSenderToolWrapperMarker enables cross-type detection of eventSenderToolWrapper
+// in generic contexts. hasUserEventSenderToolWrapper[M] receives
+// []TypedChatModelAgentMiddleware[M], so when M is *schema.AgenticMessage, a direct
+// type assertion to *eventSenderToolWrapper (which implements the *schema.Message alias)
+// would fail. The marker interface bridges this gap.
+type eventSenderToolWrapperMarker interface{ isEventSenderToolWrapper() }
+
 // NewEventSenderToolWrapper returns a ChatModelAgentMiddleware that sends tool result events.
 // By default, the framework places this before all user middlewares (outermost), so events
 // reflect the fully processed tool output. To control exactly where events are emitted,
@@ -516,7 +546,7 @@ func (w *eventSenderToolWrapper) WrapInvokableToolCall(_ context.Context, endpoi
 			event.Action = prePopAction
 		}
 
-		execCtx := getChatModelAgentExecCtx(ctx)
+		execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
 			if st.getReturnDirectlyToolCallID() == callID {
 				st.setReturnDirectlyEvent(event)
@@ -550,7 +580,7 @@ func (w *eventSenderToolWrapper) WrapStreamableToolCall(_ context.Context, endpo
 		event := EventFromMessage(nil, msgStream, schema.Tool, toolName)
 		event.Action = prePopAction
 
-		execCtx := getChatModelAgentExecCtx(ctx)
+		execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
 			if st.getReturnDirectlyToolCallID() == callID {
 				st.setReturnDirectlyEvent(event)
@@ -585,7 +615,7 @@ func (w *eventSenderToolWrapper) WrapEnhancedInvokableToolCall(_ context.Context
 			event.Action = prePopAction
 		}
 
-		execCtx := getChatModelAgentExecCtx(ctx)
+		execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
 			if st.getReturnDirectlyToolCallID() == callID {
 				st.setReturnDirectlyEvent(event)
@@ -625,7 +655,7 @@ func (w *eventSenderToolWrapper) WrapEnhancedStreamableToolCall(_ context.Contex
 		event := EventFromMessage(nil, msgStream, schema.Tool, toolName)
 		event.Action = prePopAction
 
-		execCtx := getChatModelAgentExecCtx(ctx)
+		execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx)
 		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
 			if st.getReturnDirectlyToolCallID() == callID {
 				st.setReturnDirectlyEvent(event)
@@ -639,19 +669,19 @@ func (w *eventSenderToolWrapper) WrapEnhancedStreamableToolCall(_ context.Contex
 	}, nil
 }
 
-func hasUserEventSenderToolWrapper(handlers []ChatModelAgentMiddleware) bool {
+func hasUserEventSenderToolWrapper[M messageType](handlers []TypedChatModelAgentMiddleware[M]) bool {
 	for _, handler := range handlers {
-		if _, ok := handler.(*eventSenderToolWrapper); ok {
+		if _, ok := any(handler).(eventSenderToolWrapperMarker); ok {
 			return true
 		}
 	}
 	return false
 }
 
-type stateModelWrapper struct {
-	inner               model.BaseChatModel
-	original            model.BaseChatModel
-	handlers            []ChatModelAgentMiddleware
+type typedStateModelWrapper[M messageType] struct {
+	inner               model.BaseModel[M]
+	original            model.BaseModel[M]
+	handlers            []TypedChatModelAgentMiddleware[M]
 	middlewares         []AgentMiddleware
 	toolInfos           []*schema.ToolInfo
 	modelRetryConfig    *ModelRetryConfig
@@ -659,27 +689,29 @@ type stateModelWrapper struct {
 	cancelContext       *cancelContext
 }
 
-func (w *stateModelWrapper) IsCallbacksEnabled() bool {
+type stateModelWrapper = typedStateModelWrapper[*schema.Message]
+
+func (w *typedStateModelWrapper[M]) IsCallbacksEnabled() bool {
 	return true
 }
 
-func (w *stateModelWrapper) GetType() string {
-	if typer, ok := w.original.(components.Typer); ok {
+func (w *typedStateModelWrapper[M]) GetType() string {
+	if typer, ok := any(w.original).(components.Typer); ok {
 		return typer.GetType()
 	}
 	return generic.ParseTypeName(reflect.ValueOf(w.original))
 }
 
-func (w *stateModelWrapper) hasUserEventSender() bool {
+func (w *typedStateModelWrapper[M]) hasUserEventSender() bool {
 	for _, handler := range w.handlers {
-		if _, ok := handler.(*eventSenderModelWrapper); ok {
+		if _, ok := any(handler).(*typedEventSenderModelWrapper[M]); ok {
 			return true
 		}
 	}
 	return false
 }
 
-func (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) generateEndpoint {
+func (w *typedStateModelWrapper[M]) wrapGenerateEndpoint(endpoint typedGenerateEndpoint[M]) typedGenerateEndpoint[M] {
 	hasUserEventSender := w.hasUserEventSender()
 	retryConfig := w.modelRetryConfig
 	failoverConfig := w.modelFailoverConfig
@@ -689,13 +721,14 @@ func (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) gene
 		handler := w.handlers[i]
 		innerEndpoint := endpoint
 		baseToolInfos := w.toolInfos
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 			baseOpts := &model.Options{Tools: baseToolInfos}
 			commonOpts := model.GetCommonOptions(baseOpts, opts...)
 			mc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig, cancelContext: cc}
-			wrappedModel, err := handler.WrapModel(ctx, &endpointModel{generate: innerEndpoint}, mc)
+			wrappedModel, err := handler.WrapModel(ctx, &typedEndpointModel[M]{generate: innerEndpoint}, mc)
 			if err != nil {
-				return nil, err
+				var zero M
+				return zero, err
 			}
 			return wrappedModel.Generate(ctx, input, opts...)
 		}
@@ -703,16 +736,19 @@ func (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) gene
 
 	if !hasUserEventSender {
 		innerEndpoint := endpoint
-		eventSender := NewEventSenderModelWrapper()
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
-			execCtx := getChatModelAgentExecCtx(ctx)
+		eventSender := &typedEventSenderModelWrapper[M]{
+			TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[M]{},
+		}
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) {
+			execCtx := getTypedChatModelAgentExecCtx[M](ctx)
 			if execCtx == nil || execCtx.generator == nil {
 				return innerEndpoint(ctx, input, opts...)
 			}
 			mc := &ModelContext{ModelRetryConfig: retryConfig, ModelFailoverConfig: failoverConfig, cancelContext: cc}
-			wrappedModel, err := eventSender.WrapModel(ctx, &endpointModel{generate: innerEndpoint}, mc)
+			wrappedModel, err := eventSender.WrapModel(ctx, &typedEndpointModel[M]{generate: innerEndpoint}, mc)
 			if err != nil {
-				return nil, err
+				var zero M
+				return zero, err
 			}
 			return wrappedModel.Generate(ctx, input, opts...)
 		}
@@ -720,18 +756,17 @@ func (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) gene
 
 	if w.modelRetryConfig != nil {
 		innerEndpoint := endpoint
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
-			retryWrapper := newRetryModelWrapper(&endpointModel{generate: innerEndpoint}, w.modelRetryConfig)
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) {
+			retryWrapper := newTypedRetryModelWrapper[M](&typedEndpointModel[M]{generate: innerEndpoint}, w.modelRetryConfig)
 			return retryWrapper.Generate(ctx, input, opts...)
 		}
 	}
 
-	// Needs to handle failoverWrapper after retryWrapper
 	if w.modelFailoverConfig != nil {
 		config := w.modelFailoverConfig
 		innerEndpoint := endpoint
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
-			failoverWrapper := newFailoverModelWrapper(&endpointModel{generate: innerEndpoint}, config)
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) {
+			failoverWrapper := newTypedFailoverModelWrapper[M](&typedEndpointModel[M]{generate: innerEndpoint}, config)
 			return failoverWrapper.Generate(ctx, input, opts...)
 		}
 	}
@@ -739,7 +774,7 @@ func (w *stateModelWrapper) wrapGenerateEndpoint(endpoint generateEndpoint) gene
 	return endpoint
 }
 
-func (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEndpoint {
+func (w *typedStateModelWrapper[M]) wrapStreamEndpoint(endpoint typedStreamEndpoint[M]) typedStreamEndpoint[M] {
 	hasUserEventSender := w.hasUserEventSender()
 	retryConfig := w.modelRetryConfig
 	failoverConfig := w.modelFailoverConfig
@@ -749,11 +784,11 @@ func (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEn
 		handler := w.handlers[i]
 		innerEndpoint := endpoint
 		baseToolInfos := w.toolInfos
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 			baseOpts := &model.Options{Tools: baseToolInfos}
 			commonOpts := model.GetCommonOptions(baseOpts, opts...)
 			mc := &ModelContext{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig, cancelContext: cc}
-			wrappedModel, err := handler.WrapModel(ctx, &endpointModel{stream: innerEndpoint}, mc)
+			wrappedModel, err := handler.WrapModel(ctx, &typedEndpointModel[M]{stream: innerEndpoint}, mc)
 			if err != nil {
 				return nil, err
 			}
@@ -763,14 +798,16 @@ func (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEn
 
 	if !hasUserEventSender {
 		innerEndpoint := endpoint
-		eventSender := NewEventSenderModelWrapper()
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
-			execCtx := getChatModelAgentExecCtx(ctx)
+		eventSender := &typedEventSenderModelWrapper[M]{
+			TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[M]{},
+		}
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
+			execCtx := getTypedChatModelAgentExecCtx[M](ctx)
 			if execCtx == nil || execCtx.generator == nil {
 				return innerEndpoint(ctx, input, opts...)
 			}
 			mc := &ModelContext{ModelRetryConfig: retryConfig, ModelFailoverConfig: failoverConfig, cancelContext: cc}
-			wrappedModel, err := eventSender.WrapModel(ctx, &endpointModel{stream: innerEndpoint}, mc)
+			wrappedModel, err := eventSender.WrapModel(ctx, &typedEndpointModel[M]{stream: innerEndpoint}, mc)
 			if err != nil {
 				return nil, err
 			}
@@ -780,18 +817,17 @@ func (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEn
 
 	if w.modelRetryConfig != nil {
 		innerEndpoint := endpoint
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
-			retryWrapper := newRetryModelWrapper(&endpointModel{stream: innerEndpoint}, w.modelRetryConfig)
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
+			retryWrapper := newTypedRetryModelWrapper[M](&typedEndpointModel[M]{stream: innerEndpoint}, w.modelRetryConfig)
 			return retryWrapper.Stream(ctx, input, opts...)
 		}
 	}
 
-	// Needs to handle failoverWrapper after retryWrapper
 	if w.modelFailoverConfig != nil {
 		config := w.modelFailoverConfig
 		innerEndpoint := endpoint
-		endpoint = func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
-			failoverWrapper := newFailoverModelWrapper(&endpointModel{stream: innerEndpoint}, config)
+		endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
+			failoverWrapper := newTypedFailoverModelWrapper[M](&typedEndpointModel[M]{stream: innerEndpoint}, config)
 			return failoverWrapper.Stream(ctx, input, opts...)
 		}
 	}
@@ -799,19 +835,22 @@ func (w *stateModelWrapper) wrapStreamEndpoint(endpoint streamEndpoint) streamEn
 	return endpoint
 }
 
-func (w *stateModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
-	var stateMessages []Message
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+func (w *typedStateModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
+	var stateMessages []M
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		stateMessages = st.Messages
 		return nil
 	})
 
-	state := &ChatModelAgentState{Messages: stateMessages}
+	state := &TypedChatModelAgentState[M]{Messages: stateMessages}
 
-	for _, m := range w.middlewares {
-		if m.BeforeChatModel != nil {
-			if err := m.BeforeChatModel(ctx, state); err != nil {
-				return nil, err
+	if msgState, ok := any(state).(*ChatModelAgentState); ok {
+		for _, m := range w.middlewares {
+			if m.BeforeChatModel != nil {
+				if err := m.BeforeChatModel(ctx, msgState); err != nil {
+					var zero M
+					return zero, err
+				}
 			}
 		}
 	}
@@ -823,11 +862,12 @@ func (w *stateModelWrapper) Generate(ctx context.Context, input []*schema.Messag
 		var err error
 		ctx, state, err = handler.BeforeModelRewriteState(ctx, state, mc)
 		if err != nil {
-			return nil, err
+			var zero M
+			return zero, err
 		}
 	}
 
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		st.Messages = state.Messages
 		return nil
 	})
@@ -835,14 +875,15 @@ func (w *stateModelWrapper) Generate(ctx context.Context, input []*schema.Messag
 	wrappedEndpoint := w.wrapGenerateEndpoint(w.inner.Generate)
 	result, err := wrappedEndpoint(ctx, state.Messages, opts...)
 	if err != nil {
-		return nil, err
+		var zero M
+		return zero, err
 	}
 
 	// Re-read State.Messages after Generate completes: when ShouldRetry uses
 	// PersistModifiedInputMessages, applyDecisionForRetry writes modified messages to State.
 	// We must pick up those changes before appending the model result.
 	if w.modelRetryConfig != nil && w.modelRetryConfig.ShouldRetry != nil {
-		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 			state.Messages = st.Messages
 			return nil
 		})
@@ -853,42 +894,49 @@ func (w *stateModelWrapper) Generate(ctx context.Context, input []*schema.Messag
 	for _, handler := range w.handlers {
 		ctx, state, err = handler.AfterModelRewriteState(ctx, state, mc)
 		if err != nil {
-			return nil, err
+			var zero M
+			return zero, err
 		}
 	}
 
-	for _, m := range w.middlewares {
-		if m.AfterChatModel != nil {
-			if err := m.AfterChatModel(ctx, state); err != nil {
-				return nil, err
+	if msgState, ok := any(state).(*ChatModelAgentState); ok {
+		for _, m := range w.middlewares {
+			if m.AfterChatModel != nil {
+				if err := m.AfterChatModel(ctx, msgState); err != nil {
+					var zero M
+					return zero, err
+				}
 			}
 		}
 	}
 
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		st.Messages = state.Messages
 		return nil
 	})
 
 	if len(state.Messages) == 0 {
-		return nil, errors.New("no messages left in state after model call")
+		var zero M
+		return zero, errors.New("no messages left in state after model call")
 	}
 	return state.Messages[len(state.Messages)-1], nil
 }
 
-func (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
-	var stateMessages []Message
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+func (w *typedStateModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
+	var stateMessages []M
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		stateMessages = st.Messages
 		return nil
 	})
 
-	state := &ChatModelAgentState{Messages: stateMessages}
+	state := &TypedChatModelAgentState[M]{Messages: stateMessages}
 
-	for _, m := range w.middlewares {
-		if m.BeforeChatModel != nil {
-			if err := m.BeforeChatModel(ctx, state); err != nil {
-				return nil, err
+	if msgState, ok := any(state).(*ChatModelAgentState); ok {
+		for _, m := range w.middlewares {
+			if m.BeforeChatModel != nil {
+				if err := m.BeforeChatModel(ctx, msgState); err != nil {
+					return nil, err
+				}
 			}
 		}
 	}
@@ -904,7 +952,7 @@ func (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message,
 		}
 	}
 
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		st.Messages = state.Messages
 		return nil
 	})
@@ -914,14 +962,14 @@ func (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message,
 	if err != nil {
 		return nil, err
 	}
-	result, err := schema.ConcatMessageStream(stream)
+	result, err := concatMessageStream(stream)
 	if err != nil {
 		return nil, err
 	}
 
 	// Re-read State.Messages after Stream completes: same rationale as in Generate above.
 	if w.modelRetryConfig != nil && w.modelRetryConfig.ShouldRetry != nil {
-		_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 			state.Messages = st.Messages
 			return nil
 		})
@@ -936,15 +984,17 @@ func (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message,
 		}
 	}
 
-	for _, m := range w.middlewares {
-		if m.AfterChatModel != nil {
-			if err := m.AfterChatModel(ctx, state); err != nil {
-				return nil, err
+	if msgState, ok := any(state).(*ChatModelAgentState); ok {
+		for _, m := range w.middlewares {
+			if m.AfterChatModel != nil {
+				if err := m.AfterChatModel(ctx, msgState); err != nil {
+					return nil, err
+				}
 			}
 		}
 	}
 
-	_ = compose.ProcessState(ctx, func(_ context.Context, st *State) error {
+	_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 		st.Messages = state.Messages
 		return nil
 	})
@@ -952,22 +1002,23 @@ func (w *stateModelWrapper) Stream(ctx context.Context, input []*schema.Message,
 	if len(state.Messages) == 0 {
 		return nil, errors.New("no messages left in state after model call")
 	}
-	return schema.StreamReaderFromArray([]*schema.Message{state.Messages[len(state.Messages)-1]}), nil
+	return schema.StreamReaderFromArray([]M{state.Messages[len(state.Messages)-1]}), nil
 }
 
-type endpointModel struct {
-	generate generateEndpoint
-	stream   streamEndpoint
+type typedEndpointModel[M messageType] struct {
+	generate typedGenerateEndpoint[M]
+	stream   typedStreamEndpoint[M]
 }
 
-func (m *endpointModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+func (m *typedEndpointModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
 	if m.generate != nil {
 		return m.generate(ctx, input, opts...)
 	}
-	return nil, errors.New("generate endpoint not set")
+	var zero M
+	return zero, errors.New("generate endpoint not set")
 }
 
-func (m *endpointModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+func (m *typedEndpointModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) {
 	if m.stream != nil {
 		return m.stream(ctx, input, opts...)
 	}

adk/wrappers_failover_test.go

@@ -22,6 +22,7 @@ import (
 	"sync/atomic"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/cloudwego/eino/components/model"
@@ -47,7 +48,7 @@ func TestBuildModelWrappers_FailoverProxyInner(t *testing.T) {
 		},
 	}
 
-	wrapped := buildModelWrappers(base, &modelWrapperConfig{
+	wrapped := buildModelWrappers[*schema.Message](base, &modelWrapperConfig{
 		failoverConfig: failoverCfg,
 	})
 
@@ -101,11 +102,11 @@ func TestStateModelWrapper_Generate_WithFailover(t *testing.T) {
 		},
 	}
 
-	wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+	wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 		failoverConfig: failoverCfg,
 	})
 
-	ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+	ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 		failoverLastSuccessModel: m1,
 	})
 	got, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -162,11 +163,11 @@ func TestStateModelWrapper_Stream_WithFailover(t *testing.T) {
 		},
 	}
 
-	wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+	wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 		failoverConfig: failoverCfg,
 	})
 
-	ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+	ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 		failoverLastSuccessModel: m1,
 	})
 	sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -179,3 +180,36 @@ func TestStateModelWrapper_Stream_WithFailover(t *testing.T) {
 	require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
 	require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
 }
+
+func TestFailoverAcceptsAgenticAgent(t *testing.T) {
+	ctx := context.Background()
+
+	m := &mockAgenticModel{
+		generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
+			return agenticMsg("ok"), nil
+		},
+	}
+
+	fallbackModel := &mockChatModelForAttack{
+		generateFn: func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
+			return schema.AssistantMessage("fallback", nil), nil
+		},
+	}
+
+	agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{
+		Name:        "FailoverAgent",
+		Description: "Agent with failover config",
+		Model:       m,
+		ModelFailoverConfig: &ModelFailoverConfig{
+			MaxRetries: 1,
+			ShouldFailover: func(ctx context.Context, outputMessage *schema.Message, outputErr error) bool {
+				return true
+			},
+			GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) {
+				return fallbackModel, nil, nil
+			},
+		},
+	})
+	require.NoError(t, err)
+	assert.NotNil(t, agent)
+}

adk/wrappers_retry_failover_test.go

@@ -78,12 +78,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -126,12 +126,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -177,12 +177,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -228,12 +228,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -278,12 +278,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -332,12 +332,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -386,12 +386,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -438,12 +438,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -478,12 +478,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -522,12 +522,12 @@ func TestRetryThenFailover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			retryConfig:    retryCfg,
 			failoverConfig: failoverCfg,
 		})
 
-		ctx = withChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{
+		ctx = withTypedChatModelAgentExecCtx[*schema.Message](ctx, &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -562,11 +562,11 @@ func TestErrStreamCanceled_Failover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
@@ -597,11 +597,11 @@ func TestErrStreamCanceled_Failover(t *testing.T) {
 			},
 		}
 
-		wrapped := buildModelWrappers(m1, &modelWrapperConfig{
+		wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{
 			failoverConfig: failoverCfg,
 		})
 
-		ctx := withChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
+		ctx := withTypedChatModelAgentExecCtx[*schema.Message](context.Background(), &chatModelAgentExecCtx{
 			failoverLastSuccessModel: m1,
 		})
 		_, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})

components/model/interface.go

@@ -22,7 +22,19 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-// BaseChatModel defines the core interface for all chat model implementations.
+// BaseModel is the generic base model interface parameterized by message type M.
+// It exposes two modes of interaction:
+//   - [BaseModel.Generate]: blocks until the model returns a complete response.
+//   - [BaseModel.Stream]: returns a [schema.StreamReader] that yields message
+//     chunks incrementally as the model generates them.
+type BaseModel[M any] interface {
+	Generate(ctx context.Context, input []M, opts ...Option) (M, error)
+	Stream(ctx context.Context, input []M, opts ...Option) (*schema.StreamReader[M], error)
+}
+
+// BaseChatModel is a backward-compatible type alias for BaseModel specialized
+// with *schema.Message. All existing code using model.BaseChatModel continues
+// to work without modification.
 //
 // It exposes two modes of interaction:
 //   - [BaseChatModel.Generate]: blocks until the model returns a complete response.
@@ -49,12 +61,8 @@ import (
 // Note: a [schema.StreamReader] can only be read once. If multiple consumers
 // need the stream, it must be copied before reading.
 //
-//go:generate  mockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model -source interface.go
-type BaseChatModel interface {
-	Generate(ctx context.Context, input []*schema.Message, opts ...Option) (*schema.Message, error)
-	Stream(ctx context.Context, input []*schema.Message, opts ...Option) (
-		*schema.StreamReader[*schema.Message], error)
-}
+//go:generate  mockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model github.com/cloudwego/eino/components/model BaseChatModel,ChatModel,ToolCallingChatModel
+type BaseChatModel = BaseModel[*schema.Message]
 
 // Deprecated: Use [ToolCallingChatModel] instead.
 //
@@ -85,19 +93,11 @@ type ChatModel interface {
 type ToolCallingChatModel interface {
 	BaseChatModel
 
-	// WithTools returns a new ToolCallingChatModel instance with the specified tools bound.
-	// This method does not modify the current instance, making it safer for concurrent use.
 	WithTools(tools []*schema.ToolInfo) (ToolCallingChatModel, error)
 }
 
-// AgenticModel defines the interface for agentic models that support AgenticMessage.
-// It provides methods for generating complete and streaming outputs, and supports
-// tool calling via the WithTools method.
-type AgenticModel interface {
-	Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.AgenticMessage, error)
-	Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...Option) (*schema.StreamReader[*schema.AgenticMessage], error)
-
-	// WithTools returns a new Model instance with the specified tools bound.
-	// This method does not modify the current instance, making it safer for concurrent use.
-	WithTools(tools []*schema.ToolInfo) (AgenticModel, error)
-}
+// AgenticModel is a type alias for BaseModel specialized with
+// *schema.AgenticMessage. Unlike ToolCallingChatModel, agentic models do NOT
+// expose a WithTools method; tools are passed at request time via the
+// model.WithTools option, consistent with how ChatModelAgent binds tools.
+type AgenticModel = BaseModel[*schema.AgenticMessage]

schema/agentic_message.go

@@ -17,13 +17,16 @@
 package schema
 
 import (
+	"bytes"
 	"context"
+	"encoding/gob"
 	"encoding/json"
 	"fmt"
 	"reflect"
 	"sort"
 	"strings"
 
+	"github.com/bytedance/sonic"
 	"github.com/eino-contrib/jsonschema"
 
 	"github.com/cloudwego/eino/internal"
@@ -420,6 +423,47 @@ type MCPListToolsItem struct {
 	InputSchema *jsonschema.Schema `json:"input_schema,omitempty"`
 }
 
+type mcpListToolsItemGob struct {
+	Name            string
+	Description     string
+	InputSchemaJSON []byte
+}
+
+func (m *MCPListToolsItem) GobEncode() ([]byte, error) {
+	g := mcpListToolsItemGob{
+		Name:        m.Name,
+		Description: m.Description,
+	}
+	if m.InputSchema != nil {
+		b, err := json.Marshal(m.InputSchema)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal MCPListToolsItem.InputSchema: %w", err)
+		}
+		g.InputSchemaJSON = b
+	}
+	var buf bytes.Buffer
+	if err := gob.NewEncoder(&buf).Encode(&g); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func (m *MCPListToolsItem) GobDecode(data []byte) error {
+	var g mcpListToolsItemGob
+	if err := gob.NewDecoder(bytes.NewReader(data)).Decode(&g); err != nil {
+		return err
+	}
+	m.Name = g.Name
+	m.Description = g.Description
+	if len(g.InputSchemaJSON) > 0 {
+		m.InputSchema = &jsonschema.Schema{}
+		if err := sonic.Unmarshal(g.InputSchemaJSON, m.InputSchema); err != nil {
+			return fmt.Errorf("failed to unmarshal MCPListToolsItem.InputSchema: %w", err)
+		}
+	}
+	return nil
+}
+
 type MCPToolApprovalRequest struct {
 	// ID is the approval request ID.
 	ID string `json:"id,omitempty"`
@@ -1335,7 +1379,6 @@ func concatAssistantGenTexts(texts []*AssistantGenText) (ret *AssistantGenText,
 		if err != nil {
 			return nil, err
 		}
-		ret.Extension = extensions.Interface()
 	}
 
 	if len(openaiExtensions) > 0 {
@@ -2029,7 +2072,11 @@ func (m *MCPToolResult) String() string {
 	sb.WriteString(fmt.Sprintf("      name: %s\n", m.Name))
 	sb.WriteString(fmt.Sprintf("      result: %s\n", m.Result))
 	if m.Error != nil {
-		sb.WriteString(fmt.Sprintf("      error: [%d] %s\n", *m.Error.Code, m.Error.Message))
+		if m.Error.Code != nil {
+			sb.WriteString(fmt.Sprintf("      error: [%d] %s\n", *m.Error.Code, m.Error.Message))
+		} else {
+			sb.WriteString(fmt.Sprintf("      error: %s\n", m.Error.Message))
+		}
 	}
 	return sb.String()
 }

schema/agentic_message_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestConcatAgenticMessages(t *testing.T) {
@@ -1639,3 +1640,41 @@ func TestNewContentBlock(t *testing.T) {
 		})
 	}
 }
+
+func TestNewContentBlockChunk_NilMeta(t *testing.T) {
+	require.NotPanics(t, func() {
+		block := NewContentBlockChunk(&AssistantGenText{Text: "test"}, nil)
+		require.NotNil(t, block)
+		assert.Nil(t, block.StreamingMeta)
+	}, "NewContentBlockChunk should handle nil meta without panic")
+}
+
+func TestConcatAssistantGenTexts_ExtensionOverwrite(t *testing.T) {
+	type testExtension struct {
+		Value string
+	}
+
+	texts := []*AssistantGenText{
+		{Text: "Hello ", Extension: &testExtension{Value: "ext1"}},
+		{Text: "world", Extension: &testExtension{Value: "ext2"}},
+	}
+
+	result, err := concatAssistantGenTexts(texts)
+	if err != nil {
+		t.Logf("Concat error (may be expected if ConcatSliceValue doesn't handle this type): %v", err)
+		t.Skip("Skipping: ConcatSliceValue doesn't support test type")
+	}
+	require.NotNil(t, result)
+
+	assert.Equal(t, "Hello world", result.Text)
+
+	if result.Extension != nil {
+		t.Logf("Extension type: %T, value: %v", result.Extension, result.Extension)
+		_, isSlice := result.Extension.([]*testExtension)
+		if isSlice {
+			t.Log("WARNING: Extension is a raw slice instead of a concatenated value. " +
+				"Line 1381 in agentic_message.go overwrites the ConcatSliceValue result " +
+				"with extensions.Interface(), discarding the concatenation.")
+		}
+	}
+}

schema/serialization.go

@@ -27,6 +27,8 @@ import (
 func init() {
 	RegisterName[*Message]("_eino_message")
 	RegisterName[[]*Message]("_eino_message_slice")
+	RegisterName[*AgenticMessage]("_eino_agentic_message")
+	RegisterName[[]*AgenticMessage]("_eino_agentic_message_slice")
 	RegisterName[Document]("_eino_document")
 	RegisterName[RoleType]("_eino_role_type")
 	RegisterName[ToolCall]("_eino_tool_call")

schema/tool_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/eino-contrib/jsonschema"
 	"github.com/smartystreets/goconvey/convey"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestParamsOneOfToJSONSchema(t *testing.T) {
@@ -181,3 +182,40 @@ func TestToolInfoSerialization(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, ti2, result)
 }
+
+func TestMCPToolResult_NilErrorCode(t *testing.T) {
+	result := &MCPToolResult{
+		CallID: "test-call",
+		Name:   "test-tool",
+		Result: "some result",
+		Error: &MCPToolCallError{
+			Code:    nil,
+			Message: "something went wrong",
+		},
+	}
+
+	require.NotPanics(t, func() {
+		s := result.String()
+		t.Logf("String output: %s", s)
+		assert.Contains(t, s, "something went wrong")
+	}, "BUG: MCPToolResult.String() should not panic when Error.Code is nil")
+}
+
+func TestMCPToolResult_WithErrorCode(t *testing.T) {
+	code := int64(500)
+	result := &MCPToolResult{
+		CallID: "test-call",
+		Name:   "test-tool",
+		Result: "",
+		Error: &MCPToolCallError{
+			Code:    &code,
+			Message: "internal server error",
+		},
+	}
+
+	require.NotPanics(t, func() {
+		s := result.String()
+		assert.Contains(t, s, "500")
+		assert.Contains(t, s, "internal server error")
+	})
+}

utils/callbacks/template.go

@@ -65,6 +65,7 @@ type HandlerHelper struct {
 	toolHandler             *ToolCallbackHandler
 	toolsNodeHandler        *ToolsNodeCallbackHandlers
 	agentHandler            *AgentCallbackHandler
+	agenticAgentHandler     *AgenticAgentCallbackHandler
 	agenticPromptHandler    *AgenticPromptCallbackHandler
 	agenticModelHandler     *AgenticModelCallbackHandler
 	agenticToolsNodeHandler *AgenticToolsNodeCallbackHandlers
@@ -154,6 +155,12 @@ func (c *HandlerHelper) Agent(handler *AgentCallbackHandler) *HandlerHelper {
 	return c
 }
 
+// AgenticAgent sets the agentic agent callback handler for the handler helper, which will be called when an agentic agent is executed.
+func (c *HandlerHelper) AgenticAgent(handler *AgenticAgentCallbackHandler) *HandlerHelper {
+	c.agenticAgentHandler = handler
+	return c
+}
+
 // Graph sets the graph handler for the handler helper, which will be called when the graph is executed.
 func (c *HandlerHelper) Graph(handler callbacks.Handler) *HandlerHelper {
 	c.composeTemplates[compose.ComponentOfGraph] = handler
@@ -206,6 +213,8 @@ func (c *handlerTemplate) OnStart(ctx context.Context, info *callbacks.RunInfo,
 		return c.agenticToolsNodeHandler.OnStart(ctx, info, convAgenticToolsNodeCallbackInput(input))
 	case adk.ComponentOfAgent:
 		return c.agentHandler.OnStart(ctx, info, adk.ConvAgentCallbackInput(input))
+	case adk.ComponentOfAgenticAgent:
+		return c.agenticAgentHandler.OnStart(ctx, info, adk.ConvTypedCallbackInput[*schema.AgenticMessage](input))
 	case compose.ComponentOfGraph,
 		compose.ComponentOfChain,
 		compose.ComponentOfLambda:
@@ -245,6 +254,8 @@ func (c *handlerTemplate) OnEnd(ctx context.Context, info *callbacks.RunInfo, ou
 		return c.agenticToolsNodeHandler.OnEnd(ctx, info, convAgenticToolsNodeCallbackOutput(output))
 	case adk.ComponentOfAgent:
 		return c.agentHandler.OnEnd(ctx, info, adk.ConvAgentCallbackOutput(output))
+	case adk.ComponentOfAgenticAgent:
+		return c.agenticAgentHandler.OnEnd(ctx, info, adk.ConvTypedCallbackOutput[*schema.AgenticMessage](output))
 	case compose.ComponentOfGraph,
 		compose.ComponentOfChain,
 		compose.ComponentOfLambda:
@@ -404,6 +415,10 @@ func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, t
 		if c.agentHandler != nil && c.agentHandler.Needed(ctx, info, timing) {
 			return true
 		}
+	case adk.ComponentOfAgenticAgent:
+		if c.agenticAgentHandler != nil && c.agenticAgentHandler.Needed(ctx, info, timing) {
+			return true
+		}
 	case compose.ComponentOfGraph,
 		compose.ComponentOfChain,
 		compose.ComponentOfLambda:
@@ -644,9 +659,14 @@ func convToolsNodeCallbackOutput(src callbacks.CallbackInput) []*schema.Message
 	}
 }
 
+// AgentCallbackHandler handles callbacks for agents using *schema.Message.
+// Use ComponentOfAgent to filter callback events to agent-related events.
 type AgentCallbackHandler struct {
+	// OnStart is called when an agent run begins. Return a modified context to propagate values.
 	OnStart func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context
-	OnEnd   func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context
+	// OnEnd is called when an agent run completes. The output's Events iterator should be
+	// consumed asynchronously to avoid blocking.
+	OnEnd func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context
 }
 
 func (ch *AgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
@@ -660,6 +680,27 @@ func (ch *AgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunI
 	}
 }
 
+// AgenticAgentCallbackHandler handles callbacks for agentic agents using *schema.AgenticMessage.
+// Use ComponentOfAgenticAgent to filter callback events to agentic-agent-related events.
+type AgenticAgentCallbackHandler struct {
+	// OnStart is called when an agentic agent run begins. Return a modified context to propagate values.
+	OnStart func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context
+	// OnEnd is called when an agentic agent run completes. The output's Events iterator should be
+	// consumed asynchronously to avoid blocking.
+	OnEnd func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context
+}
+
+func (ch *AgenticAgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool {
+	switch timing {
+	case callbacks.TimingOnStart:
+		return ch.OnStart != nil
+	case callbacks.TimingOnEnd:
+		return ch.OnEnd != nil
+	default:
+		return false
+	}
+}
+
 // AgenticPromptCallbackHandler is the handler for the agentic prompt callback.
 type AgenticPromptCallbackHandler struct {
 	// OnStart is the callback function for the start of the agentic prompt.

utils/callbacks/template_test.go

@@ -683,3 +683,125 @@ func TestHandlerTemplateWithAgentComponent(t *testing.T) {
 		assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart))
 	})
 }
+
+func TestAgenticAgentCallbackHandler(t *testing.T) {
+	t.Run("Needed returns correct values", func(t *testing.T) {
+		handler := &AgenticAgentCallbackHandler{
+			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
+				return ctx
+			},
+		}
+
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}
+
+		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
+		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
+	})
+
+	t.Run("Needed with OnEnd set", func(t *testing.T) {
+		handler := &AgenticAgentCallbackHandler{
+			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
+				return ctx
+			},
+		}
+
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}
+
+		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
+		assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
+	})
+
+	t.Run("Needed with nil handlers", func(t *testing.T) {
+		handler := &AgenticAgentCallbackHandler{}
+
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}
+
+		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart))
+		assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd))
+	})
+}
+
+func TestHandlerHelperWithAgenticAgent(t *testing.T) {
+	t.Run("AgenticAgent method sets handler correctly", func(t *testing.T) {
+		cnt := 0
+		tpl := NewHandlerHelper()
+		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
+			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
+				cnt++
+				return ctx
+			},
+			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
+				cnt++
+				return ctx
+			},
+		})
+
+		handler := tpl.Handler()
+		ctx := context.Background()
+		ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}, handler)
+
+		ctx = callbacks.OnStart[any](ctx, nil)
+		assert.Equal(t, 1, cnt)
+
+		callbacks.OnEnd[any](ctx, nil)
+		assert.Equal(t, 2, cnt)
+	})
+}
+
+func TestHandlerTemplateWithAgenticAgentComponent(t *testing.T) {
+	t.Run("OnStart routes to agentic agent handler", func(t *testing.T) {
+		called := false
+		tpl := NewHandlerHelper()
+		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
+			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
+				called = true
+				return ctx
+			},
+		})
+
+		handler := tpl.Handler()
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"}
+
+		handler.OnStart(ctx, info, &adk.TypedAgentCallbackInput[*schema.AgenticMessage]{})
+		assert.True(t, called)
+	})
+
+	t.Run("OnEnd routes to agentic agent handler", func(t *testing.T) {
+		called := false
+		tpl := NewHandlerHelper()
+		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
+			OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context {
+				called = true
+				return ctx
+			},
+		})
+
+		handler := tpl.Handler()
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"}
+
+		handler.OnEnd(ctx, info, &adk.TypedAgentCallbackOutput[*schema.AgenticMessage]{})
+		assert.True(t, called)
+	})
+
+	t.Run("Needed returns true for agentic agent component", func(t *testing.T) {
+		tpl := NewHandlerHelper()
+		tpl.AgenticAgent(&AgenticAgentCallbackHandler{
+			OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context {
+				return ctx
+			},
+		})
+
+		handler := tpl.Handler()
+		ctx := context.Background()
+		info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}
+
+		checker, ok := handler.(callbacks.TimingChecker)
+		assert.True(t, ok, "handler should implement TimingChecker")
+		assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart))
+	})
+}