gateway_anthropic_vertex_service_account_test.go 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. package service
  2. import (
  3. "context"
  4. "io"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/require"
  10. "github.com/tidwall/gjson"
  11. )
  12. func TestGatewayService_BuildAnthropicVertexServiceAccountRequest(t *testing.T) {
  13. gin.SetMode(gin.TestMode)
  14. rec := httptest.NewRecorder()
  15. c, _ := gin.CreateTestContext(rec)
  16. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  17. c.Request.Header.Set("Authorization", "Bearer inbound-token")
  18. c.Request.Header.Set("X-Api-Key", "inbound-api-key")
  19. c.Request.Header.Set("Anthropic-Version", "2023-06-01")
  20. c.Request.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
  21. account := &Account{
  22. ID: 301,
  23. Platform: PlatformAnthropic,
  24. Type: AccountTypeServiceAccount,
  25. Credentials: map[string]any{
  26. "project_id": "vertex-proj",
  27. "location": "us-east5",
  28. },
  29. }
  30. body := []byte(`{"model":"claude-sonnet-4-5","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hello"}]}`)
  31. svc := &GatewayService{}
  32. req, _, err := svc.buildUpstreamRequest(
  33. context.Background(),
  34. c,
  35. account,
  36. body,
  37. "vertex-token",
  38. "service_account",
  39. "claude-sonnet-4-5@20250929",
  40. false,
  41. false,
  42. )
  43. require.NoError(t, err)
  44. require.Equal(t, "https://us-east5-aiplatform.googleapis.com/v1/projects/vertex-proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929:rawPredict", req.URL.String())
  45. require.Equal(t, "Bearer vertex-token", getHeaderRaw(req.Header, "authorization"))
  46. require.Empty(t, getHeaderRaw(req.Header, "x-api-key"))
  47. require.Empty(t, getHeaderRaw(req.Header, "anthropic-version"))
  48. require.Equal(t, "interleaved-thinking-2025-05-14", getHeaderRaw(req.Header, "anthropic-beta"))
  49. got := readRequestBodyForTest(t, req)
  50. require.Equal(t, "", gjson.GetBytes(got, "model").String())
  51. require.Equal(t, vertexAnthropicVersion, gjson.GetBytes(got, "anthropic_version").String())
  52. require.Equal(t, "hello", gjson.GetBytes(got, "messages.0.content").String())
  53. }
  54. func readRequestBodyForTest(t *testing.T, req *http.Request) []byte {
  55. t.Helper()
  56. require.NotNil(t, req.Body)
  57. body, err := io.ReadAll(req.Body)
  58. require.NoError(t, err)
  59. return body
  60. }
  61. // Vertex 路径回归保护:同样需要
  62. // body↔beta header 能力维度对称。客户端 header 不带 context-management beta
  63. // 但 body 带 context_management 字段 → Vertex builder 必须 strip 字段,与 Anthropic
  64. // 直连 / Bedrock 路径保持一致。
  65. func TestGatewayService_BuildAnthropicVertexServiceAccount_StripsContextManagementWhenBetaMissing(t *testing.T) {
  66. gin.SetMode(gin.TestMode)
  67. rec := httptest.NewRecorder()
  68. c, _ := gin.CreateTestContext(rec)
  69. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  70. // 客户端 header 只带 interleaved-thinking,不带 context-management-2025-06-27
  71. c.Request.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
  72. account := &Account{
  73. ID: 302, Platform: PlatformAnthropic, Type: AccountTypeServiceAccount,
  74. Credentials: map[string]any{"project_id": "vertex-proj", "location": "us-east5"},
  75. }
  76. // body 带了 context_management 字段(客户端透传 / normalize 补齐 / mimicry 注入等场景都可能导致)
  77. body := []byte(`{"model":"claude-haiku-4-5","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]},"messages":[{"role":"user","content":"hi"}]}`)
  78. svc := &GatewayService{}
  79. req, _, err := svc.buildUpstreamRequest(
  80. context.Background(), c, account, body,
  81. "vertex-token", "service_account", "claude-haiku-4-5@20251001", false, false,
  82. )
  83. require.NoError(t, err)
  84. got := readRequestBodyForTest(t, req)
  85. require.False(t, gjson.GetBytes(got, "context_management").Exists(),
  86. "Vertex 路径下客户端 header 缺 context-management beta 时,必须 strip body 同名字段")
  87. // header 对称断言:覆盖未来某人在 Vertex builder 里加入与 sanitize 不一致的 header 处理。
  88. outBeta := getHeaderRaw(req.Header, "anthropic-beta")
  89. require.False(t, anthropicBetaTokensContains(outBeta, "context-management-2025-06-27"),
  90. "与 body 对称:outgoing anthropic-beta header 也不含 context-management beta")
  91. }
  92. // Vertex 路径反面:客户端 header 含 context-management beta 时保留字段。
  93. func TestGatewayService_BuildAnthropicVertexServiceAccount_PreservesContextManagementWhenBetaPresent(t *testing.T) {
  94. gin.SetMode(gin.TestMode)
  95. rec := httptest.NewRecorder()
  96. c, _ := gin.CreateTestContext(rec)
  97. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  98. c.Request.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14,context-management-2025-06-27")
  99. account := &Account{
  100. ID: 303, Platform: PlatformAnthropic, Type: AccountTypeServiceAccount,
  101. Credentials: map[string]any{"project_id": "vertex-proj", "location": "us-east5"},
  102. }
  103. body := []byte(`{"model":"claude-sonnet-4-6","context_management":{"edits":[{"type":"clear_thinking_20251015"}]},"messages":[]}`)
  104. svc := &GatewayService{}
  105. req, _, err := svc.buildUpstreamRequest(
  106. context.Background(), c, account, body,
  107. "vertex-token", "service_account", "claude-sonnet-4-6@20260218", false, false,
  108. )
  109. require.NoError(t, err)
  110. got := readRequestBodyForTest(t, req)
  111. require.True(t, gjson.GetBytes(got, "context_management").Exists(),
  112. "Vertex + 客户端 header 包含 context-management beta 时字段必须保留")
  113. outBeta := getHeaderRaw(req.Header, "anthropic-beta")
  114. require.True(t, anthropicBetaTokensContains(outBeta, "context-management-2025-06-27"),
  115. "与 body 对称:outgoing anthropic-beta header 同步含 context-management beta")
  116. }