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
@@ -1332,7 +1381,7 @@ func handleProviderRequest(provider schemas.Provider, req *ChannelMessage, key s
 	case schemas.TextCompletionRequest:
 		return provider.TextCompletion(req.Context, req.Model, key, *req.Input.TextCompletionInput, req.Params)
 	case schemas.ChatCompletionRequest:
-		return provider.ChatCompletion(req.Context, req.Model, key, *req.Input.ChatCompletionInput, req.Params)
+		return provider.ChatCompletion(req.Context, req.Model, key, req.Input.ChatCompletionInput, req.Params)
 	case schemas.EmbeddingRequest:
 		return provider.Embedding(req.Context, req.Model, key, req.Input.EmbeddingInput, req.Params)
 	case schemas.SpeechRequest:
@@ -1353,7 +1402,7 @@ func handleProviderRequest(provider schemas.Provider, req *ChannelMessage, key s
 func handleProviderStreamRequest(provider schemas.Provider, req *ChannelMessage, key schemas.Key, postHookRunner schemas.PostHookRunner, reqType schemas.RequestType) (chan *schemas.BifrostStream, *schemas.BifrostError) {
 	switch reqType {
 	case schemas.ChatCompletionStreamRequest:
-		return provider.ChatCompletionStream(req.Context, postHookRunner, req.Model, key, *req.Input.ChatCompletionInput, req.Params)
+		return provider.ChatCompletionStream(req.Context, postHookRunner, req.Model, key, req.Input.ChatCompletionInput, req.Params)
 	case schemas.SpeechStreamRequest:
 		return provider.SpeechStream(req.Context, postHookRunner, req.Model, key, req.Input.SpeechInput, req.Params)
 	case schemas.TranscriptionStreamRequest:
@@ -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/mcp.go

@@ -872,9 +872,9 @@ func (m *MCPManager) addMCPToolsToBifrostRequest(ctx context.Context, req *schem
 			req.Params = &schemas.ModelParameters{}
 		}
 		if req.Params.Tools == nil {
-			req.Params.Tools = &[]schemas.Tool{}
+			req.Params.Tools = []schemas.Tool{}
 		}
-		tools := *req.Params.Tools
+		tools := req.Params.Tools
 
 		// Create a map of existing tool names for O(1) lookup
 		existingToolsMap := make(map[string]bool)
@@ -890,7 +890,7 @@ func (m *MCPManager) addMCPToolsToBifrostRequest(ctx context.Context, req *schem
 				existingToolsMap[mcpTool.Function.Name] = true
 			}
 		}
-		req.Params.Tools = &tools
+		req.Params.Tools = tools
 
 	}
 	return req

core/providers/anthropic.go

@@ -210,7 +210,7 @@ const (
 	AnthropicDefaultMaxTokens = 4096
 )
 
-// mapAnthropicFinishReasonToOpenAI maps Anthropic finish reasons to OpenAI-compatible ones
+// MapAnthropicFinishReason maps Anthropic finish reasons to OpenAI-compatible ones
 func MapAnthropicFinishReason(anthropicReason string) string {
 	switch anthropicReason {
 	case "end_turn":
@@ -516,7 +516,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 					Text: *msg.Content.ContentStr,
 				})
 			} else if msg.Content.ContentBlocks != nil {
-				for _, block := range *msg.Content.ContentBlocks {
+				for _, block := range msg.Content.ContentBlocks {
 					if block.Text != nil {
 						systemMessages = append(systemMessages, BedrockAnthropicSystemMessage{
 							Text: *block.Text,
@@ -547,7 +547,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 						"text": *msg.Content.ContentStr,
 					})
 				} else if msg.Content.ContentBlocks != nil {
-					for _, block := range *msg.Content.ContentBlocks {
+					for _, block := range msg.Content.ContentBlocks {
 						if block.Text != nil {
 							toolCallResultContent = append(toolCallResultContent, map[string]interface{}{
 								"type": "text",
@@ -567,7 +567,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 						"text": *msg.Content.ContentStr,
 					})
 				} else if msg.Content.ContentBlocks != nil {
-					for _, block := range *msg.Content.ContentBlocks {
+					for _, block := range msg.Content.ContentBlocks {
 						if block.Text != nil && *block.Text != "" {
 							content = append(content, map[string]interface{}{
 								"type": "text",
@@ -596,7 +596,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 
 				// Add tool calls as content if present
 				if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
-					for _, toolCall := range *msg.AssistantMessage.ToolCalls {
+					for _, toolCall := range msg.AssistantMessage.ToolCalls {
 						if toolCall.Function.Name != nil {
 							var input map[string]interface{}
 							if toolCall.Function.Arguments != "" {
@@ -639,9 +639,9 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 	}
 
 	// Transform tools if present
-	if params != nil && params.Tools != nil && len(*params.Tools) > 0 {
+	if params != nil && params.Tools != nil && len(params.Tools) > 0 {
 		var tools []map[string]interface{}
-		for _, tool := range *params.Tools {
+		for _, tool := range params.Tools {
 			tools = append(tools, map[string]interface{}{
 				"name":         tool.Function.Name,
 				"description":  tool.Function.Description,
@@ -791,7 +791,7 @@ func parseAnthropicResponse(response *AnthropicChatResponse, bifrostResponse *sc
 	if len(toolCalls) > 0 || thinking != "" {
 		assistantMessage = &schemas.AssistantMessage{}
 		if len(toolCalls) > 0 {
-			assistantMessage.ToolCalls = &toolCalls
+			assistantMessage.ToolCalls = toolCalls
 		}
 		if thinking != "" {
 			assistantMessage.Thought = &thinking
@@ -807,7 +807,7 @@ func parseAnthropicResponse(response *AnthropicChatResponse, bifrostResponse *sc
 				Message: schemas.BifrostMessage{
 					Role: schemas.ModelChatMessageRoleAssistant,
 					Content: schemas.MessageContent{
-						ContentBlocks: &contentBlocks,
+						ContentBlocks: contentBlocks,
 					},
 					AssistantMessage: assistantMessage,
 				},