@mastra/inngest 0.11.0 → 0.11.1

package/src/index.ts

@@ -33,13 +33,17 @@ export type InngestEngineType = {
 
 export function serve({ mastra, inngest }: { mastra: Mastra; inngest: Inngest }): ReturnType<typeof inngestServe> {
   const wfs = mastra.getWorkflows();
-  const functions = Object.values(wfs).flatMap(wf => {
-    if (wf instanceof InngestWorkflow) {
-      wf.__registerMastra(mastra);
-      return wf.getFunctions();
-    }
-    return [];
-  });
+  const functions = Array.from(
+    new Set(
+      Object.values(wfs).flatMap(wf => {
+        if (wf instanceof InngestWorkflow) {
+          wf.__registerMastra(mastra);
+          return wf.getFunctions();
+        }
+        return [];
+      }),
+    ),
+  );
   return inngestServe({
     client: inngest,
     functions,
@@ -94,8 +98,14 @@ export class InngestRun<
     while (runs?.[0]?.status !== 'Completed' || runs?.[0]?.event_id !== eventId) {
       await new Promise(resolve => setTimeout(resolve, 1000));
       runs = await this.getRuns(eventId);
-      if (runs?.[0]?.status === 'Failed' || runs?.[0]?.status === 'Cancelled') {
+      if (runs?.[0]?.status === 'Failed') {
         throw new Error(`Function run ${runs?.[0]?.status}`);
+      } else if (runs?.[0]?.status === 'Cancelled') {
+        const snapshot = await this.#mastra?.storage?.loadWorkflowSnapshot({
+          workflowName: this.workflowId,
+          runId: this.runId,
+        });
+        return { output: { result: { steps: snapshot?.context, status: 'canceled' } } };
       }
     }
     return runs?.[0];
@@ -108,6 +118,30 @@ export class InngestRun<
     });
   }
 
+  async cancel() {
+    await this.inngest.send({
+      name: `cancel.workflow.${this.workflowId}`,
+      data: {
+        runId: this.runId,
+      },
+    });
+
+    const snapshot = await this.#mastra?.storage?.loadWorkflowSnapshot({
+      workflowName: this.workflowId,
+      runId: this.runId,
+    });
+    if (snapshot) {
+      await this.#mastra?.storage?.persistWorkflowSnapshot({
+        workflowName: this.workflowId,
+        runId: this.runId,
+        snapshot: {
+          ...snapshot,
+          status: 'canceled' as any,
+        },
+      });
+    }
+  }
+
   async start({
     inputData,
   }: {
@@ -147,7 +181,9 @@ export class InngestRun<
       result.error = new Error(result.error);
     }
 
-    this.cleanup?.();
+    if (result.status !== 'suspended') {
+      this.cleanup?.();
+    }
     return result;
   }
 
@@ -432,23 +468,27 @@ export class InngestWorkflow<
 
     this.runs.set(runIdToUse, run);
 
-    await this.mastra?.getStorage()?.persistWorkflowSnapshot({
-      workflowName: this.id,
-      runId: runIdToUse,
-      snapshot: {
+    const workflowSnapshotInStorage = await this.getWorkflowRunExecutionResult(runIdToUse);
+
+    if (!workflowSnapshotInStorage) {
+      await this.mastra?.getStorage()?.persistWorkflowSnapshot({
+        workflowName: this.id,
         runId: runIdToUse,
-        status: 'pending',
-        value: {},
-        context: {},
-        activePaths: [],
-        serializedStepGraph: this.serializedStepGraph,
-        suspendedPaths: {},
-        result: undefined,
-        error: undefined,
-        // @ts-ignore
-        timestamp: Date.now(),
-      },
-    });
+        snapshot: {
+          runId: runIdToUse,
+          status: 'pending',
+          value: {},
+          context: {},
+          activePaths: [],
+          serializedStepGraph: this.serializedStepGraph,
+          suspendedPaths: {},
+          result: undefined,
+          error: undefined,
+          // @ts-ignore
+          timestamp: Date.now(),
+        },
+      });
+    }
 
     return run;
   }
@@ -458,8 +498,12 @@ export class InngestWorkflow<
       return this.function;
     }
     this.function = this.inngest.createFunction(
-      // @ts-ignore
-      { id: `workflow.${this.id}`, retries: this.retryConfig?.attempts ?? 0 },
+      {
+        id: `workflow.${this.id}`,
+        // @ts-ignore
+        retries: this.retryConfig?.attempts ?? 0,
+        cancelOn: [{ event: `cancel.workflow.${this.id}` }],
+      },
       { event: `workflow.${this.id}` },
       async ({ event, step, attempt, publish }) => {
         let { inputData, runId, resume } = event.data;
@@ -508,6 +552,7 @@ export class InngestWorkflow<
           retryConfig: this.retryConfig,
           runtimeContext: new RuntimeContext(), // TODO
           resume,
+          abortController: new AbortController(),
         });
 
         return { result, runId };
@@ -630,7 +675,7 @@ export function createStep<
       outputSchema: z.object({
         text: z.string(),
       }),
