import base64
import binascii
import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast

from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request
from litellm.proxy._types import (
    LiteLLM_MCPServerTable,
    MCPApprovalStatus,
    MCPEnvVar,
    MCPEnvVarScope,
    MCPSubmissionsSummary,
    NewMCPServerRequest,
    SpecialMCPServerName,
    UpdateMCPServerRequest,
    UserAPIKeyAuth,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
    _get_salt_key,
    decrypt_value_helper,
    encrypt_value_helper,
)
from litellm.proxy.utils import PrismaClient
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
    MCPServerOAuthClientRepository,
    MCPServerRepository,
    MCPUserCredentialsRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.verification_token_repository import (
    VerificationTokenRepository,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp import MCPCredentials

if TYPE_CHECKING:
    from prisma import models as prisma_db_models
    from prisma import types as prisma_db_types

    from litellm.types.mcp_server.mcp_server_manager import MCPServer


class _UserEnvVarsTransactionClient(Protocol):
    litellm_mcpuserenvvars: "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"

    async def execute_raw(self, query: str, *args: object) -> int: ...


class _UserEnvVarsTransaction(Protocol):
    async def __aenter__(self) -> _UserEnvVarsTransactionClient: ...

    async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...


_AUTH_FLOW_SCOPED_FIELDS: Final["frozenset[str]"] = frozenset(
    {
        "issuer",
        "authorization_url",
        "token_url",
        "registration_url",
        "oauth2_flow",
        "dcr_bridge",
        "token_exchange_endpoint",
        "audience",
        "subject_token_type",
        "token_exchange_profile",
    }
)


def _blank_to_none(value: str | None) -> str | None:
    if not isinstance(value, str):
        return None
    return value.strip() or None


# Token-exchange settings with dedicated columns that also exist on
# ``MCPCredentials`` as a legacy shape (rows and REST callers that predate the
# columns). Every write lifts blob values into the columns and strips them from
# the stored blob, so the read-time ``column or blob`` fallback only serves rows
# the current code has never written — a cleared column can then never be
# silently resurrected by a stale blob copy. These keys are stored plaintext
# (endpoints/identifiers, not secrets), so values lift as-is.
_TOKEN_EXCHANGE_COLUMN_FIELDS: Final["frozenset[str]"] = frozenset(
    {
        "token_exchange_endpoint",
        "audience",
        "subject_token_type",
        "token_exchange_profile",
    }
)

# The client-forwarded token modes share one stored-credential shape: the admin-declared upstream
# OAuth app (client_id/client_secret) plus the same authorize relay, and neither mints anything the
# gateway keeps. So a switch WITHIN this class must preserve the stored app, unlike a cross-class
# switch (e.g. an oauth2 row whose client may be DCR-minted and is not reusable elsewhere).
_CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthrough", "oauth_delegate"})

# Minted token material that must never survive a client rotation on a persisted row.
_MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"})


class _OAuthCredentialAccessToken(TypedDict):
    access_token: str


class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
    type: str
    refresh_token: str
    expires_at: str
    connected_at: str
    scopes: list[str]
    server_id: str


class _OAuthTokenRefreshResponse(TypedDict, total=False):
    access_token: str
    refresh_token: str
    expires_in: int
    scope: str


def _credential_auth_class(auth_type: str | None) -> str | None:
    """Collapse the client-forwarded modes to one credential class; every other auth_type is its own
    class. Used so credential handling keys off whether the stored-credential shape actually changed,
    not off a raw auth_type inequality that treats true_passthrough<->oauth_delegate as a full reset."""
    if auth_type in _CLIENT_FORWARDED_AUTH_TYPES:
        return "client_forwarded"
    return auth_type


def _drop_stale_minted_on_client_rotation(merged: dict[str, object], new_creds: dict[str, object]) -> dict[str, object]:
    """When the update rotates the client, drop stale minted token keys it did not itself set, so an old
    app's access/refresh token never rides forward under the new client. A no-op when no client key changed."""
    if "client_id" not in new_creds and "client_secret" not in new_creds:
        return merged
    return {
        key: value for key, value in merged.items() if key not in _MINTED_TOKEN_CREDENTIAL_FIELDS or key in new_creds
    }


def _is_global_env_var_scope(scope: object) -> bool:
    """``scope="user"`` entries are placeholders the user fills in; everything
    else (including a missing scope) is an admin-supplied global value."""
    return scope != MCPEnvVarScope.user and scope != "user"


def _encrypt_global_env_var_values(env_vars: Iterable[dict[str, str]]) -> None:
    """Encrypt ``scope="global"`` env var values in place before persisting.

    Global values hold admin-supplied secrets (API keys, passwords) that get
    interpolated into headers, so they are encrypted at rest like credentials
    and the per-user ``values_b64`` column. Per-user placeholders are not
    secrets and are stored verbatim.
    """
    for entry in env_vars:
        if not _is_global_env_var_scope(entry.get("scope")):
            continue
        value = entry.get("value")
        if value:
            entry["value"] = encrypt_value_helper(value)


def decrypt_global_env_var_values(env_vars: Iterable[MCPEnvVar | dict[str, str]] | None) -> None:
    """Decrypt ``scope="global"`` env var values in place after reading the DB.

    Accepts ``MCPEnvVar`` models (``LiteLLM_MCPServerTable``) or plain dicts
    (raw rows / deserialized JSON). Global values are always stored encrypted,
    so a value that no longer decrypts (e.g. after a salt-key change) is dropped
    and a warning is logged rather than forwarding the ciphertext into upstream
    ``${NAME}`` headers, where it would silently fail.
    """
    if not env_vars:
        return
    for entry in env_vars:
        is_dict = isinstance(entry, dict)
        scope = entry.get("scope") if is_dict else getattr(entry, "scope", None)
        if not _is_global_env_var_scope(scope):
            continue
        value = entry.get("value") if is_dict else getattr(entry, "value", None)
        if not value:
            continue
        decrypted = decrypt_value_helper(
            value=value,
            key="mcp_global_env_var",
            exception_type="debug",
            return_original_value=False,
        )
        if decrypted is None:
            name = entry.get("name") if is_dict else getattr(entry, "name", None)
            verbose_proxy_logger.warning(
                "MCP global env var %s failed to decrypt (LITELLM_SALT_KEY "
                "changed?); dropping it so ciphertext is not sent upstream",
                name,
            )
            decrypted = ""
        if is_dict:
            entry["value"] = decrypted
        else:
            entry.value = decrypted


def _decrypt_env_vars_on_returned_row(row: object) -> None:
    """Decrypt ``scope="global"`` env var values on a row returned by Prisma create/update.

    Prisma may hand back ``env_vars`` either as a parsed list (the common case for
    JSONB columns) or as a raw JSON string (observed for some write paths). The
    in-place decrypt helper only mutates iterables of dicts/models, so a string
    payload would silently skip decryption and ciphertext would leak into the
    registry via ``add_server``/``update_server`` (which trust the caller).
    Parse the string back to a list so the in-place decrypt actually runs, and
    write the decrypted list back onto the row so downstream consumers see plain
    values.
    """
    env_vars = getattr(row, "env_vars", None)
    if env_vars is None:
        return
    if isinstance(env_vars, str):
        try:
            env_vars = json.loads(env_vars)
        except (json.JSONDecodeError, TypeError):
            return
        if not isinstance(env_vars, list):
            return
        try:
            setattr(row, "env_vars", env_vars)
        except (AttributeError, TypeError):
            pass
    decrypt_global_env_var_values(env_vars)


def _reencrypt_global_env_var_values(
    env_vars: str | Iterable[Mapping[str, str]] | None, new_encryption_key: str
) -> list[dict[str, str]] | None:
    """Re-encrypt ``scope="global"`` env var values for master-key rotation.

    Each global value is decrypted with the current salt key and re-encrypted
    under ``new_encryption_key``. Returns the rebuilt list when at least one
    value was rotated, else ``None`` so the caller can skip the DB write. A
    value that fails to decrypt is left untouched (and logged) so a corrupt
    entry is preserved for recovery rather than overwritten.
    """
    if not env_vars:
        return None
    entries: Iterable[Mapping[str, str]]
    if isinstance(env_vars, str):
        try:
            entries = json.loads(env_vars)
        except (json.JSONDecodeError, TypeError):
            return None
        if not entries:
            return None
    else:
        entries = env_vars
    rebuilt: Final = [dict(v) for v in entries]
    rotated = False
    for entry in rebuilt:
        if not _is_global_env_var_scope(entry.get("scope")):
            continue
        value = entry.get("value")
        if not value:
            continue
        decrypted = decrypt_value_helper(
            value=value,
            key="mcp_global_env_var",
            exception_type="debug",
            return_original_value=False,
        )
        if decrypted is None:
            verbose_proxy_logger.warning(
                "rotate_mcp_server_credentials_master_key: could not decrypt global env var %s, skipping",
                entry.get("name"),
            )
            continue
        entry["value"] = encrypt_value_helper(decrypted, new_encryption_key=new_encryption_key)
        rotated = True
    return rebuilt if rotated else None


def _prepare_mcp_server_data(
    data: NewMCPServerRequest | UpdateMCPServerRequest,
    exclude_unset: bool = False,
    fields_set: set[str] | None = None,
) -> dict[str, Any]:
    """
    Helper function to prepare MCP server data for database operations.
    Handles JSON field serialization for mcp_info and env fields.

    Args:
        data: NewMCPServerRequest or UpdateMCPServerRequest object
        exclude_unset: When True, only fields the caller explicitly provided are
            included. Used for partial updates (PUT /v1/mcp/server) so omitted
            fields keep their existing DB value instead of being silently reset
            to a Pydantic schema default. ``exclude_none`` is not enough here:
            non-Optional fields (e.g. ``transport=MCPTransport.sse``,
            ``mcp_access_groups=[]``, ``allow_all_keys=False``) are backfilled
            with their default when omitted, and a non-None default survives the
            ``exclude_none`` filter and overwrites the row.

    Returns:
        Dict with properly serialized JSON fields
    """
    from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

    # Convert model to dict.
    # - Partial update (exclude_unset): only caller-provided keys are emitted, so
    #   omitted fields are never written and keep their existing DB value.
    # - Create (exclude_none): drop None-valued fields and let DB defaults apply.
    if exclude_unset:
        if fields_set is None:
            fields_set = data.fields_set()
        data_dict = data.model_dump(exclude_unset=True)
        # ``validate_and_normalize_mcp_server_payload`` always assigns ``alias``
        # on the payload, which marks it as set even when the caller omitted it.
        # Drop it only when the original request omitted alias; an explicit
        # ``alias=None`` is a valid request to clear the stored alias.
        if data_dict.get("alias") is None and "alias" not in fields_set:
            data_dict.pop("alias", None)
        # Prisma ``allowed_tools`` is a required String[]; ``null`` is invalid.
        # The UI sends null to clear a whitelist — treat that as ``[]``.
        if "allowed_tools" in data_dict and data_dict["allowed_tools"] is None:
            data_dict["allowed_tools"] = []
        # Json map fields use ``@default("{}")``; explicit null means clear overrides.
        for json_map_field in (
            "tool_name_to_display_name",
            "tool_name_to_description",
        ):
            if json_map_field in data_dict and data_dict[json_map_field] is None:
                data_dict[json_map_field] = {}
    else:
        data_dict = data.model_dump(exclude_none=True)
        # Ensure alias is always present in the dict (even if None)
        if "alias" not in data_dict:
            data_dict["alias"] = getattr(data, "alias", None)

    # Handle credentials serialization
    credentials: Final = data_dict.get("credentials")
    if credentials is not None:
        # Lift legacy blob-shaped token-exchange settings into their dedicated
        # columns (an explicit top-level value wins, including an explicit
        # null) and strip them from the blob so it never seeds the read-time
        # fallback for rows written by current code.
        for te_field in _TOKEN_EXCHANGE_COLUMN_FIELDS:
            blob_value = credentials.pop(te_field, None)
            if blob_value is not None and te_field not in data_dict:
                data_dict[te_field] = blob_value
        data_dict["credentials"] = encrypt_credentials(credentials=credentials, encryption_key=_get_salt_key())
        data_dict["credentials"] = safe_dumps(data_dict["credentials"])

    # Serialize JSON fields from ``data_dict`` (not ``data``) so the
    # exclude_unset filter is respected. Reading back from ``data`` would
    # reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
    if data_dict.get("static_headers") is not None:
        data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])

    # env_vars is read from ``data_dict`` (not ``data``) like every other JSON
    # column so the exclude_unset filter is respected: a partial update that
    # omits env_vars never overwrites the stored value. Global values are
    # encrypted at rest before serialization.
    env_vars: Final[Sequence[Mapping[str, str]] | None] = data_dict.get("env_vars")
    if env_vars is not None:
        serialized_env_vars: Final = [dict(v) for v in env_vars]
        _encrypt_global_env_var_values(serialized_env_vars)
        data_dict["env_vars"] = safe_dumps(serialized_env_vars)

    if data_dict.get("mcp_info") is not None:
        data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])

    if data_dict.get("env") is not None:
        data_dict["env"] = safe_dumps(data_dict["env"])

    if "tool_name_to_display_name" in data_dict:
        data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {})
    if "tool_name_to_description" in data_dict:
        data_dict["tool_name_to_description"] = safe_dumps(data_dict["tool_name_to_description"] or {})

    # mcp_access_groups is already List[str], no serialization needed

    # On create, force is_byok so a False value is always written to the DB. On
    # partial update, only write it when the caller explicitly provided it.
    if not exclude_unset:
        data_dict["is_byok"] = getattr(data, "is_byok", False)

    return data_dict


