# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio

import base64

import json

import logging

import re

import sys

import warnings

from abc import abstractmethod

from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Generator, Mapping, Sequence

from contextlib import AsyncExitStack, _AsyncGeneratorContextManager, asynccontextmanager, contextmanager # type: ignore

from contextvars import ContextVar

from copy import copy

from dataclasses import dataclass

from datetime import date, datetime

from http.cookiejar import CookieJar, DefaultCookiePolicy

from inspect import currentframe, isawaitable

from typing import TYPE_CHECKING, Any, Literal, Protocol, Self, TypeAlias, TypedDict, TypeVar, cast

if sys.version_info >= (3, 13):

from warnings import deprecated # pragma: no cover

else:

from typing_extensions import deprecated # pragma: no cover

from opentelemetry import context as otel_context

from opentelemetry import propagate

from opentelemetry import trace as otel_trace

from pydantic_core import to_jsonable_python

from ._feature_stage import (

ExperimentalFeature,

ExperimentalWarning,

_warn_on_feature_use, # pyright: ignore[reportPrivateUsage]

)

from ._serialization import make_json_safe

from ._telemetry import FeatureIndex, mark_feature_used

from ._tools import (

_FUNCTION_RESULT_CARRIER_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage]

_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage]

FunctionTool,

_FunctionResultCarrier, # pyright: ignore[reportPrivateUsage]

_FunctionResultPayloadBudget, # pyright: ignore[reportPrivateUsage]

)

from ._types import (

ChatOptions,

Content,

Message,

)

from .exceptions import ToolException, ToolExecutionException

from .observability import (

OtelAttr,

create_mcp_client_span,

set_mcp_span_error,

)

if TYPE_CHECKING:

# TODO(jpalvarezl): clean up and consolidate under httpx2

import httpx2

from httpx import Request

from mcp import Client, types

from mcp.client import IncomingMessage

from mcp.client.context import ClientRequestContext

from mcp.client.session import ClientSession

from ._clients import SupportsChatGetResponse

from ._middleware import FunctionInvocationContext

logger = logging.getLogger(__name__)

_MCPListResultT = TypeVar(

"_MCPListResultT",

"types.ListToolsResult",

"types.ListPromptsResult",

)

def _mcp_exception_log_level() -> Literal[10] | None:

"""Return the automatic exception log level allowed by observability settings."""

from .observability import OBSERVABILITY_SETTINGS

return 10 if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED else None

def _log_mcp_exception(

level: int,

exception: BaseException,

sensitive_message: str,

*sensitive_args: object,

redacted_message: str,

sensitive_exc_info: bool | BaseException = False,

) -> None:

"""Log an MCP exception without exporting details unless explicitly enabled."""

from .observability import OBSERVABILITY_SETTINGS

if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED:

message = sensitive_message

args = sensitive_args

exc_info = sensitive_exc_info

else:

message = "%s (%s)."

args = (redacted_message, type(exception).__name__)

exc_info = False

if level == logging.DEBUG:

logger.debug(message, *args, exc_info=exc_info)

elif level == logging.WARNING:

logger.warning(message, *args, exc_info=exc_info)

elif level == logging.ERROR:

logger.error(message, *args, exc_info=exc_info)

else:

logger.log(level, message, *args, exc_info=exc_info)

def _mcp_tool_exception(

message: str,

inner_exception: Exception | None = None,

) -> ToolException:

"""Create an MCP tool lifecycle exception without leaking details through automatic logs."""

return ToolException(

message,

inner_exception=inner_exception,

log_level=_mcp_exception_log_level(),

)

def _mcp_tool_execution_exception(

message: str,

inner_exception: Exception | None = None,

) -> ToolExecutionException:

"""Create an MCP tool exception without leaking details through automatic logs."""

return ToolExecutionException(

message,

inner_exception=inner_exception,

log_level=_mcp_exception_log_level(),

)

if TYPE_CHECKING:

_MCPSamplingContentBlock: TypeAlias = (

types.TextContent

| types.ImageContent

| types.AudioContent

| types.EmbeddedResource

| types.ResourceLink

| types.ToolUseContent

)

else:

_MCPSamplingContentBlock = Any

class MCPSpecificApproval(TypedDict, total=False):

