@axiom-lattice/gateway 2.1.88 → 2.1.89

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/gateway",
-  "version": "2.1.88",
+  "version": "2.1.89",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
@@ -40,9 +40,9 @@
     "redis": "^5.0.1",
     "uuid": "^9.0.1",
     "zod": "3.25.76",
-    "@axiom-lattice/agent-eval": "2.1.70",
-    "@axiom-lattice/core": "2.1.76",
-    "@axiom-lattice/pg-stores": "1.0.67",
+    "@axiom-lattice/agent-eval": "2.1.71",
+    "@axiom-lattice/core": "2.1.77",
+    "@axiom-lattice/pg-stores": "1.0.68",
     "@axiom-lattice/protocols": "2.1.39",
     "@axiom-lattice/queue-redis": "1.0.38"
   },

package/src/router/MessageRouter.ts

@@ -19,6 +19,14 @@ export class BindingNotFoundError extends Error {
   }
 }
 
+/**
+ * Configuration for {@link MessageRouter}.
+ *
+ * @param middlewares - Ordered middleware chain executed before dispatch
+ * @param bindingRegistry - Resolves sender-to-agent bindings
+ * @param adapterRegistry - Registry of {@link ChannelAdapter} implementations
+ * @param installationStore - Persisted channel installation configs
+ */
 export interface MessageRouterConfig {
   middlewares: MessageMiddleware[];
   bindingRegistry: BindingRegistry;
@@ -26,12 +34,49 @@ export interface MessageRouterConfig {
   installationStore: ChannelInstallationStore;
 }
 
+/**
+ * Core message router for external channel integration.
+ *
+ * Receives normalized {@link InboundMessage} objects from channel adapters,
+ * resolves the target agent via {@link BindingRegistry}, manages thread
+ * lifecycle, and sends AI responses back to the channel via
+ * `ChannelAdapter.sendReply`.
+ *
+ * ## Reply flow
+ *
+ * The router uses a counter-based deduplication scheme on the internal
+ * {@link Agent.reply:ready} event:
+ *
+ * 1. Before `agent.addMessage()`, it subscribes to `reply:ready` on the agent.
+ * 2. Multiple concurrent dispatches on the same thread increment a counter;
+ *    only the first registers the `EventEmitter.once` listener.
+ * 3. When `reply:ready` fires, the callback extracts the last AI message from
+ *    the LangGraph state and sends it via `ChannelAdapter.sendReply`.
+ * 4. The counter is decremented; the subscription is only removed when the
+ *    counter reaches 0.
+ * 5. Stale subscriptions are cleaned up after a 1-hour timeout.
+ *
+ * The `replyTarget` is carried through by injecting `_replyTarget` into
+ * `custom_run_config` on {@link Agent.addMessage}, and retrieved from the
+ * `reply:ready` event payload.
+ *
+ * @see {@link ChannelAdapter.sendReply}
+ * @see {@link InboundMessage.replyTarget}
+ */
 export class MessageRouter {
   private middlewares: MessageMiddleware[];
   private bindingRegistry: BindingRegistry;
   private adapterRegistry: ChannelAdapterRegistry;
   private installationStore: ChannelInstallationStore;
 
+  /**
+   * Tracks reply subscriptions per thread+channel to avoid duplicate
+   * `subscribeOnce` registrations and ensure proper cleanup.
+   *
+   * Key format: `{threadId}:{adapterChannel}:reply`
+   */
+  private _replySubs: Map<string, { count: number; timer: NodeJS.Timeout }> = new Map();
+
   constructor(config: MessageRouterConfig) {
     this.middlewares = [...config.middlewares];
     this.bindingRegistry = config.bindingRegistry;
@@ -39,10 +84,28 @@ export class MessageRouter {
     this.installationStore = config.installationStore;
   }
 
+  /**
+   * Register an additional middleware at the end of the chain.
+   *
+   * @param middleware - A {@link MessageMiddleware} function
+   */
   use(middleware: MessageMiddleware): void {
     this.middlewares.push(middleware);
   }
 
+  /**
+   * Dispatch an inbound channel message to the bound agent.
+   *
+   * Full pipeline: middleware chain → binding resolution → thread lifecycle
+   * → agent.addMessage() → (async) reply via {@link ChannelAdapter.sendReply}.
+   *
+   * If the message has a {@link InboundMessage.replyTarget}, the router subscribes
+   * to the agent's `reply:ready` event before enqueuing the message, and sends
+   * the AI response back to the channel when it arrives.
+   *
+   * @param message - Normalized inbound message from a channel adapter
+   * @returns {@link DispatchResult} with success status, bindingId, and threadId
+   */
   async dispatch(message: InboundMessage): Promise<DispatchResult> {
     const ctx: MessageContext = {
       inboundMessage: message,
@@ -208,9 +271,114 @@ export class MessageRouter {
           project_id: ctx.binding.projectId || "",
         });
 
+        /**
+         * Subscribe to reply:ready before addMessage to avoid a race with
+         * an already-running queue processor.
+         *
+         * The `_replySubs` map uses a counter (not a boolean) so that multiple
+         * concurrent dispatches on the same thread each get a callback invocation.
+         * The `subscribeOnce` listener is only registered when the counter goes
+         * from 0 → 1, and removed when it goes back to 0. A 1-hour timeout
+         * protects against leaked subscriptions due to agent errors/aborts.
+         */
+        if (message.replyTarget) {
+          const replySubKey = `${threadId}:${message.replyTarget.adapterChannel}:reply`;
+          const adapter = this.adapterRegistry.get(message.replyTarget.adapterChannel);
+          if (adapter) {
+            const installation = await this.installationStore.getInstallationById(
+              message.channelInstallationId,
+            );
+            if (installation) {
+              const existing = this._replySubs.get(replySubKey);
+              if (!existing || existing.count === 0) {
+                // First subscriber for this thread+channel — register the event listener
+                const timer = setTimeout(() => {
+                  // Cleanup if reply never fires (agent aborted, errored, etc.)
+                  const entry = this._replySubs.get(replySubKey);
+                  if (entry) {
+                    this._replySubs.delete(replySubKey);
+                    console.warn({
+                      event: "dispatch:reply:timeout",
+                      threadId,
+                      channel: message.replyTarget!.adapterChannel,
+                    }, "Reply subscription timed out — no reply:ready fired within 1h");
+                  }
+                }, 3_600_000);
+
+                this._replySubs.set(replySubKey, { count: 1, timer });
+
+                console.log({
+                  event: "dispatch:reply:subscribed",
+                  threadId,
+                  channel: message.replyTarget.adapterChannel,
+                  senderId: message.sender.id,
+                }, "Subscribed to reply:ready for thread");
+
+                agent.subscribeOnce('reply:ready', (data: any) => {
+                  const messages = data.state?.values?.messages;
+                  const lastAI = messages
+                    ?.filter((m: any) => m.type === 'ai' || m.getType?.() === 'ai')
+                    .pop();
+                  const replyText = lastAI?.content ?? '';
+
+                  // Decrement counter; unsubscribe only when count reaches 0
+                  const entry = this._replySubs.get(replySubKey);
+                  if (entry) {
+                    entry.count--;
+                    if (entry.count <= 0) {
+                      clearTimeout(entry.timer);
+                      this._replySubs.delete(replySubKey);
+                    }
+                  }
+
+                  if (replyText) {
+                    console.log({
+                      event: "dispatch:reply:sending",
+                      threadId,
+                      channel: message.replyTarget!.adapterChannel,
+                      replyLength: replyText.length,
+                    }, "Sending channel reply");
+                    adapter.sendReply(message.replyTarget!, { text: replyText }, installation)
+                      .then(() => {
+                        console.log({
+                          event: "dispatch:reply:sent",
+                          threadId,
+                          channel: message.replyTarget!.adapterChannel,
+                        }, "Channel reply sent successfully");
+                      })
+                      .catch((err) => console.error({
+                        event: "dispatch:reply:failed",
+                        threadId,
+                        channel: message.replyTarget!.adapterChannel,
+                        error: err instanceof Error ? err.message : String(err),
+                      }, "Failed to send channel reply"));
+                  } else {
+                    console.warn({
+                      event: "dispatch:reply:empty",
+                      threadId,
+                      channel: message.replyTarget!.adapterChannel,
+                    }, "Agent produced no text output — skipping reply");
+                  }
+                });
+              } else {
+                // Another dispatch already subscribed — increment counter
+                existing.count++;
+                console.log({
+                  event: "dispatch:reply:incremented",
+                  threadId,
+                  channel: message.replyTarget.adapterChannel,
+                  count: existing.count,
+                }, "Incremented reply counter for thread (already subscribed)");
+              }
+            }
+          }
+        }
+
         const addResult = await agent.addMessage({
           input: { message: message.content.text },
-          custom_run_config: message.content.metadata || {},
+          custom_run_config: message.replyTarget
+            ? { ...(message.content.metadata || {}), _replyTarget: message.replyTarget }
+            : (message.content.metadata || {}),
         });
 
         console.log({
@@ -220,22 +388,6 @@ export class MessageRouter {
           messageId: (addResult as Record<string, unknown>)?.messageId,
           result: JSON.stringify(addResult),
         }, "Agent dispatch complete — messageId = " + ((addResult as Record<string, unknown>)?.messageId || "N/A"));
-
-        if (message.replyTarget) {
-          const adapter = this.adapterRegistry.get(message.replyTarget.adapterChannel);
-          if (adapter) {
-            const installation = await this.installationStore.getInstallationById(
-              message.channelInstallationId,
-            );
-            if (installation) {
-              await adapter.sendReply(
-                message.replyTarget,
-                { text: ctx.result || "" },
-                installation,
-              );
-            }
-          }
-        }
       });
 
       return {