Building a LangGraph Agent That Writes and Runs Its Own Tests

Why your agent should be writing tests, not just code

Most tutorials on LangGraph stop at "the agent writes code and returns it." That's the easy 80%. The hard 20% — the part that actually matters in production — is knowing whether the code works. If you've shipped an LLM-generated function straight into a pull request without running it, you already know the failure mode: subtly wrong logic, an off-by-one error, an exception on an edge case nobody thought to mention in the prompt.

The fix is not a smarter prompt. It's a feedback loop. A LangGraph agent that writes a function, writes tests for that function, executes those tests in a real Python process, reads the actual failure output, and revises the code — repeating until the tests pass or a retry budget runs out — behaves completely differently from a single-shot code generator. It stops treating "looks plausible" as the bar and starts treating "the assertions passed" as the bar.

This is also one of the clearest illustrations of why LangGraph exists at all. A plain chain can't do this. You need a graph with a real cycle: generate, test, evaluate, and conditionally loop back to generate again with new context. That's a state machine, not a linear pipeline. In this article we'll build exactly that agent — a "codegen-and-test" loop — from the state schema up through the graph wiring, the sandboxed test runner, and the failure-repair loop. Every piece is real code you can drop into a project today.

The core idea: a generate-test-repair cycle

Before touching code, it helps to sketch the shape of the graph, because the shape is the whole point.

The agent needs four kinds of nodes:

• A code generator node that takes a task description (and, on retries, prior failures) and produces a Python function.
• A test generator node that takes the function signature and task description and produces a pytest-style test file.
• A sandboxed executor node that actually runs the tests against the generated code and captures stdout, stderr, and pass/fail status.
• A router that looks at the execution result and decides: done, retry with feedback, or give up after too many attempts.

The cycle looks like this conceptually:

generate_code -> generate_tests -> run_tests -> [route]
                                                   |-- pass --> END
                                                   |-- fail --> generate_code (with feedback)
                                                   |-- max_retries --> END (report failure)

The critical design decision is that the failure feedback — the actual pytest output, not a summary — gets stitched back into the state and passed to the code generator on the next loop. LLMs are much better at fixing code when they see the real traceback than when they're told "the tests failed, try again."

Setting up the project

You'll need LangGraph, LangChain's OpenAI (or Anthropic) integration, and pytest available in whatever environment runs the generated tests.

pip install langgraph langchain-anthropic pytest

We'll use Anthropic's Claude models here, but the graph structure is provider-agnostic — swap the model client and everything else stays the same.

import os

os.environ["ANTHROPIC_API_KEY"] = "your-api-key-here"

Defining the state schema

LangGraph agents are built around a shared state object that every node reads from and writes to. For this agent, the state needs to carry the task description, the current code, the current tests, the last execution result, a retry counter, and a running log of attempts for debugging.

from typing import TypedDict, Optional


class CodegenState(TypedDict):
    task: str                 # natural-language description of the function to build
    function_name: str        # expected name of the function, e.g. "merge_intervals"
    code: str                 # current generated implementation
    test_code: str            # current generated pytest file
    test_output: str          # raw stdout/stderr from the last test run
    passed: bool              # whether the last test run succeeded
    attempt: int               # how many generate-test-run cycles have happened
    max_attempts: int          # retry budget
    history: list[str]         # log of what happened on each attempt, for observability

Keeping test_output as raw text rather than a parsed summary is deliberate. When we feed this back into the code-generation prompt, the model benefits from seeing the exact assertion that failed and the exact values involved — that's the same signal a human developer uses to fix a bug.

Building the code generator node

The code generator node has two modes: first attempt (just the task description) and retry (task description plus the previous code and the failure output). We handle both with one node by checking attempt.

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage, SystemMessage

llm = ChatAnthropic(model="claude-sonnet-4-5", temperature=0)

CODE_SYSTEM_PROMPT = """You are a careful Python engineer. Write a single, complete,
correct Python function that satisfies the task. Output ONLY the code, in a

be self-contained (no undefined names) and must include any necessary imports at the top of the block."""

def generate_code(state: CodegenState) -> CodegenState: if state["attempt"] == 0: user_prompt = f"""Task: {state['task']}

The function must be named {state['function_name']}.""" else: user_prompt = f"""Task: {state['task']}

The function must be named {state['function_name']}.

Your previous attempt failed its tests. Here is the code you wrote:

{state['code']}

Here is the pytest output showing what failed:

{state['test_output']}

Fix the function so all tests pass. Output only the corrected code."""

response = llm.invoke([ SystemMessage(content=CODE_SYSTEM_PROMPT), HumanMessage(content=user_prompt), ])

code = extract_code_block(response.content)

history_entry = f"Attempt {state['attempt'] + 1}: generated code ({len(code)} chars)" return { **state, "code": code, "history": state["history"] + [history_entry], }

def extract_code_block(text: str) -> str: """Pull the first ``python ... ` block out of a model response.""" if "`python" in text: start = text.index("`python") + len("`python") end = text.index("`", start) return text[start:end].strip() if "`" in text: start = text.index("`") + 3 end = text.index("``", start) return text[start:end].strip() return text.strip()

