Compare commits

...

20 Commits

Author SHA1 Message Date
Zoltán Papp
24287577c8 [client] Cover both invalid-ID orderings in the partial-error test
The partial-failure test only placed the unknown ID after the valid
exit node, so it never verified that selection processing continues
past a leading invalid ID. Run the same assertions for both orderings.
2026-07-27 22:35:45 +02:00
Zoltán Papp
7e1b35037a [client] Keep exit-node exclusivity on select-all and partial errors
Selecting all routes wipes every explicit selection, so exit nodes fall
back to management's auto-apply flags, which may mark several at once.
Reconcile immediately after select-all so at most one stays active
instead of waiting for the next network map.

A partial selection failure (e.g. an unknown ID in the request) still
selects the valid routes, so run the sibling exit-node deselection
regardless of the error and report both failures together.
2026-07-27 22:03:20 +02:00
Zoltán Papp
cc0ab3247d [client] Unify route selection in the route manager
Move route select/deselect handling from the daemon server into exported
routemanager methods (SelectRoutes, DeselectRoutes, SelectAllRoutes,
DeselectAllRoutes) so every consumer shares one implementation: v4/v6
exit-pair expansion, exit-node mutual exclusion, and selection triggering.

Previously the exit-node exclusivity lived only in the daemon's
SelectNetworks RPC, so the Android and iOS bindings could leave two exit
nodes selected until the next network map reconciliation. Both bindings
now call the shared manager methods and enforce exclusivity at toggle
time, matching the desktop behavior.
2026-07-27 21:20:38 +02:00
Maycon Santos
bab5572a74 [management, proxy] scope agent-network model allowlist per policy/group and provider (#6905)
## Describe your changes

The model-allowlist guardrail was merged into one account-wide union and
enforced flat on every request, ignoring which policy/group/provider
authorised it. With multiple policies — especially a mix of guardrailed
and un-guardrailed ones — this caused:

- **false-allow**: a model allowlisted for one group/provider leaked to
any caller; and
- **false-deny**: an un-guardrailed policy (intended unrestricted) was
blocked by another policy's allowlist.

Enforcement is now policy/group-aware, mirroring `llm_limit_check`:

- **Management (`SelectPolicyForRequest`) is authoritative.** It uses
the request model (already carried in
`CheckLLMPolicyLimitsRequest.model`, previously ignored) to keep only
applicable policies whose guardrails permit the model; no
allowlist-enabled guardrail = unrestricted. Denies
`llm_policy.model_blocked` when policies govern the (provider, groups)
but none permits the model.
- **Proxy `llm_guardrail` becomes a per-provider fail-closed backstop.**
The synthesiser emits an allowlist only for providers every authorising
policy restricts; the middleware keys off the resolved provider id and
keeps unknown-model fail-closed.
2026-07-27 20:43:59 +02:00
Zoltan Papp
9b4a5df925 [client] Use platform installer URL for manual update downloads (#6922)
The Wails UI regressed the non-enforced update download to the plain
GitHub releases page. Restore the old Fyne behavior: the tray update
item and the About card's Get installer button now open
version.DownloadUrl(), which points to the direct installer download
(pkgs.netbird.io) per OS/arch and falls back to the generic install page
where no installer exists.

Also fix the dead brew detection on darwin: exec.Command passed the
whole pipeline as a single argument to brew, so the check always failed.
Query the netbird formula and netbird-ui cask explicitly via exit codes
instead.

## Describe your changes

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [x] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)
- [ ] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature
— **OR** I have discussed it with the NetBird team beforehand (link the
issue / Slack thread in the description). See
[CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first).

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6922"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787764680&installation_model_id=427504&pr_number=6922&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6922&signature=3d9d640b25a17589114cc704e1d631895ffcd9f554947ff5790fe90d54410d2f"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added a direct installer download option for manual application
updates.
* Non-enforced updates now open the appropriate platform-specific
installer instead of the general releases page.

* **Bug Fixes**
  * Improved macOS download behavior for Homebrew installations.
* Preserved architecture-specific downloads for Intel and Apple silicon
Macs.
* Updated the tray “About” links to use the GitHub repository and
documentation instead of the releases page.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-27 20:08:35 +02:00
Misha Bragin
1816a020c4 [management, proxy] Add Claude Opus 5 (#6895) 2026-07-27 16:15:36 +02:00
Zoltan Papp
aa13928b76 [client] Export agent version info for iOS (#6918)
## Describe your changes

Export agent version info for iOS

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [ ] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [x] It is a refactor
- [ ] Created tests that fail without the change (if possible)
- [ ] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature
— **OR** I have discussed it with the NetBird team beforehand (link the
issue / Slack thread in the description). See
[CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first).

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6918"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787750541&installation_model_id=427504&pr_number=6918&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6918&signature=c2de74d89c36b01aac5866422b0be0b7525f7c6e26e46174efc280339eb864b6"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->
2026-07-27 15:53:09 +02:00
Viktor Liu
d681670a9d [misc] Restore the rootless-latest docker tag (#6914) 2026-07-27 12:07:58 +02:00
Maycon Santos
4f6247b5c3 [management, proxy] Add prompt-cache token and cost accounting to agent network usage (#6900)
Co-authored-by: braginini <bangvalo@gmail.com>
2026-07-26 21:42:41 +02:00
Riccardo Manfrin
1e5b0a5c89 [client] Make Test_ConnectPeers deterministic under Docker/eBPF kernel / Darwin CI (#6884)
## Describe your changes

Make `Test_ConnectPeers` deterministic. Two issues, both surfaced once
the
privileged suite moved into a `--privileged` Docker container (#6425):

1. The peers used `getLocalIP()` as their WireGuard endpoint, i.e. the
host's
routable NIC IP (the docker bridge IP `172.17.0.2` in CI). That address
might
not hairpin reliably inside the container, so the handshake
intermittently
timed out (flaky). Use loopback (`127.1.0.x`) instead — always
self-reachable.
2. On a Linux runner with the WG kernel module the iface uses the eBPF
proxy
factory. Its manager is a singleton with one shared XDP program +
settings
map, so bringing up the two ifaces makes the second factory overwrite
the
first's `wg_port`/`proxy_port` and the handshake is dropped. The test is
   incompatible with the eBPF factory, so disable it via
`NB_DISABLE_EBPF_WG_PROXY` (peers then handshake directly over
loopback).
Running the suite across all three modes (eBPF / UDP proxy / ICE bind)
would
   need a larger refactor.

Also fixes a typo in the `ErrSharedSockStopped` message (`socked` →
`socket`).

## Issue ticket number and link

No public issue — CI flakiness follow-up to #6871 on `Test_ConnectPeers`

(https://github.com/netbirdio/netbird/blob/main/client/iface/iface_test.go).

## Stack

<!-- branch-stack -->

### Checklist
- [x] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

Test-only change (plus a log-string typo). No public API, CLI, config,
or
behavior change.

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

N/A

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6884"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787485967&installation_model_id=427504&pr_number=6884&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6884&signature=09fb996715d039bd02775a140e8cc03deca5cb42540790b82a744527264a30ad"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Corrected the “shared socket stopped” error message for clearer
output.
* **Tests**
* Improved peer connection test reliability in privileged CI by
disabling the eBPF WireGuard proxy and using fixed loopback UDP
endpoints for deterministic setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-24 15:23:22 +02:00
Riccardo Manfrin
b65ec8b68a [client] Restores lost backup.Reset from #4935 (#6883)
## Describe your changes

Restore the management gRPC client's backoff reset after a stream
connects.

This is a regression: PR #4935 originally added this reset; it was then
lost in
commit `58daa674e` during the `withMgmtStream` refactor (Sync/Job
unification).

Compared to before the reset was done AFTER receiveEvents, whereas 
now it's done before. 

Now should cover for Sync and Job (which didn't exist in #4935. Hold for
long living stream that breaks.

Reset prevents two things

1. long living conns going wrong from waiting a long backoff time window
before to drive reconn
2. past MaxElapsedTime (3months) long lived conns failures from being
treated as "unrecoverable" (hence triggering an full restart of the
engine)

## Issue ticket number and link

Internal support case (customer agents lost data-plane connectivity
during a
management maintenance/release window). No public issue. Regression
introduced
in commit `58daa674e`, which removed the `backOff.Reset()` added by PR
#4935.

## Stack

<!-- branch-stack -->

### Checklist
- [x] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

Internal reconnection-behavior fix in the management gRPC client. No
public
NetBird CLI, API, or configuration surface changes.

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

N/A

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6883"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787484622&installation_model_id=427504&pr_number=6883&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6883&signature=525b08edbcb94ee6f4ae4226fb1fe8ddb829caa7d4489e8efa639770af9dccdf"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
  * Improved management stream retry behavior.
* Retry delays now reset after a stream is successfully established,
helping subsequent connection attempts recover more quickly.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-24 13:53:53 +02:00
Zoltan Papp
e13bcdbd44 [client] Fetch FreeBSD port files from GitHub mirror instead of cgit (#6880)
cgit.freebsd.org now sits behind an Anubis anti-bot challenge and
intermittently returns an HTML challenge page with HTTP 200 instead of
the requested file, breaking the FreeBSD port release job. Fetch the
port Makefile and distinfo from the official freebsd-ports GitHub
mirror, retry transient failures, and fail loudly if HTML is returned
instead of the expected file.

## Describe your changes

## Issue ticket number and link

## Stack

<!-- branch-stack -->

### Checklist
- [x] Is it a bug fix
- [ ] Is a typo/documentation fix
- [ ] Is a feature enhancement
- [ ] It is a refactor
- [ ] Created tests that fail without the change (if possible)
- [ ] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature
— **OR** I have discussed it with the NetBird team beforehand (link the
issue / Slack thread in the description). See
[CONTRIBUTING.md](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTING.md#discuss-changes-with-the-netbird-team-first).

> By submitting this pull request, you confirm that you have read and
agree to the terms of the [Contributor License
Agreement](https://github.com/netbirdio/netbird/blob/main/CONTRIBUTOR_LICENSE_AGREEMENT.md).

## Documentation
Select exactly one:

- [ ] I added/updated documentation for this change
- [x] Documentation is **not needed** for this change (explain why)

### Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:

https://github.com/netbirdio/docs/pull/__

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6880"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787478670&installation_model_id=427504&pr_number=6880&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6880&signature=503a5d79f1671c21cebedac66516c7d7298d401a6382d9bf483813f9e7ee7591"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved reliability when retrieving FreeBSD port metadata during
release and issue preparation.
* Added automatic retries and HTTPS-only redirect handling for safer
downloads.
* Validates fetched content to detect unexpected HTML responses and
avoids processing invalid data.
* Updated port metadata retrieval to use a more reliable mirror,
improving version extraction and the resulting comparisons and
regenerated release information.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-24 13:10:53 +02:00
Riccardo Manfrin
46568f7af8 [client] Reconcile routed allowed IPs when a lazy connection goes idle (#6863)
## Describe your changes

Under lazy connections, when a routing peer goes idle its WireGuard peer
is torn down and re-created with a wake endpoint by
the activity listener, carrying only the overlay /32
(`peerCfg.AllowedIPs`). The routed subnet prefixes are dropped from
  the device on the Connected→Idle transition.

They are meant to be restored by the route watcher, which reacts to the
peer's status change and calls `recalculateRoutes` →
  `AddAllowedIP`. Two things prevent that from healing the peer:

- `AddAllowedIP` uses `update_only`, which is a silent no-op (no error)
when the peer does not exist. While the peer is
being torn down and re-armed with its wake endpoint, it is briefly
absent, so a re-add that lands in that window is lost.
- The allowed-IP refcounter only calls its add function on a prefix's
0→1 transition. The routed prefix stays referenced
across the idle cycle, so once the device entry is gone the refcounter
does not re-push it on its own, and nothing retries.

As a result, traffic to the routed subnet is black-holed while the peer
is idle. Because the wake endpoint only fires when a
packet matches the peer's AllowedIPs, a packet to the subnet is dropped
before reaching the wake endpoint, so it cannot
wake the peer. The peer only recovers when woken by other means (e.g. a
ping to its overlay IP).

  ## Approach

This change keeps the existing Connected→Idle transition as-is and
reconciles the AllowedIPs afterwards, avoiding any
additional locking on the transition path. The peer is torn down and
re-armed with its wake endpoint as today; the routed
prefixes are then re-applied from the route manager's allowed-IP
refcounter once the wake endpoint has been (re)armed.

A single add-only method, `ReconcilePeerAllowedIPs(peerKey)`, re-applies
every routed prefix currently tracked for the peer
in the refcounter (the authoritative store; it already covers static,
dynamic and dnsinterceptor routes). It runs whenever
the peer's wake endpoint is (re)created in the lazy manager — every
point where the activity listener builds it with the
  overlay /32 only:

- **initial registration** (`AddPeer`, cold start): the route manager
may have already pushed the peer's routes before the
wake endpoint existed, so those `AddAllowedIP` calls no-op'd; the
reconcile installs them on the freshly created wake
  endpoint.
- **the two paths into idle** (`DeactivatePeer` on a remote GOAWAY,
`onPeerInactivityTimedOut` on local inactivity): the
peer is torn down and re-armed, so the routed prefixes must be
re-applied.

In every case the routed prefixes end up on the wake endpoint, so
traffic to a routed subnet can wake the peer. Arming the
wake endpoint and reconciling are wrapped in a single
`armActivityListener` helper so the two always happen together.

New helper: `refcounter.Counter.KeysMatching(pred)` to enumerate a
peer's prefixes under the counter lock.

Note on scope: the reconcile restores what the refcounter tracks. All
routed AllowedIPs currently go through it, so this
covers the routed-prefix case; it does not attempt to reconcile
AllowedIPs installed outside the refcounter. The
Idle→Connected (wake) path does not need this: the peer is not removed
there (the listener close leaves it in place and only
the endpoint is updated), so a concurrent `AddAllowedIP` lands normally.

  ## Testing

Reproduced deterministically in a local dev setup (userspace client,
`NB_WG_KERNEL_DISABLED=true`, `B_LAZY_CONN_INACTIVITY_THRESHOLD=1`
inactivity threshold 1
min). A temporary 30s sleep in the tear-down → re-arm window widens the
race so the route watcher's async `AddAllowedIP`
reliably lands while the peer is absent and no-ops (the sleep is a test
aid, not part of the change):

- **without the reconcile:** after the peer goes idle, a ping to any
routed IP — both a pre-existing route and one added
  during the window — black-holes; the peer never wakes.
  - **with the reconcile:** the same ping wakes the peer and passes.

Added unit tests: `ReconcilePeerAllowedIPs` (re-applies all of a peer's
tracked prefixes, scoped to that peer) and
  `refcounter.Counter.KeysMatching`.

Note: `netbird status -d` is not a reliable signal for this —
`AddPeerStateRoute` records the route regardless of whether
the underlying `AddAllowedIP` no-op'd, so it reflects the route
manager's intent rather than device state. The reliable
  signal is functional (ping the subnet from idle).
  
  ## Checklist

  - [x] Is it a bug fix
  - [ ] Is a typo/documentation fix
  - [ ] Is a feature enhancement
  - [ ] It is a refactor
- [x] Created tests that fail without the change (unit tests for the
reconcile + `KeysMatching`)

  ## Documentation

- [x] Documentation is **not needed** for this change (internal client
behavior, no API / gRPC / CLI / flag change)

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6863"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with Codesmith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787330960&installation_model_id=427504&pr_number=6863&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6863&signature=f3d6a97d7db82e92b3939fdd0f159c5ee74913ff88f4eb82e41e88fcb787aff4"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with Codesmith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>/codesmith</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Routed IP assignments are automatically reconciled and restored
whenever a peer’s lazy wake endpoint is armed or re-armed.
* Routed allowed IPs are re-applied after inactivity transitions and
monitoring re-initialization.
* If reconciliation can’t be performed, the client safely skips it; if
reconciliation encounters issues, failures are logged without stopping
connection monitoring.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-23 18:40:54 +02:00
Riccardo Manfrin
3358138ccc [client] Fix flaky Test_ConnectPeers busy-loop handshake wait (#6871)
## Describe your changes

  `Test_ConnectPeers` in `client/iface` is flaky in CI, timing out with
  `waiting for peer handshake timeout after 30s`.

  The wait loop polled `getPeer()` in a tight busy-loop with no sleep
(a `select` with a `default` branch), pegging a CPU core. The peers run
  userspace WireGuard (stdnet transport), so the spin starved the
wireguard-go goroutines that actually process the handshake, making the
  30s wait flaky under CI load.

Poll on a 500ms ticker instead so the CPU is yielded between checks, and
  check the handshake state before waiting. Same logic, no busy-spin.

  ## Issue ticket number and link

  N/A — CI flakiness fix.

  ### Checklist
  - [x] Is it a bug fix
  - [ ] Is a typo/documentation fix
  - [ ] Is a feature enhancement
  - [ ] It is a refactor
  - [ ] Created tests that fail without the change (if possible)
- [x] This change does **not** modify the public API, gRPC protocols,
functionality behavior, CLI / service flags, or introduce a new feature

  ## Documentation

  - [x] Documentation is **not needed** for this change (test-only fix)

<!-- codesmith:footer -->
---
<a
href="https://app.blacksmith.sh/netbirdio/codesmith/netbird/pr/6871"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-light-v2.svg"><img
alt="View with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/view-with-codesmith-dark-v2.svg"></picture></a>
<a
href="https://backend.blacksmith.sh/track/enable-autofix?expires=1787405146&installation_model_id=427504&pr_number=6871&repository=netbirdio%2Fnetbird&return_to=https%3A%2F%2Fgithub.com%2Fnetbirdio%2Fnetbird%2Fpull%2F6871&signature=a68daf1747583280b5d1d083c6b45ade4c0d441a38fd25ab9a2ca01f87511ad9"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"><source
media="(prefers-color-scheme: light)"
srcset="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-light.svg"><img
alt="Autofix with [code]smith"
src="https://pr-comments-assets.blacksmith.sh/codesmith/autofix-with-codesmith-dark.svg"></picture></a>
<sup>Need help on this PR? Tag <code>@codesmith-bot</code> with what you
need. Autofix is disabled.</sup>

<!-- codesmith:autofix:disabled -->
<!-- /codesmith:footer -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Tests**
* Improved peer connection test waiting behavior to avoid busy spinning.
* Added a timeout and periodic checks, with clearer failure handling
when a handshake does not complete.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-07-23 15:48:16 +02:00
Viktor Liu
6d15d0729a [client] Clear stale UDP checksum in eBPF XDP proxy after port rewrite (#6861) 2026-07-23 14:37:52 +02:00
Maycon Santos
0936918d24 [management, proxy] add Kimi (Moonshot AI) provider to Agent Network (#6853) 2026-07-23 13:32:32 +02:00
Bethuel Mmbaga
d4a4418969 [misc] Simplify enterprise bootstrap (#6869) 2026-07-23 12:41:13 +03:00
Pascal Fischer
31ed241a1a [management] component types (#6866) 2026-07-23 11:37:09 +02:00
Maycon Santos
178e6a8530 [misc] Probe the agent-network endpoint with a GET instead of getent (#6867)
## Describe your changes

Use http get to trigger lazy connections as e2e will have single proxy
and status checks would fail without proper lazy wake-up
2026-07-23 15:30:40 +09:00
Maycon Santos
96963b6751 [proxy] fix proxy multistage build (#6864) 2026-07-22 21:24:44 +02:00
128 changed files with 5214 additions and 1079 deletions

View File

@@ -5,6 +5,13 @@ on:
schedule:
- cron: "0 3 * * *"
workflow_dispatch:
inputs:
bedrock_model:
description: >-
Bedrock inference-profile id to drive the matrix with, exactly as
AWS issues it. Leave empty for the Sonnet 4.6 default.
required: false
default: ""
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
@@ -51,6 +58,9 @@ jobs:
# token (and URL, for gateways) is unset, so partial coverage is fine.
OPENAI_TOKEN: ${{ secrets.E2E_OPENAI_TOKEN }}
ANTHROPIC_TOKEN: ${{ secrets.E2E_ANTHROPIC_TOKEN }}
# Moonshot AI platform key (platform.kimi.ai); drives both Kimi wire
# shapes (OpenAI /v1 and Anthropic /anthropic) through kimi_api.
KIMI_TOKEN: ${{ secrets.E2E_KIMI_TOKEN }}
VERCEL_URL: ${{ secrets.E2E_VERCEL_URL }}
VERCEL_TOKEN: ${{ secrets.E2E_VERCEL_TOKEN }}
OPENROUTER_URL: ${{ secrets.E2E_OPENROUTER_URL }}
@@ -59,6 +69,8 @@ jobs:
CLOUDFLARE_TOKEN: ${{ secrets.E2E_CLOUDFLARE_TOKEN }}
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.E2E_AWS_BEARER_TOKEN_BEDROCK }}
AWS_REGION: ${{ secrets.E2E_AWS_REGION }}
# Bedrock model override: dispatch input wins, then the repo variable, else the test default.
AWS_BEDROCK_MODEL: ${{ inputs.bedrock_model || vars.E2E_AWS_BEDROCK_MODEL }}
# Vertex (Anthropic-on-Vertex): SA + project required; region defaults
# to "global", model to a pinned claude snapshot.
GOOGLE_VERTEX_SA_BASE64: ${{ secrets.E2E_GOOGLE_VERTEX_SA_BASE64 }}

View File

@@ -273,8 +273,8 @@ dockers_v2:
- netbirdio/netbird
- ghcr.io/netbirdio/netbird
tags:
- "v{{ .Version }}-rootless"
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}latest{{ end }}"
- "{{ .Version }}-rootless"
- "{{ if eq .Env.SKIP_PUBLISH \"false\" }}rootless-latest{{ end }}"
dockerfile: client/Dockerfile-rootless
extra_files:
- client/netbird-entrypoint.sh

View File

@@ -439,10 +439,6 @@ func (c *Client) RemoveConnectionListener() {
c.recorder.RemoveConnectionListener()
}
func (c *Client) toggleRoute(command routeCommand) error {
return command.toggleRoute()
}
func (c *Client) getRouteManager() (routemanager.Manager, error) {
client := c.getConnectClient()
if client == nil {
@@ -462,22 +458,22 @@ func (c *Client) getRouteManager() (routemanager.Manager, error) {
return manager, nil
}
func (c *Client) SelectRoute(route string) error {
func (c *Client) SelectRoute(id string) error {
manager, err := c.getRouteManager()
if err != nil {
return err
}
return c.toggleRoute(selectRouteCommand{route: route, manager: manager})
return manager.SelectRoutes([]route.NetID{route.NetID(id)}, true)
}
func (c *Client) DeselectRoute(route string) error {
func (c *Client) DeselectRoute(id string) error {
manager, err := c.getRouteManager()
if err != nil {
return err
}
return c.toggleRoute(deselectRouteCommand{route: route, manager: manager})
return manager.DeselectRoutes([]route.NetID{route.NetID(id)})
}
// getNetworkDomainsFromRoute extracts domains from a route and enriches each domain

View File

@@ -1,70 +0,0 @@
//go:build android
package android
import (
"fmt"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
"github.com/netbirdio/netbird/client/internal/routemanager"
"github.com/netbirdio/netbird/route"
)
func executeRouteToggle(id string, manager routemanager.Manager,
operationName string,
routeOperation func(routes []route.NetID, allRoutes []route.NetID) error) error {
netID := route.NetID(id)
routes := []route.NetID{netID}
routesMap := manager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
log.Debugf("%s with ids: %v", operationName, routes)
if err := routeOperation(routes, maps.Keys(routesMap)); err != nil {
log.Debugf("error when %s: %s", operationName, err)
return fmt.Errorf("error %s: %w", operationName, err)
}
manager.TriggerSelection(manager.GetClientRoutes())
return nil
}
type routeCommand interface {
toggleRoute() error
}
type selectRouteCommand struct {
route string
manager routemanager.Manager
}
func (s selectRouteCommand) toggleRoute() error {
routeSelector := s.manager.GetRouteSelector()
if routeSelector == nil {
return fmt.Errorf("no route selector available")
}
routeOperation := func(routes []route.NetID, allRoutes []route.NetID) error {
return routeSelector.SelectRoutes(routes, true, allRoutes)
}
return executeRouteToggle(s.route, s.manager, "selecting route", routeOperation)
}
type deselectRouteCommand struct {
route string
manager routemanager.Manager
}
func (d deselectRouteCommand) toggleRoute() error {
routeSelector := d.manager.GetRouteSelector()
if routeSelector == nil {
return fmt.Errorf("no route selector available")
}
return executeRouteToggle(d.route, d.manager, "deselecting route", routeSelector.DeselectRoutes)
}

View File

@@ -464,6 +464,8 @@ func Test_RemovePeer(t *testing.T) {
}
func Test_ConnectPeers(t *testing.T) {
t.Setenv("NB_DISABLE_EBPF_WG_PROXY", "true")
peer1ifaceName := fmt.Sprintf("utun%d", WgIntNumber+400)
peer1wgIP := netip.MustParsePrefix("10.99.99.17/30")
peer1Key, _ := wgtypes.GeneratePrivateKey()
@@ -505,12 +507,8 @@ func Test_ConnectPeers(t *testing.T) {
t.Fatal(err)
}
localIP, err := getLocalIP()
if err != nil {
t.Fatal(err)
}
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer1wgPort))
localIP1 := "127.0.0.1"
peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP1, peer1wgPort))
if err != nil {
t.Fatal(err)
}
@@ -546,7 +544,8 @@ func Test_ConnectPeers(t *testing.T) {
t.Fatal(err)
}
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer2wgPort))
localIP2 := "127.0.0.1"
peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP2, peer2wgPort))
if err != nil {
t.Fatal(err)
}
@@ -569,17 +568,17 @@ func Test_ConnectPeers(t *testing.T) {
if err != nil {
t.Fatal(err)
}
// todo: investigate why in some tests execution we need 30s
// The peers use userspace WireGuard (stdnet transport). A tight busy-loop
// here starves the wireguard-go goroutines that process the handshake, so
// poll on a ticker instead and yield the CPU between checks. WireGuard also
// only retries a lost handshake initiation every REKEY_TIMEOUT (5s), which
// is why the overall wait can occasionally stretch to tens of seconds.
timeout := 30 * time.Second
timeoutChannel := time.After(timeout)
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-timeoutChannel:
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
default:
}
peer, gpErr := getPeer(peer1ifaceName, peer2Key.PublicKey().String())
if gpErr != nil {
t.Fatal(gpErr)
@@ -588,6 +587,12 @@ func Test_ConnectPeers(t *testing.T) {
t.Log("peers successfully handshake")
break
}
select {
case <-timeoutChannel:
t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
case <-ticker.C:
}
}
}
@@ -615,28 +620,3 @@ func getPeer(ifaceName, peerPubKey string) (wgtypes.Peer, error) {
}
return wgtypes.Peer{}, fmt.Errorf("peer not found")
}
func getLocalIP() (string, error) {
// Get all interfaces
addrs, err := net.InterfaceAddrs()
if err != nil {
return "", err
}
for _, addr := range addrs {
ipNet, ok := addr.(*net.IPNet)
if !ok {
continue
}
if ipNet.IP.IsLoopback() {
continue
}
if ipNet.IP.To4() == nil {
continue
}
return ipNet.IP.String(), nil
}
return "", fmt.Errorf("no local IP found")
}

View File

@@ -49,11 +49,21 @@ type ConnMgr struct {
// engine.syncMsgMux; all other reads stay under engine.syncMsgMux only.
lazyConnMgrMu sync.RWMutex
// reconcileRoutedIPs re-applies a peer's routed allowed IPs after its lazy wake endpoint is
// (re)armed (Mode A at arm time). Injected by the engine; nil disables the reconcile.
reconcileRoutedIPs func(peerKey string) error
wg sync.WaitGroup
lazyCtx context.Context
lazyCtxCancel context.CancelFunc
}
// SetRoutedIPsReconciler injects the callback used to re-apply a peer's routed allowed IPs when
// its lazy wake endpoint is (re)armed. Must be called before the lazy manager starts.
func (e *ConnMgr) SetRoutedIPsReconciler(fn func(peerKey string) error) {
e.reconcileRoutedIPs = fn
}
func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr {
e := &ConnMgr{
peerStore: peerStore,
@@ -291,6 +301,7 @@ func (e *ConnMgr) Close() {
func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
cfg := manager.Config{
InactivityThreshold: inactivityThresholdEnv(),
ReconcileAllowedIPs: e.reconcileRoutedIPs,
}
e.lazyConnMgrMu.Lock()

View File

@@ -52,11 +52,14 @@ int xdp_dns_fwd(struct iphdr *ip, struct udphdr *udp) {
if (udp->dest == GENERAL_DNS_PORT && ip->daddr == dns_ip) {
udp->dest = dns_port;
// Clear the now-stale checksum; zero means "not computed" for IPv4.
udp->check = 0;
return XDP_PASS;
}
if (udp->source == dns_port && ip->saddr == dns_ip) {
udp->source = GENERAL_DNS_PORT;
udp->check = 0;
return XDP_PASS;
}

View File

@@ -50,5 +50,11 @@ int xdp_wg_proxy(struct iphdr *ip, struct udphdr *udp) {
__be16 new_dst_port = htons(proxy_port);
udp->dest = new_dst_port;
udp->source = new_src_port;
// The ports are covered by the UDP checksum. This is an IPv4 loopback hop
// and the payload is already integrity-protected, so clear the checksum (a
// zero UDP checksum means "not computed" for IPv4) rather than leave a
// stale value the kernel would drop as UDP_CSUM.
udp->check = 0;
return XDP_PASS;
}

View File

@@ -663,6 +663,12 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
iceCfg := e.createICEConfig()
e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface)
e.connMgr.SetRoutedIPsReconciler(func(peerKey string) error {
if e.routeManager == nil {
return nil
}
return e.routeManager.ReconcilePeerAllowedIPs(peerKey)
})
e.connMgr.Start(e.ctx)
// Wire DNS-time lazy-connection warm-up now that the connection manager

View File

@@ -29,6 +29,11 @@ type managedPeer struct {
type Config struct {
InactivityThreshold *time.Duration
// ReconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is
// armed. The activity listener creates the wake peer with the overlay /32 only; without the
// routed prefixes WireGuard would not steer subnet-bound traffic to the wake endpoint, so an
// idle routing peer could never be woken by that traffic. Optional; nil disables the reconcile.
ReconcileAllowedIPs func(peerKey string) error
}
// Manager manages lazy connections
@@ -56,6 +61,9 @@ type Manager struct {
peerToHAGroups map[string][]route.HAUniqueID // peer ID -> HA groups they belong to
haGroupToPeers map[route.HAUniqueID][]string // HA group -> peer IDs in the group
routesMu sync.RWMutex
// reconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is armed.
reconcileAllowedIPs func(peerKey string) error
}
// NewManager creates a new lazy connection manager
@@ -73,6 +81,7 @@ func NewManager(config Config, engineCtx context.Context, peerStore *peerstore.S
activityManager: activity.NewManager(wgIface),
peerToHAGroups: make(map[string][]route.HAUniqueID),
haGroupToPeers: make(map[route.HAUniqueID][]string),
reconcileAllowedIPs: config.ReconcileAllowedIPs,
}
if wgIface.IsUserspaceBind() {
@@ -201,7 +210,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
return false, nil
}
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
if err := m.armActivityListener(peerCfg); err != nil {
return false, err
}
@@ -288,7 +297,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) {
m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
if err := m.armActivityListener(*mp.peerCfg); err != nil {
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
return
}
@@ -465,6 +474,31 @@ func (m *Manager) close() {
}
// shouldDeferIdleForHA checks if peer should stay connected due to HA group requirements
// armRoutedAllowedIPs re-applies the peer's routed allowed IPs onto its freshly armed wake
// endpoint. The activity listener creates the wake peer with the overlay /32 only, so without
// this the routed prefixes would be missing and traffic to a routed subnet could not wake the
// idle routing peer. It is a no-op when no reconciler is configured.
// armActivityListener (re)arms the peer's wake endpoint via the activity manager and then
// re-applies its routed allowed IPs, so traffic to a routed subnet can wake an idle routing
// peer. The routed prefixes must be re-applied after the wake endpoint exists because the
// listener creates it with the overlay /32 only.
func (m *Manager) armActivityListener(peerCfg lazyconn.PeerConfig) error {
if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
return err
}
m.armRoutedAllowedIPs(&peerCfg)
return nil
}
func (m *Manager) armRoutedAllowedIPs(peerCfg *lazyconn.PeerConfig) {
if m.reconcileAllowedIPs == nil {
return
}
if err := m.reconcileAllowedIPs(peerCfg.PublicKey); err != nil {
peerCfg.Log.Errorf("failed to reconcile routed allowed IPs on wake endpoint: %v", err)
}
}
func (m *Manager) shouldDeferIdleForHA(inactivePeers map[string]struct{}, peerID string) bool {
m.routesMu.RLock()
defer m.routesMu.RUnlock()
@@ -577,7 +611,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
mp.peerCfg.Log.Infof("start activity monitor")
if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
if err := m.armActivityListener(*mp.peerCfg); err != nil {
mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
continue
}

View File

@@ -52,6 +52,10 @@ type Manager interface {
UpdateRoutes(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
ClassifyRoutes(newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
TriggerSelection(route.HAMap)
SelectRoutes(ids []route.NetID, appendRoute bool) error
DeselectRoutes(ids []route.NetID) error
SelectAllRoutes()
DeselectAllRoutes()
GetRouteSelector() *routeselector.RouteSelector
GetClientRoutes() route.HAMap
GetSelectedClientRoutes() route.HAMap
@@ -61,6 +65,7 @@ type Manager interface {
InitialRouteRange() []string
SetFirewall(firewall.Manager) error
SetDNSForwarderPort(port uint16)
ReconcilePeerAllowedIPs(peerKey string) error
Stop(stateManager *statemanager.Manager)
}
@@ -232,6 +237,30 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
)
}
// ReconcilePeerAllowedIPs re-applies every routed allowed IP currently tracked for the peer
// onto the WireGuard device. The allowed-IP refcounter only calls its AddFunc (which pushes to
// the device) on a prefix's 0->1 transition, so a peer whose device entry was rebuilt without a
// matching refcounter change — e.g. a lazy connection cycling through idle->wake, which recreates
// the WireGuard peer with the overlay /32 only — ends up missing routed prefixes the refcounter
// still considers installed, and nothing retries. Calling this when the peer's WireGuard entry is
// (re)created restores convergence. It is add-only and idempotent: AddAllowedIP is update-only, so
// prefixes are re-added to an existing peer and an absent peer is left untouched.
func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error {
if m.allowedIPsRefCounter == nil {
return nil
}
return m.allowedIPsRefCounter.ReapplyMatching(
func(out string) bool { return out == peerKey },
func(prefix netip.Prefix) error {
if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil {
return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)
}
return nil
},
)
}
// Init sets up the routing
func (m *DefaultManager) Init() error {
m.routeSelector = m.initSelector()
@@ -775,7 +804,7 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
var info exitNodeInfo
for haID, routes := range clientRoutes {
if !m.isExitNodeRoute(routes) {
if !isExitNodeRoutes(routes) {
continue
}
@@ -795,13 +824,6 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
return info
}
func (m *DefaultManager) isExitNodeRoute(routes []*route.Route) bool {
if len(routes) == 0 {
return false
}
return route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network)
}
func (m *DefaultManager) categorizeUserSelection(netID route.NetID, info *exitNodeInfo) {
if m.routeSelector.IsSelected(netID) {
info.userSelected = append(info.userSelected, netID)

View File

@@ -16,6 +16,8 @@ type MockManager struct {
ClassifyRoutesFunc func(routes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
UpdateRoutesFunc func(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
TriggerSelectionFunc func(haMap route.HAMap)
SelectRoutesFunc func(ids []route.NetID, appendRoute bool) error
DeselectRoutesFunc func(ids []route.NetID) error
GetRouteSelectorFunc func() *routeselector.RouteSelector
GetClientRoutesFunc func() route.HAMap
GetSelectedClientRoutesFunc func() route.HAMap
@@ -55,6 +57,30 @@ func (m *MockManager) TriggerSelection(networks route.HAMap) {
}
}
// SelectRoutes mock implementation of SelectRoutes from Manager interface
func (m *MockManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
if m.SelectRoutesFunc != nil {
return m.SelectRoutesFunc(ids, appendRoute)
}
return nil
}
// DeselectRoutes mock implementation of DeselectRoutes from Manager interface
func (m *MockManager) DeselectRoutes(ids []route.NetID) error {
if m.DeselectRoutesFunc != nil {
return m.DeselectRoutesFunc(ids)
}
return nil
}
// SelectAllRoutes mock implementation of SelectAllRoutes from Manager interface
func (m *MockManager) SelectAllRoutes() {
}
// DeselectAllRoutes mock implementation of DeselectAllRoutes from Manager interface
func (m *MockManager) DeselectAllRoutes() {
}
// GetRouteSelector mock implementation of GetRouteSelector from Manager interface
func (m *MockManager) GetRouteSelector() *routeselector.RouteSelector {
if m.GetRouteSelectorFunc != nil {
@@ -112,6 +138,11 @@ func (m *MockManager) SetFirewall(firewall.Manager) error {
func (m *MockManager) SetDNSForwarderPort(port uint16) {
}
// ReconcilePeerAllowedIPs mock implementation of ReconcilePeerAllowedIPs from Manager interface
func (m *MockManager) ReconcilePeerAllowedIPs(peerKey string) error {
return nil
}
// Stop mock implementation of Stop from Manager interface
func (m *MockManager) Stop(stateManager *statemanager.Manager) {
if m.StopFunc != nil {

View File

@@ -0,0 +1,90 @@
//go:build !windows
package routemanager
import (
"net"
"net/netip"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.zx2c4.com/wireguard/tun/netstack"
"github.com/netbirdio/netbird/client/iface/device"
"github.com/netbirdio/netbird/client/iface/wgaddr"
"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
)
// reconcileWGMock is a minimal iface.WGIface that only records AddAllowedIP calls; every other
// method is an inert stub because ReconcilePeerAllowedIPs exercises none of them.
type reconcileWGMock struct {
mu sync.Mutex
adds map[string][]netip.Prefix
}
func (m *reconcileWGMock) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
m.mu.Lock()
defer m.mu.Unlock()
if m.adds == nil {
m.adds = map[string][]netip.Prefix{}
}
m.adds[peerKey] = append(m.adds[peerKey], allowedIP)
return nil
}
func (m *reconcileWGMock) added(peerKey string) []netip.Prefix {
m.mu.Lock()
defer m.mu.Unlock()
return m.adds[peerKey]
}
func (m *reconcileWGMock) RemoveAllowedIP(string, netip.Prefix) error { return nil }
func (m *reconcileWGMock) Name() string { return "utun-test" }
func (m *reconcileWGMock) Address() wgaddr.Address { return wgaddr.Address{} }
func (m *reconcileWGMock) ToInterface() *net.Interface { return nil }
func (m *reconcileWGMock) IsUserspaceBind() bool { return false }
func (m *reconcileWGMock) GetFilter() device.PacketFilter { return nil }
func (m *reconcileWGMock) GetDevice() *device.FilteredDevice { return nil }
func (m *reconcileWGMock) GetNet() *netstack.Net { return nil }
// TestReconcilePeerAllowedIPs verifies the declarative reconcile re-applies every routed prefix
// tracked for the peer (self-heal, independent of refcount level) and stays scoped to that peer.
func TestReconcilePeerAllowedIPs(t *testing.T) {
wg := &reconcileWGMock{}
m := &DefaultManager{wgInterface: wg}
m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
func(netip.Prefix, string) error { return nil },
)
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
_, err := m.allowedIPsRefCounter.Increment(prefix, peer)
require.NoError(t, err)
}
// Extra reference: reconcile must still re-apply the prefix even though its refcount never
// hit 0 again (the exact case the plain incremental path skips).
_, err := m.allowedIPsRefCounter.Increment(peerA1, "peerA")
require.NoError(t, err)
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, wg.added("peerA"),
"reconcile must re-apply all routed prefixes of the peer")
assert.Empty(t, wg.added("peerB"), "reconcile must not touch another peer's prefixes")
}
// TestReconcilePeerAllowedIPsNoCounter verifies reconcile is a safe no-op before the refcounter is
// set up.
func TestReconcilePeerAllowedIPsNoCounter(t *testing.T) {
wg := &reconcileWGMock{}
m := &DefaultManager{wgInterface: wg}
require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
assert.Empty(t, wg.added("peerA"))
}

View File

@@ -94,6 +94,26 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
return ref, ok
}
// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the
// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect
// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its
// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied.
// pred and apply are invoked under the lock, so they must not call back into the counter.
func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error {
rm.mu.Lock()
defer rm.mu.Unlock()
var merr *multierror.Error
for key, ref := range rm.refCountMap {
if pred(ref.Out) {
if err := apply(key); err != nil {
merr = multierror.Append(merr, err)
}
}
}
return nberrors.FormatErrorOrNil(merr)
}
// Increment increments the reference count for the given key.
// If this is the first reference to the key, the AddFunc is called.
func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {

View File

@@ -0,0 +1,47 @@
package refcounter
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored
// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive
// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes.
func TestReapplyMatching(t *testing.T) {
rc := New[netip.Prefix, string, string](
func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
func(netip.Prefix, string) error { return nil },
)
peerA1 := netip.MustParsePrefix("10.0.0.0/24")
peerA2 := netip.MustParsePrefix("10.1.0.0/24")
peerB1 := netip.MustParsePrefix("10.2.0.0/24")
for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
_, err := rc.Increment(prefix, peer)
require.NoError(t, err)
}
// a second reference must not make the key applied twice
_, err := rc.Increment(peerA1, "peerA")
require.NoError(t, err)
var applied []netip.Prefix
err = rc.ReapplyMatching(
func(out string) bool { return out == "peerA" },
func(key netip.Prefix) error { applied = append(applied, key); return nil },
)
require.NoError(t, err)
assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied)
var none []netip.Prefix
err = rc.ReapplyMatching(
func(out string) bool { return out == "missing" },
func(key netip.Prefix) error { none = append(none, key); return nil },
)
require.NoError(t, err)
assert.Empty(t, none)
}

View File

@@ -0,0 +1,138 @@
package routemanager
import (
"fmt"
"slices"
"github.com/hashicorp/go-multierror"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
nberrors "github.com/netbirdio/netbird/client/errors"
"github.com/netbirdio/netbird/route"
)
// SelectRoutes selects the routes with the given network IDs and applies the
// new selection. V4/v6 exit-node pairs are expanded automatically. Exit nodes
// are mutually exclusive: if the selection activates an exit node, every other
// available exit node is deselected so two can't be active at once. With
// appendRoute=false the previous selection is replaced instead of extended.
func (m *DefaultManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
if err := m.selectRoutes(ids, appendRoute); err != nil {
return err
}
m.TriggerSelection(m.GetClientRoutes())
return nil
}
// DeselectRoutes removes the routes with the given network IDs from the
// selection and applies the change. V4/v6 exit-node pairs are expanded
// automatically.
func (m *DefaultManager) DeselectRoutes(ids []route.NetID) error {
if err := m.deselectRoutes(ids); err != nil {
return err
}
m.TriggerSelection(m.GetClientRoutes())
return nil
}
func (m *DefaultManager) deselectRoutes(ids []route.NetID) error {
routesMap := m.GetClientRoutesWithNetID()
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
log.Debugf("deselecting routes with ids: %v", routes)
if err := m.routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
return fmt.Errorf("deselect routes: %w", err)
}
return nil
}
// SelectAllRoutes selects every available route and applies the selection.
// Exit nodes stay mutually exclusive: at most one remains active.
func (m *DefaultManager) SelectAllRoutes() {
m.selectAllRoutes()
m.TriggerSelection(m.GetClientRoutes())
}
func (m *DefaultManager) selectAllRoutes() {
m.routeSelector.SelectAllRoutes()
// Select-all wipes every explicit selection, so exit nodes fall back to
// management's auto-apply flags — which may mark several at once.
// Reconcile immediately so at most one exit node stays active instead of
// waiting for the next network map to enforce it.
m.mux.Lock()
defer m.mux.Unlock()
m.updateRouteSelectorFromManagement(m.clientRoutes)
}
// DeselectAllRoutes deselects every route and applies the change.
func (m *DefaultManager) DeselectAllRoutes() {
m.routeSelector.DeselectAllRoutes()
m.TriggerSelection(m.GetClientRoutes())
}
func (m *DefaultManager) selectRoutes(ids []route.NetID, appendRoute bool) error {
routesMap := m.GetClientRoutesWithNetID()
routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
allIDs := maps.Keys(routesMap)
log.Debugf("selecting routes with ids: %v", routes)
// A partial failure (e.g. an unknown ID in the request) still selects the
// valid routes, so exclusivity below must run regardless of the error.
var merr *multierror.Error
if err := m.routeSelector.SelectRoutes(routes, appendRoute, allIDs); err != nil {
merr = multierror.Append(merr, fmt.Errorf("select routes: %w", err))
}
// Exit nodes are mutually exclusive: if this selection activates an
// exit node, deselect every other available exit node so two can't be
// selected at once. Non-exit route selections are left untouched.
if requestActivatesExitNode(routes, routesMap) {
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
if err := m.routeSelector.DeselectRoutes(others, allIDs); err != nil {
merr = multierror.Append(merr, fmt.Errorf("deselect sibling exit nodes: %w", err))
}
}
}
return nberrors.FormatErrorOrNil(merr)
}
func isExitNodeRoutes(routes []*route.Route) bool {
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
}
// requestActivatesExitNode reports whether any requested NetID maps to an exit
// node (default route) in the current route table.
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
for _, id := range requested {
if isExitNodeRoutes(routesMap[id]) {
return true
}
}
return false
}
// otherExitNodeIDs returns every available exit-node NetID that is not in the
// requested set — the siblings to deselect so a single exit node stays active.
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
keep := make(map[route.NetID]struct{}, len(requested))
for _, id := range requested {
keep[id] = struct{}{}
}
var others []route.NetID
for id, routes := range routesMap {
if !isExitNodeRoutes(routes) {
continue
}
if _, ok := keep[id]; ok {
continue
}
others = append(others, id)
}
return others
}

View File

@@ -0,0 +1,129 @@
package routemanager
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/client/internal/routeselector"
"github.com/netbirdio/netbird/route"
)
func v6ExitRoute(netID, peer string) *route.Route {
return &route.Route{
NetID: route.NetID(netID),
Network: netip.MustParsePrefix("::/0"),
Peer: peer,
}
}
func newSelectionTestManager() *DefaultManager {
return &DefaultManager{
routeSelector: routeselector.NewRouteSelector(),
clientRoutes: route.HAMap{
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", true)},
"exitA-v6|::/0": {v6ExitRoute("exitA-v6", "p1")},
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", true)},
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
},
}
}
func TestSelectRoutes_ExitNodeExclusivity(t *testing.T) {
m := newSelectionTestManager()
// Selecting an exit node selects its v6 pair and deselects the sibling.
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
assert.True(t, m.routeSelector.IsSelected("exitA"), "exitA should be selected")
assert.True(t, m.routeSelector.IsSelected("exitA-v6"), "the v6 pair follows its v4 base")
assert.False(t, m.routeSelector.IsSelected("exitB"), "the sibling exit node must be deselected")
// Switching to the sibling deselects the previous exit node and its v6 pair.
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
assert.True(t, m.routeSelector.IsSelected("exitB"), "exitB should now be selected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "the previous exit node must be deselected")
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "the previous exit node's v6 pair must be deselected")
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
// Selecting a non-exit route leaves the active exit node alone.
require.NoError(t, m.selectRoutes([]route.NetID{"lan"}, true))
assert.True(t, m.routeSelector.IsSelected("exitB"), "selecting a non-exit route keeps the exit node")
// Deselecting the active exit node turns every exit node off.
require.NoError(t, m.deselectRoutes([]route.NetID{"exitB"}))
assert.False(t, m.routeSelector.IsSelected("exitB"), "exitB should be deselected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "exitA stays deselected")
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
}
func TestSelectRoutes_PartialErrorStillEnforcesExclusivity(t *testing.T) {
// The unknown ID must be reported, but the valid exit node in the same
// request is still selected — so its sibling must still be deselected.
// Both orderings are covered: processing must continue past the invalid
// ID wherever it sits in the request.
requests := map[string][]route.NetID{
"invalid id first": {"missing", "exitB"},
"invalid id last": {"exitB", "missing"},
}
for name, ids := range requests {
t.Run(name, func(t *testing.T) {
m := newSelectionTestManager()
require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
err := m.selectRoutes(ids, true)
assert.Error(t, err, "unknown id must be reported")
assert.True(t, m.routeSelector.IsSelected("exitB"), "valid exit node from the request is selected")
assert.False(t, m.routeSelector.IsSelected("exitA"), "sibling exit node must be deselected despite the error")
assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "sibling's v6 pair must be deselected too")
})
}
}
func TestSelectAllRoutes_KeepsSingleExitNode(t *testing.T) {
// Both exit nodes are marked for auto-apply by management
// (SkipAutoApply=false), the state where select-all could turn on two at
// once without the immediate reconciliation.
m := &DefaultManager{
routeSelector: routeselector.NewRouteSelector(),
clientRoutes: route.HAMap{
"exitA|0.0.0.0/0": {exitRoute("exitA", "p1", false)},
"exitB|0.0.0.0/0": {exitRoute("exitB", "p2", false)},
"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
},
}
require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
m.selectAllRoutes()
assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit routes are all selected")
assert.True(t, m.routeSelector.IsSelected("exitA"), "the deterministic management pick stays active")
assert.False(t, m.routeSelector.IsSelected("exitB"), "select-all must not leave a second exit node active")
}
func TestSelectRoutes_UnknownRoute(t *testing.T) {
m := newSelectionTestManager()
assert.Error(t, m.selectRoutes([]route.NetID{"missing"}, true), "selecting an unavailable route must fail")
assert.Error(t, m.deselectRoutes([]route.NetID{"missing"}), "deselecting an unavailable route must fail")
}
func TestExitNodeSelectionHelpers(t *testing.T) {
routesMap := map[route.NetID][]*route.Route{
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
}
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
}

View File

@@ -13,7 +13,6 @@ import (
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/exp/maps"
"github.com/netbirdio/netbird/client/internal"
"github.com/netbirdio/netbird/client/internal/auth"
@@ -637,23 +636,18 @@ func (c *Client) SelectRoute(id string) error {
}
routeManager := engine.GetRouteManager()
routeSelector := routeManager.GetRouteSelector()
if id == "All" {
log.Debugf("select all routes")
routeSelector.SelectAllRoutes()
} else {
log.Debugf("select route with id: %s", id)
routes := toNetIDs([]string{id})
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
if err := routeSelector.SelectRoutes(routes, true, maps.Keys(routesMap)); err != nil {
log.Debugf("error when selecting routes: %s", err)
return fmt.Errorf("select routes: %w", err)
}
routeManager.SelectAllRoutes()
return nil
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
return nil
log.Debugf("select route with id: %s", id)
if err := routeManager.SelectRoutes(toNetIDs([]string{id}), true); err != nil {
log.Debugf("error when selecting routes: %s", err)
return err
}
return nil
}
func (c *Client) DeselectRoute(id string) error {
@@ -667,21 +661,17 @@ func (c *Client) DeselectRoute(id string) error {
}
routeManager := engine.GetRouteManager()
routeSelector := routeManager.GetRouteSelector()
if id == "All" {
log.Debugf("deselect all routes")
routeSelector.DeselectAllRoutes()
} else {
log.Debugf("deselect route with id: %s", id)
routes := toNetIDs([]string{id})
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
if err := routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
log.Debugf("error when deselecting routes: %s", err)
return fmt.Errorf("deselect routes: %w", err)
}
routeManager.DeselectAllRoutes()
return nil
}
log.Debugf("deselect route with id: %s", id)
if err := routeManager.DeselectRoutes(toNetIDs([]string{id})); err != nil {
log.Debugf("error when deselecting routes: %s", err)
return err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
return nil
}

View File

@@ -0,0 +1,12 @@
//go:build ios
package NetBirdSDK
import "github.com/netbirdio/netbird/version"
// GoClientVersion returns the NetBird Go client version that was baked into
// the framework at compile time via
// -ldflags "-X github.com/netbirdio/netbird/version.version=<version>".
func GoClientVersion() string {
return version.NetbirdVersion()
}

View File

@@ -8,7 +8,6 @@ import (
"sort"
"strings"
"golang.org/x/exp/maps"
"google.golang.org/grpc/codes"
gstatus "google.golang.org/grpc/status"
@@ -161,30 +160,11 @@ func (s *Server) SelectNetworks(_ context.Context, req *proto.SelectNetworksRequ
return nil, fmt.Errorf("no route manager")
}
routeSelector := routeManager.GetRouteSelector()
if req.GetAll() {
routeSelector.SelectAllRoutes()
} else {
routes := toNetIDs(req.GetNetworkIDs())
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
netIdRoutes := maps.Keys(routesMap)
if err := routeSelector.SelectRoutes(routes, req.GetAppend(), netIdRoutes); err != nil {
return nil, fmt.Errorf("select routes: %w", err)
}
// Exit nodes are mutually exclusive: if this selection activates an
// exit node, deselect every other available exit node so two can't be
// selected at once. Non-exit route selections are left untouched.
if requestActivatesExitNode(routes, routesMap) {
if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
if err := routeSelector.DeselectRoutes(others, netIdRoutes); err != nil {
return nil, fmt.Errorf("deselect sibling exit nodes: %w", err)
}
}
}
routeManager.SelectAllRoutes()
} else if err := routeManager.SelectRoutes(toNetIDs(req.GetNetworkIDs()), req.GetAppend()); err != nil {
return nil, err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
s.statusRecorder.PublishEvent(
proto.SystemEvent_INFO,
@@ -224,19 +204,11 @@ func (s *Server) DeselectNetworks(_ context.Context, req *proto.SelectNetworksRe
return nil, fmt.Errorf("no route manager")
}
routeSelector := routeManager.GetRouteSelector()
if req.GetAll() {
routeSelector.DeselectAllRoutes()
} else {
routes := toNetIDs(req.GetNetworkIDs())
routesMap := routeManager.GetClientRoutesWithNetID()
routes = route.ExpandV6ExitPairs(routes, routesMap)
netIdRoutes := maps.Keys(routesMap)
if err := routeSelector.DeselectRoutes(routes, netIdRoutes); err != nil {
return nil, fmt.Errorf("deselect routes: %w", err)
}
routeManager.DeselectAllRoutes()
} else if err := routeManager.DeselectRoutes(toNetIDs(req.GetNetworkIDs())); err != nil {
return nil, err
}
routeManager.TriggerSelection(routeManager.GetClientRoutes())
s.statusRecorder.PublishEvent(
proto.SystemEvent_INFO,
@@ -261,37 +233,3 @@ func toNetIDs(routes []string) []route.NetID {
return netIDs
}
func isExitNodeRoutes(routes []*route.Route) bool {
return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
}
// requestActivatesExitNode reports whether any requested NetID maps to an exit
// node (default route) in the current route table.
func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
for _, id := range requested {
if isExitNodeRoutes(routesMap[id]) {
return true
}
}
return false
}
// otherExitNodeIDs returns every available exit-node NetID that is not in the
// requested set — the siblings to deselect so a single exit node stays active.
func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
keep := make(map[route.NetID]struct{}, len(requested))
for _, id := range requested {
keep[id] = struct{}{}
}
var others []route.NetID
for id, routes := range routesMap {
if !isExitNodeRoutes(routes) {
continue
}
if _, ok := keep[id]; ok {
continue
}
others = append(others, id)
}
return others
}

View File

@@ -1,26 +0,0 @@
package server
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/route"
)
func TestExitNodeSelectionHelpers(t *testing.T) {
routesMap := map[route.NetID][]*route.Route{
"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
"exitB": {{Network: netip.MustParsePrefix("::/0")}},
"lan": {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
}
assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
}

View File

@@ -2,6 +2,7 @@ import { type ReactNode } from "react";
import { useTranslation } from "react-i18next";
import { Browser } from "@wailsio/runtime";
import { DownloadIcon, NotepadText } from "lucide-react";
import { Update as UpdateSvc } from "@bindings/services";
import { Button } from "@/components/buttons/Button";
import { useClientVersion } from "@/contexts/ClientVersionContext";
import { cn } from "@/lib/cn";
@@ -14,6 +15,12 @@ function openUrl(url: string) {
});
}
function openInstallerDownload() {
UpdateSvc.DownloadURL()
.then(openUrl)
.catch(() => openUrl(GITHUB_RELEASES));
}
export function UpdateVersionCard() {
const { t } = useTranslation();
const { updateVersion, enforced, triggerUpdate } = useClientVersion();
@@ -37,11 +44,7 @@ export function UpdateVersionCard() {
{t("update.card.installNow")}
</Button>
) : (
<Button
variant={"primary"}
size={"xs"}
onClick={() => openUrl(GITHUB_RELEASES)}
>
<Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
<DownloadIcon size={14} />
{t("update.card.getInstaller")}
</Button>

View File

@@ -10,6 +10,7 @@ import (
"github.com/netbirdio/netbird/client/proto"
"github.com/netbirdio/netbird/client/ui/updater"
"github.com/netbirdio/netbird/version"
)
// UpdateResult mirrors TriggerUpdateResponse.
@@ -33,6 +34,12 @@ func (s *Update) GetState() updater.State {
return s.holder.Get()
}
// DownloadURL returns the platform-appropriate installer download link for
// manual (non-enforced) updates.
func (s *Update) DownloadURL() string {
return version.DownloadUrl()
}
// Quit exits the app. Scheduled off the calling goroutine so the JS caller's
// response returns before the runtime tears down.
func (s *Update) Quit() {

View File

@@ -32,9 +32,8 @@ const (
quitDownTimeout = 5 * time.Second
urlGitHubRepo = "https://github.com/netbirdio/netbird"
urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
urlDocs = "https://docs.netbird.io"
urlGitHubRepo = "https://github.com/netbirdio/netbird"
urlDocs = "https://docs.netbird.io"
)
// TrayServices bundles the services the tray menu needs, grouped so NewTray

View File

@@ -13,6 +13,7 @@ import (
"github.com/netbirdio/netbird/client/ui/services"
"github.com/netbirdio/netbird/client/ui/updater"
"github.com/netbirdio/netbird/version"
)
// trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
@@ -76,15 +77,15 @@ func (u *trayUpdater) applyLanguage() {
u.refreshMenuItem(state)
}
// handleClick opens the GitHub releases page when not Enforced, otherwise shows
// the progress page and asks the daemon to start the installer.
// handleClick opens the installer download link when not Enforced, otherwise
// shows the progress page and asks the daemon to start the installer.
func (u *trayUpdater) handleClick() {
u.mu.Lock()
state := u.state
u.mu.Unlock()
if !state.Enforced {
_ = u.app.Browser.OpenURL(urlGitHubReleases)
_ = u.app.Browser.OpenURL(version.DownloadUrl())
return
}

View File

@@ -109,7 +109,7 @@ sequenceDiagram
Chk->>Inj: continue
Inj->>Inj: inject NetBird identity headers per provider config
Inj->>Grd: continue
Grd->>Grd: enforce model allowlist
Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
Grd->>Up: forward (over WireGuard)
Up-->>Resp: response (JSON or SSE stream)
Resp->>Resp: parse usage tokens, completion
@@ -135,6 +135,21 @@ sequenceDiagram
(`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
PII names — see `redact.go` for the full set. See
[`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
is authoritative: it resolves the policy that governs this
(provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
when no applicable policy permits the model — so an allowlist scoped to
one group/provider never leaks to another, and an un-guardrailed policy
is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
backstop: it only carries an allowlist for a provider every authorising
policy restricts, and blocks unknown/undetermined models even when
management is unreachable. Because that backstop allowlist is the UNION
of every restricting policy's models, per-group narrowing lives only in
the authoritative check: during a `CheckLLMPolicyLimits` outage
`llm_limit_check` fails open, so a caller can reach any model in the
provider's union — a group scoped to model A could reach model B if
another group restricts the same provider to B. This is the documented
fail-open trade-off; a future flag may switch it to fail-closed.
- SSE streaming requires special handling on the response side; the
parser must handle partial chunks without buffering the whole
stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).

View File

@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
| on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
| on_request | 2 | `llm_limit_check` | `{}` | |
| on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
| on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | |
| on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | |
| on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | |
| on_response | 6 | `cost_meter` | `{}` | |
| on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | |

View File

@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
| `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
| `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
| `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
| `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
| `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
| `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |

View File

@@ -9,25 +9,234 @@ import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// per1k is a model's published USD rates per 1k tokens. read is the prompt-cache read rate
// (OpenAI: the cached-input discount rate); write is the cache-creation rate where one exists.
type per1k struct{ in, out, read, write float64 }
// publishedPer1k hardcodes the vendors' PUBLISHED rates for the models the live matrix can drive,
// keyed by the normalized model id the proxy stamps. Deliberately independent of the proxy's
// pricing table so a wrong embedded rate or a broken normalization fails the run.
var publishedPer1k = map[string]per1k{
"gpt-4o-mini": {0.00015, 0.0006, 0.000075, 0},
"gpt-4o": {0.0025, 0.01, 0.00125, 0},
"claude-haiku-4-5": {0.001, 0.005, 0.0001, 0.00125},
"claude-sonnet-4-5": {0.003, 0.015, 0.0003, 0.00375},
"claude-sonnet-4-6": {0.003, 0.015, 0.0003, 0.00375},
"kimi-k3": {0.003, 0.015, 0.0003, 0.003}, // no published write rate: bills at the input rate
"anthropic.claude-haiku-4-5": {0.001, 0.005, 0.0001, 0.00125},
"anthropic.claude-sonnet-4-5": {0.003, 0.015, 0.0003, 0.00375},
"anthropic.claude-sonnet-4-6": {0.003, 0.015, 0.0003, 0.00375},
}
// rawCostVerificationSQL is the operator-facing double-check, run straight against the management
// sqlite store: recompute each usage row's expected total and cache cost from its own persisted
// token buckets and hardcoded published rates. OpenAI counts cached tokens as a subset of input;
// Anthropic-shape providers count cache buckets additively.
const rawCostVerificationSQL = `
WITH rates(model, in_rate, out_rate, read_rate, write_rate) AS (
VALUES
('gpt-4o-mini', 0.00015, 0.0006, 0.000075, 0.0),
('gpt-4o', 0.0025, 0.01, 0.00125, 0.0),
('claude-haiku-4-5', 0.001, 0.005, 0.0001, 0.00125),
('claude-sonnet-4-5', 0.003, 0.015, 0.0003, 0.00375),
('claude-sonnet-4-6', 0.003, 0.015, 0.0003, 0.00375),
('kimi-k3', 0.003, 0.015, 0.0003, 0.003),
('anthropic.claude-haiku-4-5', 0.001, 0.005, 0.0001, 0.00125),
('anthropic.claude-sonnet-4-5', 0.003, 0.015, 0.0003, 0.00375),
('anthropic.claude-sonnet-4-6', 0.003, 0.015, 0.0003, 0.00375)
)
SELECT
u.provider,
u.model,
u.input_tokens,
u.output_tokens,
u.cached_input_tokens,
u.cache_creation_tokens,
u.input_cost_usd,
u.cached_input_cost_usd,
u.cache_creation_cost_usd,
u.output_cost_usd,
-- No cost_usd / cache_cost_usd columns are stored: both are derived from the
-- four per-bucket columns above, exactly as the API renders them.
(u.input_cost_usd + u.cached_input_cost_usd + u.cache_creation_cost_usd + u.output_cost_usd) AS cost_usd,
(u.cached_input_cost_usd + u.cache_creation_cost_usd) AS cache_cost_usd,
CASE WHEN u.provider = 'openai' THEN
(u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
ELSE
u.input_tokens*r.in_rate/1000.0
END AS expected_input,
CASE WHEN u.provider = 'openai' THEN
MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
ELSE
u.cached_input_tokens*r.read_rate/1000.0
END AS expected_cached_input,
CASE WHEN u.provider = 'openai' THEN
0.0
ELSE
u.cache_creation_tokens*r.write_rate/1000.0
END AS expected_cache_creation,
u.output_tokens*r.out_rate/1000.0 AS expected_output,
CASE WHEN u.provider = 'openai' THEN
(u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
+ MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
+ u.output_tokens*r.out_rate/1000.0
ELSE
u.input_tokens*r.in_rate/1000.0 + u.cached_input_tokens*r.read_rate/1000.0
+ u.cache_creation_tokens*r.write_rate/1000.0 + u.output_tokens*r.out_rate/1000.0
END AS expected_total,
CASE WHEN u.provider = 'openai' THEN
MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
ELSE
u.cached_input_tokens*r.read_rate/1000.0 + u.cache_creation_tokens*r.write_rate/1000.0
END AS expected_cache
FROM agent_network_request_usage u
JOIN rates r ON r.model = u.model
ORDER BY u.timestamp`
// verifyUsageRowsSQL re-checks every persisted usage row directly in the management sqlite store,
// bypassing the API path — the same audit an operator can run on a production store.db.
func verifyUsageRowsSQL(t *testing.T, srv *harness.Combined) {
t.Helper()
dbPath, err := srv.SnapshotStoreDB(t.TempDir())
require.NoError(t, err, "snapshot management sqlite store")
db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
require.NoError(t, err, "open store snapshot")
sqlDB, err := db.DB()
require.NoError(t, err)
defer func() { _ = sqlDB.Close() }()
rows, err := db.Raw(rawCostVerificationSQL).Rows()
require.NoError(t, err, "run raw cost verification query")
defer func() { _ = rows.Close() }()
verified := 0
for rows.Next() {
var provider, model string
var inTok, outTok, readTok, writeTok int64
var inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost float64
var wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache float64
require.NoError(t, rows.Scan(&provider, &model, &inTok, &outTok, &readTok, &writeTok,
&inCost, &cachedInCost, &cacheCreateCost, &outCost, &cost, &cacheCost,
&wantInput, &wantCachedInput, &wantCacheCreation, &wantOutput, &wantTotal, &wantCache), "scan usage row")
t.Logf("[sql] %s/%s: in=%d out=%d cache_read=%d cache_write=%d stored in/cached/create/out=$%.6f/$%.6f/$%.6f/$%.6f total=$%.6f cache=$%.6f expected total=$%.6f cache=$%.6f",
provider, model, inTok, outTok, readTok, writeTok,
inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost, wantTotal, wantCache)
assert.InDeltaf(t, wantInput, inCost, 1e-6, "stored input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCachedInput, cachedInCost, 1e-6, "stored cached_input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCacheCreation, cacheCreateCost, 1e-6, "stored cache_creation_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantOutput, outCost, 1e-6, "stored output_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantTotal, cost, 1e-6, "derived cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, wantCache, cacheCost, 1e-6, "derived cache_cost_usd for %s/%s must match the published-rate recompute", provider, model)
assert.InDeltaf(t, inCost+cachedInCost+cacheCreateCost+outCost, cost, 1e-9,
"stored buckets must sum to the derived cost_usd for %s/%s", provider, model)
verified++
}
require.NoError(t, rows.Err(), "iterate usage rows")
require.Positive(t, verified, "raw SQL check must cover at least one usage row")
t.Logf("[sql] verified %d usage rows in store.db against published rates", verified)
gwRows, err := db.Raw(`SELECT model,
(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd) AS cost_usd
FROM agent_network_request_usage WHERE model LIKE '%/%'`).Rows()
require.NoError(t, err, "query gateway-prefixed usage rows")
defer func() { _ = gwRows.Close() }()
for gwRows.Next() {
var model string
var cost float64
require.NoError(t, gwRows.Scan(&model, &cost), "scan gateway usage row")
t.Logf("[sql] gateway %s: stored=$%.6f (must be 0 — deliberately unpriced)", model, cost)
assert.Zerof(t, cost, "gateway-prefixed model %q must store cost 0, never a guessed rate", model)
}
require.NoError(t, gwRows.Err(), "iterate gateway usage rows")
}
// validateAccessLogCost recomputes a live access-log row's expected total and cache cost from the
// published per-1k rates and the row's persisted token buckets, and asserts both stored values.
// Gateway-prefixed model ids the proxy deliberately does not price must store cost 0.
func validateAccessLogCost(t *testing.T, pc providerCase, row api.AgentNetworkAccessLog) {
t.Helper()
model := catalogModel(pc)
provider := ""
if row.Provider != nil {
provider = *row.Provider
}
t.Logf("[cost] %s: provider=%s model=%s in=%d out=%d total=%d cache_read=%d cache_write=%d cost=$%.6f cache_cost=$%.6f",
pc.name, provider, model, row.InputTokens, row.OutputTokens, row.TotalTokens,
row.CachedInputTokens, row.CacheCreationTokens, row.CostUsd, row.CacheCostUsd)
rates, known := publishedPer1k[model]
if !known {
if strings.Contains(model, "/") {
assert.Zerof(t, row.CostUsd, "gateway-prefixed model %q is not priced so the cost meter must skip (cost 0)", model)
return
}
t.Logf("[cost] %s: no published rate on file for model %q (env-overridden?); skipping cost validation", pc.name, model)
return
}
// input_tokens may legitimately be 0: Moonshot/Kimi reports fully cached prompts under the cache
// buckets only. Output and total must always be present on a priced row.
require.Positive(t, row.OutputTokens, "priced row must carry output tokens")
require.Positive(t, row.TotalTokens, "priced row must carry total tokens")
var wantInput, wantCachedInput, wantCacheCreation float64
if provider == "openai" {
cached := min(row.CachedInputTokens, row.InputTokens) // cached is a subset of input
wantInput = float64(row.InputTokens-cached) / 1000 * rates.in
wantCachedInput = float64(cached) / 1000 * rates.read
// OpenAI has no cache-write bucket; wantCacheCreation stays 0.
} else {
// Anthropic / Bedrock shape: cache buckets are additive to input_tokens.
wantInput = float64(row.InputTokens) / 1000 * rates.in
wantCachedInput = float64(row.CachedInputTokens) / 1000 * rates.read
wantCacheCreation = float64(row.CacheCreationTokens) / 1000 * rates.write
}
wantOutput := float64(row.OutputTokens) / 1000 * rates.out
wantCache := wantCachedInput + wantCacheCreation
wantTotal := wantInput + wantCache + wantOutput
t.Logf("[cost] %s: expecting input=$%.6f cached_input=$%.6f cache_creation=$%.6f output=$%.6f total=$%.6f cache=$%.6f from published rates",
pc.name, wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache)
assert.InDeltaf(t, wantInput, row.InputCostUsd, 1e-6, "stored input_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCachedInput, row.CachedInputCostUsd, 1e-6, "stored cached_input_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCacheCreation, row.CacheCreationCostUsd, 1e-6, "stored cache_creation_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantOutput, row.OutputCostUsd, 1e-6, "stored output_cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantTotal, row.CostUsd, 1e-6, "derived cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, wantCache, row.CacheCostUsd, 1e-6, "derived cache_cost_usd for %s (%s)", pc.name, model)
// The aggregates must be exactly the sum of the stored components, not an
// independently-computed figure that could drift from the breakdown.
assert.InDeltaf(t, row.InputCostUsd+row.CachedInputCostUsd+row.CacheCreationCostUsd+row.OutputCostUsd,
row.CostUsd, 1e-9, "stored buckets must sum to the derived cost_usd for %s (%s)", pc.name, model)
assert.InDeltaf(t, row.CachedInputCostUsd+row.CacheCreationCostUsd,
row.CacheCostUsd, 1e-9, "stored cache buckets must sum to the derived cache_cost_usd for %s (%s)", pc.name, model)
}
// providerCase is one entry in the live provider matrix. The same scenario runs
// for every available provider; availability is keyed off env vars so the suite
// covers whatever credentials are present (source ~/.llm-keys locally / set the
// Actions secrets in CI).
type providerCase struct {
name string
catalogID string
upstream string
apiKey string
model string // body model (chat/messages) or path model@version (vertex)
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
project string // vertex only: GCP project for the rawPredict path
region string // vertex only: GCP region for the rawPredict path
name string
catalogID string
upstream string
apiKey string
model string // body model (chat/messages) or path model@version (vertex)
kind string // harness.WireChat, harness.WireMessages, or harness.WireVertex
project string // vertex only: GCP project for the rawPredict path
region string // vertex only: GCP region for the rawPredict path
pathPrefix string // base-URL path prefix the agent carries (e.g. "/anthropic" for Kimi)
}
// availableProviders builds the matrix from the provider env vars that are set.
@@ -39,6 +248,24 @@ func availableProviders() []providerCase {
if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" {
ps = append(ps, providerCase{name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k, model: "claude-haiku-4-5", kind: harness.WireMessages})
}
if k := os.Getenv("KIMI_TOKEN"); k != "" {
// Kimi (Moonshot AI) serves two body shapes from the same key: OpenAI
// Chat Completions on the bare host (/v1/...) and the Anthropic
// Messages API under the /anthropic path prefix (the endpoint
// Moonshot's Claude Code guide uses). The provider keeps the bare
// default upstream and the AGENT carries the /anthropic prefix in
// its base URL — exactly the documented Claude Code / Kimi CLI
// setup (ANTHROPIC_BASE_URL=https://<endpoint>/anthropic) — so one
// provider serves both shapes and the prefix rides through to
// Moonshot. Run the Anthropic shape, the flagship Claude Code path;
// the OpenAI wire shape is covered live by the other chat-shaped
// matrix providers, and Kimi-over-chat passed with kimi-k3 before
// the single-model constraint surfaced (run #73 on the kimi feature
// branch). The platform serves this account exactly ONE model —
// kimi-k3 (kimi-k2-thinking and even kimi-latest return
// resource_not_found_error on both surfaces).
ps = append(ps, providerCase{name: "kimi", catalogID: "kimi_api", upstream: "https://api.moonshot.ai", apiKey: k, model: "kimi-k3", kind: harness.WireMessages, pathPrefix: "/anthropic"})
}
if k, u := os.Getenv("VERCEL_TOKEN"), os.Getenv("VERCEL_URL"); k != "" && u != "" {
ps = append(ps, providerCase{name: "vercel", catalogID: "vercel_ai_gateway", upstream: u, apiKey: k, model: "openai/gpt-4o-mini", kind: harness.WireChat})
}
@@ -84,19 +311,25 @@ func availableProviders() []providerCase {
}
}
// Bedrock: path-routed, bearer auth. Model is a cross-region inference
// profile id (distinct string from the first-party Anthropic case).
// Bedrock: path-routed, bearer auth. Model is the FULL cross-region
// inference-profile id exactly as AWS issues it — region-family prefix
// plus the date/version suffix. A bare or wrong-region id makes Bedrock
// reject the request with "The provided model identifier is invalid"
// before any inference runs. The proxy normalizes this id to the catalog
// key (anthropic.claude-haiku-4-5) for routing/pricing/allowlists.
// Defaults pair eu-central-1 with the eu.* profile; AWS_REGION overrides
// the region and the prefix follows its family.
if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
region := os.Getenv("AWS_REGION")
if region == "" {
region = "us-east-1"
region = "eu-central-1"
}
// A valid Bedrock inference-profile id (region prefix + date + version),
// overridable per account. `global.` profiles can be invoked from any
// region; set AWS_BEDROCK_MODEL to match the enabled profile for the token.
// A valid Bedrock inference-profile id, overridable per account (AWS_BEDROCK_MODEL, also the
// workflow's bedrock_model dispatch input). `global.` profiles work from any region. Defaults to
// Sonnet 4.6, whose id convention dropped the -YYYYMMDD-v1:0 suffix that Haiku 4.5 still carries.
model := os.Getenv("AWS_BEDROCK_MODEL")
if model == "" {
model = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
model = "global.anthropic.claude-sonnet-4-6"
}
ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: model, kind: harness.WireBedrock})
}
@@ -216,8 +449,7 @@ func TestProvidersMatrix(t *testing.T) {
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
@@ -233,6 +465,10 @@ func TestProvidersMatrix(t *testing.T) {
// session id and confirm the marker propagated end-to-end.
sessionID := "e2e-session-" + pc.name
// A long-form prompt so completions carry realistic token counts for cost validation;
// max_tokens in the harness bodies (2048) lets the full answer through.
const matrixPrompt = "explain GitHub workflow in 1000 words"
// Retry briefly to absorb tunnel/DNS jitter on the first call.
var code int
var body string
@@ -243,11 +479,11 @@ func TestProvidersMatrix(t *testing.T) {
var cerr error
switch pc.kind {
case harness.WireVertex:
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, matrixPrompt, sessionID)
case harness.WireBedrock:
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, matrixPrompt, sessionID)
default:
c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
c, b, cerr = cl.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, pc.kind, pc.model, matrixPrompt, sessionID)
}
if cerr == nil {
code, body = c, b
@@ -266,6 +502,7 @@ func TestProvidersMatrix(t *testing.T) {
// The session id sent as x-session-id must round-trip into the
// access-log row for this provider.
var row api.AgentNetworkAccessLog
require.Eventually(t, func() bool {
logs, lerr := srv.ListAccessLogs(ctx)
if lerr != nil {
@@ -273,11 +510,15 @@ func TestProvidersMatrix(t *testing.T) {
}
for _, r := range logs.Data {
if r.SessionId != nil && *r.SessionId == sessionID {
row = r
return true
}
}
return false
}, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row for %s", sessionID, pc.name)
// Stored total and cache cost must match the published rates applied to the row's buckets.
validateAccessLogCost(t, pc, row)
})
}
@@ -298,4 +539,7 @@ func TestProvidersMatrix(t *testing.T) {
}
return false
}, 60*time.Second, 3*time.Second, "consumption must be recorded with positive token counts after live traffic")
// Final raw-SQL audit: bypass the API and re-verify every persisted usage row in the store.
verifyUsageRowsSQL(t, srv)
}

View File

@@ -0,0 +1,209 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
// provider, its own guardrail whose allowlist holds ONLY that provider's
// allowed model, and its own policy. Each case runs in isolation (its own
// proxy + client), so the guardrail the proxy enforces contains exactly this
// provider's model — never a mixed cross-provider list.
type pathRoutedGuardrailCase struct {
name string
catalogID string // agent-network catalog provider id
wire string // harness.WireVertex | harness.WireBedrock
allowEntry string // the single model id put on the guardrail allowlist
allowModel string // model id sent that MUST be served (200)
blockModel string // model id sent that MUST be denied (403 model_blocked)
}
// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
// guard for the customer report that a model-allowlist guardrail attached to a
// policy has no effect for PATH-ROUTED providers — where the model travels in
// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
// AWS Bedrock (/model/{id}/invoke).
//
// Each provider is tested in isolation with a guardrail allowlisting a single
// model of its own: the allowed model (in the URL path) is served (200) and an
// unselected model (in the URL path) is denied 403 by the guardrail
// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
// customer verbatim — allow Sonnet, and the unselected model is the exact
// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
// sends a region-prefixed, versioned inference-profile id so URL-path model
// normalization is exercised too.
//
// The provider is catch-all (no models), so the router forwards any model and a
// 403 can only come from the guardrail, never model_not_routable. Only the
// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
// synth/reconcile, the proxy middleware chain (URL-path model extraction,
// router, guardrail) and the tunnel are all real, and the guardrail denies
// before the upstream is dialed so the mock cannot influence the block. A
// static bearer api key is used so the router injects a static Authorization
// header instead of minting a GCP token — the only reason path-routed providers
// normally need live credentials — so the test runs with none and is always on.
func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
cases := []pathRoutedGuardrailCase{
{
name: "vertex",
catalogID: "vertex_ai_api",
wire: harness.WireVertex,
allowEntry: "claude-sonnet-4-5",
allowModel: "claude-sonnet-4-5",
blockModel: "claude-opus-4-6", // the customer-reported model
},
{
name: "bedrock",
catalogID: "bedrock_api",
wire: harness.WireBedrock,
allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
blockModel: "us.anthropic.claude-opus-4-8-v1:0",
},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
runPathRoutedGuardrailCase(t, tc)
})
}
}
func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
t.Helper()
const (
vertexProject = "e2e-project"
vertexRegion = "global"
)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
require.NoError(t, err, "create group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-guardrail-" + tc.name + "-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grp.Id},
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
// Catch-all provider (no models) so the router forwards any model; a static
// bearer key means the router injects a static auth header instead of minting
// a GCP token. Bootstraps the cluster if it isn't already.
staticKey := "static-e2e-token"
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: tc.name,
ProviderId: tc.catalogID,
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create %s provider", tc.name)
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
// Guardrail allowlisting ONLY this provider's allowed model.
var gr api.AgentNetworkGuardrailRequest
gr.Name = "e2e-guardrail-" + tc.name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
guard, err := srv.CreateGuardrail(ctx, gr)
require.NoError(t, err, "create guardrail")
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
enabled := true
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-" + tc.name,
Enabled: &enabled,
SourceGroups: []string{grp.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{guard.Id},
})
require.NoError(t, err, "create policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Probe first: the GET resolves the endpoint and its first packet wakes the
// lazy proxy peer, so WaitProxyPeer then observes it connected.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
var code int
var body string
var cerr error
switch tc.wire {
case harness.WireVertex:
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
case harness.WireBedrock:
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
default:
t.Fatalf("unsupported wire %q", tc.wire)
}
require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
return code, body
}
// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
// jitter on the first call over the freshly warmed tunnel.
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(tc.allowModel)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
assert.Equal(t, 200, code,
"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
// Unselected model (in the URL path) must be blocked by the guardrail.
code, body = send(tc.blockModel)
assert.Equal(t, 403, code,
"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
}

View File

@@ -0,0 +1,205 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
// the proxy caches; the switch is invisible until that cache expires. The proxy
// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
// instead of the 5-minute default.
//
// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
// served and modelB denied; after switching the client grpA -> grpB, modelB is
// served and modelA denied. The cross-group deny comes from management's
// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
modelA = "e2e-model-a"
modelB = "e2e-model-b"
// Short tunnel-cache TTL so a group switch propagates in seconds.
// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
cacheTTL = 3 * time.Second
)
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
require.NoError(t, err, "create group A")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
require.NoError(t, err, "create group B")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-gswitch-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpA.Id}, // client starts in group A
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "gswitch",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: &[]api.AgentNetworkProviderModel{
{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
},
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
mkGuard := func(name, model string) api.AgentNetworkGuardrail {
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, gerr := srv.CreateGuardrail(ctx, gr)
require.NoError(t, gerr, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}
gA := mkGuard("e2e-gswitch-a", modelA)
gB := mkGuard("e2e-gswitch-b", modelB)
enabled := true
polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-gswitch-a",
Enabled: &enabled,
SourceGroups: []string{grpA.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gA.Id},
})
require.NoError(t, err, "create policy A")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-gswitch-b",
Enabled: &enabled,
SourceGroups: []string{grpB.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gB.Id},
})
require.NoError(t, err, "create policy B")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
})
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
var code int
var body string
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
code, body = send(model)
if code == want {
return code, body
}
time.Sleep(2 * time.Second)
}
return code, body
}
// Phase 1 — client is in group A: modelA served, modelB denied.
code, body := sendUntil(modelA, 200, 90*time.Second)
assert.Equal(t, 200, code,
"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = send(modelB)
assert.Equal(t, 403, code,
"group-B model must be denied while the client is in group A; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
// Switch the client peer from group A to group B.
peerID := clientPeerInGroup(t, ctx, grpA.Id)
_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
Name: grpB.Name,
Peers: &[]string{peerID},
})
require.NoError(t, err, "add peer to group B")
_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
Name: grpA.Name,
Peers: &[]string{},
})
require.NoError(t, err, "remove peer from group A")
// Phase 2 — after the short TTL expires the proxy re-validates the peer,
// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
code, body = sendUntil(modelB, 200, 60*time.Second)
assert.Equal(t, 200, code,
"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = send(modelA)
assert.Equal(t, 403, code,
"after the switch, the old group-A model must be denied; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
}
// clientPeerInGroup returns the id of the single peer that is a member of the
// given group — the test client. The proxy peer is never added to test groups.
func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
t.Helper()
peers, err := srv.API().Peers.List(ctx)
require.NoError(t, err, "list peers")
for _, p := range peers {
for _, g := range p.Groups {
if g.Id == groupID {
return p.Id
}
}
}
t.Fatalf("no peer found in group %s", groupID)
return ""
}

View File

@@ -0,0 +1,201 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
// modelOther denied for the grpMain client (403 model_blocked, no cross-group
// leak), and openModel on the un-guardrailed policy's provider served (200).
func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
modelSelected = "e2e-selected"
modelOther = "e2e-other"
openModel = "e2e-open"
)
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
require.NoError(t, err, "create main group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
require.NoError(t, err, "create other group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-guardrail-mp-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpMain.Id}, // client joins grpMain only
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
models := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
// pRestricted declares the two guardrailed models so routing is deterministic
// (model -> provider). Created first, so it carries the bootstrap cluster.
pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "restricted",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: models(modelSelected, modelOther),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create restricted provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "open",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: models(openModel),
})
require.NoError(t, err, "create open provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, gerr := srv.CreateGuardrail(ctx, gr)
require.NoError(t, gerr, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}
gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
enabled := true
// polMain: grpMain restricted to modelSelected on pRestricted.
polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-main",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{pRestricted.Id},
GuardrailIds: &[]string{gMain.Id},
})
require.NoError(t, err, "create main policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
// polOther: grpOther restricted to modelOther on the SAME provider. The
// client is not in grpOther, so modelOther must never be usable by it.
polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-other",
Enabled: &enabled,
SourceGroups: []string{grpOther.Id},
DestinationProviderIds: []string{pRestricted.Id},
GuardrailIds: &[]string{gOther.Id},
})
require.NoError(t, err, "create other policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-guardrail-mp-open",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{pOpen.Id},
})
require.NoError(t, err, "create open policy")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
sendUntil200 := func(model string) (int, string) {
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(model)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
return code, body
}
t.Run("selected model allowed for its group", func(t *testing.T) {
code, body := sendUntil200(modelSelected)
assert.Equal(t, 200, code,
"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
t.Run("other group's model does not leak", func(t *testing.T) {
// modelOther is allowlisted only for grpOther. The grpMain client must be
// denied by management's per-policy/group check — not waved through by an
// account-wide union. This is the security-critical wrong-ALLOW guard.
code, body := send(modelOther)
assert.Equal(t, 403, code,
"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"denial must be a model-allowlist decision; body: %s", body)
})
t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
// old account-wide union would have blocked openModel (it is on no
// allowlist); it must now be served — the false-DENY guard.
code, body := sendUntil200(openModel)
assert.Equal(t, 200, code,
"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
}

View File

@@ -0,0 +1,422 @@
//go:build e2e
package agentnetwork
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/e2e/harness"
"github.com/netbirdio/netbird/shared/management/http/api"
)
// pergroupCase describes one provider surface for the per-group allowlist matrix.
// selectedReq/otherReq are the model identifiers as they travel in the request
// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
// otherAllow are the (normalized) forms the guardrail allowlist holds — for
// Bedrock these differ from the request form so path normalization is exercised.
type pergroupCase struct {
name string
catalogID string
wire string // "chat", "messages", "vertex", "bedrock"
models *[]api.AgentNetworkProviderModel
selectedReq string
selectedAllow string
otherReq string
otherAllow string
providerID string // filled during setup
}
// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
// allowlist end to end across every always-on provider surface, including the
// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
//
// For each provider two policies target it: grpMain (the client) is allowed only
// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
// The client must get selectedReq served (200) and otherReq denied (403,
// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
// authoritative per-policy/group decision from management (the proxy per-provider
// backstop carries the union of both models), so this also confirms management
// receives the correct normalized model for path-routed providers.
func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
defer cancel()
const (
vertexProject = "e2e-project"
vertexRegion = "global"
)
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
cases := []*pergroupCase{
{
name: "openai", catalogID: "openai_api", wire: harness.WireChat,
models: priced("oai-model-a", "oai-model-b"),
selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
otherReq: "oai-model-b", otherAllow: "oai-model-b",
},
{
name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
models: priced("ant-model-a", "ant-model-b"),
selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
otherReq: "ant-model-b", otherAllow: "ant-model-b",
},
{
// Vertex catalog ids travel bare in the rawPredict path.
name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
},
{
// Bedrock request ids are region-prefixed/versioned; the parser
// normalizes them to the catalog key the allowlist holds.
name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
},
}
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
require.NoError(t, err, "create main group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
require.NoError(t, err, "create other group")
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-pergroup-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{grpMain.Id},
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
enabled := true
for i, c := range cases {
req := api.AgentNetworkProviderRequest{
Name: "e2e-pergroup-" + c.name,
ProviderId: c.catalogID,
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: c.models,
}
if i == 0 {
req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
}
prov, perr := srv.CreateProvider(ctx, req)
require.NoError(t, perr, "create provider %s", c.name)
c.providerID = prov.Id
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-pergroup-" + c.name + "-main",
Enabled: &enabled,
SourceGroups: []string{grpMain.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gSel.Id},
})
require.NoError(t, merr, "create main policy %s", c.name)
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-pergroup-" + c.name + "-other",
Enabled: &enabled,
SourceGroups: []string{grpOther.Id},
DestinationProviderIds: []string{prov.Id},
GuardrailIds: &[]string{gOth.Id},
})
require.NoError(t, oerr, "create other policy %s", c.name)
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
}
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(c *pergroupCase, model string) (int, string) {
var code int
var body string
var cerr error
switch c.wire {
case "vertex":
code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
case "bedrock":
code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
default:
code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
}
require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
return code, body
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
// the first call over the freshly warmed tunnel.
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(c, c.selectedReq)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
assert.Equal(t, 200, code,
"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
// grpOther's model must NOT leak to the grpMain client.
code, body = send(c, c.otherReq)
assert.Equal(t, 403, code,
"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
assert.Contains(t, body, "llm_policy.model_blocked",
"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
})
}
}
// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
//
// - union across the user's groups: the client is in gUX and gUY, each with
// its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
// client may use BOTH models (the union of its groups' allowlists) while a
// third, un-allowlisted model is denied.
// - an un-guardrailed group lifts the restriction: the client is in gMP and
// gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
// NO guardrail. Because one applicable policy is unrestricted, the client may
// use a model on no allowlist (mix-z) as well as mix-a.
func TestGuardrailMultiGroupUser(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
const (
unionA = "mg-union-a"
unionB = "mg-union-b"
unionC = "mg-union-c" // allowlisted by neither group
mixA = "mg-mix-a"
mixZ = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
)
priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
out := make([]api.AgentNetworkProviderModel, 0, len(ids))
for _, id := range ids {
out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
}
return &out
}
vllm, err := harness.StartVLLM(ctx, srv)
require.NoError(t, err, "start mock upstream")
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
mkGroup := func(name string) *api.Group {
g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
require.NoError(t, gerr, "create group %s", name)
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
return g
}
gUX := mkGroup("e2e-mg-union-x")
gUY := mkGroup("e2e-mg-union-y")
gMP := mkGroup("e2e-mg-mix-p")
gMQ := mkGroup("e2e-mg-mix-q")
ephemeral := false
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
Name: "e2e-mg-client",
Type: "reusable",
ExpiresIn: 86400,
UsageLimit: 0,
AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
Ephemeral: &ephemeral,
})
require.NoError(t, err, "mint setup key")
require.NotEmpty(t, sk.Key, "setup key plaintext")
staticKey := "static-e2e-token"
enabled := true
// P1 — union scenario: two restricting policies, one per group.
p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "e2e-mg-union",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: priced(unionA, unionB, unionC),
BootstrapCluster: ptr(harness.AgentNetworkCluster),
})
require.NoError(t, err, "create union provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-union-x",
Enabled: &enabled,
SourceGroups: []string{gUX.Id},
DestinationProviderIds: []string{p1.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
})
require.NoError(t, err, "create union policy X")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-union-y",
Enabled: &enabled,
SourceGroups: []string{gUY.Id},
DestinationProviderIds: []string{p1.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
})
require.NoError(t, err, "create union policy Y")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
Name: "e2e-mg-mix",
ProviderId: "openai_api",
UpstreamUrl: vllm.URL,
ApiKey: &staticKey,
Enabled: ptr(true),
Models: priced(mixA, mixZ),
})
require.NoError(t, err, "create mix provider")
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-mix-p",
Enabled: &enabled,
SourceGroups: []string{gMP.Id},
DestinationProviderIds: []string{p2.Id},
GuardrailIds: &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
})
require.NoError(t, err, "create mix policy P")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
Name: "e2e-mg-mix-q",
Enabled: &enabled,
SourceGroups: []string{gMQ.Id},
DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
})
require.NoError(t, err, "create mix policy Q")
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
settings, err := srv.GetSettings(ctx)
require.NoError(t, err, "read settings")
require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
require.NoError(t, err, "mint proxy token")
px, err := harness.StartProxy(ctx, srv, proxyToken)
require.NoError(t, err, "start proxy")
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
cl, err := harness.StartClient(ctx, srv, sk.Key)
require.NoError(t, err, "start client")
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
}
send := func(model string) (int, string) {
code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
require.NoError(t, cerr, "request must reach the proxy")
return code, body
}
sendUntil200 := func(model string) (int, string) {
var code int
var body string
deadline := time.Now().Add(90 * time.Second)
for time.Now().Before(deadline) {
code, body = send(model)
if code == 200 {
break
}
time.Sleep(5 * time.Second)
}
return code, body
}
t.Run("union across the user's groups", func(t *testing.T) {
code, body := sendUntil200(unionA)
assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = sendUntil200(unionB)
assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
code, body = send(unionC)
assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
assert.Contains(t, body, "llm_policy.model_blocked")
})
t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
code, body := sendUntil200(mixA)
assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
code, body = sendUntil200(mixZ)
assert.Equal(t, 200, code,
"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
})
}
// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
// exactly the given model, registering cleanup.
func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
t.Helper()
var gr api.AgentNetworkGuardrailRequest
gr.Name = name
gr.Checks.ModelAllowlist.Enabled = true
gr.Checks.ModelAllowlist.Models = []string{model}
g, err := srv.CreateGuardrail(ctx, gr)
require.NoError(t, err, "create guardrail %s", name)
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
return g
}

