1 files changed, 162 insertions, 33 deletions

diff --git a/src/hooks/useRequests.tsx b/src/hooks/useRequests.tsx
index dfc3b80..3135aa2 100644
--- a/src/hooks/useRequests.tsx
+++ b/src/hooks/useRequests.tsx
@@ -1,13 +1,23 @@
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { getHost } from "~/utils";
 import {
+  Call,
+  Frame,
+  Frames,
   Historic,
   ReadyState,
   Request,
-  RequestPayload,
-  RequestResponse,
+  RequestResponseType,
+  Requests,
   Response,
   ResponsePayload,
+  Responses,
+  WebsocketConnect,
+  WebsocketConnected,
+  WebsocketConnectedPayload,
+  WebsocketDisconnectPayload,
+  WebsocketInboundPayload,
+  WebsocketOutboundPayload,
   WebsocketType,
 } from "~/types";
 
@@ -16,7 +26,7 @@ export interface useRequestsProps {
 }
 
 export interface UseRequests {
-  calls: RequestResponse[];
+  calls: Call[];
   readyState: ReadyState;
   clear: () => void;
 }
@@ -27,10 +37,9 @@ export default function useRequests({
   const wsHost = useMemo(getHost, []);
 
   const [initialConnection, setInitialConnection] = useState(true);
-  const [requests, setRequests] = useState<RequestPayload[]>([]);
-  const [responses, setResponses] = useState<ResponsePayload[]>([]);
-  const [websocketFrames, setWebsocketFrames] =
-    useState<WebsocketType["payload"]>();
+  const [requests, setRequests] = useState<Requests>([]);
+  const [responses, setResponses] = useState<Responses>({});
+  const [websocketFrames, setWebsocketFrames] = useState<Frames>({});
 
   const connect = useCallback(
     () => new WebSocket(`ws://${wsHost}/inspect/`),
@@ -52,38 +61,140 @@ export default function useRequests({
       setInitialConnection(false);
       setReadyState(ws.readyState);
     };
-    const onMessage = ({ data }) => {
+    const onMessage = ({ data }: { data: string }) => {
       const { type, payload } = JSON.parse(data) as
         | Historic
-        | Request
-        | Response;
-
-      console.debug(type, payload);
+        | RequestResponseType
+        | WebsocketType;
 
       switch (type) {
         case "historic":
           if (initialConnection) {
-            const requests = (payload as (Request | Response)[]).filter(
-              ({ type }) => type === "request"
-            );
-            const responses = (payload as (Request | Response)[]).filter(
-              ({ type }) => type === "response"
-            );
+            const requests = (
+              payload as (RequestResponseType | WebsocketType)[]
+            )
+              .filter(
+                ({ type }) => type === "request" || type == "websocket_connect"
+              )
+              .map((payload) => {
+                if (
+                  (payload as WebsocketConnect).payload.method === undefined
+                ) {
+                  (payload as Request | WebsocketConnect).payload.method =
+                    "GET";
+                }
+
+                return payload;
+              });
+            const responses = (
+              payload as (RequestResponseType | WebsocketType)[]
+            )
+              .filter(
+                ({ type }) =>
+                  type === "response" || type == "websocket_connected"
+              )
+              .map(({ type, ...item }) => {
+                if (type == "websocket_connected") {
+                  (item.payload as WebsocketConnectedPayload).status = 101;
+                }
+
+                return { type, ...item };
+              })
+              .reduce<{ [id: string]: Response | WebsocketConnected }>(
+                (out, item) => ({
+                  ...out,
+                  [item.payload.id]: item as Response | WebsocketConnected,
+                }),
+                {}
+              );
+            const frames = (payload as (RequestResponseType | WebsocketType)[])
+              .filter(
+                ({ type }) =>
+                  type == "websocket_inbound" ||
+                  type == "websocket_outbound" ||
+                  type == "websocket_disconnect"
+              )
+              .reduce<Frames>((out, item) => {
+                if (!out.hasOwnProperty(item.payload.id)) {
+                  out[item.payload.id] = [];
+                }
+
+                out[item.payload.id].push(item as Frame);
+                return out;
+              }, {});
             setRequests((rqs) => [
               ...rqs,
-              ...requests.map(({ payload }) => payload as RequestPayload),
+              ...requests.map(
+                (payload) => payload as Request | WebsocketConnect
+              ),
             ]);
-            setResponses((rps) => [
+            setResponses((rps) => ({
               ...rps,
-              ...responses.map(({ payload }) => payload as ResponsePayload),
-            ]);
+              ...responses,
+            }));
+            setWebsocketFrames((frms) => ({
+              ...frms,
+              ...frames,
+            }));
           }
           break;
         case "request":
-          setRequests((rqs) => [...rqs, payload as RequestPayload]);
+        case "websocket_connect":
+          setRequests((rqs) => [
+            ...rqs,
+            {
+              type,
+              payload:
+                type === "request"
+                  ? payload
+                  : {
+                      ...payload,
+                      method: "GET",
+                    },
+            } as Request | WebsocketConnect,
+          ]);
           break;
         case "response":
-          setResponses((rps) => [...rps, payload as ResponsePayload]);
+        case "websocket_connected":
+          if (type == "websocket_connected") {
+            (payload as WebsocketConnectedPayload).status = 101;
+          }
+          setResponses((rps) => ({
+            ...rps,
+            [(payload as ResponsePayload | WebsocketConnectedPayload).id]: {
+              type,
+              payload,
+            } as Response | WebsocketConnected,
+          }));
+          break;
+        case "websocket_inbound":
+        case "websocket_outbound":
+        case "websocket_disconnect":
+          setWebsocketFrames((frms) => {
+            const id = (
+              payload as
+                | WebsocketInboundPayload
+                | WebsocketOutboundPayload
+                | WebsocketDisconnectPayload
+            ).id;
+
+            const newFrms = { ...frms };
+
+            if (!newFrms.hasOwnProperty(id)) {
+              newFrms[id] = [];
+            }
+
+            // @ts-ignore
+            newFrms[id].push({
+              type,
+              payload: payload as
+                | WebsocketInboundPayload
+                | WebsocketOutboundPayload
+                | WebsocketDisconnectPayload,
+            });
+
+            return newFrms;
+          });
           break;
       }
     };
@@ -99,19 +210,37 @@ export default function useRequests({
     };
   }, [ws]);
 
+  const calls = useMemo<Call[]>(
+    () =>
+      requests.map((request) => {
+        switch (request.type) {
+          case "request":
+            return {
+              type: "RequestResponse",
+              request: request.payload,
+              response: responses[request.payload.id]
+                ?.payload as ResponsePayload,
+            };
+          case "websocket_connect":
+            return {
+              type: "WebsocketConnection",
+              request