Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: make image attachments request-local

Browse Source
This commit is contained in:
Hare
2026-09-24 18:21:29 +09:00
parent e6608dd42f
commit e69aa03f5d
9 changed files with 514 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/compact/token_counter.rs
+27
View File
@@ -338,6 +338,33 @@ mod tests {
assert_eq!(cut.source, EstimateSource::Measured);
}
#[test]
fn retained_split_does_not_treat_image_base64_as_text_tokens() {
fn history_with_image(bytes: usize) -> Vec<Item> {
vec![
msg(&"older text ".repeat(2_000)),
Item::tool_call("image", "ViewImage", r#"{"path":"screen.png"}"#),
Item::tool_result_item_with_attachments(
"image",
"Attached image screen.png",
None,
false,
vec![agen::tool::Attachment::Image(
agen::tool::ImageAttachment::new("image/png", vec![7_u8; bytes]),
)],
),
]
}
let small = history_with_image(8);
let large = history_with_image(1_000_000);
assert_eq!(
split_for_retained_impl(&small, &[], 1_000),
split_for_retained_impl(&large, &[], 1_000)
);
}
#[test]
fn split_all_when_retained_zero() {
let history = vec![msg("a"), msg("b")];
crates/worker/src/compact/worker.rs
+72 -7
View File
@@ -93,10 +93,11 @@ struct SummaryParams {
pub text: String,
}
const MARK_DESCRIPTION: &str = "Inject a file's contents into the compacted context so the \
next session starts with it already read. Use this for files the next task needs in full. \
Optionally specify `offset` (0-based line) and `limit` (line count) to inject only a slice. \
Counts against `auto_read_budget`; overflow returns an error and the mark is not recorded.";
const MARK_DESCRIPTION: &str = "Inject a UTF-8 text file's contents into the compacted context so the \
next session starts with it already read. Images and binary files are rejected; use add_reference \
for those so they can be fetched on demand. Optionally specify `offset` (0-based line) and `limit` \
(line count) to inject only a slice. Counts against `auto_read_budget`; overflow returns an error \
and the mark is not recorded.";
const REFERENCE_DESCRIPTION: &str = "Record a Workdir-relative file path as a named reference in \
the compacted context without injecting its contents. Use for files that are contextually \
@@ -106,6 +107,14 @@ const SUMMARY_DESCRIPTION: &str = "Provide the final structured summary text. Su
replace the previous content; only the last call is used. Must be called before the compact run \
ends or compaction fails.";
fn is_supported_image(bytes: &[u8]) -> bool {
bytes.starts_with(b"\x89PNG\r\n\x1a\n")
|| bytes.starts_with(&[0xff, 0xd8, 0xff])
|| bytes.starts_with(b"GIF87a")
|| bytes.starts_with(b"GIF89a")
|| (bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP")
}
struct MarkReadRequiredTool {
session: WorkdirSessionHandle,
ctx: Arc<Mutex<CompactWorkerContext>>,
@@ -136,9 +145,25 @@ impl Tool for MarkReadRequiredTool {
})
.await
.map_err(|e| ToolError::ExecutionFailed(format!("read failed: {e}")))?;
let text = String::from_utf8_lossy(&result.bytes);
let slice = text.as_ref();
let estimated_tokens = estimate_tokens(slice.len());
if result.truncated {
return Err(ToolError::ExecutionFailed(
"file is too large for auto-read; use add_reference so the next session can fetch it on demand"
.to_string(),
));
}
if is_supported_image(&result.bytes) {
return Err(ToolError::ExecutionFailed(
"image files cannot be injected into compacted context; use add_reference so the next session can call ViewImage on demand"
.to_string(),
));
}
let text = std::str::from_utf8(&result.bytes).map_err(|_| {
ToolError::ExecutionFailed(
"binary files cannot be injected into compacted context; use add_reference so the next session can fetch them on demand"
.to_string(),
)
})?;
let estimated_tokens = estimate_tokens(text.len());
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
let budget = guard.auto_read_budget;
@@ -615,6 +640,46 @@ mod tests {
assert!(guard.auto_read_consumed <= 1_000);
}
#[tokio::test]
async fn mark_read_required_rejects_images_without_recording_auto_read() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("screen.png");
std::fs::write(&path, b"\x89PNG\r\n\x1a\nbody").unwrap();
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(1_000)));
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
session: make_fs(tmp.path()),
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": "screen.png" }).to_string();
let error = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(error.to_string().contains("use add_reference"));
let guard = ctx.lock().unwrap();
assert!(guard.read_required.is_empty());
assert_eq!(guard.auto_read_consumed, 0);
}
#[tokio::test]
async fn mark_read_required_rejects_non_utf8_binary() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("data.bin");
std::fs::write(&path, [0xff, 0xfe, 0xfd]).unwrap();
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(1_000)));
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
session: make_fs(tmp.path()),
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": "data.bin" }).to_string();
let error = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(error.to_string().contains("use add_reference"));
assert!(ctx.lock().unwrap().read_required.is_empty());
}
#[tokio::test]
async fn mark_read_required_rejects_over_budget() {
let tmp = tempfile::TempDir::new().unwrap();
crates/worker/src/worker.rs
+200 -25
View File
@@ -1,4 +1,4 @@
use std::collections::VecDeque;
use std::collections::{HashMap, VecDeque};
#[cfg(test)]
use std::path::Path;
use std::path::PathBuf;
@@ -5571,11 +5571,20 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
let history_entries = self.session.history().entries();
let retain_from = cut.index.min(history_entries.len());
let retained_history_entries = history_entries[retain_from..].to_vec();
let retained_items = retained_history_entries
let mut retained_items = retained_history_entries
.iter()
.map(|entry| entry.item.clone())
.collect::<Vec<_>>();
let entries_to_summarise = history_entries[..retain_from].to_vec();
let default_workdir_alias = (self.workdir_sessions.len() == 1)
.then(|| self.workdir_sessions.aliases().into_iter().next())
.flatten()
.map(|alias| alias.as_str().to_string());
let retained_image_references =
project_compaction_attachments(&mut retained_items, default_workdir_alias.as_deref());
// The compactor explores summaries, text details, and ViewImage call arguments. It
// never needs durable image bodies; relevant images are nominated with add_reference.
let mut entries_to_summarise = history_entries[..retain_from].to_vec();
project_compactor_capture_attachments(&mut entries_to_summarise);
let items_to_summarise = entries_to_summarise
.iter()
.map(|entry| entry.item.clone())
@@ -5840,15 +5849,16 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
"You have not called `write_summary` yet. Deliver the structured \
summary now (Completed Tasks / Active Task / Key Decisions / \
User Directives / Current Work) and nominate any files the next \
session needs with `mark_read_required`."
session needs with `mark_read_required` for UTF-8 text or \
`add_reference` for images, binary files, and on-demand reads."
.to_string(),
)
} else if snapshot.read_required.is_empty() && !default_refs.is_empty() {
Some(
"Summary received. If any of the referenced files are required \
for the next session to continue the task, call \
`mark_read_required` on them now. Otherwise reply briefly to \
close out."
"Summary received. If any of the referenced UTF-8 text files must \
be present in the next session, call `mark_read_required` on them. \
Use `add_reference` for images, binary files, or files that can be \
fetched on demand. Otherwise reply briefly to close out."
.to_string(),
)
} else {
@@ -5983,20 +5993,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
// Reference list as a single system message; omitted when empty.
let reference_message = (!final_ctx.references.is_empty()).then(|| {
let list = final_ctx
.references
.iter()
.map(|p| format!("- {}", p.display()))
.collect::<Vec<_>>()
.join("\n");
Item::system_message(format!(
"[Referenced files — read before compaction, contents not included]\n\
{list}\n\
Use read_file to access current contents if needed."
))
});
// Related files are carried by name only. Retained ViewImage payloads are
// projected out above and become explicit on-demand image references here.
let reference_message =
build_related_file_message(&final_ctx.references, &retained_image_references);
// Count surviving user_messages before consuming `retained_items`
// — needed to align `self.user_segments` after the swap below.
@@ -7276,6 +7276,124 @@ impl From<EngineResult> for WorkerRunResult {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct CompactionImageReference {
target_workdir: Option<String>,
path: String,
}
#[derive(serde::Deserialize)]
struct ViewImageCallArguments {
#[serde(default)]
target_workdir: Option<String>,
path: String,
}
/// Remove binary bodies from the immutable capture exposed to the compactor.
fn project_compactor_capture_attachments<A>(entries: &mut [HistoryEntry<A>]) -> usize {
let mut projected = 0;
for entry in entries {
if let Item::ToolResult { attachments, .. } = &mut entry.item
&& !attachments.is_empty()
{
attachments.clear();
projected += 1;
}
}
projected
}
/// Remove every durable binary attachment from the compacted tail and preserve
/// re-readable ViewImage origins as related-file references.
fn project_compaction_attachments(
items: &mut [Item],
default_workdir_alias: Option<&str>,
) -> Vec<CompactionImageReference> {
let view_image_calls = items
.iter()
.filter_map(|item| match item {
Item::ToolCall {
call_id,
name,
arguments,
..
} if name == "ViewImage" => serde_json::from_str::<ViewImageCallArguments>(arguments)
.ok()
.map(|arguments| {
(
call_id.clone(),
CompactionImageReference {
target_workdir: arguments
.target_workdir
.or_else(|| default_workdir_alias.map(str::to_string)),
path: arguments.path,
},
)
}),
_ => None,
})
.collect::<HashMap<_, _>>();
let mut references = Vec::new();
for item in items {
let Item::ToolResult {
call_id,
attachments,
..
} = item
else {
continue;
};
if attachments.is_empty() {
continue;
}
attachments.clear();
if let Some(reference) = view_image_calls.get(call_id)
&& !references.contains(reference)
{
references.push(reference.clone());
}
}
references
}
fn build_related_file_message(
references: &[PathBuf],
image_references: &[CompactionImageReference],
) -> Option<Item> {
if references.is_empty() && image_references.is_empty() {
return None;
}
let mut lines = Vec::with_capacity(references.len() + image_references.len());
for reference in references {
let path = reference.display().to_string();
if image_references
.iter()
.any(|image| image.target_workdir.is_none() && image.path == path)
{
continue;
}
lines.push(format!("- {path}"));
}
for reference in image_references {
match reference.target_workdir.as_deref() {
Some(target_workdir) => lines.push(format!(
"- {} (image; use ViewImage with target_workdir {:?})",
reference.path, target_workdir
)),
None => lines.push(format!("- {} (image; use ViewImage)", reference.path)),
}
}
Some(Item::system_message(format!(
"[Related files — contents not included]\n{}\n\
Use Read for text files and ViewImage for image references when their current contents \
are needed.",
lines.join("\n")
)))
}
#[derive(Debug, Clone, Copy)]
struct SummaryInputOptions {
overview_target_tokens: u64,
@@ -7329,9 +7447,9 @@ fn build_summary_input(
if !default_refs.is_empty() {
out.push_str(
"These files were touched recently in this session. Use `read_file` \
on them as needed, then call `mark_read_required` for any whose \
contents the next session must have, and `add_reference` for files \
it should know about by name only.\n\n## Referenced files\n",
on text files as needed, then call `mark_read_required` for UTF-8 text \
whose contents the next session must have. Call `add_reference` for images, \
binary files, and files it should know about by name only.\n\n## Referenced files\n",
);
for p in default_refs {
out.push_str("- ");
@@ -8456,6 +8574,63 @@ mod build_summary_prompt_tests {
.text
}
#[test]
fn retained_images_become_on_demand_related_files() {
let mut items = vec![
Item::tool_call_json(
"image-1",
"ViewImage",
serde_json::json!({
"path": "screenshots/result.png"
}),
),
Item::tool_result_item_with_attachments(
"image-1",
"Attached image screenshots/result.png",
None,
false,
vec![agen::tool::Attachment::Image(
agen::tool::ImageAttachment::new("image/png", vec![7_u8; 100_000]),
)],
),
];
let mut capture_entries = items
.iter()
.cloned()
.map(|item| HistoryEntry::new(item, ()))
.collect::<Vec<_>>();
assert_eq!(
project_compactor_capture_attachments(&mut capture_entries),
1
);
assert!(matches!(
&capture_entries[1].item,
Item::ToolResult { attachments, .. } if attachments.is_empty()
));
let references = project_compaction_attachments(&mut items, Some("main"));
assert_eq!(
references,
vec![CompactionImageReference {
target_workdir: Some("main".to_string()),
path: "screenshots/result.png".to_string(),
}]
);
assert!(matches!(
&items[1],
Item::ToolResult { attachments, .. } if attachments.is_empty()
));
let message = build_related_file_message(&[], &references)
.and_then(|item| item.as_text().map(str::to_string))
.unwrap();
assert!(message.contains("[Related files — contents not included]"));
assert!(message.contains("screenshots/result.png"));
assert!(message.contains("target_workdir \"main\""));
assert!(message.contains("ViewImage"));
}
#[test]
fn strips_tool_call_arguments() {
let items = vec![Item::tool_call_json(