client.go

package rig

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/openai/openai-go/v3"
	"github.com/openai/openai-go/v3/option"
	"github.com/openai/openai-go/v3/shared"
)

// defaultMaxResponseBytes caps one decompressed reply. The endpoint is
// caller-configured and the SDK buffers the whole body with io.ReadAll before
// decoding, so without a ceiling a host that answers with an endless stream --
// or a gzip bomb, since the transport asks for gzip -- exhausts the process.
const defaultMaxResponseBytes = 1 << 20

// The probe asks for one trivial tool call and nothing else, so its budget is
// small enough that a model which ignores the tool cannot burn a real one.
const (
	probeToolName     = "ping"
	probeMaxTokens    = 256
	probeSystemPrompt = `You are a connection probe. Call the ping tool. Do not answer with text.`
	probeUserPrompt   = `Call the ping tool with message "ping".`
)

// Client is one configured OpenAI-compatible endpoint. It is safe for
// concurrent use; the SDK client it wraps is immutable after construction.
type Client struct {
	sdk openai.Client
}

type config struct {
	httpClient *http.Client
	maxBytes   int64
	userAgent  string
}

// Option configures a Client at construction time. Options never read the
// environment; every value reaching the wire is supplied by the caller.
type Option func(*config)

// WithHTTPClient supplies the transport. The endpoint is caller-configured, so
// the SSRF guard, timeout and redirect policy belong to the caller's client.
func WithHTTPClient(hc *http.Client) Option {
	return func(c *config) { c.httpClient = hc }
}

// WithMaxResponseBytes sets the cap on one decompressed response body. Must be
// positive; the default is 1 MiB.
func WithMaxResponseBytes(n int64) Option {
	return func(c *config) { c.maxBytes = n }
}

// WithUserAgent overrides the SDK's User-Agent header.
func WithUserAgent(ua string) Option {
	return func(c *config) { c.userAgent = ua }
}

// NewClient builds a client for baseURL. apiKey may be empty for endpoints
// that need no credential (Ollama and friends), in which case no
// Authorization header is sent -- including one inherited from the process
// environment.
func NewClient(baseURL, apiKey string, opts ...Option) (*Client, error) {
	cfg := config{maxBytes: defaultMaxResponseBytes}
	for _, opt := range opts {
		if opt != nil {
			opt(&cfg)
		}
	}
	if cfg.maxBytes < 1 {
		return nil, errors.New("max response bytes must be positive")
	}
	base, err := normalizeBaseURL(baseURL)
	if err != nil {
		return nil, err
	}
	return &Client{sdk: openai.NewClient(sdkOptions(base, apiKey, cfg)...)}, nil
}

// normalizeBaseURL accepts only an absolute http(s) URL. Anything else is a
// configuration error rather than a request that fails later at the transport.
func normalizeBaseURL(base string) (string, error) {
	u, err := url.Parse(strings.TrimSpace(base))
	if err != nil || u.Host == "" {
		return "", errors.New("base URL must be an absolute http(s) URL")
	}
	if !strings.EqualFold(u.Scheme, "http") && !strings.EqualFold(u.Scheme, "https") {
		return "", errors.New("base URL scheme must be http or https")
	}
	return u.String(), nil
}

// sdkOptions builds the SDK client options. openai.NewClient prepends
// DefaultClientOptions unconditionally and the option that would suppress it
// is internal to the SDK, so every ambient value it can produce is overwritten
// or deleted here: the base URL and the key are set explicitly (the key
// explicitly cleared when there is none), the two fixed identifier headers and
// every name declared by OPENAI_CUSTOM_HEADERS are deleted. Retries are
// disabled -- one call is one upstream request, and a retry would multiply
// both the timeout and the load a loop can direct at a host.
func sdkOptions(base, apiKey string, cfg config) []option.RequestOption {
	opts := []option.RequestOption{
		option.WithBaseURL(base),
		option.WithMaxRetries(0),
		option.WithMiddleware(limitResponseBody(cfg.maxBytes)),
		option.WithAPIKey(apiKey),
	}
	if cfg.httpClient != nil {
		opts = append(opts, option.WithHTTPClient(cfg.httpClient))
	}
	if apiKey == "" {
		opts = append(opts, option.WithHeaderDel("Authorization"))
	}
	for _, name := range ambientHeaderNames() {
		opts = append(opts, option.WithHeaderDel(name))
	}
	if cfg.userAgent != "" {
		opts = append(opts, option.WithHeader("User-Agent", cfg.userAgent))
	}
	return opts
}

// ownedHeaders are headers the request needs to work at all. A same-named
// value inherited from the environment is already overwritten, so deleting
// them in the ambient scrub would break the call rather than contain a leak.
var ownedHeaders = map[string]bool{"authorization": true, "content-type": true, "accept": true}

// ambientHeaderNames lists the headers openai.NewClient derives from the
// process environment: OPENAI_ORG_ID and OPENAI_PROJECT_ID map to fixed names,
// and OPENAI_CUSTOM_HEADERS -- arbitrary names, and where a gateway token is
// most likely to sit -- is parsed the way the SDK parses it so every name it
// declares is deleted too.
func ambientHeaderNames() []string {
	names := []string{"OpenAI-Organization", "OpenAI-Project"}
	for _, line := range strings.Split(os.Getenv("OPENAI_CUSTOM_HEADERS"), "\n") {
		name, _, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}
		if name = strings.TrimSpace(name); name != "" && !ownedHeaders[strings.ToLower(name)] {
			names = append(names, name)
		}
	}
	return names
}

