From 55503b23d4cfcef7ad2dd0acf409143832ec01d2 Mon Sep 17 00:00:00 2001 From: Krish Chainani Date: Wed, 28 Oct 2026 20:53:53 +0000 Subject: [PATCH] Limit concurrent image availability checks to eight (#57835) Replace `join_all` with a buffered stream in `replace_unresolvable_image_files` so at most eight distinct image file IDs are checked concurrently. Collect unavailable IDs as checks complete before replacing their image content with placeholders. GitOrigin-RevId: 2a9ed8742fbb6157c99c69a15f212df3ed88810c --- .../src/context_manager/image_availability.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/codex-rs/core/src/context_manager/image_availability.rs b/codex-rs/core/src/context_manager/image_availability.rs index 32754170de..c831fd22d1 100644 --- a/codex-rs/core/src/context_manager/image_availability.rs +++ b/codex-rs/core/src/context_manager/image_availability.rs @@ -12,7 +12,10 @@ use codex_protocol::models::ContentItemMetadata; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageReference; use codex_protocol::models::ResponseItem; -use futures::future::join_all; +use futures::StreamExt; +use futures::stream; + +const MAX_CONCURRENT_IMAGE_AVAILABILITY_CHECKS: usize = 8; /// Checks distinct file IDs concurrently and replaces only unavailable image content. pub(crate) async fn replace_unresolvable_image_files<'a>( @@ -51,13 +54,14 @@ pub(crate) async fn replace_unresolvable_image_files<'a>( _ => {} } } - let results = join_all(file_ids.into_iter().map(|file_id| async move { - let resolvable = image_store.is_resolvable(&file_id).await; - (file_id, resolvable) - })) - .await; + let mut checks = stream::iter(file_ids) + .map(|file_id| async move { + let resolvable = image_store.is_resolvable(&file_id).await; + (file_id, resolvable) + }) + .buffer_unordered(MAX_CONCURRENT_IMAGE_AVAILABILITY_CHECKS); let mut unavailable_file_ids = HashSet::new(); - for (file_id, resolvable) in results { + while let Some((file_id, resolvable)) = checks.next().await { if !resolvable { unavailable_file_ids.insert(file_id); }