galaxy/app/src/remote_server/diff_state_tracker.rs

//! Server-side diff state management.
//!
//! [`RemoteDiffStateManager`] is an entity that manages per-(repo, mode)
//! `LocalDiffStateModel` instances and tracks which connections are subscribed
//! to each. It owns model creation, event subscriptions, and content reload
//! spawning. `ServerModel` subscribes to its `DiffStateUpdate` events to
//! handle proto conversion and wire delivery.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use itertools::Itertools;
use warp_util::standardized_path::StandardizedPath;
use warpui::r#async::SpawnedFutureHandle;
use warpui::{AppContext, Entity, ModelContext, ModelHandle};

use super::protocol::RequestId;
use super::server_model::ConnectionId;
use crate::code_review::diff_state::{
    BackendOrigin, DiffMetadata, DiffMode, DiffState, DiffStateModelEvent, FileDiffAndContent,
    GitDiffWithBaseContent, LocalDiffStateModel,
};

// ── Key type ────────────────────────────────────────────────────────

/// Composite key: each (repo, mode) gets its own `LocalDiffStateModel`.
#[derive(Hash, Eq, PartialEq, Clone, Debug)]
pub(super) struct DiffModelKey {
    pub repo_path: StandardizedPath,
    pub mode: DiffMode,
}

// ── Pending response tracker ────────────────────────────────────────

/// Tracks a `GetDiffState` request that arrived while the model was still loading.
/// The response is sent once `NewDiffsComputed` fires.
pub(super) struct PendingDiffStateResponse {
    pub request_id: RequestId,
    pub conn_id: ConnectionId,
}

// ── Action / outcome types ─────────────────────────────────────

/// Outcome of [`RemoteDiffStateManager::subscribe`].
#[allow(clippy::large_enum_variant)]
pub(super) enum SubscribeOutcome {
    /// Respond with this snapshot immediately.
    RespondWithSnapshot {
        key: DiffModelKey,
        state: DiffState,
        metadata: Option<DiffMetadata>,
    },
    /// An async operation is in flight (content reload or model loading).
    /// The manager tracks spawned handles internally.
    Async,
}

/// Domain-level dispatch action returned by event processing and content
/// reload completion. Proto conversion is handled by `ServerModel` at
/// dispatch time.
pub(super) enum DiffStateUpdate {
    /// Build and send a snapshot to subscribers. Entries with a `request_id`
    /// receive a `GetDiffStateResponse`; entries without receive a
    /// server-initiated push `DiffStateSnapshot`.
    Snapshot {
        repo_path: String,
        mode: DiffMode,
        state: DiffState,
        metadata: Option<DiffMetadata>,
        diffs: Option<Arc<GitDiffWithBaseContent>>,
        /// Each subscriber is a connection plus an optional request ID.
        /// `Some(request_id)` → pending `GetDiffState` response (sent with request_id).
        /// `None` → already-subscribed connection that receives a server-initiated push.
        subscribers: Vec<(ConnectionId, Option<RequestId>)>,
    },
    /// Build and send a metadata update to all subscribers.
    MetadataUpdate {
        repo_path: StandardizedPath,
        mode: DiffMode,
        metadata: DiffMetadata,
        subscribers: Vec<ConnectionId>,
    },
    /// Build and send a single-file delta to all subscribers.
    FileDelta {
        repo_path: StandardizedPath,
        mode: DiffMode,
        /// Repo-relative path for the changed file.
        path: String,
        diff: Option<Arc<FileDiffAndContent>>,
        metadata: Option<DiffMetadata>,
        subscribers: Vec<ConnectionId>,
    },
}

// ── RemoteDiffStateManager ────────────────────────────────────────

/// Manages the lifecycle of server-side `LocalDiffStateModel` instances and
/// per-connection subscription tracking.
///
/// A model is created when the first `GetDiffState` arrives for a given key
/// and dropped when the last connection unsubscribes (or disconnects).
pub(super) struct RemoteDiffStateManager {
    /// One model per (repo, mode). Mode is immutable — pinned at construction.
    states: HashMap<DiffModelKey, ModelHandle<LocalDiffStateModel>>,
    /// Per-key set of subscribed connections.
    key_to_connections: HashMap<DiffModelKey, HashSet<ConnectionId>>,
    /// Pending `GetDiffState` responses waiting for the model to finish loading.
    pending_responses: HashMap<DiffModelKey, Vec<PendingDiffStateResponse>>,
    /// In-progress content reload handles, keyed by request ID.
    in_progress: HashMap<RequestId, SpawnedFutureHandle>,
}

