dk_engine/workspace/

session_workspace.rs

//! SessionWorkspace — the isolated workspace for a single agent session.
//!
//! Each workspace owns a [`FileOverlay`] and a [`SessionGraph`], pinned to a
//! `base_commit` in the repository. Reads go through the overlay first, then
//! fall back to the Git tree at the base commit.

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use dk_core::{AgentId, RepoId, Result};
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use std::collections::HashSet;
use std::sync::Arc;
use tokio::time::Instant;
use uuid::Uuid;

use crate::git::GitRepository;
use crate::workspace::overlay::{FileOverlay, OverlayEntry};
use crate::workspace::session_graph::SessionGraph;

// ── Type aliases ─────────────────────────────────────────────────────

pub type WorkspaceId = Uuid;
pub type SessionId = Uuid;

// ── Workspace mode ───────────────────────────────────────────────────

/// Controls the lifetime semantics of a workspace.
#[derive(Debug, Clone)]
pub enum WorkspaceMode {
    /// Destroyed when the session disconnects.
    Ephemeral,
    /// Survives disconnection; optionally expires at a deadline.
    Persistent { expires_at: Option<Instant> },
}

impl WorkspaceMode {
    /// SQL label for the DB column.
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Ephemeral => "ephemeral",
            Self::Persistent { .. } => "persistent",
        }
    }
}

// ── Workspace state machine ──────────────────────────────────────────

/// Lifecycle state of a workspace.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WorkspaceState {
    Active,
    Submitted,
    Merged,
    Expired,
    Abandoned,
}

impl WorkspaceState {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Active => "active",
            Self::Submitted => "submitted",
            Self::Merged => "merged",
            Self::Expired => "expired",
            Self::Abandoned => "abandoned",
        }
    }
}

// ── File read result ─────────────────────────────────────────────────

/// Result of reading a file through the workspace layer.
#[derive(Debug, Clone)]
pub struct FileReadResult {
    pub content: Vec<u8>,
    pub hash: String,
    pub modified_in_session: bool,
}

// ── SessionWorkspace ─────────────────────────────────────────────────

/// An isolated workspace for a single agent session.
///
/// Reads resolve overlay-first, then fall through to the Git tree at
/// `base_commit`. Writes go exclusively to the overlay.
pub struct SessionWorkspace {
    pub id: WorkspaceId,
    pub session_id: SessionId,
    pub repo_id: RepoId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub changeset_id: uuid::Uuid,
    pub intent: String,
    pub base_commit: String,
    pub overlay: FileOverlay,
    pub graph: SessionGraph,
    pub mode: WorkspaceMode,
    pub state: WorkspaceState,
    pub created_at: Instant,
    pub last_active: Instant,
    /// Per-path wall-clock timestamp of the most recent `dk_file_read` in
    /// this session. Consumed by the STALE_OVERLAY pre-write check so a
    /// session whose local view predates a competing submitted changeset is
    /// forced to re-read before writing. Interior-mutable so it can be
    /// updated through `&SessionWorkspace` (mirrors how `overlay` manages
    /// its own locking).
    pub files_read: Arc<DashMap<String, DateTime<Utc>>>,
}

impl SessionWorkspace {
    /// Create a workspace without any database interaction (test-only).
    ///
    /// Uses [`FileOverlay::new_inmemory`] so writes go only to the
    /// in-memory DashMap. Suitable for unit / integration tests that
    /// verify isolation semantics without requiring PostgreSQL.
    #[doc(hidden)]
    pub fn new_test(
        session_id: SessionId,
        repo_id: RepoId,
        agent_id: AgentId,
        intent: String,
        base_commit: String,
        mode: WorkspaceMode,
    ) -> Self {
        let id = Uuid::new_v4();
        let now = Instant::now();
        let overlay = FileOverlay::new_inmemory(id);
        let graph = SessionGraph::empty();

        Self {
            id,
            session_id,
            repo_id,
            agent_id,
            agent_name: String::new(),
            changeset_id: Uuid::new_v4(),
            intent,
            base_commit,
            overlay,
            graph,
            mode,
            state: WorkspaceState::Active,
            created_at: now,
            last_active: now,
            files_read: Arc::new(DashMap::new()),
        }
    }

