Update middlewares.go

Author: github-actions[bot]

Makefile

@@ -37,7 +37,6 @@ help: ## Show this help message
 cleanup-enterprise: ## Clean up enterprise directories if present
 	@echo "$(GREEN)Cleaning up enterprise...$(NC)"
 	@if [ -d "ui/app/enterprise" ]; then rm -rf ui/app/enterprise; fi
-	@if [ -d "ui/app/enterprise" ]; then rm -rf ui/app/enterprise; fi
 	@echo "$(GREEN)Enterprise cleaned up$(NC)"
 
 install-ui: cleanup-enterprise

core/bifrost.go

@@ -294,9 +294,44 @@ func (bifrost *Bifrost) TranscriptionStreamRequest(ctx context.Context, req *sch
 	return bifrost.handleStreamRequest(ctx, req, schemas.TranscriptionStreamRequest)
 }
 
+// RemovePlugin removes a plugin from the server.
+func (bifrost *Bifrost) RemovePlugin(name string) error {
+	var pluginToCleanup schemas.Plugin
+	for {
+		oldPlugins := bifrost.plugins.Load()
+		if oldPlugins == nil {
+			return nil
+		}
+		// Create new slice with replaced plugin
+		newPlugins := make([]schemas.Plugin, len(*oldPlugins))
+		copy(newPlugins, *oldPlugins)
+		for i, p := range newPlugins {
+			if p.GetName() == name {
+				pluginToCleanup = p
+				bifrost.logger.Debug("removing plugin %s", name)
+				newPlugins = append(newPlugins[:i], newPlugins[i+1:]...)
+				break
+			}
+		}
+		// Atomic compare-and-swap
+		if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
+			// Cleanup the old plugin
+			if pluginToCleanup != nil {
+				err := pluginToCleanup.Cleanup()
+				if err != nil {
+					bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
+				}
+			}
+			return nil
+		}
+		// Retrying as swapping did not work
+	}
+}
+
 // ReloadPlugin reloads a plugin with new instance
 // During the reload - it's stop the world phase where we take a global lock on the plugin mutex
