@notionhq/custom-blocks-host 0.1.30 → 0.2.0

package/src/lifecycle/initializationController.ts

@@ -1,57 +1,34 @@
-import type { NotionContrastMode } from "@notionhq/custom-blocks-protocol/contrast.js"
-import type { NotionBlockId } from "@notionhq/custom-blocks-protocol/ids.js"
-import type { CustomBlockManifest } from "@notionhq/custom-blocks-protocol/manifest.js"
+import { DEFAULT_CONTRAST_MODE } from "@notionhq/custom-blocks-protocol/contrast.js"
 import type { ConnectMessage } from "@notionhq/custom-blocks-protocol/messages/connect.js"
-import type { HostToSandboxMessage } from "@notionhq/custom-blocks-protocol/messages/hostToSandbox.js"
 import type {
 	CustomBlockInitErrorInfo,
 	InitMessage,
 } from "@notionhq/custom-blocks-protocol/messages/init.js"
 import type { InitResultMessage } from "@notionhq/custom-blocks-protocol/messages/initResult.js"
-import type { CustomBlockPage } from "@notionhq/custom-blocks-protocol/pages/page.js"
-import type { NotionParent } from "@notionhq/custom-blocks-protocol/parent.js"
-import type { NotionTheme } from "@notionhq/custom-blocks-protocol/theme.js"
-import type { NotionUser } from "@notionhq/custom-blocks-protocol/users/user.js"
 import type { CustomBlockHostFailure } from "../messages/types.js"
+import { unreachable } from "../utils.js"
 import { initErrorForFailureReason } from "./initErrors.js"
 import { isValidBridgeProtocolVersion } from "./protocolVersion.js"
 import type {
 	ConnectSuccessMessage,
-	CustomBlockHostDataSourcesPayload,
+	CustomBlockHostInitialStateResult,
 	InitStatus,
 } from "./types.js"
 
 const DEFAULT_NO_CONNECT_TIMEOUT_MS = 5_000
+const DEFAULT_INITIAL_STATE_TIMEOUT_MS = 5_000
 const DEFAULT_NO_INIT_RESULT_TIMEOUT_MS = 5_000
 
