allow fluent syntax for defining workflows

Author: ianmacartney

example/convex/nestedWorkflow.ts

@@ -3,9 +3,12 @@ import { workflow } from "./example";
 import { internal } from "./_generated/api";
 import { internalMutation } from "./_generated/server";
 
-export const parentWorkflow = workflow.define({
-  args: { prompt: v.string() },
-  handler: async (ctx, args) => {
+export const parentWorkflow = workflow
+  .define({
+    args: { prompt: v.string() },
+    returns: v.number(),
+  })
+  .handler(async (ctx, args) => {
     console.log("Starting confirmation workflow");
     const length = await ctx.runWorkflow(
       internal.nestedWorkflow.childWorkflow,
@@ -16,12 +19,11 @@ export const parentWorkflow = workflow.define({
       foo: args.prompt,
     });
     console.log("Step result:", stepResult);
-  },
-});
+    return stepResult;
+  });
 
 export const childWorkflow = workflow.define({
   args: { foo: v.string() },
-  returns: v.number(),
   handler: async (_ctx, args) => {
     console.log("Starting nested workflow");
     return args.foo.length;

src/client/index.ts

@@ -15,12 +15,7 @@ import {
   type RegisteredMutation,
   type ReturnValueForOptionalValidator,
 } from "convex/server";
-import type {
-  Infer,
-  ObjectType,
-  PropertyValidators,
-  Validator,
-} from "convex/values";
+import type { ObjectType, PropertyValidators, Validator } from "convex/values";
 import type { Step } from "../component/schema.js";
 import type {
   EventId,
@@ -118,15 +113,59 @@ export class WorkflowManager {
       fn: "You should not call this directly, call workflow.start instead";
       args: ObjectType<ArgsValidator>;
     },
-    ReturnsValidator extends Validator<unknown, "required", string>
-      ? Infer<ReturnsValidator>
-      : void
-  > {
-    return workflowMutation(
-      this.component,
-      workflow,
-      this.options?.workpoolOptions,
-    );
+    ReturnValueForOptionalValidator<ReturnsValidator>
+  >;
+  define<
+    ArgsValidator extends PropertyValidators,
+    ReturnsValidator extends Validator<unknown, "required", string> | void,
+  >(
+    workflow: Omit<
+      WorkflowDefinition<ArgsValidator, ReturnsValidator>,
+      "handler"
+    >,
+  ): {
+    handler: (
+      handler: (
+        step: WorkflowCtx,
+        args: ObjectType<ArgsValidator>,
+      ) => Promise<ReturnValueForOptionalValidator<ReturnsValidator>>,
+    ) => RegisteredMutation<
+      "internal",
+      {
+        fn: "You should not call this directly, call workflow.start instead";
+        args: ObjectType<ArgsValidator>;
+      },
+      ReturnValueForOptionalValidator<ReturnsValidator>
+    >;
+  };
+  define<
+    ArgsValidator extends PropertyValidators,
+    ReturnsValidator extends Validator<unknown, "required", string> | void,
+  >(
+    workflow:
+      | Omit<WorkflowDefinition<ArgsValidator, ReturnsValidator>, "handler">
+      | WorkflowDefinition<ArgsValidator, ReturnsValidator>,
+  ): unknown {
+    if ("handler" in workflow) {
+      return workflowMutation(
+        this.component,
+        workflow,
+        this.options?.workpoolOptions,
+      );
+    }
+    return {
+      handler: (
+        handler: (
+          step: WorkflowCtx,
+          args: ObjectType<ArgsValidator>,
+        ) => Promise<ReturnValueForOptionalValidator<ReturnsValidator>>,
+      ) =>
+        workflowMutation(
+          this.component,
+          { ...workflow, handler },
+          this.options?.workpoolOptions,
+        ),
+    };
   }
 
   /**