@drarzter/kafka-client 0.5.6 → 0.6.3

package/dist/{envelope-LeO5e3ob.d.mts → types-zFbQH_Cy.d.mts}

@@ -1,15 +1,44 @@
+/**
+ * Context passed as the second argument to `SchemaLike.parse()`.
+ * Enables schema-registry adapters, version-aware migration, and
+ * header-driven parsing without coupling validators to Kafka internals.
+ *
+ * All fields are optional-friendly — validators that don't need the context
+ * can simply ignore the second argument.
+ */
+interface SchemaParseContext {
+    /** Topic the message was produced to / consumed from. */
+    topic: string;
+    /** Decoded message headers (envelope headers included). */
+    headers: MessageHeaders;
+    /** Value of the `x-schema-version` header, defaults to `1`. */
+    version: number;
+}
 /**
  * Any validation library with a `.parse()` method.
  * Works with Zod, Valibot, ArkType, or any custom validator.
  *
+ * The optional `ctx` argument carries topic/header/version metadata so
+ * validators can perform schema-registry lookups or version-aware migrations.
+ * Existing validators that only use the first argument continue to work
+ * unchanged — the second argument is silently ignored.
+ *
  * @example
  * ```ts
  * import { z } from 'zod';
  * const schema: SchemaLike<{ id: string }> = z.object({ id: z.string() });
+ *
+ * // Context-aware validator:
+ * const schema: SchemaLike<MyType> = {
+ *   parse(data, ctx) {
+ *     const version = ctx?.version ?? 1;
+ *     return version >= 2 ? migrateV1toV2(data) : validateV1(data);
+ *   }
+ * };
  * ```
  */
 interface SchemaLike<T = any> {
-    parse(data: unknown): T | Promise<T>;
+    parse(data: unknown, ctx?: SchemaParseContext): T | Promise<T>;
 }
 /** Infer the output type from a SchemaLike. */
 type InferSchema<S extends SchemaLike> = S extends SchemaLike<infer T> ? T : never;
