replace upload-url with server action

This commit is contained in:
arkml 2025-09-26 17:35:03 +05:30
parent aa988220da
commit 9b53e4d880
6 changed files with 88 additions and 128 deletions

View file

@ -0,0 +1,48 @@
"use server";
import { S3Client, PutObjectCommand } from '@aws-sdk/client-s3';
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
import crypto from 'crypto';
import { authCheck } from '@/app/actions/auth.actions';
import { USE_AUTH } from '@/app/lib/feature_flags';
export async function getUploadUrlForImage(mimeType: string): Promise<{ id: string; key: string; uploadUrl: string; url: string; mimeType: string }> {
// Enforce auth in server action context (supports guest mode when auth disabled)
if (USE_AUTH) {
await authCheck();
}
if (!mimeType || typeof mimeType !== 'string') {
throw new Error('mimeType is required');
}
const bucket = process.env.RAG_UPLOADS_S3_BUCKET || '';
if (!bucket) {
throw new Error('S3 bucket not configured');
}
const ext = mimeType === 'image/jpeg' ? '.jpg'
: mimeType === 'image/webp' ? '.webp'
: mimeType === 'image/png' ? '.png'
: '.bin';
const id = crypto.randomUUID();
const idWithExt = `${id}${ext}`;
const last2 = id.slice(-2).padStart(2, '0');
const dirA = last2.charAt(0);
const dirB = last2.charAt(1);
const key = `uploaded_images/${dirA}/${dirB}/${idWithExt}`;
const region = process.env.RAG_UPLOADS_S3_REGION || 'us-east-1';
const s3 = new S3Client({
region,
credentials: process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY ? {
accessKeyId: process.env.AWS_ACCESS_KEY_ID as string,
secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY as string,
} : undefined,
});
const command = new PutObjectCommand({ Bucket: bucket, Key: key, ContentType: mimeType });
const uploadUrl = await getSignedUrl(s3, command, { expiresIn: 600 });
return { id: idWithExt, key, uploadUrl, url: `/api/uploaded-images/${idWithExt}`, mimeType };
}

View file

