ops_service_redaction_test.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. package service
  2. import (
  3. "encoding/json"
  4. "testing"
  5. )
  6. func TestIsSensitiveKey_TokenBudgetKeysNotRedacted(t *testing.T) {
  7. t.Parallel()
  8. for _, key := range []string{
  9. "max_tokens",
  10. "max_output_tokens",
  11. "max_input_tokens",
  12. "max_completion_tokens",
  13. "max_tokens_to_sample",
  14. "budget_tokens",
  15. "prompt_tokens",
  16. "completion_tokens",
  17. "input_tokens",
  18. "output_tokens",
  19. "total_tokens",
  20. "token_count",
  21. } {
  22. if isSensitiveKey(key) {
  23. t.Fatalf("expected key %q to NOT be treated as sensitive", key)
  24. }
  25. }
  26. for _, key := range []string{
  27. "authorization",
  28. "Authorization",
  29. "access_token",
  30. "refresh_token",
  31. "id_token",
  32. "session_token",
  33. "token",
  34. "client_secret",
  35. "private_key",
  36. "signature",
  37. } {
  38. if !isSensitiveKey(key) {
  39. t.Fatalf("expected key %q to be treated as sensitive", key)
  40. }
  41. }
  42. }
  43. func TestSanitizeAndTrimJSONPayload_PreservesTokenBudgetFields(t *testing.T) {
  44. t.Parallel()
  45. raw := []byte(`{"model":"claude-3","max_tokens":123,"thinking":{"type":"enabled","budget_tokens":456},"access_token":"abc","messages":[{"role":"user","content":"hi"}]}`)
  46. out, _, _ := sanitizeAndTrimJSONPayload(raw, 10*1024)
  47. if out == "" {
  48. t.Fatalf("expected non-empty sanitized output")
  49. }
  50. var decoded map[string]any
  51. if err := json.Unmarshal([]byte(out), &decoded); err != nil {
  52. t.Fatalf("unmarshal sanitized output: %v", err)
  53. }
  54. if got, ok := decoded["max_tokens"].(float64); !ok || got != 123 {
  55. t.Fatalf("expected max_tokens=123, got %#v", decoded["max_tokens"])
  56. }
  57. thinking, ok := decoded["thinking"].(map[string]any)
  58. if !ok || thinking == nil {
  59. t.Fatalf("expected thinking object to be preserved, got %#v", decoded["thinking"])
  60. }
  61. if got, ok := thinking["budget_tokens"].(float64); !ok || got != 456 {
  62. t.Fatalf("expected thinking.budget_tokens=456, got %#v", thinking["budget_tokens"])
  63. }
  64. if got := decoded["access_token"]; got != "[REDACTED]" {
  65. t.Fatalf("expected access_token to be redacted, got %#v", got)
  66. }
  67. }
  68. func TestShrinkToEssentials_IncludesThinking(t *testing.T) {
  69. t.Parallel()
  70. root := map[string]any{
  71. "model": "claude-3",
  72. "max_tokens": 100,
  73. "thinking": map[string]any{
  74. "type": "enabled",
  75. "budget_tokens": 200,
  76. },
  77. "messages": []any{
  78. map[string]any{"role": "user", "content": "first"},
  79. map[string]any{"role": "user", "content": "last"},
  80. },
  81. }
  82. out := shrinkToEssentials(root)
  83. if _, ok := out["thinking"]; !ok {
  84. t.Fatalf("expected thinking to be included in essentials: %#v", out)
  85. }
  86. }