| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "strings" |
| 7 | "testing" |
| 8 | ) |
| 9 | |
| 10 | // --- SanitizeToolPairing --- |
| 11 | |
| 12 | // toolIDsAnswered reports whether every assistant tool_call id has a following |
| 13 | // tool message answering it — the contract the OpenAI/DeepSeek API enforces. |
| 14 | func toolIDsAnswered(msgs []Message) bool { |
| 15 | answered := map[string]bool{} |
| 16 | for _, m := range msgs { |
| 17 | if m.Role == RoleTool { |
| 18 | answered[m.ToolCallID] = true |
| 19 | } |
| 20 | } |
| 21 | for _, m := range msgs { |
| 22 | for _, tc := range m.ToolCalls { |
| 23 | if !answered[tc.ID] { |
| 24 | return false |
| 25 | } |
| 26 | } |
| 27 | } |
| 28 | return true |
| 29 | } |
| 30 | |
| 31 | func TestSanitizeToolPairingBackfillsDanglingCall(t *testing.T) { |
| 32 | in := []Message{ |
| 33 | {Role: RoleUser, Content: "list files"}, |
| 34 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "ls"}}}, |
| 35 | {Role: RoleUser, Content: "never mind"}, |
| 36 | } |
| 37 | out := SanitizeToolPairing(in) |
| 38 | if !toolIDsAnswered(out) { |
| 39 | t.Fatalf("dangling tool_call left unanswered: %+v", out) |
| 40 | } |
| 41 | // The backfilled result sits right after the assistant turn, keyed to its id. |
| 42 | if out[2].Role != RoleTool || out[2].ToolCallID != "c1" { |
| 43 | t.Fatalf("expected a backfilled tool result for c1 at index 2, got %+v", out[2]) |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | func TestSanitizeToolPairingKeepsCallOrderAndMultiple(t *testing.T) { |
| 48 | in := []Message{ |
| 49 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "a"}, {ID: "b"}, {ID: "c"}}}, |
| 50 | {Role: RoleTool, ToolCallID: "b", Content: "B"}, // out of order, c missing |
| 51 | {Role: RoleTool, ToolCallID: "a", Content: "A"}, |
| 52 | } |
| 53 | out := SanitizeToolPairing(in) |
| 54 | if !toolIDsAnswered(out) { |
| 55 | t.Fatalf("not all calls answered: %+v", out) |
| 56 | } |
| 57 | gotOrder := []string{out[1].ToolCallID, out[2].ToolCallID, out[3].ToolCallID} |
| 58 | want := []string{"a", "b", "c"} |
| 59 | for i := range want { |
| 60 | if gotOrder[i] != want[i] { |
| 61 | t.Fatalf("tool results out of call order: got %v want %v", gotOrder, want) |
| 62 | } |
| 63 | } |
| 64 | } |
| 65 | |
| 66 | func TestSanitizeToolPairingDropsOrphanToolMessage(t *testing.T) { |
| 67 | in := []Message{ |
| 68 | {Role: RoleUser, Content: "hi"}, |
| 69 | {Role: RoleTool, ToolCallID: "ghost", Content: "leftover"}, // no preceding call |
| 70 | {Role: RoleAssistant, Content: "hello"}, |
| 71 | } |
| 72 | out := SanitizeToolPairing(in) |
| 73 | for _, m := range out { |
| 74 | if m.Role == RoleTool { |
| 75 | t.Fatalf("orphan tool message survived: %+v", out) |
| 76 | } |
| 77 | } |
| 78 | if len(out) != 2 { |
| 79 | t.Fatalf("want 2 messages after dropping the orphan, got %d: %+v", len(out), out) |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | func TestSanitizeToolPairingLeavesWellFormedUnchanged(t *testing.T) { |
| 84 | in := []Message{ |
| 85 | {Role: RoleSystem, Content: "sys"}, |
| 86 | {Role: RoleUser, Content: "q"}, |
| 87 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "ls"}}}, |
| 88 | {Role: RoleTool, ToolCallID: "c1", Name: "ls", Content: "main.go"}, |
| 89 | {Role: RoleAssistant, Content: "done"}, |
| 90 | } |
| 91 | out := SanitizeToolPairing(in) |
| 92 | if len(out) != len(in) { |
| 93 | t.Fatalf("well-formed history changed length: %d -> %d", len(in), len(out)) |
| 94 | } |
| 95 | if &out[0] != &in[0] { |
| 96 | t.Fatalf("well-formed history should return the input slice without allocating") |
| 97 | } |
| 98 | for i := range in { |
| 99 | if out[i].Role != in[i].Role || out[i].Content != in[i].Content || out[i].ToolCallID != in[i].ToolCallID { |
| 100 | t.Fatalf("well-formed message %d mutated: %+v -> %+v", i, in[i], out[i]) |
| 101 | } |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | func TestModelMessagesAndSanitizeDropLocalOnlyInterruptedOutput(t *testing.T) { |
| 106 | local := Message{ |
| 107 | Role: RoleTool, ToolCallID: LocalOnlyToolID, Name: LocalOnlyToolName, |
| 108 | Content: "partial answer", ReasoningContent: "partial reasoning", LocalOnly: true, |
| 109 | ToolCalls: []ToolCall{{ID: "partial", Name: "write_file"}}, |
| 110 | InterruptedTurn: &InterruptedTurnRecovery{Pending: true, InterruptedTools: []string{"write_file"}}, |
| 111 | } |
| 112 | in := []Message{ |
| 113 | {Role: RoleUser, Content: "task"}, |
| 114 | local, |
| 115 | {Role: RoleUser, Content: "continue"}, |
| 116 | } |
| 117 | model := ModelMessages(in) |
| 118 | if len(model) != 2 || model[0].Content != "task" || model[1].Content != "continue" { |
| 119 | t.Fatalf("ModelMessages leaked or reordered local-only record: %+v", model) |
| 120 | } |
| 121 | wire := SanitizeToolPairing(in) |
| 122 | if len(wire) != 2 || wire[0].Content != "task" || wire[1].Content != "continue" { |
| 123 | t.Fatalf("SanitizeToolPairing leaked local-only record: %+v", wire) |
| 124 | } |
| 125 | session := NormalizeSessionMessages(in) |
| 126 | if len(session) != len(in) || !session[1].LocalOnly || session[1].Content != local.Content { |
| 127 | t.Fatalf("session normalization did not preserve local display: %+v", session) |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | func TestDecisionReceiptIsDurableButProviderExcluded(t *testing.T) { |
| 132 | receipt := &DecisionReceipt{ID: "approval-1", Kind: "tool", Tool: "write_file", Subject: "src/app.go", Outcome: "allow_once"} |
| 133 | in := []Message{ |
| 134 | {Role: RoleUser, Content: "edit the app"}, |
| 135 | {Role: RoleAssistant, LocalOnly: true, DecisionReceipt: receipt}, |
| 136 | {Role: RoleAssistant, Content: "done"}, |
| 137 | } |
| 138 | model := ModelMessages(in) |
| 139 | if len(model) != 2 || model[0].Content != "edit the app" || model[1].Content != "done" { |
| 140 | t.Fatalf("provider messages leaked decision receipt: %+v", model) |
| 141 | } |
| 142 | if len(in) != 3 || in[1].DecisionReceipt != receipt || !in[1].LocalOnly { |
| 143 | t.Fatalf("stored receipt was not preserved: %+v", in) |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | func TestAttachedDecisionReceiptPreservesCurrentAndLegacyToolPairing(t *testing.T) { |
| 148 | receipt := &DecisionReceipt{ID: "approval-1", Kind: "tool", Tool: "bash", Outcome: "allow_once"} |
| 149 | stored := []Message{ |
| 150 | {Role: RoleUser, Content: "run it"}, |
| 151 | { |
| 152 | Role: RoleAssistant, |
| 153 | ToolCalls: []ToolCall{{Name: "bash", Arguments: `{}`}}, |
| 154 | DecisionReceipts: []*DecisionReceipt{receipt}, |
| 155 | }, |
| 156 | {Role: RoleTool, Name: "bash", Content: "ok"}, |
| 157 | } |
| 158 | |
| 159 | current := SanitizeToolPairing(ModelMessages(stored)) |
| 160 | if len(current) != 3 || current[2].Content != "ok" { |
| 161 | t.Fatalf("current reader changed the valid tool turn: %+v", current) |
| 162 | } |
| 163 | if len(current[1].DecisionReceipts) != 0 { |
| 164 | t.Fatalf("provider-visible message leaked local decision metadata: %+v", current[1]) |
| 165 | } |
| 166 | |
| 167 | // Older binaries ignore the new metadata field. The remaining legacy view |
| 168 | // must still contain the same adjacent assistant/result pair, including the |
| 169 | // positional pairing used by providers that omit tool-call IDs. |
| 170 | legacy := append([]Message(nil), stored...) |
| 171 | legacy[1].DecisionReceipts = nil |
| 172 | legacy = SanitizeToolPairing(legacy) |
| 173 | if len(legacy) != 3 || legacy[1].Role != RoleAssistant || legacy[2].Role != RoleTool || legacy[2].Content != "ok" { |
| 174 | t.Fatalf("legacy reader lost the actual tool result: %+v", legacy) |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | func TestNormalizeSessionMessagesMigratesInterleavedDecisionReceipt(t *testing.T) { |
| 179 | receipt := &DecisionReceipt{ID: "approval-1", Kind: "tool", Tool: "bash", Outcome: "allow_once"} |
| 180 | old := []Message{ |
| 181 | {Role: RoleUser, Content: "run it"}, |
| 182 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "call-1", Name: "bash", Arguments: `{}`}}}, |
| 183 | {Role: RoleAssistant, LocalOnly: true, DecisionReceipt: receipt}, |
| 184 | {Role: RoleTool, ToolCallID: "call-1", Name: "bash", Content: "actual result"}, |
| 185 | } |
| 186 | |
| 187 | got := NormalizeSessionMessages(old) |
| 188 | if len(got) != 3 { |
| 189 | t.Fatalf("migrated messages = %d, want receipt folded into assistant: %+v", len(got), got) |
| 190 | } |
| 191 | if len(got[1].DecisionReceipts) != 1 || got[1].DecisionReceipts[0] != receipt { |
| 192 | t.Fatalf("migrated assistant receipt = %+v, want original receipt", got[1].DecisionReceipts) |
| 193 | } |
| 194 | if got[2].Role != RoleTool || got[2].Content != "actual result" || strings.Contains(got[2].Content, "interrupted") { |
| 195 | t.Fatalf("migrated tool result = %+v, want the actual result without a placeholder", got[2]) |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | func TestModelMessagesUsesProviderContentWithoutMutatingStoredMessage(t *testing.T) { |
| 200 | stored := []Message{{ |
| 201 | Role: RoleUser, |
| 202 | Content: "fix the bug", |
| 203 | ProviderContent: "<reasoning-language>zh</reasoning-language>\n\nfix the bug", |
| 204 | }} |
| 205 | |
| 206 | model := ModelMessages(stored) |
| 207 | if len(model) != 1 { |
| 208 | t.Fatalf("ModelMessages length = %d, want 1", len(model)) |
| 209 | } |
| 210 | if got := model[0].Content; got != stored[0].ProviderContent { |
| 211 | t.Fatalf("model content = %q, want provider content %q", got, stored[0].ProviderContent) |
| 212 | } |
| 213 | if model[0].ProviderContent != "" { |
| 214 | t.Fatalf("provider-local field leaked into model message: %q", model[0].ProviderContent) |
| 215 | } |
| 216 | if stored[0].Content != "fix the bug" || stored[0].ProviderContent == "" { |
| 217 | t.Fatalf("stored message was mutated: %+v", stored[0]) |
| 218 | } |
| 219 | } |
| 220 | |
| 221 | func TestModelMessagesStripsRawContentWithoutChangingLegacyContent(t *testing.T) { |
| 222 | const rendered = "<reasoning-language>zh</reasoning-language>\n\nfix the bug" |
| 223 | stored := []Message{{Role: RoleUser, Content: rendered, RawContent: "fix the bug"}} |
| 224 | |
| 225 | model := ModelMessages(stored) |
| 226 | if len(model) != 1 || model[0].Content != rendered { |
| 227 | t.Fatalf("provider-visible content changed: %+v", model) |
| 228 | } |
| 229 | if model[0].RawContent != "" { |
| 230 | t.Fatalf("raw display metadata leaked into provider request: %+v", model[0]) |
| 231 | } |
| 232 | if stored[0].RawContent != "fix the bug" || stored[0].Content != rendered { |
| 233 | t.Fatalf("stored message was mutated: %+v", stored[0]) |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | func TestLocalOnlySentinelIsSafeWhenNewFieldsAreIgnoredByLegacyReader(t *testing.T) { |
| 238 | legacyView := []Message{ |
| 239 | {Role: RoleUser, Content: "task"}, |
| 240 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}}}, |
| 241 | {Role: RoleTool, ToolCallID: "c1", Name: "read_file", Content: "ok"}, |
| 242 | // Simulate an older binary: unknown local_only/interrupted_turn JSON fields |
| 243 | // were ignored, leaving only the orphan tool sentinel and partial content. |
| 244 | {Role: RoleTool, ToolCallID: LocalOnlyToolID, Name: LocalOnlyToolName, Content: "partial reasoning that must not leak"}, |
| 245 | {Role: RoleUser, Content: "continue"}, |
| 246 | } |
| 247 | wire := SanitizeToolPairing(legacyView) |
| 248 | if len(wire) != 4 { |
| 249 | t.Fatalf("legacy normalization kept local sentinel: %+v", wire) |
| 250 | } |
| 251 | for _, message := range wire { |
| 252 | if message.ToolCallID == LocalOnlyToolID || strings.Contains(message.Content, "must not leak") { |
| 253 | t.Fatalf("legacy normalization leaked display-only content: %+v", wire) |
| 254 | } |
| 255 | } |
| 256 | } |
| 257 | |
| 258 | func TestNormalizeSessionMessagesPreservesStandaloneToolMessage(t *testing.T) { |
| 259 | in := []Message{ |
| 260 | {Role: RoleSystem, Content: "sys"}, |
| 261 | {Role: RoleUser, Content: "run it"}, |
| 262 | {Role: RoleTool, ToolCallID: "c1", Name: "bash", Content: "large output"}, |
| 263 | } |
| 264 | out := NormalizeSessionMessages(in) |
| 265 | if len(out) != len(in) { |
| 266 | t.Fatalf("session normalization changed length: %d -> %d", len(in), len(out)) |
| 267 | } |
| 268 | if &out[0] != &in[0] { |
| 269 | t.Fatalf("session-safe orphan tool should keep the input slice unchanged") |
| 270 | } |
| 271 | if out[2].Role != RoleTool || out[2].Content != "large output" { |
| 272 | t.Fatalf("standalone tool message was not preserved: %+v", out) |
| 273 | } |
| 274 | } |
| 275 | |
| 276 | func TestNormalizeSessionMessagesPreservesExtraToolResult(t *testing.T) { |
| 277 | in := []Message{ |
| 278 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "bash"}}}, |
| 279 | {Role: RoleTool, ToolCallID: "c1", Name: "bash", Content: "ok"}, |
| 280 | {Role: RoleTool, ToolCallID: "ghost", Name: "bash", Content: "saved extra output"}, |
| 281 | } |
| 282 | out := NormalizeSessionMessages(in) |
| 283 | if len(out) != len(in) { |
| 284 | t.Fatalf("session normalization changed length: %d -> %d", len(in), len(out)) |
| 285 | } |
| 286 | if out[2].ToolCallID != "ghost" || out[2].Content != "saved extra output" { |
| 287 | t.Fatalf("extra stored tool result was not preserved: %+v", out) |
| 288 | } |
| 289 | wire := SanitizeToolPairing(in) |
| 290 | if len(wire) != 2 { |
| 291 | t.Fatalf("wire sanitize should still drop the extra orphan result, got %+v", wire) |
| 292 | } |
| 293 | } |
| 294 | |
| 295 | func TestSanitizeToolPairingClosesTruncatedArgs(t *testing.T) { |
| 296 | cases := []struct{ in, want string }{ |
| 297 | {`{`, `{}`}, |
| 298 | {`{"time": 2`, `{"time": 2}`}, |
| 299 | {`{"command": "ls -la`, `{"command": "ls -la"}`}, |
| 300 | {`{"a": 1,`, `{"a": 1}`}, |
| 301 | {`{"a":`, `{"a":null}`}, |
| 302 | {`{"path": "C:\\tmp\`, `{"path": "C:\\tmp"}`}, |
| 303 | {`{"items": [1, 2`, `{"items": [1, 2]}`}, |
| 304 | {`total garbage`, `{}`}, |
| 305 | {`{"ok": true}`, `{"ok": true}`}, |
| 306 | {``, ``}, |
| 307 | } |
| 308 | for _, c := range cases { |
| 309 | in := []Message{ |
| 310 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "bash", Arguments: c.in}}}, |
| 311 | {Role: RoleTool, ToolCallID: "c1", Content: "r"}, |
| 312 | } |
| 313 | out := SanitizeToolPairing(in) |
| 314 | if got := out[0].ToolCalls[0].Arguments; got != c.want { |
| 315 | t.Errorf("args %q repaired to %q, want %q", c.in, got, c.want) |
| 316 | } |
| 317 | if in[0].ToolCalls[0].Arguments != c.in { |
| 318 | t.Errorf("stored history mutated for %q: %q", c.in, in[0].ToolCalls[0].Arguments) |
| 319 | } |
| 320 | } |
| 321 | } |
| 322 | |
| 323 | func TestBackfillToolCallNamesByID(t *testing.T) { |
| 324 | calls := []ToolCall{{ID: "c1"}, {ID: "c2", Name: "grep"}} |
| 325 | results := []Message{ |
| 326 | {Role: RoleTool, ToolCallID: "c2", Name: "grep"}, |
| 327 | {Role: RoleTool, ToolCallID: "c1", Name: "ls"}, // returned out of call order |
| 328 | } |
| 329 | out := backfillToolCallNames(calls, results) |
| 330 | if out[0].Name != "ls" { |
| 331 | t.Fatalf("empty name not backfilled by id: got %q want ls", out[0].Name) |
| 332 | } |
| 333 | if out[1].Name != "grep" { |
| 334 | t.Fatalf("non-empty name clobbered: got %q want grep", out[1].Name) |
| 335 | } |
| 336 | if calls[0].Name != "" { |
| 337 | t.Fatalf("input slice mutated: %+v", calls) |
| 338 | } |
| 339 | } |
| 340 | |
| 341 | func TestBackfillToolCallNamesPositional(t *testing.T) { |
| 342 | // Empty ids defeat idDistinct, so names pair by position instead. |
| 343 | calls := []ToolCall{{}, {}} |
| 344 | results := []Message{{Role: RoleTool, Name: "ls"}, {Role: RoleTool, Name: "cat"}} |
| 345 | out := backfillToolCallNames(calls, results) |
| 346 | if out[0].Name != "ls" || out[1].Name != "cat" { |
| 347 | t.Fatalf("positional backfill wrong: %+v", out) |
| 348 | } |
| 349 | } |
| 350 | |
| 351 | func TestBackfillToolCallNamesUnpairedStaysEmpty(t *testing.T) { |
| 352 | out := backfillToolCallNames([]ToolCall{{ID: "c1"}}, nil) |
| 353 | if out[0].Name != "" { |
| 354 | t.Fatalf("unpaired call should keep its empty name, got %q", out[0].Name) |
| 355 | } |
| 356 | } |
| 357 | |
| 358 | func TestBackfillToolCallNamesNoEmptyReturnsInput(t *testing.T) { |
| 359 | calls := []ToolCall{{ID: "c1", Name: "ls"}, {ID: "c2", Name: "grep"}} |
| 360 | out := backfillToolCallNames(calls, []Message{{Role: RoleTool, ToolCallID: "c1", Name: "x"}}) |
| 361 | if &out[0] != &calls[0] { |
| 362 | t.Fatalf("no empty names: want the input slice back without copying") |
| 363 | } |
| 364 | } |
| 365 | |
| 366 | func TestSanitizeToolPairingBackfillsEmptyName(t *testing.T) { |
| 367 | in := []Message{ |
| 368 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1"}}}, // old session: name lost |
| 369 | {Role: RoleTool, ToolCallID: "c1", Name: "ls", Content: "main.go"}, |
| 370 | } |
| 371 | out := SanitizeToolPairing(in) |
| 372 | if out[0].ToolCalls[0].Name != "ls" { |
| 373 | t.Fatalf("empty tool-call name not backfilled on replay: %+v", out[0].ToolCalls) |
| 374 | } |
| 375 | if in[0].ToolCalls[0].Name != "" { |
| 376 | t.Fatalf("stored history mutated: %+v", in[0].ToolCalls) |
| 377 | } |
| 378 | } |
| 379 | |
| 380 | func TestSanitizeToolPairingBackfillsMissingToolResultName(t *testing.T) { |
| 381 | in := []Message{ |
| 382 | {Role: RoleAssistant, ToolCalls: []ToolCall{{ID: "c1", Name: "ls"}}}, |
| 383 | {Role: RoleTool, ToolCallID: "c1", Content: "main.go"}, |
| 384 | } |
| 385 | out := SanitizeToolPairing(in) |
| 386 | if out[1].Name != "ls" { |
| 387 | t.Fatalf("missing tool result name not backfilled: %+v", out[1]) |
| 388 | } |
| 389 | if in[1].Name != "" { |
| 390 | t.Fatalf("stored history mutated: %+v", in[1]) |
| 391 | } |
| 392 | } |
| 393 | |
| 394 | // --- Pricing.Cost --- |
| 395 | |
| 396 | func TestPricingCostNil(t *testing.T) { |
| 397 | var p *Pricing |
| 398 | if got := p.Cost(&Usage{PromptTokens: 100}); got != 0 { |
| 399 | t.Errorf("nil Pricing.Cost = %f, want 0", got) |
| 400 | } |
| 401 | } |
| 402 | |
| 403 | func TestPricingCostNilUsage(t *testing.T) { |
| 404 | p := &Pricing{Input: 2.0, Output: 10.0} |
| 405 | if got := p.Cost(nil); got != 0 { |
| 406 | t.Errorf("nil Usage.Cost = %f, want 0", got) |
| 407 | } |
| 408 | } |
| 409 | |
| 410 | func TestPricingCostBothNil(t *testing.T) { |
| 411 | var p *Pricing |
| 412 | if got := p.Cost(nil); got != 0 { |
| 413 | t.Errorf("both nil.Cost = %f, want 0", got) |
| 414 | } |
| 415 | } |
| 416 | |
| 417 | func TestPricingCostCalculation(t *testing.T) { |
| 418 | p := &Pricing{ |
| 419 | CacheHit: 0.5, // ¥0.5 per 1M cached tokens |
| 420 | Input: 2.0, // ¥2.0 per 1M uncached tokens |
| 421 | Output: 10.0, // ¥10.0 per 1M completion tokens |
| 422 | } |
| 423 | u := &Usage{ |
| 424 | CacheHitTokens: 1_000_000, |
| 425 | CacheMissTokens: 500_000, |
| 426 | CompletionTokens: 200_000, |
| 427 | } |
| 428 | // Expected: (1M * 0.5 + 500K * 2.0 + 200K * 10.0) / 1M |
| 429 | // = (0.5 + 1.0 + 2.0) = 3.5 |
| 430 | got := p.Cost(u) |
| 431 | if got != 3.5 { |
| 432 | t.Errorf("Cost = %f, want 3.5", got) |
| 433 | } |
| 434 | } |
| 435 | |
| 436 | func TestPricingCostUsesCacheWriteBillingTier(t *testing.T) { |
| 437 | p := &Pricing{Input: 2.0} |
| 438 | u := &Usage{ |
| 439 | CacheMissTokens: 500_000, |
| 440 | CacheWriteTokens: 100_000, |
| 441 | CacheWriteBilledTokens: 200_000, // 1h write at 2x input |
| 442 | } |
| 443 | // 400K ordinary misses + 100K cache writes billed as 200K input units. |
| 444 | if got := p.Cost(u); got != 1.2 { |
| 445 | t.Errorf("Cost = %f, want 1.2", got) |
| 446 | } |
| 447 | } |
| 448 | |
| 449 | func TestPricingCostCacheWriteFieldsAreBackwardCompatible(t *testing.T) { |
| 450 | p := &Pricing{Input: 2.0} |
| 451 | |
| 452 | // Old usage records have neither cache-write field and retain the original |
| 453 | // one-input-rate calculation. |
| 454 | if got := p.Cost(&Usage{CacheMissTokens: 500_000}); got != 1.0 { |
| 455 | t.Errorf("legacy Cost = %f, want 1.0", got) |
| 456 | } |
| 457 | // A producer that reports raw write tokens without a billing tier also |
| 458 | // falls back to the ordinary input rate instead of making writes free. |
| 459 | if got := p.Cost(&Usage{CacheMissTokens: 500_000, CacheWriteTokens: 100_000}); got != 1.0 { |
| 460 | t.Errorf("unpriced write Cost = %f, want 1.0", got) |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | func TestPricingCostFallsBackToPromptTokensAsMiss(t *testing.T) { |
| 465 | p := &Pricing{Input: 2.0, Output: 10.0} |
| 466 | u := &Usage{PromptTokens: 500_000, CompletionTokens: 100_000} |
| 467 | if got := p.Cost(u); got != 2.0 { |
| 468 | t.Errorf("Cost = %f, want 2.0", got) |
| 469 | } |
| 470 | } |
| 471 | |
| 472 | func TestPricingCostZeroTokens(t *testing.T) { |
| 473 | p := &Pricing{Input: 2.0, Output: 10.0} |
| 474 | u := &Usage{} |
| 475 | if got := p.Cost(u); got != 0 { |
| 476 | t.Errorf("zero tokens Cost = %f, want 0", got) |
| 477 | } |
| 478 | } |
| 479 | |
| 480 | // --- Pricing.Symbol --- |
| 481 | |
| 482 | func TestPricingSymbolDefault(t *testing.T) { |
| 483 | p := &Pricing{} |
| 484 | if got := p.Symbol(); got != "¥" { |
| 485 | t.Errorf("empty Currency.Symbol() = %q, want ¥", got) |
| 486 | } |
| 487 | } |
| 488 | |
| 489 | func TestPricingSymbolNil(t *testing.T) { |
| 490 | var p *Pricing |
| 491 | if got := p.Symbol(); got != "¥" { |
| 492 | t.Errorf("nil.Symbol() = %q, want ¥", got) |
| 493 | } |
| 494 | } |
| 495 | |
| 496 | func TestPricingSymbolCustom(t *testing.T) { |
| 497 | p := &Pricing{Currency: "$"} |
| 498 | if got := p.Symbol(); got != "$" { |
| 499 | t.Errorf("Symbol() = %q, want $", got) |
| 500 | } |
| 501 | } |
| 502 | |
| 503 | func TestPricingSymbolNormalizesCurrencyCodes(t *testing.T) { |
| 504 | cases := []struct { |
| 505 | currency string |
| 506 | want string |
| 507 | }{ |
| 508 | {currency: "USD", want: "$"}, |
| 509 | {currency: "dollars", want: "$"}, |
| 510 | {currency: "CNY", want: "¥"}, |
| 511 | {currency: "￥", want: "¥"}, |
| 512 | {currency: "EUR", want: "€"}, |
| 513 | {currency: "₹", want: "₹"}, |
| 514 | {currency: "aud", want: "AUD "}, |
| 515 | {currency: "A$", want: "A$"}, |
| 516 | {currency: "HK$", want: "HK$"}, |
| 517 | {currency: "楼", want: "¥"}, |
| 518 | } |
| 519 | for _, tc := range cases { |
| 520 | p := &Pricing{Currency: tc.currency} |
| 521 | if got := p.Symbol(); got != tc.want { |
| 522 | t.Errorf("Currency %q Symbol() = %q, want %q", tc.currency, got, tc.want) |
| 523 | } |
| 524 | } |
| 525 | } |
| 526 | |
| 527 | // --- AuthError --- |
| 528 | |
| 529 | func TestAuthErrorWithKeyEnv(t *testing.T) { |
| 530 | e := &AuthError{Provider: "deepseek", KeyEnv: "DEEPSEEK_API_KEY", Status: 401} |
| 531 | msg := e.Error() |
| 532 | for _, want := range []string{"deepseek", "DEEPSEEK_API_KEY", "401", "invalid or expired"} { |
| 533 | if !contains(msg, want) { |
| 534 | t.Errorf("AuthError.Error() missing %q: %s", want, msg) |
| 535 | } |
| 536 | } |
| 537 | } |
| 538 | |
| 539 | func TestAuthErrorBodyStaysOutOfError(t *testing.T) { |
| 540 | // Body carries the server's reason for display layers to extract, but it |
| 541 | // must never leak into Error(): servers echo masked key fragments in auth |
| 542 | // bodies, and the ambient string flows into logs and traces. |
| 543 | e := &AuthError{Provider: "relay", Status: 401, Body: `{"error":{"message":"Your api key: ****ae54 has expired"}}`} |
| 544 | if e.Body == "" { |
| 545 | t.Fatal("Body should carry the server's reason") |
| 546 | } |
| 547 | if msg := e.Error(); contains(msg, "ae54") || contains(msg, "{") { |
| 548 | t.Errorf("AuthError.Error() must not include body content: %s", msg) |
| 549 | } |
| 550 | } |
| 551 | |
| 552 | func TestAuthErrorWithoutKeyEnv(t *testing.T) { |
| 553 | e := &AuthError{Provider: "openai", Status: 403} |
| 554 | msg := e.Error() |
| 555 | if !contains(msg, "the API key") { |
| 556 | t.Errorf("AuthError without KeyEnv should say 'the API key': %s", msg) |
| 557 | } |
| 558 | if !contains(msg, "403") { |
| 559 | t.Errorf("AuthError should include status code 403: %s", msg) |
| 560 | } |
| 561 | } |
| 562 | |
| 563 | func TestAuthErrorImplementsError(t *testing.T) { |
| 564 | var err error = &AuthError{Provider: "test", Status: 401} |
| 565 | if err.Error() == "" { |
| 566 | t.Error("AuthError.Error() should not be empty") |
| 567 | } |
| 568 | } |
| 569 | |
| 570 | // --- Registry --- |
| 571 | |
| 572 | func TestRegistryKindsSorted(t *testing.T) { |
| 573 | // The openai package self-registers via init(); we can't control that here |
| 574 | // but we can verify Kinds() returns a sorted list. |
| 575 | kinds := Kinds() |
| 576 | for i := 1; i < len(kinds); i++ { |
| 577 | if kinds[i-1] >= kinds[i] { |
| 578 | t.Errorf("Kinds() not sorted: %v", kinds) |
| 579 | break |
| 580 | } |
| 581 | } |
| 582 | } |
| 583 | |
| 584 | func TestNewUnknownKind(t *testing.T) { |
| 585 | _, err := New("nonexistent-kind-xyzzy", Config{}) |
| 586 | if err == nil { |
| 587 | t.Fatal("expected error for unknown kind") |
| 588 | } |
| 589 | if !contains(err.Error(), "unknown kind") { |
| 590 | t.Errorf("error should mention 'unknown kind': %v", err) |
| 591 | } |
| 592 | } |
| 593 | |
| 594 | func TestNewWithRegisteredKind(t *testing.T) { |
| 595 | // Register a mock factory. |
| 596 | Register("test-mock-__"+t.Name(), func(cfg Config) (Provider, error) { |
| 597 | return nil, nil |
| 598 | }) |
| 599 | // We can't easily unregister, but we can test it doesn't panic. |
| 600 | } |
| 601 | |
| 602 | func TestNewRejectsTypedNilProvider(t *testing.T) { |
| 603 | kind := "test-typed-nil-__" + t.Name() |
| 604 | Register(kind, func(cfg Config) (Provider, error) { |
| 605 | var p *mockProvider |
| 606 | return p, nil |
| 607 | }) |
| 608 | |
| 609 | _, err := New(kind, Config{}) |
| 610 | if err == nil { |
| 611 | t.Fatal("New should reject typed nil provider") |
| 612 | } |
| 613 | if !contains(err.Error(), "returned nil provider") { |
| 614 | t.Fatalf("New error = %v, want returned nil provider", err) |
| 615 | } |
| 616 | } |
| 617 | |
| 618 | // --- Role constants --- |
| 619 | |
| 620 | func TestRoleConstants(t *testing.T) { |
| 621 | if RoleSystem != "system" { |
| 622 | t.Errorf("RoleSystem = %q", RoleSystem) |
| 623 | } |
| 624 | if RoleUser != "user" { |
| 625 | t.Errorf("RoleUser = %q", RoleUser) |
| 626 | } |
| 627 | if RoleAssistant != "assistant" { |
| 628 | t.Errorf("RoleAssistant = %q", RoleAssistant) |
| 629 | } |
| 630 | if RoleTool != "tool" { |
| 631 | t.Errorf("RoleTool = %q", RoleTool) |
| 632 | } |
| 633 | } |
| 634 | |
| 635 | func TestMessageResponsesItemsRemainBackwardCompatible(t *testing.T) { |
| 636 | var legacy Message |
| 637 | if err := json.Unmarshal([]byte(`{"role":"assistant","content":"answer"}`), &legacy); err != nil { |
| 638 | t.Fatalf("unmarshal legacy message: %v", err) |
| 639 | } |
| 640 | if len(legacy.ResponsesItems) != 0 { |
| 641 | t.Fatalf("legacy ResponsesItems = %#v, want empty", legacy.ResponsesItems) |
| 642 | } |
| 643 | legacyJSON, err := json.Marshal(legacy) |
| 644 | if err != nil { |
| 645 | t.Fatalf("marshal legacy message: %v", err) |
| 646 | } |
| 647 | if strings.Contains(string(legacyJSON), "responses_items") { |
| 648 | t.Fatalf("legacy message gained responses_items: %s", legacyJSON) |
| 649 | } |
| 650 | |
| 651 | raw := json.RawMessage(`{"id":"ws_1","type":"web_search_call","status":"completed"}`) |
| 652 | current := Message{Role: RoleAssistant, Content: "answer", ResponsesItems: []json.RawMessage{raw}} |
| 653 | encoded, err := json.Marshal(current) |
| 654 | if err != nil { |
| 655 | t.Fatalf("marshal current message: %v", err) |
| 656 | } |
| 657 | var roundTrip Message |
| 658 | if err := json.Unmarshal(encoded, &roundTrip); err != nil { |
| 659 | t.Fatalf("unmarshal current message: %v", err) |
| 660 | } |
| 661 | if len(roundTrip.ResponsesItems) != 1 || string(roundTrip.ResponsesItems[0]) != string(raw) { |
| 662 | t.Fatalf("round-tripped ResponsesItems = %#v", roundTrip.ResponsesItems) |
| 663 | } |
| 664 | } |
| 665 | |
| 666 | // --- ChunkType constants --- |
| 667 | |
| 668 | func TestChunkTypeConstants(t *testing.T) { |
| 669 | types := []ChunkType{ChunkText, ChunkReasoning, ChunkToolCallStart, ChunkToolCallArgsDelta, ChunkToolCall, ChunkUsage, ChunkDone, ChunkError, ChunkResponsesItem} |
| 670 | for i, ct := range types { |
| 671 | if int(ct) != i { |
| 672 | t.Errorf("ChunkType %d: got %d", i, int(ct)) |
| 673 | } |
| 674 | } |
| 675 | } |
| 676 | |
| 677 | // --- ToolSchema --- |
| 678 | |
| 679 | func TestToolSchemaJSON(t *testing.T) { |
| 680 | ts := ToolSchema{ |
| 681 | Name: "bash", |
| 682 | Description: "Run a shell command", |
| 683 | Parameters: json.RawMessage(`{"type":"object"}`), |
| 684 | } |
| 685 | b, err := json.Marshal(ts) |
| 686 | if err != nil { |
| 687 | t.Fatalf("marshal: %v", err) |
| 688 | } |
| 689 | if !contains(string(b), "bash") { |
| 690 | t.Errorf("JSON missing name: %s", b) |
| 691 | } |
| 692 | } |
| 693 | |
| 694 | // helper |
| 695 | func contains(s, sub string) bool { |
| 696 | return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsStr(s, sub)) |
| 697 | } |
| 698 | |
| 699 | func containsStr(s, sub string) bool { |
| 700 | for i := 0; i <= len(s)-len(sub); i++ { |
| 701 | if s[i:i+len(sub)] == sub { |
| 702 | return true |
| 703 | } |
| 704 | } |
| 705 | return false |
| 706 | } |
| 707 | |
| 708 | // Ensure the Provider interface is satisfied by a minimal mock (compile-time check). |
| 709 | var _ Provider = (*mockProvider)(nil) |
| 710 | |
| 711 | type mockProvider struct{} |
| 712 | |
| 713 | func (m *mockProvider) Name() string { return "mock" } |
| 714 | func (m *mockProvider) Stream(ctx context.Context, req Request) (<-chan Chunk, error) { |
| 715 | ch := make(chan Chunk, 1) |
| 716 | ch <- Chunk{Type: ChunkDone} |
| 717 | close(ch) |
| 718 | return ch, nil |
| 719 | } |
| 720 | |
| 721 | func TestMockProviderImplementsInterface(t *testing.T) { |
| 722 | p := &mockProvider{} |
| 723 | if p.Name() != "mock" { |
| 724 | t.Errorf("Name = %q", p.Name()) |
| 725 | } |
| 726 | ch, err := p.Stream(context.Background(), Request{}) |
| 727 | if err != nil { |
| 728 | t.Fatalf("Stream: %v", err) |
| 729 | } |
| 730 | got := <-ch |
| 731 | if got.Type != ChunkDone { |
| 732 | t.Errorf("Chunk.Type = %d, want ChunkDone", got.Type) |
| 733 | } |
| 734 | } |
| 735 |