Cyrene_v2/internal/runtime/model/anthropic_test.go

package model

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// anthropicCaptured records what an AnthropicChatAdapter sent to a stub server.
type anthropicCaptured struct {
	method           string
	path             string
	apiKey           string
	anthropicVersion string
	contentType      string
	accept           string
	body             map[string]any
}

func newAnthropicServer(t *testing.T, status int, contentType, body string, ch chan<- anthropicCaptured) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var reqBody map[string]any
		if ch != nil {
			_ = json.NewDecoder(r.Body).Decode(&reqBody)
		}
		if contentType != "" {
			w.Header().Set("Content-Type", contentType)
		}
		w.WriteHeader(status)
		if body != "" {
			_, _ = io.WriteString(w, body)
		}
		if ch != nil {
			ch <- anthropicCaptured{
				method:           r.Method,
				path:             r.URL.Path,
				apiKey:           r.Header.Get("x-api-key"),
				anthropicVersion: r.Header.Get("anthropic-version"),
				contentType:      r.Header.Get("Content-Type"),
				accept:           r.Header.Get("Accept"),
				body:             reqBody,
			}
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

const anthropicTextResponse = `{
  "id":"msg_1","model":"claude-test","role":"assistant",
  "content":[{"type":"text","text":"Hello!"}],
  "usage":{"input_tokens":5,"output_tokens":2}
}`

func TestAnthropicChat_Chat_NonStreamingText(t *testing.T) {
	t.Parallel()
	ch := make(chan anthropicCaptured, 1)
	srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)

	recs := make([]CallRecord, 0, 1)
	temp := 1.0
	ad := NewAnthropicChat(config.LLMConfig{
		AnthropicBaseURL: srv.URL, AnthropicAPIKey: "ak-test", AnthropicModel: "claude-test",
	}, WithAnthropicRecorder(func(cr CallRecord) { recs = append(recs, cr) }))

	resp, err := ad.Chat(context.Background(), Request{
		Messages:    []Message{TextMessage(RoleUser, "hi")},
		Model:       "claude-test",
		Temperature: &temp,
		MaxTokens:   64,
		Purpose:     "chat",
	}, Options{})
	if err != nil {
		t.Fatalf("Chat error: %v", err)
	}

	got := <-ch
	if got.method != http.MethodPost {
		t.Errorf("method = %s, want POST", got.method)
	}
	if got.path != "/v1/messages" {
		t.Errorf("path = %s, want /v1/messages", got.path)
	}
	if got.apiKey != "ak-test" {
		t.Errorf("x-api-key = %q, want ak-test", got.apiKey)
	}
	if got.anthropicVersion != "2023-06-01" {
		t.Errorf("anthropic-version = %q, want 2023-06-01", got.anthropicVersion)
	}
	if got.contentType != "application/json" {
		t.Errorf("content-type = %s, want application/json", got.contentType)
	}
	if got.accept != "application/json" {
		t.Errorf("accept = %s, want application/json", got.accept)
	}
	if got.body["model"] != "claude-test" {
		t.Errorf("body.model = %v, want claude-test", got.body["model"])
	}
	if got.body["max_tokens"] != float64(64) {
		t.Errorf("body.max_tokens = %v, want 64", got.body["max_tokens"])
	}
	if got.body["temperature"] != 1.0 {
		t.Errorf("body.temperature = %v, want 1.0", got.body["temperature"])
	}
	if got.body["stream"] != nil && got.body["stream"] != false {
		t.Errorf("body.stream = %v, want omitted(false)", got.body["stream"])
	}
	msgs, _ := got.body["messages"].([]any)
	if len(msgs) != 1 {
		t.Fatalf("body.messages = %v, want 1 message", got.body["messages"])
	}

	if resp.Model != "claude-test" {
		t.Errorf("resp.Model = %q, want claude-test", resp.Model)
	}
	if resp.Content != "Hello!" {
		t.Errorf("resp.Content = %q, want Hello!", resp.Content)
	}
	if resp.TotalTokens != 7 || resp.PromptTokens != 5 || resp.CompletionTokens != 2 {
		t.Errorf("resp usage = (%d,%d,%d), want (5,2,7)", resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens)
	}
	if resp.Message.Role != RoleAssistant || len(resp.Message.Content) != 1 || resp.Message.Content[0].Text != "Hello!" {
		t.Errorf("resp.Message = %+v", resp.Message)
	}

	if len(recs) != 1 {
		t.Fatalf("records = %d, want 1", len(recs))
	}
	if recs[0].Model != "claude-test" || recs[0].Purpose != "chat" || recs[0].Tokens != 7 || recs[0].TookMs < 0 {
		t.Errorf("record = %+v, want model=claude-test purpose=chat tokens=7 took_ms>=0", recs[0])
	}
}

