Fix an issue where LLM requests would block forever (#18830)

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
This commit is contained in:
Mikayla Maki 2024-10-07 13:03:26 -07:00 committed by GitHub
parent 8cdb9d6b85
commit 5387a6f7f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 26 additions and 4 deletions

View file

@ -3,8 +3,7 @@ name: Publish Collab Server Image
on: on:
push: push:
tags: tags:
# Pause production deploys while we investigate an issue. - collab-production
# - collab-production
- collab-staging - collab-staging
env: env:

View file

@ -163,8 +163,12 @@ impl futures::stream::Stream for WrappedBody {
WrappedBodyInner::SyncReader(cursor) => { WrappedBodyInner::SyncReader(cursor) => {
let mut buf = Vec::new(); let mut buf = Vec::new();
match cursor.read_to_end(&mut buf) { match cursor.read_to_end(&mut buf) {
Ok(_) => { Ok(bytes) => {
return Poll::Ready(Some(Ok(Bytes::from(buf)))); if bytes == 0 {
return Poll::Ready(None);
} else {
return Poll::Ready(Some(Ok(Bytes::from(buf))));
}
} }
Err(e) => return Poll::Ready(Some(Err(e))), Err(e) => return Poll::Ready(Some(Err(e))),
} }
@ -234,3 +238,22 @@ impl http_client::HttpClient for ReqwestClient {
.boxed() .boxed()
} }
} }
#[cfg(test)]
mod test {
use core::str;
use http_client::AsyncBody;
use smol::stream::StreamExt;
use crate::WrappedBody;
#[tokio::test]
async fn test_sync_streaming_upload() {
let mut body = WrappedBody::new(AsyncBody::from("hello there".to_string())).fuse();
let result = body.next().await.unwrap().unwrap();
assert!(body.next().await.is_none());
assert_eq!(str::from_utf8(&result).unwrap(), "hello there");
}
}