
    .hC                     z   d Z ddl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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mZ ddlmZmZmZm Z m!Z!m"Z"m#Z#  ej$        e%          Z&ee'z  Z(ee(         Z)ee         Z*eg e+dz  f         Z,dZ-dZ.dZ/dZ0dZ1dZ2 G d de'          Z3 G d de3          Z4e	 G d d                      Z5 G d d          Z6ed ed           ed          dfde+d e7e+ef         dz  d!ed"ed#e8d$ee9eee'z           ee         e,f         df         fd%            Z:dS )&z
StreamableHTTP Client Transport Module

This module implements the StreamableHTTP transport for MCP clients,
providing support for HTTP POST requests with optional SSE streaming responses
and session management.
    N)AsyncGenerator	AwaitableCallable)asynccontextmanager)	dataclass)	timedelta)Any)MemoryObjectReceiveStreamMemoryObjectSendStream)EventSourceServerSentEventaconnect_sse)create_mcp_http_client)ClientMessageMetadataSessionMessage)	ErrorDataJSONRPCErrorJSONRPCMessageJSONRPCNotificationJSONRPCRequestJSONRPCResponse	RequestIdzmcp-session-idzlast-event-idzcontent-typeAcceptzapplication/jsonztext/event-streamc                       e Zd ZdZdS )StreamableHTTPErrorz3Base exception for StreamableHTTP transport errors.N__name__
__module____qualname____doc__     W/var/www/FlaskApp/flask-venv/lib/python3.11/site-packages/mcp/client/streamable_http.pyr   r   3   s        ==Dr"   r   c                       e Zd ZdZdS )ResumptionErrorz*Raised when resumption request is invalid.Nr   r!   r"   r#   r%   r%   9   s        44Dr"   r%   c                       e Zd ZU dZej        ed<   eeef         ed<   edz  ed<   e	ed<   e
dz  ed<   eed<   eed	<   dS )
RequestContextz Context for a request operation.clientheadersN
session_idsession_messagemetadataread_stream_writersse_read_timeout)r   r   r   r    httpxAsyncClient__annotations__dictstrr   r   StreamWriterr   r!   r"   r#   r'   r'   ?   s         **#s(^d
#####d****$$$$r"   r'   c                   >   e Zd ZdZd ed           ed          fdedeeef         dz  ded	ed
df
dZdeeef         d
eeef         fdZ	de
d
efdZde
d
efdZdej        d
dfdZ	 	 d)dedededz  deeged         f         dz  d
ef
dZdej        ded
dfdZded
dfdZded
dfdZdej        ded
dfdZdej        ded
dfdZdeded
dfd Zded!ed
dfd"Zdej        d#e ded$e!e"         d%eg df         d
dfd&Z#dej        d
dfd'Z$d
edz  fd(Z%dS )*StreamableHTTPTransportz/StreamableHTTP client transport implementation.N   seconds,  urlr)   timeoutr.   returnc                     || _         |pi | _        || _        || _        d| _        t
          t           dt           t          t          i| j        | _	        dS )a  Initialize the StreamableHTTP transport.

        Args:
            url: The endpoint URL.
            headers: Optional headers to include in requests.
            timeout: HTTP timeout for regular operations.
            sse_read_timeout: Timeout for SSE read operations.
        Nz, )
r;   r)   r<   r.   r*   ACCEPTJSONSSECONTENT_TYPErequest_headers)selfr;   r)   r<   r.   s        r#   __init__z StreamableHTTPTransport.__init__O   s`     }" 0&*t$$s$$$ 
 l 
