fix(sagemaker): keep AWS credentials out of the sagemaker_chat request body (#45724)

* fix(sagemaker): keep AWS credentials out of the sagemaker_chat request body

sagemaker_chat and sagemaker_nova built the container body from every optional
param, so aws_access_key_id, aws_secret_access_key, aws_session_token and the
other aws_* auth params, plus model_id and sagemaker_base_url, were sent in the
JSON body to the endpoint container. The request is still SigV4-signed with the
same credentials, which stay in optional_params for signing

* fix(sagemaker): limit the sagemaker_chat body filter to aws_ params

model_id and sagemaker_base_url are not credentials, so they stay in the
body as before and the change sticks to what the fix is about

* test(sagemaker): assert the whole sagemaker_chat request body

* fix(sagemaker): filter the sagemaker_chat body with the shared AWS auth param list

* test(sagemaker): cover the sagemaker_chat body on every endpoint, credential source, and outage

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/llms/sagemaker/chat/transformation.py

@@ -69,17 +69,20 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
         litellm_params: dict,  # mutable-ok: matches the base chat transform signature
         headers: dict,  # mutable-ok: matches the base chat transform signature
     ) -> dict:  # mutable-ok: the handler sends this body straight to httpx
-        request: Final = super().transform_request(
+        request: Final[dict[str, object]] = super().transform_request(  # mutable-ok: base transform returns a dict
             model=model,
             messages=messages,
             optional_params=optional_params,
             litellm_params=litellm_params,
             headers=headers,
         )
+        container_request: Final = {
+            key: value for key, value in request.items() if key not in self.aws_authentication_params
+        }
         served_model_name: Final = litellm_params.get("hf_model_name")
         if not isinstance(served_model_name, str):
-            return request
-        return {**request, "model": served_model_name}
+            return container_request
+        return {**container_request, "model": served_model_name}
 
     def get_complete_url(
         self,

tests/integration/providers/_sagemaker_chat_container.py

@@ -0,0 +1,132 @@
+import json
+import struct
+import uuid
+import zlib
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final
+
+from integration._support.wire import Reply, Request, Wire
+from pydantic import JsonValue, TypeAdapter
+
+ENDPOINT: Final = "integration-vllm-endpoint"
+SERVED_MODEL: Final = "integration-org/served-chat-model"
+ACCESS_KEY: Final = "AKIAINTEGRATION000003"
+SECRET_KEY: Final = "synthetic-secret-key-for-testing"
+REGION: Final = "us-east-1"
+PROMPT: Final = "synthetic inference component request"
+CONTENT: Final = ("sagemaker ", "wire control")
+NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}})
+JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MARKER_PREFIX: Final = "marker-"
+
+
+def answered_choices() -> list[JsonValue]:
+    return [
+        {
+            "finish_reason": "stop",
+            "index": 0,
+            "message": {"role": "assistant", "content": "".join(CONTENT)},
+            "provider_specific_fields": {},
+        }
+    ]
+
+
+def _completion(identity: str) -> bytes:
+    return json.dumps(
+        {
+            "id": identity,
+            "object": "chat.completion",
+            "created": 1,
+            "model": SERVED_MODEL,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": "".join(CONTENT)},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
+        }
+    ).encode()
+
+
+def _chunk(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None, **extra: JsonValue) -> bytes:
+    payload: Final = {
+        "id": identity,
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": SERVED_MODEL,
+        "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}],
+        **extra,
+    }
+    return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()
+
+
+def _event_header(name: str, value: str) -> bytes:
+    name_bytes: Final = name.encode()
+    value_bytes: Final = value.encode()
+    return struct.pack("!B", len(name_bytes)) + name_bytes + b"\x07" + struct.pack("!H", len(value_bytes)) + value_bytes
+
+
+def _payload_part(body: bytes) -> bytes:
+    headers: Final = (
+        _event_header(":event-type", "PayloadPart")
+        + _event_header(":content-type", "application/octet-stream")
+        + _event_header(":message-type", "event")
+    )
+    prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers))
+    message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body
+    return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF)
+
+
+def container_reply(identity: str, stream: bool) -> Reply:
+    if not stream:
+        return Reply(body=_completion(identity))
+    chunks: Final = (
+        _chunk(identity, {"role": "assistant", "content": CONTENT[0]}, None),
+        _chunk(identity, {"content": CONTENT[1]}, None),
+        _chunk(identity, {}, "stop", usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}),
+    )
+    return Reply(
+        chunks=tuple(_payload_part(chunk) for chunk in chunks), content_type="application/vnd.amazon.eventstream"
+    )
+
+
+def marked_prompt() -> str:
+    return f"{PROMPT} {_MARKER_PREFIX}{uuid.uuid4().hex}"
+
+
+def marker(text: str) -> str:
+    _, separator, rest = text.partition(_MARKER_PREFIX)
+    assert separator and text.count(_MARKER_PREFIX) == 1, text
+    return rest[:32]
+
+
+def container_body(url: str, prompt: str, **extra: JsonValue) -> dict[str, JsonValue]:
+    return {
+        "model": ENDPOINT,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 16,
+        "sagemaker_base_url": url,
+        **extra,
+    }
+
+
+def only_request(wire: Wire) -> Request:
+    (request,) = wire.drain()
+    assert (request.method, request.target) == ("POST", "/"), request
+    return request
+
+
+def by_marker(received: tuple[Request, ...]) -> Mapping[str, Request]:
+    forwarded: Final = {marker(request.body.decode()): request for request in received}
+    assert len(forwarded) == len(received), [request.body for request in received]
+    return forwarded
+
+
+def assert_signed(request: Request, access_key: str, region: str, session_token: str | None) -> None:
+    authorization: Final = request.headers["authorization"]
+    assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={access_key}/"), authorization
+    assert f"/{region}/sagemaker/aws4_request" in authorization, authorization
+    assert request.headers.get("x-amz-security-token") == session_token, dict(request.headers)