def encrypt_credentials(credentials: MCPCredentials, encryption_key: str | None) -> MCPCredentials:
    auth_value: Final = credentials.get("auth_value")
    if auth_value is not None:
        credentials["auth_value"] = encrypt_value_helper(
            value=auth_value,
            new_encryption_key=encryption_key,
        )
    client_id: Final = credentials.get("client_id")
    if client_id is not None:
        credentials["client_id"] = encrypt_value_helper(
            value=client_id,
            new_encryption_key=encryption_key,
        )
    client_secret: Final = credentials.get("client_secret")
    if client_secret is not None:
        credentials["client_secret"] = encrypt_value_helper(
            value=client_secret,
            new_encryption_key=encryption_key,
        )
    client_private_key: Final = credentials.get("client_private_key")
    if client_private_key is not None:
        credentials["client_private_key"] = encrypt_value_helper(
            value=client_private_key,
            new_encryption_key=encryption_key,
        )
    # AWS SigV4 credential fields
    aws_access_key_id: Final = credentials.get("aws_access_key_id")
    if aws_access_key_id is not None:
        credentials["aws_access_key_id"] = encrypt_value_helper(
            value=aws_access_key_id,
            new_encryption_key=encryption_key,
        )
    aws_secret_access_key: Final = credentials.get("aws_secret_access_key")
    if aws_secret_access_key is not None:
        credentials["aws_secret_access_key"] = encrypt_value_helper(
            value=aws_secret_access_key,
            new_encryption_key=encryption_key,
        )
    aws_session_token: Final = credentials.get("aws_session_token")
    if aws_session_token is not None:
        credentials["aws_session_token"] = encrypt_value_helper(
            value=aws_session_token,
            new_encryption_key=encryption_key,
        )
    # aws_region_name and aws_service_name are NOT secrets — stored as-is
    return credentials


def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[str, object]:
    parsed_blob: Final[dict[str, object]] = json.loads(blob) if isinstance(blob, str) else dict(blob)
    return parsed_blob


def _mcp_server_table_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
    table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
    return table


def _verification_token_table_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_VerificationToken]":
    table: Final[TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
        prisma_client
    ).table
    return table


def _team_table_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_TeamTable]":
    table: Final[TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
    return table


def _oauth_client_table_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
    table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
        prisma_client
    ).table
    return table


def _db_transaction_manager(prisma_client: PrismaClient) -> _UserEnvVarsTransaction:
    manager: Final[_UserEnvVarsTransaction] = prisma_client.db.tx()
    return manager


async def _db_find_mcp_server_rows(
    prisma_client: PrismaClient,
    where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]":
    return await _mcp_server_table_actions(prisma_client).find_many(where=where)


async def _db_find_mcp_server_row(
    prisma_client: PrismaClient, server_id: str
) -> "prisma_db_models.LiteLLM_MCPServerTable | None":
    return await _mcp_server_table_actions(prisma_client).find_unique(where={"server_id": server_id})


async def _db_update_mcp_server_row(
    prisma_client: PrismaClient,
    server_id: str,
    data: "prisma_db_types.LiteLLM_MCPServerTableUpdateInput",
) -> "prisma_db_models.LiteLLM_MCPServerTable":
    row: Final[prisma_db_models.LiteLLM_MCPServerTable | None] = await _mcp_server_table_actions(prisma_client).update(
        where={"server_id": server_id},
        data=data,
    )
    if row is None:
        raise ValueError(f"MCP server not found, passed server_id={server_id}")
    return row


