Merge branch 'pr2-add-vero' into pr3-harness-bench

Author: varunursekar

vero/src/vero/evaluation/engine.py

@@ -363,14 +363,23 @@ async def _execute_record(
             ).load()
             await asyncio.shield(self._record(cancelled))
             if charged:
-                await asyncio.shield(
-                    self.budget_ledger.refund(
-                        backend_id,
-                        request.evaluation_set,
-                        cost,
-                        principal,
+                try:
+                    await asyncio.shield(
+                        self.budget_ledger.refund(
+                            backend_id,
+                            request.evaluation_set,
+                            cost,
+                            principal,
+                        )
                     )
-                )
+                except Exception as refund_error:
+                    # The shield covers cancellation, not a refund that fails on
+                    # its own: an OSError from the ledger's write propagates here
+                    # in place of error, so the caller never sees the original
+                    # signal while the reservation stays charged -- the leak these
+                    # handlers exist to prevent. Chain it and re-raise the
+                    # original, as BudgetStore.refund already does.
+                    raise error from refund_error
             raise
         except EvaluationExecutionError as error:
             failure = EvaluationStore(
@@ -381,16 +390,25 @@ async def _execute_record(
             # refund and leak the reservation's budget.
             await asyncio.shield(self._record(failure))
             if charged:
-                await asyncio.shield(
-                    self.budget_ledger.refund(
-                        backend_id,
-                        request.evaluation_set,
-                        cost,
-                        principal,
+                try:
+                    await asyncio.shield(
+                        self.budget_ledger.refund(
+                            backend_id,
+                            request.evaluation_set,
+                            cost,
+                            principal,
+                        )
                     )
-                )
+                except Exception as refund_error:
+                    # The shield covers cancellation, not a refund that fails on
+                    # its own: an OSError from the ledger's write propagates here
+                    # in place of error, so the caller never sees the original
+                    # signal while the reservation stays charged -- the leak these
+                    # handlers exist to prevent. Chain it and re-raise the
+                    # original, as BudgetStore.refund already does.
+                    raise error from refund_error
             raise
-        except asyncio.CancelledError:
+        except asyncio.CancelledError as cancellation:
             # Last line of defence for the reservation. The evaluator's own
             # handlers convert cancellation and failure into the two typed errors
             # above, but each of them has to await a persist before it can raise,
@@ -403,14 +421,23 @@ async def _execute_record(
             # evaluation id, so there is no record to load. The evaluator already
             # persisted one on its shielded paths.
             if charged:
-                await asyncio.shield(
-                    self.budget_ledger.refund(
-                        backend_id,
-                        request.evaluation_set,
-                        cost,
-                        principal,
+                try:
+                    await asyncio.shield(
+                        self.budget_ledger.refund(
+                            backend_id,
+                            request.evaluation_set,
+                            cost,
+                            principal,
+                        )
                     )
-                )
+                except Exception as refund_error:
+                    # The shield covers cancellation, not a refund that fails on
+                    # its own: an OSError from the ledger's write propagates here
+                    # in place of cancellation, so the caller never sees the original
+                    # signal while the reservation stays charged -- the leak these
+                    # handlers exist to prevent. Chain it and re-raise the
+                    # original, as BudgetStore.refund already does.
+                    raise cancellation from refund_error
             raise
         # Shielded for the same reason as the two handlers above, and one more:
         # the budget was charged for an evaluation that has now actually run, so

vero/tests/test_v05_evaluation_runtime.py

@@ -789,6 +789,65 @@ async def test_execution_error_refunds_reservation(tmp_path: Path):
     )
 
 
+@pytest.mark.asyncio
+async def test_engine_refund_failure_preserves_the_cancellation(
+    tmp_path: Path, monkeypatch
+):
+    """A refund that fails on its own must not replace the CancelledError.
+
+    The engine's last-line-of-defence handler shields the refund against
+    cancellation, but not against the refund raising by itself. An OSError from
+    the ledger's write then propagated in place of the cancellation, so the
+    caller's structured scope never saw the teardown signal while the
+    reservation stayed charged -- the exact leak this handler exists to prevent.
+    """
+
+    class BlockingBackend(StubBackend):
+        async def evaluate(self, *, context, request):
+            self.evaluate_calls += 1
+            await asyncio.Event().wait()
+
+    workspace = StubWorkspace(tmp_path / "repo")
+    backend = BlockingBackend()
+    evaluation_set = request().evaluation_set
+    ledger = BudgetLedger(
+        [
+            EvaluationBudget(
+                backend_id="default",
+                evaluation_set_key=evaluation_set.budget_key("default"),
+                total_runs=1,
+            )
+        ],
+        path=tmp_path / "budgets.json",
+    )
+    ledger.save()
+
+    async def failing_refund(*args, **kwargs):
+        raise OSError("durable refund write failed")
+
+    monkeypatch.setattr(ledger, "refund", failing_refund)
+
+    engine = EvaluationEngine(
+        evaluator=evaluator(tmp_path, workspace),
+        backends=BackendRegistry({"default": backend}),
+        database=EvaluationDatabase(id="session"),
+        database_path=tmp_path / "database.json",
+        budget_ledger=ledger,
+        authorization_resolver=allow_all_evaluations,
+    )
+    evaluation = asyncio.create_task(
+        engine.evaluate_record(backend_id="default", request=request())
+    )
+    while backend.evaluate_calls == 0:
+        await asyncio.sleep(0)
+    evaluation.cancel()
+
+    # The cancellation wins; the refund failure is chained, not substituted.
+    with pytest.raises(asyncio.CancelledError) as raised:
+        await evaluation
+    assert isinstance(raised.value.__cause__, OSError)
+
+
 @pytest.mark.asyncio
 async def test_cancelled_evaluation_is_terminal_indexed_and_refunded(tmp_path: Path):
     class BlockingBackend(StubBackend):