unifying all image saves to /tmp to bug-proof (#14149)

image-gen feature will have the model saving to /tmp by default + at all
times
This commit is contained in:
Won Park 2026-03-10 15:13:12 -07:00 committed by Michael Bolin
parent 91ca20c7c3
commit 722e8f08e1
8 changed files with 214 additions and 82 deletions

View file

@ -6811,8 +6811,7 @@ async fn handle_assistant_item_done_in_plan_mode(
{
maybe_complete_plan_item_from_message(sess, turn_context, state, item).await;
if let Some(turn_item) =
handle_non_tool_response_item(item, true, Some(&turn_context.cwd)).await
if let Some(turn_item) = handle_non_tool_response_item(sess, turn_context, item, true).await
{
emit_turn_item_in_plan_mode(
sess,
@ -6993,8 +6992,13 @@ async fn try_run_sampling_request(
needs_follow_up |= output_result.needs_follow_up;
}
ResponseEvent::OutputItemAdded(item) => {
if let Some(turn_item) =
handle_non_tool_response_item(&item, plan_mode, Some(&turn_context.cwd)).await
if let Some(turn_item) = handle_non_tool_response_item(
sess.as_ref(),
turn_context.as_ref(),
&item,
plan_mode,
)
.await
{
let mut turn_item = turn_item;
let mut seeded_parsed: Option<ParsedAssistantTextDelta> = None;

View file

@ -154,6 +154,26 @@ fn developer_input_texts(items: &[ResponseItem]) -> Vec<&str> {
.collect()
}
fn default_image_save_developer_message_text() -> String {
let image_output_dir = crate::stream_events_utils::default_image_generation_output_dir();
format!(
"Generated images are saved to {} as {} by default.",
image_output_dir.display(),
image_output_dir.join("<image_id>.png").display(),
)
}
fn test_tool_runtime(session: Arc<Session>, turn_context: Arc<TurnContext>) -> ToolCallRuntime {
let router = Arc::new(ToolRouter::from_config(
&turn_context.tools_config,
None,
None,
turn_context.dynamic_tools.as_slice(),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
ToolCallRuntime::new(router, session, turn_context, tracker)
}
fn make_connector(id: &str, name: &str) -> AppInfo {
AppInfo {
id: id.to_string(),
@ -3123,6 +3143,114 @@ async fn build_initial_context_uses_previous_realtime_state() {
);
}
#[tokio::test]
async fn build_initial_context_omits_default_image_save_location_with_image_history() {
let (session, turn_context) = make_session_and_context().await;
session
.replace_history(
vec![ResponseItem::ImageGenerationCall {
id: "ig-test".to_string(),
status: "completed".to_string(),
revised_prompt: Some("a tiny blue square".to_string()),
result: "Zm9v".to_string(),
}],
None,
)
.await;
let initial_context = session.build_initial_context(&turn_context).await;
let developer_texts = developer_input_texts(&initial_context);
assert!(
!developer_texts
.iter()
.any(|text| text.contains("Generated images are saved to")),
"expected initial context to omit image save instructions even with image history, got {developer_texts:?}"
);
}
#[tokio::test]
async fn build_initial_context_omits_default_image_save_location_without_image_history() {
let (session, turn_context) = make_session_and_context().await;
let initial_context = session.build_initial_context(&turn_context).await;
let developer_texts = developer_input_texts(&initial_context);
assert!(
!developer_texts
.iter()
.any(|text| text.contains("Generated images are saved to")),
"expected initial context to omit image save instructions without image history, got {developer_texts:?}"
);
}
#[tokio::test]
async fn handle_output_item_done_records_image_save_message_after_successful_save() {
let (session, turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let call_id = "ig_history_records_message";
let expected_saved_path = crate::stream_events_utils::default_image_generation_output_dir()
.join(format!("{call_id}.png"));
let _ = std::fs::remove_file(&expected_saved_path);
let item = ResponseItem::ImageGenerationCall {
id: call_id.to_string(),
status: "completed".to_string(),
revised_prompt: Some("a tiny blue square".to_string()),
result: "Zm9v".to_string(),
};
let mut ctx = HandleOutputCtx {
sess: Arc::clone(&session),
turn_context: Arc::clone(&turn_context),
tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)),
cancellation_token: CancellationToken::new(),
};
handle_output_item_done(&mut ctx, item.clone(), None)
.await
.expect("image generation item should succeed");
let history = session.clone_history().await;
let expected_message: ResponseItem =
DeveloperInstructions::new(default_image_save_developer_message_text()).into();
assert_eq!(history.raw_items(), &[expected_message, item]);
assert_eq!(
std::fs::read(&expected_saved_path).expect("saved file"),
b"foo"
);
let _ = std::fs::remove_file(&expected_saved_path);
}
#[tokio::test]
async fn handle_output_item_done_skips_image_save_message_when_save_fails() {
let (session, turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let call_id = "ig_history_no_message";
let expected_saved_path = crate::stream_events_utils::default_image_generation_output_dir()
.join(format!("{call_id}.png"));
let _ = std::fs::remove_file(&expected_saved_path);
let item = ResponseItem::ImageGenerationCall {
id: call_id.to_string(),
status: "completed".to_string(),
revised_prompt: Some("broken payload".to_string()),
result: "_-8".to_string(),
};
let mut ctx = HandleOutputCtx {
sess: Arc::clone(&session),
turn_context: Arc::clone(&turn_context),
tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)),
cancellation_token: CancellationToken::new(),
};
handle_output_item_done(&mut ctx, item.clone(), None)
.await
.expect("image generation item should still complete");
let history = session.clone_history().await;
assert_eq!(history.raw_items(), &[item]);
assert!(!expected_saved_path.exists());
}
#[tokio::test]
async fn build_initial_context_uses_previous_turn_settings_for_realtime_end() {
let (session, turn_context) = make_session_and_context().await;

View file

@ -434,9 +434,6 @@ fn for_prompt_rewrites_image_generation_calls_when_images_are_supported() {
ContentItem::InputImage {
image_url: "data:image/png;base64,Zm9v".to_string(),
},
ContentItem::InputText {
text: "Saved to: CWD".to_string(),
},
],
end_turn: None,
phase: None,
@ -503,9 +500,6 @@ fn for_prompt_rewrites_image_generation_calls_when_images_are_unsupported() {
text: "image content omitted because you do not support image input"
.to_string(),
},
ContentItem::InputText {
text: "Saved to: CWD".to_string(),
},
],
end_turn: None,
phase: None,

View file

@ -242,9 +242,6 @@ pub(crate) fn rewrite_image_generation_calls_for_stateless_input(items: &mut Vec
text: format!("Prompt: {revised_prompt}"),
},
ContentItem::InputImage { image_url },
ContentItem::InputText {
text: "Saved to: CWD".to_string(),
},
],
end_turn: None,
phase: None,

View file

@ -1,4 +1,3 @@
use std::path::Path;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
@ -20,6 +19,7 @@ use crate::parse_turn_item;
use crate::state_db;
use crate::tools::parallel::ToolCallRuntime;
use crate::tools::router::ToolRouter;
use codex_protocol::models::DeveloperInstructions;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
@ -54,11 +54,7 @@ pub(crate) fn raw_assistant_output_text_from_item(item: &ResponseItem) -> Option
None
}
async fn save_image_generation_result_to_cwd(
cwd: &Path,
call_id: &str,
result: &str,
) -> Result<PathBuf> {
async fn save_image_generation_result(call_id: &str, result: &str) -> Result<PathBuf> {
let bytes = BASE64_STANDARD
.decode(result.trim().as_bytes())
.map_err(|err| {
@ -77,11 +73,15 @@ async fn save_image_generation_result_to_cwd(
if file_stem.is_empty() {
file_stem = "generated_image".to_string();
}
let path = cwd.join(format!("{file_stem}.png"));
let path = default_image_generation_output_dir().join(format!("{file_stem}.png"));
tokio::fs::write(&path, bytes).await?;
Ok(path)
}
pub(crate) fn default_image_generation_output_dir() -> PathBuf {
std::env::temp_dir()
}
/// Persist a completed model response item and record any cited memory usage.
pub(crate) async fn record_completed_response_item(
sess: &Session,
@ -189,8 +189,13 @@ pub(crate) async fn handle_output_item_done(
}
// No tool call: convert messages/reasoning into turn items and mark them as complete.
Ok(None) => {
if let Some(turn_item) =
handle_non_tool_response_item(&item, plan_mode, Some(&ctx.turn_context.cwd)).await
if let Some(turn_item) = handle_non_tool_response_item(
ctx.sess.as_ref(),
ctx.turn_context.as_ref(),
&item,
plan_mode,
)
.await
{
if previously_active_item.is_none() {
let mut started_item = turn_item.clone();
@ -276,9 +281,10 @@ pub(crate) async fn handle_output_item_done(
}
pub(crate) async fn handle_non_tool_response_item(
sess: &Session,
turn_context: &TurnContext,
item: &ResponseItem,
plan_mode: bool,
image_output_cwd: Option<&Path>,
) -> Option<TurnItem> {
debug!(?item, "Output item");
@ -300,19 +306,28 @@ pub(crate) async fn handle_non_tool_response_item(
agent_message.content =
vec![codex_protocol::items::AgentMessageContent::Text { text: stripped }];
}
if let TurnItem::ImageGeneration(image_item) = &mut turn_item
&& let Some(cwd) = image_output_cwd
{
match save_image_generation_result_to_cwd(cwd, &image_item.id, &image_item.result)
.await
{
if let TurnItem::ImageGeneration(image_item) = &mut turn_item {
match save_image_generation_result(&image_item.id, &image_item.result).await {
Ok(path) => {
image_item.saved_path = Some(path.to_string_lossy().into_owned());
let image_output_dir = default_image_generation_output_dir();
let message: ResponseItem = DeveloperInstructions::new(format!(
"Generated images are saved to {} as {} by default.",
image_output_dir.display(),
image_output_dir.join("<image_id>.png").display(),
))
.into();
sess.record_conversation_items(
turn_context,
std::slice::from_ref(&message),
)
.await;
}
Err(err) => {
let output_dir = default_image_generation_output_dir();
tracing::warn!(
call_id = %image_item.id,
cwd = %cwd.display(),
output_dir = %output_dir.display(),
"failed to save generated image: {err}"
);
}
@ -372,15 +387,16 @@ pub(crate) fn response_input_to_response_item(input: &ResponseInputItem) -> Opti
#[cfg(test)]
mod tests {
use super::default_image_generation_output_dir;
use super::handle_non_tool_response_item;
use super::last_assistant_message_from_item;
use super::save_image_generation_result_to_cwd;
use super::save_image_generation_result;
use crate::codex::make_session_and_context;
use crate::error::CodexErr;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
fn assistant_output_text(text: &str) -> ResponseItem {
ResponseItem::Message {
@ -396,12 +412,12 @@ mod tests {
#[tokio::test]
async fn handle_non_tool_response_item_strips_citations_from_assistant_message() {
let (session, turn_context) = make_session_and_context().await;
let item = assistant_output_text("hello<oai-mem-citation>doc1</oai-mem-citation> world");
let turn_item =
handle_non_tool_response_item(&item, false, Some(std::path::Path::new(".")))
.await
.expect("assistant message should parse");
let turn_item = handle_non_tool_response_item(&session, &turn_context, &item, false)
.await
.expect("assistant message should parse");
let TurnItem::AgentMessage(agent_message) = turn_item else {
panic!("expected agent message");
@ -443,26 +459,24 @@ mod tests {
}
#[tokio::test]
async fn save_image_generation_result_saves_base64_to_png_in_cwd() {
let dir = tempdir().expect("tempdir");
async fn save_image_generation_result_saves_base64_to_png_in_temp_dir() {
let expected_path = default_image_generation_output_dir().join("ig_save_base64.png");
let _ = std::fs::remove_file(&expected_path);
let saved_path = save_image_generation_result_to_cwd(dir.path(), "ig_123", "Zm9v")
let saved_path = save_image_generation_result("ig_save_base64", "Zm9v")
.await
.expect("image should be saved");
assert_eq!(
saved_path.file_name().and_then(|v| v.to_str()),
Some("ig_123.png")
);
assert_eq!(std::fs::read(saved_path).expect("saved file"), b"foo");
assert_eq!(saved_path, expected_path);
assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo");
let _ = std::fs::remove_file(&saved_path);
}
#[tokio::test]
async fn save_image_generation_result_rejects_data_url_payload() {
let dir = tempdir().expect("tempdir");
let result = "data:image/jpeg;base64,Zm9v";
let err = save_image_generation_result_to_cwd(dir.path(), "ig_456", result)
let err = save_image_generation_result("ig_456", result)
.await
.expect_err("data url payload should error");
assert!(matches!(err, CodexErr::InvalidRequest(_)));
@ -470,42 +484,35 @@ mod tests {
#[tokio::test]
async fn save_image_generation_result_overwrites_existing_file() {
let dir = tempdir().expect("tempdir");
let existing_path = dir.path().join("ig_123.png");
let existing_path = default_image_generation_output_dir().join("ig_overwrite.png");
std::fs::write(&existing_path, b"existing").expect("seed existing image");
let saved_path = save_image_generation_result_to_cwd(dir.path(), "ig_123", "Zm9v")
let saved_path = save_image_generation_result("ig_overwrite", "Zm9v")
.await
.expect("image should be saved");
assert_eq!(
saved_path.file_name().and_then(|v| v.to_str()),
Some("ig_123.png")
);
assert_eq!(std::fs::read(saved_path).expect("saved file"), b"foo");
assert_eq!(saved_path, existing_path);
assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo");
let _ = std::fs::remove_file(&saved_path);
}
#[tokio::test]
async fn save_image_generation_result_sanitizes_call_id_for_output_path() {
let dir = tempdir().expect("tempdir");
async fn save_image_generation_result_sanitizes_call_id_for_temp_dir_output_path() {
let expected_path = default_image_generation_output_dir().join("___ig___.png");
let _ = std::fs::remove_file(&expected_path);
let saved_path = save_image_generation_result_to_cwd(dir.path(), "../ig/..", "Zm9v")
let saved_path = save_image_generation_result("../ig/..", "Zm9v")
.await
.expect("image should be saved");
assert_eq!(saved_path.parent(), Some(dir.path()));
assert_eq!(
saved_path.file_name().and_then(|v| v.to_str()),
Some("___ig___.png")
);
assert_eq!(std::fs::read(saved_path).expect("saved file"), b"foo");
assert_eq!(saved_path, expected_path);
assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo");
let _ = std::fs::remove_file(&saved_path);
}
#[tokio::test]
async fn save_image_generation_result_rejects_non_standard_base64() {
let dir = tempdir().expect("tempdir");
let err = save_image_generation_result_to_cwd(dir.path(), "ig_urlsafe", "_-8")
let err = save_image_generation_result("ig_urlsafe", "_-8")
.await
.expect_err("non-standard base64 should error");
assert!(matches!(err, CodexErr::InvalidRequest(_)));
@ -513,12 +520,9 @@ mod tests {
#[tokio::test]
async fn save_image_generation_result_rejects_non_base64_data_urls() {
let dir = tempdir().expect("tempdir");
let err =
save_image_generation_result_to_cwd(dir.path(), "ig_svg", "data:image/svg+xml,<svg/>")
.await
.expect_err("non-base64 data url should error");
let err = save_image_generation_result("ig_svg", "data:image/svg+xml,<svg/>")
.await
.expect_err("non-base64 data url should error");
assert!(matches!(err, CodexErr::InvalidRequest(_)));
}
}

View file

@ -269,11 +269,14 @@ async fn image_generation_call_event_is_emitted() -> anyhow::Result<()> {
let server = start_mock_server().await;
let TestCodex { codex, cwd, .. } = test_codex().build(&server).await?;
let TestCodex { codex, .. } = test_codex().build(&server).await?;
let call_id = "ig_image_saved_to_temp_dir_default";
let expected_saved_path = std::env::temp_dir().join(format!("{call_id}.png"));
let _ = std::fs::remove_file(&expected_saved_path);
let first_response = sse(vec![
ev_response_created("resp-1"),
ev_image_generation_call("ig_123", "completed", "A tiny blue square", "Zm9v"),
ev_image_generation_call(call_id, "completed", "A tiny blue square", "Zm9v"),
ev_completed("resp-1"),
]);
mount_sse_once(&server, first_response).await;
@ -299,17 +302,17 @@ async fn image_generation_call_event_is_emitted() -> anyhow::Result<()> {
})
.await;
assert_eq!(begin.call_id, "ig_123");
assert_eq!(end.call_id, "ig_123");
assert_eq!(begin.call_id, call_id);
assert_eq!(end.call_id, call_id);
assert_eq!(end.status, "completed");
assert_eq!(end.revised_prompt, Some("A tiny blue square".to_string()));
assert_eq!(end.result, "Zm9v");
let expected_saved_path = cwd.path().join("ig_123.png");
assert_eq!(
end.saved_path,
Some(expected_saved_path.to_string_lossy().into_owned())
);
assert_eq!(std::fs::read(expected_saved_path)?, b"foo");
assert_eq!(std::fs::read(&expected_saved_path)?, b"foo");
let _ = std::fs::remove_file(&expected_saved_path);
Ok(())
}
@ -320,7 +323,9 @@ async fn image_generation_call_event_is_emitted_when_image_save_fails() -> anyho
let server = start_mock_server().await;
let TestCodex { codex, cwd, .. } = test_codex().build(&server).await?;
let TestCodex { codex, .. } = test_codex().build(&server).await?;
let expected_saved_path = std::env::temp_dir().join("ig_invalid.png");
let _ = std::fs::remove_file(&expected_saved_path);
let first_response = sse(vec![
ev_response_created("resp-1"),
@ -356,7 +361,7 @@ async fn image_generation_call_event_is_emitted_when_image_save_fails() -> anyho
assert_eq!(end.revised_prompt, Some("broken payload".to_string()));
assert_eq!(end.result, "_-8");
assert_eq!(end.saved_path, None);
assert!(!cwd.path().join("ig_invalid.png").exists());
assert!(!expected_saved_path.exists());
Ok(())
}

View file

@ -5,4 +5,4 @@ expression: combined
---
• Generated Image:
└ A tiny blue square
└ Saved to: /tmp/project
└ Saved to: /tmp

View file

@ -6289,7 +6289,7 @@ async fn image_generation_call_adds_history_cell() {
status: "completed".into(),
revised_prompt: Some("A tiny blue square".into()),
result: "Zm9v".into(),
saved_path: Some("/tmp/project/ig-1.png".into()),
saved_path: Some("/tmp/ig-1.png".into()),
}),
});