Feat/restore image generation history (#15223)

Restore image generation items in resumed thread history
This commit is contained in:
Won Park 2026-03-19 22:57:16 -07:00 committed by GitHub
parent b3a4da84da
commit 461ba012fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 213 additions and 10 deletions

View file

@ -2817,6 +2817,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -12540,6 +12540,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -10300,6 +10300,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1026,6 +1026,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1026,6 +1026,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1140,6 +1140,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1633,6 +1633,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1633,6 +1633,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1633,6 +1633,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1391,6 +1391,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1140,6 +1140,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1140,6 +1140,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -1140,6 +1140,12 @@
"null"
]
},
"savedPath": {
"type": [
"string",
"null"
]
},
"status": {
"type": "string"
},

View file

@ -97,4 +97,4 @@ reasoningEffort: ReasoningEffort | null,
/**
* 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, };

View file

@ -569,6 +569,7 @@ impl ThreadHistoryBuilder {
status: String::new(),
revised_prompt: None,
result: String::new(),
saved_path: None,
};
self.upsert_item_in_current_turn(item);
}
@ -579,6 +580,7 @@ impl ThreadHistoryBuilder {
status: payload.status.clone(),
revised_prompt: payload.revised_prompt.clone(),
result: payload.result.clone(),
saved_path: payload.saved_path.clone(),
};
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]
fn splits_reasoning_when_interleaved() {
let events = vec![

View file

@ -4256,6 +4256,9 @@ pub enum ThreadItem {
status: String,
revised_prompt: Option<String>,
result: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
saved_path: Option<String>,
},
#[serde(rename_all = "camelCase")]
#[ts(rename_all = "camelCase")]
@ -4432,6 +4435,7 @@ impl From<CoreTurnItem> for ThreadItem {
status: image.status,
revised_prompt: image.revised_prompt,
result: image.result,
saved_path: image.saved_path,
},
CoreTurnItem::ContextCompaction(compaction) => {
ThreadItem::ContextCompaction { id: compaction.id }

View file

@ -3751,7 +3751,12 @@ async fn handle_output_item_done_records_image_save_history_message() {
image_output_path.display(),
))
.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!(
std::fs::read(&expected_saved_path).expect("saved file"),
b"foo"

View file

@ -105,7 +105,8 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| EventMsg::UndoCompleted(_)
| EventMsg::TurnAborted(_)
| EventMsg::TurnStarted(_)
| EventMsg::TurnComplete(_) => Some(EventPersistenceMode::Limited),
| EventMsg::TurnComplete(_)
| EventMsg::ImageGenerationEnd(_) => Some(EventPersistenceMode::Limited),
EventMsg::ItemCompleted(event) => {
// Plan items are derived from streaming tags and are not part of the
// 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::McpToolCallEnd(_)
| EventMsg::ViewImageToolCall(_)
| EventMsg::ImageGenerationEnd(_)
| EventMsg::CollabAgentSpawnEnd(_)
| EventMsg::CollabAgentInteractionEnd(_)
| EventMsg::CollabWaitingEnd(_)
@ -183,3 +183,27 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option<EventPersistenceMode> {
| 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
));
}
}

View file

@ -372,11 +372,13 @@ pub(crate) async fn handle_non_tool_response_item(
image_output_path.display(),
))
.into();
sess.record_conversation_items(
turn_context,
std::slice::from_ref(&message),
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(),
)
.await;
.into();
sess.record_conversation_items(turn_context, &[message, copy_message])
.await;
}
Err(err) => {
let output_path = image_generation_artifact_path(

View file

@ -995,12 +995,13 @@ fn thread_item_to_core(item: &ThreadItem) -> Option<TurnItem> {
status,
revised_prompt,
result,
saved_path,
} => Some(TurnItem::ImageGeneration(ImageGenerationItem {
id: id.clone(),
status: status.clone(),
revised_prompt: revised_prompt.clone(),
result: result.clone(),
saved_path: None,
saved_path: saved_path.clone(),
})),
ThreadItem::ContextCompaction { id } => {
Some(TurnItem::ContextCompaction(ContextCompactionItem {
@ -1850,6 +1851,7 @@ mod tests {
status: "completed".to_string(),
revised_prompt: Some("diagram".to_string()),
result: "image.png".to_string(),
saved_path: None,
},
ThreadItem::ContextCompaction {
id: "compact-1".to_string(),

View file

@ -5725,13 +5725,14 @@ impl ChatWidget {
status,
revised_prompt,
result,
saved_path,
} => {
self.on_image_generation_end(ImageGenerationEndEvent {
call_id: id,
result,
revised_prompt,
status,
saved_path: None,
saved_path,
});
}
ThreadItem::EnteredReviewMode { review, .. } => {