chore: merge main into litellm_tts_true_streaming

.github/scripts/assert_ci_coverage.py

@@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
 DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
 COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
 GLOB_CHARS = frozenset("*?")
+DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
 
 # Trees whose jobs are sharded with no catch-all bucket, so every child that holds
 # tests has to be named by some shard or it runs nowhere. A child listed here is
@@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
     return frozenset(ast.literal_eval(values))
 
 
+def _owns(workflow: str, relative_path: str) -> bool:
+    owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
+    return owned_root is not None and relative_path.startswith(owned_root)
+
+
+def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
+    return frozenset[str]().union(
+        *(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
+    )
+
+
 def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
     runner: Final = repo_root / "tests/integration/run.py"
     if not runner.exists():
@@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         )
         != 1
     )
-    gha_tokens: Final = _invoked_test_tokens(
-        scalar
-        for path in (repo_root / ".github/workflows").glob("*.y*ml")
-        for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
+    workflow_tokens: Final = MappingProxyType(
+        {
+            path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
+            for path in (repo_root / ".github/workflows").glob("*.y*ml")
+        }
+    )
+    gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
+    circle_owned: Final = (
+        frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
     )
     findings: Final = (
         tuple(
-            Finding(path, "integration contract is also selected by GitHub Actions")
-            for path in paths
+            Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
+            for path in sorted(circle_owned)
             if any(_token_covers(token, path) for token in gha_tokens)
         )
         + tuple(
@@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
             for path in sorted(github_files)
             if not (repo_root / path).is_file()
         )
+        + tuple(
+            Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
+            for path in sorted(github_files)
+            if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
+        )
     )
     browser_commands: Final = tuple(
         scalar.value
@@ -688,13 +710,9 @@ def main() -> int:
     scalars = _all_scalars()
 
     integration_paths, ownership_findings = _integration_ownership()
-    test_findings = (
-        _uncovered_tests(
-            allowlist,
-            _invoked_selections(scalars)
-            + (Selection(included=integration_paths, ignored=frozenset()),),
-        )
-        + ownership_findings
+    test_findings = _uncovered_tests(
+        allowlist,
+        _invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
     )
     dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
     stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
@@ -711,13 +729,25 @@ def main() -> int:
             test_findings,
             "Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
+    if ownership_findings:
+        browser_remedy = (
+            " CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
+            "select them and the default Playwright config must exclude them."
+            if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
+            else ""
+        )
+        _report(
+            "integration contracts with the wrong CI owner",
+            ownership_findings,
+            "Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
+        )
     if dockerfile_findings:
         _report(
             "Dockerfiles that no CI job builds",
             dockerfile_findings,
             "Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
-    if stale_findings or test_findings or dockerfile_findings:
+    if stale_findings or test_findings or ownership_findings or dockerfile_findings:
         return 1
 
     _write(

.github/workflows/test-rust.yml

@@ -16,6 +16,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"
@@ -46,6 +50,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"

.github/workflows/test-unit.yml

@@ -696,7 +696,6 @@ jobs:
               tests/unit/responses/litellm_completion_transformation
               tests/unit/responses/test_additional_tools.py
               tests/unit/responses/test_custom_tool_call.py
-              tests/unit/responses/test_dispatch.py
               tests/unit/responses/test_metadata_codex_callback.py
               tests/unit/responses/test_no_duplicate_spend_logs.py
               tests/unit/responses/test_null_test_fix.py

AGENTS.md

@@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
 
 `tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
 
-End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
+End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
 
 When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
 
@@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
 
 Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
 
-If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
+Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
 
 Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
 

litellm-rust/crates/gateway-inference/src/messages.rs

@@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [
     "x-openrouter-experimental-metadata",
 ];
 const ANTHROPIC_API_HEADER_PROVIDERS: &str =
-    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
+    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek";
 
 pub async fn create(
     State(gateway): State<Arc<Gateway>>,

litellm-rust/crates/gateway-inference/tests/messages.rs

@@ -16,7 +16,7 @@ use wiremock::{
 #[rstest]
 #[case::anthropic("anthropic/test-model", "/v1/messages", true)]
 #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)]
-#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)]
+#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)]
 #[tokio::test]
 async fn messages_reaches_the_provider_and_preserves_json_or_sse(
     #[case] model: &str,

litellm/__init__.py

@@ -1177,7 +1177,9 @@ def _build_models_by_provider() -> dict:
         "edenai": edenai_models,
         "datarobot": datarobot_models,
         "vertex_ai": vertex_chat_models
+        | vertex_code_chat_models
         | vertex_text_models
+        | vertex_code_text_models
         | vertex_anthropic_models
         | vertex_vision_models
         | vertex_language_models
@@ -1185,7 +1187,7 @@ def _build_models_by_provider() -> dict:
         | vertex_minimax_models
         | vertex_moonshot_models
         | vertex_zai_models,
-        "ai21": ai21_models,
+        "ai21": ai21_models | ai21_chat_models,
         "bedrock": bedrock_models | bedrock_converse_models,
         "petals": petals_models,
         "ollama": ollama_models,
@@ -1416,6 +1418,7 @@ from .exceptions import (
     InvalidRequestError,
     BadRequestError,
     ImageFetchError,
+    MalformedRequestFieldError as MalformedRequestFieldError,
     VectorStoreSearchError,
     NotFoundError,
     PermissionDeniedError,

litellm/assistants/main.py

@@ -129,6 +129,7 @@ def get_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_assistants=aget_assistants,
         )
@@ -312,6 +313,7 @@ def create_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             create_assistant_data=create_assistant_data,
             client=client,
             async_create_assistants=async_create_assistants,
@@ -461,6 +463,7 @@ def delete_assistant(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             assistant_id=assistant_id,
             client=client,
             async_delete_assistants=async_delete_assistants,
@@ -637,6 +640,7 @@ def create_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             acreate_thread=acreate_thread,
         )
@@ -783,6 +787,7 @@ def get_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_thread=aget_thread,
         )
@@ -960,6 +965,7 @@ def add_message(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             a_add_message=a_add_message,
         )
@@ -1111,6 +1117,7 @@ def get_messages(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_messages=aget_messages,
         )
@@ -1308,6 +1315,7 @@ def run_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             arun_thread=arun_thread,
             event_handler=event_handler,

litellm/batches/batch_utils.py

@@ -19,7 +19,7 @@ from litellm.llms.vertex_ai.batches.transformation import (
 )
 from litellm.types.llms.openai import Batch
 from litellm.types.utils import ModelInfo, Usage
-from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS
+from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
 from litellm.utils import token_counter
 
 
@@ -569,6 +569,7 @@ def extract_file_access_credentials(
             # A federated deployment holds no api_key, so without these the fetch that reads a
             # finished batch's output has nothing to authenticate with and its cost is never billed.
             *sorted(ANTHROPIC_WIF_KWARGS_KEYS),
+            *sorted(OPENAI_WIF_KWARGS_KEYS),
         )
         for key in credential_keys:
             if key in litellm_params:

litellm/batches/main.py

@@ -309,6 +309,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -471,6 +472,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
             organization=organization,
             timeout=timeout,
             max_retries=optional_params.max_retries,
+            litellm_params=litellm_params,
         )
     elif custom_llm_provider == "azure":
         api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -819,6 +821,7 @@ def list_batches(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -1016,6 +1019,7 @@ def cancel_batch(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -197,7 +197,7 @@ def _drops_audio_input(model: str, litellm_params: Mapping[str, object]) -> bool
     return not _supports_audio_input(model, litellm_params)
 
 
-def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
+def responses_provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
     return MappingProxyType(
         {
             key: value
@@ -1455,7 +1455,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
         model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
-        for key, value in _provider_metadata(raw_response.model_extra).items():
+        for key, value in responses_provider_metadata(raw_response.model_extra).items():
             setattr(model_response, key, value)
 
         # Preserve hidden params from the ResponsesAPIResponse, especially the headers
@@ -2139,7 +2139,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 from litellm.responses.utils import ResponseAPILoggingUtils
 
                 usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_data.get("usage"))
-            provider_metadata: Final = _provider_metadata(response_data)
+            provider_metadata: Final = responses_provider_metadata(response_data)
             served_service_tier: Final = response_data.get("service_tier")
             return ModelResponseStream(
                 choices=[

litellm/constants.py

@@ -1859,6 +1859,7 @@ SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_L
 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
 PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
 PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
+TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS: Final = 30
 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)  # 1 minute
 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500))

litellm/containers/README.md

@@ -120,6 +120,7 @@ class AzureContainerConfig(BaseContainerConfig):
         self,
         headers: dict,
         api_key: Optional[str] = None,
+        litellm_params: Optional[GenericLiteLLMParams] = None,
     ) -> dict:
         """Azure uses api-key header instead of Bearer token."""
         import litellm

litellm/cost_calculator.py

@@ -86,6 +86,9 @@ from litellm.llms.openai.cost_calculation import (
 from litellm.llms.openai.cost_calculation import (
     video_output_cost_per_second,
 )
+from litellm.llms.openai.transcriptions.gpt_transformation import (
+    OpenAIGPTAudioTranscriptionConfig,
+)
 from litellm.llms.perplexity.cost_calculator import (
     cost_per_token as perplexity_cost_per_token,
 )
@@ -325,6 +328,15 @@ def _transcription_usage_has_token_details(
     return (prompt_tokens_val > 0) or (completion_tokens_val > 0)
 
 
+def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool:
+    if model is None or custom_llm_provider != LlmProviders.OPENAI.value:
+        return False
+    transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
+        model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI
+    )
+    return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig)
+
+
 OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"]
 
 
@@ -1739,6 +1751,12 @@ def completion_cost(
                     audio_transcription_file_duration = _hidden.get(
                         "audio_transcription_duration",
                         getattr(completion_response, "duration", 0.0),
+                    ) or (
+                        completion_response.usage.seconds
+                        if isinstance(completion_response, TranscriptionResponse)
+                        and isinstance(completion_response.usage, TranscriptionUsageDurationObject)
+                        and _bills_provider_transcription_seconds(requested_model, custom_llm_provider)
+                        else 0.0
                     )
                 elif call_type in _RERANK_CALL_TYPES:
                     if completion_response is not None and isinstance(completion_response, RerankResponse):

litellm/decisions/call.py

@@ -458,7 +458,7 @@ def prepare_call(
         ) from error
 
     try:
-        resolved_api_base, resolved_api_key = provider_config.resolve_connection(
+        resolved_api_base, static_api_key = provider_config.resolve_connection(
             dynamic_api_base or api_base, dynamic_api_key or api_key
         )
     except ValueError as error:
@@ -469,6 +469,7 @@ def prepare_call(
             model=model,
             llm_provider=provider,
         )
+    resolved_api_key: Final = provider_config.resolve_credential(static_api_key, resolved_api_base, kwargs)
     if resolved_api_key is None and provider_config.api_key_required:
         raise litellm.AuthenticationError(
             message=f"Missing API key for Decisions provider '{provider}'",

litellm/decisions/main.py

@@ -1,5 +1,7 @@
+import asyncio
 import warnings
 from collections.abc import Mapping
+from functools import partial
 from typing import Final
 
 import httpx
@@ -104,19 +106,22 @@ async def adecisions(
     videos: DecisionsVideoClips | None = None,
     **kwargs: object,
 ) -> OpenAIDecisionResponse | DecisionsResponse:
-    call: Final = _prepare_decisions_call(
-        model=model,
-        decision_input=input,
-        questions=questions,
-        safety_identifier=safety_identifier,
-        audio=audio,
-        videos=videos,
-        api_key=api_key,
-        api_base=api_base,
-        timeout=timeout,
-        custom_llm_provider=custom_llm_provider,
-        extra_headers=extra_headers,
-        kwargs=kwargs,
+    call: Final = await asyncio.to_thread(
+        partial(
+            _prepare_decisions_call,
+            model=model,
+            decision_input=input,
+            questions=questions,
+            safety_identifier=safety_identifier,
+            audio=audio,
+            videos=videos,
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            custom_llm_provider=custom_llm_provider,
+            extra_headers=extra_headers,
+            kwargs=kwargs,
+        )
     )
     return await asend(call)
 

litellm/exceptions.py

@@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError):
         )
 
 
+class MalformedRequestFieldError(BadRequestError):
+    pass
+
+
 class UnprocessableEntityError(openai.UnprocessableEntityError):
     def __init__(
         self,

litellm/files/main.py

@@ -259,6 +259,7 @@ def create_file(
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
                 create_file_data=_create_file_request,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -388,6 +389,7 @@ def file_retrieve(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -569,6 +571,7 @@ def file_delete(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -775,6 +778,7 @@ def file_list(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -972,6 +976,7 @@ def file_content(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -1093,6 +1098,7 @@ def file_content_streaming(
             organization=openai_creds.organization,
             chunk_size=chunk_size,
             client=client if isinstance(client, (OpenAI, AsyncOpenAI)) else None,
+            litellm_params=optional_params.model_dump(exclude_none=True),
         )
     elif custom_llm_provider == LlmProviders.VERTEX_AI.value:
         if not _is_async:

litellm/fine_tuning/main.py

@@ -239,6 +239,7 @@ def create_fine_tuning_job(
                 client=kwargs.get(
                     "client", None
                 ),  # note, when we add this to `GenericLiteLLMParams` it impacts a lot of other tests + linting
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -436,6 +437,7 @@ def cancel_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -591,6 +593,7 @@ def list_fine_tuning_jobs(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -730,6 +733,7 @@ def retrieve_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":

litellm/images/main.py

@@ -39,6 +39,7 @@ from openai.types.audio.transcription_create_params import FileTypes
 # BFL handlers
 from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
 from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
+from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config
 from litellm.main import (
     azure_chat_completions,
     base_llm_aiohttp_handler,
@@ -505,6 +506,7 @@ def image_generation(
                 aimg_generation=aimg_generation,
                 client=client,
                 headers=headers,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "bedrock":
             if model is None:
@@ -681,10 +683,13 @@ def image_variation(
     api_base = provider_config.get_api_base(litellm_params.get("api_base", None))
 
     if image_variation_provider == LITELLM_IMAGE_VARIATION_PROVIDERS.OPENAI:
-        if api_key is None:
-            raise ValueError("API key is required for OpenAI image variations")
         if api_base is None:
             raise ValueError("API base is required for OpenAI image variations")
+        if api_key is None and (
+            resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base, litellm_params=litellm_params)
+            is None
+        ):
+            raise ValueError("API key is required for OpenAI image variations")
 
         response = openai_image_variations.image_variations(
             model_response=model_response,

litellm/integrations/prometheus.py

@@ -27,6 +27,7 @@ from litellm.constants import (
     PROXY_REJECTED_BEFORE_ROUTING_KEY,
 )
 from litellm.exceptions import (
+    MalformedRequestFieldError,
     validate_rate_limit_category,
     validate_rate_limit_type,
 )
@@ -3104,7 +3105,11 @@ class PrometheusLogger(CustomLogger):
 
             # On LiteLLM-side rejects (no deployment picked), route request_kwargs["model"]
             # into requested_model and leave deployment-scoped labels empty.
-            deployment_selected: Final = bool(model_id) and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY)
+            deployment_selected: Final = (
+                bool(model_id)
+                and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY)
+                and not isinstance(exception, MalformedRequestFieldError)
+            )
             if deployment_selected:
                 label_litellm_model_name = litellm_model_name
                 label_model_id = model_id

litellm/litellm_core_utils/get_litellm_params.py

@@ -78,6 +78,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
             "github_copilot_auth_type",
             "github_copilot_user_session",
             "fireworks_forward_user_id",
+            "responses_api_version",
             PROVIDER_AFFINITY_HEADER_KWARG_KEY,
         }
     )

litellm/litellm_core_utils/get_model_cost_map.py

@@ -17,6 +17,7 @@ import random
 import sys
 import threading
 import time
+import urllib.parse
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass, replace
 from datetime import datetime, timezone
@@ -24,11 +25,11 @@ from importlib.resources import files
 from itertools import chain
 from pathlib import Path
 from types import MappingProxyType
-from typing import Final, Protocol
+from typing import Final, Protocol, cast  # noqa: TID251  # adapts boto3 and botocore's dynamic client responses
 
 import httpx
 from pydantic import ConfigDict, TypeAdapter
-from typing_extensions import ReadOnly, TypedDict
+from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm import verbose_logger
 from litellm.constants import (
@@ -266,6 +267,19 @@ class _SyncGetClient(Protocol):
     def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ...
 
 
+class _S3Body(Protocol):
+    def read(self) -> bytes: ...
+
+
+class _S3GetObjectResponse(TypedDict):
+    Body: ReadOnly[_S3Body]
+    ETag: NotRequired[ReadOnly[str]]
+
+
+class _S3Client(Protocol):
+    def get_object(self, *, Bucket: str, Key: str) -> _S3GetObjectResponse: ...
+
+
 _FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable
 
 
@@ -276,6 +290,20 @@ def _default_reload_client() -> _AsyncGetClient:
     return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap)
 
 
+def _default_s3_client(timeout: int) -> _S3Client:
+    import boto3
+    from botocore.config import Config
+
+    return cast(  # cast-ok: boto3's generic client factory does not expose its generated S3 response type
+        _S3Client,
+        boto3.client(  # pyright: ignore[reportUnknownMemberType]  # generic factory returns Unknown service clients
+            "s3",
+            region_name=os.getenv("AWS_REGION_NAME") or None,
+            config=Config(connect_timeout=timeout, read_timeout=timeout, retries={"total_max_attempts": 1}),
+        ),
+    )
+
+
 def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome:
     reason: Final = f"{type(error).__name__} fetching {url}: {error}"
     if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)):
@@ -283,7 +311,41 @@ def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -
     return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None)
 
 
+def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _FetchAttemptOutcome:
+    try:
+        parsed: Final[object] = json.loads(content)
+    except ValueError as e:
+        return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}")
+    if not isinstance(parsed, dict):
+        return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}")
+    model_cost_map: Final = cast(  # cast-ok: JSON object keys are strings by definition
+        dict[str, object], parsed
+    )
+    return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag)
+
+
+def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome:
+    from botocore.exceptions import BotoCoreError, ClientError
+
+    parts: Final = urllib.parse.urlsplit(url)
+    bucket: Final = parts.netloc
+    key: Final = parts.path.lstrip("/")
+    if not bucket or not key:
+        return ModelCostMapReloadUnavailable(reason=f"invalid S3 URL {url}: expected s3://bucket/key")
+
+    try:
+        client: Final = _default_s3_client(timeout)
+        response: Final = client.get_object(Bucket=bucket, Key=key)
+        content: Final = response["Body"].read()
+    except (BotoCoreError, ClientError) as error:
+        return ModelCostMapReloadUnavailable(reason=f"S3 {type(error).__name__} fetching {url}: {error}")
+    return _classify_fetch_content(content, url, response.get("ETag"))
+
+
 async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
+    parts: Final = urllib.parse.urlsplit(url)
+    if parts.scheme == "s3":
+        return await asyncio.to_thread(_attempt_fetch_s3, url, timeout)
     try:
         response: Final = await client.get(url, timeout=timeout)
     except (httpx.HTTPError, httpx.InvalidURL) as e:
@@ -292,6 +354,9 @@ async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _Fe
 
 
 def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
+    parts: Final = urllib.parse.urlsplit(url)
+    if parts.scheme == "s3":
+        return _attempt_fetch_s3(url, timeout)
     try:
         response: Final = client.get(url, timeout=timeout)
     except (httpx.HTTPError, httpx.InvalidURL) as e:
@@ -307,15 +372,7 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp
         )
     if response.is_error:
         return ModelCostMapReloadUnavailable(reason=f"HTTP {response.status_code} from {url}")
-    try:
-        parsed: Final = response.json()
-    except ValueError as e:
-        return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}")
-    if not isinstance(parsed, dict):
-        return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}")
-    return ModelCostMapReloaded(
-        model_cost_map=parsed, revision=git_blob_id(response.content), etag=response.headers.get("etag")
-    )
+    return _classify_fetch_content(response.content, url, response.headers.get("etag"))
 
 
 def _next_retry_wait(

litellm/litellm_core_utils/litellm_logging.py

@@ -2300,6 +2300,13 @@ class Logging(LiteLLMLoggingBaseClass):
         self.model_call_details["combined_usage_object"] = usage
         self.model_call_details["response_cost"] = response_cost
 
+    def attempt_scoped_copy(self) -> "Logging":
+        """A copy with its own model_call_details, so logging one attempt of a request through it leaves
+        the state the request's other attempts share untouched."""
+        attempt: Final = copy.copy(self)
+        attempt.model_call_details = dict(self.model_call_details)
+        return attempt
+
     def carry_failed_attempt_usage(self) -> None:
         """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its
         own usage while the request's failure log still bills every attempt."""
@@ -4345,6 +4352,9 @@ class Logging(LiteLLMLoggingBaseClass):
         row is built from this response rather than re-parsed from the served body."""
         self.bridged_model_response = response
 
+    def record_bridged_responses_api_response(self, response: ResponsesAPIResponse) -> None:
+        self.record_bridged_model_response(self._translate_responses_api_response_to_model_response(response))
+
     def billable_result(self, result: object) -> object:
         """The response the cost calculator prices: a /v1/messages body served through the chat
         bridge is replaced by the ModelResponse its spend row is built from, anything else is
@@ -4445,6 +4455,7 @@ class Logging(LiteLLMLoggingBaseClass):
         """
         from litellm.completion_extras.litellm_responses_transformation.transformation import (
             LiteLLMResponsesTransformationHandler,
+            responses_provider_metadata,
         )
 
         served_id: Final = _provider_response_id(result)
@@ -4469,7 +4480,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 "usage-only ModelResponse to keep the spend_logs row.",
                 str(e),
             )
-            model_response: Final = litellm.ModelResponse(id=served_id)
+            model_response: Final = litellm.ModelResponse(
+                id=served_id, **responses_provider_metadata(result.model_extra)
+            )
             model_response.model = self.model
             usage: Final = result.usage
             if usage is not None and ResponseAPILoggingUtils.is_response_api_usage(usage):

litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py

@@ -858,9 +858,9 @@ class StandardBuiltInToolCostTracking:
 
     @staticmethod
     def get_web_search_options(kwargs: Mapping[str, object]) -> WebSearchOptions | None:
-        if "web_search_options" in kwargs:
-            web_search_options: Final = cast(WebSearchOptions, kwargs.get("web_search_options", {}))
-            return WebSearchOptions(**web_search_options)
+        web_search_options: Final = kwargs.get("web_search_options")
+        if isinstance(web_search_options, dict):
+            return WebSearchOptions(**cast(WebSearchOptions, web_search_options))
 
         tools: Final = StandardBuiltInToolCostTracking._get_tools_from_kwargs(
             kwargs=kwargs, tool_type="web_search_preview"
@@ -876,10 +876,9 @@ class StandardBuiltInToolCostTracking:
     _get_web_search_options = get_web_search_options
 
     @staticmethod
-    def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> list[object] | None:
-        if "tools" in kwargs:
-            return cast(list[object], kwargs.get("tools", []))
-        return None
+    def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> Sequence[object] | None:
+        tools: Final = kwargs.get("tools")
+        return cast(Sequence[object], tools) if isinstance(tools, (list, tuple)) else None
 
     @staticmethod
     def get_file_search_tool_call(kwargs: Mapping[str, object]) -> FileSearchTool | None:

litellm/litellm_core_utils/tool_fields.py

@@ -0,0 +1,164 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import reduce
+from typing import Annotated, Final, TypeAlias
+
+from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError
+from typing_extensions import ReadOnly, TypedDict
+
+import litellm
+from litellm.litellm_core_utils.core_helpers import normalize_drop_params
+
+ROUTES_WITH_TOOLS: Final = frozenset(
+    {
+        "acompletion",
+        "aresponses",
+        "acompact_responses",
+        "anthropic_messages",
+        "agenerate_content",
+        "agenerate_content_stream",
+        "acreate_interaction",
+    }
+)
+TOOL_FIELDS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls", "toolConfig", "tool_config"})
+_TOOLS: Final = TypeAdapter(list[object])
+_FIELD_LIST: Final = TypeAdapter(list[object], config=ConfigDict(strict=True))
+_JSON_OBJECT: Final = TypeAdapter(Json[dict[str, object]])
+
+
+def _schema_as_read(value: object) -> object:
+    if not value:
+        return None
+    if not isinstance(value, str):
+        return value
+    try:
+        return _JSON_OBJECT.validate_python(value)
+    except ValidationError:
+        return value
+
+
+_Schema: TypeAlias = Annotated[dict[str, object] | None, BeforeValidator(_schema_as_read)]
+
+
+class _ToolFunction(TypedDict, total=False):
+    parameters: ReadOnly[_Schema]
+
+
+class _ToolEntry(TypedDict, total=False):
+    function: ReadOnly[_ToolFunction | None]
+    parameters: ReadOnly[_Schema]
+    input_schema: ReadOnly[_Schema]
+    vector_store_ids: ReadOnly[list[str] | None]
+
+
+_TOOL_ENTRY: Final = TypeAdapter(_ToolEntry)
+_WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True))
+
+
+@dataclass(frozen=True, slots=True)
+class DropPolicy:
+    drop_params: bool
+    additional_drop_params: frozenset[str]
+
+    def drops(self, field: str) -> bool:
+        return self.drop_params or field in self.additional_drop_params
+
+
+@dataclass(frozen=True, slots=True)
+class MalformedField:
+    field: str
+    message: str
+    dropped: Callable[[Mapping[str, object]], Mapping[str, object]]
+
+
+def _field_names(value: object) -> frozenset[str]:
+    try:
+        return frozenset(name for name in _FIELD_LIST.validate_python(value) if isinstance(name, str))
+    except ValidationError:
+        return frozenset()
+
+
+def drop_policy(params: Mapping[str, object]) -> DropPolicy:
+    return DropPolicy(
+        drop_params=litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True,
+        additional_drop_params=_field_names(params.get("additional_drop_params")),
+    )
+
+
+def first_validation_error(e: ValidationError, prefix: tuple[str | int, ...] = ()) -> str:
+    error: Final = e.errors(include_url=False)[0]
+    return f"{'.'.join(str(part) for part in (*prefix, *error['loc']))}: {error['msg']}"
+
+
+def _tool_entry_error(index: int, entry: object) -> str | None:
+    try:
+        _TOOL_ENTRY.validate_python(entry, strict=True)
+    except ValidationError as e:
+        return first_validation_error(e, prefix=("tools", index))
+    return None
+
+
+def _without_tool_fields(data: Mapping[str, object]) -> Mapping[str, object]:
+    return {name: value for name, value in data.items() if name not in TOOL_FIELDS}
+
+
+def _malformed_tools(data: Mapping[str, object]) -> MalformedField | None:
+    tools: Final = data.get("tools")
+    if tools is None:
+        return None
+    try:
+        entries: Final = _TOOLS.validate_python(tools)
+    except ValidationError as e:
+        return MalformedField("tools", first_validation_error(e, prefix=("tools",)), _without_tool_fields)
+    entry_errors: Final = tuple(_tool_entry_error(index, entry) for index, entry in enumerate(entries))
+    first_error: Final = next((error for error in entry_errors if error is not None), None)
+    if first_error is None:
+        return None
+    kept: Final = [entry for entry, error in zip(entries, entry_errors) if error is None]
+    return MalformedField(
+        "tools", first_error, lambda fields: {**fields, "tools": kept} if kept else _without_tool_fields(fields)
+    )
+
+
+def _without_web_search_options(data: Mapping[str, object]) -> Mapping[str, object]:
+    return {name: value for name, value in data.items() if name != "web_search_options"}
+
+
+def _malformed_web_search_options(data: Mapping[str, object]) -> MalformedField | None:
+    try:
+        _WEB_SEARCH_OPTIONS.validate_python(data.get("web_search_options"))
+    except ValidationError as e:
+        return MalformedField(
+            "web_search_options",
+            first_validation_error(e, prefix=("web_search_options",)),
+            _without_web_search_options,
+        )
+    return None
+
+
+def resolve_tool_fields(
+    data: Mapping[str, object], candidates: Callable[[], tuple[DropPolicy, ...]]
+) -> Mapping[str, object] | MalformedField:
+    checked: Final = (_malformed_tools(data), _malformed_web_search_options(data))
+    malformed: Final = tuple(field for field in checked if field is not None)
+    if not malformed:
+        return data
+    policies: Final = candidates()
+    rejected: Final = next(
+        (field for field in malformed if not any(policy.drops(field.field) for policy in policies)), None
+    )
+    if rejected is not None:
+        return rejected
+    return reduce(
+        lambda resolved, field: field.dropped(resolved),
+        (field for field in malformed if all(policy.drops(field.field) for policy in policies)),
+        data,
+    )
+
+
+def resolve_deployment_tool_fields(
+    kwargs: Mapping[str, object], route: object
+) -> Mapping[str, object] | MalformedField:
+    if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS:
+        return kwargs
+    return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),))

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -45,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_scan_only_tool_results_for_guardrail,
     effective_skip_system_message_for_guardrail,
     effective_skip_tool_message_for_guardrail,
+    list_field,
     merge_guardrailed_scoped_messages,
     merge_returned_tools_into_request_tools,
     scoped_structured_message_indices,
@@ -924,10 +925,10 @@ class AnthropicMessagesHandler(BaseTranslation):
             images=(),
         )
 
-    def extract_request_tool_names(self, data: dict) -> list[str]:
+    def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]:
         """Extract every tool name in an Anthropic messages request: tools[].name, plus
         tools[].function.name for OpenAI-format tools the bridge forwards verbatim."""
-        return [name for tool in data.get("tools") or [] for name in anthropic_tool_names(tool)]
+        return [name for tool in list_field(data, "tools") for name in anthropic_tool_names(tool)]
 
     @classmethod
     def _extract_input_text_and_images(

litellm/llms/anthropic/pass_through/messages/streaming_iterator.py

@@ -2,6 +2,7 @@ import asyncio
 import json
 import re
 from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
+from contextvars import ContextVar
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
 
@@ -352,6 +353,22 @@ async def aclose_if_supported(stream: object) -> None:
         await stream.aclose()
 
 
+_CLOSING_SUPERSEDED_ATTEMPT: Final[ContextVar[bool]] = ContextVar("closing_superseded_attempt", default=False)
+
+
+async def aclose_superseded_attempt(stream: object) -> None:
+    """Close an attempt a retry or fallback is taking over, so its end-of-stream logging records a failure."""
+    token: Final = _CLOSING_SUPERSEDED_ATTEMPT.set(True)
+    try:
+        await aclose_if_supported(stream)
+    finally:
+        _CLOSING_SUPERSEDED_ATTEMPT.reset(token)
+
+
+def closing_superseded_attempt() -> bool:
+    return _CLOSING_SUPERSEDED_ATTEMPT.get()
+
+
 _RESPONSE_HEADERS_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
 
 

litellm/llms/anthropic/pass_through/responses_adapters/handler.py

@@ -43,6 +43,12 @@ _BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]])
 _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter()
 
 
+def _record_bridged_response(responses_kwargs: Mapping[str, object], result: ResponsesAPIResponse) -> None:
+    logging_obj: Final = litellm_logging_obj_from_kwargs(responses_kwargs)
+    if logging_obj is not None:
+        logging_obj.record_bridged_responses_api_response(result)
+
+
 def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, object]:
     """The litellm-specific kwargs forwarded verbatim onto the Responses API request."""
     return extra_kwargs or {}
@@ -301,7 +307,9 @@ class LiteLLMMessagesToResponsesAPIHandler:
         if not isinstance(result, ResponsesAPIResponse):
             raise ValueError(f"Expected ResponsesAPIResponse, got {type(result)}")
 
-        return await with_safeguard_results(_ADAPTER.translate_response(result), safeguards_evaluator)
+        served: Final = await with_safeguard_results(_ADAPTER.translate_response(result), safeguards_evaluator)
+        _record_bridged_response(responses_kwargs, result)
+        return served
 
     @staticmethod
     def anthropic_messages_handler(

litellm/llms/azure/chat/gpt_5_transformation.py

@@ -95,7 +95,6 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
                     status_code=400,
                     message=(
                         "Azure OpenAI does not support reasoning_effort='none' for this model. "
-                        "Supported values are: 'low', 'medium', and 'high'. "
                         "To drop this parameter, set `litellm.drop_params=True` or for proxy:\n\n"
                         "`litellm_settings:\n drop_params: true`\n"
                         "Issue: https://github.com/BerriAI/litellm/issues/16704"

litellm/llms/azure/common_utils.py

@@ -2,6 +2,7 @@ import asyncio
 import hashlib
 import json
 import os
+import re
 from collections.abc import Callable, Mapping
 from functools import lru_cache
 from types import MappingProxyType
@@ -25,6 +26,7 @@ from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import add_path_to_api_base
 
 azure_ad_cache: Final = DualCache()
+_LAST_OPENAI_PATH_SEGMENT: Final = re.compile(r"^(.*/openai)(?=/|$)")
 
 
 class _AzureAdTokenJson(TypedDict, total=False):
@@ -798,6 +800,10 @@ class BaseAzureLLM(BaseOpenAILLM):
     ) -> dict[str, str]:  # mutable-ok: mirrors override contract
         return cls._base_validate_azure_environment(headers, litellm_params)
 
+    @staticmethod
+    def resolve_azure_api_base(api_base: str | None) -> str | None:
+        return api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
+
     @staticmethod
     def _get_base_azure_url(
         api_base: str | None,
@@ -815,7 +821,7 @@ class BaseAzureLLM(BaseOpenAILLM):
             default_api_version: The default API version to use if no api_version is provided. If 'latest', it will use `openai/v1/...` route.
         """
 
-        api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
+        api_base = BaseAzureLLM.resolve_azure_api_base(api_base)
         if api_base is None:
             raise ValueError(
                 f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`"
@@ -843,7 +849,7 @@ class BaseAzureLLM(BaseOpenAILLM):
             # ensure the request go to /openai/v1 and not just /openai
             if "/openai/v1" not in new_url:
                 parsed_url: Final = httpx.URL(new_url)
-                new_url = str(parsed_url.copy_with(path=parsed_url.path.replace("/openai", "/openai/v1")))
+                new_url = str(parsed_url.copy_with(path=_LAST_OPENAI_PATH_SEGMENT.sub(r"\1/v1", parsed_url.path, 1)))
 
         # Use the new query_params dictionary
         final_url: Final = httpx.URL(new_url).copy_with(params=query_params)

litellm/llms/azure/containers/transformation.py

@@ -28,6 +28,7 @@ class AzureContainerConfig(OpenAIContainerConfig):
         self,
         headers: dict,
         api_key: str | None = None,
+        litellm_params: GenericLiteLLMParams | None = None,
     ) -> dict:
         return BaseAzureLLM.base_validate_azure_environment(
             headers=headers,

litellm/llms/azure/fine_tuning/handler.py

@@ -1,4 +1,4 @@
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Mapping
 from typing import Final, cast
 
 import httpx
@@ -66,6 +66,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         self._ensure_training_type(create_fine_tuning_job_data)
 
@@ -78,6 +79,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -109,6 +111,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -119,6 +122,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -149,6 +153,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -159,6 +164,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -188,14 +194,14 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM):
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
         _is_async: bool = False,
         api_version: str | None = None,
-        litellm_params: dict | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None:
         # Override to use Azure-specific client initialization
         if isinstance(client, OpenAI) or isinstance(client, AsyncOpenAI):
             client = None
 
         return self.get_azure_openai_client(
-            litellm_params=litellm_params or {},
+            litellm_params=dict(litellm_params) if litellm_params is not None else None,
             api_key=api_key,
             api_base=api_base,
             api_version=api_version,

litellm/llms/azure/responses/transformation.py

@@ -1,4 +1,7 @@
+import re
+from collections.abc import Mapping
 from copy import deepcopy
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 import httpx
@@ -6,6 +9,7 @@ from openai.types.responses import ResponseReasoningItem
 from pydantic import TypeAdapter
 
 from litellm._logging import verbose_logger
+from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
 from litellm.llms.azure.common_utils import BaseAzureLLM
@@ -25,6 +29,34 @@ else:
 _DECODED_JSON: Final = TypeAdapter(object)
 
 
+_DATED_GA_API_VERSION: Final = re.compile(r"\d{4}-\d{2}-\d{2}")
+
+
+def _str_param(litellm_params: Mapping[str, object], key: str) -> str | None:
+    value: Final = litellm_params.get(key)
+    return value if isinstance(value, str) and value else None
+
+
+def _resolve_responses_api_version(api_base: str | None, litellm_params: Mapping[str, object]) -> str | None:
+    """
+    Dated GA api_versions (e.g. 2024-10-21) predate the Responses API on Azure and 404 there, so a deployment
+    pinned to one for chat gets the default Responses version unless it sets responses_api_version.
+    """
+    api_version: Final = _str_param(litellm_params, "api_version")
+    if api_base is not None and "api-version" in httpx.URL(api_base).params:
+        return api_version
+    responses_api_version: Final = _str_param(litellm_params, "responses_api_version")
+    if responses_api_version is not None:
+        return responses_api_version
+    if api_version is not None and _DATED_GA_API_VERSION.fullmatch(api_version):
+        return AZURE_DEFAULT_RESPONSES_API_VERSION
+    return api_version
+
+
+def _with_responses_api_version(api_base: str | None, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+    return MappingProxyType({**litellm_params, "api_version": _resolve_responses_api_version(api_base, litellm_params)})
+
+
 class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     # Parameters not supported by Azure Responses API
     AZURE_UNSUPPORTED_PARAMS = ["context_management"]
@@ -170,11 +202,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
         - A complete URL string, e.g.,
         "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
         """
-        from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION
-
+        resolved_api_base: Final = BaseAzureLLM.resolve_azure_api_base(api_base)
         return BaseAzureLLM.get_base_azure_url(
-            api_base=api_base,
-            litellm_params=litellm_params,
+            api_base=resolved_api_base,
+            litellm_params=_with_responses_api_version(resolved_api_base, litellm_params),
             route="/openai/responses",
             default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION,
         )

litellm/llms/azure_ai/chat/transformation.py

@@ -1,6 +1,7 @@
 import copy
 import enum
 import re
+from collections.abc import AsyncIterator, Iterator
 from typing import TYPE_CHECKING, Final, cast
 
 import httpx
@@ -21,12 +22,12 @@ from litellm.llms.azure_ai.common_utils import (
 from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
 from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
 from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
-from litellm.llms.openai.openai import OpenAIConfig
+from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig
 from litellm.llms.xai.chat.transformation import XAIChatConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import ModelResponse, ProviderField
+from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField
 from litellm.utils import add_path_to_api_base, supports_tool_choice
 
 if TYPE_CHECKING:
@@ -45,6 +46,18 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
 )
 
 
+def _is_grok_model(model: str) -> bool:
+    return "grok" in model
+
+
+class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator):
+    def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream:  # mutable-ok: mirrors override contract
+        usage: Final = chunk.get("usage")
+        if _is_grok_model(str(chunk.get("model") or "")) and isinstance(usage, dict):
+            XAIChatConfig.fold_reasoning_tokens_into_completion(usage)
+        return super().chunk_parser(chunk)
+
+
 class AzureAIGPT5Config(OpenAIGPT5Config):
     @classmethod
     def _model_map_lookup_name(cls, model: str) -> str:
@@ -114,7 +127,7 @@ class AzureAIStudioConfig(OpenAIConfig):
         """
         Check if the model supports stop tokens.
         """
-        if "grok" in model:
+        if _is_grok_model(model):
             # Reuse Xai method for Grok model
             xai_config: Final = XAIChatConfig()
             return xai_config.supports_stop_reason(model)
@@ -319,7 +332,7 @@ class AzureAIStudioConfig(OpenAIConfig):
         json_mode: bool | None = None,
     ) -> ModelResponse:
         model_response.model = f"azure_ai/{model}"
-        return super().transform_response(
+        response: Final = super().transform_response(
             model=model,
             raw_response=raw_response,
             model_response=model_response,
@@ -332,6 +345,21 @@ class AzureAIStudioConfig(OpenAIConfig):
             api_key=api_key,
             json_mode=json_mode,
         )
+        if _is_grok_model(model):
+            XAIChatConfig.fold_reasoning_tokens_into_completion(response)
+        return response
+
+    def get_model_response_iterator(
+        self,
+        streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
+        sync_stream: bool,
+        json_mode: bool | None = False,
+    ) -> "AzureAIChatCompletionResponseIterator":
+        return AzureAIChatCompletionResponseIterator(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
 
     def should_retry_llm_api_inside_llm_translation_on_http_error(
         self, e: httpx.HTTPStatusError, litellm_params: dict

litellm/llms/base_llm/containers/transformation.py

@@ -81,6 +81,7 @@ class BaseContainerConfig(ABC):
         self,
         headers: dict,
         api_key: str | None = None,
+        litellm_params: GenericLiteLLMParams | None = None,
     ) -> dict:
         return {}
 

litellm/llms/base_llm/decisions/transformation.py

@@ -495,6 +495,11 @@ class BaseDecisionsConfig(ABC):
     def resolve_connection(self, api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]:
         return self.resolve_api_base(api_base), self.resolve_api_key(api_key)
 
+    def resolve_credential(
+        self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object]
+    ) -> str | None:
+        return api_key
+
     @staticmethod
     def _first_secret(names: tuple[str, ...]) -> str | None:
         return next((value for value in (get_secret_str(name) for name in names) if value), None)

litellm/llms/base_llm/guardrail_translation/utils.py

@@ -4,7 +4,7 @@ import json
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from typing import Final, TypeVar, cast  # noqa: TID251  # a rebuilt chat row has no typed constructor across roles
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
@@ -292,6 +292,17 @@ def openai_tool_name(tool: object) -> str | None:
     return flat_name if isinstance(flat_name, str) else None
 
 
+_FIELDS: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(strict=True))
+_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True))
+
+
+def list_field(data: object, key: str) -> Sequence[object]:
+    try:
+        return _ITEMS.validate_python(_FIELDS.validate_python(data).get(key))
+    except ValidationError:
+        return ()
+
+
 def anthropic_tool_names(tool: object) -> tuple[str, ...]:
     """Every name a /v1/messages tool dict can act under: the flat Anthropic ``name`` plus
     ``function.name`` for OpenAI-format tools the bridge forwards verbatim. Allowlist checks

litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py

@@ -32,6 +32,7 @@ from litellm.llms.bedrock.common_utils import (
     tools_without_eager_input_streaming,
 )
 from litellm.types.llms.anthropic import (
+    ANTHROPIC_BETA_HEADER_VALUES,
     ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER,
     ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
     AnthropicThinkingParam,
@@ -270,7 +271,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages)
         is_thinking_display_updates_used: Final = self.is_thinking_display_updates_used(thinking)
 
-        user_beta_set: Final = set(get_anthropic_beta_from_headers(headers))
+        client_beta_set: Final = set(get_anthropic_beta_from_headers(headers))
+        user_beta_set: Final = (
+            client_beta_set
+            if "safeguards" in optional_params
+            else client_beta_set - {ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value}
+        )
         beta_set: Final = set(user_beta_set)
         auto_betas: Final = self.get_anthropic_beta_list(
             model=model,

litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py

@@ -90,6 +90,9 @@ class AmazonAnthropicClaudeMessagesConfig(
     def beta_headers_provider(self) -> str:
         return self.custom_llm_provider or "bedrock"
 
+    def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool:
+        return True
+
     BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys())
 
     def get_error_class(
@@ -552,8 +555,17 @@ class AmazonAnthropicClaudeMessagesConfig(
         if anthropic_model_info.is_eager_input_streaming_used(tools):
             beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER)
 
+        dangerous_tool_use_beta: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value
         if anthropic_messages_optional_request_params.get("safeguards") is not None:
-            beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value)
+            beta_set.add(dangerous_tool_use_beta)
+        elif dangerous_tool_use_beta in beta_set and self.drops_dangerous_tool_use_beta_without_safeguards():
+            beta_set.discard(dangerous_tool_use_beta)
+            verbose_logger.debug(
+                "Bedrock Invoke: dropping %s from client headers because the request has no safeguards. Bedrock "
+                "answers it with 'invalid beta flag' next to structured outputs before Opus 4.7, and on Sonnet 4.5, "
+                "Opus 4.5, and Opus 4.6 even alone.",
+                dangerous_tool_use_beta,
+            )
 
         self._filter_context_management_for_bedrock_invoke(
             anthropic_messages_request=anthropic_messages_request,

litellm/llms/bedrock/messages/mantle_transformation.py

@@ -67,6 +67,9 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
     def should_filter_anthropic_beta_headers(self) -> bool:
         return False
 
+    def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool:
+        return False
+
     def get_complete_url(
         self,
         api_base: str | None,

litellm/llms/custom_httpx/container_handler.py

@@ -5,6 +5,7 @@ This module provides a single generic handler that can process any container fil
 endpoint defined in endpoints.json, eliminating the need for individual handler methods.
 """
 
+import asyncio
 import json
 from collections.abc import Coroutine, Mapping, Sequence
 from pathlib import Path
@@ -236,6 +237,7 @@ def _request_headers(
     return container_provider_config.validate_environment(
         headers=extra_headers or {},
         api_key=litellm_params.get("api_key", None),
+        litellm_params=litellm_params,
     )
 
 
@@ -435,7 +437,7 @@ class GenericContainerHandler:
         http_client: Final = _async_http_client(client, litellm_params)
 
         # Build request
-        headers = _request_headers(container_provider_config, extra_headers, litellm_params)
+        headers = await asyncio.to_thread(_request_headers, container_provider_config, extra_headers, litellm_params)
         if extra_headers:
             headers.update(extra_headers)
 

litellm/llms/custom_httpx/llm_http_handler.py

@@ -7127,7 +7127,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = image_edit_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            image_edit_provider_config.validate_environment,
             api_key=litellm_params.api_key,
             headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
             model=model,
@@ -7586,7 +7587,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_generation_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_generation_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=video_generation_optional_request_params.get("extra_headers", {}) or {},
             model=model,
@@ -7781,7 +7783,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_content_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_content_provider_config.validate_environment,
             headers=extra_headers or {},
             model="",
             api_key=api_key,
@@ -7961,7 +7964,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_remix_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_remix_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=extra_headers or {},
             model="",
@@ -8124,7 +8128,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=extra_headers or {},
             model="",
@@ -8264,7 +8269,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=extra_headers or {},
             model="",
@@ -8436,7 +8442,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=extra_headers or {},
             model="",
@@ -8622,7 +8629,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_provider_config.validate_environment,
             api_key=api_key or litellm_params.get("api_key", None),
             headers=extra_headers or {},
             model="",
@@ -8754,7 +8762,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_list_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_list_provider_config.validate_environment,
             api_key=api_key,
             headers=extra_headers or {},
             model="",
@@ -8835,7 +8844,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_delete_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_delete_provider_config.validate_environment,
             api_key=api_key,
             headers=extra_headers or {},
             model="",
@@ -9019,7 +9029,8 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = video_status_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            video_status_provider_config.validate_environment,
             api_key=api_key,
             headers=extra_headers or {},
             model="",
@@ -9122,6 +9133,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         # Add Content-Type header for JSON requests
@@ -9197,9 +9209,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         # Add Content-Type header for JSON requests
@@ -9293,6 +9307,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9369,9 +9384,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9460,6 +9477,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9536,9 +9554,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9629,6 +9649,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9705,9 +9726,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9803,6 +9826,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9881,9 +9905,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -9972,6 +9998,7 @@ class BaseLLMHTTPHandler:
         headers: Final = container_provider_config.validate_environment(
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10047,9 +10074,11 @@ class BaseLLMHTTPHandler:
             async_httpx_client = client
 
         # Validate environment and get headers
-        headers: Final = container_provider_config.validate_environment(
+        headers: Final = await asyncio.to_thread(
+            container_provider_config.validate_environment,
             headers=extra_headers or {},
             api_key=litellm_params.get("api_key", None),
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10167,8 +10196,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10410,8 +10441,10 @@ class BaseLLMHTTPHandler:
 
         vector_store_provider_config.validate_create_vector_store()
 
-        headers: Final = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10541,8 +10574,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10663,8 +10698,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10805,8 +10842,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -10950,8 +10989,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -11068,8 +11109,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
 
         if extra_headers:
@@ -11210,8 +11253,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
         if extra_headers:
             headers.update(extra_headers)
@@ -11351,8 +11396,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
         if extra_headers:
             headers.update(extra_headers)
@@ -11482,8 +11529,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
         if extra_headers:
             headers.update(extra_headers)
@@ -11619,8 +11668,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
         if extra_headers:
             headers.update(extra_headers)
@@ -11762,8 +11813,10 @@ class BaseLLMHTTPHandler:
         else:
             async_httpx_client = client
 
-        headers: Final = vector_store_files_provider_config.validate_environment(
-            headers=extra_headers or {}, litellm_params=litellm_params
+        headers: Final = await asyncio.to_thread(
+            vector_store_files_provider_config.validate_environment,
+            headers=extra_headers or {},
+            litellm_params=litellm_params,
         )
         if extra_headers:
             headers.update(extra_headers)

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -39,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_scan_only_tool_results_for_guardrail,
     effective_skip_system_message_for_guardrail,
     effective_skip_tool_message_for_guardrail,
+    list_field,
     merge_guardrailed_scoped_messages,
     merge_returned_tools_into_request_tools,
     openai_tool_name,
@@ -156,7 +157,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
                 inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices]
             # Pass tools (function definitions) to the guardrail
             tools: Final = data.get("tools")
-            if tools and not scan_only_tool_results:
+            if isinstance(tools, list) and tools and not scan_only_tool_results:
                 inputs["tools"] = tools
             # Include model information if available
             model: Final = data.get("model")
@@ -258,19 +259,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             return "no scannable content after message scoping"
         return None if images else "no scannable content"
 
-    def extract_request_tool_names(self, data: dict) -> list[str]:
+    def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]:
         """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name,
         and the tools a replayed tool_search_output loads)."""
         names: Final[list[str]] = []
-        for tool in data.get("tools") or []:
+        for tool in list_field(data, "tools"):
             if isinstance(tool, dict) and tool.get("type") == "function":
                 fn = tool.get("function")
                 if isinstance(fn, dict) and fn.get("name"):
                     names.append(str(fn["name"]))
-        for fn in data.get("functions") or []:
+        for fn in list_field(data, "functions"):
             if isinstance(fn, dict) and fn.get("name"):
                 names.append(str(fn["name"]))
-        return [*names, *replayed_tool_search_tool_names(cast(object, data))]  # cast-ok: upcast, read field by field
+        return [*names, *replayed_tool_search_tool_names(data)]
 
     def _extract_inputs(
         self,

litellm/llms/openai/common_utils.py

@@ -407,6 +407,10 @@ class BaseOpenAILLM:
             http2=http2_enabled(),
         )
 
+    @classmethod
+    def get_sync_http_client(cls) -> httpx.Client | None:
+        return cls._get_sync_http_client()
+
 
 class OpenAICredentials(NamedTuple):
     api_base: str

litellm/llms/openai/completion/handler.py

@@ -1,9 +1,7 @@
 import json
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Final
 
-from openai import AsyncOpenAI, OpenAI
-
 import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@@ -12,7 +10,8 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser
 from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse
 from litellm.utils import ProviderConfigManager
 
-from ..common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient
+from ..common_utils import BaseOpenAILLM, OpenAIError
+from ..workload_identity import build_async_openai_client, build_openai_client
 from .transformation import OpenAITextCompletionConfig
 
 
@@ -43,7 +42,7 @@ class OpenAITextCompletion(BaseLLM):
         print_verbose: Callable | None = None,
         api_base: str | None = None,
         acompletion: bool = False,
-        litellm_params=None,
+        litellm_params: Mapping[str, object] | None = None,
         logger_fn=None,
         client=None,
         organization: str | None = None,
@@ -95,6 +94,7 @@ class OpenAITextCompletion(BaseLLM):
                         max_retries=max_retries,
                         client=client,
                         organization=organization,
+                        litellm_params=litellm_params,
                     )
                 else:
                     return self.acompletion(
@@ -109,6 +109,7 @@ class OpenAITextCompletion(BaseLLM):
                         max_retries=max_retries,
                         organization=organization,
                         client=client,
+                        litellm_params=litellm_params,
                     )
             elif optional_params.get("stream", False):
                 return self.streaming(
@@ -123,19 +124,22 @@ class OpenAITextCompletion(BaseLLM):
                     max_retries=max_retries,
                     client=client,
                     organization=organization,
+                    litellm_params=litellm_params,
                 )
             else:
-                if client is None:
-                    openai_client = OpenAI(
+                openai_client: Final = (
+                    client
+                    if client is not None
+                    else build_openai_client(
                         api_key=api_key,
-                        base_url=api_base,
-                        http_client=litellm.client_session or OpenAIHTTPClient(),
+                        api_base=api_base,
                         timeout=timeout,
                         max_retries=max_retries,
                         organization=organization,
+                        litellm_params=litellm_params,
+                        static_key_http_client_factory=lambda: litellm.client_session,
                     )
-                else:
-                    openai_client = client
+                )
 
                 raw_response: Final = openai_client.completions.with_raw_response.create(**data)
                 response: Final = raw_response.parse()
@@ -175,19 +179,22 @@ class OpenAITextCompletion(BaseLLM):
         max_retries: int,
         organization: str | None = None,
         client=None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         try:
-            if client is None:
-                openai_aclient = AsyncOpenAI(
+            openai_aclient: Final = (
+                client
+                if client is not None
+                else build_async_openai_client(
                     api_key=api_key,
-                    base_url=api_base,
-                    http_client=BaseOpenAILLM.get_async_http_client(),
+                    api_base=api_base,
                     timeout=timeout,
                     max_retries=max_retries,
                     organization=organization,
+                    litellm_params=litellm_params,
+                    static_key_http_client_factory=BaseOpenAILLM.get_async_http_client,
                 )
-            else:
-                openai_aclient = client
+            )
 
             raw_response: Final = await openai_aclient.completions.with_raw_response.create(**data)
             response: Final = raw_response.parse()
@@ -228,18 +235,21 @@ class OpenAITextCompletion(BaseLLM):
         max_retries=None,
         client=None,
         organization=None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
-        if client is None:
-            openai_client = OpenAI(
+        openai_client: Final = (
+            client
+            if client is not None
+            else build_openai_client(
                 api_key=api_key,
-                base_url=api_base,
-                http_client=litellm.client_session or OpenAIHTTPClient(),
+                api_base=api_base,
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
+                static_key_http_client_factory=lambda: litellm.client_session,
             )
-        else:
-            openai_client = client
+        )
 
         try:
             raw_response: Final = openai_client.completions.with_raw_response.create(**data)
@@ -285,18 +295,21 @@ class OpenAITextCompletion(BaseLLM):
         api_base: str | None = None,
         client=None,
         organization=None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
-        if client is None:
-            openai_client = AsyncOpenAI(
+        openai_client: Final = (
+            client
+            if client is not None
+            else build_async_openai_client(
                 api_key=api_key,
-                base_url=api_base,
-                http_client=litellm.aclient_session or OpenAIAsyncHTTPClient(),
+                api_base=api_base,
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
+                static_key_http_client_factory=lambda: litellm.aclient_session,
             )
-        else:
-            openai_client = client
+        )
 
         raw_response: Final = await openai_client.completions.with_raw_response.create(**data)
         response: Final = raw_response.parse()

litellm/llms/openai/containers/transformation.py

@@ -23,6 +23,7 @@ from litellm.types.containers.main import (
 from litellm.types.router import GenericLiteLLMParams
 
 from ...base_llm.containers.transformation import BaseContainerConfig
+from ..workload_identity import resolve_openai_bearer_token
 from .utils import join_container_api_base_path
 
 if TYPE_CHECKING:
@@ -104,13 +105,15 @@ class OpenAIContainerConfig(BaseContainerConfig):
         self,
         headers: dict,
         api_key: str | None = None,
+        litellm_params: GenericLiteLLMParams | None = None,
     ) -> dict:
-        api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
-        headers.update(
-            {
-                "Authorization": f"Bearer {api_key}",
-            },
+        params: Final = litellm_params or GenericLiteLLMParams()
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=api_key or params.api_key,
+            api_base=params.api_base,
+            litellm_params=params.model_dump(exclude_none=True),
         )
+        headers.update({"Authorization": f"Bearer {bearer_token}"})
         return headers
 
     def get_complete_url(

litellm/llms/openai/decisions/transformation.py

@@ -6,6 +6,7 @@ from typing_extensions import assert_never
 
 import litellm
 from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, decisions_text
+from litellm.llms.openai.workload_identity import resolve_openai_bearer_token
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.decisions import (
     DecisionsImage,
@@ -349,6 +350,11 @@ class OpenAIDecisionsConfig(BaseDecisionsConfig):
     def resolve_api_key(self, api_key: str | None) -> str | None:
         return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
 
+    def resolve_credential(
+        self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object]
+    ) -> str | None:
+        return resolve_openai_bearer_token(api_key=api_key, api_base=api_base, litellm_params=litellm_params)
+
     def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool:
         return False
 

litellm/llms/openai/evals/transformation.py

@@ -13,6 +13,7 @@ from litellm.llms.base_llm.evals.transformation import (
     BaseEvalsAPIConfig,
     LiteLLMLoggingObj,
 )
+from litellm.llms.openai.workload_identity import resolve_openai_bearer_token
 from litellm.types.llms.openai_evals import (
     CancelEvalResponse,
     CancelRunResponse,
@@ -44,24 +45,14 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig):
         return LlmProviders.OPENAI
 
     def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
-        """Add OpenAI-specific headers"""
-        import litellm
-        from litellm.secret_managers.main import get_secret_str
-
-        # Get API key following OpenAI pattern
-        api_key = None
-        if litellm_params:
-            api_key = litellm_params.api_key
-
-        api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
-
-        if not api_key:
+        params: Final = litellm_params or GenericLiteLLMParams()
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True)
+        )
+        if not bearer_token:
             raise ValueError("OPENAI_API_KEY is required for Evals API")
-
-        # Add required headers
-        headers["Authorization"] = f"Bearer {api_key}"
+        headers["Authorization"] = f"Bearer {bearer_token}"
         headers["Content-Type"] = "application/json"
-
         return headers
 
     def get_complete_url(

litellm/llms/openai/fine_tuning/handler.py

@@ -6,7 +6,7 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from openai.types.fine_tuning import FineTuningJob
 
 from litellm._logging import verbose_logger
-from litellm.llms.openai.common_utils import OpenAIAsyncHTTPClient, OpenAIHTTPClient
+from litellm.llms.openai.workload_identity import build_async_openai_client, build_openai_client
 from litellm.types.utils import LiteLLMFineTuningJob
 
 _AZURE_STATUS_MAP: Final[Mapping[object, str]] = {
@@ -74,27 +74,27 @@ class OpenAIFineTuningAPI:
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
         _is_async: bool = False,
         api_version: str | None = None,
-        litellm_params: dict | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None:
-        received_args: Final = locals()
-        openai_client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None
-        if client is None:
-            data: Final = {}
-            for k, v in received_args.items():
-                if k == "self" or k == "client" or k == "_is_async":
-                    pass
-                elif k == "api_base" and v is not None:
-                    data["base_url"] = v
-                elif v is not None:
-                    data[k] = v
-            if _is_async is True:
-                openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
-            else:
-                openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
-        else:
-            openai_client = client
-
-        return openai_client
+        if client is not None:
+            return client
+        if _is_async:
+            return build_async_openai_client(
+                api_key=api_key,
+                api_base=api_base,
+                timeout=timeout,
+                max_retries=max_retries,
+                organization=organization,
+                litellm_params=litellm_params,
+            )
+        return build_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+        )
 
     async def acreate_fine_tuning_job(
         self,
@@ -116,6 +116,7 @@ class OpenAIFineTuningAPI:
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -126,6 +127,7 @@ class OpenAIFineTuningAPI:
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -164,6 +166,7 @@ class OpenAIFineTuningAPI:
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -174,6 +177,7 @@ class OpenAIFineTuningAPI:
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -212,6 +216,7 @@ class OpenAIFineTuningAPI:
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
         after: str | None = None,
         limit: int | None = None,
     ):
@@ -224,6 +229,7 @@ class OpenAIFineTuningAPI:
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -263,6 +269,7 @@ class OpenAIFineTuningAPI:
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
         openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -273,6 +280,7 @@ class OpenAIFineTuningAPI:
             client=client,
             _is_async=_is_async,
             api_version=api_version,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(

litellm/llms/openai/image_edit/transformation.py

@@ -17,6 +17,7 @@ from litellm.types.router import GenericLiteLLMParams
 from litellm.utils import ImageResponse
 
 from ..common_utils import OpenAIError
+from ..workload_identity import resolve_openai_bearer_token
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@@ -160,12 +161,10 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
         litellm_params: dict | None = None,
         api_base: str | None = None,
     ) -> dict:
-        api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
-        headers.update(
-            {
-                "Authorization": f"Bearer {api_key}",
-            }
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=api_key, api_base=api_base, litellm_params=litellm_params
         )
+        headers.update({"Authorization": f"Bearer {bearer_token}"})
         return headers
 
     def get_complete_url(

litellm/llms/openai/image_variations/handler.py

@@ -2,7 +2,7 @@
 OpenAI Image Variations Handler
 """
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Final
 
 import httpx
@@ -14,27 +14,58 @@ from litellm.utils import ProviderConfigManager
 
 from ...base_llm.image_variations.transformation import BaseImageVariationConfig
 from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj
-from ..common_utils import OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient
+from ..common_utils import OpenAIError
+from ..workload_identity import build_async_openai_client, build_openai_client
 
 
 class OpenAIImageVariationsHandler:
     def get_sync_client(
         self,
         client: OpenAI | None,
-        init_client_params: dict,
-    ):
+        api_key: str | None,
+        api_base: str,
+        timeout: float | httpx.Timeout | None,
+        max_retries: int,
+        organization: str | None,
+        litellm_params: Mapping[str, object],
+    ) -> OpenAI:
         if client is not None:
             return client
-        return OpenAI(**init_client_params, http_client=litellm.client_session or OpenAIHTTPClient())
+        return build_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+            static_key_http_client_factory=lambda: litellm.client_session,
+        )
 
-    def get_async_client(self, client: AsyncOpenAI | None, init_client_params: dict) -> AsyncOpenAI:
+    def get_async_client(
+        self,
+        client: AsyncOpenAI | None,
+        api_key: str | None,
+        api_base: str,
+        timeout: float | httpx.Timeout | None,
+        max_retries: int,
+        organization: str | None,
+        litellm_params: Mapping[str, object],
+    ) -> AsyncOpenAI:
         if client is not None:
             return client
-        return AsyncOpenAI(**init_client_params, http_client=litellm.aclient_session or OpenAIAsyncHTTPClient())
+        return build_async_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+            static_key_http_client_factory=lambda: litellm.aclient_session,
+        )
 
     async def async_image_variations(
         self,
-        api_key: str,
+        api_key: str | None,
         api_base: str,
         organization: str | None,
         client: AsyncOpenAI | None,
@@ -51,15 +82,15 @@ class OpenAIImageVariationsHandler:
         provider_config: BaseImageVariationConfig,
     ) -> ImageResponse:
         try:
-            init_client_params: Final = {
-                "api_key": api_key,
-                "base_url": api_base,
-                "timeout": timeout,
-                "max_retries": max_retries,
-                "organization": organization,
-            }
-
-            client = self.get_async_client(client=client, init_client_params=init_client_params)
+            client = self.get_async_client(
+                client=client,
+                api_key=api_key,
+                api_base=api_base,
+                timeout=timeout,
+                max_retries=max_retries,
+                organization=organization,
+                litellm_params=litellm_params,
+            )
 
             raw_response: Final = await client.images.with_raw_response.create_variation(**data)
             response: Final = raw_response.parse()
@@ -103,7 +134,7 @@ class OpenAIImageVariationsHandler:
     def image_variations(
         self,
         model_response: ImageResponse,
-        api_key: str,
+        api_key: str | None,
         api_base: str,
         model: str | None,
         image: FileTypes,
@@ -167,15 +198,15 @@ class OpenAIImageVariationsHandler:
                     litellm_params=litellm_params,
                 )
 
-            init_client_params: Final = {
-                "api_key": api_key,
-                "base_url": api_base,
-                "timeout": timeout,
-                "max_retries": max_retries,
-                "organization": organization,
-            }
-
-            client = self.get_sync_client(client=client, init_client_params=init_client_params)
+            client = self.get_sync_client(
+                client=client,
+                api_key=api_key,
+                api_base=api_base,
+                timeout=timeout,
+                max_retries=max_retries,
+                organization=organization,
+                litellm_params=litellm_params,
+            )
 
             raw_response: Final = client.images.with_raw_response.create_variation(**json_data)
             response: Final = raw_response.parse()

litellm/llms/openai/openai.py

@@ -56,16 +56,19 @@ from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig
 from .chat.o_series_transformation import OpenAIOSeriesConfig
 from .common_utils import (
     BaseOpenAILLM,
-    OpenAIAsyncHTTPClient,
     OpenAIError,
-    OpenAIHTTPClient,
     build_output_token_limit_response,
     drop_params_from_unprocessable_entity_error,
     is_openai_backed_api_base,
     is_output_token_limit_error,
     model_dump_provider_response,
 )
-from .workload_identity import resolve_openai_workload_identity_config
+from .workload_identity import (
+    build_async_openai_client,
+    build_openai_client,
+    openai_static_api_key,
+    resolve_openai_workload_identity_config,
+)
 
 openaiOSeriesConfig: Final = OpenAIOSeriesConfig()
 openAIGPT5Config: Final = OpenAIGPT5Config()
@@ -429,7 +432,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                     )
                     if workload_identity_config is not None
                     else AsyncOpenAI(
-                        api_key=api_key,
+                        api_key=openai_static_api_key(api_key),
                         base_url=api_base,
                         http_client=async_http_client,
                         timeout=timeout,
@@ -451,7 +454,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                     )
                     if workload_identity_config is not None
                     else OpenAI(
-                        api_key=api_key,
+                        api_key=openai_static_api_key(api_key),
                         base_url=api_base,
                         http_client=sync_http_client,
                         timeout=timeout,
@@ -1459,10 +1462,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         max_retries=None,
         organization: str | None = None,
         headers: dict | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         try:
             openai_aclient: Final = self._get_openai_client(
                 is_async=True,
+                litellm_params=litellm_params,
                 api_key=api_key,
                 api_base=api_base,
                 timeout=timeout,
@@ -1524,6 +1529,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         aimg_generation=None,
         organization: str | None = None,
         headers: dict | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> ImageResponse:
         data = {}
         try:
@@ -1545,10 +1551,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                     max_retries=max_retries,
                     organization=organization,
                     headers=headers,
+                    litellm_params=litellm_params,
                 )
 
             openai_client: Final[OpenAI] = self._get_openai_client(
                 is_async=False,
+                litellm_params=litellm_params,
                 api_key=api_key,
                 api_base=api_base,
                 timeout=timeout,
@@ -1627,6 +1635,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         client=None,
         shared_session: Optional["ClientSession"] = None,
         stream: bool = False,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> HttpxBinaryResponseContent:
         if aspeech is not None and aspeech is True:
             return self.async_audio_speech(
@@ -1644,6 +1653,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 client=client,
                 shared_session=shared_session,
                 stream=stream,
+                litellm_params=litellm_params,
             )
 
         openai_client: Final = self._get_openai_client(
@@ -1654,6 +1664,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             max_retries=max_retries,
             client=client,
             shared_session=shared_session,
+            litellm_params=litellm_params,
         )
 
         sync_client: Final = cast(OpenAI, openai_client)
@@ -1692,6 +1703,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         client=None,
         shared_session: Optional["ClientSession"] = None,
         stream: bool = False,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> HttpxBinaryResponseContent:
         openai_client: Final = cast(
             AsyncOpenAI,
@@ -1703,6 +1715,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 max_retries=max_retries,
                 client=client,
                 shared_session=shared_session,
+                litellm_params=litellm_params,
             ),
         )
 
@@ -1749,26 +1762,27 @@ class OpenAIFilesAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
         _is_async: bool = False,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAI | AsyncOpenAI | None:
-        received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals())
-        openai_client: OpenAI | AsyncOpenAI | None = None
-        if client is None:
-            data: Final = {}
-            for k, v in received_args.items():
-                if k == "self" or k == "client" or k == "_is_async":
-                    pass
-                elif k == "api_base" and v is not None:
-                    data["base_url"] = v
-                elif v is not None:
-                    data[k] = v
-            if _is_async is True:
-                openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
-            else:
-                openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
-        else:
-            openai_client = client
-
-        return openai_client
+        if client is not None:
+            return client
+        if _is_async:
+            return build_async_openai_client(
+                api_key=api_key,
+                api_base=api_base,
+                timeout=timeout,
+                max_retries=max_retries,
+                organization=organization,
+                litellm_params=litellm_params,
+            )
+        return build_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+        )
 
     async def acreate_file(
         self,
@@ -1788,6 +1802,7 @@ class OpenAIFilesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]:
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -1797,6 +1812,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -1830,6 +1846,7 @@ class OpenAIFilesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]:
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -1839,6 +1856,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -1895,6 +1913,7 @@ class OpenAIFilesAPI(BaseLLM):
         organization: str | None,
         chunk_size: int = 1024 * 1024,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> FileContentStreamingResult:
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -1904,6 +1923,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -1958,6 +1978,7 @@ class OpenAIFilesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -1967,6 +1988,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2004,6 +2026,7 @@ class OpenAIFilesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2013,6 +2036,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2053,6 +2077,7 @@ class OpenAIFilesAPI(BaseLLM):
         organization: str | None,
         purpose: str | None = None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2062,6 +2087,7 @@ class OpenAIFilesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2107,26 +2133,27 @@ class OpenAIBatchesAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
         _is_async: bool = False,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAI | AsyncOpenAI | None:
-        received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals())
-        openai_client: OpenAI | AsyncOpenAI | None = None
-        if client is None:
-            data: Final = {}
-            for k, v in received_args.items():
-                if k == "self" or k == "client" or k == "_is_async":
-                    pass
-                elif k == "api_base" and v is not None:
-                    data["base_url"] = v
-                elif v is not None:
-                    data[k] = v
-            if _is_async is True:
-                openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
-            else:
-                openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
-        else:
-            openai_client = client
-
-        return openai_client
+        if client is not None:
+            return client
+        if _is_async:
+            return build_async_openai_client(
+                api_key=api_key,
+                api_base=api_base,
+                timeout=timeout,
+                max_retries=max_retries,
+                organization=organization,
+                litellm_params=litellm_params,
+            )
+        return build_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+        )
 
     async def acreate_batch(
         self,
@@ -2146,6 +2173,7 @@ class OpenAIBatchesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | AsyncOpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]:
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2155,6 +2183,7 @@ class OpenAIBatchesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2190,6 +2219,7 @@ class OpenAIBatchesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2199,6 +2229,7 @@ class OpenAIBatchesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2233,6 +2264,7 @@ class OpenAIBatchesAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2242,6 +2274,7 @@ class OpenAIBatchesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2282,6 +2315,7 @@ class OpenAIBatchesAPI(BaseLLM):
         after: str | None = None,
         limit: int | None = None,
         client: OpenAI | None = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client(
             api_key=api_key,
@@ -2291,6 +2325,7 @@ class OpenAIBatchesAPI(BaseLLM):
             organization=organization,
             client=client,
             _is_async=_is_async,
+            litellm_params=litellm_params,
         )
         if openai_client is None:
             raise ValueError(
@@ -2319,22 +2354,19 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: OpenAI | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAI:
-        received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals())
-        if client is None:
-            data: Final = {}
-            for k, v in received_args.items():
-                if k == "self" or k == "client":
-                    pass
-                elif k == "api_base" and v is not None:
-                    data["base_url"] = v
-                elif v is not None:
-                    data[k] = v
-            openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
-        else:
-            openai_client = client
-
-        return openai_client
+        if client is not None:
+            return client
+        return build_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+        )
 
     def async_get_openai_client(
         self,
@@ -2344,22 +2376,19 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: AsyncOpenAI | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> AsyncOpenAI:
-        received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals())
-        if client is None:
-            data: Final = {}
-            for k, v in received_args.items():
-                if k == "self" or k == "client":
-                    pass
-                elif k == "api_base" and v is not None:
-                    data["base_url"] = v
-                elif v is not None:
-                    data[k] = v
-            openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
-        else:
-            openai_client = client
-
-        return openai_client
+        if client is not None:
+            return client
+        return build_async_openai_client(
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            max_retries=max_retries,
+            organization=organization,
+            litellm_params=litellm_params,
+        )
 
     ### ASSISTANTS ###
 
