From 0617a1a93697a44dde77d67a92eed86b36d0f55a Mon Sep 17 00:00:00 2001
From: JongChan Choi <jong@chan.moe>
Date: Wed, 24 Aug 2022 15:41:10 +0900
Subject: [PATCH] add createWrpAtomSetFromSourceChannelAtom

---
 src/jotai/iframe.ts |  8 +++++---
 src/jotai/index.ts  | 22 ++++++++++++++++------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/jotai/iframe.ts b/src/jotai/iframe.ts
index 1fd9b7b..c66a369 100644
--- a/src/jotai/iframe.ts
+++ b/src/jotai/iframe.ts
@@ -1,6 +1,7 @@
 import { useEffect } from "react";
 import { atom, PrimitiveAtom, useSetAtom } from "jotai";
-import { createWrpAtomSet, WrpAtomSet } from "./index.ts";
+import { createWrpAtomSetFromSourceChannelAtom, WrpAtomSet } from "./index.ts";
+import { createWrpChannel as createWrpChannelFn } from "../channel.ts";
 import {
   default as useWrpIframeSocket,
   UseWrpIframeSocketResult,
@@ -8,14 +9,15 @@ import {
 
 export function useIframeWrpAtomSetUpdateEffect(
   primitiveWrpAtomSetAtom: PrimitiveAtom<WrpAtomSet>,
+  createWrpChannel = createWrpChannelFn,
 ): UseWrpIframeSocketResult {
   const setIframeAtomSet = useSetAtom(primitiveWrpAtomSetAtom);
   const useWrpIframeSocketResult = useWrpIframeSocket();
   const { socket } = useWrpIframeSocketResult;
   useEffect(() => {
     if (!socket) return;
-    const socketAtom = atom(async () => socket);
-    setIframeAtomSet(createWrpAtomSet(socketAtom));
+    const channelAtom = atom(createWrpChannel(socket));
+    setIframeAtomSet(createWrpAtomSetFromSourceChannelAtom(channelAtom));
   }, [socket]);
   return useWrpIframeSocketResult;
 }
diff --git a/src/jotai/index.ts b/src/jotai/index.ts
index 67c6440..c69d775 100644
--- a/src/jotai/index.ts
+++ b/src/jotai/index.ts
@@ -17,21 +17,31 @@ export interface WrpAtomSet {
   clientImplAtom: ClientImplAtom;
 }
 export function createWrpAtomSet(socketAtom: SocketAtom): WrpAtomSet {
+  const channelAtom: ChannelAtom = atom((get) => {
+    const socket = get(socketAtom);
+    if (!socket) return;
+    return createWrpChannel(socket);
+  });
+  return createWrpAtomSetFromSourceChannelAtom(channelAtom);
+}
+
+export function createWrpAtomSetFromSourceChannelAtom(
+  sourceChannelAtom: ChannelAtom,
+): WrpAtomSet {
   interface ChannelAndGuest {
     channel: WrpChannel;
     guest: Promise<WrpGuest>;
   }
   const channelAndGuestAtom = atom<Promise<ChannelAndGuest | undefined>>(
     async (get) => {
-      const socket = get(socketAtom);
-      if (!socket) return;
-      const realChannel = createWrpChannel(socket);
+      const sourceChannel = get(sourceChannelAtom);
+      if (!sourceChannel) return;
       const listeners: ((message?: WrpMessage) => void)[] = [];
       const guest = createWrpGuest({
         channel: {
-          ...realChannel,
+          ...sourceChannel,
           async *listen() {
-            for await (const message of realChannel.listen()) {
+            for await (const message of sourceChannel.listen()) {
               yield message;
               for (const listener of listeners) listener(message);
               listeners.length = 0;
@@ -40,7 +50,7 @@ export function createWrpAtomSet(socketAtom: SocketAtom): WrpAtomSet {
         },
       });
       const channel: WrpChannel = {
-        ...realChannel,
+        ...sourceChannel,
         async *listen() {
           while (true) {
             const message = await new Promise<WrpMessage | undefined>(