1190 lines · 42.6 KB · Rust
RawDownload
use super::edit_session::{    EditSession, EditSessionContext, EditSessionMode, EditSessionOutput, EditSessionResult,    initial_title_from_partial_path, run_session,};use crate::{AgentTool, Thread, ToolCallEventStream, ToolInput, ToolInputPayload};use action_log::ActionLog;use agent_client_protocol::schema as acp;use futures::FutureExt as _;use gpui::{App, AsyncApp, Entity, Task, WeakEntity};use language::LanguageRegistry;use project::Project;use schemars::JsonSchema;use serde::{Deserialize, Serialize};use std::path::PathBuf;use std::sync::Arc;use ui::SharedString; const DEFAULT_UI_TEXT: &str = "Writing file"; /// This is a tool for creating a new file or overwriting an existing file with completely new contents.////// To make granular edits to an existing file, prefer the `edit_file` tool instead.////// Before using this tool:////// 1. Verify the directory path is correct (only applicable when creating new files):///    - Use the `list_directory` tool to verify the parent directory exists and is the correct location#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]pub struct WriteFileToolInput {    /// The full path of the file to create or overwrite in the project.    ///    /// WARNING: When specifying which file path need changing, you MUST start each path with one of the project's root directories.    ///    /// The following examples assume we have two root directories in the project:    /// - /a/b/backend    /// - /c/d/frontend    ///    /// <example>    /// `backend/src/main.rs`    ///    /// Notice how the file path starts with `backend`. Without that, the path would be ambiguous and the call would fail!    /// </example>    ///    /// <example>    /// `frontend/db.js`    /// </example>    pub path: PathBuf,     /// The complete content for the file.    /// This field should contain the entire file content.    pub content: String,} #[derive(Clone, Default, Debug, Deserialize)]struct WriteFileToolPartialInput {    #[serde(default)]    path: Option<String>,    #[serde(default)]    content: Option<String>,} pub struct WriteFileTool {    session_context: Arc<EditSessionContext>,} impl WriteFileTool {    pub fn new(        project: Entity<Project>,        thread: WeakEntity<Thread>,        action_log: Entity<ActionLog>,        language_registry: Arc<LanguageRegistry>,    ) -> Self {        Self {            session_context: Arc::new(EditSessionContext::new(                project,                thread,                action_log,                language_registry,            )),        }    }     async fn process_streaming_writes(        &self,        input: &mut ToolInput<WriteFileToolInput>,        event_stream: &ToolCallEventStream,        cx: &mut AsyncApp,    ) -> EditSessionResult {        let mut session: Option<EditSession> = None;        let mut last_path: Option<String> = None;         loop {            futures::select! {                payload = input.next().fuse() => {                    match payload {                        Ok(payload) => match payload {                            ToolInputPayload::Partial(partial) => {                                if let Ok(parsed) = serde_json::from_value::<WriteFileToolPartialInput>(partial) {                                    let path_complete = parsed.path.is_some()                                        && parsed.path.as_ref() == last_path.as_ref();                                     last_path = parsed.path.clone();                                     if session.is_none()                                        && path_complete                                        && let Some(path) = parsed.path.as_ref()                                    {                                        match EditSession::new(                                            PathBuf::from(path),                                            EditSessionMode::Write,                                            Self::NAME,                                            self.session_context.clone(),                                            event_stream,                                            cx,                                        )                                        .await                                        {                                            Ok(created_session) => session = Some(created_session),                                            Err(error) => {                                                log::error!("Failed to create edit session: {}", error);                                                return EditSessionResult::Failed {                                                    error,                                                    session: None,                                                };                                            }                                        }                                    }                                     if let Some(current_session) = &mut session                                        && let Err(error) = current_session.process_write(parsed.content.as_deref(), cx)                                    {                                        log::error!("Failed to process write: {}", error);                                        return EditSessionResult::Failed { error, session };                                    }                                }                            }                            ToolInputPayload::Full(full_input) => {                                let mut session = if let Some(session) = session {                                    session                                } else {                                    match EditSession::new(                                        full_input.path.clone(),                                        EditSessionMode::Write,                                        Self::NAME,                                        self.session_context.clone(),                                        event_stream,                                        cx,                                    )                                    .await                                    {                                        Ok(created_session) => created_session,                                        Err(error) => {                                            log::error!("Failed to create edit session: {}", error);                                            return EditSessionResult::Failed {                                                error,                                                session: None,                                            };                                        }                                    }                                };                                 return match session.finalize_write(&full_input.content, cx).await {                                    Ok(()) => EditSessionResult::Completed(session),                                    Err(error) => {                                        log::error!("Failed to finalize write: {}", error);                                        EditSessionResult::Failed {                                            error,                                            session: Some(session),                                        }                                    }                                };                            }                            ToolInputPayload::InvalidJson { error_message } => {                                log::error!("Received invalid JSON: {error_message}");                                return EditSessionResult::Failed {                                    error: error_message,                                    session,                                };                            }                        },                        Err(error) => {                            return EditSessionResult::Failed {                                error: error.to_string(),                                session,                            };                        }                    }                }                _ = event_stream.cancelled_by_user().fuse() => {                    return EditSessionResult::Failed {                        error: "Write cancelled by user".to_string(),                        session,                    };                }            }        }    }} impl AgentTool for WriteFileTool {    type Input = WriteFileToolInput;    type Output = EditSessionOutput;     const NAME: &'static str = "write_file";     fn supports_input_streaming() -> bool {        true    }     fn kind() -> acp::ToolKind {        acp::ToolKind::Edit    }     fn initial_title(        &self,        input: Result<Self::Input, serde_json::Value>,        cx: &mut App,    ) -> SharedString {        match input {            Ok(input) => {                self.session_context                    .initial_title_from_path(&input.path, DEFAULT_UI_TEXT, cx)            }            Err(raw_input) => initial_title_from_partial_path::<WriteFileToolPartialInput>(                &self.session_context,                raw_input,                |partial| partial.path.clone(),                DEFAULT_UI_TEXT,                cx,            ),        }    }     fn run(        self: Arc<Self>,        mut input: ToolInput<Self::Input>,        event_stream: ToolCallEventStream,        cx: &mut App,    ) -> Task<Result<Self::Output, Self::Output>> {        cx.spawn(async move |cx: &mut AsyncApp| {            run_session(                self.process_streaming_writes(&mut input, &event_stream, cx)                    .await,                cx,            )            .await        })    }     fn replay(        &self,        _input: Self::Input,        output: Self::Output,        event_stream: ToolCallEventStream,        cx: &mut App,    ) -> anyhow::Result<()> {        self.session_context.replay_output(output, event_stream, cx)    }} #[cfg(test)]mod tests {    use super::*;    use crate::{        AgentTool, ContextServerRegistry, Templates, Thread, ToolCallEventStream, ToolInput,        ToolInputSender,    };    use acp_thread::Diff;    use action_log::ActionLog;    use fs::Fs as _;    use futures::StreamExt as _;    use gpui::{AppContext as _, Entity, TestAppContext, UpdateGlobal};    use language::language_settings::FormatOnSave;    use language_model::fake_provider::FakeLanguageModel;    use project::{Project, ProjectPath};    use prompt_store::ProjectContext;    use serde_json::json;    use settings::{Settings, SettingsStore};    use std::sync::Arc;    use util::path;    use util::rel_path::{RelPath, rel_path};     #[gpui::test]    async fn test_streaming_write_create_file(cx: &mut TestAppContext) {        let (write_tool, _project, _action_log, _fs, _thread) =            setup_test(cx, json!({"dir": {}})).await;        let result = cx            .update(|cx| {                write_tool.clone().run(                    ToolInput::resolved(WriteFileToolInput {                        path: "root/dir/new_file.txt".into(),                        content: "Hello, World!".into(),                    }),                    ToolCallEventStream::test().0,                    cx,                )            })            .await;         let EditSessionOutput::Success { new_text, diff, .. } = result.unwrap() else {            panic!("expected success");        };        assert_eq!(new_text, "Hello, World!");        assert!(!diff.is_empty());    }     #[gpui::test]    async fn test_streaming_write_overwrite_file(cx: &mut TestAppContext) {        let (write_tool, _project, _action_log, _fs, _thread) =            setup_test(cx, json!({"file.txt": "old content"})).await;        let result = cx            .update(|cx| {                write_tool.clone().run(                    ToolInput::resolved(WriteFileToolInput {                        path: "root/file.txt".into(),                        content: "new content".into(),                    }),                    ToolCallEventStream::test().0,                    cx,                )            })            .await;         let EditSessionOutput::Success {            new_text, old_text, ..        } = result.unwrap()        else {            panic!("expected success");        };        assert_eq!(new_text, "new content");        assert_eq!(*old_text, "old content");    }     #[gpui::test]    async fn test_streaming_path_completeness_heuristic(cx: &mut TestAppContext) {        let (write_tool, _project, _action_log, _fs, _thread) =            setup_test(cx, json!({"file.txt": "hello world"})).await;        let (mut sender, input) = ToolInput::<WriteFileToolInput>::test();        let (event_stream, _receiver) = ToolCallEventStream::test();        let task = cx.update(|cx| write_tool.clone().run(input, event_stream, cx));         // Send partial with path but NO mode — path should NOT be treated as complete        sender.send_partial(json!({            "path": "root/file"        }));        cx.run_until_parked();         // Now the path grows and mode appears        sender.send_partial(json!({            "path": "root/file.txt",        }));        cx.run_until_parked();         // Send final        sender.send_full(json!({            "path": "root/file.txt",            "content": "new content"        }));         let result = task.await;        let EditSessionOutput::Success { new_text, .. } = result.unwrap() else {            panic!("expected success");        };        assert_eq!(new_text, "new content");    }     #[gpui::test]    async fn test_streaming_create_file_with_partials(cx: &mut TestAppContext) {        let (write_tool, _project, _action_log, _fs, _thread) =            setup_test(cx, json!({"dir": {}})).await;        let (mut sender, input) = ToolInput::<WriteFileToolInput>::test();        let (event_stream, _receiver) = ToolCallEventStream::test();        let task = cx.update(|cx| write_tool.clone().run(input, event_stream, cx));         // Stream partials for create mode        sender.send_partial(json!({}));        cx.run_until_parked();         sender.send_partial(json!({            "path": "root/dir/new_file.txt",        }));        cx.run_until_parked();         sender.send_partial(json!({            "path": "root/dir/new_file.txt",            "content": "Hello, "        }));        cx.run_until_parked();         // Final with full content        sender.send_full(json!({            "path": "root/dir/new_file.txt",            "content": "Hello, World!"        }));         let result = task.await;        let EditSessionOutput::Success { new_text, .. } = result.unwrap() else {            panic!("expected success");        };        assert_eq!(new_text, "Hello, World!");    }     #[gpui::test]    async fn test_streaming_input_recv_drains_partials(cx: &mut TestAppContext) {        let (write_tool, _project, _action_log, _fs, _thread) =            setup_test(cx, json!({"dir": {}})).await;        // Create a channel and send multiple partials before a final, then use        // ToolInput::resolved-style immediate delivery to confirm recv() works        // when partials are already buffered.        let (mut sender, input): (ToolInputSender, ToolInput<WriteFileToolInput>) =            ToolInput::test();        let (event_stream, _event_rx) = ToolCallEventStream::test();        let task = cx.update(|cx| write_tool.clone().run(input, event_stream, cx));         // Buffer several partials before sending the final        sender.send_partial(json!({}));        sender.send_partial(json!({"path": "root/dir/new.txt"}));        sender.send_partial(json!({            "path": "root/dir/new.txt",        }));        sender.send_full(json!({            "path": "root/dir/new.txt",            "content": "streamed content"        }));         let result = task.await;        let EditSessionOutput::Success { new_text, .. } = result.unwrap() else {            panic!("expected success");        };        assert_eq!(new_text, "streamed content");    }     #[gpui::test]    async fn test_streaming_resolve_path_for_creating_file(cx: &mut TestAppContext) {        let mode = EditSessionMode::Write;         let result = test_resolve_path(&mode, "root/new.txt", cx);        assert_resolved_path_eq(result.await, rel_path("new.txt"));         let result = test_resolve_path(&mode, "new.txt", cx);        assert_resolved_path_eq(result.await, rel_path("new.txt"));         let result = test_resolve_path(&mode, "dir/new.txt", cx);        assert_resolved_path_eq(result.await, rel_path("dir/new.txt"));         let result = test_resolve_path(&mode, "root/dir/subdir/existing.txt", cx);        assert_resolved_path_eq(result.await, rel_path("dir/subdir/existing.txt"));         let result = test_resolve_path(&mode, "root/dir/subdir", cx);        assert_eq!(            result.await.unwrap_err(),            "Can't write to file: path is a directory"        );         let result = test_resolve_path(&mode, "root/dir/nonexistent_dir/new.txt", cx);        assert_eq!(            result.await.unwrap_err(),            "Can't create file: parent directory doesn't exist"        );    }     #[gpui::test]    async fn test_streaming_format_on_save(cx: &mut TestAppContext) {        init_test(cx);         let fs = project::FakeFs::new(cx.executor());        fs.insert_tree("/root", json!({"src": {}})).await;        let (write_tool, project, action_log, fs, thread) =            setup_test_with_fs(cx, fs, &[path!("/root").as_ref()]).await;         let rust_language = Arc::new(language::Language::new(            language::LanguageConfig {                name: "Rust".into(),                matcher: language::LanguageMatcher {                    path_suffixes: vec!["rs".to_string()],                    ..Default::default()                },                ..Default::default()            },            None,        ));         let language_registry = project.read_with(cx, |project, _| project.languages().clone());        language_registry.add(rust_language);         let mut fake_language_servers = language_registry.register_fake_lsp(            "Rust",            language::FakeLspAdapter {                capabilities: lsp::ServerCapabilities {                    document_formatting_provider: Some(lsp::OneOf::Left(true)),                    ..Default::default()                },                ..Default::default()            },        );         fs.save(            path!("/root/src/main.rs").as_ref(),            &"initial content".into(),            language::LineEnding::Unix,

Lines 1–500 of 1190.