From 838a029b997fa634c48587cd533462324fc0ef79 Mon Sep 17 00:00:00 2001
From: Quinn Slack <quinn@slack.org>
Date: Tue, 18 Apr 2023 19:49:44 -0700
Subject: [PATCH] add storybooks for easier local dev on vscode UI (#50772)

Running `cd client/cody && pnpm run storybook` will open storybooks for
the VS Code extension. There is a light shim in place to support
rendering the VS Code extension UI with roughly the right theme colors.

### Storybook for `<App />` (VS Code main webview UI)


![image](https://user-images.githubusercontent.com/1976/232655933-1886568f-c305-42d1-91a0-5555768e92f4.png)

![image](https://user-images.githubusercontent.com/1976/232655957-8e27e258-2214-4d81-a900-2b9f76e84f8c.png)


### Storybook for login component showing invalid state


![image](https://user-images.githubusercontent.com/1976/232655884-6632dfc8-859e-4e0e-ad1b-2fa3320154f4.png)


## Test plan

n/a; dev only

---------

Co-authored-by: Naman Kumar <naman@outlook.in>
---
 client/cody/package.json                      |  3 +-
 client/cody/webviews/App.story.tsx            | 35 ++++++++++++++++++
 client/cody/webviews/App.tsx                  | 31 +++++++++-------
 client/cody/webviews/Chat.tsx                 | 13 ++++---
 .../webviews/{Debug.css => Debug.module.css}  |  0
 client/cody/webviews/Debug.tsx                |  7 ++--
 client/cody/webviews/FileLink.tsx             |  4 +--
 .../{Header.css => Header.module.css}         |  0
 client/cody/webviews/Header.tsx               | 18 +++++-----
 client/cody/webviews/Login.story.tsx          | 27 ++++++++++++++
 .../{NavBar.css => NavBar.module.css}         |  2 +-
 client/cody/webviews/NavBar.tsx               | 18 +++++-----
 .../{Recipes.css => Recipes.module.css}       |  0
 client/cody/webviews/Recipes.tsx              | 10 +++---
 client/cody/webviews/UserHistory.tsx          |  6 ++--
 client/cody/webviews/index.tsx                |  4 ++-
 .../storybook/VSCodeStoryDecorator.module.css | 34 ++++++++++++++++++
 .../storybook/VSCodeStoryDecorator.tsx        |  8 +++++
 client/cody/webviews/utils/VSCodeApi.ts       | 36 +++++++++++--------
 client/storybook/src/main.ts                  | 16 +++++++--
 20 files changed, 206 insertions(+), 66 deletions(-)
 create mode 100644 client/cody/webviews/App.story.tsx
 rename client/cody/webviews/{Debug.css => Debug.module.css} (100%)
 rename client/cody/webviews/{Header.css => Header.module.css} (100%)
 create mode 100644 client/cody/webviews/Login.story.tsx
 rename client/cody/webviews/{NavBar.css => NavBar.module.css} (92%)
 rename client/cody/webviews/{Recipes.css => Recipes.module.css} (100%)
 create mode 100644 client/cody/webviews/storybook/VSCodeStoryDecorator.module.css
 create mode 100644 client/cody/webviews/storybook/VSCodeStoryDecorator.tsx

diff --git a/client/cody/package.json b/client/cody/package.json
index 1257ee629b0..bde89590f37 100644
--- a/client/cody/package.json
+++ b/client/cody/package.json
@@ -41,7 +41,8 @@
     "release": "ts-node ./scripts/release.ts",
     "watch": "concurrently \"pnpm watch:esbuild\" \"pnpm watch:webview\"",
     "watch:esbuild": "pnpm esbuild --sourcemap --watch",
-    "watch:webview": "vite build --mode development --watch"
+    "watch:webview": "vite build --mode development --watch",
+    "storybook": "STORIES_GLOB='client/cody/webviews/**/*.story.tsx' pnpm --filter @sourcegraph/storybook run start"
   },
   "main": "./dist/extension.js",
   "activationEvents": [
diff --git a/client/cody/webviews/App.story.tsx b/client/cody/webviews/App.story.tsx
new file mode 100644
index 00000000000..603bd55131c
--- /dev/null
+++ b/client/cody/webviews/App.story.tsx
@@ -0,0 +1,35 @@
+import { ComponentMeta, ComponentStoryObj } from '@storybook/react'
+
+import { App } from './App'
+import { VSCodeStoryDecorator } from './storybook/VSCodeStoryDecorator'
+import { VSCodeWrapper } from './utils/VSCodeApi'
+
+const meta: ComponentMeta<typeof App> = {
+    title: 'cody/App',
+    component: App,
+
+    decorators: [VSCodeStoryDecorator],
+
+    parameters: {
+        component: App,
+        chromatic: {
+            enableDarkMode: true,
+            disableSnapshot: false,
+        },
+    },
+}
+
+export default meta
+
+export const Simple: ComponentStoryObj<typeof App> = {
+    render: () => <App vscodeAPI={dummyVSCodeAPI} />,
+}
+
+const dummyVSCodeAPI: VSCodeWrapper = {
+    onMessage: cb => {
+        // Send initial message so that the component is fully rendered.
+        cb({ type: 'config', config: { debug: true, hasAccessToken: true, serverEndpoint: 'https://example.com' } })
+        return () => {}
+    },
+    postMessage: () => {},
+}
diff --git a/client/cody/webviews/App.tsx b/client/cody/webviews/App.tsx
index 330be674bec..91754613ec3 100644
--- a/client/cody/webviews/App.tsx
+++ b/client/cody/webviews/App.tsx
@@ -15,9 +15,9 @@ import { NavBar, View } from './NavBar'
 import { Recipes } from './Recipes'
 import { Settings } from './Settings'
 import { UserHistory } from './UserHistory'
-import { vscodeAPI } from './utils/VSCodeApi'
+import type { VSCodeWrapper } from './utils/VSCodeApi'
 
-export function App(): React.ReactElement {
+export const App: React.FunctionComponent<{ vscodeAPI: VSCodeWrapper }> = ({ vscodeAPI }) => {
     const [config, setConfig] = useState<Pick<Configuration, 'debug' | 'serverEndpoint'> | null>(null)
     const [debugLog, setDebugLog] = useState(['No data yet'])
     const [view, setView] = useState<View | undefined>()
@@ -74,20 +74,23 @@ export function App(): React.ReactElement {
 
         vscodeAPI.postMessage({ command: 'initialized' })
         // The dependencies array is empty to execute the callback only on component mount.
-    }, [debugLog])
+    }, [debugLog, vscodeAPI])
 
-    const onLogin = useCallback((token: string, endpoint: string) => {
-        if (!token || !endpoint) {
-            return
-        }
-        setIsValidLogin(undefined)
-        vscodeAPI.postMessage({ command: 'settings', serverEndpoint: endpoint, accessToken: token })
-    }, [])
+    const onLogin = useCallback(
+        (token: string, endpoint: string) => {
+            if (!token || !endpoint) {
+                return
+            }
+            setIsValidLogin(undefined)
+            vscodeAPI.postMessage({ command: 'settings', serverEndpoint: endpoint, accessToken: token })
+        },
+        [vscodeAPI]
+    )
 
     const onLogout = useCallback(() => {
         vscodeAPI.postMessage({ command: 'removeToken' })
         setView('login')
-    }, [setView])
+    }, [vscodeAPI])
 
     const onResetClick = useCallback(() => {
         setView('chat')
@@ -96,7 +99,7 @@ export function App(): React.ReactElement {
         setMessageInProgress(null)
         setTranscript([])
         vscodeAPI.postMessage({ command: 'reset' })
-    }, [setView, setMessageInProgress, setTranscript, setDebugLog])
+    }, [vscodeAPI])
 
     if (!view) {
         return <LoadingPage />
@@ -123,9 +126,10 @@ export function App(): React.ReactElement {
                     userHistory={userHistory}
                     setUserHistory={setUserHistory}
                     setInputHistory={setInputHistory}
+                    vscodeAPI={vscodeAPI}
                 />
             )}