impl Entity for RemoteDiffStateManager {
    type Event = DiffStateUpdate;
}

impl RemoteDiffStateManager {
    pub fn new() -> Self {
        Self {
            states: HashMap::new(),
            key_to_connections: HashMap::new(),
            pending_responses: HashMap::new(),
            in_progress: HashMap::new(),
        }
    }

    // ── Model CRUD ──────────────────────────────────────────────────

    pub fn get_model(&self, key: &DiffModelKey) -> Option<&ModelHandle<LocalDiffStateModel>> {
        self.states.get(key)
    }

    pub fn insert_model(&mut self, key: DiffModelKey, model: ModelHandle<LocalDiffStateModel>) {
        self.states.insert(key, model);
    }

    pub fn remove_model(&mut self, key: &DiffModelKey) {
        self.states.remove(key);
        self.pending_responses.remove(key);
        self.key_to_connections.remove(key);
    }

    /// Reads the current `DiffState` and cloned `DiffMetadata` from the model
    /// for `key`. Returns `None` when the model is absent.
    pub fn read_state_and_metadata(
        &self,
        key: &DiffModelKey,
        app: &AppContext,
    ) -> Option<(DiffState, Option<DiffMetadata>)> {
        self.states.get(key).map(|model| {
            let m = model.as_ref(app);
            (m.get(), m.metadata().cloned())
        })
    }

    // ── Connection subscription tracking ────────────────────────────

    /// Records that `conn_id` is subscribed to `key`.
    pub fn subscribe_connection(&mut self, key: DiffModelKey, conn_id: ConnectionId) {
        self.key_to_connections
            .entry(key)
            .or_default()
            .insert(conn_id);
    }

    /// Removes `conn_id`'s subscription for `key`.
    /// If the key has zero remaining subscribers the model is dropped inline.
    pub fn unsubscribe_connection(&mut self, key: &DiffModelKey, conn_id: ConnectionId) {
        if let Some(pending) = self.pending_responses.get_mut(key) {
            pending.retain(|p| p.conn_id != conn_id);
        }

        if let Some(connections) = self.key_to_connections.get_mut(key) {
            connections.remove(&conn_id);
            if connections.is_empty() {
                self.remove_model(key);
            }
        }
    }

    /// Removes all subscriptions for a disconnected connection.
    /// Orphaned models (no remaining subscribers) are dropped inline.
    pub fn remove_connection(&mut self, conn_id: ConnectionId) {
        let keys = self
            .key_to_connections
            .iter()
            .filter(|(_, conns)| conns.contains(&conn_id))
            .map(|(key, _)| key.clone())
            .collect_vec();

        for key in keys {
            self.unsubscribe_connection(&key, conn_id);
        }
    }

    /// Returns the connection IDs subscribed to `key`.
    pub fn subscribed_connections(&self, key: &DiffModelKey) -> Vec<ConnectionId> {
        self.key_to_connections
            .get(key)
            .map(|conns| conns.iter().copied().collect())
            .unwrap_or_default()
    }

    // ── Pending response tracking ───────────────────────────────────

    /// Returns `true` if there are pending responses queued for `key`.
    pub fn has_pending_responses(&self, key: &DiffModelKey) -> bool {
        self.pending_responses
            .get(key)
            .is_some_and(|v| !v.is_empty())
    }

    /// Registers a pending `GetDiffState` response to be sent once the model loads.
    pub fn add_pending_response(
        &mut self,
        key: DiffModelKey,
        request_id: RequestId,
        conn_id: ConnectionId,
    ) {
        self.pending_responses
            .entry(key)
            .or_default()
            .push(PendingDiffStateResponse {
                request_id,
                conn_id,
            });
    }

    /// Drains all pending responses for `key`.
    pub fn drain_pending_responses(&mut self, key: &DiffModelKey) -> Vec<PendingDiffStateResponse> {
        self.pending_responses.remove(key).unwrap_or_default()
    }

    // ── High-level operations ────────────────────────────────────

