diff --git a/public/sw.js b/public/sw.js index d24f9c6ec6..8f6017351a 100644 --- a/public/sw.js +++ b/public/sw.js @@ -138,6 +138,6 @@ self.addEventListener("notificationclick", (event) => { if (clients.openWindow) { return clients.openWindow(urlToOpen); } - }), + }) ); }); diff --git a/src/app/(dashboard)/dashboard/relay/RelayProxyClient.tsx b/src/app/(dashboard)/dashboard/relay/RelayProxyClient.tsx new file mode 100644 index 0000000000..e84011adb1 --- /dev/null +++ b/src/app/(dashboard)/dashboard/relay/RelayProxyClient.tsx @@ -0,0 +1,245 @@ +"use client"; + +import { useState, useEffect, useCallback } from "react"; +import Card from "@/shared/components/Card"; +import Badge from "@/shared/components/Badge"; +import Button from "@/shared/components/Button"; +import { useNotificationStore } from "@/store/notificationStore"; + +interface RelayToken { + id: string; + name: string; + tokenPrefix: string; + description: string; + comboId: string | null; + allowedModels: string; + maxRequestsPerMinute: number; + maxRequestsPerDay: number; + enabled: boolean; + createdAt: number; + lastUsedAt: number | null; +} + +export default function RelayProxyClient() { + const [tokens, setTokens] = useState([]); + const [loading, setLoading] = useState(true); + const [showCreate, setShowCreate] = useState(false); + const [newTokenData, setNewTokenData] = useState<{ rawToken: string; name: string } | null>(null); + const [form, setForm] = useState({ name: "", description: "", maxRpm: "60", maxRpd: "10000" }); + const addNotification = useNotificationStore((s) => s.addNotification); + + const fetchTokens = useCallback(async () => { + setLoading(true); + try { + const res = await fetch("/api/relay/tokens"); + const data = await res.json(); + setTokens(Array.isArray(data) ? data : []); + } catch { + setTokens([]); + } finally { + setLoading(false); + } + }, []); + + useEffect(() => { fetchTokens(); }, [fetchTokens]); + + const createToken = async () => { + if (!form.name.trim()) return; + try { + const res = await fetch("/api/relay/tokens", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + name: form.name, + description: form.description, + maxRequestsPerMinute: Number(form.maxRpm), + maxRequestsPerDay: Number(form.maxRpd), + }), + }); + const data = await res.json(); + if (res.ok) { + setNewTokenData({ rawToken: data.rawToken, name: data.name }); + setForm({ name: "", description: "", maxRpm: "60", maxRpd: "10000" }); + setShowCreate(false); + addNotification({ type: "success", message: "Relay token created" }); + fetchTokens(); + } else { + addNotification({ type: "error", message: data.error || "Failed to create token" }); + } + } catch { + addNotification({ type: "error", message: "Failed to create token" }); + } + }; + + const toggleToken = async (id: string, enabled: boolean) => { + try { + await fetch(`/api/relay/tokens/${id}`, { + method: "PATCH", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ enabled }), + }); + fetchTokens(); + } catch { + addNotification({ type: "error", message: "Failed to toggle token" }); + } + }; + + const deleteToken = async (id: string) => { + if (!confirm("Delete this relay token? This cannot be undone.")) return; + try { + await fetch(`/api/relay/tokens/${id}`, { method: "DELETE" }); + addNotification({ type: "success", message: "Token deleted" }); + fetchTokens(); + } catch { + addNotification({ type: "error", message: "Failed to delete token" }); + } + }; + + return ( +
+
+
+

Serverless Relay Proxies

+

+ Create public API endpoints that proxy to OmniRoute with rate limiting and access control +

+
+ +
+ + {/* Create Form */} + {showCreate && ( + +
+

Create Relay Token

+
+
+ + setForm({ ...form, name: e.target.value })} + placeholder="my-api-relay" + /> +
+
+ + setForm({ ...form, description: e.target.value })} + placeholder="For my serverless functions" + /> +
+
+ + setForm({ ...form, maxRpm: e.target.value })} + /> +
+
+ + setForm({ ...form, maxRpd: e.target.value })} + /> +
+
+ +
+ + )} + + {/* Token Display (shown once after creation) */} + {newTokenData && ( + +
+

+ Token Created — Copy it now! +

+
+

Token for {newTokenData.name}:

+ + {newTokenData.rawToken} + +
+

+ This token will not be shown again. Store it securely. +

+ +
+ + )} + + {/* Usage Guide */} + +
+

