moteva/server/internal/infrastructure/piagent/design_agent_test.go

package piagent

import (
	"context"
	"errors"
	"strconv"
	"strings"
	"sync"
	"testing"
	"unicode/utf8"

	"img_infinite_canvas/internal/domain/design"
)

type fakeRunner struct{}

func (fakeRunner) Plan(ctx context.Context, req design.AgentRequest) (design.AgentPlan, error) {
	return design.AgentPlan{
		Message: design.Message{ID: "message-1", Role: "assistant", Title: "done", Content: "done"},
		GeneratedNodes: []design.Node{
			{
				ID:      "node-1",
				Type:    design.NodeTypeFrame,
				Title:   "Website Hero",
				Width:   520,
				Height:  300,
				Content: req.Prompt,
			},
			{
				ID:      "mock-image-1",
				Type:    design.NodeTypeImage,
				Title:   "Mood x Color",
				Width:   520,
				Height:  640,
				Content: req.Prompt,
			},
			{
				ID:      "mock-image-2",
				Type:    design.NodeTypeImage,
				Title:   "Mobile Page",
				Width:   260,
				Height:  460,
				Content: req.Prompt,
			},
		},
	}, nil
}

func (fakeRunner) Replay(ctx context.Context, project design.Project) (design.AgentReplay, error) {
	return design.AgentReplay{ProjectID: project.ID, Messages: project.Messages}, nil
}

type fakeImageGenerator struct {
	mu      sync.Mutex
	called  bool
	prompt  string
	opts    ImageRequestOptions
	calls   int
	options []ImageRequestOptions
}

type failingImageGenerator struct {
	calls int
	err   error
}

type partialFailImageGenerator struct {
	mu    sync.Mutex
	calls int
}

func (g *failingImageGenerator) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
	g.calls++
	return GeneratedImage{}, g.err
}

func (g *partialFailImageGenerator) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
	g.mu.Lock()
	g.calls++
	call := g.calls
	g.mu.Unlock()
	if call == 1 {
		return GeneratedImage{URL: "https://example.com/generated-partial.png", URLs: []string{"https://example.com/generated-partial.png"}}, nil
	}
	return GeneratedImage{}, errors.New("upstream response did not return image")
}

func (g *fakeImageGenerator) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.called = true
	g.prompt = prompt
	g.opts = opts
	g.calls++
	callIndex := imageOptionIndex(prompt, g.calls)
	g.options = append(g.options, opts)
	count := opts.Count
	if count < 1 {
		count = 1
	}
	urls := make([]string, 0, count)
	for index := 0; index < count; index++ {
		urls = append(urls, "https://example.com/generated-"+strconv.Itoa(callIndex+index)+".png")
	}
	return GeneratedImage{URL: urls[0], URLs: urls}, nil
}

func imageOptionIndex(prompt string, fallback int) int {
	for index := 1; index <= 10; index++ {
		if strings.Contains(prompt, "option "+strconv.Itoa(index)+" of") {
			return index
		}
	}
	return fallback
}

func (g *fakeImageGenerator) callCount() int {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.calls
}

type fakeCreativeAgent struct{}

func (fakeCreativeAgent) Decide(ctx context.Context, req design.AgentDecisionRequest) (design.AgentDecision, error) {
	return design.AgentDecision{}, nil
}

func (fakeCreativeAgent) Respond(ctx context.Context, req design.AgentConversationRequest) (string, error) {
	return "", nil
}

func (fakeCreativeAgent) SummarizeImage(ctx context.Context, req design.AgentImageSummaryRequest) (string, error) {
	return "这张图已经生成。建议下一步检查首屏构图、商品露出和转化按钮层级。", nil
}

func (fakeCreativeAgent) BuildImagePrompt(ctx context.Context, req design.AgentImagePromptRequest) (design.AgentImagePrompt, error) {
	return design.AgentImagePrompt{
		Title:  "Still Day 首页升级",
		Prompt: "基于 still day 现有首页方向，保留高饱和撞色，升级电商首屏、商品网格和品牌页脚。",
	}, nil
}

