client_test.go

package rig

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

// probeBackend is a fake OpenAI-compatible endpoint. It records the request it
// received and answers with a scripted reply; no test reaches the network.
type probeBackend struct {
	mu     sync.Mutex
	calls  int
	path   string
	header http.Header
	body   []byte

	tool    string // when set, the reply calls this tool
	content string
	finish  string // finish_reason; empty means "stop"
	status  int    // 0 means 200
	raw     []byte // when set, written verbatim as the 200 reply body
}

func (b *probeBackend) handler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		b.mu.Lock()
		defer b.mu.Unlock()
		b.calls++
		b.path = r.URL.Path
		b.header = r.Header.Clone()
		b.body, _ = io.ReadAll(r.Body)
		if b.status != 0 {
			http.Error(w, "upstream boom", b.status)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		if b.raw != nil {
			_, _ = w.Write(b.raw)
			return
		}
		message := map[string]any{"role": "assistant", "content": b.content}
		if b.tool != "" {
			message["tool_calls"] = []any{map[string]any{
				"id":       "call-1",
				"type":     "function",
				"function": map[string]any{"name": b.tool, "arguments": `{"message":"pong"}`},
			}}
		}
		finish := b.finish
		if finish == "" {
			finish = "stop"
		}
		_, _ = w.Write(probeReplyJSON(message, finish))
	}
}

func (b *probeBackend) snapshot() (int, string, http.Header, []byte) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.calls, b.path, b.header, b.body
}

func probeReplyJSON(message map[string]any, finish string) []byte {
	out, err := json.Marshal(map[string]any{
		"choices": []any{map[string]any{"index": 0, "message": message, "finish_reason": finish}},
	})
	if err != nil {
		panic(err)
	}
	return out
}

// probeReplyOfSize builds a valid tool-calling reply of exactly size bytes by
// padding the assistant content, so the response cap can be tested on its
// boundary rather than near it.
func probeReplyOfSize(t *testing.T, size int) []byte {
	t.Helper()
	build := func(content string) []byte {
		return probeReplyJSON(map[string]any{
			"role":    "assistant",
			"content": content,
			"tool_calls": []any{map[string]any{
				"id":       "call-1",
				"type":     "function",
				"function": map[string]any{"name": probeToolName, "arguments": `{"message":"pong"}`},
			}},
		}, "tool_calls")
	}
	pad := size - len(build(""))
	if pad < 0 {
		t.Fatalf("size %d is below the %d-byte minimum reply", size, len(build("")))
	}
	out := build(strings.Repeat("x", pad))
	if len(out) != size {
		t.Fatalf("built reply of %d bytes, want %d", len(out), size)
	}
	return out
}

func newProbeClient(t *testing.T, baseURL, apiKey string, opts ...Option) *Client {
	t.Helper()
	opts = append([]Option{WithHTTPClient(&http.Client{Timeout: 5 * time.Second})}, opts...)
	c, err := NewClient(baseURL, apiKey, opts...)
	if err != nil {
		t.Fatalf("NewClient: %v", err)
	}
	return c
}

func TestNewClientValidatesConfig(t *testing.T) {
	tests := []struct {
		name    string
		baseURL string
		opts    []Option
		wantErr bool
	}{
		{name: "http", baseURL: "http://127.0.0.1:1234/v1"},
		{name: "https", baseURL: "https://api.example.test/v1"},
		{name: "no path", baseURL: "https://api.example.test"},
		{name: "trailing slash", baseURL: "https://api.example.test/v1/"},
		{name: "empty", baseURL: "", wantErr: true},
		{name: "blank", baseURL: "   ", wantErr: true},
		{name: "no host", baseURL: "not-a-url", wantErr: true},
		{name: "wrong scheme", baseURL: "ftp://api.example.test/v1", wantErr: true},
		{name: "zero cap", baseURL: "https://api.example.test/v1", opts: []Option{WithMaxResponseBytes(0)}, wantErr: true},
		{name: "negative cap", baseURL: "https://api.example.test/v1", opts: []Option{WithMaxResponseBytes(-1)}, wantErr: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			c, err := NewClient(tc.baseURL, "sk-test", tc.opts...)
			if tc.wantErr {
				if err == nil {
					t.Fatalf("NewClient(%q) = %v, want error", tc.baseURL, c)
				}
				return
			}
			if err != nil {
				t.Fatalf("NewClient(%q): %v", tc.baseURL, err)
			}
			if c == nil {
				t.Fatal("NewClient returned a nil client without an error")
			}
		})
	}
}

