molecule-ai/molecule-core
52
0
Fork 2
Code Issues 196 Pull Requests 67 Actions Packages Projects Releases Wiki Activity

fix(workspace): strip provider keys in platform-managed LLM mode #1922

Merged
hongming merged 1 commits from fix/platform-managed-provider-key-leak into main 2026-05-26 17:51:52 +00:00
Conversation 8 Commits 1 Files Changed 6
+56 -43
6 changed files with 56 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
workspace-server/internal/handlers/handlers_extended_test.go
+1 -1
View File
@@ -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")
workspace-server/internal/handlers/mcp_test.go
+3 -3
View File
@@ -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()).
workspace-server/internal/handlers/mcp_tools_test.go
+1 -1
View File
@@ -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"
)
// ─────────────────────────────────────────────────────────────────────────────
workspace-server/internal/handlers/secrets.go
+22 -2
View File
@@ -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 {
workspace-server/internal/handlers/workspace_provision.go
+9 -19
View File
@@ -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 {
workspace-server/internal/handlers/workspace_provision_shared_test.go
+20 -17
View File
@@ -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)