diff --git a/extension/.vscodeignore b/extension/.vscodeignore
index ed3f9d3..c7061af 100644
--- a/extension/.vscodeignore
+++ b/extension/.vscodeignore
@@ -1,10 +1,2 @@
-.vscode/**
-.vscode-test/**
-out/test/**
-src/**
-.gitignore
-vsc-extension-quickstart.md
-**/tsconfig.json
-**/tslint.json
-**/*.map
-**/*.ts
\ No newline at end of file
+node_modules
+ui/node_modules
\ No newline at end of file
diff --git a/extension/src/WebViews.ts b/extension/src/WebViews.ts
index ada7fb7..db4c60f 100644
--- a/extension/src/WebViews.ts
+++ b/extension/src/WebViews.ts
@@ -73,6 +73,19 @@ export function getHtml(server: Server) {
 			<body>
 				<script>
 					window.serverPort = ${server.port};
+					const api = window.VsCodeApi = acquireVsCodeApi();
+					window.addEventListener('message', event => {
+						if (event.source === window.frames[0]) {
+							if (event.data.command === "setState") {
+								console.log("setState", event.data.state);
+								api.setState(event.data.state);
+							}
+							if (event.data.command === "getState") {
+								console.log("getState, sent ", api.getState());
+								window.frames[0].postMessage({ command: "getStateResult", state: api.getState() }, "*");
+							}
+						}
+					});
 				</script>
 				${
 					isDev
@@ -83,22 +96,6 @@ export function getHtml(server: Server) {
 								server.mainBundle
 						  }"></script>`
 				}
-                <script>
-                const api = window.api = acquireVsCodeApi();
-                window.addEventListener('message', event => {
-                    debugger;
-                    if (event.source === window.frames[0]) {
-                        if (event.data.command === "setState") {
-                            console.log("setState", event.data.state);
-                            api.setState(event.data.state);
-                        }
-                        if (event.data.command === "getState") {
-                            console.log("getState, sent ", api.getState());
-                            window.frames[0].postMessage({ command: "getStateResult", state: api.getState() }, "*");
-                        }
-                    }
-                });
-                </script>
             </body>
         </html>
     `;
diff --git a/extension/ui/src/Model.ts b/extension/ui/src/Model.ts
index ae53e1c..95a0b6b 100644
--- a/extension/ui/src/Model.ts
+++ b/extension/ui/src/Model.ts
@@ -10,8 +10,14 @@ import { DataExtractorId } from "@hediet/debug-visualizer-data-extraction";
 import { Visualization, VisualizationId } from "./Visualizers/Visualizer";
 import { knownVisualizations } from "./Visualizers";
 
+interface VsCodeApi {
+	getState<TState>(): TState | undefined;
+	setState(state: unknown): void;
+}
+
 declare const window: Window & {
 	serverPort?: number;
+	VsCodeApi?: VsCodeApi;
 };
 
 interface VSCodeApi<TState> {
@@ -61,6 +67,18 @@ class IFrameVSCodeApi<TState> implements VSCodeApi<TState> {
 	}
 }
 
+class DirectVSCodeApi<TState> implements VSCodeApi<TState> {
+	constructor(private readonly vsCodeApi: VsCodeApi) {}
+
+	async getState(): Promise<TState | undefined> {
+		return this.vsCodeApi.getState<TState | undefined>();
+	}
+
+	async setState(state: TState): Promise<void> {
+		this.vsCodeApi.setState(state);
+	}
+}
+
 export class Model {
 	port: number;
 
@@ -99,7 +117,9 @@ export class Model {
 		| typeof debugVisualizerUIContract.TServerInterface
 		| undefined = undefined;
 
-	private readonly vsCodeApi = new IFrameVSCodeApi<{ expression: string }>();
+	private readonly vsCodeApi = window.VsCodeApi
+		? new DirectVSCodeApi<{ expression: string }>(window.VsCodeApi)
+		: new IFrameVSCodeApi<{ expression: string }>();
 
 	@observable private _loading = false;