fix: remove input from step context

grutt · Jan 27, 2024 · dc03647 · dc03647
1 parent c6ffa83
commit dc03647
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 20 deletions.
diff --git a/typescript-sdk/examples/dag-worker.ts b/typescript-sdk/examples/dag-worker.ts
@@ -19,7 +19,7 @@ const workflow: Workflow = {
   steps: [
     {
       name: 'dag-step1',
-      run: async (input, ctx) => {
+      run: async (ctx) => {
         console.log('executed step1!');
         await sleep(5000);
         return { step1: 'step1' };
@@ -28,23 +28,23 @@ const workflow: Workflow = {
     {
       name: 'dag-step2',
       parents: ['dag-step1'],
-      run: (input, ctx) => {
+      run: (ctx) => {
         console.log('executed step2!');
         return { step2: 'step2' };
       },
     },
     {
       name: 'dag-step3',
       parents: ['dag-step1', 'dag-step2'],
-      run: (input, ctx) => {
+      run: (ctx) => {
         console.log('executed step3!');
         return { step3: 'step3' };
       },
     },
     {
       name: 'dag-step4',
       parents: ['dag-step1', 'dag-step3'],
-      run: (input, ctx) => {
+      run: (ctx) => {
         console.log('executed step4!');
         return { step4: 'step4' };
       },

diff --git a/typescript-sdk/examples/simple-worker.ts b/typescript-sdk/examples/simple-worker.ts
@@ -17,19 +17,20 @@ const workflow: Workflow = {
   steps: [
     {
       name: 'step1',
-      run: async (input, ctx) => {
-        console.log('starting step1!');
+      run: async (ctx) => {
+        console.log('starting step1 with the following input', ctx.workflowInput());
+        console.log('waiting 5 seconds...');
         await sleep(5000);
         console.log('executed step1!');
-        return { step1: 'step1' };
+        return { step1: 'step1 results!' };
       },
     },
     {
       name: 'step2',
       parents: ['step1'],
-      run: (input, ctx) => {
-        console.log('executed step2!', ctx.workflowInput());
-        return { step2: 'step2' };
+      run: (ctx) => {
+        console.log('executed step2 after step1 returned ', ctx.stepOutput('step1'));
+        return { step2: 'step2 results!' };
       },
     },
   ],

diff --git a/typescript-sdk/hatchet/clients/worker/worker.test.ts b/typescript-sdk/hatchet/clients/worker/worker.test.ts
@@ -7,6 +7,7 @@ import { never } from 'zod';
 import sleep from '@util/sleep';
 import { ChannelCredentials } from 'nice-grpc';
 import { Worker } from './worker';
+import { Context } from '@hatchet/step';
 
 type AssignActionMock = AssignedAction | Error;
 
@@ -19,7 +20,7 @@ const mockStart: AssignActionMock = {
   stepRunId: 'runStep1',
   actionId: 'action1',
   actionType: ActionType.START_STEP_RUN,
-  actionPayload: '{"input": {"data": 1}}',
+  actionPayload: JSON.stringify('{"input": {"data": 1}}'),
 };
 
 const mockCancel: AssignActionMock = {
@@ -63,8 +64,7 @@ describe('Worker', () => {
         steps: [
           {
             name: 'step1',
-            run: (input: any, ctx: any) => {
-              console.log('step1', input, ctx);
+            run: (ctx: any) => {
               return { test: 'test' };
             },
           },
@@ -106,7 +106,7 @@ describe('Worker', () => {
       await sleep(100);
 
       expect(startSpy).toHaveBeenCalledTimes(1);
-      expect(startSpy).toHaveBeenCalledWith({ data: 1 }, expect.anything());
+
       expect(getActionEventSpy).toHaveBeenNthCalledWith(
         2,
         expect.anything(),
@@ -139,7 +139,6 @@ describe('Worker', () => {
       await sleep(100);
 
       expect(startSpy).toHaveBeenCalledTimes(1);
-      expect(startSpy).toHaveBeenCalledWith({ data: 1 }, expect.anything());
       expect(getActionEventSpy).toHaveBeenNthCalledWith(
         2,
         expect.anything(),

diff --git a/typescript-sdk/hatchet/clients/worker/worker.ts b/typescript-sdk/hatchet/clients/worker/worker.ts
@@ -88,7 +88,7 @@ export class Worker {
     }
 
     const run = async () => {
-      return step(context.workflowInput(), context);
+      return step(context);
     };
 
     const success = (result: any) => {

diff --git a/typescript-sdk/hatchet/step.ts b/typescript-sdk/hatchet/step.ts
@@ -8,11 +8,17 @@ export const CreateStepSchema = z.object({
 
 export type NextStep = { [key: string]: string };
 
-export class Context<T = any> {
-  data: T | any;
+interface ContextData<T = unknown> {
+  input: T;
+  parents: Record<string, any>;
+  triggered_by_event: string;
+}
+
+export class Context<T = unknown> {
+  data: ContextData<T>;
   constructor(payload: string) {
     try {
-      this.data = JSON.parse(payload);
+      this.data = JSON.parse(JSON.parse(payload));
     } catch (e: any) {
       throw new HatchetError(`Could not parse payload: ${e.message}`);
     }
@@ -38,5 +44,5 @@ export class Context<T = any> {
 }
 
 export interface CreateStep<T> extends z.infer<typeof CreateStepSchema> {
-  run: (input: T, ctx: Context) => Promise<NextStep> | NextStep | void;
+  run: (ctx: Context) => Promise<NextStep> | NextStep | void;
 }