引入动态流量门禁、隔离验收身份与协议级 Gemini/Volces 模拟器,覆盖双站点 API、Worker、PostgreSQL、River、账务、回调和媒体物化链路。 新增 P24/P28/P32 容量阶梯、Worker 强杀恢复、真实小流量 canary、CAS 放量和失败保持 validation 的生产编排;Worker 执行槽、连接池、媒体并发和双站点副本数改为环境配置。 验证:Go 全量测试、真实 PostgreSQL 迁移集成测试、迁移安全检查、OpenAPI 生成、ShellCheck、Kustomize、gofmt 和 git diff --check。
826 lines
25 KiB
Go
826 lines
25 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"crypto/tls"
|
|
"encoding/base64"
|
|
"encoding/binary"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"hash/crc32"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
|
|
)
|
|
|
|
const (
|
|
runHeader = "X-EasyAI-Acceptance-Run"
|
|
tokenHeader = "X-EasyAI-Acceptance-Token"
|
|
upstreamHeader = "X-EasyAI-Acceptance-Upstream"
|
|
)
|
|
|
|
var reportURLPattern = regexp.MustCompile(`(?i)(?:https?|postgres(?:ql)?):\/\/[^\s"'<>]+`)
|
|
|
|
type options struct {
|
|
gateways []string
|
|
gatewayTLSName string
|
|
emulatorURL string
|
|
apiKey string
|
|
runID string
|
|
runToken string
|
|
geminiModel string
|
|
videoModel string
|
|
profile string
|
|
reportPath string
|
|
realImageURLs []string
|
|
timeout time.Duration
|
|
}
|
|
|
|
type report struct {
|
|
RunID string `json:"runId"`
|
|
Profile string `json:"profile"`
|
|
StartedAt time.Time `json:"startedAt"`
|
|
FinishedAt time.Time `json:"finishedAt"`
|
|
Passed bool `json:"passed"`
|
|
Phases []phaseReport `json:"phases"`
|
|
Failure string `json:"failure,omitempty"`
|
|
SecretSafe bool `json:"secretSafe"`
|
|
}
|
|
|
|
type phaseReport struct {
|
|
Name string `json:"name"`
|
|
Requests int `json:"requests"`
|
|
Completed int `json:"completed"`
|
|
Failed int `json:"failed"`
|
|
DurationMS int64 `json:"durationMs"`
|
|
SubmissionDurationMS int64 `json:"submissionDurationMs,omitempty"`
|
|
P50MS float64 `json:"p50Ms"`
|
|
P95MS float64 `json:"p95Ms"`
|
|
P99MS float64 `json:"p99Ms"`
|
|
DecodedOutputBytes int64 `json:"decodedOutputBytes,omitempty"`
|
|
OutputSHA256 string `json:"outputSha256,omitempty"`
|
|
UniqueTaskIDs int `json:"uniqueTaskIds,omitempty"`
|
|
UniqueImageCombos int `json:"uniqueImageCombinations,omitempty"`
|
|
ForcedConversionTasks int `json:"forcedConversionTasks,omitempty"`
|
|
}
|
|
|
|
type phaseResult struct {
|
|
report phaseReport
|
|
latencies []time.Duration
|
|
err error
|
|
}
|
|
|
|
func main() {
|
|
opts, err := parseOptions()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "acceptance load configuration error:", err)
|
|
os.Exit(2)
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), opts.timeout)
|
|
defer cancel()
|
|
result := report{
|
|
RunID: opts.runID, Profile: opts.profile, StartedAt: time.Now().UTC(), SecretSafe: true,
|
|
}
|
|
phaseResults, runErr := run(ctx, opts)
|
|
for _, phase := range phaseResults {
|
|
phase.report.P50MS = percentileMilliseconds(phase.latencies, 0.50)
|
|
phase.report.P95MS = percentileMilliseconds(phase.latencies, 0.95)
|
|
phase.report.P99MS = percentileMilliseconds(phase.latencies, 0.99)
|
|
result.Phases = append(result.Phases, phase.report)
|
|
if runErr == nil && phase.err != nil {
|
|
runErr = phase.err
|
|
}
|
|
}
|
|
result.FinishedAt = time.Now().UTC()
|
|
result.Passed = runErr == nil
|
|
if runErr != nil {
|
|
result.Failure = redactError(runErr.Error(), opts)
|
|
}
|
|
payload, _ := json.MarshalIndent(result, "", " ")
|
|
fmt.Println(string(payload))
|
|
if opts.reportPath != "" {
|
|
if err := os.WriteFile(opts.reportPath, append(payload, '\n'), 0o600); err != nil {
|
|
fmt.Fprintln(os.Stderr, "write acceptance report failed:", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
if runErr != nil {
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func parseOptions() (options, error) {
|
|
var profile string
|
|
var reportPath string
|
|
var timeout time.Duration
|
|
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")
|
|
flag.Parse()
|
|
opts := options{
|
|
gateways: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAYS")),
|
|
gatewayTLSName: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME")),
|
|
emulatorURL: strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL")), "/"),
|
|
apiKey: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_API_KEY")),
|
|
runID: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_RUN_ID")),
|
|
runToken: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_RUN_TOKEN")),
|
|
geminiModel: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_GEMINI_MODEL")),
|
|
videoModel: strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCEPTANCE_VIDEO_MODEL")),
|
|
profile: strings.ToLower(strings.TrimSpace(profile)),
|
|
reportPath: strings.TrimSpace(reportPath),
|
|
realImageURLs: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS")),
|
|
timeout: timeout,
|
|
}
|
|
if len(opts.gateways) != 2 {
|
|
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain exactly two API base URLs")
|
|
}
|
|
for index := range opts.gateways {
|
|
opts.gateways[index] = strings.TrimRight(opts.gateways[index], "/")
|
|
if _, err := url.ParseRequestURI(opts.gateways[index]); err != nil {
|
|
return options{}, fmt.Errorf("gateway %d URL is invalid", index+1)
|
|
}
|
|
}
|
|
if opts.gatewayTLSName != "" &&
|
|
(strings.ContainsAny(opts.gatewayTLSName, "/:@") || !strings.Contains(opts.gatewayTLSName, ".")) {
|
|
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME must be a DNS hostname")
|
|
}
|
|
if opts.apiKey == "" || opts.runID == "" || opts.runToken == "" {
|
|
return options{}, errors.New("acceptance API key, Run ID, and Run Token environment variables are required")
|
|
}
|
|
if opts.geminiModel == "" || opts.videoModel == "" {
|
|
return options{}, errors.New("Gemini and video model environment variables are required")
|
|
}
|
|
if opts.timeout <= 0 {
|
|
return options{}, errors.New("timeout must be positive")
|
|
}
|
|
switch opts.profile {
|
|
case "simulated-all", "gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery":
|
|
if opts.emulatorURL == "" {
|
|
return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL is required for simulated profiles")
|
|
}
|
|
case "real-canary":
|
|
if len(opts.realImageURLs) < 3 {
|
|
return options{}, errors.New("real-canary requires at least three public AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS")
|
|
}
|
|
default:
|
|
return options{}, fmt.Errorf("unsupported profile %q", opts.profile)
|
|
}
|
|
return opts, nil
|
|
}
|
|
|
|
func run(ctx context.Context, opts options) ([]phaseResult, error) {
|
|
var tlsConfig *tls.Config
|
|
if opts.gatewayTLSName != "" {
|
|
tlsConfig = &tls.Config{
|
|
MinVersion: tls.VersionTLS12,
|
|
ServerName: opts.gatewayTLSName,
|
|
}
|
|
}
|
|
client := &http.Client{
|
|
Timeout: opts.timeout,
|
|
Transport: &http.Transport{
|
|
MaxIdleConns: 2048,
|
|
MaxIdleConnsPerHost: 1024,
|
|
MaxConnsPerHost: 1024,
|
|
TLSClientConfig: tlsConfig,
|
|
},
|
|
}
|
|
results := make([]phaseResult, 0, 5)
|
|
runPhase := func(name string) error {
|
|
var result phaseResult
|
|
if profile, ok := acceptanceworkload.GeminiProfileByName(name); ok {
|
|
result = runGemini(ctx, client, opts, name, profile.Requests, profile.InputBytes, profile.OutputBytes, false)
|
|
} else {
|
|
switch name {
|
|
case "video-throughput":
|
|
result = runVideo(ctx, client, opts, name, 1200, false, opts.emulatorFixtureURLs(), false)
|
|
case "video-recovery":
|
|
result = runVideo(ctx, client, opts, name, 96, true, opts.emulatorFixtureURLs(), false)
|
|
case "real-gemini-canary":
|
|
result = runGemini(ctx, client, opts, name, 1, 256<<10, 0, true)
|
|
case "real-video-canary":
|
|
result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true)
|
|
default:
|
|
return fmt.Errorf("unknown phase %q", name)
|
|
}
|
|
}
|
|
results = append(results, result)
|
|
return result.err
|
|
}
|
|
phases := []string{opts.profile}
|
|
if opts.profile == "simulated-all" {
|
|
phases = []string{"gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery"}
|
|
} else if opts.profile == "real-canary" {
|
|
phases = []string{"real-gemini-canary", "real-video-canary"}
|
|
}
|
|
for _, phase := range phases {
|
|
if err := runPhase(phase); err != nil {
|
|
return results, err
|
|
}
|
|
}
|
|
return results, nil
|
|
}
|
|
|
|
func runGemini(
|
|
ctx context.Context,
|
|
client *http.Client,
|
|
opts options,
|
|
name string,
|
|
requestCount int,
|
|
inputBytes int,
|
|
expectedOutputBytes int,
|
|
realUpstream bool,
|
|
) phaseResult {
|
|
startedAt := time.Now()
|
|
input := paddedPNG(inputBytes)
|
|
requestBody, _ := json.Marshal(map[string]any{
|
|
"contents": []any{map[string]any{
|
|
"role": "user",
|
|
"parts": []any{
|
|
map[string]any{"text": "将参考图编辑为蓝色赛博朋克风格,保留主体结构"},
|
|
map[string]any{"inlineData": map[string]any{
|
|
"mimeType": "image/png",
|
|
"data": base64.StdEncoding.EncodeToString(input),
|
|
}},
|
|
},
|
|
}},
|
|
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
|
|
})
|
|
expectedHash := ""
|
|
if expectedOutputBytes > 0 {
|
|
sum := sha256.Sum256(paddedPNG(expectedOutputBytes))
|
|
expectedHash = hex.EncodeToString(sum[:])
|
|
}
|
|
latencies := make([]time.Duration, requestCount)
|
|
var decodedBytes int64
|
|
var completed int
|
|
var failures int
|
|
var firstErr error
|
|
var mu sync.Mutex
|
|
slots := make(chan struct{}, requestCount)
|
|
var wg sync.WaitGroup
|
|
for index := 0; index < requestCount; index++ {
|
|
index := index
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
select {
|
|
case slots <- struct{}{}:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
defer func() { <-slots }()
|
|
requestStarted := time.Now()
|
|
endpoint := opts.gateways[index%len(opts.gateways)] + "/v1beta/models/" + url.PathEscape(opts.geminiModel) + ":generateContent"
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(requestBody))
|
|
if err == nil {
|
|
opts.setHeaders(req, realUpstream)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
var response *http.Response
|
|
response, err = client.Do(req)
|
|
if err == nil {
|
|
if response.StatusCode != http.StatusOK {
|
|
err = responseStatusError(response)
|
|
} else {
|
|
var size int64
|
|
var outputHash string
|
|
size, outputHash, err = streamGeminiImageHash(response.Body)
|
|
_ = response.Body.Close()
|
|
if err == nil && expectedOutputBytes > 0 && (size != int64(expectedOutputBytes) || outputHash != expectedHash) {
|
|
err = fmt.Errorf("Gemini output mismatch: bytes=%d hash=%s", size, outputHash)
|
|
}
|
|
if err == nil {
|
|
mu.Lock()
|
|
decodedBytes += size
|
|
if expectedHash == "" {
|
|
expectedHash = outputHash
|
|
}
|
|
mu.Unlock()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
latency := time.Since(requestStarted)
|
|
mu.Lock()
|
|
latencies[index] = latency
|
|
if err != nil {
|
|
failures++
|
|
if firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
} else {
|
|
completed++
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
elapsed := time.Since(startedAt)
|
|
result := phaseReport{
|
|
Name: name, Requests: requestCount, Completed: completed, Failed: failures,
|
|
DurationMS: elapsed.Milliseconds(), DecodedOutputBytes: decodedBytes, OutputSHA256: expectedHash,
|
|
}
|
|
if firstErr == nil && name == "gemini-baseline" && elapsed > 8*time.Minute {
|
|
firstErr = fmt.Errorf("Gemini baseline exceeded 8 minutes: %s", elapsed)
|
|
}
|
|
if firstErr == nil && completed != requestCount {
|
|
firstErr = fmt.Errorf("Gemini completed %d of %d requests", completed, requestCount)
|
|
}
|
|
return phaseResult{report: result, latencies: latencies, err: firstErr}
|
|
}
|
|
|
|
func runVideo(
|
|
ctx context.Context,
|
|
client *http.Client,
|
|
opts options,
|
|
name string,
|
|
requestCount int,
|
|
longRun bool,
|
|
imageURLs []string,
|
|
realUpstream bool,
|
|
) phaseResult {
|
|
startedAt := time.Now()
|
|
combinationCount := 128
|
|
if realUpstream {
|
|
combinationCount = 1
|
|
}
|
|
combinations := videoCombinations(imageURLs, combinationCount)
|
|
var taskIDs = make([]string, requestCount)
|
|
latencies := make([]time.Duration, requestCount)
|
|
startedTasks := make([]time.Time, requestCount)
|
|
var completed int
|
|
var failures int
|
|
var forcedConversions int
|
|
var firstErr error
|
|
var mu sync.Mutex
|
|
var wg sync.WaitGroup
|
|
submitStartedAt := time.Now()
|
|
for index := 0; index < requestCount; index++ {
|
|
index := index
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
requestStarted := time.Now()
|
|
startedTasks[index] = requestStarted
|
|
imageCount := acceptanceworkload.VideoImageCount(index)
|
|
combo := append([]string(nil), combinations[index%len(combinations)]...)
|
|
if len(combo) > imageCount {
|
|
combo = combo[:imageCount]
|
|
}
|
|
if !realUpstream && index%4 == 0 {
|
|
combo[0] = imageURLs[12+(index/4)%4]
|
|
}
|
|
content := make([]any, 0, len(combo)+1)
|
|
prompt := "多参考图生成连续运镜视频,保持人物、服装和场景一致"
|
|
if longRun {
|
|
prompt += " acceptance-long-recovery"
|
|
}
|
|
if !realUpstream && index%4 == 0 {
|
|
prompt += " acceptance-force-conversion"
|
|
}
|
|
content = append(content, map[string]any{"type": "text", "text": prompt})
|
|
for imageIndex, imageURL := range combo {
|
|
role := "reference_image"
|
|
if index%5 == 0 && imageIndex == 0 {
|
|
role = "first_frame"
|
|
}
|
|
if index%5 == 0 && len(combo) > 3 && imageIndex == len(combo)-1 {
|
|
role = "last_frame"
|
|
}
|
|
content = append(content, map[string]any{
|
|
"type": "image_url", "role": role,
|
|
"image_url": map[string]any{"url": imageURL},
|
|
})
|
|
}
|
|
body, _ := json.Marshal(map[string]any{
|
|
"model": opts.videoModel, "modelType": "omni_video", "content": content, "seed": index + 1,
|
|
"duration": 5, "ratio": "16:9", "resolution": "720p",
|
|
})
|
|
endpoint := opts.gateways[index%len(opts.gateways)] + "/api/v1/videos/generations"
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
|
if err == nil {
|
|
opts.setHeaders(req, realUpstream)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-Async", "true")
|
|
var response *http.Response
|
|
response, err = client.Do(req)
|
|
if err == nil {
|
|
if response.StatusCode != http.StatusAccepted {
|
|
err = responseStatusError(response)
|
|
} else {
|
|
var accepted struct {
|
|
TaskID string `json:"taskId"`
|
|
}
|
|
err = json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&accepted)
|
|
_ = response.Body.Close()
|
|
taskIDs[index] = strings.TrimSpace(accepted.TaskID)
|
|
if err == nil && taskIDs[index] == "" {
|
|
err = errors.New("video acceptance response did not include taskId")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
mu.Lock()
|
|
latencies[index] = time.Since(requestStarted)
|
|
if !realUpstream && index%4 == 0 {
|
|
forcedConversions++
|
|
}
|
|
if err != nil {
|
|
failures++
|
|
if firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
submissionDuration := time.Since(submitStartedAt)
|
|
if firstErr == nil && name == "video-throughput" && submissionDuration > 10*time.Second {
|
|
firstErr = fmt.Errorf("1200 video submissions exceeded 10 seconds: %s", submissionDuration)
|
|
}
|
|
uniqueTaskIDs := map[string]struct{}{}
|
|
for _, taskID := range taskIDs {
|
|
if taskID != "" {
|
|
uniqueTaskIDs[taskID] = struct{}{}
|
|
}
|
|
}
|
|
if firstErr == nil && len(uniqueTaskIDs) != requestCount {
|
|
firstErr = fmt.Errorf("video task IDs are not unique: unique=%d requests=%d", len(uniqueTaskIDs), requestCount)
|
|
}
|
|
pollSlots := make(chan struct{}, 256)
|
|
for index, taskID := range taskIDs {
|
|
if taskID == "" {
|
|
continue
|
|
}
|
|
index, taskID := index, taskID
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
select {
|
|
case pollSlots <- struct{}{}:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
defer func() { <-pollSlots }()
|
|
err := pollVideoTask(ctx, client, opts, taskID, index, realUpstream)
|
|
mu.Lock()
|
|
latencies[index] = time.Since(startedTasks[index])
|
|
if err != nil {
|
|
failures++
|
|
if firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
} else {
|
|
completed++
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
comboSet := map[string]struct{}{}
|
|
for _, combo := range combinations {
|
|
comboSet[strings.Join(combo, "\x00")] = struct{}{}
|
|
}
|
|
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),
|
|
ForcedConversionTasks: forcedConversions,
|
|
}
|
|
if firstErr == nil && completed != requestCount {
|
|
firstErr = fmt.Errorf("video completed %d of %d tasks", completed, requestCount)
|
|
}
|
|
return phaseResult{report: result, latencies: latencies, err: firstErr}
|
|
}
|
|
|
|
func pollVideoTask(ctx context.Context, client *http.Client, opts options, taskID string, index int, realUpstream bool) error {
|
|
ticker := time.NewTicker(time.Second)
|
|
defer ticker.Stop()
|
|
for {
|
|
endpoint := opts.gateways[index%len(opts.gateways)] + "/api/v1/ai/result/" + url.PathEscape(taskID)
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
opts.setHeaders(req, realUpstream)
|
|
response, err := client.Do(req)
|
|
if err == nil {
|
|
if response.StatusCode != http.StatusOK {
|
|
err = responseStatusError(response)
|
|
} else {
|
|
var payload map[string]any
|
|
err = json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&payload)
|
|
_ = response.Body.Close()
|
|
status := strings.ToLower(strings.TrimSpace(fmt.Sprint(payload["status"])))
|
|
switch status {
|
|
case "succeeded", "success":
|
|
mediaURL := findMediaURL(payload)
|
|
if mediaURL == "" {
|
|
return fmt.Errorf("video task %s succeeded without a media URL", taskID)
|
|
}
|
|
return validateVideoAsset(ctx, http.DefaultClient, mediaURL)
|
|
case "failed", "cancelled", "canceled":
|
|
return fmt.Errorf("video task %s finished with status %s", taskID, status)
|
|
}
|
|
}
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-ticker.C:
|
|
}
|
|
}
|
|
}
|
|
|
|
func findMediaURL(value any) string {
|
|
switch typed := value.(type) {
|
|
case map[string]any:
|
|
for key, item := range typed {
|
|
normalized := strings.ToLower(strings.TrimSpace(key))
|
|
if (normalized == "url" || normalized == "video_url") && strings.HasPrefix(strings.TrimSpace(fmt.Sprint(item)), "http") {
|
|
return strings.TrimSpace(fmt.Sprint(item))
|
|
}
|
|
if mediaURL := findMediaURL(item); mediaURL != "" {
|
|
return mediaURL
|
|
}
|
|
}
|
|
case []any:
|
|
for _, item := range typed {
|
|
if mediaURL := findMediaURL(item); mediaURL != "" {
|
|
return mediaURL
|
|
}
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func validateVideoAsset(ctx context.Context, client *http.Client, mediaURL string) error {
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, mediaURL, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
response, err := client.Do(request)
|
|
if err != nil {
|
|
return fmt.Errorf("download final video: %w", err)
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
return fmt.Errorf("download final video: HTTP %d", response.StatusCode)
|
|
}
|
|
payload, err := io.ReadAll(io.LimitReader(response.Body, 64<<10))
|
|
if err != nil {
|
|
return fmt.Errorf("read final video: %w", err)
|
|
}
|
|
contentType := strings.ToLower(strings.TrimSpace(response.Header.Get("Content-Type")))
|
|
hasMP4Header := len(payload) >= 12 && string(payload[4:8]) == "ftyp"
|
|
if len(payload) < 12 || (!strings.HasPrefix(contentType, "video/") && !hasMP4Header) {
|
|
return fmt.Errorf("final video asset is invalid: bytes=%d content_type=%q", len(payload), contentType)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (o options) setHeaders(request *http.Request, realUpstream bool) {
|
|
if o.gatewayTLSName != "" {
|
|
request.Host = o.gatewayTLSName
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+o.apiKey)
|
|
request.Header.Set(runHeader, o.runID)
|
|
request.Header.Set(tokenHeader, o.runToken)
|
|
if realUpstream {
|
|
request.Header.Set(upstreamHeader, "real")
|
|
}
|
|
}
|
|
|
|
func (o options) emulatorFixtureURLs() []string {
|
|
out := make([]string, 0, 16)
|
|
for index := 0; index < 4; index++ {
|
|
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.png", o.emulatorURL, index))
|
|
}
|
|
for index := 4; index < 8; index++ {
|
|
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.jpg", o.emulatorURL, index))
|
|
}
|
|
for index := 8; index < 12; index++ {
|
|
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d.webp", o.emulatorURL, index))
|
|
}
|
|
for index := 12; index < 16; index++ {
|
|
out = append(out, fmt.Sprintf("%s/fixtures/image-%02d-4k.jpg", o.emulatorURL, index))
|
|
}
|
|
return out
|
|
}
|
|
|
|
func videoCombinations(images []string, count int) [][]string {
|
|
if len(images) < 3 {
|
|
return nil
|
|
}
|
|
if len(images) < 9 {
|
|
return [][]string{append([]string(nil), images[:3]...)}
|
|
}
|
|
out := make([][]string, 0, count)
|
|
seen := map[string]struct{}{}
|
|
for seed := 0; len(out) < count; seed++ {
|
|
imageCount := acceptanceworkload.VideoImageCount(seed)
|
|
combo := make([]string, 0, imageCount)
|
|
used := map[int]struct{}{}
|
|
state := uint64(seed + 1)
|
|
for len(combo) < imageCount {
|
|
state = state*6364136223846793005 + 1442695040888963407
|
|
index := int((state >> 32) % uint64(len(images)))
|
|
if _, exists := used[index]; exists {
|
|
continue
|
|
}
|
|
used[index] = struct{}{}
|
|
combo = append(combo, images[index])
|
|
}
|
|
signature := strings.Join(combo, "\x00")
|
|
if _, exists := seen[signature]; exists {
|
|
continue
|
|
}
|
|
seen[signature] = struct{}{}
|
|
out = append(out, combo)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func streamGeminiImageHash(reader io.Reader) (int64, string, error) {
|
|
buffered := bufio.NewReaderSize(reader, 64<<10)
|
|
if err := scanUntil(buffered, []byte(`"inlineData"`), 256<<20); err != nil {
|
|
return 0, "", err
|
|
}
|
|
if err := scanUntil(buffered, []byte(`"data"`), 1<<20); err != nil {
|
|
return 0, "", err
|
|
}
|
|
for {
|
|
item, err := buffered.ReadByte()
|
|
if err != nil {
|
|
return 0, "", err
|
|
}
|
|
if item == ':' {
|
|
break
|
|
}
|
|
}
|
|
for {
|
|
item, err := buffered.ReadByte()
|
|
if err != nil {
|
|
return 0, "", err
|
|
}
|
|
if item == '"' {
|
|
break
|
|
}
|
|
if item != ' ' && item != '\t' && item != '\r' && item != '\n' {
|
|
return 0, "", errors.New("Gemini inlineData.data is not a JSON string")
|
|
}
|
|
}
|
|
jsonString := &base64JSONStringReader{reader: buffered}
|
|
digest := sha256.New()
|
|
decodedBytes, err := io.Copy(digest, base64.NewDecoder(base64.StdEncoding, jsonString))
|
|
if err != nil {
|
|
return 0, "", fmt.Errorf("decode Gemini Base64 output: %w", err)
|
|
}
|
|
if decodedBytes == 0 {
|
|
return 0, "", errors.New("Gemini Base64 output is empty")
|
|
}
|
|
return decodedBytes, hex.EncodeToString(digest.Sum(nil)), nil
|
|
}
|
|
|
|
type base64JSONStringReader struct {
|
|
reader *bufio.Reader
|
|
done bool
|
|
}
|
|
|
|
func (r *base64JSONStringReader) Read(payload []byte) (int, error) {
|
|
if r.done {
|
|
return 0, io.EOF
|
|
}
|
|
written := 0
|
|
for written < len(payload) {
|
|
item, err := r.reader.ReadByte()
|
|
if err != nil {
|
|
return written, err
|
|
}
|
|
if item == '"' {
|
|
r.done = true
|
|
if written == 0 {
|
|
return 0, io.EOF
|
|
}
|
|
return written, nil
|
|
}
|
|
if item == '\\' {
|
|
return written, errors.New("escaped Base64 JSON strings are not supported")
|
|
}
|
|
payload[written] = item
|
|
written++
|
|
}
|
|
return written, nil
|
|
}
|
|
|
|
func scanUntil(reader *bufio.Reader, target []byte, limit int64) error {
|
|
matched := 0
|
|
for read := int64(0); read < limit; read++ {
|
|
item, err := reader.ReadByte()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if item == target[matched] {
|
|
matched++
|
|
if matched == len(target) {
|
|
return nil
|
|
}
|
|
continue
|
|
}
|
|
if item == target[0] {
|
|
matched = 1
|
|
} else {
|
|
matched = 0
|
|
}
|
|
}
|
|
return errors.New("Gemini inlineData output was not found")
|
|
}
|
|
|
|
func responseStatusError(response *http.Response) error {
|
|
defer response.Body.Close()
|
|
payload, _ := io.ReadAll(io.LimitReader(response.Body, 512))
|
|
return fmt.Errorf("HTTP %d: %s", response.StatusCode, strings.TrimSpace(string(payload)))
|
|
}
|
|
|
|
func percentileMilliseconds(values []time.Duration, quantile float64) float64 {
|
|
filtered := make([]time.Duration, 0, len(values))
|
|
for _, value := range values {
|
|
if value > 0 {
|
|
filtered = append(filtered, value)
|
|
}
|
|
}
|
|
if len(filtered) == 0 {
|
|
return 0
|
|
}
|
|
sort.Slice(filtered, func(i, j int) bool { return filtered[i] < filtered[j] })
|
|
index := int(float64(len(filtered)-1) * quantile)
|
|
return float64(filtered[index].Microseconds()) / 1000
|
|
}
|
|
|
|
func paddedPNG(size int) []byte {
|
|
base, _ := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")
|
|
if size <= len(base) {
|
|
return base
|
|
}
|
|
iendOffset := len(base) - 12
|
|
paddingLength := size - len(base) - 12
|
|
if paddingLength < 0 {
|
|
paddingLength = 0
|
|
}
|
|
chunk := make([]byte, 12+paddingLength)
|
|
binary.BigEndian.PutUint32(chunk[:4], uint32(paddingLength))
|
|
copy(chunk[4:8], []byte("teST"))
|
|
binary.BigEndian.PutUint32(chunk[8+paddingLength:], crc32.ChecksumIEEE(chunk[4:8+paddingLength]))
|
|
out := make([]byte, 0, size)
|
|
out = append(out, base[:iendOffset]...)
|
|
out = append(out, chunk...)
|
|
out = append(out, base[iendOffset:]...)
|
|
return out
|
|
}
|
|
|
|
func redactError(message string, opts options) string {
|
|
for _, secret := range []string{opts.apiKey, opts.runToken} {
|
|
if secret != "" {
|
|
message = strings.ReplaceAll(message, secret, "[REDACTED]")
|
|
}
|
|
}
|
|
return reportURLPattern.ReplaceAllString(message, "[REDACTED_URL]")
|
|
}
|
|
|
|
func splitCSV(value string) []string {
|
|
items := strings.Split(value, ",")
|
|
out := make([]string, 0, len(items))
|
|
for _, item := range items {
|
|
if item = strings.TrimSpace(item); item != "" {
|
|
out = append(out, item)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func env(name string, fallback string) string {
|
|
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
|
|
return value
|
|
}
|
|
return fallback
|
|
}
|