From 05bc6b2abdf6441834b0c7e46ff54284b675e52f Mon Sep 17 00:00:00 2001
From: Michael Sloan <michael@zed.dev>
Date: Wed, 8 Jan 2025 17:25:16 -0700
Subject: [PATCH] assistant2: Split out implementation of `Context::snapshot`
 (#22878)

Release Notes:

- N/A
---
 crates/assistant2/src/context.rs | 101 ++++++++++++++++++-------------
 1 file changed, 59 insertions(+), 42 deletions(-)

diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs
index 55a0b71ba70fb..ab0def6bc4876 100644
--- a/crates/assistant2/src/context.rs
+++ b/crates/assistant2/src/context.rs
@@ -104,48 +104,10 @@ pub struct ThreadContext {
 impl Context {
     pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
         match &self {
-            Self::File(file_context) => {
-                let path = file_context.path(cx)?;
-                let full_path: SharedString = path.to_string_lossy().into_owned().into();
-                let name = match path.file_name() {
-                    Some(name) => name.to_string_lossy().into_owned().into(),
-                    None => full_path.clone(),
-                };
-                let parent = path
-                    .parent()
-                    .and_then(|p| p.file_name())
-                    .map(|p| p.to_string_lossy().into_owned().into());
-
-                Some(ContextSnapshot {
-                    id: self.id(),
-                    name,
-                    parent,
-                    tooltip: Some(full_path),
-                    kind: ContextKind::File,
-                    text: file_context.text.clone(),
-                })
-            }
-            Self::Directory(DirectoryContext { snapshot, .. }) => Some(snapshot.clone()),
-            Self::FetchedUrl(FetchedUrlContext { url, text, id }) => Some(ContextSnapshot {
-                id: *id,
-                name: url.clone(),
-                parent: None,
-                tooltip: None,
-                kind: ContextKind::FetchedUrl,
-                text: text.clone(),
-            }),
-            Self::Thread(thread_context) => {
-                let thread = thread_context.thread.read(cx);
-
-                Some(ContextSnapshot {
-                    id: self.id(),
-                    name: thread.summary().unwrap_or("New thread".into()),
-                    parent: None,
-                    tooltip: None,
-                    kind: ContextKind::Thread,
-                    text: thread_context.text.clone(),
-                })
-            }
+            Self::File(file_context) => file_context.snapshot(cx),
+            Self::Directory(directory_context) => Some(directory_context.snapshot()),
+            Self::FetchedUrl(fetched_url_context) => Some(fetched_url_context.snapshot()),
+            Self::Thread(thread_context) => Some(thread_context.snapshot(cx)),
         }
     }
 }
@@ -160,6 +122,61 @@ impl FileContext {
             None
         }
     }
+
+    pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
+        let path = self.path(cx)?;
+        let full_path: SharedString = path.to_string_lossy().into_owned().into();
+        let name = match path.file_name() {
+            Some(name) => name.to_string_lossy().into_owned().into(),
+            None => full_path.clone(),
+        };
+        let parent = path
+            .parent()
+            .and_then(|p| p.file_name())
+            .map(|p| p.to_string_lossy().into_owned().into());
+
+        Some(ContextSnapshot {
+            id: self.id,
+            name,
+            parent,
+            tooltip: Some(full_path),
+            kind: ContextKind::File,
+            text: self.text.clone(),
+        })
+    }
+}
+
+impl DirectoryContext {
+    pub fn snapshot(&self) -> ContextSnapshot {
+        self.snapshot.clone()
+    }
+}
+
+impl FetchedUrlContext {
+    pub fn snapshot(&self) -> ContextSnapshot {
+        ContextSnapshot {
+            id: self.id,
+            name: self.url.clone(),
+            parent: None,
+            tooltip: None,
+            kind: ContextKind::FetchedUrl,
+            text: self.text.clone(),
+        }
+    }
+}
+
+impl ThreadContext {
+    pub fn snapshot(&self, cx: &AppContext) -> ContextSnapshot {
+        let thread = self.thread.read(cx);
+        ContextSnapshot {
+            id: self.id,
+            name: thread.summary().unwrap_or("New thread".into()),
+            parent: None,
+            tooltip: None,
+            kind: ContextKind::Thread,
+            text: self.text.clone(),
+        }
+    }
 }
 
 pub fn attach_context_to_message(