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

fix(routing): 对齐缓存亲和力硬规则与审计

Browse Source
This commit is contained in:
wangbo
2026-07-24 23:51:27 +08:00
parent d7a0ec56f5
commit 152885bbb6
10 changed files with 540 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/httpapi/core_flow_integration_test.go
+87 -11
View File
@@ -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
}