fix(bedrock): forward aws_session_tags in batch cancel

cancel_batch resolved credentials and polled job status without the
deployment's aws_session_tags, so on a tag-gated trust policy batch
create and retrieve succeeded while cancel failed with AccessDenied.
Thread the tags through both calls and cover it with a regression test
that fakes STS behind a tag-gated trust policy.

litellm/llms/bedrock/batches/handler.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 
@@ -6,6 +6,7 @@ from openai.types.batch import BatchRequestCounts
 from openai.types.batch import Metadata as OpenAIBatchMetadata
 
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
+from litellm.types.llms.bedrock import AwsSessionTag
 from litellm.types.utils import LiteLLMBatch
 
 if TYPE_CHECKING:
@@ -116,6 +117,7 @@ class BedrockBatchesHandler:
         aws_web_identity_token: str | None = None,
         aws_sts_endpoint: str | None = None,
         aws_external_id: str | None = None,
+        aws_session_tags: Sequence[AwsSessionTag] | None = None,
         **kwargs: object,  # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim
     ) -> "LiteLLMBatch":
         try:
@@ -139,6 +141,7 @@ class BedrockBatchesHandler:
             aws_web_identity_token=aws_web_identity_token,
             aws_sts_endpoint=aws_sts_endpoint,
             aws_external_id=aws_external_id,
+            aws_session_tags=aws_session_tags,
         )
 
         client: Final = boto3.client(
@@ -163,6 +166,7 @@ class BedrockBatchesHandler:
                 aws_web_identity_token=aws_web_identity_token,
                 aws_sts_endpoint=aws_sts_endpoint,
                 aws_external_id=aws_external_id,
+                aws_session_tags=aws_session_tags,
             )
 
         try:

tests/test_litellm/llms/bedrock/batches/test_handler.py

@@ -482,10 +482,36 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3):
     assert batch.status == "cancelled"
 
 
+class _TagGatedSTSClient:
+    """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``."""
+
+    def __init__(self, tags: list[dict[str, str]], access_key_id: str) -> None:
+        self._tags = tags
+        self._access_key_id = access_key_id
+
+    def get_caller_identity(self):
+        return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
+
+    def assume_role(self, **params):
+        from botocore.exceptions import ClientError
+
+        if list(params.get("Tags", ())) != self._tags:
+            raise ClientError(
+                {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}},
+                "AssumeRole",
+            )
+        return {
+            "Credentials": {
+                "AccessKeyId": self._access_key_id,
+                "SecretAccessKey": "assumed-secret",
+                "SessionToken": "assumed-session-token",
+                "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30),
+            }
+        }
+
+
 def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_session(monkeypatch):
     """Status polling must assume the role with the deployment's session tags, like every other call."""
-    from botocore.exceptions import ClientError
-
     monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
     monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
     tags = [{"Key": "team", "Value": "genai"}]
@@ -493,28 +519,9 @@ def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_se
     fake_bedrock = MagicMock()
     fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response()
 
-    class FakeSTSClient:
-        def get_caller_identity(self):
-            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
-
-        def assume_role(self, **params):
-            if list(params.get("Tags", ())) != tags:
-                raise ClientError(
-                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}},
-                    "AssumeRole",
-                )
-            return {
-                "Credentials": {
-                    "AccessKeyId": "ASIABATCHSTATUSTAGGED",
-                    "SecretAccessKey": "assumed-secret",
-                    "SessionToken": "assumed-session-token",
-                    "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30),
-                }
-            }
-
     def boto3_client(service_name, **kwargs):
         if service_name == "sts":
-            return FakeSTSClient()
+            return _TagGatedSTSClient(tags, "ASIABATCHSTATUSTAGGED")
         bedrock_client_kwargs.append(kwargs)
         return fake_bedrock
 
@@ -530,3 +537,36 @@ def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_se
 
     assert batch.status == "completed"
     assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHSTATUSTAGGED"]
+
+
+def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatch):
+    """Cancelling on a tag-gated role must forward the deployment's session tags to both the stop and status calls."""
+    import litellm
+
+    monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
+    monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
+    tags = [{"Key": "team", "Value": "genai"}]
+    bedrock_client_kwargs: list[dict] = []
+    fake_bedrock = MagicMock()
+    fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")
+
+    def boto3_client(service_name, **kwargs):
+        if service_name == "sts":
+            return _TagGatedSTSClient(tags, "ASIABATCHCANCELTAGGED")
+        bedrock_client_kwargs.append(kwargs)
+        return fake_bedrock
+
+    with patch("boto3.client", side_effect=boto3_client):
+        batch = litellm.cancel_batch(
+            batch_id=JOB_ARN,
+            custom_llm_provider="bedrock",
+            aws_access_key_id="AKIABATCHCANCELCALLER",
+            aws_secret_access_key="pod-caller-secret",
+            aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role",
+            aws_session_name="litellm-batch-session",
+            aws_session_tags=tags,
+        )
+
+    fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
+    assert batch.status == "cancelled"
+    assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2