Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions codex-rs/core/src/codex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6636,8 +6636,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,
Expand Down Expand Up @@ -6818,8 +6817,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;
Expand Down
128 changes: 128 additions & 0 deletions codex-rs/core/src/codex_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -3116,6 +3136,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;
Expand Down
6 changes: 0 additions & 6 deletions codex-rs/core/src/context_manager/history_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 0 additions & 3 deletions codex-rs/core/src/context_manager/normalize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading