fix: don't block on creating the SSE stream (#553)

This commit is contained in:
Jack Amadeo 2025-11-19 20:20:00 -05:00 • committed by GitHub
commit da57ee7697
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -391,47 +391,51 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
} }
let mut streams = tokio::task::JoinSet::new(); let mut streams = tokio::task::JoinSet::new();
if let Some(session_id) = &session_id { if let Some(session_id) = &session_id {
match self let client = self.client.clone();
.client let uri = config.uri.clone();
.get_stream( let session_id = session_id.clone();
config.uri.clone(), let auth_header = config.auth_header.clone();
session_id.clone(), let retry_config = self.config.retry_config.clone();
None, let sse_worker_tx = sse_worker_tx.clone();
config.auth_header.clone(), let transport_task_ct = transport_task_ct.clone();
) let config_uri = config.uri.clone();
let config_auth_header = config.auth_header.clone();
streams.spawn(async move {
match client
.get_stream(uri.clone(), session_id.clone(), None, auth_header.clone())
.await .await
{ {
Ok(stream) => { Ok(stream) => {
let sse_stream = SseAutoReconnectStream::new( let sse_stream = SseAutoReconnectStream::new(
stream, stream,
StreamableHttpClientReconnect { StreamableHttpClientReconnect {
client: self.client.clone(), client: client.clone(),
session_id: session_id.clone(), session_id: session_id.clone(),
uri: config.uri.clone(), uri: config_uri,
auth_header: config.auth_header.clone(), auth_header: config_auth_header,
}, },
self.config.retry_config.clone(), retry_config,
); );
streams.spawn(Self::execute_sse_stream( Self::execute_sse_stream(
sse_stream, sse_stream,
sse_worker_tx.clone(), sse_worker_tx,
false, false,
transport_task_ct.child_token(), transport_task_ct.child_token(),
)); )
tracing::debug!("got common stream"); .await
} }
Err(StreamableHttpError::ServerDoesNotSupportSse) => { Err(StreamableHttpError::ServerDoesNotSupportSse) => {
tracing::debug!("server doesn't support sse, skip common stream"); tracing::debug!("server doesn't support sse, skip common stream");
Ok(())
} }
Err(e) => { Err(e) => {
// fail to get common stream // fail to get common stream
tracing::error!("fail to get common stream: {e}"); tracing::error!("fail to get common stream: {e}");
return Err(WorkerQuitReason::fatal( Err(e)
e,
"get general purpose event stream",
));
} }
} }
});
} }
loop { loop {
let event = tokio::select! { let event = tokio::select! {