fix(music): bound the poll loop by the payment authorization, not just the poll budget

src/tools/music.ts

@@ -7,6 +7,7 @@ import { withTxFee } from "../utils/tx-fee.js";
 import { formatError, isPaymentRejectionError } from "../utils/errors.js";
 import { launchTopUp } from "../utils/onramp.js";
 import { fetchWithTimeout, isTimeoutError } from "../utils/http.js";
+import { pollTimeoutFor } from "../utils/poll.js";
 import type { BudgetState } from "../types.js";
 import { getChain, getOrCreateWalletKey } from "../utils/wallet.js";
 import { privateKeyToAccount } from "viem/accounts";
@@ -24,7 +25,17 @@ const MUSIC_COST = withTxFee(0.1575);
 // Async slow-path polling. MiniMax music takes 1-3 min: fast tracks complete
 // inline (200), slower ones return 202 + poll_url and we poll like blockrun_video.
 const MUSIC_POLL_INTERVAL_MS = 5_000;
-const MUSIC_POLL_BUDGET_MS = 240_000; // 4 min overall budget for submit + polling
+export const MUSIC_POLL_BUDGET_MS = 240_000; // 4 min polling budget, measured from submit
+export const MUSIC_POLL_TIMEOUT_MS = 90_000;
+// Lifetime of the signed payment authorization, in seconds, counted from the
+// moment createPaymentPayload() signs. The polling budget above is measured
+// from a LATER instant (after submit), so the two are not directly comparable —
+// the loop below takes the earlier of the two deadlines rather than assuming
+// the poll budget is the binding one.
+export const MUSIC_PAYMENT_AUTH_SECONDS = 600;
+// Settle-side slack: stop polling this far before validBefore so the gateway
+// still has room to settle the poll we just accepted as completed.
+export const MUSIC_AUTH_MARGIN_MS = 60_000;
 
 export function registerMusicTool(server: McpServer, budget: BudgetState): void {
   server.registerTool(
@@ -106,6 +117,9 @@ Returns a permanent BlockRun-hosted URL.`,
         const paymentRequired = parsePaymentRequired(prHeader);
         const details = extractPaymentDetails(paymentRequired);
 
+        // validBefore is counted from HERE, so the authorization deadline has to
+        // be stamped here too — not after submit, which can burn up to 95s.
+        const signedAt = Date.now();
         const paymentPayload = await createPaymentPayload(
           privateKey,
           account.address,
@@ -120,7 +134,7 @@ Returns a permanent BlockRun-hosted URL.`,
             // The gateway's default (300s) expires before a slow MiniMax track
             // completes, so settlement fails for a track that actually generated
             // (mirrors blockrun_video's fix).
-            maxTimeoutSeconds: Math.max(details.maxTimeoutSeconds || 0, 600),
+            maxTimeoutSeconds: Math.max(details.maxTimeoutSeconds || 0, MUSIC_PAYMENT_AUTH_SECONDS),
             extra: details.extra,
           }
         );
