Add files via upload

2026-08-15 23:50:32 +02:00 · 2026-08-15 01:39:58 +08:00
parent 4fe6defa28
commit 70b01206e4
47 changed files with 5598 additions and 1191 deletions
@@ -7,6 +7,10 @@ import (
 	"strings"
 	"testing"

+	"cyberstrike-ai/internal/config"
+
+	localbk "github.com/cloudwego/eino-ext/adk/backend/local"
+	"github.com/cloudwego/eino/adk"
 	"github.com/cloudwego/eino/components/tool"
 	"github.com/cloudwego/eino/schema"
 )
@@ -28,6 +32,169 @@ func TestReductionCacheRootDir(t *testing.T) {
 	}
 }

+func TestBuildAgenticReductionMiddlewareClearsOldAgenticToolResult(t *testing.T) {
+	ctx := context.Background()
+	loc, err := localbk.NewBackend(ctx, &localbk.Config{})
+	if err != nil {
+		t.Fatalf("NewBackend: %v", err)
+	}
+	root := t.TempDir()
+	mw, err := buildAgenticReductionMiddleware(ctx, config.MultiAgentEinoMiddlewareConfig{
+		ReductionRootDir:           root,
+		ReductionMaxTokensForClear: 1,
+	}, "", "conv-1", loc, nil)
+	if err != nil {
+		t.Fatalf("buildAgenticReductionMiddleware: %v", err)
+	}
+	oldText := strings.Repeat("old-tool-output-", 20)
+	newText := strings.Repeat("new-tool-output-", 20)
+	state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
+		Messages: []*schema.AgenticMessage{
+			agenticAssistantToolCall("old-call", "execute", `{"command":"old"}`),
+			agenticToolResult("old-call", "execute", oldText),
+			agenticAssistantToolCall("new-call", "execute", `{"command":"new"}`),
+			agenticToolResult("new-call", "execute", newText),
+		},
+	}
+	_, out, err := mw.BeforeModelRewriteState(ctx, state, nil)
+	if err != nil {
+		t.Fatalf("BeforeModelRewriteState: %v", err)
+	}
+	oldGot := out.Messages[1].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text
+	newGot := out.Messages[3].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text
+	if oldGot == oldText {
+		t.Fatal("agentic reduction did not clear old oversized tool result")
+	}
+	if !strings.Contains(oldGot, "read_file") {
+		t.Fatalf("cleared content should mention read_file, got %q", oldGot)
+	}
+	if newGot != newText {
+		t.Fatalf("latest tool result should be retained, got %q", newGot)
+	}
+}
+
+func agenticAssistantToolCall(callID, name, arguments string) *schema.AgenticMessage {
+	return &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeAssistant,
+		ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{
+			CallID:    callID,
+			Name:      name,
+			Arguments: arguments,
+		})},
+	}
+}
+
+func agenticToolResult(callID, name, text string) *schema.AgenticMessage {
+	return &schema.AgenticMessage{
+		Role: schema.AgenticRoleTypeUser,
+		ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{
+			CallID: callID,
+			Name:   name,
+			Content: []*schema.FunctionToolResultContentBlock{{
+				Type: schema.FunctionToolResultContentBlockTypeText,
+				Text: &schema.UserInputText{Text: text},
+			}},
+		})},
+	}
+}
+
+func TestBuildAgenticReductionMiddlewareHandlesSingleAgenticToolResult(t *testing.T) {
+	ctx := context.Background()
+	loc, err := localbk.NewBackend(ctx, &localbk.Config{})
+	if err != nil {
+		t.Fatalf("NewBackend: %v", err)
+	}
+	mw, err := buildAgenticReductionMiddleware(ctx, config.MultiAgentEinoMiddlewareConfig{
+		ReductionRootDir:           t.TempDir(),
+		ReductionMaxTokensForClear: 1,
+	}, "", "conv-1", loc, nil)
+	if err != nil {
+		t.Fatalf("buildAgenticReductionMiddleware: %v", err)
+	}
+	state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
+		Messages: []*schema.AgenticMessage{
+			{
+				Role: schema.AgenticRoleTypeUser,
+				ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{
+					CallID: "call-1",
+					Name:   "execute",
+					Content: []*schema.FunctionToolResultContentBlock{{
+						Type: schema.FunctionToolResultContentBlockTypeText,
+						Text: &schema.UserInputText{Text: strings.Repeat("tool-output-", 20)},
+					}},
+				})},
+			},
+		},
+	}
+	_, out, err := mw.BeforeModelRewriteState(ctx, state, nil)
+	if err != nil {
+		t.Fatalf("BeforeModelRewriteState: %v", err)
+	}
+	got := out.Messages[0].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text
+	if got != strings.Repeat("tool-output-", 20) {
+		t.Fatalf("single retained tool result should not be cleared, got %q", got)
+	}
+}
+
+func TestPrependEinoAgenticMiddlewaresRespectsReductionPlacement(t *testing.T) {
+	ctx := context.Background()
+	loc, err := localbk.NewBackend(ctx, &localbk.Config{})
+	if err != nil {
+		t.Fatalf("NewBackend: %v", err)
+	}
+	patchToolCalls := false
+	mw := &config.MultiAgentEinoMiddlewareConfig{
+		ReductionEnable:            true,
+		ReductionRootDir:           t.TempDir(),
+		ReductionMaxTokensForClear: 100,
+		PatchToolCalls:             &patchToolCalls,
+	}
+	_, mainHandlers, _, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWMain, nil, loc, "", "conv-1", "", nil)
+	if err != nil {
+		t.Fatalf("prepend main: %v", err)
+	}
+	if len(mainHandlers) != 1 {
+		t.Fatalf("main handlers = %d, want reduction", len(mainHandlers))
+	}
+	_, subHandlers, _, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWSub, nil, loc, "", "conv-1", "", nil)
+	if err != nil {
+		t.Fatalf("prepend sub: %v", err)
+	}
+	if len(subHandlers) != 0 {
+		t.Fatalf("sub handlers = %d, want skipped when reduction_sub_agents=false", len(subHandlers))
+	}
+	mw.ReductionSubAgents = true
+	_, subHandlers, _, err = prependEinoAgenticMiddlewares(ctx, mw, einoMWSub, nil, loc, "", "conv-1", "", nil)
+	if err != nil {
+		t.Fatalf("prepend sub enabled: %v", err)
+	}
+	if len(subHandlers) != 1 {
+		t.Fatalf("sub handlers = %d, want reduction when reduction_sub_agents=true", len(subHandlers))
+	}
+}
+
+func TestPrependEinoAgenticMiddlewaresMountsToolSearchAndPatchToolCalls(t *testing.T) {
+	ctx := context.Background()
+	mw := &config.MultiAgentEinoMiddlewareConfig{
+		ToolSearchEnable:        true,
+		ToolSearchMinTools:      20,
+		ToolSearchAlwaysVisible: 5,
+	}
+	outTools, handlers, toolSearchActive, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWMain, stubTools(25), nil, "", "conv-test", "", nil)
+	if err != nil {
+		t.Fatalf("prependEinoAgenticMiddlewares: %v", err)
+	}
+	if !toolSearchActive {
+		t.Fatal("agentic tool_search should be active")
+	}
+	if len(outTools) != 5 {
+		t.Fatalf("mounted tools = %d, want static visible tools only", len(outTools))
+	}
+	if len(handlers) != 2 {
+		t.Fatalf("handlers = %d, want patchtoolcalls + toolsearch", len(handlers))
+	}
+}
+
 type stubTool struct{ name string }

 func (s stubTool) Info(_ context.Context) (*schema.ToolInfo, error) {