fix: proxy Codex models through API key upstreams

backend/internal/handler/openai_codex_models_handler.go

@@ -15,10 +15,9 @@ import (
 // Codex CLI and the Codex desktop app refresh their model picker from
 // GET {base_url}/models?client_version=... (custom provider mode) or
 // GET /backend-api/codex/models (chatgpt_base_url mode). Both routes land
-// here. The manifest is proxied verbatim from the ChatGPT backend with a
-// schedulable OAuth account's credentials, so clients pointed at the gateway
-// see the account's real, always-current model entitlements instead of a
-// frozen local cache.
+// here. The manifest is proxied verbatim from the selected account's ChatGPT
+// backend or custom API key upstream, so clients pointed at the gateway see an
+// always-current manifest instead of a frozen local cache.
 func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
 	apiKey, ok := middleware2.GetAPIKeyFromContext(c)
 	if !ok || apiKey.Group == nil {

backend/internal/service/openai_codex_models_service.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"net/url"
@@ -26,8 +27,8 @@ type CodexModelsManifest struct {
 	NotModified bool
 }
 
-// FetchCodexModelsManifest fetches the live Codex models manifest from the
-// ChatGPT backend using the account's OAuth credentials.
+// FetchCodexModelsManifest fetches the live Codex models manifest from either
+// the ChatGPT backend for OAuth accounts or a custom upstream for API key accounts.
 //
 // The response body is passed through verbatim: the manifest schema evolves
 // with Codex client releases, and interpreting it here would force the gateway
@@ -41,24 +42,61 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
 	if err != nil {
 		return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_CREDENTIALS_FAILED", "resolve credential account: %v", err)
 	}
-	accessToken := credAccount.GetOpenAIAccessToken()
-	if accessToken == "" {
-		return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token")
-	}
 
 	clientVersion = strings.TrimSpace(clientVersion)
 	if clientVersion == "" {
 		clientVersion = openAICodexProbeVersion
 	}
-	requestURL := chatgptCodexModelsURL + "?client_version=" + url.QueryEscape(clientVersion)
+
+	requestEndpoint := chatgptCodexModelsURL
+	authToken := ""
+	useAPIKeyUpstream := false
+	appendModelsPath := false
+	switch {
+	case credAccount.IsOpenAIOAuth():
+		authToken = strings.TrimSpace(credAccount.GetOpenAIAccessToken())
+		if authToken == "" {
+			return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token")
+		}
+	case credAccount.IsOpenAIApiKey():
+		baseURL := strings.TrimSpace(credAccount.GetCredential("base_url"))
+		if baseURL == "" || isOfficialOpenAIModelsBaseURL(baseURL) {
+			return nil, infraerrors.New(
+				http.StatusBadGateway,
+				"OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_UNSUPPORTED",
+				"Codex models manifest requires a custom API key upstream base URL",
+			)
+		}
+		authToken = strings.TrimSpace(credAccount.GetOpenAIApiKey())
+		if authToken == "" {
+			return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_MISSING", "account has no API key for the Codex models upstream")
+		}
+		normalizedBaseURL, validateErr := s.validateUpstreamBaseURL(baseURL)
+		if validateErr != nil {
+			return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID", "invalid Codex models upstream base URL: %v", validateErr)
+		}
+		requestEndpoint = normalizedBaseURL
+		useAPIKeyUpstream = true
+		appendModelsPath = true
+	default:
+		return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_ACCOUNT_TYPE_UNSUPPORTED", "account type %q cannot fetch the Codex models manifest", credAccount.Type)
+	}
+
+	requestURL, err := buildCodexModelsManifestURL(requestEndpoint, appendModelsPath, clientVersion)
+	if err != nil {
+		if useAPIKeyUpstream {
+			return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID", "invalid Codex models upstream base URL: %v", err)
+		}
+		return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "parse codex models request URL: %v", err)
+	}
 
 	reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
 	defer cancel()
-	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, requestURL, nil)
+	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, requestURL.String(), nil)
 	if err != nil {
 		return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_REQUEST_FAILED", "create codex models request: %v", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+accessToken)
+	req.Header.Set("Authorization", "Bearer "+authToken)
 	req.Header.Set("Accept", "application/json")
 	req.Header.Set("Originator", "codex_cli_rs")
 	req.Header.Set("Version", clientVersion)
@@ -66,22 +104,35 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
 	if ifNoneMatch = strings.TrimSpace(ifNoneMatch); ifNoneMatch != "" {
 		req.Header.Set("If-None-Match", ifNoneMatch)
 	}
-	setOpenAIChatGPTAccountHeaders(req.Header, credAccount)
+	if useAPIKeyUpstream {
+		credAccount.ApplyHeaderOverrides(req.Header)
+	} else {
+		setOpenAIChatGPTAccountHeaders(req.Header, credAccount)
+	}
 
 	proxyURL := ""
 	if account.ProxyID != nil && account.Proxy != nil {
 		proxyURL = account.Proxy.URL()
 	}
-	client, err := httpclient.GetClient(httpclient.Options{
-		ProxyURL:              proxyURL,
-		Timeout:               15 * time.Second,
-		ResponseHeaderTimeout: 10 * time.Second,
-	})
-	if err != nil {
-		return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_PROXY_INVALID", "invalid proxy configuration: %v", err)
-	}
 
-	resp, err := client.Do(req)
+	var resp *http.Response
+	if useAPIKeyUpstream {
+		if s.httpUpstream == nil {
+			return nil, infraerrors.New(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_UPSTREAM_NOT_CONFIGURED", "Codex models upstream HTTP client is not configured")
+		}
+		req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
+		resp, err = s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
+	} else {
+		client, clientErr := httpclient.GetClient(httpclient.Options{
+			ProxyURL:              proxyURL,
+			Timeout:               15 * time.Second,
+			ResponseHeaderTimeout: 10 * time.Second,
+		})
+		if clientErr != nil {
+			return nil, infraerrors.Newf(http.StatusInternalServerError, "OPENAI_CODEX_MODELS_PROXY_INVALID", "invalid proxy configuration: %v", clientErr)
+		}
+		resp, err = client.Do(req)
+	}
 	if err != nil {
 		return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest request failed: %v", err)
 	}
@@ -105,3 +156,35 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
 	}
 	return &CodexModelsManifest{Body: body, ETag: resp.Header.Get("ETag")}, nil
 }