-            {view === 'recipes' && <Recipes />}
+            {view === 'recipes' && <Recipes vscodeAPI={vscodeAPI} />}
             {view === 'settings' && (
                 <Settings setView={setView} onLogout={onLogout} serverEndpoint={config?.serverEndpoint} />
             )}
@@ -138,6 +142,7 @@ export function App(): React.ReactElement {
                     setFormInput={setFormInput}
                     inputHistory={inputHistory}
                     setInputHistory={setInputHistory}
+                    vscodeAPI={vscodeAPI}
                 />
             )}
         </div>
diff --git a/client/cody/webviews/Chat.tsx b/client/cody/webviews/Chat.tsx
index 5b6c7d5bc29..f286668eb45 100644
--- a/client/cody/webviews/Chat.tsx
+++ b/client/cody/webviews/Chat.tsx
@@ -9,7 +9,7 @@ import { Chat as ChatUI, ChatUISubmitButtonProps, ChatUITextAreaProps } from '@s
 import { SubmitSvg } from '@sourcegraph/cody-ui/src/utils/icons'
 
 import { FileLink } from './FileLink'
-import { vscodeAPI } from './utils/VSCodeApi'
+import { VSCodeWrapper } from './utils/VSCodeApi'
 
 import styles from './Chat.module.css'
 
