Merge pull request #62 from DingTalk-Real-AI/feat-login-optimization

feat(auth): persist OAuth credentials for reliable token refresh

internal/app/auth_command.go

@@ -185,9 +185,19 @@ func newAuthLogoutCommand() *cobra.Command {
 			defer cancel()
 			_ = authpkg.RevokeTokenRemote(revokeCtx)
 
+			// Load token data to get associated clientId before deletion
+			var storedClientID string
+			if tokenData, err := authpkg.LoadTokenData(configDir); err == nil && tokenData != nil {
+				storedClientID = tokenData.ClientID
+			}
+
 			if err := authpkg.DeleteTokenData(configDir); err != nil {
 				return apperrors.NewInternal(fmt.Sprintf("failed to clear token data: %v", err))
 			}
+			// Clean up associated client secret from keychain
+			if storedClientID != "" {
+				_ = authpkg.DeleteClientSecret(storedClientID)
+			}
 			// Clean up app credentials (app.json + keychain secret)
 			_ = authpkg.DeleteAppConfig(configDir)
 			_ = os.Remove(filepath.Join(configDir, "mcp_url"))

internal/auth/endpoints.go

@@ -140,6 +140,27 @@ func GetRevokeTokenURL() string {
 	return "" // Direct mode doesn't have revoke endpoint
 }
 
+// resolveCredentialSource determines the source of the current credentials.
+// Returns one of: "flag", "env", "app", "default".
+// This is used to track where credentials came from for token refresh.
+func resolveCredentialSource() string {
+	clientMu.RLock()
+	hasRuntimeOverride := runtimeClientID != "" || runtimeClientSecret != ""
+	clientMu.RUnlock()
+
+	if hasRuntimeOverride {
+		return "flag"
+	}
+	// Check if loaded from app config
+	if id, _ := ResolveAppCredentials(getDefaultConfigDir()); id != "" {
+		return "app"
+	}
+	if os.Getenv("DWS_CLIENT_ID") != "" || os.Getenv("DWS_CLIENT_SECRET") != "" {
+		return "env"
+	}
+	return "default"
+}
+
 // SetClientID allows runtime override of the client ID (e.g., from CLI flags).
 func SetClientID(id string) {
 	clientMu.Lock()

internal/auth/keychain_store.go

@@ -105,3 +105,45 @@ func EnsureMigration(configDir string, logger *slog.Logger) {
 func IsMigrationDone() bool {
 	return migrationDone
 }
+
+// Client credential storage functions.
+// These store the clientSecret associated with a specific clientId,
+// allowing token refresh to work even if environment variables change.
+
+const clientSecretPrefix = "client-secret:"
+
+// SaveClientSecret stores the client secret for a specific client ID.
+// This is called during login to snapshot the credentials used.
+func SaveClientSecret(clientID, clientSecret string) error {
+	if clientID == "" || clientSecret == "" {
+		return nil // Nothing to save
+	}
+	account := clientSecretPrefix + clientID
+	if err := keychain.Set(keychain.Service, account, clientSecret); err != nil {
+		return fmt.Errorf("save client secret: %w", err)
+	}
+	return nil
+}
+
+// LoadClientSecret retrieves the stored client secret for a specific client ID.
+// Returns empty string if not found.
+func LoadClientSecret(clientID string) string {
+	if clientID == "" {
+		return ""
+	}
+	account := clientSecretPrefix + clientID
+	secret, err := keychain.Get(keychain.Service, account)
+	if err != nil {
+		return ""
+	}
+	return secret
+}
+
+// DeleteClientSecret removes the stored client secret for a specific client ID.
+func DeleteClientSecret(clientID string) error {
+	if clientID == "" {
+		return nil
+	}
+	account := clientSecretPrefix + clientID
+	return keychain.Remove(keychain.Service, account)
+}

internal/auth/oauth_helpers.go

@@ -32,9 +32,11 @@ func (p *OAuthProvider) exchangeCode(ctx context.Context, code string) (*TokenDa
 		return p.exchangeCodeViaMCP(ctx, code)
 	}
 	// Direct mode with client secret
+	clientID := ClientID()
+	clientSecret := ClientSecret()
 	body := map[string]string{
-		"clientId":     ClientID(),
-		"clientSecret": ClientSecret(),
+		"clientId":     clientID,
+		"clientSecret": clientSecret,
 		"code":         code,
 		"grantType":    "authorization_code",
 	}
@@ -42,15 +44,28 @@ func (p *OAuthProvider) exchangeCode(ctx context.Context, code string) (*TokenDa
 	if err != nil {
 		return nil, err
 	}
-	return p.parseTokenResponse(resp)
+	data, err := p.parseTokenResponse(resp)
+	if err != nil {
+		return nil, err
+	}
+	// Snapshot credentials used for this token (for refresh)
+	data.ClientID = clientID
+	data.Source = resolveCredentialSource()
+	// Save clientSecret for future refresh (even if env changes)
+	if err := SaveClientSecret(clientID, clientSecret); err != nil {
+		// Log warning but don't fail login
+		fmt.Fprintf(p.Output, "Warning: failed to save client secret: %v\n", err)
+	}
+	return data, nil
 }
 
 // exchangeCodeViaMCP exchanges auth code for token via MCP proxy.
 // This is used when client secret is not available (server-side secret management).
 func (p *OAuthProvider) exchangeCodeViaMCP(ctx context.Context, code string) (*TokenData, error) {
+	clientID := ClientID()
 	url := GetMCPBaseURL() + MCPOAuthTokenPath
 	body := map[string]string{
-		"clientId":  ClientID(),
+		"clientId":  clientID,
 		"authCode":  code,
 		"grantType": "authorization_code",
 	}
@@ -58,18 +73,43 @@ func (p *OAuthProvider) exchangeCodeViaMCP(ctx context.Context, code string) (*T
 	if err != nil {
 		return nil, err
 	}
-	return p.parseMCPTokenResponse(resp)
+	data, err := p.parseMCPTokenResponse(resp)
+	if err != nil {
+		return nil, err
+	}
+	// Snapshot credentials used for this token (for refresh)
+	data.ClientID = clientID
+	data.Source = "mcp"
+	// MCP mode doesn't need to save clientSecret (server-side managed)
+	return data, nil
 }
 
 func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *TokenData) (*TokenData, error) {
-	// Use MCP mode if clientID is from MCP server
-	if IsClientIDFromMCP() {
+	// Use stored Source to determine refresh path (not current runtime state)
+	// This ensures refresh works even if environment variables changed since login
+	if data.Source == "mcp" {
 		return p.refreshViaMCP(ctx, data)
 	}
-	// Direct mode with client secret
+
+	// Direct mode: use stored clientId and load saved clientSecret
+	clientID := data.ClientID
+	if clientID == "" {
+		// Fallback for legacy tokens without stored clientId
+		clientID = ClientID()
+	}
+	clientSecret := LoadClientSecret(clientID)
+	if clientSecret == "" {
+		// Fallback: try current environment
+		clientSecret = ClientSecret()
+	}
+
+	if clientID == "" || clientSecret == "" {
+		return nil, fmt.Errorf("无法刷新 token: 缺少 clientId 或 clientSecret，请重新登录")
+	}
+
 	body := map[string]string{
-		"clientId":     ClientID(),
-		"clientSecret": ClientSecret(),
+		"clientId":     clientID,
+		"clientSecret": clientSecret,
 		"refreshToken": data.RefreshToken,
 		"grantType":    "refresh_token",
 	}
@@ -81,6 +121,9 @@ func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *Token
 	if err != nil {
 		return nil, err
 	}
+	// Preserve original credentials info
+	updated.ClientID = data.ClientID
+	updated.Source = data.Source
 	updated.PersistentCode = data.PersistentCode
 	updated.CorpID = data.CorpID
 	updated.UserID = data.UserID
@@ -95,9 +138,20 @@ func (p *OAuthProvider) refreshWithRefreshToken(ctx context.Context, data *Token
 
 // refreshViaMCP refreshes token via MCP proxy.
 func (p *OAuthProvider) refreshViaMCP(ctx context.Context, data *TokenData) (*TokenData, error) {
+	// Use stored clientId from token data
+	clientID := data.ClientID
+	if clientID == "" {
+		// Fallback for legacy tokens
+		clientID = ClientID()
+	}
+
+	if clientID == "" {
+		return nil, fmt.Errorf("无法刷新 token: 缺少 clientId，请重新登录")
+	}
+
 	url := GetMCPBaseURL() + MCPRefreshTokenPath
 	body := map[string]string{
-		"clientId":     ClientID(),
+		"clientId":     clientID,
 		"refreshToken": data.RefreshToken,
 		"grantType":    "refresh_token",
 	}
@@ -109,6 +163,9 @@ func (p *OAuthProvider) refreshViaMCP(ctx context.Context, data *TokenData) (*To
 	if err != nil {
 		return nil, err
 	}
+	// Preserve original credentials info
+	updated.ClientID = data.ClientID
+	updated.Source = data.Source
 	updated.PersistentCode = data.PersistentCode
 	updated.CorpID = data.CorpID
 	updated.UserID = data.UserID