Vectrum
Sign In
487 lines · 16.1 KB · Go
package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

const okResponse = `{"id":"msg_1","type":"message","role":"assistant","model":"claude-opus-5",
"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":2}}`

func testClient(t *testing.T, h http.HandlerFunc) *Client {
	t.Helper()
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	c := New(WithBaseURL(srv.URL), WithAPIKey("test-key"))
	c.RetryDelay = time.Millisecond // keep the retry tests fast
	return c
}

func TestCreateMessageSendsExpectedHeadersAndBody(t *testing.T) {
	var gotHeaders http.Header
	var gotBody map[string]any
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		gotHeaders = r.Header.Clone()
		body, _ := io.ReadAll(r.Body)
		json.Unmarshal(body, &gotBody)
		if r.URL.Path != "/v1/messages" {
			t.Errorf("path = %q, want /v1/messages", r.URL.Path)
		}
		io.WriteString(w, okResponse)
	})

	resp, err := c.CreateMessage(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: []ContentBlock{TextBlock("hello")}}},
	})
	if err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if resp.Text() != "hi" {
		t.Errorf("Text() = %q, want hi", resp.Text())
	}
	if gotHeaders.Get("x-api-key") != "test-key" {
		t.Errorf("x-api-key = %q", gotHeaders.Get("x-api-key"))
	}
	if gotHeaders.Get("anthropic-version") != apiVersion {
		t.Errorf("anthropic-version = %q, want %s", gotHeaders.Get("anthropic-version"), apiVersion)
	}
	if gotBody["model"] != DefaultModel {
		t.Errorf("model = %v, want the client default %s", gotBody["model"], DefaultModel)
	}
	if gotBody["max_tokens"].(float64) != DefaultMaxTokens {
		t.Errorf("max_tokens = %v, want the client default", gotBody["max_tokens"])
	}
}

func TestOAuthTokenUsesBearerAndBetaHeader(t *testing.T) {
	var got http.Header
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Clone()
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	c := New(WithBaseURL(srv.URL), WithAuthToken("oauth-tok"))
	if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if got.Get("authorization") != "Bearer oauth-tok" {
		t.Errorf("authorization = %q, want Bearer (OAuth tokens do not go on x-api-key)", got.Get("authorization"))
	}
	if got.Get("x-api-key") != "" {
		t.Errorf("x-api-key should be unset when using an OAuth token, got %q", got.Get("x-api-key"))
	}
	if got.Get("anthropic-beta") != oauthBeta {
		t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
	}
}

func TestNoCredentialsIsAClearError(t *testing.T) {
	c := New(WithBaseURL("http://127.0.0.1:1"))
	c.apiKey, c.authToken = "", ""
	_, err := c.CreateMessage(context.Background(), Request{})
	if err == nil || !strings.Contains(err.Error(), "ANTHROPIC_API_KEY") {
		t.Errorf("error = %v, want it to name the env var", err)
	}
	if c.HasCredentials() {
		t.Error("HasCredentials should be false")
	}
}

func TestRetriesTransientStatusesAndNotClientErrors(t *testing.T) {
	tests := []struct {
		name        string
		status      int
		wantCalls   int32
		wantSucceed bool
	}{
		{name: "429 is retried", status: http.StatusTooManyRequests, wantCalls: 3},
		{name: "500 is retried", status: http.StatusInternalServerError, wantCalls: 3},
		{name: "400 is not retried", status: http.StatusBadRequest, wantCalls: 1},
		{name: "401 is not retried", status: http.StatusUnauthorized, wantCalls: 1},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var calls int32
			c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
				atomic.AddInt32(&calls, 1)
				w.WriteHeader(tc.status)
				io.WriteString(w, `{"type":"error","error":{"type":"boom","message":"nope"}}`)
			})
			c.MaxRetries = 2
			c.HTTP.Timeout = 5 * time.Second

			_, err := c.CreateMessage(context.Background(), Request{})
			if err == nil {
				t.Fatal("want an error")
			}
			if got := atomic.LoadInt32(&calls); got != tc.wantCalls {
				t.Errorf("server saw %d calls, want %d", got, tc.wantCalls)
			}
			apiErr, ok := err.(*APIError)
			if !ok {
				t.Fatalf("error type = %T, want *APIError", err)
			}
			if apiErr.StatusCode != tc.status || apiErr.Message == "" {
				t.Errorf("APIError = %+v, want the status and parsed message", apiErr)
			}
		})
	}
}

