Repository: microsoft/semantic-kernel
File: python/semantic_kernel/connectors/mcp.py (Python)

Code before:
                    "Failed to connect to the MCP server. Please check your configuration."
                ) from ex
            try:
                session = await self._exit_stack.enter_async_context(
                    ClientSession(
                        read_stream=transport[0],
                        write_stream=transport[1],
                        read_timeout_seconds=timedelta(seconds=self.request_timeout),
                        message_handler=self.message_handler,
                        logging_callback=self.logging_callback,
                        sampling_callback=self.sampling_callback,
                    )
                )
            except Exception as ex:
                await self._exit_stack.aclose()
                raise KernelPluginInvalidConfigurationError(
                    "Failed to create a session. Please check your configuration."
                ) from ex
            await session.initialize()
            self.session = session
        elif self.session._request_id == 0:
            # If the session is not initialized, we need to reinitialize it
            await self.session.initialize()
        logger.debug("Connected to MCP server: %s", self.session)
        await self.load_tools()
        await self.load_prompts()

        if logger.level != logging.NOTSET:
            try:
                await self.session.set_logging_level(
                    next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
                )
            except Exception:
                logger.warning("Failed to set log level to %s", logger.level)

    async def sampling_callback(
        self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams
    ) -> types.CreateMessageResult | types.ErrorData:
        """Callback function for sampling.

        This function is called when the MCP server needs to get a message completed.

        This is a simple version of this function, it can be overridden to allow more complex sampling.
        It get's added to the session at initialization time, so overriding it is the best way to do this.
        """
        if not self.kernel or not self.kernel.services:
            return types.ErrorData(
                code=types.INTERNAL_ERROR,
                message="No services in Kernel. Please set a kernel with one or more services.",
            )
        logger.debug("Sampling callback called with params: %s", params)

Review comment (suggestion):
This call will still hang if the server returns a 400. Any timeout config will cause an issue in our `add_plugin` code where the kernel function metadata will be missing, and we'll add a blank plugin with no underlying functions.

Code after:
                transport = await self._exit_stack.enter_async_context(self.get_mcp_client())
            except Exception as ex:
                await self._exit_stack.aclose()
                raise KernelPluginInvalidConfigurationError(
                    "Failed to connect to the MCP server. Please check your configuration."
                ) from ex
            try:
                session = await self._exit_stack.enter_async_context(
                    ClientSession(
                        read_stream=transport[0],
                        write_stream=transport[1],
                        read_timeout_seconds=timedelta(seconds=self.request_timeout) if self.request_timeout else None,
                        message_handler=self.message_handler,
                        logging_callback=self.logging_callback,
                        sampling_callback=self.sampling_callback,
                    )
                )
            except Exception as ex:
                await self._exit_stack.aclose()
                raise KernelPluginInvalidConfigurationError(
                    "Failed to create a session. Please check your configuration."
                ) from ex
            await session.initialize()
            self.session = session
        elif self.session._request_id == 0:
            # If the session is not initialized, we need to reinitialize it
            await self.session.initialize()
        logger.debug("Connected to MCP server: %s", self.session)
        if self.load_tools_flag:
            await self.load_tools()
        if self.load_prompts_flag:
            await self.load_prompts()

        if logger.level != logging.NOTSET:
            try:
                await self.session.set_logging_level(
                    next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
                )
            except Exception:
                logger.warning("Failed to set log level to %s", logger.level)

    async def sampling_callback(
        self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams
    ) -> types.CreateMessageResult | types.ErrorData:
        """Callback function for sampling.

        This function is called when the MCP server needs to get a message completed.

        This is a simple version of this function, it can be overridden to allow more complex sampling.
        It get's added to the session at initialization time, so overriding it is the best way to do this.
        """