Skip to content

test(bedrock): let monkeypatch own bedrock_request_metadata_fields #37840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuneng-berri merged 1 commit into from
Aug 22, 2026
Merged

test(bedrock): let monkeypatch own bedrock_request_metadata_fields #37840

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion test-quality-budget.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
"limit": 768
},
"TQ005": {
"limit": 2832
"limit": 2810
},
"TQ006": {
"limit": 34
Expand Down
105 changes: 57 additions & 48 deletions tests/test_litellm/llms/bedrock/test_request_metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,6 @@
IDENTITY = {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform"}


@pytest.fixture(autouse=True)
def reset_setting():
previous = litellm.bedrock_request_metadata_fields
yield
litellm.bedrock_request_metadata_fields = previous


def litellm_params(metadata_key, **metadata):
return {metadata_key: dict(metadata)}

Expand Down Expand Up @@ -73,8 +66,8 @@ def converse_body_async(litellm_params_value, optional_params=None):


@pytest.mark.parametrize("setting", [None, []])
def test_feature_off_by_default_leaves_body_and_headers_untouched(setting):
litellm.bedrock_request_metadata_fields = setting
def test_feature_off_by_default_leaves_body_and_headers_untouched(setting, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting)
params = litellm_params("metadata", spend_logs_metadata={"team": "x"}, **IDENTITY)

assert "requestMetadata" not in converse_body(params)
Expand All @@ -88,18 +81,18 @@ def test_feature_off_by_default_leaves_body_and_headers_untouched(setting):


@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_resolver_reads_both_metadata_variable_names(metadata_key):
def test_resolver_reads_both_metadata_variable_names(metadata_key, monkeypatch: pytest.MonkeyPatch):
"""`/v1/chat/completions` populates `metadata`; the LITELLM_METADATA_ROUTES populate
`litellm_metadata`. Reading only one silently forwards nothing on the other route."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
params = litellm_params(metadata_key, spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)

assert converse_body(params)["requestMetadata"] == {**IDENTITY, "cost_center": "cc-1"}


@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key):
litellm.bedrock_request_metadata_fields = ALL_FIELDS
def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
params = litellm_params(metadata_key, **IDENTITY)

headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
Expand All @@ -112,10 +105,15 @@ def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key)
@pytest.mark.parametrize("reverse_client_keys", [False, True])
@pytest.mark.parametrize("field_order", [ALL_FIELDS, list(reversed(ALL_FIELDS))])
@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"])
def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, field_order, client_source):
def test_identity_survives_a_caller_filling_every_slot(
reverse_client_keys,
field_order,
client_source,
monkeypatch: pytest.MonkeyPatch,
):
"""A caller sending 16 keys of its own must not evict the identity the feature exists to
produce. Driven over every input ordering so the invariant is not an accident of one."""
litellm.bedrock_request_metadata_fields = field_order
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", field_order)
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS)]
client_pairs = {key: "v" for key in (reversed(client_keys) if reverse_client_keys else client_keys)}
if client_source == "spend_logs_metadata":
Expand All @@ -141,11 +139,14 @@ def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, fiel
["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata", "user_api_key_team_alias"],
],
)
def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field_order):
def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(
field_order,
monkeypatch: pytest.MonkeyPatch,
):
"""An operator repeating a field in YAML must not inflate the reserved count and shrink the
client budget. Asserts the client keys that should have fitted actually reach the wire, since
asserting only that identity survives passes with or without the deduplication."""
litellm.bedrock_request_metadata_fields = field_order
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", field_order)
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - 1)]
params = litellm_params("metadata", spend_logs_metadata={key: "v" for key in client_keys}, **IDENTITY)

Expand All @@ -162,11 +163,15 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field
"forged_key",
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
)
def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, client_source):
def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(
forged_key,
client_source,
monkeypatch: pytest.MonkeyPatch,
):
"""`user_api_key_org_alias` and `user_api_key_hash` are names the proxy does not set here,
so an exact-key reservation would let the forged value through under a name that reads as
proxy-authoritative in the AWS billing record."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
forged = {forged_key: "attacker-controlled"}
if client_source == "spend_logs_metadata":
params, optional_params = litellm_params("metadata", spend_logs_metadata=forged, **IDENTITY), {}
Expand All @@ -179,10 +184,10 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, clien
assert "attacker-controlled" not in resolved.values()


def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds():
def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch):
"""A team alias with an apostrophe must not turn a working request into a 400 the moment
an operator flips the setting on."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
params = litellm_params(
"metadata",
user_api_key_alias="prod-key",
Expand All @@ -196,8 +201,8 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe
assert body["messages"]


def test_caller_supplied_violation_still_raises_bad_request():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

with pytest.raises(litellm.exceptions.BadRequestError):
converse_body(
Expand All @@ -206,34 +211,34 @@ def test_caller_supplied_violation_still_raises_bad_request():
)


def test_non_string_and_absent_identity_values_are_dropped():
litellm.bedrock_request_metadata_fields = ALL_FIELDS + ["user_api_key_spend"]
def test_non_string_and_absent_identity_values_are_dropped(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS + ["user_api_key_spend"])
params = litellm_params("metadata", user_api_key_alias="prod-key", user_api_key_spend=1.25)

assert converse_body(params)["requestMetadata"] == {"user_api_key_alias": "prod-key"}


def test_email_is_separately_opt_in():
def test_email_is_separately_opt_in(monkeypatch: pytest.MonkeyPatch):
"""PII crossing into CloudTrail only when the operator names the field."""
identity_with_email = {**IDENTITY, "user_api_key_user_email": "owner@example.com"}
litellm.bedrock_request_metadata_fields = ["user_api_key_alias", "user_api_key_team_alias"]
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_alias", "user_api_key_team_alias"])
assert (
"user_api_key_user_email"
not in converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"]
)

litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
assert converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] == identity_with_email


def test_resolver_returns_none_when_nothing_survives():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
def test_resolver_returns_none_when_nothing_survives(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
assert resolve_bedrock_request_metadata(litellm_params=None) is None
assert resolve_bedrock_request_metadata(litellm_params={"metadata": {"unrelated": "x"}}) is None


def test_invoke_header_is_json_encoded_and_signed():
litellm.bedrock_request_metadata_fields = ALL_FIELDS
def test_invoke_header_is_json_encoded_and_signed(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
params = litellm_params("metadata", spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)

headers = AmazonInvokeConfig().validate_environment(
Expand All @@ -250,10 +255,10 @@ def test_invoke_header_is_json_encoded_and_signed():
assert "anthropic-version" not in signed


def test_a_caller_supplied_guardrail_header_still_wins():
def test_a_caller_supplied_guardrail_header_still_wins(monkeypatch: pytest.MonkeyPatch):
"""The no-displace rule is deliberate for the guardrail headers and must survive the
request-metadata header becoming proxy-owned."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

headers = AmazonInvokeConfig().validate_environment(
headers={"X-Amzn-Bedrock-GuardrailIdentifier": "caller-set"},
Expand Down Expand Up @@ -318,10 +323,10 @@ def metadata_header_values(headers):
return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()]


def test_converse_still_sets_the_bearer_authorization_header():
def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch):
"""Converse owns the metadata header now, and that must not disturb the api_key path its
validate_environment existed for. Closing the forgery hole cannot break authentication."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

headers = AmazonConverseConfig().validate_environment(
headers={},
Expand All @@ -341,11 +346,11 @@ def test_converse_still_sets_the_bearer_authorization_header():
"caller_header_name",
[BEDROCK_REQUEST_METADATA_HEADER, BEDROCK_REQUEST_METADATA_HEADER.lower(), "x-AMZN-bedrock-Request-METADATA"],
)
def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name):
def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name, monkeypatch: pytest.MonkeyPatch):
"""`extra_headers` puts caller-supplied names into the same dict the proxy merges into, so a
deferring merge would sign the caller's forged identity into the AWS billing record. Every
spelling must lose, or a second variant is left for the transport to choose between."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

