LangGraph for Code Generation Agents: Plan, Write, Test, Fix

Why Code Generation Agents Need a Graph, Not a Chain

Ask an LLM to "write a function that parses CSV files and handles malformed rows" and you'll usually get code that looks right and compiles wrong. Maybe it imports a module that doesn't exist. Maybe it forgets an edge case you mentioned three sentences ago. A single prompt-and-response call has no way to catch that. It just returns text and hopes for the best.

Real coding assistants don't work that way, and neither should your agent. A working code generation agent has to plan before it writes, run the code it produces, read the actual error message when something breaks, and go back to fix it. That is a loop with conditional branches, not a straight line. LangGraph exists precisely for this shape of problem: a directed graph of nodes where state flows between planning, generation, execution, and repair, with edges that can loop back based on what actually happened at runtime.

This article walks through building a Plan, Write, Test, Fix code generation agent in LangGraph. We will define the state schema, write each node function, wire up the routing logic that decides whether to retry or finish, and talk through the failure modes you will hit in production — infinite retry loops, cascading errors, and tests that pass for the wrong reason. By the end you will have a pattern you can drop into any coding-assistant project, whether it's generating unit tests, refactoring a module, or building a full feature from a spec.

The Core Idea: Treat Code Generation as a State Machine

Most tutorials treat an LLM call as the unit of work. For code generation, the better unit is a state transition. Your agent's state carries the task description, the current code draft, the test results, an error history, and a retry counter. Each node reads that state, does one job, and writes back an updated state. LangGraph routes between nodes based on conditions you define — usually "did the tests pass?" or "have we hit the retry limit?"

This gives you three things a plain chain can't:

• Observability — you can inspect the state after every node and see exactly what the agent knew at each step.
• Resumability — with a checkpointer, you can pause after a failed test run, let a human look at it, and resume.
• Bounded retries — you control exactly how many times the agent is allowed to loop before giving up, instead of an open-ended while True with an LLM inside it.

Here's the mental model before we write any code:

   ┌─────────┐
   │  plan   │
   └────┬────┘
        │
   ┌────▼────┐
   │  write  │◄──────────┐
   └────┬────┘           │
        │                │
   ┌────▼────┐           │
   │  test   │           │
   └────┬────┘           │
        │                │
   ┌────▼────┐   fail    │
   │ router  ├───────────┘
   └────┬────┘
        │ pass or max retries
   ┌────▼────┐
   │  done   │
   └─────────┘

The plan node runs once. write, test, and the router form a loop that can execute multiple times until the code passes or the agent runs out of attempts.

Defining the State Schema

LangGraph state is typically a TypedDict (or a Pydantic model) that every node reads from and writes to. For a code generation agent, keep the schema honest about what actually changes across the loop — don't cram unrelated fields in just because you can.

from typing import TypedDict, List, Optional

class CodeGenState(TypedDict):
    task: str                     # the original user request
    plan: Optional[str]           # step-by-step plan from the planner
    code: Optional[str]           # current code draft
    test_code: Optional[str]      # generated test cases
    test_output: Optional[str]    # stdout/stderr from the last test run
    passed: bool                  # did the last test run succeed
    attempts: int                 # how many write-test cycles we've run
    max_attempts: int             # hard ceiling on retries
    error_history: List[str]      # every failure message we've seen, in order

A few decisions worth calling out. error_history is a list, not a single string — you want the fix node to see the *pattern* of failures, not just the latest one. If the agent keeps hitting the same import error three times in a row, that's a signal the fix strategy itself is wrong, and the fix node should know that. attempts and max_attempts are separate fields so the router logic is a simple comparison, not something buried in a prompt.

Building the Plan Node

The plan node's job is narrow: turn a vague task into a concrete, ordered list of implementation steps. This matters more than it sounds like it should. When you skip planning and go straight to code generation, the model tends to make structural decisions on the fly — what functions to define, what the signature looks like — and those decisions are much harder to unwind later than if you'd nailed them down first.

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

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

