feat: Merge dataset label application with new label creation

app/src/components/combobox/ComboBox.tsx

@@ -33,13 +33,9 @@ export interface ComboBoxProps<T extends object>
   description?: string | null;
   errorMessage?: string | ((validation: ValidationResult) => string);
   children: React.ReactNode | ((item: T) => React.ReactNode);
-  /**
-   * The container to render the popover into. If not provided, the popover will be rendered in the body.
-   * @default document.body
-   */
-  container?: HTMLElement;
   /**
    * The width of the combobox. If not provided, the combobox will be sized to fit its contents. with a minimum width of 200px.
+   * @default "100%"
    */
   width?: string;
   /**
@@ -65,7 +61,6 @@ export function ComboBox<T extends object>({
   description,
   errorMessage,
   children,
-  container,
   size = "M",
   width,
   stopPropagation,
@@ -101,7 +96,7 @@ export function ComboBox<T extends object>({
         <Text slot="description">{description}</Text>
       ) : null}
       <FieldError>{errorMessage}</FieldError>
-      <Popover css={comboBoxPopoverCSS} UNSTABLE_portalContainer={container}>
+      <Popover css={comboBoxPopoverCSS}>
         <ListBox renderEmptyState={renderEmptyState}>{children}</ListBox>
       </Popover>
     </AriaComboBox>

app/src/components/dataset/DatasetLabelConfigButton.tsx

@@ -1,5 +1,4 @@
-import { Suspense, useMemo, useState } from "react";
-import { ModalOverlay } from "react-aria-components";
+import { Suspense, useCallback, useMemo, useState } from "react";
 import {
   ConnectionHandler,
   graphql,
@@ -10,6 +9,7 @@ import {
 import { css } from "@emotion/react";
 
 import {
+  Alert,
   Button,
   type ButtonProps,
   ColorSwatch,
@@ -24,13 +24,16 @@ import {
   ListBox,
   ListBoxItem,
   Loading,
-  Modal,
   Popover,
   PopoverArrow,
   type Selection,
   View,
 } from "@phoenix/components";
-import { NewDatasetLabelDialog } from "@phoenix/components/dataset/NewDatasetLabelDialog";
+import {
+  useDatasetLabelMutations,
+  UseDatasetLabelMutationsParams,
+} from "@phoenix/components/dataset/useDatasetLabelMutations";
+import { NewLabelForm } from "@phoenix/components/label";
 import { useNotifyError } from "@phoenix/contexts";
 import { isStringArray } from "@phoenix/typeUtils";
 import { getErrorMessagesFromRelayMutationError } from "@phoenix/utils/errorUtils";
@@ -47,68 +50,37 @@ type DatasetLabelConfigButtonProps = {
 
 export function DatasetLabelConfigButton(props: DatasetLabelConfigButtonProps) {
   const { datasetId, variant = "default" } = props;
-  const [showNewLabelDialog, setShowNewLabelDialog] = useState<boolean>(false);
-  const [isOpen, setIsOpen] = useState(false);
 
   return (
-    <>
-      <DialogTrigger isOpen={isOpen} onOpenChange={setIsOpen}>
-        <Button
-          variant={variant}
-          size="S"
-          leadingVisual={<Icon svg={<Icons.PriceTagsOutline />} />}
-          aria-label="Configure dataset labels"
-        >
-          Label
-        </Button>
-        <Popover
-          placement="bottom start"
-          shouldCloseOnInteractOutside={() => true}
-          css={css`
-            min-width: 400px;
-            max-width: 500px;
-          `}
-        >
-          <PopoverArrow />
-          <Dialog>
-            <Suspense fallback={<Loading />}>
-              <DatasetLabelSelectionDialogContent
-                datasetId={datasetId}
-                onNewLabelPress={() => {
-                  setShowNewLabelDialog(true);
-                }}
-              />
-            </Suspense>
-          </Dialog>
-        </Popover>
-      </DialogTrigger>
-      <ModalOverlay
-        isOpen={showNewLabelDialog}
-        onOpenChange={setShowNewLabelDialog}
+    <DialogTrigger>
+      <Button
+        variant={variant}
+        size="S"
+        leadingVisual={<Icon svg={<Icons.PriceTagsOutline />} />}
+        aria-label="Configure dataset labels"
       >
-        <Modal size="S">
-          <NewDatasetLabelDialog
-            updateConnectionIds={[
-              ConnectionHandler.getConnectionID(
-                "client:root",
-                "DatasetLabelConfigButtonAllLabels_datasetLabels"
-              ),
-            ]}
-            datasetId={datasetId}
-            onCompleted={() => {
-              setShowNewLabelDialog(false);
-            }}
-          />
-        </Modal>
-      </ModalOverlay>
-    </>
+        Label
+      </Button>
+      <Popover
+        placement="bottom start"
+        shouldCloseOnInteractOutside={() => true}
+        css={css`
+          min-width: 400px;
+          max-width: 500px;
+        `}
+      >
+        <PopoverArrow />
+        <Dialog>
+          <Suspense fallback={<Loading />}>
+            <DatasetLabelSelectionContent datasetId={datasetId} />
+          </Suspense>
+        </Dialog>
+      </Popover>
+    </DialogTrigger>
   );
 }
 
-function DatasetLabelSelectionDialogContent(props: {
-  datasetId: string;
-  onNewLabelPress: () => void;
-}) {
+export function DatasetLabelSelectionContent(props: { datasetId: string }) {
   const { datasetId } = props;
   const query = useLazyLoadQuery<DatasetLabelConfigButtonQuery>(
     graphql`
@@ -127,53 +99,14 @@ function DatasetLabelSelectionDialogContent(props: {
   return <DatasetLabelList query={query} dataset={query.dataset} {...props} />;
 }
 
-/**
- * Exported label selection content with integrated "Create New Label" functionality
- * Styled to match PromptLabelConfigButton
- */
-export function DatasetLabelSelectionContent(props: { datasetId: string }) {
-  const { datasetId } = props;
-  const [showNewLabelDialog, setShowNewLabelDialog] = useState<boolean>(false);
-
-  return (
-    <>
-      <DatasetLabelSelectionDialogContent
-        {...props}
-        onNewLabelPress={() => setShowNewLabelDialog(true)}
-      />
-      <ModalOverlay
-        isOpen={showNewLabelDialog}
-        onOpenChange={setShowNewLabelDialog}
-      >
-        <Modal size="S">
-          <NewDatasetLabelDialog
-            updateConnectionIds={[
-              ConnectionHandler.getConnectionID(
-                "client:root",
-                "DatasetLabelConfigButtonAllLabels_datasetLabels"
-              ),
-            ]}
-            datasetId={datasetId}
-            onCompleted={() => {
-              // Only close the create modal, keep the popover open
-              setShowNewLabelDialog(false);
-            }}
-          />
-        </Modal>
-      </ModalOverlay>
-    </>
-  );
-}
-
 function DatasetLabelList({
   query,
   dataset,
-  onNewLabelPress,
 }: {
   dataset: DatasetLabelConfigButton_datasetLabels$key;
   query: DatasetLabelConfigButton_allLabels$key;
-  onNewLabelPress: () => void;
 }) {
+  const [mode, setMode] = useState<"apply" | "create">("apply");
   const notifyError = useNotifyError();
   const datasetData = useFragment<DatasetLabelConfigButton_datasetLabels$key>(
     graphql`
@@ -270,46 +203,83 @@ function DatasetLabelList({
         borderBottomWidth="thin"
         borderColor="dark"
         minWidth={300}
+        maxWidth={300}
       >
         <Flex direction="column" gap="size-50">
           <Flex
             direction="row"
             justifyContent="space-between"
             alignItems="center"
           >
-            <Heading level={4} weight="heavy">
-              Assign labels to this dataset
-            </Heading>
-            <Button variant="quiet" size="S" onPress={onNewLabelPress}>
-              <Icon svg={<Icons.PlusOutline />} />
-            </Button>
+            <Flex direction="row" gap="size-100" alignItems="center">
+              {mode === "create" && (
+                <Button
+                  variant="quiet"
+                  size="S"
+                  leadingVisual={<Icon svg={<Icons.ChevronLeft />} />}
+                  onPress={() => setMode("apply")}
+                />
+              )}
+              <Heading level={4} weight="heavy">
+                {mode === "create"
+                  ? "Create New Label for this dataset"
+                  : "Assign labels to this dataset"}
+              </Heading>
+            </Flex>
+            {mode === "apply" && (
+              <Button
+                variant="quiet"
+                size="S"
+                leadingVisual={<Icon svg={<Icons.PlusOutline />} />}
+                onPress={() => setMode("create")}
+              />
+            )}
           </Flex>
-          <DebouncedSearch
-            autoFocus
-            aria-label="Search labels"
-            placeholder="Search labels..."
-            onChange={setSearch}
-          />
+          {mode === "apply" && (
+            <DebouncedSearch
+              autoFocus
+              aria-label="Search labels"
+              placeholder="Search labels..."
+              onChange={setSearch}
+            />
+          )}
         </Flex>
       </View>
-      <ListBox
-        aria-label="labels"
-        items={labels}
-        selectionMode="multiple"
-        selectedKeys={selected}
-        onSelectionChange={onSelectionChange}
-        css={css`
-          height: 300px;
-        `}
-        renderEmptyState={() => "No labels found"}
-      >
-        {(item) => <DatasetLabelListBoxItem key={item.id} item={item} />}
-      </ListBox>
-      <View padding="size-100" borderTopColor="dark" borderTopWidth="thin">
-        <LinkButton variant="quiet" size="S" to="/settings/datasets">
-          Edit Labels
-        </LinkButton>
-      </View>
+      {mode === "apply" && (
+        <>
+          <ListBox
+            aria-label="labels"
+            items={labels}
+            selectionMode="multiple"
+            selectedKeys={selected}
+            onSelectionChange={onSelectionChange}
+            css={css`
+              min-height: 300px;
+              max-height: 300px;
+            `}
+            renderEmptyState={() => "No labels found"}
+          >
+            {(item) => <DatasetLabelListBoxItem key={item.id} item={item} />}
+          </ListBox>
+          <View padding="size-100" borderTopColor="dark" borderTopWidth="thin">
+            <LinkButton variant="quiet" size="S" to="/settings/datasets">
+              Edit Labels
+            </LinkButton>
+          </View>
+        </>
+      )}
+      {mode === "create" && (
+        <CreateNewDatasetLabel
+          onCompleted={() => setMode("apply")}
+          updateConnectionIds={[
+            ConnectionHandler.getConnectionID(
+              "client:root",
+              "DatasetLabelConfigButtonAllLabels_datasetLabels"
+            ),
+          ]}
+          datasetId={datasetData.id}
+        />
+      )}
     </>
   );
 }
@@ -333,3 +303,36 @@ function DatasetLabelListBoxItem({
     </ListBoxItem>
   );
 }
+
+type CreateNewDatasetLabelProps = UseDatasetLabelMutationsParams & {
+  onCompleted: () => void;
+};
+
+function CreateNewDatasetLabel({
+  updateConnectionIds,
+  datasetId,
+  onCompleted,
+}: CreateNewDatasetLabelProps) {
+  const { addLabelMutation, isSubmitting, error } = useDatasetLabelMutations({
+    updateConnectionIds,
+    datasetId,
+  });
+
+  const onSubmit = useCallback(
+    (label: Parameters<typeof addLabelMutation>[0]) => {
+      addLabelMutation(label, onCompleted);
+    },
+    [addLabelMutation, onCompleted]
+  );
+
+  return (
+    <>
+      {!!error && (
+        <Alert banner variant="danger">
+          {error}
+        </Alert>
+      )}
+      <NewLabelForm onSubmit={onSubmit} isSubmitting={isSubmitting} />
+    </>
+  );
+}