Address some claude & codex suggestions. Swap to taking options rather than requiring subclassing

Author: VegetarianOrc

packages/nexus/src/index.ts

@@ -15,6 +15,7 @@ export {
 
 export {
   startWorkflow,
+  type TemporalOperationHandlerOptions,
   TemporalOperationHandler,
   TemporalOperationResult,
   type TemporalNexusClient,

packages/nexus/src/token.ts

@@ -1,6 +1,5 @@
 /**
- * OperationTokenType is used to identify the type of Operation token.
- * Currently, we only have one type of Operation token: WorkflowRun.
+ * Serializable token identifying a Nexus operation target.
  *
  * @internal
  * @hidden
@@ -28,11 +27,24 @@ export interface OperationToken {
   wid?: string;
 }
 
+/**
+ * An OperationToken that identifies a WorkflowRun operation.
+ *
+ * @internal
+ * @hidden
+ */
 export interface WorkflowRunOperationToken extends OperationToken {
   t: typeof OperationTokenType.WORKFLOW_RUN;
   wid: string;
 }
 
+/**
+ * OperationTokenType is used to identify the type of Operation token.
+ * Currently, we only have one type of Operation token: WorkflowRun.
+ *
+ * @internal
+ * @hidden
+ */
 export type OperationTokenType = (typeof OperationTokenType)[keyof typeof OperationTokenType];
 
 /**

packages/nexus/src/workflow-helpers.ts

@@ -279,16 +279,31 @@ export type TemporalOperationStartHandler<I, O> = (
   input: I
 ) => Promise<TemporalOperationResult<O>>;
 
+/**
+ * Options for customizing a {@link TemporalOperationHandler}.
+ *
+ * @experimental Nexus support in Temporal SDK is experimental.
+ */
+export interface TemporalOperationHandlerOptions {
+  cancelWorkflowRun?: (ctx: nexus.CancelOperationContext, client: Client, workflowId: string) => Promise<void>;
+}
+
 /**
  * A Nexus Operation implementation for operations that interact with Temporal.
  *
  * @experimental Nexus support in Temporal SDK is experimental.
  */
 export class TemporalOperationHandler<I, O> implements nexus.OperationHandler<I, O> {
-  constructor(readonly handler: TemporalOperationStartHandler<I, O>) {}
+  private readonly startHandler: TemporalOperationStartHandler<I, O>;
+  private readonly cancelWorkflowRunHandler: NonNullable<TemporalOperationHandlerOptions['cancelWorkflowRun']>;
+
+  constructor(options: { start: TemporalOperationStartHandler<I, O> } & TemporalOperationHandlerOptions) {
+    this.startHandler = options.start;
+    this.cancelWorkflowRunHandler = options.cancelWorkflowRun ?? defaultCancelWorkflowRun;
+  }
 
   async start(ctx: nexus.StartOperationContext, input: I): Promise<nexus.HandlerStartOperationResult<O>> {
-    const result = await this.handler(ctx, new TemporalNexusClientImpl(ctx), input);
+    const result = await this.startHandler(ctx, new TemporalNexusClientImpl(ctx), input);
     return result[toHandlerResult]();
   }
 
@@ -309,8 +324,14 @@ export class TemporalOperationHandler<I, O> implements nexus.OperationHandler<I,
     }
     switch (opToken.t) {
       case OperationTokenType.WORKFLOW_RUN:
-        assertWorkflowRunOperationToken(opToken);
-        await this.cancelWorkflowRun(ctx, opToken.wid);
+        try {
+          assertWorkflowRunOperationToken(opToken);
+        } catch (err) {
+          throw new nexus.HandlerError(nexus.HandlerErrorType.BAD_REQUEST, 'invalid workflow run operation token', {
+            cause: err,
+          });
+        }
+        await this.cancelWorkflowRunHandler(ctx, getClient(), opToken.wid);
         return;
       default:
         throw new nexus.HandlerError(
@@ -319,15 +340,8 @@ export class TemporalOperationHandler<I, O> implements nexus.OperationHandler<I,
         );
     }
   }
+}
 
-  /**
-   * Cancel the workflow backing this Nexus Operation.
-   *
-   * Override this method to customize workflow cancellation behavior.
-   *
-   * @experimental Nexus support in Temporal SDK is experimental.
-   */
-  protected async cancelWorkflowRun(_ctx: nexus.CancelOperationContext, workflowId: string): Promise<void> {
-    await getClient().workflow.getHandle(workflowId).cancel();
-  }
+async function defaultCancelWorkflowRun(_ctx: nexus.CancelOperationContext, client: Client, workflowId: string) {
+  await client.workflow.getHandle(workflowId).cancel();
 }

packages/test/src/test-nexus-temporal-operation.ts