Note the `extract_code_block` helper. Models will occasionally add a stray sentence even when told not to, and a naive `.strip()` on the raw response will break `exec()` later. Always parse fenced code blocks explicitly rather than trusting the model to return raw code with nothing else.

## Building the test generator node

The test generator node only needs to run once per attempt cycle in the common case, but we regenerate tests on every loop too — this catches cases where the first test file itself had a bug, and it's cheap.

TEST_SYSTEM_PROMPT = """You are a meticulous test engineer. Write a pytest test file for the given function. Include at least 4 test cases covering: a typical input, an edge case (empty input, boundary values, etc.), and at least one case likely to break a naive implementation. Assume the function under test is importable via from solution import {function_name}. Output ONLY the test code in a ```python fenced block."""

def generate_tests(state: CodegenState) -> CodegenState: prompt = TEST_SYSTEM_PROMPT.format(function_name=state["function_name"]) user_prompt = f"""Task the function solves: {state['task']}

Here is the function implementation for reference:

{state['code']}

Write the pytest test file now."""

response = llm.invoke([ SystemMessage(content=prompt), HumanMessage(content=user_prompt), ])

test_code = extract_code_block(response.content)

history_entry = f"Attempt {state['attempt'] + 1}: generated {test_code.count('def test_')} test functions" return { **state, "test_code": test_code, "history": state["history"] + [history_entry], }

One subtlety worth calling out: we show the model the current implementation when writing tests. This is intentional, but it does mean the test generator could theoretically write a test that matches a buggy implementation's behavior rather than the task's actual requirements. In practice, giving the test-writer the original task description as the primary source of truth (and the code only as secondary context) keeps this from happening often. If you want stricter separation, generate tests from the task description alone, before the code even exists — a "test-first" variant we'll touch on near the end.

## Running the tests in a sandbox

This is the node that makes the whole system meaningful. We write the generated code to `solution.py` and the generated tests to `test_solution.py` inside an isolated temporary directory, then invoke `pytest` as a subprocess and capture its output.

import subprocess import tempfile import os

def run_tests(state: CodegenState) -> CodegenState: with tempfile.TemporaryDirectory() as tmpdir: solution_path = os.path.join(tmpdir, "solution.py") test_path = os.path.join(tmpdir, "test_solution.py")

with open(solution_path, "w") as f: f.write(state["code"]) with open(test_path, "w") as f: f.write(state["test_code"])

result = subprocess.run( ["python", "-m", "pytest", test_path, "-v", "--tb=short"], cwd=tmpdir, capture_output=True, text=True, timeout=30, )

output = result.stdout + "\n" + result.stderr passed = result.returncode == 0

history_entry = f"Attempt {state['attempt'] + 1}: tests {'PASSED' if passed else 'FAILED'}" return { **state, "test_output": output, "passed": passed, "attempt": state["attempt"] + 1, "history": state["history"] + [history_entry], }

A few details matter here:

- **`tempfile.TemporaryDirectory()`** guarantees each attempt runs in a clean directory with no leftover state from a previous failing attempt.
- **`timeout=30`** is not optional. Generated code can contain infinite loops (a naive recursive Fibonacci with no memoization, an off-by-one that never terminates a `while` loop). Without a timeout, one bad generation hangs your whole pipeline.
- **`capture_output=True`** on `subprocess.run` gives you both streams separately if you want them, but concatenating stdout and stderr into one `output` string is usually what you want to feed back to the LLM — pytest's assertion diffs land in stdout, tracebacks sometimes land in stderr.

If you're deploying this in a multi-tenant or production setting, do not run `subprocess.run` directly against arbitrary model output on a shared host. Use a proper sandbox: a Docker container with no network access and a memory/CPU cap, gVisor, or a managed code-execution API. The subprocess approach above is fine for local development and CI-style pipelines where you control the whole box, but generated code is untrusted input the moment a user's task description can influence it.

## Wiring up the graph

With the three nodes defined, the graph itself is a small amount of code. The interesting part is the conditional edge that implements the retry logic.

from langgraph.graph import StateGraph, END

def route_after_test(state: CodegenState) -> str: if state["passed"]: return "done" if state["attempt"] >= state["max_attempts"]: return "give_up" return "retry"

graph = StateGraph(CodegenState)

graph.add_node("generate_code", generate_code) graph.add_node("generate_tests", generate_tests) graph.add_node("run_tests", run_tests)

graph.set_entry_point("generate_code") graph.add_edge("generate_code", "generate_tests") graph.add_edge("generate_tests", "run_tests")

graph.add_conditional_edges( "run_tests", route_after_test, { "done": END, "give_up": END, "retry": "generate_code", }, )

app = graph.compile()

That's the entire control flow. `run_tests` always routes through `route_after_test`, which reads `passed` and `attempt` off the state and picks one of three destinations. The cycle back to `generate_code` is what makes this a genuine agent loop rather than a pipeline — LangGraph handles the state threading across iterations for you, so `generate_code` on attempt 2 sees the `test_output` that `run_tests` wrote on attempt 1.