def _user_credential_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
    table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
        prisma_client
    ).table
    return table


def _user_env_var_actions(
    prisma_client: PrismaClient,
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
    table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
    return table


async def _db_find_user_credential_row(
    prisma_client: PrismaClient, user_id: str, server_id: str
) -> "prisma_db_models.LiteLLM_MCPUserCredentials | None":
    return await _user_credential_actions(prisma_client).find_unique(
        where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
    )


async def _db_find_user_credential_rows(
    prisma_client: PrismaClient,
    where: "prisma_db_types.LiteLLM_MCPUserCredentialsWhereInput | None" = None,
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]":
    return await _user_credential_actions(prisma_client).find_many(where=where)


async def _db_upsert_user_credential_row(
    prisma_client: PrismaClient, user_id: str, server_id: str, credential_b64: str
) -> None:
    await _user_credential_actions(prisma_client).upsert(
        where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
        data={
            "create": {
                "user_id": user_id,
                "server_id": server_id,
                "credential_b64": credential_b64,
            },
            "update": {"credential_b64": credential_b64},
        },
    )


async def _db_find_user_env_var_rows(
    prisma_client: PrismaClient,
    where: "prisma_db_types.LiteLLM_MCPUserEnvVarsWhereInput | None" = None,
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]":
    return await _user_env_var_actions(prisma_client).find_many(where=where)


def decrypt_credentials(
    credentials: MCPCredentials,
) -> MCPCredentials:
    """Decrypt all secret fields in an MCPCredentials dict using the global salt key."""
    secret_fields: Final = [
        "auth_value",
        "client_id",
        "client_secret",
        "client_private_key",
        "aws_access_key_id",
        "aws_secret_access_key",
        "aws_session_token",
    ]
    for field in secret_fields:
        value = credentials.get(field)
        if value is not None and isinstance(value, str):
            credentials[field] = decrypt_value_helper(
                value=value,
                key=field,
                exception_type="debug",
                return_original_value=True,
            )
    return credentials


async def get_all_mcp_servers(
    prisma_client: PrismaClient,
    approval_status: str | None = None,
) -> list[LiteLLM_MCPServerTable]:
    """
    Returns mcp servers from the db, optionally filtered by approval_status.
    Pass approval_status=None to return every server except drafts, which back the admin OAuth
    session flow, are addressable only by their own server_id, and must never appear in a listing.
    NULL approval_status predates the approval workflow, so those rows are kept explicitly rather
    than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them.
    """
    where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = (
        {"approval_status": approval_status}
        if approval_status is not None
        # mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop
        # NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts
        else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]}
    )
    mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)

    tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
    for table in tables:
        decrypt_global_env_var_values(table.env_vars)
    return tables


async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
    """
    Returns the matching mcp server from the db iff exists
    """
    mcp_server: Final = await _db_find_mcp_server_row(prisma_client, server_id)
    if mcp_server is None:
        return None
    table: Final = LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump())
    decrypt_global_env_var_values(table.env_vars)
    return table


async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]) -> list[LiteLLM_MCPServerTable]:
    """
    Returns the matching mcp servers from the db with the server_ids
    """
    _mcp_servers: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
        prisma_client
    ).find_many(
        where={
            "server_id": {"in": server_ids},
        }
    )
    final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
    for _mcp_server in _mcp_servers:
        table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump())
        decrypt_global_env_var_values(table.env_vars)
        final_mcp_servers.append(table)

    return final_mcp_servers


async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]:
    """
    Returns the mcp servers from the db for the verification token
    """
    verification_token_record: (
        prisma_db_models.LiteLLM_VerificationToken | None
    ) = await _verification_token_table_actions(prisma_client).find_unique(
        where={
            "token": token,
        },
        include={
            "object_permission": True,
        },
    )

    mcp_servers: list[str] | None = []
    if verification_token_record is not None and verification_token_record.object_permission is not None:
        mcp_servers = verification_token_record.object_permission.mcp_servers
    return mcp_servers or []


async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> list[str]:
    """
    Returns the mcp servers from the db for the team id
    """
    team_record: prisma_db_models.LiteLLM_TeamTable | None = await _team_table_actions(prisma_client).find_unique(
        where={
            "team_id": team_id,
        },
        include={
            "object_permission": True,
        },
    )

    mcp_servers: list[str] | None = []
    if team_record is not None and team_record.object_permission is not None:
        mcp_servers = team_record.object_permission.mcp_servers
    return mcp_servers or []


async def get_all_mcp_servers_for_user(
    prisma_client: PrismaClient,
    user: UserAPIKeyAuth,
) -> list[LiteLLM_MCPServerTable]:
    """
    Get all the mcp servers filtered by the given user has access to.

    Following Least-Privilege Principle - the requestor should only be able to see the mcp servers that they have access to.
    """

    mcp_server_ids: Final[set[str]] = set()
    mcp_servers = []

    # Get the mcp servers for the key
    if user.api_key:
        token_mcp_servers: Final = await get_mcp_servers_by_verificationtoken(prisma_client, user.api_key)
        mcp_server_ids.update(token_mcp_servers)

        # check for special team membership
        if SpecialMCPServerName.all_team_servers in mcp_server_ids and user.team_id is not None:
            team_mcp_servers: Final = await get_mcp_servers_by_team(prisma_client, user.team_id)
            mcp_server_ids.update(team_mcp_servers)

    if len(mcp_server_ids) > 0:
        mcp_servers = await get_mcp_servers(prisma_client, mcp_server_ids)

    return mcp_servers


async def get_objectpermissions_for_mcp_server(
    prisma_client: PrismaClient, mcp_server_id: str
) -> "Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]":
    """
    Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server
    """
    object_permission_records: Final[
        Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]
    ] = await ObjectPermissionRepository(prisma_client).table.find_many(
        where={
            "mcp_servers": {"has": mcp_server_id},
        },
        include={
            "teams": True,
            "verification_tokens": True,
        },
    )

    return object_permission_records


async def get_virtualkeys_for_mcp_server(
    prisma_client: PrismaClient, server_id: str
) -> "Sequence[prisma_db_models.LiteLLM_VerificationToken]":
    """
    Get all the virtual keys that have access to the mcp server
    """
    virtual_keys: Final[
        Sequence[prisma_db_models.LiteLLM_VerificationToken] | None
    ] = await VerificationTokenRepository(prisma_client).table.find_many(
        where={
            "mcp_servers": {"has": server_id},
        },
    )

    if virtual_keys is None:  # pyright: ignore[reportUnnecessaryComparison]  # unreachable per seam types; kept as-is
        return []
    return virtual_keys


async def delete_mcp_server_from_team(prisma_client: PrismaClient, server_id: str):
    """
    Remove the mcp server from the team
    """


async def delete_mcp_server_from_virtualkey():
    """
    Remove the mcp server from the virtual key
    """


async def delete_mcp_server(
    prisma_client: PrismaClient,
    server_id: str,
    invalidate_token_cache: Callable[[str, str], Awaitable[None]] | None = None,
) -> LiteLLM_MCPServerTable | None:
    """
    Delete the mcp server from the db by server_id

    The server-row delete is the commit point. Per-user credential and env var
    rows have no FK cascade, so they are cleaned up afterwards on a best-effort
    basis: a transient failure there leaves only orphaned rows pointing at a
    now-missing server and must not turn a successful delete into a
    caller-visible error. Each table is cleaned independently so a failure on one
    still attempts the other.

    Each enumerated credential row's user also gets their cached per-user token
    invalidated (legacy cache + v2 store, via invalidate_token_cache, defaulting
    to the manager's shared invalidation): the caches are keyed by
    (user_id, server_id), so without this a re-created server reusing the same
    server_id would serve tokens minted for the deleted server until TTL.

    Returns the deleted mcp server record if it exists, otherwise None
    """
    deleted_server: Final = await MCPServerRepository(prisma_client).table.delete(
        where={
            "server_id": server_id,
        },
    )
    if deleted_server is not None:
        credential_user_ids: list[str] = []
        try:
            credential_rows: Sequence[prisma_db_models.LiteLLM_MCPUserCredentials] = await _user_credential_actions(
                prisma_client
            ).find_many(where={"server_id": server_id})
            credential_user_ids = [row.user_id for row in credential_rows]
        except Exception as e:  # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL
            verbose_proxy_logger.warning(
                "MCP server %s deleted but per-user credential enumeration failed; cached tokens expire by TTL: %s",
                server_id,
                e,
            )
        for model, label in (
            (_user_credential_actions(prisma_client), "credential"),
            (_user_env_var_actions(prisma_client), "env var"),
            (_oauth_client_table_actions(prisma_client), "OAuth client"),
        ):
            try:
                await model.delete_many(where={"server_id": server_id})
            except Exception as e:
                verbose_proxy_logger.warning(
                    "MCP server %s deleted but per-user %s cleanup failed; "
                    "orphaned rows can be removed on a later delete: %s",
                    server_id,
                    label,
                    e,
                )
        if credential_user_ids:
            if invalidate_token_cache is None:
                from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                    global_mcp_server_manager,
                )

                invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache
            for user_id in credential_user_ids:
                await invalidate_token_cache(user_id, server_id)
    return deleted_server  # pyright: ignore[reportReturnType]  # prisma row, not domain LiteLLM_MCPServerTable