r"   base_headersc                 Z    |                                 }| j        r| j        |t          <   |S )z,Update headers with session ID if available.)copyr*   MCP_SESSION_ID)rD   rF   r)   s      r#   _update_headers_with_sessionz4StreamableHTTPTransport._update_headers_with_sessioni   s0     ##%%? 	6&*oGN#r"   messagec                 V    t          |j        t                    o|j        j        dk    S )z2Check if the message is an initialization request.
initialize)
isinstancerootr   methodrD   rK   s     r#   _is_initialization_requestz2StreamableHTTPTransport._is_initialization_requestr   s*     w|^44 4#|3	
r"   c                 V    t          |j        t                    o|j        j        dk    S )z4Check if the message is an initialized notification.znotifications/initialized)rN   rO   r   rP   rQ   s     r#   _is_initialized_notificationz4StreamableHTTPTransport._is_initialized_notificationy   s-     w|%899 C#'BB	
r"   responsec                     |j                             t                    }|r+|| _        t                              d| j                    dS dS )z3Extract and store session ID from response headers.zReceived session ID: N)r)   getrI   r*   loggerinfo)rD   rU   new_session_ids      r#   '_maybe_extract_session_id_from_responsez?StreamableHTTPTransport._maybe_extract_session_id_from_response   sY    
 ")--n== 	C,DOKKAAABBBBB	C 	Cr"   sser-   original_request_idresumption_callbackc                   K   |j         dk    r"	 t          j        |j                  }t                              d|            |.t          |j        t          t          z            r||j        _
        t          |          }|                    |           d{V  |j
        r|r ||j
                   d{V  t          |j        t          t          z            S # t          $ rC}t                              d|            |                    |           d{V  Y d}~dS d}~ww xY wt                              d|j                     dS )z@Handle an SSE event, returning True if the response is complete.rK   zSSE message: NzError parsing SSE message: FzUnknown SSE event: )eventr   model_validate_jsondatarX   debugrN   rO   r   r   idr   send	Exceptionerrorwarning)rD   r\   r-   r]   r^   rK   r+   excs           r#   _handle_sse_eventz)StreamableHTTPTransport._handle_sse_event   s      9	!!(<SXFF6W66777 '2zL/L"@8 82 ':GLO"0"9"9(--o>>>>>>>>> 6 61 6--cf555555555 "',,0NOOO   @3@@AAA(--c222222222uuuuu
 NN<<<===5s   CC! !
D.+8D))D.r(   c           
      p  K   	 | j         sdS |                     | j                  }t          |d| j        |t          j        | j        j        | j	        j                            4 d{V 	 }|j
                                         t                              d           |                                2 3 d{V }|                     ||           d{V  $6 	 ddd          d{V  dS # 1 d{V swxY w Y   dS # t           $ r(}t                              d|            Y d}~dS d}~ww xY w)z0Handle GET stream for server-initiated messages.NGETreadr)   r<   zGET SSE connection establishedzGET stream error (non-fatal): )r*   rJ   rC   r   r;   r/   Timeoutr<   r9   r.   rU   raise_for_statusrX   rc   	aiter_sserj   rf   )rD   r(   r-   r)   event_sourcer\   ri   s          r#   handle_get_streamz)StreamableHTTPTransport.handle_get_stream   sJ     	A? 778LMMG#L(t/D/L     J J J J J J J J %66888=>>>!-!7!7!9!9 J J J J J J J#006HIIIIIIIIII ":!9J J J J J J J J J J J J J J J J J J J J J J J J J J J J J J  	A 	A 	ALL?#??@@@@@@@@@	AsN   D A!D /AC07C=C0D 0
C::D =C:>D 
D5D00D5ctxc           
        K   |                      |j                  }|j        r!|j        j        r|j        j        |t          <   nt          d          d}t          |j        j        j	        t                    r|j        j        j	        j        }t          |j        d| j        |t          j        | j        j        |j        j                            4 d{V 	 }|j                                         t,                              d           |                                2 3 d{V }|                     ||j        ||j        r|j        j        nd           d{V }|r nB6 	 ddd          d{V  dS # 1 d{V swxY w Y   dS )z/Handle a resumption request using GET with SSE.z.Resumption request requires a resumption tokenNrl   rm   ro   z)Resumption GET SSE connection established)rJ   r)   r,   resumption_tokenLAST_EVENT_IDr%   rN   r+   rK   rO   r   rd   r   r(   r;   r/   rp   r<   r9   r.   rU   rq   rX   rc   rr   rj   r-   on_resumption_token_update)rD   ru   r)   r]   rs   r\   is_completes          r#   _handle_resumption_requestz2StreamableHTTPTransport._handle_resumption_request   sb     33CK@@< 	TCL9 	T%(\%BGM""!"RSSS #c)16GG 	F"%"5"="B"EJHM$3+?+G  
 
 
 	 	 	 	 	 	 	 	 !22444LLDEEE)3355       c$($:$:*'?B|UCL;;QU	% %        E 65	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   AE<%E'+=E<<
F	Fc           	      f  K   |                      |j                  }|j        j        }|                     |          }|j                            d| j        |                    ddd          |          4 d{V 	 }|j	        dk    r.t                              d           	 ddd          d{V  dS |j	        d	k    rYt          |j        t                    r+|                     |j        |j        j                   d{V  	 ddd          d{V  dS |                                 |r|                     |           |j                            t*          d
                                          }|                    t0                    r"|                     ||j                   d{V  nX|                    t4                    r|                     ||           d{V  n!|                     ||j                   d{V  ddd          d{V  dS # 1 d{V swxY w Y   dS )z0Handle a  POST request with response processing.POSTTjson)by_aliasmodeexclude_none)r~   r)   N   zReceived 202 Acceptedi   )rJ   r)   r+   rK   rR   r(   streamr;   
model_dumpstatus_coderX   rc   rN   rO   r   _send_session_terminated_errorr-   rd   rq   r[   rW   rB   lower
startswithr@   _handle_json_responserA   _handle_sse_response_handle_unexpected_content_type)rD   ru   r)   rK   is_initializationrU   content_types          r#   _handle_post_requestz,StreamableHTTPTransport._handle_post_request   sr     33CK@@%- ;;GDD:$$H##TT#RR	 % 
 
  	  	  	  	  	  	  	  	
 #s**4555 	  	  	  	  	  	  	  	  	  	  	  	  	  	 #s**glN;; ==.         ! 	  	  	  	  	  	  	  	  	  	  	  	  	  	$ %%'''  G<<XFFF#+//bAAGGIIL&&t,, 003;QRRRRRRRRRR((-- //#>>>>>>>>>>:: *        ; 	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	  	s    >'H 8AH C1H  
H*-H*c                 \  K   	 |                                  d{V }t          j        |          }t          |          }|                    |           d{V  dS # t
          $ rC}t                              d|            |                    |           d{V  Y d}~dS d}~ww xY w)z%Handle JSON response from the server.NzError parsing JSON response: )areadr   ra   r   re   rf   rX   rg   )rD   rU   r-   contentrK   r+   ri   s          r#   r   z-StreamableHTTPTransport._handle_json_response  s      	/$NN,,,,,,,,G$8AAG,W55O$))/::::::::::: 	/ 	/ 	/LL>>>???$))#...............	/s   AA 