View File

@@ -52,7 +52,9 @@ func catalogModel(pc providerCase) string {
func disallowedModel(pc providerCase) string {
switch pc.kind {
case harness.WireBedrock:
return "us.anthropic.claude-opus-4-8"
// Same profile prefix as the allowed model so only the model name
// differs; the guardrail must deny it before it reaches AWS.
return strings.SplitN(pc.model, ".", 2)[0] + ".anthropic.claude-opus-4-8"
case harness.WireVertex:
return "claude-opus-4-8@20250101"
default:
@@ -72,7 +74,7 @@ func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint,
case harness.WireVertex:
code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "")
default:
code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
code, _, err = cl.ChatPrefixed(ctx, endpoint, proxyIP, pc.pathPrefix, pc.kind, model, "Reply with exactly: pong", "")
}
require.NoError(t, err, "request must reach the proxy for %s", pc.name)
return code
@@ -164,8 +166,7 @@ func TestModelAllowlistEnforced(t *testing.T) {
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {

View File

@@ -104,8 +104,7 @@ func TestProviderSkipTLSVerification(t *testing.T) {
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {

View File

@@ -106,8 +106,7 @@ func TestVLLMProvider(t *testing.T) {
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
// Resolve first: the DNS lookup triggers the lazy-connection warm-up, waking
// the proxy peer so WaitProxyPeer then observes it connected.
// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
require.NoError(t, err, "resolve endpoint to proxy IP")
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {

View File

@@ -4,6 +4,7 @@ package harness
import (
"context"
"errors"
"fmt"
"io"
"os/exec"
@@ -167,22 +168,54 @@ func (cl *Client) pollStatus(ctx context.Context, timeout time.Duration, want st
return fmt.Errorf("timed out waiting for %q; last status:\n%s", want, last)
}
// ResolveProxyIP resolves the agent-network endpoint to the proxy peer's
// NetBird IP from inside the client (via magic DNS).
const (
// curlExitCouldNotResolve is curl's exit code for a DNS resolution failure, distinct from connection-level failures.
curlExitCouldNotResolve = 6
// dnsProbeRetryWindow bounds DNS-failure retries: the synthesized zone lands a beat after management connects, so early NXDOMAIN is propagation; a zone still absent after this window is a real failure.
dnsProbeRetryWindow = 30 * time.Second
dnsProbeRetryInterval = 2 * time.Second
)
// ResolveProxyIP GETs https://<endpoint>/ from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; only DNS failures retry, within dnsProbeRetryWindow. Returns the connected IP for --resolve pinning.
func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
code, reader, err := cl.container.Exec(ctx, []string{"getent", "hosts", endpoint}, tcexec.Multiplexed())
if err != nil {
return "", err
args := []string{
"run", "--rm",
"--network", "container:" + cl.container.GetContainerID(),
curlImage,
"-ksS", "-o", "/dev/null",
"--connect-timeout", "30", "--max-time", "60",
"-w", "%{remote_ip}",
"https://" + endpoint + "/",
}
out, _ := io.ReadAll(reader)
if code != 0 {
return "", fmt.Errorf("getent hosts %s exited %d", endpoint, code)
deadline := time.Now().Add(dnsProbeRetryWindow)
for {
cmd := exec.CommandContext(ctx, "docker", args...)
var stdout, stderr strings.Builder
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err := cmd.Run()
if err == nil {
ip := strings.TrimSpace(stdout.String())
if ip == "" {
return "", fmt.Errorf("got an HTTP response from %s but no remote IP", endpoint)
}
return ip, nil
}
var exitErr *exec.ExitError
if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve {
return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String()))
}
dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String()))
if time.Until(deadline) < dnsProbeRetryInterval {
return "", dnsErr
}
select {
case <-ctx.Done():
return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err())
case <-time.After(dnsProbeRetryInterval):
}
}
fields := strings.Fields(string(out))
if len(fields) == 0 {
return "", fmt.Errorf("no address for %s", endpoint)
}
return fields[0], nil
}
// Wire shapes for Chat.
@@ -206,18 +239,29 @@ const (
// the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). A non-empty
// sessionID is sent as the universal x-session-id header the proxy records.
func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) {
return cl.ChatPrefixed(ctx, endpoint, proxyIP, "", kind, model, prompt, sessionID)
}
// ChatPrefixed is Chat with a base-URL path prefix prepended to the wire
// path, mirroring agents whose base URL carries a shape-selecting prefix that
// rides through to the upstream — e.g. Claude Code against a Kimi provider
// sets ANTHROPIC_BASE_URL=https://<endpoint>/anthropic so the proxy forwards
// /anthropic/v1/messages to Moonshot's Anthropic surface while the provider's
// upstream URL stays the bare https://api.moonshot.ai. Empty prefix is plain
// Chat.
func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefix, kind, model, prompt, sessionID string) (int, string, error) {
var path, body string
var headers []string
switch kind {
case WireMessages:
path = "/v1/messages"
headers = []string{"anthropic-version: 2023-06-01"}
body = fmt.Sprintf(`{"model":%q,"max_tokens":64,"messages":[{"role":"user","content":%q}]}`, model, prompt)
body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, model, prompt)
default:
path = "/v1/chat/completions"
body = fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":%q}]}`, model, prompt)
}
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID))
return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
}
// Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
@@ -227,7 +271,7 @@ func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prom
// is sent as the universal x-session-id header the proxy records.
func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt, sessionID string) (int, string, error) {
path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", project, region, model)
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
@@ -238,7 +282,7 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region
// header the proxy records.
func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) {
path := "/model/" + model + "/invoke"
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
}
@@ -297,7 +341,7 @@ func (cl *Client) Terminate(ctx context.Context) error {
return cl.container.Terminate(ctx)
}
// containerLogs reads up to 256 KiB of a container's logs for diagnostics.
// containerLogs reads up to 4 MiB of a container's logs for diagnostics — enough for a whole provider-matrix run.
func containerLogs(ctx context.Context, c testcontainers.Container) string {
if c == nil {
return ""
@@ -307,6 +351,6 @@ func containerLogs(ctx context.Context, c testcontainers.Container) string {
return fmt.Sprintf("<logs error: %v>", err)
}
defer r.Close()
b, _ := io.ReadAll(io.LimitReader(r, 256<<10))
b, _ := io.ReadAll(io.LimitReader(r, 4<<20))
return string(b)
}

View File

@@ -221,6 +221,29 @@ func (c *Combined) CreateProxyTokenCLI(ctx context.Context, name string) (string
return "", fmt.Errorf("token not found in CLI output: %s", string(out))
}
// SnapshotStoreDB copies the management sqlite store (with WAL/SHM sidecars) out of the bind-mounted
// data dir into dstDir and returns the copy's path; reading a copy avoids locking against live writes.
func (c *Combined) SnapshotStoreDB(dstDir string) (string, error) {
src := filepath.Join(c.workDir, "data", "store.db")
if _, err := os.Stat(src); err != nil {
return "", fmt.Errorf("management store not found at %s: %w", src, err)
}
dst := filepath.Join(dstDir, "store.db")
for _, suffix := range []string{"", "-wal", "-shm"} {
data, err := os.ReadFile(src + suffix)
if err != nil {
if os.IsNotExist(err) && suffix != "" {
continue // sidecar only exists in WAL mode
}
return "", fmt.Errorf("read %s: %w", src+suffix, err)
}
if err := os.WriteFile(dst+suffix, data, 0o600); err != nil {
return "", fmt.Errorf("write %s: %w", dst+suffix, err)
}
}
return dst, nil
}
// Logs returns the combined server container logs, for diagnostics.
func (c *Combined) Logs(ctx context.Context) string {
return containerLogs(ctx, c.container)

View File

@@ -38,7 +38,11 @@ type Proxy struct {
// network, registered via the given account proxy token and serving the
// AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
// peer connectivity — callers poll management for the proxy peer.
func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
// StartProxy launches the reverse-proxy container. Optional envOverrides are
// merged into the container environment after the defaults, so callers can set
// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
// need a short authorization-cache window).
func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
root, err := repoRoot()
if err != nil {
return nil, err
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
}
for _, ov := range envOverrides {
for k, v := range ov {
req.Env[k] = v
}
}
ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
ContainerRequest: req,
Started: true,

View File

@@ -234,9 +234,6 @@ init_environment() {
NETBIRD_LICENSE_KEY=$(read_secret "Enter license key (input hidden)")
GHCR_USERNAME="netbirdExtAccess1"
GHCR_TOKEN=$(read_secret "Enter GHCR token (input hidden)")
POSTGRES_USER="netbird"
POSTGRES_DB="netbird"
POSTGRES_PASSWORD=$(rand_secret)
@@ -263,10 +260,6 @@ init_environment() {
install -m 600 /dev/null config.yaml
render_config_yaml >> config.yaml
echo "Logging in to ghcr.io ..."
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
unset GHCR_TOKEN
echo ""
echo "Pulling images ..."
$DOCKER_COMPOSE_COMMAND pull

View File

@@ -490,8 +490,6 @@ init_migration() {
echo ""
echo "Step 1: Image swap (community → Enterprise). License key required."
NB_LICENSE_KEY=$(read_secret " License key")
GHCR_USERNAME="netbirdExtAccess1"
GHCR_TOKEN=$(read_secret " GHCR token (input hidden)")
# Step 2 — optional
echo ""
@@ -588,11 +586,6 @@ apply_changes() {
fi
} >> "$ENV_FILE"
echo ""
echo "Logging in to ghcr.io ..."
printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
unset GHCR_TOKEN
echo ""
echo "Pulling enterprise images ..."
$DOCKER_COMPOSE_COMMAND pull

View File

@@ -18,21 +18,26 @@ import (
// contract between the proxy and management; management flattens them into
// queryable columns. Keep in sync with the proxy side.
const (
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDTotal = "cost.usd_total"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
metaKeyProvider = "llm.provider"
metaKeyModel = "llm.model"
metaKeyResolvedProviderID = "llm.resolved_provider_id"
metaKeySelectedPolicyID = "llm.selected_policy_id"
metaKeyPolicyDecision = "llm_policy.decision"
metaKeyPolicyReason = "llm_policy.reason"
metaKeyInputTokens = "llm.input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyOutputTokens = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyTotalTokens = "llm.total_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCachedInputTokens = "llm.cached_input_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
metaKeyCostUSDInput = "cost.usd_input"
metaKeyCostUSDCachedInput = "cost.usd_cached_input"
metaKeyCostUSDCacheCreate = "cost.usd_cache_creation"
metaKeyCostUSDOutput = "cost.usd_output"
metaKeyStream = "llm.stream"
metaKeySessionID = "llm.session_id"
metaKeyAuthorisingGroups = "llm.authorising_groups"
metaKeyRequestPrompt = "llm.request_prompt"
metaKeyResponseCompletion = "llm.response_completion"
)
// IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -108,20 +113,25 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
BytesUpload: e.BytesUpload,
BytesDownload: e.BytesDownload,
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CostUSD: parseMetaFloat(meta, metaKeyCostUSDTotal),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
Provider: meta[metaKeyProvider],
Model: meta[metaKeyModel],
SessionID: meta[metaKeySessionID],
ResolvedProviderID: meta[metaKeyResolvedProviderID],
SelectedPolicyID: meta[metaKeySelectedPolicyID],
Decision: meta[metaKeyPolicyDecision],
DenyReason: meta[metaKeyPolicyReason],
InputTokens: parseMetaInt(meta, metaKeyInputTokens),
OutputTokens: parseMetaInt(meta, metaKeyOutputTokens),
TotalTokens: parseMetaInt(meta, metaKeyTotalTokens),
CachedInputTokens: parseMetaInt(meta, metaKeyCachedInputTokens),
CacheCreationTokens: parseMetaInt(meta, metaKeyCacheCreationTokens),
InputCostUSD: parseMetaFloat(meta, metaKeyCostUSDInput),
CachedInputCostUSD: parseMetaFloat(meta, metaKeyCostUSDCachedInput),
CacheCreationCostUSD: parseMetaFloat(meta, metaKeyCostUSDCacheCreate),
OutputCostUSD: parseMetaFloat(meta, metaKeyCostUSDOutput),
Stream: parseMetaBool(meta, metaKeyStream),
RequestPrompt: meta[metaKeyRequestPrompt],
ResponseCompletion: meta[metaKeyResponseCompletion],
}
var groups []types.AgentNetworkAccessLogGroup
@@ -140,18 +150,23 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
// log's ID so the two correlate.
func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
usage := &types.AgentNetworkUsage{
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CostUSD: e.CostUSD,
ID: e.ID,
AccountID: e.AccountID,
Timestamp: e.Timestamp,
UserID: e.UserID,
ResolvedProviderID: e.ResolvedProviderID,
Provider: e.Provider,
Model: e.Model,
SessionID: e.SessionID,
InputTokens: e.InputTokens,
OutputTokens: e.OutputTokens,
TotalTokens: e.TotalTokens,
CachedInputTokens: e.CachedInputTokens,
CacheCreationTokens: e.CacheCreationTokens,
InputCostUSD: e.InputCostUSD,
CachedInputCostUSD: e.CachedInputCostUSD,
CacheCreationCostUSD: e.CacheCreationCostUSD,
OutputCostUSD: e.OutputCostUSD,
}
usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))

View File

@@ -28,17 +28,22 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
UserId: "user-1",
AgentNetwork: true,
Metadata: map[string]string{
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "150",
metaKeyCostUSDTotal: "0.0123",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
metaKeyProvider: "openai",
metaKeyModel: "gpt-5.4",
metaKeyResolvedProviderID: "prov-1",
metaKeySessionID: "sess-1",
metaKeyInputTokens: "100",
metaKeyOutputTokens: "50",
metaKeyTotalTokens: "1174",
metaKeyCachedInputTokens: "256",
metaKeyCacheCreationTokens: "768",
metaKeyCostUSDInput: "0.0071",
metaKeyCostUSDCachedInput: "0.0009",
metaKeyCostUSDCacheCreate: "0.0020",
metaKeyCostUSDOutput: "0.0023",
metaKeyStream: "true",
metaKeyRequestPrompt: "hello",
metaKeyResponseCompletion: "world",
// repeated id must be de-duplicated before the group rows insert.
metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
},
@@ -65,7 +70,19 @@ func TestIngestAccessLog_RealStore_LogCollectionOff(t *testing.T) {
require.Len(t, usage, 1, "usage row must be written even with log collection off")
assert.Equal(t, int64(100), usage[0].InputTokens, "input tokens must round-trip from metadata")
assert.Equal(t, int64(50), usage[0].OutputTokens, "output tokens must round-trip from metadata")
assert.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost must round-trip from metadata")
assert.Equal(t, int64(256), usage[0].CachedInputTokens, "cache-read tokens must round-trip from metadata")
assert.Equal(t, int64(768), usage[0].CacheCreationTokens, "cache-write tokens must round-trip from metadata")
// The per-bucket breakdown is the only cost state stored, and must survive
// the write/read cycle as real columns — usage rows are the only cost
// record for accounts with log collection off, so a dropped column here
// loses the split permanently.
assert.InDelta(t, 0.0071, usage[0].InputCostUSD, 1e-9, "input cost must round-trip from metadata")
assert.InDelta(t, 0.0009, usage[0].CachedInputCostUSD, 1e-9, "cache-read cost must round-trip from metadata")
assert.InDelta(t, 0.0020, usage[0].CacheCreationCostUSD, 1e-9, "cache-write cost must round-trip from metadata")
assert.InDelta(t, 0.0023, usage[0].OutputCostUSD, 1e-9, "output cost must round-trip from metadata")
// Aggregates are derived from the stored columns, never stored themselves.
assert.InDelta(t, 0.0123, usage[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0029, usage[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
require.NoError(t, err)
@@ -96,6 +113,14 @@ func TestIngestAccessLog_RealStore_LogCollectionOn(t *testing.T) {
require.Equal(t, int64(1), total, "exactly one access-log row expected")
require.Len(t, logs, 1, "full access-log row must be written when log collection is on")
assert.Equal(t, "gpt-5.4", logs[0].Model, "model must flatten from metadata")
assert.Equal(t, int64(256), logs[0].CachedInputTokens, "cache-read tokens must flatten from metadata")
assert.Equal(t, int64(768), logs[0].CacheCreationTokens, "cache-write tokens must flatten from metadata")
assert.InDelta(t, 0.0029, logs[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
assert.InDelta(t, 0.0123, logs[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
assert.InDelta(t, 0.0071, logs[0].InputCostUSD, 1e-9, "input cost must flatten from metadata")
assert.InDelta(t, 0.0009, logs[0].CachedInputCostUSD, 1e-9, "cache-read cost must flatten from metadata")
assert.InDelta(t, 0.0020, logs[0].CacheCreationCostUSD, 1e-9, "cache-write cost must flatten from metadata")
assert.InDelta(t, 0.0023, logs[0].OutputCostUSD, 1e-9, "output cost must flatten from metadata")
assert.Equal(t, "hello", logs[0].RequestPrompt, "prompt must be retained when log collection is on")
assert.Equal(t, "world", logs[0].ResponseCompletion, "completion must be retained when log collection is on")
assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")

View File

@@ -38,7 +38,7 @@ func accessLogRow(id, sessionID string, ts time.Time, opts ...func(*types.AgentN
InputTokens: 100,
OutputTokens: 50,
TotalTokens: 150,
CostUSD: 0.01,
InputCostUSD: 0.01,
}
for _, o := range opts {
o(e)
@@ -74,7 +74,7 @@ func withTokens(in, out, total int64, cost float64) func(*types.AgentNetworkAcce
e.InputTokens = in
e.OutputTokens = out
e.TotalTokens = total
e.CostUSD = cost
e.InputCostUSD = cost
}
}
@@ -155,7 +155,7 @@ func TestAccessLogSessions_FoldAndAggregate(t *testing.T) {
assert.Equal(t, int64(310), a.InputTokens, "input tokens summed")
assert.Equal(t, int64(135), a.OutputTokens, "output tokens summed")
assert.Equal(t, int64(445), a.TotalTokens, "total tokens summed")
assert.InDelta(t, 0.031, a.CostUSD, 1e-9, "cost summed")
assert.InDelta(t, 0.031, a.TotalCostUSD(), 1e-9, "cost summed")
assert.Equal(t, "deny", a.Decision, "any deny makes the session a deny")
assert.ElementsMatch(t, []string{"openai", "anthropic"}, a.Providers, "distinct providers")
assert.ElementsMatch(t, []string{"gpt-5.4", "claude-haiku-4-5"}, a.Models, "distinct models")

View File

@@ -420,6 +420,47 @@ var providers = []Provider{
{ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192},
},
},
{
ID: "kimi_api",
Kind: KindProvider,
Name: "Kimi (Moonshot AI) API",
Description: "Kimi K3 / K2 models via the Moonshot AI platform",
DefaultHost: "api.moonshot.ai",
AuthHeaderName: "Authorization",
AuthHeaderTemplate: "Bearer ${API_KEY}",
DefaultContentType: "application/json",
BrandColor: "#1A1A2E",
// ParserID empty on purpose: Moonshot serves two body shapes on
// the same host and key, and the proxy's URL sniffer dispatches
// both (same pattern as Bifrost). /v1/chat/completions matches
// OpenAIParser; the Anthropic-compatible endpoint the official
// Claude Code guide uses (/anthropic/v1/messages) contains
// "/v1/messages" and matches AnthropicParser. Pinning "openai"
// here would misparse the Claude Code path — the primary way
// teams consume Kimi for coding today. Both endpoints accept the
// same Moonshot key via Authorization: Bearer (Claude Code's
// ANTHROPIC_AUTH_TOKEN rides that header too).
//
// api.moonshot.ai is the international platform; mainland-China
// accounts live on api.moonshot.cn with separate billing —
// operators there override the host on the provider record. The
// kimi.com subscription coding endpoint (api.kimi.com/coding,
// model id "k3") is account-bound seat licensing rather than a
// meterable platform key, so it's deliberately not the default.
ParserID: "",
// Pricing per Moonshot's platform rates at K3 launch (July 2026):
// $3/$15 per MTok with $0.30 cached input, flat across the 1M-token
// window. kimi-k3 is the ONLY model the platform serves newer
// accounts — K2-era ids (kimi-k2-thinking) and even the kimi-latest
// alias return resource_not_found_error, verified live 2026-07-21 —
// so it's the only catalog entry. Grandfathered accounts with K2
// access can still type those ids on the provider's model rows.
// The consumer app's "K3 Swarm Max" mode is not an API SKU, so it
// doesn't appear here.
Models: []Model{
{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
},
},
{
ID: "litellm_proxy",
Kind: KindGateway,

View File

@@ -86,12 +86,17 @@ type Manager interface {
// PolicySelectionInput is the per-request selection envelope. The
// proxy populates it from CapturedData (account, user, groups) plus
// the provider llm_router resolved.
// the provider llm_router resolved and the model it extracted.
type PolicySelectionInput struct {
AccountID string
UserID string
GroupIDs []string
ProviderID string
// Model is the already-normalised upstream model id the proxy extracted
// (parser strips Bedrock region/version, Vertex @version), so a
// case-insensitive compare suffices. Empty = undetermined → not permitted
// (fail closed).
Model string
}
// PolicySelectionResult names the policy that "pays" for this request

View File

@@ -5,6 +5,7 @@ import (
"fmt"
"math"
"sort"
"strings"
"time"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
@@ -35,6 +36,10 @@ const (
denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
//nolint:gosec // account deny code label, not a credential
denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
// denyCodeModelBlocked is returned when policies govern the request's
// (provider, caller-groups) but none permits the model. Matches the proxy
// guardrail's code so both layers surface the same label.
denyCodeModelBlocked = "llm_policy.model_blocked"
)
// consumptionCache holds the consumption counters prefetched for one
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
}
candidates := filterApplicablePolicies(policies, in)
// Model-allowlist gate scoped to the matched policies: keep candidates whose
// guardrails permit the model (none enabled = unrestricted), deny when
// policies apply but none permits it. Skip the load when none has a guardrail.
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
if gErr != nil {
return nil, gErr
}
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
if len(permitted) == 0 {
return &PolicySelectionResult{
Allow: false,
DenyCode: denyCodeModelBlocked,
DenyReason: modelBlockedReason(in.Model),
}, nil
}
candidates = permitted
}
// Prefetch every consumption counter the ceiling + candidate policies will
// read, in a single store round-trip, then score against the cache.
cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
return out
}
// anyPolicyHasGuardrails reports whether any policy references at least one
// guardrail, so the selector can skip loading guardrails when none do.
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
for _, p := range policies {
if p != nil && len(p.GuardrailIDs) > 0 {
return true
}
}
return false
}
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
// model-allowlist gate to resolve each candidate policy's attached guardrails.
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
if err != nil {
return nil, fmt.Errorf("list account guardrails: %w", err)
}
byID := make(map[string]*types.Guardrail, len(guardrails))
for _, g := range guardrails {
if g != nil {
byID[g.ID] = g
}
}
return byID, nil
}
// filterModelPermittedPolicies returns the subset of policies whose guardrails
// permit the model. Order is preserved so downstream scoring is unaffected.
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
out := make([]*types.Policy, 0, len(policies))
for _, p := range policies {
if policyPermitsModel(p, byID, model) {
out = append(out, p)
}
}
return out
}
// policyPermitsModel reports whether a policy permits the model. No
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
// otherwise the model must be in the union of its allowlists, so an
// empty/undetermined model fails closed.
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
if p == nil {
return false
}
wanted := normaliseModelID(model)
restricted := false
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
if wanted == "" {
continue
}
for _, allowed := range g.Checks.ModelAllowlist.Models {
if normaliseModelID(allowed) == wanted {
return true
}
}
}
return !restricted
}
// normaliseModelID lowercases and trims a model identifier so the allowlist
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
// normaliseModel so both layers agree on what "same model" means.
func normaliseModelID(model string) string {
return strings.ToLower(strings.TrimSpace(model))
}
// modelBlockedReason builds the human-readable deny reason for a model-allowlist
// rejection. The model is quoted when known; an undetermined model is reported
// as such so the access log distinguishes "wrong model" from "no model".
func modelBlockedReason(model string) string {
if normaliseModelID(model) == "" {
return "request model could not be determined for the policy allowlist"
}
return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
}
// candidate is the per-policy intermediate the selector ranks. A
// policy that's been exhausted on any enabled cap never makes it
// into this slice; the selector's deny envelope carries the latest

View File

@@ -0,0 +1,329 @@
package agentnetwork
import (
"context"
"errors"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/store"
)
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
// to reach providerID under the given guardrails. Uncapped keeps the selector's
// headroom scoring trivial so these tests isolate the model-allowlist gate.
func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
AccountID: account,
Enabled: true,
SourceGroups: sourceGroups,
DestinationProviderIDs: []string{providerID},
GuardrailIDs: guardrailIDs,
CreatedAt: time.Now().UTC(),
}
}
// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
// carries the given models.
func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
return &types.Guardrail{
ID: id,
AccountID: account,
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
},
}
}
func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
mockStore.EXPECT().
GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
Return(policies, nil)
}
func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
Return(guardrails, nil)
}
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
// decision: a policy authorises the (provider, group) but restricts the model,
// and the requested model isn't on the list, so the request is denied.
func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
}
// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
// is on the applicable policy's allowlist, so selection proceeds normally.
func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
UserID: "user-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
// and surrounding whitespace, matching the proxy guardrail's normalisation.
func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.NoError(t, err)
assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
}
// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
// policies authorise the same (provider, group) and one has no guardrail, that
// policy makes the request unrestricted — not caught by the other's allowlist.
func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
expectPolicies(mockStore, "acc-1", restricted, open)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
}
// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
// model allowlisted only for grp-b must not be usable by a grp-a caller. The
// selector considers only policies applicable to the caller's groups.
func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
expectPolicies(mockStore, "acc-1", polA, polB)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-a"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only allowed for grp-b
})
require.NoError(t, err)
assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
// mirrors the proxy: with a restricted applicable policy and an empty model
// (e.g. a path-routed shape the parser couldn't map), the request is denied.
func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "", // undetermined
})
require.NoError(t, err)
assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
}
// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
// model allowlist is disabled imposes no model restriction, even though the
// policy references it.
func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
disabled := &types.Guardrail{
ID: "g-1",
AccountID: "acc-1",
Checks: types.GuardrailChecks{
ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
},
}
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1", disabled)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
// allowlist guardrails permits the union of their models (not just the first).
func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only in the second guardrail's list
})
require.NoError(t, err)
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
// resolving the candidate policies' guardrails surfaces as an error, not a
// silent allow/deny.
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
Return(nil, errors.New("store unavailable"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
assert.Nil(t, res)
}
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
// policy referencing a guardrail ID absent from the account's set (a stale
// reference) imposes no model restriction — same as no guardrail.
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1")
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
// gate narrows candidates before cap scoring: the permitting policy is selected
// even though the blocked one has a larger, more attractive cap.
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
polBig.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
}
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
polSmall.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
}
expectPolicies(mockStore, "acc-1", polBig, polSmall)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only pol-small's guardrail permits this
})
require.NoError(t, err)
assert.True(t, res.Allow)
assert.Equal(t, "pol-small", res.SelectedPolicyID,
"the model filter must exclude pol-big before cap scoring")
}

