feat: add connector test cases and missing attributes

core/schemas/enrichment.go

@@ -0,0 +1,140 @@
+package schemas
+
+// EnrichmentDim describes one identity/context dimension that connectors attach
+// to the telemetry they emit for a request (which team/customer/business unit/
+// virtual key/etc. the request belongs to).
+//
+// It is the single source of truth that keeps the CURATED emitters from drifting
+// apart — the ones that hand-pick a dimension list:
+//   - Prometheus labels   (OSS plugins/telemetry)
+//   - Datadog metric tags (enterprise plugins/datadog, buildMetricTags)
+//   - BigQuery columns    (enterprise plugins/bigquery, traceColumns)
+//
+// Each of those derives its list from this registry, and a per-connector
+// conformance test asserts the derived list matches — so adding a dimension in
+// one place can't silently leave the others behind.
+//
+// The GENERIC emitters (otel, kafka, pubsub) project the entire span attribute
+// map and therefore already carry every dimension; they need no derivation and
+// no conformance test here.
+//
+// `alias` and `routing_engine_used` are in the metric tier but derived only
+// post-response (once model resolution/routing has run). They are attached to
+// the span in framework/tracing and carry a normal (non-empty) SpanAttr, so a
+// record/trace-tier connector can read them like any other dimension.
+type EnrichmentDim struct {
+	// Name is the canonical short identifier, used verbatim as the Prometheus
+	// label and the Datadog metric tag key. The BigQuery column name also equals
+	// it unless Column overrides (see below).
+	Name string
+	// Column is the BigQuery column name when it differs from Name. BigQuery
+	// predates the "method" naming and stores it as "request_type"; empty means
+	// the column name equals Name.
+	Column string
+	// SpanAttr is the canonical bifrost.* span-attribute key the dimension is
+	// stored under. It is what the record/trace-tier emitters read and what a
+	// connector derives its projection from.
+	SpanAttr string
+	// MetricSafe marks a LOW-cardinality dimension eligible to become a Prometheus
+	// label / Datadog metric tag. High-cardinality dims (per-user, arrays) are
+	// false and live only on records/traces (BigQuery columns, span attributes).
+	MetricSafe bool
+	// Multi marks an array-valued dimension — governance can attach several teams/
+	// customers/business units to a single request. Array dims are never
+	// MetricSafe (they would explode metric series cardinality).
+	Multi bool
+}
+
+// EnrichmentDims is the canonical, ordered registry of identity/context
+// dimensions. Order is stable so derived lists (labels/tags/columns) are
+// deterministic. Add a dimension here once and every curated connector picks it
+// up via its derivation + conformance test.
+var EnrichmentDims = []EnrichmentDim{
+	// --- Metric tier: low-cardinality, safe as Prometheus labels / Datadog tags,
+	//     and also present on records/traces. ---
+	{Name: "provider", SpanAttr: AttrBifrostProviderName, MetricSafe: true},
+	{Name: "model", SpanAttr: AttrRequestModel, MetricSafe: true},
+	{Name: "method", Column: "request_type", SpanAttr: AttrLegacyRequestType, MetricSafe: true},
+	// alias and routing_engine_used are derived post-response and attached to the
+	// span in framework/tracing (they have no meaning until the model is resolved
+	// and routing has run), so connectors read them like any other dimension.
+	{Name: "alias", SpanAttr: AttrBifrostAlias, MetricSafe: true},
+	{Name: "routing_engine_used", SpanAttr: AttrBifrostRoutingEngineUsed, MetricSafe: true},
+	{Name: "virtual_key_id", SpanAttr: AttrBifrostVirtualKeyID, MetricSafe: true},
+	{Name: "virtual_key_name", SpanAttr: AttrBifrostVirtualKeyName, MetricSafe: true},
+	{Name: "selected_key_id", SpanAttr: AttrBifrostSelectedKeyID, MetricSafe: true},
+	{Name: "selected_key_name", SpanAttr: AttrBifrostSelectedKeyName, MetricSafe: true},
+	{Name: "routing_rule_id", SpanAttr: AttrBifrostRoutingRuleID, MetricSafe: true},
+	{Name: "routing_rule_name", SpanAttr: AttrBifrostRoutingRuleName, MetricSafe: true},
+	{Name: "team_id", SpanAttr: AttrBifrostTeamID, MetricSafe: true},
+	{Name: "team_name", SpanAttr: AttrBifrostTeamName, MetricSafe: true},
+	{Name: "customer_id", SpanAttr: AttrBifrostCustomerID, MetricSafe: true},
+	{Name: "customer_name", SpanAttr: AttrBifrostCustomerName, MetricSafe: true},
+	{Name: "business_unit_id", SpanAttr: AttrBifrostBusinessUnitID, MetricSafe: true},
+	{Name: "business_unit_name", SpanAttr: AttrBifrostBusinessUnitName, MetricSafe: true},
+	{Name: "fallback_index", SpanAttr: AttrBifrostFallbackIndex, MetricSafe: true},
+
+	// --- Record/trace tier only: high cardinality, NOT metric-safe. Present on
+	//     BigQuery columns and span attributes, never as metric labels/tags. ---
+	{Name: "user_id", SpanAttr: AttrBifrostUserID},
+	{Name: "user_name", SpanAttr: AttrBifrostUserName},
+	{Name: "team_ids", SpanAttr: AttrBifrostTeamIDs, Multi: true},
+	{Name: "team_names", SpanAttr: AttrBifrostTeamNames, Multi: true},
+	{Name: "customer_ids", SpanAttr: AttrBifrostCustomerIDs, Multi: true},
+	{Name: "customer_names", SpanAttr: AttrBifrostCustomerNames, Multi: true},
+	{Name: "business_unit_ids", SpanAttr: AttrBifrostBusinessUnitIDs, Multi: true},
+	{Name: "business_unit_names", SpanAttr: AttrBifrostBusinessUnitNames, Multi: true},
+}
+
+// ColumnName returns the BigQuery column name for the dimension — Column when
+// set, otherwise Name.
+func (d EnrichmentDim) ColumnName() string {
+	if d.Column != "" {
+		return d.Column
+	}
+	return d.Name
+}
+
+// EnrichmentDimColumnNames returns the BigQuery column name for every dimension,
+// in registry order (the record/trace-tier set).
+func EnrichmentDimColumnNames() []string {
+	out := make([]string, len(EnrichmentDims))
+	for i, d := range EnrichmentDims {
+		out[i] = d.ColumnName()
+	}
+	return out
+}
+
+// MetricSafeEnrichmentDims returns the low-cardinality dimensions eligible to be
+// Prometheus labels / Datadog metric tags, in registry order.
+func MetricSafeEnrichmentDims() []EnrichmentDim {
+	out := make([]EnrichmentDim, 0, len(EnrichmentDims))
+	for _, d := range EnrichmentDims {
+		if d.MetricSafe {
+			out = append(out, d)
+		}
+	}
+	return out
+}
+
+// EnrichmentDimNames returns every dimension name, in registry order (the
+// record/trace-tier set).
+func EnrichmentDimNames() []string {
+	out := make([]string, len(EnrichmentDims))
+	for i, d := range EnrichmentDims {
+		out[i] = d.Name
+	}
+	return out
+}
+
+// MetricSafeEnrichmentDimNames returns the names of the metric-tier dimensions,
+// in registry order.
+func MetricSafeEnrichmentDimNames() []string {
+	out := make([]string, 0, len(EnrichmentDims))
+	for _, d := range EnrichmentDims {
+		if d.MetricSafe {
+			out = append(out, d.Name)
+		}
+	}
+	return out
+}

