@usageflow/core 0.2.5 → 0.3.0

package/src/base.ts

@@ -9,19 +9,31 @@ import {
     RequestMetadata,
     UsageFlowRequest,
     UsageFlowAPIConfig,
+    BlockedEndpoint,
 } from "./types";
 import { UsageFlowSocketManger } from "./socket";
+import { UsageFlowLogger, usLogger } from "@usageflow/logger";
+import { randomUUID } from "crypto";
 
 export abstract class UsageFlowAPI {
     protected apiKey: string | null = null;
     protected usageflowUrl: string = "https://api.usageflow.io";
-    protected webServer: 'express' | 'fastify' | 'nestjs' = 'express';
+    protected webServer: "express" | "fastify" | "nestjs" = "express";
     protected apiConfigs: UsageFlowConfig[] = [];
+    protected blockedEndpoints: string[] = [];
     private configUpdateInterval: NodeJS.Timeout | null = null;
-    socketManager: UsageFlowSocketManger | null = null;
+    private socketManager: UsageFlowSocketManger | null = null;
     private applicationId: boolean = false;
-
-    constructor(config: UsageFlowAPIConfig = { apiKey: '', poolSize: 5 }) {
+    private logger: UsageFlowLogger;
+
+    constructor(config: UsageFlowAPIConfig = { apiKey: "", poolSize: 5 }) {
+        // Initialize logger with service name
+        this.logger = usLogger({ service: 'USAGEFLOW_BASE', pretty: true, silent: process.env.UF_LOGS_ENABLED !== 'true' });
+        // Set default context for all logs
+        this.logger.setContext({
+            component: 'UsageFlowAPI',
+            webServer: this.webServer,
+        });
         this.init(config.apiKey, this.usageflowUrl, config.poolSize);
     }
 
@@ -41,18 +53,30 @@ export abstract class UsageFlowAPI {
         // this.startConfigUpdater();
         this.socketManager = new UsageFlowSocketManger(apiKey, poolSize);
         // Connect the socket manager
-        this.socketManager.connect().catch((error) => {
-            console.error("[UsageFlow] Failed to establish WebSocket connection:", error);
-        }).then(() => {
-            this.startConfigUpdater();
-        });
-
+        this.socketManager
+            .connect()
+            .catch((error) => {
+                this.logger.error(
+                    "[UsageFlow] Failed to establish WebSocket connection:",
+                    error,
+                );
+            })
+            .then(() => {
+                this.logger.info("WebSocket connection established");
+                if (this.socketManager?.isConnected()) {
+                    this.startConfigUpdater();
+                } else {
+                    this.logger.error("WebSocket connection failed");
+                }
+            });
 
         return this;
     }
 
-    public setWebServer(webServer: 'express' | 'fastify' | 'nestjs'): void {
+    public setWebServer(webServer: "express" | "fastify" | "nestjs"): void {
         this.webServer = webServer;
+        // Update logger context with new web server
+        this.logger.setContext({ webServer });
     }
 
     public getApiKey(): string | null {
@@ -67,26 +91,27 @@ export abstract class UsageFlowAPI {
      * Start background config update process
      */
     private startConfigUpdater(): void {
+        this.logger.info("Starting background config update process");
         if (this.configUpdateInterval) {
             clearInterval(this.configUpdateInterval);
         }
 
-        this.fetchApiPolicies().catch(console.error);
+        this.fetchApiPolicies().catch(this.logger.error);
         this.configUpdateInterval = setInterval(async () => {
-            await this.fetchApiPolicies();
-        }, 60000);
-
-        // this.fetchApiConfig().catch(console.error);
+            await Promise.all([
+                this.fetchApiPolicies(),
+                this.fetchBlockedEndpoints(),
+            ]);
+        }, 10000);
     }
 
     public getRoutePattern(request: UsageFlowRequest): string {
-
-        if (this.webServer === 'fastify') {
-            return request?.routeOptions?.url || request.url || '';
+        if (this.webServer === "fastify") {
+            return request?.routeOptions?.url || request.url || "";
         }
 
         // For NestJS, prioritize route.path with baseUrl
-        if (this.webServer === 'nestjs') {
+        if (this.webServer === "nestjs") {
             // Method 1: Use request.route.path (available after route matching in NestJS)
             if (request.route?.path) {
                 const baseUrl = request.baseUrl || "";
@@ -102,7 +127,9 @@ export abstract class UsageFlowAPI {
                 let path = request.path || request.url?.split("?")[0] || "/";
 
                 // Split path into segments and replace matching segments with param names
-                const pathSegments = path.split("/").filter(segment => segment !== "");
+                const pathSegments = path
+                    .split("/")
+                    .filter((segment) => segment !== "");
                 const paramEntries = Object.entries(request.params);
 
                 // Replace segments that match param values with :paramName
@@ -121,21 +148,22 @@ export abstract class UsageFlowAPI {
             }
         }
 
-        const routePattern = request.route?.path || request.app._router.stack.find((route: any) => {
-            // a => a.path == request.url
-            if (!route.route) return false;
-            if (route.path) {
-                return route.path == request.url;
-            }
-
-            if (route.regexp) {
-                const patterned = route.regexp.test(request.url);
-                if (patterned) {
-                    return patterned;
+        const routePattern =
+            request.route?.path ||
+            request.app._router.stack.find((route: any) => {
+                // a => a.path == request.url
+                if (!route.route) return false;
+                if (route.path) {
+                    return route.path == request.url;
                 }
-            }
 
-        })?.route?.path
+                if (route.regexp) {
+                    const patterned = route.regexp.test(request.url);
+                    if (patterned) {
+                        return patterned;
+                    }
+                }
+            })?.route?.path;
 
         if (routePattern) {
             return routePattern;
@@ -163,7 +191,9 @@ export abstract class UsageFlowAPI {
                     if (layer.route) {
                         const route = layer.route;
                         const routePath = route.path;
-                        const routeMethods = Object.keys(route.methods).map(m => m.toLowerCase());
+                        const routeMethods = Object.keys(route.methods).map((m) =>
+                            m.toLowerCase(),
+                        );
 
                         // Check if method matches and path pattern matches
                         if (routeMethods.includes(method) || routeMethods.includes("*")) {
@@ -173,7 +203,9 @@ export abstract class UsageFlowAPI {
                             if (request.params && Object.keys(request.params).length > 0) {
                                 // Check if route path contains the param names
                                 const paramNames = Object.keys(request.params);
-                                const routeHasParams = paramNames.some(param => routePath.includes(`:${param}`));
+                                const routeHasParams = paramNames.some((param) =>
+                                    routePath.includes(`:${param}`),
+                                );
                                 if (routeHasParams) {
                                     return pattern;
                                 }
@@ -184,7 +216,11 @@ export abstract class UsageFlowAPI {
                                 }
                             }
                         }
-                    } else if (layer.name === "router" && layer.handle && layer.handle.stack) {
+                    } else if (
+                        layer.name === "router" &&
+                        layer.handle &&
+                        layer.handle.stack
+                    ) {
                         // Handle router middleware (nested routers)
                         const mountPath = layer.regexp.source
                             .replace("\\/?", "")
@@ -198,19 +234,30 @@ export abstract class UsageFlowAPI {
                             if (subLayer.route) {
                                 const route = subLayer.route;
                                 const routePath = route.path;
-                                const routeMethods = Object.keys(route.methods).map(m => m.toLowerCase());
-
-                                if (routeMethods.includes(method) || routeMethods.includes("*")) {
+                                const routeMethods = Object.keys(route.methods).map((m) =>
+                                    m.toLowerCase(),
+                                );
+
+                                if (
+                                    routeMethods.includes(method) ||
+                                    routeMethods.includes("*")
+                                ) {
                                     const fullPath = mountPath + routePath;
                                     // Check if this route matches by examining params
-                                    if (request.params && Object.keys(request.params).length > 0) {
+                                    if (
+                                        request.params &&
+                                        Object.keys(request.params).length > 0
+                                    ) {
                                         const paramNames = Object.keys(request.params);
-                                        const routeHasParams = paramNames.some(param => routePath.includes(`:${param}`));
+                                        const routeHasParams = paramNames.some((param) =>
+                                            routePath.includes(`:${param}`),
+                                        );
                                         if (routeHasParams) {
                                             return fullPath;
                                         }
                                     } else if (!routePath.includes(":")) {
-                                        const currentPath = request.path || request.url?.split("?")[0] || "";
+                                        const currentPath =
+                                            request.path || request.url?.split("?")[0] || "";
                                         if (currentPath === fullPath || currentPath === routePath) {
                                             return fullPath;
                                         }
@@ -223,7 +270,10 @@ export abstract class UsageFlowAPI {
             }
         } catch (error) {
             // Silently fail and try next method
-            console.debug("[UsageFlow] Could not extract route from router stack:", error);
+            console.debug(
+                "[UsageFlow] Could not extract route from router stack:",
+                error,
+            );
         }
 
         // Method 3: Reconstruct pattern from params and path
@@ -255,111 +305,160 @@ export abstract class UsageFlowAPI {
         return baseUrl + path || "/";
     }
 
-    public guessLedgerId(request: UsageFlowRequest, overrideUrl?: string): string {
+    public guessLedgerId(
+        request: UsageFlowRequest,
+        overrideUrl?: string,
+    ): { ledgerId: string, hasLimit: boolean } {
         const method = request.method;
         const url = overrideUrl || this.getRoutePattern(request);
-        const configs = this.apiConfigs
+        const configs = this.apiConfigs;
 
         if (!configs.length) {
-            return `${method} ${url}`;
+            return { ledgerId: `${method} ${url}`, hasLimit: false };
         }
 
-
         for (const config of configs) {
             const fieldName = config.identityFieldName!;
             const location = config.identityFieldLocation;
+            const hasLimit = config.hasRateLimit || false;
 
-            switch (location) {
-                case "path_params":
-                    if (request.params?.[fieldName]) {
-                        return `${method} ${url} ${request.params[fieldName]}`;
-                    }
-                    break;
-                case "query_params":
-                    if (request.query?.[fieldName]) {
-                        return `${method} ${url} ${request.query[fieldName]}`;
-                    }
-                    break;
-                case "body":
-                    if (request.body?.[fieldName]) {
-                        return `${method} ${url} ${request.body[fieldName]}`;
-                    }
-                    break;
-                case "bearer_token":
-                    const authHeader = this.getHeaderValue(request.headers, 'authorization');
-                    const token = this.extractBearerToken(authHeader || undefined);
-                    if (token) {
-                        const claims = this.decodeJwtUnverified(token);
-                        if (claims?.[fieldName]) {
-                            return `${method} ${url} ${claims[fieldName]}`;
+            if (method === config.method && url === config.url) {
+
+                switch (location) {
+                    case "path_params":
+                        if (request.params?.[fieldName]) {
+                            return { ledgerId: `${method} ${url} ${request.params[fieldName]}`, hasLimit };
                         }
-                    }
-                    break;
+                        break;
+                    case "query_params":
+                        if (request.query?.[fieldName]) {
+                            return { ledgerId: `${method} ${url} ${request.query[fieldName]}`, hasLimit };
+                        }
+                        break;
+                    case "body":
+                        if (request.body?.[fieldName]) {
+                            return { ledgerId: `${method} ${url} ${request.body[fieldName]}`, hasLimit };
+                        }
+                        break;
+                    case "bearer_token":
+                        const authHeader = this.getHeaderValue(
+                            request.headers,
+                            "authorization",
+                        );
+                        const token = this.extractBearerToken(authHeader || undefined);
+                        if (token) {
+                            const claims = this.decodeJwtUnverified(token);
+                            if (claims?.[fieldName]) {
+                                return { ledgerId: `${method} ${url} ${claims[fieldName]}`, hasLimit };
+                            }
+                        }
+                        break;
 
-                case "headers": {
-                    const headerValue = this.getHeaderValue(request.headers, fieldName);
-                    if (headerValue) {
-                        return `${method} ${url} ${headerValue}`;
+                    case "headers": {
+                        const headerValue = this.getHeaderValue(request.headers, fieldName);
+                        if (headerValue) {
+                            return { ledgerId: `${method} ${url} ${headerValue}`, hasLimit };
+                        }
+                        break;
                     }
-                    break;
                 }
             }
         }
 
-        return `${method} ${url}`;
+        return { ledgerId: `${method} ${url}`, hasLimit: false };
     }
 
     private async fetchApiPolicies(): Promise<void> {
         if (this.socketManager && this.socketManager.isConnected()) {
-            const response = await this.socketManager.sendAsync<{ policies: UsageFlowConfig[], total: number }>({
+            const response = await this.socketManager.sendAsync<{
+                policies: UsageFlowConfig[];
+                total: number;
+            }>({
                 type: "get_application_policies",
                 payload: null,
             });
-            if (response.type === 'success') {
+            if (response.type === "success") {
                 this.apiConfigs = response.payload?.policies || [];
             }
         }
     }
 
-    async useAllocationRequest(
-        payload: RequestForAllocation,
-    ): Promise<void> {
+
+    private async fetchBlockedEndpoints(): Promise<void> {
         if (this.socketManager && this.socketManager.isConnected()) {
-            this.socketManager.sendAsync<any>({
-                type: "use_allocation",
-                payload
-            }).catch((error) => {
-                console.error("[UsageFlow] Error sending finalization via WebSocket:", error);
-                throw error;
+            const response = await this.socketManager.sendAsync<{
+                endpoints: BlockedEndpoint[];
+                total: number;
+            }>({
+                type: "get_blocked_endpoints",
+                payload: null,
             });
+            if (response.type === "success") {
+                response.payload?.endpoints || [];
+                this.blockedEndpoints = response.payload?.endpoints.map(({ method, url, identity }) => `${method} ${url} ${identity}`) || [];
+            }
         }
     }
 
+    async useAllocationRequest(payload: RequestForAllocation): Promise<void> {
+        if (this.socketManager && this.socketManager.isConnected()) {
+            this.socketManager
+                .sendAsync<any>({
+                    type: "use_allocation",
+                    payload,
+                })
+                .catch((error) => {
+                    console.error(
+                        "[UsageFlow] Error sending finalization via WebSocket:",
+                        error,
+                    );
+                    throw error;
+                });
+        }
+    }
 
     async allocationRequest(
         request: UsageFlowRequest,
         payload: RequestForAllocation,
         metadata: RequestMetadata,
+        hasLimit: boolean,
     ): Promise<void> {
         if (this.socketManager && this.socketManager.isConnected()) {
             try {
-                const allocationResponse = await this.socketManager.sendAsync<any>({
-                    type: "request_for_allocation",
-                    payload
-                });
-
-                if (allocationResponse.type === 'error') {
-                    throw new Error(allocationResponse.message || allocationResponse.error);
-                }
-                if (allocationResponse.type === 'success') {
-                    request.usageflow!.eventId = allocationResponse.payload.allocationId;
-                    request.usageflow!.metadata = metadata;
-                    return;
+                if (hasLimit) {
+                    const allocationResponse = await this.socketManager.sendAsync<any>({
+                        type: "request_for_allocation",
+                        payload,
+                    });
+
+                    if (allocationResponse.type === "error") {
+                        throw new Error(
+                            allocationResponse.message || allocationResponse.error,
+                        );
+                    }
+                    if (allocationResponse.type === "success") {
+                        request.usageflow!.eventId = allocationResponse.payload.allocationId;
+                        request.usageflow!.metadata = metadata;
+                        return;
+                    } else {
+                        throw new Error(
+                            allocationResponse.message || "Unknown error occurred",
+                        );
+                    }
                 } else {
-                    throw new Error(allocationResponse.message || "Unknown error occurred");
+                    payload.allocationId = randomUUID();
+                    await this.socketManager.send({
+                        type: "request_for_allocation",
+                        payload,
+                    });
+                    request.usageflow!.eventId = payload.allocationId;
+                    request.usageflow!.metadata = metadata;
                 }
             } catch (error: any) {
-                console.error("[UsageFlow] WebSocket allocation failed, falling back to HTTP:", error);
+                console.error(
+                    "[UsageFlow] WebSocket allocation failed, falling back to HTTP:",
+                    error,
+                );
                 throw error;
                 // Fall through to HTTP request
             }
@@ -468,8 +567,8 @@ export abstract class UsageFlowAPI {
         }
 
         // Check if it's a bearer token
-        const parts = authHeader.split(' ');
-        if (parts.length !== 2 || parts[0].toLowerCase() !== 'bearer') {
+        const parts = authHeader.split(" ");
+        if (parts.length !== 2 || parts[0].toLowerCase() !== "bearer") {
             return null;
         }
 
@@ -479,14 +578,19 @@ export abstract class UsageFlowAPI {
     /**
      * Get header value from headers object (handles both Record and Headers types)
      */
-    private getHeaderValue(headers: Record<string, string | string[] | undefined> | Headers, key: string): string | null {
+    private getHeaderValue(
+        headers: Record<string, string | string[] | undefined> | Headers,
+        key: string,
+    ): string | null {
         // Check if it's a Headers object (from Fetch API) by checking for the 'get' method
-        if (headers && typeof (headers as any).get === 'function') {
+        if (headers && typeof (headers as any).get === "function") {
             return (headers as any).get(key) || null;
         }
         // Otherwise, treat it as a Record
-        const value = (headers as Record<string, string | string[] | undefined>)[key];
-        if (typeof value === 'string') {
+        const value = (headers as Record<string, string | string[] | undefined>)[
+            key
+        ];
+        if (typeof value === "string") {
             return value;
         }
         if (Array.isArray(value) && value.length > 0) {
@@ -503,14 +607,14 @@ export abstract class UsageFlowAPI {
     public decodeJwtUnverified(token: string): Record<string, any> | null {
         try {
             // Split the token into parts
-            const parts = token.split('.');
+            const parts = token.split(".");
             if (parts.length !== 3) {
                 return null;
             }
 
             // Decode the payload (claims)
             const payload = parts[1];
-            const decoded = Buffer.from(payload, 'base64').toString('utf-8');
+            const decoded = Buffer.from(payload, "base64").toString("utf-8");
             return JSON.parse(decoded);
         } catch (error) {
             return null;

package/src/types.ts

@@ -62,6 +62,13 @@ export interface UsageFlowConfig {
     method: string;
     identityFieldName?: string;
     identityFieldLocation?: string;
+    hasRateLimit?: boolean;
+}
+
+export interface BlockedEndpoint {
+    method: string;
+    url: string;
+    identity: string
 }
 
 export interface UsageFlowResponse {
@@ -90,7 +97,7 @@ export interface RequestForAllocation {
     duration?: number;
 }
 
-export type MessageTypes = 'request_for_allocation' | 'use_allocation' | 'get_application_policies';
+export type MessageTypes = 'request_for_allocation' | 'use_allocation' | 'get_application_policies' | 'get_blocked_endpoints';
 
 export interface UsageFlowSocketMessage {
     type: MessageTypes;