func TestAnthropicChat_Chat_ToolUseInputObject(t *testing.T) {
	t.Parallel()
	const resp = `{
	  "id":"msg_2","model":"claude-test",
	  "content":[
	    {"type":"text","text":"Let me check."},
	    {"type":"tool_use","id":"toolu_1","name":"lookup","input":{"location":"Paris","limit":3}}
	  ],
	  "usage":{"input_tokens":10,"output_tokens":12}
	}`
	ch := make(chan anthropicCaptured, 1)
	srv := newAnthropicServer(t, http.StatusOK, "application/json", resp, ch)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	req := Request{
		Messages: []Message{TextMessage(RoleUser, "hi")},
		Tools:    []ToolSpec{{Name: "lookup", Description: "look up", Parameters: map[string]any{"type": "object"}}},
	}
	out, err := ad.Chat(context.Background(), req, Options{})
	if err != nil {
		t.Fatalf("Chat error: %v", err)
	}

	got := <-ch
	tools, ok := got.body["tools"].([]any)
	if !ok || len(tools) != 1 {
		t.Fatalf("body.tools = %v, want 1 entry", got.body["tools"])
	}
	tool0, _ := tools[0].(map[string]any)
	if tool0["name"] != "lookup" {
		t.Errorf("tool name = %v, want lookup", tool0["name"])
	}
	if _, ok = tool0["input_schema"]; !ok {
		t.Errorf("tool input_schema missing from %v", tool0)
	}

	if out.Content != "Let me check." {
		t.Errorf("out.Content = %q, want Let me check.", out.Content)
	}
	if len(out.ToolCalls) != 1 {
		t.Fatalf("out.ToolCalls = %+v, want 1", out.ToolCalls)
	}
	tc := out.ToolCalls[0]
	if tc.Name != "lookup" {
		t.Errorf("tool call name = %q, want lookup", tc.Name)
	}
	if tc.Args == nil {
		t.Fatalf("tool call args = nil, want structured object")
	}
	if tc.Args["location"] != "Paris" {
		t.Errorf("tool call args.location = %v, want Paris (structured, not JSON string)", tc.Args["location"])
	}
	if tc.Args["limit"] != float64(3) {
		t.Errorf("tool call args.limit = %v, want 3", tc.Args["limit"])
	}
}

func TestAnthropicChat_Chat_SystemAndToolChoice(t *testing.T) {
	t.Parallel()
	ch := make(chan anthropicCaptured, 1)
	srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	req := Request{
		Messages: []Message{
			TextMessage(RoleSystem, "Be concise."),
			TextMessage(RoleUser, "hi"),
		},
		ToolChoice: map[string]any{"type": "auto"},
	}
	if _, err := ad.Chat(context.Background(), req, Options{}); err != nil {
		t.Fatalf("Chat error: %v", err)
	}

	got := <-ch
	if got.body["system"] != "Be concise." {
		t.Errorf("body.system = %v, want Be concise.", got.body["system"])
	}
	msgs, _ := got.body["messages"].([]any)
	if len(msgs) != 1 {
		t.Fatalf("body.messages = %v, want 1 message (system extracted)", got.body["messages"])
	}
	if _, ok := got.body["tool_choice"].(map[string]any); !ok {
		t.Errorf("body.tool_choice = %v, want an object", got.body["tool_choice"])
	}
}

