Support ChatGPT OAuth device-code login flow through proxy

internal/credentials/chatgpt.go

@@ -119,6 +119,9 @@ func (c *ChatGPTInjector) Inject(req *http.Request) bool {
 		log.Printf("ERROR chatgpt credential initialization failed")
 		return false
 	}
+	if c.document.tokens.RefreshToken == "" {
+		return false
+	}
 	if c.needsRefreshLocked() {
 		if err := c.refreshLocked(); err != nil {
 			log.Printf("ERROR chatgpt credential refresh failed")
@@ -147,16 +150,16 @@ func (c *ChatGPTInjector) ensureLoadedLocked() error {
 	}
 	c.loaded = true
 
-	if c.config.AuthPath == "" {
+	if c.config.AuthPath == "" && c.config.StatePath == "" {
 		c.loadErr = errors.New("auth source is not configured")
 		return c.loadErr
 	}
-	if c.config.StatePath != "" && samePath(c.config.AuthPath, c.config.StatePath) {
+	if c.config.AuthPath != "" && c.config.StatePath != "" && samePath(c.config.AuthPath, c.config.StatePath) {
 		c.loadErr = errors.New("auth source and state paths must differ")
 		return c.loadErr
 	}
 
-	path := c.config.AuthPath
+	path := ""
 	if c.config.StatePath != "" {
 		if _, err := os.Stat(c.config.StatePath); err == nil {
 			path = c.config.StatePath
@@ -165,7 +168,16 @@ func (c *ChatGPTInjector) ensureLoadedLocked() error {
 			return c.loadErr
 		}
 	}
+	if path == "" {
+		path = c.config.AuthPath
+	}
+	if path == "" {
+		return nil
+	}
+	return c.loadFromFileLocked(path)
+}
 
+func (c *ChatGPTInjector) loadFromFileLocked(path string) error {
 	data, err := readPrivateSecretFile(path)
 	if err != nil {
 		c.loadErr = errors.New("auth file is unavailable or has insecure permissions")
@@ -288,6 +300,66 @@ func (c *ChatGPTInjector) refreshLocked() error {
 	return nil
 }
 
+// AcceptLoginTokens processes a successful login exchange response,
+// persists the real tokens to StatePath, and updates in-memory state.
+func (c *ChatGPTInjector) AcceptLoginTokens(responseBody []byte) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	var raw struct {
+		AccessToken  string `json:"access_token"`
+		RefreshToken string `json:"refresh_token"`
+		IDToken      string `json:"id_token"`
+		ExpiresIn    int64  `json:"expires_in"`
+	}
+	if err := json.Unmarshal(responseBody, &raw); err != nil || raw.AccessToken == "" || raw.RefreshToken == "" {
+		return errors.New("login response is malformed or incomplete")
+	}
+
+	tokens := chatGPTTokens{
+		AccessToken:  raw.AccessToken,
+		RefreshToken: raw.RefreshToken,
+		IDToken:      raw.IDToken,
+	}
+	doc := chatGPTDocument{
+		fields: make(map[string]json.RawMessage),
+		tokens: tokens,
+	}
+	authMode, _ := json.Marshal(chatGPTDefaultAuthMode)
+	doc.fields["auth_mode"] = authMode
+	rawTokens, err := json.Marshal(tokens)
+	if err != nil {
+		return errors.New("marshal login tokens")
+	}
+	doc.fields["tokens"] = rawTokens
+	lastRefresh, _ := json.Marshal(c.config.Now().UTC().Format(time.RFC3339))
+	doc.fields["last_refresh"] = lastRefresh
+
+	if doc.accountID() == "" {
+		return errors.New("login response missing account id")
+	}
+
+	if c.config.StatePath == "" {
+		return errors.New("state path not configured")
+	}
+	data, err := json.Marshal(doc.fields)
+	if err != nil {
+		return errors.New("marshal login state")
+	}
+	if err := atomicWritePrivateSecret(c.config.StatePath, data); err != nil {
+		return errors.New("persist login state")
+	}
+
+	c.document = doc
+	c.expiresAt = jwtExpiry(raw.AccessToken)
+	if c.expiresAt.IsZero() && raw.ExpiresIn > 0 {
+		c.expiresAt = c.config.Now().Add(time.Duration(raw.ExpiresIn) * time.Second)
+	}
+	c.loaded = true
+	c.loadErr = nil
+	return nil
+}
+
 func (d chatGPTDocument) accountID() string {
 	if d.tokens.AccountID != "" {
 		return d.tokens.AccountID

internal/credentials/chatgpt_test.go

@@ -305,8 +305,12 @@ func TestChatGPTTokenVending(t *testing.T) {
 	if IsChatGPTTokenExchange(&http.Request{Method: http.MethodPost, URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/other"}}) {
 		t.Error("wrong path must not be treated as token exchange")
 	}
-	vendor := NewChatGPTTokenVendor()
-	response := vendor.HandleTokenExchange(&http.Request{Method: http.MethodPost, URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}})
+	vendor := NewChatGPTTokenVendor(nil)
+	response := vendor.HandleTokenExchange(&http.Request{
+		Method: http.MethodPost,
+		URL:    &url.URL{Host: "auth.openai.com", Path: "/oauth/token"},
+		Body:   io.NopCloser(strings.NewReader("grant_type=refresh_token")),
+	})
 	if response == nil || response.StatusCode != http.StatusOK {
 		t.Fatal("ChatGPT token exchange was not handled")
 	}
@@ -319,6 +323,125 @@ func TestChatGPTTokenVending(t *testing.T) {
 	}
 }
 
+func TestChatGPTStatePathOnly_NoFileYet(t *testing.T) {
+	dir := t.TempDir()
+	statePath := filepath.Join(dir, "state", "auth.json")
+
+	injector := NewChatGPTInjector("", statePath)
+	if !injector.Available() {
+		t.Fatal("StatePath-only injector should be available even before login completes")
+	}
+
+	req := &http.Request{Header: make(http.Header)}
+	req.Header.Set("Authorization", "Bearer agent-dummy")
+	if injector.Inject(req) {
+		t.Error("Inject should return false when no tokens are loaded yet")
+	}
+	if req.Header.Get("Authorization") != "Bearer agent-dummy" {
+		t.Error("failed injection should not modify the agent header")
+	}
+}
+
+func TestChatGPTStatePathOnly_FileExists(t *testing.T) {
+	dir := t.TempDir()
+	statePath := filepath.Join(dir, "auth.json")
+	access := testJWT(map[string]any{"exp": time.Now().Add(time.Hour).Unix()})
+	writePrivateAuth(t, statePath, testAuthJSON(access, "refresh", "", "account"))
+
+	injector := NewChatGPTInjector("", statePath)
+	if !injector.Available() {
+		t.Fatal("StatePath-only with existing file should be available")
+	}
+
+	req := &http.Request{Header: make(http.Header)}
+	if !injector.Inject(req) {
+		t.Fatal("Inject should succeed when StatePath has valid tokens")
+	}
+	if req.Header.Get("Authorization") != "Bearer "+access {
+		t.Error("access token was not injected")
+	}
+	if req.Header.Get(chatGPTAccountHeader) != "account" {
+		t.Error("account ID was not injected")
+	}
+}
+
+func TestChatGPTAcceptLoginTokens(t *testing.T) {
+	dir := t.TempDir()
+	statePath := filepath.Join(dir, "state", "auth.json")
+	now := time.Unix(1_700_000_000, 0)
+	access := testJWT(map[string]any{
+		"exp":                now.Add(time.Hour).Unix(),
+		"chatgpt_account_id": "logged-in-account",
+	})
+
+	injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{
+		StatePath: statePath,
+		Now:       func() time.Time { return now },
+	})
+
+	if !injector.Available() {
+		t.Fatal("StatePath-only injector should be available")
+	}
+	req := &http.Request{Header: make(http.Header)}
+	if injector.Inject(req) {
+		t.Error("Inject should return false before login")
+	}
+
+	loginResp, _ := json.Marshal(map[string]any{
+		"access_token":  access,
+		"refresh_token": "real-refresh",
+		"id_token":      testJWT(map[string]any{"chatgpt_account_id": "logged-in-account"}),
+		"expires_in":    3600,
+	})
+	if err := injector.AcceptLoginTokens(loginResp); err != nil {
+		t.Fatalf("AcceptLoginTokens failed: %v", err)
+	}
+
+	req2 := &http.Request{Header: make(http.Header)}
+	if !injector.Inject(req2) {
+		t.Fatal("Inject should succeed after AcceptLoginTokens")
+	}
+	if req2.Header.Get("Authorization") != "Bearer "+access {
+		t.Error("access token was not injected after login")
+	}
+	if req2.Header.Get(chatGPTAccountHeader) != "logged-in-account" {
+		t.Error("account ID was not injected after login")
+	}
+
+	info, err := os.Stat(statePath)
+	if err != nil {
+		t.Fatalf("state file was not persisted: %v", err)
+	}
+	if info.Mode().Perm() != 0600 {
+		t.Errorf("state permissions = %o, want 600", info.Mode().Perm())
+	}
+	persisted, _ := os.ReadFile(statePath)
+	if !bytes.Contains(persisted, []byte("real-refresh")) {
+		t.Error("refresh token was not persisted")
+	}
+}
+
+func TestChatGPTAcceptLoginTokens_Malformed(t *testing.T) {
+	injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{
+		StatePath: filepath.Join(t.TempDir(), "auth.json"),
+		Now:       time.Now,
+	})
+
+	cases := map[string][]byte{
+		"not json":           []byte("not-json"),
+		"missing access":     []byte(`{"refresh_token":"r"}`),
+		"missing refresh":    []byte(`{"access_token":"a"}`),
+		"missing account_id": []byte(`{"access_token":"opaque","refresh_token":"r"}`),
+	}
+	for name, body := range cases {
+		t.Run(name, func(t *testing.T) {
+			if err := injector.AcceptLoginTokens(body); err == nil {
+				t.Error("AcceptLoginTokens should reject malformed input")
+			}
+		})
+	}
+}
+
 func TestChatGPTNoCredentialAndNoLogLeakage(t *testing.T) {
 	if NewChatGPTInjector("", "").Available() {
 		t.Error("missing auth path should be unavailable")

internal/credentials/config.go

@@ -99,15 +99,21 @@ func BuildFromConfig(cfg *CredentialConfig) (*Store, *TokenVendor, map[string][]
 	domainMap := make(map[string][]string)
 
 	gcpADCJSON := os.Getenv("GCP_ADC_JSON")
+	chatGPTStatePath := os.Getenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE")
 
 	for _, entry := range cfg.Credentials {
 		value := os.Getenv(entry.EnvVar)
 
-		// For gcloud entries, GCP_ADC_JSON takes precedence over the file path
-		// and allows processing even if GOOGLE_APPLICATION_CREDENTIALS is unset.
-		if entry.InjectorType == "gcloud" && gcpADCJSON == "" && value == "" {
-			continue
-		} else if value == "" && entry.InjectorType != "gcloud" {
+		hasSource := value != ""
+		if !hasSource {
+			switch entry.InjectorType {
+			case "gcloud":
+				hasSource = gcpADCJSON != ""
+			case "chatgpt":
+				hasSource = chatGPTStatePath != ""
+			}
+		}
+		if !hasSource {
 			continue
 		}
 
@@ -143,14 +149,15 @@ func BuildFromConfig(cfg *CredentialConfig) (*Store, *TokenVendor, map[string][]
 			tokenVendor.googleEnabled = true
 			log.Println("Token vendor: ENABLED (returns dummy tokens for oauth2.googleapis.com/token)")
 		case "chatgpt":
-			chatGPTInjector := NewChatGPTInjector(value, os.Getenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE"))
+			chatGPTInjector := NewChatGPTInjector(value, chatGPTStatePath)
 			if !chatGPTInjector.Available() {
-				log.Printf("WARN: %s auth file is not loadable", entry.EnvVar)
+				log.Printf("WARN: chatgpt auth is not loadable")
 				continue
 			}
 			injector = chatGPTInjector
 			tokenVendor = ensureTokenVendor(tokenVendor)
 			tokenVendor.chatGPTEnabled = true
+			tokenVendor.chatGPTInjector = chatGPTInjector
 			log.Println("Token vendor: ENABLED (returns dummy tokens for auth.openai.com/oauth/token)")
 		}
 

internal/credentials/config_test.go

@@ -566,4 +566,67 @@ func TestBuildFromConfig_ChatGPT(t *testing.T) {
 	if req.Header.Get("ChatGPT-Account-ID") != "account" {
 		t.Error("ChatGPT route did not inject account ID")
 	}
+	if tokenVendor.chatGPTInjector == nil {
+		t.Error("ChatGPT injector should be wired to token vendor")
+	}
+}
+
+func TestBuildFromConfig_ChatGPT_StateFileOnly(t *testing.T) {
+	dir := t.TempDir()
+	statePath := filepath.Join(dir, "state", "auth.json")
+	t.Setenv("CHATGPT_AUTH_FILE", "")
+	t.Setenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE", statePath)
+
+	cfg := &CredentialConfig{Credentials: []CredentialEntry{
+		{
+			EnvVar:       "CHATGPT_AUTH_FILE",
+			InjectorType: "chatgpt",
+			Params:       map[string]string{"path_prefix": "/backend-api/codex"},
+			Domains:      []string{"chatgpt.com"},
+		},
+	}}
+	store, tokenVendor, domainMap := BuildFromConfig(cfg)
+	if tokenVendor == nil {
+		t.Fatal("ChatGPT token vendor should be enabled with state file only")
+	}
+	if !tokenVendor.chatGPTEnabled {
+		t.Error("chatGPTEnabled should be true")
+	}
+	if tokenVendor.chatGPTInjector == nil {
+		t.Error("chatGPTInjector should be wired to token vendor")
+	}
+	if got := domainMap["CHATGPT_AUTH_FILE"]; len(got) < 2 || got[1] != "auth.openai.com" {
+		t.Errorf("domain map = %v, want API and OAuth domains", got)
+	}
+
+	req := &http.Request{
+		Method: http.MethodPost,
+		URL:    &url.URL{Host: "chatgpt.com", Path: "/backend-api/codex/responses"},
+		Header: make(http.Header),
+	}
+	matched, injected := store.InjectCredentials(req)
+	if !matched {
+		t.Fatal("route should match even before login")
+	}
+	if injected {
+		t.Error("should not inject before login (no tokens yet)")
+	}
+}
+
+func TestBuildFromConfig_ChatGPT_NeitherSet(t *testing.T) {
+	t.Setenv("CHATGPT_AUTH_FILE", "")
+	t.Setenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE", "")
+
+	cfg := &CredentialConfig{Credentials: []CredentialEntry{
+		{
+			EnvVar:       "CHATGPT_AUTH_FILE",
+			InjectorType: "chatgpt",
+			Params:       map[string]string{"path_prefix": "/backend-api/codex"},
+			Domains:      []string{"chatgpt.com"},
+		},
+	}}
+	_, tokenVendor, _ := BuildFromConfig(cfg)
+	if tokenVendor != nil && tokenVendor.chatGPTEnabled {
+		t.Error("neither env var set should not enable ChatGPT token vending")
+	}
 }