CrossRelay/internal/anthropic/stream_test.go

package anthropic

import (
	"bytes"
	"encoding/json"
	"net/http"
	"strings"
	"testing"

	"gitea.rutua.cn/kate522/CrossRelay/internal/chatcmpl"
)

// captureWriter records everything the converter writes.
type captureWriter struct {
	header  http.Header
	body    bytes.Buffer
	status  int
	flushes int
}

func newCaptureWriter() *captureWriter {
	return &captureWriter{header: http.Header{}}
}

func (c *captureWriter) Header() http.Header         { return c.header }
func (c *captureWriter) Write(b []byte) (int, error) { return c.body.Write(b) }
func (c *captureWriter) WriteHeader(status int)      { c.status = status }
func (c *captureWriter) Flush()                      { c.flushes++ }

type sseEvent struct {
	name string
	data map[string]any
}

func parseSSE(t *testing.T, data string) []sseEvent {
	t.Helper()
	var events []sseEvent
	for _, block := range strings.Split(strings.TrimSpace(data), "\n\n") {
		if block == "" {
			continue
		}
		var name string
		var payload string
		for _, line := range strings.Split(block, "\n") {
			switch {
			case strings.HasPrefix(line, "event: "):
				name = strings.TrimPrefix(line, "event: ")
			case strings.HasPrefix(line, "data: "):
				payload = strings.TrimPrefix(line, "data: ")
			}
		}
		var parsed map[string]any
		if err := json.Unmarshal([]byte(payload), &parsed); err != nil {
			t.Fatalf("event %s payload is not JSON: %v (%s)", name, err, payload)
		}
		events = append(events, sseEvent{name: name, data: parsed})
	}
	return events
}

func eventNames(events []sseEvent) []string {
	names := make([]string, len(events))
	for i, e := range events {
		names[i] = e.name
	}
	return names
}

func textChunk(text string) *chatcmpl.StreamChunk {
	return &chatcmpl.StreamChunk{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{Content: &text}}}}
}

func finishChunk(reason string, usage *chatcmpl.Usage) *chatcmpl.StreamChunk {
	return &chatcmpl.StreamChunk{
		Choices: []chatcmpl.ChunkChoice{{FinishReason: &reason}},
		Usage:   usage,
	}
}

func runConverter(t *testing.T, chunks []*chatcmpl.StreamChunk) []sseEvent {
	t.Helper()
	w := newCaptureWriter()
	conv := NewStreamConverter(w, "claude-sonnet-4-5", 42)
	if err := conv.WriteStart(); err != nil {
		t.Fatalf("WriteStart: %v", err)
	}
	for _, c := range chunks {
		if err := conv.OnChunk(c); err != nil {
			t.Fatalf("OnChunk: %v", err)
		}
	}
	if err := conv.Finish(); err != nil {
		t.Fatalf("Finish: %v", err)
	}
	if w.status != http.StatusOK {
		t.Errorf("status = %d", w.status)
	}
	if ct := w.header.Get("Content-Type"); ct != "text/event-stream; charset=utf-8" {
		t.Errorf("content type = %q", ct)
	}
	return parseSSE(t, w.body.String())
}

func TestStreamTextOnly(t *testing.T) {
	events := runConverter(t, []*chatcmpl.StreamChunk{
		textChunk("Hello"),
		textChunk(", world"),
		finishChunk("stop", &chatcmpl.Usage{PromptTokens: 10, CompletionTokens: 3}),
	})
	wantNames := []string{
		"message_start",
		"content_block_start",
		"content_block_delta",
		"content_block_delta",
		"content_block_stop",
		"message_delta",
		"message_stop",
	}
	if got := eventNames(events); !equalStrings(got, wantNames) {
		t.Fatalf("events = %v, want %v", got, wantNames)
	}
	start := events[0].data["message"].(map[string]any)
	if start["model"] != "claude-sonnet-4-5" {
		t.Errorf("message model = %v", start["model"])
	}
	usage := start["usage"].(map[string]any)
	if usage["input_tokens"].(float64) != 42 {
		t.Errorf("input tokens = %v, want 42", usage["input_tokens"])
	}
	if events[2].data["delta"].(map[string]any)["text"] != "Hello" {
		t.Errorf("first delta = %v", events[2].data["delta"])
	}
	delta := events[5].data["delta"].(map[string]any)
	if delta["stop_reason"] != "end_turn" {
		t.Errorf("stop reason = %v", delta["stop_reason"])
	}
	if events[5].data["usage"].(map[string]any)["output_tokens"].(float64) != 3 {
		t.Errorf("output tokens = %v", events[5].data["usage"])
	}
}

