Skip to content

feat(keys): add granular API key limits #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mahdiwafy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(keys): add granular API key limits #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions open-sse/handlers/chatCore.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ export function applyLoopGuard(translatedBody, finalFormat, provider, model, log
* @param {object} options.credentials - Provider credentials
* @param {string} options.sourceFormatOverride - Override detected source format (e.g. "openai-responses")
*/
export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId, userAgent, apiKey, apiKeyName = null, ccFilterNaming, rtkEnabled, headroomEnabled, headroomUrl, headroomCompressUserMessages, cavemanEnabled, cavemanLevel, ponytailEnabled, ponytailLevel, pxpipeEnabled = false, pxpipeMinChars = 1000, pxpipeTimeoutMs = 10000, pxpipeTransform = "png", onPxpipeEvent = null, sourceFormatOverride, providerThinking, clientSignal, loopGuardEnabled = true }) {
export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId, userAgent, apiKey, apiKeyInfo = null, apiKeyName = null, ccFilterNaming, rtkEnabled, headroomEnabled, headroomUrl, headroomCompressUserMessages, cavemanEnabled, cavemanLevel, ponytailEnabled, ponytailLevel, pxpipeEnabled = false, pxpipeMinChars = 1000, pxpipeTimeoutMs = 10000, pxpipeTransform = "png", onPxpipeEvent = null, sourceFormatOverride, providerThinking, clientSignal, loopGuardEnabled = true }) {
const { provider, model, accountCount = 0 } = modelInfo;
const requestStartTime = Date.now();

Expand Down Expand Up @@ -490,7 +490,7 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
return createErrorResult(statusCode, errMsg, resetsAtMs);
}

const sharedCtx = { provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, pxpipe: pxpipeSummary };
const sharedCtx = { provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, pxpipe: pxpipeSummary };
const appendLog = (extra) => appendRequestLog({ model, provider, connectionId, ...extra }).catch(() => { });
const trackDone = () => trackPendingRequest(model, provider, connectionId, false);

Expand Down
4 changes: 2 additions & 2 deletions open-sse/handlers/chatCore/nonStreamingHandler.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ export function translateNonStreamingResponse(responseBody, targetFormat, source
/**
* Handle non-streaming response from provider.
*/
export async function handleNonStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, trackDone, appendLog, pxpipe, comboName }) {
export async function handleNonStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, trackDone, appendLog, pxpipe, comboName }) {
trackDone();
const contentType = providerResponse.headers.get("content-type") || "";
let responseBody;
Expand Down Expand Up @@ -264,7 +264,7 @@ export async function handleNonStreamingResponse({ providerResponse, provider, m

const usage = extractUsageFromResponse(responseBody);
appendLog({ tokens: usage, status: "200 OK" });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, comboName });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, comboName });

const translatedResponse = needsTranslation(targetFormat, sourceFormat)
? translateNonStreamingResponse(responseBody, targetFormat, sourceFormat)
Expand Down
9 changes: 8 additions & 1 deletion open-sse/handlers/chatCore/requestDetail.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { saveRequestUsage, appendRequestLog, saveRequestDetail } from "@/lib/usageDb.js";
import { recordApiKeyUsage } from "@/lib/db/repos/apiKeyUsageRepo.js";
import { COLORS } from "../../utils/stream.js";
import { canonicalizeUsage } from "../../utils/usageTracking.js";

Expand Down Expand Up @@ -95,7 +96,7 @@ export function buildRequestDetail(base, overrides = {}) {
};
}

