fix(wizards): use materialized cohort id response

plugins/ui/apps/wizards/src/api/__tests__/wizardCohortApi.test.ts

@@ -96,17 +96,21 @@ describe("Wizard cohort API", () => {
   });
 
   it("materializes a Wizard bookmark with the existing cohort payload", async () => {
-    mockedPost.mockResolvedValue({ data: "Cohort successfully materialized" });
+    mockedPost.mockResolvedValue({
+      data: { message: "Cohort successfully materialized", cohortDefinitionId: 42 },
+    });
     const bookmark = buildMriBookmark([], {}, { configId: "pa-config", configVersion: "7" }, "dataset-1").bookmark;
     const mriQuery = buildMriMaterializationQuery(bookmark, "dataset-1");
 
-    await materializeWizardBookmark({
-      datasetId: "dataset-1",
-      bookmarkId: "bookmark-1",
-      bookmarkName: "wizards-1783670400000",
-      description: "Generated by Wizards",
-      mriQuery,
-    });
+    await expect(
+      materializeWizardBookmark({
+        datasetId: "dataset-1",
+        bookmarkId: "bookmark-1",
+        bookmarkName: "wizards-1783670400000",
+        description: "Generated by Wizards",
+        mriQuery,
+      }),
+    ).resolves.toEqual({ cohortDefinitionId: 42 });
 
     expect(mockedPost).toHaveBeenCalledWith(
       "/d2e/analytics-svc/api/services/cohort",
@@ -123,6 +127,20 @@ describe("Wizard cohort API", () => {
     expect(decompress(submittedPayload.mriquery)).toEqual(mriQuery);
   });
 
+  it("rejects materialization when the backend omits the cohort definition id", async () => {
+    mockedPost.mockResolvedValue({ data: { message: "Cohort successfully materialized" } });
+    const bookmark = buildMriBookmark([], {}, { configId: "pa-config", configVersion: "7" }, "dataset-1").bookmark;
+
+    await expect(
+      materializeWizardBookmark({
+        datasetId: "dataset-1",
+        bookmarkId: "bookmark-1",
+        bookmarkName: "wizards-1783670400000",
+        mriQuery: buildMriMaterializationQuery(bookmark, "dataset-1"),
+      }),
+    ).rejects.toMatchObject({ operation: "materialize-cohort", code: "invalid-response" });
+  });
+
   it("rejects bookmark creation across datasets before making a request", async () => {
     const bookmark = buildMriBookmark([], {}, { configId: "pa-config", configVersion: "7" }, "dataset-2").bookmark;
 

plugins/ui/apps/wizards/src/api/wizardCohortApi.ts

@@ -45,6 +45,10 @@ export interface MaterializeWizardBookmarkInput {
   mriQuery: MriMaterializationQuery;
 }
 
+export interface MaterializeWizardBookmarkResult {
+  cohortDefinitionId: number;
+}
+
 const operationMessages: Record<WizardCohortApiOperation, string> = {
   "list-cohorts": "Unable to check previous Wizard analyses",
   "create-bookmark": "Unable to save the Wizard analysis",
@@ -126,13 +130,15 @@ export async function createWizardBookmark(input: CreateWizardBookmarkInput): Pr
   }
 }
 
-export async function materializeWizardBookmark(input: MaterializeWizardBookmarkInput): Promise<void> {
+export async function materializeWizardBookmark(
+  input: MaterializeWizardBookmarkInput,
+): Promise<MaterializeWizardBookmarkResult> {
   const operation: WizardCohortApiOperation = "materialize-cohort";
   requireValue(input.datasetId, "datasetId", operation);
   requireValue(input.bookmarkId, "bookmarkId", operation);
   requireValue(input.bookmarkName, "bookmarkName", operation);
   try {
-    await client.post(
+    const response = await client.post(
       "/d2e/analytics-svc/api/services/cohort",
       {
         datasetId: input.datasetId,
@@ -143,6 +149,12 @@ export async function materializeWizardBookmark(input: MaterializeWizardBookmark
       },
       { headers: { datasetid: input.datasetId } },
     );
+    const result = response.data as Partial<MaterializeWizardBookmarkResult> | null;
+    const cohortDefinitionId = Number(result?.cohortDefinitionId);
+    if (!Number.isInteger(cohortDefinitionId) || cohortDefinitionId <= 0) {
+      throw new WizardCohortApiError(operationMessages[operation], operation, "invalid-response");
+    }
+    return { cohortDefinitionId };
   } catch (error) {
     throw wrapApiError(error, operation);
   }

plugins/ui/apps/wizards/src/components/WizardDashboardModal.tsx

@@ -7,7 +7,6 @@ const stageMessages = {
   "awaiting-cache": "Checking your previous Wizard analyses…",
   "saving-bookmark": "Saving this Wizard analysis…",
   materializing: "Creating the cohort…",
-  "resolving-cohort": "Waiting for the cohort to become available…",
   "opening-dashboard": "Opening the dashboard…",
 } as const;
 

plugins/ui/apps/wizards/src/hooks/useWizardDashboardFlow.ts

@@ -31,7 +31,6 @@ const stageErrorMessages: Record<ActiveStage, string> = {
   "awaiting-cache": "We couldn't check your previous Wizard analyses. Please try again.",
   "saving-bookmark": "We couldn't save this Wizard analysis. Please try again.",
   materializing: "We couldn't generate the cohort. Please try again.",
-  "resolving-cohort": "The cohort is taking longer than expected. Please try again.",
   "opening-dashboard": "We couldn't open the dashboard. Please try again.",
 };
 
@@ -56,7 +55,6 @@ export function useWizardDashboardFlow({
   const abortRef = useRef<AbortController | null>(null);
   const lastInputRef = useRef<RunWizardDashboardFlowInput | null>(null);
   const pendingBookmarkRef = useRef<PendingWizardBookmark | null>(null);
-  const materializationSubmittedRef = useRef<string | null>(null);
   const activeStageRef = useRef<ActiveStage>("awaiting-cache");
   const loadInputRef = useRef<(() => Promise<OpenWizardDashboardInput>) | null>(null);
 
@@ -88,10 +86,6 @@ export function useWizardDashboardFlow({
           if (lastInputRef.current) lastInputRef.current.pendingBookmark = pendingBookmark;
           dispatch({ type: "bookmark-name", operationId, bookmarkName: pendingBookmark.bookmarkName });
         },
-        onMaterializationSubmitted: (bookmarkId) => {
-          materializationSubmittedRef.current = bookmarkId;
-          if (lastInputRef.current) lastInputRef.current.materializationSubmittedForBookmarkId = bookmarkId;
-        },
       })
         .then((result) => {
           dispatch({ type: "ready", operationId, result });
@@ -114,7 +108,6 @@ export function useWizardDashboardFlow({
     (loadInput: () => Promise<OpenWizardDashboardInput>) => {
       loadInputRef.current = loadInput;
       pendingBookmarkRef.current = null;
-      materializationSubmittedRef.current = null;
       const operationId = ++operationIdRef.current;
       abortRef.current?.abort();
       const controller = new AbortController();
@@ -166,8 +159,6 @@ export function useWizardDashboardFlow({
     execute({
       ...input,
       pendingBookmark: pendingBookmarkRef.current ?? input.pendingBookmark,
-      materializationSubmittedForBookmarkId:
-        materializationSubmittedRef.current ?? input.materializationSubmittedForBookmarkId,
     });
   }, [execute, openDashboard]);
 
@@ -180,7 +171,6 @@ export function useWizardDashboardFlow({
     abortRef.current?.abort();
     lastInputRef.current = null;
     pendingBookmarkRef.current = null;
-    materializationSubmittedRef.current = null;
     loadInputRef.current = null;
     dispatch({ type: "dataset-changed", datasetId });
   }, [datasetId]);

plugins/ui/apps/wizards/src/services/__tests__/wizardDashboardFlow.test.ts

@@ -34,10 +34,10 @@ describe("Wizard dashboard flow", () => {
     expect(materializeBookmark).not.toHaveBeenCalled();
   });
 
-  it("materializes an existing match and refreshes the bookmark list once", async () => {
+  it("uses the materialization response without waiting for the bookmark refresh", async () => {
     const createBookmark = vi.fn();
-    const materializeBookmark = vi.fn().mockResolvedValue(undefined);
-    const refreshCache = vi.fn().mockResolvedValue([bookmarkItem({ cohortDefinitionId: 42 })]);
+    const materializeBookmark = vi.fn().mockResolvedValue({ cohortDefinitionId: 42 });
+    const refreshCache = vi.fn().mockImplementation(() => new Promise<unknown>(() => undefined));
 
     const result = await runWizardDashboardFlow(baseInput, {
       ensureCache: vi.fn().mockResolvedValue([bookmarkItem()]),
@@ -52,13 +52,10 @@ describe("Wizard dashboard flow", () => {
     expect(result).toMatchObject({ cohortId: 42, cacheOutcome: "hit-unmaterialized" });
   });
 
-  it("uses the returned bookmark id, materializes, then refreshes once for the cohort id", async () => {
+  it("uses the returned bookmark and cohort ids for a new analysis", async () => {
     const createBookmark = vi.fn().mockResolvedValue({ status: "success", bmkId: "created-bookmark" });
-    const materializeBookmark = vi.fn().mockResolvedValue(undefined);
-    const refreshCache = vi
-      .fn()
-      .mockResolvedValueOnce([])
-      .mockResolvedValueOnce([bookmarkItem({ bmkId: "created-bookmark", cohortDefinitionId: 42 })]);
+    const materializeBookmark = vi.fn().mockResolvedValue({ cohortDefinitionId: 42 });
+    const refreshCache = vi.fn().mockResolvedValue([]);
     const onBookmarkCreated = vi.fn();
     const stages: string[] = [];
 
@@ -80,25 +77,19 @@ describe("Wizard dashboard flow", () => {
       bmkId: "created-bookmark",
       bookmarkName: "wizards-1783670400000",
     });
-    expect(stages).toEqual([
-      "awaiting-cache",
-      "saving-bookmark",
-      "materializing",
-      "resolving-cohort",
-      "opening-dashboard",
-    ]);
+    expect(stages).toEqual(["awaiting-cache", "saving-bookmark", "materializing", "opening-dashboard"]);
   });
 
   it("reuses a saved bookmark id on retry instead of saving again", async () => {
     const createBookmark = vi.fn();
-    const materializeBookmark = vi.fn().mockResolvedValue(undefined);
+    const materializeBookmark = vi.fn().mockResolvedValue({ cohortDefinitionId: 9 });
     const pendingBookmark = { bmkId: "created-bookmark", bookmarkName: "wizards-1783670400000" };
 
     const result = await runWizardDashboardFlow(
       { ...baseInput, pendingBookmark },
       {
         ensureCache: vi.fn().mockResolvedValue([]),
-        refreshCache: vi.fn().mockResolvedValue([bookmarkItem({ bmkId: "created-bookmark", cohortDefinitionId: 9 })]),
+        refreshCache: vi.fn().mockResolvedValue([]),
         createBookmark,
         materializeBookmark,
       },
@@ -124,33 +115,6 @@ describe("Wizard dashboard flow", () => {
     expect(materializeBookmark).not.toHaveBeenCalled();
   });
 
-  it("refreshes without materializing again after submission already completed", async () => {
-    const materializeBookmark = vi.fn();
-    const refreshCache = vi.fn().mockResolvedValue([bookmarkItem({ cohortDefinitionId: 42 })]);
-
-    await runWizardDashboardFlow(
-      { ...baseInput, materializationSubmittedForBookmarkId: "bookmark-1" },
-      {
-        ensureCache: vi.fn().mockResolvedValue([bookmarkItem()]),
-        refreshCache,
-        materializeBookmark,
-      },
-    );
-
-    expect(materializeBookmark).not.toHaveBeenCalled();
-    expect(refreshCache).toHaveBeenCalledTimes(1);
-  });
-
-  it("fails after the single refresh when no cohort id is returned", async () => {
-    await expect(
-      runWizardDashboardFlow(baseInput, {
-        ensureCache: vi.fn().mockResolvedValue([bookmarkItem()]),
-        refreshCache: vi.fn().mockResolvedValue([bookmarkItem()]),
-        materializeBookmark: vi.fn().mockResolvedValue(undefined),
-      }),
-    ).rejects.toThrow("materialized Wizard cohort was not returned");
-  });
-
   it("generates only the strict timestamp bookmark format", () => {
     expect(createWizardBookmarkName(1783670400000)).toBe("wizards-1783670400000");
     expect(() => createWizardBookmarkName(123)).toThrow();

plugins/ui/apps/wizards/src/services/wizardDashboardFlow.ts

@@ -4,6 +4,7 @@ import {
   type CreateWizardBookmarkResult,
   type CreateWizardBookmarkInput,
   type MaterializeWizardBookmarkInput,
+  type MaterializeWizardBookmarkResult,
 } from "../api/wizardCohortApi";
 import type { MriBookmark } from "../utils/mriQuery";
 import { buildMriMaterializationQuery } from "../utils/mriMaterializationQuery";
@@ -31,19 +32,17 @@ export interface RunWizardDashboardFlowInput {
   bookmark: MriBookmark;
   wizardConfig: Record<string, unknown>;
   pendingBookmark?: PendingWizardBookmark | null;
-  materializationSubmittedForBookmarkId?: string | null;
   signal?: AbortSignal;
 }
 
 export interface WizardDashboardFlowDependencies {
   ensureCache: () => Promise<unknown>;
   refreshCache: () => Promise<unknown>;
   createBookmark?: (input: CreateWizardBookmarkInput) => Promise<CreateWizardBookmarkResult>;
-  materializeBookmark?: (input: MaterializeWizardBookmarkInput) => Promise<void>;
+  materializeBookmark?: (input: MaterializeWizardBookmarkInput) => Promise<MaterializeWizardBookmarkResult>;
   now?: () => number;
   onStage?: (stage: FlowStage) => void;
   onBookmarkCreated?: (bookmark: PendingWizardBookmark) => void;
-  onMaterializationSubmitted?: (bookmarkId: string) => void;
 }
 
 export function createWizardBookmarkName(now = Date.now()): string {
@@ -120,25 +119,16 @@ export async function runWizardDashboardFlow(
 
   const mriQuery = buildMriMaterializationQuery(input.bookmark, input.datasetId);
   if (cohortDefinitionId === undefined) {
-    if (input.materializationSubmittedForBookmarkId !== bookmarkId) {
-      stage("materializing");
-      await materializeBookmark({
-        datasetId: input.datasetId,
-        bookmarkId,
-        bookmarkName,
-        mriQuery,
-      });
-      dependencies.onMaterializationSubmitted?.(bookmarkId);
-      throwIfAborted(input.signal);
-    }
-    stage("resolving-cohort");
-    const refreshedItems = await dependencies.refreshCache();
+    stage("materializing");
+    const materialized = await materializeBookmark({
+      datasetId: input.datasetId,
+      bookmarkId,
+      bookmarkName,
+      mriQuery,
+    });
     throwIfAborted(input.signal);
-    const refreshedBookmark = findWizardBookmarkById(refreshedItems, scope, bookmarkId);
-    if (refreshedBookmark?.cohortDefinitionId === undefined) {
-      throw new Error("The materialized Wizard cohort was not returned by the bookmark list");
-    }
-    cohortDefinitionId = refreshedBookmark.cohortDefinitionId;
+    cohortDefinitionId = materialized.cohortDefinitionId;
+    void dependencies.refreshCache().catch(() => undefined);
   }
 
   stage("opening-dashboard");

plugins/ui/apps/wizards/src/services/wizardDashboardState.ts

@@ -3,7 +3,6 @@ export type WizardDashboardStatus =
   | "awaiting-cache"
   | "saving-bookmark"
   | "materializing"
-  | "resolving-cohort"
   | "opening-dashboard"
   | "ready"
   | "error";