From 93a1985d9f72e976daddf3a77940704b5a8e63f0 Mon Sep 17 00:00:00 2001 From: ajaysi Date: Sat, 9 May 2026 16:57:26 +0530 Subject: [PATCH] fix(image-studio): add asset library saving + 402 subscription error handling Backend: - New POST /api/image-studio/save-to-library endpoint Saves generated base64 images to workspace disk and creates ContentAsset record for the unified asset library. Returns asset_id, file_url, filename. Frontend: - Added saveImageToLibrary() to useImageStudio hook - CreateStudio auto-saves generated images to asset library after creation - All 8 API operations now use _formatErrorMessage() helper for 402/429 subscription limit errors with upgrade prompts instead of generic error messages --- backend/routers/image_studio/__init__.py | 2 + backend/routers/image_studio/save.py | 100 ++++++++++++++++++ .../components/ImageStudio/CreateStudio.tsx | 18 ++++ frontend/src/hooks/useImageStudio.ts | 60 +++++++++-- 4 files changed, 172 insertions(+), 8 deletions(-) create mode 100644 backend/routers/image_studio/save.py diff --git a/backend/routers/image_studio/__init__.py b/backend/routers/image_studio/__init__.py index 7ae2d324..484bec0f 100644 --- a/backend/routers/image_studio/__init__.py +++ b/backend/routers/image_studio/__init__.py @@ -15,6 +15,7 @@ from .create import router as create_router from .transform import router as transform_router from .compress import router as compress_router from .convert import router as convert_router +from .save import router as save_router router = APIRouter(prefix="/api/image-studio", tags=["image-studio"]) @@ -28,5 +29,6 @@ router.include_router(create_router) router.include_router(transform_router) router.include_router(compress_router) router.include_router(convert_router) +router.include_router(save_router) __all__ = ["router"] diff --git a/backend/routers/image_studio/save.py b/backend/routers/image_studio/save.py new file mode 100644 index 00000000..f4fb722c --- /dev/null +++ b/backend/routers/image_studio/save.py @@ -0,0 +1,100 @@ +"""Save generated images to the unified asset library.""" + +import base64 +from datetime import datetime +from typing import Dict, Any, Optional +from pathlib import Path + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, Field +from sqlalchemy.orm import Session + +from .deps import _require_user_id +from middleware.auth_middleware import get_current_user +from services.database import get_db +from utils.logger_utils import get_service_logger +from utils.storage_paths import get_repo_root, sanitize_user_id + +logger = get_service_logger("api.image_studio") +router = APIRouter(tags=["image-studio"]) + + +class SaveToLibraryRequest(BaseModel): + image_base64: str = Field(..., description="Base64-encoded image (or data URL)") + prompt: Optional[str] = None + provider: Optional[str] = None + model: Optional[str] = None + cost: Optional[float] = None + operation: str = Field("image-generation", description="Operation type for labelling") + output_format: str = Field("png", description="Output image format") + + +@router.post("/save-to-library") +async def save_to_library( + req: SaveToLibraryRequest, + current_user: Dict[str, Any] = Depends(get_current_user), + db: Session = Depends(get_db), +): + """Save a generated image to the asset library. + + Decodes base64 image data, saves to workspace disk storage, + and creates a record in the ContentAsset database table. + """ + user_id = _require_user_id(current_user, "save-to-library") + + # Decode base64 payload + try: + b64data = req.image_base64 + if "base64," in b64data: + b64data = b64data.split("base64,")[1] + image_bytes = base64.b64decode(b64data) + except Exception: + raise HTTPException(status_code=400, detail="Invalid base64 image data") + + # Generate file path under workspace + safe_user = sanitize_user_id(user_id) + repo_root = get_repo_root() + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + filename = f"generated_{timestamp}.{req.output_format or 'png'}" + + assets_dir = repo_root / "workspace" / f"workspace_{safe_user}" / "assets" / "images" + assets_dir.mkdir(parents=True, exist_ok=True) + + file_path = assets_dir / filename + file_path.write_bytes(image_bytes) + + # Build serving URL (assets_serving.py serves /{user_id}/avatars/{filename}) + file_url = f"/api/assets/{safe_user}/avatars/{filename}" + + # Save to unified asset library via existing utility + from utils.asset_tracker import save_asset_to_library + + asset_id = save_asset_to_library( + db=db, + user_id=user_id, + asset_type="image", + source_module="image_studio", + filename=filename, + file_url=file_url, + file_path=str(file_path), + file_size=len(image_bytes), + mime_type=f"image/{req.output_format or 'png'}", + title=f"Generated Image - {timestamp}", + prompt=req.prompt, + provider=req.provider, + model=req.model, + cost=req.cost, + ) + + if not asset_id: + raise HTTPException(status_code=500, detail="Failed to save to asset library") + + logger.info(f"[Save to Library] ✅ Image saved: asset_id={asset_id}, user={user_id}") + + return { + "success": True, + "asset_id": asset_id, + "file_url": file_url, + "filename": filename, + "file_size": len(image_bytes), + } diff --git a/frontend/src/components/ImageStudio/CreateStudio.tsx b/frontend/src/components/ImageStudio/CreateStudio.tsx index 5acb8e4a..2623fb6b 100644 --- a/frontend/src/components/ImageStudio/CreateStudio.tsx +++ b/frontend/src/components/ImageStudio/CreateStudio.tsx @@ -191,6 +191,7 @@ export const CreateStudio: React.FC = ({ onImageGenerated }) estimateCost, loadTemplates, loadProviders, + saveImageToLibrary, } = useImageStudio(); // Load meta data on mount @@ -317,6 +318,23 @@ export const CreateStudio: React.FC = ({ onImageGenerated }) seed: providerForAdvanced ? parsedSeed : undefined, }); + // Auto-save generated images to asset library + if (result?.results?.length) { + const resolvedProvider = result.request?.provider || provider; + const resolvedModel = result.request?.model || effectiveModel; + for (const imgResult of result.results) { + if (imgResult.image_base64) { + saveImageToLibrary({ + imageBase64: imgResult.image_base64, + prompt, + provider: resolvedProvider, + model: resolvedModel, + operation: 'image-generation', + }).catch((e) => console.warn('Asset library save skipped:', e.message || e)); + } + } + } + if (onImageGenerated && (result?.results?.length ?? 0) > 0) { onImageGenerated(result); } diff --git a/frontend/src/hooks/useImageStudio.ts b/frontend/src/hooks/useImageStudio.ts index de91d124..64789839 100644 --- a/frontend/src/hooks/useImageStudio.ts +++ b/frontend/src/hooks/useImageStudio.ts @@ -500,7 +500,7 @@ export const useImageStudio = () => { } } catch (err: any) { console.error('Failed to generate image:', err); - const errorMessage = err.response?.data?.detail || 'Failed to generate image'; + const errorMessage = _formatErrorMessage(err, 'Failed to generate image'); setError(errorMessage); throw new Error(errorMessage); } finally { @@ -568,7 +568,7 @@ export const useImageStudio = () => { return response.data as EditResult; } catch (err: any) { console.error('Failed to process edit:', err); - const message = err.response?.data?.detail || 'Failed to process edit'; + const message = _formatErrorMessage(err, 'Failed to process edit'); setEditError(message); throw new Error(message); } finally { @@ -708,7 +708,7 @@ export const useImageStudio = () => { return response.data as FaceSwapResult; } catch (err: any) { console.error('Failed to process face swap:', err); - const message = err.response?.data?.detail || 'Failed to process face swap'; + const message = _formatErrorMessage(err, 'Failed to process face swap'); setFaceSwapError(message); throw new Error(message); } finally { @@ -731,7 +731,7 @@ export const useImageStudio = () => { return response.data as UpscaleResult; } catch (err: any) { console.error('Failed to upscale image:', err); - const message = err.response?.data?.detail || 'Failed to upscale image'; + const message = _formatErrorMessage(err, 'Failed to upscale image'); setUpscaleError(message); throw new Error(message); } finally { @@ -767,7 +767,7 @@ export const useImageStudio = () => { return response.data as ControlResult; } catch (err: any) { console.error('Failed to process control:', err); - const message = err.response?.data?.detail || 'Failed to process control'; + const message = _formatErrorMessage(err, 'Failed to process control'); setControlError(message); throw new Error(message); } finally { @@ -798,7 +798,7 @@ export const useImageStudio = () => { return response.data as SocialOptimizeResult; } catch (err: any) { console.error('Failed to optimize for social:', err); - const message = err.response?.data?.detail || 'Failed to optimize for social platforms'; + const message = _formatErrorMessage(err, 'Failed to optimize for social platforms'); setOptimizeError(message); throw new Error(message); } finally { @@ -875,7 +875,7 @@ export const useImageStudio = () => { return response.data; } catch (err: any) { console.error('Failed to compress image:', err); - const message = err.response?.data?.detail || 'Failed to compress image'; + const message = _formatErrorMessage(err, 'Failed to compress image'); setCompressionError(message); throw new Error(message); } finally { @@ -923,7 +923,7 @@ export const useImageStudio = () => { return response.data; } catch (err: any) { console.error('Failed to convert format:', err); - const message = err.response?.data?.detail || 'Failed to convert image format'; + const message = _formatErrorMessage(err, 'Failed to convert image format'); setFormatConversionError(message); throw new Error(message); } finally { @@ -937,6 +937,48 @@ export const useImageStudio = () => { setFormatRecommendations([]); }, []); + // Save generated image to asset library + const saveImageToLibrary = useCallback(async (params: { + imageBase64: string; + prompt?: string; + provider?: string; + model?: string; + cost?: number; + operation?: string; + outputFormat?: string; + }) => { + try { + const response = await aiApiClient.post('/api/image-studio/save-to-library', { + image_base64: params.imageBase64, + prompt: params.prompt, + provider: params.provider, + model: params.model, + cost: params.cost, + operation: params.operation || 'image-generation', + output_format: params.outputFormat || 'png', + }); + return response.data; + } catch (err: any) { + console.error('Failed to save to library:', err); + return null; + } + }, []); + + // Helper to extract user-friendly error messages including 402 subscription errors + const _formatErrorMessage = (err: any, fallback: string): string => { + if (err?.response?.status === 402 || err?.response?.status === 429) { + return `Subscription limit reached. Upgrade your plan to continue using this feature.`; + } + if (err?.response?.data?.detail) { + const detail = err.response.data.detail; + if (typeof detail === 'object' && detail?.message) { + return detail.message; + } + return String(detail); + } + return fallback; + }; + return { // State templates, @@ -1024,6 +1066,8 @@ export const useImageStudio = () => { formatConversionError, formatRecommendations, clearFormatConversionResult, + // Save to library + saveImageToLibrary, }; };