feat(worker): 实现集群限流与自适应负载

保留平台模型 RPM、TPM 和并发策略语义，增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载，避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标，并扩展本地 acceptance runner，覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证：go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过；48/48 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
2026-08-03 00:13:46 +08:00
parent 9a01fd4657
commit c28bf74230
52 changed files with 3700 additions and 272 deletions
@@ -38,8 +38,10 @@ const (
 )

 const (
-	geminiRequestPrefix = `{"contents":[{"role":"user","parts":[{"text":"将参考图编辑为蓝色赛博朋克风格，保留主体结构"},{"inlineData":{"mimeType":"image/png","data":"`
-	geminiRequestSuffix = `"}}]}],"generationConfig":{"responseModalities":["IMAGE"]}}`
+	geminiRequestPrefix     = `{"contents":[{"role":"user","parts":[{"text":"将多张参考图融合编辑为蓝色赛博朋克风格，保留主体结构"}`
+	geminiRequestImageStart = `,{"inlineData":{"mimeType":"image/png","data":"`
+	geminiRequestImageEnd   = `"}}`
+	geminiRequestSuffix     = `]}],"generationConfig":{"responseModalities":["IMAGE"]}}`
 )

 var reportURLPattern = regexp.MustCompile(`(?i)(?:https?|postgres(?:ql)?):\/\/[^\s"'<>]+`)