-      execute: async ({ inputData, [EMITTER_SYMBOL]: emitter, runtimeContext }) => {
+      execute: async ({ inputData, [EMITTER_SYMBOL]: emitter, runtimeContext, abortSignal, abort }) => {
         let streamPromise = {} as {
           promise: Promise<string>;
           resolve: (value: string) => void;
@@ -656,8 +701,13 @@ export function createStep<
           onFinish: result => {
             streamPromise.resolve(result.text);
           },
+          abortSignal,
         });
 
+        if (abortSignal.aborted) {
+          return abort();
+        }
+
         for await (const chunk of fullStream) {
           switch (chunk.type) {
             case 'text-delta':
@@ -815,6 +865,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
       delay?: number;
     };
     runtimeContext: RuntimeContext;
+    abortController: AbortController;
   }): Promise<TOutput> {
     await params.emitter.emit('watch-v2', {
       type: 'start',
@@ -914,6 +965,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     resume,
     prevOutput,
     emitter,
+    abortController,
     runtimeContext,
   }: {
     workflowId: string;
@@ -927,6 +979,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     };
     prevOutput: any;
     emitter: Emitter;
+    abortController: AbortController;
     runtimeContext: RuntimeContext;
   }): Promise<StepResult<any, any, any, any>> {
     return super.executeStep({
@@ -938,6 +991,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
       resume,
       prevOutput,
       emitter,
+      abortController,
       runtimeContext,
     });
   }
@@ -966,17 +1020,12 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     resume,
     prevOutput,
     emitter,
+    abortController,
     runtimeContext,
   }: {
     step: Step<string, any, any>;
     stepResults: Record<string, StepResult<any, any, any, any>>;
-    executionContext: {
-      workflowId: string;
-      runId: string;
-      executionPath: number[];
-      suspendedPaths: Record<string, number[]>;
-      retryConfig: { attempts: number; delay: number };
-    };
+    executionContext: ExecutionContext;
     resume?: {
       steps: string[];
       resumePayload: any;
@@ -984,6 +1033,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     };
     prevOutput: any;
     emitter: Emitter;
+    abortController: AbortController;
     runtimeContext: RuntimeContext;
   }): Promise<StepResult<any, any, any, any>> {
     const startedAt = await this.inngestStep.run(
@@ -1016,6 +1066,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
           type: 'step-start',
           payload: {
             id: step.id,
+            status: 'running',
           },
         });
 
@@ -1091,6 +1142,8 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
               payload: {
                 id: step.id,
                 status: 'failed',
+                error: result?.error,
+                payload: prevOutput,
               },
             });
 
@@ -1128,6 +1181,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
                 type: 'step-suspended',
                 payload: {
                   id: step.id,
+                  status: 'suspended',
                 },
               });
 
@@ -1187,6 +1241,15 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
             eventTimestamp: Date.now(),
           });
 
+          await emitter.emit('watch-v2', {
+            type: 'step-result',
+            payload: {
+              id: step.id,
+              status: 'success',
+              output: result?.result,
+            },
+          });
+
           await emitter.emit('watch-v2', {
             type: 'step-finish',
             payload: {
@@ -1241,6 +1304,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
           engine: {
             step: this.inngestStep,
           },
+          abortSignal: abortController.signal,
         });
         const endedAt = Date.now();
 
@@ -1307,8 +1371,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
           type: 'step-suspended',
           payload: {
             id: step.id,
-            status: execResults.status,
-            output: execResults.status === 'success' ? execResults?.output : undefined,
+            ...execResults,
           },
         });
       } else {
@@ -1316,8 +1379,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
           type: 'step-result',
           payload: {
             id: step.id,
-            status: execResults.status,
-            output: execResults.status === 'success' ? execResults?.output : undefined,
+            ...execResults,
           },
         });
 
@@ -1396,6 +1458,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     resume,
     executionContext,
     emitter,
+    abortController,
     runtimeContext,
   }: {
     workflowId: string;
@@ -1417,6 +1480,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
     };
     executionContext: ExecutionContext;
     emitter: Emitter;
+    abortController: AbortController;
     runtimeContext: RuntimeContext;
   }): Promise<StepResult<any, any, any, any>> {
     let execResults: any;
@@ -1429,6 +1493,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
                 runId,
                 mastra: this.mastra!,
                 runtimeContext,
+                runCount: -1,
                 inputData: prevOutput,
                 getInitData: () => stepResults?.input as any,
                 getStepResult: (step: any) => {
@@ -1447,10 +1512,14 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
                 // TODO: this function shouldn't have suspend probably?
                 suspend: async (_suspendPayload: any) => {},
                 bail: () => {},
+                abort: () => {
+                  abortController.abort();
+                },
                 [EMITTER_SYMBOL]: emitter,
                 engine: {
                   step: this.inngestStep,
                 },
+                abortSignal: abortController.signal,
               });
               return result ? index : null;
               // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -1482,6 +1551,7 @@ export class InngestExecutionEngine extends DefaultExecutionEngine {
             executionSpan: executionContext.executionSpan,
           },
           emitter,
+          abortController,
           runtimeContext,
         }),
       ),