headers = driver({caller_header_name: FORGED}, litellm_params("metadata", **IDENTITY))

Expand All @@ -355,11 +360,11 @@ def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header


@pytest.mark.parametrize("driver", HEADER_DRIVERS)
def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver):
def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver, monkeypatch: pytest.MonkeyPatch):
"""Forwarding enabled but nothing resolvable, which a caller can arrange by supplying values
that all fail Bedrock's rules. Owned-but-empty must mean no header on the wire, never a
fallback to the caller's."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
unresolvable = litellm_params("metadata", user_api_key_alias="O'Brien's key", user_api_key_team_alias="x" * 300)

headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, unresolvable)
Expand All @@ -373,11 +378,15 @@ def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(drive
"forged_key",
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
)
def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(forged_key, driver):
def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(
forged_key,
driver,
monkeypatch: pytest.MonkeyPatch,
):
"""The Converse body has the same fail-open shape as the header: with forwarding on and
nothing resolvable, leaving the caller's `requestMetadata` in place would keep their
reserved-prefix keys on the wire. Owned-but-empty must remove the field outright."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

body = driver(litellm_params("metadata"), {"requestMetadata": {forged_key: "FORGED"}})

Expand All @@ -386,10 +395,10 @@ def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothin


@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver):
def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver, monkeypatch: pytest.MonkeyPatch):
"""Removing the field must be scoped to the reserved keys being the only thing left, not a
blanket drop of the caller's own attribution pairs."""
litellm.bedrock_request_metadata_fields = ALL_FIELDS
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)

body = driver(
litellm_params("metadata"),
Expand All @@ -400,10 +409,10 @@ def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(dr


@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver):
def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver, monkeypatch: pytest.MonkeyPatch):
"""With the feature off the proxy does not own the field, so the pre-existing pass-through
behaviour for a caller-supplied `requestMetadata` must be unchanged."""
litellm.bedrock_request_metadata_fields = None
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)
caller_supplied = {"user_api_key_team_alias": "caller-set", "cost_center": "cc-9"}

body = driver(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_supplied})
Expand All @@ -412,10 +421,10 @@ def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver):


@pytest.mark.parametrize("driver", HEADER_DRIVERS)
def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver):
def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver, monkeypatch: pytest.MonkeyPatch):
"""The proxy only claims the name when the operator turned forwarding on; with the feature
off this is an ordinary passthrough header and stripping it would be a regression."""
litellm.bedrock_request_metadata_fields = None
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)

headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, litellm_params("metadata", **IDENTITY))

Expand Down
Loading