@databricks/appkit 0.44.0 → 0.46.0

package/dist/stream/stream-manager.js

@@ -1,9 +1,12 @@
 import { createLogger } from "../logging/logger.js";
+import { AppKitError } from "../errors/base.js";
+import { ExecutionError } from "../errors/execution.js";
 import { EventRingBuffer } from "./buffers.js";
 import { streamDefaults } from "./defaults.js";
 import { SSEErrorCode } from "./types.js";
 import { StreamValidator } from "./validator.js";
 import { SSEWriter } from "./sse-writer.js";
+import { clearGraceTimer, clearRemovalTimer } from "./timers.js";
 import { StreamRegistry } from "./stream-registry.js";
 import { randomUUID } from "node:crypto";
 import { context } from "@opentelemetry/api";
@@ -65,7 +68,8 @@ var StreamManager = class {
 				}
 			} else this.sseWriter.writeBufferOverflowWarning(res, validEventId);
 		}
-		this._clearGraceTimer(streamEntry);
+		clearGraceTimer(streamEntry);
+		clearRemovalTimer(streamEntry);
 		streamEntry.clients.add(res);
 		streamEntry.lastAccess = Date.now();
 		const combinedSignal = this._combineSignals(streamEntry.abortController.signal, options?.userSignal);
@@ -81,14 +85,14 @@ var StreamManager = class {
 			streamEntry.clients.delete(res);
 			this.activeOperations.delete(streamOperation);
 			if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) this._scheduleGraceAbort(streamEntry);
-			if (streamEntry.isCompleted && streamEntry.clients.size === 0) setTimeout(() => {
-				if (streamEntry.clients.size === 0) this.streamRegistry.remove(streamEntry.streamId);
-			}, this.bufferTTL);
+			this._scheduleRemovalAfterTTL(streamEntry);
 		});
 		if (streamEntry.isCompleted) {
 			res.end();
 			this.activeOperations.delete(streamOperation);
 			clearInterval(heartbeat);
+			streamEntry.clients.delete(res);
+			this._scheduleRemovalAfterTTL(streamEntry);
 		}
 	}
 	async _createNewStream(res, handler, options, ownerKey) {
@@ -126,6 +130,7 @@ var StreamManager = class {
 			this.activeOperations.delete(streamOperation);
 			streamEntry.clients.delete(res);
 			if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) this._scheduleGraceAbort(streamEntry);
+			this._scheduleRemovalAfterTTL(streamEntry);
 		});
 		await this._processGeneratorInBackground(streamEntry);
 		clearInterval(heartbeat);
@@ -153,35 +158,32 @@ var StreamManager = class {
 					this._broadcastEventsToClients(streamEntry, eventId, event);
 					streamEntry.lastAccess = Date.now();
 				}
-				streamEntry.isCompleted = true;
-				this._clearGraceTimer(streamEntry);
-				this._closeAllClients(streamEntry);
-				this._cleanupStream(streamEntry);
+				this._finalizeStream(streamEntry);
 			} catch (error) {
-				const errorMsg = error instanceof Error ? error.message : "Internal server error";
+				const rawMsg = error instanceof Error ? error.message : "Internal server error";
+				const clientMsg = error instanceof AppKitError ? error.clientMessage : "Internal server error";
+				const upstreamCode = error instanceof ExecutionError ? error.errorCode : void 0;
 				const errorEventId = randomUUID();
 				const errorCode = this._categorizeError(error);
 				if (errorCode === SSEErrorCode.STREAM_ABORTED) {
 					logger.info("Stream aborted by client (code=%s)", errorCode);
-					streamEntry.isCompleted = true;
-					this._clearGraceTimer(streamEntry);
-					this._closeAllClients(streamEntry);
-					this._cleanupStream(streamEntry);
+					this._finalizeStream(streamEntry);
 					return;
 				}
-				logger.error("Stream execution failed: %s (code=%s)", errorMsg, errorCode);
+				logger.error("Stream execution failed: %s (code=%s upstreamCode=%s)", rawMsg, errorCode, upstreamCode ?? "n/a");
+				const payload = {
+					error: clientMsg,
+					code: errorCode
+				};
+				if (upstreamCode) payload.errorCode = upstreamCode;
 				streamEntry.eventBuffer.add({
 					id: errorEventId,
 					type: "error",
-					data: JSON.stringify({
-						error: errorMsg,
-						code: errorCode
-					}),
+					data: JSON.stringify(payload),
 					timestamp: Date.now()
 				});
-				this._broadcastErrorToClients(streamEntry, errorEventId, errorMsg, errorCode, true);
-				streamEntry.isCompleted = true;
-				this._clearGraceTimer(streamEntry);
+				this._broadcastErrorToClients(streamEntry, errorEventId, clientMsg, errorCode, true, upstreamCode);
+				this._finalizeStream(streamEntry);
 			}
 		});
 	}
