Add files via upload

internal/multiagent/eino_turn_loop_iterator_starter_test.go

@@ -0,0 +1,188 @@
+package multiagent
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/schema"
+)
+
+type fakeTurnLoopRuntimeControl struct {
+	mu          sync.Mutex
+	runCalled   bool
+	stopIdle    bool
+	stopped     string
+	pushedNotes []string
+	pushOK      bool
+}
+
+func (f *fakeTurnLoopRuntimeControl) Run(context.Context) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.runCalled = true
+}
+
+func (f *fakeTurnLoopRuntimeControl) PushInterruptContinue(note string) bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.pushedNotes = append(f.pushedNotes, note)
+	return f.pushOK
+}
+
+func (f *fakeTurnLoopRuntimeControl) StopImmediate(cause string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.stopped = cause
+}
+
+func (f *fakeTurnLoopRuntimeControl) StopWhenIdle() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.stopIdle = true
+}
+
+func (f *fakeTurnLoopRuntimeControl) Wait() *adk.TurnLoopExitState[EinoTurnLoopItem, *schema.Message] {
+	return nil
+}
+
+func (f *fakeTurnLoopRuntimeControl) snapshot() (runCalled bool, stopIdle bool, stopped string, pushed []string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.runCalled, f.stopIdle, f.stopped, append([]string(nil), f.pushedNotes...)
+}
+
+func TestEinoTurnLoopIteratorStarterBindsRegistrarsAndProgress(t *testing.T) {
+	fakeRuntime := &fakeTurnLoopRuntimeControl{pushOK: true}
+	oldAgentCleared := false
+	oldTurnCleared := false
+	unregisterAgent := func() { oldAgentCleared = true }
+	unregisterTurn := func() { oldTurnCleared = true }
+	var interruptPush func(string) bool
+	var cancelPush func(error) bool
+	var createdCfg EinoTurnLoopRuntimeConfig
+	var events []struct {
+		eventType string
+		message   string
+		data      map[string]interface{}
+	}
+
+	iter := newEinoTurnLoopIteratorStarter(einoTurnLoopIteratorStarterConfig{
+		Context:                     context.Background(),
+		ConversationID:              "conv",
+		OrchMode:                    "deep",
+		CheckPointID:                "runner-checkpoint",
+		UnregisterAgentCancel:       &unregisterAgent,
+		UnregisterTurnLoopInterrupt: &unregisterTurn,
+		RuntimeCancelRegistrar: func(push func(error) bool) func() {
+			cancelPush = push
+			return func() {}
+		},
+		TurnLoopInterruptRegistrar: func(push func(string) bool) func() {
+			interruptPush = push
+			return func() {}
+		},
+		RuntimeFactory: func(cfg EinoTurnLoopRuntimeConfig) einoTurnLoopRuntimeControl {
+			createdCfg = cfg
+			return fakeRuntime
+		},
+		Progress: func(eventType, message string, data interface{}) {
+			item := struct {
+				eventType string
+				message   string
+				data      map[string]interface{}
+			}{eventType: eventType, message: message}
+			if m, ok := data.(map[string]interface{}); ok {
+				item.data = m
+			}
+			events = append(events, item)
+		},
+	}).Start([]adk.Message{})
+
+	if iter == nil {
+		t.Fatal("iterator should be created")
+	}
+	if !oldAgentCleared || !oldTurnCleared {
+		t.Fatalf("oldAgentCleared=%v oldTurnCleared=%v, want both true", oldAgentCleared, oldTurnCleared)
+	}
+	if interruptPush == nil {
+		t.Fatal("turn loop interrupt registrar was not bound")
+	}
+	if cancelPush == nil {
+		t.Fatal("runtime cancel registrar was not bound")
+	}
+	if createdCfg.CheckpointID != buildEinoTurnLoopCheckpointID("deep") {
+		t.Fatalf("checkpoint id = %q, want turn loop checkpoint id", createdCfg.CheckpointID)
+	}
+	if !interruptPush("  focus ssh  ") {
+		t.Fatal("interrupt push should return runtime result")
+	}
+
+	runCalled, stopIdle, _, pushed := fakeRuntime.snapshot()
+	if !runCalled || !stopIdle {
+		t.Fatalf("runCalled=%v stopIdle=%v, want both true", runCalled, stopIdle)
+	}
+	if len(pushed) != 1 || pushed[0] != "  focus ssh  " {
+		t.Fatalf("pushed notes = %#v", pushed)
+	}
+	if len(events) != 2 {
+		t.Fatalf("events = %#v, want user interrupt and progress", events)
+	}
+	if events[0].eventType != "user_interrupt_continue" || events[0].data["rawReason"] != "focus ssh" {
+		t.Fatalf("first event = %#v", events[0])
+	}
+	if events[1].eventType != "progress" {
+		t.Fatalf("second event = %#v", events[1])
+	}
+}
+
+func TestEinoTurnLoopIteratorStarterRuntimeCancel(t *testing.T) {
+	fakeRuntime := &fakeTurnLoopRuntimeControl{pushOK: true}
+	var nativeCancelCause atomic.Value
+	var cancelPush func(error) bool
+	var unregisterAgent func()
+
+	newEinoTurnLoopIteratorStarter(einoTurnLoopIteratorStarterConfig{
+		Context:               context.Background(),
+		ConversationID:        "conv",
+		OrchMode:              "eino_single",
+		NativeCancelCause:     &nativeCancelCause,
+		UnregisterAgentCancel: &unregisterAgent,
+		RuntimeCancelRegistrar: func(push func(error) bool) func() {
+			cancelPush = push
+			return func() {}
+		},
+		RuntimeFactory: func(EinoTurnLoopRuntimeConfig) einoTurnLoopRuntimeControl {
+			return fakeRuntime
+		},
+	}).Start(nil)
+	if cancelPush == nil {
+		t.Fatal("runtime cancel registrar was not bound")
+	}
+
+	if !cancelPush(ErrInterruptContinue) {
+		t.Fatal("interrupt continue cancel should be handled by TurnLoop push")
+	}
+	_, _, stopped, pushed := fakeRuntime.snapshot()
+	if stopped != "" {
+		t.Fatalf("stopped = %q, want no immediate stop for interrupt continue", stopped)
+	}
+	if len(pushed) != 1 || pushed[0] != "" {
+		t.Fatalf("pushed notes = %#v, want empty interrupt continue note", pushed)
+	}
+
+	stopErr := errors.New("stop now")
+	if !cancelPush(stopErr) {
+		t.Fatal("regular cancel should be handled")
+	}
+	_, _, stopped, _ = fakeRuntime.snapshot()
+	if stopped != "task_cancelled" {
+		t.Fatalf("stopped = %q, want task_cancelled", stopped)
+	}
+	if got, _ := nativeCancelCause.Load().(error); !errors.Is(got, stopErr) {
+		t.Fatalf("native cancel cause = %v, want %v", got, stopErr)
+	}
+}