AI workflow error handling

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

@@ -31,10 +31,18 @@ export function decodeDbWorkflowToClientWorkflow(
     return {
       id: row.id,
       timestamp: row.createdAt.getTime(),
-      status: "error",
+      status: "finished",
       inputs: {},
       steps: [],
-      result: `Invalid workflow format in the database: ${e}`,
+      result: {
+        code: "",
+        isValid: false,
+        totalPrice: 0,
+        runTimeMs: 0,
+        llmRunCount: 0,
+        logSummary: "",
+        error: `Invalid workflow format in the database: ${e}`,
+      },
     } satisfies ClientWorkflow;
   }
 }

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

@@ -94,16 +94,15 @@ export const v1WorkflowSchema = z.discriminatedUnion("status", [
   }),
 ]);
 
+// upgrading legacy workflow to new client workflow shape
 export function decodeV1_0JsonToClientWorkflow(
   json: Prisma.JsonValue
 ): ClientWorkflow {
   // `input` doesn't exist in the new ClientWorkflow shape, so we need to pull it out
   const { input, ...v1Workflow } = v1WorkflowSchema.parse(json);
 
-  // upgrading legacy workflow to new client workflow shape
-  return {
-    ...v1Workflow,
-    steps: v1Workflow.steps.map(({ outputs, ...step }) => ({
+  const steps: ClientWorkflow["steps"] = v1Workflow.steps.map(
+    ({ outputs, ...step }) => ({
       ...step,
       // modern steps in ClientWorkflow store state as an object
       state:
@@ -122,22 +121,47 @@ export function decodeV1_0JsonToClientWorkflow(
               }
             : { kind: "PENDING" },
       startTime: v1Workflow.timestamp, // old workflow steps don't have start times
-    })),
-    inputs:
-      input.type === "Create"
-        ? {
-            prompt: {
-              value: input.prompt,
-              kind: "prompt",
-              id: `${v1Workflow.id}-prompt`,
-            },
-          }
-        : {
-            source: {
-              value: input.source,
-              kind: "source",
-              id: `${v1Workflow.id}-source`,
-            },
+    })
+  );
+
+  const inputs: ClientWorkflow["inputs"] =
+    input.type === "Create"
+      ? {
+          prompt: {
+            value: input.prompt,
+            kind: "prompt",
+            id: `${v1Workflow.id}-prompt`,
+          },
+        }
+      : {
+          source: {
+            value: input.source,
+            kind: "source",
+            id: `${v1Workflow.id}-source`,
           },
+        };
+
+  if (v1Workflow.status === "error") {
+    return {
+      ...v1Workflow,
+      status: "finished",
+      result: {
+        code: "",
+        isValid: false,
+        totalPrice: 0,
+        runTimeMs: 0,
+        llmRunCount: 0,
+        logSummary: "",
+        error: v1Workflow.result,
+      },
+      steps,
+      inputs,
+    };
+  }
+
+  return {
+    ...v1Workflow,
+    steps,
+    inputs,
   };
 }

apps/hub/src/app/ai/WorkflowSummaryList/WorkflowStatusIcon.tsx

@@ -12,7 +12,7 @@ export function isWorkflowOutdated(workflow: ClientWorkflow): boolean {
 
 function getWorkflowStatusForIcon(
   workflow: ClientWorkflow
-): ClientWorkflow["status"] {
+): "loading" | "finished" | "error" {
   if (workflow.status === "loading") {
     return isWorkflowOutdated(workflow) ? "error" : "loading";
   }
@@ -36,5 +36,7 @@ export const WorkflowStatusIcon: FC<{ workflow: ClientWorkflow }> = ({
       return <CheckCircleIcon className="text-emerald-400" size={16} />;
     case "error":
       return <ErrorIcon className="text-red-400" size={16} />;
+    default:
+      throw new Error(status satisfies never);
   }
 };

apps/hub/src/app/ai/WorkflowViewer/index.tsx

@@ -3,7 +3,7 @@ import { format } from "date-fns";
 import { Children, FC } from "react";
 
 import { ClientWorkflow } from "@quri/squiggle-ai";
-import { ErrorIcon, StyledTab } from "@quri/ui";
+import { ErrorIcon, StyledTab, TextTooltip } from "@quri/ui";
 
 import { commonDateFormat } from "@/lib/constants";
 import { useAvailableHeight } from "@/lib/hooks/useAvailableHeight";
@@ -53,12 +53,26 @@ const FinishedWorkflowViewer: FC<WorkflowViewerProps<"finished">> = ({
         <Header
           workflow={workflow}
           renderLeft={() => (
-            <LineSeparatedList>
-              <span>{(workflow.result.runTimeMs / 1000).toFixed(2)}s</span>
-              <span>${workflow.result.totalPrice.toFixed(2)}</span>
-              <span>{workflow.result.llmRunCount} LLM runs</span>
-              <WorkflowDate workflow={workflow} />
-            </LineSeparatedList>
+            <div className="flex items-center gap-2">
+              <LineSeparatedList>
+                <span>{(workflow.result.runTimeMs / 1000).toFixed(2)}s</span>
+                <span>${workflow.result.totalPrice.toFixed(2)}</span>
+                <span>{workflow.result.llmRunCount} LLM runs</span>
+                <WorkflowDate workflow={workflow} />
+              </LineSeparatedList>
+              {workflow.result.error ? (
+                <div className="flex items-center gap-1">
+                  <TextTooltip text={workflow.result.error}>
+                    <span>
+                      <ErrorIcon
+                        className="cursor-pointer text-red-400"
+                        size={16}
+                      />
+                    </span>
+                  </TextTooltip>
+                </div>
+              ) : null}
+            </div>
           )}
           renderRight={() => (
             <div className="flex items-center gap-2">
@@ -137,18 +151,6 @@ export const WorkflowViewer: FC<WorkflowViewerProps> = ({
       return <FinishedWorkflowViewer {...props} workflow={workflow} />;
     case "loading":
       return <LoadingWorkflowViewer {...props} workflow={workflow} />;
-    case "error":
-      return (
-        <div className="mt-2 rounded-md border border-red-300 bg-red-50 p-4">
-          <h3 className="mb-2 text-lg font-semibold text-red-800">
-            Server Error
-          </h3>
-          <p className="mb-4 text-red-700">{workflow.result}</p>
-          <p className="text-sm text-red-600">
-            Please try refreshing the page or attempt your action again.
-          </p>
-        </div>
-      );
     default:
       throw workflow satisfies never;
   }

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

@@ -63,7 +63,7 @@ export function useSquiggleWorkflows(preloadedWorkflows: AiWorkflow[]) {
           steps: [],
         },
         author: {
-          username: session.data?.user?.name ?? "Unknown",
+          username: session.data?.user?.username ?? "Unknown",
         },
       };
       setWorkflows((workflows) => [workflow, ...workflows]);
@@ -109,8 +109,16 @@ export function useSquiggleWorkflows(preloadedWorkflows: AiWorkflow[]) {
       } catch (error) {
         updateWorkflow(id, (workflow) => ({
           ...workflow,
-          status: "error",
-          result: `Server error: ${error instanceof Error ? error.toString() : "Unknown error"}.`,
+          status: "finished",
+          result: {
+            code: "",
+            isValid: false,
+            totalPrice: 0,
+            runTimeMs: 0,
+            llmRunCount: 0,
+            logSummary: "",
+            error: `Server error: ${error instanceof Error ? error.toString() : "Unknown error"}.`,
+          },
         }));
       }
     },

internal-packages/ai/src/types.ts

@@ -93,6 +93,7 @@ const stepUpdatedSchema = stepSchema.partial().required({
 export const clientWorkflowResultSchema = z.object({
   code: z.string().describe("Squiggle code snippet"),
   isValid: z.boolean(),
+  error: z.string().optional(),
   totalPrice: z.number(),
   runTimeMs: z.number(),
   llmRunCount: z.number(),
@@ -142,11 +143,6 @@ export const clientWorkflowSchema = z.discriminatedUnion("status", [
     status: z.literal("finished"),
     result: clientWorkflowResultSchema,
   }),
-  z.object({
-    ...commonClientWorkflowFields,
-    status: z.literal("error"),
-    result: z.string(),
-  }),
 ]);
 
 export type ClientWorkflow = z.infer<typeof clientWorkflowSchema>;

internal-packages/ai/src/workflows/Workflow.ts

@@ -111,10 +111,13 @@ export class Workflow<Shape extends IOShape = IOShape> {
 
   public llmClient: LLMClient;
 
+  public error?: string;
+
   constructor(params: WorkflowInstanceParams<Shape>) {
     this.id = params.id ?? crypto.randomUUID();
     this.template = params.template;
     this.inputs = params.inputs;
+    this.error = params.error ?? undefined;
 
     this.llmConfig = params.llmConfig ?? llmConfigDefault;
     this.startTime = params.steps.at(0)?.startTime ?? Date.now();
@@ -201,9 +204,9 @@ export class Workflow<Shape extends IOShape = IOShape> {
       payload: { step },
     });
 
-    // apply the transition rule, produce the next step in PENDING state
-    // this code is inlined in this method, which guarantees that we always have one pending step
-    // (until the transition rule decides to finish)
+    // Apply the transition rule, produce the next step in PENDING state or stop the workflow.
+    // This code is inlined in this method, which guarantees that we always have one pending step
+    // (until the transition rule decides to finish).
     const result = this.transitionRule(step, new WorkflowGuardHelpers(step));
     switch (result.kind) {
       case "repeat":
@@ -213,10 +216,11 @@ export class Workflow<Shape extends IOShape = IOShape> {
         this.addStep(result.step.prepare(result.inputs));
         break;
       case "finish":
-        // no new steps to add
+        // no new steps to add, we're done
         break;
       case "fatal":
-        throw new Error(result.message);
+        this.error = result.message;
+        break;
     }
   }
 
@@ -299,7 +303,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
 
   getFinalResult(): ClientWorkflowResult {
     const finalStep = this.getRecentStepWithCode();
-    const isValid = finalStep?.step.getState().kind === "DONE";
+    const isValid = finalStep?.step.getState().kind === "DONE" && !this.error;
 
     // compute run time
     let runTimeMs: number;
@@ -326,6 +330,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
     return {
       code: finalStep?.code ?? "",
       isValid,
+      error: this.error,
       totalPrice,
       runTimeMs,
       llmRunCount,
@@ -427,6 +432,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
     return {
       id: this.id,
       templateName: this.template.name,
+      error: this.error ?? null,
       inputIds: Object.fromEntries(
         Object.entries(this.inputs).map(([key, input]) => [
           key,
@@ -454,6 +460,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
   }): Workflow<IOShape> {
     const workflow = new Workflow({
       id: node.id,
+      error: node.error ?? undefined,
       template: getWorkflowTemplateByName(node.templateName),
       inputs: Object.fromEntries(
         Object.entries(node.inputIds).map(([key, id]) => [
@@ -509,4 +516,5 @@ export type SerializedWorkflow = {
   inputIds: Record<string, number>;
   llmConfig: LlmConfig;
   stepIds: number[];
+  error: string | null;
 };

internal-packages/ai/src/workflows/WorkflowTemplate.ts

@@ -11,6 +11,7 @@ export type WorkflowInstanceParams<Shape extends IOShape> = {
   llmConfig?: LlmConfig;
   openaiApiKey?: string;
   anthropicApiKey?: string;
+  error?: string;
 };
 
 /**

internal-packages/ai/src/workflows/controllers.ts

@@ -145,6 +145,6 @@ export function getDefaultTransitionRule<Shape extends IOShape>(
       }
     }
 
-    return h.fatal("Unknown step");
+    return h.fatal(ERROR_MESSAGES.UNKNOWN_STEP);
   };
 }