## Running the agent end to end

Invoking the compiled graph is a single call. Seed the state with the task and an empty history, and let it run.

initial_state: CodegenState = { "task": "Given a list of intervals as (start, end) tuples, merge all " "overlapping intervals and return the merged list, sorted by start.", "function_name": "merge_intervals", "code": "", "test_code": "", "test_output": "", "passed": False, "attempt": 0, "max_attempts": 4, "history": [], }

final_state = app.invoke(initial_state)

print("Passed:", final_state["passed"]) print("Attempts used:", final_state["attempt"]) print("\n--- Final code ---\n") print(final_state["code"]) print("\n--- History ---") for line in final_state["history"]: print(" -", line)

On a task like interval merging, a first attempt commonly fails on the case where intervals need sorting before merging, or where a boundary-touching interval (`(1, 3)` and `(3, 5)`) should merge but a naive `<` comparison treats it as non-overlapping. The value of this architecture is that the second attempt sees the exact pytest assertion error — something like `AssertionError: assert [(1, 3), (3, 5)] == [(1, 5)]` — and that's usually enough context for the model to fix the comparison operator without any human intervention.

## Streaming intermediate steps for observability

Running `app.invoke()` and waiting for a final answer is fine for a batch job, but if you're wiring this into a CLI tool or a UI, you want to show progress as each node completes. LangGraph's `.stream()` method yields state updates as they happen.

for step in app.stream(initial_state): node_name = list(step.keys())[0] node_output = step[node_name] print(f"[{node_name}] attempt={node_output.get('attempt')} " f"passed={node_output.get('passed')}")

This is also the easiest place to add a circuit breaker for cost control — if you're paying per token and a task is clearly not converging, you can inspect `node_output["history"]` mid-stream and abort early rather than waiting for `max_attempts` to exhaust.

## Handling the failure-repair loop well

The naive version of this agent — just concatenate "it failed, try again" onto the prompt — genuinely does not work well. Two refinements make a measurable difference in how many attempts it takes to converge.

- **Truncate pytest output intelligently.** A full `-v --tb=long` output for four tests can run to several thousand tokens, most of it framework boilerplate. Switch to `--tb=short` (as shown above) or even `--tb=line`, and consider stripping the pytest header/footer banners before feeding the output back into the prompt. You want the assertion diff, not the box-drawing characters around it.
- **Cap the retry budget and report partial progress.** Don't loop forever. A `max_attempts` of 3-5 is usually enough — if the model hasn't converged by then, the task description is probably ambiguous or the function is genuinely hard, and further retries just burn tokens. When `give_up` is reached, surface the last code, the last test output, and the full `history` list so a human can see exactly where it got stuck, rather than just returning "failed."

def summarize_run(state: CodegenState) -> str: if state["passed"]: return f"Converged in {state['attempt']} attempt(s)." return ( f"Did not converge after {state['attempt']} attempt(s).\n" f"Last failure:\n{state['test_output'][-800:]}\n" f"History:\n" + "\n".join(state["history"]) )

## Variations worth trying

Once the base loop works, a few extensions are worth experimenting with:

- **Test-first generation.** Reverse the order: generate tests from the task description alone before any code exists, then generate code to satisfy those tests. This is closer to real TDD and removes the risk of tests being written to match a buggy implementation. It requires the test-writer to reason about expected behavior without an implementation to lean on, which is a harder prompt to get right but produces more trustworthy tests.
- **Static analysis as a cheap first filter.** Before spending a `pytest` subprocess call, run the generated code through `ast.parse()` to catch syntax errors for free, or through `pyflakes` to catch undefined names. This shortens the feedback loop for the cheapest class of failures.
- **Per-attempt model swapping.** Use a fast, cheap model for the first attempt and escalate to a stronger model only if the first attempt fails. Since `attempt` is already in state, this is a one-line change in `generate_code` — pick the LLM client based on `state["attempt"]`.
- **Persisting the loop with a checkpointer.** For long-running or resumable jobs, LangGraph's checkpointing lets you pause after any node and resume later — useful if `run_tests` needs to hand off to an external CI system instead of a local subprocess.

## Wrapping up

The pattern here generalizes well beyond "write a function." Any task where you can define a cheap, automatable correctness check — schema validation for generated JSON, a linter for generated SQL, a compiler for generated Terraform — benefits from the same generate-check-repair cycle. The graph shape barely changes; only the "check" node's implementation does.

What makes this worth building in LangGraph specifically, rather than hand-rolling a while loop around an LLM call, is that the graph makes the control flow explicit and inspectable. You can see the retry edge, trace exactly which node fired on which attempt, stream intermediate state to a UI, and swap in checkpointing for durability — all without restructuring the core logic. That explicitness is what turns "an LLM that sometimes gets code right" into a system you can actually trust to run unattended.

If you want to go deeper on graph design patterns like this — conditional routing, retry loops, multi-agent handoffs, and persistence — that's exactly what we cover hands-on in the **LangGraph Tutorial** course on teachyou.ai, building up from single-node graphs to full production agent architectures.