diff --git a/packages/editor-sample/src/App/InspectorDrawer/ConfigurationPanel/index.tsx b/packages/editor-sample/src/App/InspectorDrawer/ConfigurationPanel/index.tsx
index fc1b52d..3c4f144 100644
--- a/packages/editor-sample/src/App/InspectorDrawer/ConfigurationPanel/index.tsx
+++ b/packages/editor-sample/src/App/InspectorDrawer/ConfigurationPanel/index.tsx
@@ -3,7 +3,7 @@ import React from 'react';
 import { Box, Typography } from '@mui/material';
 
 import { TEditorBlock } from '../../../documents/editor/core';
-import { setEditorState, useDocument, useSelectedBlockId } from '../../../documents/editor/EditorContext';
+import { setDocument, useDocument, useSelectedBlockId } from '../../../documents/editor/EditorContext';
 
 import AvatarSidebarPanel from './input-panels/AvatarSidebarPanel';
 import ButtonSidebarPanel from './input-panels/ButtonSidebarPanel';
@@ -37,10 +37,7 @@ export default function ConfigurationPanel() {
     return renderMessage(`Block with id ${selectedBlockId} was not found. Click on a block to reset.`);
   }
 
-  const setBlock = (conf: TEditorBlock) =>
-    setEditorState({
-      document: { ...document, [selectedBlockId]: conf },
-    });
+  const setBlock = (conf: TEditorBlock) => setDocument({ [selectedBlockId]: conf });
   const { data, type } = block;
   switch (type) {
     case 'Avatar':
diff --git a/packages/editor-sample/src/App/InspectorDrawer/StylesPanel.tsx b/packages/editor-sample/src/App/InspectorDrawer/StylesPanel.tsx
index 42ba2eb..ff31c6a 100644
--- a/packages/editor-sample/src/App/InspectorDrawer/StylesPanel.tsx
+++ b/packages/editor-sample/src/App/InspectorDrawer/StylesPanel.tsx
@@ -1,12 +1,11 @@
 import React from 'react';
 
-import { setEditorState, useDocument } from '../../documents/editor/EditorContext';
+import { setDocument, useDocument } from '../../documents/editor/EditorContext';
 
 import EmailLayoutSidebarPanel from './ConfigurationPanel/input-panels/EmailLayoutSidebarPanel';
 
 export default function StylesPanel() {
   const block = useDocument().root;
-  const document = useDocument();
   if (!block) {
     return <p>Block not found</p>;
   }
@@ -15,11 +14,6 @@ export default function StylesPanel() {
   if (type !== 'EmailLayout') {
     throw new Error('Expected "root" element to be of type EmailLayout');
   }
-  return (
-    <EmailLayoutSidebarPanel
-      key="root"
-      data={data}
-      setData={(data) => setEditorState({ document: { ...document, root: { type, data } } })}
-    />
-  );
+
+  return <EmailLayoutSidebarPanel key="root" data={data} setData={(data) => setDocument({ root: { type, data } })} />;
 }
diff --git a/packages/editor-sample/src/App/InspectorDrawer/ToggleInspectorPanelButton.tsx b/packages/editor-sample/src/App/InspectorDrawer/ToggleInspectorPanelButton.tsx
index a07b39b..2dc41b9 100644
--- a/packages/editor-sample/src/App/InspectorDrawer/ToggleInspectorPanelButton.tsx
+++ b/packages/editor-sample/src/App/InspectorDrawer/ToggleInspectorPanelButton.tsx
@@ -3,13 +3,13 @@ import React from 'react';
 import { AppRegistrationOutlined, LastPageOutlined } from '@mui/icons-material';
 import { IconButton } from '@mui/material';
 
-import { setEditorState, useInspectorDrawerOpen } from '../../documents/editor/EditorContext';
+import { toggleInspectorDrawerOpen, useInspectorDrawerOpen } from '../../documents/editor/EditorContext';
 
 export default function ToggleInspectorPanelButton() {
   const inspectorDrawerOpen = useInspectorDrawerOpen();
 
   const handleClick = () => {
-    setEditorState({ inspectorDrawerOpen: !inspectorDrawerOpen });
+    toggleInspectorDrawerOpen();
   };
   if (inspectorDrawerOpen) {
     return (
diff --git a/packages/editor-sample/src/App/InspectorDrawer/index.tsx b/packages/editor-sample/src/App/InspectorDrawer/index.tsx
index 144b93f..3b2c2cc 100644
--- a/packages/editor-sample/src/App/InspectorDrawer/index.tsx
+++ b/packages/editor-sample/src/App/InspectorDrawer/index.tsx
@@ -2,7 +2,7 @@ import React from 'react';
 
 import { Box, Drawer, Tab, Tabs } from '@mui/material';
 
-import { setEditorState, useInspectorDrawerOpen, useSelectedSidebarTab } from '../../documents/editor/EditorContext';
+import { setSidebarTab, useInspectorDrawerOpen, useSelectedSidebarTab } from '../../documents/editor/EditorContext';
 
 import ConfigurationPanel from './ConfigurationPanel';
 import StylesPanel from './StylesPanel';
@@ -33,7 +33,7 @@ export default function InspectorDrawer() {
     >
       <Box sx={{ width: INSPECTOR_DRAWER_WIDTH, height: 49, borderBottom: 1, borderColor: 'divider' }}>
         <Box px={2}>
-          <Tabs value={selectedSidebarTab} onChange={(_, v) => setEditorState({ selectedSidebarTab: v })}>
+          <Tabs value={selectedSidebarTab} onChange={(_, v) => setSidebarTab(v)}>
             <Tab value="styles" label="Styles" />
             <Tab value="block-configuration" label="Inspect" />
           </Tabs>
diff --git a/packages/editor-sample/src/App/SamplesDrawer/SidebarButton.tsx b/packages/editor-sample/src/App/SamplesDrawer/SidebarButton.tsx
index 2d020d6..d551abe 100644
--- a/packages/editor-sample/src/App/SamplesDrawer/SidebarButton.tsx
+++ b/packages/editor-sample/src/App/SamplesDrawer/SidebarButton.tsx
@@ -2,13 +2,12 @@ import React from 'react';
 
 import { Button } from '@mui/material';
 
-import { setEditorState } from '../../documents/editor/EditorContext';
+import { resetDocument } from '../../documents/editor/EditorContext';
 import getConfiguration from '../../getConfiguration';
 
 export default function SidebarButton({ href, children }: { href: string; children: JSX.Element | string }) {
   const handleClick = () => {
-    const document = getConfiguration(href);
-    setEditorState({ document });
+    resetDocument(getConfiguration(href));
   };
   return (
     <Button size="small" href={href} onClick={handleClick}>
diff --git a/packages/editor-sample/src/App/SamplesDrawer/ToggleSamplesPanelButton.tsx b/packages/editor-sample/src/App/SamplesDrawer/ToggleSamplesPanelButton.tsx
index 42f847b..dbb2a23 100644
--- a/packages/editor-sample/src/App/SamplesDrawer/ToggleSamplesPanelButton.tsx
+++ b/packages/editor-sample/src/App/SamplesDrawer/ToggleSamplesPanelButton.tsx
@@ -3,24 +3,17 @@ import React from 'react';
 import { FirstPageOutlined, MenuOutlined } from '@mui/icons-material';
 import { IconButton } from '@mui/material';
 
-import { setEditorState, useSamplesDrawerOpen } from '../../documents/editor/EditorContext';
+import { toggleSamplesDrawerOpen, useSamplesDrawerOpen } from '../../documents/editor/EditorContext';
 
-export default function ToggleSamplesPanelButton() {
+function useIcon() {
   const samplesDrawerOpen = useSamplesDrawerOpen();
-
-  const handleClick = () => {
-    setEditorState({ samplesDrawerOpen: !samplesDrawerOpen });
-  };
   if (samplesDrawerOpen) {
-    return (
-      <IconButton onClick={handleClick}>
-        <FirstPageOutlined fontSize="small" />
-      </IconButton>
-    );
+    return <FirstPageOutlined fontSize="small" />;
   }
-  return (
-    <IconButton onClick={handleClick}>
-      <MenuOutlined fontSize="small" />
-    </IconButton>
-  );
+  return <MenuOutlined fontSize="small" />;
+}
+
+export default function ToggleSamplesPanelButton() {
+  const icon = useIcon();
+  return <IconButton onClick={toggleSamplesDrawerOpen}>{icon}</IconButton>;
 }
diff --git a/packages/editor-sample/src/App/TemplatePanel/MainTabsGroup.tsx b/packages/editor-sample/src/App/TemplatePanel/MainTabsGroup.tsx
new file mode 100644
index 0000000..ba8c451
--- /dev/null
+++ b/packages/editor-sample/src/App/TemplatePanel/MainTabsGroup.tsx
@@ -0,0 +1,59 @@
+import React from 'react';
+
+import { CodeOutlined, DataObjectOutlined, EditOutlined, PreviewOutlined } from '@mui/icons-material';
+import { Tab, Tabs, Tooltip } from '@mui/material';
+
+import { setSelectedMainTab, useSelectedMainTab } from '../../documents/editor/EditorContext';
+
+export default function MainTabsGroup() {
+  const selectedMainTab = useSelectedMainTab();
+  const handleChange = (_: unknown, v: unknown) => {
+    switch (v) {
+      case 'data':
+      case 'editor':
+      case 'preview':
+      case 'html':
+        setSelectedMainTab(v);
+        return;
+      default:
+        setSelectedMainTab('data');
+    }
+  };
+
+  return (
+    <Tabs value={selectedMainTab} onChange={handleChange}>
+      <Tab
+        value="editor"
+        label={
+          <Tooltip title="Edit">
+            <EditOutlined fontSize="small" />
+          </Tooltip>
+        }
+      />
+      <Tab
+        value="preview"
+        label={
+          <Tooltip title="Preview">
+            <PreviewOutlined fontSize="small" />
+          </Tooltip>
+        }
+      />
+      <Tab
+        value="html"
+        label={
+          <Tooltip title="HTML output">
+            <CodeOutlined fontSize="small" />
+          </Tooltip>
+        }
+      />
+      <Tab
+        value="data"
+        label={
+          <Tooltip title="JSON output">
+            <DataObjectOutlined fontSize="small" />
+          </Tooltip>
+        }
+      />
+    </Tabs>
+  );
+}
diff --git a/packages/editor-sample/src/App/TemplatePanel/index.tsx b/packages/editor-sample/src/App/TemplatePanel/index.tsx
index a77778b..ceae324 100644
--- a/packages/editor-sample/src/App/TemplatePanel/index.tsx
+++ b/packages/editor-sample/src/App/TemplatePanel/index.tsx
@@ -1,18 +1,11 @@
 import React from 'react';
 
-import {
-  CodeOutlined,
-  DataObjectOutlined,
-  EditOutlined,
-  MonitorOutlined,
-  PhoneIphoneOutlined,
-  PreviewOutlined,
-} from '@mui/icons-material';
-import { Box, Stack, SxProps, Tab, Tabs, ToggleButton, ToggleButtonGroup, Tooltip } from '@mui/material';
+import { MonitorOutlined, PhoneIphoneOutlined } from '@mui/icons-material';
+import { Box, Stack, SxProps, ToggleButton, ToggleButtonGroup, Tooltip } from '@mui/material';
 
 import EditorBlock from '../../documents/editor/EditorBlock';
 import {
-  setEditorState,
+  setSelectedScreenSize,
   useDocument,
   useSelectedMainTab,
   useSelectedScreenSize,
@@ -24,6 +17,7 @@ import ToggleSamplesPanelButton from '../SamplesDrawer/ToggleSamplesPanelButton'
 
 import HtmlPanel from './HtmlPanel';
 import JsonPanel from './JsonPanel';
+import MainTabsGroup from './MainTabsGroup';
 import ShareButton from './ShareButton';
 
 export default function TemplatePanel() {
@@ -46,14 +40,13 @@ export default function TemplatePanel() {
   }
 
   const handleScreenSizeChange = (_: unknown, value: unknown) => {
-    console.log(value);
     switch (value) {
       case 'mobile':
-        setEditorState({ selectedScreenSize: 'mobile' });
-        return;
       case 'desktop':
+        setSelectedScreenSize(value);
+        return;
       default:
-        setEditorState({ selectedScreenSize: 'desktop' });
+        setSelectedScreenSize('desktop');
     }
   };
 
@@ -93,40 +86,7 @@ export default function TemplatePanel() {
       >
         <ToggleSamplesPanelButton />
         <Stack px={2} direction="row" gap={2} width="100%" justifyContent="space-between" alignItems="center">
-          <Tabs value={selectedMainTab} onChange={(_, v) => setEditorState({ selectedMainTab: v })}>
-            <Tab
-              value="editor"
-              label={
-                <Tooltip title="Edit">
-                  <EditOutlined fontSize="small" />
-                </Tooltip>
-              }
-            />
-            <Tab
-              value="preview"
-              label={
-                <Tooltip title="Preview">
-                  <PreviewOutlined fontSize="small" />
-                </Tooltip>
-              }
-            />
-            <Tab
-              value="html"
-              label={
-                <Tooltip title="HTML output">
-                  <CodeOutlined fontSize="small" />
-                </Tooltip>
-              }
-            />
-            <Tab
-              value="data"
-              label={
-                <Tooltip title="JSON output">
-                  <DataObjectOutlined fontSize="small" />
-                </Tooltip>
-              }
-            />
-          </Tabs>
+          <MainTabsGroup />
           <Stack direction="row" spacing={2}>
             <ToggleButtonGroup value={selectedScreenSize} exclusive size="small" onChange={handleScreenSizeChange}>
               <ToggleButton value="desktop">
@@ -143,7 +103,6 @@ export default function TemplatePanel() {
             <ShareButton />
           </Stack>
         </Stack>
-
         <ToggleInspectorPanelButton />
       </Stack>
       <Box sx={{ height: 'calc(100vh - 49px)', overflow: 'auto', minWidth: 370 }}>
diff --git a/packages/editor-sample/src/documents/blocks/ColumnsContainer/index.tsx b/packages/editor-sample/src/documents/blocks/ColumnsContainer/index.tsx
index ea7415c..d7b7ae0 100644
--- a/packages/editor-sample/src/documents/blocks/ColumnsContainer/index.tsx
+++ b/packages/editor-sample/src/documents/blocks/ColumnsContainer/index.tsx
@@ -4,7 +4,7 @@ import { ColumnsContainer as BaseColumnsContainer } from '@usewaypoint/block-col
 
 import { TEditorBlock } from '../../editor/core';
 import { useCurrentBlockId } from '../../editor/EditorBlock';
-import { setEditorState, useDocument } from '../../editor/EditorContext';
+import { setDocument, setSelectedBlockId } from '../../editor/EditorContext';
 import ReaderBlock from '../../reader/ReaderBlock';
 import EditorChildrenIds from '../helpers/EditorChildrenIds';
 
@@ -23,7 +23,6 @@ export function ColumnsContainer({ style, props }: ColumnsContainerProps) {
 const EMPTY_COLUMNS = [{ childrenIds: [] }, { childrenIds: [] }, { childrenIds: [] }];
 
 export function EditorColumnsContainer({ style, props }: ColumnsContainerProps) {
-  const document = useDocument();
   const blockId = useCurrentBlockId();
 
   const { columns, ...restProps } = props ?? {};
@@ -59,21 +58,18 @@ export function EditorColumnsContainer({ style, props }: ColumnsContainerProps)
       return columnsCopy;
     };
 
-    setEditorState({
-      selectedBlockId: id,
-      document: {
-        ...document,
-        [id]: blockConfiguration,
-        [blockId]: {
-          type: 'ColumnsContainer',
-          data: ColumnsContainerPropsSchema.parse({
-            style,
-            props: {
-              ...restProps,
-              columns: getColumns(),
-            },
-          }),
-        },
+    setSelectedBlockId(id);
+    setDocument({
+      [id]: blockConfiguration,
+      [blockId]: {
+        type: 'ColumnsContainer',
+        data: ColumnsContainerPropsSchema.parse({
+          style,
+          props: {
+            ...restProps,
+            columns: getColumns(),
+          },
+        }),
       },
     });
   };
diff --git a/packages/editor-sample/src/documents/blocks/Container/index.tsx b/packages/editor-sample/src/documents/blocks/Container/index.tsx
index 3f94e77..d37ec3a 100644
--- a/packages/editor-sample/src/documents/blocks/Container/index.tsx
+++ b/packages/editor-sample/src/documents/blocks/Container/index.tsx
@@ -4,7 +4,7 @@ import { Container as BaseContainer } from '@usewaypoint/block-container';
 
 import { TEditorBlock } from '../../editor/core';
 import { useCurrentBlockId } from '../../editor/EditorBlock';
-import { setEditorState, useDocument } from '../../editor/EditorContext';
+import { setDocument, setSelectedBlockId, useDocument } from '../../editor/EditorContext';
 import ReaderBlock from '../../reader/ReaderBlock';
 import EditorChildrenIds from '../helpers/EditorChildrenIds';
 
@@ -36,17 +36,14 @@ export function EditorContainer({ style, props }: ContainerProps) {
       nChildrenIds = [...childrenIds.slice(0, i), id, ...childrenIds.slice(i)];
     }
 
-    setEditorState({
-      selectedBlockId: id,
-      document: {
-        ...document,
-        [id]: blockConfiguration,
-        [blockId]: {
-          type: 'Container',
-          data: {
-            ...document[blockId].data,
-            props: { childrenIds: nChildrenIds },
-          },
+    setSelectedBlockId(id);
+    setDocument({
+      [id]: blockConfiguration,
+      [blockId]: {
+        type: 'Container',
+        data: {
+          ...document[blockId].data,
+          props: { childrenIds: nChildrenIds },
         },
       },
     });
diff --git a/packages/editor-sample/src/documents/blocks/EmailLayout/index.tsx b/packages/editor-sample/src/documents/blocks/EmailLayout/index.tsx
index 6f7ed4d..58c8074 100644
--- a/packages/editor-sample/src/documents/blocks/EmailLayout/index.tsx
+++ b/packages/editor-sample/src/documents/blocks/EmailLayout/index.tsx
@@ -3,7 +3,7 @@ import { z } from 'zod';
 
 import { TEditorBlock } from '../../editor/core';
 import { useCurrentBlockId } from '../../editor/EditorBlock';
-import { setEditorState, useDocument } from '../../editor/EditorContext';
+import { setDocument, setSelectedBlockId, useDocument } from '../../editor/EditorContext';
 import ReaderBlock from '../../reader/ReaderBlock';
 import EditorChildrenIds from '../helpers/EditorChildrenIds';
 
@@ -50,30 +50,24 @@ export function EditorEmailLayout(props: EmailLayoutProps) {
       return [...childrenIds.slice(0, i), id, ...childrenIds.slice(i)];
     };
 
-    setEditorState({
-      selectedBlockId: id,
-      document: {
-        ...document,
-        [id]: blockConfiguration,
-        [blockId]: {
-          type: 'EmailLayout',
-          data: {
-            ...document[blockId].data,
-            childrenIds: getChildrenIds(),
-          },
-        } as TEditorBlock,
-      },
+    setDocument({
+      [id]: blockConfiguration,
+      [blockId]: {
+        type: 'EmailLayout',
+        data: {
+          ...document[blockId].data,
+          childrenIds: getChildrenIds(),
+        },
+      } as TEditorBlock,
     });
+    setSelectedBlockId(id);
   };
 
   return (
     <div
       style={{ height: '100%' }}
       onClick={() => {
-        setEditorState({
-          selectedSidebarTab: 'styles',
-          selectedBlockId: null,
-        });
+        setSelectedBlockId(null);
       }}
     >
       <SharedLayout {...props}>
diff --git a/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/EditorBlockWrapper.tsx b/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/EditorBlockWrapper.tsx
index 0678095..3d68ae6 100644
--- a/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/EditorBlockWrapper.tsx
+++ b/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/EditorBlockWrapper.tsx
@@ -3,7 +3,7 @@ import React, { CSSProperties, useState } from 'react';
 import { Box } from '@mui/material';
 
 import { useCurrentBlockId } from '../../../editor/EditorBlock';
-import { setEditorState, useSelectedBlockId } from '../../../editor/EditorContext';
+import { setSelectedBlockId, useSelectedBlockId } from '../../../editor/EditorContext';
 
 import TuneMenu from './TuneMenu';
 
@@ -46,11 +46,7 @@ export default function EditorBlockWrapper({ children }: TEditorBlockWrapperProp
         setMouseInside(false);
       }}
       onClick={(ev) => {
-        setEditorState({
-          inspectorDrawerOpen: true,
-          selectedSidebarTab: 'block-configuration',
-          selectedBlockId: blockId,
-        });
+        setSelectedBlockId(blockId, { inspectorDrawerOpen: true });
         ev.stopPropagation();
         ev.preventDefault();
       }}
diff --git a/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/TuneMenu.tsx b/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/TuneMenu.tsx
index 640ad94..50429f1 100644
--- a/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/TuneMenu.tsx
+++ b/packages/editor-sample/src/documents/blocks/helpers/block-wrappers/TuneMenu.tsx
@@ -4,7 +4,7 @@ import { DeleteOutlined } from '@mui/icons-material';
 import { IconButton, Paper, Stack, Tooltip } from '@mui/material';
 
 import { TEditorBlock } from '../../../editor/core';
-import { setEditorState, useDocument } from '../../../editor/EditorContext';
+import { resetDocument, useDocument } from '../../../editor/EditorContext';
 import { ColumnsContainerProps } from '../../ColumnsContainer/ColumnsContainerPropsSchema';
 
 const STYLE: CSSProperties = {
@@ -68,10 +68,7 @@ export default function TuneMenu({ blockId }: Props) {
       }
     }
     delete nDocument[blockId];
-    setEditorState({
-      document: nDocument,
-      selectedBlockId: null,
-    });
+    resetDocument(nDocument);
   };
 
   return (
diff --git a/packages/editor-sample/src/documents/editor/EditorContext.tsx b/packages/editor-sample/src/documents/editor/EditorContext.tsx
index 35a52cb..36e8abf 100644
--- a/packages/editor-sample/src/documents/editor/EditorContext.tsx
+++ b/packages/editor-sample/src/documents/editor/EditorContext.tsx
@@ -16,7 +16,7 @@ type TValue = {
   samplesDrawerOpen: boolean;
 };
 
-const useEditorState = create<TValue>(() => ({
+const editorStateStore = create<TValue>(() => ({
   document: getConfiguration(window.location.hash),
   selectedBlockId: null,
   selectedSidebarTab: 'styles',
@@ -28,33 +28,77 @@ const useEditorState = create<TValue>(() => ({
 }));
 
 export function useDocument() {
-  return useEditorState((s) => s.document);
+  return editorStateStore((s) => s.document);
 }
 
 export function useSelectedBlockId() {
-  return useEditorState((s) => s.selectedBlockId);
+  return editorStateStore((s) => s.selectedBlockId);
 }
 
 export function useSelectedScreenSize() {
-  return useEditorState((s) => s.selectedScreenSize);
+  return editorStateStore((s) => s.selectedScreenSize);
 }
 
 export function useSelectedMainTab() {
-  return useEditorState((s) => s.selectedMainTab);
+  return editorStateStore((s) => s.selectedMainTab);
+}
+
+export function setSelectedMainTab(selectedMainTab: TValue['selectedMainTab']) {
+  return editorStateStore.setState({ selectedMainTab });
 }
 
 export function useSelectedSidebarTab() {
-  return useEditorState((s) => s.selectedSidebarTab);
+  return editorStateStore((s) => s.selectedSidebarTab);
 }
 
 export function useInspectorDrawerOpen() {
-  return useEditorState((s) => s.inspectorDrawerOpen);
+  return editorStateStore((s) => s.inspectorDrawerOpen);
 }
 
 export function useSamplesDrawerOpen() {
-  return useEditorState((s) => s.samplesDrawerOpen);
+  return editorStateStore((s) => s.samplesDrawerOpen);
+}
+
+type SetSelectedBlockIdOptions = {
+  inspectorDrawerOpen: boolean;
+};
+export function setSelectedBlockId(selectedBlockId: TValue['selectedBlockId'], options?: SetSelectedBlockIdOptions) {
+  const selectedSidebarTab = selectedBlockId === null ? 'styles' : 'block-configuration';
+  return editorStateStore.setState({ selectedBlockId, selectedSidebarTab, options });
+}
+
+export function setSidebarTab(selectedSidebarTab: TValue['selectedSidebarTab']) {
+  return editorStateStore.setState({ selectedSidebarTab });
+}
+
+export function resetDocument(document: TValue['document']) {
+  return editorStateStore.setState({
+    document,
+    selectedSidebarTab: 'styles',
+    selectedBlockId: null,
+  });
+}
+
+export function setDocument(document: TValue['document']) {
+  const originalDocument = editorStateStore.getState().document;
+  return editorStateStore.setState({
+    document: {
+      ...originalDocument,
+      ...document,
+    },
+  });
+}
+
+export function toggleInspectorDrawerOpen() {
+  const inspectorDrawerOpen = !editorStateStore.getState().inspectorDrawerOpen;
+  return editorStateStore.setState({ inspectorDrawerOpen });
+}
+
+export function toggleSamplesDrawerOpen() {
+  const samplesDrawerOpen = !editorStateStore.getState().samplesDrawerOpen;
+  return editorStateStore.setState({ samplesDrawerOpen });
 }
 
-export function setEditorState(state: Partial<TValue>) {
-  useEditorState.setState({ ...state });
+export function setSelectedScreenSize(selectedScreenSize: TValue['selectedScreenSize']) {
+  return editorStateStore.setState({ selectedScreenSize });
 }