diff --git a/crates/agent2/src/agent2.rs b/crates/agent2/src/agent2.rs
index 577bf170dd..71d5ea711c 100644
--- a/crates/agent2/src/agent2.rs
+++ b/crates/agent2/src/agent2.rs
@@ -3,4 +3,7 @@ mod templates;
 mod thread;
 mod tools;
 
+#[cfg(test)]
+mod tests;
+
 pub use thread::*;
diff --git a/crates/agent2/src/thread/tests.rs b/crates/agent2/src/tests/mod.rs
similarity index 97%
rename from crates/agent2/src/thread/tests.rs
rename to crates/agent2/src/tests/mod.rs
index 9bbf95c6cf..42078a3ffa 100644
--- a/crates/agent2/src/thread/tests.rs
+++ b/crates/agent2/src/tests/mod.rs
@@ -1,5 +1,6 @@
 use super::*;
-use client::{proto::language_server_prompt_request, Client, UserStore};
+use crate::templates::Templates;
+use client::{Client, UserStore};
 use fs::FakeFs;
 use gpui::{AppContext, Entity, TestAppContext};
 use language_model::{
@@ -27,7 +28,7 @@ async fn test_echo(cx: &mut TestAppContext) {
         .await;
     agent.update(cx, |agent, _cx| {
         assert_eq!(
-            agent.messages.last().unwrap().content,
+            agent.messages().last().unwrap().content,
             vec![MessageContent::Text("Hello".to_string())]
         );
     });
@@ -74,7 +75,7 @@ async fn test_basic_tool_calls(cx: &mut TestAppContext) {
     );
     agent.update(cx, |agent, _cx| {
         assert!(agent
-            .messages
+            .messages()
             .last()
             .unwrap()
             .content
@@ -170,7 +171,7 @@ async fn test_concurrent_tool_calls(cx: &mut TestAppContext) {
     }
 
     agent.update(cx, |agent, _cx| {
-        let last_message = agent.messages.last().unwrap();
+        let last_message = agent.messages().last().unwrap();
         let text = last_message
             .content
             .iter()
diff --git a/crates/agent2/src/thread/tests/test_tools.rs b/crates/agent2/src/tests/test_tools.rs
similarity index 97%
rename from crates/agent2/src/thread/tests/test_tools.rs
rename to crates/agent2/src/tests/test_tools.rs
index 38da1bb2f2..43d0414499 100644
--- a/crates/agent2/src/thread/tests/test_tools.rs
+++ b/crates/agent2/src/tests/test_tools.rs
@@ -1,4 +1,6 @@
 use super::*;
+use anyhow::Result;
+use gpui::{App, SharedString, Task};
 
 /// A tool that echoes its input
 #[derive(JsonSchema, Serialize, Deserialize)]