CodeRescue: Budget-Calibrated Recovery Routing for Coding Agents

Technical Breakdown§

Problem Formulation§

Coding agents execute code and receive execution feedback. After a failure, the agent can either retry with a cheap model (recovery) or escalate to a stronger, more expensive model. The goal is to minimize cost while maintaining high solve rate. The paper formalizes this as a routing problem: given a failure context $x$ (including code, error, execution trace), choose an action $a \in \mathcal{A}$ (e.g., cheap recovery, escalation) with associated cost $c(a)$ and success probability $p(a|x)$. A router $\pi: \mathcal{X} \to \Delta(\mathcal{A})$ maps contexts to action probabilities. The expected cost for a given solve rate is optimized via a supervised learning approach.

Model Architecture§

The router is a neural network $f_\theta(x)$ that outputs action scores. During training, it learns from execution rollouts: for each failure, multiple actions are tried, yielding success/failure and cost. The training loss is a weighted cross-entropy where positive weight is given to successful actions and negative weight to failures, with a cost penalty term controlled by hyperparameter $\lambda$. The exact loss is:

$$ \mathcal{L}(\theta) = \mathbb{E}_{(x, a, s) \sim \mathcal{D}} \left[ - s \cdot \log \pi_\theta(a|x) + \lambda \cdot c(a) \cdot \mathbb{1}[s=1] \right] $$

where $s \in \{0,1\}$ is success indicator, $c(a)$ is cost, and $\pi_\theta$ is softmax of $f_\theta$.

Conformal Risk Control (CRC) Layer§

To make the router robust to changing budgets, a post-hoc calibration step is added. Given a held-out calibration set, CRC selects a threshold $\tau$ on the cost penalty $\lambda$ such that the expected cost of the router is controlled at level $\alpha$ with high probability. Under exchangeability of calibration data, CRC guarantees:

$$ \mathbb{P}\left( \mathbb{E}[c(\pi_\tau)] \leq \alpha \right) \geq 1 - \delta $$

where $\pi_\tau$ is the router with penalty $\tau$. Practically, they compute empirical cost for a grid of $\tau$ values and pick the smallest $\tau$ such that the empirical cost is below $\alpha$.

Implementation Details§

• Router: 2-layer MLP with hidden size 256, ReLU activations, dropout 0.1.
• Training: Adam optimizer, learning rate 1e-4, batch size 64, early stopping.
• CRC grid: 100 equally spaced $\tau$ values between 0 and 10.
• Calibration set: 200 failure instances held out from training data.

Code Snippet (PyTorch)§

import torch
import torch.nn as nn

class RecoveryRouter(nn.Module):
    def __init__(self, input_dim, num_actions):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, num_actions)
        )
    
    def forward(self, x):
        return torch.softmax(self.net(x), dim=-1)

# Training loop (simplified)
def train(model, dataloader, lambda_penalty):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    for x, action, success, cost in dataloader:
        probs = model(x)
        log_probs = torch.log(probs.gather(1, action.unsqueeze(1)))
        success = success.float()
        loss = - (success * log_probs).mean() + lambda_penalty * (cost * success).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Results§

Across five coding benchmarks (e.g., HumanEval, MBPP), the method achieves a frontier of cost vs. solve rate. One operating point under calibration uses 35% of the cost of always-escalate while achieving comparable solve rate.