tests/integration/providers/test_sagemaker_chat_credential_sources.py

@@ -0,0 +1,222 @@
+import os
+import signal
+from collections.abc import Mapping
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+from urllib.parse import parse_qs
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.client import Gateway, eventually
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration.providers._sagemaker_chat_container import (
+    ACCESS_KEY,
+    ENDPOINT,
+    JSON_OBJECT,
+    NO_CACHE,
+    REGION,
+    SECRET_KEY,
+    answered_choices,
+    assert_signed,
+    by_marker,
+    container_body,
+    container_reply,
+    marked_prompt,
+    marker,
+)
+from pydantic import JsonValue
+
+_ROLE: Final = "arn:aws:iam::123456789012:role/integration-sagemaker-chat"
+_SOURCE_KEY: Final = "AKIAINTEGRATIONSRC001"
+_ENV_KEY: Final = "AKIAINTEGRATIONENV001"
+_PROFILE_KEY: Final = "AKIAINTEGRATIONPRO001"
+_ASSUMED_KEY: Final = "ASIAINTEGRATIONROLE01"
+_WEB_IDENTITY_KEY: Final = "ASIAINTEGRATIONWEB001"
+_PROFILE: Final = "integration-sagemaker-profile"
+_WEB_IDENTITY_VARIABLE: Final = "INTEGRATION_WIF_IDENTITY_TOKEN"
+_STARTED_WORKER: Final = "Started server process ["
+_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
+_SIGNERS: Final[Mapping[str, tuple[str, str | None]]] = MappingProxyType(
+    {
+        "static": (ACCESS_KEY, None),
+        "env": (_ENV_KEY, None),
+        "profile": (_PROFILE_KEY, None),
+        "role": (_ASSUMED_KEY, f"synthetic-session-{_ASSUMED_KEY}"),
+        "web-identity": (_WEB_IDENTITY_KEY, f"synthetic-session-{_WEB_IDENTITY_KEY}"),
+    }
+)
+
+
+def _sts_reply(action: str, result: str) -> Reply:
+    return Reply(
+        content_type="text/xml",
+        body=(
+            f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/"><{action}Result>{result}'
+            f"</{action}Result><ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata>"
+            f"</{action}Response>"
+        ).encode(),
+    )
+
+
+def _assumed(action: str, key: str) -> Reply:
+    return _sts_reply(
+        action,
+        f"<Credentials><AccessKeyId>{key}</AccessKeyId><SecretAccessKey>synthetic-assumed-secret-for-testing"
+        f"</SecretAccessKey><SessionToken>synthetic-session-{key}</SessionToken><Expiration>2035-01-01T00:00:00Z"
+        "</Expiration></Credentials><AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/integration/session"
+        "</Arn><AssumedRoleId>integration:session</AssumedRoleId></AssumedRoleUser><PackedPolicySize>0"
+        "</PackedPolicySize>",
+    )
+
+
+def _sts(request: Request) -> Reply:
+    parameters: Final = parse_qs(request.body.decode())
+    action: Final = parameters["Action"][0]
+    assert request.method == "POST", request
+    match action:
+        case "GetCallerIdentity":
+            return _sts_reply(
+                action,
+                "<Arn>arn:aws:iam::123456789012:user/integration-source</Arn><UserId>integration-source</UserId>"
+                "<Account>123456789012</Account>",
+            )
+        case "AssumeRole":
+            assert parameters["RoleArn"] == [_ROLE], parameters
+            return _assumed(action, _ASSUMED_KEY)
+        case "AssumeRoleWithWebIdentity":
+            assert parameters["WebIdentityToken"] == ["synthetic.web-identity.token"], parameters
+            return _assumed(action, _WEB_IDENTITY_KEY)
+        case _:
+            raise AssertionError(parameters)
+
+
+def _container(request: Request) -> Reply:
+    return container_reply(f"sagemaker-chat-{marker(request.body.decode())}", False)
+
+
+def _config(container: Wire, authority: Wire, tmp_path: Path) -> Path:
+    common: Final = {
+        "model": f"sagemaker_chat/{ENDPOINT}",
+        "aws_region_name": REGION,
+        "sagemaker_base_url": container.url,
+    }
+    role: Final = {"aws_role_name": _ROLE, "aws_session_name": "integration-session", "aws_sts_endpoint": authority.url}
+    sources: Final[Mapping[str, Mapping[str, str]]] = {
+        "static": {"aws_access_key_id": ACCESS_KEY, "aws_secret_access_key": SECRET_KEY},
+        "env": {},
+        "profile": {"aws_profile_name": _PROFILE},
+        "role": {"aws_access_key_id": _SOURCE_KEY, "aws_secret_access_key": SECRET_KEY, **role},
+        "web-identity": {"aws_web_identity_token": f"oidc/env/{_WEB_IDENTITY_VARIABLE}", **role},
+    }
+    base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    path: Final = tmp_path / "sagemaker-chat-credential-sources.yaml"
+    path.write_text(
+        yaml.safe_dump(
+            {
+                **base,
+                "model_list": [
+                    {"model_name": f"sagemaker-chat-{source}", "litellm_params": {**common, **extra}}
+                    for source, extra in sources.items()
+                ],
+            }
+        )
+    )
+    return path
+
+
+def _environment(authority: Wire, tmp_path: Path) -> Mapping[str, str]:
+    credentials: Final = tmp_path / "aws-credentials"
+    credentials.write_text(f"[{_PROFILE}]\naws_access_key_id = {_PROFILE_KEY}\naws_secret_access_key = {SECRET_KEY}\n")
+    empty: Final = tmp_path / "aws-config"
+    empty.write_text("")
+    return {
+        "AWS_ACCESS_KEY_ID": _ENV_KEY,
+        "AWS_SECRET_ACCESS_KEY": SECRET_KEY,
+        "AWS_SHARED_CREDENTIALS_FILE": str(credentials),
+        "AWS_CONFIG_FILE": str(empty),
+        "AWS_EC2_METADATA_DISABLED": "true",
+        "AWS_ENDPOINT_URL_STS": authority.url,
+        "AWS_DEFAULT_REGION": REGION,
+        _WEB_IDENTITY_VARIABLE: "synthetic.web-identity.token",
+    }
+
+
+def _ask(owned: OwnedProxy, source: str) -> str:
+    prompt: Final = marked_prompt()
+    body: Final[Mapping[str, JsonValue]] = {
+        "model": f"sagemaker-chat-{source}",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 16,
+        **NO_CACHE,
+    }
+    with httpx.Client(base_url=owned.gateway.client.base_url, timeout=60, trust_env=False) as client:
+        response: Final = client.post(
+            "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {owned.gateway.key}"}
+        )
+    assert response.status_code == 200, (source, response.text)
+    assert JSON_OBJECT.validate_json(response.content)["choices"] == answered_choices(), (source, response.text)
+    return prompt
+
+
+def _assert_each_source_signed_its_own_clean_body(container: Wire, prompts: Mapping[str, str]) -> None:
+    forwarded: Final = by_marker(container.drain())
+    assert set(forwarded) == {marker(prompt) for prompt in prompts.values()}, sorted(forwarded)
+    for source, prompt in prompts.items():
+        _assert_forwarded(forwarded[marker(prompt)], container, source, prompt)
+
+
+def _assert_forwarded(request: Request, container: Wire, source: str, prompt: str) -> None:
+    assert JSON_OBJECT.validate_json(request.body) == container_body(container.url, prompt, stream=False), (
+        source,
+        request.body,
+    )
+    signer, session_token = _SIGNERS[source]
+    assert_signed(request, signer, REGION, session_token)
+
+
+def _worker_pids(log: str) -> tuple[int, ...]:
+    return tuple(
+        int(line.partition(_STARTED_WORKER)[2].partition("]")[0])
+        for line in log.splitlines()
+        if _STARTED_WORKER in line
+    )
+
+
+def _dead(process: psutil.Process) -> bool:
+    try:
+        return process.status() == psutil.STATUS_ZOMBIE
+    except psutil.NoSuchProcess:
+        return True
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+def test_every_aws_credential_source_keeps_its_credentials_out_of_the_body_when_a_worker_is_killed(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with wire_server(_container) as container, wire_server(_sts) as authority:
+        with owned_proxy_process(
+            gateway,
+            tmp_path,
+            _environment(authority, tmp_path),
+            config=_config(container, authority, tmp_path),
+            remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")),
+            workers=2,
+        ) as owned:
+            workers: Final = eventually(
+                lambda: _worker_pids(owned.log.read_text()), lambda pids: len(pids) == 2, seconds=30
+            )
+            _assert_each_source_signed_its_own_clean_body(
+                container, {source: _ask(owned, source) for source in _SIGNERS}
+            )
+            victim: Final = psutil.Process(workers[0])
+            victim.send_signal(signal.SIGKILL)
+            eventually(lambda: _dead(victim), bool, seconds=graceful_stop_seconds())
+            _assert_each_source_signed_its_own_clean_body(
+                container, {source: _ask(owned, source) for source in _SIGNERS}
+            )
+            assert psutil.Process(workers[1]).is_running(), workers
+            assert owned.process.poll() is None, owned.log.read_text()

tests/integration/providers/test_sagemaker_chat_wire.py

@@ -1,37 +1,215 @@
-import json
+import asyncio
 import uuid
-from typing import Final
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final, Literal, TypeAlias
 
+import httpx
 import pytest
-from integration._support.client import Gateway
-from integration._support.wire import Reply, Request, wire_server
-from pydantic import JsonValue, TypeAdapter
+from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration.providers._count_tokens_system_lift import (
+    anthropic_client,
+    async_anthropic_client,
+    async_openai_client,
+    openai_client,
+)
+from integration.providers._sagemaker_chat_container import (
+    ACCESS_KEY,
+    CONTENT,
+    ENDPOINT,
+    JSON_OBJECT,
+    NO_CACHE,
+    PROMPT,
+    REGION,
+    SECRET_KEY,
+    SERVED_MODEL,
+    answered_choices,
+    assert_signed,
+    by_marker,
+    container_body,
+    container_reply,
+    marked_prompt,
+    marker,
+    only_request,
+)
+from pydantic import JsonValue
 
-_ENDPOINT: Final = "integration-vllm-endpoint"
 _INFERENCE_COMPONENT: Final = "integration-vllm-component"
-_SERVED_MODEL: Final = "integration-org/served-chat-model"
-_ACCESS_KEY: Final = "AKIAINTEGRATION000003"
-_PROMPT: Final = "synthetic inference component request"
-_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_SESSION_TOKEN: Final = "synthetic-session-token-for-testing"
+_CALLER_ACCESS_KEY: Final = "AKIAINTEGRATIONCALLER"
+_CALLER_SECRET_KEY: Final = "synthetic-caller-secret-key-for-testing"
+_DEPLOYMENT_AUTH_KEYS: Final = frozenset({"aws_access_key_id", "aws_secret_access_key", "aws_region_name"})
+
+Endpoint: TypeAlias = Literal["chat", "messages", "responses"]
+
+_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses")
+_PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType(
+    {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"}
+)
 
 
-def _completion(identity: str) -> bytes:
-    return json.dumps(
-        {
-            "id": identity,
-            "object": "chat.completion",
-            "created": 1,
-            "model": _SERVED_MODEL,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": "sagemaker wire control"},
-                    "finish_reason": "stop",
-                }
-            ],
-            "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
-        }
-    ).encode()
+def _deployment(
+    scenario: Scenario, wire: Wire, extra: Mapping[str, JsonValue], *, provider: str = "sagemaker_chat"
+) -> str:
+    return scenario.model(
+        model_info=None,
+        model=f"{provider}/{ENDPOINT}",
+        api_key=None,
+        api_base=None,
+        aws_access_key_id=ACCESS_KEY,
+        aws_secret_access_key=SECRET_KEY,
+        aws_region_name=REGION,
+        sagemaker_base_url=wire.url,
+        **extra,
+    )
+
+
+def _chat(gateway: Gateway, model: str, prompt: str, caller: Mapping[str, JsonValue]) -> httpx.Response:
+    return gateway.request(
+        "POST",
+        "/v1/chat/completions",
+        {"model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": 16, **NO_CACHE, **caller},
+    )
+
+
+def _spend_row(call_id: str) -> Mapping[str, JsonValue]:
+    (row,) = eventually(
+        lambda: read_rows(
+            'SELECT status, model_group, row_to_json(s)::text AS stored FROM "LiteLLM_SpendLogs" s '
+            "WHERE litellm_call_id = %s",
+            (call_id,),
+        ),
+        lambda rows: len(rows) == 1,
+        seconds=70,
+    )
+    return row
+
+
+@dataclass(frozen=True, slots=True)
+class _Answer:
+    call_id: str
+    text: str
+
+
+def _chat_openai_sync(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    with openai_client(gateway) as client:
+        raw: Final = client.chat.completions.with_raw_response.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=16,
+            extra_body=dict(NO_CACHE),
+        )
+        completion: Final = raw.parse()
+    return _Answer(raw.headers["x-litellm-call-id"], completion.choices[0].message.content or "")
+
+
+async def _chat_openai_async_stream(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    async with async_openai_client(gateway) as client:
+        stream: Final = await client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=16,
+            stream=True,
+            stream_options={"include_usage": True},
+            extra_body=dict(NO_CACHE),
+        )
+        text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices])
+    return _Answer(stream.response.headers["x-litellm-call-id"], text)
+
+
+def _messages_anthropic_sync(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    with anthropic_client(gateway) as client:
+        raw: Final = client.messages.with_raw_response.create(
+            model=model,
+            max_tokens=16,
+            messages=[{"role": "user", "content": prompt}],
+            extra_body=dict(NO_CACHE),
+        )
+        message: Final = raw.parse()
+    return _Answer(
+        raw.headers["x-litellm-call-id"], "".join(block.text for block in message.content if block.type == "text")
+    )
+
+
+async def _messages_anthropic_async_stream(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    async with async_anthropic_client(gateway) as client:
+        stream: Final = await client.messages.create(
+            model=model,
+            max_tokens=16,
+            messages=[{"role": "user", "content": prompt}],
+            stream=True,
+            extra_body=dict(NO_CACHE),
+        )
+        deltas: Final = [event.delta async for event in stream if event.type == "content_block_delta"]
+    return _Answer(
+        stream.response.headers["x-litellm-call-id"],
+        "".join(delta.text for delta in deltas if delta.type == "text_delta"),
+    )
+
+
+def _responses_httpx(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    response: Final = gateway.request(
+        "POST", "/v1/responses", {"model": model, "input": prompt, "max_output_tokens": 16, **NO_CACHE}
+    )
+    assert response.status_code == 200, response.text
+    output: Final = tuple(
+        object_value(item) for item in list_value(JSON_OBJECT.validate_json(response.content)["output"])
+    )
+    (message,) = (item for item in output if item["type"] == "message")
+    text: Final = "".join(string_value(object_value(part)["text"]) for part in list_value(message["content"]))
+    return _Answer(response.headers["x-litellm-call-id"], text)
+
+
+def _responses_httpx_stream(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    with gateway.client.stream(
+        "POST",
+        "/v1/responses",
+        json={"model": model, "input": prompt, "max_output_tokens": 16, "stream": True, **NO_CACHE},
+        headers={"Authorization": f"Bearer {gateway.key}"},
+    ) as response:
+        assert response.status_code == 200, response.read()
+        lines: Final = tuple(response.iter_lines())
+    events: Final = tuple(
+        JSON_OBJECT.validate_json(line.removeprefix("data: ")) for line in lines if line.startswith("data: {")
+    )
+    assert [event["type"] for event in events if event["type"] == "response.completed"] == ["response.completed"], lines
+    text: Final = "".join(
+        string_value(event["delta"]) for event in events if event["type"] == "response.output_text.delta"
+    )
+    return _Answer(response.headers["x-litellm-call-id"], text)
+
+
+def _chat_openai_stream(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    return asyncio.run(_chat_openai_async_stream(gateway, model, prompt))
+
+
+def _messages_anthropic_stream(gateway: Gateway, model: str, prompt: str) -> _Answer:
+    return asyncio.run(_messages_anthropic_async_stream(gateway, model, prompt))
+
+
+@dataclass(frozen=True, slots=True)
+class _Driver:
+    run: Callable[[Gateway, str, str], _Answer]
+    stream: bool
+    forwarded: Mapping[str, JsonValue]
+
+
+_STREAM_WITH_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
+    {"stream": True, "stream_options": {"include_usage": True}}
+)
+_DRIVERS: Final[Mapping[str, _Driver]] = MappingProxyType(
+    {
+        "chat-openai-sync": _Driver(_chat_openai_sync, False, {"stream": False}),
+        "chat-openai-async-stream": _Driver(_chat_openai_stream, True, _STREAM_WITH_USAGE),
+        "messages-anthropic-sync": _Driver(_messages_anthropic_sync, False, {}),
+        "messages-anthropic-async-stream": _Driver(_messages_anthropic_stream, True, _STREAM_WITH_USAGE),
+        "responses-httpx": _Driver(_responses_httpx, False, {}),
+        "responses-httpx-stream": _Driver(_responses_httpx_stream, True, _STREAM_WITH_USAGE),
+    }
+)
 
 
 @pytest.mark.covers(
@@ -42,49 +220,266 @@ def test_sagemaker_chat_signs_the_inference_component_header_and_sends_hf_model_
 ) -> None:
     identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}"
 
+    with wire_server(lambda _request: container_reply(identity, False)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {"model_id": _INFERENCE_COMPONENT, "hf_model_name": SERVED_MODEL})
+        response: Final = _chat(gateway, model, PROMPT, {})
+        request: Final = only_request(wire)
+    assert response.status_code == 200, response.text
+    payload: Final = JSON_OBJECT.validate_json(response.content)
+    assert payload["id"] == identity, response.text
+    assert payload["choices"] == answered_choices(), response.text
+    assert request.headers["x-amzn-sagemaker-inference-component"] == _INFERENCE_COMPONENT, dict(request.headers)
+    assert_signed(request, ACCESS_KEY, REGION, None)
+    signed_headers: Final = next(
+        part for part in request.headers["authorization"].split(", ") if part.startswith("SignedHeaders=")
+    )
+    assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";")
+    assert JSON_OBJECT.validate_json(request.body) == {
+        **container_body(wire.url, PROMPT, stream=False, model_id=_INFERENCE_COMPONENT),
+        "model": SERVED_MODEL,
+    }, request.body
+
+
+@pytest.mark.parametrize("driver", tuple(_DRIVERS))
+def test_sagemaker_chat_keeps_the_deployment_aws_credentials_out_of_the_container_body_on_every_endpoint(
+    gateway: Gateway, driver: str
+) -> None:
+    chosen: Final = _DRIVERS[driver]
+    identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}"
+    prompt: Final = marked_prompt()
+
+    with wire_server(lambda _request: container_reply(identity, chosen.stream)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        answer: Final = chosen.run(gateway, model, prompt)
+        request: Final = only_request(wire)
+    assert answer.text == "".join(CONTENT), answer
+    assert JSON_OBJECT.validate_json(request.body) == container_body(wire.url, prompt, **chosen.forwarded), request.body
+    assert_signed(request, ACCESS_KEY, REGION, None)
+    row: Final = _spend_row(answer.call_id)
+    assert (row["status"], row["model_group"]) == ("success", model), row
+    assert SECRET_KEY not in string_value(row["stored"]), row
+
+
+@dataclass(frozen=True, slots=True)
+class _Shape:
+    provider: str
+    deployment: Mapping[str, JsonValue]
+    caller: Mapping[str, JsonValue]
+    forwarded: Mapping[str, JsonValue]
+    signer: str
+    region: str
+
+
+def _caller_flag(value: JsonValue) -> _Shape:
+    return _Shape("sagemaker_chat", {}, {"aws_custom_flag": value}, {"aws_custom_flag": value}, ACCESS_KEY, REGION)
+
+
+_SHAPES: Final[Mapping[str, _Shape]] = MappingProxyType(
+    {
+        "deployment-session-token": _Shape(
+            "sagemaker_chat", {"aws_session_token": _SESSION_TOKEN}, {}, {}, ACCESS_KEY, REGION
+        ),
+        "nova-deployment": _Shape("sagemaker_nova", {}, {}, {}, ACCESS_KEY, REGION),
+        "caller-top-k": _Shape("sagemaker_chat", {}, {"top_k": 5}, {"top_k": 5}, ACCESS_KEY, REGION),
+        "caller-aws-custom-flag": _caller_flag("keep"),
+        "caller-aws-custom-flag-int": _caller_flag(7),
+        "caller-aws-custom-flag-list": _caller_flag(["a", 1]),
+        "caller-aws-custom-flag-empty": _caller_flag(""),
+        "caller-aws-custom-flag-5kb": _caller_flag("x" * 5120),
+        "caller-aws-region-name": _Shape(
+            "sagemaker_chat", {}, {"aws_region_name": "us-west-2"}, {}, ACCESS_KEY, "us-west-2"
+        ),
+        "caller-own-credentials": _Shape(
+            "sagemaker_chat",
+            {},
+            {"aws_access_key_id": _CALLER_ACCESS_KEY, "aws_secret_access_key": _CALLER_SECRET_KEY},
+            {},
+            _CALLER_ACCESS_KEY,
+            REGION,
+        ),
+    }
+)
+
+
+@pytest.mark.parametrize("shape", tuple(_SHAPES))
+def test_sagemaker_chat_forwards_only_non_credential_params_for_each_deployment_and_caller_shape(
+    gateway: Gateway, shape: str
+) -> None:
+    case: Final = _SHAPES[shape]
+    identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}"
+    prompt: Final = marked_prompt()
+
+    with wire_server(lambda _request: container_reply(identity, False)) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, case.deployment, provider=case.provider)
+        response: Final = _chat(gateway, model, prompt, case.caller)
+        request: Final = only_request(wire)
+    assert response.status_code == 200, response.text
+    assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text
+    expected: Final = container_body(wire.url, prompt, stream=False, **case.forwarded)
+    assert JSON_OBJECT.validate_json(request.body) == (
+        expected if case.provider == "sagemaker_chat" else {k: v for k, v in expected.items() if k != "model"}
+    ), request.body
+    session_token: Final = case.deployment.get("aws_session_token")
+    assert_signed(request, case.signer, case.region, session_token if isinstance(session_token, str) else None)
+
+
+@pytest.mark.parametrize("caller", ("proxy-admin", "model-scoped-key"))
+def test_sagemaker_chat_failed_health_check_shows_no_deployment_aws_credentials(gateway: Gateway, caller: str) -> None:
+    failure: Final = Reply(status=500, body=b'{"message":"synthetic container failure"}')
+
+    with wire_server(lambda _request: failure) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        key: Final = gateway.key if caller == "proxy-admin" else scenario.key(models=[model])
+        response: Final = gateway.request("GET", "/health", params={"model": model}, key=key)
+        probes: Final = wire.drain()
+    assert response.status_code == 503, response.text
+    assert SECRET_KEY not in response.text
+    payload: Final = JSON_OBJECT.validate_json(response.content)
+    assert payload["healthy_count"] == 0, response.text
+    (unhealthy,) = list_value(payload["unhealthy_endpoints"])
+    endpoint: Final = object_value(unhealthy)
+    assert "synthetic container failure" in string_value(endpoint["error"]), response.text
+    shown_body: Final = object_value(object_value(endpoint["raw_request_typed_dict"])["raw_request_body"])
+    assert shown_body["model"] == ENDPOINT, response.text
+    assert not _DEPLOYMENT_AUTH_KEYS & set(shown_body), response.text
+    assert [JSON_OBJECT.validate_json(probe.body) for probe in probes] == [shown_body], probes
+
+
+def test_transform_request_preview_of_a_sagemaker_chat_call_shows_no_aws_credentials(gateway: Gateway) -> None:
+    prompt: Final = marked_prompt()
+    response: Final = gateway.request(
+        "POST",
+        "/utils/transform_request",
+        {
+            "call_type": "completion",
+            "request_body": {
+                "model": f"sagemaker_chat/{ENDPOINT}",
+                "messages": [{"role": "user", "content": prompt}],
+                "aws_access_key_id": _CALLER_ACCESS_KEY,
+                "aws_secret_access_key": _CALLER_SECRET_KEY,
+                "aws_region_name": REGION,
+            },
+        },
+    )
+    assert response.status_code == 200, response.text
+    assert _CALLER_SECRET_KEY not in response.text
+    payload: Final = JSON_OBJECT.validate_json(response.content)
+    assert payload["raw_request_body"] == {"model": ENDPOINT, "messages": [{"role": "user", "content": prompt}]}, (
+        response.text
+    )
+    assert string_value(payload["raw_request_api_base"]).endswith(f"/endpoints/{ENDPOINT}/invocations"), response.text
+
+
+def test_sagemaker_chat_container_error_reaches_the_caller_after_one_attempt_without_credentials(
+    gateway: Gateway,
+) -> None:
+    prompt: Final = marked_prompt()
+    failure: Final = Reply(
+        status=424, body=b'{"ErrorCode":"CLIENT_ERROR_FROM_MODEL","Message":"synthetic model error"}'
+    )
+
+    with wire_server(lambda _request: failure) as wire, gateway.scenario() as scenario:
+        model: Final = _deployment(scenario, wire, {})
+        response: Final = _chat(gateway, model, prompt, {})
+        request: Final = only_request(wire)
+    assert response.status_code == 424, response.text
+    assert "synthetic model error" in response.text, response.text
+    assert JSON_OBJECT.validate_json(request.body) == container_body(wire.url, prompt, stream=False), request.body
+    assert_signed(request, ACCESS_KEY, REGION, None)
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+    endpoint: Endpoint
+    stream: bool
+    marker: str
+
+
+@dataclass(frozen=True, slots=True)
+class _Served:
+    call: _Call
+    status: int
+    text: str
+    call_id: str
+
+
+def _calls(count: int) -> tuple[_Call, ...]:
+    return tuple(
+        _Call(endpoint=_ENDPOINTS[index % len(_ENDPOINTS)], stream=index % 2 == 1, marker=marker(marked_prompt()))
+        for index in range(count)
+    )
+
+
+def _request_body(model: str, call: _Call) -> Mapping[str, JsonValue]:
+    prompt: Final = f"{PROMPT} marker-{call.marker}"
+    common: Final[Mapping[str, JsonValue]] = {"model": model, "stream": call.stream, **NO_CACHE}
+    match call.endpoint:
+        case "chat" | "messages":
+            return {**common, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}
+        case "responses":
+            return {**common, "max_output_tokens": 16, "input": prompt}
+
+
+async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
+    async with client.stream(
+        "POST",
+        _PATHS[call.endpoint],
+        json=_request_body(model, call),
+        headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
+    ) as response:
+        raw: Final = await response.aread()
+    return _Served(call, response.status_code, raw.decode(), response.headers["x-litellm-call-id"])
+
+
+async def _burst(gateway: Gateway, model: str, calls: tuple[_Call, ...]) -> tuple[_Served, ...]:
+    async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client:
+        return tuple(await asyncio.gather(*(_send(client, gateway.key, model, call) for call in calls)))
+
+
+def _assert_answered(served: _Served, down: frozenset[str]) -> None:
+    if served.call.marker in down:
+        assert served.status == 503, (served.call, served.text)
+        assert "synthetic container outage" in served.text, (served.call, served.text)
+        return
+    assert served.status == 200, (served.call, served.text)
+    assert all(part in served.text for part in CONTENT), (served.call, served.text)
+
+
+def _assert_clean(request: Request) -> None:
+    assert not _DEPLOYMENT_AUTH_KEYS & set(JSON_OBJECT.validate_json(request.body)), request.body
+    assert SECRET_KEY.encode() not in request.body, request.body
+
+
+async def test_sagemaker_chat_container_outage_mid_burst_answers_every_call_once_without_credentials(
+    gateway: Gateway,
+) -> None:
+    burst: Final = _calls(24)
+    recovery: Final = _calls(12)
+    down: Final = frozenset(call.marker for index, call in enumerate(burst) if (index // 2) % 2 == 0)
+    streams: Final = MappingProxyType({call.marker: call.stream for call in (*burst, *recovery)})
+
     def respond(request: Request) -> Reply:
-        assert request.method == "POST", request
-        assert request.target == "/", request
-        assert request.headers["x-amzn-sagemaker-inference-component"] == _INFERENCE_COMPONENT, dict(request.headers)
-        authorization: Final = request.headers["authorization"]
-        assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), authorization
-        signed_headers: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders="))
-        assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";"), (
-            authorization
-        )
-        body: Final = _JSON_OBJECT.validate_json(request.body)
-        assert body["model"] == _SERVED_MODEL, body
-        assert body["messages"] == [{"role": "user", "content": _PROMPT}], body
-        assert body["max_tokens"] == 16, body
-        assert "hf_model_name" not in body, body
-        return Reply(body=_completion(identity))
+        found: Final = marker(request.body.decode())
+        if found in down:
+            return Reply(status=503, body=b'{"message":"synthetic container outage"}')
+        return container_reply(f"sagemaker-chat-{found}", streams[found])
 
     with wire_server(respond) as wire, gateway.scenario() as scenario:
