# demo.py
# author: Doiiars
# ... Other import statements ...
from fastapi_mcp import FastApiMCP
from fastapi_mcp.server import FastApiSseTransport
from fastapi import FastAPI, APIRouter, Request, params, Depends
from typing import Optional, Sequence
# ...
# --- Monkey Patching FastApiMCP ---
# 1. Define our own version, completely replacing the original implementation
def patched_register_mcp_connection_endpoint_sse(
 self,
 router: FastAPI | APIRouter,
 transport: FastApiSseTransport,
 mount_path: str,
 dependencies: Optional[Sequence[params.Depends]],
):
 """
 This is a patched version of _register_mcp_connection_endpoint_sse.
 It replaces @router.get with @router.api_route(methods=["GET", "POST"])
 and injects a custom dependency.
 """
 logger.info(f"--- Patched function called! Registering route for MCP endpoint at '{mount_path}' (GET/POST) ---")
 # --- Add your custom logic here ---
 # For example, we can add a custom dependency for authentication or logging
 async def custom_dependency(request: Request):
 logger.info(f"MCP connection request from: {request.client.host}")
 # You can add authentication logic here, and raise HTTPException if it fails
 # Add our dependency to the existing list of dependencies
 new_dependencies = list(dependencies) if dependencies else []
 new_dependencies.append(Depends(custom_dependency))
 
 logger.info("Custom dependency injected successfully.")
 # ------------------------------------
 # 2. Rewrite the route registration logic, using api_route and specifying methods
 @router.api_route(
 mount_path, 
 include_in_schema=False, 
 operation_id="mcp_connection", 
 dependencies=new_dependencies, # Use our modified dependencies
 methods=["GET", "POST"]
 )
 async def handle_mcp_connection(request: Request):
 async with transport.connect_sse(request.scope, request.receive, request._send) as (reader, writer):
 await self.server.run(
 reader,
 writer,
 self.server.create_initialization_options(notification_options=None, experimental_capabilities={}),
 raise_exceptions=False,
 )
 
 logger.info(f"Successfully registered MCP endpoint to '{mount_path}' using methods GET, POST.")
# 3. Apply the patch: replace the original version with our patched version
FastApiMCP._register_mcp_connection_endpoint_sse = patched_register_mcp_connection_endpoint_sse
logger.info("FastApiMCP._register_mcp_connection_endpoint_sse has been successfully patched.")
# --- End of Monkey Patching ---
mcp = FastApiMCP(app)
# Mount the MCP server directly to your FastAPI app
mcp.mount()