majordomo/provider/anthropic/stream_test.go

package anthropic

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"

	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// sse joins data payloads into an SSE body. Each payload becomes one event
// ("event:" name derived from the JSON type field is what the real API
// sends, but the client dispatches on the data, so a generic name is fine).
func sse(payloads ...string) string {
	var b strings.Builder
	for _, p := range payloads {
		b.WriteString("event: event\n")
		b.WriteString("data: ")
		b.WriteString(p)
		b.WriteString("\n\n")
	}
	return b.String()
}

func sseServer(t *testing.T, c *capture, body string) *Provider {
	t.Helper()
	return newTestProvider(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		raw, _ := io.ReadAll(r.Body)
		c.mu.Lock()
		c.hits++
		c.header = r.Header.Clone()
		c.body = raw
		c.mu.Unlock()
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = io.WriteString(w, body)
	}))
}

// drain collects all events until io.EOF, failing the test on any error.
func drain(t *testing.T, s llm.Stream) []llm.StreamEvent {
	t.Helper()
	var events []llm.StreamEvent
	for {
		ev, err := s.Next()
		if err == io.EOF {
			return events
		}
		if err != nil {
			t.Fatalf("Next: %v", err)
		}
		events = append(events, ev)
	}
}

func openStream(t *testing.T, p *Provider, modelID string) llm.Stream {
	t.Helper()
	s, err := mustModel(t, p, modelID).Stream(context.Background(),
		llm.Request{Messages: []llm.Message{llm.UserText("hi")}})
	if err != nil {
		t.Fatalf("Stream: %v", err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s
}

func TestStreamTextDeltas(t *testing.T) {
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"m","usage":{"input_tokens":10,"cache_creation_input_tokens":2,"cache_read_input_tokens":3,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"ping"}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel"}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"lo"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":" world"}}`,
		`{"type":"content_block_stop","index":1}`,
		`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":12}}`,
		`{"type":"message_stop"}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	s := openStream(t, p, "claude-test")
	events := drain(t, s)

	if len(events) != 4 {
		t.Fatalf("events = %d, want 4 (3 deltas + final response)", len(events))
	}
	for i, want := range []string{"Hel", "lo", " world"} {
		if events[i].TextDelta != want {
			t.Errorf("event[%d].TextDelta = %q, want %q", i, events[i].TextDelta, want)
		}
	}

	final := events[3].Response
	if final == nil {
		t.Fatal("last event has no Response")
	}
	if len(final.Parts) != 2 {
		t.Fatalf("final parts = %d, want 2 (one per text block)", len(final.Parts))
	}
	if final.Text() != "Hello world" {
		t.Errorf("final text = %q, want %q", final.Text(), "Hello world")
	}
	if final.FinishReason != llm.FinishStop {
		t.Errorf("finish = %q, want stop", final.FinishReason)
	}
	// Input = 10+2+3 from message_start; output = 12 from message_delta.
	if final.Usage.InputTokens != 15 || final.Usage.OutputTokens != 12 {
		t.Errorf("usage = %+v, want {15 12}", final.Usage)
	}
	if final.Model != "anthropic/claude-test" {
		t.Errorf("model = %q, want anthropic/claude-test", final.Model)
	}

	// Past EOF, Next keeps returning io.EOF.
	if _, err := s.Next(); err != io.EOF {
		t.Errorf("Next after EOF = %v, want io.EOF", err)
	}

	// The request must carry "stream": true.
	if streamFlag := c.bodyMap(t)["stream"]; streamFlag != true {
		t.Errorf("request stream = %v, want true", streamFlag)
	}
}

func TestStreamToolCallAssembly(t *testing.T) {
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":8,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Checking."}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_9","name":"get_weather","input":{}}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":""}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"location\":"}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":" \"San Francisco, CA\"}"}}`,
		`{"type":"content_block_stop","index":1}`,
		`{"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"toolu_10","name":"noop","input":{}}}`,
		`{"type":"content_block_stop","index":2}`,
		`{"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":21}}`,
		`{"type":"message_stop"}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	events := drain(t, openStream(t, p, "claude-test"))

	if len(events) != 4 {
		t.Fatalf("events = %d, want 4 (text, 2 tool calls, final)", len(events))
	}
	if events[0].TextDelta != "Checking." {
		t.Errorf("event[0] = %+v, want text delta", events[0])
	}

	call := events[1].ToolCall
	if call == nil {
		t.Fatal("event[1] has no ToolCall")
	}
	if call.ID != "toolu_9" || call.Name != "get_weather" {
		t.Errorf("tool call = %+v", call)
	}
	var args map[string]any
	if err := json.Unmarshal(call.Arguments, &args); err != nil {
		t.Fatalf("assembled arguments invalid JSON: %v (%s)", err, call.Arguments)
	}
	if args["location"] != "San Francisco, CA" {
		t.Errorf("arguments = %v", args)
	}

	empty := events[2].ToolCall
	if empty == nil || empty.ID != "toolu_10" {
		t.Fatalf("event[2] = %+v, want second tool call", events[2])
	}
	if string(empty.Arguments) != "{}" {
		t.Errorf("empty tool call arguments = %s, want {}", empty.Arguments)
	}

	final := events[3].Response
	if final == nil {
		t.Fatal("last event has no Response")
	}
	if len(final.ToolCalls) != 2 {
		t.Errorf("final tool calls = %d, want 2", len(final.ToolCalls))
	}
	if final.FinishReason != llm.FinishToolCalls {
		t.Errorf("finish = %q, want tool_calls", final.FinishReason)
	}
	if final.Text() != "Checking." {
		t.Errorf("final text = %q", final.Text())
	}
	if final.Usage.InputTokens != 8 || final.Usage.OutputTokens != 21 {
		t.Errorf("usage = %+v, want {8 21}", final.Usage)
	}
}

func TestStreamThinkingSkipped(t *testing.T) {
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":5,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"hmm"}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"hi"}}`,
		`{"type":"content_block_stop","index":1}`,
		`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}`,
		`{"type":"message_stop"}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	events := drain(t, openStream(t, p, "claude-test"))

	if len(events) != 2 {
		t.Fatalf("events = %d, want 2 (thinking produces none)", len(events))
	}
	if events[0].TextDelta != "hi" {
		t.Errorf("event[0] = %+v, want TextDelta hi", events[0])
	}
	final := events[1].Response
	if final == nil || len(final.Parts) != 1 || final.Text() != "hi" {
		t.Errorf("final = %+v, want single text part %q", final, "hi")
	}
}

func TestStreamMidStreamError(t *testing.T) {
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":5,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"par"}}`,
		`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	s := openStream(t, p, "claude-test")

	ev, err := s.Next()
	if err != nil || ev.TextDelta != "par" {
		t.Fatalf("first Next = (%+v, %v), want text delta", ev, err)
	}
	_, err = s.Next()
	if err == nil {
		t.Fatal("second Next succeeded, want mid-stream error")
	}
	apiErr, ok := errors.AsType[*llm.APIError](err)
	if !ok {
		t.Fatalf("error %T (%v), want *llm.APIError", err, err)
	}
	if apiErr.Code != "overloaded_error" || apiErr.Message != "Overloaded" || apiErr.Status != 0 {
		t.Errorf("apiErr = %+v", apiErr)
	}
	if llm.Classify(err) != llm.ClassTransient {
		t.Error("overloaded_error must classify transient")
	}
}

func TestStreamHTTPErrorBeforeEvents(t *testing.T) {
	var c capture
	p := newTestProvider(t, c.handler(529,
		`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`))
	_, err := mustModel(t, p, "claude-test").Stream(context.Background(),
		llm.Request{Messages: []llm.Message{llm.UserText("hi")}})
	if err == nil {
		t.Fatal("Stream succeeded, want APIError before any events")
	}
	apiErr, ok := errors.AsType[*llm.APIError](err)
	if !ok {
		t.Fatalf("error %T (%v), want *llm.APIError", err, err)
	}
	if apiErr.Status != 529 || apiErr.Code != "overloaded_error" {
		t.Errorf("apiErr = %+v, want 529 overloaded_error", apiErr)
	}
	if llm.Classify(err) != llm.ClassTransient {
		t.Error("529 must classify transient")
	}
}

func TestStreamTruncatedBody(t *testing.T) {
	// Stream ends without message_stop: Next must surface unexpected EOF.
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":5,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	s := openStream(t, p, "claude-test")

	if ev, err := s.Next(); err != nil || ev.TextDelta != "hi" {
		t.Fatalf("first Next = (%+v, %v)", ev, err)
	}
	if _, err := s.Next(); !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Errorf("Next on truncated stream = %v, want io.ErrUnexpectedEOF", err)
	}
}

func TestStreamCloseIsSafe(t *testing.T) {
	body := sse(
		`{"type":"message_start","message":{"id":"msg_1","usage":{"input_tokens":5,"output_tokens":1}}}`,
		`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`,
		`{"type":"content_block_stop","index":0}`,
		`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}`,
		`{"type":"message_stop"}`,
	)
	var c capture
	p := sseServer(t, &c, body)
	s := openStream(t, p, "claude-test")

	if err := s.Close(); err != nil {
		t.Errorf("first Close: %v", err)
	}
	if err := s.Close(); err != nil {
		t.Errorf("second Close: %v", err)
	}

	// After EOF, Close is still fine.
	s2 := openStream(t, p, "claude-test")
	drain(t, s2)
	if err := s2.Close(); err != nil {
		t.Errorf("Close after EOF: %v", err)
	}
}