From 58d317d223278e99ceb8ac41cb7a0ecb7f32ed02 Mon Sep 17 00:00:00 2001 From: Witroch4 Date: Wed, 17 Jun 2026 00:07:12 -0300 Subject: [PATCH] [codex] add per-key local usage command (#4034) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrated into release/v3.8.27 — per-key local @@om-usage command (cached quota, no upstream routing). Rebased onto modularized schemas/keys.ts + file-size rebaseline. Thanks @Witroch4! --- config/quality/file-size-baseline.json | 5 +- .../api-manager/ApiManagerPageClient.tsx | 56 ++ src/app/api/keys/[id]/route.ts | 3 + src/app/api/keys/route.ts | 10 +- src/i18n/messages/en.json | 3 + src/i18n/messages/pt-BR.json | 3 + src/lib/db/apiKeys.ts | 32 +- src/lib/usage/internalUsageCommand.ts | 547 ++++++++++++++++++ src/shared/utils/apiKeyPolicy.ts | 1 + src/shared/validation/schemas/keys.ts | 6 +- src/sse/handlers/chat.ts | 7 + tests/unit/api-manager-usage-command.test.ts | 33 ++ tests/unit/apikeys-usage-command.test.ts | 56 ++ tests/unit/internal-usage-command.test.ts | 240 ++++++++ 14 files changed, 994 insertions(+), 8 deletions(-) create mode 100644 src/lib/usage/internalUsageCommand.ts create mode 100644 tests/unit/api-manager-usage-command.test.ts create mode 100644 tests/unit/apikeys-usage-command.test.ts create mode 100644 tests/unit/internal-usage-command.test.ts diff --git a/config/quality/file-size-baseline.json b/config/quality/file-size-baseline.json index 29a012d522..471cb1b79c 100644 --- a/config/quality/file-size-baseline.json +++ b/config/quality/file-size-baseline.json @@ -24,6 +24,7 @@ "_rebaseline_2026_06_15_3938_perplexity_v218": "PR #3938 own growth: perplexity-web.ts 868->939 (+71 = rebuild buildPplxRequestBody to mirror the current www.perplexity.ai schematized request body — version 2.18, use_schematized_api + the full supported_block_use_cases list, dsl_query, shared requestId for frontend_uuid/client_search_results_cache_key, last_backend_uuid only on follow-ups — plus the x-perplexity-request-* / x-request-id headers replacing the stale X-App-ApiVersion pair that triggered HTTP 400). Cohesive upstream-schema sync in a single executor; not extractable.", "_rebaseline_2026_06_16_4021_context_editing": "PR #4021 own growth: base.ts 1222->1244 (+22 = inject delegated Context Editing at the single Claude pre-serialization chokepoint — applyContextEditingToBody() call gated to the genuine `claude` provider, the contextEditing field on ExecuteInput, its destructure, and a debug log) and chatCore.ts 5868->5875 (+7 = capture contextEditing.enabled at the canonical compression-settings read into a function-scoped flag, threaded to the two executor.execute() callsites). The reusable edit-builder + strategy constants live in the new small open-sse/config/contextEditing.ts (well under cap). Cohesive opt-in feature at the existing dispatch chokepoint; not extractable without hiding the Claude body-finalization boundary.", "_rebaseline_2026_06_16_4033_compression_token_saver_ui": "PR #4033 net -46 LOC overall: CompressionSettingsTab.tsx 932->974 (+42 over frozen after relocating Token Saver from Endpoint/Appearance into Compression Settings). This is a deliberate UI ownership redistribution, not aggregate code growth; keeping the control with the compression settings makes the module boundary clearer.", + "_rebaseline_2026_06_17_4034_usage_command": "PR #4034 own growth: apiKeys.ts 1633->1661 (+28 = the allow_usage_command additive column — fallback definition, parseAllowUsageCommand, prepared-statement SELECT column, and the metadata/create/update wiring, mirroring the existing disable_non_public_models accessor pattern) and chat.ts 1425->1432 (+7 = the handleInternalUsageCommand intercept hook at the existing post-auth chokepoint in handleChat). The reusable command logic lives in the new src/lib/usage/internalUsageCommand.ts (well under cap). Cohesive opt-in feature at the established column/dispatch boundaries; not extractable without splitting the api-keys domain module.", "cap": 800, "frozen": { "open-sse/config/providerRegistry.ts": 4731, @@ -101,7 +102,7 @@ "src/app/api/usage/analytics/route.ts": 941, "src/app/api/v1/models/catalog.ts": 1435, "src/lib/cloudflaredTunnel.ts": 934, - "src/lib/db/apiKeys.ts": 1633, + "src/lib/db/apiKeys.ts": 1661, "src/lib/db/core.ts": 1820, "src/lib/db/migrationRunner.ts": 1125, "src/lib/db/models.ts": 1180, @@ -126,7 +127,7 @@ "src/shared/constants/sidebarVisibility.ts": 1100, "src/shared/services/cliRuntime.ts": 1090, "src/shared/validation/schemas.ts": 2523, - "src/sse/handlers/chat.ts": 1425, + "src/sse/handlers/chat.ts": 1432, "src/sse/services/auth.ts": 2219 }, "_rebaseline_2026_06_09": "Re-baseline consciente pre-release v3.8.19: 9 arquivos cresceram durante o ciclo (features mergeadas: RequestLoggerV2 +281 request-logger rework, stream +101, combo +73, chatCore +45, catalog +32 fable-5/catalog-flag, callLogs +4, accountFallback +2, usageHistory novo 840) + core.ts +7 (fix resetAllDbModuleState, PR 3536). A catraca segue valendo destes valores — proximo crescimento falha. Decisao: encolher (esp. RequestLoggerV2/chatCore) e a issue #3501 ficam para o ciclo seguinte.", diff --git a/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx b/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx index b75e7ede89..c182ff9c84 100644 --- a/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx +++ b/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx @@ -128,6 +128,7 @@ interface ApiKey { allowedEndpoints?: string[]; streamDefaultMode?: StreamDefaultMode; disableNonPublicModels?: boolean; + allowUsageCommand?: boolean; allowedQuotas?: string[] | null; createdAt: string; } @@ -221,6 +222,7 @@ export default function ApiManagerPageClient() { const [newKeyManageEnabled, setNewKeyManageEnabled] = useState(false); const [newKeySelfUsageEnabled, setNewKeySelfUsageEnabled] = useState(true); const [newKeyAccountQuotaEnabled, setNewKeyAccountQuotaEnabled] = useState(false); + const [newKeyAllowUsageCommand, setNewKeyAllowUsageCommand] = useState(false); const [createdKey, setCreatedKey] = useState(null); const [editingKey, setEditingKey] = useState(null); const [showPermissionsModal, setShowPermissionsModal] = useState(false); @@ -542,6 +544,7 @@ export default function ApiManagerPageClient() { selfUsageEnabled: newKeySelfUsageEnabled, selfAccountQuotaEnabled: newKeyAccountQuotaEnabled, }), + allowUsageCommand: newKeyAllowUsageCommand, }), }); const data = await res.json(); @@ -553,6 +556,7 @@ export default function ApiManagerPageClient() { setNewKeyManageEnabled(false); setNewKeySelfUsageEnabled(true); setNewKeyAccountQuotaEnabled(false); + setNewKeyAllowUsageCommand(false); setShowAddModal(false); } else { setCreateError(data.error || t("failedCreateKey")); @@ -659,6 +663,7 @@ export default function ApiManagerPageClient() { allowedEndpoints: string[], streamDefaultMode: StreamDefaultMode, disableNonPublicModels: boolean, + allowUsageCommand: boolean, blockedModels: string[] ) => { if (!editingKey || !editingKey.id) return; @@ -726,6 +731,7 @@ export default function ApiManagerPageClient() { allowedEndpoints, streamDefaultMode, disableNonPublicModels, + allowUsageCommand, }), }); @@ -908,6 +914,7 @@ export default function ApiManagerPageClient() { const hasThrottle = throttleDelayMs > 0; const hasManageScope = Array.isArray(key.scopes) && key.scopes.includes("manage"); const hasJsonStreamDefault = key.streamDefaultMode === "json"; + const hasLocalUsageCommand = key.allowUsageCommand === true; const maxSessions = typeof key.maxSessions === "number" ? key.maxSessions : 0; const hasSessionLimit = maxSessions > 0; const activeSessions = sessionCounts[key.id] || 0; @@ -1033,6 +1040,12 @@ export default function ApiManagerPageClient() { {t("streamDefaultBadge")} )} + {hasLocalUsageCommand && ( + + terminal + {t("localUsageCommandBadge")} + + )} {hasSessionLimit && ( group @@ -1208,6 +1221,7 @@ export default function ApiManagerPageClient() { setNewKeyManageEnabled(false); setNewKeySelfUsageEnabled(true); setNewKeyAccountQuotaEnabled(false); + setNewKeyAllowUsageCommand(false); setNameError(null); setCreateError(null); }} @@ -1302,6 +1316,26 @@ export default function ApiManagerPageClient() { {newKeyAccountQuotaEnabled ? tc("enabled") : tc("disabled")} +
+
+

{t("localUsageCommand")}

+

{t("localUsageCommandDesc")}

+
+ +
{createError && (
@@ -1317,6 +1351,7 @@ export default function ApiManagerPageClient() { setNewKeyManageEnabled(false); setNewKeySelfUsageEnabled(true); setNewKeyAccountQuotaEnabled(false); + setNewKeyAllowUsageCommand(false); setNameError(null); setCreateError(null); }} @@ -1433,6 +1468,7 @@ const PermissionsModal = memo(function PermissionsModal({ allowedEndpoints: string[], streamDefaultMode: StreamDefaultMode, disableNonPublicModels: boolean, + allowUsageCommand: boolean, blockedModels: string[] ) => void; }) { @@ -1513,6 +1549,9 @@ const PermissionsModal = memo(function PermissionsModal({ const [disableNonPublicModels, setDisableNonPublicModels] = useState( apiKey?.disableNonPublicModels === true ); + const [usageCommandEnabled, setUsageCommandEnabled] = useState( + apiKey?.allowUsageCommand === true + ); const getModelDisplayName = useCallback( (modelId: string) => modelId === CLAUDE_CODE_DEFAULT_MODEL_ID ? CLAUDE_CODE_DEFAULT_MODEL_NAME : modelId, @@ -1696,6 +1735,7 @@ const PermissionsModal = memo(function PermissionsModal({ allowAllEndpoints ? [] : selectedEndpoints, streamDefaultMode, disableNonPublicModels, + usageCommandEnabled, blockedModels ); }, [ @@ -1727,6 +1767,7 @@ const PermissionsModal = memo(function PermissionsModal({ selectedEndpoints, streamDefaultMode, disableNonPublicModels, + usageCommandEnabled, blockedClaudeCodeFamilies, initialBlockedModels, apiKey?.scopes, @@ -2285,6 +2326,21 @@ const PermissionsModal = memo(function PermissionsModal({ {selfAccountQuotaEnabled ? tc("enabled") : tc("disabled")}

{t("sharedAccountQuotaVisibilityDesc")}

+ +

{t("localUsageCommandDesc")}

{/* Disable Non-Public Models Toggle */} diff --git a/src/app/api/keys/[id]/route.ts b/src/app/api/keys/[id]/route.ts index 08a6940751..d4161f59b6 100644 --- a/src/app/api/keys/[id]/route.ts +++ b/src/app/api/keys/[id]/route.ts @@ -82,6 +82,7 @@ export async function PATCH(request, { params }) { allowedEndpoints, streamDefaultMode, disableNonPublicModels, + allowUsageCommand, } = validation.data; const payload: Parameters[1] = {}; @@ -104,6 +105,7 @@ export async function PATCH(request, { params }) { if (streamDefaultMode !== undefined) payload.streamDefaultMode = streamDefaultMode; if (disableNonPublicModels !== undefined) payload.disableNonPublicModels = disableNonPublicModels; + if (allowUsageCommand !== undefined) payload.allowUsageCommand = allowUsageCommand; const updated = await updateApiKeyPermissions(id, payload); if (!updated) { @@ -133,6 +135,7 @@ export async function PATCH(request, { params }) { ...(allowedEndpoints !== undefined && { allowedEndpoints }), ...(streamDefaultMode !== undefined && { streamDefaultMode }), ...(disableNonPublicModels !== undefined && { disableNonPublicModels }), + ...(allowUsageCommand !== undefined && { allowUsageCommand }), }); } catch (error) { log.error("keys", "Error updating key permissions", error); diff --git a/src/app/api/keys/route.ts b/src/app/api/keys/route.ts index 0c69f50be3..32e2954f16 100644 --- a/src/app/api/keys/route.ts +++ b/src/app/api/keys/route.ts @@ -63,14 +63,17 @@ export async function POST(request) { if (isValidationFailure(validation)) { return NextResponse.json({ error: validation.error }, { status: 400 }); } - const { name, noLog, scopes } = validation.data; + const { name, noLog, scopes, allowUsageCommand } = validation.data; // Always get machineId from server const machineId = await getConsistentMachineId(); const normalizedScopes = normalizeSelfServiceScopesForCreate(scopes); const apiKey = await createApiKey(name, machineId, normalizedScopes); - if (noLog === true) { - await updateApiKeyPermissions(apiKey.id, { noLog: true }); + if (noLog === true || allowUsageCommand === true) { + await updateApiKeyPermissions(apiKey.id, { + ...(noLog === true && { noLog: true }), + ...(allowUsageCommand === true && { allowUsageCommand: true }), + }); } // Auto sync to Cloud if enabled @@ -83,6 +86,7 @@ export async function POST(request) { id: apiKey.id, machineId: apiKey.machineId, noLog: noLog === true, + allowUsageCommand: allowUsageCommand === true, streamDefaultMode: "legacy", }, { status: 201 } diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 9318b1eedf..3101992814 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -1693,6 +1693,9 @@ "ownUsageVisibilityDesc": "Allow this key to call its status endpoint for its own USD usage, budget percent, and token totals.", "sharedAccountQuotaVisibility": "Shared Account Quota", "sharedAccountQuotaVisibilityDesc": "Allow this key to see shared upstream account quota when one explicit connection is configured.", + "localUsageCommand": "Allow local usage command", + "localUsageCommandDesc": "Allows this API key to use @@om-usage to retrieve cached usage and quota information without calling an upstream provider.", + "localUsageCommandBadge": "Usage cmd", "keyCreated": "API Key Created", "keyCreatedSuccess": "Key created successfully!", "keyCreatedNote": "Copy and store this key now — it won't be shown again.", diff --git a/src/i18n/messages/pt-BR.json b/src/i18n/messages/pt-BR.json index 4c74d5c0d9..110c18ab0c 100644 --- a/src/i18n/messages/pt-BR.json +++ b/src/i18n/messages/pt-BR.json @@ -342,6 +342,9 @@ "ownUsageVisibilityDesc": "Permitir que esta chave chame seu endpoint de status para ver o próprio uso em USD, percentual do orçamento e totais de tokens.", "sharedAccountQuotaVisibility": "Cota de conta compartilhada", "sharedAccountQuotaVisibilityDesc": "Permitir que esta chave veja a cota da conta upstream compartilhada quando uma conexão explícita estiver configurada.", + "localUsageCommand": "Permitir comando local de uso", + "localUsageCommandDesc": "Permite que esta API key use @@om-usage para consultar consumo e limites em cache, sem chamar provider externo.", + "localUsageCommandBadge": "Cmd uso", "keyCreated": "Chave de API Criada", "keyCreatedSuccess": "Chave criada com sucesso!", "keyCreatedNote": "Copie e armazene esta chave agora — ela não será mostrada novamente.", diff --git a/src/lib/db/apiKeys.ts b/src/lib/db/apiKeys.ts index 090dd48f77..47b82d7784 100644 --- a/src/lib/db/apiKeys.ts +++ b/src/lib/db/apiKeys.ts @@ -67,6 +67,7 @@ interface ApiKeyMetadata { allowedEndpoints: string[]; streamDefaultMode: "legacy" | "json"; disableNonPublicModels: boolean; + allowUsageCommand: boolean; } interface ApiKeyRow extends JsonRecord { @@ -98,6 +99,8 @@ interface ApiKeyRow extends JsonRecord { proxy_id?: unknown; stream_default_mode?: unknown; streamDefaultMode?: unknown; + allow_usage_command?: unknown; + allowUsageCommand?: unknown; } interface StatementLike { @@ -140,6 +143,7 @@ interface ApiKeyView extends JsonRecord { allowedEndpoints: string[]; streamDefaultMode: "legacy" | "json"; disableNonPublicModels?: boolean; + allowUsageCommand?: boolean; } // LRU cache for API key validation (valid keys only) @@ -185,6 +189,10 @@ const API_KEY_COLUMN_FALLBACKS = [ name: "disable_non_public_models", definition: "disable_non_public_models INTEGER NOT NULL DEFAULT 0", }, + { + name: "allow_usage_command", + definition: "allow_usage_command INTEGER NOT NULL DEFAULT 0", + }, ] as const; // Cache for model permission checks @@ -502,7 +510,7 @@ function getPreparedStatements(db: ApiKeysDbLike): ApiKeysStatements { "SELECT id, expires_at, revoked_at, is_active, is_banned FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtGetKeyMetadata = db.prepare( - "SELECT id, name, machine_id, allowed_models, blocked_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?" + "SELECT id, name, machine_id, allowed_models, blocked_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models, allow_usage_command, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtInsertKey = db.prepare( "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, no_log, created_at, key_prefix, key_hash, scopes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" @@ -554,6 +562,7 @@ export async function getApiKeys() { camelRow.disableNonPublicModels = parseDisableNonPublicModels( (camelRow as JsonRecord).disableNonPublicModels ); + camelRow.allowUsageCommand = parseAllowUsageCommand((camelRow as JsonRecord).allowUsageCommand); if (typeof camelRow.id === "string" && camelRow.id.length > 0) { setNoLog(camelRow.id, camelRow.noLog === true); } @@ -584,6 +593,7 @@ export async function getApiKeyById(id: string) { camelRow.disableNonPublicModels = parseDisableNonPublicModels( (camelRow as JsonRecord).disableNonPublicModels ); + camelRow.allowUsageCommand = parseAllowUsageCommand((camelRow as JsonRecord).allowUsageCommand); if (typeof camelRow.id === "string" && camelRow.id.length > 0) { setNoLog(camelRow.id, camelRow.noLog === true); } @@ -623,6 +633,10 @@ function parseDisableNonPublicModels(value: unknown): boolean { return value === true || value === 1 || value === "1"; } +function parseAllowUsageCommand(value: unknown): boolean { + return value === true || value === 1 || value === "1"; +} + function parseIsActive(value: unknown): boolean { // DEFAULT 1 — active unless explicitly set to 0 if (value === 0 || value === "0" || value === false) return false; @@ -766,6 +780,7 @@ export async function createApiKey(name: string, machineId: string, scopes: stri allowedCombos: [], // Empty array means no explicit combo restriction allowedConnections: [], // Empty array means all connections allowed noLog: false, + allowUsageCommand: false, createdAt: now, scopes, }; @@ -850,6 +865,7 @@ export async function updateApiKeyPermissions( allowedEndpoints?: string[] | null; streamDefaultMode?: "legacy" | "json" | null; disableNonPublicModels?: boolean; + allowUsageCommand?: boolean; } ) { const db = getDbInstance() as ApiKeysDbLike; @@ -883,6 +899,7 @@ export async function updateApiKeyPermissions( .streamDefaultMode, disableNonPublicModels: (update as { disableNonPublicModels?: boolean }) .disableNonPublicModels, + allowUsageCommand: (update as { allowUsageCommand?: boolean }).allowUsageCommand, }; if ( @@ -907,7 +924,8 @@ export async function updateApiKeyPermissions( (normalized as Record).proxyId === undefined && (normalized as Record).allowedEndpoints === undefined && (normalized as Record).streamDefaultMode === undefined && - normalized.disableNonPublicModels === undefined + normalized.disableNonPublicModels === undefined && + normalized.allowUsageCommand === undefined ) { return false; } @@ -936,6 +954,7 @@ export async function updateApiKeyPermissions( proxyId?: string | null; streamDefaultMode?: "legacy" | "json"; disableNonPublicModels?: number; + allowUsageCommand?: number; } = { id }; if (normalized.name !== undefined) { @@ -1034,6 +1053,11 @@ export async function updateApiKeyPermissions( params.disableNonPublicModels = normalized.disableNonPublicModels ? 1 : 0; } + if (normalized.allowUsageCommand !== undefined) { + updates.push("allow_usage_command = @allowUsageCommand"); + params.allowUsageCommand = normalized.allowUsageCommand ? 1 : 0; + } + const maxSessionsUpdate = (normalized as Record).maxSessions; if (maxSessionsUpdate !== undefined) { updates.push("max_sessions = @maxSessions"); @@ -1414,6 +1438,7 @@ export async function getApiKeyMetadata( allowedEndpoints: [], streamDefaultMode: "legacy", disableNonPublicModels: false, + allowUsageCommand: false, }; } @@ -1484,6 +1509,9 @@ export async function getApiKeyMetadata( (record as JsonRecord).disable_non_public_models ?? (record as JsonRecord).disableNonPublicModels ), + allowUsageCommand: parseAllowUsageCommand( + (record as JsonRecord).allow_usage_command ?? (record as JsonRecord).allowUsageCommand + ), }; if (!metadata.id) { diff --git a/src/lib/usage/internalUsageCommand.ts b/src/lib/usage/internalUsageCommand.ts new file mode 100644 index 0000000000..03eae21532 --- /dev/null +++ b/src/lib/usage/internalUsageCommand.ts @@ -0,0 +1,547 @@ +import type { ProviderLimitsCacheEntry } from "@/lib/db/providerLimits"; + +export const INTERNAL_USAGE_COMMAND = "@@om-usage"; +export const USAGE_COMMAND_DISABLED_MESSAGE = "Usage command is disabled for this API key."; +const USAGE_COMMAND_AUTH_REQUIRED_MESSAGE = "Usage command requires an authenticated API key."; +const LOCAL_USAGE_MODEL = "omniroute/local-usage"; + +type JsonRecord = Record; + +interface UsageCommandApiKeyMetadata { + id: string; + name?: string; + allowedConnections?: string[] | null; + allowUsageCommand?: boolean; +} + +interface ProviderConnectionLike { + id: string; + provider: string; + isActive?: boolean; +} + +interface UsageSnapshot { + connectionId: string; + provider: string; + plan: unknown; + quotas: JsonRecord; +} + +export interface InternalUsageCommandDeps { + now?: () => number; + isValidApiKey?: (apiKey: string) => Promise; + getApiKeyMetadata?: (apiKey: string) => Promise; + getProviderConnectionById?: (connectionId: string) => Promise; + getProviderConnections?: (filter?: JsonRecord) => Promise; + getProviderLimitsCache?: (connectionId: string) => ProviderLimitsCacheEntry | null; + getAllProviderLimitsCache?: () => Record; +} + +type RequiredDeps = Required; + +async function normalizeDeps(deps: InternalUsageCommandDeps = {}): Promise { + const auth = deps.isValidApiKey ? null : await import("@/sse/services/auth"); + const apiKeys = deps.getApiKeyMetadata ? null : await import("@/lib/db/apiKeys"); + const providers = + deps.getProviderConnectionById && deps.getProviderConnections + ? null + : await import("@/lib/db/providers"); + const providerLimits = + deps.getProviderLimitsCache && deps.getAllProviderLimitsCache + ? null + : await import("@/lib/db/providerLimits"); + + return { + now: deps.now ?? Date.now, + isValidApiKey: deps.isValidApiKey ?? auth!.isValidApiKey, + getApiKeyMetadata: deps.getApiKeyMetadata ?? apiKeys!.getApiKeyMetadata, + getProviderConnectionById: + deps.getProviderConnectionById ?? providers!.getProviderConnectionById, + getProviderConnections: deps.getProviderConnections ?? providers!.getProviderConnections, + getProviderLimitsCache: deps.getProviderLimitsCache ?? providerLimits!.getProviderLimitsCache, + getAllProviderLimitsCache: + deps.getAllProviderLimitsCache ?? providerLimits!.getAllProviderLimitsCache, + }; +} + +function isRecord(value: unknown): value is JsonRecord { + return value !== null && typeof value === "object" && !Array.isArray(value); +} + +function readHeader(request: Request, name: string): string | null { + return request.headers.get(name) || request.headers.get(name.toLowerCase()); +} + +function readPathScopedToken(request: Request): string | null { + try { + const url = new URL(request.url, "http://localhost"); + const segments = url.pathname + .split("/") + .map((segment) => segment.trim()) + .filter(Boolean); + + if (segments[0] === "vscode" && segments[1]) { + return decodeURIComponent(segments[1]).trim() || null; + } + + if (segments[0] === "api" && segments[1] === "v1" && segments[2] === "vscode") { + const tokenIndex = segments[3] === "raw" || segments[3] === "combos" ? 4 : 3; + if (segments[tokenIndex]) return decodeURIComponent(segments[tokenIndex]).trim() || null; + } + } catch { + return null; + } + + return null; +} + +function extractUsageCommandApiKey(request: Request): string | null { + const authHeader = readHeader(request, "Authorization"); + if (authHeader) { + const trimmed = authHeader.trim(); + if (trimmed.toLowerCase().startsWith("bearer ")) return trimmed.slice(7).trim() || null; + } + + if (readHeader(request, "anthropic-version")) { + const xApiKey = readHeader(request, "x-api-key"); + if (xApiKey?.trim()) return xApiKey.trim(); + } + + return readPathScopedToken(request); +} + +function toNumber(value: unknown, fallback = Number.NaN): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : fallback; + } + return fallback; +} + +function textFromContent(content: unknown): string | null { + if (typeof content === "string") return content; + + if (Array.isArray(content)) { + const parts: string[] = []; + for (const part of content) { + if (typeof part === "string") { + parts.push(part); + continue; + } + if (!isRecord(part)) continue; + const text = part.text ?? part.content; + if (typeof text === "string") parts.push(text); + } + return parts.length > 0 ? parts.join("") : null; + } + + if (isRecord(content)) { + const text = content.text ?? content.content; + return typeof text === "string" ? text : null; + } + + return null; +} + +function extractLastRoleText(items: unknown, role: string): string | null { + if (!Array.isArray(items)) return null; + + for (let i = items.length - 1; i >= 0; i--) { + const item = items[i]; + if (!isRecord(item) || item.role !== role) continue; + return textFromContent(item.content); + } + + return null; +} + +export function extractLastUserText(body: unknown): string | null { + if (!isRecord(body)) return null; + + const messagesText = extractLastRoleText(body.messages, "user"); + if (messagesText !== null) return messagesText; + + if (typeof body.input === "string") return body.input; + + const inputText = extractLastRoleText(body.input, "user"); + if (inputText !== null) return inputText; + + return null; +} + +export function isInternalUsageCommand(text: string | null | undefined): boolean { + return typeof text === "string" && text.trim() === INTERNAL_USAGE_COMMAND; +} + +function connectionFromValue(value: unknown): ProviderConnectionLike | null { + if (!isRecord(value)) return null; + const id = typeof value.id === "string" ? value.id : ""; + const provider = typeof value.provider === "string" ? value.provider : ""; + if (!id || !provider || value.isActive === false) return null; + return { id, provider, isActive: value.isActive === true }; +} + +function snapshotFromConnection( + connection: ProviderConnectionLike, + cache: ProviderLimitsCacheEntry | null +): UsageSnapshot | null { + if (!cache || !isRecord(cache.quotas) || Object.keys(cache.quotas).length === 0) return null; + return { + connectionId: connection.id, + provider: connection.provider, + plan: cache.plan, + quotas: cache.quotas, + }; +} + +async function collectUsageSnapshots( + metadata: UsageCommandApiKeyMetadata, + deps: RequiredDeps +): Promise { + const allowedConnections = Array.isArray(metadata.allowedConnections) + ? metadata.allowedConnections.filter((id) => typeof id === "string" && id.trim()) + : []; + + if (allowedConnections.length > 0) { + const snapshots: UsageSnapshot[] = []; + for (const connectionId of allowedConnections) { + const connection = connectionFromValue(await deps.getProviderConnectionById(connectionId)); + if (!connection) continue; + const snapshot = snapshotFromConnection( + connection, + deps.getProviderLimitsCache(connection.id) + ); + if (snapshot) snapshots.push(snapshot); + } + return snapshots; + } + + const caches = deps.getAllProviderLimitsCache(); + const connections = await deps.getProviderConnections({ isActive: true }); + const snapshots: UsageSnapshot[] = []; + for (const rawConnection of connections) { + const connection = connectionFromValue(rawConnection); + if (!connection) continue; + const snapshot = snapshotFromConnection(connection, caches[connection.id] ?? null); + if (snapshot) snapshots.push(snapshot); + } + return snapshots; +} + +function normalizeQuotaKey(key: string): string { + return key + .toLowerCase() + .replace(/[^a-z0-9]+/g, " ") + .trim(); +} + +function findQuota(quotas: JsonRecord, kind: "session" | "weekly" | "weekly-sonnet") { + const entries = Object.entries(quotas).filter(([, value]) => isRecord(value)); + + for (const [key, value] of entries) { + const normalized = normalizeQuotaKey(key); + if (kind === "session" && (normalized.includes("session") || normalized.includes("5h"))) { + return value as JsonRecord; + } + if ( + kind === "weekly-sonnet" && + normalized.includes("weekly") && + normalized.includes("sonnet") + ) { + return value as JsonRecord; + } + if ( + kind === "weekly" && + (normalized === "weekly" || normalized.includes("weekly") || normalized.includes("7d")) && + !normalized.includes("sonnet") + ) { + return value as JsonRecord; + } + } + + return null; +} + +function getQuotaUsedPercent(quota: JsonRecord | null): number | null { + if (!quota) return null; + + const usedPercentage = toNumber(quota.usedPercentage); + if (Number.isFinite(usedPercentage)) return Math.max(0, Math.min(100, usedPercentage)); + + const used = toNumber(quota.used); + const total = toNumber(quota.total); + if (Number.isFinite(used) && Number.isFinite(total) && total > 0) { + return Math.max(0, Math.min(100, (used / total) * 100)); + } + + if (Number.isFinite(used) && used >= 0 && used <= 100) { + return used; + } + + const remainingPercentage = toNumber(quota.remainingPercentage); + if (Number.isFinite(remainingPercentage)) { + return Math.max(0, Math.min(100, 100 - remainingPercentage)); + } + + const remaining = toNumber(quota.remaining); + if (Number.isFinite(remaining) && remaining >= 0 && remaining <= 100) { + return Math.max(0, Math.min(100, 100 - remaining)); + } + + return null; +} + +function getResetAt(quota: JsonRecord | null): string | null { + if (!quota) return null; + return typeof quota.resetAt === "string" && quota.resetAt.trim() ? quota.resetAt : null; +} + +function formatPercent(percent: number | null): string { + if (percent === null || !Number.isFinite(percent)) return "Unavailable"; + return `${Math.round(percent)}%`; +} + +export function formatResetIn(resetAt: string | null, now = Date.now()): string { + if (!resetAt) return "unknown"; + const resetMs = Date.parse(resetAt); + if (!Number.isFinite(resetMs)) return "unknown"; + + const deltaMs = resetMs - now; + if (deltaMs <= 0) return "now"; + + const minuteMs = 60_000; + const hourMs = 60 * minuteMs; + const dayMs = 24 * hourMs; + + if (deltaMs < hourMs) return `${Math.max(1, Math.ceil(deltaMs / minuteMs))}m`; + if (deltaMs < dayMs) return `${Math.max(1, Math.ceil(deltaMs / hourMs))}h`; + return `${Math.max(1, Math.ceil(deltaMs / dayMs))}d`; +} + +function formatPlan(plan: unknown): string { + if (typeof plan === "string" && plan.trim()) return plan.trim(); + if (typeof plan === "number" && Number.isFinite(plan)) return String(plan); + return "Unavailable"; +} + +function snapshotScore(snapshot: UsageSnapshot): number { + let score = snapshot.provider === "claude" ? 100 : 0; + if (findQuota(snapshot.quotas, "session")) score += 10; + if (findQuota(snapshot.quotas, "weekly")) score += 10; + if (findQuota(snapshot.quotas, "weekly-sonnet")) score += 10; + if (formatPlan(snapshot.plan) !== "Unavailable") score += 1; + return score; +} + +function selectUsageSnapshot(snapshots: UsageSnapshot[]): UsageSnapshot | null { + let selected: UsageSnapshot | null = null; + let bestScore = -1; + for (const snapshot of snapshots) { + const score = snapshotScore(snapshot); + if (score > bestScore) { + selected = snapshot; + bestScore = score; + } + } + return selected; +} + +function appendQuotaBlock(lines: string[], label: string, quota: JsonRecord | null, now: number) { + lines.push(label); + lines.push(formatPercent(getQuotaUsedPercent(quota))); + lines.push(`Resets in ${formatResetIn(getResetAt(quota), now)}`); +} + +export async function buildUsageCommandText( + metadata: UsageCommandApiKeyMetadata, + deps: InternalUsageCommandDeps = {} +): Promise { + const resolvedDeps = await normalizeDeps(deps); + const snapshot = selectUsageSnapshot(await collectUsageSnapshots(metadata, resolvedDeps)); + + if (!snapshot) { + return ["Plan", "Unavailable", "", "Usage", "No cached usage data available."].join("\n"); + } + + const now = resolvedDeps.now(); + const lines = ["Plan", formatPlan(snapshot.plan), "", "Usage"]; + appendQuotaBlock(lines, "Session (5hr)", findQuota(snapshot.quotas, "session"), now); + lines.push(""); + appendQuotaBlock(lines, "Weekly (7 day)", findQuota(snapshot.quotas, "weekly"), now); + lines.push(""); + appendQuotaBlock(lines, "Weekly Sonnet", findQuota(snapshot.quotas, "weekly-sonnet"), now); + return lines.join("\n"); +} + +function getResponseModel(body: unknown): string { + return isRecord(body) && typeof body.model === "string" && body.model.trim() + ? body.model + : LOCAL_USAGE_MODEL; +} + +function isAnthropicRequest(request: Request): boolean { + if (request.headers.has("anthropic-version")) return true; + try { + return new URL(request.url).pathname.endsWith("/v1/messages"); + } catch { + return false; + } +} + +function textEncoderStream(payload: string): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(payload)); + controller.close(); + }, + }); +} + +function createOpenAITextResponse(text: string, body: unknown): Response { + const created = Math.floor(Date.now() / 1000); + const model = getResponseModel(body); + const payload = { + id: `chatcmpl_usage_${created}`, + object: "chat.completion", + created, + model, + choices: [ + { + index: 0, + message: { role: "assistant", content: text }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }; + return Response.json(payload); +} + +function createOpenAIStreamResponse(text: string, body: unknown): Response { + const created = Math.floor(Date.now() / 1000); + const model = getResponseModel(body); + const id = `chatcmpl_usage_${created}`; + const first = { + id, + object: "chat.completion.chunk", + created, + model, + choices: [{ index: 0, delta: { role: "assistant", content: text }, finish_reason: null }], + }; + const second = { + id, + object: "chat.completion.chunk", + created, + model, + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + }; + return new Response( + textEncoderStream( + `data: ${JSON.stringify(first)}\n\ndata: ${JSON.stringify(second)}\n\ndata: [DONE]\n\n` + ), + { headers: { "Content-Type": "text/event-stream; charset=utf-8" } } + ); +} + +function createAnthropicTextResponse(text: string, body: unknown): Response { + const payload = { + id: `msg_usage_${Date.now()}`, + type: "message", + role: "assistant", + model: getResponseModel(body), + content: [{ type: "text", text }], + stop_reason: "end_turn", + stop_sequence: null, + usage: { input_tokens: 0, output_tokens: 0 }, + }; + return Response.json(payload); +} + +function createAnthropicStreamResponse(text: string, body: unknown): Response { + const id = `msg_usage_${Date.now()}`; + const model = getResponseModel(body); + const events = [ + [ + "message_start", + { + type: "message_start", + message: { + id, + type: "message", + role: "assistant", + model, + content: [], + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 0, output_tokens: 0 }, + }, + }, + ], + [ + "content_block_start", + { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }, + ], + [ + "content_block_delta", + { type: "content_block_delta", index: 0, delta: { type: "text_delta", text } }, + ], + ["content_block_stop", { type: "content_block_stop", index: 0 }], + [ + "message_delta", + { + type: "message_delta", + delta: { stop_reason: "end_turn", stop_sequence: null }, + usage: { output_tokens: 0 }, + }, + ], + ["message_stop", { type: "message_stop" }], + ] as const; + const payload = events + .map(([event, data]) => `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`) + .join(""); + return new Response(textEncoderStream(payload), { + headers: { "Content-Type": "text/event-stream; charset=utf-8" }, + }); +} + +export function createLocalTextResponse(request: Request, body: unknown, text: string): Response { + const stream = isRecord(body) && body.stream === true; + if (isAnthropicRequest(request)) { + return stream + ? createAnthropicStreamResponse(text, body) + : createAnthropicTextResponse(text, body); + } + return stream ? createOpenAIStreamResponse(text, body) : createOpenAITextResponse(text, body); +} + +export async function handleInternalUsageCommand( + request: Request, + body: unknown, + deps: InternalUsageCommandDeps = {} +): Promise { + const lastUserText = extractLastUserText(body); + if (!isInternalUsageCommand(lastUserText)) return null; + + const resolvedDeps = await normalizeDeps(deps); + const apiKey = extractUsageCommandApiKey(request); + if (!apiKey || !(await resolvedDeps.isValidApiKey(apiKey))) { + return createLocalTextResponse(request, body, USAGE_COMMAND_AUTH_REQUIRED_MESSAGE); + } + + const metadata = await resolvedDeps.getApiKeyMetadata(apiKey); + if (!metadata?.id) { + return createLocalTextResponse(request, body, USAGE_COMMAND_AUTH_REQUIRED_MESSAGE); + } + + if (metadata.allowUsageCommand !== true) { + return createLocalTextResponse(request, body, USAGE_COMMAND_DISABLED_MESSAGE); + } + + return createLocalTextResponse( + request, + body, + await buildUsageCommandText(metadata, resolvedDeps) + ); +} diff --git a/src/shared/utils/apiKeyPolicy.ts b/src/shared/utils/apiKeyPolicy.ts index 270cc4fc8e..77b207a8dd 100644 --- a/src/shared/utils/apiKeyPolicy.ts +++ b/src/shared/utils/apiKeyPolicy.ts @@ -91,6 +91,7 @@ export interface ApiKeyMetadata { rateLimits?: RateLimitRule[] | null; allowedEndpoints?: string[]; disableNonPublicModels?: boolean; + allowUsageCommand?: boolean; } /** diff --git a/src/shared/validation/schemas/keys.ts b/src/shared/validation/schemas/keys.ts index 9a89a9dc66..119287297b 100644 --- a/src/shared/validation/schemas/keys.ts +++ b/src/shared/validation/schemas/keys.ts @@ -21,6 +21,7 @@ import { accessScheduleSchema } from "./misc.ts"; export const createKeySchema = z.object({ name: z.string().min(1, "Name is required").max(200), noLog: z.boolean().optional(), + allowUsageCommand: z.boolean().optional(), scopes: z.array(z.string().trim().min(1).max(64)).max(32).optional(), }); @@ -102,6 +103,7 @@ export const updateKeyPermissionsSchema = z allowedEndpoints: z.array(z.string().trim().min(1).max(64)).max(20).optional(), streamDefaultMode: z.enum(["legacy", "json"]).optional(), disableNonPublicModels: z.boolean().optional(), + allowUsageCommand: z.boolean().optional(), }) .superRefine((value, ctx) => { if ( @@ -120,7 +122,9 @@ export const updateKeyPermissionsSchema = z value.rateLimits === undefined && value.scopes === undefined && value.allowedEndpoints === undefined && - value.streamDefaultMode === undefined + value.streamDefaultMode === undefined && + value.disableNonPublicModels === undefined && + value.allowUsageCommand === undefined ) { ctx.addIssue({ code: z.ZodIssueCode.custom, diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 78dabc2919..b688c22f10 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -68,6 +68,7 @@ import { generateRequestId } from "../../shared/utils/requestId"; import { logAuditEvent } from "../../lib/compliance/index"; import { enforceApiKeyPolicy } from "../../shared/utils/apiKeyPolicy"; import { cloneLogPayload } from "@/lib/logPayloads"; +import { handleInternalUsageCommand } from "@/lib/usage/internalUsageCommand"; import { applyTaskAwareRouting, getTaskRoutingConfig, @@ -241,6 +242,12 @@ export async function handleChat(request: any, clientRawRequest: any = null) { log.debug("AUTH", "No API key provided (local mode)"); } + const internalUsageCommandResponse = await handleInternalUsageCommand(request, body); + if (internalUsageCommandResponse) { + recordTelemetry(telemetry); + return internalUsageCommandResponse; + } + const isComboLiveTest = request.headers?.get?.("x-internal-test") === "combo-health-check"; if (!modelStr) { diff --git a/tests/unit/api-manager-usage-command.test.ts b/tests/unit/api-manager-usage-command.test.ts new file mode 100644 index 0000000000..41474616a7 --- /dev/null +++ b/tests/unit/api-manager-usage-command.test.ts @@ -0,0 +1,33 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; + +const repoRoot = process.cwd(); + +function read(relativePath: string) { + return fs.readFileSync(path.join(repoRoot, relativePath), "utf8"); +} + +test("api manager exposes allowUsageCommand in create, edit, and list UI", () => { + const src = read("src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx"); + + assert.ok(src.includes("newKeyAllowUsageCommand"), "create modal must keep command state"); + assert.ok(src.includes("setUsageCommandEnabled"), "permissions modal must edit command state"); + assert.ok(src.includes("allowUsageCommand"), "API payloads must include allowUsageCommand"); + assert.ok(src.includes('t("localUsageCommand")'), "toggle must use i18n title"); + assert.ok(src.includes('t("localUsageCommandBadge")'), "key list must show enabled state"); +}); + +test("api key routes and schemas accept allowUsageCommand", () => { + const schemas = read("src/shared/validation/schemas/keys.ts"); + const createRoute = read("src/app/api/keys/route.ts"); + const updateRoute = read("src/app/api/keys/[id]/route.ts"); + + assert.ok( + schemas.includes("allowUsageCommand: z.boolean().optional()"), + "zod schemas must accept the field" + ); + assert.ok(createRoute.includes("allowUsageCommand"), "create route must persist the field"); + assert.ok(updateRoute.includes("allowUsageCommand"), "update route must persist the field"); +}); diff --git a/tests/unit/apikeys-usage-command.test.ts b/tests/unit/apikeys-usage-command.test.ts new file mode 100644 index 0000000000..479e9f6424 --- /dev/null +++ b/tests/unit/apikeys-usage-command.test.ts @@ -0,0 +1,56 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-usage-command-key-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "usage-command-test-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); + +async function resetStorage() { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(() => { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("allowUsageCommand defaults to false for new API keys", async () => { + const created = await apiKeysDb.createApiKey("Usage Command Default", "machine-usage-01"); + + const metadata = await apiKeysDb.getApiKeyMetadata(created.key); + const key = await apiKeysDb.getApiKeyById(created.id); + + assert.ok(metadata); + assert.equal(metadata.allowUsageCommand, false); + assert.equal(key?.allowUsageCommand, false); +}); + +test("allowUsageCommand can be toggled through updateApiKeyPermissions", async () => { + const created = await apiKeysDb.createApiKey("Usage Command Enabled", "machine-usage-02"); + + await apiKeysDb.updateApiKeyPermissions(created.id, { allowUsageCommand: true }); + apiKeysDb.clearApiKeyCaches(); + + const enabled = await apiKeysDb.getApiKeyMetadata(created.key); + assert.equal(enabled?.allowUsageCommand, true); + + await apiKeysDb.updateApiKeyPermissions(created.id, { allowUsageCommand: false }); + apiKeysDb.clearApiKeyCaches(); + + const disabled = await apiKeysDb.getApiKeyMetadata(created.key); + assert.equal(disabled?.allowUsageCommand, false); +}); diff --git a/tests/unit/internal-usage-command.test.ts b/tests/unit/internal-usage-command.test.ts new file mode 100644 index 0000000000..35f24ed3de --- /dev/null +++ b/tests/unit/internal-usage-command.test.ts @@ -0,0 +1,240 @@ +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("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); +});