"""Represents the specific approval mode for an MCP tool.

When using this mode, the user specifies which tools always or never require approval and may set

``default`` for every other tool and prompt.

Names match raw remote names or unambiguous prefixed names when the raw name already

matches its normalized form. Normalized-only aliases do not match. Discovery raises

``ToolExecutionException`` if a configured name matches multiple raw remote names.

A name list given as a single string is treated as one name.

A tool or prompt resolves in this order: a name in ``always_require_approval`` requires approval, otherwise

a name in ``never_require_approval`` does not, otherwise ``default`` applies. Without ``default``, a tool or

prompt named in neither list does not require approval, which is the framework default. Setting

``never_require_approval`` without ``default`` emits a ``FutureWarning`` because this fallback may change

in a future release; set ``default`` explicitly. Hosted MCP tools do not read ``default``.

Attributes:

always_require_approval: A sequence of tool names that always require approval.

never_require_approval: A sequence of tool names that never require approval.

default: The approval mode for tools and prompts named in neither list. Use ``"always_require"`` to

require approval for every tool except those in ``never_require_approval``, or ``"never_require"``

to require approval only for the tools in ``always_require_approval``. ``None`` means not set.

"""

always_require_approval: Collection[str] | None

never_require_approval: Collection[str] | None

default: Literal["always_require", "never_require"] | None

MCPToolResultContentMode = Literal[

"structured_first",

"content_first",

"content_only",

"structured_only",

"both",

]

"""How model-visible text is selected when a tool result has both ``content`` and ``structuredContent``.

MCP servers disagree on whether the two fields are duplicates or complementary (#7866).

Callers pick an explicit policy:

- ``structured_first`` (default): use ``structuredContent`` when present, else ``content``

- ``content_first``: use ``content`` when non-empty, else ``structuredContent``

- ``content_only``: ignore ``structuredContent`` for the model-visible result

- ``structured_only``: ignore ``content`` for the model-visible result

- ``both``: append serialized ``structuredContent`` after ``content`` blocks

The full MCP payload is still retained on the Host channel regardless of mode.

Custom ``parse_tool_results`` overrides this policy entirely.

"""

_MCP_REMOTE_NAME_KEY = "_mcp_remote_name"

_MCP_IS_TOOL_KEY = "_mcp_is_tool"

_MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name"

_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = "_mcp_tool_result_host_payload"

_MCP_TOOL_RESULT_SECURITY_META_KEY = "_mcp_tool_result_security_meta"

_MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools"

_MCP_PROGRESSIVE_LOAD_TOOL_NAME = "load_tool"

_MCP_PROGRESSIVE_UNLOAD_TOOL_NAME = "unload_tool"

# Reserved key in an ``additional_tool_argument_names`` mapping that applies its

# values to every tool on the server rather than a single named tool.

_MCP_GLOBAL_EXTRA_ARGS_KEY = "*"

_MCP_META_LABEL_PATTERN = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?"

_MCP_META_KEY_PATTERN = re.compile(

rf"^(?:(?:{_MCP_META_LABEL_PATTERN})(?:\.{_MCP_META_LABEL_PATTERN})*/)?"

r"[A-Za-z0-9](?:[A-Za-z0-9_.-]*[A-Za-z0-9])?$"

)

# Framework kwargs that flow through the function-invocation pipeline (via

# ``FunctionInvocationContext.kwargs``) but must never be forwarded to an MCP

# server: they are internal objects that the MCP SDK cannot serialize. They are

# dropped as a safety net when a tool declares one of them in its schema, unless

# the user explicitly opts the name back in via ``additional_tool_argument_names``

# (explicit extras always win over the denylist).

# - chat_options/tools/tool_choice/session/thread: framework runtime objects.

# - conversation_id: internal tracking ID used by services like Azure AI.

# - options: metadata/store used by AG-UI for Azure AI client requirements.

# - response_format: a Pydantic model class for structured output (not serializable).

# - _meta: reserved key extracted separately as MCP request metadata.

_MCP_FRAMEWORK_DENYLIST: frozenset[str] = frozenset({

"chat_options",

"tools",

"tool_choice",

"session",

"thread",

"conversation_id",

"options",

"response_format",

"_meta",

})

# object is a reference to the owner of the headers, so we keep track which headers belong to which tool

_mcp_call_headers: ContextVar[tuple[object, dict[str, str]] | None] = ContextVar("_mcp_call_headers")

@dataclass(frozen=True)

class _MCPToolRuntimeContext:

owner: object

runtime_kwargs: Mapping[str, Any]

result_capture: _MCPResultCapture | None = None

_mcp_tool_runtime_context: ContextVar[_MCPToolRuntimeContext | None] = ContextVar(

"_mcp_tool_runtime_context",

default=None,

)