core/schemas/enrichment_test.go

@@ -0,0 +1,30 @@
+package schemas
+
+import "testing"
+
+// TestArrayDimsAreNeverMetricSafe is the structural guard that keeps array
+// (Multi) dimensions out of the metric tier — i.e. out of Prometheus labels and
+// Datadog metric tags. An array value like "team-a,team-b,team-c" would become a
+// distinct label/tag value per team combination and explode series cardinality,
+// so a Multi dimension must never be MetricSafe. The curated connectors derive
+// their metric-tier lists from MetricSafeEnrichmentDims(), so this invariant is
+// what actually prevents arrays from ever being shared as Prometheus labels.
+func TestArrayDimsAreNeverMetricSafe(t *testing.T) {
+	for _, d := range EnrichmentDims {
+		if d.Multi && d.MetricSafe {
+			t.Errorf("dimension %q is Multi (array) AND MetricSafe — arrays must never be metric labels/tags (cardinality explosion)", d.Name)
+		}
+	}
+}
+
+// TestEnrichmentDimNamesUnique guards against a copy-paste duplicate slipping into
+// the registry, which would double-emit a label/column.
+func TestEnrichmentDimNamesUnique(t *testing.T) {
+	seen := map[string]bool{}
+	for _, d := range EnrichmentDims {
+		if seen[d.Name] {
+			t.Errorf("duplicate enrichment dimension name %q", d.Name)
+		}
+		seen[d.Name] = true
+	}
+}

