diff --git a/crates/agent/src/agent2.rs b/crates/agent/src/agent2.rs
index 73b61f76fc39d406ae5f5112b1eb903316d40836..41ce661d74d60d82f1b2c411278429d6d30453ae 100644
--- a/crates/agent/src/agent2.rs
+++ b/crates/agent/src/agent2.rs
@@ -87,13 +87,14 @@ pub enum AgentThreadAssistantMessageChunk {
 
 pub struct AgentThreadResponse {
     pub user_message_id: MessageId,
+    pub assistant_message_id: MessageId,
     pub events: BoxStream<'static, Result<AgentThreadResponseEvent>>,
 }
 
 pub trait Agent {
-    fn create_thread();
+    fn create_thread() -> BoxFuture<'static, Result<Arc<dyn AgentThread>>>;
     fn list_threads();
-    fn load_thread();
+    fn load_thread(&self, thread_id: ThreadId) -> BoxFuture<'static, Result<Arc<dyn AgentThread>>>;
 }
 
 pub trait AgentThread {
diff --git a/crates/agent/src/context_store.rs b/crates/agent/src/context_store.rs
index 662a8bb1ef5ad09de7ca2a29d6f5653bec12f23d..668f9a8879ffcee6e2852025e4a2f76b3b88d2cd 100644
--- a/crates/agent/src/context_store.rs
+++ b/crates/agent/src/context_store.rs
@@ -73,7 +73,7 @@ impl ContextStore {
         let existing_context = thread
             .messages()
             .iter()
-            .take_while(|message| exclude_messages_from_id.is_none_or(|id| message.id != id))
+            .take_while(|message| exclude_messages_from_id.is_none_or(|id| message.id() != id))
             .flat_map(|message| {
                 message
                     .loaded_context
diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs
index 296892e43acb44e641ad1e0608252a6437aa9648..0a17a8bc00e608e5543c923d3684fa6d904419bd 100644
--- a/crates/agent/src/thread.rs
+++ b/crates/agent/src/thread.rs
@@ -62,28 +62,29 @@ pub enum MessageToolCall {
     },
 }
 
-/// A message in a [`Thread`].
-pub struct Message {
-    pub id: MessageId,
-    pub role: Role,
-    pub segments: Vec<MessageSegment>,
-    pub loaded_context: LoadedContext,
-    pub creases: Vec<MessageCrease>,
-    pub is_hidden: bool, // todo!("do we need this?")
-    pub ui_only: bool,   // todo!("do we need this?")
-}
-
 pub enum Message {
     User {
+        id: MessageId,
         text: String,
         creases: Vec<MessageCrease>,
+        loaded_context: LoadedContext,
     },
     Assistant {
-        segments: Vec<MessageSegment>,
+        id: MessageId,
+        chunks: Vec<AssistantMessageChunk>,
     },
 }
 
-pub enum MessageSegment {
+impl Message {
+    pub fn id(&self) -> MessageId {
+        match self {
+            Message::User { id, .. } => *id,
+            Message::Assistant { id, .. } => *id,
+        }
+    }
+}
+
+pub enum AssistantMessageChunk {
     Text(Entity<Markdown>),
     Thinking(Entity<Markdown>),
     ToolCall(MessageToolCall),
@@ -504,12 +505,12 @@ impl Thread {
         let Some(message_ix) = self
             .messages
             .iter()
-            .rposition(|message| message.id == message_id)
+            .rposition(|message| message.id() == message_id)
         else {
             return;
         };
         for deleted_message in self.messages.drain(message_ix..) {
-            self.checkpoints_by_message.remove(&deleted_message.id);
+            self.checkpoints_by_message.remove(&deleted_message.id());
         }
         cx.notify();
     }
@@ -719,19 +720,29 @@ impl Thread {
         let mut text = String::new();
 
         for message in &self.messages {
-            text.push_str(match message.role {
-                language_model::Role::User => "User:",
-                language_model::Role::Assistant => "Agent:",
-                language_model::Role::System => "System:",
-            });
-            text.push('\n');
-
-            text.push_str("<think>");
-            text.push_str(message.thinking.read(cx).source());
-            text.push_str("</think>");
-            text.push_str(message.text.read(cx).source());
-
-            // todo!('what about tools?');
+            match message {
+                Message::User {
+                    id,
+                    text,
+                    creases,
+                    loaded_context,
+                } => todo!(),
+                Message::Assistant {
+                    id,
+                    chunks: segments,
+                } => todo!(),
+            }
+            // text.push_str(match message.role {
+            //     language_model::Role::User => "User:",
+            //     language_model::Role::Assistant => "Agent:",
+            //     language_model::Role::System => "System:",
+            // });
+            // text.push('\n');
+
+            // text.push_str("<think>");
+            // text.push_str(message.thinking.read(cx).source());
+            // text.push_str("</think>");
+            // text.push_str(message.text.read(cx).source());
 
             text.push('\n');
         }
@@ -757,7 +768,7 @@ impl Thread {
         thread_store: WeakEntity<ThreadStore>,
         cx: &mut Context<Self>,
     ) {
-        let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
+        let Some(last_message_id) = self.messages.last().map(|message| message.id()) else {
             return;
         };