traefik-llmgateway/routes_messages_test.go

package traefikllmgateway

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// newMessagesRequest builds a POST /v1/messages request, JSON-encoding
// body. authHeader/authValue let a test choose which credential header to
// send; authHeader == "" sends no credential at all.
func newMessagesRequest(t *testing.T, body map[string]any, authHeader, authValue string) *http.Request {
	t.Helper()
	b, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("marshal request body: %v", err)
	}
	req := httptest.NewRequest(http.MethodPost, messagesPath, bytes.NewReader(b))
	if authHeader != "" {
		req.Header.Set(authHeader, authValue)
	}
	return req
}

// newMessagesTestGateway builds a *Gateway wired with an "anthropic"
// provider pointed at anthropicURL (skipped when "") and an "openai"
// provider pointed at openaiURL (skipped when ""), one permissive
// "default" group, and one inline user "alice" — the shared fixture
// every handleMessages test below starts from.
func newMessagesTestGateway(t *testing.T, anthropicURL, openaiURL string) *Gateway {
	t.Helper()
	cfg := CreateConfig()
	cfg.Providers = map[string]*ProviderConfig{}
	if anthropicURL != "" {
		cfg.Providers["anthropic"] = &ProviderConfig{Type: providerTypeAnthropic, BaseURL: anthropicURL, APIKey: "sk-ant-up", Models: []string{"claude-test"}} // #nosec G101 -- test fixture literal, not a real credential
	}
	if openaiURL != "" {
		cfg.Providers["openai"] = &ProviderConfig{Type: providerTypeOpenAI, BaseURL: openaiURL, APIKey: "sk-oai-up", Models: []string{"gpt-test"}}
	}
	cfg.Groups = map[string]*GroupConfig{"default": {}}
	cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}}
	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
	h, err := New(context.Background(), next, cfg, "llmgw")
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	gw, ok := h.(*Gateway)
	if !ok {
		t.Fatal("handler is not *Gateway")
	}
	return gw
}

// TestHandleMessages_AnthropicProvider_Passthrough proves a /v1/messages
// request resolved to an anthropic-type provider is forwarded to that
// provider's own Messages endpoint and its response copied back to the
// client byte-for-byte (brief: "pass through, no translation of the
// message body"), while still authenticating upstream and accounting
// usage exactly like every other route.
func TestHandleMessages_AnthropicProvider_Passthrough(t *testing.T) {
	const anthResp = `{"id":"msg_01ABC","type":"message","role":"assistant","content":[{"type":"text","text":"hi there"}],"model":"claude-real","stop_reason":"end_turn","usage":{"input_tokens":12,"output_tokens":6}}`

	var gotBody map[string]any
	var gotAPIKey, gotVersion, gotPath string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotAPIKey = r.Header.Get("x-api-key")
		gotVersion = r.Header.Get("anthropic-version")
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(anthResp))
	}))
	defer srv.Close()

	gw := newMessagesTestGateway(t, srv.URL, "")

	body := map[string]any{
		"model":      "claude-test",
		"max_tokens": 100,
		"messages":   []any{map[string]any{"role": "user", "content": "hello"}},
	}
	req := newMessagesRequest(t, body, "x-api-key", "sk-alice")
	rec := httptest.NewRecorder()
	gw.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
	}
	if rec.Body.String() != anthResp {
		t.Errorf("body = %q, want verbatim upstream body %q (no translation)", rec.Body.String(), anthResp)
	}
	if gotPath != anthropicMessagesPath {
		t.Errorf("upstream path = %q, want %q", gotPath, anthropicMessagesPath)
	}
	if gotAPIKey != "sk-ant-up" { // #nosec G101 -- test fixture literal, not a real credential
		t.Errorf("upstream x-api-key = %q, want sk-ant-up", gotAPIKey)
	}
	if gotVersion != anthropicAPIVersion {
		t.Errorf("upstream anthropic-version = %q, want %q", gotVersion, anthropicAPIVersion)
	}
	if gotBody["model"] != "claude-test" {
		t.Errorf("upstream model = %v, want claude-test (resolved upstream id)", gotBody["model"])
	}
	if _, hasAlias := gotBody[gatewayAliasKey]; hasAlias {
		t.Error("upstream request carries gatewayAliasKey; must never reach an upstream provider")
	}

	tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
	if !ok || tokIn != 12 {
		t.Errorf("user tokin/day counter = %d (ok=%v), want 12", tokIn, ok)
	}
	tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
	if !ok || tokOut != 6 {
		t.Errorf("user tokout/day counter = %d (ok=%v), want 6", tokOut, ok)
	}
}