@@ -64,6 +66,8 @@ type options struct {
 	shardIndex     int
 	shardCount     int
 	executionID    string
+	requestCount   int
+	externalClient *http.Client
 }

 type report struct {
@@ -91,6 +95,7 @@ type phaseReport struct {
 	P99MS                 float64 `json:"p99Ms"`
 	DecodedOutputBytes    int64   `json:"decodedOutputBytes,omitempty"`
 	OutputSHA256          string  `json:"outputSha256,omitempty"`
+	InputImagesPerRequest int     `json:"inputImagesPerRequest,omitempty"`
 	UniqueTaskIDs         int     `json:"uniqueTaskIds,omitempty"`
 	UniqueImageCombos     int     `json:"uniqueImageCombinations,omitempty"`
 	ForcedConversionTasks int     `json:"forcedConversionTasks,omitempty"`
@@ -148,7 +153,7 @@ func main() {
 	ctx, cancel := context.WithTimeout(signalContext, opts.timeout)
 	defer cancel()
 	result := report{
-		SchemaVersion: "acceptance-load-report/v1",
+		SchemaVersion: "acceptance-load-report/v2",
 		RunID:         opts.runID, Profile: opts.profile, StartedAt: time.Now().UTC(), SecretSafe: true,
 	}
 	phaseResults, runErr := run(ctx, opts)
@@ -205,6 +210,7 @@ func parseOptions() (options, error) {
 	var shardIndex int
 	var shardCount int
 	var executionID string
+	var requestCount int
 	flag.StringVar(&profile, "profile", env("AI_GATEWAY_ACCEPTANCE_PROFILE", "simulated-all"), "simulated-all, Gemini profile, video profile, or real-canary")
 	flag.StringVar(&reportPath, "report", env("AI_GATEWAY_ACCEPTANCE_REPORT", ""), "secret-safe JSON report path")
 	flag.DurationVar(&timeout, "timeout", 45*time.Minute, "overall timeout")
@@ -214,6 +220,7 @@ func parseOptions() (options, error) {
 	flag.IntVar(&shardIndex, "shard-index", 0, "zero-based distributed load shard index")
 	flag.IntVar(&shardCount, "shard-count", 1, "distributed load shard count")
 	flag.StringVar(&executionID, "execution-id", "", "unique workload execution identifier")
+	flag.IntVar(&requestCount, "requests", 0, "override the fixed profile request count")
 	flag.Parse()
 	opts := options{
 		gateways:       splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAYS")),
@@ -235,6 +242,7 @@ func parseOptions() (options, error) {
 		shardIndex:     shardIndex,
 		shardCount:     shardCount,
 		executionID:    strings.ToLower(strings.TrimSpace(executionID)),
+		requestCount:   requestCount,
 	}
 	if opts.executionID == "" {
 		opts.executionID = opts.profile
@@ -245,8 +253,8 @@ func parseOptions() (options, error) {
 	if opts.shardCount < 1 || opts.shardCount > 32 || opts.shardIndex < 0 || opts.shardIndex >= opts.shardCount {
 		return options{}, fmt.Errorf("invalid load shard %d/%d", opts.shardIndex, opts.shardCount)
 	}
-	if len(opts.gateways) < 1 || len(opts.gateways) > 2 || (opts.shardCount == 1 && len(opts.gateways) != 2) {
-		return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain two API base URLs, or one URL for a distributed shard")
+	if len(opts.gateways) < 1 || len(opts.gateways) > 2 {
+		return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain one or two API base URLs")
 	}
 	for index := range opts.gateways {
 		opts.gateways[index] = strings.TrimRight(opts.gateways[index], "/")
@@ -267,8 +275,11 @@ func parseOptions() (options, error) {
 	if opts.timeout <= 0 {
 		return options{}, errors.New("timeout must be positive")
 	}
+	if opts.requestCount < 0 || opts.requestCount > 10000 {
+		return options{}, errors.New("requests must be between 0 and 10000")
+	}
 	switch opts.profile {
-	case "simulated-smoke", "simulated-all", "gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery",
+	case "simulated-smoke", "simulated-all", "gemini-baseline", "gemini-multi-image", "gemini-large", "gemini-peak", "video-throughput", "video-recovery",
 		"mixed-soak", "mixed-overload":
 		if opts.emulatorURL == "" {
 			return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL is required for simulated profiles")
@@ -322,28 +333,44 @@ func run(ctx context.Context, opts options) ([]phaseResult, error) {
 			TLSClientConfig:     tlsConfig,
 		},
 	}
+	opts.externalClient = &http.Client{
+		Timeout: opts.timeout,
+		Transport: &http.Transport{
+			MaxIdleConns:        512,
+			MaxIdleConnsPerHost: 256,
+			ForceAttemptHTTP2:   true,
+		},
+	}
 	results := make([]phaseResult, 0, 5)
 	runPhase := func(name string) error {
 		var result phaseResult
 		if profile, ok := acceptanceworkload.GeminiProfileByName(name); ok {
+			requests := profile.Requests
+			if opts.requestCount > 0 {
+				requests = opts.requestCount
+			}
 			result = runGemini(
-				ctx, client, opts, name, opts.shardRequestCount(profile.Requests),
-				profile.InputBytes, profile.OutputBytes, false, opts.shardIndex, opts.shardCount,
+				ctx, client, opts, name, opts.shardRequestCount(requests),
+				profile.InputImages, profile.InputBytes, profile.OutputBytes, false, opts.shardIndex, opts.shardCount,
 			)
 		} else {
 			switch name {
 			case "smoke-gemini":
-				result = runGemini(ctx, client, opts, name, opts.shardRequestCount(4), 256<<10, 256<<10, false, opts.shardIndex, opts.shardCount)
+				result = runGemini(ctx, client, opts, name, opts.shardRequestCount(4), 1, 256<<10, 256<<10, false, opts.shardIndex, opts.shardCount)
 			case "smoke-video":
 				result = runVideo(ctx, client, opts, name, opts.shardRequestCount(10), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
 			case "video-throughput":
-				result = runVideo(ctx, client, opts, name, opts.shardRequestCount(1200), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
+				requests := 1200
+				if opts.requestCount > 0 {
+					requests = opts.requestCount
+				}
+				result = runVideo(ctx, client, opts, name, opts.shardRequestCount(requests), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
 			case "video-recovery":
 				result = runVideo(ctx, client, opts, name, opts.shardRequestCount(96), true, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount)
 			case "mixed-soak", "mixed-overload":
 				result = runMixed(ctx, client, opts, name == "mixed-overload")
 			case "real-gemini-canary":
-				result = runGemini(ctx, client, opts, name, 1, 256<<10, 0, true, 0, 1)
+				result = runGemini(ctx, client, opts, name, 1, 1, 256<<10, 0, true, 0, 1)
 			case "real-video-canary":
 				result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true, 0, 1)
 			default:
@@ -375,6 +402,7 @@ func runGemini(
 	opts options,
 	name string,
 	requestCount int,
+	inputImages int,
 	inputBytes int,
 	expectedOutputBytes int,
 	realUpstream bool,
@@ -412,8 +440,14 @@ func runGemini(
 			defer func() { <-slots }()
 			requestStarted := time.Now()
 			endpoint := opts.gateways[logicalIndex%len(opts.gateways)] + "/v1beta/models/" + url.PathEscape(opts.geminiModel) + ":generateContent"
-			input := paddedPNGVariant(inputBytes, geminiInputVariant(opts.runID+"/"+opts.executionID, logicalIndex))
-			req, err := newGeminiRequest(ctx, endpoint, input)
+			var req *http.Request
+			var err error
+			if name == acceptanceworkload.GeminiMultiImage.Name {
+				req, err = newGeminiFileDataRequest(ctx, endpoint, geminiFixtureURLs(opts, inputImages, logicalIndex))
+			} else {
+				inputs := geminiInputs(inputImages, inputBytes, opts.runID+"/"+opts.executionID, logicalIndex)
+				req, err = newGeminiRequest(ctx, endpoint, inputs)
+			}
 			if err == nil {
 				opts.setHeaders(req, logicalIndex, realUpstream)
 				req.Header.Set("Content-Type", "application/json")
@@ -462,6 +496,7 @@ func runGemini(
 	result := phaseReport{
 		Name: name, Requests: requestCount, Completed: completed, Failed: failures,
 		DurationMS: elapsed.Milliseconds(), DecodedOutputBytes: decodedBytes, OutputSHA256: expectedHash,
+		InputImagesPerRequest: inputImages,
 	}
 	if firstErr == nil && name == "gemini-baseline" && elapsed > 8*time.Minute {
 		firstErr = fmt.Errorf("Gemini baseline exceeded 8 minutes: %s", elapsed)
@@ -472,7 +507,7 @@ func runGemini(
 	return phaseResult{report: result, latencies: latencies, err: firstErr}
 }

-func streamGeminiRequestBody(input []byte) io.ReadCloser {
+func streamGeminiRequestBody(inputs [][]byte) io.ReadCloser {
 	reader, writer := io.Pipe()
 	go func() {
 		closeWithError := func(err error) {
@@ -482,14 +517,24 @@ func streamGeminiRequestBody(input []byte) io.ReadCloser {
 			closeWithError(err)
 			return
 		}
-		encoder := base64.NewEncoder(base64.StdEncoding, writer)
-		if _, err := encoder.Write(input); err != nil {
-			closeWithError(err)
-			return
-		}
-		if err := encoder.Close(); err != nil {
-			closeWithError(err)
-			return
+		for _, input := range inputs {
+			if _, err := io.WriteString(writer, geminiRequestImageStart); err != nil {
+				closeWithError(err)
+				return
+			}
+			encoder := base64.NewEncoder(base64.StdEncoding, writer)
+			if _, err := encoder.Write(input); err != nil {
+				closeWithError(err)
+				return
+			}
+			if err := encoder.Close(); err != nil {
+				closeWithError(err)
+				return
+			}
+			if _, err := io.WriteString(writer, geminiRequestImageEnd); err != nil {
+				closeWithError(err)
+				return
+			}
 		}
 		if _, err := io.WriteString(writer, geminiRequestSuffix); err != nil {
 			closeWithError(err)
@@ -500,9 +545,9 @@ func streamGeminiRequestBody(input []byte) io.ReadCloser {
 	return reader
 }

-func newGeminiRequest(ctx context.Context, endpoint string, input []byte) (*http.Request, error) {
+func newGeminiRequest(ctx context.Context, endpoint string, inputs [][]byte) (*http.Request, error) {
 	getBody := func() (io.ReadCloser, error) {
-		return streamGeminiRequestBody(input), nil
+		return streamGeminiRequestBody(inputs), nil
 	}
 	body, _ := getBody()
 	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, body)
@@ -511,10 +556,57 @@ func newGeminiRequest(ctx context.Context, endpoint string, input []byte) (*http
 		return nil, err
 	}
 	request.GetBody = getBody
-	request.ContentLength = int64(len(geminiRequestPrefix) + base64.StdEncoding.EncodedLen(len(input)) + len(geminiRequestSuffix))
+	contentLength := len(geminiRequestPrefix) + len(geminiRequestSuffix)
+	for _, input := range inputs {
+		contentLength += len(geminiRequestImageStart) + base64.StdEncoding.EncodedLen(len(input)) + len(geminiRequestImageEnd)
+	}
+	request.ContentLength = int64(contentLength)
 	return request, nil
 }

+func newGeminiFileDataRequest(ctx context.Context, endpoint string, imageURLs []string) (*http.Request, error) {
+	parts := make([]any, 0, len(imageURLs)+1)
+	parts = append(parts, map[string]any{"text": "将多张参考图融合编辑为蓝色赛博朋克风格，保留主体结构"})
+	for _, imageURL := range imageURLs {
+		parts = append(parts, map[string]any{"fileData": map[string]any{
+			"mimeType": geminiFixtureMIMEType(imageURL),
+			"fileUri":  imageURL,
+		}})
+	}
+	body, err := json.Marshal(map[string]any{
+		"contents":         []any{map[string]any{"role": "user", "parts": parts}},
+		"generationConfig": map[string]any{"responseModalities": []string{"IMAGE"}},
+	})
+	if err != nil {
+		return nil, err
+	}
+	return http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
+}
+
+func geminiFixtureMIMEType(imageURL string) string {
+	lower := strings.ToLower(imageURL)
+	switch {
+	case strings.Contains(lower, ".webp"):
+		return "image/webp"
+	case strings.Contains(lower, ".jpg"), strings.Contains(lower, ".jpeg"):
+		return "image/jpeg"
+	default:
+		return "image/png"
+	}
+}
+
+func geminiFixtureURLs(opts options, count int, logicalIndex int) []string {
+	fixtures := opts.emulatorFixtureURLs()
+	if count > len(fixtures) {
+		count = len(fixtures)
+	}
+	urls := make([]string, 0, count)
+	for imageIndex := 0; imageIndex < count; imageIndex++ {
+		urls = append(urls, fixtures[(logicalIndex*count+imageIndex)%len(fixtures)])
+	}
+	return urls
+}
+
 func runVideo(
 	ctx context.Context,
 	client *http.Client,
@@ -629,7 +721,7 @@ func runVideo(
 	}
 	wg.Wait()
 	submissionDuration := time.Since(submitStartedAt)
-	if firstErr == nil && name == "video-throughput" && submissionDuration > 10*time.Second {
+	if firstErr == nil && name == "video-throughput" && requestCount == 1200 && submissionDuration > 10*time.Second {
 		firstErr = fmt.Errorf("1200 video submissions exceeded 10 seconds: %s", submissionDuration)
 	}
 	uniqueTaskIDs := map[string]struct{}{}
@@ -672,18 +764,25 @@ func runVideo(
 		}()
 	}
 	wg.Wait()
-	comboSet := map[string]struct{}{}
-	for _, combo := range combinations {
-		comboSet[strings.Join(combo, "\x00")] = struct{}{}
-	}
-	if firstErr == nil && !realUpstream && len(comboSet) != 128 {
-		firstErr = fmt.Errorf("video image combinations=%d, want 128", len(comboSet))
+	uniqueImageCombinations := requestedVideoCombinationCount(
+		combinations,
+		requestCount,
+		requestOffset,
+		requestStride,
+	)
+	expectedCombinations := min(requestCount, 128)
+	if firstErr == nil && !realUpstream && uniqueImageCombinations < expectedCombinations {
+		firstErr = fmt.Errorf(
+			"video image combinations=%d, want at least %d",
+			uniqueImageCombinations,
+			expectedCombinations,
+		)
 	}
 	elapsed := time.Since(startedAt)
 	result := phaseReport{
 		Name: name, Requests: requestCount, Completed: completed, Failed: failures,
 		DurationMS: elapsed.Milliseconds(), SubmissionDurationMS: submissionDuration.Milliseconds(),
-		UniqueTaskIDs: len(uniqueTaskIDs), UniqueImageCombos: len(comboSet),
+		UniqueTaskIDs: len(uniqueTaskIDs), UniqueImageCombos: uniqueImageCombinations,
 		ForcedConversionTasks: forcedConversions,
 	}
 	if firstErr == nil && completed != requestCount {
@@ -755,7 +854,7 @@ func runMixed(
 				defer func() { <-slots }()
 				var result phaseResult
 				if isImage {
-					result = runGemini(runCtx, client, opts, "mixed-image", 1, 256<<10, 256<<10, false, requestIndex, 1)
+					result = runGemini(runCtx, client, opts, "mixed-image", 1, 1, 256<<10, 256<<10, false, requestIndex, 1)
 				} else {
 					result = runVideo(
 						runCtx,
@@ -931,7 +1030,14 @@ func validateVideoAsset(ctx context.Context, client *http.Client, opts options,
 	if gatewayRequest {
 		opts.setHeaders(request, index, false)
 	}
-	response, err := client.Do(request)
+	mediaClient := client
+	if !gatewayRequest && opts.gatewayTLSName != "" {
+		mediaClient = opts.externalClient
+		if mediaClient == nil {
+			mediaClient = &http.Client{Timeout: client.Timeout}
+		}
+	}
+	response, err := mediaClient.Do(request)
 	if err != nil {
 		return fmt.Errorf("download final video: %w", err)
 	}
@@ -1056,6 +1162,31 @@ func videoCombinations(images []string, count int) [][]string {
 	return out
 }

+func requestedVideoCombinationCount(
+	combinations [][]string,
+	requestCount int,
+	requestOffset int,
+	requestStride int,
+) int {
+	if len(combinations) == 0 || requestCount <= 0 {
+		return 0
+	}
+	if requestStride < 1 {
+		requestStride = 1
+	}
+	seen := map[string]struct{}{}
+	for index := 0; index < requestCount; index++ {
+		logicalIndex := requestOffset + index*requestStride
+		imageCount := acceptanceworkload.VideoImageCount(logicalIndex)
+		combo := combinations[logicalIndex%len(combinations)]
+		if len(combo) > imageCount {
+			combo = combo[:imageCount]
+		}
+		seen[strings.Join(combo, "\x00")] = struct{}{}
+	}
+	return len(seen)
+}
+
 func streamGeminiImageHash(reader io.Reader) (int64, string, error) {
 	buffered := bufio.NewReaderSize(reader, 64<<10)
 	if err := scanUntil(buffered, []byte(`"inlineData"`), 256<<20); err != nil {
@@ -1176,6 +1307,24 @@ func paddedPNG(size int) []byte {
 	return paddedPNGVariant(size, 0)
 }

+func geminiInputs(count int, totalBytes int, executionID string, logicalIndex int) [][]byte {
+	if count < 1 {
+		count = 1
+	}
+	inputs := make([][]byte, 0, count)
+	baseSize := totalBytes / count
+	remainder := totalBytes % count
+	for imageIndex := 0; imageIndex < count; imageIndex++ {
+		size := baseSize
+		if imageIndex < remainder {
+			size++
+		}
+		variant := geminiInputVariant(executionID, logicalIndex*count+imageIndex)
+		inputs = append(inputs, paddedPNGVariant(size, variant))
+	}
+	return inputs
+}
+
 func geminiInputVariant(runID string, logicalIndex int) int {
 	digest := sha256.Sum256([]byte(fmt.Sprintf("%s:%d", strings.TrimSpace(runID), logicalIndex)))
 	return int(binary.BigEndian.Uint64(digest[:8]) & uint64(^uint(0)>>1))
@@ -17,6 +17,12 @@ import (
 	"testing"
 )

+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
+	return f(request)
+}
+
 func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
 	payload := paddedPNG(4 << 20)
 	response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
@@ -44,8 +50,8 @@ func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
 	}
 }

-func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
-	input := paddedPNGVariant(2<<20, 37)
+func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
+	inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
 	var body struct {
 		Contents []struct {
 			Parts []struct {
@@ -59,24 +65,26 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
 			ResponseModalities []string `json:"responseModalities"`
 		} `json:"generationConfig"`
 	}
-	stream := streamGeminiRequestBody(input)
+	stream := streamGeminiRequestBody(inputs)
 	defer stream.Close()
 	if err := json.NewDecoder(stream).Decode(&body); err != nil {
 		t.Fatalf("decode streamed Gemini request: %v", err)
 	}
-	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 2 || body.Contents[0].Parts[1].InlineData == nil {
+	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
 		t.Fatalf("unexpected Gemini body structure: %+v", body)
 	}
-	inlineData := body.Contents[0].Parts[1].InlineData
-	if inlineData.MIMEType != "image/png" {
-		t.Fatalf("mime type=%q", inlineData.MIMEType)
-	}
-	decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
-	if err != nil {
-		t.Fatalf("decode streamed input: %v", err)
-	}
-	if !bytes.Equal(decoded, input) {
-		t.Fatal("streamed input differs from source")
+	for index, input := range inputs {
+		inlineData := body.Contents[0].Parts[index+1].InlineData
+		if inlineData == nil || inlineData.MIMEType != "image/png" {
+			t.Fatalf("image %d inline data=%+v", index, inlineData)
+		}
+		decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
+		if err != nil {
+			t.Fatalf("decode streamed input %d: %v", index, err)
+		}
+		if !bytes.Equal(decoded, input) {
+			t.Fatalf("streamed input %d differs from source", index)
+		}
 	}
 	if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
 		t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
@@ -84,8 +92,8 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
 }

 func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
-	input := paddedPNGVariant(256<<10, 91)
-	request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", input)
+	inputs := geminiInputs(3, 768<<10, "replay-test", 91)
+	request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
 	if err != nil {
 		t.Fatalf("new Gemini request: %v", err)
 	}
@@ -114,6 +122,37 @@ func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
 	}
 }

+func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
+	request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
+		"https://fixtures.example/one.png",
+		"https://fixtures.example/two.png",
+		"https://fixtures.example/three.png",
+	})
+	if err != nil {
+		t.Fatalf("new multi-image request: %v", err)
+	}
+	var body struct {
+		Contents []struct {
+			Parts []struct {
+				FileData *struct {
+					FileURI string `json:"fileUri"`
+				} `json:"fileData"`
+			} `json:"parts"`
+		} `json:"contents"`
+	}
+	if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
+		t.Fatalf("decode multi-image request: %v", err)
+	}
+	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
+		t.Fatalf("unexpected multi-image body: %+v", body)
+	}
+	for index := 1; index < 4; index++ {
+		if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
+			t.Fatalf("missing fileData at part %d", index)
+		}
+	}
+}
+
 func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
 	images := make([]string, 16)
 	for index := range images {
@@ -133,6 +172,9 @@ func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
 	if len(seen) != 128 {
 		t.Fatalf("unique combinations=%d", len(seen))
 	}
+	if got := requestedVideoCombinationCount(combinations, 144, 0, 1); got != 131 {
+		t.Fatalf("requested combinations=%d, want 131", got)
+	}
 }

 func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
@@ -178,7 +220,7 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
 		gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
 		runToken: "token-1", geminiModel: "gemini-image-test",
 	}
-	result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 256<<10, 256<<10, false, 0, 1)
+	result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1)
 	if result.err != nil || result.report.Completed != 8 {
 		t.Fatalf("result=%+v err=%v", result.report, result.err)
 	}
@@ -254,6 +296,21 @@ func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
 	}
 }

+func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "video/mp4")
+		_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
+	}))
+	defer server.Close()
+	poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+		return nil, errors.New("gateway-only transport must not be used for external media")
+	})}
+	opts := options{gatewayTLSName: "gateway.easyai.local"}
+	if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
+		t.Fatalf("validate external video with independent transport: %v", err)
+	}
+}
+
 func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
 		if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {