diff --git a/adk/agent_tool.go b/adk/agent_tool.go index fde319cb..2c78a584 100644 --- a/adk/agent_tool.go +++ b/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, diff --git a/adk/agent_tool_test.go b/adk/agent_tool_test.go index cfedb24c..54c02ea9 100644 --- a/adk/agent_tool_test.go +++ b/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") + } +} diff --git a/adk/agentic_callback_integration_test.go b/adk/agentic_callback_integration_test.go new file mode 100644 index 00000000..689188fc --- /dev/null +++ b/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") +} diff --git a/adk/agentic_integration_test.go b/adk/agentic_integration_test.go new file mode 100644 index 00000000..eb665799 --- /dev/null +++ b/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") +} diff --git a/adk/agentic_react_test.go b/adk/agentic_react_test.go new file mode 100644 index 00000000..5896a65a --- /dev/null +++ b/adk/agentic_react_test.go @@ -0,0 +1,1143 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type agenticAgentEvent = TypedAgentEvent[*schema.AgenticMessage] + +func agenticToolCallMsg(toolName, callID, args string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{Name: toolName, CallID: callID, Arguments: args}, + }, + }, + } +} + +type sequentialAgenticModel struct { + responses []*schema.AgenticMessage + callCount int32 +} + +func (m *sequentialAgenticModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&m.callCount, 1) - 1 + if int(idx) >= len(m.responses) { + return nil, fmt.Errorf("sequentialAgenticModel: no more responses (call #%d)", idx) + } + return m.responses[idx], nil +} + +func (m *sequentialAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + result, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { defer w.Close(); w.Send(result, nil) }() + return r, nil +} + +type agenticEchoTool struct { + name string +} + +func (t *agenticEchoTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "echoes input"}, nil +} + +func (t *agenticEchoTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return "echo:" + argumentsInJSON, nil +} + +type agenticInterruptTool struct { + name string +} + +func (t *agenticInterruptTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "interrupts on first call, returns on resume"}, nil +} + +func (t *agenticInterruptTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + wasInterrupted, _, _ := tool.GetInterruptState[any](ctx) + if !wasInterrupted { + return "", tool.Interrupt(ctx, "need_approval") + } + isResume, hasData, data := tool.GetResumeContext[string](ctx) + if isResume && hasData { + return "approved:" + data, nil + } + return "resumed_no_data", nil +} + +type agenticArgCaptureTool struct { + name string + onInvoke func(args string) string +} + +func (t *agenticArgCaptureTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "captures args"}, nil +} + +func (t *agenticArgCaptureTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return t.onInvoke(argumentsInJSON), nil +} + +type agenticSignalTool struct { + name string + started chan struct{} + result string + done chan struct{} + once sync.Once +} + +func (t *agenticSignalTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "blocks until finish() is called"}, nil +} + +func (t *agenticSignalTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + t.once.Do(func() { t.done = make(chan struct{}) }) + select { + case t.started <- struct{}{}: + default: + } + <-t.done + return t.result, nil +} + +func (t *agenticSignalTool) finish() { + t.once.Do(func() { t.done = make(chan struct{}) }) + close(t.done) +} + +type agenticReactTestStore struct { + m map[string][]byte +} + +func (s *agenticReactTestStore) Set(_ context.Context, key string, value []byte) error { + s.m[key] = value + return nil +} + +func (s *agenticReactTestStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.m[key] + return v, ok, nil +} + +func newAgenticAgent(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool) TypedAgent[*schema.AgenticMessage] { + t.Helper() + config := &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: t.Name(), + Description: "test agentic agent", + Model: mdl, + } + if len(tools) > 0 { + config.ToolsConfig = ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: tools, + }, + } + } + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, config) + require.NoError(t, err) + return agent +} + +func newAgenticRunner(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool) *TypedRunner[*schema.AgenticMessage] { + t.Helper() + agent := newAgenticAgent(t, ctx, mdl, tools) + return NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) +} + +func newAgenticRunnerWithStore(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool, store CheckPointStore) *TypedRunner[*schema.AgenticMessage] { + t.Helper() + agent := newAgenticAgent(t, ctx, mdl, tools) + return NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) +} + +func drainAgenticEvents(iter *AsyncIterator[*agenticAgentEvent]) []*agenticAgentEvent { + var events []*agenticAgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + } + return events +} + +func lastAgenticEvent(events []*agenticAgentEvent) *agenticAgentEvent { + if len(events) == 0 { + return nil + } + return events[len(events)-1] +} + +func findInterruptEvent(events []*agenticAgentEvent) *agenticAgentEvent { + for _, ev := range events { + if ev.Action != nil && ev.Action.Interrupted != nil { + return ev + } + } + return nil +} + +func TestAgenticReact_BasicInvoke(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"hello"`), + agenticMsg("done: echo result received"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "test input")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "done: echo result received", agenticTextContent(last.Output.MessageOutput.Message)) + assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.callCount)) +} + +func TestAgenticReact_MultiTurnToolCalling(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"step1"`), + agenticToolCallMsg("echo", "call-2", `"step2"`), + agenticToolCallMsg("echo", "call-3", `"step3"`), + agenticMsg("all done"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "do three steps")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "all done", agenticTextContent(last.Output.MessageOutput.Message)) + assert.Equal(t, int32(4), atomic.LoadInt32(&mdl.callCount)) +} + +func TestAgenticReact_Stream(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"hello"`), + agenticMsg("stream done"), + }, + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + events := drainAgenticEvents(runner.Query(ctx, "stream test")) + + var finalText string + for _, ev := range events { + if ev.Output != nil && ev.Output.MessageOutput != nil { + msg, err := ev.Output.MessageOutput.GetMessage() + if err == nil && msg != nil { + txt := agenticTextContent(msg) + if txt != "" { + finalText = txt + } + } + } + } + + assert.Equal(t, "stream done", finalText) +} + +func TestAgenticReact_MaxIterations(t *testing.T) { + ctx := context.Background() + + t.Run("within_limit", func(t *testing.T) { + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "c1", `"1"`), + agenticToolCallMsg("echo", "c2", `"2"`), + agenticMsg("done within limit"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "go")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "done within limit", agenticTextContent(last.Output.MessageOutput.Message)) + }) + + t.Run("exceeded", func(t *testing.T) { + responses := make([]*schema.AgenticMessage, 25) + for i := range responses { + responses[i] = agenticToolCallMsg("echo", fmt.Sprintf("c%d", i), `"x"`) + } + + mdl := &sequentialAgenticModel{responses: responses} + config := &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "exceed-agent", + Description: "test max iterations exceeded", + Model: mdl, + MaxIterations: 3, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&agenticEchoTool{name: "echo"}}, + }, + }, + } + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, config) + require.NoError(t, err) + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + events := drainAgenticEvents(runner.Query(ctx, "go")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.NotNil(t, last.Err) + assert.ErrorIs(t, last.Err, ErrExceedMaxIterations) + }) +} + +func TestAgenticReact_ReturnDirectly(t *testing.T) { + t.Skip("returnDirectly for agentic agents depends on typed eventSenderToolHandler; not yet supported") +} + +func TestAgenticReact_CancelAfterChatModel(t *testing.T) { + ctx := context.Background() + + toolStarted := make(chan struct{}, 1) + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + switch count { + case 1: + return agenticToolCallMsg("slow", "c1", `"hi"`), nil + case 2: + return agenticToolCallMsg("slow", "c2", `"hi2"`), nil + default: + return agenticMsg("should not reach"), nil + } + }, + } + + slowTool := &agenticSignalTool{ + name: "slow", + started: toolStarted, + result: "slow result", + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{slowTool}) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("trigger cancel")}, + }, cancelOpt) + + <-toolStarted + + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + _ = handle.Wait() + }() + + time.Sleep(10 * time.Millisecond) + slowTool.finish() + + var capturedErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + capturedErr = ev.Err + } + } + require.Error(t, capturedErr, "expected CancelError event") + var cancelErr *CancelError + require.ErrorAs(t, capturedErr, &cancelErr) +} + +func TestAgenticReact_CancelAfterToolCalls(t *testing.T) { + ctx := context.Background() + + toolStarted := make(chan struct{}, 1) + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + if count == 1 { + return agenticToolCallMsg("slow", "c1", `"hi"`), nil + } + return agenticMsg("should not reach on second call"), nil + }, + } + + slowTool := &agenticSignalTool{ + name: "slow", + started: toolStarted, + result: "slow result", + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{slowTool}) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("trigger cancel")}, + }, cancelOpt) + + <-toolStarted + + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + _ = handle.Wait() + }() + + time.Sleep(10 * time.Millisecond) + slowTool.finish() + + var capturedErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + capturedErr = ev.Err + } + } + require.Error(t, capturedErr, "expected CancelError event") + var cancelErr *CancelError + require.ErrorAs(t, capturedErr, &cancelErr) + assert.Equal(t, int32(1), atomic.LoadInt32(&modelCallCount)) +} + +func TestAgenticReact_DoubleInterruptResume(t *testing.T) { + ctx := context.Background() + + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + switch count { + case 1: + return agenticToolCallMsg("approval_tool", "c1", `"first"`), nil + case 2: + return agenticToolCallMsg("approval_tool", "c2", `"second"`), nil + case 3: + return agenticMsg("all approved"), nil + default: + return nil, fmt.Errorf("unexpected call #%d", count) + } + }, + } + + store := &agenticReactTestStore{m: map[string][]byte{}} + runner := newAgenticRunnerWithStore(t, ctx, mdl, []tool.BaseTool{&agenticInterruptTool{name: "approval_tool"}}, store) + + events1 := drainAgenticEvents(runner.Query(ctx, "approve twice", WithCheckPointID("dbl-cp"))) + int1Event := findInterruptEvent(events1) + require.NotNil(t, int1Event, "expected first interrupt") + int1ID := int1Event.Action.Interrupted.InterruptContexts[0].ID + + iter2, err := runner.ResumeWithParams(ctx, "dbl-cp", &ResumeParams{ + Targets: map[string]any{int1ID: "approved_1"}, + }) + require.NoError(t, err) + + events2 := drainAgenticEvents(iter2) + int2Event := findInterruptEvent(events2) + require.NotNil(t, int2Event, "expected second interrupt") + int2ID := int2Event.Action.Interrupted.InterruptContexts[0].ID + + iter3, err := runner.ResumeWithParams(ctx, "dbl-cp", &ResumeParams{ + Targets: map[string]any{int2ID: "approved_2"}, + }) + require.NoError(t, err) + + events3 := drainAgenticEvents(iter3) + last := lastAgenticEvent(events3) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Contains(t, agenticTextContent(last.Output.MessageOutput.Message), "all approved") +} + +func TestAgenticReact_ChatModelAgent_NoTools(t *testing.T) { + ctx := context.Background() + + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("no tools response"), nil + }, + } + + runner := newAgenticRunner(t, ctx, mdl, nil) + events := drainAgenticEvents(runner.Query(ctx, "hello")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "no tools response", agenticTextContent(last.Output.MessageOutput.Message)) +} + +func TestAgenticReact_ChatModelAgent_ToolsReceiveArgs(t *testing.T) { + ctx := context.Background() + + var receivedArgs string + captureTool := &agenticArgCaptureTool{ + name: "capture", + onInvoke: func(args string) string { + receivedArgs = args + return "captured" + }, + } + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("capture", "c1", `{"foo":"bar"}`), + agenticMsg("done"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{captureTool}) + drainAgenticEvents(runner.Query(ctx, "call capture")) + + assert.Equal(t, `{"foo":"bar"}`, receivedArgs) +} + +func TestCoverage_AgenticReact_Streaming(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + streamFn: func(_ context.Context, input []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "streamed response"}), + }, + }, nil) + }() + return r, nil + }, + } + + echoTool := &agenticEchoTool{name: "echo"} + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "stream-react", + Description: "streaming agentic react", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{echoTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + iter := runner.Query(ctx, "stream me") + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + stream := event.Output.MessageOutput.MessageStream + for { + _, sErr := stream.Recv() + if sErr != nil { + break + } + } + } + events = append(events, event) + } + + require.NotEmpty(t, events) + assertAgenticEventRoleFields(t, events) +} + +func TestCoverage_ConcatMessageStream_Agentic(t *testing.T) { + t.Run("Success", func(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + }, nil) + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + }, nil) + }() + + result, err := concatMessageStream(r) + assert.NoError(t, err) + assert.NotNil(t, result) + }) + + t.Run("ErrorDuringRecv", func(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + w.Send(nil, fmt.Errorf("recv error")) + w.Close() + }() + + _, err := concatMessageStream(r) + assert.Error(t, err) + }) +} + +func TestCoverage_AgenticReact_InterruptResume(t *testing.T) { + ctx := context.Background() + + interruptTool := &agenticInterruptTool{name: "approval"} + + var callIdx int32 + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&callIdx, 1) + if idx == 1 { + return agenticToolCallMsg("approval", "call1", `{}`), nil + } + return agenticMsg("approved and done"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "interrupt-agent", + Description: "tests interrupt and resume", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{interruptTool}, + }, + }, + }) + require.NoError(t, err) + + store := newDTTestStore() + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("need approval"), + }, WithCheckPointID("cp-int")) + + 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, "should have interrupt event") + + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + require.NotEmpty(t, rootCauseID) + + resumeIter, err := runner.ResumeWithParams(ctx, "cp-int", &ResumeParams{ + Targets: map[string]any{rootCauseID: "approved"}, + }) + require.NoError(t, err) + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := resumeIter.Next() + if !ok { + break + } + events = append(events, event) + } + require.NotEmpty(t, events) +} + +func TestCoverage_AgenticMessageHasToolCalls(t *testing.T) { + t.Run("NilMessage", func(t *testing.T) { + assert.False(t, agenticMessageHasToolCalls(nil)) + }) + + t.Run("NoToolCalls", func(t *testing.T) { + msg := agenticMsg("just text") + assert.False(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("HasToolCalls", func(t *testing.T) { + msg := agenticToolCallMsg("tool1", "id1", `{}`) + assert.True(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("NilBlock", func(t *testing.T) { + msg := &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{nil}, + } + assert.False(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("ToolCallBlockNilFunctionToolCall", func(t *testing.T) { + msg := &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{ + {Type: schema.ContentBlockTypeFunctionToolCall, FunctionToolCall: nil}, + }, + } + assert.False(t, agenticMessageHasToolCalls(msg)) + }) +} + +func TestCoverage_ChatModelAgent_StreamError(t *testing.T) { + ctx := context.Background() + + testErr := errors.New("stream failed") + m := &mockAgenticModel{ + streamFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + return nil, testErr + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "stream-error-agent", + Description: "tests stream error", + Model: m, + }) + require.NoError(t, err) + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + iter := runner.Query(ctx, "trigger stream error") + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should propagate stream error") +} + +func TestCoverage_AgenticReact_GobStateRoundTrip(t *testing.T) { + ctx := context.Background() + + var callIdx int32 + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&callIdx, 1) + if idx == 1 { + return agenticToolCallMsg("interrupt_tool", "call1", `{}`), nil + } + return agenticMsg("completed"), nil + }, + } + + interruptTool := &agenticInterruptTool{name: "interrupt_tool"} + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "gob-test", + Description: "tests gob state round trip", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{interruptTool}, + }, + }, + }) + require.NoError(t, err) + + store := newDTTestStore() + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("test gob"), + }, WithCheckPointID("gob-cp")) + + var interrupted bool + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interrupted = true + interruptEvent = event + } + } + + if !interrupted || interruptEvent == nil { + t.Skip("no interrupt occurred, skipping gob round-trip test") + } + + _, exists, err := store.Get(ctx, "gob-cp") + assert.NoError(t, err) + assert.True(t, exists, "checkpoint should be saved") + + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + require.NotEmpty(t, rootCauseID) + + resumeIter, err := runner.ResumeWithParams(ctx, "gob-cp", &ResumeParams{ + Targets: map[string]any{rootCauseID: "approved"}, + }) + require.NoError(t, err) + + var resumed bool + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + resumed = true + } + } + assert.True(t, resumed, "should successfully resume from gob checkpoint") +} + +func TestCoverage_GetMessageFromTypedWrappedEvent_Agentic(t *testing.T) { + t.Run("NilOutput", func(t *testing.T) { + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{}, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Nil(t, msg) + }) + + t.Run("NonStreaming", func(t *testing.T) { + expected := agenticMsg("hello") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: expected, + }, + }, + }, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("StreamingAlreadyConcatenated", func(t *testing.T) { + expected := agenticMsg("already concatenated") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + concatenatedMessage: expected, + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + }, + }, + }, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("StreamingWithPriorError", func(t *testing.T) { + testErr := errors.New("prior stream error") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + }, + }, + }, + } + wrapper.StreamErr = testErr + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.Equal(t, testErr, err) + assert.Nil(t, msg) + }) +} + +func TestCoverage_GetMessageFromWrappedEvent_ErrorPaths(t *testing.T) { + t.Run("NilOutput", func(t *testing.T) { + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{}, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Nil(t, msg) + }) + + t.Run("NonStreaming", func(t *testing.T) { + expected := schema.AssistantMessage("hello", nil) + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: expected, + }, + }, + }, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("AlreadyConcatenated", func(t *testing.T) { + expected := schema.AssistantMessage("concatenated", nil) + wrapper := &agentEventWrapper{ + concatenatedMessage: expected, + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + }, + }, + }, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("PriorStreamError", func(t *testing.T) { + testErr := errors.New("prior error") + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + }, + }, + }, + } + wrapper.StreamErr = testErr + msg, err := getMessageFromWrappedEvent(wrapper) + assert.Equal(t, testErr, err) + assert.Nil(t, msg) + }) +} + +func TestCoverage_ConsumeStream_ErrorDuringRecv(t *testing.T) { + testErr := errors.New("stream recv error") + r, w := schema.Pipe[*schema.Message](2) + go func() { + w.Send(schema.AssistantMessage("partial", nil), nil) + w.Send(nil, testErr) + w.Close() + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.NotNil(t, wrapper.StreamErr) + assert.Nil(t, wrapper.concatenatedMessage) +} + +func TestCoverage_ConsumeStream_EmptyStream(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { w.Close() }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + require.NotNil(t, wrapper.StreamErr) + assert.Contains(t, wrapper.StreamErr.Error(), "no messages") +} + +func TestCoverage_ConsumeStream_MultipleMessages(t *testing.T) { + r, w := schema.Pipe[*schema.Message](3) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("chunk1", nil), nil) + w.Send(schema.AssistantMessage("chunk2", nil), nil) + w.Send(schema.AssistantMessage("chunk3", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.Nil(t, wrapper.StreamErr) + assert.NotNil(t, wrapper.concatenatedMessage) +} + +func TestCoverage_ConsumeStream_SingleMessage(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("single", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.Nil(t, wrapper.StreamErr) + require.NotNil(t, wrapper.concatenatedMessage) + assert.Equal(t, "single", wrapper.concatenatedMessage.Content) +} + +func TestCoverage_ConsumeStream_Idempotent(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("once", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + msg1 := wrapper.concatenatedMessage + + wrapper.consumeStream() + msg2 := wrapper.concatenatedMessage + + assert.Equal(t, msg1, msg2, "second call should be no-op") +} diff --git a/adk/agentic_test.go b/adk/agentic_test.go new file mode 100644 index 00000000..80f729e9 --- /dev/null +++ b/adk/agentic_test.go @@ -0,0 +1,1355 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type mockAgenticModel struct { + generateFn func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) + streamFn func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) +} + +func (m *mockAgenticModel) Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return m.generateFn(ctx, input, opts...) +} + +func (m *mockAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + if m.streamFn != nil { + return m.streamFn(ctx, input, opts...) + } + result, err := m.generateFn(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { defer w.Close(); w.Send(result, nil) }() + return r, nil +} + +type testAgenticMiddleware struct { + *TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + beforeFn func(context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) + afterFn func(context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) +} + +func (m *testAgenticMiddleware) BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m.beforeFn != nil { + return m.beforeFn(ctx, state, mc) + } + return ctx, state, nil +} + +func (m *testAgenticMiddleware) AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m.afterFn != nil { + return m.afterFn(ctx, state, mc) + } + return ctx, state, nil +} + +func TestAgenticChatModelAgentRun_NoTools(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello from agentic model"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticTestAgent", + Description: "Agentic test agent", + Instruction: "You are helpful.", + Model: m, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + } + iter := agent.Run(ctx, input) + require.NotNil(t, iter) + + event, ok := iter.Next() + assert.True(t, ok) + require.NotNil(t, event) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + assert.Equal(t, schema.AgenticRoleTypeAssistant, msg.Role) + assert.Len(t, msg.ContentBlocks, 1) + assert.Equal(t, "Hello from agentic model", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestAgenticChatModelAgentRun_WithTools(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Used tool and got result"}), + }, + } + + var receivedToolInfos []*schema.ToolInfo + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + o := model.GetCommonOptions(&model.Options{}, opts...) + receivedToolInfos = o.Tools + return agenticResponse, nil + }, + } + + dummyTool := newSlowTool("dummy_tool", 0, "ok") + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticToolAgent", + Description: "Agentic agent with tools", + Instruction: "You are helpful.", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{dummyTool}, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Call a tool"), + }, + } + iter := agent.Run(ctx, input) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + + _, ok = iter.Next() + assert.False(t, ok) + + require.Len(t, receivedToolInfos, 1) + assert.Equal(t, "dummy_tool", receivedToolInfos[0].Name) +} + +func TestAgenticChatModelAgentRun_Streaming(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: "AgenticStreamAgent", + Description: "Agentic streaming agent", + Instruction: "You are helpful.", + Model: m, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + EnableStreaming: true, + } + iter := agent.Run(ctx, input) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.NotNil(t, event.Output.MessageOutput.MessageStream) + event.Output.MessageOutput.MessageStream.Close() + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestDefaultAgenticGenModelInput(t *testing.T) { + ctx := context.Background() + + t.Run("WithInstruction", func(t *testing.T) { + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + msgs, err := newDefaultGenModelInput[*schema.AgenticMessage]()(ctx, "Be helpful", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, msgs[0].Role) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[1].Role) + }) + + t.Run("WithoutInstruction", func(t *testing.T) { + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + msgs, err := newDefaultGenModelInput[*schema.AgenticMessage]()(ctx, "", input) + assert.NoError(t, err) + assert.Len(t, msgs, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[0].Role) + }) +} + +func TestAgenticRunnerQuery(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "query response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "QueryAgent", + Description: "Query test agent", + Instruction: "Be helpful.", + Model: m, + }) + assert.NoError(t, err) + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + + iter := runner.Query(ctx, "What's up?") + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func agenticAssistantMessage(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } +} + +type mockAgenticRunnerAgent struct { + name string + description string + responses []*TypedAgentEvent[*schema.AgenticMessage] + callCount int + lastInput *TypedAgentInput[*schema.AgenticMessage] + enableStreaming bool +} + +func (a *mockAgenticRunnerAgent) Name(_ context.Context) string { return a.name } +func (a *mockAgenticRunnerAgent) Description(_ context.Context) string { return a.description } +func (a *mockAgenticRunnerAgent) Run(_ context.Context, input *TypedAgentInput[*schema.AgenticMessage], _ ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + a.callCount++ + a.lastInput = input + a.enableStreaming = input.EnableStreaming + + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + for _, event := range a.responses { + generator.Send(event) + if event.Action != nil && event.Action.Exit { + break + } + } + }() + return iterator +} + +type mockAgenticAgent struct { + name string + description string + responses []*TypedAgentEvent[*schema.AgenticMessage] +} + +func (a *mockAgenticAgent) Name(_ context.Context) string { return a.name } +func (a *mockAgenticAgent) Description(_ context.Context) string { return a.description } +func (a *mockAgenticAgent) Run(_ context.Context, _ *TypedAgentInput[*schema.AgenticMessage], _ ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + for _, event := range a.responses { + generator.Send(event) + if event.Action != nil && event.Action.Exit { + break + } + } + }() + return iterator +} + +type myAgenticAgent struct { + name string + runFn func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] + resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] +} + +func (m *myAgenticAgent) Name(_ context.Context) string { + if len(m.name) > 0 { + return m.name + } + return "myAgenticAgent" +} +func (m *myAgenticAgent) Description(_ context.Context) string { return "my agentic agent description" } +func (m *myAgenticAgent) Run(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + return m.runFn(ctx, input, options...) +} +func (m *myAgenticAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + return m.resumeFn(ctx, info, opts...) +} + +func TestAgenticChatModelAgentRun_WithMiddleware(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello from agentic agent"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + afterModelExecuted := false + + mw := &testAgenticMiddleware{ + beforeFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + state.Messages = append(state.Messages, schema.UserAgenticMessage("extra")) + return ctx, state, nil + }, + afterFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + assert.Len(t, state.Messages, 4) + afterModelExecuted = true + return ctx, state, nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticMiddlewareAgent", + Description: "Agentic agent with middleware", + Instruction: "You are helpful.", + Model: m, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{mw}, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + } + iter := agent.Run(ctx, input) + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.NotNil(t, event.Output.MessageOutput.Message) + assert.Equal(t, schema.AgenticRoleTypeAssistant, event.Output.MessageOutput.Message.Role) + _, ok = iter.Next() + assert.False(t, ok) + assert.True(t, afterModelExecuted) +} + +func TestAgenticAfterModel_NoTools_ModifyDoesNotAffectEvent(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("original content"), nil + }, + } + + var capturedMessages []*schema.AgenticMessage + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticAfterModelAgent", + Description: "Test AfterModelRewriteState", + Instruction: "You are helpful.", + Model: m, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{ + &testAgenticMiddleware{ + afterFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *ModelContext) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + capturedMessages = make([]*schema.AgenticMessage, len(state.Messages)) + copy(capturedMessages, state.Messages) + state.Messages = append(state.Messages, agenticAssistantMessage("appended content")) + return ctx, state, nil + }, + }, + }, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + assert.Equal(t, "original content", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iterator.Next() + assert.False(t, ok) + + assert.Len(t, capturedMessages, 3) +} + +func TestAgenticGetComposeOptions_WithChatModelOptions(t *testing.T) { + ctx := context.Background() + + var capturedTemperature float32 + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + options := model.GetCommonOptions(&model.Options{}, opts...) + if options.Temperature != nil { + capturedTemperature = *options.Temperature + } + return agenticAssistantMessage("response"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticOptionsAgent", + Description: "Test agent", + Model: m, + }) + assert.NoError(t, err) + + temp := float32(0.7) + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}}, + WithChatModelOptions([]model.Option{model.WithTemperature(temp)})) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, temp, capturedTemperature) +} + +func TestAgenticChatModelAgent_PrepareExecContextError(t *testing.T) { + ctx := context.Background() + + expectedErr := errors.New("tool info error") + errTool := &errorTool{infoErr: expectedErr} + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("response"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticErrToolAgent", + Description: "Test agent", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{errTool}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}}) + + event, ok := iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "tool info error") + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestAgenticChatModelAgentOutputKey(t *testing.T) { + t.Run("OutputKeyStoresInSession", func(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("Hello from agentic assistant."), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticOutputKeyAgent", + Description: "Test agent for output key", + Instruction: "You are helpful.", + Model: m, + OutputKey: "agent_output", + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + ctx, runCtx := initTypedRunCtx[*schema.AgenticMessage](ctx, "AgenticOutputKeyAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + msg := event.Output.MessageOutput.Message + assert.Equal(t, "Hello from agentic assistant.", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iterator.Next() + assert.False(t, ok) + + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "agent_output") + assert.Equal(t, "Hello from agentic assistant.", sessionValues["agent_output"]) + }) + + t.Run("OutputKeyWithStreamingStoresInSession", func(t *testing.T) { + ctx := context.Background() + + chunk1 := agenticAssistantMessage("Hello") + chunk2 := agenticAssistantMessage(", 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: "AgenticStreamOutputKeyAgent", + Description: "Test agent for streaming output key", + Instruction: "You are helpful.", + Model: m, + OutputKey: "agent_output", + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + EnableStreaming: true, + } + ctx, runCtx := initTypedRunCtx[*schema.AgenticMessage](ctx, "AgenticStreamOutputKeyAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) + }) + + t.Run("SetOutputToSessionAgenticMessage", func(t *testing.T) { + ctx := context.Background() + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}, + } + ctx, runCtx := initTypedRunCtx[*schema.AgenticMessage](ctx, "TestAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + msg := agenticAssistantMessage("Test response") + err := setOutputToSession(ctx, msg, nil, "test_output") + assert.NoError(t, err) + + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "test_output") + assert.Equal(t, "Test response", sessionValues["test_output"]) + }) +} + +func TestAgenticRunner_Run_WithStreaming(t *testing.T) { + ctx := context.Background() + + mockAgent_ := &mockAgenticRunnerAgent{ + name: "AgenticStreamRunnerAgent", + description: "Test agent for agentic runner streaming", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "AgenticStreamRunnerAgent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + agenticAssistantMessage("Streaming response"), + }), + }, + }, + }, + }, + } + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{EnableStreaming: true, Agent: mockAgent_}) + + msgs := []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello, agent!"), + } + + iterator := runner.Run(ctx, msgs) + + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, msgs, mockAgent_.lastInput.Messages) + assert.True(t, mockAgent_.enableStreaming) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "AgenticStreamRunnerAgent", event.AgentName) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestAgenticRunner_Query_WithStreaming(t *testing.T) { + ctx := context.Background() + + mockAgent_ := &mockAgenticRunnerAgent{ + name: "AgenticStreamQueryAgent", + description: "Test agent for agentic runner query streaming", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "AgenticStreamQueryAgent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + agenticAssistantMessage("Streaming query response"), + }), + }, + }, + }, + }, + } + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{EnableStreaming: true, Agent: mockAgent_}) + + iterator := runner.Query(ctx, "Test query") + + assert.Equal(t, 1, mockAgent_.callCount) + assert.Len(t, mockAgent_.lastInput.Messages, 1) + assert.True(t, mockAgent_.enableStreaming) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "AgenticStreamQueryAgent", event.AgentName) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestAgenticSimpleInterrupt(t *testing.T) { + data := "hello world" + agent := &myAgenticAgent{ + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + schema.UserAgenticMessage("hello "), + schema.UserAgenticMessage("world"), + }), + }, + }, + }) + intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, data) + intEvent.Action.Interrupted.Data = data + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + assert.True(t, info.WasInterrupted) + assert.Nil(t, info.InterruptState) + assert.True(t, info.EnableStreaming) + assert.Equal(t, data, info.Data) + + assert.True(t, info.IsResumeTarget) + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + generator.Close() + return iter + }, + } + store := newMyStore() + ctx := context.Background() + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + + 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, data, interruptEvent.Action.Interrupted.Data) + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID) + assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause) + assert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info) + assert.Equal(t, Address{{Type: AddressSegmentAgent, ID: "myAgenticAgent"}}, + interruptEvent.Action.Interrupted.InterruptContexts[0].Address) +} + +func TestCascadingFrom_NewChatModelAgentFrom(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "from response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "FromAgent", + Description: "Test cascading constructor", + Instruction: "Be helpful.", + Model: m, + }) + assert.NoError(t, err) + assert.Equal(t, "FromAgent", agent.Name(ctx)) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestCascadingTyped_TypedStatefulInterrupt(t *testing.T) { + ctx := context.Background() + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "test-agent") + + type myState struct { + Count int + } + + event := TypedStatefulInterrupt[*schema.AgenticMessage](ctx, "please confirm", &myState{Count: 42}) + require.NotNil(t, event) + require.NotNil(t, event.Action) + require.NotNil(t, event.Action.Interrupted) +} + +func TestCascadingTyped_EventFromAgenticMessage(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "hello"}), + }, + } + + event := EventFromAgenticMessage(msg, nil, schema.AgenticRoleTypeAssistant) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.Equal(t, msg, event.Output.MessageOutput.Message) + assert.False(t, event.Output.MessageOutput.IsStreaming) + assert.Equal(t, schema.RoleType(""), event.Output.MessageOutput.Role) + assert.Equal(t, schema.AgenticRoleTypeAssistant, event.Output.MessageOutput.AgenticRole) + assert.Empty(t, event.Output.MessageOutput.ToolName) +} + +// assertAgenticEventRoleFields asserts that all AgenticMessage events in the +// list have zero-valued Role and ToolName fields (which are *schema.Message-only), +// and that AgenticRole is populated with a non-zero value. +func assertAgenticEventRoleFields(t *testing.T, events []*TypedAgentEvent[*schema.AgenticMessage]) { + t.Helper() + for i, event := range events { + if event.Output == nil || event.Output.MessageOutput == nil { + continue + } + mo := event.Output.MessageOutput + assert.Equal(t, schema.RoleType(""), mo.Role, "event[%d]: AgenticMessage must have zero Role", i) + assert.Empty(t, mo.ToolName, "event[%d]: AgenticMessage must have empty ToolName", i) + assert.NotEmpty(t, mo.AgenticRole, "event[%d]: AgenticMessage must have non-zero AgenticRole", i) + } +} + +func TestCoverage_FlowAgent_ResumeNotResumable(t *testing.T) { + ctx := context.Background() + + agent := &mockAgenticAgent{ + name: "non-resumable", + description: "cannot resume", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + {Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("done"), + }, + }}, + }, + } + + fa := toTypedFlowAgent[*schema.AgenticMessage](agent) + + info := &ResumeInfo{WasInterrupted: true} + iter := fa.Resume(ctx, info) + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should get error for non-resumable agent") +} + +func TestCoverage_GenAgenticErrorIter(t *testing.T) { + testErr := errors.New("test agentic error") + iter := genAgenticErrorIter(testErr) + + event, ok := iter.Next() + require.True(t, ok) + assert.Equal(t, testErr, event.Err) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestCoverage_ChatModelAgent_OnSetSubAgents_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "freeze-test", + Description: "frozen test agent", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnSetSubAgents(ctx, []TypedAgent[*schema.AgenticMessage]{ + &mockAgenticAgent{name: "late-child"}, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_ChatModelAgent_OnSetAsSubAgent_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "freeze-child", + Description: "frozen child agent", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_ChatModelAgent_OnSetAsSubAgent_DuplicateError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "dup-child", + Description: "duplicate child agent", + Model: m, + }) + require.NoError(t, err) + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent1"}) + assert.NoError(t, err) + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent2"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already been set as a sub-agent") +} + +func TestCoverage_ChatModelAgent_OnDisallowTransferToParent_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "disallow-test", + Description: "disallow transfer test", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnDisallowTransferToParent(ctx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_TypedGetMessage_AgenticNonStreaming(t *testing.T) { + msg := agenticMsg("hello") + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: msg, + }, + }, + } + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.Equal(t, msg, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_TypedGetMessage_AgenticStreaming(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + }, nil) + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + }, nil) + }() + + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.NotNil(t, result) + require.NotNil(t, retEvent) + assert.NotNil(t, retEvent.Output.MessageOutput.MessageStream) +} + +func TestCoverage_TypedGetMessage_NilOutput(t *testing.T) { + event := &TypedAgentEvent[*schema.AgenticMessage]{} + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.Nil(t, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_GetMessage_NonStreaming(t *testing.T) { + msg := schema.AssistantMessage("hello", nil) + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: msg, + }, + }, + } + + result, retEvent, err := GetMessage(event) + assert.NoError(t, err) + assert.Equal(t, msg, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_GetMessage_Streaming(t *testing.T) { + r, w := schema.Pipe[*schema.Message](2) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("Hello ", nil), nil) + w.Send(schema.AssistantMessage("world", nil), nil) + }() + + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + result, retEvent, err := GetMessage(event) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotNil(t, retEvent) +} + +func TestCoverage_NewTypedAgentTool_Agentic(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("tool response"), nil + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "tool-agent", + Description: "agent wrapped as tool", + Model: m, + }) + require.NoError(t, err) + + agentTool := NewTypedAgentTool[*schema.AgenticMessage](ctx, agent) + + info, err := agentTool.Info(ctx) + require.NoError(t, err) + assert.Equal(t, "tool-agent", info.Name) + + result, err := agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"request":"test"}`) + require.NoError(t, err) + assert.Contains(t, result, "tool response") +} +func TestCoverage_CopyAgenticEvent(t *testing.T) { + original := &TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "agent1", + RunPath: []RunStep{{agentName: "root"}, {agentName: "agent1"}}, + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("hello"), + }, + }, + Action: &AgentAction{ + TransferToAgent: &TransferToAgentAction{DestAgentName: "agent2"}, + }, + } + + copied := copyTypedAgentEvent(original) + assert.Equal(t, original.AgentName, copied.AgentName) + assert.Equal(t, len(original.RunPath), len(copied.RunPath)) + assert.Equal(t, original.Action, copied.Action) + + copied.RunPath[0].agentName = "mutated" + assert.NotEqual(t, original.RunPath[0].agentName, copied.RunPath[0].agentName) +} + +func TestCoverage_ChatModelAgent_ModelGenerateError(t *testing.T) { + ctx := context.Background() + + testErr := errors.New("model generate failed") + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return nil, testErr + }, + } + + agent, err := NewTypedChatModelAgent[*schema.AgenticMessage](ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "error-model-agent", + Description: "tests model generate error", + Model: m, + }) + require.NoError(t, err) + + runner := NewTypedRunner[*schema.AgenticMessage](TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + + iter := runner.Query(ctx, "trigger error") + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should propagate model error") +} + +func TestCoverage_NewTypedUserMessages(t *testing.T) { + t.Run("Message", func(t *testing.T) { + msgs := newTypedUserMessages[*schema.Message]("hello") + require.Len(t, msgs, 1) + assert.Equal(t, schema.User, msgs[0].Role) + assert.Equal(t, "hello", msgs[0].Content) + }) + + t.Run("AgenticMessage", func(t *testing.T) { + msgs := newTypedUserMessages[*schema.AgenticMessage]("hello") + require.Len(t, msgs, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[0].Role) + }) +} + +func TestCoverage_TypedEndpointModel_NilEndpoints(t *testing.T) { + ctx := context.Background() + + m := &typedEndpointModel[*schema.AgenticMessage]{} + + _, err := m.Generate(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "generate endpoint not set") + + _, err = m.Stream(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "stream endpoint not set") +} + +func TestCoverage_TypedEndpointModel_WithEndpoints(t *testing.T) { + ctx := context.Background() + + expected := agenticMsg("generated") + m := &typedEndpointModel[*schema.AgenticMessage]{ + generate: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return expected, nil + }, + stream: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(expected, nil) + }() + return r, nil + }, + } + + result, err := m.Generate(ctx, nil) + assert.NoError(t, err) + assert.Equal(t, expected, result) + + stream, err := m.Stream(ctx, nil) + assert.NoError(t, err) + require.NotNil(t, stream) + msg, err := stream.Recv() + assert.NoError(t, err) + assert.Equal(t, expected, msg) + _, err = stream.Recv() + assert.Equal(t, io.EOF, err) +} + +func TestCoverage_SetAutomaticClose(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(agenticMsg("data"), nil) + }() + + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + typedSetAutomaticClose(event) +} + +func TestConcatMessageStream_AgenticClosesStream(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(agenticMsg("a"), nil) + w.Send(agenticMsg("b"), nil) + }() + + result, err := concatMessageStream(r) + require.NoError(t, err) + require.NotNil(t, result) + + _, recvErr := r.Recv() + assert.Error(t, recvErr, + "stream should be closed after concatMessageStream returns") +} diff --git a/adk/callback.go b/adk/callback.go index 19afbfc7..0b5cac87 100644 --- a/adk/callback.go +++ b/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...) +} diff --git a/adk/callback_test.go b/adk/callback_test.go index b54ea7ee..efd66f56 100644 --- a/adk/callback_test.go +++ b/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) + }) +} diff --git a/adk/cancel.go b/adk/cancel.go index 513b0cf4..72f3e109 100644 --- a/adk/cancel.go +++ b/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 diff --git a/adk/cancel_edge_test.go b/adk/cancel_edge_test.go index 248a84ee..946cd600 100644 --- a/adk/cancel_edge_test.go +++ b/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() diff --git a/adk/cancel_test.go b/adk/cancel_test.go index 97779827..e08a0f58 100644 --- a/adk/cancel_test.go +++ b/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"}) diff --git a/adk/chatmodel.go b/adk/chatmodel.go index f29493a2..2a817a06 100644 --- a/adk/chatmodel.go +++ b/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, diff --git a/adk/deterministic_transfer.go b/adk/deterministic_transfer.go index dc677a00..ce5b2009 100644 --- a/adk/deterministic_transfer.go +++ b/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) diff --git a/adk/failover_chatmodel.go b/adk/failover_chatmodel.go index 898aedd7..a0f60ea8 100644 --- a/adk/failover_chatmodel.go +++ b/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 } diff --git a/adk/failover_chatmodel_test.go b/adk/failover_chatmodel_test.go index 82866e99..75f87df3 100644 --- a/adk/failover_chatmodel_test.go +++ b/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) diff --git a/adk/flow.go b/adk/flow.go index 8edc002a..7011fa81 100644 --- a/adk/flow.go +++ b/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)) +} diff --git a/adk/handler.go b/adk/handler.go index d18abc96..255294dd 100644 --- a/adk/handler.go +++ b/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 diff --git a/adk/instruction.go b/adk/instruction.go index f02888ed..8794aff5 100644 --- a/adk/instruction.go +++ b/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, diff --git a/adk/interface.go b/adk/interface.go index e1f17eca..0a4c0bc5 100644 --- a/adk/interface.go +++ b/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") + } } diff --git a/adk/interrupt.go b/adk/interrupt.go index fce09d4c..35bfb519 100644 --- a/adk/interrupt.go +++ b/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" diff --git a/adk/prebuilt/planexecute/utils.go b/adk/prebuilt/planexecute/utils_test.go similarity index 100% rename from adk/prebuilt/planexecute/utils.go rename to adk/prebuilt/planexecute/utils_test.go diff --git a/adk/react.go b/adk/react.go index 07fdbde9..fdd224f7 100644 --- a/adk/react.go +++ b/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 +} diff --git a/adk/react_test.go b/adk/react_test.go index b0a6c398..1ac0ff5e 100644 --- a/adk/react_test.go +++ b/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") +} diff --git a/adk/retry_chatmodel.go b/adk/retry_chatmodel.go index 304e8b9b..df6fcad9 100644 --- a/adk/retry_chatmodel.go +++ b/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(¤tInput, ¤tOpts, ctx, decision) + applyDecisionForRetry(¤tInput, ¤tOpts, 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(¤tInput, ¤tOpts, ctx, decision) + applyDecisionForRetry(¤tInput, ¤tOpts, 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(¤tInput, ¤tOpts, ctx, decision) + applyDecisionForRetry(¤tInput, ¤tOpts, 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 } diff --git a/adk/runctx.go b/adk/runctx.go index 1a32f176..3c031601 100644 --- a/adk/runctx.go +++ b/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 { diff --git a/adk/runner.go b/adk/runner.go index 405b69e7..6caac130 100644 --- a/adk/runner.go +++ b/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)}) } } } diff --git a/adk/runner_test.go b/adk/runner_test.go index 6ab3f128..0eb797c8 100644 --- a/adk/runner_test.go +++ b/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") +} diff --git a/adk/turn_loop.go b/adk/turn_loop.go index df12ba40..46504c64 100644 --- a/adk/turn_loop.go +++ b/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, diff --git a/adk/turn_loop_test.go b/adk/turn_loop_test.go index 309c84f0..8a65a6c2 100644 --- a/adk/turn_loop_test.go +++ b/adk/turn_loop_test.go @@ -161,13 +161,13 @@ func (a *turnLoopStopModeProbeAgent) Run(ctx context.Context, input *AgentInput, return iter } -func newAndRunTurnLoop[T any](ctx context.Context, cfg TurnLoopConfig[T]) *TurnLoop[T] { - l := NewTurnLoop(cfg) +func newAndRunTurnLoop[T any, M messageType](ctx context.Context, cfg TurnLoopConfig[T, M]) *TurnLoop[T, M] { + l := NewTurnLoop[T, M](cfg) l.Run(ctx) return l } -func newPreemptTestLoop(t *testing.T, agent *turnLoopCancellableMockAgent) *TurnLoop[string] { +func newPreemptTestLoop(t *testing.T, agent *turnLoopCancellableMockAgent) *TurnLoop[string, *schema.Message] { t.Helper() agentStarted := make(chan struct{}) @@ -179,12 +179,12 @@ func newPreemptTestLoop(t *testing.T, agent *turnLoopCancellableMockAgent) *Turn return originalRunFunc(ctx, input) } - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -207,17 +207,17 @@ func TestTurnLoop_RunAndPush(t *testing.T) { processedItems := make([]string, 0) var mu sync.Mutex - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() processedItems = append(processedItems, items...) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -238,14 +238,14 @@ func TestTurnLoop_RunAndPush(t *testing.T) { } func TestTurnLoop_PushReturnsErrorAfterStop(t *testing.T) { - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -257,11 +257,11 @@ func TestTurnLoop_PushReturnsErrorAfterStop(t *testing.T) { } func TestTurnLoop_StopIsIdempotent(t *testing.T) { - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -275,11 +275,11 @@ func TestTurnLoop_StopIsIdempotent(t *testing.T) { } func TestTurnLoop_WaitMultipleGoroutines(t *testing.T) { - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -287,7 +287,7 @@ func TestTurnLoop_WaitMultipleGoroutines(t *testing.T) { loop.Stop() var wg sync.WaitGroup - results := make([]*TurnLoopExitState[string], 3) + results := make([]*TurnLoopExitState[string, *schema.Message], 3) for i := 0; i < 3; i++ { i := i @@ -308,17 +308,17 @@ func TestTurnLoop_UnhandledItemsOnStop(t *testing.T) { started := make(chan struct{}) blocked := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { close(started) <-blocked - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -339,11 +339,11 @@ func TestTurnLoop_UnhandledItemsOnStop(t *testing.T) { func TestTurnLoop_GenInputError(t *testing.T) { genErr := errors.New("gen input error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { return nil, genErr }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -357,11 +357,11 @@ func TestTurnLoop_GenInputError(t *testing.T) { func TestTurnLoop_GetAgentError(t *testing.T) { agentErr := errors.New("get agent error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return nil, agentErr }, }) @@ -376,19 +376,19 @@ func TestTurnLoop_BatchProcessing(t *testing.T) { var batches [][]string var mu sync.Mutex - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() batches = append(batches, items) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -409,11 +409,11 @@ func TestTurnLoop_BatchProcessing(t *testing.T) { } func TestTurnLoop_StopWithMode(t *testing.T) { - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -448,18 +448,18 @@ func TestTurnLoop_Preempt_CancelsCurrentAgent(t *testing.T) { secondGenInputCalled := make(chan struct{}) secondGenInputOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { count := atomic.AddInt32(&genInputCalls, 1) if count >= 2 { secondGenInputOnce.Do(func() { close(secondGenInputCalled) }) } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -475,7 +475,7 @@ func TestTurnLoop_Preempt_CancelsCurrentAgent(t *testing.T) { t.Fatal("agent did not start") } - loop.Push("urgent", WithPreempt[string](AnySafePoint)) + loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) select { case <-agentCancelled: @@ -528,16 +528,16 @@ func TestTurnLoop_Preempt_DiscardsConsumedItems(t *testing.T) { }, } - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() genInputResults = append(genInputResults, items) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -553,7 +553,7 @@ func TestTurnLoop_Preempt_DiscardsConsumedItems(t *testing.T) { t.Fatal("agent did not start") } - loop.Push("urgent", WithPreempt[string](AnySafePoint)) + loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) select { case <-agentDone: @@ -596,7 +596,7 @@ func TestTurnLoop_Preempt_WithAgentCancelMode(t *testing.T) { loop := newPreemptTestLoop(t, agent) - loop.Push("urgent", WithPreempt[string](AfterToolCalls)) + loop.Push("urgent", WithPreempt[string, *schema.Message](AfterToolCalls)) select { case <-cancelFuncCalled: @@ -632,7 +632,7 @@ func TestTurnLoop_PreemptAck_ClosesAfterCancelIsInitiated(t *testing.T) { loop := newPreemptTestLoop(t, agent) - ok, ack := loop.Push("urgent", WithPreempt[string](AfterToolCalls)) + ok, ack := loop.Push("urgent", WithPreempt[string, *schema.Message](AfterToolCalls)) assert.True(t, ok) assert.NotNil(t, ack) @@ -656,16 +656,16 @@ func TestTurnLoop_PreemptAck_ClosesAfterCancelIsInitiated(t *testing.T) { } func TestTurnLoop_PreemptAck_ClosesImmediatelyIfLoopNotStarted(t *testing.T) { - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) - ok, ack := loop.Push("urgent", WithPreempt[string](AnySafePoint)) + ok, ack := loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) assert.True(t, ok) assert.NotNil(t, ack) @@ -698,14 +698,14 @@ func TestTurnLoop_Preempt_EscalatesOnSecondPreempt(t *testing.T) { loop := newPreemptTestLoop(t, agent) - loop.Push("urgent1", WithPreempt[string](AfterChatModel)) + loop.Push("urgent1", WithPreempt[string, *schema.Message](AfterChatModel)) select { case <-firstCancelSeen: case <-time.After(1 * time.Second): t.Fatal("first preempt did not trigger cancel") } - loop.Push("urgent2", WithPreemptTimeout[string](AnySafePoint, time.Millisecond)) + loop.Push("urgent2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) wantMode := CancelAfterChatModel | CancelAfterToolCalls deadline := time.Now().Add(1 * time.Second) @@ -759,14 +759,14 @@ func TestTurnLoop_Preempt_JoinsSafePointModesOnSecondPreempt(t *testing.T) { loop := newPreemptTestLoop(t, agent) - loop.Push("urgent1", WithPreempt[string](AfterChatModel)) + loop.Push("urgent1", WithPreempt[string, *schema.Message](AfterChatModel)) select { case <-firstCancelSeen: case <-time.After(1 * time.Second): t.Fatal("first preempt did not trigger cancel") } - loop.Push("urgent2", WithPreempt[string](AfterToolCalls)) + loop.Push("urgent2", WithPreempt[string, *schema.Message](AfterToolCalls)) want := CancelAfterChatModel | CancelAfterToolCalls deadline := time.Now().Add(1 * time.Second) @@ -815,12 +815,12 @@ func TestTurnLoop_Push_WithoutPreempt_DoesNotCancel(t *testing.T) { }, } - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -879,12 +879,12 @@ func TestTurnLoop_PreemptDelay_NoMispreemptOnNaturalCompletion(t *testing.T) { }, } - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -900,7 +900,7 @@ func TestTurnLoop_PreemptDelay_NoMispreemptOnNaturalCompletion(t *testing.T) { t.Fatal("agent1 did not start") } - loop.Push("second", WithPreempt[string](AnySafePoint), WithPreemptDelay[string](500*time.Millisecond)) + loop.Push("second", WithPreempt[string, *schema.Message](AnySafePoint), WithPreemptDelay[string, *schema.Message](500*time.Millisecond)) select { case <-agent1Done: @@ -929,12 +929,12 @@ func TestTurnLoop_PreemptDelay_NoMispreemptOnNaturalCompletion(t *testing.T) { func TestTurnLoop_ConcurrentPush(t *testing.T) { var count int32 - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { atomic.AddInt32(&count, int32(len(items))) - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -967,14 +967,14 @@ func TestTurnLoop_StopAfterReceive_RecoverItem(t *testing.T) { receiveStarted := make(chan struct{}) cancelDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { close(receiveStarted) <-cancelDone time.Sleep(50 * time.Millisecond) - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -992,17 +992,17 @@ func TestTurnLoop_StopAfterReceive_RecoverItem(t *testing.T) { func TestTurnLoop_StopAfterGenInput_RecoverConsumed(t *testing.T) { genInputDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { close(genInputDone) time.Sleep(50 * time.Millisecond) - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { time.Sleep(100 * time.Millisecond) return &turnLoopMockAgent{name: "test"}, nil }, @@ -1023,15 +1023,15 @@ func TestTurnLoop_StopAfterGenInput_RecoverConsumed(t *testing.T) { func TestTurnLoop_GetAgentError_RecoverConsumed(t *testing.T) { agentErr := errors.New("get agent error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { return nil, agentErr }, }) @@ -1047,11 +1047,11 @@ func TestTurnLoop_GetAgentError_RecoverConsumed(t *testing.T) { func TestTurnLoop_GenInputError_RecoverItems(t *testing.T) { genErr := errors.New("gen input error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { return nil, genErr }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1069,8 +1069,8 @@ func TestTurnLoop_GenInputError_RecoverItems(t *testing.T) { func TestTurnLoop_PrepareAgentError_RecoverItemsInOrder(t *testing.T) { agentErr := errors.New("prepare agent error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { var urgent string remaining := make([]string, 0, len(items)) for _, item := range items { @@ -1081,19 +1081,19 @@ func TestTurnLoop_PrepareAgentError_RecoverItemsInOrder(t *testing.T) { } } if urgent != "" { - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: []string{urgent}, Remaining: remaining, }, nil } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return nil, agentErr }, }) @@ -1119,16 +1119,16 @@ func TestTurnLoop_ContextCancel(t *testing.T) { genInputStarted := make(chan struct{}) genInputDone := make(chan struct{}) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { close(genInputStarted) <-genInputDone if err := ctx.Err(); err != nil { return nil, err } - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1147,16 +1147,16 @@ func TestTurnLoop_ContextDeadlineExceeded(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { select { case <-time.After(100 * time.Millisecond): - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil case <-ctx.Done(): return nil, ctx.Err() } }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1171,11 +1171,11 @@ func TestTurnLoop_ContextCancelBeforeReceive(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1195,11 +1195,11 @@ func TestTurnLoop_ContextCancelDuringBlockingReceive(t *testing.T) { // the context monitoring goroutine closes the buffer, which unblocks Receive(). ctx, cancel := context.WithCancel(context.Background()) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1216,19 +1216,19 @@ func TestTurnLoop_ContextCancelAfterGenInput_RecoverItems(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) genInputCount := 0 - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { genInputCount++ if genInputCount == 1 { cancel() } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items[:1], Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], c []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { if err := ctx.Err(); err != nil { return nil, err } @@ -1249,17 +1249,17 @@ func TestTurnLoop_OnAgentEventsReceivesEvents(t *testing.T) { var receivedConsumed []string var mu sync.Mutex - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { mu.Lock() receivedConsumed = append(receivedConsumed, tc.Consumed...) mu.Unlock() @@ -1294,17 +1294,17 @@ func TestTurnLoop_OnAgentEventsReceivesEvents(t *testing.T) { func TestTurnLoop_StopDuringAgentExecution(t *testing.T) { agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) time.Sleep(200 * time.Millisecond) for { @@ -1340,15 +1340,15 @@ func TestTurnLoop_BareStop_AgentRunsToCompletion(t *testing.T) { turnsExecuted := int32(0) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "worker", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -1437,16 +1437,16 @@ func TestTurnLoop_StopCheckPointIDInCancelError(t *testing.T) { }) assert.NoError(t, err) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: checkpointID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -1490,15 +1490,15 @@ func TestTurnLoop_StopWithoutCheckpointIDDoesNotPersist(t *testing.T) { }) assert.NoError(t, err) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -1525,13 +1525,13 @@ func TestTurnLoop_StopWhileIdle_SkipsCheckpoint(t *testing.T) { } cpID := "idle-session" - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1551,13 +1551,13 @@ func TestTurnLoop_StopBetweenTurnsAndResume(t *testing.T) { store := &turnLoopCheckpointStore{m: make(map[string][]byte)} cpID := "between-turns-session" - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1572,22 +1572,22 @@ func TestTurnLoop_StopBetweenTurnsAndResume(t *testing.T) { var seen []string var mu sync.Mutex - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() seen = append([]string{}, items...) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -1633,16 +1633,16 @@ func TestTurnLoop_StopDuringAgentExecution_PersistAndResume(t *testing.T) { }) assert.NoError(t, err) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -1663,25 +1663,25 @@ func TestTurnLoop_StopDuringAgentExecution_PersistAndResume(t *testing.T) { var consumed2 []string var genResumeCalled bool var genInputCalled bool - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenResume: func(ctx context.Context, _ *TurnLoop[string], canceledItems []string, unhandledItems []string, newItems []string) (*GenResumeResult[string], error) { + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], canceledItems []string, unhandledItems []string, newItems []string) (*GenResumeResult[string, *schema.Message], error) { genResumeCalled = true - return &GenResumeResult[string]{ + return &GenResumeResult[string, *schema.Message]{ Consumed: canceledItems, Remaining: append(append([]string{}, unhandledItems...), newItems...), }, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { genInputCalled = true - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { consumed2 = append([]string{}, consumed...) return agent, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -1704,16 +1704,16 @@ func TestTurnLoop_StopDuringAgentExecution_PersistAndResume(t *testing.T) { func TestTurnLoop_CheckpointIDWithoutStore_FreshStart(t *testing.T) { ctx := context.Background() var genInputCalled bool - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ CheckpointID: "some-id", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { genInputCalled = true - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -1734,17 +1734,17 @@ func TestTurnLoop_CheckpointNotFound_FreshStart(t *testing.T) { ctx := context.Background() store := &turnLoopCheckpointStore{m: make(map[string][]byte)} var genInputCalled bool - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: "nonexistent-id", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { genInputCalled = true - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -1767,17 +1767,17 @@ func TestTurnLoop_CheckpointEmptyData_TreatedAsNoCheckpoint(t *testing.T) { store.m["cp-empty"] = nil var genInputCalled bool - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: "cp-empty", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { genInputCalled = true - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -1810,13 +1810,13 @@ func (s *errorCheckpointStore) Set(_ context.Context, _ string, _ []byte) error func TestTurnLoop_CheckpointLoadError_ReturnsError(t *testing.T) { ctx := context.Background() store := &errorCheckpointStore{getErr: fmt.Errorf("store unavailable")} - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: "cp-1", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1831,13 +1831,13 @@ func TestTurnLoop_CheckpointCorruptData_ReturnsError(t *testing.T) { ctx := context.Background() store := &turnLoopCheckpointStore{m: make(map[string][]byte)} store.m["cp-corrupt"] = []byte("not-valid-gob-data") - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: "cp-corrupt", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1869,16 +1869,16 @@ func TestTurnLoop_CheckpointSaveError_ReturnsError(t *testing.T) { }) assert.NoError(t, err) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: saveStore, CheckpointID: "cp-1", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -1897,13 +1897,13 @@ func TestTurnLoop_StaleCheckpointDeletion_OnCleanResume(t *testing.T) { store := &turnLoopCheckpointStore{m: make(map[string][]byte)} cpID := "stale-session" - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1917,19 +1917,19 @@ func TestTurnLoop_StaleCheckpointDeletion_OnCleanResume(t *testing.T) { store.mu.Unlock() assert.True(t, exists, "checkpoint should exist after first loop saves it") - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -1955,13 +1955,13 @@ func TestTurnLoop_StaleCheckpointDeletion_ContextCancel(t *testing.T) { store := &deletableCheckpointStore{turnLoopCheckpointStore: turnLoopCheckpointStore{m: make(map[string][]byte)}} cpID := "delete-on-cancel" - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -1976,19 +1976,19 @@ func TestTurnLoop_StaleCheckpointDeletion_ContextCancel(t *testing.T) { assert.True(t, exists, "checkpoint saved after loop1") ctx2, cancel2 := context.WithCancel(ctx) - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -2032,13 +2032,13 @@ func TestTurnLoop_CheckpointDeleter_CalledOnContextCancel(t *testing.T) { } cpID := "deleter-session" - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2053,19 +2053,19 @@ func TestTurnLoop_CheckpointDeleter_CalledOnContextCancel(t *testing.T) { assert.True(t, exists, "checkpoint saved after loop1") ctx2, cancel2 := context.WithCancel(ctx) - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -2111,16 +2111,16 @@ func TestTurnLoop_GenResumeNil_Error(t *testing.T) { }) assert.NoError(t, err) - loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -2129,13 +2129,13 @@ func TestTurnLoop_GenResumeNil_Error(t *testing.T) { loop1.Stop(WithImmediate()) loop1.Wait() - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2150,13 +2150,13 @@ func TestTurnLoop_SameCheckpointID_OverwritePattern(t *testing.T) { store := &turnLoopCheckpointStore{m: make(map[string][]byte)} cpID := "overwrite-session" - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2171,13 +2171,13 @@ func TestTurnLoop_SameCheckpointID_OverwritePattern(t *testing.T) { store.mu.Unlock() assert.NotEmpty(t, data1) - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2194,22 +2194,22 @@ func TestTurnLoop_SameCheckpointID_OverwritePattern(t *testing.T) { var seen []string var mu sync.Mutex - loop3 := NewTurnLoop(TurnLoopConfig[string]{ + loop3 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() seen = append([]string{}, items...) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -2243,13 +2243,13 @@ func TestTurnLoop_CheckpointHasRunnerStateButEmptyBytes(t *testing.T) { assert.NoError(t, err) store.m[cpID] = data - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2283,16 +2283,16 @@ func TestTurnLoop_GenResumeReturnsError(t *testing.T) { }) assert.NoError(t, err) - loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -2302,16 +2302,16 @@ func TestTurnLoop_GenResumeReturnsError(t *testing.T) { loop1.Wait() genResumeErr := fmt.Errorf("resume callback failed") - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - GenResume: func(ctx context.Context, _ *TurnLoop[string], canceled, unhandled, newItems []string) (*GenResumeResult[string], error) { + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], canceled, unhandled, newItems []string) (*GenResumeResult[string, *schema.Message], error) { return nil, genResumeErr }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2342,16 +2342,16 @@ func TestTurnLoop_CheckpointSaveError_MergesWithExistingError(t *testing.T) { }) assert.NoError(t, err) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: saveStore, CheckpointID: "cp-merge-err", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -2390,16 +2390,16 @@ func TestTurnLoop_ResumeWithParams(t *testing.T) { }) assert.NoError(t, err) - loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, }) @@ -2411,26 +2411,26 @@ func TestTurnLoop_ResumeWithParams(t *testing.T) { assert.True(t, errors.As(exit1.ExitReason, &ce)) var resumeParamsUsed *ResumeParams - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - GenResume: func(ctx context.Context, _ *TurnLoop[string], canceled, unhandled, newItems []string) (*GenResumeResult[string], error) { + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], canceled, unhandled, newItems []string) (*GenResumeResult[string, *schema.Message], error) { params := &ResumeParams{ Targets: map[string]any{"some-address": "user-data"}, } resumeParamsUsed = params - return &GenResumeResult[string]{ + return &GenResumeResult[string, *schema.Message]{ ResumeParams: params, Consumed: append(append(canceled, unhandled...), newItems...), }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -2451,14 +2451,14 @@ func TestTurnLoop_Stop_EscalatesCancelMode(t *testing.T) { ctx := context.Background() agentStarted := make(chan *cancelContext, 1) probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return probe, nil }, }) @@ -2491,14 +2491,14 @@ func TestTurnLoop_Stop_EscalatesCancelMode(t *testing.T) { func TestTurnLoop_DefaultOnAgentEvents_ErrorPropagation(t *testing.T) { agentErr := errors.New("agent execution error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -2519,17 +2519,17 @@ func TestTurnLoop_DefaultOnAgentEvents_ErrorPropagation(t *testing.T) { func TestTurnLoop_OnAgentEventsError(t *testing.T) { handlerErr := errors.New("event handler error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { // Drain events then return error for { _, ok := events.Next() @@ -2549,12 +2549,12 @@ func TestTurnLoop_OnAgentEventsError(t *testing.T) { func TestTurnLoop_StopCallFromGenInput(t *testing.T) { // Test that calling Stop() from within GenInput works correctly - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, loop *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { loop.Stop() - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2569,18 +2569,18 @@ func TestTurnLoop_PushFromOnAgentEvents(t *testing.T) { // Test that calling Push() from within OnAgentEvents works pushCount := int32(0) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -2613,17 +2613,17 @@ func TestNewTurnLoop_PushBeforeRun(t *testing.T) { var processedItems []string var mu sync.Mutex - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { mu.Lock() processedItems = append(processedItems, items...) mu.Unlock() - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2652,12 +2652,12 @@ func TestNewTurnLoop_PushBeforeRun(t *testing.T) { func TestNewTurnLoop_StopBeforeRun(t *testing.T) { // Stop before Run sets the stopped flag. When Run is called, the loop // exits immediately and buffered items appear as UnhandledItems. - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { t.Fatal("GenInput should not be called") return nil, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { t.Fatal("PrepareAgent should not be called") return nil, nil }, @@ -2680,16 +2680,16 @@ func TestNewTurnLoop_StopBeforeRun(t *testing.T) { func TestNewTurnLoop_WaitBeforeRun(t *testing.T) { // Wait blocks until Run is called AND the loop exits. - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) - waitDone := make(chan *TurnLoopExitState[string], 1) + waitDone := make(chan *TurnLoopExitState[string, *schema.Message], 1) go func() { waitDone <- loop.Wait() }() @@ -2718,12 +2718,12 @@ func TestNewTurnLoop_WaitBeforeRun(t *testing.T) { func TestNewTurnLoop_RunIsIdempotent(t *testing.T) { var genInputCalls int32 - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { atomic.AddInt32(&genInputCalls, 1) - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2744,12 +2744,12 @@ func TestNewTurnLoop_RunIsIdempotent(t *testing.T) { func TestNewTurnLoop_StopBeforeRun_ThenWait(t *testing.T) { // Demonstrates the full sequence: create, push, stop, run, wait. - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { t.Fatal("GenInput should not be called after Stop") return nil, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { t.Fatal("PrepareAgent should not be called after Stop") return nil, nil }, @@ -2773,12 +2773,12 @@ func TestNewTurnLoop_ConcurrentPushAndRun(t *testing.T) { for i := 0; i < 100; i++ { var count int32 - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { atomic.AddInt32(&count, int32(len(items))) - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -2820,17 +2820,17 @@ func TestTurnLoop_RunCtx_Propagation(t *testing.T) { const traceVal = "trace-123" var prepareCtxVal, agentCtxVal, eventsCtxVal string - cfg := TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, loop *TurnLoop[string], items []string) (*GenInputResult[string], error) { + cfg := TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { // Derive a new context with per-item trace data runCtx := context.WithValue(ctx, turnCtxKey{}, traceVal) - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ RunCtx: runCtx, Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, loop *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { if v, ok := ctx.Value(turnCtxKey{}).(string); ok { prepareCtxVal = v } @@ -2844,7 +2844,7 @@ func TestTurnLoop_RunCtx_Propagation(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { if v, ok := ctx.Value(turnCtxKey{}).(string); ok { eventsCtxVal = v } @@ -2873,14 +2873,14 @@ func TestTurnLoop_TurnContext_PreemptedChannel(t *testing.T) { preemptedSeen := make(chan struct{}) agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -2889,7 +2889,7 @@ func TestTurnLoop_TurnContext_PreemptedChannel(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) select { case <-tc.Preempted: @@ -2909,7 +2909,7 @@ func TestTurnLoop_TurnContext_PreemptedChannel(t *testing.T) { loop.Push("msg1") <-agentStarted - loop.Push("msg2", WithPreemptTimeout[string](AnySafePoint, time.Millisecond)) + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) select { case <-preemptedSeen: @@ -3092,13 +3092,13 @@ func TestTurnLoop_ConcurrentPreemptsDuringTurn(t *testing.T) { var genInputCount int32 - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { atomic.AddInt32(&genInputCount, 1) - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil @@ -3118,7 +3118,7 @@ func TestTurnLoop_ConcurrentPreemptsDuringTurn(t *testing.T) { wg.Add(1) go func(i int) { defer wg.Done() - ok, ack := loop.Push(fmt.Sprintf("urgent-%d", i), WithPreemptTimeout[string](AnySafePoint, 10*time.Millisecond)) + ok, ack := loop.Push(fmt.Sprintf("urgent-%d", i), WithPreemptTimeout[string, *schema.Message](AnySafePoint, 10*time.Millisecond)) if ok && ack != nil { select { case <-ack: @@ -3143,18 +3143,18 @@ func TestTurnLoop_PreemptDuringTurnTransition(t *testing.T) { firstTurnDone := make(chan struct{}) firstTurnOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "fast"}, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { count := atomic.AddInt32(&turnCount, 1) if count == 1 { firstTurnOnce.Do(func() { close(firstTurnDone) }) } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil @@ -3171,7 +3171,7 @@ func TestTurnLoop_PreemptDuringTurnTransition(t *testing.T) { time.Sleep(50 * time.Millisecond) - ok, ack := loop.Push("transitional", WithPreempt[string](AnySafePoint)) + ok, ack := loop.Push("transitional", WithPreempt[string, *schema.Message](AnySafePoint)) assert.True(t, ok, "push should succeed") if ack != nil { select { @@ -3213,18 +3213,18 @@ func TestTurnLoop_PushStrategy_DuringTurnTransition(t *testing.T) { secondTurnDone := make(chan struct{}) secondTurnOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { count := atomic.AddInt32(&genInputCount, 1) if count >= 2 { secondTurnOnce.Do(func() { close(secondTurnDone) }) } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil @@ -3243,12 +3243,12 @@ func TestTurnLoop_PushStrategy_DuringTurnTransition(t *testing.T) { var strategyTCNotNil int32 go func() { - loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { if tc != nil { atomic.StoreInt32(&strategyTCNotNil, 1) } <-strategyBlocker - return []PushOption[string]{WithPreempt[string](AnySafePoint)} + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} })) }() @@ -3288,12 +3288,12 @@ func TestTurnLoop_ConcurrentPreemptAndStop(t *testing.T) { }, } - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil @@ -3313,7 +3313,7 @@ func TestTurnLoop_ConcurrentPreemptAndStop(t *testing.T) { go func() { defer wg.Done() - _, ack := loop.Push("preempt-item", WithPreempt[string](AnySafePoint)) + _, ack := loop.Push("preempt-item", WithPreempt[string, *schema.Message](AnySafePoint)) if ack != nil { <-ack } @@ -3349,12 +3349,12 @@ func TestTurnLoop_ConcurrentPushStrategyAndStop(t *testing.T) { }, } - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{}, Consumed: items, }, nil @@ -3374,8 +3374,8 @@ func TestTurnLoop_ConcurrentPushStrategyAndStop(t *testing.T) { go func() { defer wg.Done() - _, ack := loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { - return []PushOption[string]{WithPreempt[string](AnySafePoint)} + _, ack := loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} })) if ack != nil { <-ack @@ -3397,14 +3397,14 @@ func TestTurnLoop_TurnContext_StoppedChannel(t *testing.T) { stoppedSeen := make(chan struct{}) agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -3413,7 +3413,7 @@ func TestTurnLoop_TurnContext_StoppedChannel(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) select { case <-tc.Stopped: @@ -3450,14 +3450,14 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { preemptedSeen := make(chan struct{}) agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -3466,7 +3466,7 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { close(agentStarted) select { case <-tc.Preempted: @@ -3485,7 +3485,7 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { loop.Push("msg1") <-agentStarted - loop.Push("msg2", WithPreemptTimeout[string](AnySafePoint, time.Millisecond)) + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) select { case <-preemptedSeen: @@ -3501,14 +3501,14 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { stoppedSeen := make(chan struct{}) agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -3517,7 +3517,7 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { close(agentStarted) select { case <-tc.Stopped: @@ -3544,7 +3544,7 @@ func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { t.Fatal("Stopped channel was never closed") } - loop.Push("msg2", WithPreemptTimeout[string](AnySafePoint, time.Millisecond)) + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) loop.Wait() }) } @@ -3573,18 +3573,18 @@ func TestTurnLoop_PushStrategy_DuringTurn(t *testing.T) { secondGenInputCalled := make(chan struct{}) secondGenInputOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { count := atomic.AddInt32(&genInputCalls, 1) if count >= 2 { secondGenInputOnce.Do(func() { close(secondGenInputCalled) }) } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -3602,11 +3602,11 @@ func TestTurnLoop_PushStrategy_DuringTurn(t *testing.T) { // Strategy inspects TurnContext during a running turn and decides to preempt. var strategyCalled int32 - var strategyTC *TurnContext[string] - loop.Push("urgent", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + var strategyTC *TurnContext[string, *schema.Message] + loop.Push("urgent", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { atomic.AddInt32(&strategyCalled, 1) strategyTC = tc - return []PushOption[string]{WithPreempt[string](AnySafePoint)} + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} })) select { @@ -3644,18 +3644,18 @@ func TestTurnLoop_PushStrategy_BetweenTurns(t *testing.T) { agentDone := make(chan struct{}) agentDoneOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, Remaining: nil, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -3670,7 +3670,7 @@ func TestTurnLoop_PushStrategy_BetweenTurns(t *testing.T) { }) // Push with strategy — no turn is active yet, so tc should be nil. - loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { atomic.AddInt32(&strategyCalled, 1) strategyTCWasNil = (tc == nil) return nil // plain push, no preempt @@ -3701,18 +3701,18 @@ func TestTurnLoop_PushStrategy_OverridesOtherOptions(t *testing.T) { agentDone := make(chan struct{}) agentDoneOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, Remaining: nil, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -3728,7 +3728,7 @@ func TestTurnLoop_PushStrategy_OverridesOtherOptions(t *testing.T) { // Strategy returns nil (no preempt), even though WithPreempt is also passed. // The strategy should override — so the agent should NOT be preempted. - ok, ack := loop.Push("item", WithPreempt[string](AnySafePoint), WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + ok, ack := loop.Push("item", WithPreempt[string, *schema.Message](AnySafePoint), WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { return nil // no preempt })) assert.True(t, ok) @@ -3755,18 +3755,18 @@ func TestTurnLoop_PushStrategy_NestedStrategyStripped(t *testing.T) { agentDone := make(chan struct{}) agentDoneOnce := sync.Once{} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, Remaining: nil, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { _, ok := events.Next() if !ok { @@ -3782,11 +3782,11 @@ func TestTurnLoop_PushStrategy_NestedStrategyStripped(t *testing.T) { // Strategy returns another WithPushStrategy — the nested one should be stripped. innerCalled := int32(0) - ok, ack := loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { - return []PushOption[string]{ - WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + ok, ack := loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + return []PushOption[string, *schema.Message]{ + WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { atomic.AddInt32(&innerCalled, 1) - return []PushOption[string]{WithPreempt[string](AnySafePoint)} + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} }), } })) @@ -3824,11 +3824,11 @@ func TestTurnLoop_PushStrategy_ConsumedInspection(t *testing.T) { genInputCalls := int32(0) secondGenInputItems := make(chan []string, 1) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return agent, nil }, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { count := atomic.AddInt32(&genInputCalls, 1) if count >= 2 { select { @@ -3836,7 +3836,7 @@ func TestTurnLoop_PushStrategy_ConsumedInspection(t *testing.T) { default: } } - return &GenInputResult[string]{ + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: []string{items[0]}, Remaining: items[1:], @@ -3853,9 +3853,9 @@ func TestTurnLoop_PushStrategy_ConsumedInspection(t *testing.T) { } // Strategy checks Consumed and preempts because current turn has "low-priority" items. - loop.Push("urgent-task", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string]) []PushOption[string] { + loop.Push("urgent-task", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { if tc != nil && len(tc.Consumed) > 0 && tc.Consumed[0] == "low-priority-task" { - return []PushOption[string]{WithPreempt[string](AnySafePoint)} + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} } return nil })) @@ -3875,17 +3875,17 @@ func TestTurnLoop_PushAfterStop_BufferedAsLateItems(t *testing.T) { ctx := context.Background() processed := make(chan string, 10) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -3916,11 +3916,11 @@ func TestTurnLoop_PushAfterStop_BufferedAsLateItems(t *testing.T) { func TestTurnLoop_TakeLateItems_Idempotent(t *testing.T) { ctx := context.Background() - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -3943,11 +3943,11 @@ func TestTurnLoop_TakeLateItems_Idempotent(t *testing.T) { func TestTurnLoop_PushAfterTakeLateItems_Panics(t *testing.T) { ctx := context.Background() - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -3966,11 +3966,11 @@ func TestTurnLoop_PushAfterTakeLateItems_Panics(t *testing.T) { func TestTurnLoop_TakeLateItems_NeverCalled_NoImpact(t *testing.T) { ctx := context.Background() - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -3989,13 +3989,13 @@ func TestTurnLoop_CheckpointErr_SeparateFromExitReason(t *testing.T) { ctx := context.Background() saveStore := &errorCheckpointStore{setErr: fmt.Errorf("storage unavailable")} - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: saveStore, CheckpointID: "cp-separate-err", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4014,11 +4014,11 @@ func TestTurnLoop_CheckpointErr_SeparateFromExitReason(t *testing.T) { func TestTurnLoop_CheckpointAttempted_FalseWhenNoStore(t *testing.T) { ctx := context.Background() - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4038,20 +4038,20 @@ func TestTurnLoop_CheckpointAttempted_FalseOnErrorExit(t *testing.T) { firstTurnDone := make(chan struct{}) var callCount int32 - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: "cp-err-exit", - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { n := atomic.AddInt32(&callCount, 1) if n > 1 { return nil, genInputErr } - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -4082,16 +4082,16 @@ func TestTurnLoop_StopConcurrentWithCallbackError_NoCheckpoint(t *testing.T) { stopCalled := make(chan struct{}) var prepareCount int32 - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { n := atomic.AddInt32(&prepareCount, 1) if n > 1 { // Wait until Stop() has been called so stopSig.isStopped() is true @@ -4100,7 +4100,7 @@ func TestTurnLoop_StopConcurrentWithCallbackError_NoCheckpoint(t *testing.T) { } return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -4140,13 +4140,13 @@ func TestTurnLoop_DeleteWithoutCheckPointDeleter_NoOp(t *testing.T) { cpID := "no-deleter" // First loop: save a checkpoint - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4163,19 +4163,19 @@ func TestTurnLoop_DeleteWithoutCheckPointDeleter_NoOp(t *testing.T) { // Second loop: exit via context cancel — should try to delete but store // doesn't implement CheckPointDeleter, so checkpoint persists (no-op) ctx2, cancel2 := context.WithCancel(ctx) - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { for { if _, ok := events.Next(); !ok { break @@ -4202,13 +4202,13 @@ func TestTurnLoop_StopWithSkipCheckpoint(t *testing.T) { store := &turnLoopCheckpointStore{m: make(map[string][]byte)} cpID := "skip-cp-session" - loop := NewTurnLoop(TurnLoopConfig[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4235,13 +4235,13 @@ func TestTurnLoop_StopWithSkipCheckpoint_DeletesStaleCheckpoint(t *testing.T) { } cpID := "skip-stale-session" - loop1 := NewTurnLoop(TurnLoopConfig[string]{ + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4256,13 +4256,13 @@ func TestTurnLoop_StopWithSkipCheckpoint_DeletesStaleCheckpoint(t *testing.T) { store.mu.Unlock() assert.True(t, exists, "first loop should save checkpoint") - loop2 := NewTurnLoop(TurnLoopConfig[string]{ + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4282,11 +4282,11 @@ func TestTurnLoop_StopWithStopCause(t *testing.T) { ctx := context.Background() cause := "user session timeout" - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4301,11 +4301,11 @@ func TestTurnLoop_StopWithStopCause(t *testing.T) { func TestTurnLoop_StopCause_EmptyWhenNoStop(t *testing.T) { ctx := context.Background() - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{Input: &AgentInput{}, Consumed: items}, nil + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -4320,14 +4320,14 @@ func TestTurnLoop_StopCause_InTurnContext(t *testing.T) { gotCause := make(chan string, 1) agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4336,7 +4336,7 @@ func TestTurnLoop_StopCause_InTurnContext(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) select { case <-tc.Stopped: @@ -4371,14 +4371,14 @@ func TestTurnLoop_StopCause_InTurnContext(t *testing.T) { func TestTurnLoop_StopCause_FirstNonEmptyWins(t *testing.T) { agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4387,7 +4387,7 @@ func TestTurnLoop_StopCause_FirstNonEmptyWins(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) for { if _, ok := events.Next(); !ok { @@ -4408,12 +4408,12 @@ func TestTurnLoop_StopCause_FirstNonEmptyWins(t *testing.T) { } func TestTurnLoop_StopBeforeRun_PushThenStop(t *testing.T) { - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { t.Fatal("GenInput should not be called when Stop is called before Run") return nil, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { t.Fatal("PrepareAgent should not be called when Stop is called before Run") return nil, nil }, @@ -4435,12 +4435,12 @@ func TestTurnLoop_StopBeforeRun_PushThenStop(t *testing.T) { } func TestTurnLoop_StopBeforeRun_StopThenPush(t *testing.T) { - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { t.Fatal("GenInput should not be called when Stop is called before Run") return nil, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { t.Fatal("PrepareAgent should not be called when Stop is called before Run") return nil, nil }, @@ -4468,16 +4468,16 @@ func TestTurnLoop_SkipCheckpoint_Sticky(t *testing.T) { store := &turnLoopCheckpointStore{m: make(map[string][]byte)} cpID := "sticky-skip-session" - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ Store: store, CheckpointID: cpID, - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "slow", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4486,7 +4486,7 @@ func TestTurnLoop_SkipCheckpoint_Sticky(t *testing.T) { }, }, nil }, - OnAgentEvents: func(ctx context.Context, tc *TurnContext[string], events *AsyncIterator[*AgentEvent]) error { + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { close(agentStarted) for { if _, ok := events.Next(); !ok { @@ -4520,12 +4520,12 @@ func TestWithGracefulTimeout_NonPositive_Panics(t *testing.T) { func TestWithPreempt_ZeroSafePoint_Panics(t *testing.T) { assert.PanicsWithValue(t, "adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint", - func() { WithPreempt[string](SafePoint(0)) }) + func() { WithPreempt[string, *schema.Message](SafePoint(0)) }) } func TestWithPreemptTimeout_ZeroSafePoint_Panics(t *testing.T) { assert.PanicsWithValue(t, "adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint", - func() { WithPreemptTimeout[string](SafePoint(0), time.Second) }) + func() { WithPreemptTimeout[string, *schema.Message](SafePoint(0), time.Second) }) } func TestSafePoint_ToCancelMode(t *testing.T) { @@ -4536,13 +4536,15 @@ func TestSafePoint_ToCancelMode(t *testing.T) { func TestNewTurnLoop_NilGenInput_Panics(t *testing.T) { assert.PanicsWithValue(t, "adk: NewTurnLoop: GenInput is required", func() { - NewTurnLoop(TurnLoopConfig[string]{PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { return nil, nil }}) + NewTurnLoop(TurnLoopConfig[string, *schema.Message]{PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return nil, nil + }}) }) } func TestNewTurnLoop_NilPrepareAgent_Panics(t *testing.T) { assert.PanicsWithValue(t, "adk: NewTurnLoop: PrepareAgent is required", func() { - NewTurnLoop(TurnLoopConfig[string]{GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { + NewTurnLoop(TurnLoopConfig[string, *schema.Message]{GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { return nil, nil }}) }) @@ -4556,14 +4558,14 @@ func TestDeriveChild_NilParent_ReturnsNil(t *testing.T) { func TestUntilIdleFor(t *testing.T) { t.Run("FiresAfterIdleDuration", func(t *testing.T) { turnDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4595,14 +4597,14 @@ func TestUntilIdleFor(t *testing.T) { t.Run("ResetsOnPush", func(t *testing.T) { turnCount := int32(0) turnDone := make(chan struct{}, 10) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4641,14 +4643,14 @@ func TestUntilIdleFor(t *testing.T) { t.Run("EscalatedByStopWithImmediate", func(t *testing.T) { agentStarted := make(chan *cancelContext, 1) probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return probe, nil }, }) @@ -4681,14 +4683,14 @@ func TestUntilIdleFor(t *testing.T) { t.Run("EscalatedByStopWithGraceful", func(t *testing.T) { agentStarted := make(chan struct{}) agentDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4729,14 +4731,14 @@ func TestUntilIdleFor_DoesNotCancelRunningAgent(t *testing.T) { agentCtxCanceled := int32(0) agentDone := make(chan struct{}) - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4773,14 +4775,14 @@ func TestUntilIdleFor_DoesNotCancelRunningAgent(t *testing.T) { agentCtxCanceled := int32(0) agentDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4818,14 +4820,14 @@ func TestUntilIdleFor_DoesNotCancelRunningAgent(t *testing.T) { agentCtxCanceled := int32(0) agentDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4860,14 +4862,14 @@ func TestUntilIdleFor_ContextCancelDuringIdleWait(t *testing.T) { turnDone := make(chan struct{}) ctx, cancel := context.WithCancel(context.Background()) - loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -4967,14 +4969,14 @@ func TestAttack_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { turnCount := int32(0) turnDone := make(chan struct{}, 10) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5015,14 +5017,14 @@ func TestAttack_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { func TestAttack_UntilIdleFor_MultipleStopCallsFirstWins(t *testing.T) { turnDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5056,14 +5058,14 @@ func TestAttack_BareStopOverridesUntilIdleFor(t *testing.T) { agentStarted := make(chan struct{}) agentDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5102,14 +5104,14 @@ func TestAttack_BareStopOverridesUntilIdleFor(t *testing.T) { func TestAttack_StopSignal_NilCancelOptsDoNotDeescalate(t *testing.T) { agentStarted := make(chan *cancelContext, 1) probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return probe, nil }, }) @@ -5135,14 +5137,14 @@ func TestAttack_StopSignal_NilCancelOptsDoNotDeescalate(t *testing.T) { func TestAttack_CanceledItems_EmptyWhenAgentFinishesNormally(t *testing.T) { agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5207,14 +5209,14 @@ func TestAttack_TurnBuffer_ClearWakeupPreventsSpuriousReturn(t *testing.T) { } func TestAttack_StopBeforeRun_UntilIdleFor_ExitsImmediately(t *testing.T) { - loop := NewTurnLoop(TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{name: "test"}, nil }, }) @@ -5239,14 +5241,14 @@ func TestAttack_StopBeforeRun_UntilIdleFor_ExitsImmediately(t *testing.T) { func TestAttack_PushAfterStop_UntilIdleFor_RoutedToLateItems(t *testing.T) { turnDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5273,14 +5275,14 @@ func TestAttack_PushAfterStop_UntilIdleFor_RoutedToLateItems(t *testing.T) { func TestAttack_ConcurrentStopEscalation_RaceDetector(t *testing.T) { agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5320,14 +5322,14 @@ func TestAttack_ConcurrentStopEscalation_RaceDetector(t *testing.T) { func TestAttack_StopCause_FirstNonEmptyWins_ConcurrentCallers(t *testing.T) { turnDone := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5350,14 +5352,14 @@ func TestAttack_StopCause_FirstNonEmptyWins_ConcurrentCallers(t *testing.T) { func TestAttack_SkipCheckpoint_Sticky(t *testing.T) { agentStarted := make(chan struct{}) - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return &turnLoopCancellableMockAgent{ name: "test", runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { @@ -5424,14 +5426,14 @@ func TestTurnLoop_Stop_WithImmediate_RecursivePropagation(t *testing.T) { childCCCh := make(chan *cancelContext, 1) probe := &turnLoopNestedProbeAgent{parentCCCh: parentCCCh, childCCCh: childCCCh} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return probe, nil }, }) @@ -5472,14 +5474,14 @@ func TestTurnLoop_Push_WithPreemptTimeout_RecursivePropagation(t *testing.T) { childCCCh := make(chan *cancelContext, 2) probe := &turnLoopNestedProbeAgent{parentCCCh: parentCCCh, childCCCh: childCCCh} - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string]{ - GenInput: func(ctx context.Context, _ *TurnLoop[string], items []string) (*GenInputResult[string], error) { - return &GenInputResult[string]{ + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, Consumed: items, }, nil }, - PrepareAgent: func(ctx context.Context, _ *TurnLoop[string], consumed []string) (Agent, error) { + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { return probe, nil }, }) @@ -5490,7 +5492,7 @@ func TestTurnLoop_Push_WithPreemptTimeout_RecursivePropagation(t *testing.T) { t.Cleanup(func() { child.markDone() }) // Preempt with a very short timeout so it escalates to CancelImmediate quickly. - loop.Push("urgent", WithPreemptTimeout[string](AfterChatModel, 10*time.Millisecond)) + loop.Push("urgent", WithPreemptTimeout[string, *schema.Message](AfterChatModel, 10*time.Millisecond)) // After timeout escalation, child should receive the immediate cancel // via recursive propagation. @@ -5515,15 +5517,15 @@ func TestUntilIdleFor_NonPositive_Panics(t *testing.T) { } func TestSaveTurnLoopCheckpoint_NilStore(t *testing.T) { - l := &TurnLoop[string]{config: TurnLoopConfig[string]{Store: nil}} + l := &TurnLoop[string, *schema.Message]{config: TurnLoopConfig[string, *schema.Message]{Store: nil}} err := l.saveTurnLoopCheckpoint(context.Background(), "cp-1", &turnLoopCheckpoint[string]{}) assert.Error(t, err) assert.Contains(t, err.Error(), "checkpoint store is nil") } func TestSetupBridgeStore_NilStore_Resume(t *testing.T) { - l := &TurnLoop[string]{config: TurnLoopConfig[string]{Store: nil}} - spec := &turnRunSpec[string]{isResume: true} + l := &TurnLoop[string, *schema.Message]{config: TurnLoopConfig[string, *schema.Message]{Store: nil}} + spec := &turnRunSpec[string, *schema.Message]{isResume: true} _, _, err := l.setupBridgeStore(spec, nil) assert.Error(t, err) assert.Contains(t, err.Error(), "checkpoint store is nil") diff --git a/adk/utils.go b/adk/utils.go index 89b99132..ec804f72 100644 --- a/adk/utils.go +++ b/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) +} diff --git a/adk/workflow_test.go b/adk/workflow_test.go index 298bef5c..3392187a 100644 --- a/adk/workflow_test.go +++ b/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 diff --git a/adk/wrappers.go b/adk/wrappers.go index 8d9394fc..fad96169 100644 --- a/adk/wrappers.go +++ b/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...) } diff --git a/adk/wrappers_failover_test.go b/adk/wrappers_failover_test.go index 8b14463e..92a68fe9 100644 --- a/adk/wrappers_failover_test.go +++ b/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) +} diff --git a/adk/wrappers_retry_failover_test.go b/adk/wrappers_retry_failover_test.go index 29c4b495..4e8f05d7 100644 --- a/adk/wrappers_retry_failover_test.go +++ b/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")}) diff --git a/components/model/interface.go b/components/model/interface.go index cf79785b..78eadaf2 100644 --- a/components/model/interface.go +++ b/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] diff --git a/schema/agentic_message.go b/schema/agentic_message.go index 43376c14..01a24c4a 100644 --- a/schema/agentic_message.go +++ b/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() } diff --git a/schema/agentic_message_test.go b/schema/agentic_message_test.go index 10639f73..aea4252d 100644 --- a/schema/agentic_message_test.go +++ b/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.") + } + } +} diff --git a/schema/serialization.go b/schema/serialization.go index 22fa16ad..169bf9ee 100644 --- a/schema/serialization.go +++ b/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") diff --git a/schema/tool_test.go b/schema/tool_test.go index e8f95c36..8966cde5 100644 --- a/schema/tool_test.go +++ b/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") + }) +} diff --git a/utils/callbacks/template.go b/utils/callbacks/template.go index f01a849b..850e3011 100644 --- a/utils/callbacks/template.go +++ b/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. diff --git a/utils/callbacks/template_test.go b/utils/callbacks/template_test.go index dcc0e5c7..79be157f 100644 --- a/utils/callbacks/template_test.go +++ b/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)) + }) +}