From a5ca482d7c541064cb4db7a47cebd4d034178a96 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 2 May 2026 15:30:39 +0800 Subject: [PATCH] feat(llm): implement scoped key management for providers - Introduced ScopedKey function to generate unique provider keys based on owner type (user/team) to prevent key collisions. - Updated provider creation and management tests to utilize scoped keys, ensuring consistent key handling across different contexts. - Refactored existing tests to validate the correct generation and retrieval of scoped keys, enhancing overall provider management reliability. - Adjusted OpenAPI settings to reflect the use of scoped keys in provider operations, improving API consistency. --- llmprovider/registry_test.go | 12 ++--- llmprovider/scope_test.go | 75 +++++++++++++++++++++++++++---- llmprovider/sync.go | 25 ++++++++--- openapi/setting/llm.go | 4 +- openapi/tests/setting/llm_test.go | 60 ++++++++++++++----------- 5 files changed, 129 insertions(+), 47 deletions(-) diff --git a/llmprovider/registry_test.go b/llmprovider/registry_test.go index 257b9d64..24aa7317 100644 --- a/llmprovider/registry_test.go +++ b/llmprovider/registry_test.go @@ -227,8 +227,9 @@ func TestGetMasked_EqualsGetDefault(t *testing.T) { func TestListModels_ConnectorHasRealKey(t *testing.T) { r := setupRegistry(t) + owner := llmprovider.ProviderOwner{Type: "user", UserID: "rk-user"} p := llmprovider.Provider{ - Key: "realkey-prov", + Key: llmprovider.ScopedKey(&owner, "realkey-prov"), Name: "RealKey Test", Type: "openai", APIURL: "https://api.openai.com", @@ -237,7 +238,7 @@ func TestListModels_ConnectorHasRealKey(t *testing.T) { {ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}, }, Enabled: true, - Owner: llmprovider.ProviderOwner{Type: "user", UserID: "rk-user"}, + Owner: owner, } _, err := r.Create(&p) require.NoError(t, err) @@ -290,13 +291,14 @@ func TestGetLazy(t *testing.T) { func TestList(t *testing.T) { r := setupRegistry(t) + p2Owner := llmprovider.ProviderOwner{Type: "user", UserID: "123"} providers := []llmprovider.Provider{ {Key: "p1", Name: "Provider 1", Type: "openai", Enabled: true, Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls"}, Enabled: true}}, Owner: llmprovider.ProviderOwner{Type: "system"}}, - {Key: "p2", Name: "Provider 2", Type: "anthropic", Enabled: false, + {Key: llmprovider.ScopedKey(&p2Owner, "p2"), Name: "Provider 2", Type: "anthropic", Enabled: false, Models: []llmprovider.ModelInfo{{ID: "claude-3", Name: "Claude 3", Capabilities: []string{"tool_calls"}, Enabled: true}}, - Owner: llmprovider.ProviderOwner{Type: "user", UserID: "123"}}, + Owner: p2Owner}, {Key: "p3", Name: "Provider 3", Type: "openai", Enabled: true, Models: []llmprovider.ModelInfo{{ID: "gpt-4o-mini", Name: "GPT-4o Mini", Capabilities: []string{"streaming"}, Enabled: true}}, Owner: llmprovider.ProviderOwner{Type: "system"}}, @@ -378,7 +380,7 @@ func TestList(t *testing.T) { require.NoError(t, err) found := false for _, p := range list { - if p.Key == "p2" { + if p.Name == "Provider 2" { found = true } } diff --git a/llmprovider/scope_test.go b/llmprovider/scope_test.go index fffa04f7..6f97ae97 100644 --- a/llmprovider/scope_test.go +++ b/llmprovider/scope_test.go @@ -12,6 +12,63 @@ import ( "github.com/yaoapp/yao/setting" ) +// --------------------------------------------------------------------------- +// ScopedKey +// --------------------------------------------------------------------------- + +func TestScopedKeyFormats(t *testing.T) { + assert.Equal(t, "ualice.deepseek", llmprovider.ScopedKey( + &llmprovider.ProviderOwner{Type: "user", UserID: "alice"}, "deepseek")) + assert.Equal(t, "t9253.deepseek", llmprovider.ScopedKey( + &llmprovider.ProviderOwner{Type: "team", TeamID: "9253"}, "deepseek")) + assert.Equal(t, "deepseek", llmprovider.ScopedKey( + &llmprovider.ProviderOwner{Type: "system"}, "deepseek")) + assert.Equal(t, "deepseek", llmprovider.ScopedKey( + &llmprovider.ProviderOwner{}, "deepseek")) +} + +func TestDifferentOwnerSameBaseKey(t *testing.T) { + r := setupRegistryWithSetting(t) + + ownerA := llmprovider.ProviderOwner{Type: "team", TeamID: "teamA"} + ownerB := llmprovider.ProviderOwner{Type: "team", TeamID: "teamB"} + + pA := createOwnedProvider(t, r, "deepseek", ownerA) + pB := createOwnedProvider(t, r, "deepseek", ownerB) + + assert.Equal(t, "tteamA.deepseek", pA.Key) + assert.Equal(t, "tteamB.deepseek", pB.Key) + + gotA, err := r.Get(pA.Key) + require.NoError(t, err) + assert.Equal(t, pA.Key, gotA.Key) + + gotB, err := r.Get(pB.Key) + require.NoError(t, err) + assert.Equal(t, pB.Key, gotB.Key) +} + +func TestSameOwnerDuplicateKey(t *testing.T) { + r := setupRegistryWithSetting(t) + + owner := llmprovider.ProviderOwner{Type: "user", UserID: "u1"} + _ = createOwnedProvider(t, r, "openai", owner) + + dup := llmprovider.Provider{ + Key: llmprovider.ScopedKey(&owner, "openai"), + Name: "Dup", + Type: "openai", + APIURL: "https://api.openai.com", + APIKey: "sk-dup", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Enabled: true}}, + Owner: owner, + } + _, err := r.Create(&dup) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already exists") +} + // --------------------------------------------------------------------------- // Identity interface // --------------------------------------------------------------------------- @@ -107,8 +164,9 @@ func TestListModelsBy_UserFallback(t *testing.T) { func TestListModelsExpandsMultipleModels(t *testing.T) { r := setupRegistryWithSetting(t) + owner := llmprovider.ProviderOwner{Type: "user", UserID: "multi-u"} p := llmprovider.Provider{ - Key: "multi-model-prov", + Key: llmprovider.ScopedKey(&owner, "multi-model-prov"), Name: "MultiModel", Type: "openai", APIURL: "https://api.openai.com", @@ -119,7 +177,7 @@ func TestListModelsExpandsMultipleModels(t *testing.T) { {ID: "gpt-disabled", Name: "Disabled", Enabled: false}, }, Enabled: true, - Owner: llmprovider.ProviderOwner{Type: "user", UserID: "multi-u"}, + Owner: owner, } _, err := r.Create(&p) require.NoError(t, err) @@ -143,8 +201,9 @@ func TestListModelsExpandsMultipleModels(t *testing.T) { func TestGetModelWithModelLevelCID(t *testing.T) { r := setupRegistryWithSetting(t) + owner := llmprovider.ProviderOwner{Type: "team", TeamID: "mlcid-t1"} p := llmprovider.Provider{ - Key: "mlcid-prov", + Key: llmprovider.ScopedKey(&owner, "mlcid-prov"), Name: "MLTest", Type: "openai", APIURL: "https://api.openai.com", @@ -154,7 +213,7 @@ func TestGetModelWithModelLevelCID(t *testing.T) { {ID: "gpt-4o-mini", Name: "GPT-4o Mini", Enabled: true}, }, Enabled: true, - Owner: llmprovider.ProviderOwner{Type: "team", TeamID: "mlcid-t1"}, + Owner: owner, } created, err := r.Create(&p) require.NoError(t, err) @@ -186,7 +245,7 @@ func TestGetModelByConnectorIDReverseLookup(t *testing.T) { p := createOwnedProvider(t, r, "rev-prov", llmprovider.ProviderOwner{Type: "user", UserID: "u99"}) cid := p.ConnectorID - assert.NotEqual(t, p.Key, cid, "dynamic provider ConnectorID should differ from Key") + assert.Equal(t, p.Key, cid, "dynamic provider ConnectorID should equal scoped Key") _ = connector.Unregister(cid) @@ -365,11 +424,11 @@ func TestGetRoleCapabilitiesBy(t *testing.T) { // helpers // --------------------------------------------------------------------------- -func createOwnedProvider(t *testing.T, r *llmprovider.Registry, key string, owner llmprovider.ProviderOwner) *llmprovider.Provider { +func createOwnedProvider(t *testing.T, r *llmprovider.Registry, baseKey string, owner llmprovider.ProviderOwner) *llmprovider.Provider { t.Helper() p := llmprovider.Provider{ - Key: key, - Name: "Test " + key, + Key: llmprovider.ScopedKey(&owner, baseKey), + Name: "Test " + baseKey, Type: "openai", APIURL: "https://api.openai.com", APIKey: "sk-test-owned", diff --git a/llmprovider/sync.go b/llmprovider/sync.go index 74c54718..084cc3ef 100644 --- a/llmprovider/sync.go +++ b/llmprovider/sync.go @@ -8,14 +8,29 @@ import ( goullm "github.com/yaoapp/gou/llm" ) +// ScopedKey returns a provider key prefixed with the owner scope. +// This ensures unique keys per user/team in the store. +// +// user -> "u." +// team -> "t." +// other -> baseKey (unchanged) +func ScopedKey(owner *ProviderOwner, baseKey string) string { + switch owner.Type { + case "user": + return "u" + owner.UserID + "." + baseKey + case "team": + return "t" + owner.TeamID + "." + baseKey + default: + return baseKey + } +} + // connectorID builds the runtime ID for registering into connector.Connectors. -// Dynamic providers get an owner prefix to avoid collision with builtin IDs. +// For user/team providers the Key is already scoped, so use it directly. func connectorID(p *Provider) string { switch p.Owner.Type { - case "user": - return "u" + p.Owner.UserID + "." + p.Key - case "team": - return "t" + p.Owner.TeamID + "." + p.Key + case "user", "team": + return p.Key default: return "s." + p.Key } diff --git a/openapi/setting/llm.go b/openapi/setting/llm.go index e2a3086b..203e0dc2 100644 --- a/openapi/setting/llm.go +++ b/openapi/setting/llm.go @@ -597,7 +597,7 @@ func handleLLMProviderCreate(c *gin.Context) { return } - provider.Key = presetKey + provider.Key = llmprovider.ScopedKey(owner, presetKey) provider.Name = preset.Name provider.Type = preset.Type provider.APIURL = preset.APIURL @@ -650,7 +650,7 @@ func handleLLMProviderCreate(c *gin.Context) { respondError(c, http.StatusBadRequest, "key is required for custom provider") return } - provider.Key = key + provider.Key = llmprovider.ScopedKey(owner, key) name, _ := body["name"].(string) if name == "" { diff --git a/openapi/tests/setting/llm_test.go b/openapi/tests/setting/llm_test.go index fb52b325..1efe157f 100644 --- a/openapi/tests/setting/llm_test.go +++ b/openapi/tests/setting/llm_test.go @@ -90,19 +90,21 @@ func llmBody(t *testing.T, resp *http.Response) map[string]interface{} { return body } -func createTestOpenAI(t *testing.T, serverURL, token string) { +func createTestOpenAI(t *testing.T, serverURL, token string) string { t.Helper() apiKey := requireOpenAIKey(t) - llmprovider.Global.Delete("openai") payload := map[string]interface{}{ "preset_key": "openai", "api_key": apiKey, "model_ids": []string{"gpt-4o", "gpt-4o-mini"}, } resp := llmPost(t, llmURL(serverURL, "/providers"), token, payload) - resp.Body.Close() + defer resp.Body.Close() assert.Equal(t, http.StatusCreated, resp.StatusCode, "createTestOpenAI should succeed") - t.Cleanup(func() { llmprovider.Global.Delete("openai") }) + body := llmBody(t, resp) + scopedKey, _ := body["key"].(string) + t.Cleanup(func() { llmprovider.Global.Delete(scopedKey) }) + return scopedKey } // ----------- Functional tests ----------- @@ -162,7 +164,6 @@ func TestLLMProviderCreate(t *testing.T) { initSettingRegistry(t) initLLMRegistry(t) token := obtainToken(t, serverURL) - llmprovider.Global.Delete("openai") payload := map[string]interface{}{ "preset_key": "openai", @@ -172,10 +173,12 @@ func TestLLMProviderCreate(t *testing.T) { resp := llmPost(t, llmURL(serverURL, "/providers"), token, payload) defer resp.Body.Close() assert.Equal(t, http.StatusCreated, resp.StatusCode) - t.Cleanup(func() { llmprovider.Global.Delete("openai") }) body := llmBody(t, resp) - assert.Equal(t, "openai", body["key"]) + scopedKey, _ := body["key"].(string) + t.Cleanup(func() { llmprovider.Global.Delete(scopedKey) }) + + assert.Contains(t, scopedKey, ".openai", "scoped key should end with .openai") assert.Equal(t, "OpenAI", body["name"]) assert.Equal(t, "openai", body["type"]) @@ -199,7 +202,6 @@ func TestLLMProviderCreateCustom(t *testing.T) { initSettingRegistry(t) initLLMRegistry(t) token := obtainToken(t, serverURL) - llmprovider.Global.Delete("my-custom-llm") payload := map[string]interface{}{ "key": "my-custom-llm", @@ -215,10 +217,12 @@ func TestLLMProviderCreateCustom(t *testing.T) { resp := llmPost(t, llmURL(serverURL, "/providers"), token, payload) defer resp.Body.Close() assert.Equal(t, http.StatusCreated, resp.StatusCode) - t.Cleanup(func() { llmprovider.Global.Delete("my-custom-llm") }) body := llmBody(t, resp) - assert.Equal(t, "my-custom-llm", body["key"]) + scopedKey, _ := body["key"].(string) + t.Cleanup(func() { llmprovider.Global.Delete(scopedKey) }) + + assert.Contains(t, scopedKey, ".my-custom-llm", "scoped key should end with .my-custom-llm") assert.Equal(t, "My Custom LLM", body["name"]) assert.Equal(t, true, body["is_custom"]) @@ -233,7 +237,7 @@ func TestLLMProviderUpdate(t *testing.T) { initLLMRegistry(t) token := obtainToken(t, serverURL) - createTestOpenAI(t, serverURL, token) + scopedKey := createTestOpenAI(t, serverURL, token) updatePayload := map[string]interface{}{ "name": "Updated OpenAI", @@ -242,7 +246,7 @@ func TestLLMProviderUpdate(t *testing.T) { {"id": "gpt-4o", "name": "GPT-4o Updated", "capabilities": []string{"vision", "tool_calls"}, "enabled": true}, }, } - resp := llmPut(t, llmURL(serverURL, "/providers/openai"), token, updatePayload) + resp := llmPut(t, llmURL(serverURL, "/providers/"+scopedKey), token, updatePayload) defer resp.Body.Close() assert.Equal(t, http.StatusOK, resp.StatusCode) @@ -265,19 +269,20 @@ func TestLLMProviderDelete(t *testing.T) { initSettingRegistry(t) initLLMRegistry(t) token := obtainToken(t, serverURL) - llmprovider.Global.Delete("anthropic") createPayload := map[string]interface{}{ "preset_key": "anthropic", "api_key": anthropicKey, } createResp := llmPost(t, llmURL(serverURL, "/providers"), token, createPayload) + createBody := llmBody(t, createResp) createResp.Body.Close() assert.Equal(t, http.StatusCreated, createResp.StatusCode) + scopedKey, _ := createBody["key"].(string) rolesPayload := map[string]interface{}{ "default": map[string]interface{}{ - "provider": "anthropic", + "provider": scopedKey, "model": "claude-sonnet-4-6", }, } @@ -285,7 +290,7 @@ func TestLLMProviderDelete(t *testing.T) { rolesResp.Body.Close() assert.Equal(t, http.StatusOK, rolesResp.StatusCode) - deleteResp := llmDelete(t, llmURL(serverURL, "/providers/anthropic"), token) + deleteResp := llmDelete(t, llmURL(serverURL, "/providers/"+scopedKey), token) defer deleteResp.Body.Close() assert.Equal(t, http.StatusOK, deleteResp.StatusCode) @@ -307,21 +312,22 @@ func TestLLMProviderDeleteForbidden(t *testing.T) { initLLMRegistry(t) token := obtainToken(t, serverURL) - llmprovider.Global.Delete("other-team-provider") + otherOwner := llmprovider.ProviderOwner{Type: "user", UserID: "some-other-user-999"} + scopedKey := llmprovider.ScopedKey(&otherOwner, "other-team-provider") otherProvider := &llmprovider.Provider{ - Key: "other-team-provider", + Key: scopedKey, Name: "Other Team's Provider", Type: "openai", APIURL: "https://api.example.com", Models: []llmprovider.ModelInfo{}, Enabled: true, Source: llmprovider.ProviderSourceDynamic, - Owner: llmprovider.ProviderOwner{Type: "user", UserID: "some-other-user-999"}, + Owner: otherOwner, } llmprovider.Global.Create(otherProvider) - t.Cleanup(func() { llmprovider.Global.Delete("other-team-provider") }) + t.Cleanup(func() { llmprovider.Global.Delete(scopedKey) }) - resp := llmDelete(t, llmURL(serverURL, "/providers/other-team-provider"), token) + resp := llmDelete(t, llmURL(serverURL, "/providers/"+scopedKey), token) defer resp.Body.Close() assert.Equal(t, http.StatusNotFound, resp.StatusCode, "should not be able to delete another user's provider") } @@ -333,9 +339,9 @@ func TestLLMProviderTest(t *testing.T) { initLLMRegistry(t) token := obtainToken(t, serverURL) - createTestOpenAI(t, serverURL, token) + scopedKey := createTestOpenAI(t, serverURL, token) - resp := llmPost(t, llmURL(serverURL, "/providers/openai/test"), token, nil) + resp := llmPost(t, llmURL(serverURL, "/providers/"+scopedKey+"/test"), token, nil) defer resp.Body.Close() assert.Equal(t, http.StatusOK, resp.StatusCode) @@ -351,15 +357,15 @@ func TestLLMRoles(t *testing.T) { initSettingRegistry(t) initLLMRegistry(t) token := obtainToken(t, serverURL) - createTestOpenAI(t, serverURL, token) + scopedKey := createTestOpenAI(t, serverURL, token) rolesPayload := map[string]interface{}{ "default": map[string]interface{}{ - "provider": "openai", + "provider": scopedKey, "model": "gpt-4o", }, "vision": map[string]interface{}{ - "provider": "openai", + "provider": scopedKey, "model": "gpt-4o", }, } @@ -404,11 +410,11 @@ func TestLLMRolesValidation(t *testing.T) { defer resp2.Body.Close() assert.Equal(t, http.StatusBadRequest, resp2.StatusCode, "should reject non-existent provider") - createTestOpenAI(t, serverURL, token) + scopedKey := createTestOpenAI(t, serverURL, token) resp3 := llmPut(t, llmURL(serverURL, "/roles"), token, map[string]interface{}{ "default": map[string]interface{}{ - "provider": "openai", + "provider": scopedKey, "model": "nonexistent-model", }, })