tenant.openagents/omega

1use crate::{AgentMessage, AgentMessageContent, UserMessage, UserMessageContent};
2use acp_thread::ClientUserMessageId;
3use agent_client_protocol::schema::v1 as acp;
4use agent_settings::AgentProfileId;
5use anyhow::Result;
6use chrono::{DateTime, Utc};
7use collections::{HashMap, IndexMap};
8use futures::{FutureExt, future::Shared};
9use gpui::{BackgroundExecutor, Global, Task};
10use indoc::indoc;
11use language_model::Speed;
12use parking_lot::Mutex;
13use serde::{Deserialize, Serialize};
14use sqlez::{
15    bindable::{Bind, Column},
16    connection::Connection,
17    statement::Statement,
18};
19use std::{io::ErrorKind, path::PathBuf, sync::Arc};
20use ui::{App, SharedString};
21use util::path_list::PathList;
22use zed_env_vars::ZED_STATELESS;
23
24pub type DbMessage = crate::Message;
25pub type DbSummary = crate::legacy_thread::DetailedSummaryState;
26pub type DbLanguageModel = crate::legacy_thread::SerializedLanguageModel;
27
28#[derive(Debug, Clone)]
29pub struct DbThreadMetadata {
30    pub id: acp::SessionId,
31    pub parent_session_id: Option<acp::SessionId>,
32    pub title: SharedString,
33    pub updated_at: DateTime<Utc>,
34    pub created_at: Option<DateTime<Utc>>,
35    /// The workspace folder paths this thread was created against, sorted
36    /// lexicographically. Used for grouping threads by project in the sidebar.
37    pub folder_paths: PathList,
38}
39
40impl From<&DbThreadMetadata> for acp_thread::AgentSessionInfo {
41    fn from(meta: &DbThreadMetadata) -> Self {
42        Self {
43            session_id: meta.id.clone(),
44            work_dirs: Some(meta.folder_paths.clone()),
45            title: Some(meta.title.clone()),
46            updated_at: Some(meta.updated_at),
47            created_at: meta.created_at,
48            meta: None,
49        }
50    }
51}
52
53#[derive(Debug, Serialize, Deserialize)]
54pub struct DbThread {
55    pub title: SharedString,
56    pub messages: Vec<Arc<DbMessage>>,
57    pub updated_at: DateTime<Utc>,
58    #[serde(default)]
59    pub detailed_summary: Option<SharedString>,
60    #[serde(default)]
61    pub initial_project_snapshot: Option<Arc<crate::ProjectSnapshot>>,
62    #[serde(default)]
63    pub cumulative_token_usage: language_model::TokenUsage,
64    #[serde(default)]
65    pub request_token_usage: HashMap<acp_thread::ClientUserMessageId, language_model::TokenUsage>,
66    #[serde(default)]
67    pub model: Option<DbLanguageModel>,
68    #[serde(default)]
69    pub profile: Option<AgentProfileId>,
70    #[serde(default)]
71    pub subagent_context: Option<crate::SubagentContext>,
72    #[serde(default)]
73    pub speed: Option<Speed>,
74    #[serde(default)]
75    pub thinking_enabled: bool,
76    #[serde(default)]
77    pub thinking_effort: Option<String>,
78    #[serde(default)]
79    pub draft_prompt: Option<Vec<acp::ContentBlock>>,
80    #[serde(default)]
81    pub ui_scroll_position: Option<SerializedScrollPosition>,
82    #[serde(default)]
83    pub sandboxed_terminal_temp_dir: Option<PathBuf>,
84    /// Sandbox escalations the user approved "for the rest of this thread".
85    /// Persisted so reopening a thread keeps its grants. See
86    /// [`crate::sandboxing::ThreadSandboxGrants`].
87    #[serde(default)]
88    pub sandbox_grants: DbSandboxGrants,
89}
90
91/// Serialized form of the sandbox permissions the user granted "for the rest of
92/// this thread" (the "Allow for this thread" prompt option). Stored inside the
93/// thread blob; round-trips with [`crate::sandboxing::ThreadSandboxGrants`].
94#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
95pub struct DbSandboxGrants {
96    /// Canonicalized paths granted write access; each covers its whole subtree.
97    #[serde(default)]
98    pub write_paths: Vec<PathBuf>,
99    /// Host patterns granted network access, in canonical string form (e.g.
100    /// `github.com`, `*.npmjs.org`). Parsed back into patterns on load.
101    #[serde(default)]
102    pub network_hosts: Vec<String>,
103    /// Whether arbitrary-host network access was granted.
104    #[serde(default)]
105    pub network_any_host: bool,
106    /// Whether unrestricted filesystem writes (the broad escape hatch) were
107    /// granted.
108    #[serde(default)]
109    pub allow_fs_write_all: bool,
110
111    /// Whether the model-requested fully-unsandboxed escape was granted.
112    #[serde(default)]
113    pub unsandboxed: bool,
114    /// Whether running commands unsandboxed was allowed because the OS sandbox
115    /// could not be created (the fallback prompt's "for this thread" option).
116    #[serde(default)]
117    pub sandbox_fallback: bool,
118}
119
120#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
121pub struct SerializedScrollPosition {
122    pub item_ix: usize,
123    pub offset_in_item: f32,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize)]
127pub struct SharedThread {
128    pub title: SharedString,
129    pub messages: Vec<Arc<DbMessage>>,
130    pub updated_at: DateTime<Utc>,
131    #[serde(default)]
132    pub model: Option<DbLanguageModel>,
133    pub version: String,
134}
135
136impl SharedThread {
137    pub const VERSION: &'static str = "1.0.0";
138
139    pub fn from_db_thread(thread: &DbThread) -> Self {
140        Self {
141            title: thread.title.clone(),
142            messages: thread.messages.clone(),
143            updated_at: thread.updated_at,
144            model: thread.model.clone(),
145            version: Self::VERSION.to_string(),
146        }
147    }
148
149    pub fn to_db_thread(self) -> DbThread {
150        DbThread {
151            title: format!("🔗 {}", self.title).into(),
152            messages: self.messages,
153            updated_at: self.updated_at,
154            detailed_summary: None,
155            initial_project_snapshot: None,
156            cumulative_token_usage: Default::default(),
157            request_token_usage: Default::default(),
158            model: self.model,
159            profile: None,
160            subagent_context: None,
161            speed: None,
162            thinking_enabled: false,
163            thinking_effort: None,
164            draft_prompt: None,
165            ui_scroll_position: None,
166            sandboxed_terminal_temp_dir: None,
167            sandbox_grants: DbSandboxGrants::default(),
168        }
169    }
170
171    pub fn to_bytes(&self) -> Result<Vec<u8>> {
172        const COMPRESSION_LEVEL: i32 = 3;
173        let json = serde_json::to_vec(self)?;
174        let compressed = zstd::encode_all(json.as_slice(), COMPRESSION_LEVEL)?;
175        Ok(compressed)
176    }
177
178    pub fn from_bytes(data: &[u8]) -> Result<Self> {
179        let decompressed = zstd::decode_all(data)?;
180        Ok(serde_json::from_slice(&decompressed)?)
181    }
182}
183
184impl DbThread {
185    pub const VERSION: &'static str = "0.3.0";
186
187    pub fn to_markdown(&self) -> String {
188        crate::messages_to_markdown(&self.messages)
189    }
190
191    pub fn from_json(json: &[u8]) -> Result<Self> {
192        let saved_thread_json = serde_json::from_slice::<serde_json::Value>(json)?;
193        match saved_thread_json.get("version") {
194            Some(serde_json::Value::String(version)) => match version.as_str() {
195                Self::VERSION => Ok(serde_json::from_value(saved_thread_json)?),
196                _ => Self::upgrade_from_agent_1(crate::legacy_thread::SerializedThread::from_json(
197                    json,
198                )?),
199            },
200            _ => {
201                Self::upgrade_from_agent_1(crate::legacy_thread::SerializedThread::from_json(json)?)
202            }
203        }
204    }
205
206    fn upgrade_from_agent_1(thread: crate::legacy_thread::SerializedThread) -> Result<Self> {
207        let mut messages = Vec::new();
208        let mut request_token_usage = HashMap::default();
209
210        let mut last_user_message_id = None;
211        for (ix, msg) in thread.messages.into_iter().enumerate() {
212            let message = match msg.role {
213                language_model::Role::User => {
214                    let mut content = Vec::new();
215
216                    // Convert segments to content
217                    for segment in msg.segments {
218                        match segment {
219                            crate::legacy_thread::SerializedMessageSegment::Text { text } => {
220                                content.push(UserMessageContent::Text(text));
221                            }
222                            crate::legacy_thread::SerializedMessageSegment::Thinking {
223                                text,
224                                ..
225                            } => {
226                                // User messages don't have thinking segments, but handle gracefully
227                                content.push(UserMessageContent::Text(text));
228                            }
229                            crate::legacy_thread::SerializedMessageSegment::RedactedThinking {
230                                ..
231                            } => {
232                                // User messages don't have redacted thinking, skip.
233                            }
234                        }
235                    }
236
237                    // If no content was added, add context as text if available
238                    if content.is_empty() && !msg.context.is_empty() {
239                        content.push(UserMessageContent::Text(msg.context));
240                    }
241
242                    let id = ClientUserMessageId::new();
243                    last_user_message_id = Some(id.clone());
244
245                    crate::Message::User(UserMessage {
246                        // MessageId from old format can't be meaningfully converted, so generate a new one
247                        id,
248                        content: Arc::from(content),
249                    })
250                }
251                language_model::Role::Assistant => {
252                    let mut content = Vec::new();
253
254                    // Convert segments to content
255                    for segment in msg.segments {
256                        match segment {
257                            crate::legacy_thread::SerializedMessageSegment::Text { text } => {
258                                content.push(AgentMessageContent::Text(text));
259                            }
260                            crate::legacy_thread::SerializedMessageSegment::Thinking {
261                                text,
262                                signature,
263                            } => {
264                                content.push(AgentMessageContent::Thinking { text, signature });
265                            }
266                            crate::legacy_thread::SerializedMessageSegment::RedactedThinking {
267                                data,
268                            } => {
269                                content.push(AgentMessageContent::RedactedThinking(data));
270                            }
271                        }
272                    }
273
274                    // Convert tool uses
275                    let mut tool_names_by_id = HashMap::default();
276                    for tool_use in msg.tool_uses {
277                        tool_names_by_id.insert(tool_use.id.clone(), tool_use.name.clone());
278                        content.push(AgentMessageContent::ToolUse(
279                            language_model::LanguageModelToolUse {
280                                id: tool_use.id,
281                                name: tool_use.name.into(),
282                                raw_input: serde_json::to_string(&tool_use.input)
283                                    .unwrap_or_default(),
284                                input: language_model::LanguageModelToolUseInput::Json(
285                                    tool_use.input,
286                                ),
287                                is_input_complete: true,
288                                thought_signature: None,
289                            },
290                        ));
291                    }
292
293                    // Convert tool results
294                    let mut tool_results = IndexMap::default();
295                    for tool_result in msg.tool_results {
296                        let name = tool_names_by_id
297                            .remove(&tool_result.tool_use_id)
298                            .unwrap_or_else(|| SharedString::from("unknown"));
299                        tool_results.insert(
300                            tool_result.tool_use_id.clone(),
301                            language_model::LanguageModelToolResult {
302                                tool_use_id: tool_result.tool_use_id,
303                                tool_name: name.into(),
304                                is_error: tool_result.is_error,
305                                content: vec![tool_result.content],
306                                output: tool_result.output,
307                            },
308                        );
309                    }
310
311                    if let Some(last_user_message_id) = &last_user_message_id
312                        && let Some(token_usage) = thread.request_token_usage.get(ix).copied()
313                    {
314                        request_token_usage.insert(last_user_message_id.clone(), token_usage);
315                    }
316
317                    crate::Message::Agent(AgentMessage {
318                        content,
319                        tool_results,
320                        reasoning_details: None,
321                    })
322                }
323                language_model::Role::System => {
324                    // Skip system messages as they're not supported in the new format
325                    continue;
326                }
327            };
328
329            messages.push(Arc::new(message));
330        }
331
332        Ok(Self {
333            title: thread.summary,
334            messages,
335            updated_at: thread.updated_at,
336            detailed_summary: match thread.detailed_summary_state {
337                crate::legacy_thread::DetailedSummaryState::NotGenerated
338                | crate::legacy_thread::DetailedSummaryState::Generating => None,
339                crate::legacy_thread::DetailedSummaryState::Generated { text, .. } => Some(text),
340            },
341            initial_project_snapshot: thread.initial_project_snapshot,
342            cumulative_token_usage: thread.cumulative_token_usage,
343            request_token_usage,
344            model: thread.model,
345            profile: thread.profile,
346            subagent_context: None,
347            speed: None,
348            thinking_enabled: false,
349            thinking_effort: None,
350            draft_prompt: None,
351            ui_scroll_position: None,
352            sandboxed_terminal_temp_dir: None,
353            sandbox_grants: DbSandboxGrants::default(),
354        })
355    }
356}
357
358#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
359pub enum DataType {
360    #[serde(rename = "json")]
361    Json,
362    #[serde(rename = "zstd")]
363    Zstd,
364}
365
366impl Bind for DataType {
367    fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {
368        let value = match self {
369            DataType::Json => "json",
370            DataType::Zstd => "zstd",
371        };
372        value.bind(statement, start_index)
373    }
374}
375
376impl Column for DataType {
377    fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
378        let (value, next_index) = String::column(statement, start_index)?;
379        let data_type = match value.as_str() {
380            "json" => DataType::Json,
381            "zstd" => DataType::Zstd,
382            _ => anyhow::bail!("Unknown data type: {}", value),
383        };
384        Ok((data_type, next_index))
385    }
386}
387
388pub(crate) struct ThreadsDatabase {
389    executor: BackgroundExecutor,
390    connection: Arc<Mutex<Connection>>,
391}
392
393struct GlobalThreadsDatabase(Shared<Task<Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>);
394
395impl Global for GlobalThreadsDatabase {}
396
397impl ThreadsDatabase {
398    pub fn connect(cx: &mut App) -> Shared<Task<Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>> {
399        if cx.has_global::<GlobalThreadsDatabase>() {
400            return cx.global::<GlobalThreadsDatabase>().0.clone();
401        }
402        let executor = cx.background_executor().clone();
403        let task = executor
404            .spawn({
405                let executor = executor.clone();
406                async move {
407                    match ThreadsDatabase::new(executor) {
408                        Ok(db) => Ok(Arc::new(db)),
409                        Err(err) => Err(Arc::new(err)),
410                    }
411                }
412            })
413            .shared();
414
415        cx.set_global(GlobalThreadsDatabase(task.clone()));
416        task
417    }
418
419    pub fn new(executor: BackgroundExecutor) -> Result<Self> {
420        let connection = if *ZED_STATELESS {
421            Connection::open_memory(Some("THREAD_FALLBACK_DB"))
422        } else if cfg!(any(feature = "test-support", test)) {
423            // rust stores the name of the test on the current thread.
424            // We use this to automatically create a database that will
425            // be shared within the test (for the test_retrieve_old_thread)
426            // but not with concurrent tests.
427            let thread = std::thread::current();
428            let test_name = thread.name();
429            Connection::open_memory(Some(&format!(
430                "THREAD_FALLBACK_{}",
431                test_name.unwrap_or_default()
432            )))
433        } else {
434            let threads_dir = paths::data_dir().join("threads");
435            std::fs::create_dir_all(&threads_dir)?;
436            let sqlite_path = threads_dir.join("threads.db");
437            Connection::open_file(&sqlite_path.to_string_lossy())
438        };
439
440        connection.exec(indoc! {"
441            CREATE TABLE IF NOT EXISTS threads (
442                id TEXT PRIMARY KEY,
443                summary TEXT NOT NULL,
444                updated_at TEXT NOT NULL,
445                data_type TEXT NOT NULL,
446                data BLOB NOT NULL
447            )
448        "})?()
449        .map_err(|e| e.context("Failed to create threads table"))?;
450
451        if let Ok(mut s) = connection.exec(indoc! {"
452            ALTER TABLE threads ADD COLUMN parent_id TEXT
453        "})
454        {
455            s().ok();
456        }
457
458        if let Ok(mut s) = connection.exec(indoc! {"
459            ALTER TABLE threads ADD COLUMN folder_paths TEXT;
460            ALTER TABLE threads ADD COLUMN folder_paths_order TEXT;
461        "})
462        {
463            s().ok();
464        }
465
466        if let Ok(mut s) = connection.exec(indoc! {"
467            ALTER TABLE threads ADD COLUMN created_at TEXT;
468        "})
469        {
470            if s().is_ok() {
471                connection.exec(indoc! {"
472                    UPDATE threads SET created_at = updated_at WHERE created_at IS NULL
473                "})?()?;
474            }
475        }
476
477        let db = Self {
478            executor,
479            connection: Arc::new(Mutex::new(connection)),
480        };
481
482        Ok(db)
483    }
484
485    fn save_thread_sync(
486        connection: &Arc<Mutex<Connection>>,
487        id: acp::SessionId,
488        thread: DbThread,
489        folder_paths: &PathList,
490    ) -> Result<()> {
491        const COMPRESSION_LEVEL: i32 = 3;
492
493        #[derive(Serialize)]
494        struct SerializedThread {
495            #[serde(flatten)]
496            thread: DbThread,
497            version: &'static str,
498        }
499
500        let title = thread.title.to_string();
501        let updated_at = thread.updated_at.to_rfc3339();
502        let parent_id = thread
503            .subagent_context
504            .as_ref()
505            .map(|ctx| ctx.parent_thread_id.0.clone());
506        let serialized_folder_paths = folder_paths.serialize();
507        let (folder_paths_str, folder_paths_order_str): (Option<String>, Option<String>) =
508            if folder_paths.is_empty() {
509                (None, None)
510            } else {
511                (
512                    Some(serialized_folder_paths.paths),
513                    Some(serialized_folder_paths.order),
514                )
515            };
516        let json_data = serde_json::to_string(&SerializedThread {
517            thread,
518            version: DbThread::VERSION,
519        })?;
520
521        let connection = connection.lock();
522
523        let compressed = zstd::encode_all(json_data.as_bytes(), COMPRESSION_LEVEL)?;
524        let data_type = DataType::Zstd;
525        let data = compressed;
526
527        // Use the thread's updated_at as created_at for new threads.
528        // This ensures the creation time reflects when the thread was conceptually
529        // created, not when it was saved to the database.
530        let created_at = updated_at.clone();
531
532        let mut insert = connection.exec_bound::<(Arc<str>, Option<Arc<str>>, Option<String>, Option<String>, String, String, DataType, Vec<u8>, String)>(indoc! {"
533            INSERT INTO threads (id, parent_id, folder_paths, folder_paths_order, summary, updated_at, data_type, data, created_at)
534            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
535            ON CONFLICT(id) DO UPDATE SET
536                parent_id = excluded.parent_id,
537                folder_paths = excluded.folder_paths,
538                folder_paths_order = excluded.folder_paths_order,
539                summary = excluded.summary,
540                updated_at = excluded.updated_at,
541                data_type = excluded.data_type,
542                data = excluded.data
543        "})?;
544
545        insert((
546            id.0,
547            parent_id,
548            folder_paths_str,
549            folder_paths_order_str,
550            title,
551            updated_at,
552            data_type,
553            data,
554            created_at,
555        ))?;
556
557        Ok(())
558    }
559
560    pub fn list_threads(&self) -> Task<Result<Vec<DbThreadMetadata>>> {
561        let connection = self.connection.clone();
562
563        self.executor.spawn(async move {
564            let connection = connection.lock();
565
566            let mut select = connection
567                .select_bound::<(), (Arc<str>, Option<Arc<str>>, Option<String>, Option<String>, String, String, Option<String>)>(indoc! {"
568                SELECT id, parent_id, folder_paths, folder_paths_order, summary, updated_at, created_at FROM threads ORDER BY updated_at DESC, created_at DESC
569            "})?;
570
571            let rows = select(())?;
572            let mut threads = Vec::new();
573
574            for (id, parent_id, folder_paths, folder_paths_order, summary, updated_at, created_at) in rows {
575                let folder_paths = folder_paths
576                    .map(|paths| {
577                        PathList::deserialize(&util::path_list::SerializedPathList {
578                            paths,
579                            order: folder_paths_order.unwrap_or_default(),
580                        })
581                    })
582                    .unwrap_or_default();
583                let created_at = created_at
584                    .as_deref()
585                    .map(DateTime::parse_from_rfc3339)
586                    .transpose()?
587                    .map(|dt| dt.with_timezone(&Utc));
588
589                threads.push(DbThreadMetadata {
590                    id: acp::SessionId::new(id),
591                    parent_session_id: parent_id.map(acp::SessionId::new),
592                    title: summary.into(),
593                    updated_at: DateTime::parse_from_rfc3339(&updated_at)?.with_timezone(&Utc),
594                    created_at,
595                    folder_paths,
596                });
597            }
598
599            Ok(threads)
600        })
601    }
602
603    pub fn load_thread(&self, id: acp::SessionId) -> Task<Result<Option<DbThread>>> {
604        let connection = self.connection.clone();
605
606        self.executor.spawn(async move {
607            let connection = connection.lock();
608            let mut select = connection.select_bound::<Arc<str>, (DataType, Vec<u8>)>(indoc! {"
609                SELECT data_type, data FROM threads WHERE id = ? LIMIT 1
610            "})?;
611
612            let rows = select(id.0)?;
613            if let Some((data_type, data)) = rows.into_iter().next() {
614                Ok(Some(Self::deserialize_thread(data_type, data)?))
615            } else {
616                Ok(None)
617            }
618        })
619    }
620
621    pub fn save_thread(
622        &self,
623        id: acp::SessionId,
624        thread: DbThread,
625        folder_paths: PathList,
626    ) -> Task<Result<()>> {
627        let connection = self.connection.clone();
628
629        self.executor
630            .spawn(async move { Self::save_thread_sync(&connection, id, thread, &folder_paths) })
631    }
632
633    fn deserialize_thread(data_type: DataType, data: Vec<u8>) -> Result<DbThread> {
634        let json_data = match data_type {
635            DataType::Zstd => {
636                let decompressed = zstd::decode_all(&data[..])?;
637                String::from_utf8(decompressed)?
638            }
639            DataType::Json => String::from_utf8(data)?,
640        };
641        DbThread::from_json(json_data.as_bytes())
642    }
643
644    fn sandboxed_terminal_temp_dir(data_type: DataType, data: Vec<u8>) -> Option<PathBuf> {
645        match Self::deserialize_thread(data_type, data) {
646            Ok(thread) => thread.sandboxed_terminal_temp_dir,
647            Err(error) => {
648                log::warn!("failed to deserialize thread before deleting it: {error:#}");
649                None
650            }
651        }
652    }
653
654    fn remove_sandboxed_terminal_temp_dir(temp_dir: PathBuf) {
655        match std::fs::remove_dir_all(&temp_dir) {
656            Ok(()) => {}
657            Err(error) if error.kind() == ErrorKind::NotFound => {}
658            Err(error) => {
659                log::warn!(
660                    "failed to remove sandboxed terminal temp directory {}: {error}",
661                    temp_dir.display()
662                );
663            }
664        }
665    }
666
667    pub fn delete_thread(&self, id: acp::SessionId) -> Task<Result<()>> {
668        let connection = self.connection.clone();
669
670        self.executor.spawn(async move {
671            let sandboxed_terminal_temp_dirs = {
672                let connection = connection.lock();
673
674                let mut select_children =
675                    connection.select_bound::<Arc<str>, Arc<str>>(indoc! {"
676                    SELECT id FROM threads WHERE parent_id = ?
677                "})?;
678
679                // Collect target thread together with all of its transitive
680                // subagent threads
681                let mut ids_to_delete = vec![id.0.clone()];
682                let mut frontier = vec![id.0.clone()];
683                while let Some(parent) = frontier.pop() {
684                    for child in select_children(parent)? {
685                        ids_to_delete.push(child.clone());
686                        frontier.push(child);
687                    }
688                }
689
690                let mut select =
691                    connection.select_bound::<Arc<str>, (DataType, Vec<u8>)>(indoc! {"
692                    SELECT data_type, data FROM threads WHERE id = ? LIMIT 1
693                "})?;
694
695                let mut delete = connection.exec_bound::<Arc<str>>(indoc! {"
696                    DELETE FROM threads WHERE id = ?
697                "})?;
698
699                let mut sandboxed_terminal_temp_dirs = Vec::new();
700                for thread_id in ids_to_delete {
701                    if let Some(temp_dir) = select(thread_id.clone())?.into_iter().next().and_then(
702                        |(data_type, data)| Self::sandboxed_terminal_temp_dir(data_type, data),
703                    ) {
704                        sandboxed_terminal_temp_dirs.push(temp_dir);
705                    }
706                    delete(thread_id)?;
707                }
708
709                sandboxed_terminal_temp_dirs
710            };
711
712            for temp_dir in sandboxed_terminal_temp_dirs {
713                Self::remove_sandboxed_terminal_temp_dir(temp_dir);
714            }
715
716            Ok(())
717        })
718    }
719
720    pub fn delete_threads(&self) -> Task<Result<()>> {
721        let connection = self.connection.clone();
722
723        self.executor.spawn(async move {
724            let sandboxed_terminal_temp_dirs = {
725                let connection = connection.lock();
726
727                let mut select = connection.select_bound::<(), (DataType, Vec<u8>)>(indoc! {"
728                    SELECT data_type, data FROM threads
729                "})?;
730
731                let sandboxed_terminal_temp_dirs = select(())?
732                    .into_iter()
733                    .filter_map(|(data_type, data)| {
734                        Self::sandboxed_terminal_temp_dir(data_type, data)
735                    })
736                    .collect::<Vec<_>>();
737
738                let mut delete = connection.exec_bound::<()>(indoc! {"
739                    DELETE FROM threads
740                "})?;
741
742                delete(())?;
743
744                sandboxed_terminal_temp_dirs
745            };
746
747            for temp_dir in sandboxed_terminal_temp_dirs {
748                Self::remove_sandboxed_terminal_temp_dir(temp_dir);
749            }
750
751            Ok(())
752        })
753    }
754}
755
756#[cfg(test)]
757mod tests {
758    use super::*;
759    use chrono::{DateTime, TimeZone, Utc};
760    use collections::HashMap;
761    use gpui::TestAppContext;
762    use std::sync::Arc;
763
764    #[test]
765    fn test_shared_thread_roundtrip() {
766        let original = SharedThread {
767            title: "Test Thread".into(),
768            messages: vec![],
769            updated_at: Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
770            model: None,
771            version: SharedThread::VERSION.to_string(),
772        };
773
774        let bytes = original.to_bytes().expect("Failed to serialize");
775        let restored = SharedThread::from_bytes(&bytes).expect("Failed to deserialize");
776
777        assert_eq!(restored.title, original.title);
778        assert_eq!(restored.version, original.version);
779        assert_eq!(restored.updated_at, original.updated_at);
780    }
781
782    fn session_id(value: &str) -> acp::SessionId {
783        acp::SessionId::new(Arc::<str>::from(value))
784    }
785
786    fn make_thread(title: &str, updated_at: DateTime<Utc>) -> DbThread {
787        DbThread {
788            title: title.to_string().into(),
789            messages: Vec::new(),
790            updated_at,
791            detailed_summary: None,
792            initial_project_snapshot: None,
793            cumulative_token_usage: Default::default(),
794            request_token_usage: HashMap::default(),
795            model: None,
796            profile: None,
797            subagent_context: None,
798            speed: None,
799            thinking_enabled: false,
800            thinking_effort: None,
801            draft_prompt: None,
802            ui_scroll_position: None,
803            sandboxed_terminal_temp_dir: None,
804            sandbox_grants: DbSandboxGrants::default(),
805        }
806    }
807
808    #[gpui::test]
809    async fn test_list_threads_orders_by_created_at(cx: &mut TestAppContext) {
810        let database = ThreadsDatabase::new(cx.executor()).unwrap();
811
812        let older_id = session_id("thread-a");
813        let newer_id = session_id("thread-b");
814
815        let older_thread = make_thread(
816            "Thread A",
817            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
818        );
819        let newer_thread = make_thread(
820            "Thread B",
821            Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
822        );
823
824        database
825            .save_thread(older_id.clone(), older_thread, PathList::default())
826            .await
827            .unwrap();
828        database
829            .save_thread(newer_id.clone(), newer_thread, PathList::default())
830            .await
831            .unwrap();
832
833        let entries = database.list_threads().await.unwrap();
834        assert_eq!(entries.len(), 2);
835        assert_eq!(entries[0].id, newer_id);
836        assert_eq!(entries[1].id, older_id);
837    }
838
839    #[gpui::test]
840    async fn test_save_thread_replaces_metadata(cx: &mut TestAppContext) {
841        let database = ThreadsDatabase::new(cx.executor()).unwrap();
842
843        let thread_id = session_id("thread-a");
844        let original_thread = make_thread(
845            "Thread A",
846            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
847        );
848        let updated_thread = make_thread(
849            "Thread B",
850            Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap(),
851        );
852
853        database
854            .save_thread(thread_id.clone(), original_thread, PathList::default())
855            .await
856            .unwrap();
857        database
858            .save_thread(thread_id.clone(), updated_thread, PathList::default())
859            .await
860            .unwrap();
861
862        let entries = database.list_threads().await.unwrap();
863        assert_eq!(entries.len(), 1);
864        assert_eq!(entries[0].id, thread_id);
865        assert_eq!(entries[0].title.as_ref(), "Thread B");
866        assert_eq!(
867            entries[0].updated_at,
868            Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap()
869        );
870        assert!(
871            entries[0].created_at.is_some(),
872            "created_at should be populated"
873        );
874    }
875
876    #[test]
877    fn test_subagent_context_defaults_to_none() {
878        let json = r#"{
879            "title": "Old Thread",
880            "messages": [],
881            "updated_at": "2024-01-01T00:00:00Z"
882        }"#;
883
884        let db_thread: DbThread = serde_json::from_str(json).expect("Failed to deserialize");
885
886        assert!(
887            db_thread.subagent_context.is_none(),
888            "Legacy threads without subagent_context should default to None"
889        );
890    }
891
892    #[test]
893    fn test_draft_prompt_defaults_to_none() {
894        let json = r#"{
895            "title": "Old Thread",
896            "messages": [],
897            "updated_at": "2024-01-01T00:00:00Z"
898        }"#;
899
900        let db_thread: DbThread = serde_json::from_str(json).expect("Failed to deserialize");
901
902        assert!(
903            db_thread.draft_prompt.is_none(),
904            "Legacy threads without draft_prompt field should default to None"
905        );
906    }
907
908    #[test]
909    fn test_sandboxed_terminal_temp_dir_defaults_to_none() {
910        let json = r#"{
911            "title": "Old Thread",
912            "messages": [],
913            "updated_at": "2024-01-01T00:00:00Z"
914        }"#;
915
916        let db_thread: DbThread = serde_json::from_str(json).expect("Failed to deserialize");
917
918        assert!(
919            db_thread.sandboxed_terminal_temp_dir.is_none(),
920            "Legacy threads without sandboxed_terminal_temp_dir should default to None"
921        );
922    }
923
924    #[test]
925    fn test_sandbox_grants_default_when_absent() {
926        let json = r#"{
927            "title": "Old Thread",
928            "messages": [],
929            "updated_at": "2024-01-01T00:00:00Z"
930        }"#;
931
932        let db_thread: DbThread = serde_json::from_str(json).expect("Failed to deserialize");
933
934        assert_eq!(
935            db_thread.sandbox_grants,
936            DbSandboxGrants::default(),
937            "Legacy threads without sandbox_grants should default to empty grants"
938        );
939    }
940
941    #[gpui::test]
942    async fn test_sandbox_grants_roundtrip_through_save_load(cx: &mut TestAppContext) {
943        let database = ThreadsDatabase::new(cx.executor()).unwrap();
944        let thread_id = session_id("sandbox-grants-thread");
945        let mut thread = make_thread(
946            "Sandbox Grants Thread",
947            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
948        );
949        let grants = DbSandboxGrants {
950            write_paths: vec![PathBuf::from("/tmp/build")],
951            network_hosts: vec!["github.com".to_string(), "*.npmjs.org".to_string()],
952            network_any_host: false,
953            allow_fs_write_all: false,
954            unsandboxed: true,
955            sandbox_fallback: true,
956        };
957        thread.sandbox_grants = grants.clone();
958
959        database
960            .save_thread(thread_id.clone(), thread, PathList::default())
961            .await
962            .unwrap();
963
964        let loaded = database
965            .load_thread(thread_id)
966            .await
967            .unwrap()
968            .expect("thread should exist");
969        assert_eq!(loaded.sandbox_grants, grants);
970    }
971
972    #[gpui::test]
973    async fn test_sandboxed_terminal_temp_dir_roundtrips_through_save_load(
974        cx: &mut TestAppContext,
975    ) {
976        let database = ThreadsDatabase::new(cx.executor()).unwrap();
977        let thread_id = session_id("sandbox-temp-dir-thread");
978        let temp_dir = tempfile::Builder::new()
979            .prefix("omega-agent-terminal-test-")
980            .tempdir()
981            .unwrap()
982            .keep();
983        let mut thread = make_thread(
984            "Sandbox Temp Dir Thread",
985            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
986        );
987        thread.sandboxed_terminal_temp_dir = Some(temp_dir.clone());
988
989        database
990            .save_thread(thread_id.clone(), thread, PathList::default())
991            .await
992            .unwrap();
993
994        let loaded = database
995            .load_thread(thread_id)
996            .await
997            .unwrap()
998            .expect("thread should exist");
999        assert_eq!(loaded.sandboxed_terminal_temp_dir, Some(temp_dir.clone()));
1000        std::fs::remove_dir_all(temp_dir).unwrap();
1001    }
1002
1003    #[gpui::test]
1004    async fn test_delete_thread_removes_sandboxed_terminal_temp_dir(cx: &mut TestAppContext) {
1005        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1006        let thread_id = session_id("sandbox-temp-dir-delete-thread");
1007        let temp_dir = tempfile::Builder::new()
1008            .prefix("omega-agent-terminal-test-")
1009            .tempdir()
1010            .unwrap()
1011            .keep();
1012        std::fs::write(temp_dir.join("sentinel"), b"content").unwrap();
1013        let mut thread = make_thread(
1014            "Sandbox Temp Dir Delete Thread",
1015            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1016        );
1017        thread.sandboxed_terminal_temp_dir = Some(temp_dir.clone());
1018
1019        database
1020            .save_thread(thread_id.clone(), thread, PathList::default())
1021            .await
1022            .unwrap();
1023        database.delete_thread(thread_id).await.unwrap();
1024
1025        assert!(!temp_dir.exists());
1026    }
1027
1028    #[gpui::test]
1029    async fn test_delete_thread_deletes_subagent_threads(cx: &mut TestAppContext) {
1030        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1031
1032        let parent_id = session_id("parent-thread");
1033        let child_id = session_id("child-thread");
1034        let grandchild_id = session_id("grandchild-thread");
1035        let unrelated_id = session_id("unrelated-thread");
1036
1037        let parent_thread = make_thread(
1038            "Parent Thread",
1039            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1040        );
1041
1042        let mut child_thread = make_thread(
1043            "Child Subagent Thread",
1044            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1045        );
1046        child_thread.subagent_context = Some(crate::SubagentContext {
1047            parent_thread_id: parent_id.clone(),
1048            depth: 1,
1049        });
1050
1051        let mut grandchild_thread = make_thread(
1052            "Grandchild Subagent Thread",
1053            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1054        );
1055        grandchild_thread.subagent_context = Some(crate::SubagentContext {
1056            parent_thread_id: child_id.clone(),
1057            depth: 2,
1058        });
1059
1060        let unrelated_thread = make_thread(
1061            "Unrelated Thread",
1062            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1063        );
1064
1065        for (id, thread) in [
1066            (parent_id.clone(), parent_thread),
1067            (child_id.clone(), child_thread),
1068            (grandchild_id.clone(), grandchild_thread),
1069            (unrelated_id.clone(), unrelated_thread),
1070        ] {
1071            database
1072                .save_thread(id, thread, PathList::default())
1073                .await
1074                .unwrap();
1075        }
1076
1077        database.delete_thread(parent_id.clone()).await.unwrap();
1078
1079        let remaining = database.list_threads().await.unwrap();
1080        let remaining_ids: Vec<_> = remaining.iter().map(|thread| thread.id.clone()).collect();
1081        assert_eq!(remaining_ids, vec![unrelated_id]);
1082    }
1083
1084    #[gpui::test]
1085    async fn test_subagent_context_roundtrips_through_save_load(cx: &mut TestAppContext) {
1086        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1087
1088        let parent_id = session_id("parent-thread");
1089        let child_id = session_id("child-thread");
1090
1091        let mut child_thread = make_thread(
1092            "Subagent Thread",
1093            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1094        );
1095        child_thread.subagent_context = Some(crate::SubagentContext {
1096            parent_thread_id: parent_id.clone(),
1097            depth: 2,
1098        });
1099
1100        database
1101            .save_thread(child_id.clone(), child_thread, PathList::default())
1102            .await
1103            .unwrap();
1104
1105        let loaded = database
1106            .load_thread(child_id)
1107            .await
1108            .unwrap()
1109            .expect("thread should exist");
1110
1111        let context = loaded
1112            .subagent_context
1113            .expect("subagent_context should be restored");
1114        assert_eq!(context.parent_thread_id, parent_id);
1115        assert_eq!(context.depth, 2);
1116    }
1117
1118    #[gpui::test]
1119    async fn test_non_subagent_thread_has_no_subagent_context(cx: &mut TestAppContext) {
1120        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1121
1122        let thread_id = session_id("regular-thread");
1123        let thread = make_thread(
1124            "Regular Thread",
1125            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1126        );
1127
1128        database
1129            .save_thread(thread_id.clone(), thread, PathList::default())
1130            .await
1131            .unwrap();
1132
1133        let loaded = database
1134            .load_thread(thread_id)
1135            .await
1136            .unwrap()
1137            .expect("thread should exist");
1138
1139        assert!(
1140            loaded.subagent_context.is_none(),
1141            "Regular threads should have no subagent_context"
1142        );
1143    }
1144
1145    #[gpui::test]
1146    async fn test_folder_paths_roundtrip(cx: &mut TestAppContext) {
1147        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1148
1149        let thread_id = session_id("folder-thread");
1150        let thread = make_thread(
1151            "Folder Thread",
1152            Utc.with_ymd_and_hms(2024, 6, 15, 12, 0, 0).unwrap(),
1153        );
1154
1155        let folder_paths = PathList::new(&[
1156            std::path::PathBuf::from("/home/user/project-a"),
1157            std::path::PathBuf::from("/home/user/project-b"),
1158        ]);
1159
1160        database
1161            .save_thread(thread_id.clone(), thread, folder_paths.clone())
1162            .await
1163            .unwrap();
1164
1165        let threads = database.list_threads().await.unwrap();
1166        assert_eq!(threads.len(), 1);
1167    }
1168
1169    #[gpui::test]
1170    async fn test_folder_paths_empty_when_not_set(cx: &mut TestAppContext) {
1171        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1172
1173        let thread_id = session_id("no-folder-thread");
1174        let thread = make_thread(
1175            "No Folder Thread",
1176            Utc.with_ymd_and_hms(2024, 6, 15, 12, 0, 0).unwrap(),
1177        );
1178
1179        database
1180            .save_thread(thread_id.clone(), thread, PathList::default())
1181            .await
1182            .unwrap();
1183
1184        let threads = database.list_threads().await.unwrap();
1185        assert_eq!(threads.len(), 1);
1186    }
1187
1188    #[test]
1189    fn test_scroll_position_defaults_to_none() {
1190        let json = r#"{
1191            "title": "Old Thread",
1192            "messages": [],
1193            "updated_at": "2024-01-01T00:00:00Z"
1194        }"#;
1195
1196        let db_thread: DbThread = serde_json::from_str(json).expect("Failed to deserialize");
1197
1198        assert!(
1199            db_thread.ui_scroll_position.is_none(),
1200            "Legacy threads without scroll_position field should default to None"
1201        );
1202    }
1203
1204    #[gpui::test]
1205    async fn test_scroll_position_roundtrips_through_save_load(cx: &mut TestAppContext) {
1206        let database = ThreadsDatabase::new(cx.executor()).unwrap();
1207
1208        let thread_id = session_id("thread-with-scroll");
1209
1210        let mut thread = make_thread(
1211            "Thread With Scroll",
1212            Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(),
1213        );
1214        thread.ui_scroll_position = Some(SerializedScrollPosition {
1215            item_ix: 42,
1216            offset_in_item: 13.5,
1217        });
1218
1219        database
1220            .save_thread(thread_id.clone(), thread, PathList::default())
1221            .await
1222            .unwrap();
1223
1224        let loaded = database
1225            .load_thread(thread_id)
1226            .await
1227            .unwrap()
1228            .expect("thread should exist");
1229
1230        let scroll = loaded
1231            .ui_scroll_position
1232            .expect("scroll_position should be restored");
1233        assert_eq!(scroll.item_ix, 42);
1234        assert!((scroll.offset_in_item - 13.5).abs() < f32::EPSILON);
1235    }
1236}
1237