    /// Rehydrate a workspace from existing database state without inserting a new row.
    ///
    /// Used by [`WorkspaceManager::resume`] to reconstruct an in-memory
    /// `SessionWorkspace` after the DB row has already been updated (session_id
    /// rotated, stranded_at cleared). Unlike [`SessionWorkspace::new`], this
    /// constructor does **not** insert a new `session_workspaces` row — it only
    /// wires up the in-memory structures pointing at the existing `workspace_id`.
    #[allow(clippy::too_many_arguments)]
    pub fn rehydrate(
        workspace_id: WorkspaceId,
        session_id: SessionId,
        repo_id: RepoId,
        agent_id: AgentId,
        changeset_id: Uuid,
        intent: String,
        base_commit: String,
        mode: WorkspaceMode,
        agent_name: String,
        db: PgPool,
    ) -> Self {
        let now = Instant::now();
        let overlay = FileOverlay::new(workspace_id, db);
        let graph = SessionGraph::empty();

        Self {
            id: workspace_id,
            session_id,
            repo_id,
            agent_id,
            agent_name,
            changeset_id,
            intent,
            base_commit,
            overlay,
            graph,
            mode,
            state: WorkspaceState::Active,
            created_at: now,
            last_active: now,
            files_read: Arc::new(DashMap::new()),
        }
    }

    /// Create a new workspace and persist metadata to the database.
    #[allow(clippy::too_many_arguments)]
    pub async fn new(
        session_id: SessionId,
        repo_id: RepoId,
        agent_id: AgentId,
        changeset_id: Uuid,
        intent: String,
        base_commit: String,
        mode: WorkspaceMode,
        agent_name: String,
        db: PgPool,
    ) -> Result<Self> {
        let id = Uuid::new_v4();
        let now = Instant::now();

        // Persist to DB
        sqlx::query(
            r#"
            INSERT INTO session_workspaces
                (id, session_id, repo_id, base_commit_hash, state, mode, agent_id, intent, agent_name, changeset_id)
            VALUES ($1, $2, $3, $4, 'active', $5, $6, $7, $8, $9)
            "#,
        )
        .bind(id)
        .bind(session_id)
        .bind(repo_id)
        .bind(&base_commit)
        .bind(mode.as_str())
        .bind(&agent_id)
        .bind(&intent)
        .bind(&agent_name)
        .bind(changeset_id)
        .execute(&db)
        .await?;

        let overlay = FileOverlay::new(id, db);
        let graph = SessionGraph::empty();

        Ok(Self {
            id,
            session_id,
            repo_id,
            agent_id,
            agent_name,
            changeset_id,
            intent,
            base_commit,
            overlay,
            graph,
            mode,
            state: WorkspaceState::Active,
            created_at: now,
            last_active: now,
            files_read: Arc::new(DashMap::new()),
        })
    }

    /// Read a file through the overlay-first layer.
    ///
    /// 1. If the overlay has a `Modified` or `Added` entry, return that content.
    /// 2. If the overlay has a `Deleted` entry, return a "not found" error.
    /// 3. Otherwise, read from the Git tree at `base_commit`.
    pub fn read_file(&self, path: &str, git_repo: &GitRepository) -> Result<FileReadResult> {
        if let Some(entry) = self.overlay.get(path) {
            return match entry.value() {
                OverlayEntry::Modified { content, hash } | OverlayEntry::Added { content, hash } => {
                    Ok(FileReadResult {
                        content: content.clone(),
                        hash: hash.clone(),
                        modified_in_session: true,
                    })
                }
                OverlayEntry::Deleted => Err(dk_core::Error::Git(format!(
                    "file '{path}' has been deleted in this session"
                ))),
            };
        }

        // Fall through to base tree.
        // TODO(perf): The git tree entry already stores a content-addressable
        // OID (blob hash). If GitRepository exposed the entry OID we could use
        // it directly instead of recomputing SHA-256 on every base-tree read.
        let content = git_repo.read_tree_entry(&self.base_commit, path)?;
        let hash = format!("{:x}", Sha256::digest(&content));

        Ok(FileReadResult {
            content,
            hash,
            modified_in_session: false,
        })
    }

