agent_ui: Add safeguards for deep link prompt injections (#50936)

Smit Barmase created

- Show a warning asking users to review pre-filled prompts from external
links before sending.
- Strip control characters and bidi control characters from pre-filled
prompts.
- Collapse 3+ consecutive newlines in pre-filled prompts to prevent
newline padding attacks.
- API changes make it impossible to auto-submit pre-filled prompts from
external sources.

Release Notes:

- N/A

Change summary

crates/agent_ui/src/agent_panel.rs                 |  12 
crates/agent_ui/src/agent_ui.rs                    |   9 
crates/agent_ui/src/connection_view.rs             |  74 +++++++
crates/agent_ui/src/connection_view/thread_view.rs |  47 ++++
crates/agent_ui/src/external_source_prompt.rs      | 162 ++++++++++++++++
crates/zed/src/main.rs                             |  10 
crates/zed/src/zed/open_listener.rs                |  45 +++-
7 files changed, 335 insertions(+), 24 deletions(-)

Detailed changes

crates/agent_ui/src/agent_panel.rs ๐Ÿ”—

@@ -39,7 +39,8 @@ use crate::{
     ui::EndTrialUpsell,
 };
 use crate::{
-    AgentInitialContent, ExternalAgent, NewExternalAgentThread, NewNativeAgentThreadFromSummary,
+    AgentInitialContent, ExternalAgent, ExternalSourcePrompt, NewExternalAgentThread,
+    NewNativeAgentThreadFromSummary,
 };
 use crate::{
     ExpandMessageEditor, ThreadHistory, ThreadHistoryEvent,
@@ -1994,9 +1995,9 @@ impl AgentPanel {
         }
     }
 
-    pub fn new_external_thread_with_text(
+    pub fn new_agent_thread_with_external_source_prompt(
         &mut self,
-        initial_text: Option<String>,
+        external_source_prompt: Option<ExternalSourcePrompt>,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
@@ -2005,10 +2006,7 @@ impl AgentPanel {
             None,
             None,
             None,
-            initial_text.map(|text| AgentInitialContent::ContentBlock {
-                blocks: vec![acp::ContentBlock::Text(acp::TextContent::new(text))],
-                auto_submit: false,
-            }),
+            external_source_prompt.map(AgentInitialContent::from),
             window,
             cx,
         );

crates/agent_ui/src/agent_ui.rs ๐Ÿ”—

@@ -11,6 +11,7 @@ pub(crate) mod connection_view;
 mod context;
 mod context_server_configuration;
 mod entry_view_state;
+mod external_source_prompt;
 mod favorite_models;
 mod inline_assistant;
 mod inline_prompt_editor;
@@ -66,6 +67,7 @@ use crate::agent_registry_ui::AgentRegistryPage;
 pub use crate::inline_assistant::InlineAssistant;
 pub use agent_diff::{AgentDiffPane, AgentDiffToolbar};
 pub(crate) use connection_view::ConnectionView;
+pub use external_source_prompt::ExternalSourcePrompt;
 pub(crate) use mode_selector::ModeSelector;
 pub(crate) use model_selector::ModelSelector;
 pub(crate) use model_selector_popover::ModelSelectorPopover;
@@ -250,6 +252,13 @@ pub enum AgentInitialContent {
         blocks: Vec<agent_client_protocol::ContentBlock>,
         auto_submit: bool,
     },
+    FromExternalSource(ExternalSourcePrompt),
+}
+
+impl From<ExternalSourcePrompt> for AgentInitialContent {
+    fn from(prompt: ExternalSourcePrompt) -> Self {
+        Self::FromExternalSource(prompt)
+    }
 }
 
 /// Opens the profile management interface for configuring agent tools and settings.

crates/agent_ui/src/connection_view.rs ๐Ÿ”—

@@ -2796,6 +2796,55 @@ pub(crate) mod tests {
         assert!(!weak_view.is_upgradable());
     }
 
+    #[gpui::test]
+    async fn test_external_source_prompt_requires_manual_send(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let Some(prompt) = crate::ExternalSourcePrompt::new("Write me a script") else {
+            panic!("expected prompt from external source to sanitize successfully");
+        };
+        let initial_content = AgentInitialContent::FromExternalSource(prompt);
+
+        let (thread_view, cx) = setup_thread_view_with_initial_content(
+            StubAgentServer::default_response(),
+            initial_content,
+            cx,
+        )
+        .await;
+
+        active_thread(&thread_view, cx).read_with(cx, |view, cx| {
+            assert!(view.show_external_source_prompt_warning);
+            assert_eq!(view.thread.read(cx).entries().len(), 0);
+            assert_eq!(view.message_editor.read(cx).text(cx), "Write me a script");
+        });
+    }
+
+    #[gpui::test]
+    async fn test_external_source_prompt_warning_clears_after_send(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let Some(prompt) = crate::ExternalSourcePrompt::new("Write me a script") else {
+            panic!("expected prompt from external source to sanitize successfully");
+        };
+        let initial_content = AgentInitialContent::FromExternalSource(prompt);
+
+        let (thread_view, cx) = setup_thread_view_with_initial_content(
+            StubAgentServer::default_response(),
+            initial_content,
+            cx,
+        )
+        .await;
+
+        active_thread(&thread_view, cx).update_in(cx, |view, window, cx| view.send(window, cx));
+        cx.run_until_parked();
+
+        active_thread(&thread_view, cx).read_with(cx, |view, cx| {
+            assert!(!view.show_external_source_prompt_warning);
+            assert_eq!(view.message_editor.read(cx).text(cx), "");
+            assert_eq!(view.thread.read(cx).entries().len(), 2);
+        });
+    }
+
     #[gpui::test]
     async fn test_notification_for_stop_event(cx: &mut TestAppContext) {
         init_test(cx);
@@ -3610,6 +3659,29 @@ pub(crate) mod tests {
         Entity<ConnectionView>,
         Entity<ThreadHistory>,
         &mut VisualTestContext,
+    ) {
+        setup_thread_view_with_history_and_initial_content(agent, None, cx).await
+    }
+
+    async fn setup_thread_view_with_initial_content(
+        agent: impl AgentServer + 'static,
+        initial_content: AgentInitialContent,
+        cx: &mut TestAppContext,
+    ) -> (Entity<ConnectionView>, &mut VisualTestContext) {
+        let (thread_view, _history, cx) =
+            setup_thread_view_with_history_and_initial_content(agent, Some(initial_content), cx)
+                .await;
+        (thread_view, cx)
+    }
+
+    async fn setup_thread_view_with_history_and_initial_content(
+        agent: impl AgentServer + 'static,
+        initial_content: Option<AgentInitialContent>,
+        cx: &mut TestAppContext,
+    ) -> (
+        Entity<ConnectionView>,
+        Entity<ThreadHistory>,
+        &mut VisualTestContext,
     ) {
         let fs = FakeFs::new(cx.executor());
         let project = Project::test(fs, [], cx).await;
@@ -3627,7 +3699,7 @@ pub(crate) mod tests {
                     None,
                     None,
                     None,
-                    None,
+                    initial_content,
                     workspace.downgrade(),
                     project,
                     Some(thread_store),

crates/agent_ui/src/connection_view/thread_view.rs ๐Ÿ”—

@@ -262,6 +262,7 @@ pub struct ThreadView {
     pub project: WeakEntity<Project>,
     pub recent_history_entries: Vec<AgentSessionInfo>,
     pub hovered_recent_history_item: Option<usize>,
+    pub show_external_source_prompt_warning: bool,
     pub show_codex_windows_warning: bool,
     pub history: Entity<ThreadHistory>,
     pub _history_subscription: Subscription,
@@ -324,6 +325,7 @@ impl ThreadView {
         });
 
         let mut should_auto_submit = false;
+        let mut show_external_source_prompt_warning = false;
 
         let message_editor = cx.new(|cx| {
             let mut editor = MessageEditor::new(
@@ -355,6 +357,18 @@ impl ThreadView {
                         should_auto_submit = auto_submit;
                         editor.set_message(blocks, window, cx);
                     }
+                    AgentInitialContent::FromExternalSource(prompt) => {
+                        show_external_source_prompt_warning = true;
+                        // SECURITY: Be explicit about not auto submitting prompt from external source.
+                        should_auto_submit = false;
+                        editor.set_message(
+                            vec![acp::ContentBlock::Text(acp::TextContent::new(
+                                prompt.into_string(),
+                            ))],
+                            window,
+                            cx,
+                        );
+                    }
                 }
             } else if let Some(draft) = thread.read(cx).draft_prompt() {
                 editor.set_message(draft.to_vec(), window, cx);
@@ -477,6 +491,7 @@ impl ThreadView {
             project,
             recent_history_entries,
             hovered_recent_history_item: None,
+            show_external_source_prompt_warning,
             history,
             _history_subscription: history_subscription,
             show_codex_windows_warning,
@@ -781,6 +796,13 @@ impl ThreadView {
 
     // sending
 
+    fn clear_external_source_prompt_warning(&mut self, cx: &mut Context<Self>) {
+        if self.show_external_source_prompt_warning {
+            self.show_external_source_prompt_warning = false;
+            cx.notify();
+        }
+    }
+
     pub fn send(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         let thread = &self.thread;
 
@@ -862,6 +884,7 @@ impl ThreadView {
                     .any(|command| command.name == "logout");
             if can_login && !logout_supported {
                 message_editor.update(cx, |editor, cx| editor.clear(window, cx));
+                self.clear_external_source_prompt_warning(cx);
 
                 let connection = self.thread.read(cx).connection().clone();
                 window.defer(cx, {
@@ -954,6 +977,7 @@ impl ThreadView {
             };
 
             let generation = this.update(cx, |this, cx| {
+                this.clear_external_source_prompt_warning(cx);
                 let generation = this.start_turn(cx);
                 this.in_flight_prompt = Some(contents.clone());
                 generation
@@ -7445,6 +7469,26 @@ impl ThreadView {
             )
     }
 
+    fn render_external_source_prompt_warning(&self, cx: &mut Context<Self>) -> Callout {
+        Callout::new()
+            .icon(IconName::Warning)
+            .severity(Severity::Warning)
+            .title("Review before sending")
+            .description("This prompt was pre-filled by an external link. Read it carefully before you send it.")
+            .dismiss_action(
+                IconButton::new("dismiss-external-source-prompt-warning", IconName::Close)
+                    .icon_size(IconSize::Small)
+                    .icon_color(Color::Muted)
+                    .tooltip(Tooltip::text("Dismiss Warning"))
+                    .on_click(cx.listener({
+                        move |this, _, _, cx| {
+                            this.show_external_source_prompt_warning = false;
+                            cx.notify();
+                        }
+                    })),
+            )
+    }
+
     fn render_new_version_callout(&self, version: &SharedString, cx: &mut Context<Self>) -> Div {
         let server_view = self.server_view.clone();
         v_flex().w_full().justify_end().child(
@@ -7794,6 +7838,9 @@ impl Render for ThreadView {
             .children(self.render_subagent_titlebar(cx))
             .child(conversation)
             .children(self.render_activity_bar(window, cx))
+            .when(self.show_external_source_prompt_warning, |this| {
+                this.child(self.render_external_source_prompt_warning(cx))
+            })
             .when(self.show_codex_windows_warning, |this| {
                 this.child(self.render_codex_windows_warning(cx))
             })

crates/agent_ui/src/external_source_prompt.rs ๐Ÿ”—

@@ -0,0 +1,162 @@
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ExternalSourcePrompt(String);
+
+impl ExternalSourcePrompt {
+    pub fn new(prompt: &str) -> Option<Self> {
+        sanitize(prompt).map(Self)
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+
+    pub fn into_string(self) -> String {
+        self.0
+    }
+}
+
+fn sanitize(prompt: &str) -> Option<String> {
+    let mut sanitized_prompt = String::with_capacity(prompt.len());
+    let mut consecutive_newline_count = 0;
+    let mut characters = prompt.chars().peekable();
+
+    while let Some(character) = characters.next() {
+        let character = if character == '\r' {
+            if characters.peek() == Some(&'\n') {
+                characters.next();
+            }
+            '\n'
+        } else {
+            character
+        };
+
+        if is_bidi_control_character(character) || is_disallowed_control_character(character) {
+            continue;
+        }
+
+        if character == '\n' {
+            consecutive_newline_count += 1;
+            if consecutive_newline_count > 2 {
+                continue;
+            }
+        } else {
+            consecutive_newline_count = 0;
+        }
+
+        sanitized_prompt.push(character);
+    }
+
+    if sanitized_prompt.is_empty() {
+        None
+    } else {
+        Some(sanitized_prompt)
+    }
+}
+
+fn is_disallowed_control_character(character: char) -> bool {
+    character.is_control() && !matches!(character, '\n' | '\t')
+}
+
+fn is_bidi_control_character(character: char) -> bool {
+    matches!(
+        character,
+          '\u{061C}' // ALM
+        | '\u{200E}' // LRM
+        | '\u{200F}' // RLM
+        | '\u{202A}'..='\u{202E}' // LRE, RLE, PDF, LRO, RLO
+        | '\u{2066}'..='\u{2069}' // LRI, RLI, FSI, PDI
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ExternalSourcePrompt;
+
+    #[test]
+    fn keeps_normal_prompt_text() {
+        let prompt = ExternalSourcePrompt::new("Write me a script\nThanks");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("Write me a script\nThanks")
+        );
+    }
+
+    #[test]
+    fn keeps_multilingual_text() {
+        let prompt =
+            ExternalSourcePrompt::new("ๆ—ฅๆœฌ่ชžใฎไพ้ ผใงใ™ใ€‚\nไธญๆ–‡ๆ็คบไนŸๅบ”่ฏฅไฟ็•™ใ€‚\nemoji ๐Ÿ‘ฉโ€๐Ÿ’ป");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("ๆ—ฅๆœฌ่ชžใฎไพ้ ผใงใ™ใ€‚\nไธญๆ–‡ๆ็คบไนŸๅบ”่ฏฅไฟ็•™ใ€‚\nemoji ๐Ÿ‘ฉโ€๐Ÿ’ป")
+        );
+    }
+
+    #[test]
+    fn collapses_newline_padding() {
+        let prompt = ExternalSourcePrompt::new(
+            "Review this prompt carefully.\n\nThis paragraph should stay separated.\n\n\n\n\n\n\nWrite me a script to do fizz buzz.",
+        );
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some(
+                "Review this prompt carefully.\n\nThis paragraph should stay separated.\n\nWrite me a script to do fizz buzz."
+            )
+        );
+    }
+
+    #[test]
+    fn normalizes_carriage_returns() {
+        let prompt = ExternalSourcePrompt::new("Line one\r\nLine two\rLine three");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("Line one\nLine two\nLine three")
+        );
+    }
+
+    #[test]
+    fn strips_bidi_control_characters() {
+        let prompt = ExternalSourcePrompt::new("abc\u{202E}def\u{202C}ghi");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("abcdefghi")
+        );
+    }
+
+    #[test]
+    fn strips_other_control_characters() {
+        let prompt = ExternalSourcePrompt::new("safe\u{0000}\u{001B}\u{007F}text");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("safetext")
+        );
+    }
+
+    #[test]
+    fn keeps_tabs() {
+        let prompt = ExternalSourcePrompt::new("keep\tindentation");
+
+        assert_eq!(
+            prompt.as_ref().map(ExternalSourcePrompt::as_str),
+            Some("keep\tindentation")
+        );
+    }
+
+    #[test]
+    fn drops_empty_prompt() {
+        assert_eq!(ExternalSourcePrompt::new(""), None);
+    }
+
+    #[test]
+    fn drops_prompt_with_only_removed_characters() {
+        assert_eq!(
+            ExternalSourcePrompt::new("\u{202E}\u{202C}\u{0000}\u{001B}"),
+            None
+        );
+    }
+}

crates/zed/src/main.rs ๐Ÿ”—

@@ -914,7 +914,9 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
                 })
                 .detach_and_log_err(cx);
             }
-            OpenRequestKind::AgentPanel { initial_prompt } => {
+            OpenRequestKind::AgentPanel {
+                external_source_prompt,
+            } => {
                 cx.spawn(async move |cx| {
                     let multi_workspace =
                         workspace::get_any_active_multi_workspace(app_state, cx.clone()).await?;
@@ -923,7 +925,11 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
                         multi_workspace.workspace().update(cx, |workspace, cx| {
                             if let Some(panel) = workspace.focus_panel::<AgentPanel>(window, cx) {
                                 panel.update(cx, |panel, cx| {
-                                    panel.new_external_thread_with_text(initial_prompt, window, cx);
+                                    panel.new_agent_thread_with_external_source_prompt(
+                                        external_source_prompt,
+                                        window,
+                                        cx,
+                                    );
                                 });
                             }
                         });

crates/zed/src/zed/open_listener.rs ๐Ÿ”—

@@ -1,5 +1,6 @@
 use crate::handle_open_request;
 use crate::restore_or_create_workspace;
+use agent_ui::ExternalSourcePrompt;
 use anyhow::{Context as _, Result, anyhow};
 use cli::{CliRequest, CliResponse, ipc::IpcSender};
 use cli::{IpcHandshake, ipc};
@@ -48,7 +49,7 @@ pub enum OpenRequestKind {
         extension_id: String,
     },
     AgentPanel {
-        initial_prompt: Option<String>,
+        external_source_prompt: Option<ExternalSourcePrompt>,
     },
     SharedAgentThread {
         session_id: String,
@@ -164,13 +165,14 @@ impl OpenRequest {
 
     fn parse_agent_url(&mut self, agent_path: &str) {
         // Format: "" or "?prompt=<text>"
-        let initial_prompt = agent_path.strip_prefix('?').and_then(|query| {
+        let external_source_prompt = agent_path.strip_prefix('?').and_then(|query| {
             url::form_urlencoded::parse(query.as_bytes())
                 .find_map(|(key, value)| (key == "prompt").then_some(value))
-                .filter(|s| !s.is_empty())
-                .map(|s| s.into_owned())
+                .and_then(|prompt| ExternalSourcePrompt::new(prompt.as_ref()))
+        });
+        self.kind = Some(OpenRequestKind::AgentPanel {
+            external_source_prompt,
         });
-        self.kind = Some(OpenRequestKind::AgentPanel { initial_prompt });
     }
 
     fn parse_git_clone_url(&mut self, clone_path: &str) -> Result<()> {
@@ -788,21 +790,30 @@ mod tests {
         });
 
         match request.kind {
-            Some(OpenRequestKind::AgentPanel { initial_prompt }) => {
-                assert_eq!(initial_prompt, None);
+            Some(OpenRequestKind::AgentPanel {
+                external_source_prompt,
+            }) => {
+                assert_eq!(external_source_prompt, None);
             }
             _ => panic!("Expected AgentPanel kind"),
         }
     }
 
+    fn agent_url_with_prompt(prompt: &str) -> String {
+        let mut serializer = url::form_urlencoded::Serializer::new("zed://agent?".to_string());
+        serializer.append_pair("prompt", prompt);
+        serializer.finish()
+    }
+
     #[gpui::test]
     fn test_parse_agent_url_with_prompt(cx: &mut TestAppContext) {
         let _app_state = init_test(cx);
+        let prompt = "Write me a script\nThanks";
 
         let request = cx.update(|cx| {
             OpenRequest::parse(
                 RawOpenRequest {
-                    urls: vec!["zed://agent?prompt=Write%20me%20a%20script%0AThanks".into()],
+                    urls: vec![agent_url_with_prompt(prompt)],
                     ..Default::default()
                 },
                 cx,
@@ -811,10 +822,14 @@ mod tests {
         });
 
         match request.kind {
-            Some(OpenRequestKind::AgentPanel { initial_prompt }) => {
+            Some(OpenRequestKind::AgentPanel {
+                external_source_prompt,
+            }) => {
                 assert_eq!(
-                    initial_prompt,
-                    Some("Write me a script\nThanks".to_string())
+                    external_source_prompt
+                        .as_ref()
+                        .map(ExternalSourcePrompt::as_str),
+                    Some("Write me a script\nThanks")
                 );
             }
             _ => panic!("Expected AgentPanel kind"),
@@ -828,7 +843,7 @@ mod tests {
         let request = cx.update(|cx| {
             OpenRequest::parse(
                 RawOpenRequest {
-                    urls: vec!["zed://agent?prompt=".into()],
+                    urls: vec![agent_url_with_prompt("")],
                     ..Default::default()
                 },
                 cx,
@@ -837,8 +852,10 @@ mod tests {
         });
 
         match request.kind {
-            Some(OpenRequestKind::AgentPanel { initial_prompt }) => {
-                assert_eq!(initial_prompt, None);
+            Some(OpenRequestKind::AgentPanel {
+                external_source_prompt,
+            }) => {
+                assert_eq!(external_source_prompt, None);
             }
             _ => panic!("Expected AgentPanel kind"),
         }