tenant.openagents/omega

1use anyhow::Result;
2use db::{
3    query,
4    sqlez::{
5        bindable::Column, domain::Domain, statement::Statement,
6        thread_safe_connection::ThreadSafeConnection,
7    },
8    sqlez_macros::sql,
9};
10use serde::{Deserialize, Serialize};
11use time::OffsetDateTime;
12
13#[cfg(test)]
14#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
15pub(crate) struct SerializedCommandInvocation {
16    pub(crate) command_name: String,
17    pub(crate) user_query: String,
18    pub(crate) last_invoked: OffsetDateTime,
19}
20
21#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
22pub(crate) struct SerializedCommandUsage {
23    pub(crate) command_name: String,
24    pub(crate) invocations: u16,
25    pub(crate) last_invoked: OffsetDateTime,
26}
27
28impl Column for SerializedCommandUsage {
29    fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
30        let (command_name, next_index): (String, i32) = Column::column(statement, start_index)?;
31        let (invocations, next_index): (u16, i32) = Column::column(statement, next_index)?;
32        let (last_invoked_raw, next_index): (i64, i32) = Column::column(statement, next_index)?;
33
34        let usage = Self {
35            command_name,
36            invocations,
37            last_invoked: OffsetDateTime::from_unix_timestamp(last_invoked_raw)?,
38        };
39        Ok((usage, next_index))
40    }
41}
42
43#[cfg(test)]
44impl Column for SerializedCommandInvocation {
45    fn column(statement: &mut Statement, start_index: i32) -> Result<(Self, i32)> {
46        let (command_name, next_index): (String, i32) = Column::column(statement, start_index)?;
47        let (user_query, next_index): (String, i32) = Column::column(statement, next_index)?;
48        let (last_invoked_raw, next_index): (i64, i32) = Column::column(statement, next_index)?;
49        let command_invocation = Self {
50            command_name,
51            user_query,
52            last_invoked: OffsetDateTime::from_unix_timestamp(last_invoked_raw)?,
53        };
54        Ok((command_invocation, next_index))
55    }
56}
57
58pub struct CommandPaletteDB(ThreadSafeConnection);
59
60impl Domain for CommandPaletteDB {
61    const NAME: &str = stringify!(CommandPaletteDB);
62    const MIGRATIONS: &[&str] = &[sql!(
63        CREATE TABLE IF NOT EXISTS command_invocations(
64            id INTEGER PRIMARY KEY AUTOINCREMENT,
65            command_name TEXT NOT NULL,
66            user_query TEXT NOT NULL,
67            last_invoked INTEGER DEFAULT (unixepoch())  NOT NULL
68        ) STRICT;
69    )];
70}
71
72db::static_connection!(CommandPaletteDB, []);
73
74impl CommandPaletteDB {
75    pub async fn write_command_invocation(
76        &self,
77        command_name: impl Into<String>,
78        user_query: impl Into<String>,
79    ) -> Result<()> {
80        let command_name = command_name.into();
81        let user_query = user_query.into();
82        log::debug!(
83            "Writing command invocation: command_name={command_name}, user_query={user_query}"
84        );
85        self.write_command_invocation_internal(command_name, user_query)
86            .await
87    }
88
89    #[cfg(test)]
90    query! {
91        pub(crate) fn get_last_invoked(command: &str) -> Result<Option<SerializedCommandInvocation>> {
92            SELECT
93            command_name,
94            user_query,
95            last_invoked FROM command_invocations
96            WHERE command_name=(?)
97            ORDER BY last_invoked DESC
98            LIMIT 1
99        }
100    }
101
102    #[cfg(test)]
103    query! {
104        pub(crate) async fn clear_all() -> Result<()> {
105            DELETE FROM command_invocations
106        }
107    }
108
109    query! {
110        pub fn get_command_usage(command: &str) -> Result<Option<SerializedCommandUsage>> {
111            SELECT command_name, COUNT(1), MAX(last_invoked)
112            FROM command_invocations
113            WHERE command_name=(?)
114            GROUP BY command_name
115        }
116    }
117
118    query! {
119        async fn write_command_invocation_internal(command_name: String, user_query: String) -> Result<()> {
120            INSERT INTO command_invocations (command_name, user_query) VALUES ((?), (?));
121            DELETE FROM command_invocations WHERE id IN (SELECT MIN(id) FROM command_invocations HAVING COUNT(1) > 1000);
122        }
123    }
124
125    query! {
126        pub fn list_commands_used() -> Result<Vec<SerializedCommandUsage>> {
127            SELECT command_name, COUNT(1), MAX(last_invoked)
128            FROM command_invocations
129            GROUP BY command_name
130            ORDER BY COUNT(1) DESC
131        }
132    }
133
134    query! {
135        pub fn list_recent_queries() -> Result<Vec<String>> {
136            SELECT user_query
137            FROM command_invocations
138            WHERE user_query != ""
139            GROUP BY user_query
140            ORDER BY MAX(last_invoked) ASC
141        }
142    }
143}
144
145#[cfg(test)]
146mod tests {
147
148    use crate::persistence::{CommandPaletteDB, SerializedCommandUsage};
149
150    #[gpui::test]
151    async fn test_saves_and_retrieves_command_invocation() {
152        let db =
153            CommandPaletteDB::open_test_db("test_saves_and_retrieves_command_invocation").await;
154
155        let retrieved_cmd = db.get_last_invoked("editor: backspace").unwrap();
156
157        assert!(retrieved_cmd.is_none());
158
159        db.write_command_invocation("editor: backspace", "")
160            .await
161            .unwrap();
162
163        let retrieved_cmd = db.get_last_invoked("editor: backspace").unwrap();
164
165        assert!(retrieved_cmd.is_some());
166        let retrieved_cmd = retrieved_cmd.expect("is some");
167        assert_eq!(retrieved_cmd.command_name, "editor: backspace".to_string());
168        assert_eq!(retrieved_cmd.user_query, "".to_string());
169    }
170
171    #[gpui::test]
172    async fn test_gets_usage_history() {
173        let db = CommandPaletteDB::open_test_db("test_gets_usage_history").await;
174        db.write_command_invocation("go to line: toggle", "200")
175            .await
176            .unwrap();
177        db.write_command_invocation("go to line: toggle", "201")
178            .await
179            .unwrap();
180
181        let retrieved_cmd = db.get_last_invoked("go to line: toggle").unwrap();
182
183        assert!(retrieved_cmd.is_some());
184        let retrieved_cmd = retrieved_cmd.expect("is some");
185
186        let command_usage = db.get_command_usage("go to line: toggle").unwrap();
187
188        assert!(command_usage.is_some());
189        let command_usage: SerializedCommandUsage = command_usage.expect("is some");
190
191        assert_eq!(command_usage.command_name, "go to line: toggle");
192        assert_eq!(command_usage.invocations, 2);
193        assert_eq!(command_usage.last_invoked, retrieved_cmd.last_invoked);
194    }
195
196    #[gpui::test]
197    async fn test_lists_ordered_by_usage() {
198        let db = CommandPaletteDB::open_test_db("test_lists_ordered_by_usage").await;
199
200        let empty_commands = db.list_commands_used();
201        match &empty_commands {
202            Ok(_) => (),
203            Err(e) => println!("Error: {:?}", e),
204        }
205        assert!(empty_commands.is_ok());
206        assert_eq!(empty_commands.expect("is ok").len(), 0);
207
208        db.write_command_invocation("go to line: toggle", "200")
209            .await
210            .unwrap();
211        db.write_command_invocation("editor: backspace", "")
212            .await
213            .unwrap();
214        db.write_command_invocation("editor: backspace", "")
215            .await
216            .unwrap();
217
218        let commands = db.list_commands_used();
219
220        assert!(commands.is_ok());
221        let commands = commands.expect("is ok");
222        assert_eq!(commands.len(), 2);
223        assert_eq!(commands.as_slice()[0].command_name, "editor: backspace");
224        assert_eq!(commands.as_slice()[0].invocations, 2);
225        assert_eq!(commands.as_slice()[1].command_name, "go to line: toggle");
226        assert_eq!(commands.as_slice()[1].invocations, 1);
227    }
228
229    #[gpui::test]
230    async fn test_handles_max_invocation_entries() {
231        let db = CommandPaletteDB::open_test_db("test_handles_max_invocation_entries").await;
232
233        for i in 1..=1001 {
234            db.write_command_invocation("some-command", &i.to_string())
235                .await
236                .unwrap();
237        }
238        let some_command = db.get_command_usage("some-command").unwrap();
239
240        assert!(some_command.is_some());
241        assert_eq!(some_command.expect("is some").invocations, 1000);
242    }
243}
244