feat(acceptance): 增加生产同构媒体压力验收模式
引入动态流量门禁、隔离验收身份与协议级 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:
@@ -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',
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"},
|
||||
|
||||
@@ -418,6 +418,8 @@ func cleanProviderBody(body map[string]any) map[string]any {
|
||||
"platformId",
|
||||
"platform_model_id",
|
||||
"platformModelId",
|
||||
"modelType",
|
||||
"model_type",
|
||||
} {
|
||||
delete(out, key)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()},
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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(¤tValue); err != nil {
|
||||
return GatewayTrafficMode{}, err
|
||||
}
|
||||
var current GatewayTrafficMode
|
||||
if err := json.Unmarshal(currentValue, ¤t); 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, ¤t); 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, ¤t); 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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user