diff --git a/src/context.ts b/src/context.ts
index 6c1f70f..d18faa0 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -45,6 +45,10 @@ export class Context {
     return this._tabs;
   }
 
+  currentTab(): Tab | undefined {
+    return this._currentTab;
+  }
+
   currentTabOrDie(): Tab {
     if (!this._currentTab)
       throw new Error('No open pages available. Use the "browser_navigate" tool to navigate to a page first.');
diff --git a/src/response.ts b/src/response.ts
index a075eed..d0fc25b 100644
--- a/src/response.ts
+++ b/src/response.ts
@@ -23,7 +23,6 @@ export class Response {
   private _images: { contentType: string, data: Buffer }[] = [];
   private _context: Context;
   private _includeSnapshot = false;
-  private _snapshot: string | undefined;
   private _includeTabs = false;
 
   constructor(context: Context) {
@@ -54,10 +53,6 @@ export class Response {
     return this._includeSnapshot;
   }
 
-  addSnapshot(snapshot: string) {
-    this._snapshot = snapshot;
-  }
-
   async serialize(): Promise<{ content: (TextContent | ImageContent)[] }> {
     const response: string[] = [];
 
@@ -82,8 +77,8 @@ ${this._code.join('\n')}
       response.push(...(await this._context.listTabsMarkdown(this._includeTabs)));
 
     // Add snapshot if provided.
-    if (this._snapshot)
-      response.push(this._snapshot, '');
+    if (this._includeSnapshot && this._context.currentTab())
+      response.push(await this._context.currentTabOrDie().captureSnapshot(), '');
 
     // Main response part
     const content: (TextContent | ImageContent)[] = [
diff --git a/src/tab.ts b/src/tab.ts
index 3dc393e..ca3e101 100644
--- a/src/tab.ts
+++ b/src/tab.ts
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
+import { EventEmitter } from 'events';
 import * as playwright from 'playwright';
-
 import { callOnPageNoTrace, waitForCompletion } from './tools/utils.js';
 import { logUnhandledError } from './log.js';
 import { ManualPromise } from './manualPromise.js';
@@ -23,28 +23,31 @@ import { ModalState } from './tools/tool.js';
 import { outputFile } from './config.js';
 
 import type { Context } from './context.js';
-import type { Response } from './response.js';
 
 type PageEx = playwright.Page & {
   _snapshotForAI: () => Promise<string>;
 };
 
-type PendingAction = {
-  dialogShown: ManualPromise<void>;
+export const TabEvents = {
+  modalState: 'modalState'
 };
 
-export class Tab {
+export type TabEventsInterface = {
+  [TabEvents.modalState]: [modalState: ModalState];
+};
+
+export class Tab extends EventEmitter<TabEventsInterface> {
   readonly context: Context;
   readonly page: playwright.Page;
   private _consoleMessages: ConsoleMessage[] = [];
   private _recentConsoleMessages: ConsoleMessage[] = [];
-  private _pendingAction: PendingAction | undefined;
   private _requests: Map<playwright.Request, playwright.Response | null> = new Map();
   private _onPageClose: (tab: Tab) => void;
   private _modalStates: ModalState[] = [];
   private _downloads: { download: playwright.Download, finished: boolean, outputFile: string }[] = [];
 
   constructor(context: Context, page: playwright.Page, onPageClose: (tab: Tab) => void) {
+    super();
     this.context = context;
     this.page = page;
     this._onPageClose = onPageClose;
@@ -74,6 +77,7 @@ export class Tab {
 
   setModalState(modalState: ModalState) {
     this._modalStates.push(modalState);
+    this.emit(TabEvents.modalState, modalState);
   }
 
   clearModalState(modalState: ModalState) {
@@ -97,7 +101,6 @@ export class Tab {
       description: `"${dialog.type()}" dialog with message "${dialog.message()}"`,
       dialog,
     });
-    this._pendingAction?.dialogShown.resolve();
   }
 
   private async _downloadStarted(download: playwright.Download) {
@@ -196,7 +199,7 @@ export class Tab {
     return result;
   }
 
-  async captureSnapshot(options: { omitAriaSnapshot?: boolean } = {}): Promise<string> {
+  async captureSnapshot(): Promise<string> {
     const result: string[] = [];
     if (this.modalStates().length) {
       result.push(...this.modalStatesMarkdown());
@@ -205,19 +208,19 @@ export class Tab {
 
     result.push(...this._takeRecentConsoleMarkdown());
     result.push(...this._listDownloadsMarkdown());
-    if (options.omitAriaSnapshot)
-      return result.join('\n');
 
-    const snapshot = await (this.page as PageEx)._snapshotForAI();
-    result.push(
-        `### Page state`,
-        `- Page URL: ${this.page.url()}`,
-        `- Page Title: ${await this.page.title()}`,
-        `- Page Snapshot:`,
-        '```yaml',
-        snapshot,
-        '```',
-    );
+    await this._raceAgainstModalStates(async () => {
+      const snapshot = await (this.page as PageEx)._snapshotForAI();
+      result.push(
+          `### Page state`,
+          `- Page URL: ${this.page.url()}`,
+          `- Page Title: ${await this.page.title()}`,
+          `- Page Snapshot:`,
+          '```yaml',
+          snapshot,
+          '```',
+      );
+    });
     return result.join('\n');
   }
 
@@ -225,42 +228,25 @@ export class Tab {
     return this._modalStates.some(state => state.type === 'dialog');
   }
 
-  private async _raceAgainstModalDialogs<R>(action: () => Promise<R>): Promise<R | undefined> {
-    this._pendingAction = {
-      dialogShown: new ManualPromise(),
-    };
+  private async _raceAgainstModalStates(action: () => Promise<void>): Promise<ModalState | undefined> {
+    if (this.modalStates().length)
+      return this.modalStates()[0];
 
-    let result: R | undefined;
-    try {
-      await Promise.race([
-        action().then(r => result = r),
-        this._pendingAction.dialogShown,
-      ]);
-    } finally {
-      this._pendingAction = undefined;
-    }
-    return result;
+    const promise = new ManualPromise<ModalState>();
+    const listener = (modalState: ModalState) => promise.resolve(modalState);
+    this.once(TabEvents.modalState, listener);
+
+    return await Promise.race([
+      action().then(() => {
+        this.off(TabEvents.modalState, listener);
+        return undefined;
+      }),
+      promise,
+    ]);
   }
 
-  async run(callback: () => Promise<void>, response: Response) {
-    let snapshot: string | undefined;
-    await this._raceAgainstModalDialogs(async () => {
-      try {
-        if (response.includeSnapshot())
-          await waitForCompletion(this, callback);
-        else
-          await callback();
-      } finally {
-        snapshot = await this.captureSnapshot();
-      }
-    });
-
-    if (snapshot) {
-      response.addSnapshot(snapshot);
-    } else if (response.includeSnapshot()) {
-      // We are blocked on modal dialog.
-      response.addSnapshot(await this.captureSnapshot({ omitAriaSnapshot: true }));
-    }
+  async waitForCompletion(callback: () => Promise<void>) {
+    await this._raceAgainstModalStates(() => waitForCompletion(this, callback));
   }
 
   async refLocator(params: { element: string, ref: string }): Promise<playwright.Locator> {
diff --git a/src/tools/common.ts b/src/tools/common.ts
index b559c70..ba9847a 100644
--- a/src/tools/common.ts
+++ b/src/tools/common.ts
@@ -52,9 +52,9 @@ const resize = defineTabTool({
     response.addCode(`// Resize browser window to ${params.width}x${params.height}`);
     response.addCode(`await page.setViewportSize({ width: ${params.width}, height: ${params.height} });`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await tab.page.setViewportSize({ width: params.width, height: params.height });
-    }, response);
+    });
   },
 });
 
diff --git a/src/tools/dialogs.ts b/src/tools/dialogs.ts
index a5ccf18..c8eb57b 100644
--- a/src/tools/dialogs.ts
+++ b/src/tools/dialogs.ts
@@ -39,12 +39,12 @@ const handleDialog = defineTabTool({
       throw new Error('No dialog visible');
 
     tab.clearModalState(dialogState);
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       if (params.accept)
         await dialogState.dialog.accept(params.promptText);
       else
         await dialogState.dialog.dismiss();
-    }, response);
+    });
   },
 
   clearsModalState: 'dialog',
diff --git a/src/tools/evaluate.ts b/src/tools/evaluate.ts
index 8452249..3023f39 100644
--- a/src/tools/evaluate.ts
+++ b/src/tools/evaluate.ts
@@ -49,11 +49,11 @@ const evaluate = defineTabTool({
       response.addCode(`await page.evaluate(${javascript.quote(params.function)});`);
     }
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       const receiver = locator ?? tab.page as any;
       const result = await receiver._evaluateFunction(params.function);
       response.addResult(JSON.stringify(result, null, 2) || 'undefined');
-    }, response);
+    });
   },
 });
 
diff --git a/src/tools/files.ts b/src/tools/files.ts
index 6ecf43b..939aa3c 100644
--- a/src/tools/files.ts
+++ b/src/tools/files.ts
@@ -40,10 +40,10 @@ const uploadFile = defineTabTool({
     response.addCode(`// Select files for upload`);
     response.addCode(`await fileChooser.setFiles(${JSON.stringify(params.paths)})`);
 
-    await tab.run(async () => {
+    tab.clearModalState(modalState);
+    await tab.waitForCompletion(async () => {
       await modalState.fileChooser.setFiles(params.paths);
-      tab.clearModalState(modalState);
-    }, response);
+    });
   },
   clearsModalState: 'fileChooser',
 });
diff --git a/src/tools/keyboard.ts b/src/tools/keyboard.ts
index 1e143f7..9286f5c 100644
--- a/src/tools/keyboard.ts
+++ b/src/tools/keyboard.ts
@@ -39,9 +39,9 @@ const pressKey = defineTabTool({
     response.addCode(`// Press ${params.key}`);
     response.addCode(`await page.keyboard.press('${params.key}');`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await tab.page.keyboard.press(params.key);
-    }, response);
+    });
   },
 });
 
