Cameron Mcloughlinagent: Re-enable sandbox (#61711)12a19dc3 months agoHistory
8721 lines · 292.1 KB · Rust
RawDownload
use super::*;use acp_thread::{    AgentConnection, AgentModelGroupName, AgentModelList, ClientUserMessageId, PermissionOptions,    ThreadStatus,};use agent_client_protocol::schema::v1 as acp;use agent_settings::AgentProfileId;use anyhow::Result;use client::{Client, RefreshLlmTokenListener, UserStore};use collections::IndexMap;use context_server::{ContextServer, ContextServerCommand, ContextServerId};use feature_flags::FeatureFlagAppExt as _;use fs::{FakeFs, Fs};use futures::{    FutureExt as _, StreamExt,    channel::{        mpsc::{self, UnboundedReceiver},        oneshot,    },    future::{Fuse, Shared},};use gpui::{    App, AppContext, AsyncApp, Entity, Task, TestAppContext, UpdateGlobal,    http_client::FakeHttpClient,};use indoc::indoc;use language_model::{    CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelId, LanguageModelImageExt, LanguageModelProviderId, LanguageModelProviderName,    LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelToolResult, LanguageModelToolSchemaFormat, LanguageModelToolUse, MessageContent,    Role, StopReason, TokenUsage,    fake_provider::{FakeLanguageModel, FakeLanguageModelProvider},};use pretty_assertions::assert_eq;use project::{    Project, context_server_store::ContextServerStore, project_settings::ProjectSettings,};use prompt_store::ProjectContext;use reqwest_client::ReqwestClient;use schemars::JsonSchema;use serde::{Deserialize, Serialize};use serde_json::json;use settings::{LanguageModelProviderSetting, LanguageModelSelection, Settings, SettingsStore};use std::{    path::Path,    pin::Pin,    rc::Rc,    sync::{        Arc,        atomic::{AtomicBool, AtomicUsize, Ordering},    },    time::Duration,};use util::path; mod test_tools;use test_tools::*; pub(crate) fn init_test(cx: &mut TestAppContext) {    cx.update(|cx| {        let settings_store = SettingsStore::test(cx);        cx.set_global(settings_store);    });} pub(crate) struct FakeTerminalHandle {    killed: Arc<AtomicBool>,    stopped_by_user: Arc<AtomicBool>,    exit_sender: std::cell::RefCell<Option<futures::channel::oneshot::Sender<()>>>,    wait_for_exit: Shared<Task<acp::TerminalExitStatus>>,    output: acp::TerminalOutputResponse,    id: acp::TerminalId,} impl FakeTerminalHandle {    pub(crate) fn new_never_exits(cx: &mut App) -> Self {        let killed = Arc::new(AtomicBool::new(false));        let stopped_by_user = Arc::new(AtomicBool::new(false));         let (exit_sender, exit_receiver) = futures::channel::oneshot::channel();         let wait_for_exit = cx            .spawn(async move |_cx| {                // Wait for the exit signal (sent when kill() is called)                let _ = exit_receiver.await;                acp::TerminalExitStatus::new()            })            .shared();         Self {            killed,            stopped_by_user,            exit_sender: std::cell::RefCell::new(Some(exit_sender)),            wait_for_exit,            output: acp::TerminalOutputResponse::new("partial output".to_string(), false),            id: acp::TerminalId::new("fake_terminal".to_string()),        }    }     pub(crate) fn new_with_immediate_exit(cx: &mut App, exit_code: u32) -> Self {        let killed = Arc::new(AtomicBool::new(false));        let stopped_by_user = Arc::new(AtomicBool::new(false));        let (exit_sender, _exit_receiver) = futures::channel::oneshot::channel();         let wait_for_exit = cx            .spawn(async move |_cx| acp::TerminalExitStatus::new().exit_code(exit_code))            .shared();         Self {            killed,            stopped_by_user,            exit_sender: std::cell::RefCell::new(Some(exit_sender)),            wait_for_exit,            output: acp::TerminalOutputResponse::new("command output".to_string(), false),            id: acp::TerminalId::new("fake_terminal".to_string()),        }    }     pub(crate) fn with_output(mut self, output: acp::TerminalOutputResponse) -> Self {        self.output = output;        self    }     pub(crate) fn was_killed(&self) -> bool {        self.killed.load(Ordering::SeqCst)    }     pub(crate) fn set_stopped_by_user(&self, stopped: bool) {        self.stopped_by_user.store(stopped, Ordering::SeqCst);    }     pub(crate) fn signal_exit(&self) {        if let Some(sender) = self.exit_sender.borrow_mut().take() {            let _ = sender.send(());        }    }} impl crate::TerminalHandle for FakeTerminalHandle {    fn id(&self, _cx: &AsyncApp) -> Result<acp::TerminalId> {        Ok(self.id.clone())    }     fn current_output(&self, _cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {        Ok(self.output.clone())    }     fn wait_for_exit(&self, _cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {        Ok(self.wait_for_exit.clone())    }     fn kill(&self, _cx: &AsyncApp) -> Result<()> {        self.killed.store(true, Ordering::SeqCst);        self.signal_exit();        Ok(())    }     fn was_stopped_by_user(&self, _cx: &AsyncApp) -> Result<bool> {        Ok(self.stopped_by_user.load(Ordering::SeqCst))    }} struct FakeSubagentHandle {    session_id: acp::SessionId,    send_task: Shared<Task<String>>,} impl SubagentHandle for FakeSubagentHandle {    fn id(&self) -> acp::SessionId {        self.session_id.clone()    }     fn num_entries(&self, _cx: &App) -> usize {        unimplemented!()    }     fn send(&self, _message: String, cx: &AsyncApp) -> Task<Result<String>> {        let task = self.send_task.clone();        cx.background_spawn(async move { Ok(task.await) })    }} #[derive(Default)]pub(crate) struct FakeThreadEnvironment {    terminal_handle: Option<Rc<FakeTerminalHandle>>,    subagent_handle: Option<Rc<FakeSubagentHandle>>,    terminal_creations: Arc<AtomicUsize>,    terminal_output_limits: std::cell::RefCell<Vec<Option<u64>>>,} impl FakeThreadEnvironment {    pub(crate) fn with_terminal(self, terminal_handle: FakeTerminalHandle) -> Self {        Self {            terminal_handle: Some(terminal_handle.into()),            ..self        }    }     pub(crate) fn terminal_creation_count(&self) -> usize {        self.terminal_creations.load(Ordering::SeqCst)    }     pub(crate) fn terminal_output_limits(&self) -> Vec<Option<u64>> {        self.terminal_output_limits.borrow().clone()    }} impl crate::ThreadEnvironment for FakeThreadEnvironment {    fn create_terminal(        &self,        _command: String,        _extra_env: Vec<acp::EnvVariable>,        _cwd: Option<std::path::PathBuf>,        output_byte_limit: Option<u64>,        _sandbox_wrap: Option<acp_thread::SandboxWrap>,        _cx: &mut AsyncApp,    ) -> Task<Result<Rc<dyn crate::TerminalHandle>>> {        self.terminal_creations.fetch_add(1, Ordering::SeqCst);        self.terminal_output_limits            .borrow_mut()            .push(output_byte_limit);        let handle = self            .terminal_handle            .clone()            .expect("Terminal handle not available on FakeThreadEnvironment");        Task::ready(Ok(handle as Rc<dyn crate::TerminalHandle>))    }     fn create_subagent(&self, _label: String, _cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {        Ok(self            .subagent_handle            .clone()            .expect("Subagent handle not available on FakeThreadEnvironment")            as Rc<dyn SubagentHandle>)    }} /// Environment that creates multiple independent terminal handles for testing concurrent terminals.struct MultiTerminalEnvironment {    handles: std::cell::RefCell<Vec<Rc<FakeTerminalHandle>>>,} impl MultiTerminalEnvironment {    fn new() -> Self {        Self {            handles: std::cell::RefCell::new(Vec::new()),        }    }     fn handles(&self) -> Vec<Rc<FakeTerminalHandle>> {        self.handles.borrow().clone()    }} impl crate::ThreadEnvironment for MultiTerminalEnvironment {    fn create_terminal(        &self,        _command: String,        _extra_env: Vec<acp::EnvVariable>,        _cwd: Option<std::path::PathBuf>,        _output_byte_limit: Option<u64>,        _sandbox_wrap: Option<acp_thread::SandboxWrap>,        cx: &mut AsyncApp,    ) -> Task<Result<Rc<dyn crate::TerminalHandle>>> {        let handle = Rc::new(cx.update(|cx| FakeTerminalHandle::new_never_exits(cx)));        self.handles.borrow_mut().push(handle.clone());        Task::ready(Ok(handle as Rc<dyn crate::TerminalHandle>))    }     fn create_subagent(&self, _label: String, _cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {        unimplemented!()    }} fn always_allow_tools(cx: &mut TestAppContext) {    cx.update(|cx| {        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();        settings.tool_permissions.default = settings::ToolPermissionMode::Allow;        agent_settings::AgentSettings::override_global(settings, cx);    });} /// Turns terminal sandboxing off so the non-sandboxed `TerminalTool` is the/// variant exposed to the model as `terminal`. Tests that register/// `TerminalTool` directly need this because the sandboxing feature flag is/// enabled for all, in which case `Thread::enabled_tools` would otherwise/// expose `SandboxedTerminalTool` under that name instead.fn disable_sandboxing(cx: &mut TestAppContext) {    cx.update(|cx| {        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();        settings.sandbox_permissions.allow_unsandboxed = true;        agent_settings::AgentSettings::override_global(settings, cx);    });} #[gpui::test]async fn test_echo(cx: &mut TestAppContext) {    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;    let fake_model = model.as_fake();     let events = thread        .update(cx, |thread, cx| {            thread.send(                ClientUserMessageId::new(),                ["Testing: Reply with 'Hello'"],                cx,            )        })        .unwrap();    cx.run_until_parked();    fake_model.send_last_completion_stream_text_chunk("Hello");    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));    fake_model.end_last_completion_stream();     let events = events.collect().await;    thread.update(cx, |thread, _cx| {        assert_eq!(            thread.last_received_or_pending_message().unwrap().role(),            Role::Assistant        );        assert_eq!(            thread                .last_received_or_pending_message()                .unwrap()                .to_markdown(),            "Hello\n"        )    });    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);} #[gpui::test]async fn test_terminal_tool_timeout_kills_handle(cx: &mut TestAppContext) {    init_test(cx);    always_allow_tools(cx);     let fs = FakeFs::new(cx.executor());    let project = Project::test(fs, [], cx).await;     let environment = Rc::new(cx.update(|cx| {        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))    }));    let handle = environment.terminal_handle.clone().unwrap();     #[allow(clippy::arc_with_non_send_sync)]    let tool = Arc::new(crate::TerminalTool::new(project, environment));    let (event_stream, mut rx) = crate::ToolCallEventStream::test();     let task = cx.update(|cx| {        tool.run(            ToolInput::resolved(crate::TerminalToolInput {                command: "sleep 1000".to_string(),                cd: ".".to_string(),                timeout_ms: Some(5),                ..Default::default()            }),            event_stream,            cx,        )    });     let update = rx.expect_update_fields().await;    assert!(        update.content.iter().any(|blocks| {            blocks                .iter()                .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))        }),        "expected tool call update to include terminal content"    );     let mut task_future: Pin<Box<Fuse<Task<Result<String, String>>>>> = Box::pin(task.fuse());     let deadline = std::time::Instant::now() + Duration::from_millis(500);    loop {        if let Some(result) = task_future.as_mut().now_or_never() {            let result = result.expect("terminal tool task should complete");             assert!(                handle.was_killed(),                "expected terminal handle to be killed on timeout"            );            assert!(                result.contains("partial output"),                "expected result to include terminal output, got: {result}"            );            return;        }         if std::time::Instant::now() >= deadline {            panic!("timed out waiting for terminal tool task to complete");        }         cx.run_until_parked();        cx.background_executor.timer(Duration::from_millis(1)).await;    }} #[gpui::test]#[ignore]async fn test_terminal_tool_without_timeout_does_not_kill_handle(cx: &mut TestAppContext) {    init_test(cx);    always_allow_tools(cx);     let fs = FakeFs::new(cx.executor());    let project = Project::test(fs, [], cx).await;     let environment = Rc::new(cx.update(|cx| {        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))    }));    let handle = environment.terminal_handle.clone().unwrap();     #[allow(clippy::arc_with_non_send_sync)]    let tool = Arc::new(crate::TerminalTool::new(project, environment));    let (event_stream, mut rx) = crate::ToolCallEventStream::test();     let _task = cx.update(|cx| {        tool.run(            ToolInput::resolved(crate::TerminalToolInput {                command: "sleep 1000".to_string(),                cd: ".".to_string(),                timeout_ms: None,                ..Default::default()            }),            event_stream,            cx,        )    });     let update = rx.expect_update_fields().await;    assert!(        update.content.iter().any(|blocks| {            blocks                .iter()                .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))        }),        "expected tool call update to include terminal content"    );     cx.background_executor        .timer(Duration::from_millis(25))        .await;     assert!(        !handle.was_killed(),        "did not expect terminal handle to be killed without a timeout"    );} #[gpui::test]async fn test_thinking(cx: &mut TestAppContext) {    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;    let fake_model = model.as_fake();     let events = thread        .update(cx, |thread, cx| {            thread.send(                ClientUserMessageId::new(),                [indoc! {"                    Testing:                     Generate a thinking step where you just think the word 'Think',                    and have your final answer be 'Hello'                "}],                cx,            )        })        .unwrap();    cx.run_until_parked();    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::Thinking {        text: "Think".to_string(),        signature: None,    });    fake_model.send_last_completion_stream_text_chunk("Hello");    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));    fake_model.end_last_completion_stream();     let events = events.collect().await;    thread.update(cx, |thread, _cx| {        assert_eq!(            thread.last_received_or_pending_message().unwrap().role(),            Role::Assistant        );        assert_eq!(            thread                .last_received_or_pending_message()                .unwrap()                .to_markdown(),            indoc! {"                <think>Think</think>                Hello            "}        )    });    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);} 

Lines 1–500 of 8721.