diff --git a/e2e/agentnetwork/custom_pricing_test.go b/e2e/agentnetwork/custom_pricing_test.go index 5d6e6b21b..7c13cbdde 100644 --- a/e2e/agentnetwork/custom_pricing_test.go +++ b/e2e/agentnetwork/custom_pricing_test.go @@ -669,3 +669,47 @@ func inDelta(a, b, tol float64) bool { } return d <= tol } + +// TestCustomDatedModelKeepsItsOwnPrice covers the review fix that anchored the +// release-date fallback to Claude ids. Pricing looks every model up through +// that helper, so while it matched a bare trailing date any operator id ending +// in eight digits inherited the rate of its undated sibling — a silent +// mis-bill on models NetBird knows nothing about. +func TestCustomDatedModelKeepsItsOwnPrice(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + const ( + baseModel = "internal-llm" + datedModel = "internal-llm-20250101" + baseIn = 0.010 + baseOut = 0.020 + // An order of magnitude apart, so a row billed at the wrong entry is + // unmistakable rather than a rounding argument. + datedIn = 0.100 + datedOut = 0.200 + ) + + env := provisionPricedProvider(t, ctx, "customdated", []api.AgentNetworkProviderModel{ + {Id: baseModel, InputPer1k: baseIn, OutputPer1k: baseOut}, + {Id: datedModel, InputPer1k: datedIn, OutputPer1k: datedOut}, + }) + + t.Run("the undated id bills at its own rate", func(t *testing.T) { + session := fmt.Sprintf("e2e-session-customdated-base-%d", time.Now().UnixNano()) + chatOnce(t, ctx, env, baseModel, session) + assertOpenAICostAtRates(t, findAccessLogBySession(t, ctx, session), baseIn, baseOut) + }) + + t.Run("the dated id keeps its own rate", func(t *testing.T) { + session := fmt.Sprintf("e2e-session-customdated-dated-%d", time.Now().UnixNano()) + chatOnce(t, ctx, env, datedModel, session) + row := findAccessLogBySession(t, ctx, session) + assertOpenAICostAtRates(t, row, datedIn, datedOut) + + // Spelled out because it is the regression: inheriting the sibling's + // rate would bill this request at a tenth of its price. + assert.Greater(t, row.InputCostUsd, float64(vllmPromptTokens)/1000*baseIn*2, + "a custom dated id must not inherit the undated entry's rate") + }) +} diff --git a/e2e/agentnetwork/gateway_review_test.go b/e2e/agentnetwork/gateway_review_test.go new file mode 100644 index 000000000..556bc4a53 --- /dev/null +++ b/e2e/agentnetwork/gateway_review_test.go @@ -0,0 +1,242 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "fmt" + "strings" + "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" +) + +// The cases in this file cover behaviour that arrived from code review, after +// the gateway-protocol end-to-end tests were written. Each had unit coverage +// only; none needed a new harness capability, which is why they belong here +// rather than on a manual checklist. + +// TestNonInferenceEndpointsAreAuthorised covers the two review findings on the +// endpoints that carry no body: the per-model lookup must be authorised +// against the same allowlist that bounds the listing beside it, and only a read +// method may claim the non-inference exemption that skips the token pre-flight. +func TestNonInferenceEndpointsAreAuthorised(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + env := provisionDiscoveryProvider(t, ctx) + + t.Run("lookup of an authorised model succeeds", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMModel, nil) + }, 200) + assert.Equal(t, 200, code, "an allowlisted model must remain reachable; body: %s", body) + }) + + t.Run("lookup of an unauthorised model is refused", func(t *testing.T) { + code, body, err := env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMUnlistedModel, nil) + require.NoError(t, err, "request must reach the proxy") + assert.Equal(t, 403, code, + "a model the policy does not authorise must not be confirmed by the detail lookup; body: %s", body) + }) + + // A write must not claim the exemption that lets the listing skip the token + // pre-flight. The body names no model on purpose: that is what a request + // probing for the exemption looks like, and it is the case the method gate + // exists to refuse. (A POST that does name a model is a different thing — + // it routes and meters as the inference request it is.) + for _, path := range []string{"/v1/models", "/v1/models/" + harness.VLLMModel, "/api/hello"} { + t.Run("write to "+path+" is refused", func(t *testing.T) { + code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path, + `{"messages":[{"role":"user","content":"hi"}]}`, nil) + require.NoError(t, err, "request must reach the proxy") + assert.NotEqual(t, 200, code, + "a write to a non-inference path must not be served unmetered; body: %s", body) + }) + } + + // A request carrying the sub-agent attribution headers must still be served + // and metered normally. Asserting the ids themselves is not possible yet: + // the parser lifts them onto the request's metadata, but nothing persists + // them, so they have no queryable surface to check against. + t.Run("sub-agent headers do not disturb the request", func(t *testing.T) { + sessionID := fmt.Sprintf("e2e-session-agentid-%d", time.Now().UnixNano()) + code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/chat/completions", + fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"Reply with exactly: pong"}]}`, harness.VLLMModel), + []string{ + "x-session-id: " + sessionID, + "x-claude-code-agent-id: agent-child-7", + "x-claude-code-parent-agent-id: agent-root-1", + }) + require.NoError(t, err, "request must reach the proxy") + require.Equal(t, 200, code, "the request must succeed; body: %s", body) + + row := findAccessLogBySession(t, ctx, sessionID) + assert.Positive(t, row.InputTokens, "the request must still be metered normally") + }) +} + +// TestDatedModelIdRouting covers both halves of the dated-id rule that review +// tightened: a dated id still reaches an undated registration, but a route +// pinned to one dated build must never serve a different one. +func TestDatedModelIdRouting(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + const ( + undated = "claude-sonnet-9" + datedA = "claude-sonnet-9-20250101" + datedB = "claude-sonnet-9-20250202" + ) + + t.Run("a dated id reaches its undated registration", func(t *testing.T) { + env := provisionModelProvider(t, ctx, "dated-undated", "anthropic_api", undated) + + sessionID := fmt.Sprintf("e2e-session-dated-%d", time.Now().UnixNano()) + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", sessionID) + }, 200) + require.Equal(t, 200, code, "a pinned release of a registered family must route; body: %s", body) + + row := findAccessLogBySession(t, ctx, sessionID) + assert.Positive(t, row.InputTokens, "the dated request must price at the registered rate, not zero") + }) + + t.Run("a route pinned to one dated build refuses another", func(t *testing.T) { + env := provisionModelProvider(t, ctx, "dated-pinned", "anthropic_api", datedA) + + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", "") + }, 200) + require.Equal(t, 200, code, "the exact dated id must still route; body: %s", body) + + code, body, err := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedB, "Reply with exactly: pong", "") + require.NoError(t, err, "request must reach the proxy") + assert.Equal(t, 403, code, + "a provider pinned to one dated build must not serve another; body: %s", body) + }) +} + +// TestBedrockInferenceProfilesReachTheUpstream covers the startup lookup a +// Bedrock client makes. The proxy forwards it to the configured upstream rather +// than denying it, so what comes back is the upstream's answer — never a +// NetBird policy rejection. +func TestBedrockInferenceProfilesReachTheUpstream(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + env := provisionModelProvider(t, ctx, "infprofiles", "bedrock_api", "anthropic.claude-sonnet-5") + + code, body := callUntil(t, func() (int, string, error) { + return env.client.Get(ctx, env.endpoint, env.proxyIP, "/inference-profiles", nil) + }, 200) + + assert.Equal(t, 200, code, "the lookup must reach the upstream; body: %s", body) + assert.NotContains(t, body, "llm_policy.", + "the proxy must not answer a control-plane lookup with a policy denial") + assert.Contains(t, body, "inferenceProfileSummaries", + "the upstream's own answer must come back untouched") +} + +// provisionDiscoveryProvider brings up one mock-backed provider enumerating a +// single model, with an allowlist guardrail in effect, plus a connected client. +func provisionDiscoveryProvider(t *testing.T, ctx context.Context) pricedEnv { + t.Helper() + env := provisionModelProvider(t, ctx, "noninference", "openai_api", harness.VLLMModel) + + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-noninference-allowlist-" + fmt.Sprint(time.Now().UnixNano()) + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel} + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) }) + + enabled := true + _, err = srv.UpdatePolicy(ctx, env.policyID, api.AgentNetworkPolicyRequest{ + Name: "e2e-noninference", + Enabled: &enabled, + SourceGroups: []string{env.groupID}, + DestinationProviderIds: []string{env.providerID}, + GuardrailIds: &[]string{guard.Id}, + }) + require.NoError(t, err, "attach guardrail to policy") + return env +} + +// provisionModelProvider brings up the mock, one provider under the given +// catalog id enumerating exactly one model, an authorising policy, and a +// connected proxy + client. +func provisionModelProvider(t *testing.T, ctx context.Context, name, catalogID, model string) pricedEnv { + t.Helper() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + suffix := strings.ToLower(name) + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gwr-" + suffix}) + 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-gwr-" + suffix + "-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") + + dummyKey := "sk-gwr-e2e" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-gwr-" + suffix, + ProviderId: catalogID, + UpstreamUrl: vllm.URL, + ApiKey: &dummyKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{ + {Id: model, InputPer1k: 0.001, OutputPer1k: 0.002}, + }, + }) + 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-gwr-" + suffix, + 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, "gwr-"+suffix, sk.Key) + return pricedEnv{ + providerID: prov.Id, + groupID: grp.Id, + policyID: pol.Id, + upstream: vllm.URL, + endpoint: endpoint, + proxyIP: proxyIP, + client: cl, + proxy: px, + } +} diff --git a/e2e/agentnetwork/streaming_test.go b/e2e/agentnetwork/streaming_test.go new file mode 100644 index 000000000..2f7e9e3e2 --- /dev/null +++ b/e2e/agentnetwork/streaming_test.go @@ -0,0 +1,199 @@ +//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 + 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") + assert.Greater(t, row.CostUsd, 0.0, "a streamed request must never record as free") +} + +// 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") + 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 + } + } + time.Sleep(5 * time.Second) + } + if code != 200 { + t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background())) + } + return code, body +} diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 771ccc54d..9e9e7b34a 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -293,6 +293,27 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID)) } +// ChatStream is Chat with "stream": true in the request body, so the proxy's +// request parser marks the call as streaming and its response parser takes the +// SSE accumulator rather than the buffered-body path. Pair it with a provider +// pointed at VLLM.StreamURL, which answers every request as an event stream. +func (cl *Client) ChatStream(ctx context.Context, endpoint, proxyIP, 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":2048,"stream":true,"messages":[{"role":"user","content":%q}]}`, model, prompt) + default: + path = "/v1/chat/completions" + // include_usage is what makes a real OpenAI stream emit its final usage + // frame; without it the last chunk carries no tokens at all. + body = fmt.Sprintf(`{"model":%q,"stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":%q}]}`, model, prompt) + } + return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID)) +} + // Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike // Chat, the model is carried in the request path (project/region/model), so the // proxy routes by path and mints the service-account OAuth token; the body uses diff --git a/e2e/harness/vllm.go b/e2e/harness/vllm.go index 1068a9df7..cf9316325 100644 --- a/e2e/harness/vllm.go +++ b/e2e/harness/vllm.go @@ -18,6 +18,9 @@ const ( vllmImage = "nginx:alpine" vllmAlias = "vllm" vllmPort = "8000/tcp" + // vllmStreamPort serves the same wire shapes as an SSE stream. See the + // nginx config for why streaming lives on its own listener. + vllmStreamPort = "8001/tcp" // VLLMModel is the served model id the mock advertises and echoes back. It // matches a real small model commonly served by vLLM so the provider's // enumerated model and the client's request line up. @@ -42,6 +45,20 @@ const ( VLLMMessagesOutputTokens = 3 ) +// Token counts the streaming surface reports. They differ from the +// non-streaming ones on purpose: a test that asserts these numbers proves the +// SSE accumulator ran, rather than a buffered JSON body having been parsed. +// +// Input and cache-read arrive on message_start; output arrives on +// message_delta and supersedes the seed value message_start carries. Any +// parser that cannot read message_start reports zero input tokens — which is +// exactly the bug these counts exist to catch. +const ( + VLLMStreamInputTokens = 29 + VLLMStreamOutputTokens = 5 + VLLMStreamCacheReadTokens = 7 +) + // vllmNginxConf emulates a vLLM OpenAI-compatible server over plain HTTP (vLLM's // default: no TLS, port 8000), and additionally answers the wire shapes the // other catalog surfaces speak so one mock can stand in for every provider the @@ -86,11 +103,54 @@ http { location = /api/hello { return 200; } + location = /inference-profiles { + default_type application/json; + return 200 '{"inferenceProfileSummaries":[{"inferenceProfileId":"us.anthropic.claude-sonnet-5","status":"ACTIVE"}]}'; + } location / { default_type application/json; return 200 '{"id":"chatcmpl-e2e-vllm","object":"chat.completion","created":1700000000,"model":"Qwen/Qwen2.5-0.5B-Instruct","choices":[{"index":0,"message":{"role":"assistant","content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":2,"total_tokens":13}}'; } } + + # The streaming surface, on its own port so the response content type is a + # property of the listener rather than of a per-request branch: nginx sets + # Content-Type from default_type, which cannot be varied inside an "if", and + # a second Content-Type via add_header would leave the proxy reading the + # wrong one. A provider record pointed at this port streams every answer. + # + # Input and cache-read tokens ride message_start, output rides message_delta + # — the split that makes a stream different from a buffered body, and the + # reason a parser that ignores message_start meters input as zero. + server { + listen 8001; + location = /v1/messages { + default_type text/event-stream; + return 200 'event: message_start +data: {"type":"message_start","message":{"id":"msg_e2e_stream","type":"message","role":"assistant","model":"claude-sonnet-5","content":[],"usage":{"input_tokens":29,"output_tokens":1,"cache_read_input_tokens":7}}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}} + +event: message_stop +data: {"type":"message_stop"} + +'; + } + location / { + default_type text/event-stream; + return 200 'data: {"choices":[{"delta":{"content":"pong"}}]} + +data: {"choices":[],"usage":{"prompt_tokens":29,"completion_tokens":5,"total_tokens":34}} + +data: [DONE] + +'; + } + } } ` @@ -102,6 +162,10 @@ type VLLM struct { workDir string // URL is the upstream URL the vllm provider points at (http://:8000). URL string + // StreamURL is the same mock's streaming listener. A provider pointed here + // answers every request as SSE, so the proxy's streaming accumulator runs + // instead of its buffered-body parser. + StreamURL string } // StartVLLM runs the mock vLLM server on the shared network over plain HTTP. @@ -120,14 +184,17 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) { req := testcontainers.ContainerRequest{ Image: vllmImage, - ExposedPorts: []string{vllmPort}, + ExposedPorts: []string{vllmPort, vllmStreamPort}, Networks: []string{c.network.Name}, NetworkAliases: map[string][]string{c.network.Name: {vllmAlias}}, Cmd: []string{"nginx", "-c", "/conf/nginx.conf", "-g", "daemon off;"}, HostConfigModifier: func(hc *container.HostConfig) { hc.Binds = append(hc.Binds, workDir+":/conf:ro") }, - WaitingFor: wait.ForListeningPort(vllmPort).WithStartupTimeout(60 * time.Second), + WaitingFor: wait.ForAll( + wait.ForListeningPort(vllmPort), + wait.ForListeningPort(vllmStreamPort), + ).WithStartupTimeout(60 * time.Second), } ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ @@ -139,7 +206,12 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) { return nil, fmt.Errorf("start vllm container: %w", err) } - return &VLLM{container: ctr, workDir: workDir, URL: "http://" + vllmAlias + ":8000"}, nil + return &VLLM{ + container: ctr, + workDir: workDir, + URL: "http://" + vllmAlias + ":8000", + StreamURL: "http://" + vllmAlias + ":8001", + }, nil } // Logs returns the vLLM container logs, for diagnostics on failure.