diff --git a/openapi/setting/search.go b/openapi/setting/search.go new file mode 100644 index 00000000..f7341f31 --- /dev/null +++ b/openapi/setting/search.go @@ -0,0 +1,622 @@ +package setting + +import ( + "bytes" + _ "embed" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/openapi/oauth/authorized" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/response" + "github.com/yaoapp/yao/setting" + "gopkg.in/yaml.v3" +) + +//go:embed search_presets.yml +var searchPresetsYML []byte + +var searchPresets []SearchProviderPreset + +func init() { + if err := yaml.Unmarshal(searchPresetsYML, &searchPresets); err != nil { + searchPresets = nil + } +} + +func searchFindPreset(key string) *SearchProviderPreset { + for i := range searchPresets { + if searchPresets[i].Key == key { + return &searchPresets[i] + } + } + return nil +} + +func searchScope(info *oauthTypes.AuthorizedInfo) setting.ScopeID { + if info.TeamID != "" { + return setting.ScopeID{Scope: setting.ScopeTeam, TeamID: info.TeamID} + } + return setting.ScopeID{Scope: setting.ScopeUser, UserID: info.UserID} +} + +func searchProviderNS(key string) string { + return "search.providers." + key +} + +const searchAssignmentNS = "search.tool_assignment" + +func searchPasswordFields(preset *SearchProviderPreset) map[string]bool { + m := make(map[string]bool) + for _, f := range preset.Fields { + if f.Type == "password" { + m[f.Key] = true + } + } + return m +} + +// --------------------------------------------------------------------------- +// GET /setting/search +// --------------------------------------------------------------------------- + +func handleSearchGet(c *gin.Context) { + info := authorized.GetInfo(c) + + providers := make([]SearchProviderConfig, 0, len(searchPresets)) + for _, preset := range searchPresets { + cfg := SearchProviderConfig{ + PresetKey: preset.Key, + Enabled: false, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + + if preset.IsCloud { + var cloudSaved map[string]interface{} + if setting.Global != nil { + cloudSaved, _ = setting.Global.GetMerged(info.UserID, info.TeamID, cloudNS) + } + if cloudSaved != nil { + if st, ok := cloudSaved["status"].(string); ok && st == "connected" { + cfg.Enabled = true + cfg.Status = "connected" + } + } + } else if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchProviderNS(preset.Key)) + if saved != nil { + if v, ok := saved["enabled"].(bool); ok { + cfg.Enabled = v + } + if v, ok := saved["status"].(string); ok && v != "" { + cfg.Status = v + } + pwFields := searchPasswordFields(&preset) + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + } + } + providers = append(providers, cfg) + } + + var assignment SearchToolAssignment + if setting.Global != nil { + saved, _ := setting.Global.GetMerged(info.UserID, info.TeamID, searchAssignmentNS) + if saved != nil { + if v, ok := saved["web_search"].(string); ok && v != "" { + assignment.WebSearch = &v + } + if v, ok := saved["web_scrape"].(string); ok && v != "" { + assignment.WebScrape = &v + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, SearchPageData{ + Presets: searchPresets, + Providers: providers, + ToolAssignment: assignment, + }) +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/providers/:key +// --------------------------------------------------------------------------- + +func handleSearchProviderUpdate(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + FieldValues map[string]string `json:"field_values"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + existing, _ := setting.Global.Get(scope, searchProviderNS(key)) + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + + validFields := make(map[string]bool) + for _, f := range preset.Fields { + validFields[f.Key] = true + } + + pwFields := searchPasswordFields(preset) + existingFV := map[string]interface{}{} + if fv, ok := m["field_values"].(map[string]interface{}); ok { + existingFV = fv + } + + newFV := make(map[string]interface{}) + for k, v := range existingFV { + newFV[k] = v + } + + for k, v := range body.FieldValues { + if !validFields[k] { + continue + } + if pwFields[k] { + if v == "" { + continue // keep existing + } + newFV[k] = cloudEncrypt(v) + } else { + newFV[k] = v + } + } + + m["field_values"] = newFV + if _, ok := m["enabled"]; !ok { + m["enabled"] = false + } + if _, ok := m["status"]; !ok { + m["status"] = "unconfigured" + } + + if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + cfg := SearchProviderConfig{ + PresetKey: key, + Enabled: false, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + if v, ok := m["enabled"].(bool); ok { + cfg.Enabled = v + } + if v, ok := m["status"].(string); ok && v != "" { + cfg.Status = v + } + if fv, ok := m["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/providers/:key/toggle +// --------------------------------------------------------------------------- + +func handleSearchProviderToggle(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider is managed by cloud service settings") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + Enabled bool `json:"enabled"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + existing, _ := setting.Global.Get(scope, searchProviderNS(key)) + m := make(map[string]interface{}) + for k, v := range existing { + m[k] = v + } + m["enabled"] = body.Enabled + + if _, err := setting.Global.Set(scope, searchProviderNS(key), m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + // When disabling, clear tool_assignment references + if !body.Enabled { + assignData, _ := setting.Global.Get(scope, searchAssignmentNS) + if assignData != nil { + changed := false + if v, ok := assignData["web_search"].(string); ok && v == key { + assignData["web_search"] = "" + changed = true + } + if v, ok := assignData["web_scrape"].(string); ok && v == key { + assignData["web_scrape"] = "" + changed = true + } + if changed { + setting.Global.Set(scope, searchAssignmentNS, assignData) + } + } + } + + cfg := SearchProviderConfig{ + PresetKey: key, + Enabled: body.Enabled, + FieldValues: map[string]string{}, + Status: "unconfigured", + } + if v, ok := m["status"].(string); ok && v != "" { + cfg.Status = v + } + pwFields := searchPasswordFields(preset) + if fv, ok := m["field_values"].(map[string]interface{}); ok { + for k, v := range fv { + s, _ := v.(string) + if pwFields[k] && s != "" { + cfg.FieldValues[k] = cloudMaskKey(cloudDecrypt(s)) + } else { + cfg.FieldValues[k] = s + } + } + } + + response.RespondWithSuccess(c, http.StatusOK, cfg) +} + +// --------------------------------------------------------------------------- +// POST /setting/search/providers/:key/test +// --------------------------------------------------------------------------- + +func handleSearchProviderTest(c *gin.Context) { + if !guardOwner(c) { + return + } + + key := c.Param("key") + if key == "cloud" { + respondError(c, http.StatusBadRequest, "cloud provider status is determined by cloud service configuration") + return + } + + preset := searchFindPreset(key) + if preset == nil { + respondError(c, http.StatusBadRequest, fmt.Sprintf("unknown provider: %s", key)) + return + } + + var body struct { + FieldValues map[string]string `json:"field_values"` + } + c.ShouldBindJSON(&body) + + info := authorized.GetInfo(c) + scope := searchScope(info) + + // Resolve API key: prefer body, fall back to saved + apiKey := "" + if body.FieldValues != nil { + apiKey = body.FieldValues["api_key"] + } + if apiKey == "" && setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved != nil { + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + if v, ok := fv["api_key"].(string); ok { + apiKey = cloudDecrypt(v) + } + } + } + } + + if apiKey == "" { + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: false, + Message: "API key is required", + }) + return + } + + start := time.Now() + var testErr error + + zone := "" + if body.FieldValues != nil { + zone = body.FieldValues["zone"] + } + if zone == "" && setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved != nil { + if fv, ok := saved["field_values"].(map[string]interface{}); ok { + if v, ok := fv["zone"].(string); ok { + zone = v + } + } + } + } + + switch key { + case "tavily": + testErr = searchTestTavily(apiKey) + case "serper": + testErr = searchTestSerper(apiKey) + case "brightdata": + testErr = searchTestBrightdata(apiKey, zone) + default: + respondError(c, http.StatusBadRequest, fmt.Sprintf("test not supported for provider: %s", key)) + return + } + + latency := time.Since(start).Milliseconds() + + if testErr != nil { + // Update status to disconnected + if setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved == nil { + saved = map[string]interface{}{} + } + saved["status"] = "disconnected" + setting.Global.Set(scope, searchProviderNS(key), saved) + } + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: false, + Message: testErr.Error(), + }) + return + } + + // Update status to connected + if setting.Global != nil { + saved, _ := setting.Global.Get(scope, searchProviderNS(key)) + if saved == nil { + saved = map[string]interface{}{} + } + saved["status"] = "connected" + setting.Global.Set(scope, searchProviderNS(key), saved) + } + + response.RespondWithSuccess(c, http.StatusOK, SearchTestResult{ + Success: true, + Message: "Connection successful", + LatencyMs: latency, + }) +} + +func searchTestTavily(apiKey string) error { + payload, _ := json.Marshal(map[string]interface{}{ + "api_key": apiKey, + "query": "test", + }) + client := &http.Client{Timeout: 15 * time.Second} + resp, err := client.Post("https://api.tavily.com/search", "application/json", bytes.NewReader(payload)) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +func searchTestSerper(apiKey string) error { + payload, _ := json.Marshal(map[string]string{"q": "test"}) + client := &http.Client{Timeout: 15 * time.Second} + req, err := http.NewRequest("POST", "https://google.serper.dev/search", bytes.NewReader(payload)) + if err != nil { + return fmt.Errorf("failed to build request: %s", err.Error()) + } + req.Header.Set("X-API-KEY", apiKey) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +func searchTestBrightdata(apiKey, zone string) error { + if zone == "" { + return fmt.Errorf("Zone is required") + } + client := &http.Client{Timeout: 10 * time.Second} + req, err := http.NewRequest("GET", "https://api.brightdata.com/zone/status?zone="+zone, nil) + if err != nil { + return fmt.Errorf("failed to build request: %s", err.Error()) + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("connection failed: %s", err.Error()) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return fmt.Errorf("invalid API key (HTTP %d)", resp.StatusCode) + } + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("zone '%s' not found", zone) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned HTTP %d", resp.StatusCode) + } + return nil +} + +// --------------------------------------------------------------------------- +// PUT /setting/search/tool-assignment +// --------------------------------------------------------------------------- + +func handleSearchToolAssignment(c *gin.Context) { + if !guardOwner(c) { + return + } + + var body struct { + WebSearch *string `json:"web_search"` + WebScrape *string `json:"web_scrape"` + } + if err := c.ShouldBindJSON(&body); err != nil { + respondError(c, http.StatusBadRequest, "invalid request body") + return + } + + if setting.Global == nil { + respondError(c, http.StatusInternalServerError, "setting registry not initialized") + return + } + + info := authorized.GetInfo(c) + scope := searchScope(info) + + // Validate: provider must be enabled and support the tool + validateAssignment := func(providerKey *string, toolType string) error { + if providerKey == nil || *providerKey == "" { + return nil + } + preset := searchFindPreset(*providerKey) + if preset == nil { + return fmt.Errorf("unknown provider: %s", *providerKey) + } + + hasTools := false + for _, t := range preset.Tools { + if t == toolType { + hasTools = true + break + } + } + if !hasTools { + return fmt.Errorf("provider %s does not support %s", *providerKey, toolType) + } + + if preset.IsCloud { + return nil // cloud provider enablement is implicit + } + + saved, _ := setting.Global.Get(scope, searchProviderNS(*providerKey)) + if saved != nil { + if v, ok := saved["enabled"].(bool); ok && v { + return nil + } + } + return fmt.Errorf("provider %s is not enabled", *providerKey) + } + + if err := validateAssignment(body.WebSearch, "web_search"); err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + if err := validateAssignment(body.WebScrape, "web_scrape"); err != nil { + respondError(c, http.StatusBadRequest, err.Error()) + return + } + + m := make(map[string]interface{}) + if body.WebSearch != nil { + m["web_search"] = *body.WebSearch + } else { + m["web_search"] = "" + } + if body.WebScrape != nil { + m["web_scrape"] = *body.WebScrape + } else { + m["web_scrape"] = "" + } + + if _, err := setting.Global.Set(scope, searchAssignmentNS, m); err != nil { + respondError(c, http.StatusInternalServerError, err.Error()) + return + } + + result := SearchToolAssignment{} + if v, ok := m["web_search"].(string); ok && v != "" { + result.WebSearch = &v + } + if v, ok := m["web_scrape"].(string); ok && v != "" { + result.WebScrape = &v + } + + response.RespondWithSuccess(c, http.StatusOK, result) +} diff --git a/openapi/setting/search_presets.yml b/openapi/setting/search_presets.yml new file mode 100644 index 00000000..8775c56e --- /dev/null +++ b/openapi/setting/search_presets.yml @@ -0,0 +1,61 @@ +# Search & Scrape provider presets. +# Embedded at compile time via go:embed in search.go. + +- key: cloud + name: Yao Agents + description: + zh-CN: "云服务提供的搜索与抓取能力,凭证来自云服务配置页" + en-US: "Search & scrape capabilities from cloud service, credentials from cloud config" + website: "https://yaoagents.com" + tools: [web_search, web_scrape] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + - { zh-CN: "网页抓取", en-US: "Web Scrape" } + fields: [] + is_cloud: true + +- key: tavily + name: Tavily + description: + zh-CN: "AI 优化的搜索 API,返回结构化结果,适合 Agent 使用" + en-US: "AI-optimized search API with structured results, ideal for agents" + website: "https://tavily.com" + tools: [web_search] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + +- key: serper + name: "Serper (Google)" + description: + zh-CN: "基于 Google 搜索的 API,价格实惠,结果质量高" + en-US: "Google Search API with affordable pricing and high-quality results" + website: "https://serper.dev" + tools: [web_search] + tool_labels: + - { zh-CN: "网页搜索", en-US: "Web Search" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + +- key: brightdata + name: Brightdata + description: + zh-CN: "部分网站有访问限制,启用代理可提升抓取成功率。需开通 Web Access API (Web Unlocker)。" + en-US: "Some websites have access restrictions. Enabling proxy improves scraping success rate. Requires Web Access API (Web Unlocker)." + website: "https://brightdata.com" + tools: [web_scrape] + tool_labels: + - { zh-CN: "网页抓取", en-US: "Web Scrape" } + fields: + - key: api_key + label: { zh-CN: "API Key", en-US: "API Key" } + type: password + - key: zone + label: { zh-CN: "Zone", en-US: "Zone" } + type: text + hint: { zh-CN: "Web Unlocker API 的 Zone 名称", en-US: "Zone name of your Web Unlocker API" } diff --git a/openapi/setting/setting.go b/openapi/setting/setting.go index a4a1c7c7..f3f43dba 100644 --- a/openapi/setting/setting.go +++ b/openapi/setting/setting.go @@ -45,6 +45,13 @@ func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) { 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) } // requireOwner checks that the current user is the team owner. diff --git a/openapi/setting/types.go b/openapi/setting/types.go index fef5b56b..cc2da466 100644 --- a/openapi/setting/types.go +++ b/openapi/setting/types.go @@ -92,3 +92,51 @@ type LLMPageData struct { 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"` +} diff --git a/openapi/tests/setting/search_test.go b/openapi/tests/setting/search_test.go new file mode 100644 index 00000000..28100de5 --- /dev/null +++ b/openapi/tests/setting/search_test.go @@ -0,0 +1,339 @@ +package setting_test + +import ( + "bytes" + "encoding/json" + "net/http" + "os" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// --------------------------------------------------------------------------- +// Functional tests (system:root token) +// --------------------------------------------------------------------------- + +func TestSearchGet(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := http.DefaultClient.Do(req) + if !assert.NoError(t, err) || !assert.NotNil(t, resp) { + return + } + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&body) + assert.NoError(t, err) + + assert.Contains(t, body, "presets") + assert.Contains(t, body, "providers") + assert.Contains(t, body, "tool_assignment") + + presets, ok := body["presets"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 4, len(presets), "should have 4 presets: cloud, tavily, serper, brightdata") + + providers, ok := body["providers"].([]interface{}) + assert.True(t, ok) + assert.Equal(t, 4, len(providers), "should have 4 provider configs") + + // Cloud provider should be first + first, _ := providers[0].(map[string]interface{}) + assert.Equal(t, "cloud", first["preset_key"]) +} + +func TestSearchGetUnauthenticated(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + req, err := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + assert.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestSearchProviderUpdate(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{ + "api_key": "tvly-test-key-12345", + }, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, "tavily", body["preset_key"]) + + // api_key should be masked in response + fv, _ := body["field_values"].(map[string]interface{}) + maskedKey, _ := fv["api_key"].(string) + assert.True(t, strings.Contains(maskedKey, "..."), "api_key should be masked") + + // GET should also return masked key + req2, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req2.Header.Set("Authorization", "Bearer "+token) + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + + var getData map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&getData) + providers, _ := getData["providers"].([]interface{}) + for _, p := range providers { + pm, _ := p.(map[string]interface{}) + if pm["preset_key"] == "tavily" { + tfv, _ := pm["field_values"].(map[string]interface{}) + assert.True(t, strings.Contains(tfv["api_key"].(string), "..."), "GET should return masked key") + } + } +} + +func TestSearchProviderUpdateCloud(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{}, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/cloud", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "cloud provider should be rejected") +} + +func TestSearchProviderToggle(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save tavily first + savePayload := map[string]interface{}{ + "field_values": map[string]string{"api_key": "tvly-toggle-key"}, + } + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + resp.Body.Close() + + // Enable tavily + enablePayload := map[string]interface{}{"enabled": true} + raw, _ = json.Marshal(enablePayload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusOK, resp2.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp2.Body).Decode(&body) + assert.Equal(t, true, body["enabled"]) + + // Assign tavily to web_search + assignPayload := map[string]interface{}{"web_search": "tavily"} + raw, _ = json.Marshal(assignPayload) + req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req3.Header.Set("Authorization", "Bearer "+token) + req3.Header.Set("Content-Type", "application/json") + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + resp3.Body.Close() + + // Disable tavily -- should clear tool_assignment + disablePayload := map[string]interface{}{"enabled": false} + raw, _ = json.Marshal(disablePayload) + req4, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req4.Header.Set("Authorization", "Bearer "+token) + req4.Header.Set("Content-Type", "application/json") + resp4, err := http.DefaultClient.Do(req4) + assert.NoError(t, err) + resp4.Body.Close() + + // Verify tool_assignment cleared + req5, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req5.Header.Set("Authorization", "Bearer "+token) + resp5, err := http.DefaultClient.Do(req5) + assert.NoError(t, err) + defer resp5.Body.Close() + + var getData map[string]interface{} + json.NewDecoder(resp5.Body).Decode(&getData) + ta, _ := getData["tool_assignment"].(map[string]interface{}) + assert.Nil(t, ta["web_search"], "web_search should be cleared after disabling tavily") +} + +func TestSearchToolAssignment(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Save and enable tavily + savePayload := map[string]interface{}{"field_values": map[string]string{"api_key": "tvly-assign-key"}} + raw, _ := json.Marshal(savePayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, _ := http.DefaultClient.Do(req) + resp.Body.Close() + + enablePayload := map[string]interface{}{"enabled": true} + raw, _ = json.Marshal(enablePayload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily/toggle", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+token) + req2.Header.Set("Content-Type", "application/json") + resp2, _ := http.DefaultClient.Do(req2) + resp2.Body.Close() + + // Assign tavily to web_search + assignPayload := map[string]interface{}{"web_search": "tavily"} + raw, _ = json.Marshal(assignPayload) + req3, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req3.Header.Set("Authorization", "Bearer "+token) + req3.Header.Set("Content-Type", "application/json") + resp3, err := http.DefaultClient.Do(req3) + assert.NoError(t, err) + defer resp3.Body.Close() + assert.Equal(t, http.StatusOK, resp3.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp3.Body).Decode(&body) + assert.Equal(t, "tavily", body["web_search"]) +} + +func TestSearchToolAssignmentDisabledProvider(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + // Try to assign a provider that isn't enabled + assignPayload := map[string]interface{}{"web_search": "serper"} + raw, _ := json.Marshal(assignPayload) + req, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/tool-assignment", bytes.NewReader(raw)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode, "should reject assignment to disabled provider") +} + +func TestSearchProviderTest(t *testing.T) { + apiKey := os.Getenv("TAVILY_API_KEY") + if apiKey == "" { + t.Skip("TAVILY_API_KEY not set, skipping search provider test") + } + + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + payload := map[string]interface{}{ + "field_values": map[string]string{"api_key": apiKey}, + } + raw, _ := json.Marshal(payload) + req, err := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/tavily/test", bytes.NewReader(raw)) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var body map[string]interface{} + json.NewDecoder(resp.Body).Decode(&body) + assert.Equal(t, true, body["success"]) +} + +func TestSearchProviderTestCloud(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + token := obtainToken(t, serverURL) + + req, _ := http.NewRequest("POST", serverURL+baseURL()+"/setting/search/providers/cloud/test", nil) + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +// --------------------------------------------------------------------------- +// ACL permission tests +// --------------------------------------------------------------------------- + +func TestSearchACL_ReadOnlyScopeCannotWrite(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + initSettingRegistry(t) + + readToken := obtainRestrictedToken(t, serverURL, "setting:search:read:all") + + // GET should work + req, _ := http.NewRequest("GET", serverURL+baseURL()+"/setting/search", nil) + req.Header.Set("Authorization", "Bearer "+readToken) + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "read-only scope should allow GET") + + // PUT should be denied + payload := map[string]interface{}{ + "field_values": map[string]string{"api_key": "tvly-acl-test"}, + } + raw, _ := json.Marshal(payload) + req2, _ := http.NewRequest("PUT", serverURL+baseURL()+"/setting/search/providers/tavily", bytes.NewReader(raw)) + req2.Header.Set("Authorization", "Bearer "+readToken) + req2.Header.Set("Content-Type", "application/json") + resp2, err := http.DefaultClient.Do(req2) + assert.NoError(t, err) + defer resp2.Body.Close() + assert.Equal(t, http.StatusForbidden, resp2.StatusCode, "read-only scope should deny PUT") +}