@@ -32,8 +61,8 @@ interface TopicDescriptor<N extends string = string, M extends Record<string, an
  *
  * @example
  * ```ts
- * // Without schema — type provided explicitly:
- * const OrderCreated = topic('order.created')<{ orderId: string; amount: number }>();
+ * // Without schema — explicit type via .type<T>():
+ * const OrderCreated = topic('order.created').type<{ orderId: string; amount: number }>();
  *
  * // With schema — type inferred from schema:
  * const OrderCreated = topic('order.created').schema(z.object({
@@ -50,16 +79,17 @@ interface TopicDescriptor<N extends string = string, M extends Record<string, an
  * ```
  */
 declare function topic<N extends string>(name: N): {
-    <M extends Record<string, any>>(): TopicDescriptor<N, M>;
-    schema<S extends SchemaLike<Record<string, any>>>(schema: S): TopicDescriptor<N, InferSchema<S>>;
+    /** Provide an explicit message type without a runtime schema. */
+    type: <M extends Record<string, any>>() => TopicDescriptor<N, M>;
+    schema: <S extends SchemaLike<Record<string, any>>>(schema: S) => TopicDescriptor<N, InferSchema<S>>;
 };
 /**
  * Build a topic-message map type from a union of TopicDescriptors.
  *
  * @example
  * ```ts
- * const OrderCreated = topic('order.created')<{ orderId: string }>();
- * const OrderCompleted = topic('order.completed')<{ completedAt: string }>();
+ * const OrderCreated = topic('order.created').type<{ orderId: string }>();
+ * const OrderCompleted = topic('order.completed').type<{ completedAt: string }>();
  *
  * type MyTopics = TopicsFrom<typeof OrderCreated | typeof OrderCompleted>;
  * // { 'order.created': { orderId: string }; 'order.completed': { completedAt: string } }
@@ -69,6 +99,77 @@ type TopicsFrom<D extends TopicDescriptor<any, any>> = {
     [K in D as K["__topic"]]: K["__type"];
 };
 
+declare const HEADER_EVENT_ID = "x-event-id";
+declare const HEADER_CORRELATION_ID = "x-correlation-id";
+declare const HEADER_TIMESTAMP = "x-timestamp";
+declare const HEADER_SCHEMA_VERSION = "x-schema-version";
+declare const HEADER_TRACEPARENT = "traceparent";
+/**
+ * Typed wrapper combining a parsed message payload with Kafka metadata
+ * and envelope headers.
+ *
+ * On **send**, the library auto-generates envelope headers
+ * (`x-event-id`, `x-correlation-id`, `x-timestamp`, `x-schema-version`).
+ *
+ * On **consume**, the library extracts those headers and assembles
+ * an `EventEnvelope` that is passed to the handler.
+ */
+interface EventEnvelope<T> {
+    /** Deserialized + validated message body. */
+    payload: T;
+    /** Topic the message was produced to / consumed from. */
+    topic: string;
+    /** Kafka partition (consume-side only, `-1` on send). */
+    partition: number;
+    /** Kafka offset (consume-side only, empty string on send). */
+    offset: string;
+    /** ISO-8601 timestamp set by the producer. */
+    timestamp: string;
+    /** Unique ID for this event (UUID v4). */
+    eventId: string;
+    /** Correlation ID — auto-propagated via AsyncLocalStorage. */
+    correlationId: string;
+    /** Schema version of the payload. */
+    schemaVersion: number;
+    /** W3C Trace Context `traceparent` header (set by OTel instrumentation). */
+    traceparent?: string;
+    /** All decoded Kafka headers for extensibility. */
+    headers: MessageHeaders;
+}
+interface EnvelopeCtx {
+    correlationId: string;
+    traceparent?: string;
+}
+/** Read the current envelope context (correlationId / traceparent) from ALS. */
+declare function getEnvelopeContext(): EnvelopeCtx | undefined;
+/** Execute `fn` inside an envelope context so nested sends inherit correlationId. */
+declare function runWithEnvelopeContext<R>(ctx: EnvelopeCtx, fn: () => R): R;
+/** Options accepted by `buildEnvelopeHeaders`. */
+interface EnvelopeHeaderOptions {
+    correlationId?: string;
+    schemaVersion?: number;
+    eventId?: string;
+    headers?: MessageHeaders;
+}
+/**
+ * Generate envelope headers for the send path.
+ *
+ * Priority for `correlationId`:
+ *   explicit option → ALS context → new UUID.
+ */
+declare function buildEnvelopeHeaders(options?: EnvelopeHeaderOptions): MessageHeaders;
+/**
+ * Decode kafkajs headers (`Record<string, Buffer | string | undefined>`)
+ * into plain `Record<string, string>`.
+ */
+declare function decodeHeaders(raw: Record<string, Buffer | string | (Buffer | string)[] | undefined> | undefined): MessageHeaders;
+/**
+ * Build an `EventEnvelope` from a consumed kafkajs message.
+ * Tolerates missing envelope headers — generates defaults so messages
+ * from non-envelope producers still work.
+ */
+declare function extractEnvelope<T>(payload: T, headers: MessageHeaders, topic: string, partition: number, offset: string): EventEnvelope<T>;
+
 /**
  * Mapping of topic names to their message types.
  * Define this interface to get type-safe publish/subscribe across your app.
@@ -165,6 +266,12 @@ interface ConsumerOptions<T extends TopicMapConstraint<T> = TTopicMessageMap> {
      * On exhaustion, messages go to `<topic>.dlq` (if `dlq: true`) or `onMessageLost`.
      */
     retryTopics?: boolean;
+    /**
+     * Timeout (ms) for waiting until each retry level consumer receives partition
+     * assignments after `startConsumer` connects. Default: `10000`.
+     * Increase this when the broker is slow to rebalance.
+     */
+    retryTopicAssignmentTimeoutMs?: number;
     /**
      * Log a warning if the message handler has not resolved within this window (ms).
      * The handler is not cancelled — this is a diagnostic aid to surface stuck handlers
@@ -196,6 +303,21 @@ interface ConsumerInterceptor<T extends TopicMapConstraint<T> = TTopicMessageMap
     /** Called when the message handler throws. */
     onError?(envelope: EventEnvelope<T[keyof T]>, error: Error): Promise<void> | void;
 }
+/**
+ * Return value of `KafkaInstrumentation.beforeConsume`.
+ *
+ * - `() => void` — legacy form: a cleanup function called after the handler.
+ * - Object form:
+ *   - `cleanup?()` — called after the handler (same as the legacy function form).
+ *   - `wrap?(fn)` — wraps the handler execution; call `fn()` inside the desired
+ *     async context (e.g. `context.with(spanCtx, fn)` for OpenTelemetry). Multiple
+ *     wraps from different instrumentations are composed in declaration order,
+ *     so the first instrumentation's wrap is the outermost.
+ */
+type BeforeConsumeResult = (() => void) | {
+    cleanup?(): void;
+    wrap?(fn: () => Promise<void>): Promise<void>;
+};
 /**
  * Client-wide instrumentation hooks for both send and consume paths.
  * Use this for cross-cutting concerns like tracing and metrics.
@@ -207,8 +329,13 @@ interface KafkaInstrumentation {
     beforeSend?(topic: string, headers: MessageHeaders): void;
     /** Called after a successful send. */
     afterSend?(topic: string): void;
