semantic_code_edit_mcp/
state.rs

1use std::num::NonZeroUsize;
2use std::path::PathBuf;
3use std::sync::{Arc, Mutex};
4
5use anyhow::{Result, anyhow};
6use fieldwork::Fieldwork;
7use lru::LruCache;
8use serde::{Deserialize, Serialize};
9
10use crate::editor::EditPosition;
11use crate::languages::{LanguageName, LanguageRegistry};
12use crate::selector::Selector;
13use mcplease::session::SessionStore;
14
15// Explanation for the presence of session_id that is currently unused: The intent was initially to
16// have a conversation-unique identifier of some sort in order to isolate state between
17// conversations. However, MCP provides no mechanism to distinguish between conversations, so I
18// tried adding a session_id that was provided to every tool call in order to isolate state. This
19// presents a usability concern, so I've decided to just be extra careful about switching contexts
20// until we have a better solution. I still hope to iterate towards isolated sessions, so the code
21// is still written to support that.
22
23/// Session data specific to semantic editing operations
24#[derive(Debug, Clone, Serialize, Deserialize, Default)]
25pub struct SemanticEditSessionData {
26    /// Current working context path
27    pub context_path: Option<PathBuf>,
28    /// Currently staged operation
29    pub staged_operation: Option<StagedOperation>,
30}
31
32/// Represents a staged operation that can be previewed and committed
33#[derive(Debug, Clone, Fieldwork, Serialize, Deserialize)]
34#[fieldwork(get, set, get_mut, with)]
35pub struct StagedOperation {
36    pub selector: Selector,
37    pub content: String,
38    pub file_path: PathBuf,
39    pub language_name: LanguageName,
40    pub edit_position: Option<EditPosition>,
41}
42
43impl StagedOperation {
44    pub fn retarget(&mut self, selector: Selector) {
45        self.selector = selector;
46    }
47}
48
49/// Semantic editing tools with session support
50#[derive(fieldwork::Fieldwork)]
51#[fieldwork(get)]
52pub struct SemanticEditTools {
53    #[fieldwork(get_mut)]
54    session_store: SessionStore<SemanticEditSessionData>,
55    language_registry: Arc<LanguageRegistry>,
56    file_cache: Arc<Mutex<LruCache<String, String>>>,
57    #[fieldwork(set, get_mut, option = false)]
58    commit_fn: Option<Box<(dyn Fn(PathBuf, String) + 'static)>>,
59    #[fieldwork(set, with)]
60    default_session_id: &'static str,
61}
62
63impl std::fmt::Debug for SemanticEditTools {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("SemanticEditTools")
66            .field("session_store", &self.session_store)
67            .field("language_registry", &self.language_registry)
68            .field("file_cache", &self.file_cache)
69            .field("default_session_id", &self.default_session_id)
70            .finish()
71    }
72}
73
74impl SemanticEditTools {
75    /// Create a new SemanticEditTools instance
76    pub fn new(storage_path: Option<&str>) -> Result<Self> {
77        let storage_path = storage_path.map(|s| PathBuf::from(&*shellexpand::tilde(s)));
78        let session_store = SessionStore::new(storage_path)?;
79        let language_registry = Arc::new(LanguageRegistry::new()?);
80        let file_cache = Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(50).unwrap())));
81
82        Ok(Self {
83            session_store,
84            language_registry,
85            file_cache,
86            commit_fn: None,
87            default_session_id: "default",
88        })
89    }
90
91    /// Get context for a session
92    pub fn get_context(&self, session_id: Option<&str>) -> Result<Option<PathBuf>> {
93        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
94        let session_data = self.session_store.get_or_create(session_id)?;
95        Ok(session_data.context_path)
96    }
97
98    /// Stage a new operation, replacing any existing staged operation
99    pub fn stage_operation(
100        &self,
101        session_id: Option<&str>,
102        staged_operation: Option<StagedOperation>,
103    ) -> Result<()> {
104        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
105        self.session_store.update(session_id, |data| {
106            data.staged_operation = staged_operation;
107        })
108    }
109
110    /// Get the currently staged operation, if any
111    pub fn get_staged_operation(
112        &self,
113        session_id: Option<&str>,
114    ) -> Result<Option<StagedOperation>> {
115        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
116        let session_data = self.session_store.get_or_create(session_id)?;
117        Ok(session_data.staged_operation)
118    }
119
120    /// Take the staged operation, removing it from storage
121    pub fn take_staged_operation(
122        &self,
123        session_id: Option<&str>,
124    ) -> Result<Option<StagedOperation>> {
125        let mut staged_op = None;
126        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
127        self.session_store.update(session_id, |data| {
128            staged_op = data.staged_operation.take();
129        })?;
130        Ok(staged_op)
131    }
132
133    /// Modify the staged operation in place
134    pub fn modify_staged_operation<F>(
135        &self,
136        session_id: Option<&str>,
137        fun: F,
138    ) -> Result<Option<StagedOperation>>
139    where
140        F: FnOnce(&mut StagedOperation),
141    {
142        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
143        self.session_store.update(session_id, |data| {
144            if let Some(ref mut op) = data.staged_operation {
145                fun(op);
146            }
147        })?;
148        self.get_staged_operation(Some(session_id))
149    }
150
151    /// Set context path for a session
152    pub fn set_context(&self, session_id: Option<&str>, path: PathBuf) -> Result<()> {
153        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
154
155        self.session_store.update(session_id, |data| {
156            data.context_path = Some(path);
157        })
158    }
159
160    /// Resolve a path relative to session context if needed
161    pub(crate) fn resolve_path(&self, path_str: &str, session_id: Option<&str>) -> Result<PathBuf> {
162        let path = PathBuf::from(&*shellexpand::tilde(path_str));
163
164        if path.is_absolute() {
165            return Ok(std::fs::canonicalize(path)?);
166        }
167
168        let session_id = session_id.unwrap_or_else(|| self.default_session_id());
169
170        match self.get_context(Some(session_id))? {
171            Some(context) => Ok(std::fs::canonicalize(context.join(path_str))?),
172            None => Err(anyhow!(
173                "No context found for `{session_id}`. Use set_context first or provide an absolute path.",
174            )),
175        }
176    }
177}