2026-05-25 00:14:20 +00:00
4 changed files with 136 additions and 14 deletions
@@ -30,7 +30,7 @@ func TestRefreshEnvFromCP_AppliesCPResponse(t *testing.T) {
 			t.Errorf("org id header: got %q", got)
 		}
 		w.Header().Set("Content-Type", "application/json")
-		fmt.Fprint(w, `{"MOLECULE_CP_SHARED_SECRET":"new-secret","MOLECULE_CP_URL":"https://api.moleculesai.app","DISPLAY_SESSION_SIGNING_SECRET":"display-secret"}`)
+		fmt.Fprint(w, `{"MOLECULE_CP_SHARED_SECRET":"new-secret","MOLECULE_CP_URL":"https://api.moleculesai.app","DISPLAY_SESSION_SIGNING_SECRET":"display-secret","MOLECULE_LLM_BASE_URL":"https://api.moleculesai.app/api/v1/internal/llm/openai/v1","MOLECULE_LLM_USAGE_TOKEN":"tenant-admin-token","MOLECULE_LLM_DEFAULT_MODEL":"moonshot/kimi-k2.6"}`)
 	}))
 	defer srv.Close()

@@ -48,6 +48,15 @@ func TestRefreshEnvFromCP_AppliesCPResponse(t *testing.T) {
 	if got := os.Getenv("DISPLAY_SESSION_SIGNING_SECRET"); got != "display-secret" {
 		t.Errorf("DISPLAY_SESSION_SIGNING_SECRET: want display-secret, got %q", got)
 	}
+	if got := os.Getenv("MOLECULE_LLM_BASE_URL"); got != "https://api.moleculesai.app/api/v1/internal/llm/openai/v1" {
+		t.Errorf("MOLECULE_LLM_BASE_URL: got %q", got)
+	}
+	if got := os.Getenv("MOLECULE_LLM_USAGE_TOKEN"); got != "tenant-admin-token" {
+		t.Errorf("MOLECULE_LLM_USAGE_TOKEN: got %q", got)
+	}
+	if got := os.Getenv("MOLECULE_LLM_DEFAULT_MODEL"); got != "moonshot/kimi-k2.6" {
+		t.Errorf("MOLECULE_LLM_DEFAULT_MODEL: got %q", got)
+	}
 }

 // TestRefreshEnvFromCP_CPUnreachableDoesNotFailBoot: network errors must
@@ -912,6 +912,48 @@ func applyRuntimeModelEnv(envVars map[string]string, runtime, model string) {
 	}
 }

+// applyPlatformManagedLLMEnv wires the control-plane LLM proxy into a
+// workspace only when the org is in platform-managed mode. Provider keys
+// never enter the tenant; OPENAI_API_KEY is the tenant token for the CP
+// OpenAI-compatible proxy.
+func applyPlatformManagedLLMEnv(envVars map[string]string, _ string, model string) {
+	if strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE"))) != "platform_managed" {
+		return
+	}
+	baseURL := firstNonEmptyEnv("MOLECULE_LLM_BASE_URL", "OPENAI_BASE_URL")
+	token := firstNonEmptyEnv("MOLECULE_LLM_USAGE_TOKEN", "OPENAI_API_KEY")
+	if baseURL == "" || token == "" {
+		return
+	}
+
+	envVars["MOLECULE_LLM_BILLING_MODE"] = "platform_managed"
+	envVars["MOLECULE_LLM_BASE_URL"] = baseURL
+	envVars["MOLECULE_LLM_USAGE_TOKEN"] = token
+	if usageURL := strings.TrimSpace(os.Getenv("MOLECULE_LLM_USAGE_URL")); usageURL != "" {
+		envVars["MOLECULE_LLM_USAGE_URL"] = usageURL
+	}
+
+	if strings.TrimSpace(envVars["OPENAI_API_KEY"]) == "" {
+		envVars["OPENAI_API_KEY"] = token
+		envVars["OPENAI_BASE_URL"] = baseURL
+	}
+
+	if model == "" && strings.TrimSpace(envVars["MOLECULE_MODEL"]) == "" && strings.TrimSpace(envVars["MODEL"]) == "" {
+		if defaultModel := strings.TrimSpace(os.Getenv("MOLECULE_LLM_DEFAULT_MODEL")); defaultModel != "" {
+			envVars["MOLECULE_MODEL"] = defaultModel
+		}
+	}
+}
+
+func firstNonEmptyEnv(names ...string) string {
+	for _, name := range names {
+		if v := strings.TrimSpace(os.Getenv(name)); v != "" {
+			return v
+		}
+	}
+	return ""
+}
+
 // loadWorkspaceSecrets loads global + workspace-specific secrets into a map.
 // Returns nil map + error string on decrypt failure. Shared by both Docker
 // and control plane provisioning paths to avoid duplication.
