fix(image): align cf_model proxy schema and billing

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`
 

src/cost-attribution.ts

@@ -64,11 +64,30 @@ const MODEL_QUALITY_TIER: Record<string, string> = {
   '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<string, string> = {
+  '@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, unknown>): 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';
 }
 

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',

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',

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: {

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<string, Record<string, unknown>>;
+    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<string, Record<string, unknown>>;
+    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', () => {