@comprehend/telemetry-node 0.1.4 → 0.2.1

package/dist/WebSocketConnection.d.ts

@@ -1,15 +1,33 @@
-import { ObservationInputMessage } from './wire-protocol';
+import { AttributeType, CustomMetricSpecification, ObservationInputMessage, InitAck } from './wire-protocol';
 export declare class WebSocketConnection {
     private readonly organization;
     private readonly token;
     private readonly logger?;
+    private readonly onAuthorized?;
+    private readonly onCustomMetricChange?;
     private readonly unacknowledgedObserved;
-    private readonly unacknowledgedObservations;
+    private readonly seqQueues;
+    private readonly observationsQueue;
+    private readonly timeseriesQueue;
+    private readonly cumulativeQueue;
+    private readonly traceSpansQueue;
+    private readonly dbQueryQueue;
+    private readonly contextQueue;
     private socket;
     private reconnectDelay;
     private shouldReconnect;
     private authorized;
-    constructor(organization: string, token: string, logger?: (message: string) => void);
+    private processContext;
+    private ingestionId;
+    private _seq;
+    constructor(options: {
+        organization: string;
+        token: string;
+        logger?: (message: string) => void;
+        onAuthorized?: (ack: InitAck) => void;
+        onCustomMetricChange?: (specs: CustomMetricSpecification[]) => void;
+    });
+    nextSeq(): number;
     private log;
     private connect;
     private onOpen;
@@ -17,6 +35,8 @@ export declare class WebSocketConnection {
     private onClose;
     private onError;
     private sendRaw;
+    private sendContextStart;
+    setProcessContext(serviceEntityHash: string, resources: Record<string, AttributeType>): void;
     sendMessage(message: ObservationInputMessage): void;
     close(): void;
 }

package/dist/WebSocketConnection.js

@@ -5,20 +5,57 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.WebSocketConnection = void 0;
 const ws_1 = __importDefault(require("ws"));
-const INGESTION_ENDPOINT = 'wss://ingestion.comprehend.dev';
+const crypto_1 = require("crypto");
+const util_1 = require("./util");
+const INGESTION_ENDPOINT = process.env.COMPREHEND_INGESTION_ENDPOINT || 'wss://ingestion.comprehend.dev';
+/** A queue of unacknowledged messages keyed by sequence number, with an associated ack type. */
+class SeqQueue {
+    constructor(ackType) {
+        this.pending = new Map();
+        this.ackType = ackType;
+    }
+    add(message) {
+        this.pending.set(message.seq, message);
+    }
+    ack(seq) {
+        this.pending.delete(seq);
+    }
+    values() {
+        return this.pending.values();
+    }
+}
 class WebSocketConnection {
-    constructor(organization, token, logger) {
+    constructor(options) {
         this.unacknowledgedObserved = new Map();
-        this.unacknowledgedObservations = new Map();
+        this.observationsQueue = new SeqQueue('ack-observations');
+        this.timeseriesQueue = new SeqQueue('ack-timeseries');
+        this.cumulativeQueue = new SeqQueue('ack-cumulative');
+        this.traceSpansQueue = new SeqQueue('ack-tracespans');
+        this.dbQueryQueue = new SeqQueue('ack-db-query');
+        this.contextQueue = new SeqQueue('ack-context');
         this.socket = null;
         this.reconnectDelay = 1000;
         this.shouldReconnect = true;
         this.authorized = false;
-        this.organization = organization;
-        this.token = token;
-        this.logger = logger;
+        this.processContext = null;
+        this.ingestionId = (0, crypto_1.randomUUID)();
+        this._seq = 1;
+        this.organization = options.organization;
+        this.token = options.token;
+        this.logger = options.logger;
+        this.onAuthorized = options.onAuthorized;
+        this.onCustomMetricChange = options.onCustomMetricChange;
+        // Build a lookup from ack type to queue for dispatch in onMessage
+        this.seqQueues = {};
+        for (const q of [this.observationsQueue, this.timeseriesQueue, this.cumulativeQueue,
+            this.traceSpansQueue, this.dbQueryQueue, this.contextQueue]) {
+            this.seqQueues[q.ackType] = q;
+        }
         this.connect();
     }
+    nextSeq() {
+        return this._seq++;
+    }
     log(message) {
         if (this.logger) {
             this.logger(message);
@@ -35,9 +72,10 @@ class WebSocketConnection {
     }
     onOpen() {
         this.log('WebSocket connected. Sending init/auth message.');
+        this.ingestionId = (0, crypto_1.randomUUID)();
         const init = {
             event: 'init',
-            protocolVersion: 1,
+            protocolVersion: 2,
             token: this.token,
         };
         this.sendRaw(init);
@@ -48,18 +86,40 @@ class WebSocketConnection {
             if (msg.type === 'ack-authorized') {
                 this.authorized = true;
                 this.log('Authorization acknowledged by server.');
+                if (this.onAuthorized) {
+                    this.onAuthorized(msg);
+                }
+                // Send context first if we have one
+                if (this.processContext) {
+                    this.sendContextStart();
+                }
+                // Replay entities and interactions
                 for (const message of this.unacknowledgedObserved.values()) {
                     this.sendRaw(message);
                 }
-                for (const message of this.unacknowledgedObservations.values()) {
-                    this.sendRaw(message);
+                // Replay all seq-based queues
+                for (const q of Object.values(this.seqQueues)) {
+                    if (q === this.contextQueue)
+                        continue; // already sent above
+                    for (const message of q.values()) {
+                        this.sendRaw(message);
+                    }
                 }
             }
             else if (msg.type === 'ack-observed') {
                 this.unacknowledgedObserved.delete(msg.hash);
             }
-            else if (msg.type === 'ack-observations') {
-                this.unacknowledgedObservations.delete(msg.seq);
+            else if (msg.type === 'custom-metric-change') {
+                if (this.onCustomMetricChange) {
+                    this.onCustomMetricChange(msg.customMetrics);
+                }
+            }
+            else if ('seq' in msg) {
+                // Dispatch seq-based acks to the appropriate queue
+                const queue = this.seqQueues[msg.type];
+                if (queue) {
+                    queue.ack(msg.seq);
+                }
             }
         }
         catch (e) {
@@ -81,12 +141,46 @@ class WebSocketConnection {
             this.socket.send(JSON.stringify(message));
         }
     }
+    sendContextStart() {
+        if (!this.processContext)
+            return;
+        const seq = this.nextSeq();
+        const contextMsg = {
+            event: 'context-start',
+            seq,
+            timestamp: (0, util_1.hrTimeNow)(),
+            ingestionId: this.ingestionId,
+            type: 'process',
+            serviceEntityHash: this.processContext.serviceEntityHash,
+            resources: this.processContext.resources,
+        };
+        this.contextQueue.add(contextMsg);
+        this.sendRaw(contextMsg);
+    }
+    setProcessContext(serviceEntityHash, resources) {
+        this.processContext = { serviceEntityHash, resources };
+        if (this.authorized) {
+            this.sendContextStart();
+        }
+    }
     sendMessage(message) {
         if (message.event === 'new-entity' || message.event === 'new-interaction') {
             this.unacknowledgedObserved.set(message.hash, message);
         }
         else if (message.event === 'observations') {
-            this.unacknowledgedObservations.set(message.seq, message);
+            this.observationsQueue.add(message);
+        }
+        else if (message.event === 'timeseries') {
+            this.timeseriesQueue.add(message);
+        }
+        else if (message.event === 'cumulative') {
+            this.cumulativeQueue.add(message);
+        }
+        else if (message.event === 'tracespans') {
+            this.traceSpansQueue.add(message);
+        }
+        else if (message.event === 'db-query') {
+            this.dbQueryQueue.add(message);
         }
         if (this.authorized) {
             this.sendRaw(message);