use crate::error_code::invalid_request;use crate::outgoing_message::ConnectionId;use crate::outgoing_message::OutgoingMessageSender;use codex_app_server_protocol::FsChangedNotification;use codex_app_server_protocol::FsUnwatchParams;use codex_app_server_protocol::FsUnwatchResponse;use codex_app_server_protocol::FsWatchParams;use codex_app_server_protocol::FsWatchResponse;use codex_app_server_protocol::JSONRPCErrorError;use codex_app_server_protocol::ServerNotification;use codex_file_watcher::DebouncedWatchReceiver;use codex_file_watcher::FileWatcher;use codex_file_watcher::FileWatcherSubscriber;use codex_file_watcher::WatchPath;use codex_file_watcher::WatchRegistration;use std::collections::HashMap;use std::collections::hash_map::Entry;use std::hash::Hash;use std::sync::Arc;use std::time::Duration;use tokio::sync::Mutex as AsyncMutex;#[cfg(test)]use tokio::sync::mpsc;use tokio::sync::oneshot;use tracing::warn;const FS_CHANGED_NOTIFICATION_DEBOUNCE: Duration = Duration::from_millis(200);#[derive(Clone)]pub(crate) struct FsWatchManager { outgoing: Arc<OutgoingMessageSender>, file_watcher: Arc<FileWatcher>, state: Arc<AsyncMutex<FsWatchState>>,}#[derive(Default)]struct FsWatchState { entries: HashMap<WatchKey, WatchEntry>,}struct WatchEntry { terminate_tx: oneshot::Sender<oneshot::Sender<()>>, _subscriber: FileWatcherSubscriber, _registration: WatchRegistration,}#[derive(Clone, Debug, Eq, Hash, PartialEq)]struct WatchKey { connection_id: ConnectionId, watch_id: String,}impl FsWatchManager { pub(crate) fn new(outgoing: Arc<OutgoingMessageSender>) -> Self { let file_watcher = match FileWatcher::new() { Ok(file_watcher) => Arc::new(file_watcher), Err(err) => { warn!("filesystem watch manager falling back to noop core watcher: {err}"); Arc::new(FileWatcher::noop()) } }; Self::new_with_file_watcher(outgoing, file_watcher) } fn new_with_file_watcher( outgoing: Arc<OutgoingMessageSender>, file_watcher: Arc<FileWatcher>, ) -> Self { Self { outgoing, file_watcher, state: Arc::new(AsyncMutex::new(FsWatchState::default())), } } pub(crate) async fn watch( &self, connection_id: ConnectionId, params: FsWatchParams, ) -> Result<FsWatchResponse, JSONRPCErrorError> { let watch_id = params.watch_id; let watch_key = WatchKey { connection_id, watch_id: watch_id.clone(), }; let outgoing = self.outgoing.clone(); let (subscriber, rx) = self.file_watcher.add_subscriber(); let watch_root = params.path.clone(); let registration = subscriber.register_paths(vec![WatchPath { path: params.path.to_path_buf(), recursive: false, }]); let (terminate_tx, terminate_rx) = oneshot::channel(); match self.state.lock().await.entries.entry(watch_key) { Entry::Occupied(_) => { return Err(invalid_request(format!( "watchId already exists: {watch_id}" ))); } Entry::Vacant(entry) => { entry.insert(WatchEntry { terminate_tx, _subscriber: subscriber, _registration: registration, }); } } let task_watch_id = watch_id.clone(); tokio::spawn(async move { let mut rx = DebouncedWatchReceiver::new(rx, FS_CHANGED_NOTIFICATION_DEBOUNCE); tokio::pin!(terminate_rx); loop { let event = tokio::select! { biased; _ = &mut terminate_rx => break, event = rx.recv() => match event { Some(event) => event, None => break, }, }; let mut changed_paths = event .paths .into_iter() .map(|path| watch_root.join(path)) .collect::<Vec<_>>(); changed_paths.sort_by(|left, right| left.as_path().cmp(right.as_path())); if !changed_paths.is_empty() { outgoing .send_server_notification_to_connection_and_wait( connection_id, ServerNotification::FsChanged(FsChangedNotification { watch_id: task_watch_id.clone(), changed_paths, }), ) .await; } } }); Ok(FsWatchResponse { path: params.path }) } pub(crate) async fn unwatch( &self, connection_id: ConnectionId, params: FsUnwatchParams, ) -> Result<FsUnwatchResponse, JSONRPCErrorError> { let watch_key = WatchKey { connection_id, watch_id: params.watch_id, }; let entry = self.state.lock().await.entries.remove(&watch_key); if let Some(entry) = entry { // Wait for the oneshot to be destroyed by the task to ensure that no notifications // are send after the unwatch response. let (done_tx, done_rx) = oneshot::channel(); let _ = entry.terminate_tx.send(done_tx); let _ = done_rx.await; } Ok(FsUnwatchResponse {}) } pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) { let mut state = self.state.lock().await; state .entries .extract_if(|key, _| key.connection_id == connection_id) .count(); }}#[cfg(test)]mod tests { use super::*; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use std::collections::HashSet; use std::path::PathBuf; use tempfile::TempDir; fn absolute_path(path: PathBuf) -> AbsolutePathBuf { assert!( path.is_absolute(), "path must be absolute: {}", path.display() ); AbsolutePathBuf::try_from(path).expect("path should be absolute") } fn manager_with_noop_watcher() -> FsWatchManager { const OUTGOING_BUFFER: usize = 1; let (tx, _rx) = mpsc::channel(OUTGOING_BUFFER); FsWatchManager::new_with_file_watcher( Arc::new(OutgoingMessageSender::new( tx, codex_analytics::AnalyticsEventsClient::disabled(), )), Arc::new(FileWatcher::noop()), ) } #[tokio::test] async fn watch_uses_client_id_and_tracks_the_owner_scoped_entry() { let temp_dir = TempDir::new().expect("temp dir"); let head_path = temp_dir.path().join("HEAD"); std::fs::write(&head_path, "ref: refs/heads/main\n").expect("write HEAD"); let manager = manager_with_noop_watcher(); let path = absolute_path(head_path); let watch_id = "watch-head".to_string(); let response = manager .watch( ConnectionId(1), FsWatchParams { watch_id: watch_id.clone(), path: path.clone(), }, ) .await .expect("watch should succeed"); assert_eq!(response.path, path); let state = manager.state.lock().await; assert_eq!( state.entries.keys().cloned().collect::<HashSet<_>>(), HashSet::from([WatchKey { connection_id: ConnectionId(1), watch_id, }]) ); } #[tokio::test] async fn unwatch_is_scoped_to_the_connection_that_created_the_watch() { let temp_dir = TempDir::new().expect("temp dir"); let head_path = temp_dir.path().join("HEAD"); std::fs::write(&head_path, "ref: refs/heads/main\n").expect("write HEAD"); let manager = manager_with_noop_watcher(); manager .watch( ConnectionId(1), FsWatchParams { watch_id: "watch-head".to_string(), path: absolute_path(head_path), }, ) .await .expect("watch should succeed"); let watch_key = WatchKey { connection_id: ConnectionId(1), watch_id: "watch-head".to_string(), }; manager .unwatch( ConnectionId(2), FsUnwatchParams { watch_id: "watch-head".to_string(), }, ) .await .expect("foreign unwatch should be a no-op"); assert!(manager.state.lock().await.entries.contains_key(&watch_key)); manager .unwatch( ConnectionId(1), FsUnwatchParams { watch_id: "watch-head".to_string(), }, ) .await .expect("owner unwatch should succeed"); assert!(!manager.state.lock().await.entries.contains_key(&watch_key)); } #[tokio::test] async fn watch_rejects_duplicate_id_for_the_same_connection() { let temp_dir = TempDir::new().expect("temp dir"); let head_path = temp_dir.path().join("HEAD"); let fetch_head_path = temp_dir.path().join("FETCH_HEAD"); std::fs::write(&head_path, "ref: refs/heads/main\n").expect("write HEAD"); std::fs::write(&fetch_head_path, "old-fetch\n").expect("write FETCH_HEAD"); let manager = manager_with_noop_watcher(); manager .watch( ConnectionId(1), FsWatchParams { watch_id: "watch-head".to_string(), path: absolute_path(head_path), }, ) .await .expect("first watch should succeed"); let error = manager .watch( ConnectionId(1), FsWatchParams { watch_id: "watch-head".to_string(), path: absolute_path(fetch_head_path), }, ) .await .expect_err("duplicate watch should fail"); assert_eq!(error.message, "watchId already exists: watch-head"); assert_eq!(manager.state.lock().await.entries.len(), 1); } #[tokio::test] async fn connection_closed_removes_only_that_connections_watches() { let temp_dir = TempDir::new().expect("temp dir"); let head_path = temp_dir.path().join("HEAD"); let fetch_head_path = temp_dir.path().join("FETCH_HEAD"); let packed_refs_path = temp_dir.path().join("packed-refs"); std::fs::write(&head_path, "ref: refs/heads/main\n").expect("write HEAD"); std::fs::write(&fetch_head_path, "old-fetch\n").expect("write FETCH_HEAD"); std::fs::write(&packed_refs_path, "refs\n").expect("write packed-refs"); let manager = manager_with_noop_watcher(); let response = manager .watch( ConnectionId(1), FsWatchParams { watch_id: "watch-head".to_string(), path: absolute_path(head_path.clone()), }, ) .await .expect("first watch should succeed"); manager .watch( ConnectionId(1), FsWatchParams { watch_id: "watch-fetch-head".to_string(), path: absolute_path(fetch_head_path), }, ) .await .expect("second watch should succeed"); manager .watch( ConnectionId(2), FsWatchParams { watch_id: "watch-packed-refs".to_string(), path: absolute_path(packed_refs_path), }, ) .await .expect("third watch should succeed"); manager.connection_closed(ConnectionId(1)).await; assert_eq!( manager .state .lock() .await .entries .keys() .cloned() .collect::<HashSet<_>>(), HashSet::from([WatchKey { connection_id: ConnectionId(2), watch_id: "watch-packed-refs".to_string(), }]) ); assert_eq!(response.path, absolute_path(head_path)); }}