@@ -193,6 +193,7 @@ func (h *WorkspaceHandler) prepareProvisionContext(
 	// continue to rely on workspace_secrets / org-import persona-env
 	// merge for their git auth.
 	applyAgentGitHTTPCreds(envVars, payload.Role)
+	applyPlatformManagedLLMEnv(envVars, payload.Runtime, payload.Model)
 	applyRuntimeModelEnv(envVars, payload.Runtime, payload.Model)
 	if payload.Role != "" {
 		envVars["MOLECULE_AGENT_ROLE"] = payload.Role
@@ -241,10 +241,10 @@ func TestMintWorkspaceSecrets_PersistsInboundSecretInSaaSMode(t *testing.T) {
 // inherits it automatically.
 func TestPrepareProvisionContext_ParentIDInjection(t *testing.T) {
 	cases := []struct {
-		name       string
-		parentID   *string
-		expectKey  bool
-		expectVal  string
+		name      string
+		parentID  *string
+		expectKey bool
+		expectVal string
 	}{
 		{
 			name:      "parentID nil → no PARENT_ID env",
@@ -333,11 +333,11 @@ func TestPrepareProvisionContext_InjectsGitHTTPCredsFromPersonaToken(t *testing.
 	t.Setenv("MOLECULE_PERSONA_ROOT", root)

 	cases := []struct {
-		name          string
-		role          string
-		expectInject  bool
-		expectUser    string
-		expectPass    string
+		name         string
+		role         string
+		expectInject bool
+		expectUser   string
+		expectPass   string
 	}{
 		{
 			name:         "Dev-A slug role → persona token injected as GIT_HTTP_USERNAME/PASSWORD",
@@ -505,10 +505,10 @@ func TestPrepareProvisionContext_WorkspaceSecretWinsOverPersonaToken(t *testing.
 //
 // The four branches:
 //
-//   1. Secret already present → (s, false, nil)
-//   2. Secret missing, mint succeeds → (minted, true, nil)
-//   3. Secret missing, mint fails → ("", false, mint-err)
-//   4. Read fails (non-NoInboundSecret) → ("", false, read-err)
+//  1. Secret already present → (s, false, nil)
+//  2. Secret missing, mint succeeds → (minted, true, nil)
+//  3. Secret missing, mint fails → ("", false, mint-err)
+//  4. Read fails (non-NoInboundSecret) → ("", false, read-err)
 func TestReadOrLazyHealInboundSecret(t *testing.T) {
 	t.Run("secret already present → no heal, no error", func(t *testing.T) {
 		mock := setupTestDB(t)
@@ -964,6 +964,76 @@ func TestApplyRuntimeModelEnv_SetsUniversalMODELForAllRuntimes(t *testing.T) {
 	}
 }

+func TestApplyPlatformManagedLLMEnv_DefaultsOpenAIProxyWhenNoWorkspaceKey(t *testing.T) {
+	t.Setenv("MOLECULE_LLM_BILLING_MODE", "platform_managed")
+	t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
+	t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
+	t.Setenv("MOLECULE_LLM_USAGE_URL", "https://api.example.test/api/v1/internal/llm/usage")
+	t.Setenv("MOLECULE_LLM_DEFAULT_MODEL", "moonshot/kimi-k2.6")
+
+	envVars := map[string]string{}
+	applyPlatformManagedLLMEnv(envVars, "langgraph", "")
+	applyRuntimeModelEnv(envVars, "langgraph", "")
+
+	if got := envVars["OPENAI_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/openai/v1" {
+		t.Fatalf("OPENAI_BASE_URL = %q", got)
+	}
+	if got := envVars["OPENAI_API_KEY"]; got != "tenant-admin-token" {
+		t.Fatalf("OPENAI_API_KEY = %q", got)
+	}
+	if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" {
+		t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got)
+	}
+	if got := envVars["MODEL"]; got != "moonshot/kimi-k2.6" {
+		t.Fatalf("MODEL = %q", got)
+	}
+	if got := envVars["MOLECULE_MODEL"]; got != "moonshot/kimi-k2.6" {
+		t.Fatalf("MOLECULE_MODEL = %q", got)
+	}
+}
+
+func TestApplyPlatformManagedLLMEnv_DoesNotOverrideWorkspaceOpenAIKey(t *testing.T) {
+	t.Setenv("MOLECULE_LLM_BILLING_MODE", "platform_managed")
+	t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
+	t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
+
+	envVars := map[string]string{
+		"OPENAI_API_KEY":  "user-openai-key",
+		"OPENAI_BASE_URL": "https://api.openai.com/v1",
+		"MODEL":           "openai/gpt-5.5",
+	}
+	applyPlatformManagedLLMEnv(envVars, "langgraph", "")
+
+	if got := envVars["OPENAI_API_KEY"]; got != "user-openai-key" {
+		t.Fatalf("OPENAI_API_KEY was overwritten: %q", got)
+	}
+	if got := envVars["OPENAI_BASE_URL"]; got != "https://api.openai.com/v1" {
+		t.Fatalf("OPENAI_BASE_URL was overwritten: %q", got)
+	}
+	if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" {
+		t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got)
+	}
+	if got := envVars["MODEL"]; got != "openai/gpt-5.5" {
+		t.Fatalf("MODEL = %q", got)
+	}
+}
+
+func TestApplyPlatformManagedLLMEnv_NoopsOutsidePlatformManaged(t *testing.T) {
+	t.Setenv("MOLECULE_LLM_BILLING_MODE", "byok")
+	t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
+	t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
+
+	envVars := map[string]string{}
+	applyPlatformManagedLLMEnv(envVars, "langgraph", "")
+
+	if _, ok := envVars["OPENAI_API_KEY"]; ok {
+		t.Fatalf("OPENAI_API_KEY should not be set outside platform-managed mode")
+	}
+	if _, ok := envVars["MOLECULE_LLM_USAGE_TOKEN"]; ok {
+		t.Fatalf("MOLECULE_LLM_USAGE_TOKEN should not be set outside platform-managed mode")
+	}
+}
+
 // 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,