@@ -2375,6 +2404,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         limit: int | None = 20,
         before: str | None = None,
         after: str | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> AsyncCursorPage[Assistant]:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2382,6 +2413,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
         request_params: Final = {
@@ -2409,6 +2441,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         aget_assistants: Literal[True], 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, AsyncCursorPage[Assistant]]:
         ...
 
@@ -2422,6 +2456,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | None,
         aget_assistants: Literal[False] | None, 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> SyncCursorPage[Assistant]: 
         ...
 
@@ -2440,7 +2476,9 @@ class OpenAIAssistantsAPI(BaseLLM):
         limit: int | None = 20,
         before: str | None = None,
         after: str | None = None,
-    ):
+        *,
+        litellm_params: Mapping[str, object] | None = None,
+    ) -> Coroutine[None, None, AsyncCursorPage[Assistant]] | SyncCursorPage[Assistant]:
         if aget_assistants is not None and aget_assistants is True:
             return self.async_get_assistants(
                 api_key=api_key,
@@ -2448,6 +2486,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
             )
         openai_client: Final = self.get_openai_client(
@@ -2456,6 +2495,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2483,6 +2523,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         create_assistant_data: dict,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Assistant:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2490,6 +2532,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2507,6 +2550,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         create_assistant_data: dict,
         client=None,
         async_create_assistants=None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if async_create_assistants is not None and async_create_assistants is True:
             return self.async_create_assistants(
@@ -2515,6 +2560,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
                 create_assistant_data=create_assistant_data,
             )
@@ -2524,6 +2570,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2540,6 +2587,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         assistant_id: str,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> AssistantDeleted:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2547,6 +2596,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2564,6 +2614,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         assistant_id: str,
         client=None,
         async_delete_assistants=None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if async_delete_assistants is not None and async_delete_assistants is True:
             return self.async_delete_assistant(
@@ -2572,6 +2624,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
                 assistant_id=assistant_id,
             )
@@ -2581,6 +2634,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2599,6 +2653,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: AsyncOpenAI | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAIMessage:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2606,6 +2662,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2636,6 +2693,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         a_add_message: Literal[True], 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, OpenAIMessage]:
         ...
 