core/schemas/trace.go

@@ -708,6 +708,8 @@ const (
 	AttrBifrostUserName            = "bifrost.user.name"
 	AttrBifrostRetries             = "bifrost.retries"
 	AttrBifrostFallbackIndex       = "bifrost.fallback_index"
+	AttrBifrostAlias               = "bifrost.alias"                // original requested model when it differs from the resolved model
+	AttrBifrostRoutingEngineUsed   = "bifrost.routing_engine_used"  // comma-joined routing engines that handled the request
 	AttrBifrostStopSequencesJoined = "bifrost.request.stop_sequences"
 
 	// OTel general semconv (no gen_ai prefix). Emitted alongside the legacy

framework/tracing/tracer.go

@@ -336,6 +336,23 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand
 	for k, v := range PopulateErrorAttributes(err) {
 		span.SetAttribute(k, v)
 	}
+
+	// Enrichment dimensions derivable only post-response, attached here so every
+	// connector reads them from one place (see core/schemas EnrichmentDims):
+	//   - alias: the originally requested model when it differs from the resolved
+	//     model (an alias was matched or a fallback swapped the model).
+	//   - routing_engine_used: the comma-joined set of routing engines that handled
+	//     the request; the context list is only complete once routing has run.
+	if resp != nil {
+		ef := resp.GetExtraFields()
+		if ef.ResolvedModelUsed != "" && ef.ResolvedModelUsed != ef.OriginalModelRequested && ef.OriginalModelRequested != "" {
+			span.SetAttribute(schemas.AttrBifrostAlias, ef.OriginalModelRequested)
+		}
+	}
+	if engines, ok := ctx.Value(schemas.BifrostContextKeyRoutingEnginesUsed).([]string); ok && len(engines) > 0 {
+		span.SetAttribute(schemas.AttrBifrostRoutingEngineUsed, strings.Join(engines, ","))
+	}
+
 	// Populate cost attribute using pricing manager
 	if t.pricingManager != nil && resp != nil {
 		cost := t.pricingManager.CalculateCost(resp, modelcatalog.PricingLookupScopesFromContext(ctx, string(resp.GetExtraFields().Provider)))

plugins/telemetry/main.go

@@ -227,6 +227,31 @@ var (
 )
 
 // Init creates a new PrometheusPlugin with initialized metrics.
+// defaultBifrostLabelNames is the canonical set of Prometheus labels attached to
+// bifrost.* metrics. It is a package var (not an Init local) so the connector-
+// parity conformance test can assert it against the shared enrichment registry
+// (core/schemas). Metric-tier dimensions only — no high-cardinality (user, arrays).
+var defaultBifrostLabelNames = []string{
+	"provider",
+	"model",
+	"alias",
+	"method",
+	"virtual_key_id",
+	"virtual_key_name",
+	"routing_engine_used",
+	"routing_rule_id",
+	"routing_rule_name",
+	"selected_key_id",
+	"selected_key_name",
+	"fallback_index",
+	"team_id",
+	"team_name",
+	"customer_id",
+	"customer_name",
+	"business_unit_id",
+	"business_unit_name",
+}
+
 func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger schemas.Logger) (*PrometheusPlugin, error) {
 	if config == nil {
 		return nil, fmt.Errorf("config is required")
@@ -257,24 +282,7 @@ func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger sche
 	}
 
 	defaultHTTPLabels := []string{"path", "method", "status"}
-	defaultBifrostLabels := []string{
-		"provider",
-		"model",
-		"alias",
-		"method",
-		"virtual_key_id",
-		"virtual_key_name",
-		"routing_engine_used",
-		"routing_rule_id",
-		"routing_rule_name",
-		"selected_key_id",
-		"selected_key_name",
-		"fallback_index",
-		"team_id",
-		"team_name",
-		"customer_id",
-		"customer_name",
-	}
+	defaultBifrostLabels := append([]string(nil), defaultBifrostLabelNames...)
 
 	var filteredCustomLabels []string
 	if len(config.CustomLabels) > 0 {
@@ -707,6 +715,8 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche
 	teamName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceTeamName)
 	customerID := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceCustomerID)
 	customerName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceCustomerName)