// TestHandleMessages_OpenAIProvider_TranslatesBothDirections proves a
// /v1/messages request resolved to an openai-type provider is translated
// to OpenAI's chat-completion shape going out (Anthropic's top-level
// "system" field folded into a role:"system" message) and the upstream's
// OpenAI-shaped response translated back to Anthropic's Messages shape
// coming back, with usage accounted from the translated numbers.
func TestHandleMessages_OpenAIProvider_TranslatesBothDirections(t *testing.T) {
	const oaiResp = `{"id":"chatcmpl-99","object":"chat.completion","model":"gpt-real","choices":[{"index":0,"message":{"role":"assistant","content":"hi there"},"finish_reason":"stop"}],"usage":{"prompt_tokens":8,"completion_tokens":4}}`

	var gotReq map[string]any
	var gotPath string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		_ = json.NewDecoder(r.Body).Decode(&gotReq)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(oaiResp))
	}))
	defer srv.Close()

	gw := newMessagesTestGateway(t, "", srv.URL)

	body := map[string]any{
		"model":      "gpt-test",
		"max_tokens": 100,
		"system":     "You are terse.",
		"messages":   []any{map[string]any{"role": "user", "content": "hello"}},
	}
	req := newMessagesRequest(t, body, "x-api-key", "sk-alice")
	rec := httptest.NewRecorder()
	gw.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
	}
	if gotPath != "/v1/chat/completions" {
		t.Errorf("upstream path = %q, want /v1/chat/completions", gotPath)
	}

	msgs, _ := gotReq["messages"].([]any)
	if len(msgs) != 2 {
		t.Fatalf("upstream messages = %v, want 2 (system + user)", gotReq["messages"])
	}
	first, _ := msgs[0].(map[string]any)
	if first["role"] != "system" || first["content"] != "You are terse." {
		t.Errorf("upstream messages[0] = %v, want the system message translated into OpenAI's shape", first)
	}
	if _, hasSystem := gotReq["system"]; hasSystem {
		t.Error("upstream request still carries Anthropic's own \"system\" field; want it folded into messages[]")
	}

	var got map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode client response: %v", err)
	}
	if got["type"] != "message" || got["role"] != "assistant" {
		t.Errorf("response type/role = %v/%v, want message/assistant (Anthropic shape)", got["type"], got["role"])
	}
	content, _ := got["content"].([]any)
	if len(content) != 1 {
		t.Fatalf("response content = %v, want one text block", got["content"])
	}
	block, _ := content[0].(map[string]any)
	if block["type"] != "text" || block["text"] != "hi there" {
		t.Errorf("response content[0] = %v, want {type: text, text: \"hi there\"}", block)
	}
	if got["model"] != "gpt-test" {
		t.Errorf("response model = %v, want gpt-test (the client's own requested alias echoed back)", got["model"])
	}
	if got["stop_reason"] != "end_turn" {
		t.Errorf("response stop_reason = %v, want end_turn", got["stop_reason"])
	}

	usageMap, _ := got["usage"].(map[string]any)
	if usageMap["input_tokens"] != float64(8) || usageMap["output_tokens"] != float64(4) {
		t.Errorf("response usage = %v, want input_tokens=8 output_tokens=4", usageMap)
	}

	tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
	if !ok || tokIn != 8 {
		t.Errorf("user tokin/day counter = %d (ok=%v), want 8", tokIn, ok)
	}
	tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
	if !ok || tokOut != 4 {
		t.Errorf("user tokout/day counter = %d (ok=%v), want 4", tokOut, ok)
	}
}

