Feat/restore image generation history (#15223)

Restore image generation items in resumed thread history
This commit is contained in:
Won Park authored and GitHub committed 2026-03-19 22:57:16 -07:00
1 parent b3a4da84da
commit 461ba012fc
26 files changed
+213 -10

No files matched your search

@@ -2817,6 +2817,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -12540,6 +12540,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -10300,6 +10300,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1026,6 +1026,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1026,6 +1026,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1140,6 +1140,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1633,6 +1633,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1633,6 +1633,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1633,6 +1633,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1391,6 +1391,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1140,6 +1140,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1140,6 +1140,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -1140,6 +1140,12 @@
"null" "null"
] ]
}, },
"savedPath": {
"type": [
"string",
"null"
]
},
"status": { "status": {
"type": "string" "type": "string"
}, },
@@ -97,4 +97,4 @@ reasoningEffort: ReasoningEffort | null,
/** /**
* Last known status of the target agents, when available. * Last known status of the target agents, when available.
*/ */
agentsStates: { [key in string]?: CollabAgentState }, } | { "type": "webSearch", id: string, query: string, action: WebSearchAction | null, } | { "type": "imageView", id: string, path: string, } | { "type": "imageGeneration", id: string, status: string, revisedPrompt: string | null, result: string, } | { "type": "enteredReviewMode", id: string, review: string, } | { "type": "exitedReviewMode", id: string, review: string, } | { "type": "contextCompaction", id: string, }; agentsStates: { [key in string]?: CollabAgentState }, } | { "type": "webSearch", id: string, query: string, action: WebSearchAction | null, } | { "type": "imageView", id: string, path: string, } | { "type": "imageGeneration", id: string, status: string, revisedPrompt: string | null, result: string, savedPath?: string, } | { "type": "enteredReviewMode", id: string, review: string, } | { "type": "exitedReviewMode", id: string, review: string, } | { "type": "contextCompaction", id: string, };
@@ -569,6 +569,7 @@ impl ThreadHistoryBuilder {
status: String::new(), status: String::new(),
revised_prompt: None, revised_prompt: None,
result: String::new(), result: String::new(),
saved_path: None,
}; };
self.upsert_item_in_current_turn(item); self.upsert_item_in_current_turn(item);
} }
@@ -579,6 +580,7 @@ impl ThreadHistoryBuilder {
status: payload.status.clone(), status: payload.status.clone(),
revised_prompt: payload.revised_prompt.clone(), revised_prompt: payload.revised_prompt.clone(),
result: payload.result.clone(), result: payload.result.clone(),
saved_path: payload.saved_path.clone(),
}; };
self.upsert_item_in_current_turn(item); self.upsert_item_in_current_turn(item);
} }
@@ -1385,6 +1387,61 @@ mod tests {
); );
} }
#[test]
fn replays_image_generation_end_events_into_turn_history() {
let items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
turn_id: "turn-image".into(),
model_context_window: None,
collaboration_mode_kind: Default::default(),
})),
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "generate an image".into(),
images: None,
text_elements: Vec::new(),
local_images: Vec::new(),
})),
RolloutItem::EventMsg(EventMsg::ImageGenerationEnd(ImageGenerationEndEvent {
call_id: "ig_123".into(),
status: "completed".into(),
revised_prompt: Some("final prompt".into()),
result: "Zm9v".into(),
saved_path: Some("/tmp/ig_123.png".into()),
})),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-image".into(),
last_agent_message: None,
})),
];
let turns = build_turns_from_rollout_items(&items);
assert_eq!(turns.len(), 1);
assert_eq!(
turns[0],
Turn {
id: "turn-image".into(),
status: TurnStatus::Completed,
error: None,
items: vec![
ThreadItem::UserMessage {
id: "item-1".into(),
content: vec![UserInput::Text {
text: "generate an image".into(),
text_elements: Vec::new(),
}],
},
ThreadItem::ImageGeneration {
id: "ig_123".into(),
status: "completed".into(),
revised_prompt: Some("final prompt".into()),
result: "Zm9v".into(),
saved_path: Some("/tmp/ig_123.png".into()),
},
],
}
);
}
#[test] #[test]
fn splits_reasoning_when_interleaved() { fn splits_reasoning_when_interleaved() {
let events = vec![ let events = vec![
@@ -4256,6 +4256,9 @@ pub enum ThreadItem {
status: String, status: String,
revised_prompt: Option<String>, revised_prompt: Option<String>,
result: String, result: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
saved_path: Option<String>,
}, },
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
#[ts(rename_all = "camelCase")] #[ts(rename_all = "camelCase")]
@@ -4432,6 +4435,7 @@ impl From<CoreTurnItem> for ThreadItem {
status: image.status, status: image.status,
revised_prompt: image.revised_prompt, revised_prompt: image.revised_prompt,
result: image.result, result: image.result,
saved_path: image.saved_path,
}, },
CoreTurnItem::ContextCompaction(compaction) => { CoreTurnItem::ContextCompaction(compaction) => {
ThreadItem::ContextCompaction { id: compaction.id } ThreadItem::ContextCompaction { id: compaction.id }
+6 -1
View File
@@ -3751,7 +3751,12 @@ async fn handle_output_item_done_records_image_save_history_message() {
image_output_path.display(), image_output_path.display(),
)) ))
.into(); .into();
assert_eq!(history.raw_items(), &[save_message, item]); let copy_message: ResponseItem = DeveloperInstructions::new(
"If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it."
.to_string(),
)
.into();
assert_eq!(history.raw_items(), &[save_message, copy_message, item]);
assert_eq!( assert_eq!(
std::fs::read(&expected_saved_path).expect("saved file"), std::fs::read(&expected_saved_path).expect("saved file"),
b"foo" b"foo"
+26 -2
View File
@@ -105,7 +105,8 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| EventMsg::UndoCompleted(_) | EventMsg::UndoCompleted(_)
| EventMsg::TurnAborted(_) | EventMsg::TurnAborted(_)
| EventMsg::TurnStarted(_) | EventMsg::TurnStarted(_)
| EventMsg::TurnComplete(_) => Some(EventPersistenceMode::Limited), | EventMsg::TurnComplete(_)
| EventMsg::ImageGenerationEnd(_) => Some(EventPersistenceMode::Limited),
EventMsg::ItemCompleted(event) => { EventMsg::ItemCompleted(event) => {
// Plan items are derived from streaming tags and are not part of the // Plan items are derived from streaming tags and are not part of the
// raw ResponseItem history, so we persist their completion to replay // raw ResponseItem history, so we persist their completion to replay
@@ -123,7 +124,6 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| EventMsg::PatchApplyEnd(_) | EventMsg::PatchApplyEnd(_)
| EventMsg::McpToolCallEnd(_) | EventMsg::McpToolCallEnd(_)
| EventMsg::ViewImageToolCall(_) | EventMsg::ViewImageToolCall(_)
| EventMsg::ImageGenerationEnd(_)
| EventMsg::CollabAgentSpawnEnd(_) | EventMsg::CollabAgentSpawnEnd(_)
| EventMsg::CollabAgentInteractionEnd(_) | EventMsg::CollabAgentInteractionEnd(_)
| EventMsg::CollabWaitingEnd(_) | EventMsg::CollabWaitingEnd(_)
@@ -183,3 +183,27 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| EventMsg::ImageGenerationBegin(_) => None, | EventMsg::ImageGenerationBegin(_) => None,
} }
} }
#[cfg(test)]
mod tests {
use super::EventPersistenceMode;
use super::should_persist_event_msg;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ImageGenerationEndEvent;
#[test]
fn persists_image_generation_end_events_in_limited_mode() {
let event = EventMsg::ImageGenerationEnd(ImageGenerationEndEvent {
call_id: "ig_123".into(),
status: "completed".into(),
revised_prompt: Some("final prompt".into()),
result: "Zm9v".into(),
saved_path: None,
});
assert!(should_persist_event_msg(
&event,
EventPersistenceMode::Limited
));
}
}
+6 -4
View File
@@ -372,11 +372,13 @@ pub(crate) async fn handle_non_tool_response_item(
image_output_path.display(), image_output_path.display(),
)) ))
.into(); .into();
sess.record_conversation_items( let copy_message: ResponseItem = DeveloperInstructions::new(
turn_context, "If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it."
std::slice::from_ref(&message), .to_string(),
) )
.await; .into();
sess.record_conversation_items(turn_context, &[message, copy_message])
.await;
} }
Err(err) => { Err(err) => {
let output_path = image_generation_artifact_path( let output_path = image_generation_artifact_path(
@@ -995,12 +995,13 @@ fn thread_item_to_core(item: &ThreadItem) -> Option<TurnItem> {
status, status,
revised_prompt, revised_prompt,
result, result,
saved_path,
} => Some(TurnItem::ImageGeneration(ImageGenerationItem { } => Some(TurnItem::ImageGeneration(ImageGenerationItem {
id: id.clone(), id: id.clone(),
status: status.clone(), status: status.clone(),
revised_prompt: revised_prompt.clone(), revised_prompt: revised_prompt.clone(),
result: result.clone(), result: result.clone(),
saved_path: None, saved_path: saved_path.clone(),
})), })),
ThreadItem::ContextCompaction { id } => { ThreadItem::ContextCompaction { id } => {
Some(TurnItem::ContextCompaction(ContextCompactionItem { Some(TurnItem::ContextCompaction(ContextCompactionItem {
@@ -1850,6 +1851,7 @@ mod tests {
status: "completed".to_string(), status: "completed".to_string(),
revised_prompt: Some("diagram".to_string()), revised_prompt: Some("diagram".to_string()),
result: "image.png".to_string(), result: "image.png".to_string(),
saved_path: None,
}, },
ThreadItem::ContextCompaction { ThreadItem::ContextCompaction {
id: "compact-1".to_string(), id: "compact-1".to_string(),
+2 -1
View File
@@ -5725,13 +5725,14 @@ impl ChatWidget {
status, status,
revised_prompt, revised_prompt,
result, result,
saved_path,
} => { } => {
self.on_image_generation_end(ImageGenerationEndEvent { self.on_image_generation_end(ImageGenerationEndEvent {
call_id: id, call_id: id,
result, result,
revised_prompt, revised_prompt,
status, status,
saved_path: None, saved_path,
}); });
} }
ThreadItem::EnteredReviewMode { review, .. } => { ThreadItem::EnteredReviewMode { review, .. } => {