feat: add cluster-aware log metadata and per-node usage aggregation

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

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)
 }

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
+}

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) {

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

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
+}

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