← Back to DevBytes

How to Build an Agent That Writes and Tests Its Own Code

How to Build an Agent That Writes and Tests Its Own Code

Self-coding agents represent one of the most exciting frontiers in applied AI. Instead of treating the language model as a passive code-completion tool, you give it an environment, a goal, and a feedback loop — then let it write code, run that code, read the errors, and iterate until the task is done. This tutorial walks through building such an agent from scratch using Python, the OpenAI API, and a sandboxed execution environment.

What Is a Self-Writing Code Agent?

A self-writing code agent is an autonomous loop that combines three capabilities: generation (producing source code from a natural-language prompt), execution (running that code in a real interpreter), and reflection (reading stdout, stderr, or test results and deciding what to fix). Unlike a one-shot code generator, the agent closes the loop: every failed run becomes new context for the next attempt.

Conceptually, the agent follows this cycle:

while task_not_complete:
    code = llm.generate(prompt + history)
    result = sandbox.run(code)
    history.append(result)
    if result.passed:
        return result.output

The magic is not in any single step but in the iteration. A model that gets 60% of programs right on the first try can reach 95% after five corrective rounds, because most failures are mechanical — a missing import, an off-by-one error, a wrong column name — and the error message tells the model exactly what to change.

Why It Matters

Architecture Overview

Our agent has four components:

  1. Prompt builder — assembles the task description, conversation history, and any test expectations into a single prompt.
  2. LLM client — calls the model and extracts a code block from the response.
  3. Sandbox — executes Python code safely and captures stdout, stderr, exit code, and exceptions.
  4. Controller — runs the loop, decides when to stop, and enforces a maximum iteration count.

We will keep dependencies minimal: openai for the model and the standard library for everything else.

Step 1: Building a Safe Sandbox

The sandbox is the most security-sensitive part. You should never run LLM-generated code in your main process or with access to your filesystem. The simplest robust approach is to execute code in a subprocess with a timeout, a clean working directory, and restricted environment variables.

import subprocess
import tempfile
import os
from pathlib import Path
from dataclasses import dataclass

@dataclass
class ExecutionResult:
    success: bool
    stdout: str
    stderr: str
    returncode: int

class PythonSandbox:
    def __init__(self, timeout: int = 10, workdir: str = None):
        self.timeout = timeout
        self.workdir = workdir or tempfile.mkdtemp(prefix="agent_sandbox_")

    def run(self, code: str) -> ExecutionResult:
        script_path = Path(self.workdir) / "solution.py"
        script_path.write_text(code)

        env = {
            "PATH": os.environ.get("PATH", ""),
            "PYTHONPATH": self.workdir,
            "PYTHONDONTWRITEBYTECODE": "1",
        }

        try:
            proc = subprocess.run(
                ["python3", str(script_path)],
                capture_output=True,
                text=True,
                timeout=self.timeout,
                cwd=self.workdir,
                env=env,
            )
            return ExecutionResult(
                success=(proc.returncode == 0),
                stdout=proc.stdout,
                stderr=proc.stderr,
                returncode=proc.returncode,
            )
        except subprocess.TimeoutExpired:
            return ExecutionResult(
                success=False,
                stdout="",
                stderr=f"Execution timed out after {self.timeout}s",
                returncode=-1,
            )

For production use, consider stronger isolation: Docker containers, nsjail, firejail, or a managed sandbox like E2B. The principle is the same — give the code a place to run that cannot reach your real system.

Step 2: Calling the LLM

Next, we need a thin wrapper that asks the model for code and parses the response. We instruct the model to return code inside a fenced block so we can extract it reliably.

import re
import json
from openai import OpenAI

client = OpenAI()

SYSTEM_PROMPT = """You are an expert Python programmer.
You write correct, minimal, well-tested code.

Rules:
- Return your solution inside a single python code block.
- Do not explain before the code block; explanations go after.
- The code must be self-contained and runnable with `python3 solution.py`.
- If you are fixing a previous attempt, read the error carefully and
  change only what is necessary.
"""

CODE_BLOCK_RE = re.compile(r"python\n(.*?)", re.DOTALL)

def extract_code(text: str) -> str:
    match = CODE_BLOCK_RE.search(text)
    if not match:
        # Fallback: treat the whole response as code
        return text.strip()
    return match.group(1).strip()

def generate_code(messages: list, model: str = "gpt-4o-mini") -> str:
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "system", "content": SYSTEM_PROMPT}] + messages,
        temperature=0.2,
    )
    return response.choices[0].message.content

A low temperature keeps the agent focused; we want deterministic fixes, not creative rewrites on every iteration.

Step 3: The Agent Loop

Now we wire the pieces together. The controller maintains a message history, calls the LLM, runs the result, and feeds the outcome back. It stops when execution succeeds or when it hits the iteration cap.

def run_agent(task: str, max_iterations: int = 5, model: str = "gpt-4o-mini") -> dict:
    sandbox = PythonSandbox(timeout=10)
    messages = [{"role": "user", "content": task}]

    for i in range(1, max_iterations + 1):
        print(f"\n=== Iteration {i} ===")
        raw_response = generate_code(messages, model=model)
        code = extract_code(raw_response)
        print(f"Generated code:\n{code[:500]}{'...' if len(code) > 500 else ''}")

        result = sandbox.run(code)
        print(f"Success: {result.success}, returncode: {result.returncode}")

        if result.success:
            return {
                "success": True,
                "iterations": i,
                "code": code,
                "stdout": result.stdout,
            }

        # Feed the failure back to the model
        feedback = (
            f"The previous code failed with return code {result.returncode}.\n"
            f"\nSTDOUT:\n{result.stdout}\n"
            f"\nSTDERR:\n{result.stderr}\n"
            f"\nPlease fix the code and return the complete corrected solution."
        )
        messages.append({"role": "assistant", "content": raw_response})
        messages.append({"role": "user", "content": feedback})

    return {
        "success": False,
        "iterations": max_iterations,
        "code": code,
        "stdout": result.stdout,
        "stderr": result.stderr,
    }

