fix(routing): 对齐缓存亲和力硬规则与审计
This commit is contained in:
@@ -1746,6 +1746,7 @@ WHERE m.platform_id = $1::uuid
|
||||
}
|
||||
|
||||
restartModel := "worker-restart-" + suffixText
|
||||
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel)
|
||||
createSimulationTextPlatformModel(
|
||||
t,
|
||||
server.URL,
|
||||
@@ -1864,6 +1865,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
|
||||
model := "cache-affinity-smoke-" + suffixText
|
||||
cacheKey := "cache-affinity-key-" + suffixText
|
||||
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, model)
|
||||
lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
|
||||
"rules": []map[string]any{
|
||||
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
|
||||
@@ -1880,8 +1882,31 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil)
|
||||
_ = highPlatform
|
||||
tools := []any{map[string]any{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
"description": "Look up current weather",
|
||||
"parameters": map[string]any{"type": "object", "properties": map[string]any{"city": map[string]any{"type": "string"}}},
|
||||
},
|
||||
}}
|
||||
initialToolMessages := []any{
|
||||
map[string]any{"role": "system", "content": "Use tools when helpful."},
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
}
|
||||
for index := 0; index < 3; index++ {
|
||||
detail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{
|
||||
"tools": tools,
|
||||
"messages": initialToolMessages,
|
||||
"simulationUsage": map[string]any{"inputTokens": 1000, "cachedInputTokens": 750, "outputTokens": 20},
|
||||
})
|
||||
if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
|
||||
t.Fatalf("tool-history priming request should use the only platform, got %+v", detail)
|
||||
}
|
||||
}
|
||||
|
||||
peerPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-peer-"+suffixText, "Cache Affinity Peer", model, 20, nil)
|
||||
_ = peerPlatform
|
||||
sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
|
||||
"inputTokens": 1000,
|
||||
"cachedInputTokens": 700,
|
||||
@@ -1899,12 +1924,36 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage)
|
||||
}
|
||||
|
||||
toolHistoryDetail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-history-"+suffixText, map[string]any{
|
||||
"tools": tools,
|
||||
"messages": append(append([]any{}, initialToolMessages...),
|
||||
map[string]any{"role": "assistant", "tool_calls": []any{map[string]any{
|
||||
"id": "call_weather",
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
"arguments": `{"city":"Hangzhou"}`,
|
||||
},
|
||||
}}},
|
||||
map[string]any{"role": "tool", "tool_call_id": "call_weather", "content": `{"temperature":31}`},
|
||||
),
|
||||
"simulationUsage": map[string]any{"inputTokens": 1400, "cachedInputTokens": 900, "outputTokens": 30},
|
||||
})
|
||||
if toolHistoryDetail.Status != "succeeded" || len(toolHistoryDetail.Attempts) != 1 || toolHistoryDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
|
||||
t.Fatalf("tool continuation should retain the cached platform, got %+v", toolHistoryDetail)
|
||||
}
|
||||
if !boolFromTestMap(toolHistoryDetail.Attempts[0].Metrics, "cacheAffinityMatched") ||
|
||||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityMatchedPrefixDepth"]) < 2 ||
|
||||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityCandidateCount"]) < 1 {
|
||||
t.Fatalf("tool continuation should expose prefix and candidate metrics, got %+v", toolHistoryDetail.Attempts[0].Metrics)
|
||||
}
|
||||
|
||||
newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
|
||||
"inputTokens": 1000,
|
||||
"cachedInputTokens": 0,
|
||||
"outputTokens": 20,
|
||||
})
|
||||
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
|
||||
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
|
||||
t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail)
|
||||
}
|
||||
|
||||
@@ -1914,9 +1963,12 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
"cachedInputTokens": 0,
|
||||
"outputTokens": 20,
|
||||
})
|
||||
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
|
||||
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
|
||||
t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail)
|
||||
}
|
||||
if fullAvoidedDetail.Attempts[0].Metrics["cacheAffinityOverrideReason"] != "capacity_tier_unavailable" {
|
||||
t.Fatalf("full cached platform override reason should be visible, got %+v", fullAvoidedDetail.Attempts[0].Metrics)
|
||||
}
|
||||
|
||||
var requestCount int
|
||||
var inputTokens int
|
||||
@@ -1926,7 +1978,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
|
||||
FROM gateway_cache_affinity_stats
|
||||
WHERE client_id = $1
|
||||
ORDER BY updated_at DESC
|
||||
ORDER BY request_count DESC, cached_input_tokens DESC, updated_at DESC
|
||||
LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
|
||||
t.Fatalf("read cache affinity stats: %v", err)
|
||||
}
|
||||
@@ -2214,6 +2266,20 @@ type cacheAffinityTaskDetail struct {
|
||||
} `json:"attempts"`
|
||||
}
|
||||
|
||||
func createSimulationTextBaseModel(t *testing.T, baseURL string, adminToken string, model string) {
|
||||
t.Helper()
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{
|
||||
"providerKey": "openai",
|
||||
"canonicalModelKey": "openai:" + model,
|
||||
"invocationName": model,
|
||||
"providerModelName": model,
|
||||
"modelType": []string{"text_generate"},
|
||||
"displayName": model,
|
||||
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
|
||||
"metadata": map[string]any{"source": "cache-affinity-integration-test"},
|
||||
}, http.StatusCreated, nil)
|
||||
}
|
||||
|
||||
func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken string, platformKey string, platformName string, model string, priority int, rateLimitPolicy map[string]any) (cacheAffinityPlatformFixture, string) {
|
||||
t.Helper()
|
||||
var platform cacheAffinityPlatformFixture
|
||||
@@ -2230,7 +2296,7 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
|
||||
t.Fatalf("cache affinity platform was not created: %+v", platform)
|
||||
}
|
||||
payload := map[string]any{
|
||||
"canonicalModelKey": "openai:gpt-4o-mini",
|
||||
"canonicalModelKey": "openai:" + model,
|
||||
"modelName": model,
|
||||
"providerModelName": model,
|
||||
"modelAlias": model,
|
||||
@@ -2251,17 +2317,27 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
|
||||
}
|
||||
|
||||
func runCacheAffinityChatTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, cacheAffinityKey string, marker string, simulationUsage map[string]any) cacheAffinityTaskDetail {
|
||||
t.Helper()
|
||||
return runCacheAffinityChatTaskWithBody(t, ctx, pool, baseURL, token, model, marker, map[string]any{
|
||||
"cacheAffinityKey": cacheAffinityKey,
|
||||
"simulationUsage": simulationUsage,
|
||||
"messages": []any{map[string]any{"role": "user", "content": "cache affinity route"}},
|
||||
})
|
||||
}
|
||||
|
||||
func runCacheAffinityChatTaskWithBody(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, marker string, body map[string]any) cacheAffinityTaskDetail {
|
||||
t.Helper()
|
||||
var detail cacheAffinityTaskDetail
|
||||
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{
|
||||
requestBody := map[string]any{
|
||||
"model": model,
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
"cacheAffinityKey": cacheAffinityKey,
|
||||
"simulationUsage": simulationUsage,
|
||||
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}},
|
||||
}, marker, http.StatusOK, nil, &detail)
|
||||
}
|
||||
for key, value := range body {
|
||||
requestBody[key] = value
|
||||
}
|
||||
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, requestBody, marker, http.StatusOK, nil, &detail)
|
||||
return detail
|
||||
}
|
||||
|
||||
|
||||
@@ -13,6 +13,8 @@ type cacheAffinityKeys struct {
|
||||
Record []string
|
||||
}
|
||||
|
||||
const cacheAffinityMinimumStableMessages = 2
|
||||
|
||||
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
|
||||
return buildCacheAffinityKeys(kind, modelType, body).Primary
|
||||
}
|
||||
@@ -37,48 +39,28 @@ func buildCacheAffinityKeys(kind string, modelType string, body map[string]any)
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
basePayload := map[string]any{
|
||||
"kind": kind,
|
||||
"modelType": modelType,
|
||||
"version": "v2",
|
||||
}
|
||||
if tools, ok := body["tools"]; ok {
|
||||
basePayload["tools"] = tools
|
||||
}
|
||||
if instructions, ok := body["instructions"]; ok {
|
||||
basePayload["instructions"] = instructions
|
||||
}
|
||||
messageHashes := cacheAffinityMessageHashes(body)
|
||||
if len(messageHashes) > 0 {
|
||||
if len(messageHashes) >= cacheAffinityMinimumStableMessages {
|
||||
baseHash := stableJSONHash(basePayload)
|
||||
keys := make([]string, 0, len(messageHashes))
|
||||
for index := range messageHashes {
|
||||
keyPayload := map[string]any{
|
||||
"base": baseHash,
|
||||
"messageHashes": messageHashes[:index+1],
|
||||
"prefixLength": index + 1,
|
||||
}
|
||||
keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
|
||||
}
|
||||
keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, messageHashes)
|
||||
if len(keys) == 0 {
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
lookup := make([]string, 0, len(keys))
|
||||
for index := len(keys) - 1; index >= 0; index-- {
|
||||
lookup = append(lookup, keys[index])
|
||||
}
|
||||
lookup := reverseCacheAffinityKeys(keys)
|
||||
legacyKeys := legacyCacheAffinityPrefixKeys(kind, modelType, body)
|
||||
lookup = normalizedCacheAffinityKeyList(append(lookup, reverseCacheAffinityKeys(legacyKeys)...))
|
||||
return cacheAffinityKeys{
|
||||
Primary: keys[len(keys)-1],
|
||||
Lookup: lookup,
|
||||
Record: keys,
|
||||
}
|
||||
}
|
||||
if input, ok := body["input"]; ok {
|
||||
basePayload["input"] = input
|
||||
}
|
||||
if len(basePayload) <= 2 {
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
key := "prompt:" + stableJSONHash(basePayload)
|
||||
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
|
||||
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
|
||||
@@ -95,6 +77,146 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
|
||||
}
|
||||
|
||||
func cacheAffinityMessageHashes(body map[string]any) []string {
|
||||
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
messages := cacheAffinityStableMessages(body)
|
||||
if len(messages) == 0 {
|
||||
return nil
|
||||
}
|
||||
hashes := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
hash := stableJSONHash(message)
|
||||
if hash == "" {
|
||||
continue
|
||||
}
|
||||
hashes = append(hashes, hash)
|
||||
}
|
||||
return hashes
|
||||
}
|
||||
|
||||
func cacheAffinityStableMessages(body map[string]any) []any {
|
||||
messages := make([]any, 0)
|
||||
if instructions := strings.TrimSpace(stringFromAny(body["instructions"])); instructions != "" {
|
||||
messages = append(messages, map[string]any{"role": "system", "content": instructions})
|
||||
}
|
||||
if rawMessages, ok := body["messages"].([]any); ok {
|
||||
for _, message := range rawMessages {
|
||||
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
|
||||
messages = append(messages, normalized)
|
||||
}
|
||||
}
|
||||
return messages
|
||||
}
|
||||
switch input := body["input"].(type) {
|
||||
case string:
|
||||
if strings.TrimSpace(input) != "" {
|
||||
messages = append(messages, map[string]any{"role": "user", "content": input})
|
||||
}
|
||||
case []any:
|
||||
for _, message := range input {
|
||||
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
|
||||
messages = append(messages, normalized)
|
||||
}
|
||||
}
|
||||
}
|
||||
return messages
|
||||
}
|
||||
|
||||
func normalizeCacheAffinityMessage(value any) any {
|
||||
message, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
|
||||
return map[string]any{"role": "user", "content": text}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
messageType := strings.TrimSpace(stringFromAny(message["type"]))
|
||||
switch messageType {
|
||||
case "function_call":
|
||||
return map[string]any{
|
||||
"role": "assistant",
|
||||
"tool_calls": []any{map[string]any{
|
||||
"id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": message["name"],
|
||||
"arguments": message["arguments"],
|
||||
},
|
||||
}},
|
||||
}
|
||||
case "function_call_output":
|
||||
return map[string]any{
|
||||
"role": "tool",
|
||||
"tool_call_id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
|
||||
"content": message["output"],
|
||||
}
|
||||
}
|
||||
normalized := make(map[string]any, len(message))
|
||||
for key, item := range message {
|
||||
switch key {
|
||||
case "id", "status":
|
||||
continue
|
||||
case "type":
|
||||
if messageType == "message" {
|
||||
continue
|
||||
}
|
||||
}
|
||||
normalized[key] = item
|
||||
}
|
||||
if role := strings.TrimSpace(stringFromAny(normalized["role"])); role != "" {
|
||||
normalized["role"] = role
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func firstNonEmptyCacheAffinityString(values ...any) string {
|
||||
for _, value := range values {
|
||||
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
|
||||
return text
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func cacheAffinityPrefixKeys(prefix string, baseHash string, messageHashes []string) []string {
|
||||
if len(messageHashes) < cacheAffinityMinimumStableMessages {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(messageHashes)-cacheAffinityMinimumStableMessages+1)
|
||||
for prefixLength := cacheAffinityMinimumStableMessages; prefixLength <= len(messageHashes); prefixLength++ {
|
||||
keyPayload := map[string]any{
|
||||
"base": baseHash,
|
||||
"messageHashes": messageHashes[:prefixLength],
|
||||
"prefixLength": prefixLength,
|
||||
}
|
||||
keys = append(keys, prefix+stableJSONHash(keyPayload))
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func legacyCacheAffinityPrefixKeys(kind string, modelType string, body map[string]any) []string {
|
||||
messageHashes := legacyCacheAffinityMessageHashes(body)
|
||||
if len(messageHashes) < cacheAffinityMinimumStableMessages {
|
||||
return nil
|
||||
}
|
||||
basePayload := map[string]any{
|
||||
"kind": kind,
|
||||
"modelType": modelType,
|
||||
}
|
||||
if tools, ok := body["tools"]; ok {
|
||||
basePayload["tools"] = tools
|
||||
}
|
||||
if instructions, ok := body["instructions"]; ok {
|
||||
basePayload["instructions"] = instructions
|
||||
}
|
||||
return cacheAffinityPrefixKeys("prompt_lcp:", stableJSONHash(basePayload), messageHashes)
|
||||
}
|
||||
|
||||
func legacyCacheAffinityMessageHashes(body map[string]any) []string {
|
||||
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
@@ -102,20 +224,41 @@ func cacheAffinityMessageHashes(body map[string]any) []string {
|
||||
return hashes
|
||||
}
|
||||
messages, ok := body["messages"].([]any)
|
||||
if !ok || len(messages) == 0 {
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
hashes := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
raw, err := json.Marshal(message)
|
||||
if err != nil || len(raw) == 0 {
|
||||
continue
|
||||
if err == nil && len(raw) > 0 {
|
||||
hashes = append(hashes, sha256Text(string(raw)))
|
||||
}
|
||||
hashes = append(hashes, sha256Text(string(raw)))
|
||||
}
|
||||
return hashes
|
||||
}
|
||||
|
||||
func reverseCacheAffinityKeys(keys []string) []string {
|
||||
reversed := make([]string, 0, len(keys))
|
||||
for index := len(keys) - 1; index >= 0; index-- {
|
||||
reversed = append(reversed, keys[index])
|
||||
}
|
||||
return reversed
|
||||
}
|
||||
|
||||
func normalizedCacheAffinityKeyList(keys []string) []string {
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" || seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, key)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stableJSONHash(value any) string {
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil || len(raw) == 0 {
|
||||
|
||||
@@ -39,7 +39,7 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
|
||||
|
||||
first := buildCacheAffinityKey("chat.completions", "text_generate", body)
|
||||
second := buildCacheAffinityKey("chat.completions", "text_generate", body)
|
||||
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") {
|
||||
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp_v2:") {
|
||||
t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second)
|
||||
}
|
||||
}
|
||||
@@ -54,10 +54,10 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
|
||||
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
|
||||
})
|
||||
|
||||
if !strings.HasPrefix(first.Primary, "prompt_lcp:") {
|
||||
if !strings.HasPrefix(first.Primary, "prompt_lcp_v2:") {
|
||||
t.Fatalf("expected LCP prompt key, got %+v", first)
|
||||
}
|
||||
if len(first.Record) != 2 || len(next.Lookup) != 4 {
|
||||
if len(first.Record) != 1 || len(next.Record) != 3 || len(next.Lookup) != 6 {
|
||||
t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
|
||||
}
|
||||
found := false
|
||||
@@ -71,3 +71,81 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
|
||||
t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCacheAffinityKeysRequiresTwoStableMessages(t *testing.T) {
|
||||
systemOnly := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
},
|
||||
})
|
||||
firstConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
map[string]any{"role": "user", "content": "conversation A"},
|
||||
},
|
||||
})
|
||||
secondConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
map[string]any{"role": "user", "content": "conversation B"},
|
||||
},
|
||||
})
|
||||
|
||||
if systemOnly.Primary != "" || len(systemOnly.Lookup) != 0 || len(systemOnly.Record) != 0 {
|
||||
t.Fatalf("shared system prompt alone must not create affinity, got %+v", systemOnly)
|
||||
}
|
||||
if firstConversation.Primary == "" || secondConversation.Primary == "" {
|
||||
t.Fatalf("two stable messages should create affinity, first=%+v second=%+v", firstConversation, secondConversation)
|
||||
}
|
||||
for _, firstKey := range firstConversation.Lookup {
|
||||
for _, secondKey := range secondConversation.Lookup {
|
||||
if firstKey == secondKey {
|
||||
t.Fatalf("different conversations sharing only a system prompt must not intersect: key=%s", firstKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCacheAffinityKeysNormalizesResponsesToolHistory(t *testing.T) {
|
||||
tools := []any{map[string]any{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
},
|
||||
}}
|
||||
chat := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "Use tools when helpful."},
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
},
|
||||
})
|
||||
responses := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"instructions": "Use tools when helpful.",
|
||||
"input": "Weather in Hangzhou?",
|
||||
})
|
||||
continued := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"instructions": "Use tools when helpful.",
|
||||
"input": []any{
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
map[string]any{"type": "function_call", "call_id": "call_weather", "name": "lookup_weather", "arguments": `{"city":"Hangzhou"}`},
|
||||
map[string]any{"type": "function_call_output", "call_id": "call_weather", "output": `{"temperature":31}`},
|
||||
},
|
||||
})
|
||||
|
||||
if chat.Primary == "" || chat.Primary != responses.Primary {
|
||||
t.Fatalf("chat and responses should share the V2 protocol-independent prefix, chat=%+v responses=%+v", chat, responses)
|
||||
}
|
||||
found := false
|
||||
for _, key := range continued.Lookup {
|
||||
if key == responses.Primary {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("responses tool continuation should retain its stable prefix, initial=%+v continued=%+v", responses, continued)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,6 +170,14 @@ func attemptMetrics(candidate store.RuntimeModelCandidate, attemptNo int, simula
|
||||
metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens
|
||||
metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost
|
||||
metrics["cacheAffinityApplied"] = candidate.CacheAffinity.Applied
|
||||
metrics["cacheAffinityMatched"] = candidate.CacheAffinity.Applied
|
||||
metrics["cacheAffinityCandidateCount"] = candidate.CacheAffinity.CandidateCount
|
||||
if candidate.CacheAffinity.MatchedPrefixDepth > 0 {
|
||||
metrics["cacheAffinityMatchedPrefixDepth"] = candidate.CacheAffinity.MatchedPrefixDepth
|
||||
}
|
||||
if candidate.CacheAffinity.OverrideReason != "" {
|
||||
metrics["cacheAffinityOverrideReason"] = candidate.CacheAffinity.OverrideReason
|
||||
}
|
||||
}
|
||||
if attemptNo > 0 {
|
||||
metrics["attempt"] = attemptNo
|
||||
|
||||
@@ -13,16 +13,19 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
|
||||
PlatformID: "deepseek-platform",
|
||||
PlatformPriority: 100,
|
||||
CacheAffinity: store.RuntimeCandidateCacheAffinity{
|
||||
Key: "conversation:test",
|
||||
RequestCount: 2,
|
||||
CachedInputTokens: 7296,
|
||||
EMAHitRatio: 0.91,
|
||||
LastHitRatio: 0.92,
|
||||
Confidence: 1,
|
||||
Score: 1.74,
|
||||
Boost: 20,
|
||||
AdjustedPriority: 80,
|
||||
Applied: true,
|
||||
Key: "prompt_lcp_v2:test",
|
||||
RequestCount: 2,
|
||||
CachedInputTokens: 7296,
|
||||
EMAHitRatio: 0.91,
|
||||
LastHitRatio: 0.92,
|
||||
Confidence: 1,
|
||||
Score: 1.74,
|
||||
Boost: 20,
|
||||
AdjustedPriority: 80,
|
||||
MatchedPrefixDepth: 4,
|
||||
CandidateCount: 2,
|
||||
OverrideReason: "different_effective_priority_tier",
|
||||
Applied: true,
|
||||
},
|
||||
}, 1, false)
|
||||
|
||||
@@ -35,4 +38,10 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
|
||||
if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 {
|
||||
t.Fatalf("expected hit ratios in attempt metrics, got %+v", metrics)
|
||||
}
|
||||
if metrics["cacheAffinityMatched"] != true ||
|
||||
metrics["cacheAffinityMatchedPrefixDepth"] != 4 ||
|
||||
metrics["cacheAffinityCandidateCount"] != 2 ||
|
||||
metrics["cacheAffinityOverrideReason"] != "different_effective_priority_tier" {
|
||||
t.Fatalf("expected cache affinity routing diagnostics in attempt metrics, got %+v", metrics)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -500,14 +500,15 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
|
||||
key = options.CacheAffinityKey
|
||||
}
|
||||
affinity := RuntimeCandidateCacheAffinity{
|
||||
Key: key,
|
||||
RequestCount: input.RequestCount,
|
||||
InputTokens: input.InputTokens,
|
||||
CachedInputTokens: input.CachedInputTokens,
|
||||
EMAHitRatio: boundedRatio(input.EMAHitRatio),
|
||||
LastHitRatio: boundedRatio(input.LastHitRatio),
|
||||
LastObservedUnix: input.LastObservedUnix,
|
||||
AdjustedPriority: float64(candidate.PlatformPriority),
|
||||
Key: key,
|
||||
RequestCount: input.RequestCount,
|
||||
InputTokens: input.InputTokens,
|
||||
CachedInputTokens: input.CachedInputTokens,
|
||||
EMAHitRatio: boundedRatio(input.EMAHitRatio),
|
||||
LastHitRatio: boundedRatio(input.LastHitRatio),
|
||||
LastObservedUnix: input.LastObservedUnix,
|
||||
AdjustedPriority: float64(candidate.PlatformPriority),
|
||||
MatchedPrefixDepth: cacheAffinityMatchedPrefixDepth(key, options.CacheAffinityKeys),
|
||||
}
|
||||
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
|
||||
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
|
||||
@@ -538,6 +539,32 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
|
||||
candidate.CacheAffinity = affinity
|
||||
}
|
||||
|
||||
func cacheAffinityMatchedPrefixDepth(key string, lookup []string) int {
|
||||
prefix := ""
|
||||
switch {
|
||||
case strings.HasPrefix(key, "prompt_lcp_v2:"):
|
||||
prefix = "prompt_lcp_v2:"
|
||||
case strings.HasPrefix(key, "prompt_lcp:"):
|
||||
prefix = "prompt_lcp:"
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
matchingKeys := make([]string, 0, len(lookup))
|
||||
for _, candidateKey := range lookup {
|
||||
if strings.HasPrefix(candidateKey, prefix) {
|
||||
matchingKeys = append(matchingKeys, candidateKey)
|
||||
}
|
||||
}
|
||||
for index, candidateKey := range matchingKeys {
|
||||
if candidateKey == key {
|
||||
return len(matchingKeys) - index + cacheAffinityMinimumPrefixDepth - 1
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
const cacheAffinityMinimumPrefixDepth = 2
|
||||
|
||||
func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
|
||||
if enabled, ok := policy["enabled"].(bool); ok && !enabled {
|
||||
return false
|
||||
@@ -712,6 +739,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
|
||||
if aFull != bFull {
|
||||
return !aFull
|
||||
}
|
||||
if items[i].PlatformPriority != items[j].PlatformPriority {
|
||||
return items[i].PlatformPriority < items[j].PlatformPriority
|
||||
}
|
||||
if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied {
|
||||
return items[i].CacheAffinity.Applied
|
||||
}
|
||||
@@ -746,6 +776,32 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
|
||||
}
|
||||
return false
|
||||
})
|
||||
appliedCount := 0
|
||||
for index := range items {
|
||||
if items[index].CacheAffinity.Applied {
|
||||
appliedCount++
|
||||
}
|
||||
}
|
||||
for index := range items {
|
||||
items[index].CacheAffinity.CandidateCount = appliedCount
|
||||
}
|
||||
if len(items) == 0 || items[0].CacheAffinity.Applied {
|
||||
return
|
||||
}
|
||||
for index := 1; index < len(items); index++ {
|
||||
if !items[index].CacheAffinity.Applied {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case runtimeCandidateFull(items[index]) && !runtimeCandidateFull(items[0]):
|
||||
items[0].CacheAffinity.OverrideReason = "capacity_tier_unavailable"
|
||||
case items[index].PlatformPriority != items[0].PlatformPriority:
|
||||
items[0].CacheAffinity.OverrideReason = "different_effective_priority_tier"
|
||||
}
|
||||
if items[0].CacheAffinity.OverrideReason != "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runtimeCandidateFull(candidate RuntimeModelCandidate) bool {
|
||||
|
||||
@@ -120,7 +120,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
|
||||
},
|
||||
{
|
||||
PlatformID: "cache-affinity",
|
||||
PlatformPriority: 20,
|
||||
PlatformPriority: 10,
|
||||
},
|
||||
}
|
||||
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
|
||||
@@ -148,7 +148,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) {
|
||||
func TestRuntimeCandidateSortingKeepsEffectivePriorityAheadOfCacheAffinity(t *testing.T) {
|
||||
candidates := []RuntimeModelCandidate{
|
||||
{
|
||||
PlatformID: "volces-priority",
|
||||
@@ -177,14 +177,17 @@ func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing
|
||||
|
||||
sortRuntimeModelCandidates(candidates)
|
||||
|
||||
if candidates[0].PlatformID != "deepseek-cache-hit" || !candidates[0].CacheAffinity.Applied {
|
||||
t.Fatalf("expected observed cache hit to outrank platform priority, got %+v", candidates)
|
||||
if candidates[0].PlatformID != "volces-priority" || candidates[0].CacheAffinity.Applied {
|
||||
t.Fatalf("expected effective priority tier to remain first, got %+v", candidates)
|
||||
}
|
||||
if candidates[0].CacheAffinity.Score <= 0 || candidates[0].CacheAffinity.Boost <= 0 {
|
||||
t.Fatalf("expected observed cache hit to carry score and boost, got %+v", candidates[0].CacheAffinity)
|
||||
if candidates[0].CacheAffinity.OverrideReason != "different_effective_priority_tier" {
|
||||
t.Fatalf("expected priority override reason on selected candidate, got %+v", candidates[0].CacheAffinity)
|
||||
}
|
||||
if candidates[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied {
|
||||
t.Fatalf("expected priority-only candidate second, got %+v", candidates)
|
||||
if candidates[1].PlatformID != "deepseek-cache-hit" || !candidates[1].CacheAffinity.Applied {
|
||||
t.Fatalf("expected cache affinity candidate to remain visible as fallback, got %+v", candidates)
|
||||
}
|
||||
if candidates[1].CacheAffinity.Score <= 0 || candidates[1].CacheAffinity.Boost <= 0 {
|
||||
t.Fatalf("expected observed cache hit to retain score and boost diagnostics, got %+v", candidates[1].CacheAffinity)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -196,7 +199,7 @@ func TestRuntimeCandidateSortingOrdersMultipleCacheAffinityCandidatesByScore(t *
|
||||
},
|
||||
{
|
||||
PlatformID: "higher-cache-score",
|
||||
PlatformPriority: 50,
|
||||
PlatformPriority: 1,
|
||||
},
|
||||
}
|
||||
policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}}
|
||||
@@ -260,6 +263,29 @@ func TestRuntimeCandidateSortingKeepsFullCacheAffinityCandidateAvoided(t *testin
|
||||
if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided {
|
||||
t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates)
|
||||
}
|
||||
if candidates[0].CacheAffinity.OverrideReason != "capacity_tier_unavailable" {
|
||||
t.Fatalf("expected capacity override reason on selected candidate, got %+v", candidates[0].CacheAffinity)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCacheAffinityMatchedPrefixDepthIgnoresLegacyCompatibilityKeys(t *testing.T) {
|
||||
lookup := []string{
|
||||
"prompt_lcp_v2:depth-four",
|
||||
"prompt_lcp_v2:depth-three",
|
||||
"prompt_lcp_v2:depth-two",
|
||||
"prompt_lcp:depth-four",
|
||||
"prompt_lcp:depth-three",
|
||||
"prompt_lcp:depth-two",
|
||||
}
|
||||
if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp_v2:depth-three", lookup); depth != 3 {
|
||||
t.Fatalf("V2 matched prefix depth=%d want=3", depth)
|
||||
}
|
||||
if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp:depth-two", lookup); depth != 2 {
|
||||
t.Fatalf("legacy matched prefix depth=%d want=2", depth)
|
||||
}
|
||||
if depth := cacheAffinityMatchedPrefixDepth("explicit:key", lookup); depth != 0 {
|
||||
t.Fatalf("explicit key prefix depth=%d want=0", depth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
|
||||
|
||||
@@ -178,18 +178,21 @@ type RuntimeModelCandidate struct {
|
||||
}
|
||||
|
||||
type RuntimeCandidateCacheAffinity struct {
|
||||
Key string
|
||||
RequestCount int
|
||||
InputTokens int
|
||||
CachedInputTokens int
|
||||
EMAHitRatio float64
|
||||
LastHitRatio float64
|
||||
LastObservedUnix float64
|
||||
Confidence float64
|
||||
Score float64
|
||||
Boost float64
|
||||
AdjustedPriority float64
|
||||
Applied bool
|
||||
Key string
|
||||
RequestCount int
|
||||
InputTokens int
|
||||
CachedInputTokens int
|
||||
EMAHitRatio float64
|
||||
LastHitRatio float64
|
||||
LastObservedUnix float64
|
||||
Confidence float64
|
||||
Score float64
|
||||
Boost float64
|
||||
AdjustedPriority float64
|
||||
MatchedPrefixDepth int
|
||||
CandidateCount int
|
||||
OverrideReason string
|
||||
Applied bool
|
||||
}
|
||||
|
||||
type RuntimeCandidateLoadMetrics struct {
|
||||
|
||||
@@ -128,7 +128,7 @@ WHERE task_id=$1::uuid
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
|
||||
func TestTaskAttemptKeepsOnlyCompactDiagnostics(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()}
|
||||
@@ -157,8 +157,8 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
|
||||
Status: "failed",
|
||||
Retryable: true,
|
||||
StatusCode: 503,
|
||||
Usage: map[string]any{"output_tokens": 100},
|
||||
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}},
|
||||
Usage: map[string]any{"output_tokens": 100, "raw": strings.Repeat("duplicate", 100)},
|
||||
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}, "cacheAffinityMatched": true, "cacheAffinityMatchedPrefixDepth": 3},
|
||||
ResponseSnapshot: map[string]any{"result": "duplicate"},
|
||||
ErrorCode: "upstream_error",
|
||||
ErrorMessage: strings.Repeat("错误", 4096),
|
||||
@@ -173,9 +173,16 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
|
||||
t.Fatalf("attempts=%d", len(attempts))
|
||||
}
|
||||
attempt := attempts[0]
|
||||
if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 ||
|
||||
len(attempt.Usage) != 0 || len(attempt.Metrics) != 0 || len(attempt.PricingSnapshot) != 0 {
|
||||
t.Fatalf("attempt contains duplicate JSON: %+v", attempt)
|
||||
if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || len(attempt.PricingSnapshot) != 0 {
|
||||
t.Fatalf("attempt contains duplicate snapshots: %+v", attempt)
|
||||
}
|
||||
if len(attempt.Usage) != 1 || taskAttemptMetricInt(attempt.Usage, "output_tokens") != 100 {
|
||||
t.Fatalf("attempt compact usage mismatch: %+v", attempt.Usage)
|
||||
}
|
||||
if len(attempt.Metrics) != 2 ||
|
||||
attempt.Metrics["cacheAffinityMatched"] != true ||
|
||||
taskAttemptMetricInt(attempt.Metrics, "cacheAffinityMatchedPrefixDepth") != 3 {
|
||||
t.Fatalf("attempt compact metrics mismatch: %+v", attempt.Metrics)
|
||||
}
|
||||
if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 {
|
||||
t.Fatalf("attempt status/error not preserved safely: %+v", attempt)
|
||||
|
||||
@@ -996,14 +996,16 @@ func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptIn
|
||||
if statusCode == 0 {
|
||||
statusCode = taskAttemptMetricInt(input.Metrics, "statusCode")
|
||||
}
|
||||
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage)))
|
||||
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics)))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET status = $2::text,
|
||||
retryable = $3,
|
||||
request_id = NULLIF($4::text, ''),
|
||||
status_code = NULLIF($5::int, 0),
|
||||
usage = '{}'::jsonb,
|
||||
metrics = '{}'::jsonb,
|
||||
usage = $11::jsonb,
|
||||
metrics = $12::jsonb,
|
||||
response_snapshot = '{}'::jsonb,
|
||||
pricing_snapshot = '{}'::jsonb,
|
||||
request_fingerprint = NULL,
|
||||
@@ -1024,10 +1026,45 @@ WHERE id = $1::uuid`,
|
||||
input.ResponseDurationMS,
|
||||
input.ErrorCode,
|
||||
truncateUTF8Bytes(input.ErrorMessage, 2048),
|
||||
usageJSON,
|
||||
metricsJSON,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func minimalTaskAttemptUsage(usage map[string]any) map[string]any {
|
||||
return whitelistedTaskAttemptMap(usage, []string{
|
||||
"inputTokens", "promptTokens", "input_tokens", "prompt_tokens",
|
||||
"cachedInputTokens", "cachedPromptTokens", "cached_input_tokens", "cached_tokens",
|
||||
"cachedInputTokensKnown",
|
||||
"outputTokens", "completionTokens", "output_tokens", "completion_tokens",
|
||||
"totalTokens", "total_tokens",
|
||||
})
|
||||
}
|
||||
|
||||
func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any {
|
||||
return whitelistedTaskAttemptMap(metrics, []string{
|
||||
"platformPriority", "currentPriority", "loadRatio", "loadAvoided",
|
||||
"cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples",
|
||||
"cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio",
|
||||
"cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio",
|
||||
"cacheAffinityCachedInputTokens", "cacheAffinityBoost",
|
||||
"cacheAffinityApplied", "cacheAffinityMatched",
|
||||
"cacheAffinityCandidateCount", "cacheAffinityMatchedPrefixDepth",
|
||||
"cacheAffinityOverrideReason",
|
||||
})
|
||||
}
|
||||
|
||||
func whitelistedTaskAttemptMap(input map[string]any, keys []string) map[string]any {
|
||||
out := map[string]any{}
|
||||
for _, key := range keys {
|
||||
if value, ok := input[key]; ok {
|
||||
out[key] = value
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) {
|
||||
resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result))
|
||||
if resultReport.BinaryCount > 0 {
|
||||
@@ -1041,6 +1078,8 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
|
||||
billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings))
|
||||
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage)))
|
||||
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
|
||||
attemptUsageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage)))
|
||||
attemptMetricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics)))
|
||||
billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary)))
|
||||
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
|
||||
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
|
||||
@@ -1056,8 +1095,8 @@ SET status = 'succeeded',
|
||||
retryable = false,
|
||||
request_id = NULLIF($2, ''),
|
||||
status_code = NULL,
|
||||
usage = '{}'::jsonb,
|
||||
metrics = '{}'::jsonb,
|
||||
usage = $6::jsonb,
|
||||
metrics = $7::jsonb,
|
||||
response_snapshot = '{}'::jsonb,
|
||||
pricing_snapshot = '{}'::jsonb,
|
||||
request_fingerprint = NULL,
|
||||
@@ -1072,6 +1111,7 @@ SET status = 'succeeded',
|
||||
WHERE id = $1::uuid`,
|
||||
input.AttemptID, input.RequestID,
|
||||
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
|
||||
attemptUsageJSON, attemptMetricsJSON,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user