1094 lines · 37.5 KB · Rust
RawDownload
use crate::{AgentMessage, AgentMessageContent, UserMessage, UserMessageContent};use acp_thread::UserMessageId;use agent_client_protocol::schema as acp;use agent_settings::AgentProfileId;use anyhow::{Result, anyhow};use chrono::{DateTime, Utc};use collections::{HashMap, IndexMap};use futures::{FutureExt, future::Shared};use gpui::{BackgroundExecutor, Global, Task};use indoc::indoc;use language_model::Speed;use parking_lot::Mutex;use serde::{Deserialize, Serialize};use sqlez::{    bindable::{Bind, Column},    connection::Connection,    statement::Statement,};use std::{io::ErrorKind, path::PathBuf, sync::Arc};use ui::{App, SharedString};use util::path_list::PathList;use zed_env_vars::ZED_STATELESS; pub type DbMessage = crate::Message;pub type DbSummary = crate::legacy_thread::DetailedSummaryState;pub type DbLanguageModel = crate::legacy_thread::SerializedLanguageModel; #[derive(Debug, Clone)]pub struct DbThreadMetadata {    pub id: acp::SessionId,    pub parent_session_id: Option<acp::SessionId>,    pub title: SharedString,    pub updated_at: DateTime<Utc>,    pub created_at: Option<DateTime<Utc>>,    /// The workspace folder paths this thread was created against, sorted    /// lexicographically. Used for grouping threads by project in the sidebar.    pub folder_paths: PathList,} impl From<&DbThreadMetadata> for acp_thread::AgentSessionInfo {    fn from(meta: &DbThreadMetadata) -> Self {        Self {            session_id: meta.id.clone(),            work_dirs: Some(meta.folder_paths.clone()),            title: Some(meta.title.clone()),            updated_at: Some(meta.updated_at),            created_at: meta.created_at,            meta: None,        }    }} #[derive(Debug, Serialize, Deserialize)]pub struct DbThread {    pub title: SharedString,    pub messages: Vec<Arc<DbMessage>>,    pub updated_at: DateTime<Utc>,    #[serde(default)]    pub detailed_summary: Option<SharedString>,    #[serde(default)]    pub initial_project_snapshot: Option<Arc<crate::ProjectSnapshot>>,    #[serde(default)]    pub cumulative_token_usage: language_model::TokenUsage,    #[serde(default)]    pub request_token_usage: HashMap<acp_thread::UserMessageId, language_model::TokenUsage>,    #[serde(default)]    pub model: Option<DbLanguageModel>,    #[serde(default)]    pub profile: Option<AgentProfileId>,    #[serde(default)]    pub imported: bool,    #[serde(default)]    pub subagent_context: Option<crate::SubagentContext>,    #[serde(default)]    pub speed: Option<Speed>,    #[serde(default)]    pub thinking_enabled: bool,    #[serde(default)]    pub thinking_effort: Option<String>,    #[serde(default)]    pub draft_prompt: Option<Vec<acp::ContentBlock>>,    #[serde(default)]    pub ui_scroll_position: Option<SerializedScrollPosition>,    #[serde(default)]    pub sandboxed_terminal_temp_dir: Option<PathBuf>,} #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]pub struct SerializedScrollPosition {    pub item_ix: usize,    pub offset_in_item: f32,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct SharedThread {    pub title: SharedString,    pub messages: Vec<Arc<DbMessage>>,    pub updated_at: DateTime<Utc>,    #[serde(default)]    pub model: Option<DbLanguageModel>,    pub version: String,} impl SharedThread {    pub const VERSION: &'static str = "1.0.0";     pub fn from_db_thread(thread: &DbThread) -> Self {        Self {            title: thread.title.clone(),            messages: thread.messages.clone(),            updated_at: thread.updated_at,            model: thread.model.clone(),            version: Self::VERSION.to_string(),        }    }     pub fn to_db_thread(self) -> DbThread {        DbThread {            title: format!("🔗 {}", self.title).into(),            messages: self.messages,            updated_at: self.updated_at,            detailed_summary: None,            initial_project_snapshot: None,            cumulative_token_usage: Default::default(),            request_token_usage: Default::default(),            model: self.model,            profile: None,            imported: true,            subagent_context: None,            speed: None,            thinking_enabled: false,            thinking_effort: None,            draft_prompt: None,            ui_scroll_position: None,            sandboxed_terminal_temp_dir: None,        }    }     pub fn to_bytes(&self) -> Result<Vec<u8>> {        const COMPRESSION_LEVEL: i32 = 3;        let json = serde_json::to_vec(self)?;        let compressed = zstd::encode_all(json.as_slice(), COMPRESSION_LEVEL)?;        Ok(compressed)    }     pub fn from_bytes(data: &[u8]) -> Result<Self> {        let decompressed = zstd::decode_all(data)?;        Ok(serde_json::from_slice(&decompressed)?)    }} impl DbThread {    pub const VERSION: &'static str = "0.3.0";     pub fn from_json(json: &[u8]) -> Result<Self> {        let saved_thread_json = serde_json::from_slice::<serde_json::Value>(json)?;        match saved_thread_json.get("version") {            Some(serde_json::Value::String(version)) => match version.as_str() {                Self::VERSION => Ok(serde_json::from_value(saved_thread_json)?),                _ => Self::upgrade_from_agent_1(crate::legacy_thread::SerializedThread::from_json(                    json,                )?),            },            _ => {                Self::upgrade_from_agent_1(crate::legacy_thread::SerializedThread::from_json(json)?)            }        }    }     fn upgrade_from_agent_1(thread: crate::legacy_thread::SerializedThread) -> Result<Self> {        let mut messages = Vec::new();        let mut request_token_usage = HashMap::default();         let mut last_user_message_id = None;        for (ix, msg) in thread.messages.into_iter().enumerate() {            let message = match msg.role {                language_model::Role::User => {                    let mut content = Vec::new();                     // Convert segments to content                    for segment in msg.segments {                        match segment {                            crate::legacy_thread::SerializedMessageSegment::Text { text } => {                                content.push(UserMessageContent::Text(text));                            }                            crate::legacy_thread::SerializedMessageSegment::Thinking {                                text,                                ..                            } => {                                // User messages don't have thinking segments, but handle gracefully                                content.push(UserMessageContent::Text(text));                            }                            crate::legacy_thread::SerializedMessageSegment::RedactedThinking {                                ..                            } => {                                // User messages don't have redacted thinking, skip.                            }                        }                    }                     // If no content was added, add context as text if available                    if content.is_empty() && !msg.context.is_empty() {                        content.push(UserMessageContent::Text(msg.context));                    }                     let id = UserMessageId::new();                    last_user_message_id = Some(id.clone());                     crate::Message::User(UserMessage {                        // MessageId from old format can't be meaningfully converted, so generate a new one                        id,                        content: Arc::from(content),                    })                }                language_model::Role::Assistant => {                    let mut content = Vec::new();                     // Convert segments to content                    for segment in msg.segments {                        match segment {                            crate::legacy_thread::SerializedMessageSegment::Text { text } => {                                content.push(AgentMessageContent::Text(text));                            }                            crate::legacy_thread::SerializedMessageSegment::Thinking {                                text,                                signature,                            } => {                                content.push(AgentMessageContent::Thinking { text, signature });                            }                            crate::legacy_thread::SerializedMessageSegment::RedactedThinking {                                data,                            } => {                                content.push(AgentMessageContent::RedactedThinking(data));                            }                        }                    }                     // Convert tool uses                    let mut tool_names_by_id = HashMap::default();                    for tool_use in msg.tool_uses {                        tool_names_by_id.insert(tool_use.id.clone(), tool_use.name.clone());                        content.push(AgentMessageContent::ToolUse(                            language_model::LanguageModelToolUse {                                id: tool_use.id,                                name: tool_use.name.into(),                                raw_input: serde_json::to_string(&tool_use.input)                                    .unwrap_or_default(),                                input: tool_use.input,                                is_input_complete: true,                                thought_signature: None,                            },                        ));                    }                     // Convert tool results                    let mut tool_results = IndexMap::default();                    for tool_result in msg.tool_results {                        let name = tool_names_by_id                            .remove(&tool_result.tool_use_id)                            .unwrap_or_else(|| SharedString::from("unknown"));                        tool_results.insert(                            tool_result.tool_use_id.clone(),                            language_model::LanguageModelToolResult {                                tool_use_id: tool_result.tool_use_id,                                tool_name: name.into(),                                is_error: tool_result.is_error,                                content: vec![tool_result.content],                                output: tool_result.output,                            },                        );                    }                     if let Some(last_user_message_id) = &last_user_message_id                        && let Some(token_usage) = thread.request_token_usage.get(ix).copied()                    {                        request_token_usage.insert(last_user_message_id.clone(), token_usage);                    }                     crate::Message::Agent(AgentMessage {                        content,                        tool_results,                        reasoning_details: None,                    })                }                language_model::Role::System => {                    // Skip system messages as they're not supported in the new format                    continue;                }            };             messages.push(Arc::new(message));        }         Ok(Self {            title: thread.summary,            messages,            updated_at: thread.updated_at,            detailed_summary: match thread.detailed_summary_state {                crate::legacy_thread::DetailedSummaryState::NotGenerated                | crate::legacy_thread::DetailedSummaryState::Generating => None,                crate::legacy_thread::DetailedSummaryState::Generated { text, .. } => Some(text),            },            initial_project_snapshot: thread.initial_project_snapshot,            cumulative_token_usage: thread.cumulative_token_usage,            request_token_usage,            model: thread.model,            profile: thread.profile,            imported: false,            subagent_context: None,            speed: None,            thinking_enabled: false,            thinking_effort: None,            draft_prompt: None,            ui_scroll_position: None,            sandboxed_terminal_temp_dir: None,        })    }} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]pub enum DataType {    #[serde(rename = "json")]    Json,    #[serde(rename = "zstd")]    Zstd,} impl Bind for DataType {    fn bind(&self, statement: &Statement, start_index: i32) -> Result<i32> {        let value = match self {            DataType::Json => "json",            DataType::Zstd => "zstd",        };        value.bind(statement, start_index)    }} impl Column for DataType {    fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {        let (value, next_index) = String::column(statement, start_index)?;        let data_type = match value.as_str() {            "json" => DataType::Json,            "zstd" => DataType::Zstd,            _ => anyhow::bail!("Unknown data type: {}", value),        };        Ok((data_type, next_index))    }} pub(crate) struct ThreadsDatabase {    executor: BackgroundExecutor,    connection: Arc<Mutex<Connection>>,} struct GlobalThreadsDatabase(Shared<Task<Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>); impl Global for GlobalThreadsDatabase {} impl ThreadsDatabase {    pub fn connect(cx: &mut App) -> Shared<Task<Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>> {        if cx.has_global::<GlobalThreadsDatabase>() {            return cx.global::<GlobalThreadsDatabase>().0.clone();        }        let executor = cx.background_executor().clone();        let task = executor            .spawn({                let executor = executor.clone();                async move {                    match ThreadsDatabase::new(executor) {                        Ok(db) => Ok(Arc::new(db)),                        Err(err) => Err(Arc::new(err)),                    }                }            })            .shared();         cx.set_global(GlobalThreadsDatabase(task.clone()));        task    }     pub fn new(executor: BackgroundExecutor) -> Result<Self> {        let connection = if *ZED_STATELESS {            Connection::open_memory(Some("THREAD_FALLBACK_DB"))        } else if cfg!(any(feature = "test-support", test)) {            // rust stores the name of the test on the current thread.            // We use this to automatically create a database that will            // be shared within the test (for the test_retrieve_old_thread)            // but not with concurrent tests.            let thread = std::thread::current();            let test_name = thread.name();            Connection::open_memory(Some(&format!(                "THREAD_FALLBACK_{}",                test_name.unwrap_or_default()            )))        } else {            let threads_dir = paths::data_dir().join("threads");            std::fs::create_dir_all(&threads_dir)?;            let sqlite_path = threads_dir.join("threads.db");            Connection::open_file(&sqlite_path.to_string_lossy())        };         connection.exec(indoc! {"            CREATE TABLE IF NOT EXISTS threads (                id TEXT PRIMARY KEY,                summary TEXT NOT NULL,                updated_at TEXT NOT NULL,                data_type TEXT NOT NULL,                data BLOB NOT NULL            )        "})?()        .map_err(|e| anyhow!("Failed to create threads table: {}", e))?;         if let Ok(mut s) = connection.exec(indoc! {"            ALTER TABLE threads ADD COLUMN parent_id TEXT        "})        {            s().ok();        }         if let Ok(mut s) = connection.exec(indoc! {"            ALTER TABLE threads ADD COLUMN folder_paths TEXT;            ALTER TABLE threads ADD COLUMN folder_paths_order TEXT;        "})        {            s().ok();        }         if let Ok(mut s) = connection.exec(indoc! {"            ALTER TABLE threads ADD COLUMN created_at TEXT;        "})        {            if s().is_ok() {                connection.exec(indoc! {"                    UPDATE threads SET created_at = updated_at WHERE created_at IS NULL                "})?()?;            }        }         let db = Self {            executor,            connection: Arc::new(Mutex::new(connection)),        };         Ok(db)    }     fn save_thread_sync(        connection: &Arc<Mutex<Connection>>,        id: acp::SessionId,        thread: DbThread,        folder_paths: &PathList,    ) -> Result<()> {        const COMPRESSION_LEVEL: i32 = 3;         #[derive(Serialize)]        struct SerializedThread {            #[serde(flatten)]            thread: DbThread,            version: &'static str,        }         let title = thread.title.to_string();        let updated_at = thread.updated_at.to_rfc3339();        let parent_id = thread            .subagent_context            .as_ref()            .map(|ctx| ctx.parent_thread_id.0.clone());        let serialized_folder_paths = folder_paths.serialize();        let (folder_paths_str, folder_paths_order_str): (Option<String>, Option<String>) =            if folder_paths.is_empty() {                (None, None)            } else {                (                    Some(serialized_folder_paths.paths),                    Some(serialized_folder_paths.order),                )            };        let json_data = serde_json::to_string(&SerializedThread {            thread,            version: DbThread::VERSION,        })?;         let connection = connection.lock();         let compressed = zstd::encode_all(json_data.as_bytes(), COMPRESSION_LEVEL)?;        let data_type = DataType::Zstd;        let data = compressed;         // Use the thread's updated_at as created_at for new threads.        // This ensures the creation time reflects when the thread was conceptually        // created, not when it was saved to the database.        let created_at = updated_at.clone();         let mut insert = connection.exec_bound::<(Arc<str>, Option<Arc<str>>, Option<String>, Option<String>, String, String, DataType, Vec<u8>, String)>(indoc! {"            INSERT INTO threads (id, parent_id, folder_paths, folder_paths_order, summary, updated_at, data_type, data, created_at)            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)            ON CONFLICT(id) DO UPDATE SET                parent_id = excluded.parent_id,                folder_paths = excluded.folder_paths,                folder_paths_order = excluded.folder_paths_order,

Lines 1–500 of 1094.