-    /** Called before the consumer handler. Return a cleanup function called after the handler. */
-    beforeConsume?(envelope: EventEnvelope<any>): (() => void) | void;
+    /**
+     * Called before the consumer handler.
+     * Return a cleanup function (legacy) or a `BeforeConsumeResult` object with
+     * optional `cleanup` and `wrap`. Use `wrap` to run the handler inside a
+     * specific async context (e.g. an active OpenTelemetry span).
+     */
+    beforeConsume?(envelope: EventEnvelope<any>): BeforeConsumeResult | void;
     /** Called when the consumer handler throws. */
     onConsumeError?(envelope: EventEnvelope<any>, error: Error): void;
 }
@@ -226,13 +353,19 @@ interface ConsumerHandle {
     /** Stop this consumer. Equivalent to calling `client.stopConsumer(groupId)`. */
     stop(): Promise<void>;
 }
+/** Result returned by `KafkaClient.checkStatus()`. */
+type KafkaHealthResult = {
+    status: "up";
+    clientId: string;
+    topics: string[];
+} | {
+    status: "down";
+    clientId: string;
+    error: string;
+};
 /** Interface describing all public methods of the Kafka client. */
 interface IKafkaClient<T extends TopicMapConstraint<T>> {
-    checkStatus(): Promise<{
-        status: "up";
-        clientId: string;
-        topics: string[];
-    }>;
+    checkStatus(): Promise<KafkaHealthResult>;
     /**
      * Query the consumer group lag per partition using the admin API.
      * Lag = (broker high-watermark offset) − (last committed offset).
@@ -257,8 +390,18 @@ interface IKafkaClient<T extends TopicMapConstraint<T>> {
     sendMessage<K extends keyof T>(topic: K, message: T[K], options?: SendOptions): Promise<void>;
     sendBatch<K extends keyof T>(topic: K, messages: Array<BatchMessageItem<T[K]>>): Promise<void>;
     transaction(fn: (ctx: TransactionContext<T>) => Promise<void>): Promise<void>;
-    getClientId: () => ClientId;
-    disconnect(): Promise<void>;
+    getClientId(): ClientId;
+    /**
+     * Drain in-flight handlers, then disconnect all producers, consumers, and admin.
+     * @param drainTimeoutMs Max ms to wait for in-flight handlers (default 30 000).
+     */
+    disconnect(drainTimeoutMs?: number): Promise<void>;
+    /**
+     * Register SIGTERM / SIGINT signal handlers that drain in-flight messages before
+     * disconnecting. Call once after constructing the client in non-NestJS apps.
+     * NestJS apps get drain automatically via `onModuleDestroy` → `disconnect()`.
+     */
+    enableGracefulShutdown(signals?: NodeJS.Signals[], drainTimeoutMs?: number): void;
 }
 /**
  * Logger interface for KafkaClient.
@@ -325,75 +468,4 @@ interface SubscribeRetryOptions {
     backoffMs?: number;
 }
 
-declare const HEADER_EVENT_ID = "x-event-id";
-declare const HEADER_CORRELATION_ID = "x-correlation-id";
-declare const HEADER_TIMESTAMP = "x-timestamp";
-declare const HEADER_SCHEMA_VERSION = "x-schema-version";
-declare const HEADER_TRACEPARENT = "traceparent";
-/**
- * Typed wrapper combining a parsed message payload with Kafka metadata
- * and envelope headers.
- *
- * On **send**, the library auto-generates envelope headers
- * (`x-event-id`, `x-correlation-id`, `x-timestamp`, `x-schema-version`).
- *
- * On **consume**, the library extracts those headers and assembles
- * an `EventEnvelope` that is passed to the handler.
- */
-interface EventEnvelope<T> {
-    /** Deserialized + validated message body. */
-    payload: T;
-    /** Topic the message was produced to / consumed from. */
-    topic: string;
-    /** Kafka partition (consume-side only, `-1` on send). */
-    partition: number;
-    /** Kafka offset (consume-side only, empty string on send). */
-    offset: string;
-    /** ISO-8601 timestamp set by the producer. */
-    timestamp: string;
-    /** Unique ID for this event (UUID v4). */
-    eventId: string;
-    /** Correlation ID — auto-propagated via AsyncLocalStorage. */
-    correlationId: string;
-    /** Schema version of the payload. */
-    schemaVersion: number;
-    /** W3C Trace Context `traceparent` header (set by OTel instrumentation). */
-    traceparent?: string;
-    /** All decoded Kafka headers for extensibility. */
-    headers: MessageHeaders;
-}
-interface EnvelopeCtx {
-    correlationId: string;
-    traceparent?: string;
-}
-/** Read the current envelope context (correlationId / traceparent) from ALS. */
-declare function getEnvelopeContext(): EnvelopeCtx | undefined;
-/** Execute `fn` inside an envelope context so nested sends inherit correlationId. */
-declare function runWithEnvelopeContext<R>(ctx: EnvelopeCtx, fn: () => R): R;
-/** Options accepted by `buildEnvelopeHeaders`. */
-interface EnvelopeHeaderOptions {
-    correlationId?: string;
-    schemaVersion?: number;
-    eventId?: string;
-    headers?: MessageHeaders;
-}
-/**
- * Generate envelope headers for the send path.
- *
- * Priority for `correlationId`:
- *   explicit option → ALS context → new UUID.
- */
-declare function buildEnvelopeHeaders(options?: EnvelopeHeaderOptions): MessageHeaders;
-/**
- * Decode kafkajs headers (`Record<string, Buffer | string | undefined>`)
- * into plain `Record<string, string>`.
- */
-declare function decodeHeaders(raw: Record<string, Buffer | string | (Buffer | string)[] | undefined> | undefined): MessageHeaders;
-/**
- * Build an `EventEnvelope` from a consumed kafkajs message.
- * Tolerates missing envelope headers — generates defaults so messages
- * from non-envelope producers still work.
- */
-declare function extractEnvelope<T>(payload: T, headers: MessageHeaders, topic: string, partition: number, offset: string): EventEnvelope<T>;
-
-export { type BatchMessageItem as B, type ClientId as C, type EnvelopeHeaderOptions as E, type GroupId as G, HEADER_CORRELATION_ID as H, type IKafkaClient as I, type KafkaInstrumentation as K, type MessageHeaders as M, type RetryOptions as R, type SchemaLike as S, type TopicMapConstraint as T, type KafkaClientOptions as a, type ConsumerOptions as b, type TopicDescriptor as c, type BatchMeta as d, type ConsumerHandle as e, type ConsumerInterceptor as f, type EventEnvelope as g, HEADER_EVENT_ID as h, HEADER_SCHEMA_VERSION as i, HEADER_TIMESTAMP as j, HEADER_TRACEPARENT as k, type InferSchema as l, type KafkaLogger as m, type MessageLostContext as n, type SendOptions as o, type SubscribeRetryOptions as p, type TTopicMessageMap as q, type TopicsFrom as r, type TransactionContext as s, buildEnvelopeHeaders as t, decodeHeaders as u, extractEnvelope as v, getEnvelopeContext as w, runWithEnvelopeContext as x, topic as y };
+export { runWithEnvelopeContext as A, type BatchMessageItem as B, type ClientId as C, topic as D, type EnvelopeHeaderOptions as E, type GroupId as G, HEADER_CORRELATION_ID as H, type IKafkaClient as I, type KafkaInstrumentation as K, type MessageHeaders as M, type RetryOptions as R, type SchemaLike as S, type TopicMapConstraint as T, type KafkaClientOptions as a, type ConsumerOptions as b, type TopicDescriptor as c, type KafkaHealthResult as d, type BatchMeta as e, type BeforeConsumeResult as f, type ConsumerHandle as g, type ConsumerInterceptor as h, type EventEnvelope as i, HEADER_EVENT_ID as j, HEADER_SCHEMA_VERSION as k, HEADER_TIMESTAMP as l, HEADER_TRACEPARENT as m, type InferSchema as n, type KafkaLogger as o, type MessageLostContext as p, type SchemaParseContext as q, type SendOptions as r, type SubscribeRetryOptions as s, type TTopicMessageMap as t, type TopicsFrom as u, type TransactionContext as v, buildEnvelopeHeaders as w, decodeHeaders as x, extractEnvelope as y, getEnvelopeContext as z };