View File

@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
applyAccountCollectionControls(&mergedGuardrails, settings)
guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
// The proxy guardrail is a per-provider fail-closed backstop; the
// authoritative per-policy/group decision is management's
// SelectPolicyForRequest. A provider lands in this map only when every
// authorising policy restricts models.
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
if err != nil {
return nil, err
}
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
// guardrailConfig is the JSON shape the proxy-side llm_guardrail
// middleware expects. Mirrors the proxy registration documented in
// the management→proxy contract.
// the management→proxy contract. provider_allowlists is keyed by the
// resolved provider id llm_router stamps; a provider absent from the map is
// unrestricted at the proxy layer.
type guardrailConfig struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture guardrailPromptCapture `json:"prompt_capture"`
}
type guardrailPromptCapture struct {
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
}
func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
cfg := guardrailConfig{
ModelAllowlist: merged.ModelAllowlist,
PromptCapture: guardrailPromptCapture{
Enabled: merged.PromptCapture.Enabled,
RedactPii: merged.PromptCapture.RedactPii,
},
ProviderAllowlists: providerAllowlists,
PromptCapture: guardrailPromptCapture(capture),
}
out, err := json.Marshal(cfg)
if err != nil {
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
return out, nil
}
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
// is included only when every authorising policy restricts models (their union);
// if any leaves it unrestricted it is omitted, so management decides per group.
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
type providerAcc struct {
models map[string]struct{}
anyUnrestricted bool
}
accs := make(map[string]*providerAcc)
for _, p := range policies {
if p == nil {
continue
}
restricted, models := policyModelAllowlist(p, byID)
for _, providerID := range p.DestinationProviderIDs {
if providerID == "" {
continue
}
acc, ok := accs[providerID]
if !ok {
acc = &providerAcc{models: make(map[string]struct{})}
accs[providerID] = acc
}
if !restricted {
acc.anyUnrestricted = true
continue
}
for _, m := range models {
acc.models[m] = struct{}{}
}
}
}
out := make(map[string][]string, len(accs))
for providerID, acc := range accs {
if acc.anyUnrestricted {
continue
}
models := make([]string, 0, len(acc.models))
for m := range acc.models {
models = append(models, m)
}
sort.Strings(models)
out[providerID] = models
}
return out
}
// policyModelAllowlist reports whether a policy restricts models (has an
// allowlist-enabled guardrail) and the union of allowed models. Models are
// verbatim; the proxy factory lowercases/trims them at decode time.
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
restricted := false
var models []string
for _, gID := range p.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
continue
}
restricted = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
models = append(models, m)
}
}
}
return restricted, models
}
// buildAccountService composes the per-account gateway Service. The
// target carries the noop placeholder URL — the router middleware
// rewrites every request to the matched provider's upstream before the
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
return out
}
// MergedGuardrails is the JSON shape passed to the proxy via the
// guardrail middleware's config_json. Mirrors the proxy-side
// expectations and is intentionally distinct from
// types.GuardrailChecks so we can evolve either side independently.
// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
// merged here — the model allowlist is emitted per-provider, and
// token/budget/retention moved onto Policy.Limits and account Settings.
type MergedGuardrails struct {
ModelAllowlist []string `json:"model_allowlist,omitempty"`
TokenLimits MergedTokenLimits `json:"token_limits"`
Budget MergedBudget `json:"budget"`
PromptCapture MergedPromptCapture `json:"prompt_capture"`
Retention MergedRetention `json:"retention"`
}
type MergedTokenLimits struct {
Hourly *MergedTokenWindow `json:"hourly,omitempty"`
Daily *MergedTokenWindow `json:"daily,omitempty"`
Monthly *MergedTokenWindow `json:"monthly,omitempty"`
}
type MergedTokenWindow struct {
MaxInputTokens int `json:"max_input_tokens,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
}
type MergedBudget struct {
Hourly *MergedBudgetWindow `json:"hourly,omitempty"`
Daily *MergedBudgetWindow `json:"daily,omitempty"`
Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
}
type MergedBudgetWindow struct {
SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
PromptCapture MergedPromptCapture
}
type MergedPromptCapture struct {
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
RedactPii bool `json:"redact_pii"`
}
type MergedRetention struct {
Enabled bool `json:"enabled"`
Days int `json:"days"`
}
// mergeGuardrails computes the effective guardrail spec applied at the
// proxy, given the referencing policies and the account's guardrail
// catalogue. Policy enabled-ness is the caller's responsibility — only
// enabled policies should be passed in.
// mergeGuardrails folds the referencing policies' guardrails into the
// prompt-capture decision only. The model allowlist is enforced per-policy/group
// in management and shipped per-provider; token/budget/retention live off
// guardrails now.
//
// Merge rules:
// - Model allowlist: union of allowlists across policies that enable it.
// - Token / Budget: most-restrictive (min of non-zero caps) per window.
// - Prompt capture: enabled if any policy enables it; redact_pii sticks
// if any enabling policy turns it on.
// - Retention: enabled if any enables it; smallest non-zero days wins.
// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
// sticks if any enabling policy turns it on.
func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
merged := MergedGuardrails{}
allowlist := make(map[string]struct{})
allowlistEnabled := false
for _, policy := range policies {
for _, gID := range policy.GuardrailIDs {
g, ok := byID[gID]
if !ok || g == nil {
continue
}
mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
mergeGuardrail(g, &merged)
}
}
if allowlistEnabled {
merged.ModelAllowlist = make([]string, 0, len(allowlist))
for m := range allowlist {
merged.ModelAllowlist = append(merged.ModelAllowlist, m)
}
sort.Strings(merged.ModelAllowlist)
}
return merged
}
// mergeGuardrail folds a single guardrail's enabled checks into the
// running merge: model-allowlist models join the shared set (and flip
// allowlistEnabled), and prompt-capture / redact-pii stick once any
// enabling guardrail turns them on.
//
// TokenLimits, Budget, and Retention have moved off guardrails — token
// and budget caps now live on the Policy itself (Policy.Limits) and
// retention moves to account-level Settings — so they are not merged here.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
if g.Checks.ModelAllowlist.Enabled {
*allowlistEnabled = true
for _, m := range g.Checks.ModelAllowlist.Models {
if m != "" {
allowlist[m] = struct{}{}
}
}
}
// mergeGuardrail folds a single guardrail's prompt-capture settings into the
// running merge: enabled / redact-pii stick once any enabling guardrail turns
// them on.
func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
if g.Checks.PromptCapture.Enabled {
merged.PromptCapture.Enabled = true
if g.Checks.PromptCapture.RedactPii {

View File

@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
require.Len(t, services, 1)
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
"model allowlist is a pure policy guardrail and must always reach the config")
assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
"model allowlist is a pure policy guardrail and must reach the per-provider config")
assert.False(t, cfg.PromptCapture.Enabled,
"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
}
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
require.Len(t, services, 1, "exactly one synth service expected")
cfg := decodeServiceGuardrailConfig(t, services[0])
assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
}

View File

@@ -0,0 +1,95 @@
package agentnetwork
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
)
// policyForProviders builds an enabled policy authorising the given providers
// under the given guardrails (both optional). Groups are irrelevant to
// buildProviderAllowlists, which keys purely on destination provider.
func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
return &types.Policy{
ID: id,
Enabled: true,
DestinationProviderIDs: providerIDs,
GuardrailIDs: guardrailIDs,
}
}
func TestBuildProviderAllowlists(t *testing.T) {
byID := map[string]*types.Guardrail{
"g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
"g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
}
t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
"a provider every policy restricts carries the sorted union of their models")
})
t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", nil, "prov-x"), // no guardrail
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
})
t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.NotContains(t, got, "prov-x",
"a policy whose only guardrail has a disabled allowlist is unrestricted")
})
t.Run("providers are isolated from one another", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
})
t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
}
got := buildProviderAllowlists(policies, byID)
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
})
t.Run("union across a single policy's guardrails", func(t *testing.T) {
policies := []*types.Policy{
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
}
got := buildProviderAllowlists(policies, byID)
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
"a policy's own multiple allowlist guardrails union together")
})
t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
got := buildProviderAllowlists([]*types.Policy{
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
}, empty)
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
})
}

View File

@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
var cfg guardrailConfig
require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
"model allowlist union must keep both models")
// Both policies restrict the same provider, so the per-provider backstop
// carries the union of their models — a coarse gate that management's
// per-policy/group check narrows; it only blocks models outside the union
// when management is down.
assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
"per-provider allowlist union must keep both models")
}
func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {

View File

@@ -41,8 +41,24 @@ type AgentNetworkAccessLog struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
Stream bool
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown — one column per token bucket the provider
// bills separately. These four are the only cost state stored: the total
// and the cache portion are derived on read (TotalCostUSD / CacheCostUSD)
// rather than stored alongside, so a stored aggregate can never drift out
// of step with the components it summarises.
//
// default:0 matters on upgrade: these columns are ALTER TABLE ADD COLUMN
// on an existing table, and without it every historical row holds NULL —
// which a raw SUM()/scan into float64 can't read. The default backfills
// them as 0, so pre-upgrade rows report an unknown split, not an error.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
Stream bool
// Prompt capture. Only populated when prompt collection is enabled
// (account master switch AND policy guardrail). Heavy free text.
@@ -60,19 +76,44 @@ type AgentNetworkAccessLog struct {
// the reverse-proxy AccessLogEntry table.
func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_log" }
// CostUSDSQLExpr is the SQL sum of the per-bucket cost columns — the total cost
// of a row. Used wherever a query has to sort or aggregate on total cost now
// that no cost_usd column is stored. Plain arithmetic over NOT NULL columns, so
// it stays portable across SQLite and Postgres.
const CostUSDSQLExpr = "(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd)"
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (a *AgentNetworkAccessLog) TotalCostUSD() float64 {
return a.InputCostUSD + a.CachedInputCostUSD + a.CacheCreationCostUSD + a.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets:
// cache reads plus cache writes.
func (a *AgentNetworkAccessLog) CacheCostUSD() float64 {
return a.CachedInputCostUSD + a.CacheCreationCostUSD
}
// ToAPIResponse renders the flattened entry as the API representation.
func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
out := api.AgentNetworkAccessLog{
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CostUsd: a.CostUSD,
Stream: &a.Stream,
Id: a.ID,
ServiceId: a.ServiceID,
Timestamp: a.Timestamp,
StatusCode: a.StatusCode,
DurationMs: int(a.Duration.Milliseconds()),
InputTokens: a.InputTokens,
OutputTokens: a.OutputTokens,
TotalTokens: a.TotalTokens,
CachedInputTokens: a.CachedInputTokens,
CacheCreationTokens: a.CacheCreationTokens,
InputCostUsd: a.InputCostUSD,
CachedInputCostUsd: a.CachedInputCostUSD,
CacheCreationCostUsd: a.CacheCreationCostUSD,
OutputCostUsd: a.OutputCostUSD,
CostUsd: a.TotalCostUSD(),
CacheCostUsd: a.CacheCostUSD(),
Stream: &a.Stream,
}
out.UserId = strPtr(a.UserID)
@@ -112,20 +153,36 @@ func strPtr(s string) *string {
// summary plus its ordered entries. Assembled in Go from a page of entries — it
// is not a stored table.
type AgentNetworkAccessLogSession struct {
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
SessionID string // empty for a session-less (singleton) request
UserID string
GroupIDs []string // union of the entries' authorising groups
StartedAt time.Time
EndedAt time.Time
RequestCount int
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
Providers []string // distinct vendors seen in the session
Models []string // distinct models seen in the session
Decision string // "deny" if any entry was denied, else "allow"
Entries []*AgentNetworkAccessLog
}
// TotalCostUSD is the session's total cost: the sum of the four per-bucket
// costs accumulated across its entries.
func (sess *AgentNetworkAccessLogSession) TotalCostUSD() float64 {
return sess.InputCostUSD + sess.CachedInputCostUSD + sess.CacheCreationCostUSD + sess.OutputCostUSD
}
// CacheCostUSD is the session's prompt-cache spend: cache reads plus writes.
func (sess *AgentNetworkAccessLogSession) CacheCostUSD() float64 {
return sess.CachedInputCostUSD + sess.CacheCreationCostUSD
}
// sessionKey is the grouping key for an entry: its session id, or — when the
@@ -205,7 +262,12 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
sess.InputTokens += e.InputTokens
sess.OutputTokens += e.OutputTokens
sess.TotalTokens += e.TotalTokens
sess.CostUSD += e.CostUSD
sess.CachedInputTokens += e.CachedInputTokens
sess.CacheCreationTokens += e.CacheCreationTokens
sess.InputCostUSD += e.InputCostUSD
sess.CachedInputCostUSD += e.CachedInputCostUSD
sess.CacheCreationCostUSD += e.CacheCreationCostUSD
sess.OutputCostUSD += e.OutputCostUSD
if e.Timestamp.Before(sess.StartedAt) {
sess.StartedAt = e.Timestamp
}
@@ -248,15 +310,22 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
}
out := api.AgentNetworkAccessLogSession{
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CostUsd: sess.CostUSD,
Decision: sess.Decision,
Entries: entries,
StartedAt: sess.StartedAt,
EndedAt: sess.EndedAt,
RequestCount: sess.RequestCount,
InputTokens: sess.InputTokens,
OutputTokens: sess.OutputTokens,
TotalTokens: sess.TotalTokens,
CachedInputTokens: sess.CachedInputTokens,
CacheCreationTokens: sess.CacheCreationTokens,
InputCostUsd: sess.InputCostUSD,
CachedInputCostUsd: sess.CachedInputCostUSD,
CacheCreationCostUsd: sess.CacheCreationCostUSD,
OutputCostUsd: sess.OutputCostUSD,
CostUsd: sess.TotalCostUSD(),
CacheCostUsd: sess.CacheCostUSD(),
Decision: sess.Decision,
Entries: entries,
}
out.SessionId = strPtr(sess.SessionID)
out.UserId = strPtr(sess.UserID)

View File

@@ -54,7 +54,7 @@ var accessLogSortFields = map[string]string{
"provider": "provider",
"status_code": "status_code",
"duration": "duration",
"cost_usd": "cost_usd",
"cost_usd": CostUSDSQLExpr,
"total_tokens": "total_tokens",
"user_id": "user_id",
"decision": "decision",
@@ -70,7 +70,7 @@ var accessLogSortFields = map[string]string{
var sessionSortExprs = map[string]string{ //nolint:gosec // G101 false positive: "total_tokens" sort key, not a credential
"timestamp": "MAX(timestamp)",
"started_at": "MIN(timestamp)",
"cost_usd": "SUM(cost_usd)",
"cost_usd": "SUM" + CostUSDSQLExpr,
"total_tokens": "SUM(total_tokens)",
"duration": "SUM(duration)",
"request_count": "COUNT(*)",

View File

@@ -0,0 +1,124 @@
package types
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// costRow builds an access-log entry carrying only a cost breakdown — the rest
// of the row is irrelevant to the summation identities under test.
func costRow(id, session string, ts time.Time, in, cachedIn, cacheCreate, out float64) *AgentNetworkAccessLog {
return &AgentNetworkAccessLog{
ID: id,
SessionID: session,
Timestamp: ts,
InputCostUSD: in,
CachedInputCostUSD: cachedIn,
CacheCreationCostUSD: cacheCreate,
OutputCostUSD: out,
}
}
// TestAPIResponse_CostComponentsSumToAggregates is the contract a client adding
// up an API response depends on: within a single rendered object, the four
// per-bucket fields sum to cost_usd, and the two cache fields sum to
// cache_cost_usd. Uses rates that are not exactly representable in binary
// floating point, so the identity is checked against real arithmetic rather
// than round numbers.
func TestAPIResponse_CostComponentsSumToAggregates(t *testing.T) {
row := costRow("r1", "s1", time.Now(), 0.000768, 0.0002304, 0.00192, 0.003)
api := row.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "rendered buckets must sum to the rendered cost_usd")
assert.InDelta(t, api.CachedInputCostUsd+api.CacheCreationCostUsd, api.CacheCostUsd, 1e-12,
"rendered cache buckets must sum to the rendered cache_cost_usd")
assert.InDelta(t, 0.0059184, api.CostUsd, 1e-12, "total is the exact sum, not a separately rounded figure")
assert.InDelta(t, 0.0021504, api.CacheCostUsd, 1e-12, "cache cost is the exact sum of the two cache buckets")
}
// TestSessionSummary_SumsMatchSummedEntries proves a session summary equals the
// sum of the entries it renders: a client that adds up the entries itself must
// land on the same number the summary reports, per bucket and in total.
func TestSessionSummary_SumsMatchSummedEntries(t *testing.T) {
base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
entries := []*AgentNetworkAccessLog{
costRow("r1", "s1", base, 0.000768, 0.0002304, 0.00192, 0.003),
costRow("r2", "s1", base.Add(time.Minute), 0.000625, 0.0009375, 0, 0.005),
costRow("r3", "s1", base.Add(2*time.Minute), 0.0000016, 0, 0, 0.0000032),
}
sessions := FoldAccessLogSessions([]string{"s1"}, entries)
require.Len(t, sessions, 1)
sess := sessions[0].ToAPIResponse()
var wantInput, wantCachedInput, wantCacheCreation, wantOutput float64
for _, e := range entries {
wantInput += e.InputCostUSD
wantCachedInput += e.CachedInputCostUSD
wantCacheCreation += e.CacheCreationCostUSD
wantOutput += e.OutputCostUSD
}
assert.InDelta(t, wantInput, sess.InputCostUsd, 1e-12, "session input cost is the sum of its entries")
assert.InDelta(t, wantCachedInput, sess.CachedInputCostUsd, 1e-12, "session cache-read cost is the sum of its entries")
assert.InDelta(t, wantCacheCreation, sess.CacheCreationCostUsd, 1e-12, "session cache-write cost is the sum of its entries")
assert.InDelta(t, wantOutput, sess.OutputCostUsd, 1e-12, "session output cost is the sum of its entries")
assert.InDelta(t, wantInput+wantCachedInput+wantCacheCreation+wantOutput, sess.CostUsd, 1e-12,
"session total equals the summed entry buckets")
// Summing the rendered entries must give the same answer as reading the
// summary — the property a UI relies on when it totals a table itself.
var fromEntries float64
for _, e := range sess.Entries {
fromEntries += e.CostUsd
}
assert.InDelta(t, sess.CostUsd, fromEntries, 1e-12, "summary total must match the summed rendered entries")
// The sub-microdollar row must still contribute; it would vanish under
// 6-decimal quantisation.
assert.Greater(t, sess.InputCostUsd, 0.001393, "small-cost rows must not be quantised away")
}
// TestUsageBuckets_SumsMatchSummedRows proves the same identity one level up:
// a usage bucket equals the sum of the ledger rows folded into it, and the
// buckets together equal the whole range.
func TestUsageBuckets_SumsMatchSummedRows(t *testing.T) {
day1 := time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC)
day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
rows := []*AgentNetworkUsage{
{ID: "u1", Timestamp: day1, InputCostUSD: 0.000768, CachedInputCostUSD: 0.0002304, CacheCreationCostUSD: 0.00192, OutputCostUSD: 0.003},
{ID: "u2", Timestamp: day1.Add(time.Hour), InputCostUSD: 0.000625, CachedInputCostUSD: 0.0009375, OutputCostUSD: 0.005},
{ID: "u3", Timestamp: day2, InputCostUSD: 0.0000016, OutputCostUSD: 0.0000032},
}
buckets := AggregateUsageByGranularity(rows, UsageGranularityDay)
require.Len(t, buckets, 2, "two distinct days expected")
var total, cache float64
for _, b := range buckets {
api := b.ToAPIResponse()
assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
api.CostUsd, 1e-12, "each bucket's components must sum to its cost_usd")
total += api.CostUsd
cache += api.CacheCostUsd
}
var wantTotal, wantCache float64
for _, r := range rows {
wantTotal += r.TotalCostUSD()
wantCache += r.CacheCostUSD()
}
assert.InDelta(t, wantTotal, total, 1e-12, "buckets must sum to the total across all ledger rows")
assert.InDelta(t, wantCache, cache, 1e-12, "buckets must sum to the cache spend across all ledger rows")
// A month bucket over the same rows must total identically — regrouping
// changes the partition, never the sum.
monthly := AggregateUsageByGranularity(rows, UsageGranularityMonth)
require.Len(t, monthly, 1)
assert.InDelta(t, wantTotal, monthly[0].ToAPIResponse().CostUsd, 1e-12,
"re-bucketing at a different granularity must preserve the total")
}

View File

@@ -25,8 +25,19 @@ type AgentNetworkUsage struct {
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
CreatedAt time.Time
// Prompt-cache buckets: read + write token counts.
CachedInputTokens int64
CacheCreationTokens int64
// Per-bucket cost breakdown, mirroring AgentNetworkAccessLog — the only
// cost state stored; total and cache portion are derived on read. Kept on
// the usage ledger too so spend can be attributed per bucket even for
// accounts with log collection turned off. See AgentNetworkAccessLog for
// why the columns carry a zero default.
InputCostUSD float64 `gorm:"not null;default:0"`
CachedInputCostUSD float64 `gorm:"not null;default:0"`
CacheCreationCostUSD float64 `gorm:"not null;default:0"`
OutputCostUSD float64 `gorm:"not null;default:0"`
CreatedAt time.Time
}
// TableName keeps usage records in their own stripped table. Named
@@ -34,6 +45,17 @@ type AgentNetworkUsage struct {
// agent_network_usage table in a shared database.
func (AgentNetworkUsage) TableName() string { return "agent_network_request_usage" }
// TotalCostUSD is the request's total cost: the sum of the four per-bucket
// costs. Derived rather than stored so it cannot disagree with the breakdown.
func (u *AgentNetworkUsage) TotalCostUSD() float64 {
return u.InputCostUSD + u.CachedInputCostUSD + u.CacheCreationCostUSD + u.OutputCostUSD
}
// CacheCostUSD is the portion of the total billed for prompt-cache buckets.
func (u *AgentNetworkUsage) CacheCostUSD() float64 {
return u.CachedInputCostUSD + u.CacheCreationCostUSD
}
// AgentNetworkUsageGroup is the normalised many-to-many row linking a usage
// record to one authorising group, mirroring AgentNetworkAccessLogGroup so the
// usage overview can filter by group with a `group_id IN (...)` join.

View File

@@ -33,21 +33,45 @@ func ParseUsageGranularity(s string) UsageGranularity {
// AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
// the UTC start of the bucket as YYYY-MM-DD.
type AgentNetworkUsageBucket struct {
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CostUSD float64
PeriodStart string
InputTokens int64
OutputTokens int64
TotalTokens int64
CachedInputTokens int64
CacheCreationTokens int64
InputCostUSD float64
CachedInputCostUSD float64
CacheCreationCostUSD float64
OutputCostUSD float64
}
// TotalCostUSD is the bucket's total spend: the sum of the four per-bucket
// costs. Derived rather than accumulated separately so it cannot disagree with
// the components.
func (b *AgentNetworkUsageBucket) TotalCostUSD() float64 {
return b.InputCostUSD + b.CachedInputCostUSD + b.CacheCreationCostUSD + b.OutputCostUSD
}
// CacheCostUSD is the bucket's prompt-cache spend: cache reads plus writes.
func (b *AgentNetworkUsageBucket) CacheCostUSD() float64 {
return b.CachedInputCostUSD + b.CacheCreationCostUSD
}
// ToAPIResponse renders the bucket as the API representation.
func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
return api.AgentNetworkUsageBucket{
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CostUsd: b.CostUSD,
PeriodStart: b.PeriodStart,
InputTokens: b.InputTokens,
OutputTokens: b.OutputTokens,
TotalTokens: b.TotalTokens,
CachedInputTokens: b.CachedInputTokens,
CacheCreationTokens: b.CacheCreationTokens,
InputCostUsd: b.InputCostUSD,
CachedInputCostUsd: b.CachedInputCostUSD,
CacheCreationCostUsd: b.CacheCreationCostUSD,
OutputCostUsd: b.OutputCostUSD,
CostUsd: b.TotalCostUSD(),
CacheCostUsd: b.CacheCostUSD(),
}
}
@@ -84,7 +108,12 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
b.InputTokens += r.InputTokens
b.OutputTokens += r.OutputTokens
b.TotalTokens += r.TotalTokens
b.CostUSD += r.CostUSD
b.CachedInputTokens += r.CachedInputTokens
b.CacheCreationTokens += r.CacheCreationTokens
b.InputCostUSD += r.InputCostUSD
b.CachedInputCostUSD += r.CachedInputCostUSD
b.CacheCreationCostUSD += r.CacheCreationCostUSD
b.OutputCostUSD += r.OutputCostUSD
}
out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))

View File

@@ -5,9 +5,6 @@ import (
"strconv"
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/networkmap"
@@ -166,7 +163,7 @@ func (e *componentEncoder) indexAllPeers() {
}
}
func (e *componentEncoder) appendPeer(p *nbpeer.Peer) uint32 {
func (e *componentEncoder) appendPeer(p *types.ComponentPeer) uint32 {
if idx, ok := e.peerOrder[p.ID]; ok {
return idx
}
@@ -180,7 +177,7 @@ func (e *componentEncoder) appendPeer(p *nbpeer.Peer) uint32 {
// (c.RouterPeers may contain peers not in c.Peers when validation rules drop
// them) and returns their wire indexes for the RouterPeerIndexes field. Must
// run before any encoder that resolves peer ids via e.peerOrder.
func (e *componentEncoder) indexRouterPeers(routers map[string]*nbpeer.Peer) []uint32 {
func (e *componentEncoder) indexRouterPeers(routers map[string]*types.ComponentPeer) []uint32 {
if len(routers) == 0 {
return nil
}
@@ -514,7 +511,7 @@ func encodeCustomZones(zones []nbdns.CustomZone) []*proto.CustomZone {
return out
}
func (e *componentEncoder) encodeNetworkResources(resources []*resourceTypes.NetworkResource) []*proto.NetworkResourceRaw {
func (e *componentEncoder) encodeNetworkResources(resources []*types.ComponentResource) []*proto.NetworkResourceRaw {
if len(resources) == 0 {
return nil
}
@@ -543,7 +540,7 @@ func (e *componentEncoder) encodeNetworkResources(resources []*resourceTypes.Net
return out
}
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*routerTypes.NetworkRouter) map[string]*proto.NetworkRouterList {
func (e *componentEncoder) encodeRoutersMap(routersMap map[string]map[string]*types.ComponentRouter) map[string]*proto.NetworkRouterList {
if len(routersMap) == 0 {
return nil
}
@@ -692,20 +689,20 @@ func toAccountNetwork(n *types.Network) *proto.AccountNetwork {
return out
}
func toPeerCompact(p *nbpeer.Peer) *proto.PeerCompact {
func toPeerCompact(p *types.ComponentPeer) *proto.PeerCompact {
pc := &proto.PeerCompact{
WgPubKey: decodeWgKey(p.Key),
SshPubKey: []byte(p.SSHKey),
DnsLabel: p.DNSLabel,
AgentVersion: p.Meta.WtVersion,
AddedWithSsoLogin: p.UserID != "",
AgentVersion: p.AgentVersion,
AddedWithSsoLogin: p.AddedWithSSOLogin,
LoginExpirationEnabled: p.LoginExpirationEnabled,
SshEnabled: p.SSHEnabled,
SupportsIpv6: p.SupportsIPv6(),
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
ServerSshAllowed: p.Meta.Flags.ServerSSHAllowed,
SupportsIpv6: p.SupportsIPv6,
SupportsSourcePrefixes: p.SupportsSourcePrefixes,
ServerSshAllowed: p.ServerSSHAllowed,
}
if p.LastLogin != nil {
if !p.LastLogin.IsZero() {
pc.LastLoginUnixNano = p.LastLogin.UnixNano()
}
switch {

View File

@@ -15,9 +15,6 @@ import (
goproto "google.golang.org/protobuf/proto"
nbdns "github.com/netbirdio/netbird/dns"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
nbroute "github.com/netbirdio/netbird/route"
"github.com/netbirdio/netbird/shared/management/proto"
@@ -155,29 +152,28 @@ func envelopesEquivalent(a, b *proto.NetworkMapEnvelope) bool {
}
func newTestComponents() *types.NetworkMapComponents {
peerA := &nbpeer.Peer{
ID: "peer-a",
Key: testWgKeyA,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peera",
SSHKey: "ssh-a",
Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now()},
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
peerA := &types.ComponentPeer{
ID: "peer-a",
Key: testWgKeyA,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 1}),
DNSLabel: "peera",
SSHKey: "ssh-a",
AgentVersion: "0.40.0",
}
peerB := &nbpeer.Peer{
ID: "peer-b",
Key: testWgKeyB,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
DNSLabel: "peerb",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.25.0"},
peerB := &types.ComponentPeer{
ID: "peer-b",
Key: testWgKeyB,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 2}),
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2}),
DNSLabel: "peerb",
AgentVersion: "0.25.0",
}
peerC := &nbpeer.Peer{
ID: "peer-c",
Key: testWgKeyC,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
peerC := &types.ComponentPeer{
ID: "peer-c",
Key: testWgKeyC,
IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc",
AgentVersion: "0.40.0",
}
return &types.NetworkMapComponents{
@@ -191,12 +187,12 @@ func newTestComponents() *types.NetworkMapComponents {
PeerLoginExpirationEnabled: true,
PeerLoginExpiration: 2 * time.Hour,
},
Peers: map[string]*nbpeer.Peer{
Peers: map[string]*types.ComponentPeer{
"peer-a": peerA,
"peer-b": peerB,
"peer-c": peerC,
},
Groups: map[string]*types.Group{
Groups: map[string]*types.ComponentGroup{
"group-src": {ID: "group-src", PublicID: "1", Name: "Src", Peers: []string{"peer-a"}},
"group-dst": {ID: "group-dst", PublicID: "2", Name: "Dst", Peers: []string{"peer-b", "peer-c"}},
},
@@ -215,7 +211,7 @@ func newTestComponents() *types.NetworkMapComponents {
}},
},
},
RouterPeers: map[string]*nbpeer.Peer{"peer-c": peerC},
RouterPeers: map[string]*types.ComponentPeer{"peer-c": peerC},
}
}
@@ -381,12 +377,12 @@ func TestEncodeNetworkMapEnvelope_MalformedWgKey(t *testing.T) {
func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
c := newTestComponents()
v6Only := &nbpeer.Peer{
ID: "peer-v6",
Key: testWgKeyA,
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
DNSLabel: "peerv6",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
v6Only := &types.ComponentPeer{
ID: "peer-v6",
Key: testWgKeyA,
IPv6: netip.AddrFrom16([16]byte{0xfd, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9}),
DNSLabel: "peerv6",
AgentVersion: "0.40.0",
}
c.Peers["peer-v6"] = v6Only
@@ -405,11 +401,11 @@ func TestEncodeNetworkMapEnvelope_IPv6OnlyPeer(t *testing.T) {
func TestEncodeNetworkMapEnvelope_PeerWithoutIP(t *testing.T) {
c := newTestComponents()
c.Peers["peer-noip"] = &nbpeer.Peer{
ID: "peer-noip",
Key: testWgKeyA,
DNSLabel: "peernoip",
Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
c.Peers["peer-noip"] = &types.ComponentPeer{
ID: "peer-noip",
Key: testWgKeyA,
DNSLabel: "peernoip",
AgentVersion: "0.40.0",
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -444,9 +440,9 @@ func TestEncodeNetworkMapEnvelope_EmptyInput(t *testing.T) {
func TestEncodeNetworkMapEnvelope_PeerLoginExpirationFields(t *testing.T) {
c := newTestComponents()
now := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
c.Peers["peer-a"].UserID = "user-1"
c.Peers["peer-a"].AddedWithSSOLogin = true
c.Peers["peer-a"].LoginExpirationEnabled = true
c.Peers["peer-a"].LastLogin = &now
c.Peers["peer-a"].LastLogin = now
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -557,7 +553,7 @@ func TestEncodeNetworkMapEnvelope_ResourceOnlyPolicyShippedAndIndexed(t *testing
}
// Resource must appear in components.NetworkResources with a seq id —
// encoder uses that to translate the xid map key to uint32.
c.NetworkResources = []*resourceTypes.NetworkResource{
c.NetworkResources = []*types.ComponentResource{
{ID: "resource-x", PublicID: "77", Name: "res-x", Enabled: true},
}
@@ -625,11 +621,11 @@ func TestEncodeNetworkMapEnvelope_PostureFailedPeers(t *testing.T) {
func TestEncodeNetworkMapEnvelope_RoutersMap(t *testing.T) {
c := newTestComponents()
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*routerTypes.NetworkRouter{
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
"net-1": {
"peer-c": {
ID: "router-1", PublicID: "200",
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
PublicID: "200",
Peer: "peer-c", Masquerade: true, Metric: 10, Enabled: true,
},
},
}
@@ -655,14 +651,14 @@ func TestEncodeNetworkMapEnvelope_RouterPeerNotInComponentsPeers(t *testing.T) {
// peer_index reference must still resolve.
c := newTestComponents()
delete(c.Peers, "peer-c")
routerPeer := &nbpeer.Peer{
routerPeer := &types.ComponentPeer{
ID: "peer-c", Key: testWgKeyC, IP: netip.AddrFrom4([4]byte{100, 64, 0, 3}),
DNSLabel: "peerc", Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
DNSLabel: "peerc", AgentVersion: "0.40.0",
}
c.RouterPeers = map[string]*nbpeer.Peer{"peer-c": routerPeer}
c.RouterPeers = map[string]*types.ComponentPeer{"peer-c": routerPeer}
c.NetworkXIDToPublicID = map[string]string{"net-1": "5"}
c.RoutersMap = map[string]map[string]*routerTypes.NetworkRouter{
"net-1": {"peer-c": {ID: "r-1", PublicID: "1", Peer: "peer-c", Enabled: true}},
c.RoutersMap = map[string]map[string]*types.ComponentRouter{
"net-1": {"peer-c": {PublicID: "1", Peer: "peer-c", Enabled: true}},
}
full := EncodeNetworkMapEnvelope(ComponentsEnvelopeInput{Components: c}).GetFull()
@@ -695,9 +691,9 @@ func TestToProxyPatch_EmptyInputReturnsNil(t *testing.T) {
func TestToProxyPatch_PopulatesAllFields(t *testing.T) {
nm := &types.NetworkMap{
Peers: []*nbpeer.Peer{{
Peers: []*types.ComponentPeer{{
ID: "ext-peer", Key: testWgKeyA, IP: netip.AddrFrom4([4]byte{100, 64, 0, 9}),
DNSLabel: "extpeer", Meta: nbpeer.PeerSystemMeta{WtVersion: "0.40.0"},
DNSLabel: "extpeer", AgentVersion: "0.40.0",
}},
FirewallRules: []*types.FirewallRule{{
PeerIP: "100.64.0.9", Action: "accept", Direction: 0, Protocol: "tcp",
@@ -780,6 +776,6 @@ func TestEncodeNetworkMapEnvelope_AccountSettingsAlwaysEmitted(t *testing.T) {
func emptyNetworkMapComponents() *types.NetworkMapComponents {
return types.EmptyNetworkMapComponents(
&types.NetworkMapComponents{
PeerID: "peer-id", Peers: map[string]*nbpeer.Peer{"peer-id": {}}},
PeerID: "peer-id", Peers: map[string]*types.ComponentPeer{"peer-id": {}}},
)
}

View File

@@ -19,10 +19,10 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
mkComponents := func(rule *types.PolicyRule, sshEnabled bool) (*types.NetworkMapComponents, *nbpeer.Peer) {
peer := &nbpeer.Peer{ID: targetPeerID, SSHEnabled: sshEnabled}
group := &types.Group{ID: targetGroupID, Name: "dst", Peers: []string{targetPeerID}}
group := &types.ComponentGroup{ID: targetGroupID, Name: "dst", Peers: []string{targetPeerID}}
return &types.NetworkMapComponents{
Peers: map[string]*nbpeer.Peer{targetPeerID: peer},
Groups: map[string]*types.Group{targetGroupID: group},
Peers: map[string]*types.ComponentPeer{targetPeerID: peer.ToComponent()},
Groups: map[string]*types.ComponentGroup{targetGroupID: group},
Policies: []*types.Policy{{
ID: "p",
Enabled: true,
@@ -158,8 +158,8 @@ func TestComputeSSHEnabledForPeer(t *testing.T) {
func TestComputeSSHEnabledForPeer_TargetMissingFromComponents(t *testing.T) {
peer := &nbpeer.Peer{ID: "missing", SSHEnabled: true}
c := &types.NetworkMapComponents{
Peers: map[string]*nbpeer.Peer{}, // target peer NOT present
Groups: map[string]*types.Group{
Peers: map[string]*types.ComponentPeer{}, // target peer NOT present
Groups: map[string]*types.ComponentGroup{
"g": {ID: "g", Peers: []string{"missing"}},
},
Policies: []*types.Policy{{

View File

@@ -285,6 +285,7 @@ func (s *ProxyServiceServer) CheckLLMPolicyLimits(ctx context.Context, req *prot
UserID: req.GetUserId(),
GroupIDs: req.GetGroupIds(),
ProviderID: req.GetProviderId(),
Model: req.GetModel(),
})
if err != nil {
log.WithContext(ctx).Errorf("select policy for request: %v", err)

View File

@@ -0,0 +1,138 @@
package grpc
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/codes"
grpcstatus "google.golang.org/grpc/status"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/shared/management/proto"
)
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
// and returns a pre-programmed result, so tests can assert what the handler
// forwards to the selector.
type fakeAgentNetworkLimits struct {
gotInput agentnetwork.PolicySelectionInput
result *agentnetwork.PolicySelectionResult
err error
}
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
f.gotInput = in
if f.err != nil {
return nil, f.err
}
return f.result, nil
}
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
return nil
}
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
// the model the proxy extracted must reach the selector's Model unchanged,
// alongside the account/user/group/provider fields.
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
UserId: "user-1",
GroupIds: []string{"grp-a", "grp-b"},
ProviderId: "prov-1",
Model: "claude-opus-4",
}
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
assert.Equal(t, "user-1", fake.gotInput.UserID)
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
"the request's model must be forwarded to the selector")
}
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
// model (empty string) is forwarded as-is; the selector decides how to treat it.
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
}
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
}
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
// envelope surfaces the model-allowlist deny code + reason through the response.
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
Allow: false,
DenyCode: "llm_policy.model_blocked",
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "deny", resp.Decision)
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
assert.NotEmpty(t, resp.DenyReason)
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
}
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
// failure surfaces as an Internal gRPC error rather than a silent allow.
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
}
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
// fallback: with no limits service wired the RPC returns Unimplemented.
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
s := &ProxyServiceServer{}
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Unimplemented, st.Code())
}

View File

@@ -6,6 +6,7 @@ import (
"github.com/netbirdio/netbird/management/server/account"
"github.com/netbirdio/netbird/management/server/activity"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
"github.com/netbirdio/netbird/management/server/permissions"
"github.com/netbirdio/netbird/management/server/permissions/modules"
"github.com/netbirdio/netbird/management/server/permissions/operations"
@@ -30,6 +31,10 @@ type managerImpl struct {
accountManager account.Manager
}
func eventMetaResource(group *types.Group, resource *resourceTypes.NetworkResource) map[string]any {
return map[string]any{"name": group.Name, "id": group.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
type mockManager struct {
}
@@ -109,7 +114,7 @@ func (m *managerImpl) AddResourceToGroupInTransaction(ctx context.Context, trans
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, group.EventMetaResource(networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceAddedToGroup, eventMetaResource(group, networkResource))
}
return event, nil
@@ -133,7 +138,7 @@ func (m *managerImpl) RemoveResourceFromGroupInTransaction(ctx context.Context,
}
event := func() {
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, group.EventMetaResource(networkResource))
m.accountManager.StoreEvent(ctx, userID, groupID, accountID, activity.ResourceRemovedFromGroup, eventMetaResource(group, networkResource))
}
return event, nil

View File

@@ -446,7 +446,7 @@ func (h *Handler) GetAccessiblePeers(w http.ResponseWriter, r *http.Request) {
netMap := account.GetPeerNetworkMapFromComponents(ctx, peerID, dns.CustomZone{}, nil, validPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil, account.GetActiveGroupUsers())
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, dnsDomain))
util.WriteJSONObject(ctx, w, toAccessiblePeers(netMap, account.Peers, dnsDomain))
}
func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request) {
@@ -534,15 +534,20 @@ func (h *Handler) CreateTemporaryAccess(w http.ResponseWriter, r *http.Request)
util.WriteJSONObject(r.Context(), w, resp)
}
func toAccessiblePeers(netMap *types.NetworkMap, dnsDomain string) []api.AccessiblePeer {
// toAccessiblePeers rehydrates the calculated map's component peers into the
// account's full peer objects, which carry the location/status/meta fields
// the API response needs.
func toAccessiblePeers(netMap *types.NetworkMap, accountPeers map[string]*nbpeer.Peer, dnsDomain string) []api.AccessiblePeer {
accessiblePeers := make([]api.AccessiblePeer, 0, len(netMap.Peers)+len(netMap.OfflinePeers))
for _, p := range netMap.Peers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
}
for _, p := range netMap.OfflinePeers {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(p, dnsDomain))
add := func(peers []*types.ComponentPeer) {
for _, p := range peers {
if peer := accountPeers[p.ID]; peer != nil {
accessiblePeers = append(accessiblePeers, peerToAccessiblePeer(peer, dnsDomain))
}
}
}
add(netMap.Peers)
add(netMap.OfflinePeers)
return accessiblePeers
}

View File

@@ -683,3 +683,81 @@ func BackfillPublicIDs[T any](ctx context.Context, db *gorm.DB) error {
log.WithContext(ctx).Infof("Backfill of empty public_id in table %s completed", tableName)
return nil
}
// FoldCostAggregatesIntoBuckets migrates a per-request cost table from the old
// "stored aggregate" shape (cost_usd + cache_cost_usd columns) to the per-bucket
// breakdown, where the total and cache portion are derived on read instead.
//
// The fold preserves both aggregates exactly for historical rows: the cache
// total moves into cached_input_cost_usd and the remainder into
// input_cost_usd, so a row's derived total and cache cost still match what it
// reported before the upgrade. The finer split is genuinely unknown for those
// rows — the old schema never recorded a read/write or input/output division —
// so it is lumped rather than guessed; only rows written after the upgrade
// carry a true four-way split.
//
// Dropping the columns before folding would zero every historical row's cost,
// so the update runs first and the drop only happens once it succeeds. A table
// with no cost_usd column has already been migrated (or was created fresh) and
// is skipped.
func FoldCostAggregatesIntoBuckets[T any](ctx context.Context, db *gorm.DB) error {
var model T
if !db.Migrator().HasTable(&model) {
log.WithContext(ctx).Debugf("table for %T does not exist, no cost-bucket migration needed", model)
return nil
}
if !db.Migrator().HasColumn(&model, "cost_usd") {
log.WithContext(ctx).Debugf("table for %T has no cost_usd column, cost buckets already migrated", model)
return nil
}
stmt := &gorm.Statement{DB: db}
if err := stmt.Parse(&model); err != nil {
return fmt.Errorf("parse model schema: %w", err)
}
tableName := stmt.Schema.Table
// COALESCE guards rows whose new columns were added as NULL by an earlier
// AutoMigrate run that predates the NOT NULL default.
hasCacheColumn := db.Migrator().HasColumn(&model, "cache_cost_usd")
cacheExpr := "0"
if hasCacheColumn {
cacheExpr = "COALESCE(cache_cost_usd, 0)"
}
if err := db.Transaction(func(tx *gorm.DB) error {
// Only touch rows that carry a legacy total and no breakdown yet, so
// the migration is idempotent and never overwrites a true split.
update := fmt.Sprintf(`UPDATE %s
SET input_cost_usd = COALESCE(cost_usd, 0) - %s,
cached_input_cost_usd = %s,
cache_creation_cost_usd = 0,
output_cost_usd = 0
WHERE COALESCE(cost_usd, 0) <> 0
AND COALESCE(input_cost_usd, 0) = 0
AND COALESCE(cached_input_cost_usd, 0) = 0
AND COALESCE(cache_creation_cost_usd, 0) = 0
AND COALESCE(output_cost_usd, 0) = 0`, tableName, cacheExpr, cacheExpr)
res := tx.Exec(update)
if res.Error != nil {
return fmt.Errorf("fold legacy cost aggregates in %s: %w", tableName, res.Error)
}
log.WithContext(ctx).Infof("folded legacy cost aggregates into per-bucket columns for %d rows in table %s", res.RowsAffected, tableName)
if err := tx.Migrator().DropColumn(&model, "cost_usd"); err != nil {
return fmt.Errorf("drop cost_usd from %s: %w", tableName, err)
}
if hasCacheColumn {
if err := tx.Migrator().DropColumn(&model, "cache_cost_usd"); err != nil {
return fmt.Errorf("drop cache_cost_usd from %s: %w", tableName, err)
}
}
return nil
}); err != nil {
return err
}
log.WithContext(ctx).Infof("migration of stored cost aggregates to per-bucket columns in table %s completed", tableName)
return nil
}

View File

@@ -16,6 +16,7 @@ import (
"gorm.io/driver/sqlite"
"gorm.io/gorm"
agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
"github.com/netbirdio/netbird/management/server/migration"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/testutil"
@@ -639,3 +640,99 @@ func TestCleanupOrphanedResources_SkipsWhenForeignKeyExists(t *testing.T) {
db.Model(&testChildWithFK{}).Count(&count)
assert.Equal(t, int64(2), count, "Both rows should survive — migration must skip when FK constraint exists")
}
// legacyCostRow is the pre-breakdown shape of the usage table: cost was stored
// as a total plus a cache portion, with no per-bucket columns. Used to build a
// realistic pre-upgrade table for the fold migration to run against.
type legacyCostRow struct {
ID string `gorm:"primaryKey"`
AccountID string
Model string
CostUSD float64
CacheCostUSD float64
}
func (legacyCostRow) TableName() string { return "agent_network_request_usage" }
// TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost covers the upgrade
// path: a table written under the old schema must come out with its per-row
// total and cache cost unchanged, because dropping cost_usd without folding it
// forward would silently zero every historical row's spend.
func TestFoldCostAggregatesIntoBuckets_PreservesHistoricalCost(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
// setupDatabase hands back a process-shared database, so start from a clean
// table rather than inheriting rows from another test.
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&legacyCostRow{}), "legacy table must be created")
require.NoError(t, db.Create(&legacyCostRow{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6", CostUSD: 0.0123, CacheCostUSD: 0.0029,
}).Error)
require.NoError(t, db.Create(&legacyCostRow{
ID: "u2", AccountID: "acct-1", Model: "gpt-4o", CostUSD: 0.5, CacheCostUSD: 0,
}).Error)
// A zero-cost row (denied / unpriced request) must stay zero, not be touched.
require.NoError(t, db.Create(&legacyCostRow{ID: "u3", AccountID: "acct-1", Model: "gw/unpriced"}).Error)
// AutoMigrate adds the per-bucket columns alongside the legacy ones, exactly
// as a real upgrade does before the post-auto migrations run.
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}), "new columns must be added")
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db))
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cost_usd"),
"legacy cost_usd column must be dropped once folded")
assert.False(t, db.Migrator().HasColumn(&agentNetworkTypes.AgentNetworkUsage{}, "cache_cost_usd"),
"legacy cache_cost_usd column must be dropped once folded")
var rows []*agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.Order("id").Find(&rows).Error)
require.Len(t, rows, 3)
// u1: total and cache portion both preserved; the read/write and
// input/output splits are unknowable for a legacy row, so the cache total
// lands on cached_input and the remainder on input.
assert.InDelta(t, 0.0123, rows[0].TotalCostUSD(), 1e-9, "historical total must survive the fold")
assert.InDelta(t, 0.0029, rows[0].CacheCostUSD(), 1e-9, "historical cache cost must survive the fold")
assert.InDelta(t, 0.0094, rows[0].InputCostUSD, 1e-9, "non-cache remainder lands on input")
assert.InDelta(t, 0.0029, rows[0].CachedInputCostUSD, 1e-9, "legacy cache total lands on cached input")
assert.Zero(t, rows[0].CacheCreationCostUSD, "legacy rows carry no read/write split to recover")
assert.Zero(t, rows[0].OutputCostUSD, "legacy rows carry no input/output split to recover")
// u2: no cache spend — the whole total is the non-cache remainder.
assert.InDelta(t, 0.5, rows[1].TotalCostUSD(), 1e-9, "cache-free historical total must survive")
assert.Zero(t, rows[1].CacheCostUSD(), "a cache-free row must stay cache-free")
// u3: zero stays zero rather than being rewritten.
assert.Zero(t, rows[2].TotalCostUSD(), "an unpriced row must remain unpriced")
}
// TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated proves the migration is
// safe to re-run: with no legacy column present it is a no-op that leaves a
// true four-way split untouched.
func TestFoldCostAggregatesIntoBuckets_SkipsAlreadyMigrated(t *testing.T) {
ctx := context.Background()
db := setupDatabase(t)
require.NoError(t, db.Migrator().DropTable(&agentNetworkTypes.AgentNetworkUsage{}))
require.NoError(t, db.AutoMigrate(&agentNetworkTypes.AgentNetworkUsage{}))
// Timestamp must be set explicitly: a zero time.Time serialises as
// '0000-00-00 00:00:00', which MySQL rejects under strict mode.
require.NoError(t, db.Create(&agentNetworkTypes.AgentNetworkUsage{
ID: "u1", AccountID: "acct-1", Model: "claude-sonnet-4-6",
Timestamp: time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC),
InputCostUSD: 0.001, CachedInputCostUSD: 0.002, CacheCreationCostUSD: 0.003, OutputCostUSD: 0.004,
}).Error)
require.NoError(t, migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db),
"running against an already-migrated table must be a no-op, not an error")
var row agentNetworkTypes.AgentNetworkUsage
require.NoError(t, db.First(&row, "id = ?", "u1").Error)
assert.InDelta(t, 0.001, row.InputCostUSD, 1e-9, "a true split must not be rewritten")
assert.InDelta(t, 0.002, row.CachedInputCostUSD, 1e-9)
assert.InDelta(t, 0.003, row.CacheCreationCostUSD, 1e-9)
assert.InDelta(t, 0.004, row.OutputCostUSD, 1e-9)
assert.InDelta(t, 0.01, row.TotalCostUSD(), 1e-9, "derived total sums the four buckets")
}

View File

@@ -14,6 +14,7 @@ import (
nbDomain "github.com/netbirdio/netbird/shared/management/domain"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkResourceType string
@@ -64,6 +65,27 @@ func NewNetworkResource(accountID, networkID, name, description, address string,
}, nil
}
// ToComponent converts the resource to its self-contained components
// representation. Returns nil for a nil resource.
func (n *NetworkResource) ToComponent() *sharedTypes.ComponentResource {
if n == nil {
return nil
}
return &sharedTypes.ComponentResource{
ID: n.ID,
PublicID: n.PublicID,
NetworkID: n.NetworkID,
AccountID: n.AccountID,
Name: n.Name,
Description: n.Description,
Type: sharedTypes.ComponentResourceType(n.Type),
Address: n.Address,
Domain: n.Domain,
Prefix: n.Prefix,
Enabled: n.Enabled,
}
}
func (n *NetworkResource) ToAPIResponse(groups []api.GroupMinimum) *api.NetworkResource {
addr := n.Prefix.String()
if n.Type == Domain {

View File

@@ -7,6 +7,7 @@ import (
"github.com/netbirdio/netbird/management/server/networks/types"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
type NetworkRouter struct {
@@ -21,6 +22,36 @@ type NetworkRouter struct {
Enabled bool
}
// ToComponent converts the router to its self-contained components
// representation. Returns nil for a nil router.
func (n *NetworkRouter) ToComponent() *sharedTypes.ComponentRouter {
if n == nil {
return nil
}
return &sharedTypes.ComponentRouter{
NetworkID: n.NetworkID,
PublicID: n.PublicID,
Peer: n.Peer,
PeerGroups: n.PeerGroups,
Masquerade: n.Masquerade,
Metric: n.Metric,
Enabled: n.Enabled,
}
}
// ToComponentMap converts a peer-keyed router map to its components
// representation.
func ToComponentMap(routers map[string]*NetworkRouter) map[string]*sharedTypes.ComponentRouter {
if routers == nil {
return nil
}
out := make(map[string]*sharedTypes.ComponentRouter, len(routers))
for id, r := range routers {
out[id] = r.ToComponent()
}
return out
}
func NewNetworkRouter(accountID string, networkID string, peer string, peerGroups []string, masquerade bool, metric int, enabled bool) (*NetworkRouter, error) {
r := &NetworkRouter{
ID: xid.New().String(),

View File

@@ -405,7 +405,7 @@ func (am *DefaultAccountManager) CreatePeerJob(ctx context.Context, accountID, p
return status.NewPeerNotPartOfAccountError()
}
meetMinVer, err := posture.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
meetMinVer, err := version.MeetsMinVersion(remoteJobsMinVer, p.Meta.WtVersion)
if !version.IsDevelopmentVersion(p.Meta.WtVersion) && (!meetMinVer || err != nil) {
return status.Errorf(status.PreconditionFailed, "peer version %s does not meet the minimum required version %s for remote jobs", p.Meta.WtVersion, remoteJobsMinVer)
}
@@ -1588,7 +1588,7 @@ func affectedPeerIDsFromNetworkMap(nmap *types.NetworkMap, selfPeerID string) []
}
seen := make(map[string]struct{}, len(nmap.Peers)+len(nmap.OfflinePeers))
ids := make([]string, 0, len(nmap.Peers)+len(nmap.OfflinePeers))
add := func(peers []*nbpeer.Peer) {
add := func(peers []*types.ComponentPeer) {
for _, p := range peers {
if p == nil || p.ID == "" || p.ID == selfPeerID {
continue

View File

@@ -13,6 +13,7 @@ import (
"github.com/netbirdio/netbird/management/server/util"
"github.com/netbirdio/netbird/shared/management/http/api"
sharedTypes "github.com/netbirdio/netbird/shared/management/types"
)
// Peer capability constants mirror the proto enum values.
@@ -205,6 +206,35 @@ func (p *Peer) AddedWithSSOLogin() bool {
return p.UserID != ""
}
// ToComponent converts the peer to its self-contained components
// representation, carrying exactly the subset of peer data that crosses the
// components wire format. Returns nil for a nil peer so callers can convert
// possibly-missing peers without guarding.
func (p *Peer) ToComponent() *sharedTypes.ComponentPeer {
if p == nil {
return nil
}
cp := &sharedTypes.ComponentPeer{
ID: p.ID,
Key: p.Key,
IP: p.IP,
IPv6: p.IPv6,
DNSLabel: p.DNSLabel,
SSHKey: p.SSHKey,
SSHEnabled: p.SSHEnabled,
ServerSSHAllowed: p.Meta.Flags.ServerSSHAllowed,
AgentVersion: p.Meta.WtVersion,
SupportsSourcePrefixes: p.SupportsSourcePrefixes(),
SupportsIPv6: p.SupportsIPv6(),
LoginExpirationEnabled: p.LoginExpirationEnabled,
AddedWithSSOLogin: p.AddedWithSSOLogin(),
}
if p.LastLogin != nil {
cp.LastLogin = *p.LastLogin
}
return cp
}
// HasCapability reports whether the peer has the given capability.
func (p *Peer) HasCapability(capability int32) bool {
return slices.Contains(p.Meta.Capabilities, capability)

View File

@@ -1092,14 +1092,14 @@ func TestToSyncResponse(t *testing.T) {
}
networkMap := &types.NetworkMap{
Network: &types.Network{Net: *ipnet, Serial: 1000},
Peers: []*nbpeer.Peer{{
Peers: []*types.ComponentPeer{{
IP: netip.MustParseAddr("192.168.1.2"),
IPv6: netip.MustParseAddr("fd00::2"),
Key: "peer2-key",
DNSLabel: "peer2",
SSHEnabled: true,
SSHKey: "peer2-ssh-key"}},
OfflinePeers: []*nbpeer.Peer{{
OfflinePeers: []*types.ComponentPeer{{
IP: netip.MustParseAddr("192.168.1.3"),
IPv6: netip.MustParseAddr("fd00::3"),
Key: "peer3-key",

View File

@@ -3,11 +3,9 @@ package posture
import (
"context"
"fmt"
"strings"
"github.com/hashicorp/go-version"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbversion "github.com/netbirdio/netbird/version"
)
type NBVersionCheck struct {
@@ -16,14 +14,8 @@ type NBVersionCheck struct {
var _ Check = (*NBVersionCheck)(nil)
// sanitizeVersion removes anything after the pre-release tag (e.g., "-dev", "-alpha", etc.)
func sanitizeVersion(version string) string {
parts := strings.Split(version, "-")
return parts[0]
}
func (n *NBVersionCheck) Check(ctx context.Context, peer nbpeer.Peer) (bool, error) {
meetsMin, err := MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
meetsMin, err := nbversion.MeetsMinVersion(n.MinVersion, peer.Meta.WtVersion)
if err != nil {
return false, err
}
@@ -48,21 +40,3 @@ func (n *NBVersionCheck) Validate() error {
}
return nil
}
// MeetsMinVersion checks if the peer's version meets or exceeds the minimum required version
func MeetsMinVersion(minVer, peerVer string) (bool, error) {
peerVer = sanitizeVersion(peerVer)
minVer = sanitizeVersion(minVer)
peerNBVer, err := version.NewVersion(peerVer)
if err != nil {
return false, err
}
constraints, err := version.NewConstraint(">= " + minVer)
if err != nil {
return false, err
}
return constraints.Check(peerNBVer), nil
}

View File

@@ -139,68 +139,3 @@ func TestNBVersionCheck_Validate(t *testing.T) {
})
}
}
func TestMeetsMinVersion(t *testing.T) {
tests := []struct {
name string
minVer string
peerVer string
want bool
wantErr bool
}{
{
name: "Peer version greater than min version",
minVer: "0.26.0",
peerVer: "0.60.1",
want: true,
wantErr: false,
},
{
name: "Peer version equals min version",
minVer: "1.0.0",
peerVer: "1.0.0",
want: true,
wantErr: false,
},
{
name: "Peer version less than min version",
minVer: "1.0.0",
peerVer: "0.9.9",
want: false,
wantErr: false,
},
{
name: "Peer version with pre-release tag greater than min version",
minVer: "1.0.0",
peerVer: "1.0.1-alpha",
want: true,
wantErr: false,
},
{
name: "Invalid peer version format",
minVer: "1.0.0",
peerVer: "dev",
want: false,
wantErr: true,
},
{
name: "Invalid min version format",
minVer: "invalid.version",
peerVer: "1.0.0",
want: false,
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := MeetsMinVersion(tt.minVer, tt.peerVer)
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
assert.Equal(t, tt.want, got)
})
}
}

View File

@@ -71,7 +71,7 @@ func (s *SqlStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetr
usageRow := db.Model(&agentNetworkTypes.AgentNetworkUsage{}).
Select("COALESCE(SUM(input_tokens), 0) AS input_tokens, " +
"COALESCE(SUM(output_tokens), 0) AS output_tokens, " +
"COALESCE(SUM(cost_usd), 0) AS cost_usd").Row()
"COALESCE(SUM" + agentNetworkTypes.CostUSDSQLExpr + ", 0) AS cost_usd").Row()
if err := usageRow.Scan(&m.InputTokens, &m.OutputTokens, &m.CostUSD); err != nil {
return AgentNetworkMetrics{}, fmt.Errorf("scan agent network usage metrics: %w", err)
}

View File

@@ -37,7 +37,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
InputCostUSD: 0.0231,
}
usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -71,7 +71,7 @@ func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
InputTokens: 1200,
OutputTokens: 640,
TotalTokens: 1840,
CostUSD: 0.0231,
InputCostUSD: 0.0231,
}
entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
@@ -127,7 +127,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
mk := func(id string, ts time.Time, model string, in, out int64, cost float64) *agentNetworkTypes.AgentNetworkUsage {
return &agentNetworkTypes.AgentNetworkUsage{
ID: id, AccountID: accountID, Timestamp: ts, Model: model,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, CostUSD: cost,
InputTokens: in, OutputTokens: out, TotalTokens: in + out, InputCostUSD: cost,
}
}
require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
@@ -143,7 +143,7 @@ func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
assert.Equal(t, "2026-05-05", buckets[0].PeriodStart, "oldest-first ordering")
assert.Equal(t, int64(300), buckets[0].InputTokens, "same-day input tokens summed")
assert.Equal(t, int64(130), buckets[0].OutputTokens)
assert.InDelta(t, 0.30, buckets[0].CostUSD, 1e-9, "same-day cost summed")
assert.InDelta(t, 0.30, buckets[0].TotalCostUSD(), 1e-9, "same-day cost summed")
assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
assert.Equal(t, int64(15), buckets[1].TotalTokens)
@@ -174,7 +174,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts,
UserID: user, StatusCode: 200, Provider: provider, Model: model,
SessionID: session, Decision: decision,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, CostUSD: cost,
InputTokens: 100, OutputTokens: 50, TotalTokens: 150, InputCostUSD: cost,
}
}
@@ -207,7 +207,7 @@ func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
s1 := sessions[2]
assert.Equal(t, 2, s1.RequestCount, "s1 has two requests")
assert.Equal(t, int64(300), s1.TotalTokens, "tokens summed across the session")
assert.InDelta(t, 0.30, s1.CostUSD, 1e-9, "cost summed across the session")
assert.InDelta(t, 0.30, s1.TotalCostUSD(), 1e-9, "cost summed across the session")
assert.Equal(t, "alice", s1.UserID)
assert.Equal(t, "allow", s1.Decision)
// SQLite hands times back in time.Local; normalise to UTC so the instant is

View File

@@ -650,6 +650,14 @@ func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
func(db *gorm.DB) error {
return migration.DropIndex[proxy.Proxy](ctx, db, "idx_proxy_account_id_unique")
},
// Post-auto so the per-bucket cost columns already exist when the legacy
// aggregates are folded into them and dropped.
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkAccessLog](ctx, db)
},
func(db *gorm.DB) error {
return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db)
},
}
}

View File

@@ -283,8 +283,8 @@ func (a *Account) SynthesizePrivateServiceZones(peerID string) []nbdns.CustomZon
// it, adding a single private service would black-hole every
// other name under the zone apex.
zone = &nbdns.CustomZone{
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
Domain: dns.Fqdn(serviceDomainZone),
Records: []nbdns.SimpleRecord{},
NonAuthoritative: true,
SearchDomainDisabled: true,
}
@@ -1082,6 +1082,7 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
peersExists := make(map[string]struct{})
rules := make([]*FirewallRule, 0)
peers := make([]*nbpeer.Peer, 0)
targetComponent := targetPeer.ToComponent()
return func(rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int) {
for _, peer := range groupPeers {
@@ -1117,10 +1118,10 @@ func (a *Account) connResourcesGenerator(ctx context.Context, targetPeer *nbpeer
if len(rule.Ports) == 0 && len(rule.PortRanges) == 0 {
rules = append(rules, &fr)
} else {
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetPeer)...)
rules = append(rules, ExpandPortsAndRanges(fr, rule, targetComponent)...)
}
rules = AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, FirewallRuleContext{
rules = AppendIPv6FirewallRule(rules, rulesExists, peer.ToComponent(), targetComponent, rule, FirewallRuleContext{
Direction: direction,
DirStr: strconv.Itoa(direction),
ProtocolStr: string(protocol),
@@ -1280,7 +1281,7 @@ func (a *Account) getRouteFirewallRules(ctx context.Context, peerID string, poli
return fwRules
}
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*nbpeer.Peer {
func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID string, distributionPeers map[string]struct{}, validatedPeersMap map[string]struct{}) []*ComponentPeer {
distPeersWithPolicy := make(map[string]struct{})
for _, id := range rule.Sources {
group := a.Groups[id]
@@ -1307,13 +1308,13 @@ func (a *Account) getRulePeers(rule *PolicyRule, postureChecks []string, peerID
}
}
distributionGroupPeers := make([]*nbpeer.Peer, 0, len(distPeersWithPolicy))
distributionGroupPeers := make([]*ComponentPeer, 0, len(distPeersWithPolicy))
for pID := range distPeersWithPolicy {
peer := a.Peers[pID]
if peer == nil {
continue
}
distributionGroupPeers = append(distributionGroupPeers, peer)
distributionGroupPeers = append(distributionGroupPeers, peer.ToComponent())
}
return distributionGroupPeers
}

View File

@@ -9,9 +9,7 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/internals/modules/zones"
resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/telemetry"
"github.com/netbirdio/netbird/route"
)
@@ -113,7 +111,7 @@ func (a *Account) GetPeerNetworkMapComponents(
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*nbpeer.Peer{peerID: peer},
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
@@ -126,7 +124,7 @@ func (a *Account) GetPeerNetworkMapComponents(
PeerID: peerID,
Network: a.Network.Copy(),
// must include the target peer as it's required on the client
Peers: map[string]*nbpeer.Peer{peerID: peer},
Peers: map[string]*ComponentPeer{peerID: peer.ToComponent()},
})
}
@@ -136,10 +134,10 @@ func (a *Account) GetPeerNetworkMapComponents(
NameServerGroups: make([]*nbdns.NameServerGroup, 0),
CustomZoneDomain: peersCustomZone.Domain,
ResourcePoliciesMap: make(map[string][]*Policy),
RoutersMap: make(map[string]map[string]*routerTypes.NetworkRouter),
NetworkResources: make([]*resourceTypes.NetworkResource, 0),
RoutersMap: make(map[string]map[string]*ComponentRouter),
NetworkResources: make([]*ComponentResource, 0),
PostureFailedPeers: make(map[string]map[string]struct{}, len(a.PostureChecks)),
RouterPeers: make(map[string]*nbpeer.Peer),
RouterPeers: make(map[string]*ComponentPeer),
NetworkXIDToPublicID: make(map[string]string, len(a.Networks)),
PostureCheckXIDToPublicID: make(map[string]string, len(a.PostureChecks)),
}
@@ -174,7 +172,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
components.Peers = relevantPeers
components.Groups = relevantGroups
components.Groups = GroupsToComponent(relevantGroups)
components.Policies = relevantPolicies
components.Routes = relevantRoutes
components.AllDNSRecords = filterDNSRecordsByPeers(peersCustomZone.Records, relevantPeers, peer.SupportsIPv6() && peer.IPv6.IsValid())
@@ -223,7 +221,7 @@ func (a *Account) GetPeerNetworkMapComponents(
}
for _, pID := range a.getPostureValidPeersSaveFailed(peers, policy.SourcePostureChecks, validatedPeersMap, &components.PostureFailedPeers) {
if _, exists := components.Peers[pID]; !exists {
components.Peers[pID] = a.GetPeer(pID)
components.Peers[pID] = a.GetPeer(pID).ToComponent()
}
}
} else {
@@ -256,14 +254,14 @@ func (a *Account) GetPeerNetworkMapComponents(
for _, srcGroupID := range rule.Sources {
if g := a.Groups[srcGroupID]; g != nil {
if _, exists := components.Groups[srcGroupID]; !exists {
components.Groups[srcGroupID] = g
components.Groups[srcGroupID] = g.ToComponent()
}
}
}
for _, dstGroupID := range rule.Destinations {
if g := a.Groups[dstGroupID]; g != nil {
if _, exists := components.Groups[dstGroupID]; !exists {
components.Groups[dstGroupID] = g
components.Groups[dstGroupID] = g.ToComponent()
}
}
}
@@ -278,20 +276,22 @@ func (a *Account) GetPeerNetworkMapComponents(
// network in the account — accounts with many tenants/networks
// shipped tens of unrelated peers in `peers[]` and `routers_map`.
if addSourcePeers {
components.RoutersMap[resource.NetworkID] = networkRoutingPeers
components.RoutersMap[resource.NetworkID] = routerTypes.ToComponentMap(networkRoutingPeers)
for peerIDKey := range networkRoutingPeers {
if p := a.Peers[peerIDKey]; p != nil {
if _, exists := components.RouterPeers[peerIDKey]; !exists {
components.RouterPeers[peerIDKey] = p
cp := components.RouterPeers[peerIDKey]
if cp == nil {
cp = p.ToComponent()
components.RouterPeers[peerIDKey] = cp
}
if _, exists := components.Peers[peerIDKey]; !exists {
if _, validated := validatedPeersMap[peerIDKey]; validated {
components.Peers[peerIDKey] = p
components.Peers[peerIDKey] = cp
}
}
}
}
components.NetworkResources = append(components.NetworkResources, resource)
components.NetworkResources = append(components.NetworkResources, resource.ToComponent())
}
}
@@ -312,14 +312,14 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
peerSSHEnabled bool,
validatedPeersMap map[string]struct{},
postureFailedPeers *map[string]map[string]struct{},
) (map[string]*nbpeer.Peer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*nbpeer.Peer, len(a.Peers)/4)
) (map[string]*ComponentPeer, map[string]*Group, []*Policy, []*route.Route, sshRequirements) {
relevantPeerIDs := make(map[string]*ComponentPeer, len(a.Peers)/4)
relevantGroupIDs := make(map[string]*Group, len(a.Groups)/4)
relevantPolicies := make([]*Policy, 0, len(a.Policies))
relevantRoutes := make([]*route.Route, 0, len(a.Routes))
sshReqs := sshRequirements{neededGroupIDs: make(map[string]struct{})}
relevantPeerIDs[peerID] = a.GetPeer(peerID)
relevantPeerIDs[peerID] = a.GetPeer(peerID).ToComponent()
peerGroupSet := make(map[string]struct{}, 8)
for groupID, group := range a.Groups {
@@ -384,7 +384,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if r.Peer != "" {
if _, ok := validatedPeersMap[r.Peer]; ok {
if p := a.GetPeer(r.Peer); p != nil {
relevantPeerIDs[r.Peer] = p
relevantPeerIDs[r.Peer] = p.ToComponent()
}
}
}
@@ -401,7 +401,7 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
continue
}
if p := a.GetPeer(pid); p != nil {
relevantPeerIDs[pid] = p
relevantPeerIDs[pid] = p.ToComponent()
}
}
}
@@ -458,7 +458,9 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInSources {
policyRelevant = true
for _, pid := range destinationPeers {
relevantPeerIDs[pid] = a.GetPeer(pid)
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, dstGroupID := range rule.Destinations {
relevantGroupIDs[dstGroupID] = a.GetGroup(dstGroupID)
@@ -468,7 +470,9 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
if peerInDestinations {
policyRelevant = true
for _, pid := range sourcePeers {
relevantPeerIDs[pid] = a.GetPeer(pid)
if _, exists := relevantPeerIDs[pid]; !exists {
relevantPeerIDs[pid] = a.GetPeer(pid).ToComponent()
}
}
for _, srcGroupID := range rule.Sources {
relevantGroupIDs[srcGroupID] = a.GetGroup(srcGroupID)
@@ -624,7 +628,7 @@ func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChe
// that name them. Calculate() tolerates groups with empty Peers (the inner
// loops simply iterate zero times), so retaining them is behaviourally a
// no-op for the legacy path that consumes the same NetworkMapComponents.
func filterGroupPeers(groups *map[string]*Group, peers map[string]*nbpeer.Peer) {
func filterGroupPeers(groups *map[string]*ComponentGroup, peers map[string]*ComponentPeer) {
for groupID, groupInfo := range *groups {
filteredPeers := make([]string, 0, len(groupInfo.Peers))
for _, pid := range groupInfo.Peers {
@@ -634,14 +638,14 @@ func filterGroupPeers(groups *map[string]*Group, peers map[string]*nbpeer.Peer)
}
if len(filteredPeers) != len(groupInfo.Peers) {
ng := groupInfo.Copy()
ng := *groupInfo
ng.Peers = filteredPeers
(*groups)[groupID] = ng
(*groups)[groupID] = &ng
}
}
}
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*nbpeer.Peer) {
func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}, policies []*Policy, resourcePoliciesMap map[string][]*Policy, peers map[string]*ComponentPeer) {
if len(*postureFailedPeers) == 0 {
return
}
@@ -676,7 +680,7 @@ func filterPostureFailedPeers(postureFailedPeers *map[string]map[string]struct{}
}
}
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*nbpeer.Peer, includeIPv6 bool) []nbdns.SimpleRecord {
func filterDNSRecordsByPeers(records []nbdns.SimpleRecord, peers map[string]*ComponentPeer, includeIPv6 bool) []nbdns.SimpleRecord {
if len(records) == 0 || len(peers) == 0 {
return nil
}

View File

@@ -9,7 +9,6 @@ import (
"github.com/stretchr/testify/require"
nbdns "github.com/netbirdio/netbird/dns"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
)
func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
@@ -49,7 +48,7 @@ func TestPrivateService_NetworkMap_UserPeer_AndProxyPeer(t *testing.T) {
})
}
func netmapPeerIDs(peers []*nbpeer.Peer) []string {
func netmapPeerIDs(peers []*ComponentPeer) []string {
ids := make([]string, 0, len(peers))
for _, p := range peers {
if p == nil {

View File

@@ -666,7 +666,7 @@ func Test_ExpandPortsAndRanges_SSHRuleExpansion(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer)
result := ExpandPortsAndRanges(tt.base, tt.rule, tt.peer.ToComponent())
var ports []string
for _, fr := range result {

View File

@@ -6,7 +6,6 @@ import (
"net"
"net/netip"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
nbroute "github.com/netbirdio/netbird/route"
sharedtypes "github.com/netbirdio/netbird/shared/management/types"
)
@@ -18,9 +17,6 @@ type DNSSettings = sharedtypes.DNSSettings
type FirewallRule = sharedtypes.FirewallRule
type Group = sharedtypes.Group
type GroupPeer = sharedtypes.GroupPeer
type Network = sharedtypes.Network
type NetworkMap = sharedtypes.NetworkMap
type ForwardingRule = sharedtypes.ForwardingRule
@@ -42,6 +38,18 @@ type RouteFirewallRule = sharedtypes.RouteFirewallRule
type NetworkMapComponents = sharedtypes.NetworkMapComponents
type ComponentPeer = sharedtypes.ComponentPeer
type ComponentGroup = sharedtypes.ComponentGroup
type ComponentRouter = sharedtypes.ComponentRouter
type ComponentResource = sharedtypes.ComponentResource
type ComponentResourceType = sharedtypes.ComponentResourceType
const (
ComponentResourceHost = sharedtypes.ComponentResourceHost
ComponentResourceSubnet = sharedtypes.ComponentResourceSubnet
ComponentResourceDomain = sharedtypes.ComponentResourceDomain
)
var EmptyNetworkMapComponents = sharedtypes.EmptyNetworkMapComponents
type AccountSettingsInfo = sharedtypes.AccountSettingsInfo
@@ -52,12 +60,7 @@ type NetworkMapComponentsCompact = sharedtypes.NetworkMapComponentsCompact
type LookupMap = sharedtypes.LookupMap
type FirewallRuleContext = sharedtypes.FirewallRuleContext
const (
GroupIssuedAPI = sharedtypes.GroupIssuedAPI
GroupIssuedJWT = sharedtypes.GroupIssuedJWT
GroupIssuedIntegration = sharedtypes.GroupIssuedIntegration
GroupAllName = sharedtypes.GroupAllName
)
const GroupAllName = sharedtypes.GroupAllName
// Function forwarders preserve types.X(...) call sites that previously
// resolved to package-local funcs. Plain forwarders (not var aliases) keep
@@ -67,11 +70,11 @@ func PolicyRuleImpliesLegacySSH(rule *PolicyRule) bool {
return sharedtypes.PolicyRuleImpliesLegacySSH(rule)
}
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *nbpeer.Peer) []*FirewallRule {
func ExpandPortsAndRanges(base FirewallRule, rule *PolicyRule, peer *ComponentPeer) []*FirewallRule {
return sharedtypes.ExpandPortsAndRanges(base, rule, peer)
}
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *nbpeer.Peer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
func AppendIPv6FirewallRule(rules []*FirewallRule, rulesExists map[string]struct{}, peer, targetPeer *ComponentPeer, rule *PolicyRule, rc FirewallRuleContext) []*FirewallRule {
return sharedtypes.AppendIPv6FirewallRule(rules, rulesExists, peer, targetPeer, rule, rc)
}
@@ -79,7 +82,7 @@ func CalculateNetworkMapFromComponents(ctx context.Context, components *NetworkM
return sharedtypes.CalculateNetworkMapFromComponents(ctx, components)
}
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*nbpeer.Peer, direction int, includeIPv6 bool) []*RouteFirewallRule {
func GenerateRouteFirewallRules(ctx context.Context, route *nbroute.Route, rule *PolicyRule, groupPeers []*ComponentPeer, direction int, includeIPv6 bool) []*RouteFirewallRule {
return sharedtypes.GenerateRouteFirewallRules(ctx, route, rule, groupPeers, direction, includeIPv6)
}

View File

@@ -2,7 +2,6 @@ package types
import (
"github.com/netbirdio/netbird/management/server/integration_reference"
"github.com/netbirdio/netbird/management/server/networks/resources/types"
)
const (
@@ -68,10 +67,6 @@ func (g *Group) EventMeta() map[string]any {
return map[string]any{"name": g.Name}
}
func (g *Group) EventMetaResource(resource *types.NetworkResource) map[string]any {
return map[string]any{"name": g.Name, "id": g.ID, "resource_name": resource.Name, "resource_id": resource.ID, "resource_type": resource.Type}
}
func (g *Group) Copy() *Group {
group := &Group{
ID: g.ID,
@@ -95,14 +90,39 @@ func (g *Group) HasPeers() bool {
return len(g.Peers) > 0
}
// GroupAllName is the reserved name of the default group that contains every peer in an account.
const GroupAllName = "All"
// IsGroupAll checks if the group is a default "All" group.
func (g *Group) IsGroupAll() bool {
return g.Name == GroupAllName
}
// ToComponent converts the group to its self-contained components
// representation. The Peers slice is shared, not copied — components are
// treated as immutable snapshots. Returns nil for a nil group.
func (g *Group) ToComponent() *ComponentGroup {
if g == nil {
return nil
}
return &ComponentGroup{
ID: g.ID,
PublicID: g.PublicID,
Name: g.Name,
Peers: g.Peers,
}
}
// GroupsToComponent converts an id-keyed group map to its components
// representation, preserving nil entries.
func GroupsToComponent(groups map[string]*Group) map[string]*ComponentGroup {
if groups == nil {
return nil
}
out := make(map[string]*ComponentGroup, len(groups))
for id, g := range groups {
out[id] = g.ToComponent()
}
return out
}
// AddPeer adds peerID to Peers if not present, returning true if added.
func (g *Group) AddPeer(peerID string) bool {
if peerID == "" {

View File

@@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/require"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/types"
)
func TestNetworkMapComponents_IPv6EndToEnd(t *testing.T) {
@@ -104,7 +105,7 @@ func TestNetworkMapComponents_RemotePeerWithoutCapability(t *testing.T) {
require.NotNil(t, nm)
t.Run("AllowedIPs include remote v6", func(t *testing.T) {
var dst *nbpeer.Peer
var dst *types.ComponentPeer
for _, p := range nm.Peers {
if p.ID == "peer-dst-1" {
dst = p

View File

@@ -49,7 +49,7 @@ func allPeersValidated(account *types.Account, excludePeerIDs ...string) map[str
return validated
}
func peerIDs(peers []*nbpeer.Peer) []string {
func peerIDs(peers []*types.ComponentPeer) []string {
ids := make([]string, len(peers))
for i, p := range peers {
ids[i] = p.ID

View File

@@ -19,34 +19,3 @@ func Difference(a, b []string) []string {
func ToPtr[T any](value T) *T {
return &value
}
type comparableObject[T any] interface {
Equal(other T) bool
}
func MergeUnique[T comparableObject[T]](arr1, arr2 []T) []T {
var result []T
for _, item := range arr1 {
if !contains(result, item) {
result = append(result, item)
}
}
for _, item := range arr2 {
if !contains(result, item) {
result = append(result, item)
}
}
return result
}
func contains[T comparableObject[T]](slice []T, element T) bool {
for _, item := range slice {
if item.Equal(element) {
return true
}
}
return false
}

View File

@@ -10,6 +10,7 @@ COPY encryption ./encryption
COPY flow ./flow
COPY formatter ./formatter
COPY monotime ./monotime
COPY management ./management
COPY proxy ./proxy
COPY route ./route
COPY shared ./shared

View File

@@ -221,14 +221,21 @@ func (l *Logger) allowDenyLog(serviceID types.ServiceID, reason string) bool {
// proxy/internal/middleware/keys.go — only the dimensions management needs to
// record a usage row (provider / model / tokens / cost / groups).
var usageMetadataKeys = map[string]struct{}{
"llm.provider": {},
"llm.model": {},
"llm.resolved_provider_id": {},
"llm.input_tokens": {},
"llm.output_tokens": {},
"llm.total_tokens": {},
"cost.usd_total": {},
"llm.authorising_groups": {},
"llm.provider": {},
"llm.model": {},
"llm.resolved_provider_id": {},
"llm.input_tokens": {},
"llm.output_tokens": {},
"llm.total_tokens": {},
"llm.cached_input_tokens": {},
"llm.cache_creation_tokens": {},
"cost.usd_input": {},
"cost.usd_cached_input": {},
"cost.usd_cache_creation": {},
"cost.usd_output": {},
"cost.usd_total": {},
"cost.usd_cache": {},
"llm.authorising_groups": {},
}
// stripAgentNetworkEntryForUsage returns the entry reduced to what's needed to

View File

@@ -3,20 +3,30 @@ package auth
import (
"context"
"net/netip"
"os"
"strings"
"sync"
"time"
log "github.com/sirupsen/logrus"
"golang.org/x/sync/singleflight"
"github.com/netbirdio/netbird/proxy/internal/types"
"github.com/netbirdio/netbird/shared/management/proto"
)
// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
// reused before re-fetching from management. 5 minutes balances freshness
// against management load on busy mesh networks.
// tunnelCacheTTL is the default cap on how long a positive ValidateTunnelPeer
// result is reused before re-fetching from management. 5 minutes balances
// freshness against management load on busy mesh networks. Override it with
// envTunnelCacheTTL when an account needs authorization changes to take effect
// sooner (at the cost of more ValidateTunnelPeer RPCs).
const tunnelCacheTTL = 300 * time.Second
// envTunnelCacheTTL overrides tunnelCacheTTL. The value is a Go duration string
// (e.g. "30s", "2m"); an unset, unparseable, or non-positive value keeps the
// default.
const envTunnelCacheTTL = "NB_PROXY_TUNNEL_CACHE_TTL"
// tunnelCachePerAccount caps the number of cached identities per account.
// Bounded eviction avoids memory growth in pathological cases (huge peer
// roster, brief request bursts) while staying generous for normal use.
@@ -60,16 +70,35 @@ type accountBucket struct {
order []tunnelCacheKey
}
// newTunnelValidationCache constructs a cache with default TTL and bounds.
// newTunnelValidationCache constructs a cache with the configured TTL
// (envTunnelCacheTTL override or default) and default bounds.
func newTunnelValidationCache() *tunnelValidationCache {
return &tunnelValidationCache{
entries: make(map[types.AccountID]*accountBucket),
ttl: tunnelCacheTTL,
ttl: tunnelCacheTTLFromEnv(),
maxSize: tunnelCachePerAccount,
now: time.Now,
}
}
// tunnelCacheTTLFromEnv returns the tunnel-cache TTL, honoring the
// envTunnelCacheTTL override. The override must be a positive Go duration
// string (e.g. "30s", "2m"); anything unset, unparseable, or non-positive
// falls back to tunnelCacheTTL.
func tunnelCacheTTLFromEnv() time.Duration {
raw := strings.TrimSpace(os.Getenv(envTunnelCacheTTL))
if raw == "" {
return tunnelCacheTTL
}
d, err := time.ParseDuration(raw)
if err != nil || d <= 0 {
log.Warnf("ignoring invalid %s=%q (want a positive Go duration like 30s or 2m); using default %s",
envTunnelCacheTTL, raw, tunnelCacheTTL)
return tunnelCacheTTL
}
return d
}
// get returns a cached response for the key, or nil when missing or
// expired. Expired entries are evicted lazily on read.
func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {

View File

@@ -169,3 +169,32 @@ func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
}
func TestTunnelCacheTTLFromEnv(t *testing.T) {
t.Run("unset uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("valid duration overrides", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "45s")
assert.Equal(t, 45*time.Second, tunnelCacheTTLFromEnv())
})
t.Run("whitespace trimmed", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, " 2m ")
assert.Equal(t, 2*time.Minute, tunnelCacheTTLFromEnv())
})
t.Run("unparseable uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "nonsense")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("non-positive uses default", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "0s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
t.Setenv(envTunnelCacheTTL, "-30s")
assert.Equal(t, tunnelCacheTTL, tunnelCacheTTLFromEnv())
})
t.Run("constructor honors override", func(t *testing.T) {
t.Setenv(envTunnelCacheTTL, "90s")
assert.Equal(t, 90*time.Second, newTunnelValidationCache().ttl)
})
}

View File

@@ -56,10 +56,12 @@ type bedrockResponse struct {
OutputTokens int64 `json:"output_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
// Converse — camelCase.
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
// Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens).
InputTokensCamel int64 `json:"inputTokens"`
OutputTokensCamel int64 `json:"outputTokens"`
TotalTokensCamel int64 `json:"totalTokens"`
CacheReadTokensCamel int64 `json:"cacheReadInputTokens"`
CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"`
} `json:"usage"`
}
@@ -83,16 +85,18 @@ func (BedrockParser) ParseResponse(status int, contentType string, body []byte)
}
inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel)
cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel)
total := resp.Usage.TotalTokensCamel
if total == 0 {
total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
total = inTok + outTok + cacheRead + cacheWrite
}
return Usage{
InputTokens: inTok,
OutputTokens: outTok,
TotalTokens: total,
CachedInputTokens: resp.Usage.CacheReadInputTokens,
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
CachedInputTokens: cacheRead,
CacheCreationTokens: cacheWrite,
}, nil
}