@@ -21,6 +21,7 @@ interface ChatboxProps {
     setFormInput: (input: string) => void
     inputHistory: string[]
     setInputHistory: (history: string[]) => void
+    vscodeAPI: VSCodeWrapper
 }
 
 export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>> = ({
@@ -31,10 +32,14 @@ export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>
     setFormInput,
     inputHistory,
     setInputHistory,
+    vscodeAPI,
 }) => {
-    const onSubmit = useCallback((text: string) => {
-        vscodeAPI.postMessage({ command: 'submit', text })
-    }, [])
+    const onSubmit = useCallback(
+        (text: string) => {
+            vscodeAPI.postMessage({ command: 'submit', text })
+        },
+        [vscodeAPI]
+    )
 
     return (
         <ChatUI
diff --git a/client/cody/webviews/Debug.css b/client/cody/webviews/Debug.module.css
similarity index 100%
rename from client/cody/webviews/Debug.css
rename to client/cody/webviews/Debug.module.css
diff --git a/client/cody/webviews/Debug.tsx b/client/cody/webviews/Debug.tsx
index 6f3de7e4052..b4c011c54cf 100644
--- a/client/cody/webviews/Debug.tsx
+++ b/client/cody/webviews/Debug.tsx
@@ -1,4 +1,4 @@
-import './Debug.css'
+import styles from './Debug.module.css'
 
 interface DebugProps {
     debugLog: string[]
@@ -7,9 +7,10 @@ interface DebugProps {
 export const Debug: React.FunctionComponent<React.PropsWithChildren<DebugProps>> = ({ debugLog }) => (
     <div className="inner-container">
         <div className="non-transcript-container">
-            <div className="debug-container" data-tab-target="debug">
+            <div className={styles.debugContainer} data-tab-target="debug">
                 {debugLog?.map((log, i) => (
-                    <div key={`log-${i}`} className="debug-message">
+                    // eslint-disable-next-line react/no-array-index-key
+                    <div key={i} className={styles.debugMessage}>
                         {log}
                     </div>
                 ))}
diff --git a/client/cody/webviews/FileLink.tsx b/client/cody/webviews/FileLink.tsx
index cd66ddc16aa..20d6edf1014 100644
--- a/client/cody/webviews/FileLink.tsx
+++ b/client/cody/webviews/FileLink.tsx
@@ -2,7 +2,7 @@ import React from 'react'
 
 import { FileLinkProps } from '@sourcegraph/cody-ui/src/chat/ContextFiles'
 
-import { vscodeAPI } from './utils/VSCodeApi'
+import { getVSCodeAPI } from './utils/VSCodeApi'
 
 import styles from './FileLink.module.css'
 
@@ -11,7 +11,7 @@ export const FileLink: React.FunctionComponent<FileLinkProps> = ({ path }) => (
         className={styles.linkButton}
         type="button"
         onClick={() => {
-            vscodeAPI.postMessage({ command: 'openFile', filePath: path })
+            getVSCodeAPI().postMessage({ command: 'openFile', filePath: path })
         }}
     >
         {path}
diff --git a/client/cody/webviews/Header.css b/client/cody/webviews/Header.module.css
similarity index 100%
rename from client/cody/webviews/Header.css
rename to client/cody/webviews/Header.module.css
diff --git a/client/cody/webviews/Header.tsx b/client/cody/webviews/Header.tsx
index 4b42a88cc91..ede200371f3 100644
--- a/client/cody/webviews/Header.tsx
+++ b/client/cody/webviews/Header.tsx
@@ -1,20 +1,20 @@
 import { VSCodeTag } from '@vscode/webview-ui-toolkit/react'
 
-import './Header.css'
-
 import { CodyColoredSvg } from '@sourcegraph/cody-ui/src/utils/icons'
 
+import styles from './Header.module.css'
+
 export const Header: React.FunctionComponent = () => (
-    <div className="header-container">
-        <div className="header-container-left">
-            <div className="header-logo">
+    <div className={styles.headerContainer}>
+        <div className={styles.headerContainerLeft}>
+            <div className={styles.headerLogo}>
                 <CodyColoredSvg />
             </div>
-            <div className="header-title">
-                <span className="header-cody">Cody</span>
-                <VSCodeTag className="tag-beta">experimental</VSCodeTag>
+            <div className={styles.headerTitle}>
+                <span className={styles.headerCody}>Cody</span>
+                <VSCodeTag className={styles.tagBeta}>experimental</VSCodeTag>
             </div>
         </div>
-        <div className="header-container-right" />
+        <div className={styles.headerContainerRight} />
     </div>
 )
diff --git a/client/cody/webviews/Login.story.tsx b/client/cody/webviews/Login.story.tsx
new file mode 100644
index 00000000000..41d70216ec7
--- /dev/null
+++ b/client/cody/webviews/Login.story.tsx
@@ -0,0 +1,27 @@
+import { ComponentMeta, ComponentStoryObj } from '@storybook/react'
+
+import { Login } from './Login'
+import { VSCodeStoryDecorator } from './storybook/VSCodeStoryDecorator'
+
+const meta: ComponentMeta<typeof Login> = {
+    title: 'cody/Login',
+    component: Login,
+    decorators: [VSCodeStoryDecorator],
+    parameters: {
+        component: Login,
+        chromatic: {
+            enableDarkMode: true,
+            disableSnapshot: false,
+        },
+    },
+}
+
+export default meta
+
+export const Simple: ComponentStoryObj<typeof Login> = {
+    render: () => <Login onLogin={() => {}} isValidLogin={true} />,
+}
+
+export const InvalidLogin: ComponentStoryObj<typeof Login> = {
+    render: () => <Login onLogin={() => {}} isValidLogin={false} />,
+}
diff --git a/client/cody/webviews/NavBar.css b/client/cody/webviews/NavBar.module.css
similarity index 92%
rename from client/cody/webviews/NavBar.css
rename to client/cody/webviews/NavBar.module.css
index a4c2bcffd82..fa329782bc2 100644
--- a/client/cody/webviews/NavBar.css
+++ b/client/cody/webviews/NavBar.module.css
@@ -14,7 +14,7 @@
 .tab-btn {
     cursor: pointer;
     user-select: none;
-    background-color: var(--vscode-tab-inactiveBackground);
+    background-color: transparent;
     color: var(--vscode-tab-inactiveForeground);
     border: none;
 }
diff --git a/client/cody/webviews/NavBar.tsx b/client/cody/webviews/NavBar.tsx
index ace5e121169..3eafc756ef1 100644
--- a/client/cody/webviews/NavBar.tsx
+++ b/client/cody/webviews/NavBar.tsx
@@ -1,4 +1,6 @@
-import './NavBar.css'
+import React from 'react'
+
+import styles from './NavBar.module.css'
 
 export type View = 'chat' | 'recipes' | 'login' | 'settings' | 'debug' | 'history'
 
@@ -27,16 +29,16 @@ export const NavBar: React.FunctionComponent<React.PropsWithChildren<NavBarProps
     onResetClick,
     showResetButton,
 }) => (
-    <div className="tab-menu-container">
-        <div className="tab-menu-group">
+    <div className={styles.tabMenuContainer}>
+        <div className={styles.tabMenuGroup}>
             {navBarItems.map(({ title, tab }) => (
-                <button key={title} onClick={() => setView(tab)} className="tab-btn" type="button">
-                    <p className={view === tab ? 'tab-btn-selected' : ''}>{title}</p>
+                <button key={title} onClick={() => setView(tab)} className={styles.tabBtn} type="button">
+                    <p className={view === tab ? styles.tabBtnSelected : ''}>{title}</p>
                 </button>
             ))}
             {devMode && (
-                <button onClick={() => setView('debug')} className="tab-btn" type="button">
-                    <p className={view === 'debug' ? 'tab-btn-selected' : ''}>Debug</p>
+                <button onClick={() => setView('debug')} className={styles.tabBtn} type="button">
+                    <p className={view === 'debug' ? styles.tabBtnSelected : ''}>Debug</p>
                 </button>
             )}
         </div>
@@ -44,7 +46,7 @@ export const NavBar: React.FunctionComponent<React.PropsWithChildren<NavBarProps
             {showResetButton && (
                 <button
                     onClick={() => onResetClick()}
-                    className="tab-btn"
+                    className={styles.tabBtn}
                     type="button"
                     title="Start a new conversation"
                 >
diff --git a/client/cody/webviews/Recipes.css b/client/cody/webviews/Recipes.module.css
similarity index 100%
rename from client/cody/webviews/Recipes.css
rename to client/cody/webviews/Recipes.module.css
diff --git a/client/cody/webviews/Recipes.tsx b/client/cody/webviews/Recipes.tsx
index 9d932febd27..95b0e0a7504 100644
--- a/client/cody/webviews/Recipes.tsx
+++ b/client/cody/webviews/Recipes.tsx
@@ -1,8 +1,8 @@
 import { VSCodeButton } from '@vscode/webview-ui-toolkit/react'
 
-import { vscodeAPI } from './utils/VSCodeApi'
+import { VSCodeWrapper } from './utils/VSCodeApi'
 
-import './Recipes.css'
+import styles from './Recipes.module.css'
 
 export const recipesList = {
     'explain-code-detailed': 'Explain selected code (detailed)',
@@ -16,7 +16,7 @@ export const recipesList = {
     fixup: 'Fixup code from inline instructions',
 }
 
-export function Recipes(): React.ReactElement {
+export const Recipes: React.FunctionComponent<{ vscodeAPI: VSCodeWrapper }> = ({ vscodeAPI }) => {
     const onRecipeClick = (recipeID: string): void => {
         vscodeAPI.postMessage({ command: 'executeRecipe', recipe: recipeID })
     }
@@ -24,11 +24,11 @@ export function Recipes(): React.ReactElement {
     return (
         <div className="inner-container">
             <div className="non-transcript-container">
-                <div className="recipes">
+                <div className={styles.recipes}>
                     {Object.entries(recipesList).map(([key, value]) => (
                         <VSCodeButton
                             key={key}
-                            className="recipe-button"
+                            className={styles.recipeButton}
                             type="button"
                             onClick={() => onRecipeClick(key)}
                         >
diff --git a/client/cody/webviews/UserHistory.tsx b/client/cody/webviews/UserHistory.tsx
index 3b046a3078b..ad8651e7f6c 100644
--- a/client/cody/webviews/UserHistory.tsx
+++ b/client/cody/webviews/UserHistory.tsx
@@ -11,7 +11,7 @@ import { ChatHistory, ChatMessage } from '@sourcegraph/cody-shared/src/chat/tran
 import { ContextFiles } from '@sourcegraph/cody-ui/src/chat/ContextFiles'
 
 import { FileLink } from './FileLink'
-import { vscodeAPI } from './utils/VSCodeApi'
+import { VSCodeWrapper } from './utils/VSCodeApi'
 
 import styles from './Chat.module.css'
 
@@ -19,12 +19,14 @@ interface HistoryProps {
     userHistory: ChatHistory | null
     setUserHistory: (history: ChatHistory | null) => void
     setInputHistory: (inputHistory: string[] | []) => void
+    vscodeAPI: VSCodeWrapper
 }
 
 export const UserHistory: React.FunctionComponent<React.PropsWithChildren<HistoryProps>> = ({
     userHistory,
     setUserHistory,
     setInputHistory,
+    vscodeAPI,
 }) => {
     const [chatHistory, setChatHistory] = useState('')
 
@@ -35,7 +37,7 @@ export const UserHistory: React.FunctionComponent<React.PropsWithChildren<Histor
             setUserHistory(null)
             setInputHistory([])
         }
-    }, [setInputHistory, setUserHistory, userHistory])
+    }, [setInputHistory, setUserHistory, userHistory, vscodeAPI])
 
     return (
         <div className={styles.innerContainer}>
diff --git a/client/cody/webviews/index.tsx b/client/cody/webviews/index.tsx
index 7e0c33ae8e4..3eab5f7c879 100644
--- a/client/cody/webviews/index.tsx
+++ b/client/cody/webviews/index.tsx
@@ -6,8 +6,10 @@ import { App } from './App'
 
 import './index.css'
 
+import { getVSCodeAPI } from './utils/VSCodeApi'
+
 ReactDOM.createRoot(document.querySelector('#root') as HTMLElement).render(
     <React.StrictMode>
-        <App />
+        <App vscodeAPI={getVSCodeAPI()} />
     </React.StrictMode>
 )
diff --git a/client/cody/webviews/storybook/VSCodeStoryDecorator.module.css b/client/cody/webviews/storybook/VSCodeStoryDecorator.module.css
new file mode 100644
index 00000000000..071542f2bac
--- /dev/null
+++ b/client/cody/webviews/storybook/VSCodeStoryDecorator.module.css
@@ -0,0 +1,34 @@
+/**
+* CSS styles for storybooks to define variables and otherwise resemble the VS Code.
+*/
+
+.container {
+    max-width: 600px;
+    margin: 2rem auto;
+    border: solid 1px #ffffff33;
+    font-size: var(--vscode-editor-font-size);
+    color: var(--vscode-editor-foreground);
+}
+
+.container a {
+    color: var(--vscode-link-color);
+}
+
+/**
+ * Define any --vscode-* CSS variables needed. These differ per theme and don't need to be
+ * an exact match for any particular VS Code theme. They just need to be good enough.
+ */
+
+:root {
+    --vscode-editor-foreground: var(--foreground);
+    --vscode-editor-background: #151c28;
+    --vscode-inputValidation-errorBorder: #ff0000;
+    --vscode-inputValidation-errorBackground: #ff000033;
+    --vscode-editor-font-size: 13px;
+    --vscode-link-color: #7777ff;
+    --vscode-tab-inactiveBackground: #00000033;
+    --vscode-tab-inactiveForeground: #ffffff77;
+    --vscode-tab-activeForeground: #ffffff;
+    --vscode-sideBarSectionHeader-border: #ffffff33;
+    --vscode-sideBarTitle-foreground: #ffffffcc;
+}
diff --git a/client/cody/webviews/storybook/VSCodeStoryDecorator.tsx b/client/cody/webviews/storybook/VSCodeStoryDecorator.tsx
new file mode 100644
index 00000000000..4d1e3d5b0cf
--- /dev/null
+++ b/client/cody/webviews/storybook/VSCodeStoryDecorator.tsx
@@ -0,0 +1,8 @@
+import { DecoratorFn } from '@storybook/react'
+
+import styles from './VSCodeStoryDecorator.module.css'
+
+/**
+ * A decorator for storybooks that makes them look like they're running in VS Code.
+ */
+export const VSCodeStoryDecorator: DecoratorFn = story => <div className={styles.container}>{story()}</div>
diff --git a/client/cody/webviews/utils/VSCodeApi.ts b/client/cody/webviews/utils/VSCodeApi.ts
index 23a3017f385..42c6bc0bbb1 100644
--- a/client/cody/webviews/utils/VSCodeApi.ts
+++ b/client/cody/webviews/utils/VSCodeApi.ts
@@ -8,20 +8,26 @@ interface VSCodeApi {
     postMessage: (message: unknown) => void
 }
 
-class VSCodeWrapper {
-    private readonly vscodeApi: VSCodeApi = acquireVsCodeApi()
-
-    public postMessage(message: WebviewMessage): void {
-        this.vscodeApi.postMessage(message)
-    }
-
-    public onMessage(callback: (message: ExtensionMessage) => void): () => void {
-        const listener = (event: MessageEvent<ExtensionMessage>): void => {
-            callback(event.data)
-        }
-        window.addEventListener('message', listener)
-        return () => window.removeEventListener('message', listener)
-    }
+export interface VSCodeWrapper {
+    postMessage(message: WebviewMessage): void
+    onMessage(callback: (message: ExtensionMessage) => void): () => void
 }
 
-export const vscodeAPI: VSCodeWrapper = new VSCodeWrapper()
+let api: VSCodeWrapper
+
+export function getVSCodeAPI(): VSCodeWrapper {
+    if (!api) {
+        const vsCodeApi = acquireVsCodeApi()
+        api = {
+            postMessage: message => vsCodeApi.postMessage(message),
+            onMessage: callback => {
+                const listener = (event: MessageEvent<ExtensionMessage>): void => {
+                    callback(event.data)
+                }
+                window.addEventListener('message', listener)
+                return () => window.removeEventListener('message', listener)
+            },
+        }
+    }
+    return api
+}
diff --git a/client/storybook/src/main.ts b/client/storybook/src/main.ts
index 3b563fa2f0a..f16559f92b8 100644
--- a/client/storybook/src/main.ts
+++ b/client/storybook/src/main.ts
@@ -43,9 +43,21 @@ const getStoriesGlob = (): string[] => {
     // Due to an issue with constant recompiling (https://github.com/storybookjs/storybook/issues/14342)
     // we need to make the globs more specific (`(web|shared..)` also doesn't work). Once the above issue
     // is fixed, this can be removed and watched for `client/**/*.story.tsx` again.
-    const directoriesWithStories = ['branded', 'browser', 'jetbrains/webview', 'shared', 'web', 'wildcard', 'cody-ui']
+    const directoriesWithStories = [
+        'branded',
+        'browser',
+        'jetbrains/webview',
+        'shared',
+        'web',
+        'wildcard',
+        'cody-ui',
+        'cody',
+    ]
     const storiesGlobs = directoriesWithStories.map(packageDirectory =>
-        path.resolve(ROOT_PATH, `client/${packageDirectory}/src/**/*.story.tsx`)
+        path.resolve(
+            ROOT_PATH,
+            `client/${packageDirectory}/${packageDirectory === 'cody' ? 'webviews' : 'src'}/**/*.story.tsx`
+        )
     )
 
     return [...storiesGlobs]