From 3846895310f8faff11bdbd940f956c69da92d293 Mon Sep 17 00:00:00 2001 From: Dan Piths <85949566+danpiths@users.noreply.github.com> Date: Tue, 19 May 2026 16:47:57 +0530 Subject: [PATCH] feat: add cluster-aware log metadata and per-node usage aggregation --- core/schemas/bifrost.go | 3 ++ framework/logstore/hybrid.go | 4 ++ framework/logstore/migrations.go | 63 +++++++++++++++++++++++++++ framework/logstore/rdb.go | 75 ++++++++++++++++++++++++++++++++ framework/logstore/store.go | 3 ++ framework/logstore/tables.go | 46 ++++++++++++++++++++ plugins/logging/main.go | 22 ++++++++++ 7 files changed, 216 insertions(+) diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 0ab295a5913..016800df8c6 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -228,6 +228,9 @@ const ( BifrostContextKeyHasEmittedMessageDelta BifrostContextKey = "bifrost-has-emitted-message-delta" // bool (tracks whether message_delta was already emitted during streaming - avoids duplicates) BifrostContextKeySkipDBUpdate BifrostContextKey = "bifrost-skip-db-update" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyGovernancePluginName BifrostContextKey = "governance-plugin-name" // string (name of the governance plugin that processed the request - set by bifrost) + BifrostContextKeyClusterNodeID BifrostContextKey = "bifrost-cluster-node-id" // string (cluster node ID for log attribution - set by enterprise server) + BifrostContextKeyGovernanceBudgetIDs BifrostContextKey = "bifrost-governance-budget-ids" // []string (budget IDs applicable to this request - set by governance plugin) + BifrostContextKeyGovernanceRateLimitIDs BifrostContextKey = "bifrost-governance-rate-limit-ids" // []string (rate limit IDs applicable to this request - set by governance plugin) BifrostContextKeyPromptsPluginName BifrostContextKey = "prompts-plugin-name" // string (name of the prompts plugin to use - set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyIsEnterprise BifrostContextKey = "is-enterprise" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyAvailableProviders BifrostContextKey = "available-providers" // []ModelProvider (set by bifrost - DO NOT SET THIS MANUALLY)) diff --git a/framework/logstore/hybrid.go b/framework/logstore/hybrid.go index 217900fcae1..57748b635de 100644 --- a/framework/logstore/hybrid.go +++ b/framework/logstore/hybrid.go @@ -586,6 +586,10 @@ func (h *HybridLogStore) BulkUpdateCost(ctx context.Context, updates map[string] return h.inner.BulkUpdateCost(ctx, updates) } +func (h *HybridLogStore) GetNodeUsageSince(ctx context.Context, nodeID string, since time.Time) (*NodeUsageAggregate, error) { + return h.inner.GetNodeUsageSince(ctx, nodeID, since) +} + func (h *HybridLogStore) Flush(ctx context.Context, since time.Time) error { return h.inner.Flush(ctx, since) } diff --git a/framework/logstore/migrations.go b/framework/logstore/migrations.go index ac662926dcd..a703ead5710 100644 --- a/framework/logstore/migrations.go +++ b/framework/logstore/migrations.go @@ -330,6 +330,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddDACColumnsToMCPToolLogs(ctx, db); err != nil { return err } + if err := migrationAddClusterGovernanceColumns(ctx, db); err != nil { + return err + } // migrationSplitFilterDataMatView is intentionally NOT invoked in this // release. Dropping mv_logs_filterdata while old replicas are still // serving /api/logs/filterdata from it would surface "relation does not @@ -2483,6 +2486,11 @@ var performanceIndexes = []performanceIndexDef{ name: "idx_mcp_logs_business_unit_id", sql: "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mcp_logs_business_unit_id ON mcp_tool_logs(business_unit_id)", }, + { + table: "logs", + name: "idx_logs_cluster_node_id", + sql: "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_logs_cluster_node_id ON logs(cluster_node_id, timestamp) WHERE cluster_node_id IS NOT NULL", + }, } // ensurePerformanceIndexes checks whether each performance GIN index exists and is @@ -3171,3 +3179,58 @@ func migrationAddDACColumnsToMCPToolLogs(ctx context.Context, db *gorm.DB) error } return nil } + +// migrationAddClusterGovernanceColumns adds cluster_node_id, budget_ids, and rate_limit_ids +// columns to the logs table for node usage recovery in clustered deployments. +func migrationAddClusterGovernanceColumns(ctx context.Context, db *gorm.DB) error { + opts := *migrator.DefaultOptions + opts.UseTransaction = true + m := migrator.New(db, &opts, []*migrator.Migration{{ + ID: "logs_add_cluster_governance_columns", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + if !migrator.HasColumn(&Log{}, "cluster_node_id") { + if err := migrator.AddColumn(&Log{}, "cluster_node_id"); err != nil { + return err + } + } + if !migrator.HasColumn(&Log{}, "budget_ids") { + if err := migrator.AddColumn(&Log{}, "budget_ids"); err != nil { + return err + } + } + if !migrator.HasColumn(&Log{}, "rate_limit_ids") { + if err := migrator.AddColumn(&Log{}, "rate_limit_ids"); err != nil { + return err + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + migrator := tx.Migrator() + if migrator.HasColumn(&Log{}, "cluster_node_id") { + if err := migrator.DropColumn(&Log{}, "cluster_node_id"); err != nil { + return err + } + } + if migrator.HasColumn(&Log{}, "budget_ids") { + if err := migrator.DropColumn(&Log{}, "budget_ids"); err != nil { + return err + } + } + if migrator.HasColumn(&Log{}, "rate_limit_ids") { + if err := migrator.DropColumn(&Log{}, "rate_limit_ids"); err != nil { + return err + } + } + return nil + }, + }}) + err := m.Migrate() + if err != nil { + return fmt.Errorf("error while adding cluster governance columns: %s", err.Error()) + } + return nil +} diff --git a/framework/logstore/rdb.go b/framework/logstore/rdb.go index 09f410a1c92..e68332b4f6d 100644 --- a/framework/logstore/rdb.go +++ b/framework/logstore/rdb.go @@ -352,6 +352,81 @@ func (s *RDBLogStore) BulkUpdateCost(ctx context.Context, updates map[string]flo }) } +// GetNodeUsageSince returns per-budget cost and per-rate-limit request/token usage +// for a specific cluster node from a given timestamp onwards. Usage is attributed +// to the exact budget and rate-limit IDs stored in each log entry, so callers get +// correctly scoped breakdowns rather than a single aggregate. +func (s *RDBLogStore) GetNodeUsageSince(ctx context.Context, nodeID string, since time.Time) (*NodeUsageAggregate, error) { + // Fetch only the columns needed for attribution. The row count is bounded by + // the ghost node's activity since its LastMessageAt, typically a small window. + type logRow struct { + Cost float64 `gorm:"column:cost"` + TotalTokens int64 `gorm:"column:total_tokens"` + BudgetIDs *string `gorm:"column:budget_ids"` + RateLimitIDs *string `gorm:"column:rate_limit_ids"` + } + var rows []logRow + + err := s.db.WithContext(ctx).Model(&Log{}). + Where("cluster_node_id = ?", nodeID). + Where("timestamp >= ?", since). + Where("status = ?", "success"). + Select("COALESCE(cost, 0) as cost, COALESCE(total_tokens, 0) as total_tokens, budget_ids, rate_limit_ids"). + Find(&rows).Error + if err != nil { + return nil, fmt.Errorf("failed to get node usage aggregate: %w", err) + } + + budgetCosts := make(map[string]float64) + rateLimitRequests := make(map[string]int64) + rateLimitTokens := make(map[string]int64) + + for i := range rows { + row := &rows[i] + + // Attribute cost to each budget that governed this request. + if row.BudgetIDs != nil && *row.BudgetIDs != "" { + var budgetIDs []string + if err := sonic.Unmarshal([]byte(*row.BudgetIDs), &budgetIDs); err != nil { + s.logger.Warn(fmt.Sprintf("logstore: skipping malformed budget_ids JSON in node usage aggregate: %s", err)) + } else { + // Deduplicate IDs so a row with ["b1","b1"] doesn't double-count cost. + seen := make(map[string]struct{}, len(budgetIDs)) + for _, id := range budgetIDs { + if _, dup := seen[id]; !dup { + seen[id] = struct{}{} + budgetCosts[id] += row.Cost + } + } + } + } + + // Attribute request count and tokens to each rate limit that governed this request. + if row.RateLimitIDs != nil && *row.RateLimitIDs != "" { + var rateLimitIDs []string + if err := sonic.Unmarshal([]byte(*row.RateLimitIDs), &rateLimitIDs); err != nil { + s.logger.Warn(fmt.Sprintf("logstore: skipping malformed rate_limit_ids JSON in node usage aggregate: %s", err)) + } else { + // Deduplicate IDs so a row with ["r1","r1"] doesn't double-count. + seen := make(map[string]struct{}, len(rateLimitIDs)) + for _, id := range rateLimitIDs { + if _, dup := seen[id]; !dup { + seen[id] = struct{}{} + rateLimitRequests[id]++ + rateLimitTokens[id] += row.TotalTokens + } + } + } + } + } + + return &NodeUsageAggregate{ + BudgetCosts: budgetCosts, + RateLimitRequests: rateLimitRequests, + RateLimitTokens: rateLimitTokens, + }, nil +} + // serializeLogUpdateEntry serializes parsed Log fields before passing the // update payload to GORM. Non-Log payloads are returned unchanged. func serializeLogUpdateEntry(entry any) (any, error) { diff --git a/framework/logstore/store.go b/framework/logstore/store.go index 221b8006732..d8f09e64fa0 100644 --- a/framework/logstore/store.go +++ b/framework/logstore/store.go @@ -50,6 +50,9 @@ type LogStore interface { GetDimensionTokenHistogram(ctx context.Context, filters SearchFilters, bucketSizeSeconds int64, dimension HistogramDimension) (*DimensionTokenHistogramResult, error) // GetDimensionLatencyHistogram returns time-bucketed latency percentiles grouped by the specified dimension. GetDimensionLatencyHistogram(ctx context.Context, filters SearchFilters, bucketSizeSeconds int64, dimension HistogramDimension) (*DimensionLatencyHistogramResult, error) + // GetNodeUsageSince returns cumulative cost, successful request count, and token usage + // for a specific cluster node from a given timestamp onwards. + GetNodeUsageSince(ctx context.Context, nodeID string, since time.Time) (*NodeUsageAggregate, error) Update(ctx context.Context, id string, entry any) error BulkUpdateCost(ctx context.Context, updates map[string]float64) error Flush(ctx context.Context, since time.Time) error diff --git a/framework/logstore/tables.go b/framework/logstore/tables.go index 7401fc4c951..9dbe074cb0d 100644 --- a/framework/logstore/tables.go +++ b/framework/logstore/tables.go @@ -199,6 +199,12 @@ type Log struct { IsLargePayloadResponse bool `gorm:"default:false" json:"is_large_payload_response"` HasObject bool `gorm:"default:false" json:"-"` // True when payload is stored in object storage + // Cluster governance fields - attached by the logging plugin when running in a cluster + // so that leaders can recover disconnected node usage from the logs table. + ClusterNodeID *string `gorm:"type:varchar(255)" json:"cluster_node_id,omitempty"` + BudgetIDs *string `gorm:"type:text" json:"-"` // JSON serialized []string of budget IDs applicable to this request + RateLimitIDs *string `gorm:"type:text" json:"-"` // JSON serialized []string of rate limit IDs applicable to this request + // Denormalized token fields for easier querying PromptTokens int `gorm:"default:0" json:"-"` CompletionTokens int `gorm:"default:0" json:"-"` @@ -240,6 +246,8 @@ type Log struct { VideoListOutputParsed *schemas.BifrostVideoListResponse `gorm:"-" json:"video_list_output,omitempty"` VideoDeleteOutputParsed *schemas.BifrostVideoDeleteResponse `gorm:"-" json:"video_delete_output,omitempty"` AttemptTrailParsed []schemas.KeyAttemptRecord `gorm:"-" json:"attempt_trail,omitempty"` + BudgetIDsParsed []string `gorm:"-" json:"budget_ids,omitempty"` + RateLimitIDsParsed []string `gorm:"-" json:"rate_limit_ids,omitempty"` // Populated in handlers after find using the virtual key id and key id VirtualKey *tables.TableVirtualKey `gorm:"-" json:"virtual_key,omitempty"` // redacted @@ -549,6 +557,22 @@ func (l *Log) SerializeFields() error { } } + if len(l.BudgetIDsParsed) > 0 { + if data, err := sonic.Marshal(l.BudgetIDsParsed); err != nil { + return err + } else { + l.BudgetIDs = new(string(data)) + } + } + + if len(l.RateLimitIDsParsed) > 0 { + if data, err := sonic.Marshal(l.RateLimitIDsParsed); err != nil { + return err + } else { + l.RateLimitIDs = new(string(data)) + } + } + // Build content summary for search. // Skip if already set (e.g., by the hybrid log store which builds input-only summaries). if l.ContentSummary == "" { @@ -773,6 +797,18 @@ func (l *Log) DeserializeFields() error { } } + if l.BudgetIDs != nil && *l.BudgetIDs != "" { + if err := sonic.Unmarshal([]byte(*l.BudgetIDs), &l.BudgetIDsParsed); err != nil { + l.BudgetIDsParsed = nil + } + } + + if l.RateLimitIDs != nil && *l.RateLimitIDs != "" { + if err := sonic.Unmarshal([]byte(*l.RateLimitIDs), &l.RateLimitIDsParsed); err != nil { + l.RateLimitIDsParsed = nil + } + } + if l.RoutingEnginesUsedStr != nil && *l.RoutingEnginesUsedStr != "" { // Parse comma-separated routing engines l.RoutingEnginesUsed = strings.Split(*l.RoutingEnginesUsedStr, ",") @@ -1503,3 +1539,13 @@ type UserRankingWithTrend struct { type UserRankingResult struct { Rankings []UserRankingWithTrend `json:"rankings"` } + +// NodeUsageAggregate represents aggregated usage for a specific node from the logs table, +// broken down by the budget and rate-limit IDs that each log entry was tagged with. +// This ensures usage is attributed to the correct governance resource rather than +// spread uniformly across all resources the node was tracking. +type NodeUsageAggregate struct { + BudgetCosts map[string]float64 `json:"budget_costs"` // budget_id -> cumulative cost + RateLimitRequests map[string]int64 `json:"rate_limit_requests"` // rate_limit_id -> successful request count + RateLimitTokens map[string]int64 `json:"rate_limit_tokens"` // rate_limit_id -> total tokens +} diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 02ecc07aef8..5bf138b2d39 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -292,6 +292,7 @@ type LoggerPlugin struct { writeQueue chan *writeQueueEntry // Buffered channel for batch write queue closed atomic.Bool // Set during cleanup to prevent sends on closed writeQueue deferredUsageSem chan struct{} // Limits concurrent deferred usage DB updates + clusterNodeID atomic.Value // Cluster node ID (string) for log attribution in clustered deployments } // Init creates new logger plugin with given log store @@ -350,6 +351,14 @@ func Init(ctx context.Context, config *Config, logger schemas.Logger, logsStore return plugin, nil } +// SetClusterNodeID sets the cluster node ID that will be attached to all log entries. +// Used in clustered deployments to attribute log entries to specific nodes for +// disconnected node usage recovery. Uses atomic.Value since it is written at +// startup and read concurrently from request hot paths. +func (p *LoggerPlugin) SetClusterNodeID(nodeID string) { + p.clusterNodeID.Store(nodeID) +} + // cleanupWorker periodically removes old processing logs func (p *LoggerPlugin) cleanupWorker() { defer p.wg.Done() @@ -788,6 +797,9 @@ func (p *LoggerPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas. entry.ErrorDetails = string(data) } entry.ErrorDetailsParsed = bifrostErr + if nodeID, _ := p.clusterNodeID.Load().(string); nodeID != "" { + entry.ClusterNodeID = &nodeID + } applyLargePayloadPreviewsToEntry(ctx, entry, contentLoggingEnabled) p.storeOrEnqueueEntry(ctx, entry, p.makePostWriteCallback(nil)) } else { @@ -854,6 +866,16 @@ func (p *LoggerPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas. } } applyOutputFieldsToEntry(entry, selectedKeyID, selectedKeyName, virtualKeyID, virtualKeyName, routingRuleID, routingRuleName, selectedPromptID, selectedPromptName, selectedPromptVersion, teamID, teamName, customerID, customerName, userID, userName, businessUnitID, businessUnitName, numberOfRetries, latency, attemptTrail) + // Attach cluster governance metadata for disconnected node usage recovery + if nodeID, _ := p.clusterNodeID.Load().(string); nodeID != "" { + entry.ClusterNodeID = &nodeID + } + if budgetIDs, ok := ctx.Value(schemas.BifrostContextKeyGovernanceBudgetIDs).([]string); ok && len(budgetIDs) > 0 { + entry.BudgetIDsParsed = budgetIDs + } + if rateLimitIDs, ok := ctx.Value(schemas.BifrostContextKeyGovernanceRateLimitIDs).([]string); ok && len(rateLimitIDs) > 0 { + entry.RateLimitIDsParsed = rateLimitIDs + } entry.MetadataParsed = pending.InitialData.Metadata entry.MetadataParsed = mergeRealtimeMetadata(entry.MetadataParsed, ctx) entry.RoutingEngineLogs = routingEngineLogs