@@ -2651,6 +2710,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | None,
         a_add_message: Literal[False] | None, 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> OpenAIMessage: 
         ...
 
@@ -2667,6 +2728,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client=None,
         a_add_message: bool | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if a_add_message is not None and a_add_message is True:
             return self.a_add_message(
@@ -2677,6 +2740,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
             )
         openai_client: Final = self.get_openai_client(
@@ -2685,6 +2749,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2710,6 +2775,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: AsyncOpenAI | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> AsyncCursorPage[OpenAIMessage]:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2717,6 +2784,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2737,6 +2805,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         aget_messages: Literal[True], 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, AsyncCursorPage[OpenAIMessage]]:
         ...
 
@@ -2751,6 +2821,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | None,
         aget_messages: Literal[False] | None, 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> SyncCursorPage[OpenAIMessage]: 
         ...
 
@@ -2766,6 +2838,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client=None,
         aget_messages=None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if aget_messages is not None and aget_messages is True:
             return self.async_get_messages(
@@ -2775,6 +2849,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
             )
         openai_client: Final = self.get_openai_client(
@@ -2783,6 +2858,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2802,6 +2878,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         messages: Iterable[OpenAICreateThreadParamsMessage] | None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Thread:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2809,6 +2887,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2841,6 +2920,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         messages: Iterable[OpenAICreateThreadParamsMessage] | None,
         client: AsyncOpenAI | None,
         acreate_thread: Literal[True], 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, Thread]:
         ...
 
@@ -2856,6 +2937,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         messages: Iterable[OpenAICreateThreadParamsMessage] | None,
         client: OpenAI | None,
         acreate_thread: Literal[False] | None, 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Thread: 
         ...
 
@@ -2872,6 +2955,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         messages: Iterable[OpenAICreateThreadParamsMessage] | None,
         client=None,
         acreate_thread=None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         """
         Here's an example:
@@ -2891,6 +2976,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
                 messages=messages,
             )
@@ -2900,6 +2986,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2927,6 +3014,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: AsyncOpenAI | None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Thread:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -2934,6 +3023,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -2959,6 +3049,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: AsyncOpenAI | None,
         aget_thread: Literal[True], 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, Thread]:
         ...
 
@@ -2973,6 +3065,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client: OpenAI | None,
         aget_thread: Literal[False] | None, 
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Thread: 
         ...
 
@@ -2988,6 +3082,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         organization: str | None,
         client=None,
         aget_thread=None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if aget_thread is not None and aget_thread is True:
             return self.async_get_thread(
@@ -2997,6 +3093,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
             )
         openai_client: Final = self.get_openai_client(
@@ -3005,6 +3102,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -3038,6 +3136,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         max_retries: int | None,
         organization: str | None,
         client: AsyncOpenAI | None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Run:
         openai_client: Final = self.async_get_openai_client(
             api_key=api_key,
@@ -3045,6 +3145,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 
@@ -3149,6 +3250,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         client,
         arun_thread: Literal[True], 
         event_handler: AssistantEventHandler | None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Coroutine[None, None, Run]:
         ...
 
@@ -3171,6 +3274,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         client,
         arun_thread: Literal[False] | None, 
         event_handler: AssistantEventHandler | None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ) -> Run: 
         ...
 
@@ -3194,6 +3299,8 @@ class OpenAIAssistantsAPI(BaseLLM):
         client=None,
         arun_thread=None,
         event_handler: AssistantEventHandler | None = None,
+        *,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         if arun_thread is not None and arun_thread is True:
             if stream is not None and stream is True:
@@ -3203,6 +3310,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                     timeout=timeout,
                     max_retries=max_retries,
                     organization=organization,
+                    litellm_params=litellm_params,
                     client=client,
                 )
                 return self.async_run_thread_stream(
@@ -3230,6 +3338,7 @@ class OpenAIAssistantsAPI(BaseLLM):
                 timeout=timeout,
                 max_retries=max_retries,
                 organization=organization,
+                litellm_params=litellm_params,
                 client=client,
             )
         openai_client: Final = self.get_openai_client(
@@ -3238,6 +3347,7 @@ class OpenAIAssistantsAPI(BaseLLM):
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
+            litellm_params=litellm_params,
             client=client,
         )
 

litellm/llms/openai/responses/count_tokens/token_counter.py

@@ -27,6 +27,7 @@ from litellm.llms.openai.responses.count_tokens.transformation import (
 from litellm.llms.openai.workload_identity import (
     OpenAIWorkloadIdentityConfig,
     get_workload_identity_bearer_token,
+    openai_static_api_key,
     resolve_openai_workload_identity_config,
 )
 from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
@@ -35,13 +36,18 @@ from litellm.types.utils import CountContextManagement, LlmProviders, TokenCount
 openai_count_tokens_handler: Final = OpenAICountTokensHandler()
 
 
-def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None:
+def _workload_identity_config(
+    litellm_params: Mapping[str, object],
+) -> OpenAIWorkloadIdentityConfig | OpenAIError | None:
     api_base: Final = litellm_params.get("api_base")
-    return resolve_openai_workload_identity_config(
-        api_key=None,
-        api_base=api_base if isinstance(api_base, str) else None,
-        litellm_params=litellm_params,
-    )
+    try:
+        return resolve_openai_workload_identity_config(
+            api_key=None,
+            api_base=api_base if isinstance(api_base, str) else None,
+            litellm_params=litellm_params,
+        )
+    except OpenAIError as e:
+        return e
 
 
 class OpenAITokenCounter(BaseTokenCounter):
@@ -126,7 +132,18 @@ class OpenAITokenCounter(BaseTokenCounter):
         credentials: Final = get_openai_credentials(
             api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key")
         )
-        credential: Final = credentials.api_key or _workload_identity_config(litellm_params)
+        credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params)
+
+        if isinstance(credential, OpenAIError):
+            return TokenCountResponse(
+                total_tokens=0,
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type="openai_api",
+                error=True,
+                error_message=credential.message,
+                status_code=credential.status_code,
+            )
 
         if credential is None:
             return TokenCountResponse(

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -55,6 +55,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
 from litellm.llms.base_llm.guardrail_translation.utils import (
     blocked_responses_stream_usage,
     effective_skip_system_message_for_guardrail,
+    list_field,
     merge_guardrailed_scoped_messages,
     role_out_of_guardrail_scope,
     scoped_structured_message_indices,
@@ -764,11 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation):
             data.get("input"), data.get("instructions"), structured_messages, merged
         )
 
-    def extract_request_tool_names(self, data: dict) -> list[str]:
+    def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]:
         """Extract tool names from Responses API request (tools[].name for function
         and custom, tools[].server_label for mcp)."""
         names: Final[list[str]] = []
-        for tool in data.get("tools") or []:
+        for tool in list_field(data, "tools"):
             if not isinstance(tool, dict):
                 continue
             if tool.get("type") in ("function", "custom") and tool.get("name"):

litellm/llms/openai/transcriptions/handler.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final, Optional, cast
 
 import httpx
@@ -97,6 +98,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion):
                 max_retries=max_retries,
                 logging_obj=logging_obj,
                 shared_session=shared_session,
+                litellm_params=litellm_params,
             )
 
         openai_client: Final[OpenAI] = self._get_openai_client(
@@ -106,6 +108,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion):
             timeout=timeout,
             max_retries=max_retries,
             client=client,
+            litellm_params=litellm_params,
         )
 
         ## LOGGING
@@ -159,6 +162,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion):
         client=None,
         max_retries=None,
         shared_session: Optional["ClientSession"] = None,
+        litellm_params: Mapping[str, object] | None = None,
     ):
         try:
             openai_aclient: Final[AsyncOpenAI] = self._get_openai_client(
@@ -169,6 +173,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion):
                 max_retries=max_retries,
                 client=client,
                 shared_session=shared_session,
+                litellm_params=litellm_params,
             )
 
             ## LOGGING

litellm/llms/openai/vector_store_files/transformation.py

@@ -8,6 +8,7 @@ from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.vector_store_files.transformation import (
     BaseVectorStoreFilesConfig,
 )
+from litellm.llms.openai.workload_identity import resolve_openai_bearer_token
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_store_files import (
@@ -70,11 +71,13 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig):
         headers: dict[str, str],
         litellm_params: GenericLiteLLMParams | None,
     ) -> dict[str, str]:
-        litellm_params = litellm_params or GenericLiteLLMParams()
-        api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
+        params: Final = litellm_params or GenericLiteLLMParams()
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True)
+        )
         headers.update(
             {
-                "Authorization": f"Bearer {api_key}",
+                "Authorization": f"Bearer {bearer_token}",
                 "Content-Type": "application/json",
             }
         )

litellm/llms/openai/vector_stores/transformation.py

@@ -5,6 +5,7 @@ import httpx
 import litellm
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
+from litellm.llms.openai.workload_identity import resolve_openai_bearer_token
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_stores import (
@@ -48,11 +49,13 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
         }
 
     def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
-        litellm_params = litellm_params or GenericLiteLLMParams()
-        api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
+        params: Final = litellm_params or GenericLiteLLMParams()
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True)
+        )
         headers.update(
             {
-                "Authorization": f"Bearer {api_key}",
+                "Authorization": f"Bearer {bearer_token}",
                 "Content-Type": "application/json",
             }
         )

litellm/llms/openai/videos/transformation.py

@@ -8,11 +8,10 @@ from urllib.parse import quote
 import httpx
 from httpx._types import FileContent, FileTypes, RequestFiles
 
-import litellm
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
 from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils
-from litellm.secret_managers.main import get_secret_str
+from litellm.llms.openai.workload_identity import resolve_openai_bearer_token
 from litellm.types.llms.openai import CreateVideoRequest
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.videos.main import (
@@ -78,16 +77,13 @@ class OpenAIVideoConfig(BaseVideoConfig):
         api_key: str | None = None,
         litellm_params: GenericLiteLLMParams | None = None,
     ) -> dict:
-        # Use api_key from litellm_params if available, otherwise fall back to other sources
-        if litellm_params and litellm_params.api_key:
-            api_key = api_key or litellm_params.api_key
-
-        api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
-        headers.update(
-            {
-                "Authorization": f"Bearer {api_key}",
-            }
+        params: Final = litellm_params or GenericLiteLLMParams()
+        bearer_token: Final = resolve_openai_bearer_token(
+            api_key=api_key or params.api_key,
+            api_base=params.api_base,
+            litellm_params=params.model_dump(exclude_none=True),
         )
+        headers.update({"Authorization": f"Bearer {bearer_token}"})
         return headers
 
     def get_complete_url(

litellm/llms/openai/workload_identity.py

@@ -6,10 +6,22 @@ from functools import lru_cache
 from typing import TYPE_CHECKING, Final
 from urllib.parse import urlparse
 
-import litellm
-from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
+import httpx
+from openai import DEFAULT_MAX_RETRIES as OPENAI_SDK_DEFAULT_MAX_RETRIES
+from openai import NOT_GIVEN, AsyncOpenAI, NotGiven, OpenAI
 
-from .common_utils import OpenAIError, is_openai_backed_api_base
+import litellm
+from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
+from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
+from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS
+
+from .common_utils import (
+    BaseOpenAILLM,
+    OpenAIAsyncHTTPClient,
+    OpenAIError,
+    OpenAIHTTPClient,
+    is_openai_backed_api_base,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -17,6 +29,11 @@ if TYPE_CHECKING:
     from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth
 
 OPENAI_WIF_CLIENT_ID: Final = "litellm"
+_IDENTITY_PARAM_ENV_NAMES: Final = (
+    ("openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID"),
+    ("openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID"),
+    ("openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE"),
+)
 _SDK_UPGRADE_MESSAGE: Final = (
     "OpenAI workload identity federation requires openai>=2.32.0. "
     "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / "
@@ -51,17 +68,25 @@ def resolve_openai_workload_identity_config(
     )
     if static_api_key is not None:
         return None
-    effective_api_base: Final = (
-        api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE")
-    )
-    if not _targets_openai_api(effective_api_base):
+    return _configured_openai_workload_identity(api_base=api_base, litellm_params=litellm_params)
+
+
+def _effective_openai_api_base(api_base: str | None) -> str | None:
+    return api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE")
+
+
+def _configured_openai_workload_identity(
+    api_base: str | None,
+    litellm_params: Mapping[str, object] | None,
+) -> OpenAIWorkloadIdentityConfig | None:
+    if not _targets_openai_api(_effective_openai_api_base(api_base)):
         return None
-    identity_provider_id: Final = _config_value(
-        litellm_params, "openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID"
+    identity_provider_id, service_account_id, token_file = (
+        _config_value(litellm_params, param_key, env_name) for param_key, env_name in _IDENTITY_PARAM_ENV_NAMES
     )
-    service_account_id: Final = _config_value(litellm_params, "openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID")
-    token_file: Final = _config_value(litellm_params, "openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE")
     if identity_provider_id is None or service_account_id is None or token_file is None:
+        if _carries_deployment_identity(litellm_params):
+            raise OpenAIError(status_code=401, message=_incomplete_identity_message(litellm_params))
         return None
     return OpenAIWorkloadIdentityConfig(
         identity_provider_id=identity_provider_id,
@@ -81,11 +106,200 @@ async def get_workload_identity_bearer_token_for_api_base(api_base: str) -> str
     return await _workload_identity_auth(config).get_token_async()
 
 
+def openai_static_api_key(api_key: str | None) -> str | None:
+    return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
+
+
+def resolve_openai_bearer_token(
+    api_key: str | None,
+    api_base: str | None,
+    litellm_params: Mapping[str, object] | None,
+) -> str | None:
+    static_api_key: Final = openai_static_api_key(api_key)
+    workload_identity_config: Final = resolve_openai_workload_identity_config(
+        api_key=static_api_key, api_base=api_base, litellm_params=litellm_params
+    )
+    if workload_identity_config is None:
+        return static_api_key
+    return get_workload_identity_bearer_token(workload_identity_config)
+
+
+async def resolve_openai_bearer_token_async(
+    static_api_key: str | None,
+    api_base: str | None,
+    litellm_params: Mapping[str, object] | None,
+) -> str | None:
+    workload_identity_config: Final = resolve_openai_workload_identity_config(
+        api_key=static_api_key, api_base=api_base, litellm_params=litellm_params
+    )
+    if workload_identity_config is None:
+        return static_api_key
+    return await _workload_identity_auth(workload_identity_config).get_token_async()
+
+
+def _carries_deployment_identity(litellm_params: Mapping[str, object] | None) -> bool:
+    if litellm_params is None:
+        return False
+    return any(_param_str(litellm_params, key) is not None for key in OPENAI_WIF_KWARGS_KEYS)
+
+
+def _incomplete_identity_message(litellm_params: Mapping[str, object] | None) -> str:
+    missing: Final = ", ".join(
+        param_key
+        for param_key, env_name in _IDENTITY_PARAM_ENV_NAMES
+        if _config_value(litellm_params, param_key, env_name) is None
+    )
+    return (
+        f"OpenAI workload identity federation on this deployment is missing {missing}. "
+        "Set openai_identity_provider_id, openai_service_account_id, and openai_identity_token_file "
+        "on the deployment, or the OPENAI_IDENTITY_PROVIDER_ID / OPENAI_SERVICE_ACCOUNT_ID / "
+        "OPENAI_IDENTITY_TOKEN_FILE env vars they fall back to."
+    )
+
+
+def _param_str(litellm_params: Mapping[str, object], key: str) -> str | None:
+    value: Final = litellm_params.get(key)
+    return value if isinstance(value, str) and value else None
+
+
 def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None:
-    param_value: Final = litellm_params.get(param_key) if litellm_params is not None else None
-    if isinstance(param_value, str) and param_value:
-        return param_value
-    return normalize_nonempty_secret_str(get_secret_str(env_name))
+    param_value: Final = _param_str(litellm_params, param_key) if litellm_params is not None else None
+    return param_value or normalize_nonempty_secret_str(get_secret_str(env_name))
+
+
+def build_openai_client(
+    api_key: str | None,
+    api_base: str | None,
+    timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    max_retries: int | None = None,
+    organization: str | None = None,
+    litellm_params: Mapping[str, object] | None = None,
+    static_key_http_client_factory: Callable[[], httpx.Client | None] | None = None,
+) -> OpenAI:
+    workload_identity_config: Final = resolve_openai_workload_identity_config(
+        api_key=api_key, api_base=api_base, litellm_params=litellm_params
+    )
+    retries: Final = _max_retries_or_sdk_default(max_retries)
+    if workload_identity_config is None:
+        static_key_http_client: Final = (
+            None if static_key_http_client_factory is None else static_key_http_client_factory()
+        )
+        return OpenAI(
+            api_key=openai_static_api_key(api_key),
+            base_url=api_base,
+            timeout=timeout,
+            max_retries=retries,
+            organization=organization,
+            http_client=static_key_http_client or OpenAIHTTPClient(),
+        )
+    cache_params: Final = _client_cache_params(
+        is_async=False,
+        workload_identity_config=workload_identity_config,
+        api_base=api_base,
+        timeout=timeout,
+        max_retries=retries,
+        organization=organization,
+    )
+    cached: Final = BaseOpenAILLM.get_cached_openai_client(
+        client_initialization_params=cache_params, client_type="openai"
+    )
+    if isinstance(cached, OpenAI):
+        return cached
+    process_http_client: Final = BaseOpenAILLM.get_sync_http_client()
+    client: Final = OpenAI(
+        workload_identity=workload_identity_config.to_sdk_workload_identity(),
+        base_url=api_base,
+        timeout=timeout,
+        max_retries=retries,
+        organization=organization,
+        http_client=process_http_client,
+    )
+    BaseOpenAILLM.set_cached_openai_client(
+        openai_client=client,
+        client_type="openai",
+        client_initialization_params=cache_params,
+        litellm_owned_client=BaseOpenAILLM.owns_wrapped_http_client(process_http_client),
+    )
+    return client
+
+
+def build_async_openai_client(
+    api_key: str | None,
+    api_base: str | None,
+    timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    max_retries: int | None = None,
+    organization: str | None = None,
+    litellm_params: Mapping[str, object] | None = None,
+    static_key_http_client_factory: Callable[[], httpx.AsyncClient | None] | None = None,
+) -> AsyncOpenAI:
+    workload_identity_config: Final = resolve_openai_workload_identity_config(
+        api_key=api_key, api_base=api_base, litellm_params=litellm_params
+    )
+    retries: Final = _max_retries_or_sdk_default(max_retries)
+    if workload_identity_config is None:
+        static_key_http_client: Final = (
+            None if static_key_http_client_factory is None else static_key_http_client_factory()
+        )
+        return AsyncOpenAI(
+            api_key=openai_static_api_key(api_key),
+            base_url=api_base,
+            timeout=timeout,
+            max_retries=retries,
+            organization=organization,
+            http_client=static_key_http_client or OpenAIAsyncHTTPClient(),
+        )
+    cache_params: Final = _client_cache_params(
+        is_async=True,
+        workload_identity_config=workload_identity_config,
+        api_base=api_base,
+        timeout=timeout,
+        max_retries=retries,
+        organization=organization,
+    )
+    cached: Final = BaseOpenAILLM.get_cached_openai_client(
+        client_initialization_params=cache_params, client_type="openai"
+    )
+    if isinstance(cached, AsyncOpenAI):
+        return cached
+    process_http_client: Final = BaseOpenAILLM.get_async_http_client()
+    client: Final = AsyncOpenAI(
+        workload_identity=workload_identity_config.to_sdk_workload_identity(),
+        base_url=api_base,
+        timeout=timeout,
+        max_retries=retries,
+        organization=organization,
+        http_client=process_http_client,
+    )
+    BaseOpenAILLM.set_cached_openai_client(
+        openai_client=client,
+        client_type="openai",
+        client_initialization_params=cache_params,
+        litellm_owned_client=BaseOpenAILLM.owns_wrapped_http_client(process_http_client),
+    )
+    return client
+
+
+def _client_cache_params(
+    is_async: bool,
+    workload_identity_config: OpenAIWorkloadIdentityConfig,
+    api_base: str | None,
+    timeout: float | httpx.Timeout | None | NotGiven,
+    max_retries: int,
+    organization: str | None,
+) -> dict[str, object]:
+    return {
+        "api_key": None,
+        "is_async": is_async,
+        "workload_identity_config": workload_identity_config,
+        "api_base": api_base,
+        "timeout": timeout,
+        "max_retries": max_retries,
+        "organization": organization,
+    }
+
+
+def _max_retries_or_sdk_default(max_retries: int | None) -> int:
+    return OPENAI_SDK_DEFAULT_MAX_RETRIES if max_retries is None else max_retries
 
 
 def _targets_openai_api(api_base: str | None) -> bool:
@@ -95,7 +309,7 @@ def _targets_openai_api(api_base: str | None) -> bool:
     return parsed.scheme == "https" and is_openai_backed_api_base(api_base)
 
 
-@lru_cache(maxsize=16)
+@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
 def _workload_identity_auth(config: OpenAIWorkloadIdentityConfig) -> WorkloadIdentityAuth:
     sdk_workload_identity_auth: Final = _load_sdk_workload_identity_auth()
     return sdk_workload_identity_auth(workload_identity=config.to_sdk_workload_identity())

litellm/main.py

@@ -255,6 +255,7 @@ from .llms.openai.completion.handler import OpenAITextCompletion
 from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler
 from .llms.openai.openai import OpenAIChatCompletion
 from .llms.openai.transcriptions.handler import OpenAIAudioTranscription
+from .llms.openai.workload_identity import build_openai_client
 from .llms.openai_like.chat.handler import OpenAILikeChatHandler
 from .llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler
 from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig
@@ -1212,13 +1213,18 @@ def _register_custom_pricing_for_request(
         litellm.register_model({shared_key: entry}, persist_across_reloads=False)
         return
     litellm.register_model(
-        {
-            deployment_id: entry,
-            shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry),
-        },
+        {deployment_id: entry},
         persist_across_reloads=False,
         warning_display_name=shared_key,
         custom_llm_provider=custom_llm_provider,
+        add_to_provider_model_sets=False,
+    )
+    litellm.register_model(
+        {shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry)},
+        persist_across_reloads=False,
+        warning_display_name=shared_key,
+        custom_llm_provider=custom_llm_provider,
+        add_to_provider_model_sets=False,
     )
 
 
@@ -5785,6 +5791,7 @@ def completion(
                     "github_copilot_user_session",
                     PROVIDER_AFFINITY_HEADER_KWARG_KEY,
                     "fireworks_forward_user_id",
+                    "responses_api_version",
                 )
                 if key in kwargs
             },
@@ -7860,17 +7867,14 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompl
 
 
 def moderation(input: str, model: str | None = None, api_key: str | None = None, **kwargs) -> OpenAIModerationResponse:
-    from litellm.llms.openai.common_utils import OpenAIHTTPClient
-
-    # only supports open ai for now
-    api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
-
-    # Extract api_base from kwargs
+    static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
     api_base: Final = kwargs.get("api_base", None)
 
-    openai_client = kwargs.get("client", None)
-    if openai_client is None:
-        openai_client = openai.OpenAI(api_key=api_key, base_url=api_base, http_client=OpenAIHTTPClient())
+    openai_client: Final = kwargs.get("client", None) or build_openai_client(
+        api_key=static_api_key,
+        api_base=api_base,
+        litellm_params=kwargs,
+    )
 
     if model is not None:
         response = openai_client.moderations.create(input=input, model=model)
@@ -7893,8 +7897,7 @@ async def amoderation(
 ) -> OpenAIModerationResponse:
     from openai import AsyncOpenAI
 
-    # only supports open ai for now
-    api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
+    static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
     optional_params: Final = GenericLiteLLMParams.model_validate(kwargs)
     litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None)
     _dynamic_api_base = None
@@ -7920,8 +7923,9 @@ async def amoderation(
         # _get_openai_client maintains in-memory caching logic for OpenAI clients
         _openai_client: AsyncOpenAI = openai_chat_completions.get_openai_client(
             is_async=True,
-            api_key=api_key,
+            api_key=static_api_key,
             api_base=optional_params.api_base or _dynamic_api_base,
+            litellm_params=optional_params.model_dump(exclude_none=True),
         )
     else:
         _openai_client = openai_client
@@ -7941,7 +7945,7 @@ async def amoderation(
         moderation_request: Final = {"input": input, "model": model}
         litellm_logging_obj.pre_call(
             input=input,
-            api_key=api_key,
+            api_key=static_api_key,
             additional_args={
                 "complete_input_dict": moderation_request,
                 "api_base": str(_openai_client.base_url),
@@ -8174,9 +8178,7 @@ def transcription(
             or get_secret("OPENAI_ORGANIZATION")
             or None  # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
         )
-        # set API KEY
-
-        api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY")
+        static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY")
         response = openai_audio_transcriptions.audio_transcriptions(
             model=model,
             audio_file=file,
@@ -8188,7 +8190,7 @@ def transcription(
             logging_obj=litellm_logging_obj,
             max_retries=max_retries,
             api_base=api_base,
-            api_key=api_key,
+            api_key=static_api_key,
             provider_config=provider_config,
             litellm_params=litellm_params_dict,
             shared_session=shared_session,
@@ -8514,6 +8516,7 @@ def speech(
             aspeech=aspeech,
             shared_session=shared_session,
             stream=stream_audio,
+            litellm_params=litellm_params_dict,
         )
     elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
         # Check if this is Azure Speech Service (Cognitive Services TTS)

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -1345,7 +1345,9 @@ class MCPRequestHandler:
         if auth is None:
             return None
         refreshed = auth.model_copy()
-        if auth.mcp_admitted_user_subject and auth.user_id:
+        if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id:
+            # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is
+            # sealed in the token: refresh the live user policy and keep the token's team, routes and models.
             current = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=True)
             refreshed.object_permission = current.object_permission
             refreshed.object_permission_id = current.object_permission_id

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.faults import (
     render_token_fault,
 )
 from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
+    JWT_BEARER_GRANT_TYPE,
     VendorCredentialState,
     aggregate_authorize,
     aggregate_token,
@@ -71,6 +72,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
     revoke_refresh_token,
     supported_grant_types,
 )
+from litellm.proxy._experimental.mcp_server.id_jag_grant import (
+    basic_client_id,
+    jwt_bearer_available,
+    redeem_id_jag_assertion,
+)
 from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
     exchange_idp_subject_token,
     token_exchange_available,
@@ -2396,7 +2402,7 @@ async def token_endpoint(
     grant_type: str = Form(...),
     code: str = Form(None),
     redirect_uri: str = Form(None),
-    client_id: str = Form(...),
+    client_id: str | None = Form(None),
     client_secret: str | None = Form(None),
     code_verifier: str = Form(None),
     refresh_token: str | None = Form(None),
@@ -2405,6 +2411,7 @@ async def token_endpoint(
     subject_token: str | None = Form(None),
     subject_token_type: str | None = Form(None),
     requested_token_type: str | None = Form(None),
+    assertion: str | None = Form(None),
     mcp_server_name: str | None = None,
 ):
     """