B+(8B&&B+c                 t  K   	 t          |          }|                                2 3 d{V }|                     ||j        |j        r|j        j        nd           d{V  >6 dS # t          $ rE}t                              d           |j        	                    |           d{V  Y d}~dS d}~ww xY w)z$Handle SSE response from the server.N)r^   zError reading SSE stream:)
r   rr   rj   r-   r,   ry   rf   rX   	exceptionre   )rD   rU   ru   rs   r\   es         r#   r   z,StreamableHTTPTransport._handle_sse_response(  s*     	1&x00L)3355 	 	 	 	 	 	 	c,,* <"??! -           655  	1 	1 	18999(--a000000000000000	1s!   #A( A%9A( (
B72:B22B7r   c                    K   d| }t                               |           |                    t          |                     d{V  dS )z+Handle unexpected content type in response.zUnexpected content type: N)rX   rg   re   
ValueError)rD   r   r-   	error_msgs       r#   r   z7StreamableHTTPTransport._handle_unexpected_content_type<  s[       ?>>	Y %%j&;&;<<<<<<<<<<<r"   
request_idc                    K   t          d|t          dd                    }t          t          |                    }|                    |           d{V  dS )z)Send a session terminated error response.z2.0iX  zSession terminated)coderK   )jsonrpcrd   rg   N)r   r   r   r   re   )rD   r-   r   jsonrpc_errorr+   s        r#   r   z6StreamableHTTPTransport._send_session_terminated_errorF  sv       %0DEEE
 
 

 ))F)FGG %%o66666666666r"   write_stream_readerwrite_streamstart_get_streamc                 t  K   	 |4 d{V  |2 3 d{V }|j         }t          |j        t                    r|j        nd}t	          |o|j                  }	t                              d|            |                     |          r
 |             t          || j
        | j        |||| j                  }
