Skip to main content

coder/
coder.rs

1// Code generation with test loop.
2//
3// Pipeline: planner -> coder -> tester -> (loop back to coder on failure)
4//
5// The coder agent switches its system prompt based on whether test failures
6// exist -- first pass writes from the plan, subsequent passes fix based on
7// test output.
8//
9// Run: cargo run --example coder
10// Requires an LLM (Ollama by default, or set AGENT_LINE_PROVIDER).
11
12use agent_line::{tools, Agent, Ctx, Outcome, Runner, StepResult, Workflow};
13
14// ---------------------------------------------------------------------------
15// State
16// ---------------------------------------------------------------------------
17
18#[derive(Clone, Debug)]
19struct Task {
20    description: String,
21    file_path: String,
22    code: String,
23    test_output: String,
24    attempts: u32,
25    max_attempts: u32,
26}
27
28// ---------------------------------------------------------------------------
29// Agents
30// ---------------------------------------------------------------------------
31
32struct Planner;
33impl Agent<Task> for Planner {
34    fn name(&self) -> &'static str {
35        "planner"
36    }
37    fn run(&mut self, mut state: Task, ctx: &mut Ctx) -> StepResult<Task> {
38        state.code = tools::read_file(&state.file_path).unwrap_or_default();
39
40        let plan = ctx
41            .llm()
42            .system(
43                "You are a senior developer. Create a brief implementation plan. \
44                 List the specific changes needed. Be concise. \
45                 Do not include doc comments or doc tests.",
46            )
47            .user(format!(
48                "Task: {}\n\nFile: {}\n\nCurrent code:\n{}",
49                state.description,
50                state.file_path,
51                if state.code.is_empty() {
52                    "(new file)".to_string()
53                } else {
54                    state.code.clone()
55                }
56            ))
57            .send()?;
58
59        ctx.set("plan", &plan);
60        ctx.log(format!("planner: created plan for {}", state.file_path));
61        Ok((state, Outcome::Continue))
62    }
63}
64
65struct Coder;
66impl Agent<Task> for Coder {
67    fn name(&self) -> &'static str {
68        "coder"
69    }
70    fn run(&mut self, mut state: Task, ctx: &mut Ctx) -> StepResult<Task> {
71        let is_fix = !state.test_output.is_empty();
72
73        let response = if is_fix {
74            ctx.log(format!("coder: fixing (attempt {})", state.attempts));
75
76            ctx.llm()
77                .system(
78                    "You are a developer. Fix the code based on the test failures. \
79                     Return ONLY the complete fixed file contents, no explanation. \
80                     Do not include doc comments or doc tests. \
81                     Do not wrap the output in markdown code fences.",
82                )
83                .user(format!(
84                    "Test errors:\n{}\n\nCurrent code:\n{}",
85                    state.test_output, state.code
86                ))
87                .send()?
88        } else {
89            let plan = ctx.get("plan").unwrap_or("no plan found").to_string();
90            ctx.log("coder: writing initial code");
91
92            ctx.llm()
93                .system(
94                    "You are a developer. Write the code based on the plan. \
95                     Return ONLY the complete file contents, no explanation. \
96                     Do not include doc comments or doc tests. \
97                     Do not wrap the output in markdown code fences.",
98                )
99                .user(format!(
100                    "Plan:\n{plan}\n\nFile: {}\n\nCurrent code:\n{}",
101                    state.file_path, state.code
102                ))
103                .send()?
104        };
105
106        state.code = tools::strip_code_fences(&response);
107        state.test_output.clear();
108        tools::write_file(&state.file_path, &state.code)?;
109        Ok((state, Outcome::Continue))
110    }
111}
112
113struct Tester;
114impl Agent<Task> for Tester {
115    fn name(&self) -> &'static str {
116        "tester"
117    }
118    fn run(&mut self, mut state: Task, ctx: &mut Ctx) -> StepResult<Task> {
119        let manifest = ctx.get("manifest_path").unwrap_or("Cargo.toml").to_string();
120        let result = tools::run_cmd(&format!("cargo test --manifest-path {manifest} --lib"))?;
121
122        if result.success {
123            ctx.log("tester: all passed");
124            Ok((state, Outcome::Done))
125        } else {
126            state.test_output = result.stderr;
127            state.attempts += 1;
128            ctx.log(format!("tester: failed (attempt {})", state.attempts));
129
130            if state.attempts >= state.max_attempts {
131                Ok((
132                    state,
133                    Outcome::Fail("max fix attempts reached, tests still failing".into()),
134                ))
135            } else {
136                Ok((state, Outcome::Next("coder")))
137            }
138        }
139    }
140}
141
142// ---------------------------------------------------------------------------
143// Helpers
144// ---------------------------------------------------------------------------
145
146fn scaffold_project(dir: &std::path::Path) {
147    let src = dir.join("src");
148    std::fs::create_dir_all(&src).unwrap();
149    std::fs::write(
150        dir.join("Cargo.toml"),
151        "[package]\nname = \"scratch\"\nversion = \"0.1.0\"\nedition = \"2024\"\n",
152    )
153    .unwrap();
154    std::fs::write(src.join("lib.rs"), "").unwrap();
155}
156
157// ---------------------------------------------------------------------------
158// Main
159// ---------------------------------------------------------------------------
160
161fn main() {
162    let tmp = std::env::temp_dir().join("agent-line-coder");
163    scaffold_project(&tmp);
164
165    let lib_path = tmp.join("src/lib.rs").display().to_string();
166    let manifest = tmp.join("Cargo.toml").display().to_string();
167
168    let mut ctx = Ctx::new();
169    ctx.set("manifest_path", &manifest);
170
171    let wf = Workflow::builder("coding-agent")
172        .register(Planner)
173        .register(Coder)
174        .register(Tester)
175        .start_at("planner")
176        .then("coder")
177        .then("tester")
178        .build()
179        .unwrap();
180
181    let mut runner = Runner::new(wf);
182
183    let result = runner.run(
184        Task {
185            description: "Add a function called `reverse_string` that reverses a string \
186                          and add unit tests"
187                .into(),
188            file_path: lib_path,
189            code: String::new(),
190            test_output: String::new(),
191            attempts: 0,
192            max_attempts: 3,
193        },
194        &mut ctx,
195    );
196
197    println!("=== Result ===");
198    match result {
199        Ok(task) => {
200            println!("  Success after {} fix attempts", task.attempts);
201            println!("  Final code:\n{}", task.code);
202        }
203        Err(e) => println!("  Failed: {e}"),
204    }
205
206    println!("\n=== Log ===");
207    for entry in ctx.logs() {
208        println!("  {entry}");
209    }
210}