Notice we append the full assistant response (not just the extracted code) to the history. This preserves the model's reasoning and keeps the conversation coherent across turns.

Step 4: Adding Test-Driven Verification

Running without errors is not the same as being correct. To make the agent genuinely self-testing, we should give it an explicit test to pass. The cleanest pattern is to append a hidden test suite to the generated code before execution.

def run_agent_with_tests(task: str, tests: str, max_iterations: int = 5) -> dict:
    sandbox = PythonSandbox(timeout=15)
    messages = [{"role": "user", "content": task}]

    for i in range(1, max_iterations + 1):
        raw_response = generate_code(messages)
        code = extract_code(raw_response)

        # Concatenate solution + tests, then run
        full_code = code + "\n\n# --- hidden tests ---\n" + tests
        result = sandbox.run(full_code)

        if result.success and "ALL TESTS PASSED" in result.stdout:
            return {"success": True, "iterations": i, "code": code}

        feedback = (
            f"Execution result:\n"
            f"returncode: {result.returncode}\n"
            f"STDOUT:\n{result.stdout}\n"
            f"STDERR:\n{result.stderr}\n\n"
            f"Fix the code so all tests pass. Return the full corrected solution."
        )
        messages.append({"role": "assistant", "content": raw_response})
        messages.append({"role": "user", "content": feedback})

    return {"success": False, "iterations": max_iterations, "code": code}

The test string uses assertions and prints a sentinel on success:

TESTS = '''
def test_addition():
    assert add(2, 3) == 5, f"expected 5, got {add(2,3)}"
    assert add(-1, 1) == 0

def test_edge_cases():
    assert add(0, 0) == 0
    assert add(100, -100) == 0

test_addition()
test_edge_cases()
print("ALL TESTS PASSED")
'''

Now the agent has a precise, machine-checkable definition of done.

Step 5: Putting It All Together

if __name__ == "__main__":
    task = "Write a function `add(a, b)` that returns the sum of two integers."
    tests = '''
def test_addition():
    assert add(2, 3) == 5
    assert add(-1, 1) == 0

def test_edge_cases():
    assert add(0, 0) == 0
    assert add(100, -100) == 0

test_addition()
test_edge_cases()
print("ALL TESTS PASSED")
'''

    result = run_agent_with_tests(task, tests, max_iterations=5)

    if result["success"]:
        print(f"\nāœ“ Solved in {result['iterations']} iteration(s)")
        print("Final code:")
        print(result["code"])
    else:
        print(f"\nāœ— Failed after {result['iterations']} iterations")

Run it and you will typically see the agent succeed on the first or second iteration. For harder tasks — parsing a CSV, implementing a sorting algorithm, calling an API — the loop earns its keep, recovering from import errors, type mismatches, and logic bugs that a single-shot call would leave broken.

Best Practices

Constrain the Sandbox Aggressively

LLM-generated code will eventually do something you did not expect: delete files, open network sockets, spawn subprocesses. Run in a container, set a tight timeout, drop network access, and mount a read-only filesystem for anything outside the working directory. Treat every execution as untrusted.

Keep History Bounded

As iterations accumulate, the message history grows. After three or four failed attempts, consider summarizing earlier failures into a compact bullet list rather than carrying full tracebacks. This keeps token costs predictable and prevents the model from fixating on stale errors.

Use Tests as the Stopping Criterion

"It ran without crashing" is a weak signal. Whenever possible, supply executable tests. If you cannot write tests upfront, ask the model to generate its own tests first, review them, then run the agent against those tests. This two-phase approach — generate-tests-then-solve — is the core of techniques like Self-Debugging and Reflexion.

Cap Iterations and Budget

Always set a hard maximum on iterations and total tokens. An agent stuck in a loop will happily burn through your API budget. Log every iteration so you can post-mortem failures and tune the prompt.

Separate Generation from Verification

For harder tasks, use a stronger model to generate code and a cheaper, faster model to critique it. The critic reads the code and the error output and suggests a concrete fix, which the generator applies. This division of labor often outperforms a single model talking to itself.

Log Everything

Capture each iteration's prompt, response, extracted code, stdout, stderr, and timing. These traces are invaluable for debugging the agent itself — when it fails, you want to know whether the prompt was ambiguous, the extraction regex missed the code, or the sandbox swallowed an error.

Conclusion

Building an agent that writes and tests its own code is surprisingly approachable: a sandboxed runner, a thin LLM wrapper, and a feedback loop are all you need. The real engineering work lies in the boundaries — keeping execution safe, defining success with real tests, and managing context so the model improves rather than spirals. Start with the minimal loop above, add hidden tests as your stopping criterion, harden the sandbox before you trust it with anything nontrivial, and iterate on the prompt using traces from real failures. With those foundations in place, the same pattern scales from toy arithmetic functions to genuine development tasks like data cleaning, API integration, and automated refactoring.

— Ad —

Google AdSense will appear here after approval

← Back to all articles