func TestRetrySucceedsOnSecondAttempt(t *testing.T) {
	var calls int32
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		if atomic.AddInt32(&calls, 1) == 1 {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		io.WriteString(w, okResponse)
	})
	resp, err := c.CreateMessage(context.Background(), Request{})
	if err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if resp.Text() != "hi" || atomic.LoadInt32(&calls) != 2 {
		t.Errorf("got %q after %d calls, want hi after 2", resp.Text(), calls)
	}
}

func TestContextCancellationStopsRetrying(t *testing.T) {
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := c.CreateMessage(ctx, Request{}); err == nil {
		t.Fatal("want an error from the cancelled context")
	}
}

func TestMalformedResponseBodyIsReported(t *testing.T) {
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"content":`)
	})
	_, err := c.CreateMessage(context.Background(), Request{})
	if err == nil || !strings.Contains(err.Error(), "decoding response") {
		t.Errorf("error = %v, want a decode error", err)
	}
}

func TestRetryAfterHintIsParsed(t *testing.T) {
	got, ok := retryAfter("rate limited (retry-after 7s)")
	if !ok || got != 7*time.Second {
		t.Errorf("retryAfter = (%v, %v), want (7s, true)", got, ok)
	}
	if _, ok := retryAfter("rate limited"); ok {
		t.Error("want no hint when the header was absent")
	}
}

func TestBaseURLAndCredentialResolution(t *testing.T) {
	tests := []struct {
		name          string
		env           map[string]string
		opts          []Option
		wantURL       string
		wantAPIKey    string
		wantAuthToken string
	}{
		{
			name:       "defaults",
			wantURL:    DefaultBaseURL,
			wantAPIKey: "",
		},
		{
			name:       "environment is picked up",
			env:        map[string]string{"ANTHROPIC_BASE_URL": "http://localhost:4000", "ANTHROPIC_API_KEY": "env-key"},
			wantURL:    "http://localhost:4000",
			wantAPIKey: "env-key",
		},
		{
			name:       "flags override the environment",
			env:        map[string]string{"ANTHROPIC_BASE_URL": "http://from-env", "ANTHROPIC_API_KEY": "env-key"},
			opts:       []Option{WithBaseURL("http://from-flag"), WithAPIKey("flag-key")},
			wantURL:    "http://from-flag",
			wantAPIKey: "flag-key",
		},
		{
			name:       "empty options leave the environment alone",
			env:        map[string]string{"ANTHROPIC_BASE_URL": "http://from-env", "ANTHROPIC_API_KEY": "env-key"},
			opts:       []Option{WithBaseURL(""), WithAPIKey(""), WithAuthToken(""), WithModel("")},
			wantURL:    "http://from-env",
			wantAPIKey: "env-key",
		},
		{
			name:    "trailing slashes are trimmed",
			env:     map[string]string{"ANTHROPIC_BASE_URL": "http://gateway.internal/v1/// "},
			wantURL: "http://gateway.internal/v1",
		},
		{
			name:          "api key shadows a bearer token from the environment",
			env:           map[string]string{"ANTHROPIC_API_KEY": "k", "ANTHROPIC_AUTH_TOKEN": "t"},
			wantURL:       DefaultBaseURL,
			wantAPIKey:    "k",
			wantAuthToken: "",
		},
		{
			name:          "bearer token is used when no api key is set",
			env:           map[string]string{"ANTHROPIC_AUTH_TOKEN": "t"},
			wantURL:       DefaultBaseURL,
			wantAuthToken: "t",
		},
		{
			name:          "auth-token flag displaces an api key from the environment",
			env:           map[string]string{"ANTHROPIC_API_KEY": "env-key"},
			opts:          []Option{WithAuthToken("flag-token")},
			wantURL:       DefaultBaseURL,
			wantAuthToken: "flag-token",
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			for _, k := range []string{"ANTHROPIC_BASE_URL", "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"} {
				t.Setenv(k, tc.env[k])
			}
			c := New(tc.opts...)
			if c.BaseURL != tc.wantURL {
				t.Errorf("BaseURL = %q, want %q", c.BaseURL, tc.wantURL)
			}
			if c.apiKey != tc.wantAPIKey {
				t.Errorf("apiKey = %q, want %q", c.apiKey, tc.wantAPIKey)
			}
			if c.authToken != tc.wantAuthToken {
				t.Errorf("authToken = %q, want %q", c.authToken, tc.wantAuthToken)
			}
			if wantDefault := tc.wantURL == DefaultBaseURL; c.IsDefaultEndpoint() != wantDefault {
				t.Errorf("IsDefaultEndpoint = %v, want %v", c.IsDefaultEndpoint(), wantDefault)
			}
		})
	}
}

func TestCustomBaseURLIsActuallyUsed(t *testing.T) {
	var gotPath string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	// A trailing slash must not produce a request to "//v1/messages".
	c := New(WithBaseURL(srv.URL+"/"), WithAPIKey("k"))
	if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if gotPath != "/v1/messages" {
		t.Errorf("path = %q, want /v1/messages", gotPath)
	}
}

// A token source is consulted per request, so a credential that rotates
// mid-session reaches the wire without rebuilding the client.
func TestTokenSourceIsCalledPerRequest(t *testing.T) {
	var got []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = append(got, r.Header.Get("authorization"))
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	n := 0
	c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
		n++
		return fmt.Sprintf("tok-%d", n), nil
	}))
	for i := 0; i < 2; i++ {
		if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
			t.Fatalf("CreateMessage: %v", err)
		}
	}
	if len(got) != 2 || got[0] != "Bearer tok-1" || got[1] != "Bearer tok-2" {
		t.Errorf("authorization headers = %v, want Bearer tok-1 then tok-2", got)
	}
}

func TestTokenSourceGetsTheOAuthBetaHeader(t *testing.T) {
	var got http.Header
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Clone()
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
		return "tok", nil
	}))
	if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if got.Get("anthropic-beta") != oauthBeta {
		t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
	}
	if got.Get("x-api-key") != "" {
		t.Errorf("x-api-key should be unset, got %q", got.Get("x-api-key"))
	}
}

// A credential that cannot be resolved is permanent. Retrying it three times
// with backoff just delays the error, so resolution happens before the loop.
func TestTokenSourceErrorIsNotRetried(t *testing.T) {
	var hits int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits++
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	var calls int
	c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
		calls++
		return "", errors.New("login expired")
	}))
	c.RetryDelay = time.Millisecond
	_, err := c.CreateMessage(context.Background(), Request{})
	if err == nil || !strings.Contains(err.Error(), "login expired") {
		t.Fatalf("err = %v, want the source's error", err)
	}
	if calls != 1 {
		t.Errorf("token source called %d times, want 1", calls)
	}
	if hits != 0 {
		t.Errorf("server hit %d times, want 0 — no request should have been sent", hits)
	}
}

func TestHasCredentialsCountsATokenSource(t *testing.T) {
	c := New(WithBaseURL("http://127.0.0.1:1"))
	c.apiKey, c.authToken = "", ""
	if c.HasCredentials() {
		t.Fatal("no credentials expected before the source is set")
	}
	c.tokenSource = func(context.Context) (string, error) { return "tok", nil }
	if !c.HasCredentials() {
		t.Error("a token source is a credential")
	}
}

// The streaming path must authenticate identically to the buffered one.
func TestStreamUsesTheTokenSource(t *testing.T) {
	streamOK := sse(
		`{"type":"message_start","message":{"id":"msg_1","model":"claude-opus-5","usage":{"input_tokens":100,"cache_read_input_tokens":40}}}`,
		`{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":55}}`,
		`{"type":"message_stop"}`,
	)

	var got http.Header
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Clone()
		w.Header().Set("content-type", "text/event-stream")
		io.WriteString(w, streamOK)
	}))
	defer srv.Close()

	c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
		return "stream-tok", nil
	}))
	if _, err := c.StreamMessage(context.Background(), Request{}, func(StreamEvent) {}); err != nil {
		t.Fatalf("StreamMessage: %v", err)
	}
	if got.Get("authorization") != "Bearer stream-tok" {
		t.Errorf("authorization = %q", got.Get("authorization"))
	}
	if got.Get("anthropic-beta") != oauthBeta {
		t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
	}
}

