The Death of the Static Token: How AI Coding Agents Learned to Think Before They Type

[User Request] 
      │
      ▼
[Root State: System Prompt & Context]
      ├──> Branch A (Hypothesis 1) ──> [Sandbox Execution] ──> FAIL (Pruned)
      ├──> Branch B (Hypothesis 2) ──> [Sandbox Execution] ──> PASS 
      │         │
      │         ├──> Sub-Branch B1 ──> [Sandbox Execution] ──> PASS (Selected)
      │         └──> Sub-Branch B2 ──> [Sandbox Execution] ──> FAIL (Pruned)
      ▼
[Verified Final Code Generation]

import subprocess
import tempfile
import os
from typing import List, Tuple

class CodeTreeSearchAgent:
    def __init__(self, max_depth: int = 3):
        self.max_depth = max_depth

    def _mock_llm_propose_solutions(self, prompt: str, depth: int) -> List[str]:
        """
        Simulates an LLM proposing multiple code variations (branches) 
        for a given problem state.
        """
        if depth == 1:
            return [
                "def solve(n):\n    return n + ",  # Syntax error branch
                "def solve(n):\n    return sum(range(n + 1))", # Correct branch
                "def solve(n):\n    return n * (n - 1) // 2" # Incorrect logic for some inputs
            ]
        return []

    def _run_in_sandbox(self, code_str: str, test_cases: List[Tuple[int, int]]) -> bool:
        """
        Executes generated code inside a temporary isolated file 
        to verify correctness against test assertions.
        """
        full_script = code_str + "\n\n"
        for i, (inp, expected) in enumerate(test_cases):
            full_script += f"assert solve({inp}) == {expected}, f'Test {i} failed'\n"

        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as tf:
            tf.write(full_script)
            tf_name = tf.name

        try:
            # Execute with a strict 2-second timeout to prevent infinite loops
            result = subprocess.run(
                ['python3', tf_name],
                capture_output=True,
                text=True,
                timeout=2
            )
            return result.returncode == 0
        except (subprocess.TimeoutExpired, Exception):
            return False
        finally:
            if os.path.exists(tf_name):
                os.remove(tf_name)

    def search_and_verify(self, prompt: str, test_cases: List[Tuple[int, int]]) -> str:
        """
        Executes Process-Supervised Tree Search over proposed code branches.
        """
        print(f"[*] Starting Tree Search for prompt: '{prompt}'")
        candidates = self._mock_llm_propose_solutions(prompt, depth=1)

        for idx, candidate in enumerate(candidates):
            print(f"  -> Evaluating Branch {idx+1}...")
            
            is_valid = self._run_in_sandbox(candidate, test_cases)
            
            if is_valid:
                print(f"  [✔] Branch {idx+1} PASSED sandbox verification! Pruning remaining branches.")
                return candidate
            else:
                print(f"  [✘] Branch {idx+1} FAILED execution. Pruning branch.")
                
        raise RuntimeError("All reasoning branches failed execution verification.")

# --- Execution Example for Student Builders ---
if __name__ == "__main__":
    agent = CodeTreeSearchAgent()
    
    # Problem: Calculate sum of numbers from 0 to n
    problem_prompt = "Write a python function solve(n) that returns the sum of numbers from 0 to n."
    tests = [(3, 6), (5, 15), (0, 0)] # (Input, Expected Output)

    try:
        winning_code = agent.search_and_verify(problem_prompt, tests)
        print("\n=== Verified Optimal Output Code ===")
        print(winning_code)
    except Exception as e:
        print(f"\n[!] Search failed: {e}")