Merge pull request #3465 from quantified-uncertainty/workflow-improve…

…ments

Workflow improvements

apps/hub/src/ai/data/loadWorkflows.ts

@@ -1,29 +1,72 @@
-import "server-only";
+import { Prisma } from "@prisma/client";
+
+import { ClientWorkflow } from "@quri/squiggle-ai";
 
 import { prisma } from "@/lib/server/prisma";
-import { getSessionUserOrRedirect } from "@/users/auth";
+import { Paginated } from "@/lib/types";
+import { checkRootUser, getSessionUserOrRedirect } from "@/users/auth";
 
 import { decodeDbWorkflowToClientWorkflow } from "./storage";
 
-export async function loadWorkflows({
-  limit = 20,
-}: {
-  limit?: number;
-} = {}) {
+export type AiWorkflow = {
+  workflow: ClientWorkflow;
+  author: {
+    username: string;
+  };
+};
+
+export async function loadWorkflows(
+  params: {
+    allUsers?: boolean;
+    cursor?: string;
+    limit?: number;
+  } = {}
+): Promise<Paginated<AiWorkflow>> {
   const sessionUser = await getSessionUserOrRedirect();
 
+  const limit = params.limit ?? 20;
+
+  const where: Prisma.AiWorkflowWhereInput = {};
+  if (params.allUsers) {
+    console.log("loading all workflows");
+    await checkRootUser();
+  } else {
+    where.user = { email: sessionUser.email };
+  }
+
   const rows = await prisma.aiWorkflow.findMany({
     orderBy: { createdAt: "desc" },
-    where: {
-      user: { email: sessionUser.email },
+    cursor: params.cursor ? { id: params.cursor } : undefined,
+    where,
+    include: {
+      user: {
+        select: {
+          asOwner: {
+            select: {
+              slug: true,
+            },
+          },
+        },
+      },
     },
     take: limit + 1,
   });
 
-  const workflows = rows.map((row) => decodeDbWorkflowToClientWorkflow(row));
+  // TODO - it would be good to preserve author information in the client, but this would require a new type (ClientWorkflowWithAuthor?)
+  const workflows = rows.map((row) => ({
+    workflow: decodeDbWorkflowToClientWorkflow(row),
+    author: { username: row.user.asOwner?.slug ?? "[unknown]" },
+  }));
+
+  const nextCursor = workflows[workflows.length - 1]?.workflow.id;
+
+  async function loadMore(limit: number) {
+    "use server";
+    return loadWorkflows({ ...params, cursor: nextCursor, limit });
+  }
 
   return {
-    workflows: limit ? workflows.slice(0, limit) : workflows,
-    hasMore: limit ? workflows.length > limit : false,
+    items: workflows.slice(0, limit),
+    loadMore: workflows.length > limit ? loadMore : undefined,
   };
 }

apps/hub/src/ai/data/v1_0.ts

@@ -108,10 +108,20 @@ export function decodeV1_0JsonToClientWorkflow(
       // modern steps in ClientWorkflow store state as an object
       state:
         step.state === "DONE"
-          ? ({ kind: "DONE", outputs } as const)
+          ? ({
+              kind: "DONE",
+              outputs,
+              durationMs: 0, // old workflow steps don't have durationMs
+            } as const)
           : step.state === "FAILED"
-            ? { kind: "FAILED", errorType: "CRITICAL", message: "Unknown" }
+            ? {
+                kind: "FAILED",
+                errorType: "CRITICAL",
+                message: "Unknown",
+                durationMs: 0, // old workflow steps don't have durationMs
+              }
             : { kind: "PENDING" },
+      startTime: v1Workflow.timestamp, // old workflow steps don't have start times
     })),
     inputs:
       input.type === "Create"

apps/hub/src/app/admin/layout.tsx

@@ -3,19 +3,11 @@ import { PropsWithChildren } from "react";
 import { LockIcon } from "@quri/ui";
 
 import { FullLayoutWithPadding } from "@/components/layout/FullLayoutWithPadding";
