a3s_code_core/code_intelligence/
local_provider.rs

1//! Local manifest-backed Code Intelligence provider.
2
3use std::{
4    convert::Infallible,
5    sync::{
6        atomic::{AtomicU64, Ordering},
7        Arc, Mutex as StdMutex, Weak,
8    },
9    time::Duration,
10};
11
12use async_trait::async_trait;
13use tokio::sync::{broadcast, watch, RwLock};
14use tokio_util::sync::CancellationToken;
15
16use super::{
17    project_layout::ProjectLayoutResolver,
18    registry::{
19        LocalCodeIntelligenceRegistry, RegistryAcquireError, RegistryConfig, RegistryKey,
20        RegistryKeyError, RegistryReport, RegistryShutdownError, RegistryShutdownFailure,
21        RuntimeLease,
22    },
23    workspace_runtime::WorkspaceRuntime,
24    CodeDiagnostic, CodeIntelligenceError, CodeIntelligenceResult, CodeIntelligenceState,
25    CodeIntelligenceStatus, CodeLocation, CodePosition, CodeQueryResult, DocumentSymbol,
26    NavigationKind, SymbolInformation, WorkspaceCodeIntelligence,
27};
28use crate::workspace::{
29    LocalWorkspaceManifest, LocalWorkspaceManifestSnapshot, WorkspaceFileChange,
30    WorkspaceFileSystem, WorkspacePath,
31};
32
33const DEFAULT_QUERY_TIMEOUT: Duration = Duration::from_secs(15);
34type RuntimeRegistry = LocalCodeIntelligenceRegistry<WorkspaceRuntime, Infallible>;
35type WorkspaceRuntimeLease = RuntimeLease<WorkspaceRuntime, Infallible>;
36
37/// Native local provider sharing the workspace manifest's existing watcher.
38pub struct LocalCodeIntelligence {
39    isolation_scope: String,
40    manifest: Arc<LocalWorkspaceManifest>,
41    file_system: Arc<dyn WorkspaceFileSystem>,
42    registry: RuntimeRegistry,
43    current: RwLock<Option<WorkspaceRuntimeLease>>,
44    status: watch::Sender<CodeIntelligenceStatus>,
45    generation: Arc<AtomicU64>,
46    lifetime: CancellationToken,
47    manifest_task: StdMutex<Option<tokio::task::JoinHandle<()>>>,
48    query_timeout: Duration,
49}
50
51impl std::fmt::Debug for LocalCodeIntelligence {
52    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        formatter
54            .debug_struct("LocalCodeIntelligence")
55            .field("isolation_scope", &self.isolation_scope)
56            .field("manifest_root", &self.manifest.snapshot().root)
57            .field("generation", &self.generation.load(Ordering::Acquire))
58            .finish_non_exhaustive()
59    }
60}
61
62impl LocalCodeIntelligence {
63    /// Create a provider and acquire a cheap, lazily-started runtime generation.
64    pub async fn start(
65        isolation_scope: impl Into<String>,
66        manifest: Arc<LocalWorkspaceManifest>,
67        file_system: Arc<dyn WorkspaceFileSystem>,
68    ) -> CodeIntelligenceResult<Arc<Self>> {
69        Self::start_with_timeout(
70            isolation_scope,
71            manifest,
72            file_system,
73            DEFAULT_QUERY_TIMEOUT,
74        )
75        .await
76    }
77
78    pub(crate) async fn start_with_timeout(
79        isolation_scope: impl Into<String>,
80        manifest: Arc<LocalWorkspaceManifest>,
81        file_system: Arc<dyn WorkspaceFileSystem>,
82        query_timeout: Duration,
83    ) -> CodeIntelligenceResult<Arc<Self>> {
84        let isolation_scope = isolation_scope.into();
85        let snapshot_rx = manifest.subscribe();
86        let changes_rx = manifest.subscribe_changes();
87        let registry = RuntimeRegistry::new(
88            RegistryConfig::new(Duration::ZERO, 0),
89            |runtime: Arc<WorkspaceRuntime>| async move {
90                runtime.shutdown().await;
91                Ok(())
92            },
93        );
94        let (status, _) = watch::channel(CodeIntelligenceStatus {
95            state: CodeIntelligenceState::Starting,
96            message: Some("Code Intelligence is preparing the saved workspace".to_owned()),
97            ..CodeIntelligenceStatus::default()
98        });
99        let provider = Arc::new(Self {
100            isolation_scope,
101            manifest,
102            file_system,
103            registry,
104            current: RwLock::new(None),
105            status,
106            generation: Arc::new(AtomicU64::new(0)),
107            lifetime: CancellationToken::new(),
108            manifest_task: StdMutex::new(None),
109            query_timeout,
110        });
111
112        provider
113            .refresh_snapshot(provider.manifest.snapshot())
114            .await?;
115        let weak = Arc::downgrade(&provider);
116        let task = tokio::spawn(run_manifest_updates(
117            weak,
118            snapshot_rx,
119            changes_rx,
120            provider.lifetime.clone(),
121        ));
122        *provider
123            .manifest_task
124            .lock()
125            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(task);
126        Ok(provider)
127    }
128
129    /// Stop manifest forwarding and all language processes owned by this provider.
130    pub async fn shutdown(&self) {
131        self.lifetime.cancel();
132        let task = self
133            .manifest_task
134            .lock()
135            .unwrap_or_else(std::sync::PoisonError::into_inner)
136            .take();
137        if let Some(task) = task {
138            let _ = task.await;
139        }
140        let old = self.current.write().await.take();
141        drop(old);
142        report_registry_cleanup("shutdown", self.registry.shutdown_all().await);
143        self.status.send_replace(CodeIntelligenceStatus {
144            state: CodeIntelligenceState::Unavailable,
145            message: Some("Code Intelligence is shut down".to_owned()),
146            ..CodeIntelligenceStatus::default()
147        });
148    }
149
150    async fn refresh_snapshot(
151        self: &Arc<Self>,
152        snapshot: LocalWorkspaceManifestSnapshot,
153    ) -> CodeIntelligenceResult<()> {
154        let layout = ProjectLayoutResolver::resolve(&snapshot);
155        {
156            let current = self.current.read().await;
157            if let Some(runtime) = current
158                .as_ref()
159                .filter(|runtime| runtime.layout_hash() == layout.layout_hash)
160            {
161                runtime.update_snapshot(&snapshot).await;
162                return Ok(());
163            }
164        }
165
166        let key = RegistryKey::new(
167            self.isolation_scope.clone(),
168            &snapshot.root,
169            layout.layout_hash,
170        )
171        .await
172        .map_err(map_key_error)?;
173        let canonical_root = key.canonical_root().to_path_buf();
174        let runtime_snapshot = snapshot.clone();
175        let file_system = Arc::clone(&self.file_system);
176        let timeout = self.query_timeout;
177        let lease = self
178            .registry
179            .acquire(key, move |_| async move {
180                Ok(WorkspaceRuntime::new(
181                    canonical_root,
182                    layout,
183                    &runtime_snapshot,
184                    file_system,
185                    timeout,
186                ))
187            })
188            .await
189            .map_err(map_acquire_error)?;
190
191        lease.update_snapshot(&snapshot).await;
192        let mut receiver = lease.subscribe_status();
193        let generation = self.generation.fetch_add(1, Ordering::AcqRel) + 1;
194        self.status.send_replace(receiver.borrow().clone());
195        let old = self.current.write().await.replace(lease);
196        drop(old);
197        report_registry_cleanup("layout refresh", self.registry.cleanup_idle().await);
198        self.spawn_status_forwarder(generation, &mut receiver);
199        Ok(())
200    }
201
202    fn spawn_status_forwarder(
203        &self,
204        generation: u64,
205        receiver: &mut watch::Receiver<CodeIntelligenceStatus>,
206    ) {
207        let mut receiver = receiver.clone();
208        let sender = self.status.clone();
209        let current_generation = Arc::clone(&self.generation);
210        let lifetime = self.lifetime.clone();
211        tokio::spawn(async move {
212            loop {
213                tokio::select! {
214                    _ = lifetime.cancelled() => break,
215                    changed = receiver.changed() => {
216                        if changed.is_err()
217                            || current_generation.load(Ordering::Acquire) != generation
218                        {
219                            break;
220                        }
221                        sender.send_replace(receiver.borrow().clone());
222                    }
223                }
224            }
225        });
226    }
227
228    async fn handle_changes(&self, changes: &[WorkspaceFileChange]) {
229        let current = self.current.read().await;
230        if let Some(runtime) = current.as_ref() {
231            runtime.notify_file_changes(changes).await;
232        }
233    }
234
235    async fn runtime(
236        &self,
237    ) -> CodeIntelligenceResult<tokio::sync::RwLockReadGuard<'_, Option<WorkspaceRuntimeLease>>>
238    {
239        let current = self.current.read().await;
240        if current.is_none() {
241            return Err(CodeIntelligenceError::Unavailable {
242                message: "Code Intelligence has not prepared this workspace yet".to_owned(),
243            });
244        }
245        Ok(current)
246    }
247
248    fn report_update_error(&self, error: &CodeIntelligenceError) {
249        let mut status = self.status.borrow().clone();
250        status.state = CodeIntelligenceState::Degraded;
251        status.message = Some(format!("workspace refresh failed: {error}"));
252        self.status.send_replace(status);
253    }
254}
255
256#[async_trait]
257impl WorkspaceCodeIntelligence for LocalCodeIntelligence {
258    fn subscribe_status(&self) -> watch::Receiver<CodeIntelligenceStatus> {
259        self.status.subscribe()
260    }
261
262    async fn document_symbols(
263        &self,
264        path: &WorkspacePath,
265        cancellation: CancellationToken,
266    ) -> CodeIntelligenceResult<CodeQueryResult<DocumentSymbol>> {
267        let runtime = self.runtime().await?;
268        let Some(runtime) = runtime.as_ref() else {
269            return Err(runtime_unavailable());
270        };
271        runtime.document_symbols(path, cancellation).await
272    }
273
274    async fn search_symbols(
275        &self,
276        query: &str,
277        limit: usize,
278        cancellation: CancellationToken,
279    ) -> CodeIntelligenceResult<CodeQueryResult<SymbolInformation>> {
280        let runtime = self.runtime().await?;
281        let Some(runtime) = runtime.as_ref() else {
282            return Err(runtime_unavailable());
283        };
284        runtime.search_symbols(query, limit, cancellation).await
285    }
286
287    async fn navigate(
288        &self,
289        kind: NavigationKind,
290        path: &WorkspacePath,
291        position: CodePosition,
292        cancellation: CancellationToken,
293    ) -> CodeIntelligenceResult<CodeQueryResult<CodeLocation>> {
294        let runtime = self.runtime().await?;
295        let Some(runtime) = runtime.as_ref() else {
296            return Err(runtime_unavailable());
297        };
298        runtime.navigate(kind, path, position, cancellation).await
299    }
300
301    async fn diagnostics(
302        &self,
303        path: Option<&WorkspacePath>,
304        cancellation: CancellationToken,
305    ) -> CodeIntelligenceResult<CodeQueryResult<CodeDiagnostic>> {
306        let runtime = self.runtime().await?;
307        let Some(runtime) = runtime.as_ref() else {
308            return Err(runtime_unavailable());
309        };
310        runtime.diagnostics(path, cancellation).await
311    }
312}
313
314impl Drop for LocalCodeIntelligence {
315    fn drop(&mut self) {
316        self.lifetime.cancel();
317        if let Some(task) = self
318            .manifest_task
319            .lock()
320            .unwrap_or_else(std::sync::PoisonError::into_inner)
321            .take()
322        {
323            task.abort();
324        }
325    }
326}
327
328async fn run_manifest_updates(
329    provider: Weak<LocalCodeIntelligence>,
330    mut snapshots: broadcast::Receiver<LocalWorkspaceManifestSnapshot>,
331    mut changes: broadcast::Receiver<WorkspaceFileChange>,
332    lifetime: CancellationToken,
333) {
334    loop {
335        tokio::select! {
336            _ = lifetime.cancelled() => break,
337            update = snapshots.recv() => match update {
338                Ok(snapshot) => {
339                    let Some(provider) = provider.upgrade() else { break; };
340                    if let Err(error) = provider.refresh_snapshot(snapshot).await {
341                        provider.report_update_error(&error);
342                    }
343                }
344                Err(broadcast::error::RecvError::Lagged(_)) => {
345                    let Some(provider) = provider.upgrade() else { break; };
346                    if let Err(error) = provider.refresh_snapshot(provider.manifest.snapshot()).await {
347                        provider.report_update_error(&error);
348                    }
349                }
350                Err(broadcast::error::RecvError::Closed) => break,
351            },
352            update = changes.recv() => match update {
353                Ok(change) => {
354                    let mut batch = vec![change];
355                    while let Ok(change) = changes.try_recv() {
356                        batch.push(change);
357                    }
358                    let Some(provider) = provider.upgrade() else { break; };
359                    provider.handle_changes(&batch).await;
360                }
361                Err(broadcast::error::RecvError::Lagged(count)) => {
362                    let Some(provider) = provider.upgrade() else { break; };
363                    let mut status = provider.status.borrow().clone();
364                    status.state = CodeIntelligenceState::Degraded;
365                    status.message = Some(format!(
366                        "workspace change stream skipped {count} events; saved documents will resynchronize on query"
367                    ));
368                    provider.status.send_replace(status);
369                }
370                Err(broadcast::error::RecvError::Closed) => break,
371            },
372        }
373    }
374}
375
376fn map_key_error(error: RegistryKeyError) -> CodeIntelligenceError {
377    CodeIntelligenceError::Unavailable {
378        message: error.to_string(),
379    }
380}
381
382fn map_acquire_error(error: RegistryAcquireError<Infallible>) -> CodeIntelligenceError {
383    match error {
384        RegistryAcquireError::ShuttingDown => CodeIntelligenceError::Unavailable {
385            message: "the Code Intelligence registry is shutting down".to_owned(),
386        },
387        RegistryAcquireError::Factory(error) => match *error {},
388        RegistryAcquireError::FactoryPanicked { message } => CodeIntelligenceError::Unavailable {
389            message: format!("Code Intelligence runtime initialization panicked: {message}"),
390        },
391        RegistryAcquireError::LeaseLimit => CodeIntelligenceError::Unavailable {
392            message: "the Code Intelligence runtime lease limit was exhausted".to_owned(),
393        },
394    }
395}
396
397fn runtime_unavailable() -> CodeIntelligenceError {
398    CodeIntelligenceError::Unavailable {
399        message: "Code Intelligence has not prepared this workspace yet".to_owned(),
400    }
401}
402
403fn report_registry_cleanup(context: &'static str, report: RegistryReport<Infallible>) {
404    if !report.removed.is_empty() {
405        tracing::debug!(
406            context,
407            retired = report.removed.len(),
408            "Code Intelligence retired workspace runtimes"
409        );
410    }
411    for RegistryShutdownError { key, failure } in report.errors {
412        match failure {
413            RegistryShutdownFailure::Runtime(error) => {
414                tracing::error!(
415                    context,
416                    workspace = ?key.canonical_root(),
417                    ?error,
418                    "Code Intelligence runtime cleanup returned an impossible error"
419                );
420            }
421            RegistryShutdownFailure::Panicked { message } => {
422                tracing::warn!(
423                    context,
424                    workspace = ?key.canonical_root(),
425                    %message,
426                    "Code Intelligence runtime cleanup panicked"
427                );
428            }
429        }
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use crate::workspace::{
437        LocalWorkspaceFile, LocalWorkspaceFileStatus, ManifestWorkspaceBackend,
438    };
439
440    fn file(path: &str) -> LocalWorkspaceFile {
441        LocalWorkspaceFile {
442            path: path.to_owned(),
443            size: 1,
444            modified_ms: Some(1),
445            language: None,
446            status: LocalWorkspaceFileStatus::Tracked,
447            binary: false,
448            generated: false,
449        }
450    }
451
452    async fn scanned_snapshot(manifest: &LocalWorkspaceManifest) -> LocalWorkspaceManifestSnapshot {
453        if manifest.snapshot().version > 0 {
454            return manifest.snapshot();
455        }
456        let mut snapshots = manifest.subscribe();
457        tokio::time::timeout(Duration::from_secs(5), async {
458            loop {
459                let snapshot = snapshots.recv().await.unwrap();
460                if snapshot.version > 0 {
461                    break snapshot;
462                }
463            }
464        })
465        .await
466        .expect("initial manifest scan should finish")
467    }
468
469    #[tokio::test]
470    async fn provider_reuses_unchanged_layout_and_shutdown_is_idempotent() {
471        let workspace = tempfile::tempdir().unwrap();
472        std::fs::create_dir(workspace.path().join("src")).unwrap();
473        std::fs::write(workspace.path().join("Cargo.toml"), "[workspace]\n").unwrap();
474        std::fs::write(workspace.path().join("src/lib.rs"), "pub fn saved() {}\n").unwrap();
475        let backend = ManifestWorkspaceBackend::new(workspace.path());
476        let manifest = backend.manifest();
477        let initial = scanned_snapshot(&manifest).await;
478        let file_system: Arc<dyn WorkspaceFileSystem> = backend;
479        let provider = LocalCodeIntelligence::start_with_timeout(
480            "test-session",
481            manifest,
482            file_system,
483            Duration::from_secs(1),
484        )
485        .await
486        .unwrap();
487        let generation = provider.generation.load(Ordering::Acquire);
488        assert_eq!(generation, 1);
489
490        let mut source_only = initial.clone();
491        source_only.version += 1;
492        source_only.files.push(file("src/new.rs"));
493        provider.refresh_snapshot(source_only).await.unwrap();
494        assert_eq!(provider.generation.load(Ordering::Acquire), generation);
495
496        let mut changed_layout = initial;
497        changed_layout.version += 2;
498        changed_layout
499            .files
500            .retain(|entry| entry.path != "Cargo.toml");
501        changed_layout.files.push(file("package.json"));
502        provider.refresh_snapshot(changed_layout).await.unwrap();
503        assert_eq!(provider.generation.load(Ordering::Acquire), generation + 1);
504
505        provider.shutdown().await;
506        provider.shutdown().await;
507        assert!(provider.current.read().await.is_none());
508        assert!(provider.lifetime.is_cancelled());
509        assert_eq!(
510            provider.status.borrow().state,
511            CodeIntelligenceState::Unavailable
512        );
513    }
514}
a3s_code_core/code_intelligence/local_provider.rs

a3s_code_core/code_intelligence/
local_provider.rs
