From 5800ffb8c10cf8e42475a99e07e9dbf9c07889b3 Mon Sep 17 00:00:00 2001 From: Wenzhe Xu Date: Tue, 28 Apr 2026 20:44:54 -0700 Subject: [PATCH 1/3] QUALITY-544: Lock NLD for image-attached agent-view entry When the user pastes or drags-and-drops an image into a terminal that has NLD on with autodetection enabled, entering agent view should lock the input to AI mode for the duration of the attach. Previously, two leak sites would unlock the input and let the classifier flip the buffer back to shell mode after the entry: 1. The `EnteredAgentView` subscriber in `BlocklistAIInputModel` would set `is_locked: !is_ai_autodetection_enabled` even for `ImageAdded` origins. 2. `set_input_mode_agent`'s `should_unlock` branch (which is intended for the `!`-shell-toggle clear path in fullscreen agent view with empty buffer) would re-unlock immediately after the image-add path entered agent view. This change introduces a single `has_locking_attachment()` predicate on `BlocklistAIContextModel` and uses it to gate both leak sites. It also adds an image-attachment-in-progress counter so the predicate is true during the async file-read + resize/encode pipeline, before the final `ImageContext` is appended to `pending_attachments`. Changes: - Add `pending_image_attachments_in_progress: usize` and `has_locking_attachment()` / `note_image_attachment_started()` / `note_image_attachment_completed()` to `BlocklistAIContextModel`. - Wire a handle to `BlocklistAIContextModel` into `BlocklistAIInputModel` via a `set_ai_context_model` setter, called from `TerminalView::new` after both models exist. - Add an `origin_requires_locked_ai` helper and force-lock to AI in the `EnteredAgentView` subscriber when the origin is `ImageAdded` or when `has_locking_attachment()` is true. - Defense in depth: gate `BlocklistAIInputModel::is_autodetection_enabled_for_current_context` on `!has_locking_attachment()`. - Increment / decrement the in-progress counter synchronously in `read_and_process_images_async` and `process_and_attach_images_as_ai_context` so callers see `has_locking_attachment() == true` for the entire attach pipeline. - Reorder `handle_pasted_or_dragdropped_image_filepaths` and `process_and_attach_clipboard_image` so the editor's image dispatch runs before `set_input_mode_agent`, ensuring the counter is already incremented when the lock decision is made. - Gate `set_input_mode_agent`'s `should_unlock` on `!has_locking_attachment()`. Co-Authored-By: Oz --- app/src/ai/blocklist/context_model.rs | 48 ++++++++++++++++++++++ app/src/ai/blocklist/input_model.rs | 58 +++++++++++++++++++++++++++ app/src/editor/view/mod.rs | 33 +++++++++++++++ app/src/terminal/input.rs | 46 ++++++++++++++------- app/src/terminal/view.rs | 7 ++++ 5 files changed, 177 insertions(+), 15 deletions(-) diff --git a/app/src/ai/blocklist/context_model.rs b/app/src/ai/blocklist/context_model.rs index 671b1d46..546242f7 100644 --- a/app/src/ai/blocklist/context_model.rs +++ b/app/src/ai/blocklist/context_model.rs @@ -149,6 +149,16 @@ pub struct BlocklistAIContextModel { /// instead of sending it immediately. /// Persists across exchanges in the same conversation (like fast-forward). queue_next_prompt_enabled: bool, + + /// Number of image-attachment flows currently in progress. + /// + /// Image attachment is asynchronous (file reads + resize/encode happen on a spawned task), so + /// at the moment a paste / drag-and-drop fires there are no entries in `pending_attachments` + /// yet. This counter is incremented synchronously when an image-attach pipeline starts and + /// decremented when it completes (or fails). It contributes to `has_locking_attachment` so the + /// input can be force-locked to AI mode for the entire duration of the attach, not only after + /// the actual `ImageContext` is appended. + pending_image_attachments_in_progress: usize, } pub fn block_context_from_terminal_model( @@ -315,6 +325,7 @@ impl BlocklistAIContextModel { pending_document_id: None, auto_attached_agent_view_user_block_ids: Vec::new(), queue_next_prompt_enabled: false, + pending_image_attachments_in_progress: 0, } } @@ -327,6 +338,43 @@ impl BlocklistAIContextModel { self.clear_diff_hunk_attachments(); self.set_pending_document(None, ctx); self.auto_attached_agent_view_user_block_ids.clear(); + self.pending_image_attachments_in_progress = 0; + } + + /// Returns `true` if the next AI query has any context that should force the input to be + /// locked in AI mode (skipping NLD): a pending image, a pending block, pending selected text, + /// or an in-progress image-attachment pipeline. + pub fn has_locking_attachment(&self) -> bool { + self.pending_image_attachments_in_progress > 0 + || !self.pending_context_block_ids.is_empty() + || self.pending_context_selected_text.is_some() + || self + .pending_attachments + .iter() + .any(|attachment| matches!(attachment, PendingAttachment::Image(_))) + } + + /// Marks the start of an asynchronous image-attachment pipeline. Must be paired with a call + /// to [`Self::note_image_attachment_completed`] when the pipeline finishes (or fails). + pub fn note_image_attachment_started(&mut self, ctx: &mut ModelContext) { + self.pending_image_attachments_in_progress += 1; + ctx.emit(BlocklistAIContextEvent::UpdatedPendingContext { + previous_block_ids: self.pending_context_block_ids.clone(), + requires_block_resync: false, + requires_text_resync: false, + }); + } + + /// Marks an asynchronous image-attachment pipeline as complete. Saturating-decrement so a + /// stray double-completion doesn't underflow the counter. + pub fn note_image_attachment_completed(&mut self, ctx: &mut ModelContext) { + self.pending_image_attachments_in_progress = + self.pending_image_attachments_in_progress.saturating_sub(1); + ctx.emit(BlocklistAIContextEvent::UpdatedPendingContext { + previous_block_ids: self.pending_context_block_ids.clone(), + requires_block_resync: false, + requires_text_resync: false, + }); } /// Returns the set `BlockId`s corresponding to blocks to be included as context with the next diff --git a/app/src/ai/blocklist/input_model.rs b/app/src/ai/blocklist/input_model.rs index 86182aa6..93ec8e1e 100644 --- a/app/src/ai/blocklist/input_model.rs +++ b/app/src/ai/blocklist/input_model.rs @@ -20,6 +20,7 @@ use warpui::{AppContext, Entity, EntityId, ModelContext, ModelHandle, SingletonE pub use input_classifier::InputType; use super::agent_view::{AgentViewController, AgentViewControllerEvent, AgentViewEntryOrigin}; +use super::context_model::BlocklistAIContextModel; use crate::terminal::cli_agent_sessions::{ CLIAgentInputState, CLIAgentSessionsModel, CLIAgentSessionsModelEvent, }; @@ -46,6 +47,18 @@ const HISTORY_ENTRY_MATCH_CUTOFF: f32 = 0.9; /// Duration to temporarily disable autodetection during operations like history selection. const AUTODETECTION_DISABLE_DURATION_MS: u64 = 250; +/// Returns `true` if entering agent view via this origin must force-lock the input to AI mode, +/// bypassing NLD entirely. +/// +/// Origins listed here are interaction patterns where the user's intent is unambiguously "talk to +/// the agent" and a misclassification would be visibly wrong (e.g. dropping the buffer back into +/// shell mode after pasting an image). +/// +/// Pattern C (RefineDiff) will be added here in a follow-up commit. +fn origin_requires_locked_ai(origin: &AgentViewEntryOrigin) -> bool { + matches!(origin, AgentViewEntryOrigin::ImageAdded) +} + /// Configuration for the terminal pane's input. #[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct InputConfig { @@ -148,6 +161,12 @@ pub struct BlocklistAIInputModel { agent_view_controller: ModelHandle, + /// Handle to the per-pane context model. Wired up via [`Self::set_ai_context_model`] + /// after both models are constructed in `TerminalView::new`. Used to read pending + /// attachments / blocks / selected text when deciding whether to force-lock the input + /// to AI mode (see [`BlocklistAIContextModel::has_locking_attachment`]). + ai_context_model: Option>, + terminal_view_id: EntityId, autodetect_abort_handle: Option, @@ -267,6 +286,19 @@ impl BlocklistAIInputModel { }, ctx, ); + } else if origin_requires_locked_ai(origin) || me.has_locking_attachment(ctx) { + // Interaction patterns that should fully bypass NLD on + // entry: image attachment in progress / attached, or block / selected + // text already in pending context. Force-lock to AI regardless of the + // user's NLD setting so the classifier never gets a chance to drop + // the buffer back to shell. + me.set_input_config_internal( + InputConfig { + input_type: InputType::AI, + is_locked: true, + }, + ctx, + ); } else { let is_autodetection_enabled = AISettings::as_ref(ctx).is_ai_autodetection_enabled(ctx); @@ -320,6 +352,7 @@ impl BlocklistAIInputModel { is_locked: !is_autodetection_enabled, }, agent_view_controller, + ai_context_model: None, terminal_view_id, last_ai_autodetection_ts: None, last_explicit_input_type_set_at: None, @@ -329,6 +362,22 @@ impl BlocklistAIInputModel { } } + /// Wire up the per-pane [`BlocklistAIContextModel`] handle. Must be called once after both + /// models are constructed (see `TerminalView::new`). The handle is `Option` because the + /// context model is constructed *after* the input model. + pub fn set_ai_context_model(&mut self, handle: ModelHandle) { + self.ai_context_model = Some(handle); + } + + /// Convenience wrapper around `BlocklistAIContextModel::has_locking_attachment` that no-ops + /// when the handle has not been wired up yet. + fn has_locking_attachment(&self, app: &AppContext) -> bool { + self.ai_context_model + .as_ref() + .map(|handle| handle.as_ref(app).has_locking_attachment()) + .unwrap_or(false) + } + /// Returns the InputType enum which specifies how we will handle the terminal input. pub fn input_type(&self) -> InputType { self.input_config.input_type @@ -488,6 +537,15 @@ impl BlocklistAIInputModel { return false; } + // Defense in depth: while there is a pending image (or other + // locking context), the classifier must never have a chance to flip the input + // back to shell mode, even per-keystroke. The `EnteredAgentView` subscriber + // and `set_input_mode_agent` already lock at entry; this guard protects the + // window if any future caller forgets. + if self.has_locking_attachment(app) { + return false; + } + let ai_settings = AISettings::as_ref(app); if FeatureFlag::AgentView.is_enabled() { if self.agent_view_controller.as_ref(app).is_fullscreen() { diff --git a/app/src/editor/view/mod.rs b/app/src/editor/view/mod.rs index ace141c4..ced83e9a 100644 --- a/app/src/editor/view/mod.rs +++ b/app/src/editor/view/mod.rs @@ -5089,6 +5089,16 @@ impl EditorView { return; } + // Mark image-attachment-in-progress synchronously so the input model can + // force-lock to AI mode for the duration of the async file read + processing pipeline. + // Paired with `note_image_attachment_completed` in the spawn callback below (always + // fires regardless of outcome). + if let Some(context_model) = &self.context_model { + context_model.update(ctx, |context_model, ctx| { + context_model.note_image_attachment_started(ctx); + }); + } + let window_id = ctx.window_id(); ctx.spawn( @@ -5174,6 +5184,15 @@ impl EditorView { if !images.is_empty() { this.process_and_attach_images_as_ai_context(num_images_user_attached, images, ctx); } + + // Pair with `note_image_attachment_started` above. Decrement regardless of + // outcome so the counter never leaks. If `process_and_attach_images_as_ai_context` + // ran above, it manages its own independent increment/decrement. + if let Some(context_model) = &this.context_model { + context_model.update(ctx, |context_model, ctx| { + context_model.note_image_attachment_completed(ctx); + }); + } }, ); } @@ -5214,6 +5233,16 @@ impl EditorView { ctx ); + // QUALITY-544: mark image-attachment-in-progress synchronously so callers (e.g. terminal + // input.rs) that invoke `set_input_mode_agent` immediately after this method see + // `has_locking_attachment() == true` and skip the autodetect-unlock branch. Paired with + // `note_image_attachment_completed` in the spawn callback below. + if let Some(context_model) = &self.context_model { + context_model.update(ctx, |context_model, ctx| { + context_model.note_image_attachment_started(ctx); + }); + } + self.process_attached_images_future_handle = Some(ctx.spawn( async move { let mut processed_pending_images = vec![]; @@ -5304,6 +5333,10 @@ impl EditorView { if let Some(context_model) = &this.context_model { context_model.update(ctx, |context_model, ctx| { context_model.append_pending_images(pending_images, ctx); + // Pair with `note_image_attachment_started` above. The pending images + // (if any survived processing) are now in `pending_attachments`, so + // `has_locking_attachment` continues to return true via that path. + context_model.note_image_attachment_completed(ctx); }); } diff --git a/app/src/terminal/input.rs b/app/src/terminal/input.rs index 5e5d1af8..30a2cd3f 100644 --- a/app/src/terminal/input.rs +++ b/app/src/terminal/input.rs @@ -9852,22 +9852,28 @@ impl Input { return 0; } - let is_buffer_empty = self.buffer_text(ctx).is_empty(); - let in_active_agent_view = self.agent_view_controller.as_ref(ctx).is_active(); - if is_buffer_empty || in_active_agent_view { - self.set_input_mode_agent(true, ctx); - self.update_image_context_options(ctx); - } - let paths_to_process: Vec = image_filepaths .into_iter() .take(num_images_to_attach) .collect(); - let num_paths = paths_to_process.len(); + + // Dispatch the read+attach pipeline FIRST so the context model's + // `pending_image_attachments_in_progress` counter is incremented synchronously before + // `set_input_mode_agent` runs. That way `has_locking_attachment()` is true at lock time + // and the should_unlock branch (which would otherwise re-enable autodetection in + // fullscreen agent view with NLD on) is skipped. self.editor.update(ctx, |editor, ctx| { editor.read_and_process_images_async(num_paths, paths_to_process, ctx); }); + + let is_buffer_empty = self.buffer_text(ctx).is_empty(); + let in_active_agent_view = self.agent_view_controller.as_ref(ctx).is_active(); + if is_buffer_empty || in_active_agent_view { + self.set_input_mode_agent(true, ctx); + self.update_image_context_options(ctx); + } + num_paths } @@ -9879,12 +9885,6 @@ impl Input { ) { self.maybe_enter_agent_view_for_image_add(ctx); - // Switch to AI mode with block-level lock, unless already AI-mode-locked - if !self.is_locked_in_ai_mode(ctx) { - self.set_input_mode_agent(true, ctx); - self.update_image_context_options(ctx); - } - let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() @@ -9911,9 +9911,19 @@ impl Input { file_name, }; + // QUALITY-544: dispatch image processing FIRST so the context model's + // `pending_image_attachments_in_progress` counter is incremented synchronously before + // `set_input_mode_agent` runs (which gates its should_unlock branch on + // `has_locking_attachment()`). self.editor.update(ctx, |editor, ctx| { editor.process_and_attach_images_as_ai_context(1, vec![attached_image], ctx); }); + + // Switch to AI mode with block-level lock, unless already AI-mode-locked + if !self.is_locked_in_ai_mode(ctx) { + self.set_input_mode_agent(true, ctx); + self.update_image_context_options(ctx); + } } /// Enters agent view when adding images, unless the CLI agent rich input is @@ -12846,10 +12856,16 @@ impl Input { // When AgentView is enabled, reverting to AI mode in an active agent view with an empty // buffer should unlock (re-enable autodetection) - semantically like clearing the "!". + // + // QUALITY-544: but if there is a pending image attachment (or other locking context such + // as block / selected text), do NOT unlock. The user's intent is unambiguously "talk to + // the agent"; letting the classifier flip the input back to shell mode would be a bug. + let has_locking_attachment = self.ai_context_model.as_ref(ctx).has_locking_attachment(); let should_unlock = FeatureFlag::AgentView.is_enabled() && self.agent_view_controller.as_ref(ctx).is_fullscreen() && is_input_buffer_empty - && AISettings::as_ref(ctx).is_ai_autodetection_enabled(ctx); + && AISettings::as_ref(ctx).is_ai_autodetection_enabled(ctx) + && !has_locking_attachment; if should_unlock { self.ai_input_model.update(ctx, |ai_input_model, ctx| { diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 559c407c..8f01898b 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -3311,6 +3311,13 @@ impl TerminalView { ctx, ) }); + // Now that the context model exists, wire its handle into the input model so the input + // model can consult `has_locking_attachment` for force-locking decisions (image attached, + // block context, selected text, etc.). The input model is constructed first because it + // owns lock state that the context model never reads. + ai_input_model.update(ctx, |ai_input_model, _| { + ai_input_model.set_ai_context_model(ai_context_model.clone()); + }); let ai_controller = ctx.add_model(|ctx| { BlocklistAIController::new( ai_input_model.clone(), From 1e857fa7a8f01a9ebfeb2ade76ac06db8074f6c3 Mon Sep 17 00:00:00 2001 From: Wenzhe Xu Date: Tue, 28 Apr 2026 22:07:35 -0700 Subject: [PATCH 2/3] Add unit tests for BlocklistAIContextModel locking attachment logic Cover `has_locking_attachment` (default false; true on pending block ids, selected text, or pending image; ignores file-only attachments; true while the in-progress image-attach counter is non-zero) and the `note_image_attachment_started` / `note_image_attachment_completed` counter mechanics including saturating decrement at zero and multiple concurrent in-flight pipelines. To keep the test fixture small, add `#[cfg(test)]` `new_for_test` constructors on `BlocklistAIContextModel` and `AmbientAgentViewModel` that skip subscriptions/singleton lookups (`BlocklistAIHistoryModel`, `LLMPreferences`, `CloudModel`, `UpdateManager`, `AppExecutionMode`). This follows the existing `new_for_test` convention used by `Sessions`, `TerminalModel`, `AuthManager`, `BlocklistAIHistoryModel`, etc. Tests live in `app/src/ai/blocklist/context_model_test.rs` per the repo's `${filename}_tests.rs` / `mod_test.rs` convention. Co-Authored-By: Oz --- app/src/ai/blocklist/context_model.rs | 55 ++++ app/src/ai/blocklist/context_model_test.rs | 256 +++++++++++++++++++ app/src/terminal/view/ambient_agent/model.rs | 29 +++ 3 files changed, 340 insertions(+) create mode 100644 app/src/ai/blocklist/context_model_test.rs diff --git a/app/src/ai/blocklist/context_model.rs b/app/src/ai/blocklist/context_model.rs index 546242f7..a8347d27 100644 --- a/app/src/ai/blocklist/context_model.rs +++ b/app/src/ai/blocklist/context_model.rs @@ -329,6 +329,35 @@ impl BlocklistAIContextModel { } } + /// Test-only constructor that skips every subscription and singleton lookup performed by + /// [`Self::new`], so unit tests can build a [`BlocklistAIContextModel`] without registering + /// `BlocklistAIHistoryModel`, `LLMPreferences`, `ModelEventDispatcher`, `Sessions`, or + /// `AppExecutionMode`. Callers still pass real [`TerminalModel`] and [`AgentViewController`] + /// handles to populate the struct fields, but neither needs to be functional for the + /// methods exercised by these tests. + #[cfg(test)] + pub(crate) fn new_for_test( + terminal_model: Arc>, + terminal_view_id: EntityId, + agent_view_controller: ModelHandle, + ) -> Self { + Self { + terminal_model, + directory_context: Default::default(), + pending_context_block_ids: HashSet::new(), + pending_context_selected_text: None, + pending_attachments: Default::default(), + pending_query_state: PendingQueryState::default(), + terminal_view_id, + agent_view_controller, + pending_inline_diff_hunk_attachments: Default::default(), + pending_document_id: None, + auto_attached_agent_view_user_block_ids: Vec::new(), + queue_next_prompt_enabled: false, + pending_image_attachments_in_progress: 0, + } + } + /// Resets the set of blocks to be included as context to an empty list. /// Also removes any selected text that was to be included as context. pub fn reset_context_to_default(&mut self, ctx: &mut ModelContext) { @@ -1039,3 +1068,29 @@ pub enum BlocklistAIContextEvent { impl Entity for BlocklistAIContextModel { type Event = BlocklistAIContextEvent; } + +#[cfg(test)] +#[path = "context_model_test.rs"] +mod tests; + +#[cfg(test)] +impl BlocklistAIContextModel { + pub(crate) fn pending_image_attachments_in_progress_for_test(&self) -> usize { + self.pending_image_attachments_in_progress + } + + pub(crate) fn append_pending_attachments_for_test( + &mut self, + attachments: Vec, + ) { + self.pending_attachments.extend(attachments); + } + + pub(crate) fn insert_pending_block_id_for_test(&mut self, block_id: BlockId) { + self.pending_context_block_ids.insert(block_id); + } + + pub(crate) fn set_pending_selected_text_for_test(&mut self, text: Option) { + self.pending_context_selected_text = text; + } +} diff --git a/app/src/ai/blocklist/context_model_test.rs b/app/src/ai/blocklist/context_model_test.rs new file mode 100644 index 00000000..cd08b6ad --- /dev/null +++ b/app/src/ai/blocklist/context_model_test.rs @@ -0,0 +1,256 @@ +//! Unit tests for [`BlocklistAIContextModel::has_locking_attachment`] and the +//! `note_image_attachment_started` / `note_image_attachment_completed` counter mechanics. +//! +//! These tests deliberately bypass the production [`BlocklistAIContextModel::new`] constructor +//! (which subscribes to several singletons) and instead use [`BlocklistAIContextModel::new_for_test`] +//! together with [`super::agent_view::AgentViewController::new`] backed by +//! [`crate::terminal::view::ambient_agent::AmbientAgentViewModel::new_for_test`]. That keeps the +//! fixture small enough to focus on the lock/counter logic without standing up `BlocklistAIHistoryModel`, +//! `LLMPreferences`, `CloudModel`, `UpdateManager`, or `AppExecutionMode`. + +use std::sync::Arc; + +use parking_lot::FairMutex; +use warpui::r#async::executor::Background; +use warpui::{App, EntityId, ModelHandle}; + +use super::{BlocklistAIContextModel, PendingAttachment, PendingFile}; +use crate::ai::agent::ImageContext; +use crate::ai::blocklist::agent_view::{AgentViewController, EphemeralMessageModel}; +use crate::terminal::color::{self, Colors}; +use crate::terminal::event_listener::ChannelEventListener; +use crate::terminal::model::test_utils::block_size; +use crate::terminal::model::{BlockId, TerminalModel}; +use crate::terminal::view::ambient_agent::AmbientAgentViewModel; + +/// Builds a [`BlocklistAIContextModel`] with stub dependencies. None of the dependencies are +/// exercised by the methods under test; they only need to satisfy the struct's field types. +fn build_test_context_model(app: &mut App) -> ModelHandle { + let terminal_model = Arc::new(FairMutex::new(TerminalModel::new_for_test( + block_size(), + color::List::from(&Colors::default()), + ChannelEventListener::new_for_test(), + Arc::new(Background::default()), + false, /* should_show_bootstrap_block */ + None, /* restored_blocks */ + false, /* honor_ps1 */ + false, /* is_inverted */ + None, /* session_startup_path */ + ))); + let terminal_view_id = EntityId::new(); + + let ambient_agent_view_model = app.add_model(|ctx| { + AmbientAgentViewModel::new_for_test( + terminal_view_id, + false, /* has_parent_terminal */ + ctx, + ) + }); + let ephemeral_message_model = app.add_model(|_| EphemeralMessageModel::new()); + let agent_view_controller = app.add_model(|ctx| { + AgentViewController::new( + terminal_model.clone(), + terminal_view_id, + ambient_agent_view_model, + ephemeral_message_model, + ctx, + ) + }); + + app.add_model(|_| { + BlocklistAIContextModel::new_for_test( + terminal_model, + terminal_view_id, + agent_view_controller, + ) + }) +} + +fn make_image_attachment(file_name: &str) -> PendingAttachment { + PendingAttachment::Image(ImageContext { + data: String::new(), + mime_type: "image/png".to_owned(), + file_name: file_name.to_owned(), + is_figma: false, + }) +} + +fn make_file_attachment(file_name: &str) -> PendingAttachment { + PendingAttachment::File(PendingFile { + file_name: file_name.to_owned(), + file_path: file_name.into(), + mime_type: "text/plain".to_owned(), + }) +} + +#[test] +fn has_locking_attachment_is_false_for_default_state() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.read(&app, |m, _| { + assert!(!m.has_locking_attachment()); + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 0); + }); + }); +} + +#[test] +fn has_locking_attachment_is_true_with_pending_block_id() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, _| { + m.insert_pending_block_id_for_test(BlockId::new()); + }); + + model.read(&app, |m, _| assert!(m.has_locking_attachment())); + }); +} + +#[test] +fn has_locking_attachment_is_true_with_pending_selected_text() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, _| { + m.set_pending_selected_text_for_test(Some("hello".to_owned())); + }); + + model.read(&app, |m, _| assert!(m.has_locking_attachment())); + }); +} + +#[test] +fn has_locking_attachment_is_true_with_pending_image_attachment() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, _| { + m.append_pending_attachments_for_test(vec![make_image_attachment("a.png")]); + }); + + model.read(&app, |m, _| assert!(m.has_locking_attachment())); + }); +} + +#[test] +fn has_locking_attachment_is_false_with_only_file_attachments() { + // Files are explicitly *not* locking attachments — only images, blocks, selected text, or an + // in-progress image-attach pipeline force the input into AI mode. + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, _| { + m.append_pending_attachments_for_test(vec![ + make_file_attachment("notes.txt"), + make_file_attachment("readme.md"), + ]); + }); + + model.read(&app, |m, _| assert!(!m.has_locking_attachment())); + }); +} + +#[test] +fn has_locking_attachment_ignores_non_image_attachments_when_image_present() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, _| { + m.append_pending_attachments_for_test(vec![ + make_file_attachment("notes.txt"), + make_image_attachment("a.png"), + ]); + }); + + model.read(&app, |m, _| assert!(m.has_locking_attachment())); + }); +} + +#[test] +fn note_image_attachment_started_increments_counter_and_locks_input() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, ctx| m.note_image_attachment_started(ctx)); + + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 1); + // The counter alone — without any actual `pending_attachments` entry — must lock the + // input so paste / drag-and-drop flows can't slip into NLD before the async pipeline + // appends the resulting `ImageContext`. + assert!(m.has_locking_attachment()); + }); + }); +} + +#[test] +fn note_image_attachment_completed_decrements_counter_and_unlocks_input() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, ctx| { + m.note_image_attachment_started(ctx); + m.note_image_attachment_completed(ctx); + }); + + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 0); + assert!(!m.has_locking_attachment()); + }); + }); +} + +#[test] +fn note_image_attachment_started_supports_multiple_concurrent_pipelines() { + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, ctx| { + m.note_image_attachment_started(ctx); + m.note_image_attachment_started(ctx); + m.note_image_attachment_started(ctx); + }); + + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 3); + assert!(m.has_locking_attachment()); + }); + + // One completion should not release the lock while two pipelines are still in flight. + model.update(&mut app, |m, ctx| m.note_image_attachment_completed(ctx)); + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 2); + assert!(m.has_locking_attachment()); + }); + + model.update(&mut app, |m, ctx| { + m.note_image_attachment_completed(ctx); + m.note_image_attachment_completed(ctx); + }); + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 0); + assert!(!m.has_locking_attachment()); + }); + }); +} + +#[test] +fn note_image_attachment_completed_saturates_at_zero() { + // Defensive: a stray `_completed` without a matching `_started` (or a double-completion) must + // not underflow `usize` and silently lock the input forever. + App::test((), |mut app| async move { + let model = build_test_context_model(&mut app); + + model.update(&mut app, |m, ctx| { + m.note_image_attachment_completed(ctx); + m.note_image_attachment_completed(ctx); + }); + + model.read(&app, |m, _| { + assert_eq!(m.pending_image_attachments_in_progress_for_test(), 0); + assert!(!m.has_locking_attachment()); + }); + }); +} diff --git a/app/src/terminal/view/ambient_agent/model.rs b/app/src/terminal/view/ambient_agent/model.rs index 4e6d2020..8fe15807 100644 --- a/app/src/terminal/view/ambient_agent/model.rs +++ b/app/src/terminal/view/ambient_agent/model.rs @@ -156,6 +156,35 @@ impl AmbientAgentViewModel { } } + /// Test-only constructor that skips the [`CloudModel`] subscription and the + /// [`UpdateManager`]-driven environment validation. Lets sibling tests stand up an + /// `AmbientAgentViewModel` handle (only used to satisfy `AgentViewController`'s field + /// type) without registering the `CloudModel`/`UpdateManager` singletons. + #[cfg(test)] + pub fn new_for_test( + terminal_view_id: EntityId, + has_parent_terminal: bool, + ctx: &mut ModelContext, + ) -> Self { + let ui_state = AmbientAgentProgressUIState::new(ctx); + + Self { + status: Status::NotAmbientAgent, + request: None, + terminal_view_id, + has_parent_terminal, + environment_id: None, + progress_timer_handle: None, + ui_state, + setup_commands_state: Default::default(), + task_id: None, + conversation_id: None, + harness: Harness::default(), + has_inserted_cloud_mode_user_query_block: false, + harness_command_started: false, + } + } + pub fn request(&self) -> Option<&SpawnAgentRequest> { self.request.as_ref() } From 00cfdeb8d54bcc8c055f5f506b9ff98fb2caf185 Mon Sep 17 00:00:00 2001 From: Wenzhe Xu Date: Tue, 28 Apr 2026 22:09:17 -0700 Subject: [PATCH 3/3] Fix counter leak when image processing future is aborted Switch from ctx.spawn to ctx.spawn_abortable in process_and_attach_images_as_ai_context so pending_image_attachments_in_progress is decremented on every exit path: - on_resolve (normal completion): unchanged - on_resolve early-return when handle was taken between value production and dispatch: now decrements before returning instead of leaking - on_abort (future aborted before producing a value, e.g. submit_ai_query calling abort_attached_images_future_handle while processing was in flight): new explicit handler that decrements Previously ctx.spawn's default no-op on_abort callback meant that aborting the spawned future left pending_image_attachments_in_progress nonzero, keeping has_locking_attachment() true and force-locking the input into AI mode until an unrelated context reset. Co-Authored-By: Oz --- app/src/editor/view/mod.rs | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/app/src/editor/view/mod.rs b/app/src/editor/view/mod.rs index ced83e9a..d459db43 100644 --- a/app/src/editor/view/mod.rs +++ b/app/src/editor/view/mod.rs @@ -5236,14 +5236,15 @@ impl EditorView { // QUALITY-544: mark image-attachment-in-progress synchronously so callers (e.g. terminal // input.rs) that invoke `set_input_mode_agent` immediately after this method see // `has_locking_attachment() == true` and skip the autodetect-unlock branch. Paired with - // `note_image_attachment_completed` in the spawn callback below. + // `note_image_attachment_completed` in both the resolve and abort callbacks below so the + // counter is balanced regardless of how the spawned future exits. if let Some(context_model) = &self.context_model { context_model.update(ctx, |context_model, ctx| { context_model.note_image_attachment_started(ctx); }); } - self.process_attached_images_future_handle = Some(ctx.spawn( + self.process_attached_images_future_handle = Some(ctx.spawn_abortable( async move { let mut processed_pending_images = vec![]; let mut num_oversized_images: usize = 0; @@ -5283,8 +5284,16 @@ impl EditorView { ) }, move |this, (num_oversized_images, num_unprocessed_images, pending_images), ctx| { - // Future was aborted + // If the handle was taken between the future producing a value and this callback + // running (e.g. `abort_attached_images_future_handle` fired from `submit_ai_query` + // after the result was already in the channel), discard the result but still + // decrement the counter to balance `note_image_attachment_started` above. if this.process_attached_images_future_handle.is_none() { + if let Some(context_model) = &this.context_model { + context_model.update(ctx, |context_model, ctx| { + context_model.note_image_attachment_completed(ctx); + }); + } return; } @@ -5342,6 +5351,18 @@ impl EditorView { ctx.emit(Event::ProcessingAttachedImages(false)); }, + |this, ctx| { + // Future was aborted before producing a value (e.g. `submit_ai_query` called + // `abort_attached_images_future_handle` while image processing was still in + // flight). `ctx.spawn`'s default abort callback is a no-op, so without an explicit + // decrement here `pending_image_attachments_in_progress` would leak and keep the + // input force-locked until an unrelated context reset. + if let Some(context_model) = &this.context_model { + context_model.update(ctx, |context_model, ctx| { + context_model.note_image_attachment_completed(ctx); + }); + } + }, )); ctx.emit(Event::ProcessingAttachedImages(true));