// Package server provides the HTTP server for Bifrost. package server import ( "context" "embed" "errors" "fmt" "math/rand" "net" "os" "os/signal" "strconv" "strings" "sync" "syscall" "time" "github.com/fasthttp/router" "github.com/google/uuid" "github.com/mark3labs/mcp-go/mcp" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/agent" "github.com/maximhq/bifrost/core/network" "github.com/maximhq/bifrost/core/relay" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" "github.com/maximhq/bifrost/framework/configstore/tables" "github.com/maximhq/bifrost/framework/encrypt" "github.com/maximhq/bifrost/framework/logstore" "github.com/maximhq/bifrost/framework/modelcatalog" "github.com/maximhq/bifrost/framework/openapimcp" dynamicPlugins "github.com/maximhq/bifrost/framework/plugins" "github.com/maximhq/bifrost/framework/sidekiq" "github.com/maximhq/bifrost/framework/temptoken" "github.com/maximhq/bifrost/framework/tracing" "github.com/maximhq/bifrost/framework/warp" "github.com/maximhq/bifrost/framework/webhooks" "github.com/maximhq/bifrost/plugins/governance" "github.com/maximhq/bifrost/plugins/logging" "github.com/maximhq/bifrost/plugins/otel" "github.com/maximhq/bifrost/plugins/prompts" "github.com/maximhq/bifrost/plugins/routing" "github.com/maximhq/bifrost/plugins/routing/complexity" "github.com/maximhq/bifrost/plugins/semanticcache" "github.com/maximhq/bifrost/plugins/telemetry" "github.com/maximhq/bifrost/transports/v2/bifrost-http/handlers" "github.com/maximhq/bifrost/transports/v2/bifrost-http/integrations" "github.com/maximhq/bifrost/transports/v2/bifrost-http/lib" bfws "github.com/maximhq/bifrost/transports/v2/bifrost-http/websocket" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" dto "github.com/prometheus/client_model/go" "github.com/valyala/fasthttp" "github.com/valyala/fasthttp/fasthttpadaptor" "gorm.io/gorm" ) // Constants const ( DefaultHost = "localhost" DefaultPort = "8080" DefaultAppDir = "" // Empty string means use OS-specific config directory DefaultLogLevel = string(schemas.LogLevelInfo) DefaultLogOutputStyle = string(schemas.LoggerOutputTypeJSON) ) var enterprisePlugins = []string{ "datadog", "bigquery", "pubsub", "kafka", } // ServerCallbacks is a interface that defines the callbacks for the server. type ServerCallbacks interface { // Plugins callbacks ReloadPlugin(ctx context.Context, name string, path *string, pluginConfig any, placement *schemas.PluginPlacement, order *int) error RemovePlugin(ctx context.Context, name string) error GetPluginStatus(ctx context.Context) map[string]schemas.PluginStatus GetLoadedPluginNames() []string NormalizePluginConfig(name string, config map[string]any) (map[string]any, error) ExpandPluginConfigForAPI(name string, config map[string]any) (map[string]any, error) // Auth related callbacks UpdateAuthConfig(ctx context.Context, authConfig *configstore.AuthConfig) error ValidateSetupToken(token string) bool ValidateConfiguredSetupToken(token string) bool ReloadClientConfigFromConfigStore(ctx context.Context) error // Pricing related callbacks UpdateSyncConfig(ctx context.Context) error ForceReloadPricing(ctx context.Context) error UpsertPricingOverride(ctx context.Context, override *tables.TablePricingOverride) error DeletePricingOverride(ctx context.Context, id string) error // UpsertModelPricingAttributes writes the additional_attributes JSON on // pricing rows. Enterprise wraps this so that after the local DB write // succeeds it can broadcast a peer reload via the existing pricing // EntityTypeModelCatalog/ActionReloadFromDB gossip path. UpsertModelPricingAttributes(ctx context.Context, entries []handlers.ModelPricingAttributesEntry) error // Proxy related callbacks ReloadProxyConfig(ctx context.Context, config *tables.GlobalProxyConfig) error // Client config related callbacks ReloadHeaderFilterConfig(ctx context.Context, config *tables.GlobalHeaderFilterConfig) error UpdateDropExcessRequests(ctx context.Context, value bool) // Governance related callbacks GetGovernanceData(ctx context.Context) *governance.GovernanceData ReloadTeam(ctx context.Context, id string) (*tables.TableTeam, error) RemoveTeam(ctx context.Context, id string) error ReloadCustomer(ctx context.Context, id string) (*tables.TableCustomer, error) RemoveCustomer(ctx context.Context, id string) error // Virtual key related callbacks ReloadVirtualKey(ctx context.Context, id string) (*tables.TableVirtualKey, error) // Virtual MCP cache refresh after a write, surgically per operation (mirrors the VK/team callbacks). ReloadVirtualMCP(ctx context.Context, id uint) (*tables.TableVirtualMCP, error) RemoveVirtualMCP(ctx context.Context, id uint) error AttachVirtualMCPToVirtualKeyInMemory(ctx context.Context, vkID string, id uint) error DetachVirtualMCPFromVirtualKeyInMemory(ctx context.Context, vkID string, id uint) error // ReloadVirtualKeys reloads many virtual keys exactly as ReloadVirtualKey // reloads one, from batched reads. Enterprise also propagates it to peers. ReloadVirtualKeys(ctx context.Context, ids []string) error // ResetBudgetUsageInMemory clears usage for the given budgets in the governance // store, leaving each reset boundary untouched. owner identifies the entity that // owns them so enterprise can address the cluster broadcast that propagates the // reset to peers. ResetBudgetUsageInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string) error // AdoptCalendarAlignmentInMemory re-anchors an owner's budgets and rate limits // onto their calendar boundaries after alignment was switched on, preserving // usage. owner identifies the entity so enterprise can address the cluster // broadcast, exactly as the usage reset above does. AdoptCalendarAlignmentInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string, rateLimitIDs []string) error RemoveVirtualKey(ctx context.Context, id string) error // Provider related callbacks GetModelsForProvider(provider schemas.ModelProvider) []string GetUnfilteredModelsForProvider(provider schemas.ModelProvider) []string ReloadModelConfig(ctx context.Context, id string) (*tables.TableModelConfig, error) ModelConfigIndexKey(model string, provider *string) string RemoveModelConfig(ctx context.Context, id string) error ReloadProvider(ctx context.Context, provider schemas.ModelProvider, isNew bool) (*tables.TableProvider, error) RemoveProvider(ctx context.Context, provider schemas.ModelProvider) error OnKeyAdded(ctx context.Context, provider schemas.ModelProvider, key schemas.Key) error OnKeyUpdated(ctx context.Context, provider schemas.ModelProvider, key schemas.Key) error OnKeyDeleted(ctx context.Context, provider schemas.ModelProvider, keyID string) error RefreshLiveModelsForKey(ctx context.Context, provider schemas.ModelProvider, keyID string) error RefreshLiveModelsForAllKeys(ctx context.Context, provider schemas.ModelProvider) error // Routing related callbacks ReloadRoutingRule(ctx context.Context, id string) error RemoveRoutingRule(ctx context.Context, id string) error ReloadComplexityAnalyzerConfig(ctx context.Context, config *complexity.AnalyzerConfig) error ValidateComplexityAnalyzerConfig(ctx context.Context, config *complexity.AnalyzerConfig) error GetComplexitySemanticStatus(ctx context.Context) (complexity.SemanticStatusInfo, error) RetryComplexitySemanticWarmup(ctx context.Context) (complexity.SemanticStatusInfo, bool, error) GetComplexityLLMStatus(ctx context.Context) (complexity.LLMStatusInfo, error) ListComplexityGenerations(ctx context.Context) ([]complexity.GenerationInfo, error) DeleteComplexityGeneration(ctx context.Context, namespace string) error // Prompt repository related callbacks ReloadPromptCache(ctx context.Context) error // Webhook related callbacks ReloadWebhookEndpoint(ctx context.Context, id string) error RemoveWebhookEndpoint(ctx context.Context, id string) error // MCP related callbacks 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 // 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, maxInstructionsPerClient int, maxInstructionsTotal int, codeModeLimits *schemas.MCPCodeModeLimits) error // VerifyPerUserOAuthConnection verifies an MCP server using a temporary token and discovers tools. VerifyPerUserOAuthConnection(ctx context.Context, config *schemas.MCPClientConfig, accessToken string) (*schemas.MCPDiscovery, error) // VerifyHeadersConnection verifies an MCP server using user-supplied header values and discovers tools. VerifyHeadersConnection(ctx context.Context, config *schemas.MCPClientConfig, userHeaders map[string]string) (*schemas.MCPDiscovery, error) // SetClientDiscovery installs discovery on an existing client. SetClientDiscovery(clientID string, discovery *schemas.MCPDiscovery) // RequiresPerCallConnection reports whether config resolves to a // per-call connection (true) or a persistent shared one (false), taking // auth type, connection type, and needs_session_stickiness into account // together. RequiresPerCallConnection(config *schemas.MCPClientConfig) bool ReconnectMCPClient(ctx context.Context, id string) error // RefreshMCPClientTools re-discovers a client's tools from its upstream // server on demand and reports how many it serves afterwards. Unlike // ReconnectMCPClient it applies to per-call clients too. RefreshMCPClientTools(ctx context.Context, id string) (int, error) // CloseAndMarkNeedsReauth closes a shared client's live upstream // connection and flips it to needs_reauth, without attempting a new // dial. Used after OAuth credential rotation. CloseAndMarkNeedsReauth(ctx context.Context, id string) error DisableMCPClient(ctx context.Context, id string) error EnableMCPClient(ctx context.Context, id string) error // EvictOauthTokenCacheByID drops the cached per-user MCP OAuth access // token backed by the given token row ID after a database write that // bypassed the OAuth provider's own write paths. The server also forwards // the provider's own write evictions here, since those reach only this // process. EvictOauthTokenCacheByID(ctx context.Context, tokenID string) // EvictOauthTokenCacheByMCPClient drops every cached per-user MCP OAuth // access token bound to the given MCP client after a client-level // mutation that invalidates its token rows as a set (credential // rotation, access reconciliation, client deletion). EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) // EvictMCPHeaderCredentialCacheByID drops the cached per-user MCP header // credential backed by the given credential row ID after a database // write that bypassed the headers provider's own write paths. The server // also forwards the provider's own write evictions here, since those // reach only this process. EvictMCPHeaderCredentialCacheByID(ctx context.Context, credentialID string) // EvictMCPHeaderCredentialCacheByMCPClient drops every cached per-user // MCP header credential bound to the given MCP client after a // client-level mutation that invalidates its credential rows as a set // (needs_update schema flip, access reconciliation, client deletion). EvictMCPHeaderCredentialCacheByMCPClient(ctx context.Context, mcpClientID string) // ResolveAccess answers what a request may reach, resolving it once onto the request's grant, // for the listing routes that run outside the request pipeline and so cannot wait for a hook // to resolve it. ResolveAccess(ctx *schemas.BifrostContext) (schemas.Access, error) // The models listing narrows its provider fan-out to what the request may reach. NarrowListModelsProviders(bifrostCtx *schemas.BifrostContext) } // GovernanceRouteOverridesProvider lets downstream editions replace selected OSS governance route families. type GovernanceRouteOverridesProvider interface { GetGovernanceRouteOverrides(ctx context.Context) handlers.GovernanceRouteOverrides } // LogRedactionMappingResolverProvider is implemented by servers that can attach reveal data to log-detail responses. type LogRedactionMappingResolverProvider interface { // GetLogRedactionMappingResolver returns the resolver used by the logging handler. GetLogRedactionMappingResolver() handlers.LogRedactionMappingResolver } // MCPLogRedactionMappingResolverProvider is implemented by servers that can attach reveal data to MCP log-detail responses. type MCPLogRedactionMappingResolverProvider interface { // GetMCPLogRedactionMappingResolver returns the resolver used by the logging handler. GetMCPLogRedactionMappingResolver() handlers.MCPLogRedactionMappingResolver } // BifrostHTTPServer represents a HTTP server instance. type BifrostHTTPServer struct { Ctx *schemas.BifrostContext cancel context.CancelFunc mcpDiscoveries mcpDiscoveryWriter // see PersistMCPClientDiscovery Version string UIContent embed.FS Port string Host string AppDir string LogLevel string LogOutputStyle string LogsCleaner *logstore.LogsCleaner AsyncJobCleaner *logstore.AsyncJobCleaner Client *bifrost.Bifrost Config *lib.Config // HTTPClientFactory hands out clients for outbound traffic other than provider // inference (webhooks, skills, plugin downloads, MCP, catalog sync) that honour // the global proxy for their purpose and follow proxy changes live. HTTPClientFactory *network.HTTPClientFactory Server *fasthttp.Server Router *router.Router // ShellRewriter lets a build rewrite the pre-hydration HTML shell before it // is served — the enterprise build uses it to swap in a custom logo. // nil on OSS, where the embedded document is served exactly as bundled. It // must be assigned before RegisterUIRoutes, which builds the UI handler // from it. ShellRewriter handlers.ShellRewriter WebSocketHandler *handlers.WebSocketHandler NotificationService *handlers.NotificationService RSSFeedsHandler *handlers.RSSFeedsHandler WarpHandler *handlers.WarpHandler MCPServerHandler *handlers.MCPServerHandler AgentGatewayHandler *handlers.AgentGatewayHandler devPprofHandler *handlers.DevPprofHandler skillsServingHandler *handlers.SkillsServingHandler IntegrationHandler *handlers.IntegrationHandler wsLiveHandler *handlers.WSLiveHandler webrtcLiveHandler *handlers.WebRTCLiveHandler AuthMiddleware *handlers.AuthMiddleware setupLockInstalled bool // OSS setup-lock gate guards /api (see AuthMiddleware.SetupLockMiddleware); never set on enterprise CORSMiddleware *handlers.CorsMiddleware TracingMiddleware *handlers.TracingMiddleware WSTicketStore *handlers.WSTicketStore TempTokens *temptoken.Service TempTokenSweepWorker *temptoken.SweepWorker OAuth2SweepWorker *oauth2SweepWorker // OAuth2IdentityResolver scopes a user-mode /mcp request to the user's own // tools. Optional; wired at server init when user-mode identity resolution // is available, otherwise left nil (user-mode requests fall back to the // global server). OAuth2IdentityResolver handlers.OAuth2IdentityResolver // OAuth2Extensions adds grant types and discovery metadata to Bifrost's OAuth2 // authorization server. Nil means none; set before RegisterAPIRoutes. OAuth2Extensions handlers.OAuth2Extensions // AgentGatewayIdentityResolver validates pre-authenticated Authorization // credentials for A2A HTTP and gRPC ingress. Optional; OSS leaves it nil. AgentGatewayIdentityResolver handlers.AgentGatewayIdentityResolver // ExternalQuotaBudgetResolver supplies budgets/usage for VKs whose // authoritative usage is tracked outside their own budget rows (enterprise // access-profile-managed VKs). Optional; wired at server init when available, // otherwise left nil so the quota endpoint reads the VK's own budget rows. ExternalQuotaBudgetResolver handlers.ExternalQuotaBudgetResolver // VirtualKeyAssigneeResolver supplies the user each VK is assigned to, batched // per page. Optional; wired at server init when available, otherwise left nil // so the VK read paths report no assignee (OSS has no user directory). VirtualKeyAssigneeResolver handlers.VirtualKeyAssigneeResolver // VirtualKeyBusinessUnitResolver names each VK's owning business unit on the // governance read paths. Nil on OSS builds; set by the enterprise wrapper. VirtualKeyBusinessUnitResolver handlers.VirtualKeyBusinessUnitResolver // WarpUserGovernanceReader answers Warp's describe_user_limits: what an // access profile lets one person spend. Nil on OSS builds, which have no // per-user governance, so Warp does not offer the tool; set by the // enterprise wrapper before RegisterAPIRoutes, like the resolvers above. WarpUserGovernanceReader warp.UserGovernanceReader // WarpCallerRestrictionResolver tells Warp whether row-level access control // narrows the caller's reads, so a restricted caller's totals are not // described as the whole deployment's. Nil on OSS builds, where the // request's own query scope is the only one; set by the enterprise wrapper // before RegisterAPIRoutes, like the reader above. WarpCallerRestrictionResolver warp.CallerRestrictionResolver SidekiqRunner *sidekiq.Runner // GovernanceHandler is kept so the expired-key cleanup scheduler can be started and stopped. GovernanceHandler *handlers.GovernanceHandler SidekiqDispatcherStop func() // Background live model catalog refresher. Guarded because the framework // config handler can restart it from a request goroutine while the boot // path or shutdown is touching it. liveModelRefresherMu sync.Mutex liveModelRefresherStop func() // refreshInFlight coordinates on-demand model refreshes. Distinct keys for // the same provider may refresh concurrently, while an all-keys refresh is // exclusive with every key refresh for that provider. refreshInFlightMu sync.Mutex refreshInFlight map[schemas.ModelProvider]*providerRefreshState WebhookDispatcher *webhooks.Dispatcher wsPool *bfws.Pool agentGatewayCloseOnce sync.Once // agentGatewayManager backs the gRPC binding, which needs the manager // directly rather than through the HTTP handler. agentGatewayManager *agent.Manager // relays serves /relay/{name}/{path}; built by InitializeRelays together with RelaysHandler, // the management handler a clustered server replays relay changes through. relays *relay.Manager RelaysHandler *handlers.RelaysHandler endpointSweeperCancel context.CancelFunc endpointSweeperMu sync.Mutex // serializes sweeper restarts from concurrent plugin reloads agentGatewayGRPCServer *handlers.AgentGatewayGRPCServer } var logger schemas.Logger // SetLogger sets the logger for the server. func SetLogger(l schemas.Logger) { logger = l } // NewBifrostHTTPServer creates a new instance of BifrostHTTPServer. func NewBifrostHTTPServer(version string, uiContent embed.FS) *BifrostHTTPServer { return &BifrostHTTPServer{ Version: version, UIContent: uiContent, Port: DefaultPort, Host: DefaultHost, AppDir: DefaultAppDir, LogLevel: DefaultLogLevel, LogOutputStyle: DefaultLogOutputStyle, } } type GovernanceInMemoryStore struct { Config *lib.Config } // GetConfiguredProviders returns a snapshot of the configured provider map, built under the read // lock. Provider edits write to the live map in place, and the governance plugin indexes what // this returns after the lock is released (batchBaseProvider, modelMatcher), so handing back the // live map would be a concurrent map read and write, which is fatal rather than merely stale. // The copy is shallow: the configs themselves are shared and never modified by readers. func (s *GovernanceInMemoryStore) GetConfiguredProviders() map[schemas.ModelProvider]configstore.ProviderConfig { s.Config.Mu.RLock() defer s.Config.Mu.RUnlock() snapshot := make(map[schemas.ModelProvider]configstore.ProviderConfig, len(s.Config.Providers)) for provider, config := range s.Config.Providers { snapshot[provider] = config } return snapshot } // GetConfiguredProvider returns one provider's configuration, read under the lock without copying // the map, for the per-request model check that needs only that provider. func (s *GovernanceInMemoryStore) GetConfiguredProvider(provider schemas.ModelProvider) (configstore.ProviderConfig, bool) { s.Config.Mu.RLock() defer s.Config.Mu.RUnlock() config, ok := s.Config.Providers[provider] return config, ok } // GetConfiguredProviderNames builds the name slice under the lock, because provider edits write to // the map GetConfiguredProviders hands back in place: ranging that map after the lock is released // is a concurrent iteration and write, which is fatal rather than merely stale. func (s *GovernanceInMemoryStore) GetConfiguredProviderNames() []string { providers, err := s.Config.GetAllProviders() if err != nil { return nil } names := make([]string, 0, len(providers)) for _, provider := range providers { names = append(names, string(provider)) } return names } // GetMCPClientsAllowedByDefault returns the default MCP client allowlist. func (s *GovernanceInMemoryStore) GetMCPClientsAllowedByDefault() map[string]string { return s.Config.GetMCPClientsAllowedByDefault() } // GetMCPClientNames returns configured MCP client names indexed by slug. func (s *GovernanceInMemoryStore) GetMCPClientNames() map[string]string { return s.Config.GetMCPClientNames() } // GetMCPClientBySlug resolves an MCP client configuration by slug. func (s *GovernanceInMemoryStore) GetMCPClientBySlug(slug string) (string, string, bool) { return s.Config.GetMCPClientBySlug(slug) } func (s *GovernanceInMemoryStore) GetEnabledAgents() map[string]bool { return s.Config.GetEnabledAgents() } // AddMCPClient adds a new MCP client to the in-memory store func (s *BifrostHTTPServer) AddMCPClient(ctx context.Context, clientConfig *schemas.MCPClientConfig) error { if err := s.Config.AddMCPClient(ctx, clientConfig); err != nil { return err } if err := s.MCPServerHandler.SyncMCPServer(ctx); err != nil { logger.Warn("failed to sync MCP servers after adding client: %v", err) } return nil } // ReconnectMCPClient reconnects an MCP client to the in-memory store func (s *BifrostHTTPServer) ReconnectMCPClient(ctx context.Context, id string) error { // Check if client is registered in Bifrost (can be not registered if client initialization failed) if clients, err := s.Client.GetMCPClients(); err == nil && len(clients) > 0 { for _, client := range clients { if client.Config.ID == id { if err := s.Client.ReconnectMCPClient(id); err != nil { return err } return nil } } } // Config exists in store, but not in Bifrost (can happen if client initialization failed) clientConfig, err := s.Config.GetMCPClient(id) if err != nil { return err } if err := s.Client.AddMCPClient(ctx, clientConfig); err != nil { return err } if err := s.MCPServerHandler.SyncMCPServer(ctx); err != nil { logger.Warn("failed to sync MCP servers after adding client: %v", err) } return nil } // RefreshMCPClientTools re-discovers an MCP client's tools from its upstream // server on demand, so an operator who has just changed that server does not // have to wait out the connection checker's tool-sync interval (10 minutes by // default) or restart the gateway. Applies to every client type, including // the per-call ones ReconnectMCPClient rejects. // // The discovery itself persists the new tool set and re-syncs the hosted // /mcp surface through the tools-change callback, exactly like every other // discovery path, so there is nothing to sync here. func (s *BifrostHTTPServer) RefreshMCPClientTools(ctx context.Context, id string) (int, error) { return s.Client.RefreshMCPClientTools(ctx, id) } // UpdateMCPClient updates an MCP client in the in-memory store func (s *BifrostHTTPServer) UpdateMCPClient(ctx context.Context, id string, updatedConfig *schemas.MCPClientConfig) error { if err := s.Config.UpdateMCPClient(ctx, id, updatedConfig); err != nil { return err } if err := s.MCPServerHandler.SyncMCPServer(ctx); err != nil { logger.Warn("failed to sync MCP servers after editing client: %v", err) } return nil } // mcpDiscoveryRetryBackoff is the wait before retrying a failed discovery write: 1s doubling to 1m. var mcpDiscoveryRetryBackoff = func(attempt int) time.Duration { return min(time.Second< Virtual MCP assignment in the cache. func (s *BifrostHTTPServer) AttachVirtualMCPToVirtualKeyInMemory(ctx context.Context, vkID string, id uint) error { if cache := s.governanceVirtualMCPCache(); cache != nil { cache.AttachVirtualMCPInMemory(vkID, id) } return nil } // DetachVirtualMCPFromVirtualKeyInMemory removes a VK -> Virtual MCP assignment from the cache. func (s *BifrostHTTPServer) DetachVirtualMCPFromVirtualKeyInMemory(ctx context.Context, vkID string, id uint) error { if cache := s.governanceVirtualMCPCache(); cache != nil { cache.DetachVirtualMCPInMemory(vkID, id) } return nil } // ResetBudgetUsageInMemory clears usage for the given budgets in the governance // store that enforces spend, leaving each reset boundary untouched. // // The database write happens inside the update transaction through // UpdateBudgetUsage, but the in-memory store is what enforcement actually reads, // and ReloadVirtualKey deliberately carries the cached usage forward on a config // reload. Without this step the reset would be visible in the database and // nowhere else, and the next dump tick would write the cached value back over it. // // Missing budgets are not an error: a budget can legitimately have been deleted // in the same request that asked for the reset. func (s *BifrostHTTPServer) ResetBudgetUsageInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string) error { if len(budgetIDs) == 0 { return nil } governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } store := governancePlugin.GetGovernanceStore() for _, budgetID := range budgetIDs { if _, ok := store.ResetBudgetUsageInMemory(ctx, budgetID); !ok { logger.Debug("budget %s not present in the governance store; skipping usage reset", budgetID) } } return nil } // AdoptCalendarAlignmentInMemory re-anchors an owner's budgets and rate limits // onto their calendar boundaries after alignment was switched on. // // Only meaningful on the false-to-true transition. The in-memory store is what // the reset sweep reads, and a window that opened before the current boundary // reads as already due the moment the flag flips, so without this the operator // loses the usage they were still accumulating. Adoption moves the boundary // forward instead, which every persistence path allows, and the next dump tick // carries the new boundary to the database the same way an ordinary reset does. // // Missing entries are not an error: a budget or rate limit can legitimately have // been deleted in the same request that turned alignment on. func (s *BifrostHTTPServer) AdoptCalendarAlignmentInMemory(ctx context.Context, owner handlers.BudgetUsageResetOwner, budgetIDs []string, rateLimitIDs []string) error { if len(budgetIDs) == 0 && len(rateLimitIDs) == 0 { return nil } governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } store := governancePlugin.GetGovernanceStore() now := time.Now() for _, budgetID := range budgetIDs { if !store.AdoptCalendarAlignmentInMemory(ctx, budgetID, now) { logger.Debug("budget %s needed no calendar adoption (absent, unalignable, or already current)", budgetID) } } for _, rateLimitID := range rateLimitIDs { if !store.AdoptRateLimitCalendarAlignmentInMemory(ctx, rateLimitID, now) { logger.Debug("rate limit %s needed no calendar adoption (absent, unalignable, or already current)", rateLimitID) } } return nil } // RemoveVirtualKey removes a virtual key from the in-memory store func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) error { governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } preloadedVk, err := s.Config.ConfigStore.GetVirtualKey(ctx, id) if err != nil { if !errors.Is(err, configstore.ErrNotFound) { return err } } if preloadedVk == nil { // This could be broadcast message from other server, so we will just clean up in-memory store governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(ctx, id) s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(id) return nil } governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(ctx, id) s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(id) return nil } // ReloadTeam reloads a team from the in-memory store func (s *BifrostHTTPServer) ReloadTeam(ctx context.Context, id string) (*tables.TableTeam, error) { // Load relationships for response preloadedTeam, err := s.Config.ConfigStore.GetTeam(ctx, id) if err != nil { logger.Error("failed to load relationships for created team: %v", err) return nil, err } governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, err } // Add to in-memory store governancePlugin.GetGovernanceStore().UpdateTeamInMemory(ctx, preloadedTeam, nil) return preloadedTeam, nil } // RemoveTeam removes a team from the in-memory store func (s *BifrostHTTPServer) RemoveTeam(ctx context.Context, id string) error { governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } preloadedTeam, err := s.Config.ConfigStore.GetTeam(ctx, id) if err != nil { if !errors.Is(err, configstore.ErrNotFound) { return err } } if preloadedTeam == nil { // At-least deleting from in-memory store to avoid conflicts governancePlugin.GetGovernanceStore().DeleteTeamInMemory(ctx, id) return nil } governancePlugin.GetGovernanceStore().DeleteTeamInMemory(ctx, id) return nil } // ReloadCustomer reloads a customer from the in-memory store func (s *BifrostHTTPServer) ReloadCustomer(ctx context.Context, id string) (*tables.TableCustomer, error) { preloadedCustomer, err := s.Config.ConfigStore.GetCustomer(ctx, id) if err != nil { return nil, err } governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, err } // Add to in-memory store governancePlugin.GetGovernanceStore().UpdateCustomerInMemory(ctx, preloadedCustomer, nil) return preloadedCustomer, nil } // RemoveCustomer removes a customer from the in-memory store func (s *BifrostHTTPServer) RemoveCustomer(ctx context.Context, id string) error { governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } preloadedCustomer, err := s.Config.ConfigStore.GetCustomer(ctx, id) if err != nil { if !errors.Is(err, configstore.ErrNotFound) { return err } } if preloadedCustomer == nil { // At-least deleting from in-memory store to avoid conflicts governancePlugin.GetGovernanceStore().DeleteCustomerInMemory(ctx, id) return nil } governancePlugin.GetGovernanceStore().DeleteCustomerInMemory(ctx, id) return nil } // ModelConfigIndexKey reports the spelling the governance store indexes a model config // under (see LocalGovernanceStore.ModelConfigIndexKey); without the governance plugin it // falls back to the catalog-independent canonical name. func (s *BifrostHTTPServer) ModelConfigIndexKey(model string, provider *string) string { if governancePlugin, err := s.getGovernancePlugin(); err == nil && governancePlugin != nil { return governancePlugin.GetGovernanceStore().ModelConfigIndexKey(model, provider) } return governance.CanonicalModelConfigName(model, provider) } // ReloadModelConfig reloads a model config from the database into in-memory store // If usage was modified (e.g., reset due to config change), syncs it back to DB func (s *BifrostHTTPServer) ReloadModelConfig(ctx context.Context, id string) (*tables.TableModelConfig, error) { preloadedMC, err := s.Config.ConfigStore.GetModelConfigByID(ctx, id) if err != nil { logger.Error("failed to load model config: %v", err) return nil, err } governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, err } // Update in memory and get back the potentially modified model config updatedMC := governancePlugin.GetGovernanceStore().UpdateModelConfigInMemory(ctx, preloadedMC) if updatedMC == nil { return preloadedMC, nil } // Sync updated budget usage values back to database if they changed (per budget ID, // since a model config may own multiple budgets). preloadedUsage := make(map[string]float64, len(preloadedMC.Budgets)) for i := range preloadedMC.Budgets { preloadedUsage[preloadedMC.Budgets[i].ID] = preloadedMC.Budgets[i].CurrentUsage } for i := range updatedMC.Budgets { b := &updatedMC.Budgets[i] if old, ok := preloadedUsage[b.ID]; ok && old != b.CurrentUsage { if err := s.Config.ConfigStore.UpdateBudgetUsage(ctx, b.ID, b.CurrentUsage); err != nil { logger.Error("failed to sync budget usage to database: %v", err) } } } if updatedMC.RateLimit != nil && preloadedMC.RateLimit != nil { tokenUsageChanged := updatedMC.RateLimit.TokenCurrentUsage != preloadedMC.RateLimit.TokenCurrentUsage requestUsageChanged := updatedMC.RateLimit.RequestCurrentUsage != preloadedMC.RateLimit.RequestCurrentUsage if tokenUsageChanged || requestUsageChanged { if err := s.Config.ConfigStore.UpdateRateLimitUsage(ctx, updatedMC.RateLimit.ID, updatedMC.RateLimit.TokenCurrentUsage, updatedMC.RateLimit.RequestCurrentUsage); err != nil { logger.Error("failed to sync rate limit usage to database: %v", err) } } } return updatedMC, nil } // RemoveModelConfig removes a model config from the in-memory store func (s *BifrostHTTPServer) RemoveModelConfig(ctx context.Context, id string) error { governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } governancePlugin.GetGovernanceStore().DeleteModelConfigInMemory(ctx, id) return nil } // ReloadProvider reloads persisted provider settings into the live client. isNew says whether the // write that called it added the provider rather than edited one; the reload is the same either // way, and a server that wraps this one can tell an operator's edit apart by it. func (s *BifrostHTTPServer) ReloadProvider(ctx context.Context, provider schemas.ModelProvider, isNew bool) (*tables.TableProvider, error) { s.syncIgnoreProviderCost(provider) if s.Config == nil || s.Config.ConfigStore == nil { return nil, fmt.Errorf("config store not found") } if s.Config.ModelCatalog == nil { return nil, fmt.Errorf("pricing manager not found") } if s.Client == nil { return nil, fmt.Errorf("bifrost client not found") } // Load provider from DB providerInfo, err := s.Config.ConfigStore.GetProvider(ctx, provider) if err != nil { logger.Error("failed to load provider: %v", err) return nil, err } // Initialize updatedProvider updatedProvider := providerInfo // Sync model level budgets in governance plugin (if governance is enabled) if s.Config.IsPluginLoaded(s.getGovernancePluginName()) { governancePlugin, err := s.getGovernancePlugin() if err != nil { logger.Warn("governance plugin found but failed to get: %v", err) } else { // Update in memory and get back the potentially modified provider govUpdated := governancePlugin.GetGovernanceStore().UpdateProviderInMemory(ctx, providerInfo) if govUpdated != nil { updatedProvider = govUpdated } // Sync updated usage values back to database if they changed if updatedProvider.Budget != nil && providerInfo.Budget != nil { if updatedProvider.Budget.CurrentUsage != providerInfo.Budget.CurrentUsage { if err := s.Config.ConfigStore.UpdateBudgetUsage(ctx, updatedProvider.Budget.ID, updatedProvider.Budget.CurrentUsage); err != nil { logger.Error("failed to sync budget usage to database: %v", err) } } } if updatedProvider.RateLimit != nil && providerInfo.RateLimit != nil { tokenUsageChanged := updatedProvider.RateLimit.TokenCurrentUsage != providerInfo.RateLimit.TokenCurrentUsage requestUsageChanged := updatedProvider.RateLimit.RequestCurrentUsage != providerInfo.RateLimit.RequestCurrentUsage if tokenUsageChanged || requestUsageChanged { if err := s.Config.ConfigStore.UpdateRateLimitUsage(ctx, updatedProvider.RateLimit.ID, updatedProvider.RateLimit.TokenCurrentUsage, updatedProvider.RateLimit.RequestCurrentUsage); err != nil { logger.Error("failed to sync rate limit usage to database: %v", err) } } } } } // In-memory store holds the latest schemas.Key slice after the most recent // CRUD write — read from there to avoid re-fetching + re-converting from DB. inMemoryKeys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { return nil, fmt.Errorf("failed to read provider keys for %s: %w", provider, err) } isKeylessProvider := providerInfo.CustomProviderConfig != nil && providerInfo.CustomProviderConfig.IsKeyLess hasNoKeys := len(inMemoryKeys) == 0 && !isKeylessProvider // Refresh keyconfig from the current key list, then reconcile the live // entries against it before refetching per-key. setKeyConfig(s.Config.ModelCatalog, provider, inMemoryKeys, isKeylessProvider) if hasNoKeys { // Nothing left to route through, and nothing will ever refresh these // entries, so no cached response is still meaningful. s.Config.ModelCatalog.InvalidateLiveProvider(provider) logger.Warn("model discovery skipped for provider %s: no keys configured", provider) } else { // Prune only the keys that are gone or disabled and leave the // survivors' entries in place for the refetch below to overwrite. // FetchAndStoreLiveForKey writes nothing when the upstream call fails, // so wiping the provider up front turned any transient list-models // failure (5xx, rate limit, timeout, a restarting gateway) into an // empty catalog for a provider that was perfectly healthy a moment // earlier, with no background refresher to heal it. // // The trade is deliberate: entries survive a failed refetch even when // this reload changed base_url or the custom provider config, so they // can briefly describe the previous upstream. That is strictly the // smaller blast radius, since an emptied catalog makes every model // unroutable rather than just the ones that moved. keep := make(map[string]struct{}, len(inMemoryKeys)+1) if isKeylessProvider { // Keyless providers cache only under the empty-string sentinel, // which is also what the OnKey* helpers write under. Entries under // a residual key's ID are never refetched, so they are pruned. keep[""] = struct{}{} } else { for _, key := range inMemoryKeys { // Mirror RefreshLiveModelsForProvider's skip rule: a disabled key // is never refetched, so retaining its entries would serve models // for a key core rejects at routing time. if keyEnabled(key) { keep[key.ID] = struct{}{} } } } s.Config.ModelCatalog.RetainLiveKeys(provider, keep) s.RefreshLiveModelsForProvider(ctx, provider, inMemoryKeys) } return updatedProvider, nil } // syncIgnoreProviderCost pushes the provider's ignore_provider_cost setting into the model catalog. func (s *BifrostHTTPServer) syncIgnoreProviderCost(provider schemas.ModelProvider) { if s.Config == nil || s.Config.ModelCatalog == nil { return } pc, err := s.Config.GetProviderConfigRaw(provider) s.Config.ModelCatalog.SetIgnoreProviderCost(provider, err == nil && pc != nil && pc.IgnoreProviderCost) } // RemoveProvider removes a provider from the in-memory store func (s *BifrostHTTPServer) RemoveProvider(ctx context.Context, provider schemas.ModelProvider) error { err := s.Client.RemoveProvider(provider) if err != nil && !strings.Contains(err.Error(), "not found") { logger.Error("failed to remove provider from client: %v", err) return err } err = s.Config.RemoveProvider(ctx, provider) if err != nil && !errors.Is(err, lib.ErrNotFound) { logger.Error("failed to remove provider from config: %v. Client and config may be out of sync, please restart bifrost", err) return fmt.Errorf("failed to remove provider from config: %w. Client and config may be out of sync, please restart bifrost", err) } governancePlugin, err := s.getGovernancePlugin() if err != nil { return err } governancePlugin.GetGovernanceStore().DeleteProviderInMemory(ctx, string(provider)) if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("pricing manager not found") } s.Config.ModelCatalog.InvalidateLiveProvider(provider) s.Config.ModelCatalog.RemoveKeyConfigForProvider(provider) s.Config.ModelCatalog.SetIgnoreProviderCost(provider, false) return nil } // OnKeyAdded refreshes the keyconfig snapshot and fetches list-models for the // new key only — 2 calls instead of ReloadProvider's 2×N. Called by the key // handler after a successful AddProviderKey write. func (s *BifrostHTTPServer) OnKeyAdded(ctx context.Context, provider schemas.ModelProvider, key schemas.Key) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("model catalog not found") } keys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { return fmt.Errorf("failed to read provider keys for %s: %w", provider, err) } keyless := isKeylessProvider(provider, s.Config) setKeyConfig(s.Config.ModelCatalog, provider, keys, keyless) // Keyless providers: empty keyID sentinel. keyID := key.ID if keyless { keyID = "" } // Skip the fetch for a disabled key — core rejects list-models calls // scoped to a disabled key's ID, so it would just fail and fall back // onto the (usually empty, for custom providers) static datasheet. if !keyEnabled(key) { return nil } s.FetchAndStoreLiveForKey(ctx, provider, keyID) return nil } // OnKeyUpdated invalidates the affected key's live entries (the gate may have // changed even when Value didn't), refreshes the keyconfig, then refetches // for just that key. 2 calls regardless of N keys on the provider. func (s *BifrostHTTPServer) OnKeyUpdated(ctx context.Context, provider schemas.ModelProvider, key schemas.Key) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("model catalog not found") } keys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { return fmt.Errorf("failed to read provider keys for %s: %w", provider, err) } keyless := isKeylessProvider(provider, s.Config) setKeyConfig(s.Config.ModelCatalog, provider, keys, keyless) keyID := key.ID if keyless { keyID = "" } s.Config.ModelCatalog.InvalidateLive(provider, keyID) // Skip the fetch for a disabled key — the invalidate above still clears // its stale cached entries, but re-fetching would just fail against core. if !keyEnabled(key) { return nil } s.FetchAndStoreLiveForKey(ctx, provider, keyID) return nil } // OnKeyDeleted invalidates the deleted key's live entries and refreshes the // keyconfig. No list-models calls — the provider's remaining keys' cached // entries stay valid. func (s *BifrostHTTPServer) OnKeyDeleted(ctx context.Context, provider schemas.ModelProvider, keyID string) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("model catalog not found") } keys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { return fmt.Errorf("failed to read provider keys for %s: %w", provider, err) } setKeyConfig(s.Config.ModelCatalog, provider, keys, isKeylessProvider(provider, s.Config)) s.Config.ModelCatalog.InvalidateLive(provider, keyID) return nil } // keyEnabled reports whether a key should be treated as active for // scheduling model-discovery fetches. Enabled defaults to true when nil, // matching the convention core.getAllSupportedKeys uses to filter keys for // ListModels requests — a key discovery schedules a fetch for must be one // core will actually accept, or the call is a guaranteed // "no key found with id..." failure. func keyEnabled(key schemas.Key) bool { return key.Enabled == nil || *key.Enabled } // isKeylessProvider returns true when the provider's config marks it // keyless. Used to pick the live-cache key for OnKey* helpers: keyless // providers cache under the empty-string sentinel. func isKeylessProvider(provider schemas.ModelProvider, cfg *lib.Config) bool { if cfg == nil { return false } pc, err := cfg.GetProviderConfigRaw(provider) if err != nil || pc == nil || pc.CustomProviderConfig == nil { return false } return pc.CustomProviderConfig.IsKeyLess } // setKeyConfig pushes the provider's keys into the catalog's keyconfig. Only a // keyless custom provider stays routable with zero keys; any other provider // left without a key must stop claiming models, or a bare model name can be // routed to it and fail with no_key_supports_model. func setKeyConfig(catalog *modelcatalog.ModelCatalog, provider schemas.ModelProvider, keys []schemas.Key, keyless bool) { if keyless { catalog.SetKeyConfigForKeylessProvider(provider, keys) return } catalog.SetKeyConfigForProvider(provider, keys) } // GetGovernanceData returns the governance data func (s *BifrostHTTPServer) GetGovernanceData(ctx context.Context) *governance.GovernanceData { // Use type-safe finder from Config governancePlugin, err := lib.FindPluginAs[governance.BaseGovernancePlugin](s.Config, s.getGovernancePluginName()) if err != nil { return nil } return governancePlugin.GetGovernanceStore().GetGovernanceData(ctx) } // getRoutingPlugin safely retrieves the routing plugin, or an error. // // The lookup must name *RoutingPlugin: Init registers a pointer, and FindPluginAs // type-asserts against its type parameter, which the compiler cannot check when that // parameter is generic. Asserting against the value type builds fine and fails at runtime. // Every lookup failure — the plugin missing, or registered under a different type — is a // server-side misconfiguration, so all of them are wrapped in ErrRoutingPluginUnavailable // and answered as such by the handlers, rather than as a rejection of the caller's payload. func (s *BifrostHTTPServer) getRoutingPlugin() (*routing.RoutingPlugin, error) { plugin, err := lib.FindPluginAs[*routing.RoutingPlugin](s.Config, routing.PluginName) if err != nil { return nil, fmt.Errorf("%w: %w", handlers.ErrRoutingPluginUnavailable, err) } return plugin, nil } // ValidateComplexityAnalyzerConfig checks runtime-only semantic dependencies // before a handler persists a complexity configuration. func (s *BifrostHTTPServer) ValidateComplexityAnalyzerConfig(_ context.Context, config *complexity.AnalyzerConfig) error { routingPlugin, err := s.getRoutingPlugin() if err != nil { return fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.ValidateComplexityAnalyzerConfig(config) } // GetComplexitySemanticStatus returns the current semantic classifier readiness // from the routing plugin. func (s *BifrostHTTPServer) GetComplexitySemanticStatus(_ context.Context) (complexity.SemanticStatusInfo, error) { routingPlugin, err := s.getRoutingPlugin() if err != nil { return complexity.SemanticStatusInfo{}, fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.ComplexitySemanticStatus(), nil } // RetryComplexitySemanticWarmup restarts a failed semantic warmup using its // saved configuration and reports whether the classifier accepted the retry. func (s *BifrostHTTPServer) RetryComplexitySemanticWarmup(_ context.Context) (complexity.SemanticStatusInfo, bool, error) { routingPlugin, err := s.getRoutingPlugin() if err != nil { return complexity.SemanticStatusInfo{}, false, fmt.Errorf("routing plugin not found: %w", err) } status, retried := routingPlugin.RetryComplexitySemanticWarmup() return status, retried, nil } // ListComplexityGenerations reports the exemplar generations held in the // configured vector store. func (s *BifrostHTTPServer) ListComplexityGenerations(ctx context.Context) ([]complexity.GenerationInfo, error) { routingPlugin, err := s.getRoutingPlugin() if err != nil { return nil, fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.ListComplexityGenerations(ctx) } // DeleteComplexityGeneration removes one retired exemplar generation. func (s *BifrostHTTPServer) DeleteComplexityGeneration(ctx context.Context, namespace string) error { routingPlugin, err := s.getRoutingPlugin() if err != nil { return fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.DeleteComplexityGeneration(ctx, namespace) } // GetComplexityLLMStatus returns the llm fallback classifier's readiness from // the routing plugin. func (s *BifrostHTTPServer) GetComplexityLLMStatus(_ context.Context) (complexity.LLMStatusInfo, error) { routingPlugin, err := s.getRoutingPlugin() if err != nil { return complexity.LLMStatusInfo{}, fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.ComplexityLLMStatus(), nil } // ReloadComplexityAnalyzerConfig reloads the complexity analyzer config into the routing plugin. func (s *BifrostHTTPServer) ReloadComplexityAnalyzerConfig(ctx context.Context, config *complexity.AnalyzerConfig) error { routingPlugin, err := s.getRoutingPlugin() if err != nil { return fmt.Errorf("routing plugin not found: %w", err) } return routingPlugin.ReloadComplexityAnalyzerConfig(config) } // ReloadRoutingRule reloads a routing rule from the database into the routing plugin's rule cache func (s *BifrostHTTPServer) ReloadRoutingRule(ctx context.Context, id string) error { routingPlugin, err := s.getRoutingPlugin() if err != nil { return fmt.Errorf("routing plugin not found: %w", err) } rule, err := s.Config.ConfigStore.GetRoutingRule(ctx, id) if err != nil { return fmt.Errorf("failed to get routing rule from config store: %w", err) } // Update the rule in the rule cache if err := routingPlugin.GetRuleStore().UpsertRule(ctx, rule); err != nil { return fmt.Errorf("failed to update routing rule in rule store: %w", err) } return nil } // RemoveRoutingRule removes a routing rule from the routing plugin's rule cache func (s *BifrostHTTPServer) RemoveRoutingRule(ctx context.Context, id string) error { routingPlugin, err := s.getRoutingPlugin() if err != nil { return fmt.Errorf("routing plugin not found: %w", err) } if err := routingPlugin.GetRuleStore().DeleteRule(ctx, id); err != nil { return fmt.Errorf("failed to delete routing rule from rule store: %w", err) } return nil } // ReloadWebhookEndpoint refreshes a single webhook endpoint in the in-memory // store from the database after a mutation. A clustered deployment overrides // this to also notify peers so their in-memory copies stay current. func (s *BifrostHTTPServer) ReloadWebhookEndpoint(ctx context.Context, id string) error { endpoint, err := s.Config.ConfigStore.GetWebhookEndpointByID(ctx, id) if err != nil { return err } s.Config.SetWebhookEndpoint(endpoint) return nil } // RemoveWebhookEndpoint drops a webhook endpoint from the in-memory store after // a database delete. A clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) RemoveWebhookEndpoint(ctx context.Context, id string) error { s.Config.RemoveWebhookEndpoint(id) return nil } // EvictOauthTokenCacheByID drops the cached per-user MCP OAuth access token // backed by the given token row ID from the in-memory cache after a database // write. A clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) EvictOauthTokenCacheByID(ctx context.Context, tokenID string) { if s.Config == nil || s.Config.OAuthProvider == nil { return } s.Config.OAuthProvider.EvictUserTokenByID(tokenID) } // EvictOauthTokenCacheByMCPClient drops every cached per-user MCP OAuth // access token bound to the given MCP client from the in-memory cache. A // clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) { if s.Config == nil || s.Config.OAuthProvider == nil { return } s.Config.OAuthProvider.EvictUserTokensByMCPClient(mcpClientID) } // EvictOauthTokenCacheByVirtualKey drops every cached vk-mode MCP OAuth // access token bound to the given virtual key from the in-memory cache. A // clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) EvictOauthTokenCacheByVirtualKey(ctx context.Context, virtualKeyID string) { if s.Config == nil || s.Config.OAuthProvider == nil { return } s.Config.OAuthProvider.EvictUserTokensByVirtualKey(virtualKeyID) } // FlushOauthTokenCache drops every cached per-user MCP OAuth access token // from the in-memory cache. The coarse fallback for mutations whose blast // radius cannot be scoped to one client or virtual key. A clustered // deployment overrides this to also notify peers. func (s *BifrostHTTPServer) FlushOauthTokenCache(ctx context.Context) { if s.Config == nil || s.Config.OAuthProvider == nil { return } s.Config.OAuthProvider.FlushUserTokenCache() } // EvictMCPHeaderCredentialCacheByID drops the cached per-user MCP header // credential backed by the given credential row ID from the in-memory cache // after a database write. A clustered deployment overrides this to also // notify peers. func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByID(ctx context.Context, credentialID string) { if s.Config == nil || s.Config.MCPHeadersProvider == nil { return } s.Config.MCPHeadersProvider.EvictCredentialByID(credentialID) } // EvictMCPHeaderCredentialCacheByMCPClient drops every cached per-user MCP // header credential bound to the given MCP client from the in-memory cache. // A clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByMCPClient(ctx context.Context, mcpClientID string) { if s.Config == nil || s.Config.MCPHeadersProvider == nil { return } s.Config.MCPHeadersProvider.EvictCredentialsByMCPClient(mcpClientID) } // EvictMCPHeaderCredentialCacheByVirtualKey drops every cached vk-mode MCP // header credential bound to the given virtual key from the in-memory cache. // A clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByVirtualKey(ctx context.Context, virtualKeyID string) { if s.Config == nil || s.Config.MCPHeadersProvider == nil { return } s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(virtualKeyID) } // FlushMCPHeaderCredentialCache drops every cached per-user MCP header // credential from the in-memory cache. The coarse fallback for mutations // whose blast radius cannot be scoped to one client or virtual key. A // clustered deployment overrides this to also notify peers. func (s *BifrostHTTPServer) FlushMCPHeaderCredentialCache(ctx context.Context) { if s.Config == nil || s.Config.MCPHeadersProvider == nil { return } s.Config.MCPHeadersProvider.FlushCredentialCache() } // ReloadClientConfigFromConfigStore reloads the client config from config store func (s *BifrostHTTPServer) ReloadClientConfigFromConfigStore(ctx context.Context) error { if s.Config == nil || s.Config.ConfigStore == nil { return fmt.Errorf("config store not found") } config, err := s.Config.ConfigStore.GetClientConfig(context.Background()) if err != nil { return fmt.Errorf("failed to get client config: %v", err) } if config == nil { return fmt.Errorf("client config not found") } s.Config.Mu.Lock() *s.Config.ClientConfig = *config s.Config.Mu.Unlock() if s.LogsCleaner != nil { s.LogsCleaner.UpdateRetentionDays(config.LogRetentionDays) } // This path also serves cluster peers, so one node's allowlist edit // reaches every node's OAuth client. if err := lib.ApplyMCPOAuthPrivateAllowlist(config.MCPOAuthPrivateAllowlist); err != nil { logger.Warn("ignoring invalid mcp_oauth_private_allowlist during client config reload: %v", err) } // Reloading whitelisted routes from the client config if s.AuthMiddleware != nil { s.AuthMiddleware.UpdateWhitelistedRoutes(config.WhitelistedRoutes) s.AuthMiddleware.UpdateAllowedOrigins(config.AllowedOrigins) s.AuthMiddleware.UpdateTempTokenAuthEnabled(config.MCPEnableTempTokenAuth) } // Refresh the CORS middleware's immutable snapshot so its requests pick up the // new AllowedOrigins/AllowedHeaders/DumpErrorsInConsoleLogs without racing the // in-place ClientConfig mutation above. if s.CORSMiddleware != nil { s.CORSMiddleware.UpdateConfig(s.Config) } // Reloading config in bifrost client if s.Client != nil { account := lib.NewBaseAccount(s.Config) var mcpConfig *schemas.MCPConfig if s.Config.MCPConfig != nil { mcpConfig = s.Config.MCPConfig } s.Client.ReloadConfig(schemas.BifrostConfig{ Account: account, InitialPoolSize: s.Config.ClientConfig.InitialPoolSize, DropExcessRequests: s.Config.ClientConfig.DropExcessRequests, LLMPlugins: s.Config.GetLoadedLLMPlugins(), MCPPlugins: s.Config.GetLoadedMCPPlugins(), A2APlugins: s.Config.GetLoadedA2APlugins(), RelayPlugins: s.Config.GetLoadedRelayPlugins(), MCPConfig: mcpConfig, Logger: logger, }) if err := s.Client.UpdateToolManagerConfig( s.Config.ClientConfig.MCPAgentDepth, s.Config.ClientConfig.MCPToolExecutionTimeout, s.Config.ClientConfig.MCPCodeModeBindingLevel, s.Config.ClientConfig.MCPDisableAutoToolInject, s.Config.ClientConfig.MCPMaxInstructionsPerClient, s.Config.ClientConfig.MCPMaxInstructionsTotal, ); err != nil { logger.Warn("failed to sync MCP tool manager config during client config reload: %v", err) } if err := s.Client.UpdateCodeModeLimits(s.Config.ClientConfig.MCPCodeModeLimits); err != nil { logger.Warn("failed to sync code mode limits during client config reload: %v", err) } // The global tool sync interval is client-config-backed too (minutes). // Re-time the checkers of every client that follows it, and keep the // in-memory MCPConfig aligned so a later reload carries the current // value. This path also serves cluster peers reloading the client // config, so one node's edit re-times every node. syncInterval := time.Duration(s.Config.ClientConfig.MCPToolSyncInterval) * time.Minute if s.Config.MCPConfig != nil { s.Config.MCPConfig.ToolSyncInterval = syncInterval } if err := s.Client.UpdateMCPToolSyncInterval(syncInterval); err != nil { logger.Warn("failed to sync MCP tool sync interval during client config reload: %v", err) } if s.Config.MCPConfig != nil { s.Config.MCPConfig.MaxProtocolVersion = s.Config.ClientConfig.MCPMaxProtocolVersion } if err := s.Client.UpdateMCPMaxProtocolVersion(s.Config.ClientConfig.MCPMaxProtocolVersion); err != nil { logger.Warn("failed to sync MCP max protocol version during client config reload: %v", err) } } return nil } // UpdateAuthConfig updates auth config in the config store and updates the AuthMiddleware's in-memory config func (s *BifrostHTTPServer) UpdateAuthConfig(ctx context.Context, authConfig *configstore.AuthConfig) error { if authConfig == nil { return fmt.Errorf("auth config is nil") } if s.Config == nil || s.Config.ConfigStore == nil { return fmt.Errorf("config store not found") } // Allow disabling auth without credentials, but require them when enabling if authConfig.IsEnabled && (authConfig.AdminUserName == nil || authConfig.AdminUserName.GetValue() == "" || authConfig.AdminPassword == nil || authConfig.AdminPassword.GetValue() == "") { return fmt.Errorf("username and password are required when auth is enabled") } // Update the config store if err := s.Config.ConfigStore.UpdateAuthConfig(ctx, authConfig); err != nil { return err } // Update the AuthMiddleware's in-memory config if s.AuthMiddleware != nil { // Fetch the updated config from the store to ensure we have the latest updatedAuthConfig, err := s.Config.ConfigStore.GetAuthConfig(ctx) if err != nil { logger.Warn("failed to get auth config from store after update: %v", err) // Still update with what we have s.AuthMiddleware.UpdateAuthConfig(authConfig) } else { s.AuthMiddleware.UpdateAuthConfig(updatedAuthConfig) } // The first admin account now exists (or already did) — permanently // invalidate the bootstrap setup token. No-op if already cleared. s.AuthMiddleware.ClearBootstrapToken() } return nil } // ValidateSetupToken reports whether token is a valid one-time bootstrap setup token, // as required to create the very first admin account (see AuthMiddleware.bootstrapToken). // Once an admin account exists, this always returns true — the gate closes permanently. func (s *BifrostHTTPServer) ValidateSetupToken(token string) bool { if s.AuthMiddleware == nil { return true } return s.AuthMiddleware.CheckBootstrapToken(token) } // ValidateConfiguredSetupToken reports whether token matches the operator-configured setup // token (see AuthMiddleware.CheckConfiguredSetupToken). Unlike ValidateSetupToken it does not // open up once an admin account exists, and it is false when no token is configured or no // auth middleware is installed. func (s *BifrostHTTPServer) ValidateConfiguredSetupToken(token string) bool { if s.AuthMiddleware == nil { return false } return s.AuthMiddleware.CheckConfiguredSetupToken(token) } // UpdateDropExcessRequests updates excess requests config func (s *BifrostHTTPServer) UpdateDropExcessRequests(ctx context.Context, value bool) { if s.Config == nil { return } s.Client.UpdateDropExcessRequests(value) } // UpdateMCPToolManagerConfig updates the MCP tool manager config. // Always pass the current disableAutoToolInject value so it is not reset by an update that // only meant to change something else. func (s *BifrostHTTPServer) UpdateMCPToolManagerConfig(ctx context.Context, maxAgentDepth int, toolExecutionTimeoutInSeconds int, codeModeBindingLevel string, disableAutoToolInject bool, maxInstructionsPerClient int, maxInstructionsTotal int, codeModeLimits *schemas.MCPCodeModeLimits) error { if s.Config == nil { return fmt.Errorf("config not found") } if err := s.Client.UpdateToolManagerConfig(maxAgentDepth, toolExecutionTimeoutInSeconds, codeModeBindingLevel, disableAutoToolInject, maxInstructionsPerClient, maxInstructionsTotal); err != nil { return err } return s.Client.UpdateCodeModeLimits(codeModeLimits) } // reloadObservabilityPlugins reloads all observability plugins in the tracing middleware func (s *BifrostHTTPServer) reloadObservabilityPlugins() { observabilityPlugins := s.CollectObservabilityPlugins() // Always update the tracing middleware, even with empty slice, to clear stale plugins s.TracingMiddleware.SetObservabilityPlugins(observabilityPlugins, s.CollectObservabilityLimits()) } // ReloadPricingManager reloads the pricing manager func (s *BifrostHTTPServer) UpdateSyncConfig(ctx context.Context) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("pricing manager not found") } // Snapshot the pricing pointer and release the read lock before restarting // the refresher: that path takes Config.Mu.RLock again through // LiveModelsSyncInterval, and sync.RWMutex read locks are not reentrant. // Safe to use after unlocking because the config update path publishes a // whole new FrameworkConfig rather than mutating this struct in place. s.Config.Mu.RLock() var pricing *modelcatalog.Config if s.Config.FrameworkConfig != nil { pricing = s.Config.FrameworkConfig.Pricing } s.Config.Mu.RUnlock() if pricing == nil { return fmt.Errorf("framework config not found") } // The live model refresher reads its interval from the same framework // config block, so pick up a changed (or newly disabled) value here rather // than waiting for a restart. s.RestartLiveModelRefresher(s.backgroundCtx()) return s.Config.ModelCatalog.UpdateSyncConfig(ctx, pricing) } // ResolveAccess answers what the request may reach, so a route that never enters the request // pipeline decides from the same answer the pipeline would have used: resolved once, onto the // request's grant. // // Without governance there is nothing to resolve and nothing to restrict, which is why a // missing plugin returns nil rather than access permitting nothing. func (s *BifrostHTTPServer) ResolveAccess(ctx *schemas.BifrostContext) (schemas.Access, error) { governancePlugin, err := s.getGovernancePlugin() if err != nil { logger.Warn("governance plugin not found: %v", err) return nil, nil } return governancePlugin.ResolveAccess(ctx) } // NarrowListModelsProviders implements AccessResolver, so the native models route and the // integration ones narrow their fan-out by one rule rather than each carrying a copy of it. // // A request with nothing resolved keeps the fan-out it already had rather than being narrowed to // nothing, and so does a request nothing settled who it is: publishing an empty list here would // mean "no provider may serve this", turning an unrestricted listing into one that lists nothing. func (s *BifrostHTTPServer) NarrowListModelsProviders(ctx *schemas.BifrostContext) { if ctx == nil { return } access, err := s.ResolveAccess(ctx) if err != nil || access == nil { return } granted := access.GrantedProvidersForModel("") providers := make([]schemas.ModelProvider, 0, len(granted)) for _, provider := range granted { providers = append(providers, schemas.ModelProvider(provider)) } ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, providers) } // AdmitMCPGatewayRequest runs a /mcp request through the governance funnel before any tool is listed // or called, as an MCP tool execution with no tool named yet. Evaluating is what resolves and records // the request's access, so what tools/list shows and what tools/call permits come from one answer. // // Without governance there is nothing to evaluate against and nothing to restrict, which is why a // missing plugin answers with no access rather than with access permitting nothing. func (s *BifrostHTTPServer) AdmitMCPGatewayRequest(ctx *schemas.BifrostContext) (schemas.Access, *schemas.BifrostError) { governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, nil } if _, refused := governancePlugin.Evaluate(ctx, &governance.EvaluationRequest{RequestType: schemas.MCPToolExecutionRequest}); refused != nil { return nil, refused } // Evaluating refuses a request that carries no grant, so one that was admitted carries the // access it was admitted with. return ctx.Grant().Access(), nil } // VirtualMCPToolAccess resolves what a /mcp/ request may see, delegating to the governance store // (which reads the addressable-vMCP set recorded on ctx during resolution). No governance store → the // endpoint is refused. func (s *BifrostHTTPServer) VirtualMCPToolAccess(ctx *schemas.BifrostContext, slug string, access schemas.Access) (served []string, assigned bool) { governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, false } resolver, ok := governancePlugin.GetGovernanceStore().(interface { VirtualMCPToolAccess(ctx *schemas.BifrostContext, slug string, access schemas.Access) ([]string, bool) }) if !ok { return nil, false } return resolver.VirtualMCPToolAccess(ctx, slug, access) } // MCPClientToolAccess resolves the tools a /mcp/ request may see for a single MCP client, gated // on the caller's grant (ok=false → the endpoint is refused). Sibling of VirtualMCPToolAccess. func (s *BifrostHTTPServer) MCPClientToolAccess(ctx *schemas.BifrostContext, slug string, access schemas.Access) (served []string, ok bool) { governancePlugin, err := s.getGovernancePlugin() if err != nil { return nil, false } resolver, ok := governancePlugin.GetGovernanceStore().(interface { MCPClientToolAccess(ctx *schemas.BifrostContext, slug string, access schemas.Access) ([]string, bool) }) if !ok { return nil, false } return resolver.MCPClientToolAccess(ctx, slug, access) } // backgroundCtx returns the server-lifetime context background workers should // hang off. Falls back to context.Background() before Bootstrap has set s.Ctx. func (s *BifrostHTTPServer) backgroundCtx() context.Context { if s.Ctx != nil { return s.Ctx } return context.Background() } // liveRefreshProviderConcurrency bounds how many providers a single // RefreshAllLiveModels pass fetches at once. The bootstrap seed could fan out // across every provider unbounded because nothing else was in flight yet; the // background pass runs alongside live inference traffic, and each provider it // picks up itself fans out to (enabled keys x 2) upstream calls. const liveRefreshProviderConcurrency = 4 // RefreshAllLiveModels runs a full per-provider live model refresh. // // Safe to call while providers are being edited, unlike the bootstrap seed it // replaced: the provider set and each provider's keys are read through the // Config accessors, which take Config.Mu.RLock and hand back copies, rather // than ranging Config.Providers directly. // // Deliberately does not invalidate anything first. Key and provider removals // already prune the cache through the CRUD handlers, and // FetchAndStoreLiveForKey writes nothing when a fetch fails, so leaving // entries in place keeps last-known-good models routable across a transient // upstream error instead of emptying a catalog that was healthy a moment ago. func (s *BifrostHTTPServer) RefreshAllLiveModels(ctx context.Context) { if s.Config == nil || s.Config.ModelCatalog == nil || s.Client == nil { return } providers, err := s.Config.GetAllProviders() if err != nil { logger.Warn("live model refresh skipped: failed to list providers: %v", err) return } sem := make(chan struct{}, liveRefreshProviderConcurrency) var wg sync.WaitGroup for _, provider := range providers { keys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { // Removed between GetAllProviders and this read. OnProviderRemoved // has already dropped its live entries, so there is nothing to do. continue } wg.Add(1) go func(p schemas.ModelProvider, k []schemas.Key) { defer wg.Done() sem <- struct{}{} defer func() { <-sem }() s.RefreshLiveModelsForProvider(ctx, p, k) }(provider, keys) } wg.Wait() } // liveRefreshJitterFraction spreads each refresh cycle over // interval * [1-f, 1+f]. Pods in a deployment start within seconds of each // other, so an exact interval would have every replica fan out to every // upstream at the same instant, on every cycle, forever. const liveRefreshJitterFraction = 0.1 // jitteredInterval scales d by a uniform random factor around 1. func jitteredInterval(d time.Duration) time.Duration { factor := 1 - liveRefreshJitterFraction + 2*liveRefreshJitterFraction*rand.Float64() return time.Duration(float64(d) * factor) } // LiveModelsSyncInterval returns the configured background refresh interval. // A non-positive result means the refresher is disabled. func (s *BifrostHTTPServer) LiveModelsSyncInterval() time.Duration { if s.Config == nil { return modelcatalog.DefaultLiveModelsSyncInterval } // FrameworkConfig is swapped wholesale by the config update handler, so the // pointer read needs Config.Mu like every other accessor on this struct. s.Config.Mu.RLock() defer s.Config.Mu.RUnlock() if s.Config.FrameworkConfig == nil || s.Config.FrameworkConfig.Pricing == nil { return modelcatalog.DefaultLiveModelsSyncInterval } configured := s.Config.FrameworkConfig.Pricing.LiveModelsSyncInterval if configured == nil { return modelcatalog.DefaultLiveModelsSyncInterval } return time.Duration(*configured) * time.Second } // startLiveModelRefresher runs RefreshAllLiveModels on a jittered interval and // returns a func that stops it. A non-positive interval disables the refresher // and returns a no-op, which is the documented opt-out. // // Deliberately takes no distributed lock, unlike the pricing sync worker: the // live model cache is per-process in-memory, so every node has to refresh its // own copy. Electing a single refresher would leave every other node serving // its boot-time snapshot indefinitely. func (s *BifrostHTTPServer) startLiveModelRefresher(ctx context.Context, interval time.Duration) func() { if interval <= 0 { logger.Info("background live model refresh is disabled") return func() {} } logger.Info("background live model refresh every %v (jittered +/-%d%%)", interval, int(liveRefreshJitterFraction*100)) ctx, cancel := context.WithCancel(ctx) go func() { // A timer rather than a ticker: the delay is resampled every cycle, so // replicas that happen to align on one cycle drift apart on the next. timer := time.NewTimer(jitteredInterval(interval)) defer timer.Stop() for { select { case <-ctx.Done(): return case <-timer.C: s.RefreshAllLiveModels(ctx) timer.Reset(jitteredInterval(interval)) } } }() return cancel } // RestartLiveModelRefresher stops any running refresher and starts a new one // at the currently configured interval. Called at boot and again whenever the // framework config changes, so a disabled-to-enabled edit (or vice versa) // takes effect without a restart. Exported so servers that replace the boot // path can start the refresher with its stop handle registered for the shared // cleanup and config-change restarts. func (s *BifrostHTTPServer) RestartLiveModelRefresher(ctx context.Context) { s.liveModelRefresherMu.Lock() defer s.liveModelRefresherMu.Unlock() if s.liveModelRefresherStop != nil { s.liveModelRefresherStop() } s.liveModelRefresherStop = s.startLiveModelRefresher(ctx, s.LiveModelsSyncInterval()) } // stopLiveModelRefresher halts the background refresher if one is running. func (s *BifrostHTTPServer) stopLiveModelRefresher() { s.liveModelRefresherMu.Lock() defer s.liveModelRefresherMu.Unlock() if s.liveModelRefresherStop != nil { s.liveModelRefresherStop() s.liveModelRefresherStop = nil } } type providerRefreshState struct { allKeys bool keys map[string]struct{} } // beginKeyRefresh claims an on-demand refresh slot for one provider key. // Refreshes for distinct keys may run concurrently, but the same key is // deduplicated and no key refresh may overlap an all-keys refresh. func (s *BifrostHTTPServer) beginKeyRefresh(provider schemas.ModelProvider, keyID string) (release func(), ok bool) { s.refreshInFlightMu.Lock() defer s.refreshInFlightMu.Unlock() if s.refreshInFlight == nil { s.refreshInFlight = make(map[schemas.ModelProvider]*providerRefreshState) } state := s.refreshInFlight[provider] if state == nil { state = &providerRefreshState{keys: make(map[string]struct{})} s.refreshInFlight[provider] = state } if state.allKeys { return nil, false } if _, exists := state.keys[keyID]; exists { return nil, false } state.keys[keyID] = struct{}{} return func() { s.refreshInFlightMu.Lock() defer s.refreshInFlightMu.Unlock() current := s.refreshInFlight[provider] if current != state { return } delete(state.keys, keyID) if len(state.keys) == 0 { delete(s.refreshInFlight, provider) } }, true } // beginAllKeysRefresh claims an exclusive on-demand refresh slot for a // provider. It conflicts with both another all-keys refresh and any active // per-key refresh for that provider. func (s *BifrostHTTPServer) beginAllKeysRefresh(provider schemas.ModelProvider) (release func(), ok bool) { s.refreshInFlightMu.Lock() defer s.refreshInFlightMu.Unlock() if s.refreshInFlight == nil { s.refreshInFlight = make(map[schemas.ModelProvider]*providerRefreshState) } if _, exists := s.refreshInFlight[provider]; exists { return nil, false } state := &providerRefreshState{allKeys: true} s.refreshInFlight[provider] = state return func() { s.refreshInFlightMu.Lock() defer s.refreshInFlightMu.Unlock() if s.refreshInFlight[provider] == state { delete(s.refreshInFlight, provider) } }, true } // RefreshLiveModelsForKey re-fetches list-models for one key on demand. It is // the ModelsManager entrypoint behind the per-key refresh button. func (s *BifrostHTTPServer) RefreshLiveModelsForKey(ctx context.Context, provider schemas.ModelProvider, keyID string) error { if s.Config == nil || s.Config.ModelCatalog == nil || s.Client == nil { return fmt.Errorf("model catalog is not initialized") } release, ok := s.beginKeyRefresh(provider, keyID) if !ok { return handlers.ErrRefreshInProgress } defer release() s.FetchAndStoreLiveForKey(ctx, provider, keyID) return nil } // RefreshLiveModelsForAllKeys re-fetches list-models across every enabled key // of the provider on demand. It is the ModelsManager entrypoint behind the // provider-level refresh button. func (s *BifrostHTTPServer) RefreshLiveModelsForAllKeys(ctx context.Context, provider schemas.ModelProvider) error { if s.Config == nil || s.Config.ModelCatalog == nil || s.Client == nil { return fmt.Errorf("model catalog is not initialized") } keys, err := s.Config.GetProviderKeysRaw(provider) if err != nil { return err } release, ok := s.beginAllKeysRefresh(provider) if !ok { return handlers.ErrRefreshInProgress } defer release() s.RefreshLiveModelsForProvider(ctx, provider, keys) return nil } // RefreshLiveModelsForProvider runs filtered + unfiltered list-models for the // provider, fanning out per key in parallel so the live cache ends up with // per-(provider, keyID) entries. Keyless providers cache under the "" sentinel. // // Callers are responsible for invalidating stale entries first when keys // have been removed from the provider's set. func (s *BifrostHTTPServer) RefreshLiveModelsForProvider(ctx context.Context, provider schemas.ModelProvider, keys []schemas.Key) { // Keyless providers fetch once under the "" sentinel whatever keys are // still stored: core skips key selection for them entirely, and a provider // switched from keyed to keyless keeps its old (often disabled) keys, which // must not decide whether discovery runs. if isKeylessProvider(provider, s.Config) { s.FetchAndStoreLiveForKey(ctx, provider, "") return } if len(keys) == 0 { // Empty key slice + non-keyless provider would write under the "" sentinel // reserved for keyless providers — colliding with the keyless namespace and // triggering an unauthenticated fetch for a provider that requires a key. logger.Warn("model discovery skipped for provider %s: no keys configured", provider) return } var wg sync.WaitGroup enabledCount := 0 for _, key := range keys { // Skip disabled keys — core rejects list-models calls scoped to a // disabled key's ID ("no key found with id..."), so fetching for one // is guaranteed to fail and only adds "falling back onto the static // datasheet" log noise per disabled key. if !keyEnabled(key) { continue } enabledCount++ wg.Add(1) go func(keyID string) { defer wg.Done() s.FetchAndStoreLiveForKey(ctx, provider, keyID) }(key.ID) } if enabledCount == 0 { logger.Warn("model discovery skipped for provider %s: no enabled keys configured", provider) return } wg.Wait() } // FetchAndStoreLiveForKey issues the filtered and unfiltered list-models // calls for one (provider, keyID) in parallel and writes the results into // the catalog. Errors are logged and surfaced via updateKeyStatus when the // provider returns per-key statuses, but they do not abort the other call. // keyID="" scopes to "no specific key" — used for keyless providers and as // the legacy sentinel. Always validates keys for the providers that opt into // the check (today: OpenRouter, whose /v1/models is unauthenticated) so the // routing graph is the same at boot, after a key add, and after a reload — // stale-but-routable behavior would diverge otherwise. func (s *BifrostHTTPServer) FetchAndStoreLiveForKey(ctx context.Context, provider schemas.ModelProvider, keyID string) { if s.Config == nil || s.Config.ModelCatalog == nil { return } // Skip the fetch entirely when the provider has disabled list_models via // allowed_requests — every per-(provider,keyID) call would just bounce with // "operation not allowed", wasting two goroutines and one bfCtx per attempt. if pc, err := s.Config.GetProviderConfigRaw(provider); err == nil && pc != nil && pc.CustomProviderConfig != nil && !pc.CustomProviderConfig.IsOperationAllowed(schemas.ListModelsRequest) { return } // Captured before the upstream calls go out, and checked again by the // catalog when each response commits. list-models can take seconds; a key // deleted or disabled in that window has already had its entries dropped by // the CRUD handler, and no later refresh pass fetches or prunes a key that // is no longer configured — so an unguarded commit here would re-add that // key's models and advertise them until the process restarted. gen := s.Config.ModelCatalog.LiveGeneration(provider) // One BifrostContext per goroutine. BifrostContext.SetValue mutates state // in place, so the request-scoped metadata core sets during a routing pass // (RequestID, FallbackIndex, span IDs, ...) would otherwise bleed between // the filtered and unfiltered calls and conflate them in logs/billing. newListModelsCtx := func() *schemas.BifrostContext { c := schemas.NewBifrostContext(ctx, time.Now().Add(15*time.Second)) c.SetValue(schemas.BifrostContextKeySkipPluginPipeline, true) c.SetValue(schemas.BifrostContextKeyValidateKeys, true) return c } var keyIDPtr *string if keyID != "" { keyIDPtr = &keyID } var wg sync.WaitGroup wg.Add(2) go func() { defer wg.Done() bfCtx := newListModelsCtx() defer bfCtx.Cancel() resp, bfErr := s.Client.ListModelsRequest(bfCtx, &schemas.BifrostListModelsRequest{ Provider: provider, KeyID: keyIDPtr, }) if bfErr != nil { logger.Warn("filtered list-models failed for provider %s key %s: %v: falling back onto the static datasheet", provider, keyID, bifrost.GetErrorMessage(bfErr)) if len(bfErr.ExtraFields.KeyStatuses) > 0 && s.Config.ConfigStore != nil { s.updateKeyStatus(ctx, bfErr.ExtraFields.KeyStatuses) } return } if resp == nil { return } if !s.Config.ModelCatalog.UpsertLiveFromResponseIfCurrent(provider, keyID, false, resp, gen) { // The key or provider changed while this call was in flight, so the // result describes a configuration that no longer exists. Dropping // the key statuses with it is deliberate: they are keyed by the same // key ID and would write validity back onto a row that was just // deleted or replaced. logger.Debug("discarding stale filtered list-models result for provider %s key %s: the provider's keys changed while the fetch was in flight", provider, keyID) return } if len(resp.KeyStatuses) > 0 && s.Config.ConfigStore != nil { s.updateKeyStatus(ctx, resp.KeyStatuses) } }() go func() { defer wg.Done() bfCtx := newListModelsCtx() defer bfCtx.Cancel() resp, bfErr := s.Client.ListModelsRequest(bfCtx, &schemas.BifrostListModelsRequest{ Provider: provider, KeyID: keyIDPtr, Unfiltered: true, }) if bfErr != nil { logger.Warn("unfiltered list-models failed for provider %s key %s: %v: falling back onto the static datasheet", provider, keyID, bifrost.GetErrorMessage(bfErr)) return } if resp == nil { return } if !s.Config.ModelCatalog.UpsertLiveFromResponseIfCurrent(provider, keyID, true, resp, gen) { logger.Debug("discarding stale unfiltered list-models result for provider %s key %s: the provider's keys changed while the fetch was in flight", provider, keyID) } }() wg.Wait() } // ForceReloadPricing triggers an immediate pricing sync and resets the sync // timer. No longer triggers a list-models refresh — pricing reload is now // pricing-only. func (s *BifrostHTTPServer) ForceReloadPricing(ctx context.Context) error { if s.Config == nil { return fmt.Errorf("server config not initialized") } if s.Config.ModelCatalog != nil { if err := s.Config.ModelCatalog.ForceReloadPricing(ctx); err != nil { return fmt.Errorf("failed to force reload pricing: %w", err) } } return nil } // ReloadPricingFromDBAndPopulateModelPool reloads the pricing from DB. The // list-models refresh that used to follow is gone — pricing reload is now // pricing-only. func (s *BifrostHTTPServer) ReloadPricingFromDBAndPopulateModelPool(ctx context.Context) error { if s.Config == nil { return fmt.Errorf("server config not initialized") } if s.Config.ModelCatalog != nil { if err := s.Config.ModelCatalog.ReloadFromDB(ctx); err != nil { return fmt.Errorf("failed to reload pricing from DB: %w", err) } } return nil } // UpsertPricingOverride inserts or updates a pricing override in the in-memory model catalog. func (s *BifrostHTTPServer) UpsertPricingOverride(ctx context.Context, override *tables.TablePricingOverride) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("pricing manager not found") } return s.Config.ModelCatalog.UpsertPricingOverrides(override) } // DeletePricingOverride removes a pricing override from the in-memory model catalog. func (s *BifrostHTTPServer) DeletePricingOverride(ctx context.Context, id string) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("pricing manager not found") } s.Config.ModelCatalog.DeletePricingOverride(id) return nil } // UpsertModelPricingAttributes writes the additional_attributes JSON for the // pricing rows keyed by (model, provider) for every entry in the batch. The // whole batch is wrapped in a single transaction so a missing pricing row // rolls back the lot. After a successful commit the in-memory pricing cache // is reloaded once. Enterprise overrides this method to broadcast a peer // reload after commit. func (s *BifrostHTTPServer) UpsertModelPricingAttributes(ctx context.Context, entries []handlers.ModelPricingAttributesEntry) error { if s.Config == nil || s.Config.ModelCatalog == nil { return fmt.Errorf("model catalog not initialized") } if s.Config.ConfigStore == nil { return fmt.Errorf("model catalog requires a config store") } var missing []string err := s.Config.ConfigStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { for _, e := range entries { rows, err := s.Config.ConfigStore.UpsertModelPricingAttributes(ctx, e.Model, e.Provider, e.AdditionalAttributes, tx) if err != nil { return err } if rows == 0 { missing = append(missing, fmt.Sprintf("%s/%s", e.Provider, e.Model)) } } if len(missing) > 0 { return fmt.Errorf("no pricing row for one or more (model, provider) entries: %s", strings.Join(missing, ", ")) } return nil }) if err != nil { return err } if err := s.Config.ModelCatalog.ReloadPricing(ctx); err != nil { return fmt.Errorf("failed to reload pricing cache after attribute write: %w", err) } return nil } // ReloadProxyConfig reloads the proxy configuration func (s *BifrostHTTPServer) ReloadProxyConfig(ctx context.Context, config *tables.GlobalProxyConfig) error { if s.Config == nil { return fmt.Errorf("config not found") } // Point every factory client at the new proxy, then store the config and rebuild // the providers that inherit it for inference. The config's factory is the one // registered as the process default; a server that runs its own bootstrap // (enterprise) may never set s.HTTPClientFactory, so update both. if s.Config.HTTPClientFactory != nil { s.Config.HTTPClientFactory.UpdateProxyConfig(config.ToNetwork()) } if s.HTTPClientFactory != nil && s.HTTPClientFactory != s.Config.HTTPClientFactory { s.HTTPClientFactory.UpdateProxyConfig(config.ToNetwork()) } if err := s.Config.SetGlobalProxyConfig(config); err != nil { return err } s.rebuildRelays(ctx) if config == nil { logger.Info("proxy configuration removed") return nil } logger.Info("proxy configuration reloaded: enabled=%t, type=%s", config.Enabled, config.Type) return nil } // ReloadHeaderFilterConfig reloads the header filter configuration func (s *BifrostHTTPServer) ReloadHeaderFilterConfig(ctx context.Context, config *tables.GlobalHeaderFilterConfig) error { if s.Config == nil { return fmt.Errorf("config not found") } // Store the raw header filter config in ClientConfig s.Config.ClientConfig.HeaderFilterConfig = config // Compile into optimized matcher for O(1) per-request lookups s.Config.SetHeaderMatcher(lib.NewHeaderMatcher(config)) allowlistLen := 0 denylistLen := 0 if config != nil { allowlistLen = len(config.Allowlist) denylistLen = len(config.Denylist) } logger.Info("header filter configuration reloaded: allowlist=%d, denylist=%d", allowlistLen, denylistLen) return nil } // GetModelsForProvider returns all models for a specific provider from the model catalog func (s *BifrostHTTPServer) GetModelsForProvider(provider schemas.ModelProvider) []string { if s.Config == nil || s.Config.ModelCatalog == nil { return []string{} } return s.Config.ModelCatalog.GetModelsForProvider(provider) } // GetUnfilteredModelsForProvider returns all unfiltered models for a specific provider from the model catalog func (s *BifrostHTTPServer) GetUnfilteredModelsForProvider(provider schemas.ModelProvider) []string { if s.Config == nil || s.Config.ModelCatalog == nil { return []string{} } return s.Config.ModelCatalog.GetUnfilteredModelsForProvider(provider) } // GetPluginStatus returns the status of all plugins // Delegates to Config for centralized plugin status management func (s *BifrostHTTPServer) GetPluginStatus(ctx context.Context) map[string]schemas.PluginStatus { return s.Config.GetPluginStatus() } // GetLoadedPluginNames returns the sanitized names of all currently loaded plugins, // matching the names embedded in their trace span names. func (s *BifrostHTTPServer) GetLoadedPluginNames() []string { if s.Config == nil { return []string{} } return s.Config.GetLoadedPluginNames() } // NormalizePluginConfig implements handlers.PluginsLoader. It looks up the plugin // by name in the ConfigMarshallers cache and calls MarshalConfigForStorage if found. // Returns nil, nil when the plugin is not loaded or does not implement ConfigMarshallerPlugin. func (s *BifrostHTTPServer) NormalizePluginConfig(name string, config map[string]any) (map[string]any, error) { if m := s.Config.ConfigMarshallers.Load(); m != nil { if cm, ok := (*m)[name]; ok { return cm.MarshalConfigForStorage(config) } } return nil, nil } // ExpandPluginConfigForAPI implements handlers.PluginsLoader. It looks up the plugin // by name in the ConfigMarshallers cache and calls RedactConfig if found. // Returns nil, nil when the plugin is not loaded or does not implement ConfigMarshallerPlugin. func (s *BifrostHTTPServer) ExpandPluginConfigForAPI(name string, config map[string]any) (map[string]any, error) { if m := s.Config.ConfigMarshallers.Load(); m != nil { if cm, ok := (*m)[name]; ok { return cm.RedactConfig(config) } } return nil, nil } // Helper to update error status // Uses UpdatePluginOverallStatus to create the status entry if it doesn't exist, // ensuring plugins that were never loaded can still have their error status tracked. // Always returns the original error so the actual failure reason is surfaced to the user. func (s *BifrostHTTPServer) updatePluginErrorStatus(name, step string, originalErr error) error { logs := []string{fmt.Sprintf("error %s plugin %s: %v", step, name, originalErr)} s.Config.UpdatePluginOverallStatus(name, name, schemas.PluginStatusError, logs, []schemas.PluginType{}) return originalErr } // SyncLoadedPlugin syncs a loaded plugin to the Bifrost client and updates the plugin status func (s *BifrostHTTPServer) SyncLoadedPlugin(ctx context.Context, name string, plugin schemas.BasePlugin, placement *schemas.PluginPlacement, order *int) error { // 2. Register (replaces old version atomically) if err := s.Config.ReloadPlugin(plugin); err != nil { return s.updatePluginErrorStatus(plugin.GetName(), "registering", err) } // 2b. Set order info and re-sort s.Config.SetPluginOrderInfo(plugin.GetName(), placement, order) s.Config.SortAndRebuildPlugins() // 3. Update Bifrost client if err := s.Client.ReloadPlugin(plugin, InferPluginTypes(plugin)); err != nil { return s.updatePluginErrorStatus(plugin.GetName(), "reloading bifrost config for", err) } // 3b. Sync plugin execution order from config to core s.Client.ReorderPlugins(s.Config.GetPluginOrder()) // 4. Special handling for observability plugins if _, ok := plugin.(schemas.ObservabilityPlugin); ok { s.reloadObservabilityPlugins() } // 4b. Batch accounting is wired at bootstrap against the live logging and // governance plugins. Reloading either swaps the instance — which cancels the // old sweeper and leaves the sweeper holding a torn-down usage reporter — so // it must be rewired here or batch accounting silently stops until restart. // Safe to re-run: StartBatchAccountingSweeper cancels any prior sweeper. if plugin.GetName() == logging.PluginName || plugin.GetName() == s.getGovernancePluginName() { s.WireBatchAccountingSweeper() } // 5. Update plugin status s.Config.UpdatePluginOverallStatus(plugin.GetName(), name, schemas.PluginStatusActive, []string{fmt.Sprintf("plugin %s reloaded successfully", name)}, InferPluginTypes(plugin)) return nil } // PluginInitContext returns the context to hand to a plugin constructor. A plugin outlives // the request that loaded it and derives long-lived contexts from this one, so it must be the // server-lifetime context, never the pooled *fasthttp.RequestCtx an admin reload arrives on. // Carries the bootstrap plugin-name values, which Ctx inherits. func (s *BifrostHTTPServer) PluginInitContext() context.Context { if s.Ctx == nil { return context.Background() } return s.Ctx } // ReloadPlugin reloads a plugin with new instance and updates Bifrost core. // The plugin is checked for LLM and MCP interfaces independently and registered // to the appropriate arrays based on which interfaces it implements. func (s *BifrostHTTPServer) ReloadPlugin(ctx context.Context, name string, path *string, pluginConfig any, placement *schemas.PluginPlacement, order *int) error { logger.Debug("reloading plugin %s", name) // 0. A semantic_cache config gets the same validation config.json does, here rather than when // it is saved: the plugins API loads a plugin only when it is enabled, so a cache that no // longer passes can still be disabled. A refused config changes nothing, status included. if name == semanticcache.PluginName { if err := s.Config.ValidateSemanticCacheConfig(&schemas.PluginConfig{Name: name, Config: pluginConfig}); err != nil { return fmt.Errorf("%w: %w", lib.ErrInvalidPluginConfig, err) } } // 1. Instantiate new version. When it fails, the running instance (if any) // is left in place, so the status keeps describing that instance; only a // plugin with nothing running is marked as errored. plugin, err := InstantiatePlugin(s.PluginInitContext(), name, path, pluginConfig, s.Config) if err != nil { if s.Config.IsPluginLoaded(name) { logger.Warn("plugin %s failed to load with its new config, the running instance stays in place: %v", name, err) return err } return s.updatePluginErrorStatus(name, "loading", err) } // Wire the embedding executor on the new instance before syncing. if semanticCachePlugin, ok := plugin.(*semanticcache.Plugin); ok { semanticCachePlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } // Both at once: applied separately, the classifier spends the gap between // them configured for a store it has not been given, and warms a throwaway // generation into the embedded one. if routingComplexityPlugin, ok := plugin.(routing.ComplexityWarmupDependencySetter); ok { routingComplexityPlugin.SetComplexityWarmupDependencies(s.Config.VectorStore, s.Client.EmbeddingRequest) } if routingWarmupObserverPlugin, ok := plugin.(routing.WarmupEmbedUsageObserverSetter); ok { routingWarmupObserverPlugin.SetWarmupEmbedUsageObserver(s.ObserveWarmupRoutingEmbedding) } if routingChatPlugin, ok := plugin.(interface { SetChatRequestExecutor(routing.ChatRequestExecutor) }); ok { routingChatPlugin.SetChatRequestExecutor(s.Client.ChatCompletionRequest) } if routingResponsesPlugin, ok := plugin.(routing.ResponsesExecutorSetter); ok { routingResponsesPlugin.SetResponsesRequestExecutor(s.Client.ResponsesRequest) } if routingDecisionPlugin, ok := plugin.(routing.DecisionExecutorSetter); ok { routingDecisionPlugin.SetDecisionRequestExecutor(s.Client.DecisionRequest) } if err := s.SyncLoadedPlugin(ctx, name, plugin, placement, order); err != nil { return err } // Rebind Warp's reader onto the live service rather than rebuilding the // handler: routes were registered at startup and hold the original // handler's closures, so a fresh handler would never be reached. Without // this, enabling logging at runtime leaves the chat endpoint answering 503 // until the process restarts. // // After the sync, not before: SyncLoadedPlugin replaces the configured // plugin and can then fail in Client.ReloadPlugin, which returns without // rolling back. Rebinding first left Warp researching through the new log // manager while the Bifrost client was still running the previous plugin. if loggerPlugin, ok := plugin.(*logging.LoggerPlugin); ok && s.WarpHandler != nil { s.WarpHandler.Service().SetLogReader(handlers.NewWarpLogReader(loggerPlugin.GetPluginLogManager())) } return nil } // RemovePlugin removes a plugin from the server. // The plugin is removed from both LLM and MCP arrays independently if it exists in them. func (s *BifrostHTTPServer) RemovePlugin(ctx context.Context, displayName string) error { // Get the actual plugin name from the display name name, ok := s.Config.GetPluginNameByDisplayName(displayName) if !ok { return dynamicPlugins.ErrPluginNotFound } // Check if plugin implements ObservabilityPlugin before removal var isObservability bool var err error var plugin schemas.BasePlugin if plugin, err = s.Config.FindPluginByName(name); err == nil { _, isObservability = plugin.(schemas.ObservabilityPlugin) } // 1. Unregister from config if err := s.Config.UnregisterPlugin(name); err != nil { return err } // 2. Update Bifrost client if err := s.Client.RemovePlugin(name, InferPluginTypes(plugin)); err != nil { logger.Warn("failed to reload bifrost config after plugin removal: %v", err) } // 3. Reload observability plugins if necessary if isObservability { s.reloadObservabilityPlugins() } // 3a. Warp's reader belongs to the plugin that was just removed. AddPlugin // rebinds it on the way in, and without the matching clear on the way out // Warp kept researching through a log manager nobody owns any more - and // CanChat stayed true, because it only asks whether logs is non-nil. if _, ok := plugin.(*logging.LoggerPlugin); ok && s.WarpHandler != nil { s.WarpHandler.Service().SetLogReader(nil) } // 4. Update status and marshaller if isDisabled, _ := ctx.Value(handlers.PluginDisabledKey).(bool); isDisabled { s.markPluginDisabled(name) } else { s.Config.DeletePluginOverallStatus(name) // Plugin is being permanently deleted: remove its config marshaller too. s.Config.RemoveConfigMarshaller(name) } return nil } // StartOAuth2SweepWorker creates and starts the janitor for the OAuth2 // issuance tables (expired authorize requests, aged-out revoked refresh // tokens, orphaned dynamically-registered clients). Call it once from // single-threaded bootstrap wiring, like the other worker fields on this // struct — the nil-check makes double-wiring a no-op, it is not a concurrency // guard. It is also a no-op when no config store is configured. The Start() // shutdown path stops the worker. // // shouldSweep, when non-nil, is consulted before each pass; returning false // skips that pass. Deployments running several instances against one config // store can use it to restrict sweeping to a single instance. nil means // always sweep. func (s *BifrostHTTPServer) StartOAuth2SweepWorker(ctx context.Context, shouldSweep func() bool) { if s.OAuth2SweepWorker != nil || s.Config == nil || s.Config.ConfigStore == nil { return } s.OAuth2SweepWorker = newOAuth2SweepWorker(s.Config.ConfigStore, shouldSweep) s.OAuth2SweepWorker.start(ctx) } var errAgentGatewayNotInitialized = errors.New("agent gateway is not initialized") // CreateAgentRegistration persists a registration and updates its local runtime. func (s *BifrostHTTPServer) CreateAgentRegistration(ctx context.Context, req agent.CreateRequest) (schemas.AgentRegistrationView, error) { if s.agentGatewayManager == nil { return schemas.AgentRegistrationView{}, errAgentGatewayNotInitialized } return s.agentGatewayManager.Create(ctx, req) } // UpdateAgentRegistration replaces a registration and updates its local runtime. func (s *BifrostHTTPServer) UpdateAgentRegistration(ctx context.Context, name string, req agent.UpdateRequest) (schemas.AgentRegistrationView, error) { if s.agentGatewayManager == nil { return schemas.AgentRegistrationView{}, errAgentGatewayNotInitialized } return s.agentGatewayManager.Update(ctx, name, req) } // DeleteAgentRegistration deletes a registration and retires its local runtime. func (s *BifrostHTTPServer) DeleteAgentRegistration(ctx context.Context, name string) error { if s.agentGatewayManager == nil { return errAgentGatewayNotInitialized } return s.agentGatewayManager.Delete(ctx, name) } // InitializeAgentGateway constructs the shared Agent Gateway handler before any // routes that depend on it are registered. registrationManager owns registration // mutations, while vkReloader refreshes virtual keys after grant changes. nil // values use the local OSS implementations. func (s *BifrostHTTPServer) InitializeAgentGateway(ctx context.Context, registrationManager handlers.AgentRegistrationManager, vkReloader handlers.VirtualKeyReloader) error { if s.AgentGatewayHandler != nil || s.Config == nil || s.Config.ConfigStore == nil { return nil } if vkReloader == nil { vkReloader = s } gatewayStore, ok := s.Config.ConfigStore.(agent.Store) if !ok { return errors.New("config store does not support agent gateway persistence") } managerConfig := agent.ManagerConfig{ AuthPolicy: handlers.AgentGatewayAuthPolicy(s.Config), // Read the external URL live so admin config changes (and their cluster // propagation) apply to card generation and push without a restart. ExternalURLProvider: s.Config.GetA2AExternalClientURL, RawCaptureProvider: s.Config.GetA2AStoreRawRequestResponse, PrivateAllowlistProvider: s.Config.GetA2APrivateAllowlist, } if s.Config.ServerConfig != nil { managerConfig.GRPCBaseDomain = s.Config.ServerConfig.A2AGRPCBaseDomain managerConfig.GRPCPort = s.Config.ServerConfig.A2AGRPCPort managerConfig.AllowPrivatePushCallbacks = s.Config.ServerConfig.A2AAllowPrivatePushCallbacks } if s.Client != nil { managerConfig.PluginPipelineAcquire = func() agent.PluginPipeline { return s.Client.AcquirePluginPipeline() } managerConfig.PluginPipelineRelease = func(pipeline agent.PluginPipeline) { if pluginPipeline, ok := pipeline.(*bifrost.PluginPipeline); ok { s.Client.ReleasePluginPipeline(pluginPipeline) } } } managerConfig.TracerProvider = func() schemas.Tracer { if s.TracingMiddleware == nil { return nil } return s.TracingMiddleware.GetTracer() } manager, err := agent.NewManager(ctx, gatewayStore, logger, s.Config.GetA2AExternalClientURL(), nil, managerConfig) if err != nil { return fmt.Errorf("failed to initialize agent gateway: %v", err) } s.AgentGatewayHandler = handlers.NewAgentGatewayHandler(manager, registrationManager, s.Config, vkReloader) s.agentGatewayManager = manager if err := s.AgentGatewayHandler.SyncEnabledAgents(ctx); err != nil { manager.Close() s.AgentGatewayHandler = nil s.agentGatewayManager = nil return fmt.Errorf("failed to initialize enabled agent configuration: %v", err) } agentCount := 0 if registrations, err := manager.List(ctx); err == nil { agentCount = len(registrations) } pushStatus := "disabled (external client URL not configured)" if s.Config.GetA2AExternalClientURL() != "" { pushStatus = "enabled" } logger.Info("agent gateway initialized: %d registered agent(s), push notifications %s", agentCount, pushStatus) return nil } // virtualKeyCache serves by-ID virtual key lookups from the governance in-memory // store (avoiding a per-request DB read) for the handlers that resolve an // identity to a virtual key. Best-effort: any governance store that exposes // GetVirtualKeyByID qualifies; otherwise those handlers fall back to a store read. func (s *BifrostHTTPServer) virtualKeyCache() handlers.VirtualKeyCache { // s.Ctx is nil until Bootstrap runs; plugin lookup reads it, so callers that // initialize handlers earlier simply get no cache and fall back to a store read. if s.Ctx == nil { return nil } if gp, gerr := s.getGovernancePlugin(); gerr == nil && gp != nil { if c, ok := gp.GetGovernanceStore().(handlers.VirtualKeyCache); ok { return c } } return nil } // RegisterInferenceRoutes initializes the routes for the inference handler func (s *BifrostHTTPServer) RegisterInferenceRoutes(ctx context.Context, middlewares ...schemas.BifrostHTTPMiddleware) error { // Initialize WebSocket pool and handler before integrations so it can be wired through s.wsPool = bfws.NewPool(s.Config.WebSocketConfig.Pool) wsResponsesHandler := handlers.NewWSResponsesHandler(s.Client, s.Config, s.wsPool) wsRealtimeHandler := handlers.NewWSRealtimeHandler(s.Client, s.Config, s.wsPool) s.wsLiveHandler = handlers.NewWSLiveHandler(s.Client, s.Config, s.wsPool) s.webrtcLiveHandler = handlers.NewWebRTCLiveHandler(s.Client, s.Config) webrtcRealtimeHandler := handlers.NewWebRTCRealtimeHandler(s.Client, s.Config) realtimeClientSecretsHandler := handlers.NewRealtimeClientSecretsHandler(s.Client, s.Config) inferenceHandler := handlers.NewInferenceHandler(s, s.Client, s.Config) s.IntegrationHandler = handlers.NewIntegrationHandler(s.Client, s.Config, s, wsResponsesHandler, wsRealtimeHandler, webrtcRealtimeHandler, realtimeClientSecretsHandler) mcpInferenceHandler := handlers.NewMCPInferenceHandler(s.Client, s.Config) // Serve by-ID virtual key lookups on the /mcp JWT auth path from the // governance in-memory store (avoiding a per-request DB read). Best-effort: // any store that exposes GetVirtualKeyByID qualifies; otherwise the handler // falls back to the config store. vkCache := s.virtualKeyCache() mcpServerHandler, err := handlers.NewMCPServerHandler(ctx, s.Config, s, s, s.OAuth2IdentityResolver, vkCache) if err != nil { return fmt.Errorf("failed to initialize mcp server handler: %v", err) } s.MCPServerHandler = mcpServerHandler asyncHandler := handlers.NewAsyncHandler(s.Client, s.Config) s.IntegrationHandler.RegisterRoutes(s.Router, middlewares...) s.wsLiveHandler.RegisterRoutes(s.Router, middlewares...) s.webrtcLiveHandler.RegisterRoutes(s.Router, middlewares...) handlers.NewLiveControlHandler(s.Client, s.Config).RegisterRoutes(s.Router, middlewares...) inferenceHandler.RegisterRoutes(s.Router, middlewares...) asyncHandler.RegisterRoutes(s.Router, middlewares...) mcpInferenceHandler.RegisterRoutes(s.Router, middlewares...) s.MCPServerHandler.RegisterRoutes(s.Router, middlewares...) if s.relays != nil { handlers.NewRelayGatewayHandler(s.relays, s.Config, logger).RegisterRoutes(s.Router, middlewares...) } if s.AgentGatewayHandler != nil { protocolMiddlewares := append([]schemas.BifrostHTTPMiddleware{}, middlewares...) validator, _ := s.virtualKeyCache().(handlers.AgentGatewayVirtualKeyValidator) protocolMiddlewares = append(protocolMiddlewares, handlers.AgentGatewayAuthenticationMiddleware(s.Config, validator, s.AgentGatewayIdentityResolver)) s.AgentGatewayHandler.RegisterProtocolRoutes(s.Router, protocolMiddlewares...) // The push ingress deliberately skips the Agent Gateway authentication // middleware: the upstream agent authenticates with the per-config // token Bifrost minted, not with a downstream virtual key. s.AgentGatewayHandler.RegisterPushIngressRoute(s.Router, middlewares...) // The shared gRPC binding starts here, alongside the protocol routes, // because it uses the same virtual-key validator. Routing is by dialed // authority; see handlers.AgentGatewayGRPCServer. if s.agentGatewayGRPCServer == nil && s.Config.ServerConfig != nil && s.Config.ServerConfig.A2AGRPCBaseDomain != "" && s.Config.ServerConfig.A2AGRPCPort > 0 { grpcServer := handlers.NewAgentGatewayGRPCServer(s.agentGatewayManager, s.Config, validator, s.Config.ServerConfig.A2AGRPCBaseDomain, s.AgentGatewayIdentityResolver, s.TracingMiddleware) if grpcServer != nil { // Cards always advertise A2AGRPCPort (the public endpoint). The // local bind port may differ when a load balancer or port mapping // fronts this process, e.g. multiple nodes on one host in the dev // cluster harness. listenPort := s.Config.ServerConfig.A2AGRPCPort if raw := os.Getenv("BIFROST_A2A_GRPC_LISTEN_PORT"); raw != "" { parsed, err := strconv.Atoi(raw) if err != nil || parsed <= 0 || parsed > 65535 { return fmt.Errorf("invalid BIFROST_A2A_GRPC_LISTEN_PORT %q", raw) } listenPort = parsed } if err := grpcServer.Start(listenPort); err != nil { return fmt.Errorf("failed to start agent gateway grpc listener: %v", err) } s.agentGatewayGRPCServer = grpcServer logger.Info("agent gateway grpc listener started on port %d (advertised port %d) for *.%s", listenPort, s.Config.ServerConfig.A2AGRPCPort, s.Config.ServerConfig.A2AGRPCBaseDomain) } } } return nil } // notificationPublisher returns a publisher that reads s.Config.NotificationPublisher at // call time, so a handler registered before the notification service exists (RegisterAPIRoutes // without Bootstrap) still reaches it. Publishing is a no-op while none is set. func (s *BifrostHTTPServer) notificationPublisher() schemas.NotificationPublisher { return func(ctx context.Context, input schemas.NotificationInput) (*schemas.Notification, error) { publish := s.Config.NotificationPublisher if publish == nil { return nil, nil } return publish(ctx, input) } } // staleMCPOAuthDropper is what dropStaleMCPOAuthState needs from the server. type staleMCPOAuthDropper interface { EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) CloseAndMarkNeedsReauth(ctx context.Context, id string) error } // dropStaleMCPOAuthState drops an MCP client's cached tokens and shared connection after the OAuth // provider rotated its credentials. It goes through callbacks so enterprise peers drop theirs too, and // finishes before the sign-in it rotated for starts: a late drop would close that sign-in's connection. func dropStaleMCPOAuthState(callbacks staleMCPOAuthDropper) func(ctx context.Context, mcpClientID string) { return func(ctx context.Context, mcpClientID string) { ctx = context.WithoutCancel(ctx) callbacks.EvictOauthTokenCacheByMCPClient(ctx, mcpClientID) if err := callbacks.CloseAndMarkNeedsReauth(ctx, mcpClientID); err != nil && !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { logger.Warn("failed to close MCP client %s's connection after its OAuth credentials rotated: %v", mcpClientID, err) } } } // RegisterAPIRoutes initializes the routes for the Bifrost HTTP server. func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks ServerCallbacks, middlewares ...schemas.BifrostHTTPMiddleware) error { var err error // Initializing plugin specific handlers var loggingHandler *handlers.LoggingHandler loggerPlugin, _ := lib.FindPluginAs[*logging.LoggerPlugin](s.Config, logging.PluginName) var govLogManager logging.LogManager if loggerPlugin != nil { loggingHandler = handlers.NewLoggingHandler(loggerPlugin.GetPluginLogManager(), s, s.Config) if resolverProvider, ok := callbacks.(LogRedactionMappingResolverProvider); ok { loggingHandler.SetLogRedactionMappingResolver(resolverProvider.GetLogRedactionMappingResolver()) } if resolverProvider, ok := callbacks.(MCPLogRedactionMappingResolverProvider); ok { loggingHandler.SetMCPLogRedactionMappingResolver(resolverProvider.GetMCPLogRedactionMappingResolver()) } // Wire the sidekiq runner so cost recalculation runs as a durable background // job. Registering the handler here (before RecoverIncomplete) lets a job // interrupted by a restart resume on boot. if s.SidekiqRunner != nil && s.Config != nil && s.Config.ConfigStore != nil { loggingHandler.SetSidekiqBackend(s.SidekiqRunner, s.Config.ConfigStore) } govLogManager = loggerPlugin.GetPluginLogManager() } var governanceHandler *handlers.GovernanceHandler governancePluginName := governance.PluginName if name, ok := ctx.Value(schemas.BifrostContextKeyGovernancePluginName).(string); ok && name != "" { governancePluginName = name } governancePlugin, _ := lib.FindPluginAs[schemas.LLMPlugin](s.Config, governancePluginName) if governancePlugin != nil { governanceHandler, err = handlers.NewGovernanceHandler(callbacks, s.Config.ConfigStore, govLogManager, s.ExternalQuotaBudgetResolver, s.VirtualKeyAssigneeResolver, s.VirtualKeyBusinessUnitResolver) if err != nil { return fmt.Errorf("failed to initialize governance handler: %v", err) } // Register the expired-key cleanup job before the dispatcher starts, since // any node may claim it. if s.SidekiqRunner != nil && s.Config != nil && s.Config.ConfigStore != nil { governanceHandler.SetExpiryCleanupBackend(s.SidekiqRunner, s.notificationPublisher()) } s.GovernanceHandler = governanceHandler } // Routing rules and the complexity analyzer config live in the config store, so these // endpoints have nothing to serve when persistence is disabled (initStores leaves // ConfigStore nil for config_store.enabled=false). Skip them rather than failing, which // would abort registration for every other API route too. var routingHandler *handlers.RoutingHandler if s.Config.ConfigStore != nil { routingHandler, err = handlers.NewRoutingHandler(callbacks, s.Config.ConfigStore) if err != nil { return fmt.Errorf("failed to initialize routing handler: %v", err) } } // Resolve the semantic_cache plugin per request so plugin reloads via // /api/plugins are honored — the previous boot-time capture left stale // references and (worse) skipped route registration entirely when the // plugin wasn't in config.json at startup, causing 405 on all cache-clear // endpoints for the process lifetime. cacheHandler := handlers.NewCacheHandler(func() handlers.CacheClearer { p, err := lib.FindPluginAs[*semanticcache.Plugin](s.Config, semanticcache.PluginName) if err != nil || p == nil { return nil } return p }) // Websocket handler needs to go below UI handler logger.Debug("initializing websocket server") if s.WebSocketHandler == nil { s.WebSocketHandler = handlers.NewWebSocketHandler(s.Ctx, s.Config.ClientConfig.AllowedOrigins) } if s.NotificationService == nil { s.NotificationService = handlers.NewNotificationService(s.Config.ConfigStore, s.WebSocketHandler) s.Config.NotificationPublisher = s.NotificationService.Publish s.NotificationService.Start(s.Ctx) } if s.RSSFeedsHandler == nil { s.RSSFeedsHandler = handlers.NewRSSFeedsHandler(s.Config.ConfigStore, nil) s.RSSFeedsHandler.Start(s.Ctx) } // This is the first point in Bootstrap where the logging plugin - and so // the log manager Warp's chat route needs - is known, which is why Warp is // built here rather than earlier. // // A nil log manager here is a supported deployment (logging disabled), not // a failure - Warp then serves only its config routes. // // The nil check stays even though nothing builds a prior instance today: // RegisterAPIRoutes has exactly one caller now, but if that ever changes - // a config or plugin reload re-running it - this must not leak the // previous instance's subscriptions and worker pool. if s.WarpHandler != nil { s.WarpHandler.Shutdown() } s.WarpHandler = handlers.NewWarpHandler(s.Config.ConfigStore, loggerPlugin, s.Client, s.Config.LogsStore, s.Config.VectorStore, s.SidekiqRunner, s.Config.ModelCatalog, logger, func() bool { return s.Config.FeatureFlags != nil && s.Config.FeatureFlags.IsEnabled(lib.FeatureFlagWarp) }, handlers.WarpResolvers{ ExternalQuotaBudgets: s.ExternalQuotaBudgetResolver, VirtualKeyAssignees: s.VirtualKeyAssigneeResolver, UserGovernance: s.WarpUserGovernanceReader, CallerRestriction: s.WarpCallerRestrictionResolver, }) if loggingHandler != nil { loggingHandler.SetTopicLister(s.WarpHandler.Service()) } // Read at each sweep rather than once: the cleaner is built in Bootstrap, // after this, and its retention can change at runtime. s.WarpHandler.Service().SetLogRetention(func() int { if s.LogsCleaner == nil { return 0 } return s.LogsCleaner.RetentionDays() }) // config.json's warp section saves through the service, so it is validated and provisions its // vector namespace exactly like a save from the settings page. s.Config.ApplyWarpFileConfig(s.Ctx, s.WarpHandler.Service()) // Start WebSocket heartbeat s.WebSocketHandler.StartHeartbeat() // Adding telemetry middleware // Chaining all middlewares // lib.ChainMiddlewares chains multiple middlewares together healthHandler := handlers.NewHealthHandler(s.Config) providerHandler := handlers.NewProviderHandler(callbacks, s.Config, s.Client) oauthHandler := handlers.NewOAuthHandler(s.Config.OAuthProvider, s.Client, s.Config) mcpHandler := handlers.NewMCPHandler(callbacks, callbacks, s.Client, s.Config, oauthHandler, callbacks) // Virtual MCP routes read the config store on every call; skip them when persistence is disabled // (nil ConfigStore) rather than nil-deref, like routingHandler below. var virtualMCPHandler *handlers.VirtualMCPHandler if s.Config.ConfigStore != nil { virtualMCPHandler = handlers.NewVirtualMCPHandler(s.Config, s.Client, callbacks, logger) } mcpPerUserHeadersHandler := handlers.NewMCPPerUserHeadersHandler(callbacks, s.Config, s.TempTokens) mcpSessionsHandler := handlers.NewMCPSessionsHandler(s.Config, callbacks) // The headers provider evicts a credential it writes only from this // process's cache. Forward those evictions to callbacks so a deployment // that overrides EvictMCPHeaderCredentialCacheByID reaches every process, // off the write path so a slow listener never delays the write. if s.Config.MCPHeadersProvider != nil { s.Config.MCPHeadersProvider.SetOnCredentialWritten(func(ctx context.Context, credentialID string) { go callbacks.EvictMCPHeaderCredentialCacheByID(context.WithoutCancel(ctx), credentialID) }) } // Same for token rows the OAuth provider writes. Off the write path // matters more here: the hook also fires inside refreshes that other // callers are waiting on. if s.Config.OAuthProvider != nil { s.Config.OAuthProvider.SetOnTokenWritten(func(ctx context.Context, tokenID string) { go callbacks.EvictOauthTokenCacheByID(context.WithoutCancel(ctx), tokenID) }) // A rotation the provider makes itself drops what the update API drops after its own, inline as there. s.Config.OAuthProvider.SetCredentialsRotatedHook(dropStaleMCPOAuthState(callbacks)) } configHandler := handlers.NewConfigHandler(callbacks, s.Config) pluginsHandler := handlers.NewPluginsHandler(callbacks, s.Config.ConfigStore) sessionHandler := handlers.NewSessionHandler(s.Config.ConfigStore, s.WSTicketStore) if s.setupLockInstalled && s.AuthMiddleware != nil { sessionHandler.SetSetupLock(s.AuthMiddleware) } promptsHandler := handlers.NewPromptsHandler(s.Config.ConfigStore, callbacks) featureFlagsHandler := handlers.NewFeatureFlagsHandler(s.Config.FeatureFlags, s.Config.ConfigStore) if s.AgentGatewayHandler != nil { s.AgentGatewayHandler.RegisterManagementRoutes(s.Router, middlewares...) } if s.RelaysHandler != nil { s.RelaysHandler.RegisterRoutes(s.Router, middlewares...) } // Going ahead with API handlers oauth2DiscoveryHandler := handlers.NewOAuth2DiscoveryHandler(s.Config) oauth2IssuanceHandler := handlers.NewOAuth2IssuanceHandler(s.Config, s.TempTokens, s.OAuth2IdentityResolver) if s.OAuth2Extensions != nil { oauth2DiscoveryHandler.SetExtensions(s.OAuth2Extensions) oauth2IssuanceHandler.SetExtensions(s.OAuth2Extensions) } oauth2SessionsHandler := handlers.NewOAuth2SessionsHandler(s.Config) oauth2ConsentHandler := handlers.NewOAuth2ConsentHandler(s.Config, s.TempTokens, s.OAuth2IdentityResolver) oauth2DiscoveryHandler.RegisterRoutes(s.Router, middlewares...) // No middleware needed for mcp issuance routes, they should be open oauth2IssuanceHandler.RegisterRoutes(s.Router) oauth2SessionsHandler.RegisterRoutes(s.Router, middlewares...) oauth2ConsentHandler.RegisterRoutes(s.Router, middlewares...) healthHandler.RegisterRoutes(s.Router, middlewares...) providerHandler.RegisterRoutes(s.Router, middlewares...) handlers.NewGithubCopilotSignInHandler().RegisterRoutes(s.Router, middlewares...) mcpHandler.RegisterRoutes(s.Router, middlewares...) if virtualMCPHandler != nil { virtualMCPHandler.RegisterRoutes(s.Router, middlewares...) } mcpPerUserHeadersHandler.RegisterRoutes(s.Router, middlewares...) mcpSessionsHandler.RegisterRoutes(s.Router, middlewares...) configHandler.RegisterRoutes(s.Router, middlewares...) oauthHandler.RegisterRoutes(s.Router, middlewares...) if pluginsHandler != nil { pluginsHandler.RegisterRoutes(s.Router, middlewares...) } if sessionHandler != nil { sessionHandler.RegisterRoutes(s.Router, middlewares...) } if promptsHandler != nil { promptsHandler.RegisterRoutes(s.Router, middlewares...) } skillsHandler := handlers.NewSkillsHandler(s.Config.ConfigStore, s.Config.ObjectStore) if skillsHandler != nil { skillsHandler.RegisterRoutes(s.Router, middlewares...) } webhookHandler := handlers.NewWebhookHandler(callbacks, s.Config, s.WebhookDispatcher) webhookHandler.RegisterRoutes(s.Router, middlewares...) skillsServingHandler := handlers.NewSkillsServingHandler(s.Config.ConfigStore, s.Config.ObjectStore, s.Config) if skillsServingHandler != nil { skillsServingHandler.RegisterRoutes(s.Router, middlewares...) } s.skillsServingHandler = skillsServingHandler cacheHandler.RegisterRoutes(s.Router, middlewares...) if featureFlagsHandler != nil { featureFlagsHandler.RegisterRoutes(s.Router, middlewares...) } if governanceHandler != nil { var overrides handlers.GovernanceRouteOverrides if provider, ok := callbacks.(GovernanceRouteOverridesProvider); ok { overrides = provider.GetGovernanceRouteOverrides(ctx) } governanceHandler.RegisterRoutesWithOverrides(s.Router, overrides, middlewares...) } if routingHandler != nil { routingHandler.RegisterRoutes(s.Router, middlewares...) } if loggingHandler != nil { loggingHandler.RegisterRoutes(s.Router, middlewares...) } if s.WebSocketHandler != nil { s.WebSocketHandler.RegisterRoutes(s.Router, middlewares...) } if s.NotificationService != nil { s.NotificationService.RegisterRoutes(s.Router, middlewares...) } if s.RSSFeedsHandler != nil { s.RSSFeedsHandler.RegisterRoutes(s.Router, middlewares...) } if s.WarpHandler != nil { s.WarpHandler.RegisterRoutes(s.Router, middlewares...) } if s.Config.ConfigStore != nil && s.SidekiqRunner != nil { handlers.NewSidekiqHandler(s.Config.ConfigStore, s.SidekiqRunner).RegisterRoutes(s.Router, middlewares...) } // Register dev pprof handler only in dev mode if handlers.IsDevMode() { logger.Info("dev mode enabled, registering pprof endpoints") s.devPprofHandler = handlers.NewDevPprofHandler() s.devPprofHandler.RegisterRoutes(s.Router, middlewares...) } metricsGatherer := prometheus.GathererFunc(func() ([]*dto.MetricFamily, error) { plugin, err := lib.FindPluginAs[*telemetry.PrometheusPlugin](s.Config, telemetry.PluginName) if err != nil || plugin == nil { return nil, nil } return plugin.GetMetricsGatherer().Gather() }) metricsAdapter := fasthttpadaptor.NewFastHTTPHandler(promhttp.HandlerFor(metricsGatherer, promhttp.HandlerOpts{})) metricsHandler := func(ctx *fasthttp.RequestCtx) { plugin, err := lib.FindPluginAs[*telemetry.PrometheusPlugin](s.Config, telemetry.PluginName) if err != nil || plugin == nil || !plugin.IsMetricsEnabled() { handlers.SendError(ctx, fasthttp.StatusNotFound, "Route not found: "+string(ctx.Path())) return } metricsAdapter(ctx) } s.Router.GET("/metrics", lib.ChainMiddlewares(metricsHandler, middlewares...)) // 404 handler s.Router.NotFound = func(ctx *fasthttp.RequestCtx) { handlers.SendError(ctx, fasthttp.StatusNotFound, "Route not found: "+string(ctx.Path())) } return nil } // ObserveWarmupRoutingEmbedding forwards semantic-routing warmup embedding // usage from the routing plugin to the telemetry plugin's routing overhead // counters. The telemetry plugin is resolved per call (warmup is rare — boot // and config changes only) so a reloaded telemetry instance, with its fresh // registry, is picked up without re-wiring routing. // // Exported because embedders that reimplement Bootstrap rather than calling it // must repeat this wiring themselves, and an unexported method is unreachable // from their package even through the embedded server. Warmup embeds carry no // request or response, so an embedder that misses this observer loses the usage // entirely — it cannot be recovered from the routing metadata path. func (s *BifrostHTTPServer) ObserveWarmupRoutingEmbedding(provider, model string, inputTokens int) { plugin, err := lib.FindPluginAs[*telemetry.PrometheusPlugin](s.Config, telemetry.PluginName) if err != nil || plugin == nil { return } plugin.ObserveWarmupRoutingEmbedding(provider, model, inputTokens) } // RegisterUIRoutes registers the UI handler with the specified router func (s *BifrostHTTPServer) RegisterUIRoutes(middlewares ...schemas.BifrostHTTPMiddleware) { // WARNING: This UI handler needs to be registered after all the other handlers handlers.NewUIHandler(s.UIContent, s.ShellRewriter).RegisterRoutes(s.Router, middlewares...) } // GetAllRedactedKeys gets all redacted keys from the config store func (s *BifrostHTTPServer) GetAllRedactedKeys(ctx context.Context, ids []string) []schemas.Key { if s.Config == nil || s.Config.ConfigStore == nil { return nil } redactedKeys, err := s.Config.ConfigStore.GetAllRedactedKeys(ctx, ids) if err != nil { logger.Error("failed to get all redacted keys: %v", err) return nil } return redactedKeys } // GetAllRedactedVirtualKeys gets all redacted virtual keys from the config store func (s *BifrostHTTPServer) GetAllRedactedVirtualKeys(ctx context.Context, ids []string) []tables.TableVirtualKey { if s.Config == nil || s.Config.ConfigStore == nil { return nil } virtualKeys, err := s.Config.ConfigStore.GetRedactedVirtualKeys(ctx, ids) if err != nil { logger.Error("failed to get all redacted virtual keys: %v", err) return nil } return virtualKeys } // GetAllRedactedRoutingRules gets all redacted routing rules from the config store func (s *BifrostHTTPServer) GetAllRedactedRoutingRules(ctx context.Context, ids []string) []tables.TableRoutingRule { if s.Config == nil || s.Config.ConfigStore == nil { return nil } routingRules, err := s.Config.ConfigStore.GetRedactedRoutingRules(ctx, ids) if err != nil { logger.Error("failed to get all redacted routing rules: %v", err) return nil } return routingRules } // prometheusHTTPMiddleware records Prometheus HTTP metrics. The plugin is resolved per // request rather than captured once, for the same reason as the otel middleware: a config // reload swaps in a freshly constructed instance, and a pointer captured at startup kept // recording into the orphaned one, so http_* series stopped appearing on /metrics after // any reload. Requests still pass through when no telemetry plugin is loaded. func (s *BifrostHTTPServer) prometheusHTTPMiddleware() schemas.BifrostHTTPMiddleware { return func(next fasthttp.RequestHandler) fasthttp.RequestHandler { return func(ctx *fasthttp.RequestCtx) { prometheusPlugin, err := lib.FindPluginAs[*telemetry.PrometheusPlugin](s.Config, telemetry.PluginName) if err != nil { next(ctx) return } prometheusPlugin.HTTPMiddleware(next)(ctx) } } } // PrepareCommonMiddlewares gets the common middlewares for the Bifrost HTTP server func (s *BifrostHTTPServer) PrepareCommonMiddlewares() []schemas.BifrostHTTPMiddleware { commonMiddlewares := []schemas.BifrostHTTPMiddleware{} // Copy the matched route template saved by the router (SaveMatchedRoutePath) into a // stable, router-agnostic user value so metrics middlewares below can label by route // template instead of the raw URL path. commonMiddlewares = append(commonMiddlewares, func(next fasthttp.RequestHandler) fasthttp.RequestHandler { return func(ctx *fasthttp.RequestCtx) { if route, ok := ctx.UserValue(router.MatchedRoutePathParam).(string); ok && route != "" { ctx.SetUserValue(string(schemas.BifrostContextKeyHTTPRoute), route) // Drop the router's randomized key so it doesn't leak into request PathParams. ctx.RemoveUserValue(router.MatchedRoutePathParam) } next(ctx) } }) // Preparing middlewares commonMiddlewares = append(commonMiddlewares, s.prometheusHTTPMiddleware()) // OTel HTTP metrics (http_requests_total etc., pushed via OTLP). The otel plugin is // resolved per request rather than captured here: a config reload swaps in a freshly // constructed plugin instance, and a pointer captured at startup would keep recording // against exporters whose meter provider has been shut down. commonMiddlewares = append(commonMiddlewares, func(next fasthttp.RequestHandler) fasthttp.RequestHandler { return func(ctx *fasthttp.RequestCtx) { start := time.Now() reqSize := float64(ctx.Request.Header.ContentLength()) next(ctx) otelPlugin, err := lib.FindPluginAs[*otel.OtelPlugin](s.Config, otel.PluginName) if err != nil { return } // Label by the matched route template when available (set by the middleware // above) so path params (model names, batch/file IDs) don't explode cardinality. path := string(ctx.Path()) if route, ok := ctx.UserValue(string(schemas.BifrostContextKeyHTTPRoute)).(string); ok && route != "" { path = route } otelPlugin.RecordHTTPMetrics(ctx, path, string(ctx.Method()), strconv.Itoa(ctx.Response.StatusCode()), time.Since(start).Seconds(), reqSize, float64(ctx.Response.Header.ContentLength()), ) } }) return commonMiddlewares } // shouldRun, when non-nil, is consulted once before the sweep starts; // returning false skips it entirely. Deployments running several instances // against one config store can use it to restrict the sweep to a single // instance. nil means always run. func startSkillsOrphanCleanupWorker(ctx context.Context, config *lib.Config, shouldRun func() bool) { if config == nil || config.ConfigStore == nil { return } if shouldRun != nil && !shouldRun() { return } // Run once on startup asynchronously with a 10-minute timeout so a stalled // DB or object-store call does not leave the goroutine hanging indefinitely. go func() { cleanupCtx, cancel := context.WithTimeout(ctx, 10*time.Minute) defer cancel() result, err := handlers.CleanupOrphanSkillFiles(cleanupCtx, config.ConfigStore, config.ObjectStore, false) if err != nil { logger.Warn("skills orphan cleanup failed during startup: %v", err) } else { logger.Info("skills orphan cleanup completed during startup: deleted_db_blobs=%d deleted_storage_objects=%d", result.DeletedDBBlobs, result.DeletedStorageObjects) } }() } // Bootstrap initializes the Bifrost HTTP server with all necessary components. // It: // 1. Initializes Prometheus collectors for monitoring // 2. Reads and parses configuration from the specified config file // 3. Initializes the Bifrost client with the configuration // 4. Sets up HTTP routes for text and chat completions // // The server exposes the following endpoints: // - POST /v1/text/completions: For text completion requests // - POST /v1/chat/completions: For chat completion requests // - GET /metrics: For Prometheus metrics func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { var err error ctx = context.WithValue(ctx, schemas.BifrostContextKeyRuntimeVersion, s.Version) s.Ctx, s.cancel = schemas.NewBifrostContextWithCancel(ctx) handlers.SetVersion(s.Version) configDir := GetDefaultConfigDir(s.AppDir) // Ensure app directory exists if err := os.MkdirAll(configDir, 0o755); err != nil { return fmt.Errorf("failed to create app directory %s: %v", configDir, err) } // Initialize high-performance configuration store with dedicated database s.Config, err = lib.LoadConfig(ctx, configDir) if err != nil { return fmt.Errorf("failed to load config %v", err) } if s.Config.KVStore != nil { integrations.RegisterKVDecoders(s.Config.KVStore) } // Outbound clients for non-inference traffic, built by LoadConfig on the global // proxy it loaded. s.HTTPClientFactory = s.Config.HTTPClientFactory handlers.SetSkillFetchHTTPClientFactory(s.HTTPClientFactory) // Initialize WebSocket handler early so plugins can wire event broadcasters during Init. // Log callbacks are registered later in RegisterAPIRoutes when logging plugin is available. s.WebSocketHandler = handlers.NewWebSocketHandler(s.Ctx, s.Config.ClientConfig.AllowedOrigins) s.Config.EventBroadcaster = s.WebSocketHandler.BroadcastEvent s.NotificationService = handlers.NewNotificationService(s.Config.ConfigStore, s.WebSocketHandler) s.Config.NotificationPublisher = s.NotificationService.Publish s.NotificationService.Start(s.Ctx) // Warp is built once, in RegisterAPIRoutes below (called from this same // Bootstrap, before it returns): that is the first point the logging // plugin - and so the log manager Warp's chat route needs - is known. // Building one here too used to seem necessary so the config routes would // exist early, but RegisterRoutes is only ever called once, later, on // whichever handler is current then - so a handler built here never // serves a request or gets its routes registered before being replaced. // Initializing plugin loader. Allowlist entries are validated now - a malformed entry // fails server startup rather than silently no-oping, since this is security-relaxing // config for SSRF protection on custom plugin downloads. var pluginDownloadAllowlist *network.Allowlist if s.Config.ServerConfig != nil && len(s.Config.ServerConfig.PluginDownloadPrivateAllowlist) > 0 { pluginDownloadAllowlist, err = network.NewAllowlist(s.Config.ServerConfig.PluginDownloadPrivateAllowlist) if err != nil { return fmt.Errorf("invalid server.plugin_download_private_allowlist: %w", err) } } s.Config.PluginLoader = dynamicPlugins.NewSharedObjectPluginLoader(pluginDownloadAllowlist, dynamicPlugins.WithHTTPClientFactory(s.HTTPClientFactory)) // Initialize log retention cleaner if log store is configured if s.Config.LogsStore != nil { // If log retention days remains 0, then we wont be initializing the log retention cleaner logRetentionDays := 0 if s.Config.ConfigStore != nil { // Get logs store config from config store clientConfig, err := s.Config.ConfigStore.GetClientConfig(ctx) if err != nil { logger.Warn("failed to get logs store config: %v", err) // So we wont be initializing the log retention cleaner } if clientConfig != nil { logRetentionDays = clientConfig.LogRetentionDays } } else { // We will check if the config file has the log retention days set logRetentionDays = s.Config.ClientConfig.LogRetentionDays } logger.Info("log retention days: %d", logRetentionDays) if logRetentionDays > 0 { // Type assert to get RDBLogStore (which implements LogRetentionManager) if rdbStore, ok := s.Config.LogsStore.(logstore.LogRetentionManager); ok { cleanerConfig := logstore.CleanerConfig{ RetentionDays: logRetentionDays, } s.LogsCleaner = logstore.NewLogsCleaner(rdbStore, cleanerConfig, logger) s.LogsCleaner.StartCleanupRoutine() logger.Info("log retention cleaner initialized with %d days retention", logRetentionDays) } } } // Initialize async job cleaner if log store is configured if s.Config.LogsStore != nil { s.AsyncJobCleaner = logstore.NewAsyncJobCleaner(s.Config.LogsStore, logger) s.AsyncJobCleaner.StartCleanupRoutine() } // Load all plugins if err := s.LoadPlugins(ctx); err != nil { return fmt.Errorf("failed to instantiate plugins: %v", err) } // Initialize the webhook delivery dispatcher (requires both stores; the // in-memory endpoint store on Config serves endpoint lookups). if s.Config.LogsStore != nil && s.Config.ConfigStore != nil { s.WebhookDispatcher = webhooks.NewDispatcher(ctx, "", s.Config.ClientConfig.WebhookConfig.DeliveryHistoryRetention(), s.Config.ConfigStore, s.Config.LogsStore, s.Config, logger, webhooks.WithHTTPClientFactory(s.HTTPClientFactory)) s.WebhookDispatcher.Start() logger.Info("webhook dispatcher initialized") } // Initialize async job executor (requires LogsStore + governance plugin) if s.Config.LogsStore != nil { governancePlugin, govErr := lib.FindPluginAs[governance.BaseGovernancePlugin](s.Config, s.getGovernancePluginName()) if govErr == nil { // The dispatcher interface value must stay nil when no dispatcher // exists — a typed-nil pointer would defeat the executor's nil check. var jobDispatcher logstore.WebhookDispatcher if s.WebhookDispatcher != nil { jobDispatcher = s.WebhookDispatcher } s.Config.AsyncJobExecutor = logstore.NewAsyncJobExecutor(s.Config.LogsStore, governancePlugin.GetGovernanceStore(), jobDispatcher, s.Config, logger) logger.Info("async job executor initialized") } } // openapi MCP clients synthesize their server from a spec; the framework // package does that, core only knows the hook. Registered on the config // passed to Init and again on the client below, so a manager created // lazily (no mcp block in config.json) gets it too. openAPIFactory := openapimcp.Factory(openapimcp.WithConfigDir(s.Config.ConfigDir()), openapimcp.WithMaxSpecBytes(s.Config.MCPOpenAPIMaxSpecBytes)) tableMCPConfig := s.Config.MCPConfig var mcpConfig *schemas.MCPConfig if tableMCPConfig != nil { mcpConfig = s.Config.MCPConfig if mcpConfig != nil { mcpConfig.FetchNewRequestIDFunc = func(ctx *schemas.BifrostContext) string { return uuid.New().String() } mcpConfig.InProcessServerFactory = openAPIFactory } } // Initialize bifrost client // Create account backed by the high-performance store (all processing is done in LoadFromDatabase) // The account interface now benefits from ultra-fast config access times via in-memory storage account := lib.NewBaseAccount(s.Config) s.Client, err = bifrost.Init(ctx, schemas.BifrostConfig{ Account: account, InitialPoolSize: s.Config.ClientConfig.InitialPoolSize, DropExcessRequests: s.Config.ClientConfig.DropExcessRequests, LLMPlugins: s.Config.GetLoadedLLMPlugins(), MCPPlugins: s.Config.GetLoadedMCPPlugins(), A2APlugins: s.Config.GetLoadedA2APlugins(), RelayPlugins: s.Config.GetLoadedRelayPlugins(), MCPConfig: mcpConfig, OAuth2Provider: s.Config.OAuthProvider, MCPHeadersProvider: s.Config.MCPHeadersProvider, Logger: logger, KVStore: s.Config.KVStore, ModelCatalog: s.Config.ModelCatalog, }) if err != nil { return fmt.Errorf("failed to initialize bifrost: %v", err) } s.Client.SetMCPInProcessServerFactory(openAPIFactory) logger.Info("bifrost client initialized") // Sync plugin execution order from config to core (defensive — Init receives sorted list, // but this ensures order consistency if the loading path changes in the future) s.Client.ReorderPlugins(s.Config.GetPluginOrder()) // Seed the catalog: push the initial keyconfig snapshot and fetch per-key // live models for every provider concurrently. logger.Info("listing all models and adding to model catalog") if s.Config.ModelCatalog != nil { snapshot := make(map[schemas.ModelProvider][]schemas.Key, len(s.Config.Providers)) var ignoreProviderCost, keyless []schemas.ModelProvider for provider, providerConfig := range s.Config.Providers { snapshot[provider] = providerConfig.Keys if providerConfig.IgnoreProviderCost { ignoreProviderCost = append(ignoreProviderCost, provider) } if providerConfig.CustomProviderConfig != nil && providerConfig.CustomProviderConfig.IsKeyLess { keyless = append(keyless, provider) } } s.Config.ModelCatalog.ReplaceKeyConfig(snapshot, keyless...) s.Config.ModelCatalog.ReplaceIgnoreProviderCost(ignoreProviderCost) s.RefreshAllLiveModels(ctx) } logger.Info("models added to catalog") s.Config.SetBifrostClient(s.Client) s.WireBatchAccountingSweeper() // Initialize routes s.Router = router.New() // Save the matched route template on each request // so metrics can use it as the `path` label instead of the raw URL path. s.Router.SaveMatchedRoutePath = true // Initialize CORS middleware s.CORSMiddleware = handlers.NewCorsMiddleware(s.Config) commonMiddlewares := s.PrepareCommonMiddlewares() apiMiddlewares := commonMiddlewares inferenceMiddlewares := commonMiddlewares if s.Config.ConfigStore == nil { logger.Error("auth middleware requires config store, skipping auth middleware initialization") // No auth runs in this mode, so mark every API request as bypassed; otherwise the // handlers that require genuine auth for dangerous changes see an unmarked request and // let it through. apiMiddlewares = append(apiMiddlewares, handlers.AuthBypassedMiddleware(s.Config.ClientConfig.AllowedOrigins, s.Config.ClientConfig.WhitelistedRoutes)) } else { // Use a signed (stateless) ticket store when an encryption key is configured // so tickets are verifiable across nodes; otherwise fall back to in-memory. // NewSignedWSTicketStore handles empty key by degrading to in-memory mode. s.WSTicketStore = handlers.NewSignedWSTicketStore(encrypt.Key()) // Initialize the temp-token service and register all scopes owned by the // handlers package. The service is the seam every "scoped, anonymous, // browser-only" workflow plugs into (currently just the MCP per-user OAuth // auth page s.TempTokens = temptoken.NewService(s.Config.ConfigStore, temptoken.NewRegistry()) if regErr := handlers.RegisterTempTokenScopes(s.TempTokens); regErr != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil return fmt.Errorf("failed to register temp token scopes: %v", regErr) } // Centralized janitor that reaps expired temp_tokens rows. Independent // of the per-user OAuth sweep so any future scope (not just mcp_auth) // benefits from the same cleanup loop without piggybacking on OAuth. s.TempTokenSweepWorker = temptoken.NewSweepWorker(s.TempTokens, logger) if s.TempTokenSweepWorker != nil { s.TempTokenSweepWorker.Start(s.Ctx) } s.StartOAuth2SweepWorker(s.Ctx, nil) // Hand the service to the OAuth provider so InitiateUserOAuthFlow mints // a mcp_auth token and embeds it as a URL fragment on the auth-page link. if s.Config.OAuthProvider != nil { s.Config.OAuthProvider.SetTempTokenService(s.TempTokens) } // Same wiring for the per-user-headers provider — mints // mcp_headers_auth tokens on the headers submission URL. if s.Config.MCPHeadersProvider != nil { s.Config.MCPHeadersProvider.SetTempTokenService(s.TempTokens) } s.AuthMiddleware, err = handlers.InitAuthMiddleware(s.Config.ConfigStore, s.WSTicketStore, s.TempTokens, s.Config.SetupToken) if err != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil if s.TempTokenSweepWorker != nil { s.TempTokenSweepWorker.Stop() s.TempTokenSweepWorker = nil } if s.OAuth2SweepWorker != nil { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } return fmt.Errorf("failed to initialize auth middleware: %v", err) } if ctx.Value(schemas.BifrostContextKeyIsEnterprise) == nil { // OSS only: lock /api behind the setup token while dashboard auth is not // active. It must run before APIMiddleware, whose auth-off branch lets // every request through. apiMiddlewares = append(apiMiddlewares, s.AuthMiddleware.SetupLockMiddleware(), s.AuthMiddleware.APIMiddleware()) s.setupLockInstalled = true } } // Add semantic cache plugin embedding request executor if it exists semanticCachePlugin, err := lib.FindPluginAs[*semanticcache.Plugin](s.Config, semanticcache.PluginName) if err == nil && semanticCachePlugin != nil { semanticCachePlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } // Wire the routing plugin's semantic-classification embedding path. The // executor cannot be passed at Init: the plugin is built while the bifrost // client is still being assembled. if routingPlugin, err := s.getRoutingPlugin(); err == nil { // Store and executor together: warmup depends on both, and applying them // one at a time warms a throwaway generation into the embedded store // before the second call redirects it. routingPlugin.SetComplexityWarmupDependencies(s.Config.VectorStore, s.Client.EmbeddingRequest) routingPlugin.SetWarmupEmbedUsageObserver(s.ObserveWarmupRoutingEmbedding) routingPlugin.SetChatRequestExecutor(s.Client.ChatCompletionRequest) routingPlugin.SetResponsesRequestExecutor(s.Client.ResponsesRequest) routingPlugin.SetDecisionRequestExecutor(s.Client.DecisionRequest) } // Initialize Sidekiq runner for background jobs if s.Config != nil && s.Config.ConfigStore != nil { s.SidekiqRunner = sidekiq.New(s.Config.ConfigStore, logger, 4, "") } // Initialize the Agent Gateway once before either management or protocol // routes are registered. s.InitializeRelays(s.Ctx, nil, nil) if err = s.InitializeAgentGateway(s.Ctx, nil, nil); err != nil { if s.WSTicketStore != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil } if s.TempTokenSweepWorker != nil { s.TempTokenSweepWorker.Stop() s.TempTokenSweepWorker = nil } if s.OAuth2SweepWorker != nil { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } return err } // Register routes err = s.RegisterAPIRoutes(s.Ctx, s, apiMiddlewares...) if err != nil { s.CloseAgentGateway() if s.WSTicketStore != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil } if s.TempTokenSweepWorker != nil { s.TempTokenSweepWorker.Stop() s.TempTokenSweepWorker = nil } if s.OAuth2SweepWorker != nil { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } return fmt.Errorf("failed to initialize routes: %v", err) } // Registering inference routes // The pre-auth plugin phase runs immediately before the auth middlewares so a plugin can // supply or translate the credentials they read (see schemas.HTTPTransportPreAuthHook). // Skipped entirely when no transport plugin is loaded. inferenceMiddlewares = append(inferenceMiddlewares, handlers.TransportPreAuthInterceptorMiddleware(s.Config)) if ctx.Value(schemas.BifrostContextKeyIsEnterprise) == nil && s.AuthMiddleware != nil { inferenceMiddlewares = append(inferenceMiddlewares, s.AuthMiddleware.InferenceMiddleware()) } // Once auth is done we will first add the Tracing middleware // Always add tracing middleware when tracer is enabled - it creates traces and sets traceID in context // The observability plugins are optional (can be empty if only logging is enabled) // Curating observability plugins observabilityPlugins := s.CollectObservabilityPlugins() // This enables the central streaming accumulator for both use cases // Initializing tracer with embedded streaming accumulator traceStore := tracing.NewTraceStore(60*time.Minute, logger) tracer := tracing.NewTracer(traceStore, s.Config.ModelCatalog, logger) tracer.SetObservabilityPlugins(observabilityPlugins, s.CollectObservabilityLimits()) s.Client.SetTracer(tracer) s.TracingMiddleware = handlers.NewTracingMiddleware(tracer) // Wire the settlement tracer now that it exists (WireBatchAccountingSweeper ran earlier, before it did). s.setSettlementTracer() // TransportInterceptor runs AFTER the auth middlewares so HTTPTransportPreHook observes an // authenticated request, and inside TracingMiddleware so the tracing defer runs AFTER // transport post-hooks (capturing HTTPTransportPostHook plugin logs). // Recovery sits directly inside Tracing (see handlers.InferenceOuterMiddlewares). // Order: Tracing.pre → Recovery → PreAuthInterceptor → auth → TransportInterceptor.pre → handler → // TransportInterceptor.post → Recovery.defer → Tracing.defer inferenceMiddlewares = append(inferenceMiddlewares, handlers.TransportInterceptorMiddleware(s.Config)) inferenceMiddlewares = append(handlers.InferenceOuterMiddlewares(s.TracingMiddleware, s.CORSMiddleware), inferenceMiddlewares...) err = s.RegisterInferenceRoutes(s.Ctx, inferenceMiddlewares...) if err != nil { s.CloseAgentGateway() if s.WSTicketStore != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil } if s.TempTokenSweepWorker != nil { s.TempTokenSweepWorker.Stop() s.TempTokenSweepWorker = nil } if s.OAuth2SweepWorker != nil { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } return fmt.Errorf("failed to initialize inference routes: %v", err) } // Registered before ConnectConfiguredMCPClients so even the very first // boot dial's discovered tools get persisted, not just later reconnects. s.Client.SetMCPDiscoveryChangeCallback(func(clientID, name string, discovery *schemas.MCPDiscovery) { s.PersistMCPClientDiscovery(s.Ctx, clientID, discovery) go s.SyncMCPServersAfterToolsChange(s.Ctx, clientID) }) // Dial configured MCP clients now that every plugin is registered in the core. // Construction (bifrost.Init) no longer connects MCP, so connecting here ensures // each client's PreMCPConnectionHook runs against the full plugin set rather than // the point-in-time snapshot captured at Init (which would skip plugins — e.g. // enterprise ones — registered after that snapshot, causing the client to fail // and only recover on a later health-monitor reconnect). s.Client.ConnectConfiguredMCPClients(s.Ctx) // NewMCPServerHandler (inside RegisterInferenceRoutes, above) already synced the served // server once, before any client had dialed - it necessarily served zero tools. The // callback registered above would ordinarily pick up what the boot dial just discovered, // but it is gated on a genuine change in a client's tool hash, and a client whose tools // were persisted to the config store on a previous run boots with that hash already // recorded: rediscovering the same tools now looks like no change at all, so the callback // never fires and the served server would stay at the zero tools it started with. This is // the one place that resync is not conditional on anything changing: the boot dial is // synchronous (ConnectConfiguredMCPClients blocks until every client has been attempted), // so by the time it returns every client's tools are exactly what this request should see. if err := s.MCPServerHandler.SyncMCPServer(s.Ctx); err != nil { logger.Warn("failed to sync MCP server after the boot dial: %v", err) } // A proactively refreshed shared OAuth token leaves the live MCP connection // still holding the old bearer (the credential is baked into the transport // at connect time), so recycle the connection as soon as the refresh worker // lands a fresh token instead of waiting for a call to fail. Installed only // after the boot dial above so the worker's first sweep cannot race client // construction. In-flight calls on the old connection may fail once during // the recycle; the reactive auth-failure retry covers them. Per-user // (non-shared) tokens have no persistent connection, so they are skipped. if s.Config.OAuthTokenRefreshWorker != nil { s.Config.OAuthTokenRefreshWorker.SetOnTokenRefreshed(func(mcpClientID, authMode string) { if authMode != "shared" { return } // Run in a goroutine: a reconnect can take minutes worst-case and // must not block the worker's sequential refresh loop. Expected // no-op outcomes (a reconnect already in progress, or a client // type reconnect does not apply to) surface here as errors and // are deliberately logged at Debug only. go func() { if err := s.ReconnectMCPClient(s.Ctx, mcpClientID); err != nil { logger.Debug("MCP client %s was not reconnected after a proactive token refresh: %v", mcpClientID, err) return } logger.Debug("recycled shared MCP connection for client %s after a proactive token refresh", mcpClientID) }() }) // Started here, not at construction (lib/config.go): Start's first // refresh sweep runs immediately, synchronously with launching its // goroutine. Starting any earlier — before the callback above is // wired up and before the boot dial above has run — would let a // refresh land with nothing to trigger the affected client's // reconnect, leaving a live connection on the old bearer until the // next scheduled sweep or a call happens to fail. s.Config.OAuthTokenRefreshWorker.Start(s.Ctx) } // Serve a minimal robots.txt so crawlers/CLI tools (e.g. Claude Code) don't // trigger 404 warnings when probing the host before marketplace fetches. s.Router.GET("/robots.txt", func(ctx *fasthttp.RequestCtx) { ctx.SetContentType("text/plain; charset=utf-8") ctx.SetBodyString("User-agent: *\nAllow: /\n") }) // Register UI handler s.RegisterUIRoutes() // Start the Sidekiq dispatcher: on every node it periodically claims pending and // stale (orphaned) jobs, with an atomic claim guaranteeing exactly one node runs // each job. Subsumes startup recovery of jobs left behind by a crash or restart. if s.SidekiqRunner != nil { s.SidekiqDispatcherStop = s.SidekiqRunner.StartDispatcher(sidekiq.DispatchInterval, sidekiq.StaleAfter) } if s.GovernanceHandler != nil { s.GovernanceHandler.StartExpiryCleanupScheduler(s.Ctx) } // Checking if config has server config and use it to set read buffer size logger.Debug("server read buffer size: %d", s.Config.ServerConfig.ReadBufferSize) maxRequestBodySizeMB := s.Config.ClientConfig.MaxRequestBodySizeMB lib.ApplyBodyBufferPoolLimit(s.Config.ServerConfig, maxRequestBodySizeMB*1024*1024) // Create fasthttp server instance s.Server = &fasthttp.Server{ Handler: handlers.ServerRootHandler(s.CORSMiddleware, s.Config, s.Router.Handler), MaxRequestBodySize: maxRequestBodySizeMB * 1024 * 1024, ReadBufferSize: s.Config.ServerConfig.ReadBufferSize, // Explain an oversized body (413 naming the configured limit) instead of fasthttp's generic 400. ErrorHandler: lib.RequestErrorHandler(maxRequestBodySizeMB, logger.Warn), } startSkillsOrphanCleanupWorker(s.Ctx, s.Config, nil) // Keep the live model catalog current after boot. Without this, a model an // upstream starts serving mid-uptime stays invisible until a restart or a // key edit, since nothing else re-fetches list-models. s.RestartLiveModelRefresher(s.Ctx) return nil } func (s *BifrostHTTPServer) CloseAgentGateway() { s.agentGatewayCloseOnce.Do(func() { // The gRPC listener drains first so no new operations reach the manager // while it is closing its runtimes. s.agentGatewayGRPCServer.Stop() if s.AgentGatewayHandler != nil { s.AgentGatewayHandler.Close() } }) } // Start starts the HTTP server at the specified host and port // Also watches signals and errors func (s *BifrostHTTPServer) Start() error { // Printing plugin status in a table for _, pluginStatus := range s.Config.GetPluginStatus() { logger.Info("plugin status: %s - %s", pluginStatus.Name, pluginStatus.Status) } // Create channels for signal and error handling sigChan := make(chan os.Signal, 1) errChan := make(chan error, 1) // Watching for signals signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) // Start server in a goroutine serverAddr := net.JoinHostPort(s.Host, s.Port) ln, err := net.Listen("tcp", serverAddr) if err != nil { return fmt.Errorf("failed to create listener on %s: %v", serverAddr, err) } go func() { logger.Info("successfully started bifrost, serving UI on http://%s", serverAddr) if err := s.Server.Serve(ln); err != nil { errChan <- err } }() // Wait for either termination signal or server error select { case sig := <-sigChan: logger.Info("received signal %v, initiating graceful shutdown...", sig) if s.IntegrationHandler != nil { logger.Info("closing realtime transport sessions...") s.IntegrationHandler.Close() } if s.wsLiveHandler != nil { s.wsLiveHandler.Close() } if s.webrtcLiveHandler != nil { s.webrtcLiveHandler.Close() } // Create shutdown context with timeout shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() // Perform graceful shutdown if err := s.Server.Shutdown(); err != nil { logger.Error("error during graceful shutdown: %v", err) } else { logger.Info("server gracefully shutdown") } // Cancelling main context if s.cancel != nil { s.cancel() } // Wait for shutdown to complete or timeout done := make(chan struct{}) go func() { defer close(done) // Warp first. Its indexer workers call s.Client.EmbeddingRequest (and its // chat turns s.Client.ResponsesRequest), and // LogIndexer.Close waits for them - so shutting the client down first // cancelled its context underneath work that was still being waited on, // and pending indexing failed during an orderly shutdown. if s.WarpHandler != nil { logger.Info("shutting down warp...") s.WarpHandler.Shutdown() logger.Info("warp shutdown completed") } logger.Info("closing agent gateway...") s.CloseAgentGateway() logger.Info("shutting down bifrost client...") s.Client.Shutdown() logger.Info("bifrost client shutdown completed") logger.Info("cleaning up storage engines...") // Cleanup server-specific components if s.LogsCleaner != nil { logger.Info("stopping log retention cleaner...") s.LogsCleaner.StopCleanupRoutine() } if s.AsyncJobCleaner != nil { logger.Info("stopping async job cleaner...") s.AsyncJobCleaner.StopCleanupRoutine() } if s.WebhookDispatcher != nil { logger.Info("stopping webhook dispatcher...") s.WebhookDispatcher.Stop() } if s.WSTicketStore != nil { logger.Info("stopping ws ticket store...") s.WSTicketStore.Stop() } if s.TempTokenSweepWorker != nil { logger.Info("stopping temp-token sweep worker...") s.TempTokenSweepWorker.Stop() } if s.OAuth2SweepWorker != nil { logger.Info("stopping oauth2 sweep worker...") s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } if s.GovernanceHandler != nil { s.GovernanceHandler.StopExpiryCleanupScheduler() } if s.SidekiqDispatcherStop != nil { logger.Info("stopping sidekiq dispatcher...") s.SidekiqDispatcherStop() } logger.Info("stopping live model refresher...") s.stopLiveModelRefresher() if s.SidekiqRunner != nil { logger.Info("stopping sidekiq runner...") s.SidekiqRunner.Shutdown() } if s.devPprofHandler != nil { logger.Info("stopping dev pprof handler...") s.devPprofHandler.Cleanup() } if s.skillsServingHandler != nil { s.skillsServingHandler.Close() } if s.wsPool != nil { logger.Info("closing websocket connection pool...") s.wsPool.Close() } // Cleanup Config and all its background components if s.Config != nil { s.Config.Close(shutdownCtx) } logger.Info("storage engines cleanup completed") }() select { case <-done: logger.Info("cleanup completed") case <-shutdownCtx.Done(): logger.Warn("cleanup timed out after 30 seconds") } case err := <-errChan: s.cleanupAfterServeError() return err } return nil } // cleanupAfterServeError releases what Start had set up when Serve itself // fails: open realtime sessions, the websocket pool, and the skills serving // cache, whose exported bare repositories live in temp directories that only // Close removes. The graceful-shutdown path above does the same as part of // its full teardown. func (s *BifrostHTTPServer) cleanupAfterServeError() { if s.IntegrationHandler != nil { s.IntegrationHandler.Close() } if s.wsLiveHandler != nil { s.wsLiveHandler.Close() } if s.webrtcLiveHandler != nil { s.webrtcLiveHandler.Close() } if s.wsPool != nil { s.wsPool.Close() } if s.skillsServingHandler != nil { s.skillsServingHandler.Close() } }