async def create_mcp_server(
    prisma_client: PrismaClient, data: NewMCPServerRequest, touched_by: str
) -> LiteLLM_MCPServerTable:
    """
    Create a new mcp server record in the db
    """
    if data.server_id is None:
        data.server_id = str(uuid.uuid4())

    # Use helper to prepare data with proper JSON serialization
    data_dict: Final = _prepare_mcp_server_data(data)

    # Add audit fields
    data_dict["created_by"] = touched_by
    data_dict["updated_by"] = touched_by

    new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
        data=data_dict,  # pyright: ignore[reportAssignmentType]  # prisma row, not domain LiteLLM_MCPServerTable
    )

    _decrypt_env_vars_on_returned_row(new_mcp_server)
    return new_mcp_server


async def create_draft_mcp_server(
    prisma_client: PrismaClient,
    data: NewMCPServerRequest,
    touched_by: str,
    ttl_seconds: int,
    server_id: str | None = None,
) -> LiteLLM_MCPServerTable:
    """
    Persist a short-lived draft row backing the admin OAuth "Authorize & Fetch Token" flow.

    The draft lives in the database rather than in process memory so that the /register,
    /authorize and /token legs resolve it whichever worker or replica accepts each request.

    Writing is strictly create-if-absent. Any existing row for the id is returned untouched, which
    covers both a live draft for this same session and a real server the edit form is
    re-authorizing against its own id, where writing a draft would collide on the primary key.
    Each click of Authorize mints a fresh id, so nothing is lost by never overwriting, and it is
    what makes concurrent callers sharing one id safe rather than mutually destructive.
    """
    draft_id: Final = server_id or data.server_id or str(uuid.uuid4())
    await _prune_expired_draft_mcp_servers(prisma_client, ttl_seconds)

    existing: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
    if existing is not None:
        # Already usable by every worker, whether it is a live draft for this same session or a
        # real server the edit form is re-authorizing. Either way there is nothing to write, and
        # not writing is what keeps concurrent callers for one server_id from racing each other.
        return LiteLLM_MCPServerTable.model_validate(existing.model_dump())

    draft_payload: Final = data.model_copy(update={"server_id": draft_id, "approval_status": MCPApprovalStatus.draft})
    try:
        return await create_mcp_server(prisma_client, draft_payload, touched_by)
    except Exception:
        # Lost the create race: the read above and this create are two statements, not one. The
        # winner wrote a draft for this same session, so adopt it rather than failing a caller
        # whose session is in fact ready. Anything else still raises.
        raced: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
        if raced is None or raced.approval_status != MCPApprovalStatus.draft:
            raise
        return LiteLLM_MCPServerTable.model_validate(raced.model_dump())


async def _prune_expired_draft_mcp_servers(prisma_client: PrismaClient, ttl_seconds: int) -> None:
    """Drop drafts already past ``ttl_seconds``, so abandoned OAuth sessions do not accumulate.

    Runs on each draft write rather than on a schedule, mirroring the in-memory cache this
    replaces, which pruned on every store. Expired drafts are unreadable by then anyway, so the
    only thing at stake is row count, and the work is bounded by how often admins authorize.
    """
    cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
    # Age is filtered here rather than in the query: the draft set is bounded by how many OAuth
    # authorizations are in flight, so it is a handful of rows even on a busy proxy.
    drafts: Final = await _db_find_mcp_server_rows(
        prisma_client,
        where={"approval_status": MCPApprovalStatus.draft},
    )
    for row in drafts:
        # A row without a timestamp has no age to judge, so leave it rather than guess it is stale.
        # Two workers sweeping the same row is harmless: prisma's delete returns None for a row
        # that is already gone rather than raising, so the loser of that race is a no-op.
        if row.updated_at is not None and row.updated_at < cutoff:
            await delete_mcp_server(prisma_client, row.server_id)


async def get_draft_mcp_server(
    prisma_client: PrismaClient, server_id: str, ttl_seconds: int
) -> LiteLLM_MCPServerTable | None:
    """
    Return the draft row for ``server_id`` if it has not yet aged past ``ttl_seconds``, else None.

    Age is enforced in the query rather than by a sweeper so an expired draft is unreadable the
    moment it lapses, regardless of which process last ran a cleanup.
    """
    cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
    draft_rows: Final = await _db_find_mcp_server_rows(
        prisma_client,
        where={
            "server_id": server_id,
            "approval_status": MCPApprovalStatus.draft,
            "updated_at": {"gte": cutoff},
        },
    )
    if not draft_rows:
        return None

    table: Final = LiteLLM_MCPServerTable.model_validate(draft_rows[0].model_dump())
    decrypt_global_env_var_values(table.env_vars)
    return table


async def update_mcp_server(
    prisma_client: PrismaClient,
    data: UpdateMCPServerRequest,
    touched_by: str,
    fields_set: set[str] | None = None,
) -> LiteLLM_MCPServerTable | None:
    """
    Update a new mcp server record in the db
    """
    from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

    # Use helper to prepare data with proper JSON serialization.
    # exclude_unset=True makes this a true partial update: fields the caller did
    # not provide are not written, so they keep their existing DB value instead
    # of being reset to a schema default (transport=sse, allow_all_keys=False...).
    data_dict: Final = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set)

    # Pre-fetch existing record once if we need it for auth_type, url, or credential logic
    existing = None
    has_credentials: Final = "credentials" in data_dict and data_dict["credentials"] is not None
    # An explicit token-exchange column write (set or clear) also migrates the
    # legacy blob copies below, so the existing row is needed for those updates.
    explicit_te_write: Final = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys())
    url_provided: Final = "url" in data_dict and data_dict["url"] is not None
    issuer_provided: Final = "issuer" in data_dict
    if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided:
        existing = await _db_find_mcp_server_row(prisma_client, data.server_id)

    auth_type_changed: Final = bool(
        data.auth_type
        and existing
        and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type)
    )
    # A url change re-points the server at a potentially different upstream, so any discovered or
    # trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover.
    url_changed: Final = bool(url_provided and existing and existing.url != data_dict["url"])
    old_issuer: Final = _blank_to_none(getattr(existing, "issuer", None)) if existing else None
    issuer_changed: Final = bool(
        issuer_provided and old_issuer is not None and _blank_to_none(data_dict.get("issuer")) != old_issuer
    )

    # Clear stale credentials when auth_type changes but no new credentials provided
    if auth_type_changed and "credentials" not in data_dict:
        data_dict["credentials"] = None

    if auth_type_changed or url_changed or issuer_changed:
        # Clear each auth-flow-scoped field that the caller either omitted (partial update) or
        # resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint
        # belonging to the old upstream would otherwise survive a url/auth_type change and win in the
        # resolution merge; only a genuinely new submitted value is kept.
        data_dict.update(
            {
                field: None
                for field in _AUTH_FLOW_SCOPED_FIELDS
                if field not in data_dict or data_dict[field] == getattr(existing, field, None)
            }
        )

    # An explicit column write that does not touch credentials must still migrate
    # the row's legacy blob copies: lift values for columns the caller left
    # untouched, strip every copy from the blob. Without this, clearing a column
    # (e.g. to re-enable RFC 9728/8414 discovery) would leave the blob copy in
    # place, and the next credentials update's migrate-on-write would silently
    # repopulate the column the admin just cleared. (When credentials ARE in the
    # update, the merge below performs the same migration.)
    if explicit_te_write and "credentials" not in data_dict and existing is not None and existing.credentials:
        existing_creds = _credentials_blob_to_mutable_dict(existing.credentials)
        if _TOKEN_EXCHANGE_COLUMN_FIELDS & existing_creds.keys():
            for te_field in _TOKEN_EXCHANGE_COLUMN_FIELDS:
                legacy_value = existing_creds.pop(te_field, None)
                if legacy_value is not None and te_field not in data_dict and getattr(existing, te_field, None) is None:
                    data_dict[te_field] = legacy_value
            data_dict["credentials"] = safe_dumps(existing_creds)

    # Merge credentials: preserve existing fields not present in the update.
    # Without this, a partial credential update (e.g. changing only region)
    # would wipe encrypted secrets that the UI cannot display back.
    if "credentials" in data_dict and data_dict["credentials"] is not None:
        if existing and existing.credentials:
            # Only merge when the credential CLASS is unchanged. A cross-class switch
            # (e.g. oauth2 → api_key, or oauth2 → true_passthrough) replaces credentials
            # entirely to avoid stale secrets from the previous class lingering; a switch
            # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps
            # the same declared app and so must merge, not replace.
            if not auth_type_changed:
                existing_creds = _credentials_blob_to_mutable_dict(existing.credentials)
                new_creds: Final = _credentials_blob_to_mutable_dict(data_dict["credentials"])
                # New values override existing; existing keys not in update are preserved. A client
                # rotation additionally drops the previous app's stale minted token keys.
                merged: Final = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds)
                # Migrate-on-write for legacy rows: token-exchange settings the
                # old blob shape carried move to their dedicated columns (unless
                # the caller set the column this update, or the row already has
                # one) and are never re-persisted in the blob. Stored plaintext,
                # so the merged value lifts as-is.
                for te_field in _TOKEN_EXCHANGE_COLUMN_FIELDS:
                    legacy_value = merged.pop(te_field, None)
                    if (
                        legacy_value is not None
                        and te_field not in data_dict
                        and getattr(existing, te_field, None) is None
                    ):
                        data_dict[te_field] = legacy_value
                data_dict["credentials"] = safe_dumps(merged)

    # Add audit fields
    data_dict["updated_by"] = touched_by

    # prisma-python rejects a raw ``None`` for a ``Json?`` field ("value is required but not set"); the
    # clear paths above use ``None`` as the merge-skip sentinel, so translate it here to ``Json(None)``,
    # which writes SQL null and reads back as ``None``. Done at the edge so the merge guards stay simple.
    if "credentials" in data_dict and data_dict["credentials"] is None:
        from prisma import Json  # noqa: PLC0415  # local import: prisma may be ungenerated at module load in some tools

        data_dict["credentials"] = Json(None)

    updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
        where={"server_id": data.server_id},
        data=data_dict,  # pyright: ignore[reportAssignmentType]  # prisma row, not domain LiteLLM_MCPServerTable
    )

    _decrypt_env_vars_on_returned_row(updated_mcp_server)
    return updated_mcp_server