Usage

+

+ Send requests to your relay endpoint: +

+ 
+{`curl http://localhost:20128/v1/relay/chat/completions \\
+  -H "Authorization: Bearer relay_..." \\
+  -H "Content-Type: application/json" \\
+  -d '{"model":"claude-sonnet-4","messages":[{"role":"user","content":"Hello"}]}'`}
+
+
+ + + {/* Tokens List */} + +
+

+ Relay Tokens ({tokens.length}) +

+ {loading ? ( +

Loading...

+ ) : tokens.length === 0 ? ( +

No relay tokens configured. Create one to get started.

+ ) : ( +
+ {tokens.map((t) => ( +
+
+
+
+
{t.name}
+
{t.tokenPrefix}...
+ {t.description && ( +
{t.description}
+ )} +
+
+
+ {t.maxRequestsPerMinute}/min + {t.maxRequestsPerDay}/day + + +
+
+ ))} +
+ )} +
+ +
+ ); +} diff --git a/src/app/(dashboard)/dashboard/relay/page.tsx b/src/app/(dashboard)/dashboard/relay/page.tsx new file mode 100644 index 0000000000..f9cc786e3f --- /dev/null +++ b/src/app/(dashboard)/dashboard/relay/page.tsx @@ -0,0 +1,11 @@ +import type { Metadata } from "next"; +import RelayProxyClient from "./RelayProxyClient"; + +export const metadata: Metadata = { + title: "OmniRoute — Relay Proxies", + description: "Serverless relay proxy endpoints for your AI infrastructure", +}; + +export default function RelayProxyPage() { + return ; +} diff --git a/src/app/api/relay/tokens/[id]/route.ts b/src/app/api/relay/tokens/[id]/route.ts new file mode 100644 index 0000000000..5f2b7a9cc8 --- /dev/null +++ b/src/app/api/relay/tokens/[id]/route.ts @@ -0,0 +1,51 @@ +import { NextResponse } from "next/server"; +import { getRelayToken, updateRelayToken, deleteRelayToken, toggleRelayToken, getRelayLogs, getRelayUsage } from "@/lib/db/relayProxies"; + +export async function GET(request: Request, { params }: { params: Promise<{ id: string }> }) { + const { id } = await params; + const token = getRelayToken(id); + if (!token) return NextResponse.json({ error: "Token not found" }, { status: 404 }); + + // Get usage stats + const now = Math.floor(Date.now() / 1000); + const lastHour = getRelayUsage(id, now - 3600); + const lastDay = getRelayUsage(id, now - 86400); + const logs = getRelayLogs(id, 20); + + return NextResponse.json({ + ...token, + usage: { lastHour, lastDay }, + logs, + }); +} + +export async function PATCH(request: Request, { params }: { params: Promise<{ id: string }> }) { + const { id } = await params; + const body = await request.json(); + + if (body.enabled !== undefined) { + const token = toggleRelayToken(id, body.enabled); + if (!token) return NextResponse.json({ error: "Token not found" }, { status: 404 }); + return NextResponse.json(token); + } + + const token = updateRelayToken(id, { + name: body.name, + description: body.description, + comboId: body.comboId, + allowedModels: body.allowedModels, + maxTokensPerRequest: body.maxTokensPerRequest, + maxRequestsPerMinute: body.maxRequestsPerMinute, + maxRequestsPerDay: body.maxRequestsPerDay, + maxCostPerDay: body.maxCostPerDay, + }); + + if (!token) return NextResponse.json({ error: "Token not found" }, { status: 404 }); + return NextResponse.json(token); +} + +export async function DELETE(request: Request, { params }: { params: Promise<{ id: string }> }) { + const { id } = await params; + deleteRelayToken(id); + return NextResponse.json({ success: true }); +} diff --git a/src/app/api/relay/tokens/route.ts b/src/app/api/relay/tokens/route.ts new file mode 100644 index 0000000000..e0963ad946 --- /dev/null +++ b/src/app/api/relay/tokens/route.ts @@ -0,0 +1,53 @@ +import { NextResponse } from "next/server"; +import { getRelayTokens, createRelayToken } from "@/lib/db/relayProxies"; + +export async function GET() { + const tokens = getRelayTokens(); + // Strip hash from response + const safe = tokens.map((t) => ({ + id: t.id, + name: t.name, + tokenPrefix: t.tokenPrefix, + description: t.description, + comboId: t.comboId, + allowedModels: t.allowedModels, + maxTokensPerRequest: t.maxTokensPerRequest, + maxRequestsPerMinute: t.maxRequestsPerMinute, + maxRequestsPerDay: t.maxRequestsPerDay, + maxCostPerDay: t.maxCostPerDay, + enabled: t.enabled, + createdAt: t.createdAt, + updatedAt: t.updatedAt, + expiresAt: t.expiresAt, + lastUsedAt: t.lastUsedAt, + })); + return NextResponse.json(safe); +} + +export async function POST(request: Request) { + try { + const body = await request.json(); + const token = createRelayToken({ + name: body.name, + description: body.description, + comboId: body.comboId, + allowedModels: body.allowedModels, + maxTokensPerRequest: body.maxTokensPerRequest, + maxRequestsPerMinute: body.maxRequestsPerMinute, + maxRequestsPerDay: body.maxRequestsPerDay, + maxCostPerDay: body.maxCostPerDay, + expiresAt: body.expiresAt, + metadata: body.metadata, + }); + + return NextResponse.json({ + id: token.id, + name: token.name, + rawToken: token.rawToken, + tokenPrefix: token.tokenPrefix, + }); + } catch (error) { + const message = error instanceof Error ? error.message : "Unknown error"; + return NextResponse.json({ error: message }, { status: 400 }); + } +} diff --git a/src/app/api/v1/relay/chat/completions/route.ts b/src/app/api/v1/relay/chat/completions/route.ts new file mode 100644 index 0000000000..367a91f0ae --- /dev/null +++ b/src/app/api/v1/relay/chat/completions/route.ts @@ -0,0 +1,175 @@ +/** + * POST /api/v1/relay/chat/completions + * + * Serverless Relay Proxy endpoint. + * Authenticates via relay token, applies rate limits, then proxies + * to the internal OmniRoute chat completions pipeline. + */ + +import { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors"; +import { handleChat } from "@/sse/handlers/chat"; +import { createInjectionGuard } from "@/middleware/promptInjectionGuard"; +import { getRelayTokenByHash, checkRateLimit, recordRelayUsage } from "@/lib/db/relayProxies"; +import { createHash } from "node:crypto"; + +const injectionGuard = createInjectionGuard(); + +export async function OPTIONS() { + return handleCorsOptions(); +} + +function extractToken(request: Request): string | null { + const auth = request.headers.get("authorization") || ""; + const match = auth.match(/^Bearer\s+(.+)$/i); + if (match) return match[1]; + + // Also check X-Relay-Token header + return request.headers.get("x-relay-token"); +} + +function hashToken(token: string): string { + return createHash("sha256").update(token).digest("hex"); +} + +export async function POST(request: Request) { + const startTime = Date.now(); + const clientIp = request.headers.get("x-forwarded-for")?.split(",")[0]?.trim() + || request.headers.get("x-real-ip") + || "unknown"; + const userAgent = request.headers.get("user-agent") || "unknown"; + + try { + // 1. Authenticate + const rawToken = extractToken(request); + if (!rawToken) { + return new Response( + JSON.stringify({ error: { message: "Missing relay token", type: "auth_error", code: "RELAY_AUTH_001" } }), + { status: 401, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } + + const tokenHash = hashToken(rawToken); + const token = getRelayTokenByHash(tokenHash); + if (!token) { + recordRelayUsage("unknown", { + requestId: request.headers.get("x-request-id") || undefined, + status: "auth_failed", + statusCode: 401, + latencyMs: Date.now() - startTime, + clientIp, + userAgent, + }); + return new Response( + JSON.stringify({ error: { message: "Invalid relay token", type: "auth_error", code: "RELAY_AUTH_002" } }), + { status: 401, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } + + // Check expiration + if (token.expiresAt && Math.floor(Date.now() / 1000) > token.expiresAt) { + return new Response( + JSON.stringify({ error: { message: "Relay token expired", type: "auth_error", code: "RELAY_AUTH_003" } }), + { status: 401, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } + + // 2. Rate limit check + const rateCheck = checkRateLimit(token.id); + if (!rateCheck.allowed) { + recordRelayUsage(token.id, { + requestId: request.headers.get("x-request-id") || undefined, + status: "rate_limited", + statusCode: 429, + latencyMs: Date.now() - startTime, + clientIp, + userAgent, + }); + return new Response( + JSON.stringify({ error: { message: "Rate limit exceeded", type: "rate_limited", code: "RELAY_RATE_001" } }), + { + status: 429, + headers: { + ...CORS_HEADERS, + "Content-Type": "application/json", + "Retry-After": String(rateCheck.resetIn), + "X-RateLimit-Remaining": "0", + }, + }, + ); + } + + // 3. Clone request and forward to internal handler + const cloned = request.clone(); + + // Prompt injection guard (same as main endpoint) + try { + const body = await cloned.json().catch(() => null); + if (body) { + const { blocked, result } = injectionGuard(body); + if (blocked) { + recordRelayUsage(token.id, { + requestId: request.headers.get("x-request-id") || undefined, + status: "error", + statusCode: 400, + latencyMs: Date.now() - startTime, + clientIp, + userAgent, + }); + return new Response( + JSON.stringify({ + error: { message: "Request blocked: potential prompt injection detected", type: "injection_detected", code: "SECURITY_001", detections: result.detections.length }, + }), + { status: 400, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } + + // Check allowed models + const allowedModels: string[] = JSON.parse(token.allowedModels); + if (allowedModels.length > 0 && !allowedModels.includes("*")) { + const model = (body as { model?: string }).model || ""; + const allowed = allowedModels.some( + (p) => model === p || (p.endsWith("*") && model.startsWith(p.slice(0, -1))), + ); + if (!allowed) { + return new Response( + JSON.stringify({ error: { message: `Model "${model}" not allowed by this relay token`, type: "model_not_allowed", code: "RELAY_MODEL_001" } }), + { status: 403, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } + } + } + } catch { + // Continue even if guard fails + } + + // 4. Proxy to internal handler + const originalRequest = new Request(request.url.replace("/relay/chat/completions", "/chat/completions"), request); + const response = await handleChat(originalRequest); + + // 5. Record usage (async, don't block response) + const latencyMs = Date.now() - startTime; + recordRelayUsage(token.id, { + requestId: request.headers.get("x-request-id") || undefined, + status: response.status < 500 ? "success" : "error", + statusCode: response.status, + latencyMs, + clientIp, + userAgent, + }); + + // Add relay headers + const newHeaders = new Headers(response.headers); + newHeaders.set("X-Relay-Token", token.tokenPrefix + "..."); + + return new Response(response.body, { + status: response.status, + headers: newHeaders, + }); + } catch (error) { + const message = error instanceof Error ? error.message : "Unknown error"; + return new Response( + JSON.stringify({ error: { message: `Relay error: ${message}`, type: "relay_error", code: "RELAY_ERR_001" } }), + { status: 500, headers: { ...CORS_HEADERS, "Content-Type": "application/json" } }, + ); + } +} diff --git a/src/lib/db/migrations/067_relay_proxies.sql b/src/lib/db/migrations/067_relay_proxies.sql new file mode 100644 index 0000000000..ee99376160 --- /dev/null +++ b/src/lib/db/migrations/067_relay_proxies.sql @@ -0,0 +1,56 @@ +-- Migration 066: Serverless Relay Proxies +-- Creates tables for relay tokens, rate limits, and usage tracking. + +-- Relay tokens: map external API consumers to internal OmniRoute configuration +CREATE TABLE IF NOT EXISTS relay_tokens ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + token_hash TEXT NOT NULL UNIQUE, -- bcrypt hash of the token + token_prefix TEXT NOT NULL, -- first 8 chars for display (e.g., "rl_abc123") + description TEXT DEFAULT '', + combo_id TEXT, -- optional: restrict to a specific combo + allowed_models TEXT DEFAULT '[]', -- JSON array of model patterns (e.g., ["claude-*", "gpt-*"]) + max_tokens_per_request INTEGER DEFAULT 128000, + max_requests_per_minute INTEGER DEFAULT 60, + max_requests_per_day INTEGER DEFAULT 10000, + max_cost_per_day REAL DEFAULT 0, -- 0 = unlimited + enabled INTEGER DEFAULT 1, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + expires_at INTEGER, -- optional TTL + last_used_at INTEGER, + metadata TEXT DEFAULT '{}' +); + +-- Rate limit window tracking +CREATE TABLE IF NOT EXISTS relay_rate_limits ( + token_id TEXT NOT NULL, + window_start INTEGER NOT NULL, -- unix timestamp of window start + request_count INTEGER DEFAULT 0, + cost REAL DEFAULT 0, + PRIMARY KEY (token_id, window_start), + FOREIGN KEY (token_id) REFERENCES relay_tokens(id) ON DELETE CASCADE +); + +-- Relay request logs +CREATE TABLE IF NOT EXISTS relay_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token_id TEXT NOT NULL, + request_id TEXT, -- X-Request-Id + model TEXT, + prompt_tokens INTEGER DEFAULT 0, + completion_tokens INTEGER DEFAULT 0, + cost REAL DEFAULT 0, + status TEXT DEFAULT 'success', -- success, rate_limited, auth_failed, error + status_code INTEGER DEFAULT 200, + latency_ms INTEGER DEFAULT 0, + client_ip TEXT, + user_agent TEXT, + created_at INTEGER NOT NULL, + FOREIGN KEY (token_id) REFERENCES relay_tokens(id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_relay_logs_token ON relay_logs(token_id, created_at); +CREATE INDEX IF NOT EXISTS idx_relay_logs_created ON relay_logs(created_at); +CREATE INDEX IF NOT EXISTS idx_relay_tokens_prefix ON relay_tokens(token_prefix); +CREATE INDEX IF NOT EXISTS idx_relay_rate_limits_window ON relay_rate_limits(token_id, window_start); diff --git a/src/lib/db/relayProxies.ts b/src/lib/db/relayProxies.ts new file mode 100644 index 0000000000..8e156b856b --- /dev/null +++ b/src/lib/db/relayProxies.ts @@ -0,0 +1,294 @@ +/** + * Relay Proxy DB module + * + * Manages relay tokens, rate limits, and usage tracking for serverless relay proxies. + */ + +import { randomBytes } from "node:crypto"; +import { getDbInstance } from "./core"; +import { rowToCamel } from "./core"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface RelayToken { + id: string; + name: string; + tokenHash: string; + tokenPrefix: string; + description: string; + comboId: string | null; + allowedModels: string; + maxTokensPerRequest: number; + maxRequestsPerMinute: number; + maxRequestsPerDay: number; + maxCostPerDay: number; + enabled: boolean; + createdAt: number; + updatedAt: number; + expiresAt: number | null; + lastUsedAt: number | null; + metadata: string; +} + +export interface RelayTokenRow { + id: string; + name: string; + token_hash: string; + token_prefix: string; + description: string; + combo_id: string | null; + allowed_models: string; + max_tokens_per_request: number; + max_requests_per_minute: number; + max_requests_per_day: number; + max_cost_per_day: number; + enabled: number; + created_at: number; + updated_at: number; + expires_at: number | null; + last_used_at: number | null; + metadata: string; +} + +export interface CreateRelayTokenInput { + name: string; + description?: string; + comboId?: string; + allowedModels?: string[]; + maxTokensPerRequest?: number; + maxRequestsPerMinute?: number; + maxRequestsPerDay?: number; + maxCostPerDay?: number; + expiresAt?: number; + metadata?: Record; +} + +export interface RelayTokenWithSecret extends RelayToken { + rawToken: string; // Only returned once on creation +} + +export interface RelayLogRow { + id: number; + token_id: string; + request_id: string | null; + model: string | null; + prompt_tokens: number; + completion_tokens: number; + cost: number; + status: string; + status_code: number; + latency_ms: number; + client_ip: string | null; + user_agent: string | null; + created_at: number; +} + +// ── Helpers ────────────────────────────────────────────────────────────────── + +function generateId(): string { + return "rl_" + randomBytes(16).toString("hex"); +} + +function generateToken(): string { + return "relay_" + randomBytes(24).toString("hex"); +} + +function hashToken(token: string): string { + // Simple hash for token comparison (not bcrypt-heavy for performance) + const { createHash } = require("node:crypto"); + return createHash("sha256").update(token).digest("hex"); +} + +// ── CRUD ───────────────────────────────────────────────────────────────────── + +export function createRelayToken(input: CreateRelayTokenInput): RelayTokenWithSecret { + const db = getDbInstance(); + const id = generateId(); + const rawToken = generateToken(); + const tokenHash = hashToken(rawToken); + const now = Math.floor(Date.now() / 1000); + + const prefix = "rl_" + rawToken.slice(6, 14); + + db.prepare(` + INSERT INTO relay_tokens (id, name, token_hash, token_prefix, description, combo_id, allowed_models, + max_tokens_per_request, max_requests_per_minute, max_requests_per_day, max_cost_per_day, + enabled, created_at, updated_at, expires_at, metadata) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?) + `).run( + id, input.name, tokenHash, prefix, input.description || "", input.comboId || null, + JSON.stringify(input.allowedModels || ["*"]), + input.maxTokensPerRequest || 128000, + input.maxRequestsPerMinute || 60, + input.maxRequestsPerDay || 10000, + input.maxCostPerDay || 0, + now, now, input.expiresAt || null, + JSON.stringify(input.metadata || {}), + ); + + const token = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(id) as RelayTokenRow; + return { ...rowToCamel(token), rawToken }; +} + +export function getRelayTokens(): RelayToken[] { + const db = getDbInstance(); + const rows = db.prepare("SELECT * FROM relay_tokens ORDER BY created_at DESC").all() as RelayTokenRow[]; + return rows.map((r) => ({ ...rowToCamel(r), enabled: r.enabled === 1 })); +} + +export function getRelayToken(id: string): RelayToken | null { + const db = getDbInstance(); + const row = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(id) as RelayTokenRow | undefined; + if (!row) return null; + return { ...rowToCamel(row), enabled: row.enabled === 1 }; +} + +export function getRelayTokenByHash(tokenHash: string): (RelayToken & { rawToken?: string }) | null { + const db = getDbInstance(); + const row = db.prepare("SELECT * FROM relay_tokens WHERE token_hash = ? AND enabled = 1").get(tokenHash) as RelayTokenRow | undefined; + if (!row) return null; + return { ...rowToCamel(row), enabled: row.enabled === 1 }; +} + +export function updateRelayToken(id: string, updates: Partial): RelayToken | null { + const db = getDbInstance(); + const now = Math.floor(Date.now() / 1000); + const sets: string[] = ["updated_at = ?"]; + const params: unknown[] = [now]; + + if (updates.name !== undefined) { sets.push("name = ?"); params.push(updates.name); } + if (updates.description !== undefined) { sets.push("description = ?"); params.push(updates.description); } + if (updates.comboId !== undefined) { sets.push("combo_id = ?"); params.push(updates.comboId); } + if (updates.allowedModels !== undefined) { sets.push("allowed_models = ?"); params.push(JSON.stringify(updates.allowedModels)); } + if (updates.maxTokensPerRequest !== undefined) { sets.push("max_tokens_per_request = ?"); params.push(updates.maxTokensPerRequest); } + if (updates.maxRequestsPerMinute !== undefined) { sets.push("max_requests_per_minute = ?"); params.push(updates.maxRequestsPerMinute); } + if (updates.maxRequestsPerDay !== undefined) { sets.push("max_requests_per_day = ?"); params.push(updates.maxRequestsPerDay); } + if (updates.maxCostPerDay !== undefined) { sets.push("max_cost_per_day = ?"); params.push(updates.maxCostPerDay); } + + params.push(id); + db.prepare(`UPDATE relay_tokens SET ${sets.join(", ")} WHERE id = ?`).run(...params); + return getRelayToken(id); +} + +export function deleteRelayToken(id: string): void { + const db = getDbInstance(); + db.prepare("DELETE FROM relay_tokens WHERE id = ?").run(id); +} + +export function toggleRelayToken(id: string, enabled: boolean): RelayToken | null { + const db = getDbInstance(); + const now = Math.floor(Date.now() / 1000); + db.prepare("UPDATE relay_tokens SET enabled = ?, updated_at = ? WHERE id = ?").run(enabled ? 1 : 0, now, id); + return getRelayToken(id); +} + +// ── Usage / Rate Limit ─────────────────────────────────────────────────────── + +export function checkRateLimit(tokenId: string): { allowed: boolean; remaining: number; resetIn: number } { + const db = getDbInstance(); + const token = db.prepare("SELECT * FROM relay_tokens WHERE id = ?").get(tokenId) as RelayTokenRow | undefined; + if (!token) return { allowed: false, remaining: 0, resetIn: 0 }; + + const now = Math.floor(Date.now() / 1000); + const minuteWindow = Math.floor(now / 60) * 60; + const dayWindow = Math.floor(now / 86400) * 86400; + + // Check minute rate + const minuteRow = db.prepare( + "SELECT request_count, cost FROM relay_rate_limits WHERE token_id = ? AND window_start = ?", + ).get(tokenId, minuteWindow) as { request_count: number; cost: number } | undefined; + + const minuteCount = minuteRow?.request_count || 0; + if (minuteCount >= token.max_requests_per_minute) { + return { allowed: false, remaining: 0, resetIn: 60 - (now % 60) }; + } + + // Check daily rate + const dayRow = db.prepare( + "SELECT SUM(request_count) as total FROM relay_rate_limits WHERE token_id = ? AND window_start >= ?", + ).get(tokenId, dayWindow) as { total: number } | undefined; + + const dayCount = dayRow?.total || 0; + if (dayCount >= token.max_requests_per_day) { + return { allowed: false, remaining: 0, resetIn: 86400 - (now % 86400) }; + } + + const remaining = Math.min( + token.max_requests_per_minute - minuteCount, + token.max_requests_per_day - dayCount, + ); + + return { allowed: true, remaining, resetIn: 60 - (now % 60) }; +} + +export function recordRelayUsage( + tokenId: string, + params: { + requestId?: string; + model?: string; + promptTokens?: number; + completionTokens?: number; + cost?: number; + status?: string; + statusCode?: number; + latencyMs?: number; + clientIp?: string; + userAgent?: string; + }, +): void { + const db = getDbInstance(); + const now = Math.floor(Date.now() / 1000); + const minuteWindow = Math.floor(now / 60) * 60; + + // Update rate limit window + db.prepare(` + INSERT INTO relay_rate_limits (token_id, window_start, request_count, cost) + VALUES (?, ?, 1, ?) + ON CONFLICT(token_id, window_start) DO UPDATE SET + request_count = request_count + 1, + cost = cost + ? + `).run(tokenId, minuteWindow, params.cost || 0, params.cost || 0); + + // Update last_used_at + db.prepare("UPDATE relay_tokens SET last_used_at = ? WHERE id = ?").run(now, tokenId); + + // Insert log + db.prepare(` + INSERT INTO relay_logs (token_id, request_id, model, prompt_tokens, completion_tokens, cost, + status, status_code, latency_ms, client_ip, user_agent, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `).run( + tokenId, + params.requestId || null, + params.model || null, + params.promptTokens || 0, + params.completionTokens || 0, + params.cost || 0, + params.status || "success", + params.statusCode || 200, + params.latencyMs || 0, + params.clientIp || null, + params.userAgent || null, + now, + ); +} + +export function getRelayUsage(tokenId: string, since: number): { requestCount: number; totalCost: number } { + const db = getDbInstance(); + const row = db.prepare( + "SELECT COUNT(*) as request_count, COALESCE(SUM(cost), 0) as total_cost FROM relay_logs WHERE token_id = ? AND created_at >= ?", + ).get(tokenId, since) as { request_count: number; total_cost: number }; + return { requestCount: row.request_count, totalCost: row.total_cost }; +} + +export function getRelayLogs(tokenId?: string, limit = 50): RelayLogRow[] { + const db = getDbInstance(); + if (tokenId) { + return db.prepare( + "SELECT * FROM relay_logs WHERE token_id = ? ORDER BY created_at DESC LIMIT ?", + ).all(tokenId, limit) as RelayLogRow[]; + } + return db.prepare( + "SELECT * FROM relay_logs ORDER BY created_at DESC LIMIT ?", + ).all(limit) as RelayLogRow[]; +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 23e3bd0254..b5b17b1fe2 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -433,6 +433,7 @@ export { } from "./db/contextHandoffs"; export type { HandoffPayload } from "./db/contextHandoffs"; + export { getAllMiddlewareHooks, getEnabledMiddlewareHooks, @@ -478,6 +479,7 @@ export { getRelayUsage, getRelayLogs, } from "./db/relayProxies"; + export type { RelayToken, RelayTokenRow,