|	r|                     |
           d{V  |                     |
           d{V  6 	 ddd          d{V  n# 1 d{V swxY w Y   n4# t          $ r'}t                              d|            Y d}~nd}~ww xY w|                                 d{V  |                                 d{V  dS # |                                 d{V  |                                 d{V  w xY w)z&Handle writing requests to the server.NzSending client message: )r(   r)   r*   r+   r,   r-   r.   zError in post_writer: )rK   rN   r,   r   boolrw   rX   rc   rT   r'   rC   r*   r.   r{   r   rf   rg   aclose)rD   r(   r   r-   r   r   r+   rK   r,   is_resumptionru   ri   s               r#   post_writerz#StreamableHTTPTransport.post_writerT  s     &	(* = = = = = = = =-@ = = = = = = =/-5G &o&>@UVV"00!  %))Oh6O$P$PMLL!EG!E!EFFF 88AA +((***(% $ 4#'?(7!)+=)-)>  C % ="==cBBBBBBBBBB"77<<<<<<<<<<= .A-@= = = = = = = = = = = = = = = = = = = = = = = = = = =B  	9 	9 	9LL7#7788888888	9 %++---------%%''''''''''' %++---------%%''''''''''s]   D DC1CD3D 
DD DD F 
E!E>F EF 6F7c                   K   | j         sdS 	 |                     | j                  }|                    | j        |           d{V }|j        dk    rt                              d           dS |j        dk    r$t                              d|j                    dS dS # t          $ r(}t                              d|            Y d}~dS d}~ww xY w)z2Terminate the session by sending a DELETE request.N)r)   i  z)Server does not allow session termination   zSession termination failed: )
r*   rJ   rC   deleter;   r   rX   rc   rh   rf   )rD   r(   r)   rU   ri   s        r#   terminate_sessionz)StreamableHTTPTransport.terminate_session  s      	F		A778LMMG#]]48W]EEEEEEEEH#s**HIIIII%,,Th>RTTUUUUU -, 	A 	A 	ANN?#??@@@@@@@@@	As   A!B! 0-B! !
C+CCc                     | j         S )zGet the current session ID.)r*   )rD   s    r#   get_session_idz&StreamableHTTPTransport.get_session_id  s
    r"   )NN)&r   r   r   r    r   r3   r2   r	   rE   rJ   r   r   rR   rT   r/   Responser[   r   r4   r   r   r   rj   r0   rt   r'   r{   r   r   r   r   r   StreamReaderr   r   r   r   r   r!   r"   r#   r6   r6   L   s       99
 *.&Yr222&/i&?&?&?
 

 c3h$&
 	

 $
 

 
 
 
4 cN	c3h   
. 
T 
 
 
 

N 
t 
 
 
 
C.C 
C C C C 15GK$ $$ )$ '-	$
 &seYt_&<=D$ 