@@ -2416,7 +2423,14 @@ async def token_endpoint(
     3. Return the token
     4. Return a virtual key in this response
     """
-    if mcp_server_name is None and is_gateway_dcr_client_id(client_id):
+    presented_client_id: Final = client_id or basic_client_id(request.headers.get("authorization"))
+    if presented_client_id is None:
+        raise HTTPException(
+            status_code=400, detail={"error": "invalid_request", "error_description": "client_id is required"}
+        )
+    if mcp_server_name is None and (
+        is_gateway_dcr_client_id(presented_client_id) or grant_type == JWT_BEARER_GRANT_TYPE
+    ):
         from litellm.proxy.proxy_server import (  # noqa: PLC0415  # circular import at module load
             master_key,
             user_api_key_cache,
@@ -2427,7 +2441,7 @@ async def token_endpoint(
             grant_type=grant_type,
             code=code,
             redirect_uri=redirect_uri,
-            client_id=client_id,
+            client_id=presented_client_id,
             code_verifier=code_verifier,
             refresh_token=refresh_token,
             master_key=master_key,
@@ -2439,9 +2453,12 @@ async def token_endpoint(
             subject_token_type=subject_token_type,
             requested_token_type=requested_token_type,
             exchange_subject_token=exchange_idp_subject_token,
+            assertion=assertion,
+            client_secret=client_secret,
+            redeem_assertion=redeem_id_jag_assertion,
         )
 
-    lookup_name: Final = mcp_server_name or client_id
+    lookup_name: Final = mcp_server_name or presented_client_id
     client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
     mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip)
     if mcp_server is None and mcp_server_name is None:
@@ -2454,7 +2471,7 @@ async def token_endpoint(
         grant_type=grant_type,
         code=code,
         redirect_uri=redirect_uri,
-        client_id=client_id,
+        client_id=presented_client_id,
         client_secret=client_secret,
         code_verifier=code_verifier,
         refresh_token=refresh_token,
@@ -2568,7 +2585,8 @@ async def native_client_auth_discovery(request: Request) -> JSONResponse:
     """The versioned contract a native client (``lite login --pkce``, or a CLI in any other
     language) reads to sign a user in through the browser and obtain a proxy credential."""
     return JSONResponse(
-        native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS
+        native_client_auth_contract(request, token_exchange_available(), jwt_bearer_available()),
+        headers=TOKEN_NO_CACHE_HEADERS,
     )
 
 
@@ -3135,7 +3153,7 @@ def _aggregate_header_scopes(request: Request) -> tuple[str, ...]:
 
 
 def _build_aggregate_authorization_server_response(
-    request: Request, token_exchange_available: bool
+    request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False
 ) -> dict[str, object]:
     """RFC 8414 metadata for the gateway as the aggregate authorization server.
 
@@ -3155,9 +3173,13 @@ def _build_aggregate_authorization_server_response(
         "registration_endpoint": f"{request_base_url}/register",
         "response_types_supported": ["code"],
         "scopes_supported": [],
-        "grant_types_supported": supported_grant_types(token_exchange_available),
+        "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available),
         "code_challenge_methods_supported": ["S256"],
-        "token_endpoint_auth_methods_supported": ["none", "client_secret_post"],
+        "token_endpoint_auth_methods_supported": (
+            ["none", "client_secret_post", "client_secret_basic"]
+            if jwt_bearer_available
+            else ["none", "client_secret_post"]
+        ),
     }
 
 
@@ -3193,7 +3215,7 @@ async def oauth_authorization_server_aggregate(request: Request):
     per-server row win here instead would serve an issuer of {base} against a resource that
     advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.
     """
-    return _build_aggregate_authorization_server_response(request, token_exchange_available())
+    return _build_aggregate_authorization_server_response(request, token_exchange_available(), jwt_bearer_available())
 
 
 # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
@@ -3273,6 +3295,10 @@ def _build_oauth_authorization_server_response(
         if explicitly_named
         else request_base_url
     )
+    # With no upstream resolved the document is the gateway's own: its issuer is the bare origin
+    # and its /token serves every grant the gateway does, including the ID-JAG jwt-bearer grant.
+    gateway_owned: Final = mcp_server_name is None and issuer_path is None
+    jwt_bearer: Final = gateway_owned and jwt_bearer_available()
 
     return {
         "issuer": issuer,
@@ -3280,9 +3306,15 @@ def _build_oauth_authorization_server_response(
         "token_endpoint": token_endpoint,
         "response_types_supported": ["code"],
         "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []),
-        "grant_types_supported": ["authorization_code", "refresh_token"],
+        "grant_types_supported": (
+            list(supported_grant_types(token_exchange_available(), jwt_bearer))
+            if gateway_owned
+            else ["authorization_code", "refresh_token"]
+        ),
         "code_challenge_methods_supported": ["S256"],
-        "token_endpoint_auth_methods_supported": ["client_secret_post"],
+        "token_endpoint_auth_methods_supported": (
+            ["none", "client_secret_post", "client_secret_basic"] if jwt_bearer else ["client_secret_post"]
+        ),
         # Claude expects a registration endpoint, even if we just fake it
         "registration_endpoint": (
             f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register"

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -198,20 +198,26 @@ class MintProxyCredential(Protocol):
     and mints the same credential ``lite login`` mints."""
 
     def __call__(
-        self, user_id: str, team_id: str | None, /
+        self, user_id: str, team_id: str | None, /, *, allowed_routes: Sequence[str] | None = None
     ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ...
 
 
 TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
+JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
 
 
-def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]:
+def supported_grant_types(token_exchange_available: bool, jwt_bearer_available: bool = False) -> tuple[str, ...]:
     """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed
     only where the JWT auth that proves a subject token is on, backed by a database, and
-    licensed, so a client never selects a grant the gateway would then refuse."""
-    if token_exchange_available:
-        return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)
-    return ("authorization_code", "refresh_token")
+    licensed, so a client never selects a grant the gateway would then refuse. The RFC 7523
+    jwt-bearer grant (an ID-JAG from the identity provider) is listed only where, on top of
+    that, the operator has allowlisted at least one requesting app."""
+    return (
+        "authorization_code",
+        "refresh_token",
+        *((TOKEN_EXCHANGE_GRANT_TYPE,) if token_exchange_available else ()),
+        *((JWT_BEARER_GRANT_TYPE,) if jwt_bearer_available else ()),
+    )
 
 
 """RFC 8693: a native client that already holds a token from the customer's identity
@@ -236,10 +242,29 @@ class SubjectIdentity(LiteLLMBaseModel):
 
 class SubjectTokenRefusal(LiteLLMBaseModel):
     model_config = ConfigDict(frozen=True)
-    error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"]
+    error: Literal[
+        "unsupported_grant_type",
+        "invalid_request",
+        "invalid_grant",
+        "invalid_client",
+        "invalid_scope",
+        "temporarily_unavailable",
+    ]
     description: str = Field(min_length=1)
 
 
+class AssertionIdentity(SubjectIdentity):
+    """The litellm user an RFC 7523 assertion (an ID-JAG) stands for, plus what the grant
+    tail needs from the assertion itself: the single-use key for its ``jti``, how long that
+    marker must outlive it, the ``scope`` the credential is issued for, and the route groups that
+    scope narrows the credential to (``None`` leaves it the full ``lite login`` credential)."""
+
+    replay_key: str = Field(min_length=1)
+    replay_ttl_seconds: int = Field(gt=0)
+    scope: str | None = None
+    allowed_routes: tuple[str, ...] | None = None
+
+
 class ExchangeSubjectToken(Protocol):
     """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the
     IdP token the way the proxy's own JWT auth does and names the litellm user and team it
@@ -272,7 +297,9 @@ class LookupConsentTeams(Protocol):
     def __call__(self, user_id: str, /) -> Awaitable[tuple[ConsentTeam, ...] | ReloadUserFailure]: ...
 
 
-async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCredentialMintFailure:
+async def _refuse_proxy_credential(
+    user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None
+) -> ProxyCredentialMintFailure:
     return "unresolvable"
 
 
@@ -282,6 +309,25 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject
     )
 
 
+class RedeemAssertion(Protocol):
+    """Injected RFC 7523 assertion verifier ``(assertion, client_id, client_secret, request)``:
+    proves an ID-JAG the identity provider minted for the requesting app ``client_id``,
+    authenticates that app, and names the litellm user the assertion stands for, or says why
+    this gateway will not take it."""
+
+    def __call__(
+        self, assertion: str, client_id: str, client_secret: str | None, request: Request, /
+    ) -> Awaitable[AssertionIdentity | SubjectTokenRefusal]: ...
+
+
+async def _refuse_assertion(
+    assertion: str, client_id: str, client_secret: str | None, request: Request
+) -> SubjectTokenRefusal:
+    return SubjectTokenRefusal(
+        error="unsupported_grant_type", description="this gateway is not configured to redeem identity assertions"
+    )
+
+
 async def _unavailable_vendor_credential(
     user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()
 ) -> VendorCredentialState:
@@ -797,7 +843,9 @@ class NativeClientAuthContract(TypedDict):
     revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]]
 
 
-def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract:
+def native_client_auth_contract(
+    request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False
+) -> NativeClientAuthContract:
     """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a
     native client (in any language) needs to run the sign-in without reading LiteLLM
     source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter
@@ -812,7 +860,7 @@ def native_client_auth_contract(request: Request, token_exchange_available: bool
         "revocation_endpoint": f"{base_url}/revoke",
         "resource": base_url,
         "response_types_supported": ("code",),
-        "grant_types_supported": supported_grant_types(token_exchange_available),
+        "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available),
         "code_challenge_methods_supported": ("S256",),
         "token_endpoint_auth_methods_supported": ("none",),
         "revocation_endpoint_auth_methods_supported": ("none",),
@@ -1403,10 +1451,11 @@ class _ProxyCredentialTokenResponse(TypedDict):
     access_token: ReadOnly[str]
     token_type: ReadOnly[Literal["Bearer"]]
     expires_in: ReadOnly[int]
-    refresh_token: ReadOnly[str]
+    refresh_token: NotRequired[ReadOnly[str]]
     user_id: ReadOnly[str]
     team_id: ReadOnly[str | None]
     issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]]
+    scope: NotRequired[ReadOnly[str]]
 
 
 def proxy_credential_response(
@@ -1415,6 +1464,7 @@ def proxy_credential_response(
     signing: SessionSigning,
     issued_token_type: _IssuedTokenType | None = None,
     family: str | None = None,
+    scope: str | None = None,
 ) -> Response:
     """The proxy-API token response: the access token is the very credential ``lite
     login`` stores (accepted on every proxy route with user and team attribution), and
@@ -1426,7 +1476,7 @@ def proxy_credential_response(
     refresh: Final = mint_session_refresh_token(bound_principal, signing.keys, signing.now, family=family)
     if not isinstance(refresh, MintedSessionToken):
         return _oauth_error(500, "server_error", "failed to mint the session credential")
-    credential: Final[_ProxyCredentialTokenResponse] = {
+    renewable: Final[_ProxyCredentialTokenResponse] = {
         "access_token": minted.key,
         "token_type": "Bearer",
         "expires_in": minted.expires_in,
@@ -1434,12 +1484,27 @@ def proxy_credential_response(
         "user_id": minted.user_id,
         "team_id": minted.team_id,
     }
+    credential: Final[_ProxyCredentialTokenResponse] = renewable if scope is None else {**renewable, "scope": scope}
     if issued_token_type is None:
         return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
     exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type}
     return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS)
 
 
+def assertion_credential_response(minted: MintedProxyCredential, scope: str | None) -> Response:
+    """The proxy-API token response for an assertion grant: the same access token, with the
+    scope the assertion carried, and no refresh token."""
+    unscoped: Final[_ProxyCredentialTokenResponse] = {
+        "access_token": minted.key,
+        "token_type": "Bearer",
+        "expires_in": minted.expires_in,
+        "user_id": minted.user_id,
+        "team_id": minted.team_id,
+    }
+    credential: Final[_ProxyCredentialTokenResponse] = unscoped if scope is None else {**unscoped, "scope": scope}
+    return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
+
+
 def _reload_failure_response(failure: ReloadUserFailure) -> Response:
     """Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new
     ``ReloadUserFailure`` member is a type error here rather than silently 400ing."""
@@ -1460,7 +1525,9 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response:
     match refusal.error:
         case "temporarily_unavailable":
             return _oauth_error(503, refusal.error, refusal.description)
-        case "unsupported_grant_type" | "invalid_request":
+        case "invalid_client":
+            return _oauth_error(401, refusal.error, refusal.description)
+        case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope":
             return _oauth_error(400, refusal.error, refusal.description)
         case _:
             assert_never(refusal.error)
@@ -1515,6 +1582,9 @@ async def aggregate_token(
     subject_token_type: str | None = None,
     requested_token_type: str | None = None,
     exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token,
+    assertion: str | None = None,
+    client_secret: str | None = None,
+    redeem_assertion: RedeemAssertion = _refuse_assertion,
 ) -> Response:
     """The aggregate token verb: authorization_code and refresh_token grants for the
     identity-only session pair, or for the proxy-API credential when the grant was issued
@@ -1561,10 +1631,19 @@ async def aggregate_token(
             exchange_subject_token=exchange_subject_token,
             issue=issue,
         )
+    if grant_type == JWT_BEARER_GRANT_TYPE:
+        return await _jwt_bearer_grant(
+            request=request,
+            assertion=assertion,
+            client_id=client_id,
+            client_secret=client_secret,
+            redeem_assertion=redeem_assertion,
+            issue=issue,
+        )
     return _oauth_error(
         400,
         "unsupported_grant_type",
-        f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}",
+        f"grant_type must be authorization_code, refresh_token, {TOKEN_EXCHANGE_GRANT_TYPE}, or {JWT_BEARER_GRANT_TYPE}",
     )
 
 
@@ -1690,6 +1769,48 @@ class _GrantIssuer:
             return _mint_failure_response(minted)
         return proxy_credential_response(minted, principal, self._signing, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE)
 
+    async def redeem_assertion(
+        self, assertion: str, client_id: str, client_secret: str | None, redeem_assertion: RedeemAssertion
+    ) -> Response:
+        """The RFC 7523 tail: prove the assertion and the requesting app, claim the assertion's
+        ``jti`` before anything is minted (RFC 7523 section 3 makes a replayed ``jti`` a
+        rejection, and the identity provider mints a fresh assertion cheaply, so a mint that
+        fails after the claim costs the client one more round trip, never a replay window),
+        then mint the proxy credential the user is revalidated for. No refresh token comes
+        with it (RFC 7521 section 4.1): the app renews by asking its identity provider for a
+        fresh assertion, which keeps the user's standing there in the loop."""
+        target_refusal: Final = self._assertion_target_refusal()
+        if target_refusal is not None:
+            return target_refusal
+        identity: Final = await redeem_assertion(assertion, client_id, client_secret, self._request)
+        if isinstance(identity, SubjectTokenRefusal):
+            return _subject_token_refusal_response(identity)
+        refused: Final = _claim_refusal(
+            await self._guard.claim(identity.replay_key, identity.replay_ttl_seconds),
+            _oauth_error(400, "invalid_grant", "the assertion has already been redeemed"),
+        )
+        if refused is not None:
+            return refused
+        minted: Final = await self._mint_proxy_credential(
+            identity.user_id, identity.team_id, allowed_routes=identity.allowed_routes
+        )
+        if not isinstance(minted, MintedProxyCredential):
+            return _mint_failure_response(minted)
+        return assertion_credential_response(minted, identity.scope)
+
+    def _assertion_target_refusal(self) -> Response | None:
+        """An ID-JAG names the proxy origin or its MCP resource (``{base}/mcp``) as ``aud``, and an RFC
+        8707 client sends that same value as ``resource``, so both are this grant's targets."""
+        if self._resource is None or is_proxy_api_resource(self._request, self._resource):
+            return None
+        if self._resource.rstrip("/") == f"{get_request_base_url(self._request)}/mcp":
+            return None
+        return _oauth_error(
+            400,
+            "invalid_target",
+            "resource does not match the proxy API or MCP resource this gateway issues credentials for",
+        )
+
     def _proxy_api_target_refusal(self) -> Response | None:
         if self._resource is None or is_proxy_api_resource(self._request, self._resource):
             return None
@@ -1807,6 +1928,19 @@ async def refresh_proxy_credential(
     )
 
 
