Merge pull request #1529 from trheyi/main
feat: add Settings, LLM, Search, SMTP, MCP, Sandbox module management endpoints
This commit is contained in:
commit
d4553056ef
53 changed files with 11875 additions and 3 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -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
|
||||
|
|
|
|||
5
Makefile
5
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; \
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
12
llmprovider/doc.go
Normal file
12
llmprovider/doc.go
Normal file
|
|
@ -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) }
|
||||
256
llmprovider/doc.yml
Normal file
256
llmprovider/doc.yml
Normal file
|
|
@ -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.<handler>".
|
||||
|
||||
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.<key>" for system owner, "u<user_id>.<key>" for user owner,
|
||||
"t<team_id>.<key>" 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.
|
||||
42
llmprovider/presets.go
Normal file
42
llmprovider/presets.go
Normal file
|
|
@ -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
|
||||
}
|
||||
61
llmprovider/presets.yml
Normal file
61
llmprovider/presets.yml
Normal file
|
|
@ -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: []
|
||||
170
llmprovider/process.go
Normal file
170
llmprovider/process.go
Normal file
|
|
@ -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
|
||||
}
|
||||
164
llmprovider/process_test.go
Normal file
164
llmprovider/process_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
312
llmprovider/registry.go
Normal file
312
llmprovider/registry.go
Normal file
|
|
@ -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
|
||||
}
|
||||
645
llmprovider/registry_test.go
Normal file
645
llmprovider/registry_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
274
llmprovider/store.go
Normal file
274
llmprovider/store.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
199
llmprovider/sync.go
Normal file
199
llmprovider/sync.go
Normal file
|
|
@ -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
|
||||
}
|
||||
}
|
||||
85
llmprovider/types.go
Normal file
85
llmprovider/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
12
mcpclient/doc.go
Normal file
12
mcpclient/doc.go
Normal file
|
|
@ -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) }
|
||||
189
mcpclient/doc.yml
Normal file
189
mcpclient/doc.yml
Normal file
|
|
@ -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.<handler>".
|
||||
|
||||
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.<id>" for system, "u<user_id>.<id>" for user, "t<team_id>.<id>" 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.
|
||||
120
mcpclient/process.go
Normal file
120
mcpclient/process.go
Normal file
|
|
@ -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
|
||||
}
|
||||
120
mcpclient/process_test.go
Normal file
120
mcpclient/process_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
255
mcpclient/registry.go
Normal file
255
mcpclient/registry.go
Normal file
|
|
@ -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
|
||||
}
|
||||
483
mcpclient/registry_test.go
Normal file
483
mcpclient/registry_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
179
mcpclient/store.go
Normal file
179
mcpclient/store.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
136
mcpclient/sync.go
Normal file
136
mcpclient/sync.go
Normal file
|
|
@ -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"},
|
||||
}
|
||||
}
|
||||
50
mcpclient/types.go
Normal file
50
mcpclient/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
||||
}
|
||||
|
|
|
|||
409
openapi/setting/cloud.go
Normal file
409
openapi/setting/cloud.go
Normal file
|
|
@ -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
|
||||
}
|
||||
28
openapi/setting/cloud_presets.yml
Normal file
28
openapi/setting/cloud_presets.yml
Normal file
|
|
@ -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"
|
||||
668
openapi/setting/llm.go
Normal file
668
openapi/setting/llm.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
507
openapi/setting/mcp.go
Normal file
507
openapi/setting/mcp.go
Normal file
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
42
openapi/setting/promotions.yml
Normal file
42
openapi/setting/promotions.yml
Normal file
|
|
@ -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 →"
|
||||
796
openapi/setting/sandbox.go
Normal file
796
openapi/setting/sandbox.go
Normal file
|
|
@ -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})
|
||||
}
|
||||
622
openapi/setting/search.go
Normal file
622
openapi/setting/search.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
61
openapi/setting/search_presets.yml
Normal file
61
openapi/setting/search_presets.yml
Normal file
|
|
@ -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" }
|
||||
160
openapi/setting/setting.go
Normal file
160
openapi/setting/setting.go
Normal file
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
586
openapi/setting/smtp.go
Normal file
586
openapi/setting/smtp.go
Normal file
|
|
@ -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
|
||||
}
|
||||
}
|
||||
121
openapi/setting/smtp_presets.yml
Normal file
121
openapi/setting/smtp_presets.yml
Normal file
|
|
@ -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"
|
||||
255
openapi/setting/system.go
Normal file
255
openapi/setting/system.go
Normal file
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
221
openapi/setting/types.go
Normal file
221
openapi/setting/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
330
openapi/tests/setting/cloud_test.go
Normal file
330
openapi/tests/setting/cloud_test.go
Normal file
|
|
@ -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"])
|
||||
}
|
||||
453
openapi/tests/setting/llm_test.go
Normal file
453
openapi/tests/setting/llm_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
424
openapi/tests/setting/mcp_test.go
Normal file
424
openapi/tests/setting/mcp_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
315
openapi/tests/setting/sandbox_test.go
Normal file
315
openapi/tests/setting/sandbox_test.go
Normal file
|
|
@ -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
|
||||
339
openapi/tests/setting/search_test.go
Normal file
339
openapi/tests/setting/search_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
431
openapi/tests/setting/smtp_test.go
Normal file
431
openapi/tests/setting/smtp_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
108
openapi/tests/setting/system_test.go
Normal file
108
openapi/tests/setting/system_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
12
setting/doc.go
Normal file
12
setting/doc.go
Normal file
|
|
@ -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) }
|
||||
163
setting/doc.yml
Normal file
163
setting/doc.yml
Normal file
|
|
@ -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.<handler>".
|
||||
|
||||
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.
|
||||
130
setting/process.go
Normal file
130
setting/process.go
Normal file
|
|
@ -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
|
||||
}
|
||||
109
setting/process_test.go
Normal file
109
setting/process_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
183
setting/registry.go
Normal file
183
setting/registry.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
358
setting/registry_test.go
Normal file
358
setting/registry_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
178
setting/store.go
Normal file
178
setting/store.go
Normal file
|
|
@ -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
|
||||
}
|
||||
28
setting/types.go
Normal file
28
setting/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue