diff --git a/src/frontend/src/App.css b/src/frontend/src/App.css
index 5564ca084..125282e2c 100644
--- a/src/frontend/src/App.css
+++ b/src/frontend/src/App.css
@@ -156,3 +156,8 @@ body {
 .react-flow__edge .react-flow__edge-path {
   stroke: var(--connection) !important;
 }
+
+.ag-react-container {
+  width: 100%;
+  height: 100%;
+}
diff --git a/src/frontend/src/components/tableComponent/components/loadingOverlay/index.tsx b/src/frontend/src/components/tableComponent/components/loadingOverlay/index.tsx
new file mode 100644
index 000000000..9df9144e0
--- /dev/null
+++ b/src/frontend/src/components/tableComponent/components/loadingOverlay/index.tsx
@@ -0,0 +1,9 @@
+import Loading from "../../../ui/loading";
+
+export default function LoadingOverlay() {
+  return (
+    <div className=" flex h-full w-full items-center justify-center bg-background align-middle">
+      <Loading />
+    </div>
+  );
+}
diff --git a/src/frontend/src/components/tableComponent/index.tsx b/src/frontend/src/components/tableComponent/index.tsx
index 917919f51..fdd4a5f1e 100644
--- a/src/frontend/src/components/tableComponent/index.tsx
+++ b/src/frontend/src/components/tableComponent/index.tsx
@@ -13,6 +13,7 @@ import ForwardedIconComponent from "../genericIconComponent";
 import { Alert, AlertDescription, AlertTitle } from "../ui/alert";
 import TableOptions from "./components/TableOptions";
 import resetGrid from "./utils/reset-grid-columns";
+import Loading from "../ui/loading";
 
 interface TableComponentProps extends AgGridReactProps {
   columnDefs: NonNullable<AgGridReactProps["columnDefs"]>;
@@ -76,7 +77,7 @@ const TableComponent = forwardRef<
     const dark = useDarkStore((state) => state.dark);
     const initialColumnDefs = useRef(colDef);
     const [columnStateChange, setColumnStateChange] = useState(false);
-
+    const storeReference = props.columnDefs.map((e) => e.headerName).join("_");
     const makeLastColumnNonResizable = (columnDefs) => {
       columnDefs.forEach((colDef, index) => {
         colDef.resizable = index !== columnDefs.length - 1;
@@ -89,13 +90,26 @@ const TableComponent = forwardRef<
       realRef.current = params;
       const updatedColumnDefs = makeLastColumnNonResizable([...colDef]);
       params.api.setGridOption("columnDefs", updatedColumnDefs);
+      const customInit = localStorage.getItem(storeReference);
       initialColumnDefs.current = params.api.getColumnDefs();
-      if (props.onGridReady) props.onGridReady(params);
+      if (customInit && realRef.current) {
+        realRef.current.api.applyColumnState({
+          state: JSON.parse(customInit),
+          applyOrder: true,
+        });
+      }
       setTimeout(() => {
-        setColumnStateChange(false);
-      }, 200);
+        if (customInit && realRef.current) {
+          setColumnStateChange(true);
+        } else {
+          setColumnStateChange(false);
+        }
+      }, 50);
+      setTimeout(() => {
+        realRef.current.api.hideOverlay();
+      }, 1000);
+      if (props.onGridReady) props.onGridReady(params);
     };
-
     const onColumnMoved = (params) => {
       const updatedColumnDefs = makeLastColumnNonResizable(
         params.columnApi.getAllGridColumns().map((col) => col.getColDef()),
@@ -103,7 +117,6 @@ const TableComponent = forwardRef<
       params.api.setGridOption("columnDefs", updatedColumnDefs);
       if (props.onColumnMoved) props.onColumnMoved(params);
     };
-
     if (props.rowData.length === 0) {
       return (
         <div className="flex h-full w-full items-center justify-center rounded-md border">
@@ -138,6 +151,10 @@ const TableComponent = forwardRef<
           onColumnMoved={onColumnMoved}
           onStateUpdated={(e) => {
             if (e.sources.some((source) => source.includes("column"))) {
+              localStorage.setItem(
+                storeReference,
+                JSON.stringify(realRef.current?.api?.getColumnState()),
+              );
               setColumnStateChange(true);
             }
           }}
@@ -152,6 +169,7 @@ const TableComponent = forwardRef<
               resetGrid(realRef, initialColumnDefs);
               setTimeout(() => {
                 setColumnStateChange(false);
+                localStorage.removeItem(storeReference);
               }, 100);
             }}
           />
diff --git a/src/frontend/src/modals/IOModal/components/chatView/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
index 067cf525c..876295d07 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
@@ -65,10 +65,15 @@ export default function ChatView({
       .filter((output) => output.data.message)
       .map((output, index) => {
         try {
-          const { sender, message, sender_name, stream_url, files } = output
-            .data.message as ChatOutputType;
+          console.log("output:", output);
+          const { sender, message, sender_name, stream_url, files } =
+            output.data.message.message !== "" ||
+            (output.data.message.files ?? []).length > 0
+              ? output.data.message
+              : output.data.artifacts;
 
-          const is_ai = sender === "Machine" || sender === null;
+          const is_ai =
+            sender === "Machine" || sender === null || sender === undefined;
           return {
             isSend: !is_ai,
             message: message,
diff --git a/src/frontend/src/modals/IOModal/index.tsx b/src/frontend/src/modals/IOModal/index.tsx
index 903455d5a..0af859cc3 100644
--- a/src/frontend/src/modals/IOModal/index.tsx
+++ b/src/frontend/src/modals/IOModal/index.tsx
@@ -366,7 +366,9 @@ export default function IOModal({
                                 size="md"
                                 className="block truncate"
                               >
-                                {session}
+                                {session === currentFlow?.id
+                                  ? "Default Session"
+                                  : session}
                               </Badge>
                             </div>
                           </ShadTooltip>
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 2ae8b20f6..5627401b5 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -170,7 +170,7 @@ export type VertexDataTypeAPI = {
   inactive?: boolean;
   timedelta?: number;
   duration?: string;
-  artifacts?: any;
+  artifacts?: any | ChatOutputType | ChatInputType;
   message: ChatOutputType | ChatInputType;
 };