-        model: Final = scenario.model(
-            model=f"sagemaker_chat/{_ENDPOINT}",
-            api_key=None,
-            api_base=None,
-            model_id=_INFERENCE_COMPONENT,
-            hf_model_name=_SERVED_MODEL,
-            aws_access_key_id=_ACCESS_KEY,
-            aws_secret_access_key="synthetic-secret-key-for-testing",
-            aws_region_name="us-east-1",
-            sagemaker_base_url=wire.url,
-        )
-        response: Final = gateway.request(
-            "POST",
-            "/v1/chat/completions",
-            {"model": model, "messages": [{"role": "user", "content": _PROMPT}], "max_tokens": 16},
-        )
-        assert response.status_code == 200, response.text
-        payload: Final = _JSON_OBJECT.validate_json(response.content)
-        assert payload["id"] == identity, response.text
-        assert payload["choices"] == [
-            {
-                "finish_reason": "stop",
-                "index": 0,
-                "message": {"role": "assistant", "content": "sagemaker wire control"},
-                "provider_specific_fields": {},
-            }
-        ], response.text
-        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/")], response.text
+        model: Final = _deployment(scenario, wire, {})
+        answered: Final = (*await _burst(gateway, model, burst), *await _burst(gateway, model, recovery))
+        forwarded: Final = by_marker(wire.drain())
+    assert set(forwarded) == set(streams), sorted(forwarded)
+    for request in forwarded.values():
+        _assert_clean(request)
+    for served in answered:
+        _assert_answered(served, down)
+    rows: Final = eventually(
+        lambda: read_rows('SELECT litellm_call_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)),
+        lambda found: len(found) >= len(answered),
+        seconds=70,
+    )
+    statuses: Final = {string_value(row["litellm_call_id"]): row["status"] for row in rows}
+    assert len(statuses) == len(rows) == len(answered), rows
+    assert statuses == {
+        served.call_id: "failure" if served.call.marker in down else "success" for served in answered
+    }, rows

tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py

@@ -13,6 +13,7 @@ steady provider stream into gap-then-burst delivery.
 import binascii
 import json
 import struct
+from collections.abc import Mapping
 from typing import AsyncIterator, Iterator
 from unittest.mock import MagicMock
 
@@ -258,15 +259,17 @@ class _RequestCapturingHTTPHandler(HTTPHandler):
     def __init__(self) -> None:
         super().__init__()
         self.request_headers: dict[str, str] = {}
-        self.request_body: dict = {}
+        self.request_body: dict[str, object] = {}
 
-    def post(self, url: str, headers=None, data=None, **kwargs) -> httpx.Response:
+    def post(self, url: str, headers: Mapping[str, str] | None = None, data=None, **kwargs) -> httpx.Response:
         self.request_headers = dict(headers or {})
         self.request_body = json.loads(data)
         return httpx.Response(200, json=_STUB_COMPLETION_RESPONSE, request=httpx.Request("POST", url))
 
 
-def _invoke_sagemaker_chat(monkeypatch, **extra_params) -> _RequestCapturingHTTPHandler:
+def _invoke_sagemaker_chat(
+    monkeypatch, model: str = "sagemaker_chat/my-endpoint", **extra_params
+) -> _RequestCapturingHTTPHandler:
     """Drive one sagemaker_chat completion against an injected transport.
 
     A Bedrock API key short-circuits SigV4 inside `BaseAWSLLM._sign_request`, which would hide
@@ -275,7 +278,7 @@ def _invoke_sagemaker_chat(monkeypatch, **extra_params) -> _RequestCapturingHTTP
     monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
     client = _RequestCapturingHTTPHandler()
     litellm.completion(
-        model="sagemaker_chat/my-endpoint",
+        model=model,
         messages=[{"role": "user", "content": "hi"}],
         aws_access_key_id="AKIATESTTESTTESTTEST",
         aws_secret_access_key="test-secret-key",
@@ -315,6 +318,31 @@ def test_hf_model_name_becomes_the_body_model(monkeypatch):
     assert "hf_model_name" not in client.request_body
 
 
+@pytest.mark.parametrize(
+    ("model", "expected_body"),
+    [
+        (
+            "sagemaker_chat/my-endpoint",
+            {"model": "my-endpoint", "messages": [{"role": "user", "content": "hi"}], "top_k": 5},
+        ),
+        ("sagemaker_nova/my-endpoint", {"messages": [{"role": "user", "content": "hi"}], "top_k": 5}),
+    ],
+)
+def test_body_carries_no_aws_credentials(
+    monkeypatch: pytest.MonkeyPatch, model: str, expected_body: dict[str, object]
+):
+    client = _invoke_sagemaker_chat(
+        monkeypatch,
+        model=model,
+        aws_session_token="test-session-token",
+        top_k=5,
+    )
+
+    assert client.request_body == expected_body
+    assert "Credential=AKIATESTTESTTESTTEST/" in client.request_headers["Authorization"]
+    assert client.request_headers["X-Amz-Security-Token"] == "test-session-token"
+
+
 def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch):
     """Without `hf_model_name` the body must keep the model it has today."""
     client = _invoke_sagemaker_chat(monkeypatch)