type longCreativeAgent struct{}

func (longCreativeAgent) Decide(ctx context.Context, req design.AgentDecisionRequest) (design.AgentDecision, error) {
	return design.AgentDecision{}, nil
}

func (longCreativeAgent) Respond(ctx context.Context, req design.AgentConversationRequest) (string, error) {
	return "", nil
}

func (longCreativeAgent) SummarizeImage(ctx context.Context, req design.AgentImageSummaryRequest) (string, error) {
	return "", nil
}

func (longCreativeAgent) BuildImagePrompt(ctx context.Context, req design.AgentImagePromptRequest) (design.AgentImagePrompt, error) {
	return design.AgentImagePrompt{
		Title:  "精简生图指令",
		Prompt: strings.Repeat("大量上下文不应该直接进入 gpt-image-2。", 240),
	}, nil
}

func TestDesignAgentGeneratesImageThroughPiAgent(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	plan, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "生成一张 still day 官网首页图片",
		Mode:       "image",
		ImageModel: "gpt-image-2",
		ImageSize:  "1536x1024",
	})
	if err != nil {
		t.Fatal(err)
	}
	if !imageGenerator.called {
		t.Fatal("expected image generator to be called")
	}
	if !containsAll(imageGenerator.prompt, []string{"基于 still day 现有首页方向", "高饱和撞色"}) {
		t.Fatalf("unexpected prompt: %s", imageGenerator.prompt)
	}
	if imageGenerator.opts.Model != "gpt-image-2" || imageGenerator.opts.Size != "1536x1024" {
		t.Fatalf("unexpected image options: %#v", imageGenerator.opts)
	}
	if imageGenerator.opts.Count != 1 {
		t.Fatalf("expected one image request, got %d", imageGenerator.opts.Count)
	}
	if len(plan.GeneratedNodes) != 1 {
		t.Fatalf("expected one node, got %d", len(plan.GeneratedNodes))
	}
	node := plan.GeneratedNodes[0]
	if node.Type != design.NodeTypeImage {
		t.Fatalf("expected image node, got %s", node.Type)
	}
	if node.Content != "https://example.com/generated-1.png" {
		t.Fatalf("expected generated image url, got %s", node.Content)
	}
	if node.Title != "Still Day 首页升级" {
		t.Fatalf("expected generated image title, got %s", node.Title)
	}
	if node.Width != 1536 || node.Height != 1024 {
		t.Fatalf("expected node size from request image size, got %.0fx%.0f", node.Width, node.Height)
	}
	if plan.Message.Title != "Still Day 首页升级 已生成" {
		t.Fatalf("expected contextual completion title, got %s", plan.Message.Title)
	}
	if plan.Message.Content == "" {
		t.Fatal("expected image generation summary")
	}
	if !containsAll(plan.Message.Content, []string{"首屏构图", "商品露出", "转化按钮"}) {
		t.Fatalf("expected creative agent summary, got %s", plan.Message.Content)
	}
}

func TestDesignAgentUsesDefaultImageSizeWithoutExplicitRequest(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	plan, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "生成一张 still day 官网首页图片",
		Mode:       "design",
		ImageModel: "gpt-image-2",
	})
	if err != nil {
		t.Fatal(err)
	}
	if imageGenerator.opts.Size != "1024x1024" {
		t.Fatalf("expected default image size, got %q", imageGenerator.opts.Size)
	}
	if len(plan.GeneratedNodes) != 1 {
		t.Fatalf("expected one generated node, got %d", len(plan.GeneratedNodes))
	}
	if plan.GeneratedNodes[0].Width != 1024 || plan.GeneratedNodes[0].Height != 1024 {
		t.Fatalf("expected default generated image size, got %.0fx%.0f", plan.GeneratedNodes[0].Width, plan.GeneratedNodes[0].Height)
	}
}