+async def _jwt_bearer_grant(
+    request: Request,
+    assertion: str | None,
+    client_id: str,
+    client_secret: str | None,
+    redeem_assertion: RedeemAssertion,
+    issue: _GrantIssuer,
+) -> Response:
+    if not assertion:
+        return _oauth_error(400, "invalid_request", "assertion is required")
+    return await issue.redeem_assertion(assertion, client_id, client_secret, redeem_assertion)
+
+
 async def _token_exchange_grant(
     subject_token: str | None,
     subject_token_type: str | None,

litellm/proxy/_experimental/mcp_server/id_jag_grant.py

@@ -0,0 +1,540 @@
+"""The RFC 7523 jwt-bearer grant on ``POST /token``: a requesting app that obtained an
+ID-JAG (Identity Assertion Authorization Grant, the cross-app access profile) from the
+customer's identity provider redeems it for the same proxy-API credential ``lite login``
+stores. The assertion is proven with the JWT auth issuer and JWKS plumbing the proxy
+already trusts, its claims are checked the way the ID-JAG draft prescribes, the requesting
+app must be one the operator allowlisted, and the ``jti`` is single-use until the
+assertion expires."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import hashlib
+import hmac
+import math
+from collections.abc import Awaitable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from functools import partial
+from types import MappingProxyType
+from typing import Final, Protocol
+from urllib.parse import unquote_plus
+
+import jwt
+from fastapi import Request
+
+from litellm._logging import verbose_proxy_logger
+from litellm.models.user import LiteLLM_UserTable
+from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
+    AssertionIdentity,
+    ConsentTeam,
+    SubjectIdentity,
+    SubjectTokenRefusal,
+)
+from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
+    gateway_outage_refusal,
+    log_line,
+    read_token_exchange_prerequisites,
+)
+from litellm.proxy._experimental.mcp_server.oauth_utils import proxy_base_url_from_env
+from litellm.proxy._types import IdJagRequestingApp
+from litellm.proxy.auth.handle_jwt import ID_JAG_TYP, AssertionRejection, JWTHandler
+
+ID_JAG_REQUIRED_CLAIMS: Final = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id")
+REJECTED_ASSERTION: Final = "assertion was rejected"
+TEAM_HEADER: Final = "x-litellm-team-id"
+_INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable"))
+_REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:"
+SCOPE_MCP: Final = "litellm:mcp"
+SCOPE_LLM: Final = "litellm:llm"
+SCOPE_ROUTE_GROUPS: Final = MappingProxyType(
+    {SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)}
+)
+
+
+@dataclass(frozen=True, slots=True)
+class GrantedScope:
+    """What an assertion's ``scope`` claim buys: the scope echoed in the token response, exactly the
+    tokens enforced, and the route groups the credential is narrowed to. No claim means the full
+    credential, so both are ``None``."""
+
+    scope: str | None
+    allowed_routes: tuple[str, ...] | None
+
+
+def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal:
+    """The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and
+    no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity
+    provider named that resource, so the credential is ceilinged to the MCP routes whatever the
+    claim says: ``litellm:llm`` cannot be granted and is refused."""
+    if claim is not None and not isinstance(claim, str):
+        return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION)
+    requested: Final = frozenset(claim.split()) if claim else frozenset()
+    if not requested <= SCOPE_ROUTE_GROUPS.keys():
+        return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION)
+    if mcp_resource:
+        if SCOPE_LLM in requested:
+            return SubjectTokenRefusal(error="invalid_scope", description=_MCP_RESOURCE_DESCRIPTION)
+        return GrantedScope(scope=SCOPE_MCP, allowed_routes=SCOPE_ROUTE_GROUPS[SCOPE_MCP])
+    if not requested:
+        return GrantedScope(scope=None, allowed_routes=None)
+    granted: Final = tuple(scope for scope in SCOPE_ROUTE_GROUPS if scope in requested)
+    return GrantedScope(scope=" ".join(granted), allowed_routes=tuple(_route_groups(granted)))
+
+
+def names_mcp_resource(aud: object, base_url: str | None) -> bool:
+    """Whether the verified ``aud`` names only the MCP resource, ``{base}/mcp``, and not the
+    authorization server itself (the bare origin)."""
+    audiences: Final = frozenset((aud,) if isinstance(aud, str) else aud if isinstance(aud, list) else ())
+    return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences
+
+
+def _route_groups(scopes: Sequence[str]) -> Iterator[str]:
+    for scope in scopes:
+        yield from SCOPE_ROUTE_GROUPS[scope]
+
+
+_SCOPE_DESCRIPTION: Final = (
+    f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated"
+)
+_MCP_RESOURCE_DESCRIPTION: Final = (
+    f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class IdJagPrerequisites:
+    """The deployment-level gates for redeeming an ID-JAG: the JWT auth whose issuer plumbing
+    proves the assertion must be on, backed by a database, and licensed, not mapping IdP tokens to
+    virtual keys (whose blocks and budgets a minted credential would sidestep), with a shared Redis
+    so an assertion is redeemed once across every worker and replica (the in-memory marker is per
+    process, so a second worker would mint again), ``PROXY_BASE_URL``
+    must fix the gateway URL the assertion's ``aud`` has to name (a request's ``Host`` is the
+    caller's to choose), and the operator must have allowlisted the requesting apps. Discovery advertises the grant only when every gate
+    holds, and a redemption is refused naming the first one that does not."""
+
+    jwt_auth_enabled: bool
+    has_database: bool
+    licensed: bool
+    maps_jwts_to_virtual_keys: bool
+    shared_replay_store: bool
+    base_url: str | None
+    requesting_apps: tuple[IdJagRequestingApp, ...]
+
+    @property
+    def available(self) -> bool:
+        return (
+            self.jwt_auth_enabled
+            and self.has_database
+            and self.licensed
+            and not self.maps_jwts_to_virtual_keys
+            and self.shared_replay_store
+            and self.base_url is not None
+            and bool(self.requesting_apps)
+        )
+
+    def refusal(self) -> SubjectTokenRefusal | None:
+        if not self.jwt_auth_enabled:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="JWT auth is not enabled on this gateway, so it cannot redeem identity assertions",
+            )
+        if not self.has_database:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="this gateway has no database, so it cannot redeem identity assertions",
+            )
+        if not self.licensed:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="JWT auth is an enterprise only feature; no license is set",
+            )
+        if self.maps_jwts_to_virtual_keys:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="this gateway maps IdP tokens to virtual keys, which the jwt-bearer grant does not serve",
+            )
+        if not self.shared_replay_store:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="configure Redis (cache_params or coordination_redis) so this gateway can redeem an identity assertion exactly once across its workers",
+            )
+        if self.base_url is None:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="set PROXY_BASE_URL to this gateway's public URL, the audience an identity assertion must name",
+            )
+        if not self.requesting_apps:
+            return SubjectTokenRefusal(
+                error="unsupported_grant_type",
+                description="this gateway has no requesting apps allowlisted for identity assertions",
+            )
+        return None
+
+
+def read_id_jag_prerequisites() -> IdJagPrerequisites:
+    from litellm.proxy.proxy_server import (  # noqa: PLC0415  # rebound after startup, so read them per call
+        jwt_handler,
+        redis_usage_cache,
+    )
+
+    shared: Final = read_token_exchange_prerequisites()
+    return IdJagPrerequisites(
+        jwt_auth_enabled=shared.jwt_auth_enabled,
+        has_database=shared.has_database,
+        licensed=shared.licensed,
+        maps_jwts_to_virtual_keys=shared.maps_jwts_to_virtual_keys,
+        shared_replay_store=redis_usage_cache is not None,
+        base_url=proxy_base_url_from_env(),
+        requesting_apps=_requesting_apps(jwt_handler),
+    )
+
+
+def _requesting_apps(jwt_handler: JWTHandler) -> tuple[IdJagRequestingApp, ...]:
+    if not hasattr(jwt_handler, "litellm_jwtauth"):
+        return ()
+    return tuple(jwt_handler.litellm_jwtauth.id_jag_requesting_apps or ())
+
+
+def jwt_bearer_available() -> bool:
+    return read_id_jag_prerequisites().available
+
+
+def gateway_audiences(base_url: str | None) -> tuple[str, ...]:
+    """The issuer values this gateway answers to, which an ID-JAG's ``aud`` must name: its MCP
+    resource and its authorization server, both under the fixed ``PROXY_BASE_URL``."""
+    return () if base_url is None else (f"{base_url}/mcp", base_url)
+
+
+class DecodeAssertion(Protocol):
+    """Injected signature and claim verification ``(assertion, audience, required_claims)``:
+    ``JWTHandler.decode_trusted_assertion`` in production, which returns the verified claims,
+    a rejection, or raises when the gateway could not reach the issuer's keys."""
+
+    def __call__(
+        self, assertion: str, audience: Sequence[str], required_claims: Sequence[str], /
+    ) -> Awaitable[Mapping[str, object] | AssertionRejection]: ...
+
+
+class ClaimReader(Protocol):
+    """The JWT auth user mapping as ``JWTHandler`` exposes it: the litellm user id and email the
+    verified claims name through ``user_id_jwt_field`` and ``user_email_jwt_field``."""
+
+    def get_user_id(self, token: dict[str, object], default_value: str | None) -> str | None: ...
+
+    def get_user_email(self, token: dict[str, object], default_value: str | None) -> str | None: ...
+
+    def get_all_jwt_team_ids(self, token: dict[str, object]) -> list[str]: ...
+
+    def is_allowed_domain(self, user_email: str) -> bool: ...
+
+
+class ResolveSubject(Protocol):
+    """Injected subject lookup ``(claims, team_header)``: the litellm user the verified claims
+    name through the JWT auth user mapping (``user_id_jwt_field`` or ``user_email_jwt_field``,
+    else ``sub``) and the team the credential is minted for (see ``select_team``), ``None`` when
+    the gateway knows no such user, or a refusal when the team cannot be settled. Raises on a
+    database outage."""
+
+    def __call__(
+        self, claims: Mapping[str, object], team_header: str | None, /
+    ) -> Awaitable[SubjectIdentity | SubjectTokenRefusal | None]: ...
+
+
+class LookupUser(Protocol):
+    """Injected ``(user_id, user_email)`` lookup: the stored id of the existing litellm user the
+    claims name, ``None`` when the gateway knows no such user. Raises on a database outage."""
+
+    def __call__(self, user_id: str, user_email: str | None, /) -> Awaitable[str | None]: ...
+
+
+class LookupTeams(Protocol):
+    """Injected lookup of the teams a user is on, or the failure text when they cannot be read."""
+
+    def __call__(self, user_id: str, /) -> Awaitable[Sequence[ConsentTeam] | str]: ...
+
+
+class GetUser(Protocol):
+    def __call__(
+        self, *, user_id: str, user_email: str | None, sso_user_id: str
+    ) -> Awaitable[LiteLLM_UserTable | None]: ...
+
+
+def proxy_user_lookup(get_user: GetUser) -> LookupUser:
+    """``LookupUser`` over the proxy's user table. The assertion's subject is matched as the
+    local ``user_id`` and as the ``sso_user_id`` an SSO login linked, the way the JWT and SSO
+    login paths match it, so a linked account needs no email claim."""
+
+    async def lookup_user(user_id: str, user_email: str | None) -> str | None:
+        user: Final = await get_user(user_id=user_id, user_email=user_email, sso_user_id=user_id)
+        return None if user is None else user.user_id
+
+    return lookup_user
+
+
+def subject_resolver(reader: ClaimReader, lookup_user: LookupUser, lookup_teams: LookupTeams) -> ResolveSubject:
+    """``ResolveSubject`` over the JWT auth claim mapping and the proxy's user and team lookups."""
+
+    async def resolve_subject(
+        claims: Mapping[str, object], team_header: str | None
+    ) -> SubjectIdentity | SubjectTokenRefusal | None:
+        token: Final = dict(claims)
+        subject: Final = claims.get("sub")
+        claimed_user_id: Final = reader.get_user_id(token=token, default_value=None) or (
+            subject if isinstance(subject, str) else None
+        )
+        if claimed_user_id is None:
+            return None
+        user_email: Final = reader.get_user_email(token=token, default_value=None)
+        if not reader.is_allowed_domain(user_email or ""):
+            return SubjectTokenRefusal(
+                error="invalid_grant",
+                description="the assertion's subject email is missing or outside the allowed domain",
+            )
+        user_id: Final = await lookup_user(claimed_user_id, user_email)
+        if user_id is None:
+            return None
+        teams: Final = await lookup_teams(user_id)
+        if teams in _INACTIVE_SUBJECT:
+            return SubjectTokenRefusal(
+                error="invalid_grant", description="the assertion's subject is not an active user"
+            )
+        if isinstance(teams, str):
+            return SubjectTokenRefusal(
+                error="temporarily_unavailable", description="the gateway could not read this user's teams; retry"
+            )
+        team: Final = select_team(team_header, tuple(reader.get_all_jwt_team_ids(token=token)), teams)
+        if isinstance(team, SubjectTokenRefusal):
+            return team
+        return SubjectIdentity(user_id=user_id, team_id=team)
+
+    return resolve_subject
+
+
+def select_team(
+    team_header: str | None, claim_teams: Sequence[str], teams: Sequence[ConsentTeam]
+) -> str | None | SubjectTokenRefusal:
+    """The team an ID-JAG redemption mints for, settled before the assertion is consumed since
+    the minter refuses a teamless credential for a user who has a team. The ``x-litellm-team-id``
+    header (a team id or alias) wins and must name one of the user's teams, then the assertion's
+    mapped team claim, then the user's only team. A user on several teams must name one."""
+    if team_header is not None:
+        named: Final = next((team for team in teams if team_header in (team.team_id, team.team_alias)), None)
+        if named is None:
+            return SubjectTokenRefusal(
+                error="invalid_request", description=f"{TEAM_HEADER} does not name one of this user's teams"
+            )
+        return named.team_id
+    asserted: Final = tuple(team.team_id for team in teams if team.team_id in claim_teams)
+    if len(asserted) == 1:
+        return asserted[0]
+    if len(asserted) > 1:
+        return SubjectTokenRefusal(
+            error="invalid_request",
+            description=f"the assertion names several of this user's teams; send the {TEAM_HEADER} header naming the team for this credential",
+        )
+    if len(teams) == 1:
+        return teams[0].team_id
+    if teams:
+        return SubjectTokenRefusal(
+            error="invalid_request",
+            description=f"this user belongs to several teams; send the {TEAM_HEADER} header naming the team for this credential",
+        )
+    return None
+
+
+@dataclass(frozen=True, slots=True)
+class AssertionRedemption:
+    """One ``jwt-bearer`` request as the grant sees it: the assertion, how the requesting app
+    identified and authenticated itself, the team the credential is asked for, the issuer
+    values this gateway answers to (what the assertion's ``aud`` must name), the clock, and the
+    leeway the verifier grants ``exp``, so the replay marker outlives every moment the assertion
+    would still verify."""
+
+    assertion: str
+    client_id: str
+    client_secret: str | None
+    authorization_header: str | None
+    team_id: str | None
+    audiences: tuple[str, ...]
+    now: float
+    leeway_seconds: float
+
+
+async def redeem_id_jag_assertion(
+    assertion: str, client_id: str, client_secret: str | None, request: Request
+) -> AssertionIdentity | SubjectTokenRefusal:
+    from litellm.proxy._experimental.mcp_server.proxy_api_credentials import (  # noqa: PLC0415  # circular import at module load
+        lookup_consent_teams_from_database,
+    )
+    from litellm.proxy.auth.auth_checks import get_user_object  # noqa: PLC0415  # circular import at module load
+    from litellm.proxy.proxy_server import (  # noqa: PLC0415  # rebound after startup, so read them per call
+        jwt_handler,
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
+
+    resolve_subject: Final = subject_resolver(
+        _claim_reader(jwt_handler),
+        proxy_user_lookup(
+            partial(
+                get_user_object,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                user_id_upsert=False,
+                proxy_logging_obj=proxy_logging_obj,
+                check_db_only=True,
+            )
+        ),
+        lookup_consent_teams_from_database,
+    )
+    prerequisites: Final = read_id_jag_prerequisites()
+    return await identity_from_id_jag(
+        AssertionRedemption(
+            assertion=assertion,
+            client_id=client_id,
+            client_secret=client_secret,
+            authorization_header=request.headers.get("authorization"),
+            team_id=request.headers.get(TEAM_HEADER),
+            audiences=gateway_audiences(prerequisites.base_url),
+            now=datetime.now(timezone.utc).timestamp(),
+            leeway_seconds=jwt_handler.leeway,
+        ),
+        prerequisites=prerequisites,
+        decode=jwt_handler.decode_trusted_assertion,
+        resolve_subject=resolve_subject,
+    )
+
+
+def _claim_reader(handler: ClaimReader) -> ClaimReader:
+    return handler
+
+
+async def identity_from_id_jag(
+    redemption: AssertionRedemption,
+    prerequisites: IdJagPrerequisites,
+    decode: DecodeAssertion,
+    resolve_subject: ResolveSubject,
+) -> AssertionIdentity | SubjectTokenRefusal:
+    """Prove an ID-JAG for the grant tail. The requesting app is checked first, the way a token
+    endpoint authenticates its client before it looks at the grant, so an app the operator never
+    allowlisted learns nothing about the assertion. Everything wrong with the assertion itself is
+    ``invalid_grant`` (RFC 7523 section 3.1), a check the gateway could not complete is
+    ``temporarily_unavailable``, and the reason stays in the proxy log because this endpoint is
+    public and the verifier's wording can quote the assertion."""
+    unmet: Final = prerequisites.refusal()
+    if unmet is not None:
+        return unmet
+    app: Final = next((app for app in prerequisites.requesting_apps if app.client_id == redemption.client_id), None)
+    if app is None:
+        return SubjectTokenRefusal(error="invalid_client", description="client_id is not an allowlisted requesting app")
+    if not _client_authenticated(app, redemption.client_secret, redemption.authorization_header):
+        return SubjectTokenRefusal(error="invalid_client", description="client authentication failed")
+    try:
+        header: Final = jwt.get_unverified_header(redemption.assertion)
+    except jwt.PyJWTError:
+        return SubjectTokenRefusal(error="invalid_grant", description="assertion is not a JWT")
+    if header.get("typ") != ID_JAG_TYP:
+        return SubjectTokenRefusal(error="invalid_grant", description=f"assertion typ must be {ID_JAG_TYP}")
+    try:
+        verified: Final = await decode(redemption.assertion, redemption.audiences, ID_JAG_REQUIRED_CLAIMS)
+    except Exception as denied:  # noqa: BLE001  # the verifier raises its 503 for a JWKS outage, and anything else is a bad assertion
+        return gateway_outage_refusal(denied, denied, what="an assertion", parameter="assertion") or _rejected(
+            str(denied)
+        )
+    if isinstance(verified, AssertionRejection):
+        return _rejected(verified.reason)
+    marker: Final = _replay_marker(verified, redemption)
+    if isinstance(marker, SubjectTokenRefusal):
+        return marker
+    replay_key, replay_ttl_seconds = marker
+    scope: Final = granted_scope(
+        verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url)
+    )
+    if isinstance(scope, SubjectTokenRefusal):
+        return scope
+    try:
+        subject: Final = await resolve_subject(verified, redemption.team_id)
+    except Exception as failed:  # noqa: BLE001  # the user lookup raises for a missing user and for a database outage alike
+        return gateway_outage_refusal(
+            failed, failed, what="an assertion's subject", parameter="assertion"
+        ) or _rejected(f"subject lookup failed: {failed}")
+    if subject is None:
+        return _rejected("the assertion names no user the gateway knows")
+    if isinstance(subject, SubjectTokenRefusal):
+        return subject
+    return AssertionIdentity(
+        user_id=subject.user_id,
+        team_id=subject.team_id,
+        replay_key=replay_key,
+        replay_ttl_seconds=replay_ttl_seconds,
+        scope=scope.scope,
+        allowed_routes=scope.allowed_routes,
+    )
+
+
+def _replay_marker(
+    verified: Mapping[str, object], redemption: AssertionRedemption
+) -> tuple[str, int] | SubjectTokenRefusal:
+    """The single-use key and its lifetime for a verified assertion, once the assertion is bound to
+    the presenting app, carries no proof-of-possession the gateway cannot check, and is still alive."""
+    if verified.get("client_id") != redemption.client_id:
+        return _rejected("the assertion's client_id is not the requesting app that presented it")
+    if "cnf" in verified:
+        return _rejected("the assertion is proof-bound and this gateway does not check proof of possession")
+    issuer: Final = verified.get("iss")
+    jti: Final = verified.get("jti")
+    exp: Final = verified.get("exp")
+    if not isinstance(issuer, str) or not isinstance(jti, str) or not jti or not isinstance(exp, (int, float)):
+        return _rejected("the assertion's iss, jti, or exp claim has the wrong shape")
+    replay_ttl_seconds: Final = math.ceil(exp + redemption.leeway_seconds - redemption.now)
+    if replay_ttl_seconds <= 0:
+        return _rejected("the assertion has expired")
+    return _REPLAY_KEY_PREFIX + hashlib.sha256(f"{issuer}\n{jti}".encode()).hexdigest(), replay_ttl_seconds
+
+
+def _rejected(reason: str) -> SubjectTokenRefusal:
+    verbose_proxy_logger.warning("jwt-bearer grant refused an assertion: %s", log_line(reason))
+    return SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION)
+
+
+def _client_authenticated(app: IdJagRequestingApp, form_secret: str | None, authorization_header: str | None) -> bool:
+    """A public app (no secret configured) is identified by ``client_id`` alone. A confidential
+    app must present its secret, by HTTP Basic (RFC 6749 section 2.3.1, the credentials
+    form-urlencoded inside the base64) or the ``client_secret`` form field, compared in constant
+    time."""
+    if app.client_secret is None:
+        return True
+    basic: Final = _basic_credentials(authorization_header)
+    if basic is not None:
+        return basic[0] == app.client_id and _same_secret(basic[1], app.client_secret)
+    return form_secret is not None and _same_secret(form_secret, app.client_secret)
+
+
+def _same_secret(presented: str, expected: str) -> bool:
+    return hmac.compare_digest(presented.encode(), expected.encode())
+
+
+def basic_client_id(authorization_header: str | None) -> str | None:
+    """The ``client_id`` an HTTP Basic ``Authorization`` header carries, for a confidential app
+    that leaves the form field out as RFC 6749 section 2.3.1 allows."""
+    credentials: Final = _basic_credentials(authorization_header)
+    return None if credentials is None else credentials[0]
+
+
+def _basic_credentials(authorization_header: str | None) -> tuple[str, str] | None:
+    if authorization_header is None:
+        return None
+    scheme, _, encoded = authorization_header.strip().partition(" ")
+    if scheme.lower() != "basic" or not encoded:
+        return None
+    try:
+        decoded: Final = base64.b64decode(encoded.strip(), validate=True).decode()
+    except (binascii.Error, UnicodeDecodeError):
+        return None
+    client_id, separator, secret = decoded.partition(":")
+    if not separator:
+        return None
+    return unquote_plus(client_id), unquote_plus(secret)

litellm/proxy/_experimental/mcp_server/idp_token_exchange.py

@@ -20,13 +20,15 @@ from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 
 EXCHANGE_ROUTE: Final = "/token"
 REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
-SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
-    "the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
+_CHECK_UNAVAILABLE: Final = (
+    "the gateway could not verify {parameter} because its identity provider or database is unavailable; retry"
 )
-SUBJECT_TOKEN_CHECK_FAULTED: Final = (
-    "the gateway could not verify subject_token because its database reported a fault that is not a transient "
+_CHECK_FAULTED: Final = (
+    "the gateway could not verify {parameter} because its database reported a fault that is not a transient "
     "outage; retrying will not help until the gateway deployment is repaired"
 )
+SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = _CHECK_UNAVAILABLE.format(parameter="subject_token")
+SUBJECT_TOKEN_CHECK_FAULTED: Final = _CHECK_FAULTED.format(parameter="subject_token")
 GatewayOutage = Literal["retryable", "faulted"]
 
 
@@ -182,20 +184,37 @@ async def identity_from_subject_token(
 
 
 def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
+    outage: Final = gateway_outage_refusal(denied, reason, what="a subject_token", parameter="subject_token")
+    if outage is not None:
+        return outage
+    verbose_proxy_logger.warning("token exchange refused a subject_token: %s", log_line(reason))
+    return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
+
+
+def gateway_outage_refusal(denied: Exception, reason: object, what: str, parameter: str) -> SubjectTokenRefusal | None:
+    """The ``temporarily_unavailable`` refusal when ``denied`` is the gateway failing to check
+    ``what`` (its database or the IdP's JWKS down), naming the grant ``parameter`` it could not verify,
+    else ``None`` because the token itself was bad."""
     outage: Final = _gateway_could_not_verify(denied)
     if outage is None:
-        verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
-        return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
-    verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
-    return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
+        return None
+    verbose_proxy_logger.error("token endpoint could not verify %s, %s: %s", what, outage, log_line(reason))
+    return SubjectTokenRefusal(
+        error="temporarily_unavailable", description=_check_unavailable_description(outage, parameter)
+    )
 
 
-def _check_unavailable_description(outage: GatewayOutage) -> str:
+def log_line(reason: object) -> str:
+    """``reason`` as one log line, so text a caller controls cannot forge log entries."""
+    return str(reason).replace("\r", " ").replace("\n", " ")
+
+
+def _check_unavailable_description(outage: GatewayOutage, parameter: str) -> str:
     match outage:
         case "retryable":
-            return SUBJECT_TOKEN_CHECK_UNAVAILABLE
+            return _CHECK_UNAVAILABLE.format(parameter=parameter)
         case "faulted":
-            return SUBJECT_TOKEN_CHECK_FAULTED
+            return _CHECK_FAULTED.format(parameter=parameter)
         case _:
             assert_never(outage)
 

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -455,6 +455,7 @@ class MCPServerConfig(TypedDict, total=False):
     upstream_token_header: ReadOnly[str]
     id_jag_resource_token_endpoint: str
     id_jag_resource: str
+    id_jag_subject_source: ReadOnly[str]
     id_jag_resource_client_id: ReadOnly[str]
     id_jag_resource_client_secret: ReadOnly[str]
     client_private_key: str
@@ -2789,6 +2790,7 @@ class MCPServerManager:
                 # ID-JAG fields
                 id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None),
                 id_jag_resource=server_config.get("id_jag_resource", None),
+                id_jag_subject_source=server_config.get("id_jag_subject_source", None),
                 id_jag_resource_client_id=server_config.get("id_jag_resource_client_id", None),
                 id_jag_resource_client_secret=server_config.get("id_jag_resource_client_secret", None),
                 client_private_key=server_config.get("client_private_key", None),
@@ -3380,6 +3382,7 @@ class MCPServerManager:
                 credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None
             ),
             id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None),
+            id_jag_subject_source=(credentials_dict.get("id_jag_subject_source") if credentials_dict else None),
             id_jag_resource_client_id=(credentials_dict.get("id_jag_resource_client_id") if credentials_dict else None),
             id_jag_resource_client_secret=self._decrypt_credential_field(
                 credentials_dict.get("id_jag_resource_client_secret") if credentials_dict else None,

litellm/proxy/_experimental/mcp_server/oauth_utils.py

@@ -19,6 +19,9 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
     build_token_endpoint_client_auth,
     normalize_token_endpoint_auth_method,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
+    GOOGLE_REFRESH_TOKEN_PARAMS,
+)
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path
 
@@ -59,7 +62,7 @@ _DEFAULT_NATIVE_REDIRECT_URIS: Final[list[str]] = [
 ]
 
 _GOOGLE_AUTHORIZATION_HOSTS: Final = frozenset({"accounts.google.com"})
-_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"})
+_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = GOOGLE_REFRESH_TOKEN_PARAMS
 
 _warned_invalid_proxy_base_url: str | None = None
 
@@ -128,6 +131,11 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st
     return urlunparse(parsed._replace(query=urlencode(merged)))
 
 
+def proxy_base_url_from_env() -> str | None:
+    """The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed."""
+    return _resolve_proxy_base_url_env()
+
+
 def _resolve_proxy_base_url_env() -> str | None:
     global _warned_invalid_proxy_base_url
     configured: Final = os.environ.get("PROXY_BASE_URL", "").strip()

litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py

@@ -39,6 +39,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     HeaderCarrier,
     IdJagConfig,
     IdJagResourceClient,
+    IdJagSubjectSource,
     NoneConfig,
     PassthroughConfig,
     PrivateKeyJwtAuth,