    /// Write a file through the overlay.
    ///
    /// Determines whether the file is new (not in base tree) or modified.
    pub async fn write_file(
        &self,
        path: &str,
        content: Vec<u8>,
        git_repo: &GitRepository,
    ) -> Result<String> {
        let is_new = git_repo.read_tree_entry(&self.base_commit, path).is_err();
        self.overlay.write(path, content, is_new).await
    }

    /// Delete a file in the overlay.
    pub async fn delete_file(&self, path: &str) -> Result<()> {
        self.overlay.delete(path).await
    }

    /// List files visible in this workspace.
    ///
    /// If `only_modified` is true, return only overlay entries.
    /// Otherwise, return the full base tree merged with overlay changes.
    ///
    /// When `prefix` is `Some`, only paths starting with the given prefix
    /// are included. The filter is applied early in the pipeline so that
    /// building the `HashSet` only contains relevant entries rather than
    /// the entire tree (which can be 100k+ files in large repos).
    pub fn list_files(
        &self,
        git_repo: &GitRepository,
        only_modified: bool,
        prefix: Option<&str>,
    ) -> Result<Vec<String>> {
        let matches_prefix = |p: &str| -> bool {
            match prefix {
                Some(pfx) => p.starts_with(pfx),
                None => true,
            }
        };

        if only_modified {
            return Ok(self
                .overlay
                .list_changes()
                .into_iter()
                .filter(|(path, _)| matches_prefix(path))
                .map(|(path, _)| path)
                .collect());
        }

        // Start with base tree — filter by prefix before collecting into
        // the HashSet to avoid allocating entries we will immediately discard.
        let base_files = git_repo.list_tree_files(&self.base_commit)?;
        let mut result: HashSet<String> = base_files
            .into_iter()
            .filter(|p| matches_prefix(p))
            .collect();

        // Apply overlay (only entries matching the prefix)
        for (path, entry) in self.overlay.list_changes() {
            if !matches_prefix(&path) {
                continue;
            }
            match entry {
                OverlayEntry::Added { .. } | OverlayEntry::Modified { .. } => {
                    result.insert(path);
                }
                OverlayEntry::Deleted => {
                    result.remove(&path);
                }
            }
        }

        let mut files: Vec<String> = result.into_iter().collect();
        files.sort();
        Ok(files)
    }

    /// Touch the workspace to update last-active timestamp.
    pub fn touch(&mut self) {
        self.last_active = Instant::now();
    }

    /// Record that this session just read `path` at the current wall-clock
    /// time. Called from `handle_file_read` to feed the STALE_OVERLAY
    /// pre-write check: a session whose per-path timestamp predates a
    /// competing submitted changeset touching the same path is forced to
    /// re-read before the next `dk_file_write` is accepted.
    pub fn mark_read(&self, path: &str) {
        self.files_read.insert(path.to_string(), Utc::now());
    }

    /// Return the wall-clock timestamp of the most recent `dk_file_read`
    /// for `path`, or `None` if this session has never read it.
    pub fn last_read(&self, path: &str) -> Option<DateTime<Utc>> {
        self.files_read.get(path).map(|e| *e.value())
    }

    /// Re-parse overlay file contents and rebuild the semantic graph.
    ///
    /// Called by [`WorkspaceManager::resume`] after the overlay is restored from
    /// the database. Walks every entry in the overlay, parses supported files,
    /// and updates the session graph delta accordingly:
    ///
    /// - **Deleted** entries: all session-owned symbols for that file are
    ///   removed from the delta by iterating `added_symbols` and
    ///   `modified_symbols` and dropping matches.
    /// - **Added / Modified** entries: the content is re-parsed by the
    ///   [`ParserRegistry`]. The parse result is fed into
    ///   [`SessionGraph::update_from_parse`] with an empty base (all symbols
    ///   in overlay files are session additions — there is no base-symbol set
    ///   available at this point). Unsupported file extensions are silently
    ///   skipped.
    pub async fn reindex_from_overlay(&mut self) -> dk_core::Result<()> {
        use crate::parser::ParserRegistry;
        use crate::workspace::overlay::OverlayEntry;
        use std::path::Path;

        let registry = ParserRegistry::new();
        let changes = self.overlay.list_changes();

        for (path_str, entry) in changes {
            let file_path = Path::new(&path_str);
            match entry {
                OverlayEntry::Deleted => {
                    // Remove any session-owned symbols for this file.
                    self.graph.remove_session_symbols_for_file(&path_str);
                }
                OverlayEntry::Added { content, .. } | OverlayEntry::Modified { content, .. } => {
                    if !registry.supports_file(file_path) {
                        continue;
                    }
                    let text = std::str::from_utf8(&content).map_err(|e| {
                        dk_core::Error::Internal(format!(
                            "reindex_from_overlay: non-utf8 in {path_str}: {e}"
                        ))
                    })?;
                    let analysis = match registry.parse_file(file_path, text.as_bytes()) {
                        Ok(a) => a,
                        Err(e) => {
                            tracing::warn!(
                                path = %path_str,
                                "reindex_from_overlay: parse failed, skipping: {e}"
                            );
                            continue;
                        }
                    };
                    // All overlay symbols are session additions — pass an empty
                    // base so update_from_parse classifies all new symbols as
                    // added and removes none.
                    self.graph
                        .update_from_parse(&path_str, analysis.symbols, &[]);
                }
            }
        }

        Ok(())
    }