@@ -66,7 +66,7 @@ const type = defineTabTool({
 
     const locator = await tab.refLocator(params);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       if (params.slowly) {
         response.addCode(`// Press "${params.text}" sequentially into "${params.element}"`);
         response.addCode(`await page.${await generateLocator(locator)}.pressSequentially(${javascript.quote(params.text)});`);
@@ -81,7 +81,7 @@ const type = defineTabTool({
         response.addCode(`await page.${await generateLocator(locator)}.press('Enter');`);
         await locator.press('Enter');
       }
-    }, response);
+    });
   },
 });
 
diff --git a/src/tools/mouse.ts b/src/tools/mouse.ts
index a86e90c..3889df2 100644
--- a/src/tools/mouse.ts
+++ b/src/tools/mouse.ts
@@ -38,9 +38,9 @@ const mouseMove = defineTabTool({
     response.addCode(`// Move mouse to (${params.x}, ${params.y})`);
     response.addCode(`await page.mouse.move(${params.x}, ${params.y});`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await tab.page.mouse.move(params.x, params.y);
-    }, response);
+    });
   },
 });
 
@@ -65,11 +65,11 @@ const mouseClick = defineTabTool({
     response.addCode(`await page.mouse.down();`);
     response.addCode(`await page.mouse.up();`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await tab.page.mouse.move(params.x, params.y);
       await tab.page.mouse.down();
       await tab.page.mouse.up();
-    }, response);
+    });
   },
 });
 