View File

@@ -26,6 +26,18 @@ func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
}
// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields.
func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) {
body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`)
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
require.NoError(t, err)
require.Equal(t, int64(11), u.InputTokens, "converse input tokens")
require.Equal(t, int64(3), u.OutputTokens, "converse output tokens")
require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens")
require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens")
require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens")
}
func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")

View File

@@ -28,12 +28,12 @@ func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
"ministral-8b-latest", "mistral-embed",
},
"anthropic": {
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
},
// bedrock keys are the normalized ids the request parser emits.
"bedrock": {
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",

View File

@@ -161,6 +161,16 @@ openai:
input_per_1k: 0.0001
output_per_1k: 0
# Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
# reports cache hits OpenAI-style when present; cached input is 10% of
# input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
# serves newer accounts (K2-era ids and kimi-latest 404), matching the
# management catalog.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cached_input_per_1k: 0.0003
anthropic:
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
# Pricing source: Anthropic's current published rates per million tokens,
@@ -170,6 +180,11 @@ anthropic:
output_per_1k: 0.050
cache_read_per_1k: 0.001
cache_creation_per_1k: 0.0125
claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025
@@ -206,12 +221,31 @@ anthropic:
cache_read_per_1k: 0.0001
cache_creation_per_1k: 0.00125
# Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
# (/anthropic/v1/messages — the official Claude Code setup). Same rates
# as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
# Claude Code guides set for the 1M-context alias; priced identically so
# cost metering doesn't silently skip those requests.
kimi-k3:
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
"kimi-k3[1m]":
input_per_1k: 0.003
output_per_1k: 0.015
cache_read_per_1k: 0.0003
bedrock:
# AWS Bedrock model ids, normalised by the request parser (cross-region
# inference-profile prefix + version/throughput suffix stripped), e.g.
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
anthropic.claude-opus-5:
input_per_1k: 0.005
output_per_1k: 0.025
cache_read_per_1k: 0.0005
cache_creation_per_1k: 0.00625
anthropic.claude-opus-4-8:
input_per_1k: 0.005
output_per_1k: 0.025

View File

@@ -128,6 +128,46 @@ type Table struct {
// - Other providers: cached and cacheCreation are ignored; cost is
// inTokens*InputPer1K + outTokens*OutputPer1K.
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
return c.TotalUSD, ok
}
// Costs is a per-request cost split. The four per-bucket fields are the base
// of the breakdown — one per token bucket the provider bills separately — and
// the two aggregates are derived from them:
//
// TotalUSD = InputUSD + CachedInputUSD + CacheCreationUSD + OutputUSD
// CacheUSD = CachedInputUSD + CacheCreationUSD
//
// InputUSD is always the cost of the *non-cached* input bucket, for both
// provider shapes: on OpenAI the cached subset is carved out of inTokens and
// billed as CachedInputUSD, so the two never double-count. Buckets a provider
// doesn't bill are zero, which keeps the identities above true everywhere.
type Costs struct {
InputUSD float64
CachedInputUSD float64
CacheCreationUSD float64
OutputUSD float64
TotalUSD float64
CacheUSD float64
}
// newCosts assembles a split from its per-bucket parts, deriving the two
// aggregates so TotalUSD and CacheUSD can never drift from the breakdown.
func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
return Costs{
InputUSD: input,
CachedInputUSD: cachedInput,
CacheCreationUSD: cacheCreation,
OutputUSD: output,
TotalUSD: input + cachedInput + cacheCreation + output,
CacheUSD: cachedInput + cacheCreation,
}
}
// Costs returns the estimated USD cost split for the given token counts, with
// the same semantics as Cost.
func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
// Clamp negatives to zero before any pricing math so a malformed
// upstream count can never produce a negative cost.
if inTokens < 0 {
@@ -143,15 +183,15 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
cacheCreation = 0
}
if t == nil {
return 0, false
return Costs{}, false
}
byModel, ok := t.entries[provider]
if !ok {
return 0, false
return Costs{}, false
}
entry, ok := byModel[model]
if !ok {
return 0, false
return Costs{}, false
}
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
switch provider {
@@ -168,7 +208,7 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
}
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
cached := float64(clamped) / 1000.0 * cachedRate
return nonCached + cached + output, true
return newCosts(nonCached, cached, 0, output), true
case "anthropic", "bedrock":
// Bedrock-Anthropic returns the same additive cache buckets as
// first-party Anthropic; non-Anthropic Bedrock models simply report
@@ -184,10 +224,10 @@ func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, c
input := float64(inTokens) / 1000.0 * entry.InputPer1K
read := float64(cachedInput) / 1000.0 * readRate
create := float64(cacheCreation) / 1000.0 * createRate
return input + read + create + output, true
return newCosts(input, read, create, output), true
default:
input := float64(inTokens) / 1000.0 * entry.InputPer1K
return input + output, true
return newCosts(input, 0, 0, output), true
}
}

View File

@@ -0,0 +1,329 @@
package builtin_test
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"strconv"
"testing"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
)
// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the embedded default pricing
// table and asserts exact USD amounts hardcoded from the vendors' published prices, including the cache split.
func TestCostCalculation_ProviderMatrix(t *testing.T) {
// Empty data dir → embedded defaults, like a proxy with no pricing override.
builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil)
reqMW, err := llm_request_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_request_parser")
respMW, err := llm_response_parser.Factory{}.New(nil)
require.NoError(t, err, "build llm_response_parser")
costMW, err := cost_meter.Factory{}.New(nil)
require.NoError(t, err, "build cost_meter")
t.Cleanup(func() { _ = costMW.Close() })
const jsonCT = "application/json"
const sseCT = "text/event-stream"
const awsCT = "application/vnd.amazon.eventstream"
cases := []struct {
name string
url string
reqBody []byte
respCT string
respBody []byte
wantProvider string
wantModel string
wantCost float64 // exact expected USD; ignored when wantSkip is set
wantCacheCost float64 // expected cost.usd_cache portion of wantCost
wantSkip string // expected cost.skipped reason, "" when priced
}{
{
// gpt-4o-mini $0.15/$0.60 per MTok: 1000×0.15/1M + 500×0.60/1M.
name: "openai chat completions",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"choices":[{"message":{"content":"pong"}}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// OpenAI cached tokens are a SUBSET of prompt_tokens at a discount; gpt-4o $2.50/$10 per MTok, cached $1.25/M:
// 250×2.5/1M + 750×1.25/1M + 500×10/1M.
name: "openai cached subset discount",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500,"prompt_tokens_details":{"cached_tokens":750}}}`),
wantProvider: "openai",
wantModel: "gpt-4o",
wantCost: 0.0065625,
wantCacheCost: 0.0009375,
},
{
// OpenAI streaming: usage rides the final SSE frame.
name: "openai chat SSE stream",
url: "https://api.openai.com/v1/chat/completions",
reqBody: []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"choices":[{"delta":{"content":"po"}}]}`, `{"choices":[{"delta":{"content":"ng"}}]}`, `{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500}}`, "[DONE]"),
wantProvider: "openai",
wantModel: "gpt-4o-mini",
wantCost: 0.00045,
},
{
// Mistral speaks the OpenAI shape: mistral-large-latest $0.50/$1.50 per MTok.
name: "mistral via openai shape",
url: "https://api.mistral.ai/v1/chat/completions",
reqBody: []byte(`{"model":"mistral-large-latest","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":1000}}`),
wantProvider: "openai",
wantModel: "mistral-large-latest",
wantCost: 0.002,
},
{
// The field report, minus caching: Bedrock Sonnet 4.6 $3/$15 per MTok, 3×3/1M + 1514×15/1M = $0.022719.
// Also covers inference-profile normalization of the region-prefixed versioned id in the URL.
name: "bedrock invoke — reported scenario, no cache",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":3,"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.022719,
},
{
// The field report as observed: the FIRST call of a session also wrote a 30,528-token prompt cache at
// 1.25× input ($3.75/M): 0.022719 + 30528×3.75/1M = $0.137199 — the reported $0.1372.
name: "bedrock invoke — reported scenario with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":3,"output_tokens":1514,"cache_creation_input_tokens":30528,"cache_read_input_tokens":0}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Same numbers over the InvokeModel event-stream: message_start carries input + cache, message_delta the output.
name: "bedrock invoke stream with cache write",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke-with-response-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
respCT: awsCT,
respBody: bedrockInvokeStream(t, `{"type":"message_start","message":{"usage":{"input_tokens":3,"output_tokens":1,"cache_creation_input_tokens":30528}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":1514}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-sonnet-4-6",
wantCost: 0.137199,
wantCacheCost: 0.11448,
},
{
// Converse camelCase usage incl. cache buckets. Haiku 4.5 $1/$5 per MTok, read $0.10/M, write $1.25/M:
// 50×1/1M + 100×5/1M + 2000×0.1/1M + 1000×1.25/1M = $0.002.
name: "bedrock converse with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: jsonCT,
respBody: []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// Same numbers over converse-stream: usage rides the trailing metadata frame.
name: "bedrock converse stream with cache buckets",
url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse-stream",
reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
respCT: awsCT,
respBody: bedrockConverseStream(t,
`{"delta":{"text":"pong"}}`,
`{"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`,
),
wantProvider: "bedrock",
wantModel: "anthropic.claude-haiku-4-5",
wantCost: 0.002,
wantCacheCost: 0.00145,
},
{
// First-party Anthropic, additive cache buckets. Sonnet 4.6:
// 256×3/1M + 200×15/1M + 768×0.3/1M + 512×3.75/1M.
name: "anthropic messages with cache buckets",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0059184,
wantCacheCost: 0.0021504,
},
{
// Anthropic SSE: input from message_start, output from message_delta. Haiku 4.5: 1000×1/1M + 2000×5/1M.
name: "anthropic SSE stream",
url: "https://api.anthropic.com/v1/messages",
reqBody: []byte(`{"model":"claude-haiku-4-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
respCT: sseCT,
respBody: sseBody(`{"type":"message_start","message":{"usage":{"input_tokens":1000,"output_tokens":2}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":2000}}`, `{"type":"message_stop"}`),
wantProvider: "anthropic",
wantModel: "claude-haiku-4-5",
wantCost: 0.011,
},
{
// Kimi's Anthropic-compatible endpoint: kimi-k3 $3/$15 per MTok under the anthropic table.
name: "kimi anthropic shape",
url: "https://api.moonshot.ai/anthropic/v1/messages",
reqBody: []byte(`{"model":"kimi-k3","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":1000,"output_tokens":1000}}`),
wantProvider: "anthropic",
wantModel: "kimi-k3",
wantCost: 0.018,
},
{
// Vertex path-routed model with "@version" stripped; Anthropic-on-Vertex priced under the anthropic table.
name: "vertex anthropic path-routed",
url: "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-6@20260115:rawPredict",
reqBody: []byte(`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"input_tokens":200,"output_tokens":100}}`),
wantProvider: "anthropic",
wantModel: "claude-sonnet-4-6",
wantCost: 0.0021,
},
{
// Gateway-prefixed model ids are not in the pricing table: the meter must SKIP, never guess a rate.
name: "gateway-prefixed model skips pricing",
url: "https://gateway.example.com/v1/chat/completions",
reqBody: []byte(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
respCT: jsonCT,
respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500}}`),
wantProvider: "openai",
wantModel: "openai/gpt-4o-mini",
wantSkip: "unknown_model",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
in := &middleware.Input{
Method: "POST",
URL: tc.url,
Headers: []middleware.KV{{Key: "Content-Type", Value: "application/json"}},
Body: tc.reqBody,
}
reqOut, err := reqMW.Invoke(context.Background(), in)
require.NoError(t, err, "request parser")
in.Metadata = append(in.Metadata, reqOut.Metadata...)
require.Equal(t, tc.wantProvider, metaKV(in.Metadata, middleware.KeyLLMProvider), "detected provider")
require.Equal(t, tc.wantModel, metaKV(in.Metadata, middleware.KeyLLMModel), "detected (normalized) model")
in.Status = 200
in.RespHeaders = []middleware.KV{{Key: "Content-Type", Value: tc.respCT}}
in.RespBody = tc.respBody
respOut, err := respMW.Invoke(context.Background(), in)
require.NoError(t, err, "response parser")
in.Metadata = append(in.Metadata, respOut.Metadata...)
costOut, err := costMW.Invoke(context.Background(), in)
require.NoError(t, err, "cost meter")
if tc.wantSkip != "" {
assert.Equal(t, tc.wantSkip, metaKV(costOut.Metadata, middleware.KeyCostSkipped), "expected cost skip reason")
assert.Empty(t, metaKV(costOut.Metadata, middleware.KeyCostUSDTotal), "no cost may be emitted on skip")
return
}
raw := metaKV(costOut.Metadata, middleware.KeyCostUSDTotal)
require.NotEmpty(t, raw, "cost.usd_total must be emitted; skip=%q", metaKV(costOut.Metadata, middleware.KeyCostSkipped))
got, err := strconv.ParseFloat(raw, 64)
require.NoError(t, err, "cost must be a float")
// cost.usd_total is rendered with %.6f: allow half of the last printed digit on top of float error.
assert.InDelta(t, tc.wantCost, got, 5.1e-7, "USD cost for %s", tc.name)
rawCache := metaKV(costOut.Metadata, middleware.KeyCostUSDCache)
require.NotEmpty(t, rawCache, "cost.usd_cache must be emitted next to cost.usd_total")
gotCache, err := strconv.ParseFloat(rawCache, 64)
require.NoError(t, err, "cache cost must be a float")
assert.InDelta(t, tc.wantCacheCost, gotCache, 5.1e-7, "cache USD cost for %s", tc.name)
})
}
}
// metaKV returns the value for key in kvs, or "" when absent.
func metaKV(kvs []middleware.KV, key string) string {
for _, kv := range kvs {
if kv.Key == key {
return kv.Value
}
}
return ""
}
// sseBody renders data frames as a text/event-stream body.
func sseBody(frames ...string) []byte {
var b bytes.Buffer
for _, f := range frames {
b.WriteString("data: ")
b.WriteString(f)
b.WriteString("\n\n")
}
return b.Bytes()
}
// awsFrame encodes one AWS event-stream frame with the given :event-type.
func awsFrame(t *testing.T, eventType string, payload []byte) []byte {
t.Helper()
var buf bytes.Buffer
enc := eventstream.NewEncoder()
require.NoError(t, enc.Encode(&buf, eventstream.Message{
Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}},
Payload: payload,
}), "encode event-stream frame")
return buf.Bytes()
}
// bedrockInvokeStream builds an invoke-with-response-stream body: each "chunk" frame wraps a base64 Anthropic event.
func bedrockInvokeStream(t *testing.T, events ...string) []byte {
t.Helper()
var body bytes.Buffer
for _, ev := range events {
wrap, err := json.Marshal(map[string]string{"bytes": base64.StdEncoding.EncodeToString([]byte(ev))})
require.NoError(t, err)
body.Write(awsFrame(t, "chunk", wrap))
}
return body.Bytes()
}
// bedrockConverseStream builds a converse-stream body: contentBlockDelta frames plus a trailing metadata usage frame.
func bedrockConverseStream(t *testing.T, deltas ...string) []byte {
t.Helper()
var body bytes.Buffer
for i, ev := range deltas {
eventType := "contentBlockDelta"
if i == len(deltas)-1 {
eventType = "metadata"
}
body.Write(awsFrame(t, eventType, []byte(ev)))
}
return body.Bytes()
}

View File

@@ -32,7 +32,12 @@ const (
)
var metadataKeys = []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
@@ -140,18 +145,38 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar
}
table := m.loader.Get()
cost, ok := table.Cost(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
costs, ok := table.Costs(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
if !ok {
out.Metadata = skip(skipUnknownModel)
return out, nil
}
// Per-bucket costs first: they're the base of the breakdown, and the two
// aggregates that follow are derived from exactly these four values.
out.Metadata = []middleware.KV{
{Key: middleware.KeyCostUSDTotal, Value: fmt.Sprintf("%.6f", cost)},
{Key: middleware.KeyCostUSDInput, Value: usd(costs.InputUSD)},
{Key: middleware.KeyCostUSDCachedInput, Value: usd(costs.CachedInputUSD)},
{Key: middleware.KeyCostUSDCacheCreation, Value: usd(costs.CacheCreationUSD)},
{Key: middleware.KeyCostUSDOutput, Value: usd(costs.OutputUSD)},
{Key: middleware.KeyCostUSDTotal, Value: usd(costs.TotalUSD)},
{Key: middleware.KeyCostUSDCache, Value: usd(costs.CacheUSD)},
}
return out, nil
}
// usd renders a cost as the fixed-precision string every cost.usd_* key
// carries, so the per-bucket values and the aggregates round identically.
//
// 9 decimals, not 6: these values are summed downstream — per request, per
// session, and per usage bucket — so the rounding step is applied once per
// bucket per row and then accumulated. At 6 decimals a single row loses up to
// 2e-6 across its four buckets (enough to break a 1e-6 reconciliation against
// published rates), and a bucket smaller than half a microdollar quantises to
// zero outright: 16 cache-read tokens on a cheap model is 1.6e-9, so summing
// 10k such rows reports 0.02 instead of 0.016. Nano-dollar precision keeps the
// per-row error ~1000x below the smallest realistic bucket.
func usd(v float64) string { return fmt.Sprintf("%.9f", v) }
// skip returns a single-entry metadata slice carrying the given skip
// reason under KeyCostSkipped.
func skip(reason string) []middleware.KV {

View File

@@ -67,7 +67,15 @@ func TestMiddleware_StaticSurface(t *testing.T) {
assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op")
keys := mw.MetadataKeys()
expected := []string{middleware.KeyCostUSDTotal, middleware.KeyCostSkipped}
expected := []string{
middleware.KeyCostUSDInput,
middleware.KeyCostUSDCachedInput,
middleware.KeyCostUSDCacheCreation,
middleware.KeyCostUSDOutput,
middleware.KeyCostUSDTotal,
middleware.KeyCostUSDCache,
middleware.KeyCostSkipped,
}
assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
}
@@ -105,7 +113,7 @@ func TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted for known model")
assert.Equal(t, "0.000750", value, "0.00015 + 0.0006 per 1k tokens, 6-decimal format")
assert.Equal(t, "0.000750000", value, "0.00015 + 0.0006 per 1k tokens, 9-decimal format")
}
func TestFactory_PricingPathOverride(t *testing.T) {
@@ -129,7 +137,7 @@ func TestFactory_PricingPathOverride(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted with custom pricing path")
assert.Equal(t, "0.015000", value, "2*0.0025 + 1*0.01 = 0.015 with 6-decimal format")
assert.Equal(t, "0.015000000", value, "2*0.0025 + 1*0.01 = 0.015 with 9-decimal format")
}
func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
@@ -148,7 +156,7 @@ func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cost.usd_total must be emitted")
assert.Equal(t, "0.018000", value, "0.003 + 0.015 = 0.018 with 6-decimal format")
assert.Equal(t, "0.018000000", value, "0.003 + 0.015 = 0.018 with 9-decimal format")
_, skipped := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
assert.False(t, skipped, "cost.skipped must not be set when cost is computed")
}
@@ -357,8 +365,25 @@ func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "cached subset path must produce a cost — never a skip")
// 250 non-cached at 0.0025/1k + 750 cached at 0.00125/1k + 500 output at 0.01/1k.
assert.Equal(t, "0.006563", value,
assert.Equal(t, "0.006562500", value,
"cached subset must be billed at the discount rate, non-cached at the full rate; never double-billed")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 750 cached at 0.00125/1k = 0.0009375.
assert.Equal(t, "0.000937500", cache, "cache cost is the discounted cost of the cached subset")
// Per-bucket breakdown. On OpenAI the cached subset is carved out of the
// input bucket, so input covers only the 250 non-cached tokens — the two
// must never double-count the same 750 tokens.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000625000",
"input bucket bills only the non-cached remainder")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000937500",
"cached-input bucket bills the discounted subset")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.000000000",
"OpenAI has no cache-write bucket")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.005000000",
"output bucket bills 500 tokens at 0.01/1k")
}
// TestInvoke_AnthropicCacheBucketsAdditive proves the Anthropic
@@ -384,9 +409,33 @@ func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 256 input * 0.003 + 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 + 200 output * 0.015
// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184 → "0.005918" with 6-decimal format.
assert.Equal(t, "0.005918", value,
// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184.
assert.Equal(t, "0.005918400", value,
"each Anthropic input bucket must bill at its own rate — cache_read cheap, cache_creation expensive, regular input mid")
cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
// 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 = 0.0021504.
assert.Equal(t, "0.002150400", cache, "cache cost sums the read and creation buckets")
// Per-bucket breakdown: four separately-billed buckets, each at its own rate.
assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000768000",
"input bucket bills 256 tokens at 0.003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000230400",
"cache-read bucket bills 768 tokens at the cheap 0.0003/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.001920000",
"cache-write bucket bills 512 tokens at the expensive 0.00375/1k")
assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.003000000",
"output bucket bills 200 tokens at 0.015/1k")
}
// assertBucket asserts one per-bucket cost key carries the expected
// 6-decimal value.
func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
t.Helper()
got, ok := metaValue(t, md, key)
require.Truef(t, ok, "%s must be emitted", key)
assert.Equal(t, want, got, msg)
}
// TestInvoke_CachedTokensAbsentFallsBackToBaseFormula covers the
@@ -411,7 +460,7 @@ func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok)
// 1000 input * 0.0025 + 500 output * 0.01 = 0.0025 + 0.005 = 0.0075
assert.Equal(t, "0.007500", value, "no cached metadata = same cost as before the feature landed")
assert.Equal(t, "0.007500000", value, "no cached metadata = same cost as before the feature landed")
}
// TestInvoke_UnparseableCachedTokensSkippedSilently proves the
@@ -435,7 +484,7 @@ func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) {
require.NoError(t, err)
value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
require.True(t, ok, "garbage cache metadata must NOT switch the response from a cost to a skip — fall back to 0 cached")
assert.Equal(t, "0.007500", value, "same as the no-cached-metadata path")
assert.Equal(t, "0.007500000", value, "same as the no-cached-metadata path")
}
// TestMiddleware_CloseCancelsReloader proves Close stops the per-instance

View File

@@ -10,11 +10,15 @@ import (
)
// Config is the JSON-decoded shape accepted by the factory. The
// runtime path consumes the normalised allowlist; raw config is not
// runtime path consumes the normalised allowlists; raw config is not
// retained beyond construction.
type Config struct {
ModelAllowlist []string `json:"model_allowlist"`
PromptCapture PromptCapture `json:"prompt_capture"`
// ProviderAllowlists maps a resolved provider id (KeyLLMResolvedProviderID) to
// its model allowlist. A provider present is restricted to those models; one
// absent is unrestricted. Kept per-provider so one provider's list can't leak
// onto another.
ProviderAllowlists map[string][]string `json:"provider_allowlists,omitempty"`
PromptCapture PromptCapture `json:"prompt_capture"`
}
// PromptCapture toggles the optional prompt capture + redaction step
@@ -54,21 +58,28 @@ func isEmptyJSON(raw []byte) bool {
return false
}
// normaliseConfig lowercases and trims allowlist entries so the runtime
// match is case-insensitive. Empty entries are dropped.
// normaliseConfig lowercases and trims allowlist entries for case-insensitive
// matching; empty entries drop. A provider whose entries all drop keeps an empty
// (non-nil) list — "deny every model" — distinct from an absent provider
// (unrestricted).
func normaliseConfig(cfg Config) Config {
if len(cfg.ModelAllowlist) == 0 {
if len(cfg.ProviderAllowlists) == 0 {
cfg.ProviderAllowlists = nil
return cfg
}
cleaned := make([]string, 0, len(cfg.ModelAllowlist))
for _, entry := range cfg.ModelAllowlist {
n := normaliseModel(entry)
if n == "" {
continue
cleaned := make(map[string][]string, len(cfg.ProviderAllowlists))
for provider, models := range cfg.ProviderAllowlists {
list := make([]string, 0, len(models))
for _, entry := range models {
n := normaliseModel(entry)
if n == "" {
continue
}
list = append(list, n)
}
cleaned = append(cleaned, n)
cleaned[provider] = list
}
cfg.ModelAllowlist = cleaned
cfg.ProviderAllowlists = cleaned
return cfg
}

Some files were not shown because too many files have changed in this diff Show More