From 8a36826e35dcae09bfed012a247b8a80a3c2f612 Mon Sep 17 00:00:00 2001 From: Hoid Date: Fri, 6 Mar 2026 09:09:27 +0100 Subject: [PATCH] feat: add POST /v1/screenshots/batch endpoint - Batch screenshot endpoint: take 1-10 screenshots in a single request - Concurrent processing with Promise.allSettled (partial success support) - Upfront quota check for all URLs before processing - Per-URL SSRF validation via existing takeScreenshot() - Added incrementUsage() to usage middleware for granular tracking - 10 new tests covering all edge cases - Updated OpenAPI docs (JSDoc on route) - Updated Node.js and Python SDK READMEs with batch method docs --- sdk/node/README.md | 24 +++ sdk/python/README.md | 23 +++ src/index.ts | 3 + src/middleware/usage.ts | 11 ++ src/routes/__tests__/batch.test.ts | 210 ++++++++++++++++++++++++++ src/routes/batch.ts | 235 +++++++++++++++++++++++++++++ 6 files changed, 506 insertions(+) create mode 100644 src/routes/__tests__/batch.test.ts create mode 100644 src/routes/batch.ts diff --git a/sdk/node/README.md b/sdk/node/README.md index 85f406b..8833e7e 100644 --- a/sdk/node/README.md +++ b/sdk/node/README.md @@ -216,6 +216,30 @@ Returns a `Promise` containing the screenshot image. | `css` | `string` | — | Custom CSS to inject before capture (max 5000 chars) | | `clip` | `object` | — | Crop rectangle: `{x, y, width, height}` (mutually exclusive with fullPage/selector) | +### `snap.batch(urls, options?)` + +Take multiple screenshots in a single request. Each URL counts as one screenshot toward usage limits. + +```typescript +const results = await snap.batch( + ['https://example.com', 'https://example.org'], + { format: 'jpeg', width: 1920, height: 1080 } +); + +for (const result of results) { + if (result.status === 'success') { + fs.writeFileSync(`${result.url}.jpg`, Buffer.from(result.image, 'base64')); + } else { + console.error(`Failed: ${result.url} — ${result.error}`); + } +} +``` + +- **Max 10 URLs per batch** +- All options (format, width, height, etc.) are shared across all URLs +- Returns partial results — some may succeed while others fail +- Response is always JSON with `{ results: [...] }` + ### `snap.health()` Returns API health status. diff --git a/sdk/python/README.md b/sdk/python/README.md index 1de7308..c52b265 100644 --- a/sdk/python/README.md +++ b/sdk/python/README.md @@ -225,6 +225,29 @@ except SnapAPIError as e: | `css` | `str` | — | Custom CSS to inject before capture (max 5000 chars) | | `clip` | `dict` | — | Crop rectangle: `{"x": int, "y": int, "width": int, "height": int}` (mutually exclusive with full_page/selector) | +### `snap.batch(urls, **options) -> list[dict]` + +Take multiple screenshots in a single request. Each URL counts as one screenshot toward usage limits. + +```python +results = snap.batch( + ["https://example.com", "https://example.org"], + format="jpeg", width=1920, height=1080 +) + +for result in results: + if result["status"] == "success": + with open(f"{result['url']}.jpg", "wb") as f: + f.write(base64.b64decode(result["image"])) + else: + print(f"Failed: {result['url']} — {result['error']}") +``` + +- **Max 10 URLs per batch** +- All options (format, width, height, etc.) are shared across all URLs +- Returns partial results — some may succeed while others fail +- Response is always JSON with `{ "results": [...] }` + ### `snap.health() -> dict` Returns API health status. diff --git a/src/index.ts b/src/index.ts index a113ebc..8b98405 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,6 +7,7 @@ import path from "path"; import { fileURLToPath } from "url"; import rateLimit from "express-rate-limit"; import { screenshotRouter } from "./routes/screenshot.js"; +import { batchRouter } from "./routes/batch.js"; import { healthRouter } from "./routes/health.js"; import { playgroundRouter } from "./routes/playground.js"; import { authMiddleware } from "./middleware/auth.js"; @@ -101,6 +102,7 @@ app.use("/v1/playground", playgroundRouter); // Authenticated routes app.use("/v1/usage", usageRouter); app.use("/v1/screenshot", authMiddleware, usageMiddleware, screenshotRouter); +app.use("/v1/screenshots/batch", authMiddleware, batchRouter); // API info app.get("/api", (_req, res) => { @@ -110,6 +112,7 @@ app.get("/api", (_req, res) => { endpoints: [ "POST /v1/playground — Try the API (no auth, watermarked, 5 req/hr)", "POST /v1/screenshot — Take a screenshot (requires API key)", + "POST /v1/screenshots/batch — Take multiple screenshots (requires API key)", "GET /v1/usage — Usage statistics (requires API key)", "GET /health — Health check", ], diff --git a/src/middleware/usage.ts b/src/middleware/usage.ts index f4d5e3b..e46a2f7 100644 --- a/src/middleware/usage.ts +++ b/src/middleware/usage.ts @@ -89,3 +89,14 @@ export function usageMiddleware(req: any, res: any, next: any): void { export function getUsageForKey(key: string): { count: number; monthKey: string } | undefined { return usage.get(key); } + +export function incrementUsage(key: string): void { + const monthKey = getMonthKey(); + const record = usage.get(key); + if (!record || record.monthKey !== monthKey) { + usage.set(key, { count: 1, monthKey }); + } else { + record.count++; + } + dirtyKeys.add(key); +} diff --git a/src/routes/__tests__/batch.test.ts b/src/routes/__tests__/batch.test.ts new file mode 100644 index 0000000..37a1e08 --- /dev/null +++ b/src/routes/__tests__/batch.test.ts @@ -0,0 +1,210 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' + +// Mock dependencies before imports +vi.mock('../../services/screenshot.js', () => ({ + takeScreenshot: vi.fn() +})) + +vi.mock('../../services/logger.js', () => ({ + default: { + error: vi.fn(), + info: vi.fn(), + warn: vi.fn() + } +})) + +vi.mock('../../services/keys.js', () => ({ + isValidKey: vi.fn().mockResolvedValue(true), + getKeyInfo: vi.fn().mockResolvedValue({ key: 'test_key', tier: 'pro', email: 'test@test.com' }), + getTierLimit: vi.fn().mockReturnValue(5000), + loadKeys: vi.fn(), + getAllKeys: vi.fn().mockReturnValue([]) +})) + +vi.mock('../../middleware/usage.js', () => ({ + usageMiddleware: vi.fn((req: any, res: any, next: any) => next()), + getUsageForKey: vi.fn().mockReturnValue(undefined), + loadUsageData: vi.fn(), + incrementUsage: vi.fn() +})) + +const { takeScreenshot } = await import('../../services/screenshot.js') +const { getTierLimit } = await import('../../services/keys.js') +const { getUsageForKey, incrementUsage } = await import('../../middleware/usage.js') +const mockTakeScreenshot = vi.mocked(takeScreenshot) + +import express from 'express' +import request from 'supertest' +import { batchRouter } from '../batch.js' +import { authMiddleware } from '../../middleware/auth.js' + +// Create test app +function createApp() { + const app = express() + app.use(express.json()) + app.use('/v1/screenshots/batch', authMiddleware, batchRouter) + return app +} + +describe('POST /v1/screenshots/batch', () => { + let app: express.Express + + beforeEach(() => { + vi.clearAllMocks() + app = createApp() + mockTakeScreenshot.mockResolvedValue({ + buffer: Buffer.from('fake-png'), + contentType: 'image/png' + }) + vi.mocked(getTierLimit).mockReturnValue(5000) + vi.mocked(getUsageForKey).mockReturnValue(undefined) + }) + + it('should return 401 without API key', async () => { + const res = await request(app) + .post('/v1/screenshots/batch') + .send({ urls: ['https://example.com'] }) + + expect(res.status).toBe(401) + }) + + it('should return 400 when urls field is missing', async () => { + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ format: 'png' }) + + expect(res.status).toBe(400) + expect(res.body.error).toMatch(/urls/) + }) + + it('should return 400 when urls is empty array', async () => { + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: [] }) + + expect(res.status).toBe(400) + expect(res.body.error).toMatch(/urls/) + }) + + it('should return 400 when urls has more than 10 items', async () => { + const urls = Array.from({ length: 11 }, (_, i) => `https://example${i}.com`) + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls }) + + expect(res.status).toBe(400) + expect(res.body.error).toMatch(/10/) + }) + + it('should return results for 2 valid URLs', async () => { + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: ['https://example.com', 'https://example.org'] }) + + expect(res.status).toBe(200) + expect(res.body.results).toHaveLength(2) + expect(res.body.results[0]).toMatchObject({ + url: 'https://example.com', + status: 'success' + }) + expect(res.body.results[0].image).toBeDefined() + expect(res.body.results[1]).toMatchObject({ + url: 'https://example.org', + status: 'success' + }) + }) + + it('should return partial success when 1 URL fails', async () => { + mockTakeScreenshot + .mockResolvedValueOnce({ buffer: Buffer.from('fake-png'), contentType: 'image/png' }) + .mockRejectedValueOnce(new Error('Navigation timeout')) + + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: ['https://example.com', 'https://fail.example.com'] }) + + expect(res.status).toBe(200) + expect(res.body.results).toHaveLength(2) + expect(res.body.results[0].status).toBe('success') + expect(res.body.results[1].status).toBe('error') + expect(res.body.results[1].error).toBeDefined() + }) + + it('should reject SSRF URLs per-URL, not whole batch', async () => { + mockTakeScreenshot + .mockRejectedValueOnce(new Error('URL not allowed: private IP')) + .mockResolvedValueOnce({ buffer: Buffer.from('fake-png'), contentType: 'image/png' }) + + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: ['http://169.254.169.254/metadata', 'https://example.com'] }) + + expect(res.status).toBe(200) + expect(res.body.results[0].status).toBe('error') + expect(res.body.results[1].status).toBe('success') + }) + + it('should apply shared params (format, width) to all URLs', async () => { + await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ + urls: ['https://example.com', 'https://example.org'], + format: 'jpeg', + width: 1920, + height: 1080 + }) + + expect(mockTakeScreenshot).toHaveBeenCalledTimes(2) + expect(mockTakeScreenshot).toHaveBeenCalledWith( + expect.objectContaining({ + url: 'https://example.com', + format: 'jpeg', + width: 1920, + height: 1080 + }) + ) + expect(mockTakeScreenshot).toHaveBeenCalledWith( + expect.objectContaining({ + url: 'https://example.org', + format: 'jpeg', + width: 1920, + height: 1080 + }) + ) + }) + + it('should return 429 when not enough quota for batch size', async () => { + vi.mocked(getTierLimit).mockReturnValue(100) + vi.mocked(getUsageForKey).mockReturnValue({ count: 99, monthKey: '2026-03' }) + + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: ['https://example.com', 'https://example.org'] }) + + expect(res.status).toBe(429) + expect(res.body.error).toMatch(/limit/) + }) + + it('should increment usage for each successful screenshot', async () => { + mockTakeScreenshot + .mockResolvedValueOnce({ buffer: Buffer.from('ok'), contentType: 'image/png' }) + .mockRejectedValueOnce(new Error('fail')) + + const res = await request(app) + .post('/v1/screenshots/batch') + .set('Authorization', 'Bearer test_key') + .send({ urls: ['https://example.com', 'https://fail.com'] }) + + expect(res.status).toBe(200) + // incrementUsage should be called only once (for the successful one) + expect(incrementUsage).toHaveBeenCalledTimes(1) + }) +}) diff --git a/src/routes/batch.ts b/src/routes/batch.ts new file mode 100644 index 0000000..97d8cc9 --- /dev/null +++ b/src/routes/batch.ts @@ -0,0 +1,235 @@ +import { Router } from "express"; +import { takeScreenshot } from "../services/screenshot.js"; +import { getUsageForKey, incrementUsage } from "../middleware/usage.js"; +import { getTierLimit } from "../services/keys.js"; +import logger from "../services/logger.js"; + +export const batchRouter = Router(); + +/** + * @openapi + * /v1/screenshots/batch: + * post: + * tags: [Screenshots] + * summary: Take multiple screenshots in a single request + * description: > + * Capture multiple URLs in one API call. Each URL counts as one screenshot + * toward your usage limits. All parameters except `urls` are shared across + * all screenshots. Returns partial results if some URLs fail. + * operationId: batchScreenshots + * security: + * - BearerAuth: [] + * - ApiKeyAuth: [] + * requestBody: + * required: true + * content: + * application/json: + * schema: + * type: object + * required: [urls] + * properties: + * urls: + * type: array + * items: + * type: string + * format: uri + * minItems: 1 + * maxItems: 10 + * description: Array of URLs to capture (1-10) + * example: ["https://example.com", "https://example.org"] + * format: + * type: string + * enum: [png, jpeg, webp] + * default: png + * width: + * type: integer + * minimum: 320 + * maximum: 3840 + * default: 1280 + * height: + * type: integer + * minimum: 200 + * maximum: 2160 + * default: 800 + * fullPage: + * type: boolean + * default: false + * quality: + * type: integer + * minimum: 1 + * maximum: 100 + * default: 80 + * darkMode: + * type: boolean + * default: false + * css: + * type: string + * maxLength: 5000 + * js: + * type: string + * maxLength: 5000 + * selector: + * type: string + * maxLength: 200 + * userAgent: + * type: string + * maxLength: 500 + * delay: + * type: integer + * minimum: 0 + * maximum: 5000 + * waitForSelector: + * type: string + * hideSelectors: + * oneOf: + * - type: string + * - type: array + * items: + * type: string + * maxItems: 10 + * clip: + * type: object + * properties: + * x: + * type: integer + * y: + * type: integer + * width: + * type: integer + * height: + * type: integer + * examples: + * basic: + * summary: Two URLs + * value: + * urls: ["https://example.com", "https://example.org"] + * with_options: + * summary: With shared options + * value: + * urls: ["https://example.com", "https://example.org"] + * format: jpeg + * width: 1920 + * height: 1080 + * quality: 90 + * responses: + * 200: + * description: Batch results (may include partial failures) + * content: + * application/json: + * schema: + * type: object + * properties: + * results: + * type: array + * items: + * type: object + * properties: + * url: + * type: string + * status: + * type: string + * enum: [success, error] + * image: + * type: string + * description: Base64-encoded image (only on success) + * error: + * type: string + * description: Error message (only on error) + * 400: + * description: Invalid request + * content: + * application/json: + * schema: { $ref: "#/components/schemas/Error" } + * 401: + * description: Missing API key + * 429: + * description: Usage limit exceeded + */ +batchRouter.post("/", async (req: any, res: any) => { + const { urls, ...sharedParams } = req.body; + + // Validate urls + if (!urls || !Array.isArray(urls) || urls.length === 0) { + res.status(400).json({ error: "Missing required parameter: urls (array of 1-10 URLs)" }); + return; + } + + if (urls.length > 10) { + res.status(400).json({ error: "Maximum 10 URLs per batch request" }); + return; + } + + // Check usage quota for all URLs before starting + const keyInfo = req.apiKeyInfo; + if (keyInfo) { + const key = keyInfo.key; + const limit = getTierLimit(keyInfo.tier); + const currentUsage = getUsageForKey(key); + const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`; + const currentCount = (currentUsage && currentUsage.monthKey === monthKey) ? currentUsage.count : 0; + + if (currentCount + urls.length > limit) { + res.status(429).json({ + error: `Monthly limit would be exceeded. Need ${urls.length} screenshots but only ${limit - currentCount} remaining (${limit} limit for ${keyInfo.tier} tier).`, + usage: currentCount, + limit, + }); + return; + } + } + + // Extract shared screenshot params + const { + format, width, height, fullPage, quality, darkMode, css, js, + selector, userAgent, delay, waitForSelector, hideSelectors, clip, + waitUntil, deviceScale, + } = sharedParams; + + // Process all URLs concurrently + const settled = await Promise.allSettled( + urls.map((url: string) => + takeScreenshot({ + url, + format: format || undefined, + width: width ? parseInt(width, 10) || width : undefined, + height: height ? parseInt(height, 10) || height : undefined, + fullPage, + quality: quality ? parseInt(quality, 10) || quality : undefined, + darkMode, + css, + js, + selector, + userAgent, + delay: delay ? parseInt(delay, 10) || delay : undefined, + waitForSelector, + hideSelectors, + clip, + waitUntil, + deviceScale, + }) + ) + ); + + // Build results and track usage + const results = settled.map((result, i) => { + if (result.status === "fulfilled") { + // Increment usage for successful screenshot + if (keyInfo) { + incrementUsage(keyInfo.key); + } + return { + url: urls[i], + status: "success" as const, + image: result.value.buffer.toString("base64"), + }; + } else { + return { + url: urls[i], + status: "error" as const, + error: result.reason?.message || "Unknown error", + }; + } + }); + + res.json({ results }); +});