// The model is switchable mid-session, so what a request carries is whatever
// was last set rather than what the client was built with.
func TestSetModelChangesWhatTheNextRequestAsksFor(t *testing.T) {
	var models []string
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		raw, _ := io.ReadAll(r.Body)
		json.Unmarshal(raw, &body)
		models = append(models, fmt.Sprint(body["model"]))
		io.WriteString(w, okResponse)
	})
	WithModel("claude-opus-5")(c)

	msg := Request{Messages: []Message{{Role: RoleUser, Content: []ContentBlock{TextBlock("hi")}}}}
	if _, err := c.CreateMessage(context.Background(), msg); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}

	c.SetModel("claude-haiku-4-5")
	if got := c.Model(); got != "claude-haiku-4-5" {
		t.Errorf("Model() = %q", got)
	}
	if _, err := c.CreateMessage(context.Background(), msg); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}

	if want := "claude-opus-5 claude-haiku-4-5"; strings.Join(models, " ") != want {
		t.Errorf("models sent = %v, want %q", models, want)
	}
}

// A subagent runs on its own model without disturbing the session's. Clone is
// how it gets a client that authenticates exactly as the parent does and
// differs in nothing else.
func TestCloneKeepsCredentialsAndOwnsItsModel(t *testing.T) {
	var models []string
	var keys []string
	parent := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		raw, _ := io.ReadAll(r.Body)
		json.Unmarshal(raw, &body)
		models = append(models, fmt.Sprint(body["model"]))
		keys = append(keys, r.Header.Get("x-api-key"))
		io.WriteString(w, okResponse)
	})
	parent.SetModel("claude-opus-5")

	child := parent.Clone("claude-haiku-4-5")
	msg := Request{Messages: []Message{{Role: RoleUser, Content: []ContentBlock{TextBlock("hi")}}}}
	if _, err := child.CreateMessage(context.Background(), msg); err != nil {
		t.Fatalf("child CreateMessage: %v", err)
	}
	if _, err := parent.CreateMessage(context.Background(), msg); err != nil {
		t.Fatalf("parent CreateMessage: %v", err)
	}

	if want := "claude-haiku-4-5 claude-opus-5"; strings.Join(models, " ") != want {
		t.Errorf("models sent = %v, want %q", models, want)
	}
	if keys[0] != "test-key" {
		t.Errorf("clone sent x-api-key %q, want the parent's", keys[0])
	}
	if child.BaseURL != parent.BaseURL {
		t.Errorf("clone BaseURL = %q, want %q", child.BaseURL, parent.BaseURL)
	}

	// Switching the session's model mid-run must not retarget a child that is
	// already working.
	parent.SetModel("claude-sonnet-5")
	if got := child.Model(); got != "claude-haiku-4-5" {
		t.Errorf("after parent switched, child Model() = %q", got)
	}
}

// An empty model means "whatever the parent is on", resolved at clone time.
func TestCloneWithNoModelInheritsTheParents(t *testing.T) {
	parent := New(WithAPIKey("k"), WithModel("claude-sonnet-5"))
	if got := parent.Clone("").Model(); got != "claude-sonnet-5" {
		t.Errorf("Clone(\"\").Model() = %q, want the parent's", got)
	}
}