refactor: migrate background router to trpc

apps/ai-assistant/package.json

@@ -10,13 +10,16 @@
     "@douyinfe/semi-icons": "^2.51.3",
     "@douyinfe/semi-ui": "^2.51.3",
     "@google/generative-ai": "^0.1.3",
+    "@trpc/client": "11.0.0-next-beta.248",
+    "@trpc/server": "11.0.0-next-beta.248",
     "@webx-kit/messaging": "workspace:^",
     "@webx-kit/runtime": "workspace:^",
     "@webx-kit/storage": "workspace:^",
     "clsx": "^2.1.0",
     "jotai": "^2.6.2",
     "react": "^18.2.0",
-    "react-dom": "^18.2.0"
+    "react-dom": "^18.2.0",
+    "zod": "^3.22.4"
   },
   "devDependencies": {
     "@modern-js/app-tools": "^2.46.1",

apps/ai-assistant/src/background/index.ts

@@ -1 +1,6 @@
-import './router';
+import { createTrpcHandler } from '@webx-kit/messaging/background';
+import { appRouter } from './router';
+
+createTrpcHandler({
+  router: appRouter,
+});

apps/ai-assistant/src/background/router/index.ts

@@ -1,35 +1,35 @@
-import { GoogleGenerativeAI } from '@google/generative-ai';
-import { setStreamHandler } from '@webx-kit/messaging/background';
-import { apiKeyAtom } from '@/hooks/atoms/config';
-import { atom, getDefaultStore } from 'jotai';
+import { initTRPC } from '@trpc/server';
+import { z } from 'zod';
+import { observable } from '@trpc/server/observable';
+import { genAI } from '@/background/shared';
 
-const store = getDefaultStore();
+const t = initTRPC.create();
 
-const genAIAtom = atom(async (get) => {
-  const apiKey = await get(apiKeyAtom);
-  return apiKey ? new GoogleGenerativeAI(apiKey) : null;
-});
-
-let genAI: GoogleGenerativeAI | null = null;
-const updateGenAIInstance = () => store.get(genAIAtom).then((instance) => (genAI = instance));
-updateGenAIInstance();
-store.sub(genAIAtom, updateGenAIInstance);
-
-setStreamHandler(async (message, subscriber) => {
-  const { data } = message;
-  if (data && typeof data === 'object' && 'prompt' in data && typeof data.prompt === 'string') {
-    if (!genAI) return subscriber.error('GenAI is not initialized');
-    const result = await genAI.getGenerativeModel({ model: 'gemini-pro' }).generateContentStream({
-      contents: [
-        {
-          role: 'user',
-          parts: [{ text: data.prompt }],
-        },
-      ],
+export const appRouter = t.router({
+  generateContentStream: t.procedure.input(z.object({ prompt: z.string() })).subscription(({ input }) => {
+    return observable((observer) => {
+      if (!genAI) return observer.error('GenAI is not initialized');
+      let isUnsubscribed = false;
+      (async () => {
+        const result = await genAI.getGenerativeModel({ model: 'gemini-pro' }).generateContentStream({
+          contents: [
+            {
+              role: 'user',
+              parts: [{ text: input.prompt }],
+            },
+          ],
+        });
+        for await (const token of result.stream || []) {
+          if (isUnsubscribed) break;
+          observer.next(token.text());
+        }
+        observer.complete();
+      })();
+      return () => {
+        isUnsubscribed = true;
+      };
     });
-    for await (const token of result.stream || []) {
-      subscriber.next(token.text());
-    }
-    subscriber.complete();
-  }
+  }),
 });
+
+export type AppRouter = typeof appRouter;

apps/ai-assistant/src/background/shared.ts

@@ -0,0 +1,15 @@
+import { GoogleGenerativeAI } from '@google/generative-ai';
+import { apiKeyAtom } from '@/hooks/atoms/config';
+import { atom, getDefaultStore } from 'jotai';
+
+export const store = getDefaultStore();
+
+const genAIAtom = atom(async (get) => {
+  const apiKey = await get(apiKeyAtom);
+  return apiKey ? new GoogleGenerativeAI(apiKey) : null;
+});
+
+export let genAI: GoogleGenerativeAI | null = null;
+const updateGenAIInstance = () => store.get(genAIAtom).then((instance) => (genAI = instance));
+updateGenAIInstance();
+store.sub(genAIAtom, updateGenAIInstance);

apps/ai-assistant/src/content-scripts/app.tsx

@@ -8,14 +8,17 @@ import {
   isSelectionValid,
   rangeToReference,
 } from '@webx-kit/runtime/content-scripts';
-import { client } from '@webx-kit/messaging/content-script';
+import { createTrpcHandler } from '@webx-kit/messaging/content-script';
 import clsx from 'clsx';
+import type { AppRouter } from '@/background/router';
 import { Provider } from './features/provider';
 import './global.less';
 
 // hack for missing button loading rotate keyframes
 Spin.name;
 
+const { client } = createTrpcHandler<AppRouter>({});
+
 export const App = () => {
   const [visible, setVisible] = useState(false);
   const [rootStyle, setRootStyle] = useState<React.CSSProperties>();
@@ -81,15 +84,21 @@ export const App = () => {
 
     setContent('');
     setIsLoading(true);
-    client.stream(
-      { prompt: 'Translate the following text to Chinese:\n' + text },
+    client.generateContentStream.subscribe(
+      {
+        prompt: 'Translate the following text to Chinese:\n' + text,
+      },
       {
-        next: (token) => setContent((prev) => prev + token),
-        error: (err) => {
+        onData(token) {
+          setContent((prev) => prev + token);
+        },
+        onError(err) {
           console.log('Translate Error', { err });
           setIsLoading(false);
         },
-        complete: () => setIsLoading(false),
+        onComplete() {
+          setIsLoading(false);
+        },
       }
     );
   };
@@ -101,12 +110,21 @@ export const App = () => {
 
     setContent('');
     setIsLoading(true);
-    client.stream(
-      { prompt: 'Summarize the following text to Chinese:\n' + text },
+    client.generateContentStream.subscribe(
+      {
+        prompt: 'Summarize the following text to Chinese:\n' + text,
+      },
       {
-        next: (token) => setContent((prev) => prev + token),
-        error: () => setIsLoading(false),
-        complete: () => setIsLoading(false),
+        onData(token) {
+          setContent((prev) => prev + token);
+        },
+        onError(err) {
+          console.log('Translate Error', { err });
+          setIsLoading(false);
+        },
+        onComplete() {
+          setIsLoading(false);
+        },
       }
     );
   };

packages/messaging/demo/background/index.ts

@@ -1,23 +1,24 @@
-import { connections, setRequestHandler, setStreamHandler } from '@/background';
-
-// @ts-expect-error
-globalThis.__webxConnections = connections;
+import { createCustomHandler } from '@/background';
 
 const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
 
-setRequestHandler(async (message) => {
-  return {
-    reply: 'background',
-    data: message.data,
-  };
+const { connections } = createCustomHandler({
+  async requestHandler(message) {
+    return {
+      reply: 'background',
+      data: message.data,
+    };
+  },
+  async streamHandler(_message, subscriber) {
+    await sleep(50);
+    subscriber.next(1);
+    await sleep(50);
+    subscriber.next(2);
+    await sleep(50);
+    subscriber.next(3);
+    subscriber.complete();
+  },
 });
 
-setStreamHandler(async (_message, subscriber) => {
-  await sleep(50);
-  subscriber.next(1);
-  await sleep(50);
-  subscriber.next(2);
-  await sleep(50);
-  subscriber.next(3);
-  subscriber.complete();
-});
+// @ts-expect-error
+globalThis.__webxConnections = connections;

packages/messaging/demo/content-scripts/index.ts

@@ -1,23 +1,24 @@
-import { client, setRequestHandler, setStreamHandler } from '@/content-script';
-
-// @ts-expect-error
-globalThis.__client = client;
-
-setRequestHandler((message) => {
-  return {
-    reply: 'content-script',
-    data: message.data,
-  };
-});
+import { createCustomHandler } from '@/content-script';
 
 const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
 
-setStreamHandler(async (_message, subscriber) => {
-  await sleep(50);
-  subscriber.next('content-script 1');
-  await sleep(50);
-  subscriber.next('content-script 2');
-  await sleep(50);
-  subscriber.next('content-script 3');
-  subscriber.complete();
+const { messaging } = createCustomHandler({
+  requestHandler(message) {
+    return {
+      reply: 'content-script',
+      data: message.data,
+    };
+  },
+  async streamHandler(_message, subscriber) {
+    await sleep(50);
+    subscriber.next('content-script 1');
+    await sleep(50);
+    subscriber.next('content-script 2');
+    await sleep(50);
+    subscriber.next('content-script 3');
+    subscriber.complete();
+  },
 });
+
+// @ts-expect-error
+globalThis.__client = messaging;

packages/messaging/demo/pages/options/index.ts

@@ -1,24 +1,25 @@
-import { client, setRequestHandler, setStreamHandler } from '@/options';
+import { createCustomHandler } from '@/options';
 import '../../global.css';
 
-// @ts-expect-error
-globalThis.__client = client;
-
-setRequestHandler((message) => {
-  return {
-    reply: 'options',
-    data: message.data,
-  };
-});
-
 const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
 
-setStreamHandler(async (_message, subscriber) => {
-  await sleep(50);
-  subscriber.next('options 1');
-  await sleep(50);
-  subscriber.next('options 2');
-  await sleep(50);
-  subscriber.next('options 3');
-  subscriber.complete();
+const { messaging } = createCustomHandler({
+  requestHandler(message) {
+    return {
+      reply: 'options',
+      data: message.data,
+    };
+  },
+  async streamHandler(_message, subscriber) {
+    await sleep(50);
+    subscriber.next('options 1');
+    await sleep(50);
+    subscriber.next('options 2');
+    await sleep(50);
+    subscriber.next('options 3');
+    subscriber.complete();
+  },
 });
+
+// @ts-expect-error
+globalThis.__client = messaging;

packages/messaging/demo/pages/popup/index.ts

@@ -1,24 +1,25 @@
-import { client, setRequestHandler, setStreamHandler } from '@/popup';
+import { createCustomHandler } from '@/popup';
 import '../../global.css';
 
-// @ts-expect-error
-globalThis.__client = client;
-
-setRequestHandler((message) => {
-  return {
-    reply: 'popup',
-    data: message.data,
-  };
-});
-
 const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
 
-setStreamHandler(async (_message, subscriber) => {
-  await sleep(50);
-  subscriber.next('popup 1');
-  await sleep(50);
-  subscriber.next('popup 2');
-  await sleep(50);
-  subscriber.next('popup 3');
-  subscriber.complete();
+const { messaging } = createCustomHandler({
+  requestHandler(message) {
+    return {
+      reply: 'popup',
+      data: message.data,
+    };
+  },
+  async streamHandler(_message, subscriber) {
+    await sleep(50);
+    subscriber.next('popup 1');
+    await sleep(50);
+    subscriber.next('popup 2');
+    await sleep(50);
+    subscriber.next('popup 3');
+    subscriber.complete();
+  },
 });
+
+// @ts-expect-error
+globalThis.__client = messaging;

packages/messaging/e2e/basic.spec.ts

@@ -1,14 +1,14 @@
 import { setupStaticServer } from '@webx-kit/test-utils/playwright';
 import { expect, test } from './context';
-import type { connections } from '@/background';
-import type { ensureClient } from '@/client-base';
+import type { Messaging } from '@/core';
+import type { WrappedMessaging } from '@/client-base';
 
 const getWebpageURL = setupStaticServer(test);
 
 declare module globalThis {
   /** only in background */
-  const __webxConnections: typeof connections | undefined;
-  const __client: ReturnType<typeof ensureClient>;
+  const __webxConnections: Set<Messaging> | undefined;
+  const __client: WrappedMessaging;
 }
 
 test('Background', async ({ background }) => {