anthropic-proxy/internal/proxy/billing_test.go

package proxy

import (
	"encoding/hex"
	"strings"
	"testing"
)

func TestFingerprintSaltConstant(t *testing.T) {
	if fingerprintSalt != "59cf53e54c78" {
		t.Errorf("fingerprintSalt = %q, want %q", fingerprintSalt, "59cf53e54c78")
	}
}

func TestComputeFingerprint_Deterministic(t *testing.T) {
	a := computeFingerprint("hello world test message", "1.0.0")
	b := computeFingerprint("hello world test message", "1.0.0")
	if a != b {
		t.Errorf("fingerprint not deterministic: %q != %q", a, b)
	}
}

func TestComputeFingerprint_Length(t *testing.T) {
	fp := computeFingerprint("some message here", "2.0.0")
	if len(fp) != 3 {
		t.Errorf("fingerprint length = %d, want 3", len(fp))
	}
	// Must be valid hex
	if _, err := hex.DecodeString(fp + "0"); err != nil { // pad to even length for decode
		// Check each char is hex individually
		for _, c := range fp {
			if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {
				t.Errorf("fingerprint %q contains non-hex char %c", fp, c)
			}
		}
	}
}

func TestComputeFingerprint_DifferentVersions(t *testing.T) {
	a := computeFingerprint("same message", "1.0.0")
	b := computeFingerprint("same message", "2.0.0")
	if a == b {
		t.Errorf("different versions should (almost certainly) produce different fingerprints")
	}
}

func TestComputeFingerprint_ShortMessage(t *testing.T) {
	// "hi" has only 2 chars, indices [4,7,20] all out of range → chars = "000"
	fp := computeFingerprint("hi", "1.0.0")
	if len(fp) != 3 {
		t.Errorf("short message fingerprint length = %d, want 3", len(fp))
	}
}

func TestComputeFingerprint_EmptyMessage(t *testing.T) {
	// Empty → all indices out of range → chars = "000"
	fp := computeFingerprint("", "1.0.0")
	if len(fp) != 3 {
		t.Errorf("empty message fingerprint length = %d, want 3", len(fp))
	}
	// Empty and short message with same version should produce same fingerprint
	// since both result in chars = "000"
	fpShort := computeFingerprint("hi", "1.0.0")
	if fp != fpShort {
		t.Errorf("empty and 'hi' should produce same fingerprint (both use '000'), got %q vs %q", fp, fpShort)
	}
}

func TestComputeFingerprint_Unicode(t *testing.T) {
	// Emoji: 🎉 is U+1F389, encoded as UTF-16 surrogate pair [0xD83C, 0xDF89]
	// So "abcd🎉fg" in UTF-16 is [a, b, c, d, 0xD83C, 0xDF89, f, g] = 8 uint16 values
	// indices [4,7,20]: runes[4]=0xD83C, runes[7]='g', runes[20]=out of range
	fp := computeFingerprint("abcd🎉fg", "1.0.0")
	if len(fp) != 3 {
		t.Errorf("unicode fingerprint length = %d, want 3", len(fp))
	}
}

func TestComputeFingerprint_CharExtraction(t *testing.T) {
	// "Hello, World!" UTF-16: [H,e,l,l,o,',', ,W,o,r,l,d,!]
	// indices [4,7,20]: runes[4]='o', runes[7]='W', runes[20]=out of range → "0"
	// So chars should be "oW0"
	// Verify by comparing to a message where we know the expected extracted chars
	// Two messages that extract same chars at indices should produce same fingerprint
	// "xxxxoxxWxxxxxxxxxxxx" → index 4='o', 7='W', 20=out of range → "oW0" (20 chars, index 20 out of range)
	fp1 := computeFingerprint("Hello, World!", "1.0.0")
	fp2 := computeFingerprint("xxxxoxxWxxxxxxxxxxxx", "1.0.0")
	if fp1 != fp2 {
		t.Errorf("messages with same chars at indices [4,7,20] should produce same fingerprint, got %q vs %q", fp1, fp2)
	}
}

