dk_engine/workspace/
session_workspace.rs

1//! SessionWorkspace — the isolated workspace for a single agent session.
2//!
3//! Each workspace owns a [`FileOverlay`] and a [`SessionGraph`], pinned to a
4//! `base_commit` in the repository. Reads go through the overlay first, then
5//! fall back to the Git tree at the base commit.
6
7use dk_core::{AgentId, RepoId, Result};
8use sha2::{Digest, Sha256};
9use sqlx::PgPool;
10use std::collections::HashSet;
11use tokio::time::Instant;
12use uuid::Uuid;
13
14use crate::git::GitRepository;
15use crate::workspace::overlay::{FileOverlay, OverlayEntry};
16use crate::workspace::session_graph::SessionGraph;
17
18// ── Type aliases ─────────────────────────────────────────────────────
19
20pub type WorkspaceId = Uuid;
21pub type SessionId = Uuid;
22
23// ── Workspace mode ───────────────────────────────────────────────────
24
25/// Controls the lifetime semantics of a workspace.
26#[derive(Debug, Clone)]
27pub enum WorkspaceMode {
28    /// Destroyed when the session disconnects.
29    Ephemeral,
30    /// Survives disconnection; optionally expires at a deadline.
31    Persistent { expires_at: Option<Instant> },
32}
33
34impl WorkspaceMode {
35    /// SQL label for the DB column.
36    pub fn as_str(&self) -> &'static str {
37        match self {
38            Self::Ephemeral => "ephemeral",
39            Self::Persistent { .. } => "persistent",
40        }
41    }
42}
43
44// ── Workspace state machine ──────────────────────────────────────────
45
46/// Lifecycle state of a workspace.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum WorkspaceState {
49    Active,
50    Submitted,
51    Merged,
52    Expired,
53    Abandoned,
54}
55
56impl WorkspaceState {
57    pub fn as_str(&self) -> &'static str {
58        match self {
59            Self::Active => "active",
60            Self::Submitted => "submitted",
61            Self::Merged => "merged",
62            Self::Expired => "expired",
63            Self::Abandoned => "abandoned",
64        }
65    }
66}
67
68// ── File read result ─────────────────────────────────────────────────
69
70/// Result of reading a file through the workspace layer.
71#[derive(Debug, Clone)]
72pub struct FileReadResult {
73    pub content: Vec<u8>,
74    pub hash: String,
75    pub modified_in_session: bool,
76}
77
78// ── SessionWorkspace ─────────────────────────────────────────────────
79
80/// An isolated workspace for a single agent session.
81///
82/// Reads resolve overlay-first, then fall through to the Git tree at
83/// `base_commit`. Writes go exclusively to the overlay.
84pub struct SessionWorkspace {
85    pub id: WorkspaceId,
86    pub session_id: SessionId,
87    pub repo_id: RepoId,
88    pub agent_id: AgentId,
89    pub agent_name: String,
90    pub changeset_id: uuid::Uuid,
91    pub intent: String,
92    pub base_commit: String,
93    pub overlay: FileOverlay,
94    pub graph: SessionGraph,
95    pub mode: WorkspaceMode,
96    pub state: WorkspaceState,
97    pub created_at: Instant,
98    pub last_active: Instant,
99}
100
101impl SessionWorkspace {
102    /// Create a workspace without any database interaction (test-only).
103    ///
104    /// Uses [`FileOverlay::new_inmemory`] so writes go only to the
105    /// in-memory DashMap. Suitable for unit / integration tests that
106    /// verify isolation semantics without requiring PostgreSQL.
107    #[doc(hidden)]
108    pub fn new_test(
109        session_id: SessionId,
110        repo_id: RepoId,
111        agent_id: AgentId,
112        intent: String,
113        base_commit: String,
114        mode: WorkspaceMode,
115    ) -> Self {
116        let id = Uuid::new_v4();
117        let now = Instant::now();
118        let overlay = FileOverlay::new_inmemory(id);
119        let graph = SessionGraph::empty();
120
121        Self {
122            id,
123            session_id,
124            repo_id,
125            agent_id,
126            agent_name: String::new(),
127            changeset_id: Uuid::new_v4(),
128            intent,
129            base_commit,
130            overlay,
131            graph,
132            mode,
133            state: WorkspaceState::Active,
134            created_at: now,
135            last_active: now,
136        }
137    }
138
139    /// Create a new workspace and persist metadata to the database.
140    #[allow(clippy::too_many_arguments)]
141    pub async fn new(
142        session_id: SessionId,
143        repo_id: RepoId,
144        agent_id: AgentId,
145        changeset_id: Uuid,
146        intent: String,
147        base_commit: String,
148        mode: WorkspaceMode,
149        agent_name: String,
150        db: PgPool,
151    ) -> Result<Self> {
152        let id = Uuid::new_v4();
153        let now = Instant::now();
154
155        // Persist to DB
156        sqlx::query(
157            r#"
158            INSERT INTO session_workspaces
159                (id, session_id, repo_id, base_commit_hash, state, mode, agent_id, intent, agent_name)
160            VALUES ($1, $2, $3, $4, 'active', $5, $6, $7, $8)
161            "#,
162        )
163        .bind(id)
164        .bind(session_id)
165        .bind(repo_id)
166        .bind(&base_commit)
167        .bind(mode.as_str())
168        .bind(&agent_id)
169        .bind(&intent)
170        .bind(&agent_name)
171        .execute(&db)
172        .await?;
173
174        let overlay = FileOverlay::new(id, db);
175        let graph = SessionGraph::empty();
176
177        Ok(Self {
178            id,
179            session_id,
180            repo_id,
181            agent_id,
182            agent_name,
183            changeset_id,
184            intent,
185            base_commit,
186            overlay,
187            graph,
188            mode,
189            state: WorkspaceState::Active,
190            created_at: now,
191            last_active: now,
192        })
193    }
194
195    /// Read a file through the overlay-first layer.
196    ///
197    /// 1. If the overlay has a `Modified` or `Added` entry, return that content.
198    /// 2. If the overlay has a `Deleted` entry, return a "not found" error.
199    /// 3. Otherwise, read from the Git tree at `base_commit`.
200    pub fn read_file(&self, path: &str, git_repo: &GitRepository) -> Result<FileReadResult> {
201        if let Some(entry) = self.overlay.get(path) {
202            return match entry.value() {
203                OverlayEntry::Modified { content, hash } | OverlayEntry::Added { content, hash } => {
204                    Ok(FileReadResult {
205                        content: content.clone(),
206                        hash: hash.clone(),
207                        modified_in_session: true,
208                    })
209                }
210                OverlayEntry::Deleted => Err(dk_core::Error::Git(format!(
211                    "file '{path}' has been deleted in this session"
212                ))),
213            };
214        }
215
216        // Fall through to base tree.
217        // TODO(perf): The git tree entry already stores a content-addressable
218        // OID (blob hash). If GitRepository exposed the entry OID we could use
219        // it directly instead of recomputing SHA-256 on every base-tree read.
220        let content = git_repo.read_tree_entry(&self.base_commit, path)?;
221        let hash = format!("{:x}", Sha256::digest(&content));
222
223        Ok(FileReadResult {
224            content,
225            hash,
226            modified_in_session: false,
227        })
228    }
229
230    /// Write a file through the overlay.
231    ///
232    /// Determines whether the file is new (not in base tree) or modified.
233    pub async fn write_file(
234        &self,
235        path: &str,
236        content: Vec<u8>,
237        git_repo: &GitRepository,
238    ) -> Result<String> {
239        let is_new = git_repo.read_tree_entry(&self.base_commit, path).is_err();
240        self.overlay.write(path, content, is_new).await
241    }
242
243    /// Delete a file in the overlay.
244    pub async fn delete_file(&self, path: &str) -> Result<()> {
245        self.overlay.delete(path).await
246    }
247
248    /// List files visible in this workspace.
249    ///
250    /// If `only_modified` is true, return only overlay entries.
251    /// Otherwise, return the full base tree merged with overlay changes.
252    ///
253    /// When `prefix` is `Some`, only paths starting with the given prefix
254    /// are included. The filter is applied early in the pipeline so that
255    /// building the `HashSet` only contains relevant entries rather than
256    /// the entire tree (which can be 100k+ files in large repos).
257    pub fn list_files(
258        &self,
259        git_repo: &GitRepository,
260        only_modified: bool,
261        prefix: Option<&str>,
262    ) -> Result<Vec<String>> {
263        let matches_prefix = |p: &str| -> bool {
264            match prefix {
265                Some(pfx) => p.starts_with(pfx),
266                None => true,
267            }
268        };
269
270        if only_modified {
271            return Ok(self
272                .overlay
273                .list_changes()
274                .into_iter()
275                .filter(|(path, _)| matches_prefix(path))
276                .map(|(path, _)| path)
277                .collect());
278        }
279
280        // Start with base tree — filter by prefix before collecting into
281        // the HashSet to avoid allocating entries we will immediately discard.
282        let base_files = git_repo.list_tree_files(&self.base_commit)?;
283        let mut result: HashSet<String> = base_files
284            .into_iter()
285            .filter(|p| matches_prefix(p))
286            .collect();
287
288        // Apply overlay (only entries matching the prefix)
289        for (path, entry) in self.overlay.list_changes() {
290            if !matches_prefix(&path) {
291                continue;
292            }
293            match entry {
294                OverlayEntry::Added { .. } | OverlayEntry::Modified { .. } => {
295                    result.insert(path);
296                }
297                OverlayEntry::Deleted => {
298                    result.remove(&path);
299                }
300            }
301        }
302
303        let mut files: Vec<String> = result.into_iter().collect();
304        files.sort();
305        Ok(files)
306    }
307
308    /// Touch the workspace to update last-active timestamp.
309    pub fn touch(&mut self) {
310        self.last_active = Instant::now();
311    }
312
313    /// Build the overlay vector for `commit_tree_overlay`.
314    ///
315    /// Returns `(path, Some(content))` for modified/added files and
316    /// `(path, None)` for deleted files.
317    pub fn overlay_for_tree(&self) -> Vec<(String, Option<Vec<u8>>)> {
318        self.overlay
319            .list_changes()
320            .into_iter()
321            .map(|(path, entry)| {
322                let data = match entry {
323                    OverlayEntry::Modified { content, .. }
324                    | OverlayEntry::Added { content, .. } => Some(content),
325                    OverlayEntry::Deleted => None,
326                };
327                (path, data)
328            })
329            .collect()
330    }
331}
dk_engine/workspace/session_workspace.rs

dk_engine/workspace/
session_workspace.rs
