easyai-ai-gateway/apps/api/internal/runner/proxy.go

package runner

import (
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/netproxy"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type httpClientCache struct {
	none   *http.Client
	global *http.Client
	mu     sync.Mutex
	custom map[string]*http.Client
}

const (
	providerHTTPMaxIdleConnections        = 2048
	providerHTTPMaxIdleConnectionsPerHost = 1024
)

func newHTTPClientCache() *httpClientCache {
	return &httpClientCache{
		none:   newHTTPClient(nil),
		global: newHTTPClient(http.ProxyFromEnvironment),
		custom: map[string]*http.Client{},
	}
}

func (s *Service) httpClientForCandidate(candidate store.RuntimeModelCandidate, simulated bool, kind string) (*http.Client, error) {
	var client *http.Client
	if simulated {
		client = s.httpClients.none
		return providerHTTPClientForKind(client, kind), nil
	}
	// Some Worker sites have direct provider egress but are not authorized to
	// use a platform's shared proxy. Keep this override explicitly scoped to
	// platform UUIDs so one site's routing exception cannot bypass proxies for
	// unrelated providers or platforms.
	if platformIDListed(s.cfg.PlatformProxyBypassIDs, candidate.PlatformID) {
		client = s.httpClients.none
		return providerHTTPClientForKind(client, kind), nil
	}
	config, err := netproxy.Normalize(netproxy.FromPlatformConfig(candidate.PlatformConfig))
	if err != nil {
		return nil, &clients.ClientError{Code: "invalid_proxy", Message: err.Error(), Retryable: false}
	}
	switch config.Mode {
	case netproxy.ModeGlobal:
		if strings.TrimSpace(s.cfg.GlobalHTTPProxy) != "" {
			client, err = s.httpClients.customClient(s.cfg.GlobalHTTPProxy)
			if err != nil {
				return nil, err
			}
			break
		}
		client = s.httpClients.global
	case netproxy.ModeCustom:
		client, err = s.httpClients.customClient(config.HTTPProxy)
		if err != nil {
			return nil, err
		}
	default:
		client = s.httpClients.none
	}
	return providerHTTPClientForKind(client, kind), nil
}

func platformIDListed(raw string, platformID string) bool {
	platformID = strings.TrimSpace(platformID)
	if platformID == "" {
		return false
	}
	for value := range strings.SplitSeq(raw, ",") {
		if strings.TrimSpace(value) == platformID {
			return true
		}
	}
	return false
}

func (c *httpClientCache) customClient(rawProxy string) (*http.Client, error) {
	normalized, proxyURL, err := netproxy.ParseHTTPProxy(rawProxy)
	if err != nil {
		return nil, &clients.ClientError{Code: "invalid_proxy", Message: err.Error(), Retryable: false}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if client := c.custom[normalized]; client != nil {
		return client, nil
	}
	client := newHTTPClient(http.ProxyURL(proxyURL))
	c.custom[normalized] = client
	return client, nil
}

func newHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.Proxy = proxy
	// Media workers may keep hundreds of slow upstream tasks in flight and
	// poll the same provider repeatedly. The standard library only retains two
	// idle connections per host, which turns a high-concurrency polling load
	// into avoidable TCP/TLS handshakes.
	transport.MaxIdleConns = providerHTTPMaxIdleConnections
	transport.MaxIdleConnsPerHost = providerHTTPMaxIdleConnectionsPerHost
	return &http.Client{
		Timeout:   clients.ProviderRequestTimeout(""),
		Transport: transport,
	}
}

func providerHTTPClientForKind(base *http.Client, kind string) *http.Client {
	if base == nil || base.Timeout == clients.ProviderRequestTimeout(kind) {
		return base
	}
	client := *base
	client.Timeout = clients.ProviderRequestTimeout(kind)
	return &client
}