func TestAnthropicChat_Chat_ImageContent(t *testing.T) {
	t.Parallel()
	ch := make(chan anthropicCaptured, 1)
	srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	req := Request{Messages: []Message{ImageMessage(RoleUser, "what is this", "data:image/png;base64,abc")}}
	if _, err := ad.Chat(context.Background(), req, Options{}); err != nil {
		t.Fatalf("Chat error: %v", err)
	}

	got := <-ch
	msgs, _ := got.body["messages"].([]any)
	part0, _ := msgs[0].(map[string]any)
	content, _ := part0["content"].([]any)
	if len(content) != 2 {
		t.Fatalf("content parts = %d, want 2 (text+image)", len(content))
	}
	found := false
	for _, c := range content {
		cp, _ := c.(map[string]any)
		if cp["type"] != "image" {
			continue
		}
		found = true
		src, _ := cp["source"].(map[string]any)
		if src["type"] != "base64" {
			t.Errorf("image source type = %v, want base64", src["type"])
		}
		if src["media_type"] != "image/png" {
			t.Errorf("image media_type = %v, want image/png", src["media_type"])
		}
		if src["data"] != "abc" {
			t.Errorf("image data = %v, want abc", src["data"])
		}
	}
	if !found {
		t.Fatalf("no image block found in %v", content)
	}
}

func TestAnthropicChat_Chat_UpstreamError(t *testing.T) {
	t.Parallel()
	srv := newAnthropicServer(t, http.StatusInternalServerError, "application/json", `{"error":"boom"}`, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "m"})
	_, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err == nil {
		t.Fatal("want error, got nil")
	}
	if !errors.Is(err, ErrUpstream) {
		t.Fatalf("err = %v, want ErrUpstream", err)
	}
}

func TestAnthropicChat_Chat_NoBaseURL(t *testing.T) {
	t.Parallel()
	ad := NewAnthropicChat(config.LLMConfig{AnthropicModel: "m"})
	if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) {
		t.Fatalf("err = %v, want ErrUpstream", err)
	}
}

func TestAnthropicChat_Chat_NoModel(t *testing.T) {
	t.Parallel()
	srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL})
	if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrNoModel) {
		t.Fatalf("err = %v, want ErrNoModel", err)
	}
}

