Files
OmniRoute/tests/unit/internal-usage-command.test.ts
Witroch4 70d89d2f68 feat(keys): add per-key USD usage quota controls (#4327)
* feat(keys): add per-key USD usage quotas

Adds daily and weekly API key USD caps with reset-aware weekly windows, exposes quota controls in API key permissions and costs views, and returns Claude Code-safe 400 responses when caps are exceeded.

Validations:

- node --import tsx/esm --test tests/unit/api-key-usage-limits.test.ts tests/unit/internal-usage-command.test.ts

- npm run typecheck:core

- npm run check:file-size

- npm run check:migration-numbering

- npm run lint

- Docker image build/deploy smoke test on 100.64.0.1:20128

* fix(db): renumber api_key_usage_limits migration 100->101 (avoid cli_access_tokens collision)

Migration version 100 is taken by 100_cli_access_tokens.sql on release; the
migrationRunner version-collision guard would otherwise skip one. Renumber to 101.

Co-authored-by: diegosouzapw <diegosouza.pw@gmail.com>

* chore(quality): bump apiKeys.ts file-size baseline 1661->1662 (USD quota fields)

Co-authored-by: diegosouzapw <diegosouza.pw@gmail.com>

---------

Co-authored-by: Wital <wital@example.com>
Co-authored-by: diegosouzapw <diegosouza.pw@gmail.com>
2026-06-19 22:42:54 -03:00

296 lines
8.5 KiB
TypeScript

import test from "node:test";
import assert from "node:assert/strict";
import {
buildUsageCommandText,
extractLastUserText,
handleInternalUsageCommand,
isInternalUsageCommand,
} from "../../src/lib/usage/internalUsageCommand.ts";
const NOW = Date.parse("2026-06-16T12:00:00.000Z");
test("internal usage command only matches the exact trimmed user message", () => {
assert.equal(isInternalUsageCommand("@@om-usage"), true);
assert.equal(isInternalUsageCommand(" @@om-usage "), true);
assert.equal(isInternalUsageCommand("me mostra @@om-usage"), false);
assert.equal(isInternalUsageCommand("@@om-usage agora"), false);
assert.equal(isInternalUsageCommand("/@@om-usage"), false);
assert.equal(isInternalUsageCommand("@@om-usage."), false);
assert.equal(isInternalUsageCommand("@@om-usage\nabc"), false);
assert.equal(isInternalUsageCommand("```@@om-usage```"), false);
assert.equal(isInternalUsageCommand(null), false);
});
test("extractLastUserText supports OpenAI and Anthropic text content", () => {
assert.equal(
extractLastUserText({
messages: [
{ role: "user", content: "first" },
{ role: "assistant", content: "middle" },
{ role: "user", content: [{ type: "text", text: "@@om-usage" }] },
],
}),
"@@om-usage"
);
assert.equal(
extractLastUserText({
input: [
{ role: "assistant", content: "ignored" },
{ role: "user", content: [{ type: "input_text", text: "hello" }] },
],
}),
"hello"
);
});
test("buildUsageCommandText formats cached Claude usage windows exactly", async () => {
const text = await buildUsageCommandText(
{
id: "key-1",
name: "main",
allowedConnections: ["conn-claude"],
},
{
now: () => NOW,
getProviderConnectionById: async () => ({
id: "conn-claude",
provider: "claude",
isActive: true,
}),
getProviderConnections: async () => [],
getProviderLimitsCache: () => ({
plan: "Claude Max",
quotas: {
"session (5h)": {
used: 53,
total: 100,
remaining: 47,
resetAt: new Date(NOW + 9 * 60_000).toISOString(),
},
"weekly (7d)": {
used: 72,
total: 100,
remaining: 28,
resetAt: new Date(NOW + 24 * 60 * 60_000).toISOString(),
},
"weekly sonnet (7d)": {
used: 30,
total: 100,
remaining: 70,
resetAt: new Date(NOW + 24 * 60 * 60_000).toISOString(),
},
},
message: null,
fetchedAt: new Date(NOW).toISOString(),
}),
getAllProviderLimitsCache: () => ({}),
isValidApiKey: async () => true,
getApiKeyMetadata: async () => null,
}
);
assert.equal(
text,
[
"Plan",
"Claude Max",
"",
"Usage",
"Session (5hr)",
"53%",
"Resets in 9m",
"",
"Weekly (7 day)",
"72%",
"Resets in 1d",
"",
"Weekly Sonnet",
"30%",
"Resets in 1d",
].join("\n")
);
});
test("buildUsageCommandText formats API key USD limits when fair usage is enabled", async () => {
const text = await buildUsageCommandText(
{
id: "key-limited",
name: "limited",
allowedConnections: ["conn-claude"],
usageLimitEnabled: true,
dailyUsageLimitUsd: 10,
weeklyUsageLimitUsd: 50,
},
{
now: () => NOW,
getApiKeyUsageLimitStatus: async () => ({
enabled: true,
dailyLimitUsd: 10,
weeklyLimitUsd: 50,
dailySpentUsd: 2,
weeklySpentUsd: 5.25,
dailyWindowStartIso: "2026-06-16T03:00:00.000Z",
weeklyWindowStartIso: "2026-06-09T12:00:00.000Z",
weeklyResetAtIso: "2026-06-16T12:00:00.000Z",
dailyExceeded: false,
weeklyExceeded: false,
}),
getProviderConnectionById: async () => {
throw new Error("provider connection lookup must not run for fair usage output");
},
getProviderConnections: async () => {
throw new Error("provider connection lookup must not run for fair usage output");
},
getProviderLimitsCache: () => null,
getAllProviderLimitsCache: () => {
throw new Error("provider cache lookup must not run for fair usage output");
},
isValidApiKey: async () => true,
getApiKeyMetadata: async () => null,
}
);
assert.equal(
text,
[
"Cota diaria",
"$10.00",
"Gasto diario",
"$2.00",
"",
"Cota semanal",
"$50.00",
"Gasto semanal",
"$5.25",
].join("\n")
);
});
test("handleInternalUsageCommand returns disabled response locally without provider routing", async () => {
const response = await handleInternalUsageCommand(
new Request("http://localhost/v1/chat/completions", {
method: "POST",
headers: { Authorization: "Bearer sk-disabled" },
}),
{
model: "claude-opus-4-8",
messages: [{ role: "user", content: "@@om-usage" }],
},
{
isValidApiKey: async () => true,
getApiKeyMetadata: async () => ({
id: "key-disabled",
name: "disabled",
allowedConnections: [],
allowUsageCommand: false,
}),
now: () => NOW,
getProviderConnectionById: async () => null,
getProviderConnections: async () => {
throw new Error("provider connection lookup must not run when disabled");
},
getProviderLimitsCache: () => null,
getAllProviderLimitsCache: () => {
throw new Error("provider cache lookup must not run when disabled");
},
}
);
assert.ok(response, "command should be handled locally");
assert.equal(response.status, 200);
const body = (await response.json()) as {
choices: Array<{ message: { content: string } }>;
};
assert.equal(body.choices[0].message.content, "Usage command is disabled for this API key.");
});
test("handleInternalUsageCommand returns enabled usage snapshot locally", async () => {
const response = await handleInternalUsageCommand(
new Request("http://localhost/v1/messages", {
method: "POST",
headers: {
"anthropic-version": "2023-06-01",
"x-api-key": "sk-enabled",
},
}),
{
model: "claude-opus-4-8",
messages: [{ role: "user", content: " @@om-usage " }],
},
{
isValidApiKey: async () => true,
getApiKeyMetadata: async () => ({
id: "key-enabled",
name: "enabled",
allowedConnections: ["conn-claude"],
allowUsageCommand: true,
}),
now: () => NOW,
getProviderConnectionById: async () => ({
id: "conn-claude",
provider: "claude",
isActive: true,
}),
getProviderConnections: async () => [],
getProviderLimitsCache: () => ({
plan: "Claude Max",
quotas: {
"session (5h)": {
used: 53,
total: 100,
resetAt: new Date(NOW + 9 * 60_000).toISOString(),
},
"weekly (7d)": {
used: 72,
total: 100,
resetAt: new Date(NOW + 24 * 60 * 60_000).toISOString(),
},
"weekly sonnet (7d)": {
used: 30,
total: 100,
resetAt: new Date(NOW + 24 * 60 * 60_000).toISOString(),
},
},
message: null,
fetchedAt: new Date(NOW).toISOString(),
}),
getAllProviderLimitsCache: () => ({}),
}
);
assert.ok(response, "command should be handled locally");
assert.equal(response.status, 200);
const body = (await response.json()) as {
content: Array<{ type: string; text: string }>;
};
assert.equal(body.content[0].text.includes("Weekly Sonnet\n30%\nResets in 1d"), true);
});
test("handleInternalUsageCommand ignores normal prompts", async () => {
const response = await handleInternalUsageCommand(
new Request("http://localhost/v1/chat/completions", {
method: "POST",
headers: { Authorization: "Bearer sk-enabled" },
}),
{
model: "claude-opus-4-8",
messages: [{ role: "user", content: "me mostra @@om-usage" }],
},
{
isValidApiKey: async () => {
throw new Error("auth must not run for non-exact prompts");
},
getApiKeyMetadata: async () => null,
now: () => NOW,
getProviderConnectionById: async () => null,
getProviderConnections: async () => [],
getProviderLimitsCache: () => null,
getAllProviderLimitsCache: () => ({}),
}
);
assert.equal(response, null);
});