diff --git a/internal/proxy/credential_key_parts.go b/internal/proxy/credential_key_parts.go index 1746d07b4..539f2617f 100644 --- a/internal/proxy/credential_key_parts.go +++ b/internal/proxy/credential_key_parts.go @@ -10,6 +10,9 @@ import ( // resolved upstream credential. Empty values mean the turn ran on the deployment // key or no credential was resolved. func (s *Service) credentialKeyParts(ctx context.Context) (prefix, suffix, source string) { + if managedSource := managedSubscriptionCredentialSource(ctx); managedSource != "" { + return "", "", managedSource + } creds := CredentialsFromContext(ctx) if creds == nil || len(creds.APIKey) == 0 { return "", "", "" diff --git a/internal/proxy/credential_key_parts_internal_test.go b/internal/proxy/credential_key_parts_internal_test.go index 8f1176844..fdf3a6224 100644 --- a/internal/proxy/credential_key_parts_internal_test.go +++ b/internal/proxy/credential_key_parts_internal_test.go @@ -56,6 +56,20 @@ func TestCredentialKeyParts_ShortKey(t *testing.T) { assert.Equal(t, credSourceClient, src) } +func TestCredentialKeyParts_ManagedSubscriptionSourceOutranksOuterCredential(t *testing.T) { + s := &Service{} + ctx := WithManagedSubscriptionUsage(ctxWithCreds(&Credentials{APIKey: []byte("byok-token"), Source: credSourceBYOK})) + managedCtx := context.WithValue(ctx, CredentialsContextKey{}, &Credentials{ + APIKey: []byte("managed-token"), Source: credSourceSubscription, OAuth: true, + }) + markManagedSubscriptionServed(ctx, managedCtx) + + prefix, suffix, source := s.credentialKeyParts(ctx) + assert.Empty(t, prefix, "managed access tokens must not be copied to outer telemetry") + assert.Empty(t, suffix, "managed access tokens must not be copied to outer telemetry") + assert.Equal(t, credSourceSubscription, source) +} + // These string values are a wire contract with the SQL export query; changing them breaks subscription_served. func TestCredentialSources_OAuthValuesMatchExportContract(t *testing.T) { assert.Equal(t, "subscription", credSourceSubscription) diff --git a/internal/proxy/fallback.go b/internal/proxy/fallback.go index 3439d1ab0..dacd0941e 100644 --- a/internal/proxy/fallback.go +++ b/internal/proxy/fallback.go @@ -254,7 +254,7 @@ func (s *Service) dispatchWithFallback(ctx context.Context, in failoverInputs) ( lease.Release() if attemptErr == nil { if managedAttempt { - markManagedSubscriptionServed(ctx) + markManagedSubscriptionServed(ctx, credentialCtx) } if i > 0 { log.Info("dispatchWithFallback: succeeded on fallback", diff --git a/internal/proxy/managed_subscriptions.go b/internal/proxy/managed_subscriptions.go index 1903b1989..879a7b332 100644 --- a/internal/proxy/managed_subscriptions.go +++ b/internal/proxy/managed_subscriptions.go @@ -30,7 +30,8 @@ type ManagedSubscriptionUsageContextKey struct{} // ManagedSubscriptionUsage is request-local attribution shared by the auth // middleware context and provider-specific dispatch attempt contexts. type ManagedSubscriptionUsage struct { - Served bool + Served bool + CredentialSource string } var ( @@ -78,11 +79,22 @@ func managedSubscriptionEnrollmentUnavailable(ctx context.Context) bool { return unavailable } -func markManagedSubscriptionServed(ctx context.Context) { +func markManagedSubscriptionServed(ctx context.Context, credentialCtx context.Context) { usage, _ := ctx.Value(ManagedSubscriptionUsageContextKey{}).(*ManagedSubscriptionUsage) if usage != nil { usage.Served = true + if creds := CredentialsFromContext(credentialCtx); creds != nil { + usage.CredentialSource = creds.Source + } + } +} + +func managedSubscriptionCredentialSource(ctx context.Context) string { + usage, _ := ctx.Value(ManagedSubscriptionUsageContextKey{}).(*ManagedSubscriptionUsage) + if usage == nil || !usage.Served { + return "" } + return usage.CredentialSource } func managedSubscriptionServed(ctx context.Context) bool { diff --git a/internal/proxy/managed_subscriptions_test.go b/internal/proxy/managed_subscriptions_test.go index 226f060b7..70a1a51d7 100644 --- a/internal/proxy/managed_subscriptions_test.go +++ b/internal/proxy/managed_subscriptions_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "weave-os/router/internal/auth" @@ -66,8 +67,9 @@ func TestDispatchWithFallbackUsesOnlyMatchingManagedProviderFamily(t *testing.T) recorder := httptest.NewRecorder() buffer := newPreludeBuffer(recorder) request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + ctx := managedSubscriptionContext(auth.SubscriptionProviderCodex) - _, err := svc.dispatchWithFallback(managedSubscriptionContext(auth.SubscriptionProviderCodex), failoverInputs{ + _, err := svc.dispatchWithFallback(ctx, failoverInputs{ w: recorder, buf: buffer, initialDecision: router.Decision{Model: "gpt-5.6-sol", Provider: providers.ProviderOpenAI}, bindings: []catalog.ProviderBinding{{Provider: providers.ProviderOpenAI}}, @@ -81,6 +83,7 @@ func TestDispatchWithFallbackUsesOnlyMatchingManagedProviderFamily(t *testing.T) require.NoError(t, err) require.Equal(t, []subscriptions.Provider{subscriptions.ProviderCodex}, leaser.providers) require.Equal(t, "served", recorder.Body.String()) + assert.Equal(t, credSourceCodexSubscription, managedSubscriptionCredentialSource(ctx)) } func TestDispatchWithFallbackDoesNotCrossManagedProviderFamilies(t *testing.T) {