@@ -71,6 +72,7 @@ __all__ = [
     "HeaderCarrier",
     "IdJagConfig",
     "IdJagResourceClient",
+    "IdJagSubjectSource",
     "NoOpAuth",
     "NoneConfig",
     "Ok",

litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py

@@ -34,6 +34,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     CredError,
     IdJagConfig,
     IdJagResourceClient,
+    IdJagSubjectSource,
     NoneConfig,
     PassthroughConfig,
     PrivateKeyJwtAuth,
@@ -236,7 +237,8 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None:
     if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None:
         return None
     resource_client: Final = _id_jag_resource_client(server, client_id, client_auth)
-    if resource_client is None:
+    subject_source: Final = _id_jag_subject_source(server)
+    if resource_client is None or subject_source is None:
         return None
     return ServerSpec(
         server_id=server.server_id,
@@ -249,6 +251,7 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None:
             client_auth=client_auth,
             resource_client=resource_client,
             subject_token_type=_id_jag_subject_token_type(server),
+            subject_source=subject_source,
             audience=server.audience,
             resource=server.id_jag_resource,
             scopes=tuple(server.scopes or ()),
@@ -282,6 +285,18 @@ def _id_jag_client_auth(server: MCPServer) -> ClientAuth | None:
     return None
 
 
+def _id_jag_subject_source(server: MCPServer) -> IdJagSubjectSource | None:
+    """``None`` for a value that names no source, so the half-configured server defers like a missing
+    endpoint does instead of silently presenting the wrong token."""
+    configured: Final = server.id_jag_subject_source
+    if not configured:
+        return IdJagSubjectSource.login_assertion
+    try:
+        return IdJagSubjectSource(configured)
+    except ValueError:
+        return None
+
+
 def _id_jag_subject_token_type(server: MCPServer) -> str:
     """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token;
     an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim."""

litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py

@@ -0,0 +1,216 @@
+"""Exchange the identity assertion captured at SSO login for a refresh token, at login.
+
+An ``oauth2_id_jag`` server configured with ``id_jag_subject_source: refresh_token`` does not present the
+login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued for that assertion through
+one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at
+an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at
+login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and
+the refresh token outlives the assertion by as long as the IdP decides.
+
+The servers consulted are the two configuration authorities, the pod-local config declaration and the DB
+rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a
+server added on another pod. The exchange authenticates as the ID-JAG server's own IdP client, since the
+refresh token is bound to the client that later spends it. One row holds one refresh token per user, so every server in this mode has
+to share one IdP client; a second client is logged and skipped rather than silently handed a token it
+cannot redeem. A refused or unreachable IdP never fails the login, the row keeps the plain assertion and
+the ID-JAG call reports the missing refresh token as the precondition it is.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Iterable, Mapping, Sequence
+from datetime import datetime, timezone
+from typing import Final, Protocol
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
+from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
+    SSOIdentityAssertion,
+    with_login_refresh_token,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken
+from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
+    ClientAuth,
+    CredError,
+    IdJagConfig,
+    IdJagSubjectSource,
+    ServerSpec,
+)
+from litellm.proxy._types import LiteLLM_MCPServerTable
+from litellm.types.mcp import MCPAuth
+from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
+REFRESH_TOKEN_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token"
+_LOGIN_EXCHANGE_DEADLINE_SECONDS: Final = 15.0
+
+_REFRESH_TOKEN_SCOPE: Final = "openid offline_access"
+_RUNTIME_APPROVAL_STATES: Final = (None, "active", "approved")  # the states get_runtime_mcp_server_rows serves
+
+
+class LoginExchangeCatalog(Protocol):
+    """The slice of the MCP server manager the login-time exchange reads."""
+
+    @property
+    def config_mcp_servers(self) -> Mapping[str, MCPServer]: ...
+
+    async def build_mcp_server_from_table(
+        self,
+        mcp_server: LiteLLM_MCPServerTable,
+        *,
+        credentials_are_encrypted: bool = True,
+        env_vars_are_encrypted: bool | None = None,
+        register_oauth_discovery: bool = True,
+    ) -> MCPServer: ...
+
+
+class TokenExchangeEndpoint(Protocol):
+    async def fetch(
+        self, endpoint: str, client_id: str, grant_params: Mapping[str, str], client_auth: ClientAuth
+    ) -> Result[ExchangedToken, CredError]: ...
+
+
+def _is_runtime_id_jag_row(row: LiteLLM_MCPServerTable) -> bool:
+    """An ID-JAG row the runtime registry would serve. A pending or rejected submission never gets a
+    login assertion posted to the token endpoint it names."""
+    return row.auth_type == MCPAuth.oauth2_id_jag and row.approval_status in _RUNTIME_APPROVAL_STATES
+
+
+async def login_exchange_servers(
+    catalog: LoginExchangeCatalog, db_rows: Sequence[LiteLLM_MCPServerTable]
+) -> tuple[MCPServer, ...]:
+    """Every server the exchange has to consider: the config-declared ones plus the approved ID-JAG rows
+    in the DB, built the same way the catalog builds them so their credentials are decrypted. A DB row
+    hides a config server with the same ``server_id``, exactly as the live catalog serves them."""
+    id_jag_rows: Final = tuple(row for row in db_rows if _is_runtime_id_jag_row(row))
+    db_servers: Final = tuple([await _built_or_skipped(catalog, row) for row in id_jag_rows])
+    db_by_id: Final = {server.server_id: server for server in db_servers if server is not None}
+    return tuple({**catalog.config_mcp_servers, **db_by_id}.values())
+
+
+async def _built_or_skipped(catalog: LoginExchangeCatalog, row: LiteLLM_MCPServerTable) -> MCPServer | None:
+    """One row that cannot be built (undecryptable credentials, a half-configured submission) is skipped
+    with a warning. It must not cost every other refresh-token server its login exchange."""
+    try:
+        return await catalog.build_mcp_server_from_table(row, register_oauth_discovery=False)
+    except Exception as exc:  # noqa: BLE001  # a per-row build failure is logged and isolated, not propagated
+        verbose_proxy_logger.warning(
+            "ID-JAG: MCP server row server_id=%s could not be built for the login-time refresh token exchange "
+            "and is skipped: %s",
+            row.server_id,
+            exc,
+        )
+        return None
+
+
+def refresh_token_exchange_targets(specs: Iterable[ServerSpec | None]) -> tuple[IdJagConfig, ...]:
+    """The ID-JAG configs that present a login refresh token, one per IdP client, in registration order."""
+    configs: Final = tuple(
+        spec.config
+        for spec in specs
+        if spec is not None
+        and isinstance(spec.config, IdJagConfig)
+        and spec.config.subject_source is IdJagSubjectSource.refresh_token
+    )
+    clients: Final = tuple((config.org_token_endpoint, config.client_id) for config in configs)
+    return tuple(config for index, config in enumerate(configs) if clients.index(clients[index]) == index)
+
+
+async def exchange_login_assertion_for_refresh_token(
+    assertion: SSOIdentityAssertion,
+    targets: tuple[IdJagConfig, ...],
+    token_endpoint: TokenExchangeEndpoint,
+    now: datetime,
+    deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS,
+) -> SSOIdentityAssertion:
+    """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the
+    refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token
+    does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it."""
+    if not targets or assertion.exchanged_refresh_token is not None:
+        return assertion
+    config: Final = targets[0]
+    if len(targets) > 1:
+        verbose_proxy_logger.warning(
+            "ID-JAG: %d IdP clients are configured with id_jag_subject_source refresh_token; the login refresh "
+            "token is minted for client_id=%s at %s and servers using another client cannot redeem it.",
+            len(targets),
+            config.client_id,
+            config.org_token_endpoint,
+        )
+    params: Final = {
+        "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
+        "requested_token_type": REFRESH_TOKEN_TOKEN_TYPE,
+        "subject_token": assertion.id_token.get_secret_value(),
+        "subject_token_type": assertion.subject_token_type or config.subject_token_type,
+        "scope": _REFRESH_TOKEN_SCOPE,
+    }
+    try:
+        outcome: Final = await asyncio.wait_for(
+            token_endpoint.fetch(config.org_token_endpoint, config.client_id, params, config.client_auth),
+            timeout=deadline_seconds,
+        )
+    except (TimeoutError, asyncio.TimeoutError):
+        verbose_proxy_logger.warning(
+            "ID-JAG: the IdP at %s did not answer the login-time refresh token exchange within %.0fs, so the login "
+            "finishes with the captured assertion only. Servers with id_jag_subject_source: refresh_token will ask "
+            "this user to sign in again.",
+            config.org_token_endpoint,
+            deadline_seconds,
+        )
+        return assertion
+    match outcome:
+        case Error(err):
+            verbose_proxy_logger.warning(
+                "ID-JAG: the IdP at %s refused to exchange the login assertion for a refresh token (%s); the "
+                "login goes on with the assertion alone and servers presenting a refresh token will ask the "
+                "user to sign in again.",
+                config.org_token_endpoint,
+                err.summary,
+            )
+            return assertion
+        case Ok(token):
+            verbose_proxy_logger.info(
+                "ID-JAG: exchanged the login assertion for a refresh token at %s (expires_in=%s)",
+                config.org_token_endpoint,
+                token.expires_in,
+            )
+            return with_login_refresh_token(
+                assertion,
+                token.access_token,
+                token.expires_in,
+                now,
+                endpoint=config.org_token_endpoint,
+                client_id=config.client_id,
+            )
+
+
+async def exchange_login_assertion_if_configured(
+    assertion: SSOIdentityAssertion | None,
+) -> SSOIdentityAssertion | None:
+    """The SSO-callback hook: consults the configured and DB-declared MCP servers and performs the
+    exchange when one of them presents a login refresh token. Nothing here can fail the login."""
+    if assertion is None:
+        return None
+    from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers  # noqa: PLC0415  # import cycle
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (  # noqa: PLC0415  # avoids import cycle
+        global_mcp_server_manager,
+    )
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (  # noqa: PLC0415  # import cycle
+        to_server_spec,
+    )
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (  # noqa: PLC0415  # cycle
+        TokenEndpointClient,
+    )
+    from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # runtime global
+
+    try:
+        db_rows: Final = () if prisma_client is None else await get_all_mcp_servers(prisma_client)
+        servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows)
+        targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers)
+        return await exchange_login_assertion_for_refresh_token(
+            assertion, targets, TokenEndpointClient(), datetime.now(timezone.utc)
+        )
+    except Exception as exc:  # noqa: BLE001  # the login itself must not fail on an egress-side exchange
+        verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc)
+        return assertion

litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py

@@ -20,6 +20,7 @@ v2: no imports from v1.
 from __future__ import annotations
 
 import hashlib
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
 from typing import Final
@@ -53,7 +54,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r
 from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
     AssertionStoreUnavailable,
     SSOAssertionStore,
+    SSOIdentityAssertion,
     assertion_expired,
+    exchanged_refresh_token_issued_for,
+    login_refresh_token_expired,
 )
 from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import (
     ExchangedToken,
@@ -75,6 +79,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     ClientSecretAuth,
     CredError,
     IdJagConfig,
+    IdJagSubjectSource,
     NoneConfig,
     PassthroughConfig,
     PrivateKeyJwtAuth,
@@ -88,6 +93,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
 _TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
 _JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer"
 _ID_JAG_REQUESTED_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:id-jag"
+_REFRESH_TOKEN_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token"
+
+
+@dataclass(frozen=True, slots=True)
+class _IdJagSubjectToken:
+    """What ID-JAG leg 1 asserts and the RFC 8693 type it is declared as."""
+
+    value: str
+    token_type: str
 
 
 class _NullOAuthTokenStore:
@@ -200,13 +214,15 @@ class UpstreamCredentialProvider:
     async def _id_jag(
         self, subject: Subject, server: ServerSpec, config: IdJagConfig
     ) -> Result[httpx2.Auth, CredError]:
-        match await self._id_jag_subject_token(subject):
+        match await self._id_jag_subject_token(subject, config):
             case Error(err):
                 return Error(err)
             case Ok(subject_token):
                 return await self._id_jag_exchange(subject, subject_token, server, config)
 
-    async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]:
+    async def _id_jag_subject_token(
+        self, subject: Subject, config: IdJagConfig
+    ) -> Result[_IdJagSubjectToken, CredError]:
         """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at.
 
         A caller that presents its own IdP identity token wins: that is the strongest available
@@ -223,7 +239,7 @@ class UpstreamCredentialProvider:
         a raised error so a DB blip cannot 500 the egress or the upstream-401 retry.
         """
         if subject.inbound_token is not None:
-            return Ok(subject.inbound_token.get_secret_value())
+            return Ok(_IdJagSubjectToken(subject.inbound_token.get_secret_value(), config.subject_token_type))
         if not subject.subject_id:
             return Error(
                 CredError.of_precondition_required(
@@ -232,7 +248,11 @@ class UpstreamCredentialProvider:
                 )
             )
         try:
-            assertion: Final = await self._sso_assertion_store.fetch(subject.subject_id)
+            assertion: Final = await (
+                self._sso_assertion_store.fetch_stored(subject.subject_id)
+                if config.subject_source is IdJagSubjectSource.refresh_token
+                else self._sso_assertion_store.fetch(subject.subject_id)
+            )
         except AssertionStoreUnavailable as exc:
             # The driver's message can name hosts, schemas or connection details, and this summary
             # is returned to the caller verbatim as a 503 body. Operators get it from the log.
@@ -253,6 +273,8 @@ class UpstreamCredentialProvider:
                     "Sign in through LiteLLM SSO so the gateway captures one."
                 )
             )
+        if config.subject_source is IdJagSubjectSource.refresh_token:
+            return _login_refresh_token_subject(assertion, config, datetime.now(timezone.utc))
         if assertion_expired(assertion, datetime.now(timezone.utc)):
             return Error(
                 CredError.of_precondition_required(
@@ -260,20 +282,24 @@ class UpstreamCredentialProvider:
                     "LiteLLM SSO again to capture a current one."
                 )
             )
-        return Ok(assertion.id_token.get_secret_value())
+        return Ok(
+            _IdJagSubjectToken(
+                assertion.id_token.get_secret_value(), assertion.subject_token_type or config.subject_token_type
+            )
+        )
 
     async def _id_jag_exchange(
-        self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig
+        self, subject: Subject, token: _IdJagSubjectToken, server: ServerSpec, config: IdJagConfig
     ) -> Result[httpx2.Auth, CredError]:
         slot: Final = _id_jag_slot_key(subject, server)
-        fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config)
+        fingerprint: Final = _id_jag_fingerprint(token.value, server.server_id, config)
 
         async def _exchange() -> Result[ExchangedToken, CredError]:
             leg1_params: Final = {
                 "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
                 "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE,
-                "subject_token": token,
-                "subject_token_type": config.subject_token_type,
+                "subject_token": token.value,
+                "subject_token_type": token.token_type,
                 **({"audience": config.audience} if config.audience else {}),
                 **({"resource": config.resource} if config.resource else {}),
                 **({"scope": " ".join(config.scopes)} if config.scopes else {}),
@@ -399,6 +425,41 @@ class UpstreamCredentialProvider:
             return None
 
 
+def _login_refresh_token_subject(
+    assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime
+) -> Result[_IdJagSubjectToken, CredError]:
+    """The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh
+    token the ID-JAG client exchanged the login assertion for when the user signed in. Never the SSO
+    refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry
+    does not apply here, a SAML assertion is spent minutes after login while the refresh token it
+    bought lives as long as the IdP said."""
+    if assertion.exchanged_refresh_token is None:
+        return Error(
+            CredError.of_precondition_required(
+                "This server mints its ID-JAG from a refresh token exchanged at SSO login and none is stored "
+                "for this user. Sign in through LiteLLM SSO again; if the login log shows the IdP refused the "
+                "exchange, the IdP client has to allow a refresh token for this requesting app."
+            )
+        )
+    if not exchanged_refresh_token_issued_for(assertion, config.org_token_endpoint, config.client_id):
+        return Error(
+            CredError.of_precondition_required(
+                "The refresh token exchanged at SSO login was minted for a different IdP client than this server "
+                "uses, so it cannot be presented here. Point every id_jag_subject_source: refresh_token server at "
+                "the same IdP client, then sign in through LiteLLM SSO again."
+            )
+        )
+    if login_refresh_token_expired(assertion, now):
+        return Error(
+            CredError.of_precondition_required(
+                "The refresh token exchanged at SSO login for this user has expired. Sign in through LiteLLM SSO again."
+            )
+        )
+    return Ok(
+        _IdJagSubjectToken(assertion.exchanged_refresh_token.get_secret_value(), _REFRESH_TOKEN_SUBJECT_TOKEN_TYPE)
+    )
+
+
 def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str:
     """Which cache slot this caller's bearer for this upstream lives in.
 

litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py

@@ -12,15 +12,18 @@ not the size of the user table) and single-flighted per user through the same ``
 the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent
 redemptions of one token as replay and can revoke the whole grant chain.
 
-The refresh is redeemed against the generic-OIDC client the login itself used
-(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy
-reconciles from the stored SSO row into the process environment at startup), authenticated the way
-that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the
-credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept
-only one of the two. An assertion can only exist if that client minted it, so no other client could
-redeem its refresh token, and no other method is known to be accepted. A deployment whose
-``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss
-logs the scope by name rather than failing silently.
+The refresh is redeemed against the SSO client the login itself used, resolved in the same order
+the login callback dispatches (Google, then Microsoft, then generic OIDC, from the client ids the
+proxy reconciles from the stored SSO row into the process environment at startup), authenticated the
+way that login authenticated: the Google and Microsoft logins and the non-PKCE generic path send
+HTTP Basic, while the PKCE generic path sends the credentials in the body when
+``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept only one of the two. An
+assertion can only exist if that client minted it, so no other client could redeem its refresh
+token, and no other method is known to be accepted. A login whose authorization request never asked
+for a refresh token captures none at all: the Google and Microsoft logins ask (``access_type=offline``,
+``offline_access``) only while an ``oauth2_id_jag`` server exists, and generic OIDC asks only when
+``GENERIC_SCOPE`` carries ``offline_access``, which is why that miss names the cause rather than
+failing silently.
 
 Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse
 onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged
@@ -77,7 +80,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
     SSOIdentityAssertion,
     assertion_expired,
     assertion_from_sso_login,
+    carry_exchanged_refresh_token,
     fetch_sso_identity_assertion,
+    is_saml_assertion,
     persist_sso_identity_assertion,
 )
 from litellm.types.llms.custom_http import httpxSpecialProvider
@@ -92,6 +97,10 @@ _SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion"
 # Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of
 # the exchange is replaced first. Matches the sibling per-user token store's skew.
 _DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0
+# Where fastapi-sso's Google login redeemed the authorization code, read from Google's discovery
+# document; Microsoft's mirrors ``CustomMicrosoftSSO``, override first, tenant template second.
+_GOOGLE_TOKEN_ENDPOINT: Final = "https://oauth2.googleapis.com/token"
+_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
 
 
 class AssertionRead(Protocol):
@@ -122,7 +131,7 @@ class FormPost(Protocol):
 
 @dataclass(frozen=True, slots=True)
 class SSOClientConfig:
-    """The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how."""
+    """The SSO client credentials a refresh_token grant has to authenticate as, and how."""
 
     token_endpoint: str
     client_id: str
@@ -131,32 +140,58 @@ class SSOClientConfig:
 
 
 def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None:
-    """The configured generic-OIDC client, or ``None`` when the deployment has none.
+    """The configured SSO client the login used, or ``None`` when the deployment has none.
 
-    Read from the process environment because that is where the login path reads it
-    (``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row
-    at startup, so this resolves to the same client that minted the assertion. ``None`` is an
-    ordinary state, not an error: a deployment signing in through a provider that captures no
+    Read from the process environment because that is where the login callback reads it and where
+    the proxy materializes the stored ``sso_config`` row at startup, and resolved in the callback's
+    own order (``GOOGLE_CLIENT_ID``, then ``MICROSOFT_CLIENT_ID``, then ``GENERIC_CLIENT_ID``, each
+    selecting its branch when set at all), so this is the same client that minted the assertion: a
+    deployment carrying both a Google and a generic client signed the user in through Google. ``None``
+    is an ordinary state, not an error: a deployment signing in through a provider that captures no
     assertion has nothing here to renew, and a client with no secret is not a confidential client
     that could redeem one.
 
-    ``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because
-    the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while
-    the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so
-    a constant here would authenticate the renewal differently from the sign-in that produced the
+    ``auth_method`` follows how each login authenticated at the token endpoint. The Google and
+    Microsoft logins (fastapi-sso) and the non-PKCE generic path send HTTP Basic, while the PKCE
+    generic path puts the credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set. A
+    constant here would authenticate the renewal differently from the sign-in that produced the
     refresh token and 401 against an IdP application registered for only one of the two.
     """
-    token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT")
-    client_id: Final = env.get("GENERIC_CLIENT_ID")
-    client_secret: Final = env.get("GENERIC_CLIENT_SECRET")
+    if env.get("GOOGLE_CLIENT_ID") is not None:
+        return _confidential_client(
+            _GOOGLE_TOKEN_ENDPOINT, env.get("GOOGLE_CLIENT_ID"), env.get("GOOGLE_CLIENT_SECRET"), "client_secret_basic"
+        )
+    if env.get("MICROSOFT_CLIENT_ID") is not None:
+        tenant: Final = env.get("MICROSOFT_TENANT")
+        return _confidential_client(
+            env.get("MICROSOFT_TOKEN_ENDPOINT")
+            or (_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE.format(tenant=tenant) if tenant else None),
+            env.get("MICROSOFT_CLIENT_ID"),
+            env.get("MICROSOFT_CLIENT_SECRET"),
+            "client_secret_basic",
+        )
+    includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
+    return _confidential_client(
+        env.get("GENERIC_TOKEN_ENDPOINT"),
+        env.get("GENERIC_CLIENT_ID"),
+        env.get("GENERIC_CLIENT_SECRET"),
+        "client_secret_post" if includes_client_id else "client_secret_basic",
+    )
+
+
+def _confidential_client(
+    token_endpoint: str | None,
+    client_id: str | None,
+    client_secret: str | None,
+    auth_method: MCPTokenEndpointAuthMethod,
+) -> SSOClientConfig | None:
     if not token_endpoint or not client_id or not client_secret:
         return None
-    includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
     return SSOClientConfig(
         token_endpoint=token_endpoint,
         client_id=client_id,
         client_secret=SecretStr(client_secret),
-        auth_method="client_secret_post" if includes_client_id else "client_secret_basic",
+        auth_method=auth_method,
     )
 
 
@@ -260,11 +295,19 @@ class SSOAssertionRefresher:
     async def refresh(
         self, user_id: str, assertion: SSOIdentityAssertion
     ) -> Result[SSOIdentityAssertion, RefreshFailure]:
+        if is_saml_assertion(assertion):
+            verbose_proxy_logger.warning(
+                "ID-JAG: the stored SAML assertion for user_id=%s has reached its NotOnOrAfter. A SAML assertion "
+                "cannot be renewed, so the user has to sign in through SAML SSO again.",
+                user_id,
+            )
+            return Error(RefreshFailure.of_rejected("a SAML assertion cannot be renewed; sign in again"))
         if assertion.refresh_token is None:
             verbose_proxy_logger.warning(
                 "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was "
-                "captured with it, so it cannot be renewed without another interactive sign-in. Add "
-                "'offline_access' to GENERIC_SCOPE so the SSO login captures one.",
+                "captured with it, so it cannot be renewed without another interactive sign-in. Google and Microsoft "
+                "logins ask for one once an oauth2_id_jag MCP server is registered, so a user who signed in before "
+                "that has to sign in again. Generic OIDC asks only with 'offline_access' in GENERIC_SCOPE.",
                 user_id,
             )
             return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in"))
@@ -272,11 +315,12 @@ class SSOAssertionRefresher:
         if config is None:
             verbose_proxy_logger.warning(
                 "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed "
-                "because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, "
-                "GENERIC_CLIENT_SECRET).",
+                "because the active SSO provider has no confidential client configured (GOOGLE_CLIENT_ID and "
+                "GOOGLE_CLIENT_SECRET, MICROSOFT_CLIENT_ID with MICROSOFT_CLIENT_SECRET and MICROSOFT_TENANT, or "
+                "GENERIC_TOKEN_ENDPOINT with GENERIC_CLIENT_ID and GENERIC_CLIENT_SECRET).",
                 user_id,
             )
-            return Error(RefreshFailure.of_rejected("the generic SSO client is not configured"))
+            return Error(RefreshFailure.of_rejected("the SSO client is not configured"))
 
         carried_refresh_token: Final = assertion.refresh_token.get_secret_value()
         # Whichever method the SSO login used for this client, since that is the one the IdP
@@ -323,10 +367,11 @@ class SSOAssertionRefresher:
                 user_id,
             )
             return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token"))
-        failure: Final = await self._store_renewal(user_id, previous, renewed)
+        carried: Final = carry_exchanged_refresh_token(previous, renewed)
+        failure: Final = await self._store_renewal(user_id, previous, carried)
         if failure is not None:
             return Error(failure)
-        return Ok(renewed)
+        return Ok(carried)
 
     async def _store_renewal(
         self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion
@@ -401,9 +446,12 @@ class RefreshingSSOAssertionStore:
         self._skew = timedelta(seconds=expiry_skew_seconds)
         self._clock = clock
 
+    async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
+        return await self._inner.fetch(user_id)
+
     async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
         assertion: Final = await self._inner.fetch(user_id)
-        if not self._expiring(assertion):
+        if assertion is None or is_saml_assertion(assertion) or not self._expiring(assertion):
             return assertion
         await self._coordinator().run(
             user_id,

litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py

@@ -1,9 +1,9 @@
 """Store for the enterprise IdP identity assertion captured at SSO login (EMA).
 
-The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 8693
-``subject_token``. A front-door client holds an identity-only ``llm_session_`` bearer, not an
-IdP assertion, so the assertion captured at the one SSO login is the only usable subject
-source for it. This module owns both sides of that state: the SSO callback persists here
+The ``oauth2_id_jag`` egress arm needs the user's IdP identity assertion as its RFC 8693
+``subject_token``: the OIDC ``id_token``, or the signed SAML assertion when the login was SAML.
+A front-door client holds an identity-only ``llm_session_`` bearer, not an IdP assertion, so
+the assertion captured at the one SSO login is the only usable subject source for it. This module owns both sides of that state: the SSO callback persists here
 (write-through to the DB so a login on one pod is visible to every pod) and the resolver
 seam reads back by ``user_id``. Retention is gated on an ``oauth2_id_jag`` server actually
 being registered, so a gateway with no EMA upstream never stores bearer material.
@@ -17,9 +17,11 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st
 
 from __future__ import annotations
 
+import base64
 import json
 from collections.abc import Mapping, Sequence
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol
 
 import jwt
@@ -35,7 +37,15 @@ if TYPE_CHECKING:
 
     from litellm.proxy.utils import PrismaClient
 
+"""What a Google or Microsoft authorization request has to carry for the IdP to issue a refresh token
+beside the id_token. Asked for only while an ``oauth2_id_jag`` server exists to renew an assertion for, so
+a deployment without one sees no extra consent."""
+MICROSOFT_REFRESH_TOKEN_SCOPE: Final = "offline_access"
+GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"})
+
+
 _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion"
+SAML2_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:saml2"
 _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str)
 _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
 
@@ -69,8 +79,10 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable:
 
 
 class SSOIdentityAssertion(LiteLLMBaseModel):
-    """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token,
-    ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login."""
+    """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token (the OIDC
+    id_token, or the base64 signed ``<saml:Assertion>`` when ``subject_token_type`` names SAML),
+    ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.
+    ``subject_token_type`` is ``None`` when the exchange should send the operator-configured type."""
 
     model_config = ConfigDict(frozen=True)
 
@@ -78,6 +90,11 @@ class SSOIdentityAssertion(LiteLLMBaseModel):
     refresh_token: SecretStr | None = None
     issuer: str | None = None
     expires_at: datetime | None = None
+    subject_token_type: str | None = None
+    exchanged_refresh_token: SecretStr | None = None
+    exchanged_refresh_token_expires_at: datetime | None = None
+    exchanged_refresh_token_endpoint: str | None = None
+    exchanged_refresh_token_client_id: str | None = None
 
 
 class SSOAssertionCache:
@@ -130,6 +147,11 @@ class _StoredAssertionPayload(LiteLLMBaseModel):
     refresh_token: str | None = None
     issuer: str | None = None
     expires_at: datetime | None = None
+    subject_token_type: str | None = None
+    exchanged_refresh_token: str | None = None
+    exchanged_refresh_token_expires_at: datetime | None = None
+    exchanged_refresh_token_endpoint: str | None = None
+    exchanged_refresh_token_client_id: str | None = None
 
 
 def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIdentityAssertion | None:
@@ -159,6 +181,85 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden
     )
 
 
+def assertion_from_saml_login(
+    assertion_xml: bytes, not_on_or_after: int | None, issuer: str | None
+) -> SSOIdentityAssertion:
+    """The carrier for a SAML login: the signed ``<saml:Assertion>`` element the IdP issued, base64
+    encoded the way RFC 8693 carries a ``saml2`` subject token, with its own validity window as
+    the expiry and no refresh token, since SAML has none and the user has to sign in again."""
+    return SSOIdentityAssertion(
+        id_token=SecretStr(base64.urlsafe_b64encode(assertion_xml).rstrip(b"=").decode("ascii")),
+        issuer=issuer,
+        expires_at=datetime.fromtimestamp(not_on_or_after, tz=timezone.utc) if not_on_or_after is not None else None,
+        subject_token_type=SAML2_SUBJECT_TOKEN_TYPE,
+    )
+
+
+def is_saml_assertion(assertion: SSOIdentityAssertion) -> bool:
+    return assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE
+
+
+def with_login_refresh_token(
+    assertion: SSOIdentityAssertion,
+    refresh_token: str,
+    expires_in: int | None,
+    now: datetime,
+    *,
+    endpoint: str | None = None,
+    client_id: str | None = None,
+) -> SSOIdentityAssertion:
+    """The same login, now also carrying the refresh token the ID-JAG client exchanged the assertion
+    for, with the lifetime the IdP gave that token. It is kept apart from ``refresh_token``, which an
+    OIDC login holds for the SSO client and which only that client can redeem. The assertion and its
+    own expiry stay as they were, so a server that presents the login assertion directly is unaffected."""
+    return assertion.model_copy(
+        update={
+            "exchanged_refresh_token": SecretStr(refresh_token),
+            "exchanged_refresh_token_expires_at": now + timedelta(seconds=expires_in)
+            if expires_in is not None
+            else None,
+            "exchanged_refresh_token_endpoint": endpoint,
+            "exchanged_refresh_token_client_id": client_id,
+        }
+    )
+
+
+def carry_exchanged_refresh_token(
+    previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion
+) -> SSOIdentityAssertion:
+    """The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login.
+    Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to
+    them and stays valid for as long as the IdP said."""
+    return renewed.model_copy(
+        update={
+            "exchanged_refresh_token": previous.exchanged_refresh_token,
+            "exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at,
+            "exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint,
+            "exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id,
+        }
+    )
+
+
+def exchanged_refresh_token_issued_for(assertion: SSOIdentityAssertion, endpoint: str, client_id: str) -> bool:
+    """Whether the exchanged refresh token was minted at this token endpoint for this client. A refresh
+    token is bound to the client that received it, so one minted for another server's IdP client must
+    never be posted to this server's endpoint."""
+    return (
+        assertion.exchanged_refresh_token_endpoint == endpoint
+        and assertion.exchanged_refresh_token_client_id == client_id
+    )
+
+
+def login_refresh_token_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool:
+    """Whether the refresh token exchanged at login has passed the lifetime the IdP gave it. One
+    with no recorded lifetime is left for the IdP to judge, like an assertion with no ``exp``."""
+    expires_at: Final = assertion.exchanged_refresh_token_expires_at
+    if expires_at is None:
+        return False
+    normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc)
+    return normalized <= now
+
+
 def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool:
     """Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is
     treated as usable and left for the IdP to reject, since the store records what the id_token
@@ -212,6 +313,27 @@ async def persist_sso_identity_assertion(
         **({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}),
         **({"issuer": assertion.issuer} if assertion.issuer else {}),
         **({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}),
+        **({"subject_token_type": assertion.subject_token_type} if assertion.subject_token_type else {}),
+        **(
+            {"exchanged_refresh_token": assertion.exchanged_refresh_token.get_secret_value()}
+            if assertion.exchanged_refresh_token
+            else {}
+        ),
+        **(
+            {"exchanged_refresh_token_expires_at": assertion.exchanged_refresh_token_expires_at.isoformat()}
+            if assertion.exchanged_refresh_token_expires_at
+            else {}
+        ),
+        **(
+            {"exchanged_refresh_token_endpoint": assertion.exchanged_refresh_token_endpoint}
+            if assertion.exchanged_refresh_token_endpoint
+            else {}
+        ),
+        **(
+            {"exchanged_refresh_token_client_id": assertion.exchanged_refresh_token_client_id}
+            if assertion.exchanged_refresh_token_client_id
+            else {}
+        ),
     }
     encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload)))
     await _assertion_table(prisma_client).upsert(
@@ -250,6 +372,11 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
         refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None,
         issuer=payload.issuer,
         expires_at=payload.expires_at,
+        subject_token_type=payload.subject_token_type,
+        exchanged_refresh_token=SecretStr(payload.exchanged_refresh_token) if payload.exchanged_refresh_token else None,
+        exchanged_refresh_token_expires_at=payload.exchanged_refresh_token_expires_at,
+        exchanged_refresh_token_endpoint=payload.exchanged_refresh_token_endpoint,
+        exchanged_refresh_token_client_id=payload.exchanged_refresh_token_client_id,
     )
 
 
@@ -290,6 +417,12 @@ class SSOAssertionStore(Protocol):
 
     async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ...
 
+    async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
+        """The row as stored, never renewed on the way out. For a server that presents the refresh
+        token exchanged at login, the SSO id_token's own lifetime is irrelevant and an SSO renewal
+        failing must not stand between that server and a still valid refresh token."""
+        ...
+
 
 class DbSSOAssertionStore:
     """The live store: the row the SSO callback wrote, read back by ``user_id``.