package/dist/{envelope-LeO5e3ob.d.ts → types-zFbQH_Cy.d.ts}

@@ -1,15 +1,44 @@
+/**
+ * Context passed as the second argument to `SchemaLike.parse()`.
+ * Enables schema-registry adapters, version-aware migration, and
+ * header-driven parsing without coupling validators to Kafka internals.
+ *
+ * All fields are optional-friendly — validators that don't need the context
+ * can simply ignore the second argument.
+ */
+interface SchemaParseContext {
+    /** Topic the message was produced to / consumed from. */
+    topic: string;
+    /** Decoded message headers (envelope headers included). */
+    headers: MessageHeaders;
+    /** Value of the `x-schema-version` header, defaults to `1`. */
+    version: number;
+}
 /**
  * Any validation library with a `.parse()` method.
  * Works with Zod, Valibot, ArkType, or any custom validator.
  *
+ * The optional `ctx` argument carries topic/header/version metadata so
+ * validators can perform schema-registry lookups or version-aware migrations.
+ * Existing validators that only use the first argument continue to work
+ * unchanged — the second argument is silently ignored.
+ *
  * @example
  * ```ts
  * import { z } from 'zod';
  * const schema: SchemaLike<{ id: string }> = z.object({ id: z.string() });
+ *
+ * // Context-aware validator:
+ * const schema: SchemaLike<MyType> = {
+ *   parse(data, ctx) {
+ *     const version = ctx?.version ?? 1;
+ *     return version >= 2 ? migrateV1toV2(data) : validateV1(data);
+ *   }
+ * };
  * ```
  */
 interface SchemaLike<T = any> {
-    parse(data: unknown): T | Promise<T>;
+    parse(data: unknown, ctx?: SchemaParseContext): T | Promise<T>;
 }
 /** Infer the output type from a SchemaLike. */
 type InferSchema<S extends SchemaLike> = S extends SchemaLike<infer T> ? T : never;
