fetch token from parent frame

Author: TGlide

.env.example

@@ -1,4 +1,5 @@
 PUBLIC_ENABLE_MCP=
+PUBLIC_HF_TOKEN=
 HYPERBOLIC_API_KEY=
 COHERE_API_KEY=
 TOGETHER_API_KEY=

src/lib/components/debug-menu.svelte

@@ -46,12 +46,6 @@
 				showQuotaModal();
 			},
 		},
-		{
-			label: "Show token modal",
-			cb: () => {
-				token.showModal = true;
-			},
-		},
 		{
 			label: "Test toast",
 			cb: () => {

src/lib/components/inference-playground/hf-token-modal.svelte

@@ -25,7 +25,6 @@
 	import PlaygroundConversationHeader from "./conversation-header.svelte";
 	import PlaygroundConversation from "./conversation.svelte";
 	import GenerationConfig from "./generation-config.svelte";
-	import HFTokenModal from "./hf-token-modal.svelte";
 	import MessageTextarea from "./message-textarea.svelte";
 	import ModelSelectorModal from "./model-selector-modal.svelte";
 	import ModelSelector from "./model-selector.svelte";
@@ -39,30 +38,8 @@
 
 	const systemPromptSupported = $derived(conversations.active.some(c => isSystemPromptSupported(c.model)));
 	const compareActive = $derived(conversations.active.length === 2);
-
-	function handleTokenSubmit(e: Event) {
-		const form = e.target as HTMLFormElement;
-		const formData = new FormData(form);
-		const submittedHfToken = (formData.get("hf-token") as string).trim() ?? "";
-		const RE_HF_TOKEN = /\bhf_[a-zA-Z0-9]{34}\b/;
-		if (RE_HF_TOKEN.test(submittedHfToken)) {
-			token.value = submittedHfToken;
-			// TODO: Only submit when previous action was trying to submit
-			// submit();
-		} else {
-			alert("Please provide a valid HF token.");
-		}
-	}
 </script>
 
-{#if token.showModal}
-	<HFTokenModal
-		bind:storeLocallyHfToken={token.writeToLocalStorage}
-		on:close={() => (token.showModal = false)}
-		on:submit={handleTokenSubmit}
-	/>
-{/if}
-
 <div
 	class={[
 		"motion-safe:animate-fade-in grid h-dvh divide-gray-200 overflow-hidden bg-gray-100/50",
@@ -253,26 +230,6 @@
 							<IconWaterfall class="text-xs" />
 							Metrics
 						</a>
-						<button
-							onclick={token.reset}
-							class="flex items-center gap-1 text-sm text-gray-500 underline decoration-gray-300 hover:text-gray-800 dark:text-gray-400 dark:decoration-gray-600 dark:hover:text-gray-200"
-						>
-							<svg xmlns="http://www.w3.org/2000/svg" class="text-xs" width="1em" height="1em" viewBox="0 0 32 32">
-								<path
-									fill="currentColor"
-									d="M23.216 4H26V2h-7v6h2V5.096A11.96 11.96 0 0 1 28 16c0 6.617-5.383 12-12 12v2c7.72 0 14-6.28 14-14c0-5.009-2.632-9.512-6.784-12"
-								/>
-								<path fill="currentColor" d="M16 20a1.5 1.5 0 1 0 0 3a1.5 1.5 0 0 0 0-3M15 9h2v9h-2z" /><path
-									fill="currentColor"
-									d="M16 4V2C8.28 2 2 8.28 2 16c0 4.977 2.607 9.494 6.784 12H6v2h7v-6h-2v2.903A11.97 11.97 0 0 1 4 16C4 9.383 9.383 4 16 4"
-								/>
-							</svg>
-							{#if token.value}
-								Reset token
-							{:else}
-								Set token
-							{/if}
-						</button>
 					</div>
 				</div>
 

src/lib/components/inference-playground/playground.svelte

@@ -202,7 +202,7 @@ export class ConversationClass {
 
 	genNextMessage = async () => {
 		if (!token.value) {
-			token.showModal = true;
+			token.requestTokenFromParent();
 			return;
 		}
 
@@ -437,7 +437,7 @@ class Conversations {
 
 	genNextMessages = async (conv: "left" | "right" | "both" | ConversationClass = "both") => {
 		if (!token.value) {
-			token.showModal = true;
+			token.requestTokenFromParent();
 			return;
 		}
 

src/lib/state/conversations.svelte.ts

@@ -1,17 +1,28 @@
 import { safeParse } from "$lib/utils/json.js";
+import { PUBLIC_HF_TOKEN } from "$env/static/public";
 import typia from "typia";
 
 const key = "hf_token";
 
 class Token {
 	#value = $state("");
 	writeToLocalStorage = $state(true);
-	showModal = $state(false);
 
 	constructor() {
+		if (PUBLIC_HF_TOKEN) {
+			this.#value = PUBLIC_HF_TOKEN;
+			return;
+		}
+
 		const storedHfToken = localStorage.getItem(key);
 		const parsed = safeParse(storedHfToken ?? "");
-		this.value = typia.is<string>(parsed) ? parsed : "";
+		const storedToken = typia.is<string>(parsed) ? parsed : "";
+
+		if (storedToken) {
+			this.#value = storedToken;
+		} else {
+			this.requestTokenFromParent();
+		}
 	}
 
 	get value() {
@@ -23,12 +34,34 @@ class Token {
 			localStorage.setItem(key, JSON.stringify(token));
 		}
 		this.#value = token;
-		this.showModal = !token.length;
 	}
 
+	requestTokenFromParent = (): Promise<void> => {
+		if (typeof window === "undefined") return Promise.resolve();
+
+		return new Promise(resolve => {
+			const handleMessage = (event: MessageEvent) => {
+				if (event.data.type === "INFERENCE_JWT_RESPONSE") {
+					const token = event.data.token;
+					if (token && typeof token === "string") {
+						this.value = token;
+						window.removeEventListener("message", handleMessage);
+						resolve();
+					}
+				}
+			};
+
+			window.addEventListener("message", handleMessage);
+			window.parent?.postMessage({ type: "INFERENCE_JWT_REQUEST" }, "*");
+		});
+	};
+
 	reset = () => {
 		this.value = "";
 		localStorage.removeItem(key);
+		if (!PUBLIC_HF_TOKEN) {
+			this.requestTokenFromParent();
+		}
 	};
 }
 

src/lib/state/token.svelte.ts

@@ -122,6 +122,7 @@ export async function handleStreamingResponse(
 	conversation: ConversationClass | Conversation,
 	onChunk: (content: string) => void,
 	abortController: AbortController,
+	retryCount = 0,
 ): Promise<void> {
 	const data = conversation instanceof ConversationClass ? conversation.data : conversation;
 	const model = conversation.model;
@@ -149,28 +150,37 @@ export async function handleStreamingResponse(
 		enabledMCPs: getEnabledMCPs(),
 	};
 
-	const reader = await StreamReader.fromFetch("/api/generate", {
-		method: "POST",
-		headers: {
-			"Content-Type": "application/json",
-		},
-		body: JSON.stringify(requestBody),
-		signal: abortController.signal,
-	});
+	try {
+		const reader = await StreamReader.fromFetch("/api/generate", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify(requestBody),
+			signal: abortController.signal,
+		});
 
-	let out = "";
-	for await (const chunk of reader.read()) {
-		if (chunk.type === "chunk" && chunk.content) {
-			out += chunk.content;
-			onChunk(out);
-		} else if (chunk.type === "error") {
-			throw new Error(chunk.error || "Stream error");
+		let out = "";
+		for await (const chunk of reader.read()) {
+			if (chunk.type === "chunk" && chunk.content) {
+				out += chunk.content;
+				onChunk(out);
+			} else if (chunk.type === "error") {
+				throw new Error(chunk.error || "Stream error");
+			}
+		}
+	} catch (error) {
+		if (error instanceof Error && error.message.includes("401") && retryCount === 0) {
+			await token.requestTokenFromParent();
+			return handleStreamingResponse(conversation, onChunk, abortController, retryCount + 1);
 		}
+		throw error;
 	}
 }
 
 export async function handleNonStreamingResponse(
 	conversation: ConversationClass | Conversation,
+	retryCount = 0,
 ): Promise<{ message: ChatCompletionOutputMessage; completion_tokens: number }> {
 	const data = conversation instanceof ConversationClass ? conversation.data : conversation;
 	const model = conversation.model;
@@ -207,6 +217,10 @@ export async function handleNonStreamingResponse(
 	});
 
 	if (!response.ok) {
+		if (response.status === 401 && retryCount === 0) {
+			await token.requestTokenFromParent();
+			return handleNonStreamingResponse(conversation, retryCount + 1);
+		}
 		const error = await response.json();
 		throw new Error(error.error || "Failed to generate response");
 	}

src/lib/utils/business.svelte.ts

@@ -49,6 +49,9 @@ export class StreamReader {
 	static async fromFetch(url: string, options?: RequestInit): Promise<StreamReader> {
 		const response = await fetch(url, options);
 		if (!response.ok) {
+			if (response.status === 401) {
+				throw new Error("401 Unauthorized");
+			}
 			const error = await response.json();
 			throw new Error(error.error || "Request failed");
 		}

src/lib/utils/stream.ts

@@ -7,6 +7,7 @@ import { createAdapter, type GenerationArgs } from "./adapter.js";
 import { connectToMCPServers, executeMcpTool, type MCPServerConnection } from "./mcp.js";
 import type { FinishReason, GenerateRequest } from "./types.js";
 import { debugLog } from "./utils.js";
+import { InferenceClientProviderApiError, InferenceClientHubApiError } from "@huggingface/inference";
 
 type AssistantResponse = { message: ChatCompletionMessage; finish_reason: FinishReason };
 
@@ -197,6 +198,16 @@ export const POST: RequestHandler = async ({ request }) => {
 	} catch (error) {
 		debugLog(JSON.stringify(error, null, 2));
 		console.error("Generation error:", error);
-		return json({ error: error instanceof Error ? error.message : "Unknown error occurred" }, { status: 500 });
+
+		const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
+		let status = 500;
+
+		if (error instanceof InferenceClientProviderApiError || error instanceof InferenceClientHubApiError) {
+			status = error.httpResponse.status;
+		} else if (error && typeof error === "object" && "status" in error && typeof error.status === "number") {
+			status = error.status;
+		}
+
+		return json({ error: errorMessage }, { status });
 	}
 };