@bastani/atomic 0.9.45 → 0.9.46-alpha.1

package/dist/builtin/intercom/broker/client.ts

@@ -36,6 +36,34 @@ export const PENDING_STAGE_ROUTE_ACK_TIMEOUT_MS = 10_000;
 const GROUP_REQUEST_TIMEOUT_MS = 5000;
 const PRESENCE_ACK_TIMEOUT_MS = 5000;
 const RESTORING_RECIPIENT_REQUEST_TIMEOUT_MS = 5000;
+/** #3904: retries of a connection the kernel refused before it reached the broker (a full listen backlog). */
+const BROKER_BACKLOG_RETRIES = 6;
+const BROKER_BACKLOG_RETRY_BASE_MS = 25;
+
+/** The connect attempt was cancelled by `disconnect()` before it registered. */
+function connectAborted(): IntercomClientDisconnectedError {
+	return new IntercomClientDisconnectedError({ cause: new Error("Disconnected before the broker connection registered") });
+}
+
+function abortableDelay(ms: number, signal: AbortSignal): Promise<void> {
+	return new Promise((resolveDelay, rejectDelay) => {
+		const timer = setTimeout(() => {
+			signal.removeEventListener("abort", onAbort);
+			resolveDelay();
+		}, ms);
+		const onAbort = () => {
+			clearTimeout(timer);
+			rejectDelay(connectAborted());
+		};
+		if (signal.aborted) onAbort();
+		else signal.addEventListener("abort", onAbort, { once: true });
+	});
+}
+
+function isBacklogRefusal(error: Error): boolean {
+	const code = (error as NodeJS.ErrnoException).code;
+	return code === "ECONNREFUSED" || code === "EAGAIN";
+}
 /** Extra wait past a broker hold, so the broker's own outcome for a held send always arrives first (#3760). */
 const HELD_SEND_GRACE_MS = 5000;
 
@@ -93,6 +121,8 @@ export interface PendingStageMessageRequest {
 	readonly senderRegistrationName?: string;
 	readonly senderReturnAddress?: string;
 	readonly live?: boolean;
+	/** #4006: the broker authorized the sender as the session that launched the run, not by membership. */
+	readonly invocationControl?: true;
 }
 
 /** Broker → route owner: the live stage targets a sticky broadcast was actually written to. */