def plan_node(state: CodeGenState) -> dict:
    system = SystemMessage(content=(
        "You are a senior engineer breaking a coding task into a short, "
        "concrete implementation plan. List the functions or classes needed, "
        "their signatures, and the edge cases the implementation must handle. "
        "Do not write code yet. Keep it under 10 bullet points."
    ))
    human = HumanMessage(content=f"Task: {state['task']}")

    response = llm.invoke([system, human])

    return {
        "plan": response.content,
        "attempts": 0,
        "error_history": [],
    }
}

Notice the node returns a partial dict, not the full state. LangGraph merges this into the existing state automatically. This is the pattern you'll use for every node — return only the keys you actually changed. It keeps node functions small and makes diffs between state snapshots easy to read when you're debugging a run.

The Write Node: Generating Code From the Plan (and From Failures)

The write node is where code actually gets produced. On the first pass, it works from the plan. On every subsequent pass, it needs the previous code, the test output that failed, and the accumulated error history — otherwise it will happily regenerate the exact same broken code.

def write_node(state: CodeGenState) -> dict:
    if state["attempts"] == 0:
        # first pass: generate from the plan
        prompt = (
            f"Task: {state['task']}\n\n"
            f"Plan:\n{state['plan']}\n\n"
            "Write the complete Python implementation. Return only code, "
            "no explanation, no markdown fences."
        )
    else:
        # repair pass: generate from the failure
        prompt = (
            f"Task: {state['task']}\n\n"
            f"Previous code:\n{state['code']}\n\n"
            f"Test output (failure):\n{state['test_output']}\n\n"
            f"Prior errors seen so far:\n{chr(10).join(state['error_history'])}\n\n"
            "Fix the code so it passes. Return the full corrected file, "
            "only code, no explanation."
        )

    response = llm.invoke([HumanMessage(content=prompt)])
    new_code = strip_code_fences(response.content)

    return {
        "code": new_code,
        "attempts": state["attempts"] + 1,
    }


