feat(adk): ut for supervisor supports concurrent transfer

Change-Id: I64654953e73696105bb0437d27adf707c35fe941

Author: shentongmartin

adk/deterministic_transfer.go

@@ -53,6 +53,10 @@ func (a *agentWithDeterministicTransferTo) Name(ctx context.Context) string {
 func (a *agentWithDeterministicTransferTo) Run(ctx context.Context,
 	input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 
+	if _, ok := a.agent.(*flowAgent); ok {
+		ctx = ClearRunCtx(ctx)
+	}
+
 	aIter := a.agent.Run(ctx, input, options...)
 
 	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
@@ -77,6 +81,10 @@ func (a *resumableAgentWithDeterministicTransferTo) Name(ctx context.Context) st
 func (a *resumableAgentWithDeterministicTransferTo) Run(ctx context.Context,
 	input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 
+	if _, ok := a.agent.(*flowAgent); ok {
+		ctx = ClearRunCtx(ctx)
+	}
+
 	aIter := a.agent.Run(ctx, input, options...)
 
 	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()

adk/flow.go

@@ -38,7 +38,7 @@ type HistoryEntry struct {
 
 type HistoryRewriter func(ctx context.Context, entries []*HistoryEntry) ([]Message, error)
 
-type dynamicParallelState struct {
+type flowInterruptState struct {
 	// Maps the destination agent name to the events generated in its lane before interruption.
 	// This also serves as the source of truth for which lanes need to be resumed.
 	LaneEvents map[string][]*agentEventWrapper
@@ -71,7 +71,7 @@ func (a *flowAgent) collectLaneEvents(childContexts []context.Context, agentName
 
 // createCompositeInterrupt creates a composite interrupt event with the collected state
 func (a *flowAgent) createCompositeInterrupt(ctx context.Context, laneEvents map[string][]*agentEventWrapper, subInterruptSignals []*core.InterruptSignal) *AgentEvent {
-	state := &dynamicParallelState{
+	state := &flowInterruptState{
 		LaneEvents: laneEvents,
 	}
 
@@ -85,7 +85,7 @@ func (a *flowAgent) createCompositeInterrupt(ctx context.Context, laneEvents map
 }
 
 func init() {
-	schema.RegisterName[*dynamicParallelState]("eino_adk_dynamic_parallel_state")
+	schema.RegisterName[*flowInterruptState]("eino_adk_dynamic_parallel_state")
 }
 
 type flowAgent struct {
@@ -372,7 +372,7 @@ func (a *flowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentR
 	if info.WasInterrupted {
 		// Check if we need to resume concurrent transfers
 		if info.InterruptState != nil {
-			state, ok := info.InterruptState.(*dynamicParallelState)
+			state, ok := info.InterruptState.(*flowInterruptState)
 			if ok {
 				// Delegate to resumeConcurrentLanes which will handle the type assertion
 				return a.resumeConcurrentLanes(ctx, state, info, opts...)
@@ -555,7 +555,7 @@ func (a *flowAgent) runConcurrentLanes(
 // resumeConcurrentLanes resumes execution after a concurrent transfer interruption
 func (a *flowAgent) resumeConcurrentLanes(
 	ctx context.Context,
-	state *dynamicParallelState,
+	state *flowInterruptState,
 	info *ResumeInfo,
 	opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 

adk/flow_test.go

@@ -1,3 +1,19 @@
+/*
+ * 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 (

adk/prebuilt/supervisor/supervisor_test.go

@@ -18,12 +18,14 @@ package supervisor
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 
 	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/compose"
 	mockAdk "github.com/cloudwego/eino/internal/mock/adk"
 	"github.com/cloudwego/eino/schema"
 )
@@ -167,3 +169,280 @@ func TestNewSupervisor(t *testing.T) {
 	assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role)
 	assert.Equal(t, finishMsg.Content, event.Output.MessageOutput.Message.Content)
 }
+
+// mockSupervisor is a simple supervisor that performs concurrent transfers
+type mockSupervisor struct {
+	name    string
+	targets []string
+	times   int
+}
+
+func (a *mockSupervisor) Name(_ context.Context) string {
+	return a.name
+}
+
+func (a *mockSupervisor) Description(_ context.Context) string {
+	return "mock supervisor agent"
+}
+
+func (a *mockSupervisor) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
+	iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
+	if a.times > 0 {
+		gen.Send(adk.EventFromMessage(schema.AssistantMessage("job done", nil), nil, schema.Assistant, ""))
+		gen.Close()
+		return iter
+	}
+
+	a.times++
+
+	// Create assistant message with tool call for concurrent transfer
+	toolCall := schema.ToolCall{
+		ID:   "transfer-tool-call",
+		Type: "function",
+		Function: schema.FunctionCall{
+			Name:      adk.TransferToAgentToolName,
+			Arguments: `{"agent_names":["` + a.targets[0] + `","` + a.targets[1] + `"]}`,
+		},
+	}
+	assistantMsg := schema.AssistantMessage("", []schema.ToolCall{toolCall})
+	gen.Send(adk.EventFromMessage(assistantMsg, nil, schema.Assistant, ""))
+
+	// Create tool message for the transfer
+	toolMsg := schema.ToolMessage(fmt.Sprintf("Successfully transfered to agents %v", a.targets), toolCall.ID,
+		schema.WithToolName(adk.TransferToAgentToolName))
+	transferEvent := adk.EventFromMessage(toolMsg, nil, schema.Tool, toolMsg.ToolName)
+	transferEvent.Action = &adk.AgentAction{
+		ConcurrentTransferToAgent: &adk.ConcurrentTransferToAgentAction{
+			DestAgentNames: a.targets,
+		},
+	}
+	gen.Send(transferEvent)
+	gen.Close()
+
+	return iter
+}
+
+// mockSimpleAgent is a basic agent that returns a simple message
+type mockSimpleAgent struct {
+	name string
+	msg  string
+}
+
+func (a *mockSimpleAgent) Name(_ context.Context) string {
+	return a.name
+}
+
+func (a *mockSimpleAgent) Description(_ context.Context) string {
+	return "mock simple agent"
+}
+
+func (a *mockSimpleAgent) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
+	iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
+	gen.Send(adk.EventFromMessage(schema.AssistantMessage(a.msg, nil), nil, schema.Assistant, ""))
+	gen.Close()
+	return iter
+}
+
+// mockInterruptibleResumableAgent interrupts on first run and resumes on second
+type mockInterruptibleResumableAgent struct {
+	name string
+	t    *testing.T
+}
+
+func (a *mockInterruptibleResumableAgent) Name(_ context.Context) string {
+	return a.name
+}
+
+func (a *mockInterruptibleResumableAgent) Description(_ context.Context) string {
+	return "mock interruptible/resumable agent"
+}
+
+func (a *mockInterruptibleResumableAgent) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
+	iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
+	gen.Send(adk.EventFromMessage(schema.AssistantMessage("I will interrupt", nil), nil, schema.Assistant, ""))
+	gen.Send(adk.Interrupt(ctx, "interrupt data"))
+	gen.Close()
+	return iter
+}
+
+func (a *mockInterruptibleResumableAgent) Resume(ctx context.Context, info *adk.ResumeInfo, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] {
+	assert.True(a.t, info.WasInterrupted)
+	
+	// Check if this agent is the target of the resume
+	isResumeTarget, hasData, data := compose.GetResumeContext[string](ctx)
+	if isResumeTarget && hasData {
+		assert.Equal(a.t, "resume data", data)
+	}
+	
+	iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]()
+	gen.Send(adk.EventFromMessage(schema.AssistantMessage("I have resumed", nil), nil, schema.Assistant, ""))
+	gen.Close()
+	return iter
+}
+
+// TestNestedSupervisor_ConcurrentTransfer_WithInterruptAndResume tests a complex scenario:
+// - Nested supervisor hierarchy
+// - Concurrent transfers at two levels
+// - Interrupt at grandchild level
+// - Resume with targeted data
+func TestNestedSupervisor_ConcurrentTransfer_WithInterruptAndResume(t *testing.T) {
+	ctx := context.Background()
+
+	// 1. Define the agent hierarchy
+	grandChild1 := &mockSimpleAgent{name: "GrandChild1", msg: "GrandChild1 reporting"}
+	grandChild2 := &mockInterruptibleResumableAgent{name: "GrandChild2", t: t}
+	subSupervisor := &mockSupervisor{name: "SubSupervisor", targets: []string{"GrandChild1", "GrandChild2"}}
+
+	subAgent1 := &mockSimpleAgent{name: "SubAgent1", msg: "SubAgent1 reporting"}
+	superSupervisor := &mockSupervisor{name: "SuperSupervisor", targets: []string{"SubAgent1", "SubSupervisor"}}
+
+	// 2. Build the nested supervisor hierarchy
+	nestedSupervisor, err := New(ctx, &Config{
+		Supervisor: subSupervisor,
+		SubAgents:  []adk.Agent{grandChild1, grandChild2},
+	})
+	assert.NoError(t, err)
+
+	topSupervisor, err := New(ctx, &Config{
+		Supervisor: superSupervisor,
+		SubAgents:  []adk.Agent{subAgent1, nestedSupervisor},
+	})
+	assert.NoError(t, err)
+
+	// 3. Run the top-level supervisor and expect interrupt
+	runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: topSupervisor, CheckPointStore: newMyStore()})
+	aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("start")},
+		adk.WithCheckPointID("test-checkpoint"))
+
+	var finalEvent *adk.AgentEvent
+	var events []*adk.AgentEvent
+	for event, ok := aIter.Next(); ok; event, ok = aIter.Next() {
+		var role, content string
+		var toolCalls []schema.ToolCall
+		if event.Output != nil && event.Output.MessageOutput != nil {
+			role = string(event.Output.MessageOutput.Role)
+			content = event.Output.MessageOutput.Message.Content
+			toolCalls = event.Output.MessageOutput.Message.ToolCalls
+		}
+		t.Logf("Event: Agent=%s, Role=%s, Content=%s, ToolCalls= %v, Interrupted=%v, transfer=%v, concurrentTransfer=%v",
+			event.AgentName, role, content, toolCalls,
+			event.Action != nil && event.Action.Interrupted != nil,
+			event.Action != nil && event.Action.TransferToAgent != nil,
+			event.Action != nil && event.Action.ConcurrentTransferToAgent != nil)
+		
+		events = append(events, event)
+		if event.Action != nil && event.Action.Interrupted != nil {
+			finalEvent = event
+		}
+	}
+
+	if finalEvent == nil {
+		t.Fatal("Should have received an interrupt event")
+	}
+	assert.Equal(t, "SuperSupervisor", finalEvent.AgentName, "Interrupt should propagate to top supervisor")
+
+	// 4. Verify the execution sequence - handle complex concurrent execution
+	assert.Equal(t, 8, len(events), "Should have 8 events in initial execution")
+	
+	// Check SuperSupervisor concurrent transfer (events 0-1)
+	assert.Equal(t, "SuperSupervisor", events[0].AgentName)
+	assert.Equal(t, schema.Assistant, events[0].Output.MessageOutput.Role)
+	assert.True(t, events[1].Action != nil && events[1].Action.ConcurrentTransferToAgent != nil)
+	
+	// Event 2: SubAgent1 completes (first concurrent branch)
+	assert.Equal(t, "SubAgent1", events[2].AgentName)
+	assert.Equal(t, "SubAgent1 reporting", events[2].Output.MessageOutput.Message.Content)
+	
+	// Check SubSupervisor concurrent transfer (events 3-4)  
+	assert.Equal(t, "SubSupervisor", events[3].AgentName)
+	assert.Equal(t, schema.Assistant, events[3].Output.MessageOutput.Role)
+	assert.True(t, events[4].Action != nil && events[4].Action.ConcurrentTransferToAgent != nil)
+	
+	// Events 5-6: GrandChild2 and GrandChild1 (order can vary due to concurrency)
+	agentEvents := make(map[string]string)
+	for i := 5; i <= 6; i++ {
+		if events[i].Output != nil && events[i].Output.MessageOutput != nil {
+			agentEvents[events[i].AgentName] = events[i].Output.MessageOutput.Message.Content
+		}
+	}
+	
+	// Verify both grandchild agents completed/interrupted as expected
+	assert.Equal(t, "GrandChild1 reporting", agentEvents["GrandChild1"])
+	assert.Equal(t, "I will interrupt", agentEvents["GrandChild2"])
+	
+	// Check interrupt (event 7)
+	assert.Equal(t, "SuperSupervisor", events[7].AgentName)
+	assert.True(t, events[7].Action != nil && events[7].Action.Interrupted != nil)
+
+	// 5. Resume the execution with targeted resume data
+	resumeIter, err := runner.TargetedResume(ctx, "test-checkpoint", map[string]any{
+		"GrandChild2": "resume data",
+	})
+	assert.NoError(t, err)
+
+	// 6. Verify the resume flow completes successfully
+	var resumeEvents []*adk.AgentEvent
+	for event, ok := resumeIter.Next(); ok; event, ok = resumeIter.Next() {
+		var role, content string
+		var toolCalls []schema.ToolCall
+		if event.Output != nil && event.Output.MessageOutput != nil {
+			role = string(event.Output.MessageOutput.Role)
+			content = event.Output.MessageOutput.Message.Content
+			toolCalls = event.Output.MessageOutput.Message.ToolCalls
+		}
+		t.Logf("Resume Event: Agent=%s, Role=%s, Content=%s, ToolCalls= %v, Interrupted=%v, transfer=%v, concurrentTransfer=%v",
+			event.AgentName, role, content, toolCalls,
+			event.Action != nil && event.Action.Interrupted != nil,
+			event.Action != nil && event.Action.TransferToAgent != nil,
+			event.Action != nil && event.Action.ConcurrentTransferToAgent != nil)
+		resumeEvents = append(resumeEvents, event)
+	}
+
+	assert.Equal(t, 7, len(resumeEvents), "Should have 7 events in resume execution")
+	
+	// Check GrandChild2 resume
+	assert.Equal(t, "GrandChild2", resumeEvents[0].AgentName)
+	assert.Equal(t, "I have resumed", resumeEvents[0].Output.MessageOutput.Message.Content)
+	
+	// Check SubSupervisor self-return and completion
+	assert.Equal(t, "SubSupervisor", resumeEvents[1].AgentName)
+	assert.Equal(t, schema.Assistant, resumeEvents[1].Output.MessageOutput.Role)
+	
+	assert.Equal(t, "SubSupervisor", resumeEvents[2].AgentName)
+	assert.True(t, resumeEvents[2].Action != nil && resumeEvents[2].Action.TransferToAgent != nil)
+	assert.Equal(t, "SubSupervisor", resumeEvents[2].Action.TransferToAgent.DestAgentName)
+	
+	assert.Equal(t, "SubSupervisor", resumeEvents[3].AgentName)
+	assert.Equal(t, "job done", resumeEvents[3].Output.MessageOutput.Message.Content)
+	
+	// Check SuperSupervisor self-return and completion
+	assert.Equal(t, "SuperSupervisor", resumeEvents[4].AgentName)
+	assert.Equal(t, schema.Assistant, resumeEvents[4].Output.MessageOutput.Role)
+	
+	assert.Equal(t, "SuperSupervisor", resumeEvents[5].AgentName)
+	assert.True(t, resumeEvents[5].Action != nil && resumeEvents[5].Action.TransferToAgent != nil)
+	assert.Equal(t, "SuperSupervisor", resumeEvents[5].Action.TransferToAgent.DestAgentName)
+	
+	assert.Equal(t, "SuperSupervisor", resumeEvents[6].AgentName)
+	assert.Equal(t, "job done", resumeEvents[6].Output.MessageOutput.Message.Content)
+}
+
+func newMyStore() *myStore {
+	return &myStore{
+		m: map[string][]byte{},
+	}
+}
+
+type myStore struct {
+	m map[string][]byte
+}
+
+func (m *myStore) Set(_ context.Context, key string, value []byte) error {
+	m.m[key] = value
+	return nil
+}
+
+func (m *myStore) Get(_ context.Context, key string) ([]byte, bool, error) {
+	v, ok := m.m[key]
+	return v, ok, nil
+}