wehub-resource-sync
1160 lines
38 KiB
Go
1160 lines
38 KiB
Go
/*
|
|
* 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 supervisor
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/cloudwego/eino/adk"
|
|
"github.com/cloudwego/eino/callbacks"
|
|
"github.com/cloudwego/eino/components"
|
|
"github.com/cloudwego/eino/components/tool"
|
|
"github.com/cloudwego/eino/compose"
|
|
mockAdk "github.com/cloudwego/eino/internal/mock/adk"
|
|
mockModel "github.com/cloudwego/eino/internal/mock/components/model"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
// TestNewSupervisor tests the New function
|
|
func TestNewSupervisor(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// Create a mock controller
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
// Create mock agents
|
|
supervisorAgent := mockAdk.NewMockAgent(ctrl)
|
|
subAgent1 := mockAdk.NewMockAgent(ctrl)
|
|
subAgent2 := mockAdk.NewMockAgent(ctrl)
|
|
|
|
supervisorAgent.EXPECT().Name(gomock.Any()).Return("SupervisorAgent").AnyTimes()
|
|
supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor agent description").AnyTimes()
|
|
subAgent1.EXPECT().Name(gomock.Any()).Return("SubAgent1").AnyTimes()
|
|
subAgent2.EXPECT().Name(gomock.Any()).Return("SubAgent2").AnyTimes()
|
|
|
|
aMsg, tMsg := adk.GenTransferMessages(ctx, "SubAgent1")
|
|
i, g := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
g.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, ""))
|
|
event := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)
|
|
event.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent1"}}
|
|
g.Send(event)
|
|
g.Close()
|
|
supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)
|
|
|
|
i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
subAgent1Msg := schema.AssistantMessage("SubAgent1", nil)
|
|
g.Send(adk.EventFromMessage(subAgent1Msg, nil, schema.Assistant, ""))
|
|
g.Close()
|
|
subAgent1.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)
|
|
|
|
aMsg, tMsg = adk.GenTransferMessages(ctx, "SubAgent2 message")
|
|
i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
g.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, ""))
|
|
event = adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)
|
|
event.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent2"}}
|
|
g.Send(event)
|
|
g.Close()
|
|
supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)
|
|
|
|
i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
subAgent2Msg := schema.AssistantMessage("SubAgent2 message", nil)
|
|
g.Send(adk.EventFromMessage(subAgent2Msg, nil, schema.Assistant, ""))
|
|
g.Close()
|
|
subAgent2.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)
|
|
|
|
i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
finishMsg := schema.AssistantMessage("finish", nil)
|
|
g.Send(adk.EventFromMessage(finishMsg, nil, schema.Assistant, ""))
|
|
g.Close()
|
|
supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1)
|
|
|
|
conf := &Config{
|
|
Supervisor: supervisorAgent,
|
|
SubAgents: []adk.Agent{subAgent1, subAgent2},
|
|
}
|
|
|
|
multiAgent, err := New(ctx, conf)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, multiAgent)
|
|
assert.Equal(t, "SupervisorAgent", multiAgent.Name(ctx))
|
|
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})
|
|
aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test")})
|
|
|
|
// transfer to agent1
|
|
event, ok := aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SupervisorAgent", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))
|
|
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SupervisorAgent", event.AgentName)
|
|
assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, "SubAgent1", event.Action.TransferToAgent.DestAgentName)
|
|
|
|
// agent1's output
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent1", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, subAgent1Msg.Content, event.Output.MessageOutput.Message.Content)
|
|
|
|
// transfer back to supervisor
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent1", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))
|
|
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent1", event.AgentName)
|
|
assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, "SupervisorAgent", event.Action.TransferToAgent.DestAgentName)
|
|
|
|
// transfer to agent2
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SupervisorAgent", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))
|
|
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SupervisorAgent", event.AgentName)
|
|
assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, "SubAgent2", event.Action.TransferToAgent.DestAgentName)
|
|
|
|
// agent1's output
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent2", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, subAgent2Msg.Content, event.Output.MessageOutput.Message.Content)
|
|
|
|
// transfer back to supervisor
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent2", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls))
|
|
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SubAgent2", event.AgentName)
|
|
assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, "SupervisorAgent", event.Action.TransferToAgent.DestAgentName)
|
|
|
|
// finish
|
|
event, ok = aIter.Next()
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "SupervisorAgent", event.AgentName)
|
|
assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
|
|
assert.Equal(t, finishMsg.Content, event.Output.MessageOutput.Message.Content)
|
|
}
|
|
|
|
type approvalInfo struct {
|
|
ToolName string
|
|
ArgumentsInJSON string
|
|
ToolCallID string
|
|
}
|
|
|
|
func (ai *approvalInfo) String() string {
|
|
return fmt.Sprintf("tool '%s' interrupted with arguments '%s', waiting for approval",
|
|
ai.ToolName, ai.ArgumentsInJSON)
|
|
}
|
|
|
|
type approvalResult struct {
|
|
Approved bool
|
|
DisapproveReason *string
|
|
}
|
|
|
|
func init() {
|
|
schema.Register[*approvalInfo]()
|
|
schema.Register[*approvalResult]()
|
|
}
|
|
|
|
type approvableTool struct {
|
|
name string
|
|
t *testing.T
|
|
}
|
|
|
|
func (m *approvableTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: m.name,
|
|
Desc: "A tool that requires approval before execution",
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
|
"action": {Type: schema.String, Desc: "The action to perform"},
|
|
}),
|
|
}, nil
|
|
}
|
|
|
|
func (m *approvableTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
|
|
wasInterrupted, _, storedArguments := tool.GetInterruptState[string](ctx)
|
|
if !wasInterrupted {
|
|
return "", tool.StatefulInterrupt(ctx, &approvalInfo{
|
|
ToolName: m.name,
|
|
ArgumentsInJSON: argumentsInJSON,
|
|
ToolCallID: compose.GetToolCallID(ctx),
|
|
}, argumentsInJSON)
|
|
}
|
|
|
|
isResumeTarget, hasData, data := tool.GetResumeContext[*approvalResult](ctx)
|
|
if !isResumeTarget {
|
|
return "", tool.StatefulInterrupt(ctx, &approvalInfo{
|
|
ToolName: m.name,
|
|
ArgumentsInJSON: storedArguments,
|
|
ToolCallID: compose.GetToolCallID(ctx),
|
|
}, storedArguments)
|
|
}
|
|
|
|
if !hasData {
|
|
return "", fmt.Errorf("tool '%s' resumed with no data", m.name)
|
|
}
|
|
|
|
if data.Approved {
|
|
return fmt.Sprintf("Tool '%s' executed successfully with args: %s", m.name, storedArguments), nil
|
|
}
|
|
|
|
if data.DisapproveReason != nil {
|
|
return fmt.Sprintf("Tool '%s' disapproved, reason: %s", m.name, *data.DisapproveReason), nil
|
|
}
|
|
|
|
return fmt.Sprintf("Tool '%s' disapproved", m.name), nil
|
|
}
|
|
|
|
type checkpointStore struct {
|
|
data map[string][]byte
|
|
}
|
|
|
|
func newCheckpointStore() *checkpointStore {
|
|
return &checkpointStore{data: make(map[string][]byte)}
|
|
}
|
|
|
|
func (s *checkpointStore) Set(_ context.Context, key string, value []byte) error {
|
|
s.data[key] = value
|
|
return nil
|
|
}
|
|
|
|
func (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) {
|
|
v, ok := s.data[key]
|
|
return v, ok, nil
|
|
}
|
|
|
|
type namedAgent struct {
|
|
adk.ResumableAgent
|
|
name string
|
|
description string
|
|
}
|
|
|
|
func (n *namedAgent) Name(_ context.Context) string {
|
|
return n.name
|
|
}
|
|
|
|
func (n *namedAgent) Description(_ context.Context) string {
|
|
return n.description
|
|
}
|
|
|
|
func TestNestedSupervisorInterruptResume(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
mockOuterSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
mockInnerSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
mockWorkerModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
paymentTool := &approvableTool{name: "process_payment", t: t}
|
|
|
|
userInput := []adk.Message{schema.UserMessage("Process a payment of $1000")}
|
|
|
|
mockWorkerModel.EXPECT().WithTools(gomock.Any()).Return(mockWorkerModel, nil).AnyTimes()
|
|
|
|
workerToolCallMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "call_payment_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "process_payment",
|
|
Arguments: `{"action": "process $1000 payment"}`,
|
|
},
|
|
},
|
|
})
|
|
mockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(workerToolCallMsg, nil).Times(1)
|
|
|
|
workerCompletionMsg := schema.AssistantMessage("Payment processed successfully", nil)
|
|
mockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(workerCompletionMsg, nil).AnyTimes()
|
|
|
|
workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "payment_worker",
|
|
Description: "the agent responsible for processing payments",
|
|
Instruction: "You are a payment processing worker. Use the process_payment tool to handle payments.",
|
|
Model: mockWorkerModel,
|
|
ToolsConfig: adk.ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{paymentTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
mockInnerSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockInnerSupervisorModel, nil).AnyTimes()
|
|
|
|
innerTransferMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "inner_transfer_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "transfer_to_agent",
|
|
Arguments: `{"agent_name":"payment_worker"}`,
|
|
},
|
|
},
|
|
})
|
|
mockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(innerTransferMsg, nil).Times(1)
|
|
|
|
innerFinalMsg := schema.AssistantMessage("Payment has been processed and approved.", nil)
|
|
mockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(innerFinalMsg, nil).AnyTimes()
|
|
|
|
innerSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "payment_supervisor",
|
|
Description: "the supervisor agent responsible for payment operations",
|
|
Instruction: "You are a payment supervisor. Delegate payment tasks to payment_worker.",
|
|
Model: mockInnerSupervisorModel,
|
|
Exit: &adk.ExitTool{},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
innerSupervisorAgent, err := New(ctx, &Config{
|
|
Supervisor: innerSupervisorChatAgent,
|
|
SubAgents: []adk.Agent{workerAgent},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
innerSupervisorWrapped := &namedAgent{
|
|
ResumableAgent: innerSupervisorAgent,
|
|
name: "payment_department",
|
|
description: "the department responsible for all payment-related operations",
|
|
}
|
|
|
|
mockOuterSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockOuterSupervisorModel, nil).AnyTimes()
|
|
|
|
outerTransferMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "outer_transfer_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "transfer_to_agent",
|
|
Arguments: `{"agent_name":"payment_department"}`,
|
|
},
|
|
},
|
|
})
|
|
mockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(outerTransferMsg, nil).Times(1)
|
|
|
|
outerFinalMsg := schema.AssistantMessage("The payment request has been fully processed by the payment department.", nil)
|
|
mockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(outerFinalMsg, nil).AnyTimes()
|
|
|
|
outerSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "company_coordinator",
|
|
Description: "the top-level coordinator for company operations",
|
|
Instruction: "You are the company coordinator. Route payment requests to payment_department.",
|
|
Model: mockOuterSupervisorModel,
|
|
Exit: &adk.ExitTool{},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
outerSupervisorAgent, err := New(ctx, &Config{
|
|
Supervisor: outerSupervisorChatAgent,
|
|
SubAgents: []adk.Agent{innerSupervisorWrapped},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
outerSupervisorWrapped := &namedAgent{
|
|
ResumableAgent: outerSupervisorAgent,
|
|
name: "headquarters",
|
|
description: "the company headquarters that coordinates all departments",
|
|
}
|
|
|
|
store := newCheckpointStore()
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{
|
|
Agent: outerSupervisorWrapped,
|
|
CheckPointStore: store,
|
|
})
|
|
|
|
t.Log("========================================")
|
|
t.Log("Starting Nested Supervisor Integration Test (with namedAgent wrappers)")
|
|
t.Log("Hierarchy: headquarters(wrapper) -> company_coordinator -> payment_department(wrapper) -> payment_supervisor -> payment_worker -> process_payment tool")
|
|
t.Log("========================================")
|
|
|
|
checkpointID := "test-nested-supervisor-1"
|
|
iter := runner.Run(ctx, userInput, adk.WithCheckPointID(checkpointID))
|
|
|
|
var interruptEvent *adk.AgentEvent
|
|
eventCount := 0
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
eventCount++
|
|
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
interruptEvent = event
|
|
t.Log("INTERRUPT DETECTED - Deep interrupt from tool within nested supervisor")
|
|
break
|
|
}
|
|
}
|
|
|
|
if interruptEvent == nil {
|
|
t.Fatal("Expected an interrupt event from the process_payment tool, but none was received")
|
|
}
|
|
|
|
assert.NotNil(t, interruptEvent.Action.Interrupted, "Should have interrupt info")
|
|
assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "Should have interrupt contexts")
|
|
|
|
var toolInterruptID string
|
|
for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
if intCtx.IsRootCause {
|
|
toolInterruptID = intCtx.ID
|
|
break
|
|
}
|
|
}
|
|
assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID")
|
|
|
|
t.Logf("Resuming with approval for interrupt ID: %s", toolInterruptID)
|
|
|
|
resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{
|
|
Targets: map[string]any{
|
|
toolInterruptID: &approvalResult{Approved: true},
|
|
},
|
|
})
|
|
assert.NoError(t, err, "Resume should not error")
|
|
assert.NotNil(t, resumeIter, "Resume iterator should not be nil")
|
|
|
|
var resumeEvents []*adk.AgentEvent
|
|
for {
|
|
event, ok := resumeIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
resumeEvents = append(resumeEvents, event)
|
|
}
|
|
|
|
assert.NotEmpty(t, resumeEvents, "Should have resume events after approval")
|
|
|
|
for _, event := range resumeEvents {
|
|
assert.NoError(t, event.Err, "Resume event should not have error")
|
|
}
|
|
|
|
var hasToolResponse, hasTransferBack bool
|
|
for _, event := range resumeEvents {
|
|
if event.Output != nil && event.Output.MessageOutput != nil {
|
|
msg := event.Output.MessageOutput.Message
|
|
if msg != nil && msg.Role == "tool" && strings.Contains(msg.Content, "executed successfully") {
|
|
hasToolResponse = true
|
|
}
|
|
}
|
|
if event.Action != nil && event.Action.TransferToAgent != nil {
|
|
if event.Action.TransferToAgent.DestAgentName == "company_coordinator" {
|
|
hasTransferBack = true
|
|
}
|
|
}
|
|
}
|
|
|
|
assert.True(t, hasToolResponse, "Should have tool response indicating successful payment processing")
|
|
assert.True(t, hasTransferBack, "Should have transfer back to outer supervisor indicating completion")
|
|
}
|
|
|
|
func TestSupervisorExit(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
supervisorAgent := mockAdk.NewMockAgent(ctrl)
|
|
subAgent := mockAdk.NewMockAgent(ctrl)
|
|
|
|
supervisorAgent.EXPECT().Name(gomock.Any()).Return("Supervisor").AnyTimes()
|
|
supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor description").AnyTimes()
|
|
subAgent.EXPECT().Name(gomock.Any()).Return("SubAgent").AnyTimes()
|
|
|
|
// 1. Supervisor transfers to SubAgent
|
|
aMsg, tMsg := adk.GenTransferMessages(ctx, "SubAgent")
|
|
i1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
g1.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, ""))
|
|
event1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)
|
|
event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent"}}
|
|
g1.Send(event1)
|
|
g1.Close()
|
|
supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).Times(1)
|
|
|
|
// 2. SubAgent emits Exit action
|
|
i2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
exitEvent := &adk.AgentEvent{
|
|
AgentName: "SubAgent",
|
|
Action: &adk.AgentAction{Exit: true},
|
|
Output: &adk.AgentOutput{
|
|
MessageOutput: &adk.MessageVariant{
|
|
Role: schema.Assistant,
|
|
Message: schema.AssistantMessage("Exiting...", nil),
|
|
},
|
|
},
|
|
}
|
|
g2.Send(exitEvent)
|
|
g2.Close()
|
|
subAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).Times(1)
|
|
|
|
conf := &Config{
|
|
Supervisor: supervisorAgent,
|
|
SubAgents: []adk.Agent{subAgent},
|
|
}
|
|
|
|
multiAgent, err := New(ctx, conf)
|
|
assert.NoError(t, err)
|
|
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})
|
|
aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test")})
|
|
|
|
// Collect events
|
|
var events []*adk.AgentEvent
|
|
for {
|
|
event, ok := aIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
foundExit := false
|
|
foundTransferBack := false
|
|
|
|
for _, e := range events {
|
|
if e.Action != nil {
|
|
if e.Action.Exit {
|
|
foundExit = true
|
|
}
|
|
if e.Action.TransferToAgent != nil && e.Action.TransferToAgent.DestAgentName == "Supervisor" {
|
|
foundTransferBack = true
|
|
}
|
|
}
|
|
}
|
|
|
|
assert.True(t, foundExit, "Should have found Exit action")
|
|
assert.False(t, foundTransferBack, "Should NOT have found Transfer back to Supervisor after Exit")
|
|
}
|
|
|
|
func TestNestedSupervisorExit(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
topSupervisor := mockAdk.NewMockAgent(ctrl)
|
|
midSupervisor := mockAdk.NewMockAgent(ctrl)
|
|
worker := mockAdk.NewMockAgent(ctrl)
|
|
|
|
topSupervisor.EXPECT().Name(gomock.Any()).Return("TopSupervisor").AnyTimes()
|
|
topSupervisor.EXPECT().Description(gomock.Any()).Return("Top supervisor description").AnyTimes()
|
|
midSupervisor.EXPECT().Name(gomock.Any()).Return("MidSupervisor").AnyTimes()
|
|
midSupervisor.EXPECT().Description(gomock.Any()).Return("Mid supervisor description").AnyTimes()
|
|
worker.EXPECT().Name(gomock.Any()).Return("Worker").AnyTimes()
|
|
|
|
// 1. TopSupervisor transfers to MidSupervisor
|
|
aMsg1, tMsg1 := adk.GenTransferMessages(ctx, "MidSupervisor")
|
|
i1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
g1.Send(adk.EventFromMessage(aMsg1, nil, schema.Assistant, ""))
|
|
event1 := adk.EventFromMessage(tMsg1, nil, schema.Tool, tMsg1.ToolName)
|
|
event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "MidSupervisor"}}
|
|
g1.Send(event1)
|
|
g1.Close()
|
|
topSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).AnyTimes()
|
|
|
|
// 2. MidSupervisor transfers to Worker
|
|
aMsg2, tMsg2 := adk.GenTransferMessages(ctx, "Worker")
|
|
i2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
g2.Send(adk.EventFromMessage(aMsg2, nil, schema.Assistant, ""))
|
|
event2 := adk.EventFromMessage(tMsg2, nil, schema.Tool, tMsg2.ToolName)
|
|
event2.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "Worker"}}
|
|
g2.Send(event2)
|
|
g2.Close()
|
|
midSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).AnyTimes()
|
|
|
|
// 3. Worker emits Exit action
|
|
i3, g3 := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
exitEvent := &adk.AgentEvent{
|
|
AgentName: "Worker",
|
|
Action: &adk.AgentAction{Exit: true},
|
|
Output: &adk.AgentOutput{
|
|
MessageOutput: &adk.MessageVariant{
|
|
Role: schema.Assistant,
|
|
Message: schema.AssistantMessage("Worker Exiting...", nil),
|
|
},
|
|
},
|
|
}
|
|
g3.Send(exitEvent)
|
|
g3.Close()
|
|
worker.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i3).Times(1)
|
|
|
|
// Build Nested System
|
|
// Mid System: MidSupervisor -> [Worker]
|
|
midSystem, err := New(ctx, &Config{
|
|
Supervisor: midSupervisor,
|
|
SubAgents: []adk.Agent{worker},
|
|
})
|
|
assert.NoError(t, err)
|
|
// We need to give the midSystem the name "MidSupervisor" so TopSupervisor can find it
|
|
// supervisor.New returns a ResumableAgent that delegates Name() to the supervisor agent.
|
|
// So midSystem.Name() should already be "MidSupervisor" because midSupervisor.Name() is "MidSupervisor".
|
|
|
|
// Top System: TopSupervisor -> [midSystem]
|
|
topSystem, err := New(ctx, &Config{
|
|
Supervisor: topSupervisor,
|
|
SubAgents: []adk.Agent{midSystem},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: topSystem})
|
|
aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test nested exit")})
|
|
|
|
// Collect events
|
|
var events []*adk.AgentEvent
|
|
for {
|
|
event, ok := aIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
foundExit := false
|
|
foundTransferBackToMidAfterExit := false
|
|
foundTransferBackToTopAfterExit := false
|
|
|
|
for _, e := range events {
|
|
if e.Action != nil {
|
|
if e.Action.Exit {
|
|
foundExit = true
|
|
}
|
|
if foundExit && e.Action.TransferToAgent != nil {
|
|
if e.Action.TransferToAgent.DestAgentName == "MidSupervisor" {
|
|
foundTransferBackToMidAfterExit = true
|
|
}
|
|
if e.Action.TransferToAgent.DestAgentName == "TopSupervisor" {
|
|
foundTransferBackToTopAfterExit = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
assert.True(t, foundExit, "Should have found Exit action")
|
|
assert.False(t, foundTransferBackToMidAfterExit, "Should NOT have found Transfer back to MidSupervisor after Exit")
|
|
assert.False(t, foundTransferBackToTopAfterExit, "Should NOT have found Transfer back to TopSupervisor after Exit")
|
|
}
|
|
|
|
func TestChatModelAgentInternalEventsExit(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
supervisorAgent := mockAdk.NewMockAgent(ctrl)
|
|
workerModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
innerAgent := mockAdk.NewMockAgent(ctrl)
|
|
|
|
supervisorAgent.EXPECT().Name(gomock.Any()).Return("Supervisor").AnyTimes()
|
|
supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor description").AnyTimes()
|
|
innerAgent.EXPECT().Name(gomock.Any()).Return("InnerAgent").AnyTimes()
|
|
innerAgent.EXPECT().Description(gomock.Any()).Return("Inner Agent Description").AnyTimes()
|
|
|
|
// 1. Supervisor transfers to Worker (only once, then exits when worker transfers back)
|
|
supervisorRunCount := 0
|
|
supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
|
|
func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
|
|
supervisorRunCount++
|
|
iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
go func() {
|
|
defer gen.Close()
|
|
if supervisorRunCount == 1 {
|
|
aMsg, tMsg := adk.GenTransferMessages(ctx, "Worker")
|
|
gen.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, ""))
|
|
event1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName)
|
|
event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "Worker"}}
|
|
gen.Send(event1)
|
|
} else {
|
|
exitEvent := &adk.AgentEvent{
|
|
AgentName: "Supervisor",
|
|
Action: &adk.AgentAction{Exit: true},
|
|
Output: &adk.AgentOutput{
|
|
MessageOutput: &adk.MessageVariant{
|
|
Role: schema.Assistant,
|
|
Message: schema.AssistantMessage("Supervisor done", nil),
|
|
},
|
|
},
|
|
}
|
|
gen.Send(exitEvent)
|
|
}
|
|
}()
|
|
return iter
|
|
}).AnyTimes()
|
|
|
|
// 2. Worker runs, calls AgentTool (InnerAgent)
|
|
// Mock WorkerModel behavior
|
|
workerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes()
|
|
|
|
// 2.1 Worker generates tool call
|
|
toolCallMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "call_inner_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "InnerAgent",
|
|
Arguments: `{"request": "do exit"}`,
|
|
},
|
|
},
|
|
})
|
|
workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(toolCallMsg, nil).Times(1)
|
|
|
|
// 2.2 InnerAgent runs and emits Exit
|
|
innerAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
|
|
func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
|
|
iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
|
|
go func() {
|
|
defer gen.Close()
|
|
innerExitEvent := &adk.AgentEvent{
|
|
AgentName: "InnerAgent",
|
|
Action: &adk.AgentAction{Exit: true},
|
|
RunPath: []adk.RunStep{},
|
|
Output: &adk.AgentOutput{
|
|
MessageOutput: &adk.MessageVariant{
|
|
Role: schema.Assistant,
|
|
Message: schema.AssistantMessage("Inner Exiting...", nil),
|
|
},
|
|
},
|
|
}
|
|
gen.Send(innerExitEvent)
|
|
}()
|
|
return iter
|
|
}).AnyTimes()
|
|
|
|
// 2.3 Worker receives tool result (empty string or whatever AgentTool returns on exit/interrupt)
|
|
// AgentTool implementation details: if Exit action is present, it returns whatever output is there.
|
|
// The Exit action itself is passed as internal event.
|
|
|
|
// 2.4 Worker generates final response
|
|
finalMsg := schema.AssistantMessage("Worker Finished", nil)
|
|
workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(finalMsg, nil).AnyTimes()
|
|
|
|
// Build Worker Agent
|
|
agentTool := adk.NewAgentTool(ctx, innerAgent)
|
|
workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "Worker",
|
|
Description: "Worker Agent",
|
|
Model: workerModel,
|
|
ToolsConfig: adk.ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{agentTool},
|
|
},
|
|
EmitInternalEvents: true, // Key configuration
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
// Build System
|
|
sys, err := New(ctx, &Config{
|
|
Supervisor: supervisorAgent,
|
|
SubAgents: []adk.Agent{workerAgent},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: sys})
|
|
aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("start")})
|
|
|
|
// Collect events
|
|
var events []*adk.AgentEvent
|
|
for {
|
|
event, ok := aIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
events = append(events, event)
|
|
}
|
|
|
|
foundInnerExit := false
|
|
foundTransferBack := false
|
|
|
|
for _, e := range events {
|
|
// Check for InnerAgent exit event (propagated as internal event)
|
|
if e.AgentName == "InnerAgent" && e.Action != nil && e.Action.Exit {
|
|
foundInnerExit = true
|
|
}
|
|
|
|
// Check for transfer back to Supervisor
|
|
if e.AgentName == "Worker" && e.Action != nil && e.Action.TransferToAgent != nil &&
|
|
e.Action.TransferToAgent.DestAgentName == "Supervisor" {
|
|
foundTransferBack = true
|
|
}
|
|
}
|
|
|
|
assert.True(t, foundInnerExit, "Should have captured InnerAgent Exit event")
|
|
assert.True(t, foundTransferBack, "Should have found Transfer back to Supervisor (Worker should NOT be considered exited)")
|
|
}
|
|
|
|
func TestSupervisorContainerUnifiedTracing(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
supervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
subAgentModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
supervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes()
|
|
subAgentModel.EXPECT().WithTools(gomock.Any()).Return(subAgentModel, nil).AnyTimes()
|
|
|
|
transferMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "transfer_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "transfer_to_agent",
|
|
Arguments: `{"agent_name":"SubAgent"}`,
|
|
},
|
|
},
|
|
})
|
|
supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(transferMsg, nil).Times(1)
|
|
|
|
subAgentResponse := schema.AssistantMessage("SubAgent response", nil)
|
|
subAgentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(subAgentResponse, nil).Times(1)
|
|
|
|
finalResponse := schema.AssistantMessage("Final response", nil)
|
|
supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(finalResponse, nil).Times(1)
|
|
|
|
supervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "SupervisorAgent",
|
|
Description: "Supervisor agent",
|
|
Instruction: "You are a supervisor",
|
|
Model: supervisorModel,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
subAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "SubAgent",
|
|
Description: "Sub agent",
|
|
Instruction: "You are a sub agent",
|
|
Model: subAgentModel,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
multiAgent, err := New(ctx, &Config{
|
|
Supervisor: supervisorAgent,
|
|
SubAgents: []adk.Agent{subAgent},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "SupervisorAgent", multiAgent.Name(ctx))
|
|
|
|
typer, ok := multiAgent.(components.Typer)
|
|
assert.True(t, ok, "Should implement components.Typer")
|
|
assert.Equal(t, "Supervisor", typer.GetType())
|
|
|
|
var mu sync.Mutex
|
|
var onStartCalls []string
|
|
var onEndCalls []string
|
|
|
|
handler := callbacks.NewHandlerBuilder().
|
|
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
mu.Lock()
|
|
onStartCalls = append(onStartCalls, info.Name+":"+info.Type)
|
|
mu.Unlock()
|
|
return ctx
|
|
}).
|
|
OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
mu.Lock()
|
|
onEndCalls = append(onEndCalls, info.Name+":"+info.Type)
|
|
mu.Unlock()
|
|
if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {
|
|
go func() {
|
|
for {
|
|
_, ok := agentOutput.Events.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
return ctx
|
|
}).
|
|
Build()
|
|
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent})
|
|
iter := runner.Query(ctx, "hello", adk.WithCallbacks(handler))
|
|
|
|
for {
|
|
_, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
assert.NotEmpty(t, onStartCalls, "Should have OnStart calls")
|
|
assert.Contains(t, onStartCalls, "SupervisorAgent:Supervisor", "Should have supervisor container OnStart with type 'Supervisor'")
|
|
}
|
|
|
|
type traceContextKey struct{}
|
|
|
|
func TestSupervisorContainerUnifiedTracingOnResume(t *testing.T) {
|
|
ctx := context.Background()
|
|
ctrl := gomock.NewController(t)
|
|
defer ctrl.Finish()
|
|
|
|
supervisorModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
workerModel := mockModel.NewMockToolCallingChatModel(ctrl)
|
|
|
|
supervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes()
|
|
workerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes()
|
|
|
|
paymentTool := &approvableTool{name: "process_payment", t: t}
|
|
|
|
workerToolCallMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "call_payment_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "process_payment",
|
|
Arguments: `{"action": "process $1000 payment"}`,
|
|
},
|
|
},
|
|
})
|
|
workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(workerToolCallMsg, nil).Times(1)
|
|
|
|
workerCompletionMsg := schema.AssistantMessage("Payment processed successfully", nil)
|
|
workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(workerCompletionMsg, nil).AnyTimes()
|
|
|
|
workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "Worker",
|
|
Description: "Worker agent",
|
|
Instruction: "You are a worker",
|
|
Model: workerModel,
|
|
ToolsConfig: adk.ToolsConfig{
|
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
|
Tools: []tool.BaseTool{paymentTool},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
transferMsg := schema.AssistantMessage("", []schema.ToolCall{
|
|
{
|
|
ID: "transfer_1",
|
|
Type: "function",
|
|
Function: schema.FunctionCall{
|
|
Name: "transfer_to_agent",
|
|
Arguments: `{"agent_name":"Worker"}`,
|
|
},
|
|
},
|
|
})
|
|
supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(transferMsg, nil).Times(1)
|
|
|
|
finalResponse := schema.AssistantMessage("Final response", nil)
|
|
supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).
|
|
Return(finalResponse, nil).AnyTimes()
|
|
|
|
supervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
|
Name: "SupervisorAgent",
|
|
Description: "Supervisor agent",
|
|
Instruction: "You are a supervisor",
|
|
Model: supervisorModel,
|
|
Exit: &adk.ExitTool{},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
multiAgent, err := New(ctx, &Config{
|
|
Supervisor: supervisorAgent,
|
|
SubAgents: []adk.Agent{workerAgent},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
store := newCheckpointStore()
|
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{
|
|
Agent: multiAgent,
|
|
CheckPointStore: store,
|
|
})
|
|
|
|
var mu sync.Mutex
|
|
var runOnStartCalls []string
|
|
var resumeOnStartCalls []string
|
|
var resumeParentTraceIDs []string
|
|
|
|
runHandler := callbacks.NewHandlerBuilder().
|
|
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
mu.Lock()
|
|
runOnStartCalls = append(runOnStartCalls, info.Name+":"+info.Type)
|
|
mu.Unlock()
|
|
return ctx
|
|
}).
|
|
OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {
|
|
go func() {
|
|
for {
|
|
_, ok := agentOutput.Events.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
return ctx
|
|
}).
|
|
Build()
|
|
|
|
checkpointID := "test-unified-tracing-resume"
|
|
iter := runner.Run(ctx, []adk.Message{schema.UserMessage("Process payment")}, adk.WithCallbacks(runHandler), adk.WithCheckPointID(checkpointID))
|
|
|
|
var interruptEvent *adk.AgentEvent
|
|
for {
|
|
event, ok := iter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
if event.Action != nil && event.Action.Interrupted != nil {
|
|
interruptEvent = event
|
|
break
|
|
}
|
|
}
|
|
|
|
assert.NotNil(t, interruptEvent, "Should have interrupt event")
|
|
|
|
var toolInterruptID string
|
|
for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts {
|
|
if intCtx.IsRootCause {
|
|
toolInterruptID = intCtx.ID
|
|
break
|
|
}
|
|
}
|
|
assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID")
|
|
|
|
mu.Lock()
|
|
t.Logf("Run OnStart calls: %v", runOnStartCalls)
|
|
assert.Contains(t, runOnStartCalls, "SupervisorAgent:Supervisor", "Run should have supervisor container OnStart")
|
|
mu.Unlock()
|
|
|
|
resumeHandler := callbacks.NewHandlerBuilder().
|
|
OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
mu.Lock()
|
|
resumeOnStartCalls = append(resumeOnStartCalls, info.Name+":"+info.Type)
|
|
parentID, _ := ctx.Value(traceContextKey{}).(string)
|
|
resumeParentTraceIDs = append(resumeParentTraceIDs, info.Name+":parent="+parentID)
|
|
mu.Unlock()
|
|
if info.Type == "Supervisor" {
|
|
return context.WithValue(ctx, traceContextKey{}, "supervisor-trace-id")
|
|
}
|
|
return ctx
|
|
}).
|
|
OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
|
|
if info.Component != adk.ComponentOfAgent {
|
|
return ctx
|
|
}
|
|
if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil {
|
|
go func() {
|
|
for {
|
|
_, ok := agentOutput.Events.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
return ctx
|
|
}).
|
|
Build()
|
|
|
|
resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{
|
|
Targets: map[string]any{
|
|
toolInterruptID: &approvalResult{Approved: true},
|
|
},
|
|
}, adk.WithCallbacks(resumeHandler))
|
|
assert.NoError(t, err)
|
|
|
|
for {
|
|
event, ok := resumeIter.Next()
|
|
if !ok {
|
|
break
|
|
}
|
|
assert.NoError(t, event.Err)
|
|
}
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
t.Logf("Resume OnStart calls: %v", resumeOnStartCalls)
|
|
t.Logf("Resume parent trace IDs: %v", resumeParentTraceIDs)
|
|
assert.NotEmpty(t, resumeOnStartCalls, "Should have OnStart calls during resume")
|
|
assert.Contains(t, resumeOnStartCalls, "SupervisorAgent:Supervisor", "Resume should have supervisor container OnStart with type 'Supervisor'")
|
|
|
|
foundInnerSupervisorWithParent := false
|
|
for _, entry := range resumeParentTraceIDs {
|
|
if strings.Contains(entry, "SupervisorAgent") && !strings.Contains(entry, "parent=supervisor-trace-id") && entry != "SupervisorAgent:parent=" {
|
|
if strings.Contains(resumeOnStartCalls[0], "Supervisor") {
|
|
continue
|
|
}
|
|
}
|
|
if strings.Contains(entry, "parent=supervisor-trace-id") {
|
|
foundInnerSupervisorWithParent = true
|
|
}
|
|
}
|
|
assert.True(t, foundInnerSupervisorWithParent,
|
|
"Inner agents should have parent trace from Supervisor container during Resume. Got: %v", resumeParentTraceIDs)
|
|
}
|