From f064b9df0f79ee219ff235fbdd44196646847898 Mon Sep 17 00:00:00 2001
From: Tmk <i@tmk.im>
Date: Tue, 30 Jan 2024 14:36:52 +0800
Subject: [PATCH] fix(messaging): serialize error instance

---
 .../src/background/router/index.ts            |  6 ++++--
 apps/ai-assistant/src/content-scripts/app.tsx |  5 ++++-
 .../src/core/__tests__/index.spec.ts          | 20 +++++++++++++++++++
 packages/messaging/src/core/index.ts          |  3 ++-
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/apps/ai-assistant/src/background/router/index.ts b/apps/ai-assistant/src/background/router/index.ts
index e91a3e88..b7fb8482 100644
--- a/apps/ai-assistant/src/background/router/index.ts
+++ b/apps/ai-assistant/src/background/router/index.ts
@@ -10,12 +10,14 @@ const genAIAtom = atom(async (get) => {
   return apiKey ? new GoogleGenerativeAI(apiKey) : null;
 });
 
-store.sub(genAIAtom, () => {});
+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') {
-    const genAI = await store.get(genAIAtom);
     if (!genAI) return subscriber.error('GenAI is not initialized');
     const result = await genAI.getGenerativeModel({ model: 'gemini-pro' }).generateContentStream({
       contents: [
diff --git a/apps/ai-assistant/src/content-scripts/app.tsx b/apps/ai-assistant/src/content-scripts/app.tsx
index 14e3bdb4..d79f4b5c 100644
--- a/apps/ai-assistant/src/content-scripts/app.tsx
+++ b/apps/ai-assistant/src/content-scripts/app.tsx
@@ -85,7 +85,10 @@ export const App = () => {
       { prompt: 'Translate the following text to Chinese:\n' + text },
       {
         next: (token) => setContent((prev) => prev + token),
-        error: () => setIsLoading(false),
+        error: (err) => {
+          console.log('Translate Error', { err });
+          setIsLoading(false);
+        },
         complete: () => setIsLoading(false),
       }
     );
diff --git a/packages/messaging/src/core/__tests__/index.spec.ts b/packages/messaging/src/core/__tests__/index.spec.ts
index 53818b86..f2019a13 100644
--- a/packages/messaging/src/core/__tests__/index.spec.ts
+++ b/packages/messaging/src/core/__tests__/index.spec.ts
@@ -249,3 +249,23 @@ it('should support relay stream', async () => {
   expectMessagingIsNotLeaked(relay2);
   expectMessagingIsNotLeaked(sender);
 });
+
+it('should serialize error message', async () => {
+  const { port1, port2 } = new MessageChannel();
+
+  const receiver = createMessaging(fromMessagePort(port1), {
+    async onStream(_message, _subscriber) {
+      throw new Error('Internal Error');
+    },
+  });
+  const sender = createMessaging(fromMessagePort(port2));
+
+  const { promise, resolve, reject } = withResolvers<void>();
+  sender.stream(null, { error: reject, complete: resolve });
+
+  await expect(promise).rejects.toBeTypeOf('string');
+  await expect(promise).rejects.toThrow('Internal Error');
+
+  expectMessagingIsNotLeaked(receiver);
+  expectMessagingIsNotLeaked(sender);
+});
diff --git a/packages/messaging/src/core/index.ts b/packages/messaging/src/core/index.ts
index 36401f75..c2f5cb8c 100644
--- a/packages/messaging/src/core/index.ts
+++ b/packages/messaging/src/core/index.ts
@@ -139,7 +139,8 @@ export function createMessaging(port: Port, options?: CreateMessagingOptions): M
           const cleanup = await onStream.call(context, message.d, observer);
           processingStreamCleanups.set(message.i, cleanup);
         } catch (error) {
-          terminate({ error });
+          // Error is not serializable in `chrome.runtime.Port`
+          terminate({ error: error instanceof Error ? error.message : error });
         }
         break;
       }