Skip to content

Commit 9ec2326

Browse files
AllyWAllyW
authored
{mcp} Improve MCP Proxy OAuth Token Refresh Reliability and Profile Reuse (#1282)
* fix-user-oauth * add * add test
1 parent 699bb98 commit 9ec2326

File tree

6 files changed

+871
-87
lines changed

6 files changed

+871
-87
lines changed

‎mcpproxy/mcp_profile.go‎

Lines changed: 116 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -64,13 +64,21 @@ func saveMcpProfile(profile *McpProfile) error {
6464
return fmt.Errorf("failed to create config directory %q: %w", dir, err)
6565
}
6666

67+
// log.Printf("saveMcpProfile: Before marshaling - RefreshToken length=%d, RefreshToken empty=%v",
68+
// len(profile.MCPOAuthRefreshToken), profile.MCPOAuthRefreshToken == "")
69+
6770
tempFile := mcpConfigPath + ".tmp"
6871

6972
bytes, err := json.MarshalIndent(profile, "", "\t")
7073
if err != nil {
7174
return fmt.Errorf("failed to marshal profile: %w", err)
7275
}
7376

77+
// jsonStr := string(bytes)
78+
// hasRefreshToken := strings.Contains(jsonStr, "mcp_oauth_refresh_token")
79+
// log.Printf("saveMcpProfile: After marshaling - JSON contains 'mcp_oauth_refresh_token'=%v, JSON length=%d",
80+
// hasRefreshToken, len(jsonStr))
81+
7482
if err := os.WriteFile(tempFile, bytes, 0600); err != nil {
7583
return fmt.Errorf("failed to write temp file %q: %w", tempFile, err)
7684
}
@@ -80,64 +88,121 @@ func saveMcpProfile(profile *McpProfile) error {
8088
_ = os.Remove(tempFile)
8189
return fmt.Errorf("failed to rename temp file to %q: %w", mcpConfigPath, err)
8290
}
91+
92+
log.Printf("saveMcpProfile: Successfully saved MCP profile")
8393
return nil
8494
}
8595

