mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-25 00:51:28 +02:00
acfad1a384
Three open review findings. The discovery filter treated every slash in a listed id as a gateway prefix and matched the tail against the policy. A self-hosted id carries slashes of its own, and an upstream may scope ids per tenant, so "tenant-b/claude-sonnet-5" matched a permitted "claude-sonnet-5" and reached the picker — a model the policy never named, and one the guardrail denies on sight, since enforcement compares the id as written. Strip only the namespaces a gateway is known to prepend, taken from the first slash rather than the last. The e2e retry loops slept between attempts without watching the context, so a cancelled run kept retrying calls that fail instantly and spent its remaining window sleeping between them. They now stop when the context is done. The streamed provider's setup key outlived its test: deleting the group does not delete the key that auto-joins it.
210 lines
8.4 KiB
Go
210 lines
8.4 KiB
Go
//go:build e2e
|
|
|
|
package agentnetwork
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"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"
|
|
)
|
|
|
|
// streamedModel is priced high enough that a mis-metered request is obvious in
|
|
// the recorded cost, and named so it cannot collide with another test's route.
|
|
const streamedModel = "e2e-streamed-model"
|
|
|
|
const (
|
|
streamInRate = 0.010
|
|
streamOutRate = 0.020
|
|
// The cache-read bucket is priced separately from input, so a run that
|
|
// folded the two together fails the per-bucket assertions below.
|
|
streamCacheReadRate = 0.001
|
|
)
|
|
|
|
// TestStreamingResponseMetersInputTokens is the end-to-end guard for the
|
|
// metering bug this endpoint's gateway-protocol work fixed.
|
|
//
|
|
// On a streamed answer the input-token count exists only in the opening
|
|
// message_start event; every later frame reports output. A response read with
|
|
// the wrong vendor's parser — the shape a gateway record produces when it names
|
|
// one API surface and serves another — never looks at that event, so input
|
|
// metered as zero and the bulk of the bill silently vanished. Nothing in the
|
|
// suite sent stream: true before this test, so the whole branch went unrun.
|
|
//
|
|
// The provider points at the mock's streaming listener, which answers every
|
|
// request as SSE with token counts that differ from the buffered surface. That
|
|
// difference is the point: passing these assertions is only possible if the
|
|
// stream accumulator ran.
|
|
func TestStreamingResponseMetersInputTokens(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
|
defer cancel()
|
|
|
|
env := provisionStreamingProvider(t, ctx, "anthropic_api")
|
|
|
|
sessionID := fmt.Sprintf("e2e-session-stream-%d", time.Now().UnixNano())
|
|
code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID)
|
|
require.Equal(t, 200, code, "streamed chat must succeed; body: %s", body)
|
|
assert.Contains(t, body, "message_start",
|
|
"the client must receive the event stream itself, not a buffered rewrite of it")
|
|
|
|
row := findAccessLogBySession(t, ctx, sessionID)
|
|
|
|
assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens),
|
|
"input tokens live in message_start; zero here is the bug this test exists for")
|
|
assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens),
|
|
"output tokens ride message_delta and supersede the message_start seed")
|
|
assert.Equal(t, harness.VLLMStreamCacheReadTokens, int(row.CachedInputTokens),
|
|
"the Anthropic cache bucket rides message_start too, and only its own parser reads it")
|
|
|
|
// The Anthropic surface bills cache reads additively, so the input bucket
|
|
// prices the full input count rather than a remainder.
|
|
wantInput := float64(harness.VLLMStreamInputTokens) / 1000 * streamInRate
|
|
wantOutput := float64(harness.VLLMStreamOutputTokens) / 1000 * streamOutRate
|
|
wantCacheRead := float64(harness.VLLMStreamCacheReadTokens) / 1000 * streamCacheReadRate
|
|
assert.InDelta(t, wantInput, row.InputCostUsd, 1e-6, "input cost must price the streamed input tokens")
|
|
assert.InDelta(t, wantOutput, row.OutputCostUsd, 1e-6, "output cost must price the streamed output tokens")
|
|
// The total, not merely a positive number: input and output alone are
|
|
// positive, so a cache bucket parsed and then never billed would pass any
|
|
// weaker assertion. The gap is 7e-6, well outside the delta.
|
|
assert.InDelta(t, wantInput+wantOutput+wantCacheRead, row.CostUsd, 1e-6,
|
|
"the recorded cost must be every bucket the surface bills, cache reads included")
|
|
}
|
|
|
|
// TestStreamingOnGatewayTypedProvider drives the same streamed Anthropic call
|
|
// through a provider record whose catalog id names the OpenAI surface — the
|
|
// exact misconfiguration that hid the bug, since gateway records commonly pin
|
|
// one parser while the upstream serves another shape entirely.
|
|
//
|
|
// The router must choose the parser from the request path rather than the
|
|
// record's provider id, or the Anthropic usage block goes unread and input
|
|
// meters at zero all over again.
|
|
func TestStreamingOnGatewayTypedProvider(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
|
|
defer cancel()
|
|
|
|
env := provisionStreamingProvider(t, ctx, "openai_api")
|
|
|
|
sessionID := fmt.Sprintf("e2e-session-stream-gw-%d", time.Now().UnixNano())
|
|
code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID)
|
|
require.Equal(t, 200, code, "streamed chat through a gateway record must succeed; body: %s", body)
|
|
|
|
row := findAccessLogBySession(t, ctx, sessionID)
|
|
|
|
assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens),
|
|
"a record typed openai_api must still read the Anthropic usage block it is actually serving")
|
|
assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens),
|
|
"output tokens must survive the surface mismatch too")
|
|
assert.InDelta(t, float64(harness.VLLMStreamInputTokens)/1000*streamInRate, row.InputCostUsd, 1e-6,
|
|
"the request must be priced on the surface it spoke, not the one the record names")
|
|
}
|
|
|
|
// provisionStreamingProvider brings up the mock, one provider pointed at its
|
|
// streaming listener under the given catalog id, a policy authorising it, and a
|
|
// connected proxy + client.
|
|
func provisionStreamingProvider(t *testing.T, ctx context.Context, catalogID string) pricedEnv {
|
|
t.Helper()
|
|
|
|
vllm, err := harness.StartVLLM(ctx, srv)
|
|
require.NoError(t, err, "start mock upstream")
|
|
t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
|
|
|
|
name := "stream-" + catalogID
|
|
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-" + 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-" + name + "-client",
|
|
Type: "reusable",
|
|
ExpiresIn: 86400,
|
|
UsageLimit: 0,
|
|
AutoGroups: []string{grp.Id},
|
|
Ephemeral: &ephemeral,
|
|
})
|
|
require.NoError(t, err, "mint setup key")
|
|
// Deleting the group does not delete the key it auto-joins, so the key
|
|
// needs a cleanup of its own.
|
|
t.Cleanup(func() { _ = srv.API().SetupKeys.Delete(context.Background(), sk.Id) })
|
|
require.NotEmpty(t, sk.Key, "setup key plaintext")
|
|
|
|
dummyKey := "sk-stream-e2e"
|
|
cacheRead := streamCacheReadRate
|
|
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
|
Name: name,
|
|
ProviderId: catalogID,
|
|
UpstreamUrl: vllm.StreamURL,
|
|
ApiKey: &dummyKey,
|
|
Enabled: ptr(true),
|
|
Models: &[]api.AgentNetworkProviderModel{{
|
|
Id: streamedModel,
|
|
InputPer1k: streamInRate,
|
|
OutputPer1k: streamOutRate,
|
|
CacheReadPer1k: &cacheRead,
|
|
}},
|
|
})
|
|
require.NoError(t, err, "create provider")
|
|
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
|
|
|
enabled := true
|
|
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
|
Name: "e2e-" + name,
|
|
Enabled: &enabled,
|
|
SourceGroups: []string{grp.Id},
|
|
DestinationProviderIds: []string{prov.Id},
|
|
Limits: &api.AgentNetworkPolicyLimits{
|
|
TokenLimit: api.AgentNetworkPolicyTokenLimit{
|
|
Enabled: true,
|
|
GroupCap: 10_000_000,
|
|
UserCap: 10_000_000,
|
|
WindowSeconds: 60,
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err, "create policy")
|
|
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
|
|
|
endpoint, proxyIP, cl, px := connectClient(t, ctx, name, sk.Key)
|
|
return pricedEnv{
|
|
providerID: prov.Id,
|
|
groupID: grp.Id,
|
|
policyID: pol.Id,
|
|
upstream: vllm.StreamURL,
|
|
endpoint: endpoint,
|
|
proxyIP: proxyIP,
|
|
client: cl,
|
|
proxy: px,
|
|
}
|
|
}
|
|
|
|
// chatStreamUntil drives one streamed chat, retrying to absorb the tunnel and
|
|
// DNS jitter a first call through a fresh peer can hit.
|
|
func chatStreamUntil(t *testing.T, ctx context.Context, env pricedEnv, kind, model, sessionID string) (int, string) {
|
|
t.Helper()
|
|
var code int
|
|
var body string
|
|
deadline := time.Now().Add(90 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
c, b, cerr := env.client.ChatStream(ctx, env.endpoint, env.proxyIP, kind, model, "Reply with exactly: pong", sessionID)
|
|
if cerr == nil {
|
|
code, body = c, b
|
|
if code == 200 {
|
|
break
|
|
}
|
|
}
|
|
if !waitBeforeRetry(ctx, 5*time.Second) {
|
|
break
|
|
}
|
|
}
|
|
if code != 200 {
|
|
t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background()))
|
|
}
|
|
return code, body
|
|
}
|