export function saveUsageStats({ provider, model, tokens, connectionId, apiKey, endpoint, label = "USAGE" }) {
export function saveUsageStats({ provider, model, tokens, connectionId, apiKey, apiKeyInfo, endpoint, label = "USAGE" }) {
if (!tokens || typeof tokens !== "object") return;

const inTokens = tokens.input_tokens ?? tokens.prompt_tokens ?? 0;
Expand Down Expand Up @@ -123,4 +124,10 @@ export function saveUsageStats({ provider, model, tokens, connectionId, apiKey,
apiKey: apiKey || undefined,
endpoint: endpoint || null
}).catch(() => {});

// Record per-key usage limits against the API key info (if provided).
if (apiKeyInfo) {
const totalTokens = (normalized.prompt_tokens || 0) + (normalized.completion_tokens || 0);
recordApiKeyUsage(apiKeyInfo, totalTokens);
}
}
6 changes: 3 additions & 3 deletions open-sse/handlers/chatCore/sseToJsonHandler.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ export function parseSSEToOpenAIResponse(rawSSE, fallbackModel, validToolNames =
* Handle case: provider forced streaming but client wants JSON.
* Supports both Codex/Responses API SSE and standard Chat Completions SSE.
*/
export async function handleForcedSSEToJson({ providerResponse, sourceFormat, provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, trackDone, appendLog, comboName, toolNameMap }) {
export async function handleForcedSSEToJson({ providerResponse, sourceFormat, provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, trackDone, appendLog, comboName, toolNameMap }) {
const contentType = providerResponse.headers.get("content-type") || "";
const isSSE = contentType.includes("text/event-stream") || (contentType === "" && isResponsesProvider(provider));
if (!isSSE) return null; // not handled here
Expand All @@ -147,7 +147,7 @@ export async function handleForcedSSEToJson({ providerResponse, sourceFormat, pr

const usage = jsonResponse.usage || {};
appendLog({ tokens: usage, status: "200 OK" });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, comboName });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, comboName });

const { msgItem, textContent } = pickAssistantMessageForChatCompletion(jsonResponse.output);
const totalLatency = Date.now() - requestStartTime;
Expand Down Expand Up @@ -244,7 +244,7 @@ export async function handleForcedSSEToJson({ providerResponse, sourceFormat, pr

const usage = parsed.usage || {};
appendLog({ tokens: usage, status: "200 OK" });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint });

const totalLatency = Date.now() - requestStartTime;
saveRequestDetail(buildRequestDetail({
Expand Down
11 changes: 8 additions & 3 deletions open-sse/handlers/chatCore/streamingHandler.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,12 @@ function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent,
* Includes a readiness gate: if upstream closes before any byte arrives,
* return STREAM_EARLY_EOF so the caller can retry once on the same connection.
*/
export async function handleStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, userAgent, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, streamController, onStreamComplete, streamDetailId, pxpipe }) {
export async function handleStreamingResponse({
providerResponse, provider, model, sourceFormat, targetFormat, userAgent,
body, stream, translatedBody, finalBody, requestStartTime, connectionId,
apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess,
reqLogger, toolNameMap, streamController, onStreamComplete, streamDetailId, pxpipe,
}) {
if (onRequestSuccess) {
Promise.resolve()
.then(onRequestSuccess)
Expand Down Expand Up @@ -180,7 +185,7 @@ export async function handleStreamingResponse({ providerResponse, provider, mode
/**
* Build onStreamComplete callback for streaming usage tracking.
*/
export function buildOnStreamComplete({ provider, model, connectionId, apiKey, apiKeyName, requestStartTime, body, stream, finalBody, translatedBody, clientRawRequest, pxpipe }) {
export function buildOnStreamComplete({ provider, model, connectionId, apiKey, apiKeyInfo, apiKeyName, requestStartTime, body, stream, finalBody, translatedBody, clientRawRequest, pxpipe }) {
const streamDetailId = `${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;

const onStreamComplete = (contentObj, usage, ttftAt) => {
Expand All @@ -205,7 +210,7 @@ export function buildOnStreamComplete({ provider, model, connectionId, apiKey, a
console.error("[RequestDetail] Failed to update streaming content:", err.message);
});

saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, label: "STREAM USAGE" });
saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, label: "STREAM USAGE" });
};

return { onStreamComplete, streamDetailId };
Expand Down
96 changes: 95 additions & 1 deletion src/app/(dashboard)/dashboard/endpoint/EndpointPageClient.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,16 @@ export default function APIPageClient({ machineId }) {
const [editProvidersAll, setEditProvidersAll] = useState(true);
const [editCombosAll, setEditCombosAll] = useState(true);
const [editSaving, setEditSaving] = useState(false);
// Limit edit state (empty string = unlimited)
const [editExpiresAt, setEditExpiresAt] = useState("");
const [editMaxTokens, setEditMaxTokens] = useState("");
const [editMaxTokensDaily, setEditMaxTokensDaily] = useState("");
const [editRpm, setEditRpm] = useState("");
const [editRph, setEditRph] = useState("");
const [editRpd, setEditRpd] = useState("");
const [editTokens5h, setEditTokens5h] = useState("");
const [editTokensWeekly, setEditTokensWeekly] = useState("");
const [editTokensMonthly, setEditTokensMonthly] = useState("");
const [providerList, setProviderList] = useState([]);
const [aliasMap, setAliasMap] = useState({}); // alias → provider ID
const [comboList, setComboList] = useState([]);
Expand Down Expand Up @@ -478,6 +488,15 @@ export default function APIPageClient({ machineId }) {
const handleOpenEditKey = (key) => {
setEditingKey(key);
setEditName(key.name || "");
setEditExpiresAt(key.expiresAt ? key.expiresAt.slice(0, 16) : "");
setEditMaxTokens(key.maxTokens ?? "");
setEditMaxTokensDaily(key.maxTokensDaily ?? "");
setEditRpm(key.rpm ?? "");
setEditRph(key.rph ?? "");
setEditRpd(key.rpd ?? "");
setEditTokens5h(key.tokens5h ?? "");
setEditTokensWeekly(key.tokensWeekly ?? "");
setEditTokensMonthly(key.tokensMonthly ?? "");
const ap = key.allowedProviders;
const ac = key.allowedCombos;
const ak = key.allowedKinds;
Expand Down Expand Up @@ -514,10 +533,28 @@ export default function APIPageClient({ machineId }) {
if (!editingKey) return;
setEditSaving(true);
try {
const parseNum = (v) => {
const s = String(v).trim();
if (s === "" || s === "null" || s === "undefined") return null;
const n = Number(s);
return Number.isFinite(n) && n >= 0 ? n : null;
};
const body = {
name: editName.trim() || editingKey.name,
allowedProviders: editProvidersAll ? null : editProviders,
allowedCombos: editCombosAll ? null : editCombos,
allowedKinds: editKindsAll ? null : editKinds,
limits: {
expiresAt: editExpiresAt ? new Date(editExpiresAt).toISOString() : null,
maxTokens: parseNum(editMaxTokens),
maxTokensDaily: parseNum(editMaxTokensDaily),
rpm: parseNum(editRpm),
rph: parseNum(editRph),
rpd: parseNum(editRpd),
tokens5h: parseNum(editTokens5h),
tokensWeekly: parseNum(editTokensWeekly),
tokensMonthly: parseNum(editTokensMonthly),
},
};
const res = await fetch(`/api/keys/${editingKey.id}`, {
method: "PUT",
Expand Down Expand Up @@ -1542,6 +1579,11 @@ export default function APIPageClient({ machineId }) {
</p>
{key.isActive === false && (
<p className="text-xs text-orange-500 mt-1">Paused</p>
)}
{key.expiresAt && (
<p className="text-xs text-text-muted mt-1">
Expires {new Date(key.expiresAt).toLocaleString()}
</p>
)}
{/* ACL badges */}
{(key.allowedProviders || key.allowedCombos || key.allowedKinds) && (
Expand Down Expand Up @@ -1571,6 +1613,14 @@ export default function APIPageClient({ machineId }) {
{key.allowedKinds.length === 0 ? "No kinds" : key.allowedKinds.join(", ")}
</span>
)}
{key.rpm != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-orange-500/10 text-orange-500 dark:bg-orange-500/20">{key.rpm} RPM</span>}
{key.rph != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-orange-500/10 text-orange-500 dark:bg-orange-500/20">{key.rph} RPH</span>}
{key.rpd != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-orange-500/10 text-orange-500 dark:bg-orange-500/20">{key.rpd} RPD</span>}
{key.maxTokens != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-cyan-500/10 text-cyan-500 dark:bg-cyan-500/20">max {key.maxTokens.toLocaleString()} tokens</span>}
{key.maxTokensDaily != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-cyan-500/10 text-cyan-500 dark:bg-cyan-500/20">{key.maxTokensDaily.toLocaleString()} tokens/day</span>}
{key.tokens5h != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-cyan-500/10 text-cyan-500 dark:bg-cyan-500/20">{key.tokens5h.toLocaleString()} tokens/5h</span>}
{key.tokensWeekly != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-cyan-500/10 text-cyan-500 dark:bg-cyan-500/20">{key.tokensWeekly.toLocaleString()} tokens/week</span>}
{key.tokensMonthly != null && <span className="text-[10px] px-1.5 py-0.5 rounded bg-cyan-500/10 text-cyan-500 dark:bg-cyan-500/20">{key.tokensMonthly.toLocaleString()} tokens/month</span>}
</div>
)}
</div>
Expand Down Expand Up @@ -1785,13 +1835,57 @@ export default function APIPageClient({ machineId }) {
{editCombosAll && <p className="text-xs text-text-muted">This key can access all combos.</p>}
</div>

{/* Limits */}
<div>
<label className="text-sm font-medium mb-2 block">Usage Limits (leave blank for unlimited)</label>
<div className="grid grid-cols-2 gap-3">
<div>
<label className="text-xs text-text-muted">Expiration</label>
<Input type="datetime-local" value={editExpiresAt} onChange={(e) => setEditExpiresAt(e.target.value)} />
</div>
<div>
<label className="text-xs text-text-muted">Max tokens / request</label>
<Input type="number" min="0" value={editMaxTokens} onChange={(e) => setEditMaxTokens(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Max tokens / day</label>
<Input type="number" min="0" value={editMaxTokensDaily} onChange={(e) => setEditMaxTokensDaily(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Requests / minute</label>
<Input type="number" min="0" value={editRpm} onChange={(e) => setEditRpm(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Requests / hour</label>
<Input type="number" min="0" value={editRph} onChange={(e) => setEditRph(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Requests / day</label>
<Input type="number" min="0" value={editRpd} onChange={(e) => setEditRpd(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Tokens / 5 hours</label>
<Input type="number" min="0" value={editTokens5h} onChange={(e) => setEditTokens5h(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Tokens / week</label>
<Input type="number" min="0" value={editTokensWeekly} onChange={(e) => setEditTokensWeekly(e.target.value)} placeholder="∞" />
</div>
<div>
<label className="text-xs text-text-muted">Tokens / month</label>
<Input type="number" min="0" value={editTokensMonthly} onChange={(e) => setEditTokensMonthly(e.target.value)} placeholder="∞" />
</div>
</div>
</div>

{/* ACL info */}
<div className="text-xs text-text-muted bg-surface-2 rounded-lg p-3 border border-border-subtle">
<p className="font-medium mb-1">How it works:</p>
<ul className="list-disc pl-4 space-y-0.5">
<li><strong>All allowed</strong> = unrestricted (default). No ACL filtering.</li>
<li><strong>Unchecked + empty</strong> = deny everything of that type.</li>
<li><strong>Checked items</strong> = only those items are accessible.</li>
<li><strong>Usage limits</strong> = leave blank to disable. Applies to requests made with this key.</li>
</ul>
</div>

Expand All @@ -1801,7 +1895,7 @@ export default function APIPageClient({ machineId }) {
Cancel
</Button>
<Button onClick={handleSaveEditKey} fullWidth disabled={editSaving}>
{editSaving ? "Saving..." : "Save ACL"}
{editSaving ? "Saving..." : "Save"}
</Button>
</div>
</div>
Expand Down
Loading
Loading