-type InitializationState = {
-	theme: NotionTheme
-	contrastMode: NotionContrastMode
-	parent: NotionParent
-	page: CustomBlockPage
-	currentUser: NotionUser
-}
-
-type InitializationPreparation =
-	| {
-			status: "success"
-			manifest: CustomBlockManifest
-			dataSources: CustomBlockHostDataSourcesPayload
-	  }
-	| {
-			status: "error"
-			error: CustomBlockInitErrorInfo
-	  }
-
 export type InitializationControllerOptions = {
-	blockId: NotionBlockId
-	getState: () => InitializationState
-	prepareInit: (connect: ConnectSuccessMessage) => InitializationPreparation
-	post: (message: HostToSandboxMessage) => void
+	prepareInit: (
+		connect: ConnectSuccessMessage,
+	) => Promise<CustomBlockHostInitialStateResult>
+	sendInitMessage: (message: InitMessage) => void
 	warn: (payload: unknown) => void
 	flushPendingHostStateChanges: () => void
 	minBridgeProtocolVersion: number
 	noConnectTimeoutMs?: number
+	initialStateTimeoutMs?: number
 	noInitResultTimeoutMs?: number
 	stopHostWithFailure: (error: CustomBlockHostFailure) => void
 	onInitResult?: (message: InitResultMessage) => void
@@ -80,14 +57,13 @@ export function createInitializationController(
 	options: InitializationControllerOptions,
 ): InitializationController {
 	const {
-		blockId,
-		getState,
 		prepareInit,
-		post,
+		sendInitMessage,
 		warn,
 		flushPendingHostStateChanges,
 		minBridgeProtocolVersion,
 		noConnectTimeoutMs = DEFAULT_NO_CONNECT_TIMEOUT_MS,
+		initialStateTimeoutMs = DEFAULT_INITIAL_STATE_TIMEOUT_MS,
 		noInitResultTimeoutMs = DEFAULT_NO_INIT_RESULT_TIMEOUT_MS,
 		stopHostWithFailure,
 		onInitResult,
@@ -96,6 +72,7 @@ export function createInitializationController(
 
 	let initStatus: InitStatus = "waitingForConnect"
 	let noConnectTimeoutId: number | undefined
+	let initialStateTimeoutId: number | undefined
 	let noInitResultTimeoutId: number | undefined
 	let activeInitializationId: string | undefined
 	let hasObservedIframeLoad = false
@@ -108,6 +85,13 @@ export function createInitializationController(
 		}
 	}
 
+	function clearInitialStateTimeout() {
+		if (initialStateTimeoutId !== undefined) {
+			window.clearTimeout(initialStateTimeoutId)
+			initialStateTimeoutId = undefined
+		}
+	}
+
 	function clearNoInitResultTimeout() {
 		if (noInitResultTimeoutId !== undefined) {
 			window.clearTimeout(noInitResultTimeoutId)
@@ -120,13 +104,23 @@ export function createInitializationController(
 		error: CustomBlockInitErrorInfo
 	}) {
 		const { initializationId, error } = args
-		if (initStatus !== "waitingForConnect") {
-			return
+		switch (initStatus) {
+			case "waitingForConnect":
+			case "preparingInit":
+				break
+			case "waitingForInitResult":
+			case "initialized":
+			case "error":
+			case "stopped":
+				return
+			default:
+				return unreachable(initStatus)
 		}
 		initStatus = "error"
 		activeInitializationId = initializationId
 		clearNoConnectTimeout()
-		post({
+		clearInitialStateTimeout()
+		sendInitMessage({
 			type: "init",
 			initializationId,
 			status: "error",
@@ -135,49 +129,79 @@ export function createInitializationController(
 		onFailure?.(error)
 	}
 
-	function sendInit(connect: ConnectSuccessMessage): boolean {
+	async function prepareAndSendInit(connect: ConnectSuccessMessage) {
 		if (initStatus !== "waitingForConnect") {
-			return false
-		}
-		const preparation = prepareInit(connect)
-		if (preparation.status === "error") {
-			sendInitError({
-				initializationId: connect.initializationId,
-				error: preparation.error,
-			})
-			return false
+			return
 		}
-		initStatus = "waitingForInitResult"
+		initStatus = "preparingInit"
 		activeInitializationId = connect.initializationId
 		clearNoConnectTimeout()
-		const state = getState()
-		const message: InitMessage = {
-			type: "init",
-			initializationId: connect.initializationId,
-			status: "success",
-			theme: state.theme,
-			contrastMode: state.contrastMode,
-			blockId,
-			parent: state.parent,
-			page: state.page,
-			manifest: preparation.manifest,
-			dataSources: preparation.dataSources,
-			currentUser: state.currentUser,
-		}
-		post(message)
-		noInitResultTimeoutId = window.setTimeout(() => {
-			if (initStatus !== "waitingForInitResult") {
+		initialStateTimeoutId = window.setTimeout(() => {
+			if (initStatus !== "preparingInit") {
 				return
 			}
-			initStatus = "error"
-			const error = initErrorForFailureReason("no_init_result")
+			const error = initErrorForFailureReason("initial_state_timeout")
 			warn({
-				message: "Sandbox did not send initResult",
+				...error,
 				initializationId: connect.initializationId,
 			})
-			onFailure?.(error)
-		}, noInitResultTimeoutMs)
-		return true
+			sendInitError({ initializationId: connect.initializationId, error })
+		}, initialStateTimeoutMs)
+
+		try {
+			const preparation = await prepareInit(connect)
+			if (initStatus !== "preparingInit") {
+				return
+			}
+			clearInitialStateTimeout()
+			if (preparation.status === "error") {
+				sendInitError({
+					initializationId: connect.initializationId,
+					error: preparation.error,
+				})
+				return
+			}
+			initStatus = "waitingForInitResult"
+			const message: InitMessage = {
+				type: "init",
+				initializationId: connect.initializationId,
+				status: "success",
+				theme: preparation.initialState.theme,
+				contrastMode:
+					preparation.initialState.contrastMode ?? DEFAULT_CONTRAST_MODE,
+				blockId: preparation.initialState.blockId,
+				parent: preparation.initialState.parent,
+				page: preparation.initialState.page,
+				currentUser: preparation.initialState.currentUser,
+				manifest: preparation.initialState.manifest,
+				dataSources: preparation.initialState.dataSources,
+			}
+			sendInitMessage(message)
+			noInitResultTimeoutId = window.setTimeout(() => {
+				if (initStatus !== "waitingForInitResult") {
+					return
+				}
+				initStatus = "error"
+				const error = initErrorForFailureReason("no_init_result")
+				onFailure?.(error)
+				warn({
+					...error,
+					initializationId: connect.initializationId,
+				})
+			}, noInitResultTimeoutMs)
+		} catch (error) {
+			if (initStatus !== "preparingInit") {
+				return
+			}
+			warn({
+				message: "Failed to prepare custom block initialization",
+				error,
+			})
+			sendInitError({
+				initializationId: connect.initializationId,
+				error: initErrorForFailureReason("unknown_error"),
+			})
+		}
 	}
 
 	function handleConnect(message: ConnectMessage) {
@@ -223,9 +247,8 @@ export function createInitializationController(
 			return
 		}
 
-		if (sendInit(message)) {
-			hasConnectedSandbox = true
-		}
+		hasConnectedSandbox = true
+		void prepareAndSendInit(message)
 	}
 
 	function handleInvalidConnect(args: {
@@ -323,11 +346,10 @@ export function createInitializationController(
 			return
 		}
 		clearNoInitResultTimeout()
+		initStatus = message.status === "success" ? "initialized" : "error"
 		if (message.status === "success") {
-			initStatus = "success"
 			flushPendingHostStateChanges()
 		} else {
-			initStatus = "error"
 			onFailure?.(message.error)
 		}
 		onInitResult?.(message)
@@ -336,6 +358,7 @@ export function createInitializationController(
 	function stop() {
 		initStatus = "stopped"
 		clearNoConnectTimeout()
+		clearInitialStateTimeout()
 		clearNoInitResultTimeout()
 	}
 

package/src/lifecycle/types.ts

@@ -3,6 +3,7 @@ import type { NotionDataSourceBindings } from "@notionhq/custom-blocks-protocol/
 import type { NotionBlockId } from "@notionhq/custom-blocks-protocol/ids.js"
 import type { CustomBlockManifest } from "@notionhq/custom-blocks-protocol/manifest.js"
 import type { ConnectMessage } from "@notionhq/custom-blocks-protocol/messages/connect.js"
+import type { CustomBlockInitErrorInfo } from "@notionhq/custom-blocks-protocol/messages/init.js"
 import type { CustomBlockPage } from "@notionhq/custom-blocks-protocol/pages/page.js"
 import type { NotionParent } from "@notionhq/custom-blocks-protocol/parent.js"
 import type { NotionTheme } from "@notionhq/custom-blocks-protocol/theme.js"
@@ -17,10 +18,12 @@ export type ConnectSuccessMessage = Extract<
 export type InitStatus =
 	/** The host is waiting for the sandbox's first `connect` message. */
 	| "waitingForConnect"
+	/** The host is waiting for `getInitialState` to prepare `init`. */
+	| "preparingInit"
 	/** The host is waiting for a matching `initResult` after it sends `init`. */
 	| "waitingForInitResult"
 	/** The host is initialized and can send live updates. */
-	| "success"
+	| "initialized"
 	/** The host is in an error state after an initialization failure or timeout. */
 	| "error"
 	/** The host is stopped and no longer sends or accepts sandbox messages. */
@@ -30,10 +33,6 @@ export type CustomBlockHostDataSourcesPayload = {
 	bindings: NotionDataSourceBindings
 }
 
-export type CustomBlockHostInitialDataSources =
-	| CustomBlockHostDataSourcesPayload
-	| ((connect: ConnectSuccessMessage) => CustomBlockHostDataSourcesPayload)
-
 export type CustomBlockHostInitialState = {
 	theme: NotionTheme
 	/**
@@ -44,11 +43,21 @@ export type CustomBlockHostInitialState = {
 	blockId: NotionBlockId
 	parent: NotionParent
 	page: CustomBlockPage
-	/**
-	 * The host's persisted manifest, when one exists. If omitted, the host
-	 * requires the sandbox to provide a manifest in `connect`.
-	 */
-	manifest?: CustomBlockManifest
-	dataSources: CustomBlockHostInitialDataSources
+	manifest: CustomBlockManifest
+	dataSources: CustomBlockHostDataSourcesPayload
 	currentUser: NotionUser
 }
+
+export type CustomBlockHostInitialStateResult =
+	| {
+			status: "success"
+			initialState: CustomBlockHostInitialState
+	  }
+	| {
+			status: "error"
+			error: CustomBlockInitErrorInfo
+	  }
+
+export type CustomBlockHostGetInitialState = (
+	connect: ConnectSuccessMessage,
+) => Promise<CustomBlockHostInitialStateResult>

package/src/messages/sandboxMessageRouter.ts

@@ -2,8 +2,8 @@ import type { ConnectMessage } from "@notionhq/custom-blocks-protocol/messages/c
 import type { CreatePageMessage } from "@notionhq/custom-blocks-protocol/messages/createPage.js"
 import type { GetPageMessage } from "@notionhq/custom-blocks-protocol/messages/getPage.js"
 import type { GetUserMessage } from "@notionhq/custom-blocks-protocol/messages/getUser.js"
-import type { HostToSandboxMessage } from "@notionhq/custom-blocks-protocol/messages/hostToSandbox.js"
 import type { InitResultMessage } from "@notionhq/custom-blocks-protocol/messages/initResult.js"
+import type { InvalidSandboxMessage } from "@notionhq/custom-blocks-protocol/messages/invalidSandboxMessage.js"
 import { isInitHandshakeMessage } from "@notionhq/custom-blocks-protocol/messages/isInitHandshakeMessage.js"
 import type { ListUsersMessage } from "@notionhq/custom-blocks-protocol/messages/listUsers.js"
 import type { OpenPageMessage } from "@notionhq/custom-blocks-protocol/messages/openPage.js"
@@ -35,7 +35,7 @@ export type SandboxMessageRouterOptions = {
 	sandboxOrigin: string
 	getInitStatus: () => InitStatus
 	emit: (direction: CustomBlockHostLogDirection, payload: unknown) => void
-	post: (message: HostToSandboxMessage) => void
+	sendNack: (message: InvalidSandboxMessage) => void
 	onInvalidConnect: InitializationController["handleInvalidConnect"]
 	onInvalidInitResult: InitializationController["handleInvalidInitResult"]
 	onConnect: (message: ConnectMessage) => void
@@ -61,7 +61,7 @@ export function createSandboxMessageRouter(
 		sandboxOrigin,
 		getInitStatus,
 		emit,
-		post,
+		sendNack,
 		onInvalidConnect,
 		onInvalidInitResult,
 		onConnect,
@@ -116,9 +116,9 @@ export function createSandboxMessageRouter(
 			})
 		}
 
-		// Other messages get NACKed indicating they were invalid and skipped
-		if (response.nack !== undefined) {
-			post(response.nack)
+		// Other messages get NACKed indicating they were invalid and skipped.
+		if (response.nack !== undefined && getInitStatus() !== "stopped") {
+			sendNack(response.nack)
 		}
 	}
 
@@ -139,7 +139,7 @@ export function createSandboxMessageRouter(
 		}
 
 		const message = parsed.output
-		if (!isInitHandshakeMessage(message) && getInitStatus() !== "success") {
+		if (!isInitHandshakeMessage(message) && getInitStatus() !== "initialized") {
 			emit("warn", {
 				message: `Ignored ${message.type} before initialization completed`,
 				payload: message,