Inject web search capability for Claude Console (API Key) accounts that don't natively support Anthropic's web_search tool. When a pure web_search request is detected, the gateway calls Brave Search or Tavily API directly and constructs an Anthropic-protocol-compliant SSE/JSON response without forwarding to upstream. Backend: - New `pkg/websearch/` SDK: Brave and Tavily provider implementations with io.LimitReader, proxy support, and Redis-based quota tracking (Lua atomic INCR + TTL, DECR rollback on failure) - Global config via `settings.web_search_emulation_config` (JSON) with in-process cache + singleflight, input validation, API key merge on save, and sanitized API responses - Channel-level toggle via `channels.features_config` JSONB column (DB migration 101) - Account-level toggle via `accounts.extra.web_search_emulation` - Request interception in `Forward()` with SSE streaming response construction using json.Marshal (no manual string concatenation) - Manager hot-reload: `RebuildWebSearchManager()` called on config save and startup via `SetWebSearchRedisClient()` - 70 unit tests covering providers, manager, config validation, sanitization, tool detection, query extraction, and response building Frontend: - Settings → Gateway tab: Web Search Emulation config card with global toggle, provider list (add/remove, API key, priority, quota, proxy) - Channels → Anthropic tab: web search emulation toggle with global state linkage (disabled when global off) - Account Create/Edit modals: web search emulation toggle for API Key type with Toggle component - Full i18n coverage (zh + en)
149 lines
4.5 KiB
Go
149 lines
4.5 KiB
Go
package service
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// --- validateWebSearchConfig ---
|
|
|
|
func TestValidateWebSearchConfig_Nil(t *testing.T) {
|
|
require.NoError(t, validateWebSearchConfig(nil))
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_Valid(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Enabled: true,
|
|
Providers: []WebSearchProviderConfig{
|
|
{Type: "brave", Priority: 1, QuotaLimit: 1000, QuotaRefreshInterval: "monthly"},
|
|
{Type: "tavily", Priority: 2, QuotaLimit: 500, QuotaRefreshInterval: "daily"},
|
|
},
|
|
}
|
|
require.NoError(t, validateWebSearchConfig(cfg))
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_TooManyProviders(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{Providers: make([]WebSearchProviderConfig, 11)}
|
|
for i := range cfg.Providers {
|
|
cfg.Providers[i] = WebSearchProviderConfig{Type: "brave"}
|
|
}
|
|
err := validateWebSearchConfig(cfg)
|
|
require.ErrorContains(t, err, "too many providers")
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_InvalidType(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "bing"}},
|
|
}
|
|
require.ErrorContains(t, validateWebSearchConfig(cfg), "invalid type")
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_InvalidQuotaInterval(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", QuotaRefreshInterval: "hourly"}},
|
|
}
|
|
require.ErrorContains(t, validateWebSearchConfig(cfg), "invalid quota_refresh_interval")
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_NegativeQuotaLimit(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", QuotaLimit: -1}},
|
|
}
|
|
require.ErrorContains(t, validateWebSearchConfig(cfg), "quota_limit must be >= 0")
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_DuplicateType(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{
|
|
{Type: "brave", Priority: 1},
|
|
{Type: "brave", Priority: 2},
|
|
},
|
|
}
|
|
require.ErrorContains(t, validateWebSearchConfig(cfg), "duplicate type")
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_EmptyQuotaInterval(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", QuotaRefreshInterval: ""}},
|
|
}
|
|
require.NoError(t, validateWebSearchConfig(cfg))
|
|
}
|
|
|
|
func TestValidateWebSearchConfig_ZeroQuotaLimit(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", QuotaLimit: 0}},
|
|
}
|
|
require.NoError(t, validateWebSearchConfig(cfg))
|
|
}
|
|
|
|
// --- parseWebSearchConfigJSON ---
|
|
|
|
func TestParseWebSearchConfigJSON_ValidJSON(t *testing.T) {
|
|
raw := `{"enabled":true,"providers":[{"type":"brave","api_key":"sk-xxx"}]}`
|
|
cfg := parseWebSearchConfigJSON(raw)
|
|
require.True(t, cfg.Enabled)
|
|
require.Len(t, cfg.Providers, 1)
|
|
require.Equal(t, "brave", cfg.Providers[0].Type)
|
|
}
|
|
|
|
func TestParseWebSearchConfigJSON_EmptyString(t *testing.T) {
|
|
cfg := parseWebSearchConfigJSON("")
|
|
require.False(t, cfg.Enabled)
|
|
require.Empty(t, cfg.Providers)
|
|
}
|
|
|
|
func TestParseWebSearchConfigJSON_InvalidJSON(t *testing.T) {
|
|
cfg := parseWebSearchConfigJSON("not{json")
|
|
require.False(t, cfg.Enabled)
|
|
require.Empty(t, cfg.Providers)
|
|
}
|
|
|
|
// --- SanitizeWebSearchConfig ---
|
|
|
|
func TestSanitizeWebSearchConfig_MaskAPIKey(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Enabled: true,
|
|
Providers: []WebSearchProviderConfig{
|
|
{Type: "brave", APIKey: "sk-secret-xxx"},
|
|
},
|
|
}
|
|
out := SanitizeWebSearchConfig(cfg)
|
|
require.Equal(t, "", out.Providers[0].APIKey)
|
|
require.True(t, out.Providers[0].APIKeyConfigured)
|
|
}
|
|
|
|
func TestSanitizeWebSearchConfig_NoAPIKey(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: ""}},
|
|
}
|
|
out := SanitizeWebSearchConfig(cfg)
|
|
require.Equal(t, "", out.Providers[0].APIKey)
|
|
require.False(t, out.Providers[0].APIKeyConfigured)
|
|
}
|
|
|
|
func TestSanitizeWebSearchConfig_Nil(t *testing.T) {
|
|
require.Nil(t, SanitizeWebSearchConfig(nil))
|
|
}
|
|
|
|
func TestSanitizeWebSearchConfig_PreservesOtherFields(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Enabled: true,
|
|
Providers: []WebSearchProviderConfig{
|
|
{Type: "brave", APIKey: "secret", Priority: 10, QuotaLimit: 1000},
|
|
},
|
|
}
|
|
out := SanitizeWebSearchConfig(cfg)
|
|
require.True(t, out.Enabled)
|
|
require.Equal(t, 10, out.Providers[0].Priority)
|
|
require.Equal(t, int64(1000), out.Providers[0].QuotaLimit)
|
|
}
|
|
|
|
func TestSanitizeWebSearchConfig_DoesNotMutateOriginal(t *testing.T) {
|
|
cfg := &WebSearchEmulationConfig{
|
|
Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "secret"}},
|
|
}
|
|
_ = SanitizeWebSearchConfig(cfg)
|
|
require.Equal(t, "secret", cfg.Providers[0].APIKey)
|
|
}
|