func TestProbeToolCalling(t *testing.T) {
	tests := []struct {
		name    string
		backend *probeBackend
		wantErr error // nil means the probe must pass
		wantAny bool  // any error, matching no sentinel
	}{
		{name: "tool call", backend: &probeBackend{tool: probeToolName, finish: "tool_calls"}},
		{name: "text only", backend: &probeBackend{content: "I would call ping."}, wantErr: ErrNoToolCalling},
		{name: "other tool", backend: &probeBackend{tool: "not_ping", finish: "tool_calls"}, wantErr: ErrNoToolCalling},
		{name: "no choices", backend: &probeBackend{raw: []byte(`{"choices":[]}`)}, wantErr: ErrEmptyReply},
		{name: "upstream error", backend: &probeBackend{status: http.StatusInternalServerError}, wantAny: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			srv := httptest.NewServer(tc.backend.handler())
			defer srv.Close()
			err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3")
			switch {
			case tc.wantAny:
				if err == nil {
					t.Fatal("ProbeToolCalling = nil, want an error")
				}
				if errors.Is(err, ErrNoToolCalling) {
					t.Fatalf("upstream failure reported as ErrNoToolCalling: %v", err)
				}
			case tc.wantErr != nil:
				if !errors.Is(err, tc.wantErr) {
					t.Fatalf("ProbeToolCalling = %v, want %v", err, tc.wantErr)
				}
			default:
				if err != nil {
					t.Fatalf("ProbeToolCalling: %v", err)
				}
			}
		})
	}
}

// A probe reply truncated at the 256-token budget says nothing about tool
// calling: a reasoning model can spend the whole budget on chain-of-thought
// before it emits a call. Reporting that as ErrNoToolCalling tells the caller
// to reject a model that supports tools fine.
func TestProbeTruncatedReplyIsAnOutputLimit(t *testing.T) {
	backend := &probeBackend{content: "<think>let me consider the ping tool", finish: "length"}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3")
	if !errors.Is(err, ErrOutputLimit) {
		t.Fatalf("ProbeToolCalling = %v, want ErrOutputLimit", err)
	}
	if errors.Is(err, ErrNoToolCalling) {
		t.Errorf("truncation reported as ErrNoToolCalling: %v", err)
	}
}

// A 200 whose body is JSON null decodes to a nil completion with a nil error.
func TestProbeSurvivesNullBody(t *testing.T) {
	backend := &probeBackend{raw: []byte(`null`)}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3")
	if !errors.Is(err, ErrEmptyReply) {
		t.Fatalf("ProbeToolCalling = %v, want ErrEmptyReply", err)
	}
	if errors.Is(err, ErrNoToolCalling) {
		t.Fatalf("ProbeToolCalling = %v, must not read as a tool-calling verdict", err)
	}
}

// finish_reason is an endpoint-controlled string. A caller rendering a failed
// connection test must not be handed markup by the endpoint it was testing.
func TestProbeErrorCarriesNoReplyText(t *testing.T) {
	const injected = `<img src=x onerror=alert(1)>`
	backend := &probeBackend{content: "I would call " + injected, finish: injected}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3")
	if !errors.Is(err, ErrNoToolCalling) {
		t.Fatalf("ProbeToolCalling = %v, want ErrNoToolCalling", err)
	}
	if strings.Contains(err.Error(), "img src") {
		t.Errorf("error %q carries endpoint-controlled reply text", err)
	}
}

func TestProbeRejectsEmptyModel(t *testing.T) {
	backend := &probeBackend{tool: probeToolName}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()
	if err := newProbeClient(t, srv.URL+"/v1", "").ProbeToolCalling(context.Background(), "  "); err == nil {
		t.Fatal("ProbeToolCalling with a blank model = nil, want an error")
	}
	if calls, _, _, _ := backend.snapshot(); calls != 0 {
		t.Fatalf("backend was called %d times for a blank model", calls)
	}
}

