code_system_graph_core/codegraph/
mod.rs

1mod cli;
2mod contract;
3mod mcp;
4
5use std::collections::BTreeMap;
6use std::ffi::OsString;
7use std::path::{Component, Path};
8use std::sync::Arc;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::time::Duration;
11
12use async_trait::async_trait;
13use code_system_graph_model::RepoId;
14use tokio::sync::{Mutex, OwnedSemaphorePermit, Semaphore};
15
16use self::cli::CodeGraphCli;
17use self::contract::{cli_operations, map_mcp_tools, supports_cli_contract};
18use self::mcp::{CodeGraphMcp, McpProbe};
19use crate::{
20    AffectedTestsRequest, AffectedTestsResult, LocalCodeIntelligenceProvider, LocalContextRequest, LocalContextResult, LocalImpactRequest, LocalImpactResult, LocalNeighborResult, LocalNeighborsRequest, ProviderCapability, ProviderDegradation, ProviderError, ProviderRequest, ProviderStatus, ResolveSymbolsRequest, ResolveSymbolsResult
21};
22
23static DIAGNOSTIC_SEQUENCE: AtomicU64 = AtomicU64::new(1);
24const MAX_QUERY_BYTES: usize = 4096;
25const MAX_LOCAL_DEPTH: usize = 64;
26const MAX_CHANGED_FILES: usize = 1024;
27
28/// Process and circuit-breaker settings for the production `CodeGraph` adapter.
29#[derive(Debug, Clone)]
30pub struct CodeGraphConfig {
31    /// Executable path or `PATH`-resolved binary name.
32    pub binary: OsString,
33    /// Maximum concurrent `CodeGraph` child processes.
34    pub max_concurrent_processes: usize,
35    /// Time that MCP is bypassed for one repository after an MCP timeout.
36    pub circuit_breaker_cooldown: Duration,
37}
38
39impl Default for CodeGraphConfig {
40    fn default() -> Self {
41        Self {
42            binary: OsString::from("codegraph"),
43            max_concurrent_processes: 2,
44            circuit_breaker_cooldown: Duration::from_secs(30),
45        }
46    }
47}
48
49/// Production local-intelligence provider using public `CodeGraph` MCP and CLI interfaces.
50pub struct CodeGraphProvider {
51    cli: CodeGraphCli,
52    mcp: CodeGraphMcp,
53    permits: Arc<Semaphore>,
54    circuit_breaker_cooldown: Duration,
55    mcp_circuits: Mutex<BTreeMap<RepoId, tokio::time::Instant>>,
56    cli_version: Mutex<Option<String>>,
57    closed: AtomicBool,
58}
59
60struct CliProbe {
61    version: String,
62    status: ProviderStatus,
63    degradations: Vec<ProviderDegradation>,
64}
65
66enum CliProbeOutcome {
67    Ready(CliProbe),
68    Degraded(ProviderCapability),
69}
70
71impl CodeGraphProvider {
72    /// Builds a production adapter without starting or modifying any `CodeGraph` index.
73    ///
74    /// # Errors
75    ///
76    /// Returns [`ProviderError::InvalidRequest`] when child-process concurrency is zero.
77    pub fn new(config: CodeGraphConfig) -> Result<Self, ProviderError> {
78        if config.max_concurrent_processes == 0 {
79            return Err(ProviderError::InvalidRequest(
80                "max_concurrent_processes must be greater than zero".to_owned(),
81            ));
82        }
83        Ok(Self {
84            cli: CodeGraphCli::new(config.binary.clone()),
85            mcp: CodeGraphMcp::new(config.binary),
86            permits: Arc::new(Semaphore::new(config.max_concurrent_processes)),
87            circuit_breaker_cooldown: config.circuit_breaker_cooldown,
88            mcp_circuits: Mutex::new(BTreeMap::new()),
89            cli_version: Mutex::new(None),
90            closed: AtomicBool::new(false),
91        })
92    }
93
94    /// Reads the structured local-index status without starting MCP or modifying the index.
95    ///
96    /// # Errors
97    ///
98    /// Returns [`ProviderError`] when the request is invalid, the compatible CLI cannot be
99    /// executed, or its bounded JSON status does not satisfy the validated contract.
100    pub async fn index_status(
101        &self,
102        mut request: ProviderRequest,
103    ) -> Result<ProviderStatus, ProviderError> {
104        let deadline = tokio::time::Instant::now() + request.budget.timeout;
105        let _permit = self.enter(&request, deadline).await?;
106        self.compatible_cli_version(&mut request, deadline).await?;
107        update_remaining_timeout(&mut request, deadline)?;
108        self.cli
109            .status(&request)
110            .await
111            .map(|status| status.status())
112    }
113
114    async fn enter(
115        &self,
116        request: &ProviderRequest,
117        deadline: tokio::time::Instant,
118    ) -> Result<OwnedSemaphorePermit, ProviderError> {
119        validate_request(request)?;
120        if self.closed.load(Ordering::Acquire) {
121            return Err(ProviderError::InvalidRequest(
122                "provider has already been shut down".to_owned(),
123            ));
124        }
125        tokio::select! {
126            biased;
127            () = request.cancellation.cancelled() => Err(ProviderError::Cancelled),
128            () = tokio::time::sleep_until(deadline) => Err(timeout_error()),
129            permit = self.permits.clone().acquire_owned() => permit.map_err(|_| {
130                ProviderError::InvalidRequest("provider has already been shut down".to_owned())
131            }),
132        }
133    }
134
135    async fn compatible_cli_version(
136        &self,
137        request: &mut ProviderRequest,
138        deadline: tokio::time::Instant,
139    ) -> Result<String, ProviderError> {
140        if let Some(version) = self.cli_version.lock().await.clone() {
141            if supports_cli_contract(&version) {
142                return Ok(version);
143            }
144            return Err(invalid_response(format!(
145                "CodeGraph CLI {version} has no validated structured-output adapter"
146            )));
147        }
148        update_remaining_timeout(request, deadline)?;
149        let version = self.cli.version(request).await?;
150        *self.cli_version.lock().await = Some(version.clone());
151        if !supports_cli_contract(&version) {
152            return Err(invalid_response(format!(
153                "CodeGraph CLI {version} has no validated structured-output adapter"
154            )));
155        }
156        Ok(version)
157    }
158
159    async fn mcp_is_available(&self, repo_id: &RepoId) -> bool {
160        let now = tokio::time::Instant::now();
161        let mut circuits = self.mcp_circuits.lock().await;
162        circuits.retain(|_, until| *until > now);
163        !circuits.contains_key(repo_id)
164    }
165
166    async fn open_mcp_circuit(&self, repo_id: RepoId) {
167        self.mcp_circuits.lock().await.insert(
168            repo_id,
169            tokio::time::Instant::now() + self.circuit_breaker_cooldown,
170        );
171    }
172
173    async fn probe_cli(
174        &self,
175        request: &mut ProviderRequest,
176        deadline: tokio::time::Instant,
177    ) -> Result<CliProbeOutcome, ProviderError> {
178        update_remaining_timeout(request, deadline)?;
179        let version = match self.cli.version(request).await {
180            Ok(version) => version,
181            Err(ProviderError::Cancelled) => return Err(ProviderError::Cancelled),
182            Err(error) => return Ok(CliProbeOutcome::Degraded(unavailable_capability(&error))),
183        };
184        *self.cli_version.lock().await = Some(version.clone());
185        if !supports_cli_contract(&version) {
186            return Ok(CliProbeOutcome::Ready(CliProbe {
187                status: ProviderStatus::Incompatible,
188                degradations: vec![ProviderDegradation {
189                    kind: "unsupported_cli_version".to_owned(),
190                    message: format!(
191                        "CodeGraph CLI {version} is outside the validated structured-output matrix."
192                    ),
193                    diagnostics_id: None,
194                }],
195                version,
196            }));
197        }
198        update_remaining_timeout(request, deadline)?;
199        let mut degradations = Vec::new();
200        let status = match self.cli.status(request).await {
201            Ok(status) => {
202                if status
203                    .version
204                    .as_deref()
205                    .is_some_and(|value| value != version)
206                {
207                    degradations.push(ProviderDegradation {
208                        kind: "version_mismatch".to_owned(),
209                        message:
210                            "CodeGraph version and status contracts reported different versions."
211                                .to_owned(),
212                        diagnostics_id: None,
213                    });
214                }
215                status.status()
216            }
217            Err(error) => {
218                let status = if matches!(error, ProviderError::Timeout { .. }) {
219                    ProviderStatus::TimedOut
220                } else {
221                    ProviderStatus::InvalidResponse
222                };
223                degradations.push(degradation_from_error(&error));
224                status
225            }
226        };
227        Ok(CliProbeOutcome::Ready(CliProbe {
228            version,
229            status,
230            degradations,
231        }))
232    }
233
234    async fn probe_mcp(
235        &self,
236        request: &mut ProviderRequest,
237        deadline: tokio::time::Instant,
238        degradations: &mut Vec<ProviderDegradation>,
239    ) -> Result<Option<McpProbe>, ProviderError> {
240        if !self.mcp_is_available(&request.repo_id).await {
241            degradations.push(ProviderDegradation {
242                kind: "mcp_circuit_open".to_owned(),
243                message: "MCP was bypassed after a recent timeout; compatible CLI fallback remains available.".to_owned(),
244                diagnostics_id: None,
245            });
246            return Ok(None);
247        }
248        update_remaining_timeout(request, deadline)?;
249        match self.mcp.probe(request).await {
250            Ok(probe) => Ok(Some(probe)),
251            Err(ProviderError::Cancelled) => Err(ProviderError::Cancelled),
252            Err(error) => {
253                if matches!(error, ProviderError::Timeout { .. }) {
254                    self.open_mcp_circuit(request.repo_id.clone()).await;
255                }
256                degradations.push(degradation_from_error(&error));
257                Ok(None)
258            }
259        }
260    }
261
262    fn assemble_capability(
263        request: &ProviderRequest,
264        mut cli_probe: CliProbe,
265        mcp_probe: Option<McpProbe>,
266    ) -> ProviderCapability {
267        if mcp_probe
268            .as_ref()
269            .and_then(|probe| probe.version.as_deref())
270            .is_some_and(|version| version != cli_probe.version)
271        {
272            cli_probe.degradations.push(ProviderDegradation {
273                kind: "mcp_cli_version_mismatch".to_owned(),
274                message: "CodeGraph MCP and CLI reported different public implementation versions."
275                    .to_owned(),
276                diagnostics_id: None,
277            });
278        }
279        let mut operations = BTreeMap::new();
280        if let Some(probe) = &mcp_probe {
281            for operation in map_mcp_tools(&probe.tools) {
282                operations.insert(operation.operation, operation);
283            }
284        }
285        for operation in cli_operations(&cli_probe.version) {
286            operations.entry(operation.operation).or_insert(operation);
287        }
288        let status = if operations.is_empty() && cli_probe.status == ProviderStatus::Available {
289            ProviderStatus::Incompatible
290        } else {
291            cli_probe.status
292        };
293        ProviderCapability {
294            provider: "codegraph".to_owned(),
295            version: Some(cli_probe.version),
296            status,
297            tools: mcp_tool_names(mcp_probe.as_ref()),
298            protocol_version: mcp_probe.map(|probe| probe.protocol_version),
299            operations: operations.into_values().collect(),
300            degradations: cli_probe.degradations,
301            remediation: remediation(status, &request.project_path),
302        }
303    }
304}
305
306#[async_trait]
307impl LocalCodeIntelligenceProvider for CodeGraphProvider {
308    fn provider_name(&self) -> &'static str {
309        "codegraph"
310    }
311
312    async fn probe(
313        &self,
314        mut request: ProviderRequest,
315    ) -> Result<ProviderCapability, ProviderError> {
316        let deadline = tokio::time::Instant::now() + request.budget.timeout;
317        let _permit = self.enter(&request, deadline).await?;
318        let mut cli_probe = match self.probe_cli(&mut request, deadline).await? {
319            CliProbeOutcome::Ready(probe) => probe,
320            CliProbeOutcome::Degraded(capability) => return Ok(capability),
321        };
322        let mcp_probe = self
323            .probe_mcp(&mut request, deadline, &mut cli_probe.degradations)
324            .await?;
325        Ok(Self::assemble_capability(&request, cli_probe, mcp_probe))
326    }
327
328    async fn resolve_symbols(
329        &self,
330        mut input: ResolveSymbolsRequest,
331    ) -> Result<ResolveSymbolsResult, ProviderError> {
332        validate_query(&input.query)?;
333        let deadline = tokio::time::Instant::now() + input.request.budget.timeout;
334        let _permit = self.enter(&input.request, deadline).await?;
335        self.compatible_cli_version(&mut input.request, deadline)
336            .await?;
337        update_remaining_timeout(&mut input.request, deadline)?;
338        self.cli.resolve_symbols(&input).await
339    }
340
341    async fn get_local_neighbors(
342        &self,
343        mut input: LocalNeighborsRequest,
344    ) -> Result<LocalNeighborResult, ProviderError> {
345        validate_query(&input.symbol)?;
346        let deadline = tokio::time::Instant::now() + input.request.budget.timeout;
347        let _permit = self.enter(&input.request, deadline).await?;
348        self.compatible_cli_version(&mut input.request, deadline)
349            .await?;
350        update_remaining_timeout(&mut input.request, deadline)?;
351        self.cli.local_neighbors(&input).await
352    }
353
354    async fn get_local_impact(
355        &self,
356        mut input: LocalImpactRequest,
357    ) -> Result<LocalImpactResult, ProviderError> {
358        validate_query(&input.symbol)?;
359        validate_depth(input.max_depth)?;
360        let deadline = tokio::time::Instant::now() + input.request.budget.timeout;
361        let _permit = self.enter(&input.request, deadline).await?;
362        self.compatible_cli_version(&mut input.request, deadline)
363            .await?;
364        update_remaining_timeout(&mut input.request, deadline)?;
365        self.cli.local_impact(&input).await
366    }
367
368    async fn build_local_context(
369        &self,
370        mut input: LocalContextRequest,
371    ) -> Result<LocalContextResult, ProviderError> {
372        validate_query(&input.query)?;
373        if input.max_files == 0 || input.max_files > input.request.budget.max_items {
374            return Err(ProviderError::InvalidRequest(format!(
375                "max_files must be between 1 and max_items ({})",
376                input.request.budget.max_items
377            )));
378        }
379        let deadline = tokio::time::Instant::now() + input.request.budget.timeout;
380        let _permit = self.enter(&input.request, deadline).await?;
381        let mut degradations = Vec::new();
382        if self.mcp_is_available(&input.request.repo_id).await {
383            update_remaining_timeout(&mut input.request, deadline)?;
384            match self.mcp.local_context(&input).await {
385                Ok(result) => return Ok(result),
386                Err(ProviderError::Cancelled) => return Err(ProviderError::Cancelled),
387                Err(error) => {
388                    if matches!(error, ProviderError::Timeout { .. }) {
389                        self.open_mcp_circuit(input.request.repo_id.clone()).await;
390                    }
391                    degradations.push(degradation_from_error(&error));
392                }
393            }
394        } else {
395            degradations.push(ProviderDegradation {
396                kind: "mcp_circuit_open".to_owned(),
397                message: "MCP was bypassed after a recent timeout.".to_owned(),
398                diagnostics_id: None,
399            });
400        }
401        self.compatible_cli_version(&mut input.request, deadline)
402            .await?;
403        update_remaining_timeout(&mut input.request, deadline)?;
404        let mut result = self.cli.local_context(&input).await?;
405        result.execution.degradations = degradations;
406        Ok(result)
407    }
408
409    async fn get_affected_tests(
410        &self,
411        mut input: AffectedTestsRequest,
412    ) -> Result<Option<AffectedTestsResult>, ProviderError> {
413        validate_depth(input.max_depth)?;
414        validate_changed_files(&input.changed_files)?;
415        let deadline = tokio::time::Instant::now() + input.request.budget.timeout;
416        let _permit = self.enter(&input.request, deadline).await?;
417        match self
418            .compatible_cli_version(&mut input.request, deadline)
419            .await
420        {
421            Ok(_) => {}
422            Err(ProviderError::InvalidResponse { .. }) => return Ok(None),
423            Err(error) => return Err(error),
424        }
425        update_remaining_timeout(&mut input.request, deadline)?;
426        self.cli.affected_tests(&input).await.map(Some)
427    }
428
429    async fn shutdown(&self) -> Result<(), ProviderError> {
430        self.closed.store(true, Ordering::Release);
431        self.permits.close();
432        Ok(())
433    }
434}
435
436fn unavailable_capability(error: &ProviderError) -> ProviderCapability {
437    let status = match error {
438        ProviderError::Timeout { .. } => ProviderStatus::TimedOut,
439        ProviderError::InvalidResponse { .. } => ProviderStatus::InvalidResponse,
440        _ => ProviderStatus::Unavailable,
441    };
442    ProviderCapability {
443        provider: "codegraph".to_owned(),
444        version: None,
445        status,
446        tools: Vec::new(),
447        protocol_version: None,
448        operations: Vec::new(),
449        degradations: vec![degradation_from_error(error)],
450        remediation: Some("Install CodeGraph or configure its public executable path.".to_owned()),
451    }
452}
453
454fn mcp_tool_names(probe: Option<&McpProbe>) -> Vec<String> {
455    let mut names = probe
456        .map(|probe| {
457            probe
458                .tools
459                .iter()
460                .map(|tool| tool.name.clone())
461                .collect::<Vec<_>>()
462        })
463        .unwrap_or_default();
464    names.sort();
465    names
466}
467
468fn remediation(status: ProviderStatus, project_path: &Path) -> Option<String> {
469    match status {
470        ProviderStatus::IndexMissing => Some(format!("Create the index explicitly with `codegraph init {}`.", project_path.display())),
471        ProviderStatus::Stale => Some(format!("Refresh the index explicitly with `codegraph sync {}`.", project_path.display())),
472        ProviderStatus::Incompatible => Some("Use a CodeGraph version listed in the compatibility matrix or rely on federated boundaries without local enrichment.".to_owned()),
473        _ => None,
474    }
475}
476
477fn validate_request(request: &ProviderRequest) -> Result<(), ProviderError> {
478    if request.cancellation.is_cancelled() {
479        return Err(ProviderError::Cancelled);
480    }
481    if request.budget.timeout.is_zero() {
482        return Err(ProviderError::InvalidRequest(
483            "timeout must be greater than zero".to_owned(),
484        ));
485    }
486    if request.budget.max_output_bytes == 0 {
487        return Err(ProviderError::InvalidRequest(
488            "max_output_bytes must be greater than zero".to_owned(),
489        ));
490    }
491    if request.budget.max_items == 0 {
492        return Err(ProviderError::InvalidRequest(
493            "max_items must be greater than zero".to_owned(),
494        ));
495    }
496    if !request.project_path.is_absolute() {
497        return Err(ProviderError::InvalidRequest(
498            "project_path must be absolute".to_owned(),
499        ));
500    }
501    if !request.project_path.is_dir() {
502        return Err(ProviderError::InvalidRequest(
503            "project_path must identify an existing directory".to_owned(),
504        ));
505    }
506    Ok(())
507}
508
509fn validate_query(query: &str) -> Result<(), ProviderError> {
510    if query.trim().is_empty() {
511        return Err(ProviderError::InvalidRequest(
512            "provider query must not be empty".to_owned(),
513        ));
514    }
515    if query.len() > MAX_QUERY_BYTES {
516        return Err(ProviderError::InvalidRequest(format!(
517            "provider query exceeds {MAX_QUERY_BYTES} bytes"
518        )));
519    }
520    Ok(())
521}
522
523fn validate_depth(depth: usize) -> Result<(), ProviderError> {
524    if depth == 0 || depth > MAX_LOCAL_DEPTH {
525        return Err(ProviderError::InvalidRequest(format!(
526            "local depth must be between 1 and {MAX_LOCAL_DEPTH}"
527        )));
528    }
529    Ok(())
530}
531
532fn validate_changed_files(files: &[String]) -> Result<(), ProviderError> {
533    if files.is_empty() || files.len() > MAX_CHANGED_FILES {
534        return Err(ProviderError::InvalidRequest(format!(
535            "changed_files must contain between 1 and {MAX_CHANGED_FILES} paths"
536        )));
537    }
538    for file in files {
539        let path = Path::new(file);
540        if file.is_empty()
541            || path.is_absolute()
542            || path
543                .components()
544                .any(|component| matches!(component, Component::ParentDir | Component::RootDir))
545        {
546            return Err(ProviderError::InvalidRequest(format!(
547                "changed file path must be non-empty and repository-relative: {file}"
548            )));
549        }
550    }
551    Ok(())
552}
553
554fn update_remaining_timeout(
555    request: &mut ProviderRequest,
556    deadline: tokio::time::Instant,
557) -> Result<(), ProviderError> {
558    request.budget.timeout = deadline
559        .checked_duration_since(tokio::time::Instant::now())
560        .ok_or_else(timeout_error)?;
561    if request.budget.timeout.is_zero() {
562        return Err(timeout_error());
563    }
564    Ok(())
565}
566
567pub(crate) fn diagnostic_id() -> String {
568    format!(
569        "cg-{:016x}",
570        DIAGNOSTIC_SEQUENCE.fetch_add(1, Ordering::Relaxed)
571    )
572}
573
574pub(crate) fn timeout_error() -> ProviderError {
575    ProviderError::Timeout {
576        diagnostics_id: diagnostic_id(),
577    }
578}
579
580pub(crate) fn output_limit(limit_bytes: usize) -> ProviderError {
581    ProviderError::OutputLimit {
582        limit_bytes,
583        diagnostics_id: diagnostic_id(),
584    }
585}
586
587pub(crate) fn invalid_response(message: impl Into<String>) -> ProviderError {
588    ProviderError::InvalidResponse {
589        message: message.into(),
590        diagnostics_id: diagnostic_id(),
591    }
592}
593
594pub(crate) fn transport_error(message: impl Into<String>) -> ProviderError {
595    ProviderError::Transport {
596        message: message.into(),
597        diagnostics_id: diagnostic_id(),
598    }
599}
600
601fn degradation_from_error(error: &ProviderError) -> ProviderDegradation {
602    let (kind, diagnostics_id) = match error {
603        ProviderError::Cancelled => ("cancelled", None),
604        ProviderError::Timeout { diagnostics_id } => ("timeout", Some(diagnostics_id.clone())),
605        ProviderError::OutputLimit { diagnostics_id, .. } => {
606            ("output_limit", Some(diagnostics_id.clone()))
607        }
608        ProviderError::CircuitOpen => ("circuit_open", None),
609        ProviderError::InvalidRequest(_) => ("invalid_request", None),
610        ProviderError::InvalidResponse { diagnostics_id, .. } => {
611            ("invalid_response", Some(diagnostics_id.clone()))
612        }
613        ProviderError::Transport { diagnostics_id, .. } => {
614            ("transport", Some(diagnostics_id.clone()))
615        }
616    };
617    ProviderDegradation {
618        kind: kind.to_owned(),
619        message: error.to_string(),
620        diagnostics_id,
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use std::path::PathBuf;
627
628    use code_system_graph_model::RepoId;
629    use tokio_util::sync::CancellationToken;
630
631    use super::{CodeGraphConfig, CodeGraphProvider, validate_changed_files};
632    use crate::{ProviderBudget, ProviderError, ProviderRequest};
633
634    #[test]
635    fn provider_should_reject_zero_process_budget() {
636        let result = CodeGraphProvider::new(CodeGraphConfig {
637            max_concurrent_processes: 0,
638            ..CodeGraphConfig::default()
639        });
640
641        assert!(matches!(result, Err(ProviderError::InvalidRequest(_))));
642    }
643
644    #[test]
645    fn changed_files_should_reject_parent_escape() {
646        let result = validate_changed_files(&["../outside.rs".to_owned()]);
647
648        assert!(matches!(result, Err(ProviderError::InvalidRequest(_))));
649    }
650
651    #[tokio::test]
652    async fn cancelled_probe_should_not_start_a_child() {
653        let provider =
654            CodeGraphProvider::new(CodeGraphConfig::default()).expect("valid provider config");
655        let cancellation = CancellationToken::new();
656        cancellation.cancel();
657        let result = crate::LocalCodeIntelligenceProvider::probe(
658            &provider,
659            ProviderRequest {
660                repo_id: RepoId::new("repo:test"),
661                project_path: PathBuf::from("/tmp"),
662                budget: ProviderBudget::default(),
663                cancellation,
664            },
665        )
666        .await;
667
668        assert_eq!(result, Err(ProviderError::Cancelled));
669    }
670}
code_system_graph_core/codegraph/mod.rs

code_system_graph_core/codegraph/
mod.rs