@@ -97,12 +97,12 @@ const mouseDrag = defineTabTool({
     response.addCode(`await page.mouse.move(${params.endX}, ${params.endY});`);
     response.addCode(`await page.mouse.up();`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await tab.page.mouse.move(params.startX, params.startY);
       await tab.page.mouse.down();
       await tab.page.mouse.move(params.endX, params.endY);
       await tab.page.mouse.up();
-    }, response);
+    });
   },
 });
 
diff --git a/src/tools/navigate.ts b/src/tools/navigate.ts
index afcbffc..ea09930 100644
--- a/src/tools/navigate.ts
+++ b/src/tools/navigate.ts
@@ -34,9 +34,9 @@ const navigate = defineTool({
     const tab = await context.ensureTab();
     await tab.navigate(params.url);
 
+    response.setIncludeSnapshot();
     response.addCode(`// Navigate to ${params.url}`);
     response.addCode(`await page.goto('${params.url}');`);
-    response.addSnapshot(await tab.captureSnapshot());
   },
 });
 
@@ -54,9 +54,9 @@ const goBack = defineTabTool({
     response.setIncludeSnapshot();
 
     await tab.page.goBack();
+    response.setIncludeSnapshot();
     response.addCode(`// Navigate back`);
     response.addCode(`await page.goBack();`);
-    response.addSnapshot(await tab.captureSnapshot());
   },
 });
 
