feat: budget-aware retry / graceful-degradation helper (closes #45)

CHANGELOG.md

@@ -10,6 +10,14 @@ both packages adhere to [Semantic Versioning](https://semver.org/).
 
 ## Unreleased
 
+### Added (py + js)
+
+- **Budget-aware retry helper** (`with_budget_retry` / `withBudgetRetry`,
+  issue #45): retry normally when budget is healthy, ask a caller-supplied
+  degrade callback for a cheaper retry plan when `advisory().near_limit` /
+  `nearLimit` is set, and hard-block an over-budget retry with `check()` before
+  it runs. Ships with a no-network Python example.
+
 ## py 0.8.0 — 2026-07-23
 
 ### Added (py)

README.md

@@ -149,6 +149,39 @@ ignore it; `check()` is what enforces the ceiling. See
 [`examples/budget_aware.py`](examples/budget_aware.py) for a runnable taper demo
 (no API key).
 
+### Budget-aware retry
+
+Blind retries can spend the same expensive path again right when the agent is
+running out of headroom. `with_budget_retry()` composes over the existing guard:
+retry normally while budget is healthy, ask your code for a cheaper retry plan
+when `advisory().near_limit` is true, and call `check(estimated_cost)` before
+each retry so an over-budget retry never runs.
+
+```python
+from floe_guard import BudgetGuard, RetryPlan, with_budget_retry
+
+guard = BudgetGuard(limit_usd=1.00)
+
+def premium_model():
+    return call_model("gpt-4o")
+
+def mini_model():
+    return call_model("gpt-4o-mini")
+
+result = with_budget_retry(
+    guard,
+    premium_model,
+    estimated_cost=0.20,
+    max_attempts=2,
+    on_degrade=lambda exc, adv: RetryPlan(call=mini_model, estimated_cost=0.01),
+)
+```
+
+The helper does not rank models or know provider pricing; the caller defines
+what "cheaper" means in `on_degrade`. TypeScript exposes the same pattern as
+`withBudgetRetry()`. See [`examples/budget_retry.py`](examples/budget_retry.py)
+for a no-network demo.
+
 This is the **same advisory shape** hosted Floe returns on every proxied call
 (the `X-Floe-Budget-Advisory` header), so the logic you write here ports
 unchanged — hosted just answers across *every* vendor and cap with server-truth

examples/budget_retry.py

@@ -0,0 +1,50 @@
+"""Budget-aware retry / graceful degradation demo.
+
+Run with:
+
+    python examples/budget_retry.py
+
+No API key, no network. The first call fails near the cap, so the retry helper
+asks the caller for a cheaper path and checks that the cheaper retry fits before
+running it.
+"""
+
+from __future__ import annotations
+
+from floe_guard import BudgetGuard, RetryPlan, with_budget_retry
+
+
+class TransientProviderError(RuntimeError):
+    pass
+
+
+def main() -> None:
+    guard = BudgetGuard(limit_usd=1.00, near_limit_bps=8000)
+    guard.record_tool("previous-work", 0.85)
+    calls = {"premium": 0, "mini": 0}
+
+    def premium_model() -> str:
+        calls["premium"] += 1
+        raise TransientProviderError("temporary upstream timeout")
+
+    def mini_model() -> str:
+        calls["mini"] += 1
+        return "retried on cheaper model"
+
+    def degrade(_exc: BaseException, _advisory) -> RetryPlan[str]:
+        return RetryPlan(call=mini_model, estimated_cost=0.01)
+
+    result = with_budget_retry(
+        guard,
+        premium_model,
+        estimated_cost=0.20,
+        max_attempts=2,
+        on_degrade=degrade,
+    )
+
+    print(result)
+    print(calls)
+
+
+if __name__ == "__main__":
+    main()

js/src/index.ts

@@ -33,5 +33,10 @@ export {
   resolvePrice,
   priceTokens,
 } from "./pricing.js";
+export {
+  withBudgetRetry,
+  type BudgetRetryOptions,
+  type RetryPlan,
+} from "./retry.js";
 
 export * as pricing from "./pricing.js";

js/src/retry.ts

@@ -0,0 +1,82 @@
+import { BudgetExceeded } from "./errors.js";
+import { type BudgetAdvisory, BudgetGuard } from "./guard.js";
+
+export interface RetryPlan<T> {
+  /** Operation to run for this retry attempt. */
+  call: () => T | Promise<T>;
+  /**
+   * Estimated retry cost passed to `guard.check()` before the retry runs.
+   * Leave undefined to use the guard's default last-call estimate.
+   */
+  estimatedCost?: number;
+}
+
+export interface BudgetRetryOptions<T> {
+  /** Estimated cost for retrying the original call. */
+  estimatedCost?: number;
+  /** Total attempts, including the first call. Default: 2. */
+  maxAttempts?: number;
+  /** Choose a cheaper retry plan when `guard.advisory().nearLimit` is true. */
+  onDegrade?: (
+    error: unknown,
+    advisory: BudgetAdvisory,
+  ) => RetryPlan<T> | Promise<RetryPlan<T> | undefined> | undefined;
+  /** Decide whether an error is retryable. Defaults to all non-budget errors. */
+  retryIf?: (error: unknown) => boolean;
+}
+
+function defaultRetryIf(error: unknown): boolean {
+  return !(error instanceof BudgetExceeded);
+}
+
+/**
+ * Run `call` with budget-aware retries.
+ *
+ * The first attempt runs unchanged. If it fails with a retryable error, the
+ * helper retries as-is with ample budget, asks `onDegrade` for a cheaper plan
+ * when near the limit, and always calls `guard.check(estimatedCost)` before a
+ * retry so an over-budget retry is blocked before it runs.
+ */
+export async function withBudgetRetry<T>(
+  guard: BudgetGuard,
+  call: () => T | Promise<T>,
+  options: BudgetRetryOptions<T> = {},
+): Promise<T> {
+  const maxAttempts = options.maxAttempts ?? 2;
+  if (!Number.isInteger(maxAttempts) || maxAttempts < 1) {
+    throw new RangeError(`maxAttempts must be an integer >= 1, got ${maxAttempts}`);
+  }
+  const retryIf = options.retryIf ?? defaultRetryIf;
+  let plan: RetryPlan<T> = { call, estimatedCost: options.estimatedCost };
+
+  for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+    try {
+      return await plan.call();
+    } catch (error) {
+      if (attempt >= maxAttempts || !retryIf(error)) {
+        throw error;
+      }
+      plan = await nextPlan(guard, error, plan, options.onDegrade);
+    }
+  }
+
+  throw new Error("unreachable");
+}
+
+async function nextPlan<T>(
+  guard: BudgetGuard,
+  error: unknown,
+  current: RetryPlan<T>,
+  onDegrade: BudgetRetryOptions<T>["onDegrade"],
+): Promise<RetryPlan<T>> {
+  const advisory = guard.advisory();
+  if (advisory.nearLimit && onDegrade !== undefined) {
+    const degraded = await onDegrade(error, advisory);
+    if (degraded !== undefined) {
+      guard.check(degraded.estimatedCost);
+      return degraded;
+    }
+  }
+  guard.check(current.estimatedCost);
+  return current;
+}

js/test/retry.test.ts

@@ -0,0 +1,102 @@
+import { describe, expect, it } from "vitest";
+
+import { BudgetExceeded, BudgetGuard, type RetryPlan, withBudgetRetry } from "../src/index.js";
+
+class RetryableError extends Error {}
+
+describe("withBudgetRetry", () => {
+  it("retries the same call when budget is ample", async () => {
+    const guard = new BudgetGuard(1.0);
+    let primaryCalls = 0;
+
+    const result = await withBudgetRetry(
+      guard,
+      () => {
+        primaryCalls += 1;
+        if (primaryCalls === 1) throw new RetryableError("temporary failure");
+        return "primary-ok";
+      },
+      { estimatedCost: 0.05, maxAttempts: 2 },
+    );
+
+    expect(result).toBe("primary-ok");
+    expect(primaryCalls).toBe(2);
+  });
+
+  it("uses a degraded retry plan when near the limit", async () => {
+    const guard = new BudgetGuard(1.0, { nearLimitBps: 8000 });
+    guard.recordTool("seed", 0.85);
+    let primaryCalls = 0;
+    let cheapCalls = 0;
+
+    const result = await withBudgetRetry(
+      guard,
+      () => {
+        primaryCalls += 1;
+        throw new RetryableError("temporary failure");
+      },
+      {
+        estimatedCost: 0.2,
+        maxAttempts: 2,
+        onDegrade: (error): RetryPlan<string> => {
+          expect(error).toBeInstanceOf(RetryableError);
+          return {
+            estimatedCost: 0.01,
+            call: () => {
+              cheapCalls += 1;
+              return "cheap-ok";
+            },
+          };
+        },
+      },
+    );
+
+    expect(result).toBe("cheap-ok");
+    expect(primaryCalls).toBe(1);
+    expect(cheapCalls).toBe(1);
+  });
+
+  it("aborts before a retry whose estimate would cross the budget", async () => {
+    const guard = new BudgetGuard(1.0, { onBlock: () => undefined });
+    guard.recordTool("seed", 0.95);
+    let primaryCalls = 0;
+
+    await expect(
+      withBudgetRetry(
+        guard,
+        () => {
+          primaryCalls += 1;
+          throw new RetryableError("temporary failure");
+        },
+        { estimatedCost: 0.1, maxAttempts: 2 },
+      ),
+    ).rejects.toBeInstanceOf(BudgetExceeded);
+    expect(primaryCalls).toBe(1);
+  });
+
+  it("does not retry non-retryable failures", async () => {
+    const guard = new BudgetGuard(1.0);
+    let primaryCalls = 0;
+
+    await expect(
+      withBudgetRetry(
+        guard,
+        () => {
+          primaryCalls += 1;
+          throw new TypeError("bad request");
+        },
+        {
+          estimatedCost: 0.01,
+          retryIf: (error) => !(error instanceof TypeError),
+        },
+      ),
+    ).rejects.toThrow("bad request");
+    expect(primaryCalls).toBe(1);
+  });
+
+  it("rejects invalid maxAttempts", async () => {
+    await expect(
+      withBudgetRetry(new BudgetGuard(1.0), () => "ok", { maxAttempts: 0 }),
+    ).rejects.toBeInstanceOf(RangeError);
+  });
+});

src/floe_guard/__init__.py

@@ -26,6 +26,7 @@
 from .hosted import hosted_enforcement_available, hosted_remaining_usd
 from .latency import LatencyAdvisory, LatencyBudget
 from .pricing import ManualPrice, PricedModel, price_tokens, resolve_price
+from .retry import RetryPlan, async_with_budget_retry, with_budget_retry
 from .stream import StreamGuard, guard_stream
 
 __version__ = "0.8.0"  # keep in lockstep with pyproject.toml
@@ -38,6 +39,9 @@
     "LatencyAdvisory",
     "StreamGuard",
     "guard_stream",
+    "RetryPlan",
+    "with_budget_retry",
+    "async_with_budget_retry",
     "BudgetExceeded",
     "DeadlineExceeded",
     "FloeGuardError",