mirror of
https://github.com/block/goose.git
synced 2026-06-02 06:19:33 +02:00
feat: replay acp images on session load (#9496)
Signed-off-by: Kalvin Chau <kalvin@block.xyz>
This commit is contained in:
@@ -2401,6 +2401,32 @@ fn replay_message_meta(message: &Message) -> Meta {
|
||||
meta
|
||||
}
|
||||
|
||||
fn replay_audience_annotations(audience: &[Role]) -> Annotations {
|
||||
Annotations::new().audience(
|
||||
audience
|
||||
.iter()
|
||||
.map(|role| match role {
|
||||
Role::Assistant => agent_client_protocol::schema::Role::Assistant,
|
||||
Role::User => agent_client_protocol::schema::Role::User,
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
)
|
||||
}
|
||||
|
||||
fn send_replay_content_chunk(
|
||||
cx: &ConnectionTo<Client>,
|
||||
session_id: &SessionId,
|
||||
message: &Message,
|
||||
content: ContentBlock,
|
||||
) -> std::result::Result<(), agent_client_protocol::Error> {
|
||||
let chunk = ContentChunk::new(content).meta(replay_message_meta(message));
|
||||
let update = match message.role {
|
||||
Role::User => SessionUpdate::UserMessageChunk(chunk),
|
||||
Role::Assistant => SessionUpdate::AgentMessageChunk(chunk),
|
||||
};
|
||||
cx.send_notification(SessionNotification::new(session_id.clone(), update))
|
||||
}
|
||||
|
||||
fn replay_message_goose_meta(message: &Message) -> serde_json::Map<String, serde_json::Value> {
|
||||
let mut goose = serde_json::Map::new();
|
||||
goose.insert("created".to_string(), serde_json::json!(message.created));
|
||||
@@ -2824,30 +2850,28 @@ impl GooseAcpAgent {
|
||||
MessageContent::Text(text) => {
|
||||
let mut tc = TextContent::new(text.text.clone());
|
||||
if let Some(audience) = text.audience() {
|
||||
tc = tc.annotations(
|
||||
Annotations::new().audience(
|
||||
audience
|
||||
.iter()
|
||||
.map(|r| match r {
|
||||
Role::Assistant => {
|
||||
agent_client_protocol::schema::Role::Assistant
|
||||
}
|
||||
Role::User => agent_client_protocol::schema::Role::User,
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
);
|
||||
tc = tc.annotations(replay_audience_annotations(audience));
|
||||
}
|
||||
let chunk = ContentChunk::new(ContentBlock::Text(tc))
|
||||
.meta(replay_message_meta(message));
|
||||
let update = match message.role {
|
||||
Role::User => SessionUpdate::UserMessageChunk(chunk),
|
||||
Role::Assistant => SessionUpdate::AgentMessageChunk(chunk),
|
||||
};
|
||||
cx.send_notification(SessionNotification::new(
|
||||
args.session_id.clone(),
|
||||
update,
|
||||
))?;
|
||||
send_replay_content_chunk(
|
||||
cx,
|
||||
&args.session_id,
|
||||
message,
|
||||
ContentBlock::Text(tc),
|
||||
)?;
|
||||
}
|
||||
MessageContent::Image(image) => {
|
||||
let mut image_content =
|
||||
ImageContent::new(image.data.clone(), image.mime_type.clone());
|
||||
if let Some(audience) = image.audience() {
|
||||
image_content =
|
||||
image_content.annotations(replay_audience_annotations(audience));
|
||||
}
|
||||
send_replay_content_chunk(
|
||||
cx,
|
||||
&args.session_id,
|
||||
message,
|
||||
ContentBlock::Image(image_content),
|
||||
)?;
|
||||
}
|
||||
MessageContent::ToolRequest(tool_request) => {
|
||||
// Replay-only: emit the ToolCall notification and
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 });
|
||||
|
||||
Reference in New Issue
Block a user