add more granular service control to supervisor session

Author: nicktrn

packages/core/src/v3/runEngineWorker/supervisor/session.ts

@@ -11,6 +11,8 @@ import { getDefaultWorkerHeaders } from "./util.js";
 import { HeartbeatService } from "../../utils/heartbeat.js";
 
 type SupervisorSessionOptions = SupervisorClientCommonOptions & {
+  queueConsumerEnabled?: boolean;
+  runNotificationsEnabled?: boolean;
   heartbeatIntervalSeconds?: number;
   dequeueIntervalMs?: number;
   preDequeue?: PreDequeueFn;
@@ -20,15 +22,21 @@ type SupervisorSessionOptions = SupervisorClientCommonOptions & {
 export class SupervisorSession extends EventEmitter<WorkerEvents> {
   public readonly httpClient: SupervisorHttpClient;
 
-  private socket?: Socket<WorkerServerToClientEvents, WorkerClientToServerEvents>;
+  private readonly runNotificationsEnabled: boolean;
+  private runNotificationsSocket?: Socket<WorkerServerToClientEvents, WorkerClientToServerEvents>;
 
+  private readonly queueConsumerEnabled: boolean;
   private readonly queueConsumer: RunQueueConsumer;
+
   private readonly heartbeatService: HeartbeatService;
   private readonly heartbeatIntervalSeconds: number;
 
   constructor(private opts: SupervisorSessionOptions) {
     super();
 
+    this.runNotificationsEnabled = opts.runNotificationsEnabled ?? true;
+    this.queueConsumerEnabled = opts.queueConsumerEnabled ?? true;
+
     this.httpClient = new SupervisorHttpClient(opts);
     this.queueConsumer = new RunQueueConsumer({
       client: this.httpClient,
@@ -76,12 +84,12 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
   subscribeToRunNotifications(runFriendlyIds: string[]) {
     console.log("[SupervisorSession] Subscribing to run notifications", { runFriendlyIds });
 
-    if (!this.socket) {
+    if (!this.runNotificationsSocket) {
       console.error("[SupervisorSession] Socket not connected");
       return;
     }
 
-    this.socket.emit("run:subscribe", { version: "1", runFriendlyIds });
+    this.runNotificationsSocket.emit("run:subscribe", { version: "1", runFriendlyIds });
 
     Promise.allSettled(
       runFriendlyIds.map((runFriendlyId) =>
@@ -96,12 +104,12 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
   unsubscribeFromRunNotifications(runFriendlyIds: string[]) {
     console.log("[SupervisorSession] Unsubscribing from run notifications", { runFriendlyIds });
 
-    if (!this.socket) {
+    if (!this.runNotificationsSocket) {
       console.error("[SupervisorSession] Socket not connected");
       return;
     }
 
-    this.socket.emit("run:unsubscribe", { version: "1", runFriendlyIds });
+    this.runNotificationsSocket.emit("run:unsubscribe", { version: "1", runFriendlyIds });
 
     Promise.allSettled(
       runFriendlyIds.map((runFriendlyId) =>
@@ -116,15 +124,15 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
     );
   }
 
-  private createSocket() {
+  private createRunNotificationsSocket() {
     const wsUrl = new URL(this.opts.apiUrl);
     wsUrl.pathname = "/worker";
 
-    this.socket = io(wsUrl.href, {
+    const socket = io(wsUrl.href, {
       transports: ["websocket"],
       extraHeaders: getDefaultWorkerHeaders(this.opts),
     });
-    this.socket.on("run:notify", ({ version, run }) => {
+    socket.on("run:notify", ({ version, run }) => {
       console.log("[SupervisorSession][WS] Received run notification", { version, run });
       this.emit("runNotification", { time: new Date(), run });
 
@@ -137,15 +145,17 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
           console.error("[SupervisorSession] Failed to send debug log", { error });
         });
     });
-    this.socket.on("connect", () => {
+    socket.on("connect", () => {
       console.log("[SupervisorSession][WS] Connected to platform");
     });
-    this.socket.on("connect_error", (error) => {
+    socket.on("connect_error", (error) => {
       console.error("[SupervisorSession][WS] Connection error", { error });
     });
-    this.socket.on("disconnect", (reason, description) => {
+    socket.on("disconnect", (reason, description) => {
       console.log("[SupervisorSession][WS] Disconnected from platform", { reason, description });
     });
+
+    return socket;
   }
 
   async start() {
@@ -167,14 +177,25 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
       name: workerGroup.name,
     });
 
-    this.queueConsumer.start();
-    this.heartbeatService.start();
-    this.createSocket();
+    if (this.queueConsumerEnabled) {
+      console.log("[SupervisorSession] Queue consumer enabled");
+      this.queueConsumer.start();
+      this.heartbeatService.start();
+    } else {
+      console.warn("[SupervisorSession] Queue consumer disabled");
+    }
+
+    if (this.runNotificationsEnabled) {
+      console.log("[SupervisorSession] Run notifications enabled");
+      this.runNotificationsSocket = this.createRunNotificationsSocket();
+    } else {
+      console.warn("[SupervisorSession] Run notifications disabled");
+    }
   }
 
   async stop() {
     this.heartbeatService.stop();
-    this.socket?.disconnect();
+    this.runNotificationsSocket?.disconnect();
   }
 
   private getHeartbeatBody(): WorkerApiHeartbeatRequestBody {