@@ -73,9 +73,9 @@ const goForward = defineTabTool({
     response.setIncludeSnapshot();
 
     await tab.page.goForward();
+    response.setIncludeSnapshot();
     response.addCode(`// Navigate forward`);
     response.addCode(`await page.goForward();`);
-    response.addSnapshot(await tab.captureSnapshot());
   },
 });
 
diff --git a/src/tools/snapshot.ts b/src/tools/snapshot.ts
index e855ad4..083f0a7 100644
--- a/src/tools/snapshot.ts
+++ b/src/tools/snapshot.ts
@@ -31,8 +31,8 @@ const snapshot = defineTool({
   },
 
   handle: async (context, params, response) => {
-    const tab = await context.ensureTab();
-    response.addSnapshot(await tab.captureSnapshot());
+    await context.ensureTab();
+    response.setIncludeSnapshot();
   },
 });
 
@@ -71,12 +71,12 @@ const click = defineTabTool({
       response.addCode(`await page.${await generateLocator(locator)}.click(${buttonAttr});`);
     }
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       if (params.doubleClick)
         await locator.dblclick({ button });
       else
         await locator.click({ button });
-    }, response);
+    });
   },
 });
 
@@ -103,9 +103,9 @@ const drag = defineTabTool({
       { ref: params.endRef, element: params.endElement },
     ]);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await startLocator.dragTo(endLocator);
-    }, response);
+    });
 
     response.addCode(`await page.${await generateLocator(startLocator)}.dragTo(page.${await generateLocator(endLocator)});`);
   },
@@ -127,9 +127,9 @@ const hover = defineTabTool({
     const locator = await tab.refLocator(params);
     response.addCode(`await page.${await generateLocator(locator)}.hover();`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await locator.hover();
-    }, response);
+    });
   },
 });
 
@@ -154,9 +154,9 @@ const selectOption = defineTabTool({
     response.addCode(`// Select options [${params.values.join(', ')}] in ${params.element}`);
     response.addCode(`await page.${await generateLocator(locator)}.selectOption(${javascript.formatObject(params.values)});`);
 
-    await tab.run(async () => {
+    await tab.waitForCompletion(async () => {
       await locator.selectOption(params.values);
-    }, response);
+    });
   },
 });
 
diff --git a/src/tools/tabs.ts b/src/tools/tabs.ts
index ff1a8b7..048cb28 100644
--- a/src/tools/tabs.ts
+++ b/src/tools/tabs.ts
@@ -48,9 +48,8 @@ const selectTab = defineTool({
   },
 
   handle: async (context, params, response) => {
-    const tab = await context.selectTab(params.index);
+    await context.selectTab(params.index);
     response.setIncludeSnapshot();
-    response.addSnapshot(await tab.captureSnapshot());
   },
 });
 
@@ -71,9 +70,7 @@ const newTab = defineTool({
     const tab = await context.newTab();
     if (params.url)
       await tab.navigate(params.url);
-
     response.setIncludeSnapshot();
-    response.addSnapshot(await tab.captureSnapshot());
   },
 });
 
@@ -92,10 +89,7 @@ const closeTab = defineTool({
 
   handle: async (context, params, response) => {
     await context.closeTab(params.index);
-    response.setIncludeTabs();
-    response.addCode(`await myPage.close();`);
-    if (context.tabs().length)
-      response.addSnapshot(await context.currentTabOrDie().captureSnapshot());
+    response.setIncludeSnapshot();
   },
 });
 
diff --git a/src/tools/wait.ts b/src/tools/wait.ts
index a83d022..0d5e59d 100644
--- a/src/tools/wait.ts
+++ b/src/tools/wait.ts
@@ -58,7 +58,7 @@ const wait = defineTool({
     }
 
     response.addResult(`Waited for ${params.text || params.textGone || params.time}`);
-    response.addSnapshot(await tab.captureSnapshot());
+    response.setIncludeSnapshot();
   },
 });