1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use dashmap::DashMap;
6use dk_core::{
7 CallEdge, Error, RawCallEdge, RepoId, Result, Symbol, SymbolId,
8};
9use sqlx::postgres::PgPool;
10use tokio::sync::RwLock;
11use uuid::Uuid;
12
13use crate::changeset::ChangesetStore;
14use crate::git::GitRepository;
15use crate::pipeline::PipelineStore;
16use crate::graph::{
17 CallGraphStore, DependencyStore, SearchIndex, SymbolStore, TypeInfoStore,
18};
19use crate::parser::ParserRegistry;
20use crate::workspace::cache::{NoOpCache, WorkspaceCache};
21use crate::workspace::session_manager::WorkspaceManager;
22
23#[derive(Debug, Clone)]
27pub struct CodebaseSummary {
28 pub languages: Vec<String>,
29 pub total_symbols: u64,
30 pub total_files: u64,
31}
32
33pub struct Engine {
40 pub db: PgPool,
41 pub search_index: Arc<RwLock<SearchIndex>>,
42 pub parser: Arc<ParserRegistry>,
43 pub storage_path: PathBuf,
44 symbol_store: SymbolStore,
45 call_graph_store: CallGraphStore,
46 #[allow(dead_code)]
47 dep_store: DependencyStore,
48 type_info_store: TypeInfoStore,
49 changeset_store: ChangesetStore,
50 pipeline_store: PipelineStore,
51 workspace_manager: Arc<WorkspaceManager>,
52 repo_locks: DashMap<RepoId, Arc<RwLock<()>>>,
53}
54
55impl Engine {
56 pub fn new(storage_path: PathBuf, db: PgPool) -> Result<Self> {
64 Self::with_cache(storage_path, db, Arc::new(NoOpCache))
65 }
66
67 pub fn with_cache(
73 storage_path: PathBuf,
74 db: PgPool,
75 cache: Arc<dyn WorkspaceCache>,
76 ) -> Result<Self> {
77 let search_index = SearchIndex::open(&storage_path.join("search_index"))?;
78 let parser = ParserRegistry::new();
79 let symbol_store = SymbolStore::new(db.clone());
80 let call_graph_store = CallGraphStore::new(db.clone());
81 let dep_store = DependencyStore::new(db.clone());
82 let type_info_store = TypeInfoStore::new(db.clone());
83 let changeset_store = ChangesetStore::new(db.clone());
84 let pipeline_store = PipelineStore::new(db.clone());
85 let workspace_manager = Arc::new(WorkspaceManager::with_cache(db.clone(), cache));
86
87 Ok(Self {
88 db,
89 search_index: Arc::new(RwLock::new(search_index)),
90 parser: Arc::new(parser),
91 storage_path,
92 symbol_store,
93 call_graph_store,
94 dep_store,
95 type_info_store,
96 changeset_store,
97 pipeline_store,
98 workspace_manager,
99 repo_locks: DashMap::new(),
100 })
101 }
102
103 pub fn symbol_store(&self) -> &SymbolStore {
105 &self.symbol_store
106 }
107
108 pub fn changeset_store(&self) -> &ChangesetStore {
110 &self.changeset_store
111 }
112
113 pub fn pipeline_store(&self) -> &PipelineStore {
115 &self.pipeline_store
116 }
117
118 pub fn workspace_manager(&self) -> &WorkspaceManager {
120 &self.workspace_manager
121 }
122
123 pub fn workspace_manager_arc(&self) -> Arc<WorkspaceManager> {
125 Arc::clone(&self.workspace_manager)
126 }
127
128 pub fn call_graph_store(&self) -> &CallGraphStore {
130 &self.call_graph_store
131 }
132
133 pub fn dep_store(&self) -> &DependencyStore {
135 &self.dep_store
136 }
137
138 pub fn parser(&self) -> &ParserRegistry {
140 &self.parser
141 }
142
143 pub fn repo_lock(&self, repo_id: RepoId) -> Arc<RwLock<()>> {
147 self.repo_locks
148 .entry(repo_id)
149 .or_insert_with(|| Arc::new(RwLock::new(())))
150 .clone()
151 }
152
153 pub fn remove_repo_lock(&self, repo_id: RepoId) {
155 self.repo_locks.remove(&repo_id);
156 }
157
158 pub fn spawn_gc_loop(
168 &self,
169 tick: std::time::Duration,
170 idle_ttl: std::time::Duration,
171 max_ttl: std::time::Duration,
172 stranded_ttl: std::time::Duration,
173 ) -> tokio::task::JoinHandle<()> {
174 let mgr = Arc::clone(&self.workspace_manager);
175 let db = self.db.clone();
176 tokio::spawn(async move {
177 let mut interval = tokio::time::interval(tick);
178 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
179 loop {
180 interval.tick().await;
181 let evicted = mgr.gc_expired_sessions_async(idle_ttl, max_ttl).await;
183 if !evicted.is_empty() {
184 tracing::info!(count = evicted.len(), "gc: evicted expired sessions");
185 }
186 match mgr.sweep_stranded(stranded_ttl).await {
188 Ok(n) if n > 0 => tracing::info!(count = n, "gc: abandoned stranded sessions"),
189 Ok(_) => {}
190 Err(e) => tracing::warn!("gc: sweep_stranded error: {e}"),
191 }
192 let active: i64 = sqlx::query_scalar(
194 "SELECT COUNT(*) FROM session_workspaces \
195 WHERE stranded_at IS NOT NULL AND abandoned_at IS NULL",
196 )
197 .fetch_one(&db)
198 .await
199 .unwrap_or(0);
200 crate::metrics::set_workspace_stranded_active(active);
201 }
202 })
203 }
204
205 pub async fn create_repo(&self, name: &str) -> Result<RepoId> {
213 let repo_id = Uuid::new_v4();
214 let repo_path = self.storage_path.join("repos").join(repo_id.to_string());
215
216 GitRepository::init(&repo_path)?;
217
218 sqlx::query(
219 r#"
220 INSERT INTO repositories (id, name, path)
221 VALUES ($1, $2, $3)
222 "#,
223 )
224 .bind(repo_id)
225 .bind(name)
226 .bind(repo_path.to_string_lossy().as_ref())
227 .execute(&self.db)
228 .await?;
229
230 Ok(repo_id)
231 }
232
233 pub async fn get_repo(&self, name: &str) -> Result<(RepoId, GitRepository)> {
238 let row: Option<(Uuid, String)> = sqlx::query_as(
240 "SELECT id, path FROM repositories WHERE name = $1",
241 )
242 .bind(name)
243 .fetch_optional(&self.db)
244 .await?;
245
246 let row = match row {
252 Some(r) => r,
253 None => {
254 let mut rows: Vec<(Uuid, String)> = sqlx::query_as(
255 "SELECT id, path FROM repositories \
256 WHERE split_part(name, '/', 2) = $1 \
257 OR (name = split_part($1, '/', 2) AND $1 LIKE '%/%')",
258 )
259 .bind(name)
260 .fetch_all(&self.db)
261 .await?;
262 match rows.len() {
263 0 => return Err(Error::RepoNotFound(name.to_string())),
264 1 => rows.remove(0),
265 _ => return Err(Error::AmbiguousRepoName(name.to_string())),
266 }
267 }
268 };
269
270 let (repo_id, repo_path) = row;
271 let git_repo = GitRepository::open(Path::new(&repo_path))?;
272 Ok((repo_id, git_repo))
273 }
274
275 pub async fn get_repo_by_db_id(&self, repo_id: RepoId) -> Result<(RepoId, GitRepository)> {
279 let row: (String,) = sqlx::query_as(
280 "SELECT path FROM repositories WHERE id = $1",
281 )
282 .bind(repo_id)
283 .fetch_optional(&self.db)
284 .await?
285 .ok_or_else(|| Error::RepoNotFound(repo_id.to_string()))?;
286
287 let git_repo = GitRepository::open(Path::new(&row.0))?;
288 Ok((repo_id, git_repo))
289 }
290
291 pub async fn index_repo(
299 &self,
300 repo_id: RepoId,
301 git_repo: &GitRepository,
302 ) -> Result<()> {
303 let root = git_repo.path().to_path_buf();
304 let files = collect_files(&root, &self.parser);
305
306 let mut all_symbols: Vec<Symbol> = Vec::new();
309 let mut all_raw_edges: Vec<RawCallEdge> = Vec::new();
310
311 let mut search_index = self.search_index.write().await;
313
314 for file_path in &files {
315 let relative = file_path
316 .strip_prefix(&root)
317 .unwrap_or(file_path);
318
319 let source = std::fs::read(file_path).map_err(|e| {
320 Error::Io(e)
321 })?;
322
323 let analysis = self.parser.parse_file(relative, &source)?;
324
325 for sym in &analysis.symbols {
327 self.symbol_store.upsert_symbol(repo_id, sym).await?;
328 search_index.index_symbol(repo_id, sym)?;
329 }
330
331 for ti in &analysis.types {
333 self.type_info_store.upsert_type_info(ti).await?;
334 }
335
336 all_symbols.extend(analysis.symbols);
337 all_raw_edges.extend(analysis.calls);
338 }
339
340 search_index.commit()?;
342 drop(search_index);
343
344 let edges = resolve_call_edges(&all_raw_edges, &all_symbols, repo_id);
346 for edge in &edges {
347 self.call_graph_store.insert_edge(edge).await?;
348 }
349
350 Ok(())
351 }
352
353 pub async fn update_files(
358 &self,
359 repo_id: RepoId,
360 git_repo: &GitRepository,
361 changed_files: &[PathBuf],
362 ) -> Result<()> {
363 self.update_files_by_root(repo_id, git_repo.path(), changed_files)
364 .await
365 }
366
367 pub async fn update_files_by_root(
373 &self,
374 repo_id: RepoId,
375 root: &Path,
376 changed_files: &[PathBuf],
377 ) -> Result<()> {
378 let root = root.to_path_buf();
379
380 let mut all_symbols: Vec<Symbol> = Vec::new();
381 let mut all_raw_edges: Vec<RawCallEdge> = Vec::new();
382
383 let mut search_index = self.search_index.write().await;
385
386 for file_path in changed_files {
387 let relative = file_path
388 .strip_prefix(&root)
389 .unwrap_or(file_path);
390 let rel_str = relative.to_string_lossy().to_string();
391
392 let old_symbols = self
395 .symbol_store
396 .find_by_file(repo_id, &rel_str)
397 .await?;
398 for old_sym in &old_symbols {
399 search_index.remove_symbol(old_sym.id)?;
400 }
401
402 self.call_graph_store
404 .delete_edges_for_file(repo_id, &rel_str)
405 .await?;
406 self.symbol_store
407 .delete_by_file(repo_id, &rel_str)
408 .await?;
409
410 let full_path = root.join(relative);
412 if !full_path.exists() {
413 continue;
415 }
416
417 if !self.parser.supports_file(relative) {
418 continue;
419 }
420
421 let source = std::fs::read(&full_path)?;
422 let analysis = self.parser.parse_file(relative, &source)?;
423
424 for sym in &analysis.symbols {
425 self.symbol_store.upsert_symbol(repo_id, sym).await?;
426 search_index.index_symbol(repo_id, sym)?;
427 }
428
429 for ti in &analysis.types {
430 self.type_info_store.upsert_type_info(ti).await?;
431 }
432
433 all_symbols.extend(analysis.symbols);
434 all_raw_edges.extend(analysis.calls);
435 }
436
437 search_index.commit()?;
438 drop(search_index);
439
440 let edges = resolve_call_edges(&all_raw_edges, &all_symbols, repo_id);
441 for edge in &edges {
442 self.call_graph_store.insert_edge(edge).await?;
443 }
444
445 Ok(())
446 }
447
448 pub async fn query_symbols(
455 &self,
456 repo_id: RepoId,
457 query: &str,
458 max_results: usize,
459 ) -> Result<Vec<Symbol>> {
460 let search_index = self.search_index.read().await;
461 let ids = search_index.search(repo_id, query, max_results)?;
462 drop(search_index);
463
464 self.symbol_store.get_by_ids(&ids).await
465 }
466
467 pub async fn get_call_graph(
472 &self,
473 _repo_id: RepoId,
474 symbol_id: SymbolId,
475 ) -> Result<(Vec<Symbol>, Vec<Symbol>)> {
476 let caller_edges = self.call_graph_store.find_callers(symbol_id).await?;
477 let callee_edges = self.call_graph_store.find_callees(symbol_id).await?;
478
479 let mut callers = Vec::with_capacity(caller_edges.len());
480 for edge in &caller_edges {
481 if let Some(sym) = self.symbol_store.get_by_id(edge.caller).await? {
482 callers.push(sym);
483 }
484 }
485
486 let mut callees = Vec::with_capacity(callee_edges.len());
487 for edge in &callee_edges {
488 if let Some(sym) = self.symbol_store.get_by_id(edge.callee).await? {
489 callees.push(sym);
490 }
491 }
492
493 Ok((callers, callees))
494 }
495
496 pub async fn codebase_summary(&self, repo_id: RepoId) -> Result<CodebaseSummary> {
502 let total_symbols = self.symbol_store.count(repo_id).await? as u64;
503
504 let row: (i64, Vec<String>) = sqlx::query_as(
506 r#"
507 SELECT
508 COUNT(DISTINCT file_path),
509 COALESCE(
510 array_agg(DISTINCT substring(file_path FROM '\.([^.]+)$'))
511 FILTER (WHERE substring(file_path FROM '\.([^.]+)$') IS NOT NULL),
512 ARRAY[]::text[]
513 )
514 FROM symbols
515 WHERE repo_id = $1
516 "#,
517 )
518 .bind(repo_id)
519 .fetch_one(&self.db)
520 .await?;
521
522 let total_files = row.0 as u64;
523 let mut languages = row.1;
524 languages.sort();
525
526 Ok(CodebaseSummary {
527 languages,
528 total_symbols,
529 total_files,
530 })
531 }
532}
533
534fn collect_files(root: &Path, parser: &ParserRegistry) -> Vec<PathBuf> {
539 let mut files = Vec::new();
540 collect_files_recursive(root, root, parser, &mut files);
541 files
542}
543
544fn collect_files_recursive(
545 root: &Path,
546 dir: &Path,
547 parser: &ParserRegistry,
548 out: &mut Vec<PathBuf>,
549) {
550 let entries = match std::fs::read_dir(dir) {
551 Ok(entries) => entries,
552 Err(_) => return,
553 };
554
555 for entry in entries.flatten() {
556 let path = entry.path();
557
558 if path.is_dir() {
559 if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
561 if name == ".git" || name.starts_with('.') {
562 continue;
563 }
564 }
565 collect_files_recursive(root, &path, parser, out);
566 } else if path.is_file() {
567 let relative = path.strip_prefix(root).unwrap_or(&path);
568 if parser.supports_file(relative) {
569 out.push(path);
570 }
571 }
572 }
573}
574
575fn resolve_call_edges(
578 raw_edges: &[RawCallEdge],
579 symbols: &[Symbol],
580 repo_id: RepoId,
581) -> Vec<CallEdge> {
582 let mut name_to_id: HashMap<String, SymbolId> = HashMap::new();
585 for sym in symbols {
586 name_to_id.insert(sym.name.clone(), sym.id);
587 name_to_id.insert(sym.qualified_name.clone(), sym.id);
588 }
589
590 raw_edges
591 .iter()
592 .filter_map(|raw| {
593 let caller = name_to_id.get(&raw.caller_name)?;
594 let callee = name_to_id.get(&raw.callee_name)?;
595 Some(CallEdge {
596 id: Uuid::new_v4(),
597 repo_id,
598 caller: *caller,
599 callee: *callee,
600 kind: raw.kind.clone(),
601 })
602 })
603 .collect()
604}
605
606#[cfg(test)]
607mod tests {
608 use super::*;
609
610 #[test]
611 fn test_resolve_call_edges_basic() {
612 let sym_a_id = Uuid::new_v4();
613 let sym_b_id = Uuid::new_v4();
614 let repo_id = Uuid::new_v4();
615
616 let symbols = vec![
617 Symbol {
618 id: sym_a_id,
619 name: "foo".into(),
620 qualified_name: "crate::foo".into(),
621 kind: dk_core::SymbolKind::Function,
622 visibility: dk_core::Visibility::Public,
623 file_path: "src/lib.rs".into(),
624 span: dk_core::Span {
625 start_byte: 0,
626 end_byte: 100,
627 },
628 signature: None,
629 doc_comment: None,
630 parent: None,
631 last_modified_by: None,
632 last_modified_intent: None,
633 },
634 Symbol {
635 id: sym_b_id,
636 name: "bar".into(),
637 qualified_name: "crate::bar".into(),
638 kind: dk_core::SymbolKind::Function,
639 visibility: dk_core::Visibility::Public,
640 file_path: "src/lib.rs".into(),
641 span: dk_core::Span {
642 start_byte: 100,
643 end_byte: 200,
644 },
645 signature: None,
646 doc_comment: None,
647 parent: None,
648 last_modified_by: None,
649 last_modified_intent: None,
650 },
651 ];
652
653 let raw_edges = vec![RawCallEdge {
654 caller_name: "foo".into(),
655 callee_name: "bar".into(),
656 call_site: dk_core::Span {
657 start_byte: 50,
658 end_byte: 60,
659 },
660 kind: dk_core::CallKind::DirectCall,
661 }];
662
663 let edges = resolve_call_edges(&raw_edges, &symbols, repo_id);
664 assert_eq!(edges.len(), 1);
665 assert_eq!(edges[0].caller, sym_a_id);
666 assert_eq!(edges[0].callee, sym_b_id);
667 assert_eq!(edges[0].repo_id, repo_id);
668 }
669
670 #[test]
671 fn test_resolve_call_edges_unresolved_skipped() {
672 let sym_a_id = Uuid::new_v4();
673 let repo_id = Uuid::new_v4();
674
675 let symbols = vec![Symbol {
676 id: sym_a_id,
677 name: "foo".into(),
678 qualified_name: "crate::foo".into(),
679 kind: dk_core::SymbolKind::Function,
680 visibility: dk_core::Visibility::Public,
681 file_path: "src/lib.rs".into(),
682 span: dk_core::Span {
683 start_byte: 0,
684 end_byte: 100,
685 },
686 signature: None,
687 doc_comment: None,
688 parent: None,
689 last_modified_by: None,
690 last_modified_intent: None,
691 }];
692
693 let raw_edges = vec![RawCallEdge {
695 caller_name: "foo".into(),
696 callee_name: "unknown".into(),
697 call_site: dk_core::Span {
698 start_byte: 50,
699 end_byte: 60,
700 },
701 kind: dk_core::CallKind::DirectCall,
702 }];
703
704 let edges = resolve_call_edges(&raw_edges, &symbols, repo_id);
705 assert!(edges.is_empty());
706 }
707
708 #[test]
709 fn test_resolve_call_edges_qualified_name() {
710 let sym_a_id = Uuid::new_v4();
711 let sym_b_id = Uuid::new_v4();
712 let repo_id = Uuid::new_v4();
713
714 let symbols = vec![
715 Symbol {
716 id: sym_a_id,
717 name: "foo".into(),
718 qualified_name: "crate::mod_a::foo".into(),
719 kind: dk_core::SymbolKind::Function,
720 visibility: dk_core::Visibility::Public,
721 file_path: "src/mod_a.rs".into(),
722 span: dk_core::Span {
723 start_byte: 0,
724 end_byte: 100,
725 },
726 signature: None,
727 doc_comment: None,
728 parent: None,
729 last_modified_by: None,
730 last_modified_intent: None,
731 },
732 Symbol {
733 id: sym_b_id,
734 name: "bar".into(),
735 qualified_name: "crate::mod_b::bar".into(),
736 kind: dk_core::SymbolKind::Function,
737 visibility: dk_core::Visibility::Public,
738 file_path: "src/mod_b.rs".into(),
739 span: dk_core::Span {
740 start_byte: 0,
741 end_byte: 100,
742 },
743 signature: None,
744 doc_comment: None,
745 parent: None,
746 last_modified_by: None,
747 last_modified_intent: None,
748 },
749 ];
750
751 let raw_edges = vec![RawCallEdge {
753 caller_name: "crate::mod_a::foo".into(),
754 callee_name: "crate::mod_b::bar".into(),
755 call_site: dk_core::Span {
756 start_byte: 50,
757 end_byte: 60,
758 },
759 kind: dk_core::CallKind::DirectCall,
760 }];
761
762 let edges = resolve_call_edges(&raw_edges, &symbols, repo_id);
763 assert_eq!(edges.len(), 1);
764 assert_eq!(edges[0].caller, sym_a_id);
765 assert_eq!(edges[0].callee, sym_b_id);
766 }
767
768 #[test]
769 fn test_collect_files_skips_git_dir() {
770 let dir = tempfile::tempdir().unwrap();
771 let root = dir.path();
772
773 std::fs::create_dir_all(root.join(".git")).unwrap();
775 std::fs::write(root.join(".git/config"), b"git config").unwrap();
776
777 std::fs::write(root.join("main.rs"), b"fn main() {}").unwrap();
779
780 std::fs::write(root.join("notes.txt"), b"hello").unwrap();
782
783 let parser = ParserRegistry::new();
784 let files = collect_files(root, &parser);
785
786 assert_eq!(files.len(), 1);
787 assert!(files[0].ends_with("main.rs"));
788 }
789
790 #[test]
791 fn test_codebase_summary_struct() {
792 let summary = CodebaseSummary {
793 languages: vec!["rs".into(), "py".into()],
794 total_symbols: 42,
795 total_files: 5,
796 };
797 assert_eq!(summary.languages.len(), 2);
798 assert_eq!(summary.total_symbols, 42);
799 assert_eq!(summary.total_files, 5);
800 }
801}