[agent-network] End-to-end test suite, module docs, and deployment preset

scripts/e2e/agent-network-policy/00-env.sh

@@ -0,0 +1,105 @@
+# shellcheck disable=SC2148
+# Sourced helper for the agent-network-policy e2e suite.
+# Verifies the agent-network surface: window_seconds (post-rename)
+# through the API, /api/agent-network/consumption read endpoint, and
+# the CheckLLMPolicyLimits + RecordLLMUsage gRPC RPCs.
+#
+# Do not run directly.
+
+set -euo pipefail
+
+: "${NB_API:=http://localhost:8080}"
+: "${NB_PAT_FILE:=/Users/maycon/projects/local-dev/nb-pat}"
+: "${NB_TIMEOUT_SECONDS:=60}"
+: "${NB_POLICY_NAME:=e2e-anpol}"
+: "${NB_PROVIDER_NAME:=e2e-anpol-provider}"
+: "${NB_GROUP_NAME:=e2e-anpol-engineers}"
+: "${NB_GRPC_ADDR:=localhost:8080}"
+# Proxy token shared with the tilt setup. Keep in sync with the
+# NB_PROXY_TOKEN literal in /Users/maycon/projects/local-dev/Tiltfile;
+# the management server accepts it as a registered proxy credential
+# so the e2e binary can reach the proxy_service gRPC surface.
+: "${NB_PROXY_TOKEN:=nbx_MEF9OKRhlLrWkc5TJmM3Eu2rhqigaP2yulHy}"
+: "${NB_STATE_DIR:=/tmp/nb-anpol-e2e-state}"
+
+mkdir -p "$NB_STATE_DIR"
+
+if [ ! -r "$NB_PAT_FILE" ]; then
+    echo "FAIL: cannot read PAT at $NB_PAT_FILE" >&2
+    exit 2
+fi
+NB_PAT=$(tr -d '\n\r ' <"$NB_PAT_FILE")
+if [ ${#NB_PAT} -lt 16 ]; then
+    echo "FAIL: PAT at $NB_PAT_FILE is suspiciously short (${#NB_PAT} chars)" >&2
+    exit 2
+fi
+
+if ! command -v jq >/dev/null 2>&1; then
+    echo "FAIL: jq is required (brew install jq)" >&2
+    exit 2
+fi
+
+# nb_api METHOD PATH [BODY] — wraps curl with PAT auth and JSON.
+nb_api() {
+    local method="$1" path="$2" body="${3-}"
+    if [ -n "$body" ]; then
+        curl -fsS -X "$method" \
+            -H "Authorization: Token $NB_PAT" \
+            -H "Content-Type: application/json" \
+            --data "$body" \
+            "$NB_API$path"
+    else
+        curl -fsS -X "$method" \
+            -H "Authorization: Token $NB_PAT" \
+            "$NB_API$path"
+    fi
+}
+
+# nb_api_status METHOD PATH [BODY] — returns the HTTP status code
+# rather than the body. Use for negative-path tests where a 4xx is
+# the expected outcome and curl's default -f would mask the result.
+nb_api_status() {
+    local method="$1" path="$2" body="${3-}"
+    if [ -n "$body" ]; then
+        curl -sS -o /dev/null -w '%{http_code}' -X "$method" \
+            -H "Authorization: Token $NB_PAT" \
+            -H "Content-Type: application/json" \
+            --data "$body" \
+            "$NB_API$path"
+    else
+        curl -sS -o /dev/null -w '%{http_code}' -X "$method" \
+            -H "Authorization: Token $NB_PAT" \
+            "$NB_API$path"
+    fi
+}
+
+# wait_for COND_CMD TIMEOUT_S — polls every 1s.
+wait_for() {
+    local cmd="$1" timeout="${2:-$NB_TIMEOUT_SECONDS}"
+    local i=0
+    while [ "$i" -lt "$timeout" ]; do
+        if eval "$cmd" >/dev/null 2>&1; then
+            return 0
+        fi
+        sleep 1
+        i=$((i + 1))
+    done
+    return 1
+}
+
+# resolve the absolute path of the netbird repo root from the script
+# location so the Go smoke binary can be invoked via `go run` against
+# the right module regardless of the caller's cwd.
+nb_repo_root() {
+    cd "$(dirname "$0")/../../.." && pwd
+}
+
+pass() {
+    printf 'PASS: %s\n' "$1"
+    exit 0
+}
+
+fail() {
+    printf 'FAIL: %s — %s\n' "$1" "${2:-}"
+    exit 1
+}

scripts/e2e/agent-network-policy/01-tilt-restart.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# 01-tilt-restart: re-trigger management + dashboard before each run so
+# we start from a clean process state. The on-disk store survives, so
+# the PAT, account, and any residual config persist across restarts.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+command -v tilt >/dev/null 2>&1 || fail "tilt not on PATH" "brew install tilt-dev/tap/tilt"
+
+if ! curl -fsS -o /dev/null --max-time 2 http://localhost:10350/ 2>/dev/null; then
+    fail "Tilt API not reachable at http://localhost:10350" "is 'tilt up' running in /Users/maycon/projects/local-dev?"
+fi
+
+restart_one() {
+    local name="$1"
+    if ! tilt trigger "$name" >/dev/null 2>&1; then
+        fail "tilt trigger $name failed" ""
+    fi
+    echo "triggered $name"
+}
+
+restart_one management
+restart_one dashboard
+# proxy3 needs to come back too because PR2 wired the new
+# llm_limit_check / llm_limit_record middlewares into the proxy
+# binary; without restarting it, e2e keeps exercising whatever
+# image was last loaded and silently misses any boot-order or
+# wiring regression in the chain. Cheap to restart, expensive to
+# silently miss.
+restart_one proxy3
+
+echo "waiting for management to accept requests..."
+if ! wait_for "curl -fsS -o /dev/null --max-time 2 $NB_API/oauth2/.well-known/openid-configuration" 60; then
+    fail "management did not come back up within 60s" "check 'tilt logs management'"
+fi
+echo "management is up"
+
+code=$(curl -fsS -o /dev/null -w '%{http_code}' \
+    -H "Authorization: Token $NB_PAT" \
+    "$NB_API/api/users" 2>&1) || true
+[ "$code" = "200" ] || fail "PAT auth check failed after restart (HTTP $code)" "the PAT may have been revoked"
+
+pass "Tilt resources restarted: management, dashboard"

scripts/e2e/agent-network-policy/10-policy-create.sh

@@ -0,0 +1,132 @@
+#!/usr/bin/env bash
+# 10-policy-create: end-to-end round-trip on the new window_seconds
+# field. Creates the prerequisites (group + provider with
+# bootstrap_cluster), POSTs a policy whose Limits.token_limit and
+# budget_limit both carry window_seconds, then re-fetches the policy
+# and asserts the field comes back unchanged.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+# 1. Group — sweep first so re-runs are idempotent, then POST.
+existing_group=$(nb_api GET /api/groups 2>/dev/null \
+    | jq -r --arg name "$NB_GROUP_NAME" '.[] | select(.name == $name) | .id // empty' \
+    | head -1)
+if [ -n "$existing_group" ]; then
+    echo "reusing existing group $existing_group"
+    group_id="$existing_group"
+else
+    body=$(jq -n --arg name "$NB_GROUP_NAME" '{name:$name}')
+    group_resp=$(nb_api POST /api/groups "$body")
+    group_id=$(printf '%s' "$group_resp" | jq -r '.id // ""')
+    [ -n "$group_id" ] && [ "$group_id" != "null" ] \
+        || fail "group create did not return an id" "$group_resp"
+    echo "created group $group_id"
+fi
+printf '%s' "$group_id" >"$NB_STATE_DIR/group-id"
+
+# 2. Provider — also idempotent. Bootstrap cluster pinned to the local
+#    proxy3 cluster so the management settings row resolves and the
+#    create completes (subsequent provider creates in the same account
+#    ignore bootstrap_cluster, but the FIRST one needs it).
+existing_provider=$(nb_api GET /api/agent-network/providers 2>/dev/null \
+    | jq -r --arg name "$NB_PROVIDER_NAME" '.[] | select(.name == $name) | .id // empty' \
+    | head -1)
+if [ -n "$existing_provider" ]; then
+    echo "reusing existing provider $existing_provider"
+    provider_id="$existing_provider"
+else
+    provider_body=$(jq -n \
+        --arg name "$NB_PROVIDER_NAME" \
+        '{
+            provider_id: "openai_api",
+            name: $name,
+            upstream_url: "https://api.openai.com",
+            api_key: "sk-e2e-placeholder",
+            bootstrap_cluster: "proxy.netbird.local",
+            models: []
+        }')
+    prov_resp=$(nb_api POST /api/agent-network/providers "$provider_body")
+    provider_id=$(printf '%s' "$prov_resp" | jq -r '.id // ""')
+    [ -n "$provider_id" ] && [ "$provider_id" != "null" ] \
+        || fail "provider create did not return an id" "$prov_resp"
+    echo "created provider $provider_id"
+fi
+printf '%s' "$provider_id" >"$NB_STATE_DIR/provider-id"
+
+# 3. Policy — drop any prior with the same name, then create with the
+#    NEW window_seconds field on both halves of Limits. 86400s = 24h
+#    on token, 2_592_000s = 30d on budget so the round-trip is
+#    unambiguous (no ambiguous unit-conversion artefact when we read
+#    back).
+existing_policy=$(nb_api GET /api/agent-network/policies 2>/dev/null \
+    | jq -r --arg name "$NB_POLICY_NAME" '.[] | select(.name == $name) | .id // empty' \
+    | head -1)
+if [ -n "$existing_policy" ]; then
+    echo "deleting existing policy $existing_policy"
+    nb_api DELETE "/api/agent-network/policies/$existing_policy" >/dev/null 2>&1 || true
+fi
+
+policy_body=$(jq -n \
+    --arg name "$NB_POLICY_NAME" \
+    --arg group "$group_id" \
+    --arg provider "$provider_id" \
+    '{
+        name: $name,
+        description: "agent-network e2e: window_seconds round-trip",
+        enabled: true,
+        source_groups: [$group],
+        destination_provider_ids: [$provider],
+        guardrail_ids: [],
+        limits: {
+            token_limit: {
+                enabled: true,
+                group_cap: 10000,
+                user_cap: 5000,
+                window_seconds: 86400
+            },
+            budget_limit: {
+                enabled: true,
+                group_cap_usd: 10.0,
+                user_cap_usd: 2.5,
+                window_seconds: 2592000
+            }
+        }
+    }')
+
+resp=$(nb_api POST /api/agent-network/policies "$policy_body")
+policy_id=$(printf '%s' "$resp" | jq -r '.id // ""')
+[ -n "$policy_id" ] && [ "$policy_id" != "null" ] \
+    || fail "policy create did not return an id" "$resp"
+printf '%s' "$policy_id" >"$NB_STATE_DIR/policy-id"
+echo "policy id: $policy_id"
+
+# 4. Round-trip assertions: GET the policy back and verify the
+#    window_seconds values land on both limit halves. The OLD
+#    window_hours / window_days fields must be absent.
+got=$(nb_api GET "/api/agent-network/policies/$policy_id")
+
+token_window=$(printf '%s' "$got" | jq -r '.limits.token_limit.window_seconds // empty')
+budget_window=$(printf '%s' "$got" | jq -r '.limits.budget_limit.window_seconds // empty')
+
+[ "$token_window" = "86400" ] \
+    || fail "token_limit.window_seconds did not round-trip" \
+            "expected=86400 got=$token_window body=$got"
+[ "$budget_window" = "2592000" ] \
+    || fail "budget_limit.window_seconds did not round-trip" \
+            "expected=2592000 got=$budget_window body=$got"
+
+# Negative: window_hours / window_days are legacy field names and
+# must not be present in the response at all — their presence would
+# mean the management server is still emitting the legacy shape.
+legacy_h=$(printf '%s' "$got" | jq -r '.limits.token_limit | has("window_hours")')
+legacy_d=$(printf '%s' "$got" | jq -r '.limits.token_limit | has("window_days")')
+[ "$legacy_h" = "false" ] \
+    || fail "legacy window_hours field still present in token_limit response" "$got"
+[ "$legacy_d" = "false" ] \
+    || fail "legacy window_days field still present in token_limit response" "$got"
+
+echo "token_limit.window_seconds = $token_window"
+echo "budget_limit.window_seconds = $budget_window"
+
+pass "policy persisted with window_seconds on both Limits halves"

scripts/e2e/agent-network-policy/20-policy-rejects-zero-window.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# 20-policy-rejects-zero-window: management must reject a policy whose
+# token_limit or budget_limit is enabled but carries window_seconds < 60.
+# Anything below the one-minute floor would either be a zero / negative
+# window with no reset boundary, or a sub-minute window that produces
+# untenable consumption-row volume at scale. The handler's
+# validatePolicyLimits guard owes us a 4xx with a useful message.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+[ -r "$NB_STATE_DIR/group-id" ] && [ -r "$NB_STATE_DIR/provider-id" ] \
+    || fail "missing group-id / provider-id state" \
+            "run 10-policy-create.sh first to bootstrap prerequisites"
+group_id=$(cat "$NB_STATE_DIR/group-id")
+provider_id=$(cat "$NB_STATE_DIR/provider-id")
+
+# Build a payload that is well-formed except for window_seconds=30 on
+# token_limit (under the 60s minimum). We mark the field enabled so
+# the validation path actually runs — a disabled limit is allowed to
+# carry zero.
+payload=$(jq -n \
+    --arg name "$NB_POLICY_NAME-sub-minute" \
+    --arg group "$group_id" \
+    --arg provider "$provider_id" \
+    '{
+        name: $name,
+        enabled: true,
+        source_groups: [$group],
+        destination_provider_ids: [$provider],
+        guardrail_ids: [],
+        limits: {
+            token_limit: {
+                enabled: true,
+                group_cap: 10000,
+                user_cap: 5000,
+                window_seconds: 30
+            },
+            budget_limit: {
+                enabled: false,
+                group_cap_usd: 0,
+                user_cap_usd: 0,
+                window_seconds: 0
+            }
+        }
+    }')
+
+code=$(nb_api_status POST /api/agent-network/policies "$payload")
+[ "$code" = "400" ] || [ "$code" = "422" ] \
+    || fail "expected 400/422 on enabled token_limit with window_seconds<60" \
+            "got HTTP $code"
+
+# Sweep any policy that may have been mistakenly persisted (defence
+# against a future bug; today's handler doesn't get there).
+orphan=$(nb_api GET /api/agent-network/policies 2>/dev/null \
+    | jq -r --arg name "$NB_POLICY_NAME-sub-minute" '.[] | select(.name == $name) | .id // empty' \
+    | head -1)
+if [ -n "$orphan" ]; then
+    nb_api DELETE "/api/agent-network/policies/$orphan" >/dev/null 2>&1 || true
+    fail "policy was persisted despite sub-minute window_seconds" \
+         "id=$orphan — handler validation regression"
+fi
+
+echo "POST with token_limit.window_seconds=30 rejected with HTTP $code"
+pass "validation rejects sub-minute window_seconds when limit is enabled"

scripts/e2e/agent-network-policy/30-consumption-list-empty.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# 30-consumption-list-empty: GET /api/agent-network/consumption must
+# return a JSON array (possibly empty) — never a 404 / 500. The
+# endpoint is the read side that backs the dashboard's basic counter
+# view and must always be reachable so the page can render an empty
+# state.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+resp=$(nb_api GET /api/agent-network/consumption)
+
+# Must be a JSON array.
+if ! printf '%s' "$resp" | jq -e 'type == "array"' >/dev/null 2>&1; then
+    fail "consumption endpoint did not return a JSON array" "$resp"
+fi
+
+count=$(printf '%s' "$resp" | jq 'length')
+echo "consumption rows: $count"
+
+# Stash the baseline count so 40-grpc-record-and-list can compare.
+printf '%s' "$count" >"$NB_STATE_DIR/consumption-baseline"
+
+pass "consumption read endpoint returns a JSON array (count=$count)"

scripts/e2e/agent-network-policy/40-grpc-record-and-list.sh

@@ -0,0 +1,133 @@
+#!/usr/bin/env bash
+# 40-grpc-record-and-list: drives the new RecordLLMUsage and
+# CheckLLMPolicyLimits gRPC RPCs through the e2e usage_smoke helper.
+#
+# Asserts:
+#   - CheckLLMPolicyLimits returns decision=allow, picks the lowest
+#     group id by string sort as attribution, default window of 24h.
+#   - RecordLLMUsage with both user_id and group_id ticks BOTH the
+#     user counter and the group counter exactly once each.
+#   - A second RecordLLMUsage on the same key sums the deltas server
+#     side (database upsert-increment, no read-modify-write race).
+#   - The HTTP /api/agent-network/consumption listing reflects the
+#     post-flight state.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+[ -r "$NB_STATE_DIR/group-id" ] && [ -r "$NB_STATE_DIR/provider-id" ] \
+    || fail "missing group-id / provider-id state" \
+            "run 10-policy-create.sh first to bootstrap prerequisites"
+group_id=$(cat "$NB_STATE_DIR/group-id")
+provider_id=$(cat "$NB_STATE_DIR/provider-id")
+
+# Resolve the calling account id. /api/accounts returns NetBird xids,
+# which is what the gRPC service expects (NOT IDP user UUIDs).
+account_id=$(nb_api GET /api/accounts 2>/dev/null | jq -r '.[0].id // empty')
+[ -n "$account_id" ] || fail "could not resolve account id" "GET /api/accounts returned nothing"
+
+# Pick a stable user id for the test. The dimension is per-user, so we
+# can use any unique value — a synthetic "e2e-..." prefix avoids
+# colliding with real user ids in the consumption listing.
+user_id="e2e-anpol-user-$$"
+window_seconds=86400
+echo "account=$account_id user=$user_id group=$group_id window=${window_seconds}s"
+
+# ─── 1. CheckLLMPolicyLimits ────────────────────────────────────────
+# Pass 3 group ids out of order; lowest by string sort wins.
+groups_csv="grp-zz,$group_id,grp-aa-z"
+check_resp=$(cd "$(nb_repo_root)" && go run ./scripts/e2e/agent-network-policy/cmd/usage_smoke check \
+    --account "$account_id" \
+    --user "$user_id" \
+    --groups "$groups_csv" \
+    --provider "$provider_id" \
+    --model "gpt-4o" \
+    --token "$NB_PROXY_TOKEN" \
+    --addr "$NB_GRPC_ADDR" 2>&1) \
+    || fail "CheckLLMPolicyLimits gRPC failed" "$check_resp"
+
+decision=$(printf '%s' "$check_resp" | jq -r '.decision // ""')
+attribution=$(printf '%s' "$check_resp" | jq -r '.attribution_group_id // ""')
+got_window=$(printf '%s' "$check_resp" | jq -r '.window_seconds // ""')
+
+[ "$decision" = "allow" ] \
+    || fail "Check decision must be allow under PR1 stub" "$check_resp"
+# Lowest group id of {grp-zz, $group_id, grp-aa-z} by string sort. The
+# group_id is an xid (lowercase alnum starting with a digit/letter), so
+# the sort answer depends on group_id's prefix. Compute it locally.
+expected_low=$(printf '%s\n%s\n%s' "grp-zz" "$group_id" "grp-aa-z" | sort | head -1)
+[ "$attribution" = "$expected_low" ] \
+    || fail "Check did not pick the lowest-by-sort group" \
+            "expected=$expected_low got=$attribution"
+[ "$got_window" = "86400" ] \
+    || fail "Check window_seconds stub default mismatch" "expected=86400 got=$got_window"
+
+echo "Check decision=$decision attribution=$attribution window=${got_window}s"
+
+# ─── 2. RecordLLMUsage — first increment ────────────────────────────
+(cd "$(nb_repo_root)" && go run ./scripts/e2e/agent-network-policy/cmd/usage_smoke record \
+    --account "$account_id" \
+    --user "$user_id" \
+    --group "$group_id" \
+    --window-seconds "$window_seconds" \
+    --tokens-in 100 \
+    --tokens-out 50 \
+    --cost-usd 0.0125 \
+    --token "$NB_PROXY_TOKEN" \
+    --addr "$NB_GRPC_ADDR") >/dev/null \
+    || fail "RecordLLMUsage (first increment) failed" ""
+
+# ─── 3. RecordLLMUsage — second increment, same key ─────────────────
+(cd "$(nb_repo_root)" && go run ./scripts/e2e/agent-network-policy/cmd/usage_smoke record \
+    --account "$account_id" \
+    --user "$user_id" \
+    --group "$group_id" \
+    --window-seconds "$window_seconds" \
+    --tokens-in 50 \
+    --tokens-out 25 \
+    --cost-usd 0.0025 \
+    --token "$NB_PROXY_TOKEN" \
+    --addr "$NB_GRPC_ADDR") >/dev/null \
+    || fail "RecordLLMUsage (second increment) failed" ""
+
+# ─── 4. Read-back via HTTP — sums must converge ─────────────────────
+listing=$(nb_api GET /api/agent-network/consumption)
+
+user_row=$(printf '%s' "$listing" | jq --arg u "$user_id" \
+    'map(select(.dimension_kind == "user" and .dimension_id == $u)) | .[0]')
+[ "$(printf '%s' "$user_row" | jq -r '. // empty')" != "" ] \
+    || fail "user consumption row missing after RecordLLMUsage" "$listing"
+
+user_in=$(printf '%s' "$user_row" | jq '.tokens_input')
+user_out=$(printf '%s' "$user_row" | jq '.tokens_output')
+user_cost=$(printf '%s' "$user_row" | jq '.cost_usd')
+user_window=$(printf '%s' "$user_row" | jq '.window_seconds')
+
+[ "$user_in" = "150" ] && [ "$user_out" = "75" ] \
+    || fail "user counter did not sum the two increments" \
+            "expected tokens_in=150 tokens_out=75; got in=$user_in out=$user_out"
+
+# Floating point compare with awk — drift > 1e-9 is a real bug.
+cost_ok=$(awk -v got="$user_cost" 'BEGIN { print (got > 0.0149 && got < 0.0151) ? "y" : "n" }')
+[ "$cost_ok" = "y" ] \
+    || fail "user cost did not sum to 0.015" "got=$user_cost"
+
+[ "$user_window" = "$window_seconds" ] \
+    || fail "user counter window_seconds mismatch" \
+            "expected=$window_seconds got=$user_window"
+
+# Group row gets the same deltas because RecordLLMUsage ticks both
+# dimensions on a single call.
+group_row=$(printf '%s' "$listing" | jq --arg g "$group_id" \
+    'map(select(.dimension_kind == "group" and .dimension_id == $g)) | .[0]')
+group_in=$(printf '%s' "$group_row" | jq '.tokens_input')
+group_out=$(printf '%s' "$group_row" | jq '.tokens_output')
+
+[ "$group_in" = "150" ] && [ "$group_out" = "75" ] \
+    || fail "group counter did not sum the two increments" \
+            "expected tokens_in=150 tokens_out=75; got in=$group_in out=$group_out"
+
+echo "user counter: $user_in input / $user_out output / \$$user_cost over ${user_window}s"
+echo "group counter: $group_in input / $group_out output"
+
+pass "gRPC Check + Record round-trip atomically increments user + group counters"

scripts/e2e/agent-network-policy/50-grpc-allow-record-deny.sh

@@ -0,0 +1,210 @@
+#!/usr/bin/env bash
+# 50-grpc-allow-record-deny: drives the full lifecycle of the real
+# selection algorithm landed in PR2 — initial allow, partial fills
+# below the cap, exact-at-cap, then deny once consumption reaches the
+# group_cap. Validates that:
+#
+#   - CheckLLMPolicyLimits picks the test policy as the attribution
+#     when it's the only one authorising the (group, provider) pair.
+#   - The selected_policy_id + window_seconds round-trip on the wire.
+#   - Counter-aware headroom math is wired end-to-end (RecordLLMUsage
+#     ticks counters that the next CheckLLMPolicyLimits call reads).
+#   - At-cap consumption flips decision from allow to deny with the
+#     canonical llm_policy.token_cap_exceeded code.
+#
+# The setup creates a dedicated group + provider + policy with the
+# "-tight" suffix so it's isolated from the e2e-anpol resources
+# 10-policy-create.sh seeded. 99-cleanup's prefix sweep catches both.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+# Per-run suffix so the consumption counter starts at zero on every
+# run. Without this, a successful prior run leaves 200 tokens on the
+# (group, 24h) bucket and the next run's stage 1 check sees a deny
+# from the start. Consumption rows have no delete endpoint today.
+RUN_TAG="$(date +%s)-$$"
+NB_TIGHT_GROUP_NAME="$NB_GROUP_NAME-tight-$RUN_TAG"
+NB_TIGHT_PROVIDER_NAME="$NB_PROVIDER_NAME-tight-$RUN_TAG"
+NB_TIGHT_POLICY_NAME="$NB_POLICY_NAME-tight-$RUN_TAG"
+TIGHT_CAP=200
+TIGHT_WINDOW_SECONDS=86400
+
+# 1. Group — fresh per-run so the (group, window) consumption counter
+# starts at zero. The 99-cleanup sweep catches it via the
+# `e2e-anpol-engineers-tight-` prefix.
+body=$(jq -n --arg name "$NB_TIGHT_GROUP_NAME" '{name:$name}')
+tight_group_id=$(nb_api POST /api/groups "$body" | jq -r '.id // ""')
+[ -n "$tight_group_id" ] && [ "$tight_group_id" != "null" ] \
+    || fail "tight group create failed" ""
+echo "created tight group $tight_group_id"
+printf '%s' "$tight_group_id" >"$NB_STATE_DIR/tight-group-id"
+
+# 2. Provider — fresh per-run too so the policy targets a provider
+# only this run knows about. Eliminates cross-policy interference if
+# a prior run's tight policy targeted the same provider id.
+body=$(jq -n \
+    --arg name "$NB_TIGHT_PROVIDER_NAME" \
+    '{
+        provider_id: "openai_api",
+        name: $name,
+        upstream_url: "https://api.openai.com",
+        api_key: "sk-e2e-tight-placeholder",
+        bootstrap_cluster: "proxy.netbird.local",
+        models: []
+    }')
+tight_provider_id=$(nb_api POST /api/agent-network/providers "$body" | jq -r '.id // ""')
+[ -n "$tight_provider_id" ] && [ "$tight_provider_id" != "null" ] \
+    || fail "tight provider create failed" ""
+echo "created tight provider $tight_provider_id"
+printf '%s' "$tight_provider_id" >"$NB_STATE_DIR/tight-provider-id"
+
+# 3. Policy with a tight token cap so the test can observe the deny
+# transition without burning thousands of records.
+existing_policy=$(nb_api GET /api/agent-network/policies 2>/dev/null \
+    | jq -r --arg name "$NB_TIGHT_POLICY_NAME" '.[] | select(.name == $name) | .id // empty' \
+    | head -1)
+if [ -n "$existing_policy" ]; then
+    echo "deleting existing tight policy $existing_policy so the run starts clean"
+    nb_api DELETE "/api/agent-network/policies/$existing_policy" >/dev/null 2>&1 || true
+fi
+
+policy_body=$(jq -n \
+    --arg name "$NB_TIGHT_POLICY_NAME" \
+    --arg group "$tight_group_id" \
+    --arg provider "$tight_provider_id" \
+    --argjson cap "$TIGHT_CAP" \
+    --argjson window "$TIGHT_WINDOW_SECONDS" \
+    '{
+        name: $name,
+        description: "agent-network e2e: cap-exhaust deny test",
+        enabled: true,
+        source_groups: [$group],
+        destination_provider_ids: [$provider],
+        guardrail_ids: [],
+        limits: {
+            token_limit: {
+                enabled: true,
+                group_cap: $cap,
+                user_cap: 0,
+                window_seconds: $window
+            },
+            budget_limit: {
+                enabled: false,
+                group_cap_usd: 0,
+                user_cap_usd: 0,
+                window_seconds: $window
+            }
+        }
+    }')
+tight_policy_id=$(nb_api POST /api/agent-network/policies "$policy_body" | jq -r '.id // ""')
+[ -n "$tight_policy_id" ] && [ "$tight_policy_id" != "null" ] \
+    || fail "tight policy create failed" ""
+printf '%s' "$tight_policy_id" >"$NB_STATE_DIR/tight-policy-id"
+echo "created tight policy $tight_policy_id (group_cap=$TIGHT_CAP, window=${TIGHT_WINDOW_SECONDS}s)"
+
+# 4. Resolve the calling account id; the smoke binary stamps it onto
+# every gRPC request the way the real proxy does.
+account_id=$(nb_api GET /api/accounts 2>/dev/null | jq -r '.[0].id // empty')
+[ -n "$account_id" ] || fail "could not resolve account id" ""
+
+# Test user — synthetic prefix avoids colliding with real users.
+test_user="e2e-anpol-tight-user-$$"
+
+# Pre-resolve the netbird repo root ONCE so the helpers below can
+# pushd / popd into it without `cd "$(dirname "$0")/.."` re-resolving
+# from a moved cwd between calls. nb_repo_root walks up from $0; if a
+# prior helper left cwd elsewhere the relative walk breaks.
+repo_root=$(nb_repo_root)
+[ -d "$repo_root" ] || fail "could not resolve netbird repo root" "got=$repo_root"
+
+# Run the smoke binary inside a subshell so its `cd` doesn't pollute
+# the caller's cwd between stages.
+do_check() {
+    (
+        cd "$repo_root" || exit 1
+        go run ./scripts/e2e/agent-network-policy/cmd/usage_smoke check \
+            --account "$account_id" \
+            --user "$test_user" \
+            --groups "$tight_group_id" \
+            --provider "$tight_provider_id" \
+            --model "gpt-4o" \
+            --token "$NB_PROXY_TOKEN" \
+            --addr "$NB_GRPC_ADDR"
+    )
+}
+
+do_record() {
+    local in="$1" out="$2" cost="$3"
+    (
+        cd "$repo_root" || exit 1
+        go run ./scripts/e2e/agent-network-policy/cmd/usage_smoke record \
+            --account "$account_id" \
+            --user "$test_user" \
+            --group "$tight_group_id" \
+            --window-seconds "$TIGHT_WINDOW_SECONDS" \
+            --tokens-in "$in" \
+            --tokens-out "$out" \
+            --cost-usd "$cost" \
+            --token "$NB_PROXY_TOKEN" \
+            --addr "$NB_GRPC_ADDR"
+    )
+}
+
+# Stage 1 — fresh state. Selection must allow and pick our tight
+# policy as attribution because it's the ONLY policy authorising the
+# (tight_group, tight_provider) tuple. Also confirms the wire-level
+# selected_policy_id round-trip from manager → grpc → smoke client.
+echo "stage 1: initial check (consumption=0/$TIGHT_CAP)"
+resp=$(do_check) || fail "initial check failed" "$resp"
+decision=$(printf '%s' "$resp" | jq -r '.decision // ""')
+selected=$(printf '%s' "$resp" | jq -r '.selected_policy_id // ""')
+window=$(printf '%s' "$resp" | jq -r '.window_seconds // ""')
+[ "$decision" = "allow" ] \
+    || fail "expected allow on fresh state" "$resp"
+[ "$selected" = "$tight_policy_id" ] \
+    || fail "selection did not pick the tight policy" \
+            "expected=$tight_policy_id got=$selected"
+[ "$window" = "$TIGHT_WINDOW_SECONDS" ] \
+    || fail "window_seconds mismatch on the wire" "expected=$TIGHT_WINDOW_SECONDS got=$window"
+echo "  → allow / selected=$selected / window=${window}s"
+
+# Stage 2 — book half the cap. The next check must still allow.
+echo "stage 2: record 100 input tokens (consumption=100/$TIGHT_CAP)"
+do_record 100 0 0 >/dev/null || fail "record (stage 2) failed" ""
+resp=$(do_check)
+decision=$(printf '%s' "$resp" | jq -r '.decision // ""')
+[ "$decision" = "allow" ] \
+    || fail "expected allow at half-cap" "$resp"
+echo "  → allow"
+
+# Stage 3 — push to one token below cap. Headroom shrinks but
+# decision stays allow.
+echo "stage 3: record 99 more input tokens (consumption=199/$TIGHT_CAP)"
+do_record 99 0 0 >/dev/null || fail "record (stage 3) failed" ""
+resp=$(do_check)
+decision=$(printf '%s' "$resp" | jq -r '.decision // ""')
+[ "$decision" = "allow" ] \
+    || fail "expected allow at one-below-cap" "$resp"
+echo "  → allow"
+
+# Stage 4 — exactly at cap. Selector treats consumed >= cap as
+# exhausted, so the next check must deny with the canonical token cap
+# code. The deny reason names the policy id so operators can debug
+# from the access log.
+echo "stage 4: record 1 final token (consumption=$TIGHT_CAP/$TIGHT_CAP)"
+do_record 1 0 0 >/dev/null || fail "record (stage 4) failed" ""
+resp=$(do_check) || fail "check after cap-exhaust failed" "$resp"
+decision=$(printf '%s' "$resp" | jq -r '.decision // ""')
+deny_code=$(printf '%s' "$resp" | jq -r '.deny_code // ""')
+deny_reason=$(printf '%s' "$resp" | jq -r '.deny_reason // ""')
+[ "$decision" = "deny" ] \
+    || fail "expected DENY at cap" "$resp"
+[ "$deny_code" = "llm_policy.token_cap_exceeded" ] \
+    || fail "deny_code mismatch" "expected=llm_policy.token_cap_exceeded got=$deny_code resp=$resp"
+echo "  → DENY / deny_code=$deny_code"
+echo "  → deny_reason: $deny_reason"
+[[ "$deny_reason" == *"$tight_policy_id"* ]] \
+    || fail "deny_reason must name the policy id for operator debugging" "$deny_reason"
+
+pass "selection algorithm flips allow → deny at cap-exhaust through the gRPC wire"

scripts/e2e/agent-network-policy/99-cleanup.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# 99-cleanup: idempotent teardown. Drops the policy, provider, group,
+# and any consumption rows our test created so re-runs start fresh.
+
+# shellcheck source=00-env.sh
+source "$(dirname "$0")/00-env.sh"
+
+# Policies first (FK / synth chain owns provider references). Sweep
+# every policy whose name STARTS with the e2e prefix so per-run
+# instances created by 50 (with $RUN_TAG suffixes) get cleaned up.
+nb_api GET /api/agent-network/policies 2>/dev/null \
+    | jq -r --arg pfx "$NB_POLICY_NAME" '.[] | select(.name | startswith($pfx)) | .id' \
+    | while read -r orphan; do
+        [ -n "$orphan" ] || continue
+        code=$(curl -fsS -X DELETE -H "Authorization: Token $NB_PAT" \
+            -o /dev/null -w '%{http_code}' \
+            "$NB_API/api/agent-network/policies/$orphan" 2>&1) || true
+        echo "DELETE policy/$orphan -> $code"
+    done
+
+# Providers — same prefix sweep. Both the 10-policy-create.sh's main
+# provider AND every per-run tight provider 50 minted carry the
+# NB_PROVIDER_NAME prefix.
+nb_api GET /api/agent-network/providers 2>/dev/null \
+    | jq -r --arg pfx "$NB_PROVIDER_NAME" '.[] | select(.name | startswith($pfx)) | .id' \
+    | while read -r orphan; do
+        [ -n "$orphan" ] || continue
+        code=$(curl -fsS -X DELETE -H "Authorization: Token $NB_PAT" \
+            -o /dev/null -w '%{http_code}' \
+            "$NB_API/api/agent-network/providers/$orphan" 2>&1) || true
+        echo "DELETE provider/$orphan -> $code"
+    done
+
+# Groups — prefix sweep. /api/groups doesn't error on a
+# referenced-elsewhere group; if so, the delete is a no-op and we
+# move on.
+nb_api GET /api/groups 2>/dev/null \
+    | jq -r --arg pfx "$NB_GROUP_NAME" '.[] | select(.name | startswith($pfx)) | .id' \
+    | while read -r orphan; do
+        [ -n "$orphan" ] || continue
+        code=$(curl -fsS -X DELETE -H "Authorization: Token $NB_PAT" \
+            -o /dev/null -w '%{http_code}' \
+            "$NB_API/api/groups/$orphan" 2>&1) || true
+        echo "DELETE group/$orphan -> $code"
+    done
+
+# We don't expose a consumption-delete endpoint — the rows survive
+# until the management store is recycled. Document the residue here
+# so anyone debugging "why are there old e2e rows" knows the source.
+remaining=$(nb_api GET /api/agent-network/consumption 2>/dev/null \
+    | jq --arg pfx "e2e-anpol-user-" 'map(select(.dimension_id | startswith($pfx))) | length' \
+    || echo "0")
+[ "$remaining" = "0" ] || \
+    echo "(left $remaining e2e consumption rows in the store — there's no delete endpoint yet)"
+
+rm -rf "$NB_STATE_DIR"
+
+pass "tear-down complete (idempotent)"

scripts/e2e/agent-network-policy/cmd/usage_smoke/main.go

@@ -0,0 +1,197 @@
+// usage_smoke is the e2e helper that drives the new agent-network
+// gRPC RPCs (CheckLLMPolicyLimits, RecordLLMUsage) against a local
+// management server. Run from the bash suite as a `go run` so the
+// proto types are always in sync with the management binary the
+// suite is exercising.
+//
+// Two subcommands today:
+//   - record: invokes RecordLLMUsage with the supplied tokens / cost
+//     and exits 0 on success.
+//   - check:  invokes CheckLLMPolicyLimits and prints the response as
+//     JSON on stdout so the bash test can assert on it via
+//     jq.
+//
+// Auth uses the same proxy bearer-token shape the real proxy uses
+// (see proxy/internal/grpc/auth.go); the bash suite reads the token
+// from the e2e env (NB_PROXY_TOKEN, defaulted to the Tilt literal).
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"os"
+	"strings"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// proxyTokenCreds mirrors proxy/internal/grpc.WithProxyToken's
+// PerRPCCredentials so the e2e binary can reach the gRPC service
+// using the same bearer-token shape the real proxy uses. Inlined
+// because the production helper lives behind /internal/.
+type proxyTokenCreds struct{ token string }
+
+func (c proxyTokenCreds) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return map[string]string{"authorization": "Bearer " + c.token}, nil
+}
+
+// RequireTransportSecurity is false here because Tilt's management is
+// plaintext on localhost — the e2e suite is the *only* caller of this
+// binary, never production.
+func (proxyTokenCreds) RequireTransportSecurity() bool { return false }
+
+func main() {
+	if len(os.Args) < 2 {
+		usage()
+	}
+	cmd := os.Args[1]
+	os.Args = append(os.Args[:1], os.Args[2:]...)
+
+	switch cmd {
+	case "record":
+		runRecord()
+	case "check":
+		runCheck()
+	default:
+		usage()
+	}
+}
+
+func usage() {
+	fmt.Fprintln(os.Stderr, "usage: usage_smoke <record|check> [flags]")
+	os.Exit(2)
+}
+
+func runRecord() {
+	addr := flag.String("addr", "localhost:8080", "management gRPC address")
+	token := flag.String("token", os.Getenv("NB_PROXY_TOKEN"), "proxy token")
+	accountID := flag.String("account", "", "netbird account id")
+	userID := flag.String("user", "", "netbird user id (optional)")
+	groupID := flag.String("group", "", "netbird policy attribution group id (optional)")
+	groupsCSV := flag.String("groups", "", "CSV of caller group ids (for account-rule fan-out)")
+	windowSeconds := flag.Int64("window-seconds", 86_400, "window length in seconds (0 allowed when only account rules apply)")
+	tokensIn := flag.Int64("tokens-in", 0, "input tokens to add")
+	tokensOut := flag.Int64("tokens-out", 0, "output tokens to add")
+	costUSD := flag.Float64("cost-usd", 0, "USD cost to add")
+	flag.Parse()
+
+	if strings.TrimSpace(*token) == "" {
+		exitErr("--token is required (or set NB_PROXY_TOKEN)")
+	}
+	if strings.TrimSpace(*accountID) == "" {
+		exitErr("--account is required")
+	}
+	var groupIDs []string
+	for _, g := range strings.Split(*groupsCSV, ",") {
+		g = strings.TrimSpace(g)
+		if g != "" {
+			groupIDs = append(groupIDs, g)
+		}
+	}
+	if *userID == "" && *groupID == "" && len(groupIDs) == 0 {
+		exitErr("at least one of --user, --group, or --groups must be set")
+	}
+
+	conn := dial(*addr, *token)
+	defer conn.Close()
+	client := proto.NewProxyServiceClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, err := client.RecordLLMUsage(ctx, &proto.RecordLLMUsageRequest{
+		AccountId:     *accountID,
+		UserId:        *userID,
+		GroupId:       *groupID,
+		GroupIds:      groupIDs,
+		WindowSeconds: *windowSeconds,
+		TokensInput:   *tokensIn,
+		TokensOutput:  *tokensOut,
+		CostUsd:       *costUSD,
+	})
+	if err != nil {
+		exitErr(fmt.Sprintf("RecordLLMUsage: %v", err))
+	}
+	//nolint:forbidigo // e2e helper: stdout is the contract with the bash caller
+	fmt.Println("ok")
+}
+
+func runCheck() {
+	addr := flag.String("addr", "localhost:8080", "management gRPC address")
+	token := flag.String("token", os.Getenv("NB_PROXY_TOKEN"), "proxy token")
+	accountID := flag.String("account", "", "netbird account id")
+	userID := flag.String("user", "", "netbird user id (optional)")
+	groupsCSV := flag.String("groups", "", "CSV of caller group ids")
+	providerID := flag.String("provider", "", "agent-network provider id")
+	model := flag.String("model", "gpt-4o", "upstream model identifier")
+	flag.Parse()
+
+	if strings.TrimSpace(*token) == "" {
+		exitErr("--token is required (or set NB_PROXY_TOKEN)")
+	}
+	if strings.TrimSpace(*accountID) == "" {
+		exitErr("--account is required")
+	}
+
+	var groupIDs []string
+	for _, g := range strings.Split(*groupsCSV, ",") {
+		g = strings.TrimSpace(g)
+		if g != "" {
+			groupIDs = append(groupIDs, g)
+		}
+	}
+
+	conn := dial(*addr, *token)
+	defer conn.Close()
+	client := proto.NewProxyServiceClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	resp, err := client.CheckLLMPolicyLimits(ctx, &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  *accountID,
+		UserId:     *userID,
+		GroupIds:   groupIDs,
+		ProviderId: *providerID,
+		Model:      *model,
+	})
+	if err != nil {
+		exitErr(fmt.Sprintf("CheckLLMPolicyLimits: %v", err))
+	}
+
+	out, _ := json.Marshal(map[string]any{
+		"decision":             resp.GetDecision(),
+		"selected_policy_id":   resp.GetSelectedPolicyId(),
+		"attribution_group_id": resp.GetAttributionGroupId(),
+		"window_seconds":       resp.GetWindowSeconds(),
+		"deny_code":            resp.GetDenyCode(),
+		"deny_reason":          resp.GetDenyReason(),
+	})
+	//nolint:forbidigo // e2e helper: stdout is the contract with the bash caller
+	fmt.Println(string(out))
+}
+
+// dial connects to the management gRPC over plaintext. The bearer
+// token is sent on every RPC via PerRPCCredentials matching the wire
+// format the production proxy uses.
+func dial(addr, token string) *grpc.ClientConn {
+	conn, err := grpc.NewClient(addr,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithPerRPCCredentials(proxyTokenCreds{token: token}),
+	)
+	if err != nil {
+		exitErr(fmt.Sprintf("dial %s: %v", addr, err))
+	}
+	return conn
+}
+
+func exitErr(msg string) {
+	fmt.Fprintln(os.Stderr, msg)
+	os.Exit(1)
+}