@@ -32,8 +61,8 @@ interface TopicDescriptor<N extends string = string, M extends Record<string, an
  *
  * @example
  * ```ts
- * // Without schema — type provided explicitly:
- * const OrderCreated = topic('order.created')<{ orderId: string; amount: number }>();
+ * // Without schema — explicit type via .type<T>():
+ * const OrderCreated = topic('order.created').type<{ orderId: string; amount: number }>();
  *
  * // With schema — type inferred from schema:
  * const OrderCreated = topic('order.created').schema(z.object({
@@ -50,16 +79,17 @@ interface TopicDescriptor<N extends string = string, M extends Record<string, an
  * ```
  */
 declare function topic<N extends string>(name: N): {
-    <M extends Record<string, any>>(): TopicDescriptor<N, M>;
-    schema<S extends SchemaLike<Record<string, any>>>(schema: S): TopicDescriptor<N, InferSchema<S>>;
+    /** Provide an explicit message type without a runtime schema. */
+    type: <M extends Record<string, any>>() => TopicDescriptor<N, M>;
+    schema: <S extends SchemaLike<Record<string, any>>>(schema: S) => TopicDescriptor<N, InferSchema<S>>;
 };
 /**
  * Build a topic-message map type from a union of TopicDescriptors.
  *
  * @example
  * ```ts
- * const OrderCreated = topic('order.created')<{ orderId: string }>();
- * const OrderCompleted = topic('order.completed')<{ completedAt: string }>();
+ * const OrderCreated = topic('order.created').type<{ orderId: string }>();
+ * const OrderCompleted = topic('order.completed').type<{ completedAt: string }>();
  *
  * type MyTopics = TopicsFrom<typeof OrderCreated | typeof OrderCompleted>;
  * // { 'order.created': { orderId: string }; 'order.completed': { completedAt: string } }
@@ -69,6 +99,77 @@ type TopicsFrom<D extends TopicDescriptor<any, any>> = {
     [K in D as K["__topic"]]: K["__type"];
 };
 
+declare const HEADER_EVENT_ID = "x-event-id";
+declare const HEADER_CORRELATION_ID = "x-correlation-id";
+declare const HEADER_TIMESTAMP = "x-timestamp";
+declare const HEADER_SCHEMA_VERSION = "x-schema-version";
+declare const HEADER_TRACEPARENT = "traceparent";
+/**
+ * Typed wrapper combining a parsed message payload with Kafka metadata
+ * and envelope headers.
+ *
+ * On **send**, the library auto-generates envelope headers
+ * (`x-event-id`, `x-correlation-id`, `x-timestamp`, `x-schema-version`).
+ *
+ * On **consume**, the library extracts those headers and assembles
+ * an `EventEnvelope` that is passed to the handler.
+ */
+interface EventEnvelope<T> {
+    /** Deserialized + validated message body. */
+    payload: T;
+    /** Topic the message was produced to / consumed from. */
+    topic: string;
+    /** Kafka partition (consume-side only, `-1` on send). */
+    partition: number;
+    /** Kafka offset (consume-side only, empty string on send). */
+    offset: string;
+    /** ISO-8601 timestamp set by the producer. */
+    timestamp: string;
+    /** Unique ID for this event (UUID v4). */
+    eventId: string;
+    /** Correlation ID — auto-propagated via AsyncLocalStorage. */
+    correlationId: string;
+    /** Schema version of the payload. */
+    schemaVersion: number;
+    /** W3C Trace Context `traceparent` header (set by OTel instrumentation). */
+    traceparent?: string;
+    /** All decoded Kafka headers for extensibility. */
+    headers: MessageHeaders;
+}
+interface EnvelopeCtx {
+    correlationId: string;
+    traceparent?: string;
+}
+/** Read the current envelope context (correlationId / traceparent) from ALS. */
+declare function getEnvelopeContext(): EnvelopeCtx | undefined;
+/** Execute `fn` inside an envelope context so nested sends inherit correlationId. */
+declare function runWithEnvelopeContext<R>(ctx: EnvelopeCtx, fn: () => R): R;
+/** Options accepted by `buildEnvelopeHeaders`. */
+interface EnvelopeHeaderOptions {
+    correlationId?: string;
+    schemaVersion?: number;
+    eventId?: string;
+    headers?: MessageHeaders;
+}
+/**
+ * Generate envelope headers for the send path.
+ *
+ * Priority for `correlationId`:
+ *   explicit option → ALS context → new UUID.
+ */
+declare function buildEnvelopeHeaders(options?: EnvelopeHeaderOptions): MessageHeaders;
+/**
+ * Decode kafkajs headers (`Record<string, Buffer | string | undefined>`)
+ * into plain `Record<string, string>`.
+ */
+declare function decodeHeaders(raw: Record<string, Buffer | string | (Buffer | string)[] | undefined> | undefined): MessageHeaders;
+/**
+ * Build an `EventEnvelope` from a consumed kafkajs message.
+ * Tolerates missing envelope headers — generates defaults so messages
+ * from non-envelope producers still work.
+ */
+declare function extractEnvelope<T>(payload: T, headers: MessageHeaders, topic: string, partition: number, offset: string): EventEnvelope<T>;
+
 /**
  * Mapping of topic names to their message types.
  * Define this interface to get type-safe publish/subscribe across your app.
@@ -165,6 +266,12 @@ interface ConsumerOptions<T extends TopicMapConstraint<T> = TTopicMessageMap> {
      * On exhaustion, messages go to `<topic>.dlq` (if `dlq: true`) or `onMessageLost`.
      */
     retryTopics?: boolean;