@@ -203,17 +205,24 @@ var StreamManager = class {
 	_broadcastEventsToClients(streamEntry, eventId, event) {
 		for (const client of streamEntry.clients) if (!client.writableEnded) this.sseWriter.writeEvent(client, eventId, event);
 	}
-	_broadcastErrorToClients(streamEntry, eventId, errorMessage, errorCode, closeClients = false) {
+	_broadcastErrorToClients(streamEntry, eventId, errorMessage, errorCode, closeClients = false, upstreamCode) {
 		for (const client of streamEntry.clients) if (!client.writableEnded) {
-			this.sseWriter.writeError(client, eventId, errorMessage, errorCode);
+			this.sseWriter.writeError(client, eventId, errorMessage, errorCode, upstreamCode);
 			if (closeClients) client.end();
 		}
 	}
+	_finalizeStream(streamEntry) {
+		streamEntry.isCompleted = true;
+		clearGraceTimer(streamEntry);
+		this._closeAllClients(streamEntry);
+		this._scheduleRemovalAfterTTL(streamEntry);
+	}
 	_closeAllClients(streamEntry) {
 		for (const client of streamEntry.clients) if (!client.writableEnded) client.end();
+		streamEntry.clients.clear();
 	}
 	_scheduleGraceAbort(streamEntry) {
-		this._clearGraceTimer(streamEntry);
+		clearGraceTimer(streamEntry);
 		const timer = setTimeout(() => {
 			streamEntry.disconnectGraceTimer = void 0;
 			if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) streamEntry.abortController.abort(new DOMException("Client disconnected (grace expired)", "AbortError"));
@@ -221,16 +230,15 @@ var StreamManager = class {
 		timer.unref?.();
 		streamEntry.disconnectGraceTimer = timer;
 	}
-	_clearGraceTimer(streamEntry) {
-		if (streamEntry.disconnectGraceTimer) {
-			clearTimeout(streamEntry.disconnectGraceTimer);
-			streamEntry.disconnectGraceTimer = void 0;
-		}
-	}
-	_cleanupStream(streamEntry) {
-		if (streamEntry.clients.size === 0) setTimeout(() => {
-			if (streamEntry.clients.size === 0) this.streamRegistry.remove(streamEntry.streamId);
+	_scheduleRemovalAfterTTL(streamEntry) {
+		if (!streamEntry.isCompleted || streamEntry.clients.size > 0) return;
+		clearRemovalTimer(streamEntry);
+		streamEntry.lastAccess = Date.now();
+		streamEntry.removalTimer = setTimeout(() => {
+			streamEntry.removalTimer = void 0;
+			if (streamEntry.clients.size === 0 && Date.now() - streamEntry.lastAccess >= this.bufferTTL) this.streamRegistry.remove(streamEntry.streamId);
 		}, this.bufferTTL);
+		streamEntry.removalTimer.unref?.();
 	}
 	_categorizeError(error) {
 		if (error instanceof Error) {

package/dist/stream/stream-manager.js.map

@@ -1 +1 @@
-{"version":3,"file":"stream-manager.js","names":[],"sources":["../../src/stream/stream-manager.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport { context } from \"@opentelemetry/api\";\nimport type { IAppResponse, StreamConfig } from \"shared\";\nimport { createLogger } from \"../logging/logger\";\nimport { EventRingBuffer } from \"./buffers\";\nimport { streamDefaults } from \"./defaults\";\nimport { SSEWriter } from \"./sse-writer\";\nimport { StreamRegistry } from \"./stream-registry\";\nimport { SSEErrorCode, type StreamEntry, type StreamOperation } from \"./types\";\nimport { StreamValidator } from \"./validator\";\n\nconst logger = createLogger(\"stream\");\n\n// main entry point for Server-Sent events streaming\nexport class StreamManager {\n  private activeOperations: Set<StreamOperation>;\n  private streamRegistry: StreamRegistry;\n  private sseWriter: SSEWriter;\n  private maxEventSize: number;\n  private bufferTTL: number;\n  private disconnectGraceMs: number;\n\n  constructor(options?: StreamConfig) {\n    this.streamRegistry = new StreamRegistry(\n      options?.maxActiveStreams ?? streamDefaults.maxActiveStreams,\n    );\n    this.sseWriter = new SSEWriter();\n    this.maxEventSize = options?.maxEventSize ?? streamDefaults.maxEventSize;\n    this.bufferTTL = options?.bufferTTL ?? streamDefaults.bufferTTL;\n    this.disconnectGraceMs =\n      options?.disconnectGraceMs ?? streamDefaults.disconnectGraceMs;\n    this.activeOperations = new Set();\n  }\n\n  // main streaming method - handles new connection and reconnection\n  async stream(\n    res: IAppResponse,\n    handler: (signal: AbortSignal) => AsyncGenerator<any, void, unknown>,\n    options?: StreamConfig,\n    ownerKey?: string,\n  ): Promise<void> {\n    const { streamId } = options || {};\n\n    // check if response is already closed\n    if (res.writableEnded || res.destroyed) {\n      return;\n    }\n\n    // setup SSE headers\n    this.sseWriter.setupHeaders(res);\n\n    // handle reconnection\n    if (streamId && StreamValidator.validateStreamId(streamId)) {\n      const existingStream = this.streamRegistry.get(streamId);\n      if (existingStream) {\n        // Enforce per-user binding: the stream's owner key must match the\n        // requesting caller's owner key. This prevents cross-user stream\n        // takeover via guessed/leaked stream IDs (the SSE registry was\n        // previously a global lookup with no authorization step).\n        if (existingStream.ownerKey !== ownerKey) {\n          this.sseWriter.writeError(\n            res,\n            randomUUID(),\n            \"Stream not found or access denied\",\n            SSEErrorCode.STREAM_FORBIDDEN,\n          );\n          res.end();\n          return;\n        }\n        return this._attachToExistingStream(res, existingStream, options);\n      }\n    }\n\n    // if stream does not exist, create a new one\n    return this._createNewStream(res, handler, options, ownerKey);\n  }\n\n  // abort all active operations\n  abortAll(): void {\n    // pending disconnect-grace timers are cleared by streamRegistry.clear() below\n    this.activeOperations.forEach((operation) => {\n      if (operation.heartbeat) clearInterval(operation.heartbeat);\n      operation.controller.abort(\n        new DOMException(\"Server shutdown\", \"AbortError\"),\n      );\n    });\n    this.activeOperations.clear();\n    this.streamRegistry.clear();\n  }\n\n  // get the number of active operations\n  getActiveCount(): number {\n    return this.activeOperations.size;\n  }\n\n  // attach to existing stream\n  private async _attachToExistingStream(\n    res: IAppResponse,\n    streamEntry: StreamEntry,\n    options?: StreamConfig,\n  ): Promise<void> {\n    // handle reconnection - replay missed events\n    const lastEventId = res.req?.headers[\"last-event-id\"];\n\n    if (StreamValidator.validateEventId(lastEventId)) {\n      // cast to string after validation\n      const validEventId = lastEventId as string;\n      if (streamEntry.eventBuffer.has(validEventId)) {\n        const missedEvents =\n          streamEntry.eventBuffer.getEventsSince(validEventId);\n        // broadcast missed events to client\n        for (const event of missedEvents) {\n          if (options?.userSignal?.aborted) break;\n          this.sseWriter.writeBufferedEvent(res, event);\n        }\n      } else {\n        // buffer overflow - send warning\n        this.sseWriter.writeBufferOverflowWarning(res, validEventId);\n      }\n    }\n\n    // a reconnecting client cancels the pending disconnect-grace abort\n    this._clearGraceTimer(streamEntry);\n\n    // add client to stream entry\n    streamEntry.clients.add(res);\n    streamEntry.lastAccess = Date.now();\n\n    // start heartbeat\n    const combinedSignal = this._combineSignals(\n      streamEntry.abortController.signal,\n      options?.userSignal,\n    );\n    const heartbeat = this.sseWriter.startHeartbeat(res, combinedSignal);\n\n    // track operation\n    const streamOperation: StreamOperation = {\n      controller: streamEntry.abortController,\n      type: \"stream\",\n      heartbeat,\n    };\n    this.activeOperations.add(streamOperation);\n\n    // handle client disconnect\n    res.on(\"close\", () => {\n      clearInterval(heartbeat);\n      streamEntry.clients.delete(res);\n      this.activeOperations.delete(streamOperation);\n\n      // grace-abort instead of aborting now, so a reconnect can resume\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        this._scheduleGraceAbort(streamEntry);\n      }\n\n      // cleanup if stream is completed and no clients are connected\n      if (streamEntry.isCompleted && streamEntry.clients.size === 0) {\n        setTimeout(() => {\n          if (streamEntry.clients.size === 0) {\n            this.streamRegistry.remove(streamEntry.streamId);\n          }\n        }, this.bufferTTL);\n      }\n    });\n\n    // if stream is completed, close connection\n    if (streamEntry.isCompleted) {\n      res.end();\n      // cleanup operation\n      this.activeOperations.delete(streamOperation);\n      clearInterval(heartbeat);\n    }\n  }\n  private async _createNewStream(\n    res: IAppResponse,\n    handler: (signal: AbortSignal) => AsyncGenerator<any, void, unknown>,\n    options?: StreamConfig,\n    ownerKey?: string,\n  ): Promise<void> {\n    const streamId = options?.streamId ?? randomUUID();\n\n    // abort stream if response is closed\n    if (res.writableEnded || res.destroyed) {\n      return;\n    }\n\n    const abortController = new AbortController();\n\n    // create event buffer\n    const eventBuffer = new EventRingBuffer(\n      options?.bufferSize ?? streamDefaults.bufferSize,\n    );\n\n    // setup signals and heartbeat\n    const combinedSignal = this._combineSignals(\n      abortController.signal,\n      options?.userSignal,\n    );\n    const heartbeat = this.sseWriter.startHeartbeat(res, combinedSignal);\n\n    // capture the current trace context at stream creation time\n    const traceContext = context.active();\n\n    // abort stream if response is closed\n    if (res.writableEnded || res.destroyed) {\n      clearInterval(heartbeat);\n      return;\n    }\n\n    // create stream entry\n    const streamEntry: StreamEntry = {\n      streamId,\n      ownerKey,\n      generator: handler(combinedSignal),\n      eventBuffer,\n      clients: new Set([res]),\n      isCompleted: false,\n      lastAccess: Date.now(),\n      abortController,\n      traceContext,\n    };\n    this.streamRegistry.add(streamEntry);\n\n    // track operation\n    const streamOperation: StreamOperation = {\n      controller: abortController,\n      type: \"stream\",\n      heartbeat,\n    };\n    this.activeOperations.add(streamOperation);\n\n    res.on(\"close\", () => {\n      clearInterval(heartbeat);\n      this.activeOperations.delete(streamOperation);\n      streamEntry.clients.delete(res);\n\n      // grace-abort instead of aborting now, so a reconnect can resume\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        this._scheduleGraceAbort(streamEntry);\n      }\n    });\n\n    await this._processGeneratorInBackground(streamEntry);\n\n    // cleanup\n    clearInterval(heartbeat);\n    this.activeOperations.delete(streamOperation);\n  }\n\n  private async _processGeneratorInBackground(\n    streamEntry: StreamEntry,\n  ): Promise<void> {\n    // run the entire generator processing within the captured trace context\n    return context.with(streamEntry.traceContext, async () => {\n      try {\n        // retrieve all events from generator\n        for await (const event of streamEntry.generator) {\n          if (streamEntry.abortController.signal.aborted) break;\n          const eventId = randomUUID();\n          const eventData = JSON.stringify(event);\n\n          // validate event size\n          if (eventData.length > this.maxEventSize) {\n            const errorMsg = `Event exceeds max size of ${this.maxEventSize} bytes`;\n            const errorCode = SSEErrorCode.INVALID_REQUEST;\n            // broadcast error to all connected clients\n            this._broadcastErrorToClients(\n              streamEntry,\n              eventId,\n              errorMsg,\n              errorCode,\n            );\n            continue;\n          }\n\n          // buffer event for reconnection\n          streamEntry.eventBuffer.add({\n            id: eventId,\n            type: event.type,\n            data: eventData,\n            timestamp: Date.now(),\n          });\n\n          // broadcast to all connected clients\n          this._broadcastEventsToClients(streamEntry, eventId, event);\n          streamEntry.lastAccess = Date.now();\n        }\n\n        streamEntry.isCompleted = true;\n\n        // no late grace abort should fire on a completed stream\n        this._clearGraceTimer(streamEntry);\n\n        // close all clients\n        this._closeAllClients(streamEntry);\n\n        // cleanup if no clients are connected\n        this._cleanupStream(streamEntry);\n      } catch (error) {\n        const errorMsg =\n          error instanceof Error ? error.message : \"Internal server error\";\n        const errorEventId = randomUUID();\n        const errorCode = this._categorizeError(error);\n\n        // client cancellation is a normal control-flow signal, not a failure\n        if (errorCode === SSEErrorCode.STREAM_ABORTED) {\n          logger.info(\"Stream aborted by client (code=%s)\", errorCode);\n          streamEntry.isCompleted = true;\n          this._clearGraceTimer(streamEntry);\n          this._closeAllClients(streamEntry);\n          this._cleanupStream(streamEntry);\n          return;\n        }\n\n        logger.error(\n          \"Stream execution failed: %s (code=%s)\",\n          errorMsg,\n          errorCode,\n        );\n\n        // buffer error event\n        streamEntry.eventBuffer.add({\n          id: errorEventId,\n          type: \"error\",\n          data: JSON.stringify({ error: errorMsg, code: errorCode }),\n          timestamp: Date.now(),\n        });\n\n        // send error event to all connected clients\n        this._broadcastErrorToClients(\n          streamEntry,\n          errorEventId,\n          errorMsg,\n          errorCode,\n          true,\n        );\n        streamEntry.isCompleted = true;\n        this._clearGraceTimer(streamEntry);\n      }\n    });\n  }\n\n  private _combineSignals(\n    internalSignal?: AbortSignal,\n    userSignal?: AbortSignal,\n  ): AbortSignal {\n    if (!userSignal) return internalSignal || new AbortController().signal;\n\n    const signals = [internalSignal, userSignal].filter(\n      Boolean,\n    ) as AbortSignal[];\n    const controller = new AbortController();\n\n    signals.forEach((signal) => {\n      if (signal?.aborted) {\n        controller.abort(signal.reason);\n        return;\n      }\n\n      signal?.addEventListener(\n        \"abort\",\n        () => {\n          controller.abort(signal.reason);\n        },\n        { once: true },\n      );\n    });\n    return controller.signal;\n  }\n\n  // broadcast events to all connected clients\n  private _broadcastEventsToClients(\n    streamEntry: StreamEntry,\n    eventId: string,\n    event: any,\n  ): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        this.sseWriter.writeEvent(client, eventId, event);\n      }\n    }\n  }\n\n  // broadcast error to all connected clients\n  private _broadcastErrorToClients(\n    streamEntry: StreamEntry,\n    eventId: string,\n    errorMessage: string,\n    errorCode: SSEErrorCode,\n    closeClients: boolean = false,\n  ): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        this.sseWriter.writeError(client, eventId, errorMessage, errorCode);\n        if (closeClients) {\n          client.end();\n        }\n      }\n    }\n  }\n\n  // close all connected clients\n  private _closeAllClients(streamEntry: StreamEntry): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        client.end();\n      }\n    }\n  }\n\n  // abort the generator after the grace window unless a client reconnects first\n  private _scheduleGraceAbort(streamEntry: StreamEntry): void {\n    // clear any existing timer to avoid stacking\n    this._clearGraceTimer(streamEntry);\n\n    const timer = setTimeout(() => {\n      streamEntry.disconnectGraceTimer = undefined;\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        streamEntry.abortController.abort(\n          new DOMException(\"Client disconnected (grace expired)\", \"AbortError\"),\n        );\n      }\n    }, this.disconnectGraceMs);\n\n    // never keep the process alive solely for a grace timer\n    timer.unref?.();\n    streamEntry.disconnectGraceTimer = timer;\n  }\n\n  // clear a pending disconnect-grace timer, if any\n  private _clearGraceTimer(streamEntry: StreamEntry): void {\n    if (streamEntry.disconnectGraceTimer) {\n      clearTimeout(streamEntry.disconnectGraceTimer);\n      streamEntry.disconnectGraceTimer = undefined;\n    }\n  }\n\n  // cleanup stream if no clients are connected\n  private _cleanupStream(streamEntry: StreamEntry): void {\n    if (streamEntry.clients.size === 0) {\n      setTimeout(() => {\n        if (streamEntry.clients.size === 0) {\n          this.streamRegistry.remove(streamEntry.streamId);\n        }\n      }, this.bufferTTL);\n    }\n  }\n\n  private _categorizeError(error: unknown): SSEErrorCode {\n    if (error instanceof Error) {\n      const message = error.message.toLowerCase();\n      if (message.includes(\"timeout\") || message.includes(\"timed out\")) {\n        return SSEErrorCode.TIMEOUT;\n      }\n\n      if (message.includes(\"unavailable\") || message.includes(\"econnrefused\")) {\n        return SSEErrorCode.TEMPORARY_UNAVAILABLE;\n      }\n\n      if (error.name === \"AbortError\") {\n        return SSEErrorCode.STREAM_ABORTED;\n      }\n\n      // Defense-in-depth: upstream layers (SQL client, cache) may wrap an\n      // AbortError into ExecutionError, losing `name` but keeping the message.\n      if (\n        message.includes(\"operation was aborted\") ||\n        message.includes(\"the request was aborted\") ||\n        message.includes(\"statement was canceled\")\n      ) {\n        return SSEErrorCode.STREAM_ABORTED;\n      }\n\n      // Detect upstream API errors (e.g., from Databricks SDK ApiError)\n      if (\n        \"statusCode\" in error &&\n        typeof (error as any).statusCode === \"number\"\n      ) {\n        return SSEErrorCode.UPSTREAM_ERROR;\n      }\n    }\n\n    return SSEErrorCode.INTERNAL_ERROR;\n  }\n}\n"],"mappings":";;;;;;;;;;;AAWA,MAAM,SAAS,aAAa,SAAS;AAGrC,IAAa,gBAAb,MAA2B;CACzB,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CAER,YAAY,SAAwB;AAClC,OAAK,iBAAiB,IAAI,eACxB,SAAS,oBAAoB,eAAe,iBAC7C;AACD,OAAK,YAAY,IAAI,WAAW;AAChC,OAAK,eAAe,SAAS,gBAAgB,eAAe;AAC5D,OAAK,YAAY,SAAS,aAAa,eAAe;AACtD,OAAK,oBACH,SAAS,qBAAqB,eAAe;AAC/C,OAAK,mCAAmB,IAAI,KAAK;;CAInC,MAAM,OACJ,KACA,SACA,SACA,UACe;EACf,MAAM,EAAE,aAAa,WAAW,EAAE;AAGlC,MAAI,IAAI,iBAAiB,IAAI,UAC3B;AAIF,OAAK,UAAU,aAAa,IAAI;AAGhC,MAAI,YAAY,gBAAgB,iBAAiB,SAAS,EAAE;GAC1D,MAAM,iBAAiB,KAAK,eAAe,IAAI,SAAS;AACxD,OAAI,gBAAgB;AAKlB,QAAI,eAAe,aAAa,UAAU;AACxC,UAAK,UAAU,WACb,KACA,YAAY,EACZ,qCACA,aAAa,iBACd;AACD,SAAI,KAAK;AACT;;AAEF,WAAO,KAAK,wBAAwB,KAAK,gBAAgB,QAAQ;;;AAKrE,SAAO,KAAK,iBAAiB,KAAK,SAAS,SAAS,SAAS;;CAI/D,WAAiB;AAEf,OAAK,iBAAiB,SAAS,cAAc;AAC3C,OAAI,UAAU,UAAW,eAAc,UAAU,UAAU;AAC3D,aAAU,WAAW,MACnB,IAAI,aAAa,mBAAmB,aAAa,CAClD;IACD;AACF,OAAK,iBAAiB,OAAO;AAC7B,OAAK,eAAe,OAAO;;CAI7B,iBAAyB;AACvB,SAAO,KAAK,iBAAiB;;CAI/B,MAAc,wBACZ,KACA,aACA,SACe;EAEf,MAAM,cAAc,IAAI,KAAK,QAAQ;AAErC,MAAI,gBAAgB,gBAAgB,YAAY,EAAE;GAEhD,MAAM,eAAe;AACrB,OAAI,YAAY,YAAY,IAAI,aAAa,EAAE;IAC7C,MAAM,eACJ,YAAY,YAAY,eAAe,aAAa;AAEtD,SAAK,MAAM,SAAS,cAAc;AAChC,SAAI,SAAS,YAAY,QAAS;AAClC,UAAK,UAAU,mBAAmB,KAAK,MAAM;;SAI/C,MAAK,UAAU,2BAA2B,KAAK,aAAa;;AAKhE,OAAK,iBAAiB,YAAY;AAGlC,cAAY,QAAQ,IAAI,IAAI;AAC5B,cAAY,aAAa,KAAK,KAAK;EAGnC,MAAM,iBAAiB,KAAK,gBAC1B,YAAY,gBAAgB,QAC5B,SAAS,WACV;EACD,MAAM,YAAY,KAAK,UAAU,eAAe,KAAK,eAAe;EAGpE,MAAM,kBAAmC;GACvC,YAAY,YAAY;GACxB,MAAM;GACN;GACD;AACD,OAAK,iBAAiB,IAAI,gBAAgB;AAG1C,MAAI,GAAG,eAAe;AACpB,iBAAc,UAAU;AACxB,eAAY,QAAQ,OAAO,IAAI;AAC/B,QAAK,iBAAiB,OAAO,gBAAgB;AAG7C,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,MAAK,oBAAoB,YAAY;AAIvC,OAAI,YAAY,eAAe,YAAY,QAAQ,SAAS,EAC1D,kBAAiB;AACf,QAAI,YAAY,QAAQ,SAAS,EAC/B,MAAK,eAAe,OAAO,YAAY,SAAS;MAEjD,KAAK,UAAU;IAEpB;AAGF,MAAI,YAAY,aAAa;AAC3B,OAAI,KAAK;AAET,QAAK,iBAAiB,OAAO,gBAAgB;AAC7C,iBAAc,UAAU;;;CAG5B,MAAc,iBACZ,KACA,SACA,SACA,UACe;EACf,MAAM,WAAW,SAAS,YAAY,YAAY;AAGlD,MAAI,IAAI,iBAAiB,IAAI,UAC3B;EAGF,MAAM,kBAAkB,IAAI,iBAAiB;EAG7C,MAAM,cAAc,IAAI,gBACtB,SAAS,cAAc,eAAe,WACvC;EAGD,MAAM,iBAAiB,KAAK,gBAC1B,gBAAgB,QAChB,SAAS,WACV;EACD,MAAM,YAAY,KAAK,UAAU,eAAe,KAAK,eAAe;EAGpE,MAAM,eAAe,QAAQ,QAAQ;AAGrC,MAAI,IAAI,iBAAiB,IAAI,WAAW;AACtC,iBAAc,UAAU;AACxB;;EAIF,MAAM,cAA2B;GAC/B;GACA;GACA,WAAW,QAAQ,eAAe;GAClC;GACA,SAAS,IAAI,IAAI,CAAC,IAAI,CAAC;GACvB,aAAa;GACb,YAAY,KAAK,KAAK;GACtB;GACA;GACD;AACD,OAAK,eAAe,IAAI,YAAY;EAGpC,MAAM,kBAAmC;GACvC,YAAY;GACZ,MAAM;GACN;GACD;AACD,OAAK,iBAAiB,IAAI,gBAAgB;AAE1C,MAAI,GAAG,eAAe;AACpB,iBAAc,UAAU;AACxB,QAAK,iBAAiB,OAAO,gBAAgB;AAC7C,eAAY,QAAQ,OAAO,IAAI;AAG/B,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,MAAK,oBAAoB,YAAY;IAEvC;AAEF,QAAM,KAAK,8BAA8B,YAAY;AAGrD,gBAAc,UAAU;AACxB,OAAK,iBAAiB,OAAO,gBAAgB;;CAG/C,MAAc,8BACZ,aACe;AAEf,SAAO,QAAQ,KAAK,YAAY,cAAc,YAAY;AACxD,OAAI;AAEF,eAAW,MAAM,SAAS,YAAY,WAAW;AAC/C,SAAI,YAAY,gBAAgB,OAAO,QAAS;KAChD,MAAM,UAAU,YAAY;KAC5B,MAAM,YAAY,KAAK,UAAU,MAAM;AAGvC,SAAI,UAAU,SAAS,KAAK,cAAc;MACxC,MAAM,WAAW,6BAA6B,KAAK,aAAa;MAChE,MAAM,YAAY,aAAa;AAE/B,WAAK,yBACH,aACA,SACA,UACA,UACD;AACD;;AAIF,iBAAY,YAAY,IAAI;MAC1B,IAAI;MACJ,MAAM,MAAM;MACZ,MAAM;MACN,WAAW,KAAK,KAAK;MACtB,CAAC;AAGF,UAAK,0BAA0B,aAAa,SAAS,MAAM;AAC3D,iBAAY,aAAa,KAAK,KAAK;;AAGrC,gBAAY,cAAc;AAG1B,SAAK,iBAAiB,YAAY;AAGlC,SAAK,iBAAiB,YAAY;AAGlC,SAAK,eAAe,YAAY;YACzB,OAAO;IACd,MAAM,WACJ,iBAAiB,QAAQ,MAAM,UAAU;IAC3C,MAAM,eAAe,YAAY;IACjC,MAAM,YAAY,KAAK,iBAAiB,MAAM;AAG9C,QAAI,cAAc,aAAa,gBAAgB;AAC7C,YAAO,KAAK,sCAAsC,UAAU;AAC5D,iBAAY,cAAc;AAC1B,UAAK,iBAAiB,YAAY;AAClC,UAAK,iBAAiB,YAAY;AAClC,UAAK,eAAe,YAAY;AAChC;;AAGF,WAAO,MACL,yCACA,UACA,UACD;AAGD,gBAAY,YAAY,IAAI;KAC1B,IAAI;KACJ,MAAM;KACN,MAAM,KAAK,UAAU;MAAE,OAAO;MAAU,MAAM;MAAW,CAAC;KAC1D,WAAW,KAAK,KAAK;KACtB,CAAC;AAGF,SAAK,yBACH,aACA,cACA,UACA,WACA,KACD;AACD,gBAAY,cAAc;AAC1B,SAAK,iBAAiB,YAAY;;IAEpC;;CAGJ,AAAQ,gBACN,gBACA,YACa;AACb,MAAI,CAAC,WAAY,QAAO,kBAAkB,IAAI,iBAAiB,CAAC;EAEhE,MAAM,UAAU,CAAC,gBAAgB,WAAW,CAAC,OAC3C,QACD;EACD,MAAM,aAAa,IAAI,iBAAiB;AAExC,UAAQ,SAAS,WAAW;AAC1B,OAAI,QAAQ,SAAS;AACnB,eAAW,MAAM,OAAO,OAAO;AAC/B;;AAGF,WAAQ,iBACN,eACM;AACJ,eAAW,MAAM,OAAO,OAAO;MAEjC,EAAE,MAAM,MAAM,CACf;IACD;AACF,SAAO,WAAW;;CAIpB,AAAQ,0BACN,aACA,SACA,OACM;AACN,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,cACV,MAAK,UAAU,WAAW,QAAQ,SAAS,MAAM;;CAMvD,AAAQ,yBACN,aACA,SACA,cACA,WACA,eAAwB,OAClB;AACN,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,eAAe;AACzB,QAAK,UAAU,WAAW,QAAQ,SAAS,cAAc,UAAU;AACnE,OAAI,aACF,QAAO,KAAK;;;CAOpB,AAAQ,iBAAiB,aAAgC;AACvD,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,cACV,QAAO,KAAK;;CAMlB,AAAQ,oBAAoB,aAAgC;AAE1D,OAAK,iBAAiB,YAAY;EAElC,MAAM,QAAQ,iBAAiB;AAC7B,eAAY,uBAAuB;AACnC,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,aAAY,gBAAgB,MAC1B,IAAI,aAAa,uCAAuC,aAAa,CACtE;KAEF,KAAK,kBAAkB;AAG1B,QAAM,SAAS;AACf,cAAY,uBAAuB;;CAIrC,AAAQ,iBAAiB,aAAgC;AACvD,MAAI,YAAY,sBAAsB;AACpC,gBAAa,YAAY,qBAAqB;AAC9C,eAAY,uBAAuB;;;CAKvC,AAAQ,eAAe,aAAgC;AACrD,MAAI,YAAY,QAAQ,SAAS,EAC/B,kBAAiB;AACf,OAAI,YAAY,QAAQ,SAAS,EAC/B,MAAK,eAAe,OAAO,YAAY,SAAS;KAEjD,KAAK,UAAU;;CAItB,AAAQ,iBAAiB,OAA8B;AACrD,MAAI,iBAAiB,OAAO;GAC1B,MAAM,UAAU,MAAM,QAAQ,aAAa;AAC3C,OAAI,QAAQ,SAAS,UAAU,IAAI,QAAQ,SAAS,YAAY,CAC9D,QAAO,aAAa;AAGtB,OAAI,QAAQ,SAAS,cAAc,IAAI,QAAQ,SAAS,eAAe,CACrE,QAAO,aAAa;AAGtB,OAAI,MAAM,SAAS,aACjB,QAAO,aAAa;AAKtB,OACE,QAAQ,SAAS,wBAAwB,IACzC,QAAQ,SAAS,0BAA0B,IAC3C,QAAQ,SAAS,yBAAyB,CAE1C,QAAO,aAAa;AAItB,OACE,gBAAgB,SAChB,OAAQ,MAAc,eAAe,SAErC,QAAO,aAAa;;AAIxB,SAAO,aAAa"}
+{"version":3,"file":"stream-manager.js","names":[],"sources":["../../src/stream/stream-manager.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport { context } from \"@opentelemetry/api\";\nimport type { IAppResponse, StreamConfig } from \"shared\";\nimport { AppKitError } from \"../errors/base\";\nimport { ExecutionError } from \"../errors/execution\";\nimport { createLogger } from \"../logging/logger\";\nimport { EventRingBuffer } from \"./buffers\";\nimport { streamDefaults } from \"./defaults\";\nimport { SSEWriter } from \"./sse-writer\";\nimport { StreamRegistry } from \"./stream-registry\";\nimport { clearGraceTimer, clearRemovalTimer } from \"./timers\";\nimport { SSEErrorCode, type StreamEntry, type StreamOperation } from \"./types\";\nimport { StreamValidator } from \"./validator\";\n\nconst logger = createLogger(\"stream\");\n\n// main entry point for Server-Sent events streaming\nexport class StreamManager {\n  private activeOperations: Set<StreamOperation>;\n  private streamRegistry: StreamRegistry;\n  private sseWriter: SSEWriter;\n  private maxEventSize: number;\n  private bufferTTL: number;\n  private disconnectGraceMs: number;\n\n  constructor(options?: StreamConfig) {\n    this.streamRegistry = new StreamRegistry(\n      options?.maxActiveStreams ?? streamDefaults.maxActiveStreams,\n    );\n    this.sseWriter = new SSEWriter();\n    this.maxEventSize = options?.maxEventSize ?? streamDefaults.maxEventSize;\n    this.bufferTTL = options?.bufferTTL ?? streamDefaults.bufferTTL;\n    this.disconnectGraceMs =\n      options?.disconnectGraceMs ?? streamDefaults.disconnectGraceMs;\n    this.activeOperations = new Set();\n  }\n\n  // main streaming method - handles new connection and reconnection\n  async stream(\n    res: IAppResponse,\n    handler: (signal: AbortSignal) => AsyncGenerator<any, void, unknown>,\n    options?: StreamConfig,\n    ownerKey?: string,\n  ): Promise<void> {\n    const { streamId } = options || {};\n\n    // check if response is already closed\n    if (res.writableEnded || res.destroyed) {\n      return;\n    }\n\n    // setup SSE headers\n    this.sseWriter.setupHeaders(res);\n\n    // handle reconnection\n    if (streamId && StreamValidator.validateStreamId(streamId)) {\n      const existingStream = this.streamRegistry.get(streamId);\n      if (existingStream) {\n        // Enforce per-user binding: the stream's owner key must match the\n        // requesting caller's owner key. This prevents cross-user stream\n        // takeover via guessed/leaked stream IDs (the SSE registry was\n        // previously a global lookup with no authorization step).\n        if (existingStream.ownerKey !== ownerKey) {\n          this.sseWriter.writeError(\n            res,\n            randomUUID(),\n            \"Stream not found or access denied\",\n            SSEErrorCode.STREAM_FORBIDDEN,\n          );\n          res.end();\n          return;\n        }\n        return this._attachToExistingStream(res, existingStream, options);\n      }\n    }\n\n    // if stream does not exist, create a new one\n    return this._createNewStream(res, handler, options, ownerKey);\n  }\n\n  // abort all active operations\n  abortAll(): void {\n    // pending disconnect-grace timers are cleared by streamRegistry.clear() below\n    this.activeOperations.forEach((operation) => {\n      if (operation.heartbeat) clearInterval(operation.heartbeat);\n      operation.controller.abort(\n        new DOMException(\"Server shutdown\", \"AbortError\"),\n      );\n    });\n    this.activeOperations.clear();\n    this.streamRegistry.clear();\n  }\n\n  // get the number of active operations\n  getActiveCount(): number {\n    return this.activeOperations.size;\n  }\n\n  // attach to existing stream\n  private async _attachToExistingStream(\n    res: IAppResponse,\n    streamEntry: StreamEntry,\n    options?: StreamConfig,\n  ): Promise<void> {\n    // handle reconnection - replay missed events\n    const lastEventId = res.req?.headers[\"last-event-id\"];\n\n    if (StreamValidator.validateEventId(lastEventId)) {\n      // cast to string after validation\n      const validEventId = lastEventId as string;\n      if (streamEntry.eventBuffer.has(validEventId)) {\n        const missedEvents =\n          streamEntry.eventBuffer.getEventsSince(validEventId);\n        // broadcast missed events to client\n        for (const event of missedEvents) {\n          if (options?.userSignal?.aborted) break;\n          this.sseWriter.writeBufferedEvent(res, event);\n        }\n      } else {\n        // buffer overflow - send warning\n        this.sseWriter.writeBufferOverflowWarning(res, validEventId);\n      }\n    }\n\n    // a reconnecting client cancels the pending disconnect-grace abort\n    clearGraceTimer(streamEntry);\n\n    // a reconnect cancels any pending registry removal so the entry isn't\n    // pulled out from under the newly attached client\n    clearRemovalTimer(streamEntry);\n\n    // add client to stream entry\n    streamEntry.clients.add(res);\n    streamEntry.lastAccess = Date.now();\n\n    // start heartbeat\n    const combinedSignal = this._combineSignals(\n      streamEntry.abortController.signal,\n      options?.userSignal,\n    );\n    const heartbeat = this.sseWriter.startHeartbeat(res, combinedSignal);\n\n    // track operation\n    const streamOperation: StreamOperation = {\n      controller: streamEntry.abortController,\n      type: \"stream\",\n      heartbeat,\n    };\n    this.activeOperations.add(streamOperation);\n\n    // handle client disconnect\n    res.on(\"close\", () => {\n      clearInterval(heartbeat);\n      streamEntry.clients.delete(res);\n      this.activeOperations.delete(streamOperation);\n\n      // grace-abort instead of aborting now, so a reconnect can resume\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        this._scheduleGraceAbort(streamEntry);\n      }\n\n      // cleanup if stream is completed and no clients are connected\n      this._scheduleRemovalAfterTTL(streamEntry);\n    });\n\n    // if stream is completed, close connection\n    if (streamEntry.isCompleted) {\n      res.end();\n      // cleanup operation\n      this.activeOperations.delete(streamOperation);\n      clearInterval(heartbeat);\n      // we deliberately ended this client, so drop it from the entry and\n      // schedule removal now instead of relying on the transport's `close`\n      // event (the close handler's later delete is a safe no-op)\n      streamEntry.clients.delete(res);\n      this._scheduleRemovalAfterTTL(streamEntry);\n    }\n  }\n  private async _createNewStream(\n    res: IAppResponse,\n    handler: (signal: AbortSignal) => AsyncGenerator<any, void, unknown>,\n    options?: StreamConfig,\n    ownerKey?: string,\n  ): Promise<void> {\n    const streamId = options?.streamId ?? randomUUID();\n\n    // abort stream if response is closed\n    if (res.writableEnded || res.destroyed) {\n      return;\n    }\n\n    const abortController = new AbortController();\n\n    // create event buffer\n    const eventBuffer = new EventRingBuffer(\n      options?.bufferSize ?? streamDefaults.bufferSize,\n    );\n\n    // setup signals and heartbeat\n    const combinedSignal = this._combineSignals(\n      abortController.signal,\n      options?.userSignal,\n    );\n    const heartbeat = this.sseWriter.startHeartbeat(res, combinedSignal);\n\n    // capture the current trace context at stream creation time\n    const traceContext = context.active();\n\n    // abort stream if response is closed\n    if (res.writableEnded || res.destroyed) {\n      clearInterval(heartbeat);\n      return;\n    }\n\n    // create stream entry\n    const streamEntry: StreamEntry = {\n      streamId,\n      ownerKey,\n      generator: handler(combinedSignal),\n      eventBuffer,\n      clients: new Set([res]),\n      isCompleted: false,\n      lastAccess: Date.now(),\n      abortController,\n      traceContext,\n    };\n    this.streamRegistry.add(streamEntry);\n\n    // track operation\n    const streamOperation: StreamOperation = {\n      controller: abortController,\n      type: \"stream\",\n      heartbeat,\n    };\n    this.activeOperations.add(streamOperation);\n\n    res.on(\"close\", () => {\n      clearInterval(heartbeat);\n      this.activeOperations.delete(streamOperation);\n      streamEntry.clients.delete(res);\n\n      // grace-abort instead of aborting now, so a reconnect can resume\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        this._scheduleGraceAbort(streamEntry);\n      }\n\n      // if the stream already finished (completed or errored), schedule\n      // registry removal once the buffer TTL elapses so completed streams\n      // don't accumulate in the registry forever\n      this._scheduleRemovalAfterTTL(streamEntry);\n    });\n\n    await this._processGeneratorInBackground(streamEntry);\n\n    // cleanup\n    clearInterval(heartbeat);\n    this.activeOperations.delete(streamOperation);\n  }\n\n  private async _processGeneratorInBackground(\n    streamEntry: StreamEntry,\n  ): Promise<void> {\n    // run the entire generator processing within the captured trace context\n    return context.with(streamEntry.traceContext, async () => {\n      try {\n        // retrieve all events from generator\n        for await (const event of streamEntry.generator) {\n          if (streamEntry.abortController.signal.aborted) break;\n          const eventId = randomUUID();\n          const eventData = JSON.stringify(event);\n\n          // validate event size\n          if (eventData.length > this.maxEventSize) {\n            const errorMsg = `Event exceeds max size of ${this.maxEventSize} bytes`;\n            const errorCode = SSEErrorCode.INVALID_REQUEST;\n            // broadcast error to all connected clients\n            this._broadcastErrorToClients(\n              streamEntry,\n              eventId,\n              errorMsg,\n              errorCode,\n            );\n            continue;\n          }\n\n          // buffer event for reconnection\n          streamEntry.eventBuffer.add({\n            id: eventId,\n            type: event.type,\n            data: eventData,\n            timestamp: Date.now(),\n          });\n\n          // broadcast to all connected clients\n          this._broadcastEventsToClients(streamEntry, eventId, event);\n          streamEntry.lastAccess = Date.now();\n        }\n\n        this._finalizeStream(streamEntry);\n      } catch (error) {\n        // Two distinct messages: a *raw* one for server-side logs (full\n        // detail, statement fragments, correlation IDs) and a *client*\n        // one for the SSE payload (sanitized, stable, safe to render in\n        // a UI). Mixing them leaks upstream wording to anyone connected\n        // to the stream — see CWE-209.\n        const rawMsg =\n          error instanceof Error ? error.message : \"Internal server error\";\n        const clientMsg =\n          error instanceof AppKitError\n            ? error.clientMessage\n            : \"Internal server error\";\n        // Upstream structured code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK,\n        // NOT_IMPLEMENTED). UI should branch on this, not on `error`.\n        const upstreamCode =\n          error instanceof ExecutionError ? error.errorCode : undefined;\n        const errorEventId = randomUUID();\n        const errorCode = this._categorizeError(error);\n\n        // client cancellation is a normal control-flow signal, not a failure\n        if (errorCode === SSEErrorCode.STREAM_ABORTED) {\n          logger.info(\"Stream aborted by client (code=%s)\", errorCode);\n          this._finalizeStream(streamEntry);\n          return;\n        }\n\n        logger.error(\n          \"Stream execution failed: %s (code=%s upstreamCode=%s)\",\n          rawMsg,\n          errorCode,\n          upstreamCode ?? \"n/a\",\n        );\n\n        const payload: Record<string, unknown> = {\n          error: clientMsg,\n          code: errorCode,\n        };\n        if (upstreamCode) payload.errorCode = upstreamCode;\n\n        // buffer error event\n        streamEntry.eventBuffer.add({\n          id: errorEventId,\n          type: \"error\",\n          data: JSON.stringify(payload),\n          timestamp: Date.now(),\n        });\n\n        // send error event to all connected clients\n        this._broadcastErrorToClients(\n          streamEntry,\n          errorEventId,\n          clientMsg,\n          errorCode,\n          true,\n          upstreamCode,\n        );\n        // the broadcast above already ended the connected clients; finalize\n        // still detaches them so removal is scheduled without waiting on a\n        // transport `close` event\n        this._finalizeStream(streamEntry);\n      }\n    });\n  }\n\n  private _combineSignals(\n    internalSignal?: AbortSignal,\n    userSignal?: AbortSignal,\n  ): AbortSignal {\n    if (!userSignal) return internalSignal || new AbortController().signal;\n\n    const signals = [internalSignal, userSignal].filter(\n      Boolean,\n    ) as AbortSignal[];\n    const controller = new AbortController();\n\n    signals.forEach((signal) => {\n      if (signal?.aborted) {\n        controller.abort(signal.reason);\n        return;\n      }\n\n      signal?.addEventListener(\n        \"abort\",\n        () => {\n          controller.abort(signal.reason);\n        },\n        { once: true },\n      );\n    });\n    return controller.signal;\n  }\n\n  // broadcast events to all connected clients\n  private _broadcastEventsToClients(\n    streamEntry: StreamEntry,\n    eventId: string,\n    event: any,\n  ): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        this.sseWriter.writeEvent(client, eventId, event);\n      }\n    }\n  }\n\n  // broadcast error to all connected clients\n  private _broadcastErrorToClients(\n    streamEntry: StreamEntry,\n    eventId: string,\n    errorMessage: string,\n    errorCode: SSEErrorCode,\n    closeClients: boolean = false,\n    upstreamCode?: string,\n  ): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        this.sseWriter.writeError(\n          client,\n          eventId,\n          errorMessage,\n          errorCode,\n          upstreamCode,\n        );\n        if (closeClients) {\n          client.end();\n        }\n      }\n    }\n  }\n\n  private _finalizeStream(streamEntry: StreamEntry): void {\n    streamEntry.isCompleted = true;\n    clearGraceTimer(streamEntry);\n    this._closeAllClients(streamEntry);\n    this._scheduleRemovalAfterTTL(streamEntry);\n  }\n\n  // close all connected clients and remove them from the stream entry.\n  // We are deliberately terminating these connections, so cleanup must not\n  // depend on the transport emitting a `close` event for each client (it may\n  // never fire). The close handlers' later `clients.delete(...)` calls remain\n  // safe no-ops.\n  private _closeAllClients(streamEntry: StreamEntry): void {\n    for (const client of streamEntry.clients) {\n      if (!client.writableEnded) {\n        client.end();\n      }\n    }\n    streamEntry.clients.clear();\n  }\n\n  // abort the generator after the grace window unless a client reconnects first\n  private _scheduleGraceAbort(streamEntry: StreamEntry): void {\n    // clear any existing timer to avoid stacking\n    clearGraceTimer(streamEntry);\n\n    const timer = setTimeout(() => {\n      streamEntry.disconnectGraceTimer = undefined;\n      if (streamEntry.clients.size === 0 && !streamEntry.isCompleted) {\n        streamEntry.abortController.abort(\n          new DOMException(\"Client disconnected (grace expired)\", \"AbortError\"),\n        );\n      }\n    }, this.disconnectGraceMs);\n\n    // never keep the process alive solely for a grace timer\n    timer.unref?.();\n    streamEntry.disconnectGraceTimer = timer;\n  }\n\n  // schedule registry removal once a finished (completed or errored) stream\n  // has no connected clients. The event buffer stays available for\n  // reconnect replay for `bufferTTL` after the last client disconnects;\n  // after that the stream entry is removed from the registry so it can be\n  // garbage collected.\n  private _scheduleRemovalAfterTTL(streamEntry: StreamEntry): void {\n    if (!streamEntry.isCompleted || streamEntry.clients.size > 0) {\n      return;\n    }\n\n    // at most one removal timer per stream: rescheduling replaces any\n    // pending timer instead of stacking a new one (each pending timer pins\n    // the entry's buffer/generator/trace context for the full TTL)\n    clearRemovalTimer(streamEntry);\n\n    // mark the moment the stream became idle so a reconnect during the TTL\n    // window (which refreshes lastAccess) makes the pending timer a no-op\n    streamEntry.lastAccess = Date.now();\n\n    streamEntry.removalTimer = setTimeout(() => {\n      streamEntry.removalTimer = undefined;\n      // safety net: no-op if a client reconnected during the TTL window;\n      // a fresh removal is scheduled when that client disconnects\n      if (\n        streamEntry.clients.size === 0 &&\n        Date.now() - streamEntry.lastAccess >= this.bufferTTL\n      ) {\n        this.streamRegistry.remove(streamEntry.streamId);\n      }\n    }, this.bufferTTL);\n\n    // don't keep the process alive just to clean up finished streams\n    streamEntry.removalTimer.unref?.();\n  }\n\n  private _categorizeError(error: unknown): SSEErrorCode {\n    if (error instanceof Error) {\n      const message = error.message.toLowerCase();\n      if (message.includes(\"timeout\") || message.includes(\"timed out\")) {\n        return SSEErrorCode.TIMEOUT;\n      }\n\n      if (message.includes(\"unavailable\") || message.includes(\"econnrefused\")) {\n        return SSEErrorCode.TEMPORARY_UNAVAILABLE;\n      }\n\n      if (error.name === \"AbortError\") {\n        return SSEErrorCode.STREAM_ABORTED;\n      }\n\n      // Defense-in-depth: upstream layers (SQL client, cache) may wrap an\n      // AbortError into ExecutionError, losing `name` but keeping the message.\n      if (\n        message.includes(\"operation was aborted\") ||\n        message.includes(\"the request was aborted\") ||\n        message.includes(\"statement was canceled\")\n      ) {\n        return SSEErrorCode.STREAM_ABORTED;\n      }\n\n      // Detect upstream API errors (e.g., from Databricks SDK ApiError)\n      if (\n        \"statusCode\" in error &&\n        typeof (error as any).statusCode === \"number\"\n      ) {\n        return SSEErrorCode.UPSTREAM_ERROR;\n      }\n    }\n\n    return SSEErrorCode.INTERNAL_ERROR;\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;AAcA,MAAM,SAAS,aAAa,SAAS;AAGrC,IAAa,gBAAb,MAA2B;CACzB,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CACR,AAAQ;CAER,YAAY,SAAwB;AAClC,OAAK,iBAAiB,IAAI,eACxB,SAAS,oBAAoB,eAAe,iBAC7C;AACD,OAAK,YAAY,IAAI,WAAW;AAChC,OAAK,eAAe,SAAS,gBAAgB,eAAe;AAC5D,OAAK,YAAY,SAAS,aAAa,eAAe;AACtD,OAAK,oBACH,SAAS,qBAAqB,eAAe;AAC/C,OAAK,mCAAmB,IAAI,KAAK;;CAInC,MAAM,OACJ,KACA,SACA,SACA,UACe;EACf,MAAM,EAAE,aAAa,WAAW,EAAE;AAGlC,MAAI,IAAI,iBAAiB,IAAI,UAC3B;AAIF,OAAK,UAAU,aAAa,IAAI;AAGhC,MAAI,YAAY,gBAAgB,iBAAiB,SAAS,EAAE;GAC1D,MAAM,iBAAiB,KAAK,eAAe,IAAI,SAAS;AACxD,OAAI,gBAAgB;AAKlB,QAAI,eAAe,aAAa,UAAU;AACxC,UAAK,UAAU,WACb,KACA,YAAY,EACZ,qCACA,aAAa,iBACd;AACD,SAAI,KAAK;AACT;;AAEF,WAAO,KAAK,wBAAwB,KAAK,gBAAgB,QAAQ;;;AAKrE,SAAO,KAAK,iBAAiB,KAAK,SAAS,SAAS,SAAS;;CAI/D,WAAiB;AAEf,OAAK,iBAAiB,SAAS,cAAc;AAC3C,OAAI,UAAU,UAAW,eAAc,UAAU,UAAU;AAC3D,aAAU,WAAW,MACnB,IAAI,aAAa,mBAAmB,aAAa,CAClD;IACD;AACF,OAAK,iBAAiB,OAAO;AAC7B,OAAK,eAAe,OAAO;;CAI7B,iBAAyB;AACvB,SAAO,KAAK,iBAAiB;;CAI/B,MAAc,wBACZ,KACA,aACA,SACe;EAEf,MAAM,cAAc,IAAI,KAAK,QAAQ;AAErC,MAAI,gBAAgB,gBAAgB,YAAY,EAAE;GAEhD,MAAM,eAAe;AACrB,OAAI,YAAY,YAAY,IAAI,aAAa,EAAE;IAC7C,MAAM,eACJ,YAAY,YAAY,eAAe,aAAa;AAEtD,SAAK,MAAM,SAAS,cAAc;AAChC,SAAI,SAAS,YAAY,QAAS;AAClC,UAAK,UAAU,mBAAmB,KAAK,MAAM;;SAI/C,MAAK,UAAU,2BAA2B,KAAK,aAAa;;AAKhE,kBAAgB,YAAY;AAI5B,oBAAkB,YAAY;AAG9B,cAAY,QAAQ,IAAI,IAAI;AAC5B,cAAY,aAAa,KAAK,KAAK;EAGnC,MAAM,iBAAiB,KAAK,gBAC1B,YAAY,gBAAgB,QAC5B,SAAS,WACV;EACD,MAAM,YAAY,KAAK,UAAU,eAAe,KAAK,eAAe;EAGpE,MAAM,kBAAmC;GACvC,YAAY,YAAY;GACxB,MAAM;GACN;GACD;AACD,OAAK,iBAAiB,IAAI,gBAAgB;AAG1C,MAAI,GAAG,eAAe;AACpB,iBAAc,UAAU;AACxB,eAAY,QAAQ,OAAO,IAAI;AAC/B,QAAK,iBAAiB,OAAO,gBAAgB;AAG7C,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,MAAK,oBAAoB,YAAY;AAIvC,QAAK,yBAAyB,YAAY;IAC1C;AAGF,MAAI,YAAY,aAAa;AAC3B,OAAI,KAAK;AAET,QAAK,iBAAiB,OAAO,gBAAgB;AAC7C,iBAAc,UAAU;AAIxB,eAAY,QAAQ,OAAO,IAAI;AAC/B,QAAK,yBAAyB,YAAY;;;CAG9C,MAAc,iBACZ,KACA,SACA,SACA,UACe;EACf,MAAM,WAAW,SAAS,YAAY,YAAY;AAGlD,MAAI,IAAI,iBAAiB,IAAI,UAC3B;EAGF,MAAM,kBAAkB,IAAI,iBAAiB;EAG7C,MAAM,cAAc,IAAI,gBACtB,SAAS,cAAc,eAAe,WACvC;EAGD,MAAM,iBAAiB,KAAK,gBAC1B,gBAAgB,QAChB,SAAS,WACV;EACD,MAAM,YAAY,KAAK,UAAU,eAAe,KAAK,eAAe;EAGpE,MAAM,eAAe,QAAQ,QAAQ;AAGrC,MAAI,IAAI,iBAAiB,IAAI,WAAW;AACtC,iBAAc,UAAU;AACxB;;EAIF,MAAM,cAA2B;GAC/B;GACA;GACA,WAAW,QAAQ,eAAe;GAClC;GACA,SAAS,IAAI,IAAI,CAAC,IAAI,CAAC;GACvB,aAAa;GACb,YAAY,KAAK,KAAK;GACtB;GACA;GACD;AACD,OAAK,eAAe,IAAI,YAAY;EAGpC,MAAM,kBAAmC;GACvC,YAAY;GACZ,MAAM;GACN;GACD;AACD,OAAK,iBAAiB,IAAI,gBAAgB;AAE1C,MAAI,GAAG,eAAe;AACpB,iBAAc,UAAU;AACxB,QAAK,iBAAiB,OAAO,gBAAgB;AAC7C,eAAY,QAAQ,OAAO,IAAI;AAG/B,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,MAAK,oBAAoB,YAAY;AAMvC,QAAK,yBAAyB,YAAY;IAC1C;AAEF,QAAM,KAAK,8BAA8B,YAAY;AAGrD,gBAAc,UAAU;AACxB,OAAK,iBAAiB,OAAO,gBAAgB;;CAG/C,MAAc,8BACZ,aACe;AAEf,SAAO,QAAQ,KAAK,YAAY,cAAc,YAAY;AACxD,OAAI;AAEF,eAAW,MAAM,SAAS,YAAY,WAAW;AAC/C,SAAI,YAAY,gBAAgB,OAAO,QAAS;KAChD,MAAM,UAAU,YAAY;KAC5B,MAAM,YAAY,KAAK,UAAU,MAAM;AAGvC,SAAI,UAAU,SAAS,KAAK,cAAc;MACxC,MAAM,WAAW,6BAA6B,KAAK,aAAa;MAChE,MAAM,YAAY,aAAa;AAE/B,WAAK,yBACH,aACA,SACA,UACA,UACD;AACD;;AAIF,iBAAY,YAAY,IAAI;MAC1B,IAAI;MACJ,MAAM,MAAM;MACZ,MAAM;MACN,WAAW,KAAK,KAAK;MACtB,CAAC;AAGF,UAAK,0BAA0B,aAAa,SAAS,MAAM;AAC3D,iBAAY,aAAa,KAAK,KAAK;;AAGrC,SAAK,gBAAgB,YAAY;YAC1B,OAAO;IAMd,MAAM,SACJ,iBAAiB,QAAQ,MAAM,UAAU;IAC3C,MAAM,YACJ,iBAAiB,cACb,MAAM,gBACN;IAGN,MAAM,eACJ,iBAAiB,iBAAiB,MAAM,YAAY;IACtD,MAAM,eAAe,YAAY;IACjC,MAAM,YAAY,KAAK,iBAAiB,MAAM;AAG9C,QAAI,cAAc,aAAa,gBAAgB;AAC7C,YAAO,KAAK,sCAAsC,UAAU;AAC5D,UAAK,gBAAgB,YAAY;AACjC;;AAGF,WAAO,MACL,yDACA,QACA,WACA,gBAAgB,MACjB;IAED,MAAM,UAAmC;KACvC,OAAO;KACP,MAAM;KACP;AACD,QAAI,aAAc,SAAQ,YAAY;AAGtC,gBAAY,YAAY,IAAI;KAC1B,IAAI;KACJ,MAAM;KACN,MAAM,KAAK,UAAU,QAAQ;KAC7B,WAAW,KAAK,KAAK;KACtB,CAAC;AAGF,SAAK,yBACH,aACA,cACA,WACA,WACA,MACA,aACD;AAID,SAAK,gBAAgB,YAAY;;IAEnC;;CAGJ,AAAQ,gBACN,gBACA,YACa;AACb,MAAI,CAAC,WAAY,QAAO,kBAAkB,IAAI,iBAAiB,CAAC;EAEhE,MAAM,UAAU,CAAC,gBAAgB,WAAW,CAAC,OAC3C,QACD;EACD,MAAM,aAAa,IAAI,iBAAiB;AAExC,UAAQ,SAAS,WAAW;AAC1B,OAAI,QAAQ,SAAS;AACnB,eAAW,MAAM,OAAO,OAAO;AAC/B;;AAGF,WAAQ,iBACN,eACM;AACJ,eAAW,MAAM,OAAO,OAAO;MAEjC,EAAE,MAAM,MAAM,CACf;IACD;AACF,SAAO,WAAW;;CAIpB,AAAQ,0BACN,aACA,SACA,OACM;AACN,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,cACV,MAAK,UAAU,WAAW,QAAQ,SAAS,MAAM;;CAMvD,AAAQ,yBACN,aACA,SACA,cACA,WACA,eAAwB,OACxB,cACM;AACN,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,eAAe;AACzB,QAAK,UAAU,WACb,QACA,SACA,cACA,WACA,aACD;AACD,OAAI,aACF,QAAO,KAAK;;;CAMpB,AAAQ,gBAAgB,aAAgC;AACtD,cAAY,cAAc;AAC1B,kBAAgB,YAAY;AAC5B,OAAK,iBAAiB,YAAY;AAClC,OAAK,yBAAyB,YAAY;;CAQ5C,AAAQ,iBAAiB,aAAgC;AACvD,OAAK,MAAM,UAAU,YAAY,QAC/B,KAAI,CAAC,OAAO,cACV,QAAO,KAAK;AAGhB,cAAY,QAAQ,OAAO;;CAI7B,AAAQ,oBAAoB,aAAgC;AAE1D,kBAAgB,YAAY;EAE5B,MAAM,QAAQ,iBAAiB;AAC7B,eAAY,uBAAuB;AACnC,OAAI,YAAY,QAAQ,SAAS,KAAK,CAAC,YAAY,YACjD,aAAY,gBAAgB,MAC1B,IAAI,aAAa,uCAAuC,aAAa,CACtE;KAEF,KAAK,kBAAkB;AAG1B,QAAM,SAAS;AACf,cAAY,uBAAuB;;CAQrC,AAAQ,yBAAyB,aAAgC;AAC/D,MAAI,CAAC,YAAY,eAAe,YAAY,QAAQ,OAAO,EACzD;AAMF,oBAAkB,YAAY;AAI9B,cAAY,aAAa,KAAK,KAAK;AAEnC,cAAY,eAAe,iBAAiB;AAC1C,eAAY,eAAe;AAG3B,OACE,YAAY,QAAQ,SAAS,KAC7B,KAAK,KAAK,GAAG,YAAY,cAAc,KAAK,UAE5C,MAAK,eAAe,OAAO,YAAY,SAAS;KAEjD,KAAK,UAAU;AAGlB,cAAY,aAAa,SAAS;;CAGpC,AAAQ,iBAAiB,OAA8B;AACrD,MAAI,iBAAiB,OAAO;GAC1B,MAAM,UAAU,MAAM,QAAQ,aAAa;AAC3C,OAAI,QAAQ,SAAS,UAAU,IAAI,QAAQ,SAAS,YAAY,CAC9D,QAAO,aAAa;AAGtB,OAAI,QAAQ,SAAS,cAAc,IAAI,QAAQ,SAAS,eAAe,CACrE,QAAO,aAAa;AAGtB,OAAI,MAAM,SAAS,aACjB,QAAO,aAAa;AAKtB,OACE,QAAQ,SAAS,wBAAwB,IACzC,QAAQ,SAAS,0BAA0B,IAC3C,QAAQ,SAAS,yBAAyB,CAE1C,QAAO,aAAa;AAItB,OACE,gBAAgB,SAChB,OAAQ,MAAc,eAAe,SAErC,QAAO,aAAa;;AAIxB,SAAO,aAAa"}

package/dist/stream/stream-registry.js

@@ -1,50 +1,56 @@
-import { RingBuffer } from "./buffers.js";
 import { SSEErrorCode } from "./types.js";
+import { clearGraceTimer, clearRemovalTimer } from "./timers.js";
 
 //#region src/stream/stream-registry.ts
 var StreamRegistry = class {
 	streams;
+	maxActiveStreams;
 	constructor(maxActiveStreams) {
-		this.streams = new RingBuffer(maxActiveStreams, (entry) => entry.streamId);
+		this.streams = /* @__PURE__ */ new Map();
+		this.maxActiveStreams = maxActiveStreams;
 	}
 	add(entry) {
-		if (this.streams.getSize() >= this.streams.capacity) this._evictOldestStream(entry.streamId);
-		this.streams.add(entry);
+		if (!this.streams.has(entry.streamId) && this.streams.size >= this.maxActiveStreams) this._evictOldestStream(entry.streamId);
+		this.streams.set(entry.streamId, entry);
 	}
 	get(streamId) {
-		return this.streams.get(streamId);
+		return this.streams.get(streamId) ?? null;
 	}
 	has(streamId) {
 		return this.streams.has(streamId);
 	}
 	remove(streamId) {
-		this.streams.remove(streamId);
+		this.streams.delete(streamId);
 	}
 	size() {
-		return this.streams.getSize();
+		return this.streams.size;
 	}
 	clear() {
-		const allStreams = this.streams.getAll();
-		for (const stream of allStreams) {
-			this._clearGraceTimer(stream);
+		for (const stream of this.streams.values()) {
 			stream.abortController.abort("Server shutdown");
+			clearGraceTimer(stream);
+			clearRemovalTimer(stream);
 		}
 		this.streams.clear();
 	}
-	_clearGraceTimer(stream) {
-		if (stream.disconnectGraceTimer) {
-			clearTimeout(stream.disconnectGraceTimer);
-			stream.disconnectGraceTimer = void 0;
-		}
-	}
 	_evictOldestStream(excludeStreamId) {
-		const allStreams = this.streams.getAll();
+		const allStreams = this.streams.values();
 		let oldestStream = null;
 		let oldestAccess = Infinity;
-		for (const stream of allStreams) if (stream.streamId !== excludeStreamId && stream.lastAccess < oldestAccess) {
-			oldestStream = stream;
-			oldestAccess = stream.lastAccess;
+		let oldestCompletedStream = null;
+		let oldestCompletedAccess = Infinity;
+		for (const stream of allStreams) {
+			if (stream.streamId === excludeStreamId) continue;
+			if (stream.lastAccess < oldestAccess) {
+				oldestStream = stream;
+				oldestAccess = stream.lastAccess;
+			}
+			if (stream.isCompleted && stream.lastAccess < oldestCompletedAccess) {
+				oldestCompletedStream = stream;
+				oldestCompletedAccess = stream.lastAccess;
+			}
 		}
+		oldestStream = oldestCompletedStream ?? oldestStream;
 		if (oldestStream) {
 			for (const client of oldestStream.clients) if (!client.writableEnded) try {
 				client.write(`event: error\n`);
@@ -53,9 +59,10 @@ var StreamRegistry = class {
 					code: SSEErrorCode.STREAM_EVICTED
 				})}\n\n`);
 			} catch (_error) {}
-			this._clearGraceTimer(oldestStream);
-			oldestStream.abortController.abort("Stream evicted");
-			this.streams.remove(oldestStream.streamId);
+			clearGraceTimer(oldestStream);
+			oldestStream.abortController.abort(new DOMException("Stream evicted", "AbortError"));
+			clearRemovalTimer(oldestStream);
+			this.streams.delete(oldestStream.streamId);
 		}
 	}
 };

package/dist/stream/stream-registry.js.map

@@ -1 +1 @@
-{"version":3,"file":"stream-registry.js","names":[],"sources":["../../src/stream/stream-registry.ts"],"sourcesContent":["import { RingBuffer } from \"./buffers\";\nimport { SSEErrorCode, type StreamEntry } from \"./types\";\n\nexport class StreamRegistry {\n  private streams: RingBuffer<StreamEntry>;\n\n  constructor(maxActiveStreams: number) {\n    this.streams = new RingBuffer<StreamEntry>(\n      maxActiveStreams,\n      (entry) => entry.streamId,\n    );\n  }\n\n  // add a stream to the registry\n  add(entry: StreamEntry): void {\n    // enforce hard cap\n    if (this.streams.getSize() >= this.streams.capacity) {\n      this._evictOldestStream(entry.streamId);\n    }\n\n    this.streams.add(entry);\n  }\n\n  // get a stream from the registry\n  get(streamId: string): StreamEntry | null {\n    return this.streams.get(streamId);\n  }\n\n  // check if a stream exists in the registry\n  has(streamId: string): boolean {\n    return this.streams.has(streamId);\n  }\n\n  // remove a stream from the registry\n  remove(streamId: string): void {\n    this.streams.remove(streamId);\n  }\n\n  // get the number of streams in the registry\n  size(): number {\n    return this.streams.getSize();\n  }\n\n  clear(): void {\n    const allStreams = this.streams.getAll();\n\n    for (const stream of allStreams) {\n      this._clearGraceTimer(stream);\n      stream.abortController.abort(\"Server shutdown\");\n    }\n\n    this.streams.clear();\n  }\n\n  // clear a pending grace timer so a removed stream isn't pinned until it fires\n  private _clearGraceTimer(stream: StreamEntry): void {\n    if (stream.disconnectGraceTimer) {\n      clearTimeout(stream.disconnectGraceTimer);\n      stream.disconnectGraceTimer = undefined;\n    }\n  }\n\n  // evict the oldest stream from the registry\n  private _evictOldestStream(excludeStreamId: string): void {\n    const allStreams = this.streams.getAll();\n    let oldestStream: StreamEntry | null = null;\n    let oldestAccess = Infinity;\n\n    // find the least recently accessed stream\n    for (const stream of allStreams) {\n      if (\n        stream.streamId !== excludeStreamId &&\n        stream.lastAccess < oldestAccess\n      ) {\n        oldestStream = stream;\n        oldestAccess = stream.lastAccess;\n      }\n    }\n\n    // abort the oldest stream\n    if (oldestStream) {\n      // broadcast stream eviction error to all clients\n      for (const client of oldestStream.clients) {\n        if (!client.writableEnded) {\n          try {\n            client.write(`event: error\\n`);\n            client.write(\n              `data: ${JSON.stringify({ error: \"Stream evicted\", code: SSEErrorCode.STREAM_EVICTED })}\\n\\n`,\n            );\n          } catch (_error) {\n            // ignore\n          }\n        }\n      }\n      this._clearGraceTimer(oldestStream);\n      oldestStream.abortController.abort(\"Stream evicted\");\n      this.streams.remove(oldestStream.streamId);\n    }\n  }\n}\n"],"mappings":";;;;AAGA,IAAa,iBAAb,MAA4B;CAC1B,AAAQ;CAER,YAAY,kBAA0B;AACpC,OAAK,UAAU,IAAI,WACjB,mBACC,UAAU,MAAM,SAClB;;CAIH,IAAI,OAA0B;AAE5B,MAAI,KAAK,QAAQ,SAAS,IAAI,KAAK,QAAQ,SACzC,MAAK,mBAAmB,MAAM,SAAS;AAGzC,OAAK,QAAQ,IAAI,MAAM;;CAIzB,IAAI,UAAsC;AACxC,SAAO,KAAK,QAAQ,IAAI,SAAS;;CAInC,IAAI,UAA2B;AAC7B,SAAO,KAAK,QAAQ,IAAI,SAAS;;CAInC,OAAO,UAAwB;AAC7B,OAAK,QAAQ,OAAO,SAAS;;CAI/B,OAAe;AACb,SAAO,KAAK,QAAQ,SAAS;;CAG/B,QAAc;EACZ,MAAM,aAAa,KAAK,QAAQ,QAAQ;AAExC,OAAK,MAAM,UAAU,YAAY;AAC/B,QAAK,iBAAiB,OAAO;AAC7B,UAAO,gBAAgB,MAAM,kBAAkB;;AAGjD,OAAK,QAAQ,OAAO;;CAItB,AAAQ,iBAAiB,QAA2B;AAClD,MAAI,OAAO,sBAAsB;AAC/B,gBAAa,OAAO,qBAAqB;AACzC,UAAO,uBAAuB;;;CAKlC,AAAQ,mBAAmB,iBAA+B;EACxD,MAAM,aAAa,KAAK,QAAQ,QAAQ;EACxC,IAAI,eAAmC;EACvC,IAAI,eAAe;AAGnB,OAAK,MAAM,UAAU,WACnB,KACE,OAAO,aAAa,mBACpB,OAAO,aAAa,cACpB;AACA,kBAAe;AACf,kBAAe,OAAO;;AAK1B,MAAI,cAAc;AAEhB,QAAK,MAAM,UAAU,aAAa,QAChC,KAAI,CAAC,OAAO,cACV,KAAI;AACF,WAAO,MAAM,iBAAiB;AAC9B,WAAO,MACL,SAAS,KAAK,UAAU;KAAE,OAAO;KAAkB,MAAM,aAAa;KAAgB,CAAC,CAAC,MACzF;YACM,QAAQ;AAKrB,QAAK,iBAAiB,aAAa;AACnC,gBAAa,gBAAgB,MAAM,iBAAiB;AACpD,QAAK,QAAQ,OAAO,aAAa,SAAS"}
+{"version":3,"file":"stream-registry.js","names":[],"sources":["../../src/stream/stream-registry.ts"],"sourcesContent":["import { clearGraceTimer, clearRemovalTimer } from \"./timers\";\nimport { SSEErrorCode, type StreamEntry } from \"./types\";\n\nexport class StreamRegistry {\n  // keyed storage with explicit, policy-driven eviction. A ring buffer is\n  // unsuitable here: it overwrites by insertion slot, so evicting an entry\n  // chosen by policy (e.g. a completed stream) and then adding would\n  // silently clobber an unrelated live stream sitting in the oldest slot.\n  private streams: Map<string, StreamEntry>;\n  private maxActiveStreams: number;\n\n  constructor(maxActiveStreams: number) {\n    this.streams = new Map();\n    this.maxActiveStreams = maxActiveStreams;\n  }\n\n  // add a stream to the registry\n  add(entry: StreamEntry): void {\n    // enforce hard cap, but only when inserting a genuinely new key:\n    // re-adding an existing streamId updates in place and doesn't grow the\n    // map, so evicting another stream for it would destroy an innocent one\n    if (\n      !this.streams.has(entry.streamId) &&\n      this.streams.size >= this.maxActiveStreams\n    ) {\n      this._evictOldestStream(entry.streamId);\n    }\n\n    this.streams.set(entry.streamId, entry);\n  }\n\n  // get a stream from the registry\n  get(streamId: string): StreamEntry | null {\n    return this.streams.get(streamId) ?? null;\n  }\n\n  // check if a stream exists in the registry\n  has(streamId: string): boolean {\n    return this.streams.has(streamId);\n  }\n\n  // remove a stream from the registry\n  remove(streamId: string): void {\n    this.streams.delete(streamId);\n  }\n\n  // get the number of streams in the registry\n  size(): number {\n    return this.streams.size;\n  }\n\n  clear(): void {\n    for (const stream of this.streams.values()) {\n      stream.abortController.abort(\"Server shutdown\");\n      clearGraceTimer(stream);\n      clearRemovalTimer(stream);\n    }\n\n    this.streams.clear();\n  }\n\n  // evict the oldest stream from the registry, preferring completed streams.\n  // Completed streams waiting out their buffer TTL can look recently\n  // accessed, so plain LRU could evict a live stream while dead ones\n  // survive. Prefer the oldest completed stream when one exists and fall\n  // back to LRU over all streams otherwise.\n  private _evictOldestStream(excludeStreamId: string): void {\n    const allStreams = this.streams.values();\n    let oldestStream: StreamEntry | null = null;\n    let oldestAccess = Infinity;\n    let oldestCompletedStream: StreamEntry | null = null;\n    let oldestCompletedAccess = Infinity;\n\n    // find the least recently accessed stream (overall and completed-only)\n    for (const stream of allStreams) {\n      if (stream.streamId === excludeStreamId) continue;\n\n      if (stream.lastAccess < oldestAccess) {\n        oldestStream = stream;\n        oldestAccess = stream.lastAccess;\n      }\n\n      if (stream.isCompleted && stream.lastAccess < oldestCompletedAccess) {\n        oldestCompletedStream = stream;\n        oldestCompletedAccess = stream.lastAccess;\n      }\n    }\n\n    oldestStream = oldestCompletedStream ?? oldestStream;\n\n    // abort the evicted stream\n    if (oldestStream) {\n      // broadcast stream eviction error to all clients\n      for (const client of oldestStream.clients) {\n        if (!client.writableEnded) {\n          try {\n            client.write(`event: error\\n`);\n            client.write(\n              `data: ${JSON.stringify({ error: \"Stream evicted\", code: SSEErrorCode.STREAM_EVICTED })}\\n\\n`,\n            );\n          } catch (_error) {\n            // ignore\n          }\n        }\n      }\n      clearGraceTimer(oldestStream);\n      // abort with a DOMException AbortError so the error categorizer routes\n      // eviction as an abort (matching the manager's terminal paths) rather\n      // than a stream failure\n      oldestStream.abortController.abort(\n        new DOMException(\"Stream evicted\", \"AbortError\"),\n      );\n      // a pending removal timer would otherwise pin the evicted entry\n      clearRemovalTimer(oldestStream);\n      this.streams.delete(oldestStream.streamId);\n    }\n  }\n}\n"],"mappings":";;;;AAGA,IAAa,iBAAb,MAA4B;CAK1B,AAAQ;CACR,AAAQ;CAER,YAAY,kBAA0B;AACpC,OAAK,0BAAU,IAAI,KAAK;AACxB,OAAK,mBAAmB;;CAI1B,IAAI,OAA0B;AAI5B,MACE,CAAC,KAAK,QAAQ,IAAI,MAAM,SAAS,IACjC,KAAK,QAAQ,QAAQ,KAAK,iBAE1B,MAAK,mBAAmB,MAAM,SAAS;AAGzC,OAAK,QAAQ,IAAI,MAAM,UAAU,MAAM;;CAIzC,IAAI,UAAsC;AACxC,SAAO,KAAK,QAAQ,IAAI,SAAS,IAAI;;CAIvC,IAAI,UAA2B;AAC7B,SAAO,KAAK,QAAQ,IAAI,SAAS;;CAInC,OAAO,UAAwB;AAC7B,OAAK,QAAQ,OAAO,SAAS;;CAI/B,OAAe;AACb,SAAO,KAAK,QAAQ;;CAGtB,QAAc;AACZ,OAAK,MAAM,UAAU,KAAK,QAAQ,QAAQ,EAAE;AAC1C,UAAO,gBAAgB,MAAM,kBAAkB;AAC/C,mBAAgB,OAAO;AACvB,qBAAkB,OAAO;;AAG3B,OAAK,QAAQ,OAAO;;CAQtB,AAAQ,mBAAmB,iBAA+B;EACxD,MAAM,aAAa,KAAK,QAAQ,QAAQ;EACxC,IAAI,eAAmC;EACvC,IAAI,eAAe;EACnB,IAAI,wBAA4C;EAChD,IAAI,wBAAwB;AAG5B,OAAK,MAAM,UAAU,YAAY;AAC/B,OAAI,OAAO,aAAa,gBAAiB;AAEzC,OAAI,OAAO,aAAa,cAAc;AACpC,mBAAe;AACf,mBAAe,OAAO;;AAGxB,OAAI,OAAO,eAAe,OAAO,aAAa,uBAAuB;AACnE,4BAAwB;AACxB,4BAAwB,OAAO;;;AAInC,iBAAe,yBAAyB;AAGxC,MAAI,cAAc;AAEhB,QAAK,MAAM,UAAU,aAAa,QAChC,KAAI,CAAC,OAAO,cACV,KAAI;AACF,WAAO,MAAM,iBAAiB;AAC9B,WAAO,MACL,SAAS,KAAK,UAAU;KAAE,OAAO;KAAkB,MAAM,aAAa;KAAgB,CAAC,CAAC,MACzF;YACM,QAAQ;AAKrB,mBAAgB,aAAa;AAI7B,gBAAa,gBAAgB,MAC3B,IAAI,aAAa,kBAAkB,aAAa,CACjD;AAED,qBAAkB,aAAa;AAC/B,QAAK,QAAQ,OAAO,aAAa,SAAS"}

package/dist/stream/timers.js

@@ -0,0 +1,17 @@
+//#region src/stream/timers.ts
+function clearGraceTimer(entry) {
+	if (entry.disconnectGraceTimer) {
+		clearTimeout(entry.disconnectGraceTimer);
+		entry.disconnectGraceTimer = void 0;
+	}
+}
+function clearRemovalTimer(entry) {
+	if (entry.removalTimer) {
+		clearTimeout(entry.removalTimer);
+		entry.removalTimer = void 0;
+	}
+}
+
+//#endregion
+export { clearGraceTimer, clearRemovalTimer };
+//# sourceMappingURL=timers.js.map

package/dist/stream/timers.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"timers.js","names":[],"sources":["../../src/stream/timers.ts"],"sourcesContent":["import type { StreamEntry } from \"./types\";\n\n// clear a pending disconnect-grace timer so a removed/reconnected stream\n// isn't pinned until it fires\nexport function clearGraceTimer(entry: StreamEntry): void {\n  if (entry.disconnectGraceTimer) {\n    clearTimeout(entry.disconnectGraceTimer);\n    entry.disconnectGraceTimer = undefined;\n  }\n}\n\n// clear a pending registry-removal timer so a reconnected/evicted stream\n// isn't pulled out from under a client\nexport function clearRemovalTimer(entry: StreamEntry): void {\n  if (entry.removalTimer) {\n    clearTimeout(entry.removalTimer);\n    entry.removalTimer = undefined;\n  }\n}\n"],"mappings":";AAIA,SAAgB,gBAAgB,OAA0B;AACxD,KAAI,MAAM,sBAAsB;AAC9B,eAAa,MAAM,qBAAqB;AACxC,QAAM,uBAAuB;;;AAMjC,SAAgB,kBAAkB,OAA0B;AAC1D,KAAI,MAAM,cAAc;AACtB,eAAa,MAAM,aAAa;AAChC,QAAM,eAAe"}

package/dist/stream/types.js.map

@@ -1 +1 @@
-{"version":3,"file":"types.js","names":[],"sources":["../../src/stream/types.ts"],"sourcesContent":["import type { Context } from \"@opentelemetry/api\";\nimport type { IAppResponse } from \"shared\";\nimport type { EventRingBuffer } from \"./buffers\";\n\nexport const SSEWarningCode = {\n  BUFFER_OVERFLOW_RESTART: \"BUFFER_OVERFLOW_RESTART\",\n} as const satisfies Record<string, string>;\n\nexport type SSEWarningCode =\n  (typeof SSEWarningCode)[keyof typeof SSEWarningCode];\n\nexport const SSEErrorCode = {\n  TEMPORARY_UNAVAILABLE: \"TEMPORARY_UNAVAILABLE\",\n  TIMEOUT: \"TIMEOUT\",\n  INTERNAL_ERROR: \"INTERNAL_ERROR\",\n  INVALID_REQUEST: \"INVALID_REQUEST\",\n  STREAM_ABORTED: \"STREAM_ABORTED\",\n  STREAM_EVICTED: \"STREAM_EVICTED\",\n  STREAM_FORBIDDEN: \"STREAM_FORBIDDEN\",\n  UPSTREAM_ERROR: \"UPSTREAM_ERROR\",\n} as const satisfies Record<string, string>;\n\nexport type SSEErrorCode = (typeof SSEErrorCode)[keyof typeof SSEErrorCode];\n\nexport interface SSEError {\n  error: string;\n  code: SSEErrorCode;\n}\n\nexport interface BufferedEvent {\n  id: string;\n  type: string;\n  data: string;\n  timestamp: number;\n}\n\nexport interface StreamEntry {\n  streamId: string;\n  /**\n   * Identifier of the principal that created the stream (e.g. end-user ID\n   * or service principal user ID). When set, only requests sharing the\n   * same owner key may reconnect to the stream.\n   */\n  ownerKey?: string;\n  generator: AsyncGenerator<any, void, unknown>;\n  eventBuffer: EventRingBuffer;\n  clients: Set<IAppResponse>;\n  isCompleted: boolean;\n  lastAccess: number;\n  abortController: AbortController;\n  traceContext: Context;\n  // pending grace-window abort, set while the last client is disconnected\n  disconnectGraceTimer?: NodeJS.Timeout;\n}\n\nexport interface StreamOperation {\n  controller: AbortController;\n  type: \"query\" | \"stream\";\n  heartbeat?: NodeJS.Timeout;\n}\n"],"mappings":";AAIA,MAAa,iBAAiB,EAC5B,yBAAyB,2BAC1B;AAKD,MAAa,eAAe;CAC1B,uBAAuB;CACvB,SAAS;CACT,gBAAgB;CAChB,iBAAiB;CACjB,gBAAgB;CAChB,gBAAgB;CAChB,kBAAkB;CAClB,gBAAgB;CACjB"}
+{"version":3,"file":"types.js","names":[],"sources":["../../src/stream/types.ts"],"sourcesContent":["import type { Context } from \"@opentelemetry/api\";\nimport type { IAppResponse } from \"shared\";\nimport type { EventRingBuffer } from \"./buffers\";\n\nexport const SSEWarningCode = {\n  BUFFER_OVERFLOW_RESTART: \"BUFFER_OVERFLOW_RESTART\",\n} as const satisfies Record<string, string>;\n\nexport type SSEWarningCode =\n  (typeof SSEWarningCode)[keyof typeof SSEWarningCode];\n\nexport const SSEErrorCode = {\n  TEMPORARY_UNAVAILABLE: \"TEMPORARY_UNAVAILABLE\",\n  TIMEOUT: \"TIMEOUT\",\n  INTERNAL_ERROR: \"INTERNAL_ERROR\",\n  INVALID_REQUEST: \"INVALID_REQUEST\",\n  STREAM_ABORTED: \"STREAM_ABORTED\",\n  STREAM_EVICTED: \"STREAM_EVICTED\",\n  STREAM_FORBIDDEN: \"STREAM_FORBIDDEN\",\n  UPSTREAM_ERROR: \"UPSTREAM_ERROR\",\n} as const satisfies Record<string, string>;\n\nexport type SSEErrorCode = (typeof SSEErrorCode)[keyof typeof SSEErrorCode];\n\nexport interface SSEError {\n  error: string;\n  code: SSEErrorCode;\n  /**\n   * Upstream-domain structured code (e.g. `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`,\n   * `NOT_IMPLEMENTED`). UI code should branch on this instead of parsing\n   * the human-readable `error` string.\n   */\n  errorCode?: string;\n}\n\nexport interface BufferedEvent {\n  id: string;\n  type: string;\n  data: string;\n  timestamp: number;\n}\n\nexport interface StreamEntry {\n  streamId: string;\n  /**\n   * Identifier of the principal that created the stream (e.g. end-user ID\n   * or service principal user ID). When set, only requests sharing the\n   * same owner key may reconnect to the stream.\n   */\n  ownerKey?: string;\n  generator: AsyncGenerator<any, void, unknown>;\n  eventBuffer: EventRingBuffer;\n  clients: Set<IAppResponse>;\n  isCompleted: boolean;\n  lastAccess: number;\n  abortController: AbortController;\n  traceContext: Context;\n  // pending grace-window abort, set while the last client is disconnected\n  disconnectGraceTimer?: NodeJS.Timeout;\n  /**\n   * Pending registry-removal timer. At most one removal timer exists per\n   * stream; scheduling a new one clears any previous timer first.\n   */\n  removalTimer?: NodeJS.Timeout;\n}\n\nexport interface StreamOperation {\n  controller: AbortController;\n  type: \"query\" | \"stream\";\n  heartbeat?: NodeJS.Timeout;\n}\n"],"mappings":";AAIA,MAAa,iBAAiB,EAC5B,yBAAyB,2BAC1B;AAKD,MAAa,eAAe;CAC1B,uBAAuB;CACvB,SAAS;CACT,gBAAgB;CAChB,iBAAiB;CACjB,gBAAgB;CAChB,gBAAgB;CAChB,kBAAkB;CAClB,gBAAgB;CACjB"}

package/dist/telemetry/telemetry-manager.js

@@ -19,6 +19,7 @@ var TelemetryManager = class TelemetryManager {
 	static DEFAULT_FALLBACK_APP_NAME = "databricks-app";
 	static instance;
 	sdk;
+	shutdownPromise;
 	/**
 	* Create a scoped telemetry provider for a specific plugin.
 	* The plugin's name will be used as the default tracer/meter name.
@@ -54,7 +55,6 @@ var TelemetryManager = class TelemetryManager {
 				instrumentations: this.getDefaultInstrumentations()
 			});
 			this.sdk.start();
-			this.registerShutdown();
 			logger.debug("Initialized successfully");
 		} catch (error) {
 			logger.error("Failed to initialize: %O", error);
@@ -90,21 +90,27 @@ var TelemetryManager = class TelemetryManager {
 			"@opentelemetry/instrumentation-net": { enabled: false }
 		})];
 	}
-	registerShutdown() {
-		const shutdownFn = async () => {
-			await TelemetryManager.getInstance().shutdown();
-		};
-		process.once("SIGTERM", shutdownFn);
-		process.once("SIGINT", shutdownFn);
-	}
+	/**
+	* Flush and shut down the OpenTelemetry SDK.
+	*
+	* Idempotent: the SDK reference is cleared synchronously and concurrent
+	* or repeated calls await the same in-flight flush. Awaited by the core
+	* lifecycle manager during graceful shutdown — that manager owns the
+	* process signal handlers, so telemetry no longer registers its own.
+	*/
 	async shutdown() {
-		if (!this.sdk) return;
-		try {
-			await this.sdk.shutdown();
+		if (this.sdk) {
+			const sdk = this.sdk;
 			this.sdk = void 0;
-		} catch (error) {
-			logger.error("Error shutting down: %O", error);
+			this.shutdownPromise = (async () => {
+				try {
+					await sdk.shutdown();
+				} catch (error) {
+					logger.error("Error shutting down: %O", error);
+				}
+			})();
 		}
+		return this.shutdownPromise;
 	}
 };
 

package/dist/telemetry/telemetry-manager.js.map

@@ -1 +1 @@
-{"version":3,"file":"telemetry-manager.js","names":[],"sources":["../../src/telemetry/telemetry-manager.ts"],"sourcesContent":["import { getNodeAutoInstrumentations } from \"@opentelemetry/auto-instrumentations-node\";\nimport { OTLPLogExporter } from \"@opentelemetry/exporter-logs-otlp-proto\";\nimport { OTLPMetricExporter } from \"@opentelemetry/exporter-metrics-otlp-proto\";\nimport { OTLPTraceExporter } from \"@opentelemetry/exporter-trace-otlp-proto\";\nimport {\n  type Instrumentation,\n  registerInstrumentations as otelRegisterInstrumentations,\n} from \"@opentelemetry/instrumentation\";\nimport {\n  detectResources,\n  envDetector,\n  hostDetector,\n  processDetector,\n  type Resource,\n  resourceFromAttributes,\n} from \"@opentelemetry/resources\";\nimport { BatchLogRecordProcessor } from \"@opentelemetry/sdk-logs\";\nimport { PeriodicExportingMetricReader } from \"@opentelemetry/sdk-metrics\";\nimport { NodeSDK } from \"@opentelemetry/sdk-node\";\nimport {\n  ATTR_SERVICE_NAME,\n  ATTR_SERVICE_VERSION,\n} from \"@opentelemetry/semantic-conventions\";\nimport type { TelemetryOptions } from \"shared\";\nimport { createLogger } from \"../logging/logger\";\nimport { TelemetryProvider } from \"./telemetry-provider\";\nimport { AppKitSampler } from \"./trace-sampler\";\nimport type { TelemetryConfig } from \"./types\";\n\nconst logger = createLogger(\"telemetry\");\n\nexport class TelemetryManager {\n  private static readonly DEFAULT_EXPORT_INTERVAL_MS = 10000;\n  private static readonly DEFAULT_FALLBACK_APP_NAME = \"databricks-app\";\n\n  private static instance?: TelemetryManager;\n  private sdk?: NodeSDK;\n\n  /**\n   * Create a scoped telemetry provider for a specific plugin.\n   * The plugin's name will be used as the default tracer/meter name.\n   * @param pluginName - The name of the plugin to create scoped telemetry for\n   * @param telemetryConfig - The telemetry configuration for the plugin\n   * @returns A scoped telemetry instance for the plugin\n   */\n  static getProvider(\n    pluginName: string,\n    telemetryConfig?: TelemetryOptions,\n  ): TelemetryProvider {\n    const globalManager = TelemetryManager.getInstance();\n    return new TelemetryProvider(pluginName, globalManager, telemetryConfig);\n  }\n\n  private constructor() {}\n\n  static getInstance(): TelemetryManager {\n    if (!TelemetryManager.instance) {\n      TelemetryManager.instance = new TelemetryManager();\n    }\n    return TelemetryManager.instance;\n  }\n\n  static initialize(config: Partial<TelemetryConfig> = {}): void {\n    const instance = TelemetryManager.getInstance();\n    instance._initialize(config);\n  }\n\n  private _initialize(config: Partial<TelemetryConfig>): void {\n    if (this.sdk) return;\n\n    if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT) {\n      return;\n    }\n\n    try {\n      this.sdk = new NodeSDK({\n        resource: this.createResource(config),\n        autoDetectResources: false,\n        sampler: new AppKitSampler(),\n        traceExporter: new OTLPTraceExporter({ headers: config.headers }),\n        metricReaders: [\n          new PeriodicExportingMetricReader({\n            exporter: new OTLPMetricExporter({ headers: config.headers }),\n            exportIntervalMillis:\n              config.exportIntervalMs ||\n              TelemetryManager.DEFAULT_EXPORT_INTERVAL_MS,\n          }),\n        ],\n        logRecordProcessors: [\n          new BatchLogRecordProcessor(\n            new OTLPLogExporter({ headers: config.headers }),\n          ),\n        ],\n        instrumentations: this.getDefaultInstrumentations(),\n      });\n\n      this.sdk.start();\n      this.registerShutdown();\n      logger.debug(\"Initialized successfully\");\n    } catch (error) {\n      logger.error(\"Failed to initialize: %O\", error);\n    }\n  }\n\n  /**\n   * Register OpenTelemetry instrumentations.\n   * Can be called at any time, but recommended to call in plugin constructor.\n   * @param instrumentations - Array of OpenTelemetry instrumentations to register\n   */\n  registerInstrumentations(instrumentations: Instrumentation[]): void {\n    otelRegisterInstrumentations({\n      //  global providers set by NodeSDK.start()\n      instrumentations,\n    });\n  }\n\n  private createResource(config: Partial<TelemetryConfig>): Resource {\n    const serviceName =\n      config.serviceName ||\n      process.env.OTEL_SERVICE_NAME ||\n      process.env.DATABRICKS_APP_NAME ||\n      TelemetryManager.DEFAULT_FALLBACK_APP_NAME;\n    const initialResource = resourceFromAttributes({\n      [ATTR_SERVICE_NAME]: serviceName,\n      [ATTR_SERVICE_VERSION]: config.serviceVersion ?? undefined,\n    });\n    const detectedResource = detectResources({\n      detectors: [envDetector, hostDetector, processDetector],\n    });\n    return initialResource.merge(detectedResource);\n  }\n\n  private getDefaultInstrumentations(): Instrumentation[] {\n    return [\n      ...getNodeAutoInstrumentations({\n        //\n        // enabled as a part of the server plugin\n        //\n        \"@opentelemetry/instrumentation-http\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-express\": {\n          enabled: false,\n        },\n        //\n        // reduce noise\n        //\n        \"@opentelemetry/instrumentation-fs\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-dns\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-net\": {\n          enabled: false,\n        },\n      }),\n    ];\n  }\n\n  private registerShutdown() {\n    const shutdownFn = async () => {\n      await TelemetryManager.getInstance().shutdown();\n    };\n    process.once(\"SIGTERM\", shutdownFn);\n    process.once(\"SIGINT\", shutdownFn);\n  }\n\n  private async shutdown(): Promise<void> {\n    if (!this.sdk) {\n      return;\n    }\n\n    try {\n      await this.sdk.shutdown();\n      this.sdk = undefined;\n    } catch (error) {\n      logger.error(\"Error shutting down: %O\", error);\n    }\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;AA6BA,MAAM,SAAS,aAAa,YAAY;AAExC,IAAa,mBAAb,MAAa,iBAAiB;CAC5B,OAAwB,6BAA6B;CACrD,OAAwB,4BAA4B;CAEpD,OAAe;CACf,AAAQ;;;;;;;;CASR,OAAO,YACL,YACA,iBACmB;AAEnB,SAAO,IAAI,kBAAkB,YADP,iBAAiB,aAAa,EACI,gBAAgB;;CAG1E,AAAQ,cAAc;CAEtB,OAAO,cAAgC;AACrC,MAAI,CAAC,iBAAiB,SACpB,kBAAiB,WAAW,IAAI,kBAAkB;AAEpD,SAAO,iBAAiB;;CAG1B,OAAO,WAAW,SAAmC,EAAE,EAAQ;AAE7D,EADiB,iBAAiB,aAAa,CACtC,YAAY,OAAO;;CAG9B,AAAQ,YAAY,QAAwC;AAC1D,MAAI,KAAK,IAAK;AAEd,MAAI,CAAC,QAAQ,IAAI,4BACf;AAGF,MAAI;AACF,QAAK,MAAM,IAAI,QAAQ;IACrB,UAAU,KAAK,eAAe,OAAO;IACrC,qBAAqB;IACrB,SAAS,IAAI,eAAe;IAC5B,eAAe,IAAI,kBAAkB,EAAE,SAAS,OAAO,SAAS,CAAC;IACjE,eAAe,CACb,IAAI,8BAA8B;KAChC,UAAU,IAAI,mBAAmB,EAAE,SAAS,OAAO,SAAS,CAAC;KAC7D,sBACE,OAAO,oBACP,iBAAiB;KACpB,CAAC,CACH;IACD,qBAAqB,CACnB,IAAI,wBACF,IAAI,gBAAgB,EAAE,SAAS,OAAO,SAAS,CAAC,CACjD,CACF;IACD,kBAAkB,KAAK,4BAA4B;IACpD,CAAC;AAEF,QAAK,IAAI,OAAO;AAChB,QAAK,kBAAkB;AACvB,UAAO,MAAM,2BAA2B;WACjC,OAAO;AACd,UAAO,MAAM,4BAA4B,MAAM;;;;;;;;CASnD,yBAAyB,kBAA2C;AAClE,2BAA6B,EAE3B,kBACD,CAAC;;CAGJ,AAAQ,eAAe,QAA4C;EACjE,MAAM,cACJ,OAAO,eACP,QAAQ,IAAI,qBACZ,QAAQ,IAAI,uBACZ,iBAAiB;EACnB,MAAM,kBAAkB,uBAAuB;IAC5C,oBAAoB;IACpB,uBAAuB,OAAO,kBAAkB;GAClD,CAAC;EACF,MAAM,mBAAmB,gBAAgB,EACvC,WAAW;GAAC;GAAa;GAAc;GAAgB,EACxD,CAAC;AACF,SAAO,gBAAgB,MAAM,iBAAiB;;CAGhD,AAAQ,6BAAgD;AACtD,SAAO,CACL,GAAG,4BAA4B;GAI7B,uCAAuC,EACrC,SAAS,OACV;GACD,0CAA0C,EACxC,SAAS,OACV;GAID,qCAAqC,EACnC,SAAS,OACV;GACD,sCAAsC,EACpC,SAAS,OACV;GACD,sCAAsC,EACpC,SAAS,OACV;GACF,CAAC,CACH;;CAGH,AAAQ,mBAAmB;EACzB,MAAM,aAAa,YAAY;AAC7B,SAAM,iBAAiB,aAAa,CAAC,UAAU;;AAEjD,UAAQ,KAAK,WAAW,WAAW;AACnC,UAAQ,KAAK,UAAU,WAAW;;CAGpC,MAAc,WAA0B;AACtC,MAAI,CAAC,KAAK,IACR;AAGF,MAAI;AACF,SAAM,KAAK,IAAI,UAAU;AACzB,QAAK,MAAM;WACJ,OAAO;AACd,UAAO,MAAM,2BAA2B,MAAM"}
+{"version":3,"file":"telemetry-manager.js","names":[],"sources":["../../src/telemetry/telemetry-manager.ts"],"sourcesContent":["import { getNodeAutoInstrumentations } from \"@opentelemetry/auto-instrumentations-node\";\nimport { OTLPLogExporter } from \"@opentelemetry/exporter-logs-otlp-proto\";\nimport { OTLPMetricExporter } from \"@opentelemetry/exporter-metrics-otlp-proto\";\nimport { OTLPTraceExporter } from \"@opentelemetry/exporter-trace-otlp-proto\";\nimport {\n  type Instrumentation,\n  registerInstrumentations as otelRegisterInstrumentations,\n} from \"@opentelemetry/instrumentation\";\nimport {\n  detectResources,\n  envDetector,\n  hostDetector,\n  processDetector,\n  type Resource,\n  resourceFromAttributes,\n} from \"@opentelemetry/resources\";\nimport { BatchLogRecordProcessor } from \"@opentelemetry/sdk-logs\";\nimport { PeriodicExportingMetricReader } from \"@opentelemetry/sdk-metrics\";\nimport { NodeSDK } from \"@opentelemetry/sdk-node\";\nimport {\n  ATTR_SERVICE_NAME,\n  ATTR_SERVICE_VERSION,\n} from \"@opentelemetry/semantic-conventions\";\nimport type { TelemetryOptions } from \"shared\";\nimport { createLogger } from \"../logging/logger\";\nimport { TelemetryProvider } from \"./telemetry-provider\";\nimport { AppKitSampler } from \"./trace-sampler\";\nimport type { TelemetryConfig } from \"./types\";\n\nconst logger = createLogger(\"telemetry\");\n\nexport class TelemetryManager {\n  private static readonly DEFAULT_EXPORT_INTERVAL_MS = 10000;\n  private static readonly DEFAULT_FALLBACK_APP_NAME = \"databricks-app\";\n\n  private static instance?: TelemetryManager;\n  private sdk?: NodeSDK;\n  private shutdownPromise?: Promise<void>;\n\n  /**\n   * Create a scoped telemetry provider for a specific plugin.\n   * The plugin's name will be used as the default tracer/meter name.\n   * @param pluginName - The name of the plugin to create scoped telemetry for\n   * @param telemetryConfig - The telemetry configuration for the plugin\n   * @returns A scoped telemetry instance for the plugin\n   */\n  static getProvider(\n    pluginName: string,\n    telemetryConfig?: TelemetryOptions,\n  ): TelemetryProvider {\n    const globalManager = TelemetryManager.getInstance();\n    return new TelemetryProvider(pluginName, globalManager, telemetryConfig);\n  }\n\n  private constructor() {}\n\n  static getInstance(): TelemetryManager {\n    if (!TelemetryManager.instance) {\n      TelemetryManager.instance = new TelemetryManager();\n    }\n    return TelemetryManager.instance;\n  }\n\n  static initialize(config: Partial<TelemetryConfig> = {}): void {\n    const instance = TelemetryManager.getInstance();\n    instance._initialize(config);\n  }\n\n  private _initialize(config: Partial<TelemetryConfig>): void {\n    if (this.sdk) return;\n\n    if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT) {\n      return;\n    }\n\n    try {\n      this.sdk = new NodeSDK({\n        resource: this.createResource(config),\n        autoDetectResources: false,\n        sampler: new AppKitSampler(),\n        traceExporter: new OTLPTraceExporter({ headers: config.headers }),\n        metricReaders: [\n          new PeriodicExportingMetricReader({\n            exporter: new OTLPMetricExporter({ headers: config.headers }),\n            exportIntervalMillis:\n              config.exportIntervalMs ||\n              TelemetryManager.DEFAULT_EXPORT_INTERVAL_MS,\n          }),\n        ],\n        logRecordProcessors: [\n          new BatchLogRecordProcessor(\n            new OTLPLogExporter({ headers: config.headers }),\n          ),\n        ],\n        instrumentations: this.getDefaultInstrumentations(),\n      });\n\n      this.sdk.start();\n      logger.debug(\"Initialized successfully\");\n    } catch (error) {\n      logger.error(\"Failed to initialize: %O\", error);\n    }\n  }\n\n  /**\n   * Register OpenTelemetry instrumentations.\n   * Can be called at any time, but recommended to call in plugin constructor.\n   * @param instrumentations - Array of OpenTelemetry instrumentations to register\n   */\n  registerInstrumentations(instrumentations: Instrumentation[]): void {\n    otelRegisterInstrumentations({\n      //  global providers set by NodeSDK.start()\n      instrumentations,\n    });\n  }\n\n  private createResource(config: Partial<TelemetryConfig>): Resource {\n    const serviceName =\n      config.serviceName ||\n      process.env.OTEL_SERVICE_NAME ||\n      process.env.DATABRICKS_APP_NAME ||\n      TelemetryManager.DEFAULT_FALLBACK_APP_NAME;\n    const initialResource = resourceFromAttributes({\n      [ATTR_SERVICE_NAME]: serviceName,\n      [ATTR_SERVICE_VERSION]: config.serviceVersion ?? undefined,\n    });\n    const detectedResource = detectResources({\n      detectors: [envDetector, hostDetector, processDetector],\n    });\n    return initialResource.merge(detectedResource);\n  }\n\n  private getDefaultInstrumentations(): Instrumentation[] {\n    return [\n      ...getNodeAutoInstrumentations({\n        //\n        // enabled as a part of the server plugin\n        //\n        \"@opentelemetry/instrumentation-http\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-express\": {\n          enabled: false,\n        },\n        //\n        // reduce noise\n        //\n        \"@opentelemetry/instrumentation-fs\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-dns\": {\n          enabled: false,\n        },\n        \"@opentelemetry/instrumentation-net\": {\n          enabled: false,\n        },\n      }),\n    ];\n  }\n\n  /**\n   * Flush and shut down the OpenTelemetry SDK.\n   *\n   * Idempotent: the SDK reference is cleared synchronously and concurrent\n   * or repeated calls await the same in-flight flush. Awaited by the core\n   * lifecycle manager during graceful shutdown — that manager owns the\n   * process signal handlers, so telemetry no longer registers its own.\n   */\n  async shutdown(): Promise<void> {\n    if (this.sdk) {\n      const sdk = this.sdk;\n      this.sdk = undefined;\n      this.shutdownPromise = (async () => {\n        try {\n          await sdk.shutdown();\n        } catch (error) {\n          logger.error(\"Error shutting down: %O\", error);\n        }\n      })();\n    }\n\n    return this.shutdownPromise;\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;AA6BA,MAAM,SAAS,aAAa,YAAY;AAExC,IAAa,mBAAb,MAAa,iBAAiB;CAC5B,OAAwB,6BAA6B;CACrD,OAAwB,4BAA4B;CAEpD,OAAe;CACf,AAAQ;CACR,AAAQ;;;;;;;;CASR,OAAO,YACL,YACA,iBACmB;AAEnB,SAAO,IAAI,kBAAkB,YADP,iBAAiB,aAAa,EACI,gBAAgB;;CAG1E,AAAQ,cAAc;CAEtB,OAAO,cAAgC;AACrC,MAAI,CAAC,iBAAiB,SACpB,kBAAiB,WAAW,IAAI,kBAAkB;AAEpD,SAAO,iBAAiB;;CAG1B,OAAO,WAAW,SAAmC,EAAE,EAAQ;AAE7D,EADiB,iBAAiB,aAAa,CACtC,YAAY,OAAO;;CAG9B,AAAQ,YAAY,QAAwC;AAC1D,MAAI,KAAK,IAAK;AAEd,MAAI,CAAC,QAAQ,IAAI,4BACf;AAGF,MAAI;AACF,QAAK,MAAM,IAAI,QAAQ;IACrB,UAAU,KAAK,eAAe,OAAO;IACrC,qBAAqB;IACrB,SAAS,IAAI,eAAe;IAC5B,eAAe,IAAI,kBAAkB,EAAE,SAAS,OAAO,SAAS,CAAC;IACjE,eAAe,CACb,IAAI,8BAA8B;KAChC,UAAU,IAAI,mBAAmB,EAAE,SAAS,OAAO,SAAS,CAAC;KAC7D,sBACE,OAAO,oBACP,iBAAiB;KACpB,CAAC,CACH;IACD,qBAAqB,CACnB,IAAI,wBACF,IAAI,gBAAgB,EAAE,SAAS,OAAO,SAAS,CAAC,CACjD,CACF;IACD,kBAAkB,KAAK,4BAA4B;IACpD,CAAC;AAEF,QAAK,IAAI,OAAO;AAChB,UAAO,MAAM,2BAA2B;WACjC,OAAO;AACd,UAAO,MAAM,4BAA4B,MAAM;;;;;;;;CASnD,yBAAyB,kBAA2C;AAClE,2BAA6B,EAE3B,kBACD,CAAC;;CAGJ,AAAQ,eAAe,QAA4C;EACjE,MAAM,cACJ,OAAO,eACP,QAAQ,IAAI,qBACZ,QAAQ,IAAI,uBACZ,iBAAiB;EACnB,MAAM,kBAAkB,uBAAuB;IAC5C,oBAAoB;IACpB,uBAAuB,OAAO,kBAAkB;GAClD,CAAC;EACF,MAAM,mBAAmB,gBAAgB,EACvC,WAAW;GAAC;GAAa;GAAc;GAAgB,EACxD,CAAC;AACF,SAAO,gBAAgB,MAAM,iBAAiB;;CAGhD,AAAQ,6BAAgD;AACtD,SAAO,CACL,GAAG,4BAA4B;GAI7B,uCAAuC,EACrC,SAAS,OACV;GACD,0CAA0C,EACxC,SAAS,OACV;GAID,qCAAqC,EACnC,SAAS,OACV;GACD,sCAAsC,EACpC,SAAS,OACV;GACD,sCAAsC,EACpC,SAAS,OACV;GACF,CAAC,CACH;;;;;;;;;;CAWH,MAAM,WAA0B;AAC9B,MAAI,KAAK,KAAK;GACZ,MAAM,MAAM,KAAK;AACjB,QAAK,MAAM;AACX,QAAK,mBAAmB,YAAY;AAClC,QAAI;AACF,WAAM,IAAI,UAAU;aACb,OAAO;AACd,YAAO,MAAM,2BAA2B,MAAM;;OAE9C;;AAGN,SAAO,KAAK"}

package/dist/type-generator/query-registry.js

@@ -10,6 +10,7 @@ import { WorkspaceClient } from "@databricks/sdk-experimental";
 import pc from "picocolors";
 import fs from "node:fs/promises";
 import path from "node:path";
+import { tableFromIPC } from "apache-arrow";
 
 //#region src/type-generator/query-registry.ts
 const logger = createLogger("type-generator:query-registry");
@@ -95,12 +96,41 @@ function formatParametersType(sql) {
 		return `/** ${sqlType || "any"} - use ${helper} */\n      ${p}: ${markerType}`;
 	}).join(";\n      ")};\n    }` : "Record<string, never>";
 }
+/**
+* Decode a base64 Arrow IPC attachment from a DESCRIBE QUERY response and
+* extract column metadata. Returns the same shape as rows parsed from the
+* legacy data_array path.
+*
+* IMPORTANT: a DESCRIBE QUERY response is itself a result *table* with rows
+* shaped like `(col_name, data_type, comment)` describing the user query's
+* output schema. We must read those rows — NOT `table.schema.fields`, which
+* would describe DESCRIBE QUERY's own output (`col_name`, `data_type`,
+* `comment`) and yield bogus types for every query.
+*/
+function columnsFromArrowAttachment(attachment) {
+	return tableFromIPC(Buffer.from(attachment, "base64")).toArray().map((row) => {
+		const obj = row.toJSON();
+		return {
+			name: typeof obj.col_name === "string" ? obj.col_name : "",
+			type_name: typeof obj.data_type === "string" ? obj.data_type.toUpperCase() : "STRING",
+			comment: typeof obj.comment === "string" && obj.comment !== "" ? obj.comment : void 0
+		};
+	});
+}
 function convertToQueryType(result, sql, queryName) {
-	const columns = (result.result?.data_array || []).map((row) => ({
+	let columns = (result.result?.data_array || []).map((row) => ({
 		name: row[0] || "",
 		type_name: row[1]?.toUpperCase() || "STRING",
 		comment: row[2] || void 0
 	}));
+	if (columns.length === 0 && result.result?.attachment) {
+		logger.debug("data_array empty, decoding Arrow IPC attachment for schema");
+		try {
+			columns = columnsFromArrowAttachment(result.result.attachment);
+		} catch (err) {
+			logger.warn("Failed to decode Arrow IPC attachment: %s", err instanceof Error ? err.message : String(err));
+		}
+	}
 	const paramsType = formatParametersType(sql);
 	const resultFields = columns.map((column) => {
 		const mappedType = typeMap[normalizeTypeName(column.type_name)] || "unknown";
@@ -444,7 +474,7 @@ async function generateQueriesFromDescribe(queryFolder, warehouseId, options = {
 				const result = await describeAdaptive(client, `DESCRIBE QUERY ${cleanedSql}`, warehouseId, describeFormat);
 				completed++;
 				spinner.update(`Describing ${total} ${total === 1 ? "query" : "queries"} (${completed}/${total})`);
-				logger.debug("DESCRIBE result for %s: state=%s, rows=%d", queryName, result.status.state, result.result?.data_array?.length ?? 0);
+				logger.debug("DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s", queryName, result.status.state, result.result?.data_array?.length ?? 0, !!result.result?.attachment);
 				if (result.status.state === "FAILED") {
 					const sqlError = result.status.error?.message || "Query execution failed";
 					return {