@ -1,21 +1,14 @@
import { NextRequest, NextResponse } from 'next/server'; import { NextRequest, NextResponse } from 'next/server';
import { S3Client, GetObjectCommand, HeadObjectCommand } from '@aws-sdk/client-s3'; import { S3Client, GetObjectCommand } from '@aws-sdk/client-s3';
import { GoogleGenerativeAI } from '@google/generative-ai'; import { GoogleGenerativeAI } from '@google/generative-ai';
import { UsageTracker } from '@/app/lib/billing'; import { UsageTracker } from '@/app/lib/billing';
import { logUsage } from '@/app/actions/billing.actions'; import { logUsage } from '@/app/actions/billing.actions';
import { authCheck } from '@/app/actions/auth.actions'; import { requireAuth } from '@/app/lib/auth';
import { USE_AUTH } from '@/app/lib/feature_flags';
export async function POST(request: NextRequest) { export async function POST(request: NextRequest) {
try { try {
// Require authentication if enabled // Require authentication (handles guest mode internally when auth disabled)
try { await requireAuth();
if (USE_AUTH) {
await authCheck();
}
} catch (_) {
return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
}
const { id } = await request.json(); const { id } = await request.json();
if (!id || typeof id !== 'string') { if (!id || typeof id !== 'string') {
@ -36,23 +29,13 @@ export async function POST(request: NextRequest) {
} : undefined, } : undefined,
}); });
const last2 = id.slice(-2).padStart(2, '0'); // `id` includes extension (e.g., "<uuid>.png"). Shard using the UUID part.
const lastDot = id.lastIndexOf('.');
const idWithoutExt = lastDot > 0 ? id.slice(0, lastDot) : id;
const last2 = idWithoutExt.slice(-2).padStart(2, '0');
const dirA = last2.charAt(0); const dirA = last2.charAt(0);
const dirB = last2.charAt(1); const dirB = last2.charAt(1);
const baseKey = `uploaded_images/${dirA}/${dirB}/${id}`; const key = `uploaded_images/${dirA}/${dirB}/${id}`;
const exts = ['.png', '.jpg', '.webp', '.bin'];
let foundExt: string | null = null;
for (const ext of exts) {
try {
await s3.send(new HeadObjectCommand({ Bucket: bucket, Key: `${baseKey}${ext}` }));
foundExt = ext; break;
} catch {}
}
if (!foundExt) {
return NextResponse.json({ error: 'Not found' }, { status: 404 });
}
const key = `${baseKey}${foundExt}`;
const resp = await s3.send(new GetObjectCommand({ Bucket: bucket, Key: key })); const resp = await s3.send(new GetObjectCommand({ Bucket: bucket, Key: key }));
const contentType = resp.ContentType || 'application/octet-stream'; const contentType = resp.ContentType || 'application/octet-stream';
const body = resp.Body as any; const body = resp.Body as any;

View file

@ -4,7 +4,7 @@ import crypto from 'crypto';
import { tempBinaryCache } from '@/src/application/services/temp-binary-cache'; import { tempBinaryCache } from '@/src/application/services/temp-binary-cache';
import { GoogleGenerativeAI } from '@google/generative-ai'; import { GoogleGenerativeAI } from '@google/generative-ai';
import { UsageTracker, getCustomerForUserId, logUsage as libLogUsage } from '@/app/lib/billing'; import { UsageTracker, getCustomerForUserId, logUsage as libLogUsage } from '@/app/lib/billing';
import { authCheck } from '@/app/actions/auth.actions'; import { requireAuth } from '@/app/lib/auth';
import { USE_AUTH, USE_BILLING } from '@/app/lib/feature_flags'; import { USE_AUTH, USE_BILLING } from '@/app/lib/feature_flags';
// POST /api/uploaded-images // POST /api/uploaded-images
@ -15,13 +15,13 @@ export async function POST(request: NextRequest) {
try { try {
// Require authentication if enabled // Require authentication if enabled
let currentUser: any | null = null; let currentUser: any | null = null;
try {
if (USE_AUTH) { if (USE_AUTH) {
currentUser = await authCheck(); try {
} currentUser = await requireAuth();
} catch (_) { } catch (_) {
return NextResponse.json({ error: 'Unauthorized' }, { status: 401 }); return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
} }
}
const contentType = request.headers.get('content-type') || ''; const contentType = request.headers.get('content-type') || '';
if (!contentType.includes('multipart/form-data')) { if (!contentType.includes('multipart/form-data')) {
@ -99,7 +99,7 @@ export async function POST(request: NextRequest) {
ContentType: mime, ContentType: mime,
})); }));
const url = `/api/uploaded-images/${imageId}`; const url = `/api/uploaded-images/${imageId}${ext}`;
// Log usage to billing similar to rag-worker // Log usage to billing similar to rag-worker
try { try {
@ -116,7 +116,7 @@ export async function POST(request: NextRequest) {
// ignore billing logging errors // ignore billing logging errors
} }
return NextResponse.json({ url, storage: 's3', id: imageId, mimeType: mime, description: descriptionMarkdown }); return NextResponse.json({ url, storage: 's3', id: `${imageId}${ext}`, mimeType: mime, description: descriptionMarkdown });
} }
// Otherwise store in temp cache and return temp URL // Otherwise store in temp cache and return temp URL

View file

@ -1,57 +0,0 @@
import { NextRequest, NextResponse } from 'next/server';
import { S3Client, PutObjectCommand } from '@aws-sdk/client-s3';
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
import crypto from 'crypto';
import { authCheck } from '@/app/actions/auth.actions';
import { USE_AUTH } from '@/app/lib/feature_flags';
export async function POST(request: NextRequest) {
try {
// Require authentication if enabled
try {
if (USE_AUTH) {
await authCheck();
}
} catch (_) {
return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
}
const bucket = process.env.RAG_UPLOADS_S3_BUCKET || '';
if (!bucket) {
return NextResponse.json({ error: 'S3 bucket not configured' }, { status: 500 });
}
const { mimeType } = await request.json();
if (!mimeType || typeof mimeType !== 'string') {
return NextResponse.json({ error: 'mimeType is required' }, { status: 400 });
}
const ext = mimeType === 'image/jpeg' ? '.jpg'
: mimeType === 'image/webp' ? '.webp'
: mimeType === 'image/png' ? '.png'
: '.bin';
const id = crypto.randomUUID();
const last2 = id.slice(-2).padStart(2, '0');
const dirA = last2.charAt(0);
const dirB = last2.charAt(1);
const key = `uploaded_images/${dirA}/${dirB}/${id}${ext}`;
const region = process.env.RAG_UPLOADS_S3_REGION || 'us-east-1';
const s3 = new S3Client({
region,
credentials: process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY ? {
accessKeyId: process.env.AWS_ACCESS_KEY_ID as string,
secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY as string,
} : undefined,
});
const command = new PutObjectCommand({ Bucket: bucket, Key: key, ContentType: mimeType });
const uploadUrl = await getSignedUrl(s3, command, { expiresIn: 600 });
return NextResponse.json({ id, key, uploadUrl, url: `/api/uploaded-images/${id}`, mimeType });
} catch (e) {
console.error('upload-url error', e);
return NextResponse.json({ error: 'Failed to create upload URL' }, { status: 500 });
}
}