    /// Handles a `GetDiffState` subscription request.
    ///
    /// Subscribes the connection, looks up or creates the model, and returns
    /// an outcome describing the result. When a content reload is needed it
    /// is spawned internally; when a new model is created the event
    /// subscription is wired up internally.
    pub fn subscribe(
        &mut self,
        repo_path: StandardizedPath,
        mode: DiffMode,
        request_id: &RequestId,
        conn_id: ConnectionId,
        ctx: &mut ModelContext<Self>,
    ) -> SubscribeOutcome {
        let key = DiffModelKey { repo_path, mode };
        self.subscribe_connection(key.clone(), conn_id);

        if let Some(model) = self.get_model(&key) {
            let model_ref = model.as_ref(ctx);
            let state = model_ref.get();
            match state {
                DiffState::Loaded => {
                    let already_in_flight = self.has_pending_responses(&key);
                    self.add_pending_response(key.clone(), request_id.clone(), conn_id);
                    if !already_in_flight {
                        self.spawn_content_reload(key, request_id, ctx);
                    }
                    SubscribeOutcome::Async
                }
                DiffState::Error(_) | DiffState::NotInRepository => {
                    SubscribeOutcome::RespondWithSnapshot {
                        key,
                        state,
                        metadata: model_ref.metadata().cloned(),
                    }
                }
                DiffState::Loading | DiffState::Disconnected => {
                    self.add_pending_response(key, request_id.clone(), conn_id);
                    SubscribeOutcome::Async
                }
            }
        } else {
            // Model doesn't exist — create it and wire up event subscription.
            let repo_path_str = key.repo_path.to_string();
            let mode = key.mode.clone();
            let model = ctx.add_model(|ctx| {
                let mut m =
                    LocalDiffStateModel::new(Some(repo_path_str), BackendOrigin::RemoteDaemon, ctx);
                m.set_diff_mode(mode, false, false, ctx);
                m.set_code_review_metadata_refresh_enabled(true, ctx);
                m
            });
            self.insert_model(key.clone(), model.clone());
            self.add_pending_response(key.clone(), request_id.clone(), conn_id);

            let key_for_sub = key;
            ctx.subscribe_to_model(&model, move |me, _, event, ctx| {
                me.handle_model_event(&key_for_sub, event, ctx);
            });

            SubscribeOutcome::Async
        }
    }

    /// Processes a `DiffStateModelEvent`, builds domain-level dispatch
    /// actions, and emits them as entity events for `ServerModel` to handle.
    fn handle_model_event(
        &mut self,
        key: &DiffModelKey,
        event: &DiffStateModelEvent,
        ctx: &mut ModelContext<Self>,
    ) {
        match event {
            DiffStateModelEvent::NewDiffsComputed { diffs, .. } => {
                let Some((state, metadata)) = self.read_state_and_metadata(key, ctx) else {
                    log::warn!("NewDiffsComputed for absent model key={key:?}");
                    return;
                };

                let pending = self.drain_pending_responses(key);
                let responded_conns: HashSet<ConnectionId> =
                    pending.iter().map(|p| p.conn_id).collect();
                let mut subscribers: Vec<(ConnectionId, Option<RequestId>)> = pending
                    .into_iter()
                    .map(|p| (p.conn_id, Some(p.request_id)))
                    .collect();
                subscribers.extend(
                    self.subscribed_connections(key)
                        .into_iter()
                        .filter(|c| !responded_conns.contains(c))
                        .map(|c| (c, None)),
                );

                ctx.emit(DiffStateUpdate::Snapshot {
                    repo_path: key.repo_path.to_string(),
                    mode: key.mode.clone(),
                    state,
                    metadata,
                    diffs: diffs.clone(),
                    subscribers,
                });
            }
            DiffStateModelEvent::MetadataRefreshed(metadata) => {
                ctx.emit(DiffStateUpdate::MetadataUpdate {
                    repo_path: key.repo_path.clone(),
                    mode: key.mode.clone(),
                    metadata: metadata.as_ref().clone(),
                    subscribers: self.subscribed_connections(key),
                });
            }
            DiffStateModelEvent::CurrentBranchChanged => {
                let Some(model) = self.get_model(key) else {
                    return;
                };
                let Some(metadata) = model.as_ref(ctx).metadata() else {
                    return;
                };
                ctx.emit(DiffStateUpdate::MetadataUpdate {
                    repo_path: key.repo_path.clone(),
                    mode: key.mode.clone(),
                    metadata: metadata.clone(),
                    subscribers: self.subscribed_connections(key),
                });
            }
            DiffStateModelEvent::SingleFileUpdated { path, diff } => {
                let metadata = self
                    .get_model(key)
                    .and_then(|m| m.as_ref(ctx).metadata().cloned());
                ctx.emit(DiffStateUpdate::FileDelta {
                    repo_path: key.repo_path.clone(),
                    mode: key.mode.clone(),
                    path: path.clone(),
                    diff: diff.clone(),
                    metadata,
                    subscribers: self.subscribed_connections(key),
                });
            }
            DiffStateModelEvent::ConnectionLost => {
                // Client-only event — should not occur on the server side.
                log::warn!("Unexpected ConnectionLost event on server-side model key={key:?}");
            }
            DiffStateModelEvent::BranchesReceived(_)
            | DiffStateModelEvent::GitOpCompleted(_)
            | DiffStateModelEvent::CommitMessageGenerated(_)
            | DiffStateModelEvent::BranchCommittedFilesReceived(_) => {
                // Client-only events don't go through this tracker.
            }
        }
    }

