tenant.openagents/omega

1use anyhow::Result;
2use async_recursion::async_recursion;
3use collections::HashSet;
4use futures::future::join_all;
5use gpui::{AppContext as _, AsyncWindowContext, Axis, Entity, Task, WeakEntity};
6use project::Project;
7use serde::{Deserialize, Serialize};
8use std::path::PathBuf;
9use ui::{App, Context, Window};
10use util::ResultExt as _;
11
12use db::{
13    query,
14    sqlez::{domain::Domain, statement::Statement, thread_safe_connection::ThreadSafeConnection},
15    sqlez_macros::sql,
16};
17use workspace::{
18    ItemHandle, ItemId, Member, Pane, PaneAxis, PaneGroup, SerializableItem as _, Workspace,
19    WorkspaceDb, WorkspaceId,
20};
21
22use crate::{
23    TerminalView, default_working_directory,
24    terminal_panel::{TerminalPanel, new_terminal_pane},
25};
26
27pub(crate) fn serialize_pane_group(
28    pane_group: &PaneGroup,
29    active_pane: &Entity<Pane>,
30    cx: &mut App,
31) -> SerializedPaneGroup {
32    build_serialized_pane_group(&pane_group.root, active_pane, cx)
33}
34
35fn build_serialized_pane_group(
36    pane_group: &Member,
37    active_pane: &Entity<Pane>,
38    cx: &mut App,
39) -> SerializedPaneGroup {
40    match pane_group {
41        Member::Axis(PaneAxis {
42            axis,
43            members,
44            flexes,
45            bounding_boxes: _,
46        }) => SerializedPaneGroup::Group {
47            axis: SerializedAxis(*axis),
48            children: members
49                .iter()
50                .map(|member| build_serialized_pane_group(member, active_pane, cx))
51                .collect::<Vec<_>>(),
52            flexes: Some(flexes.lock().clone()),
53        },
54        Member::Pane(pane_handle) => {
55            SerializedPaneGroup::Pane(serialize_pane(pane_handle, pane_handle == active_pane, cx))
56        }
57    }
58}
59
60fn serialize_pane(pane: &Entity<Pane>, active: bool, cx: &mut App) -> SerializedPane {
61    let mut items_to_serialize = HashSet::default();
62    let pane = pane.read(cx);
63    let children = pane
64        .items()
65        .filter_map(|item| {
66            let terminal_view = item.act_as::<TerminalView>(cx)?;
67            if terminal_view.read(cx).terminal().read(cx).task().is_some() {
68                None
69            } else {
70                let id = item.item_id().as_u64();
71                items_to_serialize.insert(id);
72                Some(id)
73            }
74        })
75        .collect::<Vec<_>>();
76    let active_item = pane
77        .active_item()
78        .map(|item| item.item_id().as_u64())
79        .filter(|active_id| items_to_serialize.contains(active_id));
80
81    let pinned_count = pane.pinned_count();
82    SerializedPane {
83        active,
84        children,
85        active_item,
86        pinned_count,
87    }
88}
89
90pub(crate) fn deserialize_terminal_panel(
91    workspace: WeakEntity<Workspace>,
92    project: Entity<Project>,
93    database_id: WorkspaceId,
94    serialized_panel: SerializedTerminalPanel,
95    window: &mut Window,
96    cx: &mut App,
97) -> Task<anyhow::Result<Entity<TerminalPanel>>> {
98    window.spawn(cx, async move |cx| {
99        let terminal_panel = workspace.update_in(cx, |workspace, window, cx| {
100            cx.new(|cx| TerminalPanel::new(workspace, window, cx))
101        })?;
102        match &serialized_panel.items {
103            SerializedItems::NoSplits(item_ids) => {
104                let items = deserialize_terminal_views(
105                    database_id,
106                    project,
107                    workspace,
108                    item_ids.as_slice(),
109                    cx,
110                )
111                .await;
112                let active_item = serialized_panel.active_item_id;
113                terminal_panel.update_in(cx, |terminal_panel, window, cx| {
114                    terminal_panel.active_pane.update(cx, |pane, cx| {
115                        populate_pane_items(pane, items, active_item, window, cx);
116                    });
117                })?;
118            }
119            SerializedItems::WithSplits(serialized_pane_group) => {
120                let center_pane = deserialize_pane_group(
121                    workspace,
122                    project,
123                    terminal_panel.clone(),
124                    database_id,
125                    serialized_pane_group,
126                    cx,
127                )
128                .await;
129                if let Some((center_group, active_pane)) = center_pane {
130                    terminal_panel.update(cx, |terminal_panel, _| {
131                        terminal_panel.center = PaneGroup::with_root(center_group);
132                        terminal_panel.active_pane =
133                            active_pane.unwrap_or_else(|| terminal_panel.center.first_pane());
134                    });
135                }
136            }
137        }
138
139        Ok(terminal_panel)
140    })
141}
142
143fn populate_pane_items(
144    pane: &mut Pane,
145    items: Vec<Entity<TerminalView>>,
146    active_item: Option<u64>,
147    window: &mut Window,
148    cx: &mut Context<Pane>,
149) {
150    let mut active_item_index = None;
151    for (item_index, item) in (pane.items_len()..).zip(items) {
152        if Some(item.item_id().as_u64()) == active_item {
153            active_item_index = Some(item_index);
154        }
155        pane.add_item(Box::new(item), false, false, None, window, cx);
156    }
157    if let Some(index) = active_item_index {
158        pane.activate_item(index, false, false, window, cx);
159    }
160}
161
162#[async_recursion(?Send)]
163async fn deserialize_pane_group(
164    workspace: WeakEntity<Workspace>,
165    project: Entity<Project>,
166    panel: Entity<TerminalPanel>,
167    workspace_id: WorkspaceId,
168    serialized: &SerializedPaneGroup,
169    cx: &mut AsyncWindowContext,
170) -> Option<(Member, Option<Entity<Pane>>)> {
171    match serialized {
172        SerializedPaneGroup::Group {
173            axis,
174            flexes,
175            children,
176        } => {
177            let mut current_active_pane = None;
178            let mut members = Vec::new();
179            for child in children {
180                if let Some((new_member, active_pane)) = deserialize_pane_group(
181                    workspace.clone(),
182                    project.clone(),
183                    panel.clone(),
184                    workspace_id,
185                    child,
186                    cx,
187                )
188                .await
189                {
190                    members.push(new_member);
191                    current_active_pane = current_active_pane.or(active_pane);
192                }
193            }
194
195            if members.is_empty() {
196                return None;
197            }
198
199            if members.len() == 1 {
200                return Some((members.remove(0), current_active_pane));
201            }
202
203            Some((
204                Member::Axis(PaneAxis::load(axis.0, members, flexes.clone())),
205                current_active_pane,
206            ))
207        }
208        SerializedPaneGroup::Pane(serialized_pane) => {
209            let active = serialized_pane.active;
210
211            let pane = panel
212                .update_in(cx, |terminal_panel, window, cx| {
213                    new_terminal_pane(
214                        workspace.clone(),
215                        project.clone(),
216                        terminal_panel.active_pane.read(cx).is_zoomed(),
217                        window,
218                        cx,
219                    )
220                })
221                .log_err()?;
222            let active_item = serialized_pane.active_item;
223            let pinned_count = serialized_pane.pinned_count;
224            let new_items = deserialize_terminal_views(
225                workspace_id,
226                project.clone(),
227                workspace.clone(),
228                serialized_pane.children.as_slice(),
229                cx,
230            );
231            cx.spawn({
232                let pane = pane.downgrade();
233                async move |cx| {
234                    let new_items = new_items.await;
235
236                    let items = pane.update_in(cx, |pane, window, cx| {
237                        populate_pane_items(pane, new_items, active_item, window, cx);
238                        pane.set_pinned_count(pinned_count.min(pane.items_len()));
239                        pane.items_len()
240                    });
241                    // Avoid blank panes in splits
242                    if items.is_ok_and(|items| items == 0) {
243                        let working_directory = workspace
244                            .update(cx, |workspace, cx| default_working_directory(workspace, cx))
245                            .ok()
246                            .flatten();
247                        let terminal = project
248                            .update(cx, |project, cx| {
249                                project.create_terminal_shell(working_directory, cx)
250                            })
251                            .await
252                            .log_err();
253                        let Some(terminal) = terminal else {
254                            return;
255                        };
256                        pane.update_in(cx, |pane, window, cx| {
257                            let terminal_view = Box::new(cx.new(|cx| {
258                                TerminalView::new(
259                                    terminal,
260                                    workspace.clone(),
261                                    Some(workspace_id),
262                                    project.downgrade(),
263                                    window,
264                                    cx,
265                                )
266                            }));
267                            pane.add_item(terminal_view, true, false, None, window, cx);
268                        })
269                        .ok();
270                    }
271                }
272            })
273            .await;
274            Some((Member::Pane(pane.clone()), active.then_some(pane)))
275        }
276    }
277}
278
279fn deserialize_terminal_views(
280    workspace_id: WorkspaceId,
281    project: Entity<Project>,
282    workspace: WeakEntity<Workspace>,
283    item_ids: &[u64],
284    cx: &mut AsyncWindowContext,
285) -> impl Future<Output = Vec<Entity<TerminalView>>> + use<> {
286    let deserialized_items = join_all(item_ids.iter().filter_map(|item_id| {
287        cx.update(|window, cx| {
288            TerminalView::deserialize(
289                project.clone(),
290                workspace.clone(),
291                workspace_id,
292                *item_id,
293                window,
294                cx,
295            )
296        })
297        .ok()
298    }));
299    async move {
300        deserialized_items
301            .await
302            .into_iter()
303            .filter_map(|item| item.log_err())
304            .collect()
305    }
306}
307
308#[derive(Debug, Serialize, Deserialize)]
309pub(crate) struct SerializedTerminalPanel {
310    pub items: SerializedItems,
311    // A deprecated field, kept for backwards compatibility for the code before terminal splits were introduced.
312    pub active_item_id: Option<u64>,
313}
314
315#[derive(Debug, Serialize, Deserialize)]
316#[serde(untagged)]
317pub(crate) enum SerializedItems {
318    // The data stored before terminal splits were introduced.
319    NoSplits(Vec<u64>),
320    WithSplits(SerializedPaneGroup),
321}
322
323#[derive(Debug, Serialize, Deserialize)]
324pub(crate) enum SerializedPaneGroup {
325    Pane(SerializedPane),
326    Group {
327        axis: SerializedAxis,
328        flexes: Option<Vec<f32>>,
329        children: Vec<SerializedPaneGroup>,
330    },
331}
332
333#[derive(Debug, Serialize, Deserialize)]
334pub(crate) struct SerializedPane {
335    pub active: bool,
336    pub children: Vec<u64>,
337    pub active_item: Option<u64>,
338    #[serde(default)]
339    pub pinned_count: usize,
340}
341
342#[derive(Debug)]
343pub(crate) struct SerializedAxis(pub Axis);
344
345impl Serialize for SerializedAxis {
346    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
347    where
348        S: serde::Serializer,
349    {
350        match self.0 {
351            Axis::Horizontal => serializer.serialize_str("horizontal"),
352            Axis::Vertical => serializer.serialize_str("vertical"),
353        }
354    }
355}
356
357impl<'de> Deserialize<'de> for SerializedAxis {
358    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
359    where
360        D: serde::Deserializer<'de>,
361    {
362        let s = String::deserialize(deserializer)?;
363        match s.as_str() {
364            "horizontal" => Ok(SerializedAxis(Axis::Horizontal)),
365            "vertical" => Ok(SerializedAxis(Axis::Vertical)),
366            invalid => Err(serde::de::Error::custom(format!(
367                "Invalid axis value: '{invalid}'"
368            ))),
369        }
370    }
371}
372
373pub struct TerminalDb(ThreadSafeConnection);
374
375impl Domain for TerminalDb {
376    const NAME: &str = stringify!(TerminalDb);
377
378    const MIGRATIONS: &[&str] = &[
379        sql!(
380            CREATE TABLE terminals (
381                workspace_id INTEGER,
382                item_id INTEGER UNIQUE,
383                working_directory BLOB,
384                PRIMARY KEY(workspace_id, item_id),
385                FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id)
386                ON DELETE CASCADE
387            ) STRICT;
388        ),
389        // Remove the unique constraint on the item_id table
390        // SQLite doesn't have a way of doing this automatically, so
391        // we have to do this silly copying.
392        sql!(
393            CREATE TABLE terminals2 (
394                workspace_id INTEGER,
395                item_id INTEGER,
396                working_directory BLOB,
397                PRIMARY KEY(workspace_id, item_id),
398                FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id)
399                ON DELETE CASCADE
400            ) STRICT;
401
402            INSERT INTO terminals2 (workspace_id, item_id, working_directory)
403            SELECT workspace_id, item_id, working_directory FROM terminals;
404
405            DROP TABLE terminals;
406
407            ALTER TABLE terminals2 RENAME TO terminals;
408        ),
409        sql! (
410            ALTER TABLE terminals ADD COLUMN working_directory_path TEXT;
411            UPDATE terminals SET working_directory_path = CAST(working_directory AS TEXT);
412        ),
413        sql! (
414            ALTER TABLE terminals ADD COLUMN custom_title TEXT;
415        ),
416    ];
417}
418
419db::static_connection!(TerminalDb, [WorkspaceDb]);
420
421impl TerminalDb {
422    query! {
423       pub async fn update_workspace_id(
424            new_id: WorkspaceId,
425            old_id: WorkspaceId,
426            item_id: ItemId
427        ) -> Result<()> {
428            UPDATE terminals
429            SET workspace_id = ?
430            WHERE workspace_id = ? AND item_id = ?
431        }
432    }
433
434    pub async fn save_working_directory(
435        &self,
436        item_id: ItemId,
437        workspace_id: WorkspaceId,
438        working_directory: PathBuf,
439    ) -> Result<()> {
440        log::debug!(
441            "Saving working directory {working_directory:?} for item {item_id} in workspace {workspace_id:?}"
442        );
443        let query =
444            "INSERT INTO terminals(item_id, workspace_id, working_directory, working_directory_path)
445            VALUES (?1, ?2, ?3, ?4)
446            ON CONFLICT DO UPDATE SET
447                item_id = ?1,
448                workspace_id = ?2,
449                working_directory = ?3,
450                working_directory_path = ?4"
451        ;
452        self.write(move |conn| {
453            let mut statement = Statement::prepare(conn, query)?;
454            let mut next_index = statement.bind(&item_id, 1)?;
455            next_index = statement.bind(&workspace_id, next_index)?;
456            next_index = statement.bind(&working_directory, next_index)?;
457            statement.bind(
458                &working_directory.to_string_lossy().into_owned(),
459                next_index,
460            )?;
461            statement.exec()
462        })
463        .await
464    }
465
466    query! {
467        pub fn get_working_directory(item_id: ItemId, workspace_id: WorkspaceId) -> Result<Option<PathBuf>> {
468            SELECT working_directory
469            FROM terminals
470            WHERE item_id = ? AND workspace_id = ?
471        }
472    }
473
474    pub async fn save_custom_title(
475        &self,
476        item_id: ItemId,
477        workspace_id: WorkspaceId,
478        custom_title: Option<String>,
479    ) -> Result<()> {
480        log::debug!(
481            "Saving custom title {:?} for item {} in workspace {:?}",
482            custom_title,
483            item_id,
484            workspace_id
485        );
486        self.write(move |conn| {
487            let query = "INSERT INTO terminals (item_id, workspace_id, custom_title)
488                VALUES (?1, ?2, ?3)
489                ON CONFLICT (workspace_id, item_id) DO UPDATE SET
490                    custom_title = excluded.custom_title";
491            let mut statement = Statement::prepare(conn, query)?;
492            let mut next_index = statement.bind(&item_id, 1)?;
493            next_index = statement.bind(&workspace_id, next_index)?;
494            statement.bind(&custom_title, next_index)?;
495            statement.exec()
496        })
497        .await
498    }
499
500    query! {
501        pub fn get_custom_title(item_id: ItemId, workspace_id: WorkspaceId) -> Result<Option<String>> {
502            SELECT custom_title
503            FROM terminals
504            WHERE item_id = ? AND workspace_id = ?
505        }
506    }
507}
508