View file

@ -89,18 +89,12 @@ export function ComposeBoxPlayground({
uploadAbortRef.current = controller; uploadAbortRef.current = controller;
let usedFallback = false; let usedFallback = false;
try { try {
// 1) Request a presigned S3 upload URL // 1) Request a presigned S3 upload URL via server action
const urlRes = await fetch('/api/uploaded-images/upload-url', { const { getUploadUrlForImage } = await import('@/app/actions/uploaded-images.actions');
method: 'POST', const urlData = await getUploadUrlForImage(file.type);
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ mimeType: file.type }),
signal: controller.signal,
});
if (!urlRes.ok) throw new Error(`Failed to get upload URL: ${urlRes.status}`);
const urlData = await urlRes.json();
const uploadUrl: string | undefined = urlData?.uploadUrl; const uploadUrl: string | undefined = urlData?.uploadUrl;
const imageId: string | undefined = urlData?.id; const imageId: string | undefined = urlData?.id; // includes extension
const imageUrl: string | undefined = urlData?.url; const imageUrl: string | undefined = urlData?.url; // points to /api/uploaded-images/<idWithExt>
if (!uploadUrl || !imageId || !imageUrl) throw new Error('Invalid upload URL response'); if (!uploadUrl || !imageId || !imageUrl) throw new Error('Invalid upload URL response');
// 2) Upload the file directly to S3 // 2) Upload the file directly to S3

View file

@ -90,21 +90,14 @@ export async function invokeGenerateImageTool(
secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY, secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
} as any : undefined, } as any : undefined,
}); });
const id = inputImageUrl.split('/api/uploaded-images/')[1]; const idWithExt = inputImageUrl.split('/api/uploaded-images/')[1];
const last2 = id.slice(-2).padStart(2, '0'); // Shard by the UUID portion (before extension)
const lastDot = idWithExt.lastIndexOf('.');
const idWithoutExt = lastDot > 0 ? idWithExt.slice(0, lastDot) : idWithExt;
const last2 = idWithoutExt.slice(-2).padStart(2, '0');
const dirA = last2.charAt(0); const dirA = last2.charAt(0);
const dirB = last2.charAt(1); const dirB = last2.charAt(1);
const baseKey = `uploaded_images/${dirA}/${dirB}/${id}`; const key = `uploaded_images/${dirA}/${dirB}/${idWithExt}`;
const exts = ['.png', '.jpg', '.webp', '.bin'];
let foundExt: string | null = null;
for (const ext of exts) {
try {
await s3.send(new HeadObjectCommand({ Bucket: bucket, Key: `${baseKey}${ext}` }));
foundExt = ext; break;
} catch {}
}
if (foundExt) {
const key = `${baseKey}${foundExt}`;
const resp = await s3.send(new GetObjectCommand({ Bucket: bucket, Key: key })); const resp = await s3.send(new GetObjectCommand({ Bucket: bucket, Key: key }));
const chunks: Buffer[] = []; const chunks: Buffer[] = [];
const body = resp.Body as any; const body = resp.Body as any;
@ -119,7 +112,6 @@ export async function invokeGenerateImageTool(
imageBuf = Buffer.concat(chunks); imageBuf = Buffer.concat(chunks);
} }
} }
}
} else if (inputImageUrl.startsWith('/api/generated-images/')) { } else if (inputImageUrl.startsWith('/api/generated-images/')) {
const bucket = process.env.RAG_UPLOADS_S3_BUCKET || ''; const bucket = process.env.RAG_UPLOADS_S3_BUCKET || '';
if (bucket) { if (bucket) {