Use xet token provided by server to save API calls (#1854)

coyotte508 · web-flow · commit 4e78fafa9413 · 2025-11-24T15:29:55.000+01:00
cc @assafvayner for viz , @Kakulukian too https://huggingface.slack.com/archives/C035N5UE2JU/p1763991238834179
diff --git a/packages/hub/src/lib/collection-info.spec.ts b/packages/hub/src/lib/collection-info.spec.ts
@@ -103,6 +103,7 @@ describe("collectionInfo", () => {
 					lastModified: "2023-04-01T15:29:38.000Z",
 					likes: expect.any(Number),
 					pinned: false,
+					featured: false,
 					position: 2,
 					private: false,
 					repoType: "space",
diff --git a/packages/hub/src/lib/commit.ts b/packages/hub/src/lib/commit.ts
@@ -22,6 +22,7 @@ import { eventToGenerator } from "../utils/eventToGenerator";
 import { base64FromBytes } from "../utils/base64FromBytes";
 import { isFrontend } from "../utils/isFrontend";
 import { createBlobs } from "../utils/createBlobs";
+import type { XetTokenParams } from "../utils/uploadShards";
 import { uploadShards } from "../utils/uploadShards";
 import { splitAsyncGenerator } from "../utils/splitAsyncGenerator";
 import { SplicedBlob } from "../utils/SplicedBlob";
@@ -350,8 +351,8 @@ export async function* commitIter(params: CommitParams): AsyncGenerator<CommitPr
 			if (useXet && json.transfer !== "xet") {
 				useXet = false;
 			}
-			let xetRefreshWriteTokenUrl: string | undefined;
-			let xetSessionId: string | undefined;
+
+			let xetParams: XetTokenParams | null = null;
 
 			if (useXet) {
 				// First get all the files that are already uploaded out of the way
@@ -377,16 +378,17 @@ export async function* commitIter(params: CommitParams): AsyncGenerator<CommitPr
 							state: "uploading",
 						};
 					} else {
-						xetRefreshWriteTokenUrl = obj.actions.upload.href;
-						// Also, obj.actions.upload.header: {
-						//   X-Xet-Cas-Url:          string;
-						//   X-Xet-Access-Token:     string;
-						//   X-Xet-Token-Expiration: string;
-						//   X-Xet-Session-Id:       string;
-						// }
 						const headers = new Headers(obj.actions.upload.header);
-						xetSessionId = headers.get("X-Xet-Session-Id") ?? undefined;
-						// todo: use other data, like x-xet-cas-url, ...
+
+						xetParams = {
+							sessionId: headers.get("X-Xet-Session-Id") ?? undefined,
+							casUrl: headers.get("X-Xet-Cas-Url") ?? undefined,
+							accessToken: headers.get("X-Xet-Access-Token") ?? undefined,
+							expiresAt: headers.get("X-Xet-Token-Expiration")
+								? new Date(parseInt(headers.get("X-Xet-Token-Expiration") ?? "0") * 1000)
+								: undefined,
+							refreshWriteTokenUrl: obj.actions.upload.href,
+						};
 					}
 				}
 				const source = (async function* () {
@@ -399,8 +401,8 @@ export async function* commitIter(params: CommitParams): AsyncGenerator<CommitPr
 						yield { content: op.content, path: op.path, sha256: obj.oid };
 					}
 				})();