_MCP_HEADER_OWNER_EXTENSION = "agent_framework.mcp_header_owner"

_MCP_INJECTED_HEADER_KEYS_EXTENSION = "agent_framework.mcp_injected_header_keys"

_MCP_REQUEST_HEADERS_EXTENSION = "agent_framework.mcp_request_headers"

_MCPHeaderIdentity: TypeAlias = tuple[tuple[str, str], ...]

_MCP_DEFAULT_READ_TIMEOUT_SECONDS = 30

MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5

_MCP_SUBSCRIPTION_RETRY_INITIAL_DELAY = 0.1

_MCP_SUBSCRIPTION_RETRY_MAX_DELAY = 2.0

_DEFAULT_MCP_DISCOVERY_MAX_PAGES = 100

_DEFAULT_MCP_DISCOVERY_MAX_ITEMS = 1_000

_DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024

_MAX_MCP_SECURITY_META_SIZE_BYTES = 4096

@dataclass

class _MCPHTTPInvocation:

"""Track a live network attempt, including callbacks that inherit its context."""

owner: object

parent: _MCPHTTPInvocation | None

runtime_kwargs: Mapping[str, Any]

connection: _MCPConnection | None = None

active_depth: int = 0

_mcp_http_invocation: ContextVar[_MCPHTTPInvocation | None] = ContextVar("_mcp_http_invocation", default=None)

class _MCPDiscoveryRestart(Exception):

"""Restart staged discovery when its connection is replaced."""

def _retain_mcp_result_security_meta(mcp_type: Any) -> dict[str, Any] | None:

"""Retain bounded, JSON-safe MCP IFC metadata independently of generic payload budgets."""

raw_meta = getattr(mcp_type, "meta", None)

if not isinstance(raw_meta, Mapping) or "ifc" not in raw_meta:

return None

raw_meta_mapping = cast(Mapping[str, Any], raw_meta)

raw_security_meta: dict[str, Any] = {"ifc": raw_meta_mapping["ifc"]}

if _json_size_exceeds(raw_security_meta, _MAX_MCP_SECURITY_META_SIZE_BYTES):

logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.")

return {"retained": False}

try:

copied = cast(

dict[str, Any],

to_jsonable_python(raw_security_meta, fallback=str),

)

except (TypeError, ValueError):

logger.warning("MCP result _meta.ifc could not be retained safely.")

return {"retained": False}

if len(json.dumps(copied).encode("utf-8")) > _MAX_MCP_SECURITY_META_SIZE_BYTES:

logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.")

return {"retained": False}

return {"retained": True, "meta": copied}

@dataclass

class _MCPResultCapture:

"""Collect one generated MCP result without affecting its public return shape."""

max_size_bytes: int | None

aggregate_budget: _FunctionResultPayloadBudget | None

apply_server_meta_to_model_items: bool

retain_generic_payload: bool = True

host_payload: dict[str, Any] | None = None

meta: dict[str, Any] | None = None

security_meta: dict[str, Any] | None = None

recorded: bool = False

meta_prepared: bool = False

security_meta_prepared: bool = False

def prepare_security_meta(self, mcp_type: Any) -> dict[str, Any] | None:

"""Retain security-relevant MCP metadata outside generic payload budgets."""

if self.security_meta_prepared:

return self.security_meta

self.security_meta_prepared = True

self.security_meta = _retain_mcp_result_security_meta(mcp_type)

return self.security_meta

def prepare_meta(self, mcp_type: Any) -> dict[str, Any] | None:

if self.meta_prepared:

return self.meta

self.meta_prepared = True

effective_limit = self.max_size_bytes

if self.aggregate_budget is not None:

remaining = self.aggregate_budget.remaining(self.max_size_bytes)

if remaining == 0:

logger.warning("Omitting MCP result _meta because the request retention budget is exhausted.")

return None

if remaining is not None:

effective_limit = min(effective_limit, remaining) if effective_limit is not None else remaining

meta = _mcp_tool_result_meta(mcp_type, max_size_bytes=effective_limit)

if meta is None:

return None

encoded_size = len(json.dumps(meta).encode("utf-8"))

if self.aggregate_budget is not None and not self.aggregate_budget.reserve(encoded_size, self.max_size_bytes):

logger.warning("Omitting MCP result _meta because the request retention budget is exhausted.")

return None

self.meta = meta

return self.meta

def record(self, mcp_type: Any) -> None:

if self.recorded:

return

self.recorded = True

self.prepare_security_meta(mcp_type)

if not self.retain_generic_payload:

