From de27e9683fc25356f32ecfda76936e3407ebaa77 Mon Sep 17 00:00:00 2001 From: Kurt Overmier Date: Sun, 14 Jun 2026 05:26:59 -0500 Subject: [PATCH] fix(image): align cf_model proxy schema and billing --- docs/api-reference.md | 3 ++- src/cost-attribution.ts | 21 +++++++++++++++++++- src/tool-registry.ts | 6 ++++-- test/cost-attribution.test.ts | 21 ++++++++++++++++++++ test/gateway.test.ts | 37 +++++++++++++++++++++++++++++++++++ test/tool-registry.test.ts | 25 +++++++++++++++++++++++ 6 files changed, 109 insertions(+), 4 deletions(-) diff --git a/docs/api-reference.md b/docs/api-reference.md index 171f37b..bd1e53f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -213,8 +213,9 @@ Generate an image from a text prompt. - `quality_tier` (string, optional) — `draft`, `standard` (default), `premium`, `ultra`, `ultra_plus` - `negative_prompt` (string, optional) — things to avoid; effective for `draft` tier only - `aspect_ratio` (string, optional) — `1:1` (default), `3:2`, `2:3`, `3:4`, `4:3`, `4:5`, `5:4`, `9:16`, `16:9`, `21:9` - - `image_size` (string, optional) — `512`, `1K` (default), `2K`, `4K` + - `image_size` (string, optional) — Gemini-only output size: `512`, `1K`, `2K`, `4K`. Use only with `model` or `ultra`/`ultra_plus`; omit when using `cf_model`. - `model` (string, optional) — `gemini-3.1-flash-image-preview` (maps to `ultra`), `gemini-3-pro-image-preview` (maps to `ultra_plus`); when set, takes billing and tier-enforcement precedence over `quality_tier` + - `cf_model` (string, optional) — Workers AI text-to-image model override. Supports FLUX.2, FLUX.1, Leonardo, SDXL, DreamShaper, and Stable Diffusion task variants. Effective billing/access tier is resolved from the selected model; `model` wins over `cf_model`, and `cf_model` wins over `quality_tier`. ### `image_list_models` diff --git a/src/cost-attribution.ts b/src/cost-attribution.ts index 264be2c..5d9c273 100644 --- a/src/cost-attribution.ts +++ b/src/cost-attribution.ts @@ -64,11 +64,30 @@ const MODEL_QUALITY_TIER: Record = { 'gemini-3-pro-image-preview': 'ultra_plus', }; -/** Resolve the effective quality tier from image_generate args. Model wins when set. */ +// Maps direct Workers AI image model overrides to their effective billing/access tier. +// model wins first, then cf_model, then quality_tier. +const CF_MODEL_QUALITY_TIER: Record = { + '@cf/black-forest-labs/flux-2-klein-9b': 'premium', + '@cf/black-forest-labs/flux-2-klein-4b': 'standard', + '@cf/black-forest-labs/flux-2-dev': 'premium', + '@cf/leonardo/lucid-origin': 'premium', + '@cf/leonardo/phoenix-1.0': 'premium', + '@cf/black-forest-labs/flux-1-schnell': 'draft', + '@cf/bytedance/stable-diffusion-xl-lightning': 'draft', + '@cf/lykon/dreamshaper-8-lcm': 'draft', + '@cf/runwayml/stable-diffusion-v1-5-img2img': 'draft', + '@cf/runwayml/stable-diffusion-v1-5-inpainting': 'draft', + '@cf/stabilityai/stable-diffusion-xl-base-1.0': 'draft', +}; + +/** Resolve the effective quality tier from image_generate args. model wins, then cf_model. */ export function resolveImageQualityTier(args?: Record): string { if (args?.model) { return MODEL_QUALITY_TIER[args.model as string] ?? (args.quality_tier as string) ?? 'standard'; } + if (args?.cf_model) { + return CF_MODEL_QUALITY_TIER[args.cf_model as string] ?? (args.quality_tier as string) ?? 'standard'; + } return (args?.quality_tier as string) ?? 'standard'; } diff --git a/src/tool-registry.ts b/src/tool-registry.ts index d2069b2..3c1eefe 100644 --- a/src/tool-registry.ts +++ b/src/tool-registry.ts @@ -129,8 +129,10 @@ const TOOL_SPECS: ToolSpec[] = [ image_size: { type: 'string', enum: ['512', '1K', '2K', '4K'], - default: '1K', - description: 'Output resolution. Defaults to 1K.', + description: + 'Gemini-only output size. Use only with model or ultra/ultra_plus tiers; ' + + 'do not send with cf_model or Cloudflare tiers. Omit for Cloudflare models; ' + + 'their dimensions are controlled by aspect_ratio and model defaults.', }, model: { type: 'string', diff --git a/test/cost-attribution.test.ts b/test/cost-attribution.test.ts index c204b13..c316925 100644 --- a/test/cost-attribution.test.ts +++ b/test/cost-attribution.test.ts @@ -49,6 +49,18 @@ describe('resolveToolCost', () => { expect(pro.baseCost).toBe(5 * 8); // ultra_plus multiplier }); + it('applies cf_model-derived tier for image_generate', () => { + const schnell = resolveToolCost('image_generate', { cf_model: '@cf/black-forest-labs/flux-1-schnell' }); + const klein4b = resolveToolCost('image_generate', { cf_model: '@cf/black-forest-labs/flux-2-klein-4b' }); + const klein9b = resolveToolCost('image_generate', { cf_model: '@cf/black-forest-labs/flux-2-klein-9b' }); + const leonardo = resolveToolCost('image_generate', { cf_model: '@cf/leonardo/lucid-origin' }); + + expect(schnell.baseCost).toBe(5 * 1); + expect(klein4b.baseCost).toBe(5 * 1); + expect(klein9b.baseCost).toBe(5 * 3); + expect(leonardo.baseCost).toBe(5 * 3); + }); + it('model wins over quality_tier when both are set', () => { const cost = resolveToolCost('image_generate', { model: 'gemini-3-pro-image-preview', @@ -57,6 +69,15 @@ describe('resolveToolCost', () => { expect(cost.baseCost).toBe(5 * 8); // ultra_plus, not standard }); + it('model wins over cf_model when both are set', () => { + const cost = resolveToolCost('image_generate', { + model: 'gemini-3-pro-image-preview', + cf_model: '@cf/black-forest-labs/flux-1-schnell', + quality_tier: 'draft', + }); + expect(cost.baseCost).toBe(5 * 8); + }); + it('falls back to quality_tier for unknown model', () => { const cost = resolveToolCost('image_generate', { model: 'unknown-model', diff --git a/test/gateway.test.ts b/test/gateway.test.ts index 5cf170b..724e268 100644 --- a/test/gateway.test.ts +++ b/test/gateway.test.ts @@ -377,6 +377,43 @@ describe('handleMcpRequest', () => { expect(text).toMatch(/pro plan|ultra_plus/i); }); + it('denies free-tier user who passes premium cf_model with standard quality_tier', async () => { + const env = makeEnv({ + AUTH_SERVICE: { + ...mockAuthService(), + validateApiKey: async () => ({ + valid: true, + tenant_id: 'tenant-1', + tier: 'free', + scopes: ['generate'], + }), + }, + }); + + const initReq = rpcRequest('initialize', { protocolVersion: '2025-03-26', capabilities: {}, clientInfo: { name: 'test' } }); + const initRes = await handleMcpRequest(initReq, env); + const sessionId = initRes.headers.get('MCP-Session-Id')!; + + const req = rpcRequest( + 'tools/call', + { + name: 'image_generate', + arguments: { + prompt: 'a product mockup', + quality_tier: 'standard', + cf_model: '@cf/black-forest-labs/flux-2-klein-9b', + }, + }, + { 'MCP-Session-Id': sessionId }, + ); + const res = await handleMcpRequest(req, env); + const body = await res.json() as any; + + expect(body.result?.isError ?? body.error).toBeTruthy(); + const text = body.result?.content?.[0]?.text ?? body.error?.message ?? ''; + expect(text).toMatch(/pro plan|premium/i); + }); + it('allows pro-tier user who passes model=gemini-3-pro-image-preview', async () => { const env = makeEnv({ AUTH_SERVICE: { diff --git a/test/tool-registry.test.ts b/test/tool-registry.test.ts index c83ce82..a79ee6a 100644 --- a/test/tool-registry.test.ts +++ b/test/tool-registry.test.ts @@ -99,6 +99,31 @@ describe('buildAggregatedCatalog', () => { expect(gen.inputSchema['required']).toContain('prompt'); }); + it('image_generate exposes the full Cloudflare text-to-image override catalog', () => { + const catalog = buildAggregatedCatalog(); + const gen = catalog.find(t => t.name === 'image_generate')!; + const props = gen.inputSchema['properties'] as Record>; + const cfModel = props['cf_model']; + + expect(cfModel).toBeDefined(); + expect(cfModel['enum']).toHaveLength(11); + expect(cfModel['enum']).toContain('@cf/black-forest-labs/flux-2-klein-9b'); + expect(cfModel['enum']).toContain('@cf/leonardo/lucid-origin'); + expect(cfModel['enum']).toContain('@cf/stabilityai/stable-diffusion-xl-base-1.0'); + }); + + it('marks image_size as Gemini-only and does not advertise a default for Cloudflare calls', () => { + const catalog = buildAggregatedCatalog(); + const gen = catalog.find(t => t.name === 'image_generate')!; + const props = gen.inputSchema['properties'] as Record>; + const imageSize = props['image_size']; + + expect(imageSize).toBeDefined(); + expect(imageSize['default']).toBeUndefined(); + expect(String(imageSize['description'])).toMatch(/Gemini-only/i); + expect(String(imageSize['description'])).toMatch(/do not send with cf_model/i); + }); + }); describe('validateToolArguments', () => {