-				if (xetRefreshWriteTokenUrl) {
-					const xetRefreshWriteTokenUrlFixed = xetRefreshWriteTokenUrl;
+				if (xetParams) {
+					const fixedXetParams = xetParams;
 					const sources = splitAsyncGenerator(source, 5);
 					yield* eventToGenerator((yieldCallback, returnCallback, rejectCallback) =>
 						Promise.all(
@@ -410,21 +412,19 @@ export async function* commitIter(params: CommitParams): AsyncGenerator<CommitPr
 									accessToken,
 									hubUrl: params.hubUrl ?? HUB_URL,
 									repo: repoId,
-									xetRefreshWriteTokenUrl: xetRefreshWriteTokenUrlFixed,
-									xetSessionId,
+									xetParams: fixedXetParams,
 									// todo: maybe leave empty if PR?
 									rev: params.branch ?? "main",
 									isPullRequest: params.isPullRequest,
 									yieldCallback: (event) => yieldCallback({ ...event, state: "uploading" }),
 								})) {
 									if (event.event === "file") {
-										// No need: uploading xorbs already sent a fileProgress event with progress 1
-										// yieldCallback({
-										// 	event: "fileProgress" as const,
-										// 	path: event.path,
-										// 	progress: 1,
-										// 	state: "uploading" as const,
-										// });
+										yieldCallback({
+											event: "fileProgress" as const,
+											path: event.path,
+											progress: 1,
+											state: "uploading" as const,
+										});
 									} else if (event.event === "fileProgress") {
 										yieldCallback({
 											event: "fileProgress" as const,
diff --git a/packages/hub/src/lib/list-collections.spec.ts b/packages/hub/src/lib/list-collections.spec.ts
@@ -115,6 +115,7 @@ describe("listCollections", () => {
 						pinned: false,
 						position: 2,
 						private: false,
+						featured: false,
 						repoType: "space",
 						runtime: {
 							hardware: {
diff --git a/packages/hub/src/lib/upload-files-with-progress.spec.ts b/packages/hub/src/lib/upload-files-with-progress.spec.ts
@@ -70,7 +70,9 @@ describe("uploadFilesWithProgress", () => {
 					// 	assert(intermediateUploadEvents.length > 0, "There should be at least one intermediate upload event");
 					// }
 					progressEvents = progressEvents.filter(
-						(e) => e.event !== "fileProgress" || e.progress === 0 || e.progress === 1
+						(e, i) =>
+							(e.event !== "fileProgress" || e.progress === 0 || e.progress === 1) &&
+							(i === 0 || JSON.stringify(e) !== JSON.stringify(progressEvents[i - 1]))
 					);
 
 					assert.deepStrictEqual(progressEvents, [
diff --git a/packages/hub/src/lib/who-am-i.spec.ts b/packages/hub/src/lib/who-am-i.spec.ts
@@ -21,6 +21,7 @@ describe("whoAmI", () => {
 			isPro: false,
 			periodEnd: null,
 			avatarUrl: "/avatars/934b830e9fdaa879487852f79eef7165.svg",
+			billingMode: "postpaid",
 			orgs: [],
 			auth: {
 				type: "access_token",
diff --git a/packages/hub/src/lib/who-am-i.ts b/packages/hub/src/lib/who-am-i.ts
@@ -20,6 +20,7 @@ export interface WhoAmIUser {
 	 * Unix timestamp in seconds
 	 */
 	periodEnd: number | null;
+	billingMode: "postpaid" | "prepaid";
 }
 
 export interface WhoAmIOrg {
diff --git a/packages/hub/src/types/api/api-collection.ts b/packages/hub/src/types/api/api-collection.ts
@@ -195,6 +195,7 @@ interface ApiCollectionItemSpace extends ApiCollectionItemBase {
 	likes: number;
 	pinned: boolean;
 	private: boolean;
+	featured: boolean;
 	repoType: "space";
 	title: string;
 	sdk?: "gradio" | "docker" | "static" | "streamlit";
diff --git a/packages/hub/src/types/api/api-who-am-i.ts b/packages/hub/src/types/api/api-who-am-i.ts
@@ -28,6 +28,7 @@ interface ApiWhoAmIUser extends ApiWhoAmIEntityBase {
 	emailVerified: boolean;
 	isPro: boolean;
 	orgs: ApiWhoAmIOrg[];
+	billingMode: "postpaid" | "prepaid";
 }
 
 interface ApiWhoAmIApp extends ApiWhoAmIBase {
diff --git a/packages/hub/src/utils/createXorbs.ts b/packages/hub/src/utils/createXorbs.ts
@@ -214,7 +214,7 @@ export async function* createXorbs(
 
 					let cacheData = chunkCache.getChunk(chunk.hash, chunkModule.compute_hmac);
 					if (cacheData === undefined && chunk.dedup && bytesSinceRemoteDedup >= INTERVAL_BETWEEN_REMOTE_DEDUP) {
-						const token = await xetWriteToken({ ...params, isPullRequest: params.isPullRequest });
+						const token = await xetWriteToken(params);
 						bytesSinceRemoteDedup = 0;
 
 						const shardResp = await (params.fetch ?? fetch)(token.casUrl + "/v1/chunks/default/" + chunk.hash, {
@@ -701,7 +701,7 @@ async function loadDedupInfoToCache(
 
 				// Try remote dedup lookup if conditions are met
 				if (chunk.dedup && bytesSinceRemoteDedup >= INTERVAL_BETWEEN_REMOTE_DEDUP) {
-					const token = await xetWriteToken({ ...params, isPullRequest: params.isPullRequest });
+					const token = await xetWriteToken(params);
 					bytesSinceRemoteDedup = 0;
 
 					const shardResp = await (params.fetch ?? fetch)(token.casUrl + "/v1/chunks/default/" + chunk.hash, {
diff --git a/packages/hub/src/utils/uploadShards.ts b/packages/hub/src/utils/uploadShards.ts
@@ -51,11 +51,18 @@ export const SHARD_MAGIC_TAG = new Uint8Array([
 	169,
 ]);
 
+export interface XetTokenParams {
+	sessionId?: string;
+	casUrl?: string;
+	accessToken?: string;
+	expiresAt?: Date;
+	refreshWriteTokenUrl: string;
+}
+
 interface UploadShardsParams {
 	accessToken: string | undefined;
 	hubUrl: string;
-	xetRefreshWriteTokenUrl: string;
-	xetSessionId: string | undefined;
+	xetParams: XetTokenParams;
 	fetch?: typeof fetch;
 	repo: RepoId;
 	rev: string;
@@ -360,14 +367,14 @@ async function uploadXorb(
 	xorb: { hash: string; xorb: Uint8Array; files: Array<{ path: string; progress: number; lastSentProgress: number }> },
 	params: UploadShardsParams
 ) {
-	const token = await xetWriteToken({ ...params, isPullRequest: params.isPullRequest });
+	const token = await xetWriteToken(params);
 
 	const resp = await (params.fetch ?? fetch)(`${token.casUrl}/v1/xorbs/default/${xorb.hash}`, {
 		method: "POST",
 		body: xorb.xorb,
 		headers: {
 			Authorization: `Bearer ${token.accessToken}`,
-			...(params.xetSessionId ? { "X-Xet-Session-Id": params.xetSessionId } : {}),
+			...(params.xetParams.sessionId ? { "X-Xet-Session-Id": params.xetParams.sessionId } : {}),
 		},
 		...{
 			progressHint: {
@@ -390,14 +397,14 @@ async function uploadXorb(
 }
 
 async function uploadShard(shard: Uint8Array, params: UploadShardsParams) {
-	const token = await xetWriteToken({ ...params, isPullRequest: params.isPullRequest });
+	const token = await xetWriteToken(params);
 
 	const resp = await (params.fetch ?? fetch)(`${token.casUrl}/v1/shards`, {
 		method: "POST",
 		body: shard,
 		headers: {
 			Authorization: `Bearer ${token.accessToken}`,
-			...(params.xetSessionId ? { "X-Xet-Session-Id": params.xetSessionId } : {}),
+			...(params.xetParams.sessionId ? { "X-Xet-Session-Id": params.xetParams.sessionId } : {}),
 		},
 	});
 
diff --git a/packages/hub/src/utils/xetWriteToken.ts b/packages/hub/src/utils/xetWriteToken.ts
@@ -1,23 +1,15 @@
 import { createApiError } from "../error";
-import type { RepoId } from "../types/public";
+import type { XetTokenParams } from "./uploadShards";
 
 export interface XetWriteTokenParams {
 	accessToken: string | undefined;
-	hubUrl: string;
 	fetch?: typeof fetch;
-	repo: RepoId;
-	rev: string;
-	isPullRequest?: boolean;
-	xetRefreshWriteTokenUrl: string | undefined;
+	xetParams: XetTokenParams;
 }
 
 const JWT_SAFETY_PERIOD = 60_000;
 const JWT_CACHE_SIZE = 1_000;
 
-function cacheKey(params: Omit<XetWriteTokenParams, "fetch">): string {
-	return JSON.stringify([params.hubUrl, params.repo, params.rev, params.accessToken, params.isPullRequest]);
-}
-
 const jwtPromises: Map<string, Promise<{ accessToken: string; casUrl: string }>> = new Map();
 /**
  * Cache to store JWTs, to avoid making many auth requests when downloading multiple files from the same repo
@@ -32,7 +24,15 @@ const jwts: Map<
 > = new Map();
 
 export async function xetWriteToken(params: XetWriteTokenParams): Promise<{ accessToken: string; casUrl: string }> {
-	const key = cacheKey(params);
+	if (
+		params.xetParams.expiresAt &&
+		params.xetParams.casUrl &&
+		params.xetParams.accessToken &&
+		params.xetParams.expiresAt > new Date(Date.now() + JWT_SAFETY_PERIOD)
+	) {
+		return { accessToken: params.xetParams.accessToken, casUrl: params.xetParams.casUrl };
+	}
+	const key = params.xetParams.refreshWriteTokenUrl;
 
 	const jwt = jwts.get(key);
 
@@ -47,19 +47,16 @@ export async function xetWriteToken(params: XetWriteTokenParams): Promise<{ acce
 	}
 
 	const promise = (async () => {
-		const resp = await (params.fetch ?? fetch)(
-			params.xetRefreshWriteTokenUrl ??
-				`${params.hubUrl}/api/${params.repo.type}s/${params.repo.name}/xet-write-token/${encodeURIComponent(
-					params.rev
-				)}` + (params.isPullRequest ? "?create_pr=1" : ""),
-			{
-				headers: params.accessToken
+		const resp = await (params.fetch ?? fetch)(params.xetParams.refreshWriteTokenUrl, {
+			headers: {
+				...(params.accessToken
 					? {
 							Authorization: `Bearer ${params.accessToken}`,
 					  }
-					: {},
-			}
-		);
+					: {}),
+				...(params.xetParams.sessionId ? { "X-Xet-Session-Id": params.xetParams.sessionId } : {}),
+			},
+		});
 
 		if (!resp.ok) {
 			throw await createApiError(resp);

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@ export interface WhoAmIUser {`
`20`	`20`	`* Unix timestamp in seconds`
`21`	`21`	`*/`
`22`	`22`	`periodEnd: number \| null;`
	`23`	`+ billingMode: "postpaid" \| "prepaid";`
`23`	`24`	`}`
`24`	`25`
`25`	`26`	`export interface WhoAmIOrg {`
Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,7 @@ interface ApiWhoAmIUser extends ApiWhoAmIEntityBase {`
`28`	`28`	`emailVerified: boolean;`
`29`	`29`	`isPro: boolean;`
`30`	`30`	`orgs: ApiWhoAmIOrg[];`
	`31`	`+ billingMode: "postpaid" \| "prepaid";`
`31`	`32`	`}`
`32`	`33`
`33`	`34`	`interface ApiWhoAmIApp extends ApiWhoAmIBase {`