From aeef5b58f33fc30f6ef8b8a9f15c9eefcbdcc2ce Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Mon, 6 Oct 2025 14:44:43 +0530 Subject: [PATCH] feat: fallbacks added to vk provider routing --- core/bifrost.go | 23 +++-- core/go.mod | 2 +- core/providers/sgl.go | 12 ++- core/schemas/bifrost.go | 1 + core/schemas/providers/gemini/chat.go | 86 ++++++++++++++++--- framework/configstore/migrations.go | 50 +++++++++++ framework/configstore/rdb.go | 2 + framework/logstore/tables.go | 1 + plugins/logging/main.go | 28 ++++-- plugins/logging/operations.go | 6 +- tests/core-providers/sgl_test.go | 11 +-- .../bifrost-http/handlers/completions.go | 12 +-- .../bifrost-http/handlers/middlewares.go | 26 +++++- transports/bifrost-http/handlers/server.go | 5 +- ui/app/not-found.tsx | 4 +- ui/app/teams-customers/page.tsx | 82 +++++++++++------- ui/app/virtual-keys/page.tsx | 70 ++++++++++----- 17 files changed, 324 insertions(+), 97 deletions(-) diff --git a/core/bifrost.go b/core/bifrost.go index 4278d1e56..b325e61a7 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -13,6 +13,7 @@ import ( "sync/atomic" "time" + "github.com/google/uuid" "github.com/maximhq/bifrost/core/providers" schemas "github.com/maximhq/bifrost/core/schemas" ) @@ -681,10 +682,11 @@ transferComplete: providerKey, providerConfig.ConcurrencyAndBufferSize.BufferSize) + waitGroupValue, _ := bifrost.waitGroups.Load(providerKey) + currentWaitGroup := waitGroupValue.(*sync.WaitGroup) + for range providerConfig.ConcurrencyAndBufferSize.Concurrency { - waitGroupValue, _ := bifrost.waitGroups.Load(providerKey) - waitGroup := waitGroupValue.(*sync.WaitGroup) - waitGroup.Add(1) + currentWaitGroup.Add(1) go bifrost.requestWorker(provider, providerConfig, newQueue) } @@ -992,10 +994,11 @@ func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, confi return fmt.Errorf("failed to create provider for the given key: %v", err) } + waitGroupValue, _ := bifrost.waitGroups.Load(providerKey) + currentWaitGroup := waitGroupValue.(*sync.WaitGroup) + for range providerConfig.ConcurrencyAndBufferSize.Concurrency { - waitGroupValue, _ := bifrost.waitGroups.Load(providerKey) - waitGroup := waitGroupValue.(*sync.WaitGroup) - waitGroup.Add(1) + currentWaitGroup.Add(1) go bifrost.requestWorker(provider, providerConfig, queue) } @@ -1182,6 +1185,8 @@ func (bifrost *Bifrost) handleRequest(ctx context.Context, req *schemas.BifrostR // Try fallbacks in order for _, fallback := range req.Fallbacks { + ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackRequestID, uuid.New().String()) + fallbackReq := bifrost.prepareFallbackRequest(req, fallback) if fallbackReq == nil { continue @@ -1190,7 +1195,7 @@ func (bifrost *Bifrost) handleRequest(ctx context.Context, req *schemas.BifrostR // Try the fallback provider result, fallbackErr := bifrost.tryRequest(fallbackReq, ctx) if fallbackErr == nil { - bifrost.logger.Info(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) + bifrost.logger.Debug(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) return result, nil } @@ -1234,6 +1239,8 @@ func (bifrost *Bifrost) handleStreamRequest(ctx context.Context, req *schemas.Bi // Try fallbacks in order for _, fallback := range req.Fallbacks { + ctx = context.WithValue(ctx, schemas.BifrostContextKeyFallbackRequestID, uuid.New().String()) + fallbackReq := bifrost.prepareFallbackRequest(req, fallback) if fallbackReq == nil { continue @@ -1242,7 +1249,7 @@ func (bifrost *Bifrost) handleStreamRequest(ctx context.Context, req *schemas.Bi // Try the fallback provider result, fallbackErr := bifrost.tryStreamRequest(fallbackReq, ctx) if fallbackErr == nil { - bifrost.logger.Info(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) + bifrost.logger.Debug(fmt.Sprintf("Successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) return result, nil } diff --git a/core/go.mod b/core/go.mod index 94b20c261..c897c92f6 100644 --- a/core/go.mod +++ b/core/go.mod @@ -8,6 +8,7 @@ require ( github.com/aws/aws-sdk-go-v2 v1.38.0 github.com/aws/aws-sdk-go-v2/config v1.31.0 github.com/bytedance/sonic v1.14.0 + github.com/google/uuid v1.6.0 github.com/mark3labs/mcp-go v0.37.0 github.com/rs/zerolog v1.34.0 github.com/valyala/fasthttp v1.65.0 @@ -33,7 +34,6 @@ require ( github.com/bytedance/sonic/loader v0.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect - github.com/google/uuid v1.6.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/klauspost/compress v1.18.0 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect diff --git a/core/providers/sgl.go b/core/providers/sgl.go index e93885c61..60eaa8bc4 100644 --- a/core/providers/sgl.go +++ b/core/providers/sgl.go @@ -114,7 +114,17 @@ func (provider *SGLProvider) Responses(ctx context.Context, key schemas.Key, req // Embedding is not supported by the SGL provider. func (provider *SGLProvider) Embedding(ctx context.Context, key schemas.Key, request *schemas.BifrostEmbeddingRequest) (*schemas.BifrostResponse, *schemas.BifrostError) { - return nil, newUnsupportedOperationError("embedding", "sgl") + return handleOpenAIEmbeddingRequest( + ctx, + provider.client, + provider.networkConfig.BaseURL+"/v1/embeddings", + request, + key, + provider.networkConfig.ExtraHeaders, + provider.GetProviderKey(), + provider.sendBackRawResponse, + provider.logger, + ) } // ChatCompletionStream performs a streaming chat completion request to the SGL API. diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 474acf071..c62c8bf55 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -107,6 +107,7 @@ type BifrostContextKey string // BifrostContextKeyRequestType is a context key for the request type. const ( BifrostContextKeyRequestID BifrostContextKey = "request-id" + BifrostContextKeyFallbackRequestID BifrostContextKey = "fallback-request-id" BifrostContextKeyVirtualKeyHeader BifrostContextKey = "x-bf-vk" BifrostContextKeyDirectKey BifrostContextKey = "bifrost-direct-key" BifrostContextKeyStreamEndIndicator BifrostContextKey = "bifrost-stream-end-indicator" diff --git a/core/schemas/providers/gemini/chat.go b/core/schemas/providers/gemini/chat.go index 74de01916..813f9fb28 100644 --- a/core/schemas/providers/gemini/chat.go +++ b/core/schemas/providers/gemini/chat.go @@ -386,19 +386,56 @@ func (r *GenerateContentResponse) ToBifrostResponse() *schemas.BifrostResponse { }, } } else if hasText && textContent != "" { - // This is a transcription response - response.Object = "audio.transcription" - response.Transcribe = &schemas.BifrostTranscribe{ - Text: textContent, - Usage: &schemas.TranscriptionUsage{ - Type: "tokens", - InputTokens: &inputTokens, - OutputTokens: &outputTokens, - TotalTokens: &totalTokens, - }, - BifrostTranscribeNonStreamResponse: &schemas.BifrostTranscribeNonStreamResponse{ - Task: schemas.Ptr("transcribe"), - }, + // Check if this is actually a transcription response by looking for transcription context + // Only treat as transcription if we have explicit transcription metadata or context + isTranscription := r.isTranscriptionResponse() + + if isTranscription { + // This is a transcription response + response.Object = "audio.transcription" + response.Transcribe = &schemas.BifrostTranscribe{ + Text: textContent, + Usage: &schemas.TranscriptionUsage{ + Type: "tokens", + InputTokens: &inputTokens, + OutputTokens: &outputTokens, + TotalTokens: &totalTokens, + }, + BifrostTranscribeNonStreamResponse: &schemas.BifrostTranscribeNonStreamResponse{ + Task: schemas.Ptr("transcribe"), + }, + } + } else { + // This is a regular chat completion response + response.Object = "chat.completion" + + // Create choice from the candidate + choice := schemas.BifrostChatResponseChoice{ + Index: 0, + BifrostNonStreamResponseChoice: &schemas.BifrostNonStreamResponseChoice{ + Message: schemas.ChatMessage{ + Role: schemas.ChatMessageRoleAssistant, + Content: schemas.ChatMessageContent{ + ContentStr: &textContent, + }, + }, + }, + } + + // Set finish reason if available + if candidate.FinishReason != "" { + finishReason := string(candidate.FinishReason) + choice.FinishReason = &finishReason + } + + response.Choices = []schemas.BifrostChatResponseChoice{choice} + + // Set usage information + response.Usage = &schemas.LLMUsage{ + PromptTokens: inputTokens, + CompletionTokens: outputTokens, + TotalTokens: totalTokens, + } } } } @@ -407,6 +444,29 @@ func (r *GenerateContentResponse) ToBifrostResponse() *schemas.BifrostResponse { return response } +// isTranscriptionResponse determines if this response is from a transcription request +// by checking for transcription-specific context and metadata +func (r *GenerateContentResponse) isTranscriptionResponse() bool { + // Check if any candidates contain audio input data in their parts + // This would indicate the original request included audio for transcription + for _, candidate := range r.Candidates { + if candidate.Content != nil { + for _, part := range candidate.Content.Parts { + if part.InlineData != nil && part.InlineData.MIMEType != "" { + // If we have audio data in the response parts, it's likely a transcription + if strings.HasPrefix(part.InlineData.MIMEType, "audio/") { + return true + } + } + } + } + } + + // Default to false - assume it's a regular chat completion + // This is safer than incorrectly classifying chat responses as transcriptions + return false +} + // FromBifrostResponse converts a BifrostResponse back to Gemini's GenerateContentResponse func ToGeminiGenerationResponse(bifrostResp *schemas.BifrostResponse) interface{} { if bifrostResp == nil { diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 3908d2aa7..a412aba2f 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -25,6 +25,12 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddOpenAIUseResponsesAPIColumn(ctx, db); err != nil { return err } + if err := migrationAddAllowedOriginsJSONColumn(ctx, db); err != nil { + return err + } + if err := migrationAddAllowDirectKeysColumn(ctx, db); err != nil { + return err + } return nil } @@ -307,3 +313,47 @@ func migrationAddOpenAIUseResponsesAPIColumn(ctx context.Context, db *gorm.DB) e } return nil } + +func migrationAddAllowedOriginsJSONColumn(ctx context.Context, db *gorm.DB) error { + m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{ + ID: "add_allowed_origins_json_column", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + if !migrator.HasColumn(&TableClientConfig{}, "allowed_origins_json") { + if err := migrator.AddColumn(&TableClientConfig{}, "allowed_origins_json"); err != nil { + return err + } + } + return nil + }, + }}) + err := m.Migrate() + if err != nil { + return fmt.Errorf("error while running db migration: %s", err.Error()) + } + return nil +} + +func migrationAddAllowDirectKeysColumn(ctx context.Context, db *gorm.DB) error { + m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{ + ID: "add_allow_direct_keys_column", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + + if !migrator.HasColumn(&TableClientConfig{}, "allow_direct_keys") { + if err := migrator.AddColumn(&TableClientConfig{}, "allow_direct_keys"); err != nil { + return err + } + } + return nil + }, + }}) + err := m.Migrate() + if err != nil { + return fmt.Errorf("error while running db migration: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 3473ce51b..5d2e45b99 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -818,6 +818,7 @@ func (s *RDBConfigStore) GetVirtualKeys(ctx context.Context) ([]TableVirtualKey, Preload("Customer"). Preload("Budget"). Preload("RateLimit"). + Preload("ProviderConfigs"). Preload("Keys", func(db *gorm.DB) *gorm.DB { return db.Select("id, key_id, models_json") }).Find(&virtualKeys).Error; err != nil { @@ -834,6 +835,7 @@ func (s *RDBConfigStore) GetVirtualKey(ctx context.Context, id string) (*TableVi Preload("Customer"). Preload("Budget"). Preload("RateLimit"). + Preload("ProviderConfigs"). Preload("Keys", func(db *gorm.DB) *gorm.DB { return db.Select("id, key_id, models_json") }).First(&virtualKey, "id = ?", id).Error; err != nil { diff --git a/framework/logstore/tables.go b/framework/logstore/tables.go index 974f1692e..b6bde627d 100644 --- a/framework/logstore/tables.go +++ b/framework/logstore/tables.go @@ -69,6 +69,7 @@ type SearchStats struct { // This is the GORM model with appropriate tags type Log struct { ID string `gorm:"primaryKey;type:varchar(255)" json:"id"` + ParentRequestID *string `gorm:"type:varchar(255)" json:"parent_request_id"` Timestamp time.Time `gorm:"index;not null" json:"timestamp"` Object string `gorm:"type:varchar(255);index;not null;column:object_type" json:"object"` // text.completion, chat.completion, or embedding Provider string `gorm:"type:varchar(255);index;not null" json:"provider"` diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 29127acd4..80cad0f7b 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -20,7 +20,7 @@ import ( ) const ( - PluginName = "bifrost-http-logging" + PluginName = "logging" ) // ContextKey is a custom type for context keys to prevent collisions @@ -37,8 +37,8 @@ const ( // Context keys for logging optimization const ( - DroppedCreateContextKey ContextKey = "bifrost-logging-dropped" - CreatedTimestampKey ContextKey = "bifrost-logging-created-timestamp" + DroppedCreateContextKey ContextKey = "logging-dropped" + CreatedTimestampKey ContextKey = "logging-created-timestamp" ) // UpdateLogData contains data for log entry updates @@ -59,7 +59,8 @@ type UpdateLogData struct { // LogMessage represents a message in the logging queue type LogMessage struct { Operation LogOperation - RequestID string + RequestID string // Unique ID for the request + ParentRequestID string // Unique ID for the parent request Timestamp time.Time // Of the preHook/postHook call InitialData *InitialLogData // For create operations SemanticCacheDebug *schemas.BifrostCacheDebug // For semantic cache operations @@ -224,6 +225,7 @@ func (p *LoggerPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest p.logger.Error("request-id not found in context or is empty") return req, nil, nil } + createdTimestamp := time.Now() // If request type is streaming we create a stream accumulator if bifrost.IsStreamRequestType(req.RequestType) { @@ -271,13 +273,22 @@ func (p *LoggerPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest // Queue the log creation message (non-blocking) - Using sync.Pool logMsg := p.getLogMessage() logMsg.Operation = LogOperationCreate - logMsg.RequestID = requestID + + // If fallback request ID is present, use it instead of the primary request ID + fallbackRequestID, ok := (*ctx).Value(schemas.BifrostContextKeyFallbackRequestID).(string) + if ok && fallbackRequestID != "" { + logMsg.RequestID = fallbackRequestID + logMsg.ParentRequestID = requestID + } else { + logMsg.RequestID = requestID + } + logMsg.Timestamp = createdTimestamp logMsg.InitialData = initialData go func(logMsg *LogMessage) { defer p.putLogMessage(logMsg) // Return to pool when done - if err := p.insertInitialLogEntry(p.ctx, logMsg.RequestID, logMsg.Timestamp, logMsg.InitialData); err != nil { + if err := p.insertInitialLogEntry(p.ctx, logMsg.RequestID, logMsg.ParentRequestID, logMsg.Timestamp, logMsg.InitialData); err != nil { p.logger.Error("failed to insert initial log entry for request %s: %v", logMsg.RequestID, err) } else { // Call callback for initial log creation (WebSocket "create" message) @@ -324,6 +335,11 @@ func (p *LoggerPlugin) PostHook(ctx *context.Context, result *schemas.BifrostRes p.logger.Error("request-id not found in context or is empty") return result, bifrostErr, nil } + // If fallback request ID is present, use it instead of the primary request ID + fallbackRequestID, ok := (*ctx).Value(schemas.BifrostContextKeyFallbackRequestID).(string) + if ok && fallbackRequestID != "" { + requestID = fallbackRequestID + } requestType, _, _ := bifrost.GetRequestFields(result, bifrostErr) // Queue the log update message (non-blocking) - use same pattern for both streaming and regular logMsg := p.getLogMessage() diff --git a/plugins/logging/operations.go b/plugins/logging/operations.go index 108a52bad..369051d0a 100644 --- a/plugins/logging/operations.go +++ b/plugins/logging/operations.go @@ -12,7 +12,7 @@ import ( ) // insertInitialLogEntry creates a new log entry in the database using GORM -func (p *LoggerPlugin) insertInitialLogEntry(ctx context.Context, requestID string, timestamp time.Time, data *InitialLogData) error { +func (p *LoggerPlugin) insertInitialLogEntry(ctx context.Context, requestID string, parentRequestID string, timestamp time.Time, data *InitialLogData) error { entry := &logstore.Log{ ID: requestID, Timestamp: timestamp, @@ -30,6 +30,10 @@ func (p *LoggerPlugin) insertInitialLogEntry(ctx context.Context, requestID stri TranscriptionInputParsed: data.TranscriptionInput, } + if parentRequestID != "" { + entry.ParentRequestID = &parentRequestID + } + return p.store.Create(ctx, entry) } diff --git a/tests/core-providers/sgl_test.go b/tests/core-providers/sgl_test.go index 00f90efaa..e9077de6a 100644 --- a/tests/core-providers/sgl_test.go +++ b/tests/core-providers/sgl_test.go @@ -18,11 +18,12 @@ func TestSGL(t *testing.T) { testConfig := config.ComprehensiveTestConfig{ Provider: schemas.SGL, - ChatModel: "Qwen2.5-VL-7B-Instruct", - TextModel: "", // SGL doesn't support text completion - EmbeddingModel: "", // SGL doesn't support embedding + ChatModel: "qwen/qwen2.5-0.5b-instruct", + VisionModel: "Qwen/Qwen2.5-VL-7B-Instruct", + TextModel: "qwen/qwen2.5-0.5b-instruct", + EmbeddingModel: "Alibaba-NLP/gte-Qwen2-1.5B-instruct", Scenarios: config.TestScenarios{ - TextCompletion: false, // Not supported + TextCompletion: true, SimpleChat: true, ChatCompletionStream: true, MultiTurnConversation: true, @@ -34,7 +35,7 @@ func TestSGL(t *testing.T) { ImageBase64: true, MultipleImages: true, CompleteEnd2End: true, - Embedding: false, + Embedding: true, }, } diff --git a/transports/bifrost-http/handlers/completions.go b/transports/bifrost-http/handlers/completions.go index 8671df07f..d926a4b5a 100644 --- a/transports/bifrost-http/handlers/completions.go +++ b/transports/bifrost-http/handlers/completions.go @@ -194,12 +194,14 @@ type TranscriptionRequest struct { // parseFallbacks extracts fallbacks from string array and converts to Fallback structs func parseFallbacks(fallbackStrings []string) ([]schemas.Fallback, error) { - fallbacks := make([]schemas.Fallback, len(fallbackStrings)) - for i, fallback := range fallbackStrings { + fallbacks := make([]schemas.Fallback, 0, len(fallbackStrings)) + for _, fallback := range fallbackStrings { fallbackProvider, fallbackModelName := schemas.ParseModelString(fallback, "") - fallbacks[i] = schemas.Fallback{ - Provider: fallbackProvider, - Model: fallbackModelName, + if fallbackProvider != "" && fallbackModelName != "" { + fallbacks = append(fallbacks, schemas.Fallback{ + Provider: fallbackProvider, + Model: fallbackModelName, + }) } } return fallbacks, nil diff --git a/transports/bifrost-http/handlers/middlewares.go b/transports/bifrost-http/handlers/middlewares.go index 166cc86d6..65dc683e8 100644 --- a/transports/bifrost-http/handlers/middlewares.go +++ b/transports/bifrost-http/handlers/middlewares.go @@ -6,6 +6,7 @@ import ( "fmt" "math/rand" "slices" + "sort" "strings" "github.com/maximhq/bifrost/core/schemas" @@ -137,7 +138,8 @@ func VKProviderRoutingMiddleware(config *lib.Config, logger schemas.Logger) Bifr next(ctx) return } - // Weighted random selection from allowed providers + + // Weighted random selection from allowed providers for the main model totalWeight := 0.0 for _, config := range allowedProviderConfigs { totalWeight += config.Weight @@ -154,8 +156,30 @@ func VKProviderRoutingMiddleware(config *lib.Config, logger schemas.Logger) Bifr break } } + // Update the model field in the request body requestBody["model"] = string(selectedProvider) + "/" + modelStr + + // Check if fallbacks field is already present + _, hasFallbacks := requestBody["fallbacks"] + if !hasFallbacks && len(allowedProviderConfigs) > 1 { + // Sort allowed provider configs by weight (descending) + sort.Slice(allowedProviderConfigs, func(i, j int) bool { + return allowedProviderConfigs[i].Weight > allowedProviderConfigs[j].Weight + }) + + // Filter out the selected provider and create fallbacks array + fallbacks := make([]string, 0, len(allowedProviderConfigs)-1) + for _, config := range allowedProviderConfigs { + if config.Provider != string(selectedProvider) { + fallbacks = append(fallbacks, string(schemas.ModelProvider(config.Provider))+"/"+modelStr) + } + } + + // Add fallbacks to request body + requestBody["fallbacks"] = fallbacks + } + // Marshal the updated request body back to JSON updatedBody, err := json.Marshal(requestBody) if err != nil { diff --git a/transports/bifrost-http/handlers/server.go b/transports/bifrost-http/handlers/server.go index 86920b029..1731af4fe 100644 --- a/transports/bifrost-http/handlers/server.go +++ b/transports/bifrost-http/handlers/server.go @@ -34,7 +34,7 @@ import ( const ( DefaultHost = "localhost" DefaultPort = "8080" - DefaultAppDir = "./bifrost-data" + DefaultAppDir = "" // Empty string means use OS-specific config directory DefaultLogLevel = string(schemas.LogLevelInfo) DefaultLogOutputStyle = string(schemas.LoggerOutputTypeJSON) ) @@ -447,6 +447,9 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { s.ctx, s.cancel = context.WithCancel(ctx) SetVersion(s.Version) configDir := GetDefaultConfigDir(s.AppDir) + + fmt.Println("configDir", configDir) + // Ensure app directory exists if err := os.MkdirAll(configDir, 0755); err != nil { return fmt.Errorf("failed to create app directory %s: %v", configDir, err) diff --git a/ui/app/not-found.tsx b/ui/app/not-found.tsx index 092a19045..8a975af41 100644 --- a/ui/app/not-found.tsx +++ b/ui/app/not-found.tsx @@ -2,14 +2,14 @@ import Link from "next/link"; export default function NotFound() { return ( -
+

404

Page not found

The page you are looking for doesn’t exist or has been moved

Go home diff --git a/ui/app/teams-customers/page.tsx b/ui/app/teams-customers/page.tsx index 631cae4ae..e1143bc13 100644 --- a/ui/app/teams-customers/page.tsx +++ b/ui/app/teams-customers/page.tsx @@ -1,7 +1,13 @@ "use client"; import FullPageLoader from "@/components/fullPageLoader"; -import { getErrorMessage, useGetCoreConfigQuery, useGetCustomersQuery, useGetTeamsQuery, useGetVirtualKeysQuery } from "@/lib/store"; +import { + getErrorMessage, + useLazyGetCoreConfigQuery, + useLazyGetCustomersQuery, + useLazyGetTeamsQuery, + useLazyGetVirtualKeysQuery, +} from "@/lib/store"; import { cn } from "@/lib/utils"; import { useEffect, useState } from "react"; import { toast } from "sonner"; @@ -10,45 +16,57 @@ import TeamsTable from "./views/teamsTable"; export default function TeamsCustomersPage() { const [activeTab, setActiveTab] = useState("teams"); + const [governanceEnabled, setGovernanceEnabled] = useState(null); - // Fetch all data with RTK Query - const { data: virtualKeysData, error: vkError, isLoading: vkLoading, refetch: refetchVirtualKeys } = useGetVirtualKeysQuery(); - const { data: teamsData, error: teamsError, isLoading: teamsLoading, refetch: refetchTeams } = useGetTeamsQuery({}); - const { data: customersData, error: customersError, isLoading: customersLoading, refetch: refetchCustomers } = useGetCustomersQuery(); - const { data: coreConfig, error: configError, isLoading: configLoading } = useGetCoreConfigQuery({ fromDB: true }); + // Lazy query hooks + const [triggerGetVirtualKeys, { data: virtualKeysData, error: vkError, isLoading: vkLoading }] = useLazyGetVirtualKeysQuery(); + const [triggerGetTeams, { data: teamsData, error: teamsError, isLoading: teamsLoading }] = useLazyGetTeamsQuery(); + const [triggerGetCustomers, { data: customersData, error: customersError, isLoading: customersLoading }] = useLazyGetCustomersQuery(); + const [triggerGetConfig] = useLazyGetCoreConfigQuery(); - const isLoading = vkLoading || teamsLoading || customersLoading || configLoading; + const isLoading = vkLoading || teamsLoading || customersLoading || governanceEnabled === null; - // Handle errors + // Check governance and trigger queries conditionally useEffect(() => { - if (configLoading) return; - if (configError) { - toast.error(`Failed to load core config: ${getErrorMessage(configError)}`); - return; - } - - if (coreConfig && !coreConfig?.client_config?.enable_governance) { - toast.error("Governance is not enabled. Please enable it in the core settings."); - return; - } - - if (vkError) { - toast.error(`Failed to load virtual keys: ${getErrorMessage(vkError)}`); - } + triggerGetConfig({ fromDB: true }).then((res) => { + if (res.data && res.data.client_config.enable_governance) { + setGovernanceEnabled(true); + // Trigger lazy queries only when governance is enabled + triggerGetVirtualKeys(); + triggerGetTeams({}); + triggerGetCustomers(); + } else { + setGovernanceEnabled(false); + toast.error("Governance is not enabled. Please enable it in the config."); + } + }); + }, [triggerGetConfig, triggerGetVirtualKeys, triggerGetTeams, triggerGetCustomers]); - if (teamsError) { - toast.error(`Failed to load teams: ${getErrorMessage(teamsError)}`); - } - - if (customersError) { - toast.error(`Failed to load customers: ${getErrorMessage(customersError)}`); + // Handle query errors - show consolidated error if all APIs fail + useEffect(() => { + if (vkError && teamsError && customersError) { + // If all three APIs fail, suggest resetting bifrost + toast.error("Failed to load governance data. Please reset Bifrost to enable governance properly."); + } else { + // Show individual errors if only some APIs fail + if (vkError) { + toast.error(`Failed to load virtual keys: ${getErrorMessage(vkError)}`); + } + if (teamsError) { + toast.error(`Failed to load teams: ${getErrorMessage(teamsError)}`); + } + if (customersError) { + toast.error(`Failed to load customers: ${getErrorMessage(customersError)}`); + } } - }, [configError, coreConfig, vkError, teamsError, customersError]); + }, [vkError, teamsError, customersError]); const handleRefresh = () => { - refetchVirtualKeys(); - refetchTeams(); - refetchCustomers(); + if (governanceEnabled) { + triggerGetVirtualKeys(); + triggerGetTeams({}); + triggerGetCustomers(); + } }; if (isLoading) { diff --git a/ui/app/virtual-keys/page.tsx b/ui/app/virtual-keys/page.tsx index e86ac2db7..4de9ed089 100644 --- a/ui/app/virtual-keys/page.tsx +++ b/ui/app/virtual-keys/page.tsx @@ -1,40 +1,68 @@ "use client"; import FullPageLoader from "@/components/fullPageLoader"; -import { getErrorMessage, useGetCustomersQuery, useGetTeamsQuery, useGetVirtualKeysQuery } from "@/lib/store"; -import { useEffect } from "react"; +import { + getErrorMessage, + useLazyGetCustomersQuery, + useLazyGetTeamsQuery, + useLazyGetVirtualKeysQuery, + useLazyGetCoreConfigQuery, +} from "@/lib/store"; +import { useEffect, useState } from "react"; import { toast } from "sonner"; import VirtualKeysTable from "./views/virtualKeysTable"; export default function VirtualKeysPage() { - const { data: virtualKeysData, error: vkError, isLoading: vkLoading, refetch: refetchVirtualKeys } = useGetVirtualKeysQuery(); - const { data: teamsData, error: teamsError, isLoading: teamsLoading, refetch: refetchTeams } = useGetTeamsQuery({}); - const { data: customersData, error: customersError, isLoading: customersLoading, refetch: refetchCustomers } = useGetCustomersQuery(); + const [governanceEnabled, setGovernanceEnabled] = useState(null); - const isLoading = vkLoading || teamsLoading || customersLoading; + const [triggerGetVirtualKeys, { data: virtualKeysData, error: vkError, isLoading: vkLoading }] = useLazyGetVirtualKeysQuery(); + const [triggerGetTeams, { data: teamsData, error: teamsError, isLoading: teamsLoading }] = useLazyGetTeamsQuery(); + const [triggerGetCustomers, { data: customersData, error: customersError, isLoading: customersLoading }] = useLazyGetCustomersQuery(); - useEffect(() => { - if (vkError) { - toast.error(`Failed to load virtual keys: ${getErrorMessage(vkError)}`); - } - }, [vkError]); + const isLoading = vkLoading || teamsLoading || customersLoading || governanceEnabled === null; + + const [triggerGetConfig] = useLazyGetCoreConfigQuery(); useEffect(() => { - if (teamsError) { - toast.error(`Failed to load teams: ${getErrorMessage(teamsError)}`); - } - }, [teamsError]); + triggerGetConfig({ fromDB: true }).then((res) => { + if (res.data && res.data.client_config.enable_governance) { + setGovernanceEnabled(true); + // Trigger lazy queries only when governance is enabled + triggerGetVirtualKeys(); + triggerGetTeams({}); + triggerGetCustomers(); + } else { + setGovernanceEnabled(false); + toast.error("Governance is not enabled. Please enable it in the config."); + } + }); + }, [triggerGetConfig, triggerGetVirtualKeys, triggerGetTeams, triggerGetCustomers]); + // Handle query errors - show consolidated error if all APIs fail useEffect(() => { - if (customersError) { - toast.error(`Failed to load customers: ${getErrorMessage(customersError)}`); + if (vkError && teamsError && customersError) { + // If all three APIs fail, suggest resetting bifrost + toast.error("Failed to load governance data. Please reset Bifrost to enable governance properly."); + } else { + // Show individual errors if only some APIs fail + if (vkError) { + toast.error(`Failed to load virtual keys: ${getErrorMessage(vkError)}`); + } + if (teamsError) { + toast.error(`Failed to load teams: ${getErrorMessage(teamsError)}`); + } + if (customersError) { + toast.error(`Failed to load customers: ${getErrorMessage(customersError)}`); + } } - }, [customersError]); + }, [vkError, teamsError, customersError]); const handleRefresh = () => { - refetchVirtualKeys(); - refetchTeams(); - refetchCustomers(); + if (governanceEnabled) { + triggerGetVirtualKeys(); + triggerGetTeams({}); + triggerGetCustomers(); + } }; if (isLoading) {