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/providers/anthropic.go

@@ -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

core/providers/bedrock.go

@@ -88,7 +88,7 @@ type BedrockMistralContent struct {
 type BedrockMistralChatMessage struct {
 	Role       schemas.ModelChatMessageRole `json:"role"`                   // Role of the message sender
 	Content    []BedrockMistralContent      `json:"content"`                // Array of message content
-	ToolCalls  *[]BedrockAnthropicToolCall  `json:"tool_calls,omitempty"`   // Optional tool calls
+	ToolCalls  []BedrockAnthropicToolCall  `json:"tool_calls,omitempty"`   // Optional tool calls
 	ToolCallID *string                      `json:"tool_call_id,omitempty"` // Optional tool call ID
 }
 
@@ -514,7 +514,7 @@ func (provider *BedrockProvider) prepareChatCompletionMessages(messages []schema
 				} else {
 					// Bedrock wants only toolUse block on content, text blocks are not allowed when tools are called.
 					if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
-						for _, toolCall := range *msg.AssistantMessage.ToolCalls {
+						for _, toolCall := range msg.AssistantMessage.ToolCalls {
 							var input map[string]interface{}
 							if toolCall.Function.Arguments != "" {
 								if err := sonic.Unmarshal([]byte(toolCall.Function.Arguments), &input); err != nil {
@@ -685,7 +685,7 @@ func (provider *BedrockProvider) prepareChatCompletionMessages(messages []schema
 
 			var filteredToolCalls []BedrockAnthropicToolCall
 			if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
-				for _, toolCall := range *msg.AssistantMessage.ToolCalls {
+				for _, toolCall := range msg.AssistantMessage.ToolCalls {
 					if toolCall.ID != nil && toolCall.Function.Name != nil {
 						// Parse the arguments to get parameters
 						var params interface{}
@@ -880,7 +880,7 @@ func (provider *BedrockProvider) extractToolsFromHistory(messages []schemas.Bifr
 		if msg.Role == schemas.ModelChatMessageRoleAssistant && msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
 			hasToolContent = true
 			// Extract tool definitions from tool calls for toolConfig
-			for _, toolCall := range *msg.AssistantMessage.ToolCalls {
+			for _, toolCall := range msg.AssistantMessage.ToolCalls {
 				if toolCall.Function.Name != nil {
 					toolName := *toolCall.Function.Name
 					if _, exists := seenTools[toolName]; !exists {
@@ -1077,7 +1077,7 @@ func (provider *BedrockProvider) ChatCompletion(ctx context.Context, model strin
 	// Create AssistantMessage if we have tool calls
 	if len(toolCalls) > 0 {
 		assistantMessage = &schemas.AssistantMessage{
-			ToolCalls: &toolCalls,
+			ToolCalls: toolCalls,
 		}
 	}
 

core/providers/cohere.go

@@ -146,7 +146,7 @@ type CohereStreamTextEvent struct {
 	Text      string `json:"text"`       // Text content being generated
 }
 
-// CohereStreamToolEvent represents the tool use event.
+// CohereStreamToolCallEvent represents the tool use event.
 type CohereStreamToolCallEvent struct {
 	EventType string `json:"event_type"` // tool-use
 	ToolCall  struct {
@@ -338,7 +338,7 @@ func (provider *CohereProvider) ChatCompletion(ctx context.Context, model string
 							ContentStr: &content,
 						},
 						AssistantMessage: &schemas.AssistantMessage{
-							ToolCalls: &toolCalls,
+							ToolCalls: toolCalls,
 						},
 					},
 				},
@@ -391,7 +391,7 @@ func prepareCohereChatRequest(messages []schemas.BifrostMessage, params *schemas
 		if msg.Role == schemas.ModelChatMessageRoleAssistant {
 			if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
 				var toolCalls []map[string]interface{}
-				for _, toolCall := range *msg.AssistantMessage.ToolCalls {
+				for _, toolCall := range msg.AssistantMessage.ToolCalls {
 					var arguments map[string]interface{}
 					var parsedJSON interface{}
 					err := sonic.Unmarshal([]byte(toolCall.Function.Arguments), &parsedJSON)
@@ -424,7 +424,7 @@ func prepareCohereChatRequest(messages []schemas.BifrostMessage, params *schemas
 					prevMsg.AssistantMessage.ToolCalls != nil {
 
 					// Search through tool calls in this assistant message
-					for _, toolCall := range *prevMsg.AssistantMessage.ToolCalls {
+					for _, toolCall := range prevMsg.AssistantMessage.ToolCalls {
 						if toolCall.ID != nil && msg.ToolMessage != nil && msg.ToolMessage.ToolCallID != nil &&
 							*toolCall.ID == *msg.ToolMessage.ToolCallID {
 
@@ -588,7 +588,7 @@ func convertChatHistory(history []struct {
 	Role      schemas.ModelChatMessageRole `json:"role"`
 	Message   string                       `json:"message"`
 	ToolCalls []CohereToolCall             `json:"tool_calls"`
-}) *[]schemas.BifrostMessage {
+}) []schemas.BifrostMessage {
 	converted := make([]schemas.BifrostMessage, len(history))
 	for i, msg := range history {
 		var toolCalls []schemas.ToolCall
@@ -617,7 +617,7 @@ func convertChatHistory(history []struct {
 				ContentStr: &msg.Message,
 			},
 			AssistantMessage: &schemas.AssistantMessage{
-				ToolCalls: &toolCalls,
+				ToolCalls: toolCalls,
 			},
 		}
 	}

core/providers/gemini.go

@@ -17,15 +17,15 @@ import (
 	"github.com/valyala/fasthttp"
 )
 
-// Response message for PredictionService.GenerateContent.
+// GenerateContentResponse represents the response message for PredictionService.GenerateContent.
 type GenerateContentResponse struct {
 	// Response variations returned by the model.
 	Candidates []*Candidate `json:"candidates,omitempty"`
 	// Usage metadata about the response(s).
 	UsageMetadata *GenerateContentResponseUsageMetadata `json:"usageMetadata,omitempty"`
 }
 
-// A response candidate generated from the model.
+// Candidate represents a response candidate generated from the model.
 type Candidate struct {
 	// Optional. Contains the multi-part content of the response.
 	Content *Content `json:"content,omitempty"`
@@ -36,7 +36,7 @@ type Candidate struct {
 	Index int32 `json:"index,omitempty"`
 }
 
-// Contains the multi-part content of a message.
+// Content represents the multi-part content of a message.
 type Content struct {
 	// Optional. List of parts that constitute a single message. Each part may have
 	// a different IANA MIME type.
@@ -47,7 +47,7 @@ type Content struct {
 	Role string `json:"role,omitempty"`
 }
 
-// A datatype containing media content.
+// Part represents a datatype containing media content.
 // Exactly one field within a Part should be set, representing the specific type
 // of content being conveyed. Using multiple fields within the same `Part`
 // instance is considered invalid.
@@ -58,13 +58,13 @@ type Part struct {
 	Text string `json:"text,omitempty"`
 }
 
-// Content blob.
+// Blob represents a content blob.
 type Blob struct {
 	// Required. Raw bytes.
 	Data []byte `json:"data,omitempty"`
 }
 
-// Usage metadata about response(s).
+// GenerateContentResponseUsageMetadata represents usage metadata about response(s).
 type GenerateContentResponseUsageMetadata struct {
 	// Number of tokens in the response(s). This includes all the generated response candidates.
 	CandidatesTokenCount int32 `json:"candidatesTokenCount,omitempty"`
@@ -198,10 +198,10 @@ func (provider *GeminiProvider) ChatCompletion(ctx context.Context, model string
 		if choice.Message.AssistantMessage == nil || choice.Message.AssistantMessage.ToolCalls == nil {
 			continue
 		}
-		for i, toolCall := range *choice.Message.AssistantMessage.ToolCalls {
+		for i, toolCall := range choice.Message.AssistantMessage.ToolCalls {
 			if (toolCall.ID == nil || *toolCall.ID == "") && toolCall.Function.Name != nil && *toolCall.Function.Name != "" {
 				id := *toolCall.Function.Name
-				(*choice.Message.AssistantMessage.ToolCalls)[i].ID = &id
+				choice.Message.AssistantMessage.ToolCalls[i].ID = &id
 			}
 		}
 	}

core/providers/openai.go

@@ -192,7 +192,7 @@ func prepareOpenAIChatRequest(messages []schemas.BifrostMessage, params *schemas
 				"content": msg.Content,
 			}
 			if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
-				assistantMessage["tool_calls"] = *msg.AssistantMessage.ToolCalls
+				assistantMessage["tool_calls"] = msg.AssistantMessage.ToolCalls
 			}
 			formattedMessages = append(formattedMessages, assistantMessage)
 		} else {