o
    ir                     @  s   d dl mZ d dlZd dlmZ d dlmZmZ d dlZd dl	Z
d dlmZmZ d dlmZmZmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ erdd dl m!Z! ee"Z#G dd deZ$G dd deeef Z%dS )    )annotationsN)AsyncExitStack)TYPE_CHECKINGAny)MemoryObjectReceiveStreamMemoryObjectSendStream)LifespanResultTNotificationOptionsRequestT)Server)InitializationOptions)ServerSession)stdio_server)SessionMessage)RequestResponder)
get_logger)FastMCPc                      s>   e Zd ZdZd fddZedddZd fddZ  ZS )MiddlewareServerSessionzMServerSession that routes initialization requests through FastMCP middleware.fastmcpr   c                   s"   t  j|i | t|| _d S N)super__init__weakrefref_fastmcp_refselfr   argskwargs	__class__ ^/var/www/html/karishye-ai-python/venv/lib/python3.10/site-packages/fastmcp/server/low_level.pyr   #   s   z MiddlewareServerSession.__init__returnc                 C     |   }|du rtd|S zGet the FastMCP instance.Nz'FastMCP instance is no longer availabler   RuntimeErrorr   r   r!   r!   r"   r   '      zMiddlewareServerSession.fastmcp	responderARequestResponder[mcp.types.ClientRequest, mcp.types.ServerResult]c              	     s   ddl }ddlm} t jjtjjrWd fdd	}|j	j
jjd
4 I dH  }| jjddd|d}j||I dH W  d  I dH  S 1 I dH sPw   Y  dS t  I dH S )z
        Override the _received_request method to route initialization requests
        through FastMCP middleware.

        These are not handled by routes that FastMCP typically overrides and
        require special handling.
        r   N)MiddlewareContextctxr,   r#   Nonec                   s   t tI d H S r   )r   r   _received_request)r-   )r    r*   r   r!   r"   call_original_handler?   s   

zHMiddlewareServerSession._received_request.<locals>.call_original_handler)r   clientrequest
initialize)messagesourcetypemethodfastmcp_context)r-   r,   r#   r.   )fastmcp.server.context$fastmcp.server.middleware.middlewarer,   
isinstancer2   rootmcptypesInitializeRequestservercontextContextr   _apply_middlewarer   r/   )r   r*   r   r,   r0   fastmcp_ctx
mw_contextr   )r*   r   r"   r/   /   s*   0z)MiddlewareServerSession._received_request)r   r   r#   r   )r*   r+   )	__name__
__module____qualname____doc__r   propertyr   r/   __classcell__r!   r!   r   r"   r       s    r   c                      sP   e Zd Zd fddZedd	d
Z		d d! fddZ		d"d#ddZ  ZS )$LowLevelServerr   r   r   r   r   c                   s2   t  j|i | t|| _tdddd| _d S )NT)prompts_changedresources_changedtools_changed)r   r   r   r   r   r	   notification_optionsr   r   r!   r"   r   Z   s   zLowLevelServer.__init__r#   c                 C  r$   r%   r&   r(   r!   r!   r"   r   f   r)   zLowLevelServer.fastmcpNrQ   NotificationOptions | Noneexperimental_capabilities dict[str, dict[str, Any]] | Noner   c                   s&   |d u r| j }t jd||d|S )N)rQ   rS   r!   )rQ   r   create_initialization_options)r   rQ   rS   r   r   r!   r"   rU   n   s   z,LowLevelServer.create_initialization_optionsFread_stream5MemoryObjectReceiveStream[SessionMessage | Exception]write_stream&MemoryObjectSendStream[SessionMessage]initialization_optionsraise_exceptionsbool	statelessc              
     s   t  4 I dH Z}|| | I dH }|t| j||||dI dH }t 4 I dH  }	|j2 z3 dH W }
|	| j	|
||| q06 W d  I dH  n1 I dH sRw   Y  W d  I dH  dS 1 I dH shw   Y  dS )zN
        Overrides the run method to use the MiddlewareServerSession.
        N)r]   )
r   enter_async_contextlifespanr   r   anyiocreate_task_groupincoming_messages
start_soon_handle_message)r   rV   rX   rZ   r[   r]   stacklifespan_contextsessiontgr4   r!   r!   r"   run}   s0   

*.zLowLevelServer.run)r   r   r   r   r   r   rF   )NN)rQ   rR   rS   rT   r   r   r#   r   )FF)
rV   rW   rX   rY   rZ   r   r[   r\   r]   r\   )	rG   rH   rI   r   rK   r   rU   ri   rL   r!   r!   r   r"   rM   Y   s    	rM   )&
__future__r   r   
contextlibr   typingr   r   r`   	mcp.typesr=   anyio.streams.memoryr   r   mcp.server.lowlevel.serverr   r	   r
   r   _Servermcp.server.modelsr   mcp.server.sessionr   mcp.server.stdior   mcp.shared.messager   mcp.shared.sessionr   fastmcp.utilities.loggingr   fastmcp.server.serverr   rG   loggerr   rM   r!   r!   r!   r"   <module>   s(    9