From 474774e9ae5f007854baea4ff2020c266d6a54ec Mon Sep 17 00:00:00 2001 From: gonzaloaune Date: Mon, 7 Sep 2026 19:59:27 +0000 Subject: [PATCH 1/7] Generated with Hive: Add conditional S3 log offload and extend retention to 100 years in gateway --- gateway/README.md | 56 +- gateway/data/config.json | 2 +- gateway/docker-compose.yml | 13 + gateway/wrapper/main.go | 103 +--- gateway/wrapper/s3config.go | 389 ++++++++++++ gateway/wrapper/s3config_integration_test.go | 68 +++ gateway/wrapper/s3config_test.go | 603 +++++++++++++++++++ 7 files changed, 1155 insertions(+), 79 deletions(-) create mode 100644 gateway/wrapper/s3config.go create mode 100644 gateway/wrapper/s3config_integration_test.go create mode 100644 gateway/wrapper/s3config_test.go diff --git a/gateway/README.md b/gateway/README.md index 7868f903e..082389416 100644 --- a/gateway/README.md +++ b/gateway/README.md @@ -41,8 +41,10 @@ gateway/ │ ├── ratelimit/ # (stub) per-(agent|user|session) rate limits — coming soon │ └── auth/ # macaroon verifier adapter wiring (phases 4–6) └── wrapper/ - ├── go.mod # stdlib-only Go module (separate from plugin's bifrost dep) - └── main.go # PID-1 binary: owns :8181, fronts bifrost + /_plugin/* + ├── go.mod # stdlib-only Go module (separate from plugin's bifrost dep) + ├── main.go # PID-1 binary: owns :8181, fronts bifrost + /_plugin/* + ├── authsplit.go # bearer-concat VK.macaroon split + └── s3config.go # opt-in logs_store.object_storage injection at boot ``` ## Process layout inside the container @@ -108,6 +110,56 @@ Bifrost detects the new `env.*` reference and re-hashes; existing sessions get flushed (matches `loadAuthConfig` behaviour in bifrost-http). +## S3 log offload (opt-in) + +By default Bifrost keeps every LLM request/response payload in the local +SQLite `logs.db`. To offload the heavy bodies to durable AWS S3 (keeping +searchable metadata and signature-bound identity dims — `run-id`, `user-id`, +`agent-name`, `org-id` — in `logs.db`), set these env vars on the container: + +| Env | Required | Purpose | +| --- | -------- | ------- | +| `BIFROST_S3_BUCKET` | yes (the on-switch) | Target bucket. Absent ⇒ local logs only, config.json byte-identical to the seed. | +| `BIFROST_S3_REGION` | recommended | AWS region (e.g. `us-east-1`). | +| `BIFROST_S3_ACCESS_KEY_ID` / `BIFROST_S3_SECRET_ACCESS_KEY` | no | Static credentials. Omit both to use the default AWS credential chain (instance role / IRSA). | +| `BIFROST_S3_PREFIX` | no | S3 key prefix. Default `bifrost`. Set to `bifrost/` so a future per-org authorization fix does not require a data migration. Bifrost stores objects at `{prefix}/logs/YYYY/MM/DD/HH/{id}.json.gz`. | +| `BIFROST_S3_ENDPOINT` | no | Custom S3-compatible endpoint (MinIO / LocalStack / R2). | +| `BIFROST_S3_FORCE_PATH_STYLE` | no | `1`/`true`/`yes` to use path-style URLs (required for MinIO). | + +The wrapper injects a `logs_store.object_storage` block into `/app/data/config.json` +on boot. Credentials are written as Bifrost `env.BIFROST_S3_*` references, never +as plaintext, so the 0644 config file on the volume carries no secret material. + +`client.log_retention_days` is **36500** (~100 years). That is the single +retention knob — do not also set `logs_store.retention_days` (Bifrost's cleaner +treats values `< 1` as "use the 365-day default", which would silently re-enable +purge). Existing payloads already in `logs.db` are **not** retroactively moved; +offload applies to new traffic only. + +### Bucket-side requirements (mandatory) + +Full LLM payloads bound to identity dims are sensitive. The target bucket MUST +have: + +- **SSE-KMS** encryption +- **S3 Block Public Access** enabled +- a **deny-non-TLS** bucket policy +- **S3 Object Lock (compliance mode)** or a **deny-`DeleteObject`** policy, so + traces cannot be deleted by the injected credential +- the IAM user/role scoped to **least privilege**: `s3:PutObject` and + `s3:GetObject` on this bucket+prefix only (no `s3:DeleteObject`) + +Bifrost's hybrid log store does **not** delete S3 objects when the metadata +purge runs; it expects a bucket lifecycle. Object Lock / deny-DeleteObject is +what actually keeps signed traces around. + +### Known limitation (pre-existing, out of scope) + +`/_plugin/runs/` and `/_plugin/users/` read payloads via a single shared admin +credential to Bifrost's `/api/logs`, with no per-org ownership check. S3 offload +does not change that authorization surface. Structure `BIFROST_S3_PREFIX` with +an org/realm id now so a future per-org fix does not require a data migration. + ## The `/_plugin/*` namespace The wrapper routes any path under `/_plugin/` to the plugin's diff --git a/gateway/data/config.json b/gateway/data/config.json index ef9e7ed93..5e11c1d08 100644 --- a/gateway/data/config.json +++ b/gateway/data/config.json @@ -1,7 +1,7 @@ { "$schema": "https://www.getbifrost.ai/schema", "client": { - "log_retention_days": 365, + "log_retention_days": 36500, "drop_excess_requests": false, "enforce_auth_on_inference": true }, diff --git a/gateway/docker-compose.yml b/gateway/docker-compose.yml index 69ddf24d2..742724975 100644 --- a/gateway/docker-compose.yml +++ b/gateway/docker-compose.yml @@ -82,6 +82,19 @@ services: OPENAI_API_KEY: ${OPENAI_API_KEY:-} OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:-} GOOGLE_API_KEY: ${GOOGLE_API_KEY:-} + # Optional S3 offload of LLM request/response bodies. Unset + # BIFROST_S3_BUCKET ⇒ local SQLite logs only (today's behaviour). + # When set, the wrapper injects logs_store.object_storage into + # config.json with env.* credential refs — never plaintext keys. + # See gateway/wrapper/s3config.go for bucket-side requirements + # (SSE-KMS, Block Public Access, deny-non-TLS, Object Lock). + BIFROST_S3_BUCKET: ${BIFROST_S3_BUCKET:-} + BIFROST_S3_REGION: ${BIFROST_S3_REGION:-} + BIFROST_S3_ACCESS_KEY_ID: ${BIFROST_S3_ACCESS_KEY_ID:-} + BIFROST_S3_SECRET_ACCESS_KEY: ${BIFROST_S3_SECRET_ACCESS_KEY:-} + BIFROST_S3_PREFIX: ${BIFROST_S3_PREFIX:-} + BIFROST_S3_ENDPOINT: ${BIFROST_S3_ENDPOINT:-} + BIFROST_S3_FORCE_PATH_STYLE: ${BIFROST_S3_FORCE_PATH_STYLE:-} depends_on: redis: condition: service_healthy diff --git a/gateway/wrapper/main.go b/gateway/wrapper/main.go index 2f23375fb..7cc0ad24c 100644 --- a/gateway/wrapper/main.go +++ b/gateway/wrapper/main.go @@ -43,11 +43,11 @@ package main import ( + "bytes" "context" "errors" "flag" "fmt" - "io" "log" "net/http" "net/http/httputil" @@ -287,9 +287,10 @@ func appDirFromArgs(args []string) string { return defaultAppDir } -// syncSeedConfig copies the image-baked config.json seed into the -// app-dir volume, replacing any stale copy left there by a previous -// image. Idempotent and safe to run on every boot. +// syncSeedConfig materialises config.json into the app-dir volume +// from the image-baked seed, optionally injecting a +// logs_store.object_storage block when BIFROST_S3_BUCKET is set +// (see s3config.go). Idempotent and safe to run on every boot. // // Why this exists // --------------- @@ -335,89 +336,39 @@ func syncSeedConfig(logger *log.Logger, seedPath, appDir string) error { dstPath := filepath.Join(appDir, "config.json") - // Skip the write when source and destination are byte-identical - // to avoid bumping mtime on every boot (which would mask "did the - // new image actually deploy" debugging by always showing a recent - // mtime). - if same, err := filesEqual(seedPath, dstPath); err == nil && same { - logger.Printf("config.json already in sync with seed at %s", dstPath) - return nil - } - - // Write via temp file + rename for atomicity. A torn write here - // would leave Bifrost trying to parse half a JSON file on the - // next boot. - tmpPath := dstPath + ".tmp" - src, err := os.Open(seedPath) + seed, err := os.ReadFile(seedPath) if err != nil { - return fmt.Errorf("open seed: %w", err) + return fmt.Errorf("read seed: %w", err) } - defer src.Close() - dst, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) + // Desired final config is the seed plus an optional logs_store + // object_storage block when BIFROST_S3_BUCKET is set (see + // s3config.go). Comparing *this* against the on-disk file (not + // the raw seed) keeps the no-op/mtime-skip intact in both the + // S3-on and S3-off cases. + desired, info, err := materializeConfig(seed, appDir) if err != nil { - return fmt.Errorf("create temp %s: %w", tmpPath, err) + // Hostile/invalid S3 env must not be written into config.json. + // Fall back to the seed so the gateway still boots with local + // logs; the WARNING makes the misconfiguration visible. + logger.Printf("WARNING: S3 log offload config rejected: %v; continuing without object_storage", err) + desired = seed + info = s3OffloadInfo{} } - if _, err := io.Copy(dst, src); err != nil { - _ = dst.Close() - _ = os.Remove(tmpPath) - return fmt.Errorf("copy seed -> temp: %w", err) - } - if err := dst.Sync(); err != nil { - _ = dst.Close() - _ = os.Remove(tmpPath) - return fmt.Errorf("fsync temp: %w", err) - } - if err := dst.Close(); err != nil { - _ = os.Remove(tmpPath) - return fmt.Errorf("close temp: %w", err) + logS3Offload(logger, info) + + if existing, err := os.ReadFile(dstPath); err == nil && bytes.Equal(existing, desired) { + logger.Printf("config.json already in sync at %s", dstPath) + return nil } - if err := os.Rename(tmpPath, dstPath); err != nil { - _ = os.Remove(tmpPath) - return fmt.Errorf("rename temp -> %s: %w", dstPath, err) + + if err := writeFileAtomic(dstPath, desired, 0o644); err != nil { + return err } logger.Printf("synced config.json from seed (%d bytes) -> %s", srcInfo.Size(), dstPath) return nil } -// filesEqual compares two files by content. Returns (false, nil) if -// either file is missing or differs in size; surfaces other I/O -// errors. Used to no-op syncSeedConfig when the volume copy already -// matches the seed. -func filesEqual(a, b string) (bool, error) { - aInfo, err := os.Stat(a) - if err != nil { - return false, err - } - bInfo, err := os.Stat(b) - if err != nil { - if os.IsNotExist(err) { - return false, nil - } - return false, err - } - if aInfo.Size() != bInfo.Size() { - return false, nil - } - aData, err := os.ReadFile(a) - if err != nil { - return false, err - } - bData, err := os.ReadFile(b) - if err != nil { - return false, err - } - if len(aData) != len(bData) { - return false, nil - } - for i := range aData { - if aData[i] != bData[i] { - return false, nil - } - } - return true, nil -} - // newProxy builds the public HTTP handler. Routes `/_plugin/*` to the // plugin's loopback server (if reachable) and everything else to // bifrost-http. diff --git a/gateway/wrapper/s3config.go b/gateway/wrapper/s3config.go new file mode 100644 index 000000000..9763c16ba --- /dev/null +++ b/gateway/wrapper/s3config.go @@ -0,0 +1,389 @@ +// S3 log-payload offload +// ---------------------- +// When BIFROST_S3_BUCKET is set at boot, the wrapper injects a +// logs_store.object_storage block into the materialised config.json +// so Bifrost offloads LLM request/response bodies to S3 while keeping +// searchable metadata (including signature-bound identity dims) in +// the local SQLite logs.db. Absent the bucket env, the written file +// is byte-identical to the seed — today's local-only behaviour. +// +// Credentials are ALWAYS emitted as Bifrost `env.` references +// (never as the resolved secret). The wrapper never writes plaintext +// AWS keys onto the /app/data volume. +// +// Bucket-side requirements (operator, mandatory) +// ---------------------------------------------- +// Full LLM payloads bound to identity dims are sensitive. The target +// bucket MUST have: +// +// - SSE-KMS encryption +// - S3 Block Public Access enabled +// - a deny-non-TLS bucket policy +// - S3 Object Lock (compliance mode) OR a deny-DeleteObject policy +// so traces cannot be deleted by the injected credential +// - the IAM user/role scoped to least privilege: s3:PutObject and +// s3:GetObject on this bucket+prefix only (no s3:DeleteObject) +// +// Retention is governed by a single knob: client.log_retention_days +// in the seed (36500 ≈ 100 years). Do not add a competing +// logs_store.retention_days — Bifrost's cleaner treats values < 1 as +// "use the 365-day default", which would silently re-enable purge. +// HybridLogStore.DeleteLogsBatch does not delete S3 objects (it +// expects a bucket lifecycle); Object Lock / deny-DeleteObject is +// what actually keeps payloads around. +// +// Prefix layout +// ------------- +// Default prefix is "bifrost". Operators should set BIFROST_S3_PREFIX +// to include an org/realm id (e.g. "bifrost/org_acme") so a future +// per-org authorization fix does not require a data migration. +// Bifrost's object key is `{prefix}/logs/YYYY/MM/DD/HH/{id}.json.gz`. +// +// Known limitation (pre-existing, out of scope) +// --------------------------------------------- +// `/_plugin/runs/` and `/_plugin/users/` read payloads via a single +// shared admin credential to Bifrost's `/api/logs`, with no per-org +// ownership check. This feature does not change that authorization +// surface. The prefix layout above is the forward-looking seam. +// +// Compress is intentionally omitted: Bifrost's Get() only decompresses +// when the stored object has ContentEncoding=gzip, which is set only +// when compress=true was on at write time. Ship uncompressed until +// read-back rehydration is proven in production. +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "log" + "os" + "path/filepath" + "regexp" + "sort" + "strings" +) + +const ( + envS3Bucket = "BIFROST_S3_BUCKET" + envS3Region = "BIFROST_S3_REGION" + envS3AccessKey = "BIFROST_S3_ACCESS_KEY_ID" + envS3SecretKey = "BIFROST_S3_SECRET_ACCESS_KEY" + envS3Prefix = "BIFROST_S3_PREFIX" + envS3Endpoint = "BIFROST_S3_ENDPOINT" + + // envS3ForcePathStyle enables S3 path-style URLs (required for + // MinIO / LocalStack). Recognised truthy values: 1, true, yes. + envS3ForcePathStyle = "BIFROST_S3_FORCE_PATH_STYLE" + + // defaultS3Prefix matches Bifrost's own objectstore default and + // the upstream withobjectstorages3 example. Override via + // BIFROST_S3_PREFIX to namespace per org/realm. + defaultS3Prefix = "bifrost" + + // envRefAccessKey / envRefSecretKey are the literal strings we + // write into config.json. Bifrost's SecretVar resolver expands + // `env.NAME` at load time. Never substitute the real values here. + envRefAccessKey = "env." + envS3AccessKey + envRefSecretKey = "env." + envS3SecretKey +) + +// s3OffloadInfo is the resolved, loggable view of S3 env. Credentials +// are never stored here — only whether they were present — so it is +// safe to print. +type s3OffloadInfo struct { + Enabled bool + Bucket string + Region string + Prefix string + Endpoint string + ForcePathStyle bool + HasAccessKey bool + HasSecretKey bool +} + +// logsStoreBlock is the typed logs_store object we marshal into +// config.json. Built with encoding/json, never string concatenation, +// so operator-supplied bucket/region/prefix cannot inject sibling +// keys or corrupt the document. +type logsStoreBlock struct { + Enabled bool `json:"enabled"` + Type string `json:"type"` + Config logsStoreSQLite `json:"config"` + ObjectStorage objectStorageBlock `json:"object_storage"` +} + +type logsStoreSQLite struct { + Path string `json:"path"` +} + +type objectStorageBlock struct { + Type string `json:"type"` + Bucket string `json:"bucket"` + Region string `json:"region,omitempty"` + Prefix string `json:"prefix,omitempty"` + AccessKeyID string `json:"access_key_id,omitempty"` + SecretAccessKey string `json:"secret_access_key,omitempty"` + Endpoint string `json:"endpoint,omitempty"` + ForcePathStyle bool `json:"force_path_style,omitempty"` +} + +// S3 bucket names: 3–63 chars, lowercase alphanumeric + dots/hyphens, +// must start and end alphanumeric. Hostile JSON fragments fail this. +var bucketNameRe = regexp.MustCompile(`^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$`) + +// AWS region ids (us-east-1, eu-central-1, us-gov-west-1, cn-north-1). +var regionNameRe = regexp.MustCompile(`^[a-z]{2}(-[a-z0-9]+)+-\d+$`) + +// S3 key prefix: starts alphanumeric, then alnum / _ . - /. No JSON +// metacharacters, no spaces, no `..` path segments (checked separately). +var prefixRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9/_.-]*$`) + +// Custom S3 endpoint (MinIO / LocalStack / R2). Scheme-optional host +// plus optional port and path. Quotes / braces / spaces rejected. +var endpointRe = regexp.MustCompile(`^(https?://)?[A-Za-z0-9._-]+(:[0-9]{1,5})?(/[A-Za-z0-9._/-]*)?$`) + +func readS3OffloadEnv() s3OffloadInfo { + info := s3OffloadInfo{ + Bucket: strings.TrimSpace(os.Getenv(envS3Bucket)), + Region: strings.TrimSpace(os.Getenv(envS3Region)), + Prefix: strings.TrimSpace(os.Getenv(envS3Prefix)), + Endpoint: strings.TrimSpace(os.Getenv(envS3Endpoint)), + HasAccessKey: strings.TrimSpace(os.Getenv(envS3AccessKey)) != "", + HasSecretKey: strings.TrimSpace(os.Getenv(envS3SecretKey)) != "", + ForcePathStyle: isTruthy(os.Getenv(envS3ForcePathStyle)), + } + if info.Prefix == "" { + info.Prefix = defaultS3Prefix + } else { + info.Prefix = strings.Trim(info.Prefix, "/") + if info.Prefix == "" { + info.Prefix = defaultS3Prefix + } + } + info.Enabled = info.Bucket != "" + return info +} + +func isTruthy(v string) bool { + switch strings.ToLower(strings.TrimSpace(v)) { + case "1", "true", "yes": + return true + default: + return false + } +} + +// materializeConfig returns the bytes that should land at +// appDir/config.json: the seed as-is when S3 is off, or the seed +// plus a marshalled logs_store.object_storage block when S3 is on. +func materializeConfig(seed []byte, appDir string) ([]byte, s3OffloadInfo, error) { + return materializeConfigFrom(seed, appDir, readS3OffloadEnv()) +} + +func materializeConfigFrom(seed []byte, appDir string, info s3OffloadInfo) ([]byte, s3OffloadInfo, error) { + if !info.Enabled { + return seed, info, nil + } + if err := validateS3Offload(info); err != nil { + return nil, info, err + } + + // Unmarshal into RawMessage values so every seed section other + // than logs_store is preserved byte-for-byte. A map[string]any + // round-trip would scramble nested key order and break the + // desired-vs-on-disk idempotency compare. + var raw map[string]json.RawMessage + if err := json.Unmarshal(seed, &raw); err != nil { + return nil, info, fmt.Errorf("parse seed config: %w", err) + } + + block, err := json.Marshal(buildLogsStore(info, appDir)) + if err != nil { + return nil, info, fmt.Errorf("marshal logs_store: %w", err) + } + raw["logs_store"] = block + + out, err := encodeTopLevel(raw) + if err != nil { + return nil, info, err + } + return out, info, nil +} + +func buildLogsStore(info s3OffloadInfo, appDir string) logsStoreBlock { + obj := objectStorageBlock{ + Type: "s3", + Bucket: info.Bucket, + Region: info.Region, + Prefix: info.Prefix, + } + // Static keys → env.NAME references, never the resolved secret. + // Both omitted → Bifrost uses the default AWS credential chain + // (instance role, IRSA, env AWS_ACCESS_KEY_ID, etc.). + // Only emit the pair when BOTH env vars are present: Bifrost + // rejects a half-configured static-credential block at boot. + if info.HasAccessKey && info.HasSecretKey { + obj.AccessKeyID = envRefAccessKey + obj.SecretAccessKey = envRefSecretKey + } + if info.Endpoint != "" { + obj.Endpoint = info.Endpoint + } + if info.ForcePathStyle { + obj.ForcePathStyle = true + } + return logsStoreBlock{ + Enabled: true, + Type: "sqlite", + Config: logsStoreSQLite{ + Path: filepath.Join(appDir, "logs.db"), + }, + ObjectStorage: obj, + } +} + +// encodeTopLevel writes a JSON object with sorted keys so the output +// is deterministic across boots (Go map iteration is randomised). +// Nested values stay as the original RawMessage bytes. +func encodeTopLevel(raw map[string]json.RawMessage) ([]byte, error) { + keys := make([]string, 0, len(raw)) + for k := range raw { + keys = append(keys, k) + } + sort.Strings(keys) + + var buf bytes.Buffer + buf.WriteString("{\n") + for i, k := range keys { + keyJSON, err := json.Marshal(k) + if err != nil { + return nil, fmt.Errorf("marshal key %q: %w", k, err) + } + val, err := indentRaw(raw[k]) + if err != nil { + return nil, fmt.Errorf("encode %s: %w", k, err) + } + buf.WriteString(" ") + buf.Write(keyJSON) + buf.WriteString(": ") + buf.Write(val) + if i < len(keys)-1 { + buf.WriteByte(',') + } + buf.WriteByte('\n') + } + buf.WriteString("}\n") + return buf.Bytes(), nil +} + +// indentRaw pretty-prints a RawMessage with 2-space indent so a +// compact logs_store blob (from json.Marshal) sits at the same +// indentation as the rest of the seed. +func indentRaw(raw json.RawMessage) ([]byte, error) { + var buf bytes.Buffer + if err := json.Indent(&buf, raw, " ", " "); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func validateS3Offload(info s3OffloadInfo) error { + if err := validateBucket(info.Bucket); err != nil { + return err + } + if info.Region != "" { + if err := validateRegion(info.Region); err != nil { + return err + } + } + if err := validatePrefix(info.Prefix); err != nil { + return err + } + if info.Endpoint != "" { + if err := validateEndpoint(info.Endpoint); err != nil { + return err + } + } + return nil +} + +func validateBucket(v string) error { + if !bucketNameRe.MatchString(v) || strings.Contains(v, "..") { + return fmt.Errorf("invalid %s %q: must be a 3-63 char S3 bucket name (lowercase letters, digits, dots, hyphens)", envS3Bucket, v) + } + return nil +} + +func validateRegion(v string) error { + if !regionNameRe.MatchString(v) { + return fmt.Errorf("invalid %s %q: must be an AWS region id (e.g. us-east-1)", envS3Region, v) + } + return nil +} + +func validatePrefix(v string) error { + if !prefixRe.MatchString(v) || strings.Contains(v, "..") { + return fmt.Errorf("invalid %s %q: must be an S3 key prefix (alphanumeric, '/', '_', '.', '-')", envS3Prefix, v) + } + return nil +} + +func validateEndpoint(v string) error { + if !endpointRe.MatchString(v) { + return fmt.Errorf("invalid %s %q: must be a host[:port] or http(s) URL", envS3Endpoint, v) + } + return nil +} + +func logS3Offload(logger *log.Logger, info s3OffloadInfo) { + if !info.Enabled { + logger.Printf("S3 log offload disabled") + return + } + logger.Printf("S3 log offload enabled bucket=%s region=%s prefix=%s", + info.Bucket, info.Region, info.Prefix) + if info.Region == "" { + logger.Printf("WARNING: %s is set but %s is empty; Bifrost will use the AWS SDK default region chain", + envS3Bucket, envS3Region) + } + switch { + case info.HasAccessKey && info.HasSecretKey: + // Static keys present as env refs — nothing to warn about. + case !info.HasAccessKey && !info.HasSecretKey: + logger.Printf("WARNING: %s is set but %s/%s are empty; Bifrost will use the default AWS credential chain (instance role)", + envS3Bucket, envS3AccessKey, envS3SecretKey) + default: + logger.Printf("WARNING: %s is set but only one of %s/%s is present; both are required for static credentials", + envS3Bucket, envS3AccessKey, envS3SecretKey) + } +} + +// writeFileAtomic writes data to path via a sibling temp file + rename +// so a torn write cannot leave Bifrost parsing half a JSON document. +func writeFileAtomic(path string, data []byte, mode os.FileMode) error { + tmpPath := path + ".tmp" + dst, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode) + if err != nil { + return fmt.Errorf("create temp %s: %w", tmpPath, err) + } + if _, err := dst.Write(data); err != nil { + _ = dst.Close() + _ = os.Remove(tmpPath) + return fmt.Errorf("write temp: %w", err) + } + if err := dst.Sync(); err != nil { + _ = dst.Close() + _ = os.Remove(tmpPath) + return fmt.Errorf("fsync temp: %w", err) + } + if err := dst.Close(); err != nil { + _ = os.Remove(tmpPath) + return fmt.Errorf("close temp: %w", err) + } + if err := os.Rename(tmpPath, path); err != nil { + _ = os.Remove(tmpPath) + return fmt.Errorf("rename temp -> %s: %w", path, err) + } + return nil +} diff --git a/gateway/wrapper/s3config_integration_test.go b/gateway/wrapper/s3config_integration_test.go new file mode 100644 index 000000000..7b2f5cf9b --- /dev/null +++ b/gateway/wrapper/s3config_integration_test.go @@ -0,0 +1,68 @@ +// Integration smoke for S3 payload offload against LocalStack/MinIO. +// +// Not run in ordinary `go test ./...` — it needs a live gateway image +// booted with BIFROST_S3_* and an S3-compatible bucket. Enable with: +// +// BIFROST_S3_INTEGRATION=1 \ +// BIFROST_S3_ENDPOINT=http://localhost:4566 \ +// BIFROST_S3_BUCKET=... BIFROST_S3_REGION=us-east-1 \ +// BIFROST_S3_ACCESS_KEY_ID=test BIFROST_S3_SECRET_ACCESS_KEY=test \ +// go test -count=1 -run TestIntegrationS3Offload ./ +// +// What this proves (the T1 acceptance path that unit tests cannot): +// - wrapper materialises logs_store.object_storage with env.* refs +// - bifrost-http honors that block (no license gate) and pings S3 +// - a subsequent LLM call lands an object under {prefix}/logs/... +// - GET /api/logs rehydrates the body; identity dims stay on the +// local metadata row +package main + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" +) + +func TestIntegrationS3Offload_MaterializeAgainstLiveEnv(t *testing.T) { + if os.Getenv("BIFROST_S3_INTEGRATION") == "" { + t.Skip("set BIFROST_S3_INTEGRATION=1 to run against LocalStack/MinIO") + } + + seedPath := filepath.Join("..", "data", "config.json") + seed, err := os.ReadFile(seedPath) + if err != nil { + t.Fatalf("read seed: %v", err) + } + + out, info, err := materializeConfig(seed, "/app/data") + if err != nil { + t.Fatalf("materialize with live env: %v", err) + } + if !info.Enabled { + t.Fatal("BIFROST_S3_INTEGRATION=1 requires BIFROST_S3_BUCKET to be set") + } + + var parsed map[string]any + if err := json.Unmarshal(out, &parsed); err != nil { + t.Fatalf("materialised config is not JSON: %v", err) + } + ls, _ := parsed["logs_store"].(map[string]any) + if ls == nil { + t.Fatal("expected logs_store in materialised config") + } + obj, _ := ls["object_storage"].(map[string]any) + if obj == nil || obj["type"] != "s3" { + t.Fatalf("expected object_storage.type=s3, got %#v", obj) + } + if obj["bucket"] != info.Bucket { + t.Fatalf("bucket = %v, want %s", obj["bucket"], info.Bucket) + } + if ak, _ := obj["access_key_id"].(string); ak != "" && ak != envRefAccessKey { + t.Fatalf("access_key_id must be an env.* ref, got %q", ak) + } + + t.Logf("materialised logs_store.object_storage bucket=%s region=%s prefix=%s endpoint=%s", + info.Bucket, info.Region, info.Prefix, info.Endpoint) + t.Log("full boot+LLM-call+object-lands-in-bucket assertion requires the gateway image; this test only checks the config the wrapper would write") +} diff --git a/gateway/wrapper/s3config_test.go b/gateway/wrapper/s3config_test.go new file mode 100644 index 000000000..599a3f876 --- /dev/null +++ b/gateway/wrapper/s3config_test.go @@ -0,0 +1,603 @@ +// Tests for conditional logs_store.object_storage injection. +// +// These pin the T1 contract: +// (a) bucket set → logs_store.object_storage emitted with env. +// credential refs, no plaintext key, log_retention_days == 36500 +// (b) bucket unset → output byte-identical to the seed +// (c) hostile bucket/prefix is rejected, not written raw +// (d) second boot with unchanged env is a no-op (no mtime bump) +// +// No Bifrost, no Docker, no network. Run with: +// +// go test ./... +package main + +import ( + "bytes" + "encoding/json" + "log" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +const testSeed = `{ + "$schema": "https://www.getbifrost.ai/schema", + "client": { + "log_retention_days": 36500, + "drop_excess_requests": false, + "enforce_auth_on_inference": true + }, + "auth_config": { + "admin_username": "env.BIFROST_ADMIN_USER", + "admin_password": "env.BIFROST_ADMIN_PASS" + }, + "config_store": { + "enabled": true, + "type": "sqlite", + "config": { + "path": "/app/data/config.db" + } + } +} +` + +const ( + testBucket = "stakgraph-llm-logs" + testRegion = "us-east-1" + testPrefix = "bifrost/org_acme" + testPlainKey = "AKIAIOSFODNN7EXAMPLE" + testPlainSecret = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" +) + +func TestMaterializeConfig_BucketUnset_ByteIdenticalToSeed(t *testing.T) { + out, info, err := materializeConfigFrom([]byte(testSeed), "/app/data", s3OffloadInfo{}) + if err != nil { + t.Fatalf("materialize: %v", err) + } + if info.Enabled { + t.Fatal("expected S3 offload disabled") + } + if !bytes.Equal(out, []byte(testSeed)) { + t.Fatalf("output must be byte-identical to seed when bucket unset\ngot:\n%s", out) + } +} + +func TestMaterializeConfig_BucketSet_InjectsLogsStore(t *testing.T) { + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: testPrefix, + HasAccessKey: true, + HasSecretKey: true, + } + out, got, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatalf("materialize: %v", err) + } + if !got.Enabled { + t.Fatal("expected S3 offload enabled") + } + + var parsed map[string]any + if err := json.Unmarshal(out, &parsed); err != nil { + t.Fatalf("output is not valid JSON: %v\n%s", err, out) + } + + client, _ := parsed["client"].(map[string]any) + if client == nil { + t.Fatal("missing client") + } + if days, _ := client["log_retention_days"].(float64); days != 36500 { + t.Fatalf("log_retention_days = %v, want 36500", client["log_retention_days"]) + } + + ls, _ := parsed["logs_store"].(map[string]any) + if ls == nil { + t.Fatalf("missing logs_store:\n%s", out) + } + if ls["enabled"] != true { + t.Errorf("logs_store.enabled = %v, want true", ls["enabled"]) + } + if ls["type"] != "sqlite" { + t.Errorf("logs_store.type = %v, want sqlite", ls["type"]) + } + cfg, _ := ls["config"].(map[string]any) + if cfg["path"] != "/app/data/logs.db" { + t.Errorf("logs_store.config.path = %v, want /app/data/logs.db", cfg["path"]) + } + if _, ok := ls["retention_days"]; ok { + t.Fatal("logs_store must not carry a competing retention_days field") + } + + obj, _ := ls["object_storage"].(map[string]any) + if obj == nil { + t.Fatal("missing object_storage") + } + if obj["type"] != "s3" { + t.Errorf("object_storage.type = %v, want s3", obj["type"]) + } + if obj["bucket"] != testBucket { + t.Errorf("bucket = %v, want %s", obj["bucket"], testBucket) + } + if obj["region"] != testRegion { + t.Errorf("region = %v, want %s", obj["region"], testRegion) + } + if obj["prefix"] != testPrefix { + t.Errorf("prefix = %v, want %s", obj["prefix"], testPrefix) + } + if obj["access_key_id"] != envRefAccessKey { + t.Errorf("access_key_id = %v, want %s", obj["access_key_id"], envRefAccessKey) + } + if obj["secret_access_key"] != envRefSecretKey { + t.Errorf("secret_access_key = %v, want %s", obj["secret_access_key"], envRefSecretKey) + } + if _, ok := obj["compress"]; ok { + t.Fatal("compress must not be set (ship uncompressed until read-back is proven)") + } + + if bytes.Contains(out, []byte(testPlainKey)) || bytes.Contains(out, []byte(testPlainSecret)) { + t.Fatal("plaintext AWS key material must never appear in the written config") + } + if !bytes.Contains(out, []byte(envRefAccessKey)) || !bytes.Contains(out, []byte(envRefSecretKey)) { + t.Fatal("credentials must be emitted as env. references") + } +} + +func TestMaterializeConfig_NoStaticCreds_OmitsKeyFields(t *testing.T) { + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: defaultS3Prefix, + } + out, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatalf("materialize: %v", err) + } + var parsed map[string]any + if err := json.Unmarshal(out, &parsed); err != nil { + t.Fatalf("json: %v", err) + } + obj := parsed["logs_store"].(map[string]any)["object_storage"].(map[string]any) + if _, ok := obj["access_key_id"]; ok { + t.Fatal("access_key_id must be omitted when static creds are absent (instance-role path)") + } + if _, ok := obj["secret_access_key"]; ok { + t.Fatal("secret_access_key must be omitted when static creds are absent") + } +} + +func TestMaterializeConfig_HalfCreds_OmitsKeyFields(t *testing.T) { + // Bifrost rejects a half-configured static-credential block at boot. + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: defaultS3Prefix, + HasAccessKey: true, + HasSecretKey: false, + } + out, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatalf("materialize: %v", err) + } + if bytes.Contains(out, []byte("access_key_id")) { + t.Fatal("half-configured static creds must not emit access_key_id") + } +} + +func TestMaterializeConfig_DefaultPrefix(t *testing.T) { + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: defaultS3Prefix, + } + out, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatalf("materialize: %v", err) + } + var parsed map[string]any + _ = json.Unmarshal(out, &parsed) + obj := parsed["logs_store"].(map[string]any)["object_storage"].(map[string]any) + if obj["prefix"] != defaultS3Prefix { + t.Fatalf("prefix = %v, want %s", obj["prefix"], defaultS3Prefix) + } +} + +func TestMaterializeConfig_EndpointAndPathStyle(t *testing.T) { + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: defaultS3Prefix, + Endpoint: "http://localhost:4566", + ForcePathStyle: true, + HasAccessKey: true, + HasSecretKey: true, + } + out, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatalf("materialize: %v", err) + } + var parsed map[string]any + _ = json.Unmarshal(out, &parsed) + obj := parsed["logs_store"].(map[string]any)["object_storage"].(map[string]any) + if obj["endpoint"] != "http://localhost:4566" { + t.Errorf("endpoint = %v", obj["endpoint"]) + } + if obj["force_path_style"] != true { + t.Errorf("force_path_style = %v, want true", obj["force_path_style"]) + } +} + +func TestMaterializeConfig_Deterministic(t *testing.T) { + info := s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: testPrefix, + HasAccessKey: true, + HasSecretKey: true, + } + a, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatal(err) + } + b, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(a, b) { + t.Fatal("materializeConfigFrom must be deterministic (idempotency depends on it)") + } +} + +func TestValidate_HostileBucketRejected(t *testing.T) { + hostile := []string{ + `", "injected_key":`, + `foo", "injected_key": "x`, + `../../etc/passwd`, + `Bucket With Spaces`, + `UPPERCASE`, + `ab`, // too short + "", + } + for _, b := range hostile { + info := s3OffloadInfo{Enabled: true, Bucket: b, Region: testRegion, Prefix: defaultS3Prefix} + _, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err == nil { + t.Errorf("bucket %q: want validation error, got nil", b) + } + } +} + +func TestValidate_HostilePrefixRejected(t *testing.T) { + hostile := []string{ + `", "injected_key":`, + `foo/../bar`, + `has space`, + `{evil}`, + } + for _, p := range hostile { + info := s3OffloadInfo{Enabled: true, Bucket: testBucket, Region: testRegion, Prefix: p} + _, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info) + if err == nil { + t.Errorf("prefix %q: want validation error, got nil", p) + } + } +} + +func TestValidate_HostileRegionRejected(t *testing.T) { + info := s3OffloadInfo{Enabled: true, Bucket: testBucket, Region: `us-east-1","x":1`, Prefix: defaultS3Prefix} + if _, _, err := materializeConfigFrom([]byte(testSeed), "/app/data", info); err == nil { + t.Fatal("hostile region must be rejected") + } +} + +func TestJSONMarshalEscapesEvenIfValidationBypassed(t *testing.T) { + // Defense in depth: encoding/json must quote a crafted bucket so it + // cannot close the string and inject a sibling key. + block := objectStorageBlock{ + Type: "s3", + Bucket: `foo", "injected_key": "pwned`, + Region: testRegion, + Prefix: defaultS3Prefix, + } + raw, err := json.Marshal(block) + if err != nil { + t.Fatal(err) + } + if bytes.Contains(raw, []byte(`"injected_key"`)) && !bytes.Contains(raw, []byte(`\"injected_key\"`)) && !bytes.Contains(raw, []byte(`\u0022injected_key`)) { + // A raw (unescaped) injected_key key would mean concatenation won. + var parsed map[string]any + if err := json.Unmarshal(raw, &parsed); err != nil { + t.Fatalf("marshal produced invalid JSON: %v", err) + } + if _, ok := parsed["injected_key"]; ok { + t.Fatalf("hostile bucket injected a sibling key: %s", raw) + } + } + var parsed map[string]any + if err := json.Unmarshal(raw, &parsed); err != nil { + t.Fatalf("marshal produced invalid JSON: %v", err) + } + if _, ok := parsed["injected_key"]; ok { + t.Fatalf("hostile bucket injected a sibling key: %s", raw) + } + if parsed["bucket"] != block.Bucket { + t.Fatalf("bucket round-trip = %v", parsed["bucket"]) + } +} + +func TestReadS3OffloadEnv(t *testing.T) { + t.Setenv(envS3Bucket, "") + t.Setenv(envS3Region, "") + t.Setenv(envS3Prefix, "") + t.Setenv(envS3AccessKey, "") + t.Setenv(envS3SecretKey, "") + t.Setenv(envS3Endpoint, "") + t.Setenv(envS3ForcePathStyle, "") + info := readS3OffloadEnv() + if info.Enabled { + t.Fatal("empty bucket must disable offload") + } + + t.Setenv(envS3Bucket, " "+testBucket+" ") + t.Setenv(envS3Region, testRegion) + t.Setenv(envS3Prefix, testPrefix) + t.Setenv(envS3AccessKey, testPlainKey) + t.Setenv(envS3SecretKey, testPlainSecret) + info = readS3OffloadEnv() + if !info.Enabled || info.Bucket != testBucket || info.Region != testRegion || info.Prefix != testPrefix { + t.Fatalf("unexpected info: %+v", info) + } + if !info.HasAccessKey || !info.HasSecretKey { + t.Fatal("static creds should be detected") + } + if info.HasAccessKey && strings.Contains(info.Bucket, testPlainKey) { + t.Fatal("info must not carry key material on Bucket") + } +} + +func TestMaterializeConfig_RealSeed(t *testing.T) { + seedPath := filepath.Join("..", "data", "config.json") + seed, err := os.ReadFile(seedPath) + if err != nil { + t.Skipf("seed not available: %v", err) + } + + off, _, err := materializeConfigFrom(seed, "/app/data", s3OffloadInfo{}) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(off, seed) { + t.Fatal("real seed must pass through byte-identical when S3 is off") + } + + on, _, err := materializeConfigFrom(seed, "/app/data", s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: testPrefix, + HasAccessKey: true, + HasSecretKey: true, + }) + if err != nil { + t.Fatal(err) + } + var parsed map[string]any + if err := json.Unmarshal(on, &parsed); err != nil { + t.Fatalf("injected real seed is not JSON: %v", err) + } + if parsed["plugins"] == nil || parsed["providers"] == nil || parsed["auth_config"] == nil { + t.Fatal("injection must preserve plugins/providers/auth_config") + } + ls := parsed["logs_store"].(map[string]any) + obj := ls["object_storage"].(map[string]any) + if obj["access_key_id"] != envRefAccessKey { + t.Fatal("real-seed injection must use env.* credential refs") + } + if bytes.Contains(on, []byte(testPlainKey)) { + t.Fatal("plaintext key in real-seed output") + } +} + +func TestSeedFile_LogRetentionDays(t *testing.T) { + seedPath := filepath.Join("..", "data", "config.json") + raw, err := os.ReadFile(seedPath) + if err != nil { + t.Skipf("seed not available: %v", err) + } + var parsed map[string]any + if err := json.Unmarshal(raw, &parsed); err != nil { + t.Fatalf("seed is not valid JSON: %v", err) + } + client, _ := parsed["client"].(map[string]any) + if days, _ := client["log_retention_days"].(float64); days != 36500 { + t.Fatalf("data/config.json client.log_retention_days = %v, want 36500", client["log_retention_days"]) + } + if _, ok := parsed["logs_store"]; ok { + t.Fatal("seed must not contain a logs_store block (injection is opt-in at boot)") + } +} + +func TestSyncSeedConfig_S3Off_ByteIdenticalAndIdempotent(t *testing.T) { + t.Setenv(envS3Bucket, "") + t.Setenv(envS3Region, "") + t.Setenv(envS3Prefix, "") + t.Setenv(envS3AccessKey, "") + t.Setenv(envS3SecretKey, "") + + dir := t.TempDir() + seedPath := filepath.Join(dir, "config.json.seed") + appDir := filepath.Join(dir, "data") + if err := os.WriteFile(seedPath, []byte(testSeed), 0o644); err != nil { + t.Fatal(err) + } + logger := log.New(os.Stderr, "[test] ", 0) + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("first sync: %v", err) + } + got, err := os.ReadFile(filepath.Join(appDir, "config.json")) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, []byte(testSeed)) { + t.Fatalf("S3-off output must equal seed\ngot:\n%s", got) + } + + dst := filepath.Join(appDir, "config.json") + past := time.Date(2020, 1, 2, 3, 4, 5, 0, time.UTC) + if err := os.Chtimes(dst, past, past); err != nil { + t.Fatal(err) + } + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("second sync: %v", err) + } + st, err := os.Stat(dst) + if err != nil { + t.Fatal(err) + } + if !st.ModTime().Equal(past) { + t.Fatalf("idempotent S3-off sync bumped mtime: got %v want %v", st.ModTime(), past) + } +} + +func TestSyncSeedConfig_S3On_EnvRefsNoPlaintextIdempotent(t *testing.T) { + t.Setenv(envS3Bucket, testBucket) + t.Setenv(envS3Region, testRegion) + t.Setenv(envS3Prefix, testPrefix) + t.Setenv(envS3AccessKey, testPlainKey) + t.Setenv(envS3SecretKey, testPlainSecret) + + dir := t.TempDir() + seedPath := filepath.Join(dir, "config.json.seed") + appDir := filepath.Join(dir, "data") + if err := os.WriteFile(seedPath, []byte(testSeed), 0o644); err != nil { + t.Fatal(err) + } + logger := log.New(os.Stderr, "[test] ", 0) + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("first sync: %v", err) + } + + dst := filepath.Join(appDir, "config.json") + got, err := os.ReadFile(dst) + if err != nil { + t.Fatal(err) + } + if bytes.Equal(got, []byte(testSeed)) { + t.Fatal("S3-on output must differ from seed (logs_store injected)") + } + if bytes.Contains(got, []byte(testPlainKey)) || bytes.Contains(got, []byte(testPlainSecret)) { + t.Fatal("plaintext key material leaked into config.json") + } + if !bytes.Contains(got, []byte(envRefAccessKey)) || !bytes.Contains(got, []byte(envRefSecretKey)) { + t.Fatal("missing env. credential references") + } + if !bytes.Contains(got, []byte(`"bucket": "`+testBucket+`"`)) { + t.Fatalf("missing bucket in output:\n%s", got) + } + + var parsed map[string]any + if err := json.Unmarshal(got, &parsed); err != nil { + t.Fatalf("written config is not JSON: %v", err) + } + client := parsed["client"].(map[string]any) + if days, _ := client["log_retention_days"].(float64); days != 36500 { + t.Fatalf("log_retention_days = %v, want 36500", client["log_retention_days"]) + } + + past := time.Date(2020, 1, 2, 3, 4, 5, 0, time.UTC) + if err := os.Chtimes(dst, past, past); err != nil { + t.Fatal(err) + } + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("second sync: %v", err) + } + st, err := os.Stat(dst) + if err != nil { + t.Fatal(err) + } + if !st.ModTime().Equal(past) { + t.Fatalf("idempotent S3-on sync bumped mtime: got %v want %v", st.ModTime(), past) + } + + // Changing env must rewrite. + t.Setenv(envS3Prefix, "bifrost/other-org") + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("third sync: %v", err) + } + st, err = os.Stat(dst) + if err != nil { + t.Fatal(err) + } + if st.ModTime().Equal(past) { + t.Fatal("changing BIFROST_S3_PREFIX should rewrite config.json") + } + rewritten, _ := os.ReadFile(dst) + if !bytes.Contains(rewritten, []byte("bifrost/other-org")) { + t.Fatalf("rewritten config missing new prefix:\n%s", rewritten) + } +} + +func TestSyncSeedConfig_HostileBucketFallsBackToSeed(t *testing.T) { + t.Setenv(envS3Bucket, `foo", "injected_key": "pwned`) + t.Setenv(envS3Region, testRegion) + t.Setenv(envS3Prefix, defaultS3Prefix) + t.Setenv(envS3AccessKey, testPlainKey) + t.Setenv(envS3SecretKey, testPlainSecret) + + dir := t.TempDir() + seedPath := filepath.Join(dir, "config.json.seed") + appDir := filepath.Join(dir, "data") + if err := os.WriteFile(seedPath, []byte(testSeed), 0o644); err != nil { + t.Fatal(err) + } + logger := log.New(os.Stderr, "[test] ", 0) + if err := syncSeedConfig(logger, seedPath, appDir); err != nil { + t.Fatalf("sync: %v", err) + } + got, err := os.ReadFile(filepath.Join(appDir, "config.json")) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, []byte(testSeed)) { + t.Fatalf("hostile bucket must fall back to seed, not write injected JSON\ngot:\n%s", got) + } + if bytes.Contains(got, []byte("injected_key")) { + t.Fatal("injected_key leaked into written config") + } +} + +func TestLogS3Offload_NeverPrintsSecrets(t *testing.T) { + var buf bytes.Buffer + logger := log.New(&buf, "", 0) + logS3Offload(logger, s3OffloadInfo{ + Enabled: true, + Bucket: testBucket, + Region: testRegion, + Prefix: testPrefix, + HasAccessKey: true, + HasSecretKey: true, + }) + s := buf.String() + if !strings.Contains(s, testBucket) || !strings.Contains(s, testRegion) { + t.Fatalf("expected bucket+region in log, got %q", s) + } + if strings.Contains(s, testPlainKey) || strings.Contains(s, testPlainSecret) { + t.Fatalf("secrets leaked in log: %q", s) + } + + buf.Reset() + logS3Offload(logger, s3OffloadInfo{Enabled: true, Bucket: testBucket}) + s = buf.String() + if !strings.Contains(s, "WARNING") { + t.Fatalf("missing WARNING for bucket-without-region/creds: %q", s) + } +} From 163d0a3dd782a1065d13b80d36092dd5ed91f9ad Mon Sep 17 00:00:00 2001 From: gonzaloaune Date: Wed, 9 Sep 2026 18:23:17 +0000 Subject: [PATCH 2/7] Generated with Hive: Remove Chrome apt sources before apt-get update to fix Hash Sum mismatch in CI --- .github/workflows/rust-test-lsp.yml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/workflows/rust-test-lsp.yml b/.github/workflows/rust-test-lsp.yml index 4b91313fc..185aad1ef 100644 --- a/.github/workflows/rust-test-lsp.yml +++ b/.github/workflows/rust-test-lsp.yml @@ -65,6 +65,11 @@ jobs: # -------------------- Java / Kotlin -------------------- - name: Install OpenJDK 17 and jq run: | + # GH ubuntu runners ship Google Chrome's apt repo, which flakes + # with Hash Sum mismatch. This job does not need Chrome. + sudo rm -f /etc/apt/sources.list.d/google-chrome.list \ + /etc/apt/sources.list.d/google-chrome*.list \ + /etc/apt/sources.list.d/google-chrome*.sources sudo apt-get update sudo apt-get install -y openjdk-17-jdk jq unzip curl @@ -107,6 +112,11 @@ jobs: # -------------------- Ruby -------------------- - name: Install Ruby and dependencies run: | + # Same Chrome-repo flake as the OpenJDK step — drop it again in + # case a later runner image re-adds the source. + sudo rm -f /etc/apt/sources.list.d/google-chrome.list \ + /etc/apt/sources.list.d/google-chrome*.list \ + /etc/apt/sources.list.d/google-chrome*.sources sudo apt-get update sudo apt-get install -y ruby-full build-essential From 4471909978390c12ca8dc36e29bb3a46527f4077 Mon Sep 17 00:00:00 2001 From: Evan Feenstra Date: Thu, 10 Sep 2026 09:27:25 -0700 Subject: [PATCH 3/7] gateway: add xAI as the fourth provider; price provider-namespaced models (#1673) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Providers are now anthropic / openai / openrouter / xai. Gemini is dropped from the seed config, compose env, plan docs, and the mcp/docs/gateway seed copy; xai reads env.XAI_API_KEY. Pricing: the accumulator looked up the bare wire model only, but bifrost's datasheet keys xAI and OpenRouter rows as "/" ("xai/grok-4.3", "openrouter/moonshotai/…"), so every Grok and OpenRouter call accumulated at $0 while logs.db priced it. pricing.Keys now yields bare → provider-prefixed → prefix-stripped candidates, shared by the catalog and the operator model_pricing table; the posthook passes RoutingInfo.Provider (falling back to the deprecated ExtraFields.Provider). Also: xai entry in the canvas provider table + a stand-in glyph, an xai call in smoke-test.sh, and Makefile BIFROST_VERSION / GO_VERSION pins brought in line with the Dockerfile (transports/v1.6.2, 1.26.4) — `make docker-build` was overriding the Dockerfile default with 1.5.2. Verified live against the rebuilt image: grok-4.3 non-streaming and streaming calls under a shadow-mode macaroon accumulate 0.00025625 and 0.000285 in Redis, matching the datasheet rate to the cent. --- gateway/Makefile | 4 +- gateway/data/config.json | 6 +- gateway/docker-compose.yml | 1 + .../adminapi/ui/src/pages/canvasTheme.ts | 15 +++++ gateway/internal/auth/pricing.go | 30 +++++---- gateway/internal/auth/pricing_test.go | 58 ++++++++++++++--- gateway/internal/hooks/llm_posthook.go | 30 ++++++--- gateway/internal/hooks/llm_posthook_test.go | 42 +++++++++++++ gateway/internal/pricing/pricing.go | 52 +++++++++++---- gateway/internal/pricing/pricing_test.go | 63 ++++++++++++++++--- gateway/plans/llm-governance-v2.md | 4 +- gateway/plans/phases/phase-1-reconciler.md | 4 +- gateway/scripts/smoke-test.sh | 6 +- mcp/docs/gateway/README.md | 6 +- mcp/docs/gateway/data/config.json | 6 +- mcp/docs/gateway/docker-compose.yml | 2 +- 16 files changed, 261 insertions(+), 68 deletions(-) diff --git a/gateway/Makefile b/gateway/Makefile index 563923ff5..9d01de629 100644 --- a/gateway/Makefile +++ b/gateway/Makefile @@ -7,8 +7,8 @@ OUTPUT = $(OUTPUT_DIR)/$(PLUGIN_NAME).so # Bifrost upstream version this plugin is built against. Must match the # bifrost-http binary it gets loaded into — Go plugins are strict. -BIFROST_VERSION ?= transports/v1.5.2 -GO_VERSION ?= 1.26.2 +BIFROST_VERSION ?= transports/v1.6.2 +GO_VERSION ?= 1.26.4 # Docker image we build (dynamic bifrost-http + this plugin). IMAGE ?= stakgraph-gateway:dev diff --git a/gateway/data/config.json b/gateway/data/config.json index 5e11c1d08..f6514fd96 100644 --- a/gateway/data/config.json +++ b/gateway/data/config.json @@ -42,11 +42,11 @@ } ] }, - "gemini": { + "xai": { "keys": [ { - "name": "gemini-key-1", - "value": "env.GOOGLE_API_KEY", + "name": "xai-key-1", + "value": "env.XAI_API_KEY", "models": ["*"], "weight": 1.0 } diff --git a/gateway/docker-compose.yml b/gateway/docker-compose.yml index 742724975..3866e4390 100644 --- a/gateway/docker-compose.yml +++ b/gateway/docker-compose.yml @@ -82,6 +82,7 @@ services: OPENAI_API_KEY: ${OPENAI_API_KEY:-} OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:-} GOOGLE_API_KEY: ${GOOGLE_API_KEY:-} + XAI_API_KEY: ${XAI_API_KEY:-} # Optional S3 offload of LLM request/response bodies. Unset # BIFROST_S3_BUCKET ⇒ local SQLite logs only (today's behaviour). # When set, the wrapper injects logs_store.object_storage into diff --git a/gateway/internal/adminapi/ui/src/pages/canvasTheme.ts b/gateway/internal/adminapi/ui/src/pages/canvasTheme.ts index e5f98dd74..b8b8b0701 100644 --- a/gateway/internal/adminapi/ui/src/pages/canvasTheme.ts +++ b/gateway/internal/adminapi/ui/src/pages/canvasTheme.ts @@ -309,6 +309,17 @@ const OPENROUTER_PATHS: IconPathData = [ }, ]; +// xAI has no simple-icons glyph, so this is a geometric stand-in +// rather than the trademark: the two strokes of an "X" in a 24-unit +// box, the right-hand stroke pulled apart into a long diagonal plus +// a short cap the way the wordmark reads. Stroked (not filled) so it +// renders at the same visual weight as the OpenRouter curves. +const XAI_PATHS: IconPathData = [ + { d: "M4 4L20 20", mode: "stroke", strokeWidth: 3 }, + { d: "M20 4L11.5 12.5", mode: "stroke", strokeWidth: 3 }, + { d: "M4 20L8.5 15.5", mode: "stroke", strokeWidth: 3 }, +]; + // Person + bot glyphs in a 24-unit source box. The `l 0.001 0` segments // are degenerate dots — at stroke-linecap='round' they render as filled // discs of diameter = stroke-width. @@ -341,6 +352,7 @@ const PROVIDER_ICON_META: Record< openai: { mode: "fill", viewBox: 24 }, gemini: { mode: "fill", viewBox: 24 }, openrouter: { mode: "fill", viewBox: 512 }, + xai: { mode: "stroke", viewBox: 24 }, }; // --------------------------------------------------------------------------- @@ -358,6 +370,7 @@ const PROVIDER_ICON_META: Record< // - OpenAI: mint, used across their docs + simple-icons // - Gemini: simple-icons "Google Gemini" (#8E75B2) // - OpenRouter: openrouter.ai brand +// - xAI: monochrome wordmark; a light neutral reads on midnight export const PROVIDER_DISPLAY: Record< string, { label: string; icon: string; color: string } @@ -366,6 +379,7 @@ export const PROVIDER_DISPLAY: Record< openai: { label: "OpenAI", icon: "openai", color: "#10A37F" }, gemini: { label: "Gemini", icon: "gemini", color: "#8E75B2" }, openrouter: { label: "OpenRouter", icon: "openrouter", color: "#6467F2" }, + xai: { label: "xAI", icon: "xai", color: "#D4D4D8" }, }; export function providerIcon(name: string): string { @@ -386,6 +400,7 @@ export const canvasTheme: CanvasTheme = resolveTheme( openai: OPENAI_PATHS, gemini: GEMINI_PATHS, openrouter: OPENROUTER_PATHS, + xai: XAI_PATHS, }, categories: { // ─── agent ─────────────────────────────────────────────────── diff --git a/gateway/internal/auth/pricing.go b/gateway/internal/auth/pricing.go index 015dbb1a1..3b825ef78 100644 --- a/gateway/internal/auth/pricing.go +++ b/gateway/internal/auth/pricing.go @@ -1,8 +1,6 @@ package auth import ( - "strings" - "github.com/stakwork/stakgraph/gateway/internal/pricing" ) @@ -20,36 +18,36 @@ import ( // bifrost prices logs.db rows from, so enforcement dollars and // reported dollars agree. // -// Model matching: exact key first, then the name with any -// "provider/" prefix stripped, so "anthropic/claude-sonnet-5" and -// "claude-sonnet-5" resolve to the same entry whichever form the -// caller holds. -func PriceCall(model string, promptTokens, completionTokens int) (float64, bool) { - if model == "" { +// Model matching follows pricing.Keys: the bare name, then +// "/", then the name with any "provider/" prefix +// stripped — so "claude-sonnet-5", "anthropic/claude-sonnet-5", and +// ("xai", "grok-4") → "xai/grok-4" all resolve against either source +// whichever form the caller holds. provider is the Bifrost provider +// id that served the call; "" is allowed and skips the second form. +func PriceCall(provider, model string, promptTokens, completionTokens int) (float64, bool) { + keys := pricing.Keys(provider, model) + if len(keys) == 0 { return 0, false } const mtok = 1_000_000 - if entry, ok := configPrice(model); ok { + if entry, ok := configPrice(keys); ok { return float64(promptTokens)*entry.InputPerMTok/mtok + float64(completionTokens)*entry.OutputPerMTok/mtok, true } - if p, ok := pricing.Lookup(model); ok { + if p, ok := pricing.Lookup(provider, model); ok { return float64(promptTokens)*p.InputPerMTok/mtok + float64(completionTokens)*p.OutputPerMTok/mtok, true } return 0, false } -func configPrice(model string) (ModelPrice, bool) { +func configPrice(keys []string) (ModelPrice, bool) { table := GetConfig().ModelPricing if len(table) == 0 { return ModelPrice{}, false } - if entry, ok := table[model]; ok { - return entry, true - } - if i := strings.LastIndexByte(model, '/'); i >= 0 { - if entry, ok := table[model[i+1:]]; ok { + for _, k := range keys { + if entry, ok := table[k]; ok { return entry, true } } diff --git a/gateway/internal/auth/pricing_test.go b/gateway/internal/auth/pricing_test.go index e8d3d3414..e2dde9d9c 100644 --- a/gateway/internal/auth/pricing_test.go +++ b/gateway/internal/auth/pricing_test.go @@ -1,6 +1,7 @@ package auth import ( + "math" "testing" "github.com/stakwork/stakgraph/gateway/internal/pricing" @@ -13,25 +14,25 @@ func TestPriceCall(t *testing.T) { t.Cleanup(func() { SetConfigForTest(Config{}) }) // 1000 in + 500 out = 0.0008 + 0.002 = 0.0028. - got, ok := PriceCall("claude-3-5-haiku-latest", 1000, 500) + got, ok := PriceCall("anthropic", "claude-3-5-haiku-latest", 1000, 500) if !ok || got != 0.0028 { t.Fatalf("PriceCall = (%v, %v), want (0.0028, true)", got, ok) } // Provider-prefixed form resolves to the same entry. - got, ok = PriceCall("anthropic/claude-3-5-haiku-latest", 1000, 500) + got, ok = PriceCall("anthropic", "anthropic/claude-3-5-haiku-latest", 1000, 500) if !ok || got != 0.0028 { t.Fatalf("prefixed PriceCall = (%v, %v), want (0.0028, true)", got, ok) } // Unpriced model: (0, false), caller logs. - if _, ok := PriceCall("gpt-4o", 10, 10); ok { + if _, ok := PriceCall("openai", "gpt-4o", 10, 10); ok { t.Fatal("unpriced model must return ok=false") } // Empty table: never ok. SetConfigForTest(Config{}) - if _, ok := PriceCall("claude-3-5-haiku-latest", 10, 10); ok { + if _, ok := PriceCall("anthropic", "claude-3-5-haiku-latest", 10, 10); ok { t.Fatal("empty pricing table must return ok=false") } } @@ -47,7 +48,7 @@ func TestPriceCall_CatalogFallback(t *testing.T) { // No config entry → catalog prices it: 1000*2/1e6 + 500*10/1e6. SetConfigForTest(Config{}) - got, ok := PriceCall("claude-sonnet-5", 1000, 500) + got, ok := PriceCall("anthropic", "claude-sonnet-5", 1000, 500) if !ok || got != 0.007 { t.Fatalf("catalog PriceCall = (%v, %v), want (0.007, true)", got, ok) } @@ -56,13 +57,56 @@ func TestPriceCall_CatalogFallback(t *testing.T) { SetConfigForTest(Config{ModelPricing: map[string]ModelPrice{ "claude-sonnet-5": {InputPerMTok: 4.0, OutputPerMTok: 20.0}, }}) - got, ok = PriceCall("claude-sonnet-5", 1000, 500) + got, ok = PriceCall("anthropic", "claude-sonnet-5", 1000, 500) if !ok || got != 0.014 { t.Fatalf("config-over-catalog PriceCall = (%v, %v), want (0.014, true)", got, ok) } // Neither source knows the model → (0, false). - if _, ok := PriceCall("mystery-model", 10, 10); ok { + if _, ok := PriceCall("anthropic", "mystery-model", 10, 10); ok { t.Fatal("model absent from both sources must return ok=false") } } + +func TestPriceCall_ProviderNamespaced(t *testing.T) { + // Catalog keyed the way bifrost's datasheet keys xAI and + // OpenRouter rows; bifrost reports the wire model bare. + pricing.SetTableForTest(map[string]pricing.Price{ + "xai/grok-4": {InputPerMTok: 3.0, OutputPerMTok: 15.0}, + "openrouter/moonshotai/kimi-k2-0905": {InputPerMTok: 1.0, OutputPerMTok: 1.0}, + }) + t.Cleanup(func() { + pricing.SetTableForTest(nil) + SetConfigForTest(Config{}) + }) + SetConfigForTest(Config{}) + + // 1000*3/1e6 + 500*15/1e6 = 0.003 + 0.0075 (float sum, so a + // tolerance rather than an exact decimal). + got, ok := PriceCall("xai", "grok-4", 1000, 500) + if !ok || math.Abs(got-0.0105) > 1e-12 { + t.Fatalf("PriceCall(xai, grok-4) = (%v, %v), want (≈0.0105, true)", got, ok) + } + got, ok = PriceCall("openrouter", "moonshotai/kimi-k2-0905", 1000, 1000) + if !ok || got != 0.002 { + t.Fatalf("PriceCall(openrouter, …kimi) = (%v, %v), want (0.002, true)", got, ok) + } + // No provider → no namespaced candidate → the bare row is absent. + if _, ok := PriceCall("", "grok-4", 10, 10); ok { + t.Fatal("bare grok-4 without a provider must miss the namespaced catalog") + } + + // The operator table accepts either spelling and still wins. + SetConfigForTest(Config{ModelPricing: map[string]ModelPrice{ + "grok-4": {InputPerMTok: 1.0, OutputPerMTok: 1.0}, + }}) + if got, ok := PriceCall("xai", "grok-4", 1000, 1000); !ok || got != 0.002 { + t.Fatalf("bare config key over catalog = (%v, %v), want (0.002, true)", got, ok) + } + SetConfigForTest(Config{ModelPricing: map[string]ModelPrice{ + "xai/grok-4": {InputPerMTok: 2.0, OutputPerMTok: 2.0}, + }}) + if got, ok := PriceCall("xai", "grok-4", 1000, 1000); !ok || got != 0.004 { + t.Fatalf("namespaced config key over catalog = (%v, %v), want (0.004, true)", got, ok) + } +} diff --git a/gateway/internal/hooks/llm_posthook.go b/gateway/internal/hooks/llm_posthook.go index a2af7a483..80c7aecfe 100644 --- a/gateway/internal/hooks/llm_posthook.go +++ b/gateway/internal/hooks/llm_posthook.go @@ -96,9 +96,12 @@ func isStreamRequest(resp *schemas.BifrostResponse) bool { // per the phase-6 accumulator design: // // 1. Provider-computed Usage.Cost.TotalCost (only some providers). -// 2. auth.PriceCall on the resolved model name — operator -// model_pricing config first, then the internal/pricing catalog -// (bifrost's own datasheet, refreshed daily). +// 2. auth.PriceCall on the (provider, resolved model) pair — +// operator model_pricing config first, then the internal/pricing +// catalog (bifrost's own datasheet, refreshed daily). The +// provider matters: the datasheet keys some providers' rows as +// "/" ("xai/grok-4") while Bifrost reports the +// wire model bare ("grok-4"). // 3. $0, with a loud log — an unpriced model must be visible in // `docker logs`, not silently guessed at. func resolveCost(chat *schemas.BifrostChatResponse, usage *schemas.BifrostLLMUsage, dims map[string]string) float64 { @@ -108,24 +111,37 @@ func resolveCost(chat *schemas.BifrostChatResponse, usage *schemas.BifrostLLMUsa if usage.Cost != nil && usage.Cost.TotalCost > 0 { return usage.Cost.TotalCost } - model := "" + model, provider := "", "" if chat != nil { if model = chat.ExtraFields.ResolvedModelUsed; model == "" { model = chat.Model } + provider = responseProvider(chat) } - if cost, ok := auth.PriceCall(model, usage.PromptTokens, usage.CompletionTokens); ok { + if cost, ok := auth.PriceCall(provider, model, usage.PromptTokens, usage.CompletionTokens); ok { return cost } if usage.TotalTokens > 0 { pluginlog.Warnf( - "accounting: no price for model=%q (run_id=%s) — %d tokens accumulated as $0; add a model_pricing entry", - model, dims[pluginctx.DimRunID], usage.TotalTokens, + "accounting: no price for provider=%q model=%q (run_id=%s) — %d tokens accumulated as $0; add a model_pricing entry", + provider, model, dims[pluginctx.DimRunID], usage.TotalTokens, ) } return 0 } +// responseProvider returns the Bifrost provider id that actually +// served a chat response. RoutingInfo.Provider is the current field; +// ExtraFields.Provider is its deprecated twin, still populated by +// core v1.6 and kept as the fallback for chunks that only carry the +// old shape. +func responseProvider(chat *schemas.BifrostChatResponse) string { + if p := chat.ExtraFields.RoutingInfo.Provider; p != "" { + return string(p) + } + return string(chat.ExtraFields.Provider) +} + // toolCallNames collects the tool names invoked in a non-streaming // chat response, feeding the tools:run history that phase-6's // tool-loop heuristic reads. Nil when the response called no tools. diff --git a/gateway/internal/hooks/llm_posthook_test.go b/gateway/internal/hooks/llm_posthook_test.go index a43fcfb60..523c04525 100644 --- a/gateway/internal/hooks/llm_posthook_test.go +++ b/gateway/internal/hooks/llm_posthook_test.go @@ -6,6 +6,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" "github.com/stakwork/stakgraph/gateway/internal/auth" + "github.com/stakwork/stakgraph/gateway/internal/pricing" ) func strPtr(s string) *string { return &s } @@ -98,3 +99,44 @@ func TestResolveCost_PrefersResolvedModel(t *testing.T) { t.Fatalf("resolved-model cost = %v, want 0.001", got) } } + +// The xAI case: bifrost resolves the wire model bare ("grok-4") and +// reports the provider on RoutingInfo; the datasheet keys the row as +// "xai/grok-4". The hook must hand the provider to the price lookup +// or every Grok call accumulates at $0. +func TestResolveCost_ProviderNamespacedCatalog(t *testing.T) { + auth.SetConfigForTest(auth.Config{}) + pricing.SetTableForTest(map[string]pricing.Price{ + "xai/grok-4": {InputPerMTok: 1.0, OutputPerMTok: 1.0}, + }) + t.Cleanup(func() { + pricing.SetTableForTest(nil) + auth.SetConfigForTest(auth.Config{}) + }) + + usage := &schemas.BifrostLLMUsage{PromptTokens: 500, CompletionTokens: 500, TotalTokens: 1000} + + chat := &schemas.BifrostChatResponse{ + Model: "grok-4", + ExtraFields: schemas.BifrostResponseExtraFields{ + ResolvedModelUsed: "grok-4", + RoutingInfo: schemas.RoutingInfo{Provider: schemas.XAI, Model: "grok-4"}, + }, + } + if got := resolveCost(chat, usage, map[string]string{}); got != 0.001 { + t.Fatalf("RoutingInfo provider cost = %v, want 0.001", got) + } + + // Older chunks only carry the deprecated ExtraFields.Provider. + chat.ExtraFields.RoutingInfo = schemas.RoutingInfo{} + chat.ExtraFields.Provider = schemas.XAI + if got := resolveCost(chat, usage, map[string]string{}); got != 0.001 { + t.Fatalf("deprecated provider field cost = %v, want 0.001", got) + } + + // No provider anywhere → the namespaced row is unreachable → $0. + chat.ExtraFields.Provider = "" + if got := resolveCost(chat, usage, map[string]string{}); got != 0 { + t.Fatalf("provider-less cost = %v, want 0", got) + } +} diff --git a/gateway/internal/pricing/pricing.go b/gateway/internal/pricing/pricing.go index d68c31b89..79d9817fc 100644 --- a/gateway/internal/pricing/pricing.go +++ b/gateway/internal/pricing/pricing.go @@ -71,21 +71,49 @@ var ( retryBaseDelay = time.Second ) -// Lookup returns the catalog price for a model: exact key first, -// then the name with any "provider/" prefix stripped (Bifrost -// resolved names are usually bare, but callers occasionally hold the -// prefixed form). ok=false when the catalog has no entry — the -// caller falls through to its next price source. -func Lookup(model string) (Price, bool) { +// Keys returns the lookup candidates for a (provider, model) pair, +// most specific first: +// +// 1. model as given — the datasheet keys most first-party rows +// bare ("claude-sonnet-5", "gpt-5.2"), which is also the wire +// model Bifrost reports in ResolvedModelUsed. +// 2. "/" — the datasheet namespaces some +// providers' rows under the provider id ("xai/grok-4", +// "openrouter/moonshotai/kimi-k2-0905") while Bifrost still +// reports the wire model bare ("grok-4"). Without this step +// every Grok and OpenRouter call priced at $0. +// 3. the last path segment of model — callers occasionally hold a +// prefixed form ("anthropic/claude-sonnet-5"). +// +// Shared by the catalog and the operator model_pricing table so the +// two sources can't disagree on what a model name means. +func Keys(provider, model string) []string { + if model == "" { + return nil + } + keys := []string{model} + if provider != "" && !strings.HasPrefix(model, provider+"/") { + keys = append(keys, provider+"/"+model) + } + if i := strings.LastIndexByte(model, '/'); i >= 0 && i+1 < len(model) { + keys = append(keys, model[i+1:]) + } + return keys +} + +// Lookup returns the catalog price for a model, trying each Keys +// candidate in order. provider is the Bifrost provider id that +// served the call ("xai", "openrouter", …); pass "" when unknown +// and only the bare and prefix-stripped forms are tried. ok=false +// when the catalog has no entry — the caller falls through to its +// next price source. +func Lookup(provider, model string) (Price, bool) { m := table.Load() - if m == nil || model == "" { + if m == nil { return Price{}, false } - if p, ok := (*m)[model]; ok { - return p, true - } - if i := strings.LastIndexByte(model, '/'); i >= 0 { - if p, ok := (*m)[model[i+1:]]; ok { + for _, k := range Keys(provider, model) { + if p, ok := (*m)[k]; ok { return p, true } } diff --git a/gateway/internal/pricing/pricing_test.go b/gateway/internal/pricing/pricing_test.go index 3af9ad7b4..1229ef666 100644 --- a/gateway/internal/pricing/pricing_test.go +++ b/gateway/internal/pricing/pricing_test.go @@ -63,20 +63,65 @@ func TestParseDatasheet_RejectsGarbage(t *testing.T) { } } -func TestLookup_PrefixStrip(t *testing.T) { - SetTableForTest(map[string]Price{"claude-sonnet-5": {InputPerMTok: 2, OutputPerMTok: 10}}) +func TestKeys_Order(t *testing.T) { + cases := []struct { + provider, model string + want []string + }{ + {"anthropic", "claude-sonnet-5", []string{"claude-sonnet-5", "anthropic/claude-sonnet-5"}}, + {"xai", "grok-4", []string{"grok-4", "xai/grok-4"}}, + {"", "grok-4", []string{"grok-4"}}, + // Already provider-prefixed: no double prefix, base form last. + {"anthropic", "anthropic/claude-sonnet-5", []string{"anthropic/claude-sonnet-5", "claude-sonnet-5"}}, + // Nested vendor path (OpenRouter): provider form before base. + {"openrouter", "moonshotai/kimi-k2-0905", []string{"moonshotai/kimi-k2-0905", "openrouter/moonshotai/kimi-k2-0905", "kimi-k2-0905"}}, + {"xai", "", nil}, + } + for _, c := range cases { + got := Keys(c.provider, c.model) + if len(got) != len(c.want) { + t.Fatalf("Keys(%q,%q) = %v, want %v", c.provider, c.model, got, c.want) + } + for i := range got { + if got[i] != c.want[i] { + t.Fatalf("Keys(%q,%q) = %v, want %v", c.provider, c.model, got, c.want) + } + } + } +} + +func TestLookup_ProviderNamespaced(t *testing.T) { + SetTableForTest(map[string]Price{ + "claude-sonnet-5": {InputPerMTok: 2, OutputPerMTok: 10}, + "xai/grok-4": {InputPerMTok: 3, OutputPerMTok: 15}, + "openrouter/moonshotai/kimi-k2-0905": {InputPerMTok: 0.5, OutputPerMTok: 2}, + }) t.Cleanup(func() { SetTableForTest(nil) }) - if _, ok := Lookup("claude-sonnet-5"); !ok { + if _, ok := Lookup("anthropic", "claude-sonnet-5"); !ok { t.Fatal("exact lookup failed") } - if _, ok := Lookup("anthropic/claude-sonnet-5"); !ok { + if _, ok := Lookup("anthropic", "anthropic/claude-sonnet-5"); !ok { t.Fatal("prefix-stripped lookup failed") } - if _, ok := Lookup("unknown-model"); ok { + // The xAI regression: bifrost reports the wire model bare, the + // datasheet keys it under the provider. + if p, ok := Lookup("xai", "grok-4"); !ok || p.InputPerMTok != 3 { + t.Fatalf("Lookup(xai, grok-4) = (%+v, %v), want the xai/grok-4 row", p, ok) + } + if _, ok := Lookup("xai", "xai/grok-4"); !ok { + t.Fatal("already-prefixed model must not be double-prefixed into a miss") + } + if _, ok := Lookup("", "grok-4"); ok { + t.Fatal("without a provider the bare grok row must miss (there is none)") + } + if p, ok := Lookup("openrouter", "moonshotai/kimi-k2-0905"); !ok || p.OutputPerMTok != 2 { + t.Fatalf("Lookup(openrouter, moonshotai/kimi-k2-0905) = (%+v, %v)", p, ok) + } + if _, ok := Lookup("anthropic", "unknown-model"); ok { t.Fatal("unknown model must miss") } - if _, ok := Lookup(""); ok { + if _, ok := Lookup("anthropic", ""); ok { t.Fatal("empty model must miss") } } @@ -93,7 +138,7 @@ func TestFetch_SwapsAndPersists(t *testing.T) { cachePath := filepath.Join(t.TempDir(), "sheet.json") FetchNowForTest(srv.URL, cachePath) - if p, ok := Lookup("gpt-5.2"); !ok || p.InputPerMTok != 1.75 { + if p, ok := Lookup("openai", "gpt-5.2"); !ok || p.InputPerMTok != 1.75 { t.Fatalf("post-fetch Lookup(gpt-5.2) = (%+v, %v)", p, ok) } raw, err := os.ReadFile(cachePath) @@ -116,7 +161,7 @@ func TestFetch_FailureKeepsLastGood(t *testing.T) { FetchNowForTest(srv.URL, filepath.Join(t.TempDir(), "sheet.json")) - if _, ok := Lookup("claude-sonnet-5"); !ok { + if _, ok := Lookup("anthropic", "claude-sonnet-5"); !ok { t.Fatal("failed fetch must keep the last-good table") } } @@ -133,7 +178,7 @@ func TestFetch_BadBodyKeepsLastGood(t *testing.T) { cachePath := filepath.Join(t.TempDir(), "sheet.json") FetchNowForTest(srv.URL, cachePath) - if _, ok := Lookup("claude-sonnet-5"); !ok { + if _, ok := Lookup("anthropic", "claude-sonnet-5"); !ok { t.Fatal("unparseable body must keep the last-good table") } if _, err := os.Stat(cachePath); err == nil { diff --git a/gateway/plans/llm-governance-v2.md b/gateway/plans/llm-governance-v2.md index ad005fd82..e0948a8ef 100644 --- a/gateway/plans/llm-governance-v2.md +++ b/gateway/plans/llm-governance-v2.md @@ -135,7 +135,7 @@ Bifrost natively enforces: | Daily $ budget | Customer (per user) | $1000/day | | Rate limit RPM | Customer (per user) | 1000 RPM | | Rate limit TPM | Customer (per user) | 5M TPM | -| Provider allowlist | VK | `[anthropic, openai, openrouter, gemini]`, all `["*"]` | +| Provider allowlist | VK | `[anthropic, openai, openrouter, xai]`, all `["*"]` | | Model allowlist | VK | `["*"]` initially | | `is_active` | Customer (per user) | `true` by default; `false` disables account org-wide | @@ -173,7 +173,7 @@ POST /api/governance/virtual-keys { provider: "anthropic", allowed_models: ["*"] }, { provider: "openai", allowed_models: ["*"] }, { provider: "openrouter",allowed_models: ["*"] }, - { provider: "gemini", allowed_models: ["*"] }, + { provider: "xai", allowed_models: ["*"] }, ] } ``` diff --git a/gateway/plans/phases/phase-1-reconciler.md b/gateway/plans/phases/phase-1-reconciler.md index b94ab93af..fbc925c8b 100644 --- a/gateway/plans/phases/phase-1-reconciler.md +++ b/gateway/plans/phases/phase-1-reconciler.md @@ -255,7 +255,7 @@ Filters available on this endpoint (lines 369-380 of governance.go): "budgets": [], "rate_limit": null } - // … openai, openrouter, gemini + // … openai, openrouter, xai ], "mcp_configs": [], "budgets": [], @@ -306,7 +306,7 @@ Content-Type: application/json { "provider": "anthropic", "allowed_models": ["*"], "key_ids": ["*"] }, { "provider": "openai", "allowed_models": ["*"], "key_ids": ["*"] }, { "provider": "openrouter", "allowed_models": ["*"], "key_ids": ["*"] }, - { "provider": "gemini", "allowed_models": ["*"], "key_ids": ["*"] } + { "provider": "xai", "allowed_models": ["*"], "key_ids": ["*"] } ] } ``` diff --git a/gateway/scripts/smoke-test.sh b/gateway/scripts/smoke-test.sh index d9aaff8ac..3f677f6a6 100755 --- a/gateway/scripts/smoke-test.sh +++ b/gateway/scripts/smoke-test.sh @@ -14,7 +14,7 @@ # - Multiple agent names (browser / coder / chat / reviewer) # - Multiple users (alice / bob / carol) and workspaces (w1 / w2) # - Spread across ≥2 minute buckets via sleeps -# - Multiple providers (anthropic / openai / openrouter) and models +# - Multiple providers (anthropic / openai / openrouter / xai) and models # - One streaming request # - One error request (bad model name) # - One request with NO dim headers (graceful absence check) @@ -140,6 +140,9 @@ M_HAIKU="anthropic/claude-haiku-4-5-20251001" M_MINI="openai/gpt-4o-mini" M_NANO="openai/gpt-4.1-nano" M_KIMI="openrouter/moonshotai/kimi-k2-0905" +# xAI retires model ids aggressively; with allowed_models ["*"] Bifrost only +# admits ids present in api.x.ai's live /v1/models, so keep this current. +M_GROK="xai/grok-4.3" # --- batch 1: minute N ---------------------------------------------------- @@ -163,6 +166,7 @@ echo "# batch 2 — minute N+1" call_llm chat u_carol w2 "$SESS_CAROL_W2" "$M_HAIKU" 'two-word greeting' call_llm reviewer u_carol w1 "$SESS_CAROL_W1" "$M_NANO" 'reply in 3 words' call_llm reviewer u_carol w1 "$SESS_CAROL_W1" "$M_KIMI" 'reply in 3 words' +call_llm reviewer u_carol w1 "$SESS_CAROL_W1" "$M_GROK" 'reply in 3 words' echo echo "# streaming request (verify final-chunk cost lands)" diff --git a/mcp/docs/gateway/README.md b/mcp/docs/gateway/README.md index c27275dbc..42af1014a 100644 --- a/mcp/docs/gateway/README.md +++ b/mcp/docs/gateway/README.md @@ -7,10 +7,10 @@ Runs Bifrost locally on `http://localhost:8181` so MCP can be tested with ## Layout - `docker-compose.yml` — Bifrost service, host port `8181 -> container 8080`. -- `data/config.json` — seed config: Anthropic + OpenAI + OpenRouter + Gemini +- `data/config.json` — seed config: Anthropic + OpenAI + OpenRouter + xAI providers, API keys read from `env.*`, `config_store` enabled (SQLite) so the - Web UI works. (Bifrost calls Google's public Gemini API `gemini`; the MCP - client side calls the same thing `google`.) + Web UI works. (xAI has no dedicated Bifrost route; the MCP client sends Grok + through `/openai/v1` with an `xai/`-prefixed model id.) - `data/config.db`, `data/logs.db` — created on first boot; gitignored. - `.env` — symlinked to `../../.env` (the MCP `.env`) so docker-compose picks up `ANTHROPIC_API_KEY` etc. without copy-paste. diff --git a/mcp/docs/gateway/data/config.json b/mcp/docs/gateway/data/config.json index 085c02a16..7743756c2 100644 --- a/mcp/docs/gateway/data/config.json +++ b/mcp/docs/gateway/data/config.json @@ -35,11 +35,11 @@ } ] }, - "gemini": { + "xai": { "keys": [ { - "name": "gemini-key-1", - "value": "env.GOOGLE_API_KEY", + "name": "xai-key-1", + "value": "env.XAI_API_KEY", "models": ["*"], "weight": 1.0 } diff --git a/mcp/docs/gateway/docker-compose.yml b/mcp/docs/gateway/docker-compose.yml index d10787d2a..0b401c04d 100644 --- a/mcp/docs/gateway/docker-compose.yml +++ b/mcp/docs/gateway/docker-compose.yml @@ -13,7 +13,7 @@ services: ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY:-} OPENAI_API_KEY: ${OPENAI_API_KEY:-} OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:-} - GOOGLE_API_KEY: ${GOOGLE_API_KEY:-} + XAI_API_KEY: ${XAI_API_KEY:-} ports: # Host 8181 -> container 8080. MCP server uses LLM_GATEWAY_URL=http://localhost:8181. - "8181:8080" From 83c5e87d961d8ea3d099fb70b75e7f69333181c4 Mon Sep 17 00:00:00 2001 From: Evan Feenstra Date: Thu, 10 Sep 2026 09:27:33 -0700 Subject: [PATCH 4/7] gateway: BIFROST_PLUGIN_ENFORCE_MACAROONS env override for enforce_macaroons (#1674) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit config.json is baked into the image and re-seeded on every boot, so the only way to move a swarm from shadow to enforce mode was a new image. The env var now overrides the plugin config block's enforce_macaroons when set (1/true/yes/on, 0/false/no/off); unset keeps the config.json value. The boot line names the winner: `auth: macaroon adapter wired enforce= source=env|config|default`. An unparseable value is deliberately not fatal. A plugin Init error does not stop bifrost-http — the wrapper waits 5s and then serves inference without the plugin, i.e. with no macaroon verification, no dim canonicalization, and /_plugin/* down, which is worse than either mode. So a typo logs at ERROR, the boot line reports source=env-invalid, and the config.json value stands. Verified live on the rebuilt image: enforce=true via env rejects a missing macaroon (401 "x-macaroon header is required") and a garbage one, and passes a valid one with mode=enforce; unset falls back to source=config; "ture" keeps the plugin active with the ERROR line. --- gateway/docker-compose.yml | 9 ++ gateway/internal/auth/config.go | 81 ++++++++++++++---- gateway/internal/auth/config_test.go | 100 ++++++++++++++++++++++ gateway/internal/auth/doc.go | 7 +- gateway/internal/env/env.go | 32 +++++++ gateway/internal/env/env_test.go | 42 +++++++++ gateway/main.go | 9 +- gateway/scripts/smoke-test-enforcement.sh | 8 +- gateway/scripts/smoke-test-phase-11.sh | 3 +- 9 files changed, 266 insertions(+), 25 deletions(-) create mode 100644 gateway/internal/env/env_test.go diff --git a/gateway/docker-compose.yml b/gateway/docker-compose.yml index 3866e4390..35ba858ee 100644 --- a/gateway/docker-compose.yml +++ b/gateway/docker-compose.yml @@ -72,6 +72,15 @@ services: # (`bifrost:login_attempts:*`). Without it, /_plugin/login # returns 503 and the admin dashboard is unreachable. BIFROST_PLUGIN_REDIS_URL: redis://redis:6379/0 + # Overrides the plugin's `enforce_macaroons` flag from + # data/config.json (baked into the image and re-seeded on every + # boot, so this is the only way to flip a running deployment + # without shipping a new image). Unset ⇒ config.json value + # (false: shadow mode — verify + log, never reject). true ⇒ + # missing/invalid macaroons short-circuit with 401. Anything + # other than 1/true/yes/on or 0/false/no/off is logged at ERROR + # and ignored (boot line shows source=env-invalid). + BIFROST_PLUGIN_ENFORCE_MACAROONS: ${BIFROST_PLUGIN_ENFORCE_MACAROONS:-} # PRODUCTION=1 forces the `Secure` attribute on the session # cookie regardless of the request scheme. Leave unset in dev # so localhost HTTP can keep its cookie; flip to 1 behind a diff --git a/gateway/internal/auth/config.go b/gateway/internal/auth/config.go index c95597e91..b44c1bfb6 100644 --- a/gateway/internal/auth/config.go +++ b/gateway/internal/auth/config.go @@ -3,6 +3,9 @@ package auth import ( "encoding/json" "sync" + + "github.com/stakwork/stakgraph/gateway/internal/env" + "github.com/stakwork/stakgraph/gateway/internal/pluginlog" ) // Config is the auth subset of the plugin's config block in @@ -25,6 +28,12 @@ import ( // Existing swarms running an old config.json will pick up the new // adapter in shadow mode without any operator action — that is the // whole point of the flag. Flip to true per-swarm as rollout proceeds. +// +// The env var BIFROST_PLUGIN_ENFORCE_MACAROONS (see internal/env) +// overrides the config.json value when set, so a swarm can be flipped +// through its environment without rebuilding the image that carries +// config.json. EnforceMacaroonsSource records which one won so the +// boot log can say so. type Config struct { // EnforceMacaroons gates whether verification failures actually // reject the request. When false (default, shadow mode), the @@ -33,6 +42,15 @@ type Config struct { // failures short-circuit with a bifrost.Error. EnforceMacaroons bool `json:"enforce_macaroons"` + // EnforceMacaroonsSource is where EnforceMacaroons came from: + // "env" (BIFROST_PLUGIN_ENFORCE_MACAROONS set and valid), "config" + // (the key is present in the plugin config block), "default" + // (neither — zero value, shadow mode), or "env-invalid" (the env + // var was set to something unparseable, an ERROR was logged, and + // the config/default value stands). Diagnostic only; never + // persisted — grep the boot line for it. + EnforceMacaroonsSource string `json:"-"` + // AgentBudgets is the per-agent windowed spend cap declared // in plugin.yaml. Phase 6's PreLLMHook reads these to gate // inference; phase 8's dashboard reads them to render the @@ -75,8 +93,11 @@ type ModelPrice struct { // subset of fields auth cares about. Bifrost passes the raw `config` // JSON object verbatim, so json.Marshal+Unmarshal round-trips through // whatever shape it actually has. +// +// EnforceMacaroons is a pointer so "key absent" (source=default) and +// "key present and false" (source=config) stay distinguishable. type pluginConfigEnvelope struct { - EnforceMacaroons bool `json:"enforce_macaroons"` + EnforceMacaroons *bool `json:"enforce_macaroons"` AgentBudgets map[string]AgentBudget `json:"agent_budgets"` ModelPricing map[string]ModelPrice `json:"model_pricing"` } @@ -94,7 +115,17 @@ var ( // to flip enforce_macaroons mid-suite without restarting the package. // // Safe to call with raw == nil; in that case the zero-value Config -// (shadow mode) is cached. +// (shadow mode) is cached — unless the env override is set, which +// applies on top of whatever the block held, nil included. +// +// Only malformed config JSON is an error. An unparseable +// BIFROST_PLUGIN_ENFORCE_MACAROONS value is deliberately NOT fatal: +// a plugin Init error does not stop bifrost-http, it just drops this +// plugin — the wrapper then serves inference with no macaroon +// verification, no claim canonicalization, and /_plugin/* down, +// which is strictly worse than shadow mode. So the typo is logged at +// ERROR, the boot line says source=env-invalid, and the config.json +// value stands. func Init(raw any) error { parsed, err := parseConfig(raw) if err != nil { @@ -125,22 +156,36 @@ func SetConfigForTest(c Config) { } func parseConfig(raw any) (Config, error) { - if raw == nil { - return Config{}, nil - } - // Round-trip through JSON so we accept whatever map shape - // Bifrost decoded the block into. No reflection on field names. - buf, err := json.Marshal(raw) - if err != nil { - return Config{}, err + cfg := Config{EnforceMacaroonsSource: "default"} + if raw != nil { + // Round-trip through JSON so we accept whatever map shape + // Bifrost decoded the block into. No reflection on field names. + buf, err := json.Marshal(raw) + if err != nil { + return Config{}, err + } + var envelope pluginConfigEnvelope + if err := json.Unmarshal(buf, &envelope); err != nil { + return Config{}, err + } + cfg.AgentBudgets = envelope.AgentBudgets + cfg.ModelPricing = envelope.ModelPricing + if envelope.EnforceMacaroons != nil { + cfg.EnforceMacaroons = *envelope.EnforceMacaroons + cfg.EnforceMacaroonsSource = "config" + } } - var env pluginConfigEnvelope - if err := json.Unmarshal(buf, &env); err != nil { - return Config{}, err + // Env override wins over the config block. A value we can't + // parse keeps the plugin alive on the config value (see Init for + // why fatal would be worse) but must be impossible to miss: an + // ERROR line here and "source=env-invalid" on the boot line. + if v, set, err := env.EnforceMacaroonsValue(); err != nil { + pluginlog.Errf("auth: %v — ignoring the override; enforce_macaroons=%t from %s stands", + err, cfg.EnforceMacaroons, cfg.EnforceMacaroonsSource) + cfg.EnforceMacaroonsSource = "env-invalid" + } else if set { + cfg.EnforceMacaroons = v + cfg.EnforceMacaroonsSource = "env" } - return Config{ - EnforceMacaroons: env.EnforceMacaroons, - AgentBudgets: env.AgentBudgets, - ModelPricing: env.ModelPricing, - }, nil + return cfg, nil } diff --git a/gateway/internal/auth/config_test.go b/gateway/internal/auth/config_test.go index f66a44969..1c9f59139 100644 --- a/gateway/internal/auth/config_test.go +++ b/gateway/internal/auth/config_test.go @@ -2,6 +2,8 @@ package auth import ( "testing" + + "github.com/stakwork/stakgraph/gateway/internal/env" ) func TestInit_NilConfig_DefaultsToShadow(t *testing.T) { @@ -61,3 +63,101 @@ func TestGetConfig_BeforeInit_ReturnsZero(t *testing.T) { t.Fatalf("pre-Init should be shadow mode, got %+v", got) } } + +// --- BIFROST_PLUGIN_ENFORCE_MACAROONS override --------------------------- + +func TestInit_Source_ConfigVsDefault(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "") + + if err := Init(map[string]any{"log_level": "info"}); err != nil { + t.Fatalf("Init: %v", err) + } + if got := GetConfig(); got.EnforceMacaroons || got.EnforceMacaroonsSource != "default" { + t.Fatalf("absent key: got enforce=%v source=%q, want false/default", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } + + if err := Init(map[string]any{"enforce_macaroons": false}); err != nil { + t.Fatalf("Init: %v", err) + } + if got := GetConfig(); got.EnforceMacaroons || got.EnforceMacaroonsSource != "config" { + t.Fatalf("explicit false: got enforce=%v source=%q, want false/config", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } +} + +func TestInit_EnvOverride_TrueOverConfigFalse(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "true") + + if err := Init(map[string]any{"enforce_macaroons": false}); err != nil { + t.Fatalf("Init: %v", err) + } + if got := GetConfig(); !got.EnforceMacaroons || got.EnforceMacaroonsSource != "env" { + t.Fatalf("got enforce=%v source=%q, want true/env", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } +} + +func TestInit_EnvOverride_FalseOverConfigTrue(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "0") + + if err := Init(map[string]any{"enforce_macaroons": true}); err != nil { + t.Fatalf("Init: %v", err) + } + if got := GetConfig(); got.EnforceMacaroons || got.EnforceMacaroonsSource != "env" { + t.Fatalf("got enforce=%v source=%q, want false/env", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } +} + +func TestInit_EnvOverride_AppliesToNilConfig(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "yes") + + if err := Init(nil); err != nil { + t.Fatalf("Init(nil): %v", err) + } + if got := GetConfig(); !got.EnforceMacaroons || got.EnforceMacaroonsSource != "env" { + t.Fatalf("nil config + env: got enforce=%v source=%q, want true/env", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } +} + +func TestInit_EnvOverride_PreservesOtherFields(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "true") + + raw := map[string]any{ + "enforce_macaroons": false, + "agent_budgets": map[string]any{"coder": map[string]any{"cap_usd": 5, "window": "1d"}}, + "model_pricing": map[string]any{"m": map[string]any{"input_per_mtok": 1, "output_per_mtok": 2}}, + } + if err := Init(raw); err != nil { + t.Fatalf("Init: %v", err) + } + got := GetConfig() + if !got.EnforceMacaroons || got.AgentBudgets["coder"].CapUSD != 5 || got.ModelPricing["m"].OutputPerMTok != 2 { + t.Fatalf("override must not drop sibling fields: %+v", got) + } +} + +// An unparseable override must NOT fail Init: a plugin that fails to +// load leaves bifrost-http serving with no verification at all, which +// is worse than either mode. It falls back to the config value and +// flags itself via the source. +func TestInit_EnvOverride_GarbageFallsBackToConfig(t *testing.T) { + t.Cleanup(func() { SetConfigForTest(Config{}) }) + t.Setenv(env.EnforceMacaroons, "ture") + + if err := Init(map[string]any{"enforce_macaroons": true}); err != nil { + t.Fatalf("Init must not fail on a bad override: %v", err) + } + if got := GetConfig(); !got.EnforceMacaroons || got.EnforceMacaroonsSource != "env-invalid" { + t.Fatalf("config true + garbage env: got enforce=%v source=%q, want true/env-invalid", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } + + if err := Init(nil); err != nil { + t.Fatalf("Init(nil) must not fail on a bad override: %v", err) + } + if got := GetConfig(); got.EnforceMacaroons || got.EnforceMacaroonsSource != "env-invalid" { + t.Fatalf("nil config + garbage env: got enforce=%v source=%q, want false/env-invalid", got.EnforceMacaroons, got.EnforceMacaroonsSource) + } +} diff --git a/gateway/internal/auth/doc.go b/gateway/internal/auth/doc.go index c2857be92..0eb305337 100644 --- a/gateway/internal/auth/doc.go +++ b/gateway/internal/auth/doc.go @@ -46,7 +46,12 @@ // // With enforce_macaroons=true the failure path becomes 401/402 with // a stable AdapterError.Code. Operators flip the flag per-swarm once -// the shadow-mode logs show no false positives. See +// the shadow-mode logs show no false positives — either in the +// config.json plugin block or, without rebuilding the image, via the +// BIFROST_PLUGIN_ENFORCE_MACAROONS env var (which wins when set; an +// unparseable value is logged at ERROR and ignored, because a plugin +// that fails Init leaves bifrost-http serving with no verification at +// all — see Init). See // gateway/plans/phases/phase-4-macaroon-shape.md ("Verifier // algorithm → Bifrost-plugin adapter"). // diff --git a/gateway/internal/env/env.go b/gateway/internal/env/env.go index ed5782763..fa9b411c9 100644 --- a/gateway/internal/env/env.go +++ b/gateway/internal/env/env.go @@ -14,6 +14,7 @@ package env import ( + "fmt" "os" "strings" ) @@ -74,6 +75,18 @@ const ( // logs.db on the same data volume. TrustPath = "BIFROST_PLUGIN_TRUST_PATH" + // EnforceMacaroons overrides the plugin config block's + // `enforce_macaroons` flag (gateway/data/config.json). config.json + // is baked into the image and re-seeded on every boot, so without + // this there is no way to flip a single swarm from shadow to + // enforce mode short of shipping a new image. Truthy: 1/true/yes/on; + // falsy: 0/false/no/off (case-insensitive). Unset or empty ⇒ the + // config.json value stands. Any other value is logged at ERROR and + // ignored (config.json value stands; boot line reports + // source=env-invalid) — see internal/auth.Init for why that beats + // failing the plugin. + EnforceMacaroons = "BIFROST_PLUGIN_ENFORCE_MACAROONS" + // RedisURL is the connection string for the macaroon-enforcement // Redis. In sphinx-swarm this points at the shared redis.sphinx // instance; in docker-compose it points at the sidecar `redis` @@ -236,6 +249,25 @@ func RedisURLValue() (string, bool) { return u, u != "" } +// EnforceMacaroonsValue parses BIFROST_PLUGIN_ENFORCE_MACAROONS. +// set=false when the variable is unset or empty (caller keeps the +// config.json value). err is non-nil for any value outside the +// recognised truthy/falsy sets — the caller decides what to do with +// a typo; this package never guesses which way it was meant. +func EnforceMacaroonsValue() (value bool, set bool, err error) { + raw := strings.TrimSpace(os.Getenv(EnforceMacaroons)) + if raw == "" { + return false, false, nil + } + switch strings.ToLower(raw) { + case "1", "true", "yes", "on": + return true, true, nil + case "0", "false", "no", "off": + return false, true, nil + } + return false, true, fmt.Errorf("%s=%q: want one of 1/true/yes/on or 0/false/no/off", EnforceMacaroons, raw) +} + // IsProduction reports whether the plugin is running in a // production-like environment. Drives a small handful of // security-defaults (currently: forcing `Secure` on the session diff --git a/gateway/internal/env/env_test.go b/gateway/internal/env/env_test.go new file mode 100644 index 000000000..cc73f4bde --- /dev/null +++ b/gateway/internal/env/env_test.go @@ -0,0 +1,42 @@ +package env + +import "testing" + +func TestEnforceMacaroonsValue(t *testing.T) { + cases := []struct { + raw string + wantVal bool + wantSet bool + wantError bool + }{ + {"", false, false, false}, + {" ", false, false, false}, + {"true", true, true, false}, + {"TRUE", true, true, false}, + {"1", true, true, false}, + {"yes", true, true, false}, + {"on", true, true, false}, + {" true ", true, true, false}, + {"false", false, true, false}, + {"0", false, true, false}, + {"no", false, true, false}, + {"OFF", false, true, false}, + // Typos must surface as errors, never as a silent default. + {"ture", false, true, true}, + {"enabled", false, true, true}, + {"2", false, true, true}, + } + for _, c := range cases { + t.Setenv(EnforceMacaroons, c.raw) + val, set, err := EnforceMacaroonsValue() + if (err != nil) != c.wantError { + t.Fatalf("%q: err=%v, wantError=%v", c.raw, err, c.wantError) + } + if set != c.wantSet { + t.Fatalf("%q: set=%v, want %v", c.raw, set, c.wantSet) + } + if err == nil && val != c.wantVal { + t.Fatalf("%q: value=%v, want %v", c.raw, val, c.wantVal) + } + } +} diff --git a/gateway/main.go b/gateway/main.go index 610a25ef9..505d458e5 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -54,8 +54,10 @@ const PluginName = "stakgraph-gateway" // and falls back to observability mode. See // gateway/plans/phases/phase-6-plugin-enforcement.md "Namespace". // 5. auth.Init + auth.SetTrustRegistry — parses the plugin's -// enforce_macaroons flag from the config block and wires the -// trust registry into the verifier. See +// enforce_macaroons flag from the config block (overridable via +// BIFROST_PLUGIN_ENFORCE_MACAROONS; an unparseable value logs an +// ERROR and is ignored) and wires the trust registry into the +// verifier. See // gateway/plans/phases/phase-4-macaroon-shape.md ("Bifrost- // plugin adapter"). // 6. adminapi.Start — boots the loopback HTTP server. @@ -87,7 +89,8 @@ func Init(config any) error { return err } auth.SetTrustRegistry(reg) - pluginlog.Logf("auth: macaroon adapter wired enforce=%t", auth.GetConfig().EnforceMacaroons) + authCfg := auth.GetConfig() + pluginlog.Logf("auth: macaroon adapter wired enforce=%t source=%s", authCfg.EnforceMacaroons, authCfg.EnforceMacaroonsSource) // Model-price catalog for the phase-6 accumulator: loads the // persisted datasheet, then fetches bifrost's published sheet in diff --git a/gateway/scripts/smoke-test-enforcement.sh b/gateway/scripts/smoke-test-enforcement.sh index 58129c596..37b5a381f 100755 --- a/gateway/scripts/smoke-test-enforcement.sh +++ b/gateway/scripts/smoke-test-enforcement.sh @@ -26,10 +26,14 @@ # in the bifrost container logs # - a logs row written with the dim-header values present # +# To exercise enforce mode instead, start the gateway with +# BIFROST_PLUGIN_ENFORCE_MACAROONS=true (env override of the config.json +# flag; see docker-compose.yml) — a bad or missing macaroon then 401s +# while the good-macaroon call below still 200s. +# # Later iterations of this script will add: # - attenuation chain (parent → child macaroons) -# - enforce-mode (BIFROST_PLUGIN_ENFORCE_MACAROONS=true → 401 on bad -# macaroon, 200 on good) +# - enforce-mode assertions (401 on bad macaroon, 200 on good) # - per-run cost cap exceeded (Redis-side accumulator check) # - kill switch (POST /_plugin/runs/:id/kill → next call 402s) # diff --git a/gateway/scripts/smoke-test-phase-11.sh b/gateway/scripts/smoke-test-phase-11.sh index 4fa81df3b..eb519f71c 100755 --- a/gateway/scripts/smoke-test-phase-11.sh +++ b/gateway/scripts/smoke-test-phase-11.sh @@ -30,7 +30,8 @@ # the local docker-compose); the log line tells us what the adapter # would have done in enforce mode. Promoting these checks to "block # the request and assert HTTP status" is a one-line flip of -# enforce_macaroons in gateway/data/config.json. +# enforce_macaroons in gateway/data/config.json, or start the gateway +# with BIFROST_PLUGIN_ENFORCE_MACAROONS=true). # # Usage # ----- From ce4b2b3b0dac6130afbf50593848afac65a5a827 Mon Sep 17 00:00:00 2001 From: Evan Feenstra Date: Thu, 10 Sep 2026 09:27:41 -0700 Subject: [PATCH 5/7] aieo: provider-agnostic web_fetch (HTTP shim off Anthropic) + aieo@0.1.37 (#1675) Sibling of the web_search shim (#1637). Anthropic keeps its native, server-executed web_fetch tool; every other provider gets a client-executed tool of the same name and result shape backed by a guarded HTTP GET plus a dependency-free HTML-to-text pass. The HTTP path runs in our process, so every URL and every redirect hop is validated before connecting: http(s) only, no embedded credentials, localhost refused, and the host must resolve exclusively to public unicast addresses (loopback, RFC 1918, link-local incl. cloud metadata, CGNAT, ULA, NAT64/6to4 and v4-mapped v6 all refused; a mixed answer is refused). allowedDomains/blockedDomains match the host and subdomains. Bodies are capped at 4 MiB and text at maxCharacters (default 40k); maxUses is enforced in-process on the shim and passed through as max_uses to Anthropic. Failures come back to the model as a readable tool error, with undici's transport cause surfaced. getProviderTool gains "webFetch"; createWebFetch returns the same handle shape as createWebSearch (tool / backend / native / results / capture). 16 unit tests cover the guard, redirects, extraction, budgets and native-result capture; try-fetch exercises both backends. --- mcp/src/aieo/package.json | 5 +- mcp/src/aieo/src/__tests__/fetch.test.ts | 423 ++++++++++++++ mcp/src/aieo/src/fetch.ts | 679 +++++++++++++++++++++++ mcp/src/aieo/src/index.ts | 1 + mcp/src/aieo/src/test/try-fetch.ts | 68 +++ mcp/src/aieo/src/tools.ts | 23 +- 6 files changed, 1188 insertions(+), 11 deletions(-) create mode 100644 mcp/src/aieo/src/__tests__/fetch.test.ts create mode 100644 mcp/src/aieo/src/fetch.ts create mode 100644 mcp/src/aieo/src/test/try-fetch.ts diff --git a/mcp/src/aieo/package.json b/mcp/src/aieo/package.json index d91227b7b..5971f8b1a 100644 --- a/mcp/src/aieo/package.json +++ b/mcp/src/aieo/package.json @@ -1,6 +1,6 @@ { "name": "aieo", - "version": "0.1.36", + "version": "0.1.37", "main": "./dist/index.js", "type": "module", "module": "./dist/index.mjs", @@ -12,12 +12,13 @@ "repository": "https://github.com/stakwork/stakgraph", "scripts": { "build": "tsup", - "test": "tsx ./src/__tests__/search.test.ts && tsx ./src/__tests__/provider.test.ts", + "test": "tsx ./src/__tests__/search.test.ts && tsx ./src/__tests__/fetch.test.ts && tsx ./src/__tests__/provider.test.ts", "try": "ts-node ./src/test/try.ts", "try-anthropic": "ts-node ./src/test/try-anthropic.ts", "try-kimi": "ts-node ./src/test/try-kimi.ts", "try-grok": "tsx ./src/test/try-grok.ts", "try-search": "tsx ./src/test/try-search.ts", + "try-fetch": "tsx ./src/test/try-fetch.ts", "cite-rate": "tsx ./src/test/cite-rate.ts" }, "keywords": [ diff --git a/mcp/src/aieo/src/__tests__/fetch.test.ts b/mcp/src/aieo/src/__tests__/fetch.test.ts new file mode 100644 index 000000000..7b4095d8e --- /dev/null +++ b/mcp/src/aieo/src/__tests__/fetch.test.ts @@ -0,0 +1,423 @@ +import { + createWebFetch, + validateFetchUrl, + isPrivateAddress, + htmlToText, + resolveFetchBackend, + captureNativeFetchResults, + WEB_FETCH_TOOL_NAME, + type WebFetchResult, +} from "../fetch.js"; + +type TestCase = { label: string; run: () => Promise | void }; + +function assert(cond: unknown, msg: string): void { + if (!cond) throw new Error(msg); +} + +async function rejects(fn: () => Promise, pattern: RegExp, label: string): Promise { + try { + await fn(); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + assert(pattern.test(msg), `${label}: rejected, but message "${msg}" !~ ${pattern}`); + return; + } + throw new Error(`${label}: expected a rejection`); +} + +const PUBLIC = async () => ["93.184.216.34"]; +const PRIVATE = async () => ["10.0.0.5"]; +const MIXED = async () => ["93.184.216.34", "10.0.0.5"]; + +/** Stub global fetch; returns the list of URLs it was called with. */ +function stubFetch(handler: (url: string) => Response | Promise): string[] { + const calls: string[] = []; + (globalThis as any).fetch = async (url: unknown) => { + calls.push(String(url)); + return handler(String(url)); + }; + return calls; +} + +const html = (body: string, headers: Record = {}) => + new Response(body, { + status: 200, + headers: { "content-type": "text/html; charset=utf-8", ...headers }, + }); + +const redirect = (to: string, status = 302) => + new Response(null, { status, headers: { location: to } }); + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +const exec = (t: any, input: any) => t.execute(input, { toolCallId: "t", messages: [] }); + +const PAGE = ` + + Hello & World + + + + + +

