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:
push:
tags:
# Pause production deploys while we investigate an issue.
# - collab-production
- collab-production
- collab-staging
env:

View file

@ -163,8 +163,12 @@ impl futures::stream::Stream for WrappedBody {
WrappedBodyInner::SyncReader(cursor) => {
let mut buf = Vec::new();
match cursor.read_to_end(&mut buf) {
Ok(_) => {
return Poll::Ready(Some(Ok(Bytes::from(buf))));
Ok(bytes) => {
if bytes == 0 {
return Poll::Ready(None);
} else {
return Poll::Ready(Some(Ok(Bytes::from(buf))));
}
}
Err(e) => return Poll::Ready(Some(Err(e))),
}
@ -234,3 +238,22 @@ impl http_client::HttpClient for ReqwestClient {
.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");
}
}