+	businessUnitID := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceBusinessUnitID)
+	businessUnitName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceBusinessUnitName)
 
 	// Extract ALL context values BEFORE spawning the goroutine.
 	labelValues := map[string]string{
@@ -726,6 +736,8 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche
 		"team_name":           teamName,
 		"customer_id":         customerID,
 		"customer_name":       customerName,
+		"business_unit_id":    businessUnitID,
+		"business_unit_name":  businessUnitName,
 	}
 
 	// Get all custom prometheus labels from context BEFORE the goroutine.

plugins/telemetry/parity_test.go

@@ -0,0 +1,54 @@
+package telemetry
+
+import (
+	"testing"
+
+	"github.com/maximhq/bifrost/core/schemas"
+)
+
+// TestPrometheusLabelsMatchEnrichmentRegistry keeps the Prometheus bifrost label
+// set in parity with the canonical enrichment registry (core/schemas): every
+// metric-safe dimension must be a label, and no record-tier (high cardinality)
+// dimension may be. Known divergences are enumerated in the allowlist below; the
+// test fails on any new drift and when a listed entry no longer applies.
+func TestPrometheusLabelsMatchEnrichmentRegistry(t *testing.T) {
+	labels := map[string]bool{}
+	for _, l := range defaultBifrostLabelNames {
+		labels[l] = true
+	}
+
+	metricSafe := map[string]bool{}
+	for _, n := range schemas.MetricSafeEnrichmentDimNames() {
+		metricSafe[n] = true
+	}
+	allDims := map[string]bool{}
+	for _, n := range schemas.EnrichmentDimNames() {
+		allDims[n] = true
+	}
+
+	// Metric-safe dims the telemetry plugin does not expose as labels. Empty: the
+	// Prometheus label set covers the full metric-safe registry.
+	knownMissing := map[string]string{}
+
+	// 1) Every metric-safe dim must be a label, unless a known gap.
+	for n := range metricSafe {
+		if !labels[n] && knownMissing[n] == "" {
+			t.Errorf("Prometheus labels are missing registry metric-safe dimension %q", n)
+		}
+	}
+	// 2) Every enrichment-dim label must be metric-safe (no high-cardinality labels).
+	for l := range labels {
+		if !allDims[l] {
+			continue // non-enrichment label (e.g. status_code) — out of scope
+		}
+		if !metricSafe[l] {
+			t.Errorf("Prometheus exposes record-tier dimension %q as a label (cardinality risk)", l)
+		}
+	}
+	// 3) Keep the allowlist honest: a gap that has closed must be removed.
+	for n := range knownMissing {
+		if labels[n] {
+			t.Errorf("dimension %q is now a label — remove it from knownMissing", n)
+		}
+	}
+}