async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None:
    """Read the persisted (encrypted) DCR OAuth client blob for a server from the
    server-scoped store, or None. Config.yaml-declared servers have no
    LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed
    by server_id. The returned value is the raw credentials blob for
    ``_get_persisted_dcr_credentials`` to parse."""
    row: Final[prisma_db_models.LiteLLM_MCPServerOAuthClient | None] = await _oauth_client_table_actions(
        prisma_client
    ).find_unique(where={"server_id": server_id})
    if row is None:
        return None
    return row.credentials


async def upsert_mcp_server_oauth_client_credentials(
    prisma_client: PrismaClient, server_id: str, credentials: MCPCredentials
) -> None:
    """Persist a server's dynamically registered OAuth client (RFC 7591 DCR) in the
    server-scoped store keyed by server_id, independent of any LiteLLM_MCPServerTable row.
    client_id/client_secret are encrypted at rest with the same salt key used for the
    server row's credentials blob, so ``_apply_persisted_dcr_credentials`` decrypts them the
    same way regardless of which store a server's client came from."""
    from litellm.litellm_core_utils.safe_json_dumps import safe_dumps

    encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key())
    blob: Final = safe_dumps(encrypted)
    await _oauth_client_table_actions(prisma_client).upsert(
        where={"server_id": server_id},
        data={
            "create": {"server_id": server_id, "credentials": blob},
            "update": {"credentials": blob},
        },
    )


def _reencrypt_mcp_credentials_blob(
    credentials: "str | Mapping[str, object] | None", new_master_key: str
) -> str | None:
    """Decrypt an at-rest MCP credentials blob with the current key and re-encrypt it under
    new_master_key, returning the serialized blob or None when there is nothing to rotate. Shared by
    every table that stores an encrypted MCP credentials blob so a master-key rotation covers them
    uniformly and cannot silently skip one."""
    if not credentials:
        return None
    from litellm.litellm_core_utils.safe_json_dumps import safe_dumps  # noqa: PLC0415  # avoids circular import

    creds_dict: Final = _credentials_blob_to_mutable_dict(credentials)
    decrypted: Final = decrypt_credentials(credentials=cast(MCPCredentials, creds_dict))
    encrypted: Final = encrypt_credentials(credentials=decrypted, encryption_key=new_master_key)
    return safe_dumps(encrypted)


async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str):
    from litellm.litellm_core_utils.safe_json_dumps import safe_dumps  # noqa: PLC0415  # avoids circular import

    mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client)

    updated = 0
    for mcp_server in mcp_servers:
        update_data: dict[str, str] = {}

        rotated_credentials = _reencrypt_mcp_credentials_blob(mcp_server.credentials, new_master_key)
        if rotated_credentials is not None:
            update_data["credentials"] = rotated_credentials

        rotated_env_vars = _reencrypt_global_env_var_values(mcp_server.env_vars, new_master_key)
        if rotated_env_vars is not None:
            update_data["env_vars"] = safe_dumps(rotated_env_vars)

        if not update_data:
            continue

        update_data["updated_by"] = touched_by
        await _mcp_server_table_actions(prisma_client).update(
            where={"server_id": mcp_server.server_id},
            data=update_data,
        )
        updated += 1

    oauth_clients: Final[Sequence[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
        prisma_client
    ).find_many()
    oauth_updated = 0
    for oauth_client in oauth_clients:
        rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key)
        if rotated_credentials is None:
            continue
        await _oauth_client_table_actions(prisma_client).update(
            where={"server_id": oauth_client.server_id},
            data={"credentials": rotated_credentials},
        )
        oauth_updated += 1

    verbose_proxy_logger.info(
        "rotate_mcp_server_credentials_master_key: rotated %d MCP server row(s) and %d OAuth-client row(s)",
        updated,
        oauth_updated,
    )


def _decode_user_credential(stored: str) -> str | None:
    """Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``.

    Tries nacl decryption first (current write format).  Falls back to a
    plain ``urlsafe_b64decode`` for rows persisted by older code that wrote
    the credential without encryption.  Returns ``None`` when neither path
    yields a valid string.
    """
    decrypted: Final = decrypt_value_helper(
        value=stored,
        key="mcp_user_credential",
        exception_type="debug",
        return_original_value=False,
    )
    if decrypted is not None:
        return decrypted
    try:
        return base64.urlsafe_b64decode(stored).decode()
    except (binascii.Error, UnicodeDecodeError, ValueError, TypeError):
        return None


def _warn_undecryptable_credential(user_id: str, server_id: str) -> None:
    """Log the one credential state that otherwise reads as "user never authorized"."""
    verbose_proxy_logger.warning(
        "MCP user credential for user=%s server=%s could not be decrypted (likely written under a "
        "previous LITELLM_SALT_KEY); the user is treated as not connected and must re-authorize.",
        user_id,
        server_id,
    )


def _parse_oauth_payload(decoded: str | None) -> OAuthCredentialPayload | None:
    """Return the OAuth2 payload dict if ``decoded`` holds one, else ``None``.

    A row is considered an OAuth2 credential iff its decoded value parses as
    a JSON object with ``"type": "oauth2"``.  Plain BYOK credentials (which
    share the same column) decode to a non-JSON string and return ``None``.

    Callers that need to tell an unreadable row from a readable non-OAuth2 one
    pass the result of :func:`_decode_user_credential` so a single decode
    answers both questions: ``None`` there means the value can be neither
    decrypted nor base64-decoded, so no caller can ever recover it.
    """
    if decoded is None:
        return None
    parsed: OAuthCredentialPayload | None
    try:
        parsed = json.loads(decoded)
    except (ValueError, TypeError):
        return None
    if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
        return parsed
    return None