Heading

+

Tom & Jerry — “quoted” 'x' done

+
  • one
  • two
+ vector junk +

Last
line

+`; + +const tests: TestCase[] = [ + { + label: "anthropic keeps the native tool; others fall to http", + run() { + assert(resolveFetchBackend("anthropic") === "anthropic", "anthropic → anthropic"); + for (const p of ["openai", "google", "openrouter", "xai"] as const) { + assert(resolveFetchBackend(p) === "http", `${p} → http`); + } + }, + }, + { + label: "anthropic path: no key → tool undefined; with key → native tool", + run() { + delete process.env.ANTHROPIC_API_KEY; + const none = createWebFetch({ provider: "anthropic" }); + assert(none.tool === undefined, "tool should be undefined without a key"); + assert(none.backend === undefined, "backend should be undefined without a key"); + const native = createWebFetch({ provider: "anthropic", apiKey: "sk-ant-fake" }); + assert(!!native.tool, "tool present with a key"); + assert(native.native === true, "anthropic is native"); + assert(native.backend === "anthropic", "backend is anthropic"); + }, + }, + { + label: "http path needs no key", + run() { + delete process.env.ANTHROPIC_API_KEY; + const wf = createWebFetch({ provider: "openai" }); + assert(!!wf.tool, "tool present"); + assert(wf.native === false, "shim is not native"); + assert(wf.backend === "http", "backend is http"); + assert(WEB_FETCH_TOOL_NAME === "web_fetch", "tool name matches Anthropic's native name"); + }, + }, + { + label: "isPrivateAddress classifies v4, v6 and mapped literals", + run() { + const blocked = [ + "127.0.0.1", "10.1.2.3", "172.16.0.1", "172.31.255.255", "192.168.1.1", + "169.254.169.254", "100.64.0.1", "0.0.0.0", "255.255.255.255", "224.0.0.1", + "::1", "::", "fe80::1", "fc00::1", "fd12:3456::1", "ff02::1", + "::ffff:127.0.0.1", "::ffff:10.0.0.1", "[::1]", "64:ff9b::a00:1", "2002:a00:1::1", + "not-an-ip", "", + ]; + for (const ip of blocked) assert(isPrivateAddress(ip) === true, `${ip} should be private`); + const allowed = ["8.8.8.8", "93.184.216.34", "172.32.0.1", "100.128.0.1", "2606:4700::1111", "::ffff:8.8.8.8"]; + for (const ip of allowed) assert(isPrivateAddress(ip) === false, `${ip} should be public`); + }, + }, + { + label: "validateFetchUrl rejects bad schemes, credentials, local names and private literals", + async run() { + const cases: Array<[string, RegExp]> = [ + ["file:///etc/passwd", /http\(s\)/], + ["ftp://example.com/x", /http\(s\)/], + ["javascript:alert(1)", /http\(s\)|valid absolute/], + ["not a url", /valid absolute/], + ["https://user:pw@example.com/", /credentials/], + ["http://localhost:3000/", /local/], + ["http://api.localhost/", /local/], + ["http://127.0.0.1/", /private or reserved/], + ["http://[::1]/", /private or reserved/], + ["http://2130706433/", /private or reserved/], + ["http://0x7f000001/", /private or reserved/], + ["http://169.254.169.254/latest/meta-data", /private or reserved/], + ["http://10.0.0.1:8080/", /private or reserved/], + ]; + for (const [url, pattern] of cases) { + await rejects(() => validateFetchUrl(url, { lookup: PUBLIC }), pattern, url); + } + }, + }, + { + label: "validateFetchUrl resolves hostnames and refuses private or mixed answers", + async run() { + const ok = await validateFetchUrl("https://example.com/page", { lookup: PUBLIC }); + assert(ok.hostname === "example.com", "public host passes"); + await rejects( + () => validateFetchUrl("https://intranet.example.com/", { lookup: PRIVATE }), + /resolves to a private/, + "private resolution", + ); + await rejects( + () => validateFetchUrl("https://mixed.example.com/", { lookup: MIXED }), + /resolves to a private/, + "mixed resolution", + ); + await rejects( + () => validateFetchUrl("https://nope.example.com/", { lookup: async () => [] }), + /Could not resolve/, + "empty resolution", + ); + await rejects( + () => validateFetchUrl("https://nx.example.com/", { lookup: async () => { throw new Error("ENOTFOUND"); } }), + /Could not resolve .*ENOTFOUND/, + "resolver error", + ); + const literal = await validateFetchUrl("http://93.184.216.34/", { + lookup: async () => { throw new Error("should not resolve a literal"); }, + }); + assert(literal.hostname === "93.184.216.34", "public literal passes without DNS"); + }, + }, + { + label: "allowed and blocked domains match the host and its subdomains", + async run() { + const allow = { lookup: PUBLIC, allowedDomains: ["example.com", "https://docs.other.org/"] }; + await validateFetchUrl("https://example.com/", allow); + await validateFetchUrl("https://deep.docs.example.com/", allow); + await validateFetchUrl("https://docs.other.org/x", allow); + await rejects(() => validateFetchUrl("https://notexample.com/", allow), /not in the allowed/, "suffix trick"); + await rejects(() => validateFetchUrl("https://other.org/", allow), /not in the allowed/, "parent of allowed sub"); + const block = { lookup: PUBLIC, blockedDomains: ["evil.com"] }; + await validateFetchUrl("https://fine.com/", block); + await rejects(() => validateFetchUrl("https://sub.evil.com/", block), /blocked/, "blocked subdomain"); + }, + }, + { + label: "http path converts html to text and records the page", + async run() { + stubFetch(() => html(PAGE)); + const wf = createWebFetch({ provider: "openai", lookup: PUBLIC }); + const out = (await exec(wf.tool, { url: "https://example.com/page" })) as WebFetchResult; + assert(out.type === "web_fetch_result", `type: ${JSON.stringify(out)}`); + assert(out.title === "Hello & World", `title: ${out.title}`); + assert(out.text!.includes("Tom & Jerry — “quoted” 'x' done"), `entities: ${out.text}`); + assert(!out.text!.includes("alert"), "script removed"); + assert(!out.text!.includes("color"), "style removed"); + assert(!out.text!.includes("vector junk"), "svg removed"); + assert(!out.text!.includes("comment"), "comment removed"); + assert(out.text!.includes("- one\n- two"), `list items: ${out.text}`); + assert(out.text!.includes("Last\nline"), `br: ${out.text}`); + assert(out.text!.startsWith("Heading"), `leading: ${out.text!.slice(0, 20)}`); + assert(out.mediaType === "text/html", `mediaType: ${out.mediaType}`); + assert(out.truncated === false, "not truncated"); + assert(out.url === "https://example.com/page", `url: ${out.url}`); + assert(!Number.isNaN(Date.parse(out.retrievedAt!)), "retrievedAt is a date"); + assert(wf.results.length === 1 && wf.results[0] === out, "recorded in results"); + }, + }, + { + label: "htmlToText shapes blocks, lists and breaks", + run() { + const { title, text } = htmlToText("

Title

a
b

  • x
  • y
"); + assert(title === null, "no title"); + assert(text === "Title\n\na\nb\n\n- x\n- y", `got: ${JSON.stringify(text)}`); + assert(htmlToText("&lt; � &#xZZ;").text === "< � &#xZZ;", "double-encoded and invalid entities survive"); + }, + }, + { + label: "text budget truncates and flags; maxContentTokens alone derives it", + async run() { + const body = `

