From 1b7d23ee9fc8ebde97092341bc2edcc84a6a4399 Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Wed, 22 Jul 2026 07:32:40 +0000 Subject: [PATCH] [e2e] cover unversioned Vertex models, versioned config ids, and count-tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the agent-network e2e to the Vertex request shapes a customer actually sends through Claude Code: - Register the Vertex provider's model under the raw "@version" id the Google console documents, so the router's candidate normalization is exercised end to end (previously the provider was created with no models array, leaving Vertex per-provider matching untested). - Allowlist the raw "@version" id in the guardrail, so entry normalization is exercised against version-stripped request models. - Drive the unversioned model id (…/models/claude-sonnet-4-5:rawPredict) through the tunnel and require 200 — the customer-reported shape. - Drive the token-count endpoint (…/models/count-tokens:rawPredict, real model in the body) via the new harness VertexCountTokens: an allowlisted body model must serve 200, a disallowed one must deny 403 before the upstream. --- e2e/agentnetwork/chat_test.go | 26 ++++++++++++-------------- e2e/agentnetwork/guardrail_test.go | 24 +++++++++++++++++++++++- e2e/harness/client.go | 9 +++++++++ 3 files changed, 44 insertions(+), 15 deletions(-) diff --git a/e2e/agentnetwork/chat_test.go b/e2e/agentnetwork/chat_test.go index 17ed40d5f..809f5be21 100644 --- a/e2e/agentnetwork/chat_test.go +++ b/e2e/agentnetwork/chat_test.go @@ -104,8 +104,9 @@ func availableProviders() []providerCase { } // providerRequest builds a create request for a matrix provider: enabled, with -// a uniquely-priced model for body-routed providers and none for the -// path-routed Vertex (whose model lives in the request path). +// a uniquely-priced model registered under the id an operator would paste — +// the normalized catalog id for Bedrock, the raw form elsewhere (including the +// "@version" Vertex id, which the router must normalize to route). func providerRequest(pc providerCase) api.AgentNetworkProviderRequest { req := api.AgentNetworkProviderRequest{ Name: pc.name, @@ -114,18 +115,15 @@ func providerRequest(pc providerCase) api.AgentNetworkProviderRequest { ApiKey: &pc.apiKey, Enabled: ptr(true), } - if pc.kind != harness.WireVertex { - // The router matches the normalized catalog id. Bedrock's request model - // travels as a region-prefixed inference-profile id in the URL path - // (us.anthropic...), which the router strips before matching, so register - // the normalized form here or routing fails as model_not_routable. - modelID := pc.model - if pc.kind == harness.WireBedrock { - modelID = catalogModel(pc) - } - req.Models = &[]api.AgentNetworkProviderModel{ - {Id: modelID, InputPer1k: 0.001, OutputPer1k: 0.002}, - } + // Register Bedrock under its normalized catalog id (the router strips the + // region prefix before matching); other providers, incl. the raw "@version" + // Vertex id, register as-is so the router's normalization is exercised. + modelID := pc.model + if pc.kind == harness.WireBedrock { + modelID = catalogModel(pc) + } + req.Models = &[]api.AgentNetworkProviderModel{ + {Id: modelID, InputPer1k: 0.001, OutputPer1k: 0.002}, } return req } diff --git a/e2e/agentnetwork/guardrail_test.go b/e2e/agentnetwork/guardrail_test.go index d4fe98dda..283f1c368 100644 --- a/e2e/agentnetwork/guardrail_test.go +++ b/e2e/agentnetwork/guardrail_test.go @@ -125,7 +125,13 @@ func TestModelAllowlistEnforced(t *testing.T) { require.NoError(t, perr, "create provider %s", pc.name) id := prov.Id ids = append(ids, id) - allowed = append(allowed, catalogModel(pc)) + // Vertex allowlists the raw "@version" id (the guardrail normalizes it); + // other providers allowlist the normalized catalog id. + if pc.kind == harness.WireVertex { + allowed = append(allowed, pc.model) + } else { + allowed = append(allowed, catalogModel(pc)) + } t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) }) } @@ -182,6 +188,22 @@ func TestModelAllowlistEnforced(t *testing.T) { // the upstream), regardless of whether it is a real catalog model. assert.Equal(t, 403, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, disallowedModel(pc)), "model outside the allowlist must be denied for %s", pc.name) + + if pc.kind != harness.WireVertex { + return + } + // Unversioned model id (the customer-reported shape) must pass the + // same allowlist entry. + assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, catalogModel(pc)), + "unversioned model id must be permitted for %s", pc.name) + // count-tokens carries the real model in the body: allowed → served. + code, _, err := cl.VertexCountTokens(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model) + require.NoError(t, err, "count-tokens request must reach the proxy for %s", pc.name) + assert.Equal(t, 200, code, "count-tokens with an allowlisted body model must be permitted for %s", pc.name) + // Disallowed body model → denied before the upstream. + code, _, err = cl.VertexCountTokens(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, disallowedModel(pc)) + require.NoError(t, err, "count-tokens request must reach the proxy for %s", pc.name) + assert.Equal(t, 403, code, "count-tokens with a body model outside the allowlist must be denied for %s", pc.name) }) } } diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 19210349f..1d6385d53 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -231,6 +231,15 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID)) } +// VertexCountTokens issues the Anthropic-on-Vertex token-count POST: the URL +// carries the "count-tokens" pseudo-model and the real model travels in the +// body, so the proxy must resolve the body model for the allowlist and routing. +func (cl *Client) VertexCountTokens(ctx context.Context, endpoint, proxyIP, project, region, model string) (int, string, error) { + path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/count-tokens:rawPredict", project, region) + body := fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hi"}]}`, model) + return cl.post(ctx, endpoint, proxyIP, path, body, nil) +} + // Bedrock issues a native AWS Bedrock InvokeModel POST over the tunnel. The // model id is carried in the request path (/model/{id}/invoke), so the proxy // routes by path; the body uses the bedrock anthropic_version rather than a