def _decode_oauth_payload(stored: str) -> OAuthCredentialPayload | None:
    """Return the OAuth2 payload dict held in ``stored``, else ``None``."""
    return _parse_oauth_payload(_decode_user_credential(stored))


async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, new_master_key: str):
    """Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``.

    Reads each ``credential_b64`` with the current salt key (falling back to
    legacy plain base64 for unmigrated rows) and writes it back encrypted
    under the new master key.  Rows that are unreadable under both paths
    are logged and skipped so one corrupt row does not abort the rotation.
    """
    rows: Final = await _db_find_user_credential_rows(prisma_client)
    rotated = 0
    skipped = 0
    for row in rows:
        plaintext = _decode_user_credential(row.credential_b64)
        if plaintext is None:
            verbose_proxy_logger.warning(
                "rotate_mcp_user_credentials_master_key: could not decode "
                "credential for user_id=%s server_id=%s, skipping",
                row.user_id,
                row.server_id,
            )
            skipped += 1
            continue
        re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key)
        await _user_credential_actions(prisma_client).update(
            where={
                "user_id_server_id": {
                    "user_id": row.user_id,
                    "server_id": row.server_id,
                }
            },
            data={"credential_b64": re_encrypted},
        )
        rotated += 1
    verbose_proxy_logger.info(
        "rotate_mcp_user_credentials_master_key: rotated %d row(s), skipped %d",
        rotated,
        skipped,
    )


async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_master_key: str):
    """Re-encrypt every ``LiteLLM_MCPUserEnvVars`` row with ``new_master_key``.

    Reads each ``values_b64`` blob with the current salt key and writes it back
    encrypted under the new master key. Rows that fail to decrypt are logged and
    skipped so one corrupt row does not abort the rotation nor overwrite values
    that may still be recoverable.
    """
    rows: Final = await _db_find_user_env_var_rows(prisma_client)
    rotated = 0
    skipped = 0
    for row in rows:
        plaintext = decrypt_value_helper(
            value=row.values_b64,
            key="mcp_user_env_vars",
            exception_type="debug",
            return_original_value=False,
        )
        if plaintext is None:
            verbose_proxy_logger.warning(
                "rotate_mcp_user_env_vars_master_key: could not decrypt env vars for user_id=%s server_id=%s, skipping",
                row.user_id,
                row.server_id,
            )
            skipped += 1
            continue
        re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key)
        await _user_env_var_actions(prisma_client).update(
            where={
                "user_id_server_id": {
                    "user_id": row.user_id,
                    "server_id": row.server_id,
                }
            },
            data={"values_b64": re_encrypted},
        )
        rotated += 1
    verbose_proxy_logger.info(
        "rotate_mcp_user_env_vars_master_key: rotated %d row(s), skipped %d",
        rotated,
        skipped,
    )


async def store_user_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
    credential: str,
) -> None:
    """Store a user credential for a BYOK MCP server."""

    encoded: Final = encrypt_value_helper(credential)
    await _db_upsert_user_credential_row(prisma_client, user_id, server_id, encoded)


async def get_user_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> str | None:
    """Return credential for a user+server pair, or None."""

    row: Final = await _db_find_user_credential_row(prisma_client, user_id, server_id)
    if row is None:
        return None
    return _decode_user_credential(row.credential_b64)


async def has_user_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> bool:
    """Return True if the user has a stored credential for this server."""
    row: Final = await _db_find_user_credential_row(prisma_client, user_id, server_id)
    return row is not None


async def delete_user_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> None:
    """Delete the user's stored credential for a BYOK MCP server."""
    await _user_credential_actions(prisma_client).delete(
        where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
    )


# ── OAuth2 user-credential helpers ────────────────────────────────────────────


async def store_user_oauth_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
    access_token: str,
    refresh_token: str | None = None,
    expires_in: int | None = None,
    scopes: list[str] | None = None,
    skip_byok_guard: bool = False,
) -> None:
    """Persist an OAuth2 access token for a user+server pair.

    The payload is JSON-serialised and stored encrypted in the same
    ``credential_b64`` column used by BYOK.  A ``"type": "oauth2"`` key
    differentiates it from plain BYOK API keys.
    """

    expires_at: str | None = None
    if expires_in is not None:
        expires_at = (datetime.now(timezone.utc) + timedelta(seconds=expires_in)).isoformat()

    payload: Final[OAuthCredentialPayload] = {
        "type": "oauth2",
        "access_token": access_token,
        "connected_at": datetime.now(timezone.utc).isoformat(),
    }
    if refresh_token:
        payload["refresh_token"] = refresh_token
    if expires_at:
        payload["expires_at"] = expires_at
    if scopes:
        payload["scopes"] = scopes

    # Guard against silently overwriting a BYOK credential with an OAuth token.
    # Skip the guard when the caller knows the row is already an OAuth2 credential
    # (e.g. during token refresh), saving an extra DB round-trip.
    if not skip_byok_guard:
        existing: Final = await _db_find_user_credential_row(prisma_client, user_id, server_id)
        decoded: Final = _decode_user_credential(existing.credential_b64) if existing is not None else None
        if existing is not None and _parse_oauth_payload(decoded) is None:
            # Refuse only while the row still holds readable content, which is a live BYOK
            # secret that overwriting would destroy.  A row that does not decode was written
            # under a different LITELLM_SALT_KEY, and one that decodes to nothing holds no
            # secret at all; refusing either preserves nothing and instead wedges the user
            # out of the OAuth flow for good, since re-authorizing is their only recovery.
            if decoded:
                raise ValueError(
                    f"Existing credential for user {user_id} and server "
                    f"{server_id} could not be verified as an OAuth2 token. "
                    f"Refusing to overwrite."
                )
            verbose_proxy_logger.warning(
                "store_user_oauth_credential: existing credential for user=%s server=%s could not be "
                "decrypted (likely written under a previous LITELLM_SALT_KEY); replacing it with the "
                "newly authorized OAuth2 token.",
                user_id,
                server_id,
            )

    encoded: Final = encrypt_value_helper(json.dumps(payload))
    await _db_upsert_user_credential_row(prisma_client, user_id, server_id, encoded)


def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: int = 0) -> bool:
    """Return True if the OAuth2 credential's access_token has expired.

    Checks the ``expires_at`` ISO-format string stored in the credential payload.
    Returns False when ``expires_at`` is absent or unparseable (treat as non-expired).
    With ``buffer_seconds`` > 0, a token that is still valid but expires within the
    buffer is also treated as expired, so callers can refresh proactively instead of
    handing back a token that may lapse mid-request.
    """
    expires_at: Final = cred.get("expires_at")
    if not expires_at:
        return False
    try:
        exp_dt = datetime.fromisoformat(expires_at)
        if exp_dt.tzinfo is None:
            exp_dt = exp_dt.replace(tzinfo=timezone.utc)
        return datetime.now(timezone.utc) + timedelta(seconds=buffer_seconds) > exp_dt
    except (ValueError, TypeError):
        return False


