feat: replay acp images on session load (#9496)

Signed-off-by: Kalvin Chau <kalvin@block.xyz>

crates/goose/tests/acp_common_tests/mod.rs

@@ -5,8 +5,8 @@
 #[path = "../acp_fixtures/mod.rs"]
 pub mod fixtures;
 use agent_client_protocol::schema::{
-    ListSessionsResponse, McpServer, McpServerHttp, ModelId, SessionInfo, SessionModeId,
-    ToolCallStatus, ToolKind,
+    ContentBlock, ListSessionsResponse, McpServer, McpServerHttp, ModelId, SessionInfo,
+    SessionModeId, SessionUpdate, ToolCallStatus, ToolKind,
 };
 use fixtures::{
     assert_notifications, Connection, FsFixture, Notification, OpenAiFixture, PermissionDecision,
@@ -622,6 +622,57 @@ pub async fn run_load_session_mcp<C: Connection>() {
     assert_eq!(output.text, FAKE_CODE, "tool call failed in loaded session");
 }
 
+pub async fn run_load_session_replays_image_attachment<C: Connection>() {
+    let expected_session_id = C::expected_session_id();
+    let openai = OpenAiFixture::new(
+        vec![(
+            r#""type":"image_url""#.into(),
+            include_str!("../acp_test_data/openai_image_attachment.txt"),
+        )],
+        expected_session_id.clone(),
+    )
+    .await;
+
+    let mut conn = C::new(TestConnectionConfig::default(), openai).await;
+    let SessionData { mut session, .. } = conn.new_session().await.unwrap();
+    expected_session_id.set(&session.session_id().0);
+    let session_id = session.session_id().0.to_string();
+
+    let output = session
+        .prompt_with_image(
+            "Describe what you see in this image",
+            TEST_IMAGE_B64,
+            "image/png",
+            PermissionDecision::Cancel,
+        )
+        .await
+        .unwrap();
+    assert!(output.text.contains("Hello Goose!"));
+    session.session_updates();
+
+    let SessionData { session, .. } = conn.load_session(&session_id, vec![]).await.unwrap();
+    let replayed_images = session
+        .session_updates()
+        .into_iter()
+        .filter_map(|update| match update {
+            SessionUpdate::UserMessageChunk(chunk) => match chunk.content {
+                ContentBlock::Image(image) => Some(image),
+                _ => None,
+            },
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+
+    assert_eq!(
+        replayed_images.len(),
+        1,
+        "expected load_session to replay the user image attachment exactly once"
+    );
+    let replayed_image = &replayed_images[0];
+    assert_eq!(replayed_image.data, TEST_IMAGE_B64);
+    assert_eq!(replayed_image.mime_type, "image/png");
+}
+
 pub async fn run_load_session_error<C: Connection>() {
     let openai = OpenAiFixture::new(vec![], C::expected_session_id()).await;
     let mut conn = C::new(TestConnectionConfig::default(), openai).await;

crates/goose/tests/acp_fixtures/mod.rs

@@ -562,6 +562,9 @@ pub trait Connection: Sized {
 pub trait Session: std::fmt::Debug {
     fn session_id(&self) -> &agent_client_protocol::schema::SessionId;
     fn work_dir(&self) -> std::path::PathBuf;
+    /// Drains and returns raw session updates collected by the fixture.
+    fn session_updates(&self) -> Vec<SessionUpdate>;
+    /// Drains and returns simplified notifications collected by the fixture.
     fn notifications(&self) -> Vec<Notification>;
     async fn prompt(
         &mut self,

crates/goose/tests/acp_fixtures/provider.rs

@@ -325,9 +325,12 @@ impl Session for AcpProviderSession {
         self.work_dir.clone()
     }
 
+    fn session_updates(&self) -> Vec<SessionUpdate> {
+        self.notification_sink.lock().unwrap().drain(..).collect()
+    }
+
     fn notifications(&self) -> Vec<super::Notification> {
-        let updates: Vec<_> = self.notification_sink.lock().unwrap().drain(..).collect();
-        super::to_notifications(&updates)
+        super::to_notifications(&self.session_updates())
     }
 
     async fn prompt(

crates/goose/tests/acp_fixtures/server.rs

@@ -53,6 +53,15 @@ impl std::fmt::Debug for AcpServerSession {
 }
 
 impl AcpServerSession {
+    pub fn session_updates(&self) -> Vec<SessionUpdate> {
+        self.updates
+            .lock()
+            .unwrap()
+            .drain(..)
+            .map(|n| n.update)
+            .collect()
+    }
+
     async fn send_prompt(
         &mut self,
         content: Vec<ContentBlock>,
@@ -464,15 +473,12 @@ impl Session for AcpServerSession {
         self._work_dir.path().to_path_buf()
     }
 
+    fn session_updates(&self) -> Vec<SessionUpdate> {
+        AcpServerSession::session_updates(self)
+    }
+
     fn notifications(&self) -> Vec<super::Notification> {
-        let updates: Vec<_> = self
-            .updates
-            .lock()
-            .unwrap()
-            .drain(..)
-            .map(|n| n.update)
-            .collect();
-        super::to_notifications(&updates)
+        super::to_notifications(&self.session_updates())
     }
 
     async fn prompt(

crates/goose/tests/acp_server_test.rs

@@ -11,12 +11,12 @@ use common_tests::{
     run_close_session, run_config_mcp, run_config_option_mode_set, run_config_option_model_set,
     run_delete_session, run_fs_read_text_file_true, run_fs_write_text_file_false,
     run_fs_write_text_file_true, run_initialize_doesnt_hit_provider, run_list_sessions,
-    run_load_mode, run_load_model, run_load_session_error, run_load_session_mcp, run_mode_set,
-    run_model_list, run_model_set, run_model_set_error_session_not_found,
-    run_new_session_returns_initial_config, run_permission_persistence, run_prompt_basic,
-    run_prompt_error, run_prompt_image, run_prompt_image_attachment, run_prompt_mcp,
-    run_prompt_model_mismatch, run_prompt_skill, run_session_name_update_notification,
-    run_shell_terminal_false, run_shell_terminal_true,
+    run_load_mode, run_load_model, run_load_session_error, run_load_session_mcp,
+    run_load_session_replays_image_attachment, run_mode_set, run_model_list, run_model_set,
+    run_model_set_error_session_not_found, run_new_session_returns_initial_config,
+    run_permission_persistence, run_prompt_basic, run_prompt_error, run_prompt_image,
+    run_prompt_image_attachment, run_prompt_mcp, run_prompt_model_mismatch, run_prompt_skill,
+    run_session_name_update_notification, run_shell_terminal_false, run_shell_terminal_true,
 };
 use goose::config::GooseMode;
 use goose::conversation::message::Message;
@@ -220,6 +220,11 @@ fn test_load_session_mcp() {
     run_test(async { run_load_session_mcp::<AcpServerConnection>().await });
 }
 
+#[test]
+fn test_load_session_replays_image_attachment() {
+    run_test(async { run_load_session_replays_image_attachment::<AcpServerConnection>().await });
+}
+
 #[test]
 fn test_mode_set() {
     run_test(async { run_mode_set::<AcpServerConnection>().await });