From 395356744d70f72b534be079ec63218b5d73a8a4 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Wed, 29 Jul 2026 10:30:08 +0530 Subject: [PATCH] fix: gate SSE OnConnectionLost on connection identity, preserve NeedsReauth --- core/mcp/clientmanager.go | 48 ++++++-- core/mcp/healthmonitor.go | 11 +- core/mcp/reauth_state_test.go | 201 ++++++++++++++++++++++++++++++++++ core/mcp/utils.go | 20 ++++ core/schemas/mcp.go | 10 ++ 5 files changed, 281 insertions(+), 9 deletions(-) create mode 100644 core/mcp/reauth_state_test.go diff --git a/core/mcp/clientmanager.go b/core/mcp/clientmanager.go index 0c895cf5c82..d75ea10bd33 100644 --- a/core/mcp/clientmanager.go +++ b/core/mcp/clientmanager.go @@ -931,13 +931,19 @@ func (m *MCPManager) EnableClient(id string) error { if err := m.connectToMCPClient(m.ctx, configCopy); err != nil { // Connection failed — leave the entry as Disconnected so the health monitor can - // recover it, but only if the client has not been disabled in the meantime. + // recover it, but only if the client has not been disabled in the meantime, and + // don't clobber NeedsReauth: connectToMCPClient already classified this failure + // as a dead OAuth2 credential (under its own lock, above), and a generic + // Disconnected here would silently erase that more specific signal. m.mu.Lock() alreadyDisabled := false if cs, exists := m.clientMap[id]; exists { - if cs.State == schemas.MCPConnectionStateDisabled { + switch cs.State { + case schemas.MCPConnectionStateDisabled: alreadyDisabled = true - } else { + case schemas.MCPConnectionStateNeedsReauth: + // preserve as-is + default: cs.State = schemas.MCPConnectionStateDisconnected } } @@ -1527,6 +1533,23 @@ func (m *MCPManager) connectToMCPClient(requestCtx context.Context, config *sche m.logger.Warn("%s Failed to close external client during cleanup: %v", MCPLogPrefix, closeErr) } } + + // A dead OAuth2 credential (refresh permanently rejected/expired) is not a + // transient connectivity problem — the entry above was already reset to + // Disconnected at the top of this function, but that's misleading here: no + // amount of automatic reconnecting will fix it, only a human reauthorizing + // the client will. Flip to NeedsReauth instead so the state itself carries + // that signal, mirroring how the success path below sets Connected under + // the same lock. Never clobber Disabled — DisableClient is authoritative, + // same invariant the health monitor's updateClientState guard preserves. + if isOAuth2TokenExpiredErrorText(gateErr.GetErrorString()) { + m.mu.Lock() + if client, exists := m.clientMap[config.ID]; exists && client.State != schemas.MCPConnectionStateDisabled { + client.State = schemas.MCPConnectionStateNeedsReauth + } + m.mu.Unlock() + } + return fmt.Errorf("failed to connect MCP client %s: %s", config.Name, gateErr.GetErrorString()) } @@ -1645,12 +1668,21 @@ func (m *MCPManager) connectToMCPClient(requestCtx context.Context, config *sche if config.ConnectionType == schemas.MCPConnectionTypeSSE && externalClient != nil { externalClient.OnConnectionLost(func(err error) { m.logger.Warn("%s SSE connection lost for MCP server '%s': %v", MCPLogPrefix, config.Name, err) - // Update state to disconnected, but never overwrite a disabled state. - // DisableClient calls Conn.Close() while holding m.mu; the SSE library - // fires OnConnectionLost after the lock is released, by which point - // State is already Disabled — do not clobber it. + // Update state to disconnected, but never overwrite a disabled or + // needs-reauth state. DisableClient calls Conn.Close() while holding + // m.mu; the SSE library fires OnConnectionLost after the lock is + // released, by which point State is already Disabled — do not clobber + // it. Also gate on client.Conn == externalClient: a reconnect can + // replace this entry's Conn (and its State, e.g. to NeedsReauth on a + // dead OAuth2 credential) before this closed connection's own + // OnConnectionLost callback gets a chance to fire, so an identity + // check is required — a plain State-value check can't tell a stale + // callback from a live one. m.mu.Lock() - if client, exists := m.clientMap[config.ID]; exists && client.State != schemas.MCPConnectionStateDisabled { + if client, exists := m.clientMap[config.ID]; exists && + client.Conn == externalClient && + client.State != schemas.MCPConnectionStateDisabled && + client.State != schemas.MCPConnectionStateNeedsReauth { client.State = schemas.MCPConnectionStateDisconnected } m.mu.Unlock() diff --git a/core/mcp/healthmonitor.go b/core/mcp/healthmonitor.go index cadf1d71561..e4a77af8d90 100644 --- a/core/mcp/healthmonitor.go +++ b/core/mcp/healthmonitor.go @@ -257,7 +257,16 @@ func (chm *ClientHealthMonitor) updateClientState(state schemas.MCPConnectionSta // Never overwrite a disabled state. DisableClient is authoritative: a health // check tick or reconnect callback that races with DisableClient must not // flip the client back to Disconnected/Connected. - if clientState.State == schemas.MCPConnectionStateDisabled { + // + // Same treatment for NeedsReauth: it means connectToMCPClient already + // determined the failure is a dead OAuth2 credential, not a routine + // connectivity blip. A ping success/failure on the same dead connection + // (or the Disconnected write performHealthCheck makes right before kicking + // off a reconnect attempt) must not silently flip that back to + // Connected/Disconnected — only a human reauthorizing the client (or a + // reconnect that succeeds because the credential was fixed) should move it + // out of this state. + if clientState.State == schemas.MCPConnectionStateDisabled || clientState.State == schemas.MCPConnectionStateNeedsReauth { chm.manager.mu.Unlock() return } diff --git a/core/mcp/reauth_state_test.go b/core/mcp/reauth_state_test.go new file mode 100644 index 00000000000..b3b43c98150 --- /dev/null +++ b/core/mcp/reauth_state_test.go @@ -0,0 +1,201 @@ +package mcp + +import ( + "context" + "fmt" + "net/http" + "testing" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestIsOAuth2TokenExpiredErrorText covers the substring-matching fallback +// connectToMCPClient relies on to recognize schemas.ErrOAuth2TokenExpired at +// the failure site: runConnectWithPluginPipeline flattens the underlying Go +// error to a string (see the doc comment on isOAuth2TokenExpiredErrorText), +// so errors.Is/As isn't usable there, and text matching is the only option +// left, mirroring isTransientError. +func TestIsOAuth2TokenExpiredErrorText(t *testing.T) { + wrappedOnce := fmt.Errorf("refresh token rejected by upstream OAuth server, re-authentication required: %w", schemas.ErrOAuth2TokenExpired) + wrappedTwice := fmt.Errorf("token expired and refresh failed: %w", wrappedOnce) + + tests := []struct { + name string + errStr string + want bool + }{ + {"exact sentinel text", schemas.ErrOAuth2TokenExpired.Error(), true}, + {"single %w wrap (RefreshAccessToken permanent-rejection path)", wrappedOnce.Error(), true}, + {"double %w wrap (GetAccessToken's refresh-failed wrapper)", wrappedTwice.Error(), true}, + {"token-not-active wrap (GetAccessToken's inactive-status path)", fmt.Errorf("oauth token is not active, status: orphaned: %w", schemas.ErrOAuth2TokenExpired).Error(), true}, + {"unrelated connectivity error", "connection refused", false}, + {"unrelated config error", "oauth2 config not found", false}, + {"empty string", "", false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, isOAuth2TokenExpiredErrorText(tc.errStr)) + }) + } +} + +// expiredOAuthCredStore simulates a shared OAuth2 credential store whose +// refresh has permanently failed: every ConnectionHeaders call returns the +// same error shape framework/oauth2's RefreshAccessToken/GetAccessToken +// produce for a dead refresh token, wrapping schemas.ErrOAuth2TokenExpired. +type expiredOAuthCredStore struct{} + +func (expiredOAuthCredStore) ConnectionHeaders(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) (http.Header, error) { + return nil, fmt.Errorf("refresh token rejected by upstream OAuth server, re-authentication required: %w", schemas.ErrOAuth2TokenExpired) +} + +func (expiredOAuthCredStore) RequestHeaders(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) (http.Header, error) { + return http.Header{}, nil +} + +func (expiredOAuthCredStore) RequiresPerCallConnection(_ *schemas.MCPClientConfig) bool { + return false +} + +// newSharedOAuthClientConfig builds a minimal shared-OAuth (auth_type=oauth, +// persistent-connection) MCP client config for the connectToMCPClient tests +// below. ConnectionType is HTTP so the failure happens at the +// credStore.ConnectionHeaders call inside connectToMCPClient's op closure, +// before any real network dial is attempted. +func newSharedOAuthClientConfig(id string) *schemas.MCPClientConfig { + oauthConfigID := "oauth-config-1" + return &schemas.MCPClientConfig{ + ID: id, + Name: "reauth-test-client-" + id, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar("https://example.invalid/mcp"), + AuthType: schemas.MCPAuthTypeOauth, + OauthConfigID: &oauthConfigID, + } +} + +// TestConnectToMCPClient_OAuth2TokenExpired_SetsNeedsReauth verifies the +// core wiring point of this change: a connect failure whose underlying cause +// is schemas.ErrOAuth2TokenExpired (a dead shared-OAuth credential) lands the +// client in MCPConnectionStateNeedsReauth, not the generic Disconnected the +// entry was initialized to at the top of connectToMCPClient. +func TestConnectToMCPClient_OAuth2TokenExpired_SetsNeedsReauth(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, expiredOAuthCredStore{}, nil, nil) + config := newSharedOAuthClientConfig("client-needs-reauth") + + // Confirmed precondition: only shared-connection auth types ( + // RequiresPerCallConnection()==false) ever reach connectToMCPClient in + // the first place — AddClient/EnableClient/UpdateClientConnection all + // special-case per-call-connection (per-user) auth types before calling + // it, and ReconnectClient refuses outright for them. expiredOAuthCredStore + // mirrors that: RequiresPerCallConnection returns false. + require.False(t, m.credStore.RequiresPerCallConnection(config)) + + err := m.connectToMCPClient(context.Background(), config) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to connect MCP client") + + m.mu.RLock() + state, exists := m.clientMap[config.ID] + m.mu.RUnlock() + require.True(t, exists) + assert.Equal(t, schemas.MCPConnectionStateNeedsReauth, state.State) +} + +// TestConnectToMCPClient_GenericFailure_StaysDisconnected is the control +// case: a connect failure that is NOT an ErrOAuth2TokenExpired-wrapped error +// (e.g. a plain connectivity error) must leave the client in the existing +// generic Disconnected state, not NeedsReauth. +type genericFailureCredStore struct{} + +func (genericFailureCredStore) ConnectionHeaders(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) (http.Header, error) { + return nil, fmt.Errorf("connection refused") +} + +func (genericFailureCredStore) RequestHeaders(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) (http.Header, error) { + return http.Header{}, nil +} + +func (genericFailureCredStore) RequiresPerCallConnection(_ *schemas.MCPClientConfig) bool { + return false +} + +func TestConnectToMCPClient_GenericFailure_StaysDisconnected(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, genericFailureCredStore{}, nil, nil) + config := newSharedOAuthClientConfig("client-generic-failure") + + err := m.connectToMCPClient(context.Background(), config) + require.Error(t, err) + + m.mu.RLock() + state, exists := m.clientMap[config.ID] + m.mu.RUnlock() + require.True(t, exists) + assert.Equal(t, schemas.MCPConnectionStateDisconnected, state.State) +} + +// TestUpdateClientState_PreservesNeedsReauth covers healthmonitor.go's +// updateClientState guard: a routine health-check ping success/failure must +// not silently flip a NeedsReauth client back to Connected/Disconnected — +// only a human reauthorizing the client (surfaced elsewhere as a fresh +// connectToMCPClient success, which sets Connected unconditionally) should +// move it out of this state. +func TestUpdateClientState_PreservesNeedsReauth(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, expiredOAuthCredStore{}, nil, nil) + config := newSharedOAuthClientConfig("client-health-cycle") + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStateNeedsReauth, + } + m.mu.Unlock() + + chm := NewClientHealthMonitor(m, config.ID, DefaultHealthCheckInterval, true, nil) + + // A failed health-check tick (performHealthCheck's failure branch) tries + // to write Disconnected first. + chm.updateClientState(schemas.MCPConnectionStateDisconnected) + m.mu.RLock() + stateAfterFailureTick := m.clientMap[config.ID].State + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStateNeedsReauth, stateAfterFailureTick, "a failed health-check tick must not clobber NeedsReauth back to Disconnected") + + // A successful ping (performHealthCheck's success branch) tries to write + // Connected — this must also be rejected: the ping succeeding against a + // stale/absent transport does not mean the credential is fixed. + chm.updateClientState(schemas.MCPConnectionStateConnected) + m.mu.RLock() + stateAfterSuccessTick := m.clientMap[config.ID].State + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStateNeedsReauth, stateAfterSuccessTick, "a successful health-check tick must not clobber NeedsReauth back to Connected") +} + +// TestUpdateClientState_StillPreservesDisabled is a regression guard for the +// pre-existing Disabled-preservation behavior updateClientState had before +// this change — the new NeedsReauth branch must be additive, not a +// replacement. +func TestUpdateClientState_StillPreservesDisabled(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, expiredOAuthCredStore{}, nil, nil) + config := newSharedOAuthClientConfig("client-disabled") + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStateDisabled, + } + m.mu.Unlock() + + chm := NewClientHealthMonitor(m, config.ID, DefaultHealthCheckInterval, true, nil) + chm.updateClientState(schemas.MCPConnectionStateConnected) + + m.mu.RLock() + state := m.clientMap[config.ID].State + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStateDisabled, state) +} diff --git a/core/mcp/utils.go b/core/mcp/utils.go index 0bee59ffd58..bf21291192b 100644 --- a/core/mcp/utils.go +++ b/core/mcp/utils.go @@ -238,6 +238,26 @@ func isTransientError(err error) bool { return true } +// isOAuth2TokenExpiredErrorText reports whether a connect-failure message +// indicates the underlying failure was schemas.ErrOAuth2TokenExpired — a +// shared client's OAuth2 credential that has permanently died (refresh +// rejected/expired, no way to silently recover) and needs a human to +// reauthorize it — as opposed to a generic connectivity failure that a +// routine reconnect can resolve on its own. +// +// connectToMCPClient's op closure returns this sentinel (wrapped via %w by +// framework/oauth2) as a plain Go error, but runConnectWithPluginPipeline +// (unlike RunWithPluginPipeline, used for tool/ping/list_tools calls) only +// carries opErr.Error() as a string on the *schemas.BifrostError it returns — +// it does not also preserve the original error on ErrorField.Error — so +// errors.Is/errors.As is not usable on the gateErr the caller receives. +// Same substring-matching technique as isTransientError above, for the same +// structural reason: the typed error info was already flattened to a string +// by the time it gets here. +func isOAuth2TokenExpiredErrorText(errStr string) bool { + return strings.Contains(errStr, schemas.ErrOAuth2TokenExpired.Error()) +} + // ExecuteWithRetry executes a function with exponential backoff retry logic. // Only retries on transient errors; permanent errors (auth, config) fail immediately. // It returns the error from the last attempt if all retries fail. diff --git a/core/schemas/mcp.go b/core/schemas/mcp.go index c0671d829a5..09bdf9b57a9 100644 --- a/core/schemas/mcp.go +++ b/core/schemas/mcp.go @@ -602,6 +602,16 @@ const ( MCPConnectionStatePendingTools MCPConnectionState = "pending_tools" // Connected but tools not yet populated MCPConnectionStatePendingVerification MCPConnectionState = "pending_verification" // Declared (typically via config.json) but the one-time auth/test flow has not been completed by an admin yet MCPConnectionStateDisabled MCPConnectionState = "disabled" // Client is intentionally disabled by the user + // MCPConnectionStateNeedsReauth means this client was previously authorized and + // connected at least once, but its credential can no longer be used — the upstream + // OAuth token died (refresh token rejected/expired with no way to silently recover, + // see ErrOAuth2TokenExpired) and a human has to reauthorize it. This is distinct from + // MCPConnectionStatePendingVerification, which means the one-time initial setup was + // never completed in the first place; NeedsReauth means setup succeeded once and the + // credential died later. Only reachable by shared-connection auth types (a persistent + // upstream connection to reconnect); per-user auth types resolve credentials per-call + // and never hold a connection this state describes. + MCPConnectionStateNeedsReauth MCPConnectionState = "needs_reauth" ) // MCPClientState represents a connected MCP client with its configuration and tools.