gateway_prompt_test.go 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537
  1. package service
  2. import (
  3. "encoding/json"
  4. "strings"
  5. "testing"
  6. "github.com/Wei-Shaw/sub2api/internal/pkg/claude"
  7. "github.com/stretchr/testify/require"
  8. "github.com/tidwall/gjson"
  9. )
  10. func TestIsClaudeCodeClient(t *testing.T) {
  11. // 合法的 legacy 格式 metadata.user_id(64位 hex + account uuid + session uuid)
  12. legacyUserID := "user_a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2_account_550e8400-e29b-41d4-a716-446655440000_session_123e4567-e89b-12d3-a456-426614174000"
  13. // 合法的 JSON 格式 metadata.user_id(2.1.78+ 版本)
  14. jsonUserID := `{"device_id":"a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2","account_uuid":"550e8400-e29b-41d4-a716-446655440000","session_id":"123e4567-e89b-12d3-a456-426614174000"}`
  15. tests := []struct {
  16. name string
  17. userAgent string
  18. metadataUserID string
  19. want bool
  20. }{
  21. {
  22. name: "Claude Code client with legacy user_id",
  23. userAgent: "claude-cli/1.0.62 (darwin; arm64)",
  24. metadataUserID: legacyUserID,
  25. want: true,
  26. },
  27. {
  28. name: "Claude Code client with JSON user_id",
  29. userAgent: "claude-cli/2.1.92 (external, cli)",
  30. metadataUserID: jsonUserID,
  31. want: true,
  32. },
  33. {
  34. name: "Claude Code case insensitive UA",
  35. userAgent: "Claude-CLI/2.0.0",
  36. metadataUserID: legacyUserID,
  37. want: true,
  38. },
  39. {
  40. name: "Missing metadata user_id",
  41. userAgent: "claude-cli/1.0.0",
  42. metadataUserID: "",
  43. want: false,
  44. },
  45. {
  46. name: "Claude CLI UA with invalid user_id format",
  47. userAgent: "claude-cli/2.0.0",
  48. metadataUserID: "fake-user-id-12345",
  49. want: false,
  50. },
  51. {
  52. name: "Different user agent with valid user_id",
  53. userAgent: "curl/7.68.0",
  54. metadataUserID: legacyUserID,
  55. want: false,
  56. },
  57. {
  58. name: "Empty user agent",
  59. userAgent: "",
  60. metadataUserID: legacyUserID,
  61. want: false,
  62. },
  63. {
  64. name: "Similar but not Claude CLI",
  65. userAgent: "claude-api/1.0.0",
  66. metadataUserID: legacyUserID,
  67. want: false,
  68. },
  69. {
  70. name: "Opencode spoofing UA with arbitrary user_id",
  71. userAgent: "claude-cli/2.1.92",
  72. metadataUserID: "session_abc",
  73. want: false,
  74. },
  75. }
  76. for _, tt := range tests {
  77. t.Run(tt.name, func(t *testing.T) {
  78. got := isClaudeCodeClient(tt.userAgent, tt.metadataUserID)
  79. require.Equal(t, tt.want, got)
  80. })
  81. }
  82. }
  83. func TestSystemIncludesClaudeCodePrompt(t *testing.T) {
  84. tests := []struct {
  85. name string
  86. system any
  87. want bool
  88. }{
  89. {
  90. name: "nil system",
  91. system: nil,
  92. want: false,
  93. },
  94. {
  95. name: "empty string",
  96. system: "",
  97. want: false,
  98. },
  99. {
  100. name: "string with Claude Code prompt",
  101. system: claudeCodeSystemPrompt,
  102. want: true,
  103. },
  104. {
  105. name: "string with different content",
  106. system: "You are a helpful assistant.",
  107. want: false,
  108. },
  109. {
  110. name: "empty array",
  111. system: []any{},
  112. want: false,
  113. },
  114. {
  115. name: "array with Claude Code prompt",
  116. system: []any{
  117. map[string]any{
  118. "type": "text",
  119. "text": claudeCodeSystemPrompt,
  120. },
  121. },
  122. want: true,
  123. },
  124. {
  125. name: "array with Claude Code prompt in second position",
  126. system: []any{
  127. map[string]any{"type": "text", "text": "First prompt"},
  128. map[string]any{"type": "text", "text": claudeCodeSystemPrompt},
  129. },
  130. want: true,
  131. },
  132. {
  133. name: "array without Claude Code prompt",
  134. system: []any{
  135. map[string]any{"type": "text", "text": "Custom prompt"},
  136. },
  137. want: false,
  138. },
  139. {
  140. name: "array with partial match (should not match)",
  141. system: []any{
  142. map[string]any{"type": "text", "text": "You are Claude"},
  143. },
  144. want: false,
  145. },
  146. // json.RawMessage cases (conversion path: ForwardAsResponses / ForwardAsChatCompletions)
  147. {
  148. name: "json.RawMessage string with Claude Code prompt",
  149. system: json.RawMessage(`"` + claudeCodeSystemPrompt + `"`),
  150. want: true,
  151. },
  152. {
  153. name: "json.RawMessage string without Claude Code prompt",
  154. system: json.RawMessage(`"You are a helpful assistant"`),
  155. want: false,
  156. },
  157. {
  158. name: "json.RawMessage nil (empty)",
  159. system: json.RawMessage(nil),
  160. want: false,
  161. },
  162. {
  163. name: "json.RawMessage empty string",
  164. system: json.RawMessage(`""`),
  165. want: false,
  166. },
  167. }
  168. for _, tt := range tests {
  169. t.Run(tt.name, func(t *testing.T) {
  170. got := systemIncludesClaudeCodePrompt(tt.system)
  171. require.Equal(t, tt.want, got)
  172. })
  173. }
  174. }
  175. func TestInjectClaudeCodePrompt(t *testing.T) {
  176. claudePrefix := strings.TrimSpace(claudeCodeSystemPrompt)
  177. tests := []struct {
  178. name string
  179. body string
  180. system any
  181. wantSystemLen int
  182. wantFirstText string
  183. wantSecondText string
  184. }{
  185. {
  186. name: "nil system",
  187. body: `{"model":"claude-3"}`,
  188. system: nil,
  189. wantSystemLen: 1,
  190. wantFirstText: claudeCodeSystemPrompt,
  191. },
  192. {
  193. name: "empty string system",
  194. body: `{"model":"claude-3"}`,
  195. system: "",
  196. wantSystemLen: 1,
  197. wantFirstText: claudeCodeSystemPrompt,
  198. },
  199. {
  200. name: "string system",
  201. body: `{"model":"claude-3"}`,
  202. system: "Custom prompt",
  203. wantSystemLen: 2,
  204. wantFirstText: claudeCodeSystemPrompt,
  205. wantSecondText: claudePrefix + "\n\nCustom prompt",
  206. },
  207. {
  208. name: "string system equals Claude Code prompt",
  209. body: `{"model":"claude-3"}`,
  210. system: claudeCodeSystemPrompt,
  211. wantSystemLen: 1,
  212. wantFirstText: claudeCodeSystemPrompt,
  213. },
  214. {
  215. name: "array system",
  216. body: `{"model":"claude-3"}`,
  217. system: []any{map[string]any{"type": "text", "text": "Custom"}},
  218. // Claude Code + Custom = 2
  219. wantSystemLen: 2,
  220. wantFirstText: claudeCodeSystemPrompt,
  221. wantSecondText: claudePrefix + "\n\nCustom",
  222. },
  223. {
  224. name: "array system with existing Claude Code prompt (should dedupe)",
  225. body: `{"model":"claude-3"}`,
  226. system: []any{
  227. map[string]any{"type": "text", "text": claudeCodeSystemPrompt},
  228. map[string]any{"type": "text", "text": "Other"},
  229. },
  230. // Claude Code at start + Other = 2 (deduped)
  231. wantSystemLen: 2,
  232. wantFirstText: claudeCodeSystemPrompt,
  233. wantSecondText: claudePrefix + "\n\nOther",
  234. },
  235. {
  236. name: "empty array",
  237. body: `{"model":"claude-3"}`,
  238. system: []any{},
  239. wantSystemLen: 1,
  240. wantFirstText: claudeCodeSystemPrompt,
  241. },
  242. // json.RawMessage cases (conversion path: ForwardAsResponses / ForwardAsChatCompletions)
  243. {
  244. name: "json.RawMessage string system",
  245. body: `{"model":"claude-3","system":"Custom prompt"}`,
  246. system: json.RawMessage(`"Custom prompt"`),
  247. wantSystemLen: 2,
  248. wantFirstText: claudeCodeSystemPrompt,
  249. wantSecondText: claudePrefix + "\n\nCustom prompt",
  250. },
  251. {
  252. name: "json.RawMessage nil system",
  253. body: `{"model":"claude-3"}`,
  254. system: json.RawMessage(nil),
  255. wantSystemLen: 1,
  256. wantFirstText: claudeCodeSystemPrompt,
  257. },
  258. {
  259. name: "json.RawMessage Claude Code prompt (should not duplicate)",
  260. body: `{"model":"claude-3","system":"` + claudeCodeSystemPrompt + `"}`,
  261. system: json.RawMessage(`"` + claudeCodeSystemPrompt + `"`),
  262. wantSystemLen: 1,
  263. wantFirstText: claudeCodeSystemPrompt,
  264. },
  265. }
  266. for _, tt := range tests {
  267. t.Run(tt.name, func(t *testing.T) {
  268. result := injectClaudeCodePrompt([]byte(tt.body), tt.system)
  269. var parsed map[string]any
  270. err := json.Unmarshal(result, &parsed)
  271. require.NoError(t, err)
  272. system, ok := parsed["system"].([]any)
  273. require.True(t, ok, "system should be an array")
  274. require.Len(t, system, tt.wantSystemLen)
  275. first, ok := system[0].(map[string]any)
  276. require.True(t, ok)
  277. require.Equal(t, tt.wantFirstText, first["text"])
  278. require.Equal(t, "text", first["type"])
  279. // Check cache_control
  280. cc, ok := first["cache_control"].(map[string]any)
  281. require.True(t, ok)
  282. require.Equal(t, "ephemeral", cc["type"])
  283. if tt.wantSecondText != "" && len(system) > 1 {
  284. second, ok := system[1].(map[string]any)
  285. require.True(t, ok)
  286. require.Equal(t, tt.wantSecondText, second["text"])
  287. }
  288. })
  289. }
  290. }
  291. func TestRewriteSystemForNonClaudeCode(t *testing.T) {
  292. tests := []struct {
  293. name string
  294. body string
  295. system any
  296. wantSystemText string // system array 第一个 block 的 text
  297. wantMessagesLen int // messages 数组长度
  298. wantFirstMsgRole string // 第一条消息的 role
  299. wantFirstMsgText string // 第一条消息的 content[0].text
  300. wantAckMsgText string // 第二条消息的 content[0].text
  301. }{
  302. {
  303. name: "nil system - no messages injected",
  304. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  305. system: nil,
  306. wantSystemText: claudeCodeSystemPrompt,
  307. wantMessagesLen: 1, // 原始 1 条消息,不注入
  308. },
  309. {
  310. name: "empty string system - no messages injected",
  311. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  312. system: "",
  313. wantSystemText: claudeCodeSystemPrompt,
  314. wantMessagesLen: 1,
  315. },
  316. {
  317. name: "custom string system - migrated to messages",
  318. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  319. system: "You are a personal assistant running inside OpenClaw.",
  320. wantSystemText: claudeCodeSystemPrompt,
  321. wantMessagesLen: 3, // instruction + ack + original
  322. wantFirstMsgRole: "user",
  323. wantFirstMsgText: "[System Instructions]\nYou are a personal assistant running inside OpenClaw.",
  324. wantAckMsgText: "Understood. I will follow these instructions.",
  325. },
  326. {
  327. name: "system equals Claude Code prompt - no messages injected",
  328. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  329. system: claudeCodeSystemPrompt,
  330. wantSystemText: claudeCodeSystemPrompt,
  331. wantMessagesLen: 1,
  332. },
  333. {
  334. name: "array system with custom blocks - text joined and migrated",
  335. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  336. system: []any{
  337. map[string]any{"type": "text", "text": "First instruction"},
  338. map[string]any{"type": "text", "text": "Second instruction"},
  339. },
  340. wantSystemText: claudeCodeSystemPrompt,
  341. wantMessagesLen: 3,
  342. wantFirstMsgRole: "user",
  343. wantFirstMsgText: "[System Instructions]\nFirst instruction\n\nSecond instruction",
  344. wantAckMsgText: "Understood. I will follow these instructions.",
  345. },
  346. {
  347. name: "empty array system - no messages injected",
  348. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  349. system: []any{},
  350. wantSystemText: claudeCodeSystemPrompt,
  351. wantMessagesLen: 1,
  352. },
  353. {
  354. name: "json.RawMessage string system",
  355. body: `{"model":"claude-3","system":"Custom prompt","messages":[{"role":"user","content":"hello"}]}`,
  356. system: json.RawMessage(`"Custom prompt"`),
  357. wantSystemText: claudeCodeSystemPrompt,
  358. wantMessagesLen: 3,
  359. wantFirstMsgRole: "user",
  360. wantFirstMsgText: "[System Instructions]\nCustom prompt",
  361. wantAckMsgText: "Understood. I will follow these instructions.",
  362. },
  363. {
  364. name: "json.RawMessage nil system",
  365. body: `{"model":"claude-3","messages":[{"role":"user","content":"hello"}]}`,
  366. system: json.RawMessage(nil),
  367. wantSystemText: claudeCodeSystemPrompt,
  368. wantMessagesLen: 1,
  369. },
  370. {
  371. name: "multiple original messages preserved",
  372. body: `{"model":"claude-3","messages":[{"role":"user","content":"msg1"},{"role":"assistant","content":"resp1"},{"role":"user","content":"msg2"}]}`,
  373. system: "Be helpful",
  374. wantSystemText: claudeCodeSystemPrompt,
  375. wantMessagesLen: 5, // 2 injected + 3 original
  376. wantFirstMsgRole: "user",
  377. wantFirstMsgText: "[System Instructions]\nBe helpful",
  378. wantAckMsgText: "Understood. I will follow these instructions.",
  379. },
  380. }
  381. for _, tt := range tests {
  382. t.Run(tt.name, func(t *testing.T) {
  383. result := rewriteSystemForNonClaudeCode([]byte(tt.body), tt.system)
  384. var parsed map[string]any
  385. err := json.Unmarshal(result, &parsed)
  386. require.NoError(t, err)
  387. // system 应为 array 格式,对齐真实 Claude Code CLI 的 3-block 形态:
  388. // [0] billing attribution block (x-anthropic-billing-header: cc_version=...;)
  389. // [1] Claude Code 身份前缀 block (不带 cache_control)
  390. // [2] 工具无关的通用提示词扩充 block (带 cache_control,作为缓存断点)
  391. systemArr, ok := parsed["system"].([]any)
  392. require.True(t, ok, "system should be an array, got %T", parsed["system"])
  393. require.Len(t, systemArr, 3, "system array should have exactly 3 blocks (billing + cc prompt + expansion)")
  394. billingBlock, ok := systemArr[0].(map[string]any)
  395. require.True(t, ok)
  396. require.Equal(t, "text", billingBlock["type"])
  397. require.Contains(t, billingBlock["text"], "x-anthropic-billing-header:")
  398. require.Contains(t, billingBlock["text"], "cc_version=")
  399. require.Contains(t, billingBlock["text"], "cc_entrypoint=cli")
  400. // 新版 CLI 已取消 cch=... 签名字段,注入的 billing block 不应再带 cch。
  401. require.NotContains(t, billingBlock["text"], "cch=")
  402. systemBlock, ok := systemArr[1].(map[string]any)
  403. require.True(t, ok)
  404. require.Equal(t, "text", systemBlock["type"])
  405. require.Equal(t, tt.wantSystemText, systemBlock["text"])
  406. _, hasCC := systemBlock["cache_control"]
  407. require.False(t, hasCC, "身份前缀 block 不应带 cache_control(断点落在扩充块)")
  408. expansionBlock, ok := systemArr[2].(map[string]any)
  409. require.True(t, ok)
  410. require.Equal(t, "text", expansionBlock["type"])
  411. require.Equal(t, claudeCodeSystemPromptExpansion, expansionBlock["text"])
  412. cc, ok := expansionBlock["cache_control"].(map[string]any)
  413. require.True(t, ok, "expansion block should have cache_control")
  414. require.Equal(t, "ephemeral", cc["type"])
  415. // 检查 messages
  416. messages, ok := parsed["messages"].([]any)
  417. require.True(t, ok, "messages should be an array")
  418. require.Len(t, messages, tt.wantMessagesLen)
  419. if tt.wantFirstMsgRole != "" && len(messages) >= 2 {
  420. // 检查注入的 instruction 消息
  421. firstMsg, ok := messages[0].(map[string]any)
  422. require.True(t, ok)
  423. require.Equal(t, tt.wantFirstMsgRole, firstMsg["role"])
  424. firstContent, ok := firstMsg["content"].([]any)
  425. require.True(t, ok)
  426. require.Len(t, firstContent, 1)
  427. firstBlock, ok := firstContent[0].(map[string]any)
  428. require.True(t, ok)
  429. require.Equal(t, tt.wantFirstMsgText, firstBlock["text"])
  430. // 检查注入的 ack 消息
  431. ackMsg, ok := messages[1].(map[string]any)
  432. require.True(t, ok)
  433. require.Equal(t, "assistant", ackMsg["role"])
  434. ackContent, ok := ackMsg["content"].([]any)
  435. require.True(t, ok)
  436. require.Len(t, ackContent, 1)
  437. ackBlock, ok := ackContent[0].(map[string]any)
  438. require.True(t, ok)
  439. require.Equal(t, tt.wantAckMsgText, ackBlock["text"])
  440. }
  441. })
  442. }
  443. }
  444. func TestRewriteSystemForNonClaudeCodeWithPrompt_UsesCustomExpansionPrompt(t *testing.T) {
  445. body := []byte(`{"model":"claude-3","system":"Project instructions","messages":[{"role":"user","content":"hello"}]}`)
  446. customPrompt := "Custom Claude OAuth expansion prompt"
  447. result := rewriteSystemForNonClaudeCodeWithPrompt(body, "Project instructions", customPrompt)
  448. system := gjson.GetBytes(result, "system")
  449. require.True(t, system.IsArray())
  450. require.Len(t, system.Array(), 3)
  451. require.Equal(t, customPrompt, system.Array()[2].Get("text").String())
  452. require.Equal(t, "ephemeral", system.Array()[2].Get("cache_control.type").String())
  453. }
  454. func TestRewriteSystemForNonClaudeCode_PreservesSystemCacheControlOnMigratedMessage(t *testing.T) {
  455. body := []byte(`{"model":"claude-3","system":[{"type":"text","text":"Stable project instructions","cache_control":{"type":"ephemeral","ttl":"1h"}}],"messages":[{"role":"user","content":"hello"}]}`)
  456. system := []any{
  457. map[string]any{
  458. "type": "text",
  459. "text": "Stable project instructions",
  460. "cache_control": map[string]any{"type": "ephemeral", "ttl": "1h"},
  461. },
  462. }
  463. result := rewriteSystemForNonClaudeCode(body, system)
  464. require.Equal(t, "[System Instructions]\nStable project instructions", gjson.GetBytes(result, "messages.0.content.0.text").String())
  465. require.Equal(t, "ephemeral", gjson.GetBytes(result, "messages.0.content.0.cache_control.type").String())
  466. require.Equal(t, "1h", gjson.GetBytes(result, "messages.0.content.0.cache_control.ttl").String())
  467. }
  468. func TestRewriteSystemForNonClaudeCode_LeavesMigratedMessageUncachedWithoutSystemBreakpoint(t *testing.T) {
  469. body := []byte(`{"model":"claude-3","system":[{"type":"text","text":"Project instructions"}],"messages":[{"role":"user","content":"hello"}]}`)
  470. system := []any{
  471. map[string]any{"type": "text", "text": "Project instructions"},
  472. }
  473. result := rewriteSystemForNonClaudeCode(body, system)
  474. require.False(t, gjson.GetBytes(result, "messages.0.content.0.cache_control").Exists())
  475. }
  476. func TestRewriteSystemForNonClaudeCodeWithPromptBlocks_UsesConfiguredBlocks(t *testing.T) {
  477. body := []byte(`{"model":"claude-3","system":"Project instructions","messages":[{"role":"user","content":"hello"}]}`)
  478. blocks := `{
  479. "blocks": [
  480. {"type":"text","text":"prefix {cc_version}.{fp}","cache_control":true},
  481. {"enabled":false,"type":"text","text":"disabled"},
  482. {"type":"text","text":"{claude_code_system_prompt}"},
  483. {"type":"text","text":"tail","cache_control":{"type":"ephemeral","ttl":"1h"}}
  484. ]
  485. }`
  486. result := rewriteSystemForNonClaudeCodeWithPromptBlocks(body, "Project instructions", "", blocks)
  487. system := gjson.GetBytes(result, "system")
  488. require.True(t, system.IsArray())
  489. arr := system.Array()
  490. require.Len(t, arr, 3)
  491. require.Contains(t, arr[0].Get("text").String(), "prefix "+claude.CLICurrentVersion+".")
  492. require.Equal(t, "ephemeral", arr[0].Get("cache_control.type").String())
  493. require.Equal(t, claude.DefaultCacheControlTTL, arr[0].Get("cache_control.ttl").String())
  494. require.Equal(t, claudeCodeSystemPrompt, arr[1].Get("text").String())
  495. require.False(t, arr[1].Get("cache_control").Exists())
  496. require.Equal(t, "tail", arr[2].Get("text").String())
  497. require.Equal(t, "1h", arr[2].Get("cache_control.ttl").String())
  498. }