${"x".repeat(500)}

`; + stubFetch(() => html(body)); + const chars = createWebFetch({ provider: "openai", lookup: PUBLIC, maxCharacters: 50 }); + const a = (await exec(chars.tool, { url: "https://example.com/" })) as WebFetchResult; + assert(a.text!.length === 50 && a.truncated === true, `chars: len=${a.text!.length} truncated=${a.truncated}`); + const tokens = createWebFetch({ provider: "openai", lookup: PUBLIC, maxContentTokens: 10 }); + const b = (await exec(tokens.tool, { url: "https://example.com/" })) as WebFetchResult; + assert(b.text!.length === 40 && b.truncated === true, `tokens: len=${b.text!.length}`); + }, + }, + { + label: "json and plain text pass through; html served as text/plain is sniffed", + async run() { + stubFetch(() => new Response('{"a":1}', { status: 200, headers: { "content-type": "application/json" } })); + const wf = createWebFetch({ provider: "google", lookup: PUBLIC }); + const j = (await exec(wf.tool, { url: "https://api.example.com/x" })) as WebFetchResult; + assert(j.text === '{"a":1}' && j.title === null && j.mediaType === "application/json", `json: ${JSON.stringify(j)}`); + + stubFetch(() => new Response("Thi", { + status: 200, headers: { "content-type": "text/plain" }, + })); + const s = (await exec(wf.tool, { url: "https://example.com/raw" })) as WebFetchResult; + assert(s.title === "T" && s.text === "hi", `sniffed: ${JSON.stringify(s)}`); + + stubFetch(() => new Response(new Uint8Array([0xe9, 0x74, 0xe9]), { + status: 200, headers: { "content-type": "text/plain; charset=iso-8859-1" }, + })); + const l = (await exec(wf.tool, { url: "https://example.com/latin" })) as WebFetchResult; + assert(l.text === "été", `charset honored: ${l.text}`); + }, + }, + { + label: "redirects are followed hop by hop and every hop is re-validated", + async run() { + const calls = stubFetch((url) => + url === "https://example.com/start" ? redirect("/next") + : url === "https://example.com/next" ? redirect("https://final.example.org/page", 301) + : html("Final"), + ); + const wf = createWebFetch({ provider: "xai", lookup: PUBLIC }); + const out = (await exec(wf.tool, { url: "https://example.com/start" })) as WebFetchResult; + assert(out.url === "https://final.example.org/page", `final url: ${out.url}`); + assert(out.title === "Final", "landed on the final page"); + assert(calls.length === 3, `hops: ${calls.length}`); + + stubFetch(() => redirect("http://169.254.169.254/latest/meta-data")); + const ssrf = (await exec(wf.tool, { url: "https://example.com/open-redirect" })) as { error?: string }; + assert(/private or reserved/.test(ssrf.error ?? ""), `redirect to metadata blocked: ${ssrf.error}`); + + const resolver = async (host: string) => (host === "internal.example.com" ? ["10.0.0.9"] : ["93.184.216.34"]); + stubFetch(() => redirect("https://internal.example.com/")); + const wf2 = createWebFetch({ provider: "xai", lookup: resolver }); + const hop = (await exec(wf2.tool, { url: "https://example.com/" })) as { error?: string }; + assert(/resolves to a private/.test(hop.error ?? ""), `redirect to private host blocked: ${hop.error}`); + + let n = 0; + stubFetch(() => redirect(`https://example.com/loop${n++}`)); + const loop = (await exec(wf.tool, { url: "https://example.com/loop" })) as { error?: string }; + assert(/Too many redirects/.test(loop.error ?? ""), `loop: ${loop.error}`); + + stubFetch(() => new Response(null, { status: 302 })); + const bare = (await exec(wf.tool, { url: "https://example.com/bare" })) as { error?: string }; + assert(/without a Location/.test(bare.error ?? ""), `bare redirect: ${bare.error}`); + }, + }, + { + label: "http errors, unsupported types and pdfs come back as errors, not throws", + async run() { + // Seven failing calls below; keep them under the budget so the last + // ones don't trip the maxUses guard instead of the case under test. + const wf = createWebFetch({ provider: "openai", lookup: PUBLIC, maxUses: 20 }); + stubFetch(() => new Response("gone", { status: 404 })); + const nf = (await exec(wf.tool, { url: "https://example.com/missing" })) as { error?: string }; + assert(/HTTP 404/.test(nf.error ?? ""), `404: ${nf.error}`); + + stubFetch(() => new Response(new Uint8Array([1, 2, 3]), { status: 200, headers: { "content-type": "image/png" } })); + const img = (await exec(wf.tool, { url: "https://example.com/a.png" })) as { error?: string }; + assert(/Unsupported content type: image\/png/.test(img.error ?? ""), `png: ${img.error}`); + + stubFetch(() => new Response("%PDF-1.4", { status: 200, headers: { "content-type": "application/pdf" } })); + const pdf = (await exec(wf.tool, { url: "https://example.com/a.pdf" })) as { error?: string }; + assert(/PDF/.test(pdf.error ?? ""), `pdf: ${pdf.error}`); + + (globalThis as any).fetch = async () => { + throw Object.assign(new TypeError("fetch failed"), { + cause: Object.assign(new Error("certificate has expired"), { code: "CERT_HAS_EXPIRED" }), + }); + }; + const down = (await exec(wf.tool, { url: "https://example.com/" })) as { error?: string }; + assert( + /Could not connect to example\.com: certificate has expired \(CERT_HAS_EXPIRED\)/.test(down.error ?? ""), + `transport cause surfaced: ${down.error}`, + ); + + (globalThis as any).fetch = async () => { + throw new DOMException("The operation was aborted due to timeout", "TimeoutError"); + }; + const slow = (await exec(wf.tool, { url: "https://example.com/" })) as { error?: string }; + assert(/timed out/.test(slow.error ?? ""), `timeout: ${slow.error}`); + + (globalThis as any).fetch = async () => { throw new TypeError("fetch failed"); }; + const bare = (await exec(wf.tool, { url: "https://example.com/" })) as { error?: string }; + assert(/fetch failed/.test(bare.error ?? ""), `no cause falls back to message: ${bare.error}`); + + const bad = (await exec(wf.tool, { url: "file:///etc/passwd" })) as { error?: string }; + assert(/http\(s\)/.test(bad.error ?? ""), `guard surfaces to the model: ${bad.error}`); + assert(wf.results.length === 0, "failures record nothing"); + }, + }, + { + label: "maxUses is enforced in-process on the http path", + async run() { + stubFetch(() => html("a")); + const wf = createWebFetch({ provider: "openai", lookup: PUBLIC, maxUses: 1 }); + await exec(wf.tool, { url: "https://example.com/1" }); + const blocked = (await exec(wf.tool, { url: "https://example.com/2" })) as { error?: string }; + assert(/budget exhausted/.test(blocked.error ?? ""), `second call refused: ${blocked.error}`); + assert(wf.results.length === 1, "refused call adds no results"); + }, + }, + { + label: "capture() is a no-op on the http path (no double-count)", + async run() { + stubFetch(() => html("a")); + const wf = createWebFetch({ provider: "openai", lookup: PUBLIC }); + await exec(wf.tool, { url: "https://example.com/" }); + wf.capture([ + { + type: "tool-result", + toolName: WEB_FETCH_TOOL_NAME, + output: { type: "web_fetch_result", url: "https://example.com/", content: { title: "a", source: { type: "text", data: "a" } } }, + }, + ]); + assert(wf.results.length === 1, `expected 1 result, got ${wf.results.length}`); + }, + }, + { + label: "captureNativeFetchResults walks output and result shapes, skips errors and junk", + run() { + const target: WebFetchResult[] = []; + captureNativeFetchResults( + [ + { + type: "tool-result", + toolName: WEB_FETCH_TOOL_NAME, + output: { + type: "web_fetch_result", + url: "https://a.com", + retrievedAt: "2026-09-10T00:00:00Z", + content: { type: "document", title: "A", source: { type: "text", mediaType: "text/plain", data: "page text" } }, + }, + }, + { + type: "tool-result", + toolName: WEB_FETCH_TOOL_NAME, + result: { + type: "web_fetch_result", + url: "https://b.com/doc.pdf", + retrieved_at: null, + content: { type: "document", title: null, source: { type: "base64", media_type: "application/pdf", data: "JVBERi0=" } }, + }, + }, + { + type: "tool-result", + toolName: WEB_FETCH_TOOL_NAME, + isError: true, + output: { type: "web_fetch_tool_result_error", errorCode: "url_not_accessible" }, + }, + { type: "tool-result", toolName: "other_tool", output: { type: "web_fetch_result", url: "https://nope.com" } }, + { type: "tool-result", toolName: WEB_FETCH_TOOL_NAME, output: "not an object" }, + { type: "text", text: "hi" }, + ], + target, + ); + assert(target.length === 2, `expected 2, got ${target.length}`); + assert(target[0].text === "page text" && target[0].title === "A", "text source captured"); + assert(target[0].mediaType === "text/plain" && target[0].retrievedAt === "2026-09-10T00:00:00Z", "fields mapped"); + assert(target[1].text === undefined, "pdf base64 is not surfaced as text"); + assert(target[1].mediaType === "application/pdf" && target[1].title === null, "snake_case shape tolerated"); + captureNativeFetchResults("not an array", target); + assert(target.length === 2, "non-array content ignored"); + }, + }, +]; + +let passed = 0; +let failed = 0; + +for (const tc of tests) { + try { + await tc.run(); + console.log(`✅ PASS: ${tc.label}`); + passed++; + } catch (err: any) { + console.error(`❌ FAIL: ${tc.label}`); + console.error(` ${err.message}`); + failed++; + } +} + +console.log(`\nResults: ${passed} passed, ${failed} failed`); +if (failed > 0) process.exit(1); diff --git a/mcp/src/aieo/src/fetch.ts b/mcp/src/aieo/src/fetch.ts new file mode 100644 index 000000000..5b3cd557f --- /dev/null +++ b/mcp/src/aieo/src/fetch.ts @@ -0,0 +1,679 @@ +import { tool } from "ai"; +import { z } from "zod"; +import { createAnthropic } from "@ai-sdk/anthropic"; +import { BlockList, isIP } from "node:net"; +import { lookup as dnsLookup } from "node:dns/promises"; +import { Provider, getGatewayBaseURL, normalizeApiKey } from "./provider.js"; + +/** + * Web fetch, uniform across providers. Sibling of `./search.js`. + * + * Anthropic ships a server-executed `web_fetch` tool: given a URL, their + * API retrieves the page (HTML or PDF), turns it into text and hands it + * to the model with no round-trip through us. Same reasons that make it + * the best option when available: no extra hop, nothing to host, + * nothing to secure. + * + * Nobody else has one. So, as with search: keep Anthropic native, and + * give every other provider a client-executed tool of the same name and + * result shape, backed by a plain HTTP GET plus an HTML-to-text pass. + * + * Two things differ from search that consumers should know: + * + * 1. On the Anthropic path the model can only fetch URLs that already + * appeared in the conversation — pasted by the user, or returned by + * an earlier `web_search` / `web_fetch`. It refuses URLs it made + * up. The HTTP path has no such memory and fetches whatever the + * model asks for; what it WON'T do is reach anything private. + * + * 2. The HTTP path runs in *our* process, so the model can point it at + * our network. Every URL — and every redirect hop — is checked + * before connecting: http(s) only, no credentials, and the host + * must resolve exclusively to public unicast addresses (loopback, + * RFC 1918, link-local incl. cloud metadata, CGNAT, ULA and the + * v4-in-v6 forms are all refused). There is a DNS-rebinding window + * between our lookup and the socket's own; pin `allowedDomains` if + * the deployment cares. + * + * Usage: + * + * const wf = createWebFetch({ provider, apiKey }); + * const tools = { ...(wf.tool ? { [WEB_FETCH_TOOL_NAME]: wf.tool } : {}) }; + * // in onStepFinish: wf.capture(step.content) + * // afterwards: wf.results — every page fetched, in order + */ + +/** Tool name registered with the model. Matches Anthropic's native name + * so consumers key UI and step-walking off one string on both paths. */ +export const WEB_FETCH_TOOL_NAME = "web_fetch"; + +/** Which implementation backs the tool. */ +export type FetchBackend = "anthropic" | "http"; + +/** + * One fetched page. Same fields on both paths; the notes say where the + * backends differ in what they put in them. + */ +export interface WebFetchResult { + /** Where the content came from. The final URL after redirects on the + * HTTP path; the URL Anthropic reports on the native path. */ + url: string; + title: string | null; + /** + * Extracted text. Absent for a PDF on the Anthropic path — their API + * returns it base64-encoded, which is only useful to their model. + */ + text?: string; + /** Original `content-type` on the HTTP path; Anthropic's normalized + * `text/plain` / `application/pdf` on the native path. */ + mediaType: string | null; + retrievedAt: string | null; + /** True when the HTTP path cut the text at `maxCharacters`. */ + truncated?: boolean; + type: "web_fetch_result"; +} + +export interface WebFetchOptions { + /** Max `web_fetch` calls per run. Default 5. Enforced in-process on + * the HTTP path, passed as `max_uses` to Anthropic. */ + maxUses?: number; + /** + * Text budget per page on the HTTP path, in characters. Default + * 40000 — about 10k tokens. When only `maxContentTokens` is given, + * derived from it at 4 chars/token so one option covers both paths. + */ + maxCharacters?: number; + /** + * Content budget per page on the Anthropic path, in tokens (their + * `max_content_tokens`). Unset means Anthropic's own default. When + * only `maxCharacters` is given, derived from it. + */ + maxContentTokens?: number; + /** Only fetch from these domains. Subdomains are included, so + * `example.com` admits `docs.example.com`. */ + allowedDomains?: string[]; + blockedDomains?: string[]; +} + +/** Resolve a hostname to every address it answers with. */ +export type HostLookup = (hostname: string) => Promise; + +export interface CreateWebFetchOptions extends WebFetchOptions { + /** LLM provider driving the run — decides the backend. */ + provider: Provider; + /** LLM API key. Only used on the Anthropic path. Falls back to env. */ + apiKey?: string; + /** Force a backend regardless of provider. `"http"` is how you A/B + * the shim against Anthropic's native tool on identical prompts. */ + backend?: FetchBackend; + abortSignal?: AbortSignal; + /** + * Override hostname resolution on the HTTP path. Tests inject a stub + * here; a deployment with its own resolver policy can too. Must return + * every address the host resolves to — the guard refuses a host if ANY + * of them is private. + */ + lookup?: HostLookup; +} + +export interface WebFetchHandle { + /** Register under {@link WEB_FETCH_TOOL_NAME}. `undefined` only when + * the Anthropic path has no API key — drop the tool rather than + * failing the request. The HTTP path needs no key. */ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + tool: any | undefined; + backend: FetchBackend | undefined; + /** True when the fetch runs server-side (Anthropic). */ + native: boolean; + /** Every page fetched during the run, in order. */ + results: WebFetchResult[]; + /** + * Feed each step's content here (AI SDK `onStepFinish`). Walks + * Anthropic tool-results into `results`; a no-op on the HTTP path, + * where `execute` already appended them. Safe to call either way. + */ + capture(stepContent: unknown): void; +} + +const DEFAULT_MAX_USES = 5; +const DEFAULT_MAX_CHARACTERS = 40_000; +const CHARS_PER_TOKEN = 4; +const DEFAULT_TIMEOUT_MS = 20_000; +const MAX_REDIRECTS = 5; +/** Raw bytes read from a response before giving up on it. HTML runs + * several times its text; 4 MiB covers any page whose text fits the + * default budget with room to spare. */ +const MAX_BODY_BYTES = 4 * 1024 * 1024; +const USER_AGENT = "aieo-web-fetch/1 (+https://github.com/stakwork/stakgraph)"; +const ACCEPT = + "text/html, application/xhtml+xml, text/plain;q=0.9, application/json;q=0.9, text/*;q=0.8, */*;q=0.5"; + +/** + * Which backend a provider gets. Anthropic keeps its native tool; + * everything else falls to the HTTP shim. + */ +export function resolveFetchBackend(provider: Provider): FetchBackend { + return provider === "anthropic" ? "anthropic" : "http"; +} + +// ── Address guard ────────────────────────────────────────────────────── + +/** + * Addresses the HTTP path never connects to. Built once. + * + * `BlockList` checks an IPv4-mapped IPv6 address (`::ffff:a.b.c.d`) + * against the v4 rules by itself. NAT64 and 6to4 embed a v4 address in + * a way it doesn't unpack, so those prefixes are refused whole — neither + * is a plausible target for a page fetch. + */ +const PRIVATE_ADDRESSES = buildBlockList(); + +function buildBlockList(): BlockList { + const b = new BlockList(); + const v4: Array<[string, number]> = [ + ["0.0.0.0", 8], // "this" network + ["10.0.0.0", 8], // RFC 1918 + ["100.64.0.0", 10], // carrier-grade NAT + ["127.0.0.0", 8], // loopback + ["169.254.0.0", 16], // link-local, incl. cloud metadata at 169.254.169.254 + ["172.16.0.0", 12], // RFC 1918 + ["192.0.0.0", 24], // IETF protocol assignments + ["192.168.0.0", 16], // RFC 1918 + ["198.18.0.0", 15], // benchmarking + ["224.0.0.0", 4], // multicast + ["240.0.0.0", 4], // reserved, incl. broadcast + ]; + for (const [net, prefix] of v4) b.addSubnet(net, prefix, "ipv4"); + b.addAddress("::", "ipv6"); // unspecified + b.addAddress("::1", "ipv6"); // loopback + const v6: Array<[string, number]> = [ + ["64:ff9b::", 96], // NAT64 + ["2002::", 16], // 6to4 + ["fc00::", 7], // unique local + ["fe80::", 10], // link-local + ["fec0::", 10], // site-local (deprecated, still routed on old gear) + ["ff00::", 8], // multicast + ]; + for (const [net, prefix] of v6) b.addSubnet(net, prefix, "ipv6"); + return b; +} + +/** + * True for an IP literal (v4 or v6, brackets tolerated) the HTTP path + * refuses to connect to. Anything that isn't an IP literal is `true` + * too: an address we can't classify isn't one we connect to. + */ +export function isPrivateAddress(ip: string): boolean { + const bare = ip.startsWith("[") && ip.endsWith("]") ? ip.slice(1, -1) : ip; + const family = isIP(bare); + try { + if (family === 4) return PRIVATE_ADDRESSES.check(bare, "ipv4"); + if (family === 6) return PRIVATE_ADDRESSES.check(bare, "ipv6"); + } catch { + // Zone ids and other oddities BlockList can't parse. + } + return true; +} + +const defaultLookup: HostLookup = async (hostname) => { + const addresses = await dnsLookup(hostname, { all: true }); + return addresses.map((a) => a.address); +}; + +function matchesDomain(host: string, domains: string[]): boolean { + return domains.some((d) => { + const dom = d + .trim() + .toLowerCase() + .replace(/^https?:\/\//, "") + .replace(/\/.*$/, "") + .replace(/^\*\./, "") + .replace(/\.$/, ""); + return !!dom && (host === dom || host.endsWith("." + dom)); + }); +} + +/** + * Validate a URL before the HTTP path connects to it. Throws a readable + * error on any rejection — the model gets the message back as the + * tool's `error` and can pick a different URL. + * + * A hostname is resolved and refused if ANY of its addresses is + * private: a host that answers with a mix is misconfigured or hostile, + * and we'd have no say in which address the socket picks. + */ +export async function validateFetchUrl( + raw: string, + opts: { allowedDomains?: string[]; blockedDomains?: string[]; lookup?: HostLookup } = {}, +): Promise { + let url: URL; + try { + url = new URL(raw); + } catch { + throw new Error(`Not a valid absolute URL: ${raw}`); + } + if (url.protocol !== "http:" && url.protocol !== "https:") { + throw new Error(`Only http(s) URLs can be fetched, got ${url.protocol}`); + } + if (url.username || url.password) { + throw new Error("URLs with embedded credentials are not fetched"); + } + // WHATWG parsing already normalized the numeric IPv4 forms + // (`http://2130706433/`, `http://0x7f.1/`) to dotted quads. + const host = url.hostname.replace(/\.$/, "").toLowerCase(); + if (!host) throw new Error("URL has no host"); + if (host === "localhost" || host.endsWith(".localhost")) { + throw new Error("Refusing to fetch a local address"); + } + if (opts.allowedDomains?.length && !matchesDomain(host, opts.allowedDomains)) { + throw new Error(`${host} is not in the allowed domains`); + } + if (opts.blockedDomains?.length && matchesDomain(host, opts.blockedDomains)) { + throw new Error(`${host} is a blocked domain`); + } + + const literal = host.startsWith("[") ? host.slice(1, -1) : host; + if (isIP(literal)) { + if (isPrivateAddress(literal)) { + throw new Error("Refusing to fetch a private or reserved address"); + } + return url; + } + + let addresses: string[]; + try { + addresses = await (opts.lookup ?? defaultLookup)(host); + } catch (err) { + throw new Error( + `Could not resolve ${host}: ${err instanceof Error ? err.message : String(err)}`, + ); + } + if (!addresses.length) throw new Error(`Could not resolve ${host}`); + if (addresses.some(isPrivateAddress)) { + throw new Error(`Refusing to fetch ${host}: it resolves to a private or reserved address`); + } + return url; +} + +// ── HTTP path ────────────────────────────────────────────────────────── + +export interface FetchUrlOptions { + maxCharacters?: number; + allowedDomains?: string[]; + blockedDomains?: string[]; + abortSignal?: AbortSignal; + timeoutMs?: number; + lookup?: HostLookup; +} + +/** + * Raw fetch-and-extract. Exposed for callers that want a page without + * an LLM in the loop (a URL enrichment pass, a link preview). + * + * Throws on any rejection or failure — {@link createWebFetch} catches + * and hands the model a readable error instead of failing the turn. + */ +export async function fetchUrl( + rawUrl: string, + opts: FetchUrlOptions = {}, +): Promise { + const maxCharacters = opts.maxCharacters ?? DEFAULT_MAX_CHARACTERS; + const signals = [AbortSignal.timeout(opts.timeoutMs ?? DEFAULT_TIMEOUT_MS)]; + if (opts.abortSignal) signals.push(opts.abortSignal); + const signal = AbortSignal.any(signals); + + let url = await validateFetchUrl(rawUrl, opts); + let res: Response; + for (let hop = 0; ; hop++) { + try { + res = await fetch(url, { + method: "GET", + redirect: "manual", + signal, + headers: { "user-agent": USER_AGENT, accept: ACCEPT }, + }); + } catch (err) { + throw new Error(`Could not connect to ${url.hostname}: ${describeTransportError(err)}`); + } + if (!isRedirect(res.status)) break; + // Every hop is re-validated: a public host that 302s to + // 169.254.169.254 is the classic SSRF bypass. + const location = res.headers.get("location"); + await discard(res); + if (!location) throw new Error(`Redirect (${res.status}) without a Location header`); + if (hop >= MAX_REDIRECTS) throw new Error(`Too many redirects (more than ${MAX_REDIRECTS})`); + let next: string; + try { + next = new URL(location, url).toString(); + } catch { + throw new Error(`Redirect to an invalid URL: ${location}`); + } + url = await validateFetchUrl(next, opts); + } + + if (!res.ok) { + await discard(res); + throw new Error(`HTTP ${res.status} fetching ${url.hostname}`); + } + const contentType = res.headers.get("content-type") ?? ""; + const mediaType = contentType.split(";")[0].trim().toLowerCase() || null; + if (mediaType === "application/pdf") { + await discard(res); + throw new Error( + "PDF documents are not supported by this fetch backend (only Anthropic's native web_fetch reads PDFs)", + ); + } + if (!isTextual(mediaType)) { + await discard(res); + throw new Error(`Unsupported content type: ${mediaType}`); + } + + const { bytes, capped } = await readCapped(res, MAX_BODY_BYTES); + const raw = decode(bytes, contentType); + const isHtml = + mediaType === "text/html" || mediaType === "application/xhtml+xml" || looksLikeHtml(raw); + const { title, text } = isHtml ? htmlToText(raw) : { title: null, text: raw.trim() }; + const truncated = capped || text.length > maxCharacters; + return { + url: url.toString(), + title, + text: truncated ? text.slice(0, maxCharacters) : text, + mediaType, + retrievedAt: new Date().toISOString(), + truncated, + type: "web_fetch_result", + }; +} + +/** + * undici wraps every transport failure as `TypeError: fetch failed` and + * puts the real reason on `cause` — an expired certificate, ECONNREFUSED, + * a reset. Surface that: "fetch failed" tells the model (and whoever + * reads the logs) nothing. + */ +function describeTransportError(err: unknown): string { + const e = err as { + name?: string; + message?: string; + cause?: { code?: string; message?: string }; + }; + if (e?.name === "TimeoutError") return "timed out"; + if (e?.name === "AbortError") return "aborted"; + const cause = e?.cause; + if (cause?.message) return cause.code ? `${cause.message} (${cause.code})` : cause.message; + return e?.message ?? String(err); +} + +function isRedirect(status: number): boolean { + return status === 301 || status === 302 || status === 303 || status === 307 || status === 308; +} + +/** Release a response we won't read, so its connection goes back to the pool. */ +async function discard(res: Response): Promise { + await res.body?.cancel().catch(() => {}); +} + +function isTextual(mediaType: string | null): boolean { + // No header at all: read it and sniff. + if (!mediaType) return true; + if (mediaType.startsWith("text/")) return true; + if (mediaType.endsWith("+json") || mediaType.endsWith("+xml")) return true; + return ( + mediaType === "application/json" || + mediaType === "application/xml" || + mediaType === "application/xhtml+xml" || + mediaType === "application/javascript" + ); +} + +async function readCapped( + res: Response, + maxBytes: number, +): Promise<{ bytes: Uint8Array; capped: boolean }> { + const body = res.body; + if (!body) return { bytes: new Uint8Array(await res.arrayBuffer()), capped: false }; + const reader = body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + let capped = false; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (!value) continue; + if (total + value.byteLength > maxBytes) { + chunks.push(value.subarray(0, maxBytes - total)); + total = maxBytes; + capped = true; + await reader.cancel().catch(() => {}); + break; + } + chunks.push(value); + total += value.byteLength; + } + const out = new Uint8Array(total); + let offset = 0; + for (const c of chunks) { + out.set(c, offset); + offset += c.byteLength; + } + return { bytes: out, capped }; +} + +function decode(bytes: Uint8Array, contentType: string): string { + const charset = /charset=["']?([\w.:-]+)/i.exec(contentType)?.[1] ?? "utf-8"; + try { + return new TextDecoder(charset).decode(bytes); + } catch { + return new TextDecoder("utf-8").decode(bytes); + } +} + +function looksLikeHtml(s: string): boolean { + return /^\s*(?:]|]|])/i.test(s.slice(0, 1024)); +} + +// ── HTML → text ──────────────────────────────────────────────────────── + +const NAMED_ENTITIES: Record = { + nbsp: " ", + lt: "<", + gt: ">", + quot: '"', + apos: "'", + mdash: "—", + ndash: "–", + hellip: "…", + copy: "©", + reg: "®", + trade: "™", + laquo: "«", + raquo: "»", + ldquo: "“", + rdquo: "”", + lsquo: "‘", + rsquo: "’", + bull: "•", + middot: "·", + times: "×", + deg: "°", +}; + +function codePoint(cp: number, fallback: string): string { + if (!Number.isFinite(cp) || cp <= 0 || cp > 0x10ffff || (cp >= 0xd800 && cp <= 0xdfff)) { + return fallback; + } + return String.fromCodePoint(cp); +} + +/** Numeric and the common named entities. `&` goes last so + * `&lt;` decodes to the literal `<` the author wrote. */ +function decodeEntities(s: string): string { + return s + .replace(/&#x([0-9a-f]{1,6});/gi, (m, hex: string) => codePoint(parseInt(hex, 16), m)) + .replace(/&#(\d{1,7});/g, (m, dec: string) => codePoint(parseInt(dec, 10), m)) + .replace(/&([a-z]+);/gi, (m, name: string) => NAMED_ENTITIES[name.toLowerCase()] ?? m) + .replace(/&/g, "&"); +} + +const BLOCK_TAGS = + "p|div|section|article|header|footer|main|aside|nav|h[1-6]|ul|ol|tr|table|thead|tbody|tfoot|blockquote|pre|hr|dd|dt|dl|figure|figcaption|form|fieldset|address|details|summary"; + +/** + * Dependency-free HTML to text. Good enough for a model to read a page; + * not a renderer. Scripts, styles and SVG go away entirely; block-level + * boundaries become line breaks; list items get a leading dash; the + * rest of the markup is dropped and entities decoded. Whitespace is + * collapsed except for line breaks, so `
` keeps its lines but not
+ * its indentation.
+ */
+export function htmlToText(html: string): { title: string | null; text: string } {
+  const titleMatch = /]*>([\s\S]*?)<\/title>/i.exec(html);
+  const title = titleMatch
+    ? decodeEntities(titleMatch[1]).replace(/\s+/g, " ").trim() || null
+    : null;
+
+  const stripped = html
+    .replace(//g, " ")
+    .replace(/<(script|style|noscript|svg|template|head)\b[\s\S]*?<\/\1\s*>/gi, " ")
+    .replace(//gi, "\n")
+    .replace(/]*>/gi, "\n- ")
+    .replace(new RegExp(`<\\/?(?:${BLOCK_TAGS})\\b[^>]*>`, "gi"), "\n")
+    .replace(/<[^>]+>/g, " ");
+
+  const text = decodeEntities(stripped)
+    .replace(/[ \t\r\f\v ]+/g, " ")
+    .replace(/ ?\n ?/g, "\n")
+    .replace(/\n{3,}/g, "\n\n")
+    .trim();
+
+  return { title, text };
+}
+
+// ── Handle ─────────────────────────────────────────────────────────────
+
+/**
+ * Build the `web_fetch` tool for a run. See the module header for the
+ * usage shape.
+ */
+export function createWebFetch(opts: CreateWebFetchOptions): WebFetchHandle {
+  const results: WebFetchResult[] = [];
+  const backend = opts.backend ?? resolveFetchBackend(opts.provider);
+  const maxUses = opts.maxUses ?? DEFAULT_MAX_USES;
+  const maxCharacters =
+    opts.maxCharacters ??
+    (opts.maxContentTokens ? opts.maxContentTokens * CHARS_PER_TOKEN : DEFAULT_MAX_CHARACTERS);
+  const maxContentTokens =
+    opts.maxContentTokens ??
+    (opts.maxCharacters ? Math.ceil(opts.maxCharacters / CHARS_PER_TOKEN) : undefined);
+
+  if (backend === "anthropic") {
+    const apiKey =
+      normalizeApiKey(opts.apiKey) || normalizeApiKey(process.env.ANTHROPIC_API_KEY);
+    if (!apiKey) {
+      return emptyFetchHandle(results);
+    }
+    const baseURL = getGatewayBaseURL("anthropic");
+    const anthropic = createAnthropic({ apiKey, ...(baseURL && { baseURL }) });
+    return {
+      tool: anthropic.tools.webFetch_20250910({
+        maxUses,
+        ...(maxContentTokens ? { maxContentTokens } : {}),
+        ...(opts.allowedDomains?.length ? { allowedDomains: opts.allowedDomains } : {}),
+        ...(opts.blockedDomains?.length ? { blockedDomains: opts.blockedDomains } : {}),
+      }),
+      backend,
+      native: true,
+      results,
+      capture: (stepContent) => captureNativeFetchResults(stepContent, results),
+    };
+  }
+
+  let uses = 0;
+  return {
+    tool: tool({
+      description:
+        "Fetch a specific web page by URL and return its text. HTML is converted to plain text; " +
+        "JSON and plain text come back as-is. Use this to read a page whose address you already " +
+        "have — one the user gave you, or one returned by web_search. It is not a search engine: " +
+        "it needs a full http(s) URL. PDFs and other binary content are not supported.",
+      inputSchema: z.object({
+        url: z.string().describe("The absolute http(s) URL to fetch."),
+      }),
+      execute: async ({ url }: { url: string }, ctx?: { abortSignal?: AbortSignal }) => {
+        if (uses >= maxUses) {
+          return {
+            error: `web_fetch budget exhausted (${maxUses} calls). Work with what you have.`,
+          };
+        }
+        uses++;
+        try {
+          const page = await fetchUrl(url, {
+            maxCharacters,
+            allowedDomains: opts.allowedDomains,
+            blockedDomains: opts.blockedDomains,
+            abortSignal: ctx?.abortSignal ?? opts.abortSignal,
+            lookup: opts.lookup,
+          });
+          results.push(page);
+          return page;
+        } catch (err) {
+          return {
+            error: `Fetch failed: ${err instanceof Error ? err.message : String(err)}`,
+          };
+        }
+      },
+    }),
+    backend,
+    native: false,
+    results,
+    // HTTP results are appended by `execute` above; walking the step
+    // would double-count them.
+    capture: () => {},
+  };
+}
+
+function emptyFetchHandle(results: WebFetchResult[]): WebFetchHandle {
+  return {
+    tool: undefined,
+    backend: undefined,
+    native: false,
+    results,
+    capture: () => {},
+  };
+}
+
+/**
+ * Walk one AI SDK step's content for `web_fetch` tool-results and
+ * append each page to `target`, in order.
+ *
+ * Tolerates both result shapes (`output` and `result`) and both key
+ * casings for the nested fields — adapters vary across AI SDK versions,
+ * and a shape we don't recognize should cost us a bookkeeping entry,
+ * not the run. Anthropic's error results (`web_fetch_tool_result_error`)
+ * are skipped: there's no page to record.
+ */
+export function captureNativeFetchResults(
+  stepContent: unknown,
+  target: WebFetchResult[],
+): void {
+  if (!Array.isArray(stepContent)) return;
+  for (const content of stepContent) {
+    if (content?.type !== "tool-result") continue;
+    if (content?.toolName !== WEB_FETCH_TOOL_NAME) continue;
+    const body = content.output ?? content.result ?? null;
+    if (!body || typeof body !== "object") continue;
+    if (body.type !== "web_fetch_result" || typeof body.url !== "string") continue;
+    const doc = body.content ?? {};
+    const source = doc.source ?? {};
+    target.push({
+      url: body.url,
+      title: doc.title ?? null,
+      ...(source.type === "text" && typeof source.data === "string"
+        ? { text: source.data }
+        : {}),
+      mediaType: source.mediaType ?? source.media_type ?? null,
+      retrievedAt: body.retrievedAt ?? body.retrieved_at ?? null,
+      type: "web_fetch_result",
+    });
+  }
+}
diff --git a/mcp/src/aieo/src/index.ts b/mcp/src/aieo/src/index.ts
index 7a10d1720..02e21511d 100644
--- a/mcp/src/aieo/src/index.ts
+++ b/mcp/src/aieo/src/index.ts
@@ -5,6 +5,7 @@ export * from "./usage.js";
 export * from "./prompt.js";
 export * from "./tools.js";
 export * from "./search.js";
+export * from "./fetch.js";
 
 export type { ModelMessage } from "ai";
 export type { Tool, ToolSet } from "ai";
diff --git a/mcp/src/aieo/src/test/try-fetch.ts b/mcp/src/aieo/src/test/try-fetch.ts
new file mode 100644
index 000000000..943bbafb4
--- /dev/null
+++ b/mcp/src/aieo/src/test/try-fetch.ts
@@ -0,0 +1,68 @@
+import { generateText, stepCountIs } from "ai";
+import * as dotenv from "dotenv";
+import { getModelDetails, getProviderOptions } from "../provider.js";
+import { createWebFetch, WEB_FETCH_TOOL_NAME } from "../fetch.js";
+
+dotenv.config({ path: "../../.env" });
+
+// `npm run try-fetch -- https://some.url` to point both backends at a page.
+const URL_TO_FETCH = process.argv.find((a) => /^https?:\/\//.test(a)) ?? "https://sphinx.chat";
+
+// The URL is in the prompt on purpose: Anthropic's native tool only
+// fetches URLs that already appeared in the conversation.
+const PROMPT =
+  `Fetch ${URL_TO_FETCH} and write exactly 3 short bullet points about what the page says. ` +
+  "Quote the page title in the first bullet.";
+
+async function run(label: string, modelName: string) {
+  console.log(`\n${"=".repeat(70)}\n${label}\n${"=".repeat(70)}`);
+  const { model, provider, apiKey, modelId } = getModelDetails(modelName);
+  const wf = createWebFetch({ provider, apiKey, maxCharacters: 12_000 });
+  console.log(`backend=${wf.backend} native=${wf.native} hasTool=${!!wf.tool}`);
+  if (!wf.tool) {
+    console.error("no tool — missing key for this backend");
+    return;
+  }
+
+  const started = Date.now();
+  const res = await generateText({
+    model,
+    tools: { [WEB_FETCH_TOOL_NAME]: wf.tool },
+    system: "You are a concise research assistant.",
+    prompt: PROMPT,
+    stopWhen: stepCountIs(4),
+    providerOptions: getProviderOptions(provider, "fast", modelId) as any,
+    onStepFinish: (step) => {
+      const calls = step.content
+        .filter((c: any) => c.type === "tool-call")
+        .map((c: any) => `${c.toolName}(${JSON.stringify(c.input)?.slice(0, 80)})`);
+      if (calls.length) console.log(`  step: ${calls.join(", ")}`);
+      const errors = step.content
+        .filter((c: any) => c.type === "tool-result" && (c.isError || c.output?.error))
+        .map((c: any) => JSON.stringify(c.output ?? c.result)?.slice(0, 200));
+      if (errors.length) console.log(`  tool errors: ${errors.join(" | ")}`);
+      wf.capture(step.content);
+    },
+  });
+  const elapsed = ((Date.now() - started) / 1000).toFixed(1);
+
+  console.log(`\nsteps=${res.steps.length} elapsed=${elapsed}s captured=${wf.results.length}`);
+  for (const r of wf.results) {
+    console.log(
+      `  ${r.url}  title=${JSON.stringify(r.title)} type=${r.mediaType} ` +
+        `${r.text ? `(${r.text.length}b text${r.truncated ? ", truncated" : ""})` : "(no text)"}`,
+    );
+  }
+  console.log(`\n--- model text ---\n${res.text}`);
+}
+
+async function main() {
+  await run("ANTHROPIC (native web_fetch)", "sonnet").catch((e) =>
+    console.error("anthropic failed:", e?.message || e),
+  );
+  await run("XAI / GROK (http shim)", "grok").catch((e) =>
+    console.error("grok failed:", e?.message || e),
+  );
+}
+
+main();
diff --git a/mcp/src/aieo/src/tools.ts b/mcp/src/aieo/src/tools.ts
index 62fd3012a..c35f7272b 100644
--- a/mcp/src/aieo/src/tools.ts
+++ b/mcp/src/aieo/src/tools.ts
@@ -1,22 +1,24 @@
 import { createAnthropic } from "@ai-sdk/anthropic";
 import { Provider, getGatewayBaseURL } from "./provider.js";
 import { createWebSearch } from "./search.js";
+import { createWebFetch } from "./fetch.js";
 
-export type ProviderTool = "webSearch" | "bash";
+export type ProviderTool = "webSearch" | "webFetch" | "bash";
 
 /**
  * Provider-native tool by name.
  *
- * `webSearch` is special: only Anthropic has a native one, so every
- * other provider gets the Exa-backed shim from `./search.js` instead of
- * an exception. That keeps the tool available (and named `web_search`)
- * on any model. This entry point returns the bare tool — for citation
- * indices and the matching prompt snippet, call `createWebSearch`
- * directly.
+ * `webSearch` and `webFetch` are special: only Anthropic has native
+ * ones, so every other provider gets a shim of the same name and result
+ * shape instead of an exception — Exa-backed search from `./search.js`,
+ * a guarded HTTP GET from `./fetch.js`. That keeps both tools available
+ * (as `web_search` / `web_fetch`) on any model. This entry point returns
+ * the bare tool — for the result bookkeeping, citation indices and
+ * prompt snippet, call `createWebSearch` / `createWebFetch` directly.
  *
- * Returns `undefined` for `webSearch` when the chosen backend has no key
+ * Returns `undefined` for those two when the chosen backend has no key
  * configured; callers should drop the tool rather than fail the request.
- * Non-`webSearch` tools still throw for unsupported providers.
+ * Other tools still throw for unsupported providers.
  */
 export function getProviderTool(
   provider: Provider,
@@ -26,6 +28,9 @@ export function getProviderTool(
   if (toolName === "webSearch") {
     return createWebSearch({ provider, apiKey }).tool;
   }
+  if (toolName === "webFetch") {
+    return createWebFetch({ provider, apiKey }).tool;
+  }
   switch (provider) {
     case "anthropic":
       return getAnthropicTool(apiKey, toolName);

From 9e682b626f1f9be83de2825bbab714ba96977b1f Mon Sep 17 00:00:00 2001
From: Evan Feenstra 
Date: Thu, 10 Sep 2026 09:39:26 -0700
Subject: [PATCH 6/7] gateway ui: render xAI in the canvas providers column
 (#1676)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR #1673 added the xai entry to PROVIDER_DISPLAY / the icon table in
canvasTheme.ts, but Canvas.tsx's hardcoded PROVIDERS list (which
drives both the provider cards and the gateway→provider edges) was
never updated, so the canvas drew no xAI card and any xai spend the
matrix endpoint returned was dropped from the column.

Add xai as the fifth provider, keeping gemini: the list is now
deliberately a superset of the seed config's providers map so a
provider absent from config.json still shows a $0 card instead of
having its historical spend vanish. Refresh the stale header
comment (it still said "google" and claimed per-provider spend was
pending a matrix-endpoint change that already landed).
---
 .../internal/adminapi/ui/src/pages/Canvas.tsx | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/gateway/internal/adminapi/ui/src/pages/Canvas.tsx b/gateway/internal/adminapi/ui/src/pages/Canvas.tsx
index e40b4d839..b1b10dd18 100644
--- a/gateway/internal/adminapi/ui/src/pages/Canvas.tsx
+++ b/gateway/internal/adminapi/ui/src/pages/Canvas.tsx
@@ -17,10 +17,9 @@
 //     requestCount are summed over all that user's pairings.
 //   - Gateway: singleton. customData.totalCost = swarm-wide total
 //     (sum of every row's cost).
-//   - Providers: hardcoded 4 — anthropic / openai / openrouter /
-//     google — pulled from gateway/docker-compose.yml env. Real
-//     per-provider spend lands when the matrix endpoint grows a
-//     provider dimension.
+//   - Providers: hardcoded 5 — anthropic / openai / openrouter /
+//     gemini / xai. customData.totalCost + requestCount come from
+//     the per-row `providers[]` slice the matrix endpoint returns.
 //
 // The canvas takes over the full shell-main area; the WindowPicker
 // floats in the top-right corner as the only chrome. No numeric
@@ -57,11 +56,16 @@ const fmtInt = (v: number) =>
   new Intl.NumberFormat("en-US").format(Math.round(v));
 
 // Hardcoded provider list — these ids match Bifrost's `provider`
-// column on every log row (`gateway/data/config.json`'s `providers`
-// map), so the per-provider rollup from the matrix endpoint
-// dispatches against `customData.icon` cleanly. The order here is
-// the vertical render order in the providers column.
-const PROVIDERS = ["anthropic", "openai", "openrouter", "gemini"];
+// column on every log row, so the per-provider rollup from the
+// matrix endpoint dispatches against `customData.icon` cleanly.
+// Deliberately a superset of `gateway/data/config.json`'s current
+// `providers` map: a provider that's been dropped from (or not yet
+// added to) the seed config still gets a card showing $0, rather
+// than having its historical spend silently vanish from the column.
+// Every id here needs a PROVIDER_DISPLAY entry in canvasTheme.ts.
+// The order here is the vertical render order in the providers
+// column.
+const PROVIDERS = ["anthropic", "openai", "openrouter", "gemini", "xai"];
 
 // Column x-coordinates (canvas-space, centered around 0). Spacing
 // picked so even the widest columns (gateway 220) don't touch their

From 15e5076bdd97b2558274b605177ca530d76d1d9c Mon Sep 17 00:00:00 2001
From: Evan Feenstra 
Date: Thu, 10 Sep 2026 09:39:35 -0700
Subject: [PATCH 7/7] gateway: restore gemini as the fifth provider alongside
 xai (#1677)

PR #1673 swapped gemini for xai in the seed config, compose env, plan
docs, and the mcp/docs/gateway seed copy. The intent was to add xAI,
not to drop Gemini: put the gemini provider back in every one of
those spots, reading env.GOOGLE_API_KEY as before, so the gateway
targets anthropic / openai / openrouter / gemini / xai.

Provider order in both config.json files matches the admin UI's
canvas column (PR #1676), which already renders all five. No Go
changes: nothing server-side hardcodes the provider list, and the
x-goog-api-key transport tests never stopped covering Gemini.
---
 gateway/data/config.json                   | 10 ++++++++++
 gateway/plans/llm-governance-v2.md         |  3 ++-
 gateway/plans/phases/phase-1-reconciler.md |  3 ++-
 mcp/docs/gateway/README.md                 | 10 ++++++----
 mcp/docs/gateway/data/config.json          | 10 ++++++++++
 mcp/docs/gateway/docker-compose.yml        |  1 +
 6 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/gateway/data/config.json b/gateway/data/config.json
index f6514fd96..a57a45a6c 100644
--- a/gateway/data/config.json
+++ b/gateway/data/config.json
@@ -42,6 +42,16 @@
         }
       ]
     },
+    "gemini": {
+      "keys": [
+        {
+          "name": "gemini-key-1",
+          "value": "env.GOOGLE_API_KEY",
+          "models": ["*"],
+          "weight": 1.0
+        }
+      ]
+    },
     "xai": {
       "keys": [
         {
diff --git a/gateway/plans/llm-governance-v2.md b/gateway/plans/llm-governance-v2.md
index e0948a8ef..0c9a0252c 100644
--- a/gateway/plans/llm-governance-v2.md
+++ b/gateway/plans/llm-governance-v2.md
@@ -135,7 +135,7 @@ Bifrost natively enforces:
 | Daily $ budget     | Customer (per user) | $1000/day                                              |
 | Rate limit RPM     | Customer (per user) | 1000 RPM                                               |
 | Rate limit TPM     | Customer (per user) | 5M TPM                                                 |
-| Provider allowlist | VK                  | `[anthropic, openai, openrouter, xai]`, all `["*"]` |
+| Provider allowlist | VK                  | `[anthropic, openai, openrouter, gemini, xai]`, all `["*"]` |
 | Model allowlist    | VK                  | `["*"]` initially                                      |
 | `is_active`        | Customer (per user) | `true` by default; `false` disables account org-wide   |
 
@@ -173,6 +173,7 @@ POST /api/governance/virtual-keys
       { provider: "anthropic", allowed_models: ["*"] },
       { provider: "openai",    allowed_models: ["*"] },
       { provider: "openrouter",allowed_models: ["*"] },
+      { provider: "gemini",    allowed_models: ["*"] },
       { provider: "xai",       allowed_models: ["*"] },
     ]
   }
diff --git a/gateway/plans/phases/phase-1-reconciler.md b/gateway/plans/phases/phase-1-reconciler.md
index fbc925c8b..7a7205fd2 100644
--- a/gateway/plans/phases/phase-1-reconciler.md
+++ b/gateway/plans/phases/phase-1-reconciler.md
@@ -255,7 +255,7 @@ Filters available on this endpoint (lines 369-380 of governance.go):
           "budgets": [],
           "rate_limit": null
         }
-        // … openai, openrouter, xai
+        // … openai, openrouter, gemini, xai
       ],
       "mcp_configs": [],
       "budgets": [],
@@ -306,6 +306,7 @@ Content-Type: application/json
     { "provider": "anthropic",  "allowed_models": ["*"], "key_ids": ["*"] },
     { "provider": "openai",     "allowed_models": ["*"], "key_ids": ["*"] },
     { "provider": "openrouter", "allowed_models": ["*"], "key_ids": ["*"] },
+    { "provider": "gemini",     "allowed_models": ["*"], "key_ids": ["*"] },
     { "provider": "xai",        "allowed_models": ["*"], "key_ids": ["*"] }
   ]
 }