@@ -308,6 +441,9 @@ class DbSSOAssertionStore:
         except Exception as exc:  # noqa: BLE001  # any driver/storage failure is an outage, not an absence
             raise AssertionStoreUnavailable(str(exc)) from exc
 
+    async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None:
+        return await self.fetch(user_id)
+
     async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None:
         try:
             return await _read_assertion_from_db(user_id)

litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py

@@ -349,6 +349,14 @@ class IdJagResourceClient(LiteLLMBaseModel):
     client_auth: ClientAuth
 
 
+class IdJagSubjectSource(str, Enum):
+    """What ID-JAG leg 1 presents as the RFC 8693 subject_token: the identity assertion captured at SSO
+    login, or the refresh token that assertion was exchanged for at the IdP when the user logged in."""
+
+    login_assertion = "login_assertion"
+    refresh_token = "refresh_token"
+
+
 class IdJagConfig(HeaderCarrier):
     """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange").
 
@@ -368,6 +376,7 @@ class IdJagConfig(HeaderCarrier):
     client_auth: ClientAuth
     resource_client: IdJagResourceClient
     subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token"
+    subject_source: IdJagSubjectSource = IdJagSubjectSource.login_assertion
     audience: str | None = None
     resource: str | None = None
     scopes: tuple[str, ...] = ()

litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py

@@ -25,7 +25,16 @@ from litellm.proxy.management_endpoints.ui_sso import (
 
 
 async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
-    user: Final = await load_active_user_by_id(user_id)
+    return await _consent_teams(await load_active_user_by_id(user_id))
+
+
+async def lookup_consent_teams_from_database(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
+    """The same teams read from the database row, for a grant that is consumed once it is settled: the
+    membership the minter will re-check must be the membership the team was picked from."""
+    return await _consent_teams(await load_active_user_by_id(user_id, source="database"))
+
+
+async def _consent_teams(user: LiteLLM_UserTable | ReloadUserFailure) -> tuple[ConsentTeam, ...] | ReloadUserFailure:
     if isinstance(user, str):
         return user
     details: Final = await _team_details(user.teams)
@@ -39,7 +48,7 @@ async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | Reload
 
 
 async def mint_proxy_credential(
-    user_id: str, team_id: str | None
+    user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None
 ) -> MintedProxyCredential | ProxyCredentialMintFailure:
     """Mint the ``lite login`` credential for a consented grant. Membership is checked
     live against the database row, so a team the user left between consent and redemption (or between refreshes)
@@ -53,7 +62,9 @@ async def mint_proxy_credential(
     user row handed to the minter carries no team list, exactly like ``lite login``'s, so
     the minter's own first-team fallback stays inert. The credential carries the role the
     proxy already enforces for the user on every request, so a row with no role (JWT auth's
-    upsert writes none) mints as an internal user instead of being refused."""
+    upsert writes none) mints as an internal user instead of being refused. ``allowed_routes``
+    narrows the credential to those route groups, the way a virtual key's ``allowed_routes`` does,
+    for a grant whose scope asked for less than the full credential."""
     user: Final = await load_active_user_by_id(user_id, source="database")
     if isinstance(user, str):
         return user
@@ -75,6 +86,7 @@ async def mint_proxy_credential(
         team_alias=selected.team_alias,
         team_models=selected.team_models,
         team_model_aliases=selected.team_model_aliases,
+        allowed_routes=allowed_routes,
     )
     return MintedProxyCredential(
         key=key,

litellm/proxy/_experimental/mcp_server/subscriptions.py

@@ -55,6 +55,7 @@ from litellm.proxy._experimental.mcp_server.catalog import (
 )
 from litellm.proxy._experimental.mcp_server.contracts import OperationContext
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
 from litellm.types.mcp import MCP_LEGACY_VERSIONS
 from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
@@ -300,6 +301,29 @@ async def _listing(event: ServerEvent, context: OperationContext) -> str | None:
     raise MCPError(code=-32603, message="Subscription catalog exceeds the pagination limit")
 
 
+async def _enforce_rate_limit(context: OperationContext, server: MCPServer) -> None:
+    await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, server)  # pyright: ignore[reportPrivateUsage]  # same admission policy as other gateway operations
+
+
+async def _unmetered(_: OperationContext, __: MCPServer) -> None:
+    return None
+
+
+async def _admitted() -> ProxyRateLimitError | None:
+    return None
+
+
+@asynccontextmanager
+async def _listener_admission(targets: tuple[SubscriptionTarget, ...]) -> AsyncGenerator[None]:
+    scope: Final = operations._mcp_server_admission_memo.set(  # pyright: ignore[reportPrivateUsage]  # the listener's own catalog refreshes reuse its admission
+        {target.server.server_id: asyncio.create_task(_admitted()) for target in targets}
+    )
+    try:
+        yield
+    finally:
+        operations._mcp_server_admission_memo.reset(scope)  # pyright: ignore[reportPrivateUsage]  # restore the caller's admission scope
+
+
 async def _recheck_authorization(check: Callable[[], Awaitable[OperationContext]]) -> None:
     while True:
         await asyncio.sleep(15)
@@ -361,7 +385,7 @@ class GatewaySubscriptions:
         if ctx.protocol_version != "2026-07-28":
             raise MCPError(code=-32601, message="Subscriptions require the modern MCP protocol")
         async with self._capacity(owner):
-            return await self._listen(ctx, params, authorize)
+            return await self._listen(ctx, params, authorize, _enforce_rate_limit)
 
     @asynccontextmanager
     async def _capacity(self, owner: str) -> AsyncGenerator[None]:
@@ -414,6 +438,7 @@ class GatewaySubscriptions:
                     ctx,
                     SubscriptionFilter(tools_list_changed=True, prompts_list_changed=True, resources_list_changed=True),
                     authorization,
+                    _unmetered,
                 )
                 if (ready := self._initializing.get(key)) is not None:
                     ready.set()
@@ -427,6 +452,7 @@ class GatewaySubscriptions:
         ctx: ServerRequestContext,
         notifications: SubscriptionFilter,
         authorization: SubscriptionAuthorization,
+        admit: Callable[[OperationContext, MCPServer], Awaitable[None]],
         scope: anyio.CancelScope | None = None,
         *,
         task_status: TaskStatus[None],
@@ -443,7 +469,7 @@ class GatewaySubscriptions:
 
             async with self._capacity(authorization.owner):
                 await self._listen(
-                    ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, ready
+                    ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, admit, ready
                 )
 
     async def legacy_resource(self, ctx: ServerRequestContext, uri: str, *, subscribe: bool) -> None:
@@ -464,6 +490,7 @@ class GatewaySubscriptions:
                     ctx,
                     SubscriptionFilter(resource_subscriptions=[uri]),
                     state.authorization,
+                    _enforce_rate_limit,
                     scope,
                 )
                 state.resources = {**state.resources, uri: scope}
@@ -476,6 +503,7 @@ class GatewaySubscriptions:
         ctx: ServerRequestContext,
         params: SubscriptionsListenRequestParams,
         authorize: Callable[[], Awaitable[OperationContext]],
+        admit: Callable[[OperationContext, MCPServer], Awaitable[None]],
         ready: Callable[[SubscriptionFilter], None] | None = None,
     ) -> SubscriptionsListenResult:
         context: Final = await authorize()
@@ -501,7 +529,7 @@ class GatewaySubscriptions:
             return fresh
 
         async def relay(target: SubscriptionTarget, *, task_status: TaskStatus[SubscriptionFilter]) -> None:
-            await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, target.server)  # pyright: ignore[reportPrivateUsage]  # same admission policy as other gateway operations
+            await admit(context, target.server)
 
             legacy: Final = LegacyNotifications()
 
@@ -567,7 +595,7 @@ class GatewaySubscriptions:
             finally:
                 self._upstream_slots.release()
 
-        async with anyio.create_task_group() as group:
+        async with _listener_admission(targets), anyio.create_task_group() as group:
             honored: Final = tuple(
                 [SubscriptionFilter.model_validate(await group.start(relay, target)) for target in targets]
             )

litellm/proxy/_lazy_features.py

@@ -155,18 +155,13 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
         module_path="litellm.proxy.management_endpoints.mcp_management_endpoints",
         path_prefixes=("/v1/mcp/",),
     ),
