refactor(genie): extract pollWaiter to simplify _handleSendMessage

Replace manual concurrency code (statusQueue, notifyGenerator,
waiterDone, waiterError, IIFE promise chain) with a reusable
pollWaiter async generator that bridges callback-based
waiter.wait({ onProgress }) into a for-await-of loop.

Signed-off-by: Jorge Calvar <jorge.calvar@databricks.com>

Author: calvarjorge

packages/appkit/src/plugins/genie/genie.ts

@@ -11,6 +11,7 @@ import { createLogger } from "../../logging";
 import { Plugin, toPlugin } from "../../plugin";
 import { genieStreamDefaults } from "./defaults";
 import { genieManifest } from "./manifest";
+import { pollWaiter } from "./poll-waiter";
 import type {
   GenieAttachmentResponse,
   GenieConversationHistoryResponse,
@@ -155,110 +156,60 @@ export class GeniePlugin extends Plugin {
         const workspaceClient = getWorkspaceClient();
 
         try {
-          // Status events queue bridging onProgress → generator
-          const statusQueue: string[] = [];
-          let notifyGenerator: () => void = () => {};
-          let waiterDone = false;
-
-          const onProgress = async (message: GenieMessage): Promise<void> => {
-            if (message.status) {
-              statusQueue.push(message.status);
-              notifyGenerator();
-            }
-          };
-
-          let resultConversationId = "";
-          let resultMessageId = "";
-          let completedMessage: GenieMessage =
-            undefined as unknown as GenieMessage;
-          let waiterError: Error | null = null;
-
-          // Launch Genie API call
-          const waiterPromise = (async () => {
-            let messageWaiter: CreateMessageWaiter;
-
-            if (conversationId) {
-              messageWaiter = await workspaceClient.genie.createMessage({
+          // Step 1: API call → get waiter + IDs
+          let messageWaiter: CreateMessageWaiter;
+          let resultConversationId: string;
+          let resultMessageId: string;
+
+          if (conversationId) {
+            messageWaiter = await workspaceClient.genie.createMessage({
+              space_id: spaceId,
+              conversation_id: conversationId,
+              content,
+            });
+            resultConversationId = conversationId;
+            resultMessageId = messageWaiter.message_id ?? "";
+          } else {
+            const startWaiter: StartConversationWaiter =
+              await workspaceClient.genie.startConversation({
                 space_id: spaceId,
-                conversation_id: conversationId,
                 content,
               });
-              resultConversationId = conversationId;
-            } else {
-              const startWaiter: StartConversationWaiter =
-                await workspaceClient.genie.startConversation({
-                  space_id: spaceId,
-                  content,
-                });
-              resultConversationId = startWaiter.conversation_id;
-              resultMessageId = startWaiter.message_id;
-              messageWaiter = startWaiter as unknown as CreateMessageWaiter;
-            }
-
-            const result = await messageWaiter.wait({ onProgress });
-            completedMessage = result;
-            resultMessageId = result.message_id;
-            return result;
-          })()
-            .catch((err: Error) => {
-              waiterError = err;
-            })
-            .finally(() => {
-              waiterDone = true;
-              notifyGenerator();
-            });
-
-          // Wait for first status or waiter completion to get IDs
-          await new Promise<void>((resolve) => {
-            notifyGenerator = resolve;
-            if (waiterDone) resolve();
-          });
-
-          // If the API call failed before anything started, yield error and exit
-          if (waiterError) {
-            throw waiterError;
+            resultConversationId = startWaiter.conversation_id;
+            resultMessageId = startWaiter.message_id;
+            messageWaiter = startWaiter as unknown as CreateMessageWaiter;
           }
 
-          // Yield message_start
+          // Step 2: Yield message_start immediately — IDs are available from API response
           yield {
             type: "message_start" as const,
             conversationId: resultConversationId,
             messageId: resultMessageId,
             spaceId,
           };
 
-          // Drain status events
-          while (!waiterDone || statusQueue.length > 0) {
-            while (statusQueue.length > 0) {
-              const status = statusQueue.shift();
-              if (status) {
-                yield { type: "status" as const, status };
+          // Step 3: Poll for status updates and completion
+          let completedMessage!: GenieMessage;
+          for await (const event of pollWaiter(messageWaiter)) {
+            if (event.type === "progress") {
+              if (event.value.status) {
+                yield { type: "status" as const, status: event.value.status };
               }
+            } else {
+              completedMessage = event.value;
+              resultMessageId = event.value.message_id;
             }
-
-            if (!waiterDone) {
-              await new Promise<void>((resolve) => {
-                notifyGenerator = resolve;
-                if (waiterDone) resolve();
-              });
-            }
-          }
-
-          // Check if waiter failed during polling
-          await waiterPromise;
-          if (waiterError) {
-            throw waiterError;
           }
 
-          // Build cleaned message response
+          // Step 4: Build cleaned message response
           const messageResponse = toMessageResponse(completedMessage);
 
           yield {
             type: "message_result" as const,
             message: messageResponse,
           };
 
-          // Fetch query results for each query attachment
+          // Step 5: Fetch query results for each query attachment
           const attachments = messageResponse.attachments ?? [];
           for (const att of attachments) {
             if (att.query?.statementId && att.attachmentId) {

packages/appkit/src/plugins/genie/poll-waiter.ts

@@ -0,0 +1,92 @@
+/**
+ * Structural interface matching the SDK's `Waiter.wait()` shape
+ * without importing the SDK directly.
+ */
+export interface Pollable<P> {
+  wait(options?: {
+    onProgress?: (p: P) => Promise<void>;
+    timeout?: unknown;
+  }): Promise<P>;
+}
+
+export type PollEvent<P> =
+  | { type: "progress"; value: P }
+  | { type: "completed"; value: P };
+
+/**
+ * Bridges a callback-based waiter into an async generator.
+ *
+ * The SDK's `waiter.wait({ onProgress })` API uses a callback to report
+ * progress and returns a promise that resolves with the final result.
+ * This function converts that push-based model into a pull-based async
+ * generator so callers can simply `for await (const event of pollWaiter(w))`.
+ *
+ * Yields `{ type: "progress", value }` for each `onProgress` callback,
+ * then `{ type: "completed", value }` for the final result.
+ * Throws if the waiter rejects.
+ */
+export async function* pollWaiter<P>(
+  waiter: Pollable<P>,
+  options?: { timeout?: unknown },
+): AsyncGenerator<PollEvent<P>> {
+  // --- shared state between the onProgress callback and the generator loop ---
+  const queue: P[] = []; // progress values waiting to be yielded
+  let notify: () => void = () => {}; // resolves the generator's "sleep" promise
+  let done = false; // true once waiter.wait() settles (success or error)
+  let result!: P;
+  let error: unknown = null;
+
+  // Start the waiter in the background (not awaited — runs concurrently
+  // with the generator loop below). The onProgress callback pushes values
+  // into the queue and wakes the generator via notify().
+  waiter
+    .wait({
+      onProgress: async (p: P) => {
+        queue.push(p);
+        notify();
+      },
+      ...(options?.timeout != null ? { timeout: options.timeout } : {}),
+    })
+    .then((r) => {
+      result = r;
+      done = true;
+      notify();
+    })
+    .catch((err) => {
+      error = err;
+      done = true;
+      notify();
+    });
+
+  // Drain progress events as they arrive. The loop exits once the waiter
+  // has settled AND the queue is empty.
+  while (!done || queue.length > 0) {
+    // Yield all queued progress values before sleeping.
+    while (queue.length > 0) {
+      const value = queue.shift() as P;
+      yield { type: "progress", value };
+    }
+
+    // Nothing in the queue yet and the waiter hasn't settled — sleep until
+    // the next onProgress call or waiter settlement wakes us via notify().
+    //
+    // Race-condition guard: after setting `notify = resolve`, we re-check
+    // `done` and `queue.length`. If either changed between the outer while
+    // check and this point (possible via microtask), we resolve immediately
+    // so the loop doesn't hang.
+    if (!done) {
+      await new Promise<void>((resolve) => {
+        notify = resolve;
+        if (done || queue.length > 0) resolve();
+      });
+    }
+  }
+
+  // The waiter settled. If it rejected, propagate the error.
+  if (error !== null) {
+    throw error;
+  }
+
+  // Final event: the completed result from waiter.wait().
+  yield { type: "completed", value: result };
+}

packages/appkit/src/plugins/genie/tests/poll-waiter.test.ts

@@ -0,0 +1,125 @@
+import { describe, expect, test, vi } from "vitest";
+import { type Pollable, type PollEvent, pollWaiter } from "../poll-waiter";
+
+function createMockWaiter<P>(opts: {
+  progressValues?: P[];
+  result: P;
+  error?: Error;
+  delay?: number;
+}): Pollable<P> {
+  return {
+    wait: vi.fn().mockImplementation(async (options: any = {}) => {
+      if (opts.progressValues) {
+        for (const value of opts.progressValues) {
+          if (opts.delay) {
+            await new Promise((r) => setTimeout(r, opts.delay));
+          }
+          if (options.onProgress) {
+            await options.onProgress(value);
+          }
+        }
+      }
+      if (opts.error) throw opts.error;
+      return opts.result;
+    }),
+  };
+}
+
+async function collect<P>(
+  gen: AsyncGenerator<PollEvent<P>>,
+): Promise<PollEvent<P>[]> {
+  const events: PollEvent<P>[] = [];
+  for await (const event of gen) {
+    events.push(event);
+  }
+  return events;
+}
+
+describe("pollWaiter", () => {
+  test("yields progress events then completed", async () => {
+    const waiter = createMockWaiter({
+      progressValues: [{ status: "A" }, { status: "B" }],
+      result: { status: "DONE" },
+    });
+
+    const events = await collect(pollWaiter(waiter));
+
+    expect(events).toEqual([
+      { type: "progress", value: { status: "A" } },
+      { type: "progress", value: { status: "B" } },
+      { type: "completed", value: { status: "DONE" } },
+    ]);
+  });
+
+  test("yields only completed when no progress events", async () => {
+    const waiter = createMockWaiter({
+      result: { value: 42 },
+    });
+
+    const events = await collect(pollWaiter(waiter));
+
+    expect(events).toEqual([{ type: "completed", value: { value: 42 } }]);
+  });
+
+  test("throws when waiter rejects", async () => {
+    const waiter = createMockWaiter({
+      result: null as any,
+      error: new Error("boom"),
+    });
+
+    const events: PollEvent<any>[] = [];
+    await expect(async () => {
+      for await (const event of pollWaiter(waiter)) {
+        events.push(event);
+      }
+    }).rejects.toThrow("boom");
+
+    expect(events).toEqual([]);
+  });
+
+  test("throws after yielding progress if waiter fails mid-poll", async () => {
+    const waiter = createMockWaiter({
+      progressValues: [{ status: "A" }],
+      result: null as any,
+      error: new Error("mid-poll failure"),
+    });
+
+    const events: PollEvent<any>[] = [];
+    await expect(async () => {
+      for await (const event of pollWaiter(waiter)) {
+        events.push(event);
+      }
+    }).rejects.toThrow("mid-poll failure");
+
+    expect(events).toEqual([{ type: "progress", value: { status: "A" } }]);
+  });
+
+  test("handles async delays between progress callbacks", async () => {
+    const waiter = createMockWaiter({
+      progressValues: [{ n: 1 }, { n: 2 }, { n: 3 }],
+      result: { n: 99 },
+      delay: 10,
+    });
+
+    const events = await collect(pollWaiter(waiter));
+
+    expect(events).toHaveLength(4);
+    expect(events[0]).toEqual({ type: "progress", value: { n: 1 } });
+    expect(events[1]).toEqual({ type: "progress", value: { n: 2 } });
+    expect(events[2]).toEqual({ type: "progress", value: { n: 3 } });
+    expect(events[3]).toEqual({ type: "completed", value: { n: 99 } });
+  });
+
+  test("passes timeout option through to waiter.wait()", async () => {
+    const waiter = createMockWaiter({
+      result: { done: true },
+    });
+
+    const timeoutValue = { ms: 5000 };
+    await collect(pollWaiter(waiter, { timeout: timeoutValue }));
+
+    expect(waiter.wait).toHaveBeenCalledWith(
+      expect.objectContaining({ timeout: timeoutValue }),
+    );
+  });
+});