diff --git a/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.test.ts b/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.test.ts index 83d7b0eb4..099c643a0 100644 --- a/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.test.ts +++ b/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.test.ts @@ -392,3 +392,94 @@ describe('ReplicateImageGenerationProvider.generate go_fast pricing', () => { expect(outputMp?.costOverride).toBe(Math.round(1.2 * 1_000_000)); }); }); + +// -- Upstream rejection handling -- + +describe('ReplicateImageGenerationProvider.generate upstream rejections', () => { + const generate = () => + withTestActor(() => + makeProvider().generate({ + model: 'black-forest-labs/flux-schnell', + prompt: 'hi', + }), + ); + + it('maps an NSFW refusal to 400 moderation_flagged without metering', async () => { + runMock.mockRejectedValueOnce( + new Error( + 'Prediction failed: Error generating image: NSFW content detected.', + ), + ); + + await expect(generate()).rejects.toMatchObject({ + statusCode: 400, + legacyCode: 'bad_request', + code: 'moderation_flagged', + message: 'Error generating image: NSFW content detected.', + fields: { provider: 'replicate' }, + }); + expect(incrementUsageSpy).not.toHaveBeenCalled(); + }); + + it('maps a sensitive-content (E005) refusal to 400 moderation_flagged', async () => { + runMock.mockRejectedValueOnce( + new Error( + 'Prediction failed: The input or output was flagged as sensitive. Please try again with different inputs. (E005)', + ), + ); + + await expect(generate()).rejects.toMatchObject({ + statusCode: 400, + code: 'moderation_flagged', + }); + }); + + it('maps any other failed prediction to 502 upstream_failed, keeping the cause', async () => { + const raw = new Error( + 'Prediction failed: q_descale must have shape (batch_size, num_heads_k)', + ); + runMock.mockRejectedValueOnce(raw); + + await expect(generate()).rejects.toMatchObject({ + statusCode: 502, + legacyCode: 'upstream_failed', + message: 'q_descale must have shape (batch_size, num_heads_k)', + fields: { provider: 'replicate' }, + cause: raw, + }); + expect(incrementUsageSpy).not.toHaveBeenCalled(); + }); + + it('strips markup and bounds the message when upstream returns an HTML error page', async () => { + const page = + '
' + + "${'x'.repeat(500)}
`; + runMock.mockRejectedValueOnce( + new Error( + `Prediction failed: Error generating image: Failed to generate: ${page}`, + ), + ); + + const err = await generate().catch((e) => e); + expect(err).toMatchObject({ + statusCode: 502, + legacyCode: 'upstream_failed', + }); + expect(err.message).toContain("Our services aren't available right now"); + expect(err.message).not.toMatch(/<|body\{/); + expect(err.message.length).toBeLessThanOrEqual(300); + }); + + it('passes a rejection that carries an HTTP status through untouched for the driver boundary', async () => { + const apiError = Object.assign( + new Error( + 'Request to https://api.replicate.com/v1/predictions failed with status 429 Too Many Requests', + ), + { name: 'ApiError', response: { status: 429 } }, + ); + runMock.mockRejectedValueOnce(apiError); + + await expect(generate()).rejects.toBe(apiError); + }); +}); diff --git a/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.ts b/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.ts index e41561aca..5114f80fa 100644 --- a/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.ts +++ b/src/backend/drivers/ai-image/providers/replicate/ReplicateImageGenerationProvider.ts @@ -33,6 +33,27 @@ import { const DEFAULT_MODEL = 'black-forest-labs/flux-schnell'; const DEFAULT_RATIO = { w: 1024, h: 1024 }; +const PREDICTION_FAILED_PREFIX = 'Prediction failed:'; +const MAX_UPSTREAM_MESSAGE_LENGTH = 300; +// Model-side content filters, as worded in failed-prediction errors. +const CONTENT_FILTER_PATTERN = + /\bnsfw\b|flagged as sensitive|sensitive content|content policy|\bE005\b/i; + +/** + * Strips markup and bounds length so an upstream HTML error page never rides + * through into a response body or an alarm signature. + */ +const sanitizeUpstreamMessage = (raw: string): string => { + const text = raw + .replace(/<(style|script)[\s\S]*?<\/\1>/gi, ' ') + .replace(/<[^>]*>/g, ' ') + .replace(/\s+/g, ' ') + .trim(); + return text.length > MAX_UPSTREAM_MESSAGE_LENGTH + ? `${text.slice(0, MAX_UPSTREAM_MESSAGE_LENGTH - 3)}...` + : text; +}; + export class ReplicateImageGenerationProvider implements IImageProvider { static readonly #CORE_PARAMS: readonly string[] = [ 'prompt', @@ -173,10 +194,15 @@ export class ReplicateImageGenerationProvider implements IImageProvider { singleImage, }); - const output = await this.#client.run( - selectedModel.replicateId as `${string}/${string}`, - { input }, - ); + let output: unknown; + try { + output = await this.#client.run( + selectedModel.replicateId as `${string}/${string}`, + { input }, + ); + } catch (err) { + throw this.#translatePredictionFailure(err); + } const url = this.#extractUrl(output); if (!url) { @@ -207,6 +233,42 @@ export class ReplicateImageGenerationProvider implements IImageProvider { return found ?? models.find((m) => m.id === DEFAULT_MODEL)!; } + /** + * A prediction that ran and ended `failed` reaches us as a plain Error with + * no HTTP status, so the driver boundary cannot classify it and it would + * surface as an unhandled 500. Content-filter refusals are the caller's to + * act on; anything else is an upstream fault. Errors that do carry a status + * (the SDK's ApiError) pass through untouched so the boundary translator + * still sees it. + */ + #translatePredictionFailure(err: unknown): unknown { + if (!(err instanceof Error)) return err; + if (!err.message.startsWith(PREDICTION_FAILED_PREFIX)) return err; + + const detail = sanitizeUpstreamMessage( + err.message.slice(PREDICTION_FAILED_PREFIX.length), + ); + const fields = { provider: 'replicate' }; + + if (CONTENT_FILTER_PATTERN.test(detail)) { + return new HttpError( + 400, + detail || 'Prompt or output was rejected by the content filter', + { + legacyCode: 'bad_request', + code: 'moderation_flagged', + fields, + cause: err, + }, + ); + } + return new HttpError(502, detail || 'Replicate prediction failed', { + legacyCode: 'upstream_failed', + fields, + cause: err, + }); + } + /** * Builds the Replicate API input payload from already-aliased+transformed * params. Image inputs and prompt/ratio are placed explicitly; everything diff --git a/src/docs/src/AI/txt2img.md b/src/docs/src/AI/txt2img.md index 5f644eeae..ff07b8997 100755 --- a/src/docs/src/AI/txt2img.md +++ b/src/docs/src/AI/txt2img.md @@ -171,6 +171,18 @@ Absolute paths (`/username/Pictures/sunset.png`) and home-relative paths (`~/Pic A `Promise` that resolves to an `HTMLImageElement`. The element’s `src` points at a data URL containing the image. +## Errors + +A rejection carries the error body as the backend sent it: `{ message, code }`, plus `errorCode` when a more specific code is available alongside a general one. + +| Code | Meaning | +| --- | --- | +| `errorCode: moderation_flagged` | The model's content filter refused the prompt or the generated image. Arrives as HTTP 400 with `code: bad_request`. Change the prompt rather than retrying it as-is. Not every provider reports refusals distinctly; when one does, this is how. | +| `upstream_failed` | The provider accepted the request but generation failed on their side. Safe to retry. | +| `insufficient_funds` | Your balance cannot cover the estimated cost of the image. Arrives as HTTP 402. | + +Other `upstream_*` codes mean the provider rejected the request or was unavailable; the `message` carries the provider's reason. + ## Examples Generate an image of a cat using AI