func TestStreamToolCall(t *testing.T) {
	events := runConverter(t, []*chatcmpl.StreamChunk{
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{{
			Index:    0,
			ID:       "call_1",
			Type:     "function",
			Function: chatcmpl.ToolCallFunction{Name: "read_file"},
		}}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{{
			Index:    0,
			Function: chatcmpl.ToolCallFunction{Arguments: `{"path":`},
		}}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{{
			Index:    0,
			Function: chatcmpl.ToolCallFunction{Arguments: `"a.txt"}`},
		}}}}}},
		finishChunk("tool_calls", nil),
	})
	wantNames := []string{
		"message_start",
		"content_block_start",
		"content_block_delta",
		"content_block_delta",
		"content_block_stop",
		"message_delta",
		"message_stop",
	}
	if got := eventNames(events); !equalStrings(got, wantNames) {
		t.Fatalf("events = %v, want %v", got, wantNames)
	}
	block := events[1].data["content_block"].(map[string]any)
	if block["type"] != "tool_use" || block["id"] != "call_1" || block["name"] != "read_file" {
		t.Fatalf("tool block = %v", block)
	}
	if events[2].data["delta"].(map[string]any)["partial_json"] != `{"path":` {
		t.Errorf("first args delta = %v", events[2].data["delta"])
	}
	if events[3].data["delta"].(map[string]any)["partial_json"] != `"a.txt"}` {
		t.Errorf("second args delta = %v", events[3].data["delta"])
	}
	// Every event for the same block carries index 0.
	for _, e := range events[1:5] {
		if e.data["index"].(float64) != 0 {
			t.Errorf("%s index = %v, want 0", e.name, e.data["index"])
		}
	}
	if events[5].data["delta"].(map[string]any)["stop_reason"] != "tool_use" {
		t.Errorf("stop reason = %v", events[5].data["delta"])
	}
	// No usage chunk from upstream: output tokens are estimated, not zero.
	if events[5].data["usage"].(map[string]any)["output_tokens"].(float64) <= 0 {
		t.Error("estimated output tokens should be positive")
	}
}

func TestStreamParallelToolCalls(t *testing.T) {
	events := runConverter(t, []*chatcmpl.StreamChunk{
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{{
			Index: 0, ID: "call_a", Type: "function",
			Function: chatcmpl.ToolCallFunction{Name: "tool_a"},
		}}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{{
			Index: 1, ID: "call_b", Type: "function",
			Function: chatcmpl.ToolCallFunction{Name: "tool_b", Arguments: `{}`},
		}}}}}},
		finishChunk("tool_calls", nil),
	})
	// Two separate blocks: index 0 for tool_a, index 1 for tool_b.
	starts := 0
	names := map[float64]string{}
	for _, e := range events {
		if e.name != "content_block_start" {
			continue
		}
		starts++
		idx := e.data["index"].(float64)
		names[idx] = e.data["content_block"].(map[string]any)["name"].(string)
	}
	if starts != 2 {
		t.Fatalf("content_block_start count = %d, want 2", starts)
	}
	if names[0] != "tool_a" || names[1] != "tool_b" {
		t.Errorf("block names = %v", names)
	}
}

func TestStreamInterleavedToolDeltas(t *testing.T) {
	// Upstream sends argument deltas for two indices in an interleaved
	// order; each delta must land on its own block index.
	events := runConverter(t, []*chatcmpl.StreamChunk{
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{
			{Index: 0, ID: "a", Type: "function", Function: chatcmpl.ToolCallFunction{Name: "tool_a"}},
		}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{
			{Index: 1, ID: "b", Type: "function", Function: chatcmpl.ToolCallFunction{Name: "tool_b"}},
		}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{
			{Index: 0, Function: chatcmpl.ToolCallFunction{Arguments: `{"x":1}`}},
		}}}}},
		{Choices: []chatcmpl.ChunkChoice{{Delta: chatcmpl.Delta{ToolCalls: []chatcmpl.ToolCall{
			{Index: 1, Function: chatcmpl.ToolCallFunction{Arguments: `{"y":2}`}},
		}}}}},
		finishChunk("tool_calls", nil),
	})
	args := map[float64]string{}
	for _, e := range events {
		if e.name != "content_block_delta" {
			continue
		}
		idx := e.data["index"].(float64)
		args[idx] += e.data["delta"].(map[string]any)["partial_json"].(string)
	}
	if args[0] != `{"x":1}` || args[1] != `{"y":2}` {
		t.Errorf("interleaved args = %v", args)
	}
}

func TestStreamEmptyUpstream(t *testing.T) {
	events := runConverter(t, []*chatcmpl.StreamChunk{})
	want := []string{"message_start", "message_delta", "message_stop"}
	if got := eventNames(events); !equalStrings(got, want) {
		t.Fatalf("events = %v, want %v", got, want)
	}
}

func equalStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}