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

code_system_graph_core/codegraph/
mod.rs