func TestAnthropicChat_ChatStream_Text(t *testing.T) {
	t.Parallel()
	sse := strings.Join([]string{
		`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}`,
		``,
		`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		``,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel"}}`,
		``,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"lo"}}`,
		``,
		`data: {"type":"content_block_stop","index":0}`,
		``,
		`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`,
		``,
		`data: {"type":"message_stop"}`,
		``,
	}, "\n")
	ch := make(chan anthropicCaptured, 1)
	srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, ch)

	recs := make([]CallRecord, 0, 1)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"},
		WithAnthropicRecorder(func(cr CallRecord) { recs = append(recs, cr) }))

	var events []StreamEvent
	resp, err := ad.ChatStream(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test", Purpose: "chat"},
		func(ev StreamEvent) { events = append(events, ev) })
	if err != nil {
		t.Fatalf("ChatStream error: %v", err)
	}

	got := <-ch
	if got.accept != "text/event-stream" {
		t.Errorf("accept = %s, want text/event-stream", got.accept)
	}
	if got.body["stream"] != true {
		t.Errorf("body.stream = %v, want true", got.body["stream"])
	}

	if resp.Content != "Hello" {
		t.Errorf("resp.Content = %q, want Hello", resp.Content)
	}
	if resp.Model != "claude-test" {
		t.Errorf("resp.Model = %q, want claude-test", resp.Model)
	}
	if resp.TotalTokens != 7 {
		t.Errorf("resp.TotalTokens = %d, want 7", resp.TotalTokens)
	}
	wantEvents := []StreamEvent{
		{Kind: StreamKindTextStart, Index: 0},
		{Kind: StreamKindTextDelta, Index: 0, Delta: "Hel"},
		{Kind: StreamKindTextDelta, Index: 0, Delta: "lo"},
		{Kind: StreamKindTextEnd, Index: 0},
	}
	if !reflect.DeepEqual(events, wantEvents) {
		t.Errorf("events = %+v, want %+v", events, wantEvents)
	}
	if len(recs) != 1 || recs[0].Tokens != 7 || recs[0].Purpose != "chat" {
		t.Errorf("records = %+v, want 1 record tokens=7 purpose=chat", recs)
	}
}

func TestAnthropicChat_ChatStream_ToolUseReassembly(t *testing.T) {
	t.Parallel()
	sse := strings.Join([]string{
		`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":10}}}`,
		``,
		`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}`,
		``,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"location\":"}}`,
		``,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"Paris\"}"}}`,
		``,
		`data: {"type":"content_block_stop","index":0}`,
		``,
		`data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":8}}`,
		``,
		`data: {"type":"message_stop"}`,
		``,
	}, "\n")
	srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	var events []StreamEvent
	resp, err := ad.ChatStream(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test"},
		func(ev StreamEvent) { events = append(events, ev) })
	if err != nil {
		t.Fatalf("ChatStream error: %v", err)
	}

	if len(resp.ToolCalls) != 1 {
		t.Fatalf("resp.ToolCalls = %+v, want 1", resp.ToolCalls)
	}
	tc := resp.ToolCalls[0]
	if tc.Name != "lookup" {
		t.Errorf("tool call name = %q, want lookup", tc.Name)
	}
	if tc.Args["location"] != "Paris" {
		t.Errorf("tool call args.location = %v, want Paris (reassembled from partial_json)", tc.Args["location"])
	}
	if resp.Content != "" {
		t.Errorf("resp.Content = %q, want empty (tool-only response)", resp.Content)
	}
	if resp.TotalTokens != 18 {
		t.Errorf("resp.TotalTokens = %d, want 18", resp.TotalTokens)
	}
	wantEvents := []StreamEvent{
		{Kind: StreamKindToolCallStart, Index: 0, ID: "toolu_1"},
		{Kind: StreamKindToolCallDelta, Index: 0, ID: "toolu_1", Delta: `{"location":`},
		{Kind: StreamKindToolCallDelta, Index: 0, ID: "toolu_1", Delta: `"Paris"}`},
		{Kind: StreamKindToolCallEnd, Index: 0, ID: "toolu_1"},
	}
	if !reflect.DeepEqual(events, wantEvents) {
		t.Errorf("events = %+v, want %+v", events, wantEvents)
	}
}

func TestAnthropicChat_ChatStream_NoMessageStop(t *testing.T) {
	t.Parallel()
	sse := strings.Join([]string{
		`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}`,
		``,
		`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
		``,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}`,
		``,
	}, "\n")
	srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
	if _, err := ad.ChatStream(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test"}, nil); !errors.Is(err, ErrStream) {
		t.Fatalf("err = %v, want ErrStream", err)
	}
}

func TestAnthropicChat_ToolResult(t *testing.T) {
	t.Parallel()
	blk := AnthropicToolResult("toolu_1", `{"ok":true}`, true)
	if blk["type"] != "tool_result" {
		t.Errorf("type = %v, want tool_result", blk["type"])
	}
	if blk["tool_use_id"] != "toolu_1" {
		t.Errorf("tool_use_id = %v, want toolu_1", blk["tool_use_id"])
	}
	if blk["content"] != `{"ok":true}` {
		t.Errorf("content = %v, want the raw result", blk["content"])
	}
	if blk["is_error"] != true {
		t.Errorf("is_error = %v, want true", blk["is_error"])
	}
}