async def get_user_oauth_credential(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> OAuthCredentialPayload | None:
    """Return the decoded OAuth2 payload dict for a user+server pair, or None."""

    row: Final = await _db_find_user_credential_row(prisma_client, user_id, server_id)
    if row is None:
        return None
    decoded: Final = _decode_user_credential(row.credential_b64)
    if decoded is None:
        _warn_undecryptable_credential(user_id, server_id)
    return _parse_oauth_payload(decoded)


async def list_user_oauth_credentials(
    prisma_client: PrismaClient,
    user_id: str,
) -> list[OAuthCredentialPayload]:
    """Return all OAuth2 credential payloads for a user, tagged with server_id."""

    rows: Final = await _db_find_user_credential_rows(prisma_client, {"user_id": user_id})
    results: Final[list[OAuthCredentialPayload]] = []
    for row in rows:
        decoded = _decode_user_credential(row.credential_b64)
        if decoded is None:
            _warn_undecryptable_credential(user_id, row.server_id)
        payload = _parse_oauth_payload(decoded)
        if payload is None:
            continue
        payload["server_id"] = row.server_id
        results.append(payload)
    return results


def _decrypted_credential_field(creds: dict[str, object], field: str) -> object:
    """Return one credential field decrypted with the global salt key; non-string and legacy
    plaintext values come back unchanged (decrypt_value_helper returns the original on failure)."""
    value: Final = creds.get(field)
    if not isinstance(value, str):
        return value
    return decrypt_value_helper(
        value=value,
        key=field,
        exception_type="debug",
        return_original_value=True,
    )


def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
    """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url, or
    spec_path for OpenAPI servers, plus the RFC 8707 upstream_resource sent on the authorize and
    token legs), the OAuth mode/grant (auth_type, oauth2_flow), the authorization-server endpoints,
    and the OAuth client + scopes. Mirrors the dashboard's getOAuthAuthorizationIdentity. When any
    of these change on a server update, previously stored per-user tokens were minted for the old
    identity and are stale. Excludes transport and delegate_auth_to_upstream, which do not affect
    what token is minted (RFC 8693).

    client_id/client_secret are compared decrypted: stored values are NaCl-encrypted with a fresh
    nonce on every write, so comparing ciphertext would flag every routine save as an identity
    change and purge tokens that are still valid."""
    creds: Final = getattr(server, "credentials", None)
    if isinstance(creds, str):
        try:
            parsed: dict[str, object] | None = json.loads(creds)
        except ValueError:
            parsed = None
    else:
        parsed = creds
    creds_dict: Final[dict[str, object]] = parsed if isinstance(parsed, dict) else {}
    return (
        getattr(server, "url", None),
        getattr(server, "spec_path", None),
        getattr(server, "auth_type", None),
        getattr(server, "oauth2_flow", None),
        getattr(server, "issuer", None),
        getattr(server, "authorization_url", None),
        getattr(server, "token_url", None),
        getattr(server, "registration_url", None),
        _decrypted_credential_field(creds_dict, "client_id"),
        _decrypted_credential_field(creds_dict, "client_secret"),
        creds_dict.get("scopes"),
        creds_dict.get("upstream_resource"),
    )


async def purge_user_oauth_credentials_for_server(
    prisma_client: PrismaClient,
    server_id: str,
    invalidate_token_cache: Callable[[str, str], Awaitable[None]] | None = None,
) -> int:
    """Delete every stored per-user OAuth token for a server and invalidate each user's cached
    token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth
    token store), so no user keeps a token minted for a superseded configuration. Called when a server
    update changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows
    removed.

    LiteLLM_MCPUserCredentials also stores BYOK API keys in the same column; only rows whose payload
    decodes as an OAuth2 credential (see _decode_oauth_payload) are deleted, because a config change
    only invalidates minted tokens, never a user's own stored key. Rows are therefore deleted per
    (user_id, server_id) pair rather than by a blanket server_id filter. An OAuth row inserted while
    the purge runs for a user not yet enumerated survives; a re-auth completing in the window for an
    already-enumerated user is deleted along with the stale row (the pair delete cannot tell them
    apart), which costs that user one extra re-auth and nothing else.

    invalidate_token_cache is injectable for tests; it defaults to the manager's shared
    invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens."""
    rows: Final = await _db_find_user_credential_rows(prisma_client, {"server_id": server_id})
    oauth_rows: Final = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None]
    if not oauth_rows:
        return 0
    deleted_count: Final = await _user_credential_actions(prisma_client).delete_many(
        where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}}
    )
    if invalidate_token_cache is None:
        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
            global_mcp_server_manager,
        )

        invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache

    for row in oauth_rows:
        await invalidate_token_cache(row.user_id, server_id)
    if deleted_count != len(oauth_rows):
        verbose_proxy_logger.warning(
            "MCP server %s: purge removed %d OAuth credential row(s) but %d were enumerated; "
            "row(s) were deleted concurrently during the purge",
            server_id,
            deleted_count,
            len(oauth_rows),
        )
    return deleted_count


async def refresh_user_oauth_token(
    prisma_client: PrismaClient,
    user_id: str,
    server: "MCPServer",
    cred: OAuthCredentialPayload,
) -> OAuthCredentialPayload | None:
    """Attempt to refresh a per-user OAuth2 token using its stored refresh_token.

    POSTs to ``server.effective_token_url`` with ``grant_type=refresh_token``.

    On success: persists the new credential via ``store_user_oauth_credential``
    and returns the updated payload dict.
    On failure (network error, invalid_grant, missing refresh_token, …): logs a
    warning and returns ``None`` — the caller is responsible for clearing the
    stale credential and triggering re-authentication.
    """
    refresh_token: Final[str | None] = cred.get("refresh_token")
    token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None)
    server_id: Final[str] = getattr(server, "server_id", "")
    client_id: Final[str | None] = getattr(server, "client_id", None)
    client_secret: Final[str | None] = getattr(server, "client_secret", None)

    if not refresh_token:
        verbose_proxy_logger.debug(
            "refresh_user_oauth_token: no refresh_token stored for user=%s server=%s",
            user_id,
            server_id,
        )
        return None
    if not token_url:
        verbose_proxy_logger.debug(
            "refresh_user_oauth_token: server=%s has no token_url configured",
            server_id,
        )
        return None

    try:
        token_request: Final = build_upstream_oauth2_token_request(
            server,
            auth_method=getattr(server, "token_endpoint_auth_method", None),
            client_id=client_id,
            client_secret=client_secret,
        )
        token_data: Final[dict[str, str]] = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
            **token_request.body,
        }
        async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
        response: Final = await async_client.post(
            token_url,
            headers={"Accept": "application/json", **token_request.headers},
            data=token_data,
        )
        response.raise_for_status()
        body: Final[_OAuthTokenRefreshResponse] = response.json()
    except Exception as exc:
        verbose_proxy_logger.warning(
            "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s",
            user_id,
            server_id,
            exc,
        )
        return None

    access_token: Final[str | None] = body.get("access_token")
    if not access_token:
        verbose_proxy_logger.warning(
            "refresh_user_oauth_token: token response missing access_token for user=%s server=%s",
            user_id,
            server_id,
        )
        return None

    expires_in: int | None = None
    raw_expires: Final = body.get("expires_in")
    try:
        expires_in = int(raw_expires) if raw_expires is not None else None
    except (TypeError, ValueError):
        pass

    # Rotate refresh token when the provider returns a new one
    new_refresh_token: Final[str | None] = body.get("refresh_token") or refresh_token

    raw_scope: Final = body.get("scope")
    scopes: list[str] | None = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get(
        "scopes"
    )

    await store_user_oauth_credential(
        prisma_client=prisma_client,
        user_id=user_id,
        server_id=server_id,
        access_token=access_token,
        refresh_token=new_refresh_token,
        expires_in=expires_in,
        scopes=scopes,
        skip_byok_guard=True,  # Row is already OAuth2; skip the extra find_unique check
    )

    verbose_proxy_logger.info(
        "refresh_user_oauth_token: refreshed token for user=%s server=%s",
        user_id,
        server_id,
    )
    return await get_user_oauth_credential(prisma_client, user_id, server_id)


async def resolve_valid_user_oauth_token(
    user_id: str,
    server: "MCPServer",
    cred: OAuthCredentialPayload | None,
    prisma_client: PrismaClient | None = None,
) -> OAuthCredentialPayload | None:
    """Return an OAuth2 credential whose access_token is good for the next request.

    Returns the credential unchanged while its token is valid for at least
    ``MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS``. Only when the token is expired (or
    expiring within that buffer) and a refresh_token is stored does it mint a new one
    via ``refresh_user_oauth_token``. Returns None when there is no usable token
    (missing token, expired with no refresh_token, or a failed refresh).

    The refresh_token is only ever sent to the server's token_url inside
    ``refresh_user_oauth_token``; it is never exposed to the caller beyond the cred
    dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh
    actually happens, so the valid-token path never requires a DB handle.
    """
    if not cred or not cred.get("access_token"):
        return None
    if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
        return cred
    if not cred.get("refresh_token"):
        return None
    if prisma_client is None:
        from litellm.proxy.utils import get_prisma_client_or_throw

        prisma_client = get_prisma_client_or_throw("Database not connected. Cannot refresh OAuth token.")
    refreshed: Final = await refresh_user_oauth_token(
        prisma_client=prisma_client,
        user_id=user_id,
        server=server,
        cred=cred,
    )
    if not refreshed or not refreshed.get("access_token"):
        return None
    return refreshed