@@ -49,8 +49,10 @@ type TemporalOpServiceHandlers = nexus.ServiceHandlerFor<typeof temporalOpServic
 type TemporalCancelOpServiceHandlers = nexus.ServiceHandlerFor<typeof temporalCancelOpService.operations>;
 
 function unusedTemporalOperationHandler<I, O>(): nexus.OperationHandler<I, O> {
-  return new temporalnexus.TemporalOperationHandler<I, O>(async () => {
-    throw new nexus.HandlerError('NOT_IMPLEMENTED', 'not used by this test');
+  return new temporalnexus.TemporalOperationHandler<I, O>({
+    async start() {
+      throw new nexus.HandlerError('NOT_IMPLEMENTED', 'not used by this test');
+    },
   });
 }
 
@@ -114,67 +116,68 @@ export async function blockingTargetWorkflow(): Promise<void> {
 // Tests
 
 test('TemporalOperationHandler infers correct output type from typed workflow function', async (t) => {
-  const _stringOp: nexus.OperationHandler<string, string> = new temporalnexus.TemporalOperationHandler(
-    async (_ctx, client, input: string) => {
+  const _stringOp: nexus.OperationHandler<string, string> = new temporalnexus.TemporalOperationHandler({
+    async start(_ctx, client, input: string) {
       return await client.startWorkflow(echoWorkflow, {
         args: [input],
         workflowId: 'test',
       });
-    }
-  );
+    },
+  });
 
   // @ts-expect-error - Output type should be string, not number
-  const _mismatchedOp: nexus.OperationHandler<string, number> = new temporalnexus.TemporalOperationHandler(
-    async (_ctx, client, input: string) => {
+  const _mismatchedOp: nexus.OperationHandler<string, number> = new temporalnexus.TemporalOperationHandler({
+    async start(_ctx, client, input: string) {
       return await client.startWorkflow(echoWorkflow, {
         args: [input],
         workflowId: 'test',
       });
-    }
-  );
+    },
+  });
 
-  const _syncOp: nexus.OperationHandler<string, string> = new temporalnexus.TemporalOperationHandler(
-    async (_ctx, _client, input: string) => {
+  const _syncOp: nexus.OperationHandler<string, string> = new temporalnexus.TemporalOperationHandler({
+    async start(_ctx, _client, input: string) {
       return temporalnexus.TemporalOperationResult.sync(input);
-    }
-  );
+    },
+  });
 
   const _explicitStringOp: nexus.OperationHandler<string, string> = new temporalnexus.TemporalOperationHandler<
     string,
     string
-  >(async (_ctx, client, input) => {
-    return await client.startWorkflow(echoWorkflow, {
-      args: [input],
-      workflowId: 'test',
-    });
+  >({
+    async start(_ctx, client, input) {
+      return await client.startWorkflow(echoWorkflow, {
+        args: [input],
+        workflowId: 'test',
+      });
+    },
   });
 
   // This test only checks for compile-time errors.
   t.pass();
 });
 
-test('TemporalOperationHandler cancel delegates to overridable cancelWorkflow', async (t) => {
+test('TemporalOperationHandler cancel delegates to provided cancelWorkflowRun handler', async (t) => {
   const { createWorker, registerNexusEndpoint } = helpers(t);
   const { endpointName } = await registerNexusEndpoint();
   const workflowId = randomUUID();
 
   let customCancelCalled = false;
 
-  class CustomCancelHandler extends temporalnexus.TemporalOperationHandler<string, void> {
-    async cancelWorkflowRun(_ctx: nexus.CancelOperationContext, workflowId: string): Promise<void> {
-      const handle = temporalnexus.getClient().workflow.getHandle(workflowId);
-      await handle.cancel();
-      customCancelCalled = true;
-    }
-  }
-
   const worker = await createWorker({
     nexusServices: [
       makeTemporalCancelOpServiceHandler({
-        blockingOp: new CustomCancelHandler(async (_ctx, client, workflowId) => {
-          return await client.startWorkflow(blockingTargetWorkflow, {
-            workflowId,
-          });
+        blockingOp: new temporalnexus.TemporalOperationHandler({
+          async start(_ctx, client, workflowId) {
+            return await client.startWorkflow(blockingTargetWorkflow, {
+              workflowId,
+            });
+          },
+          async cancelWorkflowRun(_ctx, client, workflowId) {
+            const handle = client.workflow.getHandle(workflowId);
+            await handle.cancel();
+            customCancelCalled = true;
+          },
         }),
       }),
     ],
@@ -208,14 +211,16 @@ test('TemporalOperationHandler cancel delegates to overridable cancelWorkflow',
   });
 });
 
-test('TemporalOperationHandler cancel rejects invalid operation token type', async (t) => {
-  class CustomCancelHandler extends temporalnexus.TemporalOperationHandler<void, void> {
-    async cancelWorkflowRun(_ctx: nexus.CancelOperationContext, _workflowId: string): Promise<void> {
-      throw new Error('cancelWorkflowRun should not be called');
-    }
-  }
+test('TemporalOperationHandler.cancel rejects invalid operation token type before invoking cancellation hooks', async (t) => {
+  const handler = new temporalnexus.TemporalOperationHandler({
+    async start(_ctx, _client, _input) {
+      return temporalnexus.TemporalOperationResult.sync(undefined);
+    },
 
-  const handler = new CustomCancelHandler(async () => temporalnexus.TemporalOperationResult.sync(undefined));
+    async cancelWorkflowRun(_ctx, _client, _workflowId) {
+      throw new Error('cancelWorkflowRun should not be called');
+    },
+  });
   const token = base64URLEncodeNoPadding(JSON.stringify({ t: 2, ns: 'test-namespace' }));
 
   const err = await asyncLocalStorage.run(
@@ -252,14 +257,18 @@ test('TemporalOperationHandler async and sync happy paths - caller workflow', as
   const worker = await createWorker({
     nexusServices: [
       makeTemporalOpServiceHandler({
-        asyncOp: new temporalnexus.TemporalOperationHandler<string, string>(async (_ctx, client, input) => {
-          return await client.startWorkflow(echoWorkflow, {
-            workflowId: randomUUID(),
-            args: [input],
-          });
+        asyncOp: new temporalnexus.TemporalOperationHandler<string, string>({
+          async start(_ctx, client, input) {
+            return await client.startWorkflow(echoWorkflow, {
+              workflowId: randomUUID(),
+              args: [input],
+            });
+          },
         }),
-        syncOp: new temporalnexus.TemporalOperationHandler<string, string>(async (_ctx, _client, input) => {
-          return temporalnexus.TemporalOperationResult.sync(input);
+        syncOp: new temporalnexus.TemporalOperationHandler<string, string>({
+          async start(_ctx, _client, input) {
+            return temporalnexus.TemporalOperationResult.sync(input);
+          },
         }),
       }),
     ],
@@ -285,16 +294,18 @@ test('TemporalOperationHandler rejects multiple async starts', async (t) => {
   const worker = await createWorker({
     nexusServices: [
       makeTemporalOpServiceHandler({
-        doubleStartOp: new temporalnexus.TemporalOperationHandler<string, void>(async (_ctx, client, input) => {
-          await client.startWorkflow(echoWorkflow, {
-            workflowId: randomUUID(),
-            args: [input],
-          });
-          await client.startWorkflow(echoWorkflow, {
-            workflowId: randomUUID(),
-            args: [input],
-          });
-          throw new nexus.HandlerError('INTERNAL', 'expected previous error to be thrown');
+        doubleStartOp: new temporalnexus.TemporalOperationHandler<string, void>({
+          async start(_ctx, client, input) {
+            await client.startWorkflow(echoWorkflow, {
+              workflowId: randomUUID(),
+              args: [input],
+            });
+            await client.startWorkflow(echoWorkflow, {
+              workflowId: randomUUID(),
+              args: [input],
+            });
+            throw new nexus.HandlerError('INTERNAL', 'expected previous error to be thrown');
+          },
         }),
       }),
     ],
@@ -326,8 +337,8 @@ test('TemporalOperationHandler allows retry after failed async start', async (t)
   const worker = await createWorker({
     nexusServices: [
       makeTemporalOpServiceHandler({
-        retryAfterFailedStartOp: new temporalnexus.TemporalOperationHandler<string, string>(
-          async (_ctx, client, workflowId) => {
+        retryAfterFailedStartOp: new temporalnexus.TemporalOperationHandler<string, string>({
+          async start(_ctx, client, workflowId) {
             try {
               await client.startWorkflow(blockingTargetWorkflow, {
                 workflowId,
@@ -345,8 +356,8 @@ test('TemporalOperationHandler allows retry after failed async start', async (t)
             throw new nexus.HandlerError('INTERNAL', 'Expected first workflow start to fail', {
               retryableOverride: false,
             });
-          }
-        ),
+          },
+        }),
       }),
     ],
   });
@@ -374,10 +385,12 @@ test('TemporalOperationHandler default cancelWorkflowRun cancels backing workflo
   const worker = await createWorker({
     nexusServices: [
       makeTemporalCancelOpServiceHandler({
-        blockingOp: new temporalnexus.TemporalOperationHandler<string, void>(async (_ctx, client, workflowId) => {
-          return await client.startWorkflow(blockingTargetWorkflow, {
-            workflowId,
-          });
+        blockingOp: new temporalnexus.TemporalOperationHandler<string, void>({
+          async start(_ctx, client, workflowId) {
+            return await client.startWorkflow(blockingTargetWorkflow, {
+              workflowId,
+            });
+          },
         }),
       }),
     ],