func TestDesignAgentUsesExplicitPromptImageRatio(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	plan, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "生成一张 16:9 的 still day 官网首页图片",
		Mode:       "design",
		ImageModel: "gpt-image-2",
	})
	if err != nil {
		t.Fatal(err)
	}
	if imageGenerator.opts.Size != "2048x1152" {
		t.Fatalf("expected explicit 16:9 image size, got %q", imageGenerator.opts.Size)
	}
	if len(plan.GeneratedNodes) != 1 || plan.GeneratedNodes[0].Width != 2048 || plan.GeneratedNodes[0].Height != 1152 {
		t.Fatalf("expected node size from explicit ratio, got %#v", plan.GeneratedNodes)
	}
}

func TestDesignAgentGeneratesModelPlannedImageTasks(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	plan, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "设计二次元水墨人物形象，输出角色正面、侧面、细节和色彩方案。",
		Mode:       "image",
		ImageModel: "gpt-image-2",
		ImageSize:  "1024x1024",
		TaskPlan: design.AgentTaskPlan{
			Intent:     "image",
			ImageCount: 3,
			ImageTasks: []design.AgentImageTask{
				{Title: "水墨角色_正面", Brief: "正面全身立绘"},
				{Title: "水墨角色_侧面", Brief: "侧面全身立绘"},
				{Title: "水墨角色_细节", Brief: "局部细节拆解"},
			},
		},
		Project: design.Project{
			Nodes: []design.Node{
				{ID: "slot-1", Type: design.NodeTypeImage, Status: "generating", X: 10},
				{ID: "slot-2", Type: design.NodeTypeImage, Status: "generating", X: 20},
				{ID: "slot-3", Type: design.NodeTypeImage, Status: "generating", X: 30},
			},
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if imageGenerator.callCount() != 3 {
		t.Fatalf("expected 3 image calls, got %d", imageGenerator.callCount())
	}
	idempotencyKeys := make(map[string]bool)
	for _, opts := range imageGenerator.options {
		if opts.Count != 1 {
			t.Fatalf("expected individual image request count 1, got %#v", imageGenerator.options)
		}
		if opts.IdempotencyKey == "" {
			t.Fatalf("expected idempotency key for image request, got %#v", imageGenerator.options)
		}
		if idempotencyKeys[opts.IdempotencyKey] {
			t.Fatalf("expected unique idempotency keys, got %#v", imageGenerator.options)
		}
		idempotencyKeys[opts.IdempotencyKey] = true
	}
	if len(plan.GeneratedNodes) != 3 {
		t.Fatalf("expected 3 generated nodes, got %#v", plan.GeneratedNodes)
	}
	for index, node := range plan.GeneratedNodes {
		wantID := "slot-" + strconv.Itoa(index+1)
		if node.ID != wantID || node.Content == "" || node.Status != "success" {
			t.Fatalf("unexpected generated node %d: %#v", index, node)
		}
	}
}

func TestDesignAgentGeneratesWhenTaskPlanIntentIsImageEvenInAgentMode(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	plan, err := agent.PlanIncremental(context.Background(), design.AgentRequest{
		Prompt:     "来一版",
		Mode:       "agent",
		ImageModel: "gpt-image-2",
		TaskPlan: design.AgentTaskPlan{
			Intent:     "image",
			ImageCount: 1,
			ImageTasks: []design.AgentImageTask{
				{Title: "画面方案", Brief: "生成一个高级电商主视觉，强调产品层次和留白。"},
			},
		},
		Project: design.Project{
			Nodes: []design.Node{
				{ID: "slot-1", Type: design.NodeTypeImage, Status: "generating", X: 10},
			},
		},
	}, nil)
	if err != nil {
		t.Fatal(err)
	}
	if imageGenerator.callCount() != 1 {
		t.Fatalf("expected image generation to run from task plan intent, got %d calls", imageGenerator.callCount())
	}
	if len(plan.GeneratedNodes) != 1 || plan.GeneratedNodes[0].Content == "" {
		t.Fatalf("expected generated image node, got %#v", plan.GeneratedNodes)
	}
}

func TestDesignAgentReturnsPartialFailureWhenPlannedImageIsMissing(t *testing.T) {
	imageGenerator := &partialFailImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})
	var emitted []design.Node

	plan, err := agent.PlanIncremental(context.Background(), design.AgentRequest{
		Prompt: "生成两张视觉",
		Mode:   "image",
		TaskPlan: design.AgentTaskPlan{
			Intent:     "image",
			ImageCount: 2,
			ImageTasks: []design.AgentImageTask{
				{Title: "方案一", Brief: "生成第一张视觉"},
				{Title: "方案二", Brief: "生成第二张视觉"},
			},
		},
		Project: design.Project{
			Nodes: []design.Node{
				{ID: "slot-1", Type: design.NodeTypeImage, Status: "generating"},
				{ID: "slot-2", Type: design.NodeTypeImage, Status: "generating"},
			},
		},
	}, func(result design.AgentImageGenerationResult) error {
		emitted = append(emitted, result.Node)
		return nil
	})
	if err == nil {
		t.Fatal("expected partial generation failure")
	}
	if !strings.Contains(err.Error(), "only 1 of 2 planned images returned") {
		t.Fatalf("expected partial failure count, got %v", err)
	}
	if len(emitted) != 1 || len(plan.GeneratedNodes) != 1 {
		t.Fatalf("expected one successful node to be preserved, emitted=%#v plan=%#v", emitted, plan.GeneratedNodes)
	}
}