scripts/e2e/agent-network-policy/run-all.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# run-all: every numbered test in order, halts on first FAIL,
+# prints a `=== PASS x/y ===` summary. -k keeps state (skips
+# 99-cleanup) so you can poke around the persisted policy /
+# consumption rows afterwards.
+
+set -uo pipefail
+
+KEEP=0
+[ "${1:-}" = "-k" ] && KEEP=1
+export NB_KEEP_STATE="$KEEP"
+
+cd "$(dirname "$0")"
+
+scripts=(
+    01-tilt-restart.sh
+    10-policy-create.sh
+    20-policy-rejects-zero-window.sh
+    30-consumption-list-empty.sh
+    40-grpc-record-and-list.sh
+    50-grpc-allow-record-deny.sh
+)
+[ "$KEEP" -eq 0 ] && scripts+=(99-cleanup.sh)
+
+passed=0
+total=${#scripts[@]}
+
+trap '[ "$KEEP" -eq 0 ] && bash ./99-cleanup.sh >/dev/null 2>&1 || true' EXIT
+
+for s in "${scripts[@]}"; do
+    echo
+    echo "==================== $s ===================="
+    if bash "./$s"; then
+        passed=$((passed + 1))
+    else
+        rc=$?
+        echo
+        echo "=== FAIL $passed/$total ($s exit=$rc) ==="
+        exit "$rc"
+    fi
+done
+
+echo
+echo "=== PASS $passed/$total ==="