+    /**
+     * Timeout (ms) for waiting until each retry level consumer receives partition
+     * assignments after `startConsumer` connects. Default: `10000`.
+     * Increase this when the broker is slow to rebalance.
+     */
+    retryTopicAssignmentTimeoutMs?: number;
     /**
      * Log a warning if the message handler has not resolved within this window (ms).
      * The handler is not cancelled — this is a diagnostic aid to surface stuck handlers
@@ -196,6 +303,21 @@ interface ConsumerInterceptor<T extends TopicMapConstraint<T> = TTopicMessageMap
     /** Called when the message handler throws. */
     onError?(envelope: EventEnvelope<T[keyof T]>, error: Error): Promise<void> | void;
 }
+/**
+ * Return value of `KafkaInstrumentation.beforeConsume`.
+ *
+ * - `() => void` — legacy form: a cleanup function called after the handler.
+ * - Object form:
+ *   - `cleanup?()` — called after the handler (same as the legacy function form).
+ *   - `wrap?(fn)` — wraps the handler execution; call `fn()` inside the desired
+ *     async context (e.g. `context.with(spanCtx, fn)` for OpenTelemetry). Multiple
+ *     wraps from different instrumentations are composed in declaration order,
+ *     so the first instrumentation's wrap is the outermost.
+ */
+type BeforeConsumeResult = (() => void) | {
+    cleanup?(): void;
+    wrap?(fn: () => Promise<void>): Promise<void>;
+};
 /**
  * Client-wide instrumentation hooks for both send and consume paths.
  * Use this for cross-cutting concerns like tracing and metrics.
@@ -207,8 +329,13 @@ interface KafkaInstrumentation {
     beforeSend?(topic: string, headers: MessageHeaders): void;
     /** Called after a successful send. */
     afterSend?(topic: string): void;
