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