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

package/dist/socket-manager.js

@@ -6,7 +6,6 @@ import { createISocketClient, DEFAULT_NON_GIT_BRANCH, MAX_SOCKET_PAYLOAD_BYTES,
 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";
@@ -16,7 +15,6 @@ 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>`
@@ -129,28 +127,7 @@ 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();
     /**
@@ -181,73 +158,6 @@ 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) {
@@ -415,54 +325,11 @@ 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);
                 }
-                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
                 // this signal to drain state pushes that were deferred at
@@ -495,16 +362,7 @@ 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.
-                    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;
-                            }
-                        }
-                    }
+                    aiService?.claimHeadlessTurn(socket, peerId, readHeadlessTurnId(params));
                 }
                 return await next();
             };
@@ -552,33 +410,12 @@ class SocketManager extends TracedEventEmitter {
                             if (kind !== "gateway") {
                                 return await aiService.handleAiGenerateWithQueue(aiRequest, queueMode);
                             }
-                            // 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"
+                            const turnId = aiService.beginHeadlessTurn(socket, peerId, typeof request.requestId === "string"
                                 ? request.requestId
-                                : 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);
-                            });
+                                : undefined);
                             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
@@ -586,12 +423,6 @@ 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
@@ -601,16 +432,6 @@ 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) {
@@ -1199,20 +1020,17 @@ 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) {
-                                // Retain the pre-switch state for the explicit fallback
-                                // if an atomic branch transfer fails.
-                                const wasLockedBeforeDisconnect = await lockService.isLockedWhenSettled();
+                                const wasLockedBeforeDisconnect = lockService.isLocked;
                                 try {
-                                    await lockService.switchBranch(DEFAULT_NON_GIT_BRANCH);
+                                    await lockService.switchBranch(DEFAULT_NON_GIT_BRANCH, {
+                                        reacquireLock: wasLockedBeforeDisconnect,
+                                    });
                                 }
                                 catch (err) {
                                     getLogger().warn("[git] gitDisconnect failed to switch lock branch back to main", {
@@ -1464,9 +1282,6 @@ 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);
@@ -1479,14 +1294,23 @@ 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, connectionLiveness);
+                this.emit("connect:gateway", socket, peerId);
             }
             else if (isEditor) {
                 this.editorSockets.set(socket, peerId);
-                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`.
+                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)}`);
+                    });
+                }
             }
             else {
                 this.librarySockets.add(socket);
@@ -1494,73 +1318,61 @@ 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);
-                    // 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);
+                    aiService?.endHeadlessTurn(peerId);
+                    this.emit("disconnect:gateway", peerId);
                 }
                 else if (isEditor) {
                     this.editorSockets.delete(socket);
                     const hasOtherEditorSocket = Array.from(peerSockets ?? []).some((peerSocket) => this.editorSockets.has(peerSocket));
-                    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}`);
-                                    });
-                                }
+                    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}`);
+                                });
                             }
                         }
                     }
-                    this.emit("disconnect:editor", peerId, socket, connectionLiveness);
+                    this.emit("disconnect:editor", peerId);
                 }
                 else {
                     this.librarySockets.delete(socket);
@@ -1592,48 +1404,19 @@ class SocketManager extends TracedEventEmitter {
         }
     }
     /**
-     * Any editor socket that does not belong to `exclude`.
+     * Any editor socket that does not belong to `excludePeerId`.
      *
      * 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(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;
+    getSurvivingEditorClient(excludePeerId) {
+        for (const [peer, peerId] of this.editorSockets) {
+            if (peerId !== excludePeerId) {
+                return { peer, peerId };
             }
         }
-        return pick(false);
+        return undefined;
     }
     callLibraryClients(fn) {
         for (const socket of this.librarySockets) {