Feat/restore image generation history (#15223)

Restore image generation items in resumed thread history

codex-rs/app-server-protocol/schema/json/ServerNotification.json

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

codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json

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

codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json

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

codex-rs/app-server-protocol/schema/json/v2/ItemCompletedNotification.json

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

codex-rs/app-server-protocol/schema/json/v2/ItemStartedNotification.json

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

codex-rs/app-server-protocol/schema/json/v2/ReviewStartResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json

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

codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/TurnCompletedNotification.json

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

codex-rs/app-server-protocol/schema/json/v2/TurnStartResponse.json

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

codex-rs/app-server-protocol/schema/json/v2/TurnStartedNotification.json

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

codex-rs/app-server-protocol/schema/typescript/v2/ThreadItem.ts

@@ -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, };

codex-rs/app-server-protocol/src/protocol/thread_history.rs

@@ -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![

codex-rs/app-server-protocol/src/protocol/v2.rs

@@ -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 }

codex-rs/core/src/codex_tests.rs

@@ -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"

codex-rs/core/src/rollout/policy.rs

@@ -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
+        ));
+    }
+}

codex-rs/core/src/stream_events_utils.rs

@@ -372,11 +372,13 @@ pub(crate) async fn handle_non_tool_response_item(
                             image_output_path.display(),
                         ))
                         .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) => {
                         let output_path = image_generation_artifact_path(

codex-rs/tui_app_server/src/app/app_server_adapter.rs

@@ -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(),

codex-rs/tui_app_server/src/chatwidget.rs

@@ -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, .. } => {