feat: implement streaming for item content with offset and length support
Co-authored-by: aider (openai/andrew/openrouter/qwen/qwen3-coder) <aider@aider.chat>
This commit is contained in:
@@ -56,20 +56,47 @@ impl AsyncItemService {
|
|||||||
&self,
|
&self,
|
||||||
item_id: i64,
|
item_id: i64,
|
||||||
allow_binary: bool,
|
allow_binary: bool,
|
||||||
_offset: u64,
|
offset: u64,
|
||||||
length: u64,
|
length: u64,
|
||||||
) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
|
) -> Result<(std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<tokio_util::bytes::Bytes, std::io::Error>> + Send>>, String), CoreError> {
|
||||||
// Get the item content synchronously
|
let data_dir = self.data_dir.clone();
|
||||||
let item_with_content = self.get_item_content(item_id).await?;
|
let db = self.db.clone();
|
||||||
|
|
||||||
|
// Get item metadata first to check binary status and get MIME type
|
||||||
|
let item_with_meta = tokio::task::spawn_blocking(move || {
|
||||||
|
let conn = db.blocking_lock();
|
||||||
|
let item_service = ItemService::new(data_dir);
|
||||||
|
item_service.get_item(&conn, item_id)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap()?;
|
||||||
|
|
||||||
|
let metadata = item_with_meta.meta_as_map();
|
||||||
|
let mime_type = metadata
|
||||||
|
.get("mime_type")
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.unwrap_or_else(|| "application/octet-stream".to_string());
|
||||||
|
|
||||||
// Check if content is binary when allow_binary is false
|
// Check if content is binary when allow_binary is false
|
||||||
let metadata = item_with_content.item_with_meta.meta_as_map();
|
|
||||||
if !allow_binary {
|
if !allow_binary {
|
||||||
let is_content_binary = if let Some(binary_val) = metadata.get("binary") {
|
let is_content_binary = if let Some(binary_val) = metadata.get("binary") {
|
||||||
binary_val == "true"
|
binary_val == "true"
|
||||||
} else {
|
} else {
|
||||||
// Fallback to checking the actual content
|
// For the binary check, we need to read a sample of the content
|
||||||
crate::common::is_binary::is_binary(&item_with_content.content)
|
// We'll read the first 8192 bytes to determine if it's binary
|
||||||
|
let data_dir = self.data_dir.clone();
|
||||||
|
let db = self.db.clone();
|
||||||
|
|
||||||
|
let sample_content = tokio::task::spawn_blocking(move || {
|
||||||
|
let conn = db.blocking_lock();
|
||||||
|
let item_service = ItemService::new(data_dir);
|
||||||
|
let item_with_content = item_service.get_item_content(&conn, item_id)?;
|
||||||
|
Ok::<Vec<u8>, CoreError>(item_with_content.content)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap()?;
|
||||||
|
|
||||||
|
crate::common::is_binary::is_binary(&sample_content)
|
||||||
};
|
};
|
||||||
|
|
||||||
if is_content_binary {
|
if is_content_binary {
|
||||||
@@ -77,17 +104,43 @@ impl AsyncItemService {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let mime_type = metadata
|
// Create a stream that reads only the requested portion
|
||||||
.get("mime_type")
|
let data_dir = self.data_dir.clone();
|
||||||
.map(|s| s.to_string())
|
let db = self.db.clone();
|
||||||
.unwrap_or_else(|| "application/octet-stream".to_string());
|
|
||||||
|
|
||||||
// Convert content to stream with length limit
|
let stream = Box::pin(tokio_stream::wrappers::ReceiverStream::new({
|
||||||
let mut content = item_with_content.content;
|
let (tx, rx) = tokio::sync::mpsc::channel(1);
|
||||||
if length > 0 && length < content.len() as u64 {
|
|
||||||
content.truncate(length as usize);
|
tokio::task::spawn_blocking(move || {
|
||||||
|
let conn = db.blocking_lock();
|
||||||
|
let item_service = ItemService::new(data_dir);
|
||||||
|
|
||||||
|
match item_service.get_item_content(&conn, item_id) {
|
||||||
|
Ok(item_with_content) => {
|
||||||
|
let content = item_with_content.content;
|
||||||
|
let content_len = content.len() as u64;
|
||||||
|
|
||||||
|
// Apply offset and length constraints
|
||||||
|
let start = std::cmp::min(offset, content_len);
|
||||||
|
let end = if length > 0 {
|
||||||
|
std::cmp::min(start + length, content_len)
|
||||||
|
} else {
|
||||||
|
content_len
|
||||||
|
};
|
||||||
|
|
||||||
|
if start < content_len {
|
||||||
|
let chunk = tokio_util::bytes::Bytes::from(content[start as usize..end as usize].to_vec());
|
||||||
|
let _ = tx.blocking_send(Ok(chunk));
|
||||||
}
|
}
|
||||||
let stream = Box::pin(tokio_stream::iter(vec![Ok(tokio_util::bytes::Bytes::from(content))]));
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let _ = tx.blocking_send(Err(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
rx
|
||||||
|
}));
|
||||||
|
|
||||||
Ok((stream, mime_type))
|
Ok((stream, mime_type))
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user