Resume sessions

.changeset/quiet-sessions-return.md

@@ -0,0 +1,6 @@
+---
+"@gradio/client": minor
+"gradio": 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 ResumableEvent
+} from "./utils/session";
 import { RE_SPACE_NAME, process_endpoint } from "./helpers/api_info";
 import {
 	map_names_to_ids,
@@ -56,14 +61,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;
@@ -180,6 +189,10 @@ export class Client {
 		trigger_id?: number | null,
 		all_events?: boolean
 	) => SubmitIterable<GradioEvent>;
+	resume: (
+		endpoint: string | number,
+		event_id: string
+	) => SubmitIterable<GradioEvent>;
 	predict: <T = unknown>(
 		endpoint: string | number,
 		data: unknown[] | Record<string, unknown> | undefined,
@@ -210,6 +223,17 @@ export class Client {
 		this.handle_blob = handle_blob.bind(this);
 		this.post_data = post_data.bind(this);
 		this.submit = submit.bind(this);
+		this.resume = (endpoint, event_id) =>
+			submit.call(
+				this,
+				endpoint,
+				{},
+				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);
@@ -227,9 +251,15 @@ export class Client {
 			await this.resolve_cookies();
 		}
 
-		await this._resolve_config().then(({ config }) =>
-			this._resolve_heartbeat(config)
-		);
+		const { config } = await this._resolve_config();
+		if (
+			this.restored_session_hash &&
+			typeof sessionStorage !== "undefined" &&
+			get_resumable_events(config, this.session_hash).length === 0
+		) {
+			this.session_hash = Math.random().toString(36).substring(2);
+		}
+		await this._resolve_heartbeat(config);
 
 		try {
 			this.api_info = await this.view_api();
@@ -287,13 +317,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(): ResumableEvent[] {
+		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}`
@@ -316,6 +357,10 @@ export class Client {
 
 	close(): void {
 		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,7 @@ 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 APP_ID_URL = `app_id`;
 
 export const RAW_API_INFO_URL = `info?serialize=False`;

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

@@ -16,6 +16,7 @@ import {
 } from "./test_data";
 import { initialise_server } from "./server";
 import { SPACE_METADATA_ERROR_MSG } from "../constants";
+import { track_resumable_event } from "../utils/session";
 
 const app_reference = "hmb/hello_world";
 const broken_app_reference = "hmb/bye_world";
@@ -65,6 +66,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"

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");
+	});
+});

client/js/src/test/stream.test.ts

@@ -37,7 +37,8 @@ describe("open_stream", () => {
 	});
 
 	afterEach(() => {
-		vi.clearAllMocks();
+		vi.restoreAllMocks();
+		vi.useRealTimers();
 	});
 
 	it("should throw an error if config is not defined", async () => {
@@ -62,10 +63,10 @@ describe("open_stream", () => {
 		if (!app.stream_instance?.onmessage || !app.stream_instance?.onerror) {
 			throw new Error("stream instance is not defined");
 		}
-
+		const stream = app.stream_instance;
 		const message = { msg: "hello jerry" };
 
-		app.stream_instance.onmessage({
+		stream.onmessage({
 			data: JSON.stringify(message)
 		} as MessageEvent);
 		expect(app.stream_status.open).toBe(true);
@@ -74,14 +75,73 @@ describe("open_stream", () => {
 		expect(app.pending_stream_messages).toEqual({});
 
 		const close_stream_message = { msg: "close_stream" };
-		app.stream_instance.onmessage({
+		await stream.onmessage({
 			data: JSON.stringify(close_stream_message)
 		} as MessageEvent);
 		expect(app.stream_status.open).toBe(false);
+		expect(app.stream_instance).toBeNull();
+		stream.onerror?.({
+			data: JSON.stringify("closed stream")
+		} as MessageEvent);
+		expect(app.stream_status.open).toBe(false);
+		expect(app.stream).toHaveBeenCalledTimes(1);
+	});
+
+	it("should open another stream when jobs remain after a normal close", async () => {
+		const callback = vi.fn().mockResolvedValue(undefined);
+		app.event_callbacks["event-1"] = callback;
+		app.unclosed_events.add("event-1");
+
+		await app.open_stream();
+
+		const stream = app.stream_instance;
+		if (!stream?.onmessage) {
+			throw new Error("stream instance is not defined");
+		}
+
+		await stream.onmessage({
+			data: JSON.stringify({ msg: "close_stream" })
+		} as MessageEvent);
+
+		expect(app.stream).toHaveBeenCalledTimes(2);
+		expect(app.stream_status.open).toBe(true);
+
+		stream.onerror?.({
+			data: JSON.stringify("closed stream")
+		} as MessageEvent);
+		expect(app.stream).toHaveBeenCalledTimes(2);
+	});
+
+	it("should reconnect the SSE stream after an error when jobs are active", async () => {
+		vi.useFakeTimers();
+		vi.spyOn(console, "error").mockImplementation(() => {});
+		const callback = vi.fn().mockResolvedValue(undefined);
+		app.event_callbacks["event-1"] = callback;
+		app.unclosed_events.add("event-1");
+
+		await app.open_stream();
+
+		if (!app.stream_instance?.onerror) {
+			throw new Error("stream instance is not defined");
+		}
 
 		app.stream_instance.onerror({
-			data: JSON.stringify("404")
+			data: JSON.stringify("network error")
 		} as MessageEvent);
+
 		expect(app.stream_status.open).toBe(false);
+		expect(callback).not.toHaveBeenCalled();
+		expect(app.stream).toHaveBeenCalledTimes(1);
+		expect(app.events_to_resume).toContain("event-1");
+
+		await vi.advanceTimersByTimeAsync(500);
+
+		expect(app.stream).toHaveBeenCalledTimes(2);
+		expect(app.stream_status.open).toBe(true);
+		expect(
+			(app.stream as Mock).mock.calls[1][0].searchParams.getAll(
+				"resume_event_id"
+			)
+		).toEqual(["event-1"]);
 	});
 });

client/js/src/types.ts

@@ -110,6 +110,7 @@ export interface SubmitIterable<T> extends AsyncIterable<T> {
 	event_id: () => string;
 	send_chunk: (payload: Record<string, unknown>) => void;
 	wait_for_id: () => Promise<string | null>;
+	acknowledge: () => Promise<void>;
 	close_stream: () => void;
 }
 
@@ -330,6 +331,7 @@ export interface ClientOptions {
 	headers?: Record<string, string> | Headers;
 	query_params?: Record<string, string>;
 	session_hash?: string;
+	resume_sessions?: boolean;
 	cookies?: string;
 	credentials?: RequestCredentials;
 }