Merge pull request #1194 from writer/AB-714

feat(ui): prevent socket to be closed after a background job - AB-714

Author: madeindjs

src/ui/src/builder/BuilderApp.vue

@@ -189,7 +189,9 @@ const collaborationManager = inject(injectionKeys.collaborationManager);
 const tracking = useWriterTracking(wf);
 const toasts = useToasts();
 
-provide(injectionKeys.socketTimeout, useSocketTimeout(wf, 10));
+const SOCKET_TIMEOUT_MIN =
+	Number(localStorage.getItem("socketTimeoutMin")) || 10;
+provide(injectionKeys.socketTimeout, useSocketTimeout(wf, SOCKET_TIMEOUT_MIN));
 
 const noteEl = useTemplateRef("noteEl");
 const rendererEl = useTemplateRef("rendererEl");

src/ui/src/builder/useSocketTimeout.test.ts

@@ -10,6 +10,7 @@ import {
 import { nextTick } from "vue";
 import { useSocketTimeout } from "./useSocketTimeout";
 import { buildMockCore } from "@/tests/mocks";
+import { flushPromises } from "@vue/test-utils";
 
 // Mock dependencies
 vi.mock("@/composables/useLogger", () => ({
@@ -89,7 +90,7 @@ describe("useSocketTimeout", () => {
 		expect(schedule()).toBe(true);
 	});
 
-	it("should not schedule timeout if frontend message exists", () => {
+	it("should not schedule timeout if frontend message exists", async () => {
 		mockCore.frontendMessageMap.value.set(1, { type: "event" });
 		const { schedule, onVisibilityChange } = useSocketTimeout(
 			mockCore.core,
@@ -99,6 +100,12 @@ describe("useSocketTimeout", () => {
 		onVisibilityChange();
 
 		expect(schedule()).toBe(false);
+		expect(vi.getTimerCount()).toBe(0);
+
+		// even if the task finished in background
+		mockCore.frontendMessageMap.value.clear();
+		await flushPromises();
+		expect(vi.getTimerCount()).toBe(0);
 	});
 
 	it("should clear schedule when clearSchedule is called", () => {

src/ui/src/builder/useSocketTimeout.ts

@@ -12,6 +12,7 @@ const WEBSOCKET_INACTIVITY_CLOSE_CODE = 4000;
 export function useSocketTimeout(wf: Core, timeoutMin: number) {
 	const timeoutMs = timeoutMin * 60 * 1_000;
 	const logger = useLogger();
+	const abort = useAbortController();
 
 	let timer = undefined;
 
@@ -51,8 +52,6 @@ export function useSocketTimeout(wf: Core, timeoutMin: number) {
 		() => !prevent.value && !isMessagePending.value,
 	);
 
-	const abort = useAbortController();
-
 	function schedule() {
 		if (!canCloseSocket.value || document.visibilityState !== "hidden")
 			return false;
@@ -77,13 +76,16 @@ export function useSocketTimeout(wf: Core, timeoutMin: number) {
 		timer = undefined;
 	}
 
+	// cancel closing socket timer
 	watch(canCloseSocket, () => {
 		if (!canCloseSocket.value) {
 			clearSchedule();
 		} else if (document.visibilityState === "hidden") {
 			schedule(); // Reschedule now that activity is complete
 		}
 	});
+
+	// reflect `preventTasks` changes in session storage
 	watch(
 		preventTasks,
 		() => {
@@ -92,6 +94,17 @@ export function useSocketTimeout(wf: Core, timeoutMin: number) {
 		{ deep: true },
 	);
 
+	// prevent closing socket when blueprint finish in background
+	watch(isMessagePending, (hasPending, hadPending) => {
+		const hadCompletedJob = hadPending && !hasPending;
+		if (hadCompletedJob && document.visibilityState === "hidden") {
+			logger.log(
+				`[SocketTimeout] Preventing socket to be closed due to finished task in background...`,
+			);
+			preventTasks.value.add("pendingMessageInBackground");
+		}
+	});
+
 	async function reconnect() {
 		logger.info(`[SocketTimeout] Attempting to reconnect socket...`);
 		reconnecting.value = true;
@@ -112,6 +125,7 @@ export function useSocketTimeout(wf: Core, timeoutMin: number) {
 
 	function onVisibilityChange() {
 		if (document.visibilityState === "visible") {
+			preventTasks.value.delete("pendingMessageInBackground");
 			clearSchedule();
 		} else if (canCloseSocket.value) {
 			schedule();