Resume sessions

.changeset/quiet-sessions-return.md

@@ -0,0 +1,7 @@
+---
+"@gradio/client": minor
+"gradio": minor
+"gradio_client": minor
+---
+
+feat: Resume queued jobs and restore undelivered output after temporary disconnects or page refreshes

client/js/src/client.ts

@@ -22,6 +22,11 @@ import { post_data } from "./utils/post_data";
 import { predict } from "./utils/predict";
 import { duplicate } from "./utils/duplicate";
 import { submit } from "./utils/submit";
+import {
+	get_resumable_events,
+	get_resumable_session_hash,
+	type ResumableJob
+} from "./utils/session";
 import { RE_SPACE_NAME, process_endpoint } from "./helpers/api_info";
 import {
 	map_names_to_ids,
@@ -37,6 +42,7 @@ import { open_stream, readable_stream, close_stream } from "./utils/stream";
 import {
 	API_INFO_ERROR_MSG,
 	APP_ID_URL,
+	CLOSE_URL,
 	CONFIG_ERROR_MSG,
 	HEARTBEAT_URL,
 	COMPONENT_SERVER_URL
@@ -57,14 +63,18 @@ export class Client {
 	last_status: Record<string, Status["stage"]> = {};
 
 	private cookies: string | null = null;
+	private restored_session_hash = false;
 
 	// streaming
 	stream_status = { open: false };
 	closed = false;
-	pending_stream_messages: Record<string, any[][]> = {};
+	pending_stream_messages: Record<string, any[]> = {};
 	pending_diff_streams: Record<string, any[][]> = {};
 	event_callbacks: Record<string, (data?: unknown) => Promise<void>> = {};
 	unclosed_events: Set<string> = new Set();
+	events_to_resume: Set<string> = new Set();
+	stream_reconnect_attempts = 0;
+	stream_reconnect_timer: ReturnType<typeof setTimeout> | null = null;
 	heartbeat_event: EventSource | null = null;
 	abort_controller: AbortController | null = null;
 	stream_instance: EventSource | null = null;
@@ -181,6 +191,7 @@ export class Client {
 		trigger_id?: number | null,
 		all_events?: boolean
 	) => SubmitIterable<GradioEvent>;
+	resume_jobs: (jobs?: ResumableJob[]) => SubmitIterable<GradioEvent>[];
 	predict: <T = unknown>(
 		endpoint: string | number,
 		data: unknown[] | Record<string, unknown> | undefined,
@@ -215,6 +226,21 @@ export class Client {
 		this.handle_blob = handle_blob.bind(this);
 		this.post_data = post_data.bind(this);
 		this.submit = submit.bind(this);
+		this.resume_jobs = (jobs = this.get_resumable_events()) => {
+			this.options.resume_sessions = true;
+			return jobs.map(({ fn_index, event_id }) =>
+				submit.call(
+					this,
+					fn_index,
+					{},
+					undefined,
+					undefined,
+					undefined,
+					undefined,
+					event_id
+				)
+			);
+		};
 		this.predict = predict.bind(this) as typeof this.predict;
 		this.open_stream = open_stream.bind(this);
 		this.resolve_config = resolve_config.bind(this);
@@ -232,10 +258,19 @@ export class Client {
 			await this.resolve_cookies();
 		}
 
-		await this._resolve_config().then(
-			(res: { config: Config } | undefined) =>
-				res?.config && this._resolve_heartbeat(res.config)
-		);
+		const resolved = (await this._resolve_config()) as
+			| { config: Config }
+			| undefined;
+		if (resolved?.config) {
+			if (
+				this.restored_session_hash &&
+				typeof sessionStorage !== "undefined" &&
+				get_resumable_events(resolved.config, this.session_hash).length === 0
+			) {
+				this.session_hash = Math.random().toString(36).substring(2);
+			}
+			await this._resolve_heartbeat(resolved.config);
+		}
 
 		try {
 			this.api_info = await this.view_api();
@@ -293,13 +328,24 @@ export class Client {
 		}
 	): Promise<Client> {
 		const client = new this(app_reference, options); // this refers to the class itself, not the instance
-		if (options.session_hash) {
-			client.session_hash = options.session_hash;
+		const session_hash =
+			options.session_hash ||
+			(options.resume_sessions
+				? get_resumable_session_hash(options.cookies)
+				: null);
+		if (session_hash) {
+			client.session_hash = session_hash;
+			client.restored_session_hash = !options.session_hash;
 		}
 		await client.init();
 		return client;
 	}
 
+	get_resumable_events(): ResumableJob[] {
+		if (!this.options.resume_sessions || !this.config) return [];
+		return get_resumable_events(this.config, this.session_hash);
+	}
+
 	async reconnect(): Promise<"connected" | "broken" | "changed"> {
 		const app_id_url = new URL(
 			`${this.config!.root}${this.api_prefix}/${APP_ID_URL}`
@@ -321,7 +367,31 @@ export class Client {
 	}
 
 	close(): void {
+		if (
+			!this.closed &&
+			this.options.resume_sessions &&
+			this.config &&
+			typeof window !== "undefined"
+		) {
+			const headers: Record<string, string> = {
+				"Content-Type": "application/json"
+			};
+			if (this.options.token) {
+				headers.Authorization = `Bearer ${this.options.token}`;
+			}
+			void this.fetch(`${this.config.root}${this.api_prefix}/${CLOSE_URL}`, {
+				method: "POST",
+				headers,
+				body: JSON.stringify({ session_hash: this.session_hash }),
+				credentials: this.options.credentials ?? "same-origin",
+				keepalive: true
+			}).catch(() => {});
+		}
 		this.closed = true;
+		if (this.stream_reconnect_timer) {
+			clearTimeout(this.stream_reconnect_timer);
+			this.stream_reconnect_timer = null;
+		}
 		close_stream(this.stream_status, this.abort_controller);
 	}
 

client/js/src/constants.ts

@@ -15,6 +15,8 @@ export const HEARTBEAT_URL = `heartbeat`;
 export const COMPONENT_SERVER_URL = `component_server`;
 export const RESET_URL = `reset`;
 export const CANCEL_URL = `cancel`;
+export const ACK_URL = `queue/ack`;
+export const CLOSE_URL = `queue/close`;
 export const APP_ID_URL = `app_id`;
 
 export const RAW_API_INFO_URL = `info?serialize=False`;

client/js/src/index.ts

@@ -5,6 +5,7 @@ export { submit } from "./utils/submit";
 export { upload_files } from "./utils/upload_files";
 export { FileData, upload, prepare_files } from "./upload";
 export { handle_file } from "./helpers/data";
+export type { ResumableJob } from "./utils/session";
 
 export type {
 	SpaceStatus,

client/js/src/test/init.test.ts

@@ -16,6 +16,8 @@ import {
 } from "./test_data";
 import { initialise_server } from "./server";
 import { SPACE_NOT_FOUND_MSG } from "../constants";
+import { track_resumable_event } from "../utils/session";
+import { http, HttpResponse } from "msw";
 
 const app_reference = "hmb/hello_world";
 const broken_app_reference = "hmb/bye_world";
@@ -65,6 +67,23 @@ describe("Client class", () => {
 			expect(app.config).toEqual(config_response);
 		});
 
+		test.skipIf(typeof sessionStorage === "undefined")(
+			"does not restore a session from a different app",
+			async () => {
+				track_resumable_event(
+					{ ...config_response, app_id: "another-app" },
+					"restored-session",
+					{ event_id: "event-id", fn_index: 0 }
+				);
+
+				const app = await Client.connect(direct_app_reference, {
+					resume_sessions: true
+				});
+
+				expect(app.session_hash).not.toBe("restored-session");
+			}
+		);
+
 		test("connecting successfully to a private running app with a space reference", async () => {
 			const app = await Client.connect("hmb/secret_world", {
 				token: "hf_123"
@@ -122,6 +141,59 @@ describe("Client class", () => {
 		});
 	});
 
+	describe("resume_jobs", () => {
+		test("reattaches each job to its existing queue event", async () => {
+			const app = await Client.connect(direct_app_reference);
+			app.stream_status.open = true;
+
+			const submissions = app.resume_jobs([
+				{ event_id: "event-1", fn_index: 0 },
+				{ event_id: "event-2", fn_index: 0 }
+			]);
+
+			await expect(submissions[0].wait_for_id()).resolves.toBe("event-1");
+			await expect(submissions[1].wait_for_id()).resolves.toBe("event-2");
+			expect(app.event_callbacks["event-1"]).toBeDefined();
+			expect(app.event_callbacks["event-2"]).toBeDefined();
+			expect(app.options.resume_sessions).toBe(true);
+
+			await Promise.all(submissions.map((submission) => submission.return()));
+		});
+
+		test.skipIf(typeof window === "undefined")(
+			"notifies the server when a resumable client closes",
+			async () => {
+				const app = await Client.connect(secret_direct_app_reference, {
+					token: "hf_123",
+					resume_sessions: true
+				});
+				let received_session_hash: string | undefined;
+				let received_authorization: string | null = null;
+				let resolve_request: () => void = () => {};
+				const request_received = new Promise<void>((resolve) => {
+					resolve_request = resolve;
+				});
+				server.resetHandlers(
+					http.post(
+						`${secret_direct_app_reference}/queue/close`,
+						async ({ request }) => {
+							const body = (await request.json()) as { session_hash: string };
+							received_session_hash = body.session_hash;
+							received_authorization = request.headers.get("Authorization");
+							resolve_request();
+							return HttpResponse.json({ success: true });
+						}
+					)
+				);
+				app.close();
+				await request_received;
+
+				expect(received_session_hash).toBe(app.session_hash);
+				expect(received_authorization).toBe("Bearer hf_123");
+			}
+		);
+	});
+
 	describe("duplicate", () => {
 		test("backwards compatibility of duplicate using deprecated syntax", async () => {
 			const app = await duplicate("gradio/hello_world", {

client/js/src/test/session.test.ts

@@ -0,0 +1,87 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { Config } from "../types";
+import {
+	clear_resumable_event,
+	get_resumable_events,
+	get_resumable_session_hash,
+	track_resumable_event
+} from "../utils/session";
+
+class MemoryStorage implements Storage {
+	values = new Map<string, string>();
+
+	get length(): number {
+		return this.values.size;
+	}
+
+	clear(): void {
+		this.values.clear();
+	}
+
+	getItem(key: string): string | null {
+		return this.values.get(key) ?? null;
+	}
+
+	key(index: number): string | null {
+		return [...this.values.keys()][index] ?? null;
+	}
+
+	removeItem(key: string): void {
+		this.values.delete(key);
+	}
+
+	setItem(key: string, value: string): void {
+		this.values.set(key, value);
+	}
+}
+
+const config = {
+	app_id: "app-1",
+	root:
+		typeof location === "undefined" ? "https://example.com/" : location.origin
+} as Config;
+
+describe("resumable sessions", () => {
+	beforeEach(() => {
+		if (typeof sessionStorage === "undefined") {
+			vi.stubGlobal("sessionStorage", new MemoryStorage());
+		}
+		sessionStorage.clear();
+		if (typeof document !== "undefined") {
+			document.cookie = "gradio_active_session=; Path=/; Max-Age=0";
+		}
+	});
+
+	it("stores only active events for the current app session", () => {
+		track_resumable_event(config, "session-1", {
+			event_id: "event-1",
+			fn_index: 3
+		});
+
+		expect(get_resumable_session_hash()).toBe("session-1");
+		expect(get_resumable_events(config, "session-1")).toEqual([
+			{ event_id: "event-1", fn_index: 3 }
+		]);
+
+		clear_resumable_event("event-1");
+		expect(get_resumable_events(config, "session-1")).toEqual([]);
+	});
+
+	it("drops events when the app has changed", () => {
+		track_resumable_event(config, "session-1", {
+			event_id: "event-1",
+			fn_index: 3
+		});
+
+		expect(
+			get_resumable_events({ ...config, app_id: "app-2" }, "session-1")
+		).toEqual([]);
+		expect(get_resumable_session_hash()).toBeNull();
+	});
+
+	it("reads the active session cookie during server rendering", () => {
+		expect(
+			get_resumable_session_hash("other=value; gradio_active_session=session-2")
+		).toBe("session-2");
+	});
+});