-import { NarrowPageLayout } from "@/components/layout/NarrowPageLayout";
 import { H1 } from "@/components/ui/Headers";
-import { auth } from "@/lib/server/auth";
-import { isRootEmail } from "@/users/auth";
+import { checkRootUser } from "@/users/auth";
 
 export default async function AdminLayout({ children }: PropsWithChildren) {
-  const session = await auth();
-
-  const email = session?.user.email;
-
-  if (!email || !isRootEmail(email)) {
-    return <NarrowPageLayout>Access denied.</NarrowPageLayout>;
-  }
+  await checkRootUser();
 
   return (
     <div>

apps/hub/src/app/ai/AiDashboard.tsx

@@ -2,23 +2,24 @@
 
 import { FC, useRef } from "react";
 
-import { ClientWorkflow } from "@quri/squiggle-ai";
+import { AiWorkflow } from "@/ai/data/loadWorkflows";
+import { usePaginator } from "@/lib/hooks/usePaginator";
+import { Paginated } from "@/lib/types";
 
 import { Sidebar } from "./Sidebar";
 import { useSquiggleWorkflows } from "./useSquiggleWorkflows";
 import { WorkflowViewer } from "./WorkflowViewer";
 
 type Props = {
-  initialWorkflows: ClientWorkflow[];
-  hasMoreWorkflows: boolean;
+  initialWorkflows: Paginated<AiWorkflow>;
 };
 
-export const AiDashboard: FC<Props> = ({
-  initialWorkflows,
-  hasMoreWorkflows,
-}: Props) => {
+export const AiDashboard: FC<Props> = ({ initialWorkflows }: Props) => {
+  const { items: unpaginatedWorkflows, loadNext } =
+    usePaginator(initialWorkflows);
+
   const { workflows, submitWorkflow, selectedWorkflow, selectWorkflow } =
-    useSquiggleWorkflows(initialWorkflows);
+    useSquiggleWorkflows(unpaginatedWorkflows);
 
   const sidebarRef = useRef<{ edit: (code: string) => void }>(null);
 
@@ -31,16 +32,16 @@ export const AiDashboard: FC<Props> = ({
           selectWorkflow={selectWorkflow}
           selectedWorkflow={selectedWorkflow}
           workflows={workflows}
-          hasMoreWorkflows={hasMoreWorkflows}
+          loadNext={loadNext}
           ref={sidebarRef}
         />
       </div>
       {/* Right column: Menu and SquigglePlayground */}
       {selectedWorkflow && (
         <div className="min-w-0 flex-1 overflow-x-auto">
           <WorkflowViewer
-            key={selectedWorkflow.id}
-            workflow={selectedWorkflow}
+            key={selectedWorkflow.workflow.id}
+            workflow={selectedWorkflow.workflow}
           />
         </div>
       )}

apps/hub/src/app/ai/Sidebar.tsx

@@ -10,7 +10,7 @@ import {
 } from "react";
 import { FormProvider, useForm } from "react-hook-form";
 
-import { ClientWorkflow, LlmId, MODEL_CONFIGS } from "@quri/squiggle-ai";
+import { LlmId, MODEL_CONFIGS } from "@quri/squiggle-ai";
 import {
   Button,
   NumberFormField,
@@ -20,7 +20,7 @@ import {
   TextFormField,
 } from "@quri/ui";
 
-import { LoadMoreViaSearchParam } from "@/components/LoadMoreViaSearchParam";
+import { AiWorkflow } from "@/ai/data/loadWorkflows";
 
 import { AiRequestBody } from "./utils";
 import { WorkflowSummaryList } from "./WorkflowSummaryList";
@@ -32,9 +32,9 @@ type Handle = {
 type Props = {
   submitWorkflow: (requestBody: AiRequestBody) => void;
   selectWorkflow: (id: string) => void;
-  selectedWorkflow: ClientWorkflow | undefined;
-  workflows: ClientWorkflow[];
-  hasMoreWorkflows: boolean;
+  selectedWorkflow: AiWorkflow | undefined;
+  workflows: AiWorkflow[];
+  loadNext?: (count: number) => void;
 };
 
 type FormShape = {
@@ -47,13 +47,7 @@ type FormShape = {
 };
 
 export const Sidebar = forwardRef<Handle, Props>(function Sidebar(
-  {
-    submitWorkflow,
-    selectWorkflow,
-    selectedWorkflow,
-    workflows,
-    hasMoreWorkflows,
-  },
+  { submitWorkflow, selectWorkflow, selectedWorkflow, workflows, loadNext },
   ref
 ) {
   const form = useForm<FormShape>({
@@ -82,7 +76,7 @@ Outputs:
 
   useEffect(() => {
     if (workflows.length > prevWorkflowsLengthRef.current) {
-      selectWorkflow(workflows[0].id);
+      selectWorkflow(workflows[0].workflow.id);
       prevWorkflowsLengthRef.current = workflows.length;
     }
   }, [workflows, selectWorkflow]);
@@ -228,15 +222,12 @@ Outputs:
         <Button wide onClick={handleSubmit} disabled={isSubmitDisabled}>
           Start Workflow
         </Button>
-        <div className="flex-grow overflow-y-auto">
-          <h2 className="mb-2 text-sm font-bold">Workflows</h2>
-          <WorkflowSummaryList
-            workflows={workflows}
-            selectedWorkflow={selectedWorkflow}
-            selectWorkflow={selectWorkflow}
-          />
-          {hasMoreWorkflows && <LoadMoreViaSearchParam />}
-        </div>
+        <WorkflowSummaryList
+          workflows={workflows}
+          loadNext={loadNext}
+          selectedWorkflow={selectedWorkflow}
+          selectWorkflow={selectWorkflow}
+        />
       </div>
     </FormProvider>
   );

apps/hub/src/app/ai/SquigglePlaygroundForWorkflow.tsx

@@ -2,6 +2,7 @@ import React, { useEffect, useState } from "react";
 
 import { llmLinker } from "@quri/squiggle-ai";
 import {
+  defaultSquiggleVersion,
   SquigglePlaygroundVersionPicker,
   type SquiggleVersion,
   versionedSquigglePackages,
@@ -20,7 +21,9 @@ export function SquigglePlaygroundForWorkflow({
   const [squiggle, setSquiggle] = useState<
     undefined | Awaited<ReturnType<typeof versionedSquigglePackages>>
   >();
-  const [version, setVersion] = useState<SquiggleVersion>("dev"); // Later versions are often buggy
+  const [version, setVersion] = useState<SquiggleVersion>(
+    defaultSquiggleVersion
+  );
 
   const onVersionChange = (version: SquiggleVersion) => {
     setVersion(version);
@@ -42,7 +45,7 @@ export function SquigglePlaygroundForWorkflow({
     <squiggle.components.SquigglePlayground
       height={height}
       defaultCode={defaultCode}
-      linker={llmLinker as any}
+      linker={llmLinker}
       renderExtraControls={() => (
         <div className="flex h-full items-center justify-end gap-2">
           <SquigglePlaygroundVersionPicker

apps/hub/src/app/ai/StepStatusIcon.tsx

@@ -4,6 +4,9 @@ import { ClientStep } from "@quri/squiggle-ai";
 import { CheckCircleIcon, ErrorIcon, RefreshIcon, TextTooltip } from "@quri/ui";
 
 export const StepStatusIcon: FC<{ step: ClientStep }> = ({ step }) => {
+  const ageInSeconds = (new Date().getTime() - step.startTime) / 1000;
+  const maxLoadingAge = 300;
+
   switch (step.state.kind) {
     case "PENDING":
       return <RefreshIcon className="animate-spin text-gray-400" size={16} />;