have internalMutation be a function on WorkflowManager

Author: ianmacartney

example/convex/eventTimeout.ts

@@ -10,7 +10,9 @@ import { v } from "convex/values";
 import { components, internal } from "./_generated/api";
 import { internalMutation } from "./_generated/server";
 
-const workflow = new WorkflowManager(components.workflow);
+const workflow = new WorkflowManager(components.workflow, {
+  internalMutation,
+});
 
 const approvalEvent = defineEvent({
   name: "approval",
@@ -31,7 +33,6 @@ const approvalEvent = defineEvent({
 export const eventTimeoutWorkflow = workflow.define({
   args: {},
   returns: v.string(),
-  internalMutation,
   handler: async (step): Promise<string> => {
     // 1. Create the event so we have an ID to pass to the timeout function.
     const eventId = await step.runMutation(

example/convex/example.ts

@@ -6,14 +6,15 @@ import { components } from "./_generated/api.js";
 import { vWorkflowId } from "@convex-dev/workflow";
 import { vResultValidator } from "@convex-dev/workpool";
 
-export const workflow = new WorkflowManager(components.workflow);
+export const workflow = new WorkflowManager(components.workflow, {
+  internalMutation,
+});
 
 export const myWorkflow = workflow
   .define({
     args: {
       location: v.string(),
     },
-    internalMutation,
     workpoolOptions: {
       retryActionsByDefault: true,
     },

src/client/index.ts

@@ -77,36 +77,16 @@ export type CallbackOptions = {
 export type WorkflowDefinition<
   ArgsValidator extends PropertyValidators,
   ReturnsValidator extends Validator<any, "required", any> | void = any,
-  DataModel extends GenericDataModel = GenericDataModel,
 > = {
   args?: ArgsValidator;
   returns?: ReturnsValidator;
   workpoolOptions?: WorkpoolRetryOptions;
-  /**
-   * Provide your app's `internalMutation` (from `_generated/server`) to get
-   * a fully typed `ctx` in `ctx.run()` handlers, with your data model's
-   * tables available on `ctx.db`. This also lets any custom middleware
-   * you've configured run around the workflow.
-   *
-   * ```ts
-   * import { internalMutation } from "./_generated/server";
-   * workflow.define({
-   *   internalMutation,
-   *   handler: async (ctx, args) => {
-   *     const user = await ctx.run(async (ctx) => {
-   *       return ctx.db.query("users").first(); // fully typed
-   *     });
-   *   },
-   * });
-   * ```
-   */
-  internalMutation?: MutationBuilder<DataModel, "internal">;
 };
 
 export type WorkflowHandler<
   ArgsValidator extends PropertyValidators,
   ReturnsValidator extends Validator<any, "required", any> | void,
-  DataModel extends GenericDataModel,
+  DataModel extends GenericDataModel = GenericDataModel,
 > = (
   step: WorkflowCtx<DataModel>,
   args: ObjectType<ArgsValidator>,
@@ -158,7 +138,9 @@ export function defineWorkflow<
   DM extends GenericDataModel = GenericDataModel,
 >(
   component: WorkflowComponent,
-  config: WorkflowDefinition<AV, RV, DM>,
+  config: WorkflowDefinition<AV, RV> & {
+    internalMutation?: MutationBuilder<DM, "internal">;
+  },
 ): {
   /**
    * Define the workflow handler function.
@@ -538,11 +520,27 @@ export async function cleanup(
   });
 }
 
-export class WorkflowManager {
+export class WorkflowManager<
+  DataModel extends GenericDataModel = GenericDataModel,
+> {
   constructor(
     public component: WorkflowComponent,
     public options?: {
-      workpoolOptions: WorkpoolOptions;
+      workpoolOptions?: WorkpoolOptions;
+      /**
+       * Provide your app's `internalMutation` (from `_generated/server`) to get
+       * a fully typed `ctx` in `step.run()` handlers, with your data model's
+       * tables available on `ctx.db`. This also lets any custom middleware
+       * you've configured run around the workflow.
+       *
+       * ```ts
+       * import { internalMutation } from "./_generated/server";
+       * const workflow = new WorkflowManager(components.workflow, {
+       *   internalMutation,
+       * });
+       * ```
+       */
+      internalMutation?: MutationBuilder<DataModel, "internal">;
     },
   ) {}
 
@@ -555,9 +553,8 @@ export class WorkflowManager {
   define<
     ArgsValidator extends PropertyValidators,
     ReturnsValidator extends Validator<unknown, "required", string> | void,
-    DataModel extends GenericDataModel = GenericDataModel,
   >(
-    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator, DataModel> & {
+    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator> & {
       handler: WorkflowHandler<ArgsValidator, ReturnsValidator, DataModel>;
     },
   ): RegisteredMutation<
@@ -568,9 +565,8 @@ export class WorkflowManager {
   define<
     ArgsValidator extends PropertyValidators,
     ReturnsValidator extends Validator<unknown, "required", string> | void,
-    DataModel extends GenericDataModel = GenericDataModel,
   >(
-    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator, DataModel>,
+    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator>,
   ): {
     /**
      * Define the workflow handler function.
@@ -602,21 +598,21 @@ export class WorkflowManager {
   define<
     ArgsValidator extends PropertyValidators,
     ReturnsValidator extends Validator<unknown, "required", string> | void,
-    DataModel extends GenericDataModel = GenericDataModel,
   >(
-    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator, DataModel> & {
+    workflow: WorkflowDefinition<ArgsValidator, ReturnsValidator> & {
       handler?: WorkflowHandler<ArgsValidator, ReturnsValidator, DataModel>;
     },
   ): unknown {
+    const withMutation = {
+      ...workflow,
+      internalMutation: this.options?.internalMutation,
+    };
     if (workflow.handler) {
       return workflowMutation(
         this.component,
-        workflow as WorkflowDefinition<
-          ArgsValidator,
-          ReturnsValidator,
-          DataModel
-        > & {
+        withMutation as WorkflowDefinition<ArgsValidator, ReturnsValidator> & {
           handler: WorkflowHandler<ArgsValidator, ReturnsValidator, DataModel>;
+          internalMutation?: MutationBuilder<DataModel, "internal">;
         },
         this.options?.workpoolOptions,
       );
@@ -628,7 +624,7 @@ export class WorkflowManager {
     return defineWorkflow<ArgsValidator, ReturnsValidator, DataModel>(
       this.component,
       {
-        ...workflow,
+        ...withMutation,
         workpoolOptions: {
           ...this.options?.workpoolOptions,
           ...workflow.workpoolOptions,

src/client/workflowContext.ts

@@ -116,17 +116,12 @@ export type WorkflowCtx<DataModel extends GenericDataModel = GenericDataModel> =
      * value will not be replayed.
      *
      * To get a fully typed `ctx` with your data model, provide your app's
-     * `internalMutation` in the workflow definition:
+     * `internalMutation` when creating the `WorkflowManager`:
      *
      * ```ts
      * import { internalMutation } from "./_generated/server";
-     * workflow.define({
+     * const workflow = new WorkflowManager(components.workflow, {
      *   internalMutation,
-     *   handler: async (ctx, args) => {
-     *     const user = await ctx.run(async (ctx) => {
-     *       return ctx.db.query("users").first(); // fully typed
-     *     });
-     *   },
      * });
      * ```
      *

src/client/workflowMutation.ts

@@ -6,6 +6,7 @@ import {
   internalMutationGeneric,
   makeFunctionReference,
   type GenericDataModel,
+  type MutationBuilder,
   type RegisteredMutation,
 } from "convex/server";
 import {
@@ -54,8 +55,9 @@ export function workflowMutation<
   DataModel extends GenericDataModel = GenericDataModel,
 >(
   component: WorkflowComponent,
-  registered: WorkflowDefinition<ArgsValidator, any, DataModel> & {
+  registered: WorkflowDefinition<ArgsValidator> & {
     handler: WorkflowHandler<ArgsValidator, any, DataModel>;
+    internalMutation?: MutationBuilder<DataModel, "internal">;
   },
   defaultWorkpoolOptions?: WorkpoolOptions,
   boundFn?: string,