return

self.prepare_meta(mcp_type)

effective_limit = self.max_size_bytes

if self.aggregate_budget is not None:

remaining = self.aggregate_budget.remaining(self.max_size_bytes)

if remaining == 0:

logger.warning("Omitting MCP Host payload because the request retention budget is exhausted.")

return

if remaining is not None:

effective_limit = min(effective_limit, remaining) if effective_limit is not None else remaining

host_payload = _mcp_tool_result_host_payload(mcp_type, max_size_bytes=effective_limit)

if host_payload is None:

return

encoded_size = len(json.dumps(host_payload).encode("utf-8"))

if self.aggregate_budget is not None and not self.aggregate_budget.reserve(encoded_size, self.max_size_bytes):

logger.warning("Omitting MCP Host payload because the request retention budget is exhausted.")

return

self.host_payload = host_payload

def to_carrier(self) -> _FunctionResultCarrier:

additional_properties: dict[str, Any] = {}

if self.meta is not None:

additional_properties["_meta"] = self.meta

if self.host_payload is not None:

additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] = self.host_payload

item_additional_properties = (

{"_meta": self.meta} if self.apply_server_meta_to_model_items and self.meta is not None else {}

)

return _FunctionResultCarrier(

additional_properties=additional_properties,

item_additional_properties=item_additional_properties,

private_metadata=(

{_MCP_TOOL_RESULT_SECURITY_META_KEY: self.security_meta} if self.security_meta is not None else {}

),

exclusive_outer_keys=frozenset({_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY}),

exclusive_item_keys=frozenset({

"_meta",

"_security_label_authoritative",

"_security_label_authoritative_confidentiality",

"_security_label_validated_principals",

_MCP_TOOL_RESULT_SECURITY_META_KEY,

}),

result_already_parsed=True,

)

def prepare_model_result(self, parsed: str | list[Content]) -> list[Content]:

"""Apply bounded metadata before security middleware sees generated output."""

items = [Content.from_text(parsed)] if isinstance(parsed, str) else list(parsed)

for index, item in enumerate(items):

updated_item = copy(item)

updated_item.additional_properties = dict(updated_item.additional_properties)

updated_item.additional_properties.pop(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, None)

updated_item.additional_properties.pop("_meta", None)

updated_item.additional_properties.pop(_MCP_TOOL_RESULT_SECURITY_META_KEY, None)

if self.apply_server_meta_to_model_items and self.meta is not None:

updated_item.additional_properties["_meta"] = self.meta

if self.security_meta is not None:

updated_item.additional_properties[_MCP_TOOL_RESULT_SECURITY_META_KEY] = self.security_meta

items[index] = updated_item

return items

class _EncodedSizeBudget:

"""Track JSON bytes and abort as soon as an untrusted value exceeds its budget."""

def __init__(self, limit: int) -> None:

self.remaining = limit

def consume(self, size: int) -> None:

self.remaining -= size

if self.remaining < 0:

raise OverflowError

def _consume_json_string(value: str, budget: _EncodedSizeBudget) -> None:

budget.consume(2)

for char in value:

codepoint = ord(char)

if char in {'"', "\\"} or char in {"\b", "\f", "\n", "\r", "\t"}:

budget.consume(2)

elif codepoint < 0x20:

budget.consume(6)

elif codepoint > 0xFFFF:

budget.consume(12)

elif codepoint > 0x7F:

budget.consume(6)

else:

budget.consume(1)

def _consume_json_size(value: Any, budget: _EncodedSizeBudget) -> None:

if value is None:

budget.consume(4)

return

if value is True:

budget.consume(4)

return

if value is False:

budget.consume(5)

return

if isinstance(value, str):

_consume_json_string(value, budget)

return

if isinstance(value, (bytes, bytearray)):