96+
// loadExistingMCPProfile 加载并验证已有的 MCP profile，如果有效则返回，避免重复拉起 OAuth
97+
func loadExistingMCPProfile(ctx *cli.Context, profile config.Profile, opts ProxyConfig, desiredAppName string) *McpProfile {
98+
mcpConfigPath := getMCPConfigPath()
99+
bytes, err := os.ReadFile(mcpConfigPath)
100+
if err != nil {
101+
return nil
102+
}
103+
mcpProfile, err := NewMcpProfileFromBytes(bytes)
104+
if err != nil {
105+
return nil
106+
}
107+
108+
if mcpProfile.MCPOAuthSiteType != string(opts.RegionType) {
109+
log.Printf("Region type mismatch: saved=%s, requested=%s, ignoring local profile", mcpProfile.MCPOAuthSiteType, string(opts.RegionType))
110+
return nil
111+
}
112+
113+
if mcpProfile.MCPOAuthAppName != desiredAppName {
114+
log.Printf("App name mismatch: saved=%s, requested=%s, ignoring local profile", mcpProfile.MCPOAuthAppName, desiredAppName)
115+
return nil
116+
}
117+
118+
if mcpProfile.MCPOAuthAppId == "" {
119+
log.Printf("MCP profile with AppId is empty, ignoring local profile")
120+
return nil
121+
}
122+
123+
if mcpProfile.MCPOAuthRefreshToken == "" {
124+
log.Printf("MCP profile with RefreshToken is empty, ignoring local profile")
125+
return nil
126+
}
127+
128+
if mcpProfile.MCPOAuthRefreshTokenExpire <= util.GetCurrentUnixTime() {
129+
log.Printf("MCP profile with RefreshTokenExpire is expired, ignoring local profile")
130+
return nil
131+
}
132+
133+
app, err := findOAuthApplicationById(ctx, profile, mcpProfile.MCPOAuthAppId, opts.RegionType)
134+
if err != nil {
135+
log.Printf("Failed to reuse existing MCP profile (app: %s): %v, ignoring local profile", mcpProfile.MCPOAuthAppName, err)
136+
return nil
137+
}
138+
if app == nil {
139+
log.Printf("OAuth application with AppId '%s' not found, ignoring local profile", mcpProfile.MCPOAuthAppId)
140+
return nil
141+
}
142+
143+
if err := validateOAuthApplication(app, opts.Scope, opts.Host, opts.Port); err != nil {
144+
log.Printf("Reused existing MCP profile validation failed: %v, ignoring local profile", err)
145+
return nil
146+
}
147+
148+
// 根据远端 app 信息更新 mcp profile 中的相关字段，其他字段（如 token）保持不变
149+
mcpProfile.MCPOAuthAppName = app.AppName
150+
mcpProfile.MCPOAuthAppId = app.ApplicationId
151+
mcpProfile.MCPOAuthAccessTokenValidity = app.AccessTokenValidity
152+
mcpProfile.MCPOAuthRefreshTokenValidity = app.RefreshTokenValidity
153+
154+
log.Printf("Reused existing MCP profile with app '%s' (AppId: %s)", app.AppName, app.ApplicationId)
155+
156+
return mcpProfile
157+
}
158+
86159
func getOrCreateMCPProfile(ctx *cli.Context, opts ProxyConfig) (*McpProfile, error) {
87160
profile, err := config.LoadProfileWithContext(ctx)
88161
if err != nil {
89162
return nil, fmt.Errorf("failed to load profile: %w", err)
90163
}
91164

92-
// 如果传入了 oauth-app-name，先验证该应用是否存在且合法
93-
// 如果已经验证过 oauth-app-name，直接使用验证过的 app；否则查找或创建默认的 OAuth 应用
94-
var validatedApp *OAuthApplication
95-
if opts.OAuthAppName != "" {
96-
app, err := findOAuthApplicationByName(ctx, profile, opts.RegionType, opts.OAuthAppName)
97-
if err != nil {
98-
return nil, fmt.Errorf("failed to find OAuth application '%s': %w", opts.OAuthAppName, err)
99-
}
100-
if app == nil {
101-
return nil, fmt.Errorf("OAuth application '%s' not found", opts.OAuthAppName)
102-
}
165+
// 如果未显式指定 app name，则使用默认的 MCPOAuthAppName，便于复用本地 profile
166+
desiredAppName := opts.OAuthAppName
167+
if desiredAppName == "" {
168+
desiredAppName = MCPOAuthAppName
169+
}
103170

104-
// 验证 Scopes 和 Callback URI
105-
requiredRedirectURI := buildRedirectUri(opts.Host, opts.Port)
106-
if err := validateOAuthApplication(app, opts.Scope, requiredRedirectURI); err != nil {
107-
return nil, fmt.Errorf("OAuth application validation failed: %w", err)
171+
existingMcpProfile := loadExistingMCPProfile(ctx, profile, opts, desiredAppName)
172+
if existingMcpProfile != nil {
173+
// mcpprofile might change, save it again to ensure the latest state is saved
174+
if err := saveMcpProfile(existingMcpProfile); err != nil {
175+
return nil, fmt.Errorf("failed to save mcp profile: %w", err)
108176
}
177+
return existingMcpProfile, nil
178+
}
109179

110-
validatedApp = app
111-
cli.Printf(ctx.Stdout(), "Using existing OAuth application '%s' (AppId: %s)\n", app.AppName, app.ApplicationId)
112-
} else {
113-
// 查找或创建默认的 OAuth 应用
114-
mcpConfigPath := getMCPConfigPath()
115-
if bytes, err := os.ReadFile(mcpConfigPath); err == nil {
116-
if mcpProfile, err := NewMcpProfileFromBytes(bytes); err == nil {
117-
log.Println("MCP Profile loaded from file", mcpProfile.Name, "app id", mcpProfile.MCPOAuthAppId)
118-
119-
// 检查 region type 是否匹配，因为国内和国际站的 OAuth 地址不同, Region type 不匹配则重新创建 profile
120-
if mcpProfile.MCPOAuthSiteType != string(opts.RegionType) {
121-
log.Printf("Region type mismatch: saved=%s, requested=%s, recreating profile", mcpProfile.MCPOAuthSiteType, string(opts.RegionType))
122-
} else {
123-
err = findOAuthApplicationById(ctx, profile, mcpProfile, opts.RegionType)
124-
if err == nil {
125-
return mcpProfile, nil
126-
} else {
127-
log.Println("Failed to find existing OAuth application", err.Error())
128-
}
129-
}
130-
}
180+
app, err := findOAuthApplicationByName(ctx, profile, opts.RegionType, desiredAppName)
181+
if err != nil {
182+
return nil, fmt.Errorf("failed to find OAuth application '%s': %w", desiredAppName, err)
183+
}
184+
185+
if app == nil {
186+
if opts.OAuthAppName != "" {
187+
// if user provide app name, but not found, return error
188+
return nil, fmt.Errorf("OAuth application '%s' not found", opts.OAuthAppName)
131189
}
132-
app, err := getOrCreateMCPOAuthApplication(ctx, profile, opts.RegionType, opts.Host, opts.Port, opts.Scope)
190+
cli.Printf(ctx.Stdout(), "Creating new default MCP profile '%s'...\n", DefaultMcpProfileName)
191+
app, err = createDefaultMCPOauthApplication(ctx, profile, opts.RegionType, opts.Host, opts.Port, opts.Scope)
133192
if err != nil {
134-
return nil, fmt.Errorf("failed to get or create OAuth application: %w", err)
193+
return nil, fmt.Errorf("failed to create default OAuth application: %w", err)
135194
}
136-
validatedApp = app
195+
cli.Printf(ctx.Stdout(), "Created new default OAuth application '%s' (AppId: %s)\n", app.AppName, app.ApplicationId)
196+
} else {
197+
cli.Printf(ctx.Stdout(), "Using existing OAuth application '%s' (AppId: %s)\n", app.AppName, app.ApplicationId)
137198
}
138199

139-
cli.Printf(ctx.Stdout(), "Setting up MCPOAuth profile '%s'...\n", DefaultMcpProfileName)
200+
if err := validateOAuthApplication(app, opts.Scope, opts.Host, opts.Port); err != nil {
201+
return nil, fmt.Errorf("OAuth application validation failed: %w", err)
202+
}
203+
validatedApp := app
140204

205+
cli.Printf(ctx.Stdout(), "Setting up MCPOAuth profile '%s'...\n", DefaultMcpProfileName)
141206
mcpProfile := NewMcpProfile(DefaultMcpProfileName)
142207
mcpProfile.MCPOAuthSiteType = string(opts.RegionType)
143208
mcpProfile.MCPOAuthAppId = validatedApp.ApplicationId
@@ -153,9 +218,20 @@ func getOrCreateMCPProfile(ctx *cli.Context, opts ProxyConfig) (*McpProfile, err
153218
if err != nil {
154219
return nil, fmt.Errorf("OAuth login failed: %w", err)
155220
}
221+
222+
log.Printf("OAuth flow completed: AccessToken length=%d, RefreshToken length=%d, AccessTokenExpire=%d",
223+
len(tokenResult.AccessToken), len(tokenResult.RefreshToken), tokenResult.AccessTokenExpire)
224+
if tokenResult.RefreshToken == "" {
225+
return nil, fmt.Errorf("OAuth flow returned empty RefreshToken (Region=%s, AppId=%s). "+
226+
"Please delete this application and let the system create a new NativeApp, or manually create a NativeApp",
227+
opts.RegionType, mcpProfile.MCPOAuthAppId)
228+
}
229+
156230
mcpProfile.MCPOAuthAccessToken = tokenResult.AccessToken
157231
mcpProfile.MCPOAuthRefreshToken = tokenResult.RefreshToken
158232
mcpProfile.MCPOAuthAccessTokenExpire = tokenResult.AccessTokenExpire
233+
// refresh token will be updated each time latest access token is refreshed,
234+
// however the validity and expiration time is the same as the original when finishing oauth flow
159235
mcpProfile.MCPOAuthRefreshTokenExpire = currentTime + int64(validatedApp.RefreshTokenValidity)
160236

161237
if err = saveMcpProfile(mcpProfile); err != nil {

0 commit comments

Comments
 (0)