diff --git a/mcp/docs/gateway/README.md b/mcp/docs/gateway/README.md
index 42af1014a..3513f33c2 100644
--- a/mcp/docs/gateway/README.md
+++ b/mcp/docs/gateway/README.md
@@ -7,10 +7,12 @@ Runs Bifrost locally on `http://localhost:8181` so MCP can be tested with
 ## Layout
 
 - `docker-compose.yml` — Bifrost service, host port `8181 -> container 8080`.
-- `data/config.json` — seed config: Anthropic + OpenAI + OpenRouter + xAI
-  providers, API keys read from `env.*`, `config_store` enabled (SQLite) so the
-  Web UI works. (xAI has no dedicated Bifrost route; the MCP client sends Grok
-  through `/openai/v1` with an `xai/`-prefixed model id.)
+- `data/config.json` — seed config: Anthropic + OpenAI + OpenRouter + Gemini
+  + xAI providers, API keys read from `env.*`, `config_store` enabled (SQLite)
+  so the Web UI works. (Bifrost calls Google's public Gemini API `gemini`; the
+  MCP client side calls the same thing `google`. xAI has no dedicated Bifrost
+  route; the MCP client sends Grok through `/openai/v1` with an `xai/`-prefixed
+  model id.)
 - `data/config.db`, `data/logs.db` — created on first boot; gitignored.
 - `.env` — symlinked to `../../.env` (the MCP `.env`) so docker-compose picks up
   `ANTHROPIC_API_KEY` etc. without copy-paste.
