new-api/relay/channel/claude/relay_claude_test.go

package claude

import (
	"encoding/json"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/dto"
	"github.com/stretchr/testify/require"
)

func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) {
	claudeInfo := &ClaudeResponseInfo{
		Usage: &dto.Usage{},
	}
	claudeResponse := &dto.ClaudeResponse{
		Type: "message_start",
		Message: &dto.ClaudeMediaMessage{
			Id:    "msg_123",
			Model: "claude-3-5-sonnet",
			Usage: &dto.ClaudeUsage{
				InputTokens:              100,
				OutputTokens:             1,
				CacheCreationInputTokens: 50,
				CacheReadInputTokens:     30,
			},
		},
	}

	ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
	if !ok {
		t.Fatal("expected true")
	}
	if claudeInfo.Usage.PromptTokens != 100 {
		t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
	}
	if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
		t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
	}
	if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
		t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
	}
	if claudeInfo.ResponseId != "msg_123" {
		t.Errorf("ResponseId = %s, want msg_123", claudeInfo.ResponseId)
	}
	if claudeInfo.Model != "claude-3-5-sonnet" {
		t.Errorf("Model = %s, want claude-3-5-sonnet", claudeInfo.Model)
	}
}

func TestFormatClaudeResponseInfo_MessageDelta_FullUsage(t *testing.T) {
	// message_start 先积累 usage
	claudeInfo := &ClaudeResponseInfo{
		Usage: &dto.Usage{
			PromptTokens: 100,
			PromptTokensDetails: dto.InputTokenDetails{
				CachedTokens:         30,
				CachedCreationTokens: 50,
			},
			CompletionTokens: 1,
		},
	}

	// message_delta 带完整 usage（原生 Anthropic 场景）
	claudeResponse := &dto.ClaudeResponse{
		Type: "message_delta",
		Usage: &dto.ClaudeUsage{
			InputTokens:              100,
			OutputTokens:             200,
			CacheCreationInputTokens: 50,
			CacheReadInputTokens:     30,
		},
	}

	ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
	if !ok {
		t.Fatal("expected true")
	}
	if claudeInfo.Usage.PromptTokens != 100 {
		t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
	}
	if claudeInfo.Usage.CompletionTokens != 200 {
		t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
	}
	if claudeInfo.Usage.TotalTokens != 300 {
		t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
	}
	if !claudeInfo.Done {
		t.Error("expected Done = true")
	}
}

func TestFormatClaudeResponseInfo_MessageDelta_OnlyOutputTokens(t *testing.T) {
	// 模拟 Bedrock: message_start 已积累 usage
	claudeInfo := &ClaudeResponseInfo{
		Usage: &dto.Usage{
			PromptTokens: 100,
			PromptTokensDetails: dto.InputTokenDetails{
				CachedTokens:         30,
				CachedCreationTokens: 50,
			},
			CompletionTokens:            1,
			ClaudeCacheCreation5mTokens: 10,
			ClaudeCacheCreation1hTokens: 20,
		},
	}

	// Bedrock 的 message_delta 只有 output_tokens，缺少 input_tokens 和 cache 字段
	claudeResponse := &dto.ClaudeResponse{
		Type: "message_delta",
		Usage: &dto.ClaudeUsage{
			OutputTokens: 200,
			// InputTokens, CacheCreationInputTokens, CacheReadInputTokens 都是 0
		},
	}

	ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
	if !ok {
		t.Fatal("expected true")
	}
	// PromptTokens 应保持 message_start 的值（因为 message_delta 的 InputTokens=0，不更新）
	if claudeInfo.Usage.PromptTokens != 100 {
		t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
	}
	if claudeInfo.Usage.CompletionTokens != 200 {
		t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
	}
	if claudeInfo.Usage.TotalTokens != 300 {
		t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
	}
	// cache 字段应保持 message_start 的值
	if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
		t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
	}
	if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
		t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
	}
	if claudeInfo.Usage.ClaudeCacheCreation5mTokens != 10 {
		t.Errorf("ClaudeCacheCreation5mTokens = %d, want 10", claudeInfo.Usage.ClaudeCacheCreation5mTokens)
	}
	if claudeInfo.Usage.ClaudeCacheCreation1hTokens != 20 {
		t.Errorf("ClaudeCacheCreation1hTokens = %d, want 20", claudeInfo.Usage.ClaudeCacheCreation1hTokens)
	}
	if !claudeInfo.Done {
		t.Error("expected Done = true")
	}
}

func TestFormatClaudeResponseInfo_NilClaudeInfo(t *testing.T) {
	claudeResponse := &dto.ClaudeResponse{Type: "message_start"}
	ok := FormatClaudeResponseInfo(claudeResponse, nil, nil)
	if ok {
		t.Error("expected false for nil claudeInfo")
	}
}

func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) {
	text := "hello"
	claudeInfo := &ClaudeResponseInfo{
		Usage:        &dto.Usage{},
		ResponseText: strings.Builder{},
	}
	claudeResponse := &dto.ClaudeResponse{
		Type: "content_block_delta",
		Delta: &dto.ClaudeMediaMessage{
			Text: &text,
		},
	}

	ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
	if !ok {
		t.Fatal("expected true")
	}
	if claudeInfo.ResponseText.String() != "hello" {
		t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello")
	}
}

func TestRequestOpenAI2ClaudeMessage_AutoCacheControl(t *testing.T) {
	req := dto.GeneralOpenAIRequest{
		Model: "claude-sonnet-4-20250514",
		Messages: []dto.Message{
			{Role: "user", Content: "Hello"},
		},
		MaxTokens: 1024,
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)
	require.NotNil(t, claudeReq)
	require.Equal(t, `{"type":"ephemeral"}`, string(claudeReq.CacheControl))
}

