diff --git a/core/bifrost.go b/core/bifrost.go index e84efd223e..dd016c766b 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -4204,12 +4204,12 @@ func (bifrost *Bifrost) UpdateMCPClient(id string, updatedConfig *schemas.MCPCli return bifrost.MCPManager.UpdateClient(id, updatedConfig) } -// UpdateMCPClientConnection reconnects an existing MCP client using updated headers -func (bifrost *Bifrost) UpdateMCPClientConnection(id string, newConfig *schemas.MCPClientConfig) error { +// UpdateMCPClientCredentials reconnects an existing MCP client using updated headers +func (bifrost *Bifrost) UpdateMCPClientCredentials(id string, newConfig *schemas.MCPClientConfig) error { if bifrost.MCPManager == nil { return fmt.Errorf("mcp is not configured in this bifrost instance") } - return bifrost.MCPManager.UpdateClientConnection(id, newConfig) + return bifrost.MCPManager.UpdateClientCredentials(id, newConfig) } // ReconnectMCPClient attempts to reconnect an MCP client if it is disconnected. diff --git a/core/mcp/addclient_discoveredtools_test.go b/core/mcp/addclient_discoveredtools_test.go index 383fc8ce9f..9e1a0b8622 100644 --- a/core/mcp/addclient_discoveredtools_test.go +++ b/core/mcp/addclient_discoveredtools_test.go @@ -98,3 +98,69 @@ func TestAddClient_PerUserHeaders_DiscoveredToolsStateSelection(t *testing.T) { }) } } + +// TestAddClient_PerCallConnection_StartsConnectionChecker pins the fix for +// the shared-per-call tool-discovery bug: AddClient's per-call branch (any +// auth type RequiresPerCallConnection resolves true for — per_user_headers +// here) must start a connection checker, exactly like the sticky success +// path (connectToMCPClient) already does. Without one, nothing ever +// revisits a client whose DiscoveredTools was empty/nil at add time — it +// would sit at Healthy/0-tools forever. +func TestAddClient_PerCallConnection_StartsConnectionChecker(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + + config := &schemas.MCPClientConfig{ + ID: "per-call-checker-client", + Name: "per_call_checker_client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar("https://example.invalid/mcp"), + AuthType: schemas.MCPAuthTypePerUserHeaders, + PerUserHeaderKeys: []string{"X-User-Token"}, + // Non-nil (even empty) DiscoveredTools: a nil map means admin + // verification never ran, parking in pending_verification BEFORE + // ever reaching the RequiresPerCallConnection branch this test + // targets — see TestAddClient_PerUserHeaders_DiscoveredToolsStateSelection. + DiscoveredTools: map[string]schemas.ChatTool{}, + } + + require.NoError(t, m.AddClient(context.Background(), config)) + + m.checkerManager.mu.RLock() + _, hasChecker := m.checkerManager.checkers[config.ID] + m.checkerManager.mu.RUnlock() + assert.True(t, hasChecker, "AddClient's per-call branch must start a connection checker, or this client's tools are never (re)discovered") +} + +// TestAddClient_PerCallConnection_SharedType_DiscoversToolsSynchronously pins +// the second half of the same bug report: a connection checker alone is not +// enough, because ClientConnectionChecker.Start uses the SLOW healthyInterval +// (not the fast Unstable one) for a client that already starts Healthy — so +// without a synchronous first discovery pass, a shared oauth/headers/none +// per-call client would sit at Healthy/0-tools for up to healthyInterval +// before its first real check ever ran. AddClient must discover tools +// synchronously, mirroring what per_user_headers/token_exchange already get +// from their own admin-verify-at-setup-time step. +func TestAddClient_PerCallConnection_SharedType_DiscoversToolsSynchronously(t *testing.T) { + ts, _ := buildAdminDiscoveryHTTPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + + config := &schemas.MCPClientConfig{ + ID: "shared-sync-discovery-client", + Name: "shared_sync_discovery_client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar(ts.URL), + AuthType: schemas.MCPAuthTypeHeaders, + Headers: map[string]schemas.SecretVar{"Authorization": *schemas.NewSecretVar("Bearer shared-add-token")}, + // NeedsSessionStickiness left nil: the default per-call value. + // DiscoveredTools left nil: nothing to restore, forcing the new + // synchronous-discovery path this test targets. + } + + require.NoError(t, m.AddClient(context.Background(), config)) + + state, ok := snapshotAddClientState(m, config.ID) + require.True(t, ok) + assert.Equal(t, schemas.MCPConnectionStateHealthy, state.State) + assert.Contains(t, state.ToolMap, "shared_sync_discovery_client-echo", "tools must be discovered synchronously during AddClient, not deferred entirely to the periodic checker") +} diff --git a/core/mcp/admin_tool_discovery_test.go b/core/mcp/admin_tool_discovery_test.go index 59905705fa..1ab5e98f9e 100644 --- a/core/mcp/admin_tool_discovery_test.go +++ b/core/mcp/admin_tool_discovery_test.go @@ -234,39 +234,119 @@ func TestPerformAdminToolDiscovery_PerUserHeaders_ConvertsHeadersAndDispatches(t } } -// TestPerformAdminToolDiscovery_UnsupportedAuthType_ReturnsError confirms -// every non-per-user auth type falls into the default branch and errors -// instead of attempting discovery — shared auth types have no distinct -// "admin" credential to sync with. +// TestPerformAdminToolDiscovery_UnsupportedAuthType_ReturnsError confirms a +// genuinely unknown auth type falls into the default branch and errors +// instead of attempting discovery. func TestPerformAdminToolDiscovery_UnsupportedAuthType_ReturnsError(t *testing.T) { - authTypes := []schemas.MCPAuthType{ - schemas.MCPAuthTypeNone, - schemas.MCPAuthTypeHeaders, - schemas.MCPAuthTypeOauth, - schemas.MCPAuthType("something_unknown"), + cred := &fakeAdminCredStore{headers: http.Header{"Authorization": []string{"Bearer x"}}} + m := &MCPManager{credStore: cred, logger: &MockLogger{}} + + config := &schemas.MCPClientConfig{ + ID: "client-3", + Name: "shared-client", + AuthType: schemas.MCPAuthType("something_unknown"), } - for _, authType := range authTypes { + tools, mapping, err := m.performAdminToolDiscovery(context.Background(), config) + require.Error(t, err) + require.Nil(t, tools) + require.Nil(t, mapping) + require.True(t, strings.Contains(err.Error(), "admin tool discovery not supported for auth_type")) +} + +// TestPerformAdminToolDiscovery_SharedOAuth_PerCall_ExtractsBearerTokenAndDispatches +// pins the fix for the shared-OAuth-per-call tool-discovery bug: auth_type +// "oauth" now routes through the same bearer-based dispatch as +// per_user_oauth/token_exchange, since sharedOAuthResolver.AdminConnectionHeaders +// delegates to ConnectionHeaders for a per-call client (see credstore's +// shared_oauth.go) instead of erroring. +func TestPerformAdminToolDiscovery_SharedOAuth_PerCall_ExtractsBearerTokenAndDispatches(t *testing.T) { + ts, rec := buildAdminDiscoveryHTTPServer(t) + + cred := &fakeAdminCredStore{headers: http.Header{"Authorization": []string{"Bearer shared-secret-token"}}} + m := &MCPManager{credStore: cred, logger: &MockLogger{}} + + config := &schemas.MCPClientConfig{ + ID: "client-5", + Name: "shared-oauth-client", + AuthType: schemas.MCPAuthTypeOauth, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar(ts.URL), + // NeedsSessionStickiness left nil: the default per-call value. + } + + tools, mapping, err := m.performAdminToolDiscovery(context.Background(), config) + require.NoError(t, err) + require.Contains(t, tools, "shared-oauth-client-echo") + require.Equal(t, "echo", mapping["echo"]) + require.Equal(t, 1, cred.callCount()) + + authVals := rec.headerValues("Authorization") + require.NotEmpty(t, authVals) + for _, v := range authVals { + require.Equal(t, "Bearer shared-secret-token", v) + } +} + +// TestPerformAdminToolDiscovery_SharedHeadersAndNone_PerCall_ConvertsHeadersAndDispatches +// is the shared headers/none counterpart, table-driven since both route +// through the same headers-map dispatch branch as per_user_headers. +func TestPerformAdminToolDiscovery_SharedHeadersAndNone_PerCall_ConvertsHeadersAndDispatches(t *testing.T) { + for _, authType := range []schemas.MCPAuthType{schemas.MCPAuthTypeHeaders, schemas.MCPAuthTypeNone} { t.Run(string(authType), func(t *testing.T) { - cred := &fakeAdminCredStore{headers: http.Header{"Authorization": []string{"Bearer x"}}} + ts, rec := buildAdminDiscoveryHTTPServer(t) + + cred := &fakeAdminCredStore{headers: http.Header{"X-Api-Key": []string{"shared-key-1"}}} m := &MCPManager{credStore: cred, logger: &MockLogger{}} config := &schemas.MCPClientConfig{ - ID: "client-3", - Name: "shared-client", - AuthType: authType, + ID: "client-6", + Name: "shared-headers-client", + AuthType: authType, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar(ts.URL), } tools, mapping, err := m.performAdminToolDiscovery(context.Background(), config) - require.Error(t, err) - require.Nil(t, tools) - require.Nil(t, mapping) - require.True(t, strings.Contains(err.Error(), "admin tool discovery not supported for auth_type")) - require.Zero(t, cred.callCount(), "unsupported auth types must not resolve admin credentials") + require.NoError(t, err) + require.Contains(t, tools, "shared-headers-client-echo") + require.Equal(t, "echo", mapping["echo"]) + require.Equal(t, 1, cred.callCount()) + + apiKeyVals := rec.headerValues("X-Api-Key") + require.NotEmpty(t, apiKeyVals) + for _, v := range apiKeyVals { + require.Equal(t, "shared-key-1", v) + } }) } } +// TestPerformAdminToolDiscovery_SharedHeaders_PerCall_EmptyHeadersStillDispatches +// pins the VerifyHeadersConnection guard relaxation: a shared headers/none +// client with NO resolved headers at all (auth carried entirely by other +// static config headers, or none) must still attempt discovery rather than +// erroring on "user headers are required" — that guard is per_user_headers +// specific. +func TestPerformAdminToolDiscovery_SharedHeaders_PerCall_EmptyHeadersStillDispatches(t *testing.T) { + ts, _ := buildAdminDiscoveryHTTPServer(t) + + cred := &fakeAdminCredStore{headers: http.Header{}} // resolves to nothing + m := &MCPManager{credStore: cred, logger: &MockLogger{}} + + config := &schemas.MCPClientConfig{ + ID: "client-7", + Name: "shared-none-client", + AuthType: schemas.MCPAuthTypeNone, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar(ts.URL), + } + + tools, _, err := m.performAdminToolDiscovery(context.Background(), config) + require.NoError(t, err) + require.Contains(t, tools, "shared-none-client-echo") +} + // TestPerformAdminToolDiscovery_CredStoreError_Propagates confirms a // credStore.AdminConnectionHeaders failure is wrapped and returned as-is, // without attempting any connection. diff --git a/core/mcp/clientmanager.go b/core/mcp/clientmanager.go index f97813fe7b..6baac1a9fb 100644 --- a/core/mcp/clientmanager.go +++ b/core/mcp/clientmanager.go @@ -362,8 +362,12 @@ func (m *MCPManager) CloseAndMarkNeedsReauth(id string) (retErr error) { if !ok { return fmt.Errorf("client %s not found", id) } + // Covers both genuine per-user auth types (always per-call) and a shared + // HTTP client (oauth/headers) currently running in per-call mode via + // needs_session_stickiness nil/false — either way there is no + // persistent connection here to close. if clientState.ExecutionConfig != nil && m.credStore.RequiresPerCallConnection(clientState.ExecutionConfig) { - return fmt.Errorf("per-user auth clients do not maintain a shared upstream connection (each user manages their own auth): %w", schemas.ErrMCPReconnectNotApplicable) + return fmt.Errorf("client uses per-call connections; there is no persistent connection to close: %w", schemas.ErrMCPReconnectNotApplicable) } if clientState.State == schemas.MCPConnectionStateDisabled { // DisableClient is authoritative; a rotation racing a disable must @@ -410,12 +414,14 @@ func (m *MCPManager) ReconnectClient(id string) (retErr error) { m.mu.Unlock() return fmt.Errorf("client %s not found", id) } - // Per-user auth types do not maintain a persistent upstream connection — - // auth is resolved per request/user identity, so there's nothing to - // reconnect. + // Per-call clients have no persistent upstream connection to reconnect — + // this covers both genuine per-user auth types (always per-call, auth + // resolved per request/user identity) and a shared HTTP client + // (oauth/headers) currently running in per-call mode via + // needs_session_stickiness nil/false. if client.ExecutionConfig != nil && m.credStore.RequiresPerCallConnection(client.ExecutionConfig) { m.mu.Unlock() - return fmt.Errorf("per-user auth clients do not maintain a shared upstream connection (each user manages their own auth): %w", schemas.ErrMCPReconnectNotApplicable) + return fmt.Errorf("client uses per-call connections; there is no persistent connection to reconnect: %w", schemas.ErrMCPReconnectNotApplicable) } // Clients awaiting admin OAuth authorization have no token to connect // with; a reconnect attempt would only flip them out of @@ -634,6 +640,38 @@ func (m *MCPManager) AddClient(requestCtx context.Context, config *schemas.MCPCl } } m.mu.Unlock() + // A client with nothing to restore gets its first discovery pass + // synchronously here — mirrors per_user_headers/token_exchange's own + // admin-verify-at-setup-time step (VerifyHeadersConnection / + // VerifyPerUserOAuthConnection, run by the create/verify HTTP + // handler before this ever runs). Shared oauth/headers/none have no + // such separate step, so without this they'd sit at Healthy/0-tools + // until the periodic checker's first tick — which, since Start() + // uses the slow healthyInterval (not the fast Unstable one) for a + // client that already starts Healthy, could be minutes away. + // Best-effort: a failure here just means the periodic checker's own + // retries pick it up shortly instead — not a reason to fail the + // whole add/connect. + if len(config.DiscoveredTools) == 0 { + if tools, mapping, discErr := m.performAdminToolDiscovery(requestCtx, config); discErr != nil { + m.logger.Debug("%s Initial per-call tool discovery failed for MCP client '%s': %v — the periodic checker will retry", MCPLogPrefix, config.Name, discErr) + } else { + m.SetClientTools(config.ID, tools, mapping) + } + } + // Start the connection checker so this client's tools stay current + // going forward. Released BEFORE this, mirroring connectToMCPClient's + // own success path (StartMonitoring -> Start() takes m.mu.RLock + // internally, so calling it under the write lock above would deadlock). + if config.ID != BifrostMCPClientKey { + isPingAvailable := true + if config.IsPingAvailable != nil { + isPingAvailable = *config.IsPingAvailable + } + syncInterval := ResolveToolSyncInterval(config, m.checkerManager.GetGlobalInterval()) + checker := NewClientConnectionChecker(m, config.ID, syncInterval, isPingAvailable, m.logger) + m.checkerManager.StartChecking(checker) + } return nil } @@ -798,7 +836,7 @@ func (m *MCPManager) VerifyPerUserOAuthConnection(ctx context.Context, config *s // provided user-submitted header values to verify the server is reachable // and discover available tools. The connection is closed after verification. // -// Used in two paths: +// Used in three paths: // - Admin test flow: admin enters sample values during MCP client creation, // this runs an Initialize handshake against the upstream to validate the // schema (PerUserHeaderKeys) + discover tools. The discovered tools then @@ -807,11 +845,16 @@ func (m *MCPManager) VerifyPerUserOAuthConnection(ctx context.Context, config *s // workspace submit URL surfaced inline by MCPAuthRequiredError. The // handler runs this before upserting the row so a bad submission returns // 422 immediately instead of failing on the next tool call. +// - Shared headers/none per-call discovery: performAdminToolDiscovery's +// periodic refresh for a client running per-call (needs_session_stickiness +// nil/false). userHeaders here is whatever AdminConnectionHeaders +// resolved, which may legitimately be empty for these auth types. // // Parameters: // - config: MCP client configuration (connection URL, name, PerUserHeaderKeys, etc.) -// - userHeaders: caller-supplied header_name → value map (must cover every -// PerUserHeaderKeys entry; the caller validates that before invoking). +// - userHeaders: caller-supplied header_name → value map. Must cover every +// PerUserHeaderKeys entry for auth_type=per_user_headers (the caller +// validates that before invoking); may be empty for shared headers/none. // // Returns: // - map[string]schemas.ChatTool: discovered tools keyed by prefixed name @@ -821,7 +864,15 @@ func (m *MCPManager) VerifyHeadersConnection(ctx context.Context, config *schema if config.ConnectionString == nil || config.ConnectionString.GetValue() == "" { return nil, nil, fmt.Errorf("connection URL is required for per-user headers verification") } - if len(userHeaders) == 0 { + // Non-empty userHeaders is only required for genuinely per-user auth: + // PerUserHeaderKeys declares a schema of headers every caller must + // supply, so an empty map there is always a caller bug. A shared + // headers/none client (performAdminToolDiscovery's per-call discovery + // path, added for the needs_session_stickiness feature) has no such + // schema — it may legitimately have no admin-configured Authorization + // header at all, with auth carried entirely by other static config + // headers (already layered in below) or none. + if config.AuthType == schemas.MCPAuthTypePerUserHeaders && len(userHeaders) == 0 { return nil, nil, fmt.Errorf("user headers are required for per-user headers verification") } @@ -941,34 +992,39 @@ func (m *MCPManager) VerifyHeadersConnection(ctx context.Context, config *schema return tools, toolNameMapping, nil } -// performAdminToolDiscovery resolves the retained admin bootstrap credential -// for a per_user_oauth/per_user_headers/token_exchange client (via -// credStore.AdminConnectionHeaders) and runs a one-shot ephemeral +// performAdminToolDiscovery resolves the credential for a per-call client +// (via credStore.AdminConnectionHeaders) and runs a one-shot ephemeral // connect-discover-close cycle, reusing the exact same verification // functions the one-time bootstrap flow itself uses // (VerifyPerUserOAuthConnection / VerifyHeadersConnection) — this is the // periodic connection checker's per-call counterpart to reusing a live Conn // for sticky clients. +// +// Covers two families: +// - Per-user (per_user_oauth, per_user_headers, token_exchange): always +// per-call, resolves the retained admin bootstrap credential — a +// separate credential from any individual end user's own session. +// - Shared (oauth, headers, none) running per-call via +// needs_session_stickiness nil/false: resolves the same credential a +// real tool call would (see each resolver's AdminConnectionHeaders — +// there is no separate "admin" concept for these types). func (m *MCPManager) performAdminToolDiscovery(ctx context.Context, config *schemas.MCPClientConfig) (map[string]schemas.ChatTool, map[string]string, error) { - if config.AuthType != schemas.MCPAuthTypePerUserOauth && config.AuthType != schemas.MCPAuthTypePerUserHeaders && config.AuthType != schemas.MCPAuthTypeTokenExchange { - return nil, nil, fmt.Errorf("admin tool discovery not supported for auth_type %q", config.AuthType) - } - headers, err := m.credStore.AdminConnectionHeaders(ctx, config) if err != nil { return nil, nil, fmt.Errorf("failed to resolve admin credential: %w", err) } switch config.AuthType { - case schemas.MCPAuthTypePerUserOauth, schemas.MCPAuthTypeTokenExchange: - // Both resolve to a bearer credential (the retained bootstrap token, - // or a client-credentials token for token exchange), so they share - // the bearer-based one-shot verification path. + case schemas.MCPAuthTypePerUserOauth, schemas.MCPAuthTypeTokenExchange, schemas.MCPAuthTypeOauth: + // All three resolve to a bearer credential (the retained bootstrap + // token, a client-credentials token for token exchange, or the + // shared oauth token), so they share the bearer-based one-shot + // verification path. accessToken := strings.TrimPrefix(headers.Get("Authorization"), "Bearer ") if accessToken == "" { return nil, nil, fmt.Errorf("admin credential resolved no access token") } return m.VerifyPerUserOAuthConnection(ctx, config, accessToken) - case schemas.MCPAuthTypePerUserHeaders: + case schemas.MCPAuthTypePerUserHeaders, schemas.MCPAuthTypeHeaders, schemas.MCPAuthTypeNone: userHeaders := make(map[string]string, len(headers)) for k := range headers { userHeaders[k] = headers.Get(k) @@ -992,9 +1048,7 @@ func (m *MCPManager) SetClientTools(clientID string, tools map[string]schemas.Ch defer m.mu.Unlock() if client, exists := m.clientMap[clientID]; exists { - for toolName, tool := range tools { - client.ToolMap[toolName] = tool - } + maps.Copy(client.ToolMap, tools) client.ToolNameMapping = toolNameMapping client.State = schemas.MCPConnectionStateHealthy m.logger.Debug("%s Set %d tools on client '%s'", MCPLogPrefix, len(tools), client.Name) @@ -1476,7 +1530,7 @@ func (m *MCPManager) UpdateClient(id string, updatedConfig *schemas.MCPClientCon return nil } -// UpdateClientConnection updates auth-related fields (headers) for an existing MCP client by +// UpdateClientCredentials updates auth-related fields (headers) for an existing MCP client by // closing the current connection and establishing a new one so the new credentials are verified // before being committed. Non-credential metadata (name, tools, etc.) is preserved from the // current execution config. @@ -1492,7 +1546,7 @@ func (m *MCPManager) UpdateClient(id string, updatedConfig *schemas.MCPClientCon // // Returns: // - error: Any connection error; nil on success -func (m *MCPManager) UpdateClientConnection(id string, newConfig *schemas.MCPClientConfig) (retErr error) { +func (m *MCPManager) UpdateClientCredentials(id string, newConfig *schemas.MCPClientConfig) (retErr error) { if newConfig == nil { return fmt.Errorf("newConfig must not be nil") } @@ -1511,10 +1565,74 @@ func (m *MCPManager) UpdateClientConnection(id string, newConfig *schemas.MCPCli m.mu.RUnlock() return fmt.Errorf("client %s not found", id) } - // Per-user auth clients have no persistent connection — reconnect/update is not applicable. + // Per-call clients have no persistent connection to update — this covers + // both genuine per-user auth types (always per-call) and a shared HTTP + // client (oauth/headers) currently running in per-call mode via + // needs_session_stickiness nil/false. Either way connectToMCPClient below + // (a real persistent dial) does not apply. if client.ExecutionConfig != nil && m.credStore.RequiresPerCallConnection(client.ExecutionConfig) { + // A client still in PendingVerification is completing its FIRST + // connection here — e.g. a config.json-bootstrapped OAuth client + // whose admin just finished the browser consent flow. AddClient's + // own per-call branch (see its PendingOAuthConfig early-return) is + // what would normally perform this exact state transition, but it + // never gets a second chance to run for an already-registered + // client: the entry was created once, at boot, and parked here. + // Mirror that same transition now, or this client is stuck showing + // pending_verification — with no way to retry (the pending OAuth + // stash this endpoint requires is already cleared) — until a + // restart re-runs AddClient from scratch and does it correctly. + if client.State == schemas.MCPConnectionStatePendingVerification { + m.mu.RUnlock() + m.mu.Lock() + if cs, exists := m.clientMap[id]; exists && cs.State == schemas.MCPConnectionStatePendingVerification { + cs.ExecutionConfig = newConfig + if len(newConfig.DiscoveredTools) > 0 { + maps.Copy(cs.ToolMap, newConfig.DiscoveredTools) + cs.ToolNameMapping = newConfig.DiscoveredToolNameMapping + } + cs.State = schemas.MCPConnectionStateHealthy + } + m.mu.Unlock() + // A client with nothing to restore gets its first discovery pass + // synchronously here — see AddClient's per-call branch (mirrors + // its own identical comment). Without this, a shared + // oauth/headers/none client completing its first connection here + // would sit at Healthy/0-tools until the periodic checker's + // first tick, which could be minutes away (Start() uses the slow + // healthyInterval, not the fast Unstable one, for an + // already-Healthy client). Best-effort: a failure here just + // means the periodic checker's own retries pick it up shortly. + if len(newConfig.DiscoveredTools) == 0 { + if tools, mapping, discErr := m.performAdminToolDiscovery(m.ctx, newConfig); discErr != nil { + m.logger.Debug("%s Initial per-call tool discovery failed for MCP client '%s': %v — the periodic checker will retry", MCPLogPrefix, newConfig.Name, discErr) + } else { + m.SetClientTools(id, tools, mapping) + } + } + // Start the connection checker so this client's tools stay + // current going forward — mirrors AddClient's per-call branch + // (see its own comment): without this, nothing else ever + // revisits this client. + if id != BifrostMCPClientKey { + isPingAvailable := true + if newConfig.IsPingAvailable != nil { + isPingAvailable = *newConfig.IsPingAvailable + } + syncInterval := ResolveToolSyncInterval(newConfig, m.checkerManager.GetGlobalInterval()) + checker := NewClientConnectionChecker(m, id, syncInterval, isPingAvailable, m.logger) + m.checkerManager.StartChecking(checker) + } + return nil + } + // Already past that first connection (a plain reauthorize of an + // already-Healthy client): a per-call client dials fresh on every + // call and already picks up whatever credential is current in the + // DB, so there is genuinely nothing left to do — that's a no-op, + // not a failure, hence the shared "not applicable" sentinel rather + // than an opaque error. m.mu.RUnlock() - return fmt.Errorf("connection update is not supported for per-user auth clients") + return fmt.Errorf("client uses per-call connections; there is no persistent connection to update: %w", schemas.ErrMCPReconnectNotApplicable) } if client.ExecutionConfig == nil { m.mu.RUnlock() diff --git a/core/mcp/clientmanager_test.go b/core/mcp/clientmanager_test.go index 9a0c251ff7..d139d2c242 100644 --- a/core/mcp/clientmanager_test.go +++ b/core/mcp/clientmanager_test.go @@ -273,3 +273,193 @@ func TestEnableClient_GuardLostLeavesDisabledUntouched(t *testing.T) { assert.True(t, state.ExecutionConfig.Disabled, "a failed EnableClient (guard already held elsewhere) must not have flipped Disabled — otherwise a concurrent connectToMCPClient could dial the still-Disabled-State entry") assert.Equal(t, schemas.MCPConnectionStateDisabled, state.State) } + +// TestUpdateClientCredentials_PerCallSharedOAuth_ReturnsReconnectNotApplicable +// pins the bug reported against a shared-OAuth client (auth_type='oauth') +// running in per-call mode: needs_session_stickiness nil/false — the default +// for a newly created or config.json-bootstrapped client that predates this +// field — routes RequiresPerCallConnection through the same per-call branch +// a genuine per-user client uses. UpdateClientCredentials's guard used to +// return a bare "per-user auth clients" error for this case even though the +// client is genuinely shared, just with no persistent connection to +// reconnect (the fresh OAuth token is already live for the next per-call +// dial). Callers must be able to tell this "nothing to do" case apart from a +// real failure via the same ErrMCPReconnectNotApplicable sentinel +// ReconnectClient/CloseAndMarkNeedsReauth already use for the analogous case. +func TestUpdateClientCredentials_PerCallSharedOAuth_ReturnsReconnectNotApplicable(t *testing.T) { + // nil credStore: NewMCPManager defaults to a real credstore.CredStore, + // whose RequiresPerCallConnection actually ANDs auth type with + // NeedsSessionStickiness — unlike expiredOAuthCredStore (used elsewhere + // in this file), which hardcodes false regardless of either. + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, nil, nil) + config := &schemas.MCPClientConfig{ + ID: "client-shared-percall", + Name: "shared-percall-client", + AuthType: schemas.MCPAuthTypeOauth, + ConnectionType: schemas.MCPConnectionTypeHTTP, + // NeedsSessionStickiness left nil on purpose: the default value. + } + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStateHealthy, + } + m.mu.Unlock() + + err := m.UpdateClientCredentials(config.ID, config) + require.Error(t, err) + assert.True(t, errors.Is(err, schemas.ErrMCPReconnectNotApplicable), "a per-call shared-oauth client must report the same not-applicable sentinel as a genuine per-user client, not an opaque failure") +} + +// TestUpdateClientCredentials_PerCallSharedOAuth_PendingVerification_TransitionsToHealthy +// pins the second half of the same bug report: a per-call shared-OAuth +// client's FIRST connection (still in PendingVerification — e.g. a +// config.json-bootstrapped client whose admin just completed the OAuth +// browser flow) must not be treated as "nothing to do" the way an +// already-Healthy per-call client's reauthorize is. Skipping the transition +// left the client permanently stuck showing pending_verification in the UI +// (with no way to retry, since the pending OAuth stash this endpoint +// requires is already cleared by the caller) until a restart re-ran +// AddClient from scratch and performed the transition correctly. +func TestUpdateClientCredentials_PerCallSharedOAuth_PendingVerification_TransitionsToHealthy(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, nil, nil) + config := &schemas.MCPClientConfig{ + ID: "client-pending-percall", + Name: "pending-percall-client", + AuthType: schemas.MCPAuthTypeOauth, + ConnectionType: schemas.MCPConnectionTypeHTTP, + } + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStatePendingVerification, + ToolMap: make(map[string]schemas.ChatTool), + ToolNameMapping: make(map[string]string), + } + m.mu.Unlock() + + err := m.UpdateClientCredentials(config.ID, config) + require.NoError(t, err, "the first connection out of pending_verification must succeed, not report not-applicable") + + m.mu.RLock() + state := *m.clientMap[config.ID] + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStateHealthy, state.State, "must transition out of pending_verification, mirroring AddClient's own per-call setup") + + // Without a connection checker, nothing would ever discover this + // client's tools afterward (no persistent Conn, no DiscoveredTools to + // restore for a client completing its first connection here) — see + // TestAddClient_PerCallConnection_StartsConnectionChecker for the + // AddClient-side counterpart of this same fix. + m.checkerManager.mu.RLock() + _, hasChecker := m.checkerManager.checkers[config.ID] + m.checkerManager.mu.RUnlock() + assert.True(t, hasChecker, "must start a connection checker so tools actually get discovered") + + // A second call, now that the client is already Healthy, is the plain + // reauthorize case: genuinely nothing left to do. + err = m.UpdateClientCredentials(config.ID, config) + require.Error(t, err) + assert.True(t, errors.Is(err, schemas.ErrMCPReconnectNotApplicable)) +} + +// TestUpdateClientCredentials_PerCallSharedType_PendingVerification_DiscoversToolsSynchronously +// pins the second half of the same bug report: a connection checker alone +// is not enough, because ClientConnectionChecker.Start uses the SLOW +// healthyInterval (not the fast Unstable one) for a client that already +// starts Healthy — so without a synchronous first discovery pass, a shared +// oauth/headers/none per-call client completing OAuth/verification here +// would sit at Healthy/0-tools for up to healthyInterval before its first +// real check ever ran. +func TestUpdateClientCredentials_PerCallSharedType_PendingVerification_DiscoversToolsSynchronously(t *testing.T) { + ts, _ := buildAdminDiscoveryHTTPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, nil, nil) + config := &schemas.MCPClientConfig{ + ID: "client-pending-percall-sync", + Name: "pending-percall-client-sync", + AuthType: schemas.MCPAuthTypeHeaders, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar(ts.URL), + Headers: map[string]schemas.SecretVar{"Authorization": *schemas.NewSecretVar("Bearer shared-update-token")}, + } + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStatePendingVerification, + ToolMap: make(map[string]schemas.ChatTool), + ToolNameMapping: make(map[string]string), + } + m.mu.Unlock() + + err := m.UpdateClientCredentials(config.ID, config) + require.NoError(t, err) + + m.mu.RLock() + state := *m.clientMap[config.ID] + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStateHealthy, state.State) + assert.Contains(t, state.ToolMap, "pending-percall-client-sync-echo", "tools must be discovered synchronously, not deferred entirely to the periodic checker") +} + +// TestReconnectClient_PerCallSharedOAuth_ReturnsReconnectNotApplicable pins +// the same message-wording bug as TestUpdateClientCredentials's sibling: +// ReconnectClient's per-call guard used to unconditionally say "per-user +// auth clients", even though RequiresPerCallConnection is also true for a +// genuinely shared oauth/headers client running per-call via +// needs_session_stickiness nil/false. The sentinel was always correct; only +// the message text was misleading for this case. +func TestReconnectClient_PerCallSharedOAuth_ReturnsReconnectNotApplicable(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, nil, nil) + config := &schemas.MCPClientConfig{ + ID: "client-shared-percall-reconnect", + Name: "shared-percall-client-reconnect", + AuthType: schemas.MCPAuthTypeOauth, + ConnectionType: schemas.MCPConnectionTypeHTTP, + // NeedsSessionStickiness left nil on purpose: the default value. + } + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStateHealthy, + } + m.mu.Unlock() + + err := m.ReconnectClient(config.ID) + require.Error(t, err) + assert.True(t, errors.Is(err, schemas.ErrMCPReconnectNotApplicable)) + assert.NotContains(t, err.Error(), "per-user", "a genuinely shared client must not be told it's a per-user auth client") +} + +// TestCloseAndMarkNeedsReauth_PerCallSharedOAuth_ReturnsReconnectNotApplicable +// is the CloseAndMarkNeedsReauth counterpart of the same message-wording fix. +func TestCloseAndMarkNeedsReauth_PerCallSharedOAuth_ReturnsReconnectNotApplicable(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, nil, nil) + config := &schemas.MCPClientConfig{ + ID: "client-shared-percall-close", + Name: "shared-percall-client-close", + AuthType: schemas.MCPAuthTypeOauth, + ConnectionType: schemas.MCPConnectionTypeHTTP, + } + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStateHealthy, + } + m.mu.Unlock() + + err := m.CloseAndMarkNeedsReauth(config.ID) + require.Error(t, err) + assert.True(t, errors.Is(err, schemas.ErrMCPReconnectNotApplicable)) + assert.NotContains(t, err.Error(), "per-user", "a genuinely shared client must not be told it's a per-user auth client") +} diff --git a/core/mcp/connectionchecker.go b/core/mcp/connectionchecker.go index 858156304b..51cbf434ff 100644 --- a/core/mcp/connectionchecker.go +++ b/core/mcp/connectionchecker.go @@ -205,7 +205,7 @@ func (c *ClientConnectionChecker) performCheck() time.Duration { // attempt for nothing. Keep the timer alive at the relaxed // interval so a stalled reauthorize eventually gets picked up by // something, but do no work — only an explicit reauthorize (or a - // direct UpdateClientConnection success) moves this out. + // direct UpdateClientCredentials success) moves this out. return c.healthyInterval } if config == nil { diff --git a/core/mcp/connectionchecker_test.go b/core/mcp/connectionchecker_test.go index 4207e565a2..999e95ff63 100644 --- a/core/mcp/connectionchecker_test.go +++ b/core/mcp/connectionchecker_test.go @@ -89,21 +89,27 @@ func TestPerformCheck_NilConn_PerUserHeaders_AttemptsAdminDiscovery(t *testing.T require.Equal(t, existingTools, manager.clientMap[config.ID].ToolMap) } -// TestPerformCheck_NilConn_SharedAuthType_NeverAttemptsAdminDiscovery is the -// regression-risk case: a shared-connection client (e.g. "headers" or -// "none") with Conn==nil is mid-(re)connect, not per-user. performCheck must -// never route it into admin discovery, which would be meaningless (and would -// call a credStore method shared resolvers don't meaningfully implement) — -// this branch instead attempts a direct reconnect, covered separately by -// TestPerformCheck_NilConn_SharedAuthType_TriggersReconnect below. -func TestPerformCheck_NilConn_SharedAuthType_NeverAttemptsAdminDiscovery(t *testing.T) { +// TestPerformCheck_NilConn_SharedAuthType_Sticky_NeverAttemptsAdminDiscovery +// is the regression-risk case for a genuinely sticky shared client (e.g. +// "headers" or "none") with Conn==nil because it's mid-(re)connect, not +// per-call. performCheck must never route it into admin discovery — this +// branch instead attempts a direct reconnect, covered separately by +// TestPerformCheck_NilConn_SharedAuthType_TriggersReconnect below. Uses the +// real credstore.CredStore (not a hardcoded-per-call fake) so the routing +// decision reflects actual NeedsSessionStickiness semantics: a shared type +// with needs_session_stickiness=false is legitimately per-call and DOES +// attempt admin discovery now — see +// TestPerformCheck_NilConn_SharedAuthType_PerCall_SuccessUpdatesToolMap. +func TestPerformCheck_NilConn_SharedAuthType_Sticky_NeverAttemptsAdminDiscovery(t *testing.T) { for _, authType := range []schemas.MCPAuthType{schemas.MCPAuthTypeHeaders, schemas.MCPAuthTypeNone, schemas.MCPAuthTypeOauth} { t.Run(string(authType), func(t *testing.T) { - cred := &fakeAdminCredStore{err: errors.New("must never be called for shared auth types")} - manager := NewMCPManager(context.Background(), schemas.MCPConfig{}, cred, &MockLogger{}, nil) + manager := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) existingTools := map[string]schemas.ChatTool{"shared-client-existing": {}} state, config := newSyncTestClientState("client-3", "shared-client", authType, existingTools) + config.ConnectionType = schemas.MCPConnectionTypeHTTP + config.NeedsSessionStickiness = schemas.Ptr(true) + config.ConnectionString = schemas.NewSecretVar("http://127.0.0.1:0/mcp") // unreachable — the reconnect attempt fails fast, which is fine manager.mu.Lock() manager.clientMap[config.ID] = state manager.mu.Unlock() @@ -111,17 +117,50 @@ func TestPerformCheck_NilConn_SharedAuthType_NeverAttemptsAdminDiscovery(t *test checker := NewClientConnectionChecker(manager, config.ID, time.Minute, false, &MockLogger{}) checker.performCheck() - // performCheck's default branch spawns the actual reconnect in a - // background goroutine, which mutates clientMap under manager.mu — - // read it back under the same lock rather than racing that goroutine. + require.Eventually(t, func() bool { + _, inFlightOrDone := manager.reconnectingClients.Load(config.ID) + return inFlightOrDone + }, 2*time.Second, 10*time.Millisecond, "a sticky client with no live connection must trigger a direct reconnect attempt, not admin discovery") + manager.mu.RLock() toolMap := manager.clientMap[config.ID].ToolMap - _, stillExists := manager.clientMap[config.ID] manager.mu.RUnlock() - - require.Equal(t, 0, cred.callCount(), "shared clients mid-reconnect (Conn==nil) must never be routed into admin discovery") require.Equal(t, existingTools, toolMap, "must leave the existing tool map untouched — only a real reconnect's own tools/list replaces it") - require.True(t, stillExists) + }) + } +} + +// TestPerformCheck_NilConn_SharedAuthType_PerCall_SuccessUpdatesToolMap pins +// the fix for the shared-per-call tool-discovery bug at the full +// performCheck level (not just performAdminToolDiscovery in isolation): a +// shared oauth/headers/none client running per-call +// (needs_session_stickiness nil/false) must have its tool map populated by +// the periodic checker, exactly like a per_user_oauth client already does +// (TestPerformCheck_NilConn_PerUserOAuth_SuccessUpdatesToolMap). +func TestPerformCheck_NilConn_SharedAuthType_PerCall_SuccessUpdatesToolMap(t *testing.T) { + for _, authType := range []schemas.MCPAuthType{schemas.MCPAuthTypeHeaders, schemas.MCPAuthTypeOauth} { + t.Run(string(authType), func(t *testing.T) { + ts, _ := buildAdminDiscoveryHTTPServer(t) + + cred := &fakeAdminCredStore{headers: http.Header{"Authorization": []string{"Bearer shared-token"}}} + manager := &MCPManager{ + credStore: cred, + logger: &MockLogger{}, + clientMap: map[string]*schemas.MCPClientState{}, + } + + state, config := newSyncTestClientState("client-shared-percall", "shared-client", authType, map[string]schemas.ChatTool{}) + config.ConnectionType = schemas.MCPConnectionTypeHTTP + config.ConnectionString = schemas.NewSecretVar(ts.URL) + manager.clientMap[config.ID] = state + + checker := NewClientConnectionChecker(manager, config.ID, time.Minute, false, &MockLogger{}) + checker.performCheck() + + require.Equal(t, 1, cred.callCount()) + updated := manager.clientMap[config.ID].ToolMap + require.Contains(t, updated, "shared-client-echo", "a successful admin-discovery cycle should populate the client's tool map") + require.Equal(t, schemas.MCPConnectionStateHealthy, manager.clientMap[config.ID].State) }) } } diff --git a/core/mcp/credstore/none.go b/core/mcp/credstore/none.go index 1ac1d816e0..3722d1d594 100644 --- a/core/mcp/credstore/none.go +++ b/core/mcp/credstore/none.go @@ -25,8 +25,18 @@ func (r *noneResolver) ForceRefresh(_ *schemas.BifrostContext, _ *schemas.MCPCli return nil } -// AdminConnectionHeaders is not supported for this auth type — there is no -// separate "admin" credential distinct from the one used for real calls. +// AdminConnectionHeaders delegates to ConnectionHeaders for a client +// currently running per-call (needs_session_stickiness nil/false) — there is +// no separate "admin" credential distinct from the one used for real calls +// (both are empty), so the periodic connection checker's per-call discovery +// cycle (MCPManager.performAdminToolDiscovery) can still run for a +// none-auth client without an error. A sticky client should never reach +// this method at all (it holds a persistent connection instead, discovered +// via connectToMCPClient) — erroring rather than silently delegating turns +// a caller-chain bug into an immediate, loud failure. func (r *noneResolver) AdminConnectionHeaders(ctx context.Context, config *schemas.MCPClientConfig) (http.Header, error) { - return nil, fmt.Errorf("admin connection headers not supported for auth_type %q", "none") + if !needsSessionStickiness(config) { + return r.ConnectionHeaders(schemas.NewBifrostContext(ctx, schemas.NoDeadline), config) + } + return nil, fmt.Errorf("admin connection headers not supported for auth_type %q on a sticky connection", "none") } diff --git a/core/mcp/credstore/none_test.go b/core/mcp/credstore/none_test.go index 08effe8977..6976b2b3f9 100644 --- a/core/mcp/credstore/none_test.go +++ b/core/mcp/credstore/none_test.go @@ -7,20 +7,45 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -// TestNoneResolverAdminConnectionHeadersReturnsError confirms noneResolver's -// stub AdminConnectionHeaders returns a plain (non-nil, non-panicking) error -// rather than ever being reachable in production — auth_type "none" has no -// separate admin credential, so the periodic tool syncer's per-user -// discovery path (see ClientToolSyncer.performSync) never dispatches here. -func TestNoneResolverAdminConnectionHeadersReturnsError(t *testing.T) { +// TestNoneResolverAdminConnectionHeaders_StickyReturnsError confirms +// AdminConnectionHeaders still refuses to resolve anything for a sticky +// client (the default when NeedsSessionStickiness/ConnectionType aren't +// set) — a sticky client holds a persistent connection discovered via +// connectToMCPClient, and should never reach this method at all. +func TestNoneResolverAdminConnectionHeaders_StickyReturnsError(t *testing.T) { resolver := &noneResolver{} config := &schemas.MCPClientConfig{ID: "client-1", Name: "Test Client"} headers, err := resolver.AdminConnectionHeaders(context.Background(), config) if err == nil { - t.Fatal("expected a non-nil error for auth_type \"none\"") + t.Fatal("expected a non-nil error for a sticky client") } if headers != nil { t.Errorf("expected nil headers alongside the error, got %v", headers) } } + +// TestNoneResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders +// pins the fix for the none-auth-per-call tool-discovery bug: a client +// running per-call (needs_session_stickiness nil/false on an http +// connection) must resolve successfully (empty headers, same as +// ConnectionHeaders) rather than erroring — so the periodic connection +// checker's per-call discovery cycle +// (MCPManager.performAdminToolDiscovery) can actually run for it. +func TestNoneResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders(t *testing.T) { + resolver := &noneResolver{} + config := &schemas.MCPClientConfig{ + ID: "client-1", + Name: "Test Client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + // NeedsSessionStickiness left nil: the default per-call value. + } + + headers, err := resolver.AdminConnectionHeaders(context.Background(), config) + if err != nil { + t.Fatalf("expected no error for a per-call client, got: %v", err) + } + if len(headers) != 0 { + t.Errorf("expected empty headers, got %v", headers) + } +} diff --git a/core/mcp/credstore/shared_headers.go b/core/mcp/credstore/shared_headers.go index bb151f1114..e87ed23b4c 100644 --- a/core/mcp/credstore/shared_headers.go +++ b/core/mcp/credstore/shared_headers.go @@ -43,8 +43,21 @@ func (r *sharedHeadersResolver) ForceRefresh(_ *schemas.BifrostContext, _ *schem return nil } -// AdminConnectionHeaders is not supported for this auth type — there is no -// separate "admin" credential distinct from the one used for real calls. +// AdminConnectionHeaders delegates to ConnectionHeaders for a client +// currently running per-call (needs_session_stickiness nil/false) — there is +// no separate "admin" credential distinct from the one used for real calls, +// so the periodic connection checker's per-call discovery cycle +// (MCPManager.performAdminToolDiscovery) resolves exactly the same +// Authorization header a real tool call would. Static (non-Authorization) +// config headers are layered separately by VerifyHeadersConnection, same as +// the normal call path. A sticky client should never reach this method at +// all (it holds a persistent connection instead, discovered via +// connectToMCPClient) — erroring rather than silently delegating turns a +// caller-chain bug into an immediate, loud failure instead of a +// quietly-redundant credential resolution. func (r *sharedHeadersResolver) AdminConnectionHeaders(ctx context.Context, config *schemas.MCPClientConfig) (http.Header, error) { - return nil, fmt.Errorf("admin connection headers not supported for auth_type %q", "headers") + if !needsSessionStickiness(config) { + return r.ConnectionHeaders(schemas.NewBifrostContext(ctx, schemas.NoDeadline), config) + } + return nil, fmt.Errorf("admin connection headers not supported for auth_type %q on a sticky connection", "headers") } diff --git a/core/mcp/credstore/shared_headers_test.go b/core/mcp/credstore/shared_headers_test.go index 5d6c78f544..f425ebfa98 100644 --- a/core/mcp/credstore/shared_headers_test.go +++ b/core/mcp/credstore/shared_headers_test.go @@ -7,21 +7,69 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -// TestSharedHeadersResolverAdminConnectionHeadersReturnsError confirms -// sharedHeadersResolver's stub AdminConnectionHeaders returns a plain -// (non-nil, non-panicking) error — auth_type "headers" has no separate admin -// credential distinct from the shared one, so the periodic tool syncer's -// per-user discovery path (see ClientToolSyncer.performSync) never -// dispatches here. -func TestSharedHeadersResolverAdminConnectionHeadersReturnsError(t *testing.T) { +// TestSharedHeadersResolverAdminConnectionHeaders_StickyReturnsError +// confirms AdminConnectionHeaders still refuses to resolve anything for a +// sticky client (the default when NeedsSessionStickiness/ConnectionType +// aren't set) — a sticky client holds a persistent connection discovered +// via connectToMCPClient, and should never reach this method at all. +func TestSharedHeadersResolverAdminConnectionHeaders_StickyReturnsError(t *testing.T) { resolver := &sharedHeadersResolver{} config := &schemas.MCPClientConfig{ID: "client-1", Name: "Test Client"} headers, err := resolver.AdminConnectionHeaders(context.Background(), config) if err == nil { - t.Fatal("expected a non-nil error for auth_type \"headers\"") + t.Fatal("expected a non-nil error for a sticky client") } if headers != nil { t.Errorf("expected nil headers alongside the error, got %v", headers) } } + +// TestSharedHeadersResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders +// pins the fix for the shared-headers-per-call tool-discovery bug: a client +// running per-call (needs_session_stickiness nil/false on an http +// connection) must resolve the same Authorization header ConnectionHeaders +// would — there is no separate "admin" credential for this auth type — so +// the periodic connection checker's per-call discovery cycle +// (MCPManager.performAdminToolDiscovery) can actually populate its tool map. +func TestSharedHeadersResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders(t *testing.T) { + resolver := &sharedHeadersResolver{} + config := &schemas.MCPClientConfig{ + ID: "client-1", + Name: "Test Client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + Headers: map[string]schemas.SecretVar{"Authorization": *schemas.NewSecretVar("Bearer static-token")}, + // NeedsSessionStickiness left nil: the default per-call value. + } + + headers, err := resolver.AdminConnectionHeaders(context.Background(), config) + if err != nil { + t.Fatalf("expected no error for a per-call client, got: %v", err) + } + if got := headers.Get("Authorization"); got != "Bearer static-token" { + t.Errorf("expected the resolved Authorization header, got %q", got) + } +} + +// TestSharedHeadersResolverAdminConnectionHeaders_PerCallNoAuthorizationHeaderIsFine +// covers the case CodeRabbit-adjacent reasoning flagged: a shared client +// with no Authorization header configured at all (e.g. auth is carried +// entirely via other static headers, or none) must resolve to empty +// headers, not an error — VerifyHeadersConnection's separate "userHeaders +// required" guard is what needs relaxing for this, not this method. +func TestSharedHeadersResolverAdminConnectionHeaders_PerCallNoAuthorizationHeaderIsFine(t *testing.T) { + resolver := &sharedHeadersResolver{} + config := &schemas.MCPClientConfig{ + ID: "client-1", + Name: "Test Client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + } + + headers, err := resolver.AdminConnectionHeaders(context.Background(), config) + if err != nil { + t.Fatalf("expected no error for a per-call client with no Authorization header, got: %v", err) + } + if got := headers.Get("Authorization"); got != "" { + t.Errorf("expected no Authorization header, got %q", got) + } +} diff --git a/core/mcp/credstore/shared_oauth.go b/core/mcp/credstore/shared_oauth.go index 14bfa4773e..91f532b1dd 100644 --- a/core/mcp/credstore/shared_oauth.go +++ b/core/mcp/credstore/shared_oauth.go @@ -46,8 +46,19 @@ func (r *sharedOAuthResolver) ForceRefresh(ctx *schemas.BifrostContext, config * return r.provider.ForceRefreshAccessToken(ctx, config) } -// AdminConnectionHeaders is not supported for this auth type — there is no -// separate "admin" credential distinct from the one used for real calls. +// AdminConnectionHeaders delegates to ConnectionHeaders for a client +// currently running per-call (needs_session_stickiness nil/false) — there is +// no separate "admin" credential distinct from the one used for real calls, +// so the periodic connection checker's per-call discovery cycle +// (MCPManager.performAdminToolDiscovery) resolves exactly the same shared +// token a real tool call would. A sticky client should never reach this +// method at all (it holds a persistent connection instead, discovered via +// connectToMCPClient) — erroring rather than silently delegating turns a +// caller-chain bug into an immediate, loud failure instead of a +// quietly-redundant credential resolution. func (r *sharedOAuthResolver) AdminConnectionHeaders(ctx context.Context, config *schemas.MCPClientConfig) (http.Header, error) { - return nil, fmt.Errorf("admin connection headers not supported for auth_type %q", "oauth") + if !needsSessionStickiness(config) { + return r.ConnectionHeaders(schemas.NewBifrostContext(ctx, schemas.NoDeadline), config) + } + return nil, fmt.Errorf("admin connection headers not supported for auth_type %q on a sticky connection", "oauth") } diff --git a/core/mcp/credstore/shared_oauth_test.go b/core/mcp/credstore/shared_oauth_test.go index 7d427e8a30..a216659c6c 100644 --- a/core/mcp/credstore/shared_oauth_test.go +++ b/core/mcp/credstore/shared_oauth_test.go @@ -7,21 +7,83 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -// TestSharedOAuthResolverAdminConnectionHeadersReturnsError confirms -// sharedOAuthResolver's stub AdminConnectionHeaders returns a plain -// (non-nil, non-panicking) error — auth_type "oauth" has no separate admin -// credential distinct from the shared one, so the periodic tool syncer's -// per-user discovery path (see ClientToolSyncer.performSync) never -// dispatches here. -func TestSharedOAuthResolverAdminConnectionHeadersReturnsError(t *testing.T) { +// fakeSharedOAuthProvider implements schemas.OAuth2Provider with only +// GetAccessToken configurable — the only method sharedOAuthResolver's +// ConnectionHeaders calls. Every other method is unused by these tests. +type fakeSharedOAuthProvider struct { + accessToken string +} + +func (f *fakeSharedOAuthProvider) GetAccessToken(_ context.Context, _ string) (string, error) { + return f.accessToken, nil +} +func (f *fakeSharedOAuthProvider) GetAdminAccessToken(_ context.Context, _ string) (string, error) { + return "", nil +} +func (f *fakeSharedOAuthProvider) ValidateToken(_ context.Context, _ string) (bool, error) { + return false, nil +} +func (f *fakeSharedOAuthProvider) RevokeToken(_ context.Context, _ string) error { return nil } +func (f *fakeSharedOAuthProvider) GetUserAccessTokenByMode(_ context.Context, _ schemas.MCPAuthMode, _, _ string) (string, error) { + return "", nil +} +func (f *fakeSharedOAuthProvider) InitiateUserOAuthFlow(_ context.Context, _ string, _ string, _ string, _ schemas.MCPAuthMode) (*schemas.OAuth2FlowInitiation, string, error) { + return nil, "", nil +} +func (f *fakeSharedOAuthProvider) CompleteUserOAuthFlow(_ context.Context, _ string, _ string) (string, error) { + return "", nil +} +func (f *fakeSharedOAuthProvider) RefreshAccessToken(_ context.Context, _ string) error { return nil } +func (f *fakeSharedOAuthProvider) ForceRefreshAccessToken(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) error { + return nil +} +func (f *fakeSharedOAuthProvider) TokenExchangeAvailable() bool { return false } +func (f *fakeSharedOAuthProvider) GetExchangedAccessToken(_ *schemas.BifrostContext, _ *schemas.MCPClientConfig) (string, error) { + return "", nil +} + +// TestSharedOAuthResolverAdminConnectionHeaders_StickyReturnsError confirms +// AdminConnectionHeaders still refuses to resolve anything for a sticky +// client (the default when NeedsSessionStickiness/ConnectionType aren't +// set, matching needsSessionStickiness' own default-to-sticky behavior) — a +// sticky client holds a persistent connection discovered via +// connectToMCPClient, and should never reach this method at all. +func TestSharedOAuthResolverAdminConnectionHeaders_StickyReturnsError(t *testing.T) { resolver := &sharedOAuthResolver{} config := &schemas.MCPClientConfig{ID: "client-1", Name: "Test Client"} headers, err := resolver.AdminConnectionHeaders(context.Background(), config) if err == nil { - t.Fatal("expected a non-nil error for auth_type \"oauth\"") + t.Fatal("expected a non-nil error for a sticky client") } if headers != nil { t.Errorf("expected nil headers alongside the error, got %v", headers) } } + +// TestSharedOAuthResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders +// pins the fix for the shared-OAuth-per-call tool-discovery bug: a client +// running per-call (needs_session_stickiness nil/false on an http +// connection) must resolve the same bearer token ConnectionHeaders would — +// there is no separate "admin" credential for this auth type — so the +// periodic connection checker's per-call discovery cycle +// (MCPManager.performAdminToolDiscovery) can actually populate its tool map. +func TestSharedOAuthResolverAdminConnectionHeaders_PerCallDelegatesToConnectionHeaders(t *testing.T) { + resolver := &sharedOAuthResolver{provider: &fakeSharedOAuthProvider{accessToken: "shared-token-abc"}} + oauthConfigID := "oauth-config-1" + config := &schemas.MCPClientConfig{ + ID: "client-1", + Name: "Test Client", + ConnectionType: schemas.MCPConnectionTypeHTTP, + OauthConfigID: &oauthConfigID, + // NeedsSessionStickiness left nil: the default per-call value. + } + + headers, err := resolver.AdminConnectionHeaders(context.Background(), config) + if err != nil { + t.Fatalf("expected no error for a per-call client, got: %v", err) + } + if got := headers.Get("Authorization"); got != "Bearer shared-token-abc" { + t.Errorf("expected the resolved bearer token, got %q", got) + } +} diff --git a/core/mcp/interface.go b/core/mcp/interface.go index b672609d27..0579ec34de 100644 --- a/core/mcp/interface.go +++ b/core/mcp/interface.go @@ -80,9 +80,9 @@ type MCPManagerInterface interface { // UpdateClient updates an existing MCP client configuration UpdateClient(id string, updatedConfig *schemas.MCPClientConfig) error - // UpdateClientConnection reconnects an existing MCP client using updated + // UpdateClientCredentials reconnects an existing MCP client using updated // auth-related connection fields (for example, headers and OAuth config). - UpdateClientConnection(id string, newConfig *schemas.MCPClientConfig) error + UpdateClientCredentials(id string, newConfig *schemas.MCPClientConfig) error // ReconnectClient reconnects an MCP client by ID ReconnectClient(id string) error diff --git a/core/mcp/reauth_state_test.go b/core/mcp/reauth_state_test.go index 5a630fe882..6e9d591a78 100644 --- a/core/mcp/reauth_state_test.go +++ b/core/mcp/reauth_state_test.go @@ -65,7 +65,7 @@ func TestConnectToMCPClient_OAuth2TokenExpired_SetsNeedsReauth(t *testing.T) { // Confirmed precondition: only shared-connection auth types ( // RequiresPerCallConnection()==false) ever reach connectToMCPClient in - // the first place — AddClient/EnableClient/UpdateClientConnection all + // the first place — AddClient/EnableClient/UpdateClientCredentials 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. diff --git a/core/schemas/mcp.go b/core/schemas/mcp.go index 7c3e404f18..591ce157be 100644 --- a/core/schemas/mcp.go +++ b/core/schemas/mcp.go @@ -146,7 +146,7 @@ type MCPCredentialStore interface { // ConnectionHeaders returns the headers to attach when opening an upstream // transport. Called from two sites: // - // 1. At AddClient / Reconnect / UpdateClientConnection for shared- + // 1. At AddClient / Reconnect / UpdateClientCredentials for shared- // connection auth types (none, headers, server_oauth). The caller // wraps the Bifrost lifecycle context into a synthetic BifrostContext // with no identity, so the resolver returns admin-level headers diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index f0923fee4f..daaa8b9770 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -32,8 +32,8 @@ type MCPManager interface { AddMCPClient(ctx context.Context, clientConfig *schemas.MCPClientConfig) error RemoveMCPClient(ctx context.Context, id string) error UpdateMCPClient(ctx context.Context, id string, updatedConfig *schemas.MCPClientConfig) error - // UpdateMCPClientConnection reconnects an existing MCP client using updated headers - UpdateMCPClientConnection(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error + // UpdateMCPClientCredentials reconnects an existing MCP client using updated headers + UpdateMCPClientCredentials(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error ReconnectMCPClient(ctx context.Context, id string) error // CloseAndMarkNeedsReauth closes a shared client's live upstream // connection and flips it to needs_reauth, without attempting a new @@ -2813,21 +2813,28 @@ func (h *MCPHandler) updateMCPClientWithRetry(ctx context.Context, id string, co return lastErr } -// updateMCPClientConnectionWithRetry calls mcpManager.UpdateMCPClientConnection with a short retry loop. -func (h *MCPHandler) updateMCPClientConnectionWithRetry(ctx context.Context, id string, config *schemas.MCPClientConfig) error { +// updateMCPClientCredentialsWithRetry calls mcpManager.UpdateMCPClientCredentials with a short retry loop. +func (h *MCPHandler) updateMCPClientCredentialsWithRetry(ctx context.Context, id string, config *schemas.MCPClientConfig) error { const maxAttempts = 3 const retryDelay = 500 * time.Millisecond var lastErr error for attempt := 1; attempt <= maxAttempts; attempt++ { - lastErr = h.mcpManager.UpdateMCPClientConnection(ctx, id, config) + lastErr = h.mcpManager.UpdateMCPClientCredentials(ctx, id, config) if lastErr == nil { return nil } + // Not applicable (a per-call client — no persistent connection to + // update) is never transient: retrying it wouldn't change the + // client's connection mode, so return immediately rather than + // wasting the retry budget on a client type this can never apply to. + if errors.Is(lastErr, schemas.ErrMCPReconnectNotApplicable) { + return lastErr + } if !strings.Contains(lastErr.Error(), "reconnect") || attempt == maxAttempts { return lastErr } - logger.Warn(fmt.Sprintf("[OAuth Complete] UpdateMCPClientConnection attempt %d/%d for client %s blocked by in-flight reconnect; retrying in %s: %v", + logger.Warn(fmt.Sprintf("[OAuth Complete] UpdateMCPClientCredentials attempt %d/%d for client %s blocked by in-flight reconnect; retrying in %s: %v", attempt, maxAttempts, id, retryDelay, lastErr)) time.Sleep(retryDelay) } @@ -3072,11 +3079,14 @@ func (h *MCPHandler) completeMCPClientOAuth(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to load MCP client: %v", err)) return } - if err := h.updateMCPClientConnectionWithRetry(bifrostCtx, reauthClientConfig.ID, reauthClientConfig); err != nil { + if err := h.updateMCPClientCredentialsWithRetry(bifrostCtx, reauthClientConfig.ID, reauthClientConfig); err != nil && !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { logger.Error(fmt.Sprintf("Failed to reconnect MCP client after reauthorization for client %s: %v", reauthClientConfig.ID, err)) SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("OAuth credentials refreshed but reconnecting the client failed: %v", err)) return } + // ErrMCPReconnectNotApplicable (a per-call client — no persistent + // connection to reconnect) is not an error here: the fresh + // credential is already what the next per-call dial will use. SendJSON(ctx, map[string]any{"status": "success", "message": "MCP client re-authorized and reconnected successfully"}) return } @@ -3282,13 +3292,25 @@ func (h *MCPHandler) completeMCPClientOAuth(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to update MCP config: %v", err)) return } - if err := h.updateMCPClientConnectionWithRetry(bifrostCtx, mcpClientConfig.ID, mcpClientConfig); err != nil { - if rollbackErr := h.store.ConfigStore.UpdateMCPClientConfig(ctx, mcpClientConfig.ID, &oldDBConfig); rollbackErr != nil { - logger.Error(fmt.Sprintf("Failed to rollback MCP client DB update: %v. please restart bifrost to keep core and database in sync", rollbackErr)) + if err := h.updateMCPClientCredentialsWithRetry(bifrostCtx, mcpClientConfig.ID, mcpClientConfig); err != nil { + if errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { + // The client is running in per-call mode (no persistent + // connection — e.g. a shared client with + // needs_session_stickiness nil/false), so there is nothing + // to reconnect: the fresh OAuth credential just persisted + // above is already what the next per-call dial will use. + // Not an error — fall through to the normal success path + // below instead of rolling back the DB update that already + // succeeded. + logger.Debug(fmt.Sprintf("[OAuth Complete] Client %s uses per-call connections; nothing to reconnect after OAuth update", mcpClientConfig.ID)) + } else { + if rollbackErr := h.store.ConfigStore.UpdateMCPClientConfig(ctx, mcpClientConfig.ID, &oldDBConfig); rollbackErr != nil { + logger.Error(fmt.Sprintf("Failed to rollback MCP client DB update: %v. please restart bifrost to keep core and database in sync", rollbackErr)) + } + logger.Error(fmt.Sprintf("Failed to reconnect MCP client after OAuth DB update for client %s: %v", mcpClientConfig.ID, err)) + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to reconnect MCP client with updated OAuth credentials: %v", err)) + return } - logger.Error(fmt.Sprintf("Failed to reconnect MCP client after OAuth DB update for client %s: %v", mcpClientConfig.ID, err)) - SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to reconnect MCP client with updated OAuth credentials: %v", err)) - return } } else { if h.store.ConfigStore != nil { diff --git a/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go b/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go new file mode 100644 index 0000000000..f1e45ee837 --- /dev/null +++ b/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go @@ -0,0 +1,95 @@ +package handlers + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +// fakeUpdateCredsMCPManager is a minimal MCPManager test double: only +// UpdateMCPClientCredentials has real behavior (scripted per call, tracked by +// call count); every other method is an unused no-op, since +// updateMCPClientCredentialsWithRetry only ever calls that one method. +type fakeUpdateCredsMCPManager struct { + err error // returned on every call + callCount int +} + +func (m *fakeUpdateCredsMCPManager) UpdateMCPClientCredentials(_ context.Context, _ string, _ *schemas.MCPClientConfig) error { + m.callCount++ + return m.err +} + +func (m *fakeUpdateCredsMCPManager) AddMCPClient(_ context.Context, _ *schemas.MCPClientConfig) error { + return nil +} +func (m *fakeUpdateCredsMCPManager) RemoveMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeUpdateCredsMCPManager) UpdateMCPClient(_ context.Context, _ string, _ *schemas.MCPClientConfig) error { + return nil +} +func (m *fakeUpdateCredsMCPManager) ReconnectMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeUpdateCredsMCPManager) CloseAndMarkNeedsReauth(_ context.Context, _ string) error { + return nil +} +func (m *fakeUpdateCredsMCPManager) DisableMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeUpdateCredsMCPManager) EnableMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeUpdateCredsMCPManager) VerifyPerUserOAuthConnection(_ context.Context, _ *schemas.MCPClientConfig, _ string) (map[string]schemas.ChatTool, map[string]string, error) { + return nil, nil, nil +} +func (m *fakeUpdateCredsMCPManager) VerifyHeadersConnection(_ context.Context, _ *schemas.MCPClientConfig, _ map[string]string) (map[string]schemas.ChatTool, map[string]string, error) { + return nil, nil, nil +} +func (m *fakeUpdateCredsMCPManager) SetClientTools(_ string, _ map[string]schemas.ChatTool, _ map[string]string) { +} +func (m *fakeUpdateCredsMCPManager) RequiresPerCallConnection(_ *schemas.MCPClientConfig) bool { + return false +} + +// TestUpdateMCPClientCredentialsWithRetry_NotApplicable_ReturnsImmediately pins +// the fix for the shared-OAuth-per-call reconnect bug: ErrMCPReconnectNotApplicable +// (a per-call client — no persistent connection to update) is never transient, +// so it must be returned on the first attempt, not funneled into the +// "blocked by in-flight reconnect" retry loop — which the sentinel's own +// message text would otherwise match, since it contains "reconnect" and the +// loop's fallback classification is a bare substring check. +func TestUpdateMCPClientCredentialsWithRetry_NotApplicable_ReturnsImmediately(t *testing.T) { + mgr := &fakeUpdateCredsMCPManager{ + err: fmt.Errorf("client uses per-call connections; there is no persistent connection to update: %w", schemas.ErrMCPReconnectNotApplicable), + } + h := &MCPHandler{mcpManager: mgr} + + err := h.updateMCPClientCredentialsWithRetry(context.Background(), "client-1", &schemas.MCPClientConfig{}) + if err == nil { + t.Fatal("expected the not-applicable sentinel to be returned as an error") + } + if !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { + t.Fatalf("expected errors.Is match against ErrMCPReconnectNotApplicable, got: %v", err) + } + if mgr.callCount != 1 { + t.Errorf("expected exactly 1 call (no retry for a non-transient not-applicable error), got %d", mgr.callCount) + } +} + +// TestUpdateMCPClientCredentialsWithRetry_TransientReconnectError_StillRetries +// is a regression guard for the pre-existing retry behavior: an in-flight +// "reconnect already in progress" error (transient, unrelated to +// ErrMCPReconnectNotApplicable) must still go through the retry loop as +// before. +func TestUpdateMCPClientCredentialsWithRetry_TransientReconnectError_StillRetries(t *testing.T) { + SetLogger(&mockLogger{}) // the retry loop logs a warning per retry attempt + mgr := &fakeUpdateCredsMCPManager{ + err: errors.New("reconnect already in progress for this client"), + } + h := &MCPHandler{mcpManager: mgr} + + err := h.updateMCPClientCredentialsWithRetry(context.Background(), "client-1", &schemas.MCPClientConfig{}) + if err == nil { + t.Fatal("expected an error after exhausting retries") + } + if mgr.callCount != 3 { + t.Errorf("expected all 3 attempts to run for a transient 'reconnect' error, got %d", mgr.callCount) + } +} diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index 15c423be78..8de21600c3 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -6687,9 +6687,9 @@ func (c *Config) UpdateMCPClient(ctx context.Context, id string, updatedConfig * return nil } -// UpdateMCPClientConnection updates the auth credentials (headers) for an existing MCP client. +// UpdateMCPClientCredentials updates the auth credentials (headers) for an existing MCP client. // It delegates the actual reconnection (with the new credentials) to the Bifrost client. -func (c *Config) UpdateMCPClientConnection(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error { +func (c *Config) UpdateMCPClientCredentials(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error { if c.client == nil { return fmt.Errorf("bifrost client not set") } @@ -6716,7 +6716,7 @@ func (c *Config) UpdateMCPClientConnection(ctx context.Context, id string, newCo // Attempt the credential swap on the runtime side first. // If this fails, nothing in our in-memory config has changed. - if err := c.client.UpdateMCPClientConnection(id, newConfig); err != nil { + if err := c.client.UpdateMCPClientCredentials(id, newConfig); err != nil { return fmt.Errorf("failed to update MCP client credentials: %w", err) } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 1d8757e095..0de78c8d56 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -124,8 +124,8 @@ type ServerCallbacks interface { AddMCPClient(ctx context.Context, clientConfig *schemas.MCPClientConfig) error RemoveMCPClient(ctx context.Context, id string) error UpdateMCPClient(ctx context.Context, id string, updatedConfig *schemas.MCPClientConfig) error - // UpdateMCPClientConnection reconnects an existing MCP client using updated headers - UpdateMCPClientConnection(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error + // UpdateMCPClientCredentials reconnects an existing MCP client using updated headers + UpdateMCPClientCredentials(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error UpdateMCPToolManagerConfig(ctx context.Context, maxAgentDepth int, toolExecutionTimeoutInSeconds int, codeModeBindingLevel string, disableAutoToolInject bool) error // VerifyPerUserOAuthConnection verifies an MCP server using a temporary token and discovers tools. VerifyPerUserOAuthConnection(ctx context.Context, config *schemas.MCPClientConfig, accessToken string) (map[string]schemas.ChatTool, map[string]string, error) @@ -327,9 +327,9 @@ func (s *BifrostHTTPServer) UpdateMCPClient(ctx context.Context, id string, upda return nil } -// UpdateMCPClientConnection reconnects an existing MCP client using updated headers -func (s *BifrostHTTPServer) UpdateMCPClientConnection(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error { - if err := s.Config.UpdateMCPClientConnection(ctx, id, newConfig); err != nil { +// UpdateMCPClientCredentials reconnects an existing MCP client using updated headers +func (s *BifrostHTTPServer) UpdateMCPClientCredentials(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error { + if err := s.Config.UpdateMCPClientCredentials(ctx, id, newConfig); err != nil { return err } if err := s.MCPServerHandler.SyncAllMCPServers(ctx); err != nil { diff --git a/ui/app/workspace/mcp-registry/views/mcpClientsTable.tsx b/ui/app/workspace/mcp-registry/views/mcpClientsTable.tsx index 689432f8fc..16ac0d7c56 100644 --- a/ui/app/workspace/mcp-registry/views/mcpClientsTable.tsx +++ b/ui/app/workspace/mcp-registry/views/mcpClientsTable.tsx @@ -52,6 +52,7 @@ import { import { ReactNode, useEffect, useMemo, useState } from "react"; import { IconWrap, InfoBox } from "./authorizerUi"; import MCPClientSheet from "./mcpClientSheet"; +import { canReconnectMCPClient } from "./mcpClientsTable.utils"; import { MCPHeadersAuthorizer } from "./mcpHeadersAuthorizer"; import { MCPServersEmptyState } from "./mcpServersEmptyState"; import { MCPUsageGuideSheet } from "./mcpUsageGuide"; @@ -65,7 +66,7 @@ function MCPClientActionsMenu({ isAuthorizing, isReauthorizing, isVerifyingExchange, - isPerUserAuth, + canReconnect, onEdit, onReconnect, onAuthorize, @@ -81,7 +82,7 @@ function MCPClientActionsMenu({ isAuthorizing: boolean; isReauthorizing: boolean; isVerifyingExchange: boolean; - isPerUserAuth: boolean; + canReconnect: boolean; onEdit: (client: MCPClient) => void; onReconnect: (client: MCPClient) => void; onAuthorize: (client: MCPClient) => void; @@ -149,16 +150,10 @@ function MCPClientActionsMenu({ Authorize )} - {hasUpdateAccess && ( + {hasUpdateAccess && canReconnect && ( { e.preventDefault(); onReconnect(client); @@ -298,7 +293,6 @@ export default function MCPClientsTable({ authorizeUrl: string; oauthConfigId: string; mcpClientId: string; - popup: Window | null; isPerUserOauth: boolean; } | null>(null); // Drives the MCPHeadersAuthorizer dialog for a config.json-bootstrapped @@ -363,22 +357,12 @@ export default function MCPClientsTable({ return; } const isPerUserOauth = client.config.auth_type === "per_user_oauth"; - // Open a blank popup synchronously, before the initiateVerification - // await, so the click's transient user-activation is captured here - // rather than consumed by the network round-trip — otherwise the - // browser can block OAuth2Authorizer's later window.open entirely. - // OAuth2Authorizer navigates this handle once authorize_url is known. - // Not needed for per_user_oauth: that flow shows a confirm step first - // and opens its own popup synchronously from that step's own button - // click, so pre-opening one here would just leak an unused window. - let popup: Window | null = null; - if (!isPerUserOauth) { - const width = 600; - const height = 700; - const left = window.screen.width / 2 - width / 2; - const top = window.screen.height / 2 - height / 2; - popup = window.open("", "oauth_popup", `width=${width},height=${height},left=${left},top=${top},resizable=yes,scrollbars=yes`); - } + // OAuth2Authorizer always shows a confirm step first and opens its own + // popup synchronously from that step's own "Continue" click, for both + // OAuth flavors — nothing to pre-open here. (This used to pre-open a + // blank popup for the shared-oauth case, based on a stale assumption + // that only per_user_oauth showed a confirm step; that left a stray + // blank tab open alongside the dialog on every shared-OAuth authorize.) try { setAuthorizingClients((prev) => [...prev, client.config.client_id]); const response = await initiateVerification(client.config.client_id).unwrap(); @@ -387,11 +371,9 @@ export default function MCPClientsTable({ authorizeUrl: response.authorize_url, oauthConfigId: response.oauth_config_id, mcpClientId: client.config.client_id, - popup, isPerUserOauth, }); } else { - popup?.close(); toast({ title: "Authorization failed", description: "Unexpected response from server. Please try again.", @@ -399,7 +381,6 @@ export default function MCPClientsTable({ }); } } catch (error) { - popup?.close(); toast({ title: "Authorization failed", description: getErrorMessage(error), variant: "destructive" }); } finally { setAuthorizingClients((prev) => prev.filter((id) => id !== client.config.client_id)); @@ -650,7 +631,6 @@ export default function MCPClientsTable({ authorizeUrl={bootstrapAuthorize.authorizeUrl} oauthConfigId={bootstrapAuthorize.oauthConfigId} mcpClientId={bootstrapAuthorize.mcpClientId} - initialPopup={bootstrapAuthorize.popup} isPerUserOauth={bootstrapAuthorize.isPerUserOauth} /> )} @@ -826,11 +806,15 @@ export default function MCPClientsTable({ )} -
-
- - - +
+
+
+ + Name Connection Type Auth Type @@ -865,7 +849,7 @@ export default function MCPClientsTable({ /> Status - + @@ -877,13 +861,7 @@ export default function MCPClientsTable({ ) : ( mcpClients.map((c: MCPClient) => { - // Per-user auth types (OAuth + headers) don't hold a shared - // upstream connection, so reconnect is a no-op for them — the - // backend's ReconnectClient rejects with ErrMCPReconnectNotApplicable. - const isPerUserAuth = - c.config.auth_type === "per_user_oauth" || - c.config.auth_type === "per_user_headers" || - c.config.auth_type === "token_exchange"; + const canReconnect = canReconnectMCPClient(c.config); const enabledToolsCount = c.state == "healthy" ? c.config.tools_to_execute?.includes("*") @@ -1007,7 +985,7 @@ export default function MCPClientsTable({ isAuthorizing={authorizingClients.includes(c.config.client_id)} isReauthorizing={reauthorizingClients.includes(c.config.client_id)} isVerifyingExchange={verifyingExchangeClients.includes(c.config.client_id)} - isPerUserAuth={isPerUserAuth} + canReconnect={canReconnect} onEdit={handleRowClick} onReconnect={(client) => void handleReconnect(client)} onAuthorize={(client) => void handleStartBootstrap(client)} diff --git a/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.test.ts b/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.test.ts new file mode 100644 index 0000000000..e3d6c887e6 --- /dev/null +++ b/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from "vitest"; +import { canReconnectMCPClient } from "./mcpClientsTable.utils"; + +describe("canReconnectMCPClient", () => { + it("is false for per_user_oauth (always per-call)", () => { + expect(canReconnectMCPClient({ auth_type: "per_user_oauth", connection_type: "http", needs_session_stickiness: true })).toBe(false); + }); + + it("is false for per_user_headers (always per-call)", () => { + expect(canReconnectMCPClient({ auth_type: "per_user_headers", connection_type: "http" })).toBe(false); + }); + + it("is false for token_exchange (always per-call)", () => { + expect(canReconnectMCPClient({ auth_type: "token_exchange", connection_type: "http" })).toBe(false); + }); + + it("is false for a shared oauth http client with needs_session_stickiness omitted (default per-call)", () => { + expect(canReconnectMCPClient({ auth_type: "oauth", connection_type: "http" })).toBe(false); + }); + + it("is false for a shared oauth http client with needs_session_stickiness explicitly false", () => { + expect(canReconnectMCPClient({ auth_type: "oauth", connection_type: "http", needs_session_stickiness: false })).toBe(false); + }); + + it("is false for a shared headers http client with needs_session_stickiness omitted", () => { + expect(canReconnectMCPClient({ auth_type: "headers", connection_type: "http" })).toBe(false); + }); + + it("is true for a shared oauth http client with needs_session_stickiness explicitly true", () => { + expect(canReconnectMCPClient({ auth_type: "oauth", connection_type: "http", needs_session_stickiness: true })).toBe(true); + }); + + it("is true for a shared headers http client with needs_session_stickiness explicitly true", () => { + expect(canReconnectMCPClient({ auth_type: "headers", connection_type: "http", needs_session_stickiness: true })).toBe(true); + }); + + it("is true for a shared oauth SSE client regardless of needs_session_stickiness (non-http is always sticky)", () => { + expect(canReconnectMCPClient({ auth_type: "oauth", connection_type: "sse", needs_session_stickiness: false })).toBe(true); + }); + + it("is true for a shared headers STDIO client regardless of needs_session_stickiness (non-http is always sticky)", () => { + expect(canReconnectMCPClient({ auth_type: "headers", connection_type: "stdio" })).toBe(true); + }); + + it("is true for auth_type none on a non-http connection", () => { + expect(canReconnectMCPClient({ auth_type: "none", connection_type: "stdio" })).toBe(true); + }); + + it("is false for auth_type none on http with stickiness off", () => { + expect(canReconnectMCPClient({ auth_type: "none", connection_type: "http" })).toBe(false); + }); +}); diff --git a/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.ts b/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.ts new file mode 100644 index 0000000000..b4c927e758 --- /dev/null +++ b/ui/app/workspace/mcp-registry/views/mcpClientsTable.utils.ts @@ -0,0 +1,23 @@ +import { MCPClient } from "@/lib/types/mcp"; + +/** + * canReconnectMCPClient reports whether the "Reconnect" action applies to + * this client at all — mirrors core/mcp/credstore's RequiresPerCallConnection + * exactly, so the UI never offers an action the backend rejects with + * ErrMCPReconnectNotApplicable: + * + * - per_user_oauth, per_user_headers, token_exchange never hold a shared + * upstream connection — auth is resolved per request/user identity. + * - A shared oauth/headers/none client on an http connection with + * needs_session_stickiness nil/false (the default for newly created + * clients) is also per-call: it dials fresh on every call, so there is + * no persistent connection to reconnect either. + * - Every other combination (any non-http connection type, or an http + * shared client with needs_session_stickiness === true) holds a real + * persistent connection Reconnect can act on. + */ +export function canReconnectMCPClient(config: Pick): boolean { + const alwaysPerCall = config.auth_type === "per_user_oauth" || config.auth_type === "per_user_headers" || config.auth_type === "token_exchange"; + const sharedButPerCallViaStickiness = config.connection_type === "http" && config.needs_session_stickiness !== true; + return !(alwaysPerCall || sharedButPerCallViaStickiness); +}