diff --git a/src/McpContext.ts b/src/McpContext.ts
index 6651e7cd..9952a048 100644
--- a/src/McpContext.ts
+++ b/src/McpContext.ts
@@ -617,17 +617,11 @@ export class McpContext implements Context {
     return this.#networkCollector.getIdForResource(request);
   }
 
-  waitForTextOnPage({
-    text,
-    timeout,
-  }: {
-    text: string;
-    timeout?: number | undefined;
-  }): Promise<Element> {
+  waitForTextOnPage(text: string, timeout?: number): Promise<Element> {
     const page = this.getSelectedPage();
     const frames = page.frames();
 
-    const locator = this.#locatorClass.race(
+    let locator = this.#locatorClass.race(
       frames.flatMap(frame => [
         frame.locator(`aria/${text}`),
         frame.locator(`text/${text}`),
@@ -635,7 +629,7 @@ export class McpContext implements Context {
     );
 
     if (timeout) {
-      locator.setTimeout(timeout);
+      locator = locator.setTimeout(timeout);
     }
 
     return locator.wait();
diff --git a/src/tools/ToolDefinition.ts b/src/tools/ToolDefinition.ts
index 8acebb56..1586c2d7 100644
--- a/src/tools/ToolDefinition.ts
+++ b/src/tools/ToolDefinition.ts
@@ -107,10 +107,7 @@ export type Context = Readonly<{
     filename: string,
   ): Promise<{filename: string}>;
   waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
-  waitForTextOnPage(params: {
-    text: string;
-    timeout?: number | undefined;
-  }): Promise<Element>;
+  waitForTextOnPage(text: string, timeout?: number): Promise<Element>;
   getDevToolsData(): Promise<DevToolsData>;
   /**
    * Returns a reqid for a cdpRequestId.
diff --git a/src/tools/snapshot.ts b/src/tools/snapshot.ts
index 21099f14..143d0409 100644
--- a/src/tools/snapshot.ts
+++ b/src/tools/snapshot.ts
@@ -53,7 +53,10 @@ export const waitFor = defineTool({
     ...timeoutSchema,
   },
   handler: async (request, response, context) => {
-    await context.waitForTextOnPage(request.params);
+    await context.waitForTextOnPage(
+      request.params.text,
+      request.params.timeout,
+    );
 
     response.appendResponseLine(
       `Element with text "${request.params.text}" found.`,