feat: add separate protocol for actor inspect

rivet-gg · Jan 28, 2025 · 1e1fc74 · 1e1fc74
1 parent ac2e9de
commit 1e1fc74
Show file tree

Hide file tree

Showing 13 changed files with 489 additions and 238 deletions.
diff --git a/frontend/apps/hub/src/domains/project/components/actors/actors-actor-details.tsx b/frontend/apps/hub/src/domains/project/components/actors/actors-actor-details.tsx
@@ -41,7 +41,11 @@ export function ActorsActorDetails({
 	const { data } = useSuspenseQuery(actorQueryOptions(props));
 	return (
 		<ActorDetailsSettingsProvider>
-			<ActorWorkerContextProvider enabled={!data.destroyedAt} {...props}>
+			<ActorWorkerContextProvider
+				enabled={!data.destroyedAt}
+				endpoint={data.endpoint}
+				{...props}
+			>
 				<div className="flex flex-col h-full flex-1 pt-2">
 					<Tabs
 						value={tab || "logs"}

diff --git a/frontend/apps/hub/src/domains/project/components/actors/worker/actor-repl.worker.ts b/frontend/apps/hub/src/domains/project/components/actors/worker/actor-repl.worker.ts
@@ -1,4 +1,4 @@
-import { type ActorHandle, Client } from "@rivet-gg/actor-client";
+import { ActorHandleRaw } from "@rivet-gg/actor-client";
 import type { InspectRpcResponse } from "@rivet-gg/actor-protocol/ws/to_client";
 import { fromJs } from "esast-util-from-js";
 import { toJs } from "estree-util-to-js";
@@ -104,14 +104,7 @@ const createConsole = (id: string) => {
 	);
 };
 
-interface InspectableActor {
-	internal_setState(rpc: this, state: Record<string, unknown>): Promise<void>;
-	internal_inspect(): Promise<InspectRpcResponse>;
-}
-
-let init:
-	| null
-	| ({ handle: ActorHandle<InspectableActor> } & InspectRpcResponse);
+let init: null | ({ handle: ActorHandleRaw } & InspectRpcResponse);
 
 addEventListener("message", async (event) => {
 	const { success, data } = MessageSchema.safeParse(event.data);
@@ -130,11 +123,13 @@ addEventListener("message", async (event) => {
 		}
 
 		try {
-			const cl = new Client(data.managerUrl);
-			const handle = await Promise.race([
-				cl.getWithId<InspectableActor>(data.actorId, {}),
-				wait(5000).then(() => undefined),
-			]);
+			const handle = new ActorHandleRaw(
+				`${data.endpoint}/__inspect`,
+				undefined,
+				"cbor",
+			);
+
+			handle.connect();
 
 			if (!handle) {
 				respond({
@@ -146,7 +141,7 @@ addEventListener("message", async (event) => {
 			}
 
 			const inspect = await Promise.race([
-				handle.internal_inspect(),
+				handle.rpc<[], InspectRpcResponse>("inspect"),
 				wait(5000).then(() => undefined),
 			]);
 
@@ -208,21 +203,17 @@ addEventListener("message", async (event) => {
 				data: formatted,
 			});
 
-			const rpcs = Object.fromEntries(
-				actor.rpcs.map(
-					(rpc) =>
-						[
-							rpc,
-							actor.handle[rpc as keyof typeof actor.handle],
-						] as const,
-				),
+			const exposedActor = Object.fromEntries(
+				init?.rpcs.map((rpc) => [
+					rpc,
+					actor.handle.rpc.bind(actor.handle, rpc),
+				]) ?? [],
 			);
 
 			const evaluated = await evaluateCode(data.data, {
 				console: createConsole(data.id),
 				wait,
-				actor: actor.handle,
-				...rpcs,
+				actor: exposedActor,
 			});
 			return respond({
 				type: "result",
@@ -250,7 +241,7 @@ addEventListener("message", async (event) => {
 
 		try {
 			const state = JSON.parse(data.data);
-			await actor.handle.internal_setState(state);
+			await actor.handle.rpc("setState", state);
 			return respond({
 				type: "state-change",
 				data: {

diff --git a/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-container.ts b/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-container.ts
@@ -1,7 +1,4 @@
-import {
-	actorManagerUrlQueryOptions,
-	actorQueryOptions,
-} from "@/domains/project/queries";
+import { actorQueryOptions } from "@/domains/project/queries";
 import { queryClient } from "@/queries/global";
 import ActorWorker from "./actor-repl.worker?worker";
 import {
@@ -66,11 +63,13 @@ export class ActorWorkerContainer {
 		projectNameId,
 		environmentNameId,
 		actorId,
+		endpoint,
 		signal,
 	}: {
 		projectNameId: string;
 		environmentNameId: string;
 		actorId: string;
+		endpoint: string;
 		signal: AbortSignal;
 	}) {
 		this.terminate();
@@ -79,16 +78,6 @@ export class ActorWorkerContainer {
 		this.#state.status = { type: "pending" };
 		this.#update();
 		try {
-			// To check if the actor is supported, first we need to get the actor manager URL
-			// if there's no manager URL, we can't support the actor
-			const managerUrl = await queryClient.fetchQuery(
-				actorManagerUrlQueryOptions({
-					projectNameId,
-					environmentNameId,
-				}),
-			);
-			signal.throwIfAborted();
-
 			// If we have the manager URL, next we need to check actor's runtime
 			const { actor } = await queryClient.fetchQuery(
 				actorQueryOptions({
@@ -118,7 +107,7 @@ export class ActorWorkerContainer {
 			const worker = new ActorWorker({ name: `actor-${actorId}` });
 			signal.throwIfAborted();
 			// now worker needs to check if the actor is supported
-			this.#setupWorker(worker, { actorId, managerUrl });
+			this.#setupWorker(worker, { actorId, endpoint });
 			signal.throwIfAborted();
 			return worker;
 		} catch (e) {

diff --git a/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-context.tsx b/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-context.tsx
@@ -24,6 +24,7 @@ interface ActorWorkerContextProviderProps {
 	actorId: string;
 	projectNameId: string;
 	environmentNameId: string;
+	endpoint?: string;
 	enabled?: boolean;
 	children: ReactNode;
 }
@@ -32,6 +33,7 @@ export const ActorWorkerContextProvider = ({
 	children,
 	actorId,
 	enabled,
+	endpoint,
 	projectNameId,
 	environmentNameId,
 }: ActorWorkerContextProviderProps) => {
@@ -43,11 +45,12 @@ export const ActorWorkerContextProvider = ({
 	useEffect(() => {
 		const ctrl = new AbortController();
 
-		if (enabled) {
+		if (enabled && endpoint) {
 			container.init({
 				projectNameId,
 				environmentNameId,
 				actorId,
+				endpoint,
 				signal: ctrl.signal,
 			});
 		}
@@ -56,7 +59,7 @@ export const ActorWorkerContextProvider = ({
 			ctrl.abort();
 			container.terminate();
 		};
-	}, [actorId, projectNameId, environmentNameId, enabled]);
+	}, [actorId, projectNameId, environmentNameId, endpoint, enabled]);
 
 	return (
 		<ActorWorkerContext.Provider value={container}>

diff --git a/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-schema.ts b/frontend/apps/hub/src/domains/project/components/actors/worker/actor-worker-schema.ts
@@ -14,7 +14,7 @@ const CodeMessageSchema = z.object({
 });
 const InitMessageSchema = z.object({
 	type: z.literal("init"),
-	managerUrl: z.string(),
+	endpoint: z.string(),
 	actorId: z.string(),
 });
 

diff --git a/frontend/apps/hub/src/domains/project/queries/actors/query-options.ts b/frontend/apps/hub/src/domains/project/queries/actors/query-options.ts
@@ -151,6 +151,7 @@ export const actorQueryOptions = ({
 					(arg) => arg !== "",
 				),
 			},
+			endpoint: createActorEndpoint(data.actor.network),
 		}),
 	});
 };
@@ -442,7 +443,9 @@ export const actorRegionQueryOptions = ({
 	});
 };
 
-const createActorEndpoint = (network: Rivet.actor.Network) => {
+const createActorEndpoint = (
+	network: Rivet.actor.Network,
+): string | undefined => {
 	const http = Object.values(network.ports).find(
 		(port) => port.protocol === "http" || port.protocol === "https",
 	);

diff --git a/frontend/packages/components/src/ui/button.tsx b/frontend/packages/components/src/ui/button.tsx
@@ -75,6 +75,8 @@ const Button = React.forwardRef<HTMLButtonElement, ButtonProps>(
 	) => {
 		const C = asChild ? Slot : "button";
 
+		const isIcon = size?.includes("icon");
+
 		return (
 			<C
 				className={cn(
@@ -93,9 +95,7 @@ const Button = React.forwardRef<HTMLButtonElement, ButtonProps>(
 				) : startIcon ? (
 					React.cloneElement(startIcon, startIcon.props)
 				) : null}
-				{!size?.includes("icon") && isLoading ? null : (
-					<Slottable>{children}</Slottable>
-				)}
+				{isIcon && isLoading ? null : <Slottable>{children}</Slottable>}
 				{endIcon ? React.cloneElement(endIcon, endIcon.props) : null}
 			</C>
 		);