budget.consume(2 + 4 * ((len(value) + 2) // 3))

return

if isinstance(value, (int, float)):

budget.consume(len(json.dumps(value)))

return

if isinstance(value, (datetime, date)):

_consume_json_string(value.isoformat(), budget)

return

if callable(getattr(value, "unicode_string", None)):

_consume_json_string(str(value), budget)

return

if isinstance(value, Mapping):

budget.consume(2)

for index, (key, item) in enumerate(cast(Mapping[Any, Any], value).items()):

if index:

budget.consume(2)

_consume_json_string(str(key), budget)

budget.consume(2)

_consume_json_size(item, budget)

return

if isinstance(value, Sequence):

budget.consume(2)

for index, item in enumerate(cast(Sequence[Any], value)):

if index:

budget.consume(2)

_consume_json_size(item, budget)

return

model_fields_raw = getattr(value.__class__, "model_fields", None)

if isinstance(model_fields_raw, Mapping):

model_fields = cast(Mapping[str, Any], model_fields_raw)

budget.consume(2)

field_count = 0

for name, field_info in model_fields.items():

item = getattr(value, name, None)

if item is None or getattr(field_info, "exclude", False):

continue

if field_count:

budget.consume(2)

alias = getattr(field_info, "serialization_alias", None) or getattr(field_info, "alias", None) or name

_consume_json_string(str(alias), budget)

budget.consume(2)

_consume_json_size(item, budget)

field_count += 1

model_extra_raw = getattr(value, "model_extra", None)

if isinstance(model_extra_raw, Mapping):

for key, item in cast(Mapping[Any, Any], model_extra_raw).items():

if item is None:

continue

if field_count:

budget.consume(2)

_consume_json_string(str(key), budget)

budget.consume(2)

_consume_json_size(item, budget)

field_count += 1

return

_consume_json_string(str(value), budget)

def _json_size_exceeds(value: Any, limit: int) -> bool:

try:

_consume_json_size(value, _EncodedSizeBudget(limit))

except OverflowError:

return True

return False

def _mcp_tool_result_host_payload(

mcp_type: Any,

*,

max_size_bytes: int | None,

) -> dict[str, Any] | None:

"""Return a bounded, JSON-safe copy of a complete MCP result."""

model_dump = getattr(mcp_type, "model_dump", None)

if not callable(model_dump):

return None

if max_size_bytes is not None and _json_size_exceeds(mcp_type, max_size_bytes):

logger.warning(

"Omitting MCP Host payload because its encoded size exceeds max_host_payload_size_bytes=%d.",

max_size_bytes,

)

return None

try:

dumped = model_dump(by_alias=True, exclude_none=True, mode="json", fallback=str)

except ValueError:

dumped = model_dump(by_alias=True, exclude_none=True)

if not isinstance(dumped, Mapping):

return None

host_payload = cast(dict[str, Any], make_json_safe(dict(cast(Mapping[str, Any], dumped))))

if max_size_bytes is not None and len(json.dumps(host_payload).encode("utf-8")) > max_size_bytes:

logger.warning(

"Omitting MCP Host payload because its encoded size exceeds max_host_payload_size_bytes=%d.",

max_size_bytes,

)

return None

return host_payload

def _mcp_tool_result_meta(

mcp_type: Any,

*,

max_size_bytes: int | None,

) -> dict[str, Any] | None:

"""Return a separately bounded, JSON-safe copy of MCP result metadata."""

raw_meta = getattr(mcp_type, "meta", None)

if not isinstance(raw_meta, Mapping):

return None

if max_size_bytes is not None and _json_size_exceeds(raw_meta, max_size_bytes):

logger.warning(

"Omitting MCP result _meta because its encoded size exceeds max_host_payload_size_bytes=%d.",

max_size_bytes,

)

return None

meta = cast(dict[str, Any], to_jsonable_python(dict(cast(Mapping[str, Any], raw_meta)), fallback=str))

if max_size_bytes is not None and len(json.dumps(meta).encode("utf-8")) > max_size_bytes:

logger.warning(

"Omitting MCP result _meta because its encoded size exceeds max_host_payload_size_bytes=%d.",

max_size_bytes,

)

return None

return meta

def _capture_mcp_result(owner: object, mcp_type: Any) -> None:

runtime_context = _mcp_tool_runtime_context.get()

capture = (

runtime_context.result_capture

if runtime_context is not None and runtime_context.owner is owner

else None

)

if capture is not None:

capture.record(mcp_type)

class _MCPHeaderScopedClient:

"""Attach private tool context to MCP transport requests."""

def __init__(self, client: httpx2.AsyncClient, owner: object) -> None:

self._client = client

self._owner = owner

def __getattr__(self, name: str) -> Any:

if name == "_client":

raise AttributeError(name)

if name in {"get", "post", "put", "patch", "head", "options"}:

raise RuntimeError(

f"MCP HTTP transport cannot use unscoped {name}(); use request(), stream(), send(), or delete()."

)

return getattr(self._client, name)

def _tagged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:

extensions = dict(kwargs.get("extensions") or {})

extensions[_MCP_HEADER_OWNER_EXTENSION] = self._owner

kwargs["extensions"] = extensions

return kwargs

def stream(self, *args: Any, **kwargs: Any) -> Any:

return self._client.stream(*args, **self._tagged_kwargs(kwargs))

async def request(self, *args: Any, **kwargs: Any) -> httpx2.Response:

return await self._client.request(*args, **self._tagged_kwargs(kwargs))

async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response:

request.extensions[_MCP_HEADER_OWNER_EXTENSION] = self._owner

return await self._client.send(request, **kwargs)

async def delete(self, *args: Any, **kwargs: Any) -> Any:

return await self._client.delete(*args, **self._tagged_kwargs(kwargs))

class _MCPConnection(Protocol):

"""Normalized MCP connection surface used by MCPTool."""

"""Return the high-level Client when this connection has one."""

client: Client | None

"""Return the low-level compatibility session."""

session: ClientSession | None

"""Return whether Agent Framework owns this connection."""

is_framework_owned: bool

async def call_tool(

self,

name: str,

arguments: Mapping[str, Any] | None,

*,

meta: types.RequestParamsMeta | None,

) -> types.CallToolResult:

"""Call a tool through this connection."""

...

async def get_prompt(self, name: str, arguments: Mapping[str, str] | None) -> types.GetPromptResult:

"""Get a prompt through this connection."""

...

async def list_tools_page(self, params: types.PaginatedRequestParams | None) -> types.ListToolsResult:

"""List one tools page through this connection."""

...

async def list_prompts_page(self, params: types.PaginatedRequestParams | None) -> types.ListPromptsResult:

"""List one prompts page through this connection."""

...

async def set_logging_level(self, level: types.LoggingLevel) -> None:

"""Set the legacy logging level through this connection."""

...

async def check_connection(self) -> None:

"""Check the connection with the MCP server, depending on the version."""

...

async def read_resource(self, uri: str) -> types.ReadResourceResult:

"""Read a resource through this connection."""

...

class _ClientMCPConnection:

"""MCP connection backed by a high-level Client."""

def __init__(self, client: Client, *, is_framework_owned: bool) -> None:

self.client: Client | None = client

self.session: ClientSession | None = None

self.is_framework_owned = is_framework_owned

self._client = client

async def call_tool(

self,

name: str,

arguments: Mapping[str, Any] | None,

*,

meta: types.RequestParamsMeta | None,

) -> types.CallToolResult:

"""Call a tool through the high-level Client."""

return await self._client.call_tool(

name,

arguments=dict(arguments) if arguments is not None else None,

meta=meta,

)

async def get_prompt(self, name: str, arguments: Mapping[str, str] | None) -> types.GetPromptResult:

"""Get a prompt through the high-level Client."""

return await self._client.get_prompt(name, arguments=dict(arguments) if arguments is not None else None)

async def list_tools_page(self, params: types.PaginatedRequestParams | None) -> types.ListToolsResult:

"""List one tools page while honoring server cache hints."""

return await self._client.list_tools(

cursor=params.cursor if params is not None else None,

)

async def list_prompts_page(self, params: types.PaginatedRequestParams | None) -> types.ListPromptsResult:

"""List one prompts page while honoring server cache hints."""

return await self._client.list_prompts(

cursor=params.cursor if params is not None else None,

)

async def set_logging_level(self, level: types.LoggingLevel) -> None:

"""Set legacy logging; modern request metadata is configured in the Client constructor."""

if self._client.session.initialize_result is None:

return

await self._client.set_logging_level(level) # pyright: ignore[reportDeprecated]

async def check_connection(self) -> None:

"""Check legacy Client connections; modern MCP has no ping."""

if self._client.session.initialize_result is None:

return

await self._client.send_ping() # pyright: ignore[reportDeprecated]

async def read_resource(self, uri: str) -> types.ReadResourceResult:

"""Read a resource through the high-level Client while honoring server cache hints."""

return await self._client.read_resource(uri, cache_mode="use")

class _SessionMCPConnection:

"""MCP connection backed by a low-level ClientSession."""

def __init__(self, session: ClientSession) -> None:

self.session: ClientSession | None = session

self.client: Client | None = None

self.is_framework_owned = False

self._session = session

async def call_tool(

self,

name: str,

arguments: Mapping[str, Any] | None,

*,

meta: types.RequestParamsMeta | None,

) -> types.CallToolResult:

"""Call a tool directly through the session."""

return await self._session.call_tool(

name,

arguments=dict(arguments) if arguments is not None else None,

meta=meta,

)

async def get_prompt(self, name: str, arguments: Mapping[str, str] | None) -> types.GetPromptResult:

"""Get a prompt directly through the session."""

return await self._session.get_prompt(name, arguments=dict(arguments) if arguments is not None else None)

async def list_tools_page(self, params: types.PaginatedRequestParams | None) -> types.ListToolsResult:

"""List one tools page directly through the session."""

return await self._session.list_tools(params=params)

async def list_prompts_page(self, params: types.PaginatedRequestParams | None) -> types.ListPromptsResult:

"""List one prompts page directly through the session."""

return await self._session.list_prompts(params=params)

async def set_logging_level(self, level: types.LoggingLevel) -> None:

"""Set the legacy logging level directly through the session."""

if self._session.initialize_result is None:

return

await self._session.set_logging_level(level) # pyright: ignore[reportDeprecated]

async def check_connection(self) -> None:

"""Check the connection with the MCP server through this session."""

if self._session.initialize_result is None:

return

await self._session.send_ping()

async def read_resource(self, uri: str) -> types.ReadResourceResult:

"""Read a resource through this connection."""

return await self._session.read_resource(uri)

def _as_mcp_connection( # pyright: ignore[reportUnusedFunction]

client: Client | ClientSession,

) -> _MCPConnection:

from mcp import Client as MCPClient

if isinstance(client, MCPClient):

return _ClientMCPConnection(client, is_framework_owned=False)

return _SessionMCPConnection(client)

# Default safety limits applied to server-initiated MCP sampling requests

# (``sampling/createMessage``). MCP servers are untrusted third parties, so the

# default ``sampling_callback`` denies requests unless an approval callback is

# supplied, and bounds the cost of any approved request.

# - ``_DEFAULT_SAMPLING_MAX_TOKENS`` clamps the server-requested ``maxTokens``.

# - ``_DEFAULT_SAMPLING_MAX_REQUESTS`` caps the number of sampling requests per

# session connection (the counter resets on reconnect).

_DEFAULT_SAMPLING_MAX_TOKENS = 4096

_DEFAULT_SAMPLING_MAX_REQUESTS = 25

_MCP_SAMPLING_DEPRECATION_MESSAGE = (

"MCP sampling is deprecated as of MCP specification version 2026-07-28. Under the MCP feature lifecycle, "

"it remains supported for at least twelve months before becoming eligible for removal. "

"MCP servers should call LLM provider APIs directly."

)

_MCP_SESSION_DEPRECATION_MESSAGE = (

"`MCPTool.session` is deprecated and will be removed in vNext. "

"Use a high-level `mcp.Client` via `mcp_client=`. "

"See https://github.com/microsoft/agent-framework/issues/9452."

)

# A user-supplied gate invoked before each server-initiated sampling request is

# forwarded to the chat client. It receives the raw ``CreateMessageRequestParams``

# and returns (or awaits to) a truthy value to approve the request or a falsy

# value to deny it. Both synchronous and asynchronous callables are supported.

SamplingApprovalCallback = Callable[["types.CreateMessageRequestParams"], "bool | Coroutine[Any, Any, bool]"]

_MCPFunctionLoadCallback = Callable[[FunctionTool, Any], None]

# region: Helpers

LOG_LEVEL_MAPPING: dict[str, int] = {

"debug": logging.DEBUG,

"info": logging.INFO,

"notice": logging.INFO,

"warning": logging.WARNING,

"error": logging.ERROR,

"critical": logging.CRITICAL,

"alert": logging.CRITICAL,

"emergency": logging.CRITICAL,

}

def _is_mcp_prompt_function(func: FunctionTool) -> bool:

"""Return whether a function was created from an MCP prompt."""

properties = func.additional_properties or {}

return isinstance(properties.get(_MCP_REMOTE_NAME_KEY), str) and not properties.get(_MCP_IS_TOOL_KEY)

def _to_mcp_logging_level(level: int) -> types.LoggingLevel | None:

"""Map a Python logging level to its MCP equivalent."""

if level == logging.NOTSET:

return None

return cast(

"types.LoggingLevel | None", next((name for name, value in LOG_LEVEL_MAPPING.items() if value == level), None)

)

def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> dict[str, Any]:

"""Get the input model from an MCP prompt.

Returns a JSON schema dictionary for prompt arguments.

"""

# Check if 'arguments' is missing or empty

if not prompt.arguments:

return {"type": "object", "properties": {}}

# Convert prompt arguments to JSON schema format

properties: dict[str, Any] = {}

required: list[str] = []

for prompt_argument in prompt.arguments:

# For prompts, all arguments are typically string type unless specified otherwise

# `description` is optional on PromptArgument and None when absent, which is not

# a valid JSON Schema description, so leave the key out instead.

properties[prompt_argument.name] = {"type": "string"}

if prompt_argument.description is not None:

properties[prompt_argument.name]["description"] = prompt_argument.description

if prompt_argument.required:

required.append(prompt_argument.name)

schema: dict[str, Any] = {"type": "object", "properties": properties}

if required:

schema["required"] = required

return schema

def _normalize_mcp_name(name: str) -> str:

"""Normalize MCP tool/prompt names to allowed identifier pattern (A-Za-z0-9_.-)."""

return re.sub(r"[^A-Za-z0-9_.-]", "-", name)

def _build_prefixed_mcp_name(

normalized_name: str,

tool_name_prefix: str | None,

) -> str:

"""Build the exposed MCP function name from a normalized name and optional prefix."""

if not tool_name_prefix:

return normalized_name

normalized_prefix = _normalize_mcp_name(tool_name_prefix).rstrip("_.-")

if not normalized_prefix:

return normalized_name

trimmed_name = normalized_name.lstrip("_.-")

return f"{normalized_prefix}_{trimmed_name}" if trimmed_name else normalized_prefix

@deprecated(

"MCP Tasks support was removed during the MCP SDK 2.x migration. "

"MCPTaskOptions is retained only for import compatibility.",

category=DeprecationWarning,

)

class MCPTaskOptions:

"""Deprecated compatibility shim for the removed MCP Tasks API."""

def __init__(self, *args: Any, **kwargs: Any) -> None:

del args, kwargs

def _normalize_additional_tool_argument_names(

additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None,

) -> tuple[set[str], dict[str, set[str]]]:

"""Split user-supplied extra argument names into global and per-tool sets.

Accepts either a sequence (applied to every tool) or a mapping keyed by remote

tool name, where the reserved key ``"*"`` is treated as global. Mapping values

may be a sequence or a single string. Returns a

``(global_extras, per_tool_extras)`` tuple.

"""

if additional_tool_argument_names is None:

return set(), {}

if isinstance(additional_tool_argument_names, str):

return {additional_tool_argument_names}, {}

if isinstance(additional_tool_argument_names, Mapping):

global_extras: set[str] = set()

per_tool_extras: dict[str, set[str]] = {}

for tool_name, names in additional_tool_argument_names.items():

# Treat a bare string value as a single name rather than iterating its characters.

names_set = {names} if isinstance(names, str) else set(names)

if tool_name == _MCP_GLOBAL_EXTRA_ARGS_KEY:

global_extras.update(names_set)

else:

per_tool_extras[tool_name] = names_set

return global_extras, per_tool_extras

return set(additional_tool_argument_names), {}

def _mcp_config_candidate_names(*, local_name: str, normalized_name: str, remote_name: str) -> tuple[str, ...]:

"""Return safe configuration names for MCP allow/approval matching."""

names = [remote_name]

if normalized_name == remote_name and local_name != remote_name:

names.append(local_name)

return tuple(names)

_MCP_APPROVAL_MODES = ("always_require", "never_require")

_MCP_APPROVAL_DEFAULT_WARNING = (

"approval_mode sets `never_require_approval` without `default`. Tools and prompts named in neither list use "

'"never_require" today, and this fallback may change in a future release. Set `default` to "never_require" '

'to keep the current behavior or to "always_require" to require approval for them.'

)

def _mcp_approval_names(names: Collection[str] | None) -> Collection[str]:

"""Return approval-policy names, treating a bare string as one name rather than a substring pattern."""

if names is None:

return ()

if isinstance(names, str):

return (names,)

return names

def _mcp_approval_default(

approval_mode: MCPSpecificApproval,

*,

warn_stacklevel: int | None = None,

) -> Literal["always_require", "never_require"] | None:

"""Return the validated ``default`` of an approval policy, or ``None`` when it is not set.

The per-tool resolver calls this for every discovered tool and prompt, so it stays silent unless asked.

Constructors pass ``warn_stacklevel``, counted as for ``warnings.warn`` from the caller of this function, to warn

once, at the code that built the tool, when ``never_require_approval`` is used without ``default``.

"""

default = approval_mode.get("default")

if default is None:

if warn_stacklevel is not None and approval_mode.get("never_require_approval") is not None:

warnings.warn(_MCP_APPROVAL_DEFAULT_WARNING, FutureWarning, stacklevel=warn_stacklevel + 1)