// TestHandleMessages_Auth table-drives the three credential cases the
// brief calls out explicitly: x-api-key authenticates, Authorization:
// Bearer authenticates, and a bad key gets a 401 in ANTHROPIC error
// shape ({"type":"error","error":{"type":...}}), not OpenAI's
// ({"error":{...}}).
func TestHandleMessages_Auth(t *testing.T) {
	const anthResp = `{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-real","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(anthResp))
	}))
	defer srv.Close()

	cases := []struct {
		name       string
		authHeader string
		authValue  string
		wantStatus int
	}{
		{"x-api-key header authenticates", "x-api-key", "sk-alice", http.StatusOK},
		{"Authorization Bearer header authenticates", "Authorization", "Bearer sk-alice", http.StatusOK},
		{"bad key is rejected in Anthropic error shape", "x-api-key", "sk-not-a-real-key", http.StatusUnauthorized},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			gw := newMessagesTestGateway(t, srv.URL, "")
			body := map[string]any{"model": "claude-test", "max_tokens": 10, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
			req := newMessagesRequest(t, body, tc.authHeader, tc.authValue)
			rec := httptest.NewRecorder()
			gw.ServeHTTP(rec, req)

			if rec.Code != tc.wantStatus {
				t.Fatalf("status = %d, want %d, body=%s", rec.Code, tc.wantStatus, rec.Body.String())
			}
			if tc.wantStatus != http.StatusUnauthorized {
				return
			}

			var got map[string]any
			if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
				t.Fatalf("decode error body: %v", err)
			}
			if got["type"] != "error" {
				t.Errorf(`response["type"] = %v, want "error" (Anthropic's envelope shape, not OpenAI's flat one)`, got["type"])
			}
			errObj, _ := got["error"].(map[string]any)
			if errObj["type"] != "authentication_error" {
				t.Errorf("error.type = %v, want authentication_error", errObj["type"])
			}
			if _, hasTopLevelCode := got["code"]; hasTopLevelCode {
				t.Error(`response carries a top-level "code" field — that is OpenAI's writeOAIError shape, not Anthropic's`)
			}
		})
	}
}

// TestHandleMessages_StreamingRejected proves "stream": true is rejected
// cleanly with a typed, Anthropic-shaped error before any upstream call
// is made — never silently ignored, never answered non-streamed.
func TestHandleMessages_StreamingRejected(t *testing.T) {
	called := false
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called = true
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	gw := newMessagesTestGateway(t, srv.URL, "")
	body := map[string]any{"model": "claude-test", "max_tokens": 10, "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
	req := newMessagesRequest(t, body, "x-api-key", "sk-alice")
	rec := httptest.NewRecorder()
	gw.ServeHTTP(rec, req)

	if rec.Code != http.StatusNotImplemented {
		t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String())
	}
	if called {
		t.Error("upstream was called; streaming must be rejected before any upstream request")
	}

	var got map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode error body: %v", err)
	}
	if got["type"] != "error" {
		t.Errorf(`response["type"] = %v, want "error"`, got["type"])
	}
	errObj, _ := got["error"].(map[string]any)
	msg, _ := errObj["message"].(string)
	if !strings.Contains(msg, "stream") {
		t.Errorf("error message = %q, want it to mention streaming is not supported", msg)
	}
}

// TestHandleMessages_RateLimited_Returns429 proves this route enforces
// the caller's group/user requestsPerMinute limit exactly like
// routes_unified.go's runUnified: a second request past the limit is
// refused before any upstream call, carries a Retry-After header, and
// only the first, successful request's usage is accounted.
func TestHandleMessages_RateLimited_Returns429(t *testing.T) {
	upstreamCalls := 0
	const anthResp = `{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-real","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamCalls++
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(anthResp))
	}))
	defer srv.Close()

	cfg := CreateConfig()
	cfg.Providers = map[string]*ProviderConfig{"anthropic": {Type: providerTypeAnthropic, BaseURL: srv.URL, APIKey: "sk-ant-up", Models: []string{"claude-test"}}} // #nosec G101 -- test fixture literal, not a real credential
	cfg.Groups = map[string]*GroupConfig{"default": {}}
	cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}}}}
	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
	h, err := New(context.Background(), next, cfg, "llmgw")
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	gw, ok := h.(*Gateway)
	if !ok {
		t.Fatal("handler is not *Gateway")
	}

	body := map[string]any{"model": "claude-test", "max_tokens": 10, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}

	req1 := newMessagesRequest(t, body, "x-api-key", "sk-alice")
	rec1 := httptest.NewRecorder()
	gw.ServeHTTP(rec1, req1)
	if rec1.Code != http.StatusOK {
		t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
	}

	req2 := newMessagesRequest(t, body, "x-api-key", "sk-alice")
	rec2 := httptest.NewRecorder()
	gw.ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusTooManyRequests {
		t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String())
	}
	if rec2.Header().Get("Retry-After") == "" {
		t.Error("429 response missing Retry-After header")
	}
	if upstreamCalls != 1 {
		t.Errorf("upstream called %d times, want 1 (second request refused before any upstream call)", upstreamCalls)
	}

	tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
	if !ok || tokIn != 1 {
		t.Errorf("user tokin/day counter = %d (ok=%v), want 1 (only the first, successful request accounted)", tokIn, ok)
	}
}