+
+func isOfficialOpenAIModelsBaseURL(raw string) bool {
+	parsed, err := url.Parse(strings.TrimSpace(raw))
+	if err != nil {
+		return false
+	}
+	hostname := strings.TrimSuffix(parsed.Hostname(), ".")
+	return strings.EqualFold(hostname, "api.openai.com")
+}
+
+func buildCodexModelsManifestURL(endpoint string, appendModelsPath bool, clientVersion string) (*url.URL, error) {
+	requestURL, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, err
+	}
+	if requestURL.Fragment != "" {
+		return nil, fmt.Errorf("URL fragments are not supported")
+	}
+
+	query := requestURL.Query()
+	requestURL.RawQuery = ""
+	requestURL.ForceQuery = false
+	if appendModelsPath {
+		requestURL, err = url.Parse(buildOpenAIModelsURL(requestURL.String()))
+		if err != nil {
+			return nil, err
+		}
+	}
+	query.Set("client_version", clientVersion)
+	requestURL.RawQuery = query.Encode()
+	return requestURL, nil
+}

backend/internal/service/openai_codex_models_service_test.go

@@ -2,11 +2,52 @@ package service
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
 )
 
+type codexModelsHTTPUpstreamStub struct {
+	do func(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error)
+}
+
+func (s *codexModelsHTTPUpstreamStub) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
+	return s.do(req, proxyURL, accountID, accountConcurrency)
+}
+
+func (s *codexModelsHTTPUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
+	return s.Do(req, proxyURL, accountID, accountConcurrency)
+}
+
+func newCodexModelsAPIKeyTestService(upstream HTTPUpstream) *OpenAIGatewayService {
+	return &OpenAIGatewayService{
+		cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{
+			Enabled: false,
+		}}},
+		httpUpstream: upstream,
+	}
+}
+
+func newCodexModelsAPIKeyTestAccount(baseURL string) *Account {
+	credentials := map[string]any{"api_key": "sk-upstream"}
+	if baseURL != "" {
+		credentials["base_url"] = baseURL
+	}
+	return &Account{
+		ID:          2,
+		Platform:    PlatformOpenAI,
+		Type:        AccountTypeAPIKey,
+		Credentials: credentials,
+		Concurrency: 3,
+	}
+}
+
 func newCodexModelsTestAccount() *Account {
 	return &Account{
 		ID:       1,
@@ -136,3 +177,219 @@ func TestFetchCodexModelsManifestMissingToken(t *testing.T) {
 		t.Fatal("expected error for missing access token, got nil")
 	}
 }
+
+func TestFetchCodexModelsManifestAPIKeyCustomUpstream(t *testing.T) {
+	manifestBody := `{"models":[{"slug":"gpt-5.6"}]}`
+	var gotRequest *http.Request
+	var gotProxyURL string
+	var gotAccountID int64
+	var gotConcurrency int
+	upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
+		gotRequest = req
+		gotProxyURL = proxyURL
+		gotAccountID = accountID
+		gotConcurrency = accountConcurrency
+		header := make(http.Header)
+		header.Set("ETag", `W/"api-key-manifest"`)
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     header,
+			Body:       io.NopCloser(strings.NewReader(manifestBody)),
+		}, nil
+	}}
+
+	s := newCodexModelsAPIKeyTestService(upstream)
+	manifest, err := s.FetchCodexModelsManifest(
+		context.Background(),
+		newCodexModelsAPIKeyTestAccount("https://upstream.example/v1"),
+		"0.144.0",
+		"",
+	)
+	if err != nil {
+		t.Fatalf("FetchCodexModelsManifest returned error: %v", err)
+	}
+
+	if gotRequest == nil {
+		t.Fatal("expected request to custom API key upstream")
+	}
+	if gotRequest.Method != http.MethodGet {
+		t.Errorf("method: got %q", gotRequest.Method)
+	}
+	if gotRequest.URL.String() != "https://upstream.example/v1/models?client_version=0.144.0" {
+		t.Errorf("request URL: got %q", gotRequest.URL.String())
+	}
+	if gotRequest.Header.Get("Authorization") != "Bearer sk-upstream" {
+		t.Errorf("authorization header: got %q", gotRequest.Header.Get("Authorization"))
+	}
+	if gotRequest.Header.Get("Originator") != "codex_cli_rs" {
+		t.Errorf("originator header: got %q", gotRequest.Header.Get("Originator"))
+	}
+	if gotRequest.Header.Get("Version") != "0.144.0" {
+		t.Errorf("version header: got %q", gotRequest.Header.Get("Version"))
+	}
+	if gotRequest.Header.Get("User-Agent") != codexCLIUserAgent {
+		t.Errorf("user-agent header: got %q", gotRequest.Header.Get("User-Agent"))
+	}
+	if gotRequest.Header.Get("chatgpt-account-id") != "" {
+		t.Errorf("chatgpt-account-id must not be sent to API key upstream: got %q", gotRequest.Header.Get("chatgpt-account-id"))
+	}
+	if gotProxyURL != "" || gotAccountID != 2 || gotConcurrency != 3 {
+		t.Errorf("upstream routing metadata: proxy=%q account_id=%d concurrency=%d", gotProxyURL, gotAccountID, gotConcurrency)
+	}
+	if string(manifest.Body) != manifestBody {
+		t.Errorf("body not passed through verbatim: got %q", manifest.Body)
+	}
+	if manifest.ETag != `W/"api-key-manifest"` {
+		t.Errorf("etag not passed through: got %q", manifest.ETag)
+	}
+}
+
+func TestFetchCodexModelsManifestAPIKeyNotModified(t *testing.T) {
+	var gotIfNoneMatch string
+	upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
+		gotIfNoneMatch = req.Header.Get("If-None-Match")
+		header := make(http.Header)
+		header.Set("ETag", `W/"api-key-manifest"`)
+		return &http.Response{
+			StatusCode: http.StatusNotModified,
+			Header:     header,
+			Body:       http.NoBody,
+		}, nil
+	}}
+
+	s := newCodexModelsAPIKeyTestService(upstream)
+	manifest, err := s.FetchCodexModelsManifest(
+		context.Background(),
+		newCodexModelsAPIKeyTestAccount("https://upstream.example"),
+		"0.144.0",
+		`W/"api-key-manifest"`,
+	)
+	if err != nil {
+		t.Fatalf("FetchCodexModelsManifest returned error: %v", err)
+	}
+	if !manifest.NotModified {
+		t.Error("expected NotModified to be true")
+	}
+	if manifest.ETag != `W/"api-key-manifest"` {
+		t.Errorf("etag not passed through: got %q", manifest.ETag)
+	}
+	if gotIfNoneMatch != `W/"api-key-manifest"` {
+		t.Errorf("if-none-match header: got %q", gotIfNoneMatch)
+	}
+}
+
+func TestFetchCodexModelsManifestAPIKeyPreservesBaseURLQuery(t *testing.T) {
+	var gotURL string
+	upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
+		gotURL = req.URL.String()
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     make(http.Header),
+			Body:       io.NopCloser(strings.NewReader(`{"models":[]}`)),
+		}, nil
+	}}
+
+	s := newCodexModelsAPIKeyTestService(upstream)
+	_, err := s.FetchCodexModelsManifest(
+		context.Background(),
+		newCodexModelsAPIKeyTestAccount("https://upstream.example/v1?tenant=acme"),
+		"0.144.0",
+		"",
+	)
+	if err != nil {
+		t.Fatalf("FetchCodexModelsManifest returned error: %v", err)
+	}
+	if gotURL != "https://upstream.example/v1/models?client_version=0.144.0&tenant=acme" {
+		t.Errorf("request URL: got %q", gotURL)
+	}
+}
+
+func TestFetchCodexModelsManifestAPIKeyRejectsBaseURLFragment(t *testing.T) {
+	called := false
+	upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
+		called = true
+		return &http.Response{
+			StatusCode: http.StatusOK,
+			Header:     make(http.Header),
+			Body:       io.NopCloser(strings.NewReader(`{"models":[]}`)),
+		}, nil
+	}}
+
+	s := newCodexModelsAPIKeyTestService(upstream)
+	_, err := s.FetchCodexModelsManifest(
+		context.Background(),
+		newCodexModelsAPIKeyTestAccount("https://upstream.example/v1#models"),
+		"0.144.0",
+		"",
+	)
+	if err == nil {
+		t.Fatal("expected invalid upstream base URL error, got nil")
+	}
+	if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_INVALID" {
+		t.Errorf("error reason: got %q", infraerrors.Reason(err))
+	}
+	if called {
+		t.Fatal("fragment-bearing base URL must be rejected before the upstream request")
+	}
+}
+
+func TestFetchCodexModelsManifestAPIKeyUpstreamError(t *testing.T) {
+	upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
+		return &http.Response{
+			StatusCode: http.StatusTooManyRequests,
+			Status:     "429 Too Many Requests",
+			Header:     make(http.Header),
+			Body:       io.NopCloser(strings.NewReader(`{"error":"rate limited"}`)),
+		}, nil
+	}}
+
+	s := newCodexModelsAPIKeyTestService(upstream)
+	_, err := s.FetchCodexModelsManifest(
+		context.Background(),
+		newCodexModelsAPIKeyTestAccount("https://upstream.example"),
+		"0.144.0",
+		"",
+	)
+	if err == nil {
+		t.Fatal("expected error for upstream 429, got nil")
+	}
+	if infraerrors.Code(err) != http.StatusBadGateway {
+		t.Errorf("error status: got %d, want %d", infraerrors.Code(err), http.StatusBadGateway)
+	}
+	if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_UPSTREAM_FAILED" {
+		t.Errorf("error reason: got %q", infraerrors.Reason(err))
+	}
+}
+
+func TestFetchCodexModelsManifestAPIKeyRejectsOfficialOpenAIBaseURL(t *testing.T) {
+	tests := []struct {
+		name    string
+		baseURL string
+	}{
+		{name: "missing base URL"},
+		{name: "official host", baseURL: "https://api.openai.com"},
+		{name: "official versioned URL", baseURL: "https://API.OPENAI.COM:443/v1/"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := newCodexModelsAPIKeyTestService(&codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
+				t.Fatal("official OpenAI API key must not be used as a Codex manifest upstream")
+				return nil, nil
+			}})
+
+			_, err := s.FetchCodexModelsManifest(
+				context.Background(),
+				newCodexModelsAPIKeyTestAccount(tt.baseURL),
+				"0.144.0",
+				"",
+			)
+			if err == nil {
+				t.Fatal("expected unsupported API key upstream error, got nil")
+			}
+			if infraerrors.Reason(err) != "OPENAI_CODEX_MODELS_API_KEY_UPSTREAM_UNSUPPORTED" {
+				t.Errorf("error reason: got %q", infraerrors.Reason(err))
+			}
+		})
+	}
+}