diff --git a/.gitignore b/.gitignore index 300db619..1756c0d4 100644 --- a/.gitignore +++ b/.gitignore @@ -83,3 +83,4 @@ agent/robot/ROBOT-CACHE-IMPROVEMENT.md sandbox/v2/PID-KILL-UPGRADE.md sandbox/v2/*.md POSTGRESQL_COMPAT.md +openapi/setting/*.md diff --git a/Makefile b/Makefile index c17ace96..4f85b8d2 100644 --- a/Makefile +++ b/Makefile @@ -11,6 +11,7 @@ OS := $(shell uname) # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry|agent/sandbox/v2' | awk '!/\/tests\// || /openapi\/tests/' | grep -vE 'openapi/tests/(nodes|sandbox|workspace)') +# Sandbox setting tests (openapi/tests/setting/sandbox_test.go) require Docker + Tai — skipped in CI, run locally only # Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, grpc, and integrations which require external services) TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/' | grep -vE 'openapi/tests/(nodes|sandbox|workspace)') # Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys), robot packages (tested in robot job), and agent/sandbox/v2 (WIP, has its own job) @@ -36,7 +37,7 @@ TESTTAGS ?= "" unit-test: echo "mode: count" > coverage.out for d in $(TESTFOLDER); do \ - $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' $$d > tmp.out; \ + $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_|TestSandbox' $$d > tmp.out; \ cat tmp.out; \ if grep -q "^--- FAIL" tmp.out; then \ rm tmp.out; \ @@ -68,7 +69,7 @@ unit-test: unit-test-core: echo "mode: count" > coverage.out for d in $(TESTFOLDER_CORE); do \ - $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' $$d > tmp.out; \ + $(GO) test -tags $(TESTTAGS) -v -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_|TestSandbox' $$d > tmp.out; \ cat tmp.out; \ if grep -q "^--- FAIL" tmp.out; then \ rm tmp.out; \ diff --git a/engine/load.go b/engine/load.go index 823da786..f7acefe8 100644 --- a/engine/load.go +++ b/engine/load.go @@ -33,7 +33,9 @@ import ( "github.com/yaoapp/yao/i18n" "github.com/yaoapp/yao/job" "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/llmprovider" "github.com/yaoapp/yao/mcp" + "github.com/yaoapp/yao/mcpclient" "github.com/yaoapp/yao/messenger" "github.com/yaoapp/yao/model" "github.com/yaoapp/yao/monitor" @@ -46,6 +48,7 @@ import ( sandbox "github.com/yaoapp/yao/sandbox/v2" "github.com/yaoapp/yao/schedule" "github.com/yaoapp/yao/script" + "github.com/yaoapp/yao/setting" "github.com/yaoapp/yao/share" "github.com/yaoapp/yao/store" sui "github.com/yaoapp/yao/sui/api" @@ -421,6 +424,30 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string } }() + // Initialize LLM Provider Registry + err = loadStep("LLM Provider", func() error { + return llmprovider.Init() + }, callback) + if err != nil { + warnings = append(warnings, Warning{Widget: "LLM Provider", Error: err}) + } + + // Initialize MCP Client Registry + err = loadStep("MCP Client Registry", func() error { + return mcpclient.Init() + }, callback) + if err != nil { + warnings = append(warnings, Warning{Widget: "MCP Client Registry", Error: err}) + } + + // Initialize Setting Registry + err = loadStep("Setting Registry", func() error { + return setting.Init() + }, callback) + if err != nil { + warnings = append(warnings, Warning{Widget: "Setting Registry", Error: err}) + } + for name, hook := range LoadHooks { err = hook(cfg) if err != nil { @@ -655,6 +682,45 @@ func Reload(cfg config.Config, options LoadOption) (err error) { printErr(cfg.Mode, "Agent", err) } + // Reload LLM Provider Registry + if llmprovider.Global != nil { + err = llmprovider.Global.Reload() + if err != nil { + printErr(cfg.Mode, "LLM Provider", err) + } + } else { + err = llmprovider.Init() + if err != nil { + printErr(cfg.Mode, "LLM Provider", err) + } + } + + // Reload MCP Client Registry + if mcpclient.Global != nil { + err = mcpclient.Global.Reload() + if err != nil { + printErr(cfg.Mode, "MCP Client Registry", err) + } + } else { + err = mcpclient.Init() + if err != nil { + printErr(cfg.Mode, "MCP Client Registry", err) + } + } + + // Reload Setting Registry + if setting.Global != nil { + err = setting.Global.Reload() + if err != nil { + printErr(cfg.Mode, "Setting Registry", err) + } + } else { + err = setting.Init() + if err != nil { + printErr(cfg.Mode, "Setting Registry", err) + } + } + // Load OpenAPI _, err = openapi.Load(cfg) if err != nil { diff --git a/llmprovider/doc.go b/llmprovider/doc.go new file mode 100644 index 00000000..86a0c23d --- /dev/null +++ b/llmprovider/doc.go @@ -0,0 +1,12 @@ +package llmprovider + +import ( + _ "embed" + + "github.com/yaoapp/gou/doc" +) + +//go:embed doc.yml +var docYAML []byte + +func init() { doc.LoadYAML(docYAML) } diff --git a/llmprovider/doc.yml b/llmprovider/doc.yml new file mode 100644 index 00000000..cf54adad --- /dev/null +++ b/llmprovider/doc.yml @@ -0,0 +1,256 @@ +group: llmprovider +type: process +desc: | + CRUD operations for the LLM Provider Registry. Manages provider connections + (OpenAI, Anthropic, Ollama, etc.) with persistence, API key encryption, and + lazy connector registration. + + Process names follow the pattern "llmprovider.". + + Provider structure (returned by get, getmasked, create, update; array elements from list): + - key (string): Unique identifier for this provider. Required on create. + - connector_id (string): Runtime connector ID, auto-generated. + Format: "s." for system owner, "u." for user owner, + "t." for team owner. BuiltIn providers retain their original ID. + - name (string): Display name (e.g. "OpenAI", "My Custom Provider"). + - type (string): Connector protocol type. + Values: "openai", "anthropic", "google", "ollama", "custom". + - api_url (string): Base API URL (e.g. "https://api.openai.com"). + - api_key (string): API key. Returned in full by "get"; masked by "getmasked" + and "list" (e.g. "sk-***test"). Encrypted at rest with AES-256-GCM. + - models (array of ModelInfo): Available models for this provider. + - enabled (bool): Whether the provider is active. + - status (string): Connection status. Values: "connected", "disconnected", "unconfigured". + - is_custom (bool, optional): Whether user manually configured (not from preset). + - preset_key (string, optional): Key of the preset this was created from (e.g. "openai"). + - require_key (bool): Whether an API key is required. + - source (string): Origin. Values: "dynamic" (registry-created), "builtin" (loaded from .yao DSL). + - owner (ProviderOwner): Ownership information. + + ModelInfo structure (elements of Provider.models): + - id (string): Model identifier (e.g. "gpt-4o", "claude-sonnet-4-20250514"). + - name (string): Human-readable name (e.g. "GPT-4o"). + - capabilities (array of string): Model capabilities. + Known values: "vision", "tool_calls", "streaming", "json", "reasoning". + - enabled (bool): Whether this model is active. + + ProviderOwner structure (Provider.owner): + - type (string): Scope level. Values: "system", "team", "user". + - team_id (string, optional): Required when type is "team". + - user_id (string, optional): Required when type is "user". + + ProviderFilter structure (optional argument for list): + - source (string, optional): Filter by source. + Values: "dynamic" (default when omitted), "builtin", "all". + - owner (ProviderOwner, optional): Filter by owner. Omit to include all owners. + - enabled (bool, optional): Filter by enabled status. Omit to include both. + - type (string, optional): Filter by provider type (e.g. "openai"). + - preset_key (string, optional): Filter by preset key. + - capabilities (array of string, optional): AND filter — matches providers that have + at least one model satisfying ALL listed capabilities. + - keyword (string, optional): Case-insensitive substring search in key and name. + + ProviderPreset structure (returned by getpresets, getpreset): + - key (string): Preset identifier (e.g. "openai", "anthropic", "ollama"). + - name (string): Display name. + - type (string): Connector type. + - api_url (string): Default API URL for UI auto-fill. + - require_key (bool): Whether API key is required. + - is_cloud (bool, optional): Whether this is a cloud-hosted service. + - url_editable (bool, optional): Whether the user can modify the URL. + - default_models (array of ModelInfo): Suggested models for UI pre-population. + +entries: + - name: get + desc: | + Get a provider by key, returning the full Provider object with plaintext API key. + Lazily ensures the runtime connector is registered on first access. + Throws 404 if the provider key does not exist. + args: + - name: key + type: string + required: true + desc: Provider key (e.g. "openai", "my-custom-provider"). + return: + type: object + desc: | + Full Provider object. See Provider structure above. + The api_key field contains the decrypted plaintext value. + Example: {"key":"openai","connector_id":"s.openai","name":"OpenAI","type":"openai", + "api_url":"https://api.openai.com","api_key":"sk-abc123...", + "models":[{"id":"gpt-4o","name":"GPT-4o","capabilities":["vision","streaming"],"enabled":true}], + "enabled":true,"status":"connected","source":"dynamic", + "owner":{"type":"system"}} + + - name: getmasked + desc: | + Get a provider by key with the API key masked for safe display. + Masking rule: keeps last 4 characters visible, replaces every preceding + character with "*". Example: "sk-abc123test" (14 chars) → "**********test". + Keys with 4 or fewer characters are fully replaced with "*" per character + (e.g. "abcd" → "****", "ab" → "**"). + Throws 404 if the provider key does not exist. + args: + - name: key + type: string + required: true + desc: Provider key. + return: + type: object + desc: | + Provider object with api_key masked. All other fields are identical to "get". + Example api_key value: "**********test" (for a 14-char key) + + - name: create + desc: | + Create a new LLM provider. Persists to __yao.store (with API key encrypted), + registers a runtime connector, and returns the complete Provider object. + The "source" field is automatically set to "dynamic". + The "connector_id" field is auto-generated based on owner type. + Throws 400 if key is empty or already exists. + args: + - name: data + type: object + required: true + desc: | + Provider data object with the following fields: + - key (string, required): Unique provider key. + - name (string): Display name. + - type (string): Connector type ("openai", "anthropic", etc.). + - api_url (string): Base API URL. + - api_key (string): API key (will be encrypted for storage). + - models (array of ModelInfo): Model list. + - enabled (bool): Active status (default false). + - require_key (bool): Whether API key is required. + - owner (ProviderOwner): Ownership. Defaults to {"type":"system"}. + - preset_key (string, optional): Preset key if created from template. + - is_custom (bool, optional): Custom flag. + Example: + {"key":"my-openai","name":"My OpenAI","type":"openai", + "api_url":"https://api.openai.com","api_key":"sk-abc123", + "models":[{"id":"gpt-4o","name":"GPT-4o","capabilities":["streaming","vision"],"enabled":true}], + "enabled":true,"require_key":true,"owner":{"type":"user","user_id":"42"}} + return: + type: object + desc: | + Created Provider object with connector_id and source="dynamic" populated. + The api_key in the response is the plaintext value (not encrypted). + + - name: update + desc: | + Update an existing provider by key. Replaces the stored provider with the + provided data, re-encrypts the API key, hot-replaces the runtime connector, + and returns the updated Provider object. + IMPORTANT: This is a full replacement, not a partial merge. You must provide + all fields you want to keep (name, type, api_url, api_key, models, enabled, etc.). + Only "key", "source", "connector_id", and "owner" are automatically preserved + from the existing record if omitted or zero-valued in the input. + Throws 400 if the provider key is not found. + args: + - name: key + type: string + required: true + desc: Provider key to update. + - name: data + type: object + required: true + desc: | + Full Provider data object. Same field structure as "create". + The "key" field inside data is ignored; the first argument determines + which provider to update. Fields not provided will be reset to zero values + (empty string, false, nil), except source, connector_id, and owner which + fall back to the existing record's values. + return: + type: object + desc: Updated Provider object with all fields. + + - name: delete + desc: | + Delete a provider by key. Removes from persistent store, clears cache, + and unregisters the runtime connector. + Throws 404 if the provider key does not exist. + args: + - name: key + type: string + required: true + desc: Provider key to delete. + return: + type: "null" + desc: Returns null on success. + + - name: list + desc: | + List providers matching a filter. Returns an array of Provider objects + with API keys masked. When no filter is provided, defaults to source="dynamic" + (only registry-created providers). Pass {"source":"all"} to include both + dynamic and built-in (.yao DSL) providers. + args: + - name: filter + type: object + required: false + desc: | + ProviderFilter object. All fields are optional: + - source (string): "dynamic" (default), "builtin", or "all". + - owner (ProviderOwner): {"type":"user","user_id":"42"}. + - enabled (bool): true or false. + - type (string): e.g. "openai". + - preset_key (string): e.g. "openai". + - capabilities (array of string): e.g. ["vision","streaming"]. + - keyword (string): Substring search in key and name. + Example: {"source":"all","type":"openai","capabilities":["vision"]} + Omit this argument entirely to list all dynamic providers. + return: + type: array + desc: | + Array of Provider objects with api_key masked. + May be empty if no providers match the filter. + + - name: getsetting + desc: | + Get the runtime connector setting map for a provider. This returns the + low-level connection parameters as used by the connector engine. + Throws 404 if the provider key does not exist. + args: + - name: key + type: string + required: true + desc: Provider key. + return: + type: object + desc: | + Key-value map of connector settings. Typical fields: + - host (string): API host URL. + - model (string): Default model ID. + - key (string): API key (plaintext). + Exact fields depend on connector type. + Example: {"host":"https://api.openai.com","model":"gpt-4o","key":"sk-abc123"} + + - name: getpresets + desc: | + Get all provider presets. Presets are static UI-only templates loaded from + the embedded presets.yml at compile time. They do not participate in runtime + logic — only used for UI form auto-filling when creating a new provider. + args: [] + return: + type: array + desc: | + Array of ProviderPreset objects. See ProviderPreset structure above. + Currently includes: openai, anthropic, ollama, azure, yaoagents. + Example element: + {"key":"openai","name":"OpenAI","type":"openai", + "api_url":"https://api.openai.com","require_key":true, + "default_models":[{"id":"gpt-4o","name":"GPT-4o", + "capabilities":["vision","tool_calls","streaming","json"],"enabled":true}]} + + - name: getpreset + desc: | + Get a single provider preset by key. + Throws 404 if the preset key does not exist. + args: + - name: key + type: string + required: true + desc: 'Preset key. Available keys: "openai", "anthropic", "ollama", "azure", "yaoagents".' + return: + type: object + desc: ProviderPreset object. See ProviderPreset structure above. diff --git a/llmprovider/presets.go b/llmprovider/presets.go new file mode 100644 index 00000000..f05648ec --- /dev/null +++ b/llmprovider/presets.go @@ -0,0 +1,42 @@ +package llmprovider + +import ( + _ "embed" + + "gopkg.in/yaml.v3" +) + +//go:embed presets.yml +var presetsYAML []byte + +var presets []ProviderPreset + +func init() { + presets = loadPresets() +} + +func loadPresets() []ProviderPreset { + var list []ProviderPreset + if err := yaml.Unmarshal(presetsYAML, &list); err != nil { + panic("llmprovider: failed to parse presets.yml: " + err.Error()) + } + return list +} + +// GetPresets returns a copy of the embedded preset list. +func GetPresets() []ProviderPreset { + out := make([]ProviderPreset, len(presets)) + copy(out, presets) + return out +} + +// GetPreset returns the preset for the given key, or nil if not found. +func GetPreset(key string) *ProviderPreset { + for i := range presets { + if presets[i].Key == key { + cp := presets[i] + return &cp + } + } + return nil +} diff --git a/llmprovider/presets.yml b/llmprovider/presets.yml new file mode 100644 index 00000000..8340968e --- /dev/null +++ b/llmprovider/presets.yml @@ -0,0 +1,61 @@ +- key: yaoagents + name: Yao Agents + type: openai + api_url: https://api.yaoagents.com + require_key: false + is_cloud: true + default_models: + - id: default + name: Default + capabilities: [vision, tool_calls, streaming] + enabled: true + +- key: openai + name: OpenAI + type: openai + api_url: https://api.openai.com + require_key: true + default_models: + - id: gpt-4o + name: GPT-4o + capabilities: [vision, tool_calls, streaming, json] + enabled: true + - id: gpt-4o-mini + name: GPT-4o Mini + capabilities: [tool_calls, streaming, json] + enabled: true + - id: o3-mini + name: o3-mini + capabilities: [tool_calls, streaming, reasoning] + enabled: false + +- key: anthropic + name: Anthropic + type: anthropic + api_url: https://api.anthropic.com + require_key: true + default_models: + - id: claude-sonnet-4-20250514 + name: Claude Sonnet 4 + capabilities: [vision, tool_calls, streaming, reasoning] + enabled: true + - id: claude-haiku-3-5-20241022 + name: Claude Haiku 3.5 + capabilities: [tool_calls, streaming] + enabled: true + +- key: ollama + name: Ollama + type: openai + api_url: http://localhost:11434 + require_key: false + url_editable: true + default_models: [] + +- key: azure + name: Azure OpenAI + type: openai + api_url: "" + require_key: true + url_editable: true + default_models: [] diff --git a/llmprovider/process.go b/llmprovider/process.go new file mode 100644 index 00000000..aa0a9759 --- /dev/null +++ b/llmprovider/process.go @@ -0,0 +1,170 @@ +package llmprovider + +import ( + "encoding/json" + + "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/exception" +) + +func init() { + process.RegisterGroup("llmprovider", map[string]process.Handler{ + "get": ProcessGet, + "getmasked": ProcessGetMasked, + "create": ProcessCreate, + "update": ProcessUpdate, + "delete": ProcessDelete, + "list": ProcessList, + "getsetting": ProcessGetSetting, + "getpresets": ProcessGetPresets, + "getpreset": ProcessGetPreset, + }) +} + +func requireGlobal() { + if Global == nil { + exception.New("LLM Provider Registry not initialized", 500).Throw() + } +} + +// ProcessGet retrieves a provider by key. +// Args[0] string: provider key +func ProcessGet(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + key := p.ArgsString(0) + + provider, err := Global.Get(key) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return provider +} + +// ProcessGetMasked retrieves a provider with API key masked. +// Args[0] string: provider key +func ProcessGetMasked(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + key := p.ArgsString(0) + + provider, err := Global.GetMasked(key) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return provider +} + +// ProcessCreate adds a new provider. +// Args[0] map: Provider data +func ProcessCreate(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + + var provider Provider + raw, err := json.Marshal(p.Args[0]) + if err != nil { + exception.New("invalid provider data: "+err.Error(), 400).Throw() + } + if err := json.Unmarshal(raw, &provider); err != nil { + exception.New("invalid provider data: "+err.Error(), 400).Throw() + } + + result, err := Global.Create(&provider) + if err != nil { + exception.New(err.Error(), 400).Throw() + } + return result +} + +// ProcessUpdate modifies an existing provider. +// Args[0] string: provider key +// Args[1] map: Provider data +func ProcessUpdate(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(2) + key := p.ArgsString(0) + + var provider Provider + raw, err := json.Marshal(p.Args[1]) + if err != nil { + exception.New("invalid provider data: "+err.Error(), 400).Throw() + } + if err := json.Unmarshal(raw, &provider); err != nil { + exception.New("invalid provider data: "+err.Error(), 400).Throw() + } + + result, err := Global.Update(key, &provider) + if err != nil { + exception.New(err.Error(), 400).Throw() + } + return result +} + +// ProcessDelete removes a provider by key. +// Args[0] string: provider key +func ProcessDelete(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + key := p.ArgsString(0) + + if err := Global.Delete(key); err != nil { + exception.New(err.Error(), 404).Throw() + } + return nil +} + +// ProcessList returns providers matching a filter. +// Args[0] map: ProviderFilter (optional) +func ProcessList(p *process.Process) interface{} { + requireGlobal() + + var filter *ProviderFilter + if len(p.Args) > 0 && p.Args[0] != nil { + raw, err := json.Marshal(p.Args[0]) + if err == nil { + var f ProviderFilter + if json.Unmarshal(raw, &f) == nil { + filter = &f + } + } + } + + result, err := Global.List(filter) + if err != nil { + exception.New(err.Error(), 500).Throw() + } + return result +} + +// ProcessGetSetting returns the runtime connector setting map. +// Args[0] string: provider key +func ProcessGetSetting(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + key := p.ArgsString(0) + + setting, err := Global.GetSetting(key) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return setting +} + +// ProcessGetPresets returns all provider presets. +func ProcessGetPresets(p *process.Process) interface{} { + return GetPresets() +} + +// ProcessGetPreset returns a single preset by key. +// Args[0] string: preset key +func ProcessGetPreset(p *process.Process) interface{} { + p.ValidateArgNums(1) + key := p.ArgsString(0) + + preset := GetPreset(key) + if preset == nil { + exception.New("preset "+key+" not found", 404).Throw() + } + return preset +} diff --git a/llmprovider/process_test.go b/llmprovider/process_test.go new file mode 100644 index 00000000..273fceac --- /dev/null +++ b/llmprovider/process_test.go @@ -0,0 +1,164 @@ +package llmprovider_test + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/process" +) + +func TestProcessCreate(t *testing.T) { + setupRegistry(t) + + p := process.New("llmprovider.create", map[string]interface{}{ + "key": "proc-test", + "name": "Proc Test", + "type": "openai", + "api_url": "https://api.openai.com", + "api_key": "sk-proc-test", + "enabled": true, + "require_key": true, + "models": []interface{}{map[string]interface{}{"id": "gpt-4o", "name": "GPT-4o", "capabilities": []interface{}{"streaming"}, "enabled": true}}, + "owner": map[string]interface{}{"type": "system"}, + }) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + + m := toMapResult(t, result) + assert.Equal(t, "proc-test", m["key"]) + assert.NotEmpty(t, m["connector_id"]) + assert.Equal(t, "dynamic", m["source"]) +} + +func TestProcessGet(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-get") + + p := process.New("llmprovider.get", "proc-get") + result, err := p.Exec() + require.NoError(t, err) + + m := toMapResult(t, result) + assert.Equal(t, "proc-get", m["key"]) + assert.Equal(t, "sk-proc-test", m["api_key"]) +} + +func TestProcessGetMasked(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-masked") + + p := process.New("llmprovider.getmasked", "proc-masked") + result, err := p.Exec() + require.NoError(t, err) + + m := toMapResult(t, result) + apiKey, _ := m["api_key"].(string) + assert.NotEqual(t, "sk-proc-test", apiKey) + assert.Contains(t, apiKey, "test") +} + +func TestProcessUpdate(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-upd") + + p := process.New("llmprovider.update", "proc-upd", map[string]interface{}{ + "name": "Updated Name", + "api_url": "https://custom.openai.com", + "enabled": true, + "models": []interface{}{map[string]interface{}{"id": "gpt-4o", "name": "GPT-4o", "capabilities": []interface{}{"streaming"}, "enabled": true}}, + }) + result, err := p.Exec() + require.NoError(t, err) + + m := toMapResult(t, result) + assert.Equal(t, "Updated Name", m["name"]) + assert.Equal(t, "https://custom.openai.com", m["api_url"]) +} + +func TestProcessDelete(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-del") + + p := process.New("llmprovider.delete", "proc-del") + _, err := p.Exec() + require.NoError(t, err) + + pGet := process.New("llmprovider.get", "proc-del") + _, err = pGet.Exec() + assert.Error(t, err) +} + +func TestProcessList(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-list-1") + createViaProcess(t, "proc-list-2") + + p := process.New("llmprovider.list", map[string]interface{}{ + "source": "dynamic", + }) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + t.Logf("list result type: %T", result) +} + +func TestProcessGetSetting(t *testing.T) { + setupRegistry(t) + createViaProcess(t, "proc-setting") + + p := process.New("llmprovider.getsetting", "proc-setting") + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) +} + +func TestProcessGetPresets(t *testing.T) { + p := process.New("llmprovider.getpresets") + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) +} + +func TestProcessGetPreset(t *testing.T) { + p := process.New("llmprovider.getpreset", "openai") + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + + m := toMapResult(t, result) + assert.Equal(t, "openai", m["key"]) +} + +// --- helpers --- + +func createViaProcess(t *testing.T, key string) { + t.Helper() + p := process.New("llmprovider.create", map[string]interface{}{ + "key": key, + "name": "Test " + key, + "type": "openai", + "api_url": "https://api.openai.com", + "api_key": "sk-proc-test", + "enabled": true, + "require_key": true, + "models": []interface{}{map[string]interface{}{"id": "gpt-4o", "name": "GPT-4o", "capabilities": []interface{}{"streaming"}, "enabled": true}}, + "owner": map[string]interface{}{"type": "system"}, + }) + _, err := p.Exec() + require.NoError(t, err) +} + +func toMapResult(t *testing.T, v interface{}) map[string]interface{} { + t.Helper() + if m, ok := v.(map[string]interface{}); ok { + return m + } + raw, err := json.Marshal(v) + require.NoError(t, err) + var m map[string]interface{} + require.NoError(t, json.Unmarshal(raw, &m)) + return m +} diff --git a/llmprovider/registry.go b/llmprovider/registry.go new file mode 100644 index 00000000..5e2e39af --- /dev/null +++ b/llmprovider/registry.go @@ -0,0 +1,312 @@ +package llmprovider + +import ( + "fmt" + "strings" + "sync" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/store" +) + +// Global is the singleton LLM Provider Registry. +var Global *Registry + +// Registry manages LLM providers with CRUD, persistence, cache and runtime sync. +type Registry struct { + store store.Store + cache store.Store + encKey string + mu sync.RWMutex +} + +// Init initializes the global Registry. +// Must be called after store.Load (so __yao.store and __yao.cache are available). +func Init() error { + s, err := store.Get("__yao.store") + if err != nil { + return fmt.Errorf("llmprovider.Init: %w", err) + } + c, _ := store.Get("__yao.cache") + + r := &Registry{store: s, cache: c} + Global = r + + if err := importFromConnectors(r); err != nil { + return fmt.Errorf("llmprovider.Init importFromConnectors: %w", err) + } + + return nil +} + +// SetEncryptionKey sets the key used for API key encryption at rest. +// Should be called right after Init if encryption is desired. +func (r *Registry) SetEncryptionKey(key string) { + r.mu.Lock() + defer r.mu.Unlock() + r.encKey = key +} + +// Get retrieves a provider by key. Lazily ensures its connector is registered. +func (r *Registry) Get(key string) (*Provider, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + p, err := storeGet(r.store, r.cache, key, r.encKey) + if err != nil { + return nil, err + } + + _ = ensureConnector(p) + return p, nil +} + +// GetMasked retrieves a provider with the API key masked for display. +func (r *Registry) GetMasked(key string) (*Provider, error) { + p, err := r.Get(key) + if err != nil { + return nil, err + } + cp := *p + cp.APIKey = maskAPIKey(cp.APIKey) + return &cp, nil +} + +// Create adds a new provider. Persists, caches, registers connector, and updates index. +func (r *Registry) Create(p *Provider) (*Provider, error) { + r.mu.Lock() + defer r.mu.Unlock() + + if p.Key == "" { + return nil, fmt.Errorf("provider key is required") + } + if r.store.Has(storeKey(p.Key)) { + return nil, fmt.Errorf("provider %s already exists", p.Key) + } + + if p.Source == "" { + p.Source = ProviderSourceDynamic + } + if p.ConnectorID == "" { + p.ConnectorID = connectorID(p) + } + if p.Status == "" { + p.Status = "unconfigured" + } + + if err := storeSet(r.store, r.cache, p, r.encKey); err != nil { + return nil, err + } + if err := indexAdd(r.store, r.cache, p.Key); err != nil { + return nil, err + } + + if p.Enabled { + _ = ensureConnector(p) + } + + return p, nil +} + +// Update modifies an existing provider. Hot-replaces the connector if needed. +func (r *Registry) Update(key string, p *Provider) (*Provider, error) { + r.mu.Lock() + defer r.mu.Unlock() + + old, err := storeGet(r.store, r.cache, key, r.encKey) + if err != nil { + return nil, err + } + + p.Key = key + if p.Source == "" { + p.Source = old.Source + } + if p.ConnectorID == "" { + p.ConnectorID = old.ConnectorID + } + if p.Owner == (ProviderOwner{}) { + p.Owner = old.Owner + } + + _ = unregisterConnector(old) + + if err := storeSet(r.store, r.cache, p, r.encKey); err != nil { + return nil, err + } + + if p.Enabled { + _ = ensureConnector(p) + } + + return p, nil +} + +// Delete removes a provider by key. Unregisters connector, deletes store/cache/index. +func (r *Registry) Delete(key string) error { + r.mu.Lock() + defer r.mu.Unlock() + + p, err := storeGet(r.store, r.cache, key, r.encKey) + if err != nil { + return err + } + + _ = unregisterConnector(p) + + if err := storeDel(r.store, r.cache, key); err != nil { + return err + } + return indexRemove(r.store, r.cache, key) +} + +// List returns providers matching the filter. +func (r *Registry) List(filter *ProviderFilter) ([]Provider, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + keys, err := indexGet(r.store, r.cache) + if err != nil { + return nil, err + } + + var result []Provider + for _, key := range keys { + p, err := storeGet(r.store, r.cache, key, r.encKey) + if err != nil { + continue + } + if filter != nil && !matchFilter(p, filter) { + continue + } + cp := *p + cp.APIKey = maskAPIKey(cp.APIKey) + result = append(result, cp) + } + return result, nil +} + +// Reload re-reads all providers from persistent store and rebuilds cache + connectors. +func (r *Registry) Reload() error { + r.mu.Lock() + defer r.mu.Unlock() + + keys, err := indexGet(r.store, nil) + if err != nil { + return err + } + + for _, key := range keys { + p, err := storeGet(r.store, nil, key, r.encKey) + if err != nil { + continue + } + m, err := providerToMap(p, r.encKey) + if err != nil { + continue + } + if r.cache != nil { + r.cache.Set(storeKey(key), m, 0) + } + if p.Source == ProviderSourceDynamic && p.Enabled { + _ = ensureConnector(p) + } + } + return nil +} + +// GetConnector returns the runtime connector for a given provider key. +func (r *Registry) GetConnector(key string) (connector.Connector, error) { + p, err := r.Get(key) + if err != nil { + return nil, err + } + cid := p.ConnectorID + if cid == "" { + cid = connectorID(p) + } + return connector.Select(cid) +} + +// GetSetting returns the runtime connector setting map for a given provider key. +func (r *Registry) GetSetting(key string) (map[string]interface{}, error) { + conn, err := r.GetConnector(key) + if err != nil { + return nil, err + } + return conn.Setting(), nil +} + +// matchFilter checks if a provider matches the given filter. +func matchFilter(p *Provider, f *ProviderFilter) bool { + src := f.Source + if src == "" { + src = ProviderSourceDynamic + } + if src != ProviderSourceAll && p.Source != src { + return false + } + + if f.Owner != nil { + if f.Owner.Type != "" && p.Owner.Type != f.Owner.Type { + return false + } + if f.Owner.UserID != "" && p.Owner.UserID != f.Owner.UserID { + return false + } + if f.Owner.TeamID != "" && p.Owner.TeamID != f.Owner.TeamID { + return false + } + } + + if f.Enabled != nil && p.Enabled != *f.Enabled { + return false + } + + if f.Type != nil && p.Type != *f.Type { + return false + } + + if f.PresetKey != nil && p.PresetKey != *f.PresetKey { + return false + } + + if len(f.Capabilities) > 0 && !matchCapabilities(p, f.Capabilities) { + return false + } + + if f.Keyword != "" { + kw := strings.ToLower(f.Keyword) + if !strings.Contains(strings.ToLower(p.Name), kw) && + !strings.Contains(strings.ToLower(p.Key), kw) { + return false + } + } + + return true +} + +// matchCapabilities returns true if at least one model in the provider +// satisfies ALL of the required capabilities (AND logic). +func matchCapabilities(p *Provider, required []string) bool { + for _, m := range p.Models { + if !m.Enabled { + continue + } + capSet := make(map[string]bool, len(m.Capabilities)) + for _, c := range m.Capabilities { + capSet[c] = true + } + allMatch := true + for _, req := range required { + if !capSet[req] { + allMatch = false + break + } + } + if allMatch { + return true + } + } + return false +} diff --git a/llmprovider/registry_test.go b/llmprovider/registry_test.go new file mode 100644 index 00000000..5b4cdf24 --- /dev/null +++ b/llmprovider/registry_test.go @@ -0,0 +1,645 @@ +package llmprovider_test + +import ( + "fmt" + "os" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/llmprovider" + "github.com/yaoapp/yao/test" +) + +func TestMain(m *testing.M) { + test.Prepare(nil, config.Conf) + defer test.Clean() + os.Exit(m.Run()) +} + +func setupRegistry(t *testing.T) *llmprovider.Registry { + t.Helper() + test.Prepare(t, config.Conf) + + err := llmprovider.Init() + require.NoError(t, err) + + t.Cleanup(func() { + s, _ := store.Get("__yao.store") + if s != nil { + s.Del("llmprovider:*") + } + c, _ := store.Get("__yao.cache") + if c != nil { + c.Del("llmprovider:*") + } + test.Clean() + }) + + return llmprovider.Global +} + +var testProvider = llmprovider.Provider{ + Key: "test-openai", + Name: "Test OpenAI", + Type: "openai", + APIURL: "https://api.openai.com", + APIKey: "sk-test-xxxxx", + Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls", "streaming"}, Enabled: true}}, + Enabled: true, + RequireKey: true, + Owner: llmprovider.ProviderOwner{Type: "system"}, +} + +func TestCreate(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + created, err := r.Create(&p) + require.NoError(t, err) + assert.Equal(t, "test-openai", created.Key) + assert.Equal(t, llmprovider.ProviderSourceDynamic, created.Source) + assert.NotEmpty(t, created.ConnectorID) + + // Verify store persistence + s, _ := store.Get("__yao.store") + assert.True(t, s.Has("llmprovider:p:test-openai")) + + // Verify connector registered + _, err = connector.Select(created.ConnectorID) + assert.NoError(t, err) +} + +func TestCreateDuplicate(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + _, err := r.Create(&p) + require.NoError(t, err) + + dup := testProvider + _, err = r.Create(&dup) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already exists") +} + +func TestGet(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + _, err := r.Create(&p) + require.NoError(t, err) + + got, err := r.Get("test-openai") + require.NoError(t, err) + assert.Equal(t, "Test OpenAI", got.Name) + assert.Equal(t, "openai", got.Type) + assert.Equal(t, "https://api.openai.com", got.APIURL) + assert.Len(t, got.Models, 1) + assert.Equal(t, "gpt-4o", got.Models[0].ID) +} + +func TestGetNotFound(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Get("nonexistent") + assert.Error(t, err) +} + +func TestGetMasked(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + _, err := r.Create(&p) + require.NoError(t, err) + + got, err := r.GetMasked("test-openai") + require.NoError(t, err) + assert.NotEqual(t, "sk-test-xxxxx", got.APIKey) + assert.True(t, len(got.APIKey) > 0) + // Last 4 chars should be visible + assert.Contains(t, got.APIKey, "xxxx") +} + +func TestGetLazy(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + created, err := r.Create(&p) + require.NoError(t, err) + + // Manually unregister the connector + err = connector.Unregister(created.ConnectorID) + require.NoError(t, err) + + // Verify it's gone + _, err = connector.Select(created.ConnectorID) + assert.Error(t, err) + + // Get should lazily re-register + got, err := r.Get("test-openai") + require.NoError(t, err) + assert.Equal(t, "test-openai", got.Key) + + // Connector should be back + _, err = connector.Select(got.ConnectorID) + assert.NoError(t, err) +} + +func TestList(t *testing.T) { + r := setupRegistry(t) + + 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, + Models: []llmprovider.ModelInfo{{ID: "claude-3", Name: "Claude 3", Capabilities: []string{"tool_calls"}, Enabled: true}}, + Owner: llmprovider.ProviderOwner{Type: "user", UserID: "123"}}, + {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"}}, + } + for i := range providers { + _, err := r.Create(&providers[i]) + require.NoError(t, err) + } + + t.Run("AllDynamic", func(t *testing.T) { + list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic}) + require.NoError(t, err) + assert.GreaterOrEqual(t, len(list), 3) + }) + + t.Run("FilterByType", func(t *testing.T) { + typ := "openai" + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + Type: &typ, + }) + require.NoError(t, err) + for _, p := range list { + assert.Equal(t, "openai", p.Type) + } + }) + + t.Run("FilterByEnabled", func(t *testing.T) { + enabled := true + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + Enabled: &enabled, + }) + require.NoError(t, err) + for _, p := range list { + assert.True(t, p.Enabled) + } + }) + + t.Run("FilterByOwner", func(t *testing.T) { + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + Owner: &llmprovider.ProviderOwner{Type: "user", UserID: "123"}, + }) + require.NoError(t, err) + for _, p := range list { + assert.Equal(t, "user", p.Owner.Type) + assert.Equal(t, "123", p.Owner.UserID) + } + }) + + t.Run("FilterByCapabilities", func(t *testing.T) { + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + Capabilities: []string{"vision", "tool_calls"}, + }) + require.NoError(t, err) + for _, p := range list { + found := false + for _, m := range p.Models { + capSet := map[string]bool{} + for _, c := range m.Capabilities { + capSet[c] = true + } + if capSet["vision"] && capSet["tool_calls"] { + found = true + break + } + } + assert.True(t, found, "provider %s should have model matching vision+tool_calls", p.Key) + } + }) + + t.Run("FilterByKeyword", func(t *testing.T) { + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + Keyword: "Provider 2", + }) + require.NoError(t, err) + found := false + for _, p := range list { + if p.Key == "p2" { + found = true + } + } + assert.True(t, found) + }) +} + +func TestUpdate(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + created, err := r.Create(&p) + require.NoError(t, err) + + updated := *created + updated.APIURL = "https://custom.openai.com" + updated.APIKey = "sk-new-key" + + result, err := r.Update("test-openai", &updated) + require.NoError(t, err) + assert.Equal(t, "https://custom.openai.com", result.APIURL) + + // Verify store updated + got, err := r.Get("test-openai") + require.NoError(t, err) + assert.Equal(t, "https://custom.openai.com", got.APIURL) +} + +func TestDelete(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + created, err := r.Create(&p) + require.NoError(t, err) + cid := created.ConnectorID + + err = r.Delete("test-openai") + require.NoError(t, err) + + // Verify removed from store + _, err = r.Get("test-openai") + assert.Error(t, err) + + // Verify connector unregistered + _, err = connector.Select(cid) + assert.Error(t, err) +} + +func TestReload(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + _, err := r.Create(&p) + require.NoError(t, err) + + // Clear cache to simulate stale state + c, _ := store.Get("__yao.cache") + if c != nil { + c.Del("llmprovider:*") + } + + err = r.Reload() + require.NoError(t, err) + + // Should still be able to get the provider + got, err := r.Get("test-openai") + require.NoError(t, err) + assert.Equal(t, "Test OpenAI", got.Name) +} + +func TestImportFromConnectors(t *testing.T) { + r := setupRegistry(t) + + // After Init, builtin connectors should be imported + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceAll, + }) + require.NoError(t, err) + + builtinCount := 0 + for _, p := range list { + if p.Source == llmprovider.ProviderSourceBuiltIn { + builtinCount++ + } + } + + // Should have imported some from connector.AIConnectors (if test app has connectors) + t.Logf("Imported %d builtin providers from connector.AIConnectors (total AIConnectors: %d)", builtinCount, len(connector.AIConnectors)) +} + +func TestGetPresets(t *testing.T) { + presets := llmprovider.GetPresets() + assert.Greater(t, len(presets), 0, "should have at least one preset") + + // Verify openai preset exists + var openai *llmprovider.ProviderPreset + for i := range presets { + if presets[i].Key == "openai" { + openai = &presets[i] + break + } + } + require.NotNil(t, openai, "openai preset should exist") + assert.Equal(t, "OpenAI", openai.Name) + assert.Equal(t, "openai", openai.Type) + assert.True(t, openai.RequireKey) + assert.Greater(t, len(openai.DefaultModels), 0) +} + +func TestGetPreset(t *testing.T) { + p := llmprovider.GetPreset("anthropic") + require.NotNil(t, p) + assert.Equal(t, "Anthropic", p.Name) + + none := llmprovider.GetPreset("nonexistent") + assert.Nil(t, none) +} + +func TestEncryptionRoundTrip(t *testing.T) { + r := setupRegistry(t) + r.SetEncryptionKey("my-super-secret-key-for-tests") + + p := testProvider + p.Key = "test-encrypted" + _, err := r.Create(&p) + require.NoError(t, err) + + got, err := r.Get("test-encrypted") + require.NoError(t, err) + assert.Equal(t, "sk-test-xxxxx", got.APIKey, "APIKey should be decrypted on read") + + masked, err := r.GetMasked("test-encrypted") + require.NoError(t, err) + assert.NotEqual(t, "sk-test-xxxxx", masked.APIKey) + assert.Contains(t, masked.APIKey, "xxxx") + + // Verify raw store value is encrypted + s, _ := store.Get("__yao.store") + raw, ok := s.Get("llmprovider:p:test-encrypted") + require.True(t, ok) + m := raw.(map[string]interface{}) + storedKey, _ := m["api_key"].(string) + assert.True(t, len(storedKey) > 0) + assert.NotEqual(t, "sk-test-xxxxx", storedKey, "raw stored value should be encrypted") +} + +func TestGetConnector(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + p.Key = "test-getconn" + _, err := r.Create(&p) + require.NoError(t, err) + + conn, err := r.GetConnector("test-getconn") + require.NoError(t, err) + assert.NotNil(t, conn) + + setting := conn.Setting() + assert.NotNil(t, setting) + host, _ := setting["host"].(string) + assert.Equal(t, "https://api.openai.com", host) +} + +func TestGetSetting(t *testing.T) { + r := setupRegistry(t) + + p := testProvider + p.Key = "test-getsetting" + _, err := r.Create(&p) + require.NoError(t, err) + + setting, err := r.GetSetting("test-getsetting") + require.NoError(t, err) + assert.NotNil(t, setting) + host, _ := setting["host"].(string) + assert.Equal(t, "https://api.openai.com", host) +} + +func TestGetConnectorNotFound(t *testing.T) { + r := setupRegistry(t) + _, err := r.GetConnector("not-exist") + assert.Error(t, err) +} + +func TestGetSettingNotFound(t *testing.T) { + r := setupRegistry(t) + _, err := r.GetSetting("not-exist") + assert.Error(t, err) +} + +func TestCreateEmptyKey(t *testing.T) { + r := setupRegistry(t) + p := llmprovider.Provider{Name: "No Key"} + _, err := r.Create(&p) + assert.Error(t, err) + assert.Contains(t, err.Error(), "key is required") +} + +func TestCreateDisabled(t *testing.T) { + r := setupRegistry(t) + p := llmprovider.Provider{ + Key: "test-disabled", + Name: "Disabled Provider", + Type: "openai", + APIURL: "https://api.openai.com", + Enabled: false, + Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}}, + Owner: llmprovider.ProviderOwner{Type: "system"}, + } + created, err := r.Create(&p) + require.NoError(t, err) + assert.Equal(t, "unconfigured", created.Status) + + // Disabled provider should not have its connector registered + _, err = connector.Select(created.ConnectorID) + assert.Error(t, err, "disabled provider should not register connector") +} + +func TestOwnerPrefixedIDs(t *testing.T) { + r := setupRegistry(t) + + cases := []struct { + key string + owner llmprovider.ProviderOwner + prefix string + }{ + {"owner-sys", llmprovider.ProviderOwner{Type: "system"}, "s."}, + {"owner-user", llmprovider.ProviderOwner{Type: "user", UserID: "42"}, "u42."}, + {"owner-team", llmprovider.ProviderOwner{Type: "team", TeamID: "99"}, "t99."}, + } + + for _, tc := range cases { + t.Run(tc.key, func(t *testing.T) { + p := llmprovider.Provider{ + Key: tc.key, + Name: tc.key, + Type: "openai", + APIURL: "https://api.openai.com", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "m1", Name: "M1", Capabilities: []string{"streaming"}, Enabled: true}}, + Owner: tc.owner, + } + created, err := r.Create(&p) + require.NoError(t, err) + assert.Contains(t, created.ConnectorID, tc.prefix, + "ConnectorID for %s owner should contain prefix %s", tc.owner.Type, tc.prefix) + + // Verify connector is registered with the prefixed ID + _, err = connector.Select(created.ConnectorID) + assert.NoError(t, err) + }) + } +} + +func TestListBuiltInFilter(t *testing.T) { + r := setupRegistry(t) + + builtinList, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceBuiltIn}) + require.NoError(t, err) + for _, p := range builtinList { + assert.Equal(t, llmprovider.ProviderSourceBuiltIn, p.Source) + } +} + +func TestListPresetKeyFilter(t *testing.T) { + r := setupRegistry(t) + + p := llmprovider.Provider{ + Key: "from-preset", + Name: "From Preset", + Type: "openai", + PresetKey: "openai", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}}, + Owner: llmprovider.ProviderOwner{Type: "system"}, + } + _, err := r.Create(&p) + require.NoError(t, err) + + pk := "openai" + list, err := r.List(&llmprovider.ProviderFilter{ + Source: llmprovider.ProviderSourceDynamic, + PresetKey: &pk, + }) + require.NoError(t, err) + found := false + for _, item := range list { + if item.Key == "from-preset" { + found = true + assert.Equal(t, "openai", item.PresetKey) + } + } + assert.True(t, found) +} + +func TestDefaultModelFallback(t *testing.T) { + r := setupRegistry(t) + + // Provider with no enabled models — should use first model ID as default + p := llmprovider.Provider{ + Key: "test-fallback", + Name: "Fallback", + Type: "openai", + APIURL: "https://api.openai.com", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "only-model", Name: "Only", Capabilities: []string{"streaming"}, Enabled: false}}, + Owner: llmprovider.ProviderOwner{Type: "system"}, + } + created, err := r.Create(&p) + require.NoError(t, err) + + // Connector should still be registered using the fallback model + conn, cerr := connector.Select(created.ConnectorID) + require.NoError(t, cerr) + setting := conn.Setting() + model, _ := setting["model"].(string) + assert.Equal(t, "only-model", model) +} + +func TestMaskShortKey(t *testing.T) { + r := setupRegistry(t) + + p := llmprovider.Provider{ + Key: "test-shortkey", + Name: "Short", + Type: "openai", + APIKey: "ab", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "m", Name: "M", Capabilities: []string{"streaming"}, Enabled: true}}, + Owner: llmprovider.ProviderOwner{Type: "system"}, + } + _, err := r.Create(&p) + require.NoError(t, err) + + masked, err := r.GetMasked("test-shortkey") + require.NoError(t, err) + // Short keys should be fully masked + assert.Equal(t, "**", masked.APIKey) +} + +func TestConcurrency(t *testing.T) { + r := setupRegistry(t) + + var wg sync.WaitGroup + errCh := make(chan error, 30) + + // Concurrent creates with unique keys + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + p := llmprovider.Provider{ + Key: fmt.Sprintf("conc-%d", idx), + Name: fmt.Sprintf("Concurrent %d", idx), + Type: "openai", + APIURL: "https://api.openai.com", + Enabled: true, + Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}}, + Owner: llmprovider.ProviderOwner{Type: "system"}, + } + if _, err := r.Create(&p); err != nil { + errCh <- err + } + }(i) + } + + wg.Wait() + + // Concurrent reads + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + _, err := r.Get(fmt.Sprintf("conc-%d", idx)) + if err != nil { + errCh <- err + } + }(i) + } + + wg.Wait() + + // Concurrent deletes + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + if err := r.Delete(fmt.Sprintf("conc-%d", idx)); err != nil { + errCh <- err + } + }(i) + } + + wg.Wait() + close(errCh) + + for err := range errCh { + t.Errorf("concurrent operation error: %v", err) + } +} diff --git a/llmprovider/store.go b/llmprovider/store.go new file mode 100644 index 00000000..dccd9df2 --- /dev/null +++ b/llmprovider/store.go @@ -0,0 +1,274 @@ +package llmprovider + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/yaoapp/gou/store" +) + +const ( + keyPrefix = "llmprovider:p:" + indexKey = "llmprovider:index" + maskChars = 4 + encPrefix = "enc:" +) + +func storeKey(key string) string { return keyPrefix + key } + +// providerToMap converts Provider to map[string]interface{} for store.Set. +// Encrypts APIKey before writing. +func providerToMap(p *Provider, encKey string) (map[string]interface{}, error) { + cp := *p + if cp.APIKey != "" && encKey != "" { + encrypted, err := encryptString(cp.APIKey, encKey) + if err != nil { + return nil, fmt.Errorf("encrypt api_key: %w", err) + } + cp.APIKey = encPrefix + encrypted + } + + raw, err := json.Marshal(cp) + if err != nil { + return nil, err + } + var m map[string]interface{} + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil +} + +// mapToProvider converts map[string]interface{} from store.Get back to Provider. +// Decrypts APIKey after reading. +func mapToProvider(m map[string]interface{}, encKey string) (*Provider, error) { + raw, err := json.Marshal(m) + if err != nil { + return nil, err + } + var p Provider + if err := json.Unmarshal(raw, &p); err != nil { + return nil, err + } + if strings.HasPrefix(p.APIKey, encPrefix) && encKey != "" { + decrypted, err := decryptString(strings.TrimPrefix(p.APIKey, encPrefix), encKey) + if err != nil { + return nil, fmt.Errorf("decrypt api_key: %w", err) + } + p.APIKey = decrypted + } + return &p, nil +} + +// maskAPIKey returns a masked version of the API key for display. +func maskAPIKey(key string) string { + if len(key) <= maskChars { + return strings.Repeat("*", len(key)) + } + return strings.Repeat("*", len(key)-maskChars) + key[len(key)-maskChars:] +} + +// storeGet reads a provider from cache first, then persistent store. +func storeGet(s, c store.Store, key, encKey string) (*Provider, error) { + sk := storeKey(key) + + if c != nil { + if val, ok := c.Get(sk); ok { + if m, ok := val.(map[string]interface{}); ok { + return mapToProvider(m, encKey) + } + } + } + + val, ok := s.Get(sk) + if !ok { + return nil, fmt.Errorf("provider %s not found", key) + } + m, ok := val.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("provider %s: unexpected store type %T", key, val) + } + + p, err := mapToProvider(m, encKey) + if err != nil { + return nil, err + } + + if c != nil { + c.Set(sk, m, 0) + } + return p, nil +} + +// storeSet writes a provider to both persistent store and cache. +func storeSet(s, c store.Store, p *Provider, encKey string) error { + m, err := providerToMap(p, encKey) + if err != nil { + return err + } + sk := storeKey(p.Key) + if err := s.Set(sk, m, 0); err != nil { + return err + } + if c != nil { + c.Set(sk, m, 0) + } + return nil +} + +// storeDel removes a provider from both persistent store and cache. +func storeDel(s, c store.Store, key string) error { + sk := storeKey(key) + if err := s.Del(sk); err != nil { + return err + } + if c != nil { + c.Del(sk) + } + return nil +} + +// indexGet returns all provider keys from the index. +func indexGet(s, c store.Store) ([]string, error) { + var raw interface{} + var ok bool + + if c != nil { + raw, ok = c.Get(indexKey) + } + if !ok { + raw, ok = s.Get(indexKey) + if !ok { + return nil, nil + } + if c != nil { + c.Set(indexKey, raw, 0) + } + } + + switch v := raw.(type) { + case []interface{}: + keys := make([]string, 0, len(v)) + for _, item := range v { + if str, ok := item.(string); ok { + keys = append(keys, str) + } + } + return keys, nil + case []string: + return v, nil + default: + return nil, fmt.Errorf("unexpected index type %T", raw) + } +} + +// indexSet writes the full index to both stores. +func indexSet(s, c store.Store, keys []string) error { + iface := make([]interface{}, len(keys)) + for i, k := range keys { + iface[i] = k + } + if err := s.Set(indexKey, iface, 0); err != nil { + return err + } + if c != nil { + c.Set(indexKey, iface, 0) + } + return nil +} + +// indexAdd appends a key to the index if not present. +func indexAdd(s, c store.Store, key string) error { + keys, err := indexGet(s, c) + if err != nil { + return err + } + for _, k := range keys { + if k == key { + return nil + } + } + return indexSet(s, c, append(keys, key)) +} + +// indexRemove removes a key from the index. +func indexRemove(s, c store.Store, key string) error { + keys, err := indexGet(s, c) + if err != nil { + return err + } + filtered := make([]string, 0, len(keys)) + for _, k := range keys { + if k != key { + filtered = append(filtered, k) + } + } + return indexSet(s, c, filtered) +} + +// --- AES-256-GCM encryption helpers --- + +func deriveKey(secret string) []byte { + h := sha256.Sum256([]byte(secret)) + return h[:] +} + +func encryptString(plaintext, secret string) (string, error) { + key := deriveKey(secret) + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return "", err + } + ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil) + return base64.StdEncoding.EncodeToString(ciphertext), nil +} + +func decryptString(encoded, secret string) (string, error) { + key := deriveKey(secret) + data, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", err + } + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonceSize := gcm.NonceSize() + if len(data) < nonceSize { + return "", fmt.Errorf("ciphertext too short") + } + plaintext, err := gcm.Open(nil, data[:nonceSize], data[nonceSize:], nil) + if err != nil { + return "", err + } + return string(plaintext), nil +} + +// storeCleanAll removes all llmprovider keys (for testing cleanup). +func storeCleanAll(s, c store.Store) { + _ = s.Del(keyPrefix + "*") + _ = s.Del(indexKey) + if c != nil { + _ = c.Del(keyPrefix + "*") + _ = c.Del(indexKey) + } +} diff --git a/llmprovider/sync.go b/llmprovider/sync.go new file mode 100644 index 00000000..089920a5 --- /dev/null +++ b/llmprovider/sync.go @@ -0,0 +1,199 @@ +package llmprovider + +import ( + "encoding/json" + "fmt" + + "github.com/yaoapp/gou/connector" +) + +// connectorID builds the runtime ID for registering into connector.Connectors. +// Dynamic providers get an owner prefix to avoid collision with builtin IDs. +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 + default: + return "s." + p.Key + } +} + +// defaultModel returns the first enabled model ID, or empty string. +func defaultModel(p *Provider) string { + for _, m := range p.Models { + if m.Enabled { + return m.ID + } + } + if len(p.Models) > 0 { + return p.Models[0].ID + } + return "" +} + +// marshalDSL builds a connector DSL JSON from the flat Provider fields. +func marshalDSL(p *Provider) ([]byte, error) { + dsl := map[string]interface{}{ + "type": p.Type, + "name": p.Name, + "label": p.Name, + "options": map[string]interface{}{ + "host": p.APIURL, + "key": p.APIKey, + "model": defaultModel(p), + }, + } + return json.Marshal(dsl) +} + +// ensureConnector makes sure the provider's connector is registered in the runtime. +// Builtin providers are managed by engine.Load and skipped here. +func ensureConnector(p *Provider) error { + if p.Source == ProviderSourceBuiltIn { + return nil + } + if !p.Enabled { + return nil + } + + cid := p.ConnectorID + if cid == "" { + cid = connectorID(p) + } + + if _, err := connector.Select(cid); err == nil { + return nil + } + + dslJSON, err := marshalDSL(p) + if err != nil { + return fmt.Errorf("ensureConnector %s: marshal DSL: %w", p.Key, err) + } + + _, err = connector.LoadSourceSync(dslJSON, cid, "__registry/"+cid+".conn.yao") + if err != nil { + return fmt.Errorf("ensureConnector %s: LoadSourceSync: %w", p.Key, err) + } + + return nil +} + +// unregisterConnector removes the provider's connector from the runtime. +func unregisterConnector(p *Provider) error { + if p.Source == ProviderSourceBuiltIn { + return nil + } + cid := p.ConnectorID + if cid == "" { + cid = connectorID(p) + } + return connector.Unregister(cid) +} + +// importFromConnectors scans existing AI connectors loaded by engine.Load +// and imports them as builtin providers into the Registry store. +// If a store record with the same key already exists (dynamic), it is not overwritten. +func importFromConnectors(r *Registry) error { + for _, opt := range connector.AIConnectors { + id := opt.Value + if r.store.Has(storeKey(id)) { + continue + } + + conn, err := connector.Select(id) + if err != nil { + continue + } + + p := providerFromConnector(id, conn) + m, err := providerToMap(&p, r.encKey) + if err != nil { + continue + } + sk := storeKey(id) + _ = r.store.Set(sk, m, 0) + if r.cache != nil { + _ = r.cache.Set(sk, m, 0) + } + _ = indexAdd(r.store, r.cache, id) + } + return nil +} + +// providerFromConnector builds a Provider from a runtime Connector interface. +func providerFromConnector(id string, conn connector.Connector) Provider { + meta := conn.GetMetaInfo() + setting := conn.Setting() + + name := meta.Label + if name == "" { + name = id + } + + typ := connectorType(conn) + apiURL, _ := setting["host"].(string) + apiKey, _ := setting["key"].(string) + model, _ := setting["model"].(string) + + var models []ModelInfo + if model != "" { + caps := capabilitiesFromSetting(setting) + models = []ModelInfo{{ + ID: model, + Name: model, + Capabilities: caps, + Enabled: true, + }} + } + + return Provider{ + Key: id, + ConnectorID: id, + Name: name, + Type: typ, + APIURL: apiURL, + APIKey: apiKey, + Models: models, + Enabled: true, + Status: "connected", + Source: ProviderSourceBuiltIn, + Owner: ProviderOwner{Type: "system"}, + } +} + +func connectorType(conn connector.Connector) string { + switch { + case conn.Is(6): // OPENAI + return "openai" + case conn.Is(11): // ANTHROPIC + return "anthropic" + case conn.Is(9): // FASTEMBED + return "fastembed" + case conn.Is(8): // MOAPI + return "moapi" + default: + return "custom" + } +} + +func capabilitiesFromSetting(setting map[string]interface{}) []string { + raw, ok := setting["capabilities"] + if !ok { + return nil + } + + switch caps := raw.(type) { + case map[string]interface{}: + var out []string + for k, v := range caps { + if b, ok := v.(bool); ok && b { + out = append(out, k) + } + } + return out + default: + return nil + } +} diff --git a/llmprovider/types.go b/llmprovider/types.go new file mode 100644 index 00000000..1b187770 --- /dev/null +++ b/llmprovider/types.go @@ -0,0 +1,85 @@ +package llmprovider + +// Provider represents a configured LLM provider (one vendor connection with multiple models). +// Fields align with the frontend ProviderConfig interface. +type Provider struct { + Key string `json:"key"` + ConnectorID string `json:"connector_id"` + Name string `json:"name"` + Type string `json:"type"` + APIURL string `json:"api_url"` + APIKey string `json:"api_key"` + Models []ModelInfo `json:"models"` + Enabled bool `json:"enabled"` + Status string `json:"status"` + IsCustom bool `json:"is_custom,omitempty"` + PresetKey string `json:"preset_key,omitempty"` + RequireKey bool `json:"require_key"` + Source ProviderSource `json:"source"` + Owner ProviderOwner `json:"owner"` +} + +// ModelInfo describes a single model within a provider. +// Fields align with the frontend ModelInfo interface. +type ModelInfo struct { + ID string `json:"id" yaml:"id"` + Name string `json:"name" yaml:"name"` + Capabilities []string `json:"capabilities" yaml:"capabilities"` + Enabled bool `json:"enabled" yaml:"enabled"` +} + +// ProviderOwner identifies who owns a provider. +type ProviderOwner struct { + Type string `json:"type"` + TeamID string `json:"team_id,omitempty"` + UserID string `json:"user_id,omitempty"` +} + +// ProviderSource distinguishes dynamic (registry-created) from builtin (DSL-loaded) providers. +type ProviderSource string + +const ( + ProviderSourceDynamic ProviderSource = "dynamic" + ProviderSourceBuiltIn ProviderSource = "builtin" + ProviderSourceAll ProviderSource = "all" +) + +// ProviderFilter specifies criteria for listing providers. +type ProviderFilter struct { + Owner *ProviderOwner + Enabled *bool + Source ProviderSource // defaults to "dynamic" when zero-value + Type *string + PresetKey *string + Capabilities []string // AND filter: provider matches if any model satisfies all + Keyword string +} + +// ProviderPreset is a static UI-only template for creating providers. +// Fields align with the frontend ProviderPreset interface. +type ProviderPreset struct { + Key string `json:"key" yaml:"key"` + Name string `json:"name" yaml:"name"` + Type string `json:"type" yaml:"type"` + APIURL string `json:"api_url" yaml:"api_url"` + RequireKey bool `json:"require_key" yaml:"require_key"` + IsCloud bool `json:"is_cloud,omitempty" yaml:"is_cloud,omitempty"` + URLEditable bool `json:"url_editable,omitempty" yaml:"url_editable,omitempty"` + DefaultModels []ModelInfo `json:"default_models" yaml:"default_models"` +} + +// ProviderTestResult holds the outcome of a provider connectivity test. +type ProviderTestResult struct { + Success bool `json:"success"` + Message string `json:"message"` + LatencyMs int64 `json:"latency_ms,omitempty"` +} + +// RoleAssignment maps model roles to specific provider+model pairs. +type RoleAssignment map[string]RoleTarget + +// RoleTarget identifies a provider and model for a given role. +type RoleTarget struct { + Provider string `json:"provider"` + Model string `json:"model"` +} diff --git a/mcpclient/doc.go b/mcpclient/doc.go new file mode 100644 index 00000000..2f951263 --- /dev/null +++ b/mcpclient/doc.go @@ -0,0 +1,12 @@ +package mcpclient + +import ( + _ "embed" + + "github.com/yaoapp/gou/doc" +) + +//go:embed doc.yml +var docYAML []byte + +func init() { doc.LoadYAML(docYAML) } diff --git a/mcpclient/doc.yml b/mcpclient/doc.yml new file mode 100644 index 00000000..0c665655 --- /dev/null +++ b/mcpclient/doc.yml @@ -0,0 +1,189 @@ +group: mcpclient +type: process +desc: | + CRUD operations for the MCP Client Registry. Manages MCP (Model Context Protocol) + client connections with persistence and lazy runtime registration. + + Process names follow the pattern "mcpclient.". + + Client structure (returned by get, create, update; array elements from list): + Embeds all fields from ClientDSL plus registry management fields. + + Inherited from ClientDSL: + - id (string): Unique client identifier. Required on create. + - name (string): Display name (e.g. "GitHub MCP", "File System"). + - version (string, optional): Client version. + - type (string, optional): Client type. Values: "standard", "agent", "system". + - transport (string): Transport protocol. Values: "stdio", "http", "sse", "process". + + Inherited from MetaInfo (embedded in ClientDSL): + - label (string, optional): Human-readable label for display. + - description (string, optional): Description text (markdown or plain). + - tags (array of string, optional): Categorization tags. + - readonly (bool, optional): Whether this client is read-only. + - builtin (bool, optional): Whether this is a built-in client. + + For stdio transport: + - command (string): Executable command (e.g. "npx", "python"). + - arguments (array of string): Command arguments (e.g. ["-y", "@modelcontextprotocol/server-github"]). + - env (object, optional): Environment variables as key-value pairs. + + For http/sse transport: + - url (string): Server URL. + - endpoint (string, optional): API endpoint path (e.g. "/api/mcp"). + - authorization_token (string, optional): Bearer token for authentication. + - timeout (string, optional): Request timeout (e.g. "30s", "5m"). + + For process transport: + - tools (object, optional): Tool name → process name mapping. + - prompts (object, optional): Prompt name → process name mapping. + - resources (object, optional): Resource name → process name mapping. + + Client capability flags: + - enable_sampling (bool, optional): Enable sampling capability. + - enable_roots (bool, optional): Enable roots capability. + - roots_list_changed (bool, optional): Subscribe to root change notifications. + - enable_elicitation (bool, optional): Enable elicitation capability. + + Dependencies: + - dependencies (object, optional): Other MCP clients this depends on (name → version constraint). + + Registry management fields (added by the registry): + - runtime_id (string): Runtime registration ID, auto-generated. + Format: "s." for system, "u." for user, "t." for team. + BuiltIn clients retain their original ID. + - enabled (bool): Whether the client is active. + - status (string): Connection status. Values: "connected", "disconnected", "unconfigured". + - source (string): Origin. Values: "dynamic" (registry-created), "builtin" (loaded from .yao DSL). + - tool_list (array of Tool, optional): Discovered tools from the MCP server. + Each Tool has: name (string), description (string), inputSchema (object). + - owner (ClientOwner): Ownership information. + + ClientOwner structure (Client.owner): + - type (string): Scope level. Values: "system", "team", "user". + - id (string, optional): Team ID or User ID depending on type. + + ClientFilter structure (optional argument for list): + - source (string, optional): Filter by source. + Values: "dynamic" (default when omitted), "builtin", "all". + - owner (ClientOwner, optional): Filter by owner. + - enabled (bool, optional): Filter by enabled status. Omit to include both. + - transport (string, optional): Filter by transport type ("stdio", "http", "sse", "process"). + - type (string, optional): Filter by client type ("standard", "agent", "system"). + - keyword (string, optional): Case-insensitive substring search in id, name, and label. + +entries: + - name: get + desc: | + Get an MCP client by ID, returning the full Client object. + Lazily ensures the runtime MCP client is registered on first access. + Throws 404 if the client ID does not exist. + args: + - name: id + type: string + required: true + desc: Client ID (e.g. "github-mcp", "filesystem"). + return: + type: object + desc: | + Full Client object. See Client structure above. + Example: {"id":"github-mcp","name":"GitHub MCP","type":"standard", + "transport":"stdio","command":"npx", + "arguments":["-y","@modelcontextprotocol/server-github"], + "enabled":true,"status":"connected","source":"dynamic", + "runtime_id":"s.github-mcp","owner":{"type":"system"}} + + - name: create + desc: | + Create a new MCP client. Persists to __yao.store, registers the runtime + MCP client, and returns the complete Client object. + The "source" field is automatically set to "dynamic". + The "runtime_id" field is auto-generated based on owner type. + Throws 400 if id is empty or already exists. + args: + - name: data + type: object + required: true + desc: | + Client data object. Required fields depend on transport type: + + For stdio transport: + {"id":"my-mcp","name":"My MCP","type":"standard","transport":"stdio", + "command":"npx","arguments":["-y","@some/mcp-server"], + "enabled":true,"owner":{"type":"system"}} + + For http/sse transport: + {"id":"remote-mcp","name":"Remote MCP","type":"standard","transport":"sse", + "url":"https://mcp.example.com","authorization_token":"Bearer xxx", + "enabled":true,"owner":{"type":"user","id":"42"}} + + For process transport: + {"id":"local-tools","name":"Local Tools","type":"standard","transport":"process", + "tools":{"search":"scripts.search.Run","fetch":"scripts.fetch.Run"}, + "enabled":true,"owner":{"type":"system"}} + return: + type: object + desc: Created Client object with runtime_id and source="dynamic" populated. + + - name: update + desc: | + Update an existing MCP client by ID. Replaces the stored client with the + provided data, hot-replaces the runtime client, and returns the updated object. + IMPORTANT: This is a full replacement, not a partial merge. You must provide + all fields you want to keep. Only "id", "source", "runtime_id", and "owner" + are automatically preserved from the existing record if omitted or zero-valued. + Throws 400 if the client ID is not found. + args: + - name: id + type: string + required: true + desc: Client ID to update. + - name: data + type: object + required: true + desc: | + Full Client data object. Same field structure as "create". + The "id" field inside data is ignored; the first argument determines + which client to update. Fields not provided will be reset to zero values, + except source, runtime_id, and owner which fall back to the existing values. + return: + type: object + desc: Updated Client object with all fields. + + - name: delete + desc: | + Delete an MCP client by ID. Removes from persistent store, clears cache, + and unloads the runtime MCP client. + Throws 404 if the client ID does not exist. + args: + - name: id + type: string + required: true + desc: Client ID to delete. + return: + type: "null" + desc: Returns null on success. + + - name: list + desc: | + List MCP clients matching a filter. When no filter is provided, defaults to + source="dynamic" (only registry-created clients). Pass {"source":"all"} to + include both dynamic and built-in (.yao DSL) clients. + args: + - name: filter + type: object + required: false + desc: | + ClientFilter object. All fields are optional: + - source (string): "dynamic" (default), "builtin", or "all". + - owner (ClientOwner): e.g. {"type":"user","id":"42"}. + - enabled (bool): true or false. + - transport (string): "stdio", "http", "sse", or "process". + - type (string): "standard", "agent", or "system". + - keyword (string): Substring search in id, name, and label. + Example: {"source":"all","transport":"stdio"} + Omit this argument entirely to list all dynamic clients. + return: + type: array + desc: | + Array of Client objects. May be empty if no clients match the filter. diff --git a/mcpclient/process.go b/mcpclient/process.go new file mode 100644 index 00000000..7d1995bb --- /dev/null +++ b/mcpclient/process.go @@ -0,0 +1,120 @@ +package mcpclient + +import ( + "encoding/json" + + "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/exception" +) + +func init() { + process.RegisterGroup("mcpclient", map[string]process.Handler{ + "get": ProcessGet, + "create": ProcessCreate, + "update": ProcessUpdate, + "delete": ProcessDelete, + "list": ProcessList, + }) +} + +func requireGlobal() { + if Global == nil { + exception.New("MCP Client Registry not initialized", 500).Throw() + } +} + +// ProcessGet retrieves a client by ID. +// Args[0] string: client ID +func ProcessGet(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + id := p.ArgsString(0) + + client, err := Global.Get(id) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return client +} + +// ProcessCreate adds a new MCP client. +// Args[0] map: Client data +func ProcessCreate(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + + var client Client + raw, err := json.Marshal(p.Args[0]) + if err != nil { + exception.New("invalid client data: "+err.Error(), 400).Throw() + } + if err := json.Unmarshal(raw, &client); err != nil { + exception.New("invalid client data: "+err.Error(), 400).Throw() + } + + result, err := Global.Create(&client) + if err != nil { + exception.New(err.Error(), 400).Throw() + } + return result +} + +// ProcessUpdate modifies an existing MCP client. +// Args[0] string: client ID +// Args[1] map: Client data +func ProcessUpdate(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(2) + id := p.ArgsString(0) + + var client Client + raw, err := json.Marshal(p.Args[1]) + if err != nil { + exception.New("invalid client data: "+err.Error(), 400).Throw() + } + if err := json.Unmarshal(raw, &client); err != nil { + exception.New("invalid client data: "+err.Error(), 400).Throw() + } + + result, err := Global.Update(id, &client) + if err != nil { + exception.New(err.Error(), 400).Throw() + } + return result +} + +// ProcessDelete removes a client by ID. +// Args[0] string: client ID +func ProcessDelete(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + id := p.ArgsString(0) + + if err := Global.Delete(id); err != nil { + exception.New(err.Error(), 404).Throw() + } + return nil +} + +// ProcessList returns clients matching a filter. +// Args[0] map: ClientFilter (optional) +func ProcessList(p *process.Process) interface{} { + requireGlobal() + + var filter *ClientFilter + if len(p.Args) > 0 && p.Args[0] != nil { + raw, err := json.Marshal(p.Args[0]) + if err == nil { + var f ClientFilter + if json.Unmarshal(raw, &f) == nil { + filter = &f + } + } + } + + result, err := Global.List(filter) + if err != nil { + exception.New(err.Error(), 500).Throw() + } + return result +} diff --git a/mcpclient/process_test.go b/mcpclient/process_test.go new file mode 100644 index 00000000..1f83e9c2 --- /dev/null +++ b/mcpclient/process_test.go @@ -0,0 +1,120 @@ +package mcpclient_test + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/process" +) + +func TestProcessCreate(t *testing.T) { + setupRegistry(t) + + p := process.New("mcpclient.create", map[string]interface{}{ + "id": "proc-test", + "name": "Proc Test", + "type": "standard", + "transport": "stdio", + "command": "echo", + "arguments": []interface{}{"hello"}, + "enabled": true, + "owner": map[string]interface{}{"type": "system"}, + }) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + + m := toMapResult(t, result) + assert.Equal(t, "proc-test", m["id"]) + assert.NotEmpty(t, m["runtime_id"]) + assert.Equal(t, "dynamic", m["source"]) +} + +func TestProcessGet(t *testing.T) { + setupRegistry(t) + createClientViaProcess(t, "proc-get") + + p := process.New("mcpclient.get", "proc-get") + result, err := p.Exec() + require.NoError(t, err) + + m := toMapResult(t, result) + assert.Equal(t, "proc-get", m["id"]) +} + +func TestProcessUpdate(t *testing.T) { + setupRegistry(t) + createClientViaProcess(t, "proc-upd") + + p := process.New("mcpclient.update", "proc-upd", map[string]interface{}{ + "name": "Updated MCP", + "type": "standard", + "transport": "stdio", + "command": "cat", + "enabled": true, + }) + result, err := p.Exec() + require.NoError(t, err) + + m := toMapResult(t, result) + assert.Equal(t, "Updated MCP", m["name"]) +} + +func TestProcessDelete(t *testing.T) { + setupRegistry(t) + createClientViaProcess(t, "proc-del") + + p := process.New("mcpclient.delete", "proc-del") + _, err := p.Exec() + require.NoError(t, err) + + pGet := process.New("mcpclient.get", "proc-del") + _, err = pGet.Exec() + assert.Error(t, err) +} + +func TestProcessList(t *testing.T) { + setupRegistry(t) + createClientViaProcess(t, "proc-list-1") + createClientViaProcess(t, "proc-list-2") + + p := process.New("mcpclient.list", map[string]interface{}{ + "source": "dynamic", + }) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + t.Logf("list result type: %T", result) +} + +// --- helpers --- + +func createClientViaProcess(t *testing.T, id string) { + t.Helper() + p := process.New("mcpclient.create", map[string]interface{}{ + "id": id, + "name": "Test " + id, + "type": "standard", + "transport": "stdio", + "command": "echo", + "arguments": []interface{}{"hello"}, + "enabled": true, + "owner": map[string]interface{}{"type": "system"}, + }) + _, err := p.Exec() + require.NoError(t, err) +} + +func toMapResult(t *testing.T, v interface{}) map[string]interface{} { + t.Helper() + if m, ok := v.(map[string]interface{}); ok { + return m + } + raw, err := json.Marshal(v) + require.NoError(t, err) + var m map[string]interface{} + require.NoError(t, json.Unmarshal(raw, &m)) + return m +} diff --git a/mcpclient/registry.go b/mcpclient/registry.go new file mode 100644 index 00000000..a7a24c10 --- /dev/null +++ b/mcpclient/registry.go @@ -0,0 +1,255 @@ +package mcpclient + +import ( + "fmt" + "strings" + "sync" + + "github.com/yaoapp/gou/mcp" + "github.com/yaoapp/gou/store" +) + +// Global is the singleton MCP Client Registry. +var Global *Registry + +// Registry manages MCP clients with CRUD, persistence, cache and runtime sync. +type Registry struct { + store store.Store + cache store.Store + mu sync.RWMutex +} + +// Init initializes the global Registry. +// Must be called after store.Load and mcp.Load. +func Init() error { + s, err := store.Get("__yao.store") + if err != nil { + return fmt.Errorf("mcpclient.Init: %w", err) + } + c, _ := store.Get("__yao.cache") + + r := &Registry{store: s, cache: c} + Global = r + + if err := importFromClients(r); err != nil { + return fmt.Errorf("mcpclient.Init importFromClients: %w", err) + } + + return nil +} + +// Get retrieves a client by ID. Lazily ensures its runtime client is registered. +func (r *Registry) Get(id string) (*Client, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + c, err := storeGet(r.store, r.cache, id) + if err != nil { + return nil, err + } + + _ = ensureClient(c) + return c, nil +} + +// Create adds a new client. Persists, caches, registers runtime, and updates index. +func (r *Registry) Create(c *Client) (*Client, error) { + r.mu.Lock() + defer r.mu.Unlock() + + if c.ID == "" { + return nil, fmt.Errorf("client id is required") + } + if r.store.Has(storeKey(c.ID)) { + return nil, fmt.Errorf("client %s already exists", c.ID) + } + + if c.Source == "" { + c.Source = ClientSourceDynamic + } + if c.RuntimeID == "" { + c.RuntimeID = runtimeID(c) + } + if c.Status == "" { + c.Status = "unconfigured" + } + + if err := storeSet(r.store, r.cache, c); err != nil { + return nil, err + } + if err := indexAdd(r.store, r.cache, c.ID); err != nil { + return nil, err + } + + if c.Enabled { + _ = ensureClient(c) + } + + return c, nil +} + +// Update modifies an existing client. Hot-replaces the runtime client. +func (r *Registry) Update(id string, c *Client) (*Client, error) { + r.mu.Lock() + defer r.mu.Unlock() + + old, err := storeGet(r.store, r.cache, id) + if err != nil { + return nil, err + } + + c.ID = id + if c.Source == "" { + c.Source = old.Source + } + if c.RuntimeID == "" { + c.RuntimeID = old.RuntimeID + } + if c.Owner == (ClientOwner{}) { + c.Owner = old.Owner + } + + unloadClient(old) + + if err := storeSet(r.store, r.cache, c); err != nil { + return nil, err + } + + if c.Enabled { + _ = ensureClient(c) + } + + return c, nil +} + +// Delete removes a client by ID. Unloads runtime, deletes store/cache/index. +func (r *Registry) Delete(id string) error { + r.mu.Lock() + defer r.mu.Unlock() + + c, err := storeGet(r.store, r.cache, id) + if err != nil { + return err + } + + unloadClient(c) + + if err := storeDel(r.store, r.cache, id); err != nil { + return err + } + return indexRemove(r.store, r.cache, id) +} + +// List returns clients matching the filter. +func (r *Registry) List(filter *ClientFilter) ([]Client, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + ids, err := indexGet(r.store, r.cache) + if err != nil { + return nil, err + } + + var result []Client + for _, id := range ids { + c, err := storeGet(r.store, r.cache, id) + if err != nil { + continue + } + if filter != nil && !matchFilter(c, filter) { + continue + } + result = append(result, *c) + } + return result, nil +} + +// Reload re-reads all clients from persistent store and rebuilds cache + runtime. +func (r *Registry) Reload() error { + r.mu.Lock() + defer r.mu.Unlock() + + ids, err := indexGet(r.store, nil) + if err != nil { + return err + } + + for _, id := range ids { + c, err := storeGet(r.store, nil, id) + if err != nil { + continue + } + m, err := clientToMap(c) + if err != nil { + continue + } + if r.cache != nil { + r.cache.Set(storeKey(id), m, 0) + } + if c.Source == ClientSourceDynamic && c.Enabled { + _ = ensureClient(c) + } + } + return nil +} + +// GetMCPClient returns the runtime mcp.Client for a given registry ID. +func (r *Registry) GetMCPClient(id string) (mcp.Client, error) { + c, err := r.Get(id) + if err != nil { + return nil, err + } + rid := c.RuntimeID + if rid == "" { + rid = runtimeID(c) + } + + defer func() { recover() }() + client := mcp.GetClient(rid) + if client == nil { + return nil, fmt.Errorf("runtime mcp client %s not found", rid) + } + return client, nil +} + +func matchFilter(c *Client, f *ClientFilter) bool { + src := f.Source + if src == "" { + src = ClientSourceDynamic + } + if src != ClientSourceAll && c.Source != src { + return false + } + + if f.Owner != nil { + if f.Owner.Type != "" && c.Owner.Type != f.Owner.Type { + return false + } + if f.Owner.ID != "" && c.Owner.ID != f.Owner.ID { + return false + } + } + + if f.Enabled != nil && c.Enabled != *f.Enabled { + return false + } + + if f.Transport != nil && c.ClientDSL.Transport != *f.Transport { + return false + } + + if f.Type != nil && c.ClientDSL.Type != *f.Type { + return false + } + + if f.Keyword != "" { + kw := strings.ToLower(f.Keyword) + if !strings.Contains(strings.ToLower(c.ClientDSL.Name), kw) && + !strings.Contains(strings.ToLower(c.ID), kw) && + !strings.Contains(strings.ToLower(c.ClientDSL.Label), kw) { + return false + } + } + + return true +} diff --git a/mcpclient/registry_test.go b/mcpclient/registry_test.go new file mode 100644 index 00000000..96d67a4e --- /dev/null +++ b/mcpclient/registry_test.go @@ -0,0 +1,483 @@ +package mcpclient_test + +import ( + "fmt" + "os" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/mcp" + mcpTypes "github.com/yaoapp/gou/mcp/types" + "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/mcpclient" + "github.com/yaoapp/yao/test" +) + +func TestMain(m *testing.M) { + test.Prepare(nil, config.Conf) + defer test.Clean() + os.Exit(m.Run()) +} + +func setupRegistry(t *testing.T) *mcpclient.Registry { + t.Helper() + test.Prepare(t, config.Conf) + + err := mcpclient.Init() + require.NoError(t, err) + + t.Cleanup(func() { + s, _ := store.Get("__yao.store") + if s != nil { + s.Del("mcpclient:*") + } + c, _ := store.Get("__yao.cache") + if c != nil { + c.Del("mcpclient:*") + } + test.Clean() + }) + + return mcpclient.Global +} + +func newTestClient(id string) mcpclient.Client { + return mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: id, + Name: "Test " + id, + Type: "standard", + Transport: mcpTypes.TransportStdio, + Command: "echo", + Arguments: []string{"hello"}, + }, + Enabled: true, + Owner: mcpclient.ClientOwner{Type: "system"}, + } +} + +func TestCreate(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-stdio") + created, err := r.Create(&c) + require.NoError(t, err) + assert.Equal(t, "test-stdio", created.ID) + assert.Equal(t, mcpclient.ClientSourceDynamic, created.Source) + assert.NotEmpty(t, created.RuntimeID) + + s, _ := store.Get("__yao.store") + assert.True(t, s.Has("mcpclient:c:test-stdio")) +} + +func TestCreateDuplicate(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-dup") + _, err := r.Create(&c) + require.NoError(t, err) + + dup := newTestClient("test-dup") + _, err = r.Create(&dup) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already exists") +} + +func TestGet(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-get") + _, err := r.Create(&c) + require.NoError(t, err) + + got, err := r.Get("test-get") + require.NoError(t, err) + assert.Equal(t, "Test test-get", got.Name) + assert.Equal(t, mcpTypes.TransportStdio, got.Transport) + assert.Equal(t, "echo", got.Command) +} + +func TestGetNotFound(t *testing.T) { + r := setupRegistry(t) + _, err := r.Get("nonexistent") + assert.Error(t, err) +} + +func TestGetLazy(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-lazy") + created, err := r.Create(&c) + require.NoError(t, err) + + // Manually unload the client + mcp.UnloadClient(created.RuntimeID) + assert.False(t, mcp.Exists(created.RuntimeID)) + + // Get should lazily re-register + got, err := r.Get("test-lazy") + require.NoError(t, err) + assert.Equal(t, "test-lazy", got.ID) +} + +func TestList(t *testing.T) { + r := setupRegistry(t) + + clients := []mcpclient.Client{ + { + ClientDSL: mcpTypes.ClientDSL{ID: "c1", Name: "Client 1", Type: "standard", Transport: mcpTypes.TransportStdio, Command: "echo"}, + Enabled: true, + Owner: mcpclient.ClientOwner{Type: "system"}, + }, + { + ClientDSL: mcpTypes.ClientDSL{ID: "c2", Name: "Client 2", Type: "agent", Transport: mcpTypes.TransportSSE, URL: "http://localhost:3001"}, + Enabled: false, + Owner: mcpclient.ClientOwner{Type: "user", ID: "123"}, + }, + { + ClientDSL: mcpTypes.ClientDSL{ID: "c3", Name: "Client 3", Type: "standard", Transport: mcpTypes.TransportStdio, Command: "cat"}, + Enabled: true, + Owner: mcpclient.ClientOwner{Type: "system"}, + }, + } + for i := range clients { + _, err := r.Create(&clients[i]) + require.NoError(t, err) + } + + t.Run("AllDynamic", func(t *testing.T) { + list, err := r.List(&mcpclient.ClientFilter{Source: mcpclient.ClientSourceDynamic}) + require.NoError(t, err) + assert.GreaterOrEqual(t, len(list), 3) + }) + + t.Run("FilterByTransport", func(t *testing.T) { + tp := mcpTypes.TransportSSE + list, err := r.List(&mcpclient.ClientFilter{ + Source: mcpclient.ClientSourceDynamic, + Transport: &tp, + }) + require.NoError(t, err) + for _, c := range list { + assert.Equal(t, mcpTypes.TransportSSE, c.Transport) + } + }) + + t.Run("FilterByEnabled", func(t *testing.T) { + enabled := true + list, err := r.List(&mcpclient.ClientFilter{ + Source: mcpclient.ClientSourceDynamic, + Enabled: &enabled, + }) + require.NoError(t, err) + for _, c := range list { + assert.True(t, c.Enabled) + } + }) + + t.Run("FilterByOwner", func(t *testing.T) { + list, err := r.List(&mcpclient.ClientFilter{ + Source: mcpclient.ClientSourceDynamic, + Owner: &mcpclient.ClientOwner{Type: "user", ID: "123"}, + }) + require.NoError(t, err) + for _, c := range list { + assert.Equal(t, "user", c.Owner.Type) + } + }) + + t.Run("FilterByType", func(t *testing.T) { + typ := "agent" + list, err := r.List(&mcpclient.ClientFilter{ + Source: mcpclient.ClientSourceDynamic, + Type: &typ, + }) + require.NoError(t, err) + for _, c := range list { + assert.Equal(t, "agent", c.ClientDSL.Type) + } + }) + + t.Run("FilterByKeyword", func(t *testing.T) { + list, err := r.List(&mcpclient.ClientFilter{ + Source: mcpclient.ClientSourceDynamic, + Keyword: "Client 2", + }) + require.NoError(t, err) + found := false + for _, c := range list { + if c.ID == "c2" { + found = true + } + } + assert.True(t, found) + }) +} + +func TestUpdate(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-update") + _, err := r.Create(&c) + require.NoError(t, err) + + got, err := r.Get("test-update") + require.NoError(t, err) + + updated := *got + updated.ClientDSL.Name = "Updated Name" + updated.ClientDSL.Command = "cat" + + result, err := r.Update("test-update", &updated) + require.NoError(t, err) + assert.Equal(t, "Updated Name", result.Name) + + got2, err := r.Get("test-update") + require.NoError(t, err) + assert.Equal(t, "cat", got2.Command) +} + +func TestDelete(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-delete") + _, err := r.Create(&c) + require.NoError(t, err) + + err = r.Delete("test-delete") + require.NoError(t, err) + + _, err = r.Get("test-delete") + assert.Error(t, err) +} + +func TestReload(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-reload") + _, err := r.Create(&c) + require.NoError(t, err) + + // Clear cache + cache, _ := store.Get("__yao.cache") + if cache != nil { + cache.Del("mcpclient:*") + } + + err = r.Reload() + require.NoError(t, err) + + got, err := r.Get("test-reload") + require.NoError(t, err) + assert.Equal(t, "Test test-reload", got.Name) +} + +func TestImportFromClients(t *testing.T) { + r := setupRegistry(t) + + list, err := r.List(&mcpclient.ClientFilter{Source: mcpclient.ClientSourceAll}) + require.NoError(t, err) + + builtinCount := 0 + for _, c := range list { + if c.Source == mcpclient.ClientSourceBuiltIn { + builtinCount++ + } + } + + loadedClients := mcp.ListClients() + t.Logf("Imported %d builtin clients from mcp.ListClients (total loaded: %d)", builtinCount, len(loadedClients)) +} + +func TestToolListField(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-toollist") + c.ClientDSL.Tools = map[string]string{"my-tool": "scripts.MyTool"} + c.ToolList = []mcpTypes.Tool{ + {Name: "discovered-tool", Description: "A tool discovered at runtime"}, + } + + created, err := r.Create(&c) + require.NoError(t, err) + + got, err := r.Get(created.ID) + require.NoError(t, err) + assert.Len(t, got.ToolList, 1) + assert.Equal(t, "discovered-tool", got.ToolList[0].Name) + assert.Equal(t, "scripts.MyTool", got.ClientDSL.Tools["my-tool"]) +} + +func TestGetMCPClient(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-getmcp") + _, err := r.Create(&c) + require.NoError(t, err) + + // The MCP client may or may not actually start (depends on whether "echo" is a valid MCP server), + // but we should at least exercise the code path. + _, err = r.GetMCPClient("test-getmcp") + // Either it works or returns a "not found" — both are valid for this test fixture + t.Logf("GetMCPClient result: err=%v", err) +} + +func TestGetMCPClientNotFound(t *testing.T) { + r := setupRegistry(t) + _, err := r.GetMCPClient("no-such-client") + assert.Error(t, err) +} + +func TestCreateEmptyID(t *testing.T) { + r := setupRegistry(t) + c := mcpclient.Client{ClientDSL: mcpTypes.ClientDSL{Name: "No ID"}} + _, err := r.Create(&c) + assert.Error(t, err) + assert.Contains(t, err.Error(), "id is required") +} + +func TestCreateDisabled(t *testing.T) { + r := setupRegistry(t) + c := mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ID: "test-disabled", Name: "Disabled", Type: "standard", Transport: mcpTypes.TransportStdio, Command: "echo"}, + Enabled: false, + Owner: mcpclient.ClientOwner{Type: "system"}, + } + created, err := r.Create(&c) + require.NoError(t, err) + assert.Equal(t, "unconfigured", created.Status) + + // Disabled client should not be registered at runtime + assert.False(t, mcp.Exists(created.RuntimeID), "disabled client should not be registered") +} + +func TestOwnerPrefixedRuntimeIDs(t *testing.T) { + r := setupRegistry(t) + + cases := []struct { + id string + owner mcpclient.ClientOwner + prefix string + }{ + {"owner-sys", mcpclient.ClientOwner{Type: "system"}, "s."}, + {"owner-usr", mcpclient.ClientOwner{Type: "user", ID: "42"}, "u42."}, + {"owner-team", mcpclient.ClientOwner{Type: "team", ID: "99"}, "t99."}, + {"owner-asst", mcpclient.ClientOwner{Type: "assistant", ID: "a1"}, "aa1."}, + } + + for _, tc := range cases { + t.Run(tc.id, func(t *testing.T) { + c := mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ID: tc.id, Name: tc.id, Type: "standard", Transport: mcpTypes.TransportStdio, Command: "echo"}, + Enabled: true, + Owner: tc.owner, + } + created, err := r.Create(&c) + require.NoError(t, err) + assert.Contains(t, created.RuntimeID, tc.prefix, + "RuntimeID for %s owner should contain prefix %s", tc.owner.Type, tc.prefix) + }) + } +} + +func TestListBuiltInFilter(t *testing.T) { + r := setupRegistry(t) + + builtinList, err := r.List(&mcpclient.ClientFilter{Source: mcpclient.ClientSourceBuiltIn}) + require.NoError(t, err) + for _, c := range builtinList { + assert.Equal(t, mcpclient.ClientSourceBuiltIn, c.Source) + } +} + +func TestListAllSources(t *testing.T) { + r := setupRegistry(t) + + c := newTestClient("test-all-src") + _, err := r.Create(&c) + require.NoError(t, err) + + all, err := r.List(&mcpclient.ClientFilter{Source: mcpclient.ClientSourceAll}) + require.NoError(t, err) + + hasDynamic := false + for _, item := range all { + if item.Source == mcpclient.ClientSourceDynamic { + hasDynamic = true + } + } + assert.True(t, hasDynamic) +} + +func TestUpdateNotFound(t *testing.T) { + r := setupRegistry(t) + c := newTestClient("not-exist") + _, err := r.Update("not-exist", &c) + assert.Error(t, err) +} + +func TestDeleteNotFound(t *testing.T) { + r := setupRegistry(t) + err := r.Delete("not-exist") + assert.Error(t, err) +} + +func TestConcurrency(t *testing.T) { + r := setupRegistry(t) + + var wg sync.WaitGroup + errCh := make(chan error, 30) + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + c := mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: fmt.Sprintf("conc-%d", idx), + Name: fmt.Sprintf("Concurrent %d", idx), + Type: "standard", + Transport: mcpTypes.TransportStdio, + Command: "echo", + }, + Enabled: true, + Owner: mcpclient.ClientOwner{Type: "system"}, + } + if _, err := r.Create(&c); err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + _, err := r.Get(fmt.Sprintf("conc-%d", idx)) + if err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + if err := r.Delete(fmt.Sprintf("conc-%d", idx)); err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + close(errCh) + + for err := range errCh { + t.Errorf("concurrent operation error: %v", err) + } +} diff --git a/mcpclient/store.go b/mcpclient/store.go new file mode 100644 index 00000000..93cebbbc --- /dev/null +++ b/mcpclient/store.go @@ -0,0 +1,179 @@ +package mcpclient + +import ( + "encoding/json" + "fmt" + + "github.com/yaoapp/gou/store" +) + +const ( + keyPrefix = "mcpclient:c:" + indexKey = "mcpclient:index" +) + +func storeKey(id string) string { return keyPrefix + id } + +func clientToMap(c *Client) (map[string]interface{}, error) { + raw, err := json.Marshal(c) + if err != nil { + return nil, err + } + var m map[string]interface{} + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil +} + +func mapToClient(m map[string]interface{}) (*Client, error) { + raw, err := json.Marshal(m) + if err != nil { + return nil, err + } + var c Client + if err := json.Unmarshal(raw, &c); err != nil { + return nil, err + } + return &c, nil +} + +func storeGet(s, c store.Store, id string) (*Client, error) { + sk := storeKey(id) + + if c != nil { + if val, ok := c.Get(sk); ok { + if m, ok := val.(map[string]interface{}); ok { + return mapToClient(m) + } + } + } + + val, ok := s.Get(sk) + if !ok { + return nil, fmt.Errorf("client %s not found", id) + } + m, ok := val.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("client %s: unexpected store type %T", id, val) + } + + cl, err := mapToClient(m) + if err != nil { + return nil, err + } + + if c != nil { + c.Set(sk, m, 0) + } + return cl, nil +} + +func storeSet(s, c store.Store, cl *Client) error { + m, err := clientToMap(cl) + if err != nil { + return err + } + sk := storeKey(cl.ID) + if err := s.Set(sk, m, 0); err != nil { + return err + } + if c != nil { + c.Set(sk, m, 0) + } + return nil +} + +func storeDel(s, c store.Store, id string) error { + sk := storeKey(id) + if err := s.Del(sk); err != nil { + return err + } + if c != nil { + c.Del(sk) + } + return nil +} + +func indexGet(s, c store.Store) ([]string, error) { + var raw interface{} + var ok bool + + if c != nil { + raw, ok = c.Get(indexKey) + } + if !ok { + raw, ok = s.Get(indexKey) + if !ok { + return nil, nil + } + if c != nil { + c.Set(indexKey, raw, 0) + } + } + + switch v := raw.(type) { + case []interface{}: + keys := make([]string, 0, len(v)) + for _, item := range v { + if str, ok := item.(string); ok { + keys = append(keys, str) + } + } + return keys, nil + case []string: + return v, nil + default: + return nil, fmt.Errorf("unexpected index type %T", raw) + } +} + +func indexSet(s, c store.Store, ids []string) error { + iface := make([]interface{}, len(ids)) + for i, k := range ids { + iface[i] = k + } + if err := s.Set(indexKey, iface, 0); err != nil { + return err + } + if c != nil { + c.Set(indexKey, iface, 0) + } + return nil +} + +func indexAdd(s, c store.Store, id string) error { + ids, err := indexGet(s, c) + if err != nil { + return err + } + for _, k := range ids { + if k == id { + return nil + } + } + return indexSet(s, c, append(ids, id)) +} + +func indexRemove(s, c store.Store, id string) error { + ids, err := indexGet(s, c) + if err != nil { + return err + } + filtered := make([]string, 0, len(ids)) + for _, k := range ids { + if k != id { + filtered = append(filtered, k) + } + } + return indexSet(s, c, filtered) +} + +func storeCleanAll(s, c store.Store) { + _ = s.Del(keyPrefix + "*") + _ = s.Del(indexKey) + if c != nil { + _ = c.Del(keyPrefix + "*") + _ = c.Del(indexKey) + } +} diff --git a/mcpclient/sync.go b/mcpclient/sync.go new file mode 100644 index 00000000..0056d951 --- /dev/null +++ b/mcpclient/sync.go @@ -0,0 +1,136 @@ +package mcpclient + +import ( + "encoding/json" + "fmt" + + "github.com/yaoapp/gou/mcp" + mcpTypes "github.com/yaoapp/gou/mcp/types" +) + +// runtimeID builds the runtime ID for registering into mcp.clients. +// Dynamic clients get an owner prefix to avoid collision with builtin IDs. +func runtimeID(c *Client) string { + switch c.Owner.Type { + case "user": + return "u" + c.Owner.ID + "." + c.ID + case "team": + return "t" + c.Owner.ID + "." + c.ID + case "assistant": + return "a" + c.Owner.ID + "." + c.ID + default: + return "s." + c.ID + } +} + +// ensureClient makes sure the MCP client is registered in the runtime. +// Builtin clients are managed by engine.Load and skipped here. +func ensureClient(c *Client) error { + if c.Source == ClientSourceBuiltIn { + return nil + } + if !c.Enabled { + return nil + } + + rid := c.RuntimeID + if rid == "" { + rid = runtimeID(c) + } + + if mcp.Exists(rid) { + return nil + } + + dslJSON, err := json.Marshal(c.ClientDSL) + if err != nil { + return fmt.Errorf("ensureClient %s: marshal DSL: %w", c.ID, err) + } + + clientType := c.ClientDSL.Type + _, err = mcp.LoadClientSourceWithType(string(dslJSON), rid, clientType) + if err != nil { + return fmt.Errorf("ensureClient %s: LoadClientSourceWithType: %w", c.ID, err) + } + + return nil +} + +// unloadClient removes the client from the runtime. +func unloadClient(c *Client) { + if c.Source == ClientSourceBuiltIn { + return + } + rid := c.RuntimeID + if rid == "" { + rid = runtimeID(c) + } + mcp.UnloadClient(rid) +} + +// importFromClients scans existing MCP clients loaded by engine.Load +// and imports them as builtin entries into the Registry store. +// If a store record with the same ID already exists (dynamic), it is not overwritten. +func importFromClients(r *Registry) error { + ids := mcp.ListClients() + for _, id := range ids { + if r.store.Has(storeKey(id)) { + continue + } + + cl := clientFromRuntime(id) + if cl == nil { + continue + } + + m, err := clientToMap(cl) + if err != nil { + continue + } + sk := storeKey(id) + _ = r.store.Set(sk, m, 0) + if r.cache != nil { + _ = r.cache.Set(sk, m, 0) + } + _ = indexAdd(r.store, r.cache, id) + } + return nil +} + +// clientFromRuntime builds a Client from a runtime mcp.Client interface. +// Uses Info() and GetMetaInfo() since full ClientDSL is not exposed. +func clientFromRuntime(id string) *Client { + defer func() { recover() }() + + mcpClient := mcp.GetClient(id) + if mcpClient == nil { + return nil + } + + info := mcpClient.Info() + if info == nil { + return nil + } + + meta := mcpClient.GetMetaInfo() + + name := info.Name + if name == "" { + name = id + } + + return &Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: id, + Name: name, + Type: info.Type, + Transport: info.Transport, + MetaInfo: meta, + }, + RuntimeID: id, + Enabled: true, + Status: "connected", + Source: ClientSourceBuiltIn, + Owner: ClientOwner{Type: "system"}, + } +} diff --git a/mcpclient/types.go b/mcpclient/types.go new file mode 100644 index 00000000..61e6a3cf --- /dev/null +++ b/mcpclient/types.go @@ -0,0 +1,50 @@ +package mcpclient + +import ( + mcpTypes "github.com/yaoapp/gou/mcp/types" +) + +// Client wraps mcpTypes.ClientDSL with Registry management fields. +// Uses ClientDSL.ID as the registry key. +type Client struct { + mcpTypes.ClientDSL + + RuntimeID string `json:"runtime_id"` + Enabled bool `json:"enabled"` + Status string `json:"status"` + Source ClientSource `json:"source"` + ToolList []mcpTypes.Tool `json:"tool_list,omitempty"` + Owner ClientOwner `json:"owner"` +} + +// ClientOwner identifies who owns a client entry. +type ClientOwner struct { + Type string `json:"type"` + ID string `json:"id,omitempty"` +} + +// ClientSource distinguishes registry-created from DSL-loaded clients. +type ClientSource string + +const ( + ClientSourceDynamic ClientSource = "dynamic" + ClientSourceBuiltIn ClientSource = "builtin" + ClientSourceAll ClientSource = "all" +) + +// ClientFilter specifies criteria for listing clients. +type ClientFilter struct { + Owner *ClientOwner + Enabled *bool + Source ClientSource + Transport *mcpTypes.TransportType + Type *string + Keyword string +} + +// ClientTestResult holds the outcome of a client connectivity test. +type ClientTestResult struct { + Success bool `json:"success"` + Message string `json:"message"` + LatencyMs int64 `json:"latency_ms,omitempty"` +} diff --git a/openapi/openapi.go b/openapi/openapi.go index 49f2c78f..c00cec4d 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -27,6 +27,7 @@ import ( "github.com/yaoapp/yao/openapi/otp" "github.com/yaoapp/yao/openapi/response" "github.com/yaoapp/yao/openapi/sandbox" + openAPISetting "github.com/yaoapp/yao/openapi/setting" openapiTai "github.com/yaoapp/yao/openapi/tai" "github.com/yaoapp/yao/openapi/team" openapiTrace "github.com/yaoapp/yao/openapi/trace" @@ -199,6 +200,9 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { group.POST("/tai-nodes/heartbeat", taiapi.HandleHeartbeat) group.DELETE("/tai-nodes/register/:tai_id", taiapi.HandleUnregister) + // Setting handlers (unified /setting/* endpoints) + openAPISetting.Attach(group.Group("/setting"), openapi.OAuth) + // Custom handlers (Defined by developer) } diff --git a/openapi/setting/cloud.go b/openapi/setting/cloud.go new file mode 100644 index 00000000..4194a81e --- /dev/null +++ b/openapi/setting/cloud.go @@ -0,0 +1,409 @@ +package setting + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha256" + _ "embed" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/setting" + "gopkg.in/yaml.v3" +) + +//go:embed cloud_presets.yml +var cloudPresetsYML []byte + +const ( + cloudNS = "cloud" + cloudMaskChars = 4 + cloudEncPrefix = "enc:" +) + +// cloudPresets holds the parsed region list from the embedded YAML. +type cloudPresets struct { + Regions []CloudRegion `yaml:"regions"` +} + +var cloudRegions []CloudRegion + +func init() { + var p cloudPresets + if err := yaml.Unmarshal(cloudPresetsYML, &p); err == nil { + cloudRegions = p.Regions + } +} + +func cloudDefaultRegion() CloudRegion { + for _, r := range cloudRegions { + if r.Default { + return r + } + } + if len(cloudRegions) > 0 { + return cloudRegions[0] + } + return CloudRegion{Key: "us", APIURL: "https://api-us.yao.run"} +} + +func cloudFindRegion(key string) *CloudRegion { + for i := range cloudRegions { + if cloudRegions[i].Key == key { + return &cloudRegions[i] + } + } + return nil +} + +func cloudScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID { + if info.TeamID != "" { + return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID} + } + return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID} +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +// handleCloudGet returns the cloud configuration for the current team. +// GET /setting/cloud +func handleCloudGet(c *gin.Context) { + info := authorized.GetInfo(c) + def := cloudDefaultRegion() + + var saved map[string]interface{} + if setting.Global != nil { + saved, _ = setting.Global.GetMerged(info.UserID, info.TeamID, cloudNS) + } + + data := CloudPageData{ + Regions: cloudRegions, + Region: def.Key, + APIURL: def.APIURL, + APIKey: "", + Status: "unconfigured", + } + + if saved != nil { + if v, ok := saved["region"].(string); ok && v != "" { + data.Region = v + } + if v, ok := saved["api_url"].(string); ok && v != "" { + data.APIURL = v + } + if v, ok := saved["api_key"].(string); ok && v != "" { + data.APIKey = cloudMaskKey(cloudDecrypt(v)) + } + if v, ok := saved["status"].(string); ok && v != "" { + data.Status = v + } + } + + response.RespondWithSuccess(c, http.StatusOK, data) +} + +// handleCloudUpdate saves the cloud configuration. +// When api_key is provided, validates it by calling the cloud API before saving. +// PUT /setting/cloud +func handleCloudUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := cloudScope(info) + + var body struct { + Region string `json:"region"` + APIURL string `json:"api_url"` + APIKey string `json:"api_key"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if body.Region != "" { + if r := cloudFindRegion(body.Region); r == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown region: %s", body.Region)) + return + } + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + existing, _ := setting.Global.Get(scope, cloudNS) + + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + + if body.Region != "" { + m["region"] = body.Region + } + if body.APIURL != "" { + m["api_url"] = body.APIURL + } + + // Resolve the effective api_url for key validation + apiURL := body.APIURL + if apiURL == "" { + if v, ok := m["api_url"].(string); ok { + apiURL = v + } + } + if apiURL == "" { + if body.Region != "" { + if r := cloudFindRegion(body.Region); r != nil { + apiURL = r.APIURL + } + } + if apiURL == "" { + apiURL = cloudDefaultRegion().APIURL + } + } + + if body.APIKey != "" { + if err := cloudValidateKey(apiURL, body.APIKey); err != nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("API key validation failed: %s", err.Error())) + return + } + m["api_key"] = cloudEncrypt(body.APIKey) + m["status"] = "connected" + } + + hasKey := false + if v, ok := m["api_key"].(string); ok && v != "" { + hasKey = true + } + if _, ok := m["status"].(string); !ok { + if hasKey { + m["status"] = "disconnected" + } else { + m["status"] = "unconfigured" + } + } + + if _, err := setting.Global.Set(scope, cloudNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + def := cloudDefaultRegion() + result := CloudPageData{ + Regions: cloudRegions, + Region: def.Key, + APIURL: def.APIURL, + APIKey: "", + Status: "unconfigured", + } + if v, ok := m["region"].(string); ok && v != "" { + result.Region = v + } + if v, ok := m["api_url"].(string); ok && v != "" { + result.APIURL = v + } + if v, ok := m["api_key"].(string); ok && v != "" { + result.APIKey = cloudMaskKey(cloudDecrypt(v)) + } + if v, ok := m["status"].(string); ok && v != "" { + result.Status = v + } + + response.RespondWithSuccess(c, http.StatusOK, result) +} + +// cloudValidateKey verifies the API key by calling GET {apiURL}/v1/models. +func cloudValidateKey(apiURL, apiKey string) error { + url := strings.TrimRight(apiURL, "/") + "/v1/models" + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return fmt.Errorf("failed to build request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+apiKey) + + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +// handleCloudTest tests the cloud connection by calling GET {api_url}/v1/models. +// Caller must provide api_url and api_key in the request body. +// POST /setting/cloud/test +func handleCloudTest(c *gin.Context) { + if !guardOwner(c) { + return + } + + var input struct { + APIURL string `json:"api_url"` + APIKey string `json:"api_key"` + } + if err := c.ShouldBindJSON(&input); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if input.APIURL == "" || input.APIKey == "" { + respondError(c, http.StatusBadRequest, "api_url and api_key are required") + return + } + + url := strings.TrimRight(input.APIURL, "/") + "/v1/models" + + start := time.Now() + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", url, nil) + if err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + req.Header.Set("Authorization", "Bearer "+input.APIKey) + + resp, err := client.Do(req) + latency := time.Since(start).Milliseconds() + + if err != nil { + response.RespondWithSuccess(c, http.StatusOK, CloudTestResult{ + Success: false, + Message: fmt.Sprintf("Connection failed: %s", err.Error()), + }) + return + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + response.RespondWithSuccess(c, http.StatusOK, CloudTestResult{ + Success: false, + Message: fmt.Sprintf("Server returned HTTP %d", resp.StatusCode), + }) + return + } + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + + response.RespondWithSuccess(c, http.StatusOK, CloudTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latency, + }) +} + +// --------------------------------------------------------------------------- +// Crypto helpers (AES-256-GCM, same scheme as llmprovider) +// --------------------------------------------------------------------------- + +func cloudEncrypt(plaintext string) string { + secret := config.Conf.DB.AESKey + if secret == "" { + return plaintext + } + enc, err := cloudEncryptString(plaintext, secret) + if err != nil { + return plaintext + } + return cloudEncPrefix + enc +} + +func cloudDecrypt(value string) string { + if !strings.HasPrefix(value, cloudEncPrefix) { + return value + } + secret := config.Conf.DB.AESKey + if secret == "" { + return strings.TrimPrefix(value, cloudEncPrefix) + } + dec, err := cloudDecryptString(strings.TrimPrefix(value, cloudEncPrefix), secret) + if err != nil { + return value + } + return dec +} + +func cloudMaskKey(key string) string { + if key == "" { + return "" + } + if len(key) <= cloudMaskChars { + return strings.Repeat("*", len(key)) + } + prefix := key[:3] + suffix := key[len(key)-cloudMaskChars:] + return prefix + "..." + suffix +} + +func cloudDeriveKey(secret string) []byte { + h := sha256.Sum256([]byte(secret)) + return h[:] +} + +func cloudEncryptString(plaintext, secret string) (string, error) { + key := cloudDeriveKey(secret) + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return "", err + } + ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil) + return base64.StdEncoding.EncodeToString(ciphertext), nil +} + +func cloudDecryptString(encoded, secret string) (string, error) { + key := cloudDeriveKey(secret) + data, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", err + } + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonceSize := gcm.NonceSize() + if len(data) < nonceSize { + return "", fmt.Errorf("ciphertext too short") + } + plaintext, err := gcm.Open(nil, data[:nonceSize], data[nonceSize:], nil) + if err != nil { + return "", err + } + return string(plaintext), nil +} diff --git a/openapi/setting/cloud_presets.yml b/openapi/setting/cloud_presets.yml new file mode 100644 index 00000000..5b7f04c5 --- /dev/null +++ b/openapi/setting/cloud_presets.yml @@ -0,0 +1,28 @@ +# Cloud service region presets. +# Embedded at compile time via go:embed in cloud.go. + +regions: + - key: us + label: + zh-CN: "美国" + en-US: "United States" + api_url: "https://api-us.yao.run" + default: true + + - key: cn + label: + zh-CN: "中国" + en-US: "China" + api_url: "https://api.yaoagents.cn" + + - key: ap + label: + zh-CN: "亚太" + en-US: "Asia Pacific" + api_url: "https://api-ap.yao.run" + + - key: eu + label: + zh-CN: "欧洲" + en-US: "Europe" + api_url: "https://api-eu.yao.run" diff --git a/openapi/setting/llm.go b/openapi/setting/llm.go new file mode 100644 index 00000000..566028e2 --- /dev/null +++ b/openapi/setting/llm.go @@ -0,0 +1,668 @@ +package setting + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/llmprovider" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/setting" +) + +const llmRolesNS = "llm.roles" + +func llmEnsureEncKey() { + if llmprovider.Global != nil && config.Conf.DB.AESKey != "" { + llmprovider.Global.SetEncryptionKey(config.Conf.DB.AESKey) + } +} + +func llmOwner(info *oauthTypes.AuthorizedInfo) *llmprovider.ProviderOwner { + if info.TeamID != "" { + return &llmprovider.ProviderOwner{Type: "team", TeamID: info.TeamID} + } + return &llmprovider.ProviderOwner{Type: "user", UserID: info.UserID} +} + +func llmScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID { + if info.TeamID != "" { + return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID} + } + return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID} +} + +func llmCheckOwnership(p *llmprovider.Provider, info *oauthTypes.AuthorizedInfo) error { + owner := llmOwner(info) + if p.Owner.Type != owner.Type { + return fmt.Errorf("provider not found") + } + if owner.Type == "team" && p.Owner.TeamID != owner.TeamID { + return fmt.Errorf("provider not found") + } + if owner.Type == "user" && p.Owner.UserID != owner.UserID { + return fmt.Errorf("provider not found") + } + return nil +} + +func enrichProvider(p *llmprovider.Provider) map[string]interface{} { + raw, _ := json.Marshal(p) + var m map[string]interface{} + json.Unmarshal(raw, &m) + + if p.PresetKey != "" { + if preset := llmprovider.GetPreset(p.PresetKey); preset != nil { + m["is_cloud"] = preset.IsCloud + m["url_editable"] = preset.URLEditable + } + } + + delete(m, "connector_id") + delete(m, "source") + delete(m, "owner") + + return m +} + +// llmModelsURL builds the models endpoint URL. +// Trailing slash means the user already specified the path prefix → append "models". +// No trailing slash → append "/v1/models" (standard OpenAI convention). +func llmModelsURL(apiURL string) string { + if strings.HasSuffix(apiURL, "/") { + return apiURL + "models" + } + return apiURL + "/v1/models" +} + +// llmValidateKey tests connectivity by calling GET {apiURL}/models. +// providerType controls the auth header format (anthropic uses x-api-key). +func llmValidateKey(providerType, apiURL, apiKey string) error { + url := llmModelsURL(apiURL) + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return fmt.Errorf("failed to build request: %w", err) + } + if apiKey != "" { + if providerType == "anthropic" { + req.Header.Set("x-api-key", apiKey) + req.Header.Set("anthropic-version", "2023-06-01") + } else { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + } + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +// handleLLMTest validates an API URL + Key without saving. +// POST /setting/llm/test +func handleLLMTest(c *gin.Context) { + if !guardOwner(c) { + return + } + + var input struct { + APIURL string `json:"api_url"` + APIKey string `json:"api_key"` + Type string `json:"type"` + } + if err := c.ShouldBindJSON(&input); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + if input.APIURL == "" { + respondError(c, http.StatusBadRequest, "api_url is required") + return + } + + url := llmModelsURL(input.APIURL) + start := time.Now() + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", url, nil) + if err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + if input.APIKey != "" { + if input.Type == "anthropic" { + req.Header.Set("x-api-key", input.APIKey) + req.Header.Set("anthropic-version", "2023-06-01") + } else { + req.Header.Set("Authorization", "Bearer "+input.APIKey) + } + } + + resp, err := client.Do(req) + latency := time.Since(start).Milliseconds() + + if err != nil { + response.RespondWithSuccess(c, http.StatusOK, llmprovider.ProviderTestResult{ + Success: false, + Message: fmt.Sprintf("Connection failed: %s", err.Error()), + }) + return + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + response.RespondWithSuccess(c, http.StatusOK, llmprovider.ProviderTestResult{ + Success: false, + Message: fmt.Sprintf("Server returned HTTP %d", resp.StatusCode), + }) + return + } + + response.RespondWithSuccess(c, http.StatusOK, llmprovider.ProviderTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latency, + }) +} + +// handleLLMGet returns the aggregated LLM configuration page data. +// GET /setting/llm +func handleLLMGet(c *gin.Context) { + info := authorized.GetInfo(c) + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + owner := llmOwner(info) + filter := &llmprovider.ProviderFilter{ + Owner: owner, + Source: llmprovider.ProviderSourceAll, + } + providers, err := llmprovider.Global.List(filter) + if err != nil { + providers = []llmprovider.Provider{} + } + + enriched := make([]interface{}, 0, len(providers)) + for i := range providers { + enriched = append(enriched, enrichProvider(&providers[i])) + } + + var roles map[string]interface{} + if setting.Global != nil { + roles, _ = setting.Global.GetMerged(info.UserID, info.TeamID, llmRolesNS) + } + if roles == nil { + roles = make(map[string]interface{}) + } + + presetList := llmprovider.GetPresets() + presetIface := make([]interface{}, len(presetList)) + for i, p := range presetList { + raw, _ := json.Marshal(p) + var m map[string]interface{} + json.Unmarshal(raw, &m) + presetIface[i] = m + } + + response.RespondWithSuccess(c, http.StatusOK, LLMPageData{ + Providers: enriched, + Roles: roles, + PresetProviders: presetIface, + }) +} + +// handleLLMRoles saves the role assignment (default models). +// PUT /setting/llm/roles +func handleLLMRoles(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := llmScope(info) + + var body map[string]interface{} + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if _, ok := body["default"]; !ok { + respondError(c, http.StatusBadRequest, "\"default\" role is required") + return + } + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + for roleName, target := range body { + targetMap, ok := target.(map[string]interface{}) + if !ok { + respondError(c, http.StatusBadRequest, fmt.Sprintf("invalid target for role \"%s\"", roleName)) + return + } + + providerKey, _ := targetMap["provider"].(string) + modelID, _ := targetMap["model"].(string) + if providerKey == "" || modelID == "" { + respondError(c, http.StatusBadRequest, fmt.Sprintf("role \"%s\" requires provider and model", roleName)) + return + } + + p, err := llmprovider.Global.Get(providerKey) + if err != nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("provider \"%s\" not found", providerKey)) + return + } + if !p.Enabled { + respondError(c, http.StatusBadRequest, fmt.Sprintf("provider \"%s\" is not enabled", providerKey)) + return + } + if err := llmCheckOwnership(p, info); err != nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("provider \"%s\" not found", providerKey)) + return + } + + modelFound := false + for _, m := range p.Models { + if m.ID == modelID { + modelFound = true + break + } + } + if !modelFound { + respondError(c, http.StatusBadRequest, fmt.Sprintf("model \"%s\" not found in provider \"%s\"", modelID, providerKey)) + return + } + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + if _, err := setting.Global.Set(scope, llmRolesNS, body); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + response.RespondWithSuccess(c, http.StatusOK, body) +} + +// handleLLMProviderCreate creates a new LLM provider (preset or custom). +// POST /setting/llm/providers +func handleLLMProviderCreate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + var body map[string]interface{} + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + var provider llmprovider.Provider + owner := llmOwner(info) + provider.Owner = *owner + provider.Source = llmprovider.ProviderSourceDynamic + provider.Enabled = true + + presetKey, _ := body["preset_key"].(string) + + if presetKey != "" { + preset := llmprovider.GetPreset(presetKey) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown preset: %s", presetKey)) + return + } + + provider.Key = presetKey + provider.Name = preset.Name + provider.Type = preset.Type + provider.APIURL = preset.APIURL + provider.RequireKey = preset.RequireKey + provider.PresetKey = presetKey + + if v, ok := body["api_url"].(string); ok && v != "" { + provider.APIURL = v + } + if v, ok := body["api_key"].(string); ok && v != "" { + provider.APIKey = v + } + if v, ok := body["name"].(string); ok && v != "" { + provider.Name = v + } + + modelIDs, hasModelIDs := body["model_ids"].([]interface{}) + if hasModelIDs && len(modelIDs) > 0 { + idSet := make(map[string]bool, len(modelIDs)) + for _, id := range modelIDs { + if s, ok := id.(string); ok { + idSet[s] = true + } + } + for _, m := range preset.DefaultModels { + if idSet[m.ID] { + provider.Models = append(provider.Models, m) + } + } + } else { + provider.Models = make([]llmprovider.ModelInfo, len(preset.DefaultModels)) + copy(provider.Models, preset.DefaultModels) + } + } else { + provider.IsCustom = true + + key, _ := body["key"].(string) + if key == "" { + respondError(c, http.StatusBadRequest, "key is required for custom provider") + return + } + provider.Key = key + + name, _ := body["name"].(string) + if name == "" { + respondError(c, http.StatusBadRequest, "name is required") + return + } + provider.Name = name + + typ, _ := body["type"].(string) + if typ == "" { + typ = "openai" + } + provider.Type = typ + + provider.APIURL, _ = body["api_url"].(string) + provider.APIKey, _ = body["api_key"].(string) + + if modelsRaw, ok := body["models"]; ok { + raw, _ := json.Marshal(modelsRaw) + var models []llmprovider.ModelInfo + if err := json.Unmarshal(raw, &models); err == nil { + provider.Models = models + } + } + + if v, ok := body["require_key"].(bool); ok { + provider.RequireKey = v + } + } + + if provider.Models == nil { + provider.Models = []llmprovider.ModelInfo{} + } + + if provider.RequireKey && provider.APIKey != "" && provider.APIURL != "" { + if err := llmValidateKey(provider.Type, provider.APIURL, provider.APIKey); err != nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("API key validation failed: %s", err.Error())) + return + } + } + + created, err := llmprovider.Global.Create(&provider) + if err != nil { + if strings.Contains(err.Error(), "already exists") { + respondError(c, http.StatusConflict, err.Error()) + } else { + respondError(c, http.StatusInternalServerError, err.Error()) + } + return + } + + masked, err := llmprovider.Global.GetMasked(created.Key) + if err != nil { + created.APIKey = "" + response.RespondWithSuccess(c, http.StatusCreated, enrichProvider(created)) + return + } + response.RespondWithSuccess(c, http.StatusCreated, enrichProvider(masked)) +} + +// handleLLMProviderUpdate replaces a provider's configuration. +// Full replacement: api_key empty string preserves existing value. +// PUT /setting/llm/providers/:key +func handleLLMProviderUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + key := c.Param("key") + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + existing, err := llmprovider.Global.Get(key) + if err != nil { + respondError(c, http.StatusNotFound, fmt.Sprintf("provider \"%s\" not found", key)) + return + } + if err := llmCheckOwnership(existing, info); err != nil { + respondError(c, http.StatusNotFound, err.Error()) + return + } + + var body map[string]interface{} + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + var provider llmprovider.Provider + provider.Key = key + provider.Owner = existing.Owner + provider.Source = existing.Source + provider.ConnectorID = existing.ConnectorID + provider.PresetKey = existing.PresetKey + provider.IsCustom = existing.IsCustom + + if v, ok := body["name"].(string); ok { + provider.Name = v + } else { + provider.Name = existing.Name + } + if v, ok := body["type"].(string); ok { + provider.Type = v + } else { + provider.Type = existing.Type + } + if v, ok := body["api_url"].(string); ok { + provider.APIURL = v + } else { + provider.APIURL = existing.APIURL + } + + if v, ok := body["api_key"].(string); ok && v != "" { + provider.APIKey = v + } else { + provider.APIKey = existing.APIKey + } + + if v, ok := body["enabled"].(bool); ok { + provider.Enabled = v + } else { + provider.Enabled = existing.Enabled + } + if v, ok := body["require_key"].(bool); ok { + provider.RequireKey = v + } else { + provider.RequireKey = existing.RequireKey + } + if v, ok := body["status"].(string); ok { + provider.Status = v + } else { + provider.Status = existing.Status + } + + if modelsRaw, ok := body["models"]; ok { + raw, _ := json.Marshal(modelsRaw) + var models []llmprovider.ModelInfo + if err := json.Unmarshal(raw, &models); err == nil { + provider.Models = models + } + } else { + provider.Models = existing.Models + } + if provider.Models == nil { + provider.Models = []llmprovider.ModelInfo{} + } + + if _, err = llmprovider.Global.Update(key, &provider); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + masked, err := llmprovider.Global.GetMasked(key) + if err != nil { + provider.APIKey = "" + response.RespondWithSuccess(c, http.StatusOK, enrichProvider(&provider)) + return + } + response.RespondWithSuccess(c, http.StatusOK, enrichProvider(masked)) +} + +// handleLLMProviderDelete removes a provider and cleans up role references. +// DELETE /setting/llm/providers/:key +func handleLLMProviderDelete(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + key := c.Param("key") + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + existing, err := llmprovider.Global.Get(key) + if err != nil { + respondError(c, http.StatusNotFound, fmt.Sprintf("provider \"%s\" not found", key)) + return + } + if err := llmCheckOwnership(existing, info); err != nil { + respondError(c, http.StatusNotFound, err.Error()) + return + } + + var warning string + if setting.Global != nil { + scope := llmScope(info) + roles, _ := setting.Global.Get(scope, llmRolesNS) + if roles != nil { + cleaned := false + for roleName, target := range roles { + if targetMap, ok := target.(map[string]interface{}); ok { + if provKey, _ := targetMap["provider"].(string); provKey == key { + delete(roles, roleName) + cleaned = true + } + } + } + if cleaned { + setting.Global.Set(scope, llmRolesNS, roles) + warning = fmt.Sprintf("roles referencing provider \"%s\" have been cleared", key) + } + } + } + + if err := llmprovider.Global.Delete(key); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + result := map[string]interface{}{"success": true} + if warning != "" { + result["warning"] = warning + } + response.RespondWithSuccess(c, http.StatusOK, result) +} + +// handleLLMProviderTest tests connectivity for a provider and writes back status. +// POST /setting/llm/providers/:key/test +func handleLLMProviderTest(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + key := c.Param("key") + + if llmprovider.Global == nil { + respondError(c, http.StatusInternalServerError, "LLM provider registry not initialized") + return + } + + llmEnsureEncKey() + + p, err := llmprovider.Global.Get(key) + if err != nil { + respondError(c, http.StatusNotFound, fmt.Sprintf("provider \"%s\" not found", key)) + return + } + if err := llmCheckOwnership(p, info); err != nil { + respondError(c, http.StatusNotFound, err.Error()) + return + } + + start := time.Now() + err = llmValidateKey(p.Type, p.APIURL, p.APIKey) + latency := time.Since(start).Milliseconds() + + var testResult llmprovider.ProviderTestResult + if err != nil { + testResult = llmprovider.ProviderTestResult{ + Success: false, + Message: err.Error(), + } + p.Status = "disconnected" + } else { + testResult = llmprovider.ProviderTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latency, + } + p.Status = "connected" + llmprovider.Global.Update(key, p) + } + + response.RespondWithSuccess(c, http.StatusOK, testResult) +} diff --git a/openapi/setting/mcp.go b/openapi/setting/mcp.go new file mode 100644 index 00000000..490e005a --- /dev/null +++ b/openapi/setting/mcp.go @@ -0,0 +1,507 @@ +package setting + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/gou/mcp" + mcpTypes "github.com/yaoapp/gou/mcp/types" + gouTypes "github.com/yaoapp/gou/types" + "github.com/yaoapp/yao/mcpclient" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" +) + +const mcpMaskPrefixLen = 7 + +func mcpOwner(info *oauthTypes.AuthorizedInfo) mcpclient.ClientOwner { + if info.TeamID != "" { + return mcpclient.ClientOwner{Type: "team", ID: info.TeamID} + } + return mcpclient.ClientOwner{Type: "user", ID: info.UserID} +} + +func mcpCheckOwnership(c *mcpclient.Client, info *oauthTypes.AuthorizedInfo) error { + owner := mcpOwner(info) + if c.Owner.Type != owner.Type || c.Owner.ID != owner.ID { + return fmt.Errorf("server not found") + } + return nil +} + +func mcpMaskToken(token string) string { + if token == "" { + return "" + } + plain := cloudDecrypt(token) + if len(plain) <= mcpMaskPrefixLen { + return strings.Repeat("*", len(plain)) + } + suffix := plain[len(plain)-4:] + prefix := plain[:mcpMaskPrefixLen] + return prefix + "..." + suffix +} + +func mcpClientToResponse(c *mcpclient.Client) map[string]interface{} { + resp := map[string]interface{}{ + "id": c.ID, + "name": c.Name, + "label": c.Label, + "transport": string(c.Transport), + "url": c.URL, + "enabled": c.Enabled, + "status": c.Status, + } + if c.Description != "" { + resp["description"] = c.Description + } + if c.AuthorizationToken != "" { + resp["authorization_token"] = mcpMaskToken(c.AuthorizationToken) + } + if c.Timeout != "" { + resp["timeout"] = c.Timeout + } + if len(c.Tags) > 0 { + resp["tags"] = c.Tags + } + return resp +} + +// handleMCPList returns MCP servers for the current user/team. +// Only http and sse transports are returned. +// GET /setting/mcp/servers +func handleMCPList(c *gin.Context) { + info := authorized.GetInfo(c) + owner := mcpOwner(info) + + if mcpclient.Global == nil { + respondError(c, http.StatusInternalServerError, "MCP client registry not initialized") + return + } + + all, err := mcpclient.Global.List(&mcpclient.ClientFilter{ + Owner: &owner, + Source: mcpclient.ClientSourceAll, + }) + if err != nil { + all = []mcpclient.Client{} + } + + servers := make([]map[string]interface{}, 0, len(all)) + for i := range all { + t := all[i].Transport + if t != mcpTypes.TransportHTTP && t != mcpTypes.TransportSSE { + continue + } + servers = append(servers, mcpClientToResponse(&all[i])) + } + + response.RespondWithSuccess(c, http.StatusOK, map[string]interface{}{ + "servers": servers, + }) +} + +// handleMCPCreate creates a new MCP server. +// POST /setting/mcp/servers +func handleMCPCreate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + + if mcpclient.Global == nil { + respondError(c, http.StatusInternalServerError, "MCP client registry not initialized") + return + } + + var body struct { + Name string `json:"name"` + Label string `json:"label"` + Description string `json:"description"` + Transport string `json:"transport"` + URL string `json:"url"` + AuthorizationToken string `json:"authorization_token"` + Timeout string `json:"timeout"` + Tags []string `json:"tags"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if body.Name == "" { + respondError(c, http.StatusBadRequest, "name is required") + return + } + if body.URL == "" { + respondError(c, http.StatusBadRequest, "url is required") + return + } + if _, err := url.ParseRequestURI(body.URL); err != nil { + respondError(c, http.StatusBadRequest, "invalid url format") + return + } + + transport := mcpTypes.TransportHTTP + if body.Transport == "sse" { + transport = mcpTypes.TransportSSE + } + + owner := mcpOwner(info) + + existing, _ := mcpclient.Global.List(&mcpclient.ClientFilter{ + Owner: &owner, + Source: mcpclient.ClientSourceAll, + }) + for _, ex := range existing { + if strings.EqualFold(ex.Name, body.Name) { + respondError(c, http.StatusBadRequest, fmt.Sprintf("server with name \"%s\" already exists", body.Name)) + return + } + } + + clientID := owner.Type + "." + owner.ID + "." + body.Name + client := &mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: clientID, + Name: body.Name, + Transport: transport, + URL: body.URL, + Timeout: body.Timeout, + MetaInfo: gouTypes.MetaInfo{ + Label: body.Label, + Description: body.Description, + Tags: body.Tags, + }, + }, + Enabled: true, + Status: "unconfigured", + Source: mcpclient.ClientSourceDynamic, + Owner: owner, + } + + if body.AuthorizationToken != "" { + client.AuthorizationToken = cloudEncrypt(body.AuthorizationToken) + } + if body.Timeout == "" { + client.Timeout = "30s" + } + + token := body.AuthorizationToken + status, _, errMsg := mcpProbeRaw(transport, body.URL, token, client.Timeout) + if status != "connected" { + respondError(c, http.StatusBadRequest, errMsg) + return + } + + client.Status = "connected" + created, err := mcpclient.Global.Create(client) + if err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + response.RespondWithSuccess(c, http.StatusOK, mcpClientToResponse(created)) +} + +// handleMCPUpdate updates an existing MCP server. +// PUT /setting/mcp/servers/:id +func handleMCPUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + id := c.Param("id") + + if mcpclient.Global == nil { + respondError(c, http.StatusInternalServerError, "MCP client registry not initialized") + return + } + + existing, err := mcpclient.Global.Get(id) + if err != nil { + respondError(c, http.StatusNotFound, "server not found") + return + } + if err := mcpCheckOwnership(existing, info); err != nil { + respondError(c, http.StatusNotFound, err.Error()) + return + } + + var body struct { + Name string `json:"name"` + Label string `json:"label"` + Description string `json:"description"` + Transport string `json:"transport"` + URL string `json:"url"` + AuthorizationToken string `json:"authorization_token"` + Timeout string `json:"timeout"` + Tags []string `json:"tags"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if body.URL != "" { + if _, err := url.ParseRequestURI(body.URL); err != nil { + respondError(c, http.StatusBadRequest, "invalid url format") + return + } + } + + updated := &mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: id, + Name: existing.Name, + MetaInfo: gouTypes.MetaInfo{ + Label: existing.Label, + Description: existing.Description, + Tags: existing.Tags, + }, + Transport: existing.Transport, + URL: existing.URL, + AuthorizationToken: existing.AuthorizationToken, + Timeout: existing.Timeout, + }, + Enabled: existing.Enabled, + Status: existing.Status, + Source: existing.Source, + Owner: existing.Owner, + } + + if body.Name != "" { + updated.Name = body.Name + } + if body.Label != "" { + updated.Label = body.Label + } + if body.Description != "" { + updated.Description = body.Description + } + if body.Transport != "" { + if body.Transport == "sse" { + updated.Transport = mcpTypes.TransportSSE + } else { + updated.Transport = mcpTypes.TransportHTTP + } + } + if body.URL != "" { + updated.URL = body.URL + } + if body.AuthorizationToken != "" { + updated.AuthorizationToken = cloudEncrypt(body.AuthorizationToken) + } + if body.Timeout != "" { + updated.Timeout = body.Timeout + } + if body.Tags != nil { + updated.Tags = body.Tags + } + + token := body.AuthorizationToken + if token == "" && updated.AuthorizationToken != "" { + token = cloudDecrypt(updated.AuthorizationToken) + } + probeTransport := updated.Transport + probeURL := updated.URL + status, _, errMsg := mcpProbeRaw(probeTransport, probeURL, token, updated.Timeout) + if status != "connected" { + respondError(c, http.StatusBadRequest, errMsg) + return + } + + updated.Status = "connected" + result, err := mcpclient.Global.Update(id, updated) + if err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + response.RespondWithSuccess(c, http.StatusOK, mcpClientToResponse(result)) +} + +// handleMCPDelete removes an MCP server. +// DELETE /setting/mcp/servers/:id +func handleMCPDelete(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + id := c.Param("id") + + if mcpclient.Global == nil { + respondError(c, http.StatusInternalServerError, "MCP client registry not initialized") + return + } + + existing, err := mcpclient.Global.Get(id) + if err != nil { + respondError(c, http.StatusNotFound, "server not found") + return + } + if err := mcpCheckOwnership(existing, info); err != nil { + respondError(c, http.StatusNotFound, err.Error()) + return + } + + if err := mcpclient.Global.Delete(id); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + c.Status(http.StatusNoContent) +} + +// mcpProbeRaw creates a temporary MCP client from raw config, tests Connect+Initialize+ListTools. +func mcpProbeRaw(transport mcpTypes.TransportType, urlStr, token, timeout string) (status string, latencyMs int64, errMsg string) { + if timeout == "" { + timeout = "30s" + } + tempID := fmt.Sprintf("__probe_%d", time.Now().UnixNano()) + dsl := mcpTypes.ClientDSL{ + ID: tempID, + Name: tempID, + Transport: transport, + URL: urlStr, + AuthorizationToken: token, + Timeout: timeout, + } + dslJSON, err := json.Marshal(dsl) + if err != nil { + return "disconnected", 0, fmt.Sprintf("marshal: %s", err) + } + + start := time.Now() + mcpClient, err := mcp.LoadClientSourceWithType(string(dslJSON), tempID, "") + if err != nil { + return "disconnected", 0, fmt.Sprintf("load: %s", err) + } + defer mcp.UnloadClient(tempID) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + if err := mcpClient.Connect(ctx); err != nil { + return "disconnected", time.Since(start).Milliseconds(), fmt.Sprintf("connect: %s", err) + } + defer mcpClient.Disconnect(context.Background()) + + if _, err := mcpClient.Initialize(ctx); err != nil { + return "disconnected", time.Since(start).Milliseconds(), fmt.Sprintf("initialize: %s", err) + } + + _, err = mcpClient.ListTools(ctx, "") + latencyMs = time.Since(start).Milliseconds() + if err != nil { + return "disconnected", latencyMs, fmt.Sprintf("listTools: %s", err) + } + return "connected", latencyMs, "" +} + +// handleMCPTest tests connectivity using raw config (for add/edit before save). +// Creates a temporary runtime client, tests ListTools, then cleans up. +// POST /setting/mcp/test +func handleMCPTest(c *gin.Context) { + if !guardOwner(c) { + return + } + + var body struct { + Transport string `json:"transport"` + URL string `json:"url"` + AuthorizationToken string `json:"authorization_token"` + Timeout string `json:"timeout"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + if body.URL == "" { + respondError(c, http.StatusBadRequest, "url is required") + return + } + + transport := mcpTypes.TransportHTTP + if body.Transport == "sse" { + transport = mcpTypes.TransportSSE + } + timeout := body.Timeout + if timeout == "" { + timeout = "30s" + } + + tempID := fmt.Sprintf("__test_%d", time.Now().UnixNano()) + dsl := mcpTypes.ClientDSL{ + ID: tempID, + Name: tempID, + Transport: transport, + URL: body.URL, + AuthorizationToken: body.AuthorizationToken, + Timeout: timeout, + } + + dslJSON, err := json.Marshal(dsl) + if err != nil { + respondError(c, http.StatusInternalServerError, "failed to marshal config") + return + } + + start := time.Now() + mcpClient, err := mcp.LoadClientSourceWithType(string(dslJSON), tempID, "") + if err != nil { + response.RespondWithSuccess(c, http.StatusOK, mcpclient.ClientTestResult{ + Success: false, + Message: fmt.Sprintf("Failed to load client: %s", err.Error()), + }) + return + } + defer mcp.UnloadClient(tempID) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + if err := mcpClient.Connect(ctx); err != nil { + latencyMs := time.Since(start).Milliseconds() + response.RespondWithSuccess(c, http.StatusOK, mcpclient.ClientTestResult{ + Success: false, + Message: fmt.Sprintf("Connection failed: %s", err.Error()), + LatencyMs: latencyMs, + }) + return + } + defer mcpClient.Disconnect(context.Background()) + + if _, err := mcpClient.Initialize(ctx); err != nil { + latencyMs := time.Since(start).Milliseconds() + response.RespondWithSuccess(c, http.StatusOK, mcpclient.ClientTestResult{ + Success: false, + Message: fmt.Sprintf("Initialization failed: %s", err.Error()), + LatencyMs: latencyMs, + }) + return + } + + _, err = mcpClient.ListTools(ctx, "") + latencyMs := time.Since(start).Milliseconds() + + if err != nil { + response.RespondWithSuccess(c, http.StatusOK, mcpclient.ClientTestResult{ + Success: false, + Message: fmt.Sprintf("Connection failed: %s", err.Error()), + LatencyMs: latencyMs, + }) + return + } + + response.RespondWithSuccess(c, http.StatusOK, mcpclient.ClientTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latencyMs, + }) +} diff --git a/openapi/setting/promotions.yml b/openapi/setting/promotions.yml new file mode 100644 index 00000000..a85feec9 --- /dev/null +++ b/openapi/setting/promotions.yml @@ -0,0 +1,42 @@ +# Promotions & localized labels for the System Info page. +# Embedded at compile time via go:embed. + +# --- Localized UI labels --- +labels: + deployment: + community: + zh: "社区版" + en: "Community" + starter: + zh: "入门版" + en: "Starter" + pro: + zh: "专业版" + en: "Pro" + enterprise: + zh: "企业版" + en: "Enterprise" + cloud: + zh: "Cloud" + en: "Cloud" + environment: + development: + zh: "测试环境" + en: "Development" + production: + zh: "正式环境" + en: "Production" + +# --- Promotions by deployment type --- +community: + - id: upgrade-enterprise + link: "https://yaoagents.com/enterprise?source=yao-setting" + i18n: + zh: + title: "升级到企业版" + desc: "专属支持、私有部署、完全可控，行业 Agents 方案" + label: "了解更多 →" + en: + title: "Upgrade to Enterprise" + desc: "Dedicated support, private deployment, full control, industry-specific Agents solutions" + label: "Learn more →" diff --git a/openapi/setting/sandbox.go b/openapi/setting/sandbox.go new file mode 100644 index 00000000..72fd8d12 --- /dev/null +++ b/openapi/setting/sandbox.go @@ -0,0 +1,796 @@ +package setting + +import ( + "context" + "encoding/base64" + "net/http" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + sandboxv2 "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/setting" + "github.com/yaoapp/yao/tai" + "github.com/yaoapp/yao/tai/registry" + "github.com/yaoapp/yao/tai/runtime" + taitypes "github.com/yaoapp/yao/tai/types" +) + +const sandboxRegistryNS = "sandbox.registry" + +// pullState tracks an in-progress image pull operation. +type pullState struct { + ImageRef string + NodeID string + Progress int // 0-100 + Error string // non-empty on failure + Done bool +} + +var pullTracker sync.Map // key: "nodeID:imageRef" + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +func imageRefToID(ref string) string { + return base64.RawURLEncoding.EncodeToString([]byte(ref)) +} + +func idToImageRef(id string) (string, error) { + b, err := base64.RawURLEncoding.DecodeString(id) + return string(b), err +} + +func friendlyImageError(locale string, msg string) string { + isCN := strings.HasPrefix(strings.ToLower(locale), "zh") + + if strings.Contains(msg, "conflict") || strings.Contains(msg, "must force") { + if isCN { + return "该镜像正在被运行中的沙箱使用，请先停止相关沙箱后再删除" + } + return "This image is in use by a running sandbox. Please stop the sandbox first before removing." + } + if strings.Contains(msg, "No such image") || strings.Contains(msg, "not found") { + if isCN { + return "镜像不存在或已被删除" + } + return "Image not found or already removed" + } + if strings.Contains(msg, "no matching manifest") { + if isCN { + return "该镜像不支持当前系统架构（" + msg + "）" + } + return "This image does not support the current architecture (" + msg + ")" + } + if strings.Contains(msg, "pull access denied") || strings.Contains(msg, "repository does not exist") { + if isCN { + return "镜像不存在或无拉取权限，请检查镜像名称和仓库配置" + } + return "Image not found or access denied. Please check the image name and registry config." + } + if strings.Contains(msg, "dial tcp") || strings.Contains(msg, "timeout") || strings.Contains(msg, "TLS handshake") { + if isCN { + return "无法连接镜像仓库，请检查网络连接" + } + return "Cannot connect to the image registry. Please check your network." + } + if isCN { + return "操作失败: " + msg + } + return "Operation failed: " + msg +} + +func friendlyOS(goos string) string { + switch strings.ToLower(goos) { + case "darwin": + return "macOS" + case "linux": + return "Linux" + case "windows": + return "Windows" + default: + return goos + } +} + +func getSandboxManager() *sandboxv2.Manager { + defer func() { recover() }() + return sandboxv2.M() +} + +func sandboxNodeOwnedBy(snap *taitypes.NodeMeta, authInfo *oauthTypes.AuthorizedInfo) bool { + if authInfo == nil { + return true + } + if authInfo.TeamID != "" { + return snap.Auth.TeamID == authInfo.TeamID + } + if authInfo.UserID != "" { + return snap.Auth.TeamID == "" && snap.Auth.UserID == authInfo.UserID + } + return true +} + +type dockerInfoResult struct { + Version string + MemTotal int64 + NCPU int +} + +func fetchDockerInfo(nodeID string) *dockerInfoResult { + res, ok := tai.GetResources(nodeID) + if !ok || res.Runtime == nil { + return nil + } + cli := runtime.DockerCli(res.Runtime) + if cli == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + info, err := cli.Info(ctx) + if err != nil { + return nil + } + return &dockerInfoResult{ + Version: info.ServerVersion, + MemTotal: info.MemTotal, + NCPU: info.NCPU, + } +} + +// collectAssistantImages traverses assistant cache to find sandbox image requirements. +// Returns map[imageRef][]assistantDisplayName (locale-resolved). +func collectAssistantImages(locale string) map[string][]string { + cache := assistant.GetCache() + if cache == nil { + return nil + } + result := make(map[string][]string) + for _, ast := range cache.All() { + if ast == nil { + continue + } + var imageRef string + if ast.SandboxV2 != nil && ast.SandboxV2.Computer.Image != "" { + imageRef = ast.SandboxV2.Computer.Image + } else if ast.Sandbox != nil && ast.Sandbox.Image != "" { + imageRef = ast.Sandbox.Image + } + if imageRef != "" { + name := ast.GetName(locale) + if name == "" { + name = ast.ID + } + result[imageRef] = append(result[imageRef], name) + } + } + return result +} + +// splitImageRef splits "repo/name:tag" into (imageName, tag). +func splitImageRef(ref string) (string, string) { + if idx := strings.LastIndex(ref, ":"); idx > 0 && !strings.Contains(ref[idx:], "/") { + return ref[:idx], ref[idx+1:] + } + return ref, "latest" +} + +// getNodeResources retrieves ConnResources for a node with image capability. +// Returns (resources, httpStatus, errorMessage). +func getNodeResources(nodeID string) (*tai.ConnResources, int, string) { + reg := registry.Global() + if reg == nil { + return nil, http.StatusServiceUnavailable, "tai registry not initialized" + } + meta, ok := reg.Get(nodeID) + if !ok { + return nil, http.StatusNotFound, "node not found: " + nodeID + } + if meta.Status != "online" { + return nil, http.StatusBadRequest, "node is offline: " + nodeID + } + res, ok := tai.GetResources(nodeID) + if !ok { + return nil, http.StatusBadGateway, "cannot reach node: " + nodeID + } + if res.Image == nil { + return nil, http.StatusBadRequest, "Docker not available on this node" + } + return res, 0, "" +} + +// --------------------------------------------------------------------------- +// GET /setting/sandbox +// --------------------------------------------------------------------------- + +func handleSandboxGet(c *gin.Context) { + info := authorized.GetInfo(c) + locale := strings.ToLower(c.DefaultQuery("locale", "en-us")) + + reg := registry.Global() + var snaps []taitypes.NodeMeta + if reg != nil { + snaps = reg.List() + } + + // Filter nodes by ownership + var filtered []taitypes.NodeMeta + for i := range snaps { + s := &snaps[i] + if s.Mode != "local" && !sandboxNodeOwnedBy(s, info) { + continue + } + if !s.Capabilities.Docker { + continue + } + filtered = append(filtered, *s) + } + + mgr := getSandboxManager() + + // Build nodes concurrently + nodes := make([]ComputerNode, len(filtered)) + var wg sync.WaitGroup + for i, snap := range filtered { + wg.Add(1) + go func(idx int, s taitypes.NodeMeta) { + defer wg.Done() + kind := "tai-link" + if s.Mode == "local" { + kind = "local" + } + node := ComputerNode{ + NodeID: s.TaiID, + DisplayName: s.DisplayName, + Kind: kind, + OS: friendlyOS(s.System.OS), + Arch: s.System.Arch, + CPU: s.System.NumCPU, + MemoryGB: int(s.System.TotalMem / (1024 * 1024 * 1024)), + Online: s.Status == "online", + } + if node.DisplayName == "" { + node.DisplayName = s.System.Hostname + } + if node.DisplayName == "" { + node.DisplayName = s.TaiID + } + + // Fetch Docker info for online nodes + if node.Online { + if di := fetchDockerInfo(s.TaiID); di != nil { + node.DockerVersion = di.Version + if node.MemoryGB == 0 && di.MemTotal > 0 { + node.MemoryGB = int(di.MemTotal / (1024 * 1024 * 1024)) + } + if node.CPU == 0 && di.NCPU > 0 { + node.CPU = di.NCPU + } + } + } + + // Count running sandboxes + if mgr != nil { + boxes, err := mgr.List(context.Background(), sandboxv2.ListOptions{NodeID: s.TaiID}) + if err == nil { + node.RunningSandboxes = len(boxes) + } + } + + nodes[idx] = node + }(i, snap) + } + wg.Wait() + + // Registry config + regConfig := SandboxRegistryConfig{} + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, sandboxRegistryNS) + if v, ok := saved["registry_url"].(string); ok { + regConfig.RegistryURL = v + } + if v, ok := saved["username"].(string); ok { + regConfig.Username = v + } + if v, ok := saved["password"].(string); ok && v != "" { + regConfig.Password = cloudMaskKey(cloudDecrypt(v)) + } + } + + // Collect assistant images (locale-resolved names) + assistantImages := collectAssistantImages(locale) + + // Build image list per node concurrently + images := make(map[string][]SandboxImage) + var imgWg sync.WaitGroup + var imgMu sync.Mutex + for _, node := range nodes { + if !node.Online { + imgMu.Lock() + images[node.NodeID] = []SandboxImage{} + imgMu.Unlock() + continue + } + imgWg.Add(1) + go func(nodeID string) { + defer imgWg.Done() + nodeImages := buildNodeImages(nodeID, assistantImages, locale) + imgMu.Lock() + images[nodeID] = nodeImages + imgMu.Unlock() + }(node.NodeID) + } + imgWg.Wait() + + data := SandboxPageData{ + Nodes: nodes, + Registry: regConfig, + Images: images, + } + if data.Nodes == nil { + data.Nodes = []ComputerNode{} + } + + response.RespondWithSuccess(c, http.StatusOK, data) +} + +func buildNodeImages(nodeID string, assistantImages map[string][]string, locale string) []SandboxImage { + res, ok := tai.GetResources(nodeID) + if !ok || res.Image == nil { + return []SandboxImage{} + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + localImages, err := res.Image.List(ctx) + if err != nil { + return []SandboxImage{} + } + + // Build tag index from local images + tagIndex := make(map[string]runtime.ImageInfo) + for _, img := range localImages { + for _, tag := range img.Tags { + tagIndex[tag] = img + } + } + + var result []SandboxImage + for imageRef, names := range assistantImages { + imgName, tag := splitImageRef(imageRef) + si := SandboxImage{ + ID: imageRefToID(imageRef), + AssistantNames: names, + ImageName: imgName, + Tag: tag, + Status: "not_downloaded", + } + + // Check if already downloaded + if info, ok := tagIndex[imageRef]; ok { + si.Status = "downloaded" + si.SizeMB = int(info.Size / (1024 * 1024)) + } + + trackerKey := nodeID + ":" + imageRef + if v, ok := pullTracker.Load(trackerKey); ok { + ps := v.(*pullState) + if !ps.Done { + si.Status = "downloading" + p := ps.Progress + si.Progress = &p + } else if ps.Error != "" { + si.Status = "error" + si.ErrorMessage = friendlyImageError(locale, ps.Error) + } else { + si.Status = "downloaded" + } + } + + result = append(result, si) + } + + if result == nil { + return []SandboxImage{} + } + return result +} + +// --------------------------------------------------------------------------- +// PUT /setting/sandbox/registry +// --------------------------------------------------------------------------- + +func handleSandboxRegistry(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := cloudScope(info) + + var body SandboxRegistryConfig + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + m := map[string]interface{}{ + "registry_url": body.RegistryURL, + "username": body.Username, + } + if body.Password != "" { + m["password"] = cloudEncrypt(body.Password) + } else { + existing, _ := setting.Global.Get(scope, sandboxRegistryNS) + if v, ok := existing["password"].(string); ok { + m["password"] = v + } + } + + if _, err := setting.Global.Set(scope, sandboxRegistryNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + result := SandboxRegistryConfig{ + RegistryURL: body.RegistryURL, + Username: body.Username, + } + if v, ok := m["password"].(string); ok && v != "" { + result.Password = cloudMaskKey(cloudDecrypt(v)) + } + + response.RespondWithSuccess(c, http.StatusOK, result) +} + +// --------------------------------------------------------------------------- +// POST /setting/sandbox/nodes/:nodeId/images/:imageId/pull +// --------------------------------------------------------------------------- + +func handleSandboxPull(c *gin.Context) { + if !guardOwner(c) { + return + } + + nodeID := c.Param("nodeId") + imageID := c.Param("imageId") + imageRef, err := idToImageRef(imageID) + if err != nil || imageRef == "" { + respondError(c, http.StatusBadRequest, "invalid image ID") + return + } + + res, status, errMsg := getNodeResources(nodeID) + if res == nil { + respondError(c, status, errMsg) + return + } + + pullOpts := runtime.PullOptions{} + info := authorized.GetInfo(c) + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, sandboxRegistryNS) + if regURL, ok := saved["registry_url"].(string); ok && regURL != "" { + if strings.HasPrefix(imageRef, regURL) || strings.HasPrefix(imageRef, strings.TrimPrefix(regURL, "https://")) { + user, _ := saved["username"].(string) + pass, _ := saved["password"].(string) + if user != "" { + pullOpts.Auth = &runtime.RegistryAuth{ + Username: user, + Password: cloudDecrypt(pass), + Server: regURL, + } + } + } + } + } + + trackerKey := nodeID + ":" + imageRef + log.Info("[sandbox] pull start: trackerKey=%s imageRef=%s", trackerKey, imageRef) + pullTracker.Store(trackerKey, &pullState{ + ImageRef: imageRef, + NodeID: nodeID, + Progress: 0, + }) + + ch, pullErr := res.Image.Pull(context.Background(), imageRef, pullOpts) + if pullErr != nil { + log.Error("[sandbox] pull initiate failed: %s err=%v", trackerKey, pullErr) + pullTracker.Delete(trackerKey) + respondError(c, http.StatusBadGateway, "pull failed: "+pullErr.Error()) + return + } + + if ch != nil { + log.Info("[sandbox] pull channel received, starting goroutine: %s", trackerKey) + go consumePullProgress(trackerKey, ch) + } else { + log.Info("[sandbox] pull channel is nil, marking as done: %s", trackerKey) + pullTracker.Store(trackerKey, &pullState{ + ImageRef: imageRef, + NodeID: nodeID, + Progress: 100, + Done: true, + }) + } + + imgName, tag := splitImageRef(imageRef) + p := 0 + response.RespondWithSuccess(c, http.StatusOK, SandboxImage{ + ID: imageRefToID(imageRef), + ImageName: imgName, + Tag: tag, + Status: "downloading", + Progress: &p, + }) +} + +func consumePullProgress(trackerKey string, ch <-chan runtime.PullProgress) { + log.Info("[sandbox] consumePullProgress started: %s", trackerKey) + var totalBytes int64 + var currentBytes int64 + var eventCount int + layerProgress := make(map[string]int64) + layerTotal := make(map[string]int64) + + for p := range ch { + eventCount++ + if p.Error != "" { + log.Error("[sandbox] pull error: %s err=%s", trackerKey, p.Error) + pullTracker.Store(trackerKey, &pullState{ + Done: true, + Error: p.Error, + }) + go func() { + time.Sleep(60 * time.Second) + pullTracker.Delete(trackerKey) + }() + return + } + + if p.Layer != "" && p.Total > 0 { + layerTotal[p.Layer] = p.Total + layerProgress[p.Layer] = p.Current + } + + totalBytes = 0 + currentBytes = 0 + for layer, t := range layerTotal { + totalBytes += t + currentBytes += layerProgress[layer] + } + + pct := 0 + if totalBytes > 0 { + pct = int(currentBytes * 100 / totalBytes) + if pct > 99 { + pct = 99 + } + } + + pullTracker.Store(trackerKey, &pullState{ + Progress: pct, + }) + } + + log.Info("[sandbox] pull complete (channel closed): %s events=%d", trackerKey, eventCount) + pullTracker.Store(trackerKey, &pullState{ + Progress: 100, + Done: true, + }) + go func() { + time.Sleep(60 * time.Second) + pullTracker.Delete(trackerKey) + }() +} + +// --------------------------------------------------------------------------- +// POST /setting/sandbox/nodes/:nodeId/images/pull-all +// --------------------------------------------------------------------------- + +func handleSandboxPullAll(c *gin.Context) { + if !guardOwner(c) { + return + } + + nodeID := c.Param("nodeId") + locale := strings.ToLower(c.DefaultQuery("locale", "en-us")) + + res, status, errMsg := getNodeResources(nodeID) + if res == nil { + respondError(c, status, errMsg) + return + } + + assistantImages := collectAssistantImages(locale) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + localImages, _ := res.Image.List(ctx) + tagIndex := make(map[string]bool) + for _, img := range localImages { + for _, tag := range img.Tags { + tagIndex[tag] = true + } + } + + // Build pull options + pullOpts := runtime.PullOptions{} + info := authorized.GetInfo(c) + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, sandboxRegistryNS) + if regURL, ok := saved["registry_url"].(string); ok && regURL != "" { + user, _ := saved["username"].(string) + pass, _ := saved["password"].(string) + if user != "" { + pullOpts.Auth = &runtime.RegistryAuth{ + Username: user, + Password: cloudDecrypt(pass), + Server: regURL, + } + } + } + } + + var result []SandboxImage + for imageRef, names := range assistantImages { + if tagIndex[imageRef] { + continue + } + + trackerKey := nodeID + ":" + imageRef + // Skip if already pulling + if v, ok := pullTracker.Load(trackerKey); ok { + ps := v.(*pullState) + if !ps.Done { + imgName, tag := splitImageRef(imageRef) + p := ps.Progress + result = append(result, SandboxImage{ + ID: imageRefToID(imageRef), + AssistantNames: names, + ImageName: imgName, + Tag: tag, + Status: "downloading", + Progress: &p, + }) + continue + } + } + + pullTracker.Store(trackerKey, &pullState{ + ImageRef: imageRef, + NodeID: nodeID, + Progress: 0, + }) + + ch, pullErr := res.Image.Pull(context.Background(), imageRef, pullOpts) + if pullErr != nil { + pullTracker.Delete(trackerKey) + continue + } + if ch != nil { + go consumePullProgress(trackerKey, ch) + } + + imgName, tag := splitImageRef(imageRef) + p := 0 + result = append(result, SandboxImage{ + ID: imageRefToID(imageRef), + AssistantNames: names, + ImageName: imgName, + Tag: tag, + Status: "downloading", + Progress: &p, + }) + } + + if result == nil { + result = []SandboxImage{} + } + response.RespondWithSuccess(c, http.StatusOK, result) +} + +// --------------------------------------------------------------------------- +// DELETE /setting/sandbox/nodes/:nodeId/images/:imageId +// --------------------------------------------------------------------------- + +func handleSandboxImageDelete(c *gin.Context) { + if !guardOwner(c) { + return + } + + nodeID := c.Param("nodeId") + imageID := c.Param("imageId") + imageRef, err := idToImageRef(imageID) + if err != nil || imageRef == "" { + respondError(c, http.StatusBadRequest, "invalid image ID") + return + } + + res, status, errMsg := getNodeResources(nodeID) + if res == nil { + respondError(c, status, errMsg) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if res.Runtime != nil { + containers, _ := res.Runtime.List(ctx, runtime.ListOptions{All: true}) + for _, ctr := range containers { + if ctr.Image == imageRef { + _ = res.Runtime.Remove(ctx, ctr.ID, true) + } + } + } + + if err := res.Image.Remove(ctx, imageRef, true); err != nil { + locale := strings.ToLower(c.DefaultQuery("locale", "en-us")) + respondError(c, http.StatusBadRequest, friendlyImageError(locale, err.Error())) + return + } + + response.RespondWithSuccess(c, http.StatusOK, gin.H{"success": true}) +} + +// --------------------------------------------------------------------------- +// POST /setting/sandbox/nodes/:nodeId/check-docker +// --------------------------------------------------------------------------- + +func handleSandboxCheckDocker(c *gin.Context) { + nodeID := c.Param("nodeId") + + reg := registry.Global() + if reg == nil { + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": nil, "message": "tai registry not initialized"}) + return + } + + meta, ok := reg.Get(nodeID) + if !ok { + respondError(c, http.StatusNotFound, "node not found: "+nodeID) + return + } + + if meta.Status != "online" { + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": nil, "message": "node is offline"}) + return + } + + res, ok := tai.GetResources(nodeID) + if !ok || res.Runtime == nil { + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": nil, "message": "Docker not available"}) + return + } + + cli := runtime.DockerCli(res.Runtime) + if cli == nil { + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": nil, "message": "Docker not available"}) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ver, err := cli.ServerVersion(ctx) + if err != nil { + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": nil, "message": "Docker check failed: " + err.Error()}) + return + } + + response.RespondWithSuccess(c, http.StatusOK, gin.H{"docker_version": ver.Version}) +} diff --git a/openapi/setting/search.go b/openapi/setting/search.go new file mode 100644 index 00000000..f7341f31 --- /dev/null +++ b/openapi/setting/search.go @@ -0,0 +1,622 @@ +package setting + +import ( + "bytes" + _ "embed" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/setting" + "gopkg.in/yaml.v3" +) + +//go:embed search_presets.yml +var searchPresetsYML []byte + +var searchPresets []SearchProviderPreset + +func init() { + if err := yaml.Unmarshal(searchPresetsYML, &searchPresets); err != nil { + searchPresets = nil + } +} + +func searchFindPreset(key string) *SearchProviderPreset { + for i := range searchPresets { + if searchPresets[i].Key == key { + return &searchPresets[i] + } + } + return nil +} + +func searchScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID { + if info.TeamID != "" { + return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID} + } + return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID} +} + +func searchProviderNS(key string) string { + return "search.providers." + key +} + +const searchAssignmentNS = "search.tool_assignment" + +func searchPasswordFields(preset *SearchProviderPreset) map[string]bool { + m := make(map[string]bool) + for _, f := range preset.Fields { + if f.Type == "password" { + m[f.Key] = true + } + } + return m +} + +// --------------------------------------------------------------------------- +// GET /setting/search +// --------------------------------------------------------------------------- + +func handleSearchGet(c *gin.Context) { + info := authorized.GetInfo(c) + + providers := make([]SearchProviderConfig, 0, len(searchPresets)) + for _, preset := range searchPresets { + cfg := SearchProviderConfig{ + PresetKey: preset.Key, + Enabled: false, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + + if preset.IsCloud { + var cloudSaved map[string]interface{} + if setting.Global != nil { + cloudSaved, _ = setting.Global.GetMerged(info.UserID, info.TeamID, cloudNS) + } + if cloudSaved != nil { + if st, ok := cloudSaved["status"].(string); ok && st == "connected" { + cfg.Enabled = true + cfg.Status = "connected" + } + } + } else if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchProviderNS(preset.Key)) + if saved != nil { + if v, ok := saved["enabled"].(bool); ok { + cfg.Enabled = v + } + if v, ok := saved["status"].(string); ok && v != "" { + cfg.Status = v + } + pwFields := searchPasswordFields(&preset) + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + } + } + providers = append(providers, cfg) + } + + var assignment SearchToolAssignment + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchAssignmentNS) + if saved != nil { + if v, ok := saved["web_search"].(string); ok && v != "" { + assignment.WebSearch = &v + } + if v, ok := saved["web_scrape"].(string); ok && v != "" { + assignment.WebScrape = &v + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, SearchPageData{ + Presets: searchPresets, + Providers: providers, + ToolAssignment: assignment, + }) +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/providers/:key +// --------------------------------------------------------------------------- + +func handleSearchProviderUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + FieldValues map[string]string `json:"field_values"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + existing, _ := setting.Global.Get(scope, searchProviderNS(key)) + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + + validFields := make(map[string]bool) + for _, f := range preset.Fields { + validFields[f.Key] = true + } + + pwFields := searchPasswordFields(preset) + existingFV := map[string]interface{}{} + if fv, ok := m["field_values"].(map[string]interface{}); ok { + existingFV = fv + } + + newFV := make(map[string]interface{}) + for k, v := range existingFV { + newFV[k] = v + } + + for k, v := range body.FieldValues { + if !validFields[k] { + continue + } + if pwFields[k] { + if v == "" { + continue // keep existing + } + newFV[k] = cloudEncrypt(v) + } else { + newFV[k] = v + } + } + + m["field_values"] = newFV + if _, ok := m["enabled"]; !ok { + m["enabled"] = false + } + if _, ok := m["status"]; !ok { + m["status"] = "unconfigured" + } + + if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + cfg := SearchProviderConfig{ + PresetKey: key, + Enabled: false, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + if v, ok := m["enabled"].(bool); ok { + cfg.Enabled = v + } + if v, ok := m["status"].(string); ok && v != "" { + cfg.Status = v + } + if fv, ok := m["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/providers/:key/toggle +// --------------------------------------------------------------------------- + +func handleSearchProviderToggle(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + Enabled bool `json:"enabled"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + existing, _ := setting.Global.Get(scope, searchProviderNS(key)) + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + m["enabled"] = body.Enabled + + if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + // When disabling, clear tool_assignment references + if !body.Enabled { + assignData, _ := setting.Global.Get(scope, searchAssignmentNS) + if assignData != nil { + changed := false + if v, ok := assignData["web_search"].(string); ok && v == key { + assignData["web_search"] = "" + changed = true + } + if v, ok := assignData["web_scrape"].(string); ok && v == key { + assignData["web_scrape"] = "" + changed = true + } + if changed { + setting.Global.Set(scope, searchAssignmentNS, assignData) + } + } + } + + cfg := SearchProviderConfig{ + PresetKey: key, + Enabled: body.Enabled, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + if v, ok := m["status"].(string); ok && v != "" { + cfg.Status = v + } + pwFields := searchPasswordFields(preset) + if fv, ok := m["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// POST /setting/search/providers/:key/test +// --------------------------------------------------------------------------- + +func handleSearchProviderTest(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider status is determined by cloud service configuration") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + FieldValues map[string]string `json:"field_values"` + } + c.ShouldBindJSON(&body) + + info := authorized.GetInfo(c) + scope := searchScope(info) + + // Resolve API key: prefer body, fall back to saved + apiKey := "" + if body.FieldValues != nil { + apiKey = body.FieldValues["api_key"] + } + if apiKey == "" && setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved != nil { + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + if v, ok := fv["api_key"].(string); ok { + apiKey = cloudDecrypt(v) + } + } + } + } + + if apiKey == "" { + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: false, + Message: "API key is required", + }) + return + } + + start := time.Now() + var testErr error + + zone := "" + if body.FieldValues != nil { + zone = body.FieldValues["zone"] + } + if zone == "" && setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved != nil { + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + if v, ok := fv["zone"].(string); ok { + zone = v + } + } + } + } + + switch key { + case "tavily": + testErr = searchTestTavily(apiKey) + case "serper": + testErr = searchTestSerper(apiKey) + case "brightdata": + testErr = searchTestBrightdata(apiKey, zone) + default: + respondError(c, http.StatusBadRequest, fmt.Sprintf("test not supported for provider: %s", key)) + return + } + + latency := time.Since(start).Milliseconds() + + if testErr != nil { + // Update status to disconnected + if setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved == nil { + saved = map[string]interface{}{} + } + saved["status"] = "disconnected" + setting.Global.Set(scope, searchProviderNS(key), saved) + } + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: false, + Message: testErr.Error(), + }) + return + } + + // Update status to connected + if setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved == nil { + saved = map[string]interface{}{} + } + saved["status"] = "connected" + setting.Global.Set(scope, searchProviderNS(key), saved) + } + + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latency, + }) +} + +func searchTestTavily(apiKey string) error { + payload, _ := json.Marshal(map[string]interface{}{ + "api_key": apiKey, + "query": "test", + }) + client := &http.Client{Timeout: 15 * time.Second} + resp, err := client.Post("https://api.tavily.com/search", "application/json", bytes.NewReader(payload)) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +func searchTestSerper(apiKey string) error { + payload, _ := json.Marshal(map[string]string{"q": "test"}) + client := &http.Client{Timeout: 15 * time.Second} + req, err := http.NewRequest("POST", "https://google.serper.dev/search", bytes.NewReader(payload)) + if err != nil { + return fmt.Errorf("failed to build request: %s", err.Error()) + } + req.Header.Set("X-API-KEY", apiKey) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +func searchTestBrightdata(apiKey, zone string) error { + if zone == "" { + return fmt.Errorf("Zone is required") + } + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", "https://api.brightdata.com/zone/status?zone="+zone, nil) + if err != nil { + return fmt.Errorf("failed to build request: %s", err.Error()) + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("zone '%s' not found", zone) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/tool-assignment +// --------------------------------------------------------------------------- + +func handleSearchToolAssignment(c *gin.Context) { + if !guardOwner(c) { + return + } + + var body struct { + WebSearch *string `json:"web_search"` + WebScrape *string `json:"web_scrape"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + // Validate: provider must be enabled and support the tool + validateAssignment := func(providerKey *string, toolType string) error { + if providerKey == nil || *providerKey == "" { + return nil + } + preset := searchFindPreset(*providerKey) + if preset == nil { + return fmt.Errorf("unknown provider: %s", *providerKey) + } + + hasTools := false + for _, t := range preset.Tools { + if t == toolType { + hasTools = true + break + } + } + if !hasTools { + return fmt.Errorf("provider %s does not support %s", *providerKey, toolType) + } + + if preset.IsCloud { + return nil // cloud provider enablement is implicit + } + + saved, _ := setting.Global.Get(scope, searchProviderNS(*providerKey)) + if saved != nil { + if v, ok := saved["enabled"].(bool); ok && v { + return nil + } + } + return fmt.Errorf("provider %s is not enabled", *providerKey) + } + + if err := validateAssignment(body.WebSearch, "web_search"); err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + if err := validateAssignment(body.WebScrape, "web_scrape"); err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + + m := make(map[string]interface{}) + if body.WebSearch != nil { + m["web_search"] = *body.WebSearch + } else { + m["web_search"] = "" + } + if body.WebScrape != nil { + m["web_scrape"] = *body.WebScrape + } else { + m["web_scrape"] = "" + } + + if _, err := setting.Global.Set(scope, searchAssignmentNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + result := SearchToolAssignment{} + if v, ok := m["web_search"].(string); ok && v != "" { + result.WebSearch = &v + } + if v, ok := m["web_scrape"].(string); ok && v != "" { + result.WebScrape = &v + } + + response.RespondWithSuccess(c, http.StatusOK, result) +} diff --git a/openapi/setting/search_presets.yml b/openapi/setting/search_presets.yml new file mode 100644 index 00000000..8775c56e --- /dev/null +++ b/openapi/setting/search_presets.yml @@ -0,0 +1,61 @@ +# Search & Scrape provider presets. +# Embedded at compile time via go:embed in search.go. + +- key: cloud + name: Yao Agents + description: + zh-CN: "云服务提供的搜索与抓取能力，凭证来自云服务配置页" + en-US: "Search & scrape capabilities from cloud service, credentials from cloud config" + website: "https://yaoagents.com" + tools: [web_search, web_scrape] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + - { zh-CN: "网页抓取", en-US: "Web Scrape" } + fields: [] + is_cloud: true + +- key: tavily + name: Tavily + description: + zh-CN: "AI 优化的搜索 API，返回结构化结果，适合 Agent 使用" + en-US: "AI-optimized search API with structured results, ideal for agents" + website: "https://tavily.com" + tools: [web_search] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + +- key: serper + name: "Serper (Google)" + description: + zh-CN: "基于 Google 搜索的 API，价格实惠，结果质量高" + en-US: "Google Search API with affordable pricing and high-quality results" + website: "https://serper.dev" + tools: [web_search] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + +- key: brightdata + name: Brightdata + description: + zh-CN: "部分网站有访问限制，启用代理可提升抓取成功率。需开通 Web Access API (Web Unlocker)。" + en-US: "Some websites have access restrictions. Enabling proxy improves scraping success rate. Requires Web Access API (Web Unlocker)." + website: "https://brightdata.com" + tools: [web_scrape] + tool_labels: + - { zh-CN: "网页抓取", en-US: "Web Scrape" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + - key: zone + label: { zh-CN: "Zone", en-US: "Zone" } + type: text + hint: { zh-CN: "Web Unlocker API 的 Zone 名称", en-US: "Zone name of your Web Unlocker API" } diff --git a/openapi/setting/setting.go b/openapi/setting/setting.go new file mode 100644 index 00000000..a8001501 --- /dev/null +++ b/openapi/setting/setting.go @@ -0,0 +1,160 @@ +package setting + +import ( + "fmt" + "net/http" + "time" + + "github.com/gin-gonic/gin" + gouStore "github.com/yaoapp/gou/store" + "github.com/yaoapp/kun/log" + oauth "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" +) + +const ownerCachePrefix = "setting:owner:" +const ownerCacheTTL = 5 * time.Minute + +func getCache() gouStore.Store { + c, _ := gouStore.Get("__yao.cache") + return c +} + +// Attach registers all /setting/* routes under the given group. +// Currently only System Info routes are wired; other groups will be +// added incrementally. +func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) { + group.Use(oauth.Guard) + + sys := group.Group("/system") + sys.GET("", handleSystemInfo) + sys.POST("/check-update", handleSystemCheckUpdate) + + cloud := group.Group("/cloud") + cloud.GET("", handleCloudGet) + cloud.PUT("", handleCloudUpdate) + cloud.POST("/test", handleCloudTest) + + llm := group.Group("/llm") + llm.GET("", handleLLMGet) + llm.PUT("/roles", handleLLMRoles) + llm.POST("/test", handleLLMTest) + llm.POST("/providers", handleLLMProviderCreate) + llm.PUT("/providers/:key", handleLLMProviderUpdate) + llm.DELETE("/providers/:key", handleLLMProviderDelete) + llm.POST("/providers/:key/test", handleLLMProviderTest) + + search := group.Group("/search") + search.GET("", handleSearchGet) + search.PUT("/providers/:key", handleSearchProviderUpdate) + search.PUT("/providers/:key/toggle", handleSearchProviderToggle) + search.POST("/providers/:key/test", handleSearchProviderTest) + search.PUT("/tool-assignment", handleSearchToolAssignment) + + smtpG := group.Group("/smtp") + smtpG.GET("", handleSmtpGet) + smtpG.PUT("", handleSmtpUpdate) + smtpG.PUT("/toggle", handleSmtpToggle) + smtpG.POST("/test", handleSmtpTest) + + mcpG := group.Group("/mcp") + mcpG.GET("/servers", handleMCPList) + mcpG.POST("/servers", handleMCPCreate) + mcpG.PUT("/servers/:id", handleMCPUpdate) + mcpG.DELETE("/servers/:id", handleMCPDelete) + mcpG.POST("/test", handleMCPTest) + + sb := group.Group("/sandbox") + sb.GET("", handleSandboxGet) + sb.PUT("/registry", handleSandboxRegistry) + sb.POST("/nodes/:nodeId/images/:imageId/pull", handleSandboxPull) + sb.POST("/nodes/:nodeId/images/pull-all", handleSandboxPullAll) + sb.DELETE("/nodes/:nodeId/images/:imageId", handleSandboxImageDelete) + sb.POST("/nodes/:nodeId/check-docker", handleSandboxCheckDocker) +} + +// requireOwner checks that the current user is the team owner. +// Non-team context (TeamID == ""): always allowed — user is managing their own data. +// Team context: checks cache first, then queries the member table is_owner field. +// Use this as a guard for any write operation across all /setting/* groups. +func requireOwner(c *gin.Context, info *oauthTypes.AuthorizedInfo) error { + if info == nil || info.UserID == "" { + return fmt.Errorf("authentication required") + } + if info.TeamID == "" { + return nil + } + + cacheKey := ownerCachePrefix + info.TeamID + ":" + info.UserID + + if cache := getCache(); cache != nil { + if val, ok := cache.Get(cacheKey); ok { + if isOwner, ok := val.(bool); ok { + if isOwner { + return nil + } + return fmt.Errorf("access denied: only team owner can modify settings") + } + } + } + + if oauth.OAuth == nil { + return fmt.Errorf("service not initialized") + } + provider, err := oauth.OAuth.GetUserProvider() + if err != nil { + return fmt.Errorf("service not available") + } + + member, err := provider.GetMember(c.Request.Context(), info.TeamID, info.UserID) + if err != nil { + log.Error("[setting] GetMember failed: %v", err) + return fmt.Errorf("access denied") + } + + isOwner := checkIsOwner(member["is_owner"]) + if cache := getCache(); cache != nil { + cache.Set(cacheKey, isOwner, ownerCacheTTL) + } + + if isOwner { + return nil + } + return fmt.Errorf("access denied: only team owner can modify settings") +} + +func checkIsOwner(val interface{}) bool { + switch v := val.(type) { + case bool: + return v + case int: + return v == 1 + case int64: + return v == 1 + case float64: + return v == 1 + } + return false +} + +// guardOwner is a convenience wrapper: calls requireOwner and writes 403 on failure. +// Returns true if the request should continue, false if it was aborted. +func guardOwner(c *gin.Context) bool { + info := authorized.GetInfo(c) + if err := requireOwner(c, info); err != nil { + respondError(c, http.StatusForbidden, err.Error()) + return false + } + return true +} + +// respondError is a thin helper that writes a JSON error via the shared +// response package. +func respondError(c *gin.Context, status int, msg string) { + response.RespondWithError(c, status, &response.ErrorResponse{ + Code: "server_error", + ErrorDescription: msg, + }) +} diff --git a/openapi/setting/smtp.go b/openapi/setting/smtp.go new file mode 100644 index 00000000..fc71a2e5 --- /dev/null +++ b/openapi/setting/smtp.go @@ -0,0 +1,586 @@ +package setting + +import ( + "crypto/tls" + _ "embed" + "fmt" + "net" + "net/http" + "net/smtp" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/setting" + "gopkg.in/yaml.v3" +) + +//go:embed smtp_presets.yml +var smtpPresetsYML []byte + +const smtpNS = "smtp" + +var smtpPresetsMap map[string][]SmtpPreset + +func init() { + smtpPresetsMap = make(map[string][]SmtpPreset) + if err := yaml.Unmarshal(smtpPresetsYML, &smtpPresetsMap); err != nil { + smtpPresetsMap = map[string][]SmtpPreset{} + } +} + +func smtpGetPresets(locale string) []SmtpPreset { + locale = strings.ToLower(locale) + if presets, ok := smtpPresetsMap[locale]; ok { + return presets + } + if presets, ok := smtpPresetsMap["en-us"]; ok { + return presets + } + return nil +} + +func smtpDefaultPreset(presets []SmtpPreset) *SmtpPreset { + for i := range presets { + if presets[i].Default { + return &presets[i] + } + } + if len(presets) > 0 { + return &presets[0] + } + return nil +} + +func smtpScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID { + if info.TeamID != "" { + return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID} + } + return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID} +} + +// --------------------------------------------------------------------------- +// Rate limiter: 5 test emails per minute per scope +// --------------------------------------------------------------------------- + +var ( + smtpRateMu sync.Mutex + smtpRateStore = map[string][]time.Time{} +) + +const smtpRateLimit = 5 +const smtpRateWindow = time.Minute + +func smtpCheckRateLimit(key string) bool { + smtpRateMu.Lock() + defer smtpRateMu.Unlock() + + now := time.Now() + cutoff := now.Add(-smtpRateWindow) + + var recent []time.Time + for _, t := range smtpRateStore[key] { + if t.After(cutoff) { + recent = append(recent, t) + } + } + + if len(recent) >= smtpRateLimit { + smtpRateStore[key] = recent + return false + } + + smtpRateStore[key] = append(recent, now) + return true +} + +// --------------------------------------------------------------------------- +// GET /setting/smtp +// --------------------------------------------------------------------------- + +func handleSmtpGet(c *gin.Context) { + info := authorized.GetInfo(c) + locale := c.Query("locale") + if locale == "" { + locale = "en-us" + } + + presets := smtpGetPresets(locale) + + cfg := SmtpConfig{ + Enabled: false, + PresetKey: "custom", + Host: "", + Port: 465, + Encryption: "ssl", + Username: "", + Password: "", + FromName: "", + FromEmail: "", + Status: "unconfigured", + } + + hasSaved := false + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, smtpNS) + if saved != nil { + smtpLoadConfig(&cfg, saved) + hasSaved = true + } + } + + if !hasSaved { + if def := smtpDefaultPreset(presets); def != nil { + cfg.PresetKey = def.Key + cfg.Host = def.Host + cfg.Port = def.Port + cfg.Encryption = def.Encryption + } + } + + response.RespondWithSuccess(c, http.StatusOK, SmtpPageData{ + Presets: presets, + Config: cfg, + }) +} + +// --------------------------------------------------------------------------- +// PUT /setting/smtp +// --------------------------------------------------------------------------- + +func handleSmtpUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := smtpScope(info) + + var body struct { + PresetKey string `json:"preset_key"` + Host string `json:"host"` + Port int `json:"port"` + Encryption string `json:"encryption"` + Username string `json:"username"` + Password string `json:"password"` + FromName string `json:"from_name"` + FromEmail string `json:"from_email"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + existing, _ := setting.Global.Get(scope, smtpNS) + + pwd := body.Password + if pwd == "" { + if v, ok := existing["password"].(string); ok && v != "" { + pwd = cloudDecrypt(v) + } + } + + validated := false + if body.Host != "" && body.Username != "" && pwd != "" { + if err := smtpValidateConnection(body.Host, body.Port, body.Encryption, body.Username, pwd); err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + validated = true + } + + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + + m["preset_key"] = body.PresetKey + m["host"] = body.Host + m["port"] = body.Port + m["encryption"] = body.Encryption + m["username"] = body.Username + m["from_name"] = body.FromName + m["from_email"] = body.FromEmail + + if body.Password != "" { + m["password"] = cloudEncrypt(body.Password) + } + + if validated { + m["status"] = "connected" + } + + if _, err := setting.Global.Set(scope, smtpNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + cfg := SmtpConfig{ + PresetKey: "custom", + Port: 465, + Encryption: "ssl", + Status: "unconfigured", + } + smtpLoadConfig(&cfg, m) + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// PUT /setting/smtp/toggle +// --------------------------------------------------------------------------- + +func handleSmtpToggle(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := smtpScope(info) + + var body struct { + Enabled bool `json:"enabled"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + existing, _ := setting.Global.Get(scope, smtpNS) + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + m["enabled"] = body.Enabled + if !body.Enabled { + m["status"] = "unconfigured" + } + + if _, err := setting.Global.Set(scope, smtpNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + cfg := SmtpConfig{ + PresetKey: "custom", + Port: 465, + Encryption: "ssl", + Status: "unconfigured", + } + smtpLoadConfig(&cfg, m) + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// POST /setting/smtp/test +// --------------------------------------------------------------------------- + +func handleSmtpTest(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := smtpScope(info) + + var body struct { + ToEmail string `json:"to_email"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + if strings.TrimSpace(body.ToEmail) == "" { + respondError(c, http.StatusBadRequest, "to_email is required") + return + } + + rateKey := scope.TeamID + if rateKey == "" { + rateKey = scope.UserID + } + if !smtpCheckRateLimit(rateKey) { + response.RespondWithSuccess(c, http.StatusOK, SmtpTestResult{ + Success: false, + Message: "Rate limit exceeded, please wait a moment", + }) + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + saved, _ := setting.Global.Get(scope, smtpNS) + if saved == nil { + response.RespondWithSuccess(c, http.StatusOK, SmtpTestResult{ + Success: false, + Message: "SMTP not configured", + }) + return + } + + cfg := SmtpConfig{PresetKey: "custom", Port: 465, Encryption: "ssl", Status: "unconfigured"} + smtpLoadConfig(&cfg, saved) + + if cfg.Host == "" || cfg.Username == "" { + response.RespondWithSuccess(c, http.StatusOK, SmtpTestResult{ + Success: false, + Message: "SMTP host and username are required", + }) + return + } + + password := "" + if v, ok := saved["password"].(string); ok && v != "" { + password = cloudDecrypt(v) + } + + fromAddr := cfg.FromEmail + if fromAddr == "" { + fromAddr = cfg.Username + } + + err := smtpSendTestEmail(cfg.Host, cfg.Port, cfg.Encryption, cfg.Username, password, cfg.FromName, fromAddr, body.ToEmail) + if err != nil { + saved["status"] = "disconnected" + setting.Global.Set(scope, smtpNS, saved) + response.RespondWithSuccess(c, http.StatusOK, SmtpTestResult{ + Success: false, + Message: err.Error(), + }) + return + } + + saved["status"] = "connected" + saved["last_sent_at"] = time.Now().UTC().Format(time.RFC3339) + setting.Global.Set(scope, smtpNS, saved) + + response.RespondWithSuccess(c, http.StatusOK, SmtpTestResult{ + Success: true, + Message: "Test email sent successfully", + }) +} + +// --------------------------------------------------------------------------- +// SMTP connection validation (dial + auth, no email) +// --------------------------------------------------------------------------- + +func smtpValidateConnection(host string, port int, encryption, username, password string) error { + addr := net.JoinHostPort(host, fmt.Sprintf("%d", port)) + auth := smtp.PlainAuth("", username, password, host) + + switch encryption { + case "ssl": + tlsConfig := &tls.Config{ServerName: host} + conn, err := tls.DialWithDialer(&net.Dialer{Timeout: 10 * time.Second}, "tcp", addr, tlsConfig) + if err != nil { + return fmt.Errorf("SSL connection failed: %s", err.Error()) + } + defer conn.Close() + client, err := smtp.NewClient(conn, host) + if err != nil { + return fmt.Errorf("SMTP client failed: %s", err.Error()) + } + defer client.Quit() + if err = client.Auth(auth); err != nil { + return fmt.Errorf("authentication failed: %s", err.Error()) + } + return nil + + case "tls": + conn, err := net.DialTimeout("tcp", addr, 10*time.Second) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer conn.Close() + client, err := smtp.NewClient(conn, host) + if err != nil { + return fmt.Errorf("SMTP client failed: %s", err.Error()) + } + defer client.Quit() + if err = client.StartTLS(&tls.Config{ServerName: host}); err != nil { + return fmt.Errorf("STARTTLS failed: %s", err.Error()) + } + if err = client.Auth(auth); err != nil { + return fmt.Errorf("authentication failed: %s", err.Error()) + } + return nil + + default: + conn, err := net.DialTimeout("tcp", addr, 10*time.Second) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer conn.Close() + client, err := smtp.NewClient(conn, host) + if err != nil { + return fmt.Errorf("SMTP client failed: %s", err.Error()) + } + defer client.Quit() + if err = client.Auth(auth); err != nil { + return fmt.Errorf("authentication failed: %s", err.Error()) + } + return nil + } +} + +// --------------------------------------------------------------------------- +// SMTP send helper +// --------------------------------------------------------------------------- + +func smtpSendTestEmail(host string, port int, encryption, username, password, fromName, fromEmail, toEmail string) error { + addr := net.JoinHostPort(host, fmt.Sprintf("%d", port)) + + subject := "Yao SMTP Test" + body := "This is a test email from Yao to verify your SMTP configuration." + + from := fromEmail + if fromName != "" { + from = fmt.Sprintf("%s <%s>", fromName, fromEmail) + } + + msg := fmt.Sprintf("From: %s\r\nTo: %s\r\nSubject: %s\r\nMIME-Version: 1.0\r\nContent-Type: text/plain; charset=UTF-8\r\n\r\n%s", + from, toEmail, subject, body) + + auth := smtp.PlainAuth("", username, password, host) + + switch encryption { + case "ssl": + return smtpSendSSL(addr, host, auth, fromEmail, toEmail, []byte(msg)) + case "tls": + return smtpSendStartTLS(addr, host, auth, fromEmail, toEmail, []byte(msg)) + default: + return smtp.SendMail(addr, auth, fromEmail, []string{toEmail}, []byte(msg)) + } +} + +func smtpSendSSL(addr, host string, auth smtp.Auth, from, to string, msg []byte) error { + tlsConfig := &tls.Config{ServerName: host} + conn, err := tls.DialWithDialer(&net.Dialer{Timeout: 10 * time.Second}, "tcp", addr, tlsConfig) + if err != nil { + return fmt.Errorf("SSL connection failed: %s", err.Error()) + } + defer conn.Close() + + client, err := smtp.NewClient(conn, host) + if err != nil { + return fmt.Errorf("SMTP client failed: %s", err.Error()) + } + defer client.Quit() + + if err = client.Auth(auth); err != nil { + return fmt.Errorf("authentication failed: %s", err.Error()) + } + if err = client.Mail(from); err != nil { + return fmt.Errorf("MAIL FROM failed: %s", err.Error()) + } + if err = client.Rcpt(to); err != nil { + return fmt.Errorf("RCPT TO failed: %s", err.Error()) + } + + w, err := client.Data() + if err != nil { + return fmt.Errorf("DATA failed: %s", err.Error()) + } + if _, err = w.Write(msg); err != nil { + return fmt.Errorf("write failed: %s", err.Error()) + } + return w.Close() +} + +func smtpSendStartTLS(addr, host string, auth smtp.Auth, from, to string, msg []byte) error { + conn, err := net.DialTimeout("tcp", addr, 10*time.Second) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer conn.Close() + + client, err := smtp.NewClient(conn, host) + if err != nil { + return fmt.Errorf("SMTP client failed: %s", err.Error()) + } + defer client.Quit() + + if err = client.StartTLS(&tls.Config{ServerName: host}); err != nil { + return fmt.Errorf("STARTTLS failed: %s", err.Error()) + } + if err = client.Auth(auth); err != nil { + return fmt.Errorf("authentication failed: %s", err.Error()) + } + if err = client.Mail(from); err != nil { + return fmt.Errorf("MAIL FROM failed: %s", err.Error()) + } + if err = client.Rcpt(to); err != nil { + return fmt.Errorf("RCPT TO failed: %s", err.Error()) + } + + w, err := client.Data() + if err != nil { + return fmt.Errorf("DATA failed: %s", err.Error()) + } + if _, err = w.Write(msg); err != nil { + return fmt.Errorf("write failed: %s", err.Error()) + } + return w.Close() +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +func smtpLoadConfig(cfg *SmtpConfig, m map[string]interface{}) { + if v, ok := m["enabled"].(bool); ok { + cfg.Enabled = v + } + if v, ok := m["preset_key"].(string); ok && v != "" { + cfg.PresetKey = v + } + if v, ok := m["host"].(string); ok { + cfg.Host = v + } + if v, ok := m["port"]; ok { + switch p := v.(type) { + case int: + cfg.Port = p + case float64: + cfg.Port = int(p) + case int64: + cfg.Port = int(p) + } + } + if v, ok := m["encryption"].(string); ok && v != "" { + cfg.Encryption = v + } + if v, ok := m["username"].(string); ok { + cfg.Username = v + } + if v, ok := m["password"].(string); ok && v != "" { + cfg.Password = cloudMaskKey(cloudDecrypt(v)) + } + if v, ok := m["from_name"].(string); ok { + cfg.FromName = v + } + if v, ok := m["from_email"].(string); ok { + cfg.FromEmail = v + } + if v, ok := m["status"].(string); ok && v != "" { + cfg.Status = v + } + if v, ok := m["last_sent_at"].(string); ok && v != "" { + cfg.LastSentAt = v + } +} diff --git a/openapi/setting/smtp_presets.yml b/openapi/setting/smtp_presets.yml new file mode 100644 index 00000000..209494ab --- /dev/null +++ b/openapi/setting/smtp_presets.yml @@ -0,0 +1,121 @@ +zh-cn: + - key: tencent + name: 腾讯企邮 + host: smtp.exmail.qq.com + port: 465 + encryption: ssl + default: true + url: https://exmail.qq.com/ + hint: + zh-CN: "密码需使用客户端专用密码，在企业邮设置中生成" + en-US: "Use a client-specific password generated in Tencent Exmail settings" + + - key: feishu + name: 飞书邮箱 + host: smtp.feishu.cn + port: 465 + encryption: ssl + url: https://www.feishu.cn/ + hint: + zh-CN: "需在飞书管理后台开启邮箱 SMTP 服务" + en-US: "Enable SMTP in Feishu admin console" + + - key: aliyun + name: 阿里邮箱 + host: smtp.aliyun.com + port: 465 + encryption: ssl + url: https://mail.aliyun.com/ + hint: + zh-CN: "需在阿里邮箱设置中开启 SMTP 服务" + en-US: "Enable SMTP in Aliyun Mail settings" + + - key: qq + name: QQ邮箱 + host: smtp.qq.com + port: 465 + encryption: ssl + url: https://mail.qq.com/ + hint: + zh-CN: "需在QQ邮箱设置中开启 SMTP 服务并获取授权码" + en-US: "Enable SMTP in QQ Mail settings and get authorization code" + + - key: netease163 + name: 163邮箱 + host: smtp.163.com + port: 465 + encryption: ssl + url: https://mail.163.com/ + hint: + zh-CN: "需在163邮箱设置中开启 SMTP 服务并获取授权码" + en-US: "Enable SMTP in 163 Mail settings and get authorization code" + + - key: custom + name: 自定义 + host: "" + port: 465 + encryption: ssl + hint: + zh-CN: "手动填写 SMTP 服务器信息" + en-US: "Manually enter SMTP server details" + +en-us: + - key: gmail + name: Gmail + host: smtp.gmail.com + port: 465 + encryption: ssl + default: true + url: https://myaccount.google.com/apppasswords + hint: + zh-CN: "Gmail 需要专用密码（App Password），非登录密码" + en-US: "Gmail requires an App Password, not your login password" + + - key: yahoo + name: Yahoo Mail + host: smtp.mail.yahoo.com + port: 465 + encryption: ssl + url: https://login.yahoo.com/account/security + hint: + zh-CN: "Yahoo 需要应用专用密码" + en-US: "Yahoo requires an App Password generated in account security settings" + + - key: sendgrid + name: SendGrid + host: smtp.sendgrid.net + port: 587 + encryption: tls + url: https://app.sendgrid.com/ + hint: + zh-CN: "用户名固定为 apikey，密码填 API Key" + en-US: "Username is always \"apikey\", password is your API Key" + + - key: mailgun + name: Mailgun + host: smtp.mailgun.org + port: 587 + encryption: tls + url: https://app.mailgun.com/ + hint: + zh-CN: "在 Mailgun 控制台获取 SMTP 凭证" + en-US: "Get SMTP credentials from Mailgun dashboard" + + - key: ses + name: Amazon SES + host: email-smtp.us-east-1.amazonaws.com + port: 587 + encryption: tls + url: https://console.aws.amazon.com/ses/ + hint: + zh-CN: "需在 AWS SES 控制台创建 SMTP 凭证，非 IAM 密钥" + en-US: "Create SMTP credentials in AWS SES console, not IAM keys" + + - key: custom + name: Custom + host: "" + port: 465 + encryption: ssl + hint: + zh-CN: "手动填写 SMTP 服务器信息" + en-US: "Manually enter SMTP server details" diff --git a/openapi/setting/system.go b/openapi/setting/system.go new file mode 100644 index 00000000..1a46169a --- /dev/null +++ b/openapi/setting/system.go @@ -0,0 +1,255 @@ +package setting + +import ( + _ "embed" + "encoding/json" + "fmt" + "net/http" + "runtime" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/commercial" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/share" + "gopkg.in/yaml.v3" +) + +const cdnBase = "https://get.yaoapps.com/yao" + +// update check cache (package-level, protected by mutex) +var ( + updateCache *CheckUpdateResult + updateCacheTime time.Time + updateMu sync.Mutex + cacheTTL = 10 * time.Minute +) + +// handleSystemInfo returns aggregated system information. +// GET /setting/system?locale=zh-cn +func handleSystemInfo(c *gin.Context) { + locale := strings.ToLower(c.DefaultQuery("locale", "en-us")) + + env := share.App.Option["env"] + environment, _ := env.(string) + if environment == "" { + environment = config.Conf.Mode + } + if environment == "" { + environment = "production" + } + + listen := fmt.Sprintf("%s:%d", config.Conf.Host, config.Conf.Port) + sessionStore := config.Conf.Session.Store + if sessionStore == "" { + sessionStore = "file" + } + + lang := langFromLocale(locale) + lic := commercial.License + deployment := lic.Edition + if deployment == "" { + deployment = "community" + } + + var licenseKey string + if lic.Valid && lic.SerialNumber != "" { + licenseKey = lic.SerialNumber + } + + data := SystemInfoData{ + App: AppInfo{ + Name: share.App.Name, + Short: share.App.Short, + Description: share.App.Description, + Logo: "/api/__yao/app/icons/app.png", + Version: share.App.Version, + }, + Deployment: deployment, + DeploymentLabel: resolveLabel(promFile.Labels.Deployment, deployment, lang, deployment), + LicenseKey: licenseKey, + Environment: environment, + EnvironmentLabel: resolveLabel(promFile.Labels.Environment, environment, lang, environment), + Server: VersionInfo{ + Version: share.VERSION, + BuildDate: share.PRVERSION, + CommitSHA: share.PRVERSION, + }, + Client: VersionInfo{ + Version: share.CUI, + BuildDate: share.PRCUI, + CommitSHA: share.PRCUI, + }, + Technical: TechnicalInfo{ + Listen: listen, + DBDriver: config.Conf.DB.Driver, + SessionStore: sessionStore, + }, + Promotions: buildPromotions(deployment, locale), + } + + response.RespondWithSuccess(c, http.StatusOK, data) +} + +//go:embed promotions.yml +var promotionsYML []byte + +type promotionEntry struct { + ID string `yaml:"id"` + Link string `yaml:"link"` + I18n map[string]promotionLocale `yaml:"i18n"` +} + +type promotionLocale struct { + Title string `yaml:"title"` + Desc string `yaml:"desc"` + Label string `yaml:"label"` +} + +type promotionsFile struct { + Labels struct { + Deployment map[string]map[string]string `yaml:"deployment"` + Environment map[string]map[string]string `yaml:"environment"` + } `yaml:"labels"` + Community []promotionEntry `yaml:"community"` + Enterprise []promotionEntry `yaml:"enterprise"` + Cloud []promotionEntry `yaml:"cloud"` +} + +var promFile promotionsFile + +func init() { + yaml.Unmarshal(promotionsYML, &promFile) +} + +func resolveLabel(m map[string]map[string]string, key, lang, fallback string) string { + if langs, ok := m[key]; ok { + if v, ok := langs[lang]; ok { + return v + } + if v, ok := langs["en"]; ok { + return v + } + } + return fallback +} + +func langFromLocale(locale string) string { + if strings.HasPrefix(locale, "zh") { + return "zh" + } + return "en" +} + +func buildPromotions(deployment, locale string) []Promotion { + lang := langFromLocale(locale) + + var entries []promotionEntry + switch deployment { + case "community": + entries = promFile.Community + case "enterprise": + entries = promFile.Enterprise + case "cloud": + entries = promFile.Cloud + } + if len(entries) == 0 { + return nil + } + + promos := make([]Promotion, 0, len(entries)) + for _, e := range entries { + loc, ok := e.I18n[lang] + if !ok { + loc = e.I18n["en"] + } + promos = append(promos, Promotion{ + ID: e.ID, + Title: loc.Title, + Desc: loc.Desc, + Link: e.Link, + Label: loc.Label, + }) + } + return promos +} + +// handleSystemCheckUpdate checks for a newer engine release. +// Uses the same CDN source as `yao upgrade` and yao-desktop: +// +// GET https://get.yaoapps.com/yao/latest.json +// +// POST /setting/system/check-update +func handleSystemCheckUpdate(c *gin.Context) { + updateMu.Lock() + if updateCache != nil && time.Since(updateCacheTime) < cacheTTL { + result := *updateCache + updateMu.Unlock() + response.RespondWithSuccess(c, http.StatusOK, result) + return + } + updateMu.Unlock() + + result := fetchLatestVersion() + + updateMu.Lock() + updateCache = &result + updateCacheTime = time.Now() + updateMu.Unlock() + + response.RespondWithSuccess(c, http.StatusOK, result) +} + +// cdnLatest mirrors the JSON structure of get.yaoapps.com/yao/latest.json +// (same format used by cmd/upgrade.go and yao-desktop updater.rs). +type cdnLatest struct { + Version string `json:"version"` + ReleasedAt string `json:"released_at"` + Assets map[string]string `json:"assets"` +} + +func fetchLatestVersion() CheckUpdateResult { + current := strings.TrimPrefix(share.VERSION, "v") + base := CheckUpdateResult{HasUpdate: false, CurrentVersion: current} + + url := cdnBase + "/latest.json" + client := &http.Client{Timeout: 15 * time.Second} + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return base + } + req.Header.Set("User-Agent", fmt.Sprintf("yao/%s", share.VERSION)) + + resp, err := client.Do(req) + if err != nil { + return base + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return base + } + + var data cdnLatest + if err := json.NewDecoder(resp.Body).Decode(&data); err != nil { + return base + } + + latest := strings.TrimPrefix(data.Version, "v") + if latest == "" { + return base + } + + platformKey := fmt.Sprintf("%s-%s", runtime.GOOS, runtime.GOARCH) + downloadURL := data.Assets[platformKey] + + return CheckUpdateResult{ + HasUpdate: latest != current, + CurrentVersion: current, + LatestVersion: latest, + DownloadURL: downloadURL, + } +} diff --git a/openapi/setting/types.go b/openapi/setting/types.go new file mode 100644 index 00000000..2cfd044f --- /dev/null +++ b/openapi/setting/types.go @@ -0,0 +1,221 @@ +package setting + +// SystemInfoData is the top-level response for GET /setting/system. +type SystemInfoData struct { + App AppInfo `json:"app"` + Deployment string `json:"deployment"` + DeploymentLabel string `json:"deployment_label"` + LicenseKey string `json:"license_key,omitempty"` + Server VersionInfo `json:"server"` + Client VersionInfo `json:"client"` + Environment string `json:"environment"` + EnvironmentLabel string `json:"environment_label"` + Technical TechnicalInfo `json:"technical"` + Promotions []Promotion `json:"promotions,omitempty"` +} + +// Promotion is a localized CTA banner returned by the API. +type Promotion struct { + ID string `json:"id"` + Title string `json:"title"` + Desc string `json:"desc"` + Link string `json:"link"` + Label string `json:"label"` +} + +// AppInfo describes the running application. +type AppInfo struct { + Name string `json:"name"` + Short string `json:"short"` + Description string `json:"description"` + Logo string `json:"logo"` + Version string `json:"version"` +} + +// VersionInfo carries build metadata for a component (engine / CUI). +type VersionInfo struct { + Version string `json:"version"` + BuildDate string `json:"build_date"` + CommitSHA string `json:"commit"` +} + +// TechnicalInfo contains runtime / infrastructure details. +type TechnicalInfo struct { + Listen string `json:"listen"` + DBDriver string `json:"db_driver"` + SessionStore string `json:"session_store"` +} + +// CheckUpdateResult is the response for POST /setting/system/check-update. +type CheckUpdateResult struct { + HasUpdate bool `json:"has_update"` + CurrentVersion string `json:"current_version"` + LatestVersion string `json:"latest_version,omitempty"` + DownloadURL string `json:"download_url,omitempty"` +} + +// --------------------------------------------------------------------------- +// Cloud Service +// --------------------------------------------------------------------------- + +// CloudRegion is a static entry loaded from cloud_presets.yml. +type CloudRegion struct { + Key string `json:"key" yaml:"key"` + Label map[string]string `json:"label" yaml:"label"` + APIURL string `json:"api_url" yaml:"api_url"` + Default bool `json:"default,omitempty" yaml:"default"` +} + +// CloudPageData is the response for GET /setting/cloud. +type CloudPageData struct { + Regions []CloudRegion `json:"regions"` + Region string `json:"region"` + APIURL string `json:"api_url"` + APIKey string `json:"api_key"` + Status string `json:"status"` +} + +// CloudTestResult is the response for POST /setting/cloud/test. +type CloudTestResult struct { + Success bool `json:"success"` + Message string `json:"message"` + LatencyMs int64 `json:"latency_ms,omitempty"` +} + +// --------------------------------------------------------------------------- +// LLM Providers +// --------------------------------------------------------------------------- + +// LLMPageData is the aggregated response for GET /setting/llm. +type LLMPageData struct { + Providers []interface{} `json:"providers"` + Roles map[string]interface{} `json:"roles"` + PresetProviders []interface{} `json:"preset_providers"` +} + +// --------------------------------------------------------------------------- +// Search & Scrape +// --------------------------------------------------------------------------- + +type SearchProviderField struct { + Key string `json:"key" yaml:"key"` + Label map[string]string `json:"label" yaml:"label"` + Type string `json:"type" yaml:"type"` + Default string `json:"default,omitempty" yaml:"default"` + Placeholder string `json:"placeholder,omitempty" yaml:"placeholder"` + Hint map[string]string `json:"hint,omitempty" yaml:"hint"` +} + +type SearchProviderPreset struct { + Key string `json:"key" yaml:"key"` + Name string `json:"name" yaml:"name"` + Description map[string]string `json:"description,omitempty" yaml:"description"` + Website string `json:"website,omitempty" yaml:"website"` + Tools []string `json:"tools" yaml:"tools"` + ToolLabels []map[string]string `json:"tool_labels" yaml:"tool_labels"` + Fields []SearchProviderField `json:"fields" yaml:"fields"` + IsCloud bool `json:"is_cloud,omitempty" yaml:"is_cloud"` +} + +type SearchProviderConfig struct { + PresetKey string `json:"preset_key"` + Enabled bool `json:"enabled"` + FieldValues map[string]string `json:"field_values"` + Status string `json:"status"` +} + +type SearchToolAssignment struct { + WebSearch *string `json:"web_search"` + WebScrape *string `json:"web_scrape"` +} + +type SearchPageData struct { + Presets []SearchProviderPreset `json:"presets"` + Providers []SearchProviderConfig `json:"providers"` + ToolAssignment SearchToolAssignment `json:"tool_assignment"` +} + +type SearchTestResult struct { + Success bool `json:"success"` + Message string `json:"message"` + LatencyMs int64 `json:"latency_ms,omitempty"` +} + +// --------------------------------------------------------------------------- +// SMTP +// --------------------------------------------------------------------------- + +type SmtpPreset struct { + Key string `json:"key" yaml:"key"` + Name string `json:"name" yaml:"name"` + Host string `json:"host" yaml:"host"` + Port int `json:"port" yaml:"port"` + Encryption string `json:"encryption" yaml:"encryption"` + Hint map[string]string `json:"hint,omitempty" yaml:"hint"` + URL string `json:"url,omitempty" yaml:"url"` + Default bool `json:"default,omitempty" yaml:"default"` +} + +type SmtpConfig struct { + Enabled bool `json:"enabled"` + PresetKey string `json:"preset_key"` + Host string `json:"host"` + Port int `json:"port"` + Encryption string `json:"encryption"` + Username string `json:"username"` + Password string `json:"password"` + FromName string `json:"from_name"` + FromEmail string `json:"from_email"` + Status string `json:"status"` + LastSentAt string `json:"last_sent_at,omitempty"` +} + +type SmtpPageData struct { + Presets []SmtpPreset `json:"presets"` + Config SmtpConfig `json:"config"` +} + +type SmtpTestResult struct { + Success bool `json:"success"` + Message string `json:"message"` +} + +// --------------------------------------------------------------------------- +// Sandbox +// --------------------------------------------------------------------------- + +type ComputerNode struct { + NodeID string `json:"node_id"` + DisplayName string `json:"display_name"` + Kind string `json:"kind"` + OS string `json:"os"` + Arch string `json:"arch"` + CPU int `json:"cpu"` + MemoryGB int `json:"memory_gb"` + DockerVersion string `json:"docker_version,omitempty"` + RunningSandboxes int `json:"running_sandboxes"` + Online bool `json:"online"` +} + +type SandboxRegistryConfig struct { + RegistryURL string `json:"registry_url"` + Username string `json:"username"` + Password string `json:"password"` +} + +type SandboxImage struct { + ID string `json:"id"` + AssistantNames []string `json:"assistant_names"` + ImageName string `json:"image_name"` + Tag string `json:"tag"` + SizeMB int `json:"size_mb"` + Status string `json:"status"` + Progress *int `json:"progress,omitempty"` + ErrorMessage string `json:"error_message,omitempty"` +} + +type SandboxPageData struct { + Nodes []ComputerNode `json:"nodes"` + Registry SandboxRegistryConfig `json:"registry"` + Images map[string][]SandboxImage `json:"images"` +} diff --git a/openapi/tests/setting/cloud_test.go b/openapi/tests/setting/cloud_test.go new file mode 100644 index 00000000..099d98f6 --- /dev/null +++ b/openapi/tests/setting/cloud_test.go @@ -0,0 +1,330 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/openapi/tests/testutils" + "github.com/yaoapp/yao/setting" +) + +func initSettingRegistry(t *testing.T) { + t.Helper() + if setting.Global == nil { + if err := setting.Init(); err != nil { + t.Fatalf("setting.Init: %v", err) + } + } +} + +func obtainToken(t *testing.T, serverURL string) string { + t.Helper() + client := testutils.RegisterTestClient(t, "Cloud Test", []string{"https://localhost/callback"}) + t.Cleanup(func() { testutils.CleanupTestClient(t, client.ClientID) }) + token := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + return token.AccessToken +} + +// obtainRestrictedToken creates a token with specific scope (no system:root). +// Used to test ACL permission denial. +func obtainRestrictedToken(t *testing.T, serverURL, scope string) string { + t.Helper() + client := testutils.RegisterTestClient(t, "Cloud Restricted", []string{"https://localhost/callback"}) + t.Cleanup(func() { testutils.CleanupTestClient(t, client.ClientID) }) + + oauthService := oauth.OAuth + if oauthService == nil { + t.Fatal("Global OAuth service not initialized") + } + + token := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + subject, err := oauthService.Subject(client.ClientID, token.UserID) + if err != nil { + t.Fatalf("Failed to create subject: %v", err) + } + + accessToken, err := oauthService.MakeAccessToken(client.ClientID, scope, subject, 3600) + if err != nil { + t.Fatalf("Failed to create access token: %v", err) + } + return accessToken +} + +// ----------- Functional tests (system:root token) ----------- + +func TestCloudGet(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/cloud", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + assert.Contains(t, body, "regions") + assert.Contains(t, body, "region") + assert.Contains(t, body, "api_url") + assert.Contains(t, body, "api_key") + assert.Contains(t, body, "status") + + regions, ok := body["regions"].([]interface{}) + assert.True(t, ok) + assert.GreaterOrEqual(t, len(regions), 4) + + assert.Equal(t, "unconfigured", body["status"]) + assert.Equal(t, "", body["api_key"]) +} + +func TestCloudGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/cloud", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestCloudUpdate(t *testing.T) { + apiKey := os.Getenv("CLOUD_TEST_API_KEY") + if apiKey == "" { + t.Skip("CLOUD_TEST_API_KEY not set, skipping cloud update test (key validation required)") + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "region": "us", + "api_url": "https://api-us.yao.run", + "api_key": apiKey, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + assert.Equal(t, "us", body["region"]) + assert.Equal(t, "https://api-us.yao.run", body["api_url"]) + assert.Equal(t, "connected", body["status"]) + + maskedKey, _ := body["api_key"].(string) + assert.True(t, strings.Contains(maskedKey, "..."), "masked key should use prefix...suffix format") + + // GET should also return masked key and connected status + req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/cloud", nil) + req2.Header.Set("Authorization", "Bearer "+token) + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + + var body2 map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body2) + assert.Equal(t, "us", body2["region"]) + assert.Equal(t, "connected", body2["status"]) +} + +func TestCloudUpdateInvalidKey(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "region": "us", + "api_url": "https://api-us.yao.run", + "api_key": "sk-invalid-key-that-should-fail", + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "invalid API key should be rejected") +} + +func TestCloudUpdateInvalidRegion(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "region": "mars", + "api_url": "https://api-mars.yao.run", + "api_key": "sk-test", + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +func TestCloudTest(t *testing.T) { + apiKey := os.Getenv("CLOUD_TEST_API_KEY") + if apiKey == "" { + t.Skip("CLOUD_TEST_API_KEY not set, skipping cloud connection test") + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save config first (key is validated during save) + payload := map[string]interface{}{ + "region": "us", + "api_url": "https://api-us.yao.run", + "api_key": apiKey, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + // Test connection with explicit api_url and api_key + testPayload := map[string]interface{}{ + "api_url": "https://api-us.yao.run", + "api_key": apiKey, + } + testRaw, _ := json.Marshal(testPayload) + req2, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/cloud/test", bytes.NewReader(testRaw)) + assert.NoError(t, err) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + assert.Equal(t, true, body["success"]) + assert.NotEmpty(t, body["message"]) +} + +// ----------- ACL permission tests ----------- + +func TestCloudACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + + // Token with read-only scope (no system:root, only setting:cloud:read:all) + readToken := obtainRestrictedToken(t, serverURL, "setting:cloud:read:all") + + // GET should work + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/cloud", nil) + req.Header.Set("Authorization", "Bearer "+readToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET") + + // PUT should be denied + payload := map[string]interface{}{ + "region": "cn", + "api_url": "https://api.yaoagents.cn", + "api_key": "sk-test", + } + raw, _ := json.Marshal(payload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+readToken) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny PUT") +} + +func TestCloudACL_NoScopeCannotAccess(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + + // Token with irrelevant scope (no setting scopes at all) + noSettingToken := obtainRestrictedToken(t, serverURL, "kb:collections:read:all") + + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/cloud", nil) + req.Header.Set("Authorization", "Bearer "+noSettingToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusForbidden, resp.StatusCode, "token without setting scope should be denied") +} + +func TestCloudUpdateRegionOnly(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "region": "cn", + "api_url": "https://api.yaoagents.cn", + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "update without api_key should succeed (no validation needed)") + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, "cn", body["region"]) + assert.Equal(t, "https://api.yaoagents.cn", body["api_url"]) +} diff --git a/openapi/tests/setting/llm_test.go b/openapi/tests/setting/llm_test.go new file mode 100644 index 00000000..993b3cfe --- /dev/null +++ b/openapi/tests/setting/llm_test.go @@ -0,0 +1,453 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/llmprovider" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func requireOpenAIKey(t *testing.T) string { + t.Helper() + key := os.Getenv("OPENAI_TEST_KEY") + if key == "" { + t.Skip("OPENAI_TEST_KEY not set") + } + return key +} + +func initLLMRegistry(t *testing.T) { + t.Helper() + if err := llmprovider.Init(); err != nil { + t.Fatalf("llmprovider.Init: %v", err) + } + if config.Conf.DB.AESKey != "" { + llmprovider.Global.SetEncryptionKey(config.Conf.DB.AESKey) + } +} + +func llmURL(serverURL, path string) string { + return serverURL + baseURL() + "/setting/llm" + path +} + +func llmGet(t *testing.T, url, token string) *http.Response { + t.Helper() + req, _ := http.NewRequest("GET", url, nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + return resp +} + +func llmPost(t *testing.T, url, token string, payload interface{}) *http.Response { + t.Helper() + var body io.Reader + if payload != nil { + raw, _ := json.Marshal(payload) + body = bytes.NewReader(raw) + } + req, _ := http.NewRequest("POST", url, body) + req.Header.Set("Authorization", "Bearer "+token) + if payload != nil { + req.Header.Set("Content-Type", "application/json") + } + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + return resp +} + +func llmPut(t *testing.T, url, token string, payload interface{}) *http.Response { + t.Helper() + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("PUT", url, bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + return resp +} + +func llmDelete(t *testing.T, url, token string) *http.Response { + t.Helper() + req, _ := http.NewRequest("DELETE", url, nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + return resp +} + +func llmBody(t *testing.T, resp *http.Response) map[string]interface{} { + t.Helper() + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + return body +} + +func createTestOpenAI(t *testing.T, serverURL, token 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() + assert.Equal(t, http.StatusCreated, resp.StatusCode, "createTestOpenAI should succeed") + t.Cleanup(func() { llmprovider.Global.Delete("openai") }) +} + +// ----------- Functional tests ----------- + +func TestLLMGetPageData(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + + createTestOpenAI(t, serverURL, token) + + resp := llmGet(t, llmURL(serverURL, ""), token) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + body := llmBody(t, resp) + assert.Contains(t, body, "providers") + assert.Contains(t, body, "roles") + assert.Contains(t, body, "preset_providers") + + providers, ok := body["providers"].([]interface{}) + assert.True(t, ok, "providers should be an array") + assert.GreaterOrEqual(t, len(providers), 1) + + if len(providers) > 0 { + p := providers[0].(map[string]interface{}) + assert.Contains(t, p, "key") + assert.Contains(t, p, "name") + assert.Contains(t, p, "models") + assert.NotContains(t, p, "connector_id", "internal field should be stripped") + assert.NotContains(t, p, "source", "internal field should be stripped") + assert.NotContains(t, p, "owner", "internal field should be stripped") + } + + presets, ok := body["preset_providers"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 5, len(presets), "should have 5 presets") +} + +func TestLLMGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, _ := http.NewRequest("GET", llmURL(serverURL, ""), nil) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestLLMProviderCreate(t *testing.T) { + realKey := requireOpenAIKey(t) + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + llmprovider.Global.Delete("openai") + + payload := map[string]interface{}{ + "preset_key": "openai", + "api_key": realKey, + "model_ids": []string{"gpt-4o"}, + } + 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"]) + assert.Equal(t, "OpenAI", body["name"]) + assert.Equal(t, "openai", body["type"]) + + apiKey, _ := body["api_key"].(string) + assert.NotEqual(t, realKey, apiKey, "API key should be masked") + assert.NotEmpty(t, apiKey) + + models, _ := body["models"].([]interface{}) + assert.Equal(t, 1, len(models)) +} + +func TestLLMProviderCreateCustom(t *testing.T) { + realKey := requireOpenAIKey(t) + mirror := os.Getenv("TEST_MOAPI_MIRROR") + if mirror == "" { + mirror = "https://api.openai.com" + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + llmprovider.Global.Delete("my-custom-llm") + + payload := map[string]interface{}{ + "key": "my-custom-llm", + "name": "My Custom LLM", + "type": "openai", + "api_url": mirror, + "api_key": realKey, + "models": []map[string]interface{}{ + {"id": "custom-model", "name": "Custom Model", "capabilities": []string{"streaming"}, "enabled": true}, + }, + "require_key": true, + } + 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"]) + assert.Equal(t, "My Custom LLM", body["name"]) + assert.Equal(t, true, body["is_custom"]) + + models, _ := body["models"].([]interface{}) + assert.Equal(t, 1, len(models)) +} + +func TestLLMProviderUpdate(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + + createTestOpenAI(t, serverURL, token) + + updatePayload := map[string]interface{}{ + "name": "Updated OpenAI", + "api_url": "https://api.openai.com/v2", + "models": []map[string]interface{}{ + {"id": "gpt-4o", "name": "GPT-4o Updated", "capabilities": []string{"vision", "tool_calls"}, "enabled": true}, + }, + } + resp := llmPut(t, llmURL(serverURL, "/providers/openai"), token, updatePayload) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + body := llmBody(t, resp) + assert.Equal(t, "Updated OpenAI", body["name"]) + assert.Equal(t, "https://api.openai.com/v2", body["api_url"]) + + apiKey, _ := body["api_key"].(string) + assert.NotEmpty(t, apiKey, "API key should be preserved when not sent") +} + +func TestLLMProviderDelete(t *testing.T) { + anthropicKey := os.Getenv("ANTHROPIC_API_KEY") + if anthropicKey == "" { + t.Skip("ANTHROPIC_API_KEY not set") + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + 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) + createResp.Body.Close() + assert.Equal(t, http.StatusCreated, createResp.StatusCode) + + rolesPayload := map[string]interface{}{ + "default": map[string]interface{}{ + "provider": "anthropic", + "model": "claude-sonnet-4-20250514", + }, + } + rolesResp := llmPut(t, llmURL(serverURL, "/roles"), token, rolesPayload) + rolesResp.Body.Close() + assert.Equal(t, http.StatusOK, rolesResp.StatusCode) + + deleteResp := llmDelete(t, llmURL(serverURL, "/providers/anthropic"), token) + defer deleteResp.Body.Close() + assert.Equal(t, http.StatusOK, deleteResp.StatusCode) + + body := llmBody(t, deleteResp) + assert.Equal(t, true, body["success"]) + assert.NotEmpty(t, body["warning"], "should warn about cleared roles") + + getResp := llmGet(t, llmURL(serverURL, ""), token) + defer getResp.Body.Close() + getBody := llmBody(t, getResp) + roles, _ := getBody["roles"].(map[string]interface{}) + assert.NotContains(t, roles, "default", "role referencing deleted provider should be cleared") +} + +func TestLLMProviderDeleteForbidden(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + + llmprovider.Global.Delete("other-team-provider") + otherProvider := &llmprovider.Provider{ + Key: "other-team-provider", + 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"}, + } + llmprovider.Global.Create(otherProvider) + t.Cleanup(func() { llmprovider.Global.Delete("other-team-provider") }) + + resp := llmDelete(t, llmURL(serverURL, "/providers/other-team-provider"), token) + defer resp.Body.Close() + assert.Equal(t, http.StatusNotFound, resp.StatusCode, "should not be able to delete another user's provider") +} + +func TestLLMProviderTest(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + + createTestOpenAI(t, serverURL, token) + + resp := llmPost(t, llmURL(serverURL, "/providers/openai/test"), token, nil) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + body := llmBody(t, resp) + assert.Equal(t, true, body["success"]) + assert.NotEmpty(t, body["message"]) + assert.NotNil(t, body["latency_ms"]) +} + +func TestLLMRoles(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + createTestOpenAI(t, serverURL, token) + + rolesPayload := map[string]interface{}{ + "default": map[string]interface{}{ + "provider": "openai", + "model": "gpt-4o", + }, + "vision": map[string]interface{}{ + "provider": "openai", + "model": "gpt-4o", + }, + } + resp := llmPut(t, llmURL(serverURL, "/roles"), token, rolesPayload) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + body := llmBody(t, resp) + assert.Contains(t, body, "default") + assert.Contains(t, body, "vision") + + getResp := llmGet(t, llmURL(serverURL, ""), token) + defer getResp.Body.Close() + getBody := llmBody(t, getResp) + roles, _ := getBody["roles"].(map[string]interface{}) + assert.Contains(t, roles, "default") + assert.Contains(t, roles, "vision") +} + +func TestLLMRolesValidation(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + token := obtainToken(t, serverURL) + + resp1 := llmPut(t, llmURL(serverURL, "/roles"), token, map[string]interface{}{ + "vision": map[string]interface{}{ + "provider": "openai", + "model": "gpt-4o", + }, + }) + defer resp1.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp1.StatusCode, "should require 'default' role") + + resp2 := llmPut(t, llmURL(serverURL, "/roles"), token, map[string]interface{}{ + "default": map[string]interface{}{ + "provider": "nonexistent-provider", + "model": "some-model", + }, + }) + defer resp2.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp2.StatusCode, "should reject non-existent provider") + + createTestOpenAI(t, serverURL, token) + + resp3 := llmPut(t, llmURL(serverURL, "/roles"), token, map[string]interface{}{ + "default": map[string]interface{}{ + "provider": "openai", + "model": "nonexistent-model", + }, + }) + defer resp3.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp3.StatusCode, "should reject non-existent model") +} + +// ----------- ACL permission tests ----------- + +func TestLLMACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + + readToken := obtainRestrictedToken(t, serverURL, "setting:llm:read:all") + + resp := llmGet(t, llmURL(serverURL, ""), readToken) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET") + + createPayload := map[string]interface{}{ + "preset_key": "openai", + "api_key": "sk-test", + } + resp2 := llmPost(t, llmURL(serverURL, "/providers"), readToken, createPayload) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny POST") +} + +func TestLLMACL_NoScopeCannotAccess(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initLLMRegistry(t) + + noSettingToken := obtainRestrictedToken(t, serverURL, "kb:collections:read:all") + + resp := llmGet(t, llmURL(serverURL, ""), noSettingToken) + defer resp.Body.Close() + assert.Equal(t, http.StatusForbidden, resp.StatusCode, "token without setting scope should be denied") +} diff --git a/openapi/tests/setting/mcp_test.go b/openapi/tests/setting/mcp_test.go new file mode 100644 index 00000000..d1bbf6ea --- /dev/null +++ b/openapi/tests/setting/mcp_test.go @@ -0,0 +1,424 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + mcpTypes "github.com/yaoapp/gou/mcp/types" + gouTypes "github.com/yaoapp/gou/types" + "github.com/yaoapp/yao/mcpclient" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func initMcpClientRegistry(t *testing.T) { + t.Helper() + if mcpclient.Global == nil { + if err := mcpclient.Init(); err != nil { + t.Fatalf("mcpclient.Init: %v", err) + } + } +} + +func obtainTokenInfo(t *testing.T, serverURL string) *testutils.TokenInfo { + t.Helper() + client := testutils.RegisterTestClient(t, "MCP Test", []string{"https://localhost/callback"}) + t.Cleanup(func() { testutils.CleanupTestClient(t, client.ClientID) }) + return testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") +} + +func seedMCPServer(t *testing.T, ownerID, name, url string) string { + t.Helper() + clientID := "user." + ownerID + "." + name + client := &mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: clientID, + Name: name, + Transport: mcpTypes.TransportHTTP, + URL: url, + Timeout: "30s", + MetaInfo: gouTypes.MetaInfo{Label: name}, + }, + Enabled: true, + Status: "connected", + Source: mcpclient.ClientSourceDynamic, + Owner: mcpclient.ClientOwner{Type: "user", ID: ownerID}, + } + _, err := mcpclient.Global.Create(client) + if err != nil && !strings.Contains(err.Error(), "already exists") { + t.Fatalf("seedMCPServer: %v", err) + } + return clientID +} + +// startMockMCPServer starts a minimal MCP-compatible HTTP server for testing. +// Handles JSON-RPC: initialize, notifications/initialized, tools/list. +func startMockMCPServer(t *testing.T) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, _ := io.ReadAll(r.Body) + defer r.Body.Close() + + var req struct { + JSONRPC string `json:"jsonrpc"` + ID interface{} `json:"id,omitempty"` + Method string `json:"method"` + } + json.Unmarshal(body, &req) + + w.Header().Set("Content-Type", "application/json") + + switch req.Method { + case "initialize": + json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]interface{}{ + "protocolVersion": "2025-03-26", + "serverInfo": map[string]interface{}{"name": "mock-mcp", "version": "1.0.0"}, + "capabilities": map[string]interface{}{"tools": map[string]interface{}{}}, + }, + }) + case "notifications/initialized": + w.WriteHeader(http.StatusAccepted) + case "tools/list": + json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]interface{}{ + "tools": []interface{}{ + map[string]interface{}{ + "name": "echo", + "description": "Echo tool", + "inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}, + }, + }, + }, + }) + default: + json.NewEncoder(w).Encode(map[string]interface{}{ + "jsonrpc": "2.0", + "id": req.ID, + "error": map[string]interface{}{"code": -32601, "message": "method not found"}, + }) + } + })) +} + +func TestMCPListServers(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + token := obtainToken(t, serverURL) + + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Contains(t, body, "servers") + servers, ok := body["servers"].([]interface{}) + assert.True(t, ok) + t.Logf("Listed %d MCP servers", len(servers)) +} + +func TestMCPListUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestMCPCreateServer(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + mockMCP := startMockMCPServer(t) + defer mockMCP.Close() + + payload := map[string]interface{}{ + "name": "test-create", + "label": "Test Create", + "transport": "http", + "url": mockMCP.URL, + "timeout": "10s", + } + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/mcp/servers", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + createdID, _ := body["id"].(string) + assert.NotEmpty(t, createdID) + assert.Equal(t, "test-create", body["name"]) + assert.Equal(t, "Test Create", body["label"]) + assert.Equal(t, "connected", body["status"]) + t.Logf("Created server: %s", createdID) + + // Verify in list + listReq, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + listReq.Header.Set("Authorization", "Bearer "+ti.AccessToken) + listResp, _ := http.DefaultClient.Do(listReq) + var listBody map[string]interface{} + json.NewDecoder(listResp.Body).Decode(&listBody) + listResp.Body.Close() + + found := false + for _, s := range listBody["servers"].([]interface{}) { + if s.(map[string]interface{})["id"] == createdID { + found = true + } + } + assert.True(t, found, "created server should appear in list") + + // Cleanup + mcpclient.Global.Delete(createdID) +} + +func TestMCPCreateRejectsUnreachable(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "name": "unreachable", + "label": "Unreachable", + "transport": "http", + "url": "https://192.0.2.1/mcp", + "timeout": "3s", + } + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/mcp/servers", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "create should reject unreachable URL") +} + +func TestMCPDuplicateName(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + clientID := seedMCPServer(t, ti.UserID, "dup-test", "https://example.com/mcp") + defer mcpclient.Global.Delete(clientID) + + payload := map[string]interface{}{ + "name": "dup-test", + "label": "Duplicate", + "transport": "http", + "url": "https://example.com/mcp", + } + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/mcp/servers", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +func TestMCPUpdateServer(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + mockMCP := startMockMCPServer(t) + defer mockMCP.Close() + + clientID := seedMCPServer(t, ti.UserID, "upd-test", "https://example.com/mcp") + defer mcpclient.Global.Delete(clientID) + + payload := map[string]interface{}{ + "label": "Updated Label", + "url": mockMCP.URL, + } + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/mcp/servers/"+clientID, bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + t.Logf("Update response (%d): %s", resp.StatusCode, string(respBody)) + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.Unmarshal(respBody, &body) + assert.Equal(t, "Updated Label", body["label"]) + assert.Equal(t, mockMCP.URL, body["url"]) + assert.Equal(t, "connected", body["status"]) +} + +func TestMCPUpdateRejectsUnreachable(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + clientID := seedMCPServer(t, ti.UserID, "upd-fail", "https://example.com/mcp") + defer mcpclient.Global.Delete(clientID) + + payload := map[string]interface{}{ + "url": "https://192.0.2.1/mcp", + } + raw, _ := json.Marshal(payload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/mcp/servers/"+clientID, bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "update should reject unreachable URL") +} + +func TestMCPTokenMasking(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + clientID := "user." + ti.UserID + ".mask-test" + client := &mcpclient.Client{ + ClientDSL: mcpTypes.ClientDSL{ + ID: clientID, + Name: "mask-test", + Transport: mcpTypes.TransportHTTP, + URL: "https://example.com/mcp", + AuthorizationToken: "Bearer sk-test-token-12345678", + Timeout: "30s", + MetaInfo: gouTypes.MetaInfo{Label: "Mask Test"}, + }, + Enabled: true, + Status: "connected", + Source: mcpclient.ClientSourceDynamic, + Owner: mcpclient.ClientOwner{Type: "user", ID: ti.UserID}, + } + mcpclient.Global.Create(client) + defer mcpclient.Global.Delete(clientID) + + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + resp, _ := http.DefaultClient.Do(req) + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + resp.Body.Close() + + for _, s := range body["servers"].([]interface{}) { + sm := s.(map[string]interface{}) + if sm["id"] == clientID { + maskedToken, _ := sm["authorization_token"].(string) + assert.True(t, strings.Contains(maskedToken, "..."), "token should be masked, got: %s", maskedToken) + assert.NotEqual(t, "Bearer sk-test-token-12345678", maskedToken) + } + } +} + +func TestMCPDeleteServer(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + ti := obtainTokenInfo(t, serverURL) + + clientID := seedMCPServer(t, ti.UserID, "del-test", "https://example.com/mcp") + + req, _ := http.NewRequest("DELETE", serverURL+baseURL()+"/setting/mcp/servers/"+clientID, nil) + req.Header.Set("Authorization", "Bearer "+ti.AccessToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusNoContent, resp.StatusCode) + + listReq, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + listReq.Header.Set("Authorization", "Bearer "+ti.AccessToken) + listResp, _ := http.DefaultClient.Do(listReq) + var listBody map[string]interface{} + json.NewDecoder(listResp.Body).Decode(&listBody) + listResp.Body.Close() + + for _, s := range listBody["servers"].([]interface{}) { + sm := s.(map[string]interface{}) + assert.NotEqual(t, clientID, sm["id"], "deleted server should not appear in list") + } +} + +func TestMCPACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initMcpClientRegistry(t) + + readToken := obtainRestrictedToken(t, serverURL, "setting:mcp:read:all") + + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/mcp/servers", nil) + req.Header.Set("Authorization", "Bearer "+readToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + payload := map[string]interface{}{ + "name": "acl-test", "label": "ACL Test", "transport": "http", "url": "https://example.com/mcp", + } + raw, _ := json.Marshal(payload) + req2, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/mcp/servers", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+readToken) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode) + + req3, _ := http.NewRequest("DELETE", serverURL+baseURL()+"/setting/mcp/servers/some-id", nil) + req3.Header.Set("Authorization", "Bearer "+readToken) + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + defer resp3.Body.Close() + assert.Equal(t, http.StatusForbidden, resp3.StatusCode) +} diff --git a/openapi/tests/setting/sandbox_test.go b/openapi/tests/setting/sandbox_test.go new file mode 100644 index 00000000..6bba0e3a --- /dev/null +++ b/openapi/tests/setting/sandbox_test.go @@ -0,0 +1,315 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi/tests/testutils" + "github.com/yaoapp/yao/setting" + "github.com/yaoapp/yao/tai" + "github.com/yaoapp/yao/tai/registry" +) + +func initTaiForTest(t *testing.T) { + t.Helper() + if registry.Global() == nil { + tai.InitLocal(os.Stderr, "error", "") + } +} + +func TestSandboxGet(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + initTaiForTest(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/sandbox", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var data map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&data) + assert.NoError(t, err) + + nodes, ok := data["nodes"].([]interface{}) + assert.True(t, ok, "should have nodes array") + assert.NotNil(t, nodes) + + regConfig, ok := data["registry"].(map[string]interface{}) + assert.True(t, ok, "should have registry object") + assert.NotNil(t, regConfig) + + images, ok := data["images"].(map[string]interface{}) + assert.True(t, ok, "should have images object") + assert.NotNil(t, images) + + if len(nodes) > 0 { + node := nodes[0].(map[string]interface{}) + assert.NotEmpty(t, node["node_id"]) + assert.NotEmpty(t, node["os"]) + t.Logf("Node: %s (%s, %s)", node["node_id"], node["os"], node["arch"]) + } +} + +func TestSandboxGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/sandbox", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestSandboxRegistry(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + body := map[string]string{ + "registry_url": "https://registry.example.com", + "username": "testuser", + "password": "testpass123", + } + data, _ := json.Marshal(body) + + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/sandbox/registry", bytes.NewReader(data)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var regData map[string]interface{} + json.NewDecoder(resp.Body).Decode(®Data) + + assert.Equal(t, "https://registry.example.com", regData["registry_url"]) + assert.Equal(t, "testuser", regData["username"]) + pw, _ := regData["password"].(string) + assert.NotEqual(t, "testpass123", pw, "password should be masked") + assert.Contains(t, pw, "...", "password should contain mask") + + // Verify GET returns masked password + initTaiForTest(t) + req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/sandbox", nil) + req2.Header.Set("Authorization", "Bearer "+token) + + resp2, err := http.DefaultClient.Do(req2) + if !assert.NoError(t, err) || !assert.NotNil(t, resp2) { + return + } + defer resp2.Body.Close() + + var getResult map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&getResult) + regConfig, ok := getResult["registry"].(map[string]interface{}) + if assert.True(t, ok) { + assert.Equal(t, "https://registry.example.com", regConfig["registry_url"]) + pw2, _ := regConfig["password"].(string) + assert.NotEqual(t, "testpass123", pw2) + assert.Contains(t, pw2, "...") + } +} + +func TestSandboxCheckDocker(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initTaiForTest(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/sandbox/nodes/local/check-docker", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var data map[string]interface{} + json.NewDecoder(resp.Body).Decode(&data) + + if data["docker_version"] != nil { + ver := data["docker_version"].(string) + assert.NotEmpty(t, ver, "docker_version should be a non-empty string when Docker is running") + t.Logf("Docker version: %s", ver) + } else { + t.Log("Docker not available on local node (this is OK)") + } +} + +func TestSandboxCheckDockerNotFound(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initTaiForTest(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/sandbox/nodes/nonexistent-node-id/check-docker", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusNotFound, resp.StatusCode) +} + +func TestSandboxImagePull(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initTaiForTest(t) + initSettingRegistry(t) + + reg := registry.Global() + if reg == nil { + t.Skip("tai registry not initialized") + } + meta, ok := reg.Get("local") + if !ok || !meta.Capabilities.Docker { + t.Skip("local node has no Docker capability") + } + + token := obtainToken(t, serverURL) + + imageID := "YWxwaW5lOmxhdGVzdA" // base64url("alpine:latest") + req, err := http.NewRequest("POST", + serverURL+baseURL()+"/setting/sandbox/nodes/local/images/"+imageID+"/pull", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var data map[string]interface{} + json.NewDecoder(resp.Body).Decode(&data) + assert.Equal(t, "downloading", data["status"]) + t.Logf("Pull started for alpine:latest") +} + +func TestSandboxImageDelete(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initTaiForTest(t) + initSettingRegistry(t) + + reg := registry.Global() + if reg == nil { + t.Skip("tai registry not initialized") + } + meta, ok := reg.Get("local") + if !ok || !meta.Capabilities.Docker { + t.Skip("local node has no Docker capability") + } + + token := obtainToken(t, serverURL) + + imageID := "bm9uZXhpc3RlbnQ6bGF0ZXN0" // base64url("nonexistent:latest") + req, err := http.NewRequest("DELETE", + serverURL+baseURL()+"/setting/sandbox/nodes/local/images/"+imageID, nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "deleting non-existent image should return 400") +} + +func TestSandboxRegistryKeepPassword(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + body := map[string]string{ + "registry_url": "https://registry.example.com", + "username": "user1", + "password": "secret123", + } + data, _ := json.Marshal(body) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/sandbox/registry", bytes.NewReader(data)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) { + return + } + resp.Body.Close() + + body2 := map[string]string{ + "registry_url": "https://registry2.example.com", + "username": "user2", + "password": "", + } + data2, _ := json.Marshal(body2) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/sandbox/registry", bytes.NewReader(data2)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + if !assert.NoError(t, err) || !assert.NotNil(t, resp2) { + return + } + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var regData map[string]interface{} + json.NewDecoder(resp2.Body).Decode(®Data) + + assert.Equal(t, "https://registry2.example.com", regData["registry_url"]) + assert.Equal(t, "user2", regData["username"]) + pw, _ := regData["password"].(string) + assert.NotEmpty(t, pw, "password should still be present from previous save") + assert.Contains(t, pw, "...") +} + +func TestSandboxRegistryRequiresAuth(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + body := map[string]string{"registry_url": "https://example.com"} + data, _ := json.Marshal(body) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/sandbox/registry", bytes.NewReader(data)) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +var _ = setting.Global diff --git a/openapi/tests/setting/search_test.go b/openapi/tests/setting/search_test.go new file mode 100644 index 00000000..28100de5 --- /dev/null +++ b/openapi/tests/setting/search_test.go @@ -0,0 +1,339 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// --------------------------------------------------------------------------- +// Functional tests (system:root token) +// --------------------------------------------------------------------------- + +func TestSearchGet(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + assert.Contains(t, body, "presets") + assert.Contains(t, body, "providers") + assert.Contains(t, body, "tool_assignment") + + presets, ok := body["presets"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 4, len(presets), "should have 4 presets: cloud, tavily, serper, brightdata") + + providers, ok := body["providers"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 4, len(providers), "should have 4 provider configs") + + // Cloud provider should be first + first, _ := providers[0].(map[string]interface{}) + assert.Equal(t, "cloud", first["preset_key"]) +} + +func TestSearchGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestSearchProviderUpdate(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{ + "api_key": "tvly-test-key-12345", + }, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, "tavily", body["preset_key"]) + + // api_key should be masked in response + fv, _ := body["field_values"].(map[string]interface{}) + maskedKey, _ := fv["api_key"].(string) + assert.True(t, strings.Contains(maskedKey, "..."), "api_key should be masked") + + // GET should also return masked key + req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req2.Header.Set("Authorization", "Bearer "+token) + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + + var getData map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&getData) + providers, _ := getData["providers"].([]interface{}) + for _, p := range providers { + pm, _ := p.(map[string]interface{}) + if pm["preset_key"] == "tavily" { + tfv, _ := pm["field_values"].(map[string]interface{}) + assert.True(t, strings.Contains(tfv["api_key"].(string), "..."), "GET should return masked key") + } + } +} + +func TestSearchProviderUpdateCloud(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{}, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "cloud provider should be rejected") +} + +func TestSearchProviderToggle(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save tavily first + savePayload := map[string]interface{}{ + "field_values": map[string]string{"api_key": "tvly-toggle-key"}, + } + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + + // Enable tavily + enablePayload := map[string]interface{}{"enabled": true} + raw, _ = json.Marshal(enablePayload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + assert.Equal(t, true, body["enabled"]) + + // Assign tavily to web_search + assignPayload := map[string]interface{}{"web_search": "tavily"} + raw, _ = json.Marshal(assignPayload) + req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req3.Header.Set("Authorization", "Bearer "+token) + req3.Header.Set("Content-Type", "application/json") + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + resp3.Body.Close() + + // Disable tavily -- should clear tool_assignment + disablePayload := map[string]interface{}{"enabled": false} + raw, _ = json.Marshal(disablePayload) + req4, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req4.Header.Set("Authorization", "Bearer "+token) + req4.Header.Set("Content-Type", "application/json") + resp4, err := http.DefaultClient.Do(req4) + assert.NoError(t, err) + resp4.Body.Close() + + // Verify tool_assignment cleared + req5, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req5.Header.Set("Authorization", "Bearer "+token) + resp5, err := http.DefaultClient.Do(req5) + assert.NoError(t, err) + defer resp5.Body.Close() + + var getData map[string]interface{} + json.NewDecoder(resp5.Body).Decode(&getData) + ta, _ := getData["tool_assignment"].(map[string]interface{}) + assert.Nil(t, ta["web_search"], "web_search should be cleared after disabling tavily") +} + +func TestSearchToolAssignment(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save and enable tavily + savePayload := map[string]interface{}{"field_values": map[string]string{"api_key": "tvly-assign-key"}} + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, _ := http.DefaultClient.Do(req) + resp.Body.Close() + + enablePayload := map[string]interface{}{"enabled": true} + raw, _ = json.Marshal(enablePayload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, _ := http.DefaultClient.Do(req2) + resp2.Body.Close() + + // Assign tavily to web_search + assignPayload := map[string]interface{}{"web_search": "tavily"} + raw, _ = json.Marshal(assignPayload) + req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req3.Header.Set("Authorization", "Bearer "+token) + req3.Header.Set("Content-Type", "application/json") + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + defer resp3.Body.Close() + assert.Equal(t, http.StatusOK, resp3.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp3.Body).Decode(&body) + assert.Equal(t, "tavily", body["web_search"]) +} + +func TestSearchToolAssignmentDisabledProvider(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Try to assign a provider that isn't enabled + assignPayload := map[string]interface{}{"web_search": "serper"} + raw, _ := json.Marshal(assignPayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "should reject assignment to disabled provider") +} + +func TestSearchProviderTest(t *testing.T) { + apiKey := os.Getenv("TAVILY_API_KEY") + if apiKey == "" { + t.Skip("TAVILY_API_KEY not set, skipping search provider test") + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{"api_key": apiKey}, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/tavily/test", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, true, body["success"]) +} + +func TestSearchProviderTestCloud(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/cloud/test", nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +// --------------------------------------------------------------------------- +// ACL permission tests +// --------------------------------------------------------------------------- + +func TestSearchACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + + readToken := obtainRestrictedToken(t, serverURL, "setting:search:read:all") + + // GET should work + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req.Header.Set("Authorization", "Bearer "+readToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET") + + // PUT should be denied + payload := map[string]interface{}{ + "field_values": map[string]string{"api_key": "tvly-acl-test"}, + } + raw, _ := json.Marshal(payload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+readToken) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny PUT") +} diff --git a/openapi/tests/setting/smtp_test.go b/openapi/tests/setting/smtp_test.go new file mode 100644 index 00000000..a9764354 --- /dev/null +++ b/openapi/tests/setting/smtp_test.go @@ -0,0 +1,431 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// --------------------------------------------------------------------------- +// Functional tests (system:root token) +// --------------------------------------------------------------------------- + +func TestSmtpGet(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/smtp", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + assert.Contains(t, body, "presets") + assert.Contains(t, body, "config") + + presets, ok := body["presets"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 6, len(presets), "should have 6 en-us presets: gmail, yahoo, sendgrid, mailgun, ses, custom") + + config, ok := body["config"].(map[string]interface{}) + assert.True(t, ok) + assert.Equal(t, false, config["enabled"]) + assert.Equal(t, "unconfigured", config["status"]) + assert.Equal(t, "gmail", config["preset_key"], "default preset for en-us should be gmail") + assert.Equal(t, "", config["password"], "password should be empty when unconfigured") +} + +func TestSmtpGetZhCN(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/smtp?locale=zh-cn", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + + presets, ok := body["presets"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 6, len(presets), "should have 6 zh-cn presets: tencent, feishu, aliyun, qq, netease163, custom") + + first, _ := presets[0].(map[string]interface{}) + assert.Equal(t, "tencent", first["key"]) + + config, _ := body["config"].(map[string]interface{}) + assert.Equal(t, "tencent", config["preset_key"], "default preset for zh-cn should be tencent") +} + +func TestSmtpGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/smtp", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestSmtpUpdate(t *testing.T) { + host := os.Getenv("RELIABLE_SMTP_HOST") + port := os.Getenv("RELIABLE_SMTP_PORT") + user := os.Getenv("RELIABLE_SMTP_USERNAME") + pass := os.Getenv("RELIABLE_SMTP_PASSWORD") + if host == "" || user == "" || pass == "" { + host = os.Getenv("SMTP_HOST") + port = os.Getenv("SMTP_PORT") + user = os.Getenv("SMTP_USERNAME") + pass = os.Getenv("SMTP_PASSWORD") + } + if host == "" || user == "" || pass == "" { + t.Skip("SMTP credentials not set, skipping") + } + portNum := 465 + if port != "" { + if p, err := strconv.Atoi(port); err == nil { + portNum = p + } + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "preset_key": "custom", + "host": host, + "port": portNum, + "encryption": "ssl", + "username": user, + "password": pass, + "from_name": "Test Sender", + "from_email": user, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, host, body["host"]) + assert.Equal(t, user, body["username"]) + maskedPwd, _ := body["password"].(string) + assert.True(t, strings.Contains(maskedPwd, "..."), "password should be masked: got %s", maskedPwd) + + // GET should also return masked password + req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/smtp", nil) + req2.Header.Set("Authorization", "Bearer "+token) + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + + var getData map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&getData) + config, _ := getData["config"].(map[string]interface{}) + getMasked, _ := config["password"].(string) + assert.True(t, strings.Contains(getMasked, "..."), "GET should return masked password") +} + +func TestSmtpUpdateValidationFailure(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "preset_key": "gmail", + "host": "smtp.gmail.com", + "port": 465, + "encryption": "ssl", + "username": "fake@gmail.com", + "password": "wrong-password", + "from_name": "Test", + "from_email": "fake@gmail.com", + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "should reject invalid SMTP credentials") +} + +func TestSmtpUpdateKeepPassword(t *testing.T) { + host := os.Getenv("RELIABLE_SMTP_HOST") + port := os.Getenv("RELIABLE_SMTP_PORT") + user := os.Getenv("RELIABLE_SMTP_USERNAME") + pass := os.Getenv("RELIABLE_SMTP_PASSWORD") + if host == "" || user == "" || pass == "" { + host = os.Getenv("SMTP_HOST") + port = os.Getenv("SMTP_PORT") + user = os.Getenv("SMTP_USERNAME") + pass = os.Getenv("SMTP_PASSWORD") + } + if host == "" || user == "" || pass == "" { + t.Skip("SMTP credentials not set, skipping") + } + portNum := 465 + if port != "" { + if p, err := strconv.Atoi(port); err == nil { + portNum = p + } + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // First save with password + payload1 := map[string]interface{}{ + "preset_key": "custom", + "host": host, + "port": portNum, + "encryption": "ssl", + "username": user, + "password": pass, + "from_name": "Test", + "from_email": user, + } + raw, _ := json.Marshal(payload1) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + + // Update without password — should keep original and re-validate with existing password + payload2 := map[string]interface{}{ + "preset_key": "custom", + "host": host, + "port": portNum, + "encryption": "ssl", + "username": user, + "password": "", + "from_name": "Updated", + "from_email": user, + } + raw, _ = json.Marshal(payload2) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + assert.Equal(t, "Updated", body["from_name"]) + keepMasked, _ := body["password"].(string) + assert.True(t, strings.Contains(keepMasked, "..."), "password should be masked (kept original)") +} + +func TestSmtpToggle(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save config first + savePayload := map[string]interface{}{ + "preset_key": "gmail", + "host": "smtp.gmail.com", + "port": 465, + "encryption": "ssl", + "username": "test@gmail.com", + "password": "test-pass", + "from_name": "Test", + "from_email": "test@gmail.com", + } + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + + // Enable + enablePayload := map[string]interface{}{"enabled": true} + raw, _ = json.Marshal(enablePayload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp/toggle", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + assert.Equal(t, true, body["enabled"]) + + // Disable + disablePayload := map[string]interface{}{"enabled": false} + raw, _ = json.Marshal(disablePayload) + req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp/toggle", bytes.NewReader(raw)) + req3.Header.Set("Authorization", "Bearer "+token) + req3.Header.Set("Content-Type", "application/json") + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + defer resp3.Body.Close() + assert.Equal(t, http.StatusOK, resp3.StatusCode) + + var body2 map[string]interface{} + json.NewDecoder(resp3.Body).Decode(&body2) + assert.Equal(t, false, body2["enabled"]) + assert.Equal(t, "unconfigured", body2["status"]) +} + +func TestSmtpTest(t *testing.T) { + host := os.Getenv("RELIABLE_SMTP_HOST") + port := os.Getenv("RELIABLE_SMTP_PORT") + user := os.Getenv("RELIABLE_SMTP_USERNAME") + pass := os.Getenv("RELIABLE_SMTP_PASSWORD") + if host == "" || user == "" || pass == "" { + host = os.Getenv("SMTP_HOST") + port = os.Getenv("SMTP_PORT") + user = os.Getenv("SMTP_USERNAME") + pass = os.Getenv("SMTP_PASSWORD") + } + if host == "" || user == "" || pass == "" { + t.Skip("RELIABLE_SMTP_* or SMTP_* env not set, skipping SMTP test") + } + + toEmail := os.Getenv("SMTP_TEST_TO") + if toEmail == "" { + toEmail = user + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + portNum := 465 + if port != "" { + if p, err := strconv.Atoi(port); err == nil { + portNum = p + } + } + + savePayload := map[string]interface{}{ + "preset_key": "custom", + "host": host, + "port": portNum, + "encryption": "ssl", + "username": user, + "password": pass, + "from_name": "Yao SMTP Test", + "from_email": user, + } + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + + testPayload := map[string]interface{}{"to_email": toEmail} + raw, _ = json.Marshal(testPayload) + req2, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/smtp/test", bytes.NewReader(raw)) + assert.NoError(t, err) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + t.Logf("SMTP test result: %+v", body) + assert.Equal(t, true, body["success"]) +} + +// --------------------------------------------------------------------------- +// ACL permission tests +// --------------------------------------------------------------------------- + +func TestSmtpACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + + readToken := obtainRestrictedToken(t, serverURL, "setting:smtp:read:all") + + // GET should work + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/smtp", nil) + req.Header.Set("Authorization", "Bearer "+readToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET") + + // PUT should be denied + payload := map[string]interface{}{ + "preset_key": "gmail", + "host": "smtp.gmail.com", + "port": 465, + "encryption": "ssl", + "username": "test@gmail.com", + "password": "test-pass", + "from_name": "Test", + "from_email": "test@gmail.com", + } + raw, _ := json.Marshal(payload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/smtp", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+readToken) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny PUT") +} diff --git a/openapi/tests/setting/system_test.go b/openapi/tests/setting/system_test.go new file mode 100644 index 00000000..7779e251 --- /dev/null +++ b/openapi/tests/setting/system_test.go @@ -0,0 +1,108 @@ +package setting_test + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +func baseURL() string { + if openapi.Server != nil && openapi.Server.Config != nil { + return openapi.Server.Config.BaseURL + } + return "" +} + +// TestSystemInfo verifies GET /setting/system returns the expected structure. +func TestSystemInfo(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.RegisterTestClient(t, "Setting System Test", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + token := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/system", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + // Top-level keys + assert.Contains(t, body, "app") + assert.Contains(t, body, "deployment") + assert.Contains(t, body, "server") + assert.Contains(t, body, "client") + assert.Contains(t, body, "environment") + assert.Contains(t, body, "technical") + + // app sub-fields + app, ok := body["app"].(map[string]interface{}) + assert.True(t, ok) + assert.NotEmpty(t, app["name"]) + assert.NotEmpty(t, app["version"]) + + // server sub-fields + server, ok := body["server"].(map[string]interface{}) + assert.True(t, ok) + assert.NotEmpty(t, server["version"]) + + // technical sub-fields + tech, ok := body["technical"].(map[string]interface{}) + assert.True(t, ok) + assert.NotEmpty(t, tech["listen"]) + assert.NotEmpty(t, tech["db_driver"]) + assert.NotEmpty(t, tech["session_store"]) +} + +// TestSystemInfoUnauthenticated verifies 401 when no token is provided. +func TestSystemInfoUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/system", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +// TestSystemCheckUpdate verifies POST /setting/system/check-update returns has_update. +func TestSystemCheckUpdate(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.RegisterTestClient(t, "Setting CheckUpdate Test", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + token := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/system/check-update", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token.AccessToken) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + _, exists := body["has_update"] + assert.True(t, exists, "response must contain has_update field") +} diff --git a/setting/doc.go b/setting/doc.go new file mode 100644 index 00000000..b84825de --- /dev/null +++ b/setting/doc.go @@ -0,0 +1,12 @@ +package setting + +import ( + _ "embed" + + "github.com/yaoapp/gou/doc" +) + +//go:embed doc.yml +var docYAML []byte + +func init() { doc.LoadYAML(docYAML) } diff --git a/setting/doc.yml b/setting/doc.yml new file mode 100644 index 00000000..bbad84b4 --- /dev/null +++ b/setting/doc.yml @@ -0,0 +1,163 @@ +group: setting +type: process +desc: | + Generic user personalization settings store with three-level scope hierarchy. + Stores arbitrary JSON data organized by namespace and scope, with cascading + merge support (system ← team ← user, later scope wins). + + Process names follow the pattern "setting.". + + Scoping model: + Three levels, from lowest to highest priority: + 1. system — Global defaults, shared by all users. + 2. team — Team-level overrides, shared by team members. + 3. user — Individual user preferences, highest priority. + + ScopeID structure (used as argument for get, set, delete, listnamespaces): + - scope (string, required): Scope level. Values: "system", "team", "user". + - team_id (string): Required when scope is "team". + - user_id (string): Required when scope is "user". + Examples: + System scope: {"scope":"system"} + Team scope: {"scope":"team","team_id":"99"} + User scope: {"scope":"user","user_id":"42"} + + Entry structure (returned by set): + - namespace (string): Namespace name (e.g. "preferences", "privacy", "models"). + - scope (ScopeID): The scope this entry belongs to. + - data (object): Arbitrary key-value data stored for this namespace. + - updated_at (string): ISO 8601 timestamp of last update. + + Namespace convention: + Namespaces are free-form strings chosen by the consuming module. + Typical examples: "preferences", "privacy", "models", "notifications". + Each namespace stores one JSON object (map of string → any). + The registry does not enforce any schema — the consuming module defines + the expected structure. + + Merge behavior (getmerged): + Shallow merge across three scopes: system ← team ← user. + For each top-level key, the highest-priority scope's value wins. + Example: + system: {"theme":"light","lang":"en","font_size":14} + team: {"lang":"zh-CN"} + user: {"theme":"dark"} + merged: {"theme":"dark","lang":"zh-CN","font_size":14} + If a scope has no data for the namespace, it is skipped. + Returns 404 only when no data exists at any scope. + +entries: + - name: get + desc: | + Get a namespace entry for a specific scope. Returns the raw data object + without any merging. Throws 404 if the namespace does not exist at + the given scope. + args: + - name: scope + type: object + required: true + desc: | + ScopeID object identifying the scope. + Examples: + {"scope":"system"} + {"scope":"team","team_id":"99"} + {"scope":"user","user_id":"42"} + - name: namespace + type: string + required: true + desc: 'Namespace name (e.g. "preferences", "privacy", "models").' + return: + type: object + desc: | + The namespace data as a key-value map (not wrapped in Entry). + Example: {"theme":"dark","lang":"zh-CN","font_size":14} + + - name: getmerged + desc: | + Get a namespace with three-level cascade merge: system ← team ← user. + Reads data from all three scopes and shallow-merges them, with higher-priority + scopes overriding lower ones. Pass empty string for userID or teamID to + skip that scope. Throws 404 if no data exists at any scope. + args: + - name: userID + type: string + required: true + desc: 'User ID. Pass "" (empty string) to skip user scope.' + - name: teamID + type: string + required: true + desc: 'Team ID. Pass "" (empty string) to skip team scope.' + - name: namespace + type: string + required: true + desc: Namespace name. + return: + type: object + desc: | + Shallow-merged data from all available scopes. + Example with system={"a":"sys","b":"sys"}, team={"b":"team"}, user={"a":"user"}: + Result: {"a":"user","b":"team"} + + - name: set + desc: | + Set (create or overwrite) a namespace entry for a given scope. + Persists to __yao.store and updates __yao.cache. + Overwrites any existing data for this scope + namespace combination. + args: + - name: scope + type: object + required: true + desc: 'ScopeID object. Example: {"scope":"user","user_id":"42"}' + - name: namespace + type: string + required: true + desc: Namespace name. + - name: data + type: object + required: true + desc: | + Key-value data to store. Arbitrary JSON object. + Example: {"theme":"dark","lang":"zh-CN","font_size":14} + return: + type: object + desc: | + Entry object confirming the write. Fields: + - namespace (string) + - scope (ScopeID) + - data (object): The stored data. + - updated_at (string): ISO 8601 timestamp. + Example: {"namespace":"preferences","scope":{"scope":"user","user_id":"42"}, + "data":{"theme":"dark","lang":"zh-CN"},"updated_at":"2025-01-15T10:30:00Z"} + + - name: delete + desc: | + Delete a namespace entry from a scope. Removes from __yao.store + and __yao.cache. Throws 404 if the namespace does not exist at the scope. + args: + - name: scope + type: object + required: true + desc: 'ScopeID object. Example: {"scope":"system"}' + - name: namespace + type: string + required: true + desc: Namespace name to delete. + return: + type: "null" + desc: Returns null on success. + + - name: listnamespaces + desc: | + List all namespace names stored under a scope. Returns the namespace + strings only, not the data. Use "get" to retrieve data for each namespace. + args: + - name: scope + type: object + required: true + desc: 'ScopeID object. Example: {"scope":"team","team_id":"99"}' + return: + type: array + desc: | + Array of namespace name strings. + Example: ["preferences","privacy","models"] + Returns empty array if no namespaces exist for the scope. diff --git a/setting/process.go b/setting/process.go new file mode 100644 index 00000000..289c1353 --- /dev/null +++ b/setting/process.go @@ -0,0 +1,130 @@ +package setting + +import ( + "encoding/json" + + "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/exception" +) + +func init() { + process.RegisterGroup("setting", map[string]process.Handler{ + "get": ProcessGet, + "getmerged": ProcessGetMerged, + "set": ProcessSet, + "delete": ProcessDelete, + "listnamespaces": ProcessListNamespaces, + }) +} + +func requireGlobal() { + if Global == nil { + exception.New("Setting Registry not initialized", 500).Throw() + } +} + +func parseScopeID(arg interface{}) ScopeID { + raw, err := json.Marshal(arg) + if err != nil { + exception.New("invalid scope: "+err.Error(), 400).Throw() + } + var scope ScopeID + if err := json.Unmarshal(raw, &scope); err != nil { + exception.New("invalid scope: "+err.Error(), 400).Throw() + } + return scope +} + +// ProcessGet reads a namespace entry for a given scope. +// Args[0] map: ScopeID {scope, team_id?, user_id?} +// Args[1] string: namespace +func ProcessGet(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(2) + + scope := parseScopeID(p.Args[0]) + ns := p.ArgsString(1) + + data, err := Global.Get(scope, ns) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return data +} + +// ProcessGetMerged reads a namespace with three-level cascade merge. +// Args[0] string: userID +// Args[1] string: teamID +// Args[2] string: namespace +func ProcessGetMerged(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(3) + + userID := p.ArgsString(0) + teamID := p.ArgsString(1) + ns := p.ArgsString(2) + + data, err := Global.GetMerged(userID, teamID, ns) + if err != nil { + exception.New(err.Error(), 404).Throw() + } + return data +} + +// ProcessSet writes a namespace entry for a given scope. +// Args[0] map: ScopeID +// Args[1] string: namespace +// Args[2] map: data +func ProcessSet(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(3) + + scope := parseScopeID(p.Args[0]) + ns := p.ArgsString(1) + + raw, err := json.Marshal(p.Args[2]) + if err != nil { + exception.New("invalid data: "+err.Error(), 400).Throw() + } + var data map[string]interface{} + if err := json.Unmarshal(raw, &data); err != nil { + exception.New("invalid data: "+err.Error(), 400).Throw() + } + + entry, err := Global.Set(scope, ns, data) + if err != nil { + exception.New(err.Error(), 400).Throw() + } + return entry +} + +// ProcessDelete removes a namespace entry from a given scope. +// Args[0] map: ScopeID +// Args[1] string: namespace +func ProcessDelete(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(2) + + scope := parseScopeID(p.Args[0]) + ns := p.ArgsString(1) + + if err := Global.Delete(scope, ns); err != nil { + exception.New(err.Error(), 404).Throw() + } + return nil +} + +// ProcessListNamespaces returns all namespace names under a scope. +// Args[0] map: ScopeID +func ProcessListNamespaces(p *process.Process) interface{} { + requireGlobal() + p.ValidateArgNums(1) + + scope := parseScopeID(p.Args[0]) + + ns, err := Global.ListNamespaces(scope) + if err != nil { + exception.New(err.Error(), 500).Throw() + } + return ns +} diff --git a/setting/process_test.go b/setting/process_test.go new file mode 100644 index 00000000..cf14914a --- /dev/null +++ b/setting/process_test.go @@ -0,0 +1,109 @@ +package setting_test + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/process" +) + +var sysScope = map[string]interface{}{"scope": "system"} +var teamScopeP = map[string]interface{}{"scope": "team", "team_id": "99"} +var userScopeP = map[string]interface{}{"scope": "user", "user_id": "42"} + +func TestProcessSet(t *testing.T) { + setupRegistry(t) + + p := process.New("setting.set", sysScope, "prefs", map[string]interface{}{ + "theme": "dark", "lang": "zh-CN", + }) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + + m := toMapR(t, result) + assert.Equal(t, "prefs", m["namespace"]) + assert.NotEmpty(t, m["updated_at"]) +} + +func TestProcessGet(t *testing.T) { + setupRegistry(t) + + process.New("setting.set", sysScope, "gettest", map[string]interface{}{ + "color": "blue", + }).Exec() + + p := process.New("setting.get", sysScope, "gettest") + result, err := p.Exec() + require.NoError(t, err) + + m := toMapR(t, result) + assert.Equal(t, "blue", m["color"]) +} + +func TestProcessGetMerged(t *testing.T) { + setupRegistry(t) + + process.New("setting.set", sysScope, "merged", map[string]interface{}{ + "a": "sys", "b": "sys", + }).Exec() + process.New("setting.set", teamScopeP, "merged", map[string]interface{}{ + "b": "team", + }).Exec() + process.New("setting.set", userScopeP, "merged", map[string]interface{}{ + "a": "user", + }).Exec() + + p := process.New("setting.getmerged", "42", "99", "merged") + result, err := p.Exec() + require.NoError(t, err) + + m := toMapR(t, result) + assert.Equal(t, "user", m["a"]) + assert.Equal(t, "team", m["b"]) +} + +func TestProcessDelete(t *testing.T) { + setupRegistry(t) + + process.New("setting.set", sysScope, "deltest", map[string]interface{}{ + "x": "y", + }).Exec() + + p := process.New("setting.delete", sysScope, "deltest") + _, err := p.Exec() + require.NoError(t, err) + + pGet := process.New("setting.get", sysScope, "deltest") + _, err = pGet.Exec() + assert.Error(t, err) +} + +func TestProcessListNamespaces(t *testing.T) { + setupRegistry(t) + + process.New("setting.set", sysScope, "ns-a", map[string]interface{}{"v": 1}).Exec() + process.New("setting.set", sysScope, "ns-b", map[string]interface{}{"v": 2}).Exec() + + p := process.New("setting.listnamespaces", sysScope) + result, err := p.Exec() + require.NoError(t, err) + require.NotNil(t, result) + t.Logf("namespaces: %v", result) +} + +// --- helpers --- + +func toMapR(t *testing.T, v interface{}) map[string]interface{} { + t.Helper() + if m, ok := v.(map[string]interface{}); ok { + return m + } + raw, err := json.Marshal(v) + require.NoError(t, err) + var m map[string]interface{} + require.NoError(t, json.Unmarshal(raw, &m)) + return m +} diff --git a/setting/registry.go b/setting/registry.go new file mode 100644 index 00000000..52c73c67 --- /dev/null +++ b/setting/registry.go @@ -0,0 +1,183 @@ +package setting + +import ( + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/yaoapp/gou/store" +) + +// Global is the singleton Setting Registry. +var Global *Registry + +// Registry manages namespaced settings with three-level scope cascade. +type Registry struct { + store store.Store + cache store.Store + mu sync.RWMutex +} + +// Init initializes the global Registry. +// Must be called after store.Load (so __yao.store and __yao.cache are available). +func Init() error { + s, err := store.Get("__yao.store") + if err != nil { + return fmt.Errorf("setting.Init: %w", err) + } + c, _ := store.Get("__yao.cache") + Global = &Registry{store: s, cache: c} + return nil +} + +// Get reads the raw data for a single scope+namespace. +// If one or more dest pointers are provided, the data is also unmarshalled +// into dest[0] (like json.Unmarshal). +func (r *Registry) Get(scope ScopeID, ns string, dest ...interface{}) (map[string]interface{}, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + data, err := storeGet(r.store, r.cache, scope, ns) + if err != nil { + return nil, err + } + + if len(dest) > 0 && dest[0] != nil { + if err := bindDest(data, dest[0]); err != nil { + return data, fmt.Errorf("setting bind: %w", err) + } + } + return data, nil +} + +// GetMerged reads a namespace across all three scopes and returns a shallow-merged +// result: system <- team <- user (later wins). +// If one or more dest pointers are provided, the merged data is also unmarshalled +// into dest[0]. +func (r *Registry) GetMerged(userID, teamID, ns string, dest ...interface{}) (map[string]interface{}, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + merged := make(map[string]interface{}) + + if sys, err := storeGet(r.store, r.cache, ScopeID{Scope: ScopeSystem}, ns); err == nil { + shallowMerge(merged, sys) + } + + if teamID != "" { + if team, err := storeGet(r.store, r.cache, ScopeID{Scope: ScopeTeam, TeamID: teamID}, ns); err == nil { + shallowMerge(merged, team) + } + } + + if userID != "" { + if user, err := storeGet(r.store, r.cache, ScopeID{Scope: ScopeUser, UserID: userID}, ns); err == nil { + shallowMerge(merged, user) + } + } + + if len(merged) == 0 { + return nil, fmt.Errorf("setting %s: no data found at any scope", ns) + } + + if len(dest) > 0 && dest[0] != nil { + if err := bindDest(merged, dest[0]); err != nil { + return merged, fmt.Errorf("setting bind: %w", err) + } + } + return merged, nil +} + +// Set writes (or overwrites) a namespace entry for the given scope. +func (r *Registry) Set(scope ScopeID, ns string, data map[string]interface{}) (*Entry, error) { + r.mu.Lock() + defer r.mu.Unlock() + + if ns == "" { + return nil, fmt.Errorf("namespace is required") + } + + if err := storeSet(r.store, r.cache, scope, ns, data); err != nil { + return nil, err + } + if err := indexAdd(r.store, r.cache, scope, ns); err != nil { + return nil, err + } + + return &Entry{ + Namespace: ns, + Scope: scope, + Data: data, + UpdatedAt: time.Now().UTC().Format(time.RFC3339), + }, nil +} + +// Delete removes a namespace entry from a given scope. +func (r *Registry) Delete(scope ScopeID, ns string) error { + r.mu.Lock() + defer r.mu.Unlock() + + sk := storeKey(scope, ns) + if !r.store.Has(sk) { + return fmt.Errorf("setting %s/%s not found", scopePrefix(scope), ns) + } + + if err := storeDel(r.store, r.cache, scope, ns); err != nil { + return err + } + return indexRemove(r.store, r.cache, scope, ns) +} + +// ListNamespaces returns all namespace names stored under the given scope. +func (r *Registry) ListNamespaces(scope ScopeID) ([]string, error) { + r.mu.RLock() + defer r.mu.RUnlock() + return indexGet(r.store, r.cache, scope) +} + +// Reload clears the cache and re-populates it from the persistent store. +func (r *Registry) Reload() error { + r.mu.Lock() + defer r.mu.Unlock() + + if r.cache != nil { + _ = r.cache.Del(keyPrefix + "*") + } + + for _, scope := range []ScopeID{ + {Scope: ScopeSystem}, + } { + keys, err := indexGet(r.store, nil, scope) + if err != nil { + continue + } + ik := indexKey(scope) + raw, ok := r.store.Get(ik) + if ok && r.cache != nil { + r.cache.Set(ik, raw, 0) + } + for _, ns := range keys { + if data, err := storeGet(r.store, nil, scope, ns); err == nil && r.cache != nil { + r.cache.Set(storeKey(scope, ns), data, 0) + } + } + } + return nil +} + +// shallowMerge copies all keys from src into dst (overwrites existing keys). +func shallowMerge(dst, src map[string]interface{}) { + for k, v := range src { + dst[k] = v + } +} + +// bindDest marshals data to JSON and then unmarshals into the dest pointer. +func bindDest(data map[string]interface{}, dest interface{}) error { + raw, err := json.Marshal(data) + if err != nil { + return err + } + return json.Unmarshal(raw, dest) +} diff --git a/setting/registry_test.go b/setting/registry_test.go new file mode 100644 index 00000000..2d93fca9 --- /dev/null +++ b/setting/registry_test.go @@ -0,0 +1,358 @@ +package setting_test + +import ( + "fmt" + "os" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/setting" + "github.com/yaoapp/yao/test" +) + +func TestMain(m *testing.M) { + test.Prepare(nil, config.Conf) + defer test.Clean() + os.Exit(m.Run()) +} + +func setupRegistry(t *testing.T) *setting.Registry { + t.Helper() + test.Prepare(t, config.Conf) + + err := setting.Init() + require.NoError(t, err) + + t.Cleanup(func() { + s, _ := store.Get("__yao.store") + if s != nil { + s.Del("setting:*") + } + c, _ := store.Get("__yao.cache") + if c != nil { + c.Del("setting:*") + } + test.Clean() + }) + + return setting.Global +} + +var systemScope = setting.ScopeID{Scope: setting.ScopeSystem} +var teamScope = setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "99"} +var userScope = setting.ScopeID{Scope: setting.ScopeUser, UserID: "42"} + +func TestSetAndGet(t *testing.T) { + r := setupRegistry(t) + + data := map[string]interface{}{ + "theme": "dark", + "language": "zh-CN", + "fontSize": float64(14), + } + entry, err := r.Set(systemScope, "preferences", data) + require.NoError(t, err) + assert.Equal(t, "preferences", entry.Namespace) + assert.Equal(t, systemScope, entry.Scope) + assert.NotEmpty(t, entry.UpdatedAt) + + got, err := r.Get(systemScope, "preferences") + require.NoError(t, err) + assert.Equal(t, "dark", got["theme"]) + assert.Equal(t, "zh-CN", got["language"]) + assert.Equal(t, float64(14), got["fontSize"]) +} + +func TestGetWithBind(t *testing.T) { + r := setupRegistry(t) + + data := map[string]interface{}{ + "default_chat": "gpt-4o", + "vision_model": "gpt-4o", + "embedding_enabled": true, + } + _, err := r.Set(systemScope, "models", data) + require.NoError(t, err) + + type ModelsConfig struct { + DefaultChat string `json:"default_chat"` + VisionModel string `json:"vision_model"` + EmbeddingEnabled bool `json:"embedding_enabled"` + } + + var cfg ModelsConfig + raw, err := r.Get(systemScope, "models", &cfg) + require.NoError(t, err) + + assert.Equal(t, "gpt-4o", raw["default_chat"]) + assert.Equal(t, "gpt-4o", cfg.DefaultChat) + assert.Equal(t, "gpt-4o", cfg.VisionModel) + assert.True(t, cfg.EmbeddingEnabled) +} + +func TestGetMergedWithBind(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "prefs", map[string]interface{}{ + "theme": "dark", "lang": "zh-CN", "font_size": float64(14), + }) + require.NoError(t, err) + + _, err = r.Set(teamScope, "prefs", map[string]interface{}{ + "lang": "en-US", + }) + require.NoError(t, err) + + _, err = r.Set(userScope, "prefs", map[string]interface{}{ + "theme": "light", + }) + require.NoError(t, err) + + type Prefs struct { + Theme string `json:"theme"` + Lang string `json:"lang"` + FontSize float64 `json:"font_size"` + } + + var p Prefs + _, err = r.GetMerged("42", "99", "prefs", &p) + require.NoError(t, err) + assert.Equal(t, "light", p.Theme) + assert.Equal(t, "en-US", p.Lang) + assert.Equal(t, float64(14), p.FontSize) +} + +func TestGetNotFound(t *testing.T) { + r := setupRegistry(t) + _, err := r.Get(systemScope, "nonexistent") + assert.Error(t, err) + assert.Contains(t, err.Error(), "not found") +} + +func TestGetMerged(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "theme", map[string]interface{}{ + "primary": "blue", "dark_mode": true, "font": "inter", + }) + require.NoError(t, err) + + _, err = r.Set( + setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "t1"}, + "theme", + map[string]interface{}{"dark_mode": false}, + ) + require.NoError(t, err) + + _, err = r.Set( + setting.ScopeID{Scope: setting.ScopeUser, UserID: "u1"}, + "theme", + map[string]interface{}{"primary": "red"}, + ) + require.NoError(t, err) + + merged, err := r.GetMerged("u1", "t1", "theme") + require.NoError(t, err) + assert.Equal(t, "red", merged["primary"]) + assert.Equal(t, false, merged["dark_mode"]) + assert.Equal(t, "inter", merged["font"]) +} + +func TestGetMergedPartial(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "partial", map[string]interface{}{"a": "1", "b": "2"}) + require.NoError(t, err) + + // Only system + user, no team data + _, err = r.Set(userScope, "partial", map[string]interface{}{"b": "override"}) + require.NoError(t, err) + + merged, err := r.GetMerged("42", "", "partial") + require.NoError(t, err) + assert.Equal(t, "1", merged["a"]) + assert.Equal(t, "override", merged["b"]) +} + +func TestGetMergedNoData(t *testing.T) { + r := setupRegistry(t) + _, err := r.GetMerged("42", "99", "nothing") + assert.Error(t, err) + assert.Contains(t, err.Error(), "no data found") +} + +func TestSetOverwrite(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "overwrite", map[string]interface{}{"a": "1"}) + require.NoError(t, err) + + _, err = r.Set(systemScope, "overwrite", map[string]interface{}{"a": "2", "b": "3"}) + require.NoError(t, err) + + got, err := r.Get(systemScope, "overwrite") + require.NoError(t, err) + assert.Equal(t, "2", got["a"]) + assert.Equal(t, "3", got["b"]) +} + +func TestSetEmptyNamespace(t *testing.T) { + r := setupRegistry(t) + _, err := r.Set(systemScope, "", map[string]interface{}{"a": "1"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "namespace is required") +} + +func TestDelete(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "to-delete", map[string]interface{}{"x": "y"}) + require.NoError(t, err) + + err = r.Delete(systemScope, "to-delete") + require.NoError(t, err) + + _, err = r.Get(systemScope, "to-delete") + assert.Error(t, err) +} + +func TestDeleteNotFound(t *testing.T) { + r := setupRegistry(t) + err := r.Delete(systemScope, "no-such-ns") + assert.Error(t, err) + assert.Contains(t, err.Error(), "not found") +} + +func TestListNamespaces(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "ns-a", map[string]interface{}{"v": 1}) + require.NoError(t, err) + _, err = r.Set(systemScope, "ns-b", map[string]interface{}{"v": 2}) + require.NoError(t, err) + _, err = r.Set(teamScope, "ns-c", map[string]interface{}{"v": 3}) + require.NoError(t, err) + + sysNS, err := r.ListNamespaces(systemScope) + require.NoError(t, err) + assert.Contains(t, sysNS, "ns-a") + assert.Contains(t, sysNS, "ns-b") + assert.NotContains(t, sysNS, "ns-c") + + teamNS, err := r.ListNamespaces(teamScope) + require.NoError(t, err) + assert.Contains(t, teamNS, "ns-c") +} + +func TestMultipleNamespaces(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(userScope, "alpha", map[string]interface{}{"color": "red"}) + require.NoError(t, err) + _, err = r.Set(userScope, "beta", map[string]interface{}{"color": "blue"}) + require.NoError(t, err) + + a, err := r.Get(userScope, "alpha") + require.NoError(t, err) + assert.Equal(t, "red", a["color"]) + + b, err := r.Get(userScope, "beta") + require.NoError(t, err) + assert.Equal(t, "blue", b["color"]) +} + +func TestScopeIsolation(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "shared", map[string]interface{}{"level": "system"}) + require.NoError(t, err) + _, err = r.Set(teamScope, "shared", map[string]interface{}{"level": "team"}) + require.NoError(t, err) + _, err = r.Set(userScope, "shared", map[string]interface{}{"level": "user"}) + require.NoError(t, err) + + sys, err := r.Get(systemScope, "shared") + require.NoError(t, err) + assert.Equal(t, "system", sys["level"]) + + team, err := r.Get(teamScope, "shared") + require.NoError(t, err) + assert.Equal(t, "team", team["level"]) + + user, err := r.Get(userScope, "shared") + require.NoError(t, err) + assert.Equal(t, "user", user["level"]) +} + +func TestReload(t *testing.T) { + r := setupRegistry(t) + + _, err := r.Set(systemScope, "reload-test", map[string]interface{}{"k": "v"}) + require.NoError(t, err) + + c, _ := store.Get("__yao.cache") + if c != nil { + c.Del("setting:*") + } + + err = r.Reload() + require.NoError(t, err) + + got, err := r.Get(systemScope, "reload-test") + require.NoError(t, err) + assert.Equal(t, "v", got["k"]) +} + +func TestConcurrency(t *testing.T) { + r := setupRegistry(t) + + var wg sync.WaitGroup + errCh := make(chan error, 30) + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + ns := fmt.Sprintf("conc-%d", idx) + _, err := r.Set(systemScope, ns, map[string]interface{}{"idx": idx}) + if err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + ns := fmt.Sprintf("conc-%d", idx) + _, err := r.Get(systemScope, ns) + if err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + ns := fmt.Sprintf("conc-%d", idx) + if err := r.Delete(systemScope, ns); err != nil { + errCh <- err + } + }(i) + } + wg.Wait() + close(errCh) + + for err := range errCh { + t.Errorf("concurrent operation error: %v", err) + } +} diff --git a/setting/store.go b/setting/store.go new file mode 100644 index 00000000..62df76f3 --- /dev/null +++ b/setting/store.go @@ -0,0 +1,178 @@ +package setting + +import ( + "encoding/json" + "fmt" + + "github.com/yaoapp/gou/store" +) + +const keyPrefix = "setting:" + +func scopePrefix(scope ScopeID) string { + switch scope.Scope { + case ScopeTeam: + return "t" + scope.TeamID + ":" + case ScopeUser: + return "u" + scope.UserID + ":" + default: + return "s:" + } +} + +func storeKey(scope ScopeID, ns string) string { + return keyPrefix + scopePrefix(scope) + ns +} + +func indexKey(scope ScopeID) string { + return keyPrefix + "idx:" + scopePrefix(scope) +} + +// storeGet reads a namespace entry from cache first, then persistent store. +func storeGet(s, c store.Store, scope ScopeID, ns string) (map[string]interface{}, error) { + sk := storeKey(scope, ns) + + if c != nil { + if val, ok := c.Get(sk); ok { + if m, ok := val.(map[string]interface{}); ok { + return m, nil + } + } + } + + val, ok := s.Get(sk) + if !ok { + return nil, fmt.Errorf("setting %s/%s not found", scopePrefix(scope), ns) + } + + m, err := toMap(val) + if err != nil { + return nil, fmt.Errorf("setting %s/%s: %w", scopePrefix(scope), ns, err) + } + + if c != nil { + c.Set(sk, m, 0) + } + return m, nil +} + +// storeSet writes a namespace entry to both persistent store and cache. +func storeSet(s, c store.Store, scope ScopeID, ns string, data map[string]interface{}) error { + sk := storeKey(scope, ns) + if err := s.Set(sk, data, 0); err != nil { + return err + } + if c != nil { + c.Set(sk, data, 0) + } + return nil +} + +// storeDel removes a namespace entry from both persistent store and cache. +func storeDel(s, c store.Store, scope ScopeID, ns string) error { + sk := storeKey(scope, ns) + if err := s.Del(sk); err != nil { + return err + } + if c != nil { + c.Del(sk) + } + return nil +} + +// indexGet returns all namespace names for a given scope. +func indexGet(s, c store.Store, scope ScopeID) ([]string, error) { + ik := indexKey(scope) + var raw interface{} + var ok bool + + if c != nil { + raw, ok = c.Get(ik) + } + if !ok { + raw, ok = s.Get(ik) + if !ok { + return nil, nil + } + if c != nil { + c.Set(ik, raw, 0) + } + } + + switch v := raw.(type) { + case []interface{}: + keys := make([]string, 0, len(v)) + for _, item := range v { + if str, ok := item.(string); ok { + keys = append(keys, str) + } + } + return keys, nil + case []string: + return v, nil + default: + return nil, fmt.Errorf("unexpected index type %T", raw) + } +} + +// indexSet writes the full namespace index. +func indexSet(s, c store.Store, scope ScopeID, keys []string) error { + ik := indexKey(scope) + iface := make([]interface{}, len(keys)) + for i, k := range keys { + iface[i] = k + } + if err := s.Set(ik, iface, 0); err != nil { + return err + } + if c != nil { + c.Set(ik, iface, 0) + } + return nil +} + +// indexAdd appends a namespace to the index if not already present. +func indexAdd(s, c store.Store, scope ScopeID, ns string) error { + keys, err := indexGet(s, c, scope) + if err != nil { + return err + } + for _, k := range keys { + if k == ns { + return nil + } + } + return indexSet(s, c, scope, append(keys, ns)) +} + +// indexRemove removes a namespace from the index. +func indexRemove(s, c store.Store, scope ScopeID, ns string) error { + keys, err := indexGet(s, c, scope) + if err != nil { + return err + } + filtered := make([]string, 0, len(keys)) + for _, k := range keys { + if k != ns { + filtered = append(filtered, k) + } + } + return indexSet(s, c, scope, filtered) +} + +// toMap normalizes a store value to map[string]interface{}. +// The xun store may return values that need re-serialization. +func toMap(val interface{}) (map[string]interface{}, error) { + if m, ok := val.(map[string]interface{}); ok { + return m, nil + } + raw, err := json.Marshal(val) + if err != nil { + return nil, err + } + var m map[string]interface{} + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil +} diff --git a/setting/types.go b/setting/types.go new file mode 100644 index 00000000..8bf82567 --- /dev/null +++ b/setting/types.go @@ -0,0 +1,28 @@ +package setting + +// Scope identifies the level at which a setting is stored. +type Scope string + +const ( + ScopeSystem Scope = "system" + ScopeTeam Scope = "team" + ScopeUser Scope = "user" +) + +// ScopeID fully identifies a scope instance. +// For ScopeSystem, TeamID and UserID are ignored. +// For ScopeTeam, TeamID is required. +// For ScopeUser, UserID is required. +type ScopeID struct { + Scope Scope `json:"scope"` + TeamID string `json:"team_id,omitempty"` + UserID string `json:"user_id,omitempty"` +} + +// Entry represents a single namespace's data within a scope. +type Entry struct { + Namespace string `json:"namespace"` + Scope ScopeID `json:"scope"` + Data map[string]interface{} `json:"data"` + UpdatedAt string `json:"updated_at"` +} diff --git a/tai/runtime/image_docker.go b/tai/runtime/image_docker.go index 4014bde1..3af35259 100644 --- a/tai/runtime/image_docker.go +++ b/tai/runtime/image_docker.go @@ -83,7 +83,6 @@ func (d *dockerImage) Pull(ctx context.Context, ref string, opts PullOptions) (< } pullOpts.RegistryAuth = encoded } - reader, err := d.cli.ImagePull(ctx, ref, pullOpts) if err != nil { return nil, fmt.Errorf("image pull %q: %w", ref, err)