feat(piagent): batch-generate images in one call when generators share a prompt

When multiple image-generator tasks carry the same brief, request all
images from a single GPT Image call with Count=N instead of firing one
call per image. Adds directImageBatchCount/Prompt detection, batch
(incremental) generation paths, and folds the planned count into the
idempotency key. Disables the streaming plan probe for multi-image
requests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

server/internal/infrastructure/piagent/design_agent.go

@@ -129,6 +129,17 @@ func (a *DesignAgent) planDirectImageGeneration(ctx context.Context, req design.
 	}
 	generationReq := req
 	generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt)
+	if count := directImageBatchCount(generationReq); count > 1 {
+		imagePrompt := directImageBatchPrompt(generationReq)
+		imageURLs, err := a.generateDirectImageBatch(ctx, generationReq, imagePrompt, count)
+		if err != nil {
+			return design.AgentPlan{}, fmt.Errorf("gpt image 2 generation failed: %w", err)
+		}
+		plan := directImagePlan(generationReq)
+		applyPlannedGeneratedImages(&plan, generationReq, imageURLs, []design.AgentImagePrompt{imagePrompt})
+		plan.Message.Title = generatedMessageTitle(imagePrompt.Title, req.Prompt)
+		return plan, nil
+	}
 	prompts := directImagePrompts(generationReq)
 	imageURLs, err := a.generatePlannedImages(ctx, prompts, generationReq)
 	if err != nil {
@@ -150,8 +161,23 @@ func (a *DesignAgent) planDirectImageGenerationIncremental(ctx context.Context,
 	}
 	generationReq := req
 	generationReq.ImageSize = resolveImageSizeForGeneration(req.ImageSize, req.Prompt)
-	prompts := directImagePrompts(generationReq)
 	plan := directImagePlan(generationReq)
+	if count := directImageBatchCount(generationReq); count > 1 {
+		imagePrompt := directImageBatchPrompt(generationReq)
+		nodes, err := a.generateDirectImageBatchIncremental(ctx, generationReq, imagePrompt, count, plan.GeneratedNodes, onResult)
+		if err != nil {
+			if len(nodes) > 0 {
+				plan.GeneratedNodes = nodes
+				plan.Connections = nil
+			}
+			return plan, fmt.Errorf("gpt image 2 generation failed: %w", err)
+		}
+		plan.GeneratedNodes = nodes
+		plan.Connections = nil
+		plan.Message.Title = generatedMessageTitle(imagePrompt.Title, req.Prompt)
+		return plan, nil
+	}
+	prompts := directImagePrompts(generationReq)
 	nodes, err := a.generatePlannedImagesIncremental(ctx, prompts, generationReq, plan.GeneratedNodes, onResult)
 	if err != nil {
 		if len(nodes) > 0 {
@@ -204,6 +230,49 @@ func directImagePrompts(req design.AgentRequest) []design.AgentImagePrompt {
 	return prompts
 }
 
+func directImageBatchCount(req design.AgentRequest) int {
+	count := plannedGenerationCount(req)
+	if count <= 1 || !directImageTasksSharePrompt(req) {
+		return 0
+	}
+	return count
+}
+
+func directImageTasksSharePrompt(req design.AgentRequest) bool {
+	common := ""
+	for _, task := range req.TaskPlan.ImageTasks {
+		brief := strings.TrimSpace(task.Brief)
+		if brief == "" {
+			continue
+		}
+		if common == "" {
+			common = brief
+			continue
+		}
+		if brief != common {
+			return false
+		}
+	}
+	return true
+}
+
+func directImageBatchPrompt(req design.AgentRequest) design.AgentImagePrompt {
+	prompt := strings.TrimSpace(req.Prompt)
+	if prompt == "" {
+		prompt = fallbackDirectImagePrompt(req)
+	}
+	title := normalizeGeneratedTitle("", req.Prompt)
+	if len(req.TaskPlan.ImageTasks) > 0 {
+		if taskTitle := strings.TrimSpace(req.TaskPlan.ImageTasks[0].Title); taskTitle != "" {
+			title = normalizeGeneratedTitle(taskTitle, req.Prompt)
+		}
+	}
+	return design.AgentImagePrompt{
+		Title:  title,
+		Prompt: prepareGPTImagePrompt(prompt, imageGenerationPromptMaxRunes),
+	}
+}
+
 func shouldUseDirectImageGeneration(req design.AgentRequest) bool {
 	mode := strings.ToLower(strings.TrimSpace(req.Mode))
 	return mode == "image" || mode == "pure-image" || mode == "image-generator" || strings.Contains(mode, "image-action")
@@ -295,12 +364,58 @@ func (a *DesignAgent) generatePlannedImages(ctx context.Context, prompts []desig
 	return urls, nil
 }
 
+func (a *DesignAgent) generateDirectImageBatch(ctx context.Context, req design.AgentRequest, prompt design.AgentImagePrompt, count int) ([]string, error) {
+	image, err := a.generateImageWithFallbackCount(ctx, req, prompt.Prompt, fallbackImagePrompt(req, 0), 0, count)
+	if err != nil {
+		return nil, err
+	}
+	urls := generatedImageURLs(image)
+	if len(urls) == 0 {
+		return nil, fmt.Errorf("image generation returned empty image")
+	}
+	if len(urls) < count {
+		return urls, fmt.Errorf("only %d of %d planned images returned", len(urls), count)
+	}
+	return urls[:count], nil
+}
+
 type plannedImageResult struct {
 	index int
 	node  design.Node
 	err   error
 }
 
+func (a *DesignAgent) generateDirectImageBatchIncremental(
+	ctx context.Context,
+	req design.AgentRequest,
+	prompt design.AgentImagePrompt,
+	count int,
+	fallback []design.Node,
+	onResult func(design.AgentImageGenerationResult) error,
+) ([]design.Node, error) {
+	imageURLs, err := a.generateDirectImageBatch(ctx, req, prompt, count)
+	if len(imageURLs) == 0 {
+		return nil, err
+	}
+	slots := generatedImageSlots(req, fallback, len(imageURLs))
+	nodes := make([]design.Node, 0, len(imageURLs))
+	var firstErr error
+	for index, imageURL := range imageURLs {
+		node := generatedImageNodeFromSlot(slots[index], req, imageURL, prompt.Title, index, count)
+		if onResult != nil {
+			if emitErr := onResult(design.AgentImageGenerationResult{Index: index, Node: node}); emitErr != nil && firstErr == nil {
+				firstErr = emitErr
+				continue
+			}
+		}
+		nodes = append(nodes, node)
+	}
+	if err != nil && firstErr == nil {
+		firstErr = err
+	}
+	return nodes, firstErr
+}
+
 func (a *DesignAgent) generatePlannedImagesIncremental(ctx context.Context, prompts []design.AgentImagePrompt, req design.AgentRequest, fallback []design.Node, onResult func(design.AgentImageGenerationResult) error) ([]design.Node, error) {
 	if len(prompts) == 0 {
 		prompts = []design.AgentImagePrompt{{Title: normalizeGeneratedTitle("", req.Prompt), Prompt: prepareGPTImagePrompt(fallbackDirectImagePrompt(req), imageGenerationFallbackMaxRunes)}}
@@ -377,6 +492,16 @@ func (a *DesignAgent) generatePlannedImagesIncremental(ctx context.Context, prom
 }
 
 func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design.AgentRequest, prompt string, fallback string, index int) (GeneratedImage, error) {
+	return a.generateImageWithFallbackCount(ctx, req, prompt, fallback, index, 1)
+}
+
+func (a *DesignAgent) generateImageWithFallbackCount(ctx context.Context, req design.AgentRequest, prompt string, fallback string, index int, count int) (GeneratedImage, error) {
+	if count < 1 {
+		count = 1
+	}
+	if count > 10 {
+		count = 10
+	}
 	referenceImages := referenceImageURLs(req)
 	prompt = withReferenceImageContext(prompt, req, len(referenceImages))
 	fallback = withReferenceImageContext(fallback, req, len(referenceImages))
@@ -391,7 +516,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design.
 	image, err := a.imageAgent.Generate(ctx, prompt, ImageRequestOptions{
 		Model:          req.ImageModel,
 		Size:           req.ImageSize,
-		Count:          1,
+		Count:          count,
 		Images:         referenceImages,
 		IdempotencyKey: imageGenerationIdempotencyKey(req, index, prompt, false),
 	})
@@ -407,7 +532,7 @@ func (a *DesignAgent) generateImageWithFallback(ctx context.Context, req design.
 	fallbackImage, fallbackErr := a.imageAgent.Generate(ctx, fallback, ImageRequestOptions{
 		Model:          req.ImageModel,
 		Size:           req.ImageSize,
-		Count:          1,
+		Count:          count,
 		Images:         referenceImages,
 		IdempotencyKey: imageGenerationIdempotencyKey(req, index, fallback, true),
 	})
@@ -428,6 +553,7 @@ func imageGenerationIdempotencyKey(req design.AgentRequest, index int, prompt st
 		strings.TrimSpace(req.Prompt),
 		strings.TrimSpace(req.ImageModel),
 		strings.TrimSpace(req.ImageSize),
+		strconv.Itoa(plannedGenerationCount(req)),
 		strconv.Itoa(index),
 		kind,
 		strings.Join(referenceImageURLs(req), "\n"),
@@ -776,7 +902,9 @@ func applyPlannedGeneratedImages(plan *design.AgentPlan, req design.AgentRequest
 	nodes := make([]design.Node, 0, requestedCount)
 	for index := 0; index < requestedCount; index++ {
 		title := ""
-		if index < len(prompts) {
+		if len(prompts) == 1 {
+			title = prompts[0].Title
+		} else if index < len(prompts) {
 			title = prompts[index].Title
 		}
 		nodes = append(nodes, generatedImageNodeFromSlot(slots[index], req, imageURLs[index], title, index, requestedCount))
@@ -806,6 +934,7 @@ func generatedImageNodeFromSlot(slot design.Node, req design.AgentRequest, image
 	}
 	if node.LayerRole == "image-generator" {
 		node.LayerRole = ""
+		node.ParentID = ""
 	}
 	if node.Tone == "" || node.Tone == "image-generator" {
 		node.Tone = "visual"

server/internal/infrastructure/piagent/design_agent_test.go

@@ -270,6 +270,64 @@ func TestDesignAgentImageGeneratorModeCallsImageModelDirectly(t *testing.T) {
 	}
 }
 
+func TestDesignAgentImageGeneratorModeBatchesSamePromptCount(t *testing.T) {
+	runner := &recordingRunner{}
+	creative := &recordingCreativeAgent{}
+	imageGenerator := &fakeImageGenerator{}
+	agent := NewDesignAgent(runner, imageGenerator, creative)
+	var emitted []design.Node
+
+	plan, err := agent.PlanIncremental(context.Background(), design.AgentRequest{
+		Project: design.Project{
+			ID: "project-image-generator",
+			Nodes: []design.Node{
+				{
+					ID:        "generator-1",
+					Type:      design.NodeTypeImage,
+					Tone:      "image-generator",
+					Status:    "generating",
+					LayerRole: "image-generator",
+					Width:     512,
+					Height:    512,
+				},
+			},
+		},
+		Prompt:     "画一只蓝色小鸟",
+		Mode:       "image-generator",
+		ImageModel: "gpt-image-2",
+		ImageSize:  "1024x1024",
+		TaskPlan: design.AgentTaskPlan{
+			Intent:     "image",
+			ImageCount: 2,
+			ImageTasks: []design.AgentImageTask{
+				{Title: "蓝色小鸟 方案 1", Brief: "画一只蓝色小鸟"},
+				{Title: "蓝色小鸟 方案 2", Brief: "画一只蓝色小鸟"},
+			},
+		},
+	}, func(result design.AgentImageGenerationResult) error {
+		emitted = append(emitted, result.Node)
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if runner.calls != 0 || creative.buildCalls != 0 {
+		t.Fatalf("expected direct image generator mode, runner=%d build=%d", runner.calls, creative.buildCalls)
+	}
+	if imageGenerator.callCount() != 1 {
+		t.Fatalf("expected one batched image request, got %d", imageGenerator.callCount())
+	}
+	if imageGenerator.opts.Count != 2 {
+		t.Fatalf("expected image request n/count 2, got %#v", imageGenerator.opts)
+	}
+	if len(emitted) != 2 || len(plan.GeneratedNodes) != 2 {
+		t.Fatalf("expected two generated nodes, emitted=%#v plan=%#v", emitted, plan.GeneratedNodes)
+	}
+	if plan.GeneratedNodes[0].ID != "generator-1" || plan.GeneratedNodes[1].ID == "generator-1" {
+		t.Fatalf("expected first node to replace generator and second to be separate, got %#v", plan.GeneratedNodes)
+	}
+}
+
 func (g *fakeImageGenerator) callCount() int {
 	g.mu.Lock()
 	defer g.mu.Unlock()
@@ -635,6 +693,7 @@ func TestGeneratedImageNodeFromImageGeneratorSlotBecomesRegularImage(t *testing.
 			Tone:      "image-generator",
 			Status:    "generating",
 			LayerRole: "image-generator",
+			ParentID:  "generator-parent",
 			Width:     320,
 			Height:    240,
 		},
@@ -647,6 +706,9 @@ func TestGeneratedImageNodeFromImageGeneratorSlotBecomesRegularImage(t *testing.
 	if node.LayerRole != "" || node.Tone != "visual" || node.Status != "success" {
 		t.Fatalf("expected regular completed image node, got %#v", node)
 	}
+	if node.ParentID != "" {
+		t.Fatalf("expected completed image generator node to clear parent id, got %#v", node)
+	}
 }
 
 func TestDesignAgentDoesNotRunFallbackAfterImageTimeout(t *testing.T) {

server/internal/infrastructure/piagent/image_generation.go

@@ -136,7 +136,7 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq
 
 	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
 	streamPlans := []bool{false}
-	if len(images) == 0 {
+	if len(images) == 0 && count == 1 {
 		streamPlans = []bool{true, false}
 	}
 

server/internal/infrastructure/piagent/image_generation_test.go

@@ -143,6 +143,33 @@ func TestImageClientSendsExplicitSize(t *testing.T) {
 	}
 }
 
+func TestImageClientSendsImageCount(t *testing.T) {
+	var requestBody map[string]any
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image-1.png"},{"url":"https://example.com/image-2.png"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestBody["n"] != float64(2) {
+		t.Fatalf("expected image request n=2, got %#v", requestBody["n"])
+	}
+	if requestBody["stream"] != nil {
+		t.Fatalf("expected multi-image request to avoid stream, got %#v", requestBody["stream"])
+	}
+	if len(image.URLs) != 2 {
+		t.Fatalf("expected two generated images, got %#v", image.URLs)
+	}
+}
+
 func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) {
 	var requestPath string
 	var requestContentType string