@kronos-ts/axon-server 0.5.0 → 0.6.0

package/dist/axon-server.js

@@ -1,3 +1,5 @@
+import { streamRecovery } from "./stream-recovery.js";
+import { messagingAdmission, messagingDeadline, positiveInteger } from "@kronos-ts/core";
 /**
  * The Axon Server command and query buses.
  *
@@ -15,7 +17,7 @@
  *   axonServerQueryBus(localQueryBus(unitOfWork), axon), correlation)
  * ```
  *
- * Axon-specific protocol invariants are preserved byte-for-byte:
+ * Axon-specific protocol invariants:
  *
  *   - CLIENT_SUPPORTS_STREAMING capability advertised on every dispatched
  *     query via `defaultQueryInstructions(...)`;
@@ -26,7 +28,7 @@
  *     `reestablishStreamBody`).
  */
 import { qualifiedNameToString, qualifiedNameFromString, generateIdentifier, } from "@kronos-ts/core";
-import { withRetry } from "./resilience.js";
+import {} from "./resilience.js";
 import { applySubscriptionFilter, updateHandler, runAfterCommitOrImmediately, } from "@kronos-ts/core";
 import { contextView } from "./context-view.js";
 import { metadataToProto, metadataFromProto } from "./metadata-conversion.js";
@@ -34,9 +36,8 @@ import { outboundStream } from "./outbound-stream.js";
 import { mapErrorCode, AxonServerErrorCode } from "./errors.js";
 /** Default flow control settings — aligned with Java's 5000 permits. */
 const DEFAULT_PERMITS = 5000n;
-const DEFAULT_THRESHOLD = 2500n;
 /** Default query dispatch timeout — aligned with Java's one hour. */
-const DEFAULT_QUERY_TIMEOUT_MS = 3_600_000;
+const DEFAULT_QUERY_TIMEOUT_MS = 30_000;
 /** Default command handler load factor — aligned with Java's 100. */
 const DEFAULT_LOAD_FACTOR = 100;
 // Processing instruction keys — aligned with proto ProcessingKey enum.
