diff --git a/workspace-server/internal/handlers/workspace_provision.go b/workspace-server/internal/handlers/workspace_provision.go index 9d391d7a5..3b31487d3 100644 --- a/workspace-server/internal/handlers/workspace_provision.go +++ b/workspace-server/internal/handlers/workspace_provision.go @@ -953,14 +953,24 @@ func applyPlatformManagedLLMEnv(ctx context.Context, envVars map[string]string, log.Printf("workspace_provision: resolve billing mode workspace=%s err=%v (defaulting to platform_managed)", workspaceID, resolveErr) } log.Printf("workspace_provision: billing mode workspace=%s resolved=%s source=%s org_default=%s", workspaceID, res.ResolvedMode, res.Source, res.OrgDefault) + // internal#703: MOLECULE_LLM_BILLING_MODE in the container must reflect the + // RESOLVED per-workspace mode, not a hardcoded literal. Pre-fix this var was + // only emitted (hardcoded "platform_managed") on the strip path below, so a + // byok/disabled container never carried a truthful billing-mode value — only + // MOLECULE_LLM_BILLING_MODE_RESOLVED. Emit both here, resolver-driven, for + // every mode so the value is correct on the byok/disabled early-return path + // too (and downstream consumers / debug shells see byok, not platform_managed). + envVars["MOLECULE_LLM_BILLING_MODE"] = res.ResolvedMode // Observability: surface the resolved mode in the container env so the // agent / debug shell can answer "why is my key being stripped" without // pulling logs or hitting the admin route. envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"] = res.ResolvedMode if res.ResolvedMode != LLMBillingModePlatformManaged { - // byok or disabled — DO NOT strip vendor keys, DO NOT force-route to CP. + // byok or disabled — DO NOT strip vendor keys, DO NOT force-route to CP, + // DO NOT override the workspace own ANTHROPIC_BASE_URL / OAuth token. // Leave envVars alone so CLAUDE_CODE_OAUTH_TOKEN / vendor API keys - // pulled from workspace_secrets survive into the container. + // pulled from workspace_secrets survive into the container, and the + // workspace talks to its own provider directly (internal#703). return } baseURL := firstNonEmptyEnv("MOLECULE_LLM_BASE_URL", "OPENAI_BASE_URL") @@ -971,7 +981,8 @@ func applyPlatformManagedLLMEnv(ctx context.Context, envVars map[string]string, } stripPlatformManagedLLMBypassEnv(envVars) - envVars["MOLECULE_LLM_BILLING_MODE"] = "platform_managed" + // MOLECULE_LLM_BILLING_MODE is already set to res.ResolvedMode (== + // platform_managed on this path) above (internal#703); no hardcode here. envVars["MOLECULE_LLM_BASE_URL"] = baseURL envVars["MOLECULE_LLM_USAGE_TOKEN"] = token if anthropicBaseURL != "" { diff --git a/workspace-server/internal/handlers/workspace_provision_shared_test.go b/workspace-server/internal/handlers/workspace_provision_shared_test.go index a07ee4898..fc82dc0c8 100644 --- a/workspace-server/internal/handlers/workspace_provision_shared_test.go +++ b/workspace-server/internal/handlers/workspace_provision_shared_test.go @@ -1106,6 +1106,112 @@ func TestApplyPlatformManagedLLMEnv_NoopsOutsidePlatformManaged(t *testing.T) { } } +// TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv is the +// internal#703 regression guard: a per-workspace byok override (org-level +// MOLECULE_LLM_BILLING_MODE left at the platform_managed bootstrap floor) +// must resolve to byok and leave the workspace own provider env intact — +// the CP-injected proxy ANTHROPIC_BASE_URL / usage token must NOT be forced, +// the OAuth token must NOT be stripped, and MOLECULE_LLM_BILLING_MODE in the +// container must read the RESOLVED mode (byok), not the hardcoded literal. +// +// This is the discriminating test for the byok end-to-end fix: pre-fix the +// strip path was the only emitter of MOLECULE_LLM_BILLING_MODE (hardcoded +// "platform_managed"), so a byok container carried no truthful billing mode. +func TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv(t *testing.T) { + const wsID = "77777777-7777-7777-7777-777777777777" + mock := setupTestDB(t) + mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK)) + + // Org-level env left at the bootstrap floor — the per-workspace override + // is what must flip this workspace to byok (the realistic prod shape). + 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") + + // The workspace brought its own Claude Code OAuth token (BYOK via the + // subscription provider). It must survive untouched. + envVars := map[string]string{ + "CLAUDE_CODE_OAUTH_TOKEN": "user-oauth-token", + "MODEL": "sonnet", + } + applyPlatformManagedLLMEnv(context.Background(), envVars, wsID, "claude-code", "") + + // 1. OAuth token intact — not stripped. + if got := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; got != "user-oauth-token" { + t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q, want it left intact for byok", got) + } + // 2. No CP proxy base URL / usage token forced onto the workspace. + if got, ok := envVars["ANTHROPIC_BASE_URL"]; ok { + t.Fatalf("ANTHROPIC_BASE_URL must NOT be injected for byok, got %q", got) + } + if got, ok := envVars["ANTHROPIC_API_KEY"]; ok { + t.Fatalf("ANTHROPIC_API_KEY must NOT be injected for byok, got %q", got) + } + if got, ok := envVars["MOLECULE_LLM_ANTHROPIC_BASE_URL"]; ok { + t.Fatalf("MOLECULE_LLM_ANTHROPIC_BASE_URL must NOT be injected for byok, got %q", got) + } + if got, ok := envVars["MOLECULE_LLM_USAGE_TOKEN"]; ok { + t.Fatalf("MOLECULE_LLM_USAGE_TOKEN must NOT be injected for byok, got %q", got) + } + // 3. Billing mode in the container reflects the RESOLVED mode (byok). + if got := envVars["MOLECULE_LLM_BILLING_MODE"]; got != LLMBillingModeBYOK { + t.Fatalf("MOLECULE_LLM_BILLING_MODE = %q, want %q (resolver-driven, not hardcoded)", got, LLMBillingModeBYOK) + } + if got := envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"]; got != LLMBillingModeBYOK { + t.Fatalf("MOLECULE_LLM_BILLING_MODE_RESOLVED = %q, want %q", got, LLMBillingModeBYOK) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +// TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode is the +// no-regression companion: a workspace that resolves to platform_managed must +// still strip + force the proxy AND emit MOLECULE_LLM_BILLING_MODE= +// platform_managed (now resolver-driven, internal#703). Proves the byok fix +// did not alter the platform_managed contract. +func TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode(t *testing.T) { + const wsID = "88888888-8888-8888-8888-888888888888" + mock := setupTestDB(t) + mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged)) + + 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": "user-oauth-token", + "MODEL": "sonnet", + } + applyPlatformManagedLLMEnv(context.Background(), envVars, wsID, "claude-code", "") + + // OAuth stripped, proxy forced — unchanged platform_managed contract. + if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok { + t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped for platform_managed") + } + if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic" { + t.Fatalf("ANTHROPIC_BASE_URL = %q, want proxy forced for platform_managed", got) + } + if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" { + t.Fatalf("ANTHROPIC_API_KEY = %q, want usage token for platform_managed", got) + } + if got := envVars["MOLECULE_LLM_BILLING_MODE"]; got != LLMBillingModePlatformManaged { + t.Fatalf("MOLECULE_LLM_BILLING_MODE = %q, want %q", got, LLMBillingModePlatformManaged) + } + if got := envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"]; got != LLMBillingModePlatformManaged { + t.Fatalf("MOLECULE_LLM_BILLING_MODE_RESOLVED = %q, want %q", got, LLMBillingModePlatformManaged) + } + 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,