From 6c7ac1263133373b88b8bbe59afb64a4cc73428e Mon Sep 17 00:00:00 2001 From: claude-ceo-assistant Date: Tue, 26 May 2026 10:21:08 -0700 Subject: [PATCH] fix(workspace): strip provider keys in platform-managed LLM mode --- .../handlers/handlers_extended_test.go | 2 +- .../internal/handlers/mcp_test.go | 6 +-- .../internal/handlers/mcp_tools_test.go | 2 +- workspace-server/internal/handlers/secrets.go | 24 +++++++++++- .../internal/handlers/workspace_provision.go | 28 +++++--------- .../workspace_provision_shared_test.go | 37 ++++++++++--------- 6 files changed, 56 insertions(+), 43 deletions(-) diff --git a/workspace-server/internal/handlers/handlers_extended_test.go b/workspace-server/internal/handlers/handlers_extended_test.go index 51d141063..794a9eb88 100644 --- a/workspace-server/internal/handlers/handlers_extended_test.go +++ b/workspace-server/internal/handlers/handlers_extended_test.go @@ -302,7 +302,7 @@ func TestExtended_SecretsSetRejectsHermesCustomProviderInPlatformManagedMode(t * c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: "22222222-2222-2222-2222-222222222222"}} - body := `{"key":"HERMES_CUSTOM_BASE_URL","value":"https://api.moonshot.ai/v1"}` + body := `{"key":"KIMI_API_KEY","value":"sk-test-moonshot"}` c.Request = httptest.NewRequest("POST", "/workspaces/22222222-2222-2222-2222-222222222222/secrets", bytes.NewBufferString(body)) c.Request.Header.Set("Content-Type", "application/json") diff --git a/workspace-server/internal/handlers/mcp_test.go b/workspace-server/internal/handlers/mcp_test.go index 314c2cdc7..886f5329f 100644 --- a/workspace-server/internal/handlers/mcp_test.go +++ b/workspace-server/internal/handlers/mcp_test.go @@ -14,9 +14,9 @@ import ( "errors" - "github.com/DATA-DOG/go-sqlmock" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/memory/contract" + "github.com/DATA-DOG/go-sqlmock" "github.com/gin-gonic/gin" ) @@ -196,7 +196,7 @@ func TestMCPHandler_DelegateTask_RoutesThroughPlatformA2AProxy(t *testing.T) { expectCanCommunicateSiblings(mock, callerID, targetID, parentID) mock.ExpectExec(`(?s)INSERT INTO activity_logs.*'delegation'.*'delegate'`). - WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg()). + WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg(), "pending"). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectExec(`UPDATE activity_logs`). WithArgs("dispatched", "", callerID, sqlmock.AnyArg()). @@ -241,7 +241,7 @@ func TestMCPHandler_DelegateTaskAsync_RoutesThroughPlatformA2AProxy(t *testing.T expectCanCommunicateSiblings(mock, callerID, targetID, parentID) mock.ExpectExec(`(?s)INSERT INTO activity_logs.*'delegation'.*'delegate'`). - WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg()). + WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg(), "pending"). WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectExec(`UPDATE activity_logs`). WithArgs("dispatched", "", callerID, sqlmock.AnyArg()). diff --git a/workspace-server/internal/handlers/mcp_tools_test.go b/workspace-server/internal/handlers/mcp_tools_test.go index aeb09c1c6..509e7bf30 100644 --- a/workspace-server/internal/handlers/mcp_tools_test.go +++ b/workspace-server/internal/handlers/mcp_tools_test.go @@ -5,8 +5,8 @@ import ( "encoding/json" "testing" + "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db" "github.com/DATA-DOG/go-sqlmock" - "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" ) // ───────────────────────────────────────────────────────────────────────────── diff --git a/workspace-server/internal/handlers/secrets.go b/workspace-server/internal/handlers/secrets.go index f4d59746e..a468e7a7d 100644 --- a/workspace-server/internal/handlers/secrets.go +++ b/workspace-server/internal/handlers/secrets.go @@ -19,8 +19,28 @@ import ( var uuidRegex = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`) var platformManagedDirectLLMBypassKeys = map[string]struct{}{ - "HERMES_CUSTOM_API_KEY": {}, - "HERMES_CUSTOM_BASE_URL": {}, + "AI_GATEWAY_API_KEY": {}, + "ANTHROPIC_API_KEY": {}, + "ANTHROPIC_AUTH_TOKEN": {}, + "ARCEEAI_API_KEY": {}, + "CLAUDE_CODE_OAUTH_TOKEN": {}, + "DASHSCOPE_API_KEY": {}, + "DEEPSEEK_API_KEY": {}, + "GEMINI_API_KEY": {}, + "GLM_API_KEY": {}, + "HERMES_CUSTOM_API_KEY": {}, + "HERMES_CUSTOM_BASE_URL": {}, + "HF_TOKEN": {}, + "KIMI_API_KEY": {}, + "KIMI_CN_API_KEY": {}, + "MINIMAX_API_KEY": {}, + "MINIMAX_CN_API_KEY": {}, + "NOUS_API_KEY": {}, + "OPENAI_API_KEY": {}, + "OPENAI_BASE_URL": {}, + "OPENROUTER_API_KEY": {}, + "XAI_API_KEY": {}, + "ZAI_API_KEY": {}, } func isPlatformManagedDirectLLMBypassKey(key string) bool { diff --git a/workspace-server/internal/handlers/workspace_provision.go b/workspace-server/internal/handlers/workspace_provision.go index a31ac3c3a..9e2c96184 100644 --- a/workspace-server/internal/handlers/workspace_provision.go +++ b/workspace-server/internal/handlers/workspace_provision.go @@ -935,6 +935,7 @@ func applyPlatformManagedLLMEnv(envVars map[string]string, runtime string, model if baseURL == "" || token == "" { return } + stripPlatformManagedLLMBypassEnv(envVars) envVars["MOLECULE_LLM_BILLING_MODE"] = "platform_managed" envVars["MOLECULE_LLM_BASE_URL"] = baseURL @@ -946,11 +947,11 @@ func applyPlatformManagedLLMEnv(envVars map[string]string, runtime string, model envVars["MOLECULE_LLM_USAGE_URL"] = usageURL } - if strings.TrimSpace(envVars["OPENAI_API_KEY"]) == "" && !runtimeUsesAnthropicNativeProxy(runtime) { + if !runtimeUsesAnthropicNativeProxy(runtime) { envVars["OPENAI_API_KEY"] = token envVars["OPENAI_BASE_URL"] = baseURL } - if runtimeUsesAnthropicNativeProxy(runtime) && anthropicBaseURL != "" && workspaceHasNoAnthropicAuth(envVars) { + if runtimeUsesAnthropicNativeProxy(runtime) && anthropicBaseURL != "" { envVars["ANTHROPIC_API_KEY"] = token envVars["ANTHROPIC_BASE_URL"] = anthropicBaseURL } @@ -962,25 +963,14 @@ func applyPlatformManagedLLMEnv(envVars map[string]string, runtime string, model } } -func runtimeUsesAnthropicNativeProxy(runtime string) bool { - return strings.TrimSpace(strings.ToLower(runtime)) == "claude-code" +func stripPlatformManagedLLMBypassEnv(envVars map[string]string) { + for key := range platformManagedDirectLLMBypassKeys { + delete(envVars, key) + } } -func workspaceHasNoAnthropicAuth(envVars map[string]string) bool { - for _, key := range []string{ - "CLAUDE_CODE_OAUTH_TOKEN", - "ANTHROPIC_API_KEY", - "ANTHROPIC_AUTH_TOKEN", - "MINIMAX_API_KEY", - "KIMI_API_KEY", - "GLM_API_KEY", - "DEEPSEEK_API_KEY", - } { - if strings.TrimSpace(envVars[key]) != "" { - return false - } - } - return true +func runtimeUsesAnthropicNativeProxy(runtime string) bool { + return strings.TrimSpace(strings.ToLower(runtime)) == "claude-code" } func firstNonEmptyEnv(names ...string) string { diff --git a/workspace-server/internal/handlers/workspace_provision_shared_test.go b/workspace-server/internal/handlers/workspace_provision_shared_test.go index 6fd5386b5..a77429035 100644 --- a/workspace-server/internal/handlers/workspace_provision_shared_test.go +++ b/workspace-server/internal/handlers/workspace_provision_shared_test.go @@ -992,7 +992,7 @@ func TestApplyPlatformManagedLLMEnv_NonClaudeRuntimeDefaultsOpenAIProxyWhenNoWor } } -func TestApplyPlatformManagedLLMEnv_DoesNotOverrideWorkspaceOpenAIKey(t *testing.T) { +func TestApplyPlatformManagedLLMEnv_StripsWorkspaceOpenAIKeyForClaudeCode(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") @@ -1004,11 +1004,11 @@ func TestApplyPlatformManagedLLMEnv_DoesNotOverrideWorkspaceOpenAIKey(t *testing } applyPlatformManagedLLMEnv(envVars, "claude-code", "") - if got := envVars["OPENAI_API_KEY"]; got != "user-openai-key" { - t.Fatalf("OPENAI_API_KEY was overwritten: %q", got) + if _, ok := envVars["OPENAI_API_KEY"]; ok { + t.Fatalf("OPENAI_API_KEY should be stripped for claude-code platform-managed mode") } - if got := envVars["OPENAI_BASE_URL"]; got != "https://api.openai.com/v1" { - t.Fatalf("OPENAI_BASE_URL was overwritten: %q", got) + if _, ok := envVars["OPENAI_BASE_URL"]; ok { + t.Fatalf("OPENAI_BASE_URL should be stripped for claude-code platform-managed mode") } if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" { t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got) @@ -1018,7 +1018,7 @@ func TestApplyPlatformManagedLLMEnv_DoesNotOverrideWorkspaceOpenAIKey(t *testing } } -func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyWithoutOverwritingOAuth(t *testing.T) { +func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyOverOAuth(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_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic/v1") @@ -1030,11 +1030,14 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyWithoutOverwriti } applyPlatformManagedLLMEnv(envVars, "claude-code", "") - if got := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; got != "user-oauth-token" { - t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN was overwritten: %q", got) + if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok { + t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped in platform-managed mode") } - if _, ok := envVars["ANTHROPIC_API_KEY"]; ok { - t.Fatalf("ANTHROPIC_API_KEY should not be set when Claude OAuth is present") + if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" { + t.Fatalf("ANTHROPIC_API_KEY = %q", got) + } + if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" { + t.Fatalf("ANTHROPIC_BASE_URL = %q", got) } if got := envVars["MOLECULE_LLM_ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" { t.Fatalf("MOLECULE_LLM_ANTHROPIC_BASE_URL = %q", got) @@ -1061,7 +1064,7 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeInjectsAnthropicProxyWhenNoWorkspa } } -func TestApplyPlatformManagedLLMEnv_ClaudeCodeDoesNotOverrideVendorBYOK(t *testing.T) { +func TestApplyPlatformManagedLLMEnv_ClaudeCodeStripsVendorBYOK(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_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic/v1") @@ -1073,14 +1076,14 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeDoesNotOverrideVendorBYOK(t *testi } applyPlatformManagedLLMEnv(envVars, "claude-code", "") - if got := envVars["MINIMAX_API_KEY"]; got != "user-minimax-key" { - t.Fatalf("MINIMAX_API_KEY was overwritten: %q", got) + if _, ok := envVars["MINIMAX_API_KEY"]; ok { + t.Fatalf("MINIMAX_API_KEY should be stripped in platform-managed mode") } - if _, ok := envVars["ANTHROPIC_API_KEY"]; ok { - t.Fatalf("ANTHROPIC_API_KEY should not be set when vendor BYOK is present") + if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" { + t.Fatalf("ANTHROPIC_API_KEY = %q", got) } - if _, ok := envVars["ANTHROPIC_BASE_URL"]; ok { - t.Fatalf("ANTHROPIC_BASE_URL should not be set when vendor BYOK is present") + if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" { + t.Fatalf("ANTHROPIC_BASE_URL = %q", got) } if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" { t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got) -- 2.52.0