diff --git a/crates/cli/src/ui/modern/app.rs b/crates/cli/src/ui/modern/app.rs index 71ecb684..e03fedf8 100644 --- a/crates/cli/src/ui/modern/app.rs +++ b/crates/cli/src/ui/modern/app.rs @@ -439,6 +439,28 @@ pub struct App { pub command_palette: Option, /// Ctrl+M / `/model` in-TUI model picker. pub model_picker: Option, + /// Images mentioned in the prompt, attached to the next turn. Held as + /// open descriptors from the moment their mention was validated, so + /// starting the turn needs no second look at any path. + pub pending_images: Vec, + /// The same images once read and encoded off the UI thread, waiting for + /// the turn they belong to to start. + pub pending_attachments: Vec, + /// Prompts whose images were already read when another prompt took + /// the turn, each held with its own images so it can be sent as it + /// was. Never queued as text: the queue re-expands what it holds, + /// which would resolve the mention a second time and read the file + /// again. A deque because a second one can be held before the first + /// has gone, and the order they were submitted in is the order they + /// must be sent in. + /// Set when a cancel was asked for *in order to send something else* + /// (interject, queue send-now). A staged prompt alone cannot tell the + /// difference: a slash command can stage one at any moment, including + /// while the user is pressing Ctrl+C to stop everything. + pub cancel_is_interject: bool, + #[allow(clippy::type_complexity)] + pub deferred_prompts: + std::collections::VecDeque<(String, Vec)>, /// Whether `ui.edit_mode` asked for vi bindings. pub vi_mode: bool, /// Composer mode when `vi_mode` is on. @@ -621,6 +643,10 @@ impl App { pending_task_output: None, command_palette: None, model_picker: None, + pending_images: Vec::new(), + pending_attachments: Vec::new(), + cancel_is_interject: false, + deferred_prompts: std::collections::VecDeque::new(), vi_mode: false, composer_mode: ComposerMode::Insert, vi_pending_d: false, @@ -1718,6 +1744,15 @@ impl App { /// Resolve user text into a turn: expand `/skill` invocations the same /// way slash dispatch does via `commands::execute` skill lookup. fn enqueue_turn(&mut self, text: String) { + // Attachments belong to exactly one prompt. A staged prompt can be + // replaced before it starts (two interjections while a turn is + // cancelling), and only the mention branch assigns `pending_images` + // — so clear here, or the replacement turn would carry the previous + // prompt's image and disclose a file the user did not mean to send. + // Both stages are cleared: one holds descriptors, the other the + // bytes already read from them. + self.pending_images.clear(); + self.pending_attachments.clear(); let mut mention_notes: Vec = Vec::new(); let (display, prompt) = match try_expand_skill_slash_full(&text, &self.cwd, self.disable_skill_shell) { @@ -1766,6 +1801,7 @@ impl App { ) { Some(expansion) => { mention_notes = expansion.notes; + self.pending_images = expansion.images; expansion.prompt } None => text.clone(), @@ -2025,6 +2061,7 @@ impl App { )); // enqueue_turn clears input (already empty) and sets pending_submit. self.enqueue_turn(text); + self.cancel_is_interject = true; self.request_cancel(); } else { self.enqueue_turn(text); @@ -2408,6 +2445,12 @@ impl App { pub fn new_conversation(&mut self) { self.conversation_epoch = self.conversation_epoch.wrapping_add(1); self.todos.clear(); + // Anything staged belonged to the conversation being replaced. + // Sending it into the new one would attach a file to a thread the + // user never attached it to. + self.deferred_prompts.clear(); + self.pending_images.clear(); + self.pending_attachments.clear(); self.dirty = true; } @@ -2641,6 +2684,7 @@ impl App { "queue send-now — cancelling turn…".into(), )); self.enqueue_turn(text); + self.cancel_is_interject = true; self.request_cancel(); } else { self.enqueue_turn(text); @@ -2693,6 +2737,100 @@ impl App { self.dirty = true; } + /// Take a prompt's images once they have been read and encoded. + /// + /// Another prompt can arrive while the read is running — a slash + /// command that produces a prompt is staged directly rather than + /// queued — and that one has its own attachments staged with it. + /// Overwriting it would lose it silently; keeping it and attaching + /// *these* blocks would be worse still, putting one prompt's image on + /// another's turn. So the newer prompt keeps its turn and this one is + /// held aside *with its blocks* until the turn frees up — not queued + /// as text, which would re-expand the mention and read the file a + /// second time, sending bytes the staged turn never had. + pub fn accept_encoded_attachments( + &mut self, + prompt: String, + blocks: Vec, + ) { + if self.pending_submit.is_some() { + self.transcript.push(TranscriptItem::System( + "another prompt was sent first — sending this one with its images next".into(), + )); + self.deferred_prompts.push_back((prompt, blocks)); + self.dirty = true; + return; + } + self.pending_attachments = blocks; + self.pending_submit = Some(prompt); + self.dirty = true; + } + + /// Drop anything that would send itself after a cancel. + /// + /// Interject and queue send-now cancel the live turn *in order to* + /// send something else, and say so; a held prompt then keeps its place + /// behind that. Every other cancel is a stop, and nothing may follow + /// it on its own — a prompt merely sitting in `pending_submit` proves + /// nothing, since a slash command can stage one at any moment, + /// including while the user is pressing Ctrl+C. + pub fn cancel_pending_followups(&mut self) { + if !std::mem::take(&mut self.cancel_is_interject) { + self.deferred_prompts.clear(); + } + } + + /// Send a deferred prompt once the turn it was waiting behind is gone. + /// + /// Restored with the blocks it was encoded with, so the file is never + /// read again. Held back while another prompt's descriptors are still + /// staged, so the two cannot be mixed. + pub fn rearm_deferred_prompt(&mut self) { + if self.pending_submit.is_some() || !self.pending_images.is_empty() { + return; + } + if let Some((prompt, blocks)) = self.deferred_prompts.pop_front() { + self.pending_attachments = blocks; + self.pending_submit = Some(prompt); + self.phase = Phase::Streaming; + self.dirty = true; + } + } + + /// Give up on a prompt whose attachments were still being read when the + /// user cancelled. + /// + /// No turn was ever spawned for it, so nothing else will leave the + /// streaming phase, and every later prompt would queue behind a turn + /// that does not exist. Deliberately not `mark_turn_idle`: that + /// announces a *finished* turn, and this one never started. + /// + /// Only the cancelled prompt is abandoned. Its own descriptors went + /// into the read that is being discarded, so anything staged here now + /// belongs to a prompt the user submitted *after* the cancel — an + /// interjection with its own images — and clearing that would send it + /// as a text-only turn. + pub fn abandon_staged_attachments(&mut self) { + // Held prompts are not touched here: whether they survive is the + // cancel policy's call (`cancel_pending_followups`), and an + // interject reaching this path still means only that *this* read + // is being dropped. + self.turn_live = false; + self.turn_started_at = None; + self.phase = if self.pending_submit.is_some() { + // A replacement prompt is already waiting; it still has a turn + // coming, so the streaming phase is still true. + Phase::Streaming + } else if self.modals.is_empty() { + Phase::Idle + } else { + Phase::Permission + }; + self.cancel_requested = false; + self.status_message = "cancelled".into(); + self.dirty = true; + } + pub fn tick(&mut self) { self.tick = self.tick.wrapping_add(1); if let Some((_, ref mut left)) = self.toast { @@ -5024,4 +5162,408 @@ mod tests { app.submit(); assert_eq!(app.pending_submit.as_deref(), Some("hello there")); } + + fn write_png(app: &App, name: &str) { + std::fs::write( + std::path::Path::new(&app.cwd).join(name), + [0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a], + ) + .unwrap(); + } + + #[test] + fn submitting_an_image_mention_stages_it_for_the_turn() { + let (_dir, mut app) = app_in_workspace(); + write_png(&app, "shot.png"); + type_input(&mut app, "look at @shot.png"); + app.submit(); + assert_eq!(app.pending_images.len(), 1, "image was not staged"); + } + + /// Replacing a staged prompt must drop its attachments: the second + /// prompt would otherwise ship the first one's image and disclose a + /// file the user never meant to send with it. + #[test] + fn replacing_a_staged_prompt_drops_its_images() { + let (_dir, mut app) = app_in_workspace(); + write_png(&app, "shot.png"); + type_input(&mut app, "look at @shot.png"); + app.submit(); + assert_eq!(app.pending_images.len(), 1, "precondition"); + + // Interject again with a prompt that mentions nothing. + type_input(&mut app, "actually never mind"); + app.interject(); + assert_eq!(app.pending_submit.as_deref(), Some("actually never mind")); + assert!( + app.pending_images.is_empty(), + "stale image carried onto the replacement prompt" + ); + } + + /// The skill branch never touches `pending_images` either, so a skill + /// prompt must not inherit the previous prompt's attachment. + #[test] + fn a_replacement_slash_command_drops_staged_images() { + let (_dir, mut app) = app_in_workspace(); + write_png(&app, "shot.png"); + type_input(&mut app, "look at @shot.png"); + app.submit(); + assert_eq!(app.pending_images.len(), 1, "precondition"); + + app.enqueue_turn_from_command("summarize the repo".into()); + assert!( + app.pending_images.is_empty(), + "stale image carried onto a command-produced turn" + ); + } + + /// Cancelling while the attachments are being read leaves no turn to + /// finish, so the phase has to be reset here or every later prompt + /// queues behind one that never existed. + #[test] + fn cancelling_during_attachment_encoding_returns_to_idle() { + let (_dir, mut app) = app_in_workspace(); + write_png(&app, "shot.png"); + type_input(&mut app, "look at @shot.png"); + app.submit(); + assert_eq!(app.phase, Phase::Streaming, "precondition"); + + // The run loop takes the prompt *and* its descriptors to start the + // encode; the user cancels before the read comes back. + let _ = app.pending_submit.take(); + let _ = std::mem::take(&mut app.pending_images); + app.abandon_staged_attachments(); + + assert_eq!( + app.phase, + Phase::Idle, + "stuck in streaming after a cancelled encode" + ); + assert!(app.pending_images.is_empty()); + assert!(app.pending_attachments.is_empty()); + + // And the next prompt starts a turn instead of queueing behind the + // turn that never was. + type_input(&mut app, "next"); + app.submit(); + assert_eq!(app.pending_submit.as_deref(), Some("next")); + assert!(app.queue.is_empty(), "prompt was queued, not sent"); + } + + /// A filename is attacker-controlled text that this feature puts in + /// the transcript. It reaches the screen as a `System` item, which + /// `render_item` scrubs — this pins that end to end, so the note can + /// never be routed around the sink. + #[test] + fn an_image_note_cannot_smuggle_deceptive_characters() { + let (_dir, mut app) = app_in_workspace(); + let name = "sh\u{202e}gnp.png"; + write_png(&app, name); + type_input(&mut app, &format!("look at @{name}")); + app.submit(); + + let note = app + .transcript + .iter() + .find_map(|i| match i { + TranscriptItem::System(s) if s.contains("@mentions:") => Some(i.clone()), + _ => None, + }) + .expect("no mention note in the transcript"); + let rendered: String = super::super::layout::render_item(¬e, true, false) + .iter() + .flat_map(|l| l.spans.iter().map(|sp| sp.content.to_string())) + .collect(); + assert!( + !rendered.contains('\u{202e}'), + "a filename put a bidi override on the screen: {rendered:?}" + ); + assert!(rendered.contains(""), "not escaped: {rendered:?}"); + } + + fn png_block() -> agent_code_lib::llm::message::ContentBlock { + agent_code_lib::llm::message::ContentBlock::Image { + media_type: "image/png".into(), + data: "iVBORw==".into(), + } + } + + #[test] + fn encoded_attachments_arm_their_own_prompt() { + let (_dir, mut app) = app_in_workspace(); + app.accept_encoded_attachments("look at @shot.png".into(), vec![png_block()]); + assert_eq!(app.pending_submit.as_deref(), Some("look at @shot.png")); + assert_eq!(app.pending_attachments.len(), 1); + } + + /// A slash command that produces a prompt is staged directly, not + /// queued, so it can land while the images are still being read. It + /// must not be overwritten — and must not inherit the other prompt's + /// image either. + #[test] + fn a_prompt_sent_while_encoding_is_not_overwritten() { + let (_dir, mut app) = app_in_workspace(); + app.enqueue_turn_from_command("show me the diff".into()); + assert_eq!(app.pending_submit.as_deref(), Some("show me the diff")); + + app.accept_encoded_attachments("look at @shot.png".into(), vec![png_block()]); + + assert_eq!( + app.pending_submit.as_deref(), + Some("show me the diff"), + "the newer prompt was overwritten" + ); + assert!( + app.pending_attachments.is_empty(), + "the newer prompt inherited another prompt's image" + ); + assert!( + !app.deferred_prompts.is_empty(), + "the superseded prompt was dropped" + ); + assert!( + app.queue.iter().all(|q| q != "look at @shot.png"), + "the expanded prompt was queued as text and would re-expand" + ); + } + + /// The deferred prompt comes back with the blocks it was encoded + /// with. Queuing it as text would re-run `expand_mentions`, reading + /// the file a second time and sending bytes the staged turn never had + /// — after telling the user it had already been read. + #[test] + fn a_deferred_prompt_returns_with_its_original_blocks() { + let (_dir, mut app) = app_in_workspace(); + app.enqueue_turn_from_command("show me the diff".into()); + app.accept_encoded_attachments("look at @shot.png".into(), vec![png_block()]); + + // The turn it was waiting behind starts and finishes. + let _ = app.pending_submit.take(); + app.rearm_deferred_prompt(); + + assert_eq!(app.pending_submit.as_deref(), Some("look at @shot.png")); + assert_eq!( + app.pending_attachments.len(), + 1, + "the deferred prompt lost its images" + ); + assert!( + app.deferred_prompts.is_empty(), + "deferred prompt sent twice" + ); + } + + /// Re-arming must not race another prompt's staged descriptors: those + /// belong to a different turn. + #[test] + fn a_deferred_prompt_waits_for_staged_descriptors_to_clear() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + write_png(&app, "later.png"); + type_input(&mut app, "later @later.png"); + app.submit(); + let _ = app.pending_submit.take(); + assert_eq!(app.pending_images.len(), 1, "precondition"); + + app.rearm_deferred_prompt(); + assert!( + app.pending_submit.is_none(), + "sent a deferred prompt while another's descriptors were staged" + ); + assert!(!app.deferred_prompts.is_empty(), "deferred prompt lost"); + } + + /// Ctrl+C on the turn a prompt is waiting behind must stop that prompt + /// too — otherwise the file goes out after the user stopped the flow. + #[test] + fn cancelling_a_turn_drops_a_prompt_waiting_behind_it() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + // A bare cancel: nothing staged to send. + app.cancel_pending_followups(); + assert!( + app.deferred_prompts.is_empty(), + "an image prompt would have sent itself after a cancel" + ); + } + + /// Interject cancels the live turn in order to send something else, so + /// the held prompt keeps its place behind the interjection. + #[test] + fn interjecting_keeps_a_prompt_waiting_behind_it() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + // Interject only cancels when there is a turn to cancel. + app.phase = Phase::Streaming; + type_input(&mut app, "do this instead"); + app.interject(); + assert_eq!(app.pending_submit.as_deref(), Some("do this instead")); + + app.cancel_pending_followups(); + assert!( + !app.deferred_prompts.is_empty(), + "interject dropped a prompt it only meant to go ahead of" + ); + } + + /// A staged prompt proves nothing on its own: a slash command can put + /// one there at any moment, including while the user is pressing + /// Ctrl+C to stop everything. Only an interject says it cancelled in + /// order to send. + #[test] + fn a_bare_cancel_drops_held_prompts_even_with_a_command_staged() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + // A slash command stages its prompt directly, without interjecting. + app.enqueue_turn_from_command("show me the diff".into()); + assert!(app.pending_submit.is_some(), "precondition"); + + app.cancel_pending_followups(); + assert!( + app.deferred_prompts.is_empty(), + "a bare cancel let an image prompt follow on its own" + ); + } + + /// Queue send-now is the same bargain as interject: it cancels in + /// order to send, so held prompts keep their place. + #[test] + fn queue_send_now_keeps_held_prompts() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + app.queue.push_back("queued work".into()); + app.phase = Phase::Streaming; + app.queue_send_selected(); + + app.cancel_pending_followups(); + assert!( + !app.deferred_prompts.is_empty(), + "queue send-now dropped a held prompt" + ); + } + + /// Dropping the read in flight says nothing about prompts already + /// held: whether those survive is the cancel policy's call, and an + /// interject that lands mid-encode must not lose them. + #[test] + fn abandoning_a_read_leaves_held_prompts_alone() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + app.phase = Phase::Streaming; + type_input(&mut app, "do this instead"); + app.interject(); + + app.cancel_pending_followups(); + app.abandon_staged_attachments(); + + assert_eq!( + app.deferred_prompts.len(), + 1, + "an interject during an encode lost a held prompt" + ); + } + + /// Held prompts are sent in the order they were submitted, so a second + /// one cannot overwrite the first. + #[test] + fn held_prompts_keep_their_order() { + let (_dir, mut app) = app_in_workspace(); + app.enqueue_turn_from_command("busy".into()); + app.accept_encoded_attachments("first @a.png".into(), vec![png_block()]); + app.accept_encoded_attachments("second @b.png".into(), vec![png_block()]); + assert_eq!(app.deferred_prompts.len(), 2, "a held prompt was dropped"); + + let _ = app.pending_submit.take(); + app.rearm_deferred_prompt(); + assert_eq!(app.pending_submit.as_deref(), Some("first @a.png")); + let _ = app.pending_submit.take(); + app.rearm_deferred_prompt(); + assert_eq!(app.pending_submit.as_deref(), Some("second @b.png")); + } + + /// A cleared, resumed or rewound conversation takes its attachments + /// with it: a file staged for the old thread must not surface in the + /// new one. + #[test] + fn a_replaced_conversation_drops_staged_attachments() { + let (_dir, mut app) = app_in_workspace(); + app.deferred_prompts + .push_back(("earlier @a.png".into(), vec![png_block()])); + app.pending_attachments = vec![png_block()]; + write_png(&app, "shot.png"); + type_input(&mut app, "look at @shot.png"); + app.submit(); + assert_eq!(app.pending_images.len(), 1, "precondition"); + + app.new_conversation(); + + assert!(app.deferred_prompts.is_empty(), "held prompt survived"); + assert!(app.pending_images.is_empty(), "descriptors survived"); + assert!(app.pending_attachments.is_empty(), "blocks survived"); + } + + /// A cancel abandons only the prompt that was being read for. If the + /// user interjected with another image in the meantime, that prompt is + /// still coming and must keep its own attachment — clearing it would + /// send the interjection as a text-only turn. + #[test] + fn abandoning_a_cancelled_encode_keeps_a_replacement_prompts_images() { + let (_dir, mut app) = app_in_workspace(); + write_png(&app, "first.png"); + write_png(&app, "second.png"); + type_input(&mut app, "look at @first.png"); + app.submit(); + + // The run loop takes the first prompt and its descriptors. + let _ = app.pending_submit.take(); + let _ = std::mem::take(&mut app.pending_images); + + // The user interjects with a second image-bearing prompt, then the + // discarded first read lands. + type_input(&mut app, "actually @second.png"); + app.interject(); + assert_eq!(app.pending_images.len(), 1, "precondition"); + app.abandon_staged_attachments(); + + assert_eq!( + app.pending_submit.as_deref(), + Some("actually @second.png"), + "replacement prompt lost" + ); + assert_eq!( + app.pending_images.len(), + 1, + "replacement prompt was stripped of its image" + ); + assert_eq!( + app.phase, + Phase::Streaming, + "a prompt is still waiting for its turn" + ); + } + + /// Encoded bytes are staged separately from the descriptors, so a + /// replaced prompt has to drop both — otherwise an image that had + /// already been read would still ride along with the new prompt. + #[test] + fn replacing_a_prompt_drops_already_encoded_attachments() { + let (_dir, mut app) = app_in_workspace(); + app.pending_attachments = vec![agent_code_lib::llm::message::ContentBlock::Image { + media_type: "image/png".into(), + data: "iVBORw==".into(), + }]; + type_input(&mut app, "a different prompt"); + app.submit(); + assert!( + app.pending_attachments.is_empty(), + "encoded image carried onto the replacement prompt" + ); + } } diff --git a/crates/cli/src/ui/modern/mentions.rs b/crates/cli/src/ui/modern/mentions.rs index 9e3134b5..322bb89d 100644 --- a/crates/cli/src/ui/modern/mentions.rs +++ b/crates/cli/src/ui/modern/mentions.rs @@ -26,6 +26,21 @@ pub const MAX_FILE_BYTES: usize = 64 * 1024; /// the worst case at four full-size files; the rest are skipped with a note. pub const MAX_TOTAL_BYTES: usize = 256 * 1024; +/// Per-image cap. An attachment is read whole and base64-encoded on the UI +/// thread, and base64 inflates by 4/3, so 3 MiB is the largest file that +/// still fits the 5 MB per-image payload the providers accept. +pub const MAX_IMAGE_BYTES: usize = 3 * 1024 * 1024; + +/// Cap across every image in one prompt. Images do not consume the text +/// budget — they never enter the prompt string — so they need a budget of +/// their own; without one `@*.png` over a screenshot directory can freeze +/// or OOM the CLI before the request is ever built. +pub const MAX_TOTAL_IMAGE_BYTES: usize = 8 * 1024 * 1024; + +/// Upper bound on attachments per prompt, independent of their size: each +/// one costs a full re-encode and a large block of context. +pub const MAX_IMAGES: usize = 4; + /// Upper bound on directory entries examined for one completion, so a /// pathological directory cannot stall the UI thread. pub const MAX_SCAN_ENTRIES: usize = 4_000; @@ -160,12 +175,23 @@ pub fn mention_text(candidate: &str) -> String { } /// Result of inlining `@path` mentions into a prompt. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone)] pub struct MentionExpansion { /// What the engine should receive: the user's text plus file blocks. pub prompt: String, /// Short human-readable notes about anything skipped or truncated. pub notes: Vec, + /// Images to attach to the turn, held open. An image cannot be inlined + /// as text, so `@shot.png` used to report "binary, skipped" — which is + /// never what the user meant by mentioning one. + /// + /// Descriptors, not paths. A path handed to the turn would be resolved + /// a second time, and swapping the file — or any ancestor directory — + /// for a symlink in between pointed that second lookup at a file + /// outside the workspace which had never been validated. Opening while + /// the mention is being checked leaves nothing to look up again; the + /// bytes are read later, off the UI thread, from these descriptors. + pub images: Vec, } /// Inline the contents of every `@path` mention in `text`. @@ -182,8 +208,11 @@ pub fn expand_mentions(text: &str, cwd: &Path) -> Option { let mut seen: HashSet = HashSet::new(); let mut body = String::new(); let mut notes: Vec = Vec::new(); + let mut images: Vec = Vec::new(); let mut used = 0usize; + let mut image_bytes = 0usize; let mut over_budget = 0usize; + let mut over_image_budget = 0usize; for raw in mentions { if !seen.insert(raw.clone()) { @@ -196,6 +225,28 @@ pub fn expand_mentions(text: &str, cwd: &Path) -> Option { continue; } }; + // Attached rather than inlined; the model receives it as an image + // block on the turn. Budgeted before a byte is read — an image is + // held whole and base64-encoded, so an unbounded one would freeze + // or OOM the UI thread. + if let Resolved::File(ref path) = resolved + && is_image(path) + { + if images.len() >= MAX_IMAGES { + over_image_budget += 1; + continue; + } + match stage_image(cwd, path) { + Ok((staged, len)) if image_bytes + len <= MAX_TOTAL_IMAGE_BYTES => { + image_bytes += len; + notes.push(format!("@{raw} — attached as an image")); + images.push(staged); + } + Ok(_) => over_image_budget += 1, + Err(reason) => notes.push(format!("@{raw} — {reason}")), + } + continue; + } let remaining = MAX_TOTAL_BYTES.saturating_sub(used); if remaining == 0 { over_budget += 1; @@ -234,10 +285,17 @@ pub fn expand_mentions(text: &str, cwd: &Path) -> Option { MAX_TOTAL_BYTES / 1024 )); } + if over_image_budget > 0 { + notes.push(format!( + "{over_image_budget} image(s) skipped — at most {MAX_IMAGES} images / {} MiB per prompt", + MAX_TOTAL_IMAGE_BYTES / (1024 * 1024) + )); + } Some(MentionExpansion { prompt: format!("{text}{body}"), notes, + images, }) } @@ -294,6 +352,305 @@ fn resolve_mention(cwd: &Path, raw: &str) -> Result { } } +/// Extensions the model can be shown directly. Kept to the formats the +/// API accepts, so an unsupported image still reports a clear reason +/// rather than being attached and rejected upstream. +const IMAGE_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp"]; + +fn is_image(path: &Path) -> bool { + path.extension() + .and_then(|e| e.to_str()) + .map(|e| e.to_ascii_lowercase()) + .is_some_and(|e| IMAGE_EXTENSIONS.contains(&e.as_str())) +} + +/// A validated image, held open until the turn that carries it starts. +/// +/// The descriptor *is* the validation result. Handing the turn a path +/// would mean resolving that name a second time, and a name resolved +/// twice can mean two different files — the whole point of opening here. +#[derive(Debug, Clone)] +pub struct StagedImage { + /// Kept for the media type and for error messages only; never + /// re-opened. + pub path: PathBuf, + pub file: std::sync::Arc, +} + +/// Open `path` without letting any component of it be redirected. +/// +/// `path` must already be canonical and inside `root`. Each component is +/// opened relative to the descriptor of the one before it, refusing +/// symlinks — so replacing an ancestor directory (or the file itself) +/// between validation and this open cannot walk the read outside the +/// workspace. Resolving the pathname again instead would re-run the whole +/// lookup against a tree the attacker has had time to rearrange. +/// +/// The final `fstat` is on the descriptor, so it describes the file that +/// was actually opened rather than whatever the name means afterwards. +#[cfg(unix)] +fn open_beneath(root: &Path, path: &Path) -> Result { + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt; + use std::os::unix::io::{AsRawFd, FromRawFd}; + + let rel = path + .strip_prefix(root) + .map_err(|_| "outside the workspace".to_string())?; + + // The workspace root is the trust anchor, so it is opened with the same + // suspicion as every component below it: `root` is canonical, so its + // final component is a real directory unless someone replaced it since + // — which is exactly what `O_NOFOLLOW` refuses. A session started in a + // symlinked directory still works, because canonicalization resolved + // that symlink before this point. + let mut dir = { + use std::os::unix::fs::OpenOptionsExt; + std::fs::OpenOptions::new() + .read(true) + .custom_flags(libc::O_DIRECTORY | libc::O_NOFOLLOW | libc::O_CLOEXEC) + .open(root) + .map_err(|e| format!("unreadable ({})", e.kind()))? + }; + + let components: Vec<_> = rel.components().collect(); + let Some((last, parents)) = components.split_last() else { + return Err("not a regular file".into()); + }; + for component in components.iter() { + // A canonical path relative to its own prefix has only normal + // components; anything else means the assumption broke. + if !matches!(component, std::path::Component::Normal(_)) { + return Err("not a usable path".into()); + } + } + + for component in parents { + let name = CString::new(component.as_os_str().as_bytes()) + .map_err(|_| "not a usable path".to_string())?; + // SAFETY: `dir` is an open directory descriptor and `name` is a + // valid NUL-terminated path for the duration of the call. + let fd = unsafe { + libc::openat( + dir.as_raw_fd(), + name.as_ptr(), + libc::O_RDONLY | libc::O_DIRECTORY | libc::O_NOFOLLOW | libc::O_CLOEXEC, + ) + }; + if fd < 0 { + return Err(format!( + "unreadable ({})", + std::io::Error::last_os_error().kind() + )); + } + // SAFETY: `fd` was just returned by `openat` and is owned here. + dir = unsafe { std::fs::File::from_raw_fd(fd) }; + } + + let name = + CString::new(last.as_os_str().as_bytes()).map_err(|_| "not a usable path".to_string())?; + // SAFETY: as above; `O_NONBLOCK` additionally keeps a FIFO swapped in + // for the file from blocking this call until a writer appears. + let fd = unsafe { + libc::openat( + dir.as_raw_fd(), + name.as_ptr(), + libc::O_RDONLY | libc::O_NOFOLLOW | libc::O_NONBLOCK | libc::O_CLOEXEC, + ) + }; + if fd < 0 { + return Err(format!( + "unreadable ({})", + std::io::Error::last_os_error().kind() + )); + } + // SAFETY: `fd` was just returned by `openat` and is owned here. + let file = unsafe { std::fs::File::from_raw_fd(fd) }; + if !file + .metadata() + .map_err(|e| format!("unreadable ({})", e.kind()))? + .is_file() + { + return Err("not a regular file".into()); + } + Ok(file) +} + +/// True when this entry is a reparse point of any kind. +/// +/// The attribute, not `FileType::is_symlink`: that predicate keys off the +/// name-surrogate bit in the reparse tag, which is a property of how the +/// tag is meant to be interpreted rather than of whether traversing the +/// entry leaves the workspace. Testing the attribute refuses every reparse +/// tag — junction, symlink, and whatever else the filesystem grows — +/// which is the only answer that stays correct as tags are added. +#[cfg(not(unix))] +fn is_reparse_point(meta: &std::fs::Metadata) -> bool { + #[cfg(windows)] + { + use std::os::windows::fs::MetadataExt; + const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x0000_0400; + return meta.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0; + } + #[cfg(not(windows))] + { + meta.file_type().is_symlink() + } +} + +/// Windows has no `openat`, so the ancestors are checked explicitly. +/// +/// `FILE_FLAG_OPEN_REPARSE_POINT` covers the final component: a symlink or +/// junction swapped in for the file is opened *as* the reparse point and +/// fails the regular-file check below rather than redirecting the read. +/// That flag does nothing for the directories above it, and creating a +/// junction needs no special privilege, so every ancestor between the +/// workspace root and the file is rejected outright if it has become a +/// reparse point. +/// +/// Unlike the Unix path this is not atomic: the ancestor checks and the +/// open are separate lookups, and Win32 has no handle-relative open to +/// close that gap — only `NtCreateFile` with a root directory handle can, +/// which is undocumented-adjacent FFI this crate does not otherwise carry. +/// What remains is a local attacker who can already write inside the +/// workspace and must win a race, and who could instead simply put the +/// bytes they want in an image file there and have it attached with no +/// race at all. The bound worth having — that a mention cannot *quietly* +/// reach outside the workspace — is what the checks above provide. +#[cfg(not(unix))] +fn open_beneath(root: &Path, path: &Path) -> Result { + const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x0020_0000; + let rel = path + .strip_prefix(root) + .map_err(|_| "outside the workspace".to_string())?; + + let mut ancestor = root.to_path_buf(); + for component in rel.components() { + if !matches!(component, std::path::Component::Normal(_)) { + return Err("not a usable path".into()); + } + ancestor.push(component); + if ancestor == path { + break; + } + let meta = std::fs::symlink_metadata(&ancestor) + .map_err(|e| format!("unreadable ({})", e.kind()))?; + if is_reparse_point(&meta) { + return Err("not a regular file".into()); + } + } + + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); + } + let file = options + .open(path) + .map_err(|e| format!("unreadable ({})", e.kind()))?; + let meta = file + .metadata() + .map_err(|e| format!("unreadable ({})", e.kind()))?; + // The descriptor decides, and a reparse point is judged by its + // attribute: `is_file` only excludes the tags `is_symlink` recognizes, + // so a placeholder or any other non-name-surrogate tag would otherwise + // be accepted here as an ordinary file. + if is_reparse_point(&meta) || !meta.is_file() { + return Err("not a regular file".into()); + } + Ok(file) +} + +/// Stage a validated image: open it now, and measure it from the +/// descriptor so an oversized file is refused before anything is read. +fn stage_image(cwd: &Path, path: &Path) -> Result<(StagedImage, usize), String> { + let root = cwd.canonicalize().unwrap_or_else(|_| cwd.to_path_buf()); + let file = open_beneath(&root, path)?; + let len = file + .metadata() + .map_err(|e| format!("unreadable ({})", e.kind()))? + .len(); + let len = usize::try_from(len).map_err(|_| "image too large".to_string())?; + if len > MAX_IMAGE_BYTES { + return Err(format!( + "image too large ({:.1} MiB, max {} MiB)", + len as f64 / (1024.0 * 1024.0), + MAX_IMAGE_BYTES / (1024 * 1024) + )); + } + Ok(( + StagedImage { + path: path.to_path_buf(), + file: std::sync::Arc::new(file), + }, + len, + )) +} + +/// Read and encode staged images from the descriptors already held. +/// +/// Blocking work — call it off the UI thread. No path is resolved here: +/// the files were opened when their mentions were validated, so this reads +/// the bytes of those files and nothing else, however the workspace has +/// been rearranged since. +pub fn encode_staged_images( + images: Vec, +) -> (Vec, Vec) { + use std::io::{Seek, SeekFrom}; + + let mut blocks = Vec::new(); + let mut notes = Vec::new(); + let mut total = 0usize; + for image in images.into_iter().take(MAX_IMAGES) { + let name = image.path.display(); + let mut handle: &std::fs::File = &image.file; + // Rewind: a turn that failed to spawn is retried with the same + // descriptors, and a spent offset would re-encode an empty file. + if let Err(e) = handle.seek(SeekFrom::Start(0)) { + notes.push(format!( + "could not attach {name}: unreadable ({})", + e.kind() + )); + continue; + } + let mut data = Vec::new(); + // Capped again: the file can still grow after it was measured, and + // this is the read that would actually hold the bytes. + if let Err(e) = handle + .take(MAX_IMAGE_BYTES as u64 + 1) + .read_to_end(&mut data) + { + notes.push(format!( + "could not attach {name}: unreadable ({})", + e.kind() + )); + continue; + } + if data.len() > MAX_IMAGE_BYTES { + notes.push(format!( + "could not attach {name}: image too large (over {} MiB)", + MAX_IMAGE_BYTES / (1024 * 1024) + )); + continue; + } + if total + data.len() > MAX_TOTAL_IMAGE_BYTES { + notes.push(format!( + "could not attach {name}: {} MiB total image limit reached", + MAX_TOTAL_IMAGE_BYTES / (1024 * 1024) + )); + continue; + } + total += data.len(); + blocks.push(agent_code_lib::llm::message::image_block_from_bytes( + &image.path, + &data, + )); + } + (blocks, notes) +} + /// True when `path` resolves inside `cwd`. Both sides are canonicalized. fn contained_in(cwd: &Path, path: &Path) -> bool { let cwd_canon = cwd.canonicalize().unwrap_or_else(|_| cwd.to_path_buf()); @@ -413,6 +770,7 @@ fn truncate_utf8(s: String, cap: usize) -> (String, Option) { #[cfg(test)] mod tests { use super::*; + use agent_code_lib::llm::message::ContentBlock; use std::fs; fn fixture() -> tempfile::TempDir { @@ -664,6 +1022,305 @@ mod tests { assert_eq!(out.prompt.matches(")>(); + // Image attachments are read and encoded the same way: detached, with + // the result landing in a select arm. Awaiting the work inline would + // park this loop — the only one there is — so a workspace on a slow + // mount would stop redraws and Ctrl+C until the read finished. + #[allow(clippy::type_complexity)] + let (img_tx, mut img_rx) = tokio::sync::mpsc::unbounded_channel::<( + u64, + u64, + String, + Vec, + Vec, + )>(); + // Identifies the encode in flight, if any: the turn it belongs to must + // not start without it, and a second one must not be queued behind it. + // A cancel forgets the id rather than waiting — the read cannot be + // stopped, so its result is recognised as stale when it lands and the + // staged turn is released immediately. + let mut encode_seq = 0u64; + let mut active_encode: Option = None; + // The conversation the loop last saw, so a `/clear`, `/resume` or + // `/rewind` can be noticed the moment it happens rather than when a + // read that may never finish comes back. + let mut seen_epoch = app.conversation_epoch; // Seed the pane once so tasks adopted from a previous process show // before the first turn arms the periodic poll. app.sync_background_tasks(manager_rows(&task_manager).await); @@ -584,10 +607,61 @@ pub(super) async fn event_loop( } } + // A replaced conversation invalidates a read in flight at once: + // waiting for it would hold every prompt in the new conversation + // behind an encode that belongs to a conversation nobody is + // looking at any more. + if app.conversation_epoch != seen_epoch { + seen_epoch = app.conversation_epoch; + if active_encode.take().is_some() { + app.abandon_staged_attachments(); + } + } + + // A prompt that was held aside for a turn that has since gone can + // send now, with the blocks it was already encoded with. + if turn.is_none() && active_encode.is_none() { + app.rearm_deferred_prompt(); + } + + // Hand a prompt's images to the blocking pool. The descriptors were + // opened when their mentions were validated, so nothing is resolved + // here; the result comes back through `img_rx` and re-arms the + // prompt, which keeps this loop free to redraw and to take a Ctrl+C + // while a slow mount is being read. + if turn.is_none() + && active_encode.is_none() + && !app.pending_images.is_empty() + && let Some(prompt) = app.pending_submit.take() + { + let images = std::mem::take(&mut app.pending_images); + let tx = img_tx.clone(); + encode_seq += 1; + let id = encode_seq; + active_encode = Some(id); + // Stamped with the conversation it was submitted in: the engine + // lock is free while this runs, so `/clear`, `/resume` or + // `/rewind` can replace the conversation underneath it. + let epoch = app.conversation_epoch; + tokio::task::spawn_blocking(move || { + let (blocks, notes) = super::mentions::encode_staged_images(images); + let _ = tx.send((id, epoch, prompt, blocks, notes)); + }); + } + // Start a pending turn if idle. if turn.is_none() + && active_encode.is_none() && let Some(prompt) = app.pending_submit.take() { + let blocks = std::mem::take(&mut app.pending_attachments); + // Set unconditionally, awaiting the lock rather than skipping on + // contention: an empty set clears anything a previous attempt + // staged, so no turn can inherit another turn's attachment. + { + let engine = session.engine(); + engine.lock().await.set_pending_attachments(blocks.clone()); + } let sink = ChannelSink::new(eng_tx.clone(), app.conversation_epoch); match session.spawn_turn(prompt.clone(), sink).await { Ok(handle) => { @@ -596,8 +670,10 @@ pub(super) async fn event_loop( } Err(e) => { // Should be rare: TUI serializes turns. Put the prompt - // back so the next idle loop can retry. + // and its attachments back so the next idle loop retries + // them together. app.pending_submit = Some(prompt); + app.pending_attachments = blocks; app.status_message = format!("turn busy: {e}"); app.dirty = true; } @@ -609,6 +685,17 @@ pub(super) async fn event_loop( if let Some(ref h) = turn { h.cancel(); } + // Nothing may follow on its own after a cancel — including a + // prompt held aside behind the turn being cancelled. Interject + // is the exception and says so by having staged its own prompt. + app.cancel_pending_followups(); + // A read already handed to the blocking pool cannot be stopped. + // Forget its id instead: the result is stale when it lands, and + // the staged turn is released now rather than after a read that + // may never finish. + if active_encode.take().is_some() { + app.abandon_staged_attachments(); + } app.cancel_requested = false; } @@ -682,6 +769,12 @@ pub(super) async fn event_loop( } app.mark_turn_idle(); + // A prompt held aside with its images was submitted before + // anything in the queue, so it goes first — and it goes now + // rather than waiting for whatever event next wakes the + // loop, since nothing else would arm it. + app.rearm_deferred_prompt(); + // Queue handling (plan §M5): auto-send the head on a clean // finish; on abort/error keep the queue and tell the user. // Interject leaves `pending_submit` set so we start it even @@ -828,6 +921,29 @@ pub(super) async fn event_loop( Some((id, out)) = task_out_rx.recv() => { app.show_task_output(&id, out); } + // Encoded image attachments coming back from the blocking pool. + // The prompt is re-armed with them so the turn starts on the + // next pass through the loop above. + Some((id, epoch, prompt, blocks, notes)) = img_rx.recv() => { + if active_encode != Some(id) { + // Cancelled while this read was in flight: the state it + // belonged to was released then, so the bytes are simply + // dropped rather than sent with a turn nobody asked for. + } else if epoch != app.conversation_epoch { + // The conversation it was submitted in has been cleared, + // resumed or rewound. Starting it now would attach the + // file to a thread the user never attached it to. + active_encode = None; + app.abandon_staged_attachments(); + } else { + active_encode = None; + for note in notes { + app.transcript.push(super::app::TranscriptItem::System(note)); + } + app.accept_encoded_attachments(prompt, blocks); + } + app.dirty = true; + } // Background-task rows (`&` shell jobs, workflows, monitors). // Gated on work that can still change: polling while any // rows exist at all would tick forever once a subagent row diff --git a/crates/lib/src/llm/message.rs b/crates/lib/src/llm/message.rs index 05c7e23c..28b7c4ee 100644 --- a/crates/lib/src/llm/message.rs +++ b/crates/lib/src/llm/message.rs @@ -187,6 +187,18 @@ impl ContentBlock { } } + /// The wire `type` tag of this block, for logs and hook payloads. + pub fn kind_name(&self) -> &'static str { + match self { + ContentBlock::Text { .. } => "text", + ContentBlock::ToolUse { .. } => "tool_use", + ContentBlock::ToolResult { .. } => "tool_result", + ContentBlock::Thinking { .. } => "thinking", + ContentBlock::Image { .. } => "image", + ContentBlock::Document { .. } => "document", + } + } + /// Extract tool use info, if this is a tool_use block. pub fn as_tool_use(&self) -> Option<(&str, &str, &serde_json::Value)> { match self { @@ -240,6 +252,29 @@ pub enum StopReason { } /// Helper to create a user message with text content. +/// A user message carrying attachments alongside its text. +/// +/// Images go *before* the text: a model reads the prompt as being about +/// the images it has just been shown, and the reverse order reads as an +/// afterthought. +pub fn user_message_with_attachments( + text: impl Into, + attachments: Vec, +) -> Message { + let text = text.into(); + let mut content = attachments; + if !text.is_empty() { + content.push(ContentBlock::Text { text }); + } + Message::User(UserMessage { + uuid: Uuid::new_v4(), + timestamp: chrono::Utc::now().to_rfc3339(), + content, + is_meta: false, + is_compact_summary: false, + }) +} + pub fn user_message(text: impl Into) -> Message { Message::User(UserMessage { uuid: Uuid::new_v4(), @@ -250,86 +285,105 @@ pub fn user_message(text: impl Into) -> Message { }) } +/// Describe attachment blocks for hook payloads and logs. +/// +/// Metadata, not payloads: a hook needs to know that an image of a given +/// type and size is going out — enough to audit or refuse it — and putting +/// megabytes of base64 into every hook invocation would serve nothing. +pub fn describe_attachments(blocks: &[ContentBlock]) -> Vec { + blocks + .iter() + .map(|block| match block { + ContentBlock::Image { media_type, data } => serde_json::json!({ + "type": "image", + "media_type": media_type, + "encoded_bytes": data.len(), + }), + ContentBlock::Document { + media_type, data, .. + } => serde_json::json!({ + "type": "document", + "media_type": media_type, + "encoded_bytes": data.len(), + }), + other => serde_json::json!({ "type": other.kind_name() }), + }) + .collect() +} + +/// Media type for an image path, inferred from its extension. +/// +/// Case-insensitive: callers that decide *whether* a path is an image +/// normalize the extension, so `shot.PNG` must not fall through to a +/// generic media type the provider then rejects. +pub fn image_media_type(path: &std::path::Path) -> Option<&'static str> { + let ext = path.extension()?.to_str()?.to_ascii_lowercase(); + Some(match ext.as_str() { + "png" => "image/png", + "jpg" | "jpeg" => "image/jpeg", + "gif" => "image/gif", + "webp" => "image/webp", + "svg" => "image/svg+xml", + _ => return None, + }) +} + /// Helper to create an image content block from a file path. /// /// Reads the file, base64-encodes it, and infers the media type /// from the file extension. pub fn image_block_from_file(path: &std::path::Path) -> Result { let data = std::fs::read(path).map_err(|e| format!("Failed to read image: {e}"))?; - - let media_type = match path.extension().and_then(|e| e.to_str()) { - Some("png") => "image/png", - Some("jpg" | "jpeg") => "image/jpeg", - Some("gif") => "image/gif", - Some("webp") => "image/webp", - Some("svg") => "image/svg+xml", - _ => "application/octet-stream", - }; - - use std::io::Write; - let mut encoded = String::new(); - { - let mut encoder = base64_encode_writer(&mut encoded); - encoder - .write_all(&data) - .map_err(|e| format!("base64 error: {e}"))?; - } - - Ok(ContentBlock::Image { - media_type: media_type.to_string(), - data: encoded, - }) + Ok(image_block_from_bytes(path, &data)) } -/// Simple base64 encoder (no external dependency). -fn base64_encode_writer(output: &mut String) -> Base64Writer<'_> { - Base64Writer { - output, - buffer: Vec::new(), +/// Build an image block from bytes already in hand. +/// +/// Lets a caller that must bound the read do it itself and still get the +/// same media-type inference; `path` is used only for its extension. +pub fn image_block_from_bytes(path: &std::path::Path, data: &[u8]) -> ContentBlock { + ContentBlock::Image { + media_type: image_media_type(path) + .unwrap_or("application/octet-stream") + .to_string(), + data: base64_encode(data), } } -struct Base64Writer<'a> { - output: &'a mut String, - buffer: Vec, -} - -impl<'a> std::io::Write for Base64Writer<'a> { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.buffer.extend_from_slice(buf); - Ok(buf.len()) - } - fn flush(&mut self) -> std::io::Result<()> { - const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; - let mut i = 0; - while i + 2 < self.buffer.len() { - let b0 = self.buffer[i] as usize; - let b1 = self.buffer[i + 1] as usize; - let b2 = self.buffer[i + 2] as usize; - self.output.push(CHARS[b0 >> 2] as char); - self.output.push(CHARS[((b0 & 3) << 4) | (b1 >> 4)] as char); - self.output - .push(CHARS[((b1 & 0xf) << 2) | (b2 >> 6)] as char); - self.output.push(CHARS[b2 & 0x3f] as char); - i += 3; +/// Base64 (RFC 4648, padded) without an external dependency. +/// +/// A whole-input function rather than an `io::Write` adaptor: the writer +/// this replaced only encoded from `flush()`, and every caller dropped it +/// after `write_all`, so every image block shipped with an empty payload. +fn base64_encode(data: &[u8]) -> String { + const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + let mut out = String::with_capacity(data.len().div_ceil(3) * 4); + let mut chunks = data.chunks_exact(3); + for c in &mut chunks { + let (b0, b1, b2) = (c[0] as usize, c[1] as usize, c[2] as usize); + out.push(CHARS[b0 >> 2] as char); + out.push(CHARS[((b0 & 3) << 4) | (b1 >> 4)] as char); + out.push(CHARS[((b1 & 0xf) << 2) | (b2 >> 6)] as char); + out.push(CHARS[b2 & 0x3f] as char); + } + match *chunks.remainder() { + [b0] => { + let b0 = b0 as usize; + out.push(CHARS[b0 >> 2] as char); + out.push(CHARS[(b0 & 3) << 4] as char); + out.push('='); + out.push('='); } - let remaining = self.buffer.len() - i; - if remaining == 1 { - let b0 = self.buffer[i] as usize; - self.output.push(CHARS[b0 >> 2] as char); - self.output.push(CHARS[(b0 & 3) << 4] as char); - self.output.push('='); - self.output.push('='); - } else if remaining == 2 { - let b0 = self.buffer[i] as usize; - let b1 = self.buffer[i + 1] as usize; - self.output.push(CHARS[b0 >> 2] as char); - self.output.push(CHARS[((b0 & 3) << 4) | (b1 >> 4)] as char); - self.output.push(CHARS[(b1 & 0xf) << 2] as char); - self.output.push('='); + [b0, b1] => { + let (b0, b1) = (b0 as usize, b1 as usize); + out.push(CHARS[b0 >> 2] as char); + out.push(CHARS[((b0 & 3) << 4) | (b1 >> 4)] as char); + out.push(CHARS[(b1 & 0xf) << 2] as char); + out.push('='); } - Ok(()) + _ => {} } + out } /// Helper to create a user message with an image. @@ -505,6 +559,138 @@ pub fn messages_to_api_params_cached(messages: &[Message]) -> Vec = (0..=255u8).collect(); + let encoded = base64_encode(&all); + assert_eq!(encoded.len(), 344, "256 bytes encode to 344 base64 chars"); + assert!(encoded.starts_with("AAECAwQF")); + assert!(encoded.ends_with("+/w==")); + } + + #[test] + fn an_attachment_with_no_text_carries_only_the_attachment() { + let img = ContentBlock::Image { + media_type: "image/png".into(), + data: "abc".into(), + }; + let msg = user_message_with_attachments("", vec![img]); + let Message::User(u) = msg else { + panic!("expected a user message"); + }; + assert_eq!(u.content.len(), 1, "an empty text block was appended"); + } + use super::*; #[test] diff --git a/crates/lib/src/query/mod.rs b/crates/lib/src/query/mod.rs index b0f5e12f..69ef0143 100644 --- a/crates/lib/src/query/mod.rs +++ b/crates/lib/src/query/mod.rs @@ -74,6 +74,9 @@ pub struct QueryEngine { last_seen_denial_total: usize, extraction_state: Arc>, session_allows: Arc>>, + /// Content blocks to prepend to the next user message (images from + /// the composer). Consumed by the turn that follows. + pending_attachments: Vec, /// Grants that persist across sessions. `None` until a host opts in /// via [`Self::set_persistent_grants`] — library embedders get the /// previous behaviour (session-scoped only) unless they ask for it. @@ -237,6 +240,7 @@ impl QueryEngine { crate::memory::extraction::ExtractionState::new(), )), session_allows: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())), + pending_attachments: Vec::new(), persistent_grants: None, permission_prompter: None, question_asker: None, @@ -289,6 +293,14 @@ impl QueryEngine { sink.on_context_usage(used, DEFAULT_CONTEXT_WINDOW); } + /// Attach content blocks to the next user turn. + /// + /// Replaces rather than appends, so a cancelled composer cannot + /// accumulate images across attempts. + pub fn set_pending_attachments(&mut self, blocks: Vec) { + self.pending_attachments = blocks; + } + /// Install the interactive permission prompter. /// /// Without this, an `Ask` permission decision falls through to auto-allow @@ -796,14 +808,49 @@ impl QueryEngine { user_input: &str, sink: &dyn StreamSink, ) -> crate::error::Result<()> { - // Add the user message to history. - let user_msg = user_message(user_input); + // Add the user message to history, with anything the caller + // attached for this turn. Taken rather than read: an attachment + // belongs to exactly one turn, and leaking it into the next one + // would re-send an image the user already shared. + let attachments = std::mem::take(&mut self.pending_attachments); + // Described for the hooks below before the blocks are moved into + // the message. A `UserPromptSubmit` hook exists to see everything + // the turn sends; an image that only showed up as silence there + // would slip past exactly the audit it was configured for. + let attachment_info = crate::llm::message::describe_attachments(&attachments); + let user_msg = if attachments.is_empty() { + user_message(user_input) + } else { + crate::llm::message::user_message_with_attachments(user_input, attachments) + }; self.state.push_message(user_msg); + // Images are the one block a token threshold cannot see: each is + // charged a flat vision estimate, so compaction never fires on the + // megabytes they actually add to every later request. Bound what + // history retains here instead, once the new turn's own images are + // in and counted as the most recent. + let dropped = compact::evict_old_images( + &mut self.state.messages, + compact::MAX_RETAINED_IMAGE_BYTES, + compact::MAX_RETAINED_IMAGES, + ); + if dropped > 0 { + tracing::debug!( + dropped, + "evicted image payloads beyond the retention budget" + ); + } // UserPromptSubmit fires once per user turn, as soon as the // prompt is in history and before any PreTurn / LLM work. Hooks // at this event see the FULL prompt (no truncation) so they can // do content scanning, redaction logging, or compliance audits. + // + // Observation only: unlike PreToolUse, a non-zero exit here does + // not veto the turn. Wiring that up would start blocking prompts + // for every deployment whose prompt hook happens to exit non-zero + // today — a `grep` that finds nothing is enough — so it is a + // deliberate change for its own PR, not a side effect of this one. let _user_prompt_submit_results = self .hooks .run_hooks( @@ -811,6 +858,7 @@ impl QueryEngine { None, &serde_json::json!({ "user_input": user_input, + "attachments": attachment_info, "turn": self.state.turn_count + 1, }), Some(&self.cancel), @@ -831,6 +879,7 @@ impl QueryEngine { &serde_json::json!({ "turn": self.state.turn_count + 1, "user_input_preview": user_input.chars().take(200).collect::(), + "attachments": attachment_info, }), Some(&self.cancel), ) diff --git a/crates/lib/src/services/compact.rs b/crates/lib/src/services/compact.rs index bf8e6262..af2f3c96 100644 --- a/crates/lib/src/services/compact.rs +++ b/crates/lib/src/services/compact.rs @@ -308,6 +308,86 @@ pub fn estimate_compactable_tokens(messages: &[Message], keep_recent: usize) -> freed } +/// Encoded image bytes kept in history, newest first. +/// +/// Images are the one block whose cost is measured in megabytes rather +/// than tokens: a screenshot is charged a flat vision estimate, so no +/// token threshold notices it, yet every later request reserializes and +/// resends the whole payload. A handful of screenshots in a long session +/// is enough to push requests past what a provider will accept. 12 MiB +/// keeps a few recent images available to the model while bounding what +/// the session can carry. +pub const MAX_RETAINED_IMAGE_BYTES: usize = 12 * 1024 * 1024; + +/// Image blocks kept in history, newest first. +/// +/// Bytes are not the only limit: providers also cap how many images one +/// request may carry (Anthropic accepts 100), and a session that attaches +/// small thumbnails every turn reaches that count long before it reaches +/// any byte budget. 32 leaves generous headroom under the provider cap. +pub const MAX_RETAINED_IMAGES: usize = 32; + +/// Drop image payloads once history holds more than +/// [`MAX_RETAINED_IMAGE_BYTES`] or [`MAX_RETAINED_IMAGES`] of them, +/// oldest first. +/// +/// What survives is a *contiguous run of the newest* images: as soon as +/// one does not fit, everything older goes too. Continuing to look for +/// smaller images that would fit the remaining space would keep stale +/// context alive while dropping something newer and more likely to be +/// under discussion. +/// +/// The block is replaced by text naming what was dropped, so the model +/// still knows an image was there and the conversation stays valid; +/// only the bytes go. Returns the number of images dropped. +pub fn evict_old_images( + messages: &mut [Message], + budget_bytes: usize, + budget_count: usize, +) -> usize { + // Newest first: recent images are the ones still being talked about. + let mut sites: Vec<(usize, usize, usize, String)> = Vec::new(); + for (msg_idx, msg) in messages.iter().enumerate() { + let content = match msg { + Message::User(u) => &u.content, + Message::Assistant(a) => &a.content, + Message::System(_) => continue, + }; + for (block_idx, block) in content.iter().enumerate() { + if let ContentBlock::Image { media_type, data } = block { + sites.push((msg_idx, block_idx, data.len(), media_type.clone())); + } + } + } + + let mut kept = 0usize; + let mut evicted = 0usize; + let mut kept_count = 0usize; + let mut budget_spent = false; + for (msg_idx, block_idx, len, media_type) in sites.into_iter().rev() { + if !budget_spent && kept + len <= budget_bytes && kept_count < budget_count { + kept += len; + kept_count += 1; + continue; + } + // The first image that does not fit ends retention outright. + budget_spent = true; + let placeholder = ContentBlock::Text { + text: format!( + "[image dropped from context — {media_type}, {:.1} MB]", + len as f64 / (1000.0 * 1000.0) + ), + }; + match &mut messages[msg_idx] { + Message::User(u) => u.content[block_idx] = placeholder, + Message::Assistant(a) => a.content[block_idx] = placeholder, + Message::System(_) => continue, + } + evicted += 1; + } + evicted +} + /// Perform microcompact: clear stale tool results to free tokens. /// /// Replaces the content of old tool_result blocks with a placeholder, @@ -687,6 +767,157 @@ mod tests { }) } + fn user_with_image(bytes: usize) -> Message { + crate::llm::message::user_message_with_attachments( + "look at this", + vec![ContentBlock::Image { + media_type: "image/png".into(), + data: "A".repeat(bytes), + }], + ) + } + + fn image_payload_bytes(messages: &[Message]) -> usize { + messages + .iter() + .filter_map(|m| match m { + Message::User(u) => Some(&u.content), + Message::Assistant(a) => Some(&a.content), + Message::System(_) => None, + }) + .flatten() + .filter_map(|b| match b { + ContentBlock::Image { data, .. } => Some(data.len()), + _ => None, + }) + .sum() + } + + /// Images are charged a flat vision estimate, so no token threshold + /// ever notices the megabytes they add to every later request. The + /// bytes have to be bounded on their own. + #[test] + fn old_image_payloads_are_evicted_beyond_the_budget() { + let budget = 1000; + let mut messages = vec![ + user_with_image(600), + assistant_text("first"), + user_with_image(600), + assistant_text("second"), + user_with_image(600), + ]; + let evicted = evict_old_images(&mut messages, budget, 100); + assert_eq!(evicted, 2, "older images were kept"); + assert!( + image_payload_bytes(&messages) <= budget, + "retained {} bytes over a {budget} budget", + image_payload_bytes(&messages) + ); + // The newest image survives: it is the one still being discussed. + assert!( + matches!( + messages[4].clone(), + Message::User(u) if u.content.iter().any(|b| matches!(b, ContentBlock::Image { .. })) + ), + "the newest image was dropped" + ); + } + + /// The block is replaced, not removed: the model still learns an image + /// was there, and the message keeps a valid shape. + #[test] + fn an_evicted_image_leaves_a_note_in_its_place() { + let mut messages = vec![user_with_image(500), user_with_image(500)]; + assert_eq!(evict_old_images(&mut messages, 600, 100), 1); + let Message::User(u) = &messages[0] else { + panic!("expected a user message"); + }; + let text = u + .content + .iter() + .find_map(|b| b.as_text()) + .expect("no text block replaced the image"); + assert!(text.contains("image dropped from context"), "{text}"); + assert!(text.contains("image/png"), "{text}"); + } + + #[test] + fn images_within_the_budget_are_untouched() { + let mut messages = vec![user_with_image(100), user_with_image(100)]; + let before = image_payload_bytes(&messages); + assert_eq!(evict_old_images(&mut messages, 1000, 100), 0); + assert_eq!(image_payload_bytes(&messages), before); + } + + /// Running twice must not keep rewriting history: once the payloads + /// are gone there is nothing left to evict. + #[test] + fn eviction_is_idempotent() { + let mut messages = vec![user_with_image(600), user_with_image(600)]; + assert_eq!(evict_old_images(&mut messages, 1000, 100), 1); + assert_eq!(evict_old_images(&mut messages, 1000, 100), 0); + } + + /// Retention is a contiguous run of the newest images. A greedy pass + /// would skip an image that does not fit and then keep an older, + /// smaller one behind it — preserving stale context while dropping + /// something newer. + #[test] + fn retention_stops_at_the_first_image_that_does_not_fit() { + // Oldest to newest: 4, 8, 8 with room for 12. + let mut messages = vec![ + user_with_image(4), + assistant_text("a"), + user_with_image(8), + assistant_text("b"), + user_with_image(8), + ]; + let evicted = evict_old_images(&mut messages, 12, 100); + assert_eq!(evicted, 2, "an older image was kept behind a dropped one"); + let has_image = |m: &Message| match m { + Message::User(u) => u + .content + .iter() + .any(|b| matches!(b, ContentBlock::Image { .. })), + _ => false, + }; + assert!(has_image(&messages[4]), "the newest image was dropped"); + assert!(!has_image(&messages[2]), "middle image should be gone"); + assert!( + !has_image(&messages[0]), + "the oldest image was kept after a newer one was dropped" + ); + } + + /// Providers cap how many images a request may carry, so a session of + /// small thumbnails hits that limit long before any byte budget. + #[test] + fn the_retained_image_count_is_bounded_too() { + let mut messages: Vec = (0..10).map(|_| user_with_image(10)).collect(); + let evicted = evict_old_images(&mut messages, 1_000_000, 4); + assert_eq!(evicted, 6, "count budget was not enforced"); + let remaining = messages + .iter() + .filter(|m| match m { + Message::User(u) => u + .content + .iter() + .any(|b| matches!(b, ContentBlock::Image { .. })), + _ => false, + }) + .count(); + assert_eq!(remaining, 4, "more images retained than the count budget"); + } + + /// A single image larger than the whole budget is still dropped — + /// keeping it would blow the bound it exists to enforce. + #[test] + fn an_image_larger_than_the_budget_is_evicted() { + let mut messages = vec![user_with_image(2000)]; + assert_eq!(evict_old_images(&mut messages, 1000, 100), 1); + assert_eq!(image_payload_bytes(&messages), 0); + } + fn assistant_tool_use(ids: &[&str]) -> Message { Message::Assistant(AssistantMessage { uuid: Uuid::new_v4(), diff --git a/docs/tui/KEYBINDINGS.md b/docs/tui/KEYBINDINGS.md index 32862397..e9c3fd41 100644 --- a/docs/tui/KEYBINDINGS.md +++ b/docs/tui/KEYBINDINGS.md @@ -191,3 +191,10 @@ double-press quit is still reachable. `/emacs` turns the vi bindings back off; the composer then behaves as it does everywhere else in this document. It does not add Emacs chords — the composer has none of its own, and `Ctrl+E` / `Ctrl+U` are the transcript controls listed above. + +## Images + +Mention an image the way you mention a file — `@screenshot.png` — and it is +attached to the turn as an image rather than inlined as text. Supported: +`png`, `jpg`/`jpeg`, `gif`, `webp`. Other binaries are still skipped with a +reason.