diff --git a/e2e/agentnetwork/chat_test.go b/e2e/agentnetwork/chat_test.go index 17ed40d5f..87f16afd8 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,19 @@ 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}, - } + // 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. Vertex is + // registered with the raw "@version" id the Google console documents — the + // router normalizes the registered id, so the operator can paste either + // form; keeping the raw form here guards that normalization end to end. + 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..35ab27611 100644 --- a/e2e/agentnetwork/guardrail_test.go +++ b/e2e/agentnetwork/guardrail_test.go @@ -125,7 +125,15 @@ 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 an operator copies from the + // Google console; the guardrail must normalize the entry to match the + // version-stripped request model. 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 +190,25 @@ 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 + } + // The unversioned model id — the shape an Anthropic SDK client sends + // when configured without an "@version" suffix (customer-reported) — + // must pass the same allowlist entry and resolve upstream. + assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, catalogModel(pc)), + "unversioned model id must be permitted for %s", pc.name) + // Token counting: the URL carries the literal "count-tokens" + // pseudo-model and the real model travels in the body (Claude Code + // sends one such call per request). Allowed body model → 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..b04c0ab3e 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -231,6 +231,17 @@ 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 over the +// tunnel. The URL carries the literal "count-tokens" pseudo-model and the real +// model travels in the body — the one Vertex shape whose model is body-routed, +// sent by Claude Code before each request. The proxy must resolve the body +// model for the guardrail 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