anthropic-proxy/internal/proxy/upstream.go

package proxy

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/google/uuid"
	"github.com/rs/zerolog/log"

	"github.com/fujin/anthropic-proxy/internal/auth"
	"github.com/fujin/anthropic-proxy/internal/logging"
	"github.com/fujin/anthropic-proxy/internal/transport"
	"github.com/fujin/anthropic-proxy/internal/version"
)

const messagesURL = "https://api.anthropic.com/v1/messages?beta=true"

type UpstreamClient struct {
	client    http.Client
	sessionID string
	profile   *SniffedProfile
}

func NewUpstreamClient(profile *SniffedProfile) *UpstreamClient {
	return &UpstreamClient{
		client: http.Client{
			Timeout:   0,
			Transport: transport.NewUTLS(),
		},
		sessionID: uuid.New().String(),
		profile:   profile,
	}
}

func (u *UpstreamClient) version() string {
	if u.profile != nil && u.profile.Version != "" {
		return u.profile.Version
	}
	return version.ClaudeCodeFallback
}

// applyHeaders replays sniffed headers, substituting auth + per-request IDs + accept.
func (u *UpstreamClient) applyHeaders(req *http.Request, token string, streaming bool) {
	if u.profile != nil {
		for _, h := range u.profile.Headers {
			req.Header.Set(h[0], h[1])
		}
	}

	req.Header.Del("Authorization")
	req.Header.Del("x-api-key")
	if strings.HasPrefix(token, "sk-ant-oat") {
		req.Header.Set("Authorization", "Bearer "+token)
		// OAuth tokens require this beta flag — without it the API rejects with 401
		existing := req.Header.Get("anthropic-beta")
		if !strings.Contains(existing, "oauth-2025-04-20") {
			if existing == "" {
				req.Header.Set("anthropic-beta", "oauth-2025-04-20")
			} else {
				req.Header.Set("anthropic-beta", existing+",oauth-2025-04-20")
			}
		}
	} else {
		req.Header.Set("x-api-key", token)
	}

	req.Header.Set("X-Claude-Code-Session-Id", u.sessionID)
	req.Header.Set("x-client-request-id", uuid.New().String())

	if streaming {
		req.Header.Set("Accept", "text/event-stream")
	} else {
		req.Header.Set("Accept", "application/json")
	}
	req.Header.Set("Accept-Encoding", "identity")
}

func (u *UpstreamClient) Execute(ctx context.Context, cred *auth.Credential, body []byte) ([]byte, http.Header, int, error) {
	body = injectBillingHeader(body, u.version())

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, messagesURL, bytes.NewReader(body))
	if err != nil {
		return nil, nil, 0, fmt.Errorf("build upstream request: %w", err)
	}
	u.applyHeaders(req, cred.Token(), false)

	log.Debug().
		Str("url", messagesURL).
		Str("upstream_headers", logging.RedactHeaders(req.Header)).
		Int("body_size", len(body)).
		Msg("upstream request")

	resp, err := u.client.Do(req)
	if err != nil {
		return nil, nil, 0, fmt.Errorf("upstream request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, nil, resp.StatusCode, fmt.Errorf("read upstream response: %w", err)
	}
	log.Debug().
		Int("status", resp.StatusCode).
		Str("response_headers", logging.RedactHeaders(resp.Header)).
		Int("response_size", len(respBody)).
		Msg("upstream response")

	return respBody, resp.Header, resp.StatusCode, nil
}

func (u *UpstreamClient) ExecuteStream(ctx context.Context, cred *auth.Credential, body []byte) (*http.Response, error) {
	body = injectBillingHeader(body, u.version())

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, messagesURL, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build upstream stream request: %w", err)
	}
	u.applyHeaders(req, cred.Token(), true)

	log.Debug().
		Str("url", messagesURL).
		Str("upstream_headers", logging.RedactHeaders(req.Header)).
		Int("body_size", len(body)).
		Msg("upstream stream request")

	resp, err := u.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("upstream stream request: %w", err)
	}

	log.Debug().
		Int("status", resp.StatusCode).
		Str("response_headers", logging.RedactHeaders(resp.Header)).
		Msg("upstream stream response")

	return resp, nil
}