    /// Build the overlay vector for `commit_tree_overlay`.
    ///
    /// Returns `(path, Some(content))` for modified/added files and
    /// `(path, None)` for deleted files.
    pub fn overlay_for_tree(&self) -> Vec<(String, Option<Vec<u8>>)> {
        self.overlay
            .list_changes()
            .into_iter()
            .map(|(path, entry)| {
                let data = match entry {
                    OverlayEntry::Modified { content, .. }
                    | OverlayEntry::Added { content, .. } => Some(content),
                    OverlayEntry::Deleted => None,
                };
                (path, data)
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use uuid::Uuid;

    fn make_test_workspace() -> SessionWorkspace {
        SessionWorkspace::new_test(
            Uuid::new_v4(),
            Uuid::new_v4(),
            "test-agent".to_string(),
            "test intent".to_string(),
            "abc123".to_string(),
            WorkspaceMode::Ephemeral,
        )
    }

    #[tokio::test]
    async fn reindex_from_overlay_adds_symbols_for_rust_file() {
        let mut ws = make_test_workspace();
        // Write a Rust file with a single function into the overlay.
        ws.overlay
            .write_local("x.rs", b"pub fn hello() {}".to_vec(), true);

        ws.reindex_from_overlay().await.unwrap();

        // The graph should now contain "hello" (added symbol).
        let symbols = ws.graph.changed_symbols_for_file("x.rs");
        assert!(
            symbols.iter().any(|s| s == "hello"),
            "expected 'hello' in graph symbols, got: {symbols:?}"
        );
    }

    #[tokio::test]
    async fn reindex_from_overlay_skips_unsupported_extensions() {
        let mut ws = make_test_workspace();
        ws.overlay
            .write_local("readme.txt", b"just text".to_vec(), true);

        // Should not error — unsupported extension is silently skipped.
        ws.reindex_from_overlay().await.unwrap();
        assert_eq!(ws.graph.change_count(), 0);
    }

    #[tokio::test]
    async fn reindex_from_overlay_deleted_entry_clears_symbols() {
        use dk_core::{Span, Symbol, SymbolKind, Visibility};
        use std::path::PathBuf;

        let mut ws = make_test_workspace();
        // Pre-populate the graph with a symbol for the file.
        ws.graph.add_symbol(Symbol {
            id: Uuid::new_v4(),
            name: "old_fn".to_string(),
            qualified_name: "old_fn".to_string(),
            kind: SymbolKind::Function,
            visibility: Visibility::Public,
            file_path: PathBuf::from("gone.rs"),
            span: Span { start_byte: 0, end_byte: 10 },
            signature: None,
            doc_comment: None,
            parent: None,
            last_modified_by: None,
            last_modified_intent: None,
        });
        assert_eq!(ws.graph.change_count(), 1);

        // Mark file as deleted in the overlay.
        ws.overlay.delete_local("gone.rs");

        ws.reindex_from_overlay().await.unwrap();

        // Symbol should have been removed.
        let symbols = ws.graph.changed_symbols_for_file("gone.rs");
        assert!(symbols.is_empty(), "deleted file should have no graph symbols");
    }
}