Route image extension reads through turn environments v2 (#27498)

## Why

Image generation used `std::fs::read` for referenced image paths, which
did not support environment-backed filesystems or their sandbox context.

## What changed

- Expose optional turn environments to extension tool calls.
- Include each environment’s ID, working directory, filesystem, and
sandbox context.
- Read referenced images through the selected environment filesystem.
- Keep sandbox usage at the extension call site so extensions can choose
the appropriate access mode.
- Consolidate image request construction into one async function.
- Add coverage for successful environment reads and read failures.

## Validation

- `cargo check -p codex-image-generation-extension --tests`
- `just fmt`
- `just bazel-lock-update`
- `just bazel-lock-check`

`just test -p codex-image-generation-extension` could not complete
because the build exhausted available disk space.
This commit is contained in:
Won Park authored and GitHub committed 2026-06-11 16:32:52 -07:00
1 parent b2cc01fc7c
commit 19ce6394af
14 files changed
+423 -34

No files matched your search

+32 -12
View File
@@ -10,6 +10,7 @@ use codex_core::image_generation_artifact_path;
use codex_extension_api::ExtensionTurnItem;
use codex_extension_api::FunctionCallError;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolEnvironment;
use codex_extension_api::ToolExecutor;
use codex_extension_api::ToolName;
use codex_extension_api::ToolOutput;
@@ -103,7 +104,9 @@ impl ToolExecutor<ToolCall> for ImageGenerationTool {
impl ImageGenerationTool {
async fn handle_call(&self, call: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
let args = parse_args(&call)?;
let request = request_for_args(&args, call.conversation_history.items())?;
let request =
request_for_call_args(&args, call.conversation_history.items(), &call.environments)
.await?;
call.turn_item_emitter
.emit_started(ExtensionTurnItem::ImageGeneration(ImageGenerationItem {
id: call.call_id.clone(),
@@ -154,10 +157,10 @@ enum ImageRequest {
Edit(ImageEditRequest),
}
/// Builds a generation or edit request from the mutually exclusive image selectors.
fn request_for_args(
async fn request_for_call_args(
args: &ImagegenArgs,
history: &[ResponseItem],
environments: &[ToolEnvironment],
) -> Result<ImageRequest, FunctionCallError> {
let paths = args.referenced_image_paths.as_deref().unwrap_or_default();
if paths.len() > MAX_EDIT_IMAGES {
@@ -176,7 +179,18 @@ fn request_for_args(
size: Some("auto".to_string()),
}));
}
(false, None) => paths.iter().map(image_url).collect::<Result<Vec<_>, _>>()?,
(false, None) => {
let Some(environment) = environments.first() else {
return Err(FunctionCallError::RespondToModel(
"referenced image paths are unavailable in this session".to_string(),
));
};
let mut images = Vec::with_capacity(paths.len());
for path in paths {
images.push(image_url(path, environment).await?);
}
images
}
(true, Some(count)) => {
if !(1..=MAX_EDIT_IMAGES).contains(&count) {
return Err(FunctionCallError::RespondToModel(format!(
@@ -214,7 +228,6 @@ fn request_for_args(
}))
}
/// Selects the newest requested images while preserving their conversation order.
fn recent_images(history: &[ResponseItem], count: usize) -> Vec<ImageUrl> {
let mut function_call_ids = HashSet::new();
let mut custom_tool_call_ids = HashSet::new();
@@ -307,13 +320,20 @@ fn output_image_urls(output: &FunctionCallOutputPayload) -> impl Iterator<Item =
})
}
fn image_url(path: &AbsolutePathBuf) -> Result<ImageUrl, FunctionCallError> {
let bytes = std::fs::read(path).map_err(|error| {
FunctionCallError::RespondToModel(format!(
"unable to read referenced image at `{}`: {error}",
path.display()
))
})?;
async fn image_url(
path: &AbsolutePathBuf,
environment: &ToolEnvironment,
) -> Result<ImageUrl, FunctionCallError> {
let bytes = environment
.file_system
.read_file(path, Some(&environment.file_system_sandbox_context))
.await
.map_err(|error| {
FunctionCallError::RespondToModel(format!(
"unable to read referenced image at `{}`: {error}",
path.display()
))
})?;
let image = load_for_prompt_bytes(path.as_path(), bytes, PromptImageMode::Original).map_err(
|error| {
FunctionCallError::RespondToModel(format!(