// The probe must state its own output budget and offer exactly the one tool it
// checks for; a budget-less request is truncated by the upstream default.
func TestProbeRequestShape(t *testing.T) {
	backend := &probeBackend{tool: probeToolName, finish: "tool_calls"}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()
	if err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3"); err != nil {
		t.Fatalf("ProbeToolCalling: %v", err)
	}
	_, path, _, body := backend.snapshot()
	if path != "/v1/chat/completions" {
		t.Errorf("upstream path = %q, want /v1/chat/completions", path)
	}
	var sent struct {
		Model     string `json:"model"`
		MaxTokens *int64 `json:"max_tokens"`
		Messages  []struct {
			Role string `json:"role"`
		} `json:"messages"`
		Tools []struct {
			Type     string `json:"type"`
			Function struct {
				Name       string         `json:"name"`
				Parameters map[string]any `json:"parameters"`
			} `json:"function"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(body, &sent); err != nil {
		t.Fatalf("upstream request JSON: %v (body=%s)", err, body)
	}
	if sent.Model != "qwen3" {
		t.Errorf("model = %q, want qwen3", sent.Model)
	}
	if sent.MaxTokens == nil || *sent.MaxTokens != probeMaxTokens {
		t.Errorf("max_tokens = %v, want %d", sent.MaxTokens, probeMaxTokens)
	}
	if len(sent.Messages) != 2 || sent.Messages[0].Role != "system" || sent.Messages[1].Role != "user" {
		t.Errorf("messages = %+v, want system then user", sent.Messages)
	}
	if len(sent.Tools) != 1 || sent.Tools[0].Function.Name != probeToolName {
		t.Fatalf("tools = %+v, want one %q function", sent.Tools, probeToolName)
	}
	if sent.Tools[0].Function.Parameters["type"] != "object" {
		t.Errorf("tool parameters = %+v, want a JSON Schema object", sent.Tools[0].Function.Parameters)
	}
}

// Ambient credentials in the process environment must never reach a
// caller-configured endpoint, however the SDK would otherwise pick them up.
func TestClientScrubsAmbientCredentials(t *testing.T) {
	t.Setenv("OPENAI_API_KEY", "sk-ambient-leak")
	t.Setenv("OPENAI_ORG_ID", "org-ambient")
	t.Setenv("OPENAI_PROJECT_ID", "proj-ambient")
	t.Setenv("OPENAI_CUSTOM_HEADERS", "X-Gateway-Token: ambient-gateway\nX-Tenant: ambient-tenant")
	t.Setenv("OPENAI_BASE_URL", "http://127.0.0.1:1/v1")

	backend := &probeBackend{tool: probeToolName, finish: "tool_calls"}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	c := newProbeClient(t, srv.URL+"/v1", "sk-explicit", WithUserAgent("rig-test/1"))
	if err := c.ProbeToolCalling(context.Background(), "qwen3"); err != nil {
		t.Fatalf("ProbeToolCalling: %v", err)
	}
	calls, path, header, _ := backend.snapshot()
	if calls != 1 {
		t.Fatalf("backend calls = %d, want 1 (OPENAI_BASE_URL must not win)", calls)
	}
	if path != "/v1/chat/completions" {
		t.Errorf("upstream path = %q, want /v1/chat/completions", path)
	}
	if got := header.Get("Authorization"); got != "Bearer sk-explicit" {
		t.Errorf("Authorization = %q, want the explicit key", got)
	}
	for _, name := range []string{"OpenAI-Organization", "OpenAI-Project", "X-Gateway-Token", "X-Tenant"} {
		if got := header.Get(name); got != "" {
			t.Errorf("%s reached the wire with %q, want it scrubbed", name, got)
		}
	}
	if got := header.Get("User-Agent"); got != "rig-test/1" {
		t.Errorf("User-Agent = %q, want rig-test/1", got)
	}
}

// An empty key is the Ollama case: no Authorization header at all, and in
// particular not one built from OPENAI_API_KEY.
func TestClientOmitsAuthorizationWithoutKey(t *testing.T) {
	t.Setenv("OPENAI_API_KEY", "sk-ambient-leak")

	backend := &probeBackend{tool: probeToolName, finish: "tool_calls"}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	if err := newProbeClient(t, srv.URL+"/v1", "").ProbeToolCalling(context.Background(), "qwen3"); err != nil {
		t.Fatalf("ProbeToolCalling: %v", err)
	}
	_, _, header, _ := backend.snapshot()
	if got := header.Get("Authorization"); got != "" {
		t.Fatalf("Authorization = %q, want none", got)
	}
}

func TestResponseCap(t *testing.T) {
	const limit = 4096
	tests := []struct {
		name string
		size int
		want error
	}{
		{name: "under cap", size: limit / 2},
		{name: "exactly cap", size: limit},
		{name: "one over cap", size: limit + 1, want: ErrResponseTooLarge},
		{name: "far over cap", size: limit * 4, want: ErrResponseTooLarge},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			backend := &probeBackend{raw: probeReplyOfSize(t, tc.size)}
			srv := httptest.NewServer(backend.handler())
			defer srv.Close()
			c := newProbeClient(t, srv.URL+"/v1", "sk-test", WithMaxResponseBytes(limit))
			err := c.ProbeToolCalling(context.Background(), "qwen3")
			if tc.want == nil {
				if err != nil {
					t.Fatalf("ProbeToolCalling: %v", err)
				}
				return
			}
			if !errors.Is(err, tc.want) {
				t.Fatalf("ProbeToolCalling = %v, want %v", err, tc.want)
			}
		})
	}
}

// The default cap applies when no option sets one.
func TestResponseCapDefaults(t *testing.T) {
	backend := &probeBackend{raw: probeReplyOfSize(t, defaultMaxResponseBytes+1)}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()
	err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(context.Background(), "qwen3")
	if !errors.Is(err, ErrResponseTooLarge) {
		t.Fatalf("ProbeToolCalling = %v, want ErrResponseTooLarge", err)
	}
}

// WithMaxResponseBytes(math.MaxInt64) is how a caller says "no cap"; it must
// read a reply rather than panic on the first byte of one.
func TestResponseCapAtMaxInt(t *testing.T) {
	backend := &probeBackend{tool: probeToolName, finish: "tool_calls"}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()

	c := newProbeClient(t, srv.URL+"/v1", "sk-test", WithMaxResponseBytes(math.MaxInt64))
	if err := c.ProbeToolCalling(context.Background(), "qwen3"); err != nil {
		t.Fatalf("ProbeToolCalling: %v", err)
	}
}

// A failure must not describe the endpoint: the address and the key are the
// two things an error a caller logs must never carry.
func TestBackendErrorsHideTheEndpoint(t *testing.T) {
	backend := &probeBackend{status: http.StatusUnauthorized}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()
	host := strings.TrimPrefix(srv.URL, "http://")

	err := newProbeClient(t, srv.URL+"/v1", "sk-secret-value").ProbeToolCalling(context.Background(), "qwen3")
	if err == nil {
		t.Fatal("ProbeToolCalling = nil, want an error")
	}
	msg := err.Error()
	for _, leak := range []string{host, srv.URL, "sk-secret-value"} {
		if strings.Contains(msg, leak) {
			t.Fatalf("error message %q leaks %q", msg, leak)
		}
	}
	if !strings.Contains(msg, "401") {
		t.Errorf("error message %q drops the upstream status", msg)
	}
}

func TestBackendErrorPreservesContextCause(t *testing.T) {
	backend := &probeBackend{tool: probeToolName}
	srv := httptest.NewServer(backend.handler())
	defer srv.Close()
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	err := newProbeClient(t, srv.URL+"/v1", "sk-test").ProbeToolCalling(ctx, "qwen3")
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("ProbeToolCalling = %v, want context.Canceled", err)
	}
}

func TestAmbientHeaderNames(t *testing.T) {
	tests := []struct {
		name   string
		env    string
		want   []string
		absent []string
	}{
		{name: "unset", want: []string{"OpenAI-Organization", "OpenAI-Project"}},
		{name: "single", env: "X-Gateway-Token: t", want: []string{"X-Gateway-Token"}},
		{name: "multiple", env: "X-A: 1\nX-B: 2", want: []string{"X-A", "X-B"}},
		{name: "padded names", env: "  X-Padded : 1", want: []string{"X-Padded"}},
		{name: "no colon", env: "garbage", absent: []string{"garbage"}},
		{name: "blank name", env: ": 1", absent: []string{""}},
		{name: "owned headers", env: "Authorization: a\ncontent-type: b\nAccept: c", absent: []string{"Authorization", "content-type", "Accept"}},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Setenv("OPENAI_CUSTOM_HEADERS", tc.env)
			got := ambientHeaderNames()
			index := make(map[string]bool, len(got))
			for _, name := range got {
				index[name] = true
			}
			for _, want := range tc.want {
				if !index[want] {
					t.Errorf("ambientHeaderNames() = %v, missing %q", got, want)
				}
			}
			for _, absent := range tc.absent {
				if index[absent] {
					t.Errorf("ambientHeaderNames() = %v, must not contain %q", got, absent)
				}
			}
			if !index["OpenAI-Organization"] || !index["OpenAI-Project"] {
				t.Errorf("ambientHeaderNames() = %v, want the fixed identifier headers always", got)
			}
		})
	}
}

func TestLimitedBodyStopsPastTheCap(t *testing.T) {
	tests := []struct {
		name    string
		payload int
		limit   int64
		wantErr error
	}{
		{name: "under", payload: 10, limit: 16},
		{name: "exact", payload: 16, limit: 16},
		{name: "over", payload: 17, limit: 16, wantErr: ErrResponseTooLarge},
		// A caller's way of spelling "no cap". The budget arithmetic must not
		// overflow into a negative reslice bound.
		{name: "max int cap", payload: 4096, limit: math.MaxInt64},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			body := &limitedBody{
				rc:   io.NopCloser(strings.NewReader(strings.Repeat("x", tc.payload))),
				left: tc.limit,
			}
			got, err := io.ReadAll(body)
			if !errors.Is(err, tc.wantErr) {
				t.Fatalf("ReadAll error = %v, want %v", err, tc.wantErr)
			}
			if tc.wantErr == nil && len(got) != tc.payload {
				t.Fatalf("read %d bytes, want %d", len(got), tc.payload)
			}
			if err := body.Close(); err != nil {
				t.Fatalf("Close: %v", err)
			}
		})
	}
}