-func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {	
+func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {
+	var pluginToCleanup schemas.Plugin
 	for {
 		oldPlugins := bifrost.plugins.Load()
 		if oldPlugins == nil {
@@ -308,22 +343,32 @@ func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {
 		found := false
 		for i, p := range newPlugins {
 			if p.GetName() == plugin.GetName() {
+				// Cleaning up old plugin before replacing it
+				pluginToCleanup = p
+				bifrost.logger.Debug("replacing plugin %s with new instance", plugin.GetName())
 				newPlugins[i] = plugin
 				found = true
 				break
 			}
 		}
-		if !found{
+		if !found {
 			// This means that user is adding a new plugin
+			bifrost.logger.Debug("adding new plugin %s", plugin.GetName())
 			newPlugins = append(newPlugins, plugin)
 		}
 		// Atomic compare-and-swap
 		if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
+			// Cleanup the old plugin
+			if pluginToCleanup != nil {
+				err := pluginToCleanup.Cleanup()
+				if err != nil {
+					bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
+				}
+			}
 			return nil
 		}
 		// Retrying as swapping did not work
 	}
-
 }
 
 // UpdateProviderConcurrency dynamically updates the queue size and concurrency for an existing provider.
@@ -1023,7 +1068,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 
 	msg := bifrost.getChannelMessage(*preReq, requestType)
 	msg.Context = ctx
-
+	startTime := time.Now()
 	select {
 	case queue <- *msg:
 		// Message was sent successfully
@@ -1049,6 +1094,10 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 	var resp *schemas.BifrostResponse
 	select {
 	case result = <-msg.Response:
+		latency := time.Since(startTime).Milliseconds()
+		if result.ExtraFields.Latency == nil {
+			result.ExtraFields.Latency = &latency
+		}
 		resp, bifrostErr := pipeline.RunPostHooks(&ctx, result, nil, len(*bifrost.plugins.Load()))
 		if bifrostErr != nil {
 			bifrost.releaseChannelMessage(msg)
@@ -1172,7 +1221,7 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
 		// Marking final chunk
 		ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
 		// On error we will complete post-hooks
-		recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(bifrost.plugins))
+		recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(*bifrost.plugins.Load()))
 		bifrost.releaseChannelMessage(msg)
 		if recoveredErr != nil {
 			return nil, recoveredErr
@@ -1375,6 +1424,7 @@ func (p *PluginPipeline) RunPreHooks(ctx *context.Context, req *schemas.BifrostR
 	var shortCircuit *schemas.PluginShortCircuit
 	var err error
 	for i, plugin := range p.plugins {
+		p.logger.Debug("running pre-hook for plugin %s", plugin.GetName())
 		req, shortCircuit, err = plugin.PreHook(ctx, req)
 		if err != nil {
 			p.preHookErrors = append(p.preHookErrors, err)
@@ -1391,17 +1441,19 @@ func (p *PluginPipeline) RunPreHooks(ctx *context.Context, req *schemas.BifrostR
 // RunPostHooks executes PostHooks in reverse order for the plugins whose PreHook ran.
 // Accepts the response and error, and allows plugins to transform either (e.g., recover from error, or invalidate a response).
 // Returns the final response and error after all hooks. If both are set, error takes precedence unless error is nil.
-func (p *PluginPipeline) RunPostHooks(ctx *context.Context, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError, count int) (*schemas.BifrostResponse, *schemas.BifrostError) {
+// runFrom is the count of plugins whose PreHooks ran; PostHooks will run in reverse from index (runFrom - 1) down to 0
+func (p *PluginPipeline) RunPostHooks(ctx *context.Context, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostResponse, *schemas.BifrostError) {
 	// Defensive: ensure count is within valid bounds
-	if count < 0 {
-		count = 0
+	if runFrom < 0 {
+		runFrom = 0
 	}
-	if count > len(p.plugins) {
-		count = len(p.plugins)
+	if runFrom > len(p.plugins) {
+		runFrom = len(p.plugins)
 	}
 	var err error
-	for i := count - 1; i >= 0; i-- {
+	for i := runFrom - 1; i >= 0; i-- {
 		plugin := p.plugins[i]
+		p.logger.Debug("running post-hook for plugin %s", plugin.GetName())
 		resp, bifrostErr, err = plugin.PostHook(ctx, resp, bifrostErr)
 		if err != nil {
 			p.postHookErrors = append(p.postHookErrors, err)
@@ -1618,4 +1670,5 @@ func (bifrost *Bifrost) Shutdown() {
 			bifrost.logger.Warn(fmt.Sprintf("Error cleaning up plugin: %s", err.Error()))
 		}
 	}
+	bifrost.logger.Info("all request channels closed")
 }

core/changelog.md

@@ -1,4 +1,4 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- Fix: Updates token calculation for streaming responses. #520
+- Feature: Adds dynamic reloads for plugins. This removes requirement of restarts on updating plugins.

core/schemas/bifrost.go

@@ -100,6 +100,8 @@ type BifrostContextKey string
 
 // BifrostContextKeyRequestType is a context key for the request type.
 const (
+	BifrostContextKeyRequestID          BifrostContextKey = "request-id"
+	BifrostContextKeyVirtualKeyHeader   BifrostContextKey = "x-bf-vk"
 	BifrostContextKeyDirectKey          BifrostContextKey = "bifrost-direct-key"
 	BifrostContextKeyStreamEndIndicator BifrostContextKey = "bifrost-stream-end-indicator"
 	BifrostContextKeyRequestType        BifrostContextKey = "bifrost-request-type"
@@ -792,7 +794,7 @@ type TranscriptionUsage struct {
 type BifrostResponseExtraFields struct {
 	Provider    ModelProvider      `json:"provider"`
 	Params      ModelParameters    `json:"model_params"`
-	Latency     *int64              `json:"latency,omitempty"`
+	Latency     *int64             `json:"latency,omitempty"`
 	ChatHistory *[]BifrostMessage  `json:"chat_history,omitempty"`
 	BilledUsage *BilledLLMUsage    `json:"billed_usage,omitempty"`
 	ChunkIndex  int                `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses

core/version

@@ -1 +1 @@
-1.1.38
+1.1.39

framework/changelog.md

@@ -1,4 +1,5 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- upgrade: core upgrades to 1.1.38
+- Chore: Adds ctx to each function to gracefully shutdown ongoing tasks and bring better concurrency management
+- Fix: Fixes pricing sync to make sure latest updates are synced at every restart.

framework/configstore/migrations.go

@@ -1,31 +1,33 @@
 package configstore
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/maximhq/bifrost/framework/configstore/internal/migration"
 	"gorm.io/gorm"
 )
 
 // Migrate performs the necessary database migrations.
-func triggerMigrations(db *gorm.DB) error {
-	if err := migrationInit(db); err != nil {
+func triggerMigrations(ctx context.Context, db *gorm.DB) error {
+	if err := migrationInit(ctx, db); err != nil {
 		return err
 	}
-	if err := migrationMany2ManyJoinTable(db); err != nil {
+	if err := migrationMany2ManyJoinTable(ctx, db); err != nil {
 		return err
 	}
-	if err := migrationAddCustomProviderConfigJSONColumn(db); err != nil {
+	if err := migrationAddCustomProviderConfigJSONColumn(ctx, db); err != nil {
 		return err
 	}
 	return nil
 }
 
 // migrationInit is the first migration
-func migrationInit(db *gorm.DB) error {
+func migrationInit(ctx context.Context, db *gorm.DB) error {
 	m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
 		ID: "init",
 		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
 			migrator := tx.Migrator()
 			if !migrator.HasTable(&TableConfigHash{}) {
 				if err := migrator.CreateTable(&TableConfigHash{}); err != nil {
@@ -120,6 +122,7 @@ func migrationInit(db *gorm.DB) error {
 			return nil
 		},
 		Rollback: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
 			migrator := tx.Migrator()
 			// Drop children first, then parents (adjust if your actual FKs differ)
 			if err := migrator.DropTable(&TableVirtualKey{}); err != nil {
@@ -184,10 +187,11 @@ func migrationInit(db *gorm.DB) error {
 }
 
 // createMany2ManyJoinTable creates a many-to-many join table for the given tables.
-func migrationMany2ManyJoinTable(db *gorm.DB) error {
+func migrationMany2ManyJoinTable(ctx context.Context, db *gorm.DB) error {
 	m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
 		ID: "many2manyjoin",
 		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
 			migrator := tx.Migrator()
 
 			// create the many-to-many join table for virtual keys and keys
@@ -222,10 +226,11 @@ func migrationMany2ManyJoinTable(db *gorm.DB) error {
 	return nil
 }
 
-func migrationAddCustomProviderConfigJSONColumn(db *gorm.DB) error {
+func migrationAddCustomProviderConfigJSONColumn(ctx context.Context, db *gorm.DB) error {
 	m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
 		ID: "addcustomproviderconfigjsoncolumn",
 		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
 			migrator := tx.Migrator()
 
 			if !migrator.HasColumn(&TableProvider{}, "custom_provider_config_json") {