fix: add SSE heartbeats and defer trace completion to prevent deadlock on MCP SSE streaming

.github/workflows/configs/withpostgresmcpclientsinconfig/config.json

@@ -1,7 +1,6 @@
 {
   "$schema": "https://www.getbifrost.ai/schema",
   "client": {
-    "allow_direct_keys": false,
     "allowed_origins": [
       "*"
     ],

core/bifrost.go

@@ -5795,7 +5795,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 					if len(supportedKeys) == 0 {
 						// SkipKeySelection path — keyProvider stays nil, zero Key is used.
 					} else if !canRotate {
-						// Fixed key (DirectKey, explicit ID/name, session stickiness): always
+						// Fixed key (explicit ID/name, session stickiness): always
 						// return the same key regardless of usedKeyIDs.
 						fixedKey := supportedKeys[0]
 						keyProvider = func(_ map[string]bool) (schemas.Key, error) {
@@ -7145,18 +7145,6 @@ func (bifrost *Bifrost) releaseMCPRequest(req *schemas.BifrostMCPRequest) {
 // getAllSupportedKeys retrieves all valid keys for a ListModels request.
 // allowing the provider to aggregate results from multiple keys.
 func (bifrost *Bifrost) getAllSupportedKeys(ctx *schemas.BifrostContext, providerKey schemas.ModelProvider, baseProviderType schemas.ModelProvider) ([]schemas.Key, error) {
-	// Check if key has been set in the context explicitly
-	if ctx != nil {
-		key, ok := ctx.Value(schemas.BifrostContextKeyDirectKey).(schemas.Key)
-		if ok {
-			if err := validateKey(baseProviderType, &key); err != nil {
-				return nil, fmt.Errorf("invalid direct key for provider %v: %w", baseProviderType, err)
-			}
-			// If a direct key is specified, return it as a single-element slice
-			return []schemas.Key{key}, nil
-		}
-	}
-
 	keys, err := bifrost.account.GetKeysForProvider(ctx, providerKey)
 	if err != nil {
 		return nil, err
@@ -7195,18 +7183,6 @@ func (bifrost *Bifrost) getAllSupportedKeys(ctx *schemas.BifrostContext, provide
 // For batch operations, only keys with UseForBatchAPI enabled are included.
 // Model filtering: if model is specified and key has model restrictions, only include if model is in list.
 func (bifrost *Bifrost) getKeysForBatchAndFileOps(ctx *schemas.BifrostContext, providerKey schemas.ModelProvider, baseProviderType schemas.ModelProvider, model *string, isBatchOp bool) ([]schemas.Key, error) {
-	// Check if key has been set in the context explicitly
-	if ctx != nil {
-		key, ok := ctx.Value(schemas.BifrostContextKeyDirectKey).(schemas.Key)
-		if ok {
-			if err := validateKey(baseProviderType, &key); err != nil {
-				return nil, fmt.Errorf("invalid direct key for provider %v: %w", baseProviderType, err)
-			}
-			// If a direct key is specified, return it as a single-element slice
-			return []schemas.Key{key}, nil
-		}
-	}
-
 	keys, err := bifrost.account.GetKeysForProvider(ctx, providerKey)
 	if err != nil {
 		return nil, err
@@ -7278,7 +7254,6 @@ func (bifrost *Bifrost) getKeysForBatchAndFileOps(ctx *schemas.BifrostContext, p
 // via the keyProvider closure built by the caller.
 //
 // canRotate=false is returned for cases where the caller must always use the same key:
-//   - DirectKey (caller-supplied key bypasses all selection)
 //   - SkipKeySelection (provider allows keyless requests; empty slice returned)
 //   - Explicit BifrostContextKeyAPIKeyID / APIKeyName (user pinned a specific key)
 //   - Session stickiness (key persisted in KV store for the session lifetime)
@@ -7287,15 +7262,6 @@ func (bifrost *Bifrost) getKeysForBatchAndFileOps(ctx *schemas.BifrostContext, p
 // canRotate=true is returned when there are two or more eligible keys and no pinning
 // or stickiness constraint is in effect.
 func (bifrost *Bifrost) selectKeyFromProviderForModelWithPool(ctx *schemas.BifrostContext, requestType schemas.RequestType, providerKey schemas.ModelProvider, model string, baseProviderType schemas.ModelProvider) ([]schemas.Key, bool, error) {
-	// DirectKey: caller supplied a key directly — no pool, no rotation.
-	if ctx != nil {
-		if key, ok := ctx.Value(schemas.BifrostContextKeyDirectKey).(schemas.Key); ok {
-			if err := validateKey(baseProviderType, &key); err != nil {
-				return nil, false, fmt.Errorf("invalid direct key for provider %v: %w", baseProviderType, err)
-			}
-			return []schemas.Key{key}, false, nil
-		}
-	}
 	// SkipKeySelection: provider allows keyless requests — return empty pool, no rotation.
 	if skipKeySelection, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) {
 		return []schemas.Key{}, false, nil

core/internal/llmtests/account.go

@@ -204,12 +204,6 @@ func replicateProviderTestKeys() []schemas.Key {
 	}
 }
 
-// ReplicateDirectKeyForListModels returns the key used for Replicate ListModels (deployments endpoint).
-// List-models tests set it on the context as schemas.BifrostContextKeyDirectKey so Bifrost passes only this key.
-func ReplicateDirectKeyForListModels() schemas.Key {
-	return replicateProviderTestKeys()[0]
-}
-
 // GetKeysForProvider returns the API keys and associated models for a given provider.
 func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
 	switch providerKey {

core/internal/llmtests/list_models.go

@@ -9,13 +9,14 @@ import (
 	"github.com/maximhq/bifrost/core/schemas"
 )
 
-// listModelsBifrostContext returns a context for ListModels. For Replicate, sets BifrostContextKeyDirectKey
-// so only the deployments key is used (see replicateProviderTestKeys in account.go). That key must not use an
-// empty Models allowlist, or ListModelsPipeline.ShouldEarlyExit returns no models before the API runs.
+// listModelsBifrostContext returns a context for ListModels. For Replicate, pins the deployments-endpoint
+// key by name (see replicateProviderTestKeys in account.go) so the test always exercises that specific key.
+// That key must not use an empty Models allowlist, or ListModelsPipeline.ShouldEarlyExit returns no models
+// before the API runs.
 func listModelsBifrostContext(parent context.Context, provider schemas.ModelProvider) *schemas.BifrostContext {
 	bfCtx := schemas.NewBifrostContext(parent, schemas.NoDeadline)
 	if provider == schemas.Replicate {
-		bfCtx.SetValue(schemas.BifrostContextKeyDirectKey, ReplicateDirectKeyForListModels())
+		bfCtx.SetValue(schemas.BifrostContextKeyAPIKeyName, ReplicateKeyNameListModels)
 	}
 	return bfCtx
 }

core/mcp/clientmanager.go

@@ -63,7 +63,7 @@ func (m *MCPManager) ReconnectClient(id string) error {
 	// Reconnect is not applicable because auth is resolved per request/user identity.
 	if client.ExecutionConfig != nil && client.ExecutionConfig.AuthType == schemas.MCPAuthTypePerUserOauth {
 		m.mu.Unlock()
-		return fmt.Errorf("reconnect is not supported for per_user_oauth clients")
+		return fmt.Errorf("per-user OAuth clients do not maintain a shared upstream connection (each user manages their own auth): %w", schemas.ErrMCPReconnectNotApplicable)
 	}
 	config := client.ExecutionConfig
 	m.mu.Unlock()

core/mcp/utils/utils.go

@@ -25,7 +25,11 @@ func ResolvePerUserOAuthToken(ctx *schemas.BifrostContext, client *schemas.MCPCl
 	}
 
 	accessToken, err := oauth2Provider.GetUserAccessTokenByIdentity(ctx, virtualKeyID, userID, sessionToken, client.ExecutionConfig.ID)
-	if err != nil && !errors.Is(err, schemas.ErrOAuth2TokenNotFound) {
+	// Both sentinels mean "this user must re-authenticate":
+	//   - ErrOAuth2TokenNotFound: row missing (never authed, or purged after permanent refresh failure)
+	//   - ErrOAuth2TokenExpired:  row present but tokens unusable (access expired + no refresh available)
+	// Either way, fall through to the re-auth branch below to surface an inline auth URL.
+	if err != nil && !errors.Is(err, schemas.ErrOAuth2TokenNotFound) && !errors.Is(err, schemas.ErrOAuth2TokenExpired) {
 		return "", fmt.Errorf("failed to get user access token for MCP server %s: %w", client.ExecutionConfig.Name, err)
 	}
 	if err != nil {

core/providers/utils/utils.go

@@ -2786,11 +2786,11 @@ func completeDeferredSpan(ctx *schemas.BifrostContext, result *schemas.BifrostRe
 
 // CheckAndSetDefaultProvider checks if the default provider should be used based on the context.
 // It returns the default provider if it should be used, otherwise it returns an empty string.
-// Checks if the direct key is set in the context, or if key selection is skipped.
-// Or if the available providers are set in the context and the default provider is in the list.
+// Checks if key selection is skipped, or if the available providers are set in the context
+// and the default provider is in the list.
 func CheckAndSetDefaultProvider(ctx *schemas.BifrostContext, defaultProvider schemas.ModelProvider) schemas.ModelProvider {
 	if ctx != nil {
-		if ctx.Value(schemas.BifrostContextKeyDirectKey) != nil || ctx.Value(schemas.BifrostContextKeySkipKeySelection) != nil {
+		if skip, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skip {
 			return defaultProvider
 		}
 		if ctx.Value(schemas.BifrostContextKeyAvailableProviders) != nil {

core/schemas/bifrost.go

@@ -172,7 +172,6 @@ const (
 	BifrostContextKeyAPIKeyID          BifrostContextKey = "x-bf-api-key-id"       // string (explicit key ID selection, takes priority over name)
 	BifrostContextKeyRequestID         BifrostContextKey = "request-id"            // string
 	BifrostContextKeyFallbackRequestID BifrostContextKey = "fallback-request-id"   // string
-	BifrostContextKeyDirectKey         BifrostContextKey = "bifrost-direct-key"    // Key struct
 
 	// NOTE: []string is used for both keys, and by default all clients/tools are included (when nil).
 	// If "*" is present, all clients/tools are included, and [] means no clients/tools are included.

core/schemas/context.go

@@ -16,7 +16,6 @@ var reservedKeys = []any{
 	BifrostContextKeyAPIKeyID,
 	BifrostContextKeyRequestID,
 	BifrostContextKeyFallbackRequestID,
-	BifrostContextKeyDirectKey,
 	BifrostContextKeySelectedKeyID,
 	BifrostContextKeySelectedKeyName,
 	BifrostContextKeyNumberOfRetries,

core/schemas/mcp.go

@@ -29,6 +29,11 @@ var (
 	ErrOAuth2NotPerUserSession        = errors.New("state does not match a per-user oauth session")
 	ErrOAuth2TokenNotFound            = errors.New("per-user oauth token not found for this identity and mcp server")
 	ErrPerUserOAuthPendingFlowExpired = errors.New("per-user oauth pending flow has expired")
+	// ErrMCPReconnectNotApplicable signals that the reconnect operation is not
+	// meaningful for this client type — e.g. per-user OAuth clients, where
+	// each user manages their own auth and there is no shared upstream
+	// connection to "reconnect". Distinct from "not implemented".
+	ErrMCPReconnectNotApplicable = errors.New("reconnect is not applicable for this client type")
 )
 
 // MCPUserOAuthRequiredError is returned when a per-user OAuth MCP server requires

docs/deployment-guides/config-json.mdx

@@ -136,7 +136,6 @@ A production-ready file with PostgreSQL storage, multi-provider setup, governanc
     "enable_logging": true,
     "log_retention_days": 90,
     "enforce_auth_on_inference": true,
-    "allow_direct_keys": false,
     "allowed_origins": ["https://app.yourcompany.com"]
   },
 

docs/deployment-guides/config-json/client.mdx

@@ -127,7 +127,6 @@ These settings are also configurable via the UI (**MCP Gateway → MCP Settings*
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `allowed_origins` | array | `["*"]` | CORS allowed origins (use URIs or `"*"`) |
-| `allow_direct_keys` | boolean | `false` | Allow callers to pass provider keys directly in requests |
 | `enforce_auth_on_inference` | boolean | `false` | Require auth (virtual key, API key, or user token) on `/v1/*` inference routes |
 | `max_request_body_size_mb` | integer | `100` | Maximum allowed request body size in MB |
 | `whitelisted_routes` | array of strings | `[]` | Routes that bypass auth middleware |
@@ -140,7 +139,6 @@ These settings are also configurable via the UI (**MCP Gateway → MCP Settings*
       "https://app.yourcompany.com",
       "https://admin.yourcompany.com"
     ],
-    "allow_direct_keys": false,
     "enforce_auth_on_inference": true,
     "max_request_body_size_mb": 50,
     "whitelisted_routes": ["/health", "/metrics"]
@@ -324,7 +322,6 @@ A top-level `auth_config` is also accepted for backwards compatibility, but `gov
     "mcp_external_client_url": "env.BIFROST_EXTERNAL_URL",
 
     "allowed_origins": ["https://app.yourcompany.com"],
-    "allow_direct_keys": false,
     "enforce_auth_on_inference": true,
     "max_request_body_size_mb": 100,
 

docs/deployment-guides/config-json/schema-reference.mdx

@@ -61,7 +61,6 @@ Controls the worker pool, logging pipeline, security, and SDK shims. All fields
 | `log_retention_days` | integer | `365` | Days to retain log entries |
 | `logging_headers` | array | `[]` | HTTP headers to capture in log metadata |
 | `enforce_auth_on_inference` | boolean | `false` | Require a virtual key on every `/v1/*` request |
-| `allow_direct_keys` | boolean | `false` | Allow callers to pass provider API keys directly |
 | `allowed_origins` | array | `["*"]` | CORS allowed origins |
 | `max_request_body_size_mb` | integer | `100` | Maximum request body in MB |
 | `whitelisted_routes` | array | `[]` | Routes that bypass auth middleware |

docs/deployment-guides/helm.mdx

@@ -405,7 +405,6 @@ bifrost:
     disableContentLogging: false    # set true for HIPAA/compliance
     logRetentionDays: 365
     enforceGovernanceHeader: true
-    allowDirectKeys: false
     maxRequestBodySizeMb: 100
     allowedOrigins:
       - "https://yourcompany.com"

docs/deployment-guides/helm/client.mdx

@@ -76,7 +76,6 @@ helm upgrade bifrost bifrost/bifrost \
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `bifrost.client.allowedOrigins` | CORS allowed origins | `["*"]` |
-| `bifrost.client.allowDirectKeys` | Allow callers to pass provider keys directly in requests | `false` |
 | `bifrost.client.enforceGovernanceHeader` | Require `x-bf-vk` virtual-key header on every request | `false` |
 | `bifrost.client.maxRequestBodySizeMb` | Maximum allowed request body size | `100` |
 | `bifrost.client.whitelistedRoutes` | Routes that bypass auth middleware | `[]` |
@@ -87,7 +86,6 @@ bifrost:
     allowedOrigins:
       - "https://app.yourdomain.com"
       - "https://admin.yourdomain.com"
-    allowDirectKeys: false         # Prevent callers from supplying raw provider keys
     enforceGovernanceHeader: true  # Every request must carry a virtual key
     maxRequestBodySizeMb: 50
     whitelistedRoutes:
@@ -98,8 +96,7 @@ bifrost:
 ```bash
 helm install bifrost bifrost/bifrost \
   --set image.tag=v1.4.11 \
-  --set bifrost.client.enforceGovernanceHeader=true \
-  --set bifrost.client.allowDirectKeys=false
+  --set bifrost.client.enforceGovernanceHeader=true
 ```
 
 ---
@@ -295,7 +292,6 @@ bifrost:
     disableContentLogging: false
     logRetentionDays: 90
     enforceGovernanceHeader: true
-    allowDirectKeys: false
     maxRequestBodySizeMb: 100
     headerFilterConfig:
       allowlist: []

docs/deployment-guides/helm/values.mdx

@@ -466,7 +466,6 @@ bifrost:
 
   client:
     enableLogging: true
-    allowDirectKeys: false
 
   providers:
     openai:

docs/enterprise/migration-guides/v1.4.0.mdx

@@ -26,6 +26,7 @@ Enterprise v1.4.0 ships with the full v1.5.0 OSS base, so every breaking change
 | 7 | **WhiteList validation** | Lists cannot mix `["*"]` with specific values, and cannot contain duplicates |
 | 8 | **`weight` is now nullable** | Update API consumers to handle `null` |
 | 9 | **`selected_key_id` cleared on terminal retry failures** | Read `attempt_trail` for failure attribution |
+| 10 | **Direct Key Bypass removed (HTTP + Go SDK)** | Drop `allow_direct_keys` from `config.json`; migrate header-passed provider keys to Bifrost-managed keys + virtual keys; replace any Go SDK usage of `BifrostContextKeyDirectKey` with `BifrostContextKeyAPIKeyID` / `BifrostContextKeyAPIKeyName` |
 
 The automatic database migration on startup converts existing records to the new semantics. Only `config.json` and any REST API integrations need manual updates.
 
@@ -158,7 +159,7 @@ Snapshot your config store database (Postgres dump or SQLite file copy) before s
 </Step>
 
 <Step title="Apply the OSS v1.5.0 migration steps">
-Work through the [OSS v1.5.0 Migration Guide](/migration-guides/v1.5.0) checklist: update `models`, `allowed_models`, `key_ids`, `tools_to_execute`, rename `allowed_keys` to `key_ids`, ensure every VK has at least one provider config, migrate provider key management to dedicated endpoints, and update Go SDK references.
+Work through the [OSS v1.5.0 Migration Guide](/migration-guides/v1.5.0) checklist: update `models`, `allowed_models`, `key_ids`, `tools_to_execute`, rename `allowed_keys` to `key_ids`, ensure every VK has at least one provider config, migrate provider key management to dedicated endpoints, remove `allow_direct_keys`, migrate HTTP header-key callers to Bifrost-managed keys + virtual keys, migrate Go SDK `BifrostContextKeyDirectKey` callers to `BifrostContextKeyAPIKeyID`/`BifrostContextKeyAPIKeyName`, and update Go SDK references.
 </Step>
 
 <Step title="Open the gRPC cluster port (10102/TCP) on every node">

docs/features/governance/virtual-keys.mdx

@@ -16,8 +16,6 @@ Virtual Keys are the primary governance entity in Bifrost. Users and application
 
 <Note>Old virtual keys(without `sk-bf-*` prefix) are only supported by `x-bf-vk` header.</Note>
 
-<Info>You can also use `Authorization`, `x-api-key` and `x-goog-api-key` headers to pass direct keys to the provider. Read more about it in [Direct Key Bypass](../keys-management#direct-key-bypass).</Info>
-
 **Key Features:**
 - **Access Control** - Model and provider filtering
 - **Cost Management** - Independent budgets (checked along with team/customer budgets if attached)