func TestRequestOpenAI2ClaudeMessage_SystemStringWithAutoCache(t *testing.T) {
	req := dto.GeneralOpenAIRequest{
		Model: "claude-sonnet-4-20250514",
		Messages: []dto.Message{
			{Role: "system", Content: "You are a helpful assistant."},
			{Role: "user", Content: "Hello"},
		},
		MaxTokens: 1024,
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)
	require.Equal(t, `{"type":"ephemeral"}`, string(claudeReq.CacheControl))

	systemMedia, ok := claudeReq.System.([]dto.ClaudeMediaMessage)
	require.True(t, ok, "system should be []ClaudeMediaMessage")
	require.Len(t, systemMedia, 1)
	require.Equal(t, "You are a helpful assistant.", *systemMedia[0].Text)
}

func TestRequestOpenAI2ClaudeMessage_SystemComplexWithCacheControl(t *testing.T) {
	// 通过 JSON 反序列化构造请求，模拟真实场景
	jsonReq := `{
		"model": "claude-sonnet-4-20250514",
		"messages": [
			{"role": "system", "content": [{"type": "text", "text": "System instruction", "cache_control": {"type": "ephemeral"}}]},
			{"role": "user", "content": "Hello"}
		],
		"max_tokens": 1024
	}`
	var req dto.GeneralOpenAIRequest
	require.NoError(t, json.Unmarshal([]byte(jsonReq), &req))

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)

	// system 复合内容从 JSON 反序列化后 Content 是 []any
	// ParseContent 会解析 map[string]any，但不会提取 cache_control（已知限制）
	// 所以这里验证 system 被正确传递，cache_control 在当前实现中不会透传
	require.NotNil(t, claudeReq.System)
	// 顶层 cache_control 仍然注入
	require.Equal(t, `{"type":"ephemeral"}`, string(claudeReq.CacheControl))
}

func TestRequestOpenAI2ClaudeMessage_UserMessageCacheControlPassthrough(t *testing.T) {
	// 通过 SetMediaContent 直接设置（如 OpenRouter 等非标准路径）
	msg := dto.Message{Role: "user"}
	msg.SetMediaContent([]dto.MediaContent{
		{Type: "text", Text: "Hello", CacheControl: json.RawMessage(`{"type":"ephemeral"}`)},
	})

	req := dto.GeneralOpenAIRequest{
		Model:     "claude-sonnet-4-20250514",
		Messages:  []dto.Message{msg},
		MaxTokens: 1024,
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)
	require.Len(t, claudeReq.Messages, 1)

	// content block CacheControl 透传（仅在 Content 为 []MediaContent 时有效）
	content, ok := claudeReq.Messages[0].Content.([]dto.ClaudeMediaMessage)
	require.True(t, ok)
	require.Len(t, content, 1)
	require.Equal(t, `{"type":"ephemeral"}`, string(content[0].CacheControl))
}

func TestRequestOpenAI2ClaudeMessage_ToolUseWithAutoCache(t *testing.T) {
	req := dto.GeneralOpenAIRequest{
		Model: "claude-sonnet-4-20250514",
		Messages: []dto.Message{
			{Role: "user", Content: "What's the weather?"},
			{
				Role:      "assistant",
				Content:   "",
				ToolCalls: json.RawMessage(`[{"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"NYC\"}"}}]`),
			},
			{
				Role:      "tool",
				Content:   `{"temp":72}`,
				ToolCallId: "call_1",
			},
		},
		MaxTokens: 1024,
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)
	require.Equal(t, `{"type":"ephemeral"}`, string(claudeReq.CacheControl))
	// user + assistant(with tool_use) + user(tool_result) = 3 messages
	require.Len(t, claudeReq.Messages, 3)
}

func TestRequestOpenAI2ClaudeMessage_MinimalRequestNoPanic(t *testing.T) {
	req := dto.GeneralOpenAIRequest{
		Model:    "claude-sonnet-4-20250514",
		Messages: []dto.Message{},
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)
	require.NotNil(t, claudeReq)
	require.Equal(t, `{"type":"ephemeral"}`, string(claudeReq.CacheControl))
}

func TestRequestOpenAI2ClaudeMessage_FullSerialization(t *testing.T) {
	req := dto.GeneralOpenAIRequest{
		Model: "claude-sonnet-4-20250514",
		Messages: []dto.Message{
			{Role: "system", Content: "You are helpful."},
			{Role: "user", Content: "Hello"},
		},
		MaxTokens: 1024,
	}

	claudeReq, err := RequestOpenAI2ClaudeMessage(nil, req)
	require.NoError(t, err)

	jsonBytes, err := json.Marshal(claudeReq)
	require.NoError(t, err)
	jsonStr := string(jsonBytes)

	require.Contains(t, jsonStr, `"cache_control":{"type":"ephemeral"}`)
	require.Contains(t, jsonStr, `"model":"claude-sonnet-4-20250514"`)
	require.Contains(t, jsonStr, `"max_tokens":1024`)
}

func TestClaudeRequest_EmptyCacheControl_Omitted(t *testing.T) {
	req := dto.ClaudeRequest{
		Model:     "claude-sonnet-4-20250514",
		MaxTokens: 1024,
	}

	jsonBytes, err := json.Marshal(req)
	require.NoError(t, err)
	jsonStr := string(jsonBytes)

	require.NotContains(t, jsonStr, "cache_control")
}