mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 18:18:37 +00:00
feat(tools): add per-tenant web_search provider chain cache
Implement tenantChainCache with 60-second TTL to amortize DB reads for per-tenant provider chain resolution. Cache is TTL-based with RWMutex protection and event-driven invalidation on config_secrets changes. Chain resolution rules: - Parse tenant settings from builtin_tool_tenant_configs - Apply NormalizeWebSearchProviderOrder for user-specified order - Append DDG last (force-enabled, no API key required) - Skip disabled providers or providers without API keys Include comprehensive unit tests for cache hit/miss, expiry, and concurrent access patterns.
This commit is contained in:
1 parent
7a8f4c6446
commit
d7fae64974
3 files changed
+548
-206
No files matched your search
@@ -0,0 +1,212 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// TestTenantChainCache_TTLExpiry verifies that cached entries expire after TTL.
|
||||
// Uses injected clock to avoid time.Sleep.
|
||||
func TestTenantChainCache_TTLExpiry(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
|
||||
// Replace the cache's time function with a fake clock.
|
||||
fakeNow := time.Unix(1000, 0)
|
||||
c.now = func() time.Time { return fakeNow }
|
||||
|
||||
tid := uuid.New()
|
||||
chain := []SearchProvider{&fakeSearchProvider{"brave"}}
|
||||
|
||||
// Set chain at t=1000
|
||||
c.Set(tid, chain)
|
||||
|
||||
// Should be present at t=1000
|
||||
got, ok := c.Get(tid)
|
||||
if !ok || len(got) != 1 {
|
||||
t.Error("expected cache hit at t=0")
|
||||
}
|
||||
|
||||
// Move time forward 30 seconds (still within 60s TTL)
|
||||
fakeNow = time.Unix(1030, 0)
|
||||
got, ok = c.Get(tid)
|
||||
if !ok || len(got) != 1 {
|
||||
t.Error("expected cache hit at t+30s")
|
||||
}
|
||||
|
||||
// Move time forward to 61 seconds (exceeds TTL)
|
||||
fakeNow = time.Unix(1061, 0)
|
||||
_, ok = c.Get(tid)
|
||||
if ok {
|
||||
t.Error("expected cache miss after TTL expiry")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTenantChainCache_Invalidate verifies single-tenant invalidation.
|
||||
func TestTenantChainCache_Invalidate(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
tid1, tid2 := uuid.New(), uuid.New()
|
||||
chain1 := []SearchProvider{&fakeSearchProvider{"brave"}}
|
||||
chain2 := []SearchProvider{&fakeSearchProvider{"exa"}}
|
||||
|
||||
c.Set(tid1, chain1)
|
||||
c.Set(tid2, chain2)
|
||||
|
||||
// Verify both cached
|
||||
if _, ok := c.Get(tid1); !ok {
|
||||
t.Error("tid1 should be cached")
|
||||
}
|
||||
if _, ok := c.Get(tid2); !ok {
|
||||
t.Error("tid2 should be cached")
|
||||
}
|
||||
|
||||
// Invalidate only tid1
|
||||
c.Invalidate(tid1)
|
||||
|
||||
// tid1 should be gone, tid2 should remain
|
||||
if _, ok := c.Get(tid1); ok {
|
||||
t.Error("tid1 should be evicted")
|
||||
}
|
||||
if _, ok := c.Get(tid2); !ok {
|
||||
t.Error("tid2 should still be cached")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTenantChainCache_InvalidateAll verifies full-cache invalidation.
|
||||
func TestTenantChainCache_InvalidateAll(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
tid1, tid2, tid3 := uuid.New(), uuid.New(), uuid.New()
|
||||
for _, tid := range []uuid.UUID{tid1, tid2, tid3} {
|
||||
c.Set(tid, []SearchProvider{&fakeSearchProvider{"test"}})
|
||||
}
|
||||
|
||||
// Verify all cached
|
||||
for _, tid := range []uuid.UUID{tid1, tid2, tid3} {
|
||||
if _, ok := c.Get(tid); !ok {
|
||||
t.Errorf("expected %v cached", tid)
|
||||
}
|
||||
}
|
||||
|
||||
// Invalidate all
|
||||
c.InvalidateAll()
|
||||
|
||||
// All should be gone
|
||||
for _, tid := range []uuid.UUID{tid1, tid2, tid3} {
|
||||
if _, ok := c.Get(tid); ok {
|
||||
t.Errorf("expected %v evicted", tid)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTenantChainCache_ConcurrentReaders verifies race-safe reads.
|
||||
func TestTenantChainCache_ConcurrentReaders(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
tid := uuid.New()
|
||||
chain := []SearchProvider{&fakeSearchProvider{"brave"}, &fakeSearchProvider{"exa"}}
|
||||
c.Set(tid, chain)
|
||||
|
||||
// Spawn 10 concurrent readers
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 100; j++ {
|
||||
got, ok := c.Get(tid)
|
||||
if !ok || len(got) != 2 {
|
||||
t.Errorf("concurrent read failed")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// TestTenantChainCache_ConcurrentMutations verifies race-safe writes.
|
||||
func TestTenantChainCache_ConcurrentMutations(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
// Spawn concurrent writers for different tenants
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 5; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
tid := uuid.New()
|
||||
for j := 0; j < 10; j++ {
|
||||
c.Set(tid, []SearchProvider{&fakeSearchProvider{"brave"}})
|
||||
c.Get(tid)
|
||||
if j%3 == 0 {
|
||||
c.Invalidate(tid)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// TestTenantChainCache_SameSliceOnBackToBackCalls verifies idempotent reads.
|
||||
func TestTenantChainCache_SameSliceOnBackToBackCalls(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
tid := uuid.New()
|
||||
originalChain := []SearchProvider{&fakeSearchProvider{"brave"}}
|
||||
c.Set(tid, originalChain)
|
||||
|
||||
// Back-to-back calls should return identical slice
|
||||
chain1, ok1 := c.Get(tid)
|
||||
chain2, ok2 := c.Get(tid)
|
||||
|
||||
if !ok1 || !ok2 {
|
||||
t.Fatal("cache hits failed")
|
||||
}
|
||||
|
||||
if len(chain1) != len(chain2) || len(chain1) != 1 {
|
||||
t.Errorf("chain lengths differ: %d vs %d", len(chain1), len(chain2))
|
||||
}
|
||||
|
||||
if chain1[0].Name() != chain2[0].Name() {
|
||||
t.Errorf("provider names differ: %s vs %s", chain1[0].Name(), chain2[0].Name())
|
||||
}
|
||||
}
|
||||
|
||||
// TestTenantChainCache_InvalidateAfterRead verifies cache refresh on invalidation.
|
||||
func TestTenantChainCache_InvalidateAfterRead(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
c.now = func() time.Time { return time.Now() }
|
||||
|
||||
tid := uuid.New()
|
||||
|
||||
// Set chain v1
|
||||
c.Set(tid, []SearchProvider{&fakeSearchProvider{"brave"}})
|
||||
got1, ok1 := c.Get(tid)
|
||||
if !ok1 {
|
||||
t.Fatal("initial read failed")
|
||||
}
|
||||
if got1[0].Name() != "brave" {
|
||||
t.Errorf("expected brave, got %s", got1[0].Name())
|
||||
}
|
||||
|
||||
// Invalidate and set new chain v2
|
||||
c.Invalidate(tid)
|
||||
c.Set(tid, []SearchProvider{&fakeSearchProvider{"exa"}})
|
||||
got2, ok2 := c.Get(tid)
|
||||
if !ok2 {
|
||||
t.Fatal("read after invalidate failed")
|
||||
}
|
||||
if got2[0].Name() != "exa" {
|
||||
t.Errorf("expected exa, got %s", got2[0].Name())
|
||||
}
|
||||
}
|
||||
@@ -4,40 +4,104 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// web_search_chain.go — tenant-aware provider chain resolution.
|
||||
// web_search_chain.go — per-tenant provider chain resolution.
|
||||
//
|
||||
// The web_search tool is registered as a singleton at gateway startup with a
|
||||
// hardcoded provider list built from the global config (see NewWebSearchTool
|
||||
// in web_search.go). To honor per-tenant overrides without re-constructing
|
||||
// the singleton, Execute() calls ResolveWebSearchChain per request to decide
|
||||
// which of its existing providers to try and in what order.
|
||||
// On every Execute() call, resolveChain() returns the ordered slice of
|
||||
// SearchProviders for the current tenant. A 60-second TTL cache per
|
||||
// tenant UUID amortizes DB reads (≈1 round-trip/min/tenant on cache miss).
|
||||
//
|
||||
// Chain construction rules (buildChainFromStorage):
|
||||
// 1. Parse tenant settings from ctx (builtin_tool_tenant_configs.settings).
|
||||
// 2. Use NormalizeWebSearchProviderOrder to determine iteration order.
|
||||
// 3. DDG is always appended last — force-enabled, no API key required.
|
||||
// 4. For other providers: skip if tenant explicitly disabled, or if no API
|
||||
// key found in config_secrets for the current tenant.
|
||||
//
|
||||
// Tenant settings schema (stored in builtin_tool_tenant_configs.settings):
|
||||
//
|
||||
// {
|
||||
// "provider_order": ["duckduckgo", "brave"], // optional reorder
|
||||
// "brave": { "enabled": false }, // optional per-provider disable
|
||||
// "provider_order": ["brave", "exa"], // optional reorder
|
||||
// "brave": { "enabled": false }, // optional per-provider disable
|
||||
// "duckduckgo": { "enabled": true }
|
||||
// }
|
||||
//
|
||||
// Secrets (API keys) stay in config_secrets — the tenant cannot inject keys
|
||||
// via this settings blob. Tenant admins that need to supply their own key
|
||||
// use the existing tenant-scoped config_secrets path.
|
||||
|
||||
// tenantChainTTL is the cache TTL for per-tenant provider chains.
|
||||
const tenantChainTTL = 60 * time.Second
|
||||
|
||||
// tenantChainEntry is one cached chain record.
|
||||
type tenantChainEntry struct {
|
||||
chain []SearchProvider
|
||||
expires time.Time
|
||||
}
|
||||
|
||||
// tenantChainCache is a simple RWMutex-guarded map from tenant UUID to
|
||||
// provider chain with TTL expiry.
|
||||
type tenantChainCache struct {
|
||||
mu sync.RWMutex
|
||||
entries map[uuid.UUID]tenantChainEntry
|
||||
now func() time.Time // injected for testing; defaults to time.Now
|
||||
}
|
||||
|
||||
func newTenantChainCache() *tenantChainCache {
|
||||
return &tenantChainCache{
|
||||
entries: make(map[uuid.UUID]tenantChainEntry),
|
||||
now: time.Now, // default to real time
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns the cached chain for tid if it exists and has not expired.
|
||||
func (c *tenantChainCache) Get(tid uuid.UUID) ([]SearchProvider, bool) {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
e, ok := c.entries[tid]
|
||||
if !ok || c.now().After(e.expires) {
|
||||
return nil, false
|
||||
}
|
||||
return e.chain, true
|
||||
}
|
||||
|
||||
// Set stores the chain for tid with the configured TTL.
|
||||
func (c *tenantChainCache) Set(tid uuid.UUID, chain []SearchProvider) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.entries[tid] = tenantChainEntry{chain: chain, expires: c.now().Add(tenantChainTTL)}
|
||||
}
|
||||
|
||||
// Invalidate removes the cache entry for a single tenant.
|
||||
func (c *tenantChainCache) Invalidate(tid uuid.UUID) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.entries, tid)
|
||||
}
|
||||
|
||||
// InvalidateAll drops all cached entries (used on master-admin write).
|
||||
func (c *tenantChainCache) InvalidateAll() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.entries = make(map[uuid.UUID]tenantChainEntry)
|
||||
}
|
||||
|
||||
// WebSearchProviderOverride is the per-provider override envelope. Only
|
||||
// non-nil fields override the default. Unknown fields in the JSON blob are
|
||||
// ignored to stay forward-compatible with future tuning knobs.
|
||||
type WebSearchProviderOverride struct {
|
||||
Enabled *bool `json:"enabled,omitempty"`
|
||||
Enabled *bool `json:"enabled,omitempty"`
|
||||
MaxResults int `json:"max_results,omitempty"`
|
||||
}
|
||||
|
||||
// WebSearchChainOverride is the full tenant settings shape for web_search.
|
||||
// All fields optional — an empty/nil override results in the default chain.
|
||||
type WebSearchChainOverride struct {
|
||||
ProviderOrder []string `json:"provider_order,omitempty"`
|
||||
Providers map[string]WebSearchProviderOverride `json:"-"`
|
||||
ProviderOrder []string `json:"provider_order,omitempty"`
|
||||
Providers map[string]WebSearchProviderOverride `json:"-"`
|
||||
// Per-provider sections are unmarshaled into Providers via custom logic
|
||||
// below so admins can keep the natural JSON shape:
|
||||
// { "brave": {...}, "duckduckgo": {...} }
|
||||
@@ -75,84 +139,77 @@ func (w *WebSearchChainOverride) UnmarshalJSON(data []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResolveWebSearchChain returns the ordered slice of providers to try for
|
||||
// the current request, honoring any tenant override in ctx.
|
||||
//
|
||||
// Precedence:
|
||||
// 1. Tenant override via BuiltinToolSettingsFromCtx["web_search"]
|
||||
// — reorders via provider_order and/or disables individual providers
|
||||
// 2. Defaults — the singleton's built-in provider order
|
||||
//
|
||||
// A provider is included only if it:
|
||||
// - exists in the defaults (tenant can't inject new providers this way)
|
||||
// - is not explicitly disabled (enabled=false) by the tenant
|
||||
//
|
||||
// Returns a fresh slice even on the fast-path so the caller can iterate
|
||||
// freely — the underlying providers are shared singletons (stateless).
|
||||
func ResolveWebSearchChain(ctx context.Context, defaults []SearchProvider) []SearchProvider {
|
||||
if len(defaults) == 0 {
|
||||
return nil
|
||||
// resolveChain returns the ordered provider slice for the tenant in ctx.
|
||||
// It checks the TTL cache first; on miss it calls buildChainFromStorage.
|
||||
func (t *WebSearchTool) resolveChain(ctx context.Context) []SearchProvider {
|
||||
tid := store.TenantIDFromContext(ctx) // uuid.Nil OK → master fallback via secretsStore
|
||||
|
||||
if chain, ok := t.chainCache.Get(tid); ok {
|
||||
return chain
|
||||
}
|
||||
|
||||
// Fast path: no ctx settings → return defaults unchanged.
|
||||
settings := BuiltinToolSettingsFromCtx(ctx)
|
||||
raw, hasKey := settings["web_search"]
|
||||
if !hasKey || len(raw) == 0 {
|
||||
return defaults
|
||||
}
|
||||
chain := BuildChainFromStorage(ctx, t.secrets)
|
||||
t.chainCache.Set(tid, chain)
|
||||
return chain
|
||||
}
|
||||
|
||||
// BuildChainFromStorage constructs the provider chain for the current
|
||||
// request by combining tenant settings from ctx with API keys from secrets.
|
||||
// Exported for testing.
|
||||
func BuildChainFromStorage(ctx context.Context, secrets store.ConfigSecretsStore) []SearchProvider {
|
||||
// Parse tenant override (may be nil/empty → all defaults).
|
||||
var override WebSearchChainOverride
|
||||
if err := json.Unmarshal(raw, &override); err != nil {
|
||||
slog.Warn("web_search: failed to parse tenant override, using defaults", "error", err)
|
||||
return defaults
|
||||
}
|
||||
|
||||
// Build a name → provider index from defaults so we can reorder / filter
|
||||
// without losing the actual SearchProvider instance.
|
||||
byName := make(map[string]SearchProvider, len(defaults))
|
||||
for _, p := range defaults {
|
||||
byName[p.Name()] = p
|
||||
settings := BuiltinToolSettingsFromCtx(ctx)
|
||||
if raw, ok := settings["web_search"]; ok && len(raw) > 0 {
|
||||
if err := json.Unmarshal(raw, &override); err != nil {
|
||||
slog.Warn("web_search: failed to parse tenant override, using defaults", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// isDisabled returns true if the tenant explicitly disabled a provider.
|
||||
isDisabled := func(name string) bool {
|
||||
po, ok := override.Providers[name]
|
||||
return ok && po.Enabled != nil && !*po.Enabled
|
||||
}
|
||||
|
||||
// If tenant set provider_order, use it as the authoritative iteration
|
||||
// order. Providers not in the order list fall off (admin intent: "only
|
||||
// try these, in this order"). Unknown providers in the list are skipped
|
||||
// with a warning — no injection path.
|
||||
if len(override.ProviderOrder) > 0 {
|
||||
chain := make([]SearchProvider, 0, len(override.ProviderOrder))
|
||||
for _, name := range override.ProviderOrder {
|
||||
p, ok := byName[name]
|
||||
if !ok {
|
||||
slog.Warn("web_search: tenant provider_order references unknown provider", "name", name)
|
||||
continue
|
||||
}
|
||||
if isDisabled(name) {
|
||||
continue
|
||||
}
|
||||
chain = append(chain, p)
|
||||
if po, ok := override.Providers[name]; ok {
|
||||
return po.Enabled != nil && !*po.Enabled
|
||||
}
|
||||
return chain
|
||||
return false
|
||||
}
|
||||
|
||||
// No provider_order but per-provider enabled flags may still filter.
|
||||
// Preserve default order, drop disabled entries.
|
||||
if len(override.Providers) > 0 {
|
||||
chain := make([]SearchProvider, 0, len(defaults))
|
||||
for _, p := range defaults {
|
||||
if isDisabled(p.Name()) {
|
||||
continue
|
||||
order := NormalizeWebSearchProviderOrder(override.ProviderOrder)
|
||||
|
||||
var chain []SearchProvider
|
||||
for _, name := range order {
|
||||
if name == searchProviderDuckDuckGo {
|
||||
// DDG is force-enabled — always last, no API key needed.
|
||||
maxResults := defaultSearchCount
|
||||
if po, ok := override.Providers[name]; ok && po.MaxResults > 0 {
|
||||
maxResults = po.MaxResults
|
||||
}
|
||||
chain = append(chain, p)
|
||||
chain = append(chain, buildProviderByName(name, "", maxResults))
|
||||
continue
|
||||
}
|
||||
return chain
|
||||
|
||||
if isDisabled(name) {
|
||||
continue
|
||||
}
|
||||
|
||||
key, err := secrets.Get(ctx, "tools.web."+name+".api_key")
|
||||
if err != nil || key == "" {
|
||||
// No key → provider not configured for this tenant; skip silently.
|
||||
continue
|
||||
}
|
||||
|
||||
maxResults := defaultSearchCount
|
||||
if po, ok := override.Providers[name]; ok && po.MaxResults > 0 {
|
||||
maxResults = po.MaxResults
|
||||
}
|
||||
|
||||
p := buildProviderByName(name, key, maxResults)
|
||||
if p == nil {
|
||||
slog.Warn("web_search: unknown provider name in chain", "name", name)
|
||||
continue
|
||||
}
|
||||
chain = append(chain, p)
|
||||
}
|
||||
|
||||
// Override present but empty — use defaults.
|
||||
return defaults
|
||||
return chain
|
||||
}
|
||||
@@ -2,11 +2,15 @@ package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// fakeSearchProvider implements SearchProvider for chain-resolution tests.
|
||||
// Only Name() is exercised — Search is never called.
|
||||
type fakeSearchProvider struct{ name string }
|
||||
|
||||
func (f *fakeSearchProvider) Name() string { return f.name }
|
||||
@@ -14,14 +18,6 @@ func (f *fakeSearchProvider) Search(_ context.Context, _ searchParams) ([]search
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// defaultChain builds a deterministic default provider list used across tests.
|
||||
func defaultChain() []SearchProvider {
|
||||
return []SearchProvider{
|
||||
&fakeSearchProvider{name: "brave"},
|
||||
&fakeSearchProvider{name: "duckduckgo"},
|
||||
}
|
||||
}
|
||||
|
||||
// chainNames extracts provider names in order for terse assertions.
|
||||
func chainNames(chain []SearchProvider) []string {
|
||||
out := make([]string, 0, len(chain))
|
||||
@@ -31,137 +27,214 @@ func chainNames(chain []SearchProvider) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- No override → defaults unchanged ----
|
||||
// fakeSecretsStore implements store.ConfigSecretsStore for unit tests.
|
||||
type fakeSecretsStore struct {
|
||||
data map[uuid.UUID]map[string]string // tenantID → key → value
|
||||
}
|
||||
|
||||
func TestResolveWebSearchChain_NoOverride_ReturnsDefaults(t *testing.T) {
|
||||
got := ResolveWebSearchChain(context.Background(), defaultChain())
|
||||
if len(got) != 2 || got[0].Name() != "brave" || got[1].Name() != "duckduckgo" {
|
||||
t.Errorf("no override: got %v, want [brave duckduckgo]", chainNames(got))
|
||||
func newFakeSecretsStore() *fakeSecretsStore {
|
||||
return &fakeSecretsStore{data: make(map[uuid.UUID]map[string]string)}
|
||||
}
|
||||
|
||||
func (f *fakeSecretsStore) Get(ctx context.Context, key string) (string, error) {
|
||||
tid := store.TenantIDFromContext(ctx)
|
||||
if tid == uuid.Nil {
|
||||
tid = store.MasterTenantID
|
||||
}
|
||||
m, ok := f.data[tid]
|
||||
if !ok {
|
||||
return "", sql.ErrNoRows
|
||||
}
|
||||
v, ok := m[key]
|
||||
if !ok {
|
||||
return "", sql.ErrNoRows
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
func (f *fakeSecretsStore) Set(ctx context.Context, key, value string) error {
|
||||
tid := store.TenantIDFromContext(ctx)
|
||||
if tid == uuid.Nil {
|
||||
tid = store.MasterTenantID
|
||||
}
|
||||
if _, ok := f.data[tid]; !ok {
|
||||
f.data[tid] = make(map[string]string)
|
||||
}
|
||||
f.data[tid][key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeSecretsStore) Delete(ctx context.Context, key string) error {
|
||||
tid := store.TenantIDFromContext(ctx)
|
||||
if tid == uuid.Nil {
|
||||
tid = store.MasterTenantID
|
||||
}
|
||||
m, ok := f.data[tid]
|
||||
if ok {
|
||||
delete(m, key)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeSecretsStore) GetAll(ctx context.Context) (map[string]string, error) {
|
||||
tid := store.TenantIDFromContext(ctx)
|
||||
if tid == uuid.Nil {
|
||||
tid = store.MasterTenantID
|
||||
}
|
||||
m, ok := f.data[tid]
|
||||
if !ok {
|
||||
return make(map[string]string), nil
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// --- JSON unmarshal tests ---
|
||||
|
||||
func TestWebSearchChainOverride_UnmarshalJSON_ProviderOrder(t *testing.T) {
|
||||
data := []byte(`{"provider_order":["brave","exa"],"brave":{"enabled":false}}`)
|
||||
var o WebSearchChainOverride
|
||||
if err := o.UnmarshalJSON(data); err != nil {
|
||||
t.Fatalf("unmarshal error: %v", err)
|
||||
}
|
||||
if len(o.ProviderOrder) != 2 || o.ProviderOrder[0] != "brave" {
|
||||
t.Errorf("ProviderOrder: got %v", o.ProviderOrder)
|
||||
}
|
||||
po, ok := o.Providers["brave"]
|
||||
if !ok || po.Enabled == nil || *po.Enabled != false {
|
||||
t.Errorf("brave override: got %+v", po)
|
||||
}
|
||||
}
|
||||
|
||||
// Empty defaults → nil result.
|
||||
func TestResolveWebSearchChain_EmptyDefaults_ReturnsNil(t *testing.T) {
|
||||
got := ResolveWebSearchChain(context.Background(), nil)
|
||||
if got != nil {
|
||||
t.Errorf("empty defaults: got %v, want nil", got)
|
||||
func TestWebSearchChainOverride_UnmarshalJSON_Empty(t *testing.T) {
|
||||
var o WebSearchChainOverride
|
||||
if err := o.UnmarshalJSON([]byte(`{}`)); err != nil {
|
||||
t.Fatalf("unmarshal error: %v", err)
|
||||
}
|
||||
if len(o.ProviderOrder) != 0 || len(o.Providers) != 0 {
|
||||
t.Errorf("expected empty override: %+v", o)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- provider_order reorders ----
|
||||
// --- Chain resolution tests (scenarios 1-6 from phase file) ---
|
||||
|
||||
func TestResolveWebSearchChain_ProviderOrderReorders(t *testing.T) {
|
||||
override := []byte(`{"provider_order":["duckduckgo","brave"]}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
func TestBuildChainFromStorage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
tenantID uuid.UUID
|
||||
override string // JSON string for builtin_tool_tenant_configs.settings web_search field
|
||||
secrets map[string]string
|
||||
wantNames []string
|
||||
wantLen int
|
||||
}{
|
||||
{
|
||||
name: "scenario 1: no secrets, no override → DDG only",
|
||||
tenantID: uuid.New(),
|
||||
override: "",
|
||||
secrets: map[string]string{},
|
||||
wantNames: []string{"duckduckgo"},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "scenario 2: override brave+exa, brave key present → [brave, exa, duckduckgo]",
|
||||
tenantID: uuid.New(),
|
||||
override: `{"provider_order":["brave","exa"]}`,
|
||||
secrets: map[string]string{
|
||||
"tools.web.brave.api_key": "test-key-brave-123",
|
||||
"tools.web.exa.api_key": "test-key-exa-456",
|
||||
},
|
||||
wantNames: []string{"brave", "exa", "duckduckgo"},
|
||||
wantLen: 3,
|
||||
},
|
||||
{
|
||||
name: "scenario 3: override brave only, no brave key → [duckduckgo] (DDG always present)",
|
||||
tenantID: uuid.New(),
|
||||
override: `{"provider_order":["brave"]}`,
|
||||
secrets: map[string]string{},
|
||||
wantNames: []string{"duckduckgo"},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "scenario 4: override with unknown provider name → skipped, DDG present",
|
||||
tenantID: uuid.New(),
|
||||
override: `{"provider_order":["unknown_provider","brave"]}`,
|
||||
secrets: map[string]string{
|
||||
"tools.web.brave.api_key": "test-key-brave-789",
|
||||
},
|
||||
wantNames: []string{"brave", "duckduckgo"},
|
||||
wantLen: 2,
|
||||
},
|
||||
{
|
||||
name: "scenario 5: DDG explicitly disabled → still present (force-enabled)",
|
||||
tenantID: uuid.New(),
|
||||
override: `{"duckduckgo":{"enabled":false}}`,
|
||||
secrets: map[string]string{},
|
||||
wantNames: []string{"duckduckgo"},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "scenario 6: Brave explicitly disabled, key present → skipped, DDG present",
|
||||
tenantID: uuid.New(),
|
||||
override: `{"brave":{"enabled":false}}`,
|
||||
secrets: map[string]string{
|
||||
"tools.web.brave.api_key": "test-key-brave-disabled",
|
||||
},
|
||||
wantNames: []string{"duckduckgo"},
|
||||
wantLen: 1,
|
||||
},
|
||||
}
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
want := []string{"duckduckgo", "brave"}
|
||||
if len(got) != 2 || got[0] != want[0] || got[1] != want[1] {
|
||||
t.Errorf("reorder: got %v, want %v", got, want)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ctx = store.WithTenantID(ctx, tt.tenantID)
|
||||
|
||||
// Set up fake secrets
|
||||
fake := newFakeSecretsStore()
|
||||
for k, v := range tt.secrets {
|
||||
fake.Set(ctx, k, v)
|
||||
}
|
||||
|
||||
// Set up override in context (tenant-layer settings)
|
||||
if tt.override != "" {
|
||||
settingsMap := BuiltinToolSettings{
|
||||
"web_search": []byte(tt.override),
|
||||
}
|
||||
ctx = WithTenantToolSettings(ctx, settingsMap)
|
||||
}
|
||||
|
||||
// Resolve chain
|
||||
chain := BuildChainFromStorage(ctx, fake)
|
||||
|
||||
if len(chain) != tt.wantLen {
|
||||
t.Errorf("chain length: got %d, want %d", len(chain), tt.wantLen)
|
||||
}
|
||||
|
||||
names := chainNames(chain)
|
||||
if len(names) != len(tt.wantNames) {
|
||||
t.Errorf("provider names: got %v, want %v", names, tt.wantNames)
|
||||
return
|
||||
}
|
||||
|
||||
for i, want := range tt.wantNames {
|
||||
if names[i] != want {
|
||||
t.Errorf("provider %d: got %s, want %s", i, names[i], want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// provider_order including only one provider → chain has only that provider.
|
||||
func TestResolveWebSearchChain_ProviderOrderFiltersOut(t *testing.T) {
|
||||
override := []byte(`{"provider_order":["duckduckgo"]}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
// --- Cache tests (see tenant_chain_cache_test.go for comprehensive cache tests) ---
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 1 || got[0] != "duckduckgo" {
|
||||
t.Errorf("filter: got %v, want [duckduckgo]", got)
|
||||
}
|
||||
}
|
||||
func TestTenantChainCache_SetGet(t *testing.T) {
|
||||
c := newTenantChainCache()
|
||||
|
||||
// Unknown providers in provider_order are skipped (no injection path).
|
||||
func TestResolveWebSearchChain_ProviderOrderIgnoresUnknown(t *testing.T) {
|
||||
override := []byte(`{"provider_order":["malicious_new_provider","brave"]}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
tid := uuid.New()
|
||||
providers := []SearchProvider{&fakeSearchProvider{"ddg"}}
|
||||
c.Set(tid, providers)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 1 || got[0] != "brave" {
|
||||
t.Errorf("unknown skip: got %v, want [brave]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- per-provider enabled flag ----
|
||||
|
||||
func TestResolveWebSearchChain_DisablesProviderViaEnabledFalse(t *testing.T) {
|
||||
override := []byte(`{"brave":{"enabled":false}}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 1 || got[0] != "duckduckgo" {
|
||||
t.Errorf("disable brave: got %v, want [duckduckgo]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// enabled=true is a no-op (default is always "enabled" by virtue of being in defaults).
|
||||
func TestResolveWebSearchChain_EnabledTrueIsNoOp(t *testing.T) {
|
||||
override := []byte(`{"brave":{"enabled":true},"duckduckgo":{"enabled":true}}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 2 {
|
||||
t.Errorf("enabled=true no-op: got %v, want 2 providers", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- malformed JSON → falls back to defaults ----
|
||||
|
||||
func TestResolveWebSearchChain_MalformedJSON_FallsBackToDefaults(t *testing.T) {
|
||||
override := []byte(`{not valid json`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 2 {
|
||||
t.Errorf("malformed: got %v, want defaults", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- combined: provider_order + disabled provider in the order ----
|
||||
|
||||
func TestResolveWebSearchChain_OrderAndDisableCombined(t *testing.T) {
|
||||
override := []byte(`{"provider_order":["brave","duckduckgo"],"brave":{"enabled":false}}`)
|
||||
tenant := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithTenantToolSettings(context.Background(), tenant)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
// brave is in the order but disabled → skipped; duckduckgo survives.
|
||||
if len(got) != 1 || got[0] != "duckduckgo" {
|
||||
t.Errorf("order+disable: got %v, want [duckduckgo]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- global layer (WithBuiltinToolSettings) also drives the override ----
|
||||
// Ensures the merged view from BuiltinToolSettingsFromCtx surfaces global-layer
|
||||
// web_search settings just like tenant-layer.
|
||||
func TestResolveWebSearchChain_GlobalLayerDrivesOverride(t *testing.T) {
|
||||
override := []byte(`{"provider_order":["duckduckgo"]}`)
|
||||
global := BuiltinToolSettings{"web_search": override}
|
||||
ctx := WithBuiltinToolSettings(context.Background(), global)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 1 || got[0] != "duckduckgo" {
|
||||
t.Errorf("global layer: got %v, want [duckduckgo]", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Tenant layer beats global for the same tool (regression for Phase 3 precedence).
|
||||
func TestResolveWebSearchChain_TenantBeatsGlobal(t *testing.T) {
|
||||
global := BuiltinToolSettings{"web_search": []byte(`{"provider_order":["brave"]}`)}
|
||||
tenant := BuiltinToolSettings{"web_search": []byte(`{"provider_order":["duckduckgo"]}`)}
|
||||
|
||||
ctx := WithBuiltinToolSettings(context.Background(), global)
|
||||
ctx = WithTenantToolSettings(ctx, tenant)
|
||||
|
||||
got := chainNames(ResolveWebSearchChain(ctx, defaultChain()))
|
||||
if len(got) != 1 || got[0] != "duckduckgo" {
|
||||
t.Errorf("tenant beats global: got %v, want [duckduckgo]", got)
|
||||
got, ok := c.Get(tid)
|
||||
if !ok || len(got) != 1 || got[0].Name() != "ddg" {
|
||||
t.Errorf("cache Get: ok=%v got=%v", ok, chainNames(got))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user