-    LazyFeature(
-        # Also serves /.well-known/oauth-* (OAuth metadata discovery).
-        # No /mcp/oauth prefix here: the mounted /mcp sub-app would
-        # shadow it, and there are no actual routes there anyway.
-        name="mcp_byok_oauth",
-        module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
-        path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
-    ),
     LazyFeature(
         # Serves OAuth dance endpoints (/authorize, /token, /callback,
         # /register) plus several /.well-known/ discovery URLs at the proxy
         # root — needed for MCP-over-OAuth flows even before /mcp is hit.
+        # Registered before mcp_byok_oauth, which also claims the root
+        # /.well-known/oauth-authorization-server: the root document is the
+        # gateway's own authorization server metadata (token endpoint /token).
         name="mcp_discoverable",
         module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints",
         path_prefixes=(
@@ -185,6 +180,14 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
         # Catches the /{mcp_server_name}/authorize|token|register variants.
         path_suffixes=("/authorize", "/token", "/register"),
     ),
+    LazyFeature(
+        # Also serves /.well-known/oauth-* (OAuth metadata discovery).
+        # No /mcp/oauth prefix here: the mounted /mcp sub-app would
+        # shadow it, and there are no actual routes there anyway.
+        name="mcp_byok_oauth",
+        module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
+        path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
+    ),
     LazyFeature(
         name="mcp_rest",
         module_path="litellm.proxy._experimental.mcp_server.rest_endpoints",

litellm/proxy/_lazy_openapi_snapshot.json

@@ -14261,6 +14261,19 @@
               "description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.",
               "title": "Fail On Error"
             },
+            "forward_messages_as_sent": {
+              "anyOf": [
+                {
+                  "type": "boolean"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "default": false,
+              "description": "Only honored by guardrail='generic_guardrail_api'. When True, structured_messages reaches the guardrail exactly as the client sent it, so rows can carry provider-native content blocks and keys the OpenAI chat types do not declare. Default False sends OpenAI-format chat rows with undeclared keys removed.",
+              "title": "Forward Messages As Sent"
+            },
             "guard_name": {
               "anyOf": [
                 {
@@ -16612,6 +16625,19 @@
               "description": "HTTP timeout in seconds for checking attached files. Default: 10.",
               "title": "File Guardrail Timeout"
             },
+            "forward_messages_as_sent": {
+              "anyOf": [
+                {
+                  "type": "boolean"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "default": false,
+              "description": "Only honored by guardrail='generic_guardrail_api'. When True, structured_messages reaches the guardrail exactly as the client sent it, so rows can carry provider-native content blocks and keys the OpenAI chat types do not declare. Default False sends OpenAI-format chat rows with undeclared keys removed.",
+              "title": "Forward Messages As Sent"
+            },
             "gateway_name": {
               "anyOf": [
                 {
@@ -34541,6 +34567,17 @@
               ],
               "title": "Id Jag Resource Token Endpoint"
             },
+            "id_jag_subject_source": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Id Jag Subject Source"
+            },
             "redirect_uris": {
               "anyOf": [
                 {
@@ -35771,7 +35808,7 @@
       "/.well-known/oauth-authorization-server": {
         "get": {
           "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
-          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get",
+          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2",
           "parameters": [
             {
               "in": "query",
@@ -35819,7 +35856,7 @@
       "/.well-known/oauth-authorization-server/mcp": {
         "get": {
           "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.",
-          "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get",
+          "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2",
           "responses": {
             "200": {
               "content": {
@@ -35839,7 +35876,7 @@
       "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": {
         "get": {
           "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}",
-          "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get",
+          "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2",
           "parameters": [
             {
               "in": "path",
@@ -35880,7 +35917,7 @@
       "/.well-known/oauth-authorization-server/{mcp_server_name}": {
         "get": {
           "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
-          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get",
+          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2",
           "parameters": [
             {
               "in": "path",
@@ -35928,7 +35965,7 @@
       "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": {
         "get": {
           "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.",
-          "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get",
+          "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2",
           "parameters": [
             {
               "in": "path",
@@ -35968,7 +36005,7 @@
       },
       "/.well-known/oauth-protected-resource": {
         "get": {
-          "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get",
+          "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2",
           "responses": {
             "200": {
               "content": {
@@ -36004,7 +36041,7 @@
       "/.well-known/oauth-protected-resource/mcp": {
         "get": {
           "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.",
-          "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get",
+          "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2",
           "responses": {
             "200": {
               "content": {
@@ -36024,7 +36061,7 @@
       "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": {
         "get": {
           "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.",
-          "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get",
+          "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2",
           "parameters": [
             {
               "in": "path",
@@ -36065,7 +36102,7 @@
       "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": {
         "get": {
           "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
-          "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get",
+          "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2",
           "parameters": [
             {
               "in": "path",
@@ -36194,9 +36231,27 @@
         },
         "Body_token_endpoint__mcp_server_name__token_post": {
           "properties": {
+            "assertion": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Assertion"
+            },
             "client_id": {
-              "title": "Client Id",
-              "type": "string"
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Client Id"
             },
             "client_secret": {
               "anyOf": [
@@ -36293,17 +36348,34 @@
             }
           },
           "required": [
-            "grant_type",
-            "client_id"
+            "grant_type"
           ],
           "title": "Body_token_endpoint__mcp_server_name__token_post",
           "type": "object"
         },
         "Body_token_endpoint_token_post": {
           "properties": {
+            "assertion": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Assertion"
+            },
             "client_id": {
-              "title": "Client Id",
-              "type": "string"
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Client Id"
             },
             "client_secret": {
               "anyOf": [
@@ -36400,8 +36472,7 @@
             }
           },
           "required": [
-            "grant_type",
-            "client_id"
+            "grant_type"
           ],
           "title": "Body_token_endpoint_token_post",
           "type": "object"
@@ -37324,6 +37395,17 @@
               ],
               "title": "Id Jag Resource Token Endpoint"
             },
+            "id_jag_subject_source": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Id Jag Subject Source"
+            },
             "redirect_uris": {
               "anyOf": [
                 {
@@ -38151,7 +38233,7 @@
       "/.well-known/oauth-authorization-server": {
         "get": {
           "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
-          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2",
+          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get",
           "parameters": [
             {
               "in": "query",
@@ -38199,7 +38281,7 @@
       "/.well-known/oauth-authorization-server/mcp": {
         "get": {
           "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.",
-          "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2",
+          "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get",
           "responses": {
             "200": {
               "content": {
@@ -38219,7 +38301,7 @@
       "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": {
         "get": {
           "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}",
-          "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2",
+          "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get",
           "parameters": [
             {
               "in": "path",
@@ -38260,7 +38342,7 @@
       "/.well-known/oauth-authorization-server/{mcp_server_name}": {
         "get": {
           "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.",
-          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2",
+          "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get",
           "parameters": [
             {
               "in": "path",
@@ -38308,7 +38390,7 @@
       "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": {
         "get": {
           "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.",
-          "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2",
+          "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get",
           "parameters": [
             {
               "in": "path",
@@ -38348,7 +38430,7 @@
       },
       "/.well-known/oauth-protected-resource": {
         "get": {
-          "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2",
+          "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get",
           "responses": {
             "200": {
               "content": {
@@ -38384,7 +38466,7 @@
       "/.well-known/oauth-protected-resource/mcp": {
         "get": {
           "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.",
-          "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2",
+          "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get",
           "responses": {
             "200": {
               "content": {
@@ -38404,7 +38486,7 @@
       "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": {
         "get": {
           "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.",
-          "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2",
+          "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get",
           "parameters": [
             {
               "in": "path",
@@ -38445,7 +38527,7 @@
       "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": {
         "get": {
           "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
-          "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2",
+          "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get",
           "parameters": [
             {
               "in": "path",
@@ -40729,6 +40811,17 @@
               ],
               "title": "Id Jag Resource Token Endpoint"
             },
+            "id_jag_subject_source": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Id Jag Subject Source"
+            },
             "redirect_uris": {
               "anyOf": [
                 {
@@ -44809,6 +44902,17 @@
               ],
               "title": "Id Jag Resource Token Endpoint"
             },
+            "id_jag_subject_source": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Id Jag Subject Source"
+            },
             "redirect_uris": {
               "anyOf": [
                 {

litellm/proxy/_types.py

@@ -551,6 +551,7 @@ class LiteLLMRoutes(enum.Enum):
         + list(agent_inference_routes)
         + model_info_routes
     )
+    llm_api_routes_without_mcp = sorted(set(llm_api_routes) - set(mcp_inference_routes))
     info_routes = [
         "/key/info",
         "/key/health",
@@ -4762,6 +4763,26 @@ class JWTIssuerConfig(LiteLLMBaseModel):
 DEFAULT_JWKS_STALE_TTL: Final = 3600
 
 
+class IdJagRequestingApp(LiteLLMBaseModel):
+    """A requesting application the operator lets redeem an ID-JAG (Identity Assertion
+    Authorization Grant) at the MCP gateway's ``POST /token``. ``client_id`` is the value
+    the identity provider puts in the assertion's ``client_id`` claim and the app sends as
+    the ``client_id`` form field. A ``client_secret`` makes the app a confidential client
+    that must authenticate with it (HTTP Basic or the ``client_secret`` form field); without
+    one the app is a public client."""
+
+    model_config = ConfigDict(frozen=True)
+
+    client_id: str = Field(
+        min_length=1, description="The requesting app's client_id as the identity provider asserts it."
+    )
+    client_secret: str | None = Field(
+        default=None,
+        min_length=1,
+        description="Shared secret the app must present at the gateway's token endpoint. Omit for a public client.",
+    )
+
+
 class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     """
     A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth.
@@ -4942,6 +4963,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
         default=None,
         description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.",
     )
+    id_jag_requesting_apps: tuple[IdJagRequestingApp, ...] | None = Field(
+        default=None,
+        description="Requesting apps allowed to redeem an ID-JAG at the MCP gateway's POST /token with the RFC 7523 jwt-bearer grant. Unset or empty turns that grant off.",
+    )
     #########################################################
 
     def __init__(self, **kwargs: Any) -> None:

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -6,9 +6,9 @@ from collections.abc import Mapping
 from types import MappingProxyType
 from typing import Final
 
-from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
 from fastapi.responses import JSONResponse
-from pydantic import JsonValue
+from pydantic import JsonValue, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import (
     AnthropicExceptionMapping,
 )
 from litellm.integrations.custom_guardrail import ModifyResponseException
+from litellm.litellm_core_utils.tool_fields import first_validation_error
 from litellm.llms.anthropic.pass_through.context_management import (
     AnthropicContextManagementError,
 )
@@ -32,6 +33,7 @@ from litellm.proxy.common_request_processing import (
     log_llm_api_exception,
     proxy_exception_from_http_exception,
     proxy_exception_from_route_error,
+    resolve_counted_tool_fields,
     resolve_litellm_call_id,
 )
 from litellm.proxy.common_utils.error_body_call_id import error_body_call_id
@@ -47,6 +49,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
     with_litellm_call_id,
 )
 from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli
+from litellm.router import Router
 
 router: Final = APIRouter()
 
@@ -91,6 +94,31 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO
     return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
 
 
+_TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "system")
+
+
+def _token_count_request(fields: Mapping[str, object]) -> TokenCountRequest | ProxyException:
+    try:
+        return TokenCountRequest.model_validate({name: fields[name] for name in _TOKEN_COUNT_FIELDS if name in fields})
+    except ValidationError as e:
+        return ProxyException(
+            message=first_validation_error(e),
+            type="invalid_request_error",
+            param=str(e.errors(include_url=False)[0]["loc"][0]),
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
+def _token_count_fields(
+    data: Mapping[str, object], auto_drop_params: bool, llm_router: Router | None, team_id: str | None
+) -> Mapping[str, object] | ProxyException:
+    fields: Final = {
+        name: data[name] for name in (*_TOKEN_COUNT_FIELDS, "drop_params", "additional_drop_params") if name in data
+    }
+    resolving: Final = {**fields, "drop_params": True} if auto_drop_params else fields
+    return resolve_counted_tool_fields(resolving, llm_router, team_id)
+
+
 def _strip_total_tokens_from_anthropic_response(response: object) -> None:
     """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
     injects into Anthropic /v1/messages responses.
@@ -351,7 +379,7 @@ async def count_tokens(
     
     Returns: {"input_tokens": <number>}
     """
-    from litellm.proxy.proxy_server import count_request_tokens, proxy_config
+    from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config
 
     litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
     try:
@@ -368,26 +396,21 @@ async def count_tokens(
         if not messages:
             raise HTTPException(status_code=400, detail="messages parameter is required")
 
-        # Create TokenCountRequest for the internal endpoint
-        from litellm.proxy._types import TokenCountRequest
-
-        token_request: Final = TokenCountRequest(
-            model=model_name,
-            messages=messages,
-            tools=data.get("tools"),
-            system=data.get("system"),
+        auto_drop_params: Final = should_auto_drop_params_for_agentic_cli(
+            request.headers.get("user-agent", ""), "drop_params" in data, proxy_config
         )
+        token_fields: Final = _token_count_fields(data, auto_drop_params, llm_router, user_api_key_dict.team_id)
+        token_request: Final = (
+            token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields)
+        )
+        if isinstance(token_request, ProxyException):
+            return _anthropic_error_json_response(token_request, request)
 
         token_response: Final = await count_request_tokens(
             request=token_request,
             call_endpoint=True,
             messages_format="anthropic",
-            **_count_tokens_body_options(
-                data,
-                should_auto_drop_params_for_agentic_cli(
-                    request.headers.get("user-agent", ""), "drop_params" in data, proxy_config
-                ),
-            ),
+            **_count_tokens_body_options(data, auto_drop_params),
         )
     except HTTPException as e:
         return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)

litellm/proxy/auth/auth_checks.py

@@ -3960,6 +3960,7 @@ class ExperimentalUIJWTToken:
         team_models: Sequence[str] | None = None,
         team_model_aliases: Mapping[str, str] | None = None,
         max_budget: float | None = None,
+        allowed_routes: Sequence[str] | None = None,
     ) -> str:
         """
         Generate a JWT token for CLI authentication with configurable expiration.
@@ -3973,6 +3974,8 @@ class ExperimentalUIJWTToken:
             team_alias: Team alias for the selected team, if available
             team_models: Model allowlist granted by the selected team
             team_model_aliases: Team model aliases for the selected team
+            allowed_routes: Route groups the session may call, enforced like a virtual key's
+                ``allowed_routes`` on every route; None leaves every route the role allows
 
         Returns:
             Encrypted JWT token string
@@ -4017,6 +4020,7 @@ class ExperimentalUIJWTToken:
             max_parallel_requests=None,
             user_role=LitellmUserRoles(user_info.user_role),
             is_session_token=True,
+            allowed_routes=list(allowed_routes) if allowed_routes is not None else None,
         )
 
         return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)

litellm/proxy/auth/auth_utils.py

@@ -551,6 +551,29 @@ def _iter_fallback_targets(value: object, depth: int) -> Iterator[str | Mapping[
                     yield from _iter_fallback_targets(target_list, depth + 1)
 
 
+def _reject_client_chosen_federation(litellm_params: Mapping[str, object], *, manages_deployments: bool) -> None:
+    reject_server_owned_wif_params(litellm_params)
+    if not manages_deployments:
+        reject_federated_credential_reference(litellm_params)
+
+
+def _iter_user_config_litellm_params(user_config: object) -> Iterator[Mapping[str, object]]:
+    config: Final = _coerce_metadata_to_dict(user_config)
+    if config is None:
+        return
+    if (defaults := _coerce_metadata_to_dict(config.get("default_litellm_params"))) is not None:
+        yield defaults
+    model_list: Final = config.get("model_list")
+    if not isinstance(model_list, list):
+        return
+    for deployment in model_list:
+        if (
+            isinstance(deployment, Mapping)
+            and (params := _coerce_metadata_to_dict(deployment.get("litellm_params"))) is not None
+        ):
+            yield params
+
+
 def iter_request_fallback_targets(request_body: Mapping[str, object]) -> Iterator[str | Mapping[str, object]]:
     for value in _iter_fallback_field_values(request_body):
         yield from _iter_fallback_targets(value, 0)
@@ -658,11 +681,11 @@ def is_request_body_safe(
                 _reject_url_valued_fallback_target(target_model)
         elif isinstance(target, str):
             _reject_url_valued_fallback_target(target)
+    for user_config_params in _iter_user_config_litellm_params(request_body.get("user_config")):
+        _reject_client_chosen_federation(user_config_params, manages_deployments=manages_deployments)
     litellm_params: Final = _coerce_metadata_to_dict(request_body.get("litellm_params"))
     if litellm_params is not None:
-        reject_server_owned_wif_params(litellm_params)
-        if not manages_deployments:
-            reject_federated_credential_reference(litellm_params)
+        _reject_client_chosen_federation(litellm_params, manages_deployments=manages_deployments)
         litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
         if litellm_params_metadata is not None:
             _check_banned_params(

litellm/proxy/auth/handle_jwt.py

@@ -35,6 +35,7 @@ from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
 from litellm.proxy._types import (
     DEFAULT_JWKS_STALE_TTL,
     RBAC_ROLES,
+    IdJagRequestingApp,
     JWKKeyValue,
     JWTAuthBuilderResult,
     JWTIssuerConfig,
@@ -99,6 +100,10 @@ class NoMatchingJWTPublicKeyError(Exception):
     """Raised when a JWKS endpoint returns no key matching the requested ``kid``."""
 
 
+class IdJagIsNotABearerTokenError(Exception):
+    """Raised when bearer auth is handed an ID-JAG, which only the token endpoint may redeem."""
+
+
 class JWKSUnreachableError(Exception):
     """Raised when an IdP's JWKS / OIDC discovery endpoint is unreachable and no cached copy is left to fall back on."""
 
@@ -114,9 +119,23 @@ _CACHED_VALUE: Final = TypeAdapter(object)
 _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str)
 
 
+ID_JAG_TYP: Final = "oauth-id-jag+jwt"
+
+
+@dataclass(frozen=True, slots=True)
+class AssertionRejection:
+    """Why a JWT assertion failed verification. The reason is for the proxy log only, since
+    PyJWT's wording can quote the token's claims."""
+
+    reason: str
+
+
 class _JWTAuthSettings(Protocol):
     """The JWT auth settings block this handler reads back through ``getattr``, when one is configured."""
 
+    @property
+    def id_jag_requesting_apps(self) -> Sequence[IdJagRequestingApp] | None: ...
+
     @property
     def issuers(self) -> Sequence[JWTIssuerConfig] | None: ...
 
@@ -282,7 +301,7 @@ class JWTHandler:
         return is_jwt_token(token=token)
 
     @staticmethod
-    def get_unverified_claims(token: str) -> dict | None:
+    def get_unverified_claims(token: str) -> dict[str, object] | None:
         """
         Decode JWT claims without signature verification.
         Used for routing decisions before selecting validation path.
@@ -949,7 +968,7 @@ class JWTHandler:
         keys: Final = response_json["keys"] if "keys" in response_json else response_json
         return cast(JWKKeyValue, keys)  # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material
 
-    async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> dict:
+    async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> JWTKeyItem:
         resolved_jwks_url: Final = await self._resolve_jwks_url(jwks_url)
         keys: Final = await self._cached_with_stale_fallback(
             cache_key=f"litellm_jwt_auth_keys_{resolved_jwks_url}",
@@ -960,11 +979,11 @@ class JWTHandler:
 
         public_key: Final = self.parse_keys(keys=keys, kid=kid)
         if public_key is not None:
-            return cast(dict, public_key)
+            return public_key
 
         raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={resolved_jwks_url}, kid={kid}")
 
-    async def get_public_key(self, kid: str | None) -> dict:
+    async def get_public_key(self, kid: str | None) -> JWTKeyItem:
         keys_url: Final = os.getenv("JWT_PUBLIC_KEY_URL")
 
         if keys_url is None:
@@ -1151,7 +1170,7 @@ class JWTHandler:
         # its jwks_uri, matching JWTIssuerConfig.jwks_url's documented fallback.
         return f"{issuer_config.issuer.rstrip('/')}/.well-known/openid-configuration"
 
-    def _get_claim_value_for_issuer_mapping(self, token: dict, claim_field: str) -> object:
+    def _get_claim_value_for_issuer_mapping(self, token: Mapping[str, object], claim_field: str) -> object:
         """Resolve a mapped claim from ``token``.
 
         Returns ``None`` when the field is absent or empty so that mapped claims
@@ -1168,8 +1187,10 @@ class JWTHandler:
             return None
         return claim_value
 
-    def _apply_issuer_claim_mappings(self, token: dict, issuer_config: JWTIssuerConfig) -> dict:
-        normalized: Final[dict] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS}
+    def _apply_issuer_claim_mappings(
+        self, token: Mapping[str, object], issuer_config: JWTIssuerConfig
+    ) -> dict[str, object]:
+        normalized: Final[dict[str, object]] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS}
         normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer
         claim_mappings: Final = [
             (issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM),
@@ -1192,8 +1213,8 @@ class JWTHandler:
 
         return normalized
 
-    def _get_jwk_from_public_key(self, public_key: dict) -> dict:
-        jwk: Final = {}
+    def _get_jwk_from_public_key(self, public_key: Mapping[str, object]) -> dict[str, object]:
+        jwk: Final[dict[str, object]] = {}
         for key in ["kty", "kid", "n", "e", "x", "y", "crv"]:
             if key in public_key:
                 jwk[key] = public_key[key]
@@ -1223,12 +1244,12 @@ class JWTHandler:
     def _decode_jwt_with_public_key(
         self,
         token: str,
-        public_key: dict | str,
+        public_key: JWTKeyItem | str,
         audience: str | list[str] | None,
         issuer: str | None = None,
         options: dict | None = None,
         disable_audience_validation: bool = False,
-    ) -> dict:
+    ) -> dict[str, object]:
         decode_options: Final = (
             options
             if options is not None
@@ -1298,11 +1319,64 @@ class JWTHandler:
             issuer_config=issuer_config,
         )
 
+    async def decode_trusted_assertion(
+        self, token: str, audience: Sequence[str], required_claims: Sequence[str]
+    ) -> Mapping[str, object] | AssertionRejection:
+        """Verify a JWT assertion handed to the gateway's token endpoint with the issuer plumbing
+        JWT auth already trusts, a matching ``issuers`` entry or ``JWT_ISSUER`` with
+        ``JWT_PUBLIC_KEY_URL``, but against the caller's audience instead of the API audience.
+        An assertion whose ``iss`` nothing trusts is refused rather than checked unscoped. A JWKS
+        outage propagates as the 503 ``ProxyException`` so the caller does not call the token bad."""
+        try:
+            header: Final = jwt.get_unverified_header(token)
+        except jwt.PyJWTError as malformed:
+            return AssertionRejection(reason=f"malformed assertion: {malformed}")
+        claims: Final = self.get_unverified_claims(token=token)
+        issuer_config: Final = self._get_configured_issuer(token=token)
+        trusted_issuer: Final = issuer_config.issuer if issuer_config is not None else os.getenv("JWT_ISSUER")
+        if not trusted_issuer or claims is None or claims.get("iss") != trusted_issuer:
+            return AssertionRejection(reason="iss is not an issuer this gateway trusts")
+        kid: Final = header.get("kid")
+        try:
+            public_key: Final = (
+                await self._get_public_key_from_jwks_url(self._get_jwks_url_for_issuer(issuer_config), kid)
+                if issuer_config is not None
+                else await self.get_public_key(kid=kid)
+            )
+        except JWKSUnreachableError as e:
+            raise jwks_unavailable_exception(e) from e
+        except Exception as no_key:  # noqa: BLE001  # the key lookups raise plain Exceptions for a missing key
+            return AssertionRejection(reason=f"no signing key for the assertion: {no_key}")
+        try:
+            payload: Final = self._decode_jwt_with_public_key(
+                token=token,
+                public_key=public_key,
+                audience=list(audience),
+                issuer=trusted_issuer,
+                options={"require": list(required_claims)},
+            )
+        except jwt.PyJWTError as invalid:
+            return AssertionRejection(reason=f"assertion failed verification: {invalid}")
+        verified: Final = (
+            payload
+            if issuer_config is None
+            else self._apply_issuer_claim_mappings(token=payload, issuer_config=issuer_config)
+        )
+        validate: Final = self.litellm_jwtauth.custom_validate
+        if validate is not None and not validate(dict(verified)):
+            return AssertionRejection(reason="custom_validate rejected the assertion's claims")
+        return verified
+
     async def auth_jwt(self, token: str) -> dict:
         header: Final = jwt.get_unverified_header(token)
 
         verbose_proxy_logger.debug("header: %s", header)
 
+        if header.get("typ") == ID_JAG_TYP:
+            raise IdJagIsNotABearerTokenError(
+                "an ID-JAG is an authorization grant, not a bearer token; redeem it at the token endpoint"
+            )
+
         kid: Final = header.get("kid", None)
 
         issuer_config: Final = self._get_configured_issuer(token=token)

litellm/proxy/auth/model_checks.py

@@ -332,7 +332,7 @@ def get_known_models_from_wildcard(wildcard_model: str, litellm_params: LiteLLM_
             else:
                 model = f"{wildcard_provider_prefix}/{model}"
         suffix_appended_wildcard_models.append(model)
-    return suffix_appended_wildcard_models or []
+    return list(dict.fromkeys(suffix_appended_wildcard_models))
 
 
 def expand_wildcard_deployments_for_model_info(

litellm/proxy/auth/route_checks.py

@@ -69,7 +69,8 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset(
 # paths directly because the request route carries the resolved key id.
 _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend")
 
-_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", "llm_api_routes"))
+_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp"))
+_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS))
 
 # Method-scoped because the generic credential CRUD handlers share these paths:
 # the route check cannot admit PATCH/DELETE on a name like "user_connections"
@@ -154,7 +155,7 @@ class RouteChecks:
                     ################################################
                     #  For llm_api_routes, also check registered pass-through endpoints
                     ################################################
-                    if allowed_route == "llm_api_routes":
+                    if allowed_route in _LLM_API_ROUTE_GROUPS:
                         if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET":
                             return True
 
@@ -185,7 +186,9 @@ class RouteChecks:
                         # headers, env, credentials). POST/PUT/DELETE on
                         # these paths are admin-only management writes and
                         # are intentionally not covered.
-                        if RouteChecks._is_get_mcp_server_discovery_route(route=route, request=request):
+                        if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route(
+                            route=route, request=request
+                        ):
                             return True
 
                         # Agent registry CRUD moved from llm_api_routes into
@@ -195,7 +198,7 @@ class RouteChecks:
                         # could reach these paths, so keep them reachable here;
                         # the handlers in agent_endpoints/endpoints.py still
                         # enforce proxy-admin on writes and scope reads by role.
-                        if RouteChecks.check_route_access(
+                        if allowed_route == "llm_api_routes" and RouteChecks.check_route_access(
                             route=route,
                             allowed_routes=LiteLLMRoutes.agent_management_routes.value,
                         ):

litellm/proxy/batches_endpoints/endpoints.py

@@ -17,6 +17,7 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
 from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param
+from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
@@ -171,6 +172,8 @@ def _raise_not_found_when_openai_fallback_unservable(
         return
     if data.get("api_key") or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY"):
         return
+    if resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=data) is not None:
+        return
     raise ProxyException(
         message=not_found_message,
         type="invalid_request_error",

litellm/proxy/common_request_processing.py

@@ -80,6 +80,7 @@ from litellm.litellm_core_utils.served_output_texts import (
 from litellm.litellm_core_utils.streaming_handler import (
     backfill_missing_cache_usage_fields,
 )
+from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields
 from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import (
     can_key_call_resolved_model,
@@ -129,6 +130,7 @@ from litellm.proxy.utils import (  # noqa: F401  # legacy module exports
 )
 from litellm.router import Router
 from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
+from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
 from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.router import RouterRateLimitError
@@ -633,6 +635,63 @@ def serialize_http_exception_detail(
     return str(detail), None
 
 
+_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True))
+_UNSEEN_DEPLOYMENT: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset())
+
+
+def _is_auto_router(params: Mapping[str, object]) -> bool:
+    model: Final = params.get("model")
+    return isinstance(model, str) and model.startswith(AUTO_ROUTER_MODEL_PREFIX)
+
+
+def _named_deployments(llm_router: Router, model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]:
+    by_id: Final = None if model in llm_router.model_names else llm_router.get_deployment(model_id=model)
+    if by_id is not None:
+        return (_REQUEST_FIELDS.validate_python(by_id.litellm_params.model_dump(exclude_none=True)),)
+    group: Final = llm_router.get_model_from_alias(model=model) or model
+    if group not in llm_router.model_names or llm_router.get_model_list_from_routing_groups(model_name=group):
+        return ()
+    deployments: Final = tuple(
+        _REQUEST_FIELDS.validate_python(deployment["litellm_params"])
+        for deployment in llm_router.get_all_deployments(model_name=group, team_id=team_id)
+    )
+    return () if any(_is_auto_router(params) for params in deployments) else deployments
+
+
+def _router_defaults(llm_router: Router | None) -> Mapping[str, object]:
+    configured: Final[object] = getattr(llm_router, "default_litellm_params", None)
+    if configured is None:
+        return {}
+    return {name: value for name, value in _REQUEST_FIELDS.validate_python(configured).items() if value is not None}
+
+
+def _counted_drop_policies(
+    data: Mapping[str, object], llm_router: Router | None, team_id: str | None
+) -> tuple[DropPolicy, ...]:
+    requested: Final = {**_router_defaults(llm_router), **data}
+    model: Final = data.get("model")
+    deployments: Final = (
+        _named_deployments(llm_router, model, team_id) if llm_router is not None and isinstance(model, str) else ()
+    )
+    if not deployments:
+        return (drop_policy(requested), _UNSEEN_DEPLOYMENT)
+    return tuple(drop_policy({**deployment, **requested}) for deployment in deployments)
+
+
+def resolve_counted_tool_fields(
+    data: Mapping[str, object], llm_router: Router | None, team_id: str | None
+) -> Mapping[str, object] | ProxyException:
+    resolved: Final = resolve_tool_fields(data, lambda: _counted_drop_policies(data, llm_router, team_id))
+    if isinstance(resolved, MalformedField):
+        return ProxyException(
+            message=resolved.message,
+            type="invalid_request_error",
+            param=resolved.field,
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+    return resolved
+
+
 def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException:
     raw_detail: Final = attribute_of(exc, "detail", str(exc))
     message, structured_fields = serialize_http_exception_detail(raw_detail)

litellm/proxy/db/db_transaction_queue/spend_update_queue.py

@@ -29,6 +29,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
     ) -> DBSpendUpdateTransactions:
         """Flush all updates from the queue and return all updates aggregated by entity type."""
         updates: Final = await self.flush_all_updates_from_in_memory_queue()
+        await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize())
         if len(updates) > 0:
             verbose_proxy_logger.info(
                 "Spend tracking - flushed %d spend update items from in-memory queue",
@@ -48,6 +49,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
                 "Spend update queue is full. Aggregating all entries in queue to concatenate entries."
             )
             await self.aggregate_queue_updates()
+        await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize())
 
     async def aggregate_queue_updates(self):
         """Concatenate all updates in the queue to reduce the size of in-memory queue"""
@@ -55,6 +57,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
         aggregated_updates: Final = self._get_aggregated_spend_update_queue_item(updates)
         for update in aggregated_updates:
             await self.update_queue.put(update)
+        await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize())
 
     def _get_aggregated_spend_update_queue_item(
         self, updates: list[SpendUpdateQueueItem]

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py

@@ -32,6 +32,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         additional_provider_specific_params=getattr(litellm_params, "additional_provider_specific_params", {}),
         unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"),
         fail_on_error=getattr(litellm_params, "fail_on_error", True),
+        forward_messages_as_sent=getattr(litellm_params, "forward_messages_as_sent", False),
         extra_headers=getattr(litellm_params, "extra_headers", None),
         guardrail_name=guardrail.get("guardrail_name", ""),
         event_hook=litellm_params.mode,

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -7,10 +7,12 @@
 
 import fnmatch
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from typing import TYPE_CHECKING, Final, Literal, Optional
 
 import httpx
+from pydantic import JsonValue, TypeAdapter, ValidationError
+from pydantic_core import PydanticSerializationError
 
 from litellm._logging import verbose_proxy_logger
 from litellm._version import version as litellm_version
@@ -38,6 +40,10 @@ if TYPE_CHECKING:
     from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 GUARDRAIL_NAME: Final = "generic_guardrail_api"
+_CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues)
+_SENT_ROW: Final = TypeAdapter(Mapping[str, object])
+_SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]])
+_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 
 # Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*).
 _HEADER_VALUE_ALLOWLIST: Final = frozenset(
@@ -150,23 +156,65 @@ def _extract_inbound_headers(
     return None
 
 
+def _is_contained_in(part: JsonValue, whole: JsonValue) -> bool:
+    pending: Final[list[tuple[JsonValue, JsonValue]]] = [(part, whole)]  # mutable-ok: worklist, recursion is banned
+    while pending:
+        match pending.pop():
+            case dict() as part_dict, dict() as whole_dict:
+                if not part_dict.keys() <= whole_dict.keys():
+                    return False
+                pending.extend((value, whole_dict[key]) for key, value in part_dict.items())
+            case list() as part_list, list() as whole_list:
+                if len(part_list) != len(whole_list):
+                    return False
+                pending.extend(zip(part_list, whole_list, strict=True))
+            case part_leaf, whole_leaf if part_leaf != whole_leaf:
+                return False
+            case _:
+                pass
+    return True
+
+
+def _is_shown_row(shown: AllMessageValues, returned: AllMessageValues) -> bool:
+    return returned == shown
+
+
+def _is_echo(original: AllMessageValues, returned: AllMessageValues) -> bool:
+    """A server that parses rows with its own message types and echoes them back drops the
+    keys those types do not declare. That is still an echo as long as nothing changed and
+    every key the chat completion types declare is kept."""
+    if returned == original:
+        return True
+    try:
+        sent: Final[JsonValue] = _JSON.validate_json(_SENT_ROW.dump_json(original))
+        declared: Final[JsonValue] = _JSON.validate_json(
+            _CHAT_ROW.dump_json(_CHAT_ROW.validate_python(original), warnings="error")
+        )
+        echoed: Final[JsonValue] = _JSON.validate_python(returned)
+    except (ValidationError, PydanticSerializationError):
+        return False
+    return _is_contained_in(echoed, sent) and _is_contained_in(declared, echoed)
+
+
 def _structured_rows_to_write_back(
     original_rows: Sequence[AllMessageValues] | None,
-    shown_rows: Sequence[AllMessageValues] | None,
+    compared_rows: Sequence[AllMessageValues] | None,
     returned_rows: Sequence[AllMessageValues],
+    is_echo: Callable[[AllMessageValues, AllMessageValues], bool],
 ) -> tuple[AllMessageValues, ...] | None:
-    """The request model drops row keys its message types do not declare, so a
-    row the server echoes back verbatim is restored to the original row object.
-    A server that echoes every row back unchanged has not rewritten anything
-    per row, so its answer is read from texts, as it was before rows could be
-    returned at all."""
-    if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
+    """A row the server echoes back is restored to the original row object. A server
+    that echoes every row back has not rewritten anything per row, so its answer is
+    read from texts, as it was before rows could be returned at all."""
+    if original_rows is None or compared_rows is None or len(returned_rows) != len(original_rows):
         return tuple(returned_rows)
-    if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
+    echoed: Final = tuple(
+        is_echo(compared, returned) for compared, returned in zip(compared_rows, returned_rows, strict=True)
+    )
+    if all(echoed):
         return None
     return tuple(
-        original if returned == shown else returned
-        for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
+        original if was_echoed else returned
+        for original, returned, was_echoed in zip(original_rows, returned_rows, echoed, strict=True)
     )
 
 
@@ -200,6 +248,7 @@ class GenericGuardrailAPI(CustomGuardrail):
         additional_provider_specific_params: Mapping[str, object] | None = None,
         unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
         fail_on_error: bool | None = True,
+        forward_messages_as_sent: bool | None = False,
         extra_headers: list | None = None,
         streaming_end_of_stream_only: bool | None = None,
         streaming_sampling_rate: int | None = None,
@@ -235,6 +284,7 @@ class GenericGuardrailAPI(CustomGuardrail):
         self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback
 
         self.fail_on_error: bool = True if fail_on_error is None else fail_on_error
+        self.forward_messages_as_sent: bool = bool(forward_messages_as_sent)
 
         # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook
         # via getattr(guardrail_to_apply, "streaming_*", default).
@@ -316,15 +366,17 @@ class GenericGuardrailAPI(CustomGuardrail):
         http_status_code: int | None = None,
     ) -> GenericGuardrailAPIInputs:
         status_suffix: Final = f" http_status_code={http_status_code}" if http_status_code else ""
+        error_type: Final = "HTTPStatusError" if isinstance(error, httpx.HTTPStatusError) else type(error).__name__
         verbose_proxy_logger.critical(
             "Generic Guardrail API error (fail-open). Proceeding without guardrail.%s "
-            "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s",
+            "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s error_type=%s",
             status_suffix,
             getattr(self, "guardrail_name", None),
             getattr(self, "api_base", None),
             input_type,
             getattr(logging_obj, "litellm_call_id", None) if logging_obj else None,
             getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None,
+            error_type,
             exc_info=error,
         )
         # Keep flow going - treat as action=NONE (no modifications)
@@ -339,10 +391,23 @@ class GenericGuardrailAPI(CustomGuardrail):
             headers.update(self.headers)
         return headers
 
+    def _request_body(
+        self,
+        guardrail_request: GenericGuardrailAPIRequest,
+        structured_messages: Sequence[AllMessageValues] | None,
+    ) -> Mapping[str, object]:
+        body: Final[Mapping[str, object]] = guardrail_request.model_dump(mode="json")
+        if not self.forward_messages_as_sent or structured_messages is None:
+            return body
+        return {
+            **body,
+            "structured_messages": _SENT_ROWS.dump_python(_SENT_ROWS.validate_python(structured_messages), mode="json"),
+        }
+
     def _build_guardrail_return_inputs(
         self,
         *,
-        texts: list,
+        texts: list[str],
         images: list[str] | None,
         tools: list[ChatCompletionToolParam] | None,
         structured_messages: Sequence[AllMessageValues] | None,
@@ -361,8 +426,13 @@ class GenericGuardrailAPI(CustomGuardrail):
             return_inputs["tools"] = guardrail_response.tools
         elif tools:
             return_inputs["tools"] = tools
+        compared_rows, is_echo = (
+            (structured_messages, _is_echo) if self.forward_messages_as_sent else (shown_messages, _is_shown_row)
+        )
         rows_to_write_back: Final = (
-            _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
+            _structured_rows_to_write_back(
+                structured_messages, compared_rows, guardrail_response.structured_messages, is_echo
+            )
             if guardrail_response.structured_messages
             else None
         )
@@ -465,7 +535,7 @@ class GenericGuardrailAPI(CustomGuardrail):
                 litellm_version=litellm_version,
                 images=images,
                 tools=([GuardrailToolParam.model_validate(t) for t in tools] if tools else None),
-                structured_messages=structured_messages,
+                structured_messages=None if self.forward_messages_as_sent else structured_messages,
                 tool_calls=tool_calls,
                 additional_provider_specific_params=additional_params,
                 input_type=input_type,
@@ -478,7 +548,7 @@ class GenericGuardrailAPI(CustomGuardrail):
             # Use mode="json" to ensure all iterables are converted to lists
             response: Final = await self.async_handler.post(
                 url=self.api_base,
-                json=guardrail_request.model_dump(mode="json"),
+                json=dict(self._request_body(guardrail_request, structured_messages)),
                 headers=headers,
                 timeout=self.timeout,
             )

litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py

@@ -5,6 +5,7 @@ from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import (
     OpenAIModerationGuardrail,
 )
 from litellm.types.guardrails import SupportedGuardrailIntegrations
+from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS
 
 if TYPE_CHECKING:
     from litellm.types.guardrails import Guardrail, LitellmParams
@@ -19,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
 
     openai_moderation_guardrail: Final = OpenAIModerationGuardrail(
         guardrail_name=guardrail_name,
+        workload_identity_params=litellm_params.model_dump(include=set(OPENAI_WIF_KWARGS_KEYS), exclude_none=True),
         **{
             **litellm_params.model_dump(exclude_none=True),
             "api_key": litellm_params.api_key,

litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py

@@ -3,6 +3,8 @@
 OpenAI Moderation Guardrail Integration for LiteLLM
 """
 
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Literal, Optional
 
 from fastapi import HTTPException
@@ -17,7 +19,14 @@ from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.llms.openai.workload_identity import (
+    OpenAIWorkloadIdentityConfig,
+    resolve_openai_bearer_token_async,
+    resolve_openai_workload_identity_config,
+)
 from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id
+from litellm.secret_managers.main import get_secret_str
 from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
 from litellm.types.utils import (
     GenericGuardrailAPIInputs,
@@ -32,6 +41,21 @@ if TYPE_CHECKING:
     from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 
+def _config_secret(value: object) -> object:
+    return get_secret_str(value) if isinstance(value, str) and value.startswith("os.environ/") else value
+
+
+def _guardrail_workload_identity(
+    api_base: str, workload_identity_params: Mapping[str, object]
+) -> OpenAIWorkloadIdentityConfig | None:
+    try:
+        return resolve_openai_workload_identity_config(
+            api_key=None, api_base=api_base, litellm_params=workload_identity_params
+        )
+    except OpenAIError as incomplete_identity:
+        raise ValueError(f"OpenAI Moderation: {incomplete_identity.message}") from incomplete_identity
+
+
 class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
     """
     LiteLLM Built-in Guardrail for OpenAI Content Moderation.
@@ -55,6 +79,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
         model: Literal["omni-moderation-latest", "text-moderation-latest"] | None = None,
         streaming_end_of_stream_only: bool | None = None,
         streaming_sampling_rate: int | None = None,
+        workload_identity_params: Mapping[str, object] | None = None,
         **kwargs,
     ):
         """Initialize OpenAI Moderation guardrail handler."""
@@ -78,7 +103,11 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
         )
         self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate
 
-        if not self.api_key:
+        self.workload_identity_params: Mapping[str, object] = MappingProxyType(
+            {key: _config_secret(value) for key, value in (workload_identity_params or {}).items()}
+        )
+
+        if not self.api_key and _guardrail_workload_identity(self.api_base, self.workload_identity_params) is None:
             raise ValueError(
                 "OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration."
             )
@@ -109,10 +138,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
 
         verbose_proxy_logger.debug("OpenAI Moderation guard request: %s", request_body)
 
+        bearer_token: Final = await resolve_openai_bearer_token_async(
+            static_api_key=self.api_key, api_base=self.api_base, litellm_params=self.workload_identity_params
+        )
         response: Final = await self.async_handler.post(
             url=f"{self.api_base}/moderations",
             headers={
-                "Authorization": f"Bearer {self.api_key}",
+                "Authorization": f"Bearer {bearer_token}",
                 "Content-Type": "application/json",
             },
             json=request_body,

litellm/proxy/guardrails/tool_name_extraction.py

@@ -6,23 +6,23 @@ knowledge lives in one place. Uses guardrail translation handlers where availabl
 with standalone extractors for generate_content and MCP.
 """
 
+from collections.abc import Mapping
 from typing import Any, Final
 
 from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
 from litellm.llms import load_guardrail_translation_mappings
+from litellm.llms.base_llm.guardrail_translation.utils import list_field
 from litellm.types.utils import CallTypes
 
 # Call types that have no guardrail translation handler; we use standalone extractors
 STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {}
 
 
-def _extract_generate_content_tool_names(data: dict) -> list[str]:
+def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str]:
     """Google generateContent: tools[].functionDeclarations[].name"""
     names: Final[list[str]] = []
-    for tool in data.get("tools") or []:
-        if not isinstance(tool, dict):
-            continue
-        for decl in tool.get("functionDeclarations") or []:
+    for tool in list_field(data, "tools"):
+        for decl in list_field(tool, "functionDeclarations"):
             if isinstance(decl, dict) and decl.get("name"):
                 names.append(str(decl["name"]))
     return names

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -425,7 +425,7 @@ async def add_new_user_to_default_team(
     user_id: str,
     user_email: str | None,
     user_api_key_dict: UserAPIKeyAuth,
-    teams: list[str] | list[NewUserRequestTeam],
+    teams: Sequence[str | NewUserRequestTeam],
     prisma_client: "PrismaClient",
 ):
     tasks: Final[list[Awaitable[object]]] = []

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -28,7 +28,7 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError,
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
-from litellm.constants import LITELLM_PROXY_ADMIN_NAME
+from litellm.constants import LITELLM_PROXY_ADMIN_NAME, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
 from litellm.litellm_core_utils.ptu_pricing import (
@@ -1881,7 +1881,7 @@ async def delete_team_models(
     """
     deleted_model_ids: Final[list[str]] = []
     transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db)
-    async with transaction_client.tx() as tx:
+    async with transaction_client.tx(timeout=datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
         for team_id in team_ids:
             rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
             model_ids = [row.model_id for row in rows]

litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py

@@ -1,25 +1,26 @@
 """Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion.
 
 An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693
-subject token. Only the generic OIDC login path reaches a token response the gateway retains one
-from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion
-and every store-sourced ID-JAG exchange fails for every user, however many times they sign in.
-Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows
-nothing about MCP. This module is the one shared answer both warn from.
+subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the
+gateway retains one from, and the SAML login retains the signed assertion itself as a ``saml2``
+subject token. A deployment with no SSO never signs anyone in, so every store-sourced ID-JAG
+exchange fails for every user. Neither side can see that alone: the MCP registration knows nothing
+about SSO and the login knows nothing about MCP. This module is the one shared answer both warn from.
 """
 
 from __future__ import annotations
 
 import os
 from enum import Enum
+from typing import Final
 
 from typing_extensions import assert_never
 
 from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler
 
-_GENERIC_OIDC_REMEDY = (
-    "Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token "
-    "response the gateway retains an id_token from"
+_OIDC_REMEDY: Final = (
+    "Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID) or at SAML, "
+    "the login paths the gateway retains an identity assertion from"
 )
 
 
@@ -55,18 +56,14 @@ def id_jag_assertion_capture_gap() -> str | None:
     or ``None`` when that provider does capture an assertion."""
     provider = active_sso_provider()
     match provider:
-        case ActiveSSOProvider.generic:
+        case (
+            ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic | ActiveSSOProvider.saml
+        ):
             return None
         case ActiveSSOProvider.none:
             return (
                 "no SSO provider is configured, so no IdP identity assertion is ever captured and "
-                f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}"
-            )
-        case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml:
-            return (
-                f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no "
-                "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter "
-                f"how often they sign in. {_GENERIC_OIDC_REMEDY}"
+                f"ID-JAG credential resolution fails for every user. {_OIDC_REMEDY}"
             )
         case _:
             assert_never(provider)

litellm/proxy/management_endpoints/sso/saml_sso.py

@@ -27,6 +27,7 @@ import hashlib
 import os
 import secrets
 import time
+from collections.abc import Mapping
 from typing import Final, cast
 from urllib.parse import parse_qsl
 
@@ -37,6 +38,10 @@ from pydantic import ValidationError
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
+    SSOIdentityAssertion,
+    assertion_from_saml_login,
+)
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.management_endpoints.sso_helper_utils import SSO_SESSIONS_TARGET
 from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
@@ -212,7 +217,7 @@ class SAMLAuthHandler:
         return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings)
 
     @staticmethod
-    def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]:
+    def _prepare_request_data(request: Request, post_data: Mapping[str, str] | None = None) -> dict[str, object]:
         base: Final = SAMLAuthHandler._base_url(request)
         scheme, _, host_part = base.partition("://")
         host: Final = host_part.split("/", 1)[0]
@@ -228,7 +233,7 @@ class SAMLAuthHandler:
     async def _build_auth(
         request: Request,
         cache: DualCache,
-        post_data: dict[str, str] | None = None,
+        post_data: Mapping[str, str] | None = None,
     ) -> "OneLogin_Saml2_Auth":
         if not SAML_AVAILABLE:
             raise _saml_unavailable_error()
@@ -310,7 +315,9 @@ class SAMLAuthHandler:
         return dict(parse_qsl(body.decode("utf-8", "replace")))
 
     @staticmethod
-    async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID:
+    async def handle_acs(
+        request: Request, cache: DualCache, post_data: Mapping[str, str]
+    ) -> tuple[CustomOpenID, SSOIdentityAssertion | None]:
         auth: Final = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data)
         browser_request_id: Final = request.cookies.get(_SAML_AUTHN_STATE_COOKIE)
         try:
@@ -330,7 +337,50 @@ class SAMLAuthHandler:
             )
 
         await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id)
-        return SAMLAuthHandler._result_from_auth(auth)
+        return SAMLAuthHandler._result_from_auth(auth), SAMLAuthHandler._retained_assertion(auth)
+
+    @staticmethod
+    def _retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None:
+        """The one signed ``<saml:Assertion>`` the validated response carried, serialized on its own
+        so its enveloped signature still verifies at the IdP's token endpoint. Retention is optional
+        egress material, so a failure here is logged and the validated login goes on without it."""
+        try:
+            return SAMLAuthHandler._parse_retained_assertion(auth)
+        except Exception as e:  # noqa: BLE001 - lxml and python3-saml expose no common base; capture must not block login
+            verbose_proxy_logger.warning(
+                "SAML login validated but its assertion could not be retained for oauth2_id_jag MCP servers: %s", e
+            )
+            return None
+
+    @staticmethod
+    def _parse_retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None:
+        xml: Final = cast(bytes | None, auth.get_last_response_xml())  # cast-ok: untyped python3-saml
+        if not xml:
+            verbose_proxy_logger.warning(
+                "SAML login validated but the response XML was not retained, so this user has no identity "
+                "assertion for oauth2_id_jag MCP servers."
+            )
+            return None
+        root: Final = OneLogin_Saml2_XML.to_etree(xml)
+        assertions: Final = OneLogin_Saml2_XML.query(root, "//saml:Assertion")
+        if len(assertions) != 1:
+            encrypted: Final = OneLogin_Saml2_XML.query(root, "//saml:EncryptedAssertion")
+            verbose_proxy_logger.warning(
+                "SAML login validated but the response carried %d plaintext saml:Assertion elements and %d "
+                "saml:EncryptedAssertion elements, so no identity assertion was retained for oauth2_id_jag MCP "
+                "servers. Those servers need exactly one signed, unencrypted assertion from the IdP.",
+                len(assertions),
+                len(encrypted),
+            )
+            return None
+        issuers: Final = OneLogin_Saml2_XML.query(assertions[0], "./saml:Issuer")
+        issuer: Final = cast(str | None, issuers[0].text) if issuers else None  # cast-ok: untyped python3-saml
+        not_on_or_after: Final = auth.get_last_assertion_not_on_or_after()
+        return assertion_from_saml_login(
+            assertion_xml=cast(bytes, OneLogin_Saml2_XML.to_string(assertions[0], encoding="utf-8")),  # cast-ok: saml
+            not_on_or_after=not_on_or_after if isinstance(not_on_or_after, int) else None,
+            issuer=issuer,
+        )
 
     @staticmethod
     def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int:

litellm/proxy/management_endpoints/team_endpoints.py

@@ -17,7 +17,7 @@ import traceback
 from collections.abc import Iterable, Mapping, Sequence
 from collections.abc import Set as AbstractSet
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from itertools import chain
 from types import MappingProxyType
 from typing import (
@@ -42,6 +42,7 @@ from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.caching.redis_batch import active_request_redis_batches
+from litellm.constants import TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
 from litellm.integrations.prometheus import PrometheusLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
@@ -4674,7 +4675,7 @@ async def delete_team(
     # or is still waiting on the lock, in which case its own re-read happens after this commits
     # and sees the row gone before it writes anything.
     delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
-    async with prisma_client.tx() as tx:
+    async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
         for team_id in sorted(data.team_ids):
             await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
         await tx.litellm_teamtable.delete_many(where=delete_filter)