nolo-cli 0.1.30 → 0.1.32

package/machineCommands.ts

@@ -1,4 +1,5 @@
 import type { MachineHeartbeat } from "./connector-experimental/protocol";
+import { NOLO_CLUSTER_SERVERS, normalizeKnownServerOrigin } from "./database/config";
 import { detectMachineInfo } from "./connector-experimental/machineInfo";
 import {
   type HeartbeatLoopOptions,
@@ -42,6 +43,7 @@ export {
 type MachineCommandDeps = {
   env?: EnvLike;
   output?: OutputLike;
+  signal?: AbortSignal;
   cliEntrypointPath?: string;
   maxConnectorAttempts?: number;
   sleep?: (ms: number) => Promise<void>;
@@ -60,6 +62,100 @@ type MachineCommandDeps = {
 };
 
 const defaultSleep = (ms: number) => new Promise<void>((resolve) => setTimeout(resolve, ms));
+const MACHINE_TOKEN_PREFIX = "sk_machine_";
+
+async function sleepWithAbort(ms: number, sleep: (ms: number) => Promise<void>, signal?: AbortSignal) {
+  if (!signal) {
+    await sleep(ms);
+    return;
+  }
+  if (signal.aborted) return;
+  await Promise.race([
+    sleep(ms),
+    new Promise<void>((resolve) => {
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    }),
+  ]);
+}
+
+function resolveConnectorReconnectDelayMs(env: EnvLike) {
+  const raw = Number(env.NOLO_CONNECT_RECONNECT_MS ?? "");
+  return Number.isFinite(raw) && raw >= 0 ? raw : 5_000;
+}
+
+function isMachineToken(token: string) {
+  return token.trim().startsWith(MACHINE_TOKEN_PREFIX);
+}
+
+function resolveClusterConnectServerUrls(serverUrl: string, authToken: string) {
+  const normalizedPrimary = normalizeKnownServerOrigin(serverUrl) ?? serverUrl.trim().replace(/\/+$/, "");
+  if (isMachineToken(authToken)) return [normalizedPrimary];
+  const normalizedCluster = NOLO_CLUSTER_SERVERS
+    .map((baseUrl) => normalizeKnownServerOrigin(baseUrl) ?? baseUrl)
+    .filter((value, index, list) => value && list.indexOf(value) === index);
+  if (!normalizedCluster.includes(normalizedPrimary)) return [normalizedPrimary];
+  return [normalizedPrimary, ...normalizedCluster.filter((baseUrl) => baseUrl !== normalizedPrimary)];
+}
+
+function linkAbortSignal(signal: AbortSignal | undefined, controller: AbortController) {
+  if (!signal) return;
+  if (signal.aborted) {
+    controller.abort();
+    return;
+  }
+  signal.addEventListener("abort", () => controller.abort(), { once: true });
+}
+
+async function sendMachineHeartbeat(args: {
+  serverUrl: string;
+  authToken: string;
+  fetchImpl: typeof fetch;
+  machine: MachineHeartbeat;
+}) {
+  const res = await args.fetchImpl(`${args.serverUrl}/api/machines/heartbeat`, {
+    method: "POST",
+    headers: {
+      Authorization: `Bearer ${args.authToken}`,
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(args.machine),
+  });
+
+  if (!res.ok) {
+    const text = await res.text().catch(() => "");
+    throw new Error(`HTTP ${res.status}\n${text}`);
+  }
+}
+
+async function sendHeartbeatAcrossServers(args: {
+  serverUrls: string[];
+  authToken: string;
+  fetchImpl: typeof fetch;
+  machine: MachineHeartbeat;
+  output: OutputLike;
+}) {
+  let primaryError: Error | null = null;
+  for (const [index, serverUrl] of args.serverUrls.entries()) {
+    try {
+      await sendMachineHeartbeat({
+        serverUrl,
+        authToken: args.authToken,
+        fetchImpl: args.fetchImpl,
+        machine: args.machine,
+      });
+    } catch (error) {
+      const resolvedError = error instanceof Error ? error : new Error(String(error));
+      if (index === 0) {
+        primaryError = resolvedError;
+      } else {
+        args.output.write(
+          `[nolo] secondary connector registration failed for ${serverUrl}: ${resolvedError.message}\n`
+        );
+      }
+    }
+  }
+  if (primaryError) throw primaryError;
+}
 
 async function defaultExecuteCli(
   provider: string,
@@ -75,9 +171,9 @@ function detectLaunchableMachineInfo() {
 }
 
 export async function runMachineConnectCommand(
-  args: string[],
-  deps: MachineCommandDeps = {}
-) {
+  args: string[],
+  deps: MachineCommandDeps = {}
+) {
   const env = deps.env ?? process.env;
   const output = deps.output ?? process.stdout;
   const authToken = resolveConnectAuthToken(args, env);
@@ -85,29 +181,31 @@ export async function runMachineConnectCommand(
     writeAuthMissing(output);
     return 1;
   }
-
+
   const serverUrl = resolveConnectServerUrl(args, env);
+  const serverUrls = resolveClusterConnectServerUrls(serverUrl, authToken);
   const fetchImpl = deps.fetchImpl ?? fetch;
   const machine = (deps.machineInfo ?? detectLaunchableMachineInfo)();
   const sendHeartbeat = async () => {
-    const res = await fetchImpl(`${serverUrl}/api/machines/heartbeat`, {
-    method: "POST",
-    headers: {
-      Authorization: `Bearer ${authToken}`,
-      "Content-Type": "application/json",
-    },
-    body: JSON.stringify(machine),
-    });
-
-    if (!res.ok) {
-      const text = await res.text().catch(() => "");
-      throw new Error(`HTTP ${res.status}\n${text}`);
-    }
-  };
+    await sendHeartbeatAcrossServers({
+      serverUrls,
+      authToken,
+      fetchImpl,
+      machine,
+      output,
+    });
+  };
 
   if (hasFlag(args, "--daemon") || hasFlag(args, "--background")) {
     return runMachineDaemonCommand({
-      env,
+      env: {
+        ...env,
+        NOLO_SERVER: serverUrl,
+        NOLO_SERVER_URL: serverUrl,
+        BASE_URL: serverUrl,
+        AUTH_TOKEN: authToken,
+        NOLO_MACHINE_API_KEY: authToken,
+      },
       output,
       cliEntrypointPath: deps.cliEntrypointPath,
       validateWorkspaceLinks: deps.validateWorkspaceLinks,
@@ -123,50 +221,88 @@ export async function runMachineConnectCommand(
     const maxAttempts = deps.maxConnectorAttempts ?? Infinity;
     const sleep = deps.sleep ?? defaultSleep;
     const reconnectDelayMs = resolveConnectorReconnectDelayMs(env);
-    for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
-      const runtimeServerUrl =
-        env.NOLO_CONNECT_RUNTIME_SERVER_URL ||
-        env.NOLO_RUNTIME_SERVER_URL ||
-        serverUrl;
-      const exitCode = await runMachineWsSession({
-        env,
-        output,
-        machine,
-        sendHeartbeat,
-        runHeartbeatLoop: deps.runHeartbeatLoop ?? defaultRunHeartbeatLoop,
-        connectWebSocket: deps.connectWebSocket ?? defaultConnectWebSocket,
-        serverUrl,
-        authToken,
-        resolveConnectorWebSocketTarget: (options) =>
-          resolveConnectorWebSocketTarget({
-            ...options,
-            fetchImpl,
-          }),
-        onMessage: (message, send) =>
-          handleConnectorRunMessage(
-            message,
-            send,
-            deps.executeCli ?? defaultExecuteCli,
-            {
-              ...env,
-              NOLO_SERVER: runtimeServerUrl,
-              NOLO_SERVER_URL: runtimeServerUrl,
-              BASE_URL: runtimeServerUrl,
-              AUTH_TOKEN: authToken,
-              NOLO_MACHINE_API_KEY: authToken,
-            },
-            fetchImpl,
-            {
-              buildConnectorCliPrompt,
-            }
-          ),
-      });
-      if (exitCode === 2) return 1;
-      if (attempt >= maxAttempts) return exitCode;
-      output.write(`[nolo] Connector websocket disconnected. Reconnecting in ${reconnectDelayMs}ms.\n`);
-      await sleep(reconnectDelayMs);
-    }
-    return 0;
+    const runServerSession = async (targetServerUrl: string, signal?: AbortSignal) => {
+      for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+        if (signal?.aborted) {
+          return { serverUrl: targetServerUrl, exitCode: 0 };
+        }
+        const runtimeServerUrl =
+          env.NOLO_CONNECT_RUNTIME_SERVER_URL ||
+          env.NOLO_RUNTIME_SERVER_URL ||
+          targetServerUrl;
+        const sessionResult = await runMachineWsSession({
+          env,
+          output,
+          signal,
+          machine,
+          sendHeartbeat: () =>
+            sendMachineHeartbeat({
+              serverUrl: targetServerUrl,
+              authToken,
+              fetchImpl,
+              machine,
+            }),
+          runHeartbeatLoop: deps.runHeartbeatLoop ?? defaultRunHeartbeatLoop,
+          connectWebSocket: deps.connectWebSocket ?? defaultConnectWebSocket,
+          serverUrl: targetServerUrl,
+          authToken,
+          resolveConnectorWebSocketTarget: (options) =>
+            resolveConnectorWebSocketTarget({
+              ...options,
+              fetchImpl,
+            }),
+          onMessage: (message, send) =>
+            handleConnectorRunMessage(
+              message,
+              send,
+              deps.executeCli ?? defaultExecuteCli,
+              {
+                ...env,
+                NOLO_SERVER: runtimeServerUrl,
+                NOLO_SERVER_URL: runtimeServerUrl,
+                BASE_URL: runtimeServerUrl,
+                AUTH_TOKEN: authToken,
+                NOLO_MACHINE_API_KEY: authToken,
+              },
+              fetchImpl,
+              {
+                buildConnectorCliPrompt,
+              }
+            ),
+        });
+        const exitCode = sessionResult.exitCode;
+        if (exitCode === 2) return { serverUrl: targetServerUrl, exitCode: 2 };
+        if (signal?.aborted) {
+          return { serverUrl: targetServerUrl, exitCode: 0 };
+        }
+        if (attempt >= maxAttempts) return { serverUrl: targetServerUrl, exitCode };
+        const retryDelayMs = sessionResult.retryAfterMs ?? reconnectDelayMs;
+        const reconnectContext =
+          sessionResult.reconnectReason === "core_draining" ? " during core draining" : "";
+        output.write(
+          `[nolo] Connector websocket disconnected${reconnectContext}. Reconnecting in ${retryDelayMs}ms.\n`
+        );
+        await sleepWithAbort(retryDelayMs, sleep, signal);
+      }
+      return { serverUrl: targetServerUrl, exitCode: 0 };
+    };
+
+    const sessionAbort = new AbortController();
+    linkAbortSignal(deps.signal, sessionAbort);
+
+    const sessionPromises = serverUrls.map(async (targetServerUrl, index) => {
+      const result = await runServerSession(targetServerUrl, sessionAbort.signal);
+      if (index > 0 && result.exitCode !== 0 && !sessionAbort.signal.aborted) {
+        output.write(
+          `[nolo] secondary connector registration failed for ${targetServerUrl}: exit ${result.exitCode}\n`
+        );
+      }
+      return result;
+    });
+    const results = await Promise.all(sessionPromises);
+    if (results.some((result) => result.exitCode === 0)) return 0;
+    if (results.some((result) => result.exitCode === 2)) return 1;
+    return results[0]?.exitCode ?? 1;
   }
 
   if (hasFlag(args, "--watch")) {

package/machineWsRunDispatch.ts

@@ -1,4 +1,6 @@
 import { runLocalAgentTurn, type LocalAgentToolEvent } from "./agent-runtime/localLoop";
+import { resolveHostRuntimeConfigFromEnv } from "./agent-runtime/hostRuntimeConfig";
+import { normalizeAgentRunDialogLineage } from "./ai/agent/agentRunDialogLineage";
 import {
   assertMachineRunAllowed as defaultAssertMachineRunAllowed,
   resolveMachineRunPermissionPolicy as defaultResolveMachineRunPermissionPolicy,
@@ -56,7 +58,7 @@ type ConnectorRunProgress = {
 
 type PermissionPolicy = ReturnType<typeof defaultResolveMachineRunPermissionPolicy>;
 
-type MachineWsRunDispatchDeps = {
+export type MachineWsRunDispatchDeps = {
   assertMachineRunAllowed?: typeof defaultAssertMachineRunAllowed;
   resolveMachineRunPermissionPolicy?: typeof defaultResolveMachineRunPermissionPolicy;
   resolveConnectorRunCwd?: typeof defaultResolveConnectorRunCwd;
@@ -118,6 +120,24 @@ function requestsLocalWorkspaceRuntime(policy?: ConnectorRuntimePolicy) {
   );
 }
 
+function isMachineBoundLocalCustomProvider(agentConfig: any) {
+  const machineId =
+    isRecord(agentConfig?.runtimeBinding) && typeof agentConfig.runtimeBinding.machineId === "string"
+      ? agentConfig.runtimeBinding.machineId.trim()
+      : "";
+  const providerUrl =
+    typeof agentConfig?.customProviderUrl === "string"
+      ? agentConfig.customProviderUrl.trim()
+      : "";
+  if (!machineId || !providerUrl) return false;
+  try {
+    const hostname = new URL(providerUrl).hostname.toLowerCase();
+    return hostname === "127.0.0.1" || hostname === "localhost" || hostname === "::1";
+  } catch {
+    return false;
+  }
+}
+
 function hasExplicitMachinePermissions(agentConfig: any) {
   const runtimeBinding = isRecord(agentConfig?.runtimeBinding) ? agentConfig.runtimeBinding : {};
   return Boolean(
@@ -130,6 +150,7 @@ function hasExplicitMachinePermissions(agentConfig: any) {
 
 function runtimeWorkspacePermissionPolicy(
   runtimePolicy: ConnectorRuntimePolicy | undefined,
+  hostRuntimeConfig: { capabilities: { allowShell: boolean } },
   cwd: string
 ): PermissionPolicy | null {
   if (!requestsTaskWorktree(runtimePolicy)) return null;
@@ -137,7 +158,10 @@ function runtimeWorkspacePermissionPolicy(
     mode: "full_access",
     allowFilesystemRead: true,
     allowFilesystemWrite: true,
-    allowShell: runtimePolicy?.shell?.enabled !== false,
+    allowShell:
+      hostRuntimeConfig.capabilities.allowShell &&
+      Array.isArray(runtimePolicy?.runtimeTools) &&
+      runtimePolicy.runtimeTools.includes("execShell"),
     writableRoots: [cwd],
   };
 }
@@ -175,13 +199,20 @@ function mergeToolNames(...values: unknown[]) {
 function localRuntimeEnvFromPolicy(runtimeEnv: EnvLike, policy?: ConnectorRuntimePolicy): EnvLike {
   return {
     ...runtimeEnv,
-    ...(policy?.shell?.enabled && policy.shell.mode === "worktree"
-      ? { NOLO_LOCAL_SHELL_MODE: "worktree" }
-      : {}),
     ...(policy?.workspace?.cwd ? { NOLO_LOCAL_WORKTREE: policy.workspace.cwd } : {}),
   };
 }
 
+function forwardedUserAuthToken(parsed: any): string {
+  const value = parsed?.payload?.meta?.userAuthToken;
+  return typeof value === "string" ? value.trim() : "";
+}
+
+function withForwardedUserAuthToken(runtimeEnv: EnvLike, parsed: any): EnvLike {
+  const userAuthToken = forwardedUserAuthToken(parsed);
+  return userAuthToken ? { ...runtimeEnv, AUTH_TOKEN: userAuthToken } : runtimeEnv;
+}
+
 function runtimeWorkspaceRootFromTrace(trace: unknown): string | undefined {
   if (!Array.isArray(trace)) return undefined;
   for (const message of trace) {
@@ -241,6 +272,12 @@ async function defaultRunConnectorLocalRuntimeAgent(args: {
     ? args.parsed.payload.agentConfig
     : {};
   const policy = runtimePolicyFromConnectorPayload(args.parsed);
+  const wantsTaskWorktree = requestsTaskWorktree(policy);
+  const hostRuntimeConfig = resolveHostRuntimeConfigFromEnv({
+    env: args.runtimeEnv,
+    surface: "connector-local",
+    workspaceBinding: wantsTaskWorktree ? "default-cwd" : "explicit-cwd",
+  });
   const agentRecord = {
     ...payloadAgentConfig,
     dbKey: agentKey,
@@ -262,12 +299,14 @@ async function defaultRunConnectorLocalRuntimeAgent(args: {
     [`agent-${args.runtimeEnv.NOLO_USER_ID || "local"}-${agentKey}`, agentRecord],
     [`agent-${args.runtimeEnv.NOLO_LOCAL_USER_ID || args.runtimeEnv.NOLO_USER_ID || "local"}-${agentKey}`, agentRecord],
   ]);
-  const env = localRuntimeEnvFromPolicy(args.runtimeEnv, policy);
-  const wantsTaskWorktree = requestsTaskWorktree(policy);
+  const env = localRuntimeEnvFromPolicy(
+    withForwardedUserAuthToken(args.runtimeEnv, args.parsed),
+    policy
+  );
   const adapter = createCliLocalRuntimeAdapter({
     env,
+    hostRuntimeConfig,
     cwd: args.cwd,
-    useCwdAsTaskWorkspaceBase: wantsTaskWorktree,
     ...(wantsTaskWorktree
       ? {
           prepareTaskWorktree: (input) => prepareTaskWorktree({
@@ -350,6 +389,20 @@ function normalizeConnectorRunTimeoutMs(value: unknown): number | undefined {
   return Math.floor(value);
 }
 
+function normalizeConnectorRunPayloadContinueDialogId(payload: unknown) {
+  if (!isRecord(payload)) return payload;
+  const dialogLineage = normalizeAgentRunDialogLineage({
+    continueDialogId: payload.continueDialogId,
+  });
+  const normalizedPayload = { ...payload };
+  if (dialogLineage.continueDialogId) {
+    normalizedPayload.continueDialogId = dialogLineage.continueDialogId;
+  } else {
+    delete normalizedPayload.continueDialogId;
+  }
+  return normalizedPayload;
+}
+
 export async function handleConnectorRunMessage(
   message: string,
   send: (message: string) => void,
@@ -365,6 +418,12 @@ export async function handleConnectorRunMessage(
     return;
   }
   if (parsed?.type !== "agent.run" || typeof parsed.requestId !== "string") return;
+  if (isRecord(parsed)) {
+    parsed = {
+      ...parsed,
+      payload: normalizeConnectorRunPayloadContinueDialogId(parsed.payload),
+    };
+  }
   const agentConfig = parsed.payload?.agentConfig ?? {};
   const sendProgress = (progress: ConnectorRunProgress) => {
     send(JSON.stringify({
@@ -392,6 +451,11 @@ export async function handleConnectorRunMessage(
   try {
     const machinePermissionPolicy = resolveMachineRunPermissionPolicy(agentConfig);
     const runtimePolicy = runtimePolicyFromConnectorPayload(parsed);
+    const hostRuntimeConfig = resolveHostRuntimeConfigFromEnv({
+      env: runtimeEnv,
+      surface: agentConfig.apiSource === "cli" ? "connector-cli" : "connector-local",
+      workspaceBinding: requestsTaskWorktree(runtimePolicy) ? "default-cwd" : "explicit-cwd",
+    });
     const userInput = String(parsed.payload?.userInput ?? "");
     const timeout = normalizeConnectorRunTimeoutMs(parsed.payload?.timeoutMs);
     let cwd = resolveConnectorRunCwd({ env: runtimeEnv, policy: machinePermissionPolicy });
@@ -414,7 +478,7 @@ export async function handleConnectorRunMessage(
         artifactCwd = prepared.path;
       }
       const runtimePermissionPolicy = !hasExplicitMachinePermissions(agentConfig)
-        ? runtimeWorkspacePermissionPolicy(runtimePolicy, cwd)
+        ? runtimeWorkspacePermissionPolicy(runtimePolicy, hostRuntimeConfig, cwd)
         : null;
       const effectivePermissionPolicy = scopePermissionPolicyToRuntimeWorkspace(
         runtimePermissionPolicy ?? machinePermissionPolicy,
@@ -472,17 +536,18 @@ export async function handleConnectorRunMessage(
           updatedAt: Date.now(),
         });
       }, 15_000);
+      const effectiveRuntimeEnv = withForwardedUserAuthToken(runtimeEnv, parsed);
       const result = await executeCli(provider, materializedPrompt.prompt, {
         model: agentConfig.model || undefined,
         timeout,
         cwd,
         yolo: true,
         env: {
-          NOLO_SERVER: runtimeEnv.NOLO_SERVER || runtimeEnv.NOLO_SERVER_URL || runtimeEnv.BASE_URL,
-          NOLO_SERVER_URL: runtimeEnv.NOLO_SERVER_URL || runtimeEnv.NOLO_SERVER || runtimeEnv.BASE_URL,
-          BASE_URL: runtimeEnv.BASE_URL || runtimeEnv.NOLO_SERVER || runtimeEnv.NOLO_SERVER_URL,
-          AUTH_TOKEN: runtimeEnv.AUTH_TOKEN || runtimeEnv.AUTH || runtimeEnv.NOLO_MACHINE_API_KEY,
-          NOLO_MACHINE_API_KEY: runtimeEnv.NOLO_MACHINE_API_KEY || runtimeEnv.AUTH_TOKEN || runtimeEnv.AUTH,
+          NOLO_SERVER: effectiveRuntimeEnv.NOLO_SERVER || effectiveRuntimeEnv.NOLO_SERVER_URL || effectiveRuntimeEnv.BASE_URL,
+          NOLO_SERVER_URL: effectiveRuntimeEnv.NOLO_SERVER_URL || effectiveRuntimeEnv.NOLO_SERVER || effectiveRuntimeEnv.BASE_URL,
+          BASE_URL: effectiveRuntimeEnv.BASE_URL || effectiveRuntimeEnv.NOLO_SERVER || effectiveRuntimeEnv.NOLO_SERVER_URL,
+          AUTH_TOKEN: effectiveRuntimeEnv.AUTH_TOKEN || effectiveRuntimeEnv.AUTH || effectiveRuntimeEnv.NOLO_MACHINE_API_KEY,
+          NOLO_MACHINE_API_KEY: effectiveRuntimeEnv.NOLO_MACHINE_API_KEY || effectiveRuntimeEnv.AUTH_TOKEN || effectiveRuntimeEnv.AUTH,
         },
       }).finally(() => clearInterval(heartbeat));
       logConnectorRun("agent.run.cli.result", {
@@ -518,13 +583,13 @@ export async function handleConnectorRunMessage(
       }));
       return;
     }
-    if (!requestsLocalWorkspaceRuntime(runtimePolicy)) {
+    if (!requestsLocalWorkspaceRuntime(runtimePolicy) && !isMachineBoundLocalCustomProvider(agentConfig)) {
       throw new Error("Connector can only execute non-CLI agents when runtimeToolPolicySnapshot requests a local workspace runtime.");
     }
     baseSha = await readConnectorGitHead(cwd);
     const result = await runConnectorLocalRuntimeAgent({
       parsed,
-      runtimeEnv,
+      runtimeEnv: withForwardedUserAuthToken(runtimeEnv, parsed),
       cwd,
       fetchImpl,
       onProgress: sendProgress,

package/machineWsSession.ts

@@ -1,6 +1,9 @@
 import type { HeartbeatLoopOptions } from "./connector-experimental/heartbeatLoop";
 import type { MachineHeartbeat } from "./connector-experimental/protocol";
-import { isConnectorWebSocketAuthError } from "./connectorWebSocketTarget";
+import {
+  isConnectorWebSocketAuthError,
+  isConnectorWebSocketRetryableError,
+} from "./connectorWebSocketTarget";
 
 type EnvLike = Record<string, string | undefined>;
 type OutputLike = { write(chunk: string): unknown };
@@ -9,6 +12,13 @@ export type ConnectorWebSocketOptions = {
   headers: Record<string, string>;
   onMessage: (message: string) => void | Promise<void>;
   sentMessages: string[];
+  signal?: AbortSignal;
+};
+
+export type RunMachineWsSessionResult = {
+  exitCode: number;
+  retryAfterMs?: number;
+  reconnectReason?: string;
 };
 
 const CONNECTOR_WS_KEEPALIVE_MS = 25_000;
@@ -35,6 +45,7 @@ export async function defaultConnectWebSocket(
       headers: options.headers,
     } as any);
     let opened = false;
+    let aborted = options.signal?.aborted ?? false;
     let keepalive: ReturnType<typeof setInterval> | null = null;
     const clearKeepalive = () => {
       if (keepalive) {
@@ -42,6 +53,16 @@ export async function defaultConnectWebSocket(
         keepalive = null;
       }
     };
+    const handleAbort = () => {
+      aborted = true;
+      clearKeepalive();
+      try {
+        ws.close();
+      } catch {
+        resolve();
+      }
+    };
+    options.signal?.addEventListener("abort", handleAbort, { once: true });
     ws.addEventListener(
       "open",
       () => {
@@ -58,13 +79,17 @@ export async function defaultConnectWebSocket(
     );
     ws.addEventListener("error", () => {
       clearKeepalive();
+      if (aborted) {
+        resolve();
+        return;
+      }
       reject(new Error("connector websocket failed"));
     });
     ws.addEventListener(
       "close",
       () => {
         clearKeepalive();
-        if (opened) resolve();
+        if (opened || aborted) resolve();
         else reject(new Error("connector websocket closed before opening"));
       },
       { once: true },
@@ -88,21 +113,26 @@ export async function runMachineWsSession(options: {
   machine: MachineHeartbeat;
   serverUrl: string;
   authToken: string;
+  signal?: AbortSignal;
   sendHeartbeat: () => Promise<void>;
   runHeartbeatLoop: (options: HeartbeatLoopOptions) => Promise<void>;
   resolveConnectorWebSocketTarget: (args: {
     serverUrl: string;
     machineId: string;
+    connectorSurface?: string;
     headers: Record<string, string>;
   }) => Promise<string>;
   connectWebSocket: (url: string, options: ConnectorWebSocketOptions) => Promise<void>;
   onMessage: (message: string, send: (message: string) => void) => void | Promise<void>;
-}) {
+}): Promise<RunMachineWsSessionResult> {
   const sentMessages: string[] = [];
   const heartbeatAbort = new AbortController();
   let heartbeatLoopPromise: Promise<void> | null = null;
   let websocketPromise: Promise<void> | null = null;
   try {
+    if (options.signal?.aborted) {
+      return { exitCode: 0 };
+    }
     await options.sendHeartbeat();
     options.output.write(
       `Connector websocket connected: ${options.machine.name} (${options.machine.machineId})\n`
@@ -120,20 +150,36 @@ export async function runMachineWsSession(options: {
     const wsTarget = await options.resolveConnectorWebSocketTarget({
       serverUrl: options.serverUrl,
       machineId: options.machine.machineId,
+      connectorSurface: "cli",
       headers: { Authorization: `Bearer ${options.authToken}` },
     });
     websocketPromise = options.connectWebSocket(wsTarget, {
       headers: { Authorization: `Bearer ${options.authToken}` },
+      signal: options.signal,
       sentMessages,
       onMessage: (message) =>
         options.onMessage(message, (response) => {
           sentMessages.push(response);
         }),
     });
-    await Promise.race([websocketPromise, heartbeatMonitorPromise]);
+    const abortPromise = options.signal
+      ? new Promise<"aborted">((resolve) => {
+          if (options.signal?.aborted) {
+            resolve("aborted");
+            return;
+          }
+          options.signal?.addEventListener("abort", () => resolve("aborted"), { once: true });
+        })
+      : null;
+    const raceResult = await Promise.race(
+      [websocketPromise, heartbeatMonitorPromise, abortPromise].filter(Boolean) as Promise<void | "aborted">[]
+    );
     heartbeatAbort.abort();
     await Promise.allSettled([websocketPromise, heartbeatLoopPromise]);
-    return 0;
+    if (raceResult === "aborted") {
+      return { exitCode: 0 };
+    }
+    return { exitCode: 0 };
   } catch (error) {
     heartbeatAbort.abort();
     await Promise.allSettled(
@@ -143,13 +189,25 @@ export async function runMachineWsSession(options: {
       options.output.write(
         `[nolo] Connector websocket auth failed: ${error.message}. Re-register this machine or update NOLO_MACHINE_API_KEY.\n`
       );
-      return 2;
+      return { exitCode: 2 };
+    }
+    if (isConnectorWebSocketRetryableError(error)) {
+      const description =
+        error.reason === "core_draining" ? "core draining" : error.reason;
+      options.output.write(
+        `[nolo] Connector websocket unavailable: ${description} (${error.message}).\n`
+      );
+      return {
+        exitCode: 1,
+        retryAfterMs: error.retryAfterMs,
+        reconnectReason: error.reason,
+      };
     }
     options.output.write(
       `[nolo] Connector websocket failed: ${
         error instanceof Error ? error.message : String(error)
       }\n`
     );
-    return 1;
+    return { exitCode: 1 };
   }
 }