Add files via upload

internal/multiagent/eino_main_assistant_stream_handler_test.go

@@ -0,0 +1,103 @@
+package multiagent
+
+import "testing"
+
+func TestEinoMainAssistantStreamHandlerEmitsAndRecords(t *testing.T) {
+	var eventTypes []string
+	var messages []string
+	progress := func(eventType, message string, _ interface{}) {
+		eventTypes = append(eventTypes, eventType)
+		messages = append(messages, message)
+	}
+	out := newEinoAssistantOutputAccumulator("deep")
+	runMsgs := newEinoRunMessageAccumulator(nil)
+	emitter := newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 2, progress, nil)
+	handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{
+		AgentName:       "lead",
+		Emitter:         emitter,
+		AssistantOutput: out,
+		RunMessages:     runMsgs,
+	})
+
+	if !handler.EmitDelta("he") {
+		t.Fatal("first delta should emit")
+	}
+	if !handler.EmitDelta("hello") {
+		t.Fatal("cumulative chunk should emit tail")
+	}
+	if got := handler.Finish(); got != "hello" {
+		t.Fatalf("finish = %q, want hello", got)
+	}
+
+	if len(eventTypes) != 3 || eventTypes[0] != "response_start" || eventTypes[1] != "response_delta" || eventTypes[2] != "response_delta" {
+		t.Fatalf("events = %#v", eventTypes)
+	}
+	if messages[1] != "he" || messages[2] != "llo" {
+		t.Fatalf("delta messages = %#v", messages)
+	}
+	if out.LastAssistant() != "hello" {
+		t.Fatalf("last assistant = %q", out.LastAssistant())
+	}
+	if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != "hello" {
+		t.Fatalf("run messages = %#v", msgs)
+	}
+}
+
+func TestEinoMainAssistantStreamHandlerSuppressesDuplicateExecuteStdout(t *testing.T) {
+	var eventTypes []string
+	progress := func(eventType, _ string, _ interface{}) {
+		eventTypes = append(eventTypes, eventType)
+	}
+	stdoutDup := newEinoExecuteStdoutSuppressor()
+	stdoutDup.Record("execute", "hello", false)
+	out := newEinoAssistantOutputAccumulator("deep")
+	runMsgs := newEinoRunMessageAccumulator(nil)
+	handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{
+		AgentName:        "lead",
+		Emitter:          newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 1, progress, nil),
+		StdoutSuppressor: stdoutDup,
+		AssistantOutput:  out,
+		RunMessages:      runMsgs,
+	})
+
+	if handler.EmitDelta("hello") {
+		t.Fatal("duplicate execute stdout should not emit delta")
+	}
+	if got := handler.Finish(); got != "hello" {
+		t.Fatalf("finish = %q, want hello", got)
+	}
+	if len(eventTypes) != 0 {
+		t.Fatalf("events = %#v, want none", eventTypes)
+	}
+	if stdoutDup.Peek() != "" {
+		t.Fatal("duplicate target should be cleared on finish")
+	}
+	if out.LastAssistant() != "hello" {
+		t.Fatalf("last assistant = %q", out.LastAssistant())
+	}
+	if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != "hello" {
+		t.Fatalf("run messages = %#v", msgs)
+	}
+}
+
+func TestEinoMainAssistantStreamHandlerRecordsWithoutProgress(t *testing.T) {
+	out := newEinoAssistantOutputAccumulator("plan_execute")
+	runMsgs := newEinoRunMessageAccumulator(nil)
+	handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{
+		AgentName:       "executor",
+		Emitter:         newEinoMainResponseStreamEmitter("conv-1", "plan_execute", "executor", "stream-1", 1, nil, nil),
+		AssistantOutput: out,
+		RunMessages:     runMsgs,
+	})
+
+	handler.EmitDelta(`{"response":"done"}`)
+	if got := handler.Finish(); got != `{"response":"done"}` {
+		t.Fatalf("finish = %q", got)
+	}
+	if out.LastPlanExecuteExecutor() != "done" {
+		t.Fatalf("executor output = %q", out.LastPlanExecuteExecutor())
+	}
+	if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != `{"response":"done"}` {
+		t.Fatalf("run messages = %#v", msgs)
+	}
+}