@goodandready/dsh-image-gen 0.11.30 → 0.11.32

package/lib/tools/generation.js

@@ -2,6 +2,7 @@ import { buildPresentationMeta } from '../attachment-helper.js'
 import { registerGenerationPackTool } from './generation-pack.js'
 import { defineTool } from '@deepseek-ai/dsh-tools'
 import path from 'node:path'
+import { createHash } from 'node:crypto'
 import { mkdir, writeFile } from 'node:fs/promises'
 import {
   computeGenerationHash,
@@ -15,7 +16,7 @@ import {
 } from '../providers.js'
 import { buildDualOutputMarkdown } from '../resolve-image.js'
 import { executeWithQualityGate } from '../quality-gate.js'
-import { calculateGenerationCost, assertBudgetAvailable } from '../cost-meter.js'
+import { calculateGenerationCost, assertBudgetAvailable, reserveSpend } from '../cost-meter.js'
 import { trackAndAssertLoopGuard } from '../loop-guard.js'
 import { getCachedGeneration } from '../generation-cache.js'
 import { resolveFallbackChain, executeWithFallback } from '../fallback-router.js'
@@ -316,8 +317,11 @@ export function registerGenerationTools(ctx, deps, options = {}) {
         let subscriptionImages
         try { subscriptionImages = ctx.get && ctx.get('subscriptionImages') } catch (_err) { subscriptionImages = undefined }
 
-        const source = args.source_image ? await resolveSource(ctx, exec, args.source_image) : undefined
-        const mask = args.mask ? await resolveSource(ctx, exec, args.mask) : undefined
+        const sourceInput = args.source_image || args.image
+        const source = sourceInput ? await resolveSource(ctx, exec, sourceInput) : undefined
+
+        const maskInput = args.mask || args.mask_image
+        const mask = maskInput ? await resolveSource(ctx, exec, maskInput) : undefined
 
         const {
           enhanced,
@@ -378,14 +382,28 @@ export function registerGenerationTools(ctx, deps, options = {}) {
           async (s, p) => {
             let totalAttempts = 0
             let accumulatedCost = 0
-            // #353: Pipeline evaluates raw gen.bytes first; saves artifact ONCE after acceptance
+            // #353, #355: Pipeline reserves budget per attempt, evaluates raw gen.bytes, and commits spend per attempt
+            const perAttemptCost = Number(estimatedCost / totalCount) || 0.02
             const rawGenerate = async (seedVal) => {
+              const reservation = reserveSpend(perAttemptCost, cfg.dailyBudgetUsd)
               totalAttempts++
-              const gen = await providers[k](seedVal, p)
-              guard(gen)
-              const attemptCost = Number(gen.cost) || estimatedCost
-              accumulatedCost += attemptCost
-              return gen
+              let committed = false
+              try {
+                const gen = await providers[k](seedVal, p)
+                guard(gen)
+                const attemptCost = Number(gen.cost) || perAttemptCost
+                accumulatedCost += attemptCost
+                reservation.commit(attemptCost, {
+                  ctx,
+                  meta: { provider: k, model: cfg.model || cfg.customModel, prompt: p, seed: seedVal },
+                })
+                committed = true
+                return gen
+              } finally {
+                if (!committed) {
+                  reservation.release()
+                }
+              }
             }
 
             const { generated: acceptedGen, qualityReport } = await executeWithQualityGate(
@@ -398,7 +416,7 @@ export function registerGenerationTools(ctx, deps, options = {}) {
               }
             )
 
-            // Save the accepted artifact ONCE
+            // Save the accepted artifact ONCE (spend already committed per attempt)
             const saved = await saveGeneratedImageArtifact({
               ctx,
               exec,
@@ -417,6 +435,11 @@ export function registerGenerationTools(ctx, deps, options = {}) {
               slugify,
               estimatedCost: accumulatedCost || estimatedCost,
               totalCount,
+              recordSpend: false,
+              source,
+              mask,
+              resolvedFaceRef,
+              resolvedRefImage,
             })
 
             return {
@@ -437,7 +460,11 @@ export function registerGenerationTools(ctx, deps, options = {}) {
 
           const effectiveModel = provider === 'custom'
             ? (cfg.customModel || cfg.model || 'dall-e-3')
-            : (cfg.model || cfg.customModel)
+            : (cfg.model || cfg.customModel || 'default')
+
+          const sourceBytes = source && source.bytes ? source.bytes : (args.image || args.source_image)
+          const maskBytes = mask && mask.bytes ? mask.bytes : (args.mask_image || args.mask)
+          const refBytes = (resolvedRefImage && resolvedRefImage.bytes) || (resolvedFaceRef && resolvedFaceRef.bytes) || args.reference_images || args.reference_image
 
           const h = computeGenerationHash({
             provider,
@@ -454,9 +481,9 @@ export function registerGenerationTools(ctx, deps, options = {}) {
             negativePrompt: args.negative_prompt || cfg.negativePrompt,
             quality: args.quality,
             strength: args.strength,
-            sourceImage: args.image || args.source_image,
-            maskImage: args.mask_image || args.mask,
-            reference: args.reference_images || args.reference_image,
+            sourceImage: sourceBytes,
+            maskImage: maskBytes,
+            reference: refBytes,
           })
 
           // #170: Content-addressed disk cache check (<50ms return)
@@ -500,10 +527,25 @@ export function registerGenerationTools(ctx, deps, options = {}) {
           const hPrompt = h
           const found = findCachedGeneration(historyEntries, hPrompt)
           if (found) return await cachedResult(found)
+          const srcHash = (source && source.bytes) ? createHash('sha256').update(source.bytes).digest('hex') : undefined
+          const mskHash = (mask && mask.bytes) ? createHash('sha256').update(mask.bytes).digest('hex') : undefined
+          const rfHash = (resolvedRefImage && resolvedRefImage.bytes)
+            ? createHash('sha256').update(resolvedRefImage.bytes).digest('hex')
+            : ((resolvedFaceRef && resolvedFaceRef.bytes) ? createHash('sha256').update(resolvedFaceRef.bytes).digest('hex') : undefined)
+
           const requestedCriteria = {
+            provider,
+            model: effectiveModel,
+            cacheHash: h,
             format: args.output_format || cfg.defaultFormat || 'png',
             aspectRatio: args.aspect_ratio || cfg.defaultAspectRatio,
             size: args.aspect_ratio ? 'custom' : (args.image_size || cfg.defaultSize),
+            quality: args.quality,
+            style: args.style,
+            strength: args.strength,
+            sourceHash: srcHash,
+            maskHash: mskHash,
+            refHash: rfHash,
           }
           if (cfg.cacheBySeed) {
             const bySeed = await findCached(historyEntries, seedVal, promptVal, requestedCriteria)

package/lib/tools/processing-basic.js

@@ -15,7 +15,7 @@ import {
   saveAttachmentSafe,
 } from '../providers.js'
 import { buildDualOutputMarkdown } from '../resolve-image.js'
-import { assertBudgetAvailable, recordSpend } from '../cost-meter.js'
+import { assertBudgetAvailable, reserveSpend, recordSpend } from '../cost-meter.js'
 import { trackAndAssertLoopGuard } from '../loop-guard.js'
 import { sanitizeErrorAndLogs } from '../security.js'
 
@@ -55,39 +55,50 @@ export function registerProcessingTools(ctx, deps) {
             const source = await resolveSource(ctx, exec, args.image)
             if (!source || !source.bytes) throw new Error('Source image not found: ' + args.image)
           const cfg = live()
-          assertBudgetAvailable(0.02, cfg.dailyBudgetUsd)
-          const res = await removeBackgroundFal(
-            { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
-            { imageBytes: source.bytes, mediaType: source.mediaType, model: args.model, signal: exec?.signal }
-          )
-          const stem = `${slugify(args.output_name || 'nobg')}-${Date.now().toString(36)}`
-          const name = `${stem}.png`
-          const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
-          const sessionCwd = exec?.agent?.session?.header?.cwd
-          const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
-          await mkdir(outDir, { recursive: true })
-          const filePath = path.join(outDir, name)
-          await writeFile(filePath, res.bytes)
-          recordSpend(0.02, { ctx, meta: { tool: 'remove_background', model: args.model || 'fal-ai/birefnet' } })
-          const summary = buildDualOutputMarkdown({
-            action: 'background removed',
-            filePath,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            mediaType: 'image/png',
-            provider: 'fal',
-            model: args.model || 'fal-ai/birefnet',
-            attachmentId: attachment.attachmentId,
-          })
-          return toLosslessJson({
-            summary,
-            path: filePath,
-            url: localUrl,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            format: 'png',
-            attachment,
-          })
+          const resvNobg = reserveSpend(0.02, cfg.dailyBudgetUsd)
+          let committed = false
+          try {
+            const res = await removeBackgroundFal(
+              { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
+              { imageBytes: source.bytes, mediaType: source.mediaType, model: args.model, signal: exec?.signal }
+            )
+            if (!res || !res.bytes) {
+              throw new Error('remove_background provider returned empty or invalid response')
+            }
+            const stem = `${slugify(args.output_name || 'nobg')}-${Date.now().toString(36)}`
+            const name = `${stem}.png`
+            const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
+            const sessionCwd = exec?.agent?.session?.header?.cwd
+            const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
+            await mkdir(outDir, { recursive: true })
+            const filePath = path.join(outDir, name)
+            await writeFile(filePath, res.bytes)
+            resvNobg.commit(0.02, { ctx, meta: { tool: 'remove_background', model: args.model || 'fal-ai/birefnet' } })
+            committed = true
+            const summary = buildDualOutputMarkdown({
+              action: 'background removed',
+              filePath,
+              width: res.width || attachment.width,
+              height: res.height || attachment.height,
+              mediaType: 'image/png',
+              provider: 'fal',
+              model: args.model || 'fal-ai/birefnet',
+              attachmentId: attachment.attachmentId,
+            })
+            return toLosslessJson({
+              summary,
+              path: filePath,
+              url: localUrl,
+              width: res.width || attachment.width,
+              height: res.height || attachment.height,
+              format: 'png',
+              attachment,
+            })
+          } finally {
+            if (!committed) {
+              resvNobg.release()
+            }
+          }
           } catch (err) {
             throw new Error(sanitizeErrorAndLogs(err.message || err))
           }
@@ -126,39 +137,50 @@ export function registerProcessingTools(ctx, deps) {
             const source = await resolveSource(ctx, exec, args.image)
             if (!source || !source.bytes) throw new Error('Source image not found: ' + args.image)
           const cfg = live()
-          assertBudgetAvailable(0.025, cfg.dailyBudgetUsd)
-          const res = await upscaleImageFal(
-            { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
-            { imageBytes: source.bytes, mediaType: source.mediaType, scale: args.scale, prompt: args.prompt, creativity: args.creativity, signal: exec?.signal }
-          )
-          const stem = `${slugify(args.output_name || 'upscaled')}-${Date.now().toString(36)}`
-          const name = `${stem}.png`
-          const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
-          const sessionCwd = exec?.agent?.session?.header?.cwd
-          const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
-          await mkdir(outDir, { recursive: true })
-          const filePath = path.join(outDir, name)
-          await writeFile(filePath, res.bytes)
-          recordSpend(0.025, { ctx, meta: { tool: 'upscale_image', scale: args.scale } })
-          const summary = buildDualOutputMarkdown({
-            action: `upscaled (${args.scale || 2}x)`,
-            filePath,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            mediaType: 'image/png',
-            provider: 'fal',
-            model: 'fal-ai/clarity-upscaler',
-            attachmentId: attachment.attachmentId,
-          })
-          return toLosslessJson({
-            summary,
-            path: filePath,
-            url: localUrl,
-            scale: args.scale || 2,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            attachment,
-          })
+          const resvUpscale = reserveSpend(0.025, cfg.dailyBudgetUsd)
+          let committed = false
+          try {
+            const res = await upscaleImageFal(
+              { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
+              { imageBytes: source.bytes, mediaType: source.mediaType, scale: args.scale, prompt: args.prompt, creativity: args.creativity, signal: exec?.signal }
+            )
+            if (!res || !res.bytes) {
+              throw new Error('upscale_image provider returned empty or invalid response')
+            }
+            const stem = `${slugify(args.output_name || 'upscaled')}-${Date.now().toString(36)}`
+            const name = `${stem}.png`
+            const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
+            const sessionCwd = exec?.agent?.session?.header?.cwd
+            const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
+            await mkdir(outDir, { recursive: true })
+            const filePath = path.join(outDir, name)
+            await writeFile(filePath, res.bytes)
+            resvUpscale.commit(0.025, { ctx, meta: { tool: 'upscale_image', scale: args.scale } })
+            committed = true
+            const summary = buildDualOutputMarkdown({
+              action: `upscaled (${args.scale || 2}x)`,
+              filePath,
+              width: res.width || attachment.width,
+              height: res.height || attachment.height,
+              mediaType: 'image/png',
+              provider: 'fal',
+              model: 'fal-ai/clarity-upscaler',
+              attachmentId: attachment.attachmentId,
+            })
+            return toLosslessJson({
+              summary,
+              path: filePath,
+              url: localUrl,
+              scale: args.scale || 2,
+              width: res.width || attachment.width,
+              height: res.height || attachment.height,
+              attachment,
+            })
+          } finally {
+            if (!committed) {
+              resvUpscale.release()
+            }
+          }
           } catch (err) {
             throw new Error(sanitizeErrorAndLogs(err.message || err))
           }
@@ -254,50 +276,61 @@ export function registerProcessingTools(ctx, deps) {
             const cfg = live()
             const sessionId = exec?.agent?.session?.id || exec?.agent?.session?.header?.id || 'session'
             trackAndAssertLoopGuard(sessionId, { limit: cfg.loopGuardLimit, prompt: args.prompt || '', exec })
-            assertBudgetAvailable(0.05, cfg.dailyBudgetUsd)
             const refs = Array.isArray(args.images) ? args.images : [args.images].filter(Boolean)
-          if (refs.length < 2) {
-            throw new Error('At least two images are required for blending')
-          }
-          const sources = []
-          for (const r of refs) {
-            const s = await resolveSource(ctx, exec, r)
-            if (s && s.bytes) sources.push(s)
-          }
-          if (sources.length < 2) {
-            throw new Error('At least two valid images are required to blend')
-          }
-          const res = await blendImagesFal(
-            { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
-            { images: sources, weights: args.weights, prompt: args.prompt, signal: exec?.signal }
-          )
-          const stem = `${slugify(args.output_name || 'blended')}-${Date.now().toString(36)}`
-          const name = `${stem}.png`
-          const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
-          const sessionCwd = exec?.agent?.session?.header?.cwd
-          const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
-          await mkdir(outDir, { recursive: true })
-          const filePath = path.join(outDir, name)
-          await writeFile(filePath, res.bytes)
-          recordSpend(0.05, { ctx, meta: { tool: 'blend_images' } })
-          const summary = buildDualOutputMarkdown({
-            action: 'blended composition',
-            filePath,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            mediaType: 'image/png',
-            provider: 'fal',
-            model: 'fal-ai/image-blend',
-            attachmentId: attachment.attachmentId,
-          })
-          return toLosslessJson({
-            summary,
-            path: filePath,
-            url: localUrl,
-            width: res.width || attachment.width,
-            height: res.height || attachment.height,
-            attachment,
-          })
+            if (refs.length < 2) {
+              throw new Error('At least two images are required for blending')
+            }
+            const sources = []
+            for (const r of refs) {
+              const s = await resolveSource(ctx, exec, r)
+              if (s && s.bytes) sources.push(s)
+            }
+            if (sources.length < 2) {
+              throw new Error('At least two valid images are required to blend')
+            }
+            const resvBlend = reserveSpend(0.05, cfg.dailyBudgetUsd)
+            let committed = false
+            try {
+              const res = await blendImagesFal(
+                { fetchImpl: fetch, resolveKey: (ref) => resolveApiKey(ctx, ref), cfg },
+                { images: sources, weights: args.weights, prompt: args.prompt, signal: exec?.signal }
+              )
+              if (!res || !res.bytes) {
+                throw new Error('blend_images provider returned empty or invalid response')
+              }
+              const stem = `${slugify(args.output_name || 'blended')}-${Date.now().toString(36)}`
+              const name = `${stem}.png`
+              const { attachment, localUrl } = await saveAttachmentSafe(ctx, { bytes: res.bytes, mediaType: 'image/png', name })
+              const sessionCwd = exec?.agent?.session?.header?.cwd
+              const outDir = path.resolve(sessionCwd || process.cwd(), cfg.outputDir || 'generated/images')
+              await mkdir(outDir, { recursive: true })
+              const filePath = path.join(outDir, name)
+              await writeFile(filePath, res.bytes)
+              resvBlend.commit(0.05, { ctx, meta: { tool: 'blend_images' } })
+              committed = true
+              const summary = buildDualOutputMarkdown({
+                action: 'blended composition',
+                filePath,
+                width: res.width || attachment.width,
+                height: res.height || attachment.height,
+                mediaType: 'image/png',
+                provider: 'fal',
+                model: 'fal-ai/image-blend',
+                attachmentId: attachment.attachmentId,
+              })
+              return toLosslessJson({
+                summary,
+                path: filePath,
+                url: localUrl,
+                width: res.width || attachment.width,
+                height: res.height || attachment.height,
+                attachment,
+              })
+            } finally {
+              if (!committed) {
+                resvBlend.release()
+              }
+            }
           } catch (err) {
             throw new Error(sanitizeErrorAndLogs(err.message || err))
           }

package/lib/tools/style-matrix.js

@@ -14,7 +14,7 @@ import {
   saveAttachmentSafe,
 } from '../providers.js'
 import { resolveFallbackChain, executeWithFallback } from '../fallback-router.js'
-import { assertBudgetAvailable, calculateGenerationCost, recordSpend } from '../cost-meter.js'
+import { assertBudgetAvailable, calculateGenerationCost, reserveSpend, recordSpend } from '../cost-meter.js'
 import {
   normalizeMatrixStyles,
   buildMatrixCellPrompt,
@@ -145,8 +145,9 @@ export function registerStyleMatrixTools(ctx, deps) {
             size: cfg.defaultSize || '1024x1024',
             count: cellCount,
           })
-          assertBudgetAvailable(estimatedCost, cfg.dailyBudgetUsd)
-
+          const batchReservation = reserveSpend(estimatedCost, cfg.dailyBudgetUsd)
+          let committed = false
+          try {
           const chain = resolveFallbackChain(provider, cfg.fallbackProviders, PROVIDER_KEYS)
           const quadrantIds = ['A', 'B', 'C', 'D']
           const matrixId = `matrix-${Date.now().toString(36)}`
@@ -187,7 +188,7 @@ export function registerStyleMatrixTools(ctx, deps) {
             })
 
             const cellCost = Number(gen.cost) || +(estimatedCost / cellCount).toFixed(4)
-            recordSpend(cellCost, {
+            batchReservation.commit(cellCost, {
               ctx,
               meta: { provider, model: cfg.model || cfg.customModel, prompt: cellPrompt, seed: cellSeed },
             })
@@ -236,6 +237,7 @@ export function registerStyleMatrixTools(ctx, deps) {
             cells,
           })
 
+          committed = true
           return toLosslessJson({
             matrix_id: matrixId,
             base_prompt: basePrompt,
@@ -244,6 +246,11 @@ export function registerStyleMatrixTools(ctx, deps) {
             cells,
             summary,
           })
+          } finally {
+            if (!committed) {
+              batchReservation.release()
+            }
+          }
         },
       }),
     )

package/lib/tools/theme-pair.js

@@ -13,7 +13,7 @@ import {
   saveAttachmentSafe,
 } from '../providers.js'
 import { resolveFallbackChain, executeWithFallback } from '../fallback-router.js'
-import { assertBudgetAvailable, calculateGenerationCost, recordSpend } from '../cost-meter.js'
+import { assertBudgetAvailable, calculateGenerationCost, reserveSpend, recordSpend } from '../cost-meter.js'
 import { renderToolOutput } from '../attachment-helper.js'
 import {
   THEME_STYLES,
@@ -153,8 +153,9 @@ export function registerThemePairTools(ctx, deps) {
             size: cfg.defaultSize || '1024x1024',
             count: 2,
           })
-          assertBudgetAvailable(estimatedCost, cfg.dailyBudgetUsd)
-
+          const batchReservation = reserveSpend(estimatedCost, cfg.dailyBudgetUsd)
+          let committed = false
+          try {
           const chain = resolveFallbackChain(primaryProvider, cfg.fallbackProviders, PROVIDER_KEYS)
           const baseSeed = typeof args.seed === 'number' ? args.seed : Math.floor(Math.random() * 1000000)
 
@@ -202,7 +203,7 @@ export function registerThemePairTools(ctx, deps) {
             height: lightGen.height || aspectPixels[1],
           })
           const halfCost = +(estimatedCost / 2).toFixed(4)
-          recordSpend(Number(lightGen.cost) || halfCost, {
+          batchReservation.commit(Number(lightGen.cost) || halfCost, {
             ctx,
             meta: { provider: primaryProvider, model: cfg.model || cfg.customModel, prompt: lightPrompt },
           })
@@ -245,7 +246,7 @@ export function registerThemePairTools(ctx, deps) {
             width: darkGen.width || aspectPixels[0],
             height: darkGen.height || aspectPixels[1],
           })
-          recordSpend(Number(darkGen.cost) || halfCost, {
+          batchReservation.commit(Number(darkGen.cost) || halfCost, {
             ctx,
             meta: { provider: primaryProvider, model: cfg.model || cfg.customModel, prompt: darkPrompt },
           })
@@ -277,6 +278,7 @@ export function registerThemePairTools(ctx, deps) {
             + `- Dark: ${darkFilename} (seed: ${baseSeed})\n\n`
             + `\`\`\`html\n${html}\n\`\`\``
 
+          committed = true
           return toLosslessJson({
             summary,
             prompt,
@@ -302,6 +304,11 @@ export function registerThemePairTools(ctx, deps) {
             css_snippet: css,
             attachment: lightAttachment,
           })
+          } finally {
+            if (!committed) {
+              batchReservation.release()
+            }
+          }
         },
       }),
     )

package/lib/tools/ui-asset.js

@@ -15,7 +15,7 @@ import {
   saveAttachmentSafe,
 } from '../providers.js'
 import { resolveFallbackChain, executeWithFallback } from '../fallback-router.js'
-import { assertBudgetAvailable, calculateGenerationCost, recordSpend } from '../cost-meter.js'
+import { assertBudgetAvailable, calculateGenerationCost, reserveSpend, recordSpend } from '../cost-meter.js'
 import { buildDualOutputMarkdown } from '../resolve-image.js'
 import {
   UI_ASSET_TYPES,
@@ -168,14 +168,16 @@ export function registerUiAssetTools(ctx, deps) {
             size: cfg.defaultSize || '1024x1024',
             count: 1,
           })
-          assertBudgetAvailable(estimatedCost, cfg.dailyBudgetUsd)
-
+          const reservation = reserveSpend(estimatedCost, cfg.dailyBudgetUsd)
           const chain = resolveFallbackChain(provider, cfg.fallbackProviders, PROVIDER_KEYS)
           const seedVal = args.seed ?? Math.floor(Math.random() * 100000)
 
-          const gen = await executeWithFallback(providers, chain, seedVal, synthesizedPrompt, {
-            logger: ctx.logger,
-          })
+          let gen
+          let committed = false
+          try {
+            gen = await executeWithFallback(providers, chain, seedVal, synthesizedPrompt, {
+              logger: ctx.logger,
+            })
 
           let finalBytes = gen.bytes
           let mediaType = gen.mediaType || 'image/png'
@@ -211,10 +213,11 @@ export function registerUiAssetTools(ctx, deps) {
           })
 
           const effectiveCost = Number(gen.cost) || estimatedCost
-          recordSpend(effectiveCost, {
+          reservation.commit(effectiveCost, {
             ctx,
             meta: { provider, model: cfg.model || cfg.customModel, prompt: synthesizedPrompt, seed: seedVal },
           })
+          committed = true
 
           const sessionCwd = exec?.agent?.session?.header?.cwd
           const targetDir = cfg.outputDir || 'generated/images'
@@ -269,6 +272,11 @@ export function registerUiAssetTools(ctx, deps) {
             attachment,
             _fallback: gen._fallback,
           })
+          } finally {
+            if (!committed) {
+              reservation.release()
+            }
+          }
         },
       }),
     )

package/lib/tools/vision-ocr.js

@@ -12,7 +12,7 @@ import {
 } from '../providers.js'
 import { buildDualOutputMarkdown } from '../resolve-image.js'
 import { resolveFallbackChain, executeWithFallback } from '../fallback-router.js'
-import { assertBudgetAvailable, calculateGenerationCost, recordSpend } from '../cost-meter.js'
+import { assertBudgetAvailable, calculateGenerationCost, reserveSpend, recordSpend } from '../cost-meter.js'
 import { renderToolOutput } from '../attachment-helper.js'
 import {
   normalizeBbox,
@@ -252,15 +252,17 @@ export function registerVisionOcrTools(ctx, deps) {
             size: `${imgWidth}x${imgHeight}`,
             count: 1,
           })
-          assertBudgetAvailable(estimatedCost, cfg.dailyBudgetUsd)
-
-          const gen = await executeWithFallback(
-            providers,
-            fallbackChain,
-            seed,
-            inpaintPrompt,
-            { logger: ctx.logger },
-          )
+          const reservation = reserveSpend(estimatedCost, cfg.dailyBudgetUsd)
+          let gen
+          let committed = false
+          try {
+            gen = await executeWithFallback(
+              providers,
+              fallbackChain,
+              seed,
+              inpaintPrompt,
+              { logger: ctx.logger },
+            )
 
           const inpaintedBytes = gen.bytes
           if (!inpaintedBytes) {
@@ -323,10 +325,11 @@ export function registerVisionOcrTools(ctx, deps) {
           })
 
           const effectiveCost = Number(gen.cost) || estimatedCost
-          recordSpend(effectiveCost, {
+          reservation.commit(effectiveCost, {
             ctx,
             meta: { provider: usedProvider, model: cfg.model || cfg.customModel, prompt: inpaintPrompt, seed },
           })
+          committed = true
 
           const dualOutput = buildDualOutputMarkdown({
             action: 'generated',
@@ -359,6 +362,11 @@ export function registerVisionOcrTools(ctx, deps) {
           }
 
           return toLosslessJson(response)
+          } finally {
+            if (!committed) {
+              reservation.release()
+            }
+          }
         },
       })
     )