gemini_messages_compat_service_test.go 35 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "net/http/httptest"
  10. "strings"
  11. "testing"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
  15. "github.com/gin-gonic/gin"
  16. "github.com/stretchr/testify/require"
  17. )
  18. type geminiCompatHTTPUpstreamStub struct {
  19. response *http.Response
  20. err error
  21. calls int
  22. lastReq *http.Request
  23. }
  24. func (s *geminiCompatHTTPUpstreamStub) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
  25. s.calls++
  26. s.lastReq = req
  27. if s.err != nil {
  28. return nil, s.err
  29. }
  30. if s.response == nil {
  31. return nil, fmt.Errorf("missing stub response")
  32. }
  33. resp := *s.response
  34. return &resp, nil
  35. }
  36. func (s *geminiCompatHTTPUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
  37. return s.Do(req, proxyURL, accountID, accountConcurrency)
  38. }
  39. func TestGeminiForwardAsChatCompletions_OAuthRoutesToGeminiAndReturnsChatFormat(t *testing.T) {
  40. gin.SetMode(gin.TestMode)
  41. upstreamBody := `data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello from gemini"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3}}}` + "\n\n" +
  42. "data: [DONE]\n\n"
  43. httpStub := &geminiCompatHTTPUpstreamStub{
  44. response: &http.Response{
  45. StatusCode: http.StatusOK,
  46. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  47. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  48. },
  49. }
  50. svc := &GeminiMessagesCompatService{
  51. tokenProvider: &GeminiTokenProvider{},
  52. httpUpstream: httpStub,
  53. cfg: &config.Config{},
  54. }
  55. account := &Account{
  56. ID: 101,
  57. Platform: PlatformGemini,
  58. Type: AccountTypeOAuth,
  59. Credentials: map[string]any{
  60. "access_token": "ya29.test-token",
  61. "project_id": "project-1",
  62. },
  63. Concurrency: 1,
  64. }
  65. rec := httptest.NewRecorder()
  66. c, _ := gin.CreateTestContext(rec)
  67. body := []byte(`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}]}`)
  68. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  69. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
  70. require.NoError(t, err)
  71. require.NotNil(t, result)
  72. require.Equal(t, http.StatusOK, rec.Code)
  73. require.Equal(t, "gemini-2.5-flash", result.Model)
  74. require.Equal(t, 7, result.Usage.InputTokens)
  75. require.Equal(t, 3, result.Usage.OutputTokens)
  76. require.NotNil(t, httpStub.lastReq)
  77. require.Contains(t, httpStub.lastReq.URL.String(), "/v1internal:streamGenerateContent?alt=sse")
  78. require.Equal(t, "Bearer ya29.test-token", httpStub.lastReq.Header.Get("Authorization"))
  79. require.Empty(t, httpStub.lastReq.Header.Get("x-api-key"))
  80. require.Empty(t, httpStub.lastReq.Header.Get("anthropic-version"))
  81. var sent map[string]any
  82. sentBody, err := io.ReadAll(httpStub.lastReq.Body)
  83. require.NoError(t, err)
  84. require.NoError(t, json.Unmarshal(sentBody, &sent))
  85. require.Equal(t, "gemini-2.5-flash", sent["model"])
  86. require.Equal(t, "project-1", sent["project"])
  87. require.Contains(t, fmt.Sprint(sent["request"]), "hi")
  88. var got map[string]any
  89. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
  90. require.Equal(t, "chat.completion", got["object"])
  91. require.Equal(t, "gemini-2.5-flash", got["model"])
  92. choices, ok := got["choices"].([]any)
  93. require.True(t, ok)
  94. require.NotEmpty(t, choices)
  95. choice, ok := choices[0].(map[string]any)
  96. require.True(t, ok)
  97. message, ok := choice["message"].(map[string]any)
  98. require.True(t, ok)
  99. require.Equal(t, "assistant", message["role"])
  100. require.Equal(t, "hello from gemini", message["content"])
  101. usage, ok := got["usage"].(map[string]any)
  102. require.True(t, ok)
  103. require.Equal(t, float64(7), usage["prompt_tokens"])
  104. require.Equal(t, float64(3), usage["completion_tokens"])
  105. require.Equal(t, float64(10), usage["total_tokens"])
  106. }
  107. func TestGeminiForwardAsChatCompletions_StreamsOpenAIChunksFromGeminiSSE(t *testing.T) {
  108. gin.SetMode(gin.TestMode)
  109. upstreamBody := `data: {"candidates":[{"content":{"parts":[{"text":"hel"}]}}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":1}}` + "\n\n" +
  110. `data: {"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":2}}` + "\n\n" +
  111. "data: [DONE]\n\n"
  112. httpStub := &geminiCompatHTTPUpstreamStub{
  113. response: &http.Response{
  114. StatusCode: http.StatusOK,
  115. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  116. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  117. },
  118. }
  119. svc := &GeminiMessagesCompatService{
  120. httpUpstream: httpStub,
  121. cfg: &config.Config{},
  122. }
  123. account := &Account{
  124. ID: 102,
  125. Platform: PlatformGemini,
  126. Type: AccountTypeAPIKey,
  127. Credentials: map[string]any{
  128. "api_key": "gemini-api-key",
  129. },
  130. Concurrency: 1,
  131. }
  132. rec := httptest.NewRecorder()
  133. c, _ := gin.CreateTestContext(rec)
  134. body := []byte(`{"model":"gemini-2.5-flash","stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":"hi"}]}`)
  135. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  136. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
  137. require.NoError(t, err)
  138. require.NotNil(t, result)
  139. require.Equal(t, http.StatusOK, rec.Code)
  140. require.True(t, result.Stream)
  141. require.Equal(t, 2, result.Usage.InputTokens)
  142. require.Equal(t, 2, result.Usage.OutputTokens)
  143. require.NotNil(t, httpStub.lastReq)
  144. require.Contains(t, httpStub.lastReq.URL.String(), "/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse")
  145. require.Equal(t, "gemini-api-key", httpStub.lastReq.Header.Get("x-goog-api-key"))
  146. out := rec.Body.String()
  147. require.Contains(t, out, `"object":"chat.completion.chunk"`)
  148. require.Contains(t, out, `"role":"assistant"`)
  149. require.Contains(t, out, `"content":"hel"`)
  150. require.Contains(t, out, `"content":"lo"`)
  151. require.Contains(t, out, `"usage":{"prompt_tokens":2,"completion_tokens":2,"total_tokens":4}`)
  152. require.Contains(t, out, "data: [DONE]")
  153. }
  154. func TestGeminiForwardAsChatCompletions_FunctionNamedWebSearchStaysClientSide(t *testing.T) {
  155. gin.SetMode(gin.TestMode)
  156. httpStub := &geminiCompatHTTPUpstreamStub{
  157. response: &http.Response{
  158. StatusCode: http.StatusOK,
  159. Header: http.Header{"Content-Type": []string{"application/json"}},
  160. Body: io.NopCloser(strings.NewReader(
  161. `{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":1}}`,
  162. )),
  163. },
  164. }
  165. svc := &GeminiMessagesCompatService{
  166. httpUpstream: httpStub,
  167. cfg: &config.Config{},
  168. }
  169. account := &Account{
  170. ID: 103,
  171. Platform: PlatformGemini,
  172. Type: AccountTypeAPIKey,
  173. Credentials: map[string]any{
  174. "api_key": "gemini-api-key",
  175. },
  176. Concurrency: 1,
  177. }
  178. rec := httptest.NewRecorder()
  179. c, _ := gin.CreateTestContext(rec)
  180. body := []byte(`{
  181. "model":"gemini-3.6-flash-high",
  182. "messages":[{"role":"user","content":"search and read"}],
  183. "tools":[
  184. {"type":"function","function":{"name":"web_search","description":"Search through the Hermes client","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}},
  185. {"type":"function","function":{"name":"read_file","description":"Read a local file","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}}
  186. ]
  187. }`)
  188. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  189. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
  190. require.NoError(t, err)
  191. require.NotNil(t, result)
  192. require.Equal(t, http.StatusOK, rec.Code)
  193. require.NotNil(t, httpStub.lastReq)
  194. postedBody, err := io.ReadAll(httpStub.lastReq.Body)
  195. require.NoError(t, err)
  196. var posted map[string]any
  197. require.NoError(t, json.Unmarshal(postedBody, &posted))
  198. tools, ok := posted["tools"].([]any)
  199. require.True(t, ok)
  200. require.Len(t, tools, 1, "Chat Completions function tools must not be promoted to Gemini built-ins by name")
  201. functionTool, ok := tools[0].(map[string]any)
  202. require.True(t, ok)
  203. functionDecls, ok := functionTool["functionDeclarations"].([]any)
  204. require.True(t, ok)
  205. require.Len(t, functionDecls, 2)
  206. webSearchDecl, ok := functionDecls[0].(map[string]any)
  207. require.True(t, ok)
  208. readFileDecl, ok := functionDecls[1].(map[string]any)
  209. require.True(t, ok)
  210. require.Equal(t, "web_search", webSearchDecl["name"])
  211. require.Equal(t, "read_file", readFileDecl["name"])
  212. require.NotContains(t, functionTool, "googleSearch")
  213. require.NotContains(t, functionTool, "google_search")
  214. }
  215. // TestConvertClaudeToolsToGeminiTools_CustomType 测试custom类型工具转换
  216. func TestConvertClaudeToolsToGeminiTools_CustomType(t *testing.T) {
  217. tests := []struct {
  218. name string
  219. tools any
  220. expectedLen int
  221. description string
  222. }{
  223. {
  224. name: "Standard tools",
  225. tools: []any{
  226. map[string]any{
  227. "name": "get_weather",
  228. "description": "Get weather info",
  229. "input_schema": map[string]any{"type": "object"},
  230. },
  231. },
  232. expectedLen: 1,
  233. description: "标准工具格式应该正常转换",
  234. },
  235. {
  236. name: "Custom type tool (MCP format)",
  237. tools: []any{
  238. map[string]any{
  239. "type": "custom",
  240. "name": "mcp_tool",
  241. "custom": map[string]any{
  242. "description": "MCP tool description",
  243. "input_schema": map[string]any{"type": "object"},
  244. },
  245. },
  246. },
  247. expectedLen: 1,
  248. description: "Custom类型工具应该从custom字段读取",
  249. },
  250. {
  251. name: "Mixed standard and custom tools",
  252. tools: []any{
  253. map[string]any{
  254. "name": "standard_tool",
  255. "description": "Standard",
  256. "input_schema": map[string]any{"type": "object"},
  257. },
  258. map[string]any{
  259. "type": "custom",
  260. "name": "custom_tool",
  261. "custom": map[string]any{
  262. "description": "Custom",
  263. "input_schema": map[string]any{"type": "object"},
  264. },
  265. },
  266. },
  267. expectedLen: 1,
  268. description: "混合工具应该都能正确转换",
  269. },
  270. {
  271. name: "Custom tool without custom field",
  272. tools: []any{
  273. map[string]any{
  274. "type": "custom",
  275. "name": "invalid_custom",
  276. // 缺少 custom 字段
  277. },
  278. },
  279. expectedLen: 0, // 应该被跳过
  280. description: "缺少custom字段的custom工具应该被跳过",
  281. },
  282. }
  283. for _, tt := range tests {
  284. t.Run(tt.name, func(t *testing.T) {
  285. result := convertClaudeToolsToGeminiTools(tt.tools)
  286. if tt.expectedLen == 0 {
  287. if result != nil {
  288. t.Errorf("%s: expected nil result, got %v", tt.description, result)
  289. }
  290. return
  291. }
  292. if result == nil {
  293. t.Fatalf("%s: expected non-nil result", tt.description)
  294. }
  295. if len(result) != 1 {
  296. t.Errorf("%s: expected 1 tool declaration, got %d", tt.description, len(result))
  297. return
  298. }
  299. toolDecl, ok := result[0].(map[string]any)
  300. if !ok {
  301. t.Fatalf("%s: result[0] is not map[string]any", tt.description)
  302. }
  303. funcDecls, ok := toolDecl["functionDeclarations"].([]any)
  304. if !ok {
  305. t.Fatalf("%s: functionDeclarations is not []any", tt.description)
  306. }
  307. toolsArr, _ := tt.tools.([]any)
  308. expectedFuncCount := 0
  309. for _, tool := range toolsArr {
  310. toolMap, _ := tool.(map[string]any)
  311. if toolMap["name"] != "" {
  312. // 检查是否为有效的custom工具
  313. if toolMap["type"] == "custom" {
  314. if toolMap["custom"] != nil {
  315. expectedFuncCount++
  316. }
  317. } else {
  318. expectedFuncCount++
  319. }
  320. }
  321. }
  322. if len(funcDecls) != expectedFuncCount {
  323. t.Errorf("%s: expected %d function declarations, got %d",
  324. tt.description, expectedFuncCount, len(funcDecls))
  325. }
  326. })
  327. }
  328. }
  329. func TestCleanToolSchema_NormalizesGeminiUnsupportedSchemaFields(t *testing.T) {
  330. schema := map[string]any{
  331. "type": "object",
  332. "$defs": map[string]any{
  333. "unused": map[string]any{"type": "string"},
  334. },
  335. "definitions": map[string]any{
  336. "legacy": map[string]any{"type": "number"},
  337. },
  338. "properties": map[string]any{
  339. "path": map[string]any{
  340. "type": []any{"string", "null"},
  341. },
  342. "count": map[string]any{
  343. "type": []any{"null", "integer"},
  344. },
  345. "empty": map[string]any{
  346. "type": []any{"null"},
  347. },
  348. },
  349. }
  350. cleaned, ok := cleanToolSchema(schema).(map[string]any)
  351. require.True(t, ok)
  352. require.Equal(t, "OBJECT", cleaned["type"])
  353. require.NotContains(t, cleaned, "$defs")
  354. require.NotContains(t, cleaned, "definitions")
  355. properties, ok := cleaned["properties"].(map[string]any)
  356. require.True(t, ok)
  357. pathSchema, ok := properties["path"].(map[string]any)
  358. require.True(t, ok)
  359. require.Equal(t, "STRING", pathSchema["type"])
  360. countSchema, ok := properties["count"].(map[string]any)
  361. require.True(t, ok)
  362. require.Equal(t, "INTEGER", countSchema["type"])
  363. emptySchema, ok := properties["empty"].(map[string]any)
  364. require.True(t, ok)
  365. require.NotContains(t, emptySchema, "type")
  366. }
  367. func TestCleanToolSchema_ConvertsNestedIntegerExclusiveMinimum(t *testing.T) {
  368. schema := map[string]any{
  369. "type": "object",
  370. "properties": map[string]any{
  371. "counts": map[string]any{
  372. "type": "array",
  373. "items": map[string]any{
  374. "type": "integer",
  375. "exclusiveMinimum": float64(0),
  376. },
  377. },
  378. "strict": map[string]any{
  379. "type": "integer",
  380. "exclusiveMinimum": 0,
  381. "minimum": 5,
  382. },
  383. "weak": map[string]any{
  384. "type": "integer",
  385. "exclusiveMinimum": 2,
  386. "minimum": 1,
  387. },
  388. },
  389. }
  390. cleaned, ok := cleanToolSchema(schema).(map[string]any)
  391. require.True(t, ok)
  392. properties, ok := cleaned["properties"].(map[string]any)
  393. require.True(t, ok)
  394. counts, ok := properties["counts"].(map[string]any)
  395. require.True(t, ok)
  396. items, ok := counts["items"].(map[string]any)
  397. require.True(t, ok)
  398. require.NotContains(t, items, "exclusiveMinimum")
  399. require.Equal(t, float64(1), items["minimum"])
  400. strict, ok := properties["strict"].(map[string]any)
  401. require.True(t, ok)
  402. require.NotContains(t, strict, "exclusiveMinimum")
  403. require.Equal(t, 5, strict["minimum"])
  404. weak, ok := properties["weak"].(map[string]any)
  405. require.True(t, ok)
  406. require.NotContains(t, weak, "exclusiveMinimum")
  407. require.Equal(t, 3, weak["minimum"])
  408. }
  409. func TestCleanToolSchema_DropsAmbiguousExclusiveMinimumWithoutConversion(t *testing.T) {
  410. for name, schema := range map[string]map[string]any{
  411. "number schema": {
  412. "type": "number",
  413. "exclusiveMinimum": 0,
  414. },
  415. "fractional integer bound": {
  416. "type": "integer",
  417. "exclusiveMinimum": 0.5,
  418. },
  419. } {
  420. t.Run(name, func(t *testing.T) {
  421. cleaned, ok := cleanToolSchema(schema).(map[string]any)
  422. require.True(t, ok)
  423. require.NotContains(t, cleaned, "exclusiveMinimum")
  424. require.NotContains(t, cleaned, "minimum")
  425. })
  426. }
  427. }
  428. func TestCleanToolSchema_RemovesNestedDeprecatedAndNormalizesMixedScalarEnum(t *testing.T) {
  429. schema := map[string]any{
  430. "anyOf": []any{
  431. map[string]any{
  432. "type": "string",
  433. "deprecated": true,
  434. },
  435. map[string]any{
  436. "enum": []any{"enabled", false, float64(1), nil},
  437. },
  438. },
  439. }
  440. cleaned, ok := cleanToolSchema(schema).(map[string]any)
  441. require.True(t, ok)
  442. anyOf, ok := cleaned["anyOf"].([]any)
  443. require.True(t, ok)
  444. require.Len(t, anyOf, 2)
  445. deprecatedSchema, ok := anyOf[0].(map[string]any)
  446. require.True(t, ok)
  447. require.NotContains(t, deprecatedSchema, "deprecated")
  448. enumSchema, ok := anyOf[1].(map[string]any)
  449. require.True(t, ok)
  450. require.Equal(t, []any{"enabled", "false", "1", "null"}, enumSchema["enum"])
  451. }
  452. func TestCleanToolSchema_DropsEnumWithNonScalarValue(t *testing.T) {
  453. schema := map[string]any{
  454. "enum": []any{"valid", map[string]any{"invalid": true}},
  455. }
  456. cleaned, ok := cleanToolSchema(schema).(map[string]any)
  457. require.True(t, ok)
  458. require.NotContains(t, cleaned, "enum")
  459. }
  460. func TestConvertClaudeToolsToGeminiTools_PreservesWebSearchAlongsideFunctions(t *testing.T) {
  461. tools := []any{
  462. map[string]any{
  463. "name": "get_weather",
  464. "description": "Get weather info",
  465. "input_schema": map[string]any{"type": "object"},
  466. },
  467. map[string]any{
  468. "type": "web_search_20250305",
  469. "name": "web_search",
  470. },
  471. }
  472. result := convertClaudeToolsToGeminiTools(tools)
  473. require.Len(t, result, 2)
  474. functionDecl, ok := result[0].(map[string]any)
  475. require.True(t, ok)
  476. funcDecls, ok := functionDecl["functionDeclarations"].([]any)
  477. require.True(t, ok)
  478. require.Len(t, funcDecls, 1)
  479. searchDecl, ok := result[1].(map[string]any)
  480. require.True(t, ok)
  481. googleSearch, ok := searchDecl["googleSearch"].(map[string]any)
  482. require.True(t, ok)
  483. require.Empty(t, googleSearch)
  484. }
  485. func TestGeminiHandleNativeNonStreamingResponse_DebugDisabledDoesNotEmitHeaderLogs(t *testing.T) {
  486. gin.SetMode(gin.TestMode)
  487. logSink, restore := captureStructuredLog(t)
  488. defer restore()
  489. svc := &GeminiMessagesCompatService{
  490. cfg: &config.Config{
  491. Gateway: config.GatewayConfig{
  492. GeminiDebugResponseHeaders: false,
  493. },
  494. },
  495. }
  496. w := httptest.NewRecorder()
  497. c, _ := gin.CreateTestContext(w)
  498. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  499. resp := &http.Response{
  500. StatusCode: http.StatusOK,
  501. Header: http.Header{
  502. "Content-Type": []string{"application/json"},
  503. "X-RateLimit-Limit": []string{"60"},
  504. },
  505. Body: io.NopCloser(strings.NewReader(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}`)),
  506. }
  507. usage, err := svc.handleNativeNonStreamingResponse(c, resp, false)
  508. require.NoError(t, err)
  509. require.NotNil(t, usage)
  510. require.False(t, logSink.ContainsMessage("[GeminiAPI]"), "debug 关闭时不应输出 Gemini 响应头日志")
  511. }
  512. func TestGeminiMessagesCompatServiceForward_PreservesRequestedModelAndMappedUpstreamModel(t *testing.T) {
  513. gin.SetMode(gin.TestMode)
  514. w := httptest.NewRecorder()
  515. c, _ := gin.CreateTestContext(w)
  516. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  517. httpStub := &geminiCompatHTTPUpstreamStub{
  518. response: &http.Response{
  519. StatusCode: http.StatusOK,
  520. Header: http.Header{"x-request-id": []string{"gemini-req-1"}},
  521. Body: io.NopCloser(strings.NewReader(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5}}`)),
  522. },
  523. }
  524. svc := &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}
  525. account := &Account{
  526. ID: 1,
  527. Type: AccountTypeAPIKey,
  528. Credentials: map[string]any{
  529. "api_key": "test-key",
  530. "model_mapping": map[string]any{
  531. "claude-sonnet-4": "claude-sonnet-4-20250514",
  532. },
  533. },
  534. }
  535. body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
  536. result, err := svc.Forward(context.Background(), c, account, body)
  537. require.NoError(t, err)
  538. require.NotNil(t, result)
  539. require.Equal(t, "claude-sonnet-4", result.Model)
  540. require.Equal(t, "claude-sonnet-4-20250514", result.UpstreamModel)
  541. require.Equal(t, 1, httpStub.calls)
  542. require.NotNil(t, httpStub.lastReq)
  543. require.Contains(t, httpStub.lastReq.URL.String(), "/models/claude-sonnet-4-20250514:")
  544. }
  545. func TestGeminiMessagesCompatServiceForward_NormalizesWebSearchToolForAIStudio(t *testing.T) {
  546. gin.SetMode(gin.TestMode)
  547. w := httptest.NewRecorder()
  548. c, _ := gin.CreateTestContext(w)
  549. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  550. httpStub := &geminiCompatHTTPUpstreamStub{
  551. response: &http.Response{
  552. StatusCode: http.StatusOK,
  553. Header: http.Header{"x-request-id": []string{"gemini-req-2"}},
  554. Body: io.NopCloser(strings.NewReader(`{"candidates":[{"content":{"parts":[{"text":"hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5}}`)),
  555. },
  556. }
  557. svc := &GeminiMessagesCompatService{httpUpstream: httpStub, cfg: &config.Config{}}
  558. account := &Account{
  559. ID: 1,
  560. Type: AccountTypeAPIKey,
  561. Credentials: map[string]any{
  562. "api_key": "test-key",
  563. },
  564. }
  565. body := []byte(`{"model":"claude-sonnet-4","max_tokens":16,"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"get_weather","description":"Get weather info","input_schema":{"type":"object"}},{"type":"web_search_20250305","name":"web_search"}]}`)
  566. result, err := svc.Forward(context.Background(), c, account, body)
  567. require.NoError(t, err)
  568. require.NotNil(t, result)
  569. require.NotNil(t, httpStub.lastReq)
  570. postedBody, err := io.ReadAll(httpStub.lastReq.Body)
  571. require.NoError(t, err)
  572. var posted map[string]any
  573. require.NoError(t, json.Unmarshal(postedBody, &posted))
  574. tools, ok := posted["tools"].([]any)
  575. require.True(t, ok)
  576. require.Len(t, tools, 2)
  577. searchTool, ok := tools[1].(map[string]any)
  578. require.True(t, ok)
  579. _, hasSnake := searchTool["google_search"]
  580. _, hasCamel := searchTool["googleSearch"]
  581. require.True(t, hasSnake)
  582. require.False(t, hasCamel)
  583. _, hasFuncDecl := searchTool["functionDeclarations"]
  584. require.False(t, hasFuncDecl)
  585. }
  586. func TestConvertClaudeMessagesToGeminiGenerateContent_AddsThoughtSignatureForToolUse(t *testing.T) {
  587. claudeReq := map[string]any{
  588. "model": "claude-haiku-4-5-20251001",
  589. "max_tokens": 10,
  590. "messages": []any{
  591. map[string]any{
  592. "role": "user",
  593. "content": []any{
  594. map[string]any{"type": "text", "text": "hi"},
  595. },
  596. },
  597. map[string]any{
  598. "role": "assistant",
  599. "content": []any{
  600. map[string]any{"type": "text", "text": "ok"},
  601. map[string]any{
  602. "type": "tool_use",
  603. "id": "toolu_123",
  604. "name": "default_api:write_file",
  605. "input": map[string]any{"path": "a.txt", "content": "x"},
  606. // no signature on purpose
  607. },
  608. },
  609. },
  610. },
  611. "tools": []any{
  612. map[string]any{
  613. "name": "default_api:write_file",
  614. "description": "write file",
  615. "input_schema": map[string]any{
  616. "type": "object",
  617. "properties": map[string]any{"path": map[string]any{"type": "string"}},
  618. },
  619. },
  620. },
  621. }
  622. b, _ := json.Marshal(claudeReq)
  623. out, err := convertClaudeMessagesToGeminiGenerateContent(b)
  624. if err != nil {
  625. t.Fatalf("convert failed: %v", err)
  626. }
  627. s := string(out)
  628. if !strings.Contains(s, "\"functionCall\"") {
  629. t.Fatalf("expected functionCall in output, got: %s", s)
  630. }
  631. if !strings.Contains(s, "\"thoughtSignature\":\""+geminiDummyThoughtSignature+"\"") {
  632. t.Fatalf("expected injected thoughtSignature %q, got: %s", geminiDummyThoughtSignature, s)
  633. }
  634. }
  635. func TestEnsureGeminiFunctionCallThoughtSignatures_InsertsWhenMissing(t *testing.T) {
  636. geminiReq := map[string]any{
  637. "contents": []any{
  638. map[string]any{
  639. "role": "user",
  640. "parts": []any{
  641. map[string]any{
  642. "functionCall": map[string]any{
  643. "name": "default_api:write_file",
  644. "args": map[string]any{"path": "a.txt"},
  645. },
  646. },
  647. },
  648. },
  649. },
  650. }
  651. b, _ := json.Marshal(geminiReq)
  652. out := ensureGeminiFunctionCallThoughtSignatures(b)
  653. s := string(out)
  654. if !strings.Contains(s, "\"thoughtSignature\":\""+geminiDummyThoughtSignature+"\"") {
  655. t.Fatalf("expected injected thoughtSignature %q, got: %s", geminiDummyThoughtSignature, s)
  656. }
  657. }
  658. // TestUnwrapGeminiResponse 测试 unwrapGeminiResponse 的各种输入场景
  659. // 关键区别:只有 response 为 JSON 对象/数组时才解包
  660. func TestUnwrapGeminiResponse(t *testing.T) {
  661. // 构造 >50KB 的大型 JSON 对象
  662. largePadding := strings.Repeat("x", 50*1024)
  663. largeInput := []byte(fmt.Sprintf(`{"response":{"id":"big","pad":"%s"}}`, largePadding))
  664. largeExpected := fmt.Sprintf(`{"id":"big","pad":"%s"}`, largePadding)
  665. tests := []struct {
  666. name string
  667. input []byte
  668. expected string
  669. wantErr bool
  670. }{
  671. {
  672. name: "正常 response 包装(JSON 对象)",
  673. input: []byte(`{"response":{"key":"val"}}`),
  674. expected: `{"key":"val"}`,
  675. },
  676. {
  677. name: "无包装直接返回",
  678. input: []byte(`{"key":"val"}`),
  679. expected: `{"key":"val"}`,
  680. },
  681. {
  682. name: "空 JSON",
  683. input: []byte(`{}`),
  684. expected: `{}`,
  685. },
  686. {
  687. name: "null response 返回原始 body",
  688. input: []byte(`{"response":null}`),
  689. expected: `{"response":null}`,
  690. },
  691. {
  692. name: "非法 JSON 返回原始 body",
  693. input: []byte(`not json`),
  694. expected: `not json`,
  695. },
  696. {
  697. name: "response 为基础类型 string 返回原始 body",
  698. input: []byte(`{"response":"hello"}`),
  699. expected: `{"response":"hello"}`,
  700. },
  701. {
  702. name: "嵌套 response 只解一层",
  703. input: []byte(`{"response":{"response":{"inner":true}}}`),
  704. expected: `{"response":{"inner":true}}`,
  705. },
  706. {
  707. name: "大型 JSON >50KB",
  708. input: largeInput,
  709. expected: largeExpected,
  710. },
  711. }
  712. for _, tt := range tests {
  713. t.Run(tt.name, func(t *testing.T) {
  714. got, err := unwrapGeminiResponse(tt.input)
  715. if tt.wantErr {
  716. require.Error(t, err)
  717. return
  718. }
  719. require.NoError(t, err)
  720. require.Equal(t, tt.expected, strings.TrimSpace(string(got)))
  721. })
  722. }
  723. }
  724. // ---------------------------------------------------------------------------
  725. // Task 8.1 — extractGeminiUsage 测试
  726. // ---------------------------------------------------------------------------
  727. func TestExtractGeminiUsage(t *testing.T) {
  728. tests := []struct {
  729. name string
  730. input string
  731. wantNil bool
  732. wantUsage *ClaudeUsage
  733. }{
  734. {
  735. name: "完整 usageMetadata",
  736. input: `{"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":50,"cachedContentTokenCount":20}}`,
  737. wantNil: false,
  738. wantUsage: &ClaudeUsage{
  739. InputTokens: 80,
  740. OutputTokens: 50,
  741. CacheReadInputTokens: 20,
  742. },
  743. },
  744. {
  745. name: "包含 thoughtsTokenCount",
  746. input: `{"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":20,"thoughtsTokenCount":50}}`,
  747. wantNil: false,
  748. wantUsage: &ClaudeUsage{
  749. InputTokens: 100,
  750. OutputTokens: 70,
  751. CacheReadInputTokens: 0,
  752. },
  753. },
  754. {
  755. name: "包含 thoughtsTokenCount 与缓存",
  756. input: `{"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":20,"cachedContentTokenCount":30,"thoughtsTokenCount":50}}`,
  757. wantNil: false,
  758. wantUsage: &ClaudeUsage{
  759. InputTokens: 70,
  760. OutputTokens: 70,
  761. CacheReadInputTokens: 30,
  762. },
  763. },
  764. {
  765. name: "缺失 cachedContentTokenCount",
  766. input: `{"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":50}}`,
  767. wantNil: false,
  768. wantUsage: &ClaudeUsage{
  769. InputTokens: 100,
  770. OutputTokens: 50,
  771. CacheReadInputTokens: 0,
  772. },
  773. },
  774. {
  775. name: "无 usageMetadata",
  776. input: `{"candidates":[]}`,
  777. wantNil: true,
  778. },
  779. {
  780. // gjson 对 null 返回 Exists()=true,因此函数不会返回 nil,
  781. // 而是返回全零的 ClaudeUsage。
  782. name: "null usageMetadata — gjson Exists 为 true",
  783. input: `{"usageMetadata":null}`,
  784. wantNil: false,
  785. wantUsage: &ClaudeUsage{
  786. InputTokens: 0,
  787. OutputTokens: 0,
  788. CacheReadInputTokens: 0,
  789. },
  790. },
  791. {
  792. name: "零值字段",
  793. input: `{"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"cachedContentTokenCount":0}}`,
  794. wantNil: false,
  795. wantUsage: &ClaudeUsage{
  796. InputTokens: 0,
  797. OutputTokens: 0,
  798. CacheReadInputTokens: 0,
  799. },
  800. },
  801. }
  802. for _, tt := range tests {
  803. t.Run(tt.name, func(t *testing.T) {
  804. got := extractGeminiUsage([]byte(tt.input))
  805. if tt.wantNil {
  806. if got != nil {
  807. t.Fatalf("期望返回 nil,实际返回 %+v", got)
  808. }
  809. return
  810. }
  811. if got == nil {
  812. t.Fatalf("期望返回非 nil,实际返回 nil")
  813. }
  814. if got.InputTokens != tt.wantUsage.InputTokens {
  815. t.Errorf("InputTokens: 期望 %d,实际 %d", tt.wantUsage.InputTokens, got.InputTokens)
  816. }
  817. if got.OutputTokens != tt.wantUsage.OutputTokens {
  818. t.Errorf("OutputTokens: 期望 %d,实际 %d", tt.wantUsage.OutputTokens, got.OutputTokens)
  819. }
  820. if got.CacheReadInputTokens != tt.wantUsage.CacheReadInputTokens {
  821. t.Errorf("CacheReadInputTokens: 期望 %d,实际 %d", tt.wantUsage.CacheReadInputTokens, got.CacheReadInputTokens)
  822. }
  823. // Gemini usageMetadata 只有 cachedContentTokenCount(缓存命中),没有缓存写入
  824. // 的 token 类别:cache_creation_input_tokens 恒为 0,计费侧不会产生缓存创建分项。
  825. if got.CacheCreationInputTokens != 0 {
  826. t.Errorf("CacheCreationInputTokens: 期望 0,实际 %d", got.CacheCreationInputTokens)
  827. }
  828. })
  829. }
  830. }
  831. // ---------------------------------------------------------------------------
  832. // Task 8.2 — estimateGeminiCountTokens 测试
  833. // ---------------------------------------------------------------------------
  834. func TestEstimateGeminiCountTokens(t *testing.T) {
  835. tests := []struct {
  836. name string
  837. input string
  838. wantGt0 bool // 期望结果 > 0
  839. wantExact *int // 如果非 nil,期望精确匹配
  840. }{
  841. {
  842. name: "含 systemInstruction 和 contents",
  843. input: `{
  844. "systemInstruction":{"parts":[{"text":"You are a helpful assistant."}]},
  845. "contents":[{"parts":[{"text":"Hello, how are you?"}]}]
  846. }`,
  847. wantGt0: true,
  848. },
  849. {
  850. name: "仅 contents,无 systemInstruction",
  851. input: `{
  852. "contents":[{"parts":[{"text":"Hello, how are you?"}]}]
  853. }`,
  854. wantGt0: true,
  855. },
  856. {
  857. name: "空 parts",
  858. input: `{"contents":[{"parts":[]}]}`,
  859. wantGt0: false,
  860. wantExact: intPtr(0),
  861. },
  862. {
  863. name: "非文本 parts(inlineData)",
  864. input: `{"contents":[{"parts":[{"inlineData":{"mimeType":"image/png"}}]}]}`,
  865. wantGt0: false,
  866. wantExact: intPtr(0),
  867. },
  868. {
  869. name: "空白文本",
  870. input: `{"contents":[{"parts":[{"text":" "}]}]}`,
  871. wantGt0: false,
  872. wantExact: intPtr(0),
  873. },
  874. }
  875. for _, tt := range tests {
  876. t.Run(tt.name, func(t *testing.T) {
  877. got := estimateGeminiCountTokens([]byte(tt.input))
  878. if tt.wantExact != nil {
  879. if got != *tt.wantExact {
  880. t.Errorf("期望精确值 %d,实际 %d", *tt.wantExact, got)
  881. }
  882. return
  883. }
  884. if tt.wantGt0 && got <= 0 {
  885. t.Errorf("期望返回 > 0,实际 %d", got)
  886. }
  887. if !tt.wantGt0 && got != 0 {
  888. t.Errorf("期望返回 0,实际 %d", got)
  889. }
  890. })
  891. }
  892. }
  893. // ---------------------------------------------------------------------------
  894. // Task 8.3 — ParseGeminiRateLimitResetTime 测试
  895. // ---------------------------------------------------------------------------
  896. func TestParseGeminiRateLimitResetTime(t *testing.T) {
  897. tests := []struct {
  898. name string
  899. input string
  900. wantNil bool
  901. approxDelta int64 // 预期的 (返回值 - now) 大约是多少秒
  902. }{
  903. {
  904. name: "正常 quotaResetDelay",
  905. input: `{"error":{"details":[{"metadata":{"quotaResetDelay":"12.345s"}}]}}`,
  906. wantNil: false,
  907. approxDelta: 13, // 向上取整 12.345 -> 13
  908. },
  909. {
  910. name: "daily quota",
  911. input: `{"error":{"message":"quota per day exceeded"}}`,
  912. wantNil: false,
  913. approxDelta: -1, // 不检查精确 delta,仅检查非 nil
  914. },
  915. {
  916. name: "无 details 且无 regex 匹配",
  917. input: `{"error":{"message":"rate limit"}}`,
  918. wantNil: true,
  919. },
  920. {
  921. name: "regex 回退匹配",
  922. input: `Please retry in 30s`,
  923. wantNil: false,
  924. approxDelta: 30,
  925. },
  926. {
  927. name: "完全无匹配",
  928. input: `{"error":{"code":429}}`,
  929. wantNil: true,
  930. },
  931. {
  932. name: "非法 JSON 但 regex 回退仍工作",
  933. input: `not json but Please retry in 10s`,
  934. wantNil: false,
  935. approxDelta: 10,
  936. },
  937. }
  938. for _, tt := range tests {
  939. t.Run(tt.name, func(t *testing.T) {
  940. now := time.Now().Unix()
  941. got := ParseGeminiRateLimitResetTime([]byte(tt.input))
  942. if tt.wantNil {
  943. if got != nil {
  944. t.Fatalf("期望返回 nil,实际返回 %d", *got)
  945. }
  946. return
  947. }
  948. if got == nil {
  949. t.Fatalf("期望返回非 nil,实际返回 nil")
  950. }
  951. // approxDelta == -1 表示只检查非 nil,不检查具体值(如 daily quota 场景)
  952. if tt.approxDelta == -1 {
  953. // 仅验证返回的时间戳在合理范围内(未来的某个时间)
  954. if *got < now {
  955. t.Errorf("期望返回的时间戳 >= now(%d),实际 %d", now, *got)
  956. }
  957. return
  958. }
  959. // 使用 +/-2 秒容差进行范围检查
  960. delta := *got - now
  961. if delta < tt.approxDelta-2 || delta > tt.approxDelta+2 {
  962. t.Errorf("期望 delta 约为 %d 秒(+/-2),实际 delta = %d 秒(返回值=%d, now=%d)",
  963. tt.approxDelta, delta, *got, now)
  964. }
  965. })
  966. }
  967. }
  968. // TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText guards the
  969. // tool→text ordering in the Gemini→Anthropic (messages) streaming bridge. When
  970. // Gemini emits a functionCall part followed by a text part, the tool_use content
  971. // block must be closed before the text block opens; otherwise the Anthropic SSE
  972. // stream contains overlapping content blocks. The chat-completions sibling
  973. // already enforces this via closeOpenTool().
  974. func TestGeminiMessagesHandleStreamingResponse_ClosesToolBlockBeforeText(t *testing.T) {
  975. gin.SetMode(gin.TestMode)
  976. upstreamBody := `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"city":"SF"}}}]}}]}` + "\n\n" +
  977. `data: {"candidates":[{"content":{"parts":[{"text":"All done."}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}` + "\n\n" +
  978. "data: [DONE]\n\n"
  979. resp := &http.Response{
  980. StatusCode: http.StatusOK,
  981. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  982. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  983. }
  984. rec := httptest.NewRecorder()
  985. c, _ := gin.CreateTestContext(rec)
  986. svc := &GeminiMessagesCompatService{}
  987. result, err := svc.handleStreamingResponse(c, resp, time.Now(), "claude-3-5-sonnet")
  988. require.NoError(t, err)
  989. require.NotNil(t, result)
  990. events := parseAnthropicContentBlockEvents(t, rec.Body.String())
  991. // Anthropic allows at most one content block open at a time: every
  992. // content_block_start must be matched by a content_block_stop before the
  993. // next start. Replay the lifecycle and assert there is no overlap.
  994. open := -1
  995. blockTypes := map[int]string{}
  996. textStarted := false
  997. toolClosed := false
  998. toolClosedBeforeText := false
  999. for _, ev := range events {
  1000. switch ev.event {
  1001. case "content_block_start":
  1002. require.Equalf(t, -1, open,
  1003. "content block %d opened while block %d was still open (overlapping blocks)", ev.index, open)
  1004. open = ev.index
  1005. blockTypes[ev.index] = ev.blockType
  1006. if ev.blockType == "text" {
  1007. textStarted = true
  1008. if toolClosed {
  1009. toolClosedBeforeText = true
  1010. }
  1011. }
  1012. case "content_block_stop":
  1013. require.Equalf(t, open, ev.index,
  1014. "content_block_stop index %d does not match the open block %d", ev.index, open)
  1015. if blockTypes[ev.index] == "tool_use" {
  1016. toolClosed = true
  1017. }
  1018. open = -1
  1019. }
  1020. }
  1021. require.True(t, textStarted, "expected a text content block to be emitted after the tool call")
  1022. require.True(t, toolClosedBeforeText, "tool_use block must be closed before the text block starts")
  1023. require.Equal(t, -1, open, "stream ended with a content block still open")
  1024. }
  1025. type anthropicContentBlockEvent struct {
  1026. event string
  1027. index int
  1028. blockType string
  1029. }
  1030. // parseAnthropicContentBlockEvents extracts content_block_start/stop events (with
  1031. // their index and, for starts, the content block type) from an Anthropic SSE body.
  1032. func parseAnthropicContentBlockEvents(t *testing.T, raw string) []anthropicContentBlockEvent {
  1033. t.Helper()
  1034. var events []anthropicContentBlockEvent
  1035. for _, chunk := range strings.Split(raw, "\n\n") {
  1036. var eventName, dataLine string
  1037. for _, line := range strings.Split(chunk, "\n") {
  1038. switch {
  1039. case strings.HasPrefix(line, "event:"):
  1040. eventName = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
  1041. case strings.HasPrefix(line, "data:"):
  1042. dataLine = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
  1043. }
  1044. }
  1045. if eventName != "content_block_start" && eventName != "content_block_stop" {
  1046. continue
  1047. }
  1048. var payload struct {
  1049. Index int `json:"index"`
  1050. ContentBlock struct {
  1051. Type string `json:"type"`
  1052. } `json:"content_block"`
  1053. }
  1054. require.NoError(t, json.Unmarshal([]byte(dataLine), &payload))
  1055. events = append(events, anthropicContentBlockEvent{
  1056. event: eventName,
  1057. index: payload.Index,
  1058. blockType: payload.ContentBlock.Type,
  1059. })
  1060. }
  1061. return events
  1062. }