def strip_code_fences(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        lines = lines[1:] if lines[0].startswith("```") else lines
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        return "\n".join(lines)
    return text

The strip_code_fences helper is unglamorous but necessary. Models wrap code in markdown fences constantly, even when you explicitly ask them not to, and a stray `python line will break every downstream step that tries to exec() or write this to a .py file. Handle it once, here, rather than debugging a mysterious SyntaxError three nodes later.

The Test Node: Actually Running the Code

This is the node that turns your agent from "text generator" into "code generation agent." Without execution, you're just asking the model to grade its own homework. Run it in a subprocess, capture both streams, and never let a timeout hang your graph.

import subprocess
import tempfile
import os

def test_node(state: CodeGenState) -> dict:
    with tempfile.TemporaryDirectory() as tmp_dir:
        code_path = os.path.join(tmp_dir, "solution.py")
        test_path = os.path.join(tmp_dir, "test_solution.py")

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

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

        try:
            result = subprocess.run(
                ["python", "-m", "pytest", test_path, "-v", "--tb=short"],
                cwd=tmp_dir,
                capture_output=True,
                text=True,
                timeout=30,
            )
            output = result.stdout + result.stderr
            passed = result.returncode == 0
        except subprocess.TimeoutExpired:
            output = "Execution timed out after 30 seconds. Possible infinite loop."
            passed = False

    error_history = state["error_history"]
    if not passed:
        error_history = error_history + [output[-500:]]

    return {
        "test_output": output,
        "passed": passed,
        "error_history": error_history,
    }

Three details here matter for reliability. First, timeout=30 — an LLM can absolutely generate an infinite loop, and without a timeout your graph just hangs forever. Second, running in a fresh TemporaryDirectory avoids state leaking between attempts, like a stale .pyc file or a leftover import. Third, we only append the last 500 characters of output to error_history. Stack traces from repeated failures can balloon your context window fast; you want the useful tail, not the full noise.

In production, you would run this subprocess inside a sandboxed container, not directly on the host — never execute LLM-generated code with the same privileges as your application process.

Generating Tests When None Are Provided

If the user's task doesn't come with tests attached, the agent needs to write its own before it can judge success. This is a separate concern from writing the implementation — mixing them tends to produce tests that are suspiciously easy to pass.

def generate_default_tests(state: CodeGenState) -> str:
    prompt = (
        f"Task: {state['task']}\n\n"
        f"Implementation:\n{state['code']}\n\n"
        "Write pytest test cases for this implementation, covering the happy "
        "path, at least two edge cases, and one invalid-input case. "
        "Import from solution.py. Return only code."
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    return strip_code_fences(response.content)

Call this once, cache the result in state["test_code"], and reuse the same tests across every retry attempt. If you regenerate tests on every loop iteration, the agent can "fix" the code by quietly making the tests weaker instead of making the implementation correct — a subtle failure mode that's easy to miss if you're only checking whether passed came back True.

Wiring the Graph: Nodes, Edges, and the Router

With all four node functions defined, assembling the graph is straightforward. The interesting part is the conditional edge that decides whether to loop back to write or exit.

from langgraph.graph import StateGraph, END

def router(state: CodeGenState) -> str:
    if state["passed"]:
        return "done"
    if state["attempts"] >= state["max_attempts"]:
        return "give_up"
    return "retry"

graph = StateGraph(CodeGenState)

graph.add_node("plan", plan_node)
graph.add_node("write", write_node)
graph.add_node("test", test_node)

graph.set_entry_point("plan")
graph.add_edge("plan", "write")
graph.add_edge("write", "test")

graph.add_conditional_edges(
    "test",
    router,
    {
        "done": END,
        "give_up": END,
        "retry": "write",
    },
)

app = graph.compile()

Run it with an initial state that sets max_attempts explicitly — don't let it default silently, because that's the difference between a bounded agent and one that can burn through your API budget on a task it will never solve:

result = app.invoke({
    "task": "Write a function `parse_csv_row(line: str) -> dict` that "
            "splits a CSV line into a dict using the header "
            "['name', 'age', 'email'], raising ValueError on wrong "
            "column count, and stripping whitespace from every field.",
    "plan": None,
    "code": None,
    "test_code": None,
    "test_output": None,
    "passed": False,
    "attempts": 0,
    "max_attempts": 4,
    "error_history": [],
})

print("Passed:", result["passed"])
print("Attempts used:", result["attempts"])
print(result["code"])

Notice give_up and done both route to END — the graph terminates either way, but state["passed"] tells you which one happened. Downstream, you branch on that flag: surface the code either way, but flag unresolved failures for a human to look at rather than silently shipping broken output.

Handling the Failure Modes That Actually Show Up

The happy path above works in a demo. In practice, three failure patterns show up constantly, and it's worth designing for them up front rather than patching them in after a bad run in production.

Infinite retry on an unfixable task. Sometimes the task is genuinely ambiguous or contradictory, and the agent will burn every attempt without converging. The max_attempts ceiling handles the cost side, but you should also detect repeated identical failures — if the last two entries in error_history are near-identical, the fix node is stuck in a loop and more retries won't help.

def is_stuck(error_history: List[str]) -> bool:
    if len(error_history) < 2:
        return False
    last, prev = error_history[-1], error_history[-2]
    return last[:200] == prev[:200]

Wire this into the router so a stuck agent gives up early instead of burning its full retry budget on a repeat of the same mistake:

def router(state: CodeGenState) -> str:
    if state["passed"]:
        return "done"
    if is_stuck(state["error_history"]):
        return "give_up"
    if state["attempts"] >= state["max_attempts"]:
        return "give_up"
    return "retry"

Tests that pass for the wrong reason. A generated test suite can accidentally test nothing meaningful — asserting True == True dressed up as a test, or catching an exception and treating that as success. Guard against this by asking the planner for explicit expected behaviors up front, and have the test node check that the test file actually contains assertions tied to those behaviors, not just that pytest returned exit code 0.

Cascading errors from a bad plan. If the plan itself is wrong — say it specifies the wrong function signature — every write attempt will fail against tests generated from that same wrong plan, and the loop will never converge no matter how many retries you allow. This is why error_history should be inspected by a human (or a separate "diagnose" node) when give_up fires: often the fix isn't another code attempt, it's revisiting the plan. A more advanced version of this graph adds an edge from a repeated-failure state back to plan, not just to write, giving the agent a chance to reconsider its whole approach rather than just patching the last attempt.

Adding a Checkpointer for Long-Running Tasks

For anything beyond a toy example — generating a multi-file feature, refactoring across a module — you don't want the whole loop to live in memory for one uninterrupted invoke() call. LangGraph's checkpointing lets you persist state between steps and resume later, which also gives you a natural place to insert a human review before the agent commits to another round of retries.

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()
app = graph.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "task-42"}}

result = app.invoke({
    "task": "...",
    "plan": None, "code": None, "test_code": None,
    "test_output": None, "passed": False,
    "attempts": 0, "max_attempts": 4, "error_history": [],
}, config=config)

# later, inspect or resume the same thread
state = app.get_state(config)
print(state.values["attempts"], state.values["passed"])

Swap MemorySaver for a Postgres or SQLite-backed checkpointer in production so state survives a process restart. This is also where you'd hook in a human-in-the-loop pause: interrupt before the write node on retry attempt 3, show the human the failing test output, and let them either approve another automated attempt or take over manually.

Extending the Pattern: Multi-File Projects and Static Analysis

Once the core loop works for a single function, the same skeleton extends in a few directions worth knowing about before you need them.

• Static analysis before execution — add a lightweight node that runs a linter or type checker (ruff, mypy) between write and test. Catching a NameError via static analysis is faster and cheaper than discovering it through a subprocess run, and it gives the fix node a more precise signal than a full traceback.
• Multi-file state — instead of a single code: str field, use code: Dict[str, str] mapping filenames to contents, and write each file into the temp directory before running tests. The plan node should specify the file layout explicitly so the write node isn't guessing at module boundaries.
• Parallel test execution — if the plan produces multiple independent functions, you can fan out to parallel write/test subgraphs per function using LangGraph's support for parallel branches, then join before a final integration test.
• Cost-aware routing — track token usage per attempt in the state and add a budget ceiling alongside max_attempts, so a task that requires unusually long code doesn't blow through cost limits even if it's still within its retry count.

None of these require rethinking the architecture — they're additional nodes and richer state fields layered onto the same plan-write-test-fix loop.

What This Buys You Over Ad Hoc Retry Loops

It's tempting to just wrap an LLM call in a for loop with a try/except and call it done. The reason LangGraph earns its place here is that the retry logic, the state history, and the routing decisions all become inspectable and testable in their own right. You can unit test the router function without touching an LLM. You can replay a saved state to debug exactly why attempt 3 failed differently from attempt 2. You can swap the write_node prompt without touching how tests execute. That separation is what turns "an agent that sometimes works" into something you can actually maintain.

The pattern in this article — plan once, then loop through write, test, and a conditional router until you pass or hit a ceiling — is deliberately minimal. It's meant to be a foundation you extend with static analysis, multi-file state, checkpointing, and human review as your use case demands, not a finished product.

If you want to go deeper on building graphs like this — state design, checkpointing strategies, parallel branches, and debugging non-deterministic agent runs — our LangGraph Tutorial course on teachyou.ai walks through the full arc from a single-node graph to a production-grade multi-agent system, with this exact code generation pattern as one of the worked projects.