Compute UserBlockCompleted's expensive fields lazily

app/src/ai/aws_credentials.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 pub use ai::api_keys::AwsCredentials;
@@ -7,6 +8,7 @@ use aws_credential_types::provider::ProvideCredentials;
 use aws_credential_types::provider::error::CredentialsError;
 use futures::channel::oneshot::channel;
 use futures::future::BoxFuture;
+use parking_lot::FairMutex;
 use tokio::sync::Mutex;
 use vec1::vec1;
 use warp_errors::report_error;
@@ -15,7 +17,8 @@ use warp_managed_secrets::client::IdentityTokenOptions;
 use warpui::{ModelContext, ModelHandle, SingletonEntity};
 
 use crate::settings::{AISettings, AISettingsChangedEvent};
-use crate::terminal::event::{AfterBlockCompletedEvent, BlockType, UserBlockCompleted};
+use crate::terminal::event::{AfterBlockCompletedEvent, BlockType};
+use crate::terminal::model::terminal_model::TerminalModel;
 use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher};
 use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent};
 
@@ -176,6 +179,7 @@ pub trait AwsCredentialRefresher {
     fn register_model_event_dispatcher(
         &mut self,
         model_events: &ModelHandle<ModelEventDispatcher>,
+        terminal_model: Arc<FairMutex<TerminalModel>>,
         ctx: &mut ModelContext<Self>,
     ) where
         Self: Sized;
@@ -191,14 +195,25 @@ impl AwsCredentialRefresher for ApiKeyManager {
     fn register_model_event_dispatcher(
         &mut self,
         model_events: &ModelHandle<ModelEventDispatcher>,
+        terminal_model: Arc<FairMutex<TerminalModel>>,
         ctx: &mut ModelContext<Self>,
     ) {
-        ctx.subscribe_to_model(model_events, |manager, _, event, ctx| {
+        // we cannot simply capture the strong references, or we risk having a reference cycle.
+        let terminal_model_weak = Arc::downgrade(&terminal_model);
+        ctx.subscribe_to_model(model_events, move |manager, _, event, ctx| {
+            let Some(terminal_model) = terminal_model_weak.upgrade() else {
+                return;
+            };
+
             if let ModelEvent::AfterBlockCompleted(AfterBlockCompletedEvent {
-                block_type: BlockType::User(UserBlockCompleted { command, .. }),
+                block_type: BlockType::User(user_block_completed),
                 ..
             }) = event
             {
+                let command = user_block_completed.command.get_with(|compute| {
+                    let model = terminal_model.lock();
+                    compute(model.block_list())
+                });
                 let auth_command = &AISettings::as_ref(ctx).aws_bedrock_auth_refresh_command;
                 if command.trim().starts_with(auth_command.trim()) {
                     log::debug!("Detected AWS auth command completion, refreshing credentials");

app/src/ai/block_context.rs

@@ -1,9 +1,11 @@
 use channel_versions::overrides::TargetOS;
 use chrono::{DateTime, Local};
+use parking_lot::FairMutex;
 use serde::{Deserialize, Serialize};
 use warp_core::command::ExitCode;
 
 use crate::terminal::event::UserBlockCompleted;
+use crate::terminal::model::TerminalModel;
 use crate::terminal::model::block::BlockId;
 use crate::terminal::model::terminal_model::BlockIndex;
 
@@ -59,41 +61,54 @@ pub struct BlockContext {
 }
 
 impl BlockContext {
-    /// Construct a BlockContext from a [`UserBlockCompleted`].
-    pub fn from_completed_block(block_completed: &UserBlockCompleted) -> Box<Self> {
+    /// Construct a BlockContext from a [`UserBlockCompleted`]. `model` is used to resolve the
+    /// block's lazily-computed fields (see [`UserBlockCompleted`]'s accessor methods); it's only
+    /// locked for fields that aren't already cached.
+    pub fn from_completed_block(
+        block_completed: &UserBlockCompleted,
+        model: &FairMutex<TerminalModel>,
+    ) -> Box<Self> {
+        let serialized_block = block_completed.serialized_block.get_with(|compute| {
+            let model = model.lock();
+            compute(model.block_list())
+        });
         Box::new(Self {
-            id: block_completed.serialized_block.id.clone(),
+            id: serialized_block.id.clone(),
             index: block_completed.index,
-            command: block_completed.command_with_obfuscated_secrets.clone(),
+            command: block_completed
+                .command_with_obfuscated_secrets
+                .get_with(|compute| {
+                    let model = model.lock();
+                    compute(model.block_list())
+                })
+                .to_owned(),
             output: block_completed
                 .output_truncated_with_obfuscated_secrets
-                .clone(),
-            exit_code: block_completed.serialized_block.exit_code,
+                .get_with(|compute| {
+                    let model = model.lock();
+                    compute(model.block_list())
+                })
+                .to_owned(),
+            exit_code: serialized_block.exit_code,
             is_auto_attached: false,
-            started_ts: block_completed.serialized_block.start_ts,
-            finished_ts: block_completed.serialized_block.completed_ts,
-            pwd: block_completed.serialized_block.pwd.clone(),
-            shell: block_completed
-                .serialized_block
+            started_ts: serialized_block.start_ts,
+            finished_ts: serialized_block.completed_ts,
+            pwd: serialized_block.pwd.clone(),
+            shell: serialized_block
                 .shell_host
                 .as_ref()
                 .map(|sh| sh.shell_type.name().to_owned()),
-            username: block_completed
-                .serialized_block
+            username: serialized_block
                 .shell_host
                 .as_ref()
                 .map(|sh| sh.user.clone()),
-            hostname: block_completed
-                .serialized_block
+            hostname: serialized_block
                 .shell_host
                 .as_ref()
                 .map(|sh| sh.hostname.clone()),
-            git_branch: block_completed.serialized_block.git_head.clone(),
+            git_branch: serialized_block.git_head.clone(),
             os: TargetOS::current().and_then(|os| os.name()),
-            session_id: block_completed
-                .serialized_block
-                .session_id
-                .map(|sid| sid.as_u64()),
+            session_id: serialized_block.session_id.map(|sid| sid.as_u64()),
         })
     }
 }

app/src/ai/blocklist/passive_suggestions/legacy.rs

@@ -239,9 +239,9 @@ impl PassiveSuggestionsModel {
             return;
         }
 
-        if should_generate_unit_test_suggestion(block_completed, ctx) {
+        if should_generate_unit_test_suggestion(block_completed, &self.terminal_model, ctx) {
             self.generate_unit_test_suggestion(block_completed.clone(), ctx);
-        } else if should_generate_prompt_suggestions(block_completed, ctx) {
+        } else if should_generate_prompt_suggestions(block_completed, &self.terminal_model, ctx) {
             self.generate_prompt_suggestions(block_completed.clone(), ctx);
         }
     }
@@ -251,11 +251,26 @@ impl PassiveSuggestionsModel {
         block_completed: UserBlockCompleted,
         ctx: &mut ModelContext<Self>,
     ) {
-        let block_id = block_completed.serialized_block.id.clone();
-        let command = block_completed.command.clone();
+        let block_id = block_completed
+            .serialized_block
+            .get_with(|compute| {
+                let model = self.terminal_model.lock();
+                compute(model.block_list())
+            })
+            .id
+            .clone();
+        let command = block_completed
+            .command
+            .get_with(|compute| {
+                let model = self.terminal_model.lock();
+                compute(model.block_list())
+            })
+            .to_owned();
         let start_ts_ms = Utc::now().timestamp_millis();
 
-        if let Some(suggestion) = fetch_static_prompt_suggestion(&block_completed) {
+        if let Some(suggestion) =
+            fetch_static_prompt_suggestion(&block_completed, &self.terminal_model)
+        {
             ctx.emit(PassiveSuggestionsEvent::PromptSuggestionsGenerated {
                 prompt_suggestion: suggestion.clone(),
                 block_id: block_id.clone(),
@@ -322,12 +337,11 @@ impl PassiveSuggestionsModel {
 
         #[cfg(not(target_family = "wasm"))]
         {
-            let Some(current_dir) = block_completed
-                .serialized_block
-                .pwd
-                .as_ref()
-                .map(PathBuf::from)
-            else {
+            let serialized_block = block_completed.serialized_block.get_with(|compute| {
+                let model = self.terminal_model.lock();
+                compute(model.block_list())
+            });
+            let Some(current_dir) = serialized_block.pwd.as_ref().map(PathBuf::from) else {
                 return;
             };
 
@@ -573,9 +587,18 @@ impl Entity for PassiveSuggestionsModel {
 
 fn should_generate_prompt_suggestions(
     block_completed: &UserBlockCompleted,
+    terminal_model: &FairMutex<TerminalModel>,
     ctx: &ModelContext<PassiveSuggestionsModel>,
 ) -> bool {
-    if block_completed.command.trim().is_empty() {
+    let command_is_empty = block_completed
+        .command
+        .get_with(|compute| {
+            let model = terminal_model.lock();
+            compute(model.block_list())
+        })
+        .trim()
+        .is_empty();
+    if command_is_empty {
         return false;
     }
     if !NetworkStatus::as_ref(ctx).is_online() {
@@ -588,15 +611,28 @@ fn should_generate_prompt_suggestions(
 
 fn should_generate_unit_test_suggestion(
     block_completed: &UserBlockCompleted,
+    terminal_model: &FairMutex<TerminalModel>,
     ctx: &ModelContext<PassiveSuggestionsModel>,
 ) -> bool {
     let enabled = AISettings::as_ref(ctx).is_code_suggestions_enabled(ctx)
         && UserWorkspaces::as_ref(ctx).is_code_suggestions_toggleable();
 
+    let command = block_completed.command.get_with(|compute| {
+        let model = terminal_model.lock();
+        compute(model.block_list())
+    });
+
     enabled
-        && block_completed.command.starts_with("git")
-        && block_completed.command.contains("commit")
-        && block_completed.serialized_block.exit_code.was_successful()
+        && command.starts_with("git")
+        && command.contains("commit")
+        && block_completed
+            .serialized_block
+            .get_with(|compute| {
+                let model = terminal_model.lock();
+                compute(model.block_list())
+            })
+            .exit_code
+            .was_successful()
 }
 
 fn passive_code_diffs_enabled(ctx: &ModelContext<PassiveSuggestionsModel>) -> bool {
@@ -607,28 +643,47 @@ fn passive_code_diffs_enabled(ctx: &ModelContext<PassiveSuggestionsModel>) -> bo
     is_prompt_suggestions_enabled && is_code_suggestions_enabled && is_toggleable
 }
 
-fn fetch_static_prompt_suggestion(block: &UserBlockCompleted) -> Option<AgentModePromptSuggestion> {
-    if !block.serialized_block.exit_code.was_successful() {
+fn fetch_static_prompt_suggestion(
+    block: &UserBlockCompleted,
+    terminal_model: &FairMutex<TerminalModel>,
+) -> Option<AgentModePromptSuggestion> {
+    let was_successful = block
+        .serialized_block
+        .get_with(|compute| {
+            let model = terminal_model.lock();
+            compute(model.block_list())
+        })
+        .exit_code
+        .was_successful();
+    if !was_successful {
         return None;
     }
-    static_suggested_query(&block.command).map(AgentModePromptSuggestion::Success)
+    let command = block.command.get_with(|compute| {
+        let model = terminal_model.lock();
+        compute(model.block_list())
+    });
+    static_suggested_query(command).map(AgentModePromptSuggestion::Success)
 }
 
 fn build_prompt_suggestions_request(
     block: &UserBlockCompleted,
     execution_context: WarpAiExecutionContext,
-    terminal_model: &Arc<FairMutex<TerminalModel>>,
+    terminal_model: &FairMutex<TerminalModel>,
 ) -> Option<GenerateAMQuerySuggestionsRequest> {
-    let exit_code = block.serialized_block.exit_code;
-    let working_dir = block.serialized_block.pwd.as_ref();
+    let serialized_block = block.serialized_block.get_with(|compute| {
+        let model = terminal_model.lock();
+        compute(model.block_list())
+    });
+    let exit_code = serialized_block.exit_code;
+    let working_dir = serialized_block.pwd.as_ref();
     let (processed_input, processed_output) = {
         let model = terminal_model.lock();
-        let terminal_width = model.block_list().size().columns();
-        let Some(current_block) = model.block_list().block_with_id(&block.serialized_block.id)
-        else {
+        let block_list = model.block_list();
+        let terminal_width = block_list.size().columns();
+        let Some(current_block) = block_list.block_with_id(&serialized_block.id) else {
             report_error!(
                 "Failed to fetch prompt suggestions, could not find block with ID",
-                extra: { "block_id" => ?block.serialized_block.id }
+                extra: { "block_id" => ?serialized_block.id }
             );
             return None;
         };

app/src/ai/blocklist/passive_suggestions/maa.rs

@@ -495,21 +495,30 @@ impl PassiveSuggestionsModel {
             supported_tools.push(warp_multi_agent_api::ToolType::SuggestPrompt);
         }
 
-        let block_context = BlockContext::from_completed_block(block_completed);
-        let (conversation_id, block_context) = {
+        // Note: the lock is dropped before calling `BlockContext::from_completed_block` below,
+        // since that (like `UserBlockCompleted`'s other accessors) locks `self.terminal_model`
+        // itself if needed, and `FairMutex` isn't reentrant.
+        let conversation_id = {
             let model = self.terminal_model.lock();
             let Some(block) = model.block_list().block_at(block_completed.index) else {
                 return;
             };
-
-            let conversation_id = block.agent_view_visibility().agent_view_conversation_id();
-            (conversation_id, block_context)
+            block.agent_view_visibility().agent_view_conversation_id()
         };
+        let block_context =
+            BlockContext::from_completed_block(block_completed, &self.terminal_model);
 
         // If passive code diffs are enabled, check for any files that were read.
         #[cfg(feature = "local_fs")]
         if is_passive_code_diffs_enabled
-            && let Some(current_working_directory) = block_completed.serialized_block.pwd.clone()
+            && let Some(current_working_directory) = block_completed
+                .serialized_block
+                .get_with(|compute| {
+                    let model = self.terminal_model.lock();
+                    compute(model.block_list())
+                })
+                .pwd
+                .clone()
         {
             let block_contents = format!("{}\n{}", &block_context.command, &block_context.output);
             let shell = self.active_session.as_ref(ctx).shell_launch_data(ctx);