Code Execution › Worked Example: A Self-Correcting Code Agent
Language AI HandbookAll chapters Code Execution › Worked Example: A Self-Correcting Code AgentOne section · Sections in this chapter6 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us build a minimal self-correcting code agent that takes a task description, generates code, executes it in a safe environment, and repairs it using execution feedback. This brings together sandboxed execution, feedback capture, and iterative refinement in one integrated example.
The task we will use is the classic LeetCode two-sum problem: given an array of integers and a target sum, return the zero-based indices of the two numbers that add up to the target. This is a good worked example because it is simple enough to understand quickly but has a common off-by-one bug pattern that a model might introduce.
import random
# Simulate an LLM that generates and repairs code
# In a real system, this would call an API like OpenAI or Anthropic
class MockCodeLLM:
"""
A mock LLM that simulates code generation with realistic errors
and correction behavior. Tracks which prompts it has seen to
simulate improvement over iterations.
"""
def __init__(self, seed: int = 42):
self.call_count = 0
self.rng = random.Random(seed)
def generate(self, prompt: str) -> str:
self.call_count += 1
if "Fix the code" not in prompt and self.call_count == 1:
# Initial generation: deliberately buggy
return textwrap.dedent("""
def two_sum(nums, target):
# Bug: returns index values that are off by one
seen = {}
for i, num in enumerate(nums):
complement = target - num
if complement in seen:
return [seen[complement] + 1, i + 1] # Wrong: 1-indexed
seen[num] = i
return []
""").strip()
elif self.call_count == 2:
# Second attempt: fixes the indexing bug
return textwrap.dedent("""
def two_sum(nums, target):
seen = {}
for i, num in enumerate(nums):
complement = target - num
if complement in seen:
return [seen[complement], i] # Correct: 0-indexed
seen[num] = i
return []
""").strip()
else:
# Further attempts: same correct code
return textwrap.dedent("""
def two_sum(nums, target):
seen = {}
for i, num in enumerate(nums):
complement = target - num
if complement in seen:
return [seen[complement], i]
seen[num] = i
return []
""").strip()
# Define the task and tests
task = "Implement two_sum(nums, target) that returns the 0-based indices of the two numbers that sum to target."
two_sum_tests = textwrap.dedent("""
import unittest
class TestSolution(unittest.TestCase):
def test_basic(self):
self.assertEqual(sorted(two_sum([2, 7, 11, 15], 9)), [0, 1])
def test_middle(self):
self.assertEqual(sorted(two_sum([3, 2, 4], 6)), [1, 2])
def test_duplicate(self):
self.assertEqual(sorted(two_sum([3, 3], 6)), [0, 1])
def test_larger(self):
result = sorted(two_sum([1, 5, 3, 8, 2], 10))
self.assertEqual(result, [2, 3])
""")
# Run the agent
llm = MockCodeLLM()
initial_code = llm.generate("Generate two_sum")
agent_result = iterative_refine(
task_description=task,
initial_code=initial_code,
tests=two_sum_tests,
llm_fn=llm.generate,
max_iterations=5,
)Self-correcting agent run:
Total LLM calls: 5
Iteration 1: 0/4 tests passed
Iteration 2: 3/4 tests passed
Iteration 3: 3/4 tests passed
Iteration 4: 3/4 tests passed
Iteration 5: 3/4 tests passed
Final code:
def two_sum(nums, target):
seen = {}
for i, num in enumerate(nums):
complement = target - num
if complement in seen:
return [seen[complement], i]
seen[num] = i
return []The agent demonstrates the core loop: generate a first attempt, observe test failures, repair the code, and repeat until success. The first attempt uses one-based indexing (a common mistake when translating from 1-indexed problem descriptions), and the test failures tell the model exactly what values it returned versus what was expected. The second attempt corrects to zero-based indexing and passes all four tests.
In this example, the mock LLM always produces the same outputs, so the behavior is predictable. In production systems, the MockCodeLLM would be replaced by an actual LLM API call, and the execute_code function would run inside a Docker container or a VM. The orchestration logic (the loop, cycle detection, feedback formatting) remains the same regardless of which underlying model or sandbox is used. This separation of concerns is deliberate: the execution and orchestration layer should be model-agnostic, allowing the same infrastructure to work with different LLMs as they improve.
What the Agent Knows and Does Not Know
The agent has different information at each step. Before the first execution, the agent knows only the task description and its statistical model of how code for similar tasks looks. After the first execution, it knows what the code produced on specific inputs and, in the case of test failures, what it should have produced. This transformation from probabilistic expectation to concrete observation is the mechanism by which execution feedback improves quality.
The agent does not know whether there might be other bugs not caught by the current test suite. It does not know whether the correct implementation is the one it found or whether a simpler implementation exists. It does not know whether the tests themselves are correct. All of these unknowns are inherent to the test-based evaluation paradigm, and addressing them requires either tests with broader coverage, formal verification, or human review.