Skip to main content

molo_coding/coding/
context.rs

1use crate::{RunContext, RunMetadata};
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4
5use super::git::{GitChangedFilesRequest, GitInspector, GitStatus};
6use super::instructions::{InstructionBundle, InstructionRequest, InstructionResolver};
7use super::search::{RepoSearchRequest, RepoSearchResults, RepoSearcher, SearchMode};
8use super::test_runner::VerificationResult;
9use super::workspace::{ListFilesQuery, Workspace, WorkspaceEntry, WorkspacePath};
10
11/// Context budget for repository context gathering.
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct ContextBudget {
14    /// Approximate maximum bytes across textual context fields.
15    pub max_bytes: usize,
16    /// Maximum files in tree summaries.
17    pub max_files: usize,
18    /// Maximum search matches.
19    pub max_search_matches: usize,
20}
21
22impl Default for ContextBudget {
23    fn default() -> Self {
24        Self {
25            max_bytes: 64 * 1024,
26            max_files: 500,
27            max_search_matches: 50,
28        }
29    }
30}
31
32/// Flags controlling which context sources are gathered.
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34pub struct CodingContextInclude {
35    /// Include project instructions.
36    pub instructions: bool,
37    /// Include repository tree summary.
38    pub repo_tree: bool,
39    /// Include search matches for goal terms.
40    pub search: bool,
41    /// Include git status.
42    pub git_status: bool,
43    /// Include changed files.
44    pub changed_files: bool,
45}
46
47impl Default for CodingContextInclude {
48    fn default() -> Self {
49        Self {
50            instructions: true,
51            repo_tree: true,
52            search: true,
53            git_status: true,
54            changed_files: true,
55        }
56    }
57}
58
59/// Coding context request.
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct CodingContextRequest {
62    /// User or agent goal.
63    pub goal: String,
64    /// Paths to prioritize.
65    pub focus_paths: Vec<WorkspacePath>,
66    /// Included context sources.
67    pub include: CodingContextInclude,
68    /// Context budget.
69    pub budget: ContextBudget,
70    /// Recent verification failures.
71    pub recent_test_failures: Vec<VerificationResult>,
72    /// Recent transcript summary.
73    pub transcript_summary: Option<String>,
74}
75
76impl CodingContextRequest {
77    /// Constructs a context request from a goal.
78    pub fn new(goal: impl Into<String>) -> Self {
79        Self {
80            goal: goal.into(),
81            focus_paths: Vec::new(),
82            include: CodingContextInclude::default(),
83            budget: ContextBudget::default(),
84            recent_test_failures: Vec::new(),
85            transcript_summary: None,
86        }
87    }
88}
89
90/// Context bundle returned by a [`CodingContextProvider`].
91#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
92pub struct CodingContextBundle {
93    /// Resolved project instructions.
94    pub instructions: Option<InstructionBundle>,
95    /// Repository tree entries.
96    pub repo_tree: Vec<WorkspaceEntry>,
97    /// Search results.
98    pub search_results: Vec<RepoSearchResults>,
99    /// Git status.
100    pub git_status: Option<GitStatus>,
101    /// Changed files.
102    pub changed_files: Vec<WorkspacePath>,
103    /// Focus and changed paths considered relevant by deterministic
104    /// heuristics.
105    pub relevant_files: Vec<WorkspacePath>,
106    /// Dependency manifest hints discovered in the repository tree.
107    pub dependency_metadata: Vec<DependencyMetadata>,
108    /// Recent verification failures.
109    pub recent_test_failures: Vec<VerificationResult>,
110    /// Recent transcript summary.
111    pub transcript_summary: Option<String>,
112    /// Warnings about truncation, unsupported tools, or missing context.
113    pub warnings: Vec<String>,
114    /// Approximate bytes used.
115    pub bytes_used: usize,
116    /// Whether the bundle hit a configured budget.
117    pub truncated: bool,
118    /// Host-owned metadata.
119    pub metadata: RunMetadata,
120}
121
122/// Dependency manifest metadata discovered by context gathering.
123#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
124pub struct DependencyMetadata {
125    /// Manifest path.
126    pub path: WorkspacePath,
127    /// Manifest ecosystem or format.
128    pub kind: String,
129}
130
131/// Provides repository context outside chat memory.
132#[async_trait]
133pub trait CodingContextProvider: Send + Sync {
134    /// Gathers coding context.
135    async fn gather(
136        &self,
137        request: CodingContextRequest,
138        context: &RunContext,
139    ) -> Result<CodingContextBundle, CodingContextError>;
140}
141
142/// Baseline context provider that combines workspace, search, git, and
143/// instruction primitives.
144#[derive(Debug, Clone)]
145pub struct DefaultCodingContextProvider<W, S, G, I> {
146    workspace: W,
147    searcher: S,
148    git: G,
149    instructions: I,
150}
151
152impl<W, S, G, I> DefaultCodingContextProvider<W, S, G, I> {
153    /// Constructs a default context provider.
154    pub fn new(workspace: W, searcher: S, git: G, instructions: I) -> Self {
155        Self {
156            workspace,
157            searcher,
158            git,
159            instructions,
160        }
161    }
162}
163
164#[async_trait]
165impl<W, S, G, I> CodingContextProvider for DefaultCodingContextProvider<W, S, G, I>
166where
167    W: Workspace,
168    S: RepoSearcher,
169    G: GitInspector,
170    I: InstructionResolver,
171{
172    async fn gather(
173        &self,
174        request: CodingContextRequest,
175        context: &RunContext,
176    ) -> Result<CodingContextBundle, CodingContextError> {
177        let mut warnings = Vec::new();
178        let mut bytes_used = 0usize;
179        let mut truncated = false;
180
181        let instructions = if request.include.instructions {
182            match self
183                .instructions
184                .resolve(
185                    InstructionRequest {
186                        target_paths: request.focus_paths.clone(),
187                        max_bytes: request.budget.max_bytes,
188                        ..InstructionRequest::default()
189                    },
190                    context,
191                )
192                .await
193            {
194                Ok(bundle) => {
195                    bytes_used += bundle
196                        .files
197                        .iter()
198                        .map(|file| file.content.len())
199                        .sum::<usize>();
200                    truncated |= bundle.truncated;
201                    Some(bundle)
202                }
203                Err(error) => {
204                    warnings.push(error.to_string());
205                    None
206                }
207            }
208        } else {
209            None
210        };
211
212        let repo_tree = if request.include.repo_tree {
213            let entries = self
214                .workspace
215                .list_files(ListFilesQuery {
216                    path: WorkspacePath::root(),
217                    recursive: true,
218                    max_entries: Some(request.budget.max_files),
219                    include_hidden: false,
220                    respect_gitignore: true,
221                })
222                .await
223                .map_err(|error| CodingContextError::Workspace {
224                    message: error.to_string(),
225                })?;
226            if entries.len() >= request.budget.max_files {
227                truncated = true;
228            }
229            bytes_used += entries
230                .iter()
231                .map(|entry| entry.path.display().len())
232                .sum::<usize>();
233            entries
234        } else {
235            Vec::new()
236        };
237
238        let search_results = if request.include.search && !request.goal.trim().is_empty() {
239            let query = derive_query(&request.goal);
240            match self
241                .searcher
242                .search(
243                    RepoSearchRequest {
244                        query,
245                        paths: request.focus_paths.clone(),
246                        mode: SearchMode::Literal,
247                        max_matches: request.budget.max_search_matches,
248                        context_lines: 1,
249                        include_hidden: false,
250                        respect_gitignore: true,
251                    },
252                    context,
253                )
254                .await
255            {
256                Ok(results) => {
257                    truncated |= results.truncated;
258                    bytes_used += results
259                        .matches
260                        .iter()
261                        .map(|mat| mat.line_text.len())
262                        .sum::<usize>();
263                    vec![results]
264                }
265                Err(error) => {
266                    warnings.push(error.to_string());
267                    Vec::new()
268                }
269            }
270        } else {
271            Vec::new()
272        };
273
274        let git_status = if request.include.git_status {
275            match self.git.status(Default::default(), context).await {
276                Ok(status) => {
277                    bytes_used += status.raw.len();
278                    truncated |= status.truncated;
279                    Some(status)
280                }
281                Err(error) => {
282                    warnings.push(error.to_string());
283                    None
284                }
285            }
286        } else {
287            None
288        };
289
290        let changed_files = if request.include.changed_files {
291            match self
292                .git
293                .changed_files(GitChangedFilesRequest::default(), context)
294                .await
295            {
296                Ok(files) => files.into_iter().map(|file| file.path).collect(),
297                Err(error) => {
298                    warnings.push(error.to_string());
299                    Vec::new()
300                }
301            }
302        } else {
303            Vec::new()
304        };
305        let mut relevant_files = request.focus_paths.clone();
306        for path in &changed_files {
307            if !relevant_files.contains(path) {
308                relevant_files.push(path.clone());
309            }
310        }
311        let dependency_metadata = dependency_metadata_from_tree(&repo_tree);
312
313        if bytes_used > request.budget.max_bytes {
314            truncated = true;
315            warnings.push("context byte budget exceeded".to_string());
316        }
317
318        Ok(CodingContextBundle {
319            instructions,
320            repo_tree,
321            search_results,
322            git_status,
323            changed_files,
324            relevant_files,
325            dependency_metadata,
326            recent_test_failures: request.recent_test_failures,
327            transcript_summary: request.transcript_summary,
328            warnings,
329            bytes_used,
330            truncated,
331            metadata: RunMetadata::new(),
332        })
333    }
334}
335
336fn dependency_metadata_from_tree(entries: &[WorkspaceEntry]) -> Vec<DependencyMetadata> {
337    entries
338        .iter()
339        .filter_map(|entry| {
340            let display = entry.path.display();
341            let file_name = display.rsplit('/').next().unwrap_or(display.as_str());
342            let kind = match file_name {
343                "Cargo.toml" => "cargo",
344                "package.json" => "npm",
345                "pnpm-lock.yaml" => "pnpm",
346                "yarn.lock" => "yarn",
347                "pyproject.toml" => "python",
348                "requirements.txt" => "python",
349                "go.mod" => "go",
350                "pom.xml" => "maven",
351                "build.gradle" | "build.gradle.kts" => "gradle",
352                _ => return None,
353            };
354            Some(DependencyMetadata {
355                path: entry.path.clone(),
356                kind: kind.to_string(),
357            })
358        })
359        .collect()
360}
361
362fn derive_query(goal: &str) -> String {
363    goal.split_whitespace()
364        .find(|word| word.len() >= 4)
365        .unwrap_or(goal)
366        .trim_matches(|ch: char| !ch.is_alphanumeric() && ch != '_')
367        .to_string()
368}
369
370/// Coding context errors.
371#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, Serialize, Deserialize)]
372#[non_exhaustive]
373pub enum CodingContextError {
374    /// Workspace context failed.
375    #[error("context workspace error: {message}")]
376    Workspace {
377        /// Model-safe explanation.
378        message: String,
379    },
380    /// Context source failed.
381    #[error("context source error: {message}")]
382    Source {
383        /// Model-safe explanation.
384        message: String,
385    },
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[test]
393    fn derive_query_picks_goal_term() {
394        assert_eq!(derive_query("fix parser panic"), "parser");
395    }
396}