From 45df1b1978e597546a208ecbf5b3aca56137edc5 Mon Sep 17 00:00:00 2001
From: "Yuichiro Tachibana (Tsuchiya)" <t.yic.yt@gmail.com>
Date: Mon, 23 Dec 2024 14:46:26 +0900
Subject: [PATCH] Lite: Capture stdout and stderr from the main thread (#9984)

* Add stdout and stderr events

* add changeset

* Refactoring

* Format App.tsx

* add changeset

* Add python-error event to capture Python errors occurring in the running event loop after the initial app launch

* Fix <ErrorDisplay />'s close button

* Fix <ErrorDisplay />

* Propagate python-error and initialization-error events to the controller

* Add init-code|file-run-error events

---------

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
---
 .changeset/pink-signs-fall.md      |  7 ++++
 gradio/queueing.py                 |  4 +-
 gradio/wasm_utils.py               | 32 +++++++++++++++
 js/lite/lite.html                  |  2 +-
 js/lite/src/ErrorDisplay.svelte    | 66 ++++++++++++++++++++----------
 js/lite/src/LiteIndex.svelte       | 60 ++++++++++++++++++++-------
 js/lite/src/dev/App.svelte         | 51 +++++++++++++++++++++++
 js/lite/src/index.ts               | 26 ++++++++++++
 js/wasm/src/message-types.ts       | 26 +++++++++++-
 js/wasm/src/webworker/file.test.ts |  4 +-
 js/wasm/src/webworker/file.ts      |  1 -
 js/wasm/src/webworker/index.ts     | 50 ++++++++++++++++++----
 js/wasm/src/worker-proxy.ts        | 35 ++++++++++++----
 13 files changed, 308 insertions(+), 56 deletions(-)
 create mode 100644 .changeset/pink-signs-fall.md

diff --git a/.changeset/pink-signs-fall.md b/.changeset/pink-signs-fall.md
new file mode 100644
index 0000000000000..ea94952877020
--- /dev/null
+++ b/.changeset/pink-signs-fall.md
@@ -0,0 +1,7 @@
+---
+"@gradio/lite": minor
+"@gradio/wasm": minor
+"gradio": minor
+---
+
+feat:Lite: Capture stdout and stderr from the main thread
diff --git a/gradio/queueing.py b/gradio/queueing.py
index 63df1bfa285cf..82ad5568d6fe9 100644
--- a/gradio/queueing.py
+++ b/gradio/queueing.py
@@ -13,7 +13,7 @@
 
 import fastapi
 
-from gradio import route_utils, routes
+from gradio import route_utils, routes, wasm_utils
 from gradio.data_classes import (
     PredictBodyInternal,
 )
@@ -644,6 +644,7 @@ async def process_events(
                 err = e
                 for event in awake_events:
                     content = error_payload(err, app.get_blocks().show_error)
+                    wasm_utils.send_error(err)
                     self.send_message(
                         event,
                         ProcessCompletedMessage(
@@ -736,6 +737,7 @@ async def process_events(
                     success = False
                     error = err or old_err
                     output = error_payload(error, app.get_blocks().show_error)
+                    wasm_utils.send_error(error)
                 for event in awake_events:
                     self.send_message(
                         event, ProcessCompletedMessage(output=output, success=success)
diff --git a/gradio/wasm_utils.py b/gradio/wasm_utils.py
index bd78667629a12..360f8e579c5c7 100644
--- a/gradio/wasm_utils.py
+++ b/gradio/wasm_utils.py
@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import logging
 import sys
+import traceback
 from contextlib import contextmanager
 from contextvars import ContextVar
 
+LOGGER = logging.getLogger(__name__)
+
 # See https://pyodide.org/en/stable/usage/faq.html#how-to-detect-that-code-is-run-with-pyodide
 IS_WASM = sys.platform == "emscripten"
 
@@ -57,3 +61,31 @@ def get_registered_app(app_id: str):
         raise GradioAppNotFoundError(
             f"Gradio app not found (ID: {app_id}). Forgot to call demo.launch()?"
         ) from e
+
+
+error_traceback_callback_map = {}
+
+
+def register_error_traceback_callback(app_id, callback):
+    error_traceback_callback_map[app_id] = callback
+
+
+def send_error(error: Exception | None):
+    # The callback registered by the JS process is called with the error traceback
+    # for the WebWorker process to read the traceback.
+
+    if not IS_WASM:
+        return
+    if error is None:
+        return
+
+    app_id = _app_id_context_var.get()
+    callback = error_traceback_callback_map.get(app_id)
+    if not callback:
+        LOGGER.warning(
+            f"Error callback not found for the app ID {app_id}. The error will be ignored."
+        )
+        return
+
+    tb = "".join(traceback.format_exception(type(error), error, error.__traceback__))
+    callback(tb)
diff --git a/js/lite/lite.html b/js/lite/lite.html
index 843ed059eadfb..12c62912409b9 100644
--- a/js/lite/lite.html
+++ b/js/lite/lite.html
@@ -30,6 +30,6 @@
 		<div style="flex-grow: 1; overflow: scroll; position: relative">
 			<div id="gradio-app" style="min-height: 100%"></div>
 		</div>
-		<div id="dev-app" style="height: 300px; position: relative"></div>
+		<div id="dev-app" style="height: 50%; position: relative"></div>
 	</body>
 </html>
diff --git a/js/lite/src/ErrorDisplay.svelte b/js/lite/src/ErrorDisplay.svelte
index f65ad22a37f88..a8ac1082e3534 100644
--- a/js/lite/src/ErrorDisplay.svelte
+++ b/js/lite/src/ErrorDisplay.svelte
@@ -1,38 +1,61 @@
 <script lang="ts">
 	import { StatusTracker } from "@gradio/statustracker";
+	import { Embed } from "@gradio/core";
 
+	import { createEventDispatcher } from "svelte";
 	import { _ } from "svelte-i18n";
 	import { setupi18n } from "@gradio/core";
 
 	setupi18n();
 
+	const dispatch = createEventDispatcher();
+
 	export let is_embed: boolean;
 	export let error: Error | undefined = undefined;
+	export let height: string;
+
+	// For <Embed>
+	export let container: boolean;
+	export let version: string;
+	let wrapper: HTMLDivElement;
 </script>
 
-<StatusTracker
-	i18n={$_}
-	absolute={!is_embed}
-	status="error"
-	timer={false}
-	queue_position={null}
-	queue_size={null}
-	translucent={true}
-	autoscroll={false}
+<Embed
+	display={container && is_embed}
+	{is_embed}
+	info={false}
+	{version}
+	initial_height={height}
+	loaded={false}
+	space={null}
+	fill_width={false}
+	bind:wrapper
 >
-	<div class="error" slot="error">
-		{#if error}
-			{#if error.message}
-				<p class="error-name">
-					{error.message}
-				</p>
+	<StatusTracker
+		i18n={$_}
+		absolute={!is_embed}
+		status="error"
+		timer={false}
+		queue_position={null}
+		queue_size={null}
+		translucent={true}
+		autoscroll={false}
+		on:clear_status={() => dispatch("clear_error")}
+	>
+		<div class="error" slot="error">
+			{#if error}
+				{#if error.message}
+					<p class="error-name">
+						{error.message}
+					</p>
+				{/if}
+				{#if error.stack}
+					<pre class="error-stack"><code>{error.stack}</code></pre>
+				{/if}
 			{/if}
-			{#if error.stack}
-				<pre class="error-stack"><code>{error.stack}</code></pre>
-			{/if}
-		{/if}
-	</div>
-</StatusTracker>
+		</div>
+	</StatusTracker>
+</Embed>
 
 <style>
 	.error {
@@ -40,6 +63,7 @@
 		width: 100%;
 		padding: var(--size-4);
 		color: var(--body-text-color);
+		overflow: scroll;
 		/* Status tracker sets `pointer-events: none`.
 		Override it here so the user can scroll the element with `overflow: hidden`
 		and copy and paste the error message */
diff --git a/js/lite/src/LiteIndex.svelte b/js/lite/src/LiteIndex.svelte
index 4058d414be6ec..613299040d2b5 100644
--- a/js/lite/src/LiteIndex.svelte
+++ b/js/lite/src/LiteIndex.svelte
@@ -50,6 +50,26 @@
 		requirements: requirements ?? [],
 		sharedWorkerMode: sharedWorkerMode ?? false
 	});
+
+	const dispatch = createEventDispatcher();
+
+	worker_proxy.addEventListener("modules-auto-loaded", (event) => {
+		dispatch("modules-auto-loaded", (event as CustomEvent).detail);
+	});
+	worker_proxy.addEventListener("stdout", (event) => {
+		dispatch("stdout", (event as CustomEvent).detail);
+	});
+	worker_proxy.addEventListener("stderr", (event) => {
+		dispatch("stderr", (event as CustomEvent).detail);
+	});
+	worker_proxy.addEventListener("initialization-error", (event) => {
+		error = (event as CustomEvent).detail;
+		dispatch("initialization-error", (event as CustomEvent).detail);
+	});
+	worker_proxy.addEventListener("python-error", (event) => {
+		error = (event as CustomEvent).detail;
+		dispatch("python-error", (event as CustomEvent).detail);
+	});
 	onDestroy(() => {
 		worker_proxy.terminate();
 	});
@@ -90,24 +110,18 @@
 		worker_proxy.install.bind(worker_proxy)
 	);
 
-	worker_proxy.addEventListener("initialization-error", (event) => {
-		error = (event as CustomEvent).detail;
-	});
-
-	const dispatch = createEventDispatcher();
-
-	worker_proxy.addEventListener("modules-auto-loaded", (event) => {
-		dispatch("modules-auto-loaded", (event as CustomEvent).detail);
-	});
-
 	// Internally, the execution of `runPythonCode()` or `runPythonFile()` is queued
 	// and its promise will be resolved after the Pyodide is loaded and the worker initialization is done
 	// (see the await in the `onmessage` callback in the webworker code)
 	// So we don't await this promise because we want to mount the `Index` immediately and start the app initialization asynchronously.
 	if (code != null) {
-		worker_proxy.runPythonCode(code);
+		worker_proxy.runPythonCode(code).catch((err) => {
+			dispatch("init-code-run-error", err);
+		});
 	} else if (entrypoint != null) {
-		worker_proxy.runPythonFile(entrypoint);
+		worker_proxy.runPythonFile(entrypoint).catch((err) => {
+			dispatch("init-file-run-error", err);
+		});
 	} else {
 		throw new Error("Either code or entrypoint must be provided.");
 	}
@@ -157,7 +171,16 @@
 	>
 		{#key index_component_key}
 			{#if error}
-				<ErrorDisplay {error} is_embed />
+				<ErrorDisplay
+					{error}
+					{is_embed}
+					height={initial_height}
+					{container}
+					{version}
+					on:clear_error={() => {
+						error = null;
+					}}
+				/>
 			{:else}
 				<Index
 					space={null}
@@ -182,7 +205,16 @@
 {:else}
 	{#key index_component_key}
 		{#if error}
-			<ErrorDisplay {error} {is_embed} />
+			<ErrorDisplay
+				{error}
+				{is_embed}
+				height={initial_height}
+				{container}
+				{version}
+				on:clear_error={() => {
+					error = null;
+				}}
+			/>
 		{:else}
 			<Index
 				space={null}
diff --git a/js/lite/src/dev/App.svelte b/js/lite/src/dev/App.svelte
index f7b52aa833315..654cbf2873abf 100644
--- a/js/lite/src/dev/App.svelte
+++ b/js/lite/src/dev/App.svelte
@@ -59,6 +59,9 @@ def hi(name):
 
 	const requirements = parse_requirements(requirements_txt);
 
+	let stdouts: string[] = [];
+	let stderrs: string[] = [];
+
 	let controller: ReturnType<typeof create>;
 	onMount(() => {
 		controller = create({
@@ -85,6 +88,14 @@ def hi(name):
 			requirements_txt +=
 				"\n" + packageNames.map((line) => line + "  # auto-loaded").join("\n");
 		});
+		controller.addEventListener("stdout", (event) => {
+			const message = (event as CustomEvent).detail as string;
+			stdouts = stdouts.concat(message);
+		});
+		controller.addEventListener("stderr", (event) => {
+			const message = (event as CustomEvent).detail as string;
+			stderrs = stderrs.concat(message);
+		});
 	});
 	onDestroy(() => {
 		controller.unmount();
@@ -124,6 +135,27 @@ def hi(name):
 </script>
 
 <div class="container">
+	<div class="panel">
+		<div class="log-panel-container">
+			<div class="log-panel">
+				<h4>stdout</h4>
+				<div class="log-box" id="stdout" style="color: black;">
+					{#each stdouts as stdout}
+						<pre class="log-line">{stdout}</pre>
+					{/each}
+				</div>
+			</div>
+			<div class="log-panel">
+				<h4>stdout</h4>
+				<div class="log-box" id="stderr" style="color: red;">
+					{#each stderrs as stderr}
+						<pre class="log-line">{stderr}</pre>
+					{/each}
+				</div>
+			</div>
+		</div>
+	</div>
+
 	<div class="panel">
 		When the SharedWorker mode is enabled, access the URL below (for Chrome) and
 		click the "inspect" link of the worker to show the console log emitted from
@@ -184,6 +216,25 @@ def hi(name):
 		width: 100%;
 	}
 
+	.log-panel-container {
+		height: 300px;
+		position: relative;
+		display: flex;
+		flex-direction: row;
+	}
+	.log-panel {
+		width: 50%;
+		display: flex;
+		flex-direction: column;
+	}
+	.log-box {
+		flex-grow: 1;
+		overflow: scroll;
+	}
+	.log-line {
+		margin: 0;
+	}
+
 	.cell-header {
 		display: flex;
 		flex-direction: row;
diff --git a/js/lite/src/index.ts b/js/lite/src/index.ts
index bf5fe0022b88e..e192b779807cc 100644
--- a/js/lite/src/index.ts
+++ b/js/lite/src/index.ts
@@ -33,6 +33,32 @@ export class GradioAppController extends EventTarget {
 				new CustomEvent("modules-auto-loaded", { detail: event.detail })
 			);
 		});
+		this.lite_svelte_app.$on("stdout", (event: CustomEvent) => {
+			this.dispatchEvent(new CustomEvent("stdout", { detail: event.detail }));
+		});
+		this.lite_svelte_app.$on("stderr", (event: CustomEvent) => {
+			this.dispatchEvent(new CustomEvent("stderr", { detail: event.detail }));
+		});
+		this.lite_svelte_app.$on("initialization-error", (event: CustomEvent) => {
+			this.dispatchEvent(
+				new CustomEvent("initialization-error", { detail: event.detail })
+			);
+		});
+		this.lite_svelte_app.$on("python-error", (event: CustomEvent) => {
+			this.dispatchEvent(
+				new CustomEvent("python-error", { detail: event.detail })
+			);
+		});
+		this.lite_svelte_app.$on("init-code-run-error", (event: CustomEvent) => {
+			this.dispatchEvent(
+				new CustomEvent("init-code-run-error", { detail: event.detail })
+			);
+		});
+		this.lite_svelte_app.$on("init-file-run-error", (event: CustomEvent) => {
+			this.dispatchEvent(
+				new CustomEvent("init-file-run-error", { detail: event.detail })
+			);
+		});
 	}
 
 	run_code = (code: string): Promise<void> => {
diff --git a/js/wasm/src/message-types.ts b/js/wasm/src/message-types.ts
index 24996fbe5a9af..fb3c5f7363897 100644
--- a/js/wasm/src/message-types.ts
+++ b/js/wasm/src/message-types.ts
@@ -120,5 +120,27 @@ export interface OutMessageModulesAutoLoaded extends OutMessageBase {
 		packages: PackageData[];
 	};
 }
-
-export type OutMessage = OutMessageProgressUpdate | OutMessageModulesAutoLoaded;
+export interface OutMessageStdout extends OutMessageBase {
+	type: "stdout";
+	data: {
+		output: string;
+	};
+}
+export interface OutMessageStderr extends OutMessageBase {
+	type: "stderr";
+	data: {
+		output: string;
+	};
+}
+export interface OutMessagePythonError extends OutMessageBase {
+	type: "python-error";
+	data: {
+		traceback: string;
+	};
+}
+export type OutMessage =
+	| OutMessageProgressUpdate
+	| OutMessageModulesAutoLoaded
+	| OutMessageStdout
+	| OutMessageStderr
+	| OutMessagePythonError;
diff --git a/js/wasm/src/webworker/file.test.ts b/js/wasm/src/webworker/file.test.ts
index 27d3b2141a037..4d6f56d4b2044 100644
--- a/js/wasm/src/webworker/file.test.ts
+++ b/js/wasm/src/webworker/file.test.ts
@@ -1,8 +1,8 @@
 // @vitest-environment node
 
-import path from "path";
+import path from "node:path";
 import { loadPyodide, PyodideInterface } from "pyodide";
-import { describe, it, expect, beforeEach, beforeAll } from "vitest";
+import { describe, it, expect, beforeAll } from "vitest";
 import { writeFileWithParents, renameWithParents } from "./file";
 
 describe("writeFileWithParents()", () => {
diff --git a/js/wasm/src/webworker/file.ts b/js/wasm/src/webworker/file.ts
index c19ef2242fb2f..78f8104811733 100644
--- a/js/wasm/src/webworker/file.ts
+++ b/js/wasm/src/webworker/file.ts
@@ -8,7 +8,6 @@ export const resolveAppHomeBasedPath = (
 	appId: string,
 	filePath: string
 ): string => {
-	const normalized = path.normalize(filePath);
 	return path.resolve(getAppHomeDir(appId), filePath);
 };
 
diff --git a/js/wasm/src/webworker/index.ts b/js/wasm/src/webworker/index.ts
index 218f830b1fd98..3cc3bb985522a 100644
--- a/js/wasm/src/webworker/index.ts
+++ b/js/wasm/src/webworker/index.ts
@@ -57,13 +57,15 @@ let unload_local_modules: (target_dir_path?: string) => void;
 
 async function initializeEnvironment(
 	options: InMessageInitEnv["data"],
-	updateProgress: (log: string) => void
+	updateProgress: (log: string) => void,
+	stdout: (output: string) => void,
+	stderr: (output: string) => void
 ): Promise<void> {
 	console.debug("Loading Pyodide.");
 	updateProgress("Loading Pyodide");
 	pyodide = await loadPyodide({
-		stdout: console.debug,
-		stderr: console.error
+		stdout,
+		stderr
 	});
 	console.debug("Pyodide is loaded.");
 
@@ -288,13 +290,39 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 	console.debug("Set up a new app.", { appId });
 
 	const updateProgress = (log: string): void => {
-		const message: OutMessage = {
+		receiver.postMessage({
 			type: "progress-update",
 			data: {
 				log
 			}
-		};
-		receiver.postMessage(message);
+		} satisfies OutMessage);
+	};
+	const stdout = (output: string): void => {
+		console.log(output);
+		receiver.postMessage({
+			type: "stdout",
+			data: {
+				output
+			}
+		} satisfies OutMessage);
+	};
+	const stderr = (output: string): void => {
+		console.error(output);
+		receiver.postMessage({
+			type: "stderr",
+			data: {
+				output
+			}
+		} satisfies OutMessage);
+	};
+	const onPythonError = (traceback: string): void => {
+		console.error("Python error:", traceback);
+		receiver.postMessage({
+			type: "python-error",
+			data: {
+				traceback: traceback
+			}
+		} satisfies OutMessage);
 	};
 	const onModulesAutoLoaded = (packages: PackageData[]) => {
 		const message: OutMessage = {
@@ -320,7 +348,12 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 		try {
 			if (msg.type === "init-env") {
 				if (envReadyPromise == null) {
-					envReadyPromise = initializeEnvironment(msg.data, updateProgress);
+					envReadyPromise = initializeEnvironment(
+						msg.data,
+						updateProgress,
+						stdout,
+						stderr
+					);
 				} else {
 					updateProgress(
 						"Pyodide environment initialization is ongoing in another session"
@@ -350,6 +383,9 @@ function setupMessageHandler(receiver: MessageTransceiver): void {
 			}
 			await envReadyPromise;
 
+			const gradio = pyodide.pyimport("gradio");
+			gradio.wasm_utils.register_error_traceback_callback(appId, onPythonError);
+
 			if (msg.type === "init-app") {
 				appReadyPromise = initializeApp(
 					appId,
diff --git a/js/wasm/src/worker-proxy.ts b/js/wasm/src/worker-proxy.ts
index 81e5496c41962..cdb27b4057cc4 100644
--- a/js/wasm/src/worker-proxy.ts
+++ b/js/wasm/src/worker-proxy.ts
@@ -51,16 +51,12 @@ export class WorkerProxy extends EventTarget {
 		if (sharedWorkerMode) {
 			this.postMessageTarget = (this.worker as SharedWorker).port;
 			this.postMessageTarget.start();
-			this.postMessageTarget.onmessage = (e) => {
-				this._processWorkerMessage(e.data);
-			};
 		} else {
 			this.postMessageTarget = this.worker as globalThis.Worker;
-
-			(this.worker as globalThis.Worker).onmessage = (e) => {
-				this._processWorkerMessage(e.data);
-			};
 		}
+		this.postMessageTarget.onmessage = (e) => {
+			this._processWorkerMessage(e.data);
+		};
 
 		this.postMessageAsync({
 			type: "init-env",
@@ -172,6 +168,31 @@ export class WorkerProxy extends EventTarget {
 						detail: msg.data.packages
 					})
 				);
+				break;
+			}
+			case "stdout": {
+				this.dispatchEvent(
+					new CustomEvent("stdout", {
+						detail: msg.data.output
+					})
+				);
+				break;
+			}
+			case "stderr": {
+				this.dispatchEvent(
+					new CustomEvent("stderr", {
+						detail: msg.data.output
+					})
+				);
+				break;
+			}
+			case "python-error": {
+				this.dispatchEvent(
+					new CustomEvent("python-error", {
+						detail: new Error(msg.data.traceback)
+					})
+				);
+				break;
 			}
 		}
 	}