@superblocksteam/vite-plugin-file-sync 2.0.167-next.0 → 2.0.168-next.0

package/dist/socket-manager.js

@@ -1,9 +1,9 @@
-import { randomUUID } from "node:crypto";
+import { randomUUID, timingSafeEqual } from "node:crypto";
 import { cp, mkdir, mkdtemp, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { metrics } from "@opentelemetry/api";
-import { WebSocketServer } from "ws";
+import { WebSocket, WebSocketServer } from "ws";
 import { createISocketClient, DEFAULT_NON_GIT_BRANCH, MAX_SOCKET_PAYLOAD_BYTES, SocketErrorException, SUPERBLOCKS_LIVE_GIT_BRANCH, TracedEventEmitter, TracedSocket, } from "@superblocksteam/shared";
 import { uploadCheckpointScreenshot } from "./ai-service/checkpoint-screenshot.js";
 import { uploadContextToBucketeer } from "./ai-service/context-upload.js";
@@ -18,6 +18,7 @@ import { GitConflictError, isMissingGitRevisionError, } from "./git-service/inde
 import { doesRemoteLiveBranchExist, ensureLiveBranchCheckedOut, pushLiveBranchWithAutoRecreate, rollbackLastCommit, } from "./git-service/live-branch.js";
 import { ensureProviderWorkflowFiles } from "./git-service/provider-workflow-files.js";
 import { getSyncCountsAgainstDefaultBranch } from "./git-service/status-sync.js";
+import { SplitEditAfterClark } from "./split-edit-after-clark.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");
@@ -42,6 +43,55 @@ function extractCallerUserId(peerAuthorization) {
 }
 const EDITOR_PROTOCOL = "editor";
 const GATEWAY_PROTOCOL = "gateway";