@@ -212,6 +242,12 @@ export class IntercomClient extends EventEmitter {
 	/** #3906: the broker of the current registration refuses a pending-stage route without ending the connection. */
 	private brokerRefusesPendingStageRoutesInPlace = false;
 	private brokerAnswersStatus = false;
+	/** #4006: the broker of the current registration accepts `claim_workflow_invocation`. */
+	private brokerClaimsWorkflowInvocations = false;
+	/** #3904: the connect attempt in progress, which `disconnect()` cancels. */
+	private connectAttempt: AbortController | undefined;
+	/** #3983: the broker of the current registration accepts `saving`. */
+	private brokerExtendsSaves = false;
   private socket: net.Socket | null = null;
   private _sessionId: string | null = null;
   private _supervisorSessionId: string | null = null;
@@ -231,6 +267,7 @@ export class IntercomClient extends EventEmitter {
   private pendingBrokerStatus = new Map<string, { resolve: (report: BrokerStatusReport) => void; reject: (error: Error) => void }>();
   private _groups = normalizeGroups();
   private pendingLiveWorkflowStageRoutes = new Map<string, { resolve(): void; reject(error: Error): void }>();
+  private pendingWorkflowInvocationClaims = new Map<string, { resolve(claimed: boolean): void; reject(error: Error): void }>();
   private pendingStageRouteRegistrations = new Map<string, { resolve(): void; reject(error: Error): void }>();
 
   private pendingSupervisorAuthorizations = new Map<string, {
@@ -297,6 +334,7 @@ export class IntercomClient extends EventEmitter {
     for (const pending of this.pendingPresence.values()) pending.reject(error);
     for (const pending of this.pendingMemberships.values()) pending.reject(error);
     for (const pending of this.pendingLiveWorkflowStageRoutes.values()) pending.reject(error);
+    for (const pending of this.pendingWorkflowInvocationClaims.values()) pending.reject(error);
     for (const pending of this.pendingStageRouteRegistrations.values()) pending.reject(error);
     for (const pending of this.pendingBrokerStatus.values()) pending.reject(error);
     this.pendingLists.clear();
@@ -306,6 +344,7 @@ export class IntercomClient extends EventEmitter {
     this.pendingPresence.clear();
     this.pendingMemberships.clear();
     this.pendingLiveWorkflowStageRoutes.clear();
+    this.pendingWorkflowInvocationClaims.clear();
     this.pendingStageRouteRegistrations.clear();
     this.pendingBrokerStatus.clear();
   }
@@ -345,7 +384,43 @@ export class IntercomClient extends EventEmitter {
     return socket;
   }
 
-  connect(
+  /**
+   * Connect and register. A connection the operating system refuses before it reaches the broker (`ECONNREFUSED`,
+   * `EAGAIN`) is retried a few times with a short jittered backoff: when many sessions start at once (a large pooled
+   * fan-out) the broker socket's listen backlog fills and the kernel refuses connections the broker would accept a
+   * moment later (#3904). Every other failure, and a refusal that persists past the retries, rejects as before.
+   */
+  async connect(
+    session: Omit<SessionInfo, "id">,
+    supervisor?: SupervisorRegistration,
+    supervisorOwnerToken?: string,
+    messageSource?: Message["source"],
+    registrationGroup?: string,
+  ): Promise<void> {
+    // #3904: `disconnect()` aborts this attempt, its retry delay and any later retry, so a session that finished
+    // its cleanup never ends up connected and registered by a retry still in flight.
+    const attempt = new AbortController();
+    this.connectAttempt = attempt;
+    try {
+      for (let retry = 0; ; retry++) {
+        if (attempt.signal.aborted) throw connectAborted();
+        try {
+          await this.connectOnce(attempt.signal, session, supervisor, supervisorOwnerToken, messageSource, registrationGroup);
+          return;
+        } catch (error) {
+          if (attempt.signal.aborted) throw connectAborted();
+          if (!(error instanceof Error) || !isBacklogRefusal(error) || retry >= BROKER_BACKLOG_RETRIES) throw error;
+          const backoff = BROKER_BACKLOG_RETRY_BASE_MS * 2 ** retry;
+          await abortableDelay(backoff * (0.5 + Math.random() / 2), attempt.signal);
+        }
+      }
+    } finally {
+      if (this.connectAttempt === attempt) this.connectAttempt = undefined;
+    }
+  }
+
+  private connectOnce(
+    signal: AbortSignal,
     session: Omit<SessionInfo, "id">,
     supervisor?: SupervisorRegistration,
     supervisorOwnerToken?: string,
@@ -413,6 +488,8 @@ export class IntercomClient extends EventEmitter {
         this.brokerAcknowledgesPendingStageRoutes = false;
         this.brokerRefusesPendingStageRoutesInPlace = false;
         this.brokerAnswersStatus = false;
+        this.brokerClaimsWorkflowInvocations = false;
+        this.brokerExtendsSaves = false;
         this._messageSource = undefined;
         this.disconnectError = null;
         if (connectionEstablished && !wasDisconnecting) {
@@ -459,7 +536,38 @@ export class IntercomClient extends EventEmitter {
         this.handleBrokerMessage(msg);
       }, onReaderError);
       
+      const onConnect = (): void => {
+        if (signal.aborted) {
+          onAbort();
+          return;
+        }
+        try {
+          writeMessage(socket, {
+            type: "register",
+            session,
+            ...(registrationGroup !== undefined ? { registrationGroup } : {}),
+			returnAddress: this.returnAddress,
+			...(this.preferredSessionId === undefined ? {} : { preferredSessionId: this.preferredSessionId }),
+			...(this.acknowledgesReceipt ? { acknowledgesReceipt: true as const } : {}),
+            ...(supervisor ? { supervisor } : {}),
+            ...(supervisorOwnerToken ? { supervisorOwnerToken } : {}),
+          });
+        } catch (error) {
+          cleanupConnectionAttempt();
+          cleanupSocketListeners();
+          if (this.socket === socket) {
+            this.socket = null;
+          }
+          socket.destroy();
+          reject(toError(error));
+        }
+      };
+
+      const onAbort = () => onError(connectAborted());
+
       const cleanupConnectionAttempt = () => {
+        signal.removeEventListener("abort", onAbort);
+        socket.off("connect", onConnect);
         this.off("_registered", onRegistered);
         this.off("_registration_failed", onRegistrationFailed);
         socket.off("error", onError);
@@ -480,26 +588,9 @@ export class IntercomClient extends EventEmitter {
       this.once("_registered", onRegistered);
       this.once("_registration_failed", onRegistrationFailed);
       
-      try {
-        writeMessage(socket, {
-          type: "register",
-          session,
-          ...(registrationGroup !== undefined ? { registrationGroup } : {}),
-			returnAddress: this.returnAddress,
-			...(this.preferredSessionId === undefined ? {} : { preferredSessionId: this.preferredSessionId }),
-			...(this.acknowledgesReceipt ? { acknowledgesReceipt: true as const } : {}),
-          ...(supervisor ? { supervisor } : {}),
-          ...(supervisorOwnerToken ? { supervisorOwnerToken } : {}),
-        });
-      } catch (error) {
-        cleanupConnectionAttempt();
-        cleanupSocketListeners();
-        if (this.socket === socket) {
-          this.socket = null;
-        }
-        socket.destroy();
-        reject(toError(error));
-      }
+      // Register only once connected, and only if the attempt was not cancelled meanwhile (#3904).
+      signal.addEventListener("abort", onAbort, { once: true });
+      socket.once("connect", onConnect);
     });
   }
 
@@ -523,7 +614,9 @@ export class IntercomClient extends EventEmitter {
           || (brokerMessage.acknowledgesReceipt !== undefined && brokerMessage.acknowledgesReceipt !== true)
           || (brokerMessage.acknowledgesPendingStageRoutes !== undefined && brokerMessage.acknowledgesPendingStageRoutes !== true)
           || (brokerMessage.refusesPendingStageRoutesInPlace !== undefined && brokerMessage.refusesPendingStageRoutesInPlace !== true)
-          || (brokerMessage.brokerStatus !== undefined && brokerMessage.brokerStatus !== true)) {
+          || (brokerMessage.brokerStatus !== undefined && brokerMessage.brokerStatus !== true)
+          || (brokerMessage.claimsWorkflowInvocations !== undefined && brokerMessage.claimsWorkflowInvocations !== true)
+          || (brokerMessage.extendsSaves !== undefined && brokerMessage.extendsSaves !== true)) {
           throw new Error("Invalid registered message");
         }
 
@@ -539,6 +632,8 @@ export class IntercomClient extends EventEmitter {
         this.brokerAcknowledgesPendingStageRoutes = brokerMessage.acknowledgesPendingStageRoutes === true;
         this.brokerRefusesPendingStageRoutesInPlace = brokerMessage.refusesPendingStageRoutesInPlace === true;
         this.brokerAnswersStatus = brokerMessage.brokerStatus === true;
+        this.brokerClaimsWorkflowInvocations = brokerMessage.claimsWorkflowInvocations === true;
+        this.brokerExtendsSaves = this.brokerAcknowledgesReceipt && brokerMessage.extendsSaves === true;
         this.emit("_registered", { type: "registered", sessionId: brokerMessage.sessionId });
         break;
       }
@@ -661,7 +756,8 @@ export class IntercomClient extends EventEmitter {
         break;
       }
 		case "pending_stage_message": {
-			const { requestId, from, senderRegistrationName, senderReturnAddress, runId, target, message, live } = brokerMessage;
+			const { requestId, from, senderRegistrationName, senderReturnAddress, runId, target, message, live, invocationControl } =
+				brokerMessage;
 			if (
 				typeof requestId !== "string" ||
 				!isSessionInfo(from) ||
@@ -670,7 +766,8 @@ export class IntercomClient extends EventEmitter {
 				typeof runId !== "string" ||
 				typeof target !== "string" ||
 				!isMessage(message) ||
-				(live !== undefined && typeof live !== "boolean")
+				(live !== undefined && typeof live !== "boolean") ||
+				(invocationControl !== undefined && invocationControl !== true)
 			) {
 				throw new Error("Invalid pending-stage message event");
 			}
@@ -683,6 +780,7 @@ export class IntercomClient extends EventEmitter {
 				target,
 				message,
 				...(live === true ? { live: true } : {}),
+				...(invocationControl === true ? { invocationControl: true } : {}),
 			} satisfies PendingStageMessageRequest);
 			break;
 		}
@@ -753,6 +851,15 @@ export class IntercomClient extends EventEmitter {
         pending.resolve();
         break;
       }
+      case "workflow_invocation_claimed": {
+        const { requestId, claimed } = brokerMessage;
+        if (typeof requestId !== "string" || typeof claimed !== "boolean") throw new Error("Invalid workflow invocation claim answer");
+        const pending = this.pendingWorkflowInvocationClaims.get(requestId);
+        if (!pending) return;
+        this.pendingWorkflowInvocationClaims.delete(requestId);
+        pending.resolve(claimed);
+        break;
+      }
       case "live_workflow_stage_route_registered": {
         const { requestId } = brokerMessage;
         if (typeof requestId !== "string") throw new Error("Invalid live workflow-stage route registration");
@@ -931,6 +1038,12 @@ export class IntercomClient extends EventEmitter {
     }
   }
   async disconnect(): Promise<void> {
+    // A connect still in progress (or waiting to retry) is cancelled: it closes any socket it opened and never
+    // registers (#3904).
+    if (this.connectAttempt !== undefined && this._sessionId === null) {
+      this.connectAttempt.abort();
+      return;
+    }
     const socket = this.socket;
     if (!socket) {
       return;
@@ -1194,6 +1307,7 @@ export class IntercomClient extends EventEmitter {
 	possibleStages?: WorkflowPossibleStageAnnouncement[],
 	parent?: WorkflowRunParentAnnouncement,
 	dormant?: boolean,
+	launcherCapability?: string,
   ): void {
 	writeMessage(this.requireActiveSocket(), {
 		type: "register_pending_stage_route",
@@ -1204,6 +1318,7 @@ export class IntercomClient extends EventEmitter {
 		possibleStages,
 		...(parent === undefined ? {} : { parent }),
 		...(dormant === true ? { dormant: true as const } : {}),
+		...(launcherCapability === undefined || !this.brokerClaimsWorkflowInvocations ? {} : { launcherCapability }),
 	});
   }
 
@@ -1220,6 +1335,8 @@ export class IntercomClient extends EventEmitter {
 		possibleStages?: WorkflowPossibleStageAnnouncement[],
 		parent?: WorkflowRunParentAnnouncement,
 		dormant?: boolean,
+		/** #4006: the launcher-only secret, held by the process that launched the run and never by its stages. */
+		launcherCapability?: string,
 	): Promise<void> {
 		let socket: net.Socket;
 		try {
@@ -1229,7 +1346,7 @@ export class IntercomClient extends EventEmitter {
 		}
 		if (!this.brokerAcknowledgesPendingStageRoutes) {
 			try {
-				this.registerPendingStageRoute(runId, group, capability, stages, possibleStages, parent, dormant);
+				this.registerPendingStageRoute(runId, group, capability, stages, possibleStages, parent, dormant, launcherCapability);
 			} catch (error) {
 				return Promise.reject(toError(error));
 			}
@@ -1262,6 +1379,7 @@ export class IntercomClient extends EventEmitter {
 					possibleStages,
 					...(parent === undefined ? {} : { parent }),
 					...(dormant === true ? { dormant: true as const } : {}),
+					...(launcherCapability === undefined || !this.brokerClaimsWorkflowInvocations ? {} : { launcherCapability }),
 					...(this.multiplexesPendingStageRoutes && this.brokerRefusesPendingStageRoutesInPlace
 						? { refusalKeepsConnection: true as const }
 						: {}),
@@ -1274,6 +1392,31 @@ export class IntercomClient extends EventEmitter {
 		});
 	}
 
+  /**
+   * #4006: prove with the run's launcher capability (registered with the run's route, never given to its
+   * stages) that this session launched `runId`, so it may send to the run's stages without joining
+   * `workflow:<rootRunId>`. Resolves `false` when the broker refused the claim or predates it.
+   */
+  claimWorkflowInvocation(runId: string, launcherCapability: string): Promise<boolean> {
+    let socket: net.Socket;
+    try {
+      socket = this.requireActiveSocket();
+    } catch (error) {
+      return Promise.reject(toError(error));
+    }
+    if (!this.brokerClaimsWorkflowInvocations) return Promise.resolve(false);
+    const requestId = randomUUID();
+    return new Promise((resolve, reject) => {
+      this.pendingWorkflowInvocationClaims.set(requestId, { resolve, reject });
+      try {
+        writeMessage(socket, { type: "claim_workflow_invocation", requestId, runId, capability: launcherCapability });
+      } catch (error) {
+        this.pendingWorkflowInvocationClaims.delete(requestId);
+        reject(toError(error));
+      }
+    });
+  }
+
   registerLiveWorkflowStageRoute(runId: string, stageKeys: readonly string[], capability: string): Promise<void> {
     let socket: net.Socket;
     try {
@@ -1343,6 +1486,15 @@ export class IntercomClient extends EventEmitter {
 		writeMessage(this.requireActiveSocket(), { type: "received", messageId });
 	}
 
+	/**
+	 * #3983: tell the broker this message is still being saved, so its sender keeps waiting instead of hearing the
+	 * outcome is uncertain. A no-op on a broker that did not advertise `extendsSaves`, or while disconnected.
+	 */
+	reportSaving(messageId: string): void {
+		if (!this.brokerExtendsSaves || !this.socket || this.socket.destroyed) return;
+		writeMessage(this.socket, { type: "saving", messageId });
+	}
+
 	sendPendingStageNotification(
 		runId: string,
 		capability: string,