1use std::fs;
2use std::io::Read;
3use std::path::Path;
4use std::time::UNIX_EPOCH;
5
6use anyhow::{Context, Result, bail};
7use rusqlite::{Connection, OptionalExtension, params};
8
9const MAX_COMMAND_BYTES: usize = 64 * 1024;
10const MAX_HISTORY_BYTES: u64 = 32 * 1024 * 1024;
11const MAX_HISTORY_ENTRIES: usize = 250_000;
12const SCHEMA_VERSION: i64 = 1;
13
14#[derive(Debug)]
15pub struct Store {
16 connection: Connection,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub struct ImportResult {
21 pub imported: usize,
22 pub skipped: bool,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct HistoryCandidate {
27 pub command: String,
28 pub same_cwd: bool,
29 pub uses: usize,
30}
31
32impl Store {
33 pub fn open(path: &Path) -> Result<Self> {
34 let connection = Connection::open(path)
35 .with_context(|| format!("failed to open database {}", path.display()))?;
36 Self::from_connection(connection)
37 }
38
39 pub fn in_memory() -> Result<Self> {
40 Self::from_connection(Connection::open_in_memory()?)
41 }
42
43 fn from_connection(connection: Connection) -> Result<Self> {
44 connection.busy_timeout(std::time::Duration::from_secs(2))?;
45 connection.pragma_update(None, "journal_mode", "WAL")?;
46 connection.pragma_update(None, "foreign_keys", "ON")?;
47 connection.pragma_update(None, "case_sensitive_like", "ON")?;
48 let schema_version: i64 =
49 connection.pragma_query_value(None, "user_version", |row| row.get(0))?;
50 if schema_version > SCHEMA_VERSION {
51 bail!(
52 "database schema {schema_version} is newer than supported schema {SCHEMA_VERSION}"
53 );
54 }
55 if schema_version == 0 {
56 connection.execute_batch(
57 "
58 CREATE TABLE IF NOT EXISTS command_events (
59 id INTEGER PRIMARY KEY,
60 command TEXT NOT NULL,
61 cwd TEXT NOT NULL,
62 observed_at_ms INTEGER NOT NULL,
63 exit_code INTEGER,
64 source TEXT NOT NULL,
65 source_key TEXT,
66 UNIQUE(source, source_key)
67 );
68 CREATE INDEX IF NOT EXISTS command_events_command_idx
69 ON command_events(command);
70 CREATE INDEX IF NOT EXISTS command_events_cwd_idx
71 ON command_events(cwd);
72
73 CREATE TABLE IF NOT EXISTS history_imports (
74 path TEXT PRIMARY KEY,
75 size_bytes INTEGER NOT NULL,
76 modified_ns INTEGER NOT NULL
77 );
78 PRAGMA user_version = 1;
79 ",
80 )?;
81 }
82 Ok(Self { connection })
83 }
84
85 pub fn record(
86 &self,
87 command: &str,
88 cwd: &str,
89 exit_code: i32,
90 observed_at_ms: i64,
91 session_id: &str,
92 ignore_leading_space: bool,
93 ) -> Result<bool> {
94 if !eligible_command(command, ignore_leading_space) {
95 return Ok(false);
96 }
97 self.connection.execute(
98 "INSERT INTO command_events
99 (command, cwd, observed_at_ms, exit_code, source, source_key)
100 VALUES (?1, ?2, ?3, ?4, ?5, NULL)",
101 params![
102 command,
103 cwd,
104 observed_at_ms,
105 exit_code,
106 format!("session:{session_id}")
107 ],
108 )?;
109 Ok(true)
110 }
111
112 pub fn history_candidates(
113 &self,
114 prefix: &str,
115 cwd: &str,
116 limit: usize,
117 successful_first: bool,
118 ) -> Result<Vec<HistoryCandidate>> {
119 let pattern = format!("{}%", escape_like(prefix));
120 let success_order = i64::from(successful_first);
121 let mut statement = self.connection.prepare(
122 "SELECT
123 command,
124 MAX(CASE WHEN cwd = ?2 THEN 1 ELSE 0 END) AS same_cwd,
125 SUM(CASE WHEN exit_code = 0 THEN 1 ELSE 0 END) AS successes,
126 MAX(observed_at_ms) AS latest,
127 COUNT(*) AS uses
128 FROM command_events
129 WHERE command LIKE ?1 ESCAPE '\\' AND command <> ?3
130 GROUP BY command
131 ORDER BY
132 latest DESC,
133 same_cwd DESC,
134 CASE WHEN ?5 = 1 AND successes > 0 THEN 1 ELSE 0 END DESC,
135 uses DESC,
136 command ASC
137 LIMIT ?4",
138 )?;
139 let rows = statement.query_map(
140 params![pattern, cwd, prefix, limit as i64, success_order],
141 |row| {
142 Ok(HistoryCandidate {
143 command: row.get(0)?,
144 same_cwd: row.get::<_, i64>(1)? != 0,
145 uses: row.get::<_, i64>(4)?.max(0) as usize,
146 })
147 },
148 )?;
149 rows.collect::<rusqlite::Result<Vec<HistoryCandidate>>>()
150 .context("failed to read history candidates")
151 }
152
153 pub fn history_inventory(
154 &self,
155 cwd: &str,
156 limit: usize,
157 successful_first: bool,
158 ) -> Result<Vec<HistoryCandidate>> {
159 let success_order = i64::from(successful_first);
160 let mut statement = self.connection.prepare(
161 "SELECT
162 command,
163 MAX(CASE WHEN cwd = ?1 THEN 1 ELSE 0 END) AS same_cwd,
164 SUM(CASE WHEN exit_code = 0 THEN 1 ELSE 0 END) AS successes,
165 MAX(observed_at_ms) AS latest,
166 COUNT(*) AS uses
167 FROM command_events
168 GROUP BY command
169 ORDER BY
170 latest DESC,
171 same_cwd DESC,
172 CASE WHEN ?3 = 1 AND successes > 0 THEN 1 ELSE 0 END DESC,
173 uses DESC,
174 command ASC
175 LIMIT ?2",
176 )?;
177 let rows = statement.query_map(params![cwd, limit as i64, success_order], |row| {
178 Ok(HistoryCandidate {
179 command: row.get(0)?,
180 same_cwd: row.get::<_, i64>(1)? != 0,
181 uses: row.get::<_, i64>(4)?.max(0) as usize,
182 })
183 })?;
184 rows.collect::<rusqlite::Result<Vec<HistoryCandidate>>>()
185 .context("failed to read fuzzy history inventory")
186 }
187
188 pub fn import_zsh_history(
189 &mut self,
190 path: &Path,
191 ignore_leading_space: bool,
192 ) -> Result<ImportResult> {
193 let canonical = path
194 .canonicalize()
195 .with_context(|| format!("failed to resolve history file {}", path.display()))?;
196 let metadata = fs::metadata(&canonical)
197 .with_context(|| format!("failed to inspect history file {}", canonical.display()))?;
198 if !metadata.is_file() {
199 bail!(
200 "history path is not a regular file: {}",
201 canonical.display()
202 );
203 }
204 if metadata.len() > MAX_HISTORY_BYTES {
205 bail!(
206 "history file exceeds {} MiB: {}",
207 MAX_HISTORY_BYTES / (1024 * 1024),
208 canonical.display()
209 );
210 }
211 let modified_ns = metadata
212 .modified()?
213 .duration_since(UNIX_EPOCH)
214 .unwrap_or_default()
215 .as_nanos()
216 .min(i64::MAX as u128) as i64;
217 let size_bytes = metadata.len().min(i64::MAX as u64) as i64;
218 let source = format!("zsh:{}", canonical.display());
219
220 let previous: Option<(i64, i64)> = self
221 .connection
222 .query_row(
223 "SELECT size_bytes, modified_ns FROM history_imports WHERE path = ?1",
224 params![canonical.to_string_lossy()],
225 |row| Ok((row.get(0)?, row.get(1)?)),
226 )
227 .optional()?;
228 if previous == Some((size_bytes, modified_ns)) {
229 return Ok(ImportResult {
230 imported: 0,
231 skipped: true,
232 });
233 }
234
235 let mut content = Vec::new();
236 fs::File::open(&canonical)
237 .with_context(|| format!("failed to open history file {}", canonical.display()))?
238 .take(MAX_HISTORY_BYTES + 1)
239 .read_to_end(&mut content)
240 .with_context(|| format!("failed to read history file {}", canonical.display()))?;
241 if content.len() as u64 > MAX_HISTORY_BYTES {
242 bail!("history file grew beyond the import limit while reading");
243 }
244 let content = String::from_utf8(content).context("history file is not valid UTF-8")?;
245 let entries = parse_zsh_history(&content, ignore_leading_space)?;
246
247 let transaction = self.connection.transaction()?;
248 transaction.execute(
249 "DELETE FROM command_events WHERE source = ?1",
250 params![source],
251 )?;
252 {
253 let mut insert = transaction.prepare(
254 "INSERT INTO command_events
255 (command, cwd, observed_at_ms, exit_code, source, source_key)
256 VALUES (?1, '', ?2, NULL, ?3, ?4)",
257 )?;
258 for (index, entry) in entries.iter().enumerate() {
259 insert.execute(params![
260 entry.command,
261 entry.observed_at_ms,
262 source,
263 index.to_string()
264 ])?;
265 }
266 }
267 transaction.execute(
268 "INSERT INTO history_imports (path, size_bytes, modified_ns)
269 VALUES (?1, ?2, ?3)
270 ON CONFLICT(path) DO UPDATE SET
271 size_bytes = excluded.size_bytes,
272 modified_ns = excluded.modified_ns",
273 params![canonical.to_string_lossy(), size_bytes, modified_ns],
274 )?;
275 transaction.commit()?;
276
277 Ok(ImportResult {
278 imported: entries.len(),
279 skipped: false,
280 })
281 }
282}
283
284#[derive(Debug, PartialEq, Eq)]
285struct HistoryEntry {
286 command: String,
287 observed_at_ms: i64,
288}
289
290fn parse_zsh_history(content: &str, ignore_leading_space: bool) -> Result<Vec<HistoryEntry>> {
291 let mut entries = Vec::new();
292 let mut lines = content.lines().enumerate();
293 while let Some((index, line)) = lines.next() {
294 if line.ends_with('\\') {
295 let mut continued = true;
296 while continued {
297 let Some((_, continuation)) = lines.next() else {
298 break;
299 };
300 continued = continuation.ends_with('\\');
301 }
302 continue;
303 }
304
305 let (command, observed_at_ms) = parse_zsh_history_line(line, index as i64);
306 if eligible_command(command, ignore_leading_space) {
307 entries.push(HistoryEntry {
308 command: command.to_owned(),
309 observed_at_ms,
310 });
311 if entries.len() > MAX_HISTORY_ENTRIES {
312 bail!("history file exceeds {MAX_HISTORY_ENTRIES} entries");
313 }
314 }
315 }
316 Ok(entries)
317}
318
319fn parse_zsh_history_line(line: &str, fallback_order: i64) -> (&str, i64) {
320 let Some(metadata_and_command) = line.strip_prefix(": ") else {
321 return (line, fallback_order);
322 };
323 let Some((metadata, command)) = metadata_and_command.split_once(';') else {
324 return (line, fallback_order);
325 };
326 let timestamp = metadata
327 .split(':')
328 .next()
329 .and_then(|value| value.parse::<i64>().ok())
330 .map(|seconds| seconds.saturating_mul(1_000))
331 .unwrap_or(fallback_order);
332 (command, timestamp)
333}
334
335fn eligible_command(command: &str, ignore_leading_space: bool) -> bool {
336 if command.is_empty() || command.len() > MAX_COMMAND_BYTES {
337 return false;
338 }
339 if ignore_leading_space && command.starts_with(char::is_whitespace) {
340 return false;
341 }
342 !command.chars().any(|character| {
343 character == '\0' || (character.is_control() && !matches!(character, '\t'))
344 })
345}
346
347fn escape_like(value: &str) -> String {
348 let mut escaped = String::with_capacity(value.len());
349 for character in value.chars() {
350 if matches!(character, '%' | '_' | '\\') {
351 escaped.push('\\');
352 }
353 escaped.push(character);
354 }
355 escaped
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use tempfile::tempdir;
362
363 #[test]
364 fn ranks_recent_history_before_directory_success_and_usage() {
365 let store = Store::in_memory().unwrap();
366 for observed_at_ms in 10..20 {
367 store
368 .record("git switch main", "/repo", 0, observed_at_ms, "one", true)
369 .unwrap();
370 }
371 store
372 .record("git status", "/other", 0, 300, "one", true)
373 .unwrap();
374 store
375 .record("git stash", "/repo", 1, 200, "one", true)
376 .unwrap();
377
378 let candidates = store
379 .history_candidates("git s", "/repo", 10, true)
380 .unwrap();
381 assert!(!candidates[0].same_cwd);
382 assert!(candidates[1].same_cwd);
383 assert_eq!(
384 candidates
385 .into_iter()
386 .map(|candidate| candidate.command)
387 .collect::<Vec<_>>(),
388 vec!["git status", "git stash", "git switch main"]
389 );
390 }
391
392 #[test]
393 fn fuzzy_inventory_is_not_prefix_limited() {
394 let store = Store::in_memory().unwrap();
395 store
396 .record("cargo test", "/repo", 0, 100, "one", true)
397 .unwrap();
398 store
399 .record("git status", "/other", 0, 200, "one", true)
400 .unwrap();
401
402 let candidates = store.history_inventory("/repo", 10, true).unwrap();
403 assert_eq!(candidates[0].command, "git status");
404 assert!(
405 candidates
406 .iter()
407 .any(|candidate| candidate.command == "cargo test")
408 );
409 }
410
411 #[test]
412 fn parses_extended_zsh_history() {
413 let entries =
414 parse_zsh_history(": 1700000000:4;git status\nplain command\n", true).unwrap();
415 assert_eq!(entries[0].command, "git status");
416 assert_eq!(entries[0].observed_at_ms, 1_700_000_000_000);
417 assert_eq!(entries[1].command, "plain command");
418 }
419
420 #[test]
421 fn escapes_like_metacharacters() {
422 assert_eq!(escape_like("echo 100%_done"), "echo 100\\%\\_done");
423 }
424
425 #[test]
426 fn ignores_leading_space_when_configured() {
427 assert!(!eligible_command(" secret command", true));
428 assert!(eligible_command(" secret command", false));
429 }
430
431 #[test]
432 fn imports_a_history_file_once_until_it_changes() {
433 let directory = tempdir().unwrap();
434 let history = directory.path().join("history");
435 fs::write(
436 &history,
437 ": 1700000000:0;cargo test\n: 1700000001:0;cargo check\n",
438 )
439 .unwrap();
440 let mut store = Store::in_memory().unwrap();
441
442 let first = store.import_zsh_history(&history, true).unwrap();
443 assert_eq!(first.imported, 2);
444 assert!(!first.skipped);
445
446 let second = store.import_zsh_history(&history, true).unwrap();
447 assert!(second.skipped);
448 assert_eq!(
449 store
450 .history_candidates("cargo ", "/repo", 10, true)
451 .unwrap()
452 .into_iter()
453 .map(|candidate| candidate.command)
454 .collect::<Vec<_>>(),
455 vec!["cargo check", "cargo test"]
456 );
457 }
458
459 #[test]
460 fn skips_multiline_history_entries_conservatively() {
461 let entries = parse_zsh_history(
462 ": 1700000000:0;echo first \\\n+continued\n: 1700000001:0;git status\n",
463 true,
464 )
465 .unwrap();
466 assert_eq!(entries.len(), 1);
467 assert_eq!(entries[0].command, "git status");
468 }
469
470 #[test]
471 fn rejects_invalid_utf8_history() {
472 let directory = tempdir().unwrap();
473 let history = directory.path().join("history");
474 fs::write(&history, [0xff, 0xfe]).unwrap();
475 let mut store = Store::in_memory().unwrap();
476 assert!(store.import_zsh_history(&history, true).is_err());
477 }
478
479 #[test]
480 fn rejects_newer_database_schema() {
481 let connection = Connection::open_in_memory().unwrap();
482 connection.pragma_update(None, "user_version", 2).unwrap();
483 assert!(Store::from_connection(connection).is_err());
484 }
485}