+const AI_SERVICE_PROTOCOL = "ai-service";
+const INTERNAL_METHODS = new Set([
+    "editor.internalApplyDraftChanges",
+    "editor.internalCreateDraftMarkerOnly",
+    "editor.internalHasDraftChanges",
+    "editor.internalGetBranchName",
+    "editor.internalUploadDirectory",
+    "editor.internalSyncCurrentDirectoryWithDraftsNow",
+    "editor.internalGetDraftDiff",
+    "editor.internalGetDraftChanges",
+    "editor.internalAcquireTurnLock",
+    "editor.internalReleaseTurnLock",
+]);
+const AI_SERVICE_METHODS = new Set([
+    ...INTERNAL_METHODS,
+    "editor.aiAcceptDraft",
+    "editor.aiRejectDraft",
+    "editor.aiCancel",
+    "editor.aiClearDraftState",
+]);
+function decodeDraftBinary(value) {
+    const content = Buffer.from(value, "base64");
+    if (content.toString("base64") !== value) {
+        throw new SocketErrorException(400, "Invalid draft binary content");
+    }
+    return content;
+}
+function uploadAcceptedDraftContext(aiService, applicationId, commitId, branchName, peerAuthorization) {
+    void (async () => {
+        const contextId = aiService.getContextId();
+        if (!contextId)
+            return;
+        const jwt = aiService.getClarkJwt() || peerAuthorization?.replace("Bearer ", "");
+        if (!jwt)
+            return;
+        await uploadContextToBucketeer({
+            applicationId,
+            commitId,
+            branchName,
+            appRootDirPath: aiService.appRootDirPath,
+            useServerUpload: aiService.isAiFileUploadEnabled,
+            superblocksBaseUrl: aiService.superblocksBaseUrl,
+            logger: getLogger(),
+            jwt,
+        });
+    })().catch((error) => {
+        getLogger().warn("[aiAcceptDraft] Failed to upload context to bucketeer (non-blocking)", getErrorMeta(error));
+    });
+}
 /**
  * Subprotocols the dev server answers with. The library runtime offers none at
  * all and connects on that basis; there is no `library` protocol to negotiate.
@@ -52,6 +102,7 @@ const GATEWAY_PROTOCOL = "gateway";
 const NEGOTIABLE_PROTOCOLS = [
     EDITOR_PROTOCOL,
     GATEWAY_PROTOCOL,
+    AI_SERVICE_PROTOCOL,
 ];
 function selectProtocol(offered) {
     for (const protocol of offered) {
@@ -81,6 +132,9 @@ export function classifyConnectionKind(offeredProtocols) {
     if (selected === EDITOR_PROTOCOL) {
         return "editor";
     }
+    if (selected === AI_SERVICE_PROTOCOL) {
+        return "ai-service";
+    }
     return "unsupported";
 }
 /**
@@ -132,6 +186,9 @@ export function disposeSocketManager() {
     _socketManager?.dispose();
     _socketManager = undefined;
 }
+export async function closeAllSocketClients() {
+    await _socketManager?.closeAllClientsGracefully();
+}
 class SocketManager extends TracedEventEmitter {
     peerLiveness = new PeerLiveness();
     socketLiveness = new WeakMap();
@@ -174,10 +231,13 @@ class SocketManager extends TracedEventEmitter {
         super(tracer);
     }
     socketServer;
-    startSocketServer({ server, fileSyncManager, syncService, lockService, aiService, gitService, gitBootstrapAuthFailed, gitBootstrapFailed, activateGitService, refreshGitCredentials, checkAuthorization, path, tracer, gitOperationLock, superblocksBaseUrl, maxPayloadBytes = MAX_SOCKET_PAYLOAD_BYTES, }) {
+    startSocketServer({ server, applicationId, internalToken, fileSyncManager, syncService, draftClient, lockService, aiService, gitService, gitBootstrapAuthFailed, gitBootstrapFailed, activateGitService, refreshGitCredentials, checkAuthorization, onAuthenticated, path, tracer, gitOperationLock, superblocksBaseUrl, maxPayloadBytes = MAX_SOCKET_PAYLOAD_BYTES, }) {
         if (this.socketServer) {
             throw new Error("Socket server already started");
         }
+        const splitEditAfterClark = syncService && !aiService
+            ? new SplitEditAfterClark(syncService)
+            : undefined;
         const wss = new WebSocketServer({
             path,
             noServer: true,
@@ -224,8 +284,9 @@ class SocketManager extends TracedEventEmitter {
                 return;
             }
             try {
-                await lockService?.relinquishLockForTurn(() => headlessLockOwners.some((owner) => !owner.handbackFailed &&
-                    (!ownedTurns.includes(owner) || owner.peerId !== peerId)));
+                const stillNeeded = () => headlessLockOwners.some((owner) => !owner.handbackFailed &&
+                    (!ownedTurns.includes(owner) || owner.peerId !== peerId));
+                await lockService?.relinquishLockForTurn(stillNeeded);
             }
             catch (error) {
                 await headlessLockOperationQueue.enqueue(async () => {
@@ -281,8 +342,9 @@ class SocketManager extends TracedEventEmitter {
                 return true;
             }
             try {
-                await lockService?.relinquishLockForTurn(() => headlessLockOwners.some((owner) => !owner.handbackFailed &&
-                    (owner !== ownedTurn || owner.peerId !== peerId)));
+                const stillNeeded = () => headlessLockOwners.some((owner) => !owner.handbackFailed &&
+                    (owner !== ownedTurn || owner.peerId !== peerId));
+                await lockService?.relinquishLockForTurn(stillNeeded);
             }
             catch (error) {
                 await headlessLockOperationQueue.enqueue(async () => {
@@ -411,7 +473,7 @@ class SocketManager extends TracedEventEmitter {
         // because the dev server's own auto-commit pushes to the same remote with
         // the same credential and has to be able to report a rejection too.
         clearGitAuthFailed();
-        const getActiveGitService = async () => {
+        const getActiveGitService = async (options) => {
             if (disconnectInProgress) {
                 return undefined;
             }
@@ -448,6 +510,11 @@ class SocketManager extends TracedEventEmitter {
                         aiService?.setRuntimeGitServices(activeGitService, syncService, gitOperationLock);
                         return activeGitService;
                     }
+                    // No remote. A failed startup bootstrap resolves `undefined` too;
+                    // `gitBootstrapFailed` is what keeps those apps on the retry loop.
+                    if (!options?.retryWhenUnconfigured && !gitBootstrapFailed) {
+                        return undefined;
+                    }
                 }
                 catch (error) {
                     getLogger().warn("[git] activation attempt failed", {
@@ -525,17 +592,24 @@ class SocketManager extends TracedEventEmitter {
                         ws.close(1002, "Unsupported subprotocol");
                         return;
                     }
+                    if (kind === "ai-service" &&
+                        (req.headers.origin !== undefined ||
+                            !["127.0.0.1", "::1", "::ffff:127.0.0.1"].includes(req.socket.remoteAddress ?? ""))) {
+                        ws.close(1008, "Internal access denied");
+                        return;
+                    }
                     const peerId = url.searchParams.get("peerId");
                     if (!peerId) {
                         getLogger().error("No peerId in upgrade request");
                         ws.close(1008, "No peerId found");
                         return;
                     }
-                    if (url.searchParams.get("applicationId") !== syncService?.applicationId) {
+                    const expectedApplicationId = applicationId ?? syncService?.applicationId;
+                    if (url.searchParams.get("applicationId") !== expectedApplicationId) {
                         getLogger().error("Application ID mismatch", {
                             error: {
                                 kind: "ValidationError",
-                                message: `Expected application ID ${syncService?.applicationId}, received ${url.searchParams.get("applicationId")}`,
+                                message: `Expected application ID ${expectedApplicationId}, received ${url.searchParams.get("applicationId")}`,
                             },
                         });
                         ws.close(1008, "Application ID mismatch");
@@ -560,6 +634,19 @@ class SocketManager extends TracedEventEmitter {
             if (connectionLiveness) {
                 ws.on("message", () => connectionLiveness.markInbound());
             }
+            const gatewayLivenessTimer = kind === "gateway" &&
+                !lockService &&
+                draftClient &&
+                connectionLiveness
+                ? setInterval(() => {
+                    if (!connectionLiveness.observe().available &&
+                        headlessLockOwners.some((owner) => owner.peerId === peerId)) {
+                        dropHeadlessLockOwnersInBackground(peerId, "gateway liveness lapsed");
+                        aiService?.endHeadlessTurn(peerId, connectionLiveness);
+                    }
+                }, 30_000)
+                : undefined;
+            gatewayLivenessTimer?.unref();
             /**
              * When an editor authenticates after git was enabled, the lock may
              * still be on 'main'. This calls
@@ -578,10 +665,33 @@ class SocketManager extends TracedEventEmitter {
                 ? ++nextEditorConnectionSequence
                 : 0;
             const authMiddleware = async (params, ctx, _peer, next) => {
+                if (kind === "ai-service") {
+                    const expected = internalToken?.();
+                    const provided = ctx.peerAuthorization?.startsWith("Bearer ")
+                        ? ctx.peerAuthorization.slice(7)
+                        : undefined;
+                    const expectedBytes = Buffer.from(expected ?? "");
+                    const providedBytes = Buffer.from(provided ?? "");
+                    const hasInternalToken = !!expected &&
+                        !!provided &&
+                        expectedBytes.length === providedBytes.length &&
+                        timingSafeEqual(expectedBytes, providedBytes);
+                    if (!AI_SERVICE_METHODS.has(ctx.method) ||
+                        (!hasInternalToken &&
+                            !(await checkAuthorization(ctx.peerAuthorization)).success)) {
+                        throw new SocketErrorException(403, "Internal access denied");
+                    }
+                    return await next();
+                }
+                if (INTERNAL_METHODS.has(ctx.method)) {
+                    throw new SocketErrorException(403, "Internal access denied");
+                }
                 const result = await checkAuthorization(ctx.peerAuthorization);
                 if (!result.success) {
                     throw new SocketErrorException(result.status, result.error);
                 }
+                if (ctx.peerAuthorization)
+                    onAuthenticated?.(ctx.peerAuthorization);
                 if (isEditor && this.editorSockets.has(socket)) {
                     this.authenticatedEditorSockets.add(socket);
                 }
@@ -722,6 +832,94 @@ class SocketManager extends TracedEventEmitter {
             const middlewares = [recordingMiddleware, authMiddleware].filter((m) => m != null);
             const socket = createISocketClient(new TracedSocket(ws, {
                 editor: {
+                    internalApplyDraftChanges: [
+                        async ({ changes }) => {
+                            if (!syncService) {
+                                throw new Error("Sync service not available");
+                            }
+                            await syncService.applyDraftChanges(changes.map(({ filePath, content }) => ({
+                                filePath,
+                                ...(content && {
+                                    content: content.kind === "binary"
+                                        ? decodeDraftBinary(content.value)
+                                        : content.value,
+                                }),
+                            })));
+                        },
+                    ],
+                    internalCreateDraftMarkerOnly: [
+                        async ({ filePath }) => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            await syncService.createDraftMarkerOnly(filePath);
+                        },
+                    ],
+                    internalHasDraftChanges: [
+                        async () => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            return await syncService.hasLocalDraftChanges();
+                        },
+                    ],
+                    internalGetBranchName: [
+                        async () => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            return syncService.branchName;
+                        },
+                    ],
+                    internalUploadDirectory: [
+                        async ({ source }) => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            await syncService.uploadDirectory(source);
+                        },
+                    ],
+                    internalSyncCurrentDirectoryWithDraftsNow: [
+                        async ({ source }) => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            return await syncService.syncCurrentDirectoryWithDraftsNow(source);
+                        },
+                    ],
+                    internalGetDraftDiff: [
+                        async ({ paths }) => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            return await syncService.getDraftDiff(paths);
+                        },
+                    ],
+                    internalGetDraftChanges: [
+                        async () => {
+                            if (!syncService)
+                                throw new Error("Sync service not available");
+                            return await syncService.getDraftChanges();
+                        },
+                    ],
+                    internalAcquireTurnLock: [
+                        async (request) => {
+                            if (!lockService)
+                                throw new Error("Lock service not available");
+                            await lockService.acquireLockForTurn();
+                            lockService.markGatewayAttached("ai-service");
+                            if (request?.promptId) {
+                                await splitEditAfterClark?.startTurn(request.promptId, request.chatSessionId);
+                            }
+                        },
+                    ],
+                    internalReleaseTurnLock: [
+                        async (request) => {
+                            if (!lockService)
+                                throw new Error("Lock service not available");
+                            if (request?.promptId) {
+                                void splitEditAfterClark
+                                    ?.endTurn(request.promptId)
+                                    .catch((error) => getLogger().warn("Could not finish edit-after-Clark measurement", getErrorMeta(error)));
+                            }
+                            await lockService.relinquishLockForTurn();
+                            lockService.markGatewayDetached("ai-service");
+                        },
+                    ],
                     // Renames
                     renameElement: [fileSyncManager.handleRenameElement],
                     renameApi: [fileSyncManager.handleRenameApi],
@@ -740,12 +938,18 @@ class SocketManager extends TracedEventEmitter {
                     // AI + draft handling
                     aiGenerate: [
                         async (request) => {
+                            if (kind === "ai-service") {
+                                throw new Error("Internal client cannot start AI generation");
+                            }
                             const inputSource = { kind, peerId };
                             await aiService?.handleAiGenerate(request, inputSource);
                         },
                     ],
                     aiGenerateWithQueue: [
                         async (request) => {
+                            if (kind === "ai-service") {
+                                throw new Error("Internal client cannot start AI generation");
+                            }
                             if (!aiService) {
                                 throw new Error("AI service not available");
                             }
@@ -835,12 +1039,30 @@ class SocketManager extends TracedEventEmitter {
                     ],
                     aiRejectDraft: [
                         async (request) => {
+                            if (draftClient) {
+                                await draftClient.rejectDraft(request);
+                                await aiService?.handleUserRejectedDraft(request.silently);
+                                return;
+                            }
                             await syncService?.discardLocalDraftChanges("ai:reject");
+                            splitEditAfterClark?.discardLastTurn();
                             await aiService?.handleUserRejectedDraft(request.silently);
                         },
                     ],
                     aiAcceptDraft: [
                         async (request, _context) => {
+                            if (draftClient) {
+                                const result = await draftClient.acceptDraft(request);
+                                await aiService?.handleUserAcceptedDraft(request.silently);
+                                const appId = applicationId ?? syncService?.applicationId;
+                                if (result.commitId && aiService && appId) {
+                                    uploadAcceptedDraftContext(aiService, appId, result.commitId, result.branchName, _context.peerAuthorization);
+                                }
+                                return {
+                                    commitId: result.commitId,
+                                    ...(result.warning && { warning: result.warning }),
+                                };
+                            }
                             if (!syncService) {
                                 throw new Error("Sync service not available");
                             }
@@ -887,47 +1109,8 @@ class SocketManager extends TracedEventEmitter {
                                 });
                             }
                             if (result?.commitId) {
-                                // Upload context file to bucketeer asynchronously (fire-and-forget)
-                                if (aiService && syncService) {
-                                    void (async () => {
-                                        try {
-                                            // Get contextId using the same pattern as other context operations
-                                            // This gets userId from clark.context.jwt
-                                            const contextId = aiService.getContextId();
-                                            if (!contextId) {
-                                                getLogger().warn("[aiAcceptDraft] No contextId available (JWT missing or no userId), skipping context upload");
-                                                return;
-                                            }
-                                            // Get JWT from clark context (same source as userId)
-                                            // We need to access it via a method or use peer authorization as fallback
-                                            const jwt = aiService.getClarkJwt() ||
-                                                _context.peerAuthorization?.replace("Bearer ", "");
-                                            if (!jwt) {
-                                                getLogger().warn("[aiAcceptDraft] No JWT found, skipping context upload");
-                                                return;
-                                            }
-                                            await uploadContextToBucketeer({
-                                                applicationId: syncService.applicationId,
-                                                commitId: result.commitId,
-                                                branchName: syncService.branchName,
-                                                appRootDirPath: aiService.appRootDirPath,
-                                                useServerUpload: aiService.isAiFileUploadEnabled,
-                                                superblocksBaseUrl: aiService.superblocksBaseUrl,
-                                                logger: getLogger(),
-                                                jwt,
-                                            });
-                                            getLogger().info("[aiAcceptDraft] context upload completed", {
-                                                applicationId: syncService.applicationId,
-                                                branchName: syncService.branchName,
-                                                commitId: result.commitId,
-                                                contextUserId: contextId.userId,
-                                            });
-                                        }
-                                        catch (error) {
-                                            getLogger().warn("[aiAcceptDraft] Failed to upload context to bucketeer (non-blocking)", getErrorMeta(error));
-                                            // Don't throw - this is fire-and-forget
-                                        }
-                                    })();
+                                if (aiService) {
+                                    uploadAcceptedDraftContext(aiService, syncService.applicationId, result.commitId, syncService.branchName, _context.peerAuthorization);
                                 }
                                 return draftStateCleared === true
                                     ? { commitId: result.commitId }
@@ -966,13 +1149,23 @@ class SocketManager extends TracedEventEmitter {
                     ],
                     aiClearDraftState: [
                         async () => {
-                            await syncService?.clearDraftsDirectory();
+                            if (draftClient) {
+                                await draftClient.clearDraftState();
+                            }
+                            else {
+                                await syncService?.clearDraftsDirectory();
+                            }
                         },
                     ],
                     aiCancel: [
                         async (request) => {
                             await aiService?.handleUserCanceled(request?.reason);
-                            await syncService?.reportDraftStateToServer(false);
+                            if (draftClient) {
+                                await draftClient.reportDraftStateCleared();
+                            }
+                            else {
+                                await syncService?.reportDraftStateToServer(false);
+                            }
                         },
                     ],
                     aiNotifyIntegrationSetupClosed: [
@@ -1039,8 +1232,10 @@ class SocketManager extends TracedEventEmitter {
                             const messages = await aiService?.chatSessionStore.getMessages({
                                 force: request?.force,
                             });
+                            aiService?.applyEditorContextOptions?.(request?.contextOptionsV2);
                             return {
                                 messages: messages ?? [],
+                                ...aiService?.getChatStateSnapshot(),
                             };
                         },
                     ],
@@ -1463,7 +1658,9 @@ class SocketManager extends TracedEventEmitter {
                             // would be stuck with git permanently disabled.
                             disconnectInProgress = false;
                             clearGitAuthFailed();
-                            const service = await getActiveGitService();
+                            const service = await getActiveGitService({
+                                retryWhenUnconfigured: true,
+                            });
                             return { activated: !!service };
                         }),
                     ],
@@ -1802,12 +1999,13 @@ class SocketManager extends TracedEventEmitter {
                 // Git activation is deliberately NOT started here — it runs from the
                 // auth middleware, after presence. See `activateGitForAuthedEditor`.
             }
-            else {
+            else if (kind === "library") {
                 this.librarySockets.add(socket);
                 this.emit("connect:library", socket);
             }
             existingSockets.add(socket);
             ws.on("close", (code, reason) => {
+                clearInterval(gatewayLivenessTimer);
                 connectionLiveness?.close();
                 const peerSockets = this.peerIds.get(peerId);
                 peerSockets?.delete(socket);
@@ -2133,6 +2331,23 @@ class SocketManager extends TracedEventEmitter {
             socket.close();
         }
     }
+    async closeAllClientsGracefully() {
+        await Promise.all(Array.from(this.socketServer?.clients ?? []).map((socket) => new Promise((resolve) => {
+            if (socket.readyState === WebSocket.CLOSED) {
+                resolve();
+                return;
+            }
+            const timeout = setTimeout(() => {
+                socket.terminate();
+                resolve();
+            }, 1_000);
+            socket.once("close", () => {
+                clearTimeout(timeout);
+                resolve();
+            });
+            socket.close(1001, "Server shutting down");
+        })));
+    }
     dispose() {
         this._disposer();
     }