CyberStrikeAI/internal/multiagent/eino_exit_fallback_test.go

package multiagent

import (
	"testing"

	"github.com/cloudwego/eino/schema"
)

func TestEinoExtractFallbackAssistantFromMsgs_exitToolMessage(t *testing.T) {
	u := schema.UserMessage("hi")
	tm := schema.ToolMessage("answer for user", "call-exit-1")
	tm.ToolName = "exit"
	if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{u, tm}); got != "answer for user" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_lastExitWins(t *testing.T) {
	msgs := []*schema.Message{
		schema.UserMessage("hi"),
		toolExitMsg("first", "c1"),
		toolExitMsg("second", "c2"),
	}
	if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "second" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_fromAssistantToolCalls(t *testing.T) {
	m := schema.AssistantMessage("", []schema.ToolCall{{
		ID:   "x",
		Type: "function",
		Function: schema.FunctionCall{
			Name:      "exit",
			Arguments: `{"final_result":"from args"}`,
		},
	}})
	if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{m}); got != "from args" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_prefersToolOverEarlierAssistant(t *testing.T) {
	asst := schema.AssistantMessage("", []schema.ToolCall{{
		ID:   "x",
		Type: "function",
		Function: schema.FunctionCall{
			Name:      "exit",
			Arguments: `{"final_result":"from args"}`,
		},
	}})
	tool := toolExitMsg("from tool", "c1")
	if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{asst, tool}); got != "from tool" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_plainAssistant(t *testing.T) {
	msgs := []*schema.Message{
		schema.UserMessage("hi"),
		schema.AssistantMessage("plain answer", nil),
	}
	if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "plain answer" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_finalAssistantAfterToolResult(t *testing.T) {
	msgs := []*schema.Message{
		schema.UserMessage("hi"),
		schema.AssistantMessage("", []schema.ToolCall{{
			ID:   "call-1",
			Type: "function",
			Function: schema.FunctionCall{
				Name:      "execute",
				Arguments: `{"command":"pwd"}`,
			},
		}}),
		schema.ToolMessage("/tmp", "call-1", schema.WithToolName("execute")),
		schema.AssistantMessage("final after tool", nil),
	}
	if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "final after tool" {
		t.Fatalf("got %q", got)
	}
}

func TestEinoExtractFallbackAssistantFromMsgs_doesNotUseAssistantBeforeUnfinishedToolResult(t *testing.T) {
	msgs := []*schema.Message{
		schema.UserMessage("hi"),
		schema.AssistantMessage("I will inspect that.", nil),
		schema.AssistantMessage("", []schema.ToolCall{{
			ID:   "call-1",
			Type: "function",
			Function: schema.FunctionCall{
				Name:      "execute",
				Arguments: `{"command":"pwd"}`,
			},
		}}),
		schema.ToolMessage("/tmp", "call-1", schema.WithToolName("execute")),
	}
	if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "" {
		t.Fatalf("got %q, want empty", got)
	}
}

func TestEinoRunResultBuilderFinalFallsBackToPlainAssistantTrace(t *testing.T) {
	runMessages := newEinoRunMessageAccumulator(nil)
	runMessages.Append(schema.UserMessage("hi"))
	runMessages.Append(schema.AssistantMessage("plain answer", nil))

	got := newEinoRunResultBuilder(einoRunResultBuilderConfig{
		OrchMode:    "deep",
		EmptyHint:   "empty",
		RunMessages: runMessages,
	}).BuildFinal()

	if got.Response != "plain answer" {
		t.Fatalf("response = %q, want plain answer", got.Response)
	}
}

func toolExitMsg(content, callID string) *schema.Message {
	m := schema.ToolMessage(content, callID)
	m.ToolName = "exit"
	return m
}