async def resolve_user_oauth_access_token(
    user_id: str | None,
    server: "MCPServer",
    prefetched_creds: Mapping[str, OAuthCredentialPayload] | None = None,
) -> str | None:
    """Resolve a user's valid OAuth2 access token for a server: Redis cache, else DB + refresh.

    The egress token-resolution core shared by v1's header builder and the v2 ``OAuthTokenStore``
    adapter. Redis fast-path (skipped when ``prefetched_creds`` is supplied), else a DB read through
    ``resolve_valid_user_oauth_token`` (which refreshes an expired token when a ``refresh_token`` is
    stored), re-warming the Redis cache with the per-server TTL. Returns ``None`` when there is no
    usable token; any error is swallowed to ``None`` so a transient failure reads as "not
    authorized" rather than raising.
    """
    server_id: Final[str | None] = getattr(server, "server_id", None)
    if not user_id or not server_id:
        return None
    try:
        from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
            _compute_per_user_token_ttl,
            mcp_per_user_token_cache,
        )

        if prefetched_creds is None:
            cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id)
            if cached_token is not None:
                return cached_token

        prisma_client = None
        if prefetched_creds is not None:
            cred = prefetched_creds.get(server_id)
        else:
            from litellm.proxy.utils import get_prisma_client_or_throw

            prisma_client = get_prisma_client_or_throw(
                "Database not connected. Connect a database to use OAuth2 MCP tools."
            )
            cred = await get_user_oauth_credential(prisma_client, user_id, server_id)

        if not cred or not cred.get("access_token"):
            return None

        cred = await resolve_valid_user_oauth_token(
            user_id=user_id,
            server=server,
            cred=cred,
            prisma_client=prisma_client,
        )
        if cred is None:
            # Refresh failed or token expired with no usable refresh_token — clear the stale
            # Redis entry so the next request doesn't reuse it.
            await mcp_per_user_token_cache.delete(user_id, server_id)
            return None

        access_token: Final[str] = cred["access_token"]
        if prefetched_creds is None:
            ttl: Final = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at")))
            await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
        return access_token
    except Exception as e:
        verbose_proxy_logger.warning(
            "resolve_user_oauth_access_token: failed for user=%s server=%s: %s",
            user_id,
            server_id,
            e,
        )
    return None


def _remaining_token_seconds(expires_at: str | None) -> int | None:
    """Seconds until ``expires_at`` (ISO 8601), or None when absent/past/unparseable."""
    if not expires_at:
        return None
    try:
        exp_dt = datetime.fromisoformat(expires_at)
    except (ValueError, TypeError):
        return None
    if exp_dt.tzinfo is None:
        exp_dt = exp_dt.replace(tzinfo=timezone.utc)
    remaining: Final = int((exp_dt - datetime.now(timezone.utc)).total_seconds())
    return remaining if remaining > 0 else None


async def get_active_submitted_mcp_server_ids_for_user(
    prisma_client: PrismaClient,
    user_id: str,
) -> list[str]:
    """Return active BYOM servers submitted by this user (creator visibility)."""
    if not user_id:
        return []

    rows: Final = await _db_find_mcp_server_rows(
        prisma_client,
        {
            "submitted_by": user_id,
            "approval_status": MCPApprovalStatus.active,
        },
    )
    return [row.server_id for row in rows]


async def approve_mcp_server(
    prisma_client: PrismaClient,
    server_id: str,
    touched_by: str,
) -> LiteLLM_MCPServerTable:
    """Set approval_status=active and record reviewed_at."""
    now: Final = datetime.now(timezone.utc)
    updated: Final = await _db_update_mcp_server_row(
        prisma_client,
        server_id,
        {
            "approval_status": MCPApprovalStatus.active,
            "reviewed_at": now,
            "updated_by": touched_by,
        },
    )
    table: Final = LiteLLM_MCPServerTable.model_validate(updated.model_dump())
    decrypt_global_env_var_values(table.env_vars)
    return table


async def reject_mcp_server(
    prisma_client: PrismaClient,
    server_id: str,
    touched_by: str,
    review_notes: str | None = None,
) -> LiteLLM_MCPServerTable:
    """Set approval_status=rejected, record reviewed_at and review_notes."""
    now: Final = datetime.now(timezone.utc)
    data: Final[prisma_db_types.LiteLLM_MCPServerTableUpdateInput] = {
        "approval_status": MCPApprovalStatus.rejected,
        "reviewed_at": now,
        "updated_by": touched_by,
    }
    if review_notes is not None:
        data["review_notes"] = review_notes
    updated: Final = await _db_update_mcp_server_row(prisma_client, server_id, data)
    table: Final = LiteLLM_MCPServerTable.model_validate(updated.model_dump())
    decrypt_global_env_var_values(table.env_vars)
    return table


async def get_mcp_submissions(
    prisma_client: PrismaClient,
) -> MCPSubmissionsSummary:
    """
    Returns all MCP servers that were submitted by non-admin users (submitted_at IS NOT NULL),
    along with a summary count breakdown by approval_status.
    Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
    """
    rows: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
        prisma_client
    ).find_many(
        where={"submitted_at": {"not": None}},
        order={"submitted_at": "desc"},
        take=500,  # safety cap; paginate if needed in a future iteration
    )
    items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows]
    for item in items:
        decrypt_global_env_var_values(item.env_vars)

    pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review)
    active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active)
    rejected: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.rejected)

    return MCPSubmissionsSummary(
        total=len(items),
        pending_review=pending,
        active=active,
        rejected=rejected,
        items=items,
    )


# ── Per-user MCP environment variables ────────────────────────────────────


def _decode_user_env_vars(stored: str) -> dict[str, str]:
    """Decrypt a ``values_b64`` blob and parse it as a flat ``{name: value}`` dict."""
    decrypted: Final = decrypt_value_helper(
        value=stored,
        key="mcp_user_env_vars",
        exception_type="debug",
        return_original_value=False,
    )
    if decrypted is None:
        if stored:
            verbose_proxy_logger.warning(
                "MCP per-user env vars failed to decrypt (LITELLM_SALT_KEY "
                "changed?); treating as unset so the user is prompted to "
                "re-enter them rather than silently forwarding ciphertext"
            )
        return {}
    parsed: dict[str, object] | None
    try:
        parsed = json.loads(decrypted)
    except (ValueError, TypeError):
        return {}
    if not isinstance(parsed, dict):
        return {}
    return {str(k): str(v) for k, v in parsed.items()}


async def get_user_env_vars(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> dict[str, str]:
    """Return the calling user's env var dict for ``server_id`` (empty if none)."""
    row: Final = await _user_env_var_actions(prisma_client).find_unique(
        where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
    )
    if row is None:
        return {}
    return _decode_user_env_vars(row.values_b64)


async def get_user_env_vars_bulk(
    prisma_client: PrismaClient,
    user_id: str,
    server_ids: Iterable[str],
) -> dict[str, dict[str, str]]:
    """Return ``{server_id: {var_name: value}}`` for one user across many servers.

    Servers with no stored row are simply absent from the result.
    """
    ids: Final = list(server_ids)
    if not ids:
        return {}
    rows: Final = await _db_find_user_env_var_rows(prisma_client, {"user_id": user_id, "server_id": {"in": ids}})
    return {row.server_id: _decode_user_env_vars(row.values_b64) for row in rows}


async def merge_user_env_vars(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
    updates: dict[str, str],
    allowed_names: Iterable[str],
) -> dict[str, str]:
    """Merge ``updates`` into the user's stored env vars for ``server_id`` and
    return the resulting set.

    The read-modify-write runs inside a transaction guarded by a
    ``(user_id, server_id)`` advisory lock so two concurrent writes from the
    same user can't drop one update. Names outside ``allowed_names`` are pruned,
    so an admin retiring a user-scoped variable also clears its stored value.
    """
    allowed: Final = set(allowed_names)
    lock_key: Final = int.from_bytes(
        hashlib.blake2b(f"{user_id}:{server_id}".encode(), digest_size=8).digest(),
        "big",
        signed=True,
    )
    async with _db_transaction_manager(prisma_client) as tx:
        await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
        row: Final[prisma_db_models.LiteLLM_MCPUserEnvVars | None] = await tx.litellm_mcpuserenvvars.find_unique(
            where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
        )
        existing: Final = _decode_user_env_vars(row.values_b64) if row is not None else {}
        merged: Final = {k: v for k, v in {**existing, **updates}.items() if k in allowed}
        encoded: Final = encrypt_value_helper(json.dumps(merged))
        await tx.litellm_mcpuserenvvars.upsert(
            where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
            data={
                "create": {
                    "user_id": user_id,
                    "server_id": server_id,
                    "values_b64": encoded,
                },
                "update": {"values_b64": encoded},
            },
        )
    return merged


async def delete_user_env_vars(
    prisma_client: PrismaClient,
    user_id: str,
    server_id: str,
) -> None:
    """Remove the calling user's env var values for ``server_id``.

    Uses ``delete_many`` so a missing row is a no-op; real DB errors still
    propagate to the caller instead of being silently swallowed.
    """
    await _user_env_var_actions(prisma_client).delete_many(where={"user_id": user_id, "server_id": server_id})
