@superblocksteam/vite-plugin-file-sync 2.0.156 → 2.0.157-next.1

package/dist/socket-manager.js

@@ -2,10 +2,11 @@ import { cp, mkdir, mkdtemp, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { WebSocketServer } from "ws";
-import { createISocketClient, DEFAULT_NON_GIT_BRANCH, MAX_SOCKET_PAYLOAD_BYTES, SUPERBLOCKS_LIVE_GIT_BRANCH, TracedEventEmitter, TracedSocket, } from "@superblocksteam/shared";
+import { createISocketClient, DEFAULT_NON_GIT_BRANCH, MAX_SOCKET_PAYLOAD_BYTES, SocketErrorException, SUPERBLOCKS_LIVE_GIT_BRANCH, TracedEventEmitter, TracedSocket, } from "@superblocksteam/shared";
 import { uploadContextToBucketeer } from "./ai-service/context-upload.js";
 import { createSocketRecordingMiddleware } from "./ai-service/recording/recorders/socket-recorder.js";
 import { parseJwt } from "./ai-service/util/parse-jwt.js";
+import { PeerLiveness, } from "./ai-service/util/peer-liveness.js";
 import { clearGitAuthFailed, isGitAuthFailed, markGitAuthFailed, } from "./git-auth-state.js";
 import { GIT_LOG_FORMAT, parseGitLog } from "./git-log.js";
 import { GIT_OPERATION_BUSY_MESSAGE, } from "./git-operation-lock.js";
@@ -15,6 +16,7 @@ import { doesRemoteLiveBranchExist, ensureLiveBranchCheckedOut, pushLiveBranchWi
 import { ensureProviderWorkflowFiles } from "./git-service/provider-workflow-files.js";
 import { getSyncCountsAgainstDefaultBranch } from "./git-service/status-sync.js";
 import { getErrorMeta, getLogger } from "./util/logger.js";
+import { OperationQueue } from "./util/operation-queue.js";
 const SB_WORKTREE_SNAPSHOTS_DIR = join(tmpdir(), "sb-worktree-snapshots");
 /**
  * Pulls the JWT-derived `userId` out of an `Authorization: Bearer <jwt>`
@@ -127,7 +129,28 @@ export function disposeSocketManager() {
     _socketManager = undefined;
 }
 class SocketManager extends TracedEventEmitter {
+    peerLiveness = new PeerLiveness();
+    socketLiveness = new WeakMap();
+    /**
+     * Survives `editorSockets.delete` so a just-closed socket can still tell
+     * {@link getSurvivingEditorClient} which peerId to prefer among leftovers.
+     */
+    editorPeerIdsBySocket = new WeakMap();
     peerIds = new Map();
+    /**
+     * Editor peers whose presence is registered with the lock service, and the
+     * user each was registered for.
+     *
+     * Presence is per *peer*, not per socket: one tab reconnecting produces a new
+     * socket under the same peer id. Remembering "did I register?" on the socket
+     * therefore breaks the pair, because the socket that takes presence down need
+     * not be the one that put it up — a replacement that never makes an
+     * authorized call would close without removing anything, leaving the lock
+     * claimed for a peer that is gone.
+     */
+    editorPresenceUserIdByPeerId = new Map();
+    /** Editor peers that have already triggered post-auth git activation. */
+    gitActivatedPeerIds = new Set();
     /** Attached editor sockets, each mapped to the peer that opened it. */
     editorSockets = new Map();
     /**
@@ -158,6 +181,73 @@ class SocketManager extends TracedEventEmitter {
         });
         this.socketServer = wss;
         let activeGitService = gitService;
+        const headlessLockOwners = [];
+        const headlessLockOperationQueue = new OperationQueue();
+        /**
+         * Drop every turn this peer owned, and hand the lock back if that was the
+         * last of them.
+         *
+         * Dropping the rows alone leaves the lock stranded. The lock service only
+         * hands back when *no* gateway is left, so with a replacement socket
+         * already attached that never fires — and once the owner row is gone, a
+         * later `editor.releaseTurnLock` finds nothing to release and does
+         * nothing. The AppBranchLock would then stay held for the rest of that
+         * socket's life, with no turn behind it.
+         *
+         * Queued with the release path so the two cannot both decide they are
+         * last.
+         */
+        const dropHeadlessLockOwners = async (peerId) => {
+            await headlessLockOperationQueue.enqueue(async () => {
+                const ownedAny = headlessLockOwners.some((owner) => owner.peerId === peerId);
+                if (!ownedAny) {
+                    return;
+                }
+                for (let index = headlessLockOwners.length - 1; index >= 0; index--) {
+                    if (headlessLockOwners[index]?.peerId === peerId) {
+                        headlessLockOwners.splice(index, 1);
+                    }
+                }
+                if (headlessLockOwners.length > 0) {
+                    return;
+                }
+                // Re-read after the handback's generation wait: a replacement gateway
+                // can claim a turn while the old one's write is still finishing.
+                await lockService?.relinquishLockForTurn(() => headlessLockOwners.length > 0);
+            });
+        };
+        const dropHeadlessLockOwnersInBackground = (peerId, occasion) => {
+            void dropHeadlessLockOwners(peerId).catch((err) => {
+                getLogger().warn(`[socket-manager] failed to hand the turn lock back after ${occasion}: ${err instanceof Error ? err.message : String(err)}`);
+            });
+        };
+        lockService?.setOnGatewayLivenessLapsed?.((peerId) => {
+            dropHeadlessLockOwnersInBackground(peerId, "a gateway's liveness lapsed");
+            aiService?.endHeadlessTurn(peerId);
+        });
+        const releaseHeadlessLock = async (peerId, turnId) => {
+            return await headlessLockOperationQueue.enqueue(async () => {
+                const ownerIndex = headlessLockOwners.findIndex((owner) => owner.peerId === peerId &&
+                    (turnId ? owner.turnId === turnId : owner.turnId === undefined));
+                if (ownerIndex === -1) {
+                    return false;
+                }
+                const [owner] = headlessLockOwners.splice(ownerIndex, 1);
+                if (headlessLockOwners.length > 0) {
+                    return true;
+                }
+                try {
+                    await lockService?.relinquishLockForTurn(() => headlessLockOwners.length > 0);
+                    return true;
+                }
+                catch (error) {
+                    if (owner) {
+                        headlessLockOwners.splice(ownerIndex, 0, owner);
+                    }
+                    throw error;
+                }
+            });
+        };
         // Provide sync context (branchName, lastCommitId) to the AI service
         // for remote context restore fallback from Bucketeer on fresh pods.
         if (aiService && syncService) {
@@ -325,10 +415,53 @@ class SocketManager extends TracedEventEmitter {
         };
         wss.on("connection", (ws, kind, peerId, userId) => {
             const isEditor = kind === "editor";
+            const connectionLiveness = kind === "library" ? undefined : this.peerLiveness.bindConnection();
+            if (connectionLiveness) {
+                ws.on("message", () => connectionLiveness.markInbound());
+            }
+            /**
+             * When an editor authenticates after git was enabled, the lock may
+             * still be on 'main'. This calls
+             * `ensureRuntimeDbfsBranchConsistency` -> `lockService.switchBranch`
+             * to move it to superblocks/live.
+             */
+            const activateGitForAuthedEditor = () => {
+                if (!activateGitService) {
+                    return;
+                }
+                void activateGitService().catch((err) => {
+                    getLogger().warn(`[socket-manager] post-auth git activation failed: ${err instanceof Error ? err.message : String(err)}`);
+                });
+            };
             const authMiddleware = async (params, ctx, _peer, next) => {
                 const result = await checkAuthorization(ctx.peerAuthorization);
                 if (!result.success) {
-                    throw new Error(result.error);
+                    throw new SocketErrorException(result.status, result.error);
+                }
+                if (isEditor &&
+                    userId &&
+                    !this.editorPresenceUserIdByPeerId.has(peerId) &&
+                    // The authorization above is awaited, so the socket that asked for
+                    // it can be gone by now. `peerIds` loses the peer when its last
+                    // socket closes, so this is what keeps a late answer from
+                    // registering presence nothing is left to remove.
+                    this.peerIds.has(peerId)) {
+                    this.editorPresenceUserIdByPeerId.set(peerId, userId);
+                    lockService?.addConnectedPeerIdToUser(userId, peerId);
+                }
+                // After presence, and only for a peer that has proved itself.
+                //
+                // Git activation moves the DBFS branch, which asks the lock service
+                // whether this pod holds the lock. Presence registration is what
+                // starts the attach-time acquisition, so running this first — as it
+                // did from the connect handler — asks that question before there is
+                // an acquisition to observe: `isLockedWhenSettled()` settles what is
+                // in flight, and nothing is. The switch then answers "no lock",
+                // releases the one that lands moments later, and does not take it
+                // back. Ordering them here is what keeps the two consistent.
+                if (isEditor && !this.gitActivatedPeerIds.has(peerId)) {
+                    this.gitActivatedPeerIds.add(peerId);
+                    activateGitForAuthedEditor();
                 }
                 // Notify AiService that this peer has just produced a
                 // proof-of-auth. AiService deduplicates internally and uses
@@ -362,7 +495,16 @@ class SocketManager extends TracedEventEmitter {
                     // it follows the peer rather than reporting into the socket it
                     // lost. Naming it is required — otherwise any authorized gateway
                     // session could take over another one's turn.
-                    aiService?.claimHeadlessTurn(socket, peerId, readHeadlessTurnId(params));
+                    const headlessTurnId = readHeadlessTurnId(params);
+                    if (ctx.method !== "editor.releaseTurnLock") {
+                        aiService?.claimHeadlessTurn(socket, peerId, headlessTurnId, connectionLiveness);
+                        if (headlessTurnId) {
+                            const owner = headlessLockOwners.find((candidate) => candidate.turnId === headlessTurnId);
+                            if (owner) {
+                                owner.peerId = peerId;
+                            }
+                        }
+                    }
                 }
                 return await next();
             };
@@ -410,12 +552,33 @@ class SocketManager extends TracedEventEmitter {
                             if (kind !== "gateway") {
                                 return await aiService.handleAiGenerateWithQueue(aiRequest, queueMode);
                             }
-                            const turnId = aiService.beginHeadlessTurn(socket, peerId, typeof request.requestId === "string"
+                            // Write access is taken per turn, not held for the whole
+                            // session, so the app is unlocked between prompts and a
+                            // person can open it. Before the turn is bound, and
+                            // awaited: a turn that cannot write must fail here carrying
+                            // the server's own reason — somebody is editing in the
+                            // browser — rather than bind routing and then unwind it.
+                            const requestedTurnId = typeof request.requestId === "string"
                                 ? request.requestId
-                                : undefined);
+                                : undefined;
+                            const lockOwner = {
+                                peerId,
+                                turnId: requestedTurnId,
+                            };
+                            const turnId = await headlessLockOperationQueue.enqueue(async () => {
+                                await lockService?.acquireLockForTurn();
+                                headlessLockOwners.push(lockOwner);
+                                return aiService.beginHeadlessTurn(socket, peerId, requestedTurnId, connectionLiveness);
+                            });
                             let result;
                             try {
                                 result = await aiService.handleAiGenerateWithQueue(aiRequest, queueMode);
+                                const ownerIndex = headlessLockOwners.indexOf(lockOwner);
+                                for (let index = ownerIndex - 1; index >= 0; index--) {
+                                    if (headlessLockOwners[index]?.peerId === peerId) {
+                                        headlessLockOwners.splice(index, 1);
+                                    }
+                                }
                             }
                             catch (err) {
                                 // The turn is bound before the prompt is handed over, so
@@ -423,6 +586,12 @@ class SocketManager extends TracedEventEmitter {
                                 // it. Left bound, it would capture the next turn a user
                                 // starts in the browser.
                                 aiService.abandonHeadlessTurn(peerId, turnId);
+                                try {
+                                    await releaseHeadlessLock(peerId, turnId);
+                                }
+                                catch (releaseErr) {
+                                    getLogger().warn("[socket-manager] failed to release the lock for a rejected headless turn", getErrorMeta(releaseErr));
+                                }
                                 throw err;
                             }
                             // Answering with the turn handle tells the caller two
@@ -432,6 +601,16 @@ class SocketManager extends TracedEventEmitter {
                             return { ...result, headlessTurn: { turnId } };
                         },
                     ],
+                    releaseTurnLock: [
+                        async (request) => {
+                            if (kind !== "gateway") {
+                                throw new Error("Only gateway peers can release a headless turn lock");
+                            }
+                            if (await releaseHeadlessLock(peerId, request.headlessTurnId)) {
+                                aiService?.endHeadlessTurn(peerId, request.headlessTurnId);
+                            }
+                        },
+                    ],
                     aiTriggerMigrationTurn: [
                         async (request) => {
                             if (!aiService) {
@@ -1020,17 +1199,20 @@ class SocketManager extends TracedEventEmitter {
                             activeGitService = undefined;
                             disconnectInProgress = true;
                             clearGitAuthFailed();
+                            // Keep file-sync / lock wiring; drop only the Git
+                            // service so the next Clark turn sees unconnected.
+                            aiService?.setRuntimeGitServices(undefined, syncService, gitOperationLock);
                             // Ensure DBFS sync context is no longer tied to
                             // `superblocks/live` immediately after native git disconnect.
                             syncService?.setBranchName(DEFAULT_NON_GIT_BRANCH);
                             // Native git disconnect should move runtime lock context back
                             // to default branch before clients reconnect.
                             if (lockService) {
-                                const wasLockedBeforeDisconnect = lockService.isLocked;
+                                // Retain the pre-switch state for the explicit fallback
+                                // if an atomic branch transfer fails.
+                                const wasLockedBeforeDisconnect = await lockService.isLockedWhenSettled();
                                 try {
-                                    await lockService.switchBranch(DEFAULT_NON_GIT_BRANCH, {
-                                        reacquireLock: wasLockedBeforeDisconnect,
-                                    });
+                                    await lockService.switchBranch(DEFAULT_NON_GIT_BRANCH);
                                 }
                                 catch (err) {
                                     getLogger().warn("[git] gitDisconnect failed to switch lock branch back to main", {
@@ -1282,6 +1464,9 @@ class SocketManager extends TracedEventEmitter {
                     },
                 },
             }));
+            if (connectionLiveness) {
+                this.socketLiveness.set(socket, connectionLiveness);
+            }
             // we are storing all of the sockets for a given peerId (editor+library pair of connections)
             // which lets us handle synchronization of operations between the two
             let existingSockets = this.peerIds.get(peerId);
@@ -1294,23 +1479,14 @@ class SocketManager extends TracedEventEmitter {
                 // Liveness is granted in the auth middleware, not here: the upgrade
                 // path only checks `applicationId`, and holding the sandbox open is
                 // not something an unproven socket gets to do.
-                this.emit("connect:gateway", socket, peerId);
+                this.emit("connect:gateway", socket, peerId, connectionLiveness);
             }
             else if (isEditor) {
                 this.editorSockets.set(socket, peerId);
-                this.emit("connect:editor", socket, peerId);
-                if (userId) {
-                    lockService?.addConnectedPeerIdToUser(userId, peerId);
-                }
-                // When a new editor connects after git was enabled, the lock may
-                // still be on 'main'. Trigger git activation which calls
-                // ensureRuntimeDbfsBranchConsistency → lockService.switchBranch
-                // to transfer the lock to superblocks/live.
-                if (activateGitService) {
-                    void activateGitService().catch((err) => {
-                        getLogger().warn(`[socket-manager] post-connect git activation failed: ${err instanceof Error ? err.message : String(err)}`);
-                    });
-                }
+                this.editorPeerIdsBySocket.set(socket, peerId);
+                this.emit("connect:editor", socket, peerId, connectionLiveness);
+                // Git activation is deliberately NOT started here — it runs from the
+                // auth middleware, after presence. See `activateGitForAuthedEditor`.
             }
             else {
                 this.librarySockets.add(socket);
@@ -1318,61 +1494,73 @@ class SocketManager extends TracedEventEmitter {
             }
             existingSockets.add(socket);
             ws.on("close", (code, reason) => {
+                connectionLiveness?.close();
                 const peerSockets = this.peerIds.get(peerId);
                 peerSockets?.delete(socket);
                 if (peerSockets?.size === 0) {
                     this.peerIds.delete(peerId);
+                    // The peer is gone, so a later reconnect under the same id is a
+                    // new arrival and activates git again.
+                    this.gitActivatedPeerIds.delete(peerId);
                 }
                 getLogger().info(`Socket from ${kind} closed code: ${code} reason: ${reason}`);
                 if (kind === "gateway") {
                     this.gatewaySockets.delete(socket);
                     lockService?.markGatewayDetached(peerId);
-                    aiService?.endHeadlessTurn(peerId);
-                    this.emit("disconnect:gateway", peerId);
+                    // No blanket clear when the last socket goes: every peer's rows are
+                    // dropped by its own close, and the one that empties the list is
+                    // what hands the lock back.
+                    dropHeadlessLockOwnersInBackground(peerId, "a gateway socket closed");
+                    aiService?.endHeadlessTurn(peerId, connectionLiveness);
+                    this.emit("disconnect:gateway", peerId, connectionLiveness);
                 }
                 else if (isEditor) {
                     this.editorSockets.delete(socket);
                     const hasOtherEditorSocket = Array.from(peerSockets ?? []).some((peerSocket) => this.editorSockets.has(peerSocket));
-                    if (hasOtherEditorSocket) {
-                        return;
-                    }
-                    if (userId) {
-                        lockService?.removeConnectedPeerIdFromUser(userId, peerId);
-                        if (lockService?.connectedUsers.length === 0) {
-                            if (aiService?.isBusy()) {
-                                getLogger().info("[socket-manager] All users disconnected but AI service is busy, deferring cache cleanup until generation completes");
-                                let cleaned = false;
-                                const deferredCleanup = () => {
-                                    if (cleaned)
-                                        return;
-                                    cleaned = true;
-                                    // Only clean up if no users have reconnected in the meantime
-                                    if (lockService?.connectedUsers.length === 0) {
-                                        getLogger().info("[socket-manager] Generation ended with no connected users, cleaning up AI service cache");
-                                        aiService
-                                            .removeIntegrationCache()
-                                            .catch((error) => {
-                                            getLogger().warn(`Error removing integration cache after deferred cleanup: ${error}`);
-                                        });
-                                    }
-                                    else {
-                                        getLogger().info("[socket-manager] Generation ended but users have reconnected, skipping cache cleanup");
-                                    }
-                                };
-                                aiService.once("generationCompleted", deferredCleanup);
-                                aiService.once("generationFailed", deferredCleanup);
-                            }
-                            else {
-                                getLogger().info("[socket-manager] All users disconnected, cleaning up AI service cache");
-                                aiService
-                                    ?.removeIntegrationCache()
-                                    .catch((error) => {
-                                    getLogger().warn(`Error removing integration cache after disconnect: ${error}`);
-                                });
+                    if (!hasOtherEditorSocket) {
+                        // Keyed by peer, so whichever of the peer's sockets is last out
+                        // takes the presence down — including one that never authenticated
+                        // itself.
+                        const presenceUserId = this.editorPresenceUserIdByPeerId.get(peerId);
+                        if (presenceUserId) {
+                            this.editorPresenceUserIdByPeerId.delete(peerId);
+                            lockService?.removeConnectedPeerIdFromUser(presenceUserId, peerId);
+                            if (lockService?.connectedUsers.length === 0) {
+                                if (aiService?.isBusy()) {
+                                    getLogger().info("[socket-manager] All users disconnected but AI service is busy, deferring cache cleanup until generation completes");
+                                    let cleaned = false;
+                                    const deferredCleanup = () => {
+                                        if (cleaned)
+                                            return;
+                                        cleaned = true;
+                                        // Only clean up if no users have reconnected in the meantime
+                                        if (lockService?.connectedUsers.length === 0) {
+                                            getLogger().info("[socket-manager] Generation ended with no connected users, cleaning up AI service cache");
+                                            aiService
+                                                .removeIntegrationCache()
+                                                .catch((error) => {
+                                                getLogger().warn(`Error removing integration cache after deferred cleanup: ${error}`);
+                                            });
+                                        }
+                                        else {
+                                            getLogger().info("[socket-manager] Generation ended but users have reconnected, skipping cache cleanup");
+                                        }
+                                    };
+                                    aiService.once("generationCompleted", deferredCleanup);
+                                    aiService.once("generationFailed", deferredCleanup);
+                                }
+                                else {
+                                    getLogger().info("[socket-manager] All users disconnected, cleaning up AI service cache");
+                                    aiService
+                                        ?.removeIntegrationCache()
+                                        .catch((error) => {
+                                        getLogger().warn(`Error removing integration cache after disconnect: ${error}`);
+                                    });
+                                }
                             }
                         }
                     }
-                    this.emit("disconnect:editor", peerId);
+                    this.emit("disconnect:editor", peerId, socket, connectionLiveness);
                 }
                 else {
                     this.librarySockets.delete(socket);
@@ -1404,19 +1592,48 @@ class SocketManager extends TracedEventEmitter {
         }
     }
     /**
-     * Any editor socket that does not belong to `excludePeerId`.
+     * Any editor socket that does not belong to `exclude`.
      *
      * Used when a peer disconnects: Clark's RPCs are bound to a single editor,
      * so a survivor is what keeps API execution alive for a user whose tab is
      * still open after a headless build's socket closes.
+     *
+     * Same-id sockets are preferred when `exclude` is the departing socket: a
+     * tab that reconnected keeps the bound Clark context instead of handing it
+     * (and in-flight generation / tool-permission prompts) to whichever other
+     * editor happened to connect first. A string `exclude` still means "this
+     * peer is gone" and skips every socket with that id.
      */
-    getSurvivingEditorClient(excludePeerId) {
-        for (const [peer, peerId] of this.editorSockets) {
-            if (peerId !== excludePeerId) {
-                return { peer, peerId };
+    getSurvivingEditorClient(exclude) {
+        const excludedPeerId = typeof exclude === "string" ? exclude : undefined;
+        const preferredPeerId = typeof exclude === "string"
+            ? undefined
+            : this.editorPeerIdsBySocket.get(exclude);
+        const pick = (onlyPreferred) => {
+            for (const [peer, peerId] of this.editorSockets) {
+                const excluded = excludedPeerId !== undefined
+                    ? peerId === excludedPeerId
+                    : peer === exclude;
+                if (excluded) {
+                    continue;
+                }
+                if (onlyPreferred && peerId !== preferredPeerId) {
+                    continue;
+                }
+                const liveness = this.socketLiveness.get(peer);
+                if (liveness) {
+                    return { liveness, peer, peerId };
+                }
+            }
+            return undefined;
+        };
+        if (preferredPeerId !== undefined) {
+            const sameId = pick(true);
+            if (sameId) {
+                return sameId;
             }
         }
-        return undefined;
+        return pick(false);
     }
     callLibraryClients(fn) {
         for (const socket of this.librarySockets) {