618 lines · 22.4 KB · Rust
RawDownload
use super::*;use acp_thread::UserMessageId;use action_log::ActionLog;use fs::FakeFs;use language_model::{    LanguageModelCompletionEvent, LanguageModelToolUse, MessageContent, StopReason,    fake_provider::FakeLanguageModel,};use prompt_store::ProjectContext;use serde_json::json;use std::{collections::BTreeMap, sync::Arc, time::Duration};use util::path; #[gpui::test]async fn test_edit_file_tool_in_thread_context(cx: &mut TestAppContext) {    // This test verifies that the edit_file tool works correctly when invoked    // through the full thread flow (model sends ToolUse event -> tool runs -> result sent back).    // This is different from tests that call tool.run() directly.    super::init_test(cx);    super::always_allow_tools(cx);     let fs = FakeFs::new(cx.executor());    fs.insert_tree(        path!("/project"),        json!({            "src": {                "main.rs": "fn main() {\n    println!(\"Hello, world!\");\n}\n"            }        }),    )    .await;     let project = project::Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;    let project_context = cx.new(|_cx| ProjectContext::default());    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());    let context_server_registry =        cx.new(|cx| crate::ContextServerRegistry::new(context_server_store.clone(), cx));    let model = Arc::new(FakeLanguageModel::default());    let fake_model = model.as_fake();     let thread = cx.new(|cx| {        let mut thread = crate::Thread::new(            project.clone(),            project_context,            context_server_registry,            crate::Templates::new(),            Some(model.clone()),            cx,        );        // Add just the tools we need for this test        let language_registry = project.read(cx).languages().clone();        thread.add_tool(crate::ReadFileTool::new(            cx.weak_entity(),            project.clone(),            thread.action_log().clone(),        ));        thread.add_tool(crate::EditFileTool::new(            project.clone(),            cx.weak_entity(),            language_registry,            crate::Templates::new(),        ));        thread    });     // First, read the file so the thread knows about its contents    let _events = thread        .update(cx, |thread, cx| {            thread.send(UserMessageId::new(), ["Read the file src/main.rs"], cx)        })        .unwrap();    cx.run_until_parked();     // Model calls read_file tool    let read_tool_use = LanguageModelToolUse {        id: "read_tool_1".into(),        name: "read_file".into(),        raw_input: json!({"path": "project/src/main.rs"}).to_string(),        input: json!({"path": "project/src/main.rs"}),        is_input_complete: true,        thought_signature: None,    };    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(read_tool_use));    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::ToolUse));    fake_model.end_last_completion_stream();    cx.run_until_parked();     // Wait for the read tool to complete and model to be called again    while fake_model.pending_completions().is_empty() {        cx.run_until_parked();    }     // Model responds after seeing the file content, then calls edit_file    fake_model.send_last_completion_stream_text_chunk("I'll edit the file now.");    let edit_tool_use = LanguageModelToolUse {        id: "edit_tool_1".into(),        name: "edit_file".into(),        raw_input: json!({            "display_description": "Change greeting message",            "path": "project/src/main.rs",            "mode": "edit"        })        .to_string(),        input: json!({            "display_description": "Change greeting message",            "path": "project/src/main.rs",            "mode": "edit"        }),        is_input_complete: true,        thought_signature: None,    };    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(edit_tool_use));    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::ToolUse));    fake_model.end_last_completion_stream();    cx.run_until_parked();     // The edit_file tool creates an EditAgent which makes its own model request.    // We need to respond to that request with the edit instructions.    // Wait for the edit agent's completion request    let deadline = std::time::Instant::now() + Duration::from_secs(5);    while fake_model.pending_completions().is_empty() {        if std::time::Instant::now() >= deadline {            panic!(                "Timed out waiting for edit agent completion request. Pending: {}",                fake_model.pending_completions().len()            );        }        cx.run_until_parked();        cx.background_executor            .timer(Duration::from_millis(10))            .await;    }     // Send the edit agent's response with the XML format it expects    let edit_response = "<old_text>println!(\"Hello, world!\");</old_text>\n<new_text>println!(\"Hello, Zed!\");</new_text>";    fake_model.send_last_completion_stream_text_chunk(edit_response);    fake_model.end_last_completion_stream();    cx.run_until_parked();     // Wait for the edit to complete and the thread to call the model again with tool results    let deadline = std::time::Instant::now() + Duration::from_secs(5);    while fake_model.pending_completions().is_empty() {        if std::time::Instant::now() >= deadline {            panic!("Timed out waiting for model to be called after edit completion");        }        cx.run_until_parked();        cx.background_executor            .timer(Duration::from_millis(10))            .await;    }     // Verify the file was edited    let file_content = fs        .load(path!("/project/src/main.rs").as_ref())        .await        .expect("file should exist");    assert!(        file_content.contains("Hello, Zed!"),        "File should have been edited. Content: {}",        file_content    );    assert!(        !file_content.contains("Hello, world!"),        "Old content should be replaced. Content: {}",        file_content    );     // Verify the tool result was sent back to the model    let pending = fake_model.pending_completions();    assert!(        !pending.is_empty(),        "Model should have been called with tool result"    );     let last_request = pending.last().unwrap();    let has_tool_result = last_request.messages.iter().any(|m| {        m.content            .iter()            .any(|c| matches!(c, language_model::MessageContent::ToolResult(_)))    });    assert!(        has_tool_result,        "Tool result should be in the messages sent back to the model"    );     // Complete the turn    fake_model.send_last_completion_stream_text_chunk("I've updated the greeting message.");    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));    fake_model.end_last_completion_stream();    cx.run_until_parked();     // Verify the thread completed successfully    thread.update(cx, |thread, _cx| {        assert!(            thread.is_turn_complete(),            "Thread should be complete after the turn ends"        );    });} #[gpui::test]async fn test_subagent_uses_read_file_tool(cx: &mut TestAppContext) {    // This test verifies that subagents can successfully use the read_file tool    // through the full thread flow, and that tools are properly rebound to use    // the subagent's thread ID instead of the parent's.    super::init_test(cx);    super::always_allow_tools(cx);     cx.update(|cx| {        cx.update_flags(true, vec!["subagents".to_string()]);    });     let fs = FakeFs::new(cx.executor());    fs.insert_tree(        path!("/project"),        json!({            "src": {                "lib.rs": "pub fn hello() -> &'static str {\n    \"Hello from lib!\"\n}\n"            }        }),    )    .await;     let project = project::Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;    let project_context = cx.new(|_cx| ProjectContext::default());    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());    let context_server_registry =        cx.new(|cx| crate::ContextServerRegistry::new(context_server_store.clone(), cx));    let model = Arc::new(FakeLanguageModel::default());    let fake_model = model.as_fake();     // Create subagent context    let subagent_context = crate::SubagentContext {        parent_thread_id: agent_client_protocol::SessionId::new("parent-id"),        tool_use_id: language_model::LanguageModelToolUseId::from("subagent-tool-use-id"),        depth: 1,        summary_prompt: "Summarize what you found".to_string(),        context_low_prompt: "Context low".to_string(),    };     // Create parent tools that will be passed to the subagent    // This simulates how the subagent_tool passes tools to new_subagent    let parent_tools: BTreeMap<gpui::SharedString, std::sync::Arc<dyn crate::AnyAgentTool>> = {        let action_log = cx.new(|_| ActionLog::new(project.clone()));        // Create a "fake" parent thread reference - this should get rebound        let fake_parent_thread = cx.new(|cx| {            crate::Thread::new(                project.clone(),                cx.new(|_cx| ProjectContext::default()),                cx.new(|cx| crate::ContextServerRegistry::new(context_server_store.clone(), cx)),                crate::Templates::new(),                Some(model.clone()),                cx,            )        });        let mut tools: BTreeMap<gpui::SharedString, std::sync::Arc<dyn crate::AnyAgentTool>> =            BTreeMap::new();        tools.insert(            "read_file".into(),            crate::ReadFileTool::new(fake_parent_thread.downgrade(), project.clone(), action_log)                .erase(),        );        tools    };     // Create subagent - tools should be rebound to use subagent's thread    let subagent = cx.new(|cx| {        crate::Thread::new_subagent(            project.clone(),            project_context,            context_server_registry,            crate::Templates::new(),            model.clone(),            subagent_context,            parent_tools,            cx,        )    });     // Get the subagent's thread ID    let _subagent_thread_id = subagent.read_with(cx, |thread, _| thread.id().to_string());     // Verify the subagent has the read_file tool    subagent.read_with(cx, |thread, _| {        assert!(            thread.has_registered_tool("read_file"),            "subagent should have read_file tool"        );    });     // Submit a user message to the subagent    subagent        .update(cx, |thread, cx| {            thread.submit_user_message("Read the file src/lib.rs", cx)        })        .unwrap();    cx.run_until_parked();     // Simulate the model calling the read_file tool    let read_tool_use = LanguageModelToolUse {        id: "read_tool_1".into(),        name: "read_file".into(),        raw_input: json!({"path": "project/src/lib.rs"}).to_string(),        input: json!({"path": "project/src/lib.rs"}),        is_input_complete: true,        thought_signature: None,    };    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(read_tool_use));    fake_model.end_last_completion_stream();    cx.run_until_parked();     // Wait for the tool to complete and the model to be called again with tool results    let deadline = std::time::Instant::now() + Duration::from_secs(5);    while fake_model.pending_completions().is_empty() {        if std::time::Instant::now() >= deadline {            panic!("Timed out waiting for model to be called after read_file tool completion");        }        cx.run_until_parked();        cx.background_executor            .timer(Duration::from_millis(10))            .await;    }     // Verify the tool result was sent back to the model    let pending = fake_model.pending_completions();    assert!(        !pending.is_empty(),        "Model should have been called with tool result"    );     let last_request = pending.last().unwrap();    let tool_result = last_request.messages.iter().find_map(|m| {        m.content.iter().find_map(|c| match c {            MessageContent::ToolResult(result) => Some(result),            _ => None,        })    });    assert!(        tool_result.is_some(),        "Tool result should be in the messages sent back to the model"    );     // Verify the tool result contains the file content    let result = tool_result.unwrap();    let result_text = match &result.content {        language_model::LanguageModelToolResultContent::Text(text) => text.to_string(),        _ => panic!("expected text content in tool result"),    };    assert!(        result_text.contains("Hello from lib!"),        "Tool result should contain file content, got: {}",        result_text    );     // Verify the subagent is ready for more input (tool completed, model called again)    // This test verifies the subagent can successfully use read_file tool.    // The summary flow is tested separately in test_subagent_returns_summary_on_completion.} #[gpui::test]async fn test_subagent_uses_edit_file_tool(cx: &mut TestAppContext) {    // This test verifies that subagents can successfully use the edit_file tool    // through the full thread flow, including the edit agent's model request.    // It also verifies that the edit agent uses the subagent's thread ID, not the parent's.    super::init_test(cx);    super::always_allow_tools(cx);     cx.update(|cx| {        cx.update_flags(true, vec!["subagents".to_string()]);    });     let fs = FakeFs::new(cx.executor());    fs.insert_tree(        path!("/project"),        json!({            "src": {                "config.rs": "pub const VERSION: &str = \"1.0.0\";\n"            }        }),    )    .await;     let project = project::Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;    let project_context = cx.new(|_cx| ProjectContext::default());    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());    let context_server_registry =        cx.new(|cx| crate::ContextServerRegistry::new(context_server_store.clone(), cx));    let model = Arc::new(FakeLanguageModel::default());    let fake_model = model.as_fake();     // Create a "parent" thread to simulate the real scenario where tools are inherited    let parent_thread = cx.new(|cx| {        crate::Thread::new(            project.clone(),            cx.new(|_cx| ProjectContext::default()),            cx.new(|cx| crate::ContextServerRegistry::new(context_server_store.clone(), cx)),            crate::Templates::new(),            Some(model.clone()),            cx,        )    });    let parent_thread_id = parent_thread.read_with(cx, |thread, _| thread.id().to_string());     // Create parent tools that reference the parent thread    let parent_tools: BTreeMap<gpui::SharedString, std::sync::Arc<dyn crate::AnyAgentTool>> = {        let action_log = cx.new(|_| ActionLog::new(project.clone()));        let language_registry = project.read_with(cx, |p, _| p.languages().clone());        let mut tools: BTreeMap<gpui::SharedString, std::sync::Arc<dyn crate::AnyAgentTool>> =            BTreeMap::new();        tools.insert(            "read_file".into(),            crate::ReadFileTool::new(parent_thread.downgrade(), project.clone(), action_log)                .erase(),        );        tools.insert(            "edit_file".into(),            crate::EditFileTool::new(                project.clone(),                parent_thread.downgrade(),                language_registry,                crate::Templates::new(),            )            .erase(),        );        tools    };     // Create subagent context    let subagent_context = crate::SubagentContext {        parent_thread_id: agent_client_protocol::SessionId::new("parent-id"),        tool_use_id: language_model::LanguageModelToolUseId::from("subagent-tool-use-id"),        depth: 1,        summary_prompt: "Summarize what you changed".to_string(),        context_low_prompt: "Context low".to_string(),    };     // Create subagent - tools should be rebound to use subagent's thread    let subagent = cx.new(|cx| {        crate::Thread::new_subagent(            project.clone(),            project_context,            context_server_registry,            crate::Templates::new(),            model.clone(),            subagent_context,            parent_tools,            cx,        )    });     // Get the subagent's thread ID - it should be different from parent    let subagent_thread_id = subagent.read_with(cx, |thread, _| thread.id().to_string());    assert_ne!(        parent_thread_id, subagent_thread_id,        "Subagent should have a different thread ID than parent"    );     // Verify the subagent has the tools    subagent.read_with(cx, |thread, _| {        assert!(            thread.has_registered_tool("read_file"),            "subagent should have read_file tool"        );        assert!(            thread.has_registered_tool("edit_file"),            "subagent should have edit_file tool"        );    });     // Submit a user message to the subagent    subagent        .update(cx, |thread, cx| {            thread.submit_user_message("Update the version in config.rs to 2.0.0", cx)        })        .unwrap();    cx.run_until_parked();     // First, model calls read_file to see the current content    let read_tool_use = LanguageModelToolUse {        id: "read_tool_1".into(),        name: "read_file".into(),        raw_input: json!({"path": "project/src/config.rs"}).to_string(),        input: json!({"path": "project/src/config.rs"}),        is_input_complete: true,        thought_signature: None,    };    fake_model        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(read_tool_use));    fake_model.end_last_completion_stream();    cx.run_until_parked();     // Wait for the read tool to complete and model to be called again    let deadline = std::time::Instant::now() + Duration::from_secs(5);    while fake_model.pending_completions().is_empty() {

Lines 1–500 of 618.