fix: vertex count tokens unsupported fields

core/providers/vertex/payload_ordering_test.go

@@ -13,3 +13,25 @@ func TestStripVertexGeminiUnsupportedFieldsRawPreservesOrdering(t *testing.T) {
 
 	assert.Equal(t, `{"contents":[{"role":"user","parts":[{"functionCall":{"name":"lookup","args":{"z":1,"a":2}}},{"functionResponse":{"name":"lookup","response":{"output":{"z":1,"a":2}}}}]}],"generationConfig":{"temperature":0.2}}`, string(got))
 }
+
+func TestStripVertexCountTokensUnsupportedFields(t *testing.T) {
+	t.Run("keeps fields that contribute to the count", func(t *testing.T) {
+		raw := []byte(`{"model":"gemini-3.6-flash","contents":[{"role":"user","parts":[{"text":"hi"}]}],"systemInstruction":{"parts":[{"text":"be terse"}]},"tools":[{"functionDeclarations":[{"name":"probe"}]}],"generationConfig":{"temperature":0.2}}`)
+
+		got := stripVertexCountTokensUnsupportedFields(raw)
+
+		assert.Equal(t, string(raw), string(got))
+	})
+
+	t.Run("drops fields countTokens rejects", func(t *testing.T) {
+		raw := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"systemInstruction":{"parts":[{"text":"be terse"}]},"toolConfig":{"functionCallingConfig":{"mode":"AUTO"}},"safetySettings":[{"category":"HARM_CATEGORY_HARASSMENT"}],"cachedContent":"cachedContents/abc","serviceTier":"SERVICE_TIER_STANDARD","labels":{"a":"b"}}`)
+
+		got := stripVertexCountTokensUnsupportedFields(raw)
+
+		assert.JSONEq(t, `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"systemInstruction":{"parts":[{"text":"be terse"}]}}`, string(got))
+	})
+
+	t.Run("handles empty body", func(t *testing.T) {
+		assert.Empty(t, stripVertexCountTokensUnsupportedFields(nil))
+	})
+}

core/providers/vertex/vertex.go

@@ -4025,6 +4025,33 @@ func (provider *VertexProvider) fileContentByKey(ctx *schemas.BifrostContext, ke
 	}, nil
 }
 
+// vertexCountTokensUnsupportedFields lists generateContent fields that Vertex's
+// CountTokensRequest does not define, and which it rejects with a 400.
+var vertexCountTokensUnsupportedFields = []string{
+	"toolConfig",
+	"safetySettings",
+	"cachedContent",
+	"serviceTier",
+	"labels",
+}
+
+// stripVertexCountTokensUnsupportedFields drops fields the countTokens endpoint rejects.
+// systemInstruction, tools and generationConfig are supported there and must survive —
+// they contribute to the token count.
+func stripVertexCountTokensUnsupportedFields(jsonBody []byte) []byte {
+	if len(jsonBody) == 0 {
+		return jsonBody
+	}
+
+	out := jsonBody
+	for _, field := range vertexCountTokensUnsupportedFields {
+		if updated, err := providerUtils.DeleteJSONField(out, field); err == nil {
+			out = updated
+		}
+	}
+	return out
+}
+
 // CountTokens counts the number of tokens in the provided content using Vertex AI's countTokens endpoint.
 // Supports Gemini models with both text and image content.
 func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostResponsesRequest) (*schemas.BifrostCountTokensResponse, *schemas.BifrostError) {
@@ -4071,10 +4098,7 @@ func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key sch
 		// Skip field-stripping when large payload mode is active — jsonBody is nil
 		// and the raw body will stream directly from the ingress reader.
 		if jsonBody != nil {
-			// Use sjson to delete fields directly from JSON bytes, preserving key ordering
-			jsonBody, _ = providerUtils.DeleteJSONField(jsonBody, "toolConfig")
-			jsonBody, _ = providerUtils.DeleteJSONField(jsonBody, "generationConfig")
-			jsonBody, _ = providerUtils.DeleteJSONField(jsonBody, "systemInstruction")
+			jsonBody = stripVertexCountTokensUnsupportedFields(jsonBody)
 		}
 	}