func TestComputeFingerprint_IndexBoundary(t *testing.T) {
	// Message with exactly 21 chars → index 20 is valid
	msg21 := "abcdefghijklmnopqrstu" // 21 chars
	fp21 := computeFingerprint(msg21, "1.0.0")

	// Message with exactly 20 chars → index 20 is out of range → "0"
	msg20 := "abcdefghijklmnopqrst" // 20 chars
	fp20 := computeFingerprint(msg20, "1.0.0")

	// They should differ because index 20 produces different chars
	if fp21 == fp20 {
		t.Errorf("boundary test: 21-char and 20-char messages should differ at index 20")
	}
}

func TestExtractFirstUserMessage(t *testing.T) {
	tests := []struct {
		name     string
		body     string
		expected string
	}{
		{
			name:     "simple string content",
			body:     `{"messages":[{"role":"user","content":"hello world"}]}`,
			expected: "hello world",
		},
		{
			name:     "array content with text block",
			body:     `{"messages":[{"role":"user","content":[{"type":"text","text":"from array"}]}]}`,
			expected: "from array",
		},
		{
			name:     "no user messages",
			body:     `{"messages":[{"role":"assistant","content":"I am assistant"}]}`,
			expected: "",
		},
		{
			name:     "assistant only messages",
			body:     `{"messages":[{"role":"assistant","content":"a1"},{"role":"assistant","content":"a2"}]}`,
			expected: "",
		},
		{
			name:     "user with non-text block first then text",
			body:     `{"messages":[{"role":"user","content":[{"type":"image","source":"x"},{"type":"text","text":"the text"}]}]}`,
			expected: "the text",
		},
		{
			name:     "user with only non-text blocks",
			body:     `{"messages":[{"role":"user","content":[{"type":"image","source":"x"}]}]}`,
			expected: "",
		},
		{
			name:     "no messages field",
			body:     `{"model":"claude-sonnet-4-6"}`,
			expected: "",
		},
		{
			name:     "messages not array",
			body:     `{"messages":"not array"}`,
			expected: "",
		},
		{
			name:     "empty messages array",
			body:     `{"messages":[]}`,
			expected: "",
		},
		{
			name:     "first user message used even if multiple exist",
			body:     `{"messages":[{"role":"user","content":"first"},{"role":"user","content":"second"}]}`,
			expected: "first",
		},
		{
			name:     "assistant before user",
			body:     `{"messages":[{"role":"assistant","content":"assistant msg"},{"role":"user","content":"user msg"}]}`,
			expected: "user msg",
		},
		{
			name:     "user with array content - first text block used",
			body:     `{"messages":[{"role":"user","content":[{"type":"text","text":"first text"},{"type":"text","text":"second text"}]}]}`,
			expected: "first text",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := extractFirstUserMessage([]byte(tt.body))
			if result != tt.expected {
				t.Errorf("got %q, want %q", result, tt.expected)
			}
		})
	}
}

func TestExtractFirstUserMessage_BreaksAfterFirstUser(t *testing.T) {
	// The function should break after finding the first user message,
	// even if it didn't extract text (e.g. user with only image blocks)
	body := `{"messages":[{"role":"user","content":[{"type":"image","source":"x"}]},{"role":"user","content":"second user"}]}`
	result := extractFirstUserMessage([]byte(body))
	// First user has no text blocks, function breaks, returns ""
	if result != "" {
		t.Errorf("should return empty when first user has no text, got %q", result)
	}
}

func TestBuildBillingHeader(t *testing.T) {
	body := []byte(`{"messages":[{"role":"user","content":"test message"}]}`)
	version := "1.2.3"
	header := buildBillingHeader(body, version)

	// Check format
	if !strings.HasPrefix(header, "x-anthropic-billing-header: cc_version=1.2.3.") {
		t.Errorf("header should start with 'x-anthropic-billing-header: cc_version=1.2.3.', got %q", header)
	}
	if !strings.Contains(header, "; cc_entrypoint=cli; cch=00000;") {
		t.Errorf("header should contain '; cc_entrypoint=cli; cch=00000;', got %q", header)
	}

	// Verify the fingerprint part is 3 chars
	// Format: "x-anthropic-billing-header: cc_version=1.2.3.XXX; cc_entrypoint=cli; cch=00000;"
	parts := strings.Split(header, "cc_version=")
	if len(parts) != 2 {
		t.Fatalf("unexpected header format: %q", header)
	}
	versionFP := strings.Split(parts[1], ";")[0]
	if !strings.HasPrefix(versionFP, "1.2.3.") {
		t.Errorf("version+fingerprint should start with '1.2.3.', got %q", versionFP)
	}
	fp := strings.TrimPrefix(versionFP, "1.2.3.")
	if len(fp) != 3 {
		t.Errorf("fingerprint should be 3 chars, got %q (len %d)", fp, len(fp))
	}
}

