Derive LLM billing mode from provider selection (SSOT) — internal#711/#713 #1966

Closed
hongming wants to merge 1 commits from fix/ssot-provider-selection-billing-mode-711-713 into main
7 changed files with 543 additions and 12 deletions
@@ -268,6 +268,9 @@ func TestExtended_SecretsSet(t *testing.T) {
// default (byok, set via t.Setenv above) → bypass-list check is skipped
// and the write proceeds. This pattern is the test-side mirror of the
// real-prod fall-through behavior for a fresh workspace with no override.
// internal#711/#713: the resolver reads the LLM_PROVIDER selection first;
// none set here → ErrNoRows → falls through to the llm_billing_mode chain.
expectNoProviderSelection(mock, "22222222-2222-2222-2222-222222222222")
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs("22222222-2222-2222-2222-222222222222").
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
@@ -23,6 +23,47 @@ package handlers
// ?? organizations.llm_billing_mode (org default, fetched via tenant_config)
// ?? "platform_managed" (closed default — the existing implicit default)
//
// internal#711/#713 — PROVIDER SELECTION IS THE SSOT.
//
// The chain above had a latent SSOT defect (confirmed live 2026-05-27):
// workspaces.llm_billing_mode is NULLABLE, has no default, and is NEVER
// written by any create/provider-selection flow. The org default is now
// hardcoded to platform_managed (CP tenant_config.llmBillingEnv removed the
// org-tier policy in the internal#691 follow-up). So the chain ALWAYS
// resolved to platform_managed — selecting a non-Platform provider on the
// canvas had ZERO effect on which credential a workspace used, silently
// running BYOK agents on the platform's Anthropic credits (Reno Stars SEO
// 352e3c2b: MOLECULE_LLM_BILLING_MODE=platform_managed while LLM_PROVIDER /
// MODEL_PROVIDER / PROVIDER were all empty).
//
// The fix makes the PRODUCT'S PROVIDER SELECTION the single source of truth
// for the platform-vs-byok decision, per the CTO's intent verbatim: "if in
// the config the provider is not platform, it's byok." The authoritative
// store for that selection is the `LLM_PROVIDER` row in workspace_secrets —
// written by setProviderSecret (canvas PUT /workspaces/:id/provider, derived
// from the model slug at create, propagated into /configs/config.yaml by CP
// user-data). There is NO new parallel field; we read the one that already
// holds the canvas picker's value.
//
// New resolution order (highest precedence first):
//
// LLM_PROVIDER workspace_secret (the canvas "Provider:" picker)
// - present & NOT a platform sentinel ("platform"/"platform_managed")
// → "byok" (source = provider_selection)
// - present & a platform sentinel
// → "platform_managed" (source = provider_selection)
// - absent / empty
// → fall through to the legacy chain below
// ?? workspaces.llm_billing_mode (per-workspace override, NULLABLE)
// ?? org default (platform_managed — hardcoded by CP today)
// ?? "platform_managed" (closed default — the existing implicit default)
//
// The provider selection wins because it is the only signal a human ever
// actually sets; the llm_billing_mode column stays as a defense-in-depth
// override/escape hatch but no longer the primary signal. Default stays
// platform_managed when NOTHING is selected (no provider, no override),
// preserving the existing implicit default and the #711 fail-closed.
//
// Default-closed contract — non-negotiable per the RFC Safety axis:
//
// - workspace row missing (sql.ErrNoRows) → fall through to org default
@@ -43,7 +84,9 @@ import (
"database/sql"
"errors"
"fmt"
"strings"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/crypto"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
)
@@ -58,12 +101,57 @@ const (
LLMBillingModeDisabled = "disabled"
)
// ProviderSelectionSecretKey is the workspace_secrets key that holds the
// product's "Provider:" picker value — the SINGLE SOURCE OF TRUTH for the
// platform-vs-byok decision (internal#711/#713). Written by setProviderSecret
// (canvas PUT /workspaces/:id/provider; derived from the model slug at create
// in WorkspaceHandler.Create) and surfaced by GetProvider. The value is a
// provider slug ("anthropic", "minimax", "openrouter", …) or the platform
// sentinel ("platform"); absent = no explicit selection.
const ProviderSelectionSecretKey = "LLM_PROVIDER"
// platformProviderSentinels are the LLM_PROVIDER values that mean "use the
// platform's managed LLM" rather than a specific vendor the customer brings.
// Templates declare a Platform-managed model with `provider: platform`
// (see modelSpec in templates.go + TestTemplatesList_SurfacesProviders), so
// "platform" is the canonical sentinel; "platform_managed" is accepted as a
// defensive alias in case a future caller writes the billing-mode literal
// into the provider slot. Matched case-insensitively after trimming.
var platformProviderSentinels = map[string]struct{}{
"platform": {},
"platform_managed": {},
}
// providerSelectionIsPlatform reports whether a (trimmed, lowercased) provider
// slug denotes the platform-managed provider. An empty slug is NOT platform —
// it is "no selection", which the resolver treats as undecided and falls
// through to the legacy chain (preserving the platform_managed default only
// when nothing else decides). Returning (isPlatform, decided):
//
// decided=false → no provider selected; caller falls through
// decided=true, isPlatform=true → platform_managed
// decided=true, isPlatform=false → byok (a specific vendor was picked)
func providerSelectionIsPlatform(rawProvider string) (isPlatform bool, decided bool) {
p := strings.ToLower(strings.TrimSpace(rawProvider))
if p == "" {
return false, false
}
if _, ok := platformProviderSentinels[p]; ok {
return true, true
}
return false, true
}
// BillingModeSource describes which layer of the resolution stack supplied
// the final mode. Surfaced via the admin route for operator debug
// ("why is this workspace being stripped?") per the RFC Observability axis.
type BillingModeSource string
const (
// BillingModeSourceProviderSelection means the mode was decided by the
// workspace's LLM_PROVIDER selection — the SSOT (internal#711/#713). This
// is the highest-precedence source and the one a human actually sets.
BillingModeSourceProviderSelection BillingModeSource = "provider_selection"
BillingModeSourceWorkspaceOverride BillingModeSource = "workspace_override"
BillingModeSourceOrgDefault BillingModeSource = "org_default"
BillingModeSourceConstantFallback BillingModeSource = "constant_fallback"
@@ -74,11 +162,16 @@ const (
// shape, so the resolver test asserts both the mode AND the source per case
// (catches a bug where the right mode is returned via the wrong layer).
type BillingModeResolution struct {
WorkspaceID string `json:"workspace_id"`
ResolvedMode string `json:"resolved_mode"`
WorkspaceOverride *string `json:"workspace_override"` // nil = inherit
OrgDefault string `json:"org_default"` // already default-closed by CP
Source BillingModeSource `json:"source"`
WorkspaceID string `json:"workspace_id"`
ResolvedMode string `json:"resolved_mode"`
WorkspaceOverride *string `json:"workspace_override"` // nil = inherit
OrgDefault string `json:"org_default"` // already default-closed by CP
Source BillingModeSource `json:"source"`
// ProviderSelection surfaces the workspace's LLM_PROVIDER value (the SSOT
// picker) when one is set — nil = no explicit selection. Lets the admin
// route answer "why is this workspace byok?" with the literal provider
// slug the customer chose, not just the derived mode.
ProviderSelection *string `json:"provider_selection"`
}
// isKnownBillingMode is the enum-recognizer for the resolver's default-closed
@@ -108,6 +201,39 @@ func normalizeOrgDefault(orgMode string) string {
return LLMBillingModePlatformManaged
}
// readWorkspaceProviderSelection fetches and decrypts the workspace's
// LLM_PROVIDER selection (the SSOT picker) from workspace_secrets. Returns:
//
// (slug, true, nil) — a provider was selected; slug is the decrypted value
// ("", false, nil) — no LLM_PROVIDER row (sql.ErrNoRows) → undecided
// ("", false, err) — DB or decrypt error → undecided + propagated error
//
// The "undecided" cases (no row / error) MUST NOT decide the mode — the
// resolver falls through to the legacy llm_billing_mode/org chain, which is
// default-closed to platform_managed. That keeps a transient secrets-read
// failure from silently flipping a workspace to byok (which would then strip
// the platform creds it legitimately needs). The error is surfaced for
// observability, not as a decision.
func readWorkspaceProviderSelection(ctx context.Context, workspaceID string) (string, bool, error) {
var enc []byte
var ver int
err := db.DB.QueryRowContext(ctx,
`SELECT encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id = $1 AND key = $2`,
workspaceID, ProviderSelectionSecretKey,
).Scan(&enc, &ver)
switch {
case errors.Is(err, sql.ErrNoRows):
return "", false, nil
case err != nil:
return "", false, fmt.Errorf("read LLM_PROVIDER for %s: %w", workspaceID, err)
}
dec, decErr := crypto.DecryptVersioned(enc, ver)
if decErr != nil {
return "", false, fmt.Errorf("decrypt LLM_PROVIDER for %s: %w", workspaceID, decErr)
}
return string(dec), true, nil
}
// ResolveLLMBillingMode is the canonical resolver. Every code path that
// previously gated on `os.Getenv("MOLECULE_LLM_BILLING_MODE") == "platform_managed"`
// must call this instead and gate on the returned mode. The architectural
@@ -127,7 +253,8 @@ func ResolveLLMBillingMode(ctx context.Context, workspaceID, orgMode string) (Bi
if workspaceID == "" {
// No workspace ID = pre-provision context (templating, validation).
// Resolve against the org default only, no DB read.
// Resolve against the org default only, no DB read. There is no
// provider selection to consult either (no workspace row yet).
res.ResolvedMode = res.OrgDefault
res.Source = BillingModeSourceOrgDefault
if !isKnownBillingMode(orgMode) {
@@ -139,6 +266,36 @@ func ResolveLLMBillingMode(ctx context.Context, workspaceID, orgMode string) (Bi
return res, nil
}
// SSOT (internal#711/#713): the workspace's LLM_PROVIDER selection is the
// primary signal. Read it FIRST and let it decide platform-vs-byok before
// consulting the never-written llm_billing_mode column. "if in the config
// the provider is not platform, it's byok." A non-Platform provider ⇒ byok;
// the platform sentinel ⇒ platform_managed. Only an absent selection (or a
// read/decrypt error) falls through to the legacy override/org chain below.
if provider, ok, provErr := readWorkspaceProviderSelection(ctx, workspaceID); provErr != nil {
// Surface the read error but DO NOT let it decide — fall through to the
// default-closed chain. (Propagated to the caller for observability.)
res.ResolvedMode = LLMBillingModePlatformManaged
res.Source = BillingModeSourceConstantFallback
return res, provErr
} else if ok {
sel := provider
res.ProviderSelection = &sel
if isPlatform, decided := providerSelectionIsPlatform(provider); decided {
res.Source = BillingModeSourceProviderSelection
if isPlatform {
res.ResolvedMode = LLMBillingModePlatformManaged
} else {
// A specific vendor was picked → bring-your-own-key.
res.ResolvedMode = LLMBillingModeBYOK
}
return res, nil
}
// provider row present but blank string (e.g. " ") — treat as no
// selection and fall through. Keep the surfaced ProviderSelection so
// operators can see the empty value that was stored.
}
var wsOverride sql.NullString
err := db.DB.QueryRowContext(ctx,
`SELECT llm_billing_mode FROM workspaces WHERE id = $1`,
@@ -32,6 +32,8 @@ const testWSID = "44444444-4444-4444-4444-444444444444"
func TestGetWorkspaceLLMBillingMode_HappyPath_InheritsOrgDefault(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModeBYOK)
mock := setupTestDB(t)
// No provider selected (internal#711/#713 SSOT read) → fall through.
expectNoProviderSelection(mock, testWSID)
// Workspace has no override → resolver returns org_default = byok.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
@@ -80,7 +82,9 @@ func TestPutWorkspaceLLMBillingMode_SetByok(t *testing.T) {
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = \$1 WHERE id = \$2`).
WithArgs(LLMBillingModeBYOK, testWSID).
WillReturnResult(sqlmock.NewResult(0, 1))
// Readback after write.
// Readback after write: provider read (no selection) → falls through to
// the workspace override we just wrote.
expectNoProviderSelection(mock, testWSID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -117,6 +121,8 @@ func TestPutWorkspaceLLMBillingMode_ExplicitNullClearsOverride(t *testing.T) {
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = NULL WHERE id = \$1`).
WithArgs(testWSID).
WillReturnResult(sqlmock.NewResult(0, 1))
// Readback: provider read (no selection) → falls through to org default.
expectNoProviderSelection(mock, testWSID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
@@ -11,9 +11,31 @@ import (
"errors"
"testing"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/crypto"
"github.com/DATA-DOG/go-sqlmock"
)
// expectNoProviderSelection mocks the SSOT LLM_PROVIDER read (internal#711/#713)
// returning sql.ErrNoRows — i.e. "no provider was selected" — so the resolver
// falls through to the legacy llm_billing_mode/org chain that the case under
// test is exercising. Resolver order: provider read FIRST, then workspaces row.
func expectNoProviderSelection(m sqlmock.Sqlmock, wsID string) {
m.ExpectQuery(`SELECT encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id = \$1 AND key = \$2`).
WithArgs(wsID, ProviderSelectionSecretKey).
WillReturnRows(sqlmock.NewRows([]string{"encrypted_value", "encryption_version"}))
}
// expectProviderSelection mocks the SSOT LLM_PROVIDER read returning the given
// provider slug as a plaintext-versioned (version 0) workspace_secrets row.
// crypto.DecryptVersioned treats version 0 as plaintext, so no encryption key
// is needed in unit tests.
func expectProviderSelection(m sqlmock.Sqlmock, wsID, provider string) {
m.ExpectQuery(`SELECT encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id = \$1 AND key = \$2`).
WithArgs(wsID, ProviderSelectionSecretKey).
WillReturnRows(sqlmock.NewRows([]string{"encrypted_value", "encryption_version"}).
AddRow([]byte(provider), crypto.EncryptionVersionPlaintext))
}
func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
ctx := context.Background()
const wsID = "11111111-1111-1111-1111-111111111111"
@@ -44,6 +66,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -55,6 +78,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeDisabled))
@@ -66,6 +90,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModeBYOK,
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
@@ -77,6 +102,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
@@ -93,6 +119,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
// non-PG driver in a test stub), a garbled value MUST NOT
// be honored as if it were valid. This is the default-closed
// safety axis the RFC calls out.
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("byokk"))
@@ -104,6 +131,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: "garbled-or-empty",
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("nonsense"))
@@ -117,6 +145,7 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModeBYOK,
setupMock: func(m sqlmock.Sqlmock) {
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
@@ -142,6 +171,9 @@ func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
workspaceID: wsID,
orgMode: LLMBillingModeBYOK, // org says byok but DB errored — DO NOT honor org
setupMock: func(m sqlmock.Sqlmock) {
// Provider read returns no selection (no LLM_PROVIDER row), then
// the legacy workspaces query errors → default-closed + error.
expectNoProviderSelection(m, wsID)
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnError(errors.New("connection refused"))
@@ -191,9 +223,11 @@ func TestResolveLLMBillingMode_ResolvedModeIsAlwaysValid(t *testing.T) {
ctx := context.Background()
const wsID = "22222222-2222-2222-2222-222222222222"
// Throw a pathological row at the resolver: garbled override + garbled
// org default. Resolved mode must still be a recognized enum.
// Throw a pathological row at the resolver: no provider selection (falls
// through), garbled override + garbled org default. Resolved mode must
// still be a recognized enum.
mock := setupTestDB(t)
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("totally-bogus"))
@@ -259,3 +293,203 @@ func TestSetWorkspaceLLMBillingMode_Validation(t *testing.T) {
}
})
}
// TestResolveLLMBillingMode_ProviderSelectionIsSSOT is the core regression
// guard for internal#711/#713: the workspace's LLM_PROVIDER selection — the
// product "Provider:" picker — is the single source of truth for the
// platform-vs-byok decision. CTO intent verbatim: "if in the config the
// provider is not platform, it's byok." Default stays platform_managed only
// when there is genuinely NO provider selection.
//
// The defect this catches: pre-fix the resolver read only
// workspaces.llm_billing_mode (never written) ?? org default (hardcoded
// platform_managed), so a non-Platform provider selection had ZERO effect
// and a BYOK workspace silently ran on platform creds. Each case below pins
// the provider value → resolved mode → source mapping.
func TestResolveLLMBillingMode_ProviderSelectionIsSSOT(t *testing.T) {
ctx := context.Background()
const wsID = "44444444-4444-4444-4444-444444444444"
type tc struct {
name string
provider string // value stored as the LLM_PROVIDER workspace_secret
wantMode string
wantSource BillingModeSource
wantProvider string // expected res.ProviderSelection (the surfaced slug)
}
cases := []tc{
{
name: "vendor_provider_resolves_byok_even_though_org_is_pm",
provider: "minimax",
wantMode: LLMBillingModeBYOK,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "minimax",
},
{
name: "anthropic_provider_resolves_byok",
provider: "anthropic",
wantMode: LLMBillingModeBYOK,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "anthropic",
},
{
name: "openrouter_provider_resolves_byok",
provider: "openrouter",
wantMode: LLMBillingModeBYOK,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "openrouter",
},
{
name: "platform_sentinel_resolves_platform_managed",
provider: "platform",
wantMode: LLMBillingModePlatformManaged,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "platform",
},
{
name: "platform_managed_alias_resolves_platform_managed",
provider: "platform_managed",
wantMode: LLMBillingModePlatformManaged,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "platform_managed",
},
{
name: "uppercase_platform_is_normalized",
provider: "PLATFORM",
wantMode: LLMBillingModePlatformManaged,
wantSource: BillingModeSourceProviderSelection,
wantProvider: "PLATFORM",
},
{
name: "whitespace_padded_vendor_still_byok",
provider: " zai ",
wantMode: LLMBillingModeBYOK,
wantSource: BillingModeSourceProviderSelection,
wantProvider: " zai ",
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
mock := setupTestDB(t)
// Provider selection is decisive → the legacy workspaces query
// must NOT run. Only the provider read is expected.
expectProviderSelection(mock, wsID, c.provider)
// orgMode = platform_managed (the live, hardcoded value) to prove
// the provider selection overrides it for the byok cases.
res, err := ResolveLLMBillingMode(ctx, wsID, LLMBillingModePlatformManaged)
if err != nil {
t.Fatalf("unexpected err: %v", err)
}
if res.ResolvedMode != c.wantMode {
t.Errorf("mode: got %q want %q", res.ResolvedMode, c.wantMode)
}
if res.Source != c.wantSource {
t.Errorf("source: got %q want %q", res.Source, c.wantSource)
}
if res.ProviderSelection == nil {
t.Fatalf("ProviderSelection: got nil, want %q surfaced", c.wantProvider)
}
if *res.ProviderSelection != c.wantProvider {
t.Errorf("ProviderSelection: got %q want %q", *res.ProviderSelection, c.wantProvider)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
})
}
}
// TestResolveLLMBillingMode_BlankProviderFallsThrough proves a provider row
// that is present but blank (" ") is treated as "no selection": the resolver
// falls through to the legacy chain (here: org default = platform_managed) and
// does NOT flip to byok. This is the no-regression path for a workspace whose
// provider row was cleared to empty — it must keep the default, not silently
// strip platform creds.
func TestResolveLLMBillingMode_BlankProviderFallsThrough(t *testing.T) {
ctx := context.Background()
const wsID = "55555555-5555-5555-5555-555555555555"
mock := setupTestDB(t)
expectProviderSelection(mock, wsID, " ") // blank → undecided
// Falls through to the workspaces row (NULL override) → org default.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
res, err := ResolveLLMBillingMode(ctx, wsID, LLMBillingModePlatformManaged)
if err != nil {
t.Fatalf("unexpected err: %v", err)
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Errorf("blank provider must NOT flip to byok; got mode %q", res.ResolvedMode)
}
if res.Source != BillingModeSourceOrgDefault {
t.Errorf("source: got %q want %q", res.Source, BillingModeSourceOrgDefault)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
// TestResolveLLMBillingMode_ProviderReadErrorIsDefaultClosed proves a DB error
// on the SSOT provider read defaults closed to platform_managed AND propagates
// the error — never silently flips to byok (which would strip the platform
// creds the workspace may legitimately need). Mirrors the workspaces-query
// db-error contract.
func TestResolveLLMBillingMode_ProviderReadErrorIsDefaultClosed(t *testing.T) {
ctx := context.Background()
const wsID = "66666666-6666-6666-6666-666666666666"
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id = \$1 AND key = \$2`).
WithArgs(wsID, ProviderSelectionSecretKey).
WillReturnError(errors.New("connection refused"))
// orgMode=byok to prove the error path does NOT honor org either — it
// defaults closed to platform_managed.
res, err := ResolveLLMBillingMode(ctx, wsID, LLMBillingModeBYOK)
if err == nil {
t.Fatal("expected propagated provider-read error, got nil")
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Errorf("provider-read error must default closed to platform_managed; got %q", res.ResolvedMode)
}
if res.Source != BillingModeSourceConstantFallback {
t.Errorf("source: got %q want %q", res.Source, BillingModeSourceConstantFallback)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
// TestProviderSelectionIsPlatform pins the pure classification helper — the
// decision rule for "is this provider the platform-managed one". Small, fast,
// no DB. This is where the CTO's "not platform ⇒ byok" rule lives literally.
func TestProviderSelectionIsPlatform(t *testing.T) {
cases := []struct {
in string
wantIsPlatorm bool
wantDecided bool
}{
{"platform", true, true},
{"platform_managed", true, true},
{"PLATFORM", true, true},
{" platform ", true, true},
{"minimax", false, true},
{"anthropic", false, true},
{"openrouter", false, true},
{"custom", false, true},
{"", false, false}, // no selection → undecided
{" ", false, false}, // blank → undecided
}
for _, c := range cases {
isPlatform, decided := providerSelectionIsPlatform(c.in)
if isPlatform != c.wantIsPlatorm || decided != c.wantDecided {
t.Errorf("providerSelectionIsPlatform(%q) = (%v,%v) want (%v,%v)",
c.in, isPlatform, decided, c.wantIsPlatorm, c.wantDecided)
}
}
}
@@ -867,7 +867,9 @@ func TestSecretsValues_LegacyWorkspaceGrandfathered(t *testing.T) {
AddRow("WS_KEY", []byte("ws_plainvalue"), 0))
// internal#711: Values now resolves billing mode to gate the global LLM-cred
// merge. Neither key here is a platform-managed LLM bypass key, so the mode
// is immaterial to the assertions — but the resolver query must be mocked.
// is immaterial to the assertions — but the resolver queries must be mocked.
// internal#711/#713: the SSOT provider read runs first (no selection here).
expectNoProviderSelection(mock, testWsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
@@ -951,6 +953,8 @@ func TestSecretsValues_ValidTokenReturnsDecryptedMerge(t *testing.T) {
// internal#711: billing-mode resolver query. None of these keys is a
// platform-managed LLM bypass key, so the resolved mode does not affect the
// merge assertions; platform_managed keeps the existing pass-through.
// internal#711/#713: SSOT provider read runs first (no selection here).
expectNoProviderSelection(mock, testWsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
@@ -1006,7 +1010,9 @@ func TestSecretsValues_ByokStripsGlobalLLMCred(t *testing.T) {
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("ANTHROPIC_API_KEY", []byte("CUSTOMER-OWN-ANTHROPIC-KEY"), 0))
// Resolver: this workspace is byok.
// Resolver: no provider selected (SSOT read falls through) → the workspace
// llm_billing_mode override drives byok.
expectNoProviderSelection(mock, testWsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -520,7 +520,10 @@ func TestPrepareProvisionContext_ByokWithOnlyGlobalOAuthFailsClosed(t *testing.T
mock.ExpectQuery(`SELECT key, encrypted_value, encryption_version FROM workspace_secrets`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}))
// Resolver: workspace override = byok.
// Resolver: no provider selected (SSOT read falls through) → workspace
// override = byok drives the mode. The resolver reads LLM_PROVIDER first
// (internal#711/#713), then the workspaces override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -1171,6 +1174,8 @@ func TestApplyPlatformManagedLLMEnv_NoopsOutsidePlatformManaged(t *testing.T) {
func TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv(t *testing.T) {
const wsID = "77777777-7777-7777-7777-777777777777"
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -1234,6 +1239,8 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv(t *testing
func TestApplyPlatformManagedLLMEnv_ByokStripsGlobalOriginOAuthToken(t *testing.T) {
const wsID = "352e3c2b-0546-4e9c-b487-1e2ff1cf29fc" // Reno Stars SEO agent
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -1284,6 +1291,8 @@ func TestApplyPlatformManagedLLMEnv_ByokStripsGlobalOriginOAuthToken(t *testing.
func TestApplyPlatformManagedLLMEnv_ByokKeepsWorkspaceOwnOAuthEvenWithGlobal(t *testing.T) {
const wsID = "6b66de8d-9337-4fb4-be8d-6d49dca0d809" // Reno Stars Marketing agent
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
@@ -1325,6 +1334,8 @@ func TestApplyPlatformManagedLLMEnv_ByokKeepsWorkspaceOwnOAuthEvenWithGlobal(t *
func TestApplyPlatformManagedLLMEnv_DisabledStripsGlobalButReportsNoCred(t *testing.T) {
const wsID = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeDisabled))
@@ -1360,6 +1371,8 @@ func TestApplyPlatformManagedLLMEnv_DisabledStripsGlobalButReportsNoCred(t *test
func TestApplyPlatformManagedLLMEnv_PlatformManagedStillReceivesGlobalCreds(t *testing.T) {
const wsID = "99999999-9999-9999-9999-999999999999"
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
@@ -1403,6 +1416,8 @@ func TestApplyPlatformManagedLLMEnv_PlatformManagedStillReceivesGlobalCreds(t *t
func TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode(t *testing.T) {
const wsID = "88888888-8888-8888-8888-888888888888"
mock := setupTestDB(t)
// No provider selected (SSOT read) → fall through to the workspace override.
expectNoProviderSelection(mock, wsID)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
@@ -1439,6 +1454,111 @@ func TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode(t *tes
}
}
// TestApplyPlatformManagedLLMEnv_ProviderSelectionDrivesByokStrip is the
// internal#711/#713 SSOT regression guard — the EXACT live Reno Stars leak
// shape. A workspace selected a non-Platform provider on the canvas (its
// LLM_PROVIDER workspace_secret = "minimax") but NO llm_billing_mode override
// was ever written (the column is nullable + never populated; org default is
// hardcoded platform_managed). Pre-fix the resolver saw only
// llm_billing_mode(NULL) → org(platform_managed) and resolved platform_managed,
// so the platform's scope:global CLAUDE_CODE_OAUTH_TOKEN survived and the agent
// ran on Molecule's credits. The fix makes the provider selection the SSOT:
// provider != platform ⇒ byok ⇒ the platform global token is stripped.
//
// Discriminating assertion: NO `SELECT llm_billing_mode` mock is registered —
// the provider read alone is decisive, so the resolver must NOT fall through.
func TestApplyPlatformManagedLLMEnv_ProviderSelectionDrivesByokStrip(t *testing.T) {
const wsID = "352e3c2b-0546-4e9c-b487-1e2ff1cf29fc" // Reno Stars SEO agent
mock := setupTestDB(t)
// SSOT: canvas selected a vendor provider. This alone decides byok — there
// must be NO workspaces.llm_billing_mode query after it.
expectProviderSelection(mock, wsID, "minimax")
// Org/bootstrap floor stays platform_managed (the live, hardcoded value):
// pre-fix this is what wrongly won. The provider selection must override it.
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
// Only the platform's scope:global OAuth token is present (the leak source).
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "PLATFORM-GLOBAL-OAUTH-TOKEN",
"MODEL": "opus",
}
globalKeys := map[string]struct{}{"CLAUDE_CODE_OAUTH_TOKEN": {}}
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "minimax/MiniMax-M2.7")
// 1. Provider selection drove the decision → byok.
if res.ResolvedMode != LLMBillingModeBYOK {
t.Fatalf("ResolvedMode = %q, want %q (provider=minimax is the SSOT; org=platform_managed must NOT win)", res.ResolvedMode, LLMBillingModeBYOK)
}
// 2. Platform global OAuth token STRIPPED — the leak is closed.
if got, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q present — must be stripped once provider selection resolves byok", got)
}
// 3. No CP proxy creds forced.
if got, ok := envVars["ANTHROPIC_API_KEY"]; ok {
t.Fatalf("ANTHROPIC_API_KEY must NOT be injected for byok, got %q", got)
}
// 4. No usable cred (only the stripped platform token) → caller fails closed.
if res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = true, want false (only the stripped platform global token was present)")
}
// 5. Container billing mode reflects the resolved (provider-derived) byok.
if got := envVars["MOLECULE_LLM_BILLING_MODE"]; got != LLMBillingModeBYOK {
t.Fatalf("MOLECULE_LLM_BILLING_MODE = %q, want %q (provider-derived)", got, LLMBillingModeBYOK)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_ProviderPlatformKeepsGlobalCreds is the
// no-regression companion (internal#711/#713): a workspace that explicitly
// selected the Platform provider (LLM_PROVIDER = "platform") must resolve
// platform_managed and STILL receive the platform's creds via the CP proxy.
// Proves the SSOT change does not break the platform-provider path.
//
// Discriminating: NO `SELECT llm_billing_mode` mock — the provider read alone
// decides platform_managed.
func TestApplyPlatformManagedLLMEnv_ProviderPlatformKeepsGlobalCreds(t *testing.T) {
const wsID = "352e3c2b-0546-4e9c-b487-1e2ff1cf29fc"
mock := setupTestDB(t)
expectProviderSelection(mock, wsID, "platform")
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "PLATFORM-GLOBAL-OAUTH-TOKEN",
"MODEL": "opus",
}
globalKeys := map[string]struct{}{"CLAUDE_CODE_OAUTH_TOKEN": {}}
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "")
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Fatalf("ResolvedMode = %q, want %q (provider=platform)", res.ResolvedMode, LLMBillingModePlatformManaged)
}
// Platform-managed routes through the CP proxy: OAuth replaced by proxy token.
if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped + replaced by the proxy token for platform_managed")
}
if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" {
t.Fatalf("ANTHROPIC_API_KEY = %q, want proxy usage token for platform_managed", got)
}
if !res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = false, want true (proxy token is the credential)")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyRuntimeModelEnv_PersonaEnvMODELSecretPreserved locks in the
// 2026-05-08 fix that prevents the MODEL_PROVIDER-as-slug fallback from
// silently overwriting a per-persona MODEL workspace_secret on restart,
@@ -519,6 +519,11 @@ func TestWorkspaceCreate_SecretPersistFails_RollsBack(t *testing.T) {
// resolver expects the SELECT and the mock returns no row → falls back
// to the org default (byok, set above) so the OPENAI_API_KEY write
// reaches the INSERT-and-fail path this test exercises.
// internal#711/#713: the resolver reads the LLM_PROVIDER selection first.
// The provider secret is only written AFTER commit (post-create), so at the
// secret-write gate no LLM_PROVIDER row exists → ErrNoRows → falls through.
mock.ExpectQuery(`SELECT encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id = \$1 AND key = \$2`).
WillReturnRows(sqlmock.NewRows([]string{"encrypted_value", "encryption_version"}))
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
mock.ExpectExec("INSERT INTO workspace_secrets").