-    /** Called before the consumer handler. Return a cleanup function called after the handler. */
-    beforeConsume?(envelope: EventEnvelope<any>): (() => void) | void;
+    /**
+     * Called before the consumer handler.
+     * Return a cleanup function (legacy) or a `BeforeConsumeResult` object with
+     * optional `cleanup` and `wrap`. Use `wrap` to run the handler inside a
+     * specific async context (e.g. an active OpenTelemetry span).
+     */
+    beforeConsume?(envelope: EventEnvelope<any>): BeforeConsumeResult | void;
     /** Called when the consumer handler throws. */
     onConsumeError?(envelope: EventEnvelope<any>, error: Error): void;
 }
@@ -226,13 +353,19 @@ interface ConsumerHandle {
     /** Stop this consumer. Equivalent to calling `client.stopConsumer(groupId)`. */
     stop(): Promise<void>;
 }
+/** Result returned by `KafkaClient.checkStatus()`. */
+type KafkaHealthResult = {
+    status: "up";
+    clientId: string;
+    topics: string[];
+} | {
+    status: "down";
+    clientId: string;
+    error: string;
+};
 /** Interface describing all public methods of the Kafka client. */
 interface IKafkaClient<T extends TopicMapConstraint<T>> {
-    checkStatus(): Promise<{
-        status: "up";
-        clientId: string;
-        topics: string[];
-    }>;
+    checkStatus(): Promise<KafkaHealthResult>;
     /**
      * Query the consumer group lag per partition using the admin API.
      * Lag = (broker high-watermark offset) − (last committed offset).
@@ -257,8 +390,18 @@ interface IKafkaClient<T extends TopicMapConstraint<T>> {
     sendMessage<K extends keyof T>(topic: K, message: T[K], options?: SendOptions): Promise<void>;
     sendBatch<K extends keyof T>(topic: K, messages: Array<BatchMessageItem<T[K]>>): Promise<void>;
     transaction(fn: (ctx: TransactionContext<T>) => Promise<void>): Promise<void>;
-    getClientId: () => ClientId;
-    disconnect(): Promise<void>;
+    getClientId(): ClientId;
+    /**
+     * Drain in-flight handlers, then disconnect all producers, consumers, and admin.
+     * @param drainTimeoutMs Max ms to wait for in-flight handlers (default 30 000).
+     */
+    disconnect(drainTimeoutMs?: number): Promise<void>;
+    /**
+     * Register SIGTERM / SIGINT signal handlers that drain in-flight messages before
+     * disconnecting. Call once after constructing the client in non-NestJS apps.
+     * NestJS apps get drain automatically via `onModuleDestroy` → `disconnect()`.
+     */
+    enableGracefulShutdown(signals?: NodeJS.Signals[], drainTimeoutMs?: number): void;
 }
 /**
  * Logger interface for KafkaClient.
@@ -325,75 +468,4 @@ interface SubscribeRetryOptions {
     backoffMs?: number;
 }
 
-declare const HEADER_EVENT_ID = "x-event-id";
-declare const HEADER_CORRELATION_ID = "x-correlation-id";
-declare const HEADER_TIMESTAMP = "x-timestamp";
-declare const HEADER_SCHEMA_VERSION = "x-schema-version";
-declare const HEADER_TRACEPARENT = "traceparent";
-/**
- * Typed wrapper combining a parsed message payload with Kafka metadata
- * and envelope headers.
- *
- * On **send**, the library auto-generates envelope headers
- * (`x-event-id`, `x-correlation-id`, `x-timestamp`, `x-schema-version`).
- *
- * On **consume**, the library extracts those headers and assembles
- * an `EventEnvelope` that is passed to the handler.
- */
-interface EventEnvelope<T> {
-    /** Deserialized + validated message body. */
-    payload: T;
-    /** Topic the message was produced to / consumed from. */
-    topic: string;
-    /** Kafka partition (consume-side only, `-1` on send). */
-    partition: number;
-    /** Kafka offset (consume-side only, empty string on send). */
-    offset: string;
-    /** ISO-8601 timestamp set by the producer. */
-    timestamp: string;
-    /** Unique ID for this event (UUID v4). */
-    eventId: string;
-    /** Correlation ID — auto-propagated via AsyncLocalStorage. */
-    correlationId: string;
-    /** Schema version of the payload. */
-    schemaVersion: number;
-    /** W3C Trace Context `traceparent` header (set by OTel instrumentation). */
-    traceparent?: string;
-    /** All decoded Kafka headers for extensibility. */
-    headers: MessageHeaders;
-}
-interface EnvelopeCtx {
-    correlationId: string;
-    traceparent?: string;
-}
-/** Read the current envelope context (correlationId / traceparent) from ALS. */
-declare function getEnvelopeContext(): EnvelopeCtx | undefined;
-/** Execute `fn` inside an envelope context so nested sends inherit correlationId. */
-declare function runWithEnvelopeContext<R>(ctx: EnvelopeCtx, fn: () => R): R;
-/** Options accepted by `buildEnvelopeHeaders`. */
-interface EnvelopeHeaderOptions {
-    correlationId?: string;
-    schemaVersion?: number;
-    eventId?: string;
-    headers?: MessageHeaders;
-}
-/**
- * Generate envelope headers for the send path.
- *
- * Priority for `correlationId`:
- *   explicit option → ALS context → new UUID.
- */
-declare function buildEnvelopeHeaders(options?: EnvelopeHeaderOptions): MessageHeaders;
-/**
- * Decode kafkajs headers (`Record<string, Buffer | string | undefined>`)
- * into plain `Record<string, string>`.
- */
-declare function decodeHeaders(raw: Record<string, Buffer | string | (Buffer | string)[] | undefined> | undefined): MessageHeaders;
-/**
- * Build an `EventEnvelope` from a consumed kafkajs message.
- * Tolerates missing envelope headers — generates defaults so messages
- * from non-envelope producers still work.
- */
-declare function extractEnvelope<T>(payload: T, headers: MessageHeaders, topic: string, partition: number, offset: string): EventEnvelope<T>;
-
-export { type BatchMessageItem as B, type ClientId as C, type EnvelopeHeaderOptions as E, type GroupId as G, HEADER_CORRELATION_ID as H, type IKafkaClient as I, type KafkaInstrumentation as K, type MessageHeaders as M, type RetryOptions as R, type SchemaLike as S, type TopicMapConstraint as T, type KafkaClientOptions as a, type ConsumerOptions as b, type TopicDescriptor as c, type BatchMeta as d, type ConsumerHandle as e, type ConsumerInterceptor as f, type EventEnvelope as g, HEADER_EVENT_ID as h, HEADER_SCHEMA_VERSION as i, HEADER_TIMESTAMP as j, HEADER_TRACEPARENT as k, type InferSchema as l, type KafkaLogger as m, type MessageLostContext as n, type SendOptions as o, type SubscribeRetryOptions as p, type TTopicMessageMap as q, type TopicsFrom as r, type TransactionContext as s, buildEnvelopeHeaders as t, decodeHeaders as u, extractEnvelope as v, getEnvelopeContext as w, runWithEnvelopeContext as x, topic as y };
+export { runWithEnvelopeContext as A, type BatchMessageItem as B, type ClientId as C, topic as D, type EnvelopeHeaderOptions as E, type GroupId as G, HEADER_CORRELATION_ID as H, type IKafkaClient as I, type KafkaInstrumentation as K, type MessageHeaders as M, type RetryOptions as R, type SchemaLike as S, type TopicMapConstraint as T, type KafkaClientOptions as a, type ConsumerOptions as b, type TopicDescriptor as c, type KafkaHealthResult as d, type BatchMeta as e, type BeforeConsumeResult as f, type ConsumerHandle as g, type ConsumerInterceptor as h, type EventEnvelope as i, HEADER_EVENT_ID as j, HEADER_SCHEMA_VERSION as k, HEADER_TIMESTAMP as l, HEADER_TRACEPARENT as m, type InferSchema as n, type KafkaLogger as o, type MessageLostContext as p, type SchemaParseContext as q, type SendOptions as r, type SubscribeRetryOptions as s, type TTopicMessageMap as t, type TopicsFrom as u, type TransactionContext as v, buildEnvelopeHeaders as w, decodeHeaders as x, extractEnvelope as y, getEnvelopeContext as z };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@drarzter/kafka-client",
-  "version": "0.5.6",
+  "version": "0.6.3",
   "description": "Type-safe Kafka client wrapper for NestJS with typed topic-message maps",
   "license": "MIT",
   "main": "./dist/index.js",