Enhance image attachment validation in AI query route

apps/backend/src/app/api/latest/ai/query/[mode]/route.ts

@@ -13,6 +13,7 @@ import { getVerifiedQaContext } from "@/lib/ai/verified-qa";
 import { SmartResponse } from "@/route-handlers/smart-response";
 import { createSmartRouteHandler } from "@/route-handlers/smart-route-handler";
 import { validateImageAttachments } from "@stackframe/stack-shared/dist/ai/image-limits";
+import { KnownErrors } from "@stackframe/stack-shared/dist/known-errors";
 import { yupMixed, yupObject, yupString } from "@stackframe/stack-shared/dist/schema-fields";
 import { getEnvVariable } from "@stackframe/stack-shared/dist/utils/env";
 import { StatusError } from "@stackframe/stack-shared/dist/utils/errors";
@@ -44,7 +45,15 @@ export const POST = createSmartRouteHandler({
     }
     const imageValidationResult = validateImageAttachments(messages);
     if (!imageValidationResult.ok) {
-      throw new StatusError(StatusError.BadRequest, imageValidationResult.reason);
+      const { failure } = imageValidationResult;
+      switch (failure.code) {
+        case "too_many": {
+          throw new KnownErrors.TooManyImageAttachments(failure.maxImages);
+        }
+        case "too_large": {
+          throw new KnownErrors.ImageAttachmentTooLarge(failure.maxBytes, failure.actualBytes);
+        }
+      }
     }
 
     const authenticatedApiKey = isAuthenticated

packages/stack-shared/src/ai/image-limits.ts

@@ -15,24 +15,33 @@ export function estimateBase64ByteLength(dataUrl: string): number {
   return Math.max(0, Math.floor((base64.length * 3) / 4) - padding);
 }
 
-type ValidationResult = { ok: true } | { ok: false, reason: string };
+export type ImageValidationFailure =
+  | { code: "too_many", maxImages: number }
+  | { code: "too_large", maxBytes: number, actualBytes: number };
+
+export type ImageValidationResult =
+  | { ok: true }
+  | { ok: false, failure: ImageValidationFailure, reason: string };
+
 type UnknownPart = { type?: unknown, image?: unknown };
 type MessageLike = { role?: unknown, content?: unknown };
 
-export function validateImageCount(imageCount: number): ValidationResult {
+export function validateImageCount(imageCount: number): ImageValidationResult {
   if (imageCount > MAX_IMAGES_PER_MESSAGE) {
     return {
       ok: false,
+      failure: { code: "too_many", maxImages: MAX_IMAGES_PER_MESSAGE },
       reason: `Maximum ${MAX_IMAGES_PER_MESSAGE} images per message.`,
     };
   }
   return { ok: true };
 }
 
-export function validateImageByteLength(bytes: number): ValidationResult {
+export function validateImageByteLength(bytes: number): ImageValidationResult {
   if (bytes > MAX_IMAGE_BYTES_PER_FILE) {
     return {
       ok: false,
+      failure: { code: "too_large", maxBytes: MAX_IMAGE_BYTES_PER_FILE, actualBytes: bytes },
       reason: `Image exceeds ${MAX_IMAGE_MB_PER_FILE}MB limit (${(bytes / 1024 / 1024).toFixed(2)}MB).`,
     };
   }
@@ -40,7 +49,7 @@ export function validateImageByteLength(bytes: number): ValidationResult {
 }
 
 /** Validates per-message image count and per-file size for user messages. */
-export function validateImageAttachments(messages: readonly MessageLike[]): ValidationResult {
+export function validateImageAttachments(messages: readonly MessageLike[]): ImageValidationResult {
   for (const msg of messages) {
     if (!Array.isArray(msg.content)) continue;
     let imageCount = 0;

packages/stack-shared/src/known-errors.tsx

@@ -1852,6 +1852,33 @@ const NewPurchasesBlocked = createKnownErrorConstructor(
   () => [] as const,
 );
 
+const TooManyImageAttachments = createKnownErrorConstructor(
+  KnownError,
+  "TOO_MANY_IMAGE_ATTACHMENTS",
+  (maxImages: number) => [
+    400,
+    `Maximum ${maxImages} images per message.`,
+    {
+      max_images: maxImages,
+    },
+  ] as const,
+  (json) => [json.max_images] as const,
+);
+
+const ImageAttachmentTooLarge = createKnownErrorConstructor(
+  KnownError,
+  "IMAGE_ATTACHMENT_TOO_LARGE",
+  (maxBytes: number, actualBytes: number) => [
+    400,
+    `Image exceeds ${maxBytes / (1024 * 1024)}MB limit (${(actualBytes / 1024 / 1024).toFixed(2)}MB).`,
+    {
+      max_bytes: maxBytes,
+      actual_bytes: actualBytes,
+    },
+  ] as const,
+  (json) => [json.max_bytes, json.actual_bytes] as const,
+);
+
 export type KnownErrors = {
   [K in keyof typeof KnownErrors]: InstanceType<typeof KnownErrors[K]>;
 };
@@ -2000,6 +2027,8 @@ export const KnownErrors = {
   AnalyticsQueryTimeout,
   AnalyticsQueryError,
   AnalyticsNotEnabled,
+  TooManyImageAttachments,
+  ImageAttachmentTooLarge,
 } satisfies Record<string, KnownErrorConstructor<any, any>>;