$ $ $ $LA!A )A 
	A A A A<!N !t ! ! ! !F&n & & & & &P/./ )/ 
	/ / / /11-;1	1 1 1 1(== )= 
	= = = =7(7 7 
	7 7 7 7/(!/( */( )	/(
 -^</( #2t8,/( 
/( /( /( /(bAe.? AD A A A A d
      r"   r6   r7   r8   r:   Tr;   r)   r<   r.   terminate_on_closer=   c           
       	
K   t          | |||          t          j        t          t          z           d          \  
}t          j        t                   d          \  }}t          j                    4 d{V 	 t                              d|             t          j	        t          j        j        j        j        j                            4 d{V 	 	d	
fd}                    j        	|
||           	 ||j        fW V  j        r|r                    	           d{V  j                                         nB# j        r|r                    	           d{V  j                                         w xY w	 ddd          d{V  n# 1 d{V swxY w Y   
                                 d{V  |                                 d{V  n9# 
                                 d{V  |                                 d{V  w xY w	 ddd          d{V  dS # 1 d{V swxY w Y   dS )	a  
    Client transport for StreamableHTTP.

    `sse_read_timeout` determines how long (in seconds) the client will wait for a new
    event before disconnecting. All other HTTP operations are controlled by `timeout`.

    Yields:
        Tuple containing:
            - read_stream: Stream for reading messages from the server
            - write_stream: Stream for sending messages to the server
            - get_session_id_callback: Function to retrieve the current session ID
    r   Nz'Connecting to StreamableHTTP endpoint: rm   ro   r=   c                  @                         j                    d S )N)
start_soonrt   )r(   r-   tg	transports   r#   r   z/streamablehttp_client.<locals>.start_get_stream  s.    MM!3V=O    r"   )r=   N)r6   anyiocreate_memory_object_streamr   rf   create_task_grouprX   rY   r   rC   r/   rp   r<   r9   r.   r   r   r   r*   r   cancel_scopecancelr   )r;   r)   r<   r.   r   read_streamr   r   r   r(   r-   r   r   s            @@@@r#   streamablehttp_clientr     s#     6 (Wg?OPPI&+&G"''	 '	# ).(I))	 )	%L% &(( %( %( %( %( %( %( %(B$	(KKG#GGHHH-!1%-I4N4V     - - - - - - - -
         
 )'& $  	-#$!0     !+ B0B B'99&AAAAAAAAAO**,,,, !+ B0B B'99&AAAAAAAAAO**,,,,,=- - - - - - - - - - - - - - - - - - - - - - - - - - -@ %++---------%%'''''''''' %++---------%%'''''''''''K%( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %( %(sg   >H? A!G2")F,E>F,?FF,G2,
F6	6G29F6	:G2=5H?26H((H??
I	I	);r    loggingcollections.abcr   r   r   
contextlibr   dataclassesr   datetimer   typingr	   r   r/   anyio.streams.memoryr
   r   	httpx_sser   r   r   mcp.shared._httpx_utilsr   mcp.shared.messager   r   	mcp.typesr   r   r   r   r   r   r   	getLoggerr   rX   rf   SessionMessageOrErrorr4   r   r3   GetSessionIdCallbackrI   rx   rB   r?   r@   rA   r   r%   r'   r6   r2   r   tupler   r!   r"   r#   <module>r      sM     ? ? ? ? ? ? ? ? ? ? * * * * * * ! ! ! ! ! !               R R R R R R R R @ @ @ @ @ @ @ @ @ @ : : : : : : D D D D D D D D                  
	8	$	$ '2 %&;<(8C$J/ !	 	 	 	 	 	) 	 	 		 	 	 	 	) 	 	 	 	  	  	  	  	  	  	  	 K K K K K K K K\
  &*"2..."+)F";";";#H( H(	H(#s(^d"H( H(  	H(
 H( 	!.9"<=~.	
 	
H( H( H( H( H( H(r"   