1 files changed, 170 insertions, 83 deletions

diff --git a/src/hooks/useRequests.tsx b/src/hooks/useRequests.tsx
index 1361949..3135aa2 100644
--- a/src/hooks/useRequests.tsx
+++ b/src/hooks/useRequests.tsx
@@ -1,68 +1,32 @@
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { getHost } from "~/utils";
-
-export type Headers = [string, string][];
-export type Method =
-  | "GET"
-  | "HEAD"
-  | "POST"
-  | "PUT"
-  | "DELETE"
-  | "CONNECT"
-  | "OPTIONS"
-  | "TRACE"
-  | "PATCH";
-
-export interface RequestPayload {
-  id: string;
-  timestamp: string;
-  body: string;
-  headers: Headers;
-  method: Method;
-  path: string;
-}
-
-interface Request {
-  type: "request";
-  payload: RequestPayload;
-}
-
-export interface ResponsePayload {
-  id: string;
-  timing: number;
-  body: string;
-  headers: Headers;
-  status: number;
-}
-
-interface Response {
-  type: "response";
-  payload: ResponsePayload;
-}
-
-interface Historic {
-  type: "historic";
-  payload: (Request | Response)[];
-}
-
-export interface RequestResponse {
-  request: RequestPayload;
-  response?: ResponsePayload;
-}
-
-export enum ReadyState {
-  CONNECTING = 0,
-  OPEN = 1,
-  CLOSING = 2,
-  CLOSED = 3,
-}
+import {
+  Call,
+  Frame,
+  Frames,
+  Historic,
+  ReadyState,
+  Request,
+  RequestResponseType,
+  Requests,
+  Response,
+  ResponsePayload,
+  Responses,
+  WebsocketConnect,
+  WebsocketConnected,
+  WebsocketConnectedPayload,
+  WebsocketDisconnectPayload,
+  WebsocketInboundPayload,
+  WebsocketOutboundPayload,
+  WebsocketType,
+} from "~/types";
 
 export interface useRequestsProps {
   onConnect: () => Promise<void>;
 }
 
 export interface UseRequests {
-  calls: RequestResponse[];
+  calls: Call[];
   readyState: ReadyState;
   clear: () => void;
 }
@@ -73,8 +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 [requests, setRequests] = useState<Requests>([]);
+  const [responses, setResponses] = useState<Responses>({});
+  const [websocketFrames, setWebsocketFrames] = useState<Frames>({});
 
   const connect = useCallback(
     () => new WebSocket(`ws://${wsHost}/inspect/`),
@@ -96,36 +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;
+        | 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 };
+