diff --git a/mcp/docs/gateway/data/config.json b/mcp/docs/gateway/data/config.json
index 7743756c2..6a3d8f8a1 100644
--- a/mcp/docs/gateway/data/config.json
+++ b/mcp/docs/gateway/data/config.json
@@ -35,6 +35,16 @@
         }
       ]
     },
+    "gemini": {
+      "keys": [
+        {
+          "name": "gemini-key-1",
+          "value": "env.GOOGLE_API_KEY",
+          "models": ["*"],
+          "weight": 1.0
+        }
+      ]
+    },
     "xai": {
       "keys": [
         {
diff --git a/mcp/docs/gateway/docker-compose.yml b/mcp/docs/gateway/docker-compose.yml
index 0b401c04d..73dc9123c 100644
--- a/mcp/docs/gateway/docker-compose.yml
+++ b/mcp/docs/gateway/docker-compose.yml
@@ -13,6 +13,7 @@ services:
       ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY:-}
       OPENAI_API_KEY: ${OPENAI_API_KEY:-}
       OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:-}
+      GOOGLE_API_KEY: ${GOOGLE_API_KEY:-}
       XAI_API_KEY: ${XAI_API_KEY:-}
     ports:
       # Host 8181 -> container 8080. MCP server uses LLM_GATEWAY_URL=http://localhost:8181.