func TestBuildBillingHeader_EmptyMessages(t *testing.T) {
	body := []byte(`{"messages":[]}`)
	version := "1.0.0"
	header := buildBillingHeader(body, version)
	if !strings.HasPrefix(header, "x-anthropic-billing-header: cc_version=") {
		t.Errorf("header format wrong: %q", header)
	}
}

func TestInjectBillingHeader_NoExistingSystem(t *testing.T) {
	body := []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`)
	result := injectBillingHeader(body, "1.0.0")
	resultStr := string(result)

	// Should have system field now
	if !strings.Contains(resultStr, `"system"`) {
		t.Errorf("should inject system field, got %s", resultStr)
	}
	// System should be an array with one billing block
	if !strings.Contains(resultStr, "x-anthropic-billing-header") {
		t.Errorf("should contain billing header text, got %s", resultStr)
	}
	if !strings.Contains(resultStr, `"type":"text"`) {
		t.Errorf("billing block should have type text, got %s", resultStr)
	}
}

func TestInjectBillingHeader_ExistingSystemArray(t *testing.T) {
	body := []byte(`{"system":[{"type":"text","text":"existing prompt"}],"messages":[{"role":"user","content":"hi"}]}`)
	result := injectBillingHeader(body, "1.0.0")
	resultStr := string(result)

	// Should contain both billing header and existing prompt
	if !strings.Contains(resultStr, "x-anthropic-billing-header") {
		t.Errorf("should contain billing header, got %s", resultStr)
	}
	if !strings.Contains(resultStr, "existing prompt") {
		t.Errorf("should preserve existing prompt, got %s", resultStr)
	}

	// Billing block should be FIRST (prepended)
	billingIdx := strings.Index(resultStr, "x-anthropic-billing-header")
	existingIdx := strings.Index(resultStr, "existing prompt")
	if billingIdx > existingIdx {
		t.Errorf("billing block should come before existing prompt")
	}
}

func TestInjectBillingHeader_ExistingSystemString(t *testing.T) {
	body := []byte(`{"system":"You are a helpful assistant","messages":[{"role":"user","content":"hi"}]}`)
	result := injectBillingHeader(body, "1.0.0")
	resultStr := string(result)

	// Should convert to array with billing block first, then original text
	if !strings.Contains(resultStr, "x-anthropic-billing-header") {
		t.Errorf("should contain billing header, got %s", resultStr)
	}
	if !strings.Contains(resultStr, "You are a helpful assistant") {
		t.Errorf("should preserve original system string, got %s", resultStr)
	}

	// Billing should come first
	billingIdx := strings.Index(resultStr, "x-anthropic-billing-header")
	origIdx := strings.Index(resultStr, "You are a helpful assistant")
	if billingIdx > origIdx {
		t.Errorf("billing block should come before original system text")
	}
}

func TestInjectBillingHeader_PreservesOtherFields(t *testing.T) {
	body := []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}],"max_tokens":1024}`)
	result := injectBillingHeader(body, "1.0.0")
	resultStr := string(result)

	if !strings.Contains(resultStr, `"model":"claude-sonnet-4-6"`) {
		t.Errorf("should preserve model field, got %s", resultStr)
	}
	if !strings.Contains(resultStr, `"max_tokens":1024`) {
		t.Errorf("should preserve max_tokens field, got %s", resultStr)
	}
}

func TestInjectBillingHeader_BillingBlockFormat(t *testing.T) {
	body := []byte(`{"messages":[{"role":"user","content":"test"}]}`)
	result := injectBillingHeader(body, "2.5.0")
	resultStr := string(result)

	// Verify the billing block contains the correct version
	if !strings.Contains(resultStr, "cc_version=2.5.0.") {
		t.Errorf("billing block should contain cc_version=2.5.0., got %s", resultStr)
	}
	if !strings.Contains(resultStr, "cc_entrypoint=cli") {
		t.Errorf("billing block should contain cc_entrypoint=cli, got %s", resultStr)
	}
	if !strings.Contains(resultStr, "cch=00000") {
		t.Errorf("billing block should contain cch=00000, got %s", resultStr)
	}
}