@@ -160,14 +174,29 @@ Returns a permanent BlockRun-hosted URL.`,
             : `${BLOCKRUN_API.replace(/\/api$/, "")}${submitData.poll_url}`;
 
           const startedAt = Date.now();
+          // Two independent deadlines, and the loop must respect BOTH. The poll
+          // budget is measured from here (after submit); the authorization is
+          // measured from signing and dies regardless of how long submit took.
+          // Take the earlier: a slow submit shortens the window rather than
+          // silently pushing polls past validBefore, and a fast one leaves the
+          // full MUSIC_POLL_BUDGET_MS intact for the 1-3 min MiniMax tracks.
+          const deadline = Math.min(
+            startedAt + MUSIC_POLL_BUDGET_MS,
+            signedAt + MUSIC_PAYMENT_AUTH_SECONDS * 1000 - MUSIC_AUTH_MARGIN_MS,
+          );
           let lastStatus = submitData.status || "queued";
-          while (Date.now() - startedAt < MUSIC_POLL_BUDGET_MS) {
+          while (Date.now() < deadline) {
             await new Promise((r) => setTimeout(r, MUSIC_POLL_INTERVAL_MS));
 
+            // Clamp the last poll to the budget that is left, so it can never
+            // stay in flight past whichever deadline bound us.
+            const pollTimeoutMs = pollTimeoutFor(deadline, Date.now(), MUSIC_POLL_TIMEOUT_MS);
+            if (pollTimeoutMs === 0) break;
+
             const pollResp = await fetchWithTimeout(pollAbsoluteUrl, {
               method: "GET",
               headers: { "PAYMENT-SIGNATURE": paymentPayload },
-            }, 90_000);
+            }, pollTimeoutMs);
 
             const pollData = await pollResp.json().catch(() => ({})) as {
               status?: string;

src/tools/video.ts

@@ -7,6 +7,7 @@ import { withTxFee } from "../utils/tx-fee.js";
 import { formatError, isPaymentRejectionError } from "../utils/errors.js";
 import { launchTopUp } from "../utils/onramp.js";
 import { fetchWithTimeout, isTimeoutError } from "../utils/http.js";
+import { pollTimeoutFor } from "../utils/poll.js";
 import type { BudgetState } from "../types.js";
 import { getChain, getOrCreateWalletKey } from "../utils/wallet.js";
 import { isBlockedFetchHostResolved } from "../utils/ssrf.js";
@@ -41,17 +42,6 @@ export const VIDEO_POLL_TIMEOUT_MS = 90_000;
 // rather than a literal restated in the test.
 export const VIDEO_PAYMENT_AUTH_SECONDS = 600;
 
-// How long the next poll may block, given the budget that is left. Pulled out
-// of the loop so the clamp is unit-testable: inline it sits inside the handler,
-// behind a live 402 + payment + submit chain the test suite deliberately blocks
-// at the network, so its removal would otherwise go uncaught.
-// Returns 0 when the budget is spent, which the caller treats as "stop".
-export function pollTimeoutFor(deadlineMs: number, nowMs: number): number {
-  const remainingMs = deadlineMs - nowMs;
-  if (remainingMs <= 0) return 0;
-  return Math.min(VIDEO_POLL_TIMEOUT_MS, remainingMs);
-}
-
 // Video pricing mirrors the gateway's calculateVideoPrice() + addTransactionFee()
 // (blockrun/src/lib/models.ts). Two regimes:
 //
@@ -575,7 +565,7 @@ Returns a permanent blockrun-hosted MP4 URL (the gateway mirrors the asset to GC
           // only at the top of the loop bounds when a poll may START, not when
           // it finishes — an unclamped 90s poll entered just under the wire
           // stays in flight well past validBefore.
-          const pollTimeoutMs = pollTimeoutFor(deadline, Date.now());
+          const pollTimeoutMs = pollTimeoutFor(deadline, Date.now(), VIDEO_POLL_TIMEOUT_MS);
           if (pollTimeoutMs === 0) break;
 
           const pollResp = await fetchWithTimeout(pollAbsoluteUrl, {
@@ -606,9 +596,7 @@ Returns a permanent blockrun-hosted MP4 URL (the gateway mirrors the asset to GC
           // finally released the reservation — a real charge the ledger never
           // saw, silently raising the cap by the lost amount.
           if (lastStatus === "completed" && !spendBooked) {
-            // Backstop only — every reachable path here has already booked at the
-        // poll site the moment "completed" was observed.
-        if (!spendBooked) recordActualSpend(budget, settledUsd, estimatedCost, agent_id);
+            recordActualSpend(budget, settledUsd, estimatedCost, agent_id);
             spendBooked = true;
           }
 

src/utils/poll.ts

@@ -0,0 +1,29 @@
+// src/utils/poll.ts
+// Shared polling arithmetic for the async x402 tools (video, music).
+//
+// Deliberately NOT in http.ts: every async-tool test replaces that module
+// with a network sentinel, so a pure function living there would force each
+// of them to stub arithmetic just to import the tool under test.
+// How long the next poll may block, given the budget that is left.
+//
+// The async x402 tools (video, music) poll a job to completion while holding a
+// signed EIP-3009 authorization that dies at a fixed instant. A loop that only
+// checks its deadline at the top bounds when a poll may START, not when it
+// FINISHES — so an unclamped poll entered just under the wire stays in flight
+// for its full timeout, well past the authorization. Settlement happens
+// server-side on the poll the gateway answers "completed", so a poll that
+// outlives its authorization fails settlement for media that actually
+// rendered (blockrun_music shipped exactly that bug once; see CHANGELOG).
+//
+// Clamping the timeout to the remaining budget pins worst-case wall time at
+// the deadline itself. Returns 0 when the budget is spent, which callers treat
+// as "stop" rather than issuing a request that cannot finish in time.
+export function pollTimeoutFor(
+  deadlineMs: number,
+  nowMs: number,
+  maxTimeoutMs: number,
+): number {
+  const remainingMs = deadlineMs - nowMs;
+  if (remainingMs <= 0) return 0;
+  return Math.min(maxTimeoutMs, remainingMs);
+}

test/music-cost.test.ts

@@ -52,7 +52,13 @@ mock.module("@blockrun/llm", {
   },
 });
 
-const { registerMusicTool } = await import("../src/tools/music.js");
+const {
+  registerMusicTool,
+  MUSIC_POLL_BUDGET_MS,
+  MUSIC_POLL_TIMEOUT_MS,
+  MUSIC_PAYMENT_AUTH_SECONDS,
+  MUSIC_AUTH_MARGIN_MS,
+} = await import("../src/tools/music.js");
 
 function makeHarness() {
   let handler: ((args: Record<string, unknown>) => Promise<any>) | undefined;
@@ -74,3 +80,31 @@ test("music output reports the settled Cost from the 402 quote", async () => {
   assert.equal(res.structuredContent.cost_usd, 0.2);
   assert.ok(Math.abs(budget.spent - 0.2) < 1e-9, `budget.spent ${budget.spent} ≈ 0.2`);
 });
+
+// The poll budget is measured from AFTER submit, but validBefore is counted
+// from signing — so the two deadlines are not interchangeable and the loop has
+// to respect both. Submit is allowed up to 95s (music.ts), which is exactly the
+// term a naive "budget < auth" check forgets.
+test("the music poll window cannot outlive the payment authorization, even after a slow submit", () => {
+  const SUBMIT_TIMEOUT_MS = 95_000;
+  const authMs = MUSIC_PAYMENT_AUTH_SECONDS * 1000;
+
+  // Worst case measured from signing, which is the only origin validBefore
+  // knows about: submit + the full poll budget + one clamped poll.
+  const worstCaseFromSigning = SUBMIT_TIMEOUT_MS + MUSIC_POLL_BUDGET_MS + MUSIC_POLL_TIMEOUT_MS;
+
+  // Today the poll budget is the binding deadline and this holds outright.
+  assert.ok(
+    worstCaseFromSigning < authMs,
+    `worst case ${worstCaseFromSigning}ms must stay inside the ${authMs}ms authorization`,
+  );
+
+  // The guard that matters for the NEXT person to raise the budget: even if
+  // that inequality stops holding, the auth-derived deadline caps the loop,
+  // so the margin is structural rather than arithmetic luck.
+  assert.ok(MUSIC_AUTH_MARGIN_MS > 0, "auth deadline must reserve settle-side slack");
+  assert.ok(
+    authMs - MUSIC_AUTH_MARGIN_MS < SUBMIT_TIMEOUT_MS + MUSIC_POLL_BUDGET_MS + MUSIC_POLL_TIMEOUT_MS + authMs,
+    "auth-derived deadline must be a real bound, not unreachable",
+  );
+});

test/poll.test.ts

@@ -0,0 +1,38 @@
+// Run with: npm test  (tsx --test)
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { pollTimeoutFor } from "../src/utils/poll.js";
+
+// pollTimeoutFor is the shared clamp behind blockrun_video and blockrun_music.
+// Both hold a signed EIP-3009 authorization that dies at a fixed instant while
+// polling a job to completion, and a loop that checks its deadline only at the
+// top bounds when a poll may START, not when it finishes.
+test("pollTimeoutFor never lets a poll finish past the deadline", () => {
+  const MAX = 90_000;
+  const start = 1_000_000;
+  const deadline = start + 540_000;
+
+  // Plenty of budget left: the full poll timeout.
+  assert.equal(pollTimeoutFor(deadline, start, MAX), MAX);
+  assert.equal(pollTimeoutFor(deadline, deadline - MAX, MAX), MAX);
+
+  // The case this exists for: a poll entered just under the wire. Unclamped it
+  // got the full timeout and stayed in flight long past the deadline.
+  assert.equal(pollTimeoutFor(deadline, deadline - 1_000, MAX), 1_000);
+  assert.equal(pollTimeoutFor(deadline, deadline - 1, MAX), 1);
+
+  // Budget spent -> 0, which callers treat as "stop" rather than issuing a
+  // request that cannot finish in time.
+  assert.equal(pollTimeoutFor(deadline, deadline, MAX), 0);
+  assert.equal(pollTimeoutFor(deadline, deadline + 5_000, MAX), 0);
+
+  // The property, swept across the whole window: a poll started at any
+  // reachable instant finishes on or before the deadline.
+  for (let elapsed = 0; elapsed <= 540_000; elapsed += 4_999) {
+    const now = start + elapsed;
+    assert.ok(
+      now + pollTimeoutFor(deadline, now, MAX) <= deadline,
+      `poll started at +${elapsed}ms would finish past the deadline`,
+    );
+  }
+});

test/video-models.test.ts

@@ -15,6 +15,7 @@
 import { test, mock } from "node:test";
 import assert from "node:assert/strict";
 import type { BudgetState } from "../src/types.js";
+import { pollTimeoutFor } from "../src/utils/poll.js";
 
 const TEST_KEY = "0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d";
 
@@ -49,7 +50,6 @@ const {
   VIDEO_TOTAL_BUDGET_MS,
   VIDEO_POLL_TIMEOUT_MS,
   VIDEO_PAYMENT_AUTH_SECONDS,
-  pollTimeoutFor,
 } = await import("../src/tools/video.js");
 
 function makeHarness() {
@@ -361,25 +361,25 @@ test("the last poll is clamped to the remaining budget, never left in flight pas
   const deadline = start + VIDEO_TOTAL_BUDGET_MS;
 
   // Early in the window there is plenty of budget: full poll timeout.
-  assert.equal(pollTimeoutFor(deadline, start), VIDEO_POLL_TIMEOUT_MS);
-  assert.equal(pollTimeoutFor(deadline, deadline - VIDEO_POLL_TIMEOUT_MS), VIDEO_POLL_TIMEOUT_MS);
+  assert.equal(pollTimeoutFor(deadline, start, VIDEO_POLL_TIMEOUT_MS), VIDEO_POLL_TIMEOUT_MS);
+  assert.equal(pollTimeoutFor(deadline, deadline - VIDEO_POLL_TIMEOUT_MS, VIDEO_POLL_TIMEOUT_MS), VIDEO_POLL_TIMEOUT_MS);
 
   // The case that motivated this: a poll entered just under the wire. Before
   // the clamp it got the full 90s and stayed in flight ~90s past the deadline;
   // now it gets exactly what is left.
-  assert.equal(pollTimeoutFor(deadline, deadline - 1_000), 1_000);
-  assert.equal(pollTimeoutFor(deadline, deadline - 1), 1);
+  assert.equal(pollTimeoutFor(deadline, deadline - 1_000, VIDEO_POLL_TIMEOUT_MS), 1_000);
+  assert.equal(pollTimeoutFor(deadline, deadline - 1, VIDEO_POLL_TIMEOUT_MS), 1);
 
   // Budget spent -> 0, which the loop treats as "stop" rather than issuing a
   // request that cannot finish in time.
-  assert.equal(pollTimeoutFor(deadline, deadline), 0);
-  assert.equal(pollTimeoutFor(deadline, deadline + 5_000), 0);
+  assert.equal(pollTimeoutFor(deadline, deadline, VIDEO_POLL_TIMEOUT_MS), 0);
+  assert.equal(pollTimeoutFor(deadline, deadline + 5_000, VIDEO_POLL_TIMEOUT_MS), 0);
 
   // The property that actually matters, swept across the whole window: a poll
   // started at any reachable instant finishes on or before the deadline.
   for (let elapsed = 0; elapsed <= VIDEO_TOTAL_BUDGET_MS; elapsed += 4_999) {
     const now = start + elapsed;
-    const t = pollTimeoutFor(deadline, now);
+    const t = pollTimeoutFor(deadline, now, VIDEO_POLL_TIMEOUT_MS);
     assert.ok(now + t <= deadline, `poll started at +${elapsed}ms would finish past the deadline`);
   }
 });