refactor to address comments

Author: cbrady

package.json

@@ -29,20 +29,21 @@
         "readable-stream": "^4.5.2"
     },
     "devDependencies": {
+        "@biomejs/biome": "2.3.1",
+        "@jest/globals": "^29.7.0",
+        "@types/jest": "^29.5.14",
+        "@types/node": "^18.19.70",
         "@types/node-fetch": "^2.6.12",
         "@types/readable-stream": "^4.0.18",
-        "webpack": "^5.97.1",
-        "ts-loader": "^9.5.1",
+        "@types/ws": "^8.18.1",
+        "esbuild": "^0.25.9",
         "jest": "^29.7.0",
-        "@jest/globals": "^29.7.0",
-        "@types/jest": "^29.5.14",
-        "ts-jest": "^29.3.4",
         "jest-environment-jsdom": "^29.7.0",
         "msw": "2.11.2",
-        "@types/node": "^18.19.70",
+        "ts-jest": "^29.3.4",
+        "ts-loader": "^9.5.1",
         "typescript": "~5.7.2",
-        "@biomejs/biome": "2.3.1",
-        "esbuild": "^0.25.9"
+        "webpack": "^5.97.1"
     },
     "browser": {
         "fs": false,

src/datastream/datastream-client.ts

@@ -10,24 +10,6 @@ import type { CacheProvider } from '../cache/types';
 import { LocalCache } from '../cache/local';
 import { RedisCacheProvider, type RedisOptions } from '../cache/redis';
 
-/**
- * Redis cache configuration options for DataStream client
- */
-export interface DataStreamRedisConfig {
-  /** Redis connection URL (e.g., 'redis://localhost:6379') */
-  url?: string;
-  /** Redis host (default: 'localhost') */
-  host?: string;
-  /** Redis port (default: 6379) */
-  port?: number;
-  /** Redis password */
-  password?: string;
-  /** Redis database number (default: 0) */
-  db?: number;
-  /** Redis key prefix for all cache keys (default: 'schematic:') */
-  keyPrefix?: string;
-}
-
 /**
  * Options for configuring the DataStream client
  */
@@ -41,7 +23,7 @@ export interface DataStreamClientOptions {
   /** Cache TTL in milliseconds (default: 5 minutes) */
   cacheTTL?: number;
   /** Redis configuration for all cache providers (if not provided, uses memory cache) */
-  redisConfig?: DataStreamRedisConfig;
+  redisConfig?: RedisOptions;
   /** Custom cache provider for company entities (overrides Redis config if provided) */
   companyCache?: CacheProvider<Schematic.RulesengineCompany>;
   /** Custom cache provider for user entities (overrides Redis config if provided) */
@@ -637,7 +619,7 @@ export class DataStreamClient extends EventEmitter {
   /**
    * handleMessage processes incoming datastream messages
    */
-  private async handleMessage(ctx: any, message: DataStreamResp): Promise<void> {
+  private async handleMessage(message: DataStreamResp): Promise<void> {
     this.logger.debug(`Processing datastream message: EntityType=${message.entity_type}, MessageType=${message.message_type}`);
 
     try {
@@ -754,17 +736,23 @@ export class DataStreamClient extends EventEmitter {
       return;
     }
 
-    const cacheKeys: string[] = [];
-    for (const flag of flags) {
-      if (flag?.key) {
-        const cacheKey = this.flagCacheKey(flag.key);
-        try {
+    const results = await Promise.allSettled(
+      flags
+        .filter((flag) => flag?.key)
+        .map(async (flag) => {
+          const cacheKey = this.flagCacheKey(flag.key!);
           await this.flagsCacheProvider.set(cacheKey, flag);
-          cacheKeys.push(cacheKey);
-        } catch (error) {
-          this.logger.warn(`Failed to cache flag: ${error}`);
-        }
-      }
+          return cacheKey;
+        }),
+    );
+
+    const cacheKeys = results
+      .filter((r): r is PromiseFulfilledResult<string> => r.status === 'fulfilled')
+      .map((r) => r.value);
+
+    const failures = results.filter((r) => r.status === 'rejected');
+    if (failures.length > 0) {
+      this.logger.warn(`Failed to cache ${failures.length} flag(s)`);
     }
 
     // Delete flags not in the response
@@ -845,7 +833,7 @@ export class DataStreamClient extends EventEmitter {
   /**
    * handleConnectionReady is called when the WebSocket connection is ready
    */
-  private async handleConnectionReady(ctx: any): Promise<void> {
+  private async handleConnectionReady(): Promise<void> {
     this.logger.info('DataStream connection is ready');
 
     // Request initial flag data

src/datastream/index.ts

@@ -1,8 +1,7 @@
 // Main datastream client for external use
-export { 
+export {
   DataStreamClient
 } from './datastream-client';
 export type {
   DataStreamClientOptions,
-  DataStreamRedisConfig,
 } from './datastream-client';

src/datastream/websocket-client.test.ts

@@ -25,11 +25,11 @@ describe('DatastreamWSClient', () => {
   let client: DatastreamWSClient;
   const mockLogger = new MockLogger();
 
-  const mockMessageHandler: MessageHandlerFunc = async (ctx: any, message: DataStreamResp) => {
+  const mockMessageHandler: MessageHandlerFunc = async (message: DataStreamResp) => {
     console.log('Received message:', message);
   };
 
-  const mockConnectionReadyHandler: ConnectionReadyHandlerFunc = async (ctx: any) => {
+  const mockConnectionReadyHandler: ConnectionReadyHandlerFunc = async () => {
     console.log('Connection ready');
   };
 
@@ -177,6 +177,6 @@ describe('DatastreamWSClient', () => {
       logger: mockLogger,
     });
 
-    await expect(client.sendMessage({ test: 'message' })).rejects.toThrow('WebSocket connection is not available!');
+    await expect(client.sendMessage({ data: { entity_type: 'rulesengine.Company' as any, keys: {} } })).rejects.toThrow('WebSocket connection is not available!');
   });
 });

src/datastream/websocket-client.ts

@@ -2,8 +2,9 @@
 // The ws package is required and provides the WebSocket implementation
 
 import { EventEmitter } from 'events';
-import { DataStreamResp } from './types';
+import { DataStreamResp, DataStreamBaseReq } from './types';
 import { Logger } from '../logger';
+import type WebSocket from 'ws';
 
 // Dynamic imports to avoid webpack issues
 const createWebSocket = () => {
@@ -37,12 +38,12 @@ const MAX_RECONNECT_DELAY = 30 * 1000; // 30 seconds
  * MessageHandlerFunc is a function type for handling incoming datastream messages
  * Expects parsed DataStreamResp messages
  */
-export type MessageHandlerFunc = (ctx: any, message: DataStreamResp) => Promise<void>;
+export type MessageHandlerFunc = (message: DataStreamResp) => Promise<void>;
 
 /**
  * ConnectionReadyHandlerFunc is a function type for functions that need to be called before connection is considered ready
  */
-export type ConnectionReadyHandlerFunc = (ctx: any) => Promise<void>;
+export type ConnectionReadyHandlerFunc = () => Promise<void>;
 
 /**
  * ClientOptions contains configuration for the datastream client
@@ -74,7 +75,7 @@ export interface ClientOptions {
  *   https://custom.example.com -> wss://custom.example.com/datastream
  *   http://localhost:8080 -> ws://localhost:8080/datastream
  */
-function convertAPIURLToWebSocketURL(apiURL: string): any {
+function convertAPIURLToWebSocketURL(apiURL: string): string {
   const URLClass = createURL();
   const parsedURL = new URLClass(apiURL);
 
@@ -110,7 +111,7 @@ function convertAPIURLToWebSocketURL(apiURL: string): any {
  */
 export class DatastreamWSClient extends EventEmitter {
   // Configuration
-  private readonly url: any;
+  private readonly url: string;
   private readonly headers: Record<string, string>;
   private readonly logger: Logger;
   private readonly messageHandler: MessageHandlerFunc;
@@ -120,7 +121,7 @@ export class DatastreamWSClient extends EventEmitter {
   private readonly maxReconnectDelay: number;
 
   // Connection state
-  private ws?: any;
+  private ws?: WebSocket;
   private connected: boolean = false;
   private ready: boolean = false;
 
@@ -131,9 +132,6 @@ export class DatastreamWSClient extends EventEmitter {
   private pongTimeout?: NodeJS.Timeout;
   private reconnectTimeout?: NodeJS.Timeout;
 
-  // Context
-  private ctx: any = {};
-
   constructor(options: ClientOptions) {
     super();
 
@@ -153,8 +151,7 @@ export class DatastreamWSClient extends EventEmitter {
     if (options.url.startsWith('http://') || options.url.startsWith('https://')) {
       this.url = convertAPIURLToWebSocketURL(options.url);
     } else {
-      const URLClass = createURL();
-      this.url = new URLClass(options.url);
+      this.url = options.url;
     }
 
     // Create headers with API key
@@ -198,7 +195,7 @@ export class DatastreamWSClient extends EventEmitter {
   /**
    * SendMessage sends a message through the WebSocket connection
    */
-  public async sendMessage(message: any): Promise<void> {
+  public async sendMessage(message: DataStreamBaseReq): Promise<void> {
     if (!this.isConnected() || !this.ws) {
       throw new Error('WebSocket connection is not available!');
     }
@@ -280,7 +277,7 @@ export class DatastreamWSClient extends EventEmitter {
           // Call connection ready handler if provided
           if (this.connectionReadyHandler) {
             try {
-              await this.connectionReadyHandler(this.ctx);
+              await this.connectionReadyHandler();
               this.logger.debug('Connection ready handler completed successfully');
             } catch (err) {
               this.logger.error(`Connection ready handler failed: ${err}`);
@@ -354,10 +351,10 @@ export class DatastreamWSClient extends EventEmitter {
    */
   private connect(): Promise<any> {
     return new Promise((resolve, reject) => {
-      this.logger.debug(`Connecting to WebSocket: ${this.url.toString()}`);
+      this.logger.debug(`Connecting to WebSocket: ${this.url}`);
 
       const WebSocketClass = createWebSocket();
-      const ws = new WebSocketClass(this.url.toString(), {
+      const ws = new WebSocketClass(this.url, {
         headers: this.headers,
         handshakeTimeout: 30000, // 30 seconds
       });
@@ -382,8 +379,8 @@ export class DatastreamWSClient extends EventEmitter {
   /**
    * setupWebSocketHandlers sets up message and error handlers for the WebSocket
    */
-  private setupWebSocketHandlers(ws: any): void {
-    ws.on('message', (data: any) => {
+  private setupWebSocketHandlers(ws: WebSocket): void {
+    ws.on('message', (data: Buffer | Buffer[]) => {
       this.handleMessage(data);
     });
 
@@ -403,7 +400,7 @@ export class DatastreamWSClient extends EventEmitter {
   /**
    * handleMessage processes incoming WebSocket messages
    */
-  private async handleMessage(data: any): Promise<void> {
+  private async handleMessage(data: Buffer | Buffer[]): Promise<void> {
     try {
 
       let messageStr: string;
@@ -412,7 +409,8 @@ export class DatastreamWSClient extends EventEmitter {
       } else if (Array.isArray(data)) {
         messageStr = Buffer.concat(data).toString();
       } else {
-        messageStr = data.toString();
+        // Fallback for other string-like types
+        messageStr = String(data);
       }
 
       // Parse the datastream message
@@ -427,7 +425,7 @@ export class DatastreamWSClient extends EventEmitter {
 
       // Handle the parsed message using the provided handler
       try {
-        await this.messageHandler(this.ctx, message);
+        await this.messageHandler(message);
       } catch (err) {
         this.emit('error', new Error(`Message handler error: ${err}`));
       }

src/index.ts

@@ -3,7 +3,6 @@ export { LocalCache } from "./cache/local";
 export { RedisCacheProvider } from "./cache/redis";
 export { SchematicClient } from "./wrapper";
 export type { RedisOptions } from "./cache/redis";
-export type { DataStreamRedisConfig } from "./datastream";
 export { SchematicEnvironment } from "./environments";
 export { SchematicError, SchematicTimeoutError } from "./errors";
 export { RulesEngineClient } from "./rules-engine";

src/wrapper.ts

@@ -5,7 +5,8 @@ import { type CacheProvider, LocalCache } from "./cache";
 import { ConsoleLogger, Logger } from "./logger";
 import { EventBuffer } from "./events";
 import { offlineFetcher, provideFetcher } from "./core/fetcher/custom";
-import { DataStreamClient, type DataStreamClientOptions, type DataStreamRedisConfig } from "./datastream";
+import { DataStreamClient, type DataStreamClientOptions } from "./datastream";
+import type { RedisOptions } from "./cache/redis";
 
 /**
  * Configuration options for the SchematicClient
@@ -27,7 +28,7 @@ export interface SchematicOptions {
         /** Cache TTL in milliseconds (default: 5 minutes) */
         cacheTTL?: number;
         /** Redis configuration for DataStream caching */
-        redisConfig?: DataStreamRedisConfig;
+        redisConfig?: RedisOptions;
         /** Enable replicator mode for external data synchronization */
         replicatorMode?: boolean;
         /** Health check URL for replicator mode */
@@ -174,22 +175,17 @@ export class SchematicClient extends BaseClient {
                 // Extract boolean value from DataStream response
                 const flagValue = typeof resp === 'boolean' ? resp : resp?.value;
 
-                // Track the flag check event
-                this.track({
-                    event: "flag_check",
-                    company: evalCtx.company,
-                    user: evalCtx.user,
-                    traits: {
-                        flag_key: key,
-                        value: flagValue,
-                        company_id: resp?.companyId,
-                        user_id: resp?.userId,
-                        flag_id: resp?.flagId,
-                        req_company: evalCtx.company,
-                        req_user: evalCtx.user,
-                        reason: resp?.reason,
-                    },
-                });
+                // Enqueue the flag check event
+                this.enqueueEvent(api.EventType.FlagCheck, {
+                    flagKey: key,
+                    value: flagValue ?? false,
+                    reason: resp?.reason ?? "unknown",
+                    companyId: resp?.companyId,
+                    userId: resp?.userId,
+                    flagId: resp?.flagId,
+                    reqCompany: evalCtx.company,
+                    reqUser: evalCtx.user,
+                } satisfies api.EventBodyFlagCheck);
 
                 return flagValue ?? this.getFlagDefault(key);
             } catch (err) {
@@ -433,8 +429,8 @@ export class SchematicClient extends BaseClient {
     }
 
     private async enqueueEvent(
-        eventType: "identify" | "track",
-        body: api.EventBodyIdentify | api.EventBodyTrack
+        eventType: api.EventType,
+        body: api.EventBody
     ): Promise<void> {
         try {
             this.eventBuffer.push({