@@ -146,11 +147,18 @@ function createPayloadHelpers(serializer) {
 export function axonServerCommandBus(next, conn, options = {}) {
     const { connection, serializer, metadata: axonMetadata, } = contextView(conn, options.context ?? conn.connection.config.context);
     const shutdownLatch = conn.shutdown;
+    const requestTimeoutMs = positiveInteger(options.timeoutMs ?? 30000, "timeoutMs");
+    if (requestTimeoutMs > 2_147_483_647)
+        throw new RangeError("timeoutMs exceeds the timer range");
+    const inboundAdmission = messagingAdmission("inbound handlers", options.limits?.maxConcurrentHandlers ?? 128, options.limits?.observe);
+    const outboundAdmission = messagingAdmission("pending requests", options.limits?.maxPendingRequests ?? 1024, options.limits?.observe);
     const resilience = options.resilience ?? conn.resilience;
     const metadata = axonMetadata();
     const { serializePayload, deserializePayload } = createPayloadHelpers(serializer);
-    const PERMITS = BigInt(options.flowControl?.permits ?? Number(DEFAULT_PERMITS));
-    const THRESHOLD = BigInt(options.flowControl?.refillThreshold ?? Number(DEFAULT_THRESHOLD));
+    const PERMITS = BigInt(positiveInteger(options.flowControl?.permits ?? Number(DEFAULT_PERMITS), "flowControl.permits"));
+    const THRESHOLD = BigInt(options.flowControl?.refillThreshold ?? Math.floor(Number(PERMITS) / 2));
+    if (THRESHOLD < 0n || THRESHOLD >= PERMITS)
+        throw new RangeError("refillThreshold must be between zero and permits - 1");
     const loadFactor = options.loadFactor ?? DEFAULT_LOAD_FACTOR;
     /**
      * The names this node announced to Axon Server. The handlers themselves live
@@ -162,14 +170,16 @@ export function axonServerCommandBus(next, conn, options = {}) {
     // Bidirectional stream for handler subscription + inbound command handling
     let outbound = outboundStream();
     let streamStarted = false;
+    let providerAbort = new AbortController();
+    connection.onDisconnect?.(() => { providerAbort.abort(); outbound.close(); });
     let permits = 0n;
     function ensureStreamStarted() {
         if (streamStarted)
             return;
         streamStarted = true;
         // Open stream using connection.commands (always gets current client after reconnect)
-        const inbound = connection.commands.openStream(outbound.iterable, { metadata });
-        processInboundCommands(inbound);
+        const inbound = connection.commands.openStream(outbound.iterable, { metadata, signal: providerAbort.signal });
+        void processInboundCommands(inbound, outbound);
     }
     function grantPermits() {
         outbound.send({
@@ -199,6 +209,8 @@ export function axonServerCommandBus(next, conn, options = {}) {
      * trigger a server-side stream error.
      */
     function reestablishStreamBody() {
+        providerAbort.abort();
+        providerAbort = new AbortController();
         outbound.close();
         outbound = outboundStream();
         streamStarted = false;
@@ -210,98 +222,106 @@ export function axonServerCommandBus(next, conn, options = {}) {
         // Permits AFTER subscriptions (Axon-specific ordering invariant)
         grantPermits();
     }
-    async function reestablishStreamWithRetry() {
-        if (shutdownLatch.shuttingDown)
-            return;
-        await withRetry(async () => reestablishStreamBody(), {
-            event: "reconnect",
-            ...resilience,
-        });
-    }
+    const recovery = streamRecovery(reestablishStreamBody, () => !shutdownLatch.shuttingDown && connection.state !== "closed" && connection.state !== "disconnected" && connection.state !== "reconnecting", resilience);
+    shutdownLatch.onShutdown(recovery.stop);
     // Auto-reestablish when the connection reconnects (e.g., after heartbeat timeout)
     connection.onReconnect(() => {
         if (!shutdownLatch.shuttingDown && streamStarted) {
-            reestablishStreamWithRetry().catch((err) => {
-                console.error("Axon Server command bus: reconnect retries exhausted", err);
-            });
+            recovery.restart();
         }
     });
-    async function processInboundCommands(inbound) {
+    async function handleInboundCommand(proto, responses) {
+        let activity;
+        let admission;
+        let responseSerialized;
+        let errorCode = "";
+        let errorMsg = "";
         try {
-            for await (const message of inbound) {
-                if (!message.command)
-                    continue;
-                permits--;
-                const proto = message.command;
-                const commandName = proto.name;
-                let resultPayload;
-                let errorCode = "";
-                let errorMsg = "";
-                if (subscribedNames.has(commandName)) {
-                    try {
-                        const commandMessage = {
-                            kind: "command",
-                            identifier: proto.messageIdentifier,
-                            name: qualifiedNameFromString(commandName),
-                            payload: deserializePayload(proto.payload?.data),
-                            metadata: metadataFromProto(proto.metaData),
-                            timestamp: Number(proto.timestamp),
-                        };
-                        // Through the LOCAL BUS, so the caller's unit-of-work policy runs.
-                        // AF parity is preserved: `CommandProcessingTask` runs the next
-                        // segment without re-running dispatch interceptors, and a `next`
-                        // that happens to carry `correlation` re-applies a pair of `??` seeds
-                        // that are already set.
-                        resultPayload = await next.dispatch(commandMessage);
-                    }
-                    catch (err) {
-                        errorCode = AxonServerErrorCode.COMMAND_EXECUTION_ERROR;
-                        errorMsg = err instanceof Error ? err.message : String(err);
-                    }
+            try {
+                // Remote callers have no outbound dispatch activity on this adapter.
+                // Track the entire handling, including result serialization and enqueue.
+                // Registration also rejects new work once shutdown has begun.
+                activity = shutdownLatch.registerActivity();
+                admission = inboundAdmission.enter();
+                if (subscribedNames.has(proto.name)) {
+                    const commandMessage = {
+                        kind: "command",
+                        identifier: proto.messageIdentifier,
+                        name: qualifiedNameFromString(proto.name),
+                        payload: deserializePayload(proto.payload?.data, proto.payload?.type, proto.payload?.revision),
+                        metadata: metadataFromProto(proto.metaData ?? {}),
+                        timestamp: Number(proto.timestamp),
+                    };
+                    // The local bus opens a fresh unit of work for EVERY wire command,
+                    // including children of handlers running on this same connection.
+                    const result = await next.dispatch(commandMessage);
+                    responseSerialized = result !== undefined ? serializePayload("result", result) : undefined;
                 }
                 else {
                     errorCode = AxonServerErrorCode.NO_HANDLER_FOR_COMMAND;
-                    errorMsg = `No next handler for command "${commandName}"`;
+                    errorMsg = `No next handler for command "${proto.name}"`;
                 }
-                // Send response back to Axon Server
-                outbound.send({
-                    commandResponse: {
-                        messageIdentifier: generateIdentifier(),
-                        requestIdentifier: proto.messageIdentifier,
-                        errorCode,
-                        errorMessage: errorCode
-                            ? {
-                                message: errorMsg,
-                                location: connection.config.componentName,
-                                details: [],
-                                errorCode,
-                            }
-                            : undefined,
-                        payload: resultPayload !== undefined ? serializePayload("result", resultPayload) : undefined,
-                        metaData: {},
-                        processingInstructions: [],
-                    },
-                    instructionId: "",
+            }
+            catch (err) {
+                // Decode, handler, and result-encoding failures belong to this request;
+                // none should terminate the receive loop or reconnect the stream.
+                errorCode = AxonServerErrorCode.COMMAND_EXECUTION_ERROR;
+                errorMsg = err instanceof Error ? err.message : String(err);
+            }
+            // Capture the originating stream: a late handler must not send an old
+            // request's response on a replacement stream after reconnect.
+            responses.send({
+                commandResponse: {
+                    messageIdentifier: generateIdentifier(),
+                    requestIdentifier: proto.messageIdentifier,
+                    errorCode,
+                    errorMessage: errorCode
+                        ? { message: errorMsg, location: connection.config.componentName, details: [], errorCode }
+                        : undefined,
+                    payload: responseSerialized,
+                    metaData: {},
+                    processingInstructions: [],
+                },
+                instructionId: "",
+            });
+            await responses.flush();
+        }
+        finally {
+            admission?.end();
+            activity?.end();
+        }
+    }
+    async function processInboundCommands(inbound, responses) {
+        try {
+            for await (const message of inbound) {
+                if (responses !== outbound)
+                    return;
+                recovery.received();
+                if (message.instructionId)
+                    responses.send({ ack: { instructionId: message.instructionId, success: true }, instructionId: "" });
+                permits--;
+                // Credits bound delivery batches, not unfinished handlers. Replenish
+                // on receipt: completion-based credits or a fixed handler semaphore can
+                // deadlock when every admitted parent is waiting for a queued child.
+                if (permits <= THRESHOLD && !shutdownLatch.shuttingDown)
+                    grantPermits();
+                if (!message.command)
+                    continue;
+                // Each invocation owns its response and shutdown activity. Keep reading
+                // while it awaits work so nested dispatch can return on this connection.
+                void handleInboundCommand(message.command, responses).catch((err) => {
+                    console.error("Axon Server command bus: inbound response failed", err);
                 });
-                // Refill permits when running low
-                if (permits <= THRESHOLD) {
-                    outbound.send({
-                        flowControl: { clientId: connection.config.clientId, permits: PERMITS },
-                        instructionId: "",
-                    });
-                    permits += PERMITS;
-                }
             }
+            if (responses === outbound && !shutdownLatch.shuttingDown)
+                throw new Error("Inbound provider stream ended unexpectedly");
         }
         catch (err) {
-            if (shutdownLatch.shuttingDown)
+            if (responses !== outbound || shutdownLatch.shuttingDown)
                 return;
-            if (String(err).includes("Connection dropped"))
+            if (connection.state === "reconnecting" || connection.state === "closed" || connection.state === "disconnected")
                 return;
-            console.error("Axon Server command bus: inbound stream error, attempting re-establishment via withRetry", err);
-            await reestablishStreamWithRetry().catch((retryErr) => {
-                console.error("Axon Server command bus: reconnect retries exhausted", retryErr);
-            });
+            recovery.failed(err);
         }
     }
     return {
@@ -313,7 +333,11 @@ export function axonServerCommandBus(next, conn, options = {}) {
             // instead, so the task that handles it supplies the instant.
             const message = { ...unstamped, timestamp: unstamped.timestamp ?? Date.now() };
             const activity = shutdownLatch.registerActivity();
+            let admission;
+            let deadline;
             try {
+                admission = outboundAdmission.enter(unstamped.identifier);
+                deadline = messagingDeadline(requestTimeoutMs);
                 const commandName = qualifiedNameToString(message.name);
                 const response = await connection.commands.dispatch({
                     messageIdentifier: message.identifier,
@@ -324,13 +348,15 @@ export function axonServerCommandBus(next, conn, options = {}) {
                     processingInstructions: toProtoProcessingInstructions(message.metadata?.processingInstructions),
                     clientId: connection.config.clientId,
                     componentName: connection.config.componentName,
-                }, { metadata });
+                }, { metadata, signal: deadline.signal });
                 if (response.errorCode && response.errorCode !== "") {
                     throw mapErrorCode(response.errorCode, response.errorMessage?.message ?? "Unknown error");
                 }
-                return deserializePayload(response.payload?.data);
+                return deserializePayload(response.payload?.data, response.payload?.type, response.payload?.revision);
             }
             finally {
+                deadline?.close();
+                admission?.end();
                 activity.end();
             }
         },
@@ -368,18 +394,26 @@ export function axonServerCommandBus(next, conn, options = {}) {
  * `scatterGather` and `subscriptionQuery`. Because the wrap is outside, the
  * shortcut branch gets identical correlation to the remote branch.
  *
- * KNOWN GAP: `subscriptionQuery` / `subscribeToUpdates` build their proto
- * straight from `message.metadata`, and `interceptingQueryBus` (in
- * `@kronos-ts/core`) forwards those two calls to the delegate without
- * running the dispatch chain. Closing that needs a core change.
+ * Subscription queries run the dispatch chain: `interceptingQueryBus` wraps
+ * `subscriptionQuery` / `subscribeToUpdates` with the same intercept the
+ * primary `query` gets, so the proto built from `message.metadata` already
+ * carries whatever the host's intercept stamped (pinned in core by
+ * `interception/__tests__/subscription-interception.test.ts`).
  */
 export function axonServerQueryBus(next, conn, options = {}) {
     const { connection, serializer, metadata: axonMetadata, } = contextView(conn, options.context ?? conn.connection.config.context);
     const shutdownLatch = conn.shutdown;
+    const requestTimeoutMs = positiveInteger(options.timeoutMs ?? 30000, "timeoutMs");
+    if (requestTimeoutMs > 2_147_483_647)
+        throw new RangeError("timeoutMs exceeds the timer range");
+    const inboundAdmission = messagingAdmission("inbound handlers", options.limits?.maxConcurrentHandlers ?? 128, options.limits?.observe);
+    const outboundAdmission = messagingAdmission("pending requests", options.limits?.maxPendingRequests ?? 1024, options.limits?.observe);
     const resilience = options.resilience ?? conn.resilience;
     const metadata = axonMetadata();
-    const PERMITS = BigInt(options.flowControl?.permits ?? Number(DEFAULT_PERMITS));
-    const THRESHOLD = BigInt(options.flowControl?.refillThreshold ?? Number(DEFAULT_THRESHOLD));
+    const PERMITS = BigInt(positiveInteger(options.flowControl?.permits ?? Number(DEFAULT_PERMITS), "flowControl.permits"));
+    const THRESHOLD = BigInt(options.flowControl?.refillThreshold ?? Math.floor(Number(PERMITS) / 2));
+    if (THRESHOLD < 0n || THRESHOLD >= PERMITS)
+        throw new RangeError("refillThreshold must be between zero and permits - 1");
     const shortcutQueriesToLocalHandlers = options.shortcutQueriesToLocalHandlers ?? false;
     const queryTimeoutMs = options.timeoutMs ?? DEFAULT_QUERY_TIMEOUT_MS;
     const { serializePayload, deserializePayload } = createPayloadHelpers(serializer);
@@ -394,15 +428,47 @@ export function axonServerQueryBus(next, conn, options = {}) {
     // to decide which subscriber IDs to target; the server forwards each response to the
     // exact subscriber.
     const handlerSubscriptions = new Map();
+    const responseCredits = new Map();
+    function cancelResponseCredits() {
+        for (const credit of responseCredits.values()) {
+            clearTimeout(credit.timer);
+            credit.cancel();
+        }
+        responseCredits.clear();
+    }
+    function responseCredit(identifier) {
+        const existing = responseCredits.get(identifier);
+        if (existing)
+            return existing;
+        if (responseCredits.size >= (options.limits?.maxPendingRequests ?? 1024))
+            return undefined;
+        let grant;
+        const ready = new Promise((resolve) => { grant = resolve; });
+        const credit = {
+            ready, grant, cancelled: false,
+            cancel() { this.cancelled = true; grant(); },
+            timer: setTimeout(() => {
+                credit.cancel();
+                if (responseCredits.get(identifier) === credit)
+                    responseCredits.delete(identifier);
+            }, requestTimeoutMs),
+        };
+        credit.timer.unref?.();
+        responseCredits.set(identifier, credit);
+        return credit;
+    }
+    shutdownLatch.onShutdown(() => handlerSubscriptions.clear());
     let outbound = outboundStream();
     let streamStarted = false;
+    let providerAbort = new AbortController();
+    connection.onDisconnect?.(() => { cancelResponseCredits(); providerAbort.abort(); outbound.close(); });
     let permits = 0n;
     function ensureStreamStarted() {
         if (streamStarted)
             return;
         streamStarted = true;
-        const inbound = connection.queries.openStream(outbound.iterable, { metadata });
-        processInboundQueries(inbound);
+        const inbound = connection.queries.openStream(outbound.iterable, { metadata, signal: providerAbort.signal });
+        void processInboundQueries(inbound, outbound);
     }
     function grantQueryPermits() {
         outbound.send({
@@ -431,6 +497,8 @@ export function axonServerQueryBus(next, conn, options = {}) {
      * re-emitted BEFORE the permits frame.
      */
     function reestablishStreamBody() {
+        cancelResponseCredits();
+        handlerSubscriptions.clear();
         outbound.close();
         outbound = outboundStream();
         streamStarted = false;
@@ -440,175 +508,202 @@ export function axonServerQueryBus(next, conn, options = {}) {
             sendSubscribe(queryName);
         grantQueryPermits();
     }
-    async function reestablishStreamWithRetry() {
-        if (shutdownLatch.shuttingDown)
-            return;
-        await withRetry(async () => reestablishStreamBody(), {
-            event: "reconnect",
-            ...resilience,
-        });
-    }
+    const recovery = streamRecovery(reestablishStreamBody, () => !shutdownLatch.shuttingDown && connection.state !== "closed" && connection.state !== "disconnected" && connection.state !== "reconnecting", resilience);
+    shutdownLatch.onShutdown(recovery.stop);
     // Auto-reestablish when the connection reconnects (e.g., after heartbeat timeout)
     connection.onReconnect(() => {
         if (!shutdownLatch.shuttingDown && streamStarted) {
-            reestablishStreamWithRetry().catch((err) => {
-                console.error("Axon Server query bus: reconnect retries exhausted", err);
-            });
+            recovery.restart();
         }
     });
-    async function handleSubscriptionQueryRequest(req) {
-        if (req.subscribe) {
-            const sub = req.subscribe;
-            const subId = sub.subscriptionIdentifier;
-            const proto = sub.queryRequest;
-            if (!subId || !proto)
+    async function handleInboundQuery(proto, responses, subId) {
+        let activity;
+        let admission;
+        let payload;
+        let subscriptionEntry;
+        let credit;
+        const supports = (key) => proto.processingInstructions?.some((instruction) => instruction.key === key && instruction.value?.booleanValue);
+        if (!subId && supports(7) && supports(8)) {
+            // Response credits may precede the query on Axon's provider stream.
+            // Retain those credits by request ID, within the same bounded table.
+            credit = responseCredit(proto.messageIdentifier);
+            if (!credit) {
+                responses.close();
                 return;
-            const queryName = proto.query;
-            const payload = deserializePayload(proto.payload?.data, proto.payload?.type, proto.payload?.revision);
-            handlerSubscriptions.set(subId, { queryName, payload });
-            let resultPayload;
-            let errorCode = "";
-            let errorMsg = "";
-            if (subscribedNames.has(queryName)) {
-                try {
+            }
+        }
+        let errorCode = "";
+        let errorMsg = "";
+        try {
+            try {
+                activity = shutdownLatch.registerActivity();
+                admission = inboundAdmission.enter();
+                if (subscribedNames.has(proto.query)) {
                     const queryMessage = {
                         kind: "query",
                         identifier: proto.messageIdentifier,
-                        name: qualifiedNameFromString(queryName),
-                        payload,
+                        name: qualifiedNameFromString(proto.query),
+                        payload: deserializePayload(proto.payload?.data, proto.payload?.type, proto.payload?.revision),
                         metadata: metadataFromProto(proto.metaData ?? {}),
                         timestamp: Number(proto.timestamp),
                     };
-                    resultPayload = await next.query(queryMessage);
+                    if (subId) {
+                        subscriptionEntry = handlerSubscriptions.get(subId);
+                        if (!subscriptionEntry)
+                            return;
+                    }
+                    // Every wire query enters the local bus with a fresh unit of work.
+                    const result = await next.query(queryMessage);
+                    payload = result !== undefined ? serializePayload("result", result) : undefined;
                 }
-                catch (err) {
-                    errorCode = AxonServerErrorCode.QUERY_EXECUTION_ERROR;
-                    errorMsg = err instanceof Error ? err.message : String(err);
+                else {
+                    errorCode = AxonServerErrorCode.NO_HANDLER_FOR_QUERY;
+                    errorMsg = `No next handler for query "${proto.query}"`;
                 }
             }
-            else {
-                errorCode = AxonServerErrorCode.NO_HANDLER_FOR_QUERY;
-                errorMsg = `No next handler for query "${queryName}"`;
+            catch (err) {
+                errorCode = AxonServerErrorCode.QUERY_EXECUTION_ERROR;
+                errorMsg = err instanceof Error ? err.message : String(err);
             }
-            const responseSerialized = resultPayload !== undefined ? serializePayload("result", resultPayload) : undefined;
-            outbound.send({
-                subscriptionQueryResponse: {
-                    messageIdentifier: generateIdentifier(),
-                    subscriptionIdentifier: subId,
-                    initialResult: {
-                        messageIdentifier: generateIdentifier(),
-                        requestIdentifier: proto.messageIdentifier,
-                        errorCode,
-                        errorMessage: errorCode
-                            ? {
-                                message: errorMsg,
-                                location: connection.config.componentName,
-                                details: [],
-                                errorCode,
-                            }
-                            : undefined,
-                        payload: responseSerialized,
-                        metaData: {},
-                        processingInstructions: [],
+            // An unsubscribe or completion can overtake a slow initial handler.
+            if (subId && subscriptionEntry && handlerSubscriptions.get(subId) !== subscriptionEntry)
+                return;
+            if (subId && errorCode)
+                handlerSubscriptions.delete(subId);
+            if (credit) {
+                await credit.ready;
+                if (credit.cancelled)
+                    return;
+            }
+            const response = {
+                messageIdentifier: generateIdentifier(),
+                requestIdentifier: proto.messageIdentifier,
+                errorCode,
+                errorMessage: errorCode
+                    ? { message: errorMsg, location: connection.config.componentName, details: [], errorCode }
+                    : undefined,
+                payload,
+                metaData: {},
+                processingInstructions: [],
+            };
+            if (subId) {
+                responses.send({
+                    subscriptionQueryResponse: {
+                        messageIdentifier: generateIdentifier(), subscriptionIdentifier: subId, initialResult: response,
                     },
-                },
-                instructionId: "",
-            });
-            return;
+                    instructionId: "",
+                });
+            }
+            else {
+                responses.send({ queryResponse: response, instructionId: "" });
+                responses.send({
+                    queryComplete: { messageId: generateIdentifier(), requestId: proto.messageIdentifier },
+                    instructionId: "",
+                });
+            }
+            await responses.flush();
         }
-        if (req.unsubscribe) {
-            handlerSubscriptions.delete(req.unsubscribe.subscriptionIdentifier);
+        finally {
+            if (credit)
+                clearTimeout(credit.timer);
+            if (credit && responseCredits.get(proto.messageIdentifier) === credit)
+                responseCredits.delete(proto.messageIdentifier);
+            admission?.end();
+            activity?.end();
         }
-        // flowControl + getInitialResult are not tracked per-sub; ignored for now.
     }
-    async function processInboundQueries(inbound) {
+    async function processInboundQueries(inbound, responses) {
         try {
             for await (const message of inbound) {
-                if (message.subscriptionQueryRequest) {
-                    await handleSubscriptionQueryRequest(message.subscriptionQueryRequest);
-                    continue;
-                }
-                if (!message.query)
-                    continue;
+                if (responses !== outbound)
+                    return;
+                recovery.received();
+                if (message.instructionId)
+                    responses.send({ ack: { instructionId: message.instructionId, success: true }, instructionId: "" });
+                // Every Axon query instruction consumes a provider credit, including
+                // acknowledgements. Ignoring a late subscription ack can exhaust a
+                // one-credit window between otherwise successful requests.
                 permits--;
-                const proto = message.query;
-                const queryName = proto.query;
-                let resultPayload;
-                let errorCode = "";
-                let errorMsg = "";
-                if (subscribedNames.has(queryName)) {
-                    try {
-                        const queryMessage = {
-                            kind: "query",
-                            identifier: proto.messageIdentifier,
-                            name: qualifiedNameFromString(queryName),
-                            payload: deserializePayload(proto.payload?.data),
-                            metadata: metadataFromProto(proto.metaData),
-                            timestamp: Number(proto.timestamp),
-                        };
-                        // Through the LOCAL BUS: no unit of work is handed in, so `next`
-                        // opens one under whatever policy the caller gave it.
-                        resultPayload = await next.query(queryMessage);
-                    }
-                    catch (err) {
-                        errorCode = AxonServerErrorCode.QUERY_EXECUTION_ERROR;
-                        errorMsg = err instanceof Error ? err.message : String(err);
+                if (permits <= THRESHOLD && !shutdownLatch.shuttingDown)
+                    grantQueryPermits();
+                if (message.queryFlowControl?.permits > 0n) {
+                    const identifier = message.queryFlowControl.queryReference?.requestId;
+                    if (identifier) {
+                        const credit = responseCredit(identifier);
+                        if (!credit) {
+                            responses.close();
+                            return;
+                        }
+                        credit.grant();
                     }
                 }
-                else {
-                    errorCode = AxonServerErrorCode.NO_HANDLER_FOR_QUERY;
-                    errorMsg = `No next handler for query "${queryName}"`;
+                if (message.queryCancel)
+                    responseCredits.get(message.queryCancel.requestId)?.cancel();
+                const request = message.subscriptionQueryRequest;
+                if (request) {
+                    if (request.unsubscribe)
+                        handlerSubscriptions.delete(request.unsubscribe.subscriptionIdentifier);
+                    const sub = request.subscribe;
+                    if (sub?.subscriptionIdentifier && sub.queryRequest) {
+                        try {
+                            if (shutdownLatch.shuttingDown)
+                                throw new Error("Shutdown in progress");
+                            if (handlerSubscriptions.size >= 1024 && !handlerSubscriptions.has(sub.subscriptionIdentifier))
+                                throw new Error("Provider subscription capacity exhausted");
+                            const proto = sub.queryRequest;
+                            handlerSubscriptions.set(sub.subscriptionIdentifier, {
+                                queryName: proto.query,
+                                payload: deserializePayload(proto.payload?.data, proto.payload?.type, proto.payload?.revision),
+                            });
+                        }
+                        catch (err) {
+                            responses.send({
+                                subscriptionQueryResponse: {
+                                    messageIdentifier: generateIdentifier(), subscriptionIdentifier: sub.subscriptionIdentifier,
+                                    completeExceptionally: {
+                                        errorCode: AxonServerErrorCode.QUERY_EXECUTION_ERROR,
+                                        errorMessage: { message: err instanceof Error ? err.message : String(err) },
+                                    },
+                                },
+                                instructionId: "",
+                            });
+                        }
+                    }
+                    // Axon separates update registration from requesting the initial
+                    // result. Running the handler on Subscribe answers the wrong phase.
+                    const initial = request.getInitialResult;
+                    if (initial?.subscriptionIdentifier && initial.queryRequest) {
+                        void handleInboundQuery(initial.queryRequest, responses, initial.subscriptionIdentifier).catch((err) => {
+                            console.error("Axon Server query bus: inbound subscription response failed", err);
+                        });
+                    }
+                    continue;
                 }
-                outbound.send({
-                    queryResponse: {
-                        messageIdentifier: generateIdentifier(),
-                        requestIdentifier: proto.messageIdentifier,
-                        errorCode,
-                        errorMessage: errorCode
-                            ? {
-                                message: errorMsg,
-                                location: connection.config.componentName,
-                                details: [],
-                                errorCode,
-                            }
-                            : undefined,
-                        payload: resultPayload !== undefined ? serializePayload("result", resultPayload) : undefined,
-                        metaData: {},
-                        processingInstructions: [],
-                    },
-                    instructionId: "",
-                });
-                outbound.send({
-                    queryComplete: {
-                        messageId: generateIdentifier(),
-                        requestId: proto.messageIdentifier,
-                    },
-                    instructionId: "",
+                if (!message.query)
+                    continue;
+                void handleInboundQuery(message.query, responses).catch((err) => {
+                    console.error("Axon Server query bus: inbound response failed", err);
                 });
-                if (permits <= THRESHOLD) {
-                    outbound.send({
-                        flowControl: { clientId: connection.config.clientId, permits: PERMITS },
-                        instructionId: "",
-                    });
-                    permits += PERMITS;
-                }
             }
+            if (responses === outbound && !shutdownLatch.shuttingDown)
+                throw new Error("Inbound provider stream ended unexpectedly");
         }
         catch (err) {
-            if (shutdownLatch.shuttingDown)
+            if (responses !== outbound || shutdownLatch.shuttingDown)
                 return;
-            if (String(err).includes("Connection dropped"))
+            if (connection.state === "reconnecting" || connection.state === "closed" || connection.state === "disconnected")
                 return;
-            console.error("Axon Server query bus: inbound stream error, attempting re-establishment via withRetry", err);
-            await reestablishStreamWithRetry().catch((retryErr) => {
-                console.error("Axon Server query bus: reconnect retries exhausted", retryErr);
-            });
+            recovery.failed(err);
         }
     }
     const routing = {
         async query(unstamped, uow) {
             const activity = shutdownLatch.registerActivity();
+            let admission;
+            let deadline;
             try {
+                admission = outboundAdmission.enter(unstamped.identifier);
+                deadline = messagingDeadline(requestTimeoutMs);
                 const queryName = qualifiedNameToString(unstamped.name);
                 // Local shortcut — handle locally if a handler is co-located. The
                 // caller's unit of work is passed straight through, so `next` makes the
@@ -616,7 +711,7 @@ export function axonServerQueryBus(next, conn, options = {}) {
                 // in-process read: a live unit of work handed in by `ctx.query` is
                 // reused so the consulting read shares the caller's transaction.
                 if (shortcutQueriesToLocalHandlers && subscribedNames.has(queryName)) {
-                    return next.query(unstamped, uow);
+                    return await next.query(unstamped, uow);
                 }
                 // A transport is not a task: it has no unit of work, so it has no clock.
                 // A message that reaches the wire with no instant yet gets one from system
@@ -633,16 +728,38 @@ export function axonServerQueryBus(next, conn, options = {}) {
                     processingInstructions: defaultQueryInstructions(queryTimeoutMs),
                     clientId: connection.config.clientId,
                     componentName: connection.config.componentName,
-                }, { metadata });
+                }, { metadata, signal: deadline.signal });
+                // NR_OF_RESULTS is one. Drain trailers before returning so transport
+                // failures cannot be mistaken for a successful result. The RPC deadline
+                // also bounds a stream that sends a response but never completes.
+                let received = false;
+                let result;
+                let responseError;
                 for await (const response of responseStream) {
+                    if (received)
+                        continue;
+                    received = true;
                     if (response.errorCode && response.errorCode !== "") {
-                        throw mapErrorCode(response.errorCode, response.errorMessage?.message ?? "Unknown error");
+                        responseError = mapErrorCode(response.errorCode, response.errorMessage?.message ?? "Unknown error");
+                    }
+                    else {
+                        try {
+                            result = deserializePayload(response.payload?.data, response.payload?.type, response.payload?.revision);
+                        }
+                        catch (error) {
+                            responseError = error instanceof Error ? error : new Error(String(error));
+                        }
                     }
-                    return deserializePayload(response.payload?.data);
                 }
+                if (responseError)
+                    throw responseError;
+                if (received)
+                    return result;
                 throw new Error(`No response for query "${queryName}"`);
             }
             finally {
+                deadline?.close();
+                admission?.end();
                 activity.end();
             }
         },
@@ -655,25 +772,33 @@ export function axonServerQueryBus(next, conn, options = {}) {
             grantQueryPermits();
         },
         subscriptionQuery(unstamped, bufferSize) {
+            if (shutdownLatch.shuttingDown)
+                throw new Error("Messaging shutdown in progress");
+            if (subscriptions.size >= 1024)
+                throw new Error("Subscription capacity 1024 exhausted");
             const message = { ...unstamped, timestamp: unstamped.timestamp ?? Date.now() };
             const queryId = message.identifier;
             if (subscriptions.has(queryId)) {
                 throw new Error(`Subscription query already registered for identifier "${queryId}"`);
             }
-            const handler = updateHandler(message, bufferSize);
-            subscriptions.set(queryId, handler);
+            const handler = updateHandler(message, bufferSize, () => subscriptions.delete(queryId));
             const queryName = qualifiedNameToString(message.name);
+            const serialized = serializePayload(queryName, message.payload);
             const subscriptionId = generateIdentifier();
             const outboundSub = outboundStream();
+            const window = Math.min(1024, Math.max(256, Math.floor(bufferSize ?? 256)));
+            const refillBatch = Math.max(1, Math.floor(window / 4));
+            let consumedSinceRefill = 0;
+            let subscriptionClosed = false;
             outboundSub.send({
                 subscribe: {
                     subscriptionIdentifier: subscriptionId,
-                    numberOfPermits: BigInt(bufferSize ?? 256),
+                    numberOfPermits: BigInt(window),
                     queryRequest: {
                         messageIdentifier: message.identifier,
                         query: queryName,
                         timestamp: BigInt(message.timestamp),
-                        payload: serializePayload(queryName, message.payload),
+                        payload: serialized,
                         metaData: metadataToProto(message.metadata),
                         processingInstructions: defaultQueryInstructions(queryTimeoutMs),
                         clientId: connection.config.clientId,
@@ -681,6 +806,9 @@ export function axonServerQueryBus(next, conn, options = {}) {
                     },
                 },
             });
+            // Subscribe does not grant update credits on Axon Server; a separate
+            // FlowControl frame initializes the subscription stream's update window.
+            outboundSub.send({ flowControl: { numberOfPermits: BigInt(window) } });
             outboundSub.send({
                 getInitialResult: {
                     subscriptionIdentifier: subscriptionId,
@@ -689,7 +817,7 @@ export function axonServerQueryBus(next, conn, options = {}) {
                         messageIdentifier: message.identifier,
                         query: queryName,
                         timestamp: BigInt(message.timestamp),
-                        payload: serializePayload(queryName, message.payload),
+                        payload: serialized,
                         metaData: metadataToProto(message.metadata),
                         processingInstructions: defaultQueryInstructions(queryTimeoutMs),
                         clientId: connection.config.clientId,
@@ -697,7 +825,9 @@ export function axonServerQueryBus(next, conn, options = {}) {
                     },
                 },
             });
-            const responseStream = connection.queries.subscription(outboundSub.iterable, { metadata });
+            const subscriptionController = new AbortController();
+            const responseStream = connection.queries.subscription(outboundSub.iterable, { metadata, signal: subscriptionController.signal });
+            subscriptions.set(queryId, handler);
             let resolveInitial;
             let rejectInitial;
             const initialResult = new Promise((resolve, reject) => {
@@ -705,70 +835,114 @@ export function axonServerQueryBus(next, conn, options = {}) {
                 rejectInitial = reject;
             });
             let initialSettled = false;
-            (async () => {
+            let explicitlyCompleted = false;
+            const initialTimer = setTimeout(() => closeSubscription(new Error("Subscription initial result timed out")), requestTimeoutMs);
+            const removeShutdown = shutdownLatch.onShutdown(() => closeSubscription(new Error("Messaging shutdown in progress")));
+            // Callers may consume updates without awaiting the initial result. Keep
+            // the original promise rejectable without an unhandled rejection on close.
+            void initialResult.catch(() => { });
+            function closeSubscription(error) {
+                if (subscriptionClosed)
+                    return;
+                subscriptionClosed = true;
+                clearTimeout(initialTimer);
+                removeShutdown();
+                if (!initialSettled) {
+                    rejectInitial(error ?? new Error("Subscription query closed before initial result"));
+                    initialSettled = true;
+                }
+                if (error)
+                    handler.completeExceptionally(error);
+                else
+                    handler.complete();
+                try {
+                    outboundSub.send({ unsubscribe: { subscriptionIdentifier: subscriptionId } });
+                }
+                catch { /* Broken stream; local teardown still must finish. */ }
+                outboundSub.close();
+                subscriptionController.abort();
+                subscriptions.delete(queryId);
+            }
+            void (async () => {
                 try {
                     for await (const response of responseStream) {
+                        if (subscriptionClosed)
+                            break;
                         if (response.initialResult) {
                             const initial = response.initialResult;
                             if (!initialSettled) {
-                                if (initial.errorCode && initial.errorCode !== "") {
-                                    rejectInitial(mapErrorCode(initial.errorCode, initial.errorMessage?.message ?? "Unknown error"));
-                                }
-                                else {
-                                    resolveInitial(deserializePayload(initial.payload?.data));
+                                if (initial.errorCode) {
+                                    throw mapErrorCode(initial.errorCode, initial.errorMessage?.message ?? "Unknown error");
                                 }
+                                clearTimeout(initialTimer);
+                                resolveInitial(deserializePayload(initial.payload?.data, initial.payload?.type, initial.payload?.revision));
                                 initialSettled = true;
                             }
                         }
                         else if (response.update) {
-                            const update = deserializePayload(response.update.payload?.data);
-                            handler.offer(update);
+                            const update = deserializePayload(response.update.payload?.data, response.update.payload?.type, response.update.payload?.revision);
+                            if (!handler.offer(update))
+                                throw new Error("Subscription query update buffer overflow");
+                            consumedSinceRefill++;
+                            if (consumedSinceRefill >= refillBatch) {
+                                outboundSub.send({
+                                    flowControl: { subscriptionIdentifier: subscriptionId, numberOfPermits: BigInt(consumedSinceRefill) },
+                                });
+                                consumedSinceRefill = 0;
+                            }
                         }
                         else if (response.complete) {
-                            handler.complete();
+                            explicitlyCompleted = true;
                             break;
                         }
                         else if (response.completeExceptionally) {
-                            handler.completeExceptionally(new Error(response.completeExceptionally.errorMessage?.message ??
-                                "Subscription query failed"));
-                            break;
+                            throw new Error(response.completeExceptionally.errorMessage?.message ?? "Subscription query failed");
                         }
                     }
                 }
                 catch (err) {
-                    const error = err instanceof Error ? err : new Error(String(err));
+                    closeSubscription(err instanceof Error ? err : new Error(String(err)));
+                }
+                finally {
+                    // EOF and completion frames must settle BOTH faces of a subscription.
+                    const missingInitial = !initialSettled;
                     if (!initialSettled) {
-                        rejectInitial(error);
+                        rejectInitial(new Error("Subscription stream ended before initial result"));
                         initialSettled = true;
                     }
-                    handler.completeExceptionally(error);
-                }
-                finally {
-                    subscriptions.delete(queryId);
+                    closeSubscription(!missingInitial && !explicitlyCompleted && !subscriptionClosed ? new Error("Subscription stream ended unexpectedly") : undefined);
                 }
             })();
             return {
                 initialResult,
-                updates: handler.iterable,
-                close: () => {
-                    outboundSub.send({
-                        unsubscribe: {
-                            subscriptionIdentifier: subscriptionId,
-                        },
-                    });
-                    outboundSub.close();
-                    subscriptions.delete(queryId);
-                    handler.complete();
+                updates: {
+                    [Symbol.asyncIterator]() {
+                        const iterator = handler.iterable[Symbol.asyncIterator]();
+                        return {
+                            next: () => iterator.next(),
+                            async return() {
+                                closeSubscription();
+                                return iterator.return ? iterator.return() : { value: undefined, done: true };
+                            },
+                        };
+                    },
                 },
+                close: () => closeSubscription(),
             };
         },
         subscribeToUpdates(unstamped, bufferSize) {
+            if (shutdownLatch.shuttingDown)
+                throw new Error("Messaging shutdown in progress");
+            if (subscriptions.size >= 1024)
+                throw new Error("Subscription capacity 1024 exhausted");
             const message = { ...unstamped, timestamp: unstamped.timestamp ?? Date.now() };
             const queryId = message.identifier;
             if (subscriptions.has(queryId)) {
                 throw new Error(`Subscription query already registered for identifier "${queryId}"`);
             }
-            const handler = updateHandler(message, bufferSize);
+            let removeShutdown;
+            const handler = updateHandler(message, bufferSize, () => { subscriptions.delete(queryId); removeShutdown?.(); });
+            removeShutdown = shutdownLatch.onShutdown(() => handler.completeExceptionally(new Error("Messaging shutdown in progress")));
             subscriptions.set(queryId, handler);
             return {
                 [Symbol.asyncIterator]: () => handler.iterable[Symbol.asyncIterator](),
@@ -778,7 +952,7 @@ export function axonServerQueryBus(next, conn, options = {}) {
                 },
             };
         },
-        async emitUpdate(queryName, filter, update) {
+        async emitUpdate(queryName, filter, update, uow) {
             runAfterCommitOrImmediately(() => {
                 for (const [subId, sub] of handlerSubscriptions) {
                     if (sub.queryName !== queryName)
@@ -803,9 +977,9 @@ export function axonServerQueryBus(next, conn, options = {}) {
                         instructionId: "",
                     });
                 }
-            });
+            }, uow);
         },
-        async completeSubscription(queryName, filter) {
+        async completeSubscription(queryName, filter, uow) {
             runAfterCommitOrImmediately(() => {
                 for (const [subId, sub] of handlerSubscriptions) {
                     if (sub.queryName !== queryName)
@@ -825,9 +999,9 @@ export function axonServerQueryBus(next, conn, options = {}) {
                     });
                     handlerSubscriptions.delete(subId);
                 }
-            });
+            }, uow);
         },
-        async completeSubscriptionExceptionally(queryName, error, filter) {
+        async completeSubscriptionExceptionally(queryName, error, filter, uow) {
             runAfterCommitOrImmediately(() => {
                 for (const [subId, sub] of handlerSubscriptions) {
                     if (sub.queryName !== queryName)
@@ -854,7 +1028,7 @@ export function axonServerQueryBus(next, conn, options = {}) {
                     });
                     handlerSubscriptions.delete(subId);
                 }
-            });
+            }, uow);
         },
     };
     return routing;