BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(acceptance): 增加生产同构媒体压力验收模式

Browse Source 
引入动态流量门禁、隔离验收身份与协议级 Gemini/Volces 模拟器，覆盖双站点 API、Worker、PostgreSQL、River、账务、回调和媒体物化链路。

新增 P24/P28/P32 容量阶梯、Worker 强杀恢复、真实小流量 canary、CAS 放量和失败保持 validation 的生产编排；Worker 执行槽、连接池、媒体并发和双站点副本数改为环境配置。

验证：Go 全量测试、真实 PostgreSQL 迁移集成测试、迁移安全检查、OpenAPI 生成、ShellCheck、Kustomize、gofmt 和 git diff --check。
This commit is contained in:
wangbo
2026-07-30 23:06:19 +08:00
parent f33d6d64e0
commit 3053ba4925
45 changed files with 5214 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/cmd/acceptance-emulator/main.go
+28
View File
@@ -0,0 +1,28 @@
package main
import (
"log/slog"
"net/http"
"os"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceemulator"
)
func main() {
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
if address == "" {
address = ":8090"
}
server := &http.Server{
Addr: address,
Handler: acceptanceemulator.New(acceptanceemulator.Config{}).Handler(),
ReadHeaderTimeout: 10 * time.Second,
}
slog.Info("acceptance protocol emulator started", "address", address)
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("acceptance protocol emulator stopped", "error", err)
os.Exit(1)
}
}
apps/api/cmd/acceptance-load/main.go
+825
View File
@@ -0,0 +1,825 @@
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
}
apps/api/cmd/acceptance-load/main_test.go
+127
View File
@@ -0,0 +1,127 @@
package main
import (
"bytes"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
)
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
payload := paddedPNG(4 << 20)
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
base64.StdEncoding.EncodeToString(payload))
size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
if err != nil {
t.Fatalf("stream Gemini output: %v", err)
}
expected := sha256.Sum256(payload)
if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
t.Fatalf("size=%d digest=%s", size, digest)
}
}
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
images := make([]string, 16)
for index := range images {
images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
}
combinations := videoCombinations(images, 128)
if len(combinations) != 128 {
t.Fatalf("combinations=%d", len(combinations))
}
seen := map[string]struct{}{}
for _, combination := range combinations {
seen[fmt.Sprint(combination)] = struct{}{}
if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
t.Fatalf("invalid combination size=%d", len(combination))
}
}
if len(seen) != 128 {
t.Fatalf("unique combinations=%d", len(seen))
}
}
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
output := paddedPNG(256 << 10)
encoded := base64.StdEncoding.EncodeToString(output)
var firstCalls atomic.Int64
var secondCalls atomic.Int64
newGateway := func(calls *atomic.Int64) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls.Add(1)
if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
t.Errorf("missing acceptance headers")
}
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode Gemini body: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"candidates": []any{map[string]any{"content": map[string]any{
"parts": []any{map[string]any{"inlineData": map[string]any{
"mimeType": "image/png", "data": encoded,
}}},
}}},
})
}))
}
first := newGateway(&firstCalls)
defer first.Close()
second := newGateway(&secondCalls)
defer second.Close()
opts := options{
gateways: []string{first.URL, second.URL}, apiKey: "key-1", 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)
if result.err != nil || result.report.Completed != 8 {
t.Fatalf("result=%+v err=%v", result.report, result.err)
}
if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
}
}
func TestValidateVideoAssetDownloadsFinalMedia(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()
if err := validateVideoAsset(t.Context(), server.Client(), server.URL+"/result.mp4"); err != nil {
t.Fatalf("validate video asset: %v", err)
}
if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
t.Fatalf("media URL=%q", got)
}
}
func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
got := redactError(
`token-1 failed at https://example.invalid/video.mp4?token=signed`,
options{runToken: "token-1"},
)
if got != `[REDACTED] failed at [REDACTED_URL]` {
t.Fatalf("redacted error=%q", got)
}
}
func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
opts := options{
apiKey: "key-1", runID: "run-1", runToken: "token-1",
gatewayTLSName: "ai.example.com",
}
opts.setHeaders(request, false)
if request.Host != "ai.example.com" {
t.Fatalf("Host=%q", request.Host)
}
}
apps/api/docs/swagger.json
+457
View File
@@ -2575,6 +2575,302 @@
}
}
},
"/api/admin/system/acceptance/runs": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "创建生产同构验收 Run",
"parameters": [
{
"description": "验收 Run",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.CreateAcceptanceRunInput"
}
}
],
"responses": {
"201": {
"description": "Created",
"schema": {
"$ref": "#/definitions/store.AcceptanceRun"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}": {
"get": {
"security": [
{
"BearerAuth": []
}
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "获取生产同构验收 Run",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.AcceptanceRun"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/abort": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "人工 CAS 中止验收并切回 live",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
},
{
"description": "线上 release 与 digest CAS",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.PromoteAcceptanceRunInput"
}
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.GatewayTrafficMode"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/activate": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "切换到 validation 并激活 Run",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.GatewayTrafficMode"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/finish": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "记录验收门禁结果",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
},
{
"description": "门禁结果",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.FinishAcceptanceRunInput"
}
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.AcceptanceRun"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/promote": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "通过 CAS 切回 live",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
},
{
"description": "线上 release 与 digest CAS",
"name": "body",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/store.PromoteAcceptanceRunInput"
}
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.GatewayTrafficMode"
}
}
}
}
},
"/api/admin/system/acceptance/runs/{runID}/retry": {
"post": {
"security": [
{
"BearerAuth": []
}
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "在 validation 关闭正式流量的状态下重试失败 Run",
"parameters": [
{
"type": "string",
"description": "验收 Run ID",
"name": "runID",
"in": "path",
"required": true
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.AcceptanceRun"
}
}
}
}
},
"/api/admin/system/acceptance/traffic-mode": {
"get": {
"security": [
{
"BearerAuth": []
}
],
"produces": [
"application/json"
],
"tags": [
"acceptance"
],
"summary": "获取 Gateway 流量模式",
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/store.GatewayTrafficMode"
}
}
}
}
},
"/api/admin/system/client-customization/settings": {
"get": {
"security": [
@@ -12618,6 +12914,67 @@
}
}
},
"store.AcceptanceRun": {
"type": "object",
"properties": {
"apiImageDigest": {
"type": "string"
},
"apiKeyId": {
"type": "string"
},
"callbackUrl": {
"type": "string"
},
"capacityProfile": {
"type": "string"
},
"config": {
"type": "object",
"additionalProperties": {}
},
"createdAt": {
"type": "string"
},
"emulatorBaseUrl": {
"type": "string"
},
"failureReason": {
"type": "string"
},
"finishedAt": {
"type": "string"
},
"id": {
"type": "string"
},
"promotedAt": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"report": {
"type": "object",
"additionalProperties": {}
},
"startedAt": {
"type": "string"
},
"status": {
"type": "string"
},
"updatedAt": {
"type": "string"
},
"userId": {
"type": "string"
},
"workerImageDigest": {
"type": "string"
}
}
},
"store.AccessRule": {
"type": "object",
"properties": {
@@ -12758,6 +13115,9 @@
"store.AdminGatewayTask": {
"type": "object",
"properties": {
"acceptanceRunId": {
"type": "string"
},
"adminContext": {
"$ref": "#/definitions/store.AdminTaskContext"
},
@@ -13461,6 +13821,42 @@
}
}
},
"store.CreateAcceptanceRunInput": {
"type": "object",
"properties": {
"apiImageDigest": {
"type": "string"
},
"apiKeyId": {
"type": "string"
},
"callbackUrl": {
"type": "string"
},
"capacityProfile": {
"type": "string"
},
"config": {
"type": "object",
"additionalProperties": {}
},
"emulatorBaseUrl": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"token": {
"type": "string"
},
"userId": {
"type": "string"
},
"workerImageDigest": {
"type": "string"
}
}
},
"store.CreatePlatformInput": {
"type": "object",
"properties": {
@@ -13753,9 +14149,27 @@
}
}
},
"store.FinishAcceptanceRunInput": {
"type": "object",
"properties": {
"failureReason": {
"type": "string"
},
"passed": {
"type": "boolean"
},
"report": {
"type": "object",
"additionalProperties": {}
}
}
},
"store.GatewayTask": {
"type": "object",
"properties": {
"acceptanceRunId": {
"type": "string"
},
"apiKeyId": {
"type": "string"
},
@@ -14061,6 +14475,32 @@
}
}
},
"store.GatewayTrafficMode": {
"type": "object",
"properties": {
"apiImageDigest": {
"type": "string"
},
"mode": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"revision": {
"type": "integer"
},
"runId": {
"type": "string"
},
"updatedAt": {
"type": "string"
},
"workerImageDigest": {
"type": "string"
}
}
},
"store.GatewayUser": {
"type": "object",
"properties": {
@@ -15058,6 +15498,23 @@
}
}
},
"store.PromoteAcceptanceRunInput": {
"type": "object",
"properties": {
"apiImageDigest": {
"type": "string"
},
"releaseSha": {
"type": "string"
},
"revision": {
"type": "integer"
},
"workerImageDigest": {
"type": "string"
}
}
},
"store.RateLimitMetricStatus": {
"type": "object",
"properties": {
apps/api/docs/swagger.yaml
+288
View File
@@ -1993,6 +1993,47 @@ definitions:
userId:
type: string
type: object
store.AcceptanceRun:
properties:
apiImageDigest:
type: string
apiKeyId:
type: string
callbackUrl:
type: string
capacityProfile:
type: string
config:
additionalProperties: {}
type: object
createdAt:
type: string
emulatorBaseUrl:
type: string
failureReason:
type: string
finishedAt:
type: string
id:
type: string
promotedAt:
type: string
releaseSha:
type: string
report:
additionalProperties: {}
type: object
startedAt:
type: string
status:
type: string
updatedAt:
type: string
userId:
type: string
workerImageDigest:
type: string
type: object
store.AccessRule:
properties:
conditions:
@@ -2087,6 +2128,8 @@ definitions:
type: object
store.AdminGatewayTask:
properties:
acceptanceRunId:
type: string
adminContext:
$ref: '#/definitions/store.AdminTaskContext'
apiKeyId:
@@ -2562,6 +2605,30 @@ definitions:
type: string
type: array
type: object
store.CreateAcceptanceRunInput:
properties:
apiImageDigest:
type: string
apiKeyId:
type: string
callbackUrl:
type: string
capacityProfile:
type: string
config:
additionalProperties: {}
type: object
emulatorBaseUrl:
type: string
releaseSha:
type: string
token:
type: string
userId:
type: string
workerImageDigest:
type: string
type: object
store.CreatePlatformInput:
properties:
authType:
@@ -2760,8 +2827,20 @@ definitions:
resultUploadPolicy:
type: string
type: object
store.FinishAcceptanceRunInput:
properties:
failureReason:
type: string
passed:
type: boolean
report:
additionalProperties: {}
type: object
type: object
store.GatewayTask:
properties:
acceptanceRunId:
type: string
apiKeyId:
type: string
apiKeyName:
@@ -2970,6 +3049,23 @@ definitions:
tenantKey:
type: string
type: object
store.GatewayTrafficMode:
properties:
apiImageDigest:
type: string
mode:
type: string
releaseSha:
type: string
revision:
type: integer
runId:
type: string
updatedAt:
type: string
workerImageDigest:
type: string
type: object
store.GatewayUser:
properties:
authProfile:
@@ -3642,6 +3738,17 @@ definitions:
taskId:
type: string
type: object
store.PromoteAcceptanceRunInput:
properties:
apiImageDigest:
type: string
releaseSha:
type: string
revision:
type: integer
workerImageDigest:
type: string
type: object
store.RateLimitMetricStatus:
properties:
currentValue:
@@ -5645,6 +5752,187 @@ paths:
summary: 更新 Runner 策略
tags:
- runtime
/api/admin/system/acceptance/runs:
post:
consumes:
- application/json
parameters:
- description: 验收 Run
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.CreateAcceptanceRunInput'
produces:
- application/json
responses:
"201":
description: Created
schema:
$ref: '#/definitions/store.AcceptanceRun'
security:
- BearerAuth: []
summary: 创建生产同构验收 Run
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}:
get:
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.AcceptanceRun'
security:
- BearerAuth: []
summary: 获取生产同构验收 Run
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/abort:
post:
consumes:
- application/json
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
- description: 线上 release 与 digest CAS
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.PromoteAcceptanceRunInput'
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.GatewayTrafficMode'
security:
- BearerAuth: []
summary: 人工 CAS 中止验收并切回 live
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/activate:
post:
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.GatewayTrafficMode'
security:
- BearerAuth: []
summary: 切换到 validation 并激活 Run
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/finish:
post:
consumes:
- application/json
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
- description: 门禁结果
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.FinishAcceptanceRunInput'
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.AcceptanceRun'
security:
- BearerAuth: []
summary: 记录验收门禁结果
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/promote:
post:
consumes:
- application/json
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
- description: 线上 release 与 digest CAS
in: body
name: body
required: true
schema:
$ref: '#/definitions/store.PromoteAcceptanceRunInput'
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.GatewayTrafficMode'
security:
- BearerAuth: []
summary: 通过 CAS 切回 live
tags:
- acceptance
/api/admin/system/acceptance/runs/{runID}/retry:
post:
parameters:
- description: 验收 Run ID
in: path
name: runID
required: true
type: string
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.AcceptanceRun'
security:
- BearerAuth: []
summary: 在 validation 关闭正式流量的状态下重试失败 Run
tags:
- acceptance
/api/admin/system/acceptance/traffic-mode:
get:
produces:
- application/json
responses:
"200":
description: OK
schema:
$ref: '#/definitions/store.GatewayTrafficMode'
security:
- BearerAuth: []
summary: 获取 Gateway 流量模式
tags:
- acceptance
/api/admin/system/client-customization/settings:
get:
description: 返回客户端名称、英文名称、平台名和图标路径；数据库对象尚未创建时返回默认设置。
apps/api/internal/acceptanceemulator/server.go
+658
View File
@@ -0,0 +1,658 @@
package acceptanceemulator
import (
"bytes"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/binary"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"hash/crc32"
"image"
"image/color"
"image/jpeg"
"image/png"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
_ "golang.org/x/image/webp"
)
const maxProtocolBodyBytes = 128 << 20
type Config struct {
Now func() time.Time
HTTPClient *http.Client
Wait func(context.Context, time.Duration) error
}
type Server struct {
now func() time.Time
httpClient *http.Client
wait func(context.Context, time.Duration) error
mu sync.RWMutex
tasks map[string]videoTask
fixtures map[string]fixture
nextID atomic.Uint64
report Report
callbacks map[string]map[int64]int
pngSmall string
pngLarge string
pngPeak string
}
type Report struct {
GeminiRequests int64 `json:"geminiRequests"`
GeminiInvalid int64 `json:"geminiInvalid"`
GeminiInputBytes int64 `json:"geminiInputBytes"`
GeminiOutputBytes int64 `json:"geminiOutputBytes"`
VideoSubmissions int64 `json:"videoSubmissions"`
VideoPolls int64 `json:"videoPolls"`
VideoInvalid int64 `json:"videoInvalid"`
VideoReferenceCounts map[string]int64 `json:"videoReferenceCounts"`
VideoRoleCounts map[string]int64 `json:"videoRoleCounts"`
UniqueImageHashes int `json:"uniqueImageHashes"`
UniqueVideoTasks int `json:"uniqueVideoTasks"`
ForcedConversions int64 `json:"forcedConversionRequests"`
VerifiedConversions int64 `json:"verifiedConversions"`
CallbackEvents int64 `json:"callbackEvents"`
DuplicateCallbacks int64 `json:"duplicateCallbacks"`
}
type videoTask struct {
ID string
Model string
CreatedAt time.Time
ReadyAt time.Time
ImageRefs []imageReference
}
type imageReference struct {
Role string
SHA256 string
ContentType string
Width int
Height int
}
type fixture struct {
ContentType string
Payload []byte
}
func New(config Config) *Server {
now := config.Now
if now == nil {
now = time.Now
}
client := config.HTTPClient
if client == nil {
client = &http.Client{Timeout: 30 * time.Second}
}
waiter := config.Wait
if waiter == nil {
waiter = wait
}
return &Server{
now: now, httpClient: client, wait: waiter, tasks: map[string]videoTask{}, fixtures: buildFixtures(),
callbacks: map[string]map[int64]int{},
report: Report{
VideoReferenceCounts: map[string]int64{},
VideoRoleCounts: map[string]int64{},
},
pngSmall: base64.StdEncoding.EncodeToString(paddedPNG(256 << 10)),
pngLarge: base64.StdEncoding.EncodeToString(paddedPNG(4 << 20)),
pngPeak: base64.StdEncoding.EncodeToString(paddedPNG(8 << 20)),
}
}
func (s *Server) Handler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", s.health)
mux.HandleFunc("GET /report", s.getReport)
mux.HandleFunc("POST /v1beta/models/", s.geminiGenerateContent)
mux.HandleFunc("POST /v1/models/", s.geminiGenerateContent)
mux.HandleFunc("POST /contents/generations/tasks", s.submitVideo)
mux.HandleFunc("GET /contents/generations/tasks/{taskID}", s.getVideo)
mux.HandleFunc("DELETE /contents/generations/tasks/{taskID}", s.deleteVideo)
mux.HandleFunc("GET /media/{asset}", s.getMedia)
mux.HandleFunc("GET /fixtures/{asset}", s.getFixture)
mux.HandleFunc("POST /callbacks", s.collectCallback)
return mux
}
func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
if !strings.HasSuffix(r.URL.Path, ":generateContent") {
http.NotFound(w, r)
return
}
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordGeminiInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
inputBytes, err := validateGeminiImageRequest(body)
if err != nil {
s.recordGeminiInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
outputBytes, outputBase64, delay := s.geminiProfile(inputBytes)
if err := s.wait(r.Context(), delay); err != nil {
return
}
s.mu.Lock()
s.report.GeminiRequests++
s.report.GeminiInputBytes += int64(inputBytes)
s.report.GeminiOutputBytes += int64(outputBytes)
s.mu.Unlock()
w.Header().Set("X-Request-ID", "acceptance-gemini-"+strconv.FormatUint(s.nextID.Add(1), 10))
writeJSON(w, http.StatusOK, map[string]any{
"candidates": []any{map[string]any{
"index": 0,
"content": map[string]any{
"role": "model",
"parts": []any{map[string]any{"inlineData": map[string]any{
"mimeType": "image/png",
"data": outputBase64,
}}},
},
"finishReason": "STOP",
}},
"usageMetadata": map[string]any{
"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2,
},
})
}
func (s *Server) geminiProfile(inputBytes int) (int, string, time.Duration) {
switch {
case inputBytes > acceptanceworkload.GeminiLarge.InputBytes:
profile := acceptanceworkload.GeminiPeak
return profile.OutputBytes, s.pngPeak, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
case inputBytes > acceptanceworkload.GeminiBaseline.InputBytes:
profile := acceptanceworkload.GeminiLarge
return profile.OutputBytes, s.pngLarge, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
default:
profile := acceptanceworkload.GeminiBaseline
return profile.OutputBytes, s.pngSmall, profile.DelayMin
}
}
func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordVideoInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
refs, longRun, forceConversion, err := s.validateVideoRequest(r.Context(), body)
if err != nil {
s.recordVideoInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
delay := videoDelay(body, longRun)
id := "acceptance-video-" + strconv.FormatUint(s.nextID.Add(1), 10)
task := videoTask{
ID: id, Model: strings.TrimSpace(stringValue(body["model"])),
CreatedAt: s.now(), ReadyAt: s.now().Add(delay), ImageRefs: refs,
}
s.mu.Lock()
s.tasks[id] = task
s.report.VideoSubmissions++
s.report.VideoReferenceCounts[strconv.Itoa(len(refs))]++
if forceConversion {
s.report.ForcedConversions++
s.report.VerifiedConversions++
}
for _, ref := range refs {
s.report.VideoRoleCounts[ref.Role]++
}
s.mu.Unlock()
w.Header().Set("X-Request-ID", id)
writeJSON(w, http.StatusOK, map[string]any{
"id": id, "model": task.Model, "status": "queued",
"created_at": task.CreatedAt.Unix(),
})
}
func (s *Server) getVideo(w http.ResponseWriter, r *http.Request) {
id := strings.TrimSpace(r.PathValue("taskID"))
s.mu.Lock()
task, ok := s.tasks[id]
if ok {
s.report.VideoPolls++
}
s.mu.Unlock()
if !ok {
writeProtocolError(w, http.StatusNotFound, "video task not found")
return
}
status := "running"
content := map[string]any{}
if !s.now().Before(task.ReadyAt) {
status = "succeeded"
content["video_url"] = requestOrigin(r) + "/media/" + url.PathEscape(id) + ".mp4"
}
w.Header().Set("X-Request-ID", id)
writeJSON(w, http.StatusOK, map[string]any{
"id": id, "model": task.Model, "status": status, "content": content,
"created_at": task.CreatedAt.Unix(),
"usage": map[string]any{"completion_tokens": 1, "total_tokens": 1},
})
}
func (s *Server) deleteVideo(w http.ResponseWriter, r *http.Request) {
id := strings.TrimSpace(r.PathValue("taskID"))
s.mu.Lock()
_, ok := s.tasks[id]
delete(s.tasks, id)
s.mu.Unlock()
if !ok {
writeProtocolError(w, http.StatusNotFound, "video task not found")
return
}
writeJSON(w, http.StatusOK, map[string]any{"id": id, "status": "cancelled"})
}
func (s *Server) getMedia(w http.ResponseWriter, r *http.Request) {
asset := strings.TrimSuffix(strings.TrimSpace(r.PathValue("asset")), ".mp4")
s.mu.RLock()
_, ok := s.tasks[asset]
s.mu.RUnlock()
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "video/mp4")
w.WriteHeader(http.StatusOK)
_, _ = w.Write(minimalMP4)
}
func (s *Server) getFixture(w http.ResponseWriter, r *http.Request) {
fixture, ok := s.fixtures[strings.TrimSpace(r.PathValue("asset"))]
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", fixture.ContentType)
w.Header().Set("Cache-Control", "public, max-age=3600")
w.Header().Set("X-Content-SHA256", payloadSHA256(fixture.Payload))
w.WriteHeader(http.StatusOK)
_, _ = w.Write(fixture.Payload)
}
func (s *Server) getReport(w http.ResponseWriter, _ *http.Request) {
s.mu.RLock()
report := s.report
tasks := make([]videoTask, 0, len(s.tasks))
for _, task := range s.tasks {
tasks = append(tasks, task)
}
s.mu.RUnlock()
hashes := map[string]struct{}{}
for _, task := range tasks {
for _, ref := range task.ImageRefs {
hashes[ref.SHA256] = struct{}{}
}
}
report.UniqueImageHashes = len(hashes)
report.UniqueVideoTasks = len(tasks)
writeJSON(w, http.StatusOK, report)
}
func (s *Server) collectCallback(w http.ResponseWriter, r *http.Request) {
var payload map[string]any
if err := decodeJSON(r, &payload); err != nil {
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
seq := int64(0)
switch value := payload["seq"].(type) {
case json.Number:
seq, _ = value.Int64()
case float64:
seq = int64(value)
}
if taskID == "" || seq <= 0 {
writeProtocolError(w, http.StatusBadRequest, "callback taskId and seq are required")
return
}
s.mu.Lock()
seen := s.callbacks[taskID]
if seen == nil {
seen = map[int64]int{}
s.callbacks[taskID] = seen
}
seen[seq]++
s.report.CallbackEvents++
if seen[seq] > 1 {
s.report.DuplicateCallbacks++
}
s.mu.Unlock()
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (s *Server) validateVideoRequest(ctx context.Context, body map[string]any) ([]imageReference, bool, bool, error) {
content, _ := body["content"].([]any)
refs := make([]imageReference, 0, 9)
firstFrames := 0
lastFrames := 0
longRun := false
forceConversion := false
for _, raw := range content {
item, _ := raw.(map[string]any)
switch strings.TrimSpace(stringValue(item["type"])) {
case "text":
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-long-recovery") {
longRun = true
}
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-force-conversion") {
forceConversion = true
}
case "image_url":
role := strings.TrimSpace(stringValue(item["role"]))
switch role {
case "reference_image":
case "first_frame":
firstFrames++
case "last_frame":
lastFrames++
default:
return nil, false, false, fmt.Errorf("unsupported image role %q", role)
}
nested, _ := item["image_url"].(map[string]any)
rawURL := strings.TrimSpace(stringValue(nested["url"]))
ref, err := s.fetchImageReference(ctx, role, rawURL)
if err != nil {
return nil, false, false, err
}
refs = append(refs, ref)
}
}
if len(refs) != 3 && len(refs) != 6 && len(refs) != 9 {
return nil, false, false, fmt.Errorf("expected 3, 6, or 9 image references, got %d", len(refs))
}
if firstFrames > 1 || lastFrames > 1 {
return nil, false, false, errors.New("first_frame and last_frame must be unique")
}
if forceConversion {
for _, ref := range refs {
if ref.Width >= 3840 || ref.Height >= 2160 {
return nil, false, false, errors.New("forced 4K image did not pass through automatic normalization")
}
}
}
return refs, longRun, forceConversion, nil
}
func (s *Server) fetchImageReference(ctx context.Context, role string, rawURL string) (imageReference, error) {
if rawURL == "" {
return imageReference{}, errors.New("image reference URL is required")
}
request, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return imageReference{}, err
}
response, err := s.httpClient.Do(request)
if err != nil {
return imageReference{}, fmt.Errorf("fetch image reference: %w", err)
}
defer response.Body.Close()
if response.StatusCode < 200 || response.StatusCode >= 300 {
return imageReference{}, fmt.Errorf("fetch image reference: status %d", response.StatusCode)
}
payload, err := io.ReadAll(io.LimitReader(response.Body, 32<<20))
if err != nil {
return imageReference{}, err
}
if len(payload) == 0 {
return imageReference{}, errors.New("image reference is empty")
}
config, format, err := image.DecodeConfig(bytes.NewReader(payload))
if err != nil {
return imageReference{}, fmt.Errorf("decode image reference: %w", err)
}
sum := sha256.Sum256(payload)
return imageReference{
Role: role, SHA256: hex.EncodeToString(sum[:]), ContentType: "image/" + format,
Width: config.Width, Height: config.Height,
}, nil
}
func validateGeminiImageRequest(body map[string]any) (int, error) {
generationConfig, _ := body["generationConfig"].(map[string]any)
modalities, _ := generationConfig["responseModalities"].([]any)
hasImageModality := false
for _, modality := range modalities {
if strings.EqualFold(strings.TrimSpace(stringValue(modality)), "IMAGE") {
hasImageModality = true
}
}
if !hasImageModality {
return 0, errors.New("generationConfig.responseModalities must contain IMAGE")
}
total := 0
contents, _ := body["contents"].([]any)
for _, rawContent := range contents {
content, _ := rawContent.(map[string]any)
parts, _ := content["parts"].([]any)
for _, rawPart := range parts {
part, _ := rawPart.(map[string]any)
inline, _ := part["inlineData"].(map[string]any)
if inline == nil {
inline, _ = part["inline_data"].(map[string]any)
}
encoded := strings.TrimSpace(stringValue(inline["data"]))
if encoded == "" {
continue
}
decoded, err := base64.StdEncoding.DecodeString(encoded)
if err != nil {
return 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err)
}
if len(decoded) == 0 {
return 0, errors.New("Gemini inlineData is empty")
}
total += len(decoded)
}
}
if total == 0 {
return 0, errors.New("Gemini request has no inlineData image")
}
return total, nil
}
func (s *Server) recordGeminiInvalid() {
s.mu.Lock()
s.report.GeminiInvalid++
s.mu.Unlock()
}
func (s *Server) recordVideoInvalid() {
s.mu.Lock()
s.report.VideoInvalid++
s.mu.Unlock()
}
func decodeJSON(r *http.Request, target any) error {
reader := io.LimitReader(r.Body, maxProtocolBodyBytes+1)
payload, err := io.ReadAll(reader)
if err != nil {
return err
}
if len(payload) > maxProtocolBodyBytes {
return errors.New("protocol request body is too large")
}
decoder := json.NewDecoder(bytes.NewReader(payload))
decoder.UseNumber()
return decoder.Decode(target)
}
func writeJSON(w http.ResponseWriter, status int, value any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(value)
}
func writeProtocolError(w http.ResponseWriter, status int, message string) {
writeJSON(w, status, map[string]any{"error": map[string]any{
"code": "acceptance_protocol_invalid", "message": message,
}})
}
func wait(ctx context.Context, delay time.Duration) error {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
func scaledDelay(minimum time.Duration, maximum time.Duration, seed int) time.Duration {
if maximum <= minimum {
return minimum
}
windowSeconds := int64((maximum - minimum) / time.Second)
if windowSeconds <= 0 {
return minimum
}
return minimum + time.Duration(int64(seed)%(windowSeconds+1))*time.Second
}
func videoDelay(body map[string]any, longRun bool) time.Duration {
seed := int64(0)
switch value := body["seed"].(type) {
case json.Number:
seed, _ = value.Int64()
case float64:
seed = int64(value)
case int:
seed = int64(value)
}
if seed < 0 {
seed = -seed
}
if longRun {
return 2*time.Minute + time.Duration(seed%61)*time.Second
}
return 5*time.Second + time.Duration(seed%11)*time.Second
}
func requestOrigin(r *http.Request) string {
scheme := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
if scheme == "" {
scheme = "http"
}
host := strings.TrimSpace(r.Host)
return scheme + "://" + host
}
func stringValue(value any) string {
switch typed := value.(type) {
case nil:
return ""
case string:
return typed
case json.Number:
return typed.String()
default:
return fmt.Sprint(value)
}
}
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, len(base)+len(chunk))
out = append(out, base[:iendOffset]...)
out = append(out, chunk...)
out = append(out, base[iendOffset:]...)
return out
}
func buildFixtures() map[string]fixture {
fixtures := map[string]fixture{}
for index := 0; index < 4; index++ {
imageValue := patternedImage(768, 512, index+1)
var encoded bytes.Buffer
_ = png.Encode(&encoded, imageValue)
fixtures[fmt.Sprintf("image-%02d.png", index)] = fixture{ContentType: "image/png", Payload: encoded.Bytes()}
}
for index := 0; index < 4; index++ {
imageValue := patternedImage(1024, 768, index+11)
var encoded bytes.Buffer
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 88})
fixtures[fmt.Sprintf("image-%02d.jpg", index+4)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
}
webpPayload, _ := base64.StdEncoding.DecodeString("UklGRiIAAABXRUJQVlA4IBYAAAAwAQCdASoBAAEALmk0mk0iIiIiIgBoSygABc6zbAAA")
for index := 0; index < 4; index++ {
payload := append([]byte(nil), webpPayload...)
payload = append(payload, byte(index))
fixtures[fmt.Sprintf("image-%02d.webp", index+8)] = fixture{ContentType: "image/webp", Payload: payload}
}
for index := 0; index < 4; index++ {
imageValue := patternedImage(4096, 2160, index+21)
var encoded bytes.Buffer
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 90})
fixtures[fmt.Sprintf("image-%02d-4k.jpg", index+12)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
}
return fixtures
}
func patternedImage(width int, height int, seed int) image.Image {
out := image.NewNRGBA(image.Rect(0, 0, width, height))
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
out.SetNRGBA(x, y, color.NRGBA{
R: uint8((x + seed*31) % 256),
G: uint8((y + seed*47) % 256),
B: uint8((x/8 + y/8 + seed*59) % 256),
A: 255,
})
}
}
return out
}
func payloadSHA256(payload []byte) string {
sum := sha256.Sum256(payload)
return hex.EncodeToString(sum[:])
}
var minimalMP4 = []byte{
0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm',
0, 0, 0, 0, 'i', 's', 'o', 'm', 'm', 'p', '4', '2',
0, 0, 0, 8, 'm', 'd', 'a', 't',
}
apps/api/internal/acceptanceemulator/server_test.go
+186
View File
@@ -0,0 +1,186 @@
package acceptanceemulator
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
)
func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
var mu sync.Mutex
now := time.Unix(1_800_000_000, 0)
server := New(Config{
Now: func() time.Time {
mu.Lock()
defer mu.Unlock()
return now
},
Wait: func(context.Context, time.Duration) error { return nil },
})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
input := paddedPNG(256 << 10)
geminiBody, _ := json.Marshal(map[string]any{
"contents": []any{map[string]any{"parts": []any{
map[string]any{"text": "edit image"},
map[string]any{"inlineData": map[string]any{
"mimeType": "image/png", "data": base64.StdEncoding.EncodeToString(input),
}},
}}},
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
})
response, err := http.Post(httpServer.URL+"/v1beta/models/gemini-test:generateContent", "application/json", bytes.NewReader(geminiBody))
if err != nil {
t.Fatalf("Gemini request: %v", err)
}
if response.StatusCode != http.StatusOK {
payload, _ := io.ReadAll(response.Body)
t.Fatalf("Gemini status=%d body=%s", response.StatusCode, payload)
}
var geminiResult map[string]any
if err := json.NewDecoder(response.Body).Decode(&geminiResult); err != nil {
t.Fatalf("decode Gemini response: %v", err)
}
_ = response.Body.Close()
fixtureResponse, err := http.Get(httpServer.URL + "/fixtures/image-12-4k.jpg")
if err != nil {
t.Fatalf("get 4K fixture: %v", err)
}
config, format, err := image.DecodeConfig(fixtureResponse.Body)
_ = fixtureResponse.Body.Close()
if err != nil || format != "jpeg" || config.Width != 4096 || config.Height != 2160 {
t.Fatalf("4K fixture format=%s config=%+v err=%v", format, config, err)
}
content := []any{map[string]any{"type": "text", "text": "video"}}
for index, name := range []string{"image-00.png", "image-04.jpg", "image-08.webp"} {
role := "reference_image"
if index == 0 {
role = "first_frame"
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
})
}
videoBody, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": 1})
response, err = http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(videoBody))
if err != nil {
t.Fatalf("submit video: %v", err)
}
if response.StatusCode != http.StatusOK {
t.Fatalf("submit video status=%d", response.StatusCode)
}
var submitted map[string]any
_ = json.NewDecoder(response.Body).Decode(&submitted)
_ = response.Body.Close()
taskID := strings.TrimSpace(stringValue(submitted["id"]))
if taskID == "" {
t.Fatal("video emulator returned no task ID")
}
mu.Lock()
now = now.Add(20 * time.Second)
mu.Unlock()
response, err = http.Get(httpServer.URL + "/contents/generations/tasks/" + taskID)
if err != nil {
t.Fatalf("poll video: %v", err)
}
var polled map[string]any
_ = json.NewDecoder(response.Body).Decode(&polled)
_ = response.Body.Close()
if polled["status"] != "succeeded" {
t.Fatalf("video status=%v", polled["status"])
}
response, err = http.Get(httpServer.URL + "/report")
if err != nil {
t.Fatalf("get report: %v", err)
}
var report Report
_ = json.NewDecoder(response.Body).Decode(&report)
_ = response.Body.Close()
if report.GeminiRequests != 1 || report.VideoSubmissions != 1 || report.VideoInvalid != 0 ||
report.VideoReferenceCounts["3"] != 1 || report.UniqueImageHashes != 3 {
t.Fatalf("unexpected emulator report: %+v", report)
}
}
func TestPaddedPNGIsValidAndExact(t *testing.T) {
for _, size := range []int{256 << 10, 4 << 20, 8 << 20} {
payload := paddedPNG(size)
if len(payload) != size {
t.Fatalf("PNG bytes=%d, want=%d", len(payload), size)
}
if _, _, err := image.DecodeConfig(bytes.NewReader(payload)); err != nil {
t.Fatalf("decode padded PNG: %v", err)
}
}
}
func TestScaledDelayStaysOnWholeSecondsWithinProfile(t *testing.T) {
delay := scaledDelay(8*time.Second, 15*time.Second, 2<<20)
if delay < 8*time.Second || delay > 15*time.Second || delay%time.Second != 0 {
t.Fatalf("scaled delay=%s", delay)
}
}
func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
for _, imageCount := range []int{3, 6, 9} {
content := []any{map[string]any{"type": "text", "text": "multi-reference video"}}
for index := 0; index < imageCount; index++ {
role := "reference_image"
if index == 0 {
role = "first_frame"
}
if imageCount > 3 && index == imageCount-1 {
role = "last_frame"
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{
"url": fmt.Sprintf("%s/fixtures/image-%02d.png", httpServer.URL, index%4),
},
})
}
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": imageCount})
response, err := http.Post(httpServer.URL+"/contents/generations/tasks", "application/json", bytes.NewReader(body))
if err != nil {
t.Fatalf("%d-image submit: %v", imageCount, err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("%d-image status=%d", imageCount, response.StatusCode)
}
}
response, err := http.Get(httpServer.URL + "/report")
if err != nil {
t.Fatalf("get report: %v", err)
}
defer response.Body.Close()
var report Report
if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
t.Fatalf("decode report: %v", err)
}
for _, imageCount := range []string{"3", "6", "9"} {
if report.VideoReferenceCounts[imageCount] != 1 {
t.Fatalf("reference counts=%v", report.VideoReferenceCounts)
}
}
}
apps/api/internal/acceptanceworkload/profiles.go
+52
View File
@@ -0,0 +1,52 @@
package acceptanceworkload
import "time"
type GeminiProfile struct {
Name string
Requests int
InputBytes int
OutputBytes int
DelayMin time.Duration
DelayMax time.Duration
}
var (
GeminiBaseline = GeminiProfile{
Name: "gemini-baseline", Requests: 1000, InputBytes: 256 << 10, OutputBytes: 256 << 10,
DelayMin: 4 * time.Second, DelayMax: 4 * time.Second,
}
GeminiLarge = GeminiProfile{
Name: "gemini-large", Requests: 128, InputBytes: 2 << 20, OutputBytes: 4 << 20,
DelayMin: 8 * time.Second, DelayMax: 15 * time.Second,
}
GeminiPeak = GeminiProfile{
Name: "gemini-peak", Requests: 32, InputBytes: 8 << 20, OutputBytes: 8 << 20,
DelayMin: 15 * time.Second, DelayMax: 30 * time.Second,
}
)
func GeminiProfileByName(name string) (GeminiProfile, bool) {
switch name {
case GeminiBaseline.Name:
return GeminiBaseline, true
case GeminiLarge.Name:
return GeminiLarge, true
case GeminiPeak.Name:
return GeminiPeak, true
default:
return GeminiProfile{}, false
}
}
// VideoImageCount provides a deterministic 60/30/10 split for 3/6/9-image tasks.
func VideoImageCount(index int) int {
switch index % 10 {
case 0, 1, 2, 3, 4, 5:
return 3
case 6, 7, 8:
return 6
default:
return 9
}
}
apps/api/internal/acceptanceworkload/profiles_test.go
+23
View File
@@ -0,0 +1,23 @@
package acceptanceworkload
import "testing"
func TestGeminiProfilesAndVideoDistribution(t *testing.T) {
for _, profile := range []GeminiProfile{GeminiBaseline, GeminiLarge, GeminiPeak} {
resolved, ok := GeminiProfileByName(profile.Name)
if !ok || resolved != profile {
t.Fatalf("profile %q resolved to %+v, ok=%v", profile.Name, resolved, ok)
}
if profile.Requests <= 0 || profile.InputBytes <= 0 || profile.OutputBytes <= 0 ||
profile.DelayMin <= 0 || profile.DelayMax < profile.DelayMin {
t.Fatalf("invalid Gemini profile: %+v", profile)
}
}
counts := map[int]int{}
for index := 0; index < 100; index++ {
counts[VideoImageCount(index)]++
}
if counts[3] != 60 || counts[6] != 30 || counts[9] != 10 {
t.Fatalf("video image distribution=%v", counts)
}
}
apps/api/internal/clients/clients_test.go
+2
View File
@@ -2435,6 +2435,8 @@ func TestVolcesVideoBodyAllowsOnlyTaskPayloadFields(t *testing.T) {
"tools": []any{map[string]any{"type": "web_search"}},
"task_id": "local-task-id",
"runMode": "simulation",
"modelType": "omni_video",
"model_type": "omni_video",
"fps": 24,
"content": []any{
map[string]any{"type": "text", "text": "Use <<<element_1>>> in a product reveal"},
apps/api/internal/clients/volces.go
+2
View File
@@ -418,6 +418,8 @@ func cleanProviderBody(body map[string]any) map[string]any {
"platformId",
"platform_model_id",
"platformModelId",
"modelType",
"model_type",
} {
delete(out, key)
}
apps/api/internal/httpapi/acceptance_handlers.go
+300
View File
@@ -0,0 +1,300 @@
package httpapi
import (
"encoding/json"
"errors"
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const (
acceptanceRunHeader = "X-EasyAI-Acceptance-Run"
acceptanceTokenHeader = "X-EasyAI-Acceptance-Token"
acceptanceUpstreamHeader = "X-EasyAI-Acceptance-Upstream"
)
type taskTrafficAdmission struct {
RunMode string
AcceptanceRunID string
}
type taskTrafficError struct {
Status int
Code string
Message string
Err error
}
func (e *taskTrafficError) Error() string {
return e.Message
}
func (e *taskTrafficError) ErrorCode() string {
return e.Code
}
func (s *Server) authorizeTaskTraffic(r *http.Request, user *auth.User) (taskTrafficAdmission, error) {
runID, err := s.store.AuthorizeAcceptanceTask(
r.Context(),
r.Header.Get(acceptanceRunHeader),
r.Header.Get(acceptanceTokenHeader),
user,
)
if err != nil {
switch {
case errors.Is(err, store.ErrProductionTrafficPaused):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "validation_in_progress",
Message: "new production tasks are paused while validation is running", Err: err,
}
case errors.Is(err, store.ErrAcceptanceNotAuthorized):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
Message: "acceptance credentials do not match the active run", Err: err,
}
case errors.Is(err, store.ErrAcceptanceRunNotActive):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusConflict, Code: "acceptance_run_not_active",
Message: "acceptance headers are not valid while live traffic is enabled", Err: err,
}
default:
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
Message: "task traffic gate is unavailable", Err: err,
}
}
}
if runID != "" {
runMode := "acceptance"
if strings.EqualFold(strings.TrimSpace(r.Header.Get(acceptanceUpstreamHeader)), "real") {
runMode = "acceptance_canary"
}
return taskTrafficAdmission{RunMode: runMode, AcceptanceRunID: runID}, nil
}
return taskTrafficAdmission{RunMode: "production"}, nil
}
func (s *Server) admittedTaskRunMode(admission taskTrafficAdmission, body map[string]any) (string, error) {
if admission.RunMode == "acceptance" || admission.RunMode == "acceptance_canary" {
return admission.RunMode, nil
}
requested := runModeFromRequest(body)
requestedMode := strings.ToLower(strings.TrimSpace(requested))
if requestedMode == "acceptance" || requestedMode == "acceptance_canary" {
return "", &taskTrafficError{
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
Message: "acceptance mode is available only through an active acceptance run",
}
}
if strings.EqualFold(strings.TrimSpace(s.cfg.AppEnv), "production") && requestedMode == "simulation" {
return "", &taskTrafficError{
Status: http.StatusForbidden, Code: "simulation_not_authorized",
Message: "simulation mode is available only through an active acceptance run",
}
}
return requested, nil
}
func writeTaskTrafficError(w http.ResponseWriter, err error, protocol string) {
var trafficErr *taskTrafficError
if !errors.As(err, &trafficErr) {
trafficErr = &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
Message: "task traffic gate is unavailable", Err: err,
}
}
if protocol != "" {
writeProtocolError(w, protocol, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
return
}
writeErrorWithDetails(w, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
}
// getGatewayTrafficMode godoc
// @Summary 获取 Gateway 流量模式
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/traffic-mode [get]
func (s *Server) getGatewayTrafficMode(w http.ResponseWriter, r *http.Request) {
mode, err := s.store.GetGatewayTrafficMode(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "get gateway traffic mode failed")
return
}
writeJSON(w, http.StatusOK, mode)
}
// createAcceptanceRun godoc
// @Summary 创建生产同构验收 Run
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param body body store.CreateAcceptanceRunInput true "验收 Run"
// @Success 201 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs [post]
func (s *Server) createAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.CreateAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
run, err := s.store.CreateAcceptanceRun(r.Context(), input)
if err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
}
writeJSON(w, http.StatusCreated, run)
}
// getAcceptanceRun godoc
// @Summary 获取生产同构验收 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID} [get]
func (s *Server) getAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.store.GetAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "acceptance run not found")
return
}
writeError(w, http.StatusInternalServerError, "get acceptance run failed")
return
}
writeJSON(w, http.StatusOK, run)
}
// activateAcceptanceRun godoc
// @Summary 切换到 validation 并激活 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/activate [post]
func (s *Server) activateAcceptanceRun(w http.ResponseWriter, r *http.Request) {
mode, err := s.store.ActivateAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// finishAcceptanceRun godoc
// @Summary 记录验收门禁结果
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.FinishAcceptanceRunInput true "门禁结果"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/finish [post]
func (s *Server) finishAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.FinishAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
run, err := s.store.FinishAcceptanceRun(r.Context(), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusConflict, "acceptance run is not running")
return
}
writeError(w, http.StatusInternalServerError, "finish acceptance run failed")
return
}
writeJSON(w, http.StatusOK, run)
}
// retryAcceptanceRun godoc
// @Summary 在 validation 关闭正式流量的状态下重试失败 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/retry [post]
func (s *Server) retryAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.store.RetryAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, run)
}
// promoteAcceptanceRun godoc
// @Summary 通过 CAS 切回 live
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/promote [post]
func (s *Server) promoteAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.PromoteAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
mode, err := s.store.PromoteAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// abortAcceptanceRun godoc
// @Summary 人工 CAS 中止验收并切回 live
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/abort [post]
func (s *Server) abortAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.PromoteAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
mode, err := s.store.AbortAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
func writeAcceptanceMutationError(w http.ResponseWriter, err error) {
switch {
case errors.Is(err, store.ErrAcceptanceStateConflict):
writeError(w, http.StatusConflict, err.Error())
case errors.Is(err, store.ErrAcceptancePromotionGates):
writeError(w, http.StatusPreconditionFailed, err.Error())
case store.IsNotFound(err):
writeError(w, http.StatusNotFound, "acceptance run not found")
default:
writeError(w, http.StatusInternalServerError, "acceptance state update failed")
}
}
apps/api/internal/httpapi/acceptance_handlers_test.go
+27
View File
@@ -0,0 +1,27 @@
package httpapi
import (
"errors"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
)
func TestAdmittedTaskRunModeReservesAcceptanceAndProductionSimulation(t *testing.T) {
server := &Server{cfg: config.Config{AppEnv: "production"}}
for _, runMode := range []string{"simulation", "SIMULATION", "acceptance", "acceptance_canary"} {
_, err := server.admittedTaskRunMode(taskTrafficAdmission{}, map[string]any{"runMode": runMode})
var trafficErr *taskTrafficError
if !errors.As(err, &trafficErr) || trafficErr.Status != 403 {
t.Fatalf("run mode %q error=%v", runMode, err)
}
}
for _, runMode := range []string{"acceptance", "acceptance_canary"} {
got, err := server.admittedTaskRunMode(taskTrafficAdmission{RunMode: runMode}, map[string]any{
"runMode": "production",
})
if err != nil || got != runMode {
t.Fatalf("admitted mode=%q got=%q err=%v", runMode, got, err)
}
}
}
apps/api/internal/httpapi/gateway_task_creation.go
+10 -1
View File
@@ -44,14 +44,23 @@ func (s *Server) prepareAndCreateGatewayTask(
body map[string]any,
async bool,
) (store.GatewayTask, error) {
admission, err := s.authorizeTaskTraffic(r, user)
if err != nil {
return store.GatewayTask{}, &gatewayTaskCreationError{Stage: gatewayTaskCreationPrepare, Err: err}
}
prepared, err := s.prepareTaskRequest(ctx, r, user, body)
if err != nil {
return store.GatewayTask{}, &gatewayTaskCreationError{Stage: gatewayTaskCreationPrepare, Err: err}
}
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
if err != nil {
return store.GatewayTask{}, &gatewayTaskCreationError{Stage: gatewayTaskCreationPrepare, Err: err}
}
task, err := s.store.CreateTask(ctx, store.CreateTaskInput{
Kind: kind,
Model: model,
RunMode: runModeFromRequest(prepared.Body),
RunMode: runMode,
AcceptanceRunID: admission.AcceptanceRunID,
Async: async,
Request: prepared.Body,
ConversationID: prepared.ConversationID,
apps/api/internal/httpapi/gemini_base64_stress_integration_test.go
+6 -4
View File
@@ -22,6 +22,7 @@ import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/jackc/pgx/v5/pgxpool"
@@ -37,13 +38,14 @@ func TestGeminiBase64ImageEditThousandSynchronousRequests(t *testing.T) {
}
const (
taskCount = 1000
upstreamConcurrent = 64
inputVariants = 16
mediaConcurrent = 8
)
imageBytes := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_IMAGE_BYTES", 256<<10)
upstreamDelay := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_UPSTREAM_DELAY_MS", 4000)) * time.Millisecond
baselineProfile := acceptanceworkload.GeminiBaseline
taskCount := baselineProfile.Requests
imageBytes := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_IMAGE_BYTES", baselineProfile.InputBytes)
upstreamDelay := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_UPSTREAM_DELAY_MS", int(baselineProfile.DelayMin/time.Millisecond))) * time.Millisecond
maxHeapGrowth := uint64(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_MAX_HEAP_BYTES", 1<<30))
clientConnections := stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_CLIENT_CONNECTIONS", 256)
testTimeout := time.Duration(stressEnvInt(t, "AI_GATEWAY_GEMINI_STRESS_TIMEOUT_SECONDS", 1200)) * time.Second
@@ -387,7 +389,7 @@ WHERE task.model = $1
if tasks != taskCount || succeeded != taskCount || attempts != taskCount || duplicateAttemptTasks != 0 {
t.Fatalf("task results tasks=%d succeeded=%d attempts=%d duplicate_attempt_tasks=%d", tasks, succeeded, attempts, duplicateAttemptTasks)
}
if upstreamCalls.Load() != taskCount || invalidInputs.Load() != 0 {
if upstreamCalls.Load() != int64(taskCount) || invalidInputs.Load() != 0 {
t.Fatalf("upstream calls=%d invalid_inputs=%d, want %d/0", upstreamCalls.Load(), invalidInputs.Load(), taskCount)
}
if upstreamPeak.Load() > mediaConcurrent || upstreamPeak.Load() < mediaConcurrent/2 {
apps/api/internal/httpapi/gemini_compat.go
+12 -1
View File
@@ -117,6 +117,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
writeGeminiTaskError(http.StatusUnauthorized, "unauthorized", nil, "unauthorized")
return
}
admission, err := s.authorizeTaskTraffic(r, user)
if err != nil {
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
return
}
releaseRequestBody, err := s.acquireMediaRequestBodySlot(r.Context())
if err != nil {
return
@@ -166,10 +171,16 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
mapping.Body = nil
releaseRequestBody()
releaseRequestBody = nil
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
if err != nil {
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
return
}
createInput := store.CreateTaskInput{
Kind: mapping.Kind,
Model: mapping.Model,
RunMode: runModeFromRequest(prepared.Body),
RunMode: runMode,
AcceptanceRunID: admission.AcceptanceRunID,
Async: false,
Request: prepared.Body,
ConversationID: prepared.ConversationID,
apps/api/internal/httpapi/handlers.go
+12 -1
View File
@@ -1095,6 +1095,11 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeTaskError(http.StatusUnauthorized, "unauthorized", nil)
return
}
admission, err := s.authorizeTaskTraffic(r, user)
if err != nil {
writeTaskTrafficError(w, err, targetProtocol)
return
}
body, err := s.decodeTaskRequestBody(r.Context(), w, r, kind)
if err != nil {
@@ -1147,10 +1152,16 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
return
}
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
if err != nil {
writeTaskTrafficError(w, err, targetProtocol)
return
}
createInput := store.CreateTaskInput{
Kind: kind,
Model: model,
RunMode: runModeFromRequest(prepared.Body),
RunMode: runMode,
AcceptanceRunID: admission.AcceptanceRunID,
Async: responsePlan.asyncMode,
Request: prepared.Body,
ConversationID: prepared.ConversationID,
apps/api/internal/httpapi/kling_compat.go
+25 -6
View File
@@ -112,6 +112,15 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
writeKlingCompatError(w, http.StatusForbidden, "api key scope does not allow video generation", "permission_denied")
return
}
admission, err := s.authorizeTaskTraffic(r, user)
if err != nil {
targetProtocol := clients.ProtocolKlingV1Omni
if version == "v2" {
targetProtocol = clients.ProtocolKlingV2Omni
}
writeTaskTrafficError(w, err, targetProtocol)
return
}
body, externalTaskID, err := klingCompatTaskBody(version, model, native)
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
@@ -122,13 +131,23 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
if err != nil {
targetProtocol := clients.ProtocolKlingV1Omni
if version == "v2" {
targetProtocol = clients.ProtocolKlingV2Omni
}
writeTaskTrafficError(w, err, targetProtocol)
return
}
createInput := store.CreateTaskInput{
Kind: "videos.generations",
Model: model,
ExternalTaskID: externalTaskID,
RunMode: runModeFromRequest(prepared.Body),
Async: true,
Request: prepared.Body,
Kind: "videos.generations",
Model: model,
ExternalTaskID: externalTaskID,
RunMode: runMode,
AcceptanceRunID: admission.AcceptanceRunID,
Async: true,
Request: prepared.Body,
}
if idempotencyKey, hasKey, keyErr := optionalTaskIdempotencyKey(r); keyErr != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid Idempotency-Key", "invalid_idempotency_key")
apps/api/internal/httpapi/server.go
+8
View File
@@ -217,6 +217,14 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("PATCH /api/admin/system/file-storage/settings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateFileStorageSettings)))
mux.Handle("GET /api/admin/system/client-customization/settings", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getClientCustomizationSettings)))
mux.Handle("PATCH /api/admin/system/client-customization/settings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateClientCustomizationSettings)))
mux.Handle("GET /api/admin/system/acceptance/traffic-mode", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getGatewayTrafficMode)))
mux.Handle("POST /api/admin/system/acceptance/runs", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.createAcceptanceRun)))
mux.Handle("GET /api/admin/system/acceptance/runs/{runID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/activate", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.activateAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/finish", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.finishAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/retry", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.retryAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/promote", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.promoteAcceptanceRun)))
mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/abort", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.abortAcceptanceRun)))
mux.Handle("GET /api/admin/system/identity/configuration", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityConfiguration)))
mux.Handle("POST /api/admin/system/identity/pairings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.startIdentityPairing)))
mux.Handle("GET /api/admin/system/identity/pairings/{pairingID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityPairing)))
apps/api/internal/httpapi/volces_compat_handlers.go
+5
View File
@@ -416,6 +416,11 @@ func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
if errors.As(err, &staged) {
err = staged.Err
}
var trafficErr *taskTrafficError
if errors.As(err, &trafficErr) {
writeVolcesError(w, trafficErr.Status, trafficErr.Message, trafficErr.Code)
return
}
var clientErr *clients.ClientError
if errors.As(err, &clientErr) && clientErr.StatusCode > 0 {
status = clientErr.StatusCode
apps/api/internal/httpapi/volces_compat_handlers_test.go
+25
View File
@@ -1,6 +1,9 @@
package httpapi
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
@@ -35,3 +38,25 @@ func TestVolcesCompatibleTaskBuildsProtocolResponseFromCanonicalResult(t *testin
}
}
}
func TestVolcesCompatibilityPreservesValidationGateStatus(t *testing.T) {
recorder := httptest.NewRecorder()
writeVolcesCompatibleTaskError(recorder, &gatewayTaskCreationError{
Stage: gatewayTaskCreationPrepare,
Err: &taskTrafficError{
Status: http.StatusServiceUnavailable,
Code: "validation_in_progress",
Message: "new production tasks are paused while validation is running",
},
})
if recorder.Code != http.StatusServiceUnavailable {
t.Fatalf("status=%d body=%s", recorder.Code, recorder.Body.String())
}
var payload VolcesErrorEnvelope
if err := json.NewDecoder(recorder.Body).Decode(&payload); err != nil {
t.Fatalf("decode error: %v", err)
}
if payload.Error.Code != "validation_in_progress" {
t.Fatalf("error=%+v", payload.Error)
}
}
apps/api/internal/runner/billing_mode_test.go
+11
View File
@@ -25,3 +25,14 @@ func TestBillingItemsFixedTotalKeepsNineDecimalPlaces(t *testing.T) {
t.Fatalf("total = %s", got)
}
}
func TestAcceptanceRunModesUseRealBilling(t *testing.T) {
for _, runMode := range []string{"production", "acceptance", "acceptance_canary"} {
if !isBillableRunMode(runMode) {
t.Fatalf("run mode %q must be billable", runMode)
}
}
if isBillableRunMode("simulation") {
t.Fatal("simulation must remain non-billable")
}
}
apps/api/internal/runner/service.go
+57 -6
View File
@@ -418,7 +418,7 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
preprocessingByCandidate := map[string]parameterPreprocessResult{}
reservationBillings := []any(nil)
reservationPricingSnapshot := map[string]any(nil)
if task.RunMode == "production" {
if isBillableRunMode(task.RunMode) {
billingMode := normalizedBillingEngineMode(s.cfg.BillingEngineMode)
if billingMode == "hold" {
holdErr := &clients.ClientError{Code: "billing_hold", Message: "production billing is temporarily on hold", Retryable: true}
@@ -809,7 +809,7 @@ candidatesLoop:
finalAmount := fixedAmount(0)
finalAmountText := ""
pricingSnapshot := map[string]any{"pricingVersion": pricingVersionV2, "source": "simulation"}
if task.RunMode == "production" {
if isBillableRunMode(task.RunMode) {
billingMode := normalizedBillingEngineMode(s.cfg.BillingEngineMode)
var billingErr error
if billingMode == "observe" {
@@ -1171,9 +1171,16 @@ func (s *Service) runCandidate(
admittedLeases []store.ConcurrencyLease,
) (clients.Response, error) {
var err error
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return clients.Response{}, err
if task.RunMode == "acceptance" {
candidate, err = s.acceptanceCandidate(ctx, task, candidate)
if err != nil {
return clients.Response{}, err
}
} else {
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return clients.Response{}, err
}
}
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
@@ -2085,7 +2092,14 @@ func (s *Service) emit(ctx context.Context, taskID string, eventType string, sta
if event.ID == "" {
s.observeTaskEventSkip(event.SkippedReason)
}
if s.cfg.TaskProgressCallbackEnabled {
acceptanceURL, acceptance, callbackErr := s.store.AcceptanceCallbackURL(ctx, taskID)
if callbackErr != nil {
return callbackErr
}
if acceptance && acceptanceURL != "" {
return s.store.QueueTaskCallback(ctx, event, acceptanceURL)
}
if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
return s.store.QueueTaskCallback(ctx, event, s.cfg.TaskProgressCallbackURL)
}
return nil
@@ -2230,6 +2244,43 @@ func isSimulation(task store.GatewayTask, candidate store.RuntimeModelCandidate)
return stringFromMap(candidate.Credentials, "mode") == "simulation" || boolFromMap(candidate.PlatformConfig, "testMode")
}
func isBillableRunMode(runMode string) bool {
return runMode == "production" || runMode == "acceptance" || runMode == "acceptance_canary"
}
func (s *Service) acceptanceCandidate(
ctx context.Context,
task store.GatewayTask,
candidate store.RuntimeModelCandidate,
) (store.RuntimeModelCandidate, error) {
if strings.TrimSpace(task.AcceptanceRunID) == "" {
return candidate, &clients.ClientError{
Code: "acceptance_run_missing",
Message: "acceptance task is missing its run identifier",
Retryable: false,
}
}
baseURL, credential, err := s.store.AcceptanceCandidateOverride(ctx, task.AcceptanceRunID)
if err != nil {
return candidate, &clients.ClientError{
Code: "acceptance_run_inactive",
Message: err.Error(),
Retryable: false,
}
}
candidate.BaseURL = baseURL
candidate.Credentials = map[string]any{"apiKey": credential}
platformConfig := cloneMap(candidate.PlatformConfig)
delete(platformConfig, "credentialEnv")
delete(platformConfig, "credential_env")
delete(platformConfig, "proxyMode")
delete(platformConfig, "httpProxy")
platformConfig["networkProxy"] = map[string]any{"mode": "none"}
platformConfig["acceptanceRunId"] = task.AcceptanceRunID
candidate.PlatformConfig = platformConfig
return candidate, nil
}
func retryEnabled(candidate store.RuntimeModelCandidate) bool {
policy := effectiveRetryPolicy(candidate)
if enabled, ok := policy["enabled"].(bool); ok {
apps/api/internal/securityevents/metrics.go
+12
View File
@@ -279,6 +279,12 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
return
}
}
postgresPool := store.PostgresPoolMetricsSnapshot{}
if poolProvider, ok := provider.(interface {
PostgresPoolMetrics() store.PostgresPoolMetricsSnapshot
}); ok {
postgresPool = poolProvider.PostgresPoolMetrics()
}
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
outcomeCounters(w, "easyai_gateway_ssf_receipts_total", "Received SETs by bounded outcome.", []outcomeValue{
{"accepted", m.accepted.Load()}, {"rejected", m.rejected.Load()}, {"duplicate", m.duplicate.Load()},
@@ -350,6 +356,12 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
plainGauge(w, "easyai_gateway_worker_active_instances", "Active distributed worker instances with a fresh heartbeat.", m.asyncWorkerActiveInstances.Load())
plainGauge(w, "easyai_gateway_worker_global_capacity", "Global asynchronous execution capacity before instance allocation.", m.asyncWorkerGlobalCapacity.Load())
plainGauge(w, "easyai_gateway_worker_allocated_capacity", "Asynchronous execution capacity allocated to this worker instance.", m.asyncWorkerAllocated.Load())
plainGauge(w, "easyai_gateway_postgres_pool_max_connections", "Maximum PostgreSQL connections in this process pool.", int64(postgresPool.MaxConnections))
plainGauge(w, "easyai_gateway_postgres_pool_total_connections", "Current PostgreSQL connections in this process pool.", int64(postgresPool.TotalConnections))
plainGauge(w, "easyai_gateway_postgres_pool_acquired_connections", "Currently acquired PostgreSQL connections in this process pool.", int64(postgresPool.AcquiredConnections))
plainGauge(w, "easyai_gateway_postgres_pool_idle_connections", "Currently idle PostgreSQL connections in this process pool.", int64(postgresPool.IdleConnections))
plainCounter(w, "easyai_gateway_postgres_pool_empty_acquire_total", "Pool acquires that had to wait for a PostgreSQL connection.", uint64(max(postgresPool.EmptyAcquireCount, 0)))
plainCounter(w, "easyai_gateway_postgres_pool_canceled_acquire_total", "Canceled PostgreSQL pool acquires.", uint64(max(postgresPool.CanceledAcquireCount, 0)))
outcomeCounters(w, "easyai_gateway_async_worker_resizes_total", "Asynchronous worker resize attempts by bounded outcome.", []outcomeValue{
{"success", m.asyncWorkerResizeSuccess.Load()},
{"refresh_failed", m.asyncWorkerRefreshFailed.Load()},
apps/api/internal/securityevents/metrics_test.go
+21 -1
View File
@@ -7,17 +7,24 @@ import (
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type metricsSnapshot struct {
mode string
last time.Time
pool store.PostgresPoolMetricsSnapshot
}
func (m metricsSnapshot) SecurityEventMetrics(context.Context, string, string) (string, *time.Time, error) {
return m.mode, &m.last, nil
}
func (m metricsSnapshot) PostgresPoolMetrics() store.PostgresPoolMetricsSnapshot {
return m.pool
}
func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
metrics := &Metrics{}
metrics.ObserveReceipt("accepted", 20*time.Millisecond, 2)
@@ -34,7 +41,14 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
metrics.ObserveTaskEventSkip("budget_exceeded")
recorder := httptest.NewRecorder()
metrics.Handler(metricsSnapshot{mode: "push_healthy", last: time.Now().Add(-time.Minute)}, "issuer", "audience", true).
metrics.Handler(metricsSnapshot{
mode: "push_healthy",
last: time.Now().Add(-time.Minute),
pool: store.PostgresPoolMetricsSnapshot{
MaxConnections: 32, TotalConnections: 18, AcquiredConnections: 12,
IdleConnections: 6, EmptyAcquireCount: 3, CanceledAcquireCount: 1,
},
}, "issuer", "audience", true).
ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/metrics", nil))
if recorder.Code != http.StatusOK {
t.Fatalf("metrics status=%d", recorder.Code)
@@ -56,6 +70,12 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
`easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`,
`easyai_gateway_task_events_skipped_total{outcome="duplicate"} 1`,
`easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 1`,
`easyai_gateway_postgres_pool_max_connections 32`,
`easyai_gateway_postgres_pool_total_connections 18`,
`easyai_gateway_postgres_pool_acquired_connections 12`,
`easyai_gateway_postgres_pool_idle_connections 6`,
`easyai_gateway_postgres_pool_empty_acquire_total 3`,
`easyai_gateway_postgres_pool_canceled_acquire_total 1`,
} {
if !strings.Contains(recorder.Body.String(), expected) {
t.Fatalf("missing metric %q in:\n%s", expected, recorder.Body.String())
apps/api/internal/store/acceptance.go
+533
View File
@@ -0,0 +1,533 @@
package store
import (
"context"
"crypto/sha256"
"crypto/subtle"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"net/url"
"regexp"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
)
const SystemSettingGatewayTrafficMode = "gateway_traffic_mode"
var (
acceptanceReleaseSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
acceptanceDigestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
)
var (
ErrProductionTrafficPaused = errors.New("production traffic is paused for validation")
ErrAcceptanceNotAuthorized = errors.New("acceptance request is not authorized")
ErrAcceptanceRunNotActive = errors.New("acceptance run is not active")
ErrAcceptanceStateConflict = errors.New("acceptance state changed")
ErrAcceptancePromotionGates = errors.New("acceptance promotion gates have not passed")
)
type GatewayTrafficMode struct {
Mode string `json:"mode"`
RunID string `json:"runId,omitempty"`
Revision int64 `json:"revision"`
ReleaseSHA string `json:"releaseSha,omitempty"`
APIImageDigest string `json:"apiImageDigest,omitempty"`
WorkerImageDigest string `json:"workerImageDigest,omitempty"`
UpdatedAt time.Time `json:"updatedAt"`
}
type AcceptanceRun struct {
ID string `json:"id"`
Status string `json:"status"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
APIKeyID string `json:"apiKeyId"`
UserID string `json:"userId"`
EmulatorBaseURL string `json:"emulatorBaseUrl"`
CallbackURL string `json:"callbackUrl,omitempty"`
CapacityProfile string `json:"capacityProfile"`
Config map[string]any `json:"config"`
Report map[string]any `json:"report"`
FailureReason string `json:"failureReason,omitempty"`
StartedAt string `json:"startedAt,omitempty"`
FinishedAt string `json:"finishedAt,omitempty"`
PromotedAt string `json:"promotedAt,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type CreateAcceptanceRunInput struct {
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
APIKeyID string `json:"apiKeyId"`
UserID string `json:"userId"`
Token string `json:"token"`
EmulatorBaseURL string `json:"emulatorBaseUrl"`
CallbackURL string `json:"callbackUrl,omitempty"`
CapacityProfile string `json:"capacityProfile"`
Config map[string]any `json:"config,omitempty"`
}
type FinishAcceptanceRunInput struct {
RunID string `json:"-"`
Passed bool `json:"passed"`
Report map[string]any `json:"report"`
FailureReason string `json:"failureReason,omitempty"`
}
type PromoteAcceptanceRunInput struct {
RunID string `json:"-"`
Revision int64 `json:"revision"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
}
func DefaultGatewayTrafficMode() GatewayTrafficMode {
return GatewayTrafficMode{Mode: "live"}
}
func (s *Store) GetGatewayTrafficMode(ctx context.Context) (GatewayTrafficMode, error) {
var value []byte
var updatedAt time.Time
err := s.pool.QueryRow(ctx, `
SELECT value, updated_at
FROM system_settings
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode).Scan(&value, &updatedAt)
if err != nil {
if IsNotFound(err) || IsUndefinedDatabaseObject(err) {
return DefaultGatewayTrafficMode(), nil
}
return GatewayTrafficMode{}, err
}
var mode GatewayTrafficMode
if err := json.Unmarshal(value, &mode); err != nil {
return GatewayTrafficMode{}, fmt.Errorf("decode gateway traffic mode: %w", err)
}
mode.Mode = normalizeTrafficMode(mode.Mode)
mode.UpdatedAt = updatedAt
return mode, nil
}
func (s *Store) CreateAcceptanceRun(ctx context.Context, input CreateAcceptanceRunInput) (AcceptanceRun, error) {
input.ReleaseSHA = strings.TrimSpace(input.ReleaseSHA)
input.APIImageDigest = strings.TrimSpace(input.APIImageDigest)
input.WorkerImageDigest = strings.TrimSpace(input.WorkerImageDigest)
input.APIKeyID = strings.TrimSpace(input.APIKeyID)
input.UserID = strings.TrimSpace(input.UserID)
input.Token = strings.TrimSpace(input.Token)
input.EmulatorBaseURL = strings.TrimRight(strings.TrimSpace(input.EmulatorBaseURL), "/")
input.CallbackURL = strings.TrimSpace(input.CallbackURL)
input.CapacityProfile = strings.ToUpper(strings.TrimSpace(input.CapacityProfile))
if input.CapacityProfile == "" {
input.CapacityProfile = "P24"
}
if input.ReleaseSHA == "" || input.APIImageDigest == "" || input.WorkerImageDigest == "" ||
input.APIKeyID == "" || input.UserID == "" || input.Token == "" || input.EmulatorBaseURL == "" ||
input.CallbackURL == "" {
return AcceptanceRun{}, errors.New("release SHA, image digests, API key, user, token, emulator URL, and callback URL are required")
}
if !acceptanceReleaseSHAPattern.MatchString(input.ReleaseSHA) ||
!acceptanceDigestPattern.MatchString(input.APIImageDigest) ||
!acceptanceDigestPattern.MatchString(input.WorkerImageDigest) {
return AcceptanceRun{}, errors.New("release SHA and image digests must use full immutable identifiers")
}
if len(input.Token) < 32 {
return AcceptanceRun{}, errors.New("acceptance token must contain at least 32 characters")
}
if !validAcceptanceURL(input.EmulatorBaseURL) || !validAcceptanceURL(input.CallbackURL) {
return AcceptanceRun{}, errors.New("emulator and callback URLs must be absolute HTTP URLs without user information")
}
switch input.CapacityProfile {
case "P24", "P28", "P32":
default:
return AcceptanceRun{}, errors.New("capacity profile must be P24, P28, or P32")
}
config, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(input.Config)))
tokenHash := acceptanceTokenSHA256(input.Token)
return scanAcceptanceRun(s.pool.QueryRow(ctx, `
INSERT INTO gateway_acceptance_runs (
release_sha, api_image_digest, worker_image_digest, api_key_id, user_id,
token_sha256, emulator_base_url, callback_url, capacity_profile, config
)
VALUES ($1, $2, $3, $4, $5, $6, $7, NULLIF($8, ''), $9, $10::jsonb)
RETURNING `+acceptanceRunColumns,
input.ReleaseSHA, input.APIImageDigest, input.WorkerImageDigest, input.APIKeyID, input.UserID,
tokenHash, input.EmulatorBaseURL, input.CallbackURL, input.CapacityProfile, config,
))
}
func (s *Store) GetAcceptanceRun(ctx context.Context, runID string) (AcceptanceRun, error) {
return scanAcceptanceRun(s.pool.QueryRow(ctx, `
SELECT `+acceptanceRunColumns+`
FROM gateway_acceptance_runs
WHERE id = $1::uuid`, strings.TrimSpace(runID)))
}
func (s *Store) ActivateAcceptanceRun(ctx context.Context, runID string) (GatewayTrafficMode, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return GatewayTrafficMode{}, err
}
defer rollbackTransaction(tx)
var currentValue []byte
if err := tx.QueryRow(ctx, `
SELECT value
FROM system_settings
WHERE setting_key = $1
FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&currentValue); err != nil {
return GatewayTrafficMode{}, err
}
var current GatewayTrafficMode
if err := json.Unmarshal(currentValue, &current); err != nil {
return GatewayTrafficMode{}, err
}
if normalizeTrafficMode(current.Mode) != "live" {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
run, err := scanAcceptanceRun(tx.QueryRow(ctx, `
SELECT `+acceptanceRunColumns+`
FROM gateway_acceptance_runs
WHERE id = $1::uuid
FOR UPDATE`, strings.TrimSpace(runID)))
if err != nil {
return GatewayTrafficMode{}, err
}
if run.Status != "pending" && run.Status != "failed" {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
next := GatewayTrafficMode{
Mode: "validation",
RunID: run.ID,
Revision: current.Revision + 1,
ReleaseSHA: run.ReleaseSHA,
APIImageDigest: run.APIImageDigest,
WorkerImageDigest: run.WorkerImageDigest,
}
value, _ := json.Marshal(next)
if _, err := tx.Exec(ctx, `
UPDATE system_settings
SET value = $2::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode, value); err != nil {
return GatewayTrafficMode{}, err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_acceptance_runs
SET status = 'running', started_at = COALESCE(started_at, now()),
finished_at = NULL, failure_reason = NULL, updated_at = now()
WHERE id = $1::uuid`, run.ID); err != nil {
return GatewayTrafficMode{}, err
}
if err := tx.Commit(ctx); err != nil {
return GatewayTrafficMode{}, err
}
next.UpdatedAt = time.Now()
return next, nil
}
func (s *Store) FinishAcceptanceRun(ctx context.Context, input FinishAcceptanceRunInput) (AcceptanceRun, error) {
status := "failed"
if input.Passed {
status = "passed"
}
report, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(input.Report)))
return scanAcceptanceRun(s.pool.QueryRow(ctx, `
UPDATE gateway_acceptance_runs
SET status = $2, report = $3::jsonb, failure_reason = NULLIF($4, ''),
finished_at = now(), updated_at = now()
WHERE id = $1::uuid AND status = 'running'
RETURNING `+acceptanceRunColumns,
strings.TrimSpace(input.RunID), status, report, strings.TrimSpace(input.FailureReason),
))
}
func (s *Store) RetryAcceptanceRun(ctx context.Context, runID string) (AcceptanceRun, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return AcceptanceRun{}, err
}
defer rollbackTransaction(tx)
var value []byte
if err := tx.QueryRow(ctx, `
SELECT value
FROM system_settings
WHERE setting_key = $1
FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&value); err != nil {
return AcceptanceRun{}, err
}
var mode GatewayTrafficMode
if err := json.Unmarshal(value, &mode); err != nil {
return AcceptanceRun{}, err
}
if normalizeTrafficMode(mode.Mode) != "validation" || mode.RunID != strings.TrimSpace(runID) {
return AcceptanceRun{}, ErrAcceptanceStateConflict
}
run, err := scanAcceptanceRun(tx.QueryRow(ctx, `
UPDATE gateway_acceptance_runs
SET status = 'running', report = '{}'::jsonb, failure_reason = NULL,
finished_at = NULL, updated_at = now()
WHERE id = $1::uuid AND status = 'failed'
RETURNING `+acceptanceRunColumns, strings.TrimSpace(runID)))
if err != nil {
return AcceptanceRun{}, err
}
if err := tx.Commit(ctx); err != nil {
return AcceptanceRun{}, err
}
return run, nil
}
func (s *Store) PromoteAcceptanceRun(ctx context.Context, input PromoteAcceptanceRunInput) (GatewayTrafficMode, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return GatewayTrafficMode{}, err
}
defer rollbackTransaction(tx)
var value []byte
if err := tx.QueryRow(ctx, `
SELECT value
FROM system_settings
WHERE setting_key = $1
FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&value); err != nil {
return GatewayTrafficMode{}, err
}
var current GatewayTrafficMode
if err := json.Unmarshal(value, &current); err != nil {
return GatewayTrafficMode{}, err
}
if normalizeTrafficMode(current.Mode) != "validation" ||
current.RunID != strings.TrimSpace(input.RunID) ||
current.Revision != input.Revision ||
current.ReleaseSHA != strings.TrimSpace(input.ReleaseSHA) ||
current.APIImageDigest != strings.TrimSpace(input.APIImageDigest) ||
current.WorkerImageDigest != strings.TrimSpace(input.WorkerImageDigest) {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
var status string
if err := tx.QueryRow(ctx, `
SELECT status
FROM gateway_acceptance_runs
WHERE id = $1::uuid
FOR UPDATE`, current.RunID).Scan(&status); err != nil {
return GatewayTrafficMode{}, err
}
if status != "passed" {
return GatewayTrafficMode{}, ErrAcceptancePromotionGates
}
next := GatewayTrafficMode{Mode: "live", Revision: current.Revision + 1}
nextValue, _ := json.Marshal(next)
if _, err := tx.Exec(ctx, `
UPDATE system_settings
SET value = $2::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode, nextValue); err != nil {
return GatewayTrafficMode{}, err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_acceptance_runs
SET status = 'promoted', promoted_at = now(), updated_at = now()
WHERE id = $1::uuid`, current.RunID); err != nil {
return GatewayTrafficMode{}, err
}
if err := tx.Commit(ctx); err != nil {
return GatewayTrafficMode{}, err
}
next.UpdatedAt = time.Now()
return next, nil
}
func (s *Store) AbortAcceptanceRun(ctx context.Context, input PromoteAcceptanceRunInput) (GatewayTrafficMode, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return GatewayTrafficMode{}, err
}
defer rollbackTransaction(tx)
var value []byte
if err := tx.QueryRow(ctx, `
SELECT value
FROM system_settings
WHERE setting_key = $1
FOR UPDATE`, SystemSettingGatewayTrafficMode).Scan(&value); err != nil {
return GatewayTrafficMode{}, err
}
var current GatewayTrafficMode
if err := json.Unmarshal(value, &current); err != nil {
return GatewayTrafficMode{}, err
}
if normalizeTrafficMode(current.Mode) != "validation" ||
current.RunID != strings.TrimSpace(input.RunID) ||
current.Revision != input.Revision ||
current.ReleaseSHA != strings.TrimSpace(input.ReleaseSHA) ||
current.APIImageDigest != strings.TrimSpace(input.APIImageDigest) ||
current.WorkerImageDigest != strings.TrimSpace(input.WorkerImageDigest) {
return GatewayTrafficMode{}, ErrAcceptanceStateConflict
}
next := GatewayTrafficMode{Mode: "live", Revision: current.Revision + 1}
nextValue, _ := json.Marshal(next)
if _, err := tx.Exec(ctx, `
UPDATE system_settings
SET value = $2::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode, nextValue); err != nil {
return GatewayTrafficMode{}, err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_acceptance_runs
SET status = 'aborted', finished_at = COALESCE(finished_at, now()), updated_at = now()
WHERE id = $1::uuid AND status IN ('running', 'failed', 'passed')`, current.RunID); err != nil {
return GatewayTrafficMode{}, err
}
if err := tx.Commit(ctx); err != nil {
return GatewayTrafficMode{}, err
}
next.UpdatedAt = time.Now()
return next, nil
}
func (s *Store) AuthorizeAcceptanceTask(ctx context.Context, runID string, token string, user *auth.User) (string, error) {
mode, err := s.GetGatewayTrafficMode(ctx)
if err != nil {
return "", err
}
runID = strings.TrimSpace(runID)
token = strings.TrimSpace(token)
if mode.Mode == "live" {
if runID != "" || token != "" {
return "", ErrAcceptanceRunNotActive
}
return "", nil
}
if runID == "" || token == "" || runID != mode.RunID {
return "", ErrProductionTrafficPaused
}
if user == nil || strings.TrimSpace(user.APIKeyID) == "" {
return "", ErrAcceptanceNotAuthorized
}
var activeRunID, apiKeyID, userID, tokenHash, status string
err = s.pool.QueryRow(ctx, `
SELECT id::text, api_key_id, user_id, token_sha256, status
FROM gateway_acceptance_runs
WHERE id = $1::uuid`, runID).Scan(&activeRunID, &apiKeyID, &userID, &tokenHash, &status)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return "", ErrAcceptanceNotAuthorized
}
return "", err
}
if status != "running" || apiKeyID != strings.TrimSpace(user.APIKeyID) || userID != strings.TrimSpace(user.ID) ||
subtle.ConstantTimeCompare([]byte(tokenHash), []byte(acceptanceTokenSHA256(token))) != 1 {
return "", ErrAcceptanceNotAuthorized
}
return activeRunID, nil
}
func (s *Store) AcceptanceCandidateOverride(ctx context.Context, runID string) (string, string, error) {
var baseURL, status string
err := s.pool.QueryRow(ctx, `
SELECT emulator_base_url, status
FROM gateway_acceptance_runs
WHERE id = $1::uuid`, strings.TrimSpace(runID)).Scan(&baseURL, &status)
if err != nil {
return "", "", err
}
if status != "running" && status != "passed" {
return "", "", ErrAcceptanceRunNotActive
}
return strings.TrimRight(strings.TrimSpace(baseURL), "/"), "acceptance-protocol-emulator", nil
}
func (s *Store) AcceptanceCallbackURL(ctx context.Context, taskID string) (string, bool, error) {
var callbackURL string
err := s.pool.QueryRow(ctx, `
SELECT COALESCE(run.callback_url, '')
FROM gateway_tasks task
JOIN gateway_acceptance_runs run ON run.id = task.acceptance_run_id
WHERE task.id = $1::uuid`, strings.TrimSpace(taskID)).Scan(&callbackURL)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return "", false, nil
}
return "", false, err
}
return strings.TrimSpace(callbackURL), true, nil
}
func normalizeTrafficMode(value string) string {
if strings.EqualFold(strings.TrimSpace(value), "validation") {
return "validation"
}
return "live"
}
func acceptanceTokenSHA256(token string) string {
sum := sha256.Sum256([]byte(strings.TrimSpace(token)))
return hex.EncodeToString(sum[:])
}
func validAcceptanceURL(raw string) bool {
parsed, err := url.Parse(strings.TrimSpace(raw))
if err != nil || parsed.User != nil || parsed.Host == "" {
return false
}
return parsed.Scheme == "http" || parsed.Scheme == "https"
}
func sanitizeAcceptanceMetadata(value any) any {
switch typed := value.(type) {
case map[string]any:
next := make(map[string]any, len(typed))
for key, item := range typed {
normalized := strings.ToLower(strings.NewReplacer("_", "", "-", "", ".", "").Replace(key))
if strings.Contains(normalized, "password") ||
strings.Contains(normalized, "secret") ||
strings.Contains(normalized, "credential") ||
strings.Contains(normalized, "connectionstring") ||
strings.Contains(normalized, "databaseurl") ||
(strings.Contains(normalized, "token") && !strings.HasSuffix(normalized, "count")) {
next[key] = "[REDACTED]"
continue
}
next[key] = sanitizeAcceptanceMetadata(item)
}
return next
case []any:
next := make([]any, len(typed))
for index, item := range typed {
next[index] = sanitizeAcceptanceMetadata(item)
}
return next
default:
return value
}
}
const acceptanceRunColumns = `
id::text, status, release_sha, api_image_digest, worker_image_digest,
api_key_id, user_id, emulator_base_url, COALESCE(callback_url, ''),
capacity_profile, config, report, COALESCE(failure_reason, ''),
COALESCE(started_at::text, ''), COALESCE(finished_at::text, ''),
COALESCE(promoted_at::text, ''), created_at, updated_at`
func scanAcceptanceRun(scanner taskScanner) (AcceptanceRun, error) {
var run AcceptanceRun
var config, report []byte
err := scanner.Scan(
&run.ID, &run.Status, &run.ReleaseSHA, &run.APIImageDigest, &run.WorkerImageDigest,
&run.APIKeyID, &run.UserID, &run.EmulatorBaseURL, &run.CallbackURL,
&run.CapacityProfile, &config, &report, &run.FailureReason,
&run.StartedAt, &run.FinishedAt, &run.PromotedAt, &run.CreatedAt, &run.UpdatedAt,
)
if err != nil {
return AcceptanceRun{}, err
}
run.Config = decodeObject(config)
run.Report = decodeObject(report)
return run, nil
}
apps/api/internal/store/acceptance_integration_test.go
+136
View File
@@ -0,0 +1,136 @@
package store
import (
"context"
"encoding/json"
"errors"
"os"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
func TestAcceptanceTrafficGateAndCASPromotion(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run acceptance traffic integration test")
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
applyOIDCJITTestMigrations(t, ctx, databaseURL)
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
resetTrafficMode := func() {
_, _ = db.pool.Exec(context.Background(), `
UPDATE system_settings
SET value = '{"mode":"live","revision":0}'::jsonb, updated_at = now()
WHERE setting_key = $1`, SystemSettingGatewayTrafficMode)
}
resetTrafficMode()
t.Cleanup(resetTrafficMode)
token := "acceptance-token-" + time.Now().Format("150405.000000000")
run, err := db.CreateAcceptanceRun(ctx, CreateAcceptanceRunInput{
ReleaseSHA: strings.Repeat("a", 40),
APIImageDigest: "sha256:" + strings.Repeat("b", 64),
WorkerImageDigest: "sha256:" + strings.Repeat("c", 64),
APIKeyID: "acceptance-api-key",
UserID: "acceptance-user",
Token: token,
EmulatorBaseURL: "http://acceptance-emulator:8090",
CallbackURL: "http://acceptance-emulator:8090/callbacks",
CapacityProfile: "P24",
})
if err != nil {
t.Fatalf("create acceptance run: %v", err)
}
encoded, _ := json.Marshal(run)
if strings.Contains(string(encoded), token) {
t.Fatal("acceptance run response exposed the raw token")
}
mode, err := db.ActivateAcceptanceRun(ctx, run.ID)
if err != nil {
t.Fatalf("activate acceptance run: %v", err)
}
if mode.Mode != "validation" || mode.RunID != run.ID || mode.Revision != 1 {
t.Fatalf("unexpected validation mode: %+v", mode)
}
user := &auth.User{ID: "acceptance-user", APIKeyID: "acceptance-api-key"}
if _, err := db.AuthorizeAcceptanceTask(ctx, "", "", user); !errors.Is(err, ErrProductionTrafficPaused) {
t.Fatalf("production request error=%v, want traffic paused", err)
}
if _, err := db.AuthorizeAcceptanceTask(ctx, run.ID, "wrong-token", user); !errors.Is(err, ErrAcceptanceNotAuthorized) {
t.Fatalf("wrong acceptance token error=%v", err)
}
if authorizedRunID, err := db.AuthorizeAcceptanceTask(ctx, run.ID, token, user); err != nil || authorizedRunID != run.ID {
t.Fatalf("authorize acceptance task run=%q err=%v", authorizedRunID, err)
}
if baseURL, credential, err := db.AcceptanceCandidateOverride(ctx, run.ID); err != nil ||
baseURL != "http://acceptance-emulator:8090" || credential == "" || credential == token {
t.Fatalf("candidate override base=%q credential=%q err=%v", baseURL, credential, err)
}
if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
RunID: run.ID, Passed: true, Report: map[string]any{"queueFinal": 0},
}); err != nil {
t.Fatalf("finish acceptance run: %v", err)
}
promotion := PromoteAcceptanceRunInput{
RunID: run.ID, Revision: mode.Revision,
ReleaseSHA: mode.ReleaseSHA, APIImageDigest: mode.APIImageDigest, WorkerImageDigest: mode.WorkerImageDigest,
}
stale := promotion
stale.Revision++
if _, err := db.PromoteAcceptanceRun(ctx, stale); !errors.Is(err, ErrAcceptanceStateConflict) {
t.Fatalf("stale promotion error=%v, want state conflict", err)
}
live, err := db.PromoteAcceptanceRun(ctx, promotion)
if err != nil {
t.Fatalf("promote acceptance run: %v", err)
}
if live.Mode != "live" || live.Revision != 2 {
t.Fatalf("unexpected live mode: %+v", live)
}
failedRun, err := db.CreateAcceptanceRun(ctx, CreateAcceptanceRunInput{
ReleaseSHA: strings.Repeat("d", 40),
APIImageDigest: "sha256:" + strings.Repeat("e", 64),
WorkerImageDigest: "sha256:" + strings.Repeat("f", 64),
APIKeyID: "acceptance-api-key",
UserID: "acceptance-user",
Token: token + "-retry",
EmulatorBaseURL: "http://acceptance-emulator:8090",
CallbackURL: "http://acceptance-emulator:8090/callbacks",
CapacityProfile: "P24",
})
if err != nil {
t.Fatalf("create retry acceptance run: %v", err)
}
failedMode, err := db.ActivateAcceptanceRun(ctx, failedRun.ID)
if err != nil {
t.Fatalf("activate retry acceptance run: %v", err)
}
if _, err := db.FinishAcceptanceRun(ctx, FinishAcceptanceRunInput{
RunID: failedRun.ID, Passed: false, FailureReason: "load failed",
}); err != nil {
t.Fatalf("fail acceptance run: %v", err)
}
if retried, err := db.RetryAcceptanceRun(ctx, failedRun.ID); err != nil || retried.Status != "running" {
t.Fatalf("retry acceptance run=%+v err=%v", retried, err)
}
aborted, err := db.AbortAcceptanceRun(ctx, PromoteAcceptanceRunInput{
RunID: failedRun.ID, Revision: failedMode.Revision,
ReleaseSHA: failedMode.ReleaseSHA, APIImageDigest: failedMode.APIImageDigest,
WorkerImageDigest: failedMode.WorkerImageDigest,
})
if err != nil {
t.Fatalf("abort acceptance run: %v", err)
}
if aborted.Mode != "live" || aborted.Revision != failedMode.Revision+1 {
t.Fatalf("unexpected aborted mode: %+v", aborted)
}
}
apps/api/internal/store/acceptance_test.go
+37
View File
@@ -0,0 +1,37 @@
package store
import "testing"
func TestAcceptanceMetadataRedactsSecrets(t *testing.T) {
value := sanitizeAcceptanceMetadata(map[string]any{
"token": "raw-token",
"nested": map[string]any{
"databaseUrl": "postgresql://example.invalid/database",
"taskCount": 1200,
},
})
next := value.(map[string]any)
if next["token"] != "[REDACTED]" {
t.Fatalf("token=%v", next["token"])
}
nested := next["nested"].(map[string]any)
if nested["databaseUrl"] != "[REDACTED]" || nested["taskCount"] != 1200 {
t.Fatalf("nested=%v", nested)
}
}
func TestAcceptanceURLsMustBeAbsoluteAndCredentialFree(t *testing.T) {
for _, raw := range []string{
"http://acceptance-emulator.easyai.svc.cluster.local:8090",
"https://acceptance.example.invalid/callbacks",
} {
if !validAcceptanceURL(raw) {
t.Fatalf("valid URL rejected: %s", raw)
}
}
for _, raw := range []string{"relative/path", "ftp://example.invalid/file", "https://user:pass@example.invalid"} {
if validAcceptanceURL(raw) {
t.Fatalf("invalid URL accepted: %s", raw)
}
}
}
apps/api/internal/store/postgres.go
+18 -4
View File
@@ -545,6 +545,7 @@ type CreateTaskInput struct {
Model string `json:"model"`
ExternalTaskID string `json:"externalTaskId,omitempty"`
RunMode string `json:"runMode"`
AcceptanceRunID string `json:"acceptanceRunId,omitempty"`
Async bool `json:"async"`
Request map[string]any `json:"request"`
ConversationID string `json:"conversationId"`
@@ -564,6 +565,7 @@ type GatewayTask struct {
ExternalTaskID string `json:"externalTaskId,omitempty"`
Kind string `json:"kind"`
RunMode string `json:"runMode"`
AcceptanceRunID string `json:"acceptanceRunId,omitempty"`
UserID string `json:"userId"`
GatewayUserID string `json:"gatewayUserId,omitempty"`
UserSource string `json:"userSource,omitempty"`
@@ -629,7 +631,8 @@ type GatewayTask struct {
}
const gatewayTaskColumns = `
id::text, COALESCE(external_task_id, ''), kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
id::text, COALESCE(external_task_id, ''), kind, run_mode, COALESCE(acceptance_run_id::text, ''),
user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(api_key_id, ''), COALESCE(api_key_name, ''), COALESCE(api_key_prefix, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model,
@@ -2103,15 +2106,19 @@ func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput,
task, err := scanGatewayTask(tx.QueryRow(ctx, `
INSERT INTO gateway_tasks (
external_task_id, kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
external_task_id, kind, run_mode, acceptance_run_id, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
model, requested_model, request, async_mode, status, result, billings, conversation_id, new_message_count,
idempotency_key_hash, idempotency_request_hash, finished_at
)
VALUES (NULLIF($1, ''), $2, $3, $4, NULLIF($5, '')::uuid, COALESCE(NULLIF($6, ''), 'gateway'), NULLIF($7, '')::uuid, NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, ''), NULLIF($13, '')::uuid, NULLIF($14, ''), $15, $15, $16, $17, $18, $19::jsonb, $20::jsonb, NULLIF($21, '')::uuid, $22, NULLIF($23, ''), NULLIF($24, ''), NULL)
VALUES (NULLIF($1, ''), $2, $3, NULLIF($4, '')::uuid, $5, NULLIF($6, '')::uuid, COALESCE(NULLIF($7, ''), 'gateway'), NULLIF($8, '')::uuid, NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, ''), NULLIF($13, ''), NULLIF($14, '')::uuid, NULLIF($15, ''), $16, $16, $17, $18, $19, $20::jsonb, $21::jsonb, NULLIF($22, '')::uuid, $23, NULLIF($24, ''), NULLIF($25, ''), NULL)
ON CONFLICT (user_id, idempotency_key_hash) WHERE idempotency_key_hash IS NOT NULL DO NOTHING
RETURNING `+gatewayTaskColumns,
strings.TrimSpace(input.ExternalTaskID), input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, strings.TrimSpace(input.IdempotencyKeyHash), strings.TrimSpace(input.IdempotencyRequestHash),
strings.TrimSpace(input.ExternalTaskID), input.Kind, runMode, strings.TrimSpace(input.AcceptanceRunID),
user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey,
user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey,
input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID,
input.NewMessageCount, strings.TrimSpace(input.IdempotencyKeyHash), strings.TrimSpace(input.IdempotencyRequestHash),
))
replayed := false
if errors.Is(err, pgx.ErrNoRows) && strings.TrimSpace(input.IdempotencyKeyHash) != "" {
@@ -2204,6 +2211,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
&task.ExternalTaskID,
&task.Kind,
&task.RunMode,
&task.AcceptanceRunID,
&task.UserID,
&task.GatewayUserID,
&task.UserSource,
@@ -2461,6 +2469,12 @@ func normalizeRunMode(input string, request map[string]any) string {
if mode == "simulation" {
return "simulation"
}
if mode == "acceptance" {
return "acceptance"
}
if mode == "acceptance_canary" {
return "acceptance_canary"
}
return "production"
}
apps/api/internal/store/postgres_pool_metrics.go
+25
View File
@@ -0,0 +1,25 @@
package store
type PostgresPoolMetricsSnapshot struct {
MaxConnections int32
TotalConnections int32
AcquiredConnections int32
IdleConnections int32
EmptyAcquireCount int64
CanceledAcquireCount int64
}
func (s *Store) PostgresPoolMetrics() PostgresPoolMetricsSnapshot {
if s == nil || s.pool == nil {
return PostgresPoolMetricsSnapshot{}
}
statistics := s.pool.Stat()
return PostgresPoolMetricsSnapshot{
MaxConnections: statistics.MaxConns(),
TotalConnections: statistics.TotalConns(),
AcquiredConnections: statistics.AcquiredConns(),
IdleConnections: statistics.IdleConns(),
EmptyAcquireCount: statistics.EmptyAcquireCount(),
CanceledAcquireCount: statistics.CanceledAcquireCount(),
}
}
apps/api/internal/store/tasks_runtime.go
+11 -11
View File
@@ -294,7 +294,7 @@ func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, execution
if err := tx.QueryRow(ctx, `
SELECT status = 'queued' AND next_run_at <= now(),
status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()),
run_mode = 'production',
run_mode IN ('production', 'acceptance', 'acceptance_canary'),
gateway_user_id IS NOT NULL
FROM gateway_tasks
WHERE id = $1::uuid
@@ -431,7 +431,7 @@ SELECT status = 'queued'
OR execution_lease_expires_at IS NULL
OR execution_lease_expires_at <= now()
),
run_mode = 'production',
run_mode IN ('production', 'acceptance', 'acceptance_canary'),
gateway_user_id IS NOT NULL
FROM gateway_tasks
WHERE id = $1::uuid
@@ -513,7 +513,7 @@ SET status = 'failed',
error_code = NULLIF($2, ''),
error_message = NULLIF($3, ''),
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required'
ELSE 'released'
END,
billing_updated_at = now(),
@@ -910,7 +910,7 @@ SET status = 'cancelled',
error_code = 'client_disconnected',
error_message = NULLIF($3::text, ''),
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required'
ELSE 'released'
END,
billing_updated_at = now(),
@@ -991,7 +991,7 @@ SET status = 'cancelled',
error_code = 'task_cancelled',
error_message = NULLIF($2, ''),
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
@@ -1027,7 +1027,7 @@ SELECT id, 'task.billing.release', 'release', reservation_amount, billing_curren
pricing_snapshot, $2::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON))
@@ -1569,7 +1569,7 @@ SET status = 'succeeded',
final_charge_amount = $9::numeric,
billing_version = 'effective-pricing-v2',
billing_status = CASE
WHEN run_mode = 'production' AND gateway_user_id IS NOT NULL THEN 'pending'
WHEN run_mode IN ('production', 'acceptance', 'acceptance_canary') AND gateway_user_id IS NOT NULL THEN 'pending'
ELSE 'not_required'
END,
billing_currency = $10,
@@ -1624,7 +1624,7 @@ INSERT INTO settlement_outbox (
SELECT id, 'task.billing.settle', 'settle', $2::numeric, $3, $4::jsonb, $5::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL
ON CONFLICT (task_id, event_type) DO NOTHING`,
input.TaskID, finalChargeAmount, currency, string(pricingSnapshotJSON), string(payloadJSON))
@@ -1725,7 +1725,7 @@ SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currenc
pricing_snapshot, $2::jsonb, 'manual_review', now(), $3
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON), input.Code)
return err
@@ -1902,7 +1902,7 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
response_duration_ms = $8,
result = $9::jsonb,
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
@@ -1944,7 +1944,7 @@ SELECT id, 'task.billing.release', 'release', reservation_amount, billing_curren
pricing_snapshot, $2::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON))
apps/api/migrations/0094_production_acceptance_runs.sql
+50
View File
@@ -0,0 +1,50 @@
CREATE TABLE gateway_acceptance_runs (
id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
status text NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'running', 'passed', 'failed', 'promoted', 'aborted')),
release_sha text NOT NULL,
api_image_digest text NOT NULL,
worker_image_digest text NOT NULL,
api_key_id text NOT NULL,
user_id text NOT NULL,
token_sha256 text NOT NULL,
emulator_base_url text NOT NULL,
callback_url text,
capacity_profile text NOT NULL DEFAULT 'P24'
CHECK (capacity_profile IN ('P24', 'P28', 'P32')),
config jsonb NOT NULL DEFAULT '{}'::jsonb,
report jsonb NOT NULL DEFAULT '{}'::jsonb,
failure_reason text,
started_at timestamptz,
finished_at timestamptz,
promoted_at timestamptz,
created_at timestamptz NOT NULL DEFAULT now(),
updated_at timestamptz NOT NULL DEFAULT now()
);
CREATE INDEX idx_gateway_acceptance_runs_status_created
ON gateway_acceptance_runs (status, created_at DESC);
ALTER TABLE gateway_tasks
ADD COLUMN acceptance_run_id uuid REFERENCES gateway_acceptance_runs(id),
ADD CONSTRAINT gateway_tasks_acceptance_run_mode_check CHECK (
(
run_mode IN ('acceptance', 'acceptance_canary')
AND acceptance_run_id IS DISTINCT FROM NULL
)
OR (
run_mode NOT IN ('acceptance', 'acceptance_canary')
AND acceptance_run_id IS NULL
)
) NOT VALID;
CREATE INDEX idx_gateway_tasks_acceptance_run
ON gateway_tasks (acceptance_run_id, status, created_at)
WHERE acceptance_run_id IS NOT NULL;
INSERT INTO system_settings (setting_key, value)
VALUES (
'gateway_traffic_mode',
'{"mode":"live","revision":0}'::jsonb
)
ON CONFLICT (setting_key) DO NOTHING;