From f7ec7610ae099f436409f0c4a06881a0523bb5cc Mon Sep 17 00:00:00 2001 From: Stan Date: Mon, 2 Nov 2026 16:36:26 +0500 Subject: [PATCH 1/8] feat(aimlapi): shared service key and attribution for aimlapi.com services Adds the pieces the Toolbox and MCP integrations share: - internal/aimlapi: X-AIMLAPI-Source (agent/zero), the fixed Toolbox and MCP URLs, and an exact-hostname check (https, default port, no userinfo) for api/tools/mcp.aimlapi.com before the key or attribution headers are attached. The partner id still comes from ResolvePartnerID. - internal/config: AIMLAPIServiceKey picks the key from the aimlapi profile (or AIMLAPI_API_KEY without one), reading the credential store, and withholds it when any aimlapi-keyed profile points at a proxy base URL or AIMLAPI_INFERENCE_URL is overridden. Keys that are not a single token are refused. --- internal/aimlapi/attribution.go | 57 ++++++++++ internal/aimlapi/attribution_test.go | 54 +++++++++ internal/config/aimlapi_service_key.go | 116 ++++++++++++++++++++ internal/config/aimlapi_service_key_test.go | 113 +++++++++++++++++++ 4 files changed, 340 insertions(+) create mode 100644 internal/aimlapi/attribution.go create mode 100644 internal/aimlapi/attribution_test.go create mode 100644 internal/config/aimlapi_service_key.go create mode 100644 internal/config/aimlapi_service_key_test.go diff --git a/internal/aimlapi/attribution.go b/internal/aimlapi/attribution.go new file mode 100644 index 000000000..086e3ac61 --- /dev/null +++ b/internal/aimlapi/attribution.go @@ -0,0 +1,57 @@ +package aimlapi + +import ( + "net/http" + "net/url" + "strings" +) + +const ( + // SourceHeaderName and SourceValue identify Zero as the calling agent on + // requests to aimlapi.com services (the Toolbox and the MCP server). + SourceHeaderName = "X-AIMLAPI-Source" + SourceValue = "agent/zero" + + // CanonicalInferenceBaseURL is the catalog inference endpoint. A provider + // profile whose base URL differs is treated as a proxy: its key is never sent + // to other aimlapi.com services. + CanonicalInferenceBaseURL = "https://api.aimlapi.com/v1" + // ToolboxBaseURL is the fixed AI/ML API Toolbox REST base. It is not + // configurable, so the key and attribution can only reach tools.aimlapi.com. + ToolboxBaseURL = "https://tools.aimlapi.com/v1/tools" + // MCPServerURL is the AI/ML API MCP server (streamable HTTP). + MCPServerURL = "https://mcp.aimlapi.com/mcp" +) + +// attributedHosts are the only hostnames that receive the AI/ML API key or the +// attribution headers. Matching is exact; a suffix or substring match would +// also accept hosts such as tools.aimlapi.com.example.net. +var attributedHosts = [...]string{"api.aimlapi.com", "tools.aimlapi.com", "mcp.aimlapi.com"} + +// IsAttributedURL reports whether target is an https URL on one of the +// aimlapi.com service hosts, on the default port and without userinfo. +func IsAttributedURL(target *url.URL) bool { + if target == nil || !strings.EqualFold(target.Scheme, "https") || target.User != nil { + return false + } + if port := target.Port(); port != "" && port != "443" { + return false + } + host := strings.ToLower(target.Hostname()) + for _, candidate := range attributedHosts { + if host == candidate { + return true + } + } + return false +} + +// ApplyAttributionHeaders sets the source and partner headers on header only +// when target is an aimlapi.com service URL; any other target is left as is. +func ApplyAttributionHeaders(header http.Header, target *url.URL) { + if header == nil || !IsAttributedURL(target) { + return + } + header.Set(SourceHeaderName, SourceValue) + header.Set(PartnerHeaderName, ResolvePartnerID("")) +} diff --git a/internal/aimlapi/attribution_test.go b/internal/aimlapi/attribution_test.go new file mode 100644 index 000000000..e8da9ca64 --- /dev/null +++ b/internal/aimlapi/attribution_test.go @@ -0,0 +1,54 @@ +package aimlapi + +import ( + "net/http" + "net/url" + "testing" +) + +func TestIsAttributedURLMatchesExactServiceHostsOnly(t *testing.T) { + cases := map[string]bool{ + "https://tools.aimlapi.com/v1/tools/run": true, + "https://mcp.aimlapi.com/mcp": true, + "https://api.aimlapi.com/v1": true, + "https://TOOLS.aimlapi.com:443/v1/tools": true, + "http://tools.aimlapi.com/v1/tools/run": false, + "https://tools.aimlapi.com:8443/v1/tools": false, + "https://user@tools.aimlapi.com/v1/tools": false, + "https://tools.aimlapi.com.example.net/run": false, + "https://eviltools.aimlapi.com/run": false, + "https://aimlapi.com/app": false, + "https://auth.aimlapi.com/": false, + "https://example.com/?h=tools.aimlapi.com": false, + } + for raw, want := range cases { + target, err := url.Parse(raw) + if err != nil { + t.Fatalf("parse %q: %v", raw, err) + } + if got := IsAttributedURL(target); got != want { + t.Errorf("IsAttributedURL(%q) = %v, want %v", raw, got, want) + } + } + if IsAttributedURL(nil) { + t.Fatal("IsAttributedURL(nil) = true") + } +} + +func TestApplyAttributionHeadersSetsSourceAndPartnerOnServiceHosts(t *testing.T) { + t.Setenv("AIMLAPI_PARTNER_ID", "") + header := http.Header{} + ApplyAttributionHeaders(header, &url.URL{Scheme: "https", Host: "tools.aimlapi.com", Path: "/v1/tools/run"}) + if got := header.Get(SourceHeaderName); got != "agent/zero" { + t.Fatalf("%s = %q, want agent/zero", SourceHeaderName, got) + } + if got := header.Get(PartnerHeaderName); got != DefaultPartnerID { + t.Fatalf("%s = %q, want %q", PartnerHeaderName, got, DefaultPartnerID) + } + + other := http.Header{} + ApplyAttributionHeaders(other, &url.URL{Scheme: "https", Host: "tools.aimlapi.com.example.net"}) + if len(other) != 0 { + t.Fatalf("attribution reached a non-aimlapi host: %#v", other) + } +} diff --git a/internal/config/aimlapi_service_key.go b/internal/config/aimlapi_service_key.go new file mode 100644 index 000000000..9296bea84 --- /dev/null +++ b/internal/config/aimlapi_service_key.go @@ -0,0 +1,116 @@ +package config + +import ( + "errors" + "os" + "path/filepath" + "strings" + + "github.com/Twigpine/zero/internal/aimlapi" + "github.com/Twigpine/zero/internal/providercatalog" +) + +const aimlapiKeyEnv = "AIMLAPI_API_KEY" + +var ( + // ErrAIMLAPIKeyMissing means no AI/ML API key is configured. + ErrAIMLAPIKeyMissing = errors.New("no AI/ML API key is configured; set up aimlapi.com with /provider or set AIMLAPI_API_KEY") + // ErrAIMLAPIProxyProfile means the aimlapi profile points at a proxy, so its + // key belongs to that proxy and is withheld from aimlapi.com services. + ErrAIMLAPIProxyProfile = errors.New("the aimlapi provider points at a proxy base URL; its key is not sent to aimlapi.com services") + // ErrAIMLAPIKeyInvalid means the key cannot be sent as a single Bearer token. + ErrAIMLAPIKeyInvalid = errors.New("the AI/ML API key contains whitespace, a comma or a control character") +) + +// AIMLAPIServiceKey returns the AI/ML API key that may be sent to the +// aimlapi.com Toolbox and MCP server. It is resolved again on every call, so a +// provider change is seen by the next request. +// +// The key comes from the aimlapi provider profile (the active one when it is +// aimlapi, otherwise the first), including a key held in the credential store, +// or from AIMLAPI_API_KEY when no aimlapi profile exists. The key is withheld +// when that profile, any aimlapi profile, or any profile reading +// AIMLAPI_API_KEY points at a base URL other than the canonical +// https://api.aimlapi.com/v1, and when AIMLAPI_INFERENCE_URL is overridden: +// in those cases the credential may belong to a proxy. +func AIMLAPIServiceKey(resolved ResolvedConfig, store APIKeyGetter, getenv func(string) string) (string, error) { + if getenv == nil { + getenv = os.Getenv + } + if override := strings.TrimSpace(getenv("AIMLAPI_INFERENCE_URL")); override != "" && !sameBaseURL(override, aimlapi.CanonicalInferenceBaseURL) { + return "", ErrAIMLAPIProxyProfile + } + for _, profile := range resolved.Providers { + if aimlapiKeyedProfile(profile) && !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) { + return "", ErrAIMLAPIProxyProfile + } + } + if profile, ok := aimlapiServiceProfile(resolved); ok { + if !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) { + return "", ErrAIMLAPIProxyProfile + } + if store != nil { + profile = ApplyStoredAPIKey(profile, store) + } + return validAIMLAPIServiceKey(profile.APIKey) + } + return validAIMLAPIServiceKey(getenv(aimlapiKeyEnv)) +} + +// LoadAIMLAPIServiceKey resolves the user configuration, the environment and +// the user credential store, then applies AIMLAPIServiceKey. Project config is +// not read: provider keys are user-scoped. +func LoadAIMLAPIServiceKey() (string, error) { + userConfigPath, err := DefaultUserConfigPath() + if err != nil { + return "", err + } + existing, err := existingConfigFile(userConfigPath) + if err != nil { + return "", err + } + resolved, err := Resolve(ResolveOptions{UserConfigPath: existing, ExcludeProject: true}) + if err != nil { + return "", err + } + var store APIKeyGetter + if credentials, err := ProviderKeyStoreAt(filepath.Dir(userConfigPath)); err == nil && credentials != nil { + store = credentials + } + return AIMLAPIServiceKey(resolved, store, os.Getenv) +} + +func isAIMLAPIProfile(profile ProviderProfile) bool { + return strings.EqualFold(providercatalog.NormalizeID(profile.CatalogID), providercatalog.AIMLAPIID) +} + +// aimlapiKeyedProfile reports whether profile carries an AI/ML API credential: +// an aimlapi catalog profile, or any profile that reads AIMLAPI_API_KEY. +func aimlapiKeyedProfile(profile ProviderProfile) bool { + return isAIMLAPIProfile(profile) || strings.EqualFold(strings.TrimSpace(profile.APIKeyEnv), aimlapiKeyEnv) +} + +func aimlapiServiceProfile(resolved ResolvedConfig) (ProviderProfile, bool) { + if isAIMLAPIProfile(resolved.Provider) { + return resolved.Provider, true + } + for _, profile := range resolved.Providers { + if isAIMLAPIProfile(profile) { + return profile, true + } + } + return ProviderProfile{}, false +} + +func validAIMLAPIServiceKey(raw string) (string, error) { + key := strings.TrimSpace(raw) + if key == "" { + return "", ErrAIMLAPIKeyMissing + } + for _, r := range key { + if r <= ' ' || r == 0x7f || r == ',' || (r >= 0x80 && r <= 0x9f) { + return "", ErrAIMLAPIKeyInvalid + } + } + return key, nil +} diff --git a/internal/config/aimlapi_service_key_test.go b/internal/config/aimlapi_service_key_test.go new file mode 100644 index 000000000..1fdc12696 --- /dev/null +++ b/internal/config/aimlapi_service_key_test.go @@ -0,0 +1,113 @@ +package config + +import ( + "errors" + "os" + "path/filepath" + "testing" +) + +func aimlapiServiceKeyEnv(values map[string]string) func(string) string { + return func(name string) string { return values[name] } +} + +func TestAIMLAPIServiceKeyUsesCanonicalProfileKey(t *testing.T) { + resolved := ResolvedConfig{Providers: []ProviderProfile{ + {Name: "openai", BaseURL: "https://api.openai.com/v1", APIKey: "sk-openai"}, + {Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1/", APIKey: " aiml-key "}, + }} + key, err := AIMLAPIServiceKey(resolved, nil, aimlapiServiceKeyEnv(nil)) + if err != nil || key != "aiml-key" { + t.Fatalf("AIMLAPIServiceKey = %q, %v; want aiml-key", key, err) + } +} + +func TestAIMLAPIServiceKeyReadsCredentialStore(t *testing.T) { + resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKeyStored: true}} + store := fakeKeyGetter{keys: map[string]string{"aimlapi": "stored-key"}} + key, err := AIMLAPIServiceKey(resolved, store, aimlapiServiceKeyEnv(nil)) + if err != nil || key != "stored-key" { + t.Fatalf("AIMLAPIServiceKey = %q, %v; want stored-key", key, err) + } +} + +func TestAIMLAPIServiceKeyFallsBackToEnvironmentWithoutProfile(t *testing.T) { + key, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "env-key"})) + if err != nil || key != "env-key" { + t.Fatalf("AIMLAPIServiceKey = %q, %v; want env-key", key, err) + } + if _, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIKeyMissing) { + t.Fatalf("missing key error = %v, want ErrAIMLAPIKeyMissing", err) + } +} + +func TestAIMLAPIServiceKeyWithholdsProxyCredentials(t *testing.T) { + env := aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "proxy-key"}) + cases := map[string]ResolvedConfig{ + "active aimlapi profile on a proxy": {Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://proxy.example/v1", APIKey: "proxy-key"}}, + "a second aimlapi profile on a proxy": {Providers: []ProviderProfile{ + {Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKey: "proxy-key"}, + {Name: "aimlapi-proxy", CatalogID: "aimlapi", BaseURL: "https://proxy.example/v1", APIKey: "proxy-key"}, + }}, + "a custom profile reading AIMLAPI_API_KEY": {Providers: []ProviderProfile{ + {Name: "gateway", BaseURL: "https://proxy.example/v1", APIKeyEnv: "AIMLAPI_API_KEY", APIKey: "proxy-key"}, + }}, + } + for name, resolved := range cases { + key, err := AIMLAPIServiceKey(resolved, nil, env) + if !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" { + t.Errorf("%s: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", name, key, err) + } + } + + overridden := aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "proxy-key", "AIMLAPI_INFERENCE_URL": "https://staging.example/v1"}) + if key, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, overridden); !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" { + t.Fatalf("inference URL override: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", key, err) + } +} + +func TestAIMLAPIServiceKeyRejectsKeysThatAreNotOneToken(t *testing.T) { + for _, raw := range []string{"first,second", "first\nsecond", "first\r\nsecond", "a b", "key\x00tail", "key\u0085tail"} { + resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKey: raw}} + if key, err := AIMLAPIServiceKey(resolved, nil, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIKeyInvalid) || key != "" { + t.Errorf("key %q: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIKeyInvalid", raw, key, err) + } + } +} + +func TestLoadAIMLAPIServiceKeyResolvesUserConfig(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config")) + t.Setenv("ZERO_CRED_STORAGE", "file") + t.Setenv("ZERO_PROVIDER", "") + t.Setenv("AIMLAPI_INFERENCE_URL", "") + t.Setenv("AIMLAPI_API_KEY", "env-key") + + if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" { + t.Fatalf("without config: LoadAIMLAPIServiceKey = %q, %v; want env-key", key, err) + } + + path, err := DefaultUserConfigPath() + if err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + proxy := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","baseURL":"https://proxy.example/v1","model":"m"}]}` + if err := os.WriteFile(path, []byte(proxy), 0o600); err != nil { + t.Fatal(err) + } + if key, err := LoadAIMLAPIServiceKey(); !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" { + t.Fatalf("proxy profile: LoadAIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", key, err) + } + + canonical := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi"}]}` + if err := os.WriteFile(path, []byte(canonical), 0o600); err != nil { + t.Fatal(err) + } + if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" { + t.Fatalf("canonical profile: LoadAIMLAPIServiceKey = %q, %v; want env-key", key, err) + } +} From ab392075c560df9d7f0833a5e3f782df47b5993f Mon Sep 17 00:00:00 2001 From: Stan Date: Mon, 2 Nov 2026 16:48:37 +0500 Subject: [PATCH 2/8] feat(web-search): AI/ML API Toolbox as an opt-in web_search backend ZERO_WEBSEARCH_PROVIDER=aimlapi runs the Toolbox tool web/search at the fixed https://tools.aimlapi.com/v1/tools URL. It is selected only by that setting; an AI/ML API key alone never enables it or registers web_search. - The key comes from config.LoadAIMLAPIServiceKey on every search; a missing, proxy-profile or invalid key fails before any request. ZERO_WEBSEARCH_BASE_URL and ZERO_WEBSEARCH_API_KEY are ignored. - The key and attribution headers are attached only after aimlapi.IsAttributedURL accepts the URL, and redirects are refused. - Every run carries max_cost_usd 0.025 and a fresh idempotency key; limit maps to max_results and domains to include_domains, with the local domain filter kept. - A queued run (202) is polled until it ends, honouring next_poll_hint_ms clamped to 250 ms..5 s. This backend gets a 30 s budget with wait_ms 20 s; the generic backends keep 10 s. - Non-2xx responses, is_error runs, failed/stopped/timed_out runs and malformed bodies are errors carrying the error code and request id, with the key scrubbed. --- README.md | 24 + internal/tools/web_search.go | 43 +- internal/tools/web_search_aimlapi.go | 311 ++++++++++++ internal/tools/web_search_aimlapi_test.go | 592 ++++++++++++++++++++++ 4 files changed, 964 insertions(+), 6 deletions(-) create mode 100644 internal/tools/web_search_aimlapi.go create mode 100644 internal/tools/web_search_aimlapi_test.go diff --git a/README.md b/README.md index 2d872c8b2..63ae480af 100644 --- a/README.md +++ b/README.md @@ -397,6 +397,30 @@ and can rewrite every request through it, so this is the trust you accept by configuring one. Leave the variables unset for the checks to be enforced end to end. +The built-in `web_search` tool is registered only when a search backend is +configured: `ZERO_WEBSEARCH_BASE_URL` (with an optional `ZERO_WEBSEARCH_API_KEY`) +for a generic JSON endpoint, `ZERO_WEBSEARCH_PROVIDER=searxng` with a SearXNG +base URL, or the AI/ML API Toolbox: + +```bash +export ZERO_WEBSEARCH_PROVIDER=aimlapi +``` + +The AI/ML API backend runs the Toolbox tool `web/search`, and every search is +billed to your AI/ML API balance: about $0.0013 per result, and each search is +sent with a $0.025 spending cap, which the Toolbox enforces. It is used only when +`ZERO_WEBSEARCH_PROVIDER` is set to `aimlapi`; an AI/ML API key alone never +enables it. The key is the one the aimlapi provider uses (its profile, including +a key in the credential store, or `AIMLAPI_API_KEY` without a profile), resolved +again on every search; `ZERO_WEBSEARCH_API_KEY` is not used. When an aimlapi +profile points at a base URL other than `https://api.aimlapi.com/v1`, or +`AIMLAPI_INFERENCE_URL` is overridden, the key may belong to a proxy, so it is +not sent and the search fails without a request. Requests go only to +`https://tools.aimlapi.com` (`ZERO_WEBSEARCH_BASE_URL` is ignored and redirects +are refused) and carry `X-AIMLAPI-Source: agent/zero` and an +`X-AIMLAPI-Partner-ID` header. The `domains` argument is sent to the Toolbox as +`include_domains`, and results are still filtered locally. + For local dev servers, use shell commands such as `curl` through `exec_command` so the normal sandbox and permission policy applies. Long-running commands stay attached to a background terminal session and can be listed or stopped from the diff --git a/internal/tools/web_search.go b/internal/tools/web_search.go index b7cea743b..d05cd7dba 100644 --- a/internal/tools/web_search.go +++ b/internal/tools/web_search.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/Twigpine/zero/internal/config" "github.com/Twigpine/zero/internal/redaction" zeroSandbox "github.com/Twigpine/zero/internal/sandbox" ) @@ -43,6 +44,18 @@ type searchBackend interface { Search(ctx context.Context, query string, limit int) ([]searchResult, error) } +// timedSearchBackend is implemented by a backend that needs a time budget other +// than webSearchTimeout. +type timedSearchBackend interface { + searchTimeout() time.Duration +} + +// domainScopedSearchBackend is implemented by a backend that can restrict a +// search to an allowlist of hostnames itself. Run still filters the results. +type domainScopedSearchBackend interface { + searchDomains(ctx context.Context, query string, limit int, domains []string) ([]searchResult, error) +} + type webSearchTool struct { baseTool backend searchBackend @@ -132,13 +145,22 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result { } if tool.backend == nil { - return errorResult("Error: no search backend configured. Set ZERO_WEBSEARCH_BASE_URL (and ZERO_WEBSEARCH_API_KEY) to enable web_search.") + return errorResult("Error: no search backend configured. Set ZERO_WEBSEARCH_BASE_URL (and ZERO_WEBSEARCH_API_KEY), or ZERO_WEBSEARCH_PROVIDER=aimlapi to use the AI/ML API Toolbox (billed per search), to enable web_search.") } - runCtx, cancel := context.WithTimeout(ctx, webSearchTimeout) + timeout := webSearchTimeout + if timed, ok := tool.backend.(timedSearchBackend); ok { + timeout = timed.searchTimeout() + } + runCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - results, err := tool.backend.Search(runCtx, query, limit) + var results []searchResult + if scoped, ok := tool.backend.(domainScopedSearchBackend); ok && len(domains) > 0 { + results, err = scoped.searchDomains(runCtx, query, limit, domains) + } else { + results, err = tool.backend.Search(runCtx, query, limit) + } if err != nil { return errorResult("Error performing web search: " + redactWebSearchText(err.Error())) } @@ -188,9 +210,18 @@ func formatSearchResults(results []searchResult) string { return strings.Join(lines, "\n") } -// defaultSearchBackend returns the env-configured generic backend, or nil when -// ZERO_WEBSEARCH_BASE_URL is unset (the tool then reports it as unconfigured). +// defaultSearchBackend returns the env-configured backend, or nil when none is +// configured (the tool then reports it as unconfigured). +// +// ZERO_WEBSEARCH_PROVIDER=aimlapi selects the AI/ML API Toolbox. It ignores +// ZERO_WEBSEARCH_BASE_URL and ZERO_WEBSEARCH_API_KEY: the Toolbox host is fixed +// and the key is the AI/ML API service key. Otherwise the generic backend is +// used when ZERO_WEBSEARCH_BASE_URL is set. func defaultSearchBackend() searchBackend { + provider := strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_PROVIDER")) + if strings.EqualFold(provider, aimlapiSearchProvider) { + return newAIMLAPISearchBackend(nil, config.LoadAIMLAPIServiceKey) + } baseURL := strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_BASE_URL")) if baseURL == "" { return nil @@ -202,7 +233,7 @@ func defaultSearchBackend() searchBackend { }, baseURL: baseURL, apiKey: strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_API_KEY")), - provider: strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_PROVIDER")), + provider: provider, } } diff --git a/internal/tools/web_search_aimlapi.go b/internal/tools/web_search_aimlapi.go new file mode 100644 index 000000000..987dbd7da --- /dev/null +++ b/internal/tools/web_search_aimlapi.go @@ -0,0 +1,311 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/Twigpine/zero/internal/aimlapi" + "github.com/Twigpine/zero/internal/redaction" +) + +const ( + // aimlapiSearchProvider is the ZERO_WEBSEARCH_PROVIDER value that selects the + // AI/ML API Toolbox. It is the only way to select it: every run is billed to + // the user, so a configured AI/ML API key never enables it by itself. + aimlapiSearchProvider = "aimlapi" + aimlapiSearchTool = "web/search" + // aimlapiSearchMaxCostUSD caps the spend of one search. web/search costs + // $0.0013 per result, at least $0.0091 and at most $0.0221 (20 results), so + // the cap always covers a run; the Toolbox refuses a run it would exceed + // before billing it. + aimlapiSearchMaxCostUSD = 0.025 + + // Toolbox searches typically finish in about 2-3 s, but a run can queue. + // wait_ms asks the Toolbox to hold the POST open until the run ends or that + // time passes, so most searches need no polling; the rest of the budget is + // left for polling a run that is still queued. The generic backends keep + // webSearchTimeout. + aimlapiSearchTimeout = 30 * time.Second + aimlapiSearchWaitMS = 20_000 + + aimlapiDefaultPollDelay = time.Second + aimlapiMinPollDelay = 250 * time.Millisecond + aimlapiMaxPollDelay = 5 * time.Second + aimlapiMaxErrorTextRunes = 300 +) + +// aimlapiSearchBackend runs the Toolbox tool web/search. The destination is +// the fixed aimlapi.ToolboxBaseURL (ZERO_WEBSEARCH_BASE_URL is ignored), and +// the key is the AI/ML API service key, resolved again on every search so a +// provider change is seen by the next call. +type aimlapiSearchBackend struct { + client *http.Client + loadKey func() (string, error) +} + +// newAIMLAPISearchBackend builds the backend; a nil transport means +// http.DefaultTransport. +func newAIMLAPISearchBackend(transport http.RoundTripper, loadKey func() (string, error)) *aimlapiSearchBackend { + return &aimlapiSearchBackend{ + client: &http.Client{ + Transport: transport, + Timeout: aimlapiSearchTimeout, + // Never follow a redirect: the key and the attribution headers must + // reach no host but the Toolbox. + CheckRedirect: func(*http.Request, []*http.Request) error { + return errors.New("refusing redirect from the AI/ML API Toolbox") + }, + }, + loadKey: loadKey, + } +} + +func (backend *aimlapiSearchBackend) searchTimeout() time.Duration { + return aimlapiSearchTimeout +} + +func (backend *aimlapiSearchBackend) Search(ctx context.Context, query string, limit int) ([]searchResult, error) { + return backend.searchDomains(ctx, query, limit, nil) +} + +// searchDomains passes the normalized allowlist as include_domains so the +// Toolbox filters server-side; Run still applies filterWebSearchByDomains. +func (backend *aimlapiSearchBackend) searchDomains(ctx context.Context, query string, limit int, domains []string) ([]searchResult, error) { + key, err := backend.loadKey() + if err != nil { + return nil, fmt.Errorf("AI/ML API search is unavailable: %w", err) + } + idempotencyKey, err := aimlapi.NewPaymentSessionID() + if err != nil { + return nil, fmt.Errorf("generate idempotency key: %w", err) + } + input := map[string]any{"query": query, "max_results": limit} + if len(domains) > 0 { + input["include_domains"] = domains + } + payload, err := json.Marshal(map[string]any{ + "tool": aimlapiSearchTool, + "input": input, + "max_cost_usd": aimlapiSearchMaxCostUSD, + "wait_ms": aimlapiSearchWaitMS, + // A fresh key per search: the Toolbox coalesces a retried POST onto the + // same run instead of billing a second one. + "idempotency_key": idempotencyKey, + }) + if err != nil { + return nil, fmt.Errorf("encode AI/ML API search request: %w", err) + } + + run, err := backend.do(ctx, http.MethodPost, aimlapi.ToolboxBaseURL+"/run", payload, key) + if err != nil { + return nil, err + } + for !aimlapiTerminalRunStatus(run.Status) && !run.IsError { + if strings.TrimSpace(run.ID) == "" { + return nil, errors.New("AI/ML API search returned a queued run without an id") + } + timer := time.NewTimer(aimlapiPollDelay(run.NextPollHintMS)) + select { + case <-ctx.Done(): + timer.Stop() + return nil, fmt.Errorf("AI/ML API search run %s did not finish: %w", aimlapiErrorText(run.ID, key), ctx.Err()) + case <-timer.C: + } + run, err = backend.do(ctx, http.MethodGet, aimlapi.ToolboxBaseURL+"/runs/"+url.PathEscape(run.ID), nil, key) + if err != nil { + return nil, err + } + } + return aimlapiRunResults(run, key) +} + +type aimlapiErrorBody struct { + Code string `json:"code"` + Message string `json:"message"` + RequestID string `json:"request_id"` +} + +type aimlapiRun struct { + ID string `json:"id"` + Status string `json:"status"` + IsError bool `json:"is_error"` + NextPollHintMS *float64 `json:"next_poll_hint_ms"` + Error *aimlapiErrorBody `json:"error"` + Output *struct { + Results []struct { + URL string `json:"url"` + Title string `json:"title"` + Snippet string `json:"snippet"` + } `json:"results"` + } `json:"output"` +} + +// do sends one Toolbox request and decodes the run object it returns. +func (backend *aimlapiSearchBackend) do(ctx context.Context, method, target string, payload []byte, key string) (aimlapiRun, error) { + var body io.Reader + if payload != nil { + body = bytes.NewReader(payload) + } + request, err := newAIMLAPIToolboxRequest(ctx, method, target, body, key) + if err != nil { + return aimlapiRun{}, err + } + response, err := backend.client.Do(request) + if err != nil { + return aimlapiRun{}, fmt.Errorf("AI/ML API search request failed: %s", aimlapiErrorText(err.Error(), key)) + } + defer response.Body.Close() + raw, err := io.ReadAll(io.LimitReader(response.Body, webSearchBodyLimit)) + if err != nil { + return aimlapiRun{}, fmt.Errorf("read AI/ML API search response: %s", aimlapiErrorText(err.Error(), key)) + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + return aimlapiRun{}, aimlapiHTTPError(response, raw, key) + } + var run aimlapiRun + if err := json.Unmarshal(raw, &run); err != nil { + return aimlapiRun{}, fmt.Errorf("AI/ML API search returned an invalid response (HTTP %d)", response.StatusCode) + } + return run, nil +} + +// newAIMLAPIToolboxRequest builds a Toolbox request and attaches the key and +// the attribution headers only after checking that target is an aimlapi.com +// service URL; any other target is refused before anything is sent. +func newAIMLAPIToolboxRequest(ctx context.Context, method, target string, body io.Reader, key string) (*http.Request, error) { + request, err := http.NewRequestWithContext(ctx, method, target, body) + if err != nil { + return nil, fmt.Errorf("build AI/ML API search request: %w", err) + } + if !aimlapi.IsAttributedURL(request.URL) { + return nil, fmt.Errorf("AI/ML API search refused to send credentials to %q", request.URL.Redacted()) + } + request.Header.Set("Accept", "application/json") + if body != nil { + request.Header.Set("Content-Type", "application/json") + } + request.Header.Set("User-Agent", "zero-web-search/0.1") + request.Header.Set("Authorization", "Bearer "+key) + aimlapi.ApplyAttributionHeaders(request.Header, request.URL) + return request, nil +} + +// aimlapiHTTPError reports a refused request by status, error code, message +// and request id. The raw body is never surfaced, and the key is scrubbed from +// every field taken from it. +func aimlapiHTTPError(response *http.Response, raw []byte, key string) error { + var envelope struct { + Error *aimlapiErrorBody `json:"error"` + } + _ = json.Unmarshal(raw, &envelope) + var parts aimlapiErrorBody + if envelope.Error != nil { + parts = *envelope.Error + } + text := fmt.Sprintf("AI/ML API search error HTTP %d", response.StatusCode) + if code := aimlapiErrorText(parts.Code, key); code != "" { + text += " " + code + } + switch { + case response.StatusCode == http.StatusUnauthorized: + text += " (check the AI/ML API key)" + case parts.Code == "insufficient_funds": + text += " (top up your AI/ML API balance)" + case response.StatusCode == http.StatusTooManyRequests: + if seconds, err := strconv.Atoi(strings.TrimSpace(response.Header.Get("Retry-After"))); err == nil && seconds > 0 { + text += fmt.Sprintf(" (rate limited; retry after %ds)", seconds) + } else { + text += " (rate limited)" + } + } + if message := aimlapiErrorText(parts.Message, key); message != "" { + text += ": " + message + } + if requestID := aimlapiErrorText(parts.RequestID, key); requestID != "" { + text += " (request_id: " + requestID + ")" + } + return errors.New(text) +} + +// aimlapiRunResults maps a finished run onto search results. A run that ended +// in any status but completed, or reports is_error (which a completed run can), +// is a failure. +func aimlapiRunResults(run aimlapiRun, key string) ([]searchResult, error) { + status := strings.TrimSpace(run.Status) + if status == "" { + status = "unknown" + } + if run.IsError || status != "completed" { + text := fmt.Sprintf("AI/ML API %s run %s", aimlapiSearchTool, aimlapiErrorText(status, key)) + if run.IsError && status == "completed" { + text += " with an error" + } + if run.Error != nil { + if code := aimlapiErrorText(run.Error.Code, key); code != "" { + text += " " + code + } + if message := aimlapiErrorText(run.Error.Message, key); message != "" { + text += ": " + message + } + } + if id := aimlapiErrorText(run.ID, key); id != "" { + text += " (run_id: " + id + ")" + } + return nil, errors.New(text) + } + if run.Output == nil || run.Output.Results == nil { + return nil, fmt.Errorf("AI/ML API %s returned no results list", aimlapiSearchTool) + } + out := make([]searchResult, 0, len(run.Output.Results)) + for _, item := range run.Output.Results { + if strings.TrimSpace(item.URL) == "" { + continue + } + out = append(out, searchResult{Title: item.Title, URL: item.URL, Snippet: item.Snippet}) + } + return out, nil +} + +func aimlapiTerminalRunStatus(status string) bool { + switch strings.TrimSpace(status) { + case "completed", "failed", "stopped", "timed_out": + return true + } + return false +} + +// aimlapiPollDelay honours next_poll_hint_ms, clamped to 250 ms..5 s. +func aimlapiPollDelay(hintMS *float64) time.Duration { + if hintMS == nil { + return aimlapiDefaultPollDelay + } + delay := time.Duration(*hintMS * float64(time.Millisecond)) + if delay < aimlapiMinPollDelay { + return aimlapiMinPollDelay + } + if delay > aimlapiMaxPollDelay { + return aimlapiMaxPollDelay + } + return delay +} + +// aimlapiErrorText scrubs the key from a field taken from a response, then +// collapses whitespace and clips it. Scrubbing runs first so clipping can never +// leave a fragment of the key behind. +func aimlapiErrorText(value, key string) string { + text := redaction.RedactString(value, redaction.Options{ExtraSecretValues: []string{key}}) + text = strings.Join(strings.Fields(text), " ") + if runes := []rune(text); len(runes) > aimlapiMaxErrorTextRunes { + text = string(runes[:aimlapiMaxErrorTextRunes]) + "…" + } + return text +} diff --git a/internal/tools/web_search_aimlapi_test.go b/internal/tools/web_search_aimlapi_test.go new file mode 100644 index 000000000..0e8f9e30e --- /dev/null +++ b/internal/tools/web_search_aimlapi_test.go @@ -0,0 +1,592 @@ +package tools + +import ( + "context" + "encoding/json" + "io" + "net/http" + "path/filepath" + "regexp" + "strings" + "sync" + "testing" + "time" + + "github.com/Twigpine/zero/internal/aimlapi" + "github.com/Twigpine/zero/internal/config" +) + +// Trimmed from a real 200 web/search run. +const aimlapiCompletedRun = `{"id":"h3X2expTh3D7vd4pwyxiY","tool":"web/search","status":"completed","is_error":false,"untrusted":true, +"output":{"results":[ + {"url":"https://pkg.go.dev/net/http","title":"http package - net/http - Go Packages","snippet":"For control over proxies, TLS configuration, keep-alives..."}, + {"url":"https://go.dev/src/net/http/client.go?m=text","title":"go.dev","snippet":"// Timeout specifies a time limit for requests made by this Client."}, + {"url":"https://stackoverflow.com/questions/16895294/how-to-set-timeout-for-http-get-requests-in-golang","title":"How to set timeout for http.Get() requests in Golang?","snippet":"Apparently in Go 1.3 http.Client has Timeout field","published_at":"2013-11-02T00:00:00.000Z"} +]}, +"cost":{"quoted_usd":0.0091,"actual_usd":0.0091,"over_max":false}}` + +// A real 400 upfront refusal (not billed). +const aimlapiMaxCostExceeded = `{"error":{"message":"The cheapest source costs 0.0091 USD, above the limit of 0.001 USD.","type":"invalid_request_error","code":"max_cost_exceeded","request_id":"hnbQUvsseM1EX9RrgPjBT","refused_by":"request","estimated_cost_usd":0.0091,"max_cost_usd":0.001}}` + +type toolboxCall struct { + method string + url string + header http.Header + body map[string]any + remaining time.Duration +} + +// fakeToolbox is an http.RoundTripper that records every request and answers +// from respond, so no request leaves the process. +type fakeToolbox struct { + mu sync.Mutex + calls []toolboxCall + respond func(index int, request *http.Request) *http.Response +} + +func (fake *fakeToolbox) RoundTrip(request *http.Request) (*http.Response, error) { + call := toolboxCall{method: request.Method, url: request.URL.String(), header: request.Header.Clone()} + if deadline, ok := request.Context().Deadline(); ok { + call.remaining = time.Until(deadline) + } + if request.Body != nil { + raw, _ := io.ReadAll(request.Body) + _ = request.Body.Close() + if len(raw) > 0 { + _ = json.Unmarshal(raw, &call.body) + } + } + fake.mu.Lock() + index := len(fake.calls) + fake.calls = append(fake.calls, call) + fake.mu.Unlock() + // Like a real transport, refuse a request whose context has already ended. + if err := request.Context().Err(); err != nil { + return nil, err + } + response := fake.respond(index, request) + response.Request = request + return response, nil +} + +func (fake *fakeToolbox) recorded() []toolboxCall { + fake.mu.Lock() + defer fake.mu.Unlock() + return append([]toolboxCall(nil), fake.calls...) +} + +func toolboxResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func alwaysRespond(status int, body string) func(int, *http.Request) *http.Response { + return func(int, *http.Request) *http.Response { return toolboxResponse(status, body) } +} + +func staticKey(key string) func() (string, error) { + return func() (string, error) { return key, nil } +} + +// isolateAIMLAPISearchEnv points every config, cache and credential root at a +// temporary home and clears the variables the backend selection and the key +// resolution read. +func isolateAIMLAPISearchEnv(t *testing.T) { + t.Helper() + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config")) + t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache")) + t.Setenv("APPDATA", filepath.Join(home, "AppData", "Roaming")) + t.Setenv("LOCALAPPDATA", filepath.Join(home, "AppData", "Local")) + t.Setenv("ZERO_CRED_STORAGE", "file") + t.Setenv("ZERO_PROVIDER", "") + t.Setenv("AIMLAPI_API_KEY", "") + t.Setenv("AIMLAPI_INFERENCE_URL", "") + t.Setenv("AIMLAPI_PARTNER_ID", "") + t.Setenv("ZERO_WEBSEARCH_PROVIDER", "") + t.Setenv("ZERO_WEBSEARCH_BASE_URL", "") + t.Setenv("ZERO_WEBSEARCH_API_KEY", "") +} + +func hasWebSearchTool(tools []Tool) bool { + for _, tool := range tools { + if tool.Name() == "web_search" { + return true + } + } + return false +} + +func TestWebSearchAIMLAPISelectedOnlyByExplicitProvider(t *testing.T) { + cases := []struct { + name string + provider string + baseURL string + want bool + }{ + {name: "key alone", provider: "", want: false}, + {name: "other provider", provider: "searxng", want: false}, + {name: "explicit", provider: "aimlapi", want: true}, + {name: "case-insensitive", provider: " AIMLAPI ", want: true}, + {name: "base URL ignored", provider: "aimlapi", baseURL: "https://search.example/api", want: true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + isolateAIMLAPISearchEnv(t) + t.Setenv("AIMLAPI_API_KEY", "aiml-env-key") + t.Setenv("ZERO_WEBSEARCH_PROVIDER", tc.provider) + t.Setenv("ZERO_WEBSEARCH_BASE_URL", tc.baseURL) + + _, isAIMLAPI := defaultSearchBackend().(*aimlapiSearchBackend) + if isAIMLAPI != tc.want { + t.Fatalf("defaultSearchBackend() is aimlapi = %v, want %v", isAIMLAPI, tc.want) + } + if registered := hasWebSearchTool(CoreNetworkTools()); registered != tc.want && tc.baseURL == "" { + t.Fatalf("web_search registered = %v, want %v", registered, tc.want) + } + }) + } +} + +func TestWebSearchAIMLAPIDefaultBackendUsesServiceKeyAndFixedHost(t *testing.T) { + isolateAIMLAPISearchEnv(t) + t.Setenv("ZERO_WEBSEARCH_PROVIDER", "aimlapi") + t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api") + t.Setenv("ZERO_WEBSEARCH_API_KEY", "zws-generic-key") + t.Setenv("AIMLAPI_API_KEY", "aiml-env-key") + + backend, ok := defaultSearchBackend().(*aimlapiSearchBackend) + if !ok { + t.Fatalf("defaultSearchBackend() = %T, want *aimlapiSearchBackend", defaultSearchBackend()) + } + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)} + backend.client.Transport = fake // keeps the production redirect policy and key loader + + res := newWebSearchToolWithBackend(backend).Run(context.Background(), map[string]any{"query": "go http timeout"}) + if res.Status != StatusOK { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + calls := fake.recorded() + if len(calls) != 1 { + t.Fatalf("requests = %d, want 1", len(calls)) + } + if calls[0].url != aimlapi.ToolboxBaseURL+"/run" { + t.Fatalf("request URL = %q, want the fixed Toolbox URL", calls[0].url) + } + if got := calls[0].header.Get("Authorization"); got != "Bearer aiml-env-key" { + t.Fatalf("Authorization = %q, want the AI/ML API service key", got) + } +} + +func TestAIMLAPISearchSendsExactRequest(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + for range 2 { + res := tool.Run(context.Background(), map[string]any{"query": "go http timeout", "limit": 3}) + if res.Status != StatusOK { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + for _, want := range []string{ + "1. http package - net/http - Go Packages — https://pkg.go.dev/net/http", + " For control over proxies", + "3. How to set timeout for http.Get() requests in Golang?", + } { + if !strings.Contains(res.Output, want) { + t.Fatalf("output missing %q:\n%s", want, res.Output) + } + } + } + + calls := fake.recorded() + if len(calls) != 2 { + t.Fatalf("requests = %d, want 2", len(calls)) + } + uuid := regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`) + for _, call := range calls { + if call.method != http.MethodPost || call.url != "https://tools.aimlapi.com/v1/tools/run" { + t.Fatalf("request = %s %s, want POST https://tools.aimlapi.com/v1/tools/run", call.method, call.url) + } + for header, want := range map[string]string{ + "Authorization": "Bearer test-aiml-key", + aimlapi.SourceHeaderName: "agent/zero", + aimlapi.PartnerHeaderName: aimlapi.DefaultPartnerID, + "User-Agent": "zero-web-search/0.1", + "Content-Type": "application/json", + } { + if got := call.header.Get(header); got != want { + t.Fatalf("%s = %q, want %q", header, got, want) + } + } + body := call.body + if body["tool"] != "web/search" { + t.Fatalf("tool = %v, want web/search", body["tool"]) + } + if body["max_cost_usd"] != 0.025 { + t.Fatalf("max_cost_usd = %v, want 0.025", body["max_cost_usd"]) + } + if body["wait_ms"] != float64(20000) { + t.Fatalf("wait_ms = %v, want 20000", body["wait_ms"]) + } + input, _ := body["input"].(map[string]any) + if input["query"] != "go http timeout" || input["max_results"] != float64(3) { + t.Fatalf("input = %v, want query and max_results 3", input) + } + if _, ok := input["include_domains"]; ok { + t.Fatalf("include_domains sent without a domains argument: %v", input) + } + if key, _ := body["idempotency_key"].(string); !uuid.MatchString(key) { + t.Fatalf("idempotency_key = %q, want a UUIDv4", key) + } + // The aimlapi budget replaces the generic 10 s one, and the server-side + // wait stays inside it. + if call.remaining <= webSearchTimeout || call.remaining > aimlapiSearchTimeout { + t.Fatalf("request deadline in %v, want within (%v, %v]", call.remaining, webSearchTimeout, aimlapiSearchTimeout) + } + } + if calls[0].body["idempotency_key"] == calls[1].body["idempotency_key"] { + t.Fatalf("idempotency_key reused across searches: %v", calls[0].body["idempotency_key"]) + } + if time.Duration(aimlapiSearchWaitMS)*time.Millisecond >= aimlapiSearchTimeout { + t.Fatalf("wait_ms %d must stay below the %v budget", aimlapiSearchWaitMS, aimlapiSearchTimeout) + } +} + +func TestAIMLAPISearchResolvesKeyOnEveryCall(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)} + loads := 0 + backend := newAIMLAPISearchBackend(fake, func() (string, error) { + loads++ + if loads == 1 { + return "first-key", nil + } + return "", config.ErrAIMLAPIKeyMissing + }) + tool := newWebSearchToolWithBackend(backend) + + if res := tool.Run(context.Background(), map[string]any{"query": "q"}); res.Status != StatusOK { + t.Fatalf("first search: status = %v, output = %q", res.Status, res.Output) + } + // The key was removed after the first search: the next one must see that. + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, "no AI/ML API key is configured") { + t.Fatalf("second search: status = %v, output = %q", res.Status, res.Output) + } + if calls := fake.recorded(); len(calls) != 1 || calls[0].header.Get("Authorization") != "Bearer first-key" { + t.Fatalf("requests = %+v, want exactly the first search", calls) + } +} + +func TestAIMLAPISearchSendsNothingWithoutUsableKey(t *testing.T) { + cases := []struct { + name string + env map[string]string + want string + }{ + {name: "missing", want: config.ErrAIMLAPIKeyMissing.Error()}, + {name: "proxy", env: map[string]string{"AIMLAPI_API_KEY": "proxy-key", "AIMLAPI_INFERENCE_URL": "https://proxy.example/v1"}, want: config.ErrAIMLAPIProxyProfile.Error()}, + {name: "invalid", env: map[string]string{"AIMLAPI_API_KEY": "two words"}, want: config.ErrAIMLAPIKeyInvalid.Error()}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + isolateAIMLAPISearchEnv(t) + for name, value := range tc.env { + t.Setenv(name, value) + } + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, config.LoadAIMLAPIServiceKey)) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, tc.want) { + t.Fatalf("status = %v, output = %q, want an error mentioning %q", res.Status, res.Output, tc.want) + } + if calls := fake.recorded(); len(calls) != 0 { + t.Fatalf("sent %d requests without a usable key", len(calls)) + } + }) + } +} + +func TestAIMLAPISearchPollsQueuedRun(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response { + switch index { + case 0: + return toolboxResponse(http.StatusAccepted, `{"id":"run/1","status":"queued","next_poll_hint_ms":1}`) + case 1: + return toolboxResponse(http.StatusOK, `{"id":"run/1","status":"running","is_error":false,"next_poll_hint_ms":1}`) + default: + return toolboxResponse(http.StatusOK, aimlapiCompletedRun) + } + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusOK || !strings.Contains(res.Output, "https://pkg.go.dev/net/http") { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + calls := fake.recorded() + if len(calls) != 3 { + t.Fatalf("requests = %d, want POST then two polls", len(calls)) + } + for _, call := range calls[1:] { + if call.method != http.MethodGet || call.url != "https://tools.aimlapi.com/v1/tools/runs/run%2F1" { + t.Fatalf("poll = %s %s, want GET of the run status", call.method, call.url) + } + if call.header.Get("Authorization") != "Bearer test-aiml-key" || call.header.Get(aimlapi.SourceHeaderName) != "agent/zero" { + t.Fatalf("poll headers = %v, want the key and attribution", call.header) + } + } +} + +func TestAIMLAPIPollDelayIsClamped(t *testing.T) { + ms := func(value float64) *float64 { return &value } + cases := []struct { + hint *float64 + want time.Duration + }{ + {hint: nil, want: time.Second}, + {hint: ms(1), want: 250 * time.Millisecond}, + {hint: ms(1200), want: 1200 * time.Millisecond}, + {hint: ms(60_000), want: 5 * time.Second}, + } + for _, tc := range cases { + if got := aimlapiPollDelay(tc.hint); got != tc.want { + t.Fatalf("aimlapiPollDelay(%v) = %v, want %v", tc.hint, got, tc.want) + } + } +} + +func TestAIMLAPISearchStopsPollingWhenContextEnds(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusAccepted, `{"id":"run_1","status":"queued","next_poll_hint_ms":5000}`)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + started := time.Now() + res := tool.Run(ctx, map[string]any{"query": "q"}) + if elapsed := time.Since(started); elapsed > 2*time.Second { + t.Fatalf("search took %v after the context ended", elapsed) + } + if res.Status != StatusError || !strings.Contains(res.Output, "did not finish") { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } +} + +func TestAIMLAPISearchCompletedRunWithIsErrorFails(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, `{"id":"run_9","status":"completed","is_error":true, + "error":{"code":"source_failed","message":"upstream search failed"}, + "output":{"results":[{"url":"https://example.com/partial","title":"partial"}]}}`)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError { + t.Fatalf("status = %v, output = %q, want an error", res.Status, res.Output) + } + for _, want := range []string{"source_failed", "upstream search failed", "run_9"} { + if !strings.Contains(res.Output, want) { + t.Fatalf("output missing %q: %q", want, res.Output) + } + } + if strings.Contains(res.Output, "example.com/partial") { + t.Fatalf("results of a failed run were shown: %q", res.Output) + } +} + +func TestAIMLAPISearchTerminalFailureStatuses(t *testing.T) { + isolateAIMLAPISearchEnv(t) + for _, status := range []string{"failed", "stopped", "timed_out"} { + t.Run(status, func(t *testing.T) { + fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response { + if index > 0 { // a poll would mean the status was not treated as terminal + return toolboxResponse(http.StatusInternalServerError, `{}`) + } + return toolboxResponse(http.StatusOK, `{"id":"run_2","status":"`+status+`","is_error":false}`) + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, "run "+status) { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + if calls := fake.recorded(); len(calls) != 1 { + t.Fatalf("requests = %d, want no polling after a terminal status", len(calls)) + } + }) + } +} + +func TestAIMLAPISearchMalformedResponses(t *testing.T) { + isolateAIMLAPISearchEnv(t) + cases := map[string]struct{ body, want string }{ + "not json": {body: `oops`, want: "invalid response"}, + "no results": {body: `{"id":"run_3","status":"completed","is_error":false,"output":{}}`, want: "no results list"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, tc.body)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, tc.want) { + t.Fatalf("status = %v, output = %q, want an error mentioning %q", res.Status, res.Output, tc.want) + } + }) + } +} + +func TestAIMLAPISearchMaxCostExceededReportsCodeAndRequestID(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusBadRequest, aimlapiMaxCostExceeded)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + for _, want := range []string{"HTTP 400", "max_cost_exceeded", "cheapest source costs", "request_id: hnbQUvsseM1EX9RrgPjBT"} { + if !strings.Contains(res.Output, want) { + t.Fatalf("output missing %q: %q", want, res.Output) + } + } +} + +func TestAIMLAPISearchRateLimitedReportsRetryAfter(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: func(int, *http.Request) *http.Response { + response := toolboxResponse(http.StatusTooManyRequests, `{"error":{"code":"rate_limited","message":"slow down","request_id":"req_429"}}`) + response.Header.Set("Retry-After", "7") + return response + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + for _, want := range []string{"HTTP 429 rate_limited", "retry after 7s", "request_id: req_429"} { + if !strings.Contains(res.Output, want) { + t.Fatalf("output missing %q: %q", want, res.Output) + } + } +} + +func TestAIMLAPISearchRefusesRedirect(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response { + if index == 0 { + response := toolboxResponse(http.StatusTemporaryRedirect, ``) + response.Header.Set("Location", "https://collector.example/run") + return response + } + return toolboxResponse(http.StatusOK, aimlapiCompletedRun) + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, "refusing redirect") { + t.Fatalf("status = %v, output = %q, want a refused redirect", res.Status, res.Output) + } + if calls := fake.recorded(); len(calls) != 1 { + t.Fatalf("requests = %d, want the redirect not followed", len(calls)) + } +} + +func TestAIMLAPISearchRedactsKeyEchoedInErrors(t *testing.T) { + isolateAIMLAPISearchEnv(t) + const key = "zq7Mv2Lp9Xr4Tn8Wk3" + if redacted := redactWebSearchText("key " + key); !strings.Contains(redacted, key) { + t.Fatalf("precondition: the generic redactor already hides this key shape (%q); pick another", redacted) + } + bodies := map[string]struct { + status int + body string + }{ + "http error": {status: http.StatusUnauthorized, body: `{"error":{"code":"invalid_api_key","message":"API key ` + key + ` is not valid","request_id":"req_` + key + `"}}`}, + "run error": {status: http.StatusOK, body: `{"id":"run_` + key + `","status":"failed","is_error":true,"error":{"code":"bad_auth","message":"rejected ` + key + `"}}`}, + } + for name, tc := range bodies { + t.Run(name, func(t *testing.T) { + fake := &fakeToolbox{respond: alwaysRespond(tc.status, tc.body)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey(key))) + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + if strings.Contains(res.Output, key) { + t.Fatalf("the AI/ML API key leaked into the output: %q", res.Output) + } + if !strings.Contains(res.Output, "req_") && !strings.Contains(res.Output, "run_") { + t.Fatalf("output lost the request or run id: %q", res.Output) + } + }) + } +} + +func TestAIMLAPIToolboxRequestRefusesOtherDestinations(t *testing.T) { + isolateAIMLAPISearchEnv(t) + for _, target := range []string{ + "https://tools.aimlapi.com.evil.example/v1/tools/run", + "http://tools.aimlapi.com/v1/tools/run", + "https://tools.aimlapi.com:8443/v1/tools/run", + "https://user@tools.aimlapi.com/v1/tools/run", + "https://search.example/api", + } { + request, err := newAIMLAPIToolboxRequest(context.Background(), http.MethodPost, target, strings.NewReader(`{}`), "test-aiml-key") + if err == nil || request != nil { + t.Fatalf("newAIMLAPIToolboxRequest(%q) = %v, %v; want a refusal", target, request, err) + } + if strings.Contains(err.Error(), "test-aiml-key") { + t.Fatalf("refusal leaked the key: %v", err) + } + } + request, err := newAIMLAPIToolboxRequest(context.Background(), http.MethodGet, aimlapi.ToolboxBaseURL+"/runs/x", nil, "test-aiml-key") + if err != nil { + t.Fatalf("Toolbox URL refused: %v", err) + } + if request.Header.Get("Authorization") != "Bearer test-aiml-key" || request.Header.Get(aimlapi.SourceHeaderName) != "agent/zero" { + t.Fatalf("headers = %v, want key and attribution on the Toolbox URL", request.Header) + } +} + +func TestAIMLAPISearchDomainsAreSentAndStillFiltered(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + + res := tool.Run(context.Background(), map[string]any{ + "query": "go http timeout", + "domains": []any{"https://WWW.Go.dev/doc", "pkg.go.dev", "go.dev"}, + }) + if res.Status != StatusOK { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } + calls := fake.recorded() + if len(calls) != 1 { + t.Fatalf("requests = %d, want 1", len(calls)) + } + input, _ := calls[0].body["input"].(map[string]any) + if got, _ := json.Marshal(input["include_domains"]); string(got) != `["go.dev","pkg.go.dev"]` { + t.Fatalf("include_domains = %s, want the normalized allowlist", got) + } + // The Toolbox filter is not trusted alone: a result outside the allowlist + // is still dropped. + if strings.Contains(res.Output, "stackoverflow.com") { + t.Fatalf("result outside the allowlist was shown: %q", res.Output) + } + if !strings.Contains(res.Output, "https://pkg.go.dev/net/http") || !strings.Contains(res.Output, "https://go.dev/src/net/http/client.go") { + t.Fatalf("allowed results missing: %q", res.Output) + } +} + +func TestWebSearchUnconfiguredHintNamesAIMLAPI(t *testing.T) { + res := newWebSearchToolWithBackend(nil).Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, "ZERO_WEBSEARCH_PROVIDER=aimlapi") { + t.Fatalf("status = %v, output = %q", res.Status, res.Output) + } +} From 9d84453caec781f1ed83474d1cfecff37b15c7a8 Mon Sep 17 00:00:00 2001 From: Hugo Date: Tue, 3 Nov 2026 14:31:33 +0300 Subject: [PATCH 3/8] fix(aimlapi): preserve store errors and isolate Windows config tests --- internal/config/aimlapi_service_key.go | 8 ++++- internal/config/aimlapi_service_key_test.go | 40 +++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/internal/config/aimlapi_service_key.go b/internal/config/aimlapi_service_key.go index 9296bea84..f2bf62184 100644 --- a/internal/config/aimlapi_service_key.go +++ b/internal/config/aimlapi_service_key.go @@ -2,6 +2,7 @@ package config import ( "errors" + "fmt" "os" "path/filepath" "strings" @@ -74,7 +75,12 @@ func LoadAIMLAPIServiceKey() (string, error) { return "", err } var store APIKeyGetter - if credentials, err := ProviderKeyStoreAt(filepath.Dir(userConfigPath)); err == nil && credentials != nil { + credentials, storeErr := ProviderKeyStoreAt(filepath.Dir(userConfigPath)) + if storeErr != nil { + if profile, ok := aimlapiServiceProfile(resolved); ok && profile.APIKeyStored && strings.TrimSpace(profile.APIKey) == "" { + return "", fmt.Errorf("open AI/ML API credential store: %w", storeErr) + } + } else if credentials != nil { store = credentials } return AIMLAPIServiceKey(resolved, store, os.Getenv) diff --git a/internal/config/aimlapi_service_key_test.go b/internal/config/aimlapi_service_key_test.go index 1fdc12696..1cd9a2535 100644 --- a/internal/config/aimlapi_service_key_test.go +++ b/internal/config/aimlapi_service_key_test.go @@ -4,6 +4,7 @@ import ( "errors" "os" "path/filepath" + "strings" "testing" ) @@ -78,6 +79,8 @@ func TestAIMLAPIServiceKeyRejectsKeysThatAreNotOneToken(t *testing.T) { func TestLoadAIMLAPIServiceKeyResolvesUserConfig(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + t.Setenv("APPDATA", filepath.Join(home, "AppData")) + t.Setenv("USERPROFILE", home) t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config")) t.Setenv("ZERO_CRED_STORAGE", "file") t.Setenv("ZERO_PROVIDER", "") @@ -111,3 +114,40 @@ func TestLoadAIMLAPIServiceKeyResolvesUserConfig(t *testing.T) { t.Fatalf("canonical profile: LoadAIMLAPIServiceKey = %q, %v; want env-key", key, err) } } + +func TestLoadAIMLAPIServiceKeyStoreOpenFailure(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("APPDATA", filepath.Join(home, "AppData")) + t.Setenv("USERPROFILE", home) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config")) + t.Setenv("ZERO_PROVIDER", "") + t.Setenv("AIMLAPI_INFERENCE_URL", "") + t.Setenv("AIMLAPI_API_KEY", "env-key") + t.Setenv("ZERO_CRED_STORAGE", "invalid-test-storage") + if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" { + t.Fatalf("environment key must not require a store: %q, %v", key, err) + } + path, err := DefaultUserConfigPath() + if err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + t.Setenv("AIMLAPI_API_KEY", "") + data := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKeyStored":true}]}` + if err := os.WriteFile(path, []byte(data), 0o600); err != nil { + t.Fatal(err) + } + if key, err := LoadAIMLAPIServiceKey(); key != "" || err == nil || !strings.Contains(err.Error(), "unknown storage") { + t.Fatalf("stored key must report store failure: %q, %v", key, err) + } + data = `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKeyStored":true,"apiKey":"inline-key"}]}` + if err := os.WriteFile(path, []byte(data), 0o600); err != nil { + t.Fatal(err) + } + if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "inline-key" { + t.Fatalf("inline key must not require a store: %q, %v", key, err) + } +} From 5322c855c9ca64912fc2e91d0746e7890dcb783b Mon Sep 17 00:00:00 2001 From: Hugo Date: Tue, 3 Nov 2026 14:33:51 +0300 Subject: [PATCH 4/8] fix(web-search): redact results and preserve domain errors --- internal/tools/web_search.go | 2 +- internal/tools/web_search_aimlapi.go | 7 ++- internal/tools/web_search_aimlapi_test.go | 58 +++++++++++++++++------ 3 files changed, 50 insertions(+), 17 deletions(-) diff --git a/internal/tools/web_search.go b/internal/tools/web_search.go index d05cd7dba..e2ecd67f5 100644 --- a/internal/tools/web_search.go +++ b/internal/tools/web_search.go @@ -167,7 +167,7 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result { // Filter by domains BEFORE the empty-result short-circuit so a non-empty // allowlist that ate every hit surfaces as a clear "no results matched // domains" error rather than a misleading "no results" message. - if len(results) > 0 && len(domains) > 0 { + if len(domains) > 0 { filtered, err := filterWebSearchByDomains(results, domains) if err != nil { return errorResult("Error: " + redactWebSearchText(err.Error())) diff --git a/internal/tools/web_search_aimlapi.go b/internal/tools/web_search_aimlapi.go index 987dbd7da..be5800cc8 100644 --- a/internal/tools/web_search_aimlapi.go +++ b/internal/tools/web_search_aimlapi.go @@ -270,7 +270,7 @@ func aimlapiRunResults(run aimlapiRun, key string) ([]searchResult, error) { if strings.TrimSpace(item.URL) == "" { continue } - out = append(out, searchResult{Title: item.Title, URL: item.URL, Snippet: item.Snippet}) + out = append(out, searchResult{Title: scrubAIMLAPIResult(item.Title, key), URL: scrubAIMLAPIResult(item.URL, key), Snippet: scrubAIMLAPIResult(item.Snippet, key)}) } return out, nil } @@ -309,3 +309,8 @@ func aimlapiErrorText(value, key string) string { } return text } + +// scrubAIMLAPIResult redacts the active credential without clipping page content. +func scrubAIMLAPIResult(value, key string) string { + return redaction.RedactString(value, redaction.Options{ExtraSecretValues: []string{key}}) +} diff --git a/internal/tools/web_search_aimlapi_test.go b/internal/tools/web_search_aimlapi_test.go index 0e8f9e30e..27c414414 100644 --- a/internal/tools/web_search_aimlapi_test.go +++ b/internal/tools/web_search_aimlapi_test.go @@ -478,22 +478,26 @@ func TestAIMLAPISearchRateLimitedReportsRetryAfter(t *testing.T) { func TestAIMLAPISearchRefusesRedirect(t *testing.T) { isolateAIMLAPISearchEnv(t) - fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response { - if index == 0 { - response := toolboxResponse(http.StatusTemporaryRedirect, ``) - response.Header.Set("Location", "https://collector.example/run") - return response - } - return toolboxResponse(http.StatusOK, aimlapiCompletedRun) - }} - tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) + for _, location := range []string{"https://collector.example/run", "https://tools.aimlapi.com/v1/tools/run"} { + t.Run(location, func(t *testing.T) { + fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response { + if index == 0 { + response := toolboxResponse(http.StatusTemporaryRedirect, ``) + response.Header.Set("Location", location) + return response + } + return toolboxResponse(http.StatusOK, aimlapiCompletedRun) + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key"))) - res := tool.Run(context.Background(), map[string]any{"query": "q"}) - if res.Status != StatusError || !strings.Contains(res.Output, "refusing redirect") { - t.Fatalf("status = %v, output = %q, want a refused redirect", res.Status, res.Output) - } - if calls := fake.recorded(); len(calls) != 1 { - t.Fatalf("requests = %d, want the redirect not followed", len(calls)) + res := tool.Run(context.Background(), map[string]any{"query": "q"}) + if res.Status != StatusError || !strings.Contains(res.Output, "refusing redirect") { + t.Fatalf("status = %v, output = %q, want a refused redirect", res.Status, res.Output) + } + if calls := fake.recorded(); len(calls) != 1 { + t.Fatalf("requests = %d, want the redirect not followed", len(calls)) + } + }) } } @@ -590,3 +594,27 @@ func TestWebSearchUnconfiguredHintNamesAIMLAPI(t *testing.T) { t.Fatalf("status = %v, output = %q", res.Status, res.Output) } } + +func TestAIMLAPISearchRedactsKeyInSuccessfulResults(t *testing.T) { + isolateAIMLAPISearchEnv(t) + const key = "zq7Mv2Lp9Xr4Tn8Wk3" + fake := &fakeToolbox{respond: func(_ int, _ *http.Request) *http.Response { + return toolboxResponse(http.StatusOK, `{"status":"completed","output":{"results":[{"title":"`+key+`","url":"https://example.com/`+key+`","snippet":"echo `+key+`"}]}}`) + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey(key))) + result := tool.Run(context.Background(), map[string]any{"query": "q"}) + if result.Status != StatusOK || strings.Contains(result.Output, key) { + t.Fatalf("successful result leaked key: %s", result.Output) + } +} +func TestAIMLAPISearchEmptyScopedResultsReportAllowlist(t *testing.T) { + isolateAIMLAPISearchEnv(t) + fake := &fakeToolbox{respond: func(_ int, _ *http.Request) *http.Response { + return toolboxResponse(http.StatusOK, `{"status":"completed","output":{"results":[]}}`) + }} + tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-key"))) + result := tool.Run(context.Background(), map[string]any{"query": "q", "domains": []any{"example.com"}}) + if result.Status != StatusError || !strings.Contains(result.Output, "loosen the allowlist") { + t.Fatalf("missing allowlist guidance: %s", result.Output) + } +} From e10a75142622fc8980b5f29c83b65ae5cc43f69d Mon Sep 17 00:00:00 2001 From: Hugo Date: Tue, 3 Nov 2026 14:35:33 +0300 Subject: [PATCH 5/8] docs(web-search): clarify minimum search charge --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 63ae480af..f31496cbc 100644 --- a/README.md +++ b/README.md @@ -407,8 +407,9 @@ export ZERO_WEBSEARCH_PROVIDER=aimlapi ``` The AI/ML API backend runs the Toolbox tool `web/search`, and every search is -billed to your AI/ML API balance: about $0.0013 per result, and each search is -sent with a $0.025 spending cap, which the Toolbox enforces. It is used only when +billed to your AI/ML API balance: about $0.0013 per result, with a documented +minimum of $0.0091 per search. Each search is sent with a $0.025 spending cap, +which the Toolbox enforces. It is used only when `ZERO_WEBSEARCH_PROVIDER` is set to `aimlapi`; an AI/ML API key alone never enables it. The key is the one the aimlapi provider uses (its profile, including a key in the credential store, or `AIMLAPI_API_KEY` without a profile), resolved From 7e668d69118922deec6a1c27b4252893b201f47b Mon Sep 17 00:00:00 2001 From: Hugo Date: Tue, 3 Nov 2026 14:41:15 +0300 Subject: [PATCH 6/8] fix(aimlapi): reject Unicode whitespace in service keys --- internal/config/aimlapi_service_key.go | 3 ++- internal/config/aimlapi_service_key_test.go | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/internal/config/aimlapi_service_key.go b/internal/config/aimlapi_service_key.go index f2bf62184..219b8e8d4 100644 --- a/internal/config/aimlapi_service_key.go +++ b/internal/config/aimlapi_service_key.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "strings" + "unicode" "github.com/Twigpine/zero/internal/aimlapi" "github.com/Twigpine/zero/internal/providercatalog" @@ -114,7 +115,7 @@ func validAIMLAPIServiceKey(raw string) (string, error) { return "", ErrAIMLAPIKeyMissing } for _, r := range key { - if r <= ' ' || r == 0x7f || r == ',' || (r >= 0x80 && r <= 0x9f) { + if unicode.IsSpace(r) || r <= ' ' || r == 0x7f || r == ',' || (r >= 0x80 && r <= 0x9f) { return "", ErrAIMLAPIKeyInvalid } } diff --git a/internal/config/aimlapi_service_key_test.go b/internal/config/aimlapi_service_key_test.go index 1cd9a2535..cf00f047e 100644 --- a/internal/config/aimlapi_service_key_test.go +++ b/internal/config/aimlapi_service_key_test.go @@ -68,7 +68,7 @@ func TestAIMLAPIServiceKeyWithholdsProxyCredentials(t *testing.T) { } func TestAIMLAPIServiceKeyRejectsKeysThatAreNotOneToken(t *testing.T) { - for _, raw := range []string{"first,second", "first\nsecond", "first\r\nsecond", "a b", "key\x00tail", "key\u0085tail"} { + for _, raw := range []string{"first\u00a0second", "first\u2003second", "first,second", "first\nsecond", "first\r\nsecond", "a b", "key\x00tail", "key\u0085tail"} { resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKey: raw}} if key, err := AIMLAPIServiceKey(resolved, nil, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIKeyInvalid) || key != "" { t.Errorf("key %q: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIKeyInvalid", raw, key, err) From 3d7864c1cf65db3b97ddc5fca3716c709b443284 Mon Sep 17 00:00:00 2001 From: Hugo Date: Tue, 3 Nov 2026 14:43:49 +0300 Subject: [PATCH 7/8] fix(aimlapi): report store reads after proxy eligibility checks --- internal/config/aimlapi_service_key.go | 30 ++++++++++++++------- internal/config/aimlapi_service_key_test.go | 17 ++++++++++++ 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/internal/config/aimlapi_service_key.go b/internal/config/aimlapi_service_key.go index 219b8e8d4..b38f1dd99 100644 --- a/internal/config/aimlapi_service_key.go +++ b/internal/config/aimlapi_service_key.go @@ -51,8 +51,14 @@ func AIMLAPIServiceKey(resolved ResolvedConfig, store APIKeyGetter, getenv func( if !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) { return "", ErrAIMLAPIProxyProfile } - if store != nil { - profile = ApplyStoredAPIKey(profile, store) + if store != nil && profile.APIKeyStored && strings.TrimSpace(profile.APIKey) == "" && strings.TrimSpace(profile.Name) != "" { + key, found, err := store.Get(strings.TrimSpace(profile.Name)) + if err != nil { + return "", fmt.Errorf("read AI/ML API credential store: %w", err) + } + if found { + profile.APIKey = key + } } return validAIMLAPIServiceKey(profile.APIKey) } @@ -75,14 +81,18 @@ func LoadAIMLAPIServiceKey() (string, error) { if err != nil { return "", err } - var store APIKeyGetter - credentials, storeErr := ProviderKeyStoreAt(filepath.Dir(userConfigPath)) - if storeErr != nil { - if profile, ok := aimlapiServiceProfile(resolved); ok && profile.APIKeyStored && strings.TrimSpace(profile.APIKey) == "" { - return "", fmt.Errorf("open AI/ML API credential store: %w", storeErr) - } - } else if credentials != nil { - store = credentials + // Check proxy eligibility and already-resolved keys before opening a store. + key, keyErr := AIMLAPIServiceKey(resolved, nil, os.Getenv) + if !errors.Is(keyErr, ErrAIMLAPIKeyMissing) { + return key, keyErr + } + profile, ok := aimlapiServiceProfile(resolved) + if !ok || !profile.APIKeyStored { + return key, keyErr + } + store, err := ProviderKeyStoreAt(filepath.Dir(userConfigPath)) + if err != nil { + return "", fmt.Errorf("open AI/ML API credential store: %w", err) } return AIMLAPIServiceKey(resolved, store, os.Getenv) } diff --git a/internal/config/aimlapi_service_key_test.go b/internal/config/aimlapi_service_key_test.go index cf00f047e..e168540bd 100644 --- a/internal/config/aimlapi_service_key_test.go +++ b/internal/config/aimlapi_service_key_test.go @@ -143,6 +143,11 @@ func TestLoadAIMLAPIServiceKeyStoreOpenFailure(t *testing.T) { if key, err := LoadAIMLAPIServiceKey(); key != "" || err == nil || !strings.Contains(err.Error(), "unknown storage") { t.Fatalf("stored key must report store failure: %q, %v", key, err) } + t.Setenv("AIMLAPI_INFERENCE_URL", "https://proxy.example/v1") + if key, err := LoadAIMLAPIServiceKey(); key != "" || !errors.Is(err, ErrAIMLAPIProxyProfile) { + t.Fatalf("proxy rejection must precede store open: %q, %v", key, err) + } + t.Setenv("AIMLAPI_INFERENCE_URL", "") data = `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKeyStored":true,"apiKey":"inline-key"}]}` if err := os.WriteFile(path, []byte(data), 0o600); err != nil { t.Fatal(err) @@ -151,3 +156,15 @@ func TestLoadAIMLAPIServiceKeyStoreOpenFailure(t *testing.T) { t.Fatalf("inline key must not require a store: %q, %v", key, err) } } + +func TestAIMLAPIServiceKeyReportsStoreReadFailure(t *testing.T) { + want := errors.New("test credential read failure") + resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKeyStored: true}} + if key, err := AIMLAPIServiceKey(resolved, fakeKeyGetter{err: want}, aimlapiServiceKeyEnv(nil)); key != "" || !errors.Is(err, want) { + t.Fatalf("read failure was hidden: %q, %v", key, err) + } + resolved.Provider.BaseURL = "https://proxy.example/v1" + if _, err := AIMLAPIServiceKey(resolved, fakeKeyGetter{err: want}, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIProxyProfile) { + t.Fatalf("proxy check must precede store read: %v", err) + } +} From 4f99e49e944fcfc7f99f0d355db1fee09e7b3a6d Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 4 Nov 2026 14:48:44 -0800 Subject: [PATCH 8/8] fix(web-search): keep plain no-results for unscoped backends Only a backend that applied the domain allowlist itself (the AI/ML API Toolbox) should turn an empty result into the "no results matched domains" error. A generic backend never saw the allowlist, so an empty result from it is an ordinary "No results for query", as before this branch. Also clear ZERO_WEBSEARCH_PROVIDER in the unconfigured registry test: ZERO_WEBSEARCH_PROVIDER=aimlapi alone now registers web_search, so an exported value flipped the test's oracle. Co-Authored-By: Claude Opus 5.5 --- internal/tools/registry_test.go | 1 + internal/tools/web_search.go | 8 ++++++-- internal/tools/web_search_test.go | 14 ++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/internal/tools/registry_test.go b/internal/tools/registry_test.go index 6dbe94be7..c49935d8d 100644 --- a/internal/tools/registry_test.go +++ b/internal/tools/registry_test.go @@ -213,6 +213,7 @@ func TestCoreNetworkToolsOmitWebSearchWhenUnconfigured(t *testing.T) { // No backend configured → don't offer web_search (it could only error, which // makes the model waste calls + prompts before falling back to an MCP search). t.Setenv("ZERO_WEBSEARCH_BASE_URL", "") + t.Setenv("ZERO_WEBSEARCH_PROVIDER", "") names := map[string]bool{} for _, tool := range CoreNetworkTools() { names[tool.Name()] = true diff --git a/internal/tools/web_search.go b/internal/tools/web_search.go index e2ecd67f5..ffff1c190 100644 --- a/internal/tools/web_search.go +++ b/internal/tools/web_search.go @@ -156,7 +156,9 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result { defer cancel() var results []searchResult + scopedSearch := false if scoped, ok := tool.backend.(domainScopedSearchBackend); ok && len(domains) > 0 { + scopedSearch = true results, err = scoped.searchDomains(runCtx, query, limit, domains) } else { results, err = tool.backend.Search(runCtx, query, limit) @@ -166,8 +168,10 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result { } // Filter by domains BEFORE the empty-result short-circuit so a non-empty // allowlist that ate every hit surfaces as a clear "no results matched - // domains" error rather than a misleading "no results" message. - if len(domains) > 0 { + // domains" error rather than a misleading "no results" message. A scoped + // backend applied the allowlist itself, so its empty result gets the same + // error; an unscoped backend's empty result is an ordinary "no results". + if len(domains) > 0 && (len(results) > 0 || scopedSearch) { filtered, err := filterWebSearchByDomains(results, domains) if err != nil { return errorResult("Error: " + redactWebSearchText(err.Error())) diff --git a/internal/tools/web_search_test.go b/internal/tools/web_search_test.go index c0628f8dd..7310546fa 100644 --- a/internal/tools/web_search_test.go +++ b/internal/tools/web_search_test.go @@ -300,6 +300,20 @@ func TestWebSearchDomainsFilterEmptyResultIsError(t *testing.T) { } } +func TestWebSearchDomainsUnscopedEmptyResultIsNoResults(t *testing.T) { + // A backend that cannot scope by domain never saw the allowlist, so an empty + // result is an ordinary "no results", not an allowlist that ate every hit. + tool := newWebSearchToolWithBackend(&fakeSearchBackend{}) + + res := tool.Run(context.Background(), map[string]any{ + "query": "x", + "domains": []any{"react.dev"}, + }) + if res.Status != StatusOK || !strings.Contains(res.Output, "No results for query") { + t.Fatalf("want ordinary no-results for an unscoped backend, got %v: %s", res.Status, res.Output) + } +} + func TestWebSearchDomainsFilterToleratesStringSlice(t *testing.T) { // Some providers send []string instead of []any. The tool must accept both. backend := &fakeSearchBackend{results: []searchResult{