mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-06 13:39:07 +02:00
[proxy,management] Conform the Agent Network endpoint to the LLM gateway protocol (#7154)
[proxy,management] Conform the Agent Network endpoint to the LLM gateway protocol Reviewed the proxy against Claude Code's published gateway contract. The transport layer already held up; fourteen gaps sat one layer up, in the model catalog and in the non-inference endpoints clients call. Two of them cost money. The catalog carried no claude-opus-5 or claude-sonnet-5, so an operator could not authorise the models coding agents default to — those requests denied as not-routable, or priced at zero where a catch-all carried them. And gateway records pin ParserID "openai" while the same record serves /v1/messages, so Anthropic responses were read with the OpenAI parser, which never looks at message_start where input tokens live: input metered as roughly zero on every stream and cost was skipped entirely. The rest fix requests refused for structural rather than policy reasons: model discovery denied for every account with a model allowlist, token counting denied on Bedrock and mis-parsed on Vertex, startup probes refused and written into the access log at every session start, and denials rendered in a shape no LLM client parses. Two changes are additive by design — the deny body keeps every field it had and adds the vendor's error object alongside, and body-level identity injection is now gated on the request's dialect so it stops sending OpenAI-shape fields into Anthropic bodies that reject them. The end-to-end work turned up one more: the discovery filter treated any slash in a model id as a gateway prefix, which would have dropped every self-hosted "Qwen/..." model from the picker.
This commit is contained in:
Maycon Santos
@@ -0,0 +1,209 @@
|
||||
//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
|
||||
}
|
||||
Reference in New Issue
Block a user