// limitedBody fails once the body has yielded more than the cap. left is the
// budget still under the cap; one byte past it is read so a reply of exactly
// the cap is told from a longer one, and the overrun is reported on the same
// Read that saw it -- a body whose final chunk carries io.EOF never gets a
// second Read to fail in.
type limitedBody struct {
	rc   io.ReadCloser
	left int64
}

func (b *limitedBody) Read(p []byte) (int, error) {
	if b.left < 0 {
		return 0, ErrResponseTooLarge
	}
	// Compared as b.left < len(p) rather than len(p) > b.left+1: the two select
	// the same slice, but the addition overflows at a cap of math.MaxInt64 --
	// a caller's way of saying "no cap" -- and reslices to a negative bound.
	if b.left < int64(len(p)) {
		p = p[:b.left+1]
	}
	n, err := b.rc.Read(p)
	b.left -= int64(n)
	if b.left < 0 {
		return 0, ErrResponseTooLarge
	}
	return n, err
}

func (b *limitedBody) Close() error { return b.rc.Close() }

// limitResponseBody caps what one reply may hand to the SDK's io.ReadAll. It
// runs as SDK middleware rather than in the transport because the transport's
// gzip is transparent: by the time the response gets here the body is the
// decompressed stream, which is what the allocation is made of.
func limitResponseBody(limit int64) option.Middleware {
	return func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
		res, err := next(req)
		if err != nil || res == nil || res.Body == nil {
			return res, err
		}
		res.Body = &limitedBody{rc: res.Body, left: limit}
		return res, nil
	}
}

// sendChat runs one chat completion. Every request goes through here so the
// opaque error mapping holds for all of them.
func (c *Client) sendChat(ctx context.Context, params openai.ChatCompletionNewParams) (*openai.ChatCompletion, error) {
	completion, err := c.sdk.Chat.Completions.New(ctx, params)
	if err != nil {
		return nil, backendError(err)
	}
	return completion, nil
}

// backendError maps an SDK failure to an error that carries no endpoint URL
// and no key. The SDK's own error text interpolates the request URL, so it is
// never wrapped: only the sentinel, the context cause or the upstream status
// survives, since an error a caller logs or shows a user must not turn a
// configured endpoint into a reachability oracle.
func backendError(err error) error {
	switch {
	case errors.Is(err, ErrResponseTooLarge):
		return fmt.Errorf("ai backend: %w", ErrResponseTooLarge)
	case errors.Is(err, context.Canceled):
		return fmt.Errorf("ai backend: %w", context.Canceled)
	case errors.Is(err, context.DeadlineExceeded):
		return fmt.Errorf("ai backend: %w", context.DeadlineExceeded)
	}
	var apiErr *openai.Error
	if errors.As(err, &apiErr) {
		return fmt.Errorf("ai backend: upstream returned status %d", apiErr.StatusCode)
	}
	return errors.New("ai backend: request failed")
}

// ProbeToolCalling reports whether model can call tools at all. Tool calling
// is a prerequisite, not a nice-to-have: the loop is built on it, so a model
// that only produces prose fails here rather than later on real work.
func (c *Client) ProbeToolCalling(ctx context.Context, model string) error {
	if strings.TrimSpace(model) == "" {
		return errors.New("model must not be empty")
	}
	completion, err := c.sendChat(ctx, openai.ChatCompletionNewParams{
		Model: shared.ChatModel(model),
		Messages: []openai.ChatCompletionMessageParamUnion{
			openai.SystemMessage(probeSystemPrompt),
			openai.UserMessage(probeUserPrompt),
		},
		MaxTokens: openai.Int(probeMaxTokens),
		Tools: []openai.ChatCompletionToolUnionParam{
			openai.ChatCompletionFunctionTool(shared.FunctionDefinitionParam{
				Name:        probeToolName,
				Description: openai.String("Acknowledge a connection probe."),
				Parameters: shared.FunctionParameters{
					"type": "object",
					"properties": map[string]any{
						"message": map[string]any{"type": "string", "description": "Any short string."},
					},
					"required": []string{"message"},
				},
			}),
		},
	})
	if err != nil {
		return err
	}
	// A 200 whose body decodes to JSON null leaves the completion nil with no
	// error, so it is checked before it is read.
	// An empty reply is a backend fault, not a verdict on tool calling; the two
	// are distinct sentinels so a caller does not report "model must support
	// tool calling" for a backend that answered with nothing.
	if completion == nil || len(completion.Choices) == 0 {
		return fmt.Errorf("%w: reply had no choices", ErrEmptyReply)
	}
	choice := completion.Choices[0]
	// A reply cut off at the probe budget says nothing about tool calling: a
	// reasoning model can spend all 256 tokens on chain-of-thought before it
	// emits a call, and reporting that as "no tool calling" blames the model
	// for a budget. Checked first, as the loop checks it.
	if choice.FinishReason == finishReasonLength {
		return fmt.Errorf("ai backend: %w", ErrOutputLimit)
	}
	// Content is ignored on purpose: a model that explains it would call ping
	// has not called it. Nothing from the reply -- content or finish_reason,
	// both endpoint-controlled strings -- reaches the error either, so a caller
	// that renders a failed connection test cannot be fed markup by the
	// endpoint it was testing.
	for _, call := range choice.Message.ToolCalls {
		if call.Function.Name == probeToolName {
			return nil
		}
	}
	return fmt.Errorf("%w: reply contained no call to the probe tool", ErrNoToolCalling)
}