func TestGeneratedImageSlotsIgnoreCanvasActionPlaceholders(t *testing.T) {
	slots := generatedImageSlots(design.AgentRequest{
		Project: design.Project{
			Nodes: []design.Node{
				{ID: "background-placeholder", Type: design.NodeTypeImage, Status: "generating", LayerRole: "background-removed"},
				{ID: "agent-placeholder", Type: design.NodeTypeImage, Status: "generating"},
			},
		},
	}, nil, 1)
	if len(slots) != 1 || slots[0].ID != "agent-placeholder" {
		t.Fatalf("expected agent placeholder slot, got %#v", slots)
	}
}

func TestDesignAgentDoesNotRunFallbackAfterImageTimeout(t *testing.T) {
	imageGenerator := &failingImageGenerator{err: context.DeadlineExceeded}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, fakeCreativeAgent{})

	_, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "生成一张 still day 官网首页图片",
		Mode:       "image",
		ImageModel: "gpt-image-2",
	})
	if err == nil {
		t.Fatal("expected image timeout error")
	}
	if imageGenerator.calls != 1 {
		t.Fatalf("expected timeout to skip fallback generation, got %d calls", imageGenerator.calls)
	}
}

func TestDesignAgentCompactsPromptBeforeGPTImage(t *testing.T) {
	imageGenerator := &fakeImageGenerator{}
	agent := NewDesignAgent(fakeRunner{}, imageGenerator, longCreativeAgent{})

	_, err := agent.Plan(context.Background(), design.AgentRequest{
		Prompt:     "根据当前画布生成一张电商主视觉",
		Mode:       "image",
		ImageModel: "gpt-image-2",
	})
	if err != nil {
		t.Fatal(err)
	}
	if !imageGenerator.called {
		t.Fatal("expected image generator to be called")
	}
	if utf8.RuneCountInString(imageGenerator.prompt) > imageGenerationPromptTooLargeRunes {
		t.Fatalf("expected compact image prompt, got %d runes", utf8.RuneCountInString(imageGenerator.prompt))
	}
	if containsAll(imageGenerator.prompt, []string{"任务类型", "主体描述", "风格定义", "技术参数", "输出规格"}) {
		t.Fatalf("expected compact prompt without forced framework, got %s", imageGenerator.prompt)
	}
}

func containsAll(value string, needles []string) bool {
	for _, needle := range needles {
		if !strings.Contains(value, needle) {
			return false
		}
	}
	return true
}