    /// Reads model state, drains pending responses, and emits a `Snapshot`
    /// dispatch so `ServerModel` can deliver the results to waiting clients.
    fn resolve_pending_responses(
        &mut self,
        key: &DiffModelKey,
        diffs: Option<GitDiffWithBaseContent>,
        ctx: &mut ModelContext<Self>,
    ) {
        let Some((state, metadata)) = self.read_state_and_metadata(key, ctx) else {
            log::warn!("Content reload completed for absent model key={key:?}");
            return;
        };
        let diffs_arc = diffs.map(Arc::new);
        let subscribers = self
            .drain_pending_responses(key)
            .into_iter()
            .map(|p| (p.conn_id, Some(p.request_id)))
            .collect();

        ctx.emit(DiffStateUpdate::Snapshot {
            repo_path: key.repo_path.to_string(),
            mode: key.mode.clone(),
            state,
            metadata,
            diffs: diffs_arc,
            subscribers,
        });
    }

    /// Spawns an async diff reload with `content_at_base` for late-joining subscribers.
    fn spawn_content_reload(
        &mut self,
        key: DiffModelKey,
        request_id: &RequestId,
        ctx: &mut ModelContext<Self>,
    ) {
        let diff_mode = key.mode.clone();
        let repo_path = std::path::PathBuf::from(key.repo_path.as_str());
        let resolve_id = request_id.clone();
        let abort_id = request_id.clone();
        let abort_key = key.clone();
        let handle = ctx.spawn_abortable(
            async move {
                LocalDiffStateModel::load_diffs_with_content_for_mode(diff_mode, repo_path).await
            },
            move |me, diffs, ctx| {
                me.in_progress.remove(&resolve_id);
                me.resolve_pending_responses(&key, diffs, ctx);
            },
            move |me, ctx| {
                log::info!("Request cancelled (request_id={abort_id})");
                me.in_progress.remove(&abort_id);
                // Drain pending responses with current state instead of orphaning them.
                me.resolve_pending_responses(&abort_key, None, ctx);
            },
        );
        self.in_progress.insert(request_id.clone(), handle);
    }

    /// Cancels an in-progress content reload, if one exists for this request.
    /// Returns `true` if a request was found and aborted.
    pub fn abort_request(&mut self, request_id: &RequestId) -> bool {
        if let Some(handle) = self.in_progress.remove(request_id) {
            handle.abort();
            true
        } else {
            false
        }
    }

    /// Removes a specific pending response by request_id across all keys.
    /// Called by `handle_abort` when the client times out a request.
    /// Returns `true` if a pending response was found and removed.
    pub fn abort_pending_response(&mut self, request_id: &RequestId) -> bool {
        for pending in self.pending_responses.values_mut() {
            if let Some(pos) = pending.iter().position(|p| &p.request_id == request_id) {
                pending.remove(pos);
                return true;
            }
        }
        false
    }
}

#[cfg(test)]
#[path = "diff_state_tracker_tests.rs"]
mod tests;