lean_agent_core/
context.rs

1//! Building a compact, high-signal context bundle around one line of a Lean file.
2//!
3//! The daily speedup this replaces is hand-copying imports, the enclosing
4//! theorem, the compiler errors, and the goal state into a chat prompt. Given a
5//! `FILE.lean:LINE` target, [`gather_context`] reads the source, optionally
6//! traces the file once for diagnostics and goal state, and returns a
7//! [`ContextBundle`] that serializes to JSON or renders to Markdown.
8//!
9//! The bundle builder [`build_context`] is pure: it takes already-read source
10//! text plus optional diagnostics, so it is fully testable without touching the
11//! filesystem or spawning Lean.
12
13use crate::{
14    Diagnostic, GoalState, LeanFile, Provenance, Result, TraceConfig, capture_provenance,
15    run_lean_file,
16};
17use camino::Utf8PathBuf;
18use serde::{Deserialize, Serialize};
19use std::time::Duration;
20
21/// Lean keywords that introduce a top-level declaration.
22const NAMED_DECL_KEYWORDS: &[&str] = &[
23    "theorem",
24    "lemma",
25    "def",
26    "abbrev",
27    "instance",
28    "example",
29    "structure",
30    "inductive",
31    "class",
32    "opaque",
33    "axiom",
34];
35
36/// Modifiers that may precede a declaration keyword on the same line.
37const DECL_MODIFIERS: &[&str] = &[
38    "private",
39    "protected",
40    "noncomputable",
41    "partial",
42    "unsafe",
43    "scoped",
44    "local",
45    "nonrec",
46];
47
48/// Column-zero keywords that close the current declaration's text block.
49const BOUNDARY_KEYWORDS: &[&str] = &[
50    "namespace",
51    "section",
52    "end",
53    "open",
54    "import",
55    "variable",
56    "universe",
57    "set_option",
58    "attribute",
59    "mutual",
60    "macro",
61    "macro_rules",
62    "syntax",
63    "notation",
64    "elab",
65    "deriving",
66];
67
68/// Where to center a context bundle and how much surrounding source to include.
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct ContextRequest {
71    /// Target Lean source file.
72    pub file: LeanFile,
73    /// One-based line the bundle is centered on.
74    pub line: u32,
75    /// Source lines to include before the target line.
76    pub before: usize,
77    /// Source lines to include after the target line.
78    pub after: usize,
79}
80
81impl ContextRequest {
82    /// Build a request with the default eight-line window on each side.
83    #[must_use]
84    pub const fn new(file: LeanFile, line: u32) -> Self {
85        Self {
86            file,
87            line,
88            before: 8,
89            after: 8,
90        }
91    }
92}
93
94/// Runtime options for [`gather_context`] when a live Lean trace is wanted.
95#[derive(Clone, Debug, Eq, PartialEq)]
96pub struct ContextOptions {
97    /// Run `lake lean` over the file to attach diagnostics and goal state.
98    pub run_trace: bool,
99    /// Lake workspace root used for the trace and provenance probes.
100    pub lake_root: Utf8PathBuf,
101    /// Process timeout for the trace.
102    pub timeout: Duration,
103    /// Keep warning diagnostics in the bundle.
104    pub include_warnings: bool,
105}
106
107impl ContextOptions {
108    /// Tracing options rooted at `lake_root` with a sixty-second timeout.
109    #[must_use]
110    pub fn new(lake_root: Utf8PathBuf) -> Self {
111        Self {
112            run_trace: true,
113            lake_root,
114            timeout: Duration::from_secs(60),
115            include_warnings: true,
116        }
117    }
118}
119
120/// The enclosing declaration detected around the target line.
121#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
122pub struct Declaration {
123    /// Declaration keyword, such as `theorem` or `def`.
124    pub kind: String,
125    /// Declared name when one is present (absent for `example`).
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub name: Option<String>,
128    /// One-based first line of the declaration text.
129    pub start_line: u32,
130    /// One-based last line of the declaration text.
131    pub end_line: u32,
132    /// Verbatim declaration source.
133    pub source: String,
134}
135
136/// One numbered line of the surrounding source window.
137#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
138pub struct SourceLine {
139    /// One-based line number in the file.
140    pub number: u32,
141    /// Line text without its trailing newline.
142    pub text: String,
143    /// Whether this is the requested target line.
144    pub is_target: bool,
145}
146
147/// A window of source lines centered on the target line.
148#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
149pub struct SourceWindow {
150    /// One-based first line in the window.
151    pub start_line: u32,
152    /// One-based last line in the window.
153    pub end_line: u32,
154    /// One-based target line the window is centered on.
155    pub target_line: u32,
156    /// The numbered window lines.
157    pub lines: Vec<SourceLine>,
158}
159
160/// A compact context bundle around one line of a Lean file.
161#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
162pub struct ContextBundle {
163    /// Source file the bundle describes.
164    pub file: LeanFile,
165    /// One-based target line, clamped into the file's range.
166    pub line: u32,
167    /// Total number of source lines in the file.
168    pub total_lines: usize,
169    /// Enclosing declaration, when one is detectable.
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub declaration: Option<Declaration>,
172    /// Import lines found in the file, in source order.
173    pub imports: Vec<String>,
174    /// Surrounding source window.
175    pub surrounding: SourceWindow,
176    /// Diagnostics relevant to the target line, when a trace was run.
177    pub diagnostics: Vec<Diagnostic>,
178    /// Goal state extracted from a relevant diagnostic, when present.
179    #[serde(skip_serializing_if = "Option::is_none")]
180    pub goal_state: Option<GoalState>,
181    /// Tooling provenance, when a trace was run.
182    #[serde(skip_serializing_if = "Option::is_none")]
183    pub provenance: Option<Provenance>,
184    /// Ready-to-paste prompt assembled from the fields above.
185    pub suggested_prompt: String,
186}
187
188impl ContextBundle {
189    /// Render the bundle as Markdown for human review or a chat paste.
190    #[must_use]
191    pub fn to_markdown(&self) -> String {
192        render_markdown(self)
193    }
194}
195
196/// Parse a `FILE.lean:LINE` target into its path and one-based line.
197///
198/// The split is on the final colon so Windows-style drive letters and paths
199/// with colons in earlier segments still parse.
200pub fn parse_file_line_spec(spec: &str) -> Result<(Utf8PathBuf, u32)> {
201    let (path, line) = spec
202        .rsplit_once(':')
203        .ok_or_else(|| crate::Error::InvalidLineSpec {
204            spec: spec.to_owned(),
205        })?;
206    let line: u32 = line.parse().map_err(|_| crate::Error::InvalidLineSpec {
207        spec: spec.to_owned(),
208    })?;
209    if path.is_empty() || line == 0 {
210        return Err(crate::Error::InvalidLineSpec {
211            spec: spec.to_owned(),
212        });
213    }
214    Ok((Utf8PathBuf::from(path), line))
215}
216
217/// Read the file, optionally trace it once, and assemble a [`ContextBundle`].
218pub async fn gather_context(
219    request: &ContextRequest,
220    options: &ContextOptions,
221) -> Result<ContextBundle> {
222    let source = std::fs::read_to_string(request.file.as_path())?;
223
224    let (diagnostics, provenance) = if options.run_trace {
225        let provenance = capture_provenance(options.lake_root.as_path()).await;
226        let mut config = TraceConfig::new(options.lake_root.clone()).timeout(options.timeout);
227        config.include_warnings = options.include_warnings;
228        let trace = run_lean_file(&config, &provenance, request.file.clone()).await;
229        (trace.diagnostics, Some(provenance))
230    } else {
231        (Vec::new(), None)
232    };
233
234    Ok(build_context(
235        request,
236        &source,
237        &diagnostics,
238        provenance.as_ref(),
239    ))
240}
241
242/// Assemble a bundle from already-read source and optional trace results.
243///
244/// This is the pure core: no filesystem access and no process spawning, so it
245/// is the right entry point for tests and for callers that already hold the
246/// source text and diagnostics.
247#[must_use]
248pub fn build_context(
249    request: &ContextRequest,
250    source: &str,
251    diagnostics: &[Diagnostic],
252    provenance: Option<&Provenance>,
253) -> ContextBundle {
254    let lines: Vec<&str> = source.lines().collect();
255    let total_lines = lines.len();
256    let target_line = clamp_line(request.line, total_lines);
257    let target_idx = (target_line as usize).saturating_sub(1);
258
259    let imports = collect_imports(&lines);
260    let declaration = detect_declaration(&lines, target_idx);
261    let surrounding = build_window(&lines, target_idx, request.before, request.after);
262    let (selected, goal_state) = select_diagnostics(diagnostics, target_line);
263
264    let suggested_prompt = build_prompt(
265        request.file.as_path().as_str(),
266        target_line,
267        total_lines,
268        declaration.as_ref(),
269        &imports,
270        &selected,
271        goal_state.as_ref(),
272    );
273
274    ContextBundle {
275        file: request.file.clone(),
276        line: target_line,
277        total_lines,
278        declaration,
279        imports,
280        surrounding,
281        diagnostics: selected,
282        goal_state,
283        provenance: provenance.cloned(),
284        suggested_prompt,
285    }
286}
287
288/// Clamp a one-based line into `[1, total]`, treating an empty file as line one.
289fn clamp_line(line: u32, total: usize) -> u32 {
290    if total == 0 {
291        return 1;
292    }
293    let max = total as u32;
294    line.clamp(1, max)
295}
296
297/// Collect `import` lines (those starting at column zero) in source order.
298pub fn collect_imports(lines: &[&str]) -> Vec<String> {
299    lines
300        .iter()
301        .filter(|line| is_col0(line) && first_token(line) == Some("import"))
302        .map(|line| line.trim_end().to_owned())
303        .collect()
304}
305
306/// True when a line starts a top-level command (no leading whitespace).
307fn is_col0(line: &str) -> bool {
308    !line.is_empty() && !line.starts_with(char::is_whitespace)
309}
310
311/// First whitespace-delimited token of a line, if any.
312fn first_token(line: &str) -> Option<&str> {
313    line.split_whitespace().next()
314}
315
316/// Declaration keyword for a column-zero line, skipping leading modifiers.
317fn decl_kind_of_line(line: &str) -> Option<&str> {
318    if !is_col0(line) {
319        return None;
320    }
321    let keyword = line
322        .split_whitespace()
323        .find(|token| !DECL_MODIFIERS.contains(token))?;
324    NAMED_DECL_KEYWORDS.contains(&keyword).then_some(keyword)
325}
326
327/// True when a column-zero line ends the current declaration's text block.
328fn is_boundary(line: &str) -> bool {
329    if !is_col0(line) {
330        return false;
331    }
332    if line.starts_with("@[") {
333        return true;
334    }
335    let Some(token) = first_token(line) else {
336        return false;
337    };
338    if token.starts_with('#') {
339        return true;
340    }
341    BOUNDARY_KEYWORDS.contains(&token) || decl_kind_of_line(line).is_some()
342}
343
344/// Detect the declaration enclosing `target_idx`, if any.
345///
346/// `target_idx` is a zero-based line index. The returned [`Declaration`] carries
347/// one-based start and end lines plus the verbatim source of the block.
348pub fn detect_declaration(lines: &[&str], target_idx: usize) -> Option<Declaration> {
349    if lines.is_empty() {
350        return None;
351    }
352    let scan_from = target_idx.min(lines.len() - 1);
353    let mut keyword_idx = None;
354    for idx in (0..=scan_from).rev() {
355        if decl_kind_of_line(lines[idx]).is_some() {
356            keyword_idx = Some(idx);
357            break;
358        }
359        if is_boundary(lines[idx]) {
360            // A boundary above the cursor with no enclosing decl means the
361            // cursor sits between declarations.
362            break;
363        }
364    }
365    let keyword_idx = keyword_idx?;
366    let kind = decl_kind_of_line(lines[keyword_idx])?.to_owned();
367
368    // Absorb attribute lines directly above the keyword into the declaration.
369    let mut start = keyword_idx;
370    while start > 0 {
371        let prev = lines[start - 1];
372        if is_col0(prev) && prev.starts_with("@[") {
373            start -= 1;
374        } else {
375            break;
376        }
377    }
378
379    // The block runs until the next top-level boundary or end of file.
380    let mut end = lines.len();
381    for (offset, line) in lines.iter().enumerate().skip(keyword_idx + 1) {
382        if is_boundary(line) {
383            end = offset;
384            break;
385        }
386    }
387
388    // Trim trailing blank lines from the captured block.
389    while end > keyword_idx + 1 && lines[end - 1].trim().is_empty() {
390        end -= 1;
391    }
392
393    let source = lines[start..end].join("\n");
394    Some(Declaration {
395        name: decl_name(lines[keyword_idx], &kind),
396        kind,
397        start_line: (start as u32) + 1,
398        end_line: end as u32,
399        source,
400    })
401}
402
403/// Extract the declared name from a keyword line, when one is present.
404fn decl_name(line: &str, kind: &str) -> Option<String> {
405    if kind == "example" {
406        return None;
407    }
408    let mut tokens = line.split_whitespace().skip_while(|t| t != &kind);
409    let _keyword = tokens.next()?;
410    let candidate = tokens.next()?;
411    let cut = candidate
412        .find(['(', '{', '[', ':'])
413        .unwrap_or(candidate.len());
414    let name = &candidate[..cut];
415    let valid = !name.is_empty()
416        && name
417            .chars()
418            .next()
419            .is_some_and(|c| c.is_alphabetic() || c == '_');
420    valid.then(|| name.to_owned())
421}
422
423/// Build the surrounding source window.
424fn build_window(lines: &[&str], target_idx: usize, before: usize, after: usize) -> SourceWindow {
425    if lines.is_empty() {
426        return SourceWindow {
427            start_line: 1,
428            end_line: 0,
429            target_line: 1,
430            lines: Vec::new(),
431        };
432    }
433    let start = target_idx.saturating_sub(before);
434    let end = (target_idx + after + 1).min(lines.len());
435    let window = lines[start..end]
436        .iter()
437        .enumerate()
438        .map(|(offset, text)| {
439            let number = (start + offset + 1) as u32;
440            SourceLine {
441                number,
442                text: (*text).to_owned(),
443                is_target: start + offset == target_idx,
444            }
445        })
446        .collect();
447    SourceWindow {
448        start_line: (start as u32) + 1,
449        end_line: end as u32,
450        target_line: (target_idx as u32) + 1,
451        lines: window,
452    }
453}
454
455/// Pick diagnostics for the target line, falling back to all file diagnostics.
456fn select_diagnostics(all: &[Diagnostic], line: u32) -> (Vec<Diagnostic>, Option<GoalState>) {
457    let on_line: Vec<Diagnostic> = all
458        .iter()
459        .filter(|d| d.line == Some(line))
460        .cloned()
461        .collect();
462    let chosen = if on_line.is_empty() {
463        all.to_vec()
464    } else {
465        on_line
466    };
467    let goal = chosen
468        .iter()
469        .find_map(|d| d.goal_state.clone())
470        .or_else(|| chosen.iter().find_map(|d| goal_from_message(&d.message)));
471    (chosen, goal)
472}
473
474/// Recover a goal block from a diagnostic message that the upstream parser did
475/// not tag (for example a `Tactic ... failed` error that still prints `⊢`).
476///
477/// The block is the turnstile line plus the contiguous local-context lines
478/// directly above it, stopping at the first blank line.
479fn goal_from_message(message: &str) -> Option<GoalState> {
480    let lines: Vec<&str> = message.lines().collect();
481    let goal_idx = lines
482        .iter()
483        .rposition(|line| line.trim_start().starts_with('⊢'))?;
484    let mut start = goal_idx;
485    while start > 0 && !lines[start - 1].trim().is_empty() {
486        start -= 1;
487    }
488    let block = lines[start..=goal_idx].join("\n");
489    let block = block.trim();
490    (!block.is_empty()).then(|| GoalState(block.to_owned()))
491}
492
493/// Assemble the ready-to-paste prompt from the bundle's parts.
494fn build_prompt(
495    file: &str,
496    line: u32,
497    total_lines: usize,
498    declaration: Option<&Declaration>,
499    imports: &[String],
500    diagnostics: &[Diagnostic],
501    goal_state: Option<&GoalState>,
502) -> String {
503    let mut out = String::new();
504    out.push_str("You are editing a Lean 4 proof.\n\n");
505    out.push_str(&format!("File: {file} (line {line} of {total_lines})\n"));
506
507    match declaration {
508        Some(decl) => {
509            let named = decl
510                .name
511                .as_deref()
512                .map_or_else(|| decl.kind.clone(), |name| format!("{} {name}", decl.kind));
513            out.push_str(&format!("Declaration: {named}\n"));
514        }
515        None => out.push_str("Declaration: not detected\n"),
516    }
517
518    out.push('\n');
519    if imports.is_empty() {
520        out.push_str("Imports in scope: none found in this file.\n");
521    } else {
522        out.push_str("Imports in scope:\n");
523        for import in imports {
524            out.push_str(import);
525            out.push('\n');
526        }
527    }
528
529    if let Some(decl) = declaration {
530        out.push_str("\nCurrent declaration:\n");
531        out.push_str(&decl.source);
532        out.push('\n');
533    }
534
535    if let Some(goal) = goal_state {
536        out.push_str("\nGoal state at the target line:\n");
537        out.push_str(&goal.0);
538        out.push('\n');
539    }
540
541    if diagnostics.is_empty() {
542        out.push_str("\nCompiler diagnostics: none captured.\n");
543    } else {
544        out.push_str("\nCompiler diagnostics:\n");
545        for diagnostic in diagnostics {
546            let first = diagnostic.message.lines().next().unwrap_or("").trim();
547            out.push_str(&format!("- {first}\n"));
548        }
549    }
550
551    out.push_str(
552        "\nTask: rewrite the declaration so the file compiles with no errors. \
553Return only the corrected Lean source for this declaration. Keep the existing \
554name and signature unless the error requires a change, and do not add imports \
555unless they are needed.\n",
556    );
557    out
558}
559
560/// Render a [`ContextBundle`] as Markdown.
561fn render_markdown(bundle: &ContextBundle) -> String {
562    let mut out = String::new();
563    out.push_str(&format!("# Lean context: {}\n\n", bundle.file));
564    out.push_str(&format!(
565        "Target line **{}** of {} total.\n\n",
566        bundle.line, bundle.total_lines
567    ));
568
569    if let Some(decl) = &bundle.declaration {
570        let named = decl.name.as_deref().map_or_else(
571            || decl.kind.clone(),
572            |name| format!("{} `{name}`", decl.kind),
573        );
574        out.push_str(&format!(
575            "## Declaration\n\n{named} (lines {}-{})\n\n```lean\n{}\n```\n\n",
576            decl.start_line, decl.end_line, decl.source
577        ));
578    }
579
580    out.push_str("## Imports\n\n");
581    if bundle.imports.is_empty() {
582        out.push_str("None found in this file.\n\n");
583    } else {
584        out.push_str("```lean\n");
585        for import in &bundle.imports {
586            out.push_str(import);
587            out.push('\n');
588        }
589        out.push_str("```\n\n");
590    }
591
592    out.push_str("## Surrounding source\n\n```lean\n");
593    for line in &bundle.surrounding.lines {
594        let marker = if line.is_target { ">" } else { " " };
595        out.push_str(&format!("{marker} {:>4} | {}\n", line.number, line.text));
596    }
597    out.push_str("```\n\n");
598
599    if let Some(goal) = &bundle.goal_state {
600        out.push_str(&format!("## Goal state\n\n```\n{}\n```\n\n", goal.0));
601    }
602
603    out.push_str("## Diagnostics\n\n");
604    if bundle.diagnostics.is_empty() {
605        out.push_str("None captured.\n\n");
606    } else {
607        for diagnostic in &bundle.diagnostics {
608            let first = diagnostic.message.lines().next().unwrap_or("").trim();
609            out.push_str(&format!("- {first}\n"));
610        }
611        out.push('\n');
612    }
613
614    if let Some(provenance) = &bundle.provenance {
615        out.push_str("## Provenance\n\n");
616        out.push_str(&format!(
617            "- lean: {}\n- lake: {}\n- git: {}\n\n",
618            provenance.lean_version.as_deref().unwrap_or("unknown"),
619            provenance.lake_version.as_deref().unwrap_or("unknown"),
620            provenance.git_commit.as_deref().unwrap_or("none"),
621        ));
622    }
623
624    out.push_str("## Suggested prompt\n\n```\n");
625    out.push_str(&bundle.suggested_prompt);
626    out.push_str("```\n");
627    out
628}
629
630#[cfg(test)]
631mod tests {
632    use super::*;
633    use crate::DiagnosticSeverity;
634
635    const SAMPLE: &str = "import Init\nimport Mathlib.Tactic\n\nnamespace Demo\n\n@[simp]\ntheorem foo (n : Nat) : n = n := by\n  rfl\n\ndef bar : Nat := 0\n\nend Demo\n";
636
637    fn request(line: u32) -> ContextRequest {
638        let file = LeanFile(Utf8PathBuf::from("Demo.lean"));
639        ContextRequest::new(file, line)
640    }
641
642    #[test]
643    fn parses_file_line_spec() -> Result<()> {
644        let (path, line) = parse_file_line_spec("src/Demo.lean:42")?;
645        assert_eq!(path, Utf8PathBuf::from("src/Demo.lean"));
646        assert_eq!(line, 42);
647        Ok(())
648    }
649
650    #[test]
651    fn rejects_bad_line_spec() {
652        assert!(parse_file_line_spec("no-line-here").is_err());
653        assert!(parse_file_line_spec("file.lean:0").is_err());
654        assert!(parse_file_line_spec("file.lean:abc").is_err());
655        assert!(parse_file_line_spec(":12").is_err());
656    }
657
658    #[test]
659    fn collects_imports() {
660        let bundle = build_context(&request(7), SAMPLE, &[], None);
661        assert_eq!(bundle.imports, vec!["import Init", "import Mathlib.Tactic"]);
662    }
663
664    #[test]
665    fn detects_enclosing_theorem_with_attribute() -> Result<()> {
666        let bundle = build_context(&request(8), SAMPLE, &[], None);
667        let Some(decl) = bundle.declaration else {
668            return Err(crate::Error::Todo {
669                feature: "expected a declaration",
670            });
671        };
672        assert_eq!(decl.kind, "theorem");
673        assert_eq!(decl.name.as_deref(), Some("foo"));
674        // The attribute line above the keyword is absorbed.
675        assert_eq!(decl.start_line, 6);
676        assert!(decl.source.starts_with("@[simp]"));
677        assert!(decl.source.contains("rfl"));
678        // The block stops before the next declaration.
679        assert!(!decl.source.contains("def bar"));
680        Ok(())
681    }
682
683    #[test]
684    fn detects_following_def() -> Result<()> {
685        let bundle = build_context(&request(10), SAMPLE, &[], None);
686        let Some(decl) = bundle.declaration else {
687            return Err(crate::Error::Todo {
688                feature: "expected a declaration",
689            });
690        };
691        assert_eq!(decl.kind, "def");
692        assert_eq!(decl.name.as_deref(), Some("bar"));
693        Ok(())
694    }
695
696    #[test]
697    fn window_marks_target_and_respects_bounds() {
698        let mut req = request(7);
699        req.before = 2;
700        req.after = 1;
701        let bundle = build_context(&req, SAMPLE, &[], None);
702        assert_eq!(bundle.surrounding.start_line, 5);
703        assert_eq!(bundle.surrounding.end_line, 8);
704        let target: Vec<_> = bundle
705            .surrounding
706            .lines
707            .iter()
708            .filter(|l| l.is_target)
709            .collect();
710        assert_eq!(target.len(), 1);
711        assert_eq!(target[0].number, 7);
712    }
713
714    #[test]
715    fn clamps_out_of_range_line() {
716        let bundle = build_context(&request(9999), SAMPLE, &[], None);
717        assert_eq!(bundle.total_lines, 12);
718        assert_eq!(bundle.line, 12);
719    }
720
721    #[test]
722    fn selects_diagnostics_on_target_line_and_extracts_goal() {
723        let diags = vec![
724            Diagnostic {
725                file: Some(Utf8PathBuf::from("Demo.lean")),
726                line: Some(8),
727                column: Some(2),
728                severity: DiagnosticSeverity::Error,
729                message: "error: unsolved goals\n⊢ n = n".to_owned(),
730                goal_state: Some(GoalState("⊢ n = n".to_owned())),
731            },
732            Diagnostic {
733                file: Some(Utf8PathBuf::from("Demo.lean")),
734                line: Some(99),
735                column: None,
736                severity: DiagnosticSeverity::Warning,
737                message: "warning: elsewhere".to_owned(),
738                goal_state: None,
739            },
740        ];
741        let bundle = build_context(&request(8), SAMPLE, &diags, None);
742        assert_eq!(bundle.diagnostics.len(), 1);
743        assert_eq!(bundle.diagnostics[0].line, Some(8));
744        assert_eq!(
745            bundle.goal_state.as_ref().map(|g| g.0.as_str()),
746            Some("⊢ n = n")
747        );
748        assert!(bundle.suggested_prompt.contains("Goal state"));
749    }
750
751    #[test]
752    fn recovers_goal_from_untagged_message() {
753        let diags = vec![Diagnostic {
754            file: Some(Utf8PathBuf::from("Demo.lean")),
755            line: Some(7),
756            column: Some(2),
757            severity: DiagnosticSeverity::Error,
758            message: "error: Tactic `rfl` failed: the sides differ\n  n + m\n\nn m : Nat\n⊢ n + m = m + n".to_owned(),
759            goal_state: None,
760        }];
761        let bundle = build_context(&request(7), SAMPLE, &diags, None);
762        assert_eq!(
763            bundle.goal_state.as_ref().map(|g| g.0.as_str()),
764            Some("n m : Nat\n⊢ n + m = m + n")
765        );
766    }
767
768    #[test]
769    fn empty_file_is_handled() {
770        let bundle = build_context(&request(1), "", &[], None);
771        assert_eq!(bundle.total_lines, 0);
772        assert!(bundle.declaration.is_none());
773        assert!(bundle.surrounding.lines.is_empty());
774    }
775
776    #[test]
777    fn markdown_includes_key_sections() {
778        let bundle = build_context(&request(8), SAMPLE, &[], None);
779        let md = bundle.to_markdown();
780        assert!(md.contains("# Lean context"));
781        assert!(md.contains("## Declaration"));
782        assert!(md.contains("## Imports"));
783        assert!(md.contains("## Surrounding source"));
784        assert!(md.contains("## Suggested prompt"));
785        assert!(md.contains("> "));
786    }
787
788    #[test]
789    fn prompt_opens_with_editing_lead() {
790        let bundle = build_context(&request(8), SAMPLE, &[], None);
791        assert!(
792            bundle
793                .suggested_prompt
794                .starts_with("You are editing a Lean 4 proof.")
795        );
796    }
797}
lean_agent_core/context.rs

lean_agent_core/
context.rs
