opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(llm): implement scoped key management for providers

Browse source 
- 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.
This commit is contained in:
Max 2026-05-02 15:30:39 +08:00
parent 756ff95d3f
commit a5ca482d7c
5 changed files with 129 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
llmprovider/registry_test.go
View file

@ -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
}
}

75
llmprovider/scope_test.go
View file

@ -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",

25
llmprovider/sync.go
View file

@ -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<userID>.<baseKey>"
// team -> "t<teamID>.<baseKey>"
// 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
}

4
openapi/setting/llm.go
View file

@ -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 == "" {

60
openapi/tests/setting/llm_test.go
View file

@ -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",
},
})