Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit dce4f11

Browse files
ZetkolinkZetkolink
authored
Merge pull request #2 from Zetkolink/fix/multimodal-context-messages
feat: scope images to parent messages to prevent context pollution
2 parents 56ddd1f + fdf4dc2 commit dce4f11

28 files changed

Lines changed: 509 additions & 244 deletions

‎crates/forge_app/src/compact.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,38 @@ mod tests {
510510
insta::assert_yaml_snapshot!(compacted_context);
511511
}
512512

513+
#[test]
514+
fn test_compaction_removes_images_with_parent_message() {
515+
use forge_domain::{ContextMessage, Image, Role, TextMessage};
516+
517+
let environment = test_environment();
518+
let compactor = Compactor::new(Compact::new(), environment);
519+
520+
let image = Image::new_base64("iVBORw0KGgo=".to_string(), "image/png");
521+
let msg_with_image = TextMessage::new(Role::User, "Look at this").add_image(image.clone());
522+
523+
// Context: [User+image, Assistant, User, Assistant]
524+
let context = Context::default()
525+
.add_message(ContextMessage::Text(msg_with_image))
526+
.add_message(ContextMessage::assistant("I see a cat", None, None, None))
527+
.add_message(ContextMessage::user("Write code", None))
528+
.add_message(ContextMessage::assistant("Here's code", None, None, None));
529+
530+
// Compact the first 2 messages (user+image and assistant)
531+
let actual = compactor.compress_single_sequence(context, (0, 1)).unwrap();
532+
533+
// After compaction: [U-summary, U, A]
534+
// The image from the first user message should be gone (compacted away)
535+
for msg in &actual.messages {
536+
if let ContextMessage::Text(text_msg) = &**msg {
537+
assert!(
538+
text_msg.images.is_empty(),
539+
"Images should be removed with their parent message during compaction"
540+
);
541+
}
542+
}
543+
}
544+
513545
#[test]
514546
fn test_compaction_removes_droppable_messages() {
515547
use forge_domain::{ContextMessage, Role, TextMessage};

‎crates/forge_app/src/dto/anthropic/request.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,7 @@ pub struct Message {
232232

233233
impl TryFrom<ContextMessage> for Message {
234234
type Error = anyhow::Error;
235+
#[allow(deprecated)]
235236
fn try_from(value: ContextMessage) -> std::result::Result<Self, Self::Error> {
236237
Ok(match value {
237238
ContextMessage::Text(chat_message) => {
@@ -259,6 +260,9 @@ impl TryFrom<ContextMessage> for Message {
259260
// NOTE: Anthropic does not allow empty text content.
260261
content.push(Content::Text { text: chat_message.content, cache_control: None });
261262
}
263+
for image in chat_message.images {
264+
content.push(Content::from(image));
265+
}
262266
if let Some(tool_calls) = chat_message.tool_calls {
263267
for tool_call in tool_calls {
264268
content.push(tool_call.try_into()?);

‎crates/forge_app/src/dto/google/request.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,7 @@ impl From<forge_domain::ToolDefinition> for FunctionDeclaration {
470470
}
471471

472472
impl From<ContextMessage> for Content {
473+
#[allow(deprecated)]
473474
fn from(message: ContextMessage) -> Self {
474475
match message {
475476
ContextMessage::Text(text_message) => Content::from(text_message),
@@ -499,6 +500,11 @@ impl From<forge_domain::TextMessage> for Content {
499500
});
500501
}
501502

503+
// Add image parts if present
504+
for image in text_message.images {
505+
parts.push(Part::from(image));
506+
}
507+
502508
// Add function calls if present
503509
if let Some(tool_calls) = text_message.tool_calls {
504510
parts.extend(tool_calls.into_iter().map(Part::from));

‎crates/forge_app/src/dto/openai/request.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -443,11 +443,24 @@ impl From<ToolCallFull> for ToolCall {
443443
}
444444

445445
impl From<ContextMessage> for Message {
446+
#[allow(deprecated)]
446447
fn from(value: ContextMessage) -> Self {
447448
match value {
448449
ContextMessage::Text(chat_message) => Message {
449450
role: chat_message.role.into(),
450-
content: Some(MessageContent::Text(chat_message.content)),
451+
content: Some(if chat_message.images.is_empty() {
452+
MessageContent::Text(chat_message.content)
453+
} else {
454+
let mut parts =
455+
vec![ContentPart::Text { text: chat_message.content, cache_control: None }];
456+
for image in chat_message.images {
457+
parts.push(ContentPart::ImageUrl {
458+
image_url: ImageUrl { url: image.url().clone(), detail: None },
459+
cache_control: None,
460+
});
461+
}
462+
MessageContent::Parts(parts)
463+
}),
451464
name: None,
452465
tool_call_id: None,
453466
tool_calls: chat_message

‎crates/forge_app/src/hooks/doom_loop.rs‎

Lines changed: 4 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -259,17 +259,7 @@ mod tests {
259259
use super::*;
260260

261261
fn create_assistant_message(tool_call: &ToolCallFull) -> TextMessage {
262-
TextMessage {
263-
role: Role::Assistant,
264-
content: String::new(),
265-
raw_content: None,
266-
tool_calls: Some(vec![tool_call.clone()]),
267-
thought_signature: None,
268-
model: None,
269-
reasoning_details: None,
270-
droppable: false,
271-
phase: None,
272-
}
262+
TextMessage::new(Role::Assistant, "").tool_calls(vec![tool_call.clone()])
273263
}
274264

275265
fn create_conversation_with_messages(messages: Vec<TextMessage>) -> Conversation {
@@ -395,41 +385,9 @@ mod tests {
395385

396386
#[test]
397387
fn test_extract_assistant_messages() {
398-
let assistant_msg_1 = TextMessage {
399-
role: Role::Assistant,
400-
content: "Response 1".to_string(),
401-
raw_content: None,
402-
tool_calls: None,
403-
thought_signature: None,
404-
model: None,
405-
reasoning_details: None,
406-
droppable: false,
407-
phase: None,
408-
};
409-
410-
let user_msg = TextMessage {
411-
role: Role::User,
412-
content: "Question".to_string(),
413-
raw_content: None,
414-
tool_calls: None,
415-
thought_signature: None,
416-
model: None,
417-
reasoning_details: None,
418-
droppable: false,
419-
phase: None,
420-
};
421-
422-
let assistant_msg_2 = TextMessage {
423-
role: Role::Assistant,
424-
content: "Response 2".to_string(),
425-
raw_content: None,
426-
tool_calls: None,
427-
thought_signature: None,
428-
model: None,
429-
reasoning_details: None,
430-
droppable: false,
431-
phase: None,
432-
};
388+
let assistant_msg_1 = TextMessage::new(Role::Assistant, "Response 1");
389+
let user_msg = TextMessage::new(Role::User, "Question");
390+
let assistant_msg_2 = TextMessage::new(Role::Assistant, "Response 2");
433391

434392
let messages = [
435393
ContextMessage::Text(assistant_msg_1.clone()),

‎crates/forge_app/src/user_prompt.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ impl<S: AttachmentService> UserPromptGenerator<S> {
7777
model: Some(self.agent.model.clone()),
7878
droppable: true, // Droppable so it can be removed during context compression
7979
phase: None,
80+
images: Vec::new(),
8081
};
8182
context = context.add_message(ContextMessage::Text(todo_message));
8283
}
@@ -123,6 +124,7 @@ impl<S: AttachmentService> UserPromptGenerator<S> {
123124
model: Some(self.agent.model.clone()),
124125
droppable: true, // Piped input is droppable
125126
phase: None,
127+
images: Vec::new(),
126128
};
127129
context = context.add_message(ContextMessage::Text(piped_message));
128130
}
@@ -200,6 +202,7 @@ impl<S: AttachmentService> UserPromptGenerator<S> {
200202
model: Some(self.agent.model.clone()),
201203
droppable: false,
202204
phase: None,
205+
images: Vec::new(),
203206
};
204207
context = context.add_message(ContextMessage::Text(message));
205208
}

‎crates/forge_domain/src/compact/summary.rs‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ pub struct TodoChange {
214214
pub kind: TodoChangeKind,
215215
}
216216

217+
#[allow(deprecated)]
217218
impl From<&Context> for ContextSummary {
218219
fn from(value: &Context) -> Self {
219220
let mut messages = vec![];
@@ -274,7 +275,9 @@ impl From<&Context> for ContextSummary {
274275
tool_results.insert(call_id, tool_result);
275276
}
276277
}
277-
ContextMessage::Image(_) => {}
278+
ContextMessage::Image(_) => {
279+
buffer.push(SummaryMessage::Text("[1 image(s) attached]".to_string()));
280+
}
278281
}
279282
}
280283

@@ -311,6 +314,14 @@ fn extract_summary_messages(text_msg: &TextMessage, current_todos: &[Todo]) -> V
311314
blocks.push(SummaryMessage::Text(text_msg.content.clone()));
312315
}
313316

317+
// Note image attachments in summary so LLM knows they existed
318+
if !text_msg.images.is_empty() {
319+
blocks.push(SummaryMessage::Text(format!(
320+
"[{} image(s) attached]",
321+
text_msg.images.len()
322+
)));
323+
}
324+
314325
// Add tool call blocks if present
315326
if let Some(calls) = &text_msg.tool_calls {
316327
blocks.extend(calls.iter().filter_map(|tool_call| {
@@ -873,7 +884,8 @@ mod tests {
873884
}
874885

875886
#[test]
876-
fn test_context_summary_ignores_image_messages() {
887+
#[allow(deprecated)]
888+
fn test_context_summary_includes_image_note() {
877889
let fixture = context(vec![
878890
user("User message"),
879891
ContextMessage::Image(crate::Image::new_base64(
@@ -886,7 +898,13 @@ mod tests {
886898
let actual = ContextSummary::from(&fixture);
887899

888900
let expected = ContextSummary::new(vec![
889-
SummaryBlock::new(Role::User, vec![Block::content("User message")]),
901+
SummaryBlock::new(
902+
Role::User,
903+
vec![
904+
Block::content("User message"),
905+
Block::content("[1 image(s) attached]"),
906+
],
907+
),
890908